Middleware State

Manage typed, immutable state across middleware hooks using source-generated properties.

Quick Example

1. Define your state record:

[MiddlewareState]
public sealed record MyCounterState
{
    public int Count { get; init; } = 0;
}

2. Read and update in middleware:

public class CounterMiddleware : IAgentMiddleware
{
    public Task BeforeIterationAsync(BeforeIterationContext context, CancellationToken ct)
    {
        // Update state (auto-instantiates if null!)
        context.UpdateMiddlewareState<MyCounterState>(s => s with
        {
            Count = s.Count + 1
        });

        return Task.CompletedTask;
    }
}

3. Source generator creates:

public sealed partial class MiddlewareState
{
    public MyCounterState? MyCounter { get; } // Auto-generated getter
    public MiddlewareState WithMyCounter(MyCounterState? value) { } // Auto-generated setter
}

Creating Middleware State

Step 1: Define State Record

using HPD.Agent;

[MiddlewareState(Version = 1)]
public sealed record ErrorTrackingState
{
    public int ConsecutiveFailures { get; init; } = 0;
    public DateTime? LastErrorTime { get; init; }

    // Helper methods (optional)
    public ErrorTrackingState IncrementFailures() =>
        this with
        {
            ConsecutiveFailures = ConsecutiveFailures + 1,
            LastErrorTime = DateTime.UtcNow
        };

    public ErrorTrackingState ResetFailures() =>
        this with { ConsecutiveFailures = 0, LastErrorTime = null };
}

Requirements

• Must be a record (compile error if not)
• Should be sealed for performance (compiler warning if not)
• All properties must be JSON-serializable (runtime error during checkpoint save/restore if not)
• Use { get; init; } for immutability

Reading State

⭐ Recommended: Use Analyze() Method

Use context.Analyze() to read state safely:

public Task BeforeIterationAsync(BeforeIterationContext context, CancellationToken ct)
{
    //   RECOMMENDED: Use Analyze() for safe state reads
    var consecutiveFailures = context.Analyze(s =>
        s.MiddlewareState.ErrorTracking?.ConsecutiveFailures ?? 0
    );

    if (consecutiveFailures >= 3)
    {
        context.UpdateState(s => s with { IsTerminated = true });
    }

    return Task.CompletedTask;
}

Why use Analyze()?

• Prevents stale captures: Lambda executes immediately, getting fresh state
• Async-safe: Even if you add await later, the read happens at call time
• Zero overhead: JIT inlines the lambda - identical to direct access
• Clear intent: Makes state reads explicit in your code

Extract multiple values with tuple deconstruction:

var (errors, iteration, isTerminated) = context.Analyze(s => (
    s.MiddlewareState.ErrorTracking?.ConsecutiveFailures ?? 0,
    s.Iteration,
    s.IsTerminated
));

Alternative: Read Inside UpdateState

For mutations, read state directly in the UpdateState lambda:

context.UpdateState(s =>
{
    // Read state here - always fresh!
    var errorState = s.MiddlewareState.ErrorTracking ?? new();
    var updated = errorState.IncrementFailures();

    return s with
    {
        MiddlewareState = s.MiddlewareState.WithErrorTracking(updated)
    };
});

Always provide default: State is null until first update.

Updating State

⭐ Recommended: Simplified Extension Methods

Use UpdateMiddlewareState<T>() and GetMiddlewareState<T>() for clean, concise state management:

public class ErrorTrackingMiddleware : IAgentMiddleware
{
    public Task AfterIterationAsync(AfterIterationContext context, CancellationToken ct)
    {
        var hasErrors = context.ToolResults.Any(r => r.Exception != null);

        if (hasErrors)
        {
            // Update state (auto-instantiates if null!)
            context.UpdateMiddlewareState<ErrorTrackingStateData>(s => s.IncrementFailures());

            // Read state
            var failures = context.GetMiddlewareState<ErrorTrackingStateData>()?.ConsecutiveFailures ?? 0;
            if (failures >= 3)
            {
                context.UpdateState(s => s with
                {
                    IsTerminated = true,
                    TerminationReason = "Too many consecutive failures"
                });
            }
        }
        else
        {
            // Reset on success
            context.UpdateMiddlewareState<ErrorTrackingStateData>(s => s.ResetFailures());
        }

        return Task.CompletedTask;
    }
}

Key features:

• Auto-instantiation - No ?? new() needed
• Type-safe - Compiler catches errors
• Clean syntax - Focuses on the transformation

Common patterns:

// Simple update
context.UpdateMiddlewareState<ErrorTrackingStateData>(s => s.IncrementFailures());

// Multi-field update
context.UpdateMiddlewareState<ErrorTrackingStateData>(s => s with
{
    ConsecutiveFailures = s.ConsecutiveFailures + 1,
    LastErrorTime = DateTime.UtcNow
});

// Reset to defaults
context.UpdateMiddlewareState<ErrorTrackingStateData>(_ => new ErrorTrackingStateData());

// Read state
var count = context.GetMiddlewareState<ErrorTrackingStateData>()?.ConsecutiveFailures ?? 0;

When to Use Advanced UpdateState

Use UpdateState() directly for:

• Core state changes - IsTerminated, TerminationReason, CurrentIteration
• Atomic multi-state updates - Update middleware state + core state together
• Complex transformations - Multiple state types with business logic

// Advanced: Atomic update of middleware + core state
context.UpdateState(s =>
{
    var errors = s.MiddlewareState.ErrorTracking ?? new();
    var failures = errors.ConsecutiveFailures + 1;

    return s with
    {
        MiddlewareState = s.MiddlewareState.WithErrorTracking(
            errors with { ConsecutiveFailures = failures }
        ),
        IsTerminated = failures >= 3,
        TerminationReason = $"Circuit breaker: {failures} consecutive failures"
    };
});

Updating State (Advanced)

Immediate Updates

Updates are applied immediately:

public Task BeforeIterationAsync(BeforeIterationContext context, CancellationToken ct)
{
    // Update
    context.UpdateState(s =>
    {
        var state = s.MiddlewareState.MyCounter ?? new();
        return s with
        {
            MiddlewareState = s.MiddlewareState.WithMyCounter(state with { Count = state.Count + 1 })
        };
    });

    //   Next middleware sees updated state immediately!
    var updatedCount = context.Analyze(s => s.MiddlewareState.MyCounter?.Count ?? 0);
    Console.WriteLine(updatedCount); // Shows incremented value
}

Updates are immediate - visible to all subsequent middleware.

Immutable Pattern

Always use with expressions:

//   CORRECT: Immutable update
context.UpdateState(s => s with
{
    MiddlewareState = s.MiddlewareState.WithMyState(currentState with
    {
        Count = currentState.Count + 1
    })
});

//    WRONG: Mutation
currentState.Count++; // Compile error - init-only property

Nested State Updates

public Task AfterIterationAsync(AfterIterationContext context, CancellationToken ct)
{
    var errors = context.ToolResults.Count(r => r.Exception != null);

    if (errors > 0)
    {
        context.UpdateState(s =>
        {
            var errorState = s.MiddlewareState.ErrorTracking ?? new();
            var newErrorState = errorState.IncrementFailures();

            return s with
            {
                MiddlewareState = s.MiddlewareState.WithErrorTracking(newErrorState),
                // Can update multiple fields
                IsTerminated = newErrorState.ConsecutiveFailures >= 3,
                TerminationReason = "Too many errors"
            };
        });
    }
    else
    {
        context.UpdateState(s =>
        {
            var errorState = s.MiddlewareState.ErrorTracking ?? new();
            return s with
            {
                MiddlewareState = s.MiddlewareState.WithErrorTracking(errorState.ResetFailures())
            };
        });
    }

    return Task.CompletedTask;
}

State Versioning

Version tracks breaking schema changes:

[MiddlewareState(Version = 2)] // Bumped from 1 → 2
public sealed record MyState
{
    public int Count { get; init; } // Was string in v1
}

When to Bump Version

Increment version for:

• Removing properties
• Renaming properties
• Changing property types
• Changing collection types (List → ImmutableList)

No version bump needed for:

• Adding new properties with defaults
• Adding helper methods
• Updating documentation

Complex State Example

[MiddlewareState(Version = 1)]
public sealed record CircuitBreakerState
{
    public Dictionary<string, int> CallCounts { get; init; } = new();
    public Dictionary<string, string> LastSignatures { get; init; } = new();

    public int GetCount(string toolName, string signature)
    {
        if (!LastSignatures.TryGetValue(toolName, out var lastSig))
            return 0;

        return lastSig == signature
            ? CallCounts.GetValueOrDefault(toolName, 0)
            : 0;
    }

    public CircuitBreakerState IncrementCount(string toolName, string signature)
    {
        var lastSig = LastSignatures.GetValueOrDefault(toolName);
        var count = lastSig == signature
            ? CallCounts.GetValueOrDefault(toolName, 0) + 1
            : 1;

        return this with
        {
            CallCounts = new Dictionary<string, int>(CallCounts)
            {
                [toolName] = count
            },
            LastSignatures = new Dictionary<string, string>(LastSignatures)
            {
                [toolName] = signature
            }
        };
    }

    public CircuitBreakerState Reset(string toolName)
    {
        var newCounts = new Dictionary<string, int>(CallCounts);
        newCounts.Remove(toolName);

        var newSigs = new Dictionary<string, string>(LastSignatures);
        newSigs.Remove(toolName);

        return this with
        {
            CallCounts = newCounts,
            LastSignatures = newSigs
        };
    }
}

State Sharing Between Middleware

State is global - all middleware can read/write:

// Middleware A writes
public class ErrorTrackerMiddleware : IAgentMiddleware
{
    public Task AfterIterationAsync(AfterIterationContext context, CancellationToken ct)
    {
        context.UpdateState(s =>
        {
            var state = s.MiddlewareState.ErrorTracking ?? new();
            return s with
            {
                MiddlewareState = s.MiddlewareState.WithErrorTracking(state.IncrementFailures())
            };
        });
        return Task.CompletedTask;
    }
}

// Middleware B reads
public class CircuitBreakerMiddleware : IAgentMiddleware
{
    public Task BeforeIterationAsync(BeforeIterationContext context, CancellationToken ct)
    {
        // Can read ErrorTracking state from ErrorTrackerMiddleware
        var consecutiveFailures = context.Analyze(s =>
            s.MiddlewareState.ErrorTracking?.ConsecutiveFailures ?? 0
        );

        if (consecutiveFailures >= 3)
        {
            context.UpdateState(s => s with
            {
                IsTerminated = true,
                TerminationReason = "Circuit breaker"
            });
        }

        return Task.CompletedTask;
    }
}

State Persistence

State persists across:

• Iterations within a turn
• Multiple turns in a session
• Agent restarts (if using AgentSession with checkpoint storage)

// Turn 1
await agent.RunAsync("What's 2+2?"); // State.Count = 1

// Turn 2 (same session)
await agent.RunAsync("And 3+3?"); // State.Count = 2 (persisted!)

Automatic Cross-Run Persistence

By default, middleware state resets when the agent run completes. For state that should survive across runs (like caches or user preferences), use Persistent = true:

//   Persistent state (survives across agent runs)
[MiddlewareState(Version = 1, Persistent = true)]
public sealed record HistoryReductionState
{
    public CachedReduction? LastReduction { get; init; }
}

//   Transient state (resets each run - default)
[MiddlewareState(Version = 1)]
public sealed record ErrorTrackingState
{
    public int ConsecutiveErrors { get; init; }
}

How it works:

• States marked Persistent = true are automatically saved to AgentSession
• Loaded automatically when resuming from the same session
• No manual serialization code needed

When to use Persistent = true:

• Expensive caches - HistoryReduction (avoids re-summarizing messages)
• User preferences - Permissions, settings
• Long-term tracking - Total usage metrics, user history

When to use transient (default):

• Safety metrics - Error tracking, circuit breakers (MUST reset each run)
• Per-run state - Iteration counts, current batch
• Temporary tracking - Active tool calls, pending operations

Example:

[MiddlewareState(Persistent = true)]
public sealed record UserPreferencesState
{
    public string? PreferredLanguage { get; init; }
    public bool VerboseMode { get; init; }
}

// Run 1: User sets preferences
var session = new AgentSession();
await agent.RunAsync(["Set language to Spanish"], session);

// Run 2: Preferences are restored automatically
await agent.RunAsync(["What's the weather?"], session);
// Agent remembers language = Spanish from Run 1!

Why transient by default?

Safety middlewares MUST reset between runs. If error counts or circuit breaker state persisted, subsequent runs would start with incorrect state:

//   BAD: If this persisted, next run would start pre-broken!
[MiddlewareState(Persistent = true)] // WRONG!
public sealed record CircuitBreakerState
{
    public Dictionary<string, int> FailureCounts { get; init; }
}

//   CORRECT: Resets each run
[MiddlewareState] // Transient by default
public sealed record CircuitBreakerState
{
    public Dictionary<string, int> FailureCounts { get; init; }
}

Thread Safety

State is immutable - safe for concurrent RunAsync() calls:

//   SAFE: Two parallel calls, independent state
var task1 = agent.RunAsync("First query");
var task2 = agent.RunAsync("Second query");
await Task.WhenAll(task1, task2);

Never use instance fields for state:

public class BadMiddleware : IAgentMiddleware
{
    //    WRONG: Not thread-safe!
    private int _count = 0;

    public Task BeforeIterationAsync(BeforeIterationContext context, CancellationToken ct)
    {
        _count++; // Race condition with parallel RunAsync calls
        return Task.CompletedTask;
    }
}

public class GoodMiddleware : IAgentMiddleware
{
    //   CORRECT: Use middleware state
    public Task BeforeIterationAsync(BeforeIterationContext context, CancellationToken ct)
    {
        context.UpdateState(s =>
        {
            var state = s.MiddlewareState.MyCounter ?? new();
            return s with
            {
                MiddlewareState = s.MiddlewareState.WithMyCounter(state with { Count = state.Count + 1 })
            };
        });
        return Task.CompletedTask;
    }
}

Common Patterns

Pattern 1: Error Counting

[MiddlewareState]
public sealed record ErrorCountState
{
    public int TotalErrors { get; init; }
    public int ConsecutiveErrors { get; init; }
}

public Task OnErrorAsync(ErrorContext context, CancellationToken ct)
{
    context.UpdateState(s =>
    {
        var state = s.MiddlewareState.ErrorCount ?? new();
        return s with
        {
            MiddlewareState = s.MiddlewareState.WithErrorCount(state with
            {
                TotalErrors = state.TotalErrors + 1,
                ConsecutiveErrors = state.ConsecutiveErrors + 1
            })
        };
    });

    return Task.CompletedTask;
}

Pattern 2: Rate Limiting

[MiddlewareState]
public sealed record RateLimitState
{
    public DateTime? LastCallTime { get; init; }
    public int CallsInWindow { get; init; }
}

public async Task BeforeIterationAsync(BeforeIterationContext context, CancellationToken ct)
{
    // Read state for conditional logic
    var (lastCallTime, callsInWindow) = context.Analyze(s => (
        s.MiddlewareState.RateLimit?.LastCallTime,
        s.MiddlewareState.RateLimit?.CallsInWindow ?? 0
    ));

    var now = DateTime.UtcNow;

    if (lastCallTime.HasValue &&
        (now - lastCallTime.Value) < TimeSpan.FromMinutes(1))
    {
        if (callsInWindow >= 10)
        {
            await Task.Delay(TimeSpan.FromSeconds(60), ct);
        }

        context.UpdateState(s =>
        {
            var state = s.MiddlewareState.RateLimit ?? new();
            return s with
            {
                MiddlewareState = s.MiddlewareState.WithRateLimit(state with
                {
                    CallsInWindow = state.CallsInWindow + 1
                })
            };
        });
    }
    else
    {
        context.UpdateState(s => s with
        {
            MiddlewareState = s.MiddlewareState.WithRateLimit(new RateLimitState
            {
                LastCallTime = now,
                CallsInWindow = 1
            })
        });
    }
}

Pattern 3: Batch State

[MiddlewareState]
public sealed record BatchApprovalState
{
    public HashSet<string> ApprovedFunctions { get; init; } = new();
}

public async Task BeforeParallelBatchAsync(BeforeParallelBatchContext context, CancellationToken ct)
{
    var functions = context.ParallelFunctions.Select(f => f.Name ?? "_unknown").ToList();
    var approved = await RequestApproval(functions, ct);

    context.UpdateState(s => s with
    {
        MiddlewareState = s.MiddlewareState.WithBatchApproval(new BatchApprovalState
        {
            ApprovedFunctions = approved.ToHashSet()
        })
    });
}

public Task BeforeFunctionAsync(BeforeFunctionContext context, CancellationToken ct)
{
    var isApproved = context.Analyze(s =>
        s.MiddlewareState.BatchApproval?.ApprovedFunctions.Contains(context.Function.Name) ?? false
    );

    if (!isApproved)
    {
        context.BlockExecution = true;
        context.OverrideResult = "Not approved";
    }

    return Task.CompletedTask;
}

Thread Safety and Update Patterns

Middleware state updates are thread-safe when using the correct pattern. HPD-Agent provides three layers of defense to prevent race conditions.

Defense Mechanisms

1. Fail-fast guard - Prevents Agent.cs from calling SyncState() during middleware execution
2. Reference check - Detects stale reads and background task races
3. Safe patterns - Prevent bugs by construction

⭐ Best Practice: Use Analyze() for Conditionals

For reading state to make decisions, use context.Analyze():

public Task BeforeIterationAsync(BeforeIterationContext context, CancellationToken ct)
{
    //   BEST: Use Analyze() for conditional reads
    var shouldTerminate = context.Analyze(s =>
        s.MiddlewareState.ErrorTracking?.ConsecutiveFailures >= 3
    );

    if (shouldTerminate)
    {
        context.UpdateState(s => s with
        {
            IsTerminated = true,
            TerminationReason = "Too many errors"
        });
    }

    return Task.CompletedTask;
}

Why this pattern is best:

• Async-safe: Lambda executes immediately, value captured safely
• Clear intent: Makes point-in-time reads explicit
• Zero overhead: JIT inlines the lambda
• Thread-safe: No risk of capturing stale state references

⭐ Recommended Pattern: Read Inside UpdateState

For mutations, always read state inside the UpdateState lambda:

context.UpdateState(s =>
{
    // Read current state (always fresh)
    var current = s.MiddlewareState.ErrorTracking ?? new();

    // Transform
    var updated = current.IncrementFailures();

    // Return new state
    return s with
    {
        MiddlewareState = s.MiddlewareState.WithErrorTracking(updated)
    };
});

Why this pattern is safe:

• State read happens inside lambda (always gets latest state)
• Async-safe: Can add await anywhere before UpdateState
• Readable: Local variables show intent
• Protected by generation counter guard

Example with async operations:

public async Task BeforeIterationAsync(BeforeIterationContext context, CancellationToken ct)
{
    //   SAFE: Async work happens BEFORE UpdateState
    await ValidateInputAsync();
    await LogStartAsync();

    // State read happens inside lambda - always fresh
    context.UpdateState(s =>
    {
        var current = s.MiddlewareState.MyState ?? new();
        var updated = current.Increment();
        return s with { MiddlewareState = s.MiddlewareState.WithMyState(updated) };
    });
}

🥈 Alternative: Compact Pattern

For simple one-line transforms:

context.UpdateState(s => s with
{
    CurrentIteration = s.CurrentIteration + 1
});

When to use:

• Simple field updates
• No intermediate calculations
• No complex logic

Anti-Pattern: Capturing State Outside UpdateState

Don't read state before async operations:

// DANGEROUS: Stale state capture
var errorState = /* somehow get state */;
var updated = errorState.IncrementFailures();

// If you add async work here, 'updated' becomes stale!
await SomeAsyncWork();

// UpdateState with stale data - RUNTIME ERROR
context.UpdateState(s => s with
{
    MiddlewareState = s.MiddlewareState.WithErrorTracking(updated)
});
// Throws: "State was modified during UpdateState transform"

Why this fails:

• State can change during await operations
• Using stale state leads to lost updates
• The generation counter detects this and throws

Correct patterns:

//   CORRECT: Use Analyze() for conditionals
var shouldReset = context.Analyze(s =>
    (s.MiddlewareState.ErrorTracking?.ConsecutiveFailures ?? 0) >= 3
);

// Async work can safely happen here
await SomeAsyncWork();

if (shouldReset)
{
    context.UpdateState(s => s with
    {
        MiddlewareState = s.MiddlewareState.WithErrorTracking(new())
    });
}

//   CORRECT: Or read inside UpdateState for mutations
await SomeAsyncWork();  // Async work first

context.UpdateState(s =>
{
    // Read state HERE - always fresh!
    var state = s.MiddlewareState.ErrorTracking ?? new();
    var updated = state.IncrementFailures();
    return s with
    {
        MiddlewareState = s.MiddlewareState.WithErrorTracking(updated)
    };
});

Error you'll see if you add an await:

InvalidOperationException: State was modified during UpdateState transform.

This usually indicates one of these issues:
  1. Stale read: You read context.State before an 'await', then used the old value
  2. Background task: A Task.Run() or fire-and-forget task updated state after middleware
  3. Concurrent modification: Threading bug (rare)

SOLUTION - Use the block-scoped lambda pattern:
  context.UpdateState(s =>
  {
      var current = s.MiddlewareState.MyState ?? new();
      var updated = current.Transform();
      return s with { MiddlewareState = s.MiddlewareState.WithMyState(updated) };
  });

This ensures state is read INSIDE the lambda, getting the latest value.

Understanding the Guards

HPD-Agent uses two runtime guards to catch bugs:

Guard #1: Fail-Fast for Agent.cs

Prevents SyncState() from being called during middleware execution:

// This will throw if Agent.cs has a timing bug:
internal void SyncState(AgentLoopState newState)
{
    if (_middlewareExecuting)
        throw new InvalidOperationException("SyncState() called during middleware execution");
}

You'll only see this error if there's a bug in Agent.cs - it's not a user-facing error.

Guard #2: Reference Check for Stale Reads

Detects when state changes between reading and updating:

public void UpdateState(Func<AgentLoopState, AgentLoopState> transform)
{
    var stateBefore = _state;
    var stateAfter = transform(stateBefore);

    if (!ReferenceEquals(_state, stateBefore))
        throw new InvalidOperationException("State was modified during UpdateState");
}

You'll see this error if:

• You read state outside the lambda, then await, then update
• A background task tries to update state after middleware completes

Fix: Use the block-scoped lambda pattern.

Multiple Updates in One Hook

Update multiple state fields in one atomic UpdateState call:

//   CORRECT: Single atomic update
context.UpdateState(s =>
{
    var errorState = s.MiddlewareState.ErrorTracking ?? new();
    var updatedErrors = errorState.IncrementFailures();

    return s with
    {
        CurrentIteration = s.CurrentIteration + 1,
        MiddlewareState = s.MiddlewareState.WithErrorTracking(updatedErrors),
        IsTerminated = updatedErrors.ConsecutiveFailures >= 3
    };
});

//   INCORRECT: Multiple separate updates (race window between calls)
context.UpdateState(s => s with { CurrentIteration = s.CurrentIteration + 1 });
context.UpdateState(s => s with { MiddlewareState = s.MiddlewareState.WithErrorTracking(...) });
context.UpdateState(s => s with { IsTerminated = true });

Next Steps

• 05.1 Middleware Lifecycle - Complete hook reference
• 05.3 Middleware Events - Emit events and wait for responses
• 05.4 Built-in Middleware - See state usage in real middleware
• 05.5 Custom Middleware - Build your own with state