Tag genie ask requests with databricks_cli source and disable viz

libs/genie/client.go

@@ -18,6 +18,11 @@ import (
 // "onechat" even though the CLI command is "genie", so the path keeps that name.
 const genieResponsesPath = "/api/2.0/data-rooms/tools/onechat/responses"
 
+// SourceDatabricksCLI is the source tag sent on every request so the OneChat
+// backend attributes CLI traffic distinctly (maps to OneChatEventSource.
+// DATABRICKS_CLI / ConversationSource.CONVERSATION_SOURCE_DATABRICKS_CLI).
+const SourceDatabricksCLI = "databricks_cli"
+
 // StreamingTimeoutSeconds is how long the SDK waits between body reads
 // before canceling the stream. The Genie agent can take minutes between SSE
 // events when executing multi-step tool calls (search, SQL, etc.), so this
@@ -29,6 +34,10 @@ const StreamingTimeoutSeconds = 600
 // Genie picks a default warehouse itself. A non-empty conversationID continues
 // that conversation instead of starting a new one.
 func BuildRequest(question, warehouseID, conversationID string) GenieRequest {
+	// The CLI runs in a terminal that cannot render visualizations, so viz is
+	// always disabled. Bound as a local because GenieRequest.EnableViz is a
+	// *bool (see the type's comment for why it can't be a plain bool).
+	enableViz := false
 	return GenieRequest{
 		Input: []InputItem{
 			{
@@ -41,6 +50,8 @@ func BuildRequest(question, warehouseID, conversationID string) GenieRequest {
 		},
 		WarehouseID:    warehouseID,
 		ConversationID: conversationID,
+		Source:         SourceDatabricksCLI,
+		EnableViz:      &enableViz,
 	}
 }
 

libs/genie/client_test.go

@@ -17,7 +17,9 @@ import (
 
 func TestBuildRequest_WireFormat(t *testing.T) {
 	// The backend expects camelCase warehouseId; this pins the wire format,
-	// not just the Go field values.
+	// not just the Go field values. source and enable_viz are always sent:
+	// the CLI tags its traffic (databricks_cli) and disables viz because a
+	// terminal cannot render charts.
 	b, err := json.Marshal(BuildRequest("What tables exist?", "abc123", ""))
 	require.NoError(t, err)
 	assert.JSONEq(t, `{
@@ -26,10 +28,22 @@ func TestBuildRequest_WireFormat(t *testing.T) {
 			"role": "user",
 			"content": [{"type": "input_text", "text": "What tables exist?"}]
 		}],
-		"warehouseId": "abc123"
+		"warehouseId": "abc123",
+		"source": "databricks_cli",
+		"enable_viz": false
 	}`, string(b))
 }
 
+func TestBuildRequest_TagsSourceAndDisablesViz(t *testing.T) {
+	// enable_viz:false must always be on the wire (a plain false would be
+	// dropped by omitempty, letting the backend default viz on), and source
+	// tags CLI traffic for backend attribution.
+	b, err := json.Marshal(BuildRequest("q", "", ""))
+	require.NoError(t, err)
+	assert.Contains(t, string(b), `"source":"databricks_cli"`)
+	assert.Contains(t, string(b), `"enable_viz":false`)
+}
+
 func TestBuildRequest_OmitsEmptyWarehouse(t *testing.T) {
 	b, err := json.Marshal(BuildRequest("q", "", ""))
 	require.NoError(t, err)

libs/genie/types.go

@@ -7,6 +7,15 @@ type GenieRequest struct {
 	// Tag is camelCase (conversationId), unlike the response's snake_case
 	// conversation_id; the server does not read the snake_case form on input.
 	ConversationID string `json:"conversationId,omitempty"`
+	// Source identifies this client to the OneChat backend so CLI traffic is
+	// attributed distinctly from webapp/Slack/MCP. Always SourceDatabricksCLI;
+	// the server matches it case-insensitively against a fixed allowlist.
+	Source string `json:"source,omitempty"`
+	// EnableViz is a pointer so it always serializes (a plain false would be
+	// dropped by omitempty). The CLI sends false: a terminal cannot render the
+	// chart/plot items, so we tell the agent to skip registering visualization
+	// tools. Nil would let the backend default (viz on) apply.
+	EnableViz *bool `json:"enable_viz,omitempty"`
 }
 
 // InputItem is a message in the input array.