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"}], } @pytest.mark.asyncio async def test_claude_stale_control_admission_blocks_before_reservation_or_http() -> None: client = MagicMock() client.is_closed = False client.post = AsyncMock() limiter = MagicMock() limiter.reserve_generation = AsyncMock() disabled_readback = 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", disabled_readback, ), 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 == "claude_disabled_by_runtime_control" disabled_readback.assert_awaited_once_with( "claude", run_id="run-claude-1", ) limiter.reserve_generation.assert_not_awaited() client.post.assert_not_awaited() 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", GEMINI_EXECUTION_MODE="production", GEMINI_CANARY_PERCENT=0, ) ) 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 assert request["tools"][0]["strict"] is True schema = request["tools"][0]["input_schema"] assert schema["additionalProperties"] is False assert schema["properties"]["blast_radius"]["additionalProperties"] is False assert schema["properties"]["confidence"]["minimum"] == 0.0 assert schema["properties"]["confidence"]["maximum"] == 1.0 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)