agency-swarm-cn-git/tests/test_agent_modules/send_message/test_extra_params.py

import json

import pytest
from agents import ModelSettings, RunContextWrapper
from pydantic import BaseModel, Field

from agency_swarm import Agent
from agency_swarm.context import MasterContext
from agency_swarm.tools.send_message import SendMessage
from agency_swarm.utils.thread import ThreadManager
from tests.deterministic_model import DeterministicModel


def _make_stub_agent(name: str, response: str = "ack") -> Agent:
    return Agent(
        name=name,
        instructions="stub",
        model=DeterministicModel(default_response=response),
        model_settings=ModelSettings(temperature=0.0),
    )


class NumericContextParams(BaseModel):
    count: int = Field(description="Count")
    summary: str = Field(description="Summary")


class InlineContext(BaseModel):
    value: int = Field(description="Value")


class SendMessageWithContext(SendMessage):
    extra_params_model = NumericContextParams


class BadExtra(BaseModel):
    foo: str

    @classmethod
    def model_json_schema(cls, *args, **kwargs):  # type: ignore[override]
        raise ValueError("boom")


class SendMessageBad(SendMessage):
    ExtraParams = BadExtra


# --- New pattern: inline field declarations ---


class SendMessageInline(SendMessage):
    """Inline extra params without nested class."""

    tool_name = "send_message_inline"

    priority: str = Field(description="Task priority level")
    context_summary: str = Field(description="Summary of context")


def _wrapper_with_recipient(recipient: Agent) -> RunContextWrapper[MasterContext]:
    ctx = MasterContext(
        thread_manager=ThreadManager(),
        agents={"B": recipient},
        user_context={},
        agent_runtime_state={},
        shared_instructions=None,
    )
    return RunContextWrapper(context=ctx)


@pytest.mark.asyncio
async def test_send_message_extra_params_schema_validation_and_success() -> None:
    """Extra params should be merged into schema, validate input, and still allow successful sends."""
    sender = _make_stub_agent("Sender")
    recipient = _make_stub_agent("Recipient")
    tool = SendMessageWithContext(sender, recipients={"B": recipient})

    properties = tool.params_json_schema["properties"]
    required = tool.params_json_schema["required"]
    assert properties["count"]["type"] == "integer"
    assert properties["summary"]["type"] == "string"
    assert "count" in required and "summary" in required

    valid_args = json.dumps(
        {
            "recipient_agent": "B",
            "message": "msg",
            "additional_instructions": "",
            "count": 1,
            "summary": "ok",
        }
    )
    invalid_args = json.dumps(
        {
            "recipient_agent": "B",
            "message": "msg",
            "additional_instructions": "",
            "count": "bad",
            "summary": "ok",
        }
    )

    wrapper = _wrapper_with_recipient(recipient)
    assert await tool.on_invoke_tool(wrapper, valid_args) == "ack"
    invalid_result = await tool.on_invoke_tool(wrapper, invalid_args)
    assert isinstance(invalid_result, str) and invalid_result.startswith("Error: Invalid extra parameters")


@pytest.mark.asyncio
async def test_send_message_bad_extra_params_model_falls_back_without_validation() -> None:
    """Schema generation failures in ExtraParams should keep tool usable without extra field validation."""
    sender = _make_stub_agent("Sender")
    recipient = _make_stub_agent("Recipient")
    tool = SendMessageBad(sender, recipients={"B": recipient})
    assert "foo" not in tool.params_json_schema["properties"]

    wrapper = _wrapper_with_recipient(recipient)
    base_args = {
        "recipient_agent": "B",
        "message": "m",
        "additional_instructions": "",
    }
    result_no_extra = await tool.on_invoke_tool(wrapper, json.dumps(base_args))
    result_unknown_extra = await tool.on_invoke_tool(wrapper, json.dumps({**base_args, "foo": "x"}))

    assert isinstance(result_no_extra, str) and not result_no_extra.startswith("Error: Invalid extra parameters")
    assert isinstance(result_unknown_extra, str) and not result_unknown_extra.startswith(
        "Error: Invalid extra parameters"
    )


def test_inline_fields_merged_into_schema() -> None:
    """Inline field declarations should be auto-discovered and merged into the schema."""
    sender = _make_stub_agent("Sender")
    recipient = _make_stub_agent("Recipient")
    tool = SendMessageInline(sender, recipients={"B": recipient})

    properties = tool.params_json_schema["properties"]
    required = tool.params_json_schema["required"]

    assert "priority" in properties
    assert properties["priority"]["type"] == "string"
    assert "context_summary" in properties
    assert properties["context_summary"]["type"] == "string"
    assert "priority" in required
    assert "context_summary" in required

    # Verify tool_name was applied
    assert tool.name == "send_message_inline"


def test_tool_name_class_attribute() -> None:
    """tool_name class attribute should set the tool name without __init__ override."""

    class MySendMessage(SendMessage):
        tool_name = "send_message_custom"

    sender = _make_stub_agent("Sender")
    tool = MySendMessage(sender)
    assert tool.name == "send_message_custom"


def test_tool_name_explicit_name_takes_precedence() -> None:
    """Explicit name= kwarg should override tool_name class attribute."""

    class MySendMessage(SendMessage):
        tool_name = "send_message_custom"

    sender = _make_stub_agent("Sender")
    tool = MySendMessage(sender, name="send_message_override")
    assert tool.name == "send_message_override"


def test_inline_fields_not_picked_up_when_extra_params_exists() -> None:
    """ExtraParams nested class should take priority over inline fields."""

    class MixedSendMessage(SendMessage):
        class ExtraParams(BaseModel):
            from_nested: str = Field(description="From nested")

        # This should be ignored because ExtraParams is present
        inline_field: str = Field(description="Should be ignored")

    sender = _make_stub_agent("Sender")
    tool = MixedSendMessage(sender)
    properties = tool.params_json_schema["properties"]
    assert "from_nested" in properties
    assert "inline_field" not in properties


def test_bare_annotations_are_not_treated_as_inline_fields() -> None:
    """Bare subclass annotations should remain internal typing hints."""

    class BareAnnotationSendMessage(SendMessage):
        internal_note: str

    sender = _make_stub_agent("Sender")
    tool = BareAnnotationSendMessage(sender)

    assert "internal_note" not in tool.params_json_schema["properties"]


def test_inline_fields_resolve_forward_annotations() -> None:
    """Inline fields should resolve future-style annotations from the subclass module."""

    class FutureStyleAnnotation(SendMessage):
        context: "InlineContext" = Field(description="Structured context")

    sender = _make_stub_agent("Sender")
    tool = FutureStyleAnnotation(sender)

    assert "context" in tool.params_json_schema["properties"]
    assert tool._extra_params_model is not None
    assert tool._extra_params_model.model_fields["context"].annotation is InlineContext


@pytest.mark.asyncio
async def test_inline_fields_validation() -> None:
    """Inline fields should be validated like ExtraParams fields."""
    sender = _make_stub_agent("Sender")
    recipient = _make_stub_agent("Recipient")

    class StrictInline(SendMessage):
        count: int = Field(description="Must be integer")

    tool = StrictInline(sender, recipients={"Recipient": recipient})
    wrapper = _wrapper_with_recipient(recipient)

    invalid_args = json.dumps(
        {
            "recipient_agent": "Recipient",
            "message": "test",
            "additional_instructions": "",
            "count": "not_a_number",
        }
    )
    result = await tool.on_invoke_tool(wrapper, invalid_args)
    assert isinstance(result, str) and result.startswith("Error: Invalid extra parameters")