Fix test

packages/opencode/src/control-plane/workspace-router-middleware.ts

@@ -1,23 +1,57 @@
-import type { MiddlewareHandler } from "hono"
+import { Hono, type MiddlewareHandler } from "hono"
 import { Flag } from "../flag/flag"
+import { Session } from "../session"
 import { getAdaptor } from "./adaptors"
 import { Workspace } from "./workspace"
-import { WorkspaceContext } from "./workspace-context"
 
-// This middleware forwards all non-GET requests if the workspace is a
+// These are routes that we forward to a workspace, expressed this way
-// remote. The remote workspace needs to handle session mutations
+// because we auto-infer the workspace differently for different
+// routes
+const Router = new Hono()
+  .all("/session", async (c) => {
+    if (c.req.method === "GET") return c.notFound()
+
+    const body = await c.req.json().catch(() => undefined)
+    if (!body || typeof body.workspaceID !== "string") {
+      return c.notFound()
+    }
+    return c.text(body.workspaceID)
+  })
+  .all("/session/status", async (c) => {
+    return c.notFound()
+  })
+  .all("/session/:sessionID/*", async (c) => {
+    if (c.req.method === "GET") return c.notFound()
+
+    const info = await Session.get(c.req.param("sessionID")).catch(() => undefined)
+    if (!info?.workspaceID) return c.notFound()
+    return c.text(info.workspaceID)
+  })
+
 async function routeRequest(req: Request) {
-  // Right now, we need to forward all requests to the workspace
+  let workspaceID: string | null = null
-  // because we don't have syncing. In the future all GET requests
+
-  // which don't mutate anything will be handled locally
+  const match = await Router.fetch(req.clone())
+  if (match.ok) {
+    workspaceID = await match.text()
+  } else {
+    // Fallback to a header to force routing
     //
-  // if (req.method === "GET") return
+    // This header is temporary: we allow the client to force a request
+    // to be forwarded to a workspace with it, regardless of the URL.
+    // This is only needed because we don't sync yet; when we do we can
+    // handle a lot more requests locally and the client won't have to
+    // force this
+    workspaceID = req.headers.get("x-opencode-workspace")
+  }
 
-  if (!WorkspaceContext.workspaceID) return
+  if (workspaceID == null) {
+    return
+  }
 
-  const workspace = await Workspace.get(WorkspaceContext.workspaceID)
+  const workspace = await Workspace.get(workspaceID)
   if (!workspace) {
-    return new Response(`Workspace not found: ${WorkspaceContext.workspaceID}`, {
+    return new Response(`Workspace not found: ${workspaceID}`, {
       status: 500,
       headers: {
         "content-type": "text/plain; charset=utf-8",
@@ -26,8 +60,9 @@ async function routeRequest(req: Request) {
   }
 
   const adaptor = await getAdaptor(workspace.type)
+  const url = new URL(req.url)
 
-  return adaptor.fetch(workspace, `${new URL(req.url).pathname}${new URL(req.url).search}`, {
+  return adaptor.fetch(workspace, `${url.pathname}${url.search}`, {
     method: req.method,
     body: req.method === "GET" || req.method === "HEAD" ? undefined : await req.arrayBuffer(),
     signal: req.signal,
@@ -36,7 +71,6 @@ async function routeRequest(req: Request) {
 }
 
 export const WorkspaceRouterMiddleware: MiddlewareHandler = async (c, next) => {
-  // Only available in development for now
   if (!Flag.OPENCODE_EXPERIMENTAL_WORKSPACES) {
     return next()
   }

packages/opencode/test/control-plane/workspace-router-middleware.test.ts

@@ -98,12 +98,14 @@ async function setup(state: State) {
     id2,
     app,
     async request(input: RequestInfo | URL, init?: RequestInit) {
+      const headers = new Headers(init?.headers)
+      headers.set("x-opencode-workspace", state.workspace === "first" ? id1 : id2)
       return Instance.provide({
         directory: tmp.path,
         fn: async () =>
           WorkspaceContext.provide({
             workspaceID: state.workspace === "first" ? id1 : id2,
-            fn: () => app.request(input, init),
+            fn: () => app.request(input, { ...init, headers }),
           }),
       })
     },
@@ -120,7 +122,7 @@ describe("control-plane/session-proxy-middleware", () => {
     const ctx = await setup(state)
 
     ctx.app.post("/session/foo", (c) => c.text("local", 200))
-    const response = await ctx.request("http://workspace.test/session/foo?x=1", {
+    const response = await ctx.request("http://workspace.test/session", {
       method: "POST",
       body: JSON.stringify({ hello: "world" }),
       headers: {
@@ -133,7 +135,7 @@ describe("control-plane/session-proxy-middleware", () => {
     expect(state.calls).toEqual([
       {
         method: "POST",
-        url: "/session/foo?x=1",
+        url: "/session",
         body: '{"hello":"world"}',
       },
     ])