Fix aitools/install panic in TestCountFuzz on agent-equipped machines

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,

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
+}

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 {

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,
+			},
+		})
+	})
+}