feat(sdkserver): stream workspace file writes #13

Merged

thklein merged 5 commits from feat/streaming-workspace-write into main 2026-10-02 14:00:58 +00:00

Conversation 1 Commits 5 Files changed 6 +564 -12

thklein commented 2026-10-02 12:38:19 +00:00

Owner

Copy link

Part of basepeak-ai/bpai#1220 — workspace file writes materialise the file many times over; a single large file OOM-killed a tenant.

What

Engine: engine.WithStreamedBody(ctx, toolID, contentType, r) gives one HTTP tool, by ID, a reader to send as its request body in place of the body built from its input. Nothing else about the call changes:

• daemon start and supervision
• X-GPTScript-Daemon-Token
• the forwarded X-GPTScript-Env values and X-GPTScript-Tool-Name
• the user-cancel check
• error formatting

The body can be sent only once; a second call of the same tool fails rather than sending a drained stream. Sent() reports whether the call was made. This uses the same context-value pattern as WithToolCategory.

sdkserver: new route POST /workspaces/write-file-stream. The body is a JSON header (the write-file fields minus contents) followed immediately by the raw file bytes. The route runs the workspace tool as usual, with the content base64-encoded on the fly as the streamed body. The workspace daemon already decodes its body as a stream. The route fails if the run returned without sending the body, so a write that never happened can't report success. A workspace tool that is not an HTTP tool falls back to the buffered run.

Both write routes now build the tool input with json.Marshal instead of Sprintf, so quotes and backslashes in a file path are escaped correctly.

History of this branch

The first commit resolved the daemon URL with a blocking run and made the request itself. That skipped the token and the env headers that runHTTP sends; today's provider ignores them, but a future one might not. The second commit moves the streaming into runHTTP instead. The third fixes a regression from the first commit: omitempty on body made empty-file writes fail on the existing route (illegal base64 data at input byte 0), because an HTTP tool with no body argument sends its whole JSON input as the body.

Validation

• Unit tests, all passing under -race:
  • pkg/engine/http_stream_test.go: the stream replaces the body, keeps the content type, tool name and env headers, applies only to its own tool, and is sent at most once.
  • pkg/sdkserver/workspaces_test.go: the header/content split, the escaped tool input (including body present when empty), and the base64 reader: round-trip, read-error propagation, early close.
• Mutation check: breaking each engine guard (tool-ID match, single-use) fails exactly the test written for it.
• Parity with the existing route: a recording stand-in for the workspace daemon logged what each route sends for a small text file, an empty file, and 3 MB of binary with createRevision=false and latestRevision=3. Path, query, header set, daemon-token match and decoded-content SHA-256 are identical. The only difference is chunked transfer encoding on the stream.
• go-gptscript integration (basepeak-ai/go-gptscript#2) against a binary of this branch, with the real workspace-provider: all 13 directory-workspace tests pass (S3 skipped, no credentials). That includes the stream round-trip, the conflict mapping and the new empty-file test on both routes. Against the first commit of this branch, the empty-file test fails on the buffered route, which is the regression above.
• make test (./pkg/...) and ./integration/... pass. golangci-lint reports 32 issues, the same 32 as on main; none are in files this PR touches.

Measured

Peak RSS of this process (as the sdkserver) while writing:

route	1 GiB	4 GiB
write-file (buffered)	12 336 MiB	—
write-file-stream	36 MiB	36 MiB

Companion PRs: basepeak-ai/go-gptscript#2 (client), basepeak-ai/bpai-tools#360 (smb connector), basepeak-ai/bpai#1234 (pin bump).

Part of basepeak-ai/bpai#1220 — workspace file writes materialise the file many times over; a single large file OOM-killed a tenant. ## What **Engine:** `engine.WithStreamedBody(ctx, toolID, contentType, r)` gives one HTTP tool, by ID, a reader to send as its request body in place of the body built from its input. Nothing else about the call changes: - daemon start and supervision - `X-GPTScript-Daemon-Token` - the forwarded `X-GPTScript-Env` values and `X-GPTScript-Tool-Name` - the user-cancel check - error formatting The body can be sent only once; a second call of the same tool fails rather than sending a drained stream. `Sent()` reports whether the call was made. This uses the same context-value pattern as `WithToolCategory`. **sdkserver:** new route `POST /workspaces/write-file-stream`. The body is a JSON header (the `write-file` fields minus `contents`) followed immediately by the raw file bytes. The route runs the workspace tool as usual, with the content base64-encoded on the fly as the streamed body. The workspace daemon already decodes its body as a stream. The route fails if the run returned without sending the body, so a write that never happened can't report success. A workspace tool that is not an HTTP tool falls back to the buffered run. Both write routes now build the tool input with `json.Marshal` instead of `Sprintf`, so quotes and backslashes in a file path are escaped correctly. ### History of this branch The first commit resolved the daemon URL with a blocking run and made the request itself. That skipped the token and the env headers that `runHTTP` sends; today's provider ignores them, but a future one might not. The second commit moves the streaming into `runHTTP` instead. The third fixes a regression from the first commit: `omitempty` on `body` made **empty-file writes fail on the existing route** (`illegal base64 data at input byte 0`), because an HTTP tool with no body argument sends its whole JSON input as the body. ## Validation - **Unit tests**, all passing under `-race`: - `pkg/engine/http_stream_test.go`: the stream replaces the body, keeps the content type, tool name and env headers, applies only to its own tool, and is sent at most once. - `pkg/sdkserver/workspaces_test.go`: the header/content split, the escaped tool input (including `body` present when empty), and the base64 reader: round-trip, read-error propagation, early close. - **Mutation check:** breaking each engine guard (tool-ID match, single-use) fails exactly the test written for it. - **Parity with the existing route:** a recording stand-in for the workspace daemon logged what each route sends for a small text file, an empty file, and 3 MB of binary with `createRevision=false` and `latestRevision=3`. Path, query, header set, daemon-token match and decoded-content SHA-256 are identical. The only difference is chunked transfer encoding on the stream. - **go-gptscript integration** (basepeak-ai/go-gptscript#2) against a binary of this branch, with the real workspace-provider: all 13 directory-workspace tests pass (S3 skipped, no credentials). That includes the stream round-trip, the conflict mapping and the new empty-file test on both routes. Against the first commit of this branch, the empty-file test fails on the buffered route, which is the regression above. - `make test` (`./pkg/...`) and `./integration/...` pass. `golangci-lint` reports 32 issues, the same 32 as on main; none are in files this PR touches. ## Measured Peak RSS of this process (as the sdkserver) while writing: | route | 1 GiB | 4 GiB | |---|---|---| | `write-file` (buffered) | 12 336 MiB | — | | `write-file-stream` | 36 MiB | 36 MiB | Companion PRs: basepeak-ai/go-gptscript#2 (client), basepeak-ai/bpai-tools#360 (smb connector), basepeak-ai/bpai#1234 (pin bump).

thklein added 1 commit 2026-10-02 12:38:19 +00:00

feat(sdkserver): stream workspace file writes ... 77302e62e5

POST /workspaces/write-file takes the file as a base64 string inside a
JSON body, then Sprintf's it into a second JSON document as a tool
argument, which the engine copies again into an environment map and a
request body. Writing a 1 GiB file peaks at ~11 GiB in this process
alone.

The new POST /workspaces/write-file-stream reads a JSON header off the
front of the body and treats the rest as raw content. It runs the
workspace tool's write blocking, which starts its daemon if needed and
returns the URL it would have called, then base64-encodes the content
on the fly into a request to that URL. The daemon already decodes its
body as a stream. Peak stays at ~35 MiB for 1 GiB.

A workspace tool that is not an HTTP tool falls back to the buffered
call. The tool input for both routes is now built with json.Marshal
instead of Sprintf, so a quote in a file path no longer breaks it.

Refs basepeak-ai/bpai#1220

thklein referenced this pull request from basepeak-ai/go-gptscript 2026-10-02 12:38:32 +00:00

feat: add WriteFileInWorkspaceStream #2

thklein added 2 commits 2026-10-02 13:23:31 +00:00

refactor(engine): stream an HTTP tool's request body ... 2604e1bc52

The streaming write route ran the workspace tool blocking to learn its
daemon's URL, then made the request itself. That request skipped what
every HTTP tool call sends: the daemon token, the forwarded
X-GPTScript-Env values and X-GPTScript-Tool-Name. Today's workspace
provider ignores them, but a provider relying on them would have
worked on one route and failed on the other.

engine.WithStreamedBody now lets a caller give one HTTP tool, by ID, a
reader to send as its request body in place of the body built from
its input. The rest of the call is unchanged. The body can be sent
only once, and Sent() reports whether it was, so the route can reject
a run that returned without making the call.

The route now runs the tool normally with that context, and the
separate HTTP client is gone.

Refs basepeak-ai/bpai#1220

fix(sdkserver): keep the body argument for empty files ... 5a0f88e8ec

Building the tool input with json.Marshal and `omitempty` dropped the
body argument whenever the file was empty. An HTTP tool with no body
argument sends its whole input as the request body, so the daemon got
JSON instead of base64 and failed with "illegal base64 data at input
byte 0". The Sprintf this replaced always sent `"body": ""`.

Refs basepeak-ai/bpai#1220

thklein added 2 commits 2026-10-02 14:00:18 +00:00

fix(sdkserver): encode the streamed body inside Read ... 572b96b134

The base64 encoder ran in a goroutine feeding an io.Pipe. When the
daemon answered early (a conflict) or the run failed, the handler
returned while that goroutine could still be blocked reading the
client's request body, which net/http does not allow after ServeHTTP
returns, and it lived as long as a stalled client kept the connection.

The encoder now works a chunk at a time inside Read, so the body is
only ever read by the request that owns it. Adds end-to-end tests of
the route through the real loader, runner and engine: a 2 MiB binary
write, a conflict with bodies up to 32 MiB that the daemon refuses
before reading (checked against go-gptscript's conflict regex), and a
client that goes away mid-upload.

Refs basepeak-ai/bpai#1220

fix(sdkserver): fail a streamed write the tool did not read in full ... e065f3a8c9

Sent() only says the tool's server answered. A tool that answers
success before reading all of its request body would have had a cut
file reported as written. The route now also checks that the content
was read to its end. workspace-provider only answers after reading to
EOF, so this guards other workspace tools and future changes.

Refs basepeak-ai/bpai#1220

thklein commented 2026-10-02 14:00:39 +00:00

Author

Owner

Copy link

Review (fresh-context reviewers) — outcome

Fixed

• The encoder goroutine could read the request body after the handler returned when the daemon answered early. The base64 encoding now happens inside Read, so there is no goroutine (572b96b).
• Sent() only said the daemon had answered. The route now also requires the content to have been read to EOF, so a tool that answers success early can't have a cut file reported as written (e065f3a). There is a test with a fake daemon that ignores the body, and a mutation check confirms it.
• No end-to-end test of the route: added, through the real loader, runner and engine. It covers a 2 MiB binary write, a conflict refused before reading with bodies up to 32 MiB (checked against go-gptscript's conflict regex), a client aborting mid-upload (nothing is committed), and the early-answer case above.

Dismissed

• Also stream when the workspace tool is a #!sys.daemon tool: no such provider exists, and the buffered fallback is the safer default for an unknown tool shape.

Follow-up, not in this PR (workspace-provider)

• directoryProvider.WriteFile (also s3 and azure) writes the revision bookkeeping before it reads the body. A write that fails mid-upload therefore leaves CurrentID bumped, and a retry with the old latestRevisionID gets a spurious conflict. This only happens with createRevision on; the knowledge connectors write with it off.
• The s3 and azure providers io.ReadAll the body, so on those backends the daemon still holds one copy of the file.

Re-validated on the final commit

• engine and sdkserver tests: pass three times under -race.
• go-gptscript integration matrix against real binaries: new client × new server and old client × new server pass. New client × old server passes everything except the tests of the stream call, which can't pass without the route.
• 1 GiB write: ~34 MiB peak in the server.

CI note: this repo's CI is red on main too (private module fetch fails at make tidy), so these local runs are the evidence.

## Review (fresh-context reviewers) — outcome **Fixed** - **The encoder goroutine could read the request body after the handler returned** when the daemon answered early. The base64 encoding now happens inside `Read`, so there is no goroutine (`572b96b`). - **`Sent()` only said the daemon had answered.** The route now also requires the content to have been read to EOF, so a tool that answers success early can't have a cut file reported as written (`e065f3a`). There is a test with a fake daemon that ignores the body, and a mutation check confirms it. - **No end-to-end test of the route:** added, through the real loader, runner and engine. It covers a 2 MiB binary write, a conflict refused before reading with bodies up to 32 MiB (checked against go-gptscript's conflict regex), a client aborting mid-upload (nothing is committed), and the early-answer case above. **Dismissed** - *Also stream when the workspace tool is a `#!sys.daemon` tool:* no such provider exists, and the buffered fallback is the safer default for an unknown tool shape. **Follow-up, not in this PR (workspace-provider)** - `directoryProvider.WriteFile` (also s3 and azure) writes the revision bookkeeping before it reads the body. A write that fails mid-upload therefore leaves `CurrentID` bumped, and a retry with the old `latestRevisionID` gets a spurious conflict. This only happens with `createRevision` on; the knowledge connectors write with it off. - The s3 and azure providers `io.ReadAll` the body, so on those backends the daemon still holds one copy of the file. **Re-validated on the final commit** - engine and sdkserver tests: pass three times under `-race`. - go-gptscript integration matrix against real binaries: new client × new server and old client × new server pass. New client × old server passes everything except the tests of the stream call, which can't pass without the route. - 1 GiB write: ~34 MiB peak in the server. CI note: this repo's CI is red on `main` too (private module fetch fails at `make tidy`), so these local runs are the evidence.

thklein merged commit 5530f433ff into main 2026-10-02 14:00:58 +00:00

thklein deleted branch feat/streaming-workspace-write 2026-10-02 14:00:58 +00:00

thklein referenced this pull request from a commit 2026-10-02 14:01:00 +00:00

feat(sdkserver): stream workspace file writes (#13)

Sign in to join this conversation.

Reviewers

No reviewers

Labels

Clear labels   

area/engine

  

clarity/clear

  

enhancement

  

security

  

size/M

No labels

area/engine

clarity/clear

enhancement

security

size/M

Milestone

Clear milestone

No items

No milestone

Projects

Clear projects

No items

No project

Assignees

Clear assignees

No assignees

1 participant

Notifications

Subscribe

Due date

The due date is invalid or out of range. Please use the format "yyyy-mm-dd".

No due date set.

Dependencies

No dependencies set

Reference

basepeak-ai/gptscript!13

No description provided.