Read the Ollama Cloud key omp keeps in its agent.db

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

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`:
 

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.