perf(host): cache auth headers and parallelize status payloads

[TomeHirata]

Related issue

N/A

Summary

Two follow-on speedups for omni host status building on #4031.

1. Cache _remote_headers() per base URL

_host_http_json creates a new httpx.Client on every call and passes freshly-resolved headers. For Databricks URLs, resolving headers calls into _resolve_databricks_auth_for_host which constructs a Databricks SDK Config object — this shells out to the Databricks CLI and takes ~3 s each time. Within a single CLI invocation the resulting token is valid for its full lifetime, so resolving once and caching is safe.

A threading.Lock with a double-checked pattern serialises concurrent first-time resolution for the same URL; subsequent callers in other threads hit the cache without taking the lock.

2. Build daemon status payloads in parallel

The payload loop in host_status was a sequential list comprehension. With the dead-process skip from #4031, only live daemons issue HTTP calls — but each live remote daemon still takes ~4 s (auth + RTT). Switching to ThreadPoolExecutor.map lets calls to independent servers overlap instead of stack. The header cache ensures each server's credentials are only resolved once across all threads.

Combined effect on a workstation with 39 daemon records (2 live): ~14 s → ~5 s (from #4031) → same ~5 s for a single remote, but would scale down linearly with multiple live remote daemons.

Test Plan

• Existing 16 tests in tests/host/test_cli_host.py pass, including the new test from perf(host): skip host-status HTTP call for dead daemon processes #4031.
• Manually timed omni host status before/after on a workstation with 1 live remote daemon + 1 live local server.

Demo

N/A

Type of change

• Bug fix
• Feature
• UI / frontend change
• Refactor / chore
• Docs
• Test / CI
• Breaking change

Test coverage

• Unit tests added / updated
• Integration tests added / updated
• E2E tests added / updated
• Manual verification completed
• Existing tests cover this change
• Not applicable

Coverage notes

Changelog

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[TomeHirata]

/review

[omnigent-ci]

Polly AI Review

Blocking issues

Header resolution moved outside the try/except in _host_http_json — narrows error handling on a failure path.

On main, _remote_headers(server_url=base_url) is evaluated inside the try block, so any exception it raises is caught by except (httpx.HTTPError, OSError) and converted into a graceful _HostHttpResult(status_code=0, body="…"):

try:
    with httpx.Client(
        base_url=base_url,
        headers=_remote_headers(server_url=base_url),   # inside try
        timeout=timeout_s,
    ) as client:

The diff hoists resolution before the try:

if base_url not in _host_http_headers_cache:
    with _host_http_headers_lock:
        if base_url not in _host_http_headers_cache:
            _host_http_headers_cache[base_url] = _remote_headers(server_url=base_url)
headers = _host_http_headers_cache[base_url]

try:
    with httpx.Client(base_url=base_url, headers=headers, timeout=timeout_s) as client:

_remote_headers does file I/O (load_token, _read_databrickscfg) and can mint a Databricks workspace token via the SDK (the exact ~3 s path this PR is optimizing) — all of which can raise OSError (and SDK/CLI failures). Previously such a failure produced a per-daemon status_code=0 result that callers surface as a clean error string. Now it propagates as an unhandled exception. Because host_status iterates pool.map(...), the first raising daemon aborts the entire omni host status listing with a traceback — which is precisely the diagnostic-under-failure scenario the command exists to handle gracefully. This also affects every other caller of _host_http_json (e.g. _stop_session_on_server).

Fix: perform the cache-populating _remote_headers call inside a try that maps failures to an _HostHttpResult(status_code=0, …) (or move the resolution back inside the existing try while keeping the cache lookup), so header-resolution errors stay recoverable.

Security vulnerabilities

None. Caching bearer headers in a process-global dict for the lifetime of a short-lived CLI invocation is consistent with the existing token model and does not weaken a boundary.

Non-blocking notes

• Per-URL retry storm on a failing credential: with the double-checked lock, if _remote_headers raises for a given base_url, nothing is cached, so every other thread that was blocked on the lock will re-check, re-acquire, and re-attempt the expensive (and failing) resolution serially. Harmless functionally but negates the "resolve once" benefit for the failing URL. Consider caching a sentinel/negative result or resolving failure once.
• Unbounded ThreadPoolExecutor: default max_workers scales with CPU count; with 39 records that's fine, but a cap (e.g. max_workers=min(len(records), 16)) would bound concurrent SDK/HTTP work if record counts grow.
• Cache is never invalidated; correct for a single CLI invocation as the comment argues, but note it would be stale if _host_http_json were ever reused inside a long-lived process.

Summary

The perf approach (per-URL header cache + parallel payload build) is sound and preserves output ordering via pool.map. The one real problem is that hoisting _remote_headers out of the try block silently narrows the existing error handling: auth/file-resolution failures that were previously caught and reported per-daemon now raise and, under the new parallel map, abort the whole host status command. Wrap the resolution so those failures remain recoverable and this is good to merge.

---

_{Automated review by Polly · workflow run}

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[TomeHirata]

/review

[omnigent-ci]

Polly AI Review

Blocking issues

None. The header cache is correctly populated only on success (the assignment is inside the if base_url not in cache block and inside the try, so an OSError/HTTPError from _remote_headers leaves the cache untouched and is converted to status_code=0). ThreadPoolExecutor.map preserves input order, so payloads ordering is unchanged, and exception propagation from a failed build matches the previous sequential comprehension (no regression). The double-checked locking pattern is safe here because CPython dict membership/read/write are atomic under the GIL.

Security vulnerabilities

None. Caching a resolved bearer per base URL for the lifetime of a single short-lived CLI invocation does not widen any trust boundary — the token was already resolved and sent on every call before this change. Headers are keyed on base_url, so credentials for one server are never reused for another.

Non-blocking notes

• Test docstring is slightly inaccurate. The new test's rationale ("Before the fix the cache-populating call was outside the try/except, so such a failure propagated unhandled") doesn't match reality: _remote_headers() was already invoked inside the same try block before this PR (as the headers= argument to httpx.Client(...)), so an OSError from auth resolution was already caught and returned as status_code=0. The test is still a worthwhile guard, but the narrative implies a bug that this diff isn't actually fixing.
• Cache is process-global and never invalidated. Fine for the short-lived omni host status invocation as documented, but _host_http_json is a general helper. If it were ever reused from a long-lived process, a cached token could outlive its validity. A one-line comment noting the cache is intended only for single-invocation lifetime (or scoping it to the command) would prevent future misuse.
• Unbounded default ThreadPoolExecutor. With 39 records the default worker count (min(32, cpu+4)) is acceptable since dead-process records return without network I/O, but capping max_workers to the number of records (or a small constant) would avoid spinning up more threads than useful.

Summary

A clean, well-scoped performance change. The auth-header cache and parallel payload construction are both implemented correctly and thread-safely, with failures still degrading gracefully to status_code=0 and no change to output ordering or the trust boundary. The only real nit is a misleading test docstring; the caching/parallelization logic itself is sound and ready to merge.

---

_{Automated review by Polly · workflow run}

[github-actions]

🏷️ Doc impact: no-doc-update

Internal performance change adding header caching and parallel host status fetching plus error handling in the CLI, with no change to user-facing commands, flags, or documented behavior.

_{Auto-classified on merge. Set the label manually before merging to override. · run}