From e86a8636592b57b1242c575bd7c45453dd1fb75f Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 11 Jun 2026 14:23:47 +0000 Subject: [PATCH 1/2] Fix code-review findings: JSON list output, agent/TTS robustness, auth-flow structure Correctness: - aai agent --list-voices and aai llm --list-models now route through run_command, so --json emits a machine-readable array instead of plain text. - File-driven agent runs fail loudly (APIError) when the session-ready wait times out, instead of silently consuming and dropping the clip's audio. - Malformed TTS Audio frames (missing or undecodable payload) map to a clear APIError; base64 is now strictly validated so junk bytes can't corrupt PCM. Structure: - New aai_cli/ws.py deduplicates the WebSocket auth-error classification and websockets logger silencing that agent/session.py and tts/session.py carried as identical copies. - NotAuthenticated gains a structured rejected_key flag; auto-login suppression keys off it instead of matching error message text. - login/logout opt out of auto-login at their run_command call sites (auto_login=False) instead of a hardcoded command-name set in context.py. - TTS PCM duration formula extracted into one helper. https://claude.ai/code/session_01RriYLpxVebHxots7yyPUo3 --- .importlinter | 1 + aai_cli/agent/session.py | 68 ++++++++--------------------- aai_cli/commands/agent.py | 10 ++++- aai_cli/commands/llm.py | 10 ++++- aai_cli/commands/login.py | 7 ++- aai_cli/context.py | 15 +++---- aai_cli/errors.py | 8 +++- aai_cli/tts/session.py | 61 +++++++++++--------------- aai_cli/ws.py | 57 ++++++++++++++++++++++++ tests/test_agent_command.py | 15 +++++++ tests/test_agent_session.py | 16 +++++++ tests/test_agent_session_run.py | 9 ++-- tests/test_context.py | 26 +---------- tests/test_errors.py | 3 ++ tests/test_llm_command.py | 14 ++++++ tests/test_login.py | 37 ++++++++++++++++ tests/test_tts_session.py | 24 ++++++++++ tests/test_ws.py | 77 +++++++++++++++++++++++++++++++++ 18 files changed, 327 insertions(+), 131 deletions(-) create mode 100644 aai_cli/ws.py create mode 100644 tests/test_ws.py diff --git a/.importlinter b/.importlinter index f333b96c..c8bdbae7 100644 --- a/.importlinter +++ b/.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 diff --git a/aai_cli/agent/session.py b/aai_cli/agent/session.py index b311aeaf..efecfbf4 100644 --- a/aai_cli/agent/session.py +++ b/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 @@ -181,8 +181,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 +201,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 +247,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 +290,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() diff --git a/aai_cli/commands/agent.py b/aai_cli/commands/agent.py index a7cd7d3e..013006c0 100644 --- a/aai_cli/commands/agent.py +++ b/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) diff --git a/aai_cli/commands/llm.py b/aai_cli/commands/llm.py index b971d660..d5b54a14 100644 --- a/aai_cli/commands/llm.py +++ b/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) diff --git a/aai_cli/commands/login.py b/aai_cli/commands/login.py index 15b4f986..66ce390e 100644 --- a/aai_cli/commands/login.py +++ b/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( diff --git a/aai_cli/context.py b/aai_cli/context.py index 3b411618..3bac1a04 100644 --- a/aai_cli/context.py +++ b/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,10 +137,7 @@ 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 @@ -148,8 +145,10 @@ def _should_auto_login(ctx: typer.Context, err: NotAuthenticated) -> bool: 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) diff --git a/aai_cli/errors.py b/aai_cli/errors.py index 36055fd8..9a251acc 100644 --- a/aai_cli/errors.py +++ b/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 + ) diff --git a/aai_cli/tts/session.py b/aai_cli/tts/session.py index 01d2767f..8a6e26ee 100644 --- a/aai_cli/tts/session.py +++ b/aai_cli/tts/session.py @@ -1,16 +1,17 @@ from __future__ import annotations import base64 +import binascii import contextlib import json -import logging from collections.abc import Callable from dataclasses import dataclass from typing import Protocol from urllib.parse import urlencode from aai_cli import environments -from aai_cli.errors import APIError, CLIError, auth_failure, is_auth_failure +from aai_cli import ws as wsutil +from aai_cli.errors import APIError, CLIError from aai_cli.tts import audio @@ -27,14 +28,6 @@ def close(self) -> None: ... # matches structurally; tests inject a fake with the same surface. _Connect = Callable[..., _WebSocket] -# A pre-upgrade HTTP 403 on the WebSocket handshake is NOT a rejected key (it also -# covers WAF/region/plan blocks) — mirrors how agent/stream classify handshakes. -_HTTP_FORBIDDEN = 403 - -# websockets' sync client logs teardown errors through these; silence them so an -# internal "stream ended" traceback never lands on stderr next to our clean error. -_WEBSOCKETS_LOGGERS = ("websockets", "websockets.client") - # The streaming-TTS server synthesizes at 24 kHz unless a sample_rate is requested, # and echoes the resolved value back in the Begin frame's configuration. Audio frames # carry only the PCM payload, so Begin is the single source of truth for the rate; this @@ -86,26 +79,24 @@ def ws_url(params: dict[str, str]) -> str: return f"{base}?{urlencode(params)}" if params else base -def _silence_websockets_logging() -> None: - """Keep websockets' internal teardown logging off the user's stderr.""" - for name in _WEBSOCKETS_LOGGERS: - logging.getLogger(name).setLevel(logging.CRITICAL) - +def _pcm_duration_seconds(pcm: bytes | bytearray, sample_rate: int) -> float: + """Seconds of audio in 16-bit mono PCM: two bytes per sample.""" + return len(pcm) / 2 / sample_rate -def _is_rejected_key(exc: Exception) -> bool: - """Auth-shaped failure (the key was rejected) vs a generic block. A plain HTTP - 403 on the upgrade is NOT a rejected key (WAF/region/plan also 403).""" - 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: - """A rejected key becomes auth_failure(); anything else becomes an APIError.""" - if _is_rejected_key(exc): - return auth_failure() - return APIError(f"{message}: {exc}") +def _decode_audio_frame(msg: dict[str, object]) -> bytes: + """The PCM payload of an Audio frame, with a malformed frame (missing or + undecodable ``audio``) mapped to an APIError that names the defect — not a + bare KeyError/binascii.Error that surfaces as a cryptic "TTS session failed".""" + encoded = msg.get("audio") + if not isinstance(encoded, str): + raise APIError("TTS service sent an Audio frame without an audio payload.") + try: + # validate=True: junk characters must fail loudly, not be silently dropped + # (the default discards them, corrupting the PCM byte stream). + return base64.b64decode(encoded, validate=True) + except (binascii.Error, ValueError) as exc: + raise APIError(f"TTS service sent an Audio frame that is not valid base64: {exc}") from exc def _default_connect( @@ -126,7 +117,7 @@ def _open_ws(connect: _Connect, api_key: str, url: str) -> _WebSocket: max_size=None, ) except Exception as exc: - raise _auth_or_api_error(exc, "Could not connect to the TTS service") from exc + raise wsutil.auth_or_api_error(exc, "Could not connect to the TTS service") from exc def _run_protocol( @@ -146,7 +137,7 @@ def _run_protocol( msg = json.loads(ws.recv()) mtype = msg.get("type") if mtype == "Audio": - pcm.extend(base64.b64decode(msg["audio"])) + pcm.extend(_decode_audio_frame(msg)) if msg.get("is_final"): break elif mtype == "Error": @@ -159,8 +150,7 @@ def _run_protocol( with contextlib.suppress(Exception): ws.send(json.dumps({"type": "Terminate"})) - duration = len(pcm) / 2 / sample_rate - return SpeakResult(bytes(pcm), sample_rate, duration) + return SpeakResult(bytes(pcm), sample_rate, _pcm_duration_seconds(pcm, sample_rate)) def synthesize( @@ -175,7 +165,7 @@ def synthesize( ``connect`` defaults to websockets' synchronous client; injectable for tests. Connect/session failures map to a clean CLIError (a rejected key -> exit 4). """ - _silence_websockets_logging() + wsutil.silence_websockets_logging() if connect is None: connect = _default_connect @@ -185,7 +175,7 @@ def synthesize( except (CLIError, KeyboardInterrupt, BrokenPipeError): raise # clean CLI errors, Ctrl-C, and a closed pipe are handled upstream except Exception as exc: - raise _auth_or_api_error(exc, "TTS session failed") from exc + raise wsutil.auth_or_api_error(exc, "TTS session failed") from exc finally: with contextlib.suppress(Exception): ws.close() @@ -215,5 +205,4 @@ def synthesize_dialogue( pcm.extend(audio.silence(result.sample_rate, _INTER_TURN_SILENCE_SECONDS)) pcm.extend(result.pcm) sample_rate_out = result.sample_rate - duration = len(pcm) / 2 / sample_rate_out - return SpeakResult(bytes(pcm), sample_rate_out, duration) + return SpeakResult(bytes(pcm), sample_rate_out, _pcm_duration_seconds(pcm, sample_rate_out)) diff --git a/aai_cli/ws.py b/aai_cli/ws.py new file mode 100644 index 00000000..8433dc82 --- /dev/null +++ b/aai_cli/ws.py @@ -0,0 +1,57 @@ +"""Shared plumbing for the CLI's WebSocket-backed sessions (voice agent, streaming TTS). + +Both sessions classify connect/session failures the same way and silence the same +library loggers; keeping that here means a change to either behavior lands in +`aai agent` and `aai speak` together instead of drifting apart. +""" + +from __future__ import annotations + +import logging + +from aai_cli.errors import APIError, CLIError, auth_failure, is_auth_failure + +# A pre-upgrade HTTP 403 on the WebSocket handshake is NOT a rejected key (it also +# covers WAF/region/plan blocks) — mirrors how `stream` classifies handshakes. +_HTTP_FORBIDDEN = 403 + +# The sync websockets client logs through these; both are silenced for a 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}") diff --git a/tests/test_agent_command.py b/tests/test_agent_command.py index d093fba0..b1828ba9 100644 --- a/tests/test_agent_command.py +++ b/tests/test_agent_command.py @@ -3,6 +3,7 @@ from typer.testing import CliRunner from aai_cli import config +from aai_cli.agent.voices import VOICES, format_voice_list from aai_cli.auth.flow import LoginResult from aai_cli.main import app @@ -36,9 +37,23 @@ def fake_run_session(*a, **k): result = runner.invoke(app, ["agent", "--list-voices"]) assert result.exit_code == 0 assert "ivy" in result.output + # Human mode prints the bare list, not a JSON array. + assert result.output.strip() == format_voice_list() assert called["ran"] is False +def test_list_voices_json_emits_machine_readable_array(monkeypatch): + monkeypatch.setattr( + "aai_cli.commands.agent.run_session", + lambda *a, **k: (_ for _ in ()).throw(AssertionError("must not connect")), + ) + result = runner.invoke(app, ["agent", "--list-voices", "--json"]) + assert result.exit_code == 0 + voices = json.loads(result.output) + assert voices == VOICES # the whole list, as a machine-readable array + assert "ivy" in voices + + def test_agent_unauthenticated_runs_login(monkeypatch): monkeypatch.setattr("aai_cli.context._interactive_session", lambda: True) monkeypatch.setattr("aai_cli.context.run_login_flow", _login_result) diff --git a/tests/test_agent_session.py b/tests/test_agent_session.py index ec51844e..fe232abc 100644 --- a/tests/test_agent_session.py +++ b/tests/test_agent_session.py @@ -285,3 +285,19 @@ def wait(self, timeout=None): s.ready = True _send_audio_loop(_RecordingWS(), s, [b"\x01\x02"]) assert seen["timeout"] == 10 + + +def test_send_audio_loop_fails_loudly_when_ready_times_out(): + # A server that never sends session.ready must fail the run, not silently + # consume (and drop) the finite file source frame by frame. + class _NeverReadyEvent: + def wait(self, timeout=None): + return False + + s = _session(exit_after_reply=True, ready_event=_NeverReadyEvent()) + s.ready = True # even an open gate must not be reached past a timed-out wait + ws = _RecordingWS() + with pytest.raises(APIError) as exc: + _send_audio_loop(ws, s, [b"\x01\x02"]) + assert "ready" in exc.value.message + assert ws.sent == [] # no audio was consumed or sent diff --git a/tests/test_agent_session_run.py b/tests/test_agent_session_run.py index 7254223f..7118a533 100644 --- a/tests/test_agent_session_run.py +++ b/tests/test_agent_session_run.py @@ -9,8 +9,9 @@ import pytest -from aai_cli.agent.session import _WEBSOCKETS_LOGGERS, AgentRunConfig, run_session +from aai_cli.agent.session import AgentRunConfig, run_session from aai_cli.errors import APIError, CLIError, NotAuthenticated +from aai_cli.ws import WEBSOCKETS_LOGGERS class FakeRenderer: @@ -281,7 +282,7 @@ def close(self): def test_run_session_silences_websockets_loggers(): # websockets' sync reader thread logs teardown errors (EOFError tracebacks) via # its own loggers; run_session must mute them so they never hit the user's stderr. - loggers = [logging.getLogger(name) for name in _WEBSOCKETS_LOGGERS] + loggers = [logging.getLogger(name) for name in WEBSOCKETS_LOGGERS] previous = [lg.level for lg in loggers] try: for lg in loggers: @@ -298,8 +299,8 @@ def test_run_session_silences_websockets_loggers(): def test_websockets_logger_names_cover_the_sync_client(): # The sync client logs through "websockets.client"; pin that the silenced set # covers it (and the parent, for any future child loggers). - assert "websockets.client" in _WEBSOCKETS_LOGGERS - assert "websockets" in _WEBSOCKETS_LOGGERS + assert "websockets.client" in WEBSOCKETS_LOGGERS + assert "websockets" in WEBSOCKETS_LOGGERS def test_run_session_non_auth_failure_stays_api_error(): diff --git a/tests/test_context.py b/tests/test_context.py index d70bc72d..39a19a8b 100644 --- a/tests/test_context.py +++ b/tests/test_context.py @@ -227,30 +227,6 @@ def body(state, json_mode): assert result.exit_code == 4 -def test_run_command_never_auto_logs_in_login_command(monkeypatch): - _force_interactive(monkeypatch) - monkeypatch.setattr( - "aai_cli.context.run_login_flow", - lambda: (_ for _ in ()).throw(AssertionError("login command must not auto-login")), - ) - - app = typer.Typer() - - @app.callback() - def cb(ctx: typer.Context): - ctx.obj = AppState() - - @app.command(name="login") - def login_cmd(ctx: typer.Context): - def body(state, json_mode): - raise NotAuthenticated() - - run_command(ctx, body) - - result = runner.invoke(app, ["login"]) - assert result.exit_code == 4 - - def test_run_command_runs_body_on_success(): seen = {} @@ -372,7 +348,7 @@ def fake_login(): monkeypatch.setattr("aai_cli.context.run_login_flow", fake_login) def body(state, json_mode): - raise NotAuthenticated() # message != REJECTED_KEY_MESSAGE + raise NotAuthenticated() # rejected_key is False: not a key rejection result = runner.invoke(_make_app(body), ["go"]) assert ran["login"] == 1 # auto-login was attempted despite the env key diff --git a/tests/test_errors.py b/tests/test_errors.py index 4d2da228..29a69b53 100644 --- a/tests/test_errors.py +++ b/tests/test_errors.py @@ -6,6 +6,7 @@ def test_not_authenticated_defaults(): assert err.exit_code == 4 assert err.error_type == "not_authenticated" assert err.message == "You're not signed in." + assert err.rejected_key is False # "no credential" is not a key rejection assert err.suggestion == ( "Run 'aai onboard' for guided setup, 'aai login' if you have an account, " "or set ASSEMBLYAI_API_KEY." @@ -58,6 +59,8 @@ def test_auth_failure_splits_message_and_suggestion(): assert err.message == "Your API key was rejected." assert "aai login" in (err.suggestion or "") assert "ASSEMBLYAI_API_KEY" in (err.suggestion or "") + # The structured marker auto-login keys off (instead of matching message text). + assert err.rejected_key is True def test_is_auth_failure_matches_credential_signals(): diff --git a/tests/test_llm_command.py b/tests/test_llm_command.py index 6cde0a3b..ed0a121a 100644 --- a/tests/test_llm_command.py +++ b/tests/test_llm_command.py @@ -5,6 +5,7 @@ from aai_cli import config from aai_cli.auth.flow import LoginResult +from aai_cli.llm import KNOWN_MODELS from aai_cli.main import app runner = CliRunner() @@ -44,9 +45,22 @@ def test_llm_list_models_exits_without_network(monkeypatch): result = runner.invoke(app, ["llm", "--list-models"]) assert result.exit_code == 0 assert "claude-sonnet-4-6" in result.output + # Human mode prints the bare newline-joined list, not a JSON array. + assert result.output.strip() == "\n".join(KNOWN_MODELS) assert called["ran"] is False +def test_llm_list_models_json_emits_machine_readable_array(monkeypatch): + monkeypatch.setattr( + "aai_cli.commands.llm.gateway.complete", + lambda *a, **k: (_ for _ in ()).throw(AssertionError("must not call the gateway")), + ) + result = runner.invoke(app, ["llm", "--list-models", "--json"]) + assert result.exit_code == 0 + models = json.loads(result.output) + assert models == list(KNOWN_MODELS) # the whole list, as a machine-readable array + + def test_llm_sends_prompt_and_prints_output(monkeypatch): _auth() seen = {} diff --git a/tests/test_login.py b/tests/test_login.py index 33e9af57..9bf9813d 100644 --- a/tests/test_login.py +++ b/tests/test_login.py @@ -384,3 +384,40 @@ def test_whoami_honors_env_api_key(monkeypatch, mocker): data = json.loads(result.output) assert data["reachable"] is True assert "sk_env_1234567890" not in result.output # still masked + + +def test_login_failure_never_auto_logs_in_again(monkeypatch): + # The login command owns sign-in (auto_login=False): a NotAuthenticated from its + # own browser flow must surface as exit 4, not trigger run_command's auto-login + # retry — even in an interactive session. + from aai_cli.errors import NotAuthenticated + + calls = {"n": 0} + + def timed_out(): + calls["n"] += 1 + raise NotAuthenticated("Login timed out waiting for the browser.") + + monkeypatch.setattr("aai_cli.context._interactive_session", lambda: True) + monkeypatch.setattr("aai_cli.context.run_login_flow", timed_out) + result = runner.invoke(app, ["login"]) + assert result.exit_code == 4 + assert calls["n"] == 1 # the command's own attempt only; no auto-login retry + + +def test_logout_never_auto_logs_in(monkeypatch): + # Signing out must never start a sign-in flow (auto_login=False), even if the + # body surfaces a NotAuthenticated and the session is interactive. + from aai_cli.errors import NotAuthenticated + + monkeypatch.setattr("aai_cli.context._interactive_session", lambda: True) + monkeypatch.setattr( + "aai_cli.context.run_login_flow", + lambda: (_ for _ in ()).throw(AssertionError("logout must never start a login")), + ) + monkeypatch.setattr( + "aai_cli.commands.login.config.clear_api_key", + lambda _profile: (_ for _ in ()).throw(NotAuthenticated()), + ) + result = runner.invoke(app, ["logout"]) + assert result.exit_code == 4 diff --git a/tests/test_tts_session.py b/tests/test_tts_session.py index d086a938..5d0d5c14 100644 --- a/tests/test_tts_session.py +++ b/tests/test_tts_session.py @@ -167,6 +167,30 @@ def test_synthesize_raises_on_missing_begin(): session.synthesize("k", session.SpeakConfig(text="hi"), connect=lambda *a, **k: ws) +def test_synthesize_maps_audio_frame_without_payload_to_api_error(): + # A malformed Audio frame (no "audio" key) must surface as a clean APIError + # naming the defect, not a bare KeyError read as "TTS session failed: 'audio'". + ws = FakeWS( + [ + json.dumps({"type": "Begin", "id": "s", "expires_at": 1}), + json.dumps({"type": "Audio", "is_final": True}), + ] + ) + with pytest.raises(APIError, match="without an audio payload"): + session.synthesize("k", session.SpeakConfig(text="hi"), connect=lambda *a, **k: ws) + + +def test_synthesize_maps_undecodable_audio_frame_to_api_error(): + ws = FakeWS( + [ + json.dumps({"type": "Begin", "id": "s", "expires_at": 1}), + json.dumps({"type": "Audio", "audio": "!!!", "is_final": True}), + ] + ) + with pytest.raises(APIError, match="not valid base64"): + session.synthesize("k", session.SpeakConfig(text="hi"), connect=lambda *a, **k: ws) + + def test_synthesize_maps_error_frame_to_api_error(): ws = FakeWS( [ diff --git a/tests/test_ws.py b/tests/test_ws.py new file mode 100644 index 00000000..3c923cb3 --- /dev/null +++ b/tests/test_ws.py @@ -0,0 +1,77 @@ +"""Direct tests for the shared WebSocket-session helpers (aai_cli/ws.py). + +The agent and TTS sessions both route through these; session-level behavior is +covered in test_agent_session_run.py / test_tts_session.py. +""" + +from __future__ import annotations + +import logging +import types + +from aai_cli.errors import APIError, NotAuthenticated +from aai_cli.ws import ( + WEBSOCKETS_LOGGERS, + auth_or_api_error, + is_rejected_key, + silence_websockets_logging, +) + + +class _HandshakeRejected(Exception): + """Mimics websockets' InvalidStatus: a structured HTTP status on ``.response``.""" + + def __init__(self, status: int) -> None: + super().__init__(f"server rejected WebSocket connection: HTTP {status}") + self.response = types.SimpleNamespace(status_code=status) + + +def test_is_rejected_key_false_for_handshake_403(): + # 403 also covers WAF/region/plan blocks, so it must NOT read as a rejected key. + assert is_rejected_key(_HandshakeRejected(403)) is False + + +def test_is_rejected_key_true_for_handshake_401(): + assert is_rejected_key(_HandshakeRejected(401)) is True + + +def test_is_rejected_key_true_for_auth_worded_message(): + assert is_rejected_key(RuntimeError("Unauthorized")) is True + + +def test_is_rejected_key_false_for_generic_failure(): + assert is_rejected_key(RuntimeError("network unreachable")) is False + + +def test_auth_or_api_error_maps_rejected_key_to_not_authenticated(): + err = auth_or_api_error(RuntimeError("Unauthorized"), "Could not connect") + assert isinstance(err, NotAuthenticated) + assert err.exit_code == 4 + + +def test_auth_or_api_error_wraps_other_failures_with_context(): + err = auth_or_api_error(RuntimeError("network unreachable"), "Could not connect") + assert isinstance(err, APIError) + assert err.message == "Could not connect: network unreachable" + + +def test_silence_websockets_logging_raises_both_loggers_to_critical(): + loggers = [logging.getLogger(name) for name in WEBSOCKETS_LOGGERS] + previous = [lg.level for lg in loggers] + try: + for lg in loggers: + lg.setLevel(logging.NOTSET) + silence_websockets_logging() + for lg in loggers: + assert lg.level == logging.CRITICAL + assert not lg.isEnabledFor(logging.ERROR) + finally: + for lg, level in zip(loggers, previous, strict=True): + lg.setLevel(level) + + +def test_websockets_logger_names_cover_the_sync_client(): + # The sync client logs through "websockets.client"; the parent covers any + # future child loggers. + assert "websockets.client" in WEBSOCKETS_LOGGERS + assert "websockets" in WEBSOCKETS_LOGGERS From 213a8d71281e15d9abd269a06cada24bae91945d Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 11 Jun 2026 14:29:42 +0000 Subject: [PATCH 2/2] Mark voice-agent session auth rejection as rejected_key MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Security review follow-up: a session.error with UNAUTHORIZED/FORBIDDEN means a credential was presented and refused, so it must carry rejected_key=True — otherwise an interactive `aai agent` run with a bad ASSEMBLYAI_API_KEY would trigger a browser auto-login that cannot fix the failure (the env key still wins on rerun) and clobbers the profile's stored key. https://claude.ai/code/session_01RriYLpxVebHxots7yyPUo3 --- aai_cli/agent/session.py | 3 +++ tests/test_agent_session.py | 2 ++ 2 files changed, 5 insertions(+) diff --git a/aai_cli/agent/session.py b/aai_cli/agent/session.py index efecfbf4..b6d2f135 100644 --- a/aai_cli/agent/session.py +++ b/aai_cli/agent/session.py @@ -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}") diff --git a/tests/test_agent_session.py b/tests/test_agent_session.py index fe232abc..d8a797b8 100644 --- a/tests/test_agent_session.py +++ b/tests/test_agent_session.py @@ -130,6 +130,8 @@ def test_unauthorized_error_raises_not_authenticated_exit_4(): assert excinfo.value.exit_code == 4 assert "bad key" in str(excinfo.value) # the server message wins over code/fallback assert excinfo.value.suggestion is not None and "aai login" in excinfo.value.suggestion + # A presented-and-refused key: auto-login must not retry over a bad env key. + assert excinfo.value.rejected_key is True def test_other_session_error_raises_api_error():