ai-video-fullstack/backend/tests/test_tool_organization_v4.py

from __future__ import annotations

import importlib
import unittest

from fastapi import HTTPException
from models import RuntimeTool, RuntimeToolBinding
from routes.tools import _payload
from schemas import AssistantOut, AssistantUpsert, ToolUpsert
from services.runtime_variables import DynamicVariableStore
from services.state_expressions import (
    StateExpressionError,
    evaluate_update_rules,
    validate_update_rules,
)
from services.system_tools import runtime_system_tools
from services.tool_executor import ToolExecutionError, ToolExecutor


class ToolOrganizationContractTests(unittest.TestCase):
    def test_system_is_not_a_resource_type(self):
        with self.assertRaises(ValueError):
            ToolUpsert(
                name="结束对话",
                functionName="end_conversation",
                definition={
                    "schemaVersion": 2,
                    "type": "system",
                    "config": {},
                },
            )

    def test_prompt_keeps_resources_and_builtins_separate(self):
        assistant = AssistantUpsert(
            name="客服",
            type="prompt",
            toolBindings=[
                {
                    "toolId": "tool_lookup",
                    "argumentBindings": {"order_id": {"source": "model"}},
                }
            ],
            builtInTools={"endConversation": {"captureReason": True}},
        )

        self.assertEqual(assistant.tool_bindings[0].tool_id, "tool_lookup")
        tools = runtime_system_tools(assistant.built_in_tools.model_dump())
        self.assertEqual([tool.function_name for tool in tools], ["end_conversation"])

    def test_realtime_prompt_request_and_response_preserve_builtins(self):
        data = {
            "name": "历史 Realtime 助手",
            "type": "prompt",
            "runtimeMode": "realtime",
            "builtInTools": {"endConversation": {"captureReason": True}},
        }

        request = AssistantUpsert(**data)
        self.assertIsNotNone(request.built_in_tools.end_conversation)

        response = AssistantOut(id="asst_realtime", **data)
        self.assertIsNotNone(response.built_in_tools.end_conversation)

    def test_workflow_camel_case_builtin_config_is_normalized_for_runtime(self):
        [tool] = runtime_system_tools(
            {
                "endConversation": {
                    "messageType": "custom",
                    "customMessage": "再见",
                    "captureReason": False,
                    "allowEndInterruption": True,
                }
            }
        )

        config = tool.definition["config"]
        self.assertEqual(config["message_type"], "custom")
        self.assertEqual(config["custom_message"], "再见")
        self.assertFalse(config["capture_reason"])
        self.assertTrue(config["allow_end_interruption"])

    def test_workflow_public_boundary_is_v4(self):
        assistant = AssistantUpsert(
            name="工作流",
            type="workflow",
            graph={"specVersion": 4, "settings": {}, "nodes": [], "edges": []},
        )
        self.assertEqual(assistant.graph["specVersion"], 4)

    def test_http_resource_derives_explicit_template_inputs(self):
        body = ToolUpsert(
            name="订单查询",
            functionName="query_order",
            definition={
                "schemaVersion": 2,
                "type": "http",
                "config": {
                    "url": "https://example.com/{{input__tenant}}/orders",
                    "headers": {},
                    "body": {},
                },
            },
        )

        payload = _payload(body)

        self.assertEqual(
            payload["definition"]["config"]["template_inputs"][0]["name"],
            "tenant",
        )

    def test_http_resource_rejects_assistant_variable_placeholder(self):
        body = ToolUpsert(
            name="订单查询",
            functionName="query_order",
            definition={
                "schemaVersion": 2,
                "type": "http",
                "config": {
                    "url": "https://example.com/{{tenant}}/orders",
                    "headers": {},
                    "body": {},
                },
            },
        )

        with self.assertRaisesRegex(HTTPException, "input__name"):
            _payload(body)

    def test_http_resource_rejects_non_local_secret_name(self):
        body = ToolUpsert(
            name="订单查询",
            functionName="query_order",
            definition={
                "schemaVersion": 2,
                "type": "http",
                "config": {
                    "url": "https://example.com/orders",
                    "headers": {},
                    "body": {},
                },
            },
            secrets={"dynamic_variables": {"api_token": "secret"}},
        )

        with self.assertRaisesRegex(HTTPException, "secret__"):
            _payload(body)

    def test_http_resource_keeps_secret_placeholders_in_headers(self):
        body = ToolUpsert(
            name="订单查询",
            functionName="query_order",
            definition={
                "schemaVersion": 2,
                "type": "http",
                "config": {
                    "url": "https://example.com/orders",
                    "headers": {"Authorization": "Bearer {{secret__token}}"},
                    "body": {},
                },
            },
            secrets={"dynamic_variables": {"secret__token": "secret"}},
        )

        payload = _payload(body)

        self.assertEqual(
            payload["definition"]["config"]["headers"]["Authorization"],
            "Bearer {{secret__token}}",
        )

    def test_migration_preserves_inherited_update_state_scope(self):
        migration = importlib.import_module(
            "migrations.versions.20260813_0021_organize_tool_resources"
        )
        settings = {
            "runtimeMode": "pipeline",
            "globalPrompt": "全局规则",
            "defaultLlmResourceId": "llm_default",
            "toolBindings": [{"toolId": "tool_lookup"}],
            "builtInTools": {
                "updateState": {
                    "updates": [
                        {
                            "variableName": "name",
                            "expression": {"type": "constant", "value": "all"},
                        },
                        {
                            "variableName": "city",
                            "expression": {"type": "constant", "value": "all"},
                        },
                    ]
                }
            },
        }
        data = {"prompt": "节点任务", "inheritGlobalConfig": True}

        migration._freeze_inherited_agent_for_update_state(
            data,
            settings,
            {"name": {"type": "string"}, "city": {"type": "string"}},
            ["city"],
        )

        self.assertFalse(data["inheritGlobalConfig"])
        self.assertEqual(data["llmResourceId"], "llm_default")
        self.assertEqual(data["prompt"], "全局规则\n\n节点任务")
        self.assertEqual(
            [
                rule["variableName"]
                for rule in data["builtInTools"]["updateState"]["updates"]
            ],
            ["city"],
        )
        self.assertEqual(len(settings["builtInTools"]["updateState"]["updates"]), 2)

    def test_realtime_migration_does_not_create_model_overrides(self):
        migration = importlib.import_module(
            "migrations.versions.20260813_0021_organize_tool_resources"
        )
        settings = {
            "runtimeMode": "realtime",
            "defaultLlmResourceId": "unused",
            "enableInterrupt": True,
            "toolBindings": [],
            "builtInTools": {"updateState": {"updates": []}},
        }
        data = {"inheritGlobalConfig": True}

        migration._freeze_inherited_agent_for_update_state(
            data,
            settings,
            {"name": {"type": "string"}},
            ["name"],
        )

        self.assertNotIn("llmResourceId", data)
        self.assertNotIn("enableInterrupt", data)


class ToolBindingTests(unittest.IsolatedAsyncioTestCase):
    async def test_only_model_sources_are_visible_to_llm(self):
        tool = RuntimeTool(
            id="tool_order",
            name="查询订单",
            function_name="query_order",
            type="http",
            definition={
                "config": {
                    "allow_interruptions": True,
                    "execution_mode": "immediate",
                    "parameters": [
                        {"name": "order_id", "type": "string", "required": True},
                        {"name": "tenant", "type": "string", "required": True},
                    ]
                }
            },
        )
        binding = RuntimeToolBinding(
            toolId=tool.id,
            argumentBindings={
                "order_id": {"source": "model"},
                "tenant": {"source": "variable", "variableName": "tenant"},
            },
        )

        properties, required = ToolExecutor.schema_parts(tool, binding)

        self.assertEqual(set(properties), {"order_id"})
        self.assertEqual(required, ["order_id"])
        executor = ToolExecutor(DynamicVariableStore({"tenant": "school-a"}))
        self.assertEqual(
            executor.resolve_arguments(tool, binding, {"order_id": "A-1"}),
            {"order_id": "A-1", "tenant": "school-a"},
        )

    async def test_model_cannot_override_non_model_binding(self):
        tool = RuntimeTool(
            id="tool_order",
            name="查询订单",
            function_name="query_order",
            type="http",
            definition={
                "config": {
                    "allow_interruptions": True,
                    "execution_mode": "immediate",
                    "parameters": [
                        {"name": "tenant", "type": "string", "required": True}
                    ]
                }
            },
        )
        binding = RuntimeToolBinding(
            toolId=tool.id,
            argumentBindings={
                "tenant": {"source": "variable", "variableName": "tenant"}
            },
        )
        executor = ToolExecutor(DynamicVariableStore({"tenant": "school-a"}))

        with self.assertRaisesRegex(ToolExecutionError, "未授权"):
            executor.resolve_arguments(tool, binding, {"tenant": "attacker"})


class UpdateStateExpressionTests(unittest.TestCase):
    def test_expression_tree_reads_one_snapshot(self):
        definitions = {
            "visits": {"type": "number"},
            "vip": {"type": "boolean"},
        }
        rules = [
            {
                "variableName": "visits",
                "expression": {
                    "type": "add",
                    "operands": [
                        {"type": "variable", "variableName": "visits"},
                        {"type": "constant", "value": 1},
                    ],
                },
            },
            {
                "variableName": "vip",
                "expression": {
                    "type": "gte",
                    "operands": [
                        {"type": "variable", "variableName": "visits"},
                        {"type": "constant", "value": 3},
                    ],
                },
            },
        ]

        self.assertEqual(
            validate_update_rules(rules, definitions, allow_llm=False),
            {},
        )
        self.assertEqual(
            evaluate_update_rules(rules, {"visits": 3, "vip": False}),
            {"visits": 4, "vip": True},
        )

    def test_deterministic_node_rejects_llm_leaf(self):
        with self.assertRaisesRegex(StateExpressionError, "不能使用 LLM"):
            validate_update_rules(
                [
                    {
                        "variableName": "name",
                        "expression": {
                            "type": "llm",
                            "parameterName": "name",
                            "valueType": "string",
                            "description": "姓名",
                        },
                    }
                ],
                {"name": {"type": "string"}},
                allow_llm=False,
            )

    def test_llm_parameter_cannot_drive_multiple_update_targets(self):
        expression = {
            "type": "llm",
            "parameterName": "shared",
            "valueType": "string",
            "description": "提取文本",
        }
        with self.assertRaisesRegex(StateExpressionError, "多个变量"):
            validate_update_rules(
                [
                    {"variableName": "first", "expression": expression},
                    {"variableName": "second", "expression": expression},
                ],
                {
                    "first": {"type": "string"},
                    "second": {"type": "string"},
                },
                allow_llm=True,
            )


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