gemini-compat-endpoint

import fs from "fs";
import path from "path";
import crypto from "crypto";
import { type Chat } from "@lmstudio/sdk";
import { sanitizeToolName } from "./tools";

export interface ThoughtSignature {
    signature: string;
    contentHash: string;
    timestamp: number;
}

export interface SignatureState {
    signatures: ThoughtSignature[];
}

const FILENAME = "thought_signatures.json";

function getFilePath(cwd: string): string {
    return path.join(cwd, FILENAME);
}

// Simple normalization: trim and collapse whitespace before hashing
export function computeContentHash(text: string): string {
    return crypto.createHash("sha256").update(text ? text.trim().replace(/\s+/g, " ") : "").digest("hex");
}

export async function loadSignatures(cwd: string): Promise<SignatureState> {
    try {
        const fp = getFilePath(cwd);
        const content = await fs.promises.readFile(fp, "utf-8");
        return JSON.parse(content);
    } catch {
        return { signatures: [] };
    }
}

export async function saveSignatures(cwd: string, state: SignatureState): Promise<void> {
    const fp = getFilePath(cwd);
    await fs.promises.writeFile(fp, JSON.stringify(state, null, 2), "utf-8");
}

export async function appendSignature(cwd: string, signature: string, contentHash: string) {
    const state = await loadSignatures(cwd);
    // Avoid duplicates for the exact same content hash (idempotency)
    if (!state.signatures.find(s => s.contentHash === contentHash && s.signature === signature)) {
        state.signatures.push({
            signature,
            contentHash,
            timestamp: Date.now(),
        });
        await saveSignatures(cwd, state);
    }
}

export async function pruneSignatures(cwd: string, history: Chat) {
    const state = await loadSignatures(cwd);
    if (state.signatures.length === 0) return;

    const validHashes = new Set<string>();

    for (const msg of history) {
        if (msg.getRole() === "assistant") {
            const text = msg.getText();
            if (text) {
                validHashes.add(computeContentHash(text));
            }
            const calls = (msg as any).getToolCallRequests?.() as any[] | undefined;
            if (Array.isArray(calls)) {
                for (const call of calls) {
                    const orig = String(call?.name || "tool");
                    // We assume no collisions for pruning purposes to keep it simple.
                    // If there's a collision, we might delete a signature, which is a trade-off.
                    const safe = sanitizeToolName(orig);
                    const args = call?.arguments ?? {};
                    const id = `${safe}:${JSON.stringify(args)}`;
                    validHashes.add(computeContentHash(id));
                }
            }
        }
    }

    const newSignatures = state.signatures.filter(s => validHashes.has(s.contentHash));

    if (newSignatures.length !== state.signatures.length) {
        state.signatures = newSignatures;
        await saveSignatures(cwd, state);
    }
}