[SSH] Pass Codex context via developer_instructions (#7197)

## Changes

Use `-c developer_instructions="..."` instead of `AGENTS.md` to pass
system context to Codex

## Why

This avoids issues with stale/user defined `AGENT.md` files

## Tests

tests updated

Author: rclarey

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) {

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

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,45 +689,24 @@ 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
 }
 
 // put the Databricks context where the agent reads it, returning any extra argv.
 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.

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