cleveragents-core/features/steps/mcp_adapter_steps.py

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

from __future__ import annotations

import json
import threading
from typing import Any

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

from cleveragents.mcp.adapter import (
    MCPServerConfig,
    MCPToolAdapter,
    MCPToolFilter,
)
from cleveragents.tool.registry import ToolRegistry
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 {},
    }


# ---------------------------------------------------------------
# Adapter Lifecycle
# ---------------------------------------------------------------


@given('an MCP server config for "{name}" with stdio transport')
def step_mcp_config_stdio(context: Context, name: str) -> None:
    context.mcp_config = MCPServerConfig(
        name=name,
        transport="stdio",
        command="echo",
        args=["hello"],
    )


@given('an MCP server config for "{name}" with sse transport and url "{url}"')
def step_mcp_config_sse(context: Context, name: str, url: str) -> None:
    context.mcp_config = MCPServerConfig(name=name, transport="sse", url=url)


@given("an MCP adapter with a mock transport")
def step_mcp_adapter_mock(context: Context) -> None:
    config = MCPServerConfig(name="test-server", transport="stdio", command="echo")
    context.mcp_transport = MockMCPTransport()
    context.mcp_adapter = MCPToolAdapter(config=config, transport=context.mcp_transport)


@given("a connected MCP adapter with a mock transport")
def step_mcp_adapter_connected(context: Context) -> None:
    config = MCPServerConfig(name="test-server", transport="stdio", command="echo")
    context.mcp_transport = MockMCPTransport()
    context.mcp_adapter = MCPToolAdapter(config=config, transport=context.mcp_transport)
    context.mcp_adapter.connect()


@given("a connected MCP adapter with {count:d} mock tools")
def step_mcp_adapter_with_tools(context: Context, count: int) -> None:
    tools = [_mock_tool(f"tool_{i}") for i in range(count)]
    config = MCPServerConfig(name="test-server", transport="stdio", command="echo")
    context.mcp_transport = MockMCPTransport(tools=tools)
    context.mcp_adapter = MCPToolAdapter(config=config, transport=context.mcp_transport)
    context.mcp_adapter.connect()


@given('a connected MCP adapter with a callable mock tool "{tool_name}"')
def step_mcp_adapter_callable_tool(context: Context, tool_name: str) -> None:
    tools = [_mock_tool(tool_name)]
    results = {tool_name: {"id": 42, "status": "created"}}
    config = MCPServerConfig(name="test-server", transport="stdio", command="echo")
    context.mcp_transport = MockMCPTransport(tools=tools, invoke_results=results)
    context.mcp_adapter = MCPToolAdapter(config=config, transport=context.mcp_transport)
    context.mcp_adapter.connect()
    context.mcp_adapter.discover_tools()


@given('a connected MCP adapter with a schema-validated mock tool "{tool_name}"')
def step_mcp_adapter_schema_tool(context: Context, tool_name: str) -> None:
    schema: dict[str, Any] = {
        "type": "object",
        "properties": {
            "title": {"type": "string"},
            "body": {"type": "string"},
        },
        "required": ["title", "body"],
    }
    tools = [_mock_tool(tool_name, schema=schema)]
    results = {tool_name: {"id": 1}}
    config = MCPServerConfig(name="test-server", transport="stdio", command="echo")
    context.mcp_transport = MockMCPTransport(tools=tools, invoke_results=results)
    context.mcp_adapter = MCPToolAdapter(config=config, transport=context.mcp_transport)
    context.mcp_adapter.connect()
    context.mcp_adapter.discover_tools()


@given('a connected MCP adapter with a failing mock tool "{tool_name}"')
def step_mcp_adapter_failing_tool(context: Context, tool_name: str) -> None:
    tools = [_mock_tool(tool_name)]
    errors = {tool_name: "Internal server error"}
    config = MCPServerConfig(name="test-server", transport="stdio", command="echo")
    context.mcp_transport = MockMCPTransport(tools=tools, invoke_errors=errors)
    context.mcp_adapter = MCPToolAdapter(config=config, transport=context.mcp_transport)
    context.mcp_adapter.connect()
    context.mcp_adapter.discover_tools()


@given('a tool filter including only "{tool_name}"')
def step_tool_filter_include(context: Context, tool_name: str) -> None:
    context.mcp_tool_filter = MCPToolFilter(include=[tool_name])


@given('a tool filter excluding "{tool_name}"')
def step_tool_filter_exclude(context: Context, tool_name: str) -> None:
    context.mcp_tool_filter = MCPToolFilter(exclude=[tool_name])


@given("an empty MCP tool registry")
def step_empty_mcp_registry(context: Context) -> None:
    context.mcp_registry = ToolRegistry()


@given("an MCP adapter with a transport that fails to connect")
def step_mcp_adapter_fail_connect(context: Context) -> None:
    config = MCPServerConfig(name="bad-server", transport="stdio", command="echo")
    context.mcp_transport = MockMCPTransport(fail_connect=True)
    context.mcp_adapter = MCPToolAdapter(config=config, transport=context.mcp_transport)


@given("a connected MCP adapter with a tool that times out")
def step_mcp_adapter_timeout_tool(context: Context) -> None:
    tools = [_mock_tool("slow_tool")]
    config = MCPServerConfig(name="test-server", transport="stdio", command="echo")
    context.mcp_transport = MockMCPTransport(tools=tools, timeout_tools={"slow_tool"})
    context.mcp_adapter = MCPToolAdapter(config=config, transport=context.mcp_transport)
    context.mcp_adapter.connect()
    context.mcp_adapter.discover_tools()


@given("a connected MCP adapter with a transport that errors on close")
def step_mcp_adapter_close_error(context: Context) -> None:
    config = MCPServerConfig(name="test-server", transport="stdio", command="echo")

    class _ErrorCloseTransport(MockMCPTransport):
        def close(self) -> None:
            msg = "Close failed"
            raise OSError(msg)

    context.mcp_transport = _ErrorCloseTransport()
    context.mcp_adapter = MCPToolAdapter(config=config, transport=context.mcp_transport)
    context.mcp_adapter.connect()


@given('a connected MCP adapter with a raw list-content mock tool "{tool_name}"')
def step_mcp_adapter_raw_list_content_tool(context: Context, tool_name: str) -> None:
    """Adapter whose transport returns a raw MCP 1.4.0 list-format content response."""
    tools = [_mock_tool(tool_name)]
    config = MCPServerConfig(name="test-server", transport="stdio", command="echo")

    class _ListContentTransport(MockMCPTransport):
        def call(self, method: str, params: dict[str, Any]) -> dict[str, Any]:
            if method == "tools/call":
                return {
                    "content": [{"type": "text", "text": "fetched data"}],
                }
            return super().call(method, params)

    context.mcp_transport = _ListContentTransport(tools=tools)
    context.mcp_adapter = MCPToolAdapter(config=config, transport=context.mcp_transport)
    context.mcp_adapter.connect()
    context.mcp_adapter.discover_tools()


@given('a connected MCP adapter with a transport-error mock tool "{tool_name}"')
def step_mcp_adapter_transport_error_tool(context: Context, tool_name: str) -> None:
    tools = [_mock_tool(tool_name)]
    config = MCPServerConfig(name="test-server", transport="stdio", command="echo")

    class _CrashTransport(MockMCPTransport):
        def call(self, method: str, params: dict[str, Any]) -> dict[str, Any]:
            if method == "tools/call":
                msg = "Transport crashed unexpectedly"
                raise RuntimeError(msg)
            return super().call(method, params)

    context.mcp_transport = _CrashTransport(tools=tools)
    context.mcp_adapter = MCPToolAdapter(config=config, transport=context.mcp_transport)
    context.mcp_adapter.connect()
    context.mcp_adapter.discover_tools()


@given('an MCP server config for "{name}" with stdio transport but no command')
def step_mcp_config_stdio_no_cmd(context: Context, name: str) -> None:
    context.mcp_config = MCPServerConfig(name=name, transport="stdio")


@given('an MCP server config for "{name}" with sse transport but no url')
def step_mcp_config_sse_no_url(context: Context, name: str) -> None:
    context.mcp_config = MCPServerConfig(name=name, transport="sse")


@given("a connected MCP adapter with a transport that fails to reconnect")
def step_mcp_adapter_fail_reconnect(context: Context) -> None:
    """Adapter starts connected but the transport will refuse the next connect."""
    config = MCPServerConfig(name="flaky-server", transport="stdio", command="echo")
    call_count: list[int] = [0]

    class _FlakeyTransport(MockMCPTransport):
        def connect(self, cfg: MCPServerConfig) -> dict[str, Any]:
            call_count[0] += 1
            if call_count[0] > 1:
                msg = "Server unavailable after restart"
                raise ConnectionRefusedError(msg)
            return super().connect(cfg)

    context.mcp_transport = _FlakeyTransport()
    context.mcp_adapter = MCPToolAdapter(config=config, transport=context.mcp_transport)
    context.mcp_adapter.connect()


@given("an MCP adapter with a transport that hangs on connect")
def step_mcp_adapter_hang_connect(context: Context) -> None:
    """Adapter whose transport blocks indefinitely during connect."""
    config = MCPServerConfig(name="hang-server", transport="stdio", command="echo")

    _hang_event = threading.Event()

    class _HangTransport(MockMCPTransport):
        def connect(self, cfg: MCPServerConfig) -> dict[str, Any]:
            _hang_event.wait(timeout=10)
            return super().connect(cfg)

    context.mcp_transport = _HangTransport()
    context.mcp_adapter = MCPToolAdapter(config=config, transport=context.mcp_transport)
    context._hang_event = _hang_event


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


@when("I create an MCP adapter from the config")
def step_create_adapter(context: Context) -> None:
    context.mcp_error = None
    try:
        context.mcp_adapter = MCPToolAdapter.from_server_config(context.mcp_config)
    except Exception as exc:
        context.mcp_error = str(exc)


@when("I create an MCP adapter from the config expecting validation error")
def step_create_adapter_fail(context: Context) -> None:
    context.mcp_error = None
    try:
        context.mcp_adapter = MCPToolAdapter.from_server_config(context.mcp_config)
    except Exception as exc:
        context.mcp_error = str(exc)


@when("I connect the adapter")
def step_connect_adapter(context: Context) -> None:
    context.mcp_error = None
    context.mcp_adapter.connect()


@when("I connect the adapter expecting error")
def step_connect_adapter_fail(context: Context) -> None:
    context.mcp_error = None
    try:
        context.mcp_adapter.connect()
    except Exception as exc:
        context.mcp_error = str(exc)


@when("I disconnect the adapter")
def step_disconnect_adapter(context: Context) -> None:
    context.mcp_adapter.disconnect()


@when("I reconnect the adapter")
def step_reconnect_adapter(context: Context) -> None:
    context.mcp_error = None
    context.mcp_adapter.reconnect()


@when("I reconnect the adapter expecting error")
def step_reconnect_adapter_fail(context: Context) -> None:
    context.mcp_error = None
    try:
        context.mcp_adapter.reconnect()
    except Exception as exc:
        context.mcp_error = str(exc)


@when("I connect the adapter with timeout {secs:f} expecting error")
def step_connect_with_timeout_fail(context: Context, secs: float) -> None:
    context.mcp_error = None
    try:
        context.mcp_adapter.connect(timeout=secs)
    except Exception as exc:
        context.mcp_error = str(exc)


@when("I discover tools from the adapter")
def step_discover_tools(context: Context) -> None:
    context.mcp_discovered = context.mcp_adapter.discover_tools()


@when("I discover tools from the adapter with filter")
def step_discover_tools_filtered(context: Context) -> None:
    context.mcp_discovered = context.mcp_adapter.discover_tools(
        tool_filter=context.mcp_tool_filter
    )


@when("I discover tools expecting an error")
def step_discover_tools_error(context: Context) -> None:
    context.mcp_error = None
    try:
        context.mcp_adapter.discover_tools()
    except Exception as exc:
        context.mcp_error = str(exc)


@when('I invoke "{tool_name}" with arguments {args_json}')
def step_invoke_tool(context: Context, tool_name: str, args_json: str) -> None:
    arguments = json.loads(args_json)
    context.mcp_invoke_result = context.mcp_adapter.invoke(tool_name, arguments)


@when('I invoke MCP tool "{tool_name}" while disconnected')
def step_invoke_not_connected(context: Context, tool_name: str) -> None:
    context.mcp_error = None
    try:
        context.mcp_adapter.invoke(tool_name, {})
    except Exception as exc:
        context.mcp_error = str(exc)


@when('I register MCP tools in the registry with namespace "{namespace}"')
def step_register_tools(context: Context, namespace: str) -> None:
    context.mcp_registered = context.mcp_adapter.register_tools(
        registry=context.mcp_registry,
        namespace=namespace,
    )


@when('the MCP server adds a new tool "{tool_name}"')
def step_server_adds_tool(context: Context, tool_name: str) -> None:
    context.mcp_transport.add_tool(_mock_tool(tool_name))


@when('I re-register MCP tools in the registry with namespace "{namespace}"')
def step_reregister_tools(context: Context, namespace: str) -> None:
    context.mcp_registered = context.mcp_adapter.register_tools(
        registry=context.mcp_registry,
        namespace=namespace,
    )


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


@then('the adapter server name should be "{name}"')
def step_adapter_name(context: Context, name: str) -> None:
    assert context.mcp_adapter.server_name == name


@then('the adapter transport should be "{transport}"')
def step_adapter_transport(context: Context, transport: str) -> None:
    assert context.mcp_adapter.transport_type == transport


@then("the adapter should not be connected")
def step_adapter_not_connected(context: Context) -> None:
    assert not context.mcp_adapter.is_connected


@then("the adapter should be connected")
def step_adapter_connected(context: Context) -> None:
    assert context.mcp_adapter.is_connected


@then("the adapter server capabilities should be set")
def step_adapter_capabilities(context: Context) -> None:
    assert context.mcp_adapter.capabilities


@then("the adapter should have {count:d} discovered tools")
def step_adapter_tool_count(context: Context, count: int) -> None:
    assert len(context.mcp_discovered) == count, (
        f"Expected {count} tools, got {len(context.mcp_discovered)}"
    )


@then("discovered tool {idx:d} should have a name")
def step_discovered_tool_name(context: Context, idx: int) -> None:
    assert context.mcp_discovered[idx].name


@then("discovered tool {idx:d} should have an input schema")
def step_discovered_tool_schema(context: Context, idx: int) -> None:
    assert context.mcp_discovered[idx].input_schema is not None


@then("the invocation should succeed")
def step_invoke_success(context: Context) -> None:
    assert context.mcp_invoke_result.success, (
        f"Expected success, got error: {context.mcp_invoke_result.error}"
    )


@then("the invocation should fail")
def step_invoke_fail(context: Context) -> None:
    assert not context.mcp_invoke_result.success


@then('the invocation result should contain key "{key}"')
def step_invoke_result_key(context: Context, key: str) -> None:
    assert key in context.mcp_invoke_result.data, (
        f"Key '{key}' not in result: {context.mcp_invoke_result.data}"
    )


@then('the invocation result data should have key "{key}"')
def step_invoke_result_data_key(context: Context, key: str) -> None:
    assert isinstance(context.mcp_invoke_result.data, dict), (
        f"Expected data to be a dict, got {type(context.mcp_invoke_result.data)}"
    )
    assert key in context.mcp_invoke_result.data, (
        f"Key '{key}' not in result data: {context.mcp_invoke_result.data}"
    )


@then("the invocation result content should be a list")
def step_invoke_result_content_is_list(context: Context) -> None:
    content = context.mcp_invoke_result.data.get("content")
    assert isinstance(content, list), (
        f"Expected 'content' to be a list, got {type(content)}: {content}"
    )


@then("the invocation result content list should have {count:d} item")
def step_invoke_result_content_count(context: Context, count: int) -> None:
    content = context.mcp_invoke_result.data.get("content")
    assert isinstance(content, list), (
        f"Expected 'content' to be a list, got {type(content)}"
    )
    assert len(content) == count, (
        f"Expected {count} content item(s), got {len(content)}: {content}"
    )


@then('the invocation error should mention "{text}"')
def step_invoke_error_mention(context: Context, text: str) -> None:
    assert context.mcp_invoke_result.error is not None
    assert text.lower() in context.mcp_invoke_result.error.lower(), (
        f"Error '{context.mcp_invoke_result.error}' does not mention '{text}'"
    )


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


@then("the registry should have {count:d} tools")
def step_registry_count(context: Context, count: int) -> None:
    all_tools = context.mcp_registry.list_tools()
    assert len(all_tools) == count, f"Expected {count} tools, got {len(all_tools)}"


@then('registry tool {idx:d} name should start with "{prefix}"')
def step_registry_tool_prefix(context: Context, idx: int, prefix: str) -> None:
    all_tools = context.mcp_registry.list_tools()
    assert all_tools[idx].name.startswith(prefix)


@then('calling registry tool "{tool_name}" should succeed')
def step_registry_tool_callable(context: Context, tool_name: str) -> None:
    spec = context.mcp_registry.get(tool_name)
    assert spec is not None, f"Tool '{tool_name}' not in registry"
    result = spec.handler()
    assert isinstance(result, dict)


@then('calling registry tool "{tool_name}" should return an error')
def step_registry_tool_error(context: Context, tool_name: str) -> None:
    spec = context.mcp_registry.get(tool_name)
    assert spec is not None, f"Tool '{tool_name}' not in registry"
    result = spec.handler()
    assert isinstance(result, dict)
    assert "error" in result


@then("the adapter discovered_tools property should have {count:d} items")
def step_adapter_discovered_tools_property(context: Context, count: int) -> None:
    dt = context.mcp_adapter.discovered_tools
    assert len(dt) == count, f"Expected {count}, got {len(dt)}"


# ---------------------------------------------------------------
# Capability Inference Steps
# ---------------------------------------------------------------


@when('I infer capabilities for tool name "{tool_name}"')
def step_infer_capabilities(context: Context, tool_name: str) -> None:
    context.inferred_caps = MCPToolAdapter.infer_capabilities(tool_name)


@then("the inferred capabilities should have read_only true")
def step_inferred_read_only_true(context: Context) -> None:
    assert context.inferred_caps["read_only"] is True


@then("the inferred capabilities should have read_only false")
def step_inferred_read_only_false(context: Context) -> None:
    assert context.inferred_caps["read_only"] is False


@then("the inferred capabilities should have writes true")
def step_inferred_writes_true(context: Context) -> None:
    assert context.inferred_caps["writes"] is True


@then("the inferred capabilities should have writes false")
def step_inferred_writes_false(context: Context) -> None:
    assert context.inferred_caps["writes"] is False


@then('registry tool "{tool_name}" should have read_only capability')
def step_registry_tool_read_only(context: Context, tool_name: str) -> None:
    spec = context.mcp_registry.get(tool_name)
    assert spec is not None, f"Tool '{tool_name}' not in registry"
    assert spec.capabilities.read_only is True


@then('registry tool "{tool_name}" should have writes capability')
def step_registry_tool_writes(context: Context, tool_name: str) -> None:
    spec = context.mcp_registry.get(tool_name)
    assert spec is not None, f"Tool '{tool_name}' not in registry"
    assert spec.capabilities.writes is True


@then('registry tool "{tool_name}" should have source "{source}"')
def step_registry_tool_source(context: Context, tool_name: str, source: str) -> None:
    spec = context.mcp_registry.get(tool_name)
    assert spec is not None, f"Tool '{tool_name}' not in registry"
    assert spec.source == source, f"Expected source '{source}', got '{spec.source}'"


@then('registry tool "{tool_name}" should have checkpointable false')
def step_registry_tool_checkpointable_false(context: Context, tool_name: str) -> None:
    spec = context.mcp_registry.get(tool_name)
    assert spec is not None, f"Tool '{tool_name}' not in registry"
    assert spec.capabilities.checkpointable is False, (
        f"Expected checkpointable=False, got {spec.capabilities.checkpointable}"
    )