feat(shell): add blocking wait option to job_output tool (#2189)

Sean Porter created 1 week ago

Change summary

internal/agent/common_test.go      |  9 +++------
internal/agent/tools/job_output.go |  5 +++++
internal/agent/tools/job_output.md |  3 +++
internal/shell/background.go       |  9 +++++++++
internal/shell/background_test.go  | 25 +++++++++++++++++++++++++
5 files changed, 45 insertions(+), 6 deletions(-)

Detailed changes

internal/agent/common_test.go

@@ -182,12 +182,9 @@ func coderAgent(r *vcr.Recorder, env fakeEnv, large, small fantasy.LanguageModel
 		GeneratedWith: true,
 	}
 
-	// Clear skills paths to ensure test reproducibility - user's skills
-	// would be included in prompt and break VCR cassette matching.
-	cfg.Options.SkillsPaths = []string{}
-
-	// Clear LSP config to ensure test reproducibility - user's LSP config
-	// would be included in prompt and break VCR cassette matching.
+	// Clear some fields to avoid issues with VCR cassette matching.
+	cfg.Options.SkillsPaths = nil
+	cfg.Options.ContextPaths = nil
 	cfg.LSP = nil
 
 	systemPrompt, err := prompt.Build(context.TODO(), large.Provider(), large.Model(), *cfg)

internal/agent/tools/job_output.go

@@ -19,6 +19,7 @@ var jobOutputDescription []byte
 
 type JobOutputParams struct {
 	ShellID string `json:"shell_id" description:"The ID of the background shell to retrieve output from"`
+	Wait    bool   `json:"wait" description:"If true, block until the background shell completes before returning output"`
 }
 
 type JobOutputResponseMetadata struct {
@@ -44,6 +45,10 @@ func NewJobOutputTool() fantasy.AgentTool {
 				return fantasy.NewTextErrorResponse(fmt.Sprintf("background shell not found: %s", params.ShellID)), nil
 			}
 
+			if params.Wait {
+				bgShell.WaitContext(ctx)
+			}
+
 			stdout, stderr, done, err := bgShell.GetOutput()
 
 			var outputParts []string

internal/agent/tools/job_output.md

@@ -4,16 +4,19 @@ Retrieves the current output from a background shell.
 - Provide the shell ID returned from a background bash execution
 - Returns the current stdout and stderr output
 - Indicates whether the shell has completed execution
+- Set wait=true to block until the shell completes or the request context is done
 </usage>
 
 <features>
 - View output from running background processes
 - Check if background process has completed
 - Get cumulative output from process start
+- Optionally wait for process completion (returns early on context cancel)
 </features>
 
 <tips>
 - Use this to monitor long-running processes
 - Check the 'done' status to see if process completed
 - Can be called multiple times to view incremental output
+- Use wait=true when you need the final output and exit status (or current output if the request cancels)
 </tips>

internal/shell/background.go

@@ -234,3 +234,12 @@ func (bs *BackgroundShell) IsDone() bool {
 func (bs *BackgroundShell) Wait() {
 	<-bs.done
 }
+
+func (bs *BackgroundShell) WaitContext(ctx context.Context) bool {
+	select {
+	case <-bs.done:
+		return true
+	case <-ctx.Done():
+		return false
+	}
+}

internal/shell/background_test.go

@@ -307,3 +307,28 @@ func TestBackgroundShellManager_KillAll_Timeout(t *testing.T) {
 	// Must return promptly after timeout, not hang for 60 seconds.
 	require.Less(t, elapsed, 2*time.Second)
 }
+
+func TestBackgroundShell_WaitContext_Completed(t *testing.T) {
+	t.Parallel()
+
+	done := make(chan struct{})
+	close(done)
+
+	bgShell := &BackgroundShell{done: done}
+
+	ctx, cancel := context.WithTimeout(t.Context(), time.Second)
+	t.Cleanup(cancel)
+
+	require.True(t, bgShell.WaitContext(ctx))
+}
+
+func TestBackgroundShell_WaitContext_Canceled(t *testing.T) {
+	t.Parallel()
+
+	bgShell := &BackgroundShell{done: make(chan struct{})}
+
+	ctx, cancel := context.WithCancel(t.Context())
+	cancel()
+
+	require.False(t, bgShell.WaitContext(ctx))
+}