diff --git a/cycode/cli/apps/ai_guardrails/scan/handlers.py b/cycode/cli/apps/ai_guardrails/scan/handlers.py index e82c61d7..4e773cf1 100644 --- a/cycode/cli/apps/ai_guardrails/scan/handlers.py +++ b/cycode/cli/apps/ai_guardrails/scan/handlers.py @@ -334,6 +334,7 @@ def build_ai_guardrails_scan_parameters( 'device_hostname': get_hostname(), 'conversation_id': payload.conversation_id, 'generation_id': payload.generation_id, + 'hook_event_id': payload.hook_event_id, 'ide_user_email': payload.ide_user_email, 'mcp_server_name': payload.mcp_server_name, 'mcp_tool_name': payload.mcp_tool_name, diff --git a/cycode/cli/apps/ai_guardrails/scan/payload.py b/cycode/cli/apps/ai_guardrails/scan/payload.py index 19845601..b131f6a8 100644 --- a/cycode/cli/apps/ai_guardrails/scan/payload.py +++ b/cycode/cli/apps/ai_guardrails/scan/payload.py @@ -5,7 +5,8 @@ respective IDE class. """ -from dataclasses import dataclass +import uuid +from dataclasses import dataclass, field from typing import Optional @@ -18,6 +19,11 @@ class AIHookPayload: conversation_id: Optional[str] = None generation_id: Optional[str] = None + # Minted here rather than by the server: the guardrail scan and the hook event are reported in two + # separate requests, and both have to name the same event. A generation id can't stand in for it - the + # IDE mints one per prompt, so several hook events share it, and some IDEs don't supply one at all. + hook_event_id: str = field(default_factory=lambda: str(uuid.uuid4())) + # User and IDE information ide_user_email: Optional[str] = None model: Optional[str] = None diff --git a/cycode/cyclient/ai_security_manager_client.py b/cycode/cyclient/ai_security_manager_client.py index 5dee7f2c..62f5618b 100644 --- a/cycode/cyclient/ai_security_manager_client.py +++ b/cycode/cyclient/ai_security_manager_client.py @@ -73,6 +73,7 @@ def create_event( return body = { + 'id': payload.hook_event_id, 'conversation_id': conversation_id, 'event_type': event_type, 'outcome': outcome, diff --git a/tests/cli/commands/ai_guardrails/scan/test_handlers.py b/tests/cli/commands/ai_guardrails/scan/test_handlers.py index 401482ac..7c008bfc 100644 --- a/tests/cli/commands/ai_guardrails/scan/test_handlers.py +++ b/tests/cli/commands/ai_guardrails/scan/test_handlers.py @@ -557,6 +557,8 @@ def test_build_ai_guardrails_scan_parameters( 'device_hostname': 'test-host', 'conversation_id': 'test-conv-id', 'generation_id': 'test-gen-id', + # The same id the hook event is reported under, so the detection can point at that one event + 'hook_event_id': mock_payload.hook_event_id, 'ide_user_email': 'test@example.com', 'mcp_server_name': None, 'mcp_tool_name': None, diff --git a/tests/cli/commands/ai_guardrails/scan/test_payload.py b/tests/cli/commands/ai_guardrails/scan/test_payload.py new file mode 100644 index 00000000..5010c433 --- /dev/null +++ b/tests/cli/commands/ai_guardrails/scan/test_payload.py @@ -0,0 +1,24 @@ +from uuid import UUID + +from cycode.cli.apps.ai_guardrails.scan.payload import AIHookPayload + + +def test_hook_event_id_is_a_uuid() -> None: + payload = AIHookPayload(event_name='Prompt') + + # Round-trips through UUID, so ai-security-manager can store it as the hook event's primary key + assert UUID(payload.hook_event_id) + + +def test_hook_event_id_is_unique_per_payload() -> None: + """One payload is one hook event. Sharing an id across events is what generation ids already do wrong.""" + first = AIHookPayload(event_name='Prompt', generation_id='same-gen') + second = AIHookPayload(event_name='Prompt', generation_id='same-gen') + + assert first.hook_event_id != second.hook_event_id + + +def test_hook_event_id_survives_an_explicit_value() -> None: + payload = AIHookPayload(event_name='Prompt', hook_event_id='fixed-id') + + assert payload.hook_event_id == 'fixed-id' diff --git a/tests/cyclient/test_ai_security_manager_client.py b/tests/cyclient/test_ai_security_manager_client.py new file mode 100644 index 00000000..c9450675 --- /dev/null +++ b/tests/cyclient/test_ai_security_manager_client.py @@ -0,0 +1,50 @@ +from unittest.mock import MagicMock + +from cycode.cli.apps.ai_guardrails.scan.payload import AIHookPayload +from cycode.cli.apps.ai_guardrails.scan.types import AiHookEventType, AIHookOutcome +from cycode.cyclient.ai_security_manager_client import AISecurityManagerClient + + +def _build_client() -> tuple[AISecurityManagerClient, MagicMock]: + http_client = MagicMock() + service_config = MagicMock() + service_config.get_service_name.return_value = None + + return AISecurityManagerClient(http_client, service_config), http_client + + +def _posted_body(http_client: MagicMock) -> dict: + return http_client.post.call_args.kwargs['body'] + + +def test_create_event_reports_the_payload_hook_event_id_as_the_event_id() -> None: + """The CLI owns the id so the guardrail detection, reported separately, can name this exact event.""" + client, http_client = _build_client() + payload = AIHookPayload(event_name='Prompt', conversation_id='conv-1', generation_id='gen-1') + + client.create_event(payload, AiHookEventType.PROMPT, AIHookOutcome.ALLOWED) + + assert _posted_body(http_client)['id'] == payload.hook_event_id + + +def test_create_event_reports_a_distinct_id_per_hook_event() -> None: + client, http_client = _build_client() + conversation_id = 'conv-1' + # Two hooks of the same prompt: the generation id is shared, the hook event id must not be + first = AIHookPayload(event_name='Prompt', conversation_id=conversation_id, generation_id='gen-1') + second = AIHookPayload(event_name='FileRead', conversation_id=conversation_id, generation_id='gen-1') + + client.create_event(first, AiHookEventType.PROMPT, AIHookOutcome.ALLOWED) + client.create_event(second, AiHookEventType.FILE_READ, AIHookOutcome.ALLOWED) + + reported_ids = [call.kwargs['body']['id'] for call in http_client.post.call_args_list] + assert reported_ids == [first.hook_event_id, second.hook_event_id] + assert len(set(reported_ids)) == 2 + + +def test_create_event_without_a_conversation_posts_nothing() -> None: + client, http_client = _build_client() + + client.create_event(AIHookPayload(event_name='Prompt'), AiHookEventType.PROMPT, AIHookOutcome.ALLOWED) + + http_client.post.assert_not_called()