diff --git a/assistant/src/__tests__/host-transfer-proxy.test.ts b/assistant/src/__tests__/host-transfer-proxy.test.ts index 28188e3ce66..ecf54902809 100644 --- a/assistant/src/__tests__/host-transfer-proxy.test.ts +++ b/assistant/src/__tests__/host-transfer-proxy.test.ts @@ -341,6 +341,110 @@ describe("HostTransferProxy", () => { }); }); + describe("takeJustConsumedTransferMetadata", () => { + test("returns size+sha256 immediately after getTransferContent", async () => { + setup(); + tempDir = await mkdtemp(join(tmpdir(), "htp-test-")); + const srcPath = join(tempDir, "source.txt"); + const fileContent = "header metadata test"; + await globalThis.Bun.write(srcPath, fileContent); + + const resultPromise = proxy.requestToHost({ + sourcePath: srcPath, + destPath: "/host/dest.txt", + overwrite: true, + conversationId: "conv-meta-1", + }); + + await waitForMessages(sentMessages, 1); + const sent = sentMessages[0] as Record; + const transferId = sent.transferId as string; + + const expectedSize = Buffer.from(fileContent).length; + const expectedSha = createHash("sha256") + .update(Buffer.from(fileContent)) + .digest("hex"); + + // Handler consumes content; metadata should now be available for the + // GET-content route's responseHeaders resolver. + const content = proxy.getTransferContent(transferId); + expect(content).not.toBeNull(); + + const meta = proxy.takeJustConsumedTransferMetadata(transferId); + expect(meta).toEqual({ sizeBytes: expectedSize, sha256: expectedSha }); + + // Resolve the transfer to avoid hanging + const requestId = sent.requestId as string; + proxy.resolveTransferResult(requestId, { + isError: false, + bytesWritten: expectedSize, + }); + await resultPromise; + }); + + test("single-use: second take returns null", async () => { + setup(); + tempDir = await mkdtemp(join(tmpdir(), "htp-test-")); + const srcPath = join(tempDir, "source.txt"); + await globalThis.Bun.write(srcPath, "x"); + + const resultPromise = proxy.requestToHost({ + sourcePath: srcPath, + destPath: "/host/dest.txt", + overwrite: true, + conversationId: "conv-meta-2", + }); + + await waitForMessages(sentMessages, 1); + const sent = sentMessages[0] as Record; + const transferId = sent.transferId as string; + + proxy.getTransferContent(transferId); + expect(proxy.takeJustConsumedTransferMetadata(transferId)).not.toBeNull(); + expect(proxy.takeJustConsumedTransferMetadata(transferId)).toBeNull(); + + const requestId = sent.requestId as string; + proxy.resolveTransferResult(requestId, { isError: false }); + await resultPromise; + }); + + test("returns null when getTransferContent was never called", () => { + setup(); + expect( + proxy.takeJustConsumedTransferMetadata("never-consumed-id"), + ).toBeNull(); + }); + + test("returns null for unknown transfer ID even after a different transfer was consumed", async () => { + setup(); + tempDir = await mkdtemp(join(tmpdir(), "htp-test-")); + const srcPath = join(tempDir, "source.txt"); + await globalThis.Bun.write(srcPath, "x"); + + const resultPromise = proxy.requestToHost({ + sourcePath: srcPath, + destPath: "/host/dest.txt", + overwrite: true, + conversationId: "conv-meta-3", + }); + + await waitForMessages(sentMessages, 1); + const sent = sentMessages[0] as Record; + const transferId = sent.transferId as string; + + proxy.getTransferContent(transferId); + + // Different transferId should still return null + expect( + proxy.takeJustConsumedTransferMetadata("other-id"), + ).toBeNull(); + + const requestId = sent.requestId as string; + proxy.resolveTransferResult(requestId, { isError: false }); + await resultPromise; + }); + }); + describe("timeout behavior", () => { /** Use a very short real timeout instead of fake timers to avoid deadlocks in Bun. */ const SHORT_TIMEOUT_MS = 150; diff --git a/assistant/src/daemon/host-transfer-proxy.ts b/assistant/src/daemon/host-transfer-proxy.ts index 373f1b624f0..2ddbeade46d 100644 --- a/assistant/src/daemon/host-transfer-proxy.ts +++ b/assistant/src/daemon/host-transfer-proxy.ts @@ -83,6 +83,20 @@ export class HostTransferProxy { /** Pending transfers keyed by transferId (for content endpoint lookups). */ private transfers = new Map(); + /** + * Briefly retains size/sha256 of a just-consumed transfer so the GET-content + * route's `resolveResponseHeaders` callback (which the HTTP adapter invokes + * AFTER the request handler) can still set `Content-Length` and + * `X-Transfer-SHA256` headers. Without this, the handler's `getTransferContent` + * call deletes the entry before the header resolver runs, and the resolver + * silently falls back to default headers — meaning the documented response + * headers were never actually sent. Entries here self-clear on read; a 30s + * fallback timer prevents long-term retention if the resolver never runs. + */ + private justConsumedMetadata = new Map< + string, + { sizeBytes: number; sha256: string } + >(); /** * Whether a client with `host_file` capability is connected. @@ -451,6 +465,21 @@ export class HostTransferProxy { ) { return null; } + // Stash size/sha256 so the GET-content route's `resolveResponseHeaders` + // callback can still set `Content-Length` and `X-Transfer-SHA256` on the + // response. The HTTP adapter invokes the handler (this method) BEFORE the + // response-header resolver, so without this stash the resolver sees a + // deleted entry and silently falls back to default headers. + this.justConsumedMetadata.set(transferId, { + sizeBytes: entry.sizeBytes, + sha256: entry.sha256, + }); + // Fallback cleanup: if the resolver never reads (e.g., handler error after + // consume, request abort), drop the metadata after a short grace window. + // `unref()` so the timer never holds the process open. + setTimeout(() => { + this.justConsumedMetadata.delete(transferId); + }, 30_000).unref?.(); this.transfers.delete(transferId); return { buffer: entry.fileBuffer, @@ -459,6 +488,23 @@ export class HostTransferProxy { }; } + /** + * Returns and clears the size/sha256 metadata for a transfer that was just + * consumed by `getTransferContent`. Intended for use by the GET-content + * route's `resolveResponseHeaders` callback to populate `Content-Length` and + * `X-Transfer-SHA256` response headers. Returns null if no metadata is + * cached (e.g., transfer was never consumed, or already read by a previous + * resolver call). + */ + takeJustConsumedTransferMetadata( + transferId: string, + ): { sizeBytes: number; sha256: string } | null { + const meta = this.justConsumedMetadata.get(transferId); + if (!meta) return null; + this.justConsumedMetadata.delete(transferId); + return meta; + } + /** * Receive file content from the client for a to_sandbox transfer (the PUT content endpoint). * diff --git a/assistant/src/runtime/routes/host-transfer-routes.ts b/assistant/src/runtime/routes/host-transfer-routes.ts index dbbf22d7556..90d74382edc 100644 --- a/assistant/src/runtime/routes/host-transfer-routes.ts +++ b/assistant/src/runtime/routes/host-transfer-routes.ts @@ -59,8 +59,11 @@ function handleTransferContentGet({ /** * Resolve Content-Length and X-Transfer-SHA256 response headers for the - * GET transfer content endpoint. Called by the HTTP adapter before - * sending the response. + * GET transfer content endpoint. Called by the HTTP adapter AFTER the handler + * runs (`http-adapter.ts:107-125`), so the entry has already been consumed by + * `getTransferContent`. We read the size/sha256 from + * `takeJustConsumedTransferMetadata`, which the proxy populates synchronously + * during the handler's `getTransferContent` call. */ function resolveTransferContentGetHeaders({ pathParams = {}, @@ -70,16 +73,14 @@ function resolveTransferContentGetHeaders({ const transferId = pathParams?.transferId; if (!transferId) return { "Content-Type": "application/octet-stream" }; - const match = findProxyByTransferId(transferId); - if (!match) return { "Content-Type": "application/octet-stream" }; - - const content = match.proxy.getTransferContent(transferId); - if (!content) return { "Content-Type": "application/octet-stream" }; + const meta = + HostTransferProxy.instance.takeJustConsumedTransferMetadata(transferId); + if (!meta) return { "Content-Type": "application/octet-stream" }; return { "Content-Type": "application/octet-stream", - "Content-Length": content.sizeBytes.toString(), - "X-Transfer-SHA256": content.sha256, + "Content-Length": meta.sizeBytes.toString(), + "X-Transfer-SHA256": meta.sha256, }; }