diff --git a/experimental/ssh/internal/agentlauncher/agentlauncher.go b/experimental/ssh/internal/agentlauncher/agentlauncher.go index 2c7daac8ade..d13bd7e005a 100644 --- a/experimental/ssh/internal/agentlauncher/agentlauncher.go +++ b/experimental/ssh/internal/agentlauncher/agentlauncher.go @@ -7,8 +7,11 @@ import ( "context" "errors" "fmt" + "os" "path/filepath" + "strings" + "github.com/BurntSushi/toml" "github.com/databricks/cli/libs/atomicfile" "github.com/databricks/cli/libs/env" "github.com/databricks/cli/libs/shellquote" @@ -24,21 +27,48 @@ const ( // WorkspaceHomeEnv passes the user's workspace home to the shim for the agent context. WorkspaceHomeEnv = "DATABRICKS_WORKSPACE_HOME" + + // ContextFile holds the Databricks session context, under RootDir. + ContextFile = "agent-system-context.md" ) type AgentSpec struct { Name string TelemetryName protos.SshAgentShimAgent - // 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 + // ContextArg puts the system context where the agent reads it, returning the + // argv to pass to the agent binary. + ContextArg func(home, systemContext string) ([]string, error) } // Currently limited to the agents whose Unity Gateway CLI command supports --workspace var SupportedAgents = []AgentSpec{ - {Name: "claude", TelemetryName: protos.SshAgentShimAgentClaudeCode, ContextFlag: "--append-system-prompt-file"}, - {Name: "codex", TelemetryName: protos.SshAgentShimAgentCodex, ContextHomeFile: ".codex/AGENTS.md"}, + {Name: "claude", TelemetryName: protos.SshAgentShimAgentClaudeCode, ContextArg: claudeContextArg}, + {Name: "codex", TelemetryName: protos.SshAgentShimAgentCodex, ContextArg: codexContextArg}, +} + +func claudeContextArg(home, systemContext string) ([]string, error) { + // 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, RootDir) + 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{"--append-system-prompt-file", f}, nil +} + +func codexContextArg(home, systemContext string) ([]string, error) { + // The override value is parsed as TOML, so encode the multi-line context as + // a TOML string rather than relying on codex's raw-string fallback. + // https://github.com/openai/codex/blob/rust-v0.145.0/codex-rs/utils/cli/src/config_override.rs + kv, err := toml.Marshal(map[string]string{"developer_instructions": systemContext}) + if err != nil { + return nil, fmt.Errorf("failed to encode agent context: %w", err) + } + return []string{"-c", strings.TrimSpace(string(kv))}, nil } func AgentByName(name string) (AgentSpec, bool) { diff --git a/experimental/ssh/internal/agentlauncher/agentlauncher_test.go b/experimental/ssh/internal/agentlauncher/agentlauncher_test.go index b502e604f92..5bd8d03ff26 100644 --- a/experimental/ssh/internal/agentlauncher/agentlauncher_test.go +++ b/experimental/ssh/internal/agentlauncher/agentlauncher_test.go @@ -6,6 +6,7 @@ import ( "runtime" "testing" + "github.com/BurntSushi/toml" "github.com/databricks/cli/libs/env" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -62,20 +63,40 @@ func TestInstallContinuesPastFailure(t *testing.T) { } func TestSupportedAgents(t *testing.T) { - // 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"}, - } - require.Len(t, SupportedAgents, len(expect)) - for name, want := range expect { + for _, name := range []string{"claude", "codex"} { 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) + assert.NotNil(t, a.ContextArg, "%s ContextArg", name) } + assert.Len(t, SupportedAgents, 2) _, ok := AgentByName("not-an-agent") assert.False(t, ok) } + +func TestClaudeContextArg(t *testing.T) { + home := t.TempDir() + args, err := claudeContextArg(home, "the context") + require.NoError(t, err) + path := filepath.Join(home, RootDir, ContextFile) + assert.Equal(t, []string{"--append-system-prompt-file", path}, args) + data, err := os.ReadFile(path) + require.NoError(t, err) + assert.Equal(t, "the context", string(data)) +} + +func TestCodexContextArg(t *testing.T) { + home := t.TempDir() + const systemContext = "line \"one\"\nline two" + args, err := codexContextArg(home, systemContext) + require.NoError(t, err) + require.Len(t, args, 2) + assert.Equal(t, "-c", args[0]) + // codex splits the override on the first '=' and parses the rest as TOML + var got map[string]string + require.NoError(t, toml.Unmarshal([]byte(args[1]), &got)) + assert.Equal(t, map[string]string{"developer_instructions": systemContext}, got) + entries, err := os.ReadDir(home) + require.NoError(t, err) + assert.Empty(t, entries) +} diff --git a/experimental/ssh/internal/client/agentshim.go b/experimental/ssh/internal/client/agentshim.go index 29974efd32d..e52fbbf11eb 100644 --- a/experimental/ssh/internal/client/agentshim.go +++ b/experimental/ssh/internal/client/agentshim.go @@ -13,6 +13,7 @@ import ( "path/filepath" "regexp" "runtime" + "slices" "strings" "time" @@ -77,7 +78,7 @@ const ( workspaceHomeEnv = agentlauncher.WorkspaceHomeEnv // the file holding the Databricks session context. - contextFile = "agent-system-context.md" + contextFile = agentlauncher.ContextFile // lock file to serialize first-run toolchain setup across // concurrent SSH clients sharing this driver's $HOME. @@ -688,8 +689,16 @@ func prepareAgentLaunch(ctx context.Context, home string, agent agentlauncher.Ag return nil, nil, categorize(protos.SshAgentShimErrorCategoryAgentContextSetupFailed, err) } argv = []string{ugExecutable, agent.Name, "--"} - argv = append(argv, contextArgs...) - argv = append(argv, agentArgs...) + // If the user passed `--` then inject context args just before, otherwise add + // them after the user's args. This is needed so we don't interfere with `ug`'s + // heuristics around positional args + i := slices.Index(agentArgs, "--") + if i == -1 { + argv = append(argv, agentArgs...) + argv = append(argv, contextArgs...) + } else { + argv = slices.Concat(argv, agentArgs[:i], contextArgs, agentArgs[i:]) + } return argv, os.Environ(), nil } @@ -697,36 +706,7 @@ func prepareAgentLaunch(ctx context.Context, home string, agent agentlauncher.Ag func injectAgentContext(ctx context.Context, home string, agent agentlauncher.AgentSpec) ([]string, error) { serverless := dbr.RuntimeVersion(ctx).Type == dbr.ClusterTypeServerless systemContext := agentSystemContext(env.Get(ctx, workspaceHomeEnv), serverless) - switch { - case agent.ContextFlag != "": - // 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, agentRootDir) - 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 - 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(systemContext), 0o644); err != nil { - return nil, fmt.Errorf("failed to write %s: %w", f, err) - } - return nil, nil - default: - return nil, nil - } + return agent.ContextArg(home, systemContext) } // fetch url to a temp file, failing unless its SHA256 matches the expected checksum. diff --git a/experimental/ssh/internal/client/agentshim_test.go b/experimental/ssh/internal/client/agentshim_test.go index a175f8ddb5d..fc17e1dc772 100644 --- a/experimental/ssh/internal/client/agentshim_test.go +++ b/experimental/ssh/internal/client/agentshim_test.go @@ -20,10 +20,12 @@ import ( "testing" "time" + "github.com/BurntSushi/toml" "github.com/databricks/cli/experimental/ssh/internal/agentlauncher" "github.com/databricks/cli/libs/cmdio" "github.com/databricks/cli/libs/dbr" + "github.com/databricks/cli/libs/env" "github.com/databricks/cli/libs/telemetry/protos" "github.com/databricks/databricks-sdk-go" "github.com/databricks/databricks-sdk-go/config" @@ -360,9 +362,8 @@ func TestEnsureToolchain(t *testing.T) { func TestInjectAgentContext(t *testing.T) { claude, _ := agentlauncher.AgentByName("claude") codex, _ := agentlauncher.AgentByName("codex") - noContext := agentlauncher.AgentSpec{Name: "none"} // an agent with neither mechanism - t.Run("flag agent writes a scratch file and returns the flag", func(t *testing.T) { + t.Run("claude gets the context as a scratch file", func(t *testing.T) { home := t.TempDir() args, err := injectAgentContext(dbr.MockRuntime(t.Context(), serverlessRuntime), home, claude) require.NoError(t, err) @@ -373,27 +374,20 @@ func TestInjectAgentContext(t *testing.T) { assert.Contains(t, string(data), "databricks ssh connect") }) - t.Run("home-file agent writes the global instructions file, no args", func(t *testing.T) { + t.Run("codex gets the context as a TOML config override", func(t *testing.T) { home := t.TempDir() - args, err := injectAgentContext(dbr.MockRuntime(t.Context(), serverlessRuntime), home, codex) + ctx := env.Set(dbr.MockRuntime(t.Context(), serverlessRuntime), workspaceHomeEnv, "/Workspace/Users/a\"b@example.com") + args, err := injectAgentContext(ctx, 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(dbr.MockRuntime(t.Context(), serverlessRuntime), home, codex) - require.NoError(t, err) - assert.Nil(t, args) - data, err := os.ReadFile(target) + require.Len(t, args, 2) + assert.Equal(t, "-c", args[0]) + // codex splits the override on the first '=' and parses the rest as TOML + var got map[string]string + require.NoError(t, toml.Unmarshal([]byte(args[1]), &got)) + assert.Equal(t, map[string]string{"developer_instructions": agentSystemContext("/Workspace/Users/a\"b@example.com", true)}, got) + entries, err := os.ReadDir(home) require.NoError(t, err) - assert.Equal(t, "user's own instructions", string(data)) + assert.Empty(t, entries) }) t.Run("context describes the compute the shim runs on", func(t *testing.T) { @@ -413,16 +407,6 @@ func TestInjectAgentContext(t *testing.T) { assert.Contains(t, string(data), c.want) } }) - - t.Run("agent without a mechanism writes nothing", func(t *testing.T) { - home := t.TempDir() - args, err := injectAgentContext(dbr.MockRuntime(t.Context(), serverlessRuntime), 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) { @@ -796,12 +780,21 @@ func TestPrepareAgentLaunch(t *testing.T) { agent, ok := agentlauncher.AgentByName("claude") require.True(t, ok) - t.Run("passes the context args before the user's agent args", func(t *testing.T) { + t.Run("passes the context args after the user's args", func(t *testing.T) { home := t.TempDir() argv, environ, err := prepareAgentLaunch(dbr.MockRuntime(t.Context(), serverlessRuntime), home, agent, []string{"--resume"}) require.NoError(t, err) require.NotEmpty(t, argv) - assert.Equal(t, []string{"claude", "--", "--append-system-prompt-file", filepath.Join(home, agentRootDir, contextFile), "--resume"}, argv[1:]) + assert.Equal(t, []string{"claude", "--", "--resume", "--append-system-prompt-file", filepath.Join(home, agentRootDir, contextFile)}, argv[1:]) + assert.Equal(t, os.Environ(), environ) + }) + + t.Run("passes the context args before a '--' arg", func(t *testing.T) { + home := t.TempDir() + argv, environ, err := prepareAgentLaunch(dbr.MockRuntime(t.Context(), serverlessRuntime), home, agent, []string{"--max-budget-usd", "1", "--", "hello there"}) + require.NoError(t, err) + require.NotEmpty(t, argv) + assert.Equal(t, []string{"claude", "--", "--max-budget-usd", "1", "--append-system-prompt-file", filepath.Join(home, agentRootDir, contextFile), "--", "hello there"}, argv[1:]) assert.Equal(t, os.Environ(), environ) })