codex-memory 3.0.15

#![allow(clippy::uninlined_format_args)]

use clap::{Parser, Subcommand};
use codex_memory::{Config, MCPServer, Storage};
use std::fs;
use std::process;
use std::sync::Arc;
use tokio::signal;
use tracing::{error, info, warn};
use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};

#[derive(Parser)]
#[command(name = "codex-memory")]
#[command(about = "Simple text storage service with MCP interface")]
#[command(version = env!("CARGO_PKG_VERSION"))]
struct Cli {
    #[command(subcommand)]
    command: Commands,
}

#[derive(Subcommand)]
enum Commands {
    /// Run MCP server in stdio mode for Claude Desktop
    Mcp,

    /// Setup local database (create database, user, and run migrations)
    Setup {
        /// Skip database/user creation (only run migrations)
        #[arg(long)]
        skip_db_creation: bool,
    },

    /// Store text content
    Store {
        /// Content to store
        content: String,

        /// Context about what is being stored (required)
        #[arg(long)]
        context: String,

        /// Summary of the content - 120 words or less (required)
        #[arg(long)]
        summary: String,

        /// Optional tags (comma-separated)
        #[arg(long)]
        tags: Option<String>,
    },

    /// Get content by ID
    Get {
        /// UUID of the content to retrieve
        id: String,
    },

    /// Get storage statistics
    Stats,
}

const MCP_PID_FILE: &str = "/tmp/codex-memory-mcp.pid";

/// Ensure only one MCP instance runs at a time
fn ensure_singleton() -> Result<(), Box<dyn std::error::Error>> {
    let current_pid = process::id();

    // Check if PID file exists
    if let Ok(pid_content) = fs::read_to_string(MCP_PID_FILE) {
        if let Ok(existing_pid) = pid_content.trim().parse::<u32>() {
            // Check if process is actually running
            #[cfg(unix)]
            {
                use std::process::Command;
                let output = Command::new("kill")
                    .args(["-0", &existing_pid.to_string()])
                    .output();

                match output {
                    Ok(result) if result.status.success() => {
                        // Another instance is running, but Claude Desktop may need multiple instances
                        eprintln!(
                            "MCP server already running with PID {}, allowing additional instance",
                            existing_pid
                        );
                        // Allow this instance to continue with a different PID file
                        let alt_pid_file = format!("/tmp/codex-memory-mcp-{}.pid", current_pid);
                        fs::write(&alt_pid_file, current_pid.to_string())?;
                        info!(
                            "MCP server starting with PID {} (alternative instance)",
                            current_pid
                        );
                        return Ok(());
                    }
                    _ => {
                        // Process not running, clean up stale PID file
                        eprintln!("Cleaning up stale PID file for process {}", existing_pid);
                        let _ = fs::remove_file(MCP_PID_FILE);
                    }
                }
            }
        }
    }

    // Write our PID to the file
    fs::write(MCP_PID_FILE, current_pid.to_string())?;
    info!("MCP server starting with PID {}", current_pid);
    Ok(())
}

/// Clean up PID file on exit
fn cleanup_pid_file() {
    if let Err(e) = fs::remove_file(MCP_PID_FILE) {
        error!("Failed to remove PID file: {}", e);
    } else {
        info!("Cleaned up PID file");
    }
}

/// Setup comprehensive signal handlers for graceful shutdown
/// Implements CODEX-RUST-006 graceful shutdown requirements
async fn setup_signal_handlers() -> Result<(), Box<dyn std::error::Error>> {
    use signal::unix::{SignalKind, signal};
    
    let mut sigterm = signal(SignalKind::terminate())?;
    let mut sigint = signal(SignalKind::interrupt())?;
    let mut sigquit = signal(SignalKind::quit())?;
    let mut sighup = signal(SignalKind::hangup())?;

    tokio::select! {
        _ = sigterm.recv() => {
            info!("Received SIGTERM, initiating graceful shutdown...");
        }
        _ = sigint.recv() => {
            info!("Received SIGINT, initiating graceful shutdown...");
        }
        _ = sigquit.recv() => {
            info!("Received SIGQUIT, initiating graceful shutdown...");
        }
        _ = sighup.recv() => {
            info!("Received SIGHUP, initiating graceful shutdown...");
        }
        _ = signal::ctrl_c() => {
            info!("Received Ctrl+C, initiating graceful shutdown...");
        }
    }

    // Give time for cleanup operations
    info!("Starting graceful shutdown sequence...");
    
    // Wait a brief moment for any in-flight operations to complete
    tokio::time::sleep(tokio::time::Duration::from_millis(500)).await;
    
    cleanup_pid_file();
    info!("Graceful shutdown complete");
    
    Ok(())
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let cli = Cli::parse();
    let config = Config::from_env()?;

    // Initialize logging - use stderr for MCP mode to keep stdout clean for JSON-RPC
    if matches!(cli.command, Commands::Mcp) {
        // For MCP: log to stderr so stdout remains clean for JSON-RPC protocol
        tracing_subscriber::registry()
            .with(tracing_subscriber::EnvFilter::from_default_env())
            .with(
                tracing_subscriber::fmt::layer()
                    .with_writer(std::io::stderr)
                    .with_target(false),
            )
            .init();
    } else {
        // For CLI: normal logging to stdout
        tracing_subscriber::registry()
            .with(tracing_subscriber::EnvFilter::from_default_env())
            .with(tracing_subscriber::fmt::layer())
            .init();
    }

    match cli.command {
        Commands::Setup { skip_db_creation } => {
            // Setup database
            if !skip_db_creation {
                codex_memory::database::setup_local_database().await?;
            }

            // Run migrations
            println!("\n📝 Running database migrations...");
            let pool = codex_memory::database::create_pool(&config.database_url).await?;
            codex_memory::database::run_migrations(&pool).await?;
            println!("✅ Migrations complete!");

            println!("\n🎉 Setup complete! You can now run:");
            println!("  codex-memory mcp    # Start MCP server");
            println!("  codex-memory store \"Your text here\" --context \"example\" --summary \"test\"    # Store content");
            println!("  codex-memory stats    # View statistics");
        }
        Commands::Mcp => {
            // Ensure only one MCP instance runs
            ensure_singleton()?;

            // Set up graceful shutdown handling with timeout
            let shutdown_handler = tokio::spawn(async {
                if let Err(e) = setup_signal_handlers().await {
                    error!("Signal handler setup failed: {}", e);
                }
            });

            // Create database pool and run migrations
            info!("Initializing database connection and running migrations...");
            let pool = codex_memory::database::create_pool(&config.database_url).await?;
            codex_memory::database::run_migrations(&pool).await?;
            let storage = Arc::new(Storage::new(pool));

            // Run MCP server in stdio mode with graceful shutdown
            info!("Starting MCP server in stdio mode...");
            let server = MCPServer::new(config, storage);

            tokio::select! {
                result = server.run_stdio() => {
                    info!("MCP server finished normally");
                    cleanup_pid_file();
                    result?;
                }
                _ = shutdown_handler => {
                    info!("Graceful shutdown initiated by signal");
                    cleanup_pid_file();
                }
                // Add timeout fallback for forced shutdown
                _ = tokio::time::sleep(tokio::time::Duration::from_secs(30)) => {
                    warn!("Shutdown timeout exceeded, forcing exit");
                    cleanup_pid_file();
                    std::process::exit(1);
                }
            }
        }
        Commands::Store {
            content,
            context,
            summary,
            tags,
        } => {
            // Create database pool and run migrations
            let pool = codex_memory::database::create_pool(&config.database_url).await?;
            codex_memory::database::run_migrations(&pool).await?;
            let storage = Arc::new(Storage::new(pool));
            // Parse tags if provided
            let tag_list = tags.map(|t| {
                t.split(',')
                    .map(|s| s.trim().to_string())
                    .collect::<Vec<_>>()
            });

            let id = storage
                .store(&content, context.clone(), summary.clone(), tag_list)
                .await?;
            println!("Stored with ID: {}", id);
            println!("Context: {}", context);
            println!("Summary: {}", summary);
        }
        Commands::Get { id } => {
            // Create database pool and run migrations
            let pool = codex_memory::database::create_pool(&config.database_url).await?;
            codex_memory::database::run_migrations(&pool).await?;
            let storage = Arc::new(Storage::new(pool));

            let uuid = id.parse::<uuid::Uuid>()?;
            match storage.get(uuid).await? {
                Some(memory) => {
                    println!("Content: {}", memory.content);
                    println!("Context: {}", memory.context);
                    println!("Summary: {}", memory.summary);
                    println!("Created: {}", memory.created_at);
                    if !memory.tags.is_empty() {
                        println!("Tags: {:?}", memory.tags);
                    }
                }
                None => {
                    println!("Content not found");
                }
            }
        }
        Commands::Stats => {
            // Create database pool and run migrations
            let pool = codex_memory::database::create_pool(&config.database_url).await?;
            codex_memory::database::run_migrations(&pool).await?;
            let storage = Arc::new(Storage::new(pool));

            let stats = storage.stats().await?;
            println!("Total memories: {}", stats.total_memories);
            println!("Table size: {}", stats.table_size);
            if let Some(last) = stats.last_memory_created {
                println!("Last created: {}", last);
            }
        }
    }

    Ok(())
}