Merge pull request #18 from litentry/fix/issue-17

fix(cli): #17 revoke command — self-revoke + revoke-by-wallet

Author: hanwencheng

Cargo.lock

@@ -26,8 +26,10 @@ reqwest = { version = "0.12", features = ["json"] }
 assert_cmd = "2"
 predicates = "3"
 agentkeys-mock-server = { path = "../agentkeys-mock-server" }
+agentkeys-types = { workspace = true }
 tokio = { workspace = true }
 reqwest = { version = "0.12", features = ["json"] }
 axum = { version = "0.7", features = ["json"] }
 rusqlite = { version = "0.31", features = ["bundled"] }
 serde_json = { workspace = true }
+tempfile = "3"

crates/agentkeys-cli/Cargo.toml

@@ -200,32 +200,58 @@ pub async fn cmd_run(ctx: &CommandContext, agent: &str, cmd: &[String]) -> Resul
     Ok(String::new())
 }
 
-pub async fn cmd_revoke(ctx: &CommandContext, agent: &str) -> Result<String> {
+pub async fn cmd_revoke(ctx: &CommandContext, agent: Option<&str>) -> Result<String> {
     let session = ctx.load_session().context("load session (run `agentkeys init` first)")?;
 
-    let target_session = Session {
-        token: agent.to_string(),
-        wallet: WalletAddress(agent.to_string()),
-        scope: None,
-        created_at: 0,
-        ttl_seconds: 0,
-    };
-
     if ctx.verbose {
         eprintln!("[verbose] POST {}/session/revoke", ctx.backend_url);
-        eprintln!("[verbose] target: {}", agent);
     }
 
-    ctx.backend()
-        .revoke_session(&session, &target_session)
-        .await
-        .map_err(wrap_backend_error)?;
-
-    let now = std::time::SystemTime::now()
-        .duration_since(std::time::UNIX_EPOCH)
-        .map(|d| d.as_secs())
-        .unwrap_or(0);
-    Ok(format!("Revoked agent={} at timestamp={}", agent, now))
+    match agent {
+        None => {
+            let wallet_display = session.wallet.0.clone();
+            ctx.backend()
+                .revoke_session(&session, &session)
+                .await
+                .map_err(wrap_backend_error)?;
+            session_store::clear_session().context("clear local session")?;
+            Ok(format!(
+                "Revoked current session for wallet={}. Local session wiped. Run `agentkeys init` to re-pair.",
+                wallet_display
+            ))
+        }
+        Some(target_wallet_str) => {
+            if ctx.verbose {
+                eprintln!("[verbose] target wallet: {}", target_wallet_str);
+            }
+            let target_wallet = WalletAddress(target_wallet_str.to_string());
+            ctx.backend()
+                .revoke_by_wallet(&session, &target_wallet)
+                .await
+                .map_err(wrap_backend_error)?;
+
+            // If the target wallet IS the caller's own wallet, the just-revoked
+            // session matches the locally-cached one. Wipe local state too so
+            // subsequent commands fail cleanly with "no session" instead of
+            // loading the stale revoked token (codex P2 from the original review,
+            // tracked at issue-17 review thread).
+            //
+            // Wallet addresses are compared case-insensitively because the EVM
+            // canonical form (EIP-55 mixed case) can differ from the lowercase
+            // form returned by the mock backend.
+            let revoked_self = session.wallet.0.eq_ignore_ascii_case(target_wallet_str);
+            if revoked_self {
+                session_store::clear_session()
+                    .context("clear local session after self-revoke")?;
+                Ok(format!(
+                    "Revoked agent={} (was your own session — local state wiped, run `agentkeys init` to re-pair).",
+                    target_wallet_str
+                ))
+            } else {
+                Ok(format!("Revoked agent={}", target_wallet_str))
+            }
+        }
+    }
 }
 
 pub async fn cmd_teardown(ctx: &CommandContext, agent: &str) -> Result<String> {

crates/agentkeys-cli/src/lib.rs

@@ -73,12 +73,12 @@ enum Commands {
     },
 
     #[command(
-        about = "Revoke an agent's session",
-        long_about = "Immediately invalidate an agent's session token.\n\nExamples:\n  agentkeys revoke 0xAGENT"
+        about = "Revoke a session",
+        long_about = "Revoke a session. Without arguments, revokes the current session and wipes the local keychain entry (you must run `agentkeys init` again). With a wallet address, revokes all active sessions for that child agent (ownership check enforced).\n\nExamples:\n  agentkeys revoke\n  agentkeys revoke 0xCHILD_WALLET"
     )]
     Revoke {
-        #[arg(help = "Agent wallet address or session token to revoke")]
-        agent: String,
+        #[arg(help = "Child agent wallet address to revoke (omit to revoke your own current session)", required = false)]
+        agent: Option<String>,
     },
 
     #[command(
@@ -155,7 +155,7 @@ async fn main() {
         Commands::Store { agent, service, key } => cmd_store(&ctx, agent, service, key).await,
         Commands::Read { agent, service } => cmd_read(&ctx, agent, service).await,
         Commands::Run { agent, cmd } => cmd_run(&ctx, agent, cmd).await,
-        Commands::Revoke { agent } => cmd_revoke(&ctx, agent).await,
+        Commands::Revoke { agent } => cmd_revoke(&ctx, agent.as_deref()).await,
         Commands::Teardown { agent } => cmd_teardown(&ctx, agent).await,
         Commands::Usage { agent, json } => {
             cmd_usage(&ctx, agent.as_deref(), *json).await

crates/agentkeys-cli/src/main.rs

@@ -5,7 +5,7 @@ use std::path::PathBuf;
 const KEYRING_SERVICE: &str = "agentkeys";
 const KEYRING_USER: &str = "session";
 
-fn fallback_path() -> PathBuf {
+pub fn fallback_path() -> PathBuf {
     let home = std::env::var("HOME")
         .or_else(|_| std::env::var("USERPROFILE"))
         .unwrap_or_else(|_| ".".to_string());
@@ -91,3 +91,57 @@ fn load_from_file() -> Result<Session> {
         .with_context(|| format!("read session file at {}", path.display()))?;
     serde_json::from_str(&json).context("deserialize session from file")
 }
+
+/// Delete the locally stored session from both keyring and file.
+/// Best-effort: ignores "not found" errors. Returns Err only if both
+/// attempts failed with non-NotFound errors.
+///
+/// When `AGENTKEYS_SESSION_STORE=file` is set, the keyring branch is skipped
+/// entirely (no 2-second timeout, no chance of spurious keyring errors).
+pub fn clear_session() -> Result<()> {
+    let keyring_result = if should_skip_keyring() {
+        Ok(())
+    } else {
+        try_keyring_delete()
+    };
+    let file_result = delete_session_file();
+
+    match (keyring_result, file_result) {
+        (Err(ke), Err(fe)) => {
+            Err(anyhow::anyhow!("could not clear session: keyring: {}; file: {}", ke, fe))
+        }
+        _ => Ok(()),
+    }
+}
+
+fn try_keyring_delete() -> Result<()> {
+    let (tx, rx) = std::sync::mpsc::channel::<Result<()>>();
+    std::thread::spawn(move || {
+        let result = keyring::Entry::new(KEYRING_SERVICE, KEYRING_USER)
+            .map_err(|e| anyhow::anyhow!("{}", e))
+            .and_then(|e| e.delete_password().map_err(|ke| anyhow::anyhow!("{}", ke)));
+        let _ = tx.send(result);
+    });
+    match rx.recv_timeout(std::time::Duration::from_secs(2)) {
+        Ok(Ok(())) => Ok(()),
+        Ok(Err(e)) => {
+            let msg = e.to_string().to_lowercase();
+            if msg.contains("not found") || msg.contains("no such") || msg.contains("no password") {
+                Ok(())
+            } else {
+                Err(e)
+            }
+        }
+        Err(_) => Ok(()),
+    }
+}
+
+fn delete_session_file() -> Result<()> {
+    let path = fallback_path();
+    match std::fs::remove_file(&path) {
+        Ok(()) => Ok(()),
+        Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(()),
+        Err(e) => Err(anyhow::anyhow!("remove {}: {}", path.display(), e)),
+    }
+}
+