4-26 树rust-2

wolai-frontend/src/app/api/mnote-web/stream/route.ts

@@ -1,76 +1,72 @@
+import { NextResponse } from "next/server";
+import { isConvexEnabled } from "@/lib/convex/enabled";
 import { getAuthedConvexClient } from "@/lib/convex/route";
 import {
-  buildDocumentBridgeContext,
+  buildDocumentBridgeContextWithActor,
   buildDocumentQueryEnvelope,
-  documentBridgeErrorResponse,
 } from "@/lib/documents/bridge";
 import {
   executeRustBridgeQueryTransport,
   resolveRustBridgeQueryPlan,
 } from "@/lib/documents/rust-runtime";
-import { mapSidebarDatasetListQueryResultToInitialData } from "@/lib/sidebar-data";
+import type { SidebarDatasetListQueryResult } from "@/lib/sidebar-data";
+import { attachKernelFileTreeProjection, resolveKernelFileTreeProjection } from "@/lib/server/kernel-file-tree";
 import {
-  attachKernelFileTreeProjection,
-  resolveKernelFileTreeProjection,
-} from "@/lib/server/kernel-file-tree";
+  streamTreeFrames,
+  type TreeStreamOverview,
+  type TreeStreamSnapshotPayload,
+} from "@/lib/tree-stream/server";
 
 export const dynamic = "force-dynamic";
+export const runtime = "nodejs";
 
-function toSseFrame(event: string, data: unknown) {
-  return `event: ${event}\ndata: ${JSON.stringify(data ?? null)}\n\n`;
+function readNumberParam(url: URL, name: string): number | null {
+  const raw = url.searchParams.get(name);
+  if (!raw?.trim()) {
+    return null;
+  }
+  const parsed = Number(raw);
+  return Number.isFinite(parsed) ? parsed : null;
+}
+
+function encodeSseFrame(event: string, payload: unknown) {
+  return `event: ${event}\ndata: ${JSON.stringify(payload)}\n\n`;
 }
 
 export async function GET(request: Request) {
-  try {
-    const requestUrl = new URL(request.url);
-    const workspaceId = String(requestUrl.searchParams.get("workspaceId") || "").trim();
-    const cursor = String(requestUrl.searchParams.get("cursor") || "").trim() || null;
+  if (!isConvexEnabled()) {
+    return NextResponse.json({ error: "当前仅支持 Convex 模式" }, { status: 501 });
+  }
 
-    if (!workspaceId) {
-      return Response.json({ error: "缺少 workspaceId" }, { status: 400 });
-    }
+  const requestUrl = new URL(request.url);
+  const workspaceId = requestUrl.searchParams.get("workspaceId")?.trim();
+  if (!workspaceId) {
+    return NextResponse.json({ error: "缺少 workspaceId" }, { status: 400 });
+  }
 
-    const { auth, client } = await getAuthedConvexClient();
-    const context = await buildDocumentBridgeContext({
-      request,
-      workspaceId,
-    });
+  const { auth, client } = await getAuthedConvexClient();
+  const actor = {
+    actorType: "user",
+    actorId: auth.userId,
+    sessionId: null,
+  };
+  const context = buildDocumentBridgeContextWithActor({
+    request,
+    actor,
+    workspaceId,
+    source: {
+      channel: "next_mnote_web_stream",
+      client: "wolai-frontend",
+    },
+  });
 
-    const sidebarEnvelope = buildDocumentQueryEnvelope({
-      name: "sidebar.dataset.list",
-      payload: {
-        workspaceId,
-      },
-    });
-    const sidebarPlan = await resolveRustBridgeQueryPlan({
-      context,
-      envelope: sidebarEnvelope,
-    });
-    const sidebarDataset = await executeRustBridgeQueryTransport({
-      client,
-      plan: sidebarPlan,
-    });
-    const sidebarDatasetWithFileTree = attachKernelFileTreeProjection({
-      dataset: sidebarDataset,
-      projection: await resolveKernelFileTreeProjection({
-        client,
-        request,
-        workspaceId,
-        actor: {
-          actorType: "user",
-          actorId: auth.userId,
-          sessionId: null,
-        },
-        dataset: sidebarDataset,
-      }),
-    });
-
-    const overviewEnvelope = buildDocumentQueryEnvelope({
+  const loadOverview = async (): Promise<TreeStreamOverview> => {
+    const envelope = buildDocumentQueryEnvelope({
       name: "bridge.workspace.overview",
       payload: {
         workspaceId,
-        limit: 20,
-        cursor,
+        limit: 50,
+        cursor: null,
         commandStatus: null,
         eventStatus: null,
         targetPageId: null,
@@ -79,43 +75,88 @@ export async function GET(request: Request) {
         aggregateId: null,
       },
     });
-    const overviewPlan = await resolveRustBridgeQueryPlan({
+    const plan = await resolveRustBridgeQueryPlan({
       context,
-      envelope: overviewEnvelope,
+      envelope,
     });
-    const overview = await executeRustBridgeQueryTransport({
+    return executeRustBridgeQueryTransport<TreeStreamOverview>({
       client,
-      plan: overviewPlan,
+      plan,
     });
+  };
 
-    const payload = {
-      kind: "snapshot",
-      stream: "workspace",
-      projection: "sidebar_tree",
-      workspaceId,
-      rootNodeId: null,
-      cursor,
+  const loadSnapshot = async (): Promise<TreeStreamSnapshotPayload> => {
+    const envelope = buildDocumentQueryEnvelope({
+      name: "sidebar.dataset.list",
+      payload: {
+        workspaceId,
+      },
+    });
+    const plan = await resolveRustBridgeQueryPlan({
+      context,
+      envelope,
+    });
+    const dataset = await executeRustBridgeQueryTransport<SidebarDatasetListQueryResult>({
+      client,
+      plan,
+    });
+    const datasetWithFileTree = attachKernelFileTreeProjection({
+      dataset,
+      projection: await resolveKernelFileTreeProjection({
+        client,
+        request,
+        workspaceId,
+        actor,
+        dataset,
+        rootNodeId: requestUrl.searchParams.get("rootNodeId")?.trim() || null,
+        depth: readNumberParam(requestUrl, "depth"),
+      }),
+    });
+    return {
       requestId: context.requestId,
       traceId: context.traceId,
-      data: mapSidebarDatasetListQueryResultToInitialData(sidebarDatasetWithFileTree),
+      data: datasetWithFileTree,
       snapshot: {
-        dataset: sidebarDatasetWithFileTree,
-        tree:
-          sidebarDatasetWithFileTree.kernel_sidebar_projection ??
-          sidebarDatasetWithFileTree.kernelSidebarProjection ??
-          null,
+        dataset: datasetWithFileTree,
       },
-      overview,
     };
+  };
 
-    return new Response(toSseFrame("snapshot", payload), {
-      status: 200,
-      headers: {
-        "content-type": "text/event-stream; charset=utf-8",
-        "cache-control": "no-store",
-      },
-    });
-  } catch (error) {
-    return documentBridgeErrorResponse(error);
-  }
+  const encoder = new TextEncoder();
+  const stream = new ReadableStream<Uint8Array>({
+    async start(controller) {
+      try {
+        for await (const frame of streamTreeFrames({
+          workspaceId,
+          rootNodeId: requestUrl.searchParams.get("rootNodeId"),
+          initialCursor: requestUrl.searchParams.get("cursor"),
+          pollMs: readNumberParam(requestUrl, "pollMs") ?? undefined,
+          maxPolls: readNumberParam(requestUrl, "maxPolls"),
+          loadOverview,
+          loadSnapshot,
+        })) {
+          if (request.signal.aborted) {
+            break;
+          }
+          controller.enqueue(encoder.encode(encodeSseFrame(frame.event, frame.payload)));
+        }
+        controller.close();
+      } catch (error) {
+        controller.error(error);
+      }
+    },
+    cancel() {
+      return undefined;
+    },
+  });
+
+  return new NextResponse(stream, {
+    status: 200,
+    headers: {
+      "content-type": "text/event-stream; charset=utf-8",
+      "cache-control": "no-store",
+      connection: "keep-alive",
+      "x-upstream": "next-tree-stream",
+    },
+  });
 }