test_shell_tool.py — langchain Source File

Source Code

from __future__ import annotations

import gc
import tempfile
import time
from pathlib import Path
from typing import cast

import pytest
from langchain_core.messages import ToolMessage
from langchain_core.tools.base import ToolException
from langgraph.runtime import Runtime

from langchain.agents.middleware.shell_tool import (
    HostExecutionPolicy,
    RedactionRule,
    ShellToolMiddleware,
    ShellToolState,
    _SessionResources,
    _ShellToolInput,
)


def _empty_state() -> ShellToolState:
    return {"messages": []}


def test_executes_command_and_persists_state(tmp_path: Path) -> None:
    workspace = tmp_path / "workspace"
    middleware = ShellToolMiddleware(workspace_root=workspace)
    runtime = Runtime()
    state = _empty_state()
    try:
        updates = middleware.before_agent(state, runtime)
        if updates:
            state.update(cast("ShellToolState", updates))
        resources = middleware._get_or_create_resources(state)

        middleware._run_shell_tool(resources, {"command": "cd /"}, tool_call_id=None)
        result = middleware._run_shell_tool(resources, {"command": "pwd"}, tool_call_id=None)
        assert isinstance(result, str)
        assert result.strip() == "/"
        echo_result = middleware._run_shell_tool(
            resources, {"command": "echo ready"}, tool_call_id=None
        )
        assert "ready" in echo_result
    finally:
        middleware.after_agent(state, runtime)


def test_restart_resets_session_environment(tmp_path: Path) -> None:
    middleware = ShellToolMiddleware(workspace_root=tmp_path / "workspace")
    runtime = Runtime()
    state = _empty_state()
    try:
        updates = middleware.before_agent(state, runtime)
        if updates:
            state.update(cast("ShellToolState", updates))
        resources = middleware._get_or_create_resources(state)

// ... (489 more lines)