From 93de39fd93c3d895c6ab37c2c40ddb4a44a5d1c1 Mon Sep 17 00:00:00 2001
From: Pieter Noordhuis <pieter.noordhuis@databricks.com>
Date: Mon, 20 Jul 2026 12:57:10 +0000
Subject: [PATCH] Fix aitools/install panic in TestCountFuzz on agent-equipped
 machines

TestCountFuzz/aitools/install panics with `telemetry logger not found`,
but only on machines with a supported coding agent installed.

Two causes:

1. aitools is hand-written but was being fuzzed. It was missing from the
   fuzz harness manualRoots blocklist added in #5102 (the test targets
   auto-generated commands, guarding against codegen regressions like #5070).
2. telemetry.Log panicked when no logger was on the context, which
   aitools install started triggering in #5862 (telemetry for the command).

CI never caught it: install's deferred logInstallEvent only runs once an
agent is detected, and the panic is unreachable via the real CLI anyway
(cmd/root.Execute always installs the logger; the fuzz harness bypasses it).

Co-authored-by: Isaac
---
 cmd/fuzz_panic_test.go        |  1 +
 libs/telemetry/context.go     |  8 ++++++++
 libs/telemetry/logger.go      | 11 ++++++++++-
 libs/telemetry/logger_test.go | 11 +++++++++++
 4 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/cmd/fuzz_panic_test.go b/cmd/fuzz_panic_test.go
index 5ac344d0341..970ca17f0a8 100644
--- a/cmd/fuzz_panic_test.go
+++ b/cmd/fuzz_panic_test.go
@@ -198,6 +198,7 @@ func isAutoGenerated(leaf leafCommand) bool {
 	// "bundle", "auth", "sync", "fs", etc. The heuristic: anything whose
 	// root isn't in this block list is auto-generated.
 	manualRoots := map[string]bool{
+		"aitools":      true,
 		"bundle":       true,
 		"auth":         true,
 		"sync":         true,
diff --git a/libs/telemetry/context.go b/libs/telemetry/context.go
index e556e462cdf..f0eb6ad60e8 100644
--- a/libs/telemetry/context.go
+++ b/libs/telemetry/context.go
@@ -23,3 +23,11 @@ func fromContext(ctx context.Context) *logger {
 
 	return v.(*logger)
 }
+
+// loggerFromContext returns the telemetry logger, or false if none was
+// installed. Unlike fromContext it does not panic, so callers on paths that may
+// run without telemetry setup can drop events instead of crashing.
+func loggerFromContext(ctx context.Context) (*logger, bool) {
+	v, ok := ctx.Value(telemetryLoggerKey).(*logger)
+	return v, ok
+}
diff --git a/libs/telemetry/logger.go b/libs/telemetry/logger.go
index cddc407c62e..c92b7470f3a 100644
--- a/libs/telemetry/logger.go
+++ b/libs/telemetry/logger.go
@@ -29,7 +29,16 @@ const (
 )
 
 func Log(ctx context.Context, event protos.DatabricksCliLog) {
-	fromContext(ctx).log(event)
+	// A missing logger means telemetry was never initialized on this context
+	// (e.g. a command invoked outside the normal cmd/root setup). Dropping the
+	// event is the right call: telemetry is best-effort and must never crash a
+	// command.
+	l, ok := loggerFromContext(ctx)
+	if !ok {
+		log.Debugf(ctx, "telemetry logger not found in the context; dropping event")
+		return
+	}
+	l.log(event)
 }
 
 type logger struct {
diff --git a/libs/telemetry/logger_test.go b/libs/telemetry/logger_test.go
index 59633403b3d..e4ff717e6d1 100644
--- a/libs/telemetry/logger_test.go
+++ b/libs/telemetry/logger_test.go
@@ -152,3 +152,14 @@ func TestTelemetryUploadMaxRetries(t *testing.T) {
 	assert.EqualError(t, err, "failed to upload telemetry logs after three attempts")
 	assert.Equal(t, 3, count)
 }
+
+func TestLogWithoutLoggerDropsEvent(t *testing.T) {
+	// A context without a telemetry logger must not panic; the event is dropped.
+	assert.NotPanics(t, func() {
+		Log(t.Context(), protos.DatabricksCliLog{
+			CliTestEvent: &protos.CliTestEvent{
+				Name: protos.DummyCliEnumValue1,
+			},
+		})
+	})
+}