Fix code-review findings: JSON list output, agent/TTS robustness, auth-flow structure

.importlinter

@@ -28,6 +28,7 @@ source_modules =
     aai_cli.streaming
     aai_cli.theme
     aai_cli.transcribe_render
+    aai_cli.ws
     aai_cli.youtube
 forbidden_modules =
     aai_cli.commands

aai_cli/agent/session.py

@@ -3,14 +3,14 @@
 import base64
 import contextlib
 import json
-import logging
 import threading
 from collections.abc import Callable
 from dataclasses import dataclass
 from typing import Any
 
 from aai_cli import environments
-from aai_cli.errors import APIError, CLIError, NotAuthenticated, auth_failure, is_auth_failure
+from aai_cli import ws as wsutil
+from aai_cli.errors import APIError, CLIError, NotAuthenticated
 
 
 def ws_url() -> str:
@@ -33,8 +33,8 @@ def ws_url() -> str:
 # NotAuthenticated code every other rejected-credential path across the CLI uses.
 _AUTH_ERROR_CODES = {"UNAUTHORIZED", "FORBIDDEN"}
 
-# A pre-upgrade HTTP 403 on the WebSocket handshake (see _is_rejected_key).
-_HTTP_FORBIDDEN = 403
+# How long a file-driven run waits for the server's session.ready before giving up.
+_READY_TIMEOUT_SECONDS = 10
 
 # The websocket connection, the `connect` factory, and the renderer/player/mic I/O
 # objects come from libraries/modules with no usable type stubs. Alias that untyped
@@ -155,9 +155,12 @@ def raise_error(self, event: dict[str, Any]) -> None:
         code = event.get("code", "")
         message = event.get("message") or code or "Voice agent error."
         if code in _AUTH_ERROR_CODES:
+            # rejected_key: a credential was presented and the server refused it, so
+            # auto-login must not retry when the key came from ASSEMBLYAI_API_KEY.
             raise NotAuthenticated(
                 f"Voice agent rejected the connection: {message}",
                 suggestion="Run 'aai login' with a valid key, or set ASSEMBLYAI_API_KEY.",
+                rejected_key=True,
             )
         raise APIError(f"Voice agent error ({code}): {message}")
 
@@ -181,8 +184,16 @@ def _send_audio_loop(ws: _WebSocket, session: VoiceAgentSession, mic: _IO) -> No
     """Forward mic PCM as input.audio while the session gate allows it."""
     # File-driven runs wait for session.ready before consuming the source, so a
     # finite clip isn't partly drained (and dropped) before the server accepts it.
-    if session.ready_event is not None:
-        session.ready_event.wait(timeout=10)
+    # A server that never becomes ready fails the run loudly: continuing would
+    # silently discard the start of the clip frame by frame.
+    if session.ready_event is not None and not session.ready_event.wait(
+        timeout=_READY_TIMEOUT_SECONDS
+    ):
+        raise APIError(
+            f"Voice agent session was not ready after {_READY_TIMEOUT_SECONDS}s; "
+            "no audio was sent.",
+            suggestion="Try again; if it persists, check https://status.assemblyai.com.",
+        )
     for chunk in mic:
         if not session.should_send_audio():
             continue  # half-duplex: drop frames while the agent is speaking
@@ -193,54 +204,12 @@ def _send_audio_loop(ws: _WebSocket, session: VoiceAgentSession, mic: _IO) -> No
             return
 
 
-# The sync websockets client logs through these; both are silenced for the session
-# (the parent covers any future child logger, the client logger is the one that fires).
-_WEBSOCKETS_LOGGERS = ("websockets", "websockets.client")
-
-
-def _silence_websockets_logging() -> None:
-    """Keep websockets' internal logging off the user's stderr for the session.
-
-    The sync client's background reader thread logs unhandled teardown errors (e.g.
-    ``EOFError: stream ended``) as "unexpected internal error" + traceback through the
-    ``websockets.client`` logger, which would land on stderr right next to our clean
-    CLIError. Those internals are never user-actionable from the CLI, so raise the
-    loggers above every level they emit at. Idempotent: re-setting the level is a no-op.
-    """
-    for name in _WEBSOCKETS_LOGGERS:
-        logging.getLogger(name).setLevel(logging.CRITICAL)
-
-
-def _is_rejected_key(exc: Exception) -> bool:
-    """Is this connect/session failure auth-shaped (the key itself was rejected)?
-
-    Mirrors how `stream` classifies handshake failures: a plain HTTP 403 on the
-    WebSocket upgrade stays an API error there ("Streaming error: WebSocket handshake
-    rejected (HTTP 403)"), so it must not become "Your API key was rejected" here —
-    403 also covers non-credential blocks (WAF, region, plan). Only 401, the Voice
-    Agent's 1008 policy-violation close, or an explicitly auth-worded message
-    (`is_auth_failure`'s text hints) count as a rejected key.
-    """
-    status = getattr(getattr(exc, "response", None), "status_code", None)
-    if status == _HTTP_FORBIDDEN:
-        return False
-    return is_auth_failure(exc)
-
-
-def _auth_or_api_error(exc: Exception, message: str) -> CLIError:
-    """Map a connect/session exception to the right CLIError: a rejected key becomes
-    auth_failure(), anything else becomes APIError(f"{message}: {exc}")."""
-    if _is_rejected_key(exc):
-        return auth_failure()
-    return APIError(f"{message}: {exc}")
-
-
 def _open_ws(connect: _Connect, api_key: str) -> _WebSocket:
     """Open the Voice Agent socket, mapping a connect failure to a clean CLIError."""
     try:
         return connect(ws_url(), additional_headers={"Authorization": f"Bearer {api_key}"})
     except Exception as exc:
-        raise _auth_or_api_error(exc, "Could not connect to the voice agent") from exc
+        raise wsutil.auth_or_api_error(exc, "Could not connect to the voice agent") from exc
 
 
 def _session_update_message(config: AgentRunConfig) -> str:
@@ -281,7 +250,7 @@ def run_session(
     the agent's first reply to the spoken input and the capture thread waits for
     session.ready before streaming the source.
     """
-    _silence_websockets_logging()
+    wsutil.silence_websockets_logging()
     if connect is None:
         from websockets.sync.client import connect
 
@@ -324,7 +293,7 @@ def _capture() -> None:
     except Exception as exc:
         if capture_error:
             raise capture_error[0] from exc  # a mic-open failure is the real cause
-        raise _auth_or_api_error(exc, "Voice agent session failed") from exc
+        raise wsutil.auth_or_api_error(exc, "Voice agent session failed") from exc
     finally:
         with contextlib.suppress(Exception):
             ws.close()

aai_cli/commands/agent.py

@@ -66,6 +66,12 @@ def _open_audio(
     return duplex.mic, duplex.player
 
 
+def _emit_voice_list(_state: AppState, json_mode: bool) -> None:
+    """--list-voices body, routed through run_command so --json yields a
+    machine-readable array instead of the human list; needs no auth."""
+    output.emit(VOICES, lambda _voices: format_voice_list(), json_mode=json_mode)
+
+
 @app.command(
     rich_help_panel=help_panels.TRANSCRIPTION,
     epilog=examples_epilog(
@@ -129,8 +135,8 @@ def agent(
     """
 
     if list_voices:
-        typer.echo(format_voice_list())
-        raise typer.Exit(code=0)
+        run_command(ctx, _emit_voice_list, json=json_out)
+        return
 
     def body(state: AppState, json_mode: bool) -> None:
         validate_output_flags(json_mode=json_mode, output_field=output_field)

aai_cli/commands/llm.py

@@ -42,6 +42,12 @@ def _validate_follow_args(
     return prompt
 
 
+def _emit_model_list(_state: AppState, json_mode: bool) -> None:
+    """--list-models body, routed through run_command so --json yields a
+    machine-readable array instead of the human list; needs no auth."""
+    output.emit(list(gateway.KNOWN_MODELS), "\n".join, json_mode=json_mode)
+
+
 @app.command(
     rich_help_panel=help_panels.TRANSCRIPTION,
     epilog=examples_epilog(
@@ -100,8 +106,8 @@ def llm(
     """
 
     if list_models:
-        typer.echo("\n".join(gateway.KNOWN_MODELS))
-        raise typer.Exit(code=0)
+        run_command(ctx, _emit_model_list, json=json_out)
+        return
 
     def follow_body(state: AppState, json_mode: bool) -> None:
         prompt_text = _validate_follow_args(prompt, output_field, transcript_id)

aai_cli/commands/login.py

@@ -83,7 +83,9 @@ def render(_d: object) -> str:
             json_mode=json_mode,
         )
 
-    run_command(ctx, body, json=json_out)
+    # auto_login=False: this command owns sign-in; an auth failure here (e.g. the
+    # browser flow timing out) must surface, not trigger a second login attempt.
+    run_command(ctx, body, json=json_out, auto_login=False)
 
 
 @app.command(
@@ -114,7 +116,8 @@ def body(state: AppState, json_mode: bool) -> None:
             json_mode=json_mode,
         )
 
-    run_command(ctx, body, json=json_out)
+    # auto_login=False: signing out while signed out must not start a sign-in.
+    run_command(ctx, body, json=json_out, auto_login=False)
 
 
 @app.command(

aai_cli/context.py

@@ -12,7 +12,7 @@
 from aai_cli import config, environments, output
 from aai_cli.auth import run_login_flow
 from aai_cli.environments import Environment
-from aai_cli.errors import REJECTED_KEY_MESSAGE, APIError, CLIError, NotAuthenticated
+from aai_cli.errors import APIError, CLIError, NotAuthenticated
 
 
 @dataclass
@@ -137,19 +137,18 @@ def _interactive_session() -> bool:
     return sys.stdin.isatty() and sys.stderr.isatty() and not output.is_agentic()
 
 
-def _should_auto_login(ctx: typer.Context, err: NotAuthenticated) -> bool:
-    command_name = ctx.command.name if ctx.command else None
-    if command_name in {"login", "logout"}:
-        return False
+def _should_auto_login(err: NotAuthenticated) -> bool:
     # CI/pipelines/agents have no human to finish a browser sign-in; starting one
     # would bind a loopback port and block for up to two minutes. Surface the
     # original NotAuthenticated (with its 'aai login' / ASSEMBLYAI_API_KEY
     # suggestion) instead.
     if not _interactive_session():
         return False
     # An invalid ASSEMBLYAI_API_KEY would still take precedence after browser login,
-    # so retrying cannot fix that case.
-    return not (os.environ.get(config.ENV_API_KEY) and err.message == REJECTED_KEY_MESSAGE)
+    # so retrying cannot fix that case. `rejected_key` is the structured marker set
+    # by auth_failure(); auth-owning commands (login/logout) opt out at their
+    # run_command call site with auto_login=False instead of being name-matched here.
+    return not (os.environ.get(config.ENV_API_KEY) and err.rejected_key)
 
 
 def _auto_login_and_exit(state: AppState, *, json_mode: bool) -> NoReturn:
@@ -191,7 +190,7 @@ def run_command(
     try:
         fn(state, json_mode)
     except NotAuthenticated as err:
-        if not auto_login or not _should_auto_login(ctx, err):
+        if not auto_login or not _should_auto_login(err):
             output.emit_error(err, json_mode=json_mode)
             raise typer.Exit(code=err.exit_code) from None
         _auto_login_and_exit(state, json_mode=json_mode)

aai_cli/errors.py

@@ -61,10 +61,14 @@ def __init__(
             "Run 'aai onboard' for guided setup, 'aai login' if you have an account, "
             "or set ASSEMBLYAI_API_KEY."
         ),
+        rejected_key: bool = False,
     ) -> None:
         super().__init__(
             message, error_type="not_authenticated", exit_code=4, suggestion=suggestion
         )
+        # Structured marker for "a key was presented and the server rejected it" (vs
+        # "no credential at all"), so callers like auto-login don't match message text.
+        self.rejected_key = rejected_key
 
 
 class APIError(CLIError):
@@ -142,4 +146,6 @@ def is_auth_failure(exc: object) -> bool:
 
 def auth_failure() -> NotAuthenticated:
     """A NotAuthenticated for the 'key present but rejected by the server' case."""
-    return NotAuthenticated(REJECTED_KEY_MESSAGE, suggestion=REJECTED_KEY_SUGGESTION)
+    return NotAuthenticated(
+        REJECTED_KEY_MESSAGE, suggestion=REJECTED_KEY_SUGGESTION, rejected_key=True
+    )