@m6d/cortex-server 1.8.0 → 2.0.1

package/src/lib/adapters/database/postgres/attachments.ts

@@ -0,0 +1,88 @@
+// fallow-ignore-file code-duplication -- mirrors ../mssql/attachments.ts by design;
+// Drizzle types query builders per dialect. See DatabaseAdapter in ../index.ts.
+import { and, asc, count, eq, isNull, or, sql } from "drizzle-orm";
+import { attachments, threads } from "@/db/schema.pg";
+import type { DatabaseAdapter } from "@/adapters/database/index";
+import type { PostgresDb } from "./client";
+
+export function createAttachmentsRepository(db: PostgresDb) {
+    return {
+        async createPending(row, limit) {
+            return await db.transaction(async function (tx) {
+                await tx.execute(
+                    sql`SELECT 1 FROM ${threads} WHERE ${threads.id} = ${row.threadId} FOR UPDATE`,
+                );
+                const pending = await tx
+                    .select({ value: count() })
+                    .from(attachments)
+                    .where(
+                        and(
+                            eq(attachments.threadId, row.threadId),
+                            eq(attachments.userId, row.userId),
+                            isNull(attachments.messageId),
+                        ),
+                    )
+                    .execute()
+                    .then((rows) => rows[0]!.value);
+                if (pending >= limit) return undefined;
+
+                return await tx
+                    .insert(attachments)
+                    .values(row)
+                    .returning()
+                    .then((rows) => rows[0]!);
+            });
+        },
+
+        async getById(id, userId) {
+            return await db
+                .select()
+                .from(attachments)
+                .where(and(eq(attachments.id, id), eq(attachments.userId, userId)))
+                .limit(1)
+                .execute()
+                .then((rows) => rows[0]);
+        },
+
+        async listByThread(threadId, userId) {
+            return await db
+                .select()
+                .from(attachments)
+                .where(and(eq(attachments.threadId, threadId), eq(attachments.userId, userId)))
+                .orderBy(asc(attachments.createdAt), asc(attachments.id))
+                .execute();
+        },
+
+        async claimForMessage(threadId, userId, messageId) {
+            // ponytail: retries only this message's pending rows; add claim timestamps for cross-message recovery.
+            return await db
+                .update(attachments)
+                .set({ messageId })
+                .where(
+                    and(
+                        eq(attachments.threadId, threadId),
+                        eq(attachments.userId, userId),
+                        eq(attachments.status, "pending"),
+                        or(isNull(attachments.messageId), eq(attachments.messageId, messageId)),
+                    ),
+                )
+                .returning();
+        },
+
+        async setDescription(id, userId, status, description) {
+            await db
+                .update(attachments)
+                .set(description === undefined ? { status } : { status, description })
+                .where(and(eq(attachments.id, id), eq(attachments.userId, userId)))
+                .execute();
+        },
+
+        async remove(id, userId) {
+            return await db
+                .delete(attachments)
+                .where(and(eq(attachments.id, id), eq(attachments.userId, userId)))
+                .returning()
+                .then((rows) => rows[0]);
+        },
+    } satisfies DatabaseAdapter["attachments"];
+}

package/src/lib/adapters/database/postgres/client.ts

@@ -0,0 +1,11 @@
+import { drizzle } from "drizzle-orm/node-postgres";
+
+/**
+ * Global snake_case mapping means schema columns are never spelled out twice —
+ * the TypeScript name is the source of truth and Drizzle derives the column.
+ */
+export function createPostgresDb(connectionString: string) {
+    return drizzle(connectionString, { casing: "snake_case" });
+}
+
+export type PostgresDb = ReturnType<typeof createPostgresDb>;

package/src/lib/adapters/database/postgres/index.ts

@@ -0,0 +1,23 @@
+import type { DatabaseAdapter } from "@/adapters/database/index";
+import type { StorageAdapter } from "@/adapters/storage/index";
+import { createPostgresDb } from "./client";
+import { createThreadsRepository } from "./threads";
+import { createMessagesRepository } from "./messages";
+import { createLlmRequestsRepository } from "./llm-requests";
+import { createAttachmentsRepository } from "./attachments";
+
+/**
+ * The Postgres implementation of the database contract. One repository per
+ * table; only threads takes storage, because deleting a thread is the one
+ * operation that also has to reach into object storage.
+ */
+export function createPostgresAdapter(connectionString: string, storage?: StorageAdapter) {
+    const db = createPostgresDb(connectionString);
+
+    return {
+        threads: createThreadsRepository(db, storage),
+        messages: createMessagesRepository(db),
+        llmRequests: createLlmRequestsRepository(db),
+        attachments: createAttachmentsRepository(db),
+    } satisfies DatabaseAdapter;
+}

package/src/lib/adapters/database/postgres/llm-requests.ts

@@ -0,0 +1,45 @@
+// fallow-ignore-file code-duplication -- mirrors ../mssql/llm-requests.ts by design;
+// Drizzle types query builders per dialect. See DatabaseAdapter in ../index.ts.
+import { asc, eq, type InferInsertModel } from "drizzle-orm";
+import { llmRequests } from "@/db/schema.pg";
+import type { DatabaseAdapter } from "@/adapters/database/index";
+import type { PostgresDb } from "./client";
+
+export function createLlmRequestsRepository(db: PostgresDb) {
+    return {
+        async insert(requests) {
+            if (!requests.length) return;
+
+            await db
+                .insert(llmRequests)
+                .values(
+                    requests.map(
+                        (r) =>
+                            ({
+                                messageId: r.messageId,
+                                stepNumber: r.stepNumber,
+                                prompt: r.prompt,
+                                output: r.output,
+                                tokenUsage: r.tokenUsage,
+                            }) satisfies InferInsertModel<typeof llmRequests>,
+                    ),
+                )
+                .execute();
+        },
+
+        /** Ordered by step so the inspector replays a turn in the order it ran. */
+        async listByMessageId(messageId) {
+            return await db
+                .select({
+                    id: llmRequests.id,
+                    prompt: llmRequests.prompt,
+                    output: llmRequests.output,
+                    tokenUsage: llmRequests.tokenUsage,
+                })
+                .from(llmRequests)
+                .where(eq(llmRequests.messageId, messageId))
+                .orderBy(asc(llmRequests.stepNumber))
+                .execute();
+        },
+    } satisfies DatabaseAdapter["llmRequests"];
+}

package/src/lib/adapters/database/postgres/messages.ts

@@ -0,0 +1,92 @@
+// fallow-ignore-file code-duplication -- mirrors ../mssql/messages.ts by design;
+// Drizzle types query builders per dialect. See DatabaseAdapter in ../index.ts.
+import { and, desc, eq, getColumns, inArray, type InferInsertModel } from "drizzle-orm";
+import { threads, messages } from "@/db/schema.pg";
+import type { ChatMessage } from "@/types";
+import type { DatabaseAdapter } from "@/adapters/database/index";
+import {
+    mergeStoredContent,
+    textOf,
+    withOwnedAttachments,
+} from "@/adapters/database/message-content";
+import type { PostgresDb } from "./client";
+
+const DEFAULT_MESSAGE_LIMIT = 100;
+
+export function createMessagesRepository(db: PostgresDb) {
+    return {
+        /**
+         * Fetches the newest N and reverses, so a limit keeps the most recent
+         * messages while callers still receive them oldest-first.
+         */
+        async list(userId, threadId, opts) {
+            return await db
+                .select(getColumns(messages))
+                .from(messages)
+                .innerJoin(threads, eq(threads.id, messages.threadId))
+                .where(and(eq(threads.userId, userId), eq(threads.id, threadId)))
+                .orderBy(desc(messages.createdAt), desc(messages.ordinal))
+                .limit(opts?.limit ?? DEFAULT_MESSAGE_LIMIT)
+                .then((x) => x.reverse());
+        },
+
+        /**
+         * Ids are client-generated, so a turn mixes new messages with edits to
+         * ones already stored. Splitting on existence lets the new ones go in as
+         * a single insert; the rest are updated individually.
+         */
+        async upsert(
+            threadId: string,
+            messagesToUpsert: ChatMessage[],
+            options?: { replaceAttachments?: boolean },
+        ) {
+            const incomingMessages = withOwnedAttachments(
+                messagesToUpsert,
+                options?.replaceAttachments,
+            );
+            const existingMessages = await db
+                .select({ id: messages.id, content: messages.content })
+                .from(messages)
+                .where(
+                    inArray(
+                        messages.id,
+                        incomingMessages.map((x) => x.id),
+                    ),
+                )
+                .execute();
+            const existingById = new Map(
+                existingMessages.map((message) => [message.id, message.content]),
+            );
+
+            const newMessages = incomingMessages.filter((x) => !existingById.has(x.id));
+            if (newMessages.length) {
+                await db
+                    .insert(messages)
+                    .values(
+                        newMessages.map(
+                            (x, ordinal) =>
+                                ({
+                                    id: x.id,
+                                    role: x.role,
+                                    threadId,
+                                    content: x,
+                                    text: textOf(x),
+                                    ordinal,
+                                }) satisfies InferInsertModel<typeof messages>,
+                        ),
+                    )
+                    .execute();
+            }
+
+            for (const message of incomingMessages.filter((x) => existingById.has(x.id))) {
+                const content = mergeStoredContent(existingById.get(message.id)!, message);
+
+                await db
+                    .update(messages)
+                    .set({ content, text: textOf(content) })
+                    .where(eq(messages.id, message.id))
+                    .execute();
+            }
+        },
+    } satisfies DatabaseAdapter["messages"];
+}

package/src/lib/adapters/database/postgres/threads.ts

@@ -0,0 +1,84 @@
+// fallow-ignore-file code-duplication -- mirrors ../mssql/threads.ts by design;
+// Drizzle types query builders per dialect. See DatabaseAdapter in ../index.ts.
+import { and, desc, eq } from "drizzle-orm";
+import { attachments, threads } from "@/db/schema.pg";
+import type { Thread } from "@/types";
+import type { ThreadContextMeta } from "@/ai/context/types";
+import type { DatabaseAdapter } from "@/adapters/database/index";
+import type { PostgresDb } from "./client";
+import type { StorageAdapter } from "@/adapters/storage/index";
+
+export function createThreadsRepository(db: PostgresDb, storage?: StorageAdapter) {
+    return {
+        async list(userId, agentId) {
+            return await db
+                .select()
+                .from(threads)
+                .where(and(eq(threads.userId, userId), eq(threads.agentId, agentId)))
+                .orderBy(desc(threads.updatedAt), desc(threads.createdAt));
+        },
+
+        async getById(userId, threadId) {
+            const result = await db
+                .select()
+                .from(threads)
+                .where(and(eq(threads.id, threadId), eq(threads.userId, userId)))
+                .limit(1)
+                .execute();
+            return result.length ? (result[0] as Thread) : null;
+        },
+
+        async create(userId, agentId) {
+            const result = await db.insert(threads).values({ userId, agentId }).returning();
+            return result[0] as Thread;
+        },
+
+        /** Object storage has no foreign keys, so attachment paths must be swept before row cascade. */
+        async delete(userId, threadId) {
+            const attachmentPaths = await db
+                .select({ storageKey: attachments.storageKey })
+                .from(attachments)
+                .where(and(eq(attachments.threadId, threadId), eq(attachments.userId, userId)))
+                .execute()
+                .then((rows) => rows.map((row) => row.storageKey));
+
+            if (storage && attachmentPaths.length) {
+                await storage.delete(attachmentPaths);
+            }
+
+            // attachments.message_id has no cascade, so those rows go first.
+            await db.delete(attachments).where(eq(attachments.threadId, threadId)).execute();
+
+            await db
+                .delete(threads)
+                .where(and(eq(threads.id, threadId), eq(threads.userId, userId)))
+                .execute();
+        },
+
+        async touch(threadId) {
+            const result = await db
+                .update(threads)
+                .set({ updatedAt: new Date() })
+                .where(eq(threads.id, threadId))
+                .returning();
+
+            return result[0] as Thread;
+        },
+
+        async updateTitle(threadId, title) {
+            await db.update(threads).set({ title }).where(eq(threads.id, threadId)).execute();
+        },
+
+        async updateSession(threadId, session) {
+            await db.update(threads).set({ session }).where(eq(threads.id, threadId)).execute();
+        },
+
+        async updateContextMeta(threadId, meta: ThreadContextMeta) {
+            await db
+                .update(threads)
+                .set({ contextMeta: meta })
+                .where(eq(threads.id, threadId))
+                .execute();
+        },
+    } satisfies DatabaseAdapter["threads"];
+}

package/{dist/src/lib/adapters/storage.d.ts → src/lib/adapters/storage/index.ts}

@@ -2,5 +2,5 @@ export type StorageAdapter = {
     put(path: string, data: string): Promise<boolean>;
     get(path: string): Promise<string | false>;
     delete(paths: string[]): Promise<boolean>;
-    stream(path: string): Promise<ReadableStream>;
+    stream(path: string): Promise<ReadableStream<Uint8Array>>;
 };

package/src/lib/adapters/{minio.ts → storage/minio.ts}

@@ -1,6 +1,6 @@
 import * as Minio from "minio";
 import { Readable } from "node:stream";
-import type { StorageAdapter } from "./storage";
+import type { StorageAdapter } from "./index";
 
 export type MinioConfig = {
     endPoint: string;
@@ -53,11 +53,11 @@ export function createMinioAdapter(config: MinioConfig) {
                 return await new Promise<string>((resolve, reject) => {
                     const chunks: Uint8Array[] = [];
 
-                    stream.on("error", () => {
-                        reject();
+                    stream.on("error", (error: Error) => {
+                        reject(error);
                     });
 
-                    stream.on("data", (chunk) => chunks.push(chunk));
+                    stream.on("data", (chunk: Buffer) => chunks.push(chunk));
                     stream.on("end", () => {
                         const buffer = Buffer.concat(chunks);
                         resolve(buffer.toString("base64"));
@@ -82,7 +82,7 @@ export function createMinioAdapter(config: MinioConfig) {
 
         async stream(path: string) {
             const nodeStream = await client.getObject(bucketName, path);
-            return Readable.toWeb(nodeStream) as unknown as ReadableStream;
+            return Readable.toWeb(nodeStream) as ReadableStream<Uint8Array>;
         },
     };
 

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

@@ -0,0 +1,157 @@
+// fallow-ignore-file duplicate-export -- coordinator-facing wrapper intentionally shares redis-stream-log's helper name
+import { memoryStream, type StreamDurability } from "@tanstack/ai";
+
+/**
+ * The live turn on a thread. Generation runs detached from the HTTP request that
+ * started it, so this registry — not the request — owns the run's lifetime: it is
+ * how abort reaches a producer and how a reconnecting client finds the run to
+ * rejoin.
+ */
+export type ActiveRun = {
+    runId: string;
+    abortController: AbortController;
+};
+
+/**
+ * One backend owns both halves of a live run: the thread→run claim registry and
+ * the durability log the run streams through. In-memory is the single-process
+ * default; `createCortex({ redis })` swaps in the Redis coordinator so claims,
+ * aborts, and logs survive a producer restart and reach every instance.
+ */
+export type RunCoordinator = {
+    /** Claims the thread for a new run, aborting whatever was running on it. */
+    startRun: (threadId: string) => Promise<ActiveRun>;
+    getRunId: (threadId: string) => Promise<string | undefined>;
+    /** Which of `threadIds` have a live run — one round trip, for list views. */
+    runningThreadIds: (threadIds: string[]) => Promise<Set<string>>;
+    /** Aborts the thread's live run. False when there was none to abort — the
+     *  caller must not report an abort that never happened. */
+    abortRun: (threadId: string) => Promise<boolean>;
+    /** Releases the thread, but only for the run that still holds it — a
+     *  superseded run finishing late must not evict its replacement. */
+    endRun: (threadId: string, runId: string) => Promise<void>;
+    /** The producer's sink: `resumeFrom()` is null, so appends flow into the log. */
+    producerLog: (runId: string) => StreamDurability;
+    /** A reader on the run's log, replaying from `offset` and tailing until close.
+     *  Takes the thread too: the thread's claim is how a reader can tell a live
+     *  producer from a dead one. */
+    joinLog: (threadId: string, runId: string, offset?: string) => StreamDurability;
+    /**
+     * Whether a client-supplied resume offset belongs to `runId`. The offset is
+     * the adapter's own opaque token and names the log it replays — so an offset
+     * naming another run would replay that run's log to whoever presents it.
+     */
+    offsetBelongsToRun: (offset: string, runId: string) => boolean;
+};
+
+/**
+ * How long a joiner waits for a run's first chunk. The adapter arms this on any
+ * empty log, and it cannot tell "the producer died" from "the producer has not
+ * reached the model yet" — so it has to clear the worst case: a compression
+ * middleware's extra summarization call plus the model's time to first token.
+ * It is generous because it is only a backstop: a producer that fails writes its
+ * own terminal event into the log, and a run that ended is gone from the
+ * registry, so neither ever reaches this timer.
+ */
+export const JOIN_FIRST_CHUNK_TIMEOUT_MS = 5 * 60_000;
+
+/**
+ * `memoryStream` keys its log by the run id it reads off a request, so a
+ * synthetic request is how a detached producer names the log it writes to and
+ * how a joiner names the one it reads from. The URL is never fetched.
+ */
+function logRequest(runId: string, offset?: string) {
+    const query = offset === undefined ? "" : `&offset=${offset}`;
+    return new Request(`http://cortex.invalid/run?runId=${encodeURIComponent(runId)}${query}`);
+}
+
+function createMemoryRuns() {
+    const runs = new Map<string, ActiveRun>();
+
+    function abort(threadId: string) {
+        const run = runs.get(threadId);
+        if (!run) return false;
+        runs.delete(threadId);
+        run.abortController.abort();
+        return true;
+    }
+
+    return {
+        startRun(threadId) {
+            abort(threadId);
+            const run: ActiveRun = {
+                runId: crypto.randomUUID(),
+                abortController: new AbortController(),
+            };
+            runs.set(threadId, run);
+            return Promise.resolve(run);
+        },
+        getRunId(threadId) {
+            return Promise.resolve(runs.get(threadId)?.runId);
+        },
+        runningThreadIds(threadIds) {
+            return Promise.resolve(new Set(threadIds.filter((threadId) => runs.has(threadId))));
+        },
+        abortRun(threadId) {
+            return Promise.resolve(abort(threadId));
+        },
+        endRun(threadId, runId) {
+            if (runs.get(threadId)?.runId === runId) runs.delete(threadId);
+            return Promise.resolve();
+        },
+        producerLog(runId) {
+            return memoryStream(logRequest(runId));
+        },
+        joinLog(_threadId, runId, offset = "-1") {
+            return memoryStream(logRequest(runId, offset), {
+                firstChunkDeadlineMs: JOIN_FIRST_CHUNK_TIMEOUT_MS,
+            });
+        },
+        offsetBelongsToRun(offset, runId) {
+            return offset.startsWith(`memory:v1:${encodeURIComponent(runId)}:`);
+        },
+    } satisfies RunCoordinator;
+}
+
+let coordinator: RunCoordinator = createMemoryRuns();
+
+/** `createCortex` swaps the backend in once at startup; everything below delegates. */
+export function configureRunCoordination(backend: RunCoordinator) {
+    coordinator = backend;
+}
+
+export function startRun(threadId: string) {
+    return coordinator.startRun(threadId);
+}
+
+export function getRunId(threadId: string) {
+    return coordinator.getRunId(threadId);
+}
+
+export async function isRunning(threadId: string) {
+    return (await coordinator.getRunId(threadId)) !== undefined;
+}
+
+export function runningThreadIds(threadIds: string[]) {
+    return coordinator.runningThreadIds(threadIds);
+}
+
+export function abortRun(threadId: string) {
+    return coordinator.abortRun(threadId);
+}
+
+export function endRun(threadId: string, runId: string) {
+    return coordinator.endRun(threadId, runId);
+}
+
+export function producerLog(runId: string) {
+    return coordinator.producerLog(runId);
+}
+
+export function joinLog(threadId: string, runId: string, offset?: string) {
+    return coordinator.joinLog(threadId, runId, offset);
+}
+
+export function offsetBelongsToRun(offset: string, runId: string) {
+    return coordinator.offsetBelongsToRun(offset, runId);
+}