summaryrefslogtreecommitdiff
path: root/packages/ai/src/cache.ts
blob: 5d59163e422764ddd71880c6b76744ca929743a3 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
// memoize.ts (Bun-compatible, no Node Buffers)
import { mkdir } from "node:fs/promises";
import path from "node:path";

type MemoOpts<V> = {
  ttlMs?: number; // time-to-live for entries
  maxEntries?: number; // cap; oldest (LRU) evicted
  persistDir?: string; // set to enable disk cache (e.g. ".cache/memo")
  keyFn?: (...args: any[]) => string; // custom key if you need it
  cacheErrors?: boolean; // default false
};

type Entry<V> = {
  v: V;
  exp: number | null;
  at: number; // last hit (LRU)
};

const enc = new TextEncoder();
const dec = new TextDecoder();

const stableStringify = (x: any): string => {
  const seen = new WeakSet();
  const S = (v: any): any => {
    if (v && typeof v === "object") {
      if (seen.has(v)) return "[Circular]";
      seen.add(v);
      if (Array.isArray(v)) return v.map(S);
      return Object.fromEntries(
        Object.keys(v)
          .sort()
          .map((k) => [k, S(v[k])]),
      );
    }
    if (typeof v === "function") return `[Function:${v.name || "anon"}]`;
    if (typeof v === "undefined") return "__undefined__";
    return v;
  };
  return JSON.stringify(S(x));
};

async function sha256Hex(s: string) {
  const h = await crypto.subtle.digest("SHA-256", enc.encode(s));
  return Array.from(new Uint8Array(h))
    .map((b) => b.toString(16).padStart(2, "0"))
    .join("");
}

function now() {
  return Date.now();
}

export function memoize<
  F extends (...args: any[]) => any,
  V = Awaited<ReturnType<F>>,
>(fn: F, opts: MemoOpts<V> = {}): F {
  const ttl = opts.ttlMs ?? 0;
  const max = opts.maxEntries ?? 0;
  const dir = opts.persistDir ? path.resolve(opts.persistDir) : null;

  const mem = new Map<string, Entry<V>>();
  const inflight = new Map<string, Promise<V>>();

  async function keyOf(args: any[]): Promise<string> {
    const base = opts.keyFn ? opts.keyFn(...args) : stableStringify(args);
    return dir ? await sha256Hex(base) : base; // hash when persisting (safe filename)
  }

  async function readDisk(k: string): Promise<Entry<V> | null> {
    if (!dir) throw new Error("no dir!");
    const f = Bun.file(path.join(dir, `${k}.json`));
    if (!(await f.exists())) return null;
    try {
      const obj = JSON.parse(await f.text());
      return obj as Entry<V>;
    } catch {
      return null;
    }
  }

  async function writeDisk(k: string, e: Entry<V>) {
    if (!dir) throw new Error("no dir!");
    await Bun.write(path.join(dir, `${k}.json`), JSON.stringify(e));
  }

  function evictLRU() {
    if (!max || mem.size <= max) return;
    const arr = [...mem.entries()].sort((a, b) => a[1].at - b[1].at);
    for (let i = 0; i < mem.size - max; i++) mem.delete(arr[i][0]);
  }

  async function getOrCall(args: any[]): Promise<V> {
    const k = await keyOf(args);
    const t = now();

    // in-flight coalescing
    if (inflight.has(k)) return inflight.get(k)!;

    // memory hit
    const m = mem.get(k);
    if (m && (!m.exp || t < m.exp)) {
      m.at = t;
      return m.v;
    }

    // disk hit
    const d = await readDisk(k);
    if (d && (!d.exp || t < d.exp)) {
      d.at = t;
      mem.set(k, d);
      evictLRU();
      return d.v;
    }

    // miss → call underlying
    const call = (async () => {
      try {
        const r = fn.apply(undefined, args);
        const v: V = r instanceof Promise ? await r : (r as V);
        const e: Entry<V> = { v, exp: ttl ? t + ttl : null, at: t };
        mem.set(k, e);
        evictLRU();
        await writeDisk(k, e);
        return v;
      } catch (err) {
        if (opts.cacheErrors) {
          const e: Entry<any> = { v: err, exp: ttl ? t + ttl : null, at: t };
          mem.set(k, e);
          await writeDisk(k, e as Entry<V>);
        }
        throw err;
      } finally {
        inflight.delete(k);
      }
    })();

    inflight.set(k, call);
    return call;
  }

  // Wrap preserving arity & `this` for methods
  const wrapped = function (this: any, ...args: any[]) {
    const maybe = getOrCall(args).then((v) => v);
    // If original fn is sync (per your signature), unwrap to sync only when it's truly sync.
    // We can detect by calling without awaiting once—dangerous—so be conservative:
    return maybe as unknown as ReturnType<F>;
  } as any as F;

  return wrapped;
}