Middleware chain and planning tool

Overview

Introduce the extension architecture every remaining harness feature is built on: an ordered, data-driven middleware chain with BeforeRun, WrapModelCall, WrapToolCall, ToolProvider and PromptContributor hooks. Refactor the phase-5 loop onto it without changing observable behaviour, then prove the chain with two small middlewares: write_todos planning (state field + tool + prompt fragment) and repeat-call detection moved out of the loop.

Requirements

Functional

• agent.Middleware is a marker interface; capability interfaces are optional and discovered by type assertion at chain construction (not per call).
• Chain order is produced by one constructor from the agent’s configuration; there is no merge algorithm and no user-supplied ordering.
• WrapModelCall is an onion: a middleware may modify the request for one call, retry, or replace the response. It must not mutate persisted history.
• WrapToolCall is an onion around tool.View.Call.
• ToolProvider middlewares contribute tools; tool order presented to the model follows chain order.
• **The system prompt is assembled inside the chain, at model-call time.** `PromptContributor.SystemBlocks(s *State) []provider.Block` returns blocks, not a string, and the base model handler concatenates the agent's own system blocks followed by each contributor's blocks in chain order to produce `Request.System`. Assembly happens on every model call (it is cheap and must reflect current state), never once before the loop. This is what makes the cache-breakpoint and memory-ordering constraints expressible at all: `PromptCacheMiddleware` (phase 10) sets `CacheControl` on a block **by index** in that assembled slice, and `MemoryMiddleware` (phase 11) contributes blocks that land after the breakpoint index because it sits later in the chain. Ordering assertions are made on the breakpoint index, not on a slice of strings.
• BeforeRun runs once per run, before the first model call, and may return state updates.
• TodoMiddleware: write_todos(todos) whole-list replace, state stored on the run (persisted so a resumed run keeps the plan), rendered into the system prompt, never echoed back as a message.
• RepeatCallMiddleware: hashes tool name + normalized args; on the third identical call within a run, returns a nudge tool result instead of executing.

Non-functional

• Refactoring the loop must not change any phase-5 or phase-6 test outcome; those tests are the regression gate.
• Zero allocation-heavy work per delta: the chain wraps model calls and tool calls only, never individual stream events.
• Middleware private state lives in `State.Private map[string]json.RawMessage` keyed by middleware name. Each middleware owns a typed struct and marshals/unmarshals it through two small accessors (`loadX(s) (X, bool, error)` / `storeX(s, X) error`), so persisting run state to `runs.state_json` is a plain marshal of `State` and a resumed run rebuilds typed state instead of silently getting `map[string]any` where it expected a struct. A key that is missing on resume is an explicit, handled condition: the accessor returns `found=false`, the middleware initializes its default, and records a span event `middleware.state_missing` with `hango.middleware.name` — never a silent zero value. Typed top-level fields exist only for cross-middleware contracts (`Todos`).
• Per-batch scratch state (the mutated-path set used by the filesystem middleware in phase 8, the repeat-call counters) is explicitly cleared at the end of each tool-call batch by the loop, so a long run cannot accumulate it and a resumed run does not inherit a stale batch.

Architecture

// internal/agent/middleware.go
type Middleware interface{ Name() string }

type State struct {
	RunID, SessionID, AgentID, TenantID uuid.UUID
	Messages []provider.Message // the loop's working history (persisted separately)
	Todos    []Todo
	// Red team: F10 — per-middleware typed state, marshalled by its owner.
	Private  map[string]json.RawMessage // keyed by middleware name
}

type ModelRequest struct {
	Req   *provider.Request
	State *State
}

type ModelResponse struct {
	Message    provider.Message
	StopReason provider.StopReason
	Usage      provider.Usage
}

type ModelHandler func(context.Context, *ModelRequest) (*ModelResponse, error)

type ModelWrapper interface {
	WrapModelCall(ctx context.Context, req *ModelRequest, next ModelHandler) (*ModelResponse, error)
}

type ToolCallRequest struct {
	Call  provider.ToolUseBlock
	State *State
}

type ToolHandler func(context.Context, *ToolCallRequest) (tool.Result, error)

type ToolWrapper interface {
	WrapToolCall(ctx context.Context, req *ToolCallRequest, next ToolHandler) (tool.Result, error)
}

type RunHook interface {
	BeforeRun(ctx context.Context, s *State) error
}

type ToolProvider interface {
	ProvideTools(ctx context.Context, s *State) []tool.Tool
}

// Red team: F6 — contributors return blocks so cache breakpoints can be placed by index.
type PromptContributor interface {
	SystemBlocks(s *State) []provider.Block
}

// internal/agent/chain.go
type Chain struct{} // built once per run from an ordered []Middleware

func BuildChain(mws []Middleware) *Chain
func (c *Chain) ModelCall(ctx context.Context, req *ModelRequest, base ModelHandler) (*ModelResponse, error)
func (c *Chain) ToolCall(ctx context.Context, req *ToolCallRequest, base ToolHandler) (tool.Result, error)
func (c *Chain) BeforeRun(ctx context.Context, s *State) error
func (c *Chain) Tools(ctx context.Context, s *State) []tool.Tool

// Red team: F6 — called by the base model handler on every model call, inside ModelCall.
func (c *Chain) SystemBlocks(def store.Agent, s *State) []provider.Block

// internal/agent/mwstack.go — the single place order is decided
func DefaultStack(cfg StackConfig) []Middleware
// order: Skills -> Filesystem -> Subagent -> Summarization -> RepairToolCalls ->
//        Todo -> RepeatCall -> PromptCache -> Memory -> HumanInTheLoop -> ToolAllowlist
// (later phases fill slots; this phase ships RepairToolCalls, Todo, RepeatCall, ToolAllowlist)

// internal/agent/mwtodo.go
type Todo struct {
	Content string `json:"content"`
	Status  string `json:"status"` // pending|in_progress|completed
}

type TodoMiddleware struct{} // ToolProvider + PromptContributor

-- migrations/00005_run_state.sql
alter table runs add column state_json jsonb not null default '{}'::jsonb;
-- holds Todos and serializable middleware state so a resumed or approved run keeps its plan

Instrumentation

• Chain construction is recorded once per run as a span event agent.middleware_stack on the invoke_agent span, with the ordered middleware names as an attribute — so “which stack ran” is never a guess.
• Each WrapModelCall and WrapToolCall middleware that changes behaviour records a span event on the enclosing turn span: middleware.request_rewritten, middleware.short_circuited, middleware.retried, each with hango.middleware.name and a reason.
• BeforeRun hooks share one INTERNAL span agent.before_run with a child span per hook (hango.middleware.name), so a slow load-once hook (memory, skills in later phases) is immediately visible.
• Todo state is recorded as a span event agent.todos_updated with counts by status (never the content, which can be user data) and mirrored into langfuse.observation.metadata.todos_total/todos_completed.
• Repeat-call blocks emit span event middleware.repeat_call_blocked and increment hango.middleware.repeat_blocked (attribute gen_ai.tool.name).
• Metrics: hango.middleware.before_run.duration histogram (attribute hango.middleware.name), hango.middleware.tool_short_circuits counter.

Related Code Files

Create

• internal/agent/middleware.go, internal/agent/chain.go, internal/agent/chain_test.go.
• internal/agent/mwstack.go, internal/agent/mwstack_test.go (asserts the exact order and the tool order the model sees).
• internal/agent/mwtodo.go, internal/agent/mwtodo_test.go.
• internal/agent/mwrepeat.go, internal/agent/mwrepeat_test.go.
• internal/agent/mwallowlist.go — enforces the agent’s tool allowlist last so no middleware can re-add a banned tool.
• internal/agent/state.go — State, persistence helpers to/from runs.state_json.
• migrations/00005_run_state.sql.
• docs/learning/07-middleware-architecture.md.

Modify

• internal/agent/loop.go — model calls go through Chain.ModelCall, tool calls through Chain.ToolCall, the system prompt assembled inside the base model handler from Chain.SystemBlocks, tools from Chain.Tools merged with the registry view.
• internal/agent/repair.go — becomes RepairToolCallsMiddleware implementing BeforeRun.
• internal/run/worker.go — build the stack from agent config, load/save State around the run.
• internal/store/pg/run.go — read/write state_json.
• internal/config/config.go — HANGO_REPEAT_CALL_THRESHOLD.
• docs/system-architecture.md, docs/learning/README.md.

Implementation Steps

1. Write internal/agent/middleware.go with the interfaces above and internal/agent/state.go with State, MarshalState/UnmarshalState for runs.state_json, and the generic private-state accessors (LoadPrivate[T](s, name) (T, bool, error) / StorePrivate[T](s, name, v) error) that every middleware uses instead of touching the map.
2. Write chain.go: at BuildChain, partition the slice once into []ModelWrapper, []ToolWrapper, []RunHook, []ToolProvider, []PromptContributor by type assertion. ModelCall composes wrappers so index 0 is outermost, and the base handler performs the provider call plus accumulation.
3. Write mwstack.go with DefaultStack, the single ordered constructor. Document each slot with the reason it sits there: allowlist last so nothing can re-add a banned tool; prompt cache before memory so memory edits do not invalidate the cache prefix; summarization after filesystem so it sees already-offloaded tool results. Leave explicit nil-skipped slots for unimplemented phases.
4. Refactor loop.go: replace the direct provider call with chain.ModelCall and the direct Tools.Call with chain.ToolCall. Prompt assembly moves into the base model handler as req.System = chain.SystemBlocks(def, state), so it runs inside the onion on every model call and every WrapModelCall middleware sees (and may rewrite) the assembled blocks. Run chain.BeforeRun once before the loop. Clear per-batch scratch state at each batch end. Do not change event emission.
5. Convert the phase-5 repair function into RepairToolCallsMiddleware with BeforeRun, keeping its tests.
6. Write mwallowlist.go: a ToolWrapper that rejects any call whose name is not in the agent’s effective allowlist with an error tool result, plus a filter applied to the tool list handed to the provider.
7. Write mwtodo.go: a ToolProvider exposing write_todos whose description carries the policy text (when to use, when not to, exactly one in_progress, marking complete is not an answer), a handler that validates statuses and replaces State.Todos wholesale, and a PromptContributor returning one text block with the current list. Todos are never appended to State.Messages.
8. Write mwrepeat.go: a ToolWrapper keeping a typed per-run counter struct in State.Private through the accessors; hash is fnv64(name + canonicalJSON(args)); at the threshold return tool.Errorf("this exact call was already made N times; use the previous result or change approach") without executing.
9. Persist State to runs.state_json at every turn boundary so a cancelled or approval-suspended run resumes with its plan.
10. Tests: chain order test (a recording middleware asserts outermost/innermost); system-block assembly order asserted by block index, with a placeholder cache-breakpoint index to pin the contract phases 10 and 11 depend on ; a model wrapper that retries once; a tool wrapper that short-circuits; todo tool replace semantics and prompt rendering; repeat-call threshold; allowlist enforcement including a middleware that tries to re-add a banned tool; private-state round trip — marshal State with three middlewares’ typed state, reload it, and assert each middleware reads back its struct, plus a missing-key case asserting the default is initialized and middleware.state_missing is recorded ; per-batch scratch state is empty at the start of each batch; full regression run of the phase-5 and phase-6 suites.
11. Update docs and write the learning doc.

Learning Doc Outline

docs/learning/07-middleware-architecture.md must contain:

• Concept: the onion/decorator pattern applied to model and tool calls; why “every feature is a middleware” is the load-bearing decision of the whole harness.
• Design choices: marker interface plus optional capability interfaces, one-time type partitioning, non-mutating request rewriting, State.Private keyed by name, ordered stack from one constructor, persisted run state.
• Order rationale, slot by slot, including the two non-obvious constraints (allowlist last, prompt cache before memory) and what breaks if violated.
• Trade-offs: middleware chain vs a stage pipeline vs plain function composition; type assertion vs five separate slices in the constructor; map[string]json.RawMessage plus generic accessors vs map[string]any (the latter cannot survive a database round trip, which is the whole point of persisted run state).
• Prompt assembly as a chain-time operation: why contributing `[]provider.Block` rather than strings is what makes cache breakpoints and memory ordering expressible, and why order assertions must be made on the breakpoint index.
• goclaw comparison: the V3 Stage pipeline with Continue/BreakLoop/AbortRun, plus the XxxAware setter-interface DI hack, and why hango uses constructor arguments instead.
• deepagents comparison: AgentMiddleware with before_agent/wrap_model_call/wrap_tool_call, the exact default order in graph.py:855-905, _apply_custom_middleware/excluded_middleware splicing machinery, write_todos as whole-list replace with OmitFromInput, and the fact that todos are opt-in since 0.7.
• How to debug this layer: the agent.middleware_stack span event tells you which middlewares ran and in what order for that exact run; middleware.* span events tell you which one rewrote the request, short-circuited a tool, or blocked a repeat call; hango.middleware.before_run.duration tells you which load-once hook is slow.
• Why hango diverges: no user-supplied middleware splicing (order is code), no dynamic state-schema merging, todos always available but gated by the agent’s tool allowlist.
• Exercises: write a middleware that logs every model request; reorder two slots and observe the test failure.
• Time budget: 11h implementation + 2h learning doc.

Success Criteria

• go build ./... && go vet ./... && go test ./... pass, including every unchanged phase-5 and phase-6 test.
• mwstack_test.go asserts the exact middleware order and the exact tool order presented to the provider.
• Chain test proves index 0 is outermost for both model and tool calls.
• A middleware attempting to re-add a disallowed tool is defeated by the allowlist middleware (test).
• write_todos replaces the list, appears in the system prompt, and never appears in State.Messages (test).
• Repeat-call middleware blocks the third identical call and returns the nudge (test).
• runs.state_json holds the todo list after a run; reloading the run restores it (integration test).
• Private-state round trip: every middleware’s typed state survives marshal → runs.state_json → unmarshal; a missing key initializes the documented default and emits middleware.state_missing.
• The assembled Request.System block order is asserted by index inside ModelCall, and a contributor added later in the chain lands at a higher index than one added earlier.
• In Tempo, the run’s invoke_agent span carries an agent.middleware_stack event listing the exact ordered stack, and blocked repeat calls appear as middleware.repeat_call_blocked events.

Risk Assessment

Rollback

goose down one step drops runs.state_json. The chain is additive around the phase-5 loop: constructing DefaultStack with an empty slice yields the phase-5 behaviour exactly, which is the fallback if the refactor misbehaves.