From 59fd06579cadacf36e0ad0353b7131bca91f7607 Mon Sep 17 00:00:00 2001 From: Rushil Patel Date: Tue, 19 Aug 2025 12:35:59 -0700 Subject: [PATCH] fix: create an persist agent session before claude code --- src/codegen/cli/commands/claude/hooks.py | 19 ++----------------- src/codegen/cli/commands/claude/main.py | 20 +++++++++++++++++++- 2 files changed, 21 insertions(+), 18 deletions(-) diff --git a/src/codegen/cli/commands/claude/hooks.py b/src/codegen/cli/commands/claude/hooks.py index ac993d20b..6ecc45838 100644 --- a/src/codegen/cli/commands/claude/hooks.py +++ b/src/codegen/cli/commands/claude/hooks.py @@ -36,8 +36,6 @@ def ensure_claude_hook() -> bool: CLAUDE_CONFIG_DIR.mkdir(exist_ok=True) # Build the shell command that will create session via API and write session data - start_hook_script_path = Path(__file__).parent / "config" / "claude_session_hook.py" - start_hook_command = f"python3 {start_hook_script_path} > {SESSION_FILE}" # Build the stop hook command to mark session COMPLETE stop_hook_script_path = Path(__file__).parent / "config" / "claude_session_stop_hook.py" @@ -63,28 +61,23 @@ def ensure_claude_hook() -> bool: # Ensure proper structure exists if "hooks" not in hooks_config: hooks_config["hooks"] = {} - if "SessionStart" not in hooks_config["hooks"]: - hooks_config["hooks"]["SessionStart"] = [] if "Stop" not in hooks_config["hooks"]: hooks_config["hooks"]["Stop"] = [] if "UserPromptSubmit" not in hooks_config["hooks"]: hooks_config["hooks"]["UserPromptSubmit"] = [] # Get existing hooks - session_start_hooks = hooks_config["hooks"]["SessionStart"] stop_hooks = hooks_config["hooks"]["Stop"] active_hooks = hooks_config["hooks"]["UserPromptSubmit"] # Check if we're replacing existing hooks - replaced_existing = (len(session_start_hooks) > 0) or (len(stop_hooks) > 0) or (len(active_hooks) > 0) + replaced_existing = (len(stop_hooks) > 0) or (len(active_hooks) > 0) # Create the new hook structures (following Claude's format) - new_start_hook_group = {"hooks": [{"type": "command", "command": start_hook_command}]} new_stop_hook_group = {"hooks": [{"type": "command", "command": stop_hook_command}]} new_active_hook_group = {"hooks": [{"type": "command", "command": active_hook_command}]} # Replace all existing hooks with our single hook per event - hooks_config["hooks"]["SessionStart"] = [new_start_hook_group] hooks_config["hooks"]["Stop"] = [new_stop_hook_group] hooks_config["hooks"]["UserPromptSubmit"] = [new_active_hook_group] @@ -97,7 +90,6 @@ def ensure_claude_hook() -> bool: console.print("✅ Replaced existing Claude hooks (SessionStart, Stop)", style="green") else: console.print("✅ Registered new Claude hooks (SessionStart, Stop)", style="green") - console.print(f" Start hook: {start_hook_command}", style="dim") console.print(f" Stop hook: {stop_hook_command}", style="dim") console.print(f" Active hook:{' ' if len('Active hook:') < 1 else ''} {active_hook_command}", style="dim") @@ -106,13 +98,6 @@ def ensure_claude_hook() -> bool: with open(HOOKS_CONFIG_FILE) as f: verify_config = json.load(f) - # Check that our hooks are in the config - found_start_hook = False - for hook_group in verify_config.get("hooks", {}).get("SessionStart", []): - for hook in hook_group.get("hooks", []): - if SESSION_FILE.name in hook.get("command", ""): - found_start_hook = True - break found_stop_hook = False for hook_group in verify_config.get("hooks", {}).get("Stop", []): for hook in hook_group.get("hooks", []): @@ -126,7 +111,7 @@ def ensure_claude_hook() -> bool: found_active_hook = True break - if found_start_hook and found_stop_hook and found_active_hook: + if found_stop_hook and found_active_hook: console.print("✅ Hook configuration verified", style="dim") else: console.print("⚠️ Hook was written but verification failed", style="yellow") diff --git a/src/codegen/cli/commands/claude/main.py b/src/codegen/cli/commands/claude/main.py index ed3375963..71baf98d5 100644 --- a/src/codegen/cli/commands/claude/main.py +++ b/src/codegen/cli/commands/claude/main.py @@ -1,5 +1,6 @@ """Claude Code command with session tracking.""" +import json import os import signal import subprocess @@ -20,7 +21,7 @@ create_claude_session, ) from codegen.cli.commands.claude.config.mcp_setup import add_codegen_mcp_server, cleanup_codegen_mcp_server -from codegen.cli.commands.claude.hooks import cleanup_claude_hook, ensure_claude_hook, get_codegen_url +from codegen.cli.commands.claude.hooks import cleanup_claude_hook, ensure_claude_hook, get_codegen_url, SESSION_FILE from codegen.cli.commands.claude.quiet_console import console from codegen.cli.rich.spinners import create_spinner from codegen.cli.utils.org import resolve_org_id @@ -98,12 +99,29 @@ def _run_claude_interactive(resolved_org_id: int, no_mcp: bool | None) -> None: else: console.print("⚠️ Could not create backend session at startup (will rely on hooks)", style="yellow") except Exception as e: + agent_run_id = None console.print(f"⚠️ Session creation error at startup: {e}", style="yellow") # Set up Claude hook for session tracking if not ensure_claude_hook(): console.print("⚠️ Failed to set up session tracking hook", style="yellow") + # Write session context file for downstream hooks and tools (after hook setup) + try: + SESSION_FILE.parent.mkdir(exist_ok=True) + session_payload = { + "session_id": session_id, + "agent_run_id": agent_run_id, + "org_id": resolved_org_id, + "hook_event": "Startup", + } + with open(SESSION_FILE, "w") as f: + json.dump(session_payload, f, indent=2) + f.write("\n") + console.print("📝 Wrote session file to ~/.codegen/claude-session.json", style="dim") + except Exception as e: + console.print(f"⚠️ Could not write session file: {e}", style="yellow") + # Initialize log watcher manager log_watcher_manager = ClaudeLogWatcherManager()