diff --git a/.nextchanges/cli/experimental-genie-cli.md b/.nextchanges/cli/experimental-genie-cli.md new file mode 100644 index 00000000000..5449e2936bf --- /dev/null +++ b/.nextchanges/cli/experimental-genie-cli.md @@ -0,0 +1 @@ +Added the experimental `databricks experimental genie-cli` command, which starts an interactive AI coding agent preconfigured for your Databricks workspace: it routes model requests through the Databricks AI Gateway, preloads the Databricks skills plugin, registers Databricks MCP tools, and primes the agent with a Databricks-aware system prompt scoped to the resolved workspace host and profile. Use `--harness` to choose the coding agent (default `codex`; `opencode` is also primed with the Databricks system prompt). diff --git a/cmd/experimental/experimental.go b/cmd/experimental/experimental.go index d87c893abc5..d49399c2b07 100644 --- a/cmd/experimental/experimental.go +++ b/cmd/experimental/experimental.go @@ -4,6 +4,7 @@ import ( aircmd "github.com/databricks/cli/experimental/air/cmd" aitoolscmd "github.com/databricks/cli/experimental/aitools/cmd" geniecmd "github.com/databricks/cli/experimental/genie/cmd" + genieclicmd "github.com/databricks/cli/experimental/geniecli/cmd" postgrescmd "github.com/databricks/cli/experimental/postgres/cmd" "github.com/spf13/cobra" ) @@ -26,6 +27,7 @@ development. They may change or be removed in future versions without notice.`, cmd.AddCommand(aircmd.New()) cmd.AddCommand(aitoolscmd.NewAitoolsCmd()) cmd.AddCommand(geniecmd.NewGenieCmd()) + cmd.AddCommand(genieclicmd.New()) cmd.AddCommand(postgrescmd.New()) cmd.AddCommand(newWorkspaceOpenCommand()) diff --git a/experimental/geniecli/cmd/geniecli.go b/experimental/geniecli/cmd/geniecli.go new file mode 100644 index 00000000000..f9fe6e95a2c --- /dev/null +++ b/experimental/geniecli/cmd/geniecli.go @@ -0,0 +1,97 @@ +package genieclicmd + +import ( + "github.com/databricks/cli/cmd/root" + "github.com/databricks/cli/libs/cmdctx" + "github.com/databricks/cli/libs/cmdio" + "github.com/spf13/cobra" +) + +// New returns the `genie-cli` command: an out-of-the-box interactive AI coding +// agent preconfigured for Databricks. It routes model traffic through the +// workspace's Databricks AI Gateway (no separate API keys) and preloads the +// Databricks skills plugin so the agent works effectively against Databricks. +// +// The agent runtime is an implementation detail intentionally kept out of the +// user-facing surface: the user only ever interacts with `genie-cli`. +func New() *cobra.Command { + var noSystemPrompt bool + var harness string + + cmd := &cobra.Command{ + Use: "genie-cli [-- AGENT_ARGS...]", + Hidden: true, + Short: "Start an interactive AI coding agent configured for Databricks", + Long: `Start an interactive AI coding agent that is preconfigured for your +Databricks workspace. + +The agent authenticates through your Databricks workspace and routes model +requests through the Databricks AI Gateway, so no separate model API keys are +needed. Databricks coding skills are installed and kept up to date on every +launch, Databricks MCP tools are registered, and the agent is primed with a +system prompt so it works effectively against Databricks out of the box. + +Use --harness to pick which coding agent runs (default codex; opencode is also +primed with the Databricks system prompt). Any arguments after "--" are +forwarded to the underlying agent, e.g.: + + databricks experimental genie-cli -- --full-auto + databricks experimental genie-cli --harness opencode`, + // The agent takes over the terminal, so no bundle is loaded. + PreRunE: func(cmd *cobra.Command, args []string) error { + cmd.SetContext(root.SkipLoadBundle(cmd.Context())) + return root.MustWorkspaceClient(cmd, args) + }, + RunE: func(cmd *cobra.Command, args []string) error { + ctx := cmd.Context() + // Use the identity from the authenticated workspace client, which + // reflects the profile/host the invocation resolved (--profile, + // DATABRICKS_CONFIG_PROFILE, the default profile, or env-var auth). + cfg := cmdctx.WorkspaceClient(ctx).Config + + // Ensure the launcher runtime is present, installing it if needed. + ucodePath, err := ensureUcode(ctx) + if err != nil { + return err + } + + // Configure the harness against this workspace. This also installs the + // agent binary and the Databricks CLI the runtime depends on, so it + // must run before the plugin step (which drives the agent's own CLI). + if err := configureAgent(ctx, ucodePath, harness, cfg); err != nil { + return err + } + + // Install the Databricks skills plugin for the harness, then refresh + // it. No-op for harnesses without a headless plugin (e.g. opencode). + // Refresh failures (offline, GitHub unreachable) are non-fatal: a stale + // plugin still works, so we warn and launch anyway. + if err := ensureDatabricksPlugin(ctx, harness); err != nil { + return err + } + refreshDatabricksPlugin(ctx) + + // Register Databricks MCP tools for the session (best-effort). + registerMCP(ctx, ucodePath) + + // Prime the agent to operate as a Databricks CLI assistant, injected + // per-session only (no user config is written). Carries the resolved + // host and profile so the agent acts against this workspace. + var inj injection + if !noSystemPrompt { + inj, err = buildInjection(ctx, harness, cfg.Host, cfg.Profile) + if err != nil { + return err + } + } + + cmdio.LogString(ctx, "Starting the Databricks agent...") + return launchAgent(ctx, ucodePath, harness, inj, args) + }, + } + + cmd.Flags().StringVar(&harness, "harness", defaultHarness, "Coding agent to run (e.g. codex, opencode)") + cmd.Flags().BoolVar(&noSystemPrompt, "no-system-prompt", false, "Do not prime the agent with the default Databricks system prompt") + + return cmd +} diff --git a/experimental/geniecli/cmd/harness.go b/experimental/geniecli/cmd/harness.go new file mode 100644 index 00000000000..5941e93fd6d --- /dev/null +++ b/experimental/geniecli/cmd/harness.go @@ -0,0 +1,92 @@ +package genieclicmd + +import ( + "context" + "encoding/json" + "fmt" + "os" + "path/filepath" + + "github.com/databricks/cli/libs/env" +) + +// defaultHarness is the coding agent genie-cli drives unless the user selects +// another with --harness. It matches an agent name ucode understands and (for +// codex) the agent registry name in libs/aitools/agents. +const defaultHarness = "codex" + +// harnessCodex and harnessOpenCode are the harnesses genie-cli primes with the +// Databricks system prompt. Any other ucode agent name is still accepted and +// launched, just without prompt injection (see buildInjection). +const ( + harnessCodex = "codex" + harnessOpenCode = "opencode" +) + +// injection describes how the Databricks system prompt is delivered to a +// harness for one session. Delivery differs per agent: Codex takes an inline +// `-c` override, OpenCode reads an instructions file referenced through an +// environment variable. Fields are empty when a harness gets no prompt. +type injection struct { + // forwardArgs are appended after the agent's "--" separator (Codex). + forwardArgs []string + // env are extra KEY=value entries for the launched process (OpenCode). + env []string +} + +// buildInjection renders the system prompt for the harness and returns how to +// deliver it for this session, without writing to the user's agent config. +// +// - codex: inline `-c developer_instructions=` (additive developer-role +// message). +// - opencode: the prompt is written to a genie-cli-managed file and referenced +// via OPENCODE_CONFIG_CONTENT={"instructions":[""]}, whose entries +// OpenCode appends to the model context. Uses an env var (not a flag) +// because OpenCode has no inline system-prompt flag. +// - any other harness: no injection (returns a zero injection, nil error). +func buildInjection(ctx context.Context, harness, host, profile string) (injection, error) { + prompt := buildSystemPrompt(host, profile) + + switch harness { + case harnessCodex: + return injection{forwardArgs: []string{"-c", "developer_instructions=" + tomlQuote(prompt)}}, nil + case harnessOpenCode: + path, err := writeOpenCodePromptFile(ctx, prompt) + if err != nil { + return injection{}, err + } + content, err := json.Marshal(map[string]any{"instructions": []string{path}}) + if err != nil { + return injection{}, fmt.Errorf("failed to encode OpenCode config: %w", err) + } + return injection{env: []string{"OPENCODE_CONFIG_CONTENT=" + string(content)}}, nil + default: + return injection{}, nil + } +} + +// openCodePromptPath returns the genie-cli-managed file that holds the OpenCode +// instructions. It lives under the user's home in a stable, overwritten +// location so there is no temp file to clean up after the exec handoff. +func openCodePromptPath(ctx context.Context) (string, error) { + home, err := env.UserHomeDir(ctx) + if err != nil { + return "", err + } + return filepath.Join(home, ".databricks", "genie-cli", "opencode-instructions.md"), nil +} + +// writeOpenCodePromptFile writes the prompt to the managed path and returns it. +func writeOpenCodePromptFile(ctx context.Context, prompt string) (string, error) { + path, err := openCodePromptPath(ctx) + if err != nil { + return "", err + } + if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil { + return "", fmt.Errorf("failed to create genie-cli directory: %w", err) + } + if err := os.WriteFile(path, []byte(prompt), 0o644); err != nil { + return "", fmt.Errorf("failed to write the OpenCode system prompt: %w", err) + } + return path, nil +} diff --git a/experimental/geniecli/cmd/harness_test.go b/experimental/geniecli/cmd/harness_test.go new file mode 100644 index 00000000000..dac303afa4f --- /dev/null +++ b/experimental/geniecli/cmd/harness_test.go @@ -0,0 +1,63 @@ +package genieclicmd + +import ( + "encoding/json" + "os" + "strings" + "testing" + + "github.com/databricks/cli/libs/env" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestBuildInjectionCodex(t *testing.T) { + inj, err := buildInjection(t.Context(), harnessCodex, "https://myworkspace.databricks.com", "prod") + require.NoError(t, err) + assert.Empty(t, inj.env) + require.Len(t, inj.forwardArgs, 2) + assert.Equal(t, "-c", inj.forwardArgs[0]) + + // The prompt is delivered as a single quoted TOML basic string. + value := inj.forwardArgs[1] + assert.True(t, strings.HasPrefix(value, `developer_instructions="`)) + assert.True(t, strings.HasSuffix(value, `"`)) + body := strings.TrimPrefix(value, "developer_instructions=") + assert.NotContains(t, body[1:len(body)-1], "\n") + assert.Contains(t, body, `\n`) + assert.Contains(t, value, "myworkspace.databricks.com") +} + +func TestBuildInjectionOpenCode(t *testing.T) { + home := t.TempDir() + ctx := env.Set(t.Context(), "HOME", home) + ctx = env.Set(ctx, "USERPROFILE", home) + + inj, err := buildInjection(ctx, harnessOpenCode, "https://myworkspace.databricks.com", "prod") + require.NoError(t, err) + assert.Empty(t, inj.forwardArgs) + require.Len(t, inj.env, 1) + + // The env var carries JSON pointing at the managed prompt file, which must + // exist and contain the rendered prompt. + content := strings.TrimPrefix(inj.env[0], "OPENCODE_CONFIG_CONTENT=") + assert.NotEqual(t, inj.env[0], content) // prefix was present + + var cfg struct { + Instructions []string `json:"instructions"` + } + require.NoError(t, json.Unmarshal([]byte(content), &cfg)) + require.Len(t, cfg.Instructions, 1) + + data, err := os.ReadFile(cfg.Instructions[0]) + require.NoError(t, err) + assert.Contains(t, string(data), "myworkspace.databricks.com") + assert.Contains(t, string(data), "databricks experimental genie ask") +} + +func TestBuildInjectionUnknownHarnessIsEmpty(t *testing.T) { + inj, err := buildInjection(t.Context(), "gemini", "https://myworkspace.databricks.com", "prod") + require.NoError(t, err) + assert.Empty(t, inj.forwardArgs) + assert.Empty(t, inj.env) +} diff --git a/experimental/geniecli/cmd/plugin.go b/experimental/geniecli/cmd/plugin.go new file mode 100644 index 00000000000..24feb6158bb --- /dev/null +++ b/experimental/geniecli/cmd/plugin.go @@ -0,0 +1,95 @@ +package genieclicmd + +import ( + "context" + "fmt" + + "github.com/databricks/cli/libs/aitools/agents" + "github.com/databricks/cli/libs/aitools/installer" + "github.com/databricks/cli/libs/log" +) + +// pluginScope is the CLI install scope for the Databricks plugin. genie-cli is +// a global, workspace-wide agent, so it installs at global scope. +const pluginScope = installer.ScopeGlobal + +// nativeScope is the agent-native plugin scope Codex installs into. Codex +// ignores --scope, but the installer records it, so "user" matches how a +// global CLI install maps to an agent's user-level scope. +const nativeScope = "user" + +// ensureDatabricksPlugin installs the Databricks skills plugin for the harness +// if the CLI has not already recorded installing it. The harness binary must +// already be on PATH here (the configure step installs it) because the plugin +// installs through the agent's own CLI. +// +// It is a no-op for a harness that has no headless plugin (e.g. OpenCode, which +// is skills-only) or that the aitools registry doesn't know: those harnesses +// rely on the injected system prompt for Databricks priming instead. Codex +// exposes no plugin manifest the CLI can read back, so prior installs are +// detected from the CLI's own install state rather than from the agent (unlike +// Claude Code). +func ensureDatabricksPlugin(ctx context.Context, harness string) error { + agent := agents.ByName(harness) + if agent == nil || agent.Plugin == nil { + return nil + } + + installed, err := pluginRecorded(ctx, agent.Name) + if err != nil { + return err + } + if installed { + return nil + } + + ref, _, err := installer.GetSkillsRef(ctx) + if err != nil { + return fmt.Errorf("failed to resolve the Databricks skills version: %w", err) + } + + rec, err := installer.InstallPluginForAgent(ctx, agent, nativeScope, ref) + if err != nil { + return fmt.Errorf("failed to install the Databricks skills plugin: %w", err) + } + + // Record the install so the refresh path (and `databricks aitools`) can act + // on exactly where we installed. + records := map[string]installer.PluginRecord{agent.Name: rec} + if err := installer.RecordPluginInstalls(ctx, pluginScope, records, ref); err != nil { + return fmt.Errorf("failed to record the Databricks skills plugin install: %w", err) + } + return nil +} + +// pluginRecorded reports whether the CLI's install state already records the +// Databricks plugin for the named agent at the global scope. +func pluginRecorded(ctx context.Context, name string) (bool, error) { + dir, err := installer.GlobalSkillsDir(ctx) + if err != nil { + return false, err + } + state, err := installer.LoadState(dir) + if err != nil { + return false, fmt.Errorf("failed to read the Databricks skills install state: %w", err) + } + if state == nil { + return false, nil + } + _, ok := state.Plugins[name] + return ok, nil +} + +// refreshDatabricksPlugin updates the Databricks skills plugin to the latest +// release before the agent starts. Failures are non-fatal: a stale plugin still +// works, so an offline or unreachable-GitHub run warns and launches anyway. +func refreshDatabricksPlugin(ctx context.Context) { + ref, _, err := installer.GetSkillsRef(ctx) + if err != nil { + log.Warnf(ctx, "Skipping Databricks skills update: %v", err) + return + } + if _, err := installer.UpdateInstalledPlugins(ctx, pluginScope, ref); err != nil { + log.Warnf(ctx, "Skipping Databricks skills update: %v", err) + } +} diff --git a/experimental/geniecli/cmd/plugin_test.go b/experimental/geniecli/cmd/plugin_test.go new file mode 100644 index 00000000000..78d21aee555 --- /dev/null +++ b/experimental/geniecli/cmd/plugin_test.go @@ -0,0 +1,52 @@ +package genieclicmd + +import ( + "testing" + + "github.com/databricks/cli/libs/aitools/installer" + "github.com/databricks/cli/libs/env" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestPluginRecorded(t *testing.T) { + home := t.TempDir() + ctx := env.Set(t.Context(), "HOME", home) + // On Windows env.UserHomeDir reads USERPROFILE. + ctx = env.Set(ctx, "USERPROFILE", home) + + // No state yet. + got, err := pluginRecorded(ctx, "codex") + require.NoError(t, err) + assert.False(t, got) + + dir, err := installer.GlobalSkillsDir(ctx) + require.NoError(t, err) + + // State without the agent's plugin. + require.NoError(t, installer.SaveState(dir, &installer.InstallState{ + Plugins: map[string]installer.PluginRecord{"claude-code": {Plugin: "databricks"}}, + })) + got, err = pluginRecorded(ctx, "codex") + require.NoError(t, err) + assert.False(t, got) + + // State with the agent's plugin recorded. + require.NoError(t, installer.SaveState(dir, &installer.InstallState{ + Plugins: map[string]installer.PluginRecord{"codex": {Plugin: "databricks"}}, + })) + got, err = pluginRecorded(ctx, "codex") + require.NoError(t, err) + assert.True(t, got) +} + +func TestEnsureDatabricksPluginSkipsNonPluginHarness(t *testing.T) { + home := t.TempDir() + ctx := env.Set(t.Context(), "HOME", home) + ctx = env.Set(ctx, "USERPROFILE", home) + + // opencode is skills-only (no headless plugin) and gemini is unknown to the + // aitools registry; both must be no-ops rather than errors or installs. + require.NoError(t, ensureDatabricksPlugin(ctx, "opencode")) + require.NoError(t, ensureDatabricksPlugin(ctx, "gemini")) +} diff --git a/experimental/geniecli/cmd/prompt.go b/experimental/geniecli/cmd/prompt.go new file mode 100644 index 00000000000..348e26b00b5 --- /dev/null +++ b/experimental/geniecli/cmd/prompt.go @@ -0,0 +1,107 @@ +package genieclicmd + +import ( + "fmt" + "strings" +) + +// systemPromptTemplate is the default developer instruction genie-cli injects +// into the agent so it operates as a Databricks-native CLI assistant out of the +// box. It is adapted for a terminal coding agent (Codex) that has the real +// `databricks` CLI, the Databricks skills, and the Databricks MCP servers +// available, rather than the web/notebook harness the original Genie Code agent +// targets. +// +// The %s placeholders are the resolved workspace host and CLI profile so the +// agent acts against the same workspace the genie-cli command authenticated as. +const systemPromptTemplate = `You are the Databricks CLI assistant: a coding agent that helps users get work done on Databricks directly from the terminal. + +# Workspace +You are operating against Databricks workspace %s. +%s +Use this workspace and profile for every Databricks operation. Do not switch workspaces or profiles unless the user explicitly asks. + +# Primary objective +Analyze the user's request and decide where the work belongs. The single most +important decision: **does the task depend on the user's local filesystem +(their code, local data files, a repo, a bundle, notebooks on disk)?** + +- **No local dependency** (a question or task purely about workspace data — metrics, analytics, data science, "what data do I have", exploring tables, computing a result): **do the whole task through hosted Genie** with ` + "`genie ask`" + ` (see below). Do NOT just use Genie to find tables and then hand-write and run SQL locally — send Genie the actual task and let it produce the answer. +- **Depends on the local filesystem** (editing code, a DAB/bundle, files in the repo, local data, or generating artifacts on disk): do it locally with your own tools and the ` + "`databricks`" + ` CLI, pulling data in via Genie when you need workspace facts. +- **Workspace operations** (jobs, pipelines, clusters, apps, Lakebase, Unity Catalog management): use the Databricks CLI. +- For general help, gather what you need before answering; do not run destructive or expensive operations just to explain something. + +# Answering data questions: do the work in hosted Genie +Hosted Databricks Genie is grounded in this workspace's data. For any data +discovery, analytics, or data-science task that does not depend on the user's +local files, send the whole task to Genie rather than writing and running SQL +yourself: + + databricks experimental genie ask "" + +Prefer Genie for the complete task, not just table discovery, because: +- It produces a better, workspace-grounded answer than locally hand-written SQL. +- The result is easily shareable — it lives in a Genie conversation the user can open, re-run, and share. +- It still returns the SQL it ran and the tables it used, so you can read those back and explore further locally when the task genuinely needs it. + +Guidance: +- Send the real task in natural language (e.g. ` + "`genie ask \"what is monthly active revenue by region for the last 6 months?\"`" + `), not a narrow "find me a table" query. Let Genie find the tables and write the SQL. +- Use ` + "`--include-sql`" + ` to see the queries Genie ran, and ` + "`--output json`" + ` when you need to parse the answer, SQL, or table references programmatically for follow-up work. +- Continue a line of questioning in one conversation with ` + "`-s `" + `, e.g. ` + "`genie ask -s sales \"break that down by region\"`" + `. +- Add ` + "`--warehouse-id `" + ` only if the user names a specific warehouse. +- Only fall back to hand-writing and running SQL yourself when Genie cannot answer, the user explicitly asks for hand-written SQL, or the task is about authoring/editing a query or pipeline on disk rather than getting an answer. + +# Other tools available to you +- **The ` + "`databricks`" + ` CLI** is installed and pre-authenticated for this workspace. Prefer it for workspace resource management (jobs, pipelines, clusters, apps, secrets, Unity Catalog, Lakebase, filesystem) and for hosted Genie (above). Always target the resolved profile above. Run ` + "`databricks --help`" + ` to discover commands rather than guessing. +- **Databricks skills** are installed for this session under the Databricks plugin. Consult the relevant skill before a non-trivial Databricks task (for example the DABs, jobs, pipelines, or SQL skills) so you follow current best practices. +- **Databricks MCP servers** may be registered (for example Unity Catalog functions and Genie spaces). When an MCP tool fits the task, prefer it over shelling out. + +# Core principles +1. **Discover before concluding.** Assume the data or resource exists in the workspace until you have actively looked and found nothing. Never claim you "don't have access" or that something is "unavailable" without first asking hosted Genie or searching with the CLI or MCP tools. +2. **Act with sensible defaults.** When the user directs a concrete operation, attempt it immediately with reasonable assumptions. Only ask a clarifying question if the attempt fails or a critical detail is genuinely ambiguous. +3. **Read before modifying.** Always read a resource's current definition before editing it. +4. **Plan multi-step work.** For tasks with dependencies (investigate -> change -> verify), track steps explicitly and finish them before ending your turn. Skip the ceremony for simple single-action fixes. +5. **Stay simple.** Solve what was asked. Do not add extra features or speculative complexity. +6. **Be persistent.** Keep working autonomously until the request is resolved. If one approach is blocked, try another or report precisely what is blocking you. +7. **Be safe.** When generating SQL or code, avoid injection and other vulnerabilities. Confirm before running anything destructive (dropping data, deleting resources, force-deploying over existing state). + +# Response format +- Keep answers concise: a few short paragraphs. +- Incorporate tool results into your analysis; do not restate raw command output or code verbatim. +- Use standard Markdown tables for structured or comparative data. +` + +// buildSystemPrompt renders the default developer instruction for the resolved +// workspace. host is the workspace URL; profile is the resolved CLI profile +// name, which may be empty when auth came from environment variables rather +// than a named profile. +func buildSystemPrompt(host, profile string) string { + hostText := host + if hostText == "" { + hostText = "(the workspace resolved by the Databricks CLI)" + } + + var profileLine string + if profile != "" { + profileLine = fmt.Sprintf("Authenticate with the Databricks CLI profile %q; pass `--profile %s` to every `databricks` command.", profile, profile) + } else { + // No named profile: auth resolved from environment/host, which the CLI + // already picks up, so the agent must not invent a --profile flag. + profileLine = "Authentication is resolved from the environment; run `databricks` commands without a --profile flag." + } + + return fmt.Sprintf(systemPromptTemplate, hostText, profileLine) +} + +// tomlQuote renders s as a TOML basic string: wrapped in double quotes with +// backslashes, double quotes, and newlines escaped. +func tomlQuote(s string) string { + r := strings.NewReplacer( + `\`, `\\`, + `"`, `\"`, + "\n", `\n`, + "\t", `\t`, + "\r", `\r`, + ) + return `"` + r.Replace(s) + `"` +} diff --git a/experimental/geniecli/cmd/prompt_test.go b/experimental/geniecli/cmd/prompt_test.go new file mode 100644 index 00000000000..fefdbc03880 --- /dev/null +++ b/experimental/geniecli/cmd/prompt_test.go @@ -0,0 +1,30 @@ +package genieclicmd + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestBuildSystemPromptWithProfile(t *testing.T) { + got := buildSystemPrompt("https://myworkspace.databricks.com", "prod") + assert.Contains(t, got, "https://myworkspace.databricks.com") + assert.Contains(t, got, `--profile prod`) + assert.Contains(t, got, "Databricks CLI assistant") + // Data questions should route to hosted Genie. + assert.Contains(t, got, "databricks experimental genie ask") +} + +func TestBuildSystemPromptWithoutProfile(t *testing.T) { + got := buildSystemPrompt("https://myworkspace.databricks.com", "") + assert.Contains(t, got, "https://myworkspace.databricks.com") + // With no named profile the agent is told auth comes from the environment + // and must not be told to pass a specific `--profile `. + assert.Contains(t, got, "resolved from the environment") + assert.NotContains(t, got, `--profile "`) +} + +func TestTomlQuoteEscapes(t *testing.T) { + got := tomlQuote("a\"b\nc\\d") + assert.Equal(t, `"a\"b\nc\\d"`, got) +} diff --git a/experimental/geniecli/cmd/ucode.go b/experimental/geniecli/cmd/ucode.go new file mode 100644 index 00000000000..642caa03fa5 --- /dev/null +++ b/experimental/geniecli/cmd/ucode.go @@ -0,0 +1,151 @@ +package genieclicmd + +import ( + "context" + "errors" + "fmt" + "os" + "os/exec" + + "github.com/databricks/cli/libs/auth" + "github.com/databricks/cli/libs/env" + "github.com/databricks/cli/libs/execv" + "github.com/databricks/cli/libs/log" + "github.com/databricks/cli/libs/process" + "github.com/databricks/databricks-sdk-go/config" +) + +const ( + // ucodeBin is the launcher that wires a coding agent to the Databricks AI + // Gateway. It, the agent binary, and the underlying model routing are + // implementation details hidden behind the genie-cli command. + ucodeBin = "ucode" + // uvBin installs ucode; it is the only supported installer for the launcher + // (a Python tool published as a uv tool). See ucodeInstallSpec. + uvBin = "uv" + // ucodeInstallSpec is the source uv installs the launcher from. + // https://github.com/databricks/ucode + ucodeInstallSpec = "git+https://github.com/databricks/ucode" + + // mcpLocation is the Unity Catalog schema whose MCP services are registered + // for the agent. system.ai hosts the built-in Databricks MCP services. + mcpLocation = "system.ai" +) + +// lookPath is a package-level seam so tests can resolve binaries without +// depending on the host PATH. +var lookPath = exec.LookPath + +// launch replaces the current process with the given command; a package-level +// seam so tests can assert the launch argv without exec'ing. +var launch = execv.Execv + +// ensureUcode makes sure the launcher is installed, installing it through uv if +// it is missing. It returns the absolute path to the launcher binary. +func ensureUcode(ctx context.Context) (string, error) { + if path, err := lookPath(ucodeBin); err == nil { + return path, nil + } + + // The launcher is only distributed as a uv tool, so uv is a hard prerequisite + // rather than something we can work around. + if _, err := lookPath(uvBin); err != nil { + return "", fmt.Errorf("genie-cli requires uv to install its runtime, but uv was not found on PATH\n"+ + "Install uv (https://docs.astral.sh/uv/getting-started/installation) and re-run, "+ + "or install the runtime manually with: uv tool install %s", ucodeInstallSpec) + } + + log.Infof(ctx, "Installing the genie-cli runtime with uv") + if err := process.Forwarded(ctx, []string{uvBin, "tool", "install", ucodeInstallSpec}, os.Stdin, os.Stdout, os.Stderr); err != nil { + return "", fmt.Errorf("failed to install the genie-cli runtime: %w", err) + } + + // uv installs into its tool bin dir (e.g. ~/.local/bin), which may not be on + // the current process's PATH even after a successful install. + path, err := lookPath(ucodeBin) + if err != nil { + return "", fmt.Errorf("the genie-cli runtime was installed but %q is not on PATH\n"+ + "Add uv's tool bin directory to PATH (run: uv tool update-shell) and re-run", ucodeBin) + } + return path, nil +} + +// configureArgs builds the launcher's non-interactive configure argv for the +// chosen harness and resolved workspace config. The launcher needs an explicit +// workspace identity or it prompts interactively, which would break the +// non-interactive setup. +func configureArgs(ucodePath, harness string, cfg *config.Config) ([]string, error) { + args := []string{ucodePath, "configure", "--agents", harness, "--skip-validate", "--skip-upgrade"} + + switch { + case cfg.Profile != "": + args = append(args, "--profiles", cfg.Profile) + // --use-pat reuses the profile's stored token with no interactive login; + // it is only valid for PAT profiles and requires --profiles. + if cfg.AuthType == auth.AuthTypePat { + args = append(args, "--use-pat") + } + case cfg.Host != "": + args = append(args, "--workspaces", cfg.Host) + default: + return nil, errors.New("no Databricks profile or host resolved; run `databricks auth login` first") + } + return args, nil +} + +// configureAgent points the harness at this workspace's AI Gateway. This also +// installs the agent binary and the Databricks CLI the launcher relies on. +func configureAgent(ctx context.Context, ucodePath, harness string, cfg *config.Config) error { + args, err := configureArgs(ucodePath, harness, cfg) + if err != nil { + return err + } + if err := process.Forwarded(ctx, args, os.Stdin, os.Stdout, os.Stderr); err != nil { + return fmt.Errorf("failed to configure the Databricks agent: %w", err) + } + return nil +} + +// registerMCP registers the Databricks MCP services from the configured Unity +// Catalog location for the agent, so the session has Databricks MCP tools. It +// is best-effort: a workspace without MCP services in that schema (or an +// unreachable one) should not block the agent, so failures are logged and +// swallowed, matching the plugin-refresh policy. +func registerMCP(ctx context.Context, ucodePath string) { + args := []string{ucodePath, "configure", "mcp", "--location", mcpLocation} + if err := process.Forwarded(ctx, args, os.Stdin, os.Stdout, os.Stderr); err != nil { + log.Warnf(ctx, "Skipping Databricks MCP setup: %v", err) + } +} + +// launchAgent replaces the current process with the harness session. +// --skip-preflight trusts the configure step we just ran and skips a redundant +// auth + gateway re-validation. The system prompt is delivered per the harness's +// injection (forwarded args and/or extra env); userArgs are forwarded to the +// underlying agent after ours. +func launchAgent(ctx context.Context, ucodePath, harness string, inj injection, userArgs []string) error { + // Everything after "--" is forwarded verbatim to the underlying agent. + forwarded := append([]string{}, inj.forwardArgs...) + forwarded = append(forwarded, userArgs...) + + args := []string{ucodePath, harness, "--skip-preflight"} + if len(forwarded) > 0 { + args = append(args, "--") + args = append(args, forwarded...) + } + return launch(execv.Options{ + Args: args, + Env: append(envSlice(ctx), inj.env...), + }) +} + +// envSlice returns the process environment as a "KEY=value" slice, pulled +// through libs/env so tests can override it per context. +func envSlice(ctx context.Context) []string { + all := env.All(ctx) + out := make([]string, 0, len(all)) + for k, v := range all { + out = append(out, k+"="+v) + } + return out +} diff --git a/experimental/geniecli/cmd/ucode_test.go b/experimental/geniecli/cmd/ucode_test.go new file mode 100644 index 00000000000..fc8e163332c --- /dev/null +++ b/experimental/geniecli/cmd/ucode_test.go @@ -0,0 +1,170 @@ +package genieclicmd + +import ( + "testing" + + "github.com/databricks/cli/libs/auth" + "github.com/databricks/cli/libs/execv" + "github.com/databricks/cli/libs/process" + "github.com/databricks/databricks-sdk-go/config" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestConfigureArgs(t *testing.T) { + tests := []struct { + name string + harness string + cfg *config.Config + want []string + }{ + { + name: "pat profile uses --use-pat", + harness: "codex", + cfg: &config.Config{Profile: "prod", AuthType: auth.AuthTypePat}, + want: []string{"ucode", "configure", "--agents", "codex", "--skip-validate", "--skip-upgrade", "--profiles", "prod", "--use-pat"}, + }, + { + name: "oauth profile omits --use-pat", + harness: "codex", + cfg: &config.Config{Profile: "prod", AuthType: "databricks-cli"}, + want: []string{"ucode", "configure", "--agents", "codex", "--skip-validate", "--skip-upgrade", "--profiles", "prod"}, + }, + { + name: "host without profile falls back to --workspaces", + harness: "codex", + cfg: &config.Config{Host: "https://myworkspace.databricks.com"}, + want: []string{"ucode", "configure", "--agents", "codex", "--skip-validate", "--skip-upgrade", "--workspaces", "https://myworkspace.databricks.com"}, + }, + { + name: "harness is forwarded to --agents", + harness: "opencode", + cfg: &config.Config{Profile: "prod", AuthType: "databricks-cli"}, + want: []string{"ucode", "configure", "--agents", "opencode", "--skip-validate", "--skip-upgrade", "--profiles", "prod"}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, err := configureArgs("ucode", tt.harness, tt.cfg) + require.NoError(t, err) + assert.Equal(t, tt.want, got) + }) + } +} + +func TestConfigureArgsNoIdentity(t *testing.T) { + _, err := configureArgs("ucode", "codex", &config.Config{}) + require.Error(t, err) + assert.Contains(t, err.Error(), "databricks auth login") +} + +func TestEnsureUcodeAlreadyOnPath(t *testing.T) { + restore := lookPath + t.Cleanup(func() { lookPath = restore }) + lookPath = func(file string) (string, error) { + assert.Equal(t, ucodeBin, file) + return "/usr/local/bin/ucode", nil + } + + path, err := ensureUcode(t.Context()) + require.NoError(t, err) + assert.Equal(t, "/usr/local/bin/ucode", path) +} + +func TestEnsureUcodeMissingUvErrors(t *testing.T) { + restore := lookPath + t.Cleanup(func() { lookPath = restore }) + lookPath = func(file string) (string, error) { + return "", assert.AnError + } + + _, err := ensureUcode(t.Context()) + require.Error(t, err) + assert.Contains(t, err.Error(), "uv") +} + +func TestEnsureUcodeInstallsWhenMissing(t *testing.T) { + restore := lookPath + t.Cleanup(func() { lookPath = restore }) + // ucode absent on the first probe, uv present, ucode present after install. + seen := map[string]int{} + lookPath = func(file string) (string, error) { + seen[file]++ + switch file { + case uvBin: + return "/usr/local/bin/uv", nil + case ucodeBin: + if seen[ucodeBin] == 1 { + return "", assert.AnError + } + return "/home/user/.local/bin/ucode", nil + } + return "", assert.AnError + } + + ctx, stub := process.WithStub(t.Context()) + stub.WithStdoutFor("uv tool install", "installed") + + path, err := ensureUcode(ctx) + require.NoError(t, err) + assert.Equal(t, "/home/user/.local/bin/ucode", path) + assert.Len(t, stub.Commands(), 1) + assert.Contains(t, stub.Commands()[0], "uv tool install") +} + +func TestLaunchAgent(t *testing.T) { + tests := []struct { + name string + harness string + inj injection + userArgs []string + wantArgs []string + wantEnv []string // env entries that must be present + }{ + { + name: "no prompt, no args", + harness: "codex", + wantArgs: []string{"/usr/local/bin/ucode", "codex", "--skip-preflight"}, + }, + { + name: "user args only", + harness: "codex", + userArgs: []string{"--full-auto"}, + wantArgs: []string{"/usr/local/bin/ucode", "codex", "--skip-preflight", "--", "--full-auto"}, + }, + { + name: "codex injection forwards -c before user args", + harness: "codex", + inj: injection{forwardArgs: []string{"-c", `developer_instructions="hi"`}}, + userArgs: []string{"--full-auto"}, + wantArgs: []string{"/usr/local/bin/ucode", "codex", "--skip-preflight", "--", "-c", `developer_instructions="hi"`, "--full-auto"}, + }, + { + name: "opencode injection sets env, no forwarded args", + harness: "opencode", + inj: injection{env: []string{`OPENCODE_CONFIG_CONTENT={"instructions":["/x"]}`}}, + wantArgs: []string{"/usr/local/bin/ucode", "opencode", "--skip-preflight"}, + wantEnv: []string{`OPENCODE_CONFIG_CONTENT={"instructions":["/x"]}`}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + restore := launch + t.Cleanup(func() { launch = restore }) + var got execv.Options + launch = func(opts execv.Options) error { + got = opts + return nil + } + + err := launchAgent(t.Context(), "/usr/local/bin/ucode", tt.harness, tt.inj, tt.userArgs) + require.NoError(t, err) + assert.Equal(t, tt.wantArgs, got.Args) + for _, e := range tt.wantEnv { + assert.Contains(t, got.Env, e) + } + }) + } +}