AI Agent Pipeline Design: Splitting the Execution Loop into 8 Stages
Most people build their first LLM agent the same way: take a message, build a prompt, call the model, run the tools, feed the results back, repeat. Pack that into one function and it works beautifully at around 200 lines. Then the requirements arrive. You need session summaries. You need to trim when tokens overflow. You need hooks in the middle. Some agents must skip tool execution entirely. At some point that function is 1,500 lines and nobody will touch it.
The fix is a familiar one: split the loop into stages. This article covers a design that breaks the agent loop into eight stages, and the problems you inevitably hit when you do.
What the Eight Stages Are
Logically, one agent turn flows like this.
- context — settle the workspace, config files, and user information
- history — load prior conversation and the session summary
- prompt — assemble the system prompt
- think — call the LLM
- act — execute whatever tools the model requested
- observe — feed tool results back into the conversation
- memory — persist what this turn should leave behind
- summarize — summarize the session and wrap up
Here's the first trap. These eight do not all run on every iteration. context, history, and prompt need to happen once at the start of a turn; think, act, and observe must keep cycling while the model asks for more tools; memory and summarize run once after the loop. So the real structure has three zones.
type Pipeline struct {
setup []Stage // once at turn start (context/history/prompt)
iteration []Stage // repeats per tool loop (think/act/observe)
finalize []Stage // once at turn end (memory/summarize)
}
func NewDefaultPipeline(deps Deps) *Pipeline {
return &Pipeline{
setup: []Stage{NewContextStage(&deps)},
iteration: []Stage{
NewThinkStage(&deps),
NewPruneStage(&deps), // trim when over the token budget
NewToolStage(&deps), // act
NewObserveStage(&deps),
NewCheckpointStage(&deps),
},
finalize: []Stage{NewFinalizeStage(&deps)},
}
}Keep the Stage Interface Tiny
Two methods are enough: a name and an execute. Then let only the stages that need flow control optionally implement one more interface.
type StageResult int
const (
Continue StageResult = iota // proceed to the next stage
BreakLoop // exit the iteration loop normally
AbortRun // abort the entire run
)
type Stage interface {
Name() string // for logging and tracing
Execute(ctx context.Context, st *RunState) error
}
// Only stages that must steer the loop implement this.
// Stages that don't are treated as Continue by the pipeline.
type StageWithResult interface {
Stage
Result() StageResult
}This is Go's optional interface pattern. Force every stage to implement Result() and most of them end up with a single boilerplate return Continue. Let the ones that need it opt in, and have the pipeline check with a type assertion.
for _, stage := range p.iteration {
if err := stage.Execute(ctx, state); err != nil {
return nil, fmt.Errorf("iter %d %s: %w", state.Iteration, stage.Name(), err)
}
// An abort signal takes effect immediately — skip the remaining stages
if swr, ok := stage.(StageWithResult); ok && swr.Result() == AbortRun {
state.ExitCode = AbortRun
break
}
}The Core Rule: Stateless Stages, State in One Place
This is the center of the design. Stage objects hold no state; all mutable state lives in a single struct passed by pointer.
// Shared mutable state for one run. Passed by pointer through every stage.
type RunState struct {
// Set at run start, immutable thereafter
Input *RunInput
Model string
Provider Provider
// Message buffer that several stages read and write
Messages *MessageBuffer
// Per-stage substates — the name documents ownership
Context ContextState
Think ThinkState
Tool ToolState
Observe ObserveState
Iteration int
ExitCode StageResult
}The reason matters: if a stage keeps state in its own fields, you can't run the same pipeline instance twice. The moment you handle two conversations concurrently, values from one run leak into another. Keep stages stateless and you can build the pipeline once and create a fresh RunState per request.
Grouping substates under stage names like Think and Tool is deliberate too. Flatten those fields and six months later you can't trace who writes what. The name documents ownership.
Break Coupling Between Stages with Callbacks
The context stage resolves a workspace, reads files, assembles a prompt, and loads history. Implement all of that inside the stage and your pipeline package now depends on your store, your filesystem, and your prompt builder. Instead, inject dependencies as function fields.
type Deps struct {
ResolveWorkspace func(ctx context.Context, in *RunInput) (*Workspace, error)
LoadContextFiles func(ctx context.Context, userID string) ([]ContextFile, bool)
LoadHistory func(ctx context.Context, key string) ([]Message, string)
BuildMessages func(ctx context.Context, in *RunInput, h []Message, sum string) ([]Message, error)
CallLLM func(ctx context.Context, st *RunState, req ChatRequest) (*ChatResponse, error)
Config Config
}
// Anything not injected is simply skipped — partial assembly becomes possible in tests
if s.deps.LoadHistory != nil && state.Input.SessionKey != "" {
history, summary := s.deps.LoadHistory(ctx, state.Input.SessionKey)
state.Messages.SetHistory(history)
state.Context.Summary = summary
}The nil checks look untidy but they earn their keep. You can build a pipeline with only the callbacks you care about, so verifying the think stage doesn't require standing up a whole store mock.
What the think Stage Actually Does
The stage that looks simplest is the messiest, because LLM calls fail in a variety of ways.
func (s *ThinkStage) Execute(ctx context.Context, st *RunState) error {
s.result = Continue
resp, err := s.deps.CallLLM(ctx, st, buildRequest(st))
if err != nil {
return fmt.Errorf("llm call: %w", err)
}
// Accumulate usage — always include cache tokens in the tally
if resp.Usage != nil {
st.Think.TotalUsage.PromptTokens += resp.Usage.PromptTokens
st.Think.TotalUsage.CompletionTokens += resp.Usage.CompletionTokens
st.Think.TotalUsage.CacheReadTokens += resp.Usage.CacheReadTokens
}
// Truncation: retry only when tool call arguments were cut off.
// Text-only truncation with no tool calls is a long answer, not an error.
if resp.FinishReason == "length" && len(resp.ToolCalls) > 0 {
st.Think.TruncRetries++
if st.Think.TruncRetries >= maxTruncRetries {
s.result = AbortRun
return nil
}
st.Messages.AppendPending(hintMessage())
return nil // retry on the next iteration
}
// No tool calls means this is the final answer — end the loop
if len(resp.ToolCalls) == 0 {
s.result = BreakLoop
return nil
}
st.Messages.AppendPending(assistantMessage(resp))
return nil
}Note that it distinguishes two kinds of truncation. Cut-off tool call arguments leave broken JSON you can't execute, so you retry. Plain text cut off with no tool calls just means the answer was long, and it should go to the user as-is. Skip that distinction and you'll retry pointlessly on every long response.
Treating BreakLoop and AbortRun Differently
There's a reason for two flow-control signals.
- BreakLoop (normal completion) exits after finishing every remaining stage in this iteration. The observe stage still has to collect the final response and the checkpoint stage still has to persist state. Leave early and the answer meant for the user disappears.
- AbortRun exits immediately. The state is already unrecoverable, so running the remaining stages is pointless or unsafe.
And the finalize zone must run even on a cancelled context. If the user cancelled midway, the conversation so far still needs to be saved.
// Run finalize stages on a context with cancellation detached
finalizeCtx := context.WithoutCancel(ctx)
for _, stage := range p.finalize {
if err := stage.Execute(finalizeCtx, state); err != nil {
// A finalize failure logs and still returns a result
slog.Warn("finalize stage error", "stage", stage.Name(), "err", err)
}
}What This Structure Actually Buys You
- Adding a stage is one line in a list. Permission checks, cost guards, prompt-injection defenses — none of them require touching existing code.
- Tracing comes free. Every stage has a
Name(), so recording duration and errors per stage is written once in the pipeline. "Which stage is slow?" becomes answerable immediately. - Stages are unit-testable. Hand-build a
RunState, run one stage, inspect the result. - You can assemble different pipelines per agent type. A simple response agent with no tools just gets assembled without the tool stage.
Summary
- The agent loop splits into setup / iteration / finalize — don't lay all eight stages out in a straight line
- Stages are stateless; mutable state lives in one
RunStatepassed by pointer - Flow control goes through an optional interface — only the stages that need it implement it
- BreakLoop finishes the iteration; AbortRun leaves immediately
- Break external dependencies with injected callbacks to keep the pipeline package light
- finalize has to run even on a cancelled context
The real payoff of a pipeline isn't performance, it's changeability. Six months from now, when a new requirement lands, can you add one stage and slot it into a list instead of reading a 1,500-line function? That difference is the whole thing.
댓글
댓글 쓰기