wehub-resource-sync
252 lines
7.6 KiB
Python
252 lines
7.6 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
from dataclasses import replace
|
|
from typing import Any
|
|
|
|
from opensquilla.engine import Agent, AgentConfig
|
|
from opensquilla.sandbox.config import SandboxSettings
|
|
from opensquilla.sandbox.integration import configure_runtime, reset_runtime
|
|
from opensquilla.tool_boundary import ToolCall
|
|
from opensquilla.tools.dispatch import build_tool_handler
|
|
from opensquilla.tools.registry import ToolRegistry, get_default_registry
|
|
from opensquilla.tools.types import (
|
|
InteractionMode,
|
|
ToolContext,
|
|
ToolSpec,
|
|
current_tool_context,
|
|
)
|
|
|
|
|
|
class _Provider:
|
|
provider_name = "fake"
|
|
|
|
def chat(self, messages, tools=None, config=None):
|
|
raise AssertionError("provider should not be used")
|
|
|
|
async def list_models(self) -> list[Any]:
|
|
return []
|
|
|
|
|
|
def test_agent_wires_tool_context_runtime_events(tmp_path) -> None:
|
|
runtime_events_path = tmp_path / "runtime_events.jsonl"
|
|
agent = Agent(
|
|
provider=_Provider(),
|
|
config=AgentConfig(
|
|
runtime_events_path=str(runtime_events_path),
|
|
metadata={"agent_id": "configured-agent"},
|
|
),
|
|
session_key="agent-session",
|
|
tool_context=ToolContext(),
|
|
)
|
|
|
|
assert agent._tool_context is not None
|
|
assert agent._tool_context.on_runtime_event is not None
|
|
|
|
agent._tool_context.on_runtime_event(
|
|
{
|
|
"feature": "fresh_read_guard",
|
|
"name": "fresh_read_guard.blocked",
|
|
"tool": "edit_file",
|
|
"outcome": "blocked",
|
|
}
|
|
)
|
|
|
|
events = [
|
|
json.loads(line)
|
|
for line in runtime_events_path.read_text(encoding="utf-8").splitlines()
|
|
]
|
|
assert len(events) == 1
|
|
assert events[0]["feature"] == "fresh_read_guard"
|
|
assert events[0]["name"] == "fresh_read_guard.blocked"
|
|
assert events[0]["tool"] == "edit_file"
|
|
assert events[0]["outcome"] == "blocked"
|
|
assert events[0]["session_key"] == "agent-session"
|
|
assert events[0]["agent_id"] == "configured-agent"
|
|
assert "created_at" in events[0]
|
|
assert "timestamp" in events[0]
|
|
|
|
|
|
def test_agent_preserves_tool_context_runtime_event_identity(tmp_path) -> None:
|
|
runtime_events_path = tmp_path / "runtime_events.jsonl"
|
|
agent = Agent(
|
|
provider=_Provider(),
|
|
config=AgentConfig(
|
|
runtime_events_path=str(runtime_events_path),
|
|
metadata={"agent_id": "configured-agent"},
|
|
),
|
|
session_key="agent-session",
|
|
tool_context=ToolContext(),
|
|
)
|
|
|
|
assert agent._tool_context is not None
|
|
assert agent._tool_context.on_runtime_event is not None
|
|
|
|
agent._tool_context.on_runtime_event(
|
|
{
|
|
"feature": "fresh_read_guard",
|
|
"name": "fresh_read_guard.blocked",
|
|
"agent_id": "tool-agent",
|
|
"session_key": "tool-session",
|
|
"outcome": "blocked",
|
|
}
|
|
)
|
|
|
|
event = json.loads(runtime_events_path.read_text(encoding="utf-8"))
|
|
assert event["agent_id"] == "tool-agent"
|
|
assert event["session_key"] == "tool-session"
|
|
|
|
|
|
def test_agent_backfills_none_tool_context_runtime_event_identity(tmp_path) -> None:
|
|
runtime_events_path = tmp_path / "runtime_events.jsonl"
|
|
agent = Agent(
|
|
provider=_Provider(),
|
|
config=AgentConfig(
|
|
runtime_events_path=str(runtime_events_path),
|
|
metadata={"agent_id": "configured-agent"},
|
|
),
|
|
session_key="agent-session",
|
|
tool_context=ToolContext(),
|
|
)
|
|
|
|
assert agent._tool_context is not None
|
|
assert agent._tool_context.on_runtime_event is not None
|
|
|
|
agent._tool_context.on_runtime_event(
|
|
{
|
|
"feature": "fresh_read_guard",
|
|
"name": "fresh_read_guard.blocked",
|
|
"agent_id": None,
|
|
"session_key": None,
|
|
"outcome": "blocked",
|
|
}
|
|
)
|
|
|
|
event = json.loads(runtime_events_path.read_text(encoding="utf-8"))
|
|
assert event["agent_id"] == "configured-agent"
|
|
assert event["session_key"] == "agent-session"
|
|
|
|
|
|
def test_agent_runtime_event_callback_reaches_prebuilt_tool_handler(tmp_path) -> None:
|
|
runtime_events_path = tmp_path / "runtime_events.jsonl"
|
|
ctx = ToolContext()
|
|
registry = ToolRegistry()
|
|
|
|
async def emit_event() -> str:
|
|
active = current_tool_context.get()
|
|
assert active is ctx
|
|
assert active.on_runtime_event is not None
|
|
active.on_runtime_event(
|
|
{
|
|
"feature": "tool_runtime",
|
|
"name": "tool_runtime.probe",
|
|
"outcome": "used",
|
|
}
|
|
)
|
|
return "ok"
|
|
|
|
registry.register(
|
|
ToolSpec(
|
|
name="emit_event",
|
|
description="Emit a runtime event.",
|
|
parameters={},
|
|
required=[],
|
|
),
|
|
emit_event,
|
|
)
|
|
handler = build_tool_handler(registry, ctx)
|
|
|
|
agent = Agent(
|
|
provider=_Provider(),
|
|
config=AgentConfig(
|
|
runtime_events_path=str(runtime_events_path),
|
|
metadata={"agent_id": "configured-agent"},
|
|
),
|
|
tool_handler=handler,
|
|
session_key="agent-session",
|
|
tool_context=ctx,
|
|
)
|
|
|
|
assert agent._tool_context is ctx
|
|
assert ctx.on_runtime_event is not None
|
|
|
|
import asyncio
|
|
|
|
result = asyncio.run(handler(ToolCall("call-1", "emit_event", {})))
|
|
assert result.content == "ok"
|
|
|
|
event = json.loads(runtime_events_path.read_text(encoding="utf-8"))
|
|
assert event["feature"] == "tool_runtime"
|
|
assert event["name"] == "tool_runtime.probe"
|
|
assert event["agent_id"] == "configured-agent"
|
|
assert event["session_key"] == "agent-session"
|
|
|
|
|
|
def test_agent_runtime_event_context_reaches_replaced_builtin_tool_handler(
|
|
tmp_path,
|
|
) -> None:
|
|
runtime_events_path = tmp_path / "runtime_events.jsonl"
|
|
target = tmp_path / "src" / "app.py"
|
|
target.parent.mkdir()
|
|
target.write_text(
|
|
"def run():\n"
|
|
" if enabled:\n"
|
|
" return 1\n",
|
|
encoding="utf-8",
|
|
)
|
|
|
|
configure_runtime(
|
|
SandboxSettings(
|
|
sandbox=False,
|
|
security_grading=False,
|
|
allow_legacy_mode=True,
|
|
),
|
|
workspace=tmp_path,
|
|
)
|
|
try:
|
|
handler_ctx = ToolContext(
|
|
is_owner=True,
|
|
interaction_mode=InteractionMode.UNATTENDED,
|
|
workspace_dir=str(tmp_path),
|
|
elevated="bypass",
|
|
)
|
|
handler = build_tool_handler(get_default_registry(), handler_ctx)
|
|
agent_ctx = replace(handler_ctx, session_key="agent-session")
|
|
agent = Agent(
|
|
provider=_Provider(),
|
|
config=AgentConfig(
|
|
runtime_events_path=str(runtime_events_path),
|
|
metadata={"agent_id": "configured-agent"},
|
|
),
|
|
tool_handler=handler,
|
|
session_key="agent-session",
|
|
tool_context=agent_ctx,
|
|
)
|
|
|
|
assert agent.tool_handler is not None
|
|
result = asyncio.run(
|
|
agent.tool_handler(
|
|
ToolCall(
|
|
"call-1",
|
|
"edit_file",
|
|
{
|
|
"path": str(target),
|
|
"old_text": "if enabled:\n return 1\n",
|
|
"new_text": "if enabled:\n return 2\n",
|
|
},
|
|
)
|
|
)
|
|
)
|
|
finally:
|
|
reset_runtime()
|
|
|
|
assert result.is_error is False
|
|
events = [
|
|
json.loads(line)
|
|
for line in runtime_events_path.read_text(encoding="utf-8").splitlines()
|
|
]
|
|
assert any(event["name"] == "edit_file.flexible_match_used" for event in events)
|
|
assert events[0]["agent_id"] == "main"
|
|
assert events[0]["session_key"] == "agent-session"
|