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

Commitf903423

Browse files
committed
Move file read buffering to toolsdk
1 parenta042936 commitf903423

File tree

5 files changed

+97
-68
lines changed

5 files changed

+97
-68
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
defercancel()
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
iftt.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-
ifexpectedMimeType=="" {
202-
expectedMimeType="application/octet-stream"
203-
}
204-
require.Equal(t,expectedMimeType,mimeType)
208+
deferreader.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
MimeTypestring`json:"mimeType"`
13811381
}
13821382

1383+
constmaxFileLimit=1<<20// 1MiB
1384+
13831385
varWorkspaceReadFile=Tool[WorkspaceReadFileArgs,WorkspaceReadFileResponse]{
13841386
Tool: aisdk.Tool{
13851387
Name:ToolNameWorkspaceReadFile,
@@ -1414,12 +1416,28 @@ var WorkspaceReadFile = Tool[WorkspaceReadFileArgs, WorkspaceReadFileResponse]{
14141416
}
14151417
deferconn.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 limit to compensate.
1422+
limit:=args.Limit
1423+
iflimit==0 {
1424+
limit=maxFileLimit
1425+
}elseiflimit>maxFileLimit {
1426+
returnWorkspaceReadFileResponse{},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,limit)
14181430
iferr!=nil {
14191431
returnWorkspaceReadFileResponse{},err
14201432
}
1433+
deferreader.Close()
1434+
1435+
bs,err:=io.ReadAll(reader)
1436+
iferr!=nil {
1437+
returnWorkspaceReadFileResponse{},xerrors.Errorf("read response body: %w",err)
1438+
}
14211439

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

‎codersdk/toolsdk/toolsdk_test.go‎

Lines changed: 20 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -471,6 +471,12 @@ func TestTools(t *testing.T) {
471471
err=afero.WriteFile(fs,filePath, []byte("content"),0o644)
472472
require.NoError(t,err)
473473

474+
largeFilePath:=filepath.Join(tmpdir,"large")
475+
largeFile,err:=fs.Create(largeFilePath)
476+
require.NoError(t,err)
477+
err=largeFile.Truncate(1<<21)
478+
require.NoError(t,err)
479+
474480
imagePath:=filepath.Join(tmpdir,"file.png")
475481
err=afero.WriteFile(fs,imagePath, []byte("not really an image"),0o644)
476482
require.NoError(t,err)
@@ -482,6 +488,7 @@ func TestTools(t *testing.T) {
482488
offsetint64
483489
mimeTypestring
484490
bytes []byte
491+
lengthint
485492
errorstring
486493
}{
487494
{
@@ -504,7 +511,13 @@ func TestTools(t *testing.T) {
504511
mimeType:"application/octet-stream",
505512
},
506513
{
507-
name:"MaxLimit",
514+
name:"DefaultMaxLimit",
515+
path:largeFilePath,
516+
length:1<<20,
517+
mimeType:"application/octet-stream",
518+
},
519+
{
520+
name:"ExceedMaxLimit",
508521
path:filePath,
509522
limit:1<<21,
510523
error:"limit must be 1048576 or less, got 2097152",
@@ -532,7 +545,12 @@ func TestTools(t *testing.T) {
532545
require.Contains(t,err.Error(),tt.error)
533546
}else {
534547
require.NoError(t,err)
535-
require.Equal(t,tt.bytes,resp.Content)
548+
iftt.length!=0 {
549+
require.Len(t,resp.Content,tt.length)
550+
}
551+
iftt.bytes!=nil {
552+
require.Equal(t,tt.bytes,resp.Content)
553+
}
536554
require.Equal(t,tt.mimeType,resp.MimeType)
537555
}
538556
})

‎codersdk/workspacesdk/agentconn.go‎

Lines changed: 7 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,heightuint16,widthuint16,commandstring,initOpts...AgentReconnectingPTYInitOption) (net.Conn,error)
6262
RecreateDevcontainer(ctx context.Context,devcontainerIDstring) (codersdk.Response,error)
63-
ReadFile(ctx context.Context,pathstring,offset,limitint64) ([]byte,string,error)
63+
ReadFile(ctx context.Context,pathstring,offset,limitint64) (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,portuint16) (*ssh.Client,error)
@@ -477,43 +477,28 @@ func (c *agentConn) RecreateDevcontainer(ctx context.Context, devcontainerID str
477477
returnm,nil
478478
}
479479

480-
constmaxFileLimit=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,pathstring,offset,limitint64) ([]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,pathstring,offset,limitint64) (io.ReadCloser,string,error) {
485483
ctx,span:=tracing.StartSpan(ctx)
486484
deferspan.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-
iflimit==0 {
492-
limit=maxFileLimit
493-
}elseiflimit>maxFileLimit {
494-
returnnil,"",xerrors.Errorf("limit must be %d or less, got %d",maxFileLimit,limit)
495-
}
496-
486+
//nolint:bodyclose // we want to return the body so the caller can stream.
497487
res,err:=c.apiRequest(ctx,http.MethodGet,fmt.Sprintf("/api/v0/read-file?path=%s&offset=%d&limit=%d",path,offset,limit),nil)
498488
iferr!=nil {
499489
returnnil,"",xerrors.Errorf("do request: %w",err)
500490
}
501-
deferres.Body.Close()
502491
ifres.StatusCode!=http.StatusOK {
492+
// codersdk.ReadBodyAsError will close the body.
503493
returnnil,"",codersdk.ReadBodyAsError(res)
504494
}
505495

506-
bs,err:=io.ReadAll(res.Body)
507-
iferr!=nil {
508-
returnnil,"",xerrors.Errorf("read response body: %w",err)
509-
}
510-
511496
mimeType:=res.Header.Get("Content-Type")
512497
ifmimeType=="" {
513498
mimeType="application/octet-stream"
514499
}
515500

516-
returnbs,mimeType,nil
501+
returnres.Body,mimeType,nil
517502
}
518503

519504
// 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 aboutcustomizing how changed files appear on GitHub.

0 commit comments

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