66from pathlib import Path
77from typing import Optional
88
9+ from cycode .cli .apps .ai_guardrails .consts import AIIDEType
910from cycode .cli .apps .ai_guardrails .scan .types import (
1011 CLAUDE_CODE_EVENT_MAPPING ,
1112 CLAUDE_CODE_EVENT_NAMES ,
@@ -47,7 +48,7 @@ def _reverse_readline(path: Path, buf_size: int = 8192) -> Iterator[str]:
4748 if newline_pos == - 1 :
4849 break
4950 # Yield the line after this newline
50- line = buffer [newline_pos + 1 :]
51+ line = buffer [newline_pos + 1 :]
5152 buffer = buffer [: newline_pos + 1 ]
5253 if line .strip ():
5354 yield line .decode ('utf-8' , errors = 'replace' )
@@ -57,8 +58,20 @@ def _reverse_readline(path: Path, buf_size: int = 8192) -> Iterator[str]:
5758 yield buffer .decode ('utf-8' , errors = 'replace' )
5859
5960
60- def _extract_from_claude_transcript ( # noqa: C901
61- transcript_path : str ,
61+ def _extract_model (entry : dict ) -> Optional [str ]:
62+ """Extract model from a transcript entry (top level or nested in message)."""
63+ return entry .get ('model' ) or (entry .get ('message' ) or {}).get ('model' )
64+
65+
66+ def _extract_generation_id (entry : dict ) -> Optional [str ]:
67+ """Extract generation ID from a user-type transcript entry."""
68+ if entry .get ('type' ) == 'user' :
69+ return entry .get ('uuid' )
70+ return None
71+
72+
73+ def _extract_from_claude_transcript (
74+ transcript_path : str ,
6275) -> tuple [Optional [str ], Optional [str ], Optional [str ]]:
6376 """Extract IDE version, model, and latest generation ID from Claude Code transcript file.
6477
@@ -90,16 +103,11 @@ def _extract_from_claude_transcript( # noqa: C901
90103 continue
91104 try :
92105 entry = json .loads (line )
93- if ide_version is None and 'version' in entry :
94- ide_version = entry ['version' ]
95- # Model can be at top level or nested in message.model
96- if model is None :
97- model = entry .get ('model' ) or (entry .get ('message' ) or {}).get ('model' )
98- # Get the latest user message UUID as generation_id
99- if generation_id is None and entry .get ('type' ) == 'user' and entry .get ('uuid' ):
100- generation_id = entry ['uuid' ]
101- # Stop early if we found all values
102- if ide_version is not None and model is not None and generation_id is not None :
106+ ide_version = ide_version or entry .get ('version' )
107+ model = model or _extract_model (entry )
108+ generation_id = generation_id or _extract_generation_id (entry )
109+
110+ if ide_version and model and generation_id :
103111 break
104112 except json .JSONDecodeError :
105113 continue
@@ -121,7 +129,7 @@ class AIHookPayload:
121129 # User and IDE information
122130 ide_user_email : Optional [str ] = None
123131 model : Optional [str ] = None
124- ide_provider : str = None # e.g., 'cursor', 'claude-code'
132+ ide_provider : str = None # AIIDEType value ( e.g., 'cursor', 'claude-code')
125133 ide_version : Optional [str ] = None
126134
127135 # Event-specific data
@@ -147,7 +155,7 @@ def from_cursor_payload(cls, payload: dict) -> 'AIHookPayload':
147155 generation_id = payload .get ('generation_id' ),
148156 ide_user_email = payload .get ('user_email' ),
149157 model = payload .get ('model' ),
150- ide_provider = 'cursor' ,
158+ ide_provider = AIIDEType . CURSOR ,
151159 ide_version = payload .get ('cursor_version' ),
152160 prompt = payload .get ('prompt' , '' ),
153161 file_path = payload .get ('file_path' ) or payload .get ('path' ),
@@ -205,7 +213,7 @@ def from_claude_code_payload(cls, payload: dict) -> 'AIHookPayload':
205213 generation_id = generation_id ,
206214 ide_user_email = None , # Claude Code doesn't provide this in hook payload
207215 model = model ,
208- ide_provider = 'claude-code' ,
216+ ide_provider = AIIDEType . CLAUDE_CODE ,
209217 ide_version = ide_version ,
210218 prompt = payload .get ('prompt' , '' ),
211219 file_path = file_path ,
@@ -224,37 +232,37 @@ def is_payload_for_ide(payload: dict, ide: str) -> bool:
224232
225233 Args:
226234 payload: The raw payload from the IDE
227- ide: The IDE name (e.g., 'cursor', 'claude-code')
235+ ide: The IDE name or AIIDEType enum value
228236
229237 Returns:
230238 True if the payload matches the IDE, False otherwise.
231239 """
232240 hook_event_name = payload .get ('hook_event_name' , '' )
233241
234- if ide == 'claude-code' :
242+ if ide == AIIDEType . CLAUDE_CODE :
235243 return hook_event_name in CLAUDE_CODE_EVENT_NAMES
236- if ide == 'cursor' :
244+ if ide == AIIDEType . CURSOR :
237245 return hook_event_name in CURSOR_EVENT_NAMES
238246
239247 # Unknown IDE, allow processing
240248 return True
241249
242250 @classmethod
243- def from_payload (cls , payload : dict , tool : str = 'cursor' ) -> 'AIHookPayload' :
251+ def from_payload (cls , payload : dict , tool : str = AIIDEType . CURSOR ) -> 'AIHookPayload' :
244252 """Create AIHookPayload from any tool's payload.
245253
246254 Args:
247255 payload: The raw payload from the IDE
248- tool: The IDE/tool name (e.g., 'cursor', 'claude-code')
256+ tool: The IDE/tool name or AIIDEType enum value
249257
250258 Returns:
251259 AIHookPayload instance
252260
253261 Raises:
254262 ValueError: If the tool is not supported
255263 """
256- if tool == 'cursor' :
264+ if tool == AIIDEType . CURSOR :
257265 return cls .from_cursor_payload (payload )
258- if tool == 'claude-code' :
266+ if tool == AIIDEType . CLAUDE_CODE :
259267 return cls .from_claude_code_payload (payload )
260268 raise ValueError (f'Unsupported IDE/tool: { tool } )
0 commit comments