Merge pull request #86 from redis/feature/token-cli-json-output

Add JSON output for token CLI token commands

Author: abrookins

agent_memory_server/cli.py

@@ -3,6 +3,7 @@
 """
 
 import importlib
+import json
 import sys
 from datetime import UTC, datetime, timedelta
 
@@ -305,7 +306,26 @@ def token():
 @token.command()
 @click.option("--description", "-d", required=True, help="Token description")
 @click.option("--expires-days", "-e", type=int, help="Token expiration in days")
-def add(description: str, expires_days: int | None):
+@click.option(
+    "--format",
+    "output_format",
+    type=click.Choice(["text", "json"]),
+    default="text",
+    show_default=True,
+    help="Output format.",
+)
+@click.option(
+    "--token",
+    "provided_token",
+    type=str,
+    help="Use a pre-generated token instead of generating a new one.",
+)
+def add(
+    description: str,
+    expires_days: int | None,
+    output_format: str,
+    provided_token: str | None,
+) -> None:
     """Add a new authentication token."""
     import asyncio
 
@@ -315,9 +335,9 @@ def add(description: str, expires_days: int | None):
     async def create_token():
         redis = await get_redis_conn()
 
-        # Generate token
-        token = generate_token()
-        token_hash = hash_token(token)
+        # Determine token value
+        token_value = provided_token or generate_token()
+        token_hash = hash_token(token_value)
 
         # Calculate expiration
         now = datetime.now(UTC)
@@ -344,20 +364,43 @@ async def create_token():
         list_key = Keys.auth_tokens_list_key()
         await redis.sadd(list_key, token_hash)
 
+        return token_value, token_info
+
+    token, token_info = asyncio.run(create_token())
+
+    if output_format == "json":
+        data = {
+            "token": token,
+            "description": token_info.description,
+            "created_at": token_info.created_at.isoformat(),
+            "expires_at": token_info.expires_at.isoformat()
+            if token_info.expires_at
+            else None,
+            "hash": token_info.token_hash,
+        }
+        click.echo(json.dumps(data))
+    else:
+        expires_at = token_info.expires_at
         click.echo("Token created successfully!")
         click.echo(f"Token: {token}")
-        click.echo(f"Description: {description}")
+        click.echo(f"Description: {token_info.description}")
         if expires_at:
             click.echo(f"Expires: {expires_at.isoformat()}")
         else:
             click.echo("Expires: Never")
         click.echo("\nWARNING: Save this token securely. It will not be shown again.")
 
-    asyncio.run(create_token())
-
 
 @token.command()
-def list():
+@click.option(
+    "--format",
+    "output_format",
+    type=click.Choice(["text", "json"]),
+    default="text",
+    show_default=True,
+    help="Output format.",
+)
+def list(output_format: str):
     """List all authentication tokens."""
     import asyncio
 
@@ -371,12 +414,16 @@ async def list_tokens():
         list_key = Keys.auth_tokens_list_key()
         token_hashes = await redis.smembers(list_key)
 
+        tokens_data = []
+
         if not token_hashes:
-            click.echo("No tokens found.")
-            return
+            if output_format == "text":
+                click.echo("No tokens found.")
+            return tokens_data
 
-        click.echo("Authentication Tokens:")
-        click.echo("=" * 50)
+        if output_format == "text":
+            click.echo("Authentication Tokens:")
+            click.echo("=" * 50)
 
         for token_hash in token_hashes:
             key = Keys.auth_token_key(token_hash)
@@ -390,27 +437,62 @@ async def list_tokens():
             try:
                 token_info = TokenInfo.model_validate_json(token_data)
 
-                # Mask the token hash for display
-                masked_hash = token_hash[:8] + "..." + token_hash[-8:]
-
-                click.echo(f"Token: {masked_hash}")
-                click.echo(f"Description: {token_info.description}")
-                click.echo(f"Created: {token_info.created_at.isoformat()}")
-                if token_info.expires_at:
-                    click.echo(f"Expires: {token_info.expires_at.isoformat()}")
-                else:
-                    click.echo("Expires: Never")
-                click.echo("-" * 30)
+                tokens_data.append(
+                    {
+                        "hash": token_hash,
+                        "description": token_info.description,
+                        "created_at": token_info.created_at.isoformat(),
+                        "expires_at": token_info.expires_at.isoformat()
+                        if token_info.expires_at
+                        else None,
+                        "status": (
+                            "Never Expires"
+                            if not token_info.expires_at
+                            else (
+                                "EXPIRED"
+                                if datetime.now(UTC) > token_info.expires_at
+                                else "Active"
+                            )
+                        ),
+                    }
+                )
+
+                if output_format == "text":
+                    # Mask the token hash for display
+                    masked_hash = token_hash[:8] + "..." + token_hash[-8:]
+
+                    click.echo(f"Token: {masked_hash}")
+                    click.echo(f"Description: {token_info.description}")
+                    click.echo(f"Created: {token_info.created_at.isoformat()}")
+                    if token_info.expires_at:
+                        click.echo(f"Expires: {token_info.expires_at.isoformat()}")
+                    else:
+                        click.echo("Expires: Never")
+                    click.echo("-" * 30)
 
             except Exception as e:
-                click.echo(f"Error processing token {token_hash}: {e}")
+                if output_format == "text":
+                    click.echo(f"Error processing token {token_hash}: {e}")
+
+        return tokens_data
+
+    tokens_data = asyncio.run(list_tokens())
 
-    asyncio.run(list_tokens())
+    if output_format == "json":
+        click.echo(json.dumps(tokens_data))
 
 
 @token.command()
 @click.argument("token_hash")
-def show(token_hash: str):
+@click.option(
+    "--format",
+    "output_format",
+    type=click.Choice(["text", "json"]),
+    default="text",
+    show_default=True,
+    help="Output format.",
+)
+def show(token_hash: str, output_format: str):
     """Show details for a specific token."""
     import asyncio
 
@@ -430,45 +512,90 @@ async def show_token():
             matching_hashes = [h for h in token_hashes if h.startswith(token_hash)]
 
             if not matching_hashes:
-                click.echo(f"No token found matching '{token_hash}'")
-                return
+                if output_format == "json":
+                    click.echo(
+                        json.dumps({"error": f"No token found matching '{token_hash}'"})
+                    )
+                    sys.exit(1)
+                else:
+                    click.echo(f"No token found matching '{token_hash}'")
+                    sys.exit(1)
             if len(matching_hashes) > 1:
-                click.echo(f"Multiple tokens match '{token_hash}':")
-                for h in matching_hashes:
-                    click.echo(f"  {h[:8]}...{h[-8:]}")
-                return
+                if output_format == "json":
+                    click.echo(
+                        json.dumps(
+                            {
+                                "error": f"Multiple tokens match '{token_hash}'",
+                                "matches": [
+                                    f"{h[:8]}...{h[-8:]}" for h in matching_hashes
+                                ],
+                            }
+                        )
+                    )
+                    sys.exit(1)
+                else:
+                    click.echo(f"Multiple tokens match '{token_hash}':")
+                    for h in matching_hashes:
+                        click.echo(f"  {h[:8]}...{h[-8:]}")
+                    sys.exit(1)
             token_hash = matching_hashes[0]
 
         key = Keys.auth_token_key(token_hash)
         token_data = await redis.get(key)
 
         if not token_data:
-            click.echo(f"Token not found: {token_hash}")
-            return
+            if output_format == "json":
+                click.echo(json.dumps({"error": f"Token not found: {token_hash}"}))
+                sys.exit(1)
+            else:
+                click.echo(f"Token not found: {token_hash}")
+                sys.exit(1)
 
         try:
             token_info = TokenInfo.model_validate_json(token_data)
 
-            click.echo("Token Details:")
-            click.echo("=" * 30)
-            click.echo(f"Hash: {token_hash}")
-            click.echo(f"Description: {token_info.description}")
-            click.echo(f"Created: {token_info.created_at.isoformat()}")
-            if token_info.expires_at:
-                click.echo(f"Expires: {token_info.expires_at.isoformat()}")
-                # Check if expired
-                if datetime.now(UTC) > token_info.expires_at:
-                    click.echo("Status: EXPIRED")
-                else:
-                    click.echo("Status: Active")
+            if token_info.expires_at and datetime.now(UTC) > token_info.expires_at:
+                status = "EXPIRED"
             else:
-                click.echo("Expires: Never")
-                click.echo("Status: Active")
+                status = "Active"
 
-        except Exception as e:
-            click.echo(f"Error processing token: {e}")
+            return token_hash, token_info, status
 
-    asyncio.run(show_token())
+        except Exception as e:
+            if output_format == "json":
+                click.echo(
+                    json.dumps({"error": f"Failed to parse token data: {str(e)}"})
+                )
+                sys.exit(1)
+            else:
+                click.echo(f"Error processing token: {e}", err=True)
+                sys.exit(1)
+
+    result = asyncio.run(show_token())
+    token_hash, token_info, status = result
+
+    if output_format == "json":
+        data = {
+            "hash": token_hash,
+            "description": token_info.description,
+            "created_at": token_info.created_at.isoformat(),
+            "expires_at": token_info.expires_at.isoformat()
+            if token_info.expires_at
+            else None,
+            "status": status,
+        }
+        click.echo(json.dumps(data))
+    else:
+        click.echo("Token Details:")
+        click.echo("=" * 30)
+        click.echo(f"Hash: {token_hash}")
+        click.echo(f"Description: {token_info.description}")
+        click.echo(f"Created: {token_info.created_at.isoformat()}")
+        if token_info.expires_at:
+            click.echo(f"Expires: {token_info.expires_at.isoformat()}")
+        else:
+            click.echo("Expires: Never")
+        click.echo(f"Status: {status}")
 
 
 @token.command()