Agent/tests/test_application_runtime.py

import asyncio
import json
import unittest

from pydantic import ValidationError

from cyber_agent.application import (
    AgentApplication,
    AgentRole,
    ApplicationRegistry,
    ApplicationServices,
    ProviderRef,
    RoleRunLimits,
    RunLimits,
    SkillSpec,
    ToolSet,
    ToolSpec,
)
from cyber_agent.tools.registry import ToolRegistry


def application(*, prompt="write carefully", provider_version="1", limit=20):
    return AgentApplication(
        application_id="creative.reference",
        application_version="0.1.0",
        root_role="writer",
        roles=(AgentRole(
            role_id="writer",
            model="fake-model",
            system_prompt=prompt,
            skills=(SkillSpec(
                name="evidence-writing",
                version="1",
                content="Use checked evidence.",
            ),),
            tool_sets=("writing",),
            allowed_child_roles=("writer",),
            run_limits=RoleRunLimits(max_iterations=limit),
        ),),
        tool_sets=(ToolSet(
            tool_set_id="writing",
            tools=(ToolSpec(
                tool_name="save",
                implementation_version="1",
                side_effect_class="write",
                idempotent=True,
            ),),
        ),),
        context_provider_ref=ProviderRef(
            provider_id="context",
            provider_version=provider_version,
        ),
        run_limits=RunLimits(max_iterations=30),
    )


def registry(output):
    tools = ToolRegistry()

    async def save(value: str):
        return {"application": output, "value": value}

    tools.register(save, groups=["application"])
    return tools


class ApplicationDeclarationTest(unittest.TestCase):
    def test_role_graph_and_tool_set_references_fail_closed(self):
        with self.assertRaisesRegex(ValidationError, "unknown child roles"):
            AgentApplication(
                application_id="broken",
                application_version="1",
                root_role="root",
                roles=(AgentRole(
                    role_id="root",
                    model="fake",
                    system_prompt="root",
                    allowed_child_roles=("missing",),
                ),),
            )
        with self.assertRaisesRegex(ValidationError, "unknown tool sets"):
            AgentApplication(
                application_id="broken",
                application_version="1",
                root_role="root",
                roles=(AgentRole(
                    role_id="root",
                    model="fake",
                    system_prompt="root",
                    tool_sets=("missing",),
                ),),
            )

    def test_registry_freezes_private_tool_copy_and_resets_stats(self):
        source = registry("A")
        source._stats["save"].call_count = 7
        binding = ApplicationRegistry().register(
            application(),
            ApplicationServices(
                tool_registry=source,
                context_provider=object(),
            ),
        )

        self.assertTrue(binding.tool_registry.frozen)
        self.assertEqual(0, binding.tool_registry.get_stats("save")["save"]["call_count"])
        with self.assertRaisesRegex(RuntimeError, "frozen"):
            binding.tool_registry.register(lambda: None)
        self.assertFalse(source.frozen)

    def test_same_tool_name_is_isolated_between_application_bindings(self):
        app_a = application()
        app_b = app_a.model_copy(update={
            "application_id": "creative.other",
            "roles": (
                app_a.roles[0].model_copy(update={"system_prompt": "other prompt"}),
            ),
        })
        binding_a = ApplicationRegistry().register(
            app_a,
            ApplicationServices(
                tool_registry=registry("A"),
                context_provider=object(),
            ),
        )
        binding_b = ApplicationRegistry().register(
            app_b,
            ApplicationServices(
                tool_registry=registry("B"),
                context_provider=object(),
            ),
        )

        self.assertIsNot(binding_a.tool_registry, binding_b.tool_registry)
        self.assertNotEqual(
            binding_a.application_ref.config_hash,
            binding_b.application_ref.config_hash,
        )
        self.assertNotEqual(
            binding_a.role("writer").role_hash,
            binding_b.role("writer").role_hash,
        )

    def test_hash_covers_prompt_provider_limits_and_tool_schema(self):
        base = ApplicationRegistry().register(
            application(),
            ApplicationServices(
                tool_registry=registry("base"),
                context_provider=object(),
            ),
        ).application_ref.config_hash
        variants = [
            application(prompt="changed prompt"),
            application(provider_version="2"),
            application(limit=10),
        ]
        hashes = []
        for item in variants:
            hashes.append(ApplicationRegistry().register(
                item,
                ApplicationServices(
                    tool_registry=registry("base"),
                    context_provider=object(),
                ),
            ).application_ref.config_hash)
        self.assertTrue(all(item != base for item in hashes))
        changed_schema = registry("base")
        definition = changed_schema._tools["save"]
        definition["schema"] = {
            **definition["schema"],
            "function": {
                **definition["schema"]["function"],
                "description": "changed schema",
            },
        }
        schema_hash = ApplicationRegistry().register(
            application(),
            ApplicationServices(
                tool_registry=changed_schema,
                context_provider=object(),
            ),
        ).application_ref.config_hash
        self.assertNotEqual(base, schema_hash)


class ApplicationToolIsolationTest(unittest.IsolatedAsyncioTestCase):
    async def test_same_named_tools_execute_in_private_registries_concurrently(self):
        app_a = application()
        app_b = app_a.model_copy(update={"application_id": "creative.other"})
        binding_a = ApplicationRegistry().register(
            app_a,
            ApplicationServices(
                tool_registry=registry("A"),
                context_provider=object(),
            ),
        )
        binding_b = ApplicationRegistry().register(
            app_b,
            ApplicationServices(
                tool_registry=registry("B"),
                context_provider=object(),
            ),
        )

        async def execute(binding, expected, index):
            raw = await binding.tool_registry.execute(
                "save",
                {"value": str(index)},
                uid="user",
                allowed_tool_names={"save"},
            )
            payload = json.loads(raw) if isinstance(raw, str) else raw
            self.assertEqual(expected, payload["application"])
            self.assertEqual(str(index), payload["value"])

        await asyncio.gather(*(
            execute(binding, expected, index)
            for binding, expected in ((binding_a, "A"), (binding_b, "B"))
            for index in range(20)
        ))
        self.assertEqual(
            20,
            binding_a.tool_registry.get_stats("save")["save"]["call_count"],
        )
        self.assertEqual(
            20,
            binding_b.tool_registry.get_stats("save")["save"]["call_count"],
        )

    def test_runtime_service_identity_is_not_part_of_config_hash(self):
        first = ApplicationRegistry().register(
            application(),
            ApplicationServices(
                tool_registry=registry("same"),
                context_provider=object(),
            ),
        )
        second = ApplicationRegistry().register(
            application(),
            ApplicationServices(
                tool_registry=registry("same"),
                context_provider=object(),
            ),
        )
        self.assertEqual(first.application_ref, second.application_ref)

    def test_missing_declared_provider_is_rejected(self):
        with self.assertRaisesRegex(ValueError, "context_provider"):
            ApplicationRegistry().register(
                application(),
                ApplicationServices(tool_registry=registry("A")),
            )


if __name__ == "__main__":
    unittest.main()