Agent 中间件与安全护栏：限流、脱敏、重试与降级实测

"""Test doubles shared by the middleware labs: no network access and no API key."""
from typing import Any, List, Optional, Sequence
 
from langchain_core.language_models.chat_models import BaseChatModel
from langchain_core.messages import AIMessage, BaseMessage, ToolMessage
from langchain_core.outputs import ChatGeneration, ChatResult
from langchain_core.tools import tool
 
TOOL_RUNS: List[str] = []  # one entry per real execution of the tool body
 
 
@tool
def search(query: str) -> str:
    """Search the knowledge base."""
    TOOL_RUNS.append(query)
    return f"result for {query}"
 
 
class LoopingModel(BaseChatModel):
    """Requests the `search` tool on every call.
 
    With stop_on_error=True it gives a final answer as soon as the last message is a
    ToolMessage with status "error", which imitates a model that reacts to error feedback.
    """
 
    stop_on_error: bool = False
    calls: int = 0
 
    @property
    def _llm_type(self) -> str:
        return "looping"
 
    def bind_tools(self, tools: Sequence[Any], **kwargs: Any):
        return self
 
    def _generate(
        self,
        messages: List[BaseMessage],
        stop: Optional[List[str]] = None,
        run_manager: Any = None,
        **kwargs: Any,
    ) -> ChatResult:
        self.calls += 1
        last = messages[-1]
        if self.stop_on_error and isinstance(last, ToolMessage) and last.status == "error":
            message = AIMessage(content="giving up after the tool reported an error")
        else:
            message = AIMessage(
                content="",
                tool_calls=[{"name": "search", "args": {"query": f"q{self.calls}"}, "id": f"call_{self.calls}"}],
            )
        return ChatResult(generations=[ChatGeneration(message=message)])
 
 
class AnsweringModel(BaseChatModel):
    """Always answers directly, and remembers every message list it was shown."""
 
    seen: List[List[BaseMessage]] = []
    calls: int = 0
 
    @property
    def _llm_type(self) -> str:
        return "answering"
 
    def bind_tools(self, tools: Sequence[Any], **kwargs: Any):
        return self
 
    def _generate(
        self,
        messages: List[BaseMessage],
        stop: Optional[List[str]] = None,
        run_manager: Any = None,
        **kwargs: Any,
    ) -> ChatResult:
        self.calls += 1
        self.seen.append(list(messages))
        return ChatResult(generations=[ChatGeneration(message=AIMessage(content=f"answer {self.calls}"))])
 
 
class FlakyModel(BaseChatModel):
    """Fails the first `failures` calls with `error`, then answers with its own name."""
 
    label: str = "model"
    failures: int = 0
    error: Any = ConnectionError
    attempts: int = 0
 
    @property
    def _llm_type(self) -> str:
        return "flaky"
 
    def bind_tools(self, tools: Sequence[Any], **kwargs: Any):
        return self
 
    def _generate(
        self,
        messages: List[BaseMessage],
        stop: Optional[List[str]] = None,
        run_manager: Any = None,
        **kwargs: Any,
    ) -> ChatResult:
        self.attempts += 1
        if self.attempts <= self.failures:
            raise self.error(f"{self.label} attempt {self.attempts} failed")
        return ChatResult(generations=[ChatGeneration(message=AIMessage(content=f"answered by {self.label}"))])

"""How do ModelCallLimitMiddleware and ToolCallLimitMiddleware end a runaway tool loop?"""
from common import TOOL_RUNS, AnsweringModel, LoopingModel
from langchain.agents import create_agent
from langchain.agents.middleware import ModelCallLimitMiddleware, ToolCallLimitMiddleware
from langchain_core.messages import HumanMessage
from langgraph.checkpoint.memory import InMemorySaver
 
from common import search
 
 
def run(label, middleware, stop_on_error=False):
    TOOL_RUNS.clear()
    model = LoopingModel(stop_on_error=stop_on_error)
    agent = create_agent(model, tools=[search], middleware=middleware)
    try:
        result = agent.invoke({"messages": [HumanMessage("find it")]}, {"recursion_limit": 20})
        last = result["messages"][-1]
        outcome = f"finished, last message: {last.type}: {last.content!r}"
    except Exception as exc:
        outcome = f"raised {type(exc).__name__}"
    print(f"{label:<44} model_calls={model.calls:<2} tool_runs={len(TOOL_RUNS)}  {outcome}")
 
 
run("no limits (recursion_limit=20)", [])
run("ModelCallLimit run_limit=3, end", [ModelCallLimitMiddleware(run_limit=3)])
run("ModelCallLimit run_limit=3, error", [ModelCallLimitMiddleware(run_limit=3, exit_behavior="error")])
run("ToolCallLimit run_limit=2, continue", [ToolCallLimitMiddleware(run_limit=2)])
run("ToolCallLimit run_limit=2, continue, model reacts", [ToolCallLimitMiddleware(run_limit=2)], stop_on_error=True)
run("ToolCallLimit run_limit=2, error", [ToolCallLimitMiddleware(run_limit=2, exit_behavior="error")])
 
print()
print("thread_limit=2 with a checkpointer, one thread:")
model = AnsweringModel()
agent = create_agent(
    model,
    tools=[],
    checkpointer=InMemorySaver(),
    middleware=[ModelCallLimitMiddleware(thread_limit=2)],
)
for turn in range(1, 4):
    before = model.calls
    result = agent.invoke({"messages": [HumanMessage(f"turn {turn}")]}, {"configurable": {"thread_id": "t1"}})
    last = result["messages"][-1]
    print(f"  turn {turn}: model_calls_this_turn={model.calls - before} last={last.content!r}")
before = model.calls
result = agent.invoke({"messages": [HumanMessage("other thread")]}, {"configurable": {"thread_id": "t2"}})
print(f"  new thread t2: model_calls_this_turn={model.calls - before} last={result['messages'][-1].content!r}")

no limits (recursion_limit=20)               model_calls=10 tool_runs=10  raised GraphRecursionError
ModelCallLimit run_limit=3, end              model_calls=3  tool_runs=3  finished, last message: ai: 'Model call limits exceeded: run limit (3/3)'
ModelCallLimit run_limit=3, error            model_calls=3  tool_runs=3  raised ModelCallLimitExceededError
ToolCallLimit run_limit=2, continue          model_calls=9  tool_runs=2  raised GraphRecursionError
ToolCallLimit run_limit=2, continue, model reacts model_calls=4  tool_runs=2  finished, last message: ai: 'giving up after the tool reported an error'
ToolCallLimit run_limit=2, error             model_calls=3  tool_runs=2  raised ToolCallLimitExceededError

thread_limit=2 with a checkpointer, one thread:
  turn 1: model_calls_this_turn=1 last='answer 1'
  turn 2: model_calls_this_turn=1 last='answer 2'
  turn 3: model_calls_this_turn=0 last='Model call limits exceeded: thread limit (2/2)'
  new thread t2: model_calls_this_turn=1 last='answer 3'

"""How many graph steps does one model -> tool cycle cost once middleware is attached?"""
import sys
from pathlib import Path
 
sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "agent-vs-langgraph"))
 
from common import search
from langchain.agents import create_agent
from langchain.agents.middleware import (
    ModelCallLimitMiddleware,
    ModelRetryMiddleware,
    PIIMiddleware,
    ToolCallLimitMiddleware,
    ToolRetryMiddleware,
    wrap_tool_call,
)
from langchain_core.messages import AIMessage, HumanMessage
from langgraph.errors import GraphRecursionError
from shared import ScriptedChatModel
 
CYCLES = 3
 
 
@wrap_tool_call
def passthrough(request, handler):
    return handler(request)
 
 
def script():
    calls = [
        AIMessage(content="", tool_calls=[{"name": "search", "args": {"query": f"q{i}"}, "id": f"c{i}"}])
        for i in range(CYCLES)
    ]
    return calls + [AIMessage(content="done")]
 
 
def build(middleware):
    return create_agent(ScriptedChatModel(script=script()), tools=[search], middleware=middleware)
 
 
def count_steps(middleware):
    chunks = list(build(middleware).stream({"messages": [HumanMessage("go")]}, stream_mode="updates"))
    return len(chunks)
 
 
def smallest_passing_limit(middleware):
    """Find the smallest recursion_limit under which the whole run completes."""
    for limit in range(1, 100):
        try:
            build(middleware).invoke({"messages": [HumanMessage("go")]}, {"recursion_limit": limit})
            return limit
        except GraphRecursionError:
            continue
    raise RuntimeError("no passing recursion_limit below 100")
 
 
CASES = [
    ("no middleware", []),
    ("ModelCallLimit", [ModelCallLimitMiddleware(run_limit=50)]),
    ("ToolCallLimit", [ToolCallLimitMiddleware(run_limit=50)]),
    ("PII (email, input only)", [PIIMiddleware("email")]),
    ("ModelRetry + ToolRetry (wrap hooks only)", [ModelRetryMiddleware(), ToolRetryMiddleware()]),
    ("custom @wrap_tool_call", [passthrough]),
    ("ModelCallLimit + PII", [ModelCallLimitMiddleware(run_limit=50), PIIMiddleware("email")]),
]
 
print(f"{CYCLES} tool cycles plus the final answer")
for label, middleware in CASES:
    steps = count_steps(middleware)
    limit = smallest_passing_limit(middleware)
    print(f"{label:<42} node_runs={steps:<3} smallest_passing_recursion_limit={limit}")
 
print()
print("nodes that run for one tool cycle plus the final answer")
for label, middleware in (
    ("ModelCallLimit", [ModelCallLimitMiddleware(run_limit=50)]),
    ("ToolCallLimit", [ToolCallLimitMiddleware(run_limit=50)]),
    ("PII", [PIIMiddleware("email")]),
):
    model = ScriptedChatModel(script=[script()[0], script()[-1]])
    agent = create_agent(model, tools=[search], middleware=middleware)
    names = [next(iter(chunk)) for chunk in agent.stream({"messages": [HumanMessage("go")]}, stream_mode="updates")]
    print(f"{label}: {' > '.join(name.split('.')[-1] for name in names)}")

3 tool cycles plus the final answer
no middleware                              node_runs=7   smallest_passing_recursion_limit=8
ModelCallLimit                             node_runs=15  smallest_passing_recursion_limit=16
ToolCallLimit                              node_runs=11  smallest_passing_recursion_limit=12
PII (email, input only)                    node_runs=15  smallest_passing_recursion_limit=16
ModelRetry + ToolRetry (wrap hooks only)   node_runs=7   smallest_passing_recursion_limit=8
custom @wrap_tool_call                     node_runs=7   smallest_passing_recursion_limit=8
ModelCallLimit + PII                       node_runs=23  smallest_passing_recursion_limit=24

nodes that run for one tool cycle plus the final answer
ModelCallLimit: before_model > model > after_model > tools > before_model > model > after_model
ToolCallLimit: model > after_model > tools > model > after_model
PII: before_model > model > after_model > tools > before_model > model > after_model

"""What does PIIMiddleware actually change, and which paths does it leave alone?"""
import sys
from pathlib import Path
from typing import Any, List, Optional
 
sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "agent-vs-langgraph"))
 
from langchain.agents import create_agent
from langchain.agents.middleware import PIIMiddleware
from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, ToolMessage
from langchain_core.outputs import ChatResult
from langchain_core.tools import tool
from shared import ScriptedChatModel
 
 
class RecordingModel(ScriptedChatModel):
    """Scripted model that remembers the message list of every call."""
 
    seen: List[List[BaseMessage]] = []
 
    def _generate(self, messages: List[BaseMessage], stop: Optional[List[str]] = None, run_manager: Any = None, **kwargs: Any) -> ChatResult:
        self.seen.append(list(messages))
        return super()._generate(messages, stop, run_manager, **kwargs)
 
 
@tool
def lookup_customer(customer_id: str) -> str:
    """Return the customer record."""
    return "customer 7 contact: [email protected]"
 
 
def answer_only(text: str) -> RecordingModel:
    return RecordingModel(script=[AIMessage(content=text)])
 
 
def show_input(label: str, middleware, user_text: str) -> None:
    model = answer_only("ok")
    agent = create_agent(model, tools=[], middleware=middleware)
    try:
        result = agent.invoke({"messages": [HumanMessage(user_text)]})
    except Exception as exc:
        print(f"{label:<34} raised {type(exc).__name__}")
        return
    seen = model.seen[0][0].content
    stored = result["messages"][0].content
    print(f"{label:<34} model saw: {seen!r} | stored: {stored!r}")
 
 
print("input handling")
mail = "write to [email protected] please"
show_input("email, redact", [PIIMiddleware("email", strategy="redact")], mail)
show_input("email, mask", [PIIMiddleware("email", strategy="mask")], mail)
show_input("email, hash", [PIIMiddleware("email", strategy="hash")], mail)
show_input("email, block", [PIIMiddleware("email", strategy="block")], mail)
show_input("credit card, mask", [PIIMiddleware("credit_card", strategy="mask")], "card 4111111111111111 expires soon")
show_input(
    "custom regex EMP-\\d{6}, redact",
    [PIIMiddleware("employee_id", detector=r"EMP-\d{6}", strategy="redact")],
    "ticket for EMP-123456",
)
show_input("email rule, text without email", [PIIMiddleware("email")], "no personal data here")
 
print()
print("tool results")
for flag in (False, True):
    model = RecordingModel(
        script=[
            AIMessage(content="", tool_calls=[{"name": "lookup_customer", "args": {"customer_id": "7"}, "id": "c1"}]),
            AIMessage(content="done"),
        ]
    )
    agent = create_agent(
        model,
        tools=[lookup_customer],
        middleware=[PIIMiddleware("email", strategy="redact", apply_to_tool_results=flag)],
    )
    agent.invoke({"messages": [HumanMessage("who is customer 7?")]})
    tool_message = next(m for m in model.seen[1] if isinstance(m, ToolMessage))
    print(f"apply_to_tool_results={flag!s:<5} model saw tool result: {tool_message.content!r}")
 
print()
print("model output")
for flag in (False, True):
    model = answer_only("reach [email protected] for details")
    agent = create_agent(model, tools=[], middleware=[PIIMiddleware("email", strategy="redact", apply_to_output=flag)])
    result = agent.invoke({"messages": [HumanMessage("who can help?")]})
    print(f"apply_to_output={flag!s:<5} final message: {result['messages'][-1].content!r}")

input handling
email, redact                      model saw: 'write to [REDACTED_EMAIL] please' | stored: 'write to [REDACTED_EMAIL] please'
email, mask                        model saw: 'write to alice@****.com please' | stored: 'write to alice@****.com please'
email, hash                        model saw: 'write to <email_hash:ff8d9819> please' | stored: 'write to <email_hash:ff8d9819> please'
email, block                       raised PIIDetectionError
credit card, mask                  model saw: 'card ************1111 expires soon' | stored: 'card ************1111 expires soon'
custom regex EMP-\d{6}, redact     model saw: 'ticket for [REDACTED_EMPLOYEE_ID]' | stored: 'ticket for [REDACTED_EMPLOYEE_ID]'
email rule, text without email     model saw: 'no personal data here' | stored: 'no personal data here'

tool results
apply_to_tool_results=False model saw tool result: 'customer 7 contact: [email protected]'
apply_to_tool_results=True  model saw tool result: 'customer 7 contact: [REDACTED_EMAIL]'

model output
apply_to_output=False final message: 'reach [email protected] for details'
apply_to_output=True  final message: 'reach [REDACTED_EMAIL] for details'

"""Model retry, model fallback and tool retry: attempts, ordering and side effects."""
import sys
from pathlib import Path
 
sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "agent-vs-langgraph"))
 
from common import FlakyModel
from langchain.agents import create_agent
from langchain.agents.middleware import (
    ModelFallbackMiddleware,
    ModelRetryMiddleware,
    ToolErrorMiddleware,
    ToolRetryMiddleware,
)
from langchain_core.messages import AIMessage, HumanMessage
from langchain_core.tools import tool
from shared import ScriptedChatModel
 
ASK = {"messages": [HumanMessage("hello")]}
NO_DELAY = dict(initial_delay=0.0, jitter=False)
 
 
def run_model(label, models, middleware):
    agent = create_agent(models[0], tools=[], middleware=middleware)
    try:
        last = agent.invoke(ASK)["messages"][-1].content
        outcome = f"finished: {last!r}"
    except Exception as exc:
        outcome = f"raised {type(exc).__name__}"
    attempts = ", ".join(f"{m.label}={m.attempts}" for m in models)
    print(f"{label:<44} attempts[{attempts}]  {outcome}")
 
 
print("model retry")
m = FlakyModel(label="primary", failures=2)
run_model("no middleware, 2 failures", [m], [])
m = FlakyModel(label="primary", failures=2)
run_model("ModelRetry max_retries=2, 2 failures", [m], [ModelRetryMiddleware(max_retries=2, **NO_DELAY)])
m = FlakyModel(label="primary", failures=3)
run_model("ModelRetry max_retries=2, 3 failures", [m], [ModelRetryMiddleware(max_retries=2, **NO_DELAY)])
m = FlakyModel(label="primary", failures=3)
run_model(
    "ModelRetry on_failure=error, 3 failures",
    [m],
    [ModelRetryMiddleware(max_retries=2, on_failure="error", **NO_DELAY)],
)
m = FlakyModel(label="primary", failures=9, error=ValueError)
run_model("ModelRetry default retry_on, ValueError", [m], [ModelRetryMiddleware(max_retries=2, on_failure="error", **NO_DELAY)])
m = FlakyModel(label="primary", failures=9, error=ValueError)
run_model(
    "ModelRetry retry_on=(ConnectionError,), ValueError",
    [m],
    [ModelRetryMiddleware(max_retries=2, retry_on=(ConnectionError,), on_failure="error", **NO_DELAY)],
)
 
print()
print("model fallback and ordering (primary always fails)")
p, f = FlakyModel(label="primary", failures=99), FlakyModel(label="fallback")
run_model("ModelFallback only", [p, f], [ModelFallbackMiddleware(f)])
p, f = FlakyModel(label="primary", failures=99), FlakyModel(label="fallback")
run_model(
    "[ModelRetry(1), ModelFallback]",
    [p, f],
    [ModelRetryMiddleware(max_retries=1, **NO_DELAY), ModelFallbackMiddleware(f)],
)
p, f = FlakyModel(label="primary", failures=99), FlakyModel(label="fallback")
run_model(
    "[ModelFallback, ModelRetry(1)]",
    [p, f],
    [ModelFallbackMiddleware(f), ModelRetryMiddleware(max_retries=1, **NO_DELAY)],
)
p, f = FlakyModel(label="primary", failures=99), FlakyModel(label="fallback")
run_model(
    "[ModelFallback, ModelRetry(1, on_failure=error)]",
    [p, f],
    [ModelFallbackMiddleware(f), ModelRetryMiddleware(max_retries=1, on_failure="error", **NO_DELAY)],
)
p, f1, f2 = (FlakyModel(label="primary", failures=99), FlakyModel(label="fallback1", failures=99), FlakyModel(label="fallback2"))
run_model("ModelFallback with two fallbacks", [p, f1, f2], [ModelFallbackMiddleware(f1, f2)])
p, f = FlakyModel(label="primary", failures=99), FlakyModel(label="fallback", failures=99)
run_model("ModelFallback, every model fails", [p, f], [ModelFallbackMiddleware(f)])
 
SIDE_EFFECTS = []
 
 
@tool
def charge(order_id: str) -> str:
    """Charge the customer for an order."""
    SIDE_EFFECTS.append(order_id)  # the side effect happens before the failure
    if len(SIDE_EFFECTS) <= FAIL_FIRST:
        raise RuntimeError("payment gateway timeout")
    return f"charged {order_id}"
 
 
def run_tool(label, middleware, fail_first):
    global FAIL_FIRST
    FAIL_FIRST = fail_first
    SIDE_EFFECTS.clear()
    model = ScriptedChatModel(
        script=[
            AIMessage(content="", tool_calls=[{"name": "charge", "args": {"order_id": "A1"}, "id": "c1"}]),
            AIMessage(content="done"),
        ]
    )
    agent = create_agent(model, tools=[charge], middleware=middleware)
    try:
        messages = agent.invoke(ASK)["messages"]
        tool_message = next(m for m in messages if m.type == "tool")
        outcome = f"finished, tool message ({tool_message.status}): {tool_message.content!r}"
    except Exception as exc:
        outcome = f"raised {type(exc).__name__}"
    print(f"{label:<46} real_executions={len(SIDE_EFFECTS)}  {outcome}")
 
 
print()
print("tool retry (one logical call, side effect before the failure)")
run_tool("no middleware, fails once", [], fail_first=1)
run_tool("ToolRetry max_retries=2, fails twice", [ToolRetryMiddleware(max_retries=2, **NO_DELAY)], fail_first=2)
run_tool("ToolRetry max_retries=2, always fails", [ToolRetryMiddleware(max_retries=2, **NO_DELAY)], fail_first=99)
run_tool(
    "ToolRetry on_failure=error, always fails",
    [ToolRetryMiddleware(max_retries=2, on_failure="error", **NO_DELAY)],
    fail_first=99,
)
run_tool(
    "ToolRetry max_retries=0, always fails",
    [ToolRetryMiddleware(max_retries=0, **NO_DELAY)],
    fail_first=99,
)
 
 
def on_error(exc, request):
    return f"{request.tool_call['name']} failed with {type(exc).__name__}"
 
 
run_tool(
    "[ToolError, ToolRetry(on_failure=error)], fails",
    [ToolErrorMiddleware(on_error), ToolRetryMiddleware(max_retries=2, on_failure="error", **NO_DELAY)],
    fail_first=99,
)

model retry
no middleware, 2 failures                    attempts[primary=1]  raised ConnectionError
ModelRetry max_retries=2, 2 failures         attempts[primary=3]  finished: 'answered by primary'
ModelRetry max_retries=2, 3 failures         attempts[primary=3]  finished: 'Model call failed after 3 attempts with ConnectionError: primary attempt 3 failed'
ModelRetry on_failure=error, 3 failures      attempts[primary=3]  raised ConnectionError
ModelRetry default retry_on, ValueError      attempts[primary=3]  raised ValueError
ModelRetry retry_on=(ConnectionError,), ValueError attempts[primary=1]  raised ValueError

model fallback and ordering (primary always fails)
ModelFallback only                           attempts[primary=1, fallback=1]  finished: 'answered by fallback'
[ModelRetry(1), ModelFallback]               attempts[primary=1, fallback=1]  finished: 'answered by fallback'
[ModelFallback, ModelRetry(1)]               attempts[primary=2, fallback=0]  finished: 'Model call failed after 2 attempts with ConnectionError: primary attempt 2 failed'
[ModelFallback, ModelRetry(1, on_failure=error)] attempts[primary=2, fallback=1]  finished: 'answered by fallback'
ModelFallback with two fallbacks             attempts[primary=1, fallback1=1, fallback2=1]  finished: 'answered by fallback2'
ModelFallback, every model fails             attempts[primary=1, fallback=1]  raised ConnectionError

tool retry (one logical call, side effect before the failure)
no middleware, fails once                      real_executions=1  raised RuntimeError
ToolRetry max_retries=2, fails twice           real_executions=3  finished, tool message (success): 'charged A1'
ToolRetry max_retries=2, always fails          real_executions=3  finished, tool message (error): "Tool 'charge' failed after 3 attempts with RuntimeError: payment gateway timeout. Please try again."
ToolRetry on_failure=error, always fails       real_executions=3  raised RuntimeError
ToolRetry max_retries=0, always fails          real_executions=1  finished, tool message (error): "Tool 'charge' failed after 1 attempt with RuntimeError: payment gateway timeout. Please try again."
[ToolError, ToolRetry(on_failure=error)], fails real_executions=3  finished, tool message (error): 'charge failed with RuntimeError'