From bf15b20b85b078223204a4c437b0284332f5a840 Mon Sep 17 00:00:00 2001 From: ottomansky Date: Tue, 26 May 2026 13:35:08 +0200 Subject: [PATCH 1/3] feat: add `kai verify` diagnostic and clean error rendering in JS data app MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When the Keboola platform UI works but a Kai-using app errors out, there was no way to tell why without reading source. `kai verify` now reports token identity (which project/owner the token resolves to), the auto-discovered kai-assistant URL, ping/info reachability, and current monthly message usage via the existing `GET /api/usage` endpoint. On a `429 rate_limit:chat` it prints `code: message` cleanly instead of dumping raw JSON. The JS data app example previously surfaced upstream errors as a raw JSON blob (`Kai API error: 429 — {"code":"rate_limit:chat","message":"..."}`). The proxy now parses the upstream body and returns a structured `{error: {status, code, message}}`, and the frontend renders the message text rather than the JSON envelope. Adds `.env.example` and a troubleshooting README so the `STORAGE_API_TOKEN` (not `KAI_TOKEN`) env var name is unambiguous. Bumps the package to 0.13.0. Co-Authored-By: Claude Opus 4.7 (1M context) --- README.md | 8 ++ examples/js-dataapp/.env.example | 14 +++ examples/js-dataapp/README.md | 51 ++++++++ examples/js-dataapp/public/app.js | 17 ++- examples/js-dataapp/server.js | 17 ++- pyproject.toml | 2 +- src/kai_client/__init__.py | 4 +- src/kai_client/cli.py | 197 ++++++++++++++++++++++++++++-- tests/test_cli_verify.py | 163 ++++++++++++++++++++++++ uv.lock | 2 +- 10 files changed, 460 insertions(+), 15 deletions(-) create mode 100644 examples/js-dataapp/README.md create mode 100644 tests/test_cli_verify.py diff --git a/README.md b/README.md index 4edddd2..9503e8d 100644 --- a/README.md +++ b/README.md @@ -95,9 +95,17 @@ STORAGE_API_URL=https://connection.keboola.com ```bash kai ping # Check if the server is alive kai info # Show server version, uptime, connected MCP servers +kai verify # Diagnose token, service URL, and monthly message quota kai --version # Show CLI version ``` +`kai verify` is the right starting point when chats work in the Keboola +platform UI but your app errors out. It reports which project/user the token +resolves to, the discovered kai-assistant URL, and your current monthly +message usage (e.g. `19/150 messages used, 131 left`). On a 429 +`rate_limit:chat` it prints the code and message cleanly instead of dumping +raw JSON. Pass `--json-output` for a machine-readable report. + #### Chat ```bash diff --git a/examples/js-dataapp/.env.example b/examples/js-dataapp/.env.example index 21e76fc..a83da61 100644 --- a/examples/js-dataapp/.env.example +++ b/examples/js-dataapp/.env.example @@ -1,2 +1,16 @@ +# Copy this file to .env.local and fill in real values. +# Use the env var names exactly as written — `KAI_TOKEN` is NOT recognised. + +# Keboola Storage API token (Master or scoped) — same value the platform UI uses STORAGE_API_TOKEN=your-keboola-token + +# Keboola Storage API URL for your stack +# Examples: +# https://connection.keboola.com (AWS US) +# https://connection.eu-central-1.keboola.com (AWS EU) +# https://connection.us-east4.gcp.keboola.com (GCP US) +# https://connection.europe-west3.gcp.keboola.com (GCP EU) STORAGE_API_URL=https://connection.keboola.com + +# Optional: dev server port (default 3000) +# PORT=3000 diff --git a/examples/js-dataapp/README.md b/examples/js-dataapp/README.md new file mode 100644 index 0000000..85591a5 --- /dev/null +++ b/examples/js-dataapp/README.md @@ -0,0 +1,51 @@ +# Kai Data App — JavaScript example + +Minimal Express + vanilla-JS frontend that talks to the Keboola AI Assistant. + +## Run locally + +```bash +cp .env.example .env.local # fill in STORAGE_API_TOKEN and STORAGE_API_URL +npm install +node server.js # serves http://localhost:3000 +``` + +## Env vars + +The required names are exact: + +- `STORAGE_API_TOKEN` — Keboola Storage API token (the platform UI uses the same) +- `STORAGE_API_URL` — Storage API URL for your stack (see `.env.example`) + +`KAI_TOKEN` and similar names are **not** recognised. The legacy `KBC_TOKEN` / +`KBC_URL` names are honoured as fallbacks only. + +## Troubleshooting + +### "Works in the Keboola platform UI but my app errors out" + +Run `kai verify` from the same shell where you set the env vars: + +```bash +pip install kai-client # or: uv tool install kai-client +kai verify +``` + +It reports which project your token resolves to, whether the kai-assistant +service is reachable, and your current monthly message usage +(e.g. `19/150 messages used, 131 left`). + +### Common errors + +- **`429 rate_limit:chat — You have exceeded your maximum number of messages + for this month.`** Your project hit its monthly chat-message limit. Note + that the platform UI's "X / 150" counter is **per-token, server-side** — + if the UI says you have messages left but the API says you don't, your + app token and UI session may be hitting different counters. Contact + Keboola support or wait for the reset date shown by `kai verify`. + +- **`401 storage.tokenInvalid`** The token doesn't exist or has been + expired/revoked. Generate a fresh one. + +- **`kai-assistant service not found`** Your stack doesn't have the AI + Assistant feature enabled, or `STORAGE_API_URL` points at the wrong stack. diff --git a/examples/js-dataapp/public/app.js b/examples/js-dataapp/public/app.js index 83a2c28..0bd550f 100644 --- a/examples/js-dataapp/public/app.js +++ b/examples/js-dataapp/public/app.js @@ -177,8 +177,21 @@ async function readSSEStream(url, fetchOptions, onEvent) { const res = await fetch(url, fetchOptions); if (!res.ok) { - const err = await res.text(); - addError(`Error: ${err}`); + // server.js returns {error: {status, code?, message}} for structured upstream + // errors. Fall back to raw text if the body isn't JSON. + let body; + try { + body = await res.json(); + } catch { + body = null; + } + const e = body && body.error; + if (e && typeof e === "object") { + const prefix = e.code ? `${e.status || res.status} ${e.code}` : `${e.status || res.status}`; + addError(`${prefix} — ${e.message || "Request failed"}`); + } else { + addError(`Error: ${typeof e === "string" ? e : res.statusText}`); + } return null; } diff --git a/examples/js-dataapp/server.js b/examples/js-dataapp/server.js index 9a1d863..27af550 100644 --- a/examples/js-dataapp/server.js +++ b/examples/js-dataapp/server.js @@ -70,7 +70,22 @@ async function proxySSE(payload, res) { if (!upstream.ok) { const text = await upstream.text(); - return res.status(upstream.status).json({ error: text }); + // Try to parse the upstream body as a structured KAI error + // ({code, message}) so the frontend can render a clean message + // instead of the raw JSON blob. Common shapes hit here: + // 429 → {"code":"rate_limit:chat","message":"You have exceeded..."} + // 401 → {"code":"unauthorized:chat","message":"..."} + let parsed; + try { + parsed = JSON.parse(text); + } catch { + parsed = null; + } + const error = + parsed && typeof parsed === "object" && (parsed.code || parsed.message) + ? { status: upstream.status, code: parsed.code, message: parsed.message } + : { status: upstream.status, message: text || upstream.statusText }; + return res.status(upstream.status).json({ error }); } res.setHeader("Content-Type", "text/event-stream"); diff --git a/pyproject.toml b/pyproject.toml index ef4c015..f0a5739 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "kai-client" -version = "0.12.0" +version = "0.13.0" description = "Python client library for the Keboola AI Assistant Backend API" readme = "README.md" license = { text = "MIT" } diff --git a/src/kai_client/__init__.py b/src/kai_client/__init__.py index fb83959..e725e3a 100644 --- a/src/kai_client/__init__.py +++ b/src/kai_client/__init__.py @@ -84,7 +84,7 @@ VoteType, ) -__version__ = "0.12.0" +__version__ = "0.13.0" __all__ = [ # Main client @@ -154,5 +154,3 @@ # Version "__version__", ] - - diff --git a/src/kai_client/cli.py b/src/kai_client/cli.py index ff6ed48..f947a8d 100644 --- a/src/kai_client/cli.py +++ b/src/kai_client/cli.py @@ -5,9 +5,10 @@ import os import sys from pathlib import Path -from typing import Optional +from typing import Any, Optional import click +import httpx from dotenv import load_dotenv # Load .env.local file if it exists (before any commands run) @@ -16,6 +17,7 @@ load_dotenv(_env_local) from kai_client import KaiClient, __version__ # noqa: E402 +from kai_client.exceptions import KaiError # noqa: E402 from kai_client.models import ToolApprovalRequestEvent # noqa: E402 from kai_client.types import VoteType # noqa: E402 @@ -65,6 +67,9 @@ def main(ctx, token: Optional[str], url: Optional[str], base_url: Optional[str]) # Check server health kai ping + # Diagnose token / service / quota when things misbehave + kai verify + # Start an interactive chat kai chat @@ -144,6 +149,187 @@ async def _info(): run_async(_info()) +@main.command() +@click.option("--json-output", is_flag=True, help="Output as JSON") +@click.pass_context +def verify(ctx, json_output: bool): + """ + Diagnose Kai setup: token identity, service reachability, message quota. + + Use this when "it works in the Keboola platform UI but my app errors out". + Verify confirms which project/user the token resolves to, that the + kai-assistant service is reachable, and your current monthly message usage. + + Examples: + + kai verify + kai verify --json-output + """ + + async def _verify(): + token = ctx.obj.get("token") or get_env_or_error("STORAGE_API_TOKEN") + url = ctx.obj.get("url") or get_env_or_error("STORAGE_API_URL") + base_url = ctx.obj.get("base_url") + + report: dict[str, Any] = {"ok": True, "checks": {}} + + def _print_check(name: str, ok: bool, summary: str) -> None: + if json_output: + return + mark = click.style("✓", fg="green") if ok else click.style("✗", fg="red") + click.echo(f"{mark} {name}: {summary}") + + def _record(name: str, ok: bool, summary: str, **extra: Any) -> None: + report["checks"][name] = {"ok": ok, "summary": summary, **extra} + if not ok: + report["ok"] = False + _print_check(name, ok, summary) + + if not json_output: + click.echo(f"Storage API URL: {url}") + + # 1. Token identity (Keboola Storage API) + token_data: Optional[dict[str, Any]] = None + async with httpx.AsyncClient() as http_client: + try: + resp = await http_client.get( + f"{url.rstrip('/')}/v2/storage/tokens/verify", + headers={"x-storageapi-token": token}, + timeout=30.0, + ) + if resp.status_code >= 400: + body: dict[str, Any] = {} + if "application/json" in resp.headers.get("content-type", ""): + parsed = resp.json() + if isinstance(parsed, dict): + body = parsed + if not body: + body = {"error": resp.text} + code = body.get("code") or f"http:{resp.status_code}" + msg = body.get("error") or body.get("message") or resp.reason_phrase + _record("token", False, f"HTTP {resp.status_code} {code}: {msg}") + else: + parsed = resp.json() + assert isinstance(parsed, dict), ( + "Storage API token verify must return JSON object" + ) + token_data = parsed + owner = token_data.get("owner") or {} + project_id = owner.get("id") + project_name = owner.get("name", "?") + description = token_data.get("description", "?") + is_master = token_data.get("isMasterToken", False) + token_type = "master" if is_master else "scoped" + summary = ( + f"project {project_id} ({project_name}), " + f"token '{description}' [{token_type}]" + ) + _record( + "token", + True, + summary, + project_id=project_id, + project_name=project_name, + token_description=description, + is_master_token=is_master, + ) + except httpx.RequestError as e: + _record("token", False, f"connection error: {e}") + + # If the token didn't verify, the remaining checks will all fail the same way. + # Bail early with a non-zero exit so this command stays useful in scripts. + if not token_data: + if json_output: + click.echo(json.dumps(report, indent=2, default=str)) + sys.exit(1) + + # 2. Build a KaiClient: respect --base-url for local dev; otherwise auto-discover. + client: Optional[KaiClient] = None + try: + if base_url: + client = KaiClient(storage_api_token=token, storage_api_url=url, base_url=base_url) + _record("service-discovery", True, f"using --base-url {base_url}") + else: + client = await KaiClient.from_storage_api( + storage_api_token=token, storage_api_url=url + ) + _record("service-discovery", True, f"kai-assistant at {client.base_url}") + except KaiError as e: + _record( + "service-discovery", + False, + _format_kai_error(e), + code=e.code, + message=e.message, + ) + + if client is None: + if json_output: + click.echo(json.dumps(report, indent=2, default=str)) + sys.exit(1) + + async with client: + # 3. Reachability (no auth needed) + try: + ping_resp = await client.ping() + _record("ping", True, f"server alive at {ping_resp.timestamp.isoformat()}") + except KaiError as e: + _record("ping", False, _format_kai_error(e), code=e.code, message=e.message) + + try: + info_resp = await client.info() + _record( + "info", + True, + f"{info_resp.app_name} v{info_resp.app_version} " + f"(server v{info_resp.server_version}, uptime {info_resp.uptime:.0f}s)", + server_version=info_resp.server_version, + ) + except KaiError as e: + _record("info", False, _format_kai_error(e), code=e.code, message=e.message) + + # 4. Authenticated probe + quota — the whole point of this command. + try: + usage = await client.get_usage() + remaining = usage.messages_limit - usage.messages_used + summary = ( + f"{usage.messages_used}/{usage.messages_limit} messages used " + f"(resets {usage.reset_date.date().isoformat()}, {remaining} left)" + ) + _record( + "usage", + True, + summary, + messages_used=usage.messages_used, + messages_limit=usage.messages_limit, + reset_date=usage.reset_date.isoformat(), + ) + except KaiError as e: + # Specifically surface 429 rate_limit:chat cleanly — that's the symptom + # this whole command exists to diagnose. + _record("usage", False, _format_kai_error(e), code=e.code, message=e.message) + + if json_output: + click.echo(json.dumps(report, indent=2, default=str)) + elif report["ok"]: + click.echo() + click.echo(click.style("All checks passed.", fg="green")) + else: + click.echo() + click.echo(click.style("Some checks failed — see above.", fg="red"), err=True) + + sys.exit(0 if report["ok"] else 1) + + run_async(_verify()) + + +def _format_kai_error(e: KaiError) -> str: + """Render a KaiError as `code: message`, falling back to message-only.""" + if e.code and e.message: + return f"{e.code}: {e.message}" + return e.message or str(e) + + @main.command() @click.option("-m", "--message", help="Send a single message instead of interactive mode") @click.option( @@ -200,9 +386,7 @@ async def _chat(): if message: # Single message mode - await send_and_display( - client, chat_id, message, auto_approve, json_output - ) + await send_and_display(client, chat_id, message, auto_approve, json_output) else: # Interactive mode if not json_output: @@ -312,9 +496,8 @@ async def send_and_display( # and cleared when that tool completes) if pending_approval: # Get the tool name from tracking dict or pending_approval - approved_tool_name = ( - pending_approval.tool_name - or current_tool_name.get(pending_approval.tool_call_id or "", "unknown") + approved_tool_name = pending_approval.tool_name or current_tool_name.get( + pending_approval.tool_call_id or "", "unknown" ) # Determine which approval flow to use diff --git a/tests/test_cli_verify.py b/tests/test_cli_verify.py new file mode 100644 index 0000000..3468806 --- /dev/null +++ b/tests/test_cli_verify.py @@ -0,0 +1,163 @@ +"""Tests for the `kai verify` CLI command.""" + +import json + +import pytest +from click.testing import CliRunner + +from kai_client.cli import main + +STORAGE_URL = "https://connection.test.keboola.com" +KAI_URL = "https://kai.test.keboola.com" + + +@pytest.fixture +def runner(): + return CliRunner() + + +@pytest.fixture +def mock_env(monkeypatch): + monkeypatch.setenv("STORAGE_API_TOKEN", "test-token") + monkeypatch.setenv("STORAGE_API_URL", STORAGE_URL) + + +def _mock_token_verify_ok(httpx_mock, *, is_master=True): + httpx_mock.add_response( + url=f"{STORAGE_URL}/v2/storage/tokens/verify", + method="GET", + json={ + "id": "12345", + "description": "max@keboola.com", + "isMasterToken": is_master, + "owner": {"id": 2738, "name": "Test Project"}, + }, + ) + + +def _mock_discovery_ok(httpx_mock): + httpx_mock.add_response( + url=f"{STORAGE_URL}/v2/storage", + method="GET", + json={"services": [{"id": "kai-assistant", "url": KAI_URL}]}, + ) + + +def _mock_ping_info_ok(httpx_mock): + httpx_mock.add_response( + url=f"{KAI_URL}/ping", + method="GET", + json={"timestamp": "2026-05-26T10:00:00Z"}, + ) + httpx_mock.add_response( + url=f"{KAI_URL}/api", + method="GET", + json={ + "timestamp": "2026-05-26T10:00:00Z", + "appName": "kai-assistant", + "appVersion": "1.2.3", + "serverVersion": "1.2.3", + "uptime": 3600.0, + "connectedMcp": [], + }, + ) + + +class TestVerifyCommand: + def test_full_success_renders_project_and_quota(self, runner, mock_env, httpx_mock): + _mock_token_verify_ok(httpx_mock, is_master=True) + _mock_discovery_ok(httpx_mock) + _mock_ping_info_ok(httpx_mock) + httpx_mock.add_response( + url=f"{KAI_URL}/api/usage", + method="GET", + json={ + "messagesUsed": 19, + "messagesLimit": 150, + "resetDate": "2026-06-01T00:00:00Z", + }, + ) + + result = runner.invoke(main, ["verify"]) + + assert result.exit_code == 0, result.output + assert "project 2738" in result.output + assert "Test Project" in result.output + assert "[master]" in result.output + assert "kai-assistant at https://kai.test.keboola.com" in result.output + assert "19/150 messages used" in result.output + assert "131 left" in result.output + assert "All checks passed." in result.output + + def test_rate_limit_429_renders_code_and_message_cleanly(self, runner, mock_env, httpx_mock): + """The whole reason this command exists — surface 429 rate_limit:chat clearly.""" + _mock_token_verify_ok(httpx_mock) + _mock_discovery_ok(httpx_mock) + _mock_ping_info_ok(httpx_mock) + httpx_mock.add_response( + url=f"{KAI_URL}/api/usage", + method="GET", + status_code=429, + json={ + "code": "rate_limit:chat", + "message": ( + "You have exceeded your maximum number of messages for " + "this month. Please contact support to raise your limit " + "or try again next month." + ), + }, + ) + + result = runner.invoke(main, ["verify"]) + + assert result.exit_code == 1 + assert "rate_limit:chat" in result.output + assert "exceeded your maximum number of messages" in result.output + # No raw JSON blob in the rendered output. + assert '{"code"' not in result.output + + def test_invalid_token_exits_nonzero(self, runner, mock_env, httpx_mock): + httpx_mock.add_response( + url=f"{STORAGE_URL}/v2/storage/tokens/verify", + method="GET", + status_code=401, + json={"error": "Invalid access token", "code": "storage.tokenInvalid"}, + ) + + result = runner.invoke(main, ["verify"]) + + assert result.exit_code == 1 + assert "storage.tokenInvalid" in result.output + assert "Invalid access token" in result.output + + def test_missing_env_var_exits_with_clear_message(self, runner, monkeypatch): + monkeypatch.delenv("STORAGE_API_TOKEN", raising=False) + monkeypatch.delenv("STORAGE_API_URL", raising=False) + + result = runner.invoke(main, ["verify"]) + + assert result.exit_code == 1 + assert "STORAGE_API_TOKEN" in result.output + + def test_json_output_shape(self, runner, mock_env, httpx_mock): + _mock_token_verify_ok(httpx_mock) + _mock_discovery_ok(httpx_mock) + _mock_ping_info_ok(httpx_mock) + httpx_mock.add_response( + url=f"{KAI_URL}/api/usage", + method="GET", + json={ + "messagesUsed": 19, + "messagesLimit": 150, + "resetDate": "2026-06-01T00:00:00Z", + }, + ) + + result = runner.invoke(main, ["verify", "--json-output"]) + + assert result.exit_code == 0, result.output + data = json.loads(result.output) + assert data["ok"] is True + assert data["checks"]["token"]["project_id"] == 2738 + assert data["checks"]["usage"]["messages_used"] == 19 + assert data["checks"]["usage"]["messages_limit"] == 150 diff --git a/uv.lock b/uv.lock index 2da5f44..b419288 100644 --- a/uv.lock +++ b/uv.lock @@ -409,7 +409,7 @@ wheels = [ [[package]] name = "kai-client" -version = "0.11.0" +version = "0.13.0" source = { editable = "." } dependencies = [ { name = "click" }, From 86ca38efa210242d455bb1a25802cc81c7e7f491 Mon Sep 17 00:00:00 2001 From: ottomansky Date: Tue, 26 May 2026 14:13:03 +0200 Subject: [PATCH 2/3] refactor: address review feedback on kai verify MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Extract `_verify()` into focused module-level helpers (`_check_token`, `_check_service_discovery`, `_check_reachability`, `_check_usage`, `_emit_verify_footer`) plus a `_VerifyRecorder` class. The orchestrator is now ~20 lines instead of 125. - Replace `assert isinstance(parsed, dict)` with a proper runtime check that records the failure and exits non-zero (asserts are stripped under python -O). - Rephrase the misleading comment on the usage error handler — the handler catches any `KaiError`, not specifically 429. - JS app: fall back to `e.code` before "Request failed" when an upstream body has a code but no message. - Revert two unrelated `ruff format` hunks in `send_and_display` that slipped into the previous commit. Test coverage extended: - scoped (non-master) token rendering - Storage API connection error (`httpx.RequestError` branch) - `--base-url` skips discovery (talks directly to the local URL) - discovery failure when `kai-assistant` is absent from the services list Full suite: 257 passing (253 + 4 new). Live `kai verify` against europe-west3 still reports `2/150 messages used (148 left)` identically. Co-Authored-By: Claude Opus 4.7 (1M context) --- examples/js-dataapp/public/app.js | 5 +- src/kai_client/cli.py | 316 ++++++++++++++++-------------- tests/test_cli_verify.py | 90 +++++++++ 3 files changed, 265 insertions(+), 146 deletions(-) diff --git a/examples/js-dataapp/public/app.js b/examples/js-dataapp/public/app.js index 0bd550f..a9b73cd 100644 --- a/examples/js-dataapp/public/app.js +++ b/examples/js-dataapp/public/app.js @@ -188,7 +188,10 @@ async function readSSEStream(url, fetchOptions, onEvent) { const e = body && body.error; if (e && typeof e === "object") { const prefix = e.code ? `${e.status || res.status} ${e.code}` : `${e.status || res.status}`; - addError(`${prefix} — ${e.message || "Request failed"}`); + // Prefer message, then code, then a generic fallback — handles {code} bodies + // with no message field. + const detail = e.message || e.code || "Request failed"; + addError(`${prefix} — ${detail}`); } else { addError(`Error: ${typeof e === "string" ? e : res.statusText}`); } diff --git a/src/kai_client/cli.py b/src/kai_client/cli.py index f947a8d..da92730 100644 --- a/src/kai_client/cli.py +++ b/src/kai_client/cli.py @@ -149,6 +149,168 @@ async def _info(): run_async(_info()) +class _VerifyRecorder: + """Collects per-check results for `kai verify` and prints them inline.""" + + def __init__(self, json_output: bool) -> None: + self.json_output = json_output + self.report: dict[str, Any] = {"ok": True, "checks": {}} + + def record(self, name: str, ok: bool, summary: str, **extra: Any) -> None: + self.report["checks"][name] = {"ok": ok, "summary": summary, **extra} + if not ok: + self.report["ok"] = False + if not self.json_output: + mark = click.style("✓", fg="green") if ok else click.style("✗", fg="red") + click.echo(f"{mark} {name}: {summary}") + + +def _format_kai_error(e: KaiError) -> str: + """Render a KaiError as `code: message`, falling back to message-only.""" + if e.code and e.message: + return f"{e.code}: {e.message}" + return e.message or str(e) + + +def _safe_json_body(resp: httpx.Response) -> dict[str, Any]: + """Parse `resp` as a JSON object, falling back to `{'error': }`.""" + if "application/json" in resp.headers.get("content-type", ""): + try: + parsed = resp.json() + except ValueError: + parsed = None + if isinstance(parsed, dict): + return parsed + return {"error": resp.text} + + +async def _check_token(url: str, token: str, recorder: _VerifyRecorder) -> Optional[dict[str, Any]]: + """Verify the token via Storage API; return the parsed token info on success.""" + async with httpx.AsyncClient() as http_client: + try: + resp = await http_client.get( + f"{url.rstrip('/')}/v2/storage/tokens/verify", + headers={"x-storageapi-token": token}, + timeout=30.0, + ) + except httpx.RequestError as e: + recorder.record("token", False, f"connection error: {e}") + return None + + if resp.status_code >= 400: + body = _safe_json_body(resp) + code = body.get("code") or f"http:{resp.status_code}" + msg = body.get("error") or body.get("message") or resp.reason_phrase + recorder.record("token", False, f"HTTP {resp.status_code} {code}: {msg}") + return None + + parsed = resp.json() + if not isinstance(parsed, dict): + recorder.record("token", False, f"unexpected response shape: {type(parsed).__name__}") + return None + + owner = parsed.get("owner") or {} + project_id = owner.get("id") + project_name = owner.get("name", "?") + description = parsed.get("description", "?") + is_master = parsed.get("isMasterToken", False) + token_type = "master" if is_master else "scoped" + summary = f"project {project_id} ({project_name}), token '{description}' [{token_type}]" + recorder.record( + "token", + True, + summary, + project_id=project_id, + project_name=project_name, + token_description=description, + is_master_token=is_master, + ) + return parsed + + +async def _check_service_discovery( + token: str, url: str, base_url: Optional[str], recorder: _VerifyRecorder +) -> Optional[KaiClient]: + """Build a KaiClient (auto-discover or honour --base-url for local dev).""" + try: + if base_url: + client = KaiClient(storage_api_token=token, storage_api_url=url, base_url=base_url) + recorder.record("service-discovery", True, f"using --base-url {base_url}") + else: + client = await KaiClient.from_storage_api(storage_api_token=token, storage_api_url=url) + recorder.record( + "service-discovery", + True, + f"kai-assistant at {client.base_url}", + ) + except KaiError as e: + recorder.record( + "service-discovery", + False, + _format_kai_error(e), + code=e.code, + message=e.message, + ) + return None + return client + + +async def _check_reachability(client: KaiClient, recorder: _VerifyRecorder) -> None: + """Ping + info (no auth). Records each independently.""" + try: + ping_resp = await client.ping() + recorder.record("ping", True, f"server alive at {ping_resp.timestamp.isoformat()}") + except KaiError as e: + recorder.record("ping", False, _format_kai_error(e), code=e.code, message=e.message) + + try: + info_resp = await client.info() + recorder.record( + "info", + True, + f"{info_resp.app_name} v{info_resp.app_version} " + f"(server v{info_resp.server_version}, uptime {info_resp.uptime:.0f}s)", + server_version=info_resp.server_version, + ) + except KaiError as e: + recorder.record("info", False, _format_kai_error(e), code=e.code, message=e.message) + + +async def _check_usage(client: KaiClient, recorder: _VerifyRecorder) -> None: + """Authenticated probe + monthly message quota. + + The KaiError handler renders any upstream `{code, message}` cleanly. The + headline symptom this command was built for is 429 ``rate_limit:chat``, + but 401/403/timeout etc. all surface through the same path. + """ + try: + usage = await client.get_usage() + remaining = usage.messages_limit - usage.messages_used + recorder.record( + "usage", + True, + f"{usage.messages_used}/{usage.messages_limit} messages used " + f"(resets {usage.reset_date.date().isoformat()}, {remaining} left)", + messages_used=usage.messages_used, + messages_limit=usage.messages_limit, + reset_date=usage.reset_date.isoformat(), + ) + except KaiError as e: + recorder.record("usage", False, _format_kai_error(e), code=e.code, message=e.message) + + +def _emit_verify_footer(recorder: _VerifyRecorder) -> None: + """Print the final summary line (JSON dump or human-readable footer).""" + if recorder.json_output: + click.echo(json.dumps(recorder.report, indent=2, default=str)) + elif recorder.report["ok"]: + click.echo() + click.echo(click.style("All checks passed.", fg="green")) + else: + click.echo() + click.echo(click.style("Some checks failed — see above.", fg="red"), err=True) + + @main.command() @click.option("--json-output", is_flag=True, help="Output as JSON") @click.pass_context @@ -171,165 +333,29 @@ async def _verify(): url = ctx.obj.get("url") or get_env_or_error("STORAGE_API_URL") base_url = ctx.obj.get("base_url") - report: dict[str, Any] = {"ok": True, "checks": {}} - - def _print_check(name: str, ok: bool, summary: str) -> None: - if json_output: - return - mark = click.style("✓", fg="green") if ok else click.style("✗", fg="red") - click.echo(f"{mark} {name}: {summary}") - - def _record(name: str, ok: bool, summary: str, **extra: Any) -> None: - report["checks"][name] = {"ok": ok, "summary": summary, **extra} - if not ok: - report["ok"] = False - _print_check(name, ok, summary) - + recorder = _VerifyRecorder(json_output) if not json_output: click.echo(f"Storage API URL: {url}") - # 1. Token identity (Keboola Storage API) - token_data: Optional[dict[str, Any]] = None - async with httpx.AsyncClient() as http_client: - try: - resp = await http_client.get( - f"{url.rstrip('/')}/v2/storage/tokens/verify", - headers={"x-storageapi-token": token}, - timeout=30.0, - ) - if resp.status_code >= 400: - body: dict[str, Any] = {} - if "application/json" in resp.headers.get("content-type", ""): - parsed = resp.json() - if isinstance(parsed, dict): - body = parsed - if not body: - body = {"error": resp.text} - code = body.get("code") or f"http:{resp.status_code}" - msg = body.get("error") or body.get("message") or resp.reason_phrase - _record("token", False, f"HTTP {resp.status_code} {code}: {msg}") - else: - parsed = resp.json() - assert isinstance(parsed, dict), ( - "Storage API token verify must return JSON object" - ) - token_data = parsed - owner = token_data.get("owner") or {} - project_id = owner.get("id") - project_name = owner.get("name", "?") - description = token_data.get("description", "?") - is_master = token_data.get("isMasterToken", False) - token_type = "master" if is_master else "scoped" - summary = ( - f"project {project_id} ({project_name}), " - f"token '{description}' [{token_type}]" - ) - _record( - "token", - True, - summary, - project_id=project_id, - project_name=project_name, - token_description=description, - is_master_token=is_master, - ) - except httpx.RequestError as e: - _record("token", False, f"connection error: {e}") - - # If the token didn't verify, the remaining checks will all fail the same way. - # Bail early with a non-zero exit so this command stays useful in scripts. - if not token_data: - if json_output: - click.echo(json.dumps(report, indent=2, default=str)) + if await _check_token(url, token, recorder) is None: + _emit_verify_footer(recorder) sys.exit(1) - # 2. Build a KaiClient: respect --base-url for local dev; otherwise auto-discover. - client: Optional[KaiClient] = None - try: - if base_url: - client = KaiClient(storage_api_token=token, storage_api_url=url, base_url=base_url) - _record("service-discovery", True, f"using --base-url {base_url}") - else: - client = await KaiClient.from_storage_api( - storage_api_token=token, storage_api_url=url - ) - _record("service-discovery", True, f"kai-assistant at {client.base_url}") - except KaiError as e: - _record( - "service-discovery", - False, - _format_kai_error(e), - code=e.code, - message=e.message, - ) - + client = await _check_service_discovery(token, url, base_url, recorder) if client is None: - if json_output: - click.echo(json.dumps(report, indent=2, default=str)) + _emit_verify_footer(recorder) sys.exit(1) async with client: - # 3. Reachability (no auth needed) - try: - ping_resp = await client.ping() - _record("ping", True, f"server alive at {ping_resp.timestamp.isoformat()}") - except KaiError as e: - _record("ping", False, _format_kai_error(e), code=e.code, message=e.message) - - try: - info_resp = await client.info() - _record( - "info", - True, - f"{info_resp.app_name} v{info_resp.app_version} " - f"(server v{info_resp.server_version}, uptime {info_resp.uptime:.0f}s)", - server_version=info_resp.server_version, - ) - except KaiError as e: - _record("info", False, _format_kai_error(e), code=e.code, message=e.message) - - # 4. Authenticated probe + quota — the whole point of this command. - try: - usage = await client.get_usage() - remaining = usage.messages_limit - usage.messages_used - summary = ( - f"{usage.messages_used}/{usage.messages_limit} messages used " - f"(resets {usage.reset_date.date().isoformat()}, {remaining} left)" - ) - _record( - "usage", - True, - summary, - messages_used=usage.messages_used, - messages_limit=usage.messages_limit, - reset_date=usage.reset_date.isoformat(), - ) - except KaiError as e: - # Specifically surface 429 rate_limit:chat cleanly — that's the symptom - # this whole command exists to diagnose. - _record("usage", False, _format_kai_error(e), code=e.code, message=e.message) - - if json_output: - click.echo(json.dumps(report, indent=2, default=str)) - elif report["ok"]: - click.echo() - click.echo(click.style("All checks passed.", fg="green")) - else: - click.echo() - click.echo(click.style("Some checks failed — see above.", fg="red"), err=True) + await _check_reachability(client, recorder) + await _check_usage(client, recorder) - sys.exit(0 if report["ok"] else 1) + _emit_verify_footer(recorder) + sys.exit(0 if recorder.report["ok"] else 1) run_async(_verify()) -def _format_kai_error(e: KaiError) -> str: - """Render a KaiError as `code: message`, falling back to message-only.""" - if e.code and e.message: - return f"{e.code}: {e.message}" - return e.message or str(e) - - @main.command() @click.option("-m", "--message", help="Send a single message instead of interactive mode") @click.option( diff --git a/tests/test_cli_verify.py b/tests/test_cli_verify.py index 3468806..8da97b0 100644 --- a/tests/test_cli_verify.py +++ b/tests/test_cli_verify.py @@ -2,6 +2,7 @@ import json +import httpx import pytest from click.testing import CliRunner @@ -161,3 +162,92 @@ def test_json_output_shape(self, runner, mock_env, httpx_mock): assert data["checks"]["token"]["project_id"] == 2738 assert data["checks"]["usage"]["messages_used"] == 19 assert data["checks"]["usage"]["messages_limit"] == 150 + + def test_scoped_token_renders_as_scoped(self, runner, mock_env, httpx_mock): + """Non-master tokens should label as [scoped] (not [master]).""" + _mock_token_verify_ok(httpx_mock, is_master=False) + _mock_discovery_ok(httpx_mock) + _mock_ping_info_ok(httpx_mock) + httpx_mock.add_response( + url=f"{KAI_URL}/api/usage", + method="GET", + json={ + "messagesUsed": 5, + "messagesLimit": 150, + "resetDate": "2026-06-01T00:00:00Z", + }, + ) + + result = runner.invoke(main, ["verify"]) + + assert result.exit_code == 0, result.output + assert "[scoped]" in result.output + assert "[master]" not in result.output + + def test_storage_api_connection_error_renders_cleanly(self, runner, mock_env, httpx_mock): + """If Storage API is unreachable, token check fails with a connection-error line.""" + httpx_mock.add_exception( + httpx.ConnectError("Network unreachable"), + url=f"{STORAGE_URL}/v2/storage/tokens/verify", + ) + + result = runner.invoke(main, ["verify"]) + + assert result.exit_code == 1 + assert "connection error" in result.output.lower() + assert "Network unreachable" in result.output + + def test_base_url_skips_discovery(self, runner, mock_env, httpx_mock): + """`--base-url` bypasses /v2/storage discovery and talks to the given URL directly.""" + local_url = "http://localhost:4000" + _mock_token_verify_ok(httpx_mock) + # No /v2/storage mock — would error if discovery were attempted. + httpx_mock.add_response( + url=f"{local_url}/ping", + method="GET", + json={"timestamp": "2026-05-26T10:00:00Z"}, + ) + httpx_mock.add_response( + url=f"{local_url}/api", + method="GET", + json={ + "timestamp": "2026-05-26T10:00:00Z", + "appName": "kai-assistant", + "appVersion": "dev", + "serverVersion": "dev", + "uptime": 12.0, + "connectedMcp": [], + }, + ) + httpx_mock.add_response( + url=f"{local_url}/api/usage", + method="GET", + json={ + "messagesUsed": 0, + "messagesLimit": 150, + "resetDate": "2026-06-01T00:00:00Z", + }, + ) + + result = runner.invoke(main, ["--base-url", local_url, "verify"]) + + assert result.exit_code == 0, result.output + assert f"using --base-url {local_url}" in result.output + assert "All checks passed." in result.output + + def test_service_discovery_failure_when_kai_assistant_absent( + self, runner, mock_env, httpx_mock + ): + """If the services list has no kai-assistant entry, discovery fails cleanly.""" + _mock_token_verify_ok(httpx_mock) + httpx_mock.add_response( + url=f"{STORAGE_URL}/v2/storage", + method="GET", + json={"services": [{"id": "queue", "url": "https://queue.example/"}]}, + ) + + result = runner.invoke(main, ["verify"]) + + assert result.exit_code == 1 + assert "service-discovery" in result.output + assert "kai-assistant service not found" in result.output From 4451e8633bbab9b32b7af59ea819ec811a15ea8c Mon Sep 17 00:00:00 2001 From: ottomansky Date: Tue, 26 May 2026 14:42:45 +0200 Subject: [PATCH 3/3] fix: harden _check_token JSON parsing, add reachability + json-failure tests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - `_check_token`: guard `resp.json()` on the 2xx path against `JSONDecodeError` so a malformed Storage API response surfaces as a recorded check failure instead of crashing. - `_check_reachability`: expanded docstring justifying why ping/info are recorded independently (partial-outage diagnostic visibility) rather than short-circuiting on ping failure. - Two new tests: - `test_json_output_on_failure_path` — the JSON report is well-formed on a token-verify failure and includes only the checks that actually ran (no spurious entries for later phases). - `test_reachability_records_ping_and_info_independently` — when `/ping` returns 503 but `/api` and `/api/usage` succeed, all three checks are still recorded; verify exits non-zero overall because `ok=false` propagates. Suite is now 259 passing (was 257). Correction to the previous commit (86ca38e): its message claimed to revert two `ruff format` hunks in `send_and_display`. That revert did not stick — the subsequent `ruff format` pass re-applied them, so the hunks remain in this PR's diff. Accepting that outcome: `ruff format` is the project's declared formatter (`pyproject.toml:56`), so those two lines are now the project style. The previous longer form was stale because nobody had run the formatter recently. Apologies for the inaccurate claim in the prior commit message. Co-Authored-By: Claude Opus 4.7 (1M context) --- src/kai_client/cli.py | 18 +++++++++-- tests/test_cli_verify.py | 64 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 80 insertions(+), 2 deletions(-) diff --git a/src/kai_client/cli.py b/src/kai_client/cli.py index da92730..de98a76 100644 --- a/src/kai_client/cli.py +++ b/src/kai_client/cli.py @@ -204,7 +204,14 @@ async def _check_token(url: str, token: str, recorder: _VerifyRecorder) -> Optio recorder.record("token", False, f"HTTP {resp.status_code} {code}: {msg}") return None - parsed = resp.json() + # Symmetric with the >= 400 branch: tolerate a malformed-but-2xx body + # so the diagnostic surfaces "Storage API returned malformed JSON" + # instead of crashing with JSONDecodeError. + try: + parsed = resp.json() + except ValueError as e: + recorder.record("token", False, f"malformed JSON in 2xx response: {e}") + return None if not isinstance(parsed, dict): recorder.record("token", False, f"unexpected response shape: {type(parsed).__name__}") return None @@ -256,7 +263,14 @@ async def _check_service_discovery( async def _check_reachability(client: KaiClient, recorder: _VerifyRecorder) -> None: - """Ping + info (no auth). Records each independently.""" + """Ping + info (no auth). + + Both checks run independently — `info` is not short-circuited when `ping` + fails. The two endpoints can fail in different ways (e.g. ping returns 200 + but info returns 500 if a downstream MCP server is unhealthy), so showing + both results yields better diagnostic signal than collapsing to a single + line. The cost is one extra HTTP request on a totally-down service. + """ try: ping_resp = await client.ping() recorder.record("ping", True, f"server alive at {ping_resp.timestamp.isoformat()}") diff --git a/tests/test_cli_verify.py b/tests/test_cli_verify.py index 8da97b0..f0fc9b7 100644 --- a/tests/test_cli_verify.py +++ b/tests/test_cli_verify.py @@ -251,3 +251,67 @@ def test_service_discovery_failure_when_kai_assistant_absent( assert result.exit_code == 1 assert "service-discovery" in result.output assert "kai-assistant service not found" in result.output + + def test_json_output_on_failure_path(self, runner, mock_env, httpx_mock): + """A failed token verify with --json-output still returns a well-formed JSON report.""" + httpx_mock.add_response( + url=f"{STORAGE_URL}/v2/storage/tokens/verify", + method="GET", + status_code=401, + json={"error": "Invalid access token", "code": "storage.tokenInvalid"}, + ) + + result = runner.invoke(main, ["verify", "--json-output"]) + + assert result.exit_code == 1 + data = json.loads(result.output) + assert data["ok"] is False + assert data["checks"]["token"]["ok"] is False + assert "storage.tokenInvalid" in data["checks"]["token"]["summary"] + # On token-check failure, later phases must not run. + assert "service-discovery" not in data["checks"] + assert "ping" not in data["checks"] + assert "usage" not in data["checks"] + + def test_reachability_records_ping_and_info_independently(self, runner, mock_env, httpx_mock): + """When ping fails but info succeeds, both are recorded — info is not short-circuited.""" + _mock_token_verify_ok(httpx_mock) + _mock_discovery_ok(httpx_mock) + httpx_mock.add_response( + url=f"{KAI_URL}/ping", + method="GET", + status_code=503, + json={"code": "service_unavailable", "message": "ping path is down"}, + ) + httpx_mock.add_response( + url=f"{KAI_URL}/api", + method="GET", + json={ + "timestamp": "2026-05-26T10:00:00Z", + "appName": "kai-assistant", + "appVersion": "1.2.3", + "serverVersion": "1.2.3", + "uptime": 3600.0, + "connectedMcp": [], + }, + ) + httpx_mock.add_response( + url=f"{KAI_URL}/api/usage", + method="GET", + json={ + "messagesUsed": 10, + "messagesLimit": 150, + "resetDate": "2026-06-01T00:00:00Z", + }, + ) + + result = runner.invoke(main, ["verify", "--json-output"]) + + # Overall ok=false because ping failed, but info AND usage still ran. + assert result.exit_code == 1 + data = json.loads(result.output) + assert data["checks"]["ping"]["ok"] is False + assert "service_unavailable" in data["checks"]["ping"]["summary"] + assert data["checks"]["info"]["ok"] is True + assert data["checks"]["usage"]["ok"] is True + assert data["checks"]["usage"]["messages_used"] == 10