Skip to content

[SSH] Pass Codex context via developer_instructions #7197

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
rclarey merged 2 commits into from
Oct 28, 2026
+105 −81
Merged

[SSH] Pass Codex context via developer_instructions #7197

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
92791d5
ssh: pass codex context via developer_instructions
rclarey Oct 27, 2026
9f28a77
put context args after user args
rclarey Oct 28, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
42 changes: 36 additions & 6 deletions experimental/ssh/internal/agentlauncher/agentlauncher.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -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) {

@rugpanov rugpanov Oct 27, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[blocking] ug codex -- -c developer_instructions=... <args> always gives ug two or more passthrough args. At the pinned ugCommit, _smart_routing_launch_shape in src/ucode/cli.py treats a Codex launch as interactive only when it has no args or one explicit prompt. So _disable_smart_routing_for_subcommand turns off smart routing for every Codex session that starts through the shim, also for workspaces that turned it on. The tool_args[:1] == ["update"] checks in src/ucode/agents/codex.py also no longer match codex update. Agree with the Unity Gateway team on a supported way to add developer instructions, or make ug skip leading -c overrides when it picks the launch shape. Then add a Codex case to TestPrepareAgentLaunch.

@rclarey rclarey Oct 28, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

At the pinned ugCommit, _smart_routing_launch_shape in src/ucode/cli.py treats a Codex launch as interactive only when it has no args or one explicit prompt. So _disable_smart_routing_for_subcommand turns off smart routing for every Codex session that starts through the shim, also for workspaces that turned it on

Yeah this is a bad heuristic. There are plenty of flags a normal user would pass to an interactive session, and with this they are opted out of smart routing.

The tool_args[:1] == ["update"] checks in src/ucode/agents/codex.py also no longer match codex update

Will fix 👍

// 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

@anton-107 anton-107 Oct 27, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Blocking: this silently drops the user's own developer_instructions.

The old ContextHomeFile path deliberately didn't clobber (// already present — don't clobber), and a test covered that. A -c override is Codex's highest-priority config layer and it replaces the key, so any developer_instructions the user set in ~/.codex/config.toml (or a project .codex/config.toml) vanishes once they launch through the shim. No error or warning. This PR also removes the "does not clobber an existing file" test without a replacement.

Repro with a real codex (0.164.0), using the exact arg this function produces (arg below):

D=$(mktemp -d); W=$(mktemp -d)
printf 'developer_instructions = "USER-OWN-INSTRUCTIONS"\n' > $D/config.toml
cd $W && git init -q
CODEX_HOME=$D codex debug prompt-input | grep -c USER-OWN-INSTRUCTIONS                 # 1
CODEX_HOME=$D codex -c "$(cat arg)" debug prompt-input | grep -c USER-OWN-INSTRUCTIONS  # 0

The second run shows only the Databricks context. Could we combine this with the user's existing value, or use a mechanism that adds rather than replaces, and bring back a test that the user's own instructions survive?

@rclarey rclarey Oct 27, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@anton-107 This was intentional (for now at least) since I don't want to reimplement Codex's config resolution in order to read the existing value (if it exists).

I thought this was a reasonable trade off since these are agents running in a managed setting (ug configured in the SSH tunnel) and users can use AGENTS.md (global or project scope) to add additional context for their projects.

If this ends up being an issue in practice we can always add the config resolution to merge our instructions with the user's

@anton-107 anton-107 Oct 27, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fair, I'm withdrawing "blocking". Merging would mean reimplementing Codex's layered config resolution, and AGENTS.md still gets loaded alongside the override (I checked), so users have a way to add context. The old path also dropped our context completely whenever the user had their own AGENTS.md, so this is a net improvement. One optional nit: a line in the codexContextArg comment saying the override replaces any user-set developer_instructions, and why that's acceptable, so nobody "fixes" it later.

}

func AgentByName(name string) (AgentSpec, bool) {
Expand Down
41 changes: 31 additions & 10 deletions experimental/ssh/internal/agentlauncher/agentlauncher_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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)
}
46 changes: 13 additions & 33 deletions experimental/ssh/internal/client/agentshim.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@ import (
"path/filepath"
"regexp"
"runtime"
"slices"
"strings"
"time"

Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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)
Comment thread
anton-107 marked this conversation as resolved.
}

// fetch url to a temp file, failing unless its SHA256 matches the expected checksum.
Expand Down
57 changes: 25 additions & 32 deletions experimental/ssh/internal/client/agentshim_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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)
Expand All @@ -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) {
Expand All @@ -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) {
Expand Down Expand Up @@ -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)
})

Expand Down
Loading