Jack1808/Claude_Code
0
1"""Claude Code CLI session management."""2 3import asyncio4import json5import os6from collections.abc import AsyncGenerator7from typing import Any8 9from loguru import logger10 11from .process_registry import register_pid, unregister_pid12 13 14class CLISession:15 """Manages a single persistent Claude Code CLI subprocess."""16 17 def __init__(18 self,19 workspace_path: str,20 api_url: str,21 allowed_dirs: list[str] | None = None,22 plans_directory: str | None = None,23 ):24 self.workspace = os.path.normpath(os.path.abspath(workspace_path))25 self.api_url = api_url26 self.allowed_dirs = [os.path.normpath(d) for d in (allowed_dirs or [])]27 self.plans_directory = plans_directory28 self.process: asyncio.subprocess.Process | None = None29 self.current_session_id: str | None = None30 self._is_busy = False31 self._cli_lock = asyncio.Lock()32 33 @property34 def is_busy(self) -> bool:35 """Check if a task is currently running."""36 return self._is_busy37 38 async def start_task(39 self, prompt: str, session_id: str | None = None, fork_session: bool = False40 ) -> AsyncGenerator[dict]:41 """42 Start a new task or continue an existing session.43 44 Args:45 prompt: The user's message/prompt46 session_id: Optional session ID to resume47 48 Yields:49 Event dictionaries from the CLI50 """51 async with self._cli_lock:52 self._is_busy = True53 env = os.environ.copy()54 55 if "ANTHROPIC_API_KEY" not in env:56 env["ANTHROPIC_API_KEY"] = "sk-placeholder-key-for-proxy"57 58 env["ANTHROPIC_API_URL"] = self.api_url59 if self.api_url.endswith("/v1"):60 env["ANTHROPIC_BASE_URL"] = self.api_url[:-3]61 else:62 env["ANTHROPIC_BASE_URL"] = self.api_url63 64 env["TERM"] = "dumb"65 env["PYTHONIOENCODING"] = "utf-8"66 67 # Build command68 if session_id and not session_id.startswith("pending_"):69 cmd = [70 "claude",71 "--resume",72 session_id,73 ]74 if fork_session:75 cmd.append("--fork-session")76 cmd += [77 "-p",78 prompt,79 "--output-format",80 "stream-json",81 "--dangerously-skip-permissions",82 "--verbose",83 ]84 logger.info(f"Resuming Claude session {session_id}")85 else:86 cmd = [87 "claude",88 "-p",89 prompt,90 "--output-format",91 "stream-json",92 "--dangerously-skip-permissions",93 "--verbose",94 ]95 logger.info("Starting new Claude session")96 97 if self.allowed_dirs:98 for d in self.allowed_dirs:99 cmd.extend(["--add-dir", d])100 101 if self.plans_directory is not None:102 settings_json = json.dumps({"plansDirectory": self.plans_directory})103 cmd.extend(["--settings", settings_json])104 105 try:106 self.process = await asyncio.create_subprocess_exec(107 *cmd,108 stdout=asyncio.subprocess.PIPE,109 stderr=asyncio.subprocess.PIPE,110 cwd=self.workspace,111 env=env,112 )113 if self.process and self.process.pid:114 register_pid(self.process.pid)115 116 if not self.process or not self.process.stdout:117 yield {"type": "exit", "code": 1}118 return119 120 session_id_extracted = False121 buffer = bytearray()122 123 try:124 while True:125 chunk = await self.process.stdout.read(65536)126 if not chunk:127 if buffer:128 line_str = buffer.decode(129 "utf-8", errors="replace"130 ).strip()131 if line_str:132 async for event in self._handle_line_gen(133 line_str, session_id_extracted134 ):135 if event.get("type") == "session_info":136 session_id_extracted = True137 yield event138 break139 140 buffer.extend(chunk)141 142 while True:143 newline_pos = buffer.find(b"\n")144 if newline_pos == -1:145 break146 147 line = buffer[:newline_pos]148 buffer = buffer[newline_pos + 1 :]149 150 line_str = line.decode("utf-8", errors="replace").strip()151 if line_str:152 async for event in self._handle_line_gen(153 line_str, session_id_extracted154 ):155 if event.get("type") == "session_info":156 session_id_extracted = True157 yield event158 except asyncio.CancelledError:159 # Cancelling the handler task should not leave a Claude CLI160 # subprocess running in the background.161 try:162 await asyncio.shield(self.stop())163 finally:164 raise165 166 stderr_text = None167 if self.process.stderr:168 stderr_output = await self.process.stderr.read()169 if stderr_output:170 stderr_text = stderr_output.decode(171 "utf-8", errors="replace"172 ).strip()173 logger.error(f"Claude CLI Stderr: {stderr_text}")174 # Yield stderr as error event so it shows in UI175 if stderr_text:176 logger.info("CLI_SESSION: Yielding error event from stderr")177 yield {"type": "error", "error": {"message": stderr_text}}178 179 return_code = await self.process.wait()180 logger.info(181 f"Claude CLI exited with code {return_code}, stderr_present={bool(stderr_text)}"182 )183 if return_code != 0 and not stderr_text:184 logger.warning(185 f"CLI_SESSION: Process exited with code {return_code} but no stderr captured"186 )187 yield {188 "type": "exit",189 "code": return_code,190 "stderr": stderr_text,191 }192 finally:193 self._is_busy = False194 if self.process and self.process.pid:195 unregister_pid(self.process.pid)196 197 async def _handle_line_gen(198 self, line_str: str, session_id_extracted: bool199 ) -> AsyncGenerator[dict]:200 """Process a single line and yield events."""201 try:202 event = json.loads(line_str)203 if not session_id_extracted:204 extracted_id = self._extract_session_id(event)205 if extracted_id:206 self.current_session_id = extracted_id207 logger.info(f"Extracted session ID: {extracted_id}")208 yield {"type": "session_info", "session_id": extracted_id}209 210 yield event211 except json.JSONDecodeError:212 logger.debug(f"Non-JSON output: {line_str}")213 yield {"type": "raw", "content": line_str}214 215 def _extract_session_id(self, event: Any) -> str | None:216 """Extract session ID from CLI event."""217 if not isinstance(event, dict):218 return None219 220 if "session_id" in event:221 return event["session_id"]222 if "sessionId" in event:223 return event["sessionId"]224 225 for key in ["init", "system", "result", "metadata"]:226 if key in event and isinstance(event[key], dict):227 nested = event[key]228 if "session_id" in nested:229 return nested["session_id"]230 if "sessionId" in nested:231 return nested["sessionId"]232 233 if "conversation" in event and isinstance(event["conversation"], dict):234 conv = event["conversation"]235 if "id" in conv:236 return conv["id"]237 238 return None239 240 async def stop(self):241 """Stop the CLI process."""242 if self.process and self.process.returncode is None:243 try:244 logger.info(f"Stopping Claude CLI process {self.process.pid}")245 self.process.terminate()246 try:247 await asyncio.wait_for(self.process.wait(), timeout=5.0)248 except TimeoutError:249 self.process.kill()250 await self.process.wait()251 if self.process and self.process.pid:252 unregister_pid(self.process.pid)253 return True254 except Exception as e:255 logger.error(f"Error stopping process: {e}")256 return False257 return False258