wehub-resource-sync
494 lines
15 KiB
Python
494 lines
15 KiB
Python
from __future__ import annotations
|
|
|
|
from collections.abc import Awaitable, Callable, Iterable, Sequence
|
|
from dataclasses import dataclass
|
|
from typing import Any, cast
|
|
|
|
from openai.types.responses import ResponseFunctionToolCall
|
|
|
|
from agents import Agent, Runner, RunResult, RunResultStreaming
|
|
from agents.items import ToolApprovalItem, ToolCallOutputItem, TResponseOutputItem
|
|
from agents.run_context import RunContextWrapper
|
|
from agents.run_internal.run_loop import NextStepInterruption, SingleStepResult
|
|
from agents.run_state import RunState as RunStateClass
|
|
|
|
from ..fake_model import FakeModel
|
|
|
|
HITL_REJECTION_MSG = "Tool execution was not approved."
|
|
|
|
|
|
@dataclass
|
|
class ApprovalScenario:
|
|
"""Container for approval-driven tool scenarios."""
|
|
|
|
tool: Any
|
|
raw_call: TResponseOutputItem
|
|
final_output: TResponseOutputItem
|
|
assert_result: Callable[[RunResult], None]
|
|
|
|
|
|
@dataclass
|
|
class PendingScenario:
|
|
"""Container for scenarios with pending approvals."""
|
|
|
|
tool: Any
|
|
raw_call: TResponseOutputItem
|
|
assert_result: Callable[[RunResult], None] | None = None
|
|
|
|
|
|
async def roundtrip_interruptions_via_run(
|
|
agent: Agent[Any],
|
|
model: FakeModel,
|
|
raw_call: Any,
|
|
*,
|
|
user_input: str = "test",
|
|
) -> list[ToolApprovalItem]:
|
|
"""Run once with a tool call, serialize state, and deserialize it."""
|
|
model.set_next_output([raw_call])
|
|
result = await Runner.run(agent, user_input)
|
|
assert result.interruptions, "expected an interruption"
|
|
state = result.to_state()
|
|
deserialized_state = await RunStateClass.from_json(agent, state.to_json())
|
|
return deserialized_state.get_interruptions()
|
|
|
|
|
|
async def assert_roundtrip_tool_name(
|
|
agent: Agent[Any],
|
|
model: FakeModel,
|
|
raw_call: TResponseOutputItem,
|
|
expected_tool_name: str,
|
|
*,
|
|
user_input: str,
|
|
) -> None:
|
|
"""Assert that deserialized interruptions keep the tool name intact."""
|
|
interruptions = await roundtrip_interruptions_via_run(
|
|
agent, model, raw_call, user_input=user_input
|
|
)
|
|
assert interruptions, "Interruptions should be preserved after deserialization"
|
|
assert interruptions[0].tool_name == expected_tool_name, (
|
|
f"{expected_tool_name} tool approval should be preserved, not converted to function"
|
|
)
|
|
|
|
|
|
def make_state_with_interruptions(
|
|
agent: Agent[Any],
|
|
interruptions: list[ToolApprovalItem],
|
|
*,
|
|
original_input: str = "test",
|
|
max_turns: int = 10,
|
|
) -> RunStateClass[Any, Agent[Any]]:
|
|
"""Create a RunState primed with interruptions."""
|
|
context = make_context_wrapper()
|
|
state = RunStateClass(
|
|
context=context,
|
|
original_input=original_input,
|
|
starting_agent=agent,
|
|
max_turns=max_turns,
|
|
)
|
|
state._current_step = NextStepInterruption(interruptions=interruptions)
|
|
return state
|
|
|
|
|
|
async def assert_tool_output_roundtrip(
|
|
agent: Agent[Any],
|
|
raw_output: Any,
|
|
expected_type: str,
|
|
*,
|
|
output: Any = "command output",
|
|
) -> None:
|
|
"""Ensure tool outputs keep their type through serialization and deserialization."""
|
|
context = make_context_wrapper()
|
|
state = RunStateClass(context=context, original_input="test", starting_agent=agent, max_turns=3)
|
|
state._generated_items = [
|
|
ToolCallOutputItem(
|
|
agent=agent,
|
|
raw_item=raw_output,
|
|
output=output,
|
|
)
|
|
]
|
|
|
|
json_data = state.to_json()
|
|
|
|
generated_items_json = json_data.get("generated_items", [])
|
|
assert len(generated_items_json) == 1, f"{expected_type} item should be serialized"
|
|
serialized_type = generated_items_json[0].get("raw_item", {}).get("type")
|
|
|
|
assert serialized_type == expected_type, (
|
|
f"Expected {expected_type} in serialized JSON, but got {serialized_type}. "
|
|
"Serialization should not coerce tool outputs."
|
|
)
|
|
|
|
deserialized_state = await RunStateClass.from_json(agent, json_data)
|
|
|
|
assert len(deserialized_state._generated_items) == 1, (
|
|
f"{expected_type} item should be deserialized."
|
|
)
|
|
deserialized_item = deserialized_state._generated_items[0]
|
|
assert isinstance(deserialized_item, ToolCallOutputItem)
|
|
|
|
raw_item = deserialized_item.raw_item
|
|
output_type = raw_item.get("type") if isinstance(raw_item, dict) else raw_item.type
|
|
|
|
assert output_type == expected_type, (
|
|
f"Expected {expected_type}, but got {output_type}. "
|
|
"Serialization should preserve the tool output type."
|
|
)
|
|
|
|
|
|
async def run_and_resume(
|
|
agent: Agent[Any],
|
|
model: Any,
|
|
raw_call: Any,
|
|
*,
|
|
user_input: str,
|
|
) -> RunResult:
|
|
"""Run once, then resume from the produced state."""
|
|
model.set_next_output([raw_call])
|
|
first = await Runner.run(agent, user_input)
|
|
return await Runner.run(agent, first.to_state())
|
|
|
|
|
|
def approve_first_interruption(
|
|
result: Any,
|
|
*,
|
|
always_approve: bool = False,
|
|
) -> RunStateClass[Any, Agent[Any]]:
|
|
"""Approve the first interruption on the result and return the updated state."""
|
|
assert getattr(result, "interruptions", None), "expected an approval interruption"
|
|
state = cast(RunStateClass[Any, Agent[Any]], result.to_state())
|
|
state.approve(result.interruptions[0], always_approve=always_approve)
|
|
return state
|
|
|
|
|
|
async def resume_after_first_approval(
|
|
agent: Agent[Any],
|
|
result: Any,
|
|
*,
|
|
always_approve: bool = False,
|
|
) -> RunResult:
|
|
"""Approve the first interruption and resume the run."""
|
|
state = approve_first_interruption(result, always_approve=always_approve)
|
|
return await Runner.run(agent, state)
|
|
|
|
|
|
async def resume_streamed_after_first_approval(
|
|
agent: Agent[Any],
|
|
result: Any,
|
|
*,
|
|
always_approve: bool = False,
|
|
) -> RunResultStreaming:
|
|
"""Approve the first interruption and resume a streamed run to completion."""
|
|
state = approve_first_interruption(result, always_approve=always_approve)
|
|
resumed = Runner.run_streamed(agent, state)
|
|
await consume_stream(resumed)
|
|
return resumed
|
|
|
|
|
|
async def run_and_resume_after_approval(
|
|
agent: Agent[Any],
|
|
model: Any,
|
|
raw_call: Any,
|
|
final_output: Any,
|
|
*,
|
|
user_input: str,
|
|
) -> RunResult:
|
|
"""Run, approve the first interruption, and resume."""
|
|
model.set_next_output([raw_call])
|
|
first = await Runner.run(agent, user_input)
|
|
state = approve_first_interruption(first, always_approve=True)
|
|
model.set_next_output([final_output])
|
|
return await Runner.run(agent, state)
|
|
|
|
|
|
def collect_tool_outputs(
|
|
items: Iterable[Any],
|
|
*,
|
|
output_type: str,
|
|
) -> list[ToolCallOutputItem]:
|
|
"""Return ToolCallOutputItems matching a raw_item type."""
|
|
return [
|
|
item
|
|
for item in items
|
|
if isinstance(item, ToolCallOutputItem)
|
|
and isinstance(item.raw_item, dict)
|
|
and item.raw_item.get("type") == output_type
|
|
]
|
|
|
|
|
|
async def consume_stream(result: Any) -> None:
|
|
"""Drain all stream events to completion."""
|
|
async for _ in result.stream_events():
|
|
pass
|
|
|
|
|
|
def assert_single_approval_interruption(
|
|
result: SingleStepResult,
|
|
*,
|
|
tool_name: str | None = None,
|
|
) -> ToolApprovalItem:
|
|
"""Assert the result contains exactly one approval interruption and return it."""
|
|
assert isinstance(result.next_step, NextStepInterruption)
|
|
assert len(result.next_step.interruptions) == 1
|
|
interruption = result.next_step.interruptions[0]
|
|
assert isinstance(interruption, ToolApprovalItem)
|
|
if tool_name:
|
|
assert interruption.tool_name == tool_name
|
|
return interruption
|
|
|
|
|
|
async def require_approval(
|
|
_ctx: Any | None = None, _params: Any = None, _call_id: str | None = None
|
|
) -> bool:
|
|
"""Approval helper that always requires a HITL decision."""
|
|
return True
|
|
|
|
|
|
class RecordingEditor:
|
|
"""Editor that records operations for testing."""
|
|
|
|
def __init__(self) -> None:
|
|
self.operations: list[Any] = []
|
|
|
|
def create_file(self, operation: Any) -> Any:
|
|
self.operations.append(operation)
|
|
return {"output": f"Created {operation.path}", "status": "completed"}
|
|
|
|
def update_file(self, operation: Any) -> Any:
|
|
self.operations.append(operation)
|
|
return {"output": f"Updated {operation.path}", "status": "completed"}
|
|
|
|
def delete_file(self, operation: Any) -> Any:
|
|
self.operations.append(operation)
|
|
return {"output": f"Deleted {operation.path}", "status": "completed"}
|
|
|
|
|
|
def make_shell_call(
|
|
call_id: str,
|
|
*,
|
|
id_value: str | None = None,
|
|
commands: list[str] | None = None,
|
|
status: str = "in_progress",
|
|
) -> TResponseOutputItem:
|
|
"""Build a shell_call payload with optional overrides."""
|
|
return cast(
|
|
TResponseOutputItem,
|
|
{
|
|
"type": "shell_call",
|
|
"id": id_value or call_id,
|
|
"call_id": call_id,
|
|
"status": status,
|
|
"action": {"type": "exec", "commands": commands or ["echo test"], "timeout_ms": 1000},
|
|
},
|
|
)
|
|
|
|
|
|
def make_apply_patch_dict(call_id: str, diff: str = "-a\n+b\n") -> TResponseOutputItem:
|
|
"""Create an apply_patch_call dict payload."""
|
|
return cast(
|
|
TResponseOutputItem,
|
|
{
|
|
"type": "apply_patch_call",
|
|
"call_id": call_id,
|
|
"operation": {"type": "update_file", "path": "test.md", "diff": diff},
|
|
},
|
|
)
|
|
|
|
|
|
def make_function_tool_call(
|
|
name: str,
|
|
*,
|
|
call_id: str = "call-1",
|
|
arguments: str = "{}",
|
|
namespace: str | None = None,
|
|
) -> ResponseFunctionToolCall:
|
|
"""Create a ResponseFunctionToolCall for HITL scenarios."""
|
|
if namespace is None:
|
|
return ResponseFunctionToolCall(
|
|
type="function_call",
|
|
name=name,
|
|
call_id=call_id,
|
|
arguments=arguments,
|
|
)
|
|
return ResponseFunctionToolCall(
|
|
type="function_call",
|
|
name=name,
|
|
call_id=call_id,
|
|
arguments=arguments,
|
|
namespace=namespace,
|
|
)
|
|
|
|
|
|
def queue_function_call_and_text(
|
|
model: FakeModel,
|
|
function_call: TResponseOutputItem,
|
|
*,
|
|
first_turn_extra: Sequence[TResponseOutputItem] | None = None,
|
|
followup: Sequence[TResponseOutputItem] | None = None,
|
|
) -> None:
|
|
"""Queue a function call turn followed by a follow-up turn on the fake model."""
|
|
raw_type = (
|
|
function_call.get("type")
|
|
if isinstance(function_call, dict)
|
|
else getattr(function_call, "type", None)
|
|
)
|
|
assert raw_type == "function_call", "queue_function_call_and_text expects a function call item"
|
|
model.add_multiple_turn_outputs(
|
|
[
|
|
[function_call, *(first_turn_extra or [])],
|
|
list(followup or []),
|
|
]
|
|
)
|
|
|
|
|
|
async def run_and_resume_with_mutation(
|
|
agent: Agent[Any],
|
|
model: Any,
|
|
turn_outputs: Sequence[Sequence[Any]],
|
|
*,
|
|
user_input: str,
|
|
mutate_state: Callable[[RunStateClass[Any, Agent[Any]], ToolApprovalItem], None] | None = None,
|
|
) -> tuple[RunResult, RunResult]:
|
|
"""Run until interruption, optionally mutate state, then resume."""
|
|
model.add_multiple_turn_outputs(turn_outputs)
|
|
first = await Runner.run(agent, input=user_input)
|
|
assert first.interruptions, "expected an approval interruption"
|
|
state = first.to_state()
|
|
if mutate_state and first.interruptions:
|
|
mutate_state(state, first.interruptions[0])
|
|
resumed = await Runner.run(agent, input=state)
|
|
return first, resumed
|
|
|
|
|
|
async def assert_pending_resume(
|
|
tool: Any,
|
|
model: Any,
|
|
raw_call: TResponseOutputItem,
|
|
*,
|
|
user_input: str,
|
|
output_type: str,
|
|
) -> RunResult:
|
|
"""Run, resume, and assert pending approvals stay pending."""
|
|
agent = make_agent(model=model, tools=[tool])
|
|
|
|
resumed = await run_and_resume(agent, model, raw_call, user_input=user_input)
|
|
|
|
assert resumed.interruptions, "pending approval should remain after resuming"
|
|
assert any(
|
|
isinstance(item, ToolApprovalItem) and item.tool_name == tool.name
|
|
for item in resumed.interruptions
|
|
)
|
|
assert not collect_tool_outputs(resumed.new_items, output_type=output_type), (
|
|
f"{output_type} should not execute without approval"
|
|
)
|
|
return resumed
|
|
|
|
|
|
def make_mcp_raw_item(
|
|
*,
|
|
call_id: str = "call_mcp_1",
|
|
include_provider_data: bool = True,
|
|
tool_name: str = "test_mcp_tool",
|
|
provider_data: dict[str, Any] | None = None,
|
|
include_name: bool = True,
|
|
use_call_id: bool = True,
|
|
) -> dict[str, Any]:
|
|
"""Build a hosted MCP tool call payload for approvals."""
|
|
|
|
raw_item: dict[str, Any] = {"type": "hosted_tool_call"}
|
|
if include_name:
|
|
raw_item["name"] = tool_name
|
|
if include_provider_data:
|
|
if use_call_id:
|
|
raw_item["call_id"] = call_id
|
|
else:
|
|
raw_item["id"] = call_id
|
|
raw_item["provider_data"] = provider_data or {
|
|
"type": "mcp_approval_request",
|
|
"id": "req-1",
|
|
"server_label": "test_server",
|
|
}
|
|
else:
|
|
raw_item["id"] = call_id
|
|
return raw_item
|
|
|
|
|
|
def make_mcp_approval_item(
|
|
agent: Agent[Any],
|
|
*,
|
|
call_id: str = "call_mcp_1",
|
|
include_provider_data: bool = True,
|
|
tool_name: str | None = "test_mcp_tool",
|
|
provider_data: dict[str, Any] | None = None,
|
|
include_name: bool = True,
|
|
use_call_id: bool = True,
|
|
) -> ToolApprovalItem:
|
|
"""Create a ToolApprovalItem for MCP or hosted tool calls."""
|
|
|
|
raw_item = make_mcp_raw_item(
|
|
call_id=call_id,
|
|
include_provider_data=include_provider_data,
|
|
tool_name=tool_name or "unknown_mcp_tool",
|
|
provider_data=provider_data,
|
|
include_name=include_name,
|
|
use_call_id=use_call_id,
|
|
)
|
|
return ToolApprovalItem(agent=agent, raw_item=raw_item, tool_name=tool_name)
|
|
|
|
|
|
def make_context_wrapper() -> RunContextWrapper[dict[str, Any]]:
|
|
"""Create an empty RunContextWrapper for HITL tests."""
|
|
return RunContextWrapper(context={})
|
|
|
|
|
|
def make_agent(
|
|
*,
|
|
model: Any | None = None,
|
|
tools: Sequence[Any] | None = None,
|
|
name: str = "TestAgent",
|
|
) -> Agent[Any]:
|
|
"""Build a test Agent with optional model and tools."""
|
|
return Agent(name=name, model=model, tools=list(tools or []))
|
|
|
|
|
|
def make_model_and_agent(
|
|
*,
|
|
tools: Sequence[Any] | None = None,
|
|
name: str = "TestAgent",
|
|
) -> tuple[FakeModel, Agent[Any]]:
|
|
"""Build a FakeModel with a paired Agent for HITL tests."""
|
|
model = FakeModel()
|
|
agent = make_agent(model=model, tools=tools, name=name)
|
|
return model, agent
|
|
|
|
|
|
def reject_tool_call(
|
|
context_wrapper: RunContextWrapper[Any],
|
|
agent: Agent[Any],
|
|
raw_item: Any,
|
|
tool_name: str,
|
|
*,
|
|
rejection_message: str | None = None,
|
|
) -> ToolApprovalItem:
|
|
"""Reject a tool call in the context and return the approval item used."""
|
|
approval_item = ToolApprovalItem(agent=agent, raw_item=raw_item, tool_name=tool_name)
|
|
context_wrapper.reject_tool(approval_item, rejection_message=rejection_message)
|
|
return approval_item
|
|
|
|
|
|
def make_on_approval_callback(
|
|
approve: bool,
|
|
*,
|
|
reason: str | None = None,
|
|
) -> Callable[[RunContextWrapper[Any], ToolApprovalItem], Awaitable[Any]]:
|
|
"""Build an on_approval callback that always approves or rejects."""
|
|
|
|
async def on_approval(
|
|
_ctx: RunContextWrapper[Any], _approval_item: ToolApprovalItem
|
|
) -> dict[str, Any]:
|
|
payload: dict[str, Any] = {"approve": approve}
|
|
if reason:
|
|
payload["reason"] = reason
|
|
return payload
|
|
|
|
return on_approval
|