hotfix/cli with server support python protocol (#172)

* Fix python protocol bug in cli, Add --with-server to cli

* bump version

* Revert rogue dir

Author: yuval-qf

VERSION

@@ -1 +1 @@
-0.6.3
+0.6.4

rogue/__main__.py

@@ -3,6 +3,7 @@
 import sys
 import time
 from argparse import ArgumentParser, Namespace
+from ipaddress import ip_address
 from pathlib import Path
 
 import platformdirs
@@ -100,6 +101,13 @@ def parse_args() -> Namespace:
         help="Run in non-interactive CLI mode",
         parents=[common_parser()],
     )
+    cli_parser.add_argument(
+        "--with-server",
+        action="store_true",
+        default=False,
+        help="Start the rogue server alongside the CLI",
+    )
+    set_server_args(cli_parser)
     set_cli_args(cli_parser)
 
     # TUI mode
@@ -119,6 +127,20 @@ def parse_args() -> Namespace:
     return parser.parse_args()
 
 
+def build_local_server_url(host: str, port: int) -> str:
+    """Build an HTTP URL for a server bound to host:port, normalizing wildcards."""
+    http_host = host
+    if http_host in ("0.0.0.0", "::"):  # nosec B104
+        http_host = "127.0.0.1"
+    else:
+        try:
+            if ip_address(http_host).version == 6:
+                http_host = f"[{http_host}]"
+        except ValueError:
+            pass  # hostname, leave as-is
+    return f"http://{http_host}:{port}"
+
+
 def start_example_agent(
     example_name: str,
     host: str,
@@ -279,7 +301,32 @@ def main() -> None:
         if args.mode == "server":
             run_server(args, background=False)
         elif args.mode == "cli":
-            exit_code = asyncio.run(run_cli(args))
+            server_process = None
+            if args.with_server:
+                server_process = run_server(
+                    args,
+                    background=True,
+                    log_file=log_file_path,
+                )
+                if not server_process:
+                    logger.error("Failed to start rogue server. Exiting.")
+                    sys.exit(1)
+                # Point the CLI at the embedded server.
+                args.rogue_server_url = build_local_server_url(
+                    args.host,
+                    args.port,
+                )
+                logger.info(
+                    "Rogue server started for CLI",
+                    extra={"rogue_server_url": args.rogue_server_url},
+                )
+
+            try:
+                exit_code = asyncio.run(run_cli(args))
+            finally:
+                if server_process:
+                    server_process.terminate()
+                    server_process.join()
             sys.exit(exit_code)
         elif args.mode == "tui":
             if not RogueTuiInstaller().install_rogue_tui():

sdks/python/rogue_sdk/types.py

@@ -178,7 +178,7 @@ class RedTeamConfig(BaseModel):
     frameworks: List[str] = Field(
         default_factory=list,
         description=(
-            "Framework IDs for report mapping " "(e.g., 'owasp-llm', 'mitre-atlas')"
+            "Framework IDs for report mapping (e.g., 'owasp-llm', 'mitre-atlas')"
         ),
     )
     random_seed: Optional[int] = Field(
@@ -293,7 +293,7 @@ class AgentConfig(BaseModel):
     """Configuration for the agent being evaluated."""
 
     protocol: Protocol = Protocol.A2A
-    transport: Transport = None  # type: ignore # fixed in model_post_init
+    transport: Transport | None = None  # type: ignore # fixed in model_post_init
     evaluated_agent_url: Optional[HttpUrl] = None  # Optional for PYTHON protocol
     evaluated_agent_auth_type: AuthType = Field(
         default=AuthType.NO_AUTH,
@@ -337,6 +337,27 @@ def check_auth_credentials(self) -> "AgentConfig":
             )
         return self
 
+    @model_validator(mode="after")
+    def validate_transport(self) -> "AgentConfig":
+        if self.protocol == Protocol.PYTHON:
+            if self.transport is not None:
+                raise ValueError(
+                    "Transport cannot be provided when protocol is PYTHON",
+                )
+            else:
+                return self
+
+        if self.transport is None:
+            raise ValueError(
+                "Transport is required when protocol is not PYTHON",
+            )
+
+        if not self.transport.is_valid_for_protocol(self.protocol):  # type: ignore
+            raise ValueError(
+                "Transport is not valid for the selected protocol",
+            )
+        return self
+
     @model_validator(mode="after")
     def check_protocol_requirements(self) -> "AgentConfig":
         """Validate protocol-specific requirements."""