Skip to content

Commit 1ef7e95

Browse files
committed
feat/batch-changes: support codingAgent steps in batch exec
Vendors lib/batches/codingagent + lib/batches/codex from sourcegraph/ sourcegraph, adds the v3 codingAgent step + image alias to the spec + schema, and rewrites step.Run via codingagent.RenderRunCommand when step.CodingAgent is set. SRC_BATCHES_MODEL_PROVIDER_TOKEN and SRC_BATCHES_JOB_ID are forwarded from src's own env into the user container so the agent CLI can reach the server's model-provider proxy. Demo/MVP for v1 codingAgent support. Not for merge as-is.
1 parent ed8850d commit 1ef7e95

13 files changed

Lines changed: 362 additions & 10 deletions

File tree

.gitignore

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,3 +15,4 @@ bazel-src-cli
1515
.DS_Store
1616
samples
1717
.amp
18+
bin/

cmd/src/batch_exec.go

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -278,6 +278,7 @@ func convertWorkspace(w batcheslib.WorkspacesExecutionInput) *executor.Task {
278278
BatchChangeAttributes: &w.BatchChangeAttributes,
279279
CachedStepResultFound: w.CachedStepResultFound,
280280
CachedStepResult: w.CachedStepResult,
281+
ModelProviderURL: w.ModelProviderURL,
281282
}
282283

283284
return task

internal/batches/executor/run_steps.go

Lines changed: 67 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,8 @@ import (
1414
"time"
1515

1616
batcheslib "github.com/sourcegraph/sourcegraph/lib/batches"
17+
"github.com/sourcegraph/sourcegraph/lib/batches/codingagent"
18+
codingagenttypes "github.com/sourcegraph/sourcegraph/lib/batches/codingagent/types"
1719
"github.com/sourcegraph/sourcegraph/lib/batches/execution"
1820
"github.com/sourcegraph/sourcegraph/lib/batches/git"
1921
"github.com/sourcegraph/sourcegraph/lib/batches/template"
@@ -272,12 +274,32 @@ func executeSingleStep(
272274
return bytes.Buffer{}, bytes.Buffer{}, err
273275
}
274276

275-
runScriptFile, runScript, cleanup, err := createRunScriptFile(ctx, opts.TempDir, step.Run, stepContext)
277+
var (
278+
runScriptFile string
279+
runScript string
280+
runScriptCleanup func()
281+
)
282+
if step.CodingAgent != nil {
283+
if opts.Task.ModelProviderURL == "" {
284+
err = errors.New("codingAgent step requires WorkspacesExecutionInput.ModelProviderURL to be set")
285+
opts.UI.StepPreparingFailed(stepIdx+1, err)
286+
return bytes.Buffer{}, bytes.Buffer{}, err
287+
}
288+
runScript, err = codingagent.RenderRunCommand(step.CodingAgent, opts.Task.ModelProviderURL, stepContext)
289+
if err != nil {
290+
err = errors.Wrap(err, "rendering codingAgent step")
291+
opts.UI.StepPreparingFailed(stepIdx+1, err)
292+
return bytes.Buffer{}, bytes.Buffer{}, err
293+
}
294+
runScriptFile, runScriptCleanup, err = writeRunScriptFile(opts.TempDir, runScript)
295+
} else {
296+
runScriptFile, runScript, runScriptCleanup, err = createRunScriptFile(ctx, opts.TempDir, step.Run, stepContext)
297+
}
276298
if err != nil {
277299
opts.UI.StepPreparingFailed(stepIdx+1, err)
278300
return bytes.Buffer{}, bytes.Buffer{}, err
279301
}
280-
defer cleanup()
302+
defer runScriptCleanup()
281303

282304
// Parse and render the step.Files.
283305
filesToMount, cleanup, err := createFilesToMount(opts.TempDir, step, stepContext)
@@ -303,6 +325,21 @@ func executeSingleStep(
303325
return bytes.Buffer{}, bytes.Buffer{}, err
304326
}
305327

328+
// For codingAgent steps, forward the model-provider auth env from the
329+
// executor's environment (injected via CliStep.Env and JobTokenEnvVar on
330+
// the server) into the user container so the agent CLI can talk to the
331+
// /model-provider/batches proxy.
332+
if step.CodingAgent != nil {
333+
for _, key := range []string{codingagenttypes.ModelProviderTokenEnvVar, codingagenttypes.JobIDEnvVar} {
334+
for _, e := range opts.GlobalEnv {
335+
if v, ok := strings.CutPrefix(e, key+"="); ok {
336+
env[key] = v
337+
break
338+
}
339+
}
340+
}
341+
}
342+
306343
opts.UI.StepPreparingSuccess(stepIdx + 1)
307344

308345
// ----------
@@ -573,6 +610,34 @@ func createRunScriptFile(ctx context.Context, tempDir string, stepRun string, st
573610
return runScriptFile.Name(), runScript.String(), cleanup, nil
574611
}
575612

613+
// writeRunScriptFile writes a pre-rendered run script (e.g. a codingAgent
614+
// step desugared by the codingagent registry) verbatim to a temp file, with
615+
// the same permission semantics as createRunScriptFile. Unlike that helper,
616+
// this one does NOT pass the content through template.RenderStepTemplate:
617+
// the script is treated as opaque shell text so embedded sequences like
618+
// `{{` inside a shell-quoted prompt are not re-parsed as templates.
619+
func writeRunScriptFile(tempDir, script string) (string, func(), error) {
620+
runScriptFile, err := os.CreateTemp(tempDir, "")
621+
if err != nil {
622+
return "", nil, errors.Wrap(err, "creating temporary file")
623+
}
624+
cleanup := func() { os.Remove(runScriptFile.Name()) }
625+
626+
if _, err := runScriptFile.WriteString(script); err != nil {
627+
cleanup()
628+
return "", nil, errors.Wrap(err, "writing temporary file")
629+
}
630+
if err := runScriptFile.Close(); err != nil {
631+
cleanup()
632+
return "", nil, errors.Wrap(err, "closing temporary file")
633+
}
634+
if err := os.Chmod(runScriptFile.Name(), 0644); err != nil {
635+
cleanup()
636+
return "", nil, errors.Wrap(err, "setting permissions on the temporary file")
637+
}
638+
return runScriptFile.Name(), cleanup, nil
639+
}
640+
576641
// createCidFile creates a temporary file that will contain the container ID
577642
// when executing steps.
578643
// It returns the location of the file and a function that cleans up the

internal/batches/executor/task.go

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,9 @@ type Task struct {
2929
// When this field is true, CachedStepResult is also populated.
3030
CachedStepResultFound bool
3131
CachedStepResult execution.AfterStepResult
32+
// ModelProviderURL is the resolved proxy base URL for coding-agent
33+
// steps; empty unless the spec contains at least one codingAgent step.
34+
ModelProviderURL string
3235
}
3336

3437
func (t *Task) ArchivePathToFetch() string {

lib/batches/batch_spec.go

Lines changed: 33 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -90,13 +90,29 @@ func (oqor *OnQueryOrRepository) GetBranches() ([]string, error) {
9090
}
9191

9292
type Step struct {
93-
Run string `json:"run,omitempty" yaml:"run"`
94-
Container string `json:"container,omitempty" yaml:"container"`
95-
Env env.Environment `json:"env" yaml:"env"`
96-
Files map[string]string `json:"files,omitempty" yaml:"files,omitempty"`
97-
Outputs Outputs `json:"outputs,omitempty" yaml:"outputs,omitempty"`
98-
Mount []Mount `json:"mount,omitempty" yaml:"mount,omitempty"`
99-
If any `json:"if,omitempty" yaml:"if,omitempty"`
93+
Run string `json:"run,omitempty" yaml:"run"`
94+
CodingAgent *CodingAgentStep `json:"codingAgent,omitempty" yaml:"codingAgent,omitempty"`
95+
Container string `json:"container,omitempty" yaml:"container"`
96+
Image string `json:"image,omitempty" yaml:"image,omitempty"`
97+
Env env.Environment `json:"env" yaml:"env"`
98+
Files map[string]string `json:"files,omitempty" yaml:"files,omitempty"`
99+
Outputs Outputs `json:"outputs,omitempty" yaml:"outputs,omitempty"`
100+
Mount []Mount `json:"mount,omitempty" yaml:"mount,omitempty"`
101+
If any `json:"if,omitempty" yaml:"if,omitempty"`
102+
}
103+
104+
// CodingAgentType identifies a registered coding-agent implementation.
105+
type CodingAgentType string
106+
107+
const (
108+
CodingAgentTypeCodex CodingAgentType = "codex"
109+
)
110+
111+
// CodingAgentStep is a v3-spec step that delegates the step's work to a
112+
// coding agent CLI invoked via the server-side model-provider proxy.
113+
type CodingAgentStep struct {
114+
Type CodingAgentType `json:"type,omitempty" yaml:"type"`
115+
Prompt string `json:"prompt,omitempty" yaml:"prompt"`
100116
}
101117

102118
func (s *Step) IfCondition() string {
@@ -161,6 +177,16 @@ func parseBatchSpec(schema string, data []byte) (*BatchSpec, error) {
161177
return nil, err
162178
}
163179

180+
if spec.Version == 3 {
181+
// Mirror v3 `image:` into `container:` so executor consumers that
182+
// read step.Container keep working.
183+
for i := range spec.Steps {
184+
if spec.Steps[i].Image != "" {
185+
spec.Steps[i].Container = spec.Steps[i].Image
186+
}
187+
}
188+
}
189+
164190
var errs error
165191

166192
if len(spec.Steps) != 0 && spec.ChangesetTemplate == nil {

lib/batches/codex/codex.go

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,43 @@
1+
// Package codex implements the codex coding agent run-command rewrite.
2+
package codex
3+
4+
import (
5+
"fmt"
6+
7+
"github.com/kballard/go-shellquote"
8+
9+
batcheslib "github.com/sourcegraph/sourcegraph/lib/batches"
10+
"github.com/sourcegraph/sourcegraph/lib/batches/codingagent/types"
11+
)
12+
13+
const model = "gpt-5.4"
14+
15+
var routes = []types.ModelProviderRoute{
16+
{WirePath: "/responses", UpstreamPath: "/v1/completions/openai-responses"},
17+
}
18+
19+
type Agent struct{}
20+
21+
func (Agent) Type() batcheslib.CodingAgentType { return batcheslib.CodingAgentTypeCodex }
22+
func (Agent) ModelProviderRoutes() []types.ModelProviderRoute { return routes }
23+
func (Agent) ImageRequirements() []string { return []string{"codex"} }
24+
25+
func (Agent) RunCommand(prompt, modelProviderURL string) string {
26+
return shellquote.Join(
27+
"codex",
28+
"exec",
29+
"--json",
30+
"--sandbox", "danger-full-access",
31+
"--ephemeral",
32+
"--model", model,
33+
"-c", `approval_policy="never"`,
34+
"-c", `model_reasoning_effort="medium"`,
35+
"-c", `model_provider="sourcegraph"`,
36+
"-c", `model_providers.sourcegraph.name="Sourcegraph"`,
37+
"-c", fmt.Sprintf(`model_providers.sourcegraph.base_url=%q`, modelProviderURL),
38+
"-c", fmt.Sprintf(`model_providers.sourcegraph.env_key=%q`, types.ModelProviderTokenEnvVar),
39+
"-c", fmt.Sprintf(`model_providers.sourcegraph.env_http_headers={%q=%q}`, types.JobIDHeaderName, types.JobIDEnvVar),
40+
"-c", `model_providers.sourcegraph.wire_api="responses"`,
41+
prompt,
42+
)
43+
}
Lines changed: 87 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,87 @@
1+
// Package codingagent rewrites v3 batch spec coding-agent steps into the
2+
// shell commands that drive them. Register new agents in the agents list.
3+
package codingagent
4+
5+
import (
6+
"bytes"
7+
"fmt"
8+
"strings"
9+
10+
"github.com/kballard/go-shellquote"
11+
12+
batcheslib "github.com/sourcegraph/sourcegraph/lib/batches"
13+
"github.com/sourcegraph/sourcegraph/lib/batches/codex"
14+
"github.com/sourcegraph/sourcegraph/lib/batches/codingagent/types"
15+
"github.com/sourcegraph/sourcegraph/lib/batches/template"
16+
"github.com/sourcegraph/sourcegraph/lib/errors"
17+
)
18+
19+
// ErrUnknownType is returned when a codingAgent step references a type that
20+
// has no registered Agent.
21+
var ErrUnknownType = errors.New("unknown codingAgent type")
22+
23+
// RenderRunCommand returns the shell-quoted command that runs agentStep. The
24+
// prompt is rendered before quoting; reversing would let templated values
25+
// break out of the shell quoting.
26+
func RenderRunCommand(agentStep *batcheslib.CodingAgentStep, modelProviderURL string, stepCtx *template.StepContext) (string, error) {
27+
a, ok := agents[agentStep.Type]
28+
if !ok {
29+
return "", errors.Wrapf(ErrUnknownType, "codingAgent type %q", agentStep.Type)
30+
}
31+
var renderedPrompt bytes.Buffer
32+
if err := template.RenderStepTemplate("codingagent-prompt", agentStep.Prompt, &renderedPrompt, stepCtx); err != nil {
33+
return "", errors.Wrap(err, "rendering codingAgent.prompt")
34+
}
35+
prefixed := strings.TrimRight(modelProviderURL, "/") + "/" + string(a.Type())
36+
37+
var b strings.Builder
38+
for _, binary := range a.ImageRequirements() {
39+
b.WriteString(failIfMissing(a.Type(), binary))
40+
}
41+
b.WriteString(a.RunCommand(renderedPrompt.String(), prefixed))
42+
return b.String(), nil
43+
}
44+
45+
// failIfMissing returns a shell snippet that, when prepended to the step
46+
// script, writes a message to stderr and exits the container with status 1
47+
// if binary isn't on PATH.
48+
func failIfMissing(agentType batcheslib.CodingAgentType, binary string) string {
49+
msg := fmt.Sprintf(
50+
"codingAgent %q requires %q on PATH in the run container; ensure your image includes the %s binary",
51+
agentType, binary, binary,
52+
)
53+
return fmt.Sprintf("command -v %s >/dev/null 2>&1 || { echo %s >&2; exit 1; }\n",
54+
binary,
55+
shellquote.Join(msg),
56+
)
57+
}
58+
59+
var agents = func() map[batcheslib.CodingAgentType]types.Agent {
60+
out := map[batcheslib.CodingAgentType]types.Agent{}
61+
for _, a := range []types.Agent{
62+
codex.Agent{},
63+
} {
64+
if _, exists := out[a.Type()]; exists {
65+
panic("duplicate codingagent agent for " + a.Type())
66+
}
67+
out[a.Type()] = a
68+
}
69+
return out
70+
}()
71+
72+
// RegisteredModelProviderRoutes returns the model-provider proxy routes
73+
// contributed by every registered Agent, each WirePath prefixed with
74+
// its agent type.
75+
func RegisteredModelProviderRoutes() []types.ModelProviderRoute {
76+
var out []types.ModelProviderRoute
77+
for _, a := range agents {
78+
prefix := "/" + string(a.Type())
79+
for _, route := range a.ModelProviderRoutes() {
80+
out = append(out, types.ModelProviderRoute{
81+
WirePath: prefix + route.WirePath,
82+
UpstreamPath: route.UpstreamPath,
83+
})
84+
}
85+
}
86+
return out
87+
}
Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,52 @@
1+
package codingagent_test
2+
3+
import (
4+
"errors"
5+
"testing"
6+
7+
"github.com/kballard/go-shellquote"
8+
9+
batcheslib "github.com/sourcegraph/sourcegraph/lib/batches"
10+
"github.com/sourcegraph/sourcegraph/lib/batches/codingagent"
11+
"github.com/sourcegraph/sourcegraph/lib/batches/template"
12+
)
13+
14+
func TestRenderRunCommand_unknownType(t *testing.T) {
15+
agentStep := &batcheslib.CodingAgentStep{Type: "nope", Prompt: "x"}
16+
_, err := codingagent.RenderRunCommand(agentStep, "https://example/", &template.StepContext{})
17+
if err == nil {
18+
t.Fatal("expected error, got nil")
19+
}
20+
if !errors.Is(err, codingagent.ErrUnknownType) {
21+
t.Fatalf("expected ErrUnknownType, got %v", err)
22+
}
23+
}
24+
25+
// TestRenderRunCommand_promptShellQuoting ensures the rendered prompt is
26+
// shell-quoted as a single argument even when it contains shell
27+
// metacharacters (apostrophes, semicolons), so prompt text can't break out
28+
// into the host shell.
29+
func TestRenderRunCommand_promptShellQuoting(t *testing.T) {
30+
const repoName = `github.com/sourcegraph/sourcegraph`
31+
prompt := "You're working in the ${{ repository.name }} repository.\n" +
32+
"Add a README section describing the project; don't touch existing files."
33+
agentStep := &batcheslib.CodingAgentStep{
34+
Type: batcheslib.CodingAgentTypeCodex,
35+
Prompt: prompt,
36+
}
37+
stepCtx := &template.StepContext{Repository: template.Repository{Name: repoName}}
38+
cmd, err := codingagent.RenderRunCommand(agentStep, "https://example/", stepCtx)
39+
if err != nil {
40+
t.Fatal(err)
41+
}
42+
43+
tokens, err := shellquote.Split(cmd)
44+
if err != nil {
45+
t.Fatal(err)
46+
}
47+
wantPrompt := "You're working in the " + repoName + " repository.\n" +
48+
"Add a README section describing the project; don't touch existing files."
49+
if got := tokens[len(tokens)-1]; got != wantPrompt {
50+
t.Fatalf("prompt mismatch:\n got: %q\n want: %q", got, wantPrompt)
51+
}
52+
}

0 commit comments

Comments
 (0)