Skip to content

Commit 836a7f3

Browse files
authored
feat: add organization and user rules to agent (#1278)
# Motivation <!-- Why is this change necessary? --> # Content <!-- Please include a summary of the change --> # Testing <!-- How was the change tested? --> # Please check the following before marking your PR as ready for review - [x] I have added tests for my changes - [x] I have updated the documentation or added new documentation as needed
1 parent 12387e7 commit 836a7f3

File tree

2 files changed

+93
-3
lines changed

2 files changed

+93
-3
lines changed

src/codegen/cli/commands/claude/claude_session_api.py

Lines changed: 57 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,12 +5,13 @@
55
from typing import Optional
66

77
import requests
8-
from .quiet_console import console
98

109
from codegen.cli.api.endpoints import API_ENDPOINT
1110
from codegen.cli.auth.token_manager import get_current_token
1211
from codegen.cli.utils.org import resolve_org_id
1312

13+
from .quiet_console import console
14+
1415

1516
class ClaudeSessionAPIError(Exception):
1617
"""Exception raised for Claude session API errors."""
@@ -192,6 +193,61 @@ def send_claude_session_log(session_id: str, log_entry: dict, org_id: Optional[i
192193
return False
193194

194195

196+
def get_cli_rules(org_id: Optional[int] = None) -> Optional[dict]:
197+
"""Fetch CLI rules from the API endpoint.
198+
199+
Args:
200+
org_id: Organization ID (will be resolved if None)
201+
202+
Returns:
203+
Dictionary containing organization_rules and user_custom_prompt, or None if failed
204+
"""
205+
try:
206+
# Resolve org_id
207+
resolved_org_id = resolve_org_id(org_id)
208+
if resolved_org_id is None:
209+
console.print("⚠️ Could not resolve organization ID for CLI rules", style="yellow")
210+
return None
211+
212+
# Get authentication token
213+
token = get_current_token()
214+
if not token:
215+
console.print("⚠️ No authentication token found for CLI rules", style="yellow")
216+
return None
217+
218+
# Prepare API request
219+
url = f"{API_ENDPOINT.rstrip('/')}/v1/organizations/{resolved_org_id}/cli/rules"
220+
headers = {"Authorization": f"Bearer {token}"}
221+
222+
# Make API request
223+
response = requests.get(url, headers=headers, timeout=30)
224+
225+
if response.status_code == 200:
226+
try:
227+
result = response.json()
228+
return result
229+
except json.JSONDecodeError as e:
230+
console.print(f"⚠️ Invalid response format from CLI rules: {e}", style="yellow")
231+
return None
232+
else:
233+
error_msg = f"HTTP {response.status_code}"
234+
try:
235+
error_detail = response.json().get("detail", response.text)
236+
error_msg = f"{error_msg}: {error_detail}"
237+
except Exception:
238+
error_msg = f"{error_msg}: {response.text}"
239+
240+
console.print(f"⚠️ Failed to fetch CLI rules: {error_msg}", style="yellow")
241+
return None
242+
243+
except requests.RequestException as e:
244+
console.print(f"⚠️ Network error fetching CLI rules: {e}", style="yellow")
245+
return None
246+
except Exception as e:
247+
console.print(f"⚠️ Unexpected error fetching CLI rules: {e}", style="yellow")
248+
return None
249+
250+
195251
def write_session_hook_data(session_id: str, org_id: Optional[int] = None) -> str:
196252
"""Write session data for Claude hook and create session via API.
197253

src/codegen/cli/commands/claude/main.py

Lines changed: 36 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
from codegen.cli.commands.claude.claude_session_api import (
2020
create_claude_session,
2121
generate_session_id,
22+
get_cli_rules,
2223
update_claude_session_status,
2324
)
2425
from codegen.cli.commands.claude.config.mcp_setup import add_codegen_mcp_server, cleanup_codegen_mcp_server
@@ -221,13 +222,46 @@ def _run_claude_interactive(resolved_org_id: int, no_mcp: bool | None) -> None:
221222
console.print("🔵 Starting Claude Code session...", style="blue")
222223

223224
try:
225+
# Fetch CLI rules for system prompt
226+
console.print("📋 Fetching CLI rules...", style="blue")
227+
cli_rules = get_cli_rules(resolved_org_id)
228+
229+
# Build Claude command
230+
claude_cmd = [claude_path, "--session-id", session_id]
231+
232+
# Add system prompt if CLI rules were fetched successfully
233+
if cli_rules:
234+
system_prompt_parts = []
235+
236+
# Add organization rules if available
237+
if cli_rules.get("organization_rules"):
238+
system_prompt_parts.append("Organization Rules:")
239+
system_prompt_parts.append(cli_rules["organization_rules"])
240+
241+
# Add user custom prompt if available
242+
if cli_rules.get("user_custom_prompt"):
243+
if system_prompt_parts: # Add separator if we already have org rules
244+
system_prompt_parts.append("\n")
245+
system_prompt_parts.append("User Custom Prompt:")
246+
system_prompt_parts.append(cli_rules["user_custom_prompt"])
247+
248+
# Combine all parts into system prompt
249+
if system_prompt_parts:
250+
system_prompt = "\n".join(system_prompt_parts)
251+
claude_cmd.extend(["--append-system-prompt", system_prompt])
252+
console.print("✅ Added CLI rules to system prompt", style="green")
253+
else:
254+
console.print("⚠️ CLI rules response was empty", style="yellow")
255+
else:
256+
console.print("⚠️ Could not fetch CLI rules, continuing without system prompt", style="yellow")
257+
224258
# Launch Claude Code with our session ID
225259
console.print(f"🚀 Launching Claude Code with session ID: {session_id[:8]}...", style="blue")
226260

227261
url = get_codegen_url(session_id)
228262
console.print(f"\n🔵 Codegen URL: {url}\n", style="bold blue")
229263

230-
process = subprocess.Popen([claude_path, "--session-id", session_id])
264+
process = subprocess.Popen(claude_cmd)
231265

232266
# Start log watcher for the session
233267
console.print("📋 Starting log watcher...", style="blue")
@@ -288,7 +322,7 @@ def signal_handler(signum, frame):
288322
console.print("✅ Claude Code finished successfully", style="green")
289323

290324
except FileNotFoundError:
291-
logger.error(
325+
logger.exception(
292326
"Claude Code executable not found",
293327
extra={"operation": "claude.interactive", "org_id": resolved_org_id, "claude_session_id": session_id, "error_type": "claude_executable_not_found", **_get_session_context()},
294328
)

0 commit comments

Comments
 (0)