Skip to content

Commit 43b9858

Browse files
zpzjzjclaude
andcommitted
refactor(custom-engine): extract customTransport interface with localTransport (#51)
Introduce an internal `customTransport` seam in the agent package so a future HTTP transport drops in without rewriting the local-transport logic. `CustomAgent.Run` now does three transport-agnostic steps: `selectTransport` (local vs the http not-yet-implemented stub), `prepareRun` (build the timeout, template vars, and SessionInput), and `assembleResult` (map the outcome to a SessionResult with the shared exit-code/error semantics). The local-specific work — resolving I/O files, persisting the input, clearing a stale output, building+executing the command, and reading the result payload — moves behind `localTransport.run` in custom_local.go, which returns a uniform `transportOutcome`. Behavior-preserving: - Setup errors (pre-exec) still return errorResult(0); exec-phase failures are carried in transportOutcome.execErr so a partial result is still assembled, including the post-timeout fresh-context output read. - transportOutcome keeps `raw` and `stdout` distinct so a bookkeeping output_file is never graded as a text-format answer. - Framework input/output file registration keeps the "input always, output iff produced-or-cleared" rule; registerFrameworkIO becomes appendFrameworkFiles (no flag parameter, so the revive suppression is dropped). - The http transport still returns the exact "not yet implemented" error. All internal/agent tests are unchanged and pass (the safety net); make test, make verify (0 issues), and the e2e custom-engine suite are green. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
1 parent 57d057f commit 43b9858

2 files changed

Lines changed: 240 additions & 140 deletions

File tree

internal/agent/custom.go

Lines changed: 129 additions & 140 deletions
Original file line numberDiff line numberDiff line change
@@ -92,38 +92,103 @@ func (a *CustomAgent) Check(_ context.Context, _ Runtime) error { return nil }
9292
// explicitly via ${api_key}, so skill-up does not pre-validate them.
9393
func (a *CustomAgent) CheckCredentials(_ context.Context) error { return nil }
9494

95-
// Run executes the custom engine for a single case.
95+
// customTransport executes a prepared custom-engine run and returns a uniform
96+
// transportOutcome that the shared assembler (CustomAgent.assembleResult) maps
97+
// to a *SessionResult. The local transport reads result files and the process
98+
// exit; a future http transport will map an HTTP response body and status onto
99+
// the same shape.
100+
//
101+
// run's returned error is a SETUP error raised before the engine executes (e.g.
102+
// the command failed to render, or the input file could not be written): the
103+
// run is abandoned and CustomAgent.Run reports an errorResult. An
104+
// execution-phase failure — the engine ran but exited non-zero, timed out, or
105+
// kept stdout open past WaitDelay — is carried in transportOutcome.execErr
106+
// instead, so a partial result can still be assembled.
107+
type customTransport interface {
108+
run(ctx context.Context, rt Runtime, opts ExecOptions, prep *customRunPrep) (*transportOutcome, error)
109+
}
110+
111+
// customRunPrep is the transport-agnostic state CustomAgent.Run builds once,
112+
// before dispatching to a transport. Every transport consumes the same prepared
113+
// session input and template variable set. start is captured at preparation
114+
// time so the reported duration spans preparation through result assembly.
115+
type customRunPrep struct {
116+
custom *config.CustomEngineConfig
117+
vars map[string]string
118+
sessJSON []byte
119+
timeoutSec int
120+
messages []transcript.Message
121+
start time.Time
122+
}
123+
124+
// transportOutcome is what a transport hands back to the shared assembler.
125+
// exitCode and stderr are the process exit code and stderr for the local
126+
// transport (a future http transport maps HTTP status and error body onto
127+
// them). raw is the result payload graded for session_result responses; stdout
128+
// is the text-format final-message source (the local transport keeps them
129+
// distinct so a bookkeeping output_file is never graded as a text answer).
130+
// frameworkFiles are the framework-written paths (input/output files) to record
131+
// in GeneratedFiles so the workspace-diff collector excludes them.
132+
type transportOutcome struct {
133+
raw string
134+
stdout string
135+
exitCode int
136+
stderr string
137+
execErr error
138+
frameworkFiles []string
139+
}
140+
141+
// Run executes the custom engine for a single case: it builds the shared
142+
// session input, selects the transport, runs it, and assembles the result.
96143
func (a *CustomAgent) Run(ctx context.Context, rt Runtime, opts ExecOptions, messages []transcript.Message) (*SessionResult, error) {
97144
custom := a.Cfg.Custom
98145
if custom == nil {
99146
return a.errorResult(0), errors.New("custom engine config is missing")
100147
}
148+
transport, err := a.selectTransport(custom)
149+
if err != nil {
150+
return a.errorResult(0), err
151+
}
152+
prep, err := a.prepareRun(rt, opts, messages, custom)
153+
if err != nil {
154+
return a.errorResult(0), err
155+
}
156+
outcome, err := transport.run(ctx, rt, opts, prep)
157+
if err != nil {
158+
return a.errorResult(0), err
159+
}
160+
return a.assembleResult(ctx, rt, opts, prep, outcome)
161+
}
162+
163+
// selectTransport maps engine.custom.transport to its implementation. The http
164+
// transport is designed but not yet implemented; selecting it returns the
165+
// explicit error rather than a transport. When http lands, this case returns an
166+
// &httpTransport{} and CustomAgent.Run is otherwise unchanged.
167+
func (a *CustomAgent) selectTransport(custom *config.CustomEngineConfig) (customTransport, error) {
101168
switch custom.Transport {
102169
case customTransportLocal:
103-
return a.runLocal(ctx, rt, opts, messages, custom)
170+
return &localTransport{a: a}, nil
104171
case customTransportHTTP:
105-
return a.errorResult(0), errors.New("custom engine http transport is not yet implemented")
172+
return nil, errors.New("custom engine http transport is not yet implemented")
106173
default:
107-
return a.errorResult(0), fmt.Errorf("custom engine transport %q is not supported", custom.Transport)
174+
return nil, fmt.Errorf("custom engine transport %q is not supported", custom.Transport)
108175
}
109176
}
110177

111-
func (a *CustomAgent) runLocal(ctx context.Context, rt Runtime, opts ExecOptions, messages []transcript.Message, custom *config.CustomEngineConfig) (*SessionResult, error) {
112-
// Guard against a nil local block: validation skips engine.custom for
113-
// built-in engines, so a `--engine` override can reach here unvalidated.
114-
if custom.Local == nil {
115-
return a.errorResult(0), errors.New("engine.custom.local.command is required when transport is local")
116-
}
117-
178+
// prepareRun builds the transport-agnostic run state: the effective timeout, the
179+
// full template variable set, and the marshaled SessionInput. start is captured
180+
// here so the reported duration spans preparation through result assembly,
181+
// matching the pre-refactor behavior.
182+
func (a *CustomAgent) prepareRun(rt Runtime, opts ExecOptions, messages []transcript.Message, custom *config.CustomEngineConfig) (*customRunPrep, error) {
118183
start := time.Now()
119184
// engine.custom.timeout_seconds is the per-engine deadline; opts.TimeoutSec
120185
// is the outer case deadline (the evaluator already wraps ctx with it).
121-
// When both are positive, the effective deadline is the smaller of the
122-
// two — exceeding the case deadline would have the case context kill the
123-
// process while ${timeout_seconds} and SessionInput still advertised a
124-
// longer timeout to the agent, producing premature termination with
125-
// inconsistent metadata. Take the min so the value handed to the agent
126-
// reflects the real wall-clock budget.
186+
// When both are positive, the effective deadline is the smaller of the two
187+
// — exceeding the case deadline would have the case context kill the process
188+
// while ${timeout_seconds} and SessionInput still advertised a longer
189+
// timeout to the agent, producing premature termination with inconsistent
190+
// metadata. Take the min so the value handed to the agent reflects the real
191+
// wall-clock budget.
127192
timeoutSec := custom.TimeoutSeconds
128193
switch {
129194
case timeoutSec <= 0:
@@ -132,87 +197,37 @@ func (a *CustomAgent) runLocal(ctx context.Context, rt Runtime, opts ExecOptions
132197
timeoutSec = opts.TimeoutSec
133198
}
134199

135-
// Build the full template variable set first: kwargs may reference
136-
// built-in variables (e.g. ${case_id}), and the I/O path templates may in
137-
// turn reference ${kwargs.<key>}, so kwargs must be resolved before paths.
200+
// Build the full template variable set first: kwargs may reference built-in
201+
// variables (e.g. ${case_id}), and the I/O path templates may in turn
202+
// reference ${kwargs.<key>}, so kwargs must be resolved before paths.
138203
baseVars := a.buildBaseVars(rt, opts, messages, timeoutSec)
139204

140205
renderedKwargs, err := renderTemplateMap(custom.Kwargs, baseVars)
141206
if err != nil {
142-
return a.errorResult(0), fmt.Errorf("render custom.kwargs: %w", err)
207+
return nil, fmt.Errorf("render custom.kwargs: %w", err)
143208
}
144209

145210
sess := a.buildSessionInput(rt, opts, messages, renderedKwargs, timeoutSec)
146211
// Marshal the session input once and share the bytes between the template
147-
// expansion (${session_input}) and the on-disk input file, avoiding a
148-
// second copy of what for long transcripts can be tens of MB per case.
212+
// expansion (${session_input}) and the on-disk input file, avoiding a second
213+
// copy of what for long transcripts can be tens of MB per case.
149214
sessJSON, err := json.Marshal(sess)
150215
if err != nil {
151-
return a.errorResult(0), fmt.Errorf("marshal session input: %w", err)
216+
return nil, fmt.Errorf("marshal session input: %w", err)
152217
}
153218
vars, err := a.completeTemplateVars(baseVars, renderedKwargs, sess, sessJSON)
154219
if err != nil {
155-
return a.errorResult(0), err
220+
return nil, err
156221
}
157222

158-
// Resolve the I/O paths with the full variable set (so they may use
159-
// ${kwargs.<key>}), then expose the resolved paths to command rendering.
160-
inputFile, outputFile, err := resolveCustomIOFiles(rt, custom, vars)
161-
if err != nil {
162-
return a.errorResult(0), err
163-
}
164-
vars["input_file"], vars["output_file"] = inputFile, outputFile
165-
166-
if err := persistRuntimeArtifact(ctx, rt, inputFile, string(sessJSON)); err != nil {
167-
return a.errorResult(0), fmt.Errorf("write session input: %w", err)
168-
}
169-
170-
// Remove any stale output file so a result left by a fixture or a previous
171-
// run is never mistaken for this invocation's output. Only an explicitly
172-
// configured output_file is cleared — the default ${output_file} path is
173-
// left untouched, as it may be ordinary fixture input the agent reads.
174-
// Also skip the clear when output_file resolves to the same path as
175-
// input_file, so the SessionInput just written is not self-deleted.
176-
clearedStaleOutput := false
177-
if custom.Local.OutputFile != "" && filepath.Clean(outputFile) != filepath.Clean(inputFile) {
178-
clearedStaleOutput = a.clearStaleOutputFile(ctx, rt, outputFile)
179-
}
180-
181-
cmd, execOpts, err := a.buildLocalExec(ctx, rt, opts, custom, vars, timeoutSec)
182-
if err != nil {
183-
return a.errorResult(0), err
184-
}
185-
186-
// Enforce the custom timeout via a context deadline: some runtimes
187-
// (e.g. NoneRuntime) do not honor ExecOptions.TimeoutSec directly.
188-
execCtx := ctx
189-
if timeoutSec > 0 {
190-
var cancel context.CancelFunc
191-
execCtx, cancel = context.WithTimeout(ctx, time.Duration(timeoutSec)*time.Second)
192-
defer cancel()
193-
}
194-
195-
result, execErr := rt.Exec(execCtx, cmd, execOpts)
196-
return a.finishLocal(ctx, rt, opts, custom, result, execErr, localRunContext{
197-
InputFile: inputFile,
198-
OutputFile: outputFile,
199-
ClearedStaleOutput: clearedStaleOutput,
200-
Start: start,
201-
Messages: messages,
202-
})
203-
}
204-
205-
// localRunContext bundles the local-run bookkeeping that finishLocal needs:
206-
// the resolved input/output paths, whether the pre-run cleanup actually
207-
// removed a stale output file, the run start time (for duration), and the
208-
// caller's transcript messages. Carrying these as one struct keeps
209-
// finishLocal's signature manageable as more bookkeeping accrues.
210-
type localRunContext struct {
211-
InputFile string
212-
OutputFile string
213-
ClearedStaleOutput bool
214-
Start time.Time
215-
Messages []transcript.Message
223+
return &customRunPrep{
224+
custom: custom,
225+
vars: vars,
226+
sessJSON: sessJSON,
227+
timeoutSec: timeoutSec,
228+
messages: messages,
229+
start: start,
230+
}, nil
216231
}
217232

218233
// clearStaleOutputFile removes a pre-existing output file before the run and
@@ -291,72 +306,53 @@ func (a *CustomAgent) buildLocalExec(ctx context.Context, rt Runtime, opts ExecO
291306
return strings.Join(parts, " "), execOpts, nil
292307
}
293308

294-
// finishLocal turns a runtime exec result into a SessionResult.
295-
func (a *CustomAgent) finishLocal(ctx context.Context, rt Runtime, opts ExecOptions, custom *config.CustomEngineConfig, result ExecResult, execErr error, rc localRunContext) (*SessionResult, error) {
296-
durationMs := time.Since(rc.Start).Milliseconds()
297-
298-
// The command may already have emitted a valid result even when execErr
299-
// is non-nil — a timeout, exec.ErrWaitDelay (a backgrounded child kept
300-
// stdout open past WaitDelay), or a non-zero exit after printing output.
301-
// Always attempt the read so judges and expect checks can inspect the
302-
// partial answer; if nothing was produced, raw stays empty and the
303-
// parse-error path below takes over.
304-
var (
305-
raw string
306-
outputFileProduced bool
307-
)
308-
readCtx := ctx
309-
if execErr != nil {
310-
// ctx may have been canceled (timeout/cancel) or unrelated to the
311-
// failure; use a fresh context so a partial result is still
312-
// recoverable on remote runtimes whose DownloadFile honors ctx.
313-
var cancel context.CancelFunc
314-
readCtx, cancel = context.WithTimeout(context.WithoutCancel(ctx), customArtifactReadTimeout)
315-
defer cancel()
316-
}
317-
raw, outputFileProduced = a.readRawResult(readCtx, rt, custom, result, rc.OutputFile)
318-
319-
res, parseErr := a.buildResult(ctx, rt, opts, custom, raw, result, durationMs, rc.Messages)
309+
// assembleResult maps a transport's outcome onto a *SessionResult, applying the
310+
// shared exit-code and error semantics. It is transport-agnostic: the local and
311+
// (future) http transports differ only in how they produce the
312+
// transportOutcome. The raw-result read happens inside the transport, so this
313+
// shared step never touches the output file or the process directly.
314+
func (a *CustomAgent) assembleResult(ctx context.Context, rt Runtime, opts ExecOptions, prep *customRunPrep, o *transportOutcome) (*SessionResult, error) {
315+
durationMs := time.Since(prep.start).Milliseconds()
320316

321-
usedOutputFile := outputFileProduced || rc.ClearedStaleOutput
317+
res, parseErr := a.buildResult(ctx, rt, opts, prep.custom, o.raw, o.stdout, o.exitCode, o.stderr, durationMs, prep.messages)
322318

323-
if execErr != nil {
319+
if o.execErr != nil {
324320
switch {
325321
case parseErr != nil:
326322
// Nothing usable was produced before the failure.
327-
res = a.errorResult(result.ExitCode)
323+
res = a.errorResult(o.exitCode)
328324
res.DurationMs = durationMs
329-
res.Transcript = minimalCustomTranscript(rc.Messages, "")
325+
res.Transcript = minimalCustomTranscript(prep.messages, "")
330326
case res.ExitCode == 0:
331327
// The run was interrupted; a parsed exit_code 0 must not let the
332328
// evaluator treat an interrupted run as a success.
333-
res.ExitCode = result.ExitCode
329+
res.ExitCode = o.exitCode
334330
if res.ExitCode == 0 {
335331
res.ExitCode = 1
336332
}
337333
}
338334
if res.Stderr == "" {
339-
res.Stderr = a.maskAPIKey(result.Stderr)
335+
res.Stderr = a.maskAPIKey(o.stderr)
340336
} else {
341337
res.Stderr = a.maskAPIKey(res.Stderr)
342338
}
343-
a.registerFrameworkIO(res, rc.InputFile, rc.OutputFile, usedOutputFile)
344-
return res, fmt.Errorf("custom engine run failed: %w", execErr)
339+
a.appendFrameworkFiles(res, o.frameworkFiles)
340+
return res, fmt.Errorf("custom engine run failed: %w", o.execErr)
345341
}
346342

347-
a.registerFrameworkIO(res, rc.InputFile, rc.OutputFile, usedOutputFile)
343+
a.appendFrameworkFiles(res, o.frameworkFiles)
348344
// A non-zero process exit is a failed run even when the engine's JSON
349345
// reports exit_code 0 (e.g. a wrapper that crashed after printing output)
350346
// or never emitted parseable JSON at all. Reflect the real process
351347
// exit/stderr so reports surface the actual command failure.
352-
if result.ExitCode != 0 {
353-
res.ExitCode = result.ExitCode
348+
if o.exitCode != 0 {
349+
res.ExitCode = o.exitCode
354350
if res.Stderr == "" {
355-
res.Stderr = a.maskAPIKey(result.Stderr)
351+
res.Stderr = a.maskAPIKey(o.stderr)
356352
} else {
357353
res.Stderr = a.maskAPIKey(res.Stderr)
358354
}
359-
return res, fmt.Errorf("custom engine command exited %d: %s", result.ExitCode, a.maskAPIKey(result.Stderr))
355+
return res, fmt.Errorf("custom engine command exited %d: %s", o.exitCode, a.maskAPIKey(o.stderr))
360356
}
361357
if parseErr != nil {
362358
res.Stderr = a.maskAPIKey(res.Stderr)
@@ -374,43 +370,36 @@ func (a *CustomAgent) finishLocal(ctx context.Context, rt Runtime, opts ExecOpti
374370
// configured response_format. The text format never errors and always grades
375371
// stdout (never the output file); session_result returns a parse error when
376372
// the payload is missing or malformed.
377-
func (a *CustomAgent) buildResult(ctx context.Context, rt Runtime, opts ExecOptions, custom *config.CustomEngineConfig, raw string, result ExecResult, durationMs int64, messages []transcript.Message) (*SessionResult, error) {
373+
func (a *CustomAgent) buildResult(ctx context.Context, rt Runtime, opts ExecOptions, custom *config.CustomEngineConfig, raw, stdout string, exitCode int, stderr string, durationMs int64, messages []transcript.Message) (*SessionResult, error) {
378374
if customResponseFormat(custom) == customResponseText {
379375
// Per the contract, text responses come from stdout; an output file
380376
// produced for bookkeeping must not be graded as the final answer.
381377
// Mask before constructing the transcript so the assistant-reply
382378
// message embedded in it never carries an unmasked value.
383-
finalMsg := a.maskAPIKey(strings.TrimSpace(result.Stdout))
379+
finalMsg := a.maskAPIKey(strings.TrimSpace(stdout))
384380
return &SessionResult{
385381
Engine: a.Name(),
386-
ExitCode: result.ExitCode,
382+
ExitCode: exitCode,
387383
DurationMs: durationMs,
388384
FinalMessage: finalMsg,
389-
Stderr: a.maskAPIKey(result.Stderr),
385+
Stderr: a.maskAPIKey(stderr),
390386
Transcript: minimalCustomTranscript(messages, finalMsg),
391387
Artifacts: &SessionArtifacts{},
392388
}, nil
393389
}
394390
return a.parseSessionResult(ctx, rt, opts, raw, durationMs, messages)
395391
}
396392

397-
// registerFrameworkIO records the framework-written input/output files in
398-
// GeneratedFiles so the workspace-diff collector excludes them (they would
399-
// otherwise show up as user changes) and they are archived for debugging.
400-
// The bool gates registration of outputFile to match the "produced or cleared"
401-
// rule the caller computes.
402-
//
403-
//revive:disable-next-line:flag-parameter
404-
func (a *CustomAgent) registerFrameworkIO(res *SessionResult, inputFile, outputFile string, usedOutputFile bool) {
393+
// appendFrameworkFiles records the framework-written files in GeneratedFiles so
394+
// the workspace-diff collector excludes them (they would otherwise show up as
395+
// user changes) and they are archived for debugging. The transport decides
396+
// which files belong in the slice — for the local transport, the input file
397+
// always, and the output file only when it was produced or cleared.
398+
func (a *CustomAgent) appendFrameworkFiles(res *SessionResult, files []string) {
405399
if res.Artifacts == nil {
406400
res.Artifacts = &SessionArtifacts{}
407401
}
408-
if inputFile != "" {
409-
res.Artifacts.GeneratedFiles = append(res.Artifacts.GeneratedFiles, inputFile)
410-
}
411-
if usedOutputFile && outputFile != "" {
412-
res.Artifacts.GeneratedFiles = append(res.Artifacts.GeneratedFiles, outputFile)
413-
}
402+
res.Artifacts.GeneratedFiles = append(res.Artifacts.GeneratedFiles, files...)
414403
}
415404

416405
// readRawResult reads the result payload from the output file, or from stdout.

0 commit comments

Comments
 (0)