From 8cfc0050bfbe2a1ab84d0cd4dd06edd674499a21 Mon Sep 17 00:00:00 2001
From: Vincent Koc <vincentkoc@ieee.org>
Date: Tue, 5 May 2026 01:11:49 -0700
Subject: [PATCH] fix(code): clean proxied workbench html

---
 internal/cli/code.go      | 11 +++++++++++
 internal/cli/code_test.go | 11 +++++++++++
 worker/src/fleet.ts       |  3 +++
 worker/test/fleet.test.ts |  2 ++
 4 files changed, 27 insertions(+)

diff --git a/internal/cli/code.go b/internal/cli/code.go
index d424172..447337f 100644
--- a/internal/cli/code.go
+++ b/internal/cli/code.go
@@ -272,6 +272,9 @@ func (b *codeBridge) handleHTTP(ctx context.Context, msg codeProxyMessage) {
 		_ = b.writeJSON(ctx, codeProxyMessage{Type: "http", ID: msg.ID, Status: 502, Error: err.Error()})
 		return
 	}
+	if isCodeHTML(resp.Header.Get("content-type")) {
+		respBody = rewriteCodeHTML(respBody)
+	}
 	headers := map[string]string{}
 	for key, values := range resp.Header {
 		if len(values) > 0 {
@@ -408,6 +411,14 @@ func codeUpstreamPath(path string) string {
 	return u.RequestURI()
 }
 
+func isCodeHTML(contentType string) bool {
+	return strings.HasPrefix(strings.ToLower(contentType), "text/html")
+}
+
+func rewriteCodeHTML(body []byte) []byte {
+	return bytes.ReplaceAll(body, []byte(`<script type="module" src=""></script>`), nil)
+}
+
 func availableLocalCodePort() string {
 	for port := 8081; port <= 8180; port++ {
 		ln, err := net.Listen("tcp", fmt.Sprintf("127.0.0.1:%d", port))
diff --git a/internal/cli/code_test.go b/internal/cli/code_test.go
index f83e729..27af46a 100644
--- a/internal/cli/code_test.go
+++ b/internal/cli/code_test.go
@@ -48,3 +48,14 @@ func TestStartCodeServerCommand(t *testing.T) {
 		t.Fatalf("startCodeServerCommand should not use pkill -f:\n%s", got)
 	}
 }
+
+func TestRewriteCodeHTMLRemovesEmptyLocaleScript(t *testing.T) {
+	input := []byte(`<script type="module" src=""></script><script type="module" src="workbench.js"></script>`)
+	got := string(rewriteCodeHTML(input))
+	if strings.Contains(got, `src=""`) {
+		t.Fatalf("rewriteCodeHTML left empty script:\n%s", got)
+	}
+	if !strings.Contains(got, `src="workbench.js"`) {
+		t.Fatalf("rewriteCodeHTML removed non-empty script:\n%s", got)
+	}
+}
diff --git a/worker/src/fleet.ts b/worker/src/fleet.ts
index 48588e4..cf472e8 100644
--- a/worker/src/fleet.ts
+++ b/worker/src/fleet.ts
@@ -1888,6 +1888,9 @@ export function codeResponseHeaders(values: Record<string, string>): Headers {
     }
     headers.set(key, value);
   }
+  if ((headers.get("content-type") || "").toLowerCase().startsWith("text/html")) {
+    headers.set("cache-control", "no-store, no-transform");
+  }
   headers.set("content-security-policy", codePortalContentSecurityPolicy);
   return headers;
 }
diff --git a/worker/test/fleet.test.ts b/worker/test/fleet.test.ts
index b7b8733..7c70ef5 100644
--- a/worker/test/fleet.test.ts
+++ b/worker/test/fleet.test.ts
@@ -573,6 +573,7 @@ describe("fleet lease identity and idle", () => {
       "content-security-policy": "default-src 'none'; script-src 'self'",
       "content-length": "123",
       "content-type": "text/html",
+      "cache-control": "public, max-age=31536000",
     });
 
     const csp = headers.get("content-security-policy") || "";
@@ -581,6 +582,7 @@ describe("fleet lease identity and idle", () => {
     expect(csp).toContain("worker-src 'self' data: blob:");
     expect(headers.get("content-length")).toBeNull();
     expect(headers.get("content-type")).toBe("text/html");
+    expect(headers.get("cache-control")).toBe("no-store, no-transform");
   });
 
   it("serves WebVNC pages only for desktop leases and requires an agent upgrade", async () => {