Skip to content

Commit 6ddbb00

Browse files
stefan-ernststefan-ernst
committed
refactor: consolidate run claim preparation
1 parent e25f747 commit 6ddbb00

2 files changed

Lines changed: 172 additions & 166 deletions

File tree

‎internal/services/run_service_inprocess.go‎

Lines changed: 150 additions & 137 deletions
Original file line numberDiff line numberDiff line change
@@ -51,10 +51,18 @@ func runRepos(req RunRequest) []*repoprep.RepoSpec {
5151
// workspace layout (WI-449): the slug's last segment ("owner/core-tests" ->
5252
// "core-tests"), disambiguated with a numeric suffix if two repos share a name.
5353
func repoDirNames(repos []*repoprep.RepoSpec) []string {
54-
names := make([]string, len(repos))
55-
seen := make(map[string]int, len(repos))
56-
for i, r := range repos {
57-
base := r.RepoSlug
54+
slugs := make([]string, len(repos))
55+
for i, repo := range repos {
56+
slugs[i] = repo.RepoSlug
57+
}
58+
return repoWorkspaceDirNames(slugs)
59+
}
60+
61+
func repoWorkspaceDirNames(slugs []string) []string {
62+
names := make([]string, len(slugs))
63+
seen := make(map[string]int, len(slugs))
64+
for i, slug := range slugs {
65+
base := slug
5866
if idx := strings.LastIndex(base, "/"); idx >= 0 && idx < len(base)-1 {
5967
base = base[idx+1:]
6068
}
@@ -85,156 +93,161 @@ func queueBuffer(capacity int) int {
8593
// claimNext admits queued work; preamble failures finalize in place, and shutdown drains queued runs.
8694
func (s *RunService) claimNext() *ClaimedJob {
8795
for {
88-
var job queuedJob
89-
select {
90-
case job = <-s.queue:
91-
case <-s.shutdownCh:
92-
// Drain still-queued runs as canceled, then stop. Channel
93-
// receive is safe across workers; whichever worker wins
94-
// finalizes the run.
95-
for {
96-
select {
97-
case j := <-s.queue:
98-
s.finalize(j.runID, models.AgentRunStatusCanceled, "shutdown before admission")
99-
s.wg.Done()
100-
default:
101-
return nil
102-
}
103-
}
96+
job, ok := s.dequeueClaimJob()
97+
if !ok {
98+
return nil
10499
}
105100

106-
// Per-run context, wired to shutdown so RunService.Cancel and
107-
// process shutdown both reach the in-flight runner.
108-
runCtx, cancel := context.WithCancel(context.Background())
109-
s.registerCancel(job.runID, cancel)
110-
go func() {
111-
select {
112-
case <-s.shutdownCh:
113-
cancel()
114-
case <-runCtx.Done():
115-
}
116-
}()
117-
118-
now := s.now()
119-
transitioned, err := s.repo.MarkRunningIfQueued(runCtx, job.runID, "", now)
120-
if err != nil {
121-
s.logger.Printf("run service: mark running run=%d: %v", job.runID, err)
122-
s.failClaim(job, cancel, fmt.Sprintf("mark running: %v", err), false)
101+
runCtx, cancel, ok := s.beginClaim(job)
102+
if !ok {
123103
continue
124104
}
125-
if !transitioned {
126-
// The row left 'queued' while the job sat on the in-memory
127-
// channel — canceled via the API (WI-341) or otherwise already
128-
// terminal. The terminal status (and its lifecycle event) is
129-
// recorded by whoever made that transition; just release the
130-
// run's accounting and move on instead of executing it anyway.
131-
s.logger.Printf("run service: skipping run=%d: no longer queued at dequeue", job.runID)
132-
cancel()
133-
s.unregisterCancel(job.runID)
134-
s.wg.Done()
105+
106+
st := claimState{req: job.req, ephemeral: job.req.Ephemeral, cancel: cancel}
107+
if err := s.prepareClaimRepos(runCtx, job, &st); err != nil {
108+
s.failClaim(job, cancel, "prepare checkout failed", true)
135109
continue
136110
}
137-
if err := s.repo.AppendEvent(runCtx, job.runID, "lifecycle", `{"phase":"running"}`); err != nil {
138-
s.logger.Printf("run service: append running event run=%d: %v", job.runID, err)
139-
}
140111

141-
st := claimState{req: job.req, ephemeral: job.req.Ephemeral, cancel: cancel}
112+
env, err := s.buildClaimEnv(runCtx, job, &st)
113+
if err != nil {
114+
s.logger.Printf("run service: mint ws token run=%d: %v", job.runID, err)
115+
s.failClaim(job, cancel, fmt.Sprintf("mint ws token: %v", err), false)
116+
continue
117+
}
118+
return s.finishClaim(runCtx, job, &st, env)
119+
}
120+
}
142121

143-
repos := runRepos(job.req)
144-
if len(repos) > 0 {
145-
// Single repo → the checkout dir itself is the agent's cwd (the
146-
// pre-WI-449 layout, unchanged). Multiple repos → each is checked
147-
// out as a sibling dir under a shared per-run workspace root that
148-
// becomes the cwd, so the agent sees every bound repo at once.
149-
multi := len(repos) > 1
150-
if multi {
151-
st.workspaceRoot = s.preparer.RunWorkspaceDir(job.runID)
152-
}
153-
dirNames := repoDirNames(repos)
154-
prepFailed := false
155-
for i, rspec := range repos {
156-
spec := *rspec
157-
if multi {
158-
spec.DestDir = filepath.Join(st.workspaceRoot, dirNames[i])
159-
}
160-
pw, err := s.preparer.Prepare(runCtx, spec, job.runID)
161-
if err != nil {
162-
s.logger.Printf("run service: prepare checkout run=%d repo=%s: %v", job.runID, rspec.RepoSlug, err)
163-
prepFailed = true
164-
break
165-
}
166-
st.repos = append(st.repos, rspec)
167-
st.checkouts = append(st.checkouts, pw)
168-
_ = s.repo.AppendEvent(runCtx, job.runID, "lifecycle", fmt.Sprintf(
169-
`{"phase":"worktree_ready","repo":%q,"path":%q,"branch":%q,"base_commit":%q}`,
170-
rspec.RepoSlug, pw.Path, pw.Branch, pw.BaseCommit))
171-
}
172-
if prepFailed {
173-
// A partial multi-repo checkout is unusable; clean up what we
174-
// prepared and fail visibly. Checkout-prep failure fires the
175-
// post-run hook (matches the prior inline behavior).
176-
for _, pw := range st.checkouts {
177-
_ = s.preparer.Cleanup(context.Background(), pw)
178-
}
179-
if st.workspaceRoot != "" {
180-
s.preparer.CleanupWorkspaceDir(job.runID)
181-
}
182-
s.failClaim(job, cancel, "prepare checkout failed", true)
183-
continue
184-
}
185-
// The primary checkout (index 0) drives the single-repo-compatible
186-
// path: its branch is the grant ref, and it's the cwd for a
187-
// single-repo run.
188-
primary := st.checkouts[0]
189-
st.branch = primary.Branch
190-
st.baseCommit = primary.BaseCommit
191-
if multi {
192-
st.path = st.workspaceRoot
193-
} else {
194-
st.path = primary.Path
122+
func (s *RunService) dequeueClaimJob() (queuedJob, bool) {
123+
select {
124+
case job := <-s.queue:
125+
return job, true
126+
case <-s.shutdownCh:
127+
for {
128+
select {
129+
case job := <-s.queue:
130+
s.finalize(job.runID, models.AgentRunStatusCanceled, "shutdown before admission")
131+
s.wg.Done()
132+
default:
133+
return queuedJob{}, false
195134
}
196135
}
136+
}
137+
}
197138

198-
// Caller-supplied env first; the orchestrator's own injections
199-
// (WS_TOKEN) overwrite on conflict so a confused caller cannot
200-
// smuggle in its own token. The token mint + grant snapshot (bound to
201-
// the minted token, git ref = the prepared worktree branch) is the
202-
// shared preamble the remote claim path also runs (WI-195).
203-
env := make(map[string]string, len(job.req.Env)+1)
204-
for k, v := range job.req.Env {
205-
env[k] = v
139+
func (s *RunService) beginClaim(job queuedJob) (runCtx context.Context, cancel context.CancelFunc, claimed bool) {
140+
runCtx, cancel = context.WithCancel(context.Background())
141+
s.registerCancel(job.runID, cancel)
142+
go func() {
143+
select {
144+
case <-s.shutdownCh:
145+
cancel()
146+
case <-runCtx.Done():
206147
}
207-
if job.req.Token != nil {
208-
// Per-repo push refs: each grant may push only its prepared branch
209-
// (WI-449). Built from the parallel repos/checkouts slices.
210-
refByRepo := make(map[string]string, len(st.checkouts))
211-
for i, pw := range st.checkouts {
212-
refByRepo[st.repos[i].RepoSlug] = pw.Branch
148+
}()
149+
150+
transitioned, err := s.repo.MarkRunningIfQueued(runCtx, job.runID, "", s.now())
151+
if err != nil {
152+
s.logger.Printf("run service: mark running run=%d: %v", job.runID, err)
153+
s.failClaim(job, cancel, fmt.Sprintf("mark running: %v", err), false)
154+
return nil, nil, false
155+
}
156+
if !transitioned {
157+
s.logger.Printf("run service: skipping run=%d: no longer queued at dequeue", job.runID)
158+
cancel()
159+
s.unregisterCancel(job.runID)
160+
s.wg.Done()
161+
return nil, nil, false
162+
}
163+
if err := s.repo.AppendEvent(runCtx, job.runID, "lifecycle", `{"phase":"running"}`); err != nil {
164+
s.logger.Printf("run service: append running event run=%d: %v", job.runID, err)
165+
}
166+
return runCtx, cancel, true
167+
}
168+
169+
func (s *RunService) prepareClaimRepos(ctx context.Context, job queuedJob, state *claimState) error {
170+
repos := runRepos(job.req)
171+
if len(repos) == 0 {
172+
return nil
173+
}
174+
175+
multi := len(repos) > 1
176+
if multi {
177+
state.workspaceRoot = s.preparer.RunWorkspaceDir(job.runID)
178+
}
179+
dirNames := repoDirNames(repos)
180+
for i, repo := range repos {
181+
spec := *repo
182+
if multi {
183+
spec.DestDir = filepath.Join(state.workspaceRoot, dirNames[i])
184+
}
185+
prepared, err := s.preparer.Prepare(ctx, spec, job.runID)
186+
if err != nil {
187+
s.logger.Printf("run service: prepare checkout run=%d repo=%s: %v", job.runID, repo.RepoSlug, err)
188+
for _, checkout := range state.checkouts {
189+
_ = s.preparer.Cleanup(context.Background(), checkout)
213190
}
214-
token, err := s.mintTokenAndGrants(runCtx, job.runID, *job.req.Token, job.req.Grants, refByRepo)
215-
if err != nil {
216-
s.logger.Printf("run service: mint ws token run=%d: %v", job.runID, err)
217-
// Token-mint failure does not fire the hook (matches the
218-
// prior inline behavior).
219-
s.failClaim(job, cancel, fmt.Sprintf("mint ws token: %v", err), false)
220-
continue
191+
if state.workspaceRoot != "" {
192+
s.preparer.CleanupWorkspaceDir(job.runID)
221193
}
222-
env["WS_TOKEN"] = token
223-
applyLLMProxyEnv(env, job.req.Grants, job.runID, token)
194+
return err
224195
}
196+
state.repos = append(state.repos, repo)
197+
state.checkouts = append(state.checkouts, prepared)
198+
_ = s.repo.AppendEvent(ctx, job.runID, "lifecycle", fmt.Sprintf(
199+
`{"phase":"worktree_ready","repo":%q,"path":%q,"branch":%q,"base_commit":%q}`,
200+
repo.RepoSlug, prepared.Path, prepared.Branch, prepared.BaseCommit))
201+
}
202+
203+
primary := state.checkouts[0]
204+
state.branch = primary.Branch
205+
state.baseCommit = primary.BaseCommit
206+
state.path = primary.Path
207+
if multi {
208+
state.path = state.workspaceRoot
209+
}
210+
return nil
211+
}
225212

226-
s.claimsMu.Lock()
227-
s.claims[job.runID] = &st
228-
s.claimsMu.Unlock()
213+
func (s *RunService) buildClaimEnv(ctx context.Context, job queuedJob, state *claimState) (map[string]string, error) {
214+
env := make(map[string]string, len(job.req.Env)+1)
215+
for key, value := range job.req.Env {
216+
env[key] = value
217+
}
218+
if job.req.Token == nil {
219+
return env, nil
220+
}
229221

230-
initialPrompt := s.initialPrompt
231-
if job.req.InitialPrompt != "" {
232-
initialPrompt = job.req.InitialPrompt
233-
}
234-
// Per-binding instructions + skills index ride as a suffix on top of
235-
// whichever base prompt won (WI-258).
236-
initialPrompt += job.req.InitialPromptSuffix
237-
return &ClaimedJob{Spec: JobSpec{RunID: job.runID, WorkspacePath: st.path, Env: env, InitialPrompt: initialPrompt, Kind: job.req.JobKind, Image: job.req.JobImage}, Ctx: runCtx}
222+
refByRepo := make(map[string]string, len(state.checkouts))
223+
for i, checkout := range state.checkouts {
224+
refByRepo[state.repos[i].RepoSlug] = checkout.Branch
225+
}
226+
token, err := s.mintTokenAndGrants(ctx, job.runID, *job.req.Token, job.req.Grants, refByRepo)
227+
if err != nil {
228+
return nil, err
229+
}
230+
env["WS_TOKEN"] = token
231+
applyLLMProxyEnv(env, job.req.Grants, job.runID, token)
232+
return env, nil
233+
}
234+
235+
func (s *RunService) finishClaim(ctx context.Context, job queuedJob, state *claimState, env map[string]string) *ClaimedJob {
236+
s.claimsMu.Lock()
237+
s.claims[job.runID] = state
238+
s.claimsMu.Unlock()
239+
240+
initialPrompt := s.initialPrompt
241+
if job.req.InitialPrompt != "" {
242+
initialPrompt = job.req.InitialPrompt
243+
}
244+
initialPrompt += job.req.InitialPromptSuffix
245+
return &ClaimedJob{
246+
Spec: JobSpec{
247+
RunID: job.runID, WorkspacePath: state.path, Env: env,
248+
InitialPrompt: initialPrompt, Kind: job.req.JobKind, Image: job.req.JobImage,
249+
},
250+
Ctx: ctx,
238251
}
239252
}
240253

‎internal/services/triage_runner.go‎

Lines changed: 22 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ package services
33
import (
44
"context"
55
"encoding/json"
6+
"errors"
67
"fmt"
78
"log"
89
"os"
@@ -56,12 +57,9 @@ func (t *TriageRunner) Run(ctx context.Context, input RunInput, emit EventSink)
5657
if input.Repo == nil {
5758
return t.Inner.Run(ctx, input, emit)
5859
}
59-
if t.TriageBin == "" || t.CacheRoot == "" || t.APIBase == "" {
60-
return RunnerResult{Status: models.AgentRunStatusFailed, Error: "triage runner: TriageBin, CacheRoot and APIBase are required for repo runs"}
61-
}
62-
token := input.Env["WS_TOKEN"]
63-
if token == "" {
64-
return RunnerResult{Status: models.AgentRunStatusFailed, Error: "triage runner: no WS_TOKEN for git-proxy auth"}
60+
token, err := t.repoRunToken(input)
61+
if err != nil {
62+
return RunnerResult{Status: models.AgentRunStatusFailed, Error: err.Error()}
6563
}
6664
owner, repo, ok := splitSlug(input.Repo.Slug)
6765
if !ok {
@@ -128,12 +126,9 @@ func (t *TriageRunner) Run(ctx context.Context, input RunInput, emit EventSink)
128126
// run_service multi-repo path; the broker authorizes each push against that
129127
// repo's grant.
130128
func (t *TriageRunner) runMulti(ctx context.Context, input RunInput, emit EventSink) RunnerResult {
131-
if t.TriageBin == "" || t.CacheRoot == "" || t.APIBase == "" {
132-
return RunnerResult{Status: models.AgentRunStatusFailed, Error: "triage runner: TriageBin, CacheRoot and APIBase are required for repo runs"}
133-
}
134-
token := input.Env["WS_TOKEN"]
135-
if token == "" {
136-
return RunnerResult{Status: models.AgentRunStatusFailed, Error: "triage runner: no WS_TOKEN for git-proxy auth"}
129+
token, err := t.repoRunToken(input)
130+
if err != nil {
131+
return RunnerResult{Status: models.AgentRunStatusFailed, Error: err.Error()}
137132
}
138133
tokenFile, cleanupToken, err := writeTokenFile(token, input.RunID)
139134
if err != nil {
@@ -208,24 +203,22 @@ func (t *TriageRunner) runMulti(ctx context.Context, input RunInput, emit EventS
208203
// triageDirNames mirrors repoDirNames for the remote runner: a unique sibling
209204
// dir name per repo (last slug segment, numeric-suffixed on collision).
210205
func triageDirNames(repos []JobRepo) []string {
211-
names := make([]string, len(repos))
212-
seen := map[string]int{}
213-
for i, r := range repos {
214-
base := r.Slug
215-
if idx := strings.LastIndex(base, "/"); idx >= 0 && idx < len(base)-1 {
216-
base = base[idx+1:]
217-
}
218-
if base == "" {
219-
base = fmt.Sprintf("repo%d", i)
220-
}
221-
name := base
222-
if n, ok := seen[base]; ok {
223-
name = fmt.Sprintf("%s-%d", base, n+1)
224-
}
225-
seen[base]++
226-
names[i] = name
206+
slugs := make([]string, len(repos))
207+
for i, repo := range repos {
208+
slugs[i] = repo.Slug
209+
}
210+
return repoWorkspaceDirNames(slugs)
211+
}
212+
213+
func (t *TriageRunner) repoRunToken(input RunInput) (string, error) {
214+
if t.TriageBin == "" || t.CacheRoot == "" || t.APIBase == "" {
215+
return "", errors.New("triage runner: TriageBin, CacheRoot and APIBase are required for repo runs")
216+
}
217+
token := input.Env["WS_TOKEN"]
218+
if token == "" {
219+
return "", errors.New("triage runner: no WS_TOKEN for git-proxy auth")
227220
}
228-
return names
221+
return token, nil
229222
}
230223

231224
func (t *TriageRunner) prepare(ctx context.Context, jr JobRepo, runID int, proxyURL, tokenFile, destDir string) (triagePrepareOut, error) {

0 commit comments

Comments
 (0)