@mathew-macharia/black-mamba 1.0.4 → 1.0.6

package/src/session/planner.js

@@ -0,0 +1,177 @@
+// Mapping planner: decides, for one pi request, exactly which DeepSeek chat/parent to attach to and what text to
+// send. Every outcome is a named strategy with a reason. There is no implicit fallback: either the parent node
+// exists and the context pi sent verifiably extends it (-> delta), or DeepSeek has no matching context and a
+// fresh chat is seeded with the full history (-> seed). Both outcomes reproduce pi's context exactly.
+//
+// Strategies
+//   stateless_seed           no x-session-id: full history into a throwaway DeepSeek chat, nothing recorded
+//   root_seed                stateful request with no usable parent node (first turn, unknown session, /tree to root)
+//   delta                    parent node found and verified: send only the messages DeepSeek has not seen
+//   reseed_context_diverged  parent node found but pi's context no longer extends it (compaction, context_edit, ...)
+//   reseed_upstream_rejected DeepSeek rejected the stored chat/parent (chat deleted or expired on their side)
+
+import { BadRequestError, HttpError } from "../errors.js";
+import {
+  assistantFingerprintFromMessage,
+  contentToText,
+  fingerprintsEqual,
+  hasLeadingSystem,
+  isCompactionSummaryMessage,
+  normalizeRole,
+  prefixHash,
+  systemHash,
+} from "../pi/messages.js";
+import { renderDeltaPrompt, renderSeedPrompt, renderSystemPrompt } from "../pi/prompt.js";
+
+export const STRATEGY = Object.freeze({
+  STATELESS_SEED: "stateless_seed",
+  ROOT_SEED: "root_seed",
+  DELTA: "delta",
+  RESEED_DIVERGED: "reseed_context_diverged",
+  RESEED_UPSTREAM: "reseed_upstream_rejected",
+});
+
+function contextFingerprint(messages) {
+  return { messageCount: messages.length, systemHash: systemHash(messages), prefixHash: prefixHash(messages) };
+}
+
+/**
+ * Checks that `messages` is exactly "the parent's request context + the parent's reply + new messages".
+ * @returns {{ ok: true } | { ok: false, reason: string, details: object }}
+ */
+export function verifyContextExtendsNode(messages, node) {
+  const parentCount = node.context.messageCount;
+  const firstBody = hasLeadingSystem(messages) ? 1 : 0;
+  // pi replaces the history with a summary message on compaction; the context then no longer extends any earlier node.
+  const startsWithCompaction = messages.length > firstBody && isCompactionSummaryMessage(messages[firstBody]);
+
+  if (messages.length <= parentCount) {
+    return {
+      ok: false,
+      reason: startsWithCompaction ? "compaction" : "context_shorter_than_parent",
+      details: { messageCount: messages.length, parentMessageCount: parentCount },
+    };
+  }
+  const expectedPrefix = prefixHash(messages, parentCount);
+  if (expectedPrefix !== node.context.prefixHash) {
+    return {
+      ok: false,
+      reason: startsWithCompaction ? "compaction" : "context_prefix_changed",
+      details: { parentMessageCount: parentCount, expectedPrefixHash: node.context.prefixHash, actualPrefixHash: expectedPrefix },
+    };
+  }
+  const reply = messages[parentCount];
+  if (normalizeRole(reply.role) !== "assistant") {
+    return { ok: false, reason: "parent_reply_missing", details: { indexOfReply: parentCount, roleFound: normalizeRole(reply.role) } };
+  }
+  if (!fingerprintsEqual(assistantFingerprintFromMessage(reply), node.context.assistant)) {
+    return { ok: false, reason: "parent_reply_modified", details: { indexOfReply: parentCount } };
+  }
+  return { ok: true };
+}
+
+function seedPlan({ request, strategy, reason, details = null, session = null, parentNode = null }) {
+  const { messages, tools, routing } = request;
+  return Object.freeze({
+    strategy,
+    reason,
+    details,
+    session,
+    parentNode,
+    prompt: renderSeedPrompt(messages, tools),
+    deepseek: { sessionId: null, parentMessageId: null }, // null sessionId => the handler creates a fresh DeepSeek chat
+    cachedTokens: 0,
+    deltaMessageCount: messages.length,
+    record: routing
+      ? { piEntryId: routing.entryId, piParentEntryId: routing.parentEntryId, context: contextFingerprint(messages) }
+      : null,
+  });
+}
+
+/** Resolves the parent node, adopting a forked session's graph when pi tells us where it was forked from. */
+async function resolveParent({ graph, routing, logger }) {
+  let session = await graph.getSession(routing.sessionId);
+  let parentNode = null;
+  let adoptedFrom = null;
+
+  if (routing.parentEntryId) {
+    parentNode = graph.getNode(session, routing.parentEntryId);
+
+    if (!parentNode && routing.parentSessionId && routing.parentSessionId !== routing.sessionId) {
+      const parentSession = await graph.getSession(routing.parentSessionId);
+      const candidate = graph.getNode(parentSession, routing.parentEntryId);
+      if (candidate) {
+        session = session ?? (await graph.createSession(routing.sessionId));
+        await graph.adoptNodes(session, parentSession);
+        parentNode = graph.getNode(session, routing.parentEntryId);
+        adoptedFrom = routing.parentSessionId;
+        logger.info(`Fork detected: session ${routing.sessionId} adopted ${graph.nodeCount(parentSession)} node(s) from ${routing.parentSessionId}.`);
+      }
+    }
+  }
+
+  if (!session) session = await graph.createSession(routing.sessionId);
+  return { session, parentNode, adoptedFrom };
+}
+
+/**
+ * @param {{ request: ReturnType<import("../pi/request.js").parseChatRequest>, graph: import("./graph.js").SessionGraph, logger: object }} deps
+ */
+export async function planCompletion({ request, graph, logger }) {
+  const { messages, tools, routing } = request;
+
+  if (!routing) {
+    return seedPlan({ request, strategy: STRATEGY.STATELESS_SEED, reason: "no_session_routing" });
+  }
+
+  const { session, parentNode, adoptedFrom } = await resolveParent({ graph, routing, logger });
+
+  if (!parentNode) {
+    let reason = "no_parent_entry";
+    if (routing.parentEntryId) reason = graph.nodeCount(session) === 0 ? "session_unknown" : "parent_entry_unknown";
+    return seedPlan({ request, strategy: STRATEGY.ROOT_SEED, reason, session, details: { parentEntryId: routing.parentEntryId } });
+  }
+
+  const verification = verifyContextExtendsNode(messages, parentNode);
+  if (!verification.ok) {
+    return seedPlan({
+      request,
+      strategy: STRATEGY.RESEED_DIVERGED,
+      reason: verification.reason,
+      details: { ...verification.details, parentEntryId: parentNode.piEntryId },
+      session,
+      parentNode,
+    });
+  }
+
+  const deltaMessages = messages.slice(parentNode.context.messageCount + 1);
+  if (deltaMessages.length === 0) {
+    throw new HttpError(409, "Request contains no messages after the parent node's reply; nothing to send to DeepSeek.", {
+      code: "no_new_messages",
+      details: { parentEntryId: parentNode.piEntryId, messageCount: messages.length },
+    });
+  }
+
+  const fingerprint = contextFingerprint(messages);
+  const systemChanged = fingerprint.systemHash !== parentNode.context.systemHash;
+  const systemUpdate = systemChanged ? renderSystemPrompt(hasLeadingSystem(messages) ? contentToText(messages[0].content) : "", tools) : null;
+
+  return Object.freeze({
+    strategy: STRATEGY.DELTA,
+    reason: adoptedFrom ? `forked_from:${adoptedFrom}` : systemChanged ? "system_prompt_changed" : null,
+    details: null,
+    session,
+    parentNode,
+    prompt: renderDeltaPrompt(deltaMessages, tools, { systemUpdate }),
+    deepseek: { sessionId: parentNode.deepseek.sessionId, parentMessageId: parentNode.deepseek.messageId },
+    cachedTokens: graph.branchCachedTokens(session, parentNode),
+    deltaMessageCount: deltaMessages.length,
+    record: { piEntryId: routing.entryId, piParentEntryId: parentNode.piEntryId, context: fingerprint },
+  });
+}
+
+/** Explicit re-seed after DeepSeek rejected the chat/parent a delta plan referenced. */
+export function reseedAfterUpstreamRejection(request, plan, details) {
+  if (!plan.session) throw new BadRequestError("Cannot re-seed a stateless request", "reseed_without_session");
+  return seedPlan({ request, strategy: STRATEGY.RESEED_UPSTREAM, reason: "upstream_context_rejected", details, session: plan.session, parentNode: plan.parentNode });
+}

package/src/session/store.js

@@ -0,0 +1,100 @@
+// Disk persistence for session graphs: one JSON file per pi session, atomic writes, TTL pruning.
+
+import crypto from "node:crypto";
+import path from "node:path";
+import { existsSync } from "node:fs";
+import { mkdir, readFile, readdir, rename, stat as statFs, unlink, writeFile } from "node:fs/promises";
+
+export const SESSION_SCHEMA_VERSION = 2;
+
+export function sanitizeSessionFilename(sessionId) {
+  const safePrefix = String(sessionId).replace(/[^a-zA-Z0-9_-]/g, "_").slice(0, 48);
+  const hash = crypto.createHash("sha256").update(String(sessionId)).digest("hex").slice(0, 16);
+  return `${safePrefix}_${hash}.json`;
+}
+
+export class SessionStore {
+  constructor({ dir, ttlMs, logger }) {
+    this.dir = dir;
+    this.ttlMs = ttlMs;
+    this.logger = logger;
+  }
+
+  filePath(sessionId) {
+    return path.join(this.dir, sanitizeSessionFilename(sessionId));
+  }
+
+  async init() {
+    if (!existsSync(this.dir)) await mkdir(this.dir, { recursive: true });
+  }
+
+  /**
+   * Returns the stored session or null. Expired files are removed. Files written by the pre-v2 schema are
+   * ignored (their node keys mixed pi and DeepSeek ids and cannot be trusted); they are reported once and
+   * left to TTL pruning.
+   */
+  async read(sessionId, now = Date.now()) {
+    const filePath = this.filePath(sessionId);
+    if (!existsSync(filePath)) return null;
+
+    let parsed;
+    try {
+      parsed = JSON.parse(await readFile(filePath, "utf8"));
+    } catch (err) {
+      this.logger.warn(`Session file ${path.basename(filePath)} is unreadable (${err.message}); treating session as unknown.`);
+      return null;
+    }
+    if (!parsed || typeof parsed !== "object") return null;
+
+    if (parsed.schemaVersion !== SESSION_SCHEMA_VERSION) {
+      this.logger.warn(
+        `Session file ${path.basename(filePath)} uses schema ${parsed.schemaVersion ?? 1} (expected ${SESSION_SCHEMA_VERSION}); ignoring it. The branch will be re-seeded.`
+      );
+      return null;
+    }
+    if (now - (parsed.lastAccessedAt || parsed.createdAt || 0) > this.ttlMs) {
+      await unlink(filePath).catch(() => {});
+      return null;
+    }
+    return parsed;
+  }
+
+  async write(sessionId, session) {
+    await this.init();
+    const filePath = this.filePath(sessionId);
+    const tmpPath = `${filePath}.tmp.${process.pid}.${Date.now()}`;
+    const json = JSON.stringify(session, null, 2);
+    await writeFile(tmpPath, json, "utf8");
+    try {
+      await rename(tmpPath, filePath);
+    } catch (err) {
+      // Windows can refuse to replace a file that another process has open; write in place instead.
+      this.logger.warn(`Atomic rename failed for ${path.basename(filePath)} (${err.message}); writing in place.`);
+      await writeFile(filePath, json, "utf8");
+      await unlink(tmpPath).catch(() => {});
+    }
+  }
+
+  async remove(sessionId) {
+    await unlink(this.filePath(sessionId)).catch(() => {});
+  }
+
+  async pruneOlderThan(now = Date.now()) {
+    if (!existsSync(this.dir)) return 0;
+    let removed = 0;
+    for (const file of await readdir(this.dir)) {
+      if (!file.endsWith(".json")) continue;
+      const filePath = path.join(this.dir, file);
+      try {
+        const info = await statFs(filePath);
+        if (now - info.mtimeMs > this.ttlMs) {
+          await unlink(filePath);
+          removed++;
+        }
+      } catch {
+        // file vanished concurrently
+      }
+    }
+    return removed;
+  }
+}

package/src/utils/hash.js

@@ -0,0 +1,27 @@
+import crypto from "node:crypto";
+
+export function sha256Hex(input) {
+  return crypto.createHash("sha256").update(input, "utf8").digest("hex");
+}
+
+/** Short (16 hex chars) digest used for fingerprints that are compared, never used as a key. */
+export function shortHash(input) {
+  return sha256Hex(input).slice(0, 16);
+}
+
+/** Deterministic JSON: object keys sorted recursively so semantically equal values hash equally. */
+export function canonicalJson(value) {
+  return JSON.stringify(sortKeys(value));
+}
+
+function sortKeys(value) {
+  if (Array.isArray(value)) return value.map(sortKeys);
+  if (value && typeof value === "object") {
+    const out = {};
+    for (const key of Object.keys(value).sort()) {
+      out[key] = sortKeys(value[key]);
+    }
+    return out;
+  }
+  return value;
+}

package/src/completions.js

@@ -1,149 +0,0 @@
-// POST /v1/chat/completions handler:
-// Stateless prompt dump -> rotational account selection -> human pacing -> DeepSeek turn -> OpenAI stream/JSON.
-
-import { CONFIG } from "./config.js";
-import { createLogger } from "./logger.js";
-import { consumeDeepSeekSseStream } from "./deepseek/sse.js";
-import { renderSeedPrompt } from "./prompt.js";
-import { parseChatRequest } from "./request.js";
-import {
-  buildUsageStats,
-  createCompletionMetadata,
-  formatCompletionResponse,
-  formatStreamDeltaEvent,
-  formatStreamFinishEvent,
-  formatStreamReasoningEvent,
-  formatStreamRoleEvent,
-  formatStreamToolCallArgsEvent,
-  formatStreamToolCallStartEvent,
-  formatStreamUsageEvent,
-} from "./response.js";
-import { createToolCallStreamProcessor, parseToolCallsFromText } from "./tool-calls.js";
-import { estimateTokenCount } from "./tokens.js";
-import { readJsonBody, sendJson, startSse, writeSse } from "./http.js";
-
-/**
- * @param {object} deps
- * @param {ReturnType<import("./accounts.js").createAccountPool>} deps.pool
- * @param {object} [deps.logger]
- */
-export function createCompletionHandler({ pool, logger = createLogger("Completions"), keepAliveIntervalMs = CONFIG.server.keepAliveIntervalMs }) {
-  async function respondJson({ res, request, account, streamBody, prompt, meta }) {
-    const result = await consumeDeepSeekSseStream(streamBody);
-    let text = result.reply;
-    let toolCalls = [];
-
-    if (request.tools.length > 0) {
-      const parsed = parseToolCallsFromText(result.reply);
-      text = parsed.cleanText;
-      toolCalls = parsed.toolCalls;
-    }
-
-    const turnTokens = result.turnTokens > 0
-      ? result.turnTokens
-      : estimateTokenCount(prompt) + estimateTokenCount(result.reply);
-
-    const usage = buildUsageStats({ prompt, reply: result.reply, turnTokens });
-    const payload = formatCompletionResponse({
-      ...meta,
-      model: request.model,
-      content: text,
-      toolCalls,
-      usage,
-    });
-    sendJson(res, 200, payload);
-  }
-
-  async function respondStream({ req, res, request, account, streamBody, prompt, meta }) {
-    startSse(res, req);
-    writeSse(res, formatStreamRoleEvent(meta));
-    const keepAlive = setInterval(() => writeSse(res, ": keep-alive\n\n"), keepAliveIntervalMs);
-    keepAlive.unref?.();
-
-    let emittedText = "";
-    const emittedToolCalls = [];
-    const emitText = (content) => {
-      emittedText += content;
-      writeSse(res, formatStreamDeltaEvent(meta, content));
-    };
-    const emitToolCall = ({ index, id, name, argumentsString }) => {
-      emittedToolCalls.push({ name, argumentsString });
-      writeSse(res, formatStreamToolCallStartEvent(meta, { toolCallIndex: index, id, name }));
-      writeSse(res, formatStreamToolCallArgsEvent(meta, { toolCallIndex: index, argumentsChunk: argumentsString }));
-    };
-
-    const processor = request.tools.length > 0
-      ? createToolCallStreamProcessor({ onTextDelta: emitText, onToolCall: emitToolCall })
-      : null;
-
-    let result;
-    try {
-      result = await consumeDeepSeekSseStream(streamBody, (chunk) => {
-        if (chunk.isThinking) writeSse(res, formatStreamReasoningEvent(meta, chunk.text));
-        else if (processor) processor.processChunk(chunk.text);
-        else emitText(chunk.text);
-      });
-      processor?.flush();
-    } finally {
-      clearInterval(keepAlive);
-    }
-
-    const turnTokens = result.turnTokens > 0
-      ? result.turnTokens
-      : estimateTokenCount(prompt) + estimateTokenCount(result.reply);
-
-    writeSse(
-      res,
-      formatStreamFinishEvent(meta, {
-        finishReason: emittedToolCalls.length > 0 ? "tool_calls" : "stop",
-      })
-    );
-
-    if (request.includeUsage) {
-      writeSse(res, formatStreamUsageEvent(meta, buildUsageStats({ prompt, reply: result.reply, turnTokens })));
-    }
-    writeSse(res, "data: [DONE]\n\n");
-    res.end();
-  }
-
-  return async function handleChatCompletions(req, res) {
-    const body = await readJsonBody(req);
-    const request = parseChatRequest(body);
-    const prompt = renderSeedPrompt(request.messages, request.tools);
-
-    // Rotational account selection ("Wi-Fi router" hub)
-    const account = pool.nextAccount();
-    logger.info(`Request [${request.model}] -> [${account.name}] (prompt: ${prompt.length} chars, tools: ${request.tools.length})`);
-
-    const abortController = new AbortController();
-    let finished = false;
-    if (typeof res.on === "function") {
-      res.on("close", () => {
-        if (!finished && !res.writableFinished) abortController.abort();
-      });
-    }
-
-    // Apply per-account human pacing
-    await account.pacer.apply();
-
-    try {
-      const { streamBody } = await account.deepseek.openCompletionStream({
-        prompt,
-        thinkingEnabled: request.thinkingEnabled,
-        searchEnabled: request.searchEnabled,
-        signal: abortController.signal,
-      });
-
-      const meta = { ...createCompletionMetadata(), model: request.model };
-      if (request.stream) {
-        await respondStream({ req, res, request, account, streamBody, prompt, meta });
-      } else {
-        await respondJson({ res, request, account, streamBody, prompt, meta });
-      }
-    } finally {
-      finished = true;
-      account.pacer.recordCompletionFinished();
-    }
-  };
-}
-