molecule-ai/molecule-core
35
0
Fork 2
Code Issues 16 Pull Requests 2 Actions 49 Packages Projects Releases Wiki Activity

refactor(chat_files): extract streamWorkspaceResponse helper for Upload+Download

Browse Source 
The "do request → check err → defer close → forward headers → set
status → io.Copy → log mid-stream errors" tail was duplicated between
Upload and Download. Each handler had ~12 lines that differed only in:

  - the op label in log messages ("upload" vs "download")
  - the set of response headers to forward verbatim
    (Upload: Content-Type only; Download: Content-Type +
    Content-Length + Content-Disposition)

Hoist into ChatFilesHandler.streamWorkspaceResponse(c, op,
workspaceID, forwardURL, req, forwardHeaders). Each call site
reduces to one line. Future changes — request-id forwarding,
observability metric, response-size cap, bytes-streamed log —
go in ONE place rather than two.

Same drift-prevention rationale as resolveWorkspaceForwardCreds
(#2372) and readOrLazyHealInboundSecret (#2376), applied to the
response-streaming layer of the same handlers.

Behavior preserved: existing TestChatUpload_* and TestChatDownload_*
integration tests (8 across both handlers) all pass unchanged. The
log message format is consistent across both handlers now (single
"chat_files {op}: ..." string template) — operators can grep one
prefix for both features instead of separate prefixes per handler.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
Hongming Wang 2026-04-30 08:27:45 -07:00
parent 7cb8b476ad
commit 830e4aa548
1 changed files with 24 additions and 26 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

50
workspace-server/internal/handlers/chat_files.go
View File

@ -259,24 +259,7 @@ func (h *ChatFilesHandler) Upload(c *gin.Context) {
req.ContentLength = c.Request.ContentLength
}
resp, err := h.httpClient.Do(req)
if err != nil {
log.Printf("chat_files Upload: forward to %s failed: %v", forwardURL, err)
c.JSON(http.StatusBadGateway, gin.H{"error": "workspace unreachable"})
return
}
defer resp.Body.Close()
// Stream response back. Copy headers we know are safe + the body.
if ct := resp.Header.Get("Content-Type"); ct != "" {
c.Header("Content-Type", ct)
}
c.Status(resp.StatusCode)
if _, err := io.Copy(c.Writer, resp.Body); err != nil {
// Mid-stream failure — too late to write a JSON error, just
// log so ops can correlate with the workspace's logs.
log.Printf("chat_files Upload: stream response back failed for %s: %v", workspaceID, err)
}
h.streamWorkspaceResponse(c, "upload", workspaceID, forwardURL, req, []string{"Content-Type"})
}
// Download handles GET /workspaces/:id/chat/download?path=<abs path>.
@ -351,27 +334,42 @@ func (h *ChatFilesHandler) Download(c *gin.Context) {
}
req.Header.Set("Authorization", "Bearer "+secret)
h.streamWorkspaceResponse(c, "download", workspaceID, forwardURL, req,
[]string{"Content-Type", "Content-Length", "Content-Disposition"})
}
// streamWorkspaceResponse executes the prepared forward request and
// streams the workspace's response back to the inbound caller.
// Forwards the named response headers verbatim. Centralizes the
// "do request → check err → defer close → copy headers → set status →
// io.Copy" tail that's identical between Upload and Download.
//
// op is the human-readable feature label ("upload"/"download") used
// in log messages so operators can distinguish which feature ran.
func (h *ChatFilesHandler) streamWorkspaceResponse(
c *gin.Context,
op, workspaceID, forwardURL string,
req *http.Request,
forwardHeaders []string,
) {
resp, err := h.httpClient.Do(req)
if err != nil {
log.Printf("chat_files Download: forward to %s failed: %v", forwardURL, err)
log.Printf("chat_files %s: forward to %s failed: %v", op, forwardURL, err)
c.JSON(http.StatusBadGateway, gin.H{"error": "workspace unreachable"})
return
}
defer resp.Body.Close()
// Stream response back, including the workspace's headers so the
// client gets the correct Content-Type + Content-Disposition (the
// workspace constructs them from the actual file's extension +
// basename — keeping that logic on the workspace side avoids a
// double-source-of-truth on filename encoding rules).
for _, hdr := range []string{"Content-Type", "Content-Length", "Content-Disposition"} {
for _, hdr := range forwardHeaders {
if v := resp.Header.Get(hdr); v != "" {
c.Header(hdr, v)
}
}
c.Status(resp.StatusCode)
if _, err := io.Copy(c.Writer, resp.Body); err != nil {
log.Printf("chat_files Download: stream response back failed for %s: %v", workspaceID, err)
// Mid-stream failure — too late to write a JSON error, just
// log so ops can correlate with the workspace's logs.
log.Printf("chat_files %s: stream response back failed for %s: %v", op, workspaceID, err)
}
}