feat: cleanup cli (#69)

This commit is contained in:
Ishaan Sehgal
2025-08-11 11:03:18 -07:00
committed by GitHub
parent 67fc8a5312
commit 4295bf2739

View File

@@ -1,8 +1,9 @@
"""Omnara Main Entry Point
This is the main entry point for the omnara command that dispatches to either:
- MCP stdio server (default or with --stdio)
- Claude Code webhook server (with --claude-code-webhook)
This is the main entry point for the omnara command that supports:
- Default (no subcommand): Claude chat integration
- serve: Webhook server with tunnel options
- mcp: MCP stdio server
"""
import argparse
@@ -134,6 +135,7 @@ class AuthCallbackHandler(BaseHTTPRequestHandler):
<html>
<head>
<title>Omnara CLI - Authentication Successful</title>
<meta http-equiv="refresh" content="1;url=https://omnara.com/dashboard">
<style>
body {
margin: 0;
@@ -217,15 +219,21 @@ class AuthCallbackHandler(BaseHTTPRequestHandler):
<h1>Authentication Successful!</h1>
<p>Your CLI has been authorized to access Omnara.</p>
<p class="close-hint">Redirecting to dashboard in a moment...</p>
<p style="margin-top: 20px; font-size: 12px;">
If you are not redirected automatically,
<a href="https://omnara.com/dashboard" style="color: #86efac;">click here</a>.
</p>
</div>
<script>
setTimeout(() => {
window.location.href = 'https://omnara.com/dashboard';
}, 2000);
}, 500);
</script>
</body>
</html>
""")
# Give the browser time to receive the response
self.wfile.flush()
return
else:
# Invalid or missing state parameter
@@ -296,6 +304,10 @@ def authenticate_via_browser(auth_url="https://omnara.com"):
while not server.api_key and (time.time() - start_time) < 300:
time.sleep(0.1)
# If we got the API key, wait a bit for the browser to process the redirect
if server.api_key:
time.sleep(1.5) # Give browser time to receive response and start redirect
# Shutdown server in a separate thread to avoid deadlock
def shutdown_server():
server.shutdown()
@@ -312,55 +324,33 @@ def authenticate_via_browser(auth_url="https://omnara.com"):
raise Exception("Authentication failed - no API key received")
def run_stdio_server(args):
"""Run the MCP stdio server with the provided arguments"""
cmd = [
sys.executable,
"-m",
"servers.mcp_server.stdio_server",
"--api-key",
args.api_key,
]
if args.base_url:
cmd.extend(["--base-url", args.base_url])
if (
hasattr(args, "claude_code_permission_tool")
and args.claude_code_permission_tool
):
cmd.append("--claude-code-permission-tool")
if hasattr(args, "git_diff") and args.git_diff:
cmd.append("--git-diff")
if hasattr(args, "agent_instance_id") and args.agent_instance_id:
cmd.extend(["--agent-instance-id", args.agent_instance_id])
def ensure_api_key(args):
"""Ensure API key is available, authenticate if needed"""
# Check if API key is provided via argument
if hasattr(args, "api_key") and args.api_key:
return args.api_key
subprocess.run(cmd)
# Try to load from storage
api_key = load_stored_api_key()
if api_key:
return api_key
# Authenticate via browser
print("No API key found. Starting authentication...")
auth_url = getattr(args, "auth_url", "https://omnara.com")
try:
api_key = authenticate_via_browser(auth_url)
save_api_key(api_key)
print("Authentication successful! API key saved.")
return api_key
except Exception as e:
raise Exception(f"Authentication failed: {str(e)}")
def run_webhook_server(
cloudflare_tunnel=False, dangerously_skip_permissions=False, port=None
):
"""Run the Claude Code webhook FastAPI server"""
cmd = [
sys.executable,
"-m",
"webhooks.claude_code",
]
def run_claude_chat(args, unknown_args):
"""Run the Claude chat integration (default behavior)"""
api_key = ensure_api_key(args)
if dangerously_skip_permissions:
cmd.append("--dangerously-skip-permissions")
if cloudflare_tunnel:
cmd.append("--cloudflare-tunnel")
if port is not None:
cmd.extend(["--port", str(port)])
print("[INFO] Starting Claude Code webhook server...")
subprocess.run(cmd)
def run_claude_wrapper(api_key, base_url=None, claude_args=None):
"""Run the Claude wrapper V3 for Omnara integration"""
# Import and run directly instead of subprocess
from webhooks.claude_wrapper_v3 import main as claude_wrapper_main
@@ -368,12 +358,12 @@ def run_claude_wrapper(api_key, base_url=None, claude_args=None):
original_argv = sys.argv
new_argv = ["claude_wrapper_v3", "--api-key", api_key]
if base_url:
new_argv.extend(["--base-url", base_url])
if hasattr(args, "base_url") and args.base_url:
new_argv.extend(["--base-url", args.base_url])
# Add any additional Claude arguments
if claude_args:
new_argv.extend(claude_args)
if unknown_args:
new_argv.extend(unknown_args)
try:
sys.argv = new_argv
@@ -382,80 +372,63 @@ def run_claude_wrapper(api_key, base_url=None, claude_args=None):
sys.argv = original_argv
def main():
"""Main entry point that dispatches based on command line arguments"""
parser = argparse.ArgumentParser(
description="Omnara - AI Agent Dashboard and Tools",
formatter_class=argparse.RawDescriptionHelpFormatter,
epilog="""
Examples:
# Run Claude wrapper (default)
omnara --api-key YOUR_API_KEY
def cmd_serve(args):
"""Handle the 'serve' subcommand"""
# Run the webhook server with appropriate tunnel configuration
cmd = [
sys.executable,
"-m",
"webhooks.claude_code",
]
# Run Claude wrapper with custom base URL
omnara --api-key YOUR_API_KEY --base-url http://localhost:8000
if args.skip_permissions:
cmd.append("--dangerously-skip-permissions")
# Run MCP stdio server
omnara --stdio --api-key YOUR_API_KEY
# Handle tunnel configuration
if not args.no_tunnel:
# Default: use Cloudflare tunnel
cmd.append("--cloudflare-tunnel")
print("[INFO] Starting webhook server with Cloudflare tunnel...")
else:
# Local only, no tunnel
print("[INFO] Starting local webhook server (no tunnel)...")
# Run Claude Code webhook server
omnara --claude-code-webhook
if args.port is not None:
cmd.extend(["--port", str(args.port)])
# Run webhook server with Cloudflare tunnel
omnara --claude-code-webhook --cloudflare-tunnel
subprocess.run(cmd)
# Run webhook server on custom port
omnara --claude-code-webhook --port 8080
# Run with custom API base URL
omnara --stdio --api-key YOUR_API_KEY --base-url http://localhost:8000
def cmd_mcp(args):
"""Handle the 'mcp' subcommand"""
api_key = ensure_api_key(args)
# Run with custom frontend URL for authentication
omnara --auth-url http://localhost:3000
cmd = [
sys.executable,
"-m",
"servers.mcp_server.stdio_server",
"--api-key",
api_key,
]
# Run with git diff capture enabled
omnara --stdio --api-key YOUR_API_KEY --git-diff
""",
)
if args.base_url:
cmd.extend(["--base-url", args.base_url])
if args.permission_tool:
cmd.append("--claude-code-permission-tool")
if args.git_diff:
cmd.append("--git-diff")
if args.agent_instance_id:
cmd.extend(["--agent-instance-id", args.agent_instance_id])
# Add mutually exclusive group for server modes
mode_group = parser.add_mutually_exclusive_group()
mode_group.add_argument(
"--stdio",
action="store_true",
help="Run the MCP stdio server",
)
mode_group.add_argument(
"--claude-code-webhook",
action="store_true",
help="Run the Claude Code webhook server",
)
mode_group.add_argument(
"--claude",
action="store_true",
help="Run the Claude wrapper V3 for Omnara integration (default if no mode specified)",
)
subprocess.run(cmd)
# Arguments for webhook mode
def add_global_arguments(parser):
"""Add global arguments that work across all subcommands"""
parser.add_argument(
"--cloudflare-tunnel",
"--auth",
action="store_true",
help="Run Cloudflare tunnel for the webhook server (webhook mode only)",
)
parser.add_argument(
"--dangerously-skip-permissions",
action="store_true",
help="Skip permission prompts in Claude Code (webhook mode only) - USE WITH CAUTION",
)
parser.add_argument(
"--port",
type=int,
help="Port to run the webhook server on (webhook mode only, default: 6662)",
)
# Arguments for stdio mode
parser.add_argument(
"--api-key", help="API key for authentication (uses stored key if not provided)"
help="Authenticate or re-authenticate with Omnara",
)
parser.add_argument(
"--reauth",
@@ -463,100 +436,130 @@ Examples:
help="Force re-authentication even if API key exists",
)
parser.add_argument(
"--version", action="store_true", help="Show the current version of omnara"
"--version", action="store_true", help="Show version information"
)
parser.add_argument(
"--api-key", help="API key for authentication (uses stored key if not provided)"
)
parser.add_argument(
"--base-url",
default="https://agent-dashboard-mcp.onrender.com",
help="Base URL of the Omnara API server (stdio mode only)",
help="Base URL of the Omnara API server",
)
parser.add_argument(
"--auth-url",
default="https://omnara.com",
help="Base URL of the Omnara frontend for authentication (default: https://omnara.com)",
help="Base URL of the Omnara frontend for authentication",
)
parser.add_argument(
"--claude-code-permission-tool",
def main():
"""Main entry point with subcommand support"""
# Create main parser
parser = argparse.ArgumentParser(
description="Omnara - AI Agent Dashboard and Tools",
formatter_class=argparse.RawDescriptionHelpFormatter,
epilog="""
Examples:
# Start Claude chat (default)
omnara
omnara --api-key YOUR_API_KEY
# Start webhook server with Cloudflare tunnel
omnara serve
# Start local webhook server (no tunnel)
omnara serve --no-tunnel
omnara serve --no-tunnel --port 8080
# Run MCP stdio server
omnara mcp
omnara mcp --git-diff
# Authenticate
omnara --auth
# Show version
omnara --version
""",
)
# Add global arguments
add_global_arguments(parser)
# Create subparsers
subparsers = parser.add_subparsers(dest="command", help="Available commands")
# 'serve' subcommand
serve_parser = subparsers.add_parser(
"serve", help="Start webhook server for Claude Code integration"
)
serve_parser.add_argument(
"--no-tunnel",
action="store_true",
help="Enable Claude Code permission prompt tool for handling tool execution approvals (stdio mode only)",
help="Run locally without tunnel (default: uses Cloudflare tunnel)",
)
parser.add_argument(
serve_parser.add_argument(
"--port", type=int, help="Port to run the webhook server on (default: 6662)"
)
serve_parser.add_argument(
"--skip-permissions",
action="store_true",
help="Skip permission prompts in Claude Code - USE WITH CAUTION",
)
# 'mcp' subcommand
mcp_parser = subparsers.add_parser("mcp", help="Run MCP stdio server")
mcp_parser.add_argument(
"--permission-tool",
action="store_true",
help="Enable Claude Code permission prompt tool",
)
mcp_parser.add_argument(
"--git-diff",
action="store_true",
help="Enable git diff capture for log_step and ask_question (stdio mode only)",
help="Enable git diff capture for log_step and ask_question",
)
parser.add_argument(
mcp_parser.add_argument(
"--agent-instance-id",
type=str,
help="Pre-existing agent instance ID to use for this session (stdio mode only)",
help="Pre-existing agent instance ID to use for this session",
)
# Use parse_known_args to capture remaining args for Claude
# Parse arguments
args, unknown_args = parser.parse_known_args()
# Handle --version flag
# Handle version flag
if args.version:
print(f"omnara version {get_current_version()}")
sys.exit(0)
# Check for updates (only when running actual commands, not --version)
check_for_updates()
if args.cloudflare_tunnel and not args.claude_code_webhook:
parser.error("--cloudflare-tunnel can only be used with --claude-code-webhook")
if args.port is not None and not args.claude_code_webhook:
parser.error("--port can only be used with --claude-code-webhook")
# Handle re-authentication
if args.reauth:
# Handle auth flag
if args.auth or args.reauth:
try:
print("Re-authenticating...")
if args.reauth:
print("Re-authenticating...")
else:
print("Starting authentication...")
api_key = authenticate_via_browser(args.auth_url)
save_api_key(api_key)
args.api_key = api_key
print("Re-authentication successful! API key saved.")
print("Authentication successful! API key saved.")
sys.exit(0)
except Exception as e:
parser.error(f"Re-authentication failed: {str(e)}")
else:
# Load API key from storage if not provided
api_key = args.api_key
if not api_key and (args.stdio or not args.claude_code_webhook):
api_key = load_stored_api_key()
print(f"Authentication failed: {str(e)}")
sys.exit(1)
# Update args with the loaded API key
if api_key and not args.api_key:
args.api_key = api_key
# Check for updates
check_for_updates()
if args.claude_code_webhook:
run_webhook_server(
cloudflare_tunnel=args.cloudflare_tunnel,
dangerously_skip_permissions=args.dangerously_skip_permissions,
port=args.port,
)
elif args.stdio:
if not args.api_key:
try:
print("No API key found. Starting authentication...")
api_key = authenticate_via_browser(args.auth_url)
save_api_key(api_key)
args.api_key = api_key
print("Authentication successful! API key saved.")
except Exception as e:
parser.error(f"Authentication failed: {str(e)}")
run_stdio_server(args)
# Handle subcommands
if args.command == "serve":
cmd_serve(args)
elif args.command == "mcp":
cmd_mcp(args)
else:
# Default to Claude mode when no mode is specified
if not args.api_key:
try:
print("No API key found. Starting authentication...")
api_key = authenticate_via_browser(args.auth_url)
save_api_key(api_key)
args.api_key = api_key
print("Authentication successful! API key saved.")
except Exception as e:
parser.error(f"Authentication failed: {str(e)}")
run_claude_wrapper(args.api_key, args.base_url, unknown_args)
# Default behavior: run Claude chat
run_claude_chat(args, unknown_args)
if __name__ == "__main__":