Movatterモバイル変換

[0]ホーム
URL:

Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly

 Sign up
Appearance settings

Commit00dd127

Browse files
committed
feat: add timeout support to workspace bash tool
Change-Id: I996cbde4a50debb54a0a95ca5a067781719fa25aSigned-off-by: Thomas Kosiewski <tk@coder.com>
1 parent5c31b98 commit00dd127

File tree

2 files changed

+321
-11
lines changed

2 files changed

+321
-11
lines changed

‎codersdk/toolsdk/bash.go‎

Lines changed: 141 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,14 @@
11
package toolsdk
22

33
import (
4+
"bytes"
45
"context"
56
"errors"
67
"fmt"
78
"io"
89
"strings"
10+
"sync"
11+
"time"
912

1013
gossh"golang.org/x/crypto/ssh"
1114
"golang.org/x/xerrors"
@@ -20,6 +23,7 @@ import (
2023
typeWorkspaceBashArgsstruct {
2124
Workspacestring`json:"workspace"`
2225
Commandstring`json:"command"`
26+
TimeoutMsint`json:"timeout_ms,omitempty"`
2327
}
2428

2529
typeWorkspaceBashResultstruct {
@@ -43,9 +47,12 @@ The workspace parameter supports various formats:
4347
- workspace.agent (specific agent)
4448
- owner/workspace.agent
4549
50+
The timeout_ms parameter specifies the command timeout in milliseconds (defaults to 60000ms, maximum of 300000ms).
51+
If the command times out, all output captured up to that point is returned with a cancellation message.
52+
4653
Examples:
4754
- workspace: "my-workspace", command: "ls -la"
48-
- workspace: "john/dev-env", command: "git status"
55+
- workspace: "john/dev-env", command: "git status", timeout_ms: 30000
4956
- workspace: "my-workspace.main", command: "docker ps"`,
5057
Schema: aisdk.Schema{
5158
Properties:map[string]any{
@@ -57,18 +64,27 @@ Examples:
5764
"type":"string",
5865
"description":"The bash command to execute in the workspace.",
5966
},
67+
"timeout_ms":map[string]any{
68+
"type":"integer",
69+
"description":"Command timeout in milliseconds. Defaults to 60000ms (60 seconds) if not specified.",
70+
"default":60000,
71+
"minimum":1,
72+
},
6073
},
6174
Required: []string{"workspace","command"},
6275
},
6376
},
64-
Handler:func(ctx context.Context,depsDeps,argsWorkspaceBashArgs) (WorkspaceBashResult,error) {
77+
Handler:func(ctx context.Context,depsDeps,argsWorkspaceBashArgs) (resWorkspaceBashResult,errerror) {
6578
ifargs.Workspace=="" {
6679
returnWorkspaceBashResult{},xerrors.New("workspace name cannot be empty")
6780
}
6881
ifargs.Command=="" {
6982
returnWorkspaceBashResult{},xerrors.New("command cannot be empty")
7083
}
7184

85+
ctx,cancel:=context.WithTimeoutCause(ctx,5*time.Minute,xerrors.New("MCP handler timeout after 5 min"))
86+
defercancel()
87+
7288
// Normalize workspace input to handle various formats
7389
workspaceName:=NormalizeWorkspaceInput(args.Workspace)
7490

@@ -119,23 +135,42 @@ Examples:
119135
}
120136
defersession.Close()
121137

122-
// Execute command and capture output
123-
output,err:=session.CombinedOutput(args.Command)
138+
// Set default timeout if not specified (60 seconds)
139+
timeoutMs:=args.TimeoutMs
140+
iftimeoutMs<=0 {
141+
timeoutMs=60000
142+
}
143+
144+
// Create context with timeout
145+
ctx,cancel=context.WithTimeout(ctx,time.Duration(timeoutMs)*time.Millisecond)
146+
defercancel()
147+
148+
// Execute command with timeout handling
149+
output,err:=executeCommandWithTimeout(ctx,session,args.Command)
124150
outputStr:=strings.TrimSpace(string(output))
125151

152+
// Handle command execution results
126153
iferr!=nil {
127-
// Check ifit's an SSH exit error to gettheexit code
128-
varexitErr*gossh.ExitError
129-
iferrors.As(err,&exitErr) {
154+
// Check if thecommand timed out
155+
iferrors.Is(context.Cause(ctx),context.DeadlineExceeded) {
156+
outputStr+="\nCommand canceled due to timeout"
130157
returnWorkspaceBashResult{
131158
Output:outputStr,
132-
ExitCode:exitErr.ExitStatus(),
159+
ExitCode:124,
133160
},nil
134161
}
135-
// For other errors, return exit code 1
162+
163+
// Extract exit code from SSH error if available
164+
exitCode:=1
165+
varexitErr*gossh.ExitError
166+
iferrors.As(err,&exitErr) {
167+
exitCode=exitErr.ExitStatus()
168+
}
169+
170+
// For other errors, use standard timeout or generic error code
136171
returnWorkspaceBashResult{
137172
Output:outputStr,
138-
ExitCode:1,
173+
ExitCode:exitCode,
139174
},nil
140175
}
141176

@@ -292,3 +327,99 @@ func NormalizeWorkspaceInput(input string) string {
292327

293328
returnnormalized
294329
}
330+
331+
// executeCommandWithTimeout executes a command with timeout support
332+
funcexecuteCommandWithTimeout(ctx context.Context,session*gossh.Session,commandstring) ([]byte,error) {
333+
// Set up pipes to capture output
334+
stdoutPipe,err:=session.StdoutPipe()
335+
iferr!=nil {
336+
returnnil,xerrors.Errorf("failed to create stdout pipe: %w",err)
337+
}
338+
339+
stderrPipe,err:=session.StderrPipe()
340+
iferr!=nil {
341+
returnnil,xerrors.Errorf("failed to create stderr pipe: %w",err)
342+
}
343+
344+
// Start the command
345+
iferr:=session.Start(command);err!=nil {
346+
returnnil,xerrors.Errorf("failed to start command: %w",err)
347+
}
348+
349+
// Create a thread-safe buffer for combined output
350+
varoutput bytes.Buffer
351+
varmu sync.Mutex
352+
safeWriter:=&syncWriter{w:&output,mu:&mu}
353+
354+
// Use io.MultiWriter to combine stdout and stderr
355+
multiWriter:=io.MultiWriter(safeWriter)
356+
357+
// Channel to signal when command completes
358+
done:=make(chanerror,1)
359+
360+
// Start goroutine to copy output and wait for completion
361+
gofunc() {
362+
// Copy stdout and stderr concurrently
363+
varwg sync.WaitGroup
364+
wg.Add(2)
365+
366+
gofunc() {
367+
deferwg.Done()
368+
_,_=io.Copy(multiWriter,stdoutPipe)
369+
}()
370+
371+
gofunc() {
372+
deferwg.Done()
373+
_,_=io.Copy(multiWriter,stderrPipe)
374+
}()
375+
376+
// Wait for all output to be copied
377+
wg.Wait()
378+
379+
// Wait for the command to complete
380+
done<-session.Wait()
381+
}()
382+
383+
// Wait for either completion or context cancellation
384+
select {
385+
caseerr:=<-done:
386+
// Command completed normally
387+
returnsafeWriter.Bytes(),err
388+
case<-ctx.Done():
389+
// Context was canceled (timeout or other cancellation)
390+
// Close the session to stop the command
391+
_=session.Close()
392+
393+
// Give a brief moment to collect any remaining output
394+
timer:=time.NewTimer(50*time.Millisecond)
395+
defertimer.Stop()
396+
397+
select {
398+
case<-timer.C:
399+
// Timer expired, return what we have
400+
caseerr:=<-done:
401+
// Command finished during grace period
402+
returnsafeWriter.Bytes(),err
403+
}
404+
405+
returnsafeWriter.Bytes(),context.Cause(ctx)
406+
}
407+
}
408+
409+
// syncWriter is a thread-safe writer
410+
typesyncWriterstruct {
411+
w*bytes.Buffer
412+
mu*sync.Mutex
413+
}
414+
415+
func (sw*syncWriter)Write(p []byte) (nint,errerror) {
416+
sw.mu.Lock()
417+
defersw.mu.Unlock()
418+
returnsw.w.Write(p)
419+
}
420+
421+
func (sw*syncWriter)Bytes() []byte {
422+
sw.mu.Lock()
423+
defersw.mu.Unlock()
424+
returnsw.w.Bytes()
425+
}

0 commit comments

Comments
 (0)
[8]ページ先頭
©2009-2025 Movatter.jp