Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions cycode/cli/apps/ai_guardrails/scan/handlers.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
8 changes: 7 additions & 1 deletion cycode/cli/apps/ai_guardrails/scan/payload.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,8 @@
respective IDE class.
"""

from dataclasses import dataclass
import uuid
from dataclasses import dataclass, field
from typing import Optional


Expand All @@ -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
Expand Down
1 change: 1 addition & 0 deletions cycode/cyclient/ai_security_manager_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -73,6 +73,7 @@ def create_event(
return

body = {
'id': payload.hook_event_id,
'conversation_id': conversation_id,
'event_type': event_type,
'outcome': outcome,
Expand Down
2 changes: 2 additions & 0 deletions tests/cli/commands/ai_guardrails/scan/test_handlers.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
24 changes: 24 additions & 0 deletions tests/cli/commands/ai_guardrails/scan/test_payload.py
Original file line number Diff line number Diff line change
@@ -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'
50 changes: 50 additions & 0 deletions tests/cyclient/test_ai_security_manager_client.py
Original file line number Diff line number Diff line change
@@ -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()