Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions .github/workflows/devflow-pr-review.yml
Original file line number Diff line number Diff line change
Expand Up @@ -246,10 +246,10 @@ jobs:
python-version: "3.13"

- name: Set up uv
uses: astral-sh/setup-uv@c771a70e6277c0a99b617c7a806ffedaca235ff9 # v9.0.0
uses: astral-sh/setup-uv@20cfd1bf945f4377ade1205e4dbc17946fc9a30d # v10.0.1
with:
version: "0.11.x"
enable-cache: true
enable-cache: auto

- name: Install DevFlow dependencies
working-directory: ${{ env.DEVFLOW_PATH }}
Expand Down
2 changes: 1 addition & 1 deletion .github/workflows/issue-triage.yml
Original file line number Diff line number Diff line change
Expand Up @@ -190,7 +190,7 @@ jobs:
python-version: "3.13"

- name: Set up uv
uses: astral-sh/setup-uv@c771a70e6277c0a99b617c7a806ffedaca235ff9 # v9.0.0
uses: astral-sh/setup-uv@20cfd1bf945f4377ade1205e4dbc17946fc9a30d # v10.0.1
with:
version: "0.11.x"
enable-cache: true
Expand Down
4 changes: 2 additions & 2 deletions .github/workflows/python-docs.yml
Original file line number Diff line number Diff line change
Expand Up @@ -26,10 +26,10 @@ jobs:
steps:
- uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
- name: Set up uv
uses: astral-sh/setup-uv@c771a70e6277c0a99b617c7a806ffedaca235ff9 # v9.0.0
uses: astral-sh/setup-uv@20cfd1bf945f4377ade1205e4dbc17946fc9a30d # v10.0.1
with:
version-file: "python/pyproject.toml"
enable-cache: true
enable-cache: auto
cache-suffix: ${{ runner.os }}-${{ env.UV_PYTHON }}
cache-dependency-glob: "python/uv.lock"
- name: Install dependencies
Expand Down
167 changes: 167 additions & 0 deletions python/packages/tools/tests/test_local_shell_tool.py
Original file line number Diff line number Diff line change
@@ -1,16 +1,30 @@
# Copyright (c) Microsoft. All rights reserved.

import asyncio
import json
import os
import sys
from collections.abc import Awaitable, Mapping, Sequence
from typing import Any
from unittest.mock import AsyncMock, patch

import pytest
from agent_framework import Agent, AgentSession, BaseChatClient, ChatResponse, Content, FunctionInvocationLayer, Message
from agent_framework.security import (
ContentLabel,
IntegrityLabel,
LabelTrackingFunctionMiddleware,
PolicyEnforcementFunctionMiddleware,
)

from agent_framework_tools._feature_usage import FeatureIndex
from agent_framework_tools.shell import LocalShellTool, ShellCommandError, ShellPolicy
from agent_framework_tools.shell._executor import _popen_kwargs_for_group, run_stateless

_TEST_SHELL = "agent-framework-test-shell"
_APPROVED_COMMAND = "printf '%s' approved-value"
_ALTERNATE_COMMAND = "printf '%s' alternate-value"


class _FakeExecProcess:
def __init__(
Expand All @@ -32,6 +46,69 @@ async def communicate(self) -> tuple[bytes, bytes]:
return stdout, stderr


class _ScriptedFunctionClient(FunctionInvocationLayer, BaseChatClient):
def __init__(self, responses: Sequence[ChatResponse]) -> None:
super().__init__()
self._responses = list(responses)

def _inner_get_response(
self,
*,
messages: Sequence[Message],
stream: bool,
options: Mapping[str, Any],
**kwargs: Any,
) -> Awaitable[ChatResponse]:
assert not stream

async def get_response() -> ChatResponse:
if self._responses:
return self._responses.pop(0)
return ChatResponse(messages=Message(role="assistant", contents=["done"]))

return get_response()


async def _request_variable_shell_approval(
session_id: str,
) -> tuple[Agent[Any], AgentSession, Content, str, LabelTrackingFunctionMiddleware]:
tracker = LabelTrackingFunctionMiddleware()
session = AgentSession(session_id=session_id)
variable_id = tracker.get_variable_store(session).store(
_APPROVED_COMMAND,
ContentLabel(integrity=IntegrityLabel.UNTRUSTED),
)
function_call = Content.from_function_call(
call_id="provider-shell-call",
name="run_shell",
arguments={"command": f"[{variable_id}]"},
id="shell-call-occurrence",
)
client = _ScriptedFunctionClient([
ChatResponse(messages=Message(role="assistant", contents=[function_call])),
])
# Isolate the variable-aware policy approval from the tool's independent blanket approval.
shell = LocalShellTool(
mode="stateless",
shell=[_TEST_SHELL],
approval_mode="never_require",
acknowledge_unsafe=True,
)
agent = Agent(
client=client,
tools=[shell.as_function()],
middleware=[
tracker,
PolicyEnforcementFunctionMiddleware(approval_on_violation=True),
],
)

response = await agent.run("Run the hidden command", session=session)

assert len(response.user_input_requests) == 1
return agent, session, response.user_input_requests[0], variable_id, tracker


async def test_stateless_echo() -> None:
tool = LocalShellTool(mode="stateless", approval_mode="never_require", acknowledge_unsafe=True)
cmd = "Write-Output hello" if sys.platform == "win32" else "echo hello"
Expand Down Expand Up @@ -293,6 +370,96 @@ async def test_as_function_wires_kind_and_approval() -> None:
assert ft.approval_mode == "always_require"


async def test_variable_shell_approval_executes_only_the_reviewed_command() -> None:
agent, session, request, variable_id, _ = await _request_variable_shell_approval("exact-shell-approval")

assert request.id == "shell-call-occurrence"
assert request.function_call is not None
assert request.function_call.call_id == "provider-shell-call"
assert request.function_call.name == "run_shell"
assert request.function_call.parse_arguments() == {"command": f"[{variable_id}]"}

create_process = AsyncMock(return_value=_FakeExecProcess(communicate_results=[(b"approved-value", b"")]))
with patch("agent_framework_tools.shell._executor.asyncio.create_subprocess_exec", create_process):
await agent.run(request.to_function_approval_response(True), session=session)

assert create_process.await_count == 1
assert create_process.await_args is not None
assert create_process.await_args.args == (_TEST_SHELL, "-c", _APPROVED_COMMAND)


async def test_variable_shell_approval_blocks_changed_resolved_command() -> None:
agent, session, request, variable_id, _ = await _request_variable_shell_approval("changed-shell-variable")
security_state = session.state["__agent_framework_fides_security__"]
security_state["variables"][variable_id]["content"] = json.dumps(_ALTERNATE_COMMAND)

create_process = AsyncMock()
with patch("agent_framework_tools.shell._executor.asyncio.create_subprocess_exec", create_process):
response = await agent.run(request.to_function_approval_response(True), session=session)

create_process.assert_not_awaited()
replacement = response.user_input_requests
assert len(replacement) == 1
assert replacement[0].id != request.id
assert replacement[0].function_call is not None
assert replacement[0].function_call.parse_arguments() == {"command": f"[{variable_id}]"}


@pytest.mark.parametrize("mutation", ["reference", "command", "arguments", "request_id", "call_id"])
async def test_variable_shell_approval_substitutions_cannot_change_executed_argv(mutation: str) -> None:
agent, session, request, _, tracker = await _request_variable_shell_approval(f"substituted-shell-{mutation}")
response = Content.from_dict(request.to_function_approval_response(True).to_dict())
assert response.function_call is not None

if mutation == "reference":
alternate_variable_id = tracker.get_variable_store(session).store(
_ALTERNATE_COMMAND,
ContentLabel(integrity=IntegrityLabel.UNTRUSTED),
)
response.function_call.arguments = {"command": f"[{alternate_variable_id}]"}
elif mutation == "command":
response.function_call.arguments = {"command": _ALTERNATE_COMMAND}
elif mutation == "arguments":
response.function_call.arguments = {
"command": _ALTERNATE_COMMAND,
"unexpected": "substituted",
}
elif mutation == "request_id":
response.id = "different-request"
else:
response.function_call.call_id = "different-provider-call"

create_process = AsyncMock(
side_effect=lambda *_args, **_kwargs: _FakeExecProcess(communicate_results=[(b"approved-value", b"")])
)
with patch("agent_framework_tools.shell._executor.asyncio.create_subprocess_exec", create_process):
await agent.run(response, session=session)

if mutation == "request_id":
create_process.assert_not_awaited()
else:
assert create_process.await_count == 1
assert create_process.await_args is not None
assert create_process.await_args.args == (_TEST_SHELL, "-c", _APPROVED_COMMAND)


async def test_variable_shell_approval_cannot_be_replayed_on_later_turn_or_session() -> None:
agent, session, request, _, _ = await _request_variable_shell_approval("shell-approval-replay")
approval = request.to_function_approval_response(True)
create_process = AsyncMock(
side_effect=lambda *_args, **_kwargs: _FakeExecProcess(communicate_results=[(b"approved-value", b"")])
)

with patch("agent_framework_tools.shell._executor.asyncio.create_subprocess_exec", create_process):
await agent.run(approval, session=session)
await agent.run(approval, session=session)
await agent.run(approval, session=AgentSession(session_id="different-shell-session"))

assert create_process.await_count == 1
assert create_process.await_args is not None
assert create_process.await_args.args == (_TEST_SHELL, "-c", _APPROVED_COMMAND)


@pytest.mark.skipif(sys.platform == "win32", reason="POSIX persistent reanchor test")
async def test_persistent_confines_workdir_by_default(tmp_path: os.PathLike[str]) -> None:
"""With the default ``confine_workdir=True``, a ``cd`` in one call
Expand Down
Loading