Skip to content

Commit 1d4cd78

Browse files
code-ashercode-asher
committed
Move file read buffering to toolsdk
1 parent a042936 commit 1d4cd78

File tree

4 files changed

+75
-66
lines changed

4 files changed

+75
-66
lines changed

‎agent/files_test.go

Lines changed: 48 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ package agent_test
22

33
import (
44
"context"
5+
"io"
56
"net/http"
67
"os"
78
"path/filepath"
@@ -125,53 +126,61 @@ func TestReadFile(t *testing.T) {
125126
error: "permission denied",
126127
},
127128
{
128-
name: "Defaults",
129-
path: filePath,
130-
bytes: []byte("content"),
129+
name: "Defaults",
130+
path: filePath,
131+
bytes: []byte("content"),
132+
mimeType: "application/octet-stream",
131133
},
132134
{
133-
name: "Limit1",
134-
path: filePath,
135-
limit: 1,
136-
bytes: []byte("c"),
135+
name: "Limit1",
136+
path: filePath,
137+
limit: 1,
138+
bytes: []byte("c"),
139+
mimeType: "application/octet-stream",
137140
},
138141
{
139-
name: "Offset1",
140-
path: filePath,
141-
offset: 1,
142-
bytes: []byte("ontent"),
142+
name: "Offset1",
143+
path: filePath,
144+
offset: 1,
145+
bytes: []byte("ontent"),
146+
mimeType: "application/octet-stream",
143147
},
144148
{
145-
name: "Limit1Offset2",
146-
path: filePath,
147-
limit: 1,
148-
offset: 2,
149-
bytes: []byte("n"),
149+
name: "Limit1Offset2",
150+
path: filePath,
151+
limit: 1,
152+
offset: 2,
153+
bytes: []byte("n"),
154+
mimeType: "application/octet-stream",
150155
},
151156
{
152-
name: "Limit7Offset0",
153-
path: filePath,
154-
limit: 7,
155-
offset: 0,
156-
bytes: []byte("content"),
157+
name: "Limit7Offset0",
158+
path: filePath,
159+
limit: 7,
160+
offset: 0,
161+
bytes: []byte("content"),
162+
mimeType: "application/octet-stream",
157163
},
158164
{
159-
name: "Limit100",
160-
path: filePath,
161-
limit: 100,
162-
bytes: []byte("content"),
165+
name: "Limit100",
166+
path: filePath,
167+
limit: 100,
168+
bytes: []byte("content"),
169+
mimeType: "application/octet-stream",
163170
},
164171
{
165-
name: "Offset7",
166-
path: filePath,
167-
offset: 7,
168-
bytes: []byte{},
172+
name: "Offset7",
173+
path: filePath,
174+
offset: 7,
175+
bytes: []byte{},
176+
mimeType: "application/octet-stream",
169177
},
170178
{
171-
name: "Offset100",
172-
path: filePath,
173-
offset: 100,
174-
bytes: []byte{},
179+
name: "Offset100",
180+
path: filePath,
181+
offset: 100,
182+
bytes: []byte{},
183+
mimeType: "application/octet-stream",
175184
},
176185
{
177186
name: "MimeTypePng",
@@ -188,20 +197,19 @@ func TestReadFile(t *testing.T) {
188197
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
189198
defer cancel()
190199

191-
b, mimeType, err := conn.ReadFile(ctx, tt.path, tt.offset, tt.limit)
200+
reader, mimeType, err := conn.ReadFile(ctx, tt.path, tt.offset, tt.limit)
192201
if tt.errCode != 0 {
193202
require.Error(t, err)
194203
cerr := coderdtest.SDKError(t, err)
195204
require.Contains(t, cerr.Error(), tt.error)
196205
require.Equal(t, tt.errCode, cerr.StatusCode())
197206
} else {
198207
require.NoError(t, err)
199-
require.Equal(t, tt.bytes, b)
200-
expectedMimeType := tt.mimeType
201-
if expectedMimeType == "" {
202-
expectedMimeType = "application/octet-stream"
203-
}
204-
require.Equal(t, expectedMimeType, mimeType)
208+
defer reader.Close()
209+
bytes, err := io.ReadAll(reader)
210+
require.NoError(t, err)
211+
require.Equal(t, tt.bytes, bytes)
212+
require.Equal(t, tt.mimeType, mimeType)
205213
}
206214
})
207215
}

‎codersdk/toolsdk/toolsdk.go

Lines changed: 20 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1380,6 +1380,8 @@ type WorkspaceReadFileResponse struct {
13801380
MimeType string `json:"mimeType"`
13811381
}
13821382

1383+
const maxFileLimit = 1 << 20 // 1MiB
1384+
13831385
var WorkspaceReadFile = Tool[WorkspaceReadFileArgs, WorkspaceReadFileResponse]{
13841386
Tool: aisdk.Tool{
13851387
Name: ToolNameWorkspaceReadFile,
@@ -1414,12 +1416,28 @@ var WorkspaceReadFile = Tool[WorkspaceReadFileArgs, WorkspaceReadFileResponse]{
14141416
}
14151417
defer conn.Close()
14161418

1417-
bytes, mimeType, err := conn.ReadFile(ctx, args.Path, args.Offset, args.Limit)
1419+
// Ideally we could stream this all the way back, but it looks like the MCP
1420+
// interfaces only allow returning full responses which means the whole
1421+
// thing has to be read into memory. So, add a maximum to compensate.
1422+
limit := args.Limit
1423+
if limit == 0 {
1424+
limit = maxFileLimit
1425+
} else if limit > maxFileLimit {
1426+
return WorkspaceReadFileResponse{}, xerrors.Errorf("limit must be %d or less, got %d", maxFileLimit, limit)
1427+
}
1428+
1429+
reader, mimeType, err := conn.ReadFile(ctx, args.Path, args.Offset, args.Limit)
14181430
if err != nil {
14191431
return WorkspaceReadFileResponse{}, err
14201432
}
1433+
defer reader.Close()
1434+
1435+
bs, err := io.ReadAll(reader)
1436+
if err != nil {
1437+
return WorkspaceReadFileResponse{}, xerrors.Errorf("read response body: %w", err)
1438+
}
14211439

1422-
return WorkspaceReadFileResponse{Content: bytes, MimeType: mimeType}, nil
1440+
return WorkspaceReadFileResponse{Content: bs, MimeType: mimeType}, nil
14231441
},
14241442
}
14251443

‎codersdk/workspacesdk/agentconn.go

Lines changed: 5 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -60,7 +60,7 @@ type AgentConn interface {
6060
PrometheusMetrics(ctx context.Context) ([]byte, error)
6161
ReconnectingPTY(ctx context.Context, id uuid.UUID, height uint16, width uint16, command string, initOpts ...AgentReconnectingPTYInitOption) (net.Conn, error)
6262
RecreateDevcontainer(ctx context.Context, devcontainerID string) (codersdk.Response, error)
63-
ReadFile(ctx context.Context, path string, offset, limit int64) ([]byte, string, error)
63+
ReadFile(ctx context.Context, path string, offset, limit int64) (io.ReadCloser, string, error)
6464
SSH(ctx context.Context) (*gonet.TCPConn, error)
6565
SSHClient(ctx context.Context) (*ssh.Client, error)
6666
SSHClientOnPort(ctx context.Context, port uint16) (*ssh.Client, error)
@@ -477,43 +477,26 @@ func (c *agentConn) RecreateDevcontainer(ctx context.Context, devcontainerID str
477477
return m, nil
478478
}
479479

480-
const maxFileLimit = 1 << 20 // 1MiB
481-
482-
// ReadFile reads from a file from the workspace, returning the file's
483-
// (potentially partial) bytes and the mime type.
484-
func (c *agentConn) ReadFile(ctx context.Context, path string, offset, limit int64) ([]byte, string, error) {
480+
// ReadFile reads from a file from the workspace, returning a file reader and
481+
// the mime type.
482+
func (c *agentConn) ReadFile(ctx context.Context, path string, offset, limit int64) (io.ReadCloser, string, error) {
485483
ctx, span := tracing.StartSpan(ctx)
486484
defer span.End()
487485

488-
// Ideally we could stream this all the way back, but it looks like the MCP
489-
// interfaces only allow returning full responses which means the whole thing
490-
// has to be read into memory. So, add a maximum to compensate.
491-
if limit == 0 {
492-
limit = maxFileLimit
493-
} else if limit > maxFileLimit {
494-
return nil, "", xerrors.Errorf("limit must be %d or less, got %d", maxFileLimit, limit)
495-
}
496-
497486
res, err := c.apiRequest(ctx, http.MethodGet, fmt.Sprintf("/api/v0/read-file?path=%s&offset=%d&limit=%d", path, offset, limit), nil)
498487
if err != nil {
499488
return nil, "", xerrors.Errorf("do request: %w", err)
500489
}
501-
defer res.Body.Close()
502490
if res.StatusCode != http.StatusOK {
503491
return nil, "", codersdk.ReadBodyAsError(res)
504492
}
505493

506-
bs, err := io.ReadAll(res.Body)
507-
if err != nil {
508-
return nil, "", xerrors.Errorf("read response body: %w", err)
509-
}
510-
511494
mimeType := res.Header.Get("Content-Type")
512495
if mimeType == "" {
513496
mimeType = "application/octet-stream"
514497
}
515498

516-
return bs, mimeType, nil
499+
return res.Body, mimeType, nil
517500
}
518501

519502
// apiRequest makes a request to the workspace agent's HTTP API server.

‎codersdk/workspacesdk/agentconnmock/agentconnmock.go

Lines changed: 2 additions & 2 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

0 commit comments

Comments
 (0)