clawmatrix 0.1.0

package/src/model-proxy.ts

@@ -0,0 +1,441 @@
+import type { PeerManager } from "./peer-manager.ts";
+import type { ClawMatrixConfig } from "./config.ts";
+import type {
+  ModelRequest,
+  ModelResponse,
+  ModelStreamChunk,
+} from "./types.ts";
+
+const MODEL_TIMEOUT = 120_000; // 2 minutes
+
+interface PendingModelReq {
+  resolve: (value: unknown) => void;
+  reject: (error: Error) => void;
+  timer: ReturnType<typeof setTimeout>;
+  stream: boolean;
+  controller?: ReadableStreamDefaultController;
+  encoder?: TextEncoder;
+}
+
+export class ModelProxy {
+  private config: ClawMatrixConfig;
+  private peerManager: PeerManager;
+  private pending = new Map<string, PendingModelReq>();
+  private httpServer: ReturnType<typeof Bun.serve> | null = null;
+
+  constructor(config: ClawMatrixConfig, peerManager: PeerManager) {
+    this.config = config;
+    this.peerManager = peerManager;
+  }
+
+  /** Start the local HTTP proxy server for OpenAI-compatible requests. */
+  start() {
+    this.httpServer = Bun.serve({
+      port: this.config.proxyPort,
+      hostname: "127.0.0.1",
+      routes: {
+        "/v1/chat/completions": {
+          POST: async (req) => this.handleChatCompletion(req),
+        },
+        "/v1/models": {
+          GET: () => this.handleListModels(),
+        },
+      },
+      fetch() {
+        return new Response("Not Found", { status: 404 });
+      },
+    });
+  }
+
+  stop() {
+    if (this.httpServer) {
+      this.httpServer.stop();
+      this.httpServer = null;
+    }
+    for (const [, pending] of this.pending) {
+      clearTimeout(pending.timer);
+      pending.reject(new Error("Shutting down"));
+    }
+    this.pending.clear();
+  }
+
+  // ── HTTP handlers ──────────────────────────────────────────────
+  private async handleChatCompletion(req: Request): Promise<Response> {
+    let body: {
+      model: string;
+      messages: unknown[];
+      stream?: boolean;
+      temperature?: number;
+      max_tokens?: number;
+    };
+
+    try {
+      body = (await req.json()) as typeof body;
+    } catch {
+      return new Response(JSON.stringify({ error: "Invalid JSON" }), {
+        status: 400,
+        headers: { "Content-Type": "application/json" },
+      });
+    }
+
+    const modelId = body.model;
+    const route = this.peerManager.router.resolveModel(modelId);
+    if (!route) {
+      return new Response(
+        JSON.stringify({ error: { message: `Model "${modelId}" not found in cluster` } }),
+        { status: 404, headers: { "Content-Type": "application/json" } },
+      );
+    }
+
+    const stream = body.stream ?? false;
+    const requestId = crypto.randomUUID();
+
+    const frame: ModelRequest = {
+      type: "model_req",
+      id: requestId,
+      from: this.config.nodeId,
+      to: route.nodeId,
+      timestamp: Date.now(),
+      payload: {
+        model: modelId,
+        messages: body.messages,
+        temperature: body.temperature,
+        maxTokens: body.max_tokens,
+        stream,
+      },
+    };
+
+    if (stream) {
+      return this.handleStreamRequest(requestId, route.nodeId, frame);
+    } else {
+      return this.handleNonStreamRequest(requestId, route.nodeId, frame);
+    }
+  }
+
+  private handleStreamRequest(
+    requestId: string,
+    targetNodeId: string,
+    frame: ModelRequest,
+  ): Response {
+    const encoder = new TextEncoder();
+
+    const readable = new ReadableStream({
+      start: (controller) => {
+        const timer = setTimeout(() => {
+          this.pending.delete(requestId);
+          this.peerManager.router.markFailed(requestId);
+          try {
+            controller.enqueue(
+              encoder.encode(`data: ${JSON.stringify({ error: "timeout" })}\n\n`),
+            );
+            controller.enqueue(encoder.encode("data: [DONE]\n\n"));
+            controller.close();
+          } catch {
+            // controller may already be closed
+          }
+        }, MODEL_TIMEOUT);
+
+        this.pending.set(requestId, {
+          resolve: () => {},
+          reject: () => {},
+          timer,
+          stream: true,
+          controller,
+          encoder,
+        });
+
+        const sent = this.peerManager.sendTo(targetNodeId, frame);
+        if (!sent) {
+          this.pending.delete(requestId);
+          clearTimeout(timer);
+          controller.enqueue(
+            encoder.encode(
+              `data: ${JSON.stringify({ error: "Cannot reach model node" })}\n\n`,
+            ),
+          );
+          controller.enqueue(encoder.encode("data: [DONE]\n\n"));
+          controller.close();
+        }
+      },
+    });
+
+    return new Response(readable, {
+      headers: {
+        "Content-Type": "text/event-stream",
+        "Cache-Control": "no-cache",
+        Connection: "keep-alive",
+      },
+    });
+  }
+
+  private async handleNonStreamRequest(
+    requestId: string,
+    targetNodeId: string,
+    frame: ModelRequest,
+  ): Promise<Response> {
+    try {
+      const result = await new Promise<ModelResponse["payload"]>(
+        (resolve, reject) => {
+          const timer = setTimeout(() => {
+            this.pending.delete(requestId);
+            this.peerManager.router.markFailed(requestId);
+            reject(new Error("Model request timed out"));
+          }, MODEL_TIMEOUT);
+
+          this.pending.set(requestId, {
+            resolve: resolve as (v: unknown) => void,
+            reject,
+            timer,
+            stream: false,
+          });
+
+          const sent = this.peerManager.sendTo(targetNodeId, frame);
+          if (!sent) {
+            this.pending.delete(requestId);
+            clearTimeout(timer);
+            reject(new Error("Cannot reach model node"));
+          }
+        },
+      );
+
+      if (!result.success) {
+        return new Response(
+          JSON.stringify({ error: { message: result.error } }),
+          { status: 502, headers: { "Content-Type": "application/json" } },
+        );
+      }
+
+      // OpenAI-compatible response
+      return new Response(
+        JSON.stringify({
+          id: `chatcmpl-${requestId}`,
+          object: "chat.completion",
+          created: Math.floor(Date.now() / 1000),
+          model: frame.payload.model,
+          choices: [
+            {
+              index: 0,
+              message: { role: "assistant", content: result.content },
+              finish_reason: "stop",
+            },
+          ],
+          usage: result.usage
+            ? {
+                prompt_tokens: result.usage.inputTokens,
+                completion_tokens: result.usage.outputTokens,
+                total_tokens: result.usage.inputTokens + result.usage.outputTokens,
+              }
+            : undefined,
+        }),
+        { headers: { "Content-Type": "application/json" } },
+      );
+    } catch (err) {
+      return new Response(
+        JSON.stringify({ error: { message: err instanceof Error ? err.message : String(err) } }),
+        { status: 502, headers: { "Content-Type": "application/json" } },
+      );
+    }
+  }
+
+  private handleListModels(): Response {
+    const models = this.peerManager.router
+      .getAllPeers()
+      .flatMap((p) =>
+        p.models.map((m) => ({
+          id: m.id,
+          object: "model",
+          created: 0,
+          owned_by: m.provider,
+        })),
+      );
+
+    return new Response(
+      JSON.stringify({ object: "list", data: models }),
+      { headers: { "Content-Type": "application/json" } },
+    );
+  }
+
+  // ── Incoming frame handlers ────────────────────────────────────
+  handleModelResponse(frame: ModelResponse) {
+    if (this.peerManager.router.isFailed(frame.id)) return;
+    const pending = this.pending.get(frame.id);
+    if (!pending || pending.stream) return;
+
+    clearTimeout(pending.timer);
+    this.pending.delete(frame.id);
+    pending.resolve(frame.payload);
+  }
+
+  handleModelStream(frame: ModelStreamChunk) {
+    if (this.peerManager.router.isFailed(frame.id)) return;
+    const pending = this.pending.get(frame.id);
+    if (!pending?.stream || !pending.controller || !pending.encoder) return;
+
+    try {
+      if (frame.payload.done) {
+        // Send final chunk with usage if available
+        if (frame.payload.usage) {
+          const usageChunk = {
+            id: `chatcmpl-${frame.id}`,
+            object: "chat.completion.chunk",
+            choices: [{ index: 0, delta: {}, finish_reason: "stop" }],
+            usage: {
+              prompt_tokens: frame.payload.usage.inputTokens,
+              completion_tokens: frame.payload.usage.outputTokens,
+              total_tokens: frame.payload.usage.inputTokens + frame.payload.usage.outputTokens,
+            },
+          };
+          pending.controller.enqueue(
+            pending.encoder.encode(`data: ${JSON.stringify(usageChunk)}\n\n`),
+          );
+        }
+        pending.controller.enqueue(pending.encoder.encode("data: [DONE]\n\n"));
+        pending.controller.close();
+        clearTimeout(pending.timer);
+        this.pending.delete(frame.id);
+      } else {
+        const chunk = {
+          id: `chatcmpl-${frame.id}`,
+          object: "chat.completion.chunk",
+          choices: [
+            {
+              index: 0,
+              delta: { content: frame.payload.delta },
+              finish_reason: null,
+            },
+          ],
+        };
+        pending.controller.enqueue(
+          pending.encoder.encode(`data: ${JSON.stringify(chunk)}\n\n`),
+        );
+      }
+    } catch {
+      // controller may be closed
+      clearTimeout(pending.timer);
+      this.pending.delete(frame.id);
+    }
+  }
+
+  /** Handle model_req locally: forward to OpenClaw's configured model provider. */
+  async handleModelRequest(frame: ModelRequest): Promise<void> {
+    const { id, from, payload } = frame;
+
+    // Find the model in our local config
+    const model = this.config.models.find((m) => m.id === payload.model);
+    if (!model) {
+      this.peerManager.sendTo(from, {
+        type: "model_res",
+        id,
+        from: this.config.nodeId,
+        to: from,
+        timestamp: Date.now(),
+        payload: { success: false, error: `Model "${payload.model}" not available locally` },
+      } satisfies ModelResponse);
+      return;
+    }
+
+    try {
+      // Use OpenClaw's API to run the model locally
+      // This is done by calling the local gateway's /v1/chat/completions endpoint
+      // The local OpenClaw gateway handles provider routing
+      const localUrl = `http://127.0.0.1:${process.env.OPENCLAW_GATEWAY_PORT ?? 3000}/v1/chat/completions`;
+
+      const response = await fetch(localUrl, {
+        method: "POST",
+        headers: { "Content-Type": "application/json" },
+        body: JSON.stringify({
+          model: `${model.provider}/${model.id}`,
+          messages: payload.messages,
+          temperature: payload.temperature,
+          max_tokens: payload.maxTokens,
+          stream: payload.stream,
+        }),
+      });
+
+      if (payload.stream) {
+        // Read SSE stream and forward as model_stream frames
+        const reader = response.body?.getReader();
+        if (!reader) throw new Error("No response body");
+
+        const decoder = new TextDecoder();
+        let buffer = "";
+
+        while (true) {
+          const { done, value } = await reader.read();
+          if (done) break;
+
+          buffer += decoder.decode(value, { stream: true });
+          const lines = buffer.split("\n");
+          buffer = lines.pop()!;
+
+          for (const line of lines) {
+            if (!line.startsWith("data: ")) continue;
+            const data = line.slice(6).trim();
+            if (data === "[DONE]") {
+              this.peerManager.sendTo(from, {
+                type: "model_stream",
+                id,
+                from: this.config.nodeId,
+                to: from,
+                timestamp: Date.now(),
+                payload: { delta: "", done: true },
+              } satisfies ModelStreamChunk);
+              break;
+            }
+
+            try {
+              const parsed = JSON.parse(data);
+              const delta = parsed.choices?.[0]?.delta?.content ?? "";
+              if (delta) {
+                this.peerManager.sendTo(from, {
+                  type: "model_stream",
+                  id,
+                  from: this.config.nodeId,
+                  to: from,
+                  timestamp: Date.now(),
+                  payload: { delta, done: false },
+                } satisfies ModelStreamChunk);
+              }
+            } catch {
+              // skip malformed chunks
+            }
+          }
+        }
+      } else {
+        const result = (await response.json()) as {
+          choices?: { message?: { content?: string } }[];
+          usage?: { prompt_tokens: number; completion_tokens: number };
+        };
+        const content = result.choices?.[0]?.message?.content ?? "";
+        const usage = result.usage;
+
+        this.peerManager.sendTo(from, {
+          type: "model_res",
+          id,
+          from: this.config.nodeId,
+          to: from,
+          timestamp: Date.now(),
+          payload: {
+            success: true,
+            content,
+            usage: usage
+              ? { inputTokens: usage.prompt_tokens, outputTokens: usage.completion_tokens }
+              : undefined,
+          },
+        } satisfies ModelResponse);
+      }
+    } catch (err) {
+      this.peerManager.sendTo(from, {
+        type: "model_res",
+        id,
+        from: this.config.nodeId,
+        to: from,
+        timestamp: Date.now(),
+        payload: {
+          success: false,
+          error: err instanceof Error ? err.message : String(err),
+        },
+      } satisfies ModelResponse);
+    }
+  }
+}