Files
awoooi/apps/api/tests/test_claude_paid_fallback.py

706 lines
23 KiB
Python

from __future__ import annotations
import json
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from src.plugins.mcp.interfaces import MCPTool, MCPToolResult
from src.services import ai_control as ai_control_module
from src.services.ai_providers.claude import ClaudeProvider
from src.services.ai_providers.gemini import GeminiProvider
from src.services.ai_providers.interfaces import cloud_context_block_reason
from src.services.ai_rate_limiter import (
AIGenerationReservation,
get_claude_pricing_policy,
)
def _settings(*, mode: str = "production", canary_percent: int = 0) -> SimpleNamespace:
return SimpleNamespace(
CLAUDE_API_KEY="configured-test-value",
CLAUDE_MODEL="claude-sonnet-5",
CLAUDE_MAX_OUTPUT_TOKENS=4096,
CLAUDE_EXECUTION_MODE=mode,
CLAUDE_CANARY_PERCENT=canary_percent,
)
def _reservation(
*,
allowed: bool = True,
reason: str = "allowed",
receipt_id: str = "gen-claude-test-1",
) -> AIGenerationReservation:
pricing = get_claude_pricing_policy("claude-sonnet-5")
assert pricing is not None
return AIGenerationReservation(
provider="claude",
receipt_id=receipt_id,
allowed=allowed,
reason=reason,
estimated_prompt_tokens=20,
estimated_completion_tokens=4096,
estimated_cost_usd=pricing.cost_usd(20, 4096),
date="2026-07-15",
model=pricing.model,
pricing_source=pricing.source,
pricing_version=pricing.version,
pricing_checked_at=pricing.checked_at,
input_usd_per_million=pricing.input_usd_per_million,
output_usd_per_million=pricing.output_usd_per_million,
trace_id="trace-claude-1",
run_id="run-claude-1",
work_item_id="work-claude-1",
)
def _context(**overrides: object) -> dict[str, object]:
return {
"trace_id": "trace-claude-1",
"run_id": "run-claude-1",
"work_item_id": "work-claude-1",
"data_classification": "sanitized",
**overrides,
}
class _ClaudeResponse:
def raise_for_status(self) -> None:
return None
def json(self) -> dict:
return {
"stop_reason": "tool_use",
"usage": {"input_tokens": 12, "output_tokens": 8},
"content": [
{
"type": "tool_use",
"name": "submit_analysis",
"input": {
"action_title": "bounded action",
"suggested_action": "NO_ACTION",
},
}
],
}
class _ClaudeToolUseResponse:
def raise_for_status(self) -> None:
return None
def json(self) -> dict:
return {
"stop_reason": "tool_use",
"usage": {"input_tokens": 12, "output_tokens": 8},
"content": [
{
"type": "tool_use",
"id": "tool-use-1",
"name": "kubernetes__kubectl_get",
"input": {"resource": "pods"},
}
],
}
class _ClaudeTerminalResponse:
def raise_for_status(self) -> None:
return None
def json(self) -> dict:
return {
"stop_reason": "end_turn",
"usage": {"input_tokens": 10, "output_tokens": 4},
"content": [{"type": "text", "text": "bounded terminal result"}],
}
class _ClaudeRefusalResponse:
def raise_for_status(self) -> None:
return None
def json(self) -> dict:
return {
"stop_reason": "refusal",
"usage": {"input_tokens": 12, "output_tokens": 8},
"content": [{"type": "text", "text": "not returned"}],
}
class _ClaudeMissingContentResponse:
def raise_for_status(self) -> None:
return None
def json(self) -> dict:
return {
"stop_reason": "end_turn",
"usage": {"input_tokens": 12, "output_tokens": 8},
"content": [],
}
def _mcp_tool() -> MCPTool:
return MCPTool(
name="kubectl_get",
description="read one Kubernetes resource",
input_schema={"type": "object", "properties": {}},
server_name="kubernetes",
)
def test_claude_sonnet_5_pricing_is_exact_and_conservative() -> None:
pricing = get_claude_pricing_policy("claude-sonnet-5")
assert pricing is not None
assert pricing.input_usd_per_million == 3.0
assert pricing.output_usd_per_million == 15.0
assert pricing.version == "anthropic-standard-ceiling-2026-07-15"
assert get_claude_pricing_policy("claude-3-5-sonnet") is None
assert get_claude_pricing_policy("claude-sonnet-latest") is None
@pytest.mark.parametrize(
("context", "reason"),
[
({}, "cloud_context_classification_missing"),
({"data_classification": "private"}, "cloud_context_data_classification_blocked"),
(
{"data_classification": "sanitized", "contains_secret": True},
"cloud_context_sensitive_marker_blocked",
),
(
{"data_classification": "sanitized", "secret_value_exposed": True},
"cloud_context_sensitive_marker_blocked",
),
(
{"data_classification": "sanitized", "raw_log_payload": "raw line"},
"cloud_context_raw_log_blocked",
),
(
{
"dataClassification": "sanitized",
"headers": {"Authorization": "opaque-test-value"},
},
"cloud_context_sensitive_marker_blocked",
),
(
{"data_classification": "sanitized", "nested": {"password": "x"}},
"cloud_context_sensitive_marker_blocked",
),
(
{"data_classification": "sanitized", "nested": {"Cookie": "x"}},
"cloud_context_sensitive_marker_blocked",
),
(
{
"data_classification": "sanitized",
"nested": {"accessToken": "x"},
},
"cloud_context_sensitive_marker_blocked",
),
(
{"data_classification": "sanitized", "nested": {"apiKey": "x"}},
"cloud_context_sensitive_marker_blocked",
),
(
{"data_classification": "sanitized", "nested": {"rawLog": "x"}},
"cloud_context_raw_log_blocked",
),
(
{"data_classification": "sanitized", "require_local": True},
"cloud_context_requires_local",
),
],
)
def test_cloud_context_gate_uses_structured_metadata_not_prompt_regex(
context: dict[str, object],
reason: str,
) -> None:
assert cloud_context_block_reason(context) == reason
@pytest.mark.asyncio
async def test_claude_shadow_mode_performs_zero_auth_reservation_or_http_calls() -> None:
provider = ClaudeProvider()
provider._get_client = AsyncMock() # type: ignore[method-assign]
disabled_readback = AsyncMock(return_value=False)
limiter_factory = MagicMock()
with patch("src.services.ai_providers.claude.settings", _settings(mode="shadow")), patch(
"src.services.ai_providers.claude.is_provider_enabled_by_env",
return_value=True,
), patch.object(
ai_control_module,
"is_provider_disabled",
disabled_readback,
), patch(
"src.services.ai_rate_limiter.get_ai_rate_limiter",
limiter_factory,
):
result = await provider.analyze("sanitized metadata", context=_context())
health = await provider.health_check()
assert result.success is False
assert result.error == "claude_shadow_metadata_only"
assert health is False
disabled_readback.assert_not_awaited()
limiter_factory.assert_not_called()
provider._get_client.assert_not_awaited() # type: ignore[attr-defined]
@pytest.mark.asyncio
@pytest.mark.parametrize(
"provider_factory",
[ClaudeProvider, GeminiProvider],
)
async def test_paid_provider_sensitive_context_blocks_before_control_cost_or_http(
provider_factory,
) -> None:
provider = provider_factory()
provider._get_client = AsyncMock() # type: ignore[method-assign]
disabled_readback = AsyncMock(return_value=False)
limiter_factory = MagicMock()
provider_module = (
"src.services.ai_providers.claude"
if provider.name == "claude"
else "src.services.ai_providers.gemini"
)
provider_settings = (
_settings(mode="production")
if provider.name == "claude"
else SimpleNamespace(GEMINI_API_KEY="configured-test-value")
)
with patch(f"{provider_module}.settings", provider_settings), patch(
f"{provider_module}.is_provider_enabled_by_env",
return_value=True,
), patch.object(
ai_control_module,
"is_provider_disabled",
disabled_readback,
), patch(
"src.services.ai_rate_limiter.get_ai_rate_limiter",
limiter_factory,
):
result = await provider.analyze(
"opaque prompt not inspected by this gate",
context=_context(contains_secret=True),
)
assert result.success is False
assert result.error == "cloud_context_sensitive_marker_blocked"
disabled_readback.assert_not_awaited()
limiter_factory.assert_not_called()
provider._get_client.assert_not_awaited() # type: ignore[attr-defined]
@pytest.mark.asyncio
async def test_claude_cost_guard_blocks_without_http_call() -> None:
client = MagicMock()
client.is_closed = False
client.post = AsyncMock()
limiter = MagicMock()
limiter.reserve_generation = AsyncMock(
return_value=_reservation(allowed=False, reason="daily_cost_limit")
)
provider = ClaudeProvider()
provider._http_client = client
with patch("src.services.ai_providers.claude.settings", _settings()), patch(
"src.services.ai_providers.claude.is_provider_enabled_by_env",
return_value=True,
), patch.object(
ai_control_module,
"is_provider_disabled",
AsyncMock(return_value=False),
), patch(
"src.services.ai_rate_limiter.get_ai_rate_limiter",
return_value=limiter,
):
result = await provider.analyze("sanitized metadata", context=_context())
assert result.success is False
assert result.error == "cost_guard_blocked:daily_cost_limit"
assert result.audit_metadata["generation_receipt_id"] == "gen-claude-test-1"
client.post.assert_not_awaited()
@pytest.mark.asyncio
async def test_claude_success_uses_exact_model_and_no_manual_sampling_or_thinking() -> None:
client = MagicMock()
client.is_closed = False
client.post = AsyncMock(return_value=_ClaudeResponse())
limiter = MagicMock()
limiter.reserve_generation = AsyncMock(return_value=_reservation())
limiter.finalize_generation = AsyncMock(return_value=True)
provider = ClaudeProvider()
provider._http_client = client
context = _context()
with patch("src.services.ai_providers.claude.settings", _settings()), patch(
"src.services.ai_providers.claude.is_provider_enabled_by_env",
return_value=True,
), patch.object(
ai_control_module,
"is_provider_disabled",
AsyncMock(return_value=False),
), patch(
"src.services.ai_rate_limiter.get_ai_rate_limiter",
return_value=limiter,
):
result = await provider.analyze("sanitized metadata", context=context)
assert result.success is True
assert json.loads(result.raw_response)["suggested_action"] == "NO_ACTION"
assert result.tokens == 20
assert result.cost_usd == _reservation().actual_cost_usd(12, 8)
assert result.audit_metadata == {
"schema_version": "paid_provider_audit_v1",
"paid_provider": True,
"provider": "claude",
"fallback_position": 4,
"execution_mode": "production",
"generation_receipt_id": "gen-claude-test-1",
"generation_receipt_ids": [],
"pricing_version": "anthropic-standard-ceiling-2026-07-15",
"canary_bucket": None,
}
request = client.post.await_args.kwargs["json"]
assert request["model"] == "claude-sonnet-5"
assert request["max_tokens"] == 4096
for forbidden in ("temperature", "top_p", "top_k", "thinking"):
assert forbidden not in request
limiter.reserve_generation.assert_awaited_once_with(
"claude",
"sanitized metadata",
model="claude-sonnet-5",
max_output_tokens=4096,
context=context,
)
limiter.finalize_generation.assert_awaited_once_with(
_reservation(),
success=True,
prompt_tokens=12,
completion_tokens=8,
)
@pytest.mark.asyncio
async def test_claude_safe_structured_tool_result_permits_next_iteration() -> None:
client = MagicMock()
client.is_closed = False
client.post = AsyncMock(
side_effect=[_ClaudeToolUseResponse(), _ClaudeTerminalResponse()]
)
limiter = MagicMock()
limiter.reserve_generation = AsyncMock(
side_effect=[
_reservation(receipt_id="gen-claude-tool-1"),
_reservation(receipt_id="gen-claude-tool-2"),
]
)
limiter.finalize_generation = AsyncMock(return_value=True)
tool_executor = AsyncMock(
return_value=MCPToolResult(
success=True,
execution_id="safe-tool-receipt",
output={
"data_classification": "sanitized",
"status": "healthy",
},
)
)
provider = ClaudeProvider()
provider._http_client = client
with patch("src.services.ai_providers.claude.settings", _settings()), patch(
"src.services.ai_providers.claude.is_provider_enabled_by_env",
return_value=True,
), patch.object(
ai_control_module,
"is_provider_disabled",
AsyncMock(return_value=False),
), patch(
"src.services.ai_rate_limiter.get_ai_rate_limiter",
return_value=limiter,
):
result = await provider.analyze_with_tools(
"sanitized investigation",
[_mcp_tool()],
tool_executor,
context=_context(),
)
assert result.success is True
assert result.raw_response == "bounded terminal result"
assert client.post.await_count == 2
second_request = client.post.await_args_list[1].kwargs["json"]
serialized_messages = json.dumps(second_request["messages"])
assert "safe-tool-receipt" in serialized_messages
assert "data_classification" in serialized_messages
assert limiter.finalize_generation.await_count == 2
@pytest.mark.asyncio
@pytest.mark.parametrize(
("response", "expected_error", "receipt_error_code"),
[
(
_ClaudeRefusalResponse(),
"claude_provider_refusal",
"provider_refusal",
),
(
_ClaudeMissingContentResponse(),
"claude_response_missing_content",
"response_missing_content",
),
],
)
async def test_claude_invalid_terminal_response_finalizes_failure_once(
response: object,
expected_error: str,
receipt_error_code: str,
) -> None:
client = MagicMock()
client.is_closed = False
client.post = AsyncMock(return_value=response)
limiter = MagicMock()
reservation = _reservation(receipt_id="gen-claude-invalid-terminal")
limiter.reserve_generation = AsyncMock(return_value=reservation)
limiter.finalize_generation = AsyncMock(return_value=True)
provider = ClaudeProvider()
provider._http_client = client
with patch("src.services.ai_providers.claude.settings", _settings()), patch(
"src.services.ai_providers.claude.is_provider_enabled_by_env",
return_value=True,
), patch.object(
ai_control_module,
"is_provider_disabled",
AsyncMock(return_value=False),
), patch(
"src.services.ai_rate_limiter.get_ai_rate_limiter",
return_value=limiter,
):
result = await provider.analyze("sanitized metadata", context=_context())
assert result.success is False
assert result.error == expected_error
limiter.finalize_generation.assert_awaited_once_with(
reservation,
success=False,
prompt_tokens=12,
completion_tokens=8,
error_code=receipt_error_code,
)
@pytest.mark.asyncio
@pytest.mark.parametrize(
("response", "expected_error", "receipt_error_code"),
[
(
_ClaudeRefusalResponse(),
"claude_provider_refusal",
"provider_refusal",
),
(
_ClaudeMissingContentResponse(),
"claude_agent_loop_response_missing_content",
"response_missing_content",
),
],
)
async def test_claude_agent_loop_invalid_response_finalizes_failure_once(
response: object,
expected_error: str,
receipt_error_code: str,
) -> None:
client = MagicMock()
client.is_closed = False
client.post = AsyncMock(return_value=response)
limiter = MagicMock()
reservation = _reservation(receipt_id="gen-claude-agent-invalid")
limiter.reserve_generation = AsyncMock(return_value=reservation)
limiter.finalize_generation = AsyncMock(return_value=True)
tool_executor = AsyncMock()
provider = ClaudeProvider()
provider._http_client = client
with patch("src.services.ai_providers.claude.settings", _settings()), patch(
"src.services.ai_providers.claude.is_provider_enabled_by_env",
return_value=True,
), patch.object(
ai_control_module,
"is_provider_disabled",
AsyncMock(return_value=False),
), patch(
"src.services.ai_rate_limiter.get_ai_rate_limiter",
return_value=limiter,
):
result = await provider.analyze_with_tools(
"sanitized investigation",
[_mcp_tool()],
tool_executor,
context=_context(),
)
assert result.success is False
assert result.error == expected_error
assert client.post.await_count == 1
tool_executor.assert_not_awaited()
limiter.finalize_generation.assert_awaited_once_with(
reservation,
success=False,
prompt_tokens=12,
completion_tokens=8,
error_code=receipt_error_code,
)
@pytest.mark.asyncio
@pytest.mark.parametrize(
("tool_output", "reason", "leak_marker"),
[
(
{"status": "LEAK_UNCLASSIFIED_TOOL_RESULT"},
"cloud_context_classification_missing",
"LEAK_UNCLASSIFIED_TOOL_RESULT",
),
(
{
"data_classification": "sanitized",
"headers": {"Authorization": "LEAK_AUTHORIZATION_TOOL_RESULT"},
},
"cloud_context_sensitive_marker_blocked",
"LEAK_AUTHORIZATION_TOOL_RESULT",
),
(
{
"data_classification": "sanitized",
"rawLog": "LEAK_RAW_LOG_TOOL_RESULT",
},
"cloud_context_raw_log_blocked",
"LEAK_RAW_LOG_TOOL_RESULT",
),
],
)
async def test_claude_tool_result_privacy_block_stops_before_second_api_call(
tool_output: dict[str, object],
reason: str,
leak_marker: str,
) -> None:
client = MagicMock()
client.is_closed = False
client.post = AsyncMock(return_value=_ClaudeToolUseResponse())
limiter = MagicMock()
reservation = _reservation(receipt_id="gen-claude-blocked-tool")
limiter.reserve_generation = AsyncMock(return_value=reservation)
limiter.finalize_generation = AsyncMock(return_value=True)
tool_executor = AsyncMock(
return_value=MCPToolResult(
success=True,
execution_id="blocked-tool-receipt",
output=tool_output,
)
)
provider = ClaudeProvider()
provider._http_client = client
with patch("src.services.ai_providers.claude.settings", _settings()), patch(
"src.services.ai_providers.claude.is_provider_enabled_by_env",
return_value=True,
), patch.object(
ai_control_module,
"is_provider_disabled",
AsyncMock(return_value=False),
), patch(
"src.services.ai_rate_limiter.get_ai_rate_limiter",
return_value=limiter,
), patch("src.services.ai_providers.claude.logger.warning") as warning:
result = await provider.analyze_with_tools(
"sanitized investigation",
[_mcp_tool()],
tool_executor,
context=_context(),
)
assert result.success is False
assert result.raw_response == ""
assert result.error == f"claude_tool_result_blocked:{reason}"
assert client.post.await_count == 1
assert limiter.reserve_generation.await_count == 1
privacy_receipt = result.audit_metadata["tool_result_privacy_receipt"]
assert privacy_receipt == {
"schema_version": "cloud_tool_result_privacy_v1",
"status": "blocked_no_forward",
"reason": reason,
"iteration": 1,
"provider_call_count": 1,
"paid_iteration_terminal": "failed",
"receipt_finalized": True,
"content_logged": False,
"content_forwarded": False,
}
limiter.finalize_generation.assert_awaited_once_with(
reservation,
success=False,
prompt_tokens=12,
completion_tokens=8,
error_code="tool_result_privacy_blocked",
)
assert leak_marker not in json.dumps(result.to_dict())
assert leak_marker not in str(warning.call_args)
assert leak_marker not in str(client.post.await_args_list)
@pytest.mark.asyncio
async def test_claude_tool_executor_failure_finalizes_once_without_error_leak() -> None:
leak_marker = "LEAK_TOOL_EXECUTOR_EXCEPTION_BODY"
client = MagicMock()
client.is_closed = False
client.post = AsyncMock(return_value=_ClaudeToolUseResponse())
limiter = MagicMock()
reservation = _reservation(receipt_id="gen-claude-tool-executor-failed")
limiter.reserve_generation = AsyncMock(return_value=reservation)
limiter.finalize_generation = AsyncMock(return_value=True)
tool_executor = AsyncMock(side_effect=RuntimeError(leak_marker))
provider = ClaudeProvider()
provider._http_client = client
with patch("src.services.ai_providers.claude.settings", _settings()), patch(
"src.services.ai_providers.claude.is_provider_enabled_by_env",
return_value=True,
), patch.object(
ai_control_module,
"is_provider_disabled",
AsyncMock(return_value=False),
), patch(
"src.services.ai_rate_limiter.get_ai_rate_limiter",
return_value=limiter,
), patch("src.services.ai_providers.claude.logger.warning") as warning:
result = await provider.analyze_with_tools(
"sanitized investigation",
[_mcp_tool()],
tool_executor,
context=_context(),
)
assert result.success is False
assert result.raw_response == ""
assert result.error == "claude_tool_executor_failed"
assert client.post.await_count == 1
limiter.finalize_generation.assert_awaited_once_with(
reservation,
success=False,
prompt_tokens=12,
completion_tokens=8,
error_code="tool_executor_failed",
)
assert leak_marker not in json.dumps(result.to_dict())
assert leak_marker not in str(warning.call_args)