railtownai/railtracks/middleware

Add Middleware to a Railtracks Agent or Node

The user wants to add middleware to a railtracks node/agent: $ARGUMENTS

Quick reference

• All five accept optional keyword-only name= and work bare or called: @rt.pre_llm, @rt.post_node(name="print_after").
• wrap_node/wrap_llm require async def (TypeError otherwise). pre_llm/post_llm/post_node accept sync or async.
• middleware= wraps the whole node call (an agent's full tool-calling loop, including every tool node it invokes). model_middleware= wraps one raw LLM call inside that loop.
• Attach at creation — rt.agent_node(..., middleware=[...], model_middleware=[...]) — or after the fact without mutating the original — rt.couple(node, middleware=[...], model_middleware=[...]).
• List order = wrapping order: middleware=[A, B, C] → A wraps B wraps C wraps the node. rt.couple-ing more onto an already-middlewared node adds the new ones outermost; existing middleware stays innermost, closest to the node.

Best practices

1. Exception handling & retries

• pre_llm, post_llm, post_node only run on success — an exception skips them entirely. If you need to react to failure (cleanup, translation, fallback), use wrap_node/wrap_llm and put your own try/except around await call(...).
• Where you put a retry changes what it re-executes:
• model_middleware= retry only re-issues the raw LLM call — safe by default, no side effects beyond the model provider request.
• middleware= retry on an agent re-runs the entire node, including every tool call that already succeeded in that attempt. Only retry at this slot if the tools involved are safe to re-run (idempotent, or guarded by an idempotency key) — otherwise a transient failure after a non-idempotent tool call (e.g. "send email", "charge card") causes it to fire twice.
• Default to retrying in model_middleware= for provider flakiness (rate limits, timeouts). Reach for middleware= retry only when you've confirmed the tools are safe to repeat.
• Prefer railtracks.prebuilt.middleware.Retry(max_tries=..., approach=..., retry_on=...) over hand-rolling — it raises once attempts are exhausted rather than swallowing the failure, and lets you scope retry_on to specific exception types instead of catching everything.
• Don't swallow exceptions silently inside a wrap_node/wrap_llm middleware (except Exception: return default) — that hides the failure from callers and from anything else in the chain. If you must return a fallback, do it deliberately (log it, tag it in the result) rather than as an unannounced side effect.

2. Middleware vs. guardrails

• rt.input_guard / rt.output_guard are middleware under the hood (BaseGuardrail extends Middleware), specialized for one thing: "should this content pass?" The guard function receives an LLMGuardrailEvent and returns a GuardrailDecision (allow / block / transform).
• Use a guard, not hand-rolled middleware, whenever the logic is a content-policy decision (PII/secrets detection, profanity, jailbreak checks, redaction). You get for free: structured GuardrailTrace records (rail_name, phase, action, reason), fail_open= to decide whether a guard's own bug blocks or passes traffic, and consistent block semantics other tooling can rely on.
• To keep a domain-specific agent from answering off-topic requests, use the prebuilt ScopeGuard(scope=..., model=..., sensitivity=...) from railtracks.prebuilt.guardrails; a judge model labels the request before the agent's own model runs.
• Use plain middleware for everything that isn't a pass/block decision on content: retries, timing, logging, context injection, auth/telemetry.

3. Ordering & composition

• Two different things happen depending on middleware shape — know which one you're ordering:
• Input-transforming middleware (pre_llm-style: mutates message_history/schema/tools before calling onward) — earlier entries in the list run first and hand their result to later entries. Put transformations before validations: model_middleware=[ContextInjection(), my_guard] so the guard sees the filled-in prompt, not the raw template.
• Span-wrapping middleware (retry, timing, logging that brackets a whole call) — the outermost entry sees the aggregate outcome; the innermost is what actually gets re-invoked on each attempt. middleware=[Retry(3), log_after] → the log fires once per retry attempt (it's inside the retry). middleware=[log_after, Retry(3)] → the log fires once, for the final outcome only (it's outside the retry). Pick based on whether you want per-attempt or per-call visibility.
• When rt.couple-ing middleware onto a node that already has some, the new list is outermost — it runs before the existing middleware, so it can't be placed inside an existing auth check or retry.

4. Naming & observability

• name= sets Middleware.name, which is recorded with the middleware's events and shown in the beta visualizer (railtracks viz --beta) — the same way guardrails write their name into every GuardrailTrace.rail_name. Give your middleware a real name= so it's identifiable in traces/logs instead of falling back to the wrapper function's __name__.
• Set name= when the same wrapper function is reused in more than one slot (e.g. one generic retry_node attached to several agents) and you need to tell instances apart.
• Keep one concern per middleware function. A function that both logs and mutates the prompt makes the ordering question in §3 ambiguous — split it into two so each one has one clear rule for where it belongs in the list.

5. State

• A middleware instance is shared, never copied: every node it is attached to (both slots), every node rt.couple builds from those, and every run, including concurrent runs, call the same object.
• Keep per-run state in rt.context under a key the middleware owns (one key per instance, e.g. a uuid suffix), as ConversationMemory does. Each run gets its own copy of the context dictionary, so concurrent runs stay apart.
• Keep only process-wide resources on self or in a closure: clients, connection pools, semaphores.
• For a budget, use railtracks.prebuilt.middleware.MaxCalls, which keeps a separate count per session. To serialize calls that share state, put railtracks.prebuilt.middleware.Lock outside the middleware that uses it.

Patterns

Retry: prebuilt vs. hand-rolled, and where to put it

import railtracks as rt
from railtracks.prebuilt import middleware
Agent = rt.agent_node(
    "Agent",
    tool_nodes=[send_email],  # has a side effect — don't retry the whole node
    llm=rt.llm.OpenAILLM(model_name="gpt-6-luna"),
    model_middleware=[middleware.Retry(max_tries=3)],  # retries only the raw LLM call
)

Guardrail vs. plain middleware for the same-shaped problem

import railtracks as rt
# Content-policy decision -> guardrail, not middleware
@rt.input_guard
def block_secrets(event: rt.guardrails.LLMGuardrailEvent) -> rt.guardrails.GuardrailDecision:
    if any("SECRET" in m.content for m in event.messages if isinstance(m.content, str)):
        return rt.guardrails.GuardrailDecision.block(reason="secret leaked")
    return rt.guardrails.GuardrailDecision.allow()
# Not a policy decision -> plain middleware
@rt.pre_llm
def log_prompt(message_history, schema, tools):
    print(message_history)
    return message_history, schema, tools

Running a guarded agent

Run it through a Flow like any agent. A block raises GuardrailBlockedError, so catch it where you want to show the user a refusal.

import railtracks as rt
SupportAgent = rt.agent_node(
    "Support Agent",
    llm=rt.llm.OpenAILLM("gpt-6-luna"),
    model_middleware=[block_secrets],
)
flow = rt.Flow(name="Support", entry_point=SupportAgent)
try:
    result = flow.invoke("user message here")
    print(result.content)
except rt.guardrails.GuardrailBlockedError as e:
    print(f"Blocked: {e}")

Ordering a retry against a logger to control per-attempt vs. per-call visibility

import railtracks as rt
@rt.post_node(name="print_after")
async def log_result(result):
    print("Finished:", result)
    return result
retry = rt.prebuilt.middleware.Retry(max_tries=3)
PerAttemptLogging = rt.couple(Agent, middleware=[retry, log_result])   # logs every attempt
PerCallLogging = rt.couple(Agent, middleware=[log_result, retry])     # logs only the final outcome

Attaching middleware after the node already exists

import railtracks as rt
Adjusted = rt.couple(
    Agent,
    middleware=[log_result],
    model_middleware=[log_prompt],
)

Rules

• Retry at model_middleware= by default. Retry at middleware= only when every tool the agent calls is safe to run twice, since a retry there re-runs every tool call in the attempt.
• Write content allow/block checks as input_guard/output_guard, so they show up in GuardrailTrace with the rest of your rails.
• Write wrap_node/wrap_llm functions as async def.
• Put failure-aware logic in wrap_llm/wrap_node; post_llm/post_node run only on success.
• Keep per-run middleware state in rt.context, not on self; one instance serves every run, including concurrent ones.