CM-62381-add-session-start-hook

cycode/cli/apps/ai_guardrails/__init__.py

@@ -1,10 +1,10 @@
 import typer
 
-from cycode.cli.apps.ai_guardrails.ensure_auth_command import ensure_auth_command
 from cycode.cli.apps.ai_guardrails.install_command import install_command
+from cycode.cli.apps.ai_guardrails.session_start_command import session_start_command
 from cycode.cli.apps.ai_guardrails.scan.scan_command import scan_command
 from cycode.cli.apps.ai_guardrails.status_command import status_command
 from cycode.cli.apps.ai_guardrails.uninstall_command import uninstall_command
 
 app = typer.Typer(name='ai-guardrails', no_args_is_help=True, hidden=True)
 
@@ -18,6 +18,9 @@
     name='scan',
     short_help='Scan content from AI IDE hooks for secrets (reads JSON from stdin).',
 )(scan_command)
-app.command(hidden=True, name='ensure-auth', short_help='Ensure authentication, triggering auth if needed.')(
-    ensure_auth_command
+app.command(hidden=True, name='session-start', short_help='Handle session start: auth, conversation, data flow.')(
+    session_start_command
+)
+app.command(hidden=True, name='ensure-auth', short_help='[Deprecated] Alias for session-start.')(
+    session_start_command
 )

cycode/cli/apps/ai_guardrails/consts.py

@@ -84,15 +84,15 @@ def _get_claude_code_hooks_dir() -> Path:
 
 # Command used in hooks
 CYCODE_SCAN_PROMPT_COMMAND = 'cycode ai-guardrails scan'
-CYCODE_ENSURE_AUTH_COMMAND = 'cycode ai-guardrails ensure-auth'
+CYCODE_SESSION_START_COMMAND = 'cycode ai-guardrails session-start'
 
 
 def _get_cursor_hooks_config(async_mode: bool = False) -> dict:
     """Get Cursor-specific hooks configuration."""
     config = IDE_CONFIGS[AIIDEType.CURSOR]
     command = f'{CYCODE_SCAN_PROMPT_COMMAND} &' if async_mode else CYCODE_SCAN_PROMPT_COMMAND
     hooks = {event: [{'command': command}] for event in config.hook_events}
-    hooks['sessionStart'] = [{'command': CYCODE_ENSURE_AUTH_COMMAND}]
+    hooks['sessionStart'] = [{'command': f'{CYCODE_SESSION_START_COMMAND} --ide cursor'}]
 
     return {
         'version': 1,
@@ -119,7 +119,7 @@ def _get_claude_code_hooks_config(async_mode: bool = False) -> dict:
             'SessionStart': [
                 {
                     'matcher': 'startup',
-                    'hooks': [{'type': 'command', 'command': CYCODE_ENSURE_AUTH_COMMAND}],
+                    'hooks': [{'type': 'command', 'command': f'{CYCODE_SESSION_START_COMMAND} --ide claude-code'}],
                 }
             ],
             'UserPromptSubmit': [

cycode/cli/apps/ai_guardrails/scan/handlers.py

@@ -42,7 +42,6 @@ def handle_before_submit_prompt(ctx: typer.Context, payload: AIHookPayload, poli
     response_builder = get_response_builder(ide)
 
     prompt_config = get_policy_value(policy, 'prompt', default={})
-    ai_client.create_conversation(payload)
     if not get_policy_value(prompt_config, 'enabled', default=True):
         ai_client.create_event(payload, AiHookEventType.PROMPT, AIHookOutcome.ALLOWED)
         return response_builder.allow_prompt()
@@ -100,7 +99,6 @@ def handle_before_read_file(ctx: typer.Context, payload: AIHookPayload, policy:
     response_builder = get_response_builder(ide)
 
     file_read_config = get_policy_value(policy, 'file_read', default={})
-    ai_client.create_conversation(payload)
     if not get_policy_value(file_read_config, 'enabled', default=True):
         ai_client.create_event(payload, AiHookEventType.FILE_READ, AIHookOutcome.ALLOWED)
         return response_builder.allow_permission()
@@ -187,7 +185,6 @@ def handle_before_mcp_execution(ctx: typer.Context, payload: AIHookPayload, poli
     response_builder = get_response_builder(ide)
 
     mcp_config = get_policy_value(policy, 'mcp', default={})
-    ai_client.create_conversation(payload)
     if not get_policy_value(mcp_config, 'enabled', default=True):
         ai_client.create_event(payload, AiHookEventType.MCP_EXECUTION, AIHookOutcome.ALLOWED)
         return response_builder.allow_permission()

cycode/cli/apps/ai_guardrails/session_start_command.py

@@ -0,0 +1,106 @@
+import sys
+from typing import Annotated
+
+import typer
+
+from cycode.cli.apps.ai_guardrails.consts import AIIDEType
+from cycode.cli.apps.ai_guardrails.scan.claude_config import get_mcp_servers, get_user_email, load_claude_config
+from cycode.cli.apps.ai_guardrails.scan.payload import AIHookPayload, _extract_from_claude_transcript
+from cycode.cli.apps.ai_guardrails.scan.utils import safe_json_parse
+from cycode.cli.apps.auth.auth_common import get_authorization_info
+from cycode.cli.apps.auth.auth_manager import AuthManager
+from cycode.cli.exceptions.handle_auth_errors import handle_auth_exception
+from cycode.cli.utils.get_api_client import get_ai_security_manager_client
+from cycode.logger import get_logger
+
+logger = get_logger('AI Guardrails')
+
+
+def _build_session_payload(payload: dict, ide: str) -> AIHookPayload:
+    """Build an AIHookPayload from a session-start stdin payload."""
+    if ide == AIIDEType.CLAUDE_CODE:
+        ide_version, model, _ = _extract_from_claude_transcript(payload.get('transcript_path'))
+        claude_config = load_claude_config()
+        ide_user_email = get_user_email(claude_config) if claude_config else None
+
+        return AIHookPayload(
+            event_name='session_start',
+            conversation_id=payload.get('session_id'),
+            ide_user_email=ide_user_email,
+            model=payload.get('model') or model,
+            ide_provider=AIIDEType.CLAUDE_CODE.value,
+            ide_version=ide_version,
+        )
+
+    # Cursor
+    return AIHookPayload(
+        event_name='session_start',
+        conversation_id=payload.get('conversation_id'),
+        ide_user_email=payload.get('user_email'),
+        model=payload.get('model'),
+        ide_provider=AIIDEType.CURSOR.value,
+        ide_version=payload.get('cursor_version'),
+    )
+
+
+def session_start_command(
+    ctx: typer.Context,
+    ide: Annotated[
+        str,
+        typer.Option(
+            '--ide',
+            help='IDE that triggered the session start.',
+            hidden=True,
+        ),
+    ] = AIIDEType.CURSOR.value,
+) -> None:
+    """Handle session start: ensure auth, create conversation, report data flow."""
+    # Step 1: Ensure authentication
+    auth_info = get_authorization_info(ctx)
+    if auth_info is None:
+        logger.debug('Not authenticated, starting authentication')
+        try:
+            auth_manager = AuthManager()
+            auth_manager.authenticate()
+        except Exception as err:
+            handle_auth_exception(ctx, err)
+            return
+    else:
+        logger.debug('Already authenticated')
+
+    # Step 2: Read stdin payload (backward compat: old hooks pipe no stdin)
+    if sys.stdin.isatty():
+        logger.debug('No stdin payload (TTY), skipping session initialization')
+        return
+
+    stdin_data = sys.stdin.read().strip()
+    payload = safe_json_parse(stdin_data)
+    if not payload:
+        logger.debug('Empty or invalid stdin payload, skipping session initialization')
+        return
+
+    # Step 3: Build session payload and initialize API client
+    session_payload = _build_session_payload(payload, ide)
+
+    try:
+        ai_client = get_ai_security_manager_client(ctx)
+    except Exception as e:
+        logger.debug('Failed to initialize AI security client', exc_info=e)
+        return
+
+    # Step 4: Create conversation
+    try:
+        ai_client.create_conversation(session_payload)
+    except Exception as e:
+        logger.debug('Failed to create conversation during session start', exc_info=e)
+
+    # Step 5: Report data flow (MCP servers, Claude Code only)
+    if ide == AIIDEType.CLAUDE_CODE:
+        claude_config = load_claude_config()
+        if claude_config:
+            mcp_servers = get_mcp_servers(claude_config)
+            if mcp_servers:
+                try:
+                    ai_client.report_data_flow(mcp_servers)
+                except Exception as e:
+                    logger.debug('Failed to report MCP servers', exc_info=e)

cycode/cyclient/ai_security_manager_client.py

@@ -17,6 +17,7 @@ class AISecurityManagerClient:
 
     _CONVERSATIONS_PATH = 'v4/ai-security/interactions/conversations'
     _EVENTS_PATH = 'v4/ai-security/interactions/events'
+    _DATA_FLOW_PATH = 'v4/ai-security/interactions/data-flow'
 
     def __init__(self, client: CycodeClientBase, service_config: 'AISecurityManagerServiceConfigBase') -> None:
         self.client = client
@@ -88,3 +89,15 @@ def create_event(
         except Exception as e:
             logger.debug('Failed to create AI hook event', exc_info=e)
             # Don't fail the hook if tracking fails
+
+    def report_data_flow(self, mcp_servers: Optional[dict] = None) -> None:
+        """Report session data flow to the backend."""
+        body: dict = {
+            'mcp_servers': mcp_servers,
+        }
+
+        try:
+            self.client.post(self._build_endpoint_path(self._DATA_FLOW_PATH), body=body)
+        except Exception as e:
+            logger.debug('Failed to report data flow', exc_info=e)
+            # Don't fail the session if reporting fails

tests/cli/commands/ai_guardrails/scan/test_handlers.py

@@ -67,6 +67,7 @@ def test_handle_before_submit_prompt_disabled(
 
     assert result == {'continue': True}
     mock_ctx.obj['ai_security_client'].create_event.assert_called_once()
+    mock_ctx.obj['ai_security_client'].create_conversation.assert_not_called()
 
 
 @patch('cycode.cli.apps.ai_guardrails.scan.handlers._scan_text_for_secrets')
@@ -80,6 +81,7 @@ def test_handle_before_submit_prompt_no_secrets(
 
     assert result == {'continue': True}
     mock_ctx.obj['ai_security_client'].create_event.assert_called_once()
+    mock_ctx.obj['ai_security_client'].create_conversation.assert_not_called()
     call_args = mock_ctx.obj['ai_security_client'].create_event.call_args
     # outcome is arg[2], scan_id and block_reason are kwargs
     assert call_args.args[2] == AIHookOutcome.ALLOWED

tests/cli/commands/ai_guardrails/test_hooks_manager.py

@@ -6,8 +6,8 @@
 from pyfakefs.fake_filesystem import FakeFilesystem
 
 from cycode.cli.apps.ai_guardrails.consts import (
-    CYCODE_ENSURE_AUTH_COMMAND,
     CYCODE_SCAN_PROMPT_COMMAND,
+    CYCODE_SESSION_START_COMMAND,
     AIIDEType,
     PolicyMode,
     get_hooks_config,
@@ -88,12 +88,13 @@ def test_get_hooks_config_cursor_async() -> None:
 
 
 def test_get_hooks_config_cursor_session_start() -> None:
-    """Test Cursor hooks config includes sessionStart auth check."""
+    """Test Cursor hooks config includes sessionStart with --ide flag."""
     config = get_hooks_config(AIIDEType.CURSOR)
     assert 'sessionStart' in config['hooks']
     entries = config['hooks']['sessionStart']
     assert len(entries) == 1
-    assert entries[0]['command'] == CYCODE_ENSURE_AUTH_COMMAND
+    assert CYCODE_SESSION_START_COMMAND in entries[0]['command']
+    assert '--ide cursor' in entries[0]['command']
 
 
 def test_get_hooks_config_claude_code_sync() -> None:
@@ -118,12 +119,13 @@ def test_get_hooks_config_claude_code_async() -> None:
 
 
 def test_get_hooks_config_claude_code_session_start() -> None:
-    """Test Claude Code hooks config includes SessionStart auth check."""
+    """Test Claude Code hooks config includes SessionStart with --ide flag."""
     config = get_hooks_config(AIIDEType.CLAUDE_CODE)
     assert 'SessionStart' in config['hooks']
     entries = config['hooks']['SessionStart']
     assert len(entries) == 1
-    assert entries[0]['hooks'][0]['command'] == CYCODE_ENSURE_AUTH_COMMAND
+    assert CYCODE_SESSION_START_COMMAND in entries[0]['hooks'][0]['command']
+    assert '--ide claude-code' in entries[0]['hooks'][0]['command']
 
 
 def test_create_policy_file_warn(fs: FakeFilesystem) -> None: