@cogitator-ai/openai-compat 21.2.13 → 21.3.0

package/dist/client/openai-adapter.js

@@ -8,8 +8,9 @@ import { CogitatorError, } from '@cogitator-ai/types';
 import { EventEmitter } from 'events';
 import { nanoid } from 'nanoid';
 import { z } from 'zod';
-import { ThreadManager } from './thread-manager.js';
+import { ThreadManager, assistantMessage, } from './thread-manager.js';
 import { InvalidRequestError } from './errors.js';
+import { RunStepRecorder, isSettledStep, settledEvent } from './run-steps.js';
 export const COGITATOR_MODEL_ID = 'cogitator';
 const ACTIVE_STATUSES = [
     'queued',
@@ -20,6 +21,15 @@ const ACTIVE_STATUSES = [
 const RUN_TTL_SECONDS = 600;
 const CHARS_PER_TOKEN = 4;
 const TRANSCRIPT_HEADER = 'Conversation so far:\n\n';
+function newRunMessage() {
+    return {
+        id: `msg_${nanoid()}`,
+        createdAt: nowSeconds(),
+        announced: false,
+        written: false,
+        text: '',
+    };
+}
 class RunTerminatedError extends Error {
     status;
     constructor(status) {
@@ -40,6 +50,14 @@ function nowSeconds() {
 function isActive(status) {
     return ACTIVE_STATUSES.includes(status);
 }
+/**
+ * Whether `requester` may use a thread or run of `owner`. Everything is reachable without a
+ * requester (no API keys configured, or a direct adapter call), and a thread stored before key
+ * scoping has no owner, so every key reaches it.
+ */
+function canAccess(owner, requester) {
+    return owner === undefined || requester === undefined || owner === requester;
+}
 /** The agent response format and extra instructions for an OpenAI `response_format`. */
 export function toAgentResponseFormat(format) {
     if (!format || format === 'auto' || format.type === 'text')
@@ -201,45 +219,70 @@ export class OpenAIAdapter {
             response_format: stored.response_format,
         };
     }
-    async createThread(metadata) {
-        return this.threadManager.createThread(metadata);
+    /**
+     * Create a thread. With an `owner` (the API key that created it, see the server's `apiKeys`),
+     * only that owner reaches the thread, its messages and its runs.
+     *
+     * The thread, message and run methods below take the `owner` of the caller last: a thread or
+     * run of another owner is treated as missing. Without one, nothing is scoped.
+     */
+    async createThread(metadata, owner) {
+        return this.threadManager.createThread(metadata, owner);
     }
-    async getThread(id) {
-        return this.threadManager.getThread(id);
+    async getThread(id, owner) {
+        return (await this.accessibleThread(id, owner))?.thread;
     }
-    async updateThread(id, updates) {
+    async updateThread(id, updates, owner) {
+        if (!(await this.accessibleThread(id, owner)))
+            return undefined;
         return this.threadManager.updateThread(id, updates);
     }
-    async deleteThread(id) {
+    async deleteThread(id, owner) {
+        if (!(await this.accessibleThread(id, owner)))
+            return false;
         return this.threadManager.deleteThread(id);
     }
-    async addMessage(threadId, params) {
+    async addMessage(threadId, params, owner) {
+        if (!(await this.accessibleThread(threadId, owner)))
+            return undefined;
         return this.threadManager.addMessage(threadId, params);
     }
-    async getMessage(threadId, messageId) {
-        return this.threadManager.getMessage(threadId, messageId);
+    async getMessage(threadId, messageId, owner) {
+        const stored = await this.accessibleThread(threadId, owner);
+        return stored?.messages.find((message) => message.id === messageId);
     }
-    async listMessages(threadId, options) {
+    async listMessages(threadId, options, owner) {
+        if (!(await this.accessibleThread(threadId, owner)))
+            return [];
         return this.threadManager.listMessages(threadId, options);
     }
+    async accessibleThread(id, owner) {
+        const stored = await this.threadManager.getStoredThread(id);
+        return stored && canAccess(stored.owner, owner) ? stored : undefined;
+    }
+    accessibleRun(threadId, runId, owner) {
+        const state = this.runs.get(runId);
+        return state?.run.thread_id === threadId && canAccess(state.owner, owner) ? state : undefined;
+    }
     /**
      * Create a run and start executing it in the background.
      *
-     * Rejects when the assistant or thread does not exist, or when the thread
-     * already has an active run.
+     * Rejects when the assistant or thread does not exist (or the thread belongs to another
+     * `owner`), or when the thread already has an active run.
      */
-    async createRun(threadId, request) {
+    async createRun(threadId, request, owner) {
         const assistant = await this.threadManager.getAssistant(request.assistant_id);
         if (!assistant) {
             throw new InvalidRequestError(`Assistant ${request.assistant_id} not found`, 'assistant_id');
         }
-        const thread = await this.threadManager.getThread(threadId);
+        const thread = await this.accessibleThread(threadId, owner);
         if (!thread) {
             throw new InvalidRequestError(`Thread ${threadId} not found`);
         }
         const activeRunId = this.activeRunByThread.get(threadId);
         if (activeRunId) {
-            throw new InvalidRequestError(`Thread ${threadId} already has an active run ${activeRunId}.`);
+            const visible = this.accessibleRun(threadId, activeRunId, owner) !== undefined;
+            throw new InvalidRequestError(`Thread ${threadId} already has an active run${visible ? ` ${activeRunId}` : ''}.`);
         }
         if (request.max_prompt_tokens !== undefined && !isPositiveInteger(request.max_prompt_tokens)) {
             throw new InvalidRequestError('max_prompt_tokens must be a positive integer', 'max_prompt_tokens');
@@ -277,11 +320,16 @@ export class OpenAIAdapter {
         };
         const state = {
             run,
+            ...(owner !== undefined && { owner }),
+            steps: new RunStepRecorder(run, (event, data) => this.emitStep(state, event, data)),
             abortController: new AbortController(),
             emitter: new EventEmitter(),
             events: [],
             ended: false,
             clientOutputs: new Map(),
+            messageWrites: [],
+            outbox: Promise.resolve(),
+            queued: 0,
         };
         this.runs.set(runId, state);
         this.activeRunByThread.set(threadId, runId);
@@ -304,28 +352,35 @@ export class OpenAIAdapter {
     /**
      * Get a run by ID
      */
-    getRun(threadId, runId) {
-        const state = this.runs.get(runId);
-        if (state?.run.thread_id === threadId) {
-            return structuredClone(state.run);
-        }
-        return undefined;
+    getRun(threadId, runId, owner) {
+        const state = this.accessibleRun(threadId, runId, owner);
+        return state ? structuredClone(state.run) : undefined;
     }
     /**
      * List runs of a thread, newest first
      */
-    listRuns(threadId) {
+    listRuns(threadId, owner) {
         return Array.from(this.runs.values())
-            .filter((state) => state.run.thread_id === threadId)
+            .filter((state) => state.run.thread_id === threadId && canAccess(state.owner, owner))
             .map((state) => structuredClone(state.run))
             .reverse();
     }
+    /**
+     * The steps of a run in the order they were created, or undefined when there is no such run.
+     * Steps are kept with their run (see `maxStoredRuns`).
+     */
+    listRunSteps(threadId, runId, owner) {
+        return this.accessibleRun(threadId, runId, owner)?.steps.list();
+    }
+    getRunStep(threadId, runId, stepId, owner) {
+        return this.accessibleRun(threadId, runId, owner)?.steps.get(stepId);
+    }
     /**
      * Cancel a run. Throws when the run already finished.
      */
-    cancelRun(threadId, runId) {
-        const state = this.runs.get(runId);
-        if (state?.run.thread_id !== threadId) {
+    cancelRun(threadId, runId, owner) {
+        const state = this.accessibleRun(threadId, runId, owner);
+        if (!state) {
             return undefined;
         }
         if (!isActive(state.run.status) || state.run.status === 'cancelling') {
@@ -341,9 +396,9 @@ export class OpenAIAdapter {
      * tool call must be provided: the result of a client-side function, or the decision on a
      * server tool that needs approval (see `toolOutputToDecision`).
      */
-    async submitToolOutputs(threadId, runId, request) {
-        const state = this.runs.get(runId);
-        if (state?.run.thread_id !== threadId) {
+    async submitToolOutputs(threadId, runId, request, owner) {
+        const state = this.accessibleRun(threadId, runId, owner);
+        if (!state) {
             return undefined;
         }
         const pending = state.pending;
@@ -378,7 +433,8 @@ export class OpenAIAdapter {
      * stream only events emitted after this point.
      */
     getRunEventCursor(runId) {
-        return this.runs.get(runId)?.events.length ?? 0;
+        const state = this.runs.get(runId);
+        return state ? state.events.length + state.queued : 0;
     }
     /**
      * Stream run events as server-sent-event payloads, starting at `fromIndex`
@@ -425,10 +481,51 @@ export class OpenAIAdapter {
             event,
             data: typeof data === 'string' ? data : structuredClone(data),
         };
+        if (state.queued === 0) {
+            this.deliver(state, payload);
+            return;
+        }
+        this.enqueue(state, async () => this.deliver(state, payload));
+    }
+    /**
+     * Emit a step event. The settle event of a step that waits in the outbox is named by the
+     * status the step has when it goes out, so a message step whose message was not stored after
+     * all goes out failed.
+     */
+    emitStep(state, event, data) {
+        if (state.queued === 0 || !isSettledStep(data)) {
+            this.emit(state, event, data);
+            return;
+        }
+        this.enqueue(state, async () => this.deliver(state, { event: settledEvent(data), data: structuredClone(data) }));
+    }
+    /** Send an event to the run's listeners: one that throws is logged and never stops the run */
+    deliver(state, payload) {
         state.events.push(payload);
-        state.emitter.emit('event', payload.event, payload.data);
+        try {
+            state.emitter.emit('event', payload.event, payload.data);
+        }
+        catch (error) {
+            console.error(`[CogitatorOpenAI] A listener of run ${state.run.id} failed:`, error);
+        }
+    }
+    /** Runs `step` once every event queued before it is out, holding back the events after it */
+    enqueue(state, step) {
+        state.queued++;
+        state.outbox = state.outbox
+            .then(step)
+            .catch((error) => {
+            console.error(`[CogitatorOpenAI] Run ${state.run.id} could not send an event:`, error);
+        })
+            .finally(() => {
+            state.queued--;
+        });
     }
     end(state) {
+        if (state.queued > 0) {
+            void state.outbox.then(() => this.end(state));
+            return;
+        }
         if (state.ended)
             return;
         state.ended = true;
@@ -444,6 +541,7 @@ export class OpenAIAdapter {
             state.expiryTimer = undefined;
         }
     }
+    /** Stops the run with `reason`: a cancel, an expiry, or a store of its messages that failed */
     terminate(state, reason) {
         this.clearExpiry(state);
         const pending = state.pending;
@@ -522,11 +620,16 @@ export class OpenAIAdapter {
     }
     /**
      * Put the run in `requires_action` with `calls` and wait for the client's outputs, or for the
-     * run to be cancelled or to expire (`RUN_TTL_SECONDS`).
+     * run to be cancelled or to expire (`RUN_TTL_SECONDS`). The messages before the pause are
+     * stored and their events sent first, so a client that sees `requires_action` has all of them,
+     * and a store that failed fails the run instead.
      */
-    awaitToolOutputs(state, calls) {
-        if (state.abortController.signal.aborted) {
-            return Promise.reject(new RunTerminatedError('cancelled'));
+    async awaitToolOutputs(state, calls) {
+        await Promise.all(state.messageWrites);
+        await state.outbox;
+        const { signal } = state.abortController;
+        if (signal.aborted) {
+            throw signal.reason instanceof Error ? signal.reason : new RunTerminatedError('cancelled');
         }
         return new Promise((resolve, reject) => {
             state.pending = { calls, resolve, reject };
@@ -604,6 +707,31 @@ export class OpenAIAdapter {
         const { run } = state;
         const threadId = run.thread_id;
         const signal = state.abortController.signal;
+        let message = newRunMessage();
+        const announceMessage = () => {
+            if (message.announced)
+                return;
+            message.announced = true;
+            message.createdAt = nowSeconds();
+            state.steps.messageStarted(message.id);
+            const inProgress = {
+                id: message.id,
+                object: 'thread.message',
+                created_at: message.createdAt,
+                thread_id: threadId,
+                status: 'in_progress',
+                completed_at: null,
+                incomplete_at: null,
+                role: 'assistant',
+                content: [],
+                assistant_id: assistant.id,
+                run_id: run.id,
+                attachments: [],
+                metadata: {},
+            };
+            this.emit(state, 'thread.message.created', inProgress);
+            this.emit(state, 'thread.message.in_progress', inProgress);
+        };
         try {
             run.status = 'in_progress';
             run.started_at = nowSeconds();
@@ -638,31 +766,12 @@ export class OpenAIAdapter {
                 responseFormat: format.responseFormat,
                 tools,
             });
-            const messageId = `msg_${nanoid()}`;
-            let messageAnnounced = false;
-            const announceMessage = () => {
-                if (messageAnnounced)
-                    return;
-                messageAnnounced = true;
-                const inProgress = {
-                    id: messageId,
-                    object: 'thread.message',
-                    created_at: nowSeconds(),
-                    thread_id: threadId,
-                    status: 'in_progress',
-                    completed_at: null,
-                    incomplete_at: null,
-                    role: 'assistant',
-                    content: [],
-                    assistant_id: assistant.id,
-                    run_id: run.id,
-                    attachments: [],
-                    metadata: {},
-                };
-                this.emit(state, 'thread.message.created', inProgress);
-                this.emit(state, 'thread.message.in_progress', inProgress);
-            };
-            let accumulated = '';
+            /**
+             * Text of the latest model turn, as every middleware left it. A turn that writes text and
+             * calls tools becomes its own message, which a run without `stream: true` gets no tokens
+             * of, so it takes the text from here when the first call of the turn arrives.
+             */
+            let turnText;
             const turnOptions = {
                 signal,
                 stream: !!request.stream,
@@ -671,15 +780,35 @@ export class OpenAIAdapter {
                     if (signal.aborted)
                         return;
                     announceMessage();
-                    accumulated += token;
+                    message.text += token;
                     const delta = {
-                        id: messageId,
+                        id: message.id,
                         object: 'thread.message.delta',
                         delta: { content: [{ index: 0, type: 'text', text: { value: token } }] },
                     };
                     this.emit(state, 'thread.message.delta', delta);
                 },
                 onApproval: (approval) => clientToolNames.has(approval.toolName) ? 'pause' : undefined,
+                onModelTurn: (turn) => {
+                    turnText = turn.content;
+                },
+                onToolCall: (call) => {
+                    if (signal.aborted)
+                        return;
+                    if (turnText !== undefined && (message.announced || turnText)) {
+                        announceMessage();
+                        message.text = turnText;
+                    }
+                    turnText = undefined;
+                    if (message.announced && !message.written) {
+                        void this.sealMessage(state, assistant.id, message, message.text);
+                        state.steps.messageFinished(null);
+                        message = newRunMessage();
+                    }
+                    state.steps.toolCalled(call);
+                },
+                onToolResult: (result) => state.steps.toolReturned(result),
+                onSpan: (span) => state.steps.recordSpan(span),
             };
             let result = await this.cogitator.run(agent, {
                 ...turnOptions,
@@ -701,7 +830,14 @@ export class OpenAIAdapter {
             if (signal.aborted) {
                 throw signal.reason instanceof Error ? signal.reason : new RunTerminatedError('cancelled');
             }
-            const finalContent = result.output || accumulated;
+            run.usage = result.usage
+                ? {
+                    prompt_tokens: result.usage.inputTokens,
+                    completion_tokens: result.usage.outputTokens,
+                    total_tokens: result.usage.totalTokens,
+                }
+                : null;
+            const finalContent = result.output || message.text;
             const incompleteReason = result.truncated
                 ? 'max_tokens'
                 : result.blocked === 'content_filter'
@@ -709,18 +845,11 @@ export class OpenAIAdapter {
                     : undefined;
             if (finalContent || incompleteReason) {
                 announceMessage();
-                const message = await this.threadManager.addAssistantMessage(threadId, finalContent, assistant.id, run.id, messageId, incompleteReason);
-                if (message) {
-                    this.emit(state, incompleteReason ? 'thread.message.incomplete' : 'thread.message.completed', message);
-                }
+                void this.sealMessage(state, assistant.id, message, finalContent, incompleteReason);
             }
-            run.usage = result.usage
-                ? {
-                    prompt_tokens: result.usage.inputTokens,
-                    completion_tokens: result.usage.outputTokens,
-                    total_tokens: result.usage.totalTokens,
-                }
-                : null;
+            await Promise.all(state.messageWrites);
+            state.steps.messageFinished(run.usage);
+            state.steps.finish({ status: 'completed' });
             if (result.truncated) {
                 run.status = 'incomplete';
                 run.incomplete_details = { reason: 'max_completion_tokens' };
@@ -733,9 +862,13 @@ export class OpenAIAdapter {
             }
         }
         catch (error) {
+            await Promise.allSettled(state.messageWrites);
             const reason = signal.aborted ? signal.reason : error;
             if (reason instanceof RunTerminatedError) {
-                if (reason.status === 'cancelled') {
+                const cancelled = reason.status === 'cancelled';
+                await this.keepUnfinishedMessage(state, assistant.id, message, cancelled ? 'run_cancelled' : 'run_expired');
+                state.steps.finish({ status: reason.status });
+                if (cancelled) {
                     run.status = 'cancelled';
                     run.cancelled_at = nowSeconds();
                     this.emit(state, 'thread.run.cancelled', run);
@@ -746,14 +879,21 @@ export class OpenAIAdapter {
                 }
             }
             else if (reason instanceof PromptBudgetExceededError) {
+                state.steps.finish({ status: 'completed' });
                 run.status = 'incomplete';
                 run.incomplete_details = { reason: 'max_prompt_tokens' };
                 this.emit(state, 'thread.run.incomplete', run);
             }
             else {
+                const lastError = {
+                    code: 'server_error',
+                    message: this.clientErrorMessage(run, reason),
+                };
+                await this.keepUnfinishedMessage(state, assistant.id, message, 'run_failed');
+                state.steps.finish({ status: 'failed', error: lastError });
                 run.status = 'failed';
                 run.failed_at = nowSeconds();
-                run.last_error = { code: 'server_error', message: this.clientErrorMessage(run, error) };
+                run.last_error = lastError;
                 this.emit(state, 'thread.run.failed', run);
             }
         }
@@ -764,6 +904,59 @@ export class OpenAIAdapter {
             this.end(state);
         }
     }
+    /**
+     * Finish a message of the run: store it and emit its final event once it is stored, with the
+     * events of the next steps held back until then. Stores to one thread keep the order of the
+     * calls, and the run waits for them (`messageWrites`) before it completes. A store that fails,
+     * or finds the thread deleted, stops the run right away, and the message ends as incomplete
+     * (`run_failed`) instead of completed, as it was never kept.
+     */
+    sealMessage(state, assistantId, message, content, incompleteReason) {
+        const threadId = state.run.thread_id;
+        const sealedAs = (reason) => assistantMessage({
+            threadId,
+            content,
+            assistantId,
+            runId: state.run.id,
+            messageId: message.id,
+            incompleteReason: reason,
+            createdAt: message.createdAt,
+        });
+        const sealed = sealedAs(incompleteReason);
+        message.written = true;
+        const write = this.threadManager.appendMessage(threadId, sealed).then((stored) => {
+            if (!stored) {
+                throw new InvalidRequestError(`Thread ${threadId} was deleted while the run was writing to it`);
+            }
+        });
+        this.enqueue(state, () => write.then(() => this.deliver(state, {
+            event: incompleteReason ? 'thread.message.incomplete' : 'thread.message.completed',
+            data: structuredClone(sealed),
+        }), () => {
+            state.steps.messageLost(message.id, {
+                code: 'server_error',
+                message: 'The message could not be stored',
+            });
+            this.deliver(state, { event: 'thread.message.incomplete', data: sealedAs('run_failed') });
+        }));
+        write.catch((error) => this.terminate(state, error instanceof Error ? error : new Error(String(error))));
+        state.messageWrites.push(write);
+        return write;
+    }
+    /**
+     * Keep what the model wrote of a message the run announced but did not finish, as an
+     * incomplete message, so the `message_creation` step points at a message that exists.
+     */
+    async keepUnfinishedMessage(state, assistantId, message, reason) {
+        if (!message.announced || message.written)
+            return;
+        try {
+            await this.sealMessage(state, assistantId, message, message.text, reason);
+        }
+        catch (error) {
+            console.error(`[CogitatorOpenAI] Run ${state.run.id} could not keep its unfinished message:`, error);
+        }
+    }
 }
 /**
  * Create an OpenAI adapter for Cogitator