mirror of
https://github.com/Sea-Haven-Industries/open-swe.git
synced 2026-09-30 23:13:15 +00:00
19eb89b010
Co-authored-by: langsmith-fleet[bot] <langsmith-fleet[bot]@users.noreply.github.com> (cherry picked from commit 75fb8b487852003916c4984504a13ee7226b2ceb)
93 lines
3.7 KiB
Python
93 lines
3.7 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
from typing import Any
|
|
|
|
from langchain_core.messages import ToolMessage
|
|
|
|
from agent.middleware.pr_creation_guard import (
|
|
PullRequestCreationGuardMiddleware,
|
|
is_pr_creation_fallback_command,
|
|
)
|
|
|
|
|
|
class _Request:
|
|
def __init__(self, command: str) -> None:
|
|
self.tool_call = {
|
|
"name": "execute",
|
|
"args": {"command": command},
|
|
"id": "call-1",
|
|
}
|
|
|
|
|
|
async def _handler(_request: Any) -> ToolMessage:
|
|
return ToolMessage(content="allowed", tool_call_id="call-1")
|
|
|
|
|
|
def test_detects_pr_creation_fallback_commands() -> None:
|
|
assert is_pr_creation_fallback_command("GH_TOKEN=dummy gh pr create --draft")
|
|
assert is_pr_creation_fallback_command(
|
|
"gh api repos/langchain-ai/open-swe/pulls -X POST -f title=x"
|
|
)
|
|
assert is_pr_creation_fallback_command(
|
|
"gh api -X POST repos/langchain-ai/open-swe/pulls -f title=x"
|
|
)
|
|
assert is_pr_creation_fallback_command(
|
|
"GH_TOKEN=dummy gh api -X POST repos/langchain-ai/open-swe/pulls -f title=x"
|
|
)
|
|
assert is_pr_creation_fallback_command(
|
|
"curl -X POST https://api.github.com/repos/langchain-ai/open-swe/pulls -d '{}'"
|
|
)
|
|
assert is_pr_creation_fallback_command("/usr/bin/gh pr create --draft")
|
|
assert is_pr_creation_fallback_command(
|
|
"/usr/bin/curl -X POST https://api.github.com/repos/langchain-ai/open-swe/pulls -d '{}'"
|
|
)
|
|
assert is_pr_creation_fallback_command("bash -c 'gh pr create --draft'")
|
|
assert is_pr_creation_fallback_command("GH_TOKEN=dummy sh -c 'gh pr create --draft'")
|
|
assert is_pr_creation_fallback_command(
|
|
"zsh -lc 'gh api repos/langchain-ai/open-swe/pulls -X POST -f title=x'"
|
|
)
|
|
assert is_pr_creation_fallback_command(
|
|
"bash -c \"curl -X POST https://api.github.com/repos/langchain-ai/open-swe/pulls -d '{}'\""
|
|
)
|
|
assert is_pr_creation_fallback_command(
|
|
'bash -c \'sh -c "dash -c \\"zsh -c \\\\\\"gh pr create --draft\\\\\\"\\""\''
|
|
)
|
|
|
|
|
|
def test_allows_safe_pr_commands() -> None:
|
|
assert not is_pr_creation_fallback_command("GH_TOKEN=dummy gh pr view 1 --json url")
|
|
assert not is_pr_creation_fallback_command("gh pr list --head open-swe/foo")
|
|
assert not is_pr_creation_fallback_command("gh pr edit 1 --add-label ready")
|
|
assert not is_pr_creation_fallback_command("gh pr comment 1 --body done")
|
|
assert not is_pr_creation_fallback_command("bash -c 'gh pr view 1 --json url'")
|
|
assert not is_pr_creation_fallback_command("/usr/bin/gh pr view 1 --json url")
|
|
|
|
|
|
async def test_middleware_blocks_execute_pr_creation_fallbacks() -> None:
|
|
for command in (
|
|
"GH_TOKEN=dummy gh pr create --draft",
|
|
"gh api repos/langchain-ai/open-swe/pulls -X POST -f title=x",
|
|
"GH_TOKEN=dummy gh api -X POST repos/langchain-ai/open-swe/pulls -f title=x",
|
|
"curl -X POST https://api.github.com/repos/langchain-ai/open-swe/pulls -d '{}'",
|
|
):
|
|
result = await PullRequestCreationGuardMiddleware().awrap_tool_call(
|
|
_Request(command), _handler
|
|
)
|
|
|
|
assert isinstance(result, ToolMessage)
|
|
assert result.status == "error"
|
|
payload = json.loads(str(result.content))
|
|
assert payload["code"] == "pr_creation_fallback_blocked"
|
|
assert payload["recoverable_by_agent"] is False
|
|
assert "open_pull_request" in payload["error"]
|
|
assert payload["blocked_command"] == command
|
|
|
|
|
|
async def test_middleware_allows_safe_pr_view() -> None:
|
|
result = await PullRequestCreationGuardMiddleware().awrap_tool_call(
|
|
_Request("GH_TOKEN=dummy gh pr view 1 --json url"), _handler
|
|
)
|
|
|
|
assert isinstance(result, ToolMessage)
|
|
assert result.content == "allowed"
|