cleveragents-core/features/steps/mcp_lifecycle_steps.py

"""Step definitions for features/mcp_lifecycle.feature."""

from __future__ import annotations

import json
import time
from typing import Any

from behave import given, then, when
from behave.runner import Context

from cleveragents.mcp.adapter import MCPServerConfig
from cleveragents.mcp.client import McpClient, McpClientConfig
from cleveragents.mcp.registry import McpRegistry
from features.mocks.mock_mcp_transport import MockMCPTransport


def _mock_tool(name: str, desc: str = "", schema: dict | None = None) -> dict:
    return {
        "name": name,
        "description": desc or f"Mock tool {name}",
        "inputSchema": schema or {},
    }


# ---------------------------------------------------------------
# Given Steps
# ---------------------------------------------------------------


@given('an MCP client config for "{name}" with lazy start enabled')
def step_mcp_client_config_lazy(context: Context, name: str) -> None:
    context.mcp_server_config = MCPServerConfig(
        name=name, transport="stdio", command="echo"
    )
    context.mcp_client_config = McpClientConfig(
        server=context.mcp_server_config,
        lazy_start=True,
        idle_timeout_seconds=0,
        health_check_interval_seconds=0,
    )


@given('an MCP client config for "{name}" with lazy start disabled')
def step_mcp_client_config_no_lazy(context: Context, name: str) -> None:
    context.mcp_server_config = MCPServerConfig(
        name=name, transport="stdio", command="echo"
    )
    context.mcp_client_config = McpClientConfig(
        server=context.mcp_server_config,
        lazy_start=False,
        idle_timeout_seconds=0,
        health_check_interval_seconds=0,
    )


@given('an MCP client config for "{name}" with idle timeout {secs:f} seconds')
def step_mcp_client_config_idle(context: Context, name: str, secs: float) -> None:
    context.mcp_server_config = MCPServerConfig(
        name=name, transport="stdio", command="echo"
    )
    context.mcp_client_config = McpClientConfig(
        server=context.mcp_server_config,
        lazy_start=True,
        idle_timeout_seconds=secs,
        health_check_interval_seconds=0,
    )


@given('an MCP client config for "{name}" with health check interval {secs:f} seconds')
def step_mcp_client_config_health(context: Context, name: str, secs: float) -> None:
    context.mcp_server_config = MCPServerConfig(
        name=name, transport="stdio", command="echo"
    )
    context.mcp_client_config = McpClientConfig(
        server=context.mcp_server_config,
        lazy_start=True,
        idle_timeout_seconds=0,
        health_check_interval_seconds=secs,
        auto_restart=True,
    )


@given('a mock MCP transport with tool "{tool_name}"')
def step_mock_transport_with_tool(context: Context, tool_name: str) -> None:
    tools = [_mock_tool(tool_name)]
    results = {tool_name: {"result": "ok"}}
    context.mcp_lifecycle_transport = MockMCPTransport(
        tools=tools, invoke_results=results
    )


@given('a mock MCP transport with tool "{tool_name}" that crashes after {count:d} call')
def step_mock_transport_crash(context: Context, tool_name: str, count: int) -> None:
    """Transport that works for N calls then starts failing health checks."""
    tools = [_mock_tool(tool_name)]
    results = {tool_name: {"result": "ok"}}
    call_counter: list[int] = [0]
    crash_active: list[bool] = [False]

    class _CrashAfterNTransport(MockMCPTransport):
        def connect(self, config: MCPServerConfig) -> dict[str, Any]:
            # Reset crash state and counter on reconnect so the
            # server can recover after a restart.
            crash_active[0] = False
            call_counter[0] = 0
            return super().connect(config)

        def call(self, method: str, params: dict[str, Any]) -> dict[str, Any]:
            if method == "tools/list":
                call_counter[0] += 1
                if call_counter[0] > count:
                    crash_active[0] = True
                if crash_active[0]:
                    msg = "Server crashed"
                    raise ConnectionError(msg)
            return super().call(method, params)

    context.mcp_lifecycle_transport = _CrashAfterNTransport(
        tools=tools, invoke_results=results
    )


@given("a mock MCP transport with full capabilities")
def step_mock_transport_full_caps(context: Context) -> None:
    tools = [_mock_tool("test_tool")]

    class _FullCapsTransport(MockMCPTransport):
        def connect(self, config: MCPServerConfig) -> dict[str, Any]:
            self._connected = True
            return {
                "capabilities": {
                    "tools": True,
                    "resources": True,
                    "prompts": True,
                }
            }

    context.mcp_lifecycle_transport = _FullCapsTransport(tools=tools)


@given("an MCP registry")
def step_mcp_registry(context: Context) -> None:
    context.mcp_lifecycle_registry = McpRegistry()


@given("an MCP registry with {count:d} running servers")
def step_mcp_registry_running(context: Context, count: int) -> None:
    registry = McpRegistry()
    for i in range(count):
        server_name = f"server-{i}"
        tools = [_mock_tool(f"tool_{i}")]
        transport = MockMCPTransport(tools=tools)
        config = McpClientConfig(
            server=MCPServerConfig(name=server_name, transport="stdio", command="echo"),
            lazy_start=True,
            idle_timeout_seconds=0,
            health_check_interval_seconds=0,
        )
        client = registry.register(config, namespace=f"ns-{i}", transport=transport)
        client.start()
    context.mcp_lifecycle_registry = registry


# ---------------------------------------------------------------
# When Steps
# ---------------------------------------------------------------


@when("I create an MCP client")
def step_create_mcp_client(context: Context) -> None:
    context.mcp_lifecycle_client = McpClient(
        config=context.mcp_client_config,
        transport=context.mcp_lifecycle_transport,
    )
    context.mcp_lifecycle_error = None


@when('I call MCP tool "{tool_name}" with arguments {args_json}')
def step_call_mcp_tool(context: Context, tool_name: str, args_json: str) -> None:
    arguments = json.loads(args_json)
    context.mcp_lifecycle_result = context.mcp_lifecycle_client.call_tool(
        tool_name, arguments
    )


@when('I call MCP tool "{tool_name}" expecting a runtime error')
def step_call_mcp_tool_error(context: Context, tool_name: str) -> None:
    context.mcp_lifecycle_error = None
    try:
        context.mcp_lifecycle_client.call_tool(tool_name, {})
    except Exception as exc:
        context.mcp_lifecycle_error = str(exc)


@when("I explicitly start the MCP client")
def step_start_mcp_client(context: Context) -> None:
    context.mcp_lifecycle_client.start()


@when("I shutdown the MCP client")
def step_shutdown_mcp_client(context: Context) -> None:
    context.mcp_lifecycle_client.shutdown()


@when("I wait {secs:f} seconds")
def step_wait_seconds(context: Context, secs: float) -> None:
    # Use the original sleep if the test environment has patched
    # time.sleep (e.g. _install_fast_sleep_patch caps at 10 ms).
    # Timer-based lifecycle tests need real wall-clock advancement.
    original_sleep = getattr(time, "_original_sleep", time.sleep)
    original_sleep(secs)


@when('I register the MCP client in namespace "{namespace}"')
def step_register_mcp_client(context: Context, namespace: str) -> None:
    context.mcp_lifecycle_registry.register(
        config=context.mcp_client_config,
        namespace=namespace,
        transport=context.mcp_lifecycle_transport,
    )


@when('I register another MCP client config for "{name}" in namespace "{namespace}"')
def step_register_another_mcp_client(
    context: Context, name: str, namespace: str
) -> None:
    config = McpClientConfig(
        server=MCPServerConfig(name=name, transport="stdio", command="echo"),
        lazy_start=True,
        idle_timeout_seconds=0,
        health_check_interval_seconds=0,
    )
    tools = [_mock_tool("default_tool")]
    transport = MockMCPTransport(tools=tools)
    context.mcp_lifecycle_registry.register(
        config=config, namespace=namespace, transport=transport
    )


@when(
    'I call tool "{tool_name}" on MCP namespace "{namespace}" with arguments {args_json}'
)
def step_registry_call_tool(
    context: Context, tool_name: str, namespace: str, args_json: str
) -> None:
    arguments = json.loads(args_json)
    context.mcp_lifecycle_result = context.mcp_lifecycle_registry.call_tool(
        namespace, tool_name, arguments
    )


@when("I shutdown all MCP servers")
def step_shutdown_all(context: Context) -> None:
    context.mcp_lifecycle_registry.shutdown_all()


# ---------------------------------------------------------------
# Then Steps
# ---------------------------------------------------------------


@then('the MCP client state should be "{state}"')
def step_mcp_client_state(context: Context, state: str) -> None:
    actual = context.mcp_lifecycle_client.state
    assert actual == state, f"Expected state '{state}', got '{actual}'"


@then("the MCP client should not be connected")
def step_mcp_client_not_connected(context: Context) -> None:
    assert not context.mcp_lifecycle_client.is_connected


@then("the MCP client should be connected")
def step_mcp_client_connected(context: Context) -> None:
    assert context.mcp_lifecycle_client.is_connected


@then('the MCP client error should mention "{text}"')
def step_mcp_client_error_mention(context: Context, text: str) -> None:
    assert context.mcp_lifecycle_error is not None, "Expected error but none occurred"
    assert text.lower() in context.mcp_lifecycle_error.lower(), (
        f"Error '{context.mcp_lifecycle_error}' does not mention '{text}'"
    )


@then("the MCP client restart count should be greater than {count:d}")
def step_mcp_restart_count_gt(context: Context, count: int) -> None:
    actual = context.mcp_lifecycle_client.restart_count
    assert actual > count, f"Expected restart count > {count}, got {actual}"


@then("the MCP client health check failures should be {count:d}")
def step_mcp_health_failures(context: Context, count: int) -> None:
    actual = context.mcp_lifecycle_client.health_check_failures
    assert actual == count, f"Expected {count} health failures, got {actual}"


@then("the MCP registry should have {count:d} namespaces")
def step_registry_namespace_count(context: Context, count: int) -> None:
    actual = len(context.mcp_lifecycle_registry.list_namespaces())
    assert actual == count, f"Expected {count} namespaces, got {actual}"


@then('MCP namespace "{namespace}" should not be connected')
def step_namespace_not_connected(context: Context, namespace: str) -> None:
    client = context.mcp_lifecycle_registry.get(namespace)
    assert client is not None, f"Namespace '{namespace}' not found"
    assert not client.is_connected


@then('MCP namespace "{namespace}" should be connected')
def step_namespace_connected(context: Context, namespace: str) -> None:
    client = context.mcp_lifecycle_registry.get(namespace)
    assert client is not None, f"Namespace '{namespace}' not found"
    assert client.is_connected


@then("all MCP servers should be stopped")
def step_all_stopped(context: Context) -> None:
    for ns in context.mcp_lifecycle_registry.list_namespaces():
        client = context.mcp_lifecycle_registry.get(ns)
        assert client is not None
        assert not client.is_connected, f"Server in namespace '{ns}' still connected"


@then("the MCP adapter capability metadata should have tools true")
def step_cap_tools_true(context: Context) -> None:
    meta = context.mcp_lifecycle_client.adapter.capability_metadata
    assert meta.tools is True, f"Expected tools=True, got {meta.tools}"


@then("the MCP adapter capability metadata should have resources true")
def step_cap_resources_true(context: Context) -> None:
    meta = context.mcp_lifecycle_client.adapter.capability_metadata
    assert meta.resources is True, f"Expected resources=True, got {meta.resources}"


@then("the MCP adapter capability metadata should have prompts true")
def step_cap_prompts_true(context: Context) -> None:
    meta = context.mcp_lifecycle_client.adapter.capability_metadata
    assert meta.prompts is True, f"Expected prompts=True, got {meta.prompts}"