cleveragents-core/features/steps/mcp_sandbox_rewrite_steps.py

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

from __future__ import annotations

from typing import Any

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

from cleveragents.mcp.sandbox import SandboxPathRewriter

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


@given(
    'a sandbox path rewriter with host root "{host_root}" and sandbox root "{sandbox_root}"'
)
def step_sandbox_rewriter(context: Context, host_root: str, sandbox_root: str) -> None:
    context.sandbox_rewriter = SandboxPathRewriter(
        host_root=host_root, sandbox_root=sandbox_root
    )


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


@when('I rewrite tool arguments with key "{key}" and value "{value}"')
def step_rewrite_arguments_kv(context: Context, key: str, value: str) -> None:
    arguments = {key: value}
    context.sandbox_rewritten_args = context.sandbox_rewriter.rewrite_arguments(
        arguments
    )


@when(
    'I rewrite nested tool arguments with outer key "{outer}" inner key "{inner}" and value "{value}"'
)
def step_rewrite_nested_arguments(
    context: Context, outer: str, inner: str, value: str
) -> None:
    arguments = {outer: {inner: value}}
    context.sandbox_rewritten_args = context.sandbox_rewriter.rewrite_arguments(
        arguments
    )


@when('I rewrite tool arguments with key "{key}" and path list "{paths}"')
def step_rewrite_arguments_list(context: Context, key: str, paths: str) -> None:
    path_list = paths.split(",")
    arguments = {key: path_list}
    context.sandbox_rewritten_args = context.sandbox_rewriter.rewrite_arguments(
        arguments
    )


@when('I rewrite tool response with key "{key}" and value "{value}"')
def step_rewrite_response_kv(context: Context, key: str, value: str) -> None:
    response = {key: value}
    context.sandbox_rewritten_resp = context.sandbox_rewriter.rewrite_response(response)


@when(
    'I rewrite nested tool response with outer key "{outer}" inner key "{inner}" and value "{value}"'
)
def step_rewrite_nested_response(
    context: Context, outer: str, inner: str, value: str
) -> None:
    response = {outer: {inner: value}}
    context.sandbox_rewritten_resp = context.sandbox_rewriter.rewrite_response(response)


@when('I roundtrip rewrite tool arguments with key "{key}" and value "{value}"')
def step_roundtrip_rewrite(context: Context, key: str, value: str) -> None:
    arguments = {key: value}
    rewritten = context.sandbox_rewriter.rewrite_arguments(arguments)
    # Simulate the sandbox returning the rewritten path
    context.sandbox_roundtrip_result = context.sandbox_rewriter.rewrite_response(
        rewritten
    )


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


@then('the rewritten arguments should have "{key}" equal to "{expected}"')
def step_rewritten_arg_equal(context: Context, key: str, expected: str) -> None:
    actual = context.sandbox_rewritten_args.get(key)
    assert actual == expected, f"Expected '{expected}', got '{actual}'"


@then('the rewritten nested argument "{path}" should equal "{expected}"')
def step_rewritten_nested_arg(context: Context, path: str, expected: str) -> None:
    parts = path.split(".")
    value: Any = context.sandbox_rewritten_args
    for part in parts:
        value = value[part]
    assert value == expected, f"Expected '{expected}', got '{value}'"


@then('the rewritten argument "{key}" should be a list with {count:d} rewritten paths')
def step_rewritten_arg_list(context: Context, key: str, count: int) -> None:
    value = context.sandbox_rewritten_args.get(key)
    assert isinstance(value, list), f"Expected list, got {type(value)}"
    assert len(value) == count, f"Expected {count} items, got {len(value)}"
    for item in value:
        assert item.startswith("/workspace"), f"Path not rewritten: {item}"


@then('the rewritten response should have "{key}" equal to "{expected}"')
def step_rewritten_resp_equal(context: Context, key: str, expected: str) -> None:
    actual = context.sandbox_rewritten_resp.get(key)
    assert actual == expected, f"Expected '{expected}', got '{actual}'"


@then('the rewritten nested response "{path}" should equal "{expected}"')
def step_rewritten_nested_resp(context: Context, path: str, expected: str) -> None:
    parts = path.split(".")
    value: Any = context.sandbox_rewritten_resp
    for part in parts:
        value = value[part]
    assert value == expected, f"Expected '{expected}', got '{value}'"


@then('the roundtrip path should equal "{expected}"')
def step_roundtrip_equal(context: Context, expected: str) -> None:
    result = context.sandbox_roundtrip_result
    for v in result.values():
        if isinstance(v, str):
            assert v == expected, f"Expected '{expected}', got '{v}'"
            return
    msg = f"No string value found in roundtrip result: {result}"
    raise AssertionError(msg)