tests/test_claude_sdk_step.py

"""Tests for the Claude SDK integration step."""

from __future__ import annotations

import logging
from dataclasses import dataclass
from types import SimpleNamespace
from typing import Any, List, Optional
from unittest.mock import MagicMock, patch

import pytest
from pydantic import ValidationError

from ace.core.context import ACEStepContext, SkillbookView
from ace.core.skillbook import Skillbook
from ace.integrations.claude_sdk import (
    ClaudeSDKExecuteStep,
    ClaudeSDKResult,
    ClaudeSDKToTrace,
    ToolCall,
)

# ------------------------------------------------------------------ #
# Helpers — mock Anthropic API objects
# ------------------------------------------------------------------ #


def _make_text_block(text: str) -> SimpleNamespace:
    return SimpleNamespace(type="text", text=text)


def _make_tool_block(
    tool_id: str = "toolu_01", name: str = "calculator", inp: Any = None
) -> SimpleNamespace:
    return SimpleNamespace(type="tool_use", id=tool_id, name=name, input=inp or {})


def _make_usage(input_tokens: int = 100, output_tokens: int = 50) -> SimpleNamespace:
    return SimpleNamespace(input_tokens=input_tokens, output_tokens=output_tokens)


def _make_response(
    content: list | None = None,
    model: str = "claude-sonnet-4-20250514",
    stop_reason: str = "end_turn",
    input_tokens: int = 100,
    output_tokens: int = 50,
) -> SimpleNamespace:
    if content is None:
        content = [_make_text_block("Hello, world!")]
    return SimpleNamespace(
        content=content,
        model=model,
        stop_reason=stop_reason,
        usage=_make_usage(input_tokens, output_tokens),
    )


def _make_mock_client(response: SimpleNamespace | None = None) -> MagicMock:
    client = MagicMock()
    client.messages.create.return_value = response or _make_response()
    return client


@dataclass
class FakeSample:
    question: str = "What is 2+2?"
    context: str = "math quiz"
    ground_truth: str = "4"
    metadata: dict = None  # type: ignore[assignment]

    def __post_init__(self) -> None:
        if self.metadata is None:
            self.metadata = {}


# ------------------------------------------------------------------ #
# ClaudeSDKResult
# ------------------------------------------------------------------ #


class TestClaudeSDKResult:
    def test_defaults(self):
        r = ClaudeSDKResult(task="test", success=True)
        assert r.task == "test"
        assert r.success is True
        assert r.output == ""
        assert r.error is None
        assert r.input_tokens == 0
        assert r.output_tokens == 0
        assert r.total_tokens == 0
        assert r.latency_seconds == 0.0
        assert r.tool_calls == []
        assert r.cited_skill_ids == []
        assert r.raw_response is None

    def test_full(self):
        tc = ToolCall(id="toolu_01", name="calc", input={"expr": "1+1"})
        r = ClaudeSDKResult(
            task="hello",
            success=True,
            output="world",
            model="claude-sonnet-4-20250514",
            stop_reason="end_turn",
            input_tokens=100,
            output_tokens=50,
            total_tokens=150,
            latency_seconds=1.5,
            tool_calls=[tc],
            cited_skill_ids=["math-001"],
        )
        assert r.total_tokens == 150
        assert r.tool_calls[0].name == "calc"

    def test_auto_compute_total_tokens(self):
        """total_tokens is auto-computed from input + output when left at 0."""
        r = ClaudeSDKResult(
            task="test", success=True, input_tokens=100, output_tokens=50
        )
        assert r.total_tokens == 150

    def test_explicit_total_not_overwritten(self):
        """An explicitly set total_tokens is preserved."""
        r = ClaudeSDKResult(
            task="test",
            success=True,
            input_tokens=100,
            output_tokens=50,
            total_tokens=999,
        )
        assert r.total_tokens == 999

    def test_negative_tokens_rejected(self):
        """Negative token counts should be rejected by validation."""
        with pytest.raises(Exception):
            ClaudeSDKResult(task="test", success=True, input_tokens=-1)

    def test_negative_latency_rejected(self):
        """Negative latency should be rejected by validation."""
        with pytest.raises(Exception):
            ClaudeSDKResult(task="test", success=True, latency_seconds=-0.1)

    def test_serialization(self):
        """Result should serialise to dict/JSON (raw_response excluded)."""
        r = ClaudeSDKResult(
            task="test",
            success=True,
            input_tokens=10,
            output_tokens=5,
            raw_response=object(),
        )
        d = r.model_dump()
        assert d["task"] == "test"
        assert "raw_response" not in d  # excluded

    def test_tool_call_validation(self):
        """ToolCall should validate required fields."""
        tc = ToolCall(id="toolu_01", name="calc")
        assert tc.input == {}

        with pytest.raises(Exception):
            ToolCall(name="calc")  # missing id


# ------------------------------------------------------------------ #
# ClaudeSDKExecuteStep — contracts
# ------------------------------------------------------------------ #


class TestClaudeSDKExecuteStepContracts:
    def test_requires_and_provides(self):
        client = _make_mock_client()
        step = ClaudeSDKExecuteStep(client=client)
        assert "sample" in step.requires
        assert "skillbook" in step.requires
        assert "trace" in step.provides

    def test_not_available_raises_without_client(self):
        with patch("ace.integrations.claude_sdk.ANTHROPIC_SDK_AVAILABLE", False):
            with pytest.raises(ImportError, match="anthropic SDK not installed"):
                ClaudeSDKExecuteStep()

    def test_injected_client_skips_availability_check(self):
        with patch("ace.integrations.claude_sdk.ANTHROPIC_SDK_AVAILABLE", False):
            step = ClaudeSDKExecuteStep(client=_make_mock_client())
            assert "sample" in step.requires

    def test_invalid_max_tokens_rejected(self):
        with pytest.raises(ValidationError):
            ClaudeSDKExecuteStep(client=_make_mock_client(), max_tokens=0)

    def test_invalid_temperature_rejected(self):
        with pytest.raises(ValidationError):
            ClaudeSDKExecuteStep(client=_make_mock_client(), temperature=1.5)


# ------------------------------------------------------------------ #
# ClaudeSDKExecuteStep — execution
# ------------------------------------------------------------------ #


class TestClaudeSDKExecuteStepExecution:
    def test_basic_call(self):
        response = _make_response(
            content=[_make_text_block("The answer is 4")],
            input_tokens=80,
            output_tokens=20,
        )
        client = _make_mock_client(response)
        step = ClaudeSDKExecuteStep(client=client, model="claude-sonnet-4-20250514")

        ctx = ACEStepContext(sample="What is 2+2?", skillbook=None)
        result_ctx = step(ctx)

        r: ClaudeSDKResult = result_ctx.trace  # type: ignore[assignment]
        assert r.success is True
        assert r.output == "The answer is 4"
        assert r.input_tokens == 80
        assert r.output_tokens == 20
        assert r.total_tokens == 100
        assert r.latency_seconds >= 0
        assert r.model == "claude-sonnet-4-20250514"
        assert r.stop_reason == "end_turn"

        # Verify API was called correctly
        call_kwargs = client.messages.create.call_args[1]
        assert call_kwargs["model"] == "claude-sonnet-4-20250514"
        assert call_kwargs["messages"] == [{"role": "user", "content": "What is 2+2?"}]

    def test_with_system_prompt(self):
        client = _make_mock_client()
        step = ClaudeSDKExecuteStep(
            client=client,
            system_prompt="You are a math tutor.",
        )

        ctx = ACEStepContext(sample="What is 2+2?", skillbook=None)
        step(ctx)

        call_kwargs = client.messages.create.call_args[1]
        assert call_kwargs["system"] == "You are a math tutor."

    def test_skillbook_injection(self):
        client = _make_mock_client()
        step = ClaudeSDKExecuteStep(client=client, inject_skillbook=True)

        sb = Skillbook()
        sb.add_skill("math", "Always show your work")
        ctx = ACEStepContext(sample="What is 2+2?", skillbook=SkillbookView(sb))
        step(ctx)

        call_kwargs = client.messages.create.call_args[1]
        assert "system" in call_kwargs
        assert "Strategic Knowledge" in call_kwargs["system"]

    def test_skillbook_injection_with_system_prompt(self):
        client = _make_mock_client()
        step = ClaudeSDKExecuteStep(
            client=client,
            system_prompt="You are a tutor.",
            inject_skillbook=True,
        )

        sb = Skillbook()
        sb.add_skill("math", "Show work")
        ctx = ACEStepContext(sample="test", skillbook=SkillbookView(sb))
        step(ctx)

        call_kwargs = client.messages.create.call_args[1]
        system = call_kwargs["system"]
        assert "You are a tutor." in system
        assert "Strategic Knowledge" in system

    def test_skillbook_injection_disabled(self):
        client = _make_mock_client()
        step = ClaudeSDKExecuteStep(client=client, inject_skillbook=False)

        sb = Skillbook()
        sb.add_skill("math", "Show work")
        ctx = ACEStepContext(sample="test", skillbook=SkillbookView(sb))
        step(ctx)

        call_kwargs = client.messages.create.call_args[1]
        assert "system" not in call_kwargs

    def test_empty_skillbook_no_system(self):
        client = _make_mock_client()
        step = ClaudeSDKExecuteStep(client=client, inject_skillbook=True)

        sb = Skillbook()
        ctx = ACEStepContext(sample="test", skillbook=SkillbookView(sb))
        step(ctx)

        call_kwargs = client.messages.create.call_args[1]
        assert "system" not in call_kwargs

    def test_with_tools(self):
        tools = [
            {
                "name": "calculator",
                "description": "A calculator",
                "input_schema": {
                    "type": "object",
                    "properties": {"expr": {"type": "string"}},
                },
            }
        ]
        response = _make_response(
            content=[
                _make_tool_block("toolu_01", "calculator", {"expr": "2+2"}),
                _make_text_block("The result is 4"),
            ]
        )
        client = _make_mock_client(response)
        step = ClaudeSDKExecuteStep(client=client, tools=tools)

        ctx = ACEStepContext(sample="Calculate 2+2", skillbook=None)
        result_ctx = step(ctx)

        r: ClaudeSDKResult = result_ctx.trace  # type: ignore[assignment]
        assert r.success is True
        assert r.output == "The result is 4"
        assert len(r.tool_calls) == 1
        assert r.tool_calls[0].name == "calculator"
        assert r.tool_calls[0].input == {"expr": "2+2"}

        call_kwargs = client.messages.create.call_args[1]
        assert call_kwargs["tools"] == tools

    def test_api_error_handled(self):
        client = _make_mock_client()
        client.messages.create.side_effect = RuntimeError("API down")
        step = ClaudeSDKExecuteStep(client=client)

        ctx = ACEStepContext(sample="test", skillbook=None)
        result_ctx = step(ctx)

        r: ClaudeSDKResult = result_ctx.trace  # type: ignore[assignment]
        assert r.success is False
        assert "API down" in r.error
        assert r.latency_seconds >= 0

    def test_temperature_and_max_tokens(self):
        client = _make_mock_client()
        step = ClaudeSDKExecuteStep(client=client, temperature=0.7, max_tokens=1024)

        ctx = ACEStepContext(sample="test", skillbook=None)
        step(ctx)

        call_kwargs = client.messages.create.call_args[1]
        assert call_kwargs["temperature"] == 0.7
        assert call_kwargs["max_tokens"] == 1024


# ------------------------------------------------------------------ #
# ClaudeSDKExecuteStep — task extraction
# ------------------------------------------------------------------ #


class TestClaudeSDKTaskExtraction:
    def test_string_sample(self):
        assert ClaudeSDKExecuteStep._extract_task("hello") == "hello"

    def test_sample_with_question(self):
        sample = FakeSample(question="What is 2+2?", context="")
        result = ClaudeSDKExecuteStep._extract_task(sample)
        assert result == "What is 2+2?"

    def test_sample_with_context(self):
        sample = FakeSample(question="What is 2+2?", context="math quiz")
        result = ClaudeSDKExecuteStep._extract_task(sample)
        assert "What is 2+2?" in result
        assert "Context: math quiz" in result

    def test_arbitrary_object(self):
        result = ClaudeSDKExecuteStep._extract_task(42)
        assert result == "42"


# ------------------------------------------------------------------ #
# ClaudeSDKExecuteStep — skill ID extraction
# ------------------------------------------------------------------ #


class TestClaudeSDKSkillExtraction:
    def test_extracts_skill_ids(self):
        response = _make_response(
            content=[
                _make_text_block(
                    "Following [math-00001], the answer is 4. "
                    "Also [general-00042] applies."
                )
            ]
        )
        client = _make_mock_client(response)
        step = ClaudeSDKExecuteStep(client=client)

        ctx = ACEStepContext(sample="test", skillbook=None)
        result_ctx = step(ctx)

        r: ClaudeSDKResult = result_ctx.trace  # type: ignore[assignment]
        assert "math-00001" in r.cited_skill_ids
        assert "general-00042" in r.cited_skill_ids

    def test_no_skill_ids(self):
        response = _make_response(content=[_make_text_block("No citations here")])
        client = _make_mock_client(response)
        step = ClaudeSDKExecuteStep(client=client)

        ctx = ACEStepContext(sample="test", skillbook=None)
        result_ctx = step(ctx)

        r: ClaudeSDKResult = result_ctx.trace  # type: ignore[assignment]
        assert r.cited_skill_ids == []


# ------------------------------------------------------------------ #
# ClaudeSDKExecuteStep — observability logging
# ------------------------------------------------------------------ #


class TestClaudeSDKObservability:
    def test_auto_instruments_anthropic_when_logfire_configured(self):
        mock_logfire = MagicMock()
        mock_ctx = MagicMock()
        mock_logfire.instrument_anthropic.return_value = mock_ctx

        with (
            patch("ace.observability.is_configured", return_value=True),
            patch.dict("sys.modules", {"logfire": mock_logfire}),
        ):
            step = ClaudeSDKExecuteStep(client=_make_mock_client())

        mock_logfire.instrument_anthropic.assert_called_once_with(step._client)
        mock_ctx.__enter__.assert_not_called()

    def test_logs_metrics(self, caplog):
        response = _make_response(input_tokens=200, output_tokens=100)
        client = _make_mock_client(response)
        step = ClaudeSDKExecuteStep(client=client)

        with caplog.at_level(logging.INFO, logger="ace.integrations.claude_sdk"):
            ctx = ACEStepContext(sample="test", skillbook=None)
            step(ctx)

        assert "ClaudeSDK:" in caplog.text
        assert "tokens=" in caplog.text

    def test_logs_error(self, caplog):
        client = _make_mock_client()
        client.messages.create.side_effect = RuntimeError("boom")
        step = ClaudeSDKExecuteStep(client=client)

        with caplog.at_level(logging.ERROR, logger="ace.integrations.claude_sdk"):
            ctx = ACEStepContext(sample="test", skillbook=None)
            step(ctx)

        assert "failed" in caplog.text.lower()

    def test_logfire_span_on_success(self):
        """When Logfire is configured, __call__ opens a span with attributes."""
        mock_span = MagicMock()
        mock_logfire = MagicMock()
        mock_logfire.span.return_value.__enter__ = MagicMock(return_value=mock_span)
        mock_logfire.span.return_value.__exit__ = MagicMock(return_value=False)

        response = _make_response(input_tokens=50, output_tokens=25)
        client = _make_mock_client(response)
        step = ClaudeSDKExecuteStep(client=client)

        with (
            patch(
                "ace.integrations.claude_sdk._get_logfire", return_value=mock_logfire
            ),
        ):
            ctx = ACEStepContext(sample="What is 2+2?", skillbook=None)
            result_ctx = step(ctx)

        # Span was opened
        mock_logfire.span.assert_called_once()
        call_kwargs = mock_logfire.span.call_args
        assert call_kwargs[0][0] == "ClaudeSDKExecuteStep"
        assert call_kwargs[1]["model"] == "claude-sonnet-4-20250514"

        # Attributes were set on the span
        attr_calls = {c[0][0]: c[0][1] for c in mock_span.set_attribute.call_args_list}
        assert attr_calls["success"] is True
        assert attr_calls["input_tokens"] == 50
        assert attr_calls["output_tokens"] == 25
        assert attr_calls["total_tokens"] == 75
        assert "error" not in attr_calls

        # logfire.info was called with metrics
        mock_logfire.info.assert_called_once()
        info_kwargs = mock_logfire.info.call_args[1]
        assert info_kwargs["input_tokens"] == 50
        assert info_kwargs["output_tokens"] == 25

    def test_logfire_span_on_failure(self):
        """On API error, span captures error attribute and logfire.error is called."""
        mock_span = MagicMock()
        mock_logfire = MagicMock()
        mock_logfire.span.return_value.__enter__ = MagicMock(return_value=mock_span)
        mock_logfire.span.return_value.__exit__ = MagicMock(return_value=False)

        client = _make_mock_client()
        client.messages.create.side_effect = RuntimeError("rate limited")
        step = ClaudeSDKExecuteStep(client=client)

        with (
            patch(
                "ace.integrations.claude_sdk._get_logfire", return_value=mock_logfire
            ),
        ):
            ctx = ACEStepContext(sample="test", skillbook=None)
            step(ctx)

        # Span captured the error attribute
        attr_calls = {c[0][0]: c[0][1] for c in mock_span.set_attribute.call_args_list}
        assert attr_calls["success"] is False
        assert "rate limited" in attr_calls["error"]

        # logfire.error was called
        mock_logfire.error.assert_called_once()
        error_kwargs = mock_logfire.error.call_args[1]
        assert "rate limited" in error_kwargs["error"]

    def test_no_logfire_noop(self):
        """When Logfire is not configured, execution proceeds without spans."""
        client = _make_mock_client()
        step = ClaudeSDKExecuteStep(client=client)

        with patch("ace.integrations.claude_sdk._get_logfire", return_value=None):
            ctx = ACEStepContext(sample="test", skillbook=None)
            result_ctx = step(ctx)

        r: ClaudeSDKResult = result_ctx.trace  # type: ignore[assignment]
        assert r.success is True


# ------------------------------------------------------------------ #
# ClaudeSDKToTrace
# ------------------------------------------------------------------ #


class TestClaudeSDKToTrace:
    def test_requires_and_provides(self):
        step = ClaudeSDKToTrace()
        assert "trace" in step.requires
        assert "trace" in step.provides

    def test_success_trace(self):
        r = ClaudeSDKResult(
            task="What is 2+2?",
            success=True,
            output="4",
            model="claude-sonnet-4-20250514",
            stop_reason="end_turn",
            input_tokens=100,
            output_tokens=50,
            total_tokens=150,
            latency_seconds=1.2,
            cited_skill_ids=["math-001"],
        )
        ctx = ACEStepContext(trace=r)
        result_ctx = ClaudeSDKToTrace()(ctx)

        trace = result_ctx.trace
        assert trace["question"] == "What is 2+2?"
        assert trace["answer"] == "4"
        assert trace["skill_ids"] == ["math-001"]
        assert "succeeded" in trace["reasoning"]
        assert "claude-sonnet-4-20250514" in trace["reasoning"]
        assert "100" in trace["reasoning"]  # input tokens
        assert "50" in trace["reasoning"]  # output tokens
        assert "1.2" in trace["reasoning"]  # latency
        assert "succeeded" in trace["feedback"]
        assert trace["ground_truth"] is None

    def test_failure_trace(self):
        r = ClaudeSDKResult(
            task="fail",
            success=False,
            error="API timeout",
            model="claude-sonnet-4-20250514",
            latency_seconds=30.0,
        )
        ctx = ACEStepContext(trace=r)
        result_ctx = ClaudeSDKToTrace()(ctx)

        trace = result_ctx.trace
        assert trace["question"] == "fail"
        assert trace["answer"] == ""
        assert "failed" in trace["reasoning"]
        assert "API timeout" in trace["reasoning"]
        assert "failed" in trace["feedback"]
        assert "API timeout" in trace["feedback"]

    def test_tool_calls_in_reasoning(self):
        r = ClaudeSDKResult(
            task="calc",
            success=True,
            output="4",
            model="claude-sonnet-4-20250514",
            tool_calls=[
                ToolCall(id="toolu_01", name="calculator", input={"expr": "2+2"}),
                ToolCall(id="toolu_02", name="formatter"),
            ],
        )
        ctx = ACEStepContext(trace=r)
        result_ctx = ClaudeSDKToTrace()(ctx)

        trace = result_ctx.trace
        assert "Tool calls (2)" in trace["reasoning"]
        assert "calculator" in trace["reasoning"]
        assert "formatter" in trace["reasoning"]