"""Tests for Task A31: Preflight Dependency Agent, Workflow Run State, Local Concurrency & Context Window, PII Masking, and Cost Controller Honesty. """ from __future__ import annotations import json from pathlib import Path from unittest.mock import MagicMock, patch import pytest from antigravity_provider import paths from antigravity_provider.router.action_handler import ActionExecutor from antigravity_provider.router.adapters.local_adapter import LocalLLMAdapter from antigravity_provider.router.preflight_service import PreflightCheckService, PreflightItem, PreflightReport from antigravity_provider.router.role_registry import CANONICAL_ROLES, RoleRegistry from antigravity_provider.router.router_config import ( RouterConfig, RouterProfileConfig, get_default_router_config, load_router_config, ) from antigravity_provider.router.settings_service import ( DEFAULT_SETTINGS, get_hub_settings, invalidate_settings_cache, save_hub_settings, ) from antigravity_provider.router.telemetry_service import ( TelemetryAggregates, TelemetryRecord, TelemetryService, format_token_count, ) from antigravity_provider.router.web.server import sanitize_snapshot from antigravity_provider.router.workflow_service import ( AgentDefinition, WorkflowDefinition, WorkflowExecutionService, WorkflowService, get_last_run_state, sanitize_run_data, ) # ============================================================================ # P0-1: Preflight Dependency Agent # ============================================================================ def test_dependency_agent_role_registered(): """Verify 13th role 'dependency-agent' and its canonical aliases in RoleRegistry.""" assert len(CANONICAL_ROLES) == 13 assert "dependency-agent" in CANONICAL_ROLES role_def = CANONICAL_ROLES["dependency-agent"] assert role_def.role_id == "dependency-agent" assert role_def.display_name_ru == "Проверяющий готовность" assert role_def.short_name_ru == "Готовность" assert role_def.is_implemented is True assert "preflight" in role_def.capabilities # Test alias resolution aliases = [ "dependency-agent", "dependency_agent", "preflight", "проверяющий готовность", "агент зависимостей", "готовность", "dependency", ] for alias in aliases: canonical = RoleRegistry.resolve_role_name(alias) assert canonical == "dependency-agent", f"Alias '{alias}' resolved to '{canonical}'" def test_preflight_service_cli_and_environment(): """Verify CLI tools and environment checks with controlled discovery.""" service = PreflightCheckService.get() with patch("antigravity_provider.router.preflight_service.shutil.which", return_value="/opt/agy"), \ patch("antigravity_provider.router.preflight_service.importlib.util.find_spec", return_value=object()): cli_items = service.check_cli_dependencies() assert len(cli_items) >= 3 ids = {item.check_id for item in cli_items} assert "cli_agy" in ids assert "pkg_fastapi" in ids assert "pkg_uvicorn" in ids assert all(item.status == "PASS" for item in cli_items) env_items = service.check_system_environment() assert len(env_items) >= 3 env_ids = {item.check_id for item in env_items} assert "env_hermes_home" in env_ids assert "env_config_writable" in env_ids assert "env_logs_writable" in env_ids def _fake_preflight_items() -> list: return [ PreflightItem(check_id="cli_agy", name="CLI", status="PASS", message="mocked"), PreflightItem(check_id="env_hermes_home", name="HOME", status="PASS", message="mocked"), PreflightItem(check_id="auth_local-1", name="AUTH", status="WARN", message="mocked"), PreflightItem(check_id="local_srv_local-1", name="LLM", status="PASS", message="mocked"), ] def test_preflight_service_run_all_and_action(): """Verify run_all_checks orchestration without probing this machine.""" service = PreflightCheckService.get() cli, env, auth, local = ( [_fake_preflight_items()[0]], [_fake_preflight_items()[1]], [_fake_preflight_items()[2]], [_fake_preflight_items()[3]], ) with patch.object(service, "check_cli_dependencies", return_value=cli) as mock_cli, \ patch.object(service, "check_system_environment", return_value=env) as mock_env, \ patch.object(service, "check_auth_credentials", return_value=auth) as mock_auth, \ patch.object(service, "check_local_servers", return_value=local) as mock_local, \ patch("antigravity_provider.router.preflight_service.urllib.request.urlopen") as mock_urlopen: report = service.run_all_checks() action_res = ActionExecutor.execute("run_preflight", {}) mock_cli.assert_called() mock_env.assert_called() mock_auth.assert_called() mock_local.assert_called() mock_urlopen.assert_not_called() assert isinstance(report, PreflightReport) assert report.passed_count == 3 assert report.warn_count == 1 assert report.failed_count == 0 assert len(report.checks) == 4 report_dict = report.to_dict() assert "success" in report_dict assert "checks" in report_dict assert isinstance(report_dict["checks"], list) assert "ok" in action_res assert "message" in action_res assert "data" in action_res assert "checks" in action_res["data"] assert mock_urlopen.call_count == 0 @pytest.mark.live def test_preflight_live_local_servers(): """Optional live probe of configured local servers. Not part of hermetic pytest.""" service = PreflightCheckService.get() items = service.check_local_servers() assert isinstance(items, list) assert items # ============================================================================ # P0-2: Workflow Run State Manager # ============================================================================ def test_workflow_run_state_sanitization(): """Verify recursive secret stripping in workflow run state.""" raw_state = { "run_id": "test-run-123", "status": "RUNNING", "api_key": "sk-1234567890abcdef", "token": "gho_secrettoken123456", "nested": { "password": "supersecretpass", "auth_status": "ok", "message": "Bearer eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIn0", "extra_url": "https://example.com/callback?access_token=secret12345&foo=bar", }, "step_list": [ {"account": "acc-1", "client_secret": "my-client-secret-123"}, {"safe_field": "public_data"}, ], } sanitized = sanitize_run_data(raw_state) assert sanitized["api_key"] == "***" assert sanitized["token"] == "***" assert sanitized["nested"]["password"] == "***" assert sanitized["nested"]["auth_status"] == "ok" assert "Bearer ***" in sanitized["nested"]["message"] assert "access_token=***" in sanitized["nested"]["extra_url"] assert sanitized["step_list"][0]["client_secret"] == "***" assert sanitized["step_list"][1]["safe_field"] == "public_data" def test_workflow_run_state_interrupted_on_startup(tmp_path: Path): """Verify that a RUNNING state in workflow_run_state.json transitions to INTERRUPTED on reload.""" state_file = tmp_path / "workflow_state.json" run_state_file = tmp_path / "workflow_run_state.json" # Pre-populate run state with RUNNING status initial_run_state = { "run_id": "run-crash-test", "status": "RUNNING", "started_at": "2026-08-26T00:00:00Z", "updated_at": "2026-08-26T00:00:00Z", "current_step_index": 2, "current_agent_id": "developer-1", "iteration_count": 1, "completed_steps": [ {"step_index": 0, "agent_id": "manager", "status": "SUCCESS"}, {"step_index": 1, "agent_id": "developer-1", "status": "WORKING"}, ], "interruption_reason": None, } run_state_file.write_text(json.dumps(initial_run_state), encoding="utf-8") # Initialize WorkflowService service = WorkflowService(state_path=state_file, run_state_path=run_state_file) # Check that state transitioned to INTERRUPTED last_state = service.get_last_run_state() assert last_state is not None assert last_state["status"] == "INTERRUPTED" assert last_state["interruption_reason"] == "Прогон был прерван перезапуском сервера или сбоем процесса" assert len(last_state["completed_steps"]) == 2 # Verify top-level function assert get_last_run_state(run_state_file)["status"] == "INTERRUPTED" def test_workflow_execution_service_alias(): """Verify WorkflowExecutionService is an alias of WorkflowService.""" assert WorkflowExecutionService is WorkflowService # ============================================================================ # P0-3: Local Concurrency & Context Window # ============================================================================ def test_local_profile_max_concurrency_is_one(): """Verify all local provider profiles have max_concurrency = 1.""" config = get_default_router_config() for pid, pcfg in config.profiles.items(): if pcfg.provider == "local": assert pcfg.max_concurrency == 1, f"Local profile {pid} has max_concurrency={pcfg.max_concurrency}" # Verify loaded config also enforces max_concurrency = 1 for local profiles loaded = load_router_config() for pid, pcfg in loaded.profiles.items(): if pcfg.provider == "local": assert pcfg.max_concurrency == 1 def test_local_adapter_get_context_window(): """Verify LocalLLMAdapter retrieves context window accurately without hallucinating defaults.""" adapter = LocalLLMAdapter() # Profile with explicit context_window in auth_config prof_with_cfg = RouterProfileConfig( profile_id="local-test-1", provider="local", account_id="acc-1", auth_config={"context_window": 8192}, ) assert adapter.get_context_window(prof_with_cfg) == 8192 # Profile without context length and with non-responding server prof_empty = RouterProfileConfig( profile_id="local-test-2", provider="local", account_id="acc-2", custom_base_url="http://127.0.0.1:9999/v1", ) # Must return None instead of inventing fake numbers assert adapter.get_context_window(prof_empty) is None def test_local_adapter_context_truncation_guard(): """Verify context truncation guard protects against VRAM overflow when context_window is known.""" adapter = LocalLLMAdapter() prof = RouterProfileConfig( profile_id="local-small-ctx", provider="local", account_id="acc-1", auth_config={"context_window": 500}, custom_base_url="http://127.0.0.1:12345/v1", ) # Huge prompt exceeding 500 tokens long_middle_content = "important historical dialogue step " * 100 messages = [ {"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": "Initial prompt 1"}, {"role": "assistant", "content": long_middle_content}, {"role": "user", "content": "Initial prompt 2"}, {"role": "assistant", "content": long_middle_content}, {"role": "user", "content": "Latest user task to execute."}, ] mock_resp = { "choices": [{"message": {"role": "assistant", "content": "Truncated prompt executed successfully."}}], "usage": {"prompt_tokens": 200, "completion_tokens": 10}, } with patch("urllib.request.urlopen") as mock_urlopen: mock_cm = MagicMock() mock_cm.read.return_value = json.dumps(mock_resp).encode("utf-8") mock_urlopen.return_value.__enter__.return_value = mock_cm res = adapter.invoke(prof, {"messages": messages, "max_tokens": 100}) assert res["choices"][0]["message"]["content"] == "Truncated prompt executed successfully." # Verify sent payload messages were truncated args, kwargs = mock_urlopen.call_args sent_req = args[0] sent_body = json.loads(sent_req.data.decode("utf-8")) sent_messages = sent_body["messages"] assert sent_messages[0]["role"] == "system" assert sent_messages[-1]["content"] == "Latest user task to execute." # Total count of messages should be pruned assert len(sent_messages) < len(messages) # ============================================================================ # P0-4: PII Email Masking # ============================================================================ def test_settings_email_masking_mode(): """Verify email_masking_mode in default settings and persistence.""" assert DEFAULT_SETTINGS["email_masking_mode"] == "none" settings = get_hub_settings() assert settings.get("email_masking_mode") in ("none", "partial", "full") def test_sanitize_snapshot_email_masking_modes(): """Verify email masking behavior across 'none', 'partial', and 'full' modes.""" snapshot_data = { "user_email": "vasya.pupkin@example.com", "account_id": "google-user-1", "api_key": "sk-secret123456789", "nested": { "developer": "developer.one@domain.org", "reviewer": "r@test.com", }, } # 1. Mode: none (emails unchanged, secrets masked) san_none = sanitize_snapshot(snapshot_data, email_masking_mode="none") assert san_none["user_email"] == "vasya.pupkin@example.com" assert san_none["nested"]["developer"] == "developer.one@domain.org" assert san_none["nested"]["reviewer"] == "r@test.com" assert "api_key" not in san_none # 2. Mode: partial (preserves first and last char of local part + domain for differentiation) san_partial = sanitize_snapshot(snapshot_data, email_masking_mode="partial") assert san_partial["user_email"] == "v***n@example.com" assert san_partial["nested"]["developer"] == "d***e@domain.org" assert san_partial["nested"]["reviewer"] == "r***@test.com" assert "api_key" not in san_partial # 3. Mode: full (***@***.***) san_full = sanitize_snapshot(snapshot_data, email_masking_mode="full") assert san_full["user_email"] == "***@***.***" assert san_full["nested"]["developer"] == "***@***.***" assert san_full["nested"]["reviewer"] == "***@***.***" assert "api_key" not in san_full # ============================================================================ # P0-5: Cost Controller Token Honesty # ============================================================================ def test_telemetry_measured_vs_estimated_tokens(tmp_path: Path): """Verify telemetry distinguishes measured exact tokens from estimated tokens with ~.""" log_file = tmp_path / "telemetry_test.jsonl" service = TelemetryService(log_path=log_file) # 1. Record measured call rec1 = service.record_call( role="developer-1", profile_id="ag-w1", provider="antigravity", model="claude-3-7-sonnet", outcome="success", latency_seconds=1.25, prompt_tokens_measured=500, completion_tokens_measured=150, is_estimated=False, ) assert rec1.prompt_tokens_measured == 500 assert rec1.prompt_tokens_estimated is None assert rec1.is_estimated is False assert rec1.total_tokens == 650 # 2. Record estimated call rec2 = service.record_call( role="tester", profile_id="local-1", provider="local", model="Qwen3.8-27B-Q4_K_M.gguf", outcome="success", latency_seconds=0.85, prompt_tokens_estimated=300, completion_tokens_estimated=50, is_estimated=True, ) assert rec2.prompt_tokens_measured is None assert rec2.prompt_tokens_estimated == 300 assert rec2.is_estimated is True assert rec2.total_tokens == 350 # 3. Aggregates for measured only agg_measured = service.get_aggregates(profile_id="ag-w1") assert agg_measured.total_tokens_measured == 650 assert agg_measured.tokens_display == "650" assert agg_measured.has_estimated_tokens is False # 4. Aggregates for estimated only agg_est = service.get_aggregates(profile_id="local-1") assert agg_est.total_tokens_estimated == 350 assert agg_est.tokens_display == "~350" assert agg_est.has_estimated_tokens is True def test_format_token_count(): """Verify format_token_count formatting helper.""" assert format_token_count(1250, None) == "1250" assert format_token_count(None, 1250) == "~1250" assert format_token_count(1000, 250) == "1000" assert format_token_count(None, None) is None