cleveragents-core/features/steps/session_tell_llm_steps.py

"""Step definitions for session tell real LLM invocation tests (issue #5784).

Tests verify that ``agents session tell`` routes through
:class:`~cleveragents.application.services.session_workflow.SessionWorkflow`
with a real (mock) :class:`LLMCaller`, persists the response, and records
token usage.

All LLM interactions are handled via a mock ``LLMCaller`` so no real API
keys or network access is required.
"""

from __future__ import annotations

import contextlib
import json
import re
from collections.abc import Iterator
from typing import Any
from unittest.mock import MagicMock, patch

from behave import given, then, when
from behave.runner import Context
from typer.testing import CliRunner

from cleveragents.application.services.session_workflow import SessionWorkflow
from cleveragents.cli.commands.session import app as session_app
from cleveragents.domain.models.core.session import (
    MessageRole,
    Session,
    SessionMessage,
    SessionService,
    SessionTokenUsage,
)

_ULID = "01KJ1N8YDN05N29P5RZV4SPDVZ"
_STUB_RESPONSE = "Here is what I can help you with."


class _StubResp:
    """Minimal LLM response stub for _StubLLM.invoke()."""

    def __init__(self, content: str) -> None:
        self.content = content
        self.tool_calls: list[Any] = []
        self.response_metadata: dict[str, Any] = {
            "usage": {"input_tokens": 10, "output_tokens": 5}
        }


class _StubChunk:
    """Minimal streaming chunk for _StubLLM.stream()."""

    def __init__(self, content: str) -> None:
        self.content = content


runner = CliRunner()


class _StubLLM:
    """Minimal LangChain-compatible LLM stub for LangChainSessionCaller tests."""

    def __init__(self, response: str = _STUB_RESPONSE) -> None:
        self._response = response
        self.invoked_messages: list[Any] = []

    def invoke(self, messages: Any, **kwargs: Any) -> _StubResp:
        self.invoked_messages = list(messages)
        return _StubResp(self._response)

    def stream(self, messages: Any, **kwargs: Any) -> Iterator[_StubChunk]:
        self.invoked_messages = list(messages)
        words = self._response.split()
        for i, word in enumerate(words):
            # Reconstruct spaces between words so the concatenated
            # output matches the original response string verbatim.
            token = word if i == 0 else " " + word
            yield _StubChunk(token)


def _make_session(
    session_id: str = _ULID,
    actor_name: str | None = None,
) -> Session:
    return Session(
        session_id=session_id,
        actor_name=actor_name,
        namespace="local",
        messages=[],
        token_usage=SessionTokenUsage(),
    )


def _make_mock_service(session: Session) -> MagicMock:
    """Build a MagicMock SessionService pre-configured for session tell."""
    svc = MagicMock(spec=SessionService)
    svc.get.return_value = session
    svc.get_messages.return_value = []
    svc.append_message.return_value = MagicMock(spec=SessionMessage)
    svc.update_token_usage.return_value = None
    return svc


# ---------------------------------------------------------------------------
# Background
# ---------------------------------------------------------------------------


@given("a session tell LLM mock environment")
def step_setup_llm_mock_env(context: Context) -> None:
    """Set up runner, stub LLM, and cleanup handlers."""
    context.runner = runner
    context.stub_llm = _StubLLM()
    context.captured_actor_name: str | None = None
    context._cleanup_handlers: list[Any] = []

    def cleanup() -> None:
        for stop in reversed(context._cleanup_handlers):
            with contextlib.suppress(Exception):
                stop()

    context.add_cleanup(cleanup)


# ---------------------------------------------------------------------------
# Given steps
# ---------------------------------------------------------------------------


@given('a session with actor "{actor}" exists')
def step_session_with_actor(context: Context, actor: str) -> None:
    session = _make_session(actor_name=actor)
    context.session = session
    context.svc = _make_mock_service(session)
    # Patch _get_session_service instead of module-level _service so tests
    # are resilient to changes in the internal caching strategy (m4).
    svc_patcher = patch(
        "cleveragents.cli.commands.session._get_session_service",
        return_value=context.svc,
    )
    svc_patcher.start()
    context._cleanup_handlers.append(svc_patcher.stop)


@given("a session with no actor exists")
def step_session_no_actor(context: Context) -> None:
    session = _make_session(actor_name=None)
    context.session = session
    context.svc = _make_mock_service(session)
    svc_patcher = patch(
        "cleveragents.cli.commands.session._get_session_service",
        return_value=context.svc,
    )
    svc_patcher.start()
    context._cleanup_handlers.append(svc_patcher.stop)


# ---------------------------------------------------------------------------
# When steps
# ---------------------------------------------------------------------------


@when('I invoke session tell with prompt "{prompt}"')
def step_invoke_tell(context: Context, prompt: str) -> None:
    """Invoke tell routing through _facade_dispatch.

    The non-streaming CLI path calls :func:`_facade_dispatch`, which in
    production delegates to ``A2aLocalFacade``.  The test intercepts
    ``_facade_dispatch`` with a thin adapter that runs the real
    ``SessionWorkflow.tell()`` (backed by the test's stub LLM) and
    returns a dict matching the facade's response envelope.  This
    keeps the workflow assertions (append_message, token usage) working
    while also testing the CLI's facade-routing code path.
    """
    context.prompt = prompt
    context.invoked_actor_name = None
    stub_llm = context.stub_llm
    svc = context.svc

    def _llm_factory(actor: str) -> Any:
        context.invoked_actor_name = actor
        return stub_llm

    wf = SessionWorkflow(session_service=svc, llm_factory=_llm_factory)

    def _mock_dispatch(operation: str, params: dict[str, Any]) -> dict[str, Any]:
        if operation == "message/send":
            result = wf.tell(
                session_id=params["session_id"],
                prompt=params["message"],
                actor_override=params.get("actor"),
            )
            return {
                "session_id": result.session_id,
                "assistant_message": result.assistant_message,
                "usage": result.usage_dict(),
            }
        raise RuntimeError(f"Unknown operation: {operation}")

    dispatch_patcher = patch(
        "cleveragents.cli.commands.session._facade_dispatch",
        side_effect=_mock_dispatch,
    )
    dispatch_patcher.start()
    context._cleanup_handlers.append(dispatch_patcher.stop)

    context.result = runner.invoke(
        session_app,
        ["tell", "--session", _ULID, prompt],
    )


@when('I invoke session tell with --stream and prompt "{prompt}"')
def step_invoke_tell_stream(context: Context, prompt: str) -> None:
    """Invoke tell --stream with a real SessionWorkflow backed by _StubLLM."""
    context.prompt = prompt
    stub_llm = context.stub_llm
    svc = context.svc

    def _workflow_factory() -> SessionWorkflow:
        return SessionWorkflow(
            session_service=svc,
            llm_factory=lambda actor: stub_llm,
        )

    wf_patcher = patch(
        "cleveragents.cli.commands.session._build_session_workflow",
        side_effect=_workflow_factory,
    )
    wf_patcher.start()
    context._cleanup_handlers.append(wf_patcher.stop)

    context.result = runner.invoke(
        session_app,
        ["tell", "--session", _ULID, "--stream", prompt],
    )


@when('I invoke session tell with --actor "{actor}" and prompt "{prompt}"')
def step_invoke_tell_with_actor(context: Context, actor: str, prompt: str) -> None:
    """Invoke tell with --actor override, routing through _facade_dispatch."""
    context.prompt = prompt
    context.override_actor = actor
    stub_llm = context.stub_llm
    svc = context.svc

    def _llm_factory(actor_name: str) -> Any:
        context.invoked_actor_name = actor_name
        return stub_llm

    wf = SessionWorkflow(session_service=svc, llm_factory=_llm_factory)

    def _mock_dispatch(operation: str, params: dict[str, Any]) -> dict[str, Any]:
        if operation == "message/send":
            result = wf.tell(
                session_id=params["session_id"],
                prompt=params["message"],
                actor_override=params.get("actor"),
            )
            return {
                "session_id": result.session_id,
                "assistant_message": result.assistant_message,
                "usage": result.usage_dict(),
            }
        raise RuntimeError(f"Unknown operation: {operation}")

    dispatch_patcher = patch(
        "cleveragents.cli.commands.session._facade_dispatch",
        side_effect=_mock_dispatch,
    )
    dispatch_patcher.start()
    context._cleanup_handlers.append(dispatch_patcher.stop)

    context.result = runner.invoke(
        session_app,
        ["tell", "--session", _ULID, "--actor", actor, prompt],
    )


@when('I invoke session tell with --format json and prompt "{prompt}"')
def step_invoke_tell_json(context: Context, prompt: str) -> None:
    """Invoke tell --format json, routing through _facade_dispatch."""
    context.prompt = prompt
    stub_llm = context.stub_llm
    svc = context.svc

    wf = SessionWorkflow(
        session_service=svc,
        llm_factory=lambda actor: stub_llm,
    )

    def _mock_dispatch(operation: str, params: dict[str, Any]) -> dict[str, Any]:
        if operation == "message/send":
            result = wf.tell(
                session_id=params["session_id"],
                prompt=params["message"],
                actor_override=params.get("actor"),
            )
            return {
                "session_id": result.session_id,
                "assistant_message": result.assistant_message,
                "usage": result.usage_dict(),
            }
        raise RuntimeError(f"Unknown operation: {operation}")

    dispatch_patcher = patch(
        "cleveragents.cli.commands.session._facade_dispatch",
        side_effect=_mock_dispatch,
    )
    dispatch_patcher.start()
    context._cleanup_handlers.append(dispatch_patcher.stop)

    context.result = runner.invoke(
        session_app,
        ["tell", "--session", _ULID, "--format", "json", prompt],
    )

    def _mock_dispatch(operation: str, params: dict[str, Any]) -> dict[str, Any]:
        if operation == "message/send":
            result = wf.tell(
                session_id=params["session_id"],
                prompt=params["message"],
                actor_override=params.get("actor"),
            )
            return {
                "session_id": result.session_id,
                "assistant_message": result.assistant_message,
                "usage": result.usage_dict(),
            }
        raise RuntimeError(f"Unknown operation: {operation}")

    dispatch_patcher = patch(
        "cleveragents.cli.commands.session._facade_dispatch",
        side_effect=_mock_dispatch,
    )
    dispatch_patcher.start()
    context._cleanup_handlers.append(dispatch_patcher.stop)

    context.result = runner.invoke(
        session_app,
        ["tell", "--session", _ULID, "--format", "json", prompt],
    )


@when('I invoke session tell with --stream --format json and prompt "{prompt}"')
def step_invoke_tell_stream_json(context: Context, prompt: str) -> None:
    """Invoke tell --stream --format json with a real SessionWorkflow backed by _StubLLM."""
    context.prompt = prompt
    stub_llm = context.stub_llm
    svc = context.svc

    def _workflow_factory() -> SessionWorkflow:
        return SessionWorkflow(
            session_service=svc,
            llm_factory=lambda actor: stub_llm,
        )

    wf_patcher = patch(
        "cleveragents.cli.commands.session._build_session_workflow",
        side_effect=_workflow_factory,
    )
    wf_patcher.start()
    context._cleanup_handlers.append(wf_patcher.stop)

    context.result = runner.invoke(
        session_app,
        ["tell", "--session", _ULID, "--stream", "--format", "json", prompt],
    )


# ---------------------------------------------------------------------------
# Then steps
# ---------------------------------------------------------------------------


@then("the tell command returns the LLM response")
def step_tell_returns_llm_response(context: Context) -> None:
    assert context.result.exit_code == 0, (
        f"Expected exit code 0, got {context.result.exit_code}.\n"
        f"Output:\n{context.result.output}"
    )
    # The stub LLM response should appear in the output
    assert _STUB_RESPONSE in context.result.output, (
        f"LLM response not found in output:\n{context.result.output}"
    )


@then("the assistant message is persisted to the session")
def step_assistant_message_persisted(context: Context) -> None:
    # append_message should be called exactly twice:
    # once for the user message and once for the assistant response.
    call_count = context.svc.append_message.call_count
    assert call_count == 2, (
        f"Expected append_message to be called exactly 2 times (user + assistant), "
        f"got {call_count}"
    )
    # The second call should be for the ASSISTANT role
    calls = context.svc.append_message.call_args_list
    assistant_calls = [
        c
        for c in calls
        if c.kwargs.get("role") == MessageRole.ASSISTANT
        or (len(c.args) >= 2 and c.args[1] == MessageRole.ASSISTANT)
    ]
    assert len(assistant_calls) == 1, (
        f"Expected exactly 1 ASSISTANT persist, got {len(assistant_calls)}. "
        f"Calls: {calls}"
    )
    # The second call should be for the ASSISTANT role
    calls = context.svc.append_message.call_args_list
    assistant_calls = [
        c
        for c in calls
        if c.kwargs.get("role") == MessageRole.ASSISTANT
        or (len(c.args) >= 2 and c.args[1] == MessageRole.ASSISTANT)
    ]
    assert len(assistant_calls) == 1, (
        f"Expected exactly 1 ASSISTANT persist, got {len(assistant_calls)}. "
        f"Calls: {calls}"
    )


@then("token usage is recorded")
def step_token_usage_recorded(context: Context) -> None:
    assert context.svc.update_token_usage.called, (
        "update_token_usage was not called — token usage was not recorded"
    )
    assert context.svc.update_token_usage.call_count == 1, (
        f"Expected update_token_usage to be called exactly once, "
        f"got {context.svc.update_token_usage.call_count}"
    )
    # Verify exact token counts match the metadata from the stub LLM (m5).
    call_kwargs = context.svc.update_token_usage.call_args.kwargs
    assert call_kwargs.get("input_tokens") == 10, (
        f"Expected input_tokens=10, got {call_kwargs.get('input_tokens')}"
    )
    assert call_kwargs.get("output_tokens") == 5, (
        f"Expected output_tokens=5, got {call_kwargs.get('output_tokens')}"
    )
    assert context.svc.update_token_usage.call_count == 1, (
        f"Expected update_token_usage to be called exactly once, "
        f"got {context.svc.update_token_usage.call_count}"
    )
    # Verify exact token counts match the metadata from the stub LLM (m5).
    call_kwargs = context.svc.update_token_usage.call_args.kwargs
    assert call_kwargs.get("input_tokens") == 10, (
        f"Expected input_tokens=10, got {call_kwargs.get('input_tokens')}"
    )
    assert call_kwargs.get("output_tokens") == 5, (
        f"Expected output_tokens=5, got {call_kwargs.get('output_tokens')}"
    )


@then("the streamed output contains the LLM response tokens")
def step_streamed_output_contains_tokens(context: Context) -> None:
    assert context.result.exit_code == 0, (
        f"Expected exit code 0, got {context.result.exit_code}.\n"
        f"Output:\n{context.result.output}"
    )
    # Strip ANSI escape sequences for text matching
    output_clean = re.sub(r"\x1b\[[0-9;]*m", "", context.result.output)
    # When streamed through the A2A facade (Option B: fallback to
    # non-streaming), verify the full response is present in the output
    # rather than checking token-by-token word positions (C5).
    assert _STUB_RESPONSE in output_clean, (
        f"Expected full response '{_STUB_RESPONSE}' in streamed output, "
        f"but it was not found.\nOutput:\n{context.result.output}"
    )


@then("the assistant message is persisted after streaming")
def step_assistant_persisted_after_stream(context: Context) -> None:
    call_count = context.svc.append_message.call_count
    assert call_count == 2, (
        f"Expected append_message to be called exactly 2 times (user + assistant), "
        f"got {call_count}"
    )
    # Verify that the last persisted message is an ASSISTANT role.
    calls = context.svc.append_message.call_args_list
    assert calls[-1].kwargs.get("role") == MessageRole.ASSISTANT, (
        f"Last persisted message role is not ASSISTANT. Calls: {calls}"
    )
    # Verify that the last persisted message is an ASSISTANT role.
    calls = context.svc.append_message.call_args_list
    assert calls[-1].kwargs.get("role") == MessageRole.ASSISTANT, (
        f"Last persisted message role is not ASSISTANT. Calls: {calls}"
    )


@then("the tell command exits with code 1")
def step_tell_exits_code_1(context: Context) -> None:
    assert context.result.exit_code == 1, (
        f"Expected exit code 1, got {context.result.exit_code}.\n"
        f"Output:\n{context.result.output}"
    )


@then("the error output mentions actor configuration")
def step_error_mentions_actor(context: Context) -> None:
    output = context.result.output.lower()
    assert "actor" in output, (
        f"Expected 'actor' in error output.\nGot:\n{context.result.output}"
    )


@then('the tell command uses the override actor "{actor}"')
def step_tell_uses_override_actor(context: Context, actor: str) -> None:
    invoked = getattr(context, "invoked_actor_name", None)
    assert invoked is not None, "Actor name was never captured from _resolve_llm"
    assert invoked == actor, f"Expected actor '{actor}' to be used, but got '{invoked}'"


@then("the tell output is valid JSON with a data envelope")
def step_tell_output_is_valid_json(context: Context) -> None:
    assert context.result.exit_code == 0, (
        f"Expected exit code 0, got {context.result.exit_code}.\n"
        f"Output:\n{context.result.output}"
    )
    try:
        raw = json.loads(context.result.output.strip())
    except json.JSONDecodeError as exc:
        raise AssertionError(
            f"Output is not valid JSON: {exc}\nOutput:\n{context.result.output}"
        ) from exc
    # format_output wraps data in a spec-required envelope
    context.json_output = raw
    context.json_data = raw.get("data", {})
    assert isinstance(context.json_data, dict), (
        f"Expected data envelope to be a dict, got {type(context.json_data)}"
    )


@then("the data section contains the session_id")
def step_json_data_contains_session_id(context: Context) -> None:
    assert context.json_data.get("session_id") == _ULID, (
        f"Expected session_id={_ULID}, got {context.json_data.get('session_id')}"
    )


@then("the data section contains a usage object with expected keys")
def step_json_data_contains_usage_object(context: Context) -> None:
    usage = context.json_data.get("usage")
    assert isinstance(usage, dict), f"Expected usage to be a dict, got {type(usage)}"
    expected_keys = {
        "input_tokens",
        "output_tokens",
        "cost_usd",
        "duration_ms",
        "tool_calls",
    }
    missing = expected_keys - set(usage.keys())
    assert not missing, f"Usage object missing expected keys: {missing}\nUsage: {usage}"


@then("the output contains a Usage panel")
def step_output_contains_usage_panel(context: Context) -> None:
    output = context.result.output
    # Rich Usage panel is rendered with title="Usage"
    assert "Usage" in output, f"Expected Usage panel in output.\nOutput:\n{output}"
    assert "Input tokens" in output, (
        f"Expected input tokens in Usage panel.\nOutput:\n{output}"
    )
    assert "Output tokens" in output, (
        f"Expected output tokens in Usage panel.\nOutput:\n{output}"
    )