agentkit-sdk-python 0.6.2__py3-none-any.whl → 0.6.4__py3-none-any.whl

agentkit/toolkit/cli/cli_invoke.py

@@ -16,6 +16,8 @@
 
 from pathlib import Path
 from typing import Optional, Any
+import base64
+import binascii
 import json
 import typer
 from typer.core import TyperGroup
@@ -627,6 +629,160 @@ def build_harness_overrides(
     return overrides
 
 
+# Fixed ADK app name for the run_sse path. The harness loader serves its single
+# agent under any app name, so a stable constant keeps the CLI decoupled from the
+# deployed HARNESS_NAME.
+_HARNESS_RUN_SSE_APP = "harness"
+
+
+def _user_id_from_token(token: str) -> str | None:
+    """Return the OIDC ``sub`` claim from a JWT bearer token, else ``None``.
+
+    A custom_jwt harness is called with an OIDC id_token whose ``sub`` is the
+    authenticated user's stable id — use it as the run's user_id so sessions are
+    tied to the real identity. A key_auth token is an opaque api key (not a JWT,
+    no ``sub``), so this returns ``None`` and the caller falls back to a random id.
+    """
+    parts = token.split(".")
+    if len(parts) != 3:
+        return None  # not a JWT (e.g. a key_auth api key)
+    payload = parts[1]
+    payload += "=" * (-len(payload) % 4)  # restore base64 padding
+    try:
+        claims = json.loads(base64.urlsafe_b64decode(payload))
+    except (ValueError, binascii.Error):
+        return None  # malformed JWT payload → fall back to random
+    sub = claims.get("sub")
+    return sub if isinstance(sub, str) and sub else None
+
+
+def _harness_run_sse(
+    *,
+    base_url: str,
+    token: str,
+    prompt: str,
+    session_id: str,
+    overrides: dict,
+    raw: bool,
+) -> Any:
+    """Invoke a deployed harness via the ADK ``/run_sse`` endpoint (streaming).
+
+    app_name is the fixed ``"harness"``; user_id is the JWT ``sub`` when the token
+    is an OIDC id_token, else a random id; session_id is the caller's. When
+    ``overrides`` is non-empty it is sent as the ``harness`` field so the runtime
+    streams a spawned (overridden) agent; otherwise the base agent.
+    """
+    import requests
+
+    app_name = _HARNESS_RUN_SSE_APP
+    sub = _user_id_from_token(token)
+    user_id = sub or f"u-{uuid.uuid4().hex[:12]}"
+    user_id_origin = "jwt sub" if sub else "random"
+    headers = {"Content-Type": "application/json"}
+    if token:
+        headers["Authorization"] = f"Bearer {token}"
+    console.print(
+        f"[blue]run_sse: app_name={app_name}, user_id={user_id} ({user_id_origin}), "
+        f"session_id={session_id}[/blue]"
+    )
+
+    # ADK /run_sse requires an existing session; create it (ignore "exists").
+    session_url = f"{base_url}/apps/{app_name}/users/{user_id}/sessions/{session_id}"
+    try:
+        sr = requests.post(session_url, json={}, headers=headers, timeout=60)
+    except requests.RequestException as e:
+        console.print(f"[red]❌ Session create failed: {e}[/red]")
+        raise typer.Exit(1)
+    if sr.status_code not in (200, 400, 409):
+        console.print(
+            f"[red]❌ Session create HTTP {sr.status_code}: {sr.text[:300]}[/red]"
+        )
+        raise typer.Exit(1)
+
+    body: dict[str, Any] = {
+        "app_name": app_name,
+        "user_id": user_id,
+        "session_id": session_id,
+        "new_message": {"role": "user", "parts": [{"text": prompt}]},
+        "streaming": True,
+    }
+    if overrides:
+        body["harness"] = overrides
+        console.print(f"[blue]Using one-time overrides: {overrides}[/blue]")
+    try:
+        resp = requests.post(
+            f"{base_url}/run_sse",
+            json=body,
+            headers=headers,
+            timeout=300,
+            stream=True,
+        )
+    except requests.RequestException as e:
+        console.print(f"[red]❌ run_sse request failed: {e}[/red]")
+        raise typer.Exit(1)
+    if resp.status_code != 200:
+        console.print(
+            f"[red]❌ run_sse HTTP {resp.status_code}: {resp.text[:500]}[/red]"
+        )
+        raise typer.Exit(1)
+
+    def _answer_text(event: dict) -> str:
+        # Only the answer text; the model's "thought" (reasoning) parts stay hidden
+        # behind the thinking spinner.
+        parts = (event.get("content") or {}).get("parts") or []
+        return "".join(
+            p["text"]
+            for p in parts
+            if isinstance(p, dict) and p.get("text") and not p.get("thought")
+        )
+
+    streamed = []
+    final_answer = ""
+    answer_open = False  # first answer token seen → spinner stopped, reply started
+    # A spinner covers the wait (model latency + reasoning) until the answer starts.
+    status = None if raw else console.status("thinking", spinner="dots")
+    if status is not None:
+        status.start()
+    try:
+        for line in resp.iter_lines(decode_unicode=True):
+            if not line:
+                continue
+            if raw:
+                print(line)  # builtin print: no rich wrapping/markup on raw JSON
+                continue
+            event = _normalize_stream_event(line)
+            if not isinstance(event, dict):
+                continue
+            if event.get("error"):
+                if status is not None and not answer_open:
+                    status.stop()
+                console.print(f"\n[red]Error: {event['error']}[/red]")
+                continue
+            answer = _answer_text(event)
+            # `partial=False` is the final aggregate (repeats everything) — keep it
+            # as a fallback but don't print it; partial events stream the deltas.
+            if event.get("partial") is False:
+                final_answer = answer or final_answer
+                continue
+            if answer:
+                if not answer_open:
+                    if status is not None:
+                        status.stop()  # leave the thinking phase
+                    console.print("")  # blank line before the reply
+                    answer_open = True
+                console.print(answer, end="", style="green")
+                streamed.append(answer)
+    finally:
+        if status is not None and not answer_open:
+            status.stop()
+
+    if not streamed and final_answer:
+        console.print("")
+        console.print(final_answer, end="", style="green")
+    console.print("")
+    return "".join(streamed) or final_answer
+
+
 @invoke_app.command("harness")
 def harness_command(
     name: str = typer.Argument(
@@ -640,7 +796,9 @@ def harness_command(
         "agentkit_user", "--user-id", help="user_id for the run."
     ),
     session_id: str = typer.Option(
-        "agentkit_sample_session", "--session-id", help="session_id for the run."
+        None,
+        "--session-id",
+        help="session_id for the run (random if unset).",
     ),
     max_llm_calls: int = typer.Option(
         None,
@@ -675,7 +833,12 @@ def harness_command(
         help="Bearer token override (e.g. OAuth JWT for custom_jwt harnesses).",
     ),
     raw: bool = typer.Option(
-        False, "--raw", help="Print the raw InvokeHarnessResponse JSON."
+        False, "--raw", help="Print the raw response (InvokeHarnessResponse / SSE)."
+    ),
+    protocol: str = typer.Option(
+        "run_sse",
+        "--protocol",
+        help="Transport: 'run_sse' (ADK /run_sse, default) or 'invoke' (POST /harness/invoke).",
     ),
 ) -> Any:
     """Invoke a deployed harness by name (resolved via the harness.json registry).
@@ -712,8 +875,69 @@ def harness_command(
         )
         raise typer.Exit(1)
 
-    invoke_url = entry["url"].rstrip("/") + "/harness/invoke"
-    token = apikey or entry.get("key") or ""
+    if protocol not in ("invoke", "run_sse"):
+        console.print(
+            f"[red]Error: --protocol must be 'invoke' or 'run_sse', got '{protocol}'.[/red]"
+        )
+        raise typer.Exit(1)
+
+    base_url = entry["url"].rstrip("/")
+
+    # Inbound credential selection is auth-type aware — the registry records how each
+    # harness was deployed (harness/deploy.py _record_harness): a custom_jwt harness is
+    # {auth_type:"custom_jwt", no "key"}; a key_auth harness is {"key": ...}.
+    #   1. --apikey always wins (explicit manual override).
+    #   2. custom_jwt harness → the `agentkit login` id_token (OIDC JWT, auto-refreshed) —
+    #      the data plane's JWT path; refresh failure errors (re-login), never silent.
+    #   3. key_auth harness → the static "key"; a key_auth authorizer would reject a JWT,
+    #      so we never force the id_token onto it.
+    # Both transports (invoke and run_sse) share this resolution.
+    from agentkit.auth.errors import AuthError
+    from agentkit.auth.sso import load_session
+
+    is_jwt_harness = entry.get("auth_type") == "custom_jwt" or not entry.get("key")
+    auth_session = None  # set only when the login id_token is the credential (enables 401 refresh)
+    token = apikey or ""
+    if not token and is_jwt_harness:
+        try:
+            auth_session = load_session()
+        except Exception:
+            auth_session = None
+        if auth_session is not None:
+            try:
+                token = auth_session.valid_id_token()
+            except AuthError as e:
+                console.print(f"[red]❌ {e}[/red]")
+                if getattr(e, "hint", None):
+                    console.print(f"[yellow]{e.hint}[/yellow]")
+                raise typer.Exit(1)
+    if not token:
+        token = entry.get("key") or ""
+
+    # No session given → mint a random one (both transports behave identically;
+    # creating it is idempotent).
+    session_id = session_id or f"s-{uuid.uuid4().hex[:12]}"
+
+    if protocol == "run_sse":
+        # run_sse supports the same overrides (sent as the `harness` field); only
+        # --max-llm-calls is invoke-only (not part of the ADK run_sse request).
+        if max_llm_calls is not None:
+            console.print(
+                "[yellow]Note: --max-llm-calls is ignored with --protocol "
+                "run_sse.[/yellow]"
+            )
+        return _harness_run_sse(
+            base_url=base_url,
+            token=token,
+            prompt=message,
+            session_id=session_id,
+            overrides=build_harness_overrides(
+                system_prompt, model_name, tools, skills, runtime
+            ),
+            raw=raw,
+        )
+
+    invoke_url = base_url + "/harness/invoke"
     req_headers = {"Content-Type": "application/json"}
     if token:
         req_headers["Authorization"] = f"Bearer {token}"
@@ -739,6 +963,18 @@ def harness_command(
 
     try:
         resp = requests.post(invoke_url, json=body, headers=req_headers, timeout=300)
+        # If the harness rejects a locally-valid JWT (clock skew, mid-flight rotation,
+        # server-side revocation), force one refresh and retry exactly once.
+        if resp.status_code == 401 and auth_session is not None and not apikey:
+            try:
+                token = auth_session.valid_id_token(force_refresh=True)
+            except AuthError as e:
+                console.print(f"[red]❌ session refresh failed: {e}[/red]")
+                if getattr(e, "hint", None):
+                    console.print(f"[yellow]{e.hint}[/yellow]")
+                raise typer.Exit(1)
+            req_headers["Authorization"] = f"Bearer {token}"
+            resp = requests.post(invoke_url, json=body, headers=req_headers, timeout=300)
     except requests.RequestException as e:
         console.print(f"[red]❌ Request to {invoke_url} failed: {e}[/red]")
         raise typer.Exit(1)
@@ -755,10 +991,20 @@ def harness_command(
         console.print(f"[red]❌ Non-JSON response: {resp.text}[/red]")
         raise typer.Exit(1)
 
+    # The harness returns failures (unsupported tool, skill load failure, or a
+    # runtime error) in `error`; surface it verbatim instead of as normal output.
+    error = data.get("error") if isinstance(data, dict) else None
+
     if raw:
         console.print(json.dumps(data, ensure_ascii=False, indent=2))
+        if error:
+            raise typer.Exit(1)
         return str(data)
 
+    if error:
+        console.print(f"[red]❌ Harness error: {error}[/red]")
+        raise typer.Exit(1)
+
     console.print("[green]✅ Invocation successful[/green]")
     if data.get("overwrite"):
         console.print("[yellow](served with one-time overrides)[/yellow]")

agentkit/toolkit/cli/cli_list.py

@@ -46,6 +46,39 @@ def _is_harness_runtime(runtime: rt.AgentKitRuntimesForListRuntimes) -> bool:
     return False
 
 
+# Fixed ADK app name the harness loader serves its agent under (the deployed
+# HARNESS_NAME is irrelevant to the ADK app path).
+_HARNESS_APP = "harness"
+
+
+def _user_id_from_token(token: str) -> Optional[str]:
+    """Return the OIDC ``sub`` claim from a JWT bearer token, else ``None``.
+
+    A custom_jwt harness is reached with an OIDC id_token whose ``sub`` is the
+    authenticated user's stable id. A key_auth token is an opaque api key (not a
+    JWT, no ``sub``), so this returns ``None`` and the caller must be given an
+    explicit ``--user-id``.
+
+    NOTE: this mirrors ``cli_invoke._user_id_from_token``; consolidate the two
+    into one shared helper once the ``invoke harness`` run_sse work lands.
+    """
+    import base64
+    import binascii
+    import json
+
+    parts = token.split(".")
+    if len(parts) != 3:
+        return None  # not a JWT (e.g. a key_auth api key)
+    payload = parts[1]
+    payload += "=" * (-len(payload) % 4)  # restore base64 padding
+    try:
+        claims = json.loads(base64.urlsafe_b64decode(payload))
+    except (ValueError, binascii.Error):
+        return None  # malformed JWT payload
+    sub = claims.get("sub")
+    return sub if isinstance(sub, str) and sub else None
+
+
 @list_app.command("harness")
 def list_harness_command(
     region: Optional[str] = typer.Option(
@@ -117,6 +150,7 @@ def list_harness_command(
         ("RuntimeId", "RuntimeId", "cyan"),
         ("Name", "Name", "white"),
         ("Status", "Status", "green"),
+        ("CurrentVersionNumber", "Version", "blue"),
         ("ProjectName", "ProjectName", "yellow"),
         ("UpdatedAt", "UpdatedAt", "magenta"),
     ]
@@ -128,6 +162,173 @@ def list_harness_command(
     local_console.print(table)
 
 
+def _resolve_harness_endpoint(harness: str, directory: str, apikey: Optional[str]):
+    """Resolve a deployed harness's data-plane base URL and bearer token.
+
+    Mirrors ``invoke harness``'s auth-type-aware credential selection:
+      1. ``--apikey`` always wins (explicit override);
+      2. a custom_jwt harness (no stored ``key``) → the ``agentkit login`` OIDC
+         id_token (auto-refreshed), whose ``sub`` identifies the caller;
+      3. a key_auth harness → its static ``key`` (an opaque api key, not a JWT).
+
+    Returns ``(base_url, token)``. Fast-fails when the harness is not in the
+    registry, or when a needed login id_token cannot be refreshed.
+    """
+    from agentkit.toolkit.harness import load_harness_registry
+
+    registry = load_harness_registry(directory)
+    entry = registry.get(harness)
+    if not isinstance(entry, dict) or not entry.get("url"):
+        from pathlib import Path
+
+        registry_path = Path(directory).resolve() / "harness.json"
+        console.print(
+            f"[red]Error: harness '{harness}' not found in registry {registry_path}. "
+            f"Deploy it first with `agentkit deploy --harness {harness}`.[/red]"
+        )
+        raise typer.Exit(1)
+
+    base_url = entry["url"].rstrip("/")
+
+    from agentkit.auth.errors import AuthError
+    from agentkit.auth.sso import load_session
+
+    is_jwt_harness = entry.get("auth_type") == "custom_jwt" or not entry.get("key")
+    token = apikey or ""
+    if not token and is_jwt_harness:
+        try:
+            auth_session = load_session()
+        except Exception:
+            auth_session = None
+        if auth_session is not None:
+            try:
+                token = auth_session.valid_id_token()
+            except AuthError as e:
+                console.print(f"[red]❌ {e}[/red]")
+                if getattr(e, "hint", None):
+                    console.print(f"[yellow]{e.hint}[/yellow]")
+                raise typer.Exit(1)
+    if not token:
+        token = entry.get("key") or ""
+
+    return base_url, token
+
+
+@list_app.command("sessions")
+def list_sessions_command(
+    harness: str = typer.Option(
+        ..., "--harness", help="Harness name (resolved via the harness.json registry)."
+    ),
+    user_id: Optional[str] = typer.Option(
+        None,
+        "--user-id",
+        help=(
+            "user_id whose sessions to list. If omitted, it is taken from the "
+            "JWT `sub` of the harness credential; a key_auth harness has no JWT, "
+            "so --user-id is required there."
+        ),
+    ),
+    apikey: Optional[str] = typer.Option(
+        None,
+        "--apikey",
+        "-ak",
+        help="Bearer token (e.g. a custom_jwt harness's OIDC JWT) overriding the registry credential.",
+    ),
+    directory: str = typer.Option(
+        ".", "--directory", help="Directory containing harness.json."
+    ),
+    output: str = typer.Option(
+        "table", "--output", help="Output format: table|json|yaml"
+    ),
+    quiet: bool = typer.Option(
+        False, "--quiet", "-q", help="Print only session ids"
+    ),
+    no_color: bool = typer.Option(
+        False, "--no-color", "-nc", help="Disable colored output for tables/panels"
+    ),
+):
+    """List a user's conversation sessions on a deployed harness runtime.
+
+    Calls the harness runtime's ADK endpoint
+    ``GET /apps/harness/users/{user_id}/sessions``. ADK requires a user_id (there
+    is no cross-user listing), so it must be given explicitly via --user-id or be
+    derivable from the JWT `sub` of the harness credential.
+    """
+    import json
+    from datetime import datetime
+
+    import requests
+    from rich.table import Table
+
+    local_console = console if not no_color else Console(no_color=True)
+
+    base_url, token = _resolve_harness_endpoint(harness, directory, apikey)
+
+    resolved_user_id = user_id or _user_id_from_token(token)
+    if not resolved_user_id:
+        local_console.print(
+            "[red]Error: cannot determine user_id — no --user-id given and the "
+            "harness credential is not a JWT with a `sub` claim (e.g. a key_auth "
+            "harness). Pass --user-id explicitly.[/red]"
+        )
+        raise typer.Exit(1)
+
+    url = f"{base_url}/apps/{_HARNESS_APP}/users/{resolved_user_id}/sessions"
+    headers = {}
+    if token:
+        headers["Authorization"] = f"Bearer {token}"
+
+    with local_console.status("[cyan]Fetching sessions...[/cyan]", spinner="dots"):
+        try:
+            resp = requests.get(url, headers=headers, timeout=60)
+        except requests.RequestException as e:
+            local_console.print(f"[red]❌ List sessions request failed: {e}[/red]")
+            raise typer.Exit(1)
+
+    if resp.status_code != 200:
+        local_console.print(
+            f"[red]❌ List sessions HTTP {resp.status_code}: {resp.text[:300]}[/red]"
+        )
+        raise typer.Exit(1)
+
+    sessions = resp.json() or []
+
+    if quiet:
+        for s in sessions:
+            local_console.print(s.get("id", ""))
+        return
+
+    if output.lower() == "json":
+        local_console.print(json.dumps(sessions, indent=2, ensure_ascii=False))
+        return
+    if output.lower() == "yaml":
+        import yaml
+
+        local_console.print(yaml.safe_dump(sessions, sort_keys=False, allow_unicode=True))
+        return
+
+    table = Table(
+        title=f"Sessions for user '{resolved_user_id}' (Count: {len(sessions)})",
+        show_lines=False,
+    )
+    table.add_column("SessionId", style="cyan")
+    table.add_column("Events", style="green")
+    table.add_column("LastUpdate", style="magenta")
+    for s in sessions:
+        ts = s.get("lastUpdateTime")
+        last_update = (
+            datetime.fromtimestamp(ts).strftime("%Y-%m-%d %H:%M:%S")
+            if isinstance(ts, (int, float)) and ts
+            else ""
+        )
+        table.add_row(
+            str(s.get("id", "")),
+            str(len(s.get("events", []))),
+            last_update,
+        )
+    local_console.print(table)
+
+
 @list_app.command("credentials")
 def list_credentials_command(
     region: Optional[str] = typer.Option(