langchain 1.0.5__py3-none-any.whl → 1.2.4__py3-none-any.whl

langchain/agents/middleware/_execution.py

@@ -15,8 +15,10 @@ from pathlib import Path
 
 try:  # pragma: no cover - optional dependency on POSIX platforms
     import resource
+
+    _HAS_RESOURCE = True
 except ImportError:  # pragma: no cover - non-POSIX systems
-    resource = None  # type: ignore[assignment]
+    _HAS_RESOURCE = False
 
 
 SHELL_TEMP_PREFIX = "langchain-shell-"
@@ -56,11 +58,12 @@ class BaseExecutionPolicy(abc.ABC):
     """Configuration contract for persistent shell sessions.
 
     Concrete subclasses encapsulate how a shell process is launched and constrained.
+
     Each policy documents its security guarantees and the operating environments in
-    which it is appropriate. Use :class:`HostExecutionPolicy` for trusted, same-host
-    execution; :class:`CodexSandboxExecutionPolicy` when the Codex CLI sandbox is
-    available and you want additional syscall restrictions; and
-    :class:`DockerExecutionPolicy` for container-level isolation using Docker.
+    which it is appropriate. Use `HostExecutionPolicy` for trusted, same-host execution;
+    `CodexSandboxExecutionPolicy` when the Codex CLI sandbox is available and you want
+    additional syscall restrictions; and `DockerExecutionPolicy` for container-level
+    isolation using Docker.
     """
 
     command_timeout: float = 30.0
@@ -91,13 +94,13 @@ class HostExecutionPolicy(BaseExecutionPolicy):
 
     This policy is best suited for trusted or single-tenant environments (CI jobs,
     developer workstations, pre-sandboxed containers) where the agent must access the
-    host filesystem and tooling without additional isolation. It enforces optional CPU
-    and memory limits to prevent runaway commands but offers **no** filesystem or network
+    host filesystem and tooling without additional isolation. Enforces optional CPU and
+    memory limits to prevent runaway commands but offers **no** filesystem or network
     sandboxing; commands can modify anything the process user can reach.
 
-    On Linux platforms resource limits are applied with ``resource.prlimit`` after the
-    shell starts. On macOS, where ``prlimit`` is unavailable, limits are set in a
-    ``preexec_fn`` before ``exec``. In both cases the shell runs in its own process group
+    On Linux platforms resource limits are applied with `resource.prlimit` after the
+    shell starts. On macOS, where `prlimit` is unavailable, limits are set in a
+    `preexec_fn` before `exec`. In both cases the shell runs in its own process group
     so timeouts can terminate the full subtree.
     """
 
@@ -118,7 +121,7 @@ class HostExecutionPolicy(BaseExecutionPolicy):
         self._limits_requested = any(
             value is not None for value in (self.cpu_time_seconds, self.memory_bytes)
         )
-        if self._limits_requested and resource is None:
+        if self._limits_requested and not _HAS_RESOURCE:
             msg = (
                 "HostExecutionPolicy cpu/memory limits require the Python 'resource' module. "
                 "Either remove the limits or run on a POSIX platform."
@@ -162,11 +165,9 @@ class HostExecutionPolicy(BaseExecutionPolicy):
     def _apply_post_spawn_limits(self, process: subprocess.Popen[str]) -> None:
         if not self._limits_requested or not self._can_use_prlimit():
             return
-        if resource is None:  # pragma: no cover - defensive
+        if not _HAS_RESOURCE:  # pragma: no cover - defensive
             return
         pid = process.pid
-        if pid is None:
-            return
         try:
             prlimit = typing.cast("typing.Any", resource).prlimit
             if self.cpu_time_seconds is not None:
@@ -183,11 +184,7 @@ class HostExecutionPolicy(BaseExecutionPolicy):
 
     @staticmethod
     def _can_use_prlimit() -> bool:
-        return (
-            resource is not None
-            and hasattr(resource, "prlimit")
-            and sys.platform.startswith("linux")
-        )
+        return _HAS_RESOURCE and hasattr(resource, "prlimit") and sys.platform.startswith("linux")
 
 
 @dataclass
@@ -199,9 +196,9 @@ class CodexSandboxExecutionPolicy(BaseExecutionPolicy):
     (Linux) profiles. Commands still run on the host, but within the sandbox requested by
     the CLI. If the Codex binary is unavailable or the runtime lacks the required
     kernel features (e.g., Landlock inside some containers), process startup fails with a
-    :class:`RuntimeError`.
+    `RuntimeError`.
 
-    Configure sandbox behaviour via ``config_overrides`` to align with your Codex CLI
+    Configure sandbox behavior via `config_overrides` to align with your Codex CLI
     profile. This policy does not add its own resource limits; combine it with
     host-level guards (cgroups, container resource limits) as needed.
     """
@@ -250,9 +247,9 @@ class CodexSandboxExecutionPolicy(BaseExecutionPolicy):
             return self.platform
         if sys.platform.startswith("linux"):
             return "linux"
-        if sys.platform == "darwin":
+        if sys.platform == "darwin":  # type: ignore[unreachable, unused-ignore]
             return "macos"
-        msg = (
+        msg = (  # type: ignore[unreachable, unused-ignore]
             "Codex sandbox policy could not determine a supported platform; "
             "set 'platform' explicitly."
         )
@@ -271,17 +268,17 @@ class DockerExecutionPolicy(BaseExecutionPolicy):
     """Run the shell inside a dedicated Docker container.
 
     Choose this policy when commands originate from untrusted users or you require
-    strong isolation between sessions. By default the workspace is bind-mounted only when
-    it refers to an existing non-temporary directory; ephemeral sessions run without a
-    mount to minimise host exposure. The container's network namespace is disabled by
-    default (``--network none``) and you can enable further hardening via
-    ``read_only_rootfs`` and ``user``.
+    strong isolation between sessions. By default the workspace is bind-mounted only
+    when it refers to an existing non-temporary directory; ephemeral sessions run
+    without a mount to minimise host exposure. The container's network namespace is
+    disabled by default (`--network none`) and you can enable further hardening via
+    `read_only_rootfs` and `user`.
 
     The security guarantees depend on your Docker daemon configuration. Run the agent on
-    a host where Docker is locked down (rootless mode, AppArmor/SELinux, etc.) and review
-    any additional volumes or capabilities passed through ``extra_run_args``. The default
-    image is ``python:3.12-alpine3.19``; supply a custom image if you need preinstalled
-    tooling.
+    a host where Docker is locked down (rootless mode, AppArmor/SELinux, etc.) and
+    review any additional volumes or capabilities passed through ``extra_run_args``. The
+    default image is `python:3.12-alpine3.19`; supply a custom image if you need
+    preinstalled tooling.
     """
 
     binary: str = "docker"

langchain/agents/middleware/_redaction.py

@@ -4,6 +4,7 @@ from __future__ import annotations
 
 import hashlib
 import ipaddress
+import operator
 import re
 from collections.abc import Callable, Sequence
 from dataclasses import dataclass
@@ -47,7 +48,14 @@ Detector = Callable[[str], list[PIIMatch]]
 
 
 def detect_email(content: str) -> list[PIIMatch]:
-    """Detect email addresses in content."""
+    """Detect email addresses in content.
+
+    Args:
+        content: The text content to scan for email addresses.
+
+    Returns:
+        A list of detected email matches.
+    """
     pattern = r"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b"
     return [
         PIIMatch(
@@ -61,7 +69,14 @@ def detect_email(content: str) -> list[PIIMatch]:
 
 
 def detect_credit_card(content: str) -> list[PIIMatch]:
-    """Detect credit card numbers in content using Luhn validation."""
+    """Detect credit card numbers in content using Luhn validation.
+
+    Args:
+        content: The text content to scan for credit card numbers.
+
+    Returns:
+        A list of detected credit card matches.
+    """
     pattern = r"\b\d{4}[\s-]?\d{4}[\s-]?\d{4}[\s-]?\d{4}\b"
     matches = []
 
@@ -81,7 +96,14 @@ def detect_credit_card(content: str) -> list[PIIMatch]:
 
 
 def detect_ip(content: str) -> list[PIIMatch]:
-    """Detect IPv4 or IPv6 addresses in content."""
+    """Detect IPv4 or IPv6 addresses in content.
+
+    Args:
+        content: The text content to scan for IP addresses.
+
+    Returns:
+        A list of detected IP address matches.
+    """
     matches: list[PIIMatch] = []
     ipv4_pattern = r"\b(?:[0-9]{1,3}\.){3}[0-9]{1,3}\b"
 
@@ -104,7 +126,14 @@ def detect_ip(content: str) -> list[PIIMatch]:
 
 
 def detect_mac_address(content: str) -> list[PIIMatch]:
-    """Detect MAC addresses in content."""
+    """Detect MAC addresses in content.
+
+    Args:
+        content: The text content to scan for MAC addresses.
+
+    Returns:
+        A list of detected MAC address matches.
+    """
     pattern = r"\b([0-9A-Fa-f]{2}[:-]){5}[0-9A-Fa-f]{2}\b"
     return [
         PIIMatch(
@@ -118,7 +147,14 @@ def detect_mac_address(content: str) -> list[PIIMatch]:
 
 
 def detect_url(content: str) -> list[PIIMatch]:
-    """Detect URLs in content using regex and stdlib validation."""
+    """Detect URLs in content using regex and stdlib validation.
+
+    Args:
+        content: The text content to scan for URLs.
+
+    Returns:
+        A list of detected URL matches.
+    """
     matches: list[PIIMatch] = []
 
     # Pattern 1: URLs with scheme (http:// or https://)
@@ -127,7 +163,7 @@ def detect_url(content: str) -> list[PIIMatch]:
     for match in re.finditer(scheme_pattern, content):
         url = match.group()
         result = urlparse(url)
-        if result.scheme in ("http", "https") and result.netloc:
+        if result.scheme in {"http", "https"} and result.netloc:
             matches.append(
                 PIIMatch(
                     type="url",
@@ -179,11 +215,14 @@ BUILTIN_DETECTORS: dict[str, Detector] = {
 }
 """Registry of built-in detectors keyed by type name."""
 
+_CARD_NUMBER_MIN_DIGITS = 13
+_CARD_NUMBER_MAX_DIGITS = 19
+
 
 def _passes_luhn(card_number: str) -> bool:
     """Validate credit card number using the Luhn checksum."""
     digits = [int(d) for d in card_number if d.isdigit()]
-    if not 13 <= len(digits) <= 19:
+    if not _CARD_NUMBER_MIN_DIGITS <= len(digits) <= _CARD_NUMBER_MAX_DIGITS:
         return False
 
     checksum = 0
@@ -191,7 +230,7 @@ def _passes_luhn(card_number: str) -> bool:
         value = digit
         if index % 2 == 1:
             value *= 2
-            if value > 9:
+            if value > 9:  # noqa: PLR2004
                 value -= 9
         checksum += value
     return checksum % 10 == 0
@@ -199,24 +238,28 @@ def _passes_luhn(card_number: str) -> bool:
 
 def _apply_redact_strategy(content: str, matches: list[PIIMatch]) -> str:
     result = content
-    for match in sorted(matches, key=lambda item: item["start"], reverse=True):
+    for match in sorted(matches, key=operator.itemgetter("start"), reverse=True):
         replacement = f"[REDACTED_{match['type'].upper()}]"
         result = result[: match["start"]] + replacement + result[match["end"] :]
     return result
 
 
+_UNMASKED_CHAR_NUMBER = 4
+_IPV4_PARTS_NUMBER = 4
+
+
 def _apply_mask_strategy(content: str, matches: list[PIIMatch]) -> str:
     result = content
-    for match in sorted(matches, key=lambda item: item["start"], reverse=True):
+    for match in sorted(matches, key=operator.itemgetter("start"), reverse=True):
         value = match["value"]
         pii_type = match["type"]
         if pii_type == "email":
             parts = value.split("@")
-            if len(parts) == 2:
+            if len(parts) == 2:  # noqa: PLR2004
                 domain_parts = parts[1].split(".")
                 masked = (
                     f"{parts[0]}@****.{domain_parts[-1]}"
-                    if len(domain_parts) >= 2
+                    if len(domain_parts) > 1
                     else f"{parts[0]}@****"
                 )
             else:
@@ -225,12 +268,15 @@ def _apply_mask_strategy(content: str, matches: list[PIIMatch]) -> str:
             digits_only = "".join(c for c in value if c.isdigit())
             separator = "-" if "-" in value else " " if " " in value else ""
             if separator:
-                masked = f"****{separator}****{separator}****{separator}{digits_only[-4:]}"
+                masked = (
+                    f"****{separator}****{separator}****{separator}"
+                    f"{digits_only[-_UNMASKED_CHAR_NUMBER:]}"
+                )
             else:
-                masked = f"************{digits_only[-4:]}"
+                masked = f"************{digits_only[-_UNMASKED_CHAR_NUMBER:]}"
         elif pii_type == "ip":
             octets = value.split(".")
-            masked = f"*.*.*.{octets[-1]}" if len(octets) == 4 else "****"
+            masked = f"*.*.*.{octets[-1]}" if len(octets) == _IPV4_PARTS_NUMBER else "****"
         elif pii_type == "mac_address":
             separator = ":" if ":" in value else "-"
             masked = (
@@ -239,14 +285,18 @@ def _apply_mask_strategy(content: str, matches: list[PIIMatch]) -> str:
         elif pii_type == "url":
             masked = "[MASKED_URL]"
         else:
-            masked = f"****{value[-4:]}" if len(value) > 4 else "****"
+            masked = (
+                f"****{value[-_UNMASKED_CHAR_NUMBER:]}"
+                if len(value) > _UNMASKED_CHAR_NUMBER
+                else "****"
+            )
         result = result[: match["start"]] + masked + result[match["end"] :]
     return result
 
 
 def _apply_hash_strategy(content: str, matches: list[PIIMatch]) -> str:
     result = content
-    for match in sorted(matches, key=lambda item: item["start"], reverse=True):
+    for match in sorted(matches, key=operator.itemgetter("start"), reverse=True):
         digest = hashlib.sha256(match["value"].encode()).hexdigest()[:8]
         replacement = f"<{match['type']}_hash:{digest}>"
         result = result[: match["start"]] + replacement + result[match["end"] :]
@@ -258,7 +308,20 @@ def apply_strategy(
     matches: list[PIIMatch],
     strategy: RedactionStrategy,
 ) -> str:
-    """Apply the configured strategy to matches within content."""
+    """Apply the configured strategy to matches within content.
+
+    Args:
+        content: The content to apply strategy to.
+        matches: List of detected PII matches.
+        strategy: The redaction strategy to apply.
+
+    Returns:
+        The content with the strategy applied.
+
+    Raises:
+        PIIDetectionError: If the strategy is `'block'` and matches are found.
+        ValueError: If the strategy is unknown.
+    """
     if not matches:
         return content
     if strategy == "redact":
@@ -269,12 +332,24 @@ def apply_strategy(
         return _apply_hash_strategy(content, matches)
     if strategy == "block":
         raise PIIDetectionError(matches[0]["type"], matches)
-    msg = f"Unknown redaction strategy: {strategy}"
+    msg = f"Unknown redaction strategy: {strategy}"  # type: ignore[unreachable]
     raise ValueError(msg)
 
 
 def resolve_detector(pii_type: str, detector: Detector | str | None) -> Detector:
-    """Return a callable detector for the given configuration."""
+    """Return a callable detector for the given configuration.
+
+    Args:
+        pii_type: The PII type name.
+        detector: Optional custom detector or regex pattern. If `None`, a built-in detector
+            for the given PII type will be used.
+
+    Returns:
+        The resolved detector.
+
+    Raises:
+        ValueError: If an unknown PII type is specified without a custom detector or regex.
+    """
     if detector is None:
         if pii_type not in BUILTIN_DETECTORS:
             msg = (
@@ -310,7 +385,11 @@ class RedactionRule:
     detector: Detector | str | None = None
 
     def resolve(self) -> ResolvedRedactionRule:
-        """Resolve runtime detector and return an immutable rule."""
+        """Resolve runtime detector and return an immutable rule.
+
+        Returns:
+            The resolved redaction rule.
+        """
         resolved_detector = resolve_detector(self.pii_type, self.detector)
         return ResolvedRedactionRule(
             pii_type=self.pii_type,
@@ -328,7 +407,14 @@ class ResolvedRedactionRule:
     detector: Detector
 
     def apply(self, content: str) -> tuple[str, list[PIIMatch]]:
-        """Apply this rule to content, returning new content and matches."""
+        """Apply this rule to content, returning new content and matches.
+
+        Args:
+            content: The text content to scan and redact.
+
+        Returns:
+            A tuple of (updated content, list of detected matches).
+        """
         matches = self.detector(content)
         if not matches:
             return content, []

langchain/agents/middleware/_retry.py

@@ -0,0 +1,123 @@
+"""Shared retry utilities for agent middleware.
+
+This module contains common constants, utilities, and logic used by both
+model and tool retry middleware implementations.
+"""
+
+from __future__ import annotations
+
+import random
+from collections.abc import Callable
+from typing import Literal
+
+# Type aliases
+RetryOn = tuple[type[Exception], ...] | Callable[[Exception], bool]
+"""Type for specifying which exceptions to retry on.
+
+Can be either:
+- A tuple of exception types to retry on (based on `isinstance` checks)
+- A callable that takes an exception and returns `True` if it should be retried
+"""
+
+OnFailure = Literal["error", "continue"] | Callable[[Exception], str]
+"""Type for specifying failure handling behavior.
+
+Can be either:
+- A literal action string (`'error'` or `'continue'`)
+    - `'error'`: Re-raise the exception, stopping agent execution.
+    - `'continue'`: Inject a message with the error details, allowing the agent to continue.
+       For tool retries, a `ToolMessage` with the error details will be injected.
+       For model retries, an `AIMessage` with the error details will be returned.
+- A callable that takes an exception and returns a string for error message content
+"""
+
+
+def validate_retry_params(
+    max_retries: int,
+    initial_delay: float,
+    max_delay: float,
+    backoff_factor: float,
+) -> None:
+    """Validate retry parameters.
+
+    Args:
+        max_retries: Maximum number of retry attempts.
+        initial_delay: Initial delay in seconds before first retry.
+        max_delay: Maximum delay in seconds between retries.
+        backoff_factor: Multiplier for exponential backoff.
+
+    Raises:
+        ValueError: If any parameter is invalid (negative values).
+    """
+    if max_retries < 0:
+        msg = "max_retries must be >= 0"
+        raise ValueError(msg)
+    if initial_delay < 0:
+        msg = "initial_delay must be >= 0"
+        raise ValueError(msg)
+    if max_delay < 0:
+        msg = "max_delay must be >= 0"
+        raise ValueError(msg)
+    if backoff_factor < 0:
+        msg = "backoff_factor must be >= 0"
+        raise ValueError(msg)
+
+
+def should_retry_exception(
+    exc: Exception,
+    retry_on: RetryOn,
+) -> bool:
+    """Check if an exception should trigger a retry.
+
+    Args:
+        exc: The exception that occurred.
+        retry_on: Either a tuple of exception types to retry on, or a callable
+            that takes an exception and returns `True` if it should be retried.
+
+    Returns:
+        `True` if the exception should be retried, `False` otherwise.
+    """
+    if callable(retry_on):
+        return retry_on(exc)
+    return isinstance(exc, retry_on)
+
+
+def calculate_delay(
+    retry_number: int,
+    *,
+    backoff_factor: float,
+    initial_delay: float,
+    max_delay: float,
+    jitter: bool,
+) -> float:
+    """Calculate delay for a retry attempt with exponential backoff and optional jitter.
+
+    Args:
+        retry_number: The retry attempt number (0-indexed).
+        backoff_factor: Multiplier for exponential backoff.
+
+            Set to `0.0` for constant delay.
+        initial_delay: Initial delay in seconds before first retry.
+        max_delay: Maximum delay in seconds between retries.
+
+            Caps exponential backoff growth.
+        jitter: Whether to add random jitter to delay to avoid thundering herd.
+
+    Returns:
+        Delay in seconds before next retry.
+    """
+    if backoff_factor == 0.0:
+        delay = initial_delay
+    else:
+        delay = initial_delay * (backoff_factor**retry_number)
+
+    # Cap at max_delay
+    delay = min(delay, max_delay)
+
+    if jitter and delay > 0:
+        jitter_amount = delay * 0.25  # ±25% jitter
+        delay += random.uniform(-jitter_amount, jitter_amount)  # noqa: S311
+        # Ensure delay is not negative after jitter
+        delay = max(0, delay)
+
+    return delay

langchain/agents/middleware/context_editing.py

@@ -1,14 +1,16 @@
 """Context editing middleware.
 
-This middleware mirrors Anthropic's context editing capabilities by clearing
-older tool results once the conversation grows beyond a configurable token
-threshold. The implementation is intentionally model-agnostic so it can be used
-with any LangChain chat model.
+Mirrors Anthropic's context editing capabilities by clearing older tool results once the
+conversation grows beyond a configurable token threshold.
+
+The implementation is intentionally model-agnostic so it can be used with any LangChain
+chat model.
 """
 
 from __future__ import annotations
 
 from collections.abc import Awaitable, Callable, Iterable, Sequence
+from copy import deepcopy
 from dataclasses import dataclass
 from typing import Literal
 
@@ -16,7 +18,6 @@ from langchain_core.messages import (
     AIMessage,
     AnyMessage,
     BaseMessage,
-    SystemMessage,
     ToolMessage,
 )
 from langchain_core.messages.utils import count_tokens_approximately
@@ -151,8 +152,8 @@ class ClearToolUsesEdit(ContextEdit):
 
         return
 
+    @staticmethod
     def _build_cleared_tool_input_message(
-        self,
         message: AIMessage,
         tool_call_id: str,
     ) -> AIMessage:
@@ -182,11 +183,13 @@ class ClearToolUsesEdit(ContextEdit):
 
 
 class ContextEditingMiddleware(AgentMiddleware):
-    """Automatically prunes tool results to manage context size.
+    """Automatically prune tool results to manage context size.
 
-    The middleware applies a sequence of edits when the total input token count
-    exceeds configured thresholds. Currently the `ClearToolUsesEdit` strategy is
-    supported, aligning with Anthropic's `clear_tool_uses_20250919` behaviour.
+    The middleware applies a sequence of edits when the total input token count exceeds
+    configured thresholds.
+
+    Currently the `ClearToolUsesEdit` strategy is supported, aligning with Anthropic's
+    `clear_tool_uses_20250919` behavior [(read more)](https://platform.claude.com/docs/en/agents-and-tools/tool-use/memory-tool).
     """
 
     edits: list[ContextEdit]
@@ -198,11 +201,12 @@ class ContextEditingMiddleware(AgentMiddleware):
         edits: Iterable[ContextEdit] | None = None,
         token_count_method: Literal["approximate", "model"] = "approximate",  # noqa: S107
     ) -> None:
-        """Initializes a context editing middleware instance.
+        """Initialize an instance of context editing middleware.
 
         Args:
-            edits: Sequence of edit strategies to apply. Defaults to a single
-                `ClearToolUsesEdit` mirroring Anthropic defaults.
+            edits: Sequence of edit strategies to apply.
+
+                Defaults to a single `ClearToolUsesEdit` mirroring Anthropic defaults.
             token_count_method: Whether to use approximate token counting
                 (faster, less accurate) or exact counting implemented by the
                 chat model (potentially slower, more accurate).
@@ -216,7 +220,16 @@ class ContextEditingMiddleware(AgentMiddleware):
         request: ModelRequest,
         handler: Callable[[ModelRequest], ModelResponse],
     ) -> ModelCallResult:
-        """Apply context edits before invoking the model via handler."""
+        """Apply context edits before invoking the model via handler.
+
+        Args:
+            request: Model request to execute (includes state and runtime).
+            handler: Async callback that executes the model request and returns
+                `ModelResponse`.
+
+        Returns:
+            The result of invoking the handler with potentially edited messages.
+        """
         if not request.messages:
             return handler(request)
 
@@ -224,27 +237,36 @@ class ContextEditingMiddleware(AgentMiddleware):
 
             def count_tokens(messages: Sequence[BaseMessage]) -> int:
                 return count_tokens_approximately(messages)
+
         else:
-            system_msg = (
-                [SystemMessage(content=request.system_prompt)] if request.system_prompt else []
-            )
+            system_msg = [request.system_message] if request.system_message else []
 
             def count_tokens(messages: Sequence[BaseMessage]) -> int:
                 return request.model.get_num_tokens_from_messages(
                     system_msg + list(messages), request.tools
                 )
 
+        edited_messages = deepcopy(list(request.messages))
         for edit in self.edits:
-            edit.apply(request.messages, count_tokens=count_tokens)
+            edit.apply(edited_messages, count_tokens=count_tokens)
 
-        return handler(request)
+        return handler(request.override(messages=edited_messages))
 
     async def awrap_model_call(
         self,
         request: ModelRequest,
         handler: Callable[[ModelRequest], Awaitable[ModelResponse]],
     ) -> ModelCallResult:
-        """Apply context edits before invoking the model via handler (async version)."""
+        """Apply context edits before invoking the model via handler.
+
+        Args:
+            request: Model request to execute (includes state and runtime).
+            handler: Async callback that executes the model request and returns
+                `ModelResponse`.
+
+        Returns:
+            The result of invoking the handler with potentially edited messages.
+        """
         if not request.messages:
             return await handler(request)
 
@@ -252,20 +274,20 @@ class ContextEditingMiddleware(AgentMiddleware):
 
             def count_tokens(messages: Sequence[BaseMessage]) -> int:
                 return count_tokens_approximately(messages)
+
         else:
-            system_msg = (
-                [SystemMessage(content=request.system_prompt)] if request.system_prompt else []
-            )
+            system_msg = [request.system_message] if request.system_message else []
 
             def count_tokens(messages: Sequence[BaseMessage]) -> int:
                 return request.model.get_num_tokens_from_messages(
                     system_msg + list(messages), request.tools
                 )
 
+        edited_messages = deepcopy(list(request.messages))
         for edit in self.edits:
-            edit.apply(request.messages, count_tokens=count_tokens)
+            edit.apply(edited_messages, count_tokens=count_tokens)
 
-        return await handler(request)
+        return await handler(request.override(messages=edited_messages))
 
 
 __all__ = [