wehub-resource-sync
530 lines
17 KiB
Python
530 lines
17 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
from collections.abc import AsyncIterator
|
|
from typing import Any
|
|
|
|
import pytest
|
|
|
|
from opensquilla.engine import Agent, AgentConfig, ToolResult
|
|
from opensquilla.engine.types import ToolCall, ToolResultEvent
|
|
from opensquilla.gateway.approval_queue import get_approval_queue, reset_approval_queue
|
|
from opensquilla.provider import ChatConfig, Message, ToolDefinition, ToolInputSchema
|
|
from opensquilla.provider import DoneEvent as ProviderDone
|
|
from opensquilla.provider import TextDeltaEvent as ProviderTextDelta
|
|
from opensquilla.provider import ToolUseEndEvent as ProviderToolUseEnd
|
|
from opensquilla.provider import ToolUseStartEvent as ProviderToolUseStart
|
|
|
|
|
|
class _OneApprovalToolProvider:
|
|
provider_name = "fake"
|
|
|
|
def __init__(self) -> None:
|
|
self.calls: list[list[Message]] = []
|
|
|
|
def chat(
|
|
self,
|
|
messages: list[Message],
|
|
tools: list[Any] | None = None,
|
|
config: ChatConfig | None = None,
|
|
) -> AsyncIterator[Any]:
|
|
self.calls.append(messages)
|
|
call_number = len(self.calls)
|
|
return self._stream(call_number)
|
|
|
|
async def _stream(self, call_number: int) -> AsyncIterator[Any]:
|
|
if call_number > 1:
|
|
yield ProviderTextDelta(text="done")
|
|
yield ProviderDone(stop_reason="stop", input_tokens=1, output_tokens=1)
|
|
return
|
|
yield ProviderToolUseStart(tool_use_id="tool-1", tool_name="exec_command")
|
|
yield ProviderToolUseEnd(
|
|
tool_use_id="tool-1",
|
|
tool_name="exec_command",
|
|
arguments={"command": "pip install demo"},
|
|
)
|
|
yield ProviderDone(stop_reason="tool_use", input_tokens=1, output_tokens=1)
|
|
|
|
async def list_models(self) -> list[Any]:
|
|
return []
|
|
|
|
|
|
class _DeniedApprovalThenAnswerProvider:
|
|
provider_name = "fake"
|
|
|
|
def __init__(self) -> None:
|
|
self.calls: list[list[Message]] = []
|
|
|
|
def chat(
|
|
self,
|
|
messages: list[Message],
|
|
tools: list[Any] | None = None,
|
|
config: ChatConfig | None = None,
|
|
) -> AsyncIterator[Any]:
|
|
self.calls.append(messages)
|
|
call_number = len(self.calls)
|
|
return self._stream(call_number)
|
|
|
|
async def _stream(self, call_number: int) -> AsyncIterator[Any]:
|
|
if call_number > 1:
|
|
yield ProviderTextDelta(text="用户拒绝访问，所以我无法查看该路径。")
|
|
yield ProviderDone(stop_reason="stop", input_tokens=1, output_tokens=1)
|
|
return
|
|
yield ProviderToolUseStart(tool_use_id="tool-deny", tool_name="exec_command")
|
|
yield ProviderToolUseEnd(
|
|
tool_use_id="tool-deny",
|
|
tool_name="exec_command",
|
|
arguments={"command": "ls /outside"},
|
|
)
|
|
yield ProviderDone(stop_reason="tool_use", input_tokens=1, output_tokens=1)
|
|
|
|
async def list_models(self) -> list[Any]:
|
|
return []
|
|
|
|
|
|
class _PendingApprovalShouldNotAnswerProvider:
|
|
provider_name = "fake"
|
|
|
|
def __init__(self) -> None:
|
|
self.calls: list[list[Message]] = []
|
|
|
|
def chat(
|
|
self,
|
|
messages: list[Message],
|
|
tools: list[Any] | None = None,
|
|
config: ChatConfig | None = None,
|
|
) -> AsyncIterator[Any]:
|
|
self.calls.append(messages)
|
|
call_number = len(self.calls)
|
|
return self._stream(call_number)
|
|
|
|
async def _stream(self, call_number: int) -> AsyncIterator[Any]:
|
|
if call_number > 1:
|
|
yield ProviderTextDelta(text="fallback from training knowledge")
|
|
yield ProviderDone(stop_reason="stop", input_tokens=1, output_tokens=1)
|
|
return
|
|
yield ProviderToolUseStart(tool_use_id="tool-pending", tool_name="exec_command")
|
|
yield ProviderToolUseEnd(
|
|
tool_use_id="tool-pending",
|
|
tool_name="exec_command",
|
|
arguments={"command": "web_fetch https://example.com"},
|
|
)
|
|
yield ProviderDone(stop_reason="tool_use", input_tokens=1, output_tokens=1)
|
|
|
|
async def list_models(self) -> list[Any]:
|
|
return []
|
|
|
|
|
|
class _DoneProvider:
|
|
provider_name = "fake"
|
|
|
|
def chat(
|
|
self,
|
|
messages: list[Message],
|
|
tools: list[Any] | None = None,
|
|
config: ChatConfig | None = None,
|
|
) -> AsyncIterator[Any]:
|
|
return self._stream()
|
|
|
|
async def _stream(self) -> AsyncIterator[Any]:
|
|
yield ProviderTextDelta(text="done")
|
|
yield ProviderDone(stop_reason="stop", input_tokens=1, output_tokens=1)
|
|
|
|
async def list_models(self) -> list[Any]:
|
|
return []
|
|
|
|
|
|
def _exec_definition() -> ToolDefinition:
|
|
return ToolDefinition(
|
|
name="exec_command",
|
|
description="Execute command.",
|
|
input_schema=ToolInputSchema(
|
|
properties={
|
|
"command": {"type": "string"},
|
|
"approval_id": {"type": "string"},
|
|
},
|
|
required=["command"],
|
|
),
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_agent_run_turn_clears_sandbox_approval_denials_for_session(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
cleared: list[str | None] = []
|
|
|
|
def _clear(session_key: str | None = None) -> None:
|
|
cleared.append(session_key)
|
|
|
|
monkeypatch.setattr(
|
|
"opensquilla.sandbox.escalation.clear_sandbox_approval_denials",
|
|
_clear,
|
|
)
|
|
agent = Agent(
|
|
provider=_DoneProvider(),
|
|
config=AgentConfig(max_iterations=1),
|
|
session_key="agent:main:webchat:abc",
|
|
)
|
|
|
|
_ = [event async for event in agent.run_turn("hello")]
|
|
|
|
assert cleared == ["agent:main:webchat:abc"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_interactive_approval_result_is_waited_and_retried_before_model_continues() -> None:
|
|
approval_prompt_seen = asyncio.Event()
|
|
allow_retry = asyncio.Event()
|
|
tool_calls: list[dict[str, Any]] = []
|
|
|
|
async def _handler(call: ToolCall) -> ToolResult:
|
|
tool_calls.append(dict(call.arguments))
|
|
approval_id = call.arguments.get("approval_id")
|
|
if approval_id is None:
|
|
return ToolResult(
|
|
tool_use_id=call.tool_use_id,
|
|
tool_name=call.tool_name,
|
|
content=json.dumps(
|
|
{
|
|
"status": "approval_required",
|
|
"approval_id": "approve-1",
|
|
"command": call.arguments["command"],
|
|
"warning": "command requires approval",
|
|
}
|
|
),
|
|
)
|
|
assert approval_id == "approve-1"
|
|
return ToolResult(
|
|
tool_use_id=call.tool_use_id,
|
|
tool_name=call.tool_name,
|
|
content="exit_code=0\ninstalled\n",
|
|
)
|
|
|
|
provider = _OneApprovalToolProvider()
|
|
agent = Agent(
|
|
provider=provider,
|
|
config=AgentConfig(max_iterations=2),
|
|
tool_definitions=[_exec_definition()],
|
|
tool_handler=_handler,
|
|
)
|
|
|
|
events: list[Any] = []
|
|
|
|
async def _drive() -> None:
|
|
async for event in agent.run_turn("install package"):
|
|
events.append(event)
|
|
if isinstance(event, ToolResultEvent) and "approval_required" in event.result:
|
|
approval_prompt_seen.set()
|
|
await allow_retry.wait()
|
|
|
|
task = asyncio.create_task(_drive())
|
|
await asyncio.wait_for(approval_prompt_seen.wait(), timeout=2.0)
|
|
assert len(provider.calls) == 1
|
|
assert tool_calls == [{"command": "pip install demo"}]
|
|
|
|
allow_retry.set()
|
|
await asyncio.wait_for(task, timeout=2.0)
|
|
|
|
assert tool_calls == [
|
|
{"command": "pip install demo"},
|
|
{"command": "pip install demo", "approval_id": "approve-1"},
|
|
]
|
|
assert len(provider.calls) == 2
|
|
assert any(
|
|
isinstance(event, ToolResultEvent) and event.result.startswith("exit_code=0")
|
|
for event in events
|
|
)
|
|
second_provider_request = provider.calls[1]
|
|
tool_result_messages = [
|
|
msg
|
|
for msg in second_provider_request
|
|
if any(getattr(block, "type", None) == "tool_result" for block in msg.content)
|
|
]
|
|
assert len(tool_result_messages) == 1
|
|
block = next(
|
|
block
|
|
for block in tool_result_messages[0].content
|
|
if getattr(block, "type", None) == "tool_result"
|
|
)
|
|
assert block.content == "exit_code=0\ninstalled\n"
|
|
assert "approval_required" not in block.content
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_agent_waits_for_approval_resolution_before_retry_result_reaches_model() -> None:
|
|
reset_approval_queue()
|
|
approval_prompt_seen = asyncio.Event()
|
|
tool_calls: list[dict[str, Any]] = []
|
|
|
|
async def _handler(call: ToolCall) -> ToolResult:
|
|
tool_calls.append(dict(call.arguments))
|
|
approval_id = call.arguments.get("approval_id")
|
|
if approval_id is None:
|
|
approval_id = get_approval_queue().request(
|
|
"exec",
|
|
{
|
|
"toolName": call.tool_name,
|
|
"command": call.arguments["command"],
|
|
"args": dict(call.arguments),
|
|
},
|
|
)
|
|
return ToolResult(
|
|
tool_use_id=call.tool_use_id,
|
|
tool_name=call.tool_name,
|
|
content=json.dumps(
|
|
{
|
|
"status": "approval_required",
|
|
"approval_id": approval_id,
|
|
"command": call.arguments["command"],
|
|
"warning": "command requires approval",
|
|
}
|
|
),
|
|
)
|
|
|
|
entry = get_approval_queue().get(str(approval_id))
|
|
if not entry.resolved:
|
|
return ToolResult(
|
|
tool_use_id=call.tool_use_id,
|
|
tool_name=call.tool_name,
|
|
content=json.dumps(
|
|
{
|
|
"status": "approval_pending",
|
|
"approval_id": approval_id,
|
|
"command": call.arguments["command"],
|
|
"warning": "command requires approval",
|
|
}
|
|
),
|
|
)
|
|
return ToolResult(
|
|
tool_use_id=call.tool_use_id,
|
|
tool_name=call.tool_name,
|
|
content="exit_code=0\napproved\n",
|
|
)
|
|
|
|
provider = _OneApprovalToolProvider()
|
|
agent = Agent(
|
|
provider=provider,
|
|
config=AgentConfig(
|
|
max_iterations=2,
|
|
metadata={"approval_wait_timeout_seconds": 1.0},
|
|
),
|
|
tool_definitions=[_exec_definition()],
|
|
tool_handler=_handler,
|
|
)
|
|
|
|
events: list[Any] = []
|
|
|
|
async def _drive() -> None:
|
|
async for event in agent.run_turn("install package"):
|
|
events.append(event)
|
|
if isinstance(event, ToolResultEvent) and "approval_required" in event.result:
|
|
approval_prompt_seen.set()
|
|
|
|
try:
|
|
task = asyncio.create_task(_drive())
|
|
await asyncio.wait_for(approval_prompt_seen.wait(), timeout=2.0)
|
|
await asyncio.sleep(0.05)
|
|
|
|
assert len(provider.calls) == 1
|
|
assert len(tool_calls) == 1
|
|
|
|
approval_event = next(
|
|
event
|
|
for event in events
|
|
if isinstance(event, ToolResultEvent) and "approval_required" in event.result
|
|
)
|
|
approval_id = json.loads(approval_event.result)["approval_id"]
|
|
get_approval_queue().resolve(approval_id, True)
|
|
await asyncio.wait_for(task, timeout=2.0)
|
|
|
|
assert len(tool_calls) == 2
|
|
second_provider_request = provider.calls[1]
|
|
tool_result_blocks = [
|
|
block
|
|
for msg in second_provider_request
|
|
for block in msg.content
|
|
if getattr(block, "type", None) == "tool_result"
|
|
]
|
|
assert [block.content for block in tool_result_blocks] == [
|
|
"exit_code=0\napproved\n"
|
|
]
|
|
assert all(
|
|
"approval_pending" not in event.result
|
|
for event in events
|
|
if isinstance(event, ToolResultEvent)
|
|
)
|
|
finally:
|
|
reset_approval_queue()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_unresolved_approval_pauses_turn_without_model_fallback(
|
|
tmp_path: Any,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
from opensquilla.application import approval_queue as approval_queue_mod
|
|
|
|
monkeypatch.setattr(
|
|
approval_queue_mod,
|
|
"_DEFAULT_APPROVAL_QUEUE_PATH",
|
|
tmp_path / "approval_queue.sqlite",
|
|
)
|
|
reset_approval_queue()
|
|
|
|
async def _handler(call: ToolCall) -> ToolResult:
|
|
approval_id = call.arguments.get("approval_id")
|
|
if approval_id is None:
|
|
approval_id = get_approval_queue().request(
|
|
"exec",
|
|
{
|
|
"toolName": call.tool_name,
|
|
"command": call.arguments["command"],
|
|
"args": dict(call.arguments),
|
|
},
|
|
)
|
|
status = "approval_required"
|
|
else:
|
|
status = "approval_pending"
|
|
return ToolResult(
|
|
tool_use_id=call.tool_use_id,
|
|
tool_name=call.tool_name,
|
|
content=json.dumps(
|
|
{
|
|
"status": status,
|
|
"approval_id": approval_id,
|
|
"command": call.arguments["command"],
|
|
"warning": "network target requires approval",
|
|
}
|
|
),
|
|
)
|
|
|
|
provider = _PendingApprovalShouldNotAnswerProvider()
|
|
agent = Agent(
|
|
provider=provider,
|
|
config=AgentConfig(
|
|
max_iterations=2,
|
|
metadata={"approval_wait_timeout_seconds": 0.01},
|
|
),
|
|
tool_definitions=[_exec_definition()],
|
|
tool_handler=_handler,
|
|
)
|
|
|
|
events = [event async for event in agent.run_turn("what is on this page?")]
|
|
|
|
assert len(provider.calls) == 1
|
|
assert not any(
|
|
getattr(event, "kind", "") == "text_delta"
|
|
and "fallback from training knowledge" in event.text
|
|
for event in events
|
|
)
|
|
assert any(
|
|
isinstance(event, ToolResultEvent) and "approval_pending" in event.result
|
|
for event in events
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_denied_approval_result_reaches_model_for_final_answer(
|
|
tmp_path: Any,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
from opensquilla.application import approval_queue as approval_queue_mod
|
|
|
|
monkeypatch.setattr(
|
|
approval_queue_mod,
|
|
"_DEFAULT_APPROVAL_QUEUE_PATH",
|
|
tmp_path / "approval_queue.sqlite",
|
|
)
|
|
reset_approval_queue()
|
|
approval_prompt_seen = asyncio.Event()
|
|
|
|
async def _handler(call: ToolCall) -> ToolResult:
|
|
approval_id = call.arguments.get("approval_id")
|
|
if approval_id is None:
|
|
approval_id = get_approval_queue().request(
|
|
"exec",
|
|
{
|
|
"toolName": call.tool_name,
|
|
"command": call.arguments["command"],
|
|
"args": dict(call.arguments),
|
|
},
|
|
)
|
|
return ToolResult(
|
|
tool_use_id=call.tool_use_id,
|
|
tool_name=call.tool_name,
|
|
content=json.dumps(
|
|
{
|
|
"status": "approval_required",
|
|
"approval_id": approval_id,
|
|
"command": call.arguments["command"],
|
|
"warning": "command requires approval",
|
|
}
|
|
),
|
|
)
|
|
|
|
entry = get_approval_queue().get(str(approval_id))
|
|
assert entry.resolved is True
|
|
assert entry.approved is False
|
|
return ToolResult(
|
|
tool_use_id=call.tool_use_id,
|
|
tool_name=call.tool_name,
|
|
content=json.dumps(
|
|
{
|
|
"status": "approval_denied",
|
|
"approval_id": approval_id,
|
|
"message": "The user denied access. Explain that the path cannot be inspected.",
|
|
}
|
|
),
|
|
)
|
|
|
|
provider = _DeniedApprovalThenAnswerProvider()
|
|
agent = Agent(
|
|
provider=provider,
|
|
config=AgentConfig(
|
|
max_iterations=3,
|
|
metadata={"approval_wait_timeout_seconds": 1.0},
|
|
),
|
|
tool_definitions=[_exec_definition()],
|
|
tool_handler=_handler,
|
|
)
|
|
|
|
events: list[Any] = []
|
|
|
|
async def _drive() -> None:
|
|
async for event in agent.run_turn("can you inspect /outside?"):
|
|
events.append(event)
|
|
if isinstance(event, ToolResultEvent) and "approval_required" in event.result:
|
|
approval_prompt_seen.set()
|
|
|
|
try:
|
|
task = asyncio.create_task(_drive())
|
|
await asyncio.wait_for(approval_prompt_seen.wait(), timeout=2.0)
|
|
approval_event = next(
|
|
event
|
|
for event in events
|
|
if isinstance(event, ToolResultEvent) and "approval_required" in event.result
|
|
)
|
|
approval_id = json.loads(approval_event.result)["approval_id"]
|
|
get_approval_queue().resolve(approval_id, False)
|
|
await asyncio.wait_for(task, timeout=2.0)
|
|
|
|
assert len(provider.calls) == 2
|
|
second_provider_request = provider.calls[1]
|
|
tool_result_blocks = [
|
|
block
|
|
for msg in second_provider_request
|
|
for block in msg.content
|
|
if getattr(block, "type", None) == "tool_result"
|
|
]
|
|
assert len(tool_result_blocks) == 1
|
|
assert "approval_denied" in tool_result_blocks[0].content
|
|
assert any(
|
|
getattr(event, "kind", "") == "text_delta"
|
|
and "用户拒绝访问" in event.text
|
|
for event in events
|
|
)
|
|
finally:
|
|
reset_approval_queue()
|