From d3f37f785a8be82176e32a55ba04b9377c47488d Mon Sep 17 00:00:00 2001 From: Russell Clarey Date: Wed, 2 Sep 2026 15:10:20 +0200 Subject: [PATCH 1/6] experimental/ssh: add the `agent-shim` command to launch a ucode agent Adds a hidden `databricks ssh agent-shim ` command that runs on the serverless driver: it probes the workspace AI Gateway, sets up PATH, and execs a ucode-configured coding agent (Claude Code, Codex) with the Databricks session context injected. This first commit assumes the toolchain (uv/ucode/node) is already present on the driver; first-run bootstrap of that toolchain follows in the next commit. Co-authored-by: Isaac --- experimental/ssh/cmd/agent_shim.go | 43 ++ experimental/ssh/cmd/ssh.go | 1 + experimental/ssh/internal/client/agentshim.go | 436 ++++++++++++++++++ .../ssh/internal/client/agentshim_test.go | 257 +++++++++++ experimental/ssh/internal/client/exec_unix.go | 11 + .../ssh/internal/client/exec_windows.go | 10 + 6 files changed, 758 insertions(+) create mode 100644 experimental/ssh/cmd/agent_shim.go create mode 100644 experimental/ssh/internal/client/agentshim.go create mode 100644 experimental/ssh/internal/client/agentshim_test.go create mode 100644 experimental/ssh/internal/client/exec_unix.go create mode 100644 experimental/ssh/internal/client/exec_windows.go diff --git a/experimental/ssh/cmd/agent_shim.go b/experimental/ssh/cmd/agent_shim.go new file mode 100644 index 00000000000..fedd1589c3b --- /dev/null +++ b/experimental/ssh/cmd/agent_shim.go @@ -0,0 +1,43 @@ +package ssh + +import ( + "github.com/databricks/cli/cmd/root" + "github.com/databricks/cli/experimental/ssh/internal/client" + "github.com/databricks/cli/libs/cmdctx" + "github.com/spf13/cobra" +) + +func newAgentShimCommand() *cobra.Command { + cmd := &cobra.Command{ + Use: "agent-shim", + Short: "Launch a ucode-configured coding agent (invoked on the remote)", + Hidden: true, + } + for _, agent := range client.SupportedAgentNames() { + cmd.AddCommand(newAgentShimAgentCommand(agent)) + } + return cmd +} + +func newAgentShimAgentCommand(agent string) *cobra.Command { + cmd := &cobra.Command{ + Use: agent, + Short: "Launch the ucode-configured " + agent + " agent", + // Disable flag parsing: forward everything after the agent name to the agent verbatim. + DisableFlagParsing: true, + } + + cmd.PreRunE = func(cmd *cobra.Command, args []string) error { + // Runs on the driver with injected env auth; no bundle, no prompt. + cmd.SetContext(root.SkipLoadBundle(cmd.Context())) + cmd.SetContext(root.SkipPrompt(cmd.Context())) + return root.MustWorkspaceClient(cmd, args) + } + + cmd.RunE = func(cmd *cobra.Command, args []string) error { + ctx := cmd.Context() + return client.RunAgentShim(ctx, cmdctx.WorkspaceClient(ctx), agent, args) + } + + return cmd +} diff --git a/experimental/ssh/cmd/ssh.go b/experimental/ssh/cmd/ssh.go index 9939d2f7050..19517cd6204 100644 --- a/experimental/ssh/cmd/ssh.go +++ b/experimental/ssh/cmd/ssh.go @@ -22,6 +22,7 @@ Use ` + "`databricks ssh connect --help`" + ` to see all available flags.`, cmd.AddCommand(newSetupCommand()) cmd.AddCommand(newConnectCommand()) cmd.AddCommand(newServerCommand()) + cmd.AddCommand(newAgentShimCommand()) return cmd } diff --git a/experimental/ssh/internal/client/agentshim.go b/experimental/ssh/internal/client/agentshim.go new file mode 100644 index 00000000000..0783f9d14dd --- /dev/null +++ b/experimental/ssh/internal/client/agentshim.go @@ -0,0 +1,436 @@ +package client + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "net/url" + "os" + "os/exec" + "path/filepath" + "regexp" + "strings" + "time" + + "github.com/databricks/cli/libs/log" + "github.com/databricks/databricks-sdk-go" +) + +const ( + // agentDir (home-relative) is the shim's working area. + agentDir = ".agent-shim" + + // shimDir holds the per-agent wrappers; must match remoteShimDir in client.go. + shimDir = agentDir + "/bin" + + // workspaceHomeEnv passes the user's workspace home to the shim for the agent context. + workspaceHomeEnv = "DATABRICKS_WORKSPACE_HOME" + + // contextFile is the scratch file holding the Databricks session context. + contextFile = "agent-system-context.md" + + // readyFile marks that first-run toolchain setup completed. + readyFile = ".ready" +) + +type agentSpec struct { + name string + // contextFlag passes the context file via this CLI flag (Claude's --append-system-prompt-file). + contextFlag string + // contextHomeFile writes the context to this $HOME instructions file (codex's AGENTS.md). + contextHomeFile string +} + +// supportedAgents are the agents the shim launches — limited to those whose ucode +// command accepts --workspace (how we target the workspace headlessly). +var supportedAgents = []agentSpec{ + {name: "claude", contextFlag: "--append-system-prompt-file"}, + {name: "codex", contextHomeFile: ".codex/AGENTS.md"}, +} + +// SupportedAgentNames lists the agents the ssh command registers a subcommand for. +func SupportedAgentNames() []string { + names := make([]string, len(supportedAgents)) + for i, a := range supportedAgents { + names[i] = a.name + } + return names +} + +func agentByName(name string) (agentSpec, bool) { + for _, a := range supportedAgents { + if a.name == name { + return a, true + } + } + return agentSpec{}, false +} + +func RunAgentShim(ctx context.Context, client *databricks.WorkspaceClient, agentName string, agentArgs []string) error { + agent, ok := agentByName(agentName) + if !ok { + return fmt.Errorf("unsupported agent %q", agentName) + } + // Probe first: fail fast before the slow first-run bootstrap if the gateway is off. + if err := probeAIGateway(ctx, client); err != nil { + return err + } + // ucode targets this workspace via --workspace, so it configures without prompting. + workspace := strings.TrimRight(client.Config.Host, "/") + return bootstrapAndLaunchAgent(ctx, agent, workspace, agentArgs) +} + +// --- Unity AI Gateway preflight --- +// +// Kept in lockstep with ucode's probe_unity_gateway_capabilities +// (databricks/ucode, src/ucode/databricks.py): probe the Unity Catalog +// model-services API (v3) first, paging through it, then fall back to the legacy +// AI Gateway endpoints API (v2). A path counts only when its JSON body actually +// lists a usable resource — a 200 with an empty collection is "reachable but +// empty", not "enabled". We fail fast with the same actionable guidance ucode +// surfaces so the shim's preflight matches what ucode itself checks. + +const ( + modelServicesPath = "/api/2.1/unity-catalog/model-services" + legacyEndpointsPath = "/api/ai-gateway/v2/endpoints" + modelServiceProbePageSize = 50 + modelServiceProbeMaxPages = 20 + aiGatewayDocsURL = "https://docs.databricks.com/aws/en/ai-gateway/overview-beta" + + // modelServiceEmptyDetail matches ucode's wording for a reachable model-services + // API that lists nothing the caller can use — almost always a UC grant gap. + modelServiceEmptyDetail = "reachable, no accessible model services returned; " + + "check USE CATALOG on system, and USE SCHEMA and EXECUTE on system.ai" +) + +// gatewayProbe is the outcome of probing one AI Gateway API. resourceAvailable +// means the API returned at least one usable resource; conclusive is false when +// paging couldn't be completed, so "no resources" was never actually confirmed. +type gatewayProbe struct { + reachable bool + detail string + resourceAvailable bool + conclusive bool +} + +// probeAIGateway fails fast if the workspace's Unity AI Gateway can't run an +// agent. It returns nil when the gateway is usable — logging a warning when it's +// reachable but exposes no model services to this caller — and an actionable +// error otherwise. +func probeAIGateway(ctx context.Context, client *databricks.WorkspaceClient) error { + host := strings.TrimRight(client.Config.Host, "/") + + modelSvc := probeModelServices(ctx, client, host) + // A 401 (or a 400 "invalid token") can't be rescued by trying another API, + // so surface it before the fallback probe. + if !modelSvc.reachable && looksLikeDefinitiveAuthFailure(modelSvc.detail) { + return aiGatewayAuthError(host, modelSvc.detail) + } + if modelSvc.resourceAvailable { + return nil + } + + legacy := probeLegacyEndpoints(ctx, client, host) + switch { + case legacy.reachable: + // The legacy endpoints API answered, so the gateway is enabled even if no + // model services are visible to this caller. + log.Warnf(ctx, "Unity AI Gateway model service check: %s", modelSvc.detail) + return nil + case modelSvc.reachable && !modelSvc.conclusive: + // v3 answered but couldn't be paged to completion; give it the benefit of + // the doubt rather than blocking on an unconfirmed "empty". + return nil + case looksLikeDefinitiveAuthFailure(legacy.detail): + return aiGatewayAuthError(host, legacy.detail) + case looksLikeScopeFailure(modelSvc.detail): + return aiGatewayScopeError(host, modelSvc.detail) + case looksLikeScopeFailure(legacy.detail): + return aiGatewayScopeError(host, legacy.detail) + case looksLikePermissionFailure(modelSvc.detail): + return fmt.Errorf("Databricks Unity AI Gateway model service access could not be verified on %s (%s). The legacy endpoint fallback also failed (%s). The model service probe requires permission to list Unity Catalog model services. Verify USE CATALOG on `system`, and USE SCHEMA and EXECUTE on `system.ai`", host, modelSvc.detail, legacy.detail) + case looksLikePermissionFailure(legacy.detail): + return fmt.Errorf("Databricks Unity AI Gateway legacy endpoint access could not be verified on %s (%s). The model service probe also failed (%s). Verify the caller's workspace permissions for the legacy endpoints listing", host, legacy.detail, modelSvc.detail) + default: + return fmt.Errorf("Databricks Unity AI Gateway is not enabled on this workspace (%s): neither model services (%s) nor legacy endpoints (%s) are available. See %s", host, modelSvc.detail, legacy.detail, aiGatewayDocsURL) + } +} + +// probeModelServices probes the UC model-services API (v3), paging until it finds +// an accessible model service, exhausts the cursor, or hits the page cap. +func probeModelServices(ctx context.Context, client *databricks.WorkspaceClient, host string) gatewayProbe { + pageToken := "" + for page := 0; page < modelServiceProbeMaxPages; page++ { + reqURL := fmt.Sprintf("%s%s?page_size=%d", host, modelServicesPath, modelServiceProbePageSize) + if pageToken != "" { + reqURL += "&page_token=" + url.QueryEscape(pageToken) + } + payload, reason := gatewayGetJSON(ctx, client, reqURL) + if payload == nil { + // A first-page failure is the real reachability signal; a later page + // failing still means the API answered at least once, so treat it as + // reachable-but-inconclusive rather than a confirmed "empty". + if page == 0 { + return gatewayProbe{detail: versionNeutralGatewayDetail(reason), conclusive: true} + } + return gatewayProbe{reachable: true, detail: "reachable"} + } + if hasNonEmptyCollection(payload, "model_services") { + return gatewayProbe{reachable: true, detail: "reachable, accessible model service returned", resourceAvailable: true, conclusive: true} + } + pageToken = stringField(payload, "next_page_token") + if pageToken == "" { + return gatewayProbe{reachable: true, detail: modelServiceEmptyDetail, conclusive: true} + } + } + return gatewayProbe{reachable: true, detail: "reachable"} +} + +// probeLegacyEndpoints probes the legacy AI Gateway endpoints API (v2). +func probeLegacyEndpoints(ctx context.Context, client *databricks.WorkspaceClient, host string) gatewayProbe { + payload, reason := gatewayGetJSON(ctx, client, host+legacyEndpointsPath+"?page_size=1") + if payload == nil { + return gatewayProbe{detail: versionNeutralGatewayDetail(reason), conclusive: true} + } + if hasNonEmptyCollection(payload, "endpoints") { + return gatewayProbe{reachable: true, detail: "reachable, accessible endpoint returned", resourceAvailable: true, conclusive: true} + } + return gatewayProbe{reachable: true, detail: "reachable, no accessible endpoints returned", conclusive: true} +} + +// gatewayGetJSON issues an authenticated GET expecting JSON. It returns the +// decoded payload on HTTP 200, or (nil, reason) on any failure — mirroring +// ucode's _http_get_json, including a short response-body excerpt so gateway auth +// errors (e.g. a 400 whose body is "Invalid Token") stay visible in the reason. +func gatewayGetJSON(ctx context.Context, client *databricks.WorkspaceClient, reqURL string) (any, string) { + req, err := http.NewRequestWithContext(ctx, http.MethodGet, reqURL, nil) + if err != nil { + return nil, fmt.Sprintf("network error: %v", err) + } + req.Header.Set("Accept", "application/json") + if err := client.Config.Authenticate(req); err != nil { + return nil, fmt.Sprintf("network error: %v", err) + } + resp, err := (&http.Client{Transport: client.Config.HTTPTransport, Timeout: 10 * time.Second}).Do(req) + if err != nil { + return nil, fmt.Sprintf("network error: %v", err) + } + defer resp.Body.Close() + body, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<20)) + if resp.StatusCode == http.StatusOK { + var payload any + if err := json.Unmarshal(body, &payload); err != nil { + return nil, fmt.Sprintf("response was not valid JSON (%v)", err) + } + return payload, "" + } + reason := fmt.Sprintf("HTTP %d %s", resp.StatusCode, http.StatusText(resp.StatusCode)) + if excerpt := strings.TrimSpace(string(body)); excerpt != "" { + if len(excerpt) > 200 { + excerpt = excerpt[:200] + } + reason += ": " + excerpt + } + return nil, reason +} + +var ( + gatewayDetailV3 = regexp.MustCompile(`(?i)\bv3\b`) + gatewayDetailV2 = regexp.MustCompile(`(?i)\bv2\b`) +) + +// versionNeutralGatewayDetail rewrites the internal v2/v3 API labels in a failure +// reason into user-facing terms, matching ucode so the error text stays in sync. +func versionNeutralGatewayDetail(detail string) string { + detail = gatewayDetailV3.ReplaceAllString(detail, "model service") + return gatewayDetailV2.ReplaceAllString(detail, "legacy endpoint") +} + +// looksLikeDefinitiveAuthFailure is true when retrying another workspace API +// can't rescue the token. A bare 403 is left out: it can be endpoint-specific +// authorization, so the preflight still tries the fallback before giving up. +func looksLikeDefinitiveAuthFailure(reason string) bool { + if strings.Contains(reason, "HTTP 401") { + return true + } + return strings.Contains(reason, "HTTP 400") && strings.Contains(strings.ToLower(reason), "invalid token") +} + +// looksLikeScopeFailure matches a 403 that reports the OAuth token is missing a +// required scope (which re-login can fix), as opposed to a plain permission 403. +func looksLikeScopeFailure(reason string) bool { + l := strings.ToLower(reason) + return strings.Contains(l, "http 403") && strings.Contains(l, "oauth token") && strings.Contains(l, "required scopes") +} + +func looksLikePermissionFailure(reason string) bool { + return strings.Contains(reason, "HTTP 403") +} + +func aiGatewayAuthError(host, reason string) error { + return fmt.Errorf("Databricks rejected the access token for %s (%s). Try:\n databricks auth logout --host %s\n databricks auth login --host %s", host, reason, host, host) +} + +func aiGatewayScopeError(host, reason string) error { + return fmt.Errorf("the access token for %s is missing an OAuth scope required by the AI Gateway APIs (%s). Re-authenticate to mint a token with the needed scopes:\n databricks auth login --host %s", host, reason, host) +} + +// hasNonEmptyCollection reports whether payload is a JSON object with a non-empty +// array at key. +func hasNonEmptyCollection(payload any, key string) bool { + obj, ok := payload.(map[string]any) + if !ok { + return false + } + arr, ok := obj[key].([]any) + return ok && len(arr) > 0 +} + +// stringField returns payload[key] when payload is a JSON object holding a string +// there, else "". +func stringField(payload any, key string) string { + obj, ok := payload.(map[string]any) + if !ok { + return "" + } + s, _ := obj[key].(string) + return s +} + +func bootstrapAndLaunchAgent(ctx context.Context, agent agentSpec, workspace string, agentArgs []string) error { + home, err := os.UserHomeDir() + if err != nil { + return fmt.Errorf("failed to resolve home directory: %w", err) + } + shimDir := filepath.Join(home, shimDir) + readyFile := filepath.Join(home, agentDir, readyFile) + + // Fast path: first-run setup recorded the resolved PATH; restore it and launch. + if data, err := os.ReadFile(readyFile); err == nil { + _ = os.Setenv("PATH", strings.TrimSpace(string(data))) + return launchUcodeAgent(agent, workspace, agentArgs) + } + + // Build the PATH tooling runs under. Drop the shim dir so ucode installs and + // execs the real agent rather than this wrapper; prepend ucode's install bin + // and the uploaded databricks CLI's dir (so ucode finds it and skips its own + // install). + removePath(shimDir) + prependPath(filepath.Join(home, ".local", "bin")) + if self, err := os.Executable(); err == nil { + prependPath(filepath.Dir(self)) + } + + // Record the resolved PATH so later launches take the fast path. + if err := os.MkdirAll(filepath.Dir(readyFile), 0o755); err != nil { + return fmt.Errorf("failed to create %s: %w", filepath.Dir(readyFile), err) + } + if err := os.WriteFile(readyFile, []byte(os.Getenv("PATH")), 0o644); err != nil { + return fmt.Errorf("failed to record setup completion: %w", err) + } + + return launchUcodeAgent(agent, workspace, agentArgs) +} + +// launchUcodeAgent launches the ucode-configured agent, replacing this process. +func launchUcodeAgent(agent agentSpec, workspace string, agentArgs []string) error { + home, err := os.UserHomeDir() + if err != nil { + return fmt.Errorf("failed to resolve home directory: %w", err) + } + ucodePath, err := exec.LookPath("ucode") + if err != nil { + return fmt.Errorf("ucode not found on PATH after setup: %w", err) + } + contextArgs, err := injectAgentContext(home, agent) + if err != nil { + return err + } + argv := []string{"ucode", agent.name} + if workspace != "" { + // Every supported agent accepts --workspace (the gate for supportedAgents). + argv = append(argv, "--workspace", workspace) + } + argv = append(argv, contextArgs...) + argv = append(argv, agentArgs...) + // Pass the session token as DATABRICKS_BEARER so stock ucode authenticates headlessly. + if os.Getenv("DATABRICKS_BEARER") == "" { + if token := os.Getenv("DATABRICKS_TOKEN"); token != "" { + _ = os.Setenv("DATABRICKS_BEARER", token) + } + } + return execProcess(ucodePath, argv, os.Environ()) +} + +// injectAgentContext puts the Databricks context where the agent reads it, returning any extra argv. +func injectAgentContext(home string, agent agentSpec) ([]string, error) { + context := agentSystemContext(os.Getenv(workspaceHomeEnv)) + switch { + case agent.contextFlag != "": + f := filepath.Join(home, contextFile) + if err := os.WriteFile(f, []byte(context), 0o644); err != nil { + return nil, fmt.Errorf("failed to write agent context file: %w", err) + } + return []string{agent.contextFlag, f}, nil + case agent.contextHomeFile != "": + f := filepath.Join(home, agent.contextHomeFile) + if _, err := os.Stat(f); err == nil { + return nil, nil // already present — don't clobber + } else if !errors.Is(err, os.ErrNotExist) { + return nil, fmt.Errorf("failed to stat %s: %w", f, err) + } + if err := os.MkdirAll(filepath.Dir(f), 0o755); err != nil { + return nil, fmt.Errorf("failed to create %s: %w", filepath.Dir(f), err) + } + if err := os.WriteFile(f, []byte(context), 0o644); err != nil { + return nil, fmt.Errorf("failed to write %s: %w", f, err) + } + return nil, nil + default: + return nil, nil + } +} + +// prependPath puts dir at the front of $PATH, dropping any later duplicate. +func prependPath(dir string) { + dirs := []string{dir} + for _, d := range filepath.SplitList(os.Getenv("PATH")) { + if d != dir { + dirs = append(dirs, d) + } + } + _ = os.Setenv("PATH", strings.Join(dirs, string(os.PathListSeparator))) +} + +// removePath drops every occurrence of dir from $PATH. +func removePath(dir string) { + var dirs []string + for _, d := range filepath.SplitList(os.Getenv("PATH")) { + if d != dir { + dirs = append(dirs, d) + } + } + _ = os.Setenv("PATH", strings.Join(dirs, string(os.PathListSeparator))) +} + +// agentSystemContext returns the Databricks session context; wsHome names the working directory. +func agentSystemContext(wsHome string) string { + cwd := "the user's Databricks workspace home directory" + if wsHome != "" { + cwd = wsHome + } + return fmt.Sprintf(`You are running inside a "databricks ssh connect" session on the driver node of a +Databricks serverless cluster. +- The "databricks" CLI is installed and already authenticated: DATABRICKS_HOST and + DATABRICKS_TOKEN are set in the environment, so "databricks ..." commands work with no + "databricks auth login". The same token governs Unity Catalog and serving-endpoint access. +- This container is ephemeral; only paths under /Workspace persist across sessions. Your + working directory is %s. +- DATABRICKS_TOKEN is a static session token that may expire during a long session; if + "databricks" calls start failing with auth errors, the session likely needs reconnecting. +- You can run shell commands and use the "databricks" CLI to explore the workspace + (clusters, jobs, Unity Catalog, DBFS, etc.).`, cwd) +} diff --git a/experimental/ssh/internal/client/agentshim_test.go b/experimental/ssh/internal/client/agentshim_test.go new file mode 100644 index 00000000000..621e1ce1131 --- /dev/null +++ b/experimental/ssh/internal/client/agentshim_test.go @@ -0,0 +1,257 @@ +package client + +import ( + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/config" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestPrependPath(t *testing.T) { + t.Setenv("PATH", strings.Join([]string{"/a", "/b"}, string(os.PathListSeparator))) + prependPath("/b") // existing entry moves to the front (deduped) + assert.Equal(t, []string{"/b", "/a"}, filepath.SplitList(os.Getenv("PATH"))) + prependPath("/new") + assert.Equal(t, []string{"/new", "/b", "/a"}, filepath.SplitList(os.Getenv("PATH"))) +} + +func TestRemovePath(t *testing.T) { + t.Setenv("PATH", strings.Join([]string{"/a", "/b", "/a"}, string(os.PathListSeparator))) + removePath("/a") + assert.Equal(t, []string{"/b"}, filepath.SplitList(os.Getenv("PATH"))) +} + +func TestSupportedAgents(t *testing.T) { + names := SupportedAgentNames() + // Limited to the agents whose ucode command accepts --workspace. + assert.Equal(t, []string{"claude", "codex"}, names) + + // Each agent injects context by exactly one mechanism (flag OR home file); + // the two are never both set. + expect := map[string]struct{ flag, home string }{ + "claude": {flag: "--append-system-prompt-file"}, + "codex": {home: ".codex/AGENTS.md"}, + } + for name, want := range expect { + a, ok := agentByName(name) + require.True(t, ok) + assert.Equal(t, want.flag, a.contextFlag, "%s contextFlag", name) + assert.Equal(t, want.home, a.contextHomeFile, "%s contextHomeFile", name) + } + + _, ok := agentByName("not-an-agent") + assert.False(t, ok) +} + +func TestInjectAgentContext(t *testing.T) { + claude, _ := agentByName("claude") + codex, _ := agentByName("codex") + noContext := agentSpec{name: "none"} // an agent with neither mechanism + + t.Run("flag agent writes a scratch file and returns the flag", func(t *testing.T) { + home := t.TempDir() + args, err := injectAgentContext(home, claude) + require.NoError(t, err) + require.Len(t, args, 2) + assert.Equal(t, "--append-system-prompt-file", args[0]) + data, err := os.ReadFile(args[1]) + require.NoError(t, err) + assert.Contains(t, string(data), "databricks ssh connect") + }) + + t.Run("home-file agent writes the global instructions file, no args", func(t *testing.T) { + home := t.TempDir() + args, err := injectAgentContext(home, codex) + require.NoError(t, err) + assert.Nil(t, args) + data, err := os.ReadFile(filepath.Join(home, ".codex", "AGENTS.md")) + require.NoError(t, err) + assert.Contains(t, string(data), "databricks ssh connect") + }) + + t.Run("home-file agent does not clobber an existing file", func(t *testing.T) { + home := t.TempDir() + target := filepath.Join(home, ".codex", "AGENTS.md") + require.NoError(t, os.MkdirAll(filepath.Dir(target), 0o755)) + require.NoError(t, os.WriteFile(target, []byte("user's own instructions"), 0o644)) + args, err := injectAgentContext(home, codex) + require.NoError(t, err) + assert.Nil(t, args) + data, err := os.ReadFile(target) + require.NoError(t, err) + assert.Equal(t, "user's own instructions", string(data)) + }) + + t.Run("agent without a mechanism writes nothing", func(t *testing.T) { + home := t.TempDir() + args, err := injectAgentContext(home, noContext) + require.NoError(t, err) + assert.Nil(t, args) + entries, err := os.ReadDir(home) + require.NoError(t, err) + assert.Empty(t, entries) + }) +} + +func TestAgentSystemContext(t *testing.T) { + assert.Contains(t, agentSystemContext("/Workspace/Users/me@example.com"), "working directory is /Workspace/Users/me@example.com") + // Falls back to a generic phrase when the workspace home is unknown. + assert.Contains(t, agentSystemContext(""), "working directory is the user's Databricks workspace home directory") +} + +func newProbeClient(t *testing.T, host string) *databricks.WorkspaceClient { + t.Helper() + w, err := databricks.NewWorkspaceClient((*databricks.Config)(&config.Config{Host: host, Token: "test-token"})) + require.NoError(t, err) + return w +} + +// newGatewayServer routes the model-services (v3) and legacy endpoints (v2) probe +// requests to per-API handlers, 404ing anything else (e.g. the SDK's config +// discovery). A nil handler means that API is not served (404). +func newGatewayServer(t *testing.T, modelServices, endpoints http.HandlerFunc) *httptest.Server { + t.Helper() + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch { + case strings.HasPrefix(r.URL.Path, "/api/2.1/unity-catalog/model-services") && modelServices != nil: + assert.Equal(t, "Bearer test-token", r.Header.Get("Authorization")) + modelServices(w, r) + case strings.HasPrefix(r.URL.Path, "/api/ai-gateway/v2/endpoints") && endpoints != nil: + assert.Equal(t, "Bearer test-token", r.Header.Get("Authorization")) + endpoints(w, r) + default: + w.WriteHeader(http.StatusNotFound) + _, _ = w.Write([]byte("{}")) + } + })) +} + +// jsonHandler replies with status and body for every request it receives. +func jsonHandler(status int, body string) http.HandlerFunc { + return func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(status) + _, _ = w.Write([]byte(body)) + } +} + +func TestProbeAIGateway(t *testing.T) { + const modelServicesBody = `{"model_services":[{"name":"model-services/system.ai.gpt-5"}]}` + const endpointsBody = `{"endpoints":[{"name":"databricks-gpt-5"}]}` + const emptyBody = `{}` + + t.Run("model service available: connected, legacy not probed", func(t *testing.T) { + var legacyCalls int + srv := newGatewayServer(t, + jsonHandler(http.StatusOK, modelServicesBody), + func(w http.ResponseWriter, r *http.Request) { + legacyCalls++ + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(emptyBody)) + }, + ) + defer srv.Close() + require.NoError(t, probeAIGateway(t.Context(), newProbeClient(t, srv.URL))) + assert.Zero(t, legacyCalls, "legacy endpoints must not be probed once a model service is found") + }) + + t.Run("model service found across pages", func(t *testing.T) { + srv := newGatewayServer(t, func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + if r.URL.Query().Get("page_token") == "" { + _, _ = w.Write([]byte(`{"next_page_token":"cursor-1"}`)) + return + } + _, _ = w.Write([]byte(modelServicesBody)) + }, nil) + defer srv.Close() + require.NoError(t, probeAIGateway(t.Context(), newProbeClient(t, srv.URL))) + }) + + t.Run("empty model service but legacy reachable: proceeds", func(t *testing.T) { + // v3 200-empty and v2 200-empty: both reachable, no resources → proceed (warn). + srv := newGatewayServer(t, jsonHandler(http.StatusOK, emptyBody), jsonHandler(http.StatusOK, emptyBody)) + defer srv.Close() + assert.NoError(t, probeAIGateway(t.Context(), newProbeClient(t, srv.URL))) + }) + + t.Run("legacy-only workspace (v3 404, v2 reachable): proceeds", func(t *testing.T) { + srv := newGatewayServer(t, + jsonHandler(http.StatusNotFound, `{"message":"not found"}`), + jsonHandler(http.StatusOK, endpointsBody), + ) + defer srv.Close() + assert.NoError(t, probeAIGateway(t.Context(), newProbeClient(t, srv.URL))) + }) + + t.Run("auth failure (401) fails fast without probing legacy", func(t *testing.T) { + var legacyCalls int + srv := newGatewayServer(t, + jsonHandler(http.StatusUnauthorized, `{"message":"unauthorized"}`), + func(w http.ResponseWriter, r *http.Request) { + legacyCalls++ + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(endpointsBody)) + }, + ) + defer srv.Close() + err := probeAIGateway(t.Context(), newProbeClient(t, srv.URL)) + require.Error(t, err) + assert.Contains(t, err.Error(), "rejected the access token") + assert.Zero(t, legacyCalls, "a definitive auth failure must not fall back to the legacy probe") + }) + + t.Run("missing OAuth scope on both paths routes to re-auth", func(t *testing.T) { + scope := jsonHandler(http.StatusForbidden, "Provided OAuth token does not have required scopes: unity-catalog") + srv := newGatewayServer(t, scope, scope) + defer srv.Close() + err := probeAIGateway(t.Context(), newProbeClient(t, srv.URL)) + require.Error(t, err) + assert.Contains(t, err.Error(), "missing an OAuth scope") + assert.Contains(t, err.Error(), "databricks auth login") + }) + + t.Run("plain 403 on both paths is a permission failure", func(t *testing.T) { + forbidden := jsonHandler(http.StatusForbidden, `{"message":"forbidden"}`) + srv := newGatewayServer(t, forbidden, forbidden) + defer srv.Close() + err := probeAIGateway(t.Context(), newProbeClient(t, srv.URL)) + require.Error(t, err) + assert.Contains(t, err.Error(), "model service access could not be verified") + }) + + t.Run("neither gateway available: not enabled", func(t *testing.T) { + notFound := jsonHandler(http.StatusNotFound, `{"message":"not found"}`) + srv := newGatewayServer(t, notFound, notFound) + defer srv.Close() + err := probeAIGateway(t.Context(), newProbeClient(t, srv.URL)) + require.Error(t, err) + assert.Contains(t, err.Error(), "not enabled on this workspace") + }) +} + +func TestProbeModelServicesInconclusive(t *testing.T) { + // A later page failing (after the first page succeeded) is reachable-but- + // inconclusive, never a confirmed "empty". + srv := newGatewayServer(t, func(w http.ResponseWriter, r *http.Request) { + if r.URL.Query().Get("page_token") == "" { + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`{"next_page_token":"cursor-1"}`)) + return + } + w.WriteHeader(http.StatusInternalServerError) + _, _ = w.Write([]byte("boom")) + }, nil) + defer srv.Close() + + probe := probeModelServices(t.Context(), newProbeClient(t, srv.URL), strings.TrimRight(srv.URL, "/")) + assert.True(t, probe.reachable) + assert.False(t, probe.conclusive) + assert.False(t, probe.resourceAvailable) +} diff --git a/experimental/ssh/internal/client/exec_unix.go b/experimental/ssh/internal/client/exec_unix.go new file mode 100644 index 00000000000..8c024b980bf --- /dev/null +++ b/experimental/ssh/internal/client/exec_unix.go @@ -0,0 +1,11 @@ +//go:build !windows + +package client + +import "syscall" + +// execProcess replaces the current process (execve(2)) so the ssh PTY connects +// straight to the agent and its exit code propagates. +func execProcess(argv0 string, argv, env []string) error { + return syscall.Exec(argv0, argv, env) +} diff --git a/experimental/ssh/internal/client/exec_windows.go b/experimental/ssh/internal/client/exec_windows.go new file mode 100644 index 00000000000..3cfc5b2e1ed --- /dev/null +++ b/experimental/ssh/internal/client/exec_windows.go @@ -0,0 +1,10 @@ +//go:build windows + +package client + +import "errors" + +// execProcess is unsupported on Windows (no syscall.Exec); the shim only runs on the Linux driver. +func execProcess(argv0 string, argv, env []string) error { + return errors.New("process exec is not supported on Windows") +} From 9a97b990f63bf53c0d32aa6b2df08b4cb5143c05 Mon Sep 17 00:00:00 2001 From: Russell Clarey Date: Thu, 3 Sep 2026 11:18:06 +0200 Subject: [PATCH 2/6] experimental/ssh: bootstrap the agent toolchain on first run On first launch the shim now installs the toolchain the agent needs before execing ucode: uv (into ~/.local/bin), a pinned ucode release via uv, and Node/npm (downloaded and SHA256-verified) when the serverless image ships none. It also silences npm's update-notifier and records the resolved PATH so later launches skip straight to the agent. Co-authored-by: Isaac --- experimental/ssh/internal/client/agentshim.go | 388 ++++++++++++++++-- .../ssh/internal/client/agentshim_test.go | 99 ++++- 2 files changed, 436 insertions(+), 51 deletions(-) diff --git a/experimental/ssh/internal/client/agentshim.go b/experimental/ssh/internal/client/agentshim.go index 0783f9d14dd..e5ed11eac13 100644 --- a/experimental/ssh/internal/client/agentshim.go +++ b/experimental/ssh/internal/client/agentshim.go @@ -2,6 +2,8 @@ package client import ( "context" + "crypto/sha256" + "encoding/hex" "encoding/json" "errors" "fmt" @@ -12,9 +14,13 @@ import ( "os/exec" "path/filepath" "regexp" + "runtime" "strings" "time" + "github.com/databricks/cli/libs/auth" + "github.com/databricks/cli/libs/cmdio" + "github.com/databricks/cli/libs/env" "github.com/databricks/cli/libs/log" "github.com/databricks/databricks-sdk-go" ) @@ -26,14 +32,30 @@ const ( // shimDir holds the per-agent wrappers; must match remoteShimDir in client.go. shimDir = agentDir + "/bin" + // ucodeRepo is the GitHub repo the shim installs ucode from. + ucodeRepo = "databricks/ucode" + + // ucodeVersion pins the ucode release tag the shim installs, so a new upstream + // release can't change the toolchain out from under an SSH session unexpectedly. + ucodeVersion = "v0.1.0" + // workspaceHomeEnv passes the user's workspace home to the shim for the agent context. workspaceHomeEnv = "DATABRICKS_WORKSPACE_HOME" + // depsDir (home-relative) holds toolchain deps (Node/npm) fetched when the image ships none. + depsDir = agentDir + "/deps" + // contextFile is the scratch file holding the Databricks session context. contextFile = "agent-system-context.md" - // readyFile marks that first-run toolchain setup completed. - readyFile = ".ready" + // setupLockName (under agentDir) serializes first-run toolchain setup across + // concurrent SSH clients sharing this driver's $HOME. + setupLockName = ".setup.lock" + + // setupLockStaleAfter reclaims a setup lock left behind by a process that died + // mid-install. It is deliberately generous: a cold first run downloads uv, + // ucode, and Node, which can legitimately take minutes. + setupLockStaleAfter = 10 * time.Minute ) type agentSpec struct { @@ -150,12 +172,16 @@ func probeAIGateway(ctx context.Context, client *databricks.WorkspaceClient) err return aiGatewayScopeError(host, modelSvc.detail) case looksLikeScopeFailure(legacy.detail): return aiGatewayScopeError(host, legacy.detail) + case looksLikeTransient(modelSvc.detail) || looksLikeTransient(legacy.detail): + // A rate-limit/5xx/network blip is not "disabled" — tell the user to retry + // rather than sending them to the enablement docs. + return fmt.Errorf("could not verify the Databricks Unity AI Gateway on %s: the probe hit a transient error (model services: %s; legacy endpoints: %s). Retry in a moment", host, modelSvc.detail, legacy.detail) case looksLikePermissionFailure(modelSvc.detail): - return fmt.Errorf("Databricks Unity AI Gateway model service access could not be verified on %s (%s). The legacy endpoint fallback also failed (%s). The model service probe requires permission to list Unity Catalog model services. Verify USE CATALOG on `system`, and USE SCHEMA and EXECUTE on `system.ai`", host, modelSvc.detail, legacy.detail) + return fmt.Errorf("model service access could not be verified on %s (%s). The legacy endpoint fallback also failed (%s). The model service probe requires permission to list Unity Catalog model services. Verify USE CATALOG on `system`, and USE SCHEMA and EXECUTE on `system.ai`", host, modelSvc.detail, legacy.detail) case looksLikePermissionFailure(legacy.detail): - return fmt.Errorf("Databricks Unity AI Gateway legacy endpoint access could not be verified on %s (%s). The model service probe also failed (%s). Verify the caller's workspace permissions for the legacy endpoints listing", host, legacy.detail, modelSvc.detail) + return fmt.Errorf("legacy endpoint access could not be verified on %s (%s). The model service probe also failed (%s). Verify the caller's workspace permissions for the legacy endpoints listing", host, legacy.detail, modelSvc.detail) default: - return fmt.Errorf("Databricks Unity AI Gateway is not enabled on this workspace (%s): neither model services (%s) nor legacy endpoints (%s) are available. See %s", host, modelSvc.detail, legacy.detail, aiGatewayDocsURL) + return fmt.Errorf("the Databricks Unity AI Gateway is not enabled on this workspace (%s): neither model services (%s) nor legacy endpoints (%s) are available. See %s", host, modelSvc.detail, legacy.detail, aiGatewayDocsURL) } } @@ -211,6 +237,12 @@ func gatewayGetJSON(ctx context.Context, client *databricks.WorkspaceClient, req return nil, fmt.Sprintf("network error: %v", err) } req.Header.Set("Accept", "application/json") + // Set X-Databricks-Workspace-Id (and suppress the legacy workspace_id=none + // sentinel) so these hand-written workspace-routed GETs reach the workspace + // plane on unified/SPOG hosts instead of the account plane. + for k, v := range auth.WorkspaceIDHeaders(client.Config) { + req.Header.Set(k, v) + } if err := client.Config.Authenticate(req); err != nil { return nil, fmt.Sprintf("network error: %v", err) } @@ -270,8 +302,18 @@ func looksLikePermissionFailure(reason string) bool { return strings.Contains(reason, "HTTP 403") } +// gatewayTransientStatus matches the HTTP statuses worth retrying (429 + any 5xx). +var gatewayTransientStatus = regexp.MustCompile(`HTTP (429|5\d\d)`) + +// looksLikeTransient is true for a probe failure that is likely temporary (rate +// limit, server error, or a network/transport error) rather than a stable +// "disabled"/"unauthorized" verdict. +func looksLikeTransient(reason string) bool { + return strings.HasPrefix(reason, "network error") || gatewayTransientStatus.MatchString(reason) +} + func aiGatewayAuthError(host, reason string) error { - return fmt.Errorf("Databricks rejected the access token for %s (%s). Try:\n databricks auth logout --host %s\n databricks auth login --host %s", host, reason, host, host) + return fmt.Errorf("the Databricks workspace %s rejected the access token (%s). Try:\n databricks auth logout --host %s\n databricks auth login --host %s", host, reason, host, host) } func aiGatewayScopeError(host, reason string) error { @@ -301,51 +343,141 @@ func stringField(payload any, key string) string { } func bootstrapAndLaunchAgent(ctx context.Context, agent agentSpec, workspace string, agentArgs []string) error { - home, err := os.UserHomeDir() + home, err := env.UserHomeDir(ctx) if err != nil { return fmt.Errorf("failed to resolve home directory: %w", err) } - shimDir := filepath.Join(home, shimDir) - readyFile := filepath.Join(home, agentDir, readyFile) - - // Fast path: first-run setup recorded the resolved PATH; restore it and launch. - if data, err := os.ReadFile(readyFile); err == nil { - _ = os.Setenv("PATH", strings.TrimSpace(string(data))) - return launchUcodeAgent(agent, workspace, agentArgs) - } - // Build the PATH tooling runs under. Drop the shim dir so ucode installs and - // execs the real agent rather than this wrapper; prepend ucode's install bin - // and the uploaded databricks CLI's dir (so ucode finds it and skips its own - // install). - removePath(shimDir) - prependPath(filepath.Join(home, ".local", "bin")) + // Reconstruct the PATH tooling runs under on every launch rather than caching + // it: every entry is a known location, so deriving it can't restore a stale + // path (e.g. a versioned CLI dir from an older session). Drop the shim dir so + // ucode execs the real agent rather than this wrapper, then prepend uv/ucode's + // bin, this databricks CLI's own dir (so ucode finds it and skips its own + // install), and the Node bin ensureNode installs into. + removePath(ctx, filepath.Join(home, shimDir)) + prependPath(ctx, filepath.Join(home, ".local", "bin")) if self, err := os.Executable(); err == nil { - prependPath(filepath.Dir(self)) + prependPath(ctx, filepath.Dir(self)) } + prependPath(ctx, filepath.Join(home, depsDir, "node", "bin")) - // Record the resolved PATH so later launches take the fast path. - if err := os.MkdirAll(filepath.Dir(readyFile), 0o755); err != nil { - return fmt.Errorf("failed to create %s: %w", filepath.Dir(readyFile), err) + if err := ensureToolchain(ctx, home); err != nil { + return err } - if err := os.WriteFile(readyFile, []byte(os.Getenv("PATH")), 0o644); err != nil { - return fmt.Errorf("failed to record setup completion: %w", err) + + // npm is on PATH now (shipped or just installed); silence its update-notifier + // box so it doesn't clutter the agent session. + disableNpmUpdateNotifier(ctx) + + // Put npm's global bin on PATH so an npm-installed agent binary resolves after + // ucode installs it. + if prefix := npmGlobalPrefix(ctx); prefix != "" { + prependPath(ctx, filepath.Join(prefix, "bin")) } - return launchUcodeAgent(agent, workspace, agentArgs) + return launchUcodeAgent(ctx, home, agent, workspace, agentArgs) } -// launchUcodeAgent launches the ucode-configured agent, replacing this process. -func launchUcodeAgent(agent agentSpec, workspace string, agentArgs []string) error { - home, err := os.UserHomeDir() +// toolchainReady reports whether the agent toolchain is already installed, derived +// purely from PATH (which the caller has already set up). ucode launches the agent +// and npm provides the Node runtime agents like Claude Code need. +func toolchainReady() bool { + _, ucodeErr := exec.LookPath("ucode") + _, npmErr := exec.LookPath("npm") + return ucodeErr == nil && npmErr == nil +} + +// ensureToolchain installs uv, ucode, and Node/npm when missing. It is a no-op +// once the toolchain is present, and holds a machine-local lock while installing +// so two first-run SSH clients sharing this $HOME don't concurrently install into +// (and corrupt) the same paths. +func ensureToolchain(ctx context.Context, home string) error { + if toolchainReady() { + return nil + } + unlock, err := acquireSetupLock(ctx, home) if err != nil { - return fmt.Errorf("failed to resolve home directory: %w", err) + return err + } + defer unlock() + // Another client may have finished the install while we waited for the lock. + if toolchainReady() { + return nil + } + return installToolchain(ctx, home) +} + +// installToolchain performs the actual installs. The caller holds the setup lock. +func installToolchain(ctx context.Context, home string) error { + // 1. uv (installs into ~/.local/bin). + if _, err := exec.LookPath("uv"); err != nil { + cmdio.LogString(ctx, "Installing uv...") + if err := runShell(ctx, "curl -LsSf https://astral.sh/uv/install.sh | sh"); err != nil { + return fmt.Errorf("failed to install uv: %w", err) + } } + + // 2. ucode (pinned stock upstream release). + if _, err := exec.LookPath("ucode"); err != nil { + cmdio.LogString(ctx, "Installing ucode...") + if err := runCommand(ctx, "uv", "tool", "install", "git+https://github.com/"+ucodeRepo+"@"+ucodeVersion); err != nil { + return fmt.Errorf("failed to install ucode: %w", err) + } + } + + // 3. Node/npm. Some agents (e.g. Claude Code) need it and serverless images + // don't ship it. It installs into depsDir/node/bin, already prepended to PATH. + if _, err := exec.LookPath("npm"); err != nil { + cmdio.LogString(ctx, "Installing npm...") + if _, err := ensureNode(ctx, home); err != nil { + return err + } + } + + return nil +} + +// acquireSetupLock takes a machine-local lock (an O_EXCL sentinel file, matching +// the pattern in libs/versioncheck) serializing first-run setup across SSH +// clients. It blocks until the lock is free, reclaiming one left behind by a dead +// process after setupLockStaleAfter. The returned func releases the lock. +func acquireSetupLock(ctx context.Context, home string) (func(), error) { + lockPath := filepath.Join(home, agentDir, setupLockName) + if err := os.MkdirAll(filepath.Dir(lockPath), 0o755); err != nil { + return nil, fmt.Errorf("failed to create %s: %w", filepath.Dir(lockPath), err) + } + for { + f, err := os.OpenFile(lockPath, os.O_CREATE|os.O_EXCL|os.O_WRONLY, 0o644) + if err == nil { + _ = f.Close() + return func() { _ = os.Remove(lockPath) }, nil + } + if !errors.Is(err, os.ErrExist) { + return nil, fmt.Errorf("failed to acquire setup lock: %w", err) + } + // Held by another process: reclaim it if stale, otherwise wait and retry. + if info, statErr := os.Stat(lockPath); statErr == nil && time.Since(info.ModTime()) > setupLockStaleAfter { + log.Warnf(ctx, "reclaiming stale agent-shim setup lock at %s", lockPath) + _ = os.Remove(lockPath) + continue + } + log.Infof(ctx, "waiting for a concurrent agent-shim setup to finish") + select { + case <-ctx.Done(): + return nil, ctx.Err() + case <-time.After(time.Second): + } + } +} + +// launchUcodeAgent launches the ucode-configured agent, replacing this process. +// home is the already-resolved user home directory. +func launchUcodeAgent(ctx context.Context, home string, agent agentSpec, workspace string, agentArgs []string) error { ucodePath, err := exec.LookPath("ucode") if err != nil { return fmt.Errorf("ucode not found on PATH after setup: %w", err) } - contextArgs, err := injectAgentContext(home, agent) + contextArgs, err := injectAgentContext(ctx, home, agent) if err != nil { return err } @@ -357,8 +489,8 @@ func launchUcodeAgent(agent agentSpec, workspace string, agentArgs []string) err argv = append(argv, contextArgs...) argv = append(argv, agentArgs...) // Pass the session token as DATABRICKS_BEARER so stock ucode authenticates headlessly. - if os.Getenv("DATABRICKS_BEARER") == "" { - if token := os.Getenv("DATABRICKS_TOKEN"); token != "" { + if env.Get(ctx, "DATABRICKS_BEARER") == "" { + if token := env.Get(ctx, "DATABRICKS_TOKEN"); token != "" { _ = os.Setenv("DATABRICKS_BEARER", token) } } @@ -366,12 +498,18 @@ func launchUcodeAgent(agent agentSpec, workspace string, agentArgs []string) err } // injectAgentContext puts the Databricks context where the agent reads it, returning any extra argv. -func injectAgentContext(home string, agent agentSpec) ([]string, error) { - context := agentSystemContext(os.Getenv(workspaceHomeEnv)) +func injectAgentContext(ctx context.Context, home string, agent agentSpec) ([]string, error) { + systemContext := agentSystemContext(env.Get(ctx, workspaceHomeEnv)) switch { case agent.contextFlag != "": - f := filepath.Join(home, contextFile) - if err := os.WriteFile(f, []byte(context), 0o644); err != nil { + // Keep the scratch context file under the shim's own directory so it never + // clobbers an unrelated file the user happens to have in their home. + dir := filepath.Join(home, agentDir) + if err := os.MkdirAll(dir, 0o755); err != nil { + return nil, fmt.Errorf("failed to create %s: %w", dir, err) + } + f := filepath.Join(dir, contextFile) + if err := os.WriteFile(f, []byte(systemContext), 0o644); err != nil { return nil, fmt.Errorf("failed to write agent context file: %w", err) } return []string{agent.contextFlag, f}, nil @@ -385,7 +523,7 @@ func injectAgentContext(home string, agent agentSpec) ([]string, error) { if err := os.MkdirAll(filepath.Dir(f), 0o755); err != nil { return nil, fmt.Errorf("failed to create %s: %w", filepath.Dir(f), err) } - if err := os.WriteFile(f, []byte(context), 0o644); err != nil { + if err := os.WriteFile(f, []byte(systemContext), 0o644); err != nil { return nil, fmt.Errorf("failed to write %s: %w", f, err) } return nil, nil @@ -394,10 +532,172 @@ func injectAgentContext(home string, agent agentSpec) ([]string, error) { } } +// ensureNode downloads the latest Krypton LTS Node into deps/node once, returning +// its bin. Linux-only by design: the shim runs on the serverless driver, so the +// tarball name is hardcoded to linux while nodeDownloadArch guards the arch. +func ensureNode(ctx context.Context, home string) (string, error) { + depsRoot := filepath.Join(home, depsDir) + nodeDir := filepath.Join(depsRoot, "node") + nodeBin := filepath.Join(nodeDir, "bin") + if _, err := os.Stat(filepath.Join(nodeBin, "npm")); err == nil { + return nodeBin, nil + } + + arch := nodeDownloadArch(runtime.GOARCH) + if arch == "" { + return "", fmt.Errorf("unsupported architecture for Node download: %s", runtime.GOARCH) + } + const base = "https://nodejs.org/dist/latest-krypton" + tarName, wantSum, err := latestNodeTarball(ctx, base+"/SHASUMS256.txt", arch) + if err != nil { + return "", err + } + if err := os.MkdirAll(depsRoot, 0o755); err != nil { + return "", fmt.Errorf("failed to create %s: %w", depsRoot, err) + } + // Download to a temp file, verifying its SHA256 from SHASUMS256.txt before use. + tarball, err := downloadVerified(ctx, base+"/"+tarName, wantSum) + if err != nil { + return "", err + } + defer os.Remove(tarball) + // Extract into a sibling temp dir and atomically rename it into place, so an + // interrupted extraction never leaves a half-populated node dir that the + // os.Stat(npm) check above would then wrongly accept as a finished install. + tmpDir, err := os.MkdirTemp(depsRoot, "node-*") + if err != nil { + return "", fmt.Errorf("failed to create temp dir: %w", err) + } + defer os.RemoveAll(tmpDir) // no-op once renamed; cleans up a failed extraction + // Node's .tar.xz is the smallest download; extract it with the system tar. + if err := runCommand(ctx, "tar", "-xJf", tarball, "--strip-components=1", "-C", tmpDir); err != nil { + return "", fmt.Errorf("failed to extract Node.js: %w", err) + } + // Clear any partial leftover from a previously-interrupted run, then publish + // atomically (same-filesystem rename, since tmpDir is a sibling of nodeDir). + if err := os.RemoveAll(nodeDir); err != nil { + return "", fmt.Errorf("failed to remove %s: %w", nodeDir, err) + } + if err := os.Rename(tmpDir, nodeDir); err != nil { + return "", fmt.Errorf("failed to install Node.js: %w", err) + } + return nodeBin, nil +} + +// downloadVerified fetches url to a temp file, failing unless its SHA256 matches +// wantSum. The caller is responsible for removing the returned file. +func downloadVerified(ctx context.Context, url, wantSum string) (string, error) { + req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil) + if err != nil { + return "", err + } + resp, err := http.DefaultClient.Do(req) + if err != nil { + return "", fmt.Errorf("failed to download %s: %w", url, err) + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + return "", fmt.Errorf("failed to download %s: HTTP %d", url, resp.StatusCode) + } + + f, err := os.CreateTemp("", "node-*.tar.xz") + if err != nil { + return "", err + } + sum := sha256.New() + if _, err := io.Copy(io.MultiWriter(f, sum), resp.Body); err != nil { + f.Close() + os.Remove(f.Name()) + return "", fmt.Errorf("failed to download %s: %w", url, err) + } + if err := f.Close(); err != nil { + os.Remove(f.Name()) + return "", err + } + if got := hex.EncodeToString(sum.Sum(nil)); !strings.EqualFold(got, wantSum) { + os.Remove(f.Name()) + return "", fmt.Errorf("checksum mismatch for %s: got %s, want %s", url, got, wantSum) + } + return f.Name(), nil +} + +// nodeDownloadArch maps a Go arch to Node's release arch token, or "" if unsupported. +func nodeDownloadArch(goarch string) string { + switch goarch { + case "amd64": + return "x64" + case "arm64": + return "arm64" + default: + return "" + } +} + +// latestNodeTarball returns the linux tarball name and its SHA256 for arch from +// the SHASUMS256.txt listing (lines of " "). +func latestNodeTarball(ctx context.Context, shasumsURL, arch string) (name, sum string, err error) { + req, err := http.NewRequestWithContext(ctx, http.MethodGet, shasumsURL, nil) + if err != nil { + return "", "", err + } + resp, err := http.DefaultClient.Do(req) + if err != nil { + return "", "", fmt.Errorf("failed to fetch Node checksums: %w", err) + } + defer resp.Body.Close() + body, err := io.ReadAll(io.LimitReader(resp.Body, 1<<20)) + if err != nil { + return "", "", fmt.Errorf("failed to read Node checksums: %w", err) + } + re := regexp.MustCompile(`^([0-9a-f]{64})\s+(node-v[0-9.]+-linux-` + regexp.QuoteMeta(arch) + `\.tar\.xz)$`) + for _, line := range strings.Split(string(body), "\n") { + if m := re.FindStringSubmatch(strings.TrimSpace(line)); m != nil { + return m[2], m[1], nil + } + } + return "", "", fmt.Errorf("no linux-%s Node tarball found in %s", arch, shasumsURL) +} + +// disableNpmUpdateNotifier turns off npm's "new version available" box so it +// doesn't clutter the agent session. Best-effort: it's cosmetic, so a failure +// (e.g. npm not runnable) is logged and ignored rather than blocking the launch. +func disableNpmUpdateNotifier(ctx context.Context) { + if err := exec.CommandContext(ctx, "npm", "config", "set", "update-notifier", "false").Run(); err != nil { + log.Debugf(ctx, "failed to disable npm update-notifier: %v", err) + } +} + +// npmGlobalPrefix returns `npm prefix -g`, or "" if npm isn't runnable. +func npmGlobalPrefix(ctx context.Context) string { + out, err := exec.CommandContext(ctx, "npm", "prefix", "-g").Output() + if err != nil { + return "" + } + return strings.TrimSpace(string(out)) +} + +// runCommand runs name with the given args, inheriting stdio and the process env. +func runCommand(ctx context.Context, name string, args ...string) error { + cmd := exec.CommandContext(ctx, name, args...) + cmd.Stdin = os.Stdin + cmd.Stdout = os.Stdout + cmd.Stderr = os.Stderr + return cmd.Run() +} + +// runShell runs a shell snippet — for the curl|sh / curl|tar pipelines. +func runShell(ctx context.Context, script string) error { + cmd := exec.CommandContext(ctx, "sh", "-c", script) + cmd.Stdin = os.Stdin + cmd.Stdout = os.Stdout + cmd.Stderr = os.Stderr + return cmd.Run() +} + // prependPath puts dir at the front of $PATH, dropping any later duplicate. -func prependPath(dir string) { +func prependPath(ctx context.Context, dir string) { dirs := []string{dir} - for _, d := range filepath.SplitList(os.Getenv("PATH")) { + for _, d := range filepath.SplitList(env.Get(ctx, "PATH")) { if d != dir { dirs = append(dirs, d) } @@ -406,9 +706,9 @@ func prependPath(dir string) { } // removePath drops every occurrence of dir from $PATH. -func removePath(dir string) { +func removePath(ctx context.Context, dir string) { var dirs []string - for _, d := range filepath.SplitList(os.Getenv("PATH")) { + for _, d := range filepath.SplitList(env.Get(ctx, "PATH")) { if d != dir { dirs = append(dirs, d) } diff --git a/experimental/ssh/internal/client/agentshim_test.go b/experimental/ssh/internal/client/agentshim_test.go index 621e1ce1131..75b12dadd01 100644 --- a/experimental/ssh/internal/client/agentshim_test.go +++ b/experimental/ssh/internal/client/agentshim_test.go @@ -1,12 +1,15 @@ package client import ( + "crypto/sha256" + "encoding/hex" "net/http" "net/http/httptest" "os" "path/filepath" "strings" "testing" + "time" "github.com/databricks/databricks-sdk-go" "github.com/databricks/databricks-sdk-go/config" @@ -16,18 +19,90 @@ import ( func TestPrependPath(t *testing.T) { t.Setenv("PATH", strings.Join([]string{"/a", "/b"}, string(os.PathListSeparator))) - prependPath("/b") // existing entry moves to the front (deduped) + prependPath(t.Context(), "/b") // existing entry moves to the front (deduped) assert.Equal(t, []string{"/b", "/a"}, filepath.SplitList(os.Getenv("PATH"))) - prependPath("/new") + prependPath(t.Context(), "/new") assert.Equal(t, []string{"/new", "/b", "/a"}, filepath.SplitList(os.Getenv("PATH"))) } func TestRemovePath(t *testing.T) { t.Setenv("PATH", strings.Join([]string{"/a", "/b", "/a"}, string(os.PathListSeparator))) - removePath("/a") + removePath(t.Context(), "/a") assert.Equal(t, []string{"/b"}, filepath.SplitList(os.Getenv("PATH"))) } +func TestAcquireSetupLock(t *testing.T) { + home := t.TempDir() + lockPath := filepath.Join(home, agentDir, setupLockName) + + // Acquire creates the sentinel; release removes it. + unlock, err := acquireSetupLock(t.Context(), home) + require.NoError(t, err) + _, statErr := os.Stat(lockPath) + require.NoError(t, statErr, "lock sentinel should exist while held") + unlock() + _, statErr = os.Stat(lockPath) + assert.True(t, os.IsNotExist(statErr), "lock sentinel should be gone after release") + + // A lock left behind by a dead process is reclaimed once stale, not waited on + // forever. + require.NoError(t, os.WriteFile(lockPath, nil, 0o644)) + stale := time.Now().Add(-2 * setupLockStaleAfter) + require.NoError(t, os.Chtimes(lockPath, stale, stale)) + unlock2, err := acquireSetupLock(t.Context(), home) + require.NoError(t, err) + unlock2() +} + +func TestNodeDownloadArch(t *testing.T) { + assert.Equal(t, "x64", nodeDownloadArch("amd64")) + assert.Equal(t, "arm64", nodeDownloadArch("arm64")) + assert.Empty(t, nodeDownloadArch("mips")) +} + +func TestLatestNodeTarball(t *testing.T) { + gzSum, x64Sum, armSum := strings.Repeat("0", 64), strings.Repeat("a", 64), strings.Repeat("b", 64) + body := gzSum + " node-v24.1.0-linux-x64.tar.gz\n" + // .gz is skipped (want .xz) + x64Sum + " node-v24.1.0-linux-x64.tar.xz\n" + + armSum + " node-v24.1.0-linux-arm64.tar.xz\n" + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, _ = w.Write([]byte(body)) + })) + defer srv.Close() + + name, sum, err := latestNodeTarball(t.Context(), srv.URL, "x64") + require.NoError(t, err) + assert.Equal(t, "node-v24.1.0-linux-x64.tar.xz", name) + assert.Equal(t, x64Sum, sum) + + _, _, err = latestNodeTarball(t.Context(), srv.URL, "ppc64le") + assert.Error(t, err) +} + +func TestDownloadVerified(t *testing.T) { + payload := []byte("fake node tarball") + sum := sha256.Sum256(payload) + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, _ = w.Write(payload) + })) + defer srv.Close() + + t.Run("matching checksum writes the file", func(t *testing.T) { + path, err := downloadVerified(t.Context(), srv.URL, hex.EncodeToString(sum[:])) + require.NoError(t, err) + defer os.Remove(path) + got, err := os.ReadFile(path) + require.NoError(t, err) + assert.Equal(t, payload, got) + }) + + t.Run("mismatched checksum errors and leaves no file", func(t *testing.T) { + _, err := downloadVerified(t.Context(), srv.URL, strings.Repeat("0", 64)) + require.Error(t, err) + assert.Contains(t, err.Error(), "checksum mismatch") + }) +} + func TestSupportedAgents(t *testing.T) { names := SupportedAgentNames() // Limited to the agents whose ucode command accepts --workspace. @@ -57,7 +132,7 @@ func TestInjectAgentContext(t *testing.T) { t.Run("flag agent writes a scratch file and returns the flag", func(t *testing.T) { home := t.TempDir() - args, err := injectAgentContext(home, claude) + args, err := injectAgentContext(t.Context(), home, claude) require.NoError(t, err) require.Len(t, args, 2) assert.Equal(t, "--append-system-prompt-file", args[0]) @@ -68,7 +143,7 @@ func TestInjectAgentContext(t *testing.T) { t.Run("home-file agent writes the global instructions file, no args", func(t *testing.T) { home := t.TempDir() - args, err := injectAgentContext(home, codex) + args, err := injectAgentContext(t.Context(), home, codex) require.NoError(t, err) assert.Nil(t, args) data, err := os.ReadFile(filepath.Join(home, ".codex", "AGENTS.md")) @@ -81,7 +156,7 @@ func TestInjectAgentContext(t *testing.T) { target := filepath.Join(home, ".codex", "AGENTS.md") require.NoError(t, os.MkdirAll(filepath.Dir(target), 0o755)) require.NoError(t, os.WriteFile(target, []byte("user's own instructions"), 0o644)) - args, err := injectAgentContext(home, codex) + args, err := injectAgentContext(t.Context(), home, codex) require.NoError(t, err) assert.Nil(t, args) data, err := os.ReadFile(target) @@ -91,7 +166,7 @@ func TestInjectAgentContext(t *testing.T) { t.Run("agent without a mechanism writes nothing", func(t *testing.T) { home := t.TempDir() - args, err := injectAgentContext(home, noContext) + args, err := injectAgentContext(t.Context(), home, noContext) require.NoError(t, err) assert.Nil(t, args) entries, err := os.ReadDir(home) @@ -234,6 +309,16 @@ func TestProbeAIGateway(t *testing.T) { require.Error(t, err) assert.Contains(t, err.Error(), "not enabled on this workspace") }) + + t.Run("transient failure (5xx both) is retryable, not disabled", func(t *testing.T) { + unavailable := jsonHandler(http.StatusServiceUnavailable, `{"message":"try later"}`) + srv := newGatewayServer(t, unavailable, unavailable) + defer srv.Close() + err := probeAIGateway(t.Context(), newProbeClient(t, srv.URL)) + require.Error(t, err) + assert.Contains(t, err.Error(), "transient error") + assert.NotContains(t, err.Error(), "not enabled") + }) } func TestProbeModelServicesInconclusive(t *testing.T) { From d68088ea8ca407aa7e99cb7f98fd4bf5e97c39df Mon Sep 17 00:00:00 2001 From: Russell Clarey Date: Thu, 17 Sep 2026 13:43:45 +0000 Subject: [PATCH 3/6] cleanup --- experimental/ssh/cmd/agent_shim.go | 4 +- experimental/ssh/internal/client/agentshim.go | 194 +++++++----------- .../ssh/internal/client/agentshim_test.go | 1 - 3 files changed, 73 insertions(+), 126 deletions(-) diff --git a/experimental/ssh/cmd/agent_shim.go b/experimental/ssh/cmd/agent_shim.go index fedd1589c3b..4eb97189201 100644 --- a/experimental/ssh/cmd/agent_shim.go +++ b/experimental/ssh/cmd/agent_shim.go @@ -10,7 +10,7 @@ import ( func newAgentShimCommand() *cobra.Command { cmd := &cobra.Command{ Use: "agent-shim", - Short: "Launch a ucode-configured coding agent (invoked on the remote)", + Short: "Launch a coding agent configured against Unity AI Gateway (invoked on the remote)", Hidden: true, } for _, agent := range client.SupportedAgentNames() { @@ -22,7 +22,7 @@ func newAgentShimCommand() *cobra.Command { func newAgentShimAgentCommand(agent string) *cobra.Command { cmd := &cobra.Command{ Use: agent, - Short: "Launch the ucode-configured " + agent + " agent", + Short: "Launch " + agent + " configured against Unity AI Gateway", // Disable flag parsing: forward everything after the agent name to the agent verbatim. DisableFlagParsing: true, } diff --git a/experimental/ssh/internal/client/agentshim.go b/experimental/ssh/internal/client/agentshim.go index e5ed11eac13..ea0ad9580de 100644 --- a/experimental/ssh/internal/client/agentshim.go +++ b/experimental/ssh/internal/client/agentshim.go @@ -26,36 +26,35 @@ import ( ) const ( - // agentDir (home-relative) is the shim's working area. + // home-relative directory to hold files/dependencies for the shim. agentDir = ".agent-shim" - // shimDir holds the per-agent wrappers; must match remoteShimDir in client.go. - shimDir = agentDir + "/bin" + // directory which holds the per-agent wrappers; must match remoteShimDir in client.go. + binDir = agentDir + "/bin" - // ucodeRepo is the GitHub repo the shim installs ucode from. - ucodeRepo = "databricks/ucode" + // the GitHub repo the shim installs the Unity Gateway CLI from. + ugRepo = "databricks/unity-gateway" - // ucodeVersion pins the ucode release tag the shim installs, so a new upstream - // release can't change the toolchain out from under an SSH session unexpectedly. - ucodeVersion = "v0.1.0" + // pins the Unity Gateway CLI release tag the shim installs. + ugVersion = "v0.1.0" - // workspaceHomeEnv passes the user's workspace home to the shim for the agent context. + // environment variable used to pass the user's workspace home to the shim for the agent context. workspaceHomeEnv = "DATABRICKS_WORKSPACE_HOME" - // depsDir (home-relative) holds toolchain deps (Node/npm) fetched when the image ships none. + // directory which holds toolchain deps (Node/npm) fetched when the image ships none. depsDir = agentDir + "/deps" - // contextFile is the scratch file holding the Databricks session context. + // the file holding the Databricks session context. contextFile = "agent-system-context.md" - // setupLockName (under agentDir) serializes first-run toolchain setup across + // lock file to serialize first-run toolchain setup across // concurrent SSH clients sharing this driver's $HOME. setupLockName = ".setup.lock" - // setupLockStaleAfter reclaims a setup lock left behind by a process that died - // mid-install. It is deliberately generous: a cold first run downloads uv, - // ucode, and Node, which can legitimately take minutes. - setupLockStaleAfter = 10 * time.Minute + // timeout to reclaim a setup lock left behind by a process that died + // mid-install. It is deliberately generous: a cold first run can download uv, + // ug, and Node + setupLockStaleAfter = 2 * time.Minute ) type agentSpec struct { @@ -66,8 +65,7 @@ type agentSpec struct { contextHomeFile string } -// supportedAgents are the agents the shim launches — limited to those whose ucode -// command accepts --workspace (how we target the workspace headlessly). +// Currently limited to the agents whose Unity Gateway CLI command supports --workspace var supportedAgents = []agentSpec{ {name: "claude", contextFlag: "--append-system-prompt-file"}, {name: "codex", contextHomeFile: ".codex/AGENTS.md"}, @@ -100,20 +98,20 @@ func RunAgentShim(ctx context.Context, client *databricks.WorkspaceClient, agent if err := probeAIGateway(ctx, client); err != nil { return err } - // ucode targets this workspace via --workspace, so it configures without prompting. workspace := strings.TrimRight(client.Config.Host, "/") return bootstrapAndLaunchAgent(ctx, agent, workspace, agentArgs) } // --- Unity AI Gateway preflight --- // -// Kept in lockstep with ucode's probe_unity_gateway_capabilities -// (databricks/ucode, src/ucode/databricks.py): probe the Unity Catalog -// model-services API (v3) first, paging through it, then fall back to the legacy -// AI Gateway endpoints API (v2). A path counts only when its JSON body actually -// lists a usable resource — a 200 with an empty collection is "reachable but -// empty", not "enabled". We fail fast with the same actionable guidance ucode -// surfaces so the shim's preflight matches what ucode itself checks. +// Kept in lockstep with Unity Gateway CLI's probe_unity_gateway_capabilities +// (https://github.com/databricks/unity-gateway/blob/main/src/ucode/databricks.py) +// Probe the Unity Catalog model-services API (v3) first, paging through it, then +// fall back to the legacy AI Gateway endpoints API (v2). A path counts only when +// its JSON body actually lists a usable resource — a 200 with an empty collection +// is "reachable but empty", not "enabled". We fail fast with the same actionable +// guidance Unity Gateway CLI surfaces so the shim's preflight matches what Unity +// Gateway CLI itself checks. const ( modelServicesPath = "/api/2.1/unity-catalog/model-services" @@ -121,27 +119,19 @@ const ( modelServiceProbePageSize = 50 modelServiceProbeMaxPages = 20 aiGatewayDocsURL = "https://docs.databricks.com/aws/en/ai-gateway/overview-beta" - - // modelServiceEmptyDetail matches ucode's wording for a reachable model-services - // API that lists nothing the caller can use — almost always a UC grant gap. modelServiceEmptyDetail = "reachable, no accessible model services returned; " + "check USE CATALOG on system, and USE SCHEMA and EXECUTE on system.ai" ) -// gatewayProbe is the outcome of probing one AI Gateway API. resourceAvailable -// means the API returned at least one usable resource; conclusive is false when -// paging couldn't be completed, so "no resources" was never actually confirmed. type gatewayProbe struct { reachable bool detail string + // true if the API returned at least one usable resource resourceAvailable bool + // false when paging couldn't be completed, so "no resources" was never actually confirmed conclusive bool } -// probeAIGateway fails fast if the workspace's Unity AI Gateway can't run an -// agent. It returns nil when the gateway is usable — logging a warning when it's -// reachable but exposes no model services to this caller — and an actionable -// error otherwise. func probeAIGateway(ctx context.Context, client *databricks.WorkspaceClient) error { host := strings.TrimRight(client.Config.Host, "/") @@ -185,8 +175,6 @@ func probeAIGateway(ctx context.Context, client *databricks.WorkspaceClient) err } } -// probeModelServices probes the UC model-services API (v3), paging until it finds -// an accessible model service, exhausts the cursor, or hits the page cap. func probeModelServices(ctx context.Context, client *databricks.WorkspaceClient, host string) gatewayProbe { pageToken := "" for page := 0; page < modelServiceProbeMaxPages; page++ { @@ -215,7 +203,6 @@ func probeModelServices(ctx context.Context, client *databricks.WorkspaceClient, return gatewayProbe{reachable: true, detail: "reachable"} } -// probeLegacyEndpoints probes the legacy AI Gateway endpoints API (v2). func probeLegacyEndpoints(ctx context.Context, client *databricks.WorkspaceClient, host string) gatewayProbe { payload, reason := gatewayGetJSON(ctx, client, host+legacyEndpointsPath+"?page_size=1") if payload == nil { @@ -227,10 +214,6 @@ func probeLegacyEndpoints(ctx context.Context, client *databricks.WorkspaceClien return gatewayProbe{reachable: true, detail: "reachable, no accessible endpoints returned", conclusive: true} } -// gatewayGetJSON issues an authenticated GET expecting JSON. It returns the -// decoded payload on HTTP 200, or (nil, reason) on any failure — mirroring -// ucode's _http_get_json, including a short response-body excerpt so gateway auth -// errors (e.g. a 400 whose body is "Invalid Token") stay visible in the reason. func gatewayGetJSON(ctx context.Context, client *databricks.WorkspaceClient, reqURL string) (any, string) { req, err := http.NewRequestWithContext(ctx, http.MethodGet, reqURL, nil) if err != nil { @@ -274,16 +257,15 @@ var ( gatewayDetailV2 = regexp.MustCompile(`(?i)\bv2\b`) ) -// versionNeutralGatewayDetail rewrites the internal v2/v3 API labels in a failure -// reason into user-facing terms, matching ucode so the error text stays in sync. +// rewrite the internal v2/v3 API labels in a failure reason into user-facing terms func versionNeutralGatewayDetail(detail string) string { detail = gatewayDetailV3.ReplaceAllString(detail, "model service") return gatewayDetailV2.ReplaceAllString(detail, "legacy endpoint") } -// looksLikeDefinitiveAuthFailure is true when retrying another workspace API -// can't rescue the token. A bare 403 is left out: it can be endpoint-specific -// authorization, so the preflight still tries the fallback before giving up. +// returns true when retrying another workspace API can't rescue the token. +// A bare 403 is left out: it can be endpoint-specific authorization, so the +// preflight still tries the fallback before giving up. func looksLikeDefinitiveAuthFailure(reason string) bool { if strings.Contains(reason, "HTTP 401") { return true @@ -291,8 +273,8 @@ func looksLikeDefinitiveAuthFailure(reason string) bool { return strings.Contains(reason, "HTTP 400") && strings.Contains(strings.ToLower(reason), "invalid token") } -// looksLikeScopeFailure matches a 403 that reports the OAuth token is missing a -// required scope (which re-login can fix), as opposed to a plain permission 403. +// matches a 403 that reports the OAuth token is missing a required scope (which +// re-login can fix), as opposed to a plain permission 403. func looksLikeScopeFailure(reason string) bool { l := strings.ToLower(reason) return strings.Contains(l, "http 403") && strings.Contains(l, "oauth token") && strings.Contains(l, "required scopes") @@ -302,11 +284,10 @@ func looksLikePermissionFailure(reason string) bool { return strings.Contains(reason, "HTTP 403") } -// gatewayTransientStatus matches the HTTP statuses worth retrying (429 + any 5xx). var gatewayTransientStatus = regexp.MustCompile(`HTTP (429|5\d\d)`) -// looksLikeTransient is true for a probe failure that is likely temporary (rate -// limit, server error, or a network/transport error) rather than a stable +// returns true for a probe failure that is likely temporary (rate limit, +// server error, or a network/transport error) rather than a stable // "disabled"/"unauthorized" verdict. func looksLikeTransient(reason string) bool { return strings.HasPrefix(reason, "network error") || gatewayTransientStatus.MatchString(reason) @@ -320,8 +301,6 @@ func aiGatewayScopeError(host, reason string) error { return fmt.Errorf("the access token for %s is missing an OAuth scope required by the AI Gateway APIs (%s). Re-authenticate to mint a token with the needed scopes:\n databricks auth login --host %s", host, reason, host) } -// hasNonEmptyCollection reports whether payload is a JSON object with a non-empty -// array at key. func hasNonEmptyCollection(payload any, key string) bool { obj, ok := payload.(map[string]any) if !ok { @@ -331,8 +310,6 @@ func hasNonEmptyCollection(payload any, key string) bool { return ok && len(arr) > 0 } -// stringField returns payload[key] when payload is a JSON object holding a string -// there, else "". func stringField(payload any, key string) string { obj, ok := payload.(map[string]any) if !ok { @@ -342,6 +319,8 @@ func stringField(payload any, key string) string { return s } +// --- end Unity AI Gateway preflight --- + func bootstrapAndLaunchAgent(ctx context.Context, agent agentSpec, workspace string, agentArgs []string) error { home, err := env.UserHomeDir(ctx) if err != nil { @@ -351,11 +330,11 @@ func bootstrapAndLaunchAgent(ctx context.Context, agent agentSpec, workspace str // Reconstruct the PATH tooling runs under on every launch rather than caching // it: every entry is a known location, so deriving it can't restore a stale // path (e.g. a versioned CLI dir from an older session). Drop the shim dir so - // ucode execs the real agent rather than this wrapper, then prepend uv/ucode's - // bin, this databricks CLI's own dir (so ucode finds it and skips its own - // install), and the Node bin ensureNode installs into. - removePath(ctx, filepath.Join(home, shimDir)) - prependPath(ctx, filepath.Join(home, ".local", "bin")) + // Unity Gateway CLI execs the real agent rather than this wrapper, then prepend + // uv/ug's bin, this databricks CLI's own dir (so it's used by ug/the spawned agent), + // and the Node bin ensureNode installs into. + removePath(ctx, filepath.Join(home, binDir)) + prependPath(ctx, filepath.Join(home, ".local", "bin")) // uv if self, err := os.Executable(); err == nil { prependPath(ctx, filepath.Dir(self)) } @@ -370,27 +349,20 @@ func bootstrapAndLaunchAgent(ctx context.Context, agent agentSpec, workspace str disableNpmUpdateNotifier(ctx) // Put npm's global bin on PATH so an npm-installed agent binary resolves after - // ucode installs it. + // Unity Gateway CLI installs it. if prefix := npmGlobalPrefix(ctx); prefix != "" { prependPath(ctx, filepath.Join(prefix, "bin")) } - return launchUcodeAgent(ctx, home, agent, workspace, agentArgs) + return launchAgent(ctx, home, agent, workspace, agentArgs) } -// toolchainReady reports whether the agent toolchain is already installed, derived -// purely from PATH (which the caller has already set up). ucode launches the agent -// and npm provides the Node runtime agents like Claude Code need. func toolchainReady() bool { - _, ucodeErr := exec.LookPath("ucode") + _, ugErr := exec.LookPath("ucode") _, npmErr := exec.LookPath("npm") - return ucodeErr == nil && npmErr == nil + return ugErr == nil && npmErr == nil } -// ensureToolchain installs uv, ucode, and Node/npm when missing. It is a no-op -// once the toolchain is present, and holds a machine-local lock while installing -// so two first-run SSH clients sharing this $HOME don't concurrently install into -// (and corrupt) the same paths. func ensureToolchain(ctx context.Context, home string) error { if toolchainReady() { return nil @@ -404,11 +376,7 @@ func ensureToolchain(ctx context.Context, home string) error { if toolchainReady() { return nil } - return installToolchain(ctx, home) -} -// installToolchain performs the actual installs. The caller holds the setup lock. -func installToolchain(ctx context.Context, home string) error { // 1. uv (installs into ~/.local/bin). if _, err := exec.LookPath("uv"); err != nil { cmdio.LogString(ctx, "Installing uv...") @@ -417,16 +385,15 @@ func installToolchain(ctx context.Context, home string) error { } } - // 2. ucode (pinned stock upstream release). + // 2. Unity Gateway CLI (pinned stock upstream release). if _, err := exec.LookPath("ucode"); err != nil { - cmdio.LogString(ctx, "Installing ucode...") - if err := runCommand(ctx, "uv", "tool", "install", "git+https://github.com/"+ucodeRepo+"@"+ucodeVersion); err != nil { - return fmt.Errorf("failed to install ucode: %w", err) + cmdio.LogString(ctx, "Installing Unity Gateway CLI...") + if err := runCommand(ctx, "uv", "tool", "install", "git+https://github.com/"+ugRepo+"@"+ugVersion); err != nil { + return fmt.Errorf("failed to install Unity Gateway CLI: %w", err) } } - // 3. Node/npm. Some agents (e.g. Claude Code) need it and serverless images - // don't ship it. It installs into depsDir/node/bin, already prepended to PATH. + // 3. Node/npm (installs into depsDir/node/bin) if _, err := exec.LookPath("npm"); err != nil { cmdio.LogString(ctx, "Installing npm...") if _, err := ensureNode(ctx, home); err != nil { @@ -437,10 +404,9 @@ func installToolchain(ctx context.Context, home string) error { return nil } -// acquireSetupLock takes a machine-local lock (an O_EXCL sentinel file, matching -// the pattern in libs/versioncheck) serializing first-run setup across SSH -// clients. It blocks until the lock is free, reclaiming one left behind by a dead -// process after setupLockStaleAfter. The returned func releases the lock. +// take a machine-local lock (an O_EXCL sentinel file) serializing first-run setup +// across SSH clients. It blocks until the lock is free, reclaiming one left behind +// by a dead process after setupLockStaleAfter. The returned func releases the lock. func acquireSetupLock(ctx context.Context, home string) (func(), error) { lockPath := filepath.Join(home, agentDir, setupLockName) if err := os.MkdirAll(filepath.Dir(lockPath), 0o755); err != nil { @@ -461,7 +427,7 @@ func acquireSetupLock(ctx context.Context, home string) (func(), error) { _ = os.Remove(lockPath) continue } - log.Infof(ctx, "waiting for a concurrent agent-shim setup to finish") + log.Infof(ctx, "waiting for a concurrent setup to finish") select { case <-ctx.Done(): return nil, ctx.Err() @@ -470,12 +436,10 @@ func acquireSetupLock(ctx context.Context, home string) (func(), error) { } } -// launchUcodeAgent launches the ucode-configured agent, replacing this process. -// home is the already-resolved user home directory. -func launchUcodeAgent(ctx context.Context, home string, agent agentSpec, workspace string, agentArgs []string) error { - ucodePath, err := exec.LookPath("ucode") +func launchAgent(ctx context.Context, home string, agent agentSpec, workspace string, agentArgs []string) error { + ugPath, err := exec.LookPath("ucode") if err != nil { - return fmt.Errorf("ucode not found on PATH after setup: %w", err) + return fmt.Errorf("Unity Gateway CLI not found on PATH after setup: %w", err) } contextArgs, err := injectAgentContext(ctx, home, agent) if err != nil { @@ -483,21 +447,20 @@ func launchUcodeAgent(ctx context.Context, home string, agent agentSpec, workspa } argv := []string{"ucode", agent.name} if workspace != "" { - // Every supported agent accepts --workspace (the gate for supportedAgents). argv = append(argv, "--workspace", workspace) } argv = append(argv, contextArgs...) argv = append(argv, agentArgs...) - // Pass the session token as DATABRICKS_BEARER so stock ucode authenticates headlessly. + // Pass the session token as DATABRICKS_BEARER so Unity Gateway CLI authenticates headlessly. if env.Get(ctx, "DATABRICKS_BEARER") == "" { if token := env.Get(ctx, "DATABRICKS_TOKEN"); token != "" { _ = os.Setenv("DATABRICKS_BEARER", token) } } - return execProcess(ucodePath, argv, os.Environ()) + return execProcess(ugPath, argv, os.Environ()) } -// injectAgentContext puts the Databricks context where the agent reads it, returning any extra argv. +// put the Databricks context where the agent reads it, returning any extra argv. func injectAgentContext(ctx context.Context, home string, agent agentSpec) ([]string, error) { systemContext := agentSystemContext(env.Get(ctx, workspaceHomeEnv)) switch { @@ -532,9 +495,9 @@ func injectAgentContext(ctx context.Context, home string, agent agentSpec) ([]st } } -// ensureNode downloads the latest Krypton LTS Node into deps/node once, returning -// its bin. Linux-only by design: the shim runs on the serverless driver, so the -// tarball name is hardcoded to linux while nodeDownloadArch guards the arch. +// download the latest Krypton LTS Node into deps/node once, returning its bin. +// Linux-only by design: the shim runs on the serverless driver, so the tarball +// name is hardcoded to linux while nodeDownloadArch guards the arch. func ensureNode(ctx context.Context, home string) (string, error) { depsRoot := filepath.Join(home, depsDir) nodeDir := filepath.Join(depsRoot, "node") @@ -584,8 +547,8 @@ func ensureNode(ctx context.Context, home string) (string, error) { return nodeBin, nil } -// downloadVerified fetches url to a temp file, failing unless its SHA256 matches -// wantSum. The caller is responsible for removing the returned file. +// fetch url to a temp file, failing unless its SHA256 matches wantSum. The caller +// is responsible for removing the returned file. func downloadVerified(ctx context.Context, url, wantSum string) (string, error) { req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil) if err != nil { @@ -621,7 +584,6 @@ func downloadVerified(ctx context.Context, url, wantSum string) (string, error) return f.Name(), nil } -// nodeDownloadArch maps a Go arch to Node's release arch token, or "" if unsupported. func nodeDownloadArch(goarch string) string { switch goarch { case "amd64": @@ -633,8 +595,6 @@ func nodeDownloadArch(goarch string) string { } } -// latestNodeTarball returns the linux tarball name and its SHA256 for arch from -// the SHASUMS256.txt listing (lines of " "). func latestNodeTarball(ctx context.Context, shasumsURL, arch string) (name, sum string, err error) { req, err := http.NewRequestWithContext(ctx, http.MethodGet, shasumsURL, nil) if err != nil { @@ -658,16 +618,15 @@ func latestNodeTarball(ctx context.Context, shasumsURL, arch string) (name, sum return "", "", fmt.Errorf("no linux-%s Node tarball found in %s", arch, shasumsURL) } -// disableNpmUpdateNotifier turns off npm's "new version available" box so it -// doesn't clutter the agent session. Best-effort: it's cosmetic, so a failure -// (e.g. npm not runnable) is logged and ignored rather than blocking the launch. +// turn off npm's "new version available" box so it doesn't clutter the installation +// progress. Best-effort: it's cosmetic, so a failure (e.g. npm not runnable) is +// logged and ignored rather than blocking the launch. func disableNpmUpdateNotifier(ctx context.Context) { if err := exec.CommandContext(ctx, "npm", "config", "set", "update-notifier", "false").Run(); err != nil { log.Debugf(ctx, "failed to disable npm update-notifier: %v", err) } } -// npmGlobalPrefix returns `npm prefix -g`, or "" if npm isn't runnable. func npmGlobalPrefix(ctx context.Context) string { out, err := exec.CommandContext(ctx, "npm", "prefix", "-g").Output() if err != nil { @@ -676,7 +635,6 @@ func npmGlobalPrefix(ctx context.Context) string { return strings.TrimSpace(string(out)) } -// runCommand runs name with the given args, inheriting stdio and the process env. func runCommand(ctx context.Context, name string, args ...string) error { cmd := exec.CommandContext(ctx, name, args...) cmd.Stdin = os.Stdin @@ -685,7 +643,6 @@ func runCommand(ctx context.Context, name string, args ...string) error { return cmd.Run() } -// runShell runs a shell snippet — for the curl|sh / curl|tar pipelines. func runShell(ctx context.Context, script string) error { cmd := exec.CommandContext(ctx, "sh", "-c", script) cmd.Stdin = os.Stdin @@ -694,7 +651,6 @@ func runShell(ctx context.Context, script string) error { return cmd.Run() } -// prependPath puts dir at the front of $PATH, dropping any later duplicate. func prependPath(ctx context.Context, dir string) { dirs := []string{dir} for _, d := range filepath.SplitList(env.Get(ctx, "PATH")) { @@ -705,7 +661,6 @@ func prependPath(ctx context.Context, dir string) { _ = os.Setenv("PATH", strings.Join(dirs, string(os.PathListSeparator))) } -// removePath drops every occurrence of dir from $PATH. func removePath(ctx context.Context, dir string) { var dirs []string for _, d := range filepath.SplitList(env.Get(ctx, "PATH")) { @@ -716,7 +671,6 @@ func removePath(ctx context.Context, dir string) { _ = os.Setenv("PATH", strings.Join(dirs, string(os.PathListSeparator))) } -// agentSystemContext returns the Databricks session context; wsHome names the working directory. func agentSystemContext(wsHome string) string { cwd := "the user's Databricks workspace home directory" if wsHome != "" { @@ -724,13 +678,7 @@ func agentSystemContext(wsHome string) string { } return fmt.Sprintf(`You are running inside a "databricks ssh connect" session on the driver node of a Databricks serverless cluster. -- The "databricks" CLI is installed and already authenticated: DATABRICKS_HOST and - DATABRICKS_TOKEN are set in the environment, so "databricks ..." commands work with no - "databricks auth login". The same token governs Unity Catalog and serving-endpoint access. -- This container is ephemeral; only paths under /Workspace persist across sessions. Your - working directory is %s. -- DATABRICKS_TOKEN is a static session token that may expire during a long session; if - "databricks" calls start failing with auth errors, the session likely needs reconnecting. -- You can run shell commands and use the "databricks" CLI to explore the workspace - (clusters, jobs, Unity Catalog, DBFS, etc.).`, cwd) +- The "databricks" CLI is installed and already authenticated: DATABRICKS_HOST and DATABRICKS_TOKEN are set in the environment, so "databricks ..." commands work with no "databricks auth login". The same token governs Unity Catalog and serving-endpoint access. +- This container is ephemeral; only paths under /Workspace and /Volumes persist across sessions. Your working directory is %s. +- You can run shell commands and use the "databricks" CLI to explore the workspace.`, cwd) } diff --git a/experimental/ssh/internal/client/agentshim_test.go b/experimental/ssh/internal/client/agentshim_test.go index 75b12dadd01..36369c9bdc3 100644 --- a/experimental/ssh/internal/client/agentshim_test.go +++ b/experimental/ssh/internal/client/agentshim_test.go @@ -105,7 +105,6 @@ func TestDownloadVerified(t *testing.T) { func TestSupportedAgents(t *testing.T) { names := SupportedAgentNames() - // Limited to the agents whose ucode command accepts --workspace. assert.Equal(t, []string{"claude", "codex"}, names) // Each agent injects context by exactly one mechanism (flag OR home file); From 174c5e5b990c9ca816ec12f4dcb3b295e9a91880 Mon Sep 17 00:00:00 2001 From: Russell Clarey Date: Thu, 17 Sep 2026 14:34:40 +0000 Subject: [PATCH 4/6] lint --- experimental/ssh/internal/client/agentshim.go | 14 +++++++------- experimental/ssh/internal/client/agentshim_test.go | 3 ++- 2 files changed, 9 insertions(+), 8 deletions(-) diff --git a/experimental/ssh/internal/client/agentshim.go b/experimental/ssh/internal/client/agentshim.go index ea0ad9580de..aa6f461dc51 100644 --- a/experimental/ssh/internal/client/agentshim.go +++ b/experimental/ssh/internal/client/agentshim.go @@ -119,17 +119,17 @@ const ( modelServiceProbePageSize = 50 modelServiceProbeMaxPages = 20 aiGatewayDocsURL = "https://docs.databricks.com/aws/en/ai-gateway/overview-beta" - modelServiceEmptyDetail = "reachable, no accessible model services returned; " + + modelServiceEmptyDetail = "reachable, no accessible model services returned; " + "check USE CATALOG on system, and USE SCHEMA and EXECUTE on system.ai" ) type gatewayProbe struct { - reachable bool - detail string + reachable bool + detail string // true if the API returned at least one usable resource resourceAvailable bool // false when paging couldn't be completed, so "no resources" was never actually confirmed - conclusive bool + conclusive bool } func probeAIGateway(ctx context.Context, client *databricks.WorkspaceClient) error { @@ -177,7 +177,7 @@ func probeAIGateway(ctx context.Context, client *databricks.WorkspaceClient) err func probeModelServices(ctx context.Context, client *databricks.WorkspaceClient, host string) gatewayProbe { pageToken := "" - for page := 0; page < modelServiceProbeMaxPages; page++ { + for page := range modelServiceProbeMaxPages { reqURL := fmt.Sprintf("%s%s?page_size=%d", host, modelServicesPath, modelServiceProbePageSize) if pageToken != "" { reqURL += "&page_token=" + url.QueryEscape(pageToken) @@ -439,7 +439,7 @@ func acquireSetupLock(ctx context.Context, home string) (func(), error) { func launchAgent(ctx context.Context, home string, agent agentSpec, workspace string, agentArgs []string) error { ugPath, err := exec.LookPath("ucode") if err != nil { - return fmt.Errorf("Unity Gateway CLI not found on PATH after setup: %w", err) + return fmt.Errorf("the Unity Gateway CLI was not found on PATH after setup: %w", err) } contextArgs, err := injectAgentContext(ctx, home, agent) if err != nil { @@ -610,7 +610,7 @@ func latestNodeTarball(ctx context.Context, shasumsURL, arch string) (name, sum return "", "", fmt.Errorf("failed to read Node checksums: %w", err) } re := regexp.MustCompile(`^([0-9a-f]{64})\s+(node-v[0-9.]+-linux-` + regexp.QuoteMeta(arch) + `\.tar\.xz)$`) - for _, line := range strings.Split(string(body), "\n") { + for line := range strings.SplitSeq(string(body), "\n") { if m := re.FindStringSubmatch(strings.TrimSpace(line)); m != nil { return m[2], m[1], nil } diff --git a/experimental/ssh/internal/client/agentshim_test.go b/experimental/ssh/internal/client/agentshim_test.go index 36369c9bdc3..658f73b1a66 100644 --- a/experimental/ssh/internal/client/agentshim_test.go +++ b/experimental/ssh/internal/client/agentshim_test.go @@ -3,6 +3,7 @@ package client import ( "crypto/sha256" "encoding/hex" + "io/fs" "net/http" "net/http/httptest" "os" @@ -42,7 +43,7 @@ func TestAcquireSetupLock(t *testing.T) { require.NoError(t, statErr, "lock sentinel should exist while held") unlock() _, statErr = os.Stat(lockPath) - assert.True(t, os.IsNotExist(statErr), "lock sentinel should be gone after release") + assert.ErrorIs(t, statErr, fs.ErrNotExist, "lock sentinel should be gone after release") // A lock left behind by a dead process is reclaimed once stale, not waited on // forever. From f58ffa30a03b7655c3bb7ee2c250864d587e5aca Mon Sep 17 00:00:00 2001 From: Russell Clarey Date: Wed, 23 Sep 2026 10:55:32 +0000 Subject: [PATCH 5/6] use SDK --- experimental/ssh/internal/client/agentshim.go | 224 +++++++++--------- .../ssh/internal/client/agentshim_test.go | 28 ++- 2 files changed, 129 insertions(+), 123 deletions(-) diff --git a/experimental/ssh/internal/client/agentshim.go b/experimental/ssh/internal/client/agentshim.go index aa6f461dc51..5dd6149d829 100644 --- a/experimental/ssh/internal/client/agentshim.go +++ b/experimental/ssh/internal/client/agentshim.go @@ -4,12 +4,10 @@ import ( "context" "crypto/sha256" "encoding/hex" - "encoding/json" "errors" "fmt" "io" "net/http" - "net/url" "os" "os/exec" "path/filepath" @@ -23,6 +21,12 @@ import ( "github.com/databricks/cli/libs/env" "github.com/databricks/cli/libs/log" "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/apierr" + "github.com/databricks/databricks-sdk-go/client" + "github.com/databricks/databricks-sdk-go/config" + "github.com/databricks/databricks-sdk-go/httpclient" + "github.com/databricks/databricks-sdk-go/listing" + "github.com/databricks/databricks-sdk-go/service/catalog" ) const ( @@ -114,142 +118,126 @@ func RunAgentShim(ctx context.Context, client *databricks.WorkspaceClient, agent // Gateway CLI itself checks. const ( - modelServicesPath = "/api/2.1/unity-catalog/model-services" legacyEndpointsPath = "/api/ai-gateway/v2/endpoints" modelServiceProbePageSize = 50 modelServiceProbeMaxPages = 20 + modelServiceProbeMaxItems = modelServiceProbePageSize * modelServiceProbeMaxPages aiGatewayDocsURL = "https://docs.databricks.com/aws/en/ai-gateway/overview-beta" - modelServiceEmptyDetail = "reachable, no accessible model services returned; " + - "check USE CATALOG on system, and USE SCHEMA and EXECUTE on system.ai" ) type gatewayProbe struct { reachable bool - detail string + err error // true if the API returned at least one usable resource resourceAvailable bool - // false when paging couldn't be completed, so "no resources" was never actually confirmed - conclusive bool } -func probeAIGateway(ctx context.Context, client *databricks.WorkspaceClient) error { - host := strings.TrimRight(client.Config.Host, "/") +func probeAIGateway(ctx context.Context, wsclient *databricks.WorkspaceClient) error { + host := strings.TrimRight(wsclient.Config.Host, "/") - modelSvc := probeModelServices(ctx, client, host) + // build a client with retries disabled so a transient 429/503 fails fast + // instead of looping on the SDK's default 5-minute retry budget before the + // slow first-run bootstrap even starts. + apiClient, err := newProbeAPIClient(wsclient.Config) + if err != nil { + // this probe is best effort, so don't block + return nil + } + + modelSvc := probeModelServices(ctx, apiClient) // A 401 (or a 400 "invalid token") can't be rescued by trying another API, // so surface it before the fallback probe. - if !modelSvc.reachable && looksLikeDefinitiveAuthFailure(modelSvc.detail) { - return aiGatewayAuthError(host, modelSvc.detail) + if !modelSvc.reachable && looksLikeDefinitiveAuthFailure(modelSvc.err) { + return aiGatewayAuthError(host, modelSvc.err.Error()) } if modelSvc.resourceAvailable { return nil } - legacy := probeLegacyEndpoints(ctx, client, host) + legacy := probeLegacyEndpoints(ctx, apiClient) switch { case legacy.reachable: // The legacy endpoints API answered, so the gateway is enabled even if no // model services are visible to this caller. - log.Warnf(ctx, "Unity AI Gateway model service check: %s", modelSvc.detail) + if modelSvc.err != nil { + log.Warnf(ctx, "Unity AI Gateway model service check: %s", modelSvc.err) + } return nil - case modelSvc.reachable && !modelSvc.conclusive: - // v3 answered but couldn't be paged to completion; give it the benefit of - // the doubt rather than blocking on an unconfirmed "empty". + case modelSvc.reachable: + // v3 answered; give it the benefit of the doubt rather than blocking on an + // unconfirmed "empty". return nil - case looksLikeDefinitiveAuthFailure(legacy.detail): - return aiGatewayAuthError(host, legacy.detail) - case looksLikeScopeFailure(modelSvc.detail): - return aiGatewayScopeError(host, modelSvc.detail) - case looksLikeScopeFailure(legacy.detail): - return aiGatewayScopeError(host, legacy.detail) - case looksLikeTransient(modelSvc.detail) || looksLikeTransient(legacy.detail): + case looksLikeDefinitiveAuthFailure(legacy.err): + return aiGatewayAuthError(host, legacy.err.Error()) + case looksLikeScopeFailure(modelSvc.err): + return aiGatewayScopeError(host, modelSvc.err.Error()) + case looksLikeScopeFailure(legacy.err): + return aiGatewayScopeError(host, legacy.err.Error()) + case looksLikeTransient(modelSvc.err) || looksLikeTransient(legacy.err): // A rate-limit/5xx/network blip is not "disabled" — tell the user to retry // rather than sending them to the enablement docs. - return fmt.Errorf("could not verify the Databricks Unity AI Gateway on %s: the probe hit a transient error (model services: %s; legacy endpoints: %s). Retry in a moment", host, modelSvc.detail, legacy.detail) - case looksLikePermissionFailure(modelSvc.detail): - return fmt.Errorf("model service access could not be verified on %s (%s). The legacy endpoint fallback also failed (%s). The model service probe requires permission to list Unity Catalog model services. Verify USE CATALOG on `system`, and USE SCHEMA and EXECUTE on `system.ai`", host, modelSvc.detail, legacy.detail) - case looksLikePermissionFailure(legacy.detail): - return fmt.Errorf("legacy endpoint access could not be verified on %s (%s). The model service probe also failed (%s). Verify the caller's workspace permissions for the legacy endpoints listing", host, legacy.detail, modelSvc.detail) + return versionNeutralGatewayError( + fmt.Sprintf("could not verify the Databricks Unity AI Gateway on %s: the probe hit a transient error (model services: %s; legacy endpoints: %s). Retry in a moment", host, modelSvc.err.Error(), legacy.err.Error()), + ) + case errors.Is(modelSvc.err, apierr.ErrPermissionDenied): + return versionNeutralGatewayError( + fmt.Sprintf("model service access could not be verified on %s (%s). The legacy endpoint fallback also failed (%s). The model service probe requires permission to list Unity Catalog model services. Verify USE CATALOG on `system`, and USE SCHEMA and EXECUTE on `system.ai`", host, modelSvc.err.Error(), legacy.err.Error()), + ) + case errors.Is(legacy.err, apierr.ErrPermissionDenied): + return versionNeutralGatewayError( + fmt.Sprintf("legacy endpoint access could not be verified on %s (%s). The model service probe also failed (%s). Verify the caller's workspace permissions for the legacy endpoints listing", host, legacy.err.Error(), modelSvc.err.Error()), + ) default: - return fmt.Errorf("the Databricks Unity AI Gateway is not enabled on this workspace (%s): neither model services (%s) nor legacy endpoints (%s) are available. See %s", host, modelSvc.detail, legacy.detail, aiGatewayDocsURL) + return versionNeutralGatewayError( + fmt.Sprintf("the Databricks Unity AI Gateway is not enabled on this workspace (%s): neither model services (%s) nor legacy endpoints (%s) are available. See %s", host, modelSvc.err.Error(), legacy.err.Error(), aiGatewayDocsURL), + ) } } -func probeModelServices(ctx context.Context, client *databricks.WorkspaceClient, host string) gatewayProbe { - pageToken := "" - for page := range modelServiceProbeMaxPages { - reqURL := fmt.Sprintf("%s%s?page_size=%d", host, modelServicesPath, modelServiceProbePageSize) - if pageToken != "" { - reqURL += "&page_token=" + url.QueryEscape(pageToken) - } - payload, reason := gatewayGetJSON(ctx, client, reqURL) - if payload == nil { - // A first-page failure is the real reachability signal; a later page - // failing still means the API answered at least once, so treat it as - // reachable-but-inconclusive rather than a confirmed "empty". - if page == 0 { - return gatewayProbe{detail: versionNeutralGatewayDetail(reason), conclusive: true} - } - return gatewayProbe{reachable: true, detail: "reachable"} - } - if hasNonEmptyCollection(payload, "model_services") { - return gatewayProbe{reachable: true, detail: "reachable, accessible model service returned", resourceAvailable: true, conclusive: true} - } - pageToken = stringField(payload, "next_page_token") - if pageToken == "" { - return gatewayProbe{reachable: true, detail: modelServiceEmptyDetail, conclusive: true} - } +// newProbeAPIClient builds a low-level client from cfg with the SDK's retries +// turned off, so the gateway preflight is one-shot. Clearing ErrorRetriable and +// TransientErrors makes a retriable status (429, 503, ...) surface on the first +// attempt rather than looping until the retry budget expires; it mirrors +// client.New, which builds the retrying client the rest of the CLI uses. +func newProbeAPIClient(cfg *config.Config) (*client.DatabricksClient, error) { + clientCfg, err := config.HTTPClientConfigFromConfig(cfg) + if err != nil { + return nil, err } - return gatewayProbe{reachable: true, detail: "reachable"} + clientCfg.ErrorRetriable = func(context.Context, error) bool { return false } + clientCfg.TransientErrors = nil + return client.NewWithClient(cfg, httpclient.NewApiClient(clientCfg)) } -func probeLegacyEndpoints(ctx context.Context, client *databricks.WorkspaceClient, host string) gatewayProbe { - payload, reason := gatewayGetJSON(ctx, client, host+legacyEndpointsPath+"?page_size=1") - if payload == nil { - return gatewayProbe{detail: versionNeutralGatewayDetail(reason), conclusive: true} - } - if hasNonEmptyCollection(payload, "endpoints") { - return gatewayProbe{reachable: true, detail: "reachable, accessible endpoint returned", resourceAvailable: true, conclusive: true} +func probeModelServices(ctx context.Context, apiClient *client.DatabricksClient) gatewayProbe { + it := catalog.NewAiGateway(apiClient).ListModelServices(ctx, catalog.ListModelServicesRequest{PageSize: modelServiceProbePageSize}) + _, err := it.Next(ctx) + if err != nil { + if err == listing.ErrNoMoreItems { + // endpoint is reachable but there are no results + return gatewayProbe{reachable: true} + } + return gatewayProbe{err: err} } - return gatewayProbe{reachable: true, detail: "reachable, no accessible endpoints returned", conclusive: true} + return gatewayProbe{reachable: true, resourceAvailable: true} } -func gatewayGetJSON(ctx context.Context, client *databricks.WorkspaceClient, reqURL string) (any, string) { - req, err := http.NewRequestWithContext(ctx, http.MethodGet, reqURL, nil) - if err != nil { - return nil, fmt.Sprintf("network error: %v", err) - } - req.Header.Set("Accept", "application/json") - // Set X-Databricks-Workspace-Id (and suppress the legacy workspace_id=none - // sentinel) so these hand-written workspace-routed GETs reach the workspace - // plane on unified/SPOG hosts instead of the account plane. - for k, v := range auth.WorkspaceIDHeaders(client.Config) { - req.Header.Set(k, v) - } - if err := client.Config.Authenticate(req); err != nil { - return nil, fmt.Sprintf("network error: %v", err) - } - resp, err := (&http.Client{Transport: client.Config.HTTPTransport, Timeout: 10 * time.Second}).Do(req) +type legacyEndpointsResponse struct { + Endpoints []any `json:"endpoints"` +} + +func probeLegacyEndpoints(ctx context.Context, apiClient *client.DatabricksClient) gatewayProbe { + var resp legacyEndpointsResponse + err := apiClient.Do(ctx, http.MethodGet, legacyEndpointsPath, auth.WorkspaceIDHeaders(apiClient.Config), map[string]any{ + "page_size": "1", + }, nil, &resp) if err != nil { - return nil, fmt.Sprintf("network error: %v", err) - } - defer resp.Body.Close() - body, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<20)) - if resp.StatusCode == http.StatusOK { - var payload any - if err := json.Unmarshal(body, &payload); err != nil { - return nil, fmt.Sprintf("response was not valid JSON (%v)", err) - } - return payload, "" + return gatewayProbe{err: err} } - reason := fmt.Sprintf("HTTP %d %s", resp.StatusCode, http.StatusText(resp.StatusCode)) - if excerpt := strings.TrimSpace(string(body)); excerpt != "" { - if len(excerpt) > 200 { - excerpt = excerpt[:200] - } - reason += ": " + excerpt + if len(resp.Endpoints) > 0 { + return gatewayProbe{reachable: true, resourceAvailable: true} } - return nil, reason + return gatewayProbe{reachable: true} } var ( @@ -258,47 +246,52 @@ var ( ) // rewrite the internal v2/v3 API labels in a failure reason into user-facing terms -func versionNeutralGatewayDetail(detail string) string { - detail = gatewayDetailV3.ReplaceAllString(detail, "model service") - return gatewayDetailV2.ReplaceAllString(detail, "legacy endpoint") +func versionNeutralGatewayError(reason string) error { + reason = gatewayDetailV3.ReplaceAllString(reason, "model service") + reason = gatewayDetailV2.ReplaceAllString(reason, "legacy endpoint") + return errors.New(reason) } // returns true when retrying another workspace API can't rescue the token. // A bare 403 is left out: it can be endpoint-specific authorization, so the // preflight still tries the fallback before giving up. -func looksLikeDefinitiveAuthFailure(reason string) bool { - if strings.Contains(reason, "HTTP 401") { +func looksLikeDefinitiveAuthFailure(err error) bool { + // HTTP 401 + if errors.Is(err, apierr.ErrUnauthenticated) { return true } - return strings.Contains(reason, "HTTP 400") && strings.Contains(strings.ToLower(reason), "invalid token") + return errors.Is(err, apierr.ErrBadRequest) && strings.Contains(strings.ToLower(err.Error()), "invalid token") } // matches a 403 that reports the OAuth token is missing a required scope (which // re-login can fix), as opposed to a plain permission 403. -func looksLikeScopeFailure(reason string) bool { - l := strings.ToLower(reason) - return strings.Contains(l, "http 403") && strings.Contains(l, "oauth token") && strings.Contains(l, "required scopes") +func looksLikeScopeFailure(err error) bool { + reason := strings.ToLower(err.Error()) + return errors.Is(err, apierr.ErrPermissionDenied) && strings.Contains(reason, "oauth token") && strings.Contains(reason, "required scopes") } func looksLikePermissionFailure(reason string) bool { return strings.Contains(reason, "HTTP 403") } -var gatewayTransientStatus = regexp.MustCompile(`HTTP (429|5\d\d)`) - -// returns true for a probe failure that is likely temporary (rate limit, -// server error, or a network/transport error) rather than a stable -// "disabled"/"unauthorized" verdict. -func looksLikeTransient(reason string) bool { - return strings.HasPrefix(reason, "network error") || gatewayTransientStatus.MatchString(reason) +func looksLikeTransient(err error) bool { + return strings.HasPrefix(err.Error(), "network error") || + errors.Is(err, apierr.ErrTooManyRequests) || // HTTP 429 + errors.Is(err, apierr.ErrInternalError) || // HTTP 500 + errors.Is(err, apierr.ErrTemporarilyUnavailable) || // HTTP 503 + errors.Is(err, apierr.ErrDeadlineExceeded) // HTTP 504 } func aiGatewayAuthError(host, reason string) error { - return fmt.Errorf("the Databricks workspace %s rejected the access token (%s). Try:\n databricks auth logout --host %s\n databricks auth login --host %s", host, reason, host, host) + return versionNeutralGatewayError( + fmt.Sprintf("the Databricks workspace %s rejected the access token (%s). Try:\n databricks auth logout --host %s\n databricks auth login --host %s", host, reason, host, host), + ) } func aiGatewayScopeError(host, reason string) error { - return fmt.Errorf("the access token for %s is missing an OAuth scope required by the AI Gateway APIs (%s). Re-authenticate to mint a token with the needed scopes:\n databricks auth login --host %s", host, reason, host) + return versionNeutralGatewayError( + fmt.Sprintf("the access token for %s is missing an OAuth scope required by the AI Gateway APIs (%s). Re-authenticate to mint a token with the needed scopes:\n databricks auth login --host %s", host, reason, host), + ) } func hasNonEmptyCollection(payload any, key string) bool { @@ -449,6 +442,7 @@ func launchAgent(ctx context.Context, home string, agent agentSpec, workspace st if workspace != "" { argv = append(argv, "--workspace", workspace) } + argv = append(argv, "--") argv = append(argv, contextArgs...) argv = append(argv, agentArgs...) // Pass the session token as DATABRICKS_BEARER so Unity Gateway CLI authenticates headlessly. diff --git a/experimental/ssh/internal/client/agentshim_test.go b/experimental/ssh/internal/client/agentshim_test.go index 658f73b1a66..3cae09d978c 100644 --- a/experimental/ssh/internal/client/agentshim_test.go +++ b/experimental/ssh/internal/client/agentshim_test.go @@ -1,6 +1,7 @@ package client import ( + "context" "crypto/sha256" "encoding/hex" "io/fs" @@ -183,7 +184,15 @@ func TestAgentSystemContext(t *testing.T) { func newProbeClient(t *testing.T, host string) *databricks.WorkspaceClient { t.Helper() - w, err := databricks.NewWorkspaceClient((*databricks.Config)(&config.Config{Host: host, Token: "test-token"})) + w, err := databricks.NewWorkspaceClient((*databricks.Config)(&config.Config{ + Host: host, + Token: "test-token", + // The fake gateway server has no /.well-known/databricks-config, so skip + // the metadata fetch that would otherwise log a noisy resolve warning. + HostMetadataResolver: func(context.Context, string) (*config.HostMetadata, error) { + return nil, nil + }, + })) require.NoError(t, err) return w } @@ -283,7 +292,7 @@ func TestProbeAIGateway(t *testing.T) { }) t.Run("missing OAuth scope on both paths routes to re-auth", func(t *testing.T) { - scope := jsonHandler(http.StatusForbidden, "Provided OAuth token does not have required scopes: unity-catalog") + scope := jsonHandler(http.StatusForbidden, `{"error_code":"PERMISSION_DENIED","message":"Provided OAuth token does not have required scopes: unity-catalog"}`) srv := newGatewayServer(t, scope, scope) defer srv.Close() err := probeAIGateway(t.Context(), newProbeClient(t, srv.URL)) @@ -321,9 +330,10 @@ func TestProbeAIGateway(t *testing.T) { }) } -func TestProbeModelServicesInconclusive(t *testing.T) { - // A later page failing (after the first page succeeded) is reachable-but- - // inconclusive, never a confirmed "empty". +func TestProbeModelServicesPageError(t *testing.T) { + // A failure while paging surfaces as a probe error: the iterator can't tell a + // first-page failure from a later-page one, so any page error is not treated + // as a reachable "empty". srv := newGatewayServer(t, func(w http.ResponseWriter, r *http.Request) { if r.URL.Query().Get("page_token") == "" { w.WriteHeader(http.StatusOK) @@ -335,8 +345,10 @@ func TestProbeModelServicesInconclusive(t *testing.T) { }, nil) defer srv.Close() - probe := probeModelServices(t.Context(), newProbeClient(t, srv.URL), strings.TrimRight(srv.URL, "/")) - assert.True(t, probe.reachable) - assert.False(t, probe.conclusive) + apiClient, err := newProbeAPIClient(newProbeClient(t, srv.URL).Config) + require.NoError(t, err) + probe := probeModelServices(t.Context(), apiClient) + assert.False(t, probe.reachable) assert.False(t, probe.resourceAvailable) + assert.Error(t, probe.err) } From 692c2210b79bf2bdaebdf562cb3fe16518195ac1 Mon Sep 17 00:00:00 2001 From: Russell Clarey Date: Wed, 23 Sep 2026 11:33:22 +0000 Subject: [PATCH 6/6] lint --- experimental/ssh/internal/client/agentshim.go | 23 +------------------ 1 file changed, 1 insertion(+), 22 deletions(-) diff --git a/experimental/ssh/internal/client/agentshim.go b/experimental/ssh/internal/client/agentshim.go index 5dd6149d829..ead57dc3ebe 100644 --- a/experimental/ssh/internal/client/agentshim.go +++ b/experimental/ssh/internal/client/agentshim.go @@ -141,6 +141,7 @@ func probeAIGateway(ctx context.Context, wsclient *databricks.WorkspaceClient) e apiClient, err := newProbeAPIClient(wsclient.Config) if err != nil { // this probe is best effort, so don't block + log.Debugf(ctx, "Skipping Unity AI Gateway check: %s", err) return nil } @@ -270,10 +271,6 @@ func looksLikeScopeFailure(err error) bool { return errors.Is(err, apierr.ErrPermissionDenied) && strings.Contains(reason, "oauth token") && strings.Contains(reason, "required scopes") } -func looksLikePermissionFailure(reason string) bool { - return strings.Contains(reason, "HTTP 403") -} - func looksLikeTransient(err error) bool { return strings.HasPrefix(err.Error(), "network error") || errors.Is(err, apierr.ErrTooManyRequests) || // HTTP 429 @@ -294,24 +291,6 @@ func aiGatewayScopeError(host, reason string) error { ) } -func hasNonEmptyCollection(payload any, key string) bool { - obj, ok := payload.(map[string]any) - if !ok { - return false - } - arr, ok := obj[key].([]any) - return ok && len(arr) > 0 -} - -func stringField(payload any, key string) string { - obj, ok := payload.(map[string]any) - if !ok { - return "" - } - s, _ := obj[key].(string) - return s -} - // --- end Unity AI Gateway preflight --- func bootstrapAndLaunchAgent(ctx context.Context, agent agentSpec, workspace string, agentArgs []string) error {