Skip to content

Commit d13de92

Browse files
anton-107Isaac
andcommitted
aitools: bound best-effort telemetry auth
`aitools install` resolved auth for telemetry in PreRunE, before doing any work, so an unreachable workspace blocked the install for as long as the SDK retried. It now resolves auth in the background while the install runs and waits for it at most 3 s after the command started. Nothing is resolved when `DATABRICKS_CLI_DISABLE_TELEMETRY` is set (new `telemetry.Disabled`), and a run that logged no event does not wait. ```diff aitools install - PreRunE: tryConfigureAuth - MustAnyClient(cmd) # blocks while the SDK retries host metadata - RunE: install skills + PreRunE: startTelemetryAuth # returns at once + if telemetry.Disabled → skip + go MustAnyClient(shadow cmd) # shares cmd's flags (--profile), not its context + # SkipTokenRefresh: never refreshes an OAuth token + RunE wrapper + runE: install skills # runs while auth resolves + pending.apply(cmd) + no event logged → return + wait until start + 3 s, at most + succeeded → SetConfigUsed(cmd ctx) + failed or not finished → no config, no upload root: telemetry.Upload ``` The background resolution never refreshes an OAuth token. The process may exit before it finishes, and the server rotates the single-use refresh token, so an abandoned refresh could lose the new one and log the user out. The new `root.SkipTokenRefresh` makes `MustWorkspaceClient` and `MustAccountClient` use `auth.CredentialsWithoutRefresh`, whose databricks-cli strategy passes the new `u2m.WithoutRefresh`: the cached token is used only while it is valid, and an expired one fails auth, so nothing is uploaded. The shadow command logs through `debugOnlyHandler`, so a background warning such as `[hostmetadata] failed to fetch host metadata ...` shows only with `--debug` and does not interrupt the install's output or prompts. This relies on root's host-metadata resolver logging through the command's logger (#7390). A panic in the background resolution is recovered and logged at debug level; the install does not fail. Changelog fragment: `.nextchanges/cli/aitools-install-telemetry-auth.md`. Co-authored-by: Isaac <no-reply@databricks.com>
1 parent 0f407e7 commit d13de92

25 files changed

Lines changed: 781 additions & 16 deletions
Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
* Fix `databricks aitools install` waiting a long time on an unreachable or slow workspace while it resolves authentication for telemetry; that wait is now limited to a few seconds. ([#7468](https://github.com/databricks/cli/pull/7468))

‎acceptance/cmd/aitools/install-host-metadata/home/.claude/.gitkeep‎

Whitespace-only changes.
Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
Cloud = false
Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
2+
=== workspace unreachable: the background auth resolution prints nothing
3+
4+
>>> [CLI] aitools install --skills-only --scope=global --agents=claude-code
5+
Installing Databricks skills for Claude Code...
6+
Using skills version test-ref
7+
Fetching skills manifest...
8+
Installed 1 skill.
9+
10+
=== --debug still shows the hostmetadata line, at debug level
Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
# Isolate HOME so the real machine's agents and skills are not seen. The
2+
# committed home/.claude is what agent detection finds.
3+
sethome home
4+
5+
# The workspace host does not resolve, so the background telemetry auth
6+
# resolution fails to fetch host metadata and logs hostmetadata's warning.
7+
# The skills still come from the mock server (DATABRICKS_SKILLS_BASE_URL).
8+
export DATABRICKS_HOST=https://unreachable.test
9+
10+
title "workspace unreachable: the background auth resolution prints nothing"
11+
trace $CLI aitools install --skills-only --scope=global --agents=claude-code
12+
13+
title "--debug still shows the hostmetadata line, at debug level"
14+
# A fresh cache, so the first run's negative-cache entry does not skip the fetch.
15+
export DATABRICKS_CACHE_DIR="$PWD/debug-cache"
16+
trace $CLI aitools install --skills-only --scope=global --agents=claude-code --debug &> LOG.debug
17+
contains.py 'Debug: [hostmetadata] failed to fetch host metadata for https://unreachable.test' '!Warn: [hostmetadata]' < LOG.debug > /dev/null
Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,36 @@
1+
# Mock server replaces raw.githubusercontent.com for the skills manifest and files.
2+
Env.DATABRICKS_SKILLS_BASE_URL = "$DATABRICKS_HOST"
3+
Env.DATABRICKS_SKILLS_REF = "test-ref"
4+
# The case needs the telemetry auth resolution to run, whatever the shell sets.
5+
Env.DATABRICKS_CLI_DISABLE_TELEMETRY = ""
6+
7+
Ignore = [
8+
"home",
9+
"debug-cache",
10+
]
11+
12+
[[Server]]
13+
Pattern = "GET /test-ref/manifest.json"
14+
Response.Body = '''
15+
{
16+
"version": "2",
17+
"updated_at": "2026-01-01T00:00:00Z",
18+
"skills": {
19+
"test-stable": {
20+
"version": "1.0.0",
21+
"description": "Stable test skill",
22+
"files": ["SKILL.md"],
23+
"repo_dir": "skills"
24+
}
25+
}
26+
}
27+
'''
28+
29+
[[Server]]
30+
Pattern = "GET /test-ref/skills/test-stable/SKILL.md"
31+
Response.Body = '''---
32+
name: test-stable
33+
---
34+
35+
# Test stable skill
36+
'''

‎acceptance/cmd/aitools/telemetry/install/home/.claude/.gitkeep‎

Whitespace-only changes.
Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
Cloud = false
Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
2+
=== install resolves auth in the background: its install event is uploaded
3+
4+
>>> [CLI] aitools install --skills-only --scope=global --agents=claude-code
5+
Installing Databricks skills for Claude Code...
6+
Using skills version test-ref
7+
Fetching skills manifest...
8+
Installed 1 skill.
9+
10+
>>> print_telemetry
11+
{"command":"aitools_install","exit_code":0,"aitools_install_event":{"agents":["CLAUDE_CODE"],"scope":"GLOBAL","error_category":"TYPE_UNSPECIFIED"}}
12+
13+
=== databricks-cli auth with a valid cached token: the event is uploaded
14+
15+
>>> [CLI] aitools install --skills-only --scope=global --agents=claude-code --profile test-profile
16+
Installing Databricks skills for Claude Code...
17+
Using skills version test-ref
18+
Fetching skills manifest...
19+
Installed 1 skill.
20+
21+
>>> print_telemetry
22+
{"command":"aitools_install","exit_code":0,"aitools_install_event":{"agents":["CLAUDE_CODE"],"scope":"GLOBAL","error_category":"TYPE_UNSPECIFIED"}}
23+
24+
=== databricks-cli auth with an expired cached token: no refresh, nothing is uploaded
25+
26+
>>> [CLI] aitools install --skills-only --scope=global --agents=claude-code --profile test-profile
27+
Installing Databricks skills for Claude Code...
28+
Using skills version test-ref
29+
Fetching skills manifest...
30+
Installed 1 skill.
31+
32+
>>> print_requests.py //oidc/v1/token //telemetry-ext
33+
34+
>>> gron.py [TEST_TMP_DIR]/home/.databricks/token-cache.json
35+
json.tokens.test-profile.access_token = "expired-access-token";
Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
# Isolate HOME so the real machine's agents and skills are not seen. The
2+
# committed home/.claude is what agent detection finds.
3+
sethome home
4+
5+
title "install resolves auth in the background: its install event is uploaded"
6+
trace $CLI aitools install --skills-only --scope=global --agents=claude-code
7+
trace print_telemetry
8+
9+
# A databricks-cli profile whose token comes from the CLI's token cache.
10+
cat > "$HOME/.databrickscfg" <<ENDCFG
11+
[test-profile]
12+
host = $DATABRICKS_HOST
13+
auth_type = databricks-cli
14+
ENDCFG
15+
unset DATABRICKS_HOST DATABRICKS_TOKEN
16+
write_token_cache() {
17+
mkdir -p "$HOME/.databricks"
18+
cat > "$HOME/.databricks/token-cache.json" <<ENDCACHE
19+
{
20+
"version": 1,
21+
"tokens": {
22+
"test-profile": {
23+
"access_token": "$1",
24+
"token_type": "Bearer",
25+
"refresh_token": "test-refresh-token",
26+
"expiry": "$2"
27+
}
28+
}
29+
}
30+
ENDCACHE
31+
}
32+
33+
title "databricks-cli auth with a valid cached token: the event is uploaded"
34+
write_token_cache valid-access-token 2099-01-01T00:00:00Z
35+
trace $CLI aitools install --skills-only --scope=global --agents=claude-code --profile test-profile
36+
trace print_telemetry
37+
38+
title "databricks-cli auth with an expired cached token: no refresh, nothing is uploaded"
39+
write_token_cache expired-access-token 2000-01-01T00:00:00Z
40+
trace $CLI aitools install --skills-only --scope=global --agents=claude-code --profile test-profile
41+
trace print_requests.py //oidc/v1/token //telemetry-ext
42+
trace gron.py "$HOME/.databricks/token-cache.json" | grep access_token

0 commit comments

Comments
 (0)