temp/tests/features/steps/composite_agent_coverage_steps.py

"""Step definitions for comprehensive composite agent coverage tests."""

import asyncio
import tempfile
import threading
from unittest.mock import AsyncMock, MagicMock

from behave import given, then, when

from cleveragents.agents.base import Agent
from cleveragents.agents.composite import CompositeAgent
from cleveragents.core.exceptions import ConfigurationError, ExecutionError
from cleveragents.langgraph.bridge import RxPyLangGraphBridge
from cleveragents.reactive.stream_router import ReactiveStreamRouter
from cleveragents.templates.renderer import TemplateRenderer


@given("I have a composite agent test environment")
def step_composite_agent_test_environment(context):
    """Set up test environment for composite agent testing."""
    # Initialize test data on context
    context.test_data = {
        "temp_dir": tempfile.mkdtemp(),
        "mock_template_renderer": MagicMock(spec=TemplateRenderer),
        "mock_stream_router": MagicMock(spec=ReactiveStreamRouter),
        "mock_langgraph_bridge": MagicMock(spec=RxPyLangGraphBridge),
        "mock_agents": {},
        "mock_graphs": {},
        "mock_streams": {},
        "test_results": {},
        "raised_exception": None,
        "config": {},
        "composite_agent": None,
    }

    # Initialize direct context attributes for convenience
    context.mock_agents = {}
    context.mock_graphs = {}
    context.mock_streams = {}

    # Set up async event loop for testing
    try:
        loop = asyncio.get_event_loop()
    except RuntimeError:
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
    context.test_data["loop"] = loop


@given("I have a basic composite agent configuration")
def step_basic_composite_agent_config(context):
    """Create basic composite agent configuration."""
    context.test_data["config"] = {"components": {}, "routing": {}, "expose_params": {}}


@given("I have a legacy strategy-based configuration")
def step_legacy_strategy_config(context):
    """Create legacy strategy-based configuration."""
    context.test_data["config"] = {
        "strategy": "parallel",
        "agents": ["agent1", "agent2"],
    }


@given("I have a composite agent")
def step_have_composite_agent(context):
    """Create a basic composite agent."""
    config = {"components": {}, "routing": {}, "expose_params": {}}
    context.composite_agent = CompositeAgent(
        name="test_composite",
        config=config,
        template_renderer=context.test_data["mock_template_renderer"],
        stream_router=context.test_data["mock_stream_router"],
        langgraph_bridge=context.test_data["mock_langgraph_bridge"],
    )


@given("I have mock child agents")
def step_mock_child_agents(context):
    """Create mock child agents."""
    for i in range(3):
        mock_agent = MagicMock(spec=Agent)
        mock_agent.name = f"child_agent_{i}"
        mock_agent.process = AsyncMock(return_value=f"Response from agent {i}")
        mock_agent.get_capabilities = MagicMock(return_value=[f"capability_{i}"])
        context.mock_agents[f"child_agent_{i}"] = mock_agent


@given("I have mock LangGraph instances")
def step_mock_langgraph_instances(context):
    """Create mock LangGraph instances."""
    for i in range(2):
        mock_graph = MagicMock()
        mock_graph.execute = AsyncMock(return_value=f"Graph result {i}")
        context.mock_graphs[f"graph_{i}"] = mock_graph


@given("I have mock stream configurations")
def step_mock_stream_configurations(context):
    """Create mock stream configurations."""
    for i in range(2):
        mock_stream = {
            "type": "cold",
            "operators": [{"type": "map", "params": {"function": f"stream_func_{i}"}}],
        }
        context.mock_streams[f"stream_{i}"] = mock_stream


@when("I create a composite agent")
def step_create_composite_agent(context):
    """Create a composite agent with the configured settings."""
    try:
        context.composite_agent = CompositeAgent(
            name="test_composite",
            config=context.test_data["config"],
            template_renderer=context.test_data["mock_template_renderer"],
            stream_router=context.test_data["mock_stream_router"],
            langgraph_bridge=context.test_data["mock_langgraph_bridge"],
        )
        context.creation_successful = True
    except Exception as e:
        context.raised_exception = e
        context.creation_successful = False


@when("I attempt to create a composite agent with legacy config")
def step_attempt_create_with_legacy_config(context):
    """Attempt to create composite agent with legacy configuration."""
    try:
        context.composite_agent = CompositeAgent(
            name="test_composite",
            config=context.test_data["config"],
            template_renderer=context.test_data["mock_template_renderer"],
        )
        context.creation_successful = True
    except Exception as e:
        context.raised_exception = e
        context.creation_successful = False


@when("I add agents to the composite agent")
def step_add_agents_to_composite(context):
    """Add mock agents to the composite agent."""
    for name, agent in context.mock_agents.items():
        context.composite_agent.add_agent(name, agent)


@when("I add graphs to the composite agent")
def step_add_graphs_to_composite(context):
    """Add mock graphs to the composite agent."""
    for name, graph in context.mock_graphs.items():
        context.composite_agent.add_graph(name, graph)


@when("I add streams to the composite agent")
def step_add_streams_to_composite(context):
    """Add mock streams to the composite agent."""
    for name, stream in context.mock_streams.items():
        context.composite_agent.add_stream(name, stream)


@when("I set various parameters on the agent")
def step_set_parameters_on_agent(context):
    """Set various parameters on the composite agent."""
    context.test_params = {
        "temperature": 0.7,
        "max_tokens": 1000,
        "custom_param": "test_value",
    }
    for param, value in context.test_params.items():
        context.composite_agent.set_param(param, value)


@when("I call process_message on the composite agent")
def step_call_process_message(context):
    """Call process_message method on composite agent."""
    # Call process_message using the child agent that was already added
    context.process_message_result = context.test_data["loop"].run_until_complete(
        context.composite_agent.process_message("test message", {"context": "test"})
    )


@when("I process a message through the composite agent")
def step_process_message_through_composite(context):
    """Process a message through the composite agent."""
    try:
        context.process_result = context.test_data["loop"].run_until_complete(
            context.composite_agent.process("test message", {"context": "test"})
        )
        context.process_successful = True
    except Exception as e:
        context.raised_exception = e
        context.process_successful = False


@when("I get the capabilities of the composite agent")
def step_get_capabilities(context):
    """Get capabilities from the composite agent."""
    context.capabilities = context.composite_agent.get_capabilities()


@then("the composite agent should be initialized correctly")
def step_composite_agent_initialized_correctly(context):
    """Verify composite agent is initialized correctly."""
    assert context.creation_successful
    assert context.composite_agent.name == "test_composite"
    assert hasattr(context.composite_agent, "components")
    assert hasattr(context.composite_agent, "routing")
    assert hasattr(context.composite_agent, "expose_params")


@then("the agent should have empty component collections")
def step_agent_has_empty_collections(context):
    """Verify agent has empty component collections."""
    assert len(context.composite_agent.agents) == 0
    assert len(context.composite_agent.graphs) == 0
    assert len(context.composite_agent.streams) == 0


@then("the agent should have default routing configuration")
def step_agent_has_default_routing(context):
    """Verify agent has default routing configuration."""
    assert context.composite_agent.routing == {}
    assert context.composite_agent.expose_params == {}


@then("a ConfigurationError should be raised for composite agent")
def step_configuration_error_raised_composite(context):
    """Verify a ConfigurationError was raised."""
    assert (
        not context.creation_successful if hasattr(context, "creation_successful") else not context.process_successful
    )
    assert isinstance(context.raised_exception, ConfigurationError)


@then("an ExecutionError should be raised for composite agent")
def step_execution_error_raised_composite(context):
    """Verify an ExecutionError was raised."""
    assert not context.process_successful
    assert isinstance(context.raised_exception, ExecutionError)


@then("the error should mention deprecated strategy-based configuration")
def step_error_mentions_deprecated_strategy(context):
    """Verify error mentions deprecated strategy configuration."""
    assert "deprecated strategy-based configuration" in str(context.raised_exception)


@then("the agents should be stored in the agents collection")
def step_agents_stored_in_collection(context):
    """Verify agents are stored in the agents collection."""
    for name in context.mock_agents.keys():
        assert name in context.composite_agent.agents
        assert context.composite_agent.agents[name] == context.mock_agents[name]


@then("the agents should be added to components configuration")
def step_agents_added_to_components(context):
    """Verify agents are added to components configuration."""
    assert "agents" in context.composite_agent.components
    for name in context.mock_agents.keys():
        assert name in context.composite_agent.components["agents"]


@then("the agents collection should contain the correct agents")
def step_agents_collection_correct(context):
    """Verify agents collection contains correct agents."""
    assert len(context.composite_agent.agents) == len(context.mock_agents)
    for name, agent in context.mock_agents.items():
        assert context.composite_agent.agents[name] == agent


@then("the graphs should be stored in the graphs collection")
def step_graphs_stored_in_collection(context):
    """Verify graphs are stored in the graphs collection."""
    for name in context.mock_graphs.keys():
        assert name in context.composite_agent.graphs
        assert context.composite_agent.graphs[name] == context.mock_graphs[name]


@then("the graphs should be added to components configuration")
def step_graphs_added_to_components(context):
    """Verify graphs are added to components configuration."""
    assert "graphs" in context.composite_agent.components
    for name in context.mock_graphs.keys():
        assert name in context.composite_agent.components["graphs"]


@then("the graphs collection should contain the correct graphs")
def step_graphs_collection_correct(context):
    """Verify graphs collection contains correct graphs."""
    assert len(context.composite_agent.graphs) == len(context.mock_graphs)
    for name, graph in context.mock_graphs.items():
        assert context.composite_agent.graphs[name] == graph


@then("the streams should be stored in the streams collection")
def step_streams_stored_in_collection(context):
    """Verify streams are stored in the streams collection."""
    for name in context.mock_streams.keys():
        assert name in context.composite_agent.streams
        assert context.composite_agent.streams[name] == context.mock_streams[name]


@then("the streams should be added to components configuration")
def step_streams_added_to_components(context):
    """Verify streams are added to components configuration."""
    assert "streams" in context.composite_agent.components
    for name in context.mock_streams.keys():
        assert name in context.composite_agent.components["streams"]


@then("the streams collection should contain the correct streams")
def step_streams_collection_correct(context):
    """Verify streams collection contains correct streams."""
    assert len(context.composite_agent.streams) == len(context.mock_streams)
    for name, stream in context.mock_streams.items():
        assert context.composite_agent.streams[name] == stream


@then("the parameters should be stored in expose_params")
def step_parameters_stored_in_expose_params(context):
    """Verify parameters are stored in expose_params."""
    for param, value in context.test_params.items():
        assert context.composite_agent.expose_params[param] == value


@then("the parameters should be available for propagation")
def step_parameters_available_for_propagation(context):
    """Verify parameters are available for propagation."""
    # Parameters should be accessible through expose_params
    assert len(context.composite_agent.expose_params) == len(context.test_params)


@then("it should delegate to the process method")
def step_should_delegate_to_process(context):
    """Verify process_message delegates to process method."""
    assert context.process_message_result == "Child agent response"


@then("return the same result as process")
def step_return_same_result_as_process(context):
    """Verify same result is returned."""
    # The result should match what the mock agent returns
    assert "response" in context.process_message_result.lower()


# Additional step definitions for specific routing scenarios


@given("I have a composite agent with agents but no routing")
def step_composite_with_agents_no_routing(context):
    """Create composite agent with agents but no routing."""
    config = {"components": {}, "routing": {}}
    context.composite_agent = CompositeAgent(
        name="test_composite",
        config=config,
        template_renderer=context.test_data["mock_template_renderer"],
    )

    # Add a mock agent
    mock_agent = MagicMock(spec=Agent)
    mock_agent.process = AsyncMock(return_value="Agent response")
    context.composite_agent.add_agent("first_agent", mock_agent)


@given("I have a composite agent with graphs but no routing")
def step_composite_with_graphs_no_routing(context):
    """Create composite agent with graphs but no routing."""
    config = {"components": {}, "routing": {}}
    context.composite_agent = CompositeAgent(
        name="test_composite",
        config=config,
        template_renderer=context.test_data["mock_template_renderer"],
        langgraph_bridge=context.test_data["mock_langgraph_bridge"],
    )

    # Add a mock graph
    mock_graph = MagicMock()
    mock_graph.execute = AsyncMock(return_value="Graph response")
    context.composite_agent.add_graph("first_graph", mock_graph)


@given("I have a composite agent with streams but no routing")
def step_composite_with_streams_no_routing(context):
    """Create composite agent with streams but no routing."""
    config = {"components": {}, "routing": {}}
    context.composite_agent = CompositeAgent(
        name="test_composite",
        config=config,
        template_renderer=context.test_data["mock_template_renderer"],
        stream_router=context.test_data["mock_stream_router"],
    )

    # Add a mock stream with simplified behavior
    mock_stream = {"type": "cold"}
    context.composite_agent.add_stream("first_stream", mock_stream)

    # Mock the stream router with proper observable behavior
    mock_observable = MagicMock()
    mock_subscription = MagicMock()
    mock_observable.subscribe.return_value = mock_subscription
    context.test_data["mock_stream_router"].observables = {"__output__": mock_observable}

    # Mock send_message to trigger the observer callback
    def mock_send_message(stream_name, message, context_data):
        # Simulate message processing and callback
        def delayed_callback():
            # Find the observer callback from the subscribe call
            if mock_observable.subscribe.called:
                observer = mock_observable.subscribe.call_args[0][0]
                if hasattr(observer, "on_next"):
                    observer.on_next("Stream processed immediately")

        # Call callback with slight delay to simulate async processing
        timer = threading.Timer(0.01, delayed_callback)
        timer.start()

    context.test_data["mock_stream_router"].send_message = mock_send_message


@given("I have a composite agent with no components")
def step_composite_with_no_components(context):
    """Create composite agent with no components."""
    config = {"components": {}, "routing": {}}
    context.composite_agent = CompositeAgent(
        name="test_composite",
        config=config,
        template_renderer=context.test_data["mock_template_renderer"],
    )


@given("I have a composite agent with explicit agent routing")
def step_composite_with_agent_routing(context):
    """Create composite agent with explicit agent routing."""
    config = {
        "components": {},
        "routing": {
            "input": {"type": "agent", "name": "target_agent"},
            "output": {"name": "output_stream"},
        },
    }
    context.composite_agent = CompositeAgent(
        name="test_composite",
        config=config,
        template_renderer=context.test_data["mock_template_renderer"],
    )

    # Add the target agent
    mock_agent = MagicMock(spec=Agent)
    mock_agent.process = AsyncMock(return_value="Routed agent response")
    context.composite_agent.add_agent("target_agent", mock_agent)


@given("I have a composite agent with explicit graph routing")
def step_composite_with_graph_routing(context):
    """Create composite agent with explicit graph routing."""
    config = {
        "components": {},
        "routing": {"input": {"type": "graph", "name": "target_graph"}},
    }
    context.composite_agent = CompositeAgent(
        name="test_composite",
        config=config,
        template_renderer=context.test_data["mock_template_renderer"],
        langgraph_bridge=context.test_data["mock_langgraph_bridge"],
    )

    # Add the target graph
    mock_graph = MagicMock()
    mock_result = MagicMock()
    mock_result.messages = [{"content": "Graph routed response"}]
    mock_graph.execute = AsyncMock(return_value=mock_result)
    context.composite_agent.add_graph("target_graph", mock_graph)


@given("I have a composite agent with explicit stream routing")
def step_composite_with_stream_routing(context):
    """Create composite agent with explicit stream routing."""
    config = {
        "components": {},
        "routing": {
            "input": {"type": "stream", "name": "target_stream"},
            "output": {"name": "custom_output"},
        },
    }
    context.composite_agent = CompositeAgent(
        name="test_composite",
        config=config,
        template_renderer=context.test_data["mock_template_renderer"],
        stream_router=context.test_data["mock_stream_router"],
    )

    # Mock the stream router behavior
    mock_observable = MagicMock()
    mock_subscription = MagicMock()
    mock_observable.subscribe.return_value = mock_subscription
    context.test_data["mock_stream_router"].observables = {
        "custom_output": mock_observable,
        "__output__": mock_observable,
    }

    # Mock send_message to trigger the observer callback
    def mock_send_message(stream_name, message, context_data):
        # Simulate message processing and callback
        def delayed_callback():
            # Find the observer callback from the subscribe call
            if mock_observable.subscribe.called:
                observer = mock_observable.subscribe.call_args[0][0]
                if hasattr(observer, "on_next"):
                    observer.on_next("Stream routing response")

        # Call callback with slight delay to simulate async processing
        timer = threading.Timer(0.01, delayed_callback)
        timer.start()

    context.test_data["mock_stream_router"].send_message = mock_send_message

    # Add the target stream
    context.composite_agent.add_stream("target_stream", {"type": "cold"})


@given("I have a composite agent with unknown routing type")
def step_composite_with_unknown_routing(context):
    """Create composite agent with unknown routing type."""
    config = {
        "components": {},
        "routing": {"input": {"type": "unknown_type", "name": "target"}},
    }
    context.composite_agent = CompositeAgent(
        name="test_composite",
        config=config,
        template_renderer=context.test_data["mock_template_renderer"],
    )


@given("I have a composite agent with agent routing to non-existent agent")
def step_composite_with_missing_agent_routing(context):
    """Create composite agent with routing to non-existent agent."""
    config = {
        "components": {},
        "routing": {"input": {"type": "agent", "name": "missing_agent"}},
    }
    context.composite_agent = CompositeAgent(
        name="test_composite",
        config=config,
        template_renderer=context.test_data["mock_template_renderer"],
    )


@given("I have a composite agent with graph routing but no bridge")
def step_composite_with_graph_no_bridge(context):
    """Create composite agent with graph routing but no bridge."""
    config = {
        "components": {},
        "routing": {"input": {"type": "graph", "name": "some_graph"}},
    }
    context.composite_agent = CompositeAgent(
        name="test_composite",
        config=config,
        template_renderer=context.test_data["mock_template_renderer"],
        langgraph_bridge=None,
    )


@given("I have a composite agent with graph routing to non-existent graph")
def step_composite_with_missing_graph_routing(context):
    """Create composite agent with routing to non-existent graph."""
    config = {
        "components": {},
        "routing": {"input": {"type": "graph", "name": "missing_graph"}},
    }
    context.composite_agent = CompositeAgent(
        name="test_composite",
        config=config,
        template_renderer=context.test_data["mock_template_renderer"],
        langgraph_bridge=context.test_data["mock_langgraph_bridge"],
    )

    # Mock bridge to return None for missing graph
    context.test_data["mock_langgraph_bridge"].get_graph.return_value = None


@given("I have a composite agent with LangGraph bridge")
def step_composite_with_langgraph_bridge(context):
    """Create composite agent with LangGraph bridge."""
    config = {
        "components": {},
        "routing": {"input": {"type": "graph", "name": "bridge_graph"}},
    }
    context.composite_agent = CompositeAgent(
        name="test_composite",
        config=config,
        template_renderer=context.test_data["mock_template_renderer"],
        langgraph_bridge=context.test_data["mock_langgraph_bridge"],
    )


@given("the bridge has a graph not in local collection")
def step_bridge_has_external_graph(context):
    """Set up bridge to have graph not in local collection."""
    mock_graph = MagicMock()
    mock_graph.execute = AsyncMock(return_value="Bridge graph response")
    context.test_data["mock_langgraph_bridge"].get_graph.return_value = mock_graph


@given("I have a composite agent with LangGraph that returns messages")
def step_composite_with_message_returning_graph(context):
    """Create composite agent with graph that returns messages."""
    config = {
        "components": {},
        "routing": {"input": {"type": "graph", "name": "message_graph"}},
    }
    context.composite_agent = CompositeAgent(
        name="test_composite",
        config=config,
        template_renderer=context.test_data["mock_template_renderer"],
        langgraph_bridge=context.test_data["mock_langgraph_bridge"],
    )

    # Create mock graph with messages result
    mock_graph = MagicMock()
    mock_result = MagicMock()
    mock_result.messages = [
        {"content": "First message"},
        {"content": "Last message content"},
    ]
    mock_graph.execute = AsyncMock(return_value=mock_result)
    context.composite_agent.add_graph("message_graph", mock_graph)


@given("I have a composite agent with LangGraph that returns non-message result")
def step_composite_with_non_message_graph(context):
    """Create composite agent with graph that returns non-message result."""
    config = {
        "components": {},
        "routing": {"input": {"type": "graph", "name": "non_message_graph"}},
    }
    context.composite_agent = CompositeAgent(
        name="test_composite",
        config=config,
        template_renderer=context.test_data["mock_template_renderer"],
        langgraph_bridge=context.test_data["mock_langgraph_bridge"],
    )

    # Create mock graph with non-message result
    mock_graph = MagicMock()
    mock_graph.execute = AsyncMock(return_value={"result": "graph_output"})
    context.composite_agent.add_graph("non_message_graph", mock_graph)


@given("I have a composite agent with stream routing but no router")
def step_composite_with_stream_no_router(context):
    """Create composite agent with stream routing but no router."""
    config = {
        "components": {},
        "routing": {"input": {"type": "stream", "name": "some_stream"}},
    }
    context.composite_agent = CompositeAgent(
        name="test_composite",
        config=config,
        template_renderer=context.test_data["mock_template_renderer"],
        stream_router=None,
    )


@given("I have a composite agent with stream routing and custom output")
def step_composite_with_custom_stream_output(context):
    """Create composite agent with custom stream output."""
    config = {
        "components": {},
        "routing": {
            "input": {"type": "stream", "name": "input_stream"},
            "output": {"name": "custom_output_stream"},
        },
    }
    context.composite_agent = CompositeAgent(
        name="test_composite",
        config=config,
        template_renderer=context.test_data["mock_template_renderer"],
        stream_router=context.test_data["mock_stream_router"],
    )

    # Set up mock stream router with custom output
    mock_observable = MagicMock()
    mock_subscription = MagicMock()
    mock_observable.subscribe.return_value = mock_subscription

    # Mock the custom output stream
    context.test_data["mock_stream_router"].observables = {
        "custom_output_stream": mock_observable,
        "__output__": mock_observable,
    }

    # Mock send_message to trigger the observer callback
    def mock_send_message(stream_name, message, context_data):
        # Simulate message processing and callback
        def delayed_callback():
            # Find the observer callback from the subscribe call
            if mock_observable.subscribe.called:
                observer = mock_observable.subscribe.call_args[0][0]
                if hasattr(observer, "on_next"):
                    observer.on_next("Custom stream response")

        # Call callback with slight delay to simulate async processing
        timer = threading.Timer(0.01, delayed_callback)
        timer.start()

    context.test_data["mock_stream_router"].send_message = mock_send_message


@given("I have a composite agent with stream routing and default output")
def step_composite_with_default_stream_output(context):
    """Create composite agent with default stream output."""
    config = {
        "components": {},
        "routing": {"input": {"type": "stream", "name": "input_stream"}},
    }
    context.composite_agent = CompositeAgent(
        name="test_composite",
        config=config,
        template_renderer=context.test_data["mock_template_renderer"],
        stream_router=context.test_data["mock_stream_router"],
    )

    # Set up mock stream router with default output
    mock_observable = MagicMock()
    mock_subscription = MagicMock()
    mock_observable.subscribe.return_value = mock_subscription

    context.test_data["mock_stream_router"].observables = {"__output__": mock_observable}

    # Mock send_message to trigger the observer callback
    def mock_send_message(stream_name, message, context_data):
        def delayed_callback():
            if mock_observable.subscribe.called:
                observer = mock_observable.subscribe.call_args[0][0]
                if hasattr(observer, "on_next"):
                    observer.on_next("Default stream response")

        timer = threading.Timer(0.01, delayed_callback)
        timer.start()

    context.test_data["mock_stream_router"].send_message = mock_send_message


@given("I have a composite agent with stream that returns message objects")
def step_composite_with_message_object_stream(context):
    """Create composite agent with stream that returns message objects."""
    config = {
        "components": {},
        "routing": {"input": {"type": "stream", "name": "message_stream"}},
    }
    context.composite_agent = CompositeAgent(
        name="test_composite",
        config=config,
        template_renderer=context.test_data["mock_template_renderer"],
        stream_router=context.test_data["mock_stream_router"],
    )

    # Set up mock stream router
    mock_observable = MagicMock()
    mock_subscription = MagicMock()
    mock_observable.subscribe.return_value = mock_subscription

    context.test_data["mock_stream_router"].observables = {"__output__": mock_observable}

    # Mock send_message to return message object
    def mock_send_message(stream_name, message, context_data):
        def delayed_callback():
            if mock_observable.subscribe.called:
                observer = mock_observable.subscribe.call_args[0][0]
                if hasattr(observer, "on_next"):
                    # Create mock message object with content attribute
                    mock_message = MagicMock()
                    mock_message.content = "Message object content"
                    observer.on_next(mock_message)

        timer = threading.Timer(0.01, delayed_callback)
        timer.start()

    context.test_data["mock_stream_router"].send_message = mock_send_message


@given("I have a composite agent with slow stream processing")
def step_composite_with_slow_stream(context):
    """Create composite agent with slow stream processing."""
    config = {
        "components": {},
        "routing": {"input": {"type": "stream", "name": "slow_stream"}},
    }
    context.composite_agent = CompositeAgent(
        name="test_composite",
        config=config,
        template_renderer=context.test_data["mock_template_renderer"],
        stream_router=context.test_data["mock_stream_router"],
    )

    # Set up mock stream router that processes slowly
    mock_observable = MagicMock()
    mock_subscription = MagicMock()
    mock_observable.subscribe.return_value = mock_subscription

    context.test_data["mock_stream_router"].observables = {"__output__": mock_observable}

    # Mock send_message that takes longer than timeout
    def mock_send_message(stream_name, message, context_data):
        # Don't call the callback, simulating timeout
        pass

    context.test_data["mock_stream_router"].send_message = mock_send_message


@given("I have a composite agent with agents, graphs, and streams")
def step_composite_with_all_components(context):
    """Create composite agent with all component types."""
    config = {"components": {}, "routing": {}}
    context.composite_agent = CompositeAgent(
        name="test_composite",
        config=config,
        template_renderer=context.test_data["mock_template_renderer"],
        stream_router=context.test_data["mock_stream_router"],
        langgraph_bridge=context.test_data["mock_langgraph_bridge"],
    )

    # Add agents with capabilities
    for i in range(2):
        mock_agent = MagicMock(spec=Agent)
        mock_agent.get_capabilities.return_value = [
            f"agent_capability_{i}",
            "shared_capability",
        ]
        context.composite_agent.add_agent(f"agent_{i}", mock_agent)

    # Add graphs
    context.composite_agent.add_graph("graph_1", MagicMock())

    # Add streams
    context.composite_agent.add_stream("stream_1", {"type": "cold"})


@given("I have a composite agent with legacy strategy attribute")
def step_composite_with_legacy_strategy(context):
    """Create composite agent with legacy strategy attribute."""
    config = {"components": {}, "routing": {}}
    context.composite_agent = CompositeAgent(
        name="test_composite",
        config=config,
        template_renderer=context.test_data["mock_template_renderer"],
    )

    # Manually add legacy strategy attribute
    context.composite_agent.strategy = "parallel"


@given("I have a composite agent with a mock child agent")
def step_composite_with_mock_child_agent(context):
    """Create composite agent with a single mock child agent."""
    config = {"components": {}, "routing": {}, "expose_params": {}}
    context.composite_agent = CompositeAgent(
        name="test_composite",
        config=config,
        template_renderer=context.test_data["mock_template_renderer"],
    )

    # Add a single mock child agent
    mock_agent = MagicMock(spec=Agent)
    mock_agent.process = AsyncMock(return_value="Child agent response")
    mock_agent.get_capabilities = MagicMock(return_value=["child_capability"])
    context.composite_agent.add_agent("child_agent", mock_agent)


@given("I have a composite agent with exposed parameters")
def step_composite_with_exposed_parameters(context):
    """Create composite agent with exposed parameters."""
    config = {
        "components": {},
        "routing": {},
        "expose_params": {"temperature": 0.5, "model": "gpt-4"},
    }
    context.composite_agent = CompositeAgent(
        name="test_composite",
        config=config,
        template_renderer=context.test_data["mock_template_renderer"],
    )

    # Add a mock agent for processing
    mock_agent = MagicMock(spec=Agent)
    mock_agent.process = AsyncMock(return_value="Processed with context")
    context.composite_agent.add_agent("test_agent", mock_agent)


@when("I process a message via graph routing")
def step_process_via_graph_routing(context):
    """Process message via graph routing."""
    try:
        context.process_result = asyncio.run(context.composite_agent.process("test message"))
        context.process_successful = True
    except Exception as e:
        context.raised_exception = e
        context.process_successful = False


@when("I process a message with additional context")
def step_process_with_additional_context(context):
    """Process message with additional context."""
    additional_context = {
        "user_id": "123",
        "temperature": 0.9,  # This should be overridden by exposed params
        "new_param": "additional_value",
    }

    # Capture the context that gets passed to the agent
    def capture_context(message, ctx):
        context.captured_context = ctx
        future = asyncio.Future()
        future.set_result("Processed with captured context")
        return future

    context.composite_agent.agents["test_agent"].process.side_effect = capture_context

    try:
        context.test_data["loop"].run_until_complete(
            context.composite_agent.process("test message", additional_context)
        )
    except Exception:
        pass  # We're just capturing the context


# Then step definitions for validation


@then("it should use the first available agent")
def step_should_use_first_agent(context):
    """Verify it uses the first available agent."""
    assert context.process_successful
    assert "response" in context.process_result.lower()


@then("return the processed message")
def step_return_processed_message(context):
    """Verify processed message is returned."""
    assert context.process_successful
    assert context.process_result is not None


@then("it should use the first available graph")
def step_should_use_first_graph(context):
    """Verify it uses the first available graph."""
    assert context.process_successful
    assert "response" in context.process_result.lower()


@then("return the processed message from graph")
def step_return_processed_from_graph(context):
    """Verify processed message from graph is returned."""
    assert context.process_successful
    assert context.process_result is not None


@then("it should use the first available stream")
def step_should_use_first_stream(context):
    """Verify it uses the first available stream."""
    assert context.process_successful
    assert context.process_result is not None


@then("return the processed message from stream")
def step_return_processed_from_stream(context):
    """Verify processed message from stream is returned."""
    assert context.process_successful
    assert context.process_result is not None


@then("the error should mention no components to process with")
def step_error_mentions_no_components(context):
    """Verify error mentions no components."""
    assert "no components to process with" in str(context.raised_exception)


@then("it should route to the specified agent")
def step_should_route_to_specified_agent(context):
    """Verify routing to specified agent."""
    assert context.process_successful
    assert "routed agent response" in context.process_result.lower()


@then("return the agent's processed response")
def step_return_agent_processed_response(context):
    """Verify agent's processed response is returned."""
    assert "routed" in context.process_result.lower()


@then("it should route to the specified graph")
def step_should_route_to_specified_graph(context):
    """Verify routing to specified graph."""
    assert context.process_successful


@then("return the graph's processed response")
def step_return_graph_processed_response(context):
    """Verify graph's processed response is returned."""
    assert "graph routed response" in context.process_result.lower()


@then("it should route to the specified stream")
def step_should_route_to_specified_stream(context):
    """Verify routing to specified stream."""
    assert context.process_successful


@then("return the stream's processed response")
def step_return_stream_processed_response(context):
    """Verify stream's processed response is returned."""
    assert context.process_result is not None


@then("the error should mention unknown input type")
def step_error_mentions_unknown_input_type(context):
    """Verify error mentions unknown input type."""
    assert "unknown input type" in str(context.raised_exception).lower()


@then("the error should mention agent not found")
def step_error_mentions_agent_not_found(context):
    """Verify error mentions agent not found."""
    assert "agent" in str(context.raised_exception).lower()
    assert "not found" in str(context.raised_exception).lower()


@then("the error should mention LangGraph bridge not available")
def step_error_mentions_bridge_not_available(context):
    """Verify error mentions LangGraph bridge not available."""
    assert "langgraph bridge not available" in str(context.raised_exception).lower()


@then("the error should mention graph not found")
def step_error_mentions_graph_not_found(context):
    """Verify error mentions graph not found."""
    assert "graph" in str(context.raised_exception).lower()
    assert "not found" in str(context.raised_exception).lower()


@then("it should retrieve the graph from the bridge")
def step_should_retrieve_from_bridge(context):
    """Verify graph is retrieved from bridge."""
    assert context.process_successful
    context.test_data["mock_langgraph_bridge"].get_graph.assert_called_once()


@then("execute the graph successfully")
def step_execute_graph_successfully(context):
    """Verify graph executes successfully."""
    assert context.process_result is not None


@then("it should extract content from the last message")
def step_should_extract_last_message_content(context):
    """Verify content is extracted from last message."""
    assert context.process_successful
    assert "last message content" in context.process_result.lower()


@then("return the message content")
def step_return_message_content(context):
    """Verify message content is returned."""
    assert "content" in context.process_result.lower()


@then("it should convert the result to string")
def step_should_convert_to_string(context):
    """Verify result is converted to string."""
    assert context.process_successful
    assert isinstance(context.process_result, str)


@then("return the string representation")
def step_return_string_representation(context):
    """Verify string representation is returned."""
    assert context.process_result is not None


@then("the error should mention stream router not available")
def step_error_mentions_stream_router_not_available(context):
    """Verify error mentions stream router not available."""
    assert "stream router not available" in str(context.raised_exception).lower()


@then("it should subscribe to the custom output stream")
def step_should_subscribe_to_custom_output(context):
    """Verify subscription to custom output stream."""
    assert context.process_successful
    # Verify the mock observable was used
    context.test_data["mock_stream_router"].observables["custom_output_stream"].subscribe.assert_called_once()


@then("return the processed message from custom stream")
def step_return_from_custom_stream(context):
    """Verify message from custom stream is returned."""
    assert "custom stream response" in context.process_result.lower()


@then("it should subscribe to the default __output__ stream")
def step_should_subscribe_to_default_output(context):
    """Verify subscription to default output stream."""
    assert context.process_successful
    context.test_data["mock_stream_router"].observables["__output__"].subscribe.assert_called_once()


@then("it should extract content from message object")
def step_should_extract_from_message_object(context):
    """Verify content is extracted from message object."""
    assert context.process_successful
    assert "message object content" in context.process_result.lower()


@then("it should handle timeout appropriately")
def step_should_handle_timeout(context):
    """Verify timeout is handled appropriately."""
    # The process should fail due to timeout
    assert not context.process_successful
    assert "timeout" in str(context.raised_exception).lower() or isinstance(
        context.raised_exception, asyncio.TimeoutError
    )


@then("dispose of the subscription properly")
def step_should_dispose_subscription(context):
    """Verify subscription is disposed properly."""
    # Even on timeout, subscription should be disposed
    # This is handled in the finally block
    assert True  # If we get here, the finally block executed


@then("it should include composite capability")
def step_should_include_composite_capability(context):
    """Verify composite capability is included."""
    assert "composite" in context.capabilities


@then("it should include capabilities from all child agents")
def step_should_include_child_capabilities(context):
    """Verify capabilities from child agents are included."""
    # Should include capabilities from both agents
    assert "agent_capability_0" in context.capabilities
    assert "agent_capability_1" in context.capabilities


@then("it should include stateful-workflow for graphs")
def step_should_include_stateful_workflow(context):
    """Verify stateful-workflow capability for graphs."""
    assert "stateful-workflow" in context.capabilities


@then("it should include reactive-processing for streams")
def step_should_include_reactive_processing(context):
    """Verify reactive-processing capability for streams."""
    assert "reactive-processing" in context.capabilities


@then("it should remove duplicate capabilities")
def step_should_remove_duplicate_capabilities(context):
    """Verify duplicate capabilities are removed."""
    # shared_capability appears in both agents, should only appear once
    capability_count = context.capabilities.count("shared_capability")
    assert capability_count == 1


@then("it should include the legacy strategy capability")
def step_should_include_legacy_strategy(context):
    """Verify legacy strategy capability is included."""
    assert "parallel" in context.capabilities


@then("the context should be merged with exposed parameters")
def step_context_merged_with_exposed_params(context):
    """Verify context is merged with exposed parameters."""
    # Both original context and exposed params should be present
    assert hasattr(context, "captured_context")
    captured = context.captured_context
    assert "user_id" in captured  # From additional context
    assert "model" in captured  # From exposed params


@then("exposed parameters should take precedence")
def step_exposed_params_take_precedence(context):
    """Verify exposed parameters take precedence."""
    captured = context.captured_context
    # temperature from exposed_params (0.5) should override additional context (0.9)
    assert captured["temperature"] == 0.5


@then("the merged context should be used for processing")
def step_merged_context_used_for_processing(context):
    """Verify merged context is used for processing."""
    captured = context.captured_context
    # Should have all parameters
    assert "user_id" in captured  # From additional context
    assert "temperature" in captured  # From exposed params (overridden)
    assert "model" in captured  # From exposed params
    assert "new_param" in captured  # From additional context