Support machine-readable token CLI output

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,15 @@ 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.",
+)
+def add(description: str, expires_days: int | None, output_format: str):
     """Add a new authentication token."""
     import asyncio
 
@@ -344,20 +353,43 @@ async def create_token():
         list_key = Keys.auth_tokens_list_key()
         await redis.sadd(list_key, token_hash)
 
+        return token, 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 +403,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 +426,57 @@ 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,
+                        "expired": bool(
+                            token_info.expires_at
+                            and datetime.now(UTC) > token_info.expires_at
+                        ),
+                    }
+                )
+
+                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 +496,69 @@ 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 == "text":
+                    click.echo(f"No token found matching '{token_hash}'")
+                return None
             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 == "text":
+                    click.echo(f"Multiple tokens match '{token_hash}':")
+                    for h in matching_hashes:
+                        click.echo(f"  {h[:8]}...{h[-8:]}")
+                return None
             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 == "text":
+                click.echo(f"Token not found: {token_hash}")
+            return None
 
         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 == "text":
+                click.echo(f"Error processing token: {e}")
+            return None
+
+    result = asyncio.run(show_token())
+
+    if result is None:
+        return
+
+    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()

tests/test_token_cli.py

@@ -60,6 +60,41 @@ def test_token_add_command_no_expiry(self, mock_get_redis, mock_redis, cli_runne
         mock_redis.expire.assert_not_called()  # No expiry set
         mock_redis.sadd.assert_called_once()
 
+    @patch("agent_memory_server.cli.get_redis_conn")
+    def test_token_add_command_json_output(
+        self, mock_get_redis, mock_redis, cli_runner
+    ):
+        """Test token add command with JSON output."""
+        mock_get_redis.return_value = mock_redis
+
+        import json
+
+        result = cli_runner.invoke(
+            token,
+            [
+                "add",
+                "--description",
+                "Test token",
+                "--expires-days",
+                "30",
+                "--format",
+                "json",
+            ],
+        )
+
+        assert result.exit_code == 0
+
+        data = json.loads(result.output)
+        assert data["description"] == "Test token"
+        assert data["token"]
+        assert data["hash"]
+        assert data["expires_at"] is not None
+
+        # Verify Redis calls
+        mock_redis.set.assert_called_once()
+        mock_redis.expire.assert_called_once()
+        mock_redis.sadd.assert_called_once()
+
     @patch("agent_memory_server.cli.get_redis_conn")
     def test_token_list_command_empty(self, mock_get_redis, mock_redis, cli_runner):
         """Test token list command with no tokens."""
@@ -97,6 +132,38 @@ def test_token_list_command_with_tokens(
         assert "Test token" in result.output
         assert "test_has...34567890" in result.output  # Masked hash
 
+    @patch("agent_memory_server.cli.get_redis_conn")
+    def test_token_list_command_json_output(
+        self, mock_get_redis, mock_redis, cli_runner
+    ):
+        """Test token list command with JSON output."""
+        mock_get_redis.return_value = mock_redis
+
+        import json
+
+        # Create sample token data
+        token_hash = "test_hash_123456789012345678901234567890"
+        token_info = TokenInfo(
+            description="Test token",
+            created_at=datetime.now(UTC),
+            expires_at=datetime.now(UTC) + timedelta(days=30),
+            token_hash=token_hash,
+        )
+
+        mock_redis.smembers.return_value = {token_hash}
+        mock_redis.get.return_value = token_info.model_dump_json()
+
+        result = cli_runner.invoke(token, ["list", "--format", "json"])
+
+        assert result.exit_code == 0
+
+        data = json.loads(result.output)
+        assert isinstance(data, list)
+        assert len(data) == 1
+        item = data[0]
+        assert item["hash"] == token_hash
+        assert item["description"] == "Test token"
+
     @patch("agent_memory_server.cli.get_redis_conn")
     def test_token_show_command(self, mock_get_redis, mock_redis, cli_runner):
         """Test token show command."""
@@ -120,6 +187,43 @@ def test_token_show_command(self, mock_get_redis, mock_redis, cli_runner):
         assert "Test token" in result.output
         assert "Status: Active" in result.output
 
+    @patch("agent_memory_server.cli.get_redis_conn")
+    def test_token_show_command_json_output(
+        self, mock_get_redis, mock_redis, cli_runner
+    ):
+        """Test token show command with JSON output."""
+        mock_get_redis.return_value = mock_redis
+
+        import json
+
+        # Create sample token data
+        token_hash = "test_hash_123456789012345678901234567890"
+        token_info = TokenInfo(
+            description="Test token",
+            created_at=datetime.now(UTC),
+            expires_at=datetime.now(UTC) + timedelta(days=30),
+            token_hash=token_hash,
+        )
+
+        mock_redis.get.return_value = token_info.model_dump_json()
+
+        result = cli_runner.invoke(
+            token,
+            [
+                "show",
+                "--format",
+                "json",
+                token_hash,
+            ],
+        )
+
+        assert result.exit_code == 0
+
+        data = json.loads(result.output)
+        assert data["hash"] == token_hash
+        assert data["description"] == "Test token"
+        assert data["status"] == "Active"
+
     @patch("agent_memory_server.cli.get_redis_conn")
     def test_token_show_command_partial_hash(
         self, mock_get_redis, mock_redis, cli_runner