@m6d/cortex-server 1.8.0 → 2.0.1

package/src/lib/ai/redis-runs.ts

@@ -0,0 +1,167 @@
+import type { RedisConnections } from "@/redis";
+import type { ActiveRun, RunCoordinator } from "./active-runs";
+import { JOIN_FIRST_CHUNK_TIMEOUT_MS } from "./active-runs";
+import { offsetBelongsToRun, redisStreamLog } from "./redis-stream-log";
+
+/**
+ * The claim is the producer's liveness signal: a heartbeat refreshes its short
+ * TTL while the run is live, so a claim that expired means the producer died —
+ * that is how joiners on an orphaned run learn to stop tailing, and how a
+ * crashed instance's threads free up within a minute instead of fifteen. The
+ * stream TTL governs how long a run's chunk log stays replayable.
+ */
+const CLAIM_TTL_SECONDS = 60;
+const CLAIM_HEARTBEAT_MS = 20_000;
+const DEFAULT_STREAM_TTL_SECONDS = 24 * 60 * 60;
+
+/**
+ * Cross-instance abort: the producer's AbortController lives in-process on
+ * whichever instance runs the turn, but a stop may land on any instance — so
+ * abort publishes the run id and every instance listens, aborting the run if
+ * it owns the controller. Single-instance works too: the publisher's own
+ * subscriber receives the message.
+ */
+const ABORT_CHANNEL = "cortex:chat:abort";
+
+/**
+ * The claim doubles as mutex and discovery pointer: a joiner that only knows
+ * the thread id finds the live run here.
+ */
+function claimKey(threadId: string) {
+    return `cortex:chat:active:${threadId}`;
+}
+
+/**
+ * Atomically take the claim and learn who held it. Two concurrent starts on
+ * one thread each see the other exactly once, so the loser is always aborted —
+ * a GET-then-SET pair could interleave and leave both producers running.
+ */
+const CLAIM_SWAP = `local previous = redis.call("get", KEYS[1])
+redis.call("set", KEYS[1], ARGV[1], "EX", ARGV[2])
+return previous`;
+
+/** Release only if still held by this run — a superseded run finishing late
+ *  must not evict its replacement. GET-then-DEL would race; the script cannot. */
+const RELEASE_IF_HELD = `if redis.call("get", KEYS[1]) == ARGV[1] then return redis.call("del", KEYS[1]) else return 0 end`;
+
+/** The heartbeat must never resurrect a claim that abort or supersede released. */
+const REFRESH_IF_HELD = `if redis.call("get", KEYS[1]) == ARGV[1] then return redis.call("expire", KEYS[1], ARGV[2]) else return 0 end`;
+
+type RedisRunOptions = {
+    streamTtlSeconds?: number;
+};
+
+export function createRedisRuns(connections: RedisConnections, options: RedisRunOptions = {}) {
+    const { publisher, subscriber } = connections;
+    const streamTtlMs = (options.streamTtlSeconds ?? DEFAULT_STREAM_TTL_SECONDS) * 1000;
+    const logOptions = { streamTtlMs, firstChunkDeadlineMs: JOIN_FIRST_CHUNK_TIMEOUT_MS };
+    const aborters = new Map<string, AbortController>();
+    const heartbeats = new Map<string, ReturnType<typeof setInterval>>();
+
+    subscriber.on("message", function (channel: string, runId: string) {
+        if (channel === ABORT_CHANNEL) aborters.get(runId)?.abort();
+    });
+
+    /**
+     * Lazy and retried, never memoized-rejected: caching one boot-time
+     * subscribe promise would leave every future startRun rethrowing a single
+     * transient outage forever. Once it resolves, ioredis re-subscribes by
+     * itself across reconnects.
+     */
+    let abortSubscription: Promise<unknown> | undefined;
+    function ensureAbortSubscription() {
+        abortSubscription ??= subscriber.subscribe(ABORT_CHANNEL).catch((error: unknown) => {
+            abortSubscription = undefined;
+            throw error;
+        });
+        return abortSubscription;
+    }
+
+    async function abortRun(threadId: string) {
+        const runId = await publisher.get(claimKey(threadId));
+        if (!runId) return false;
+        // The local abort is immediate; the publish reaches the other instances.
+        aborters.get(runId)?.abort();
+        await publisher.publish(ABORT_CHANNEL, runId);
+        await publisher.eval(RELEASE_IF_HELD, 1, claimKey(threadId), runId);
+        return true;
+    }
+
+    function startHeartbeat(threadId: string, runId: string) {
+        const timer = setInterval(function () {
+            void publisher
+                .eval(REFRESH_IF_HELD, 1, claimKey(threadId), runId, String(CLAIM_TTL_SECONDS))
+                .catch((error: unknown) =>
+                    console.error("[cortex-server] redis claim heartbeat:", error),
+                );
+        }, CLAIM_HEARTBEAT_MS);
+        heartbeats.set(runId, timer);
+    }
+
+    function stopHeartbeat(runId: string) {
+        const timer = heartbeats.get(runId);
+        if (timer !== undefined) clearInterval(timer);
+        heartbeats.delete(runId);
+    }
+
+    return {
+        async startRun(threadId) {
+            // A producer must be able to hear aborts before it may exist, or a
+            // stop from another instance could slip by while we boot.
+            await ensureAbortSubscription();
+            const run: ActiveRun = {
+                runId: crypto.randomUUID(),
+                abortController: new AbortController(),
+            };
+            // Registered before the claim exists, so an abort cannot land in
+            // between; unregistered again if the claim never materializes.
+            aborters.set(run.runId, run.abortController);
+            try {
+                const superseded = (await publisher.eval(
+                    CLAIM_SWAP,
+                    1,
+                    claimKey(threadId),
+                    run.runId,
+                    String(CLAIM_TTL_SECONDS),
+                )) as string | null;
+                if (superseded !== null) {
+                    aborters.get(superseded)?.abort();
+                    await publisher.publish(ABORT_CHANNEL, superseded);
+                }
+            } catch (error) {
+                aborters.delete(run.runId);
+                throw error;
+            }
+            startHeartbeat(threadId, run.runId);
+            return run;
+        },
+        async getRunId(threadId) {
+            return (await publisher.get(claimKey(threadId))) ?? undefined;
+        },
+        async runningThreadIds(threadIds) {
+            if (threadIds.length === 0) return new Set<string>();
+            const claims = await publisher.mget(threadIds.map(claimKey));
+            return new Set(threadIds.filter((_, index) => claims[index] !== null));
+        },
+        abortRun,
+        async endRun(threadId, runId) {
+            stopHeartbeat(runId);
+            aborters.delete(runId);
+            await publisher.eval(RELEASE_IF_HELD, 1, claimKey(threadId), runId);
+        },
+        producerLog(runId) {
+            return redisStreamLog(publisher, runId, { resumeOffset: null, ...logOptions });
+        },
+        joinLog(threadId, runId, offset = "-1") {
+            return redisStreamLog(publisher, runId, {
+                resumeOffset: offset,
+                // The claim is the producer's pulse: while this run holds it the
+                // reader may park between chunks indefinitely, and once it is
+                // gone with the log still unclosed, the producer died.
+                isLive: async () => (await publisher.get(claimKey(threadId))) === runId,
+                ...logOptions,
+            });
+        },
+        offsetBelongsToRun,
+    } satisfies RunCoordinator;
+}

package/src/lib/ai/redis-stream-log.ts

@@ -0,0 +1,219 @@
+import type { Redis } from "ioredis";
+import type { StreamChunk, StreamDurability } from "@tanstack/ai";
+
+/**
+ * A `StreamDurability` adapter over a Redis Stream, one stream per run. The
+ * entry id is naturally ordered, so — wrapped in a run-scoped prefix — it is
+ * the resumable offset. Redis holds the actual chunk log rather than acting as
+ * a pub/sub rendezvous, so resume survives a producer restart and works from
+ * any instance.
+ */
+
+const OFFSET_PREFIX = "redis:v1:";
+
+/** How long a blocked read waits before re-checking its deadline and abort signal. */
+const READ_BLOCK_MS = 5_000;
+
+function streamKey(runId: string) {
+    return `cortex:stream:${runId}`;
+}
+
+function runPrefix(runId: string) {
+    return `${OFFSET_PREFIX}${encodeURIComponent(runId)}:`;
+}
+
+/**
+ * Whether a client-supplied resume offset belongs to `runId`. The offset names
+ * the log it replays, so an offset naming another run would replay that run's
+ * log to whoever presents it — matching the prefix this adapter encodes is what
+ * binds a resume to the thread the caller owns.
+ */
+export function offsetBelongsToRun(offset: string, runId: string) {
+    return offset.startsWith(runPrefix(runId));
+}
+
+function encodeOffset(runId: string, entryId: string) {
+    return `${runPrefix(runId)}${entryId}`;
+}
+
+function decodeOffset(offset: string, runId: string) {
+    const entryId = offset.slice(runPrefix(runId).length);
+    if (!offsetBelongsToRun(offset, runId) || !/^\d+-\d+$/.test(entryId)) {
+        throw new Error(`Invalid redis stream offset for run ${JSON.stringify(runId)}: ${offset}`);
+    }
+    return entryId;
+}
+
+/** An entry is either `chunk <json>` or the `end` marker `close()` writes. */
+function fieldValue(fields: string[], name: string) {
+    for (let index = 0; index < fields.length; index += 2) {
+        if (fields[index] === name) return fields[index + 1];
+    }
+    return undefined;
+}
+
+/**
+ * Where a read starts: from the beginning, or the entry id a concrete offset
+ * names. The contract's "now" sentinel never reaches this adapter — cortex's
+ * join paths always pass "-1" or a validated concrete offset — so it falls
+ * through to decodeOffset's rejection instead of carrying dead code.
+ */
+function resolveStartId(runId: string, offset: string) {
+    return offset === "-1" ? "0" : decodeOffset(offset, runId);
+}
+
+const READ_ABORTED = Symbol("read aborted");
+
+/** What an idle (null) reply means: keep parking, or the run will never end. */
+function idlePolicy(runId: string, options: RedisStreamLogOptions) {
+    const deadline = Date.now() + options.firstChunkDeadlineMs;
+    let producerGoneTicks = 0;
+    return {
+        sawData() {
+            producerGoneTicks = 0;
+        },
+        async onIdle(logHasData: boolean) {
+            if (options.isLive && !(await options.isLive())) {
+                // One more blocking read before failing: an abort releases the
+                // claim moments before the terminal marker lands, and that
+                // close must win the race.
+                producerGoneTicks += 1;
+                if (producerGoneTicks > 1) {
+                    throw new Error(
+                        `Redis stream run was abandoned by its producer: ${JSON.stringify(runId)}`,
+                    );
+                }
+                return;
+            }
+            if (!logHasData && Date.now() >= deadline) {
+                throw new Error(
+                    `Redis stream run produced no data within ${options.firstChunkDeadlineMs}ms: ${JSON.stringify(runId)}`,
+                );
+            }
+        },
+    };
+}
+
+async function blockingRead(reader: Redis, key: string, lastId: string, signal?: AbortSignal) {
+    try {
+        return await reader.xread("BLOCK", READ_BLOCK_MS, "STREAMS", key, lastId);
+    } catch (error) {
+        // disconnect() rejects the in-flight read; that is the abort path.
+        if (signal?.aborted) return READ_ABORTED;
+        throw error;
+    }
+}
+
+function parseBatch(runId: string, entries: [id: string, fields: string[]][], lastId: string) {
+    const items: { offset: string; chunk: StreamChunk }[] = [];
+    let closed = false;
+    for (const [entryId, fields] of entries) {
+        lastId = entryId;
+        if (fieldValue(fields, "end") !== undefined) {
+            closed = true;
+            break;
+        }
+        const chunk = fieldValue(fields, "chunk");
+        if (chunk !== undefined) {
+            items.push({
+                offset: encodeOffset(runId, entryId),
+                chunk: JSON.parse(chunk) as StreamChunk,
+            });
+        }
+    }
+    return { items, closed, lastId };
+}
+
+type RedisStreamLogOptions = {
+    /** `null` builds a producer (appends flow into the log); an offset builds a reader. */
+    resumeOffset: string | null;
+    /** How long the log stays replayable after its last append or `close()`. */
+    streamTtlMs: number;
+    /** How long a reader waits for the run's first entry before failing. */
+    firstChunkDeadlineMs: number;
+    /**
+     * Whether the run's producer is still alive. While it is, a caught-up
+     * reader parks indefinitely between chunks — the producer owns termination.
+     * Once it reports dead with the log still unclosed, the run was orphaned
+     * (producer crashed before `close()`), and the read fails so the joiner
+     * gets a terminal error instead of tailing forever.
+     */
+    isLive?: () => Promise<boolean>;
+};
+
+export function redisStreamLog(redis: Redis, runId: string, options: RedisStreamLogOptions) {
+    const key = streamKey(runId);
+
+    return {
+        resumeFrom: () => options.resumeOffset,
+
+        append: async (chunks) => {
+            if (chunks.length === 0) return [];
+            const pipeline = redis.pipeline();
+            for (const chunk of chunks) pipeline.xadd(key, "*", "chunk", JSON.stringify(chunk));
+            // Refreshed on every append so a crashed producer's log — which
+            // never reaches close() — still expires instead of leaking forever.
+            pipeline.pexpire(key, options.streamTtlMs);
+            const results = (await pipeline.exec()) ?? [];
+            return results.slice(0, chunks.length).map(function ([error, entryId]) {
+                if (error) throw error;
+                return encodeOffset(runId, entryId as string);
+            });
+        },
+
+        snapshot: async () => {
+            // XRANGE on a missing key returns [], satisfying the contract's
+            // empty-run case without the existence probe read() needs.
+            const entries = await redis.xrange(key, "-", "+");
+            return parseBatch(runId, entries, "0").items;
+        },
+
+        close: async () => {
+            // The marker is what ends a tailing read — a terminal chunk cannot,
+            // since the agent loop emits one per iteration. The TTL then bounds
+            // how long the finished log lingers for late replays.
+            await redis.xadd(key, "*", "end", "1");
+            await redis.pexpire(key, options.streamTtlMs);
+        },
+
+        read: async function* (offset: string, signal?: AbortSignal) {
+            let lastId = resolveStartId(runId, offset);
+            // Armed while the log is empty, matching memoryStream: once a run
+            // has produced anything the producer owns termination, and a
+            // caught-up reader may legitimately park between chunks — a slow
+            // tool call must not read as "no data".
+            let logHasData = (await redis.exists(key)) === 1;
+            if (offset !== "-1" && !logHasData) {
+                // A concrete offset names entries that existed, so a missing
+                // key means the log expired — fail rather than park forever.
+                throw new Error(`Unknown or expired redis stream run: ${JSON.stringify(runId)}`);
+            }
+
+            // A blocked XREAD stalls every command queued behind it, so each
+            // reader tails on a connection of its own for as long as it lives.
+            const reader = redis.duplicate();
+            const onAbort = () => reader.disconnect();
+            signal?.addEventListener("abort", onAbort, { once: true });
+            const idle = idlePolicy(runId, options);
+            try {
+                while (!signal?.aborted) {
+                    const reply = await blockingRead(reader, key, lastId, signal);
+                    if (reply === READ_ABORTED) return;
+                    if (reply === null) {
+                        await idle.onIdle(logHasData);
+                        continue;
+                    }
+                    idle.sawData();
+                    logHasData = true;
+                    const batch = parseBatch(runId, reply[0]?.[1] ?? [], lastId);
+                    lastId = batch.lastId;
+                    for (const item of batch.items) yield item;
+                    if (batch.closed) return;
+                }
+            } finally {
+                signal?.removeEventListener("abort", onAbort);
+                reader.disconnect();
+            }
+        },
+    } satisfies StreamDurability;
+}