Files
2026-07-13 13:12:33 +08:00

236 lines
7.2 KiB
Python

from __future__ import annotations
import ast
from pathlib import Path
from typing import Any
import pytest
from opensquilla.gateway.auth import Principal
from opensquilla.gateway.config import GatewayConfig
from opensquilla.gateway.rpc import RpcContext, get_dispatcher
from opensquilla.gateway.scopes import READ_SCOPE
from opensquilla.tools.registry import ToolRegistry
from opensquilla.tools.types import ToolSpec
async def _handler() -> str:
return "ok"
def _ctx(*, tool_registry: Any, is_owner: bool) -> RpcContext:
return RpcContext(
conn_id="test",
config=GatewayConfig(),
principal=Principal(
role="operator",
scopes=frozenset({READ_SCOPE}),
is_owner=is_owner,
authenticated=True,
),
tool_registry=tool_registry,
session_manager=object(),
task_runtime=object(),
)
def _tool_names(payload: dict[str, Any]) -> set[str]:
return {tool["name"] for tool in payload["tools"]}
def test_tools_rpc_delegates_payloads_to_tools_boundary() -> None:
from opensquilla.gateway import rpc_tools
from opensquilla.tools import registry
source = Path(rpc_tools.__file__).read_text(encoding="utf-8")
tree = ast.parse(source)
registry_tree = ast.parse(Path(registry.__file__).read_text(encoding="utf-8"))
boundary_path = Path(registry.__file__).with_name("rpc_payload.py")
assert boundary_path.exists()
boundary_tree = ast.parse(boundary_path.read_text(encoding="utf-8"))
imports = {
(node.module, alias.name)
for node in ast.walk(tree)
if isinstance(node, ast.ImportFrom) and node.module
for alias in node.names
}
registry_defs = {
node.name
for node in registry_tree.body
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))
}
boundary_defs = {
node.name
for node in boundary_tree.body
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))
}
assert {
("opensquilla.tools.rpc_payload", "tools_catalog_payload"),
("opensquilla.tools.rpc_payload", "tools_effective_payload"),
} <= imports
assert {
"tools_catalog_payload",
"tools_effective_payload",
} <= registry_defs
assert {
"tool_rpc_params",
"tool_surface_capabilities_for_runtime",
"tools_catalog_payload",
"tools_effective_payload",
} <= boundary_defs
def _registry_with_owner_only_probe() -> ToolRegistry:
registry = ToolRegistry()
registry.register(
ToolSpec(
name="ordinary_probe",
description="ordinary probe",
parameters={},
),
_handler,
)
registry.register(
ToolSpec(
name="owner_probe",
description="owner probe",
parameters={},
owner_only=True,
),
_handler,
)
return registry
@pytest.mark.asyncio
@pytest.mark.parametrize("method", ["tools.catalog", "tools.effective"])
async def test_tools_rpc_visibility_respects_principal_ownership(method: str) -> None:
registry = _registry_with_owner_only_probe()
non_owner = await get_dispatcher().dispatch(
"r1",
method,
{"callerKind": "agent"},
_ctx(tool_registry=registry, is_owner=False),
)
owner = await get_dispatcher().dispatch(
"r2",
method,
{"callerKind": "agent"},
_ctx(tool_registry=registry, is_owner=True),
)
assert non_owner.error is None, non_owner.error
assert owner.error is None, owner.error
assert _tool_names(non_owner.payload) == {"ordinary_probe"}
assert _tool_names(owner.payload) == {"ordinary_probe", "owner_probe"}
@pytest.mark.asyncio
async def test_tools_catalog_without_runtime_params_respects_principal_ownership() -> None:
registry = _registry_with_owner_only_probe()
non_owner = await get_dispatcher().dispatch(
"r1",
"tools.catalog",
{},
_ctx(tool_registry=registry, is_owner=False),
)
assert non_owner.error is None, non_owner.error
assert _tool_names(non_owner.payload) == {"ordinary_probe"}
@pytest.mark.asyncio
@pytest.mark.parametrize(
"params",
[
{"callerKind": "subagent"},
{"sessionKey": "subagent:test"},
],
)
@pytest.mark.parametrize("method", ["tools.catalog", "tools.effective"])
async def test_tools_rpc_subagent_visibility_respects_principal_ownership(
method: str,
params: dict[str, str],
) -> None:
registry = _registry_with_owner_only_probe()
non_owner = await get_dispatcher().dispatch(
"r1",
method,
params,
_ctx(tool_registry=registry, is_owner=False),
)
assert non_owner.error is None, non_owner.error
assert _tool_names(non_owner.payload) == {"ordinary_probe"}
@pytest.mark.asyncio
@pytest.mark.parametrize("method", ["tools.catalog", "tools.effective"])
async def test_default_tools_rpc_hides_owner_only_tools_from_non_owner(method: str) -> None:
import opensquilla.tools.builtin # noqa: F401
from opensquilla.gateway.config import ImageGenerationConfig, LlmProviderConfig
from opensquilla.tools.builtin.media import configure_image_generation
from opensquilla.tools.registry import get_default_registry
configure_image_generation(
ImageGenerationConfig(enabled=True),
llm_config=LlmProviderConfig(provider="openrouter", api_key="sk-or-configured"),
)
try:
non_owner = await get_dispatcher().dispatch(
"r1",
method,
{"callerKind": "agent"},
_ctx(tool_registry=get_default_registry(), is_owner=False),
)
owner = await get_dispatcher().dispatch(
"r2",
method,
{"callerKind": "agent"},
_ctx(tool_registry=get_default_registry(), is_owner=True),
)
finally:
configure_image_generation(ImageGenerationConfig())
assert non_owner.error is None, non_owner.error
assert owner.error is None, owner.error
non_owner_names = _tool_names(non_owner.payload)
owner_names = _tool_names(owner.payload)
assert "http_request" not in non_owner_names
assert "git_commit" not in non_owner_names
assert {"http_request", "git_commit"} <= owner_names
assert {"image_generate", "sessions_spawn", "sessions_send"} <= owner_names
assert "spawn_subagent" not in owner_names
assert "send_message" not in owner_names
assert "generate_image" not in owner_names
@pytest.mark.asyncio
@pytest.mark.parametrize("method", ["tools.catalog", "tools.effective"])
async def test_default_channel_tools_rpc_exposes_structured_file_authoring(method: str) -> None:
import opensquilla.tools.builtin # noqa: F401
from opensquilla.tools.registry import get_default_registry
result = await get_dispatcher().dispatch(
"r1",
method,
{"callerKind": "channel"},
_ctx(tool_registry=get_default_registry(), is_owner=False),
)
assert result.error is None, result.error
names = _tool_names(result.payload)
assert {"create_csv", "create_xlsx", "create_pdf_report", "create_pptx"} <= names
assert "write_file" not in names
assert "execute_code" not in names
assert "apply_patch" not in names