feat: add coder_workspace_write_file MCP tool (#19591)

This commit is contained in:
Asher
2025-09-11 12:17:15 -08:00
committed by GitHub
parent eec6c8c120
commit d5a02d570f
7 changed files with 346 additions and 2 deletions
+1
View File
@@ -61,6 +61,7 @@ func (a *agent) apiHandler() http.Handler {
r.Get("/api/v0/netcheck", a.HandleNetcheck)
r.Post("/api/v0/list-directory", a.HandleLS)
r.Get("/api/v0/read-file", a.HandleReadFile)
r.Post("/api/v0/write-file", a.HandleWriteFile)
r.Get("/debug/logs", a.HandleHTTPDebugLogs)
r.Get("/debug/magicsock", a.HandleHTTPDebugMagicsock)
r.Get("/debug/magicsock/debug-logging/{state}", a.HandleHTTPMagicsockDebugLoggingState)
+72 -1
View File
@@ -3,12 +3,14 @@ package agent
import (
"context"
"errors"
"fmt"
"io"
"mime"
"net/http"
"os"
"path/filepath"
"strconv"
"syscall"
"golang.org/x/xerrors"
@@ -17,6 +19,8 @@ import (
"github.com/coder/coder/v2/codersdk"
)
type HTTPResponseCode = int
func (a *agent) HandleReadFile(rw http.ResponseWriter, r *http.Request) {
ctx := r.Context()
@@ -43,7 +47,7 @@ func (a *agent) HandleReadFile(rw http.ResponseWriter, r *http.Request) {
}
}
func (a *agent) streamFile(ctx context.Context, rw http.ResponseWriter, path string, offset, limit int64) (int, error) {
func (a *agent) streamFile(ctx context.Context, rw http.ResponseWriter, path string, offset, limit int64) (HTTPResponseCode, error) {
if !filepath.IsAbs(path) {
return http.StatusBadRequest, xerrors.Errorf("file path must be absolute: %q", path)
}
@@ -94,3 +98,70 @@ func (a *agent) streamFile(ctx context.Context, rw http.ResponseWriter, path str
return 0, nil
}
func (a *agent) HandleWriteFile(rw http.ResponseWriter, r *http.Request) {
ctx := r.Context()
query := r.URL.Query()
parser := httpapi.NewQueryParamParser().RequiredNotEmpty("path")
path := parser.String(query, "", "path")
parser.ErrorExcessParams(query)
if len(parser.Errors) > 0 {
httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
Message: "Query parameters have invalid values.",
Validations: parser.Errors,
})
return
}
status, err := a.writeFile(ctx, r, path)
if err != nil {
httpapi.Write(ctx, rw, status, codersdk.Response{
Message: err.Error(),
})
return
}
httpapi.Write(ctx, rw, http.StatusOK, codersdk.Response{
Message: fmt.Sprintf("Successfully wrote to %q", path),
})
}
func (a *agent) writeFile(ctx context.Context, r *http.Request, path string) (HTTPResponseCode, error) {
if !filepath.IsAbs(path) {
return http.StatusBadRequest, xerrors.Errorf("file path must be absolute: %q", path)
}
dir := filepath.Dir(path)
err := a.filesystem.MkdirAll(dir, 0o755)
if err != nil {
status := http.StatusInternalServerError
switch {
case errors.Is(err, os.ErrPermission):
status = http.StatusForbidden
case errors.Is(err, syscall.ENOTDIR):
status = http.StatusBadRequest
}
return status, err
}
f, err := a.filesystem.Create(path)
if err != nil {
status := http.StatusInternalServerError
switch {
case errors.Is(err, os.ErrPermission):
status = http.StatusForbidden
case errors.Is(err, syscall.EISDIR):
status = http.StatusBadRequest
}
return status, err
}
defer f.Close()
_, err = io.Copy(f, r.Body)
if err != nil && !errors.Is(err, io.EOF) && ctx.Err() == nil {
a.logger.Error(ctx, "workspace agent write file", slog.Error(err))
}
return 0, nil
}
+163 -1
View File
@@ -1,11 +1,14 @@
package agent_test
import (
"bytes"
"context"
"io"
"net/http"
"os"
"path/filepath"
"runtime"
"syscall"
"testing"
"github.com/spf13/afero"
@@ -38,6 +41,56 @@ func (fs *testFs) Open(name string) (afero.File, error) {
return fs.Fs.Open(name)
}
func (fs *testFs) Create(name string) (afero.File, error) {
if err := fs.intercept("create", name); err != nil {
return nil, err
}
// Unlike os, afero lets you create files where directories already exist and
// lets you nest them underneath files, somehow.
stat, err := fs.Fs.Stat(name)
if err == nil && stat.IsDir() {
return nil, &os.PathError{
Op: "open",
Path: name,
Err: syscall.EISDIR,
}
}
stat, err = fs.Fs.Stat(filepath.Dir(name))
if err == nil && !stat.IsDir() {
return nil, &os.PathError{
Op: "open",
Path: name,
Err: syscall.ENOTDIR,
}
}
return fs.Fs.Create(name)
}
func (fs *testFs) MkdirAll(name string, mode os.FileMode) error {
if err := fs.intercept("mkdirall", name); err != nil {
return err
}
// Unlike os, afero lets you create directories where files already exist and
// lets you nest them underneath files somehow.
stat, err := fs.Fs.Stat(filepath.Dir(name))
if err == nil && !stat.IsDir() {
return &os.PathError{
Op: "mkdir",
Path: name,
Err: syscall.ENOTDIR,
}
}
stat, err = fs.Fs.Stat(name)
if err == nil && !stat.IsDir() {
return &os.PathError{
Op: "mkdir",
Path: name,
Err: syscall.ENOTDIR,
}
}
return fs.Fs.MkdirAll(name, mode)
}
func TestReadFile(t *testing.T) {
t.Parallel()
@@ -82,7 +135,7 @@ func TestReadFile(t *testing.T) {
error: "\"path\" is required",
},
{
name: "RelativePath",
name: "RelativePathDotSlash",
path: "./relative",
errCode: http.StatusBadRequest,
error: "file path must be absolute",
@@ -214,3 +267,112 @@ func TestReadFile(t *testing.T) {
})
}
}
func TestWriteFile(t *testing.T) {
t.Parallel()
tmpdir := os.TempDir()
noPermsFilePath := filepath.Join(tmpdir, "no-perms-file")
noPermsDirPath := filepath.Join(tmpdir, "no-perms-dir")
//nolint:dogsled
conn, _, _, fs, _ := setupAgent(t, agentsdk.Manifest{}, 0, func(_ *agenttest.Client, opts *agent.Options) {
opts.Filesystem = newTestFs(opts.Filesystem, func(call, file string) error {
if file == noPermsFilePath || file == noPermsDirPath {
return os.ErrPermission
}
return nil
})
})
dirPath := filepath.Join(tmpdir, "directory")
err := fs.MkdirAll(dirPath, 0o755)
require.NoError(t, err)
filePath := filepath.Join(tmpdir, "file")
err = afero.WriteFile(fs, filePath, []byte("content"), 0o644)
require.NoError(t, err)
notDirErr := "not a directory"
if runtime.GOOS == "windows" {
notDirErr = "cannot find the path"
}
tests := []struct {
name string
path string
bytes []byte
errCode int
error string
}{
{
name: "NoPath",
path: "",
errCode: http.StatusBadRequest,
error: "\"path\" is required",
},
{
name: "RelativePathDotSlash",
path: "./relative",
errCode: http.StatusBadRequest,
error: "file path must be absolute",
},
{
name: "RelativePath",
path: "also-relative",
errCode: http.StatusBadRequest,
error: "file path must be absolute",
},
{
name: "NonExistent",
path: filepath.Join(tmpdir, "/nested/does-not-exist"),
bytes: []byte("now it does exist"),
},
{
name: "IsDir",
path: dirPath,
errCode: http.StatusBadRequest,
error: "is a directory",
},
{
name: "IsNotDir",
path: filepath.Join(filePath, "file2"),
errCode: http.StatusBadRequest,
error: notDirErr,
},
{
name: "NoPermissionsFile",
path: noPermsFilePath,
errCode: http.StatusForbidden,
error: "permission denied",
},
{
name: "NoPermissionsDir",
path: filepath.Join(noPermsDirPath, "within-no-perm-dir"),
errCode: http.StatusForbidden,
error: "permission denied",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
t.Parallel()
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
defer cancel()
reader := bytes.NewReader(tt.bytes)
err := conn.WriteFile(ctx, tt.path, reader)
if tt.errCode != 0 {
require.Error(t, err)
cerr := coderdtest.SDKError(t, err)
require.Contains(t, cerr.Error(), tt.error)
require.Equal(t, tt.errCode, cerr.StatusCode())
} else {
require.NoError(t, err)
b, err := afero.ReadFile(fs, tt.path)
require.NoError(t, err)
require.Equal(t, tt.bytes, b)
}
})
}
}
+50
View File
@@ -43,6 +43,7 @@ const (
ToolNameChatGPTSearch = "search"
ToolNameChatGPTFetch = "fetch"
ToolNameWorkspaceReadFile = "coder_workspace_read_file"
ToolNameWorkspaceWriteFile = "coder_workspace_write_file"
)
func NewDeps(client *codersdk.Client, opts ...func(*Deps)) (Deps, error) {
@@ -211,6 +212,7 @@ var All = []GenericTool{
ChatGPTSearch.Generic(),
ChatGPTFetch.Generic(),
WorkspaceReadFile.Generic(),
WorkspaceWriteFile.Generic(),
}
type ReportTaskArgs struct {
@@ -1441,6 +1443,54 @@ var WorkspaceReadFile = Tool[WorkspaceReadFileArgs, WorkspaceReadFileResponse]{
},
}
type WorkspaceWriteFileArgs struct {
Workspace string `json:"workspace"`
Path string `json:"path"`
Content []byte `json:"content"`
}
var WorkspaceWriteFile = Tool[WorkspaceWriteFileArgs, codersdk.Response]{
Tool: aisdk.Tool{
Name: ToolNameWorkspaceWriteFile,
Description: `Write a file in a workspace.`,
Schema: aisdk.Schema{
Properties: map[string]any{
"workspace": map[string]any{
"type": "string",
"description": "The workspace name in the format [owner/]workspace[.agent]. If an owner is not specified, the authenticated user is used.",
},
"path": map[string]any{
"type": "string",
"description": "The absolute path of the file to write in the workspace.",
},
"content": map[string]any{
"type": "string",
"description": "The base64-encoded bytes to write to the file.",
},
},
Required: []string{"path", "workspace", "content"},
},
},
UserClientOptional: true,
Handler: func(ctx context.Context, deps Deps, args WorkspaceWriteFileArgs) (codersdk.Response, error) {
conn, err := newAgentConn(ctx, deps.coderClient, args.Workspace)
if err != nil {
return codersdk.Response{}, err
}
defer conn.Close()
reader := bytes.NewReader(args.Content)
err = conn.WriteFile(ctx, args.Path, reader)
if err != nil {
return codersdk.Response{}, err
}
return codersdk.Response{
Message: "File written successfully.",
}, nil
},
}
// NormalizeWorkspaceInput converts workspace name input to standard format.
// Handles the following input formats:
// - workspace → workspace
+24
View File
@@ -555,6 +555,30 @@ func TestTools(t *testing.T) {
})
}
})
t.Run("WorkspaceWriteFile", func(t *testing.T) {
t.Parallel()
client, workspace, agentToken := setupWorkspaceForAgent(t)
fs := afero.NewMemMapFs()
_ = agenttest.New(t, client.URL, agentToken, func(opts *agent.Options) {
opts.Filesystem = fs
})
coderdtest.NewWorkspaceAgentWaiter(t, client, workspace.ID).Wait()
tb, err := toolsdk.NewDeps(client)
require.NoError(t, err)
_, err = testTool(t, toolsdk.WorkspaceWriteFile, tb, toolsdk.WorkspaceWriteFileArgs{
Workspace: workspace.Name,
Path: "/test/some/path",
Content: []byte("content"),
})
require.NoError(t, err)
b, err := afero.ReadFile(fs, "/test/some/path")
require.NoError(t, err)
require.Equal(t, []byte("content"), b)
})
}
// TestedTools keeps track of which tools have been tested.
+22
View File
@@ -61,6 +61,7 @@ type AgentConn interface {
ReconnectingPTY(ctx context.Context, id uuid.UUID, height uint16, width uint16, command string, initOpts ...AgentReconnectingPTYInitOption) (net.Conn, error)
RecreateDevcontainer(ctx context.Context, devcontainerID string) (codersdk.Response, error)
ReadFile(ctx context.Context, path string, offset, limit int64) (io.ReadCloser, string, error)
WriteFile(ctx context.Context, path string, reader io.Reader) error
SSH(ctx context.Context) (*gonet.TCPConn, error)
SSHClient(ctx context.Context) (*ssh.Client, error)
SSHClientOnPort(ctx context.Context, port uint16) (*ssh.Client, error)
@@ -501,6 +502,27 @@ func (c *agentConn) ReadFile(ctx context.Context, path string, offset, limit int
return res.Body, mimeType, nil
}
// WriteFile writes to a file in the workspace.
func (c *agentConn) WriteFile(ctx context.Context, path string, reader io.Reader) error {
ctx, span := tracing.StartSpan(ctx)
defer span.End()
res, err := c.apiRequest(ctx, http.MethodPost, fmt.Sprintf("/api/v0/write-file?path=%s", path), reader)
if err != nil {
return xerrors.Errorf("do request: %w", err)
}
defer res.Body.Close()
if res.StatusCode != http.StatusOK {
return codersdk.ReadBodyAsError(res)
}
var m codersdk.Response
if err := json.NewDecoder(res.Body).Decode(&m); err != nil {
return xerrors.Errorf("decode response body: %w", err)
}
return nil
}
// apiRequest makes a request to the workspace agent's HTTP API server.
func (c *agentConn) apiRequest(ctx context.Context, method, path string, body io.Reader) (*http.Response, error) {
ctx, span := tracing.StartSpan(ctx)
@@ -387,3 +387,17 @@ func (mr *MockAgentConnMockRecorder) WatchContainers(ctx, logger any) *gomock.Ca
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "WatchContainers", reflect.TypeOf((*MockAgentConn)(nil).WatchContainers), ctx, logger)
}
// WriteFile mocks base method.
func (m *MockAgentConn) WriteFile(ctx context.Context, path string, reader io.Reader) error {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "WriteFile", ctx, path, reader)
ret0, _ := ret[0].(error)
return ret0
}
// WriteFile indicates an expected call of WriteFile.
func (mr *MockAgentConnMockRecorder) WriteFile(ctx, path, reader any) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "WriteFile", reflect.TypeOf((*MockAgentConn)(nil).WriteFile), ctx, path, reader)
}