langchain 1.1.2__tar.gz → 1.2.2__tar.gz

{langchain-1.1.2 → langchain-1.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: langchain
-Version: 1.1.2
+Version: 1.2.2
 Summary: Building applications with LLMs through composability
 Project-URL: Homepage, https://docs.langchain.com/
 Project-URL: Documentation, https://reference.langchain.com/python/langchain/langchain/
@@ -12,7 +12,7 @@ Project-URL: Reddit, https://www.reddit.com/r/LangChain/
 License: MIT
 License-File: LICENSE
 Requires-Python: <4.0.0,>=3.10.0
-Requires-Dist: langchain-core<2.0.0,>=1.1.0
+Requires-Dist: langchain-core<2.0.0,>=1.2.1
 Requires-Dist: langgraph<1.1.0,>=1.0.2
 Requires-Dist: pydantic<3.0.0,>=2.7.4
 Provides-Extra: anthropic

{langchain-1.1.2 → langchain-1.2.2}/langchain/__init__.py

@@ -1,3 +1,3 @@
 """Main entrypoint into LangChain."""
 
-__version__ = "1.1.2"
+__version__ = "1.2.2"

{langchain-1.1.2 → langchain-1.2.2}/langchain/agents/factory.py

@@ -20,9 +20,7 @@ from langgraph._internal._runnable import RunnableCallable
 from langgraph.constants import END, START
 from langgraph.graph.state import StateGraph
 from langgraph.prebuilt.tool_node import ToolCallWithContext, ToolNode
-from langgraph.runtime import Runtime  # noqa: TC002
 from langgraph.types import Command, Send
-from langgraph.typing import ContextT  # noqa: TC002
 from typing_extensions import NotRequired, Required, TypedDict
 
 from langchain.agents.middleware.types import (
@@ -56,8 +54,10 @@ if TYPE_CHECKING:
     from langchain_core.runnables import Runnable
     from langgraph.cache.base import BaseCache
     from langgraph.graph.state import CompiledStateGraph
+    from langgraph.runtime import Runtime
     from langgraph.store.base import BaseStore
     from langgraph.types import Checkpointer
+    from langgraph.typing import ContextT
 
     from langchain.agents.middleware.types import ToolCallRequest, ToolCallWrapper
 
@@ -314,7 +314,7 @@ def _resolve_schema(schemas: set[type], schema_name: str, omit_flag: str | None
 def _extract_metadata(type_: type) -> list:
     """Extract metadata from a field type, handling Required/NotRequired and Annotated wrappers."""
     # Handle Required[Annotated[...]] or NotRequired[Annotated[...]]
-    if get_origin(type_) in (Required, NotRequired):
+    if get_origin(type_) in {Required, NotRequired}:
         inner_type = get_args(type_)[0]
         if get_origin(inner_type) is Annotated:
             return list(get_args(inner_type)[1:])
@@ -538,7 +538,7 @@ def _chain_async_tool_call_wrappers(
     return result
 
 
-def create_agent(  # noqa: PLR0915
+def create_agent(
     model: str | BaseChatModel,
     tools: Sequence[BaseTool | Callable | dict[str, Any]] | None = None,
     *,
@@ -791,9 +791,9 @@ def create_agent(  # noqa: PLR0915
         default_tools = list(built_in_tools)
 
     # validate middleware
-    assert len({m.name for m in middleware}) == len(middleware), (  # noqa: S101
-        "Please remove duplicate middleware instances."
-    )
+    if len({m.name for m in middleware}) != len(middleware):
+        msg = "Please remove duplicate middleware instances."
+        raise AssertionError(msg)
     middleware_w_before_agent = [
         m
         for m in middleware
@@ -886,12 +886,12 @@ def create_agent(  # noqa: PLR0915
                 )
                 try:
                     structured_response = provider_strategy_binding.parse(output)
-                except Exception as exc:  # noqa: BLE001
+                except Exception as exc:
                     schema_name = getattr(
                         effective_response_format.schema_spec.schema, "__name__", "response_format"
                     )
                     validation_error = StructuredOutputValidationError(schema_name, exc, output)
-                    raise validation_error
+                    raise validation_error from exc
                 else:
                     return {"messages": [output], "structured_response": structured_response}
             return {"messages": [output]}
@@ -937,8 +937,7 @@ def create_agent(  # noqa: PLR0915
 
                     tool_message_content = (
                         effective_response_format.tool_message_content
-                        if effective_response_format.tool_message_content
-                        else f"Returning structured response: {structured_response}"
+                        or f"Returning structured response: {structured_response}"
                     )
 
                     return {
@@ -952,13 +951,13 @@ def create_agent(  # noqa: PLR0915
                         ],
                         "structured_response": structured_response,
                     }
-                except Exception as exc:  # noqa: BLE001
+                except Exception as exc:
                     exception = StructuredOutputValidationError(tool_call["name"], exc, output)
                     should_retry, error_message = _handle_structured_output_error(
                         exception, effective_response_format
                     )
                     if not should_retry:
-                        raise exception
+                        raise exception from exc
 
                     return {
                         "messages": [
@@ -1100,6 +1099,8 @@ def create_agent(  # noqa: PLR0915
             messages = [request.system_message, *messages]
 
         output = model_.invoke(messages)
+        if name:
+            output.name = name
 
         # Handle model output to get messages and structured_response
         handled_output = _handle_model_output(output, effective_response_format)
@@ -1153,6 +1154,8 @@ def create_agent(  # noqa: PLR0915
             messages = [request.system_message, *messages]
 
         output = await model_.ainvoke(messages)
+        if name:
+            output.name = name
 
         # Handle model output to get messages and structured_response
         handled_output = _handle_model_output(output, effective_response_format)

{langchain-1.1.2 → langchain-1.2.2}/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
@@ -127,7 +128,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 +180,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 +195,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 +203,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 +233,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 +250,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"] :]

{langchain-1.1.2 → langchain-1.2.2}/langchain/agents/middleware/_retry.py

@@ -116,7 +116,7 @@ def calculate_delay(
 
     if jitter and delay > 0:
         jitter_amount = delay * 0.25  # ±25% jitter
-        delay = delay + random.uniform(-jitter_amount, jitter_amount)  # noqa: S311
+        delay += random.uniform(-jitter_amount, jitter_amount)  # noqa: S311
         # Ensure delay is not negative after jitter
         delay = max(0, delay)
 

{langchain-1.1.2 → langchain-1.2.2}/langchain/agents/middleware/context_editing.py

@@ -228,6 +228,7 @@ class ContextEditingMiddleware(AgentMiddleware):
 
             def count_tokens(messages: Sequence[BaseMessage]) -> int:
                 return count_tokens_approximately(messages)
+
         else:
             system_msg = [request.system_message] if request.system_message else []
 
@@ -255,6 +256,7 @@ class ContextEditingMiddleware(AgentMiddleware):
 
             def count_tokens(messages: Sequence[BaseMessage]) -> int:
                 return count_tokens_approximately(messages)
+
         else:
             system_msg = [request.system_message] if request.system_message else []
 

{langchain-1.1.2 → langchain-1.2.2}/langchain/agents/middleware/model_call_limit.py

@@ -6,7 +6,7 @@ from typing import TYPE_CHECKING, Annotated, Any, Literal
 
 from langchain_core.messages import AIMessage
 from langgraph.channels.untracked_value import UntrackedValue
-from typing_extensions import NotRequired
+from typing_extensions import NotRequired, override
 
 from langchain.agents.middleware.types import (
     AgentMiddleware,
@@ -148,7 +148,7 @@ class ModelCallLimitMiddleware(AgentMiddleware[ModelCallLimitState, Any]):
             msg = "At least one limit must be specified (thread_limit or run_limit)"
             raise ValueError(msg)
 
-        if exit_behavior not in ("end", "error"):
+        if exit_behavior not in {"end", "error"}:
             msg = f"Invalid exit_behavior: {exit_behavior}. Must be 'end' or 'error'"
             raise ValueError(msg)
 
@@ -157,7 +157,8 @@ class ModelCallLimitMiddleware(AgentMiddleware[ModelCallLimitState, Any]):
         self.exit_behavior = exit_behavior
 
     @hook_config(can_jump_to=["end"])
-    def before_model(self, state: ModelCallLimitState, runtime: Runtime) -> dict[str, Any] | None:  # noqa: ARG002
+    @override
+    def before_model(self, state: ModelCallLimitState, runtime: Runtime) -> dict[str, Any] | None:
         """Check model call limits before making a model call.
 
         Args:
@@ -222,7 +223,8 @@ class ModelCallLimitMiddleware(AgentMiddleware[ModelCallLimitState, Any]):
         """
         return self.before_model(state, runtime)
 
-    def after_model(self, state: ModelCallLimitState, runtime: Runtime) -> dict[str, Any] | None:  # noqa: ARG002
+    @override
+    def after_model(self, state: ModelCallLimitState, runtime: Runtime) -> dict[str, Any] | None:
         """Increment model call counts after a model call.
 
         Args:

{langchain-1.1.2 → langchain-1.2.2}/langchain/agents/middleware/model_fallback.py

@@ -87,14 +87,14 @@ class ModelFallbackMiddleware(AgentMiddleware):
         last_exception: Exception
         try:
             return handler(request)
-        except Exception as e:  # noqa: BLE001
+        except Exception as e:
             last_exception = e
 
         # Try fallback models
         for fallback_model in self.models:
             try:
                 return handler(request.override(model=fallback_model))
-            except Exception as e:  # noqa: BLE001
+            except Exception as e:
                 last_exception = e
                 continue
 
@@ -121,14 +121,14 @@ class ModelFallbackMiddleware(AgentMiddleware):
         last_exception: Exception
         try:
             return await handler(request)
-        except Exception as e:  # noqa: BLE001
+        except Exception as e:
             last_exception = e
 
         # Try fallback models
         for fallback_model in self.models:
             try:
                 return await handler(request.override(model=fallback_model))
-            except Exception as e:  # noqa: BLE001
+            except Exception as e:
                 last_exception = e
                 continue
 

{langchain-1.1.2 → langchain-1.2.2}/langchain/agents/middleware/model_retry.py

@@ -223,7 +223,7 @@ class ModelRetryMiddleware(AgentMiddleware):
         for attempt in range(self.max_retries + 1):
             try:
                 return handler(request)
-            except Exception as exc:  # noqa: BLE001
+            except Exception as exc:
                 attempts_made = attempt + 1  # attempt is 0-indexed
 
                 # Check if we should retry this exception
@@ -270,7 +270,7 @@ class ModelRetryMiddleware(AgentMiddleware):
         for attempt in range(self.max_retries + 1):
             try:
                 return await handler(request)
-            except Exception as exc:  # noqa: BLE001
+            except Exception as exc:
                 attempts_made = attempt + 1  # attempt is 0-indexed
 
                 # Check if we should retry this exception

{langchain-1.1.2 → langchain-1.2.2}/langchain/agents/middleware/pii.py

@@ -5,6 +5,7 @@ from __future__ import annotations
 from typing import TYPE_CHECKING, Any, Literal
 
 from langchain_core.messages import AIMessage, AnyMessage, HumanMessage, ToolMessage
+from typing_extensions import override
 
 from langchain.agents.middleware._redaction import (
     PIIDetectionError,
@@ -92,6 +93,8 @@ class PIIMiddleware(AgentMiddleware):
 
     def __init__(
         self,
+        # From a typing point of view, the literals are covered by 'str'.
+        # Nonetheless, we escape PYI051 to keep hints and autocompletion for the caller.
         pii_type: Literal["email", "credit_card", "ip", "mac_address", "url"] | str,  # noqa: PYI051
         *,
         strategy: Literal["block", "redact", "mask", "hash"] = "redact",
@@ -158,10 +161,11 @@ class PIIMiddleware(AgentMiddleware):
         return sanitized, matches
 
     @hook_config(can_jump_to=["end"])
+    @override
     def before_model(
         self,
         state: AgentState,
-        runtime: Runtime,  # noqa: ARG002
+        runtime: Runtime,
     ) -> dict[str, Any] | None:
         """Check user messages and tool results for PII before model invocation.
 
@@ -273,10 +277,11 @@ class PIIMiddleware(AgentMiddleware):
         """
         return self.before_model(state, runtime)
 
+    @override
     def after_model(
         self,
         state: AgentState,
-        runtime: Runtime,  # noqa: ARG002
+        runtime: Runtime,
     ) -> dict[str, Any] | None:
         """Check AI messages for PII after model invocation.
 
@@ -355,6 +360,7 @@ class PIIMiddleware(AgentMiddleware):
 
 __all__ = [
     "PIIDetectionError",
+    "PIIMatch",
     "PIIMiddleware",
     "detect_credit_card",
     "detect_email",

{langchain-1.1.2 → langchain-1.2.2}/langchain/agents/middleware/shell_tool.py

@@ -22,7 +22,7 @@ from langchain_core.tools.base import ToolException
 from langgraph.channels.untracked_value import UntrackedValue
 from pydantic import BaseModel, model_validator
 from pydantic.json_schema import SkipJsonSchema
-from typing_extensions import NotRequired
+from typing_extensions import NotRequired, override
 
 from langchain.agents.middleware._execution import (
     SHELL_TEMP_PREFIX,
@@ -78,10 +78,10 @@ class _SessionResources:
     session: ShellSession
     tempdir: tempfile.TemporaryDirectory[str] | None
     policy: BaseExecutionPolicy
-    _finalizer: weakref.finalize = field(init=False, repr=False)
+    finalizer: weakref.finalize = field(init=False, repr=False)
 
     def __post_init__(self) -> None:
-        self._finalizer = weakref.finalize(
+        self.finalizer = weakref.finalize(
             self,
             _cleanup_resources,
             self.session,
@@ -211,9 +211,14 @@ class ShellSession:
         with self._lock:
             self._drain_queue()
             payload = command if command.endswith("\n") else f"{command}\n"
-            self._stdin.write(payload)
-            self._stdin.write(f"printf '{marker} %s\\n' $?\n")
-            self._stdin.flush()
+            try:
+                self._stdin.write(payload)
+                self._stdin.write(f"printf '{marker} %s\\n' $?\n")
+                self._stdin.flush()
+            except (BrokenPipeError, OSError):
+                # The shell exited before we could write the marker command.
+                # This happens when commands like 'exit 1' terminate the shell.
+                return self._collect_output_after_exit(deadline)
 
             return self._collect_output(marker, deadline, timeout)
 
@@ -248,6 +253,10 @@ class ShellSession:
             if source == "stdout" and data.startswith(marker):
                 _, _, status = data.partition(" ")
                 exit_code = self._safe_int(status.strip())
+                # Drain any remaining stderr that may have arrived concurrently.
+                # The stderr reader thread runs independently, so output might
+                # still be in flight when the stdout marker arrives.
+                self._drain_remaining_stderr(collected, deadline)
                 break
 
             total_lines += 1
@@ -300,6 +309,80 @@ class ShellSession:
             total_bytes=total_bytes,
         )
 
+    def _collect_output_after_exit(self, deadline: float) -> CommandExecutionResult:
+        """Collect output after the shell exited unexpectedly.
+
+        Called when a `BrokenPipeError` occurs while writing to stdin, indicating the
+        shell process terminated (e.g., due to an 'exit' command).
+
+        Args:
+            deadline: Absolute time by which collection must complete.
+
+        Returns:
+            `CommandExecutionResult` with collected output and the process exit code.
+        """
+        collected: list[str] = []
+        total_lines = 0
+        total_bytes = 0
+        truncated_by_lines = False
+        truncated_by_bytes = False
+
+        # Give reader threads a brief moment to enqueue any remaining output.
+        drain_timeout = 0.1
+        drain_deadline = min(time.monotonic() + drain_timeout, deadline)
+
+        while True:
+            remaining = drain_deadline - time.monotonic()
+            if remaining <= 0:
+                break
+            try:
+                source, data = self._queue.get(timeout=remaining)
+            except queue.Empty:
+                break
+
+            if data is None:
+                # EOF marker from a reader thread; continue draining.
+                continue
+
+            total_lines += 1
+            encoded = data.encode("utf-8", "replace")
+            total_bytes += len(encoded)
+
+            if total_lines > self._policy.max_output_lines:
+                truncated_by_lines = True
+                continue
+
+            if (
+                self._policy.max_output_bytes is not None
+                and total_bytes > self._policy.max_output_bytes
+            ):
+                truncated_by_bytes = True
+                continue
+
+            if source == "stderr":
+                stripped = data.rstrip("\n")
+                collected.append(f"[stderr] {stripped}")
+                if data.endswith("\n"):
+                    collected.append("\n")
+            else:
+                collected.append(data)
+
+        # Get exit code from the terminated process.
+        exit_code: int | None = None
+        if self._process:
+            exit_code = self._process.poll()
+
+        output = "".join(collected)
+        return CommandExecutionResult(
+            output=output,
+            exit_code=exit_code,
+            timed_out=False,
+            truncated_by_lines=truncated_by_lines,
+            truncated_by_bytes=truncated_by_bytes,
+            total_lines=total_lines,
+            total_bytes=total_bytes,
+        )
+
     def _kill_process(self) -> None:
         if not self._process:
             return
@@ -323,6 +406,37 @@ class ShellSession:
             except queue.Empty:
                 break
 
+    def _drain_remaining_stderr(
+        self, collected: list[str], deadline: float, drain_timeout: float = 0.05
+    ) -> None:
+        """Drain any stderr output that arrived concurrently with the done marker.
+
+        The stdout and stderr reader threads run independently. When a command writes to
+        stderr just before exiting, the stderr output may still be in transit when the
+        done marker arrives on stdout. This method briefly polls the queue to capture
+        such output.
+
+        Args:
+            collected: The list to append collected stderr lines to.
+            deadline: The original command deadline (used as an upper bound).
+            drain_timeout: Maximum time to wait for additional stderr output.
+        """
+        drain_deadline = min(time.monotonic() + drain_timeout, deadline)
+        while True:
+            remaining = drain_deadline - time.monotonic()
+            if remaining <= 0:
+                break
+            try:
+                source, data = self._queue.get(timeout=remaining)
+            except queue.Empty:
+                break
+            if data is None or source != "stderr":
+                continue
+            stripped = data.rstrip("\n")
+            collected.append(f"[stderr] {stripped}")
+            if data.endswith("\n"):
+                collected.append("\n")
+
     @staticmethod
     def _safe_int(value: str) -> int | None:
         with contextlib.suppress(ValueError):
@@ -405,6 +519,12 @@ class ShellToolMiddleware(AgentMiddleware[ShellToolState, Any]):
                 Defaults to `HostExecutionPolicy` for native execution.
             redaction_rules: Optional redaction rules to sanitize command output before
                 returning it to the model.
+
+                !!! warning
+                    Redaction rules are applied post execution and do not prevent
+                    exfiltration of secrets or sensitive data when using
+                    `HostExecutionPolicy`.
+
             tool_description: Optional override for the registered shell tool
                 description.
             tool_name: Name for the registered shell tool.
@@ -489,7 +609,8 @@ class ShellToolMiddleware(AgentMiddleware[ShellToolState, Any]):
             normalized[key] = str(value)
         return normalized
 
-    def before_agent(self, state: ShellToolState, runtime: Runtime) -> dict[str, Any] | None:  # noqa: ARG002
+    @override
+    def before_agent(self, state: ShellToolState, runtime: Runtime) -> dict[str, Any] | None:
         """Start the shell session and run startup commands."""
         resources = self._get_or_create_resources(state)
         return {"shell_session_resources": resources}
@@ -498,7 +619,8 @@ class ShellToolMiddleware(AgentMiddleware[ShellToolState, Any]):
         """Async start the shell session and run startup commands."""
         return self.before_agent(state, runtime)
 
-    def after_agent(self, state: ShellToolState, runtime: Runtime) -> None:  # noqa: ARG002
+    @override
+    def after_agent(self, state: ShellToolState, runtime: Runtime) -> None:
         """Run shutdown commands and release resources when an agent completes."""
         resources = state.get("shell_session_resources")
         if not isinstance(resources, _SessionResources):
@@ -507,7 +629,7 @@ class ShellToolMiddleware(AgentMiddleware[ShellToolState, Any]):
         try:
             self._run_shutdown_commands(resources.session)
         finally:
-            resources._finalizer()
+            resources.finalizer()
 
     async def aafter_agent(self, state: ShellToolState, runtime: Runtime) -> None:
         """Async run shutdown commands and release resources when an agent completes."""
@@ -568,7 +690,7 @@ class ShellToolMiddleware(AgentMiddleware[ShellToolState, Any]):
             return
         for command in self._startup_commands:
             result = session.execute(command, timeout=self._execution_policy.startup_timeout)
-            if result.timed_out or (result.exit_code not in (0, None)):
+            if result.timed_out or (result.exit_code not in {0, None}):
                 msg = f"Startup command '{command}' failed with exit code {result.exit_code}"
                 raise RuntimeError(msg)
 
@@ -580,7 +702,7 @@ class ShellToolMiddleware(AgentMiddleware[ShellToolState, Any]):
                 result = session.execute(command, timeout=self._execution_policy.command_timeout)
                 if result.timed_out:
                     LOGGER.warning("Shutdown command '%s' timed out.", command)
-                elif result.exit_code not in (0, None):
+                elif result.exit_code not in {0, None}:
                     LOGGER.warning(
                         "Shutdown command '%s' exited with %s.", command, result.exit_code
                     )
@@ -671,7 +793,7 @@ class ShellToolMiddleware(AgentMiddleware[ShellToolState, Any]):
                 f"(observed {result.total_bytes})."
             )
 
-        if result.exit_code not in (0, None):
+        if result.exit_code not in {0, None}:
             sanitized_output = f"{sanitized_output.rstrip()}\n\nExit code: {result.exit_code}"
             final_status: Literal["success", "error"] = "error"
         else: