diff --git a/bin/omarchy-agent-usage-ollama b/bin/omarchy-agent-usage-ollama index 230d07c22e7..9d4415f2844 100755 --- a/bin/omarchy-agent-usage-ollama +++ b/bin/omarchy-agent-usage-ollama @@ -17,8 +17,10 @@ formats or endpoints. from __future__ import annotations import argparse +import json import math import os +import sqlite3 import sys import urllib.request from datetime import datetime @@ -55,10 +57,10 @@ AGENT_NAME = "Ollama Cloud" PROVIDER = "ollama-cloud" BALANCE_ENDPOINT = "https://ollama.com/api/balance" ME_ENDPOINT = "https://ollama.com/api/me" -AUTH_HELP_NO_KEY = "Set OLLAMA_API_KEY, or sign in to pi with an Ollama Cloud key, to show the account limits." +AUTH_HELP_NO_KEY = "Set OLLAMA_API_KEY, or sign in to pi or omp with an Ollama Cloud key, to show the account limits." AUTH_HELP_BAD_KEY = ( "Ollama Cloud rejected the API key. Update OLLAMA_API_KEY or ~/.config/omarchy/agents/ollama.json," - " or sign in to pi again. Local stats are still shown." + " or sign in to pi or omp again. Local stats are still shown." ) # GET /api/balance (docs.ollama.com/api/balance) answers a legacy plan with # `included.session` / `included.weekly` ({remaining_percent 0..100, @@ -80,6 +82,46 @@ def pi_agent_dir() -> Path: return Path(os.path.expanduser(configured)) if configured else Path.home() / ".pi" / "agent" +# omp moves its whole agent tree when PI_CODING_AGENT_DIR is set, as pi does. +def omp_agent_dir() -> Path: + return pi_agent_dir() if os.environ.get("PI_CODING_AGENT_DIR", "").strip() else Path.home() / ".omp" / "agent" + + +def literal_key(entry: Any) -> str: + # Only a literal key is taken: a `$VAR` or `!command` entry is resolved by + # pi alone. + key = str(entry.get("key") or "").strip() if isinstance(entry, dict) else "" + return key if not key.startswith(("!", "$")) else "" + + +# omp keeps the keys it signed in with in its agent.db, never in an +# auth.json; a key it disabled is passed over. Read-only, like every store a +# collector reads. +def omp_database_key() -> str: + database = omp_agent_dir() / "agent.db" + if not database.is_file(): + return "" + try: + conn = sqlite3.connect(database.resolve().as_uri() + "?mode=ro", uri=True, timeout=2) + try: + rows = conn.execute( + "SELECT data FROM auth_credentials WHERE provider = ? AND credential_type = 'api_key' AND disabled_cause IS NULL ORDER BY id DESC", + (PROVIDER,), + ).fetchall() + finally: + conn.close() + except sqlite3.Error: + return "" + for (data,) in rows: + try: + key = literal_key(json.loads(data)) + except (TypeError, ValueError): + continue + if key: + return key + return "" + + def api_key() -> str: env = os.environ.get("OLLAMA_API_KEY", "").strip() if env: @@ -88,14 +130,11 @@ def api_key() -> str: if configured: return configured # pi and omp store the Ollama Cloud key they signed in with; either one - # proves the account, and the balance endpoint accepts the same key. Only a - # literal key is taken: a `$VAR` or `!command` entry is resolved by pi alone. - for auth_path in (pi_agent_dir() / "auth.json", Path.home() / ".omp" / "agent" / "auth.json"): - entry = read_json(auth_path).get(PROVIDER) - key = str(entry.get("key") or "").strip() if isinstance(entry, dict) else "" - if key and not key.startswith(("!", "$")): - return key - return "" + # proves the account, and the balance endpoint accepts the same key. + pi_key = literal_key(read_json(pi_agent_dir() / "auth.json").get(PROVIDER)) + if pi_key: + return pi_key + return omp_database_key() or literal_key(read_json(Path.home() / ".omp" / "agent" / "auth.json").get(PROVIDER)) # ---------------------------------------------------------------- local scan diff --git a/shell/plugins/agents/README.md b/shell/plugins/agents/README.md index 581e3852243..8ac630bee35 100644 --- a/shell/plugins/agents/README.md +++ b/shell/plugins/agents/README.md @@ -212,7 +212,7 @@ Z.ai limits need a key that carries a GLM Coding Plan. A pay-as-you-go key repor ### Ollama Cloud -Ollama Cloud reads `OLLAMA_API_KEY` first, then `apiKey` in `~/.config/omarchy/agents/ollama.json`, then the key pi or omp signed in with in `~/.pi/agent/auth.json` (or under `PI_CODING_AGENT_DIR`) or `~/.omp/agent/auth.json`. Each key keeps its own last good limits and plan. +Ollama Cloud reads `OLLAMA_API_KEY` first, then `apiKey` in `~/.config/omarchy/agents/ollama.json`, then the key pi saved in `~/.pi/agent/auth.json`, then the one omp keeps in `~/.omp/agent/agent.db` (both under `PI_CODING_AGENT_DIR` when it is set; a key omp disabled is passed over), then an older omp's `~/.omp/agent/auth.json`. Each key keeps its own last good limits and plan. Ollama's balance endpoint (`/api/balance`) reports a legacy plan's session and weekly windows as the percent remaining with their reset times, and a current plan's included credits as a balance against the monthly allowance, shown as an Included credits meter that resets when the period ends. Credits bought on top of the plan (`purchased.balance_usd`, the unexpired ones left) show as a balance of what remains; the answer says nothing of how much was bought, so it raises no low-balance alarm. The limits cover the whole account — every model the key can reach — while the local stats count pi, omp, and opencode sessions on the `ollama-cloud` provider, and cloud-tagged models (`glm-4.6:cloud`, `gpt-oss:120b-cloud`) reached through the local daemon's `ollama` provider, which run on Ollama's servers too, from Codex's `--oss` rollouts as well (with the prompts typed in them, and any model on a Codex provider named `ollama-cloud`), and Claude Code turns sent to `ollama.com` or answered by a cloud-tagged model, by the [Claude Code](#claude-code) rule. To count only the Claude models Ollama Cloud serves, set a model prefix in `~/.config/omarchy/agents/ollama.json`: diff --git a/test/shell.d/agent-usage-ollama-scanner-test.sh b/test/shell.d/agent-usage-ollama-scanner-test.sh index 5280e9e3e0d..2ca50ddd8b4 100755 --- a/test/shell.d/agent-usage-ollama-scanner-test.sh +++ b/test/shell.d/agent-usage-ollama-scanner-test.sh @@ -259,6 +259,30 @@ pi_auth(None, "ollama-cloud", "$OLLAMA_KEY") expect("Ollama passes over pi's $VAR entry to omp's key", load_collector("ollama").api_key(), "omp-key") pi_auth(home() / ".omp" / "agent", "ollama-cloud", "!pass show ollama") expect("Ollama never sends a pi $VAR or !command entry as a key", load_collector("ollama").api_key(), "") + +# omp keeps the key it signed in with in its agent.db, never in an auth.json; +# a key it disabled is passed over, and PI_CODING_AGENT_DIR moves the +# database with the rest of omp's tree. +def omp_key(directory, *rows): + database = Path(directory) / "agent.db" + database.parent.mkdir(parents=True, exist_ok=True) + conn = sqlite3.connect(database) + conn.execute("CREATE TABLE auth_credentials (id INTEGER PRIMARY KEY, provider TEXT, credential_type TEXT, data TEXT, disabled_cause TEXT)") + for provider, key, disabled in rows: + conn.execute("INSERT INTO auth_credentials (provider, credential_type, data, disabled_cause) VALUES (?, 'api_key', ?, ?)", + (provider, json.dumps({"key": key, "source": "login"}), disabled)) + conn.commit() + conn.close() + + +sandbox() +omp_key(home() / ".omp" / "agent", ("ollama-cloud", "omp-db-key", None), ("ollama-cloud", "revoked-key", "revoked"), ("zai", "zai-key", None)) +pi_auth(home() / ".omp" / "agent", "ollama-cloud", "omp-json-key") +expect("Ollama asks with the key omp keeps in its agent.db, passing over a disabled one", load_collector("ollama").api_key(), "omp-db-key") +sandbox() +omp_key(home() / "moved", ("ollama-cloud", "moved-key", None)) +os.environ["PI_CODING_AGENT_DIR"] = str(home() / "moved") +expect("Ollama finds omp's agent.db where PI_CODING_AGENT_DIR puts it", load_collector("ollama").api_key(), "moved-key") # A MiniMax Code custom provider on the local daemon that answered with a # cloud-tagged model ran on Ollama's servers, so it is Ollama's, prompt and # all, and no local work.