diff --git a/tegg/core/agent-runtime/src/AgentRuntime.ts b/tegg/core/agent-runtime/src/AgentRuntime.ts index 8b63a484fe..8fa314a467 100644 --- a/tegg/core/agent-runtime/src/AgentRuntime.ts +++ b/tegg/core/agent-runtime/src/AgentRuntime.ts @@ -1,38 +1,91 @@ +import { EventEmitter } from 'node:events'; +import { appendFileSync, createReadStream, mkdirSync, existsSync } from 'node:fs'; +import { tmpdir } from 'node:os'; +import { join } from 'node:path'; +import { createInterface } from 'node:readline'; + import type { CreateRunInput, + CreateThreadOptions, + GetThreadOptions, ThreadObject, ThreadObjectWithMessages, RunObject, - MessageObject, - MessageDeltaObject, - MessageContentBlock, - AgentStreamMessage, + AgentMessage, AgentStore, + StreamEvent, +} from '@eggjs/tegg-types/agent-runtime'; +import { + RunStatus, + AgentObjectType, + AgentConflictError, + AgentInvalidRequestError, + AgentNotFoundError, + AgentTimeoutError, } from '@eggjs/tegg-types/agent-runtime'; -import { RunStatus, AgentSSEEvent, AgentObjectType } from '@eggjs/tegg-types/agent-runtime'; -import { AgentConflictError } from '@eggjs/tegg-types/agent-runtime'; import type { EggLogger } from 'egg-logger'; -import { newMsgId } from './AgentStoreUtils.ts'; import { MessageConverter } from './MessageConverter.ts'; import { RunBuilder } from './RunBuilder.ts'; -import type { RunUsage } from './RunBuilder.ts'; import type { SSEWriter } from './SSEWriter.ts'; +const HEARTBEAT_INTERVAL_MS = 10_000; +const EVENT_DIR = join(tmpdir(), 'agent-runtime-events'); +const DEFAULT_CANCEL_COMMIT_TIMEOUT_MS = 30_000; + +function validateMetadata(value: unknown): Record | undefined { + if (value === undefined) return undefined; + if (value === null || typeof value !== 'object' || Array.isArray(value)) { + throw new AgentInvalidRequestError("'metadata' must be an object"); + } + return value as Record; +} + +interface RunEventBuffer { + filePath: string; + lastSeq: number; + done: boolean; + emitter: EventEmitter; +} + export const AGENT_RUNTIME: unique symbol = Symbol('agentRuntime'); /** - * The executor interface — only requires execRun so the runtime can delegate + * The executor interface — execRun is required so the runtime can delegate * execution back through the controller's prototype chain (AOP/mock friendly). + * + * `isSessionCommitted` is an optional hook that lets the executor tell the + * runtime when its underlying session has been persisted to storage (e.g. the + * Claude Code SDK jsonl file). The runtime uses this to decide when a pending + * `cancelRun` can safely abort and persist the thread. See AgentHandler.ts + * for the semantics and the default heuristic used when this hook is absent. */ export interface AgentExecutor { - execRun(input: CreateRunInput, signal?: AbortSignal): AsyncGenerator; + execRun(input: CreateRunInput, signal?: AbortSignal): AsyncGenerator; + isSessionCommitted?(msg: AgentMessage, history: AgentMessage[]): boolean | Promise; } export interface AgentRuntimeOptions { executor: AgentExecutor; store: AgentStore; logger: EggLogger; + /** + * How long cancelRun should wait for the executor's session to become + * committed before giving up and marking the run as failed. Defaults to + * 30 seconds. + */ + cancelCommitTimeoutMs?: number; +} + +interface RunTaskState { + promise: Promise; + abortController: AbortController; + /** True once the executor has reported (or the heuristic has detected) that + * its session is safely persisted and the run can be cancelled cleanly. */ + committed: boolean; + /** Emits 'commit' the first time committed flips to true, and 'end' when the + * task's execution finally finishes (success, failure, or abort). */ + emitter: EventEmitter; } export class AgentRuntime { @@ -43,10 +96,25 @@ export class AgentRuntime { RunStatus.Expired, ]); + // Statuses that must short-circuit the "write Completed" path in the + // execution loops. Covers a TOCTOU window where another actor (most + // notably cancelRun's commit-timeout watchdog, which writes Failed) sets + // a terminal state while this worker has just exited the for-await loop + // but hasn't yet written rb.complete(usage). Completed is intentionally + // excluded so the normal success path is not routed through here. + private static readonly POST_LOOP_TERMINAL_STATUSES = new Set([ + RunStatus.Cancelling, + RunStatus.Cancelled, + RunStatus.Failed, + RunStatus.Expired, + ]); + private store: AgentStore; - private runningTasks: Map; abortController: AbortController }>; + private runningTasks: Map; + private runBuffers: Map; private executor: AgentExecutor; private logger: EggLogger; + private cancelCommitTimeoutMs: number; constructor(options: AgentRuntimeOptions) { this.executor = options.executor; @@ -55,11 +123,13 @@ export class AgentRuntime { throw new Error('AgentRuntimeOptions.logger is required'); } this.logger = options.logger; + this.cancelCommitTimeoutMs = options.cancelCommitTimeoutMs ?? DEFAULT_CANCEL_COMMIT_TIMEOUT_MS; this.runningTasks = new Map(); + this.runBuffers = new Map(); } - async createThread(): Promise { - const thread = await this.store.createThread(); + async createThread(options?: CreateThreadOptions): Promise { + const thread = await this.store.createThread(options?.metadata); return { id: thread.id, object: AgentObjectType.Thread, @@ -68,8 +138,8 @@ export class AgentRuntime { }; } - async getThread(threadId: string): Promise { - const thread = await this.store.getThread(threadId); + async getThread(threadId: string, options?: GetThreadOptions): Promise { + const thread = await this.store.getThread(threadId, options); return { id: thread.id, object: AgentObjectType.Thread, @@ -79,12 +149,33 @@ export class AgentRuntime { }; } + /** + * Resolve the thread for a run and persist the run's `metadata` onto the + * thread. The same `metadata` is also stored on the run record (see + * {@link AgentStore.createRun}); here it initializes an auto-created thread or + * is shallow-merged into an existing thread's `meta.json`. + */ private async ensureThread(input: CreateRunInput): Promise<{ threadId: string; input: CreateRunInput }> { + const metadata = validateMetadata(input.metadata); if (input.threadId) { - return { threadId: input.threadId, input }; + const thread = await this.store.getThread(input.threadId); + if (metadata && Object.keys(metadata).length > 0) { + if (!this.store.updateThreadMetadata) { + throw new Error('AgentStore does not support updating thread metadata'); + } + // Best-effort: the same metadata is already persisted on the run record, + // so a failure to mirror it onto the thread must not fail run creation. + try { + await this.store.updateThreadMetadata(input.threadId, metadata); + } catch (err) { + this.logger.error('[AgentRuntime] failed to persist metadata onto thread threadId=%s:', input.threadId, err); + } + } + const isResume = thread.messages.length > 0; + return { threadId: input.threadId, input: { ...input, isResume } }; } - const thread = await this.store.createThread(); - return { threadId: thread.id, input: { ...input, threadId: thread.id } }; + const thread = await this.store.createThread(metadata); + return { threadId: thread.id, input: { ...input, threadId: thread.id, isResume: false } }; } async syncRun(input: CreateRunInput, signal?: AbortSignal): Promise { @@ -105,47 +196,79 @@ export class AgentRuntime { } // Register in runningTasks so cancelRun can find and await this run. - // Use a real pending promise (not Promise.resolve()) so cancelRun's - // `await task.promise` blocks until syncRun's try/finally completes. let resolveTask!: () => void; const taskPromise = new Promise((r) => { resolveTask = r; }); - this.runningTasks.set(run.id, { promise: taskPromise, abortController }); - + const task: RunTaskState = { + promise: taskPromise, + abortController, + committed: false, + emitter: new EventEmitter(), + }; + this.runningTasks.set(run.id, task); + + const streamMessages: AgentMessage[] = []; + // Persist a turn's transcript to the thread at most once, even if a later + // store call (e.g. updateRun) throws after a successful append and routes + // us through a catch-block persist. Guarded by task.committed so we never + // write a thread the executor has not persisted to its own session. + let messagesPersisted = false; + const persistPartialOnce = async (): Promise => { + if (!task.committed || messagesPersisted) return; + messagesPersisted = true; + await this.persistPartialMessages(threadId, input, streamMessages); + }; try { await this.store.updateRun(run.id, rb.start()); - const streamMessages: AgentStreamMessage[] = []; for await (const msg of this.executor.execRun(input, abortController.signal)) { if (abortController.signal.aborted) { - // Run was cancelled externally — re-read store for the latest state + await persistPartialOnce(); + await this.finaliseAbortedRun(run.id); const latest = await this.store.getRun(run.id); return RunBuilder.fromRecord(latest).snapshot(); } streamMessages.push(msg); + await this.markCommittedIfNeeded(task, msg, streamMessages); + } + + // TOCTOU: another worker (e.g. cancelRun, or its commit-timeout + // watchdog which writes Failed) may have terminated this run while + // we were finishing the last iterator.next(). Respect the already-set + // terminal state instead of overwriting it with Completed. + const currentRun = await this.store.getRun(run.id); + if (AgentRuntime.POST_LOOP_TERMINAL_STATUSES.has(currentRun.status)) { + await persistPartialOnce(); + await this.finaliseAbortedRun(run.id); + const latest = await this.store.getRun(run.id); + return RunBuilder.fromRecord(latest).snapshot(); } - const { output, usage } = MessageConverter.extractFromStreamMessages(streamMessages, run.id); + const usage = MessageConverter.extractUsage(streamMessages); - // Append messages first so that if updateRun fails the run stays in_progress - // and can be retried, rather than showing completed with missing thread history. - // TODO(atomicity): for full consistency, add an aggregate store method - // (e.g. completeRunWithMessages) that wraps both writes in a single transaction. + // Append input messages + stream messages to thread (excluding stream_event deltas) await this.store.appendMessages(threadId, [ - ...MessageConverter.toInputMessageObjects(input.input.messages, threadId), - ...output, + ...MessageConverter.toAgentMessages(input.input.messages), + ...MessageConverter.filterForStorage(streamMessages), ]); + messagesPersisted = true; - await this.store.updateRun(run.id, rb.complete(output, usage)); + await this.store.updateRun(run.id, rb.complete(usage)); return rb.snapshot(); } catch (err: unknown) { if (abortController.signal.aborted) { - // Cancelled — re-read store for the latest state + await persistPartialOnce(); + await this.finaliseAbortedRun(run.id); const latest = await this.store.getRun(run.id); return RunBuilder.fromRecord(latest).snapshot(); } + // Non-abort failure (e.g. upstream stream terminated mid-turn). Persist + // the partial transcript so the thread history keeps the user turn and + // any committed assistant output instead of silently dropping the run. + // No-op if the success path already appended (avoids duplicate history). + await persistPartialOnce(); try { await this.store.updateRun(run.id, rb.fail(err as Error)); } catch (storeErr) { @@ -153,6 +276,7 @@ export class AgentRuntime { } throw err; } finally { + task.emitter.emit('end'); resolveTask(); this.runningTasks.delete(run.id); } @@ -171,56 +295,81 @@ export class AgentRuntime { const queuedSnapshot = rb.snapshot(); // Register in runningTasks before the IIFE starts executing to avoid a race - // where the IIFE's finally block deletes the entry before it is set. let resolveTask!: () => void; const taskPromise = new Promise((r) => { resolveTask = r; }); - this.runningTasks.set(run.id, { promise: taskPromise, abortController }); + const task: RunTaskState = { + promise: taskPromise, + abortController, + committed: false, + emitter: new EventEmitter(), + }; + this.runningTasks.set(run.id, task); (async () => { + const streamMessages: AgentMessage[] = []; + // Persist a turn's transcript to the thread at most once (see syncRun). + let messagesPersisted = false; + const persistPartialOnce = async (): Promise => { + if (!task.committed || messagesPersisted) return; + messagesPersisted = true; + await this.persistPartialMessages(threadId, input, streamMessages); + }; try { await this.store.updateRun(run.id, rb.start()); - const streamMessages: AgentStreamMessage[] = []; for await (const msg of this.executor.execRun(input, abortController.signal)) { - if (abortController.signal.aborted) return; + if (abortController.signal.aborted) { + await persistPartialOnce(); + await this.finaliseAbortedRun(run.id); + return; + } streamMessages.push(msg); + await this.markCommittedIfNeeded(task, msg, streamMessages); } - // Check if another worker has cancelled this run before writing final state + // TOCTOU: respect any terminal-ish status set by another worker + // (cancelRun, its commit-timeout watchdog which writes Failed, or + // an external expiration) instead of overwriting it with Completed. const currentRun = await this.store.getRun(run.id); - if (currentRun.status === RunStatus.Cancelling || currentRun.status === RunStatus.Cancelled) { + if (AgentRuntime.POST_LOOP_TERMINAL_STATUSES.has(currentRun.status)) { + await persistPartialOnce(); return; } - const { output, usage } = MessageConverter.extractFromStreamMessages(streamMessages, run.id); + const usage = MessageConverter.extractUsage(streamMessages); - // Append messages before marking run as completed — see syncRun comment. - // TODO(atomicity): add aggregate store method for full transactional guarantee. + // Append input messages + stream messages to thread (excluding stream_event deltas) await this.store.appendMessages(threadId, [ - ...MessageConverter.toInputMessageObjects(input.input.messages, threadId), - ...output, + ...MessageConverter.toAgentMessages(input.input.messages), + ...MessageConverter.filterForStorage(streamMessages), ]); + messagesPersisted = true; - await this.store.updateRun(run.id, rb.complete(output, usage)); + await this.store.updateRun(run.id, rb.complete(usage)); } catch (err: unknown) { if (!abortController.signal.aborted) { - // Check store before writing failed state — another worker may have cancelled + // Non-abort failure (e.g. upstream stream terminated mid-turn). + // Persist the partial transcript before marking the run failed so + // the thread history is not silently dropped. No-op if the success + // path already appended (avoids duplicate history). + await persistPartialOnce(); try { const currentRun = await this.store.getRun(run.id); if (currentRun.status !== RunStatus.Cancelling && currentRun.status !== RunStatus.Cancelled) { await this.store.updateRun(run.id, rb.fail(err as Error)); } } catch (storeErr) { - // TODO: need a background expiry mechanism to clean up runs stuck in non-terminal states - // (e.g. in_progress or cancelling) when store writes fail persistently. this.logger.error('[AgentRuntime] failed to update run status after error:', storeErr); } } else { + await persistPartialOnce(); + await this.finaliseAbortedRun(run.id); this.logger.error('[AgentRuntime] execRun error during abort:', err); } } finally { + task.emitter.emit('end'); resolveTask(); this.runningTasks.delete(run.id); } @@ -229,155 +378,393 @@ export class AgentRuntime { return queuedSnapshot; } + /** + * Start a streaming run with background execution. + * The task continues running even if the SSE client disconnects. + * Events are persisted to a JSONL file for reconnection support. + */ async streamRun(input: CreateRunInput, writer: SSEWriter): Promise { - // Abort execRun generator when client disconnects - const abortController = new AbortController(); - writer.onClose(() => abortController.abort()); - const { threadId, input: resolvedInput } = await this.ensureThread(input); input = resolvedInput; const run = await this.store.createRun(input.input.messages, threadId, input.config, input.metadata); const rb = RunBuilder.create(run, threadId); - // Register in runningTasks so cancelRun/destroy can manage streaming runs. + // Create event buffer for this run (events persisted to JSONL file) + if (!existsSync(EVENT_DIR)) { + mkdirSync(EVENT_DIR, { recursive: true }); + } + const buffer: RunEventBuffer = { + filePath: join(EVENT_DIR, `${run.id}.jsonl`), + lastSeq: 0, + done: false, + emitter: new EventEmitter(), + }; + this.runBuffers.set(run.id, buffer); + + // Emit initial lifecycle event + this.pushEvent(buffer, 'run_created', { runId: run.id, threadId }); + + // Start background execution (not tied to SSE connection) + const abortController = new AbortController(); let resolveTask!: () => void; const taskPromise = new Promise((r) => { resolveTask = r; }); - this.runningTasks.set(run.id, { promise: taskPromise, abortController }); + const task: RunTaskState = { + promise: taskPromise, + abortController, + committed: false, + emitter: new EventEmitter(), + }; + this.runningTasks.set(run.id, task); - // event: thread.run.created - writer.writeEvent(AgentSSEEvent.ThreadRunCreated, rb.snapshot()); + this.executeStreamBackground(input, run.id, threadId, rb, buffer, task).finally(() => { + task.emitter.emit('end'); + resolveTask(); + this.runningTasks.delete(run.id); + this.runBuffers.delete(run.id); + buffer.emitter.removeAllListeners(); + }); - // event: thread.run.in_progress - await this.store.updateRun(run.id, rb.start()); - writer.writeEvent(AgentSSEEvent.ThreadRunInProgress, rb.snapshot()); + // Stream events to the current client + await this.streamEventsToWriter(buffer, writer, 0); + } - const msgId = newMsgId(); + /** + * Reconnect to a running or completed run's event stream. + * Replays events after lastSeq, then continues real-time if still running. + */ + async getRunStream(runId: string, writer: SSEWriter, lastSeq = 0): Promise { + const buffer = this.runBuffers.get(runId); + if (buffer) { + await this.streamEventsToWriter(buffer, writer, lastSeq); + return; + } - // event: thread.message.created - const msgObj = MessageConverter.createStreamMessage(msgId, run.id); - writer.writeEvent(AgentSSEEvent.ThreadMessageCreated, msgObj); + // Task already finished — replay from JSONL file directly + const filePath = join(EVENT_DIR, `${runId}.jsonl`); + if (!existsSync(filePath)) { + throw new AgentNotFoundError(`Run event stream not found: ${runId}`); + } + for await (const event of this.readEventsFromFile(filePath, lastSeq)) { + if (writer.closed) return; + writer.writeEvent(event.type, event); + } + if (!writer.closed) writer.end(); + } + private pushEvent(buffer: RunEventBuffer, type: string, data: unknown): void { + const event: StreamEvent = { + seq: ++buffer.lastSeq, + type, + data, + ts: Date.now(), + }; + appendFileSync(buffer.filePath, JSON.stringify(event) + '\n'); + buffer.emitter.emit('event', event); + } + + /** + * Execute the run in the background, persisting events to JSONL file. + * AgentMessage objects are passed through directly as event data. + */ + private async executeStreamBackground( + input: CreateRunInput, + runId: string, + threadId: string, + rb: RunBuilder, + buffer: RunEventBuffer, + task: RunTaskState, + ): Promise { + const abortController = task.abortController; + const streamMessages: AgentMessage[] = []; + // Persist a turn's transcript to the thread at most once (see syncRun). + let messagesPersisted = false; + const persistPartialOnce = async (): Promise => { + if (!task.committed || messagesPersisted) return; + messagesPersisted = true; + await this.persistPartialMessages(threadId, input, streamMessages); + }; try { - const { content, usage, aborted } = await this.consumeStreamMessages( - input, - abortController.signal, - writer, - msgId, - ); - - if (aborted) { - // Skip intermediate cancelling store write — no external observer between the - // two states since the SSE client has already disconnected. - rb.cancelling(); - try { - await this.store.updateRun(run.id, rb.cancel()); - } catch (storeErr) { - this.logger.error('[AgentRuntime] failed to write cancelled status during stream abort:', storeErr); - } - if (!writer.closed) { - writer.writeEvent(AgentSSEEvent.ThreadRunCancelled, rb.snapshot()); + await this.store.updateRun(runId, rb.start()); + + for await (const msg of this.executor.execRun(input, abortController.signal)) { + if (abortController.signal.aborted) { + await persistPartialOnce(); + await this.finaliseAbortedRun(runId); + this.pushEvent(buffer, 'error', { message: 'cancelled', runId }); + return; } - return; + + streamMessages.push(msg); + + // Pass through SDK message directly as event data + const eventType = msg.type || 'message'; + this.pushEvent(buffer, eventType, msg); + + await this.markCommittedIfNeeded(task, msg, streamMessages); } - // event: thread.message.completed - const completedMsg = MessageConverter.completeMessage(msgObj, content); - writer.writeEvent(AgentSSEEvent.ThreadMessageCompleted, completedMsg); + // TOCTOU: respect any terminal-ish status set by another worker + // (cancelRun, its commit-timeout watchdog which writes Failed, or + // an external expiration) instead of overwriting it with Completed. + const currentRun = await this.store.getRun(runId); + if (AgentRuntime.POST_LOOP_TERMINAL_STATUSES.has(currentRun.status)) { + await persistPartialOnce(); + this.pushEvent(buffer, 'error', { message: currentRun.status, runId }); + return; + } - // Persist and emit completion — append messages before marking run as completed - // so a failure leaves the run in_progress (retryable) instead of completed-but-incomplete. - // TODO(atomicity): add aggregate store method for full transactional guarantee. - const output: MessageObject[] = content.length > 0 ? [completedMsg] : []; + // Persist to store (excluding stream_event deltas) + const usage = MessageConverter.extractUsage(streamMessages); await this.store.appendMessages(threadId, [ - ...MessageConverter.toInputMessageObjects(input.input.messages, threadId), - ...output, + ...MessageConverter.toAgentMessages(input.input.messages), + ...MessageConverter.filterForStorage(streamMessages), ]); - await this.store.updateRun(run.id, rb.complete(output, usage)); + messagesPersisted = true; + await this.store.updateRun(runId, rb.complete(usage)); - // event: thread.run.completed - writer.writeEvent(AgentSSEEvent.ThreadRunCompleted, rb.snapshot()); + this.pushEvent(buffer, 'done', { result: 'success', runId }); } catch (err: unknown) { - if (abortController.signal.aborted) { - // Client disconnected or cancelRun fired — mark as cancelled, not failed - rb.cancelling(); - try { - await this.store.updateRun(run.id, rb.cancel()); - } catch (storeErr) { - this.logger.error('[AgentRuntime] failed to write cancelled status during stream error:', storeErr); - } - if (!writer.closed) { - writer.writeEvent(AgentSSEEvent.ThreadRunCancelled, rb.snapshot()); - } - } else { + if (!abortController.signal.aborted) { + // Non-abort failure (e.g. upstream stream terminated mid-turn). + // Persist the partial transcript before marking the run failed so the + // thread history is not silently dropped. No-op if the success path + // already appended (avoids duplicate history). + await persistPartialOnce(); try { - await this.store.updateRun(run.id, rb.fail(err as Error)); + const currentRun = await this.store.getRun(runId); + if (currentRun.status !== RunStatus.Cancelling && currentRun.status !== RunStatus.Cancelled) { + await this.store.updateRun(runId, rb.fail(err as Error)); + } } catch (storeErr) { this.logger.error('[AgentRuntime] failed to update run status after error:', storeErr); } - - // event: thread.run.failed - if (!writer.closed) { - writer.writeEvent(AgentSSEEvent.ThreadRunFailed, rb.snapshot()); - } + this.pushEvent(buffer, 'error', { message: (err as Error).message, runId }); + } else { + await persistPartialOnce(); + await this.finaliseAbortedRun(runId); + this.logger.error('[AgentRuntime] execRun error during abort:', err); + this.pushEvent(buffer, 'error', { message: 'cancelled', runId }); } } finally { - resolveTask(); - this.runningTasks.delete(run.id); + buffer.done = true; + buffer.emitter.emit('event'); + } + } - // event: done - if (!writer.closed) { - writer.writeEvent(AgentSSEEvent.Done, '[DONE]'); - writer.end(); - } + /** + * Flip the task's `committed` flag the first time the executor's current + * message indicates its session has been persisted to storage. Uses the + * executor's `isSessionCommitted` hook when available, otherwise a default + * heuristic where any message with `type !== 'system'` counts as committed + * (the Claude Code SDK writes the jsonl around the first non-system event). + */ + private async markCommittedIfNeeded(task: RunTaskState, msg: AgentMessage, history: AgentMessage[]): Promise { + if (task.committed) return; + let committed: boolean; + try { + committed = + typeof this.executor.isSessionCommitted === 'function' + ? await this.executor.isSessionCommitted(msg, history) + : msg.type !== 'system'; + } catch (err) { + this.logger.error('[AgentRuntime] isSessionCommitted threw, treating as not committed:', err); + committed = false; + } + if (committed) { + task.committed = true; + task.emitter.emit('commit'); } } /** - * Consume the execRun async generator, emitting SSE message.delta events - * for each chunk and accumulating content blocks and token usage. + * Wait until the task reports that its session is committed, or the task + * finishes on its own, or the timeout elapses. Rejects with + * AgentTimeoutError on timeout. Resolves without error when the task ends + * before committing — in that case the caller should re-read the run's + * terminal status rather than trying to cancel further. */ - private async consumeStreamMessages( + private waitForCommitted(task: RunTaskState, timeoutMs: number): Promise { + if (task.committed) return Promise.resolve(); + return new Promise((resolve, reject) => { + // Handlers need to reference each other (cleanup must off() all of + // commit / end / timer), which would force a forward reference if the + // arrow functions referred to each other by name. Stash them on a + // shared container so cleanup can read them by property access and the + // source order stays linear. + const refs: { + timer?: ReturnType; + onCommit?: () => void; + onEnd?: () => void; + } = {}; + const cleanup = (): void => { + if (refs.timer) clearTimeout(refs.timer); + if (refs.onCommit) task.emitter.off('commit', refs.onCommit); + if (refs.onEnd) task.emitter.off('end', refs.onEnd); + }; + refs.onCommit = () => { + cleanup(); + resolve(); + }; + refs.onEnd = () => { + cleanup(); + resolve(); + }; + refs.timer = globalThis.setTimeout(() => { + cleanup(); + reject( + new AgentTimeoutError(`Timed out waiting ${timeoutMs}ms for executor session to be committed before cancel`), + ); + }, timeoutMs); + task.emitter.once('commit', refs.onCommit); + task.emitter.once('end', refs.onEnd); + }); + } + + /** + * Persist input + collected stream messages to the thread when a run ends + * on a non-success path — either an abort/cancel, or a mid-turn failure + * (e.g. the upstream stream is terminated). Keeping the thread in sync with + * any partial state that the executor has already written (e.g. Claude CLI + * session file) is what allows subsequent resume requests to continue from a + * consistent history instead of diverging and failing at executor startup, + * and prevents a failed turn from vanishing entirely from thread history. + * + * Callers must check `task.committed` before invoking this; if the + * executor never reached a committed state the thread should be left + * untouched so the next run starts fresh instead of trying to resume a + * session that was never created on disk. + * + * Errors are swallowed here so a store failure cannot mask the original + * abort/failure or prevent the run status from being finalised. + */ + private async persistPartialMessages( + threadId: string, input: CreateRunInput, - signal: AbortSignal, - writer: SSEWriter, - msgId: string, - ): Promise<{ content: MessageContentBlock[]; usage?: RunUsage; aborted: boolean }> { - const content: MessageContentBlock[] = []; - let promptTokens = 0; - let completionTokens = 0; - let hasUsage = false; - - for await (const msg of this.executor.execRun(input, signal)) { - if (signal.aborted) { - return { content, usage: undefined, aborted: true as const }; + streamMessages: AgentMessage[], + ): Promise { + try { + await this.store.appendMessages(threadId, [ + ...MessageConverter.toAgentMessages(input.input.messages), + ...MessageConverter.filterForStorage(streamMessages), + ]); + } catch (err) { + this.logger.error('[AgentRuntime] failed to persist messages on abort:', err); + } + } + + /** + * Push an aborted run to a terminal `cancelled` state when nobody else + * will. Abort can be driven either by `cancelRun` — which already owns + * the `in_progress → cancelling → cancelled` transition — or by an + * external `AbortSignal` / `destroy()`, where the run would otherwise + * stay stuck in `in_progress` forever. + * + * Behaviour: + * - terminal status (completed/failed/cancelled/expired): no-op. + * - `cancelling`: no-op, let `cancelRun` finish the transition. + * - `in_progress` / `queued`: write `cancelling` then `cancelled`. + * + * Errors are swallowed so a store failure cannot mask the abort. + */ + private async finaliseAbortedRun(runId: string): Promise { + try { + const current = await this.store.getRun(runId); + if (AgentRuntime.TERMINAL_RUN_STATUSES.has(current.status)) return; + if (current.status === RunStatus.Cancelling) return; + + const rb = RunBuilder.fromRecord(current); + await this.store.updateRun(runId, rb.cancelling()); + await this.store.updateRun(runId, rb.cancel()); + } catch (err) { + this.logger.error('[AgentRuntime] failed to finalise aborted run:', err); + } + } + + private async *readEventsFromFile(filePath: string, afterSeq: number): AsyncGenerator { + if (!existsSync(filePath)) return; + const rl = createInterface({ input: createReadStream(filePath) }); + for await (const line of rl) { + if (!line.trim()) continue; + try { + const event = JSON.parse(line) as StreamEvent; + if (event.seq > afterSeq) { + yield event; + } + } catch { + // skip malformed lines } - if (msg.message) { - const contentBlocks = MessageConverter.toContentBlocks(msg.message); - content.push(...contentBlocks); - - // event: thread.message.delta - const delta: MessageDeltaObject = { - id: msgId, - object: AgentObjectType.ThreadMessageDelta, - delta: { content: contentBlocks }, - }; - writer.writeEvent(AgentSSEEvent.ThreadMessageDelta, delta); + } + } + + private async streamEventsToWriter(buffer: RunEventBuffer, writer: SSEWriter, lastSeq: number): Promise { + // Phase 1: Replay from JSONL file + let lastWrittenSeq = lastSeq; + for await (const event of this.readEventsFromFile(buffer.filePath, lastSeq)) { + if (writer.closed) return; + writer.writeEvent(event.type, event); + lastWrittenSeq = event.seq; + } + + if (buffer.done) { + if (!writer.closed) writer.end(); + return; + } + + // Phase 2: Real-time events via EventEmitter + heartbeat + const queue: StreamEvent[] = []; + let waitResolve: (() => void) | null = null; + + function onEvent(event?: StreamEvent): void { + if (event) queue.push(event); + waitResolve?.(); + } + + buffer.emitter.on('event', onEvent); + + try { + // Catch-up: drain any events that arrived during Phase 1 file read + for await (const event of this.readEventsFromFile(buffer.filePath, lastWrittenSeq)) { + if (writer.closed) return; + if (event.seq > lastWrittenSeq) { + writer.writeEvent(event.type, event); + lastWrittenSeq = event.seq; + } } - if (msg.usage) { - hasUsage = true; - promptTokens += msg.usage.promptTokens ?? 0; - completionTokens += msg.usage.completionTokens ?? 0; + + const waitForEvent = () => + new Promise<'event' | 'heartbeat'>((resolve) => { + waitResolve = () => resolve('event'); + setTimeout(() => resolve('heartbeat'), HEARTBEAT_INTERVAL_MS); + }); + + while (!buffer.done || queue.length > 0) { + while (queue.length > 0) { + const event = queue.shift()!; + if (event.seq > lastWrittenSeq) { + if (writer.closed) return; + writer.writeEvent(event.type, event); + lastWrittenSeq = event.seq; + } + } + + if (buffer.done) break; + if (writer.closed) return; + + const reason = await waitForEvent(); + waitResolve = null; + if (reason === 'heartbeat' && queue.length === 0 && !buffer.done) { + if (writer.closed) return; + writer.writeComment('keepalive'); + } } + } finally { + buffer.emitter.off('event', onEvent); } - return { - content, - usage: hasUsage ? { promptTokens, completionTokens, totalTokens: promptTokens + completionTokens } : undefined, - aborted: false as const, - }; + if (!writer.closed) writer.end(); } async getRun(runId: string): Promise { @@ -385,42 +772,86 @@ export class AgentRuntime { return RunBuilder.fromRecord(run).snapshot(); } + /** + * Resolve the most recent run created on a thread. Returns `{ runId: null }` + * when the thread exists but has no recorded run (e.g. threads created + * before run tracking, or with no runs yet). Throws AgentNotFoundError when + * the thread does not exist. + */ + async getLatestRunId(threadId: string): Promise<{ threadId: string; runId: string | null }> { + const runId = await this.store.getLatestRunId(threadId); + return { threadId, runId }; + } + + /** + * Cancel a running task. The call blocks until either (a) the executor + * reports its session is safely committed to storage, and the task has + * been aborted and the thread persisted, or (b) the commit watchdog times + * out, in which case the run is marked `failed` (not `cancelled`) and + * AgentTimeoutError is thrown to the caller. + * + * The hold is there to guarantee that whatever user input the thread + * records on abort is also present in the executor's own persistent + * session (e.g. Claude Code SDK jsonl), so a subsequent resume request + * on the same thread doesn't diverge from a session that was never + * actually written. + */ async cancelRun(runId: string): Promise { - // 1. Check current status — reject if already terminal const run = await this.store.getRun(runId); if (AgentRuntime.TERMINAL_RUN_STATUSES.has(run.status)) { throw new AgentConflictError(`Cannot cancel run with status '${run.status}'`); } const rb = RunBuilder.fromRecord(run); - - // 2. Write "cancelling" to store first — visible to all workers await this.store.updateRun(runId, rb.cancelling()); - // 3. If the task is running locally, abort it for immediate effect const task = this.runningTasks.get(runId); if (task) { + if (!task.committed) { + this.logger.info( + '[AgentRuntime] cancelRun %s holding up to %dms for executor session to commit', + runId, + this.cancelCommitTimeoutMs, + ); + try { + await this.waitForCommitted(task, this.cancelCommitTimeoutMs); + } catch (err) { + // Commit watchdog timed out. Mark the run as failed *before* + // aborting so the execution path's finaliseAbortedRun sees a + // terminal status and skips the cancelled transition. The thread + // is left untouched because task.committed is still false. + this.logger.error( + '[AgentRuntime] cancelRun %s timed out after %dms waiting for executor to commit; marking run failed and leaving thread untouched', + runId, + this.cancelCommitTimeoutMs, + ); + try { + await this.store.updateRun(runId, rb.fail(err as Error)); + } catch (storeErr) { + this.logger.error('[AgentRuntime] failed to mark run failed after cancel timeout:', storeErr); + } + task.abortController.abort(); + await task.promise.catch(() => { + /* ignore */ + }); + throw err; + } + } task.abortController.abort(); await task.promise.catch(() => { /* ignore */ }); } - // 4. Re-read store to mitigate TOCTOU: if the run completed/failed between - // steps 2 and 4, do not overwrite the terminal state. - // TODO: For full atomicity, use CAS / ETag-based conditional writes. const freshRun = await this.store.getRun(runId); if (AgentRuntime.TERMINAL_RUN_STATUSES.has(freshRun.status)) { - // Run reached a terminal state while we were cancelling — return as-is return RunBuilder.fromRecord(freshRun).snapshot(); } - // 5. Transition to final "cancelled" state try { await this.store.updateRun(runId, rb.cancel()); } catch (err) { this.logger.error('[AgentRuntime] failed to write cancelled state after cancelling:', err); - // Return best-effort snapshot from store const fallback = await this.store.getRun(runId); return RunBuilder.fromRecord(fallback).snapshot(); } @@ -428,7 +859,6 @@ export class AgentRuntime { return rb.snapshot(); } - /** Wait for all in-flight background tasks to complete naturally (without aborting). */ async waitForPendingTasks(): Promise { if (this.runningTasks.size) { const pending = Array.from(this.runningTasks.values()).map((t) => t.promise); @@ -437,19 +867,21 @@ export class AgentRuntime { } async destroy(): Promise { - // Abort all in-flight background tasks, then wait for them to settle for (const task of this.runningTasks.values()) { task.abortController.abort(); } await this.waitForPendingTasks(); - // Destroy store + for (const buffer of this.runBuffers.values()) { + buffer.emitter.removeAllListeners(); + } + this.runBuffers.clear(); + if (this.store.destroy) { await this.store.destroy(); } } - /** Factory method — avoids the spread-arg type issue with dynamic delegation. */ static create(options: AgentRuntimeOptions): AgentRuntime { return new AgentRuntime(options); } diff --git a/tegg/core/agent-runtime/src/AgentStoreUtils.ts b/tegg/core/agent-runtime/src/AgentStoreUtils.ts index 775219b7f0..43b57937f4 100644 --- a/tegg/core/agent-runtime/src/AgentStoreUtils.ts +++ b/tegg/core/agent-runtime/src/AgentStoreUtils.ts @@ -15,3 +15,31 @@ export function newThreadId(): string { export function newRunId(): string { return `run_${crypto.randomUUID()}`; } + +/** + * Upper bound for 13-digit millisecond timestamps. The time complement + * `TS_MAX_MS - ms` sorts newest-first in ascending key order. + */ +export const TS_MAX_MS = 9_999_999_999_999; + +const REV_MS_WIDTH = String(TS_MAX_MS).length; + +/** + * Encode a millisecond timestamp so ascending key order is newest-first. + */ +export function reverseMs(ms: number): string { + if (!Number.isInteger(ms) || ms < 0 || ms > TS_MAX_MS) { + throw new RangeError(`reverseMs: ms must be an integer in [0, ${TS_MAX_MS}], got ${ms}`); + } + return String(TS_MAX_MS - ms).padStart(REV_MS_WIDTH, '0'); +} + +/** + * Format a Unix-millisecond timestamp as a UTC `YYYY-MM-DD` bucket. + */ +export function dateBucket(ms: number): string { + if (!Number.isInteger(ms) || ms < 0) { + throw new RangeError(`dateBucket: ms must be a nonnegative integer Unix-millisecond timestamp, got ${ms}`); + } + return new Date(ms).toISOString().slice(0, 10); +} diff --git a/tegg/core/agent-runtime/src/HttpSSEWriter.ts b/tegg/core/agent-runtime/src/HttpSSEWriter.ts index 3adf93a8bf..8ed92fd2b4 100644 --- a/tegg/core/agent-runtime/src/HttpSSEWriter.ts +++ b/tegg/core/agent-runtime/src/HttpSSEWriter.ts @@ -36,6 +36,12 @@ export class HttpSSEWriter implements SSEWriter { this.res.write(`event: ${event}\ndata: ${JSON.stringify(data)}\n\n`); } + writeComment(text: string): void { + if (this._closed) return; + this.ensureHeaders(); + this.res.write(`: ${text}\n\n`); + } + get closed(): boolean { return this._closed; } diff --git a/tegg/core/agent-runtime/src/MessageConverter.ts b/tegg/core/agent-runtime/src/MessageConverter.ts index 78ded93e24..b3d6e1d8ca 100644 --- a/tegg/core/agent-runtime/src/MessageConverter.ts +++ b/tegg/core/agent-runtime/src/MessageConverter.ts @@ -1,129 +1,57 @@ -import type { - CreateRunInput, - MessageObject, - MessageContentBlock, - AgentStreamMessage, - AgentStreamMessagePayload, -} from '@eggjs/tegg-types/agent-runtime'; -import { AgentObjectType, MessageRole, MessageStatus, ContentBlockType } from '@eggjs/tegg-types/agent-runtime'; +import type { AgentMessage, InputMessage, SDKResultMessage } from '@eggjs/tegg-types/agent-runtime'; -import { nowUnix, newMsgId } from './AgentStoreUtils.ts'; import type { RunUsage } from './RunBuilder.ts'; export class MessageConverter { /** - * Convert an AgentStreamMessage's message payload into OpenAI MessageContentBlock[]. + * Extract accumulated usage from AgentMessage objects. + * Only `result` type messages carry usage information. */ - static toContentBlocks(msg: AgentStreamMessagePayload): MessageContentBlock[] { - if (!msg) return []; - const content = msg.content; - if (typeof content === 'string') { - return [{ type: ContentBlockType.Text, text: { value: content, annotations: [] } }]; - } - if (Array.isArray(content)) { - return content - .filter((part) => part.type === ContentBlockType.Text) - .map((part) => ({ type: ContentBlockType.Text, text: { value: part.text, annotations: [] } })); - } - return []; - } - - /** - * Build a completed MessageObject from an AgentStreamMessage payload. - */ - static toMessageObject(msg: AgentStreamMessagePayload, runId?: string): MessageObject { - return { - id: newMsgId(), - object: AgentObjectType.ThreadMessage, - createdAt: nowUnix(), - runId, - role: MessageRole.Assistant, - status: MessageStatus.Completed, - content: MessageConverter.toContentBlocks(msg), - }; - } - - /** - * Extract MessageObjects and accumulated usage from AgentStreamMessage objects. - */ - static extractFromStreamMessages( - messages: AgentStreamMessage[], - runId?: string, - ): { - output: MessageObject[]; - usage?: RunUsage; - } { - const output: MessageObject[] = []; + static extractUsage(messages: AgentMessage[]): RunUsage | undefined { let promptTokens = 0; let completionTokens = 0; let hasUsage = false; for (const msg of messages) { - if (msg.message) { - output.push(MessageConverter.toMessageObject(msg.message, runId)); + if (msg.type === 'result') { + const resultMsg = msg as SDKResultMessage; + if (resultMsg.usage) { + hasUsage = true; + promptTokens += resultMsg.usage.input_tokens ?? 0; + completionTokens += resultMsg.usage.output_tokens ?? 0; + } } - if (msg.usage) { - hasUsage = true; - promptTokens += msg.usage.promptTokens ?? 0; - completionTokens += msg.usage.completionTokens ?? 0; - } - } - - let usage: RunUsage | undefined; - if (hasUsage) { - usage = { - promptTokens, - completionTokens, - totalTokens: promptTokens + completionTokens, - }; } - return { output, usage }; - } - - /** - * Produce a completed copy of a streaming MessageObject with final content. - */ - static completeMessage(msg: MessageObject, content: MessageContentBlock[]): MessageObject { - return { ...msg, status: MessageStatus.Completed, content }; + if (!hasUsage) return undefined; + return { + promptTokens, + completionTokens, + totalTokens: promptTokens + completionTokens, + }; } /** - * Create an in-progress MessageObject for streaming (before content is known). + * Filter out stream_event messages before persisting to thread storage. + * Stream events are incremental deltas (one per token) only useful during + * real-time streaming; the final assistant message already contains the + * complete response. */ - static createStreamMessage(msgId: string, runId: string): MessageObject { - return { - id: msgId, - object: AgentObjectType.ThreadMessage, - createdAt: nowUnix(), - runId, - role: MessageRole.Assistant, - status: MessageStatus.InProgress, - content: [], - }; + static filterForStorage(messages: AgentMessage[]): AgentMessage[] { + return messages.filter((m) => m.type !== 'stream_event'); } /** - * Convert input messages to MessageObjects for thread history. - * System messages are filtered out — they are transient instructions, not conversation history. + * Convert input messages to AgentMessage format for thread history. + * System messages are filtered out — they are transient instructions, + * not conversation history. */ - static toInputMessageObjects(messages: CreateRunInput['input']['messages'], threadId?: string): MessageObject[] { + static toAgentMessages(messages: InputMessage[]): AgentMessage[] { return messages - .filter( - (m): m is typeof m & { role: Exclude } => - m.role !== MessageRole.System, - ) + .filter((m) => m.role !== 'system') .map((m) => ({ - id: newMsgId(), - object: AgentObjectType.ThreadMessage, - createdAt: nowUnix(), - threadId, - role: m.role, - status: MessageStatus.Completed, - content: - typeof m.content === 'string' - ? [{ type: ContentBlockType.Text, text: { value: m.content, annotations: [] } }] - : m.content.map((p) => ({ type: ContentBlockType.Text, text: { value: p.text, annotations: [] } })), + type: m.role as 'user' | 'assistant', + message: { role: m.role, content: m.content }, })); } } diff --git a/tegg/core/agent-runtime/src/OSSAgentStore.ts b/tegg/core/agent-runtime/src/OSSAgentStore.ts index 1221fa4a57..43ac31cc5a 100644 --- a/tegg/core/agent-runtime/src/OSSAgentStore.ts +++ b/tegg/core/agent-runtime/src/OSSAgentStore.ts @@ -1,113 +1,158 @@ import type { + AgentMessage, AgentRunConfig, AgentStore, + GetThreadOptions, InputMessage, - MessageObject, RunRecord, ThreadRecord, + ObjectStorageClient, } from '@eggjs/tegg-types/agent-runtime'; -import { AgentObjectType, RunStatus } from '@eggjs/tegg-types/agent-runtime'; -import { AgentNotFoundError } from '@eggjs/tegg-types/agent-runtime'; -import type { ObjectStorageClient } from '@eggjs/tegg-types/agent-runtime'; +import { AgentObjectType, RunStatus, AgentNotFoundError } from '@eggjs/tegg-types/agent-runtime'; -import { nowUnix, newThreadId, newRunId } from './AgentStoreUtils.ts'; +import { dateBucket, newRunId, newThreadId, nowUnix, reverseMs } from './AgentStoreUtils.ts'; + +/** + * Warn logger used when a background thread activity-index write fails. + * `console` and `egg-logger` both satisfy this shape. + */ +export interface OSSAgentStoreWarnLogger { + warn(message: string, ...args: unknown[]): void; +} export interface OSSAgentStoreOptions { client: ObjectStorageClient; prefix?: string; + /** + * Background activity-index PUT failures are logged here and never + * propagated to store callers. + */ + logger?: OSSAgentStoreWarnLogger; } /** * Thread metadata stored as a JSON object (excludes messages). - * Messages are stored separately in a JSONL file for append-friendly writes. + * Messages are stored separately in a JSONL file for append-friendly + * writes. */ type ThreadMetadata = Omit; /** - * AgentStore implementation backed by an ObjectStorageClient (OSS, S3, etc.). + * `AgentStore` implementation backed by an `ObjectStorageClient` (OSS, + * S3, etc.). * * ## Storage layout * * ``` - * {prefix}threads/{id}/meta.json — Thread metadata (JSON) - * {prefix}threads/{id}/messages.jsonl — Messages (JSONL, one JSON object per line) - * {prefix}runs/{id}.json — Run record (JSON) + * {prefix}threads/{id}/meta.json — Thread metadata (JSON, source of truth) + * {prefix}threads/{id}/messages.jsonl — Messages (JSONL, one AgentMessage per line) + * {prefix}runs/{id}.json — Run record (JSON) + * {prefix}index/threads-by-updated-date/{YYYY-MM-DD}/{revMs13}_{threadId} — Activity-time index sidecar (JSON body) * ``` * - * ### Why split threads into two keys? - * - * Thread messages are append-only: new messages are added at the end but never - * modified or deleted. Storing them as a JSONL file allows us to leverage the - * OSS AppendObject API (or similar) to write new messages without reading the - * entire thread first. This is much more efficient than read-modify-write for - * long conversations. - * - * If the underlying ObjectStorageClient provides an `append()` method, it will - * be used for O(1) message writes. Otherwise, the store falls back to - * get-concat-put (which is NOT atomic and may lose data under concurrent - * writers — acceptable for single-writer scenarios). - * - * ### Atomicity note - * - * Run updates still use read-modify-write because run fields are mutated - * (status, timestamps, output, etc.) — they cannot be modelled as append-only. - * For multi-writer safety, consider a database-backed AgentStore or ETag-based - * conditional writes with retry. + * The activity-time index sidecar is best-effort and write-only. It is not read + * by `getThread`, and `destroy()` drains any in-flight index writes. */ export class OSSAgentStore implements AgentStore { private readonly client: ObjectStorageClient; private readonly prefix: string; + private readonly logger: OSSAgentStoreWarnLogger; + private readonly pendingIndexWrites = new Set>(); + private readonly threadMetaWriteTails = new Map>(); constructor(options: OSSAgentStoreOptions) { this.client = options.client; - // Normalize: ensure non-empty prefix ends with '/' const raw = options.prefix ?? ''; this.prefix = raw && !raw.endsWith('/') ? raw + '/' : raw; + this.logger = options.logger ?? { warn: (message, ...args) => console.warn(message, ...args) }; } - // ── Key helpers ────────────────────────────────────────────────────── - - /** Key for thread metadata (JSON). */ private threadMetaKey(threadId: string): string { return `${this.prefix}threads/${threadId}/meta.json`; } - /** Key for thread messages (JSONL, one message per line). */ private threadMessagesKey(threadId: string): string { return `${this.prefix}threads/${threadId}/messages.jsonl`; } - /** Key for run record (JSON). */ private runKey(runId: string): string { return `${this.prefix}runs/${runId}.json`; } - // ── Lifecycle ──────────────────────────────────────────────────────── + private threadActivityIndexKey(nowMs: number, threadId: string): string { + const date = dateBucket(nowMs); + const rev = reverseMs(nowMs); + return `${this.prefix}index/threads-by-updated-date/${date}/${rev}_${threadId}`; + } + + private writeThreadActivityIndex( + threadId: string, + createdAt: number, + updatedAtMs: number, + metadata: Record, + ): void { + const indexKey = this.threadActivityIndexKey(updatedAtMs, threadId); + const indexBody = JSON.stringify({ + threadId, + createdAt, + updatedAt: Math.floor(updatedAtMs / 1000), + metadata, + }); + const tracked: Promise = this.client + .put(indexKey, indexBody) + .catch((err: unknown) => { + const errForLog: Error = err instanceof Error ? err : new Error(String(err)); + this.logger.warn( + '[OSSAgentStore] failed to write thread activity index threadId=%s key=%s', + threadId, + indexKey, + errForLog, + ); + }) + .finally(() => { + this.pendingIndexWrites.delete(tracked); + }); + this.pendingIndexWrites.add(tracked); + } async init(): Promise { await this.client.init?.(); } async destroy(): Promise { + await this.awaitPendingWrites(); await this.client.destroy?.(); } - // ── Thread operations ──────────────────────────────────────────────── + /** + * Wait for in-flight background activity-index writes. Individual write + * failures are logged when scheduled, so this method never rejects + * for index-write failures. The set size is bounded by in-flight index PUTs. + */ + async awaitPendingWrites(): Promise { + while (this.pendingIndexWrites.size > 0) { + await Promise.allSettled(this.pendingIndexWrites); + } + } async createThread(metadata?: Record): Promise { const threadId = newThreadId(); + const nowMs = Date.now(); + const createdAt = Math.floor(nowMs / 1000); const meta: ThreadMetadata = { id: threadId, object: AgentObjectType.Thread, metadata: metadata ?? {}, - createdAt: nowUnix(), + createdAt, }; await this.client.put(this.threadMetaKey(threadId), JSON.stringify(meta)); - // Messages file is created lazily on first appendMessages call. + + this.writeThreadActivityIndex(threadId, createdAt, nowMs, meta.metadata); + return { ...meta, messages: [] }; } - async getThread(threadId: string): Promise { + async getThread(threadId: string, options?: GetThreadOptions): Promise { const [metaData, messagesData] = await Promise.all([ this.client.get(this.threadMetaKey(threadId)), this.client.get(this.threadMessagesKey(threadId)), @@ -117,50 +162,100 @@ export class OSSAgentStore implements AgentStore { } const meta = JSON.parse(metaData) as ThreadMetadata; - // Parse messages JSONL — may not exist yet if no messages were appended. - const messages: MessageObject[] = messagesData + let messages: AgentMessage[] = messagesData ? messagesData .trim() .split('\n') .filter((line) => line.length > 0) - .map((line) => JSON.parse(line) as MessageObject) + .map((line) => JSON.parse(line) as AgentMessage) : []; + // By default only return the conversation-shape messages (user + + // assistant), matching the SDK's `getSessionMessages` semantics. + // The JSONL file is the full event log including framework-level + // entries (system events, tool results, etc.); the filter + // narrows the visible set to the application-level conversation. + if (!options?.includeAllMessages) { + messages = messages.filter((m) => m.type === 'user' || m.type === 'assistant'); + } + return { ...meta, messages }; } /** - * Append messages to a thread. + * Serialize read-modify-write operations on a single thread's `meta.json` + * within this process. Both {@link updateThreadMetadata} (business metadata + * merge) and {@link recordLatestRunId} (the `latestRunId` pointer) mutate the + * same `meta.json` with no compare-and-swap, so without a shared per-thread + * lock a concurrent run could clobber the other's write. Cross-process writes + * remain last-writer-wins. * - * Each message is serialized as a single JSON line (JSONL format). - * When the underlying client supports `append()`, this is a single - * O(1) write — no need to read the existing messages first. + * The lock is a per-thread promise chain: `current` is always resolved in the + * `finally` (decoupled from `fn`'s success/failure), so a rejecting `fn` never + * leaves the chain stuck for subsequent waiters. */ - async appendMessages(threadId: string, messages: MessageObject[]): Promise { - // Verify the thread exists before writing messages (or returning early), - // so callers always get AgentNotFoundError for invalid threadIds. + private async runExclusiveThreadMetaWrite(threadId: string, fn: () => Promise): Promise { + const previous = this.threadMetaWriteTails.get(threadId) ?? Promise.resolve(); + let release!: () => void; + const current = new Promise((resolve) => { + release = resolve; + }); + const tail = previous.then(() => current); + this.threadMetaWriteTails.set(threadId, tail); + + await previous; + try { + return await fn(); + } finally { + release(); + if (this.threadMetaWriteTails.get(threadId) === tail) { + this.threadMetaWriteTails.delete(threadId); + } + } + } + + async updateThreadMetadata(threadId: string, metadata: Record): Promise { + if (Object.keys(metadata).length === 0) return; + + await this.runExclusiveThreadMetaWrite(threadId, async () => { + const metaData = await this.client.get(this.threadMetaKey(threadId)); + if (!metaData) { + throw new AgentNotFoundError(`Thread ${threadId} not found`); + } + const meta = JSON.parse(metaData) as ThreadMetadata; + const mergedMetadata = { ...meta.metadata, ...metadata }; + await this.client.put( + this.threadMetaKey(threadId), + JSON.stringify({ + ...meta, + metadata: mergedMetadata, + }), + ); + this.writeThreadActivityIndex(threadId, meta.createdAt, Date.now(), mergedMetadata); + }); + } + + async appendMessages(threadId: string, messages: AgentMessage[]): Promise { const metaData = await this.client.get(this.threadMetaKey(threadId)); if (!metaData) { throw new AgentNotFoundError(`Thread ${threadId} not found`); } if (messages.length === 0) return; + const meta = JSON.parse(metaData) as ThreadMetadata; + const nowMs = Date.now(); const lines = messages.map((m) => JSON.stringify(m)).join('\n') + '\n'; const messagesKey = this.threadMessagesKey(threadId); if (this.client.append) { - // Fast path: use the native append API (e.g., OSS AppendObject). await this.client.append(messagesKey, lines); } else { - // Slow path: read-modify-write fallback. - // NOTE: Not atomic — concurrent appends may lose data. const existing = (await this.client.get(messagesKey)) ?? ''; await this.client.put(messagesKey, existing + lines); } + this.writeThreadActivityIndex(threadId, meta.createdAt, nowMs, meta.metadata); } - // ── Run operations ─────────────────────────────────────────────────── - async createRun( input: InputMessage[], threadId?: string, @@ -179,9 +274,62 @@ export class OSSAgentStore implements AgentStore { createdAt: nowUnix(), }; await this.client.put(this.runKey(runId), JSON.stringify(record)); + if (threadId) { + // Fire-and-forget: keep the latestRunId pointer write off the run-creation + // hot path (it would otherwise add a GET+PUT roundtrip to every run + // start). Tracked in pendingIndexWrites — same as writeThreadActivityIndex + // — so destroy() drains it before shutdown. recordLatestRunId never rejects + // (it logs and swallows), so no unhandled rejection can leak here. + const tracked: Promise = this.recordLatestRunId(threadId, runId).finally(() => { + this.pendingIndexWrites.delete(tracked); + }); + this.pendingIndexWrites.add(tracked); + } return record; } + /** + * Record `runId` as the thread's most recent run by merging `latestRunId` + * into the thread metadata (a read-modify-write on `meta.json` that + * preserves all existing fields). + * + * Runs in the background (scheduled fire-and-forget by `createRun` and + * tracked in `pendingIndexWrites`) so it never adds latency to run creation; + * `awaitPendingWrites()` / `destroy()` drain it. + * + * Best-effort: a missing thread meta or any storage error is logged and + * swallowed. When it is skipped, `getLatestRunId` simply degrades to + * returning the previously recorded value (or `null`). + * + * Weak consistency: because the write is asynchronous and an unconditional + * read-modify-write with no compare-and-swap, `getLatestRunId` is eventually + * consistent — a read racing a just-created run may briefly see the prior + * value, and concurrent runs in different processes are last-writer-wins. + * Within one process it shares {@link runExclusiveThreadMetaWrite} with + * `updateThreadMetadata`, so the two never clobber each other's `meta.json`. + */ + private async recordLatestRunId(threadId: string, runId: string): Promise { + try { + await this.runExclusiveThreadMetaWrite(threadId, async () => { + const metaData = await this.client.get(this.threadMetaKey(threadId)); + if (!metaData) { + this.logger.warn( + '[OSSAgentStore] skip latestRunId write: thread meta not found threadId=%s runId=%s', + threadId, + runId, + ); + return; + } + const meta = JSON.parse(metaData) as ThreadMetadata; + meta.latestRunId = runId; + await this.client.put(this.threadMetaKey(threadId), JSON.stringify(meta)); + }); + } catch (err: unknown) { + const errForLog: Error = err instanceof Error ? err : new Error(String(err)); + this.logger.warn('[OSSAgentStore] failed to record latestRunId threadId=%s runId=%s', threadId, runId, errForLog); + } + } + async getRun(runId: string): Promise { const data = await this.client.get(this.runKey(runId)); if (!data) { @@ -190,12 +338,20 @@ export class OSSAgentStore implements AgentStore { return JSON.parse(data) as RunRecord; } - // TODO: read-modify-write is NOT atomic. Concurrent updates may lose data. - // Acceptable for single-writer scenarios; for multi-writer, consider ETag-based - // conditional writes with retry, or use a database-backed AgentStore instead. + async getLatestRunId(threadId: string): Promise { + const metaData = await this.client.get(this.threadMetaKey(threadId)); + if (!metaData) { + throw new AgentNotFoundError(`Thread ${threadId} not found`); + } + const meta = JSON.parse(metaData) as ThreadMetadata; + return meta.latestRunId ?? null; + } + async updateRun(runId: string, updates: Partial): Promise { const run = await this.getRun(runId); - const { id: _, object: __, ...safeUpdates } = updates; + const safeUpdates = { ...updates }; + delete safeUpdates.id; + delete (safeUpdates as any).object; Object.assign(run, safeUpdates); await this.client.put(this.runKey(runId), JSON.stringify(run)); } diff --git a/tegg/core/agent-runtime/src/RunBuilder.ts b/tegg/core/agent-runtime/src/RunBuilder.ts index 5b071b1ed8..7b67304047 100644 --- a/tegg/core/agent-runtime/src/RunBuilder.ts +++ b/tegg/core/agent-runtime/src/RunBuilder.ts @@ -1,6 +1,10 @@ -import type { MessageObject, RunObject, RunRecord, AgentRunConfig } from '@eggjs/tegg-types/agent-runtime'; -import { RunStatus, AgentErrorCode, AgentObjectType } from '@eggjs/tegg-types/agent-runtime'; -import { InvalidRunStateTransitionError } from '@eggjs/tegg-types/agent-runtime'; +import type { RunObject, RunRecord, AgentRunConfig } from '@eggjs/tegg-types/agent-runtime'; +import { + RunStatus, + AgentErrorCode, + AgentObjectType, + InvalidRunStateTransitionError, +} from '@eggjs/tegg-types/agent-runtime'; import { nowUnix } from './AgentStoreUtils.ts'; @@ -29,7 +33,6 @@ export class RunBuilder { private failedAt?: number; private lastError?: { code: string; message: string } | null; private usage?: RunUsage; - private output?: MessageObject[]; private constructor( id: string, @@ -60,7 +63,6 @@ export class RunBuilder { rb.cancelledAt = run.cancelledAt ?? undefined; rb.failedAt = run.failedAt ?? undefined; rb.lastError = run.lastError ?? undefined; - rb.output = run.output; if (run.usage) { rb.usage = { ...run.usage }; } @@ -78,25 +80,33 @@ export class RunBuilder { } /** in_progress -> completed. Returns store update. */ - complete(output: MessageObject[], usage?: RunUsage): Partial { + complete(usage?: RunUsage): Partial { if (this.status !== RunStatus.InProgress) { throw new InvalidRunStateTransitionError(this.status, RunStatus.Completed); } this.status = RunStatus.Completed; this.completedAt = nowUnix(); - this.output = output; this.usage = usage; return { status: this.status, - output, usage, completedAt: this.completedAt, }; } - /** queued/in_progress -> failed. Returns store update. */ + /** + * queued/in_progress/cancelling -> failed. Returns store update. + * + * `cancelling -> failed` covers the case where AgentRuntime has initiated + * a cancel but the watchdog times out before the executor commits — the + * run is treated as a failed startup rather than a successful cancel. + */ fail(error: Error): Partial { - if (this.status !== RunStatus.InProgress && this.status !== RunStatus.Queued) { + if ( + this.status !== RunStatus.InProgress && + this.status !== RunStatus.Queued && + this.status !== RunStatus.Cancelling + ) { throw new InvalidRunStateTransitionError(this.status, RunStatus.Failed); } this.status = RunStatus.Failed; @@ -149,7 +159,6 @@ export class RunBuilder { failedAt: this.failedAt ?? null, usage: this.usage ?? null, metadata: this.metadata, - output: this.output, config: this.config, }; } diff --git a/tegg/core/agent-runtime/src/SSEWriter.ts b/tegg/core/agent-runtime/src/SSEWriter.ts index bfd29b4848..f4fc401695 100644 --- a/tegg/core/agent-runtime/src/SSEWriter.ts +++ b/tegg/core/agent-runtime/src/SSEWriter.ts @@ -5,6 +5,8 @@ export interface SSEWriter { /** Write an SSE event with the given name and JSON-serializable data. */ writeEvent(event: string, data: unknown): void; + /** Write an SSE comment (e.g. `: keepalive`). Used for heartbeat signals. */ + writeComment(text: string): void; /** Whether the underlying connection has been closed. */ readonly closed: boolean; /** End the SSE stream. */ diff --git a/tegg/core/agent-runtime/test/AgentRuntime.metadata.test.ts b/tegg/core/agent-runtime/test/AgentRuntime.metadata.test.ts new file mode 100644 index 0000000000..cf982efbb2 --- /dev/null +++ b/tegg/core/agent-runtime/test/AgentRuntime.metadata.test.ts @@ -0,0 +1,246 @@ +import assert from 'node:assert'; + +import type { AgentMessage, CreateRunInput, StreamEvent } from '@eggjs/tegg-types/agent-runtime'; +import { AgentObjectType, RunStatus } from '@eggjs/tegg-types/agent-runtime'; +import { describe, it, beforeEach, afterEach } from 'vitest'; + +import { AgentRuntime } from '../src/AgentRuntime.ts'; +import type { AgentExecutor, AgentRuntimeOptions } from '../src/AgentRuntime.ts'; +import { OSSAgentStore } from '../src/OSSAgentStore.ts'; +import type { SSEWriter } from '../src/SSEWriter.ts'; +import { MapStorageClient } from './helpers.ts'; + +class MockSSEWriter implements SSEWriter { + events: Array<{ event: string; data: unknown }> = []; + comments: string[] = []; + closed = false; + private closeCallbacks: Array<() => void> = []; + + writeEvent(event: string, data: unknown): void { + this.events.push({ event, data }); + } + + writeComment(text: string): void { + this.comments.push(text); + } + + end(): void { + this.closed = true; + } + + onClose(callback: () => void): void { + this.closeCallbacks.push(callback); + } + + simulateClose(): void { + this.closed = true; + for (const cb of this.closeCallbacks) cb(); + } +} + +describe('test/AgentRuntime.metadata.test.ts', () => { + let runtime: AgentRuntime; + let store: OSSAgentStore; + let client: MapStorageClient; + let executor: AgentExecutor; + + beforeEach(() => { + client = new MapStorageClient(); + store = new OSSAgentStore({ client }); + executor = { + async *execRun(input: CreateRunInput): AsyncGenerator { + const messages = input.input.messages; + yield { + type: 'assistant', + message: { + role: 'assistant', + content: [{ type: 'text', text: `Hello ${messages.length} messages` }], + }, + }; + }, + }; + runtime = new AgentRuntime({ + executor, + store, + logger: { + error() { + /* noop */ + }, + } as unknown as AgentRuntimeOptions['logger'], + }); + }); + + afterEach(async () => { + await runtime.destroy(); + }); + + describe('createThread', () => { + it('should accept metadata and persist it on the thread record', async () => { + const meta = { agentName: 'foo', traceId: 't-1' }; + const result = await runtime.createThread({ metadata: meta }); + assert.equal(result.object, AgentObjectType.Thread); + assert.deepStrictEqual(result.metadata, meta); + + const stored = await store.getThread(result.id); + assert.deepStrictEqual(stored.metadata, meta); + }); + + it('should default to empty metadata when not provided', async () => { + const result = await runtime.createThread(); + assert.deepStrictEqual(result.metadata, {}); + }); + }); + + describe('syncRun metadata handling', () => { + it('should store metadata on the run and initialize an auto-created thread', async () => { + const metadata = { + bizId: 'order_123', + nested: { source: 'customer_service' }, + tags: ['vip', 7], + enabled: true, + nullable: null, + }; + const result = await runtime.syncRun({ + input: { messages: [{ role: 'user', content: 'Hi' }] }, + metadata, + }); + + // The run record keeps the metadata verbatim... + assert.deepStrictEqual(result.metadata, metadata); + assert.equal(result.status, RunStatus.Completed); + // ...and the auto-created thread is initialized with the same metadata. + const thread = await store.getThread(result.threadId); + assert.deepStrictEqual(thread.metadata, metadata); + }); + + it('should shallow-merge metadata into an existing thread', async () => { + const thread = await runtime.createThread({ + metadata: { + bizId: 'order_123', + source: 'customer_service', + nested: { old: true }, + }, + }); + + const result = await runtime.syncRun({ + threadId: thread.id, + input: { messages: [{ role: 'user', content: 'Hi' }] }, + metadata: { + source: 'operator_console', + nested: { replacement: true }, + }, + }); + + // The run keeps exactly what was passed in... + assert.deepStrictEqual(result.metadata, { + source: 'operator_console', + nested: { replacement: true }, + }); + // ...while the thread metadata is shallow-merged (bizId preserved). + const stored = await store.getThread(thread.id); + assert.deepStrictEqual(stored.metadata, { + bizId: 'order_123', + source: 'operator_console', + nested: { replacement: true }, + }); + }); + + it('should leave existing thread metadata unchanged when metadata is empty or omitted', async () => { + const original = { bizId: 'order_123', source: 'customer_service' }; + const thread = await runtime.createThread({ metadata: original }); + + await runtime.syncRun({ + threadId: thread.id, + input: { messages: [{ role: 'user', content: 'First' }] }, + metadata: {}, + }); + await runtime.syncRun({ + threadId: thread.id, + input: { messages: [{ role: 'user', content: 'Second' }] }, + }); + + assert.deepStrictEqual((await store.getThread(thread.id)).metadata, original); + }); + + it('should reject invalid metadata before creating a thread or run', async () => { + const assertInvalid = async (invalid: unknown): Promise => { + await assert.rejects( + () => + runtime.syncRun({ + input: { messages: [{ role: 'user', content: 'Hi' }] }, + metadata: invalid, + } as unknown as CreateRunInput), + (err: unknown) => { + assert.equal((err as { status?: number }).status, 400); + assert.match((err as Error).message, /metadata/); + return true; + }, + ); + }; + + for (const invalid of [null, [], 'invalid', 1, true]) { + await assertInvalid(invalid); + } + assert.deepStrictEqual(client.keysWithPrefix('threads/'), []); + assert.deepStrictEqual(client.keysWithPrefix('runs/'), []); + }); + + it('should not fail run creation when persisting metadata onto an existing thread fails', async () => { + const thread = await runtime.createThread({ metadata: { a: 1 } }); + // Force the thread-side metadata write to fail. + store.updateThreadMetadata = async () => { + throw new Error('boom'); + }; + + const result = await runtime.syncRun({ + threadId: thread.id, + input: { messages: [{ role: 'user', content: 'Hi' }] }, + metadata: { b: 2 }, + }); + + // The run still completes and keeps its metadata on the run record... + assert.equal(result.status, RunStatus.Completed); + assert.deepStrictEqual(result.metadata, { b: 2 }); + // ...while the thread metadata was left untouched by the failed write. + assert.deepStrictEqual((await store.getThread(thread.id)).metadata, { a: 1 }); + }); + }); + + describe('asyncRun metadata handling', () => { + it('should store metadata on the run and persist it on the thread', async () => { + const metadata = { bizId: 'async_123' }; + const result = await runtime.asyncRun({ + input: { messages: [{ role: 'user', content: 'Hi' }] }, + metadata, + }); + await runtime.waitForPendingTasks(); + + assert.deepStrictEqual(result.metadata, metadata); + const thread = await store.getThread(result.threadId); + assert.deepStrictEqual(thread.metadata, metadata); + }); + }); + + describe('streamRun metadata handling', () => { + it('should store metadata on the run and persist it on the thread', async () => { + const meta = { agentName: 'stream', source: 'sse' }; + const writer = new MockSSEWriter(); + + await runtime.streamRun( + { + input: { messages: [{ role: 'user', content: 'Hi' }] }, + metadata: meta, + }, + writer, + ); + + const runCreatedEvent = writer.events.find((e) => e.event === 'run_created')!.data as StreamEvent; + const runId = (runCreatedEvent.data as { runId: string }).runId; + const threadId = (runCreatedEvent.data as { threadId: string }).threadId; + + const run = await store.getRun(runId); + assert.deepStrictEqual(run.metadata, meta); + assert.deepStrictEqual((await store.getThread(threadId)).metadata, meta); + }); + }); +}); diff --git a/tegg/core/agent-runtime/test/AgentRuntime.test.ts b/tegg/core/agent-runtime/test/AgentRuntime.test.ts index 2154467ca7..3bd0459d5e 100644 --- a/tegg/core/agent-runtime/test/AgentRuntime.test.ts +++ b/tegg/core/agent-runtime/test/AgentRuntime.test.ts @@ -1,16 +1,20 @@ import assert from 'node:assert'; import { setTimeout } from 'node:timers/promises'; +import type { + RunRecord, + CreateRunInput, + AgentMessage, + SDKResultMessage, + StreamEvent, +} from '@eggjs/tegg-types/agent-runtime'; import { RunStatus, - AgentSSEEvent, AgentObjectType, - MessageRole, - MessageStatus, - ContentBlockType, + AgentNotFoundError, + AgentConflictError, + AgentTimeoutError, } from '@eggjs/tegg-types/agent-runtime'; -import type { RunRecord, RunObject, CreateRunInput, AgentStreamMessage } from '@eggjs/tegg-types/agent-runtime'; -import { AgentNotFoundError, AgentConflictError } from '@eggjs/tegg-types/agent-runtime'; import { describe, it, beforeEach, afterEach } from 'vitest'; import { AgentRuntime } from '../src/AgentRuntime.ts'; @@ -21,6 +25,7 @@ import { MapStorageClient } from './helpers.ts'; class MockSSEWriter implements SSEWriter { events: Array<{ event: string; data: unknown }> = []; + comments: string[] = []; closed = false; private closeCallbacks: Array<() => void> = []; @@ -28,6 +33,10 @@ class MockSSEWriter implements SSEWriter { this.events.push({ event, data }); } + writeComment(text: string): void { + this.comments.push(text); + } + end(): void { this.closed = true; } @@ -57,8 +66,8 @@ async function waitForRunStatus( throw new Error(`Run ${runId} did not reach status '${expectedStatus}' within ${timeoutMs}ms`); } -function createSlowExecRun(chunks: AgentStreamMessage[], onYielded?: () => void): AgentExecutor['execRun'] { - return async function* (_input: CreateRunInput, signal?: AbortSignal): AsyncGenerator { +function createSlowExecRun(chunks: AgentMessage[], onYielded?: () => void): AgentExecutor['execRun'] { + return async function* (_input: CreateRunInput, signal?: AbortSignal): AsyncGenerator { for (const chunk of chunks) { yield chunk; } @@ -79,11 +88,8 @@ function createSlowExecRun(chunks: AgentStreamMessage[], onYielded?: () => void) }; } -function createBlockingExecRun( - resolveRef: { resolve?: () => void }, - chunks: AgentStreamMessage[], -): AgentExecutor['execRun'] { - return async function* (_input: CreateRunInput, signal?: AbortSignal): AsyncGenerator { +function createBlockingExecRun(resolveRef: { resolve?: () => void }, chunks: AgentMessage[]): AgentExecutor['execRun'] { + return async function* (_input: CreateRunInput, signal?: AbortSignal): AsyncGenerator { await new Promise((resolve, reject) => { resolveRef.resolve = resolve; if (signal) { @@ -104,23 +110,29 @@ describe('test/AgentRuntime.test.ts', () => { beforeEach(() => { store = new OSSAgentStore({ client: new MapStorageClient() }); executor = { - async *execRun(input: CreateRunInput): AsyncGenerator { + async *execRun(input: CreateRunInput): AsyncGenerator { const messages = input.input.messages; yield { + type: 'assistant', message: { - role: MessageRole.Assistant, + role: 'assistant', content: [{ type: 'text', text: `Hello ${messages.length} messages` }], }, }; yield { - usage: { promptTokens: 10, completionTokens: 5 }, - }; + type: 'result', + subtype: 'success', + usage: { input_tokens: 10, output_tokens: 5 }, + } as SDKResultMessage; }, }; runtime = new AgentRuntime({ executor, store, logger: { + info() { + /* noop */ + }, error() { /* noop */ }, @@ -176,14 +188,6 @@ describe('test/AgentRuntime.test.ts', () => { assert.equal(result.status, RunStatus.Completed); assert(result.threadId); assert(result.threadId.startsWith('thread_')); - assert.equal(result.output!.length, 1); - assert.equal(result.output![0].object, AgentObjectType.ThreadMessage); - assert.equal(result.output![0].role, MessageRole.Assistant); - assert.equal(result.output![0].status, MessageStatus.Completed); - const content = result.output![0].content; - assert.equal(content[0].type, ContentBlockType.Text); - assert.equal(content[0].text.value, 'Hello 1 messages'); - assert(Array.isArray(content[0].text.annotations)); assert.equal(result.usage!.promptTokens, 10); assert.equal(result.usage!.completionTokens, 5); assert.equal(result.usage!.totalTokens, 15); @@ -220,9 +224,10 @@ describe('test/AgentRuntime.test.ts', () => { }); const updated = await runtime.getThread(thread.id); + // Should have: user input message + assistant message + result message assert.equal(updated.messages.length, 2); - assert.equal(updated.messages[0].role, MessageRole.User); - assert.equal(updated.messages[1].role, MessageRole.Assistant); + assert.equal(updated.messages[0].type, 'user'); + assert.equal(updated.messages[1].type, 'assistant'); }); it('should auto-create thread and append messages when threadId not provided', async () => { @@ -234,12 +239,60 @@ describe('test/AgentRuntime.test.ts', () => { const thread = await runtime.getThread(result.threadId); assert.equal(thread.messages.length, 2); - assert.equal(thread.messages[0].role, MessageRole.User); - assert.equal(thread.messages[1].role, MessageRole.Assistant); + assert.equal(thread.messages[0].type, 'user'); + assert.equal(thread.messages[1].type, 'assistant'); + }); + + it('should set isResume=false when no threadId provided (auto-create)', async () => { + let capturedInput: CreateRunInput | undefined; + executor.execRun = async function* (input: CreateRunInput): AsyncGenerator { + capturedInput = input; + yield { type: 'assistant', message: { role: 'assistant', content: [{ type: 'text', text: 'hi' }] } }; + }; + + await runtime.syncRun({ + input: { messages: [{ role: 'user', content: 'Hi' }] }, + }); + assert.equal(capturedInput!.isResume, false); + }); + + it('should set isResume=false when threadId provided but thread has no messages', async () => { + let capturedInput: CreateRunInput | undefined; + executor.execRun = async function* (input: CreateRunInput): AsyncGenerator { + capturedInput = input; + yield { type: 'assistant', message: { role: 'assistant', content: [{ type: 'text', text: 'hi' }] } }; + }; + + const thread = await runtime.createThread(); + await runtime.syncRun({ + threadId: thread.id, + input: { messages: [{ role: 'user', content: 'Hi' }] }, + }); + assert.equal(capturedInput!.isResume, false); + }); + + it('should set isResume=true when threadId provided and thread has messages', async () => { + let capturedInput: CreateRunInput | undefined; + executor.execRun = async function* (input: CreateRunInput): AsyncGenerator { + capturedInput = input; + yield { type: 'assistant', message: { role: 'assistant', content: [{ type: 'text', text: 'hi' }] } }; + }; + + // First run creates thread with messages + const result = await runtime.syncRun({ + input: { messages: [{ role: 'user', content: 'Hi' }] }, + }); + + // Second run on the same thread — now it has history + await runtime.syncRun({ + threadId: result.threadId, + input: { messages: [{ role: 'user', content: 'Hello again' }] }, + }); + assert.equal(capturedInput!.isResume, true); }); it('should not throw when store.updateRun fails in catch block', async () => { - executor.execRun = async function* (): AsyncGenerator { + executor.execRun = async function* (): AsyncGenerator { throw new Error('exec failed'); }; @@ -262,6 +315,63 @@ describe('test/AgentRuntime.test.ts', () => { }, ); }); + + it('should persist the partial transcript and mark Failed when execRun throws after committing', async () => { + const thread = await runtime.createThread(); + executor.execRun = async function* (): AsyncGenerator { + yield { type: 'assistant', message: { role: 'assistant', content: [{ type: 'text', text: 'partial' }] } }; + throw new Error('stream terminated'); + }; + + await assert.rejects( + () => runtime.syncRun({ threadId: thread.id, input: { messages: [{ role: 'user', content: 'Hi' }] } }), + (err: unknown) => { + assert(err instanceof Error); + assert.equal(err.message, 'stream terminated'); + return true; + }, + ); + + // The failed turn must not vanish: user + committed assistant are persisted. + const persisted = await store.getThread(thread.id); + assert.equal(persisted.messages.length, 2); + assert.equal(persisted.messages[0].type, 'user'); + assert.equal(persisted.messages[1].type, 'assistant'); + }); + + it('should not duplicate thread messages when updateRun fails after a successful append', async () => { + const thread = await runtime.createThread(); + // Executor finishes normally (commits), so the success path appends, then + // the completing updateRun throws — routing into the catch block. + executor.execRun = async function* (): AsyncGenerator { + yield { type: 'assistant', message: { role: 'assistant', content: [{ type: 'text', text: 'done' }] } }; + yield { + type: 'result', + subtype: 'success', + usage: { input_tokens: 1, output_tokens: 1 }, + } as SDKResultMessage; + }; + // Fail the 2nd updateRun (1=rb.start, 2=rb.complete) — i.e. right after + // the success-path appendMessages has already written the transcript. + let calls = 0; + const origUpdateRun = store.updateRun.bind(store); + store.updateRun = async (runId: string, updates: Partial) => { + calls++; + if (calls === 2) throw new Error('updateRun failed'); + return origUpdateRun(runId, updates); + }; + + await assert.rejects(() => + runtime.syncRun({ threadId: thread.id, input: { messages: [{ role: 'user', content: 'Hi' }] } }), + ); + + // Messages must be appended exactly once — the catch-block persist must be + // a no-op because the success path already persisted (no duplicate history). + const persisted = await store.getThread(thread.id); + assert.equal(persisted.messages.length, 2); + assert.equal(persisted.messages[0].type, 'user'); + assert.equal(persisted.messages[1].type, 'assistant'); + }); }); describe('asyncRun', () => { @@ -285,8 +395,6 @@ describe('test/AgentRuntime.test.ts', () => { const run = await store.getRun(result.id); assert.equal(run.status, RunStatus.Completed); - const outputContent = run.output![0].content; - assert.equal(outputContent[0].text.value, 'Hello 1 messages'); }); it('should auto-create thread and append messages when threadId not provided', async () => { @@ -299,8 +407,8 @@ describe('test/AgentRuntime.test.ts', () => { const thread = await store.getThread(result.threadId); assert.equal(thread.messages.length, 2); - assert.equal(thread.messages[0].role, MessageRole.User); - assert.equal(thread.messages[1].role, MessageRole.Assistant); + assert.equal(thread.messages[0].type, 'user'); + assert.equal(thread.messages[1].type, 'assistant'); }); it('should pass metadata through to store and return it', async () => { @@ -316,100 +424,317 @@ describe('test/AgentRuntime.test.ts', () => { const run = await store.getRun(result.id); assert.deepStrictEqual(run.metadata, meta); }); + + it('should persist the partial transcript and mark Failed when execRun throws after committing', async () => { + executor.execRun = async function* (): AsyncGenerator { + yield { type: 'assistant', message: { role: 'assistant', content: [{ type: 'text', text: 'partial' }] } }; + throw new Error('stream terminated'); + }; + + const result = await runtime.asyncRun({ input: { messages: [{ role: 'user', content: 'Hi' }] } }); + await runtime.waitForPendingTasks(); + + const run = await store.getRun(result.id); + assert.equal(run.status, RunStatus.Failed); + + const thread = await store.getThread(result.threadId!); + assert.equal(thread.messages.length, 2); + assert.equal(thread.messages[0].type, 'user'); + assert.equal(thread.messages[1].type, 'assistant'); + }); }); describe('streamRun', () => { - it('should emit correct SSE event sequence for normal flow', async () => { + it('should emit StreamEvent sequence: run_created, message events, done', async () => { const writer = new MockSSEWriter(); await runtime.streamRun({ input: { messages: [{ role: 'user', content: 'Hi' }] } }, writer); - const eventNames = writer.events.map((e) => e.event); - assert(eventNames.includes(AgentSSEEvent.ThreadRunCreated)); - assert(eventNames.includes(AgentSSEEvent.ThreadRunInProgress)); - assert(eventNames.includes(AgentSSEEvent.ThreadMessageCreated)); - assert(eventNames.includes(AgentSSEEvent.ThreadMessageDelta)); - assert(eventNames.includes(AgentSSEEvent.ThreadMessageCompleted)); - assert(eventNames.includes(AgentSSEEvent.ThreadRunCompleted)); - assert(eventNames.includes(AgentSSEEvent.Done)); - assert(writer.closed); + const eventTypes = writer.events.map((e) => e.event); + assert(eventTypes.includes('run_created'), 'should have run_created event'); + assert(eventTypes.includes('done'), 'should have done event'); + assert(writer.closed, 'writer should be closed'); + + // Verify order: run_created comes first, done comes last + const createdIdx = eventTypes.indexOf('run_created'); + const doneIdx = eventTypes.indexOf('done'); + assert(createdIdx === 0, 'run_created should be first'); + assert(doneIdx === eventTypes.length - 1, 'done should be last'); + + // Verify StreamEvent format: { seq, type, data, ts } + for (const ev of writer.events) { + const streamEvent = ev.data as StreamEvent; + assert(typeof streamEvent.seq === 'number', 'seq should be a number'); + assert(typeof streamEvent.type === 'string', 'type should be a string'); + assert(typeof streamEvent.ts === 'number', 'ts should be a number'); + assert('data' in streamEvent, 'should have data field'); + } + + // Verify sequential seq numbers + const seqs = writer.events.map((e) => (e.data as StreamEvent).seq); + for (let i = 1; i < seqs.length; i++) { + assert(seqs[i] === seqs[i - 1] + 1, `seq should be sequential: ${seqs[i]} after ${seqs[i - 1]}`); + } + + // Verify run_created data has runId and threadId + const runCreatedEvent = writer.events[0].data as StreamEvent; + assert((runCreatedEvent.data as any).runId, 'run_created should have runId'); + assert((runCreatedEvent.data as any).threadId, 'run_created should have threadId'); - // Verify order: created < in_progress < message.created < delta < message.completed < run.completed < done - const createdIdx = eventNames.indexOf(AgentSSEEvent.ThreadRunCreated); - const progressIdx = eventNames.indexOf(AgentSSEEvent.ThreadRunInProgress); - const msgCreatedIdx = eventNames.indexOf(AgentSSEEvent.ThreadMessageCreated); - const deltaIdx = eventNames.indexOf(AgentSSEEvent.ThreadMessageDelta); - const msgCompletedIdx = eventNames.indexOf(AgentSSEEvent.ThreadMessageCompleted); - const runCompletedIdx = eventNames.indexOf(AgentSSEEvent.ThreadRunCompleted); - const doneIdx = eventNames.indexOf(AgentSSEEvent.Done); - assert(createdIdx < progressIdx); - assert(progressIdx < msgCreatedIdx); - assert(msgCreatedIdx < deltaIdx); - assert(deltaIdx < msgCompletedIdx); - assert(msgCompletedIdx < runCompletedIdx); - assert(runCompletedIdx < doneIdx); - - // Verify messages persisted to thread (consistent with syncRun/asyncRun tests) - const runCreatedEvent = writer.events.find((e) => e.event === AgentSSEEvent.ThreadRunCreated); - const threadId = (runCreatedEvent!.data as RunObject).threadId; + // Verify messages persisted to thread + const threadId = (runCreatedEvent.data as any).threadId; const thread = await runtime.getThread(threadId); assert.equal(thread.messages.length, 2); - assert.equal(thread.messages[0]['role'], MessageRole.User); - assert.equal(thread.messages[1]['role'], MessageRole.Assistant); + assert.equal(thread.messages[0].type, 'user'); + assert.equal(thread.messages[1].type, 'assistant'); }); - it('should emit cancelled event on client disconnect', async () => { + it('should continue background execution on client disconnect', async () => { let resolveYielded!: () => void; const yieldedPromise = new Promise((r) => { resolveYielded = r; }); - executor.execRun = async function* ( - _input: CreateRunInput, - signal?: AbortSignal, - ): AsyncGenerator { - yield { message: { role: MessageRole.Assistant, content: [{ type: 'text', text: 'start' }] } }; + executor.execRun = async function* (_input: CreateRunInput, signal?: AbortSignal): AsyncGenerator { + yield { type: 'assistant', message: { role: 'assistant', content: [{ type: 'text', text: 'start' }] } }; resolveYielded(); - await new Promise((resolve) => { - const timer = globalThis.setTimeout(resolve, 5000); - if (signal) { - signal.addEventListener( - 'abort', - () => { - clearTimeout(timer); - resolve(); - }, - { once: true }, - ); - } - }); + // Simulate more work after client disconnects + await setTimeout(50); + if (!signal?.aborted) { + yield { type: 'assistant', message: { role: 'assistant', content: [{ type: 'text', text: ' end' }] } }; + } }; const writer = new MockSSEWriter(); - const streamPromise = runtime.streamRun({ input: { messages: [{ role: 'user', content: 'Hi' }] } }, writer); + // Wait for first chunk to be yielded, then give the writer time to receive it await yieldedPromise; + await setTimeout(50); writer.simulateClose(); - await streamPromise; - const eventNames = writer.events.map((e) => e.event); - assert(eventNames.includes(AgentSSEEvent.ThreadRunCreated)); - assert(eventNames.includes(AgentSSEEvent.ThreadRunInProgress)); + // Writer should have received at least run_created before disconnect + assert(writer.events.length >= 1, 'should have at least run_created event'); + + // Background task should complete — wait for it + await runtime.waitForPendingTasks(); + + // Verify the run completed in the store (not cancelled) + const runCreatedEvent = writer.events[0].data as StreamEvent; + const runId = (runCreatedEvent.data as any).runId; + const run = await runtime.getRun(runId); + assert.equal(run.status, RunStatus.Completed, 'run should complete despite client disconnect'); + }); + + it('should forward SDK message types as event types', async () => { + executor.execRun = async function* (): AsyncGenerator { + yield { type: 'system', subtype: 'init', session_id: 'sess-1' }; + yield { type: 'stream_event', event: { type: 'content_block_delta' } }; + yield { type: 'assistant', message: { role: 'assistant', content: [{ type: 'text', text: 'Hello' }] } }; + }; + + const writer = new MockSSEWriter(); + await runtime.streamRun({ input: { messages: [{ role: 'user', content: 'Hi' }] } }, writer); + + const eventTypes = writer.events.map((e) => e.event); + assert(eventTypes.includes('system'), 'should forward system event'); + assert(eventTypes.includes('stream_event'), 'should forward stream_event'); + assert(eventTypes.includes('assistant'), 'should forward assistant event'); + }); + + it('should pass through SDK message directly as event data', async () => { + const sdkMsg: AgentMessage = { + type: 'assistant', + message: { role: 'assistant', content: [{ type: 'text', text: 'raw hello' }] }, + }; + executor.execRun = async function* (): AsyncGenerator { + yield sdkMsg; + }; + + const writer = new MockSSEWriter(); + await runtime.streamRun({ input: { messages: [{ role: 'user', content: 'Hi' }] } }, writer); + + const assistantEvent = writer.events.find((e) => e.event === 'assistant'); + assert.ok(assistantEvent); + const streamEvent = assistantEvent.data as StreamEvent; + assert.deepStrictEqual(streamEvent.data, sdkMsg, 'should pass SDK message directly as data'); }); - it('should emit failed event when execRun throws', async () => { - executor.execRun = async function* (): AsyncGenerator { + it('should use "message" as default event type when type is not set', async () => { + executor.execRun = async function* (): AsyncGenerator { + yield { type: '', message: { content: 'Hello' } } as unknown as AgentMessage; + }; + + const writer = new MockSSEWriter(); + await runtime.streamRun({ input: { messages: [{ role: 'user', content: 'Hi' }] } }, writer); + + const eventTypes = writer.events.map((e) => e.event); + assert(eventTypes.includes('message'), 'should fallback to "message" for empty type'); + }); + + it('should emit error event when execRun throws', async () => { + executor.execRun = async function* (): AsyncGenerator { throw new Error('model unavailable'); }; const writer = new MockSSEWriter(); await runtime.streamRun({ input: { messages: [{ role: 'user', content: 'Hi' }] } }, writer); - const eventNames = writer.events.map((e) => e.event); - assert(eventNames.includes(AgentSSEEvent.ThreadRunFailed)); - assert(eventNames.includes(AgentSSEEvent.Done)); + const errorEvent = writer.events.find((e) => e.event === 'error'); + assert.ok(errorEvent, 'should have error event'); + const streamEvent = errorEvent.data as StreamEvent; + assert.equal((streamEvent.data as any).message, 'model unavailable'); assert(writer.closed); + + // Verify run is marked as failed in store + const runCreatedEvent = writer.events[0].data as StreamEvent; + const runId = (runCreatedEvent.data as any).runId; + const run = await runtime.getRun(runId); + assert.equal(run.status, RunStatus.Failed); + }); + + it('should persist the partial transcript when execRun throws after committing', async () => { + executor.execRun = async function* (): AsyncGenerator { + yield { type: 'assistant', message: { role: 'assistant', content: [{ type: 'text', text: 'partial' }] } }; + throw new Error('stream terminated'); + }; + + const writer = new MockSSEWriter(); + await runtime.streamRun({ input: { messages: [{ role: 'user', content: 'Hi' }] } }, writer); + + const runCreatedEvent = writer.events[0].data as StreamEvent; + const runId = (runCreatedEvent.data as any).runId; + const threadId = (runCreatedEvent.data as any).threadId; + + const run = await runtime.getRun(runId); + assert.equal(run.status, RunStatus.Failed); + + // The failed turn must not vanish from thread history. + const thread = await runtime.getThread(threadId); + assert.equal(thread.messages.length, 2); + assert.equal(thread.messages[0].type, 'user'); + assert.equal(thread.messages[1].type, 'assistant'); + }); + + it('should persist usage to store on completion', async () => { + executor.execRun = async function* (): AsyncGenerator { + yield { type: 'assistant', message: { content: 'Hi' } }; + yield { + type: 'result', + subtype: 'success', + usage: { input_tokens: 10, output_tokens: 8 }, + } as SDKResultMessage; + }; + + const writer = new MockSSEWriter(); + await runtime.streamRun({ input: { messages: [{ role: 'user', content: 'Hi' }] } }, writer); + + const runCreatedEvent = writer.events[0].data as StreamEvent; + const runId = (runCreatedEvent.data as any).runId; + const run = await runtime.getRun(runId); + assert.equal(run.status, RunStatus.Completed); + assert.equal(run.usage!.promptTokens, 10); + assert.equal(run.usage!.completionTokens, 8); + assert.equal(run.usage!.totalTokens, 18); + }); + }); + + describe('getRunStream', () => { + it('should replay all events on reconnect with lastSeq=0', async () => { + const writer1 = new MockSSEWriter(); + await runtime.streamRun({ input: { messages: [{ role: 'user', content: 'Hi' }] } }, writer1); + + // Get runId from the first event + const runCreatedEvent = writer1.events[0].data as StreamEvent; + const runId = (runCreatedEvent.data as any).runId; + + // Reconnect and get all events + const writer2 = new MockSSEWriter(); + await runtime.getRunStream(runId, writer2, 0); + + // Should get the same events + assert.equal(writer2.events.length, writer1.events.length); + for (let i = 0; i < writer1.events.length; i++) { + const se1 = writer1.events[i].data as StreamEvent; + const se2 = writer2.events[i].data as StreamEvent; + assert.equal(se1.seq, se2.seq); + assert.equal(se1.type, se2.type); + } + }); + + it('should replay only events after lastSeq', async () => { + const writer1 = new MockSSEWriter(); + await runtime.streamRun({ input: { messages: [{ role: 'user', content: 'Hi' }] } }, writer1); + + const runCreatedEvent = writer1.events[0].data as StreamEvent; + const runId = (runCreatedEvent.data as any).runId; + const totalEvents = writer1.events.length; + + // Reconnect from seq 2 (skip first 2 events) + const writer2 = new MockSSEWriter(); + await runtime.getRunStream(runId, writer2, 2); + + assert.equal(writer2.events.length, totalEvents - 2); + const firstReplayedSeq = (writer2.events[0].data as StreamEvent).seq; + assert.equal(firstReplayedSeq, 3); + }); + + it('should throw AgentNotFoundError for unknown runId', async () => { + const writer = new MockSSEWriter(); + await assert.rejects( + () => runtime.getRunStream('run_nonexistent', writer), + (err: unknown) => { + assert(err instanceof AgentNotFoundError); + return true; + }, + ); + }); + + it('should stream real-time events during reconnect to running task', async () => { + let resolveExec!: () => void; + const execPromise = new Promise((r) => { + resolveExec = r; + }); + + executor.execRun = async function* (_input: CreateRunInput, signal?: AbortSignal): AsyncGenerator { + yield { type: 'assistant', message: { content: 'chunk1' } }; + // Wait for reconnect to happen + await execPromise; + if (!signal?.aborted) { + yield { type: 'assistant', message: { content: 'chunk2' } }; + } + }; + + // Start streaming and disconnect immediately after first events + const writer1 = new MockSSEWriter(); + const streamPromise = runtime.streamRun({ input: { messages: [{ role: 'user', content: 'Hi' }] } }, writer1); + + // Wait a bit for the first chunk to be buffered + await setTimeout(50); + writer1.simulateClose(); + await streamPromise; + + const runId = ((writer1.events[0].data as StreamEvent).data as any).runId; + + // Reconnect — should get replayed events then real-time + const writer2 = new MockSSEWriter(); + const reconnectPromise = runtime.getRunStream(runId, writer2, 0); + + // Let the background task continue + await setTimeout(20); + resolveExec(); + + await runtime.waitForPendingTasks(); + // Give streamEventsToWriter time to process the final events + await setTimeout(20); + + // Close writer2 to end reconnect + writer2.simulateClose(); + await reconnectPromise; + + // writer2 should have received all events including chunk2 and done + const eventTypes = writer2.events.map((e) => e.event); + assert(eventTypes.includes('done'), 'reconnected stream should receive done event'); }); }); @@ -438,12 +763,35 @@ describe('test/AgentRuntime.test.ts', () => { }); }); + describe('getLatestRunId', () => { + it('should resolve the most recent run id for a thread', async () => { + const created = await runtime.createThread(); + const run = await runtime.syncRun({ + threadId: created.id, + input: { messages: [{ role: 'user', content: 'Hi' }] }, + }); + + // latestRunId is recorded in the background; drain before asserting. + await store.awaitPendingWrites(); + const result = await runtime.getLatestRunId(created.id); + assert.deepStrictEqual(result, { threadId: created.id, runId: run.id }); + }); + + it('should return runId null for a thread with no run', async () => { + const created = await runtime.createThread(); + const result = await runtime.getLatestRunId(created.id); + assert.deepStrictEqual(result, { threadId: created.id, runId: null }); + }); + + it('should reject for a non-existent thread', async () => { + await assert.rejects(() => runtime.getLatestRunId('thread_non_existent')); + }); + }); + describe('cancelRun', () => { it('should cancel a run', async () => { executor.execRun = createSlowExecRun([ - { - message: { role: MessageRole.Assistant, content: [{ type: 'text', text: 'start' }] }, - }, + { type: 'assistant', message: { role: 'assistant', content: [{ type: 'text', text: 'start' }] } }, ]); const result = await runtime.asyncRun({ @@ -464,9 +812,7 @@ describe('test/AgentRuntime.test.ts', () => { it('should write cancelling then cancelled to store', async () => { executor.execRun = createSlowExecRun([ - { - message: { role: MessageRole.Assistant, content: [{ type: 'text', text: 'start' }] }, - }, + { type: 'assistant', message: { role: 'assistant', content: [{ type: 'text', text: 'start' }] } }, ]); const statusHistory: string[] = []; @@ -511,12 +857,12 @@ describe('test/AgentRuntime.test.ts', () => { it('should not overwrite cancelling status with completed (cross-worker scenario)', async () => { const resolveRef: { resolve?: () => void } = {}; executor.execRun = createBlockingExecRun(resolveRef, [ + { type: 'assistant', message: { role: 'assistant', content: [{ type: 'text', text: 'done' }] } }, { - message: { role: MessageRole.Assistant, content: [{ type: 'text', text: 'done' }] }, - }, - { - usage: { promptTokens: 1, completionTokens: 1 }, - }, + type: 'result', + subtype: 'success', + usage: { input_tokens: 1, output_tokens: 1 }, + } as SDKResultMessage, ]); const result = await runtime.asyncRun({ @@ -534,13 +880,268 @@ describe('test/AgentRuntime.test.ts', () => { assert.equal(run.status, RunStatus.Cancelling); }); + it('should hold until the executor commits before aborting and persisting', async () => { + // Simulates the Claude Code SDK startup window: the executor emits a + // non-committing `system/init` right away and then a committing + // `assistant` message later. cancelRun is called in between and must + // wait for the assistant chunk before it aborts, otherwise it would + // persist a user message against a session file that does not yet + // exist on disk. + let resolveGate!: () => void; + const gate = new Promise((r) => { + resolveGate = r; + }); + executor.execRun = async function* (_input: CreateRunInput, signal?: AbortSignal): AsyncGenerator { + yield { type: 'system', subtype: 'init' } as AgentMessage; + await gate; + yield { type: 'assistant', message: { role: 'assistant', content: [{ type: 'text', text: 'now committed' }] } }; + await new Promise((_resolve, reject) => { + if (signal?.aborted) { + reject(new Error('aborted')); + return; + } + signal?.addEventListener('abort', () => reject(new Error('aborted')), { once: true }); + }); + }; + + const thread = await runtime.createThread(); + const result = await runtime.asyncRun({ + threadId: thread.id, + input: { messages: [{ role: 'user', content: 'Hi' }] }, + }); + await waitForRunStatus(store, result.id, RunStatus.InProgress); + + // Give the executor time to yield the first system message; the task + // must still be considered uncommitted because type === 'system'. + await setTimeout(30); + + const cancelPromise = runtime.cancelRun(result.id); + let cancelResolved = false; + cancelPromise.then( + () => { + cancelResolved = true; + }, + () => { + cancelResolved = true; + }, + ); + + // Hold for a bit — cancel should not resolve while we are blocked in + // the gate, because no committing message has been yielded yet. + await setTimeout(100); + assert.equal(cancelResolved, false, 'cancelRun must block until executor commits'); + + // Release the gate so the executor yields a non-system message. + resolveGate(); + const cancelled = await cancelPromise; + assert.equal(cancelled.status, RunStatus.Cancelled); + + // Thread should contain both the user input and the partial assistant + // reply — i.e. it is in sync with what the SDK jsonl would contain. + const updated = await runtime.getThread(thread.id); + assert.equal(updated.messages.length, 2); + assert.equal(updated.messages[0].type, 'user'); + assert.equal(updated.messages[1].type, 'assistant'); + }); + + it('should fail the run and leave the thread empty when cancel waits past the commit timeout', async () => { + // Rebuild the runtime with a short commit timeout so the watchdog + // fires before the blocking executor ever yields. + await runtime.destroy(); + const resolveRef: { resolve?: () => void } = {}; + executor.execRun = createBlockingExecRun(resolveRef, []); + runtime = new AgentRuntime({ + executor, + store, + logger: { + info() { + /* noop */ + }, + error() { + /* noop */ + }, + } as unknown as AgentRuntimeOptions['logger'], + cancelCommitTimeoutMs: 50, + }); + + const thread = await runtime.createThread(); + const result = await runtime.asyncRun({ + threadId: thread.id, + input: { messages: [{ role: 'user', content: 'Hi' }] }, + }); + await waitForRunStatus(store, result.id, RunStatus.InProgress); + + await assert.rejects( + () => runtime.cancelRun(result.id), + (err: unknown) => { + assert(err instanceof AgentTimeoutError, `expected AgentTimeoutError, got ${err}`); + return true; + }, + ); + + const persisted = await store.getRun(result.id); + assert.equal(persisted.status, RunStatus.Failed); + + const updated = await runtime.getThread(thread.id); + assert.equal(updated.messages.length, 0, 'thread must stay empty when executor never committed'); + + // Unblock the executor so runtime.destroy() in afterEach completes cleanly. + resolveRef.resolve?.(); + }); + + it('should not overwrite watchdog-set Failed with Completed when executor finishes naturally without committing', async () => { + // Regression test for a TOCTOU race on the commit-timeout path: + // 1. cancelRun's watchdog fires at cancelCommitTimeoutMs and writes Failed. + // 2. The executor does not listen to the abort signal and finishes + // naturally (no committing message ever yielded). + // 3. asyncRun's IIFE then enters its post-loop status check; without + // the Failed/Expired guard it would fall through to rb.complete() + // and overwrite the watchdog-set Failed with Completed. + await runtime.destroy(); + let resolveGate!: () => void; + const gate = new Promise((r) => { + resolveGate = r; + }); + executor = { + async *execRun(): AsyncGenerator { + yield { type: 'system', subtype: 'init' } as AgentMessage; + // Intentionally does NOT listen to the abort signal — simulates an + // executor that keeps running to natural completion after the + // runtime has already declared the run Failed. + await gate; + }, + }; + runtime = new AgentRuntime({ + executor, + store, + logger: { + info() { + /* noop */ + }, + error() { + /* noop */ + }, + } as unknown as AgentRuntimeOptions['logger'], + cancelCommitTimeoutMs: 50, + }); + + const thread = await runtime.createThread(); + const result = await runtime.asyncRun({ + threadId: thread.id, + input: { messages: [{ role: 'user', content: 'Hi' }] }, + }); + await waitForRunStatus(store, result.id, RunStatus.InProgress); + + const cancelPromise = runtime.cancelRun(result.id); + + // Wait well past the 50ms watchdog so Failed has definitely been + // written, then release the gate to let the executor finish naturally. + await setTimeout(150); + resolveGate(); + + await assert.rejects(cancelPromise, (err: unknown) => { + assert(err instanceof AgentTimeoutError, `expected AgentTimeoutError, got ${err}`); + return true; + }); + + const persisted = await store.getRun(result.id); + assert.equal( + persisted.status, + RunStatus.Failed, + 'watchdog-set Failed must not be overwritten by post-loop rb.complete', + ); + + const updated = await runtime.getThread(thread.id); + assert.equal(updated.messages.length, 0, 'thread must stay empty when executor never committed'); + }); + + it('should not hold cancelRun when the executor has already committed', async () => { + executor.execRun = createSlowExecRun([ + { type: 'system', subtype: 'init' } as AgentMessage, + { type: 'assistant', message: { role: 'assistant', content: [{ type: 'text', text: 'committed' }] } }, + ]); + + const result = await runtime.asyncRun({ + input: { messages: [{ role: 'user', content: 'Hi' }] }, + }); + await waitForRunStatus(store, result.id, RunStatus.InProgress); + // Wait long enough for the assistant chunk to land and mark committed. + await setTimeout(50); + + const start = Date.now(); + const cancelled = await runtime.cancelRun(result.id); + const elapsed = Date.now() - start; + assert.equal(cancelled.status, RunStatus.Cancelled); + assert(elapsed < 1000, `cancelRun should return quickly when already committed, took ${elapsed}ms`); + }); + + it('should consult a custom isSessionCommitted hook instead of the default heuristic', async () => { + // The executor yields three assistant messages; the hook treats only + // the third as committing. cancelRun must wait for the third. + const decisions: Array<{ text: string; committed: boolean }> = []; + let resolveGate!: () => void; + const gate = new Promise((r) => { + resolveGate = r; + }); + executor.execRun = async function* (_input: CreateRunInput, signal?: AbortSignal): AsyncGenerator { + yield { type: 'assistant', message: { role: 'assistant', content: [{ type: 'text', text: 'first' }] } }; + yield { type: 'assistant', message: { role: 'assistant', content: [{ type: 'text', text: 'second' }] } }; + await gate; + yield { type: 'assistant', message: { role: 'assistant', content: [{ type: 'text', text: 'committed' }] } }; + await new Promise((_resolve, reject) => { + if (signal?.aborted) { + reject(new Error('aborted')); + return; + } + signal?.addEventListener('abort', () => reject(new Error('aborted')), { once: true }); + }); + }; + executor.isSessionCommitted = (msg: AgentMessage) => { + const text = (msg as { message?: { content?: Array<{ text?: string }> } }).message?.content?.[0]?.text ?? ''; + const committed = text === 'committed'; + decisions.push({ text, committed }); + return committed; + }; + + const result = await runtime.asyncRun({ + input: { messages: [{ role: 'user', content: 'Hi' }] }, + }); + await waitForRunStatus(store, result.id, RunStatus.InProgress); + // Let the first two non-committing chunks land before cancelling. + await setTimeout(30); + + const cancelPromise = runtime.cancelRun(result.id); + let cancelResolved = false; + cancelPromise.then( + () => { + cancelResolved = true; + }, + () => { + cancelResolved = true; + }, + ); + + await setTimeout(80); + assert.equal(cancelResolved, false, 'cancelRun must wait for the hook to accept a message as committed'); + + resolveGate(); + const cancelled = await cancelPromise; + assert.equal(cancelled.status, RunStatus.Cancelled); + assert( + decisions.some((d) => d.committed), + 'hook must have observed the committing chunk', + ); + }); + it('should not overwrite terminal state when run completes during cancellation (TOCTOU)', async () => { const resolveRef: { resolve?: () => void } = {}; executor.execRun = createBlockingExecRun(resolveRef, [ + { type: 'assistant', message: { role: 'assistant', content: [{ type: 'text', text: 'done' }] } }, { - message: { role: MessageRole.Assistant, content: [{ type: 'text', text: 'done' }] }, - }, - { usage: { promptTokens: 1, completionTokens: 1 } }, + type: 'result', + subtype: 'success', + usage: { input_tokens: 1, output_tokens: 1 }, + } as SDKResultMessage, ]); const result = await runtime.asyncRun({ @@ -563,4 +1164,275 @@ describe('test/AgentRuntime.test.ts', () => { assert.equal(cancelResult.status, RunStatus.Completed); }); }); + + describe('abort message persistence', () => { + // Rationale: when an executor supports resuming from a session file + // (e.g. Claude CLI session), aborts leave partial state in that session. + // If the thread is NOT updated with the same partial state, any + // subsequent resume request diverges from the executor's view of history + // and can fail at executor startup. These tests pin down that abort + // writes the same messages to the thread that the executor has already + // observed. + + async function waitUntil(cond: () => boolean, timeoutMs = 2000): Promise { + const start = Date.now(); + while (!cond()) { + if (Date.now() - start > timeoutMs) throw new Error('waitUntil timeout'); + await setTimeout(10); + } + } + + it('syncRun: should persist user + partial assistant messages when aborted via signal', async () => { + let yielded = false; + executor.execRun = createSlowExecRun( + [{ type: 'assistant', message: { role: 'assistant', content: [{ type: 'text', text: 'partial' }] } }], + () => { + yielded = true; + }, + ); + + const thread = await runtime.createThread(); + const ac = new AbortController(); + const syncPromise = runtime.syncRun( + { threadId: thread.id, input: { messages: [{ role: 'user', content: 'Hi' }] } }, + ac.signal, + ); + + await waitUntil(() => yielded); + ac.abort(); + const result = await syncPromise; + assert.equal(result.threadId, thread.id); + + const updated = await runtime.getThread(thread.id); + assert.equal(updated.messages.length, 2); + assert.equal(updated.messages[0].type, 'user'); + assert.equal(updated.messages[1].type, 'assistant'); + assert.deepStrictEqual((updated.messages[1].message as { content: unknown }).content, [ + { type: 'text', text: 'partial' }, + ]); + }); + + it('syncRun: should NOT persist the user message when aborted before the executor commits', async () => { + // When an external signal aborts before the executor has yielded any + // non-system message, the executor's underlying session (e.g. the + // Claude Code SDK jsonl file) was never created on disk. Persisting + // the user message to the thread in that state would cause the next + // run to diverge from a session that does not exist, so the thread + // must be left empty instead. + const resolveRef: { resolve?: () => void } = {}; + executor.execRun = createBlockingExecRun(resolveRef, [ + { type: 'assistant', message: { role: 'assistant', content: [{ type: 'text', text: 'never' }] } }, + ]); + + const thread = await runtime.createThread(); + const ac = new AbortController(); + const syncPromise = runtime.syncRun( + { threadId: thread.id, input: { messages: [{ role: 'user', content: 'Hi' }] } }, + ac.signal, + ); + + // Give syncRun time to enter the await on the executor + await setTimeout(20); + ac.abort(); + await syncPromise; + + const updated = await runtime.getThread(thread.id); + assert.equal(updated.messages.length, 0); + }); + + it('asyncRun: should persist user + partial assistant messages when aborted via cancelRun', async () => { + let yielded = false; + executor.execRun = createSlowExecRun( + [{ type: 'assistant', message: { role: 'assistant', content: [{ type: 'text', text: 'partial' }] } }], + () => { + yielded = true; + }, + ); + + const thread = await runtime.createThread(); + const result = await runtime.asyncRun({ + threadId: thread.id, + input: { messages: [{ role: 'user', content: 'Hi' }] }, + }); + + await waitUntil(() => yielded); + await runtime.cancelRun(result.id); + + const updated = await runtime.getThread(thread.id); + assert.equal(updated.messages.length, 2); + assert.equal(updated.messages[0].type, 'user'); + assert.equal(updated.messages[1].type, 'assistant'); + }); + + it('streamRun: should persist user + partial assistant messages when aborted via cancelRun', async () => { + let yielded = false; + executor.execRun = createSlowExecRun( + [{ type: 'assistant', message: { role: 'assistant', content: [{ type: 'text', text: 'partial' }] } }], + () => { + yielded = true; + }, + ); + + const thread = await runtime.createThread(); + const writer = new MockSSEWriter(); + const streamPromise = runtime.streamRun( + { threadId: thread.id, input: { messages: [{ role: 'user', content: 'Hi' }] } }, + writer, + ); + + await waitUntil(() => yielded); + // Wait for run_created event to land so we can read the runId + await waitUntil(() => writer.events.some((e) => e.event === 'run_created')); + const runCreated = writer.events.find((e) => e.event === 'run_created')!; + const runId = ((runCreated.data as StreamEvent).data as { runId: string }).runId; + + await runtime.cancelRun(runId); + writer.simulateClose(); + await streamPromise; + + const updated = await runtime.getThread(thread.id); + assert.equal(updated.messages.length, 2); + assert.equal(updated.messages[0].type, 'user'); + assert.equal(updated.messages[1].type, 'assistant'); + }); + + it('should set isResume=true on the next run after an abort (regression: abort+continue)', async () => { + // Reproduces the bug: user aborts turn N, then sends turn N+1. Before the fix, turn N's + // messages were never persisted, so the second call's isResume depended only on earlier + // completed turns — and any divergence from the executor's session caused subsequent + // resume attempts to fail. + let yielded = false; + executor.execRun = createSlowExecRun( + [{ type: 'assistant', message: { role: 'assistant', content: [{ type: 'text', text: 'partial' }] } }], + () => { + yielded = true; + }, + ); + + const thread = await runtime.createThread(); + const ac = new AbortController(); + const firstPromise = runtime.syncRun( + { threadId: thread.id, input: { messages: [{ role: 'user', content: 'Hi' }] } }, + ac.signal, + ); + await waitUntil(() => yielded); + ac.abort(); + await firstPromise; + + let capturedInput: CreateRunInput | undefined; + executor.execRun = async function* (input: CreateRunInput): AsyncGenerator { + capturedInput = input; + yield { type: 'assistant', message: { role: 'assistant', content: [{ type: 'text', text: 'ok' }] } }; + }; + const secondResult = await runtime.syncRun({ + threadId: thread.id, + input: { messages: [{ role: 'user', content: 'continue' }] }, + }); + assert.equal(capturedInput!.isResume, true); + assert.equal(capturedInput!.input.messages[0].content, 'continue'); + assert.equal(secondResult.status, RunStatus.Completed); + }); + + it('syncRun: should finalise run status to Cancelled when aborted via external signal (no cancelRun)', async () => { + let yielded = false; + executor.execRun = createSlowExecRun( + [{ type: 'assistant', message: { role: 'assistant', content: [{ type: 'text', text: 'partial' }] } }], + () => { + yielded = true; + }, + ); + + const thread = await runtime.createThread(); + const ac = new AbortController(); + const syncPromise = runtime.syncRun( + { threadId: thread.id, input: { messages: [{ role: 'user', content: 'Hi' }] } }, + ac.signal, + ); + await waitUntil(() => yielded); + ac.abort(); + const snapshot = await syncPromise; + + // Snapshot reflects the live store state after finalisation + assert.equal(snapshot.status, RunStatus.Cancelled); + const persisted = await store.getRun(snapshot.id); + assert.equal(persisted.status, RunStatus.Cancelled); + assert(persisted.cancelledAt); + }); + + it('asyncRun: should finalise run status to Cancelled when destroy() aborts in-flight runs', async () => { + const resolveRef: { resolve?: () => void } = {}; + executor.execRun = createBlockingExecRun(resolveRef, [ + { type: 'assistant', message: { role: 'assistant', content: [{ type: 'text', text: 'never' }] } }, + ]); + + const result = await runtime.asyncRun({ + input: { messages: [{ role: 'user', content: 'Hi' }] }, + }); + await waitForRunStatus(store, result.id, RunStatus.InProgress); + + await runtime.destroy(); + + const persisted = await store.getRun(result.id); + assert.equal(persisted.status, RunStatus.Cancelled); + }); + + it('should preserve cancelling → cancelled ordering when cancelRun drives the abort', async () => { + // Regression guard for the finaliseAbortedRun addition: our in-branch + // finaliser must NOT write when cancelRun has already set `cancelling`, + // otherwise cancelRun's own `cancel()` transition would throw. + executor.execRun = createSlowExecRun([ + { type: 'assistant', message: { role: 'assistant', content: [{ type: 'text', text: 'partial' }] } }, + ]); + + const statusHistory: string[] = []; + const origUpdateRun = store.updateRun.bind(store); + store.updateRun = async (runId: string, updates: Partial) => { + if (updates.status) statusHistory.push(updates.status); + return origUpdateRun(runId, updates); + }; + + const asyncResult = await runtime.asyncRun({ + input: { messages: [{ role: 'user', content: 'Hi' }] }, + }); + await waitForRunStatus(store, asyncResult.id, RunStatus.InProgress); + statusHistory.length = 0; + + await runtime.cancelRun(asyncResult.id); + + const cancellingWrites = statusHistory.filter((s) => s === RunStatus.Cancelling).length; + const cancelledWrites = statusHistory.filter((s) => s === RunStatus.Cancelled).length; + assert.equal(cancellingWrites, 1, 'cancelling should be written exactly once (by cancelRun)'); + assert.equal(cancelledWrites, 1, 'cancelled should be written exactly once (by cancelRun)'); + }); + + it('should swallow store.appendMessages errors during abort cleanup', async () => { + let yielded = false; + executor.execRun = createSlowExecRun( + [{ type: 'assistant', message: { role: 'assistant', content: [{ type: 'text', text: 'partial' }] } }], + () => { + yielded = true; + }, + ); + + const thread = await runtime.createThread(); + const origAppend = store.appendMessages.bind(store); + store.appendMessages = async () => { + throw new Error('store down'); + }; + + const ac = new AbortController(); + const syncPromise = runtime.syncRun( + { threadId: thread.id, input: { messages: [{ role: 'user', content: 'Hi' }] } }, + ac.signal, + ); + await waitUntil(() => yielded); + ac.abort(); + // Should resolve with a snapshot instead of rejecting + const result = await syncPromise; + assert(result.id.startsWith('run_')); + + // Thread append failed — state is empty, but the abort path completed cleanly + store.appendMessages = origAppend; + }); + }); }); diff --git a/tegg/core/agent-runtime/test/AgentStoreUtils.test.ts b/tegg/core/agent-runtime/test/AgentStoreUtils.test.ts new file mode 100644 index 0000000000..9a0cab4f74 --- /dev/null +++ b/tegg/core/agent-runtime/test/AgentStoreUtils.test.ts @@ -0,0 +1,81 @@ +import assert from 'node:assert'; + +import { describe, it } from 'vitest'; + +import { TS_MAX_MS, dateBucket, reverseMs } from '../src/index.ts'; + +describe('test/AgentStoreUtils.test.ts', () => { + describe('TS_MAX_MS', () => { + it('is the 13-digit millisecond timestamp ceiling', () => { + assert.strictEqual(TS_MAX_MS, 9_999_999_999_999); + assert.strictEqual(String(TS_MAX_MS).length, 13); + const isoOfCeiling = new Date(TS_MAX_MS).toISOString(); + assert.match(isoOfCeiling, /^2286-/); + }); + }); + + describe('reverseMs', () => { + it('produces a 13-character zero-padded decimal string', () => { + assert.strictEqual(reverseMs(0), '9999999999999'); + assert.strictEqual(reverseMs(TS_MAX_MS), '0000000000000'); + assert.strictEqual(reverseMs(1), String(TS_MAX_MS - 1).padStart(13, '0')); + assert.strictEqual(reverseMs(1).length, 13); + }); + + it('matches the worked example timestamp from the PR description', () => { + const knownMs = Date.UTC(2025, 10, 13, 8, 0, 0, 0); + assert.strictEqual(knownMs, 1_763_020_800_000); + assert.strictEqual(reverseMs(knownMs), '8236979199999'); + }); + + it('throws RangeError for non-integer, negative, or out-of-range ms', () => { + assert.throws(() => reverseMs(-1), RangeError); + assert.throws(() => reverseMs(TS_MAX_MS + 1), RangeError); + assert.throws(() => reverseMs(1.5), RangeError); + assert.throws(() => reverseMs(Number.NaN), RangeError); + assert.throws(() => reverseMs(Number.POSITIVE_INFINITY), RangeError); + assert.throws(() => reverseMs(Number.NEGATIVE_INFINITY), RangeError); + }); + + it('is strictly monotonically decreasing in lex order versus the input', () => { + const pairs: Array<[number, number]> = [ + [0, 1], + [0, TS_MAX_MS], + [1, 2], + [1_000, 1_001], + [1_762_992_000_000, 1_763_078_400_000], // one day apart in November 2025 + [Date.UTC(2025, 0, 1), Date.UTC(2026, 0, 1)], + [Math.floor(Date.now() / 2), Date.now()], + [TS_MAX_MS - 1, TS_MAX_MS], + ]; + for (const [a, b] of pairs) { + assert.ok(a < b, `precondition: ${a} < ${b}`); + const ra = reverseMs(a); + const rb = reverseMs(b); + assert.strictEqual(ra.length, 13); + assert.strictEqual(rb.length, 13); + assert.ok(ra > rb, `monotonicity broken for a=${a} b=${b}: reverseMs(a)=${ra} should be > reverseMs(b)=${rb}`); + } + }); + }); + + describe('dateBucket', () => { + it('uses UTC and produces YYYY-MM-DD', () => { + assert.strictEqual(dateBucket(0), '1970-01-01'); + assert.strictEqual(dateBucket(Date.UTC(2025, 10, 13, 0, 0, 0, 0)), '2025-11-13'); + assert.strictEqual(dateBucket(Date.UTC(2025, 10, 13, 12, 34, 56, 789)), '2025-11-13'); + assert.strictEqual(dateBucket(Date.UTC(2025, 10, 13, 23, 59, 59, 999)), '2025-11-13'); + assert.strictEqual(dateBucket(Date.UTC(2025, 10, 13, 23, 59, 59, 999) + 1), '2025-11-14'); + }); + + it('rejects non-finite, non-integer, or negative ms inputs with a RangeError', () => { + assert.throws(() => dateBucket(Number.NaN), RangeError); + assert.throws(() => dateBucket(Number.POSITIVE_INFINITY), RangeError); + assert.throws(() => dateBucket(Number.NEGATIVE_INFINITY), RangeError); + assert.throws(() => dateBucket(1.5), RangeError, 'a fractional ms is out of contract'); + assert.throws(() => dateBucket(-1), RangeError, 'a pre-epoch negative ms is out of contract'); + assert.doesNotThrow(() => dateBucket(-0), 'signed -0 collapses to the epoch instant'); + assert.strictEqual(dateBucket(-0), dateBucket(0)); + }); + }); +}); diff --git a/tegg/core/agent-runtime/test/HttpSSEWriter.test.ts b/tegg/core/agent-runtime/test/HttpSSEWriter.test.ts index 6dd3f4535e..331364eafe 100644 --- a/tegg/core/agent-runtime/test/HttpSSEWriter.test.ts +++ b/tegg/core/agent-runtime/test/HttpSSEWriter.test.ts @@ -36,7 +36,6 @@ describe('test/HttpSSEWriter.test.ts', () => { }); it('should delay headers until first writeEvent', () => { - // eslint-disable-next-line @typescript-eslint/no-explicit-any const writer = new HttpSSEWriter(res as any); // Headers not sent yet after construction @@ -51,18 +50,16 @@ describe('test/HttpSSEWriter.test.ts', () => { }); it('should use lowercase header keys', () => { - // eslint-disable-next-line @typescript-eslint/no-explicit-any const writer = new HttpSSEWriter(res as any); writer.writeEvent('ping', {}); assert.ok(res.writtenHead); assert.equal(res.writtenHead.headers['content-type'], 'text/event-stream'); assert.equal(res.writtenHead.headers['cache-control'], 'no-cache'); - assert.equal(res.writtenHead.headers['connection'], 'keep-alive'); + assert.equal(res.writtenHead.headers.connection, 'keep-alive'); }); it('should format SSE events correctly', () => { - // eslint-disable-next-line @typescript-eslint/no-explicit-any const writer = new HttpSSEWriter(res as any); writer.writeEvent('message', { text: 'hello' }); @@ -71,7 +68,6 @@ describe('test/HttpSSEWriter.test.ts', () => { }); it('should not write after connection closes', () => { - // eslint-disable-next-line @typescript-eslint/no-explicit-any const writer = new HttpSSEWriter(res as any); // Simulate client disconnect @@ -86,7 +82,6 @@ describe('test/HttpSSEWriter.test.ts', () => { }); it('should trigger onClose callbacks when connection closes', () => { - // eslint-disable-next-line @typescript-eslint/no-explicit-any const writer = new HttpSSEWriter(res as any); const calls: number[] = []; @@ -99,7 +94,6 @@ describe('test/HttpSSEWriter.test.ts', () => { }); it('should handle end() idempotently', () => { - // eslint-disable-next-line @typescript-eslint/no-explicit-any const writer = new HttpSSEWriter(res as any); assert.equal(writer.closed, false); @@ -115,7 +109,6 @@ describe('test/HttpSSEWriter.test.ts', () => { }); it('should write multiple events sequentially', () => { - // eslint-disable-next-line @typescript-eslint/no-explicit-any const writer = new HttpSSEWriter(res as any); writer.writeEvent('event1', { n: 1 }); @@ -132,8 +125,35 @@ describe('test/HttpSSEWriter.test.ts', () => { }); it('should start with closed=false', () => { - // eslint-disable-next-line @typescript-eslint/no-explicit-any const writer = new HttpSSEWriter(res as any); assert.equal(writer.closed, false); }); + + it('should format SSE comments correctly', () => { + const writer = new HttpSSEWriter(res as any); + writer.writeComment('keepalive'); + + assert.equal(res.chunks.length, 1); + assert.equal(res.chunks[0], ': keepalive\n\n'); + }); + + it('should not write comment after connection closes', () => { + const writer = new HttpSSEWriter(res as any); + res.emit('close'); + + writer.writeComment('keepalive'); + + assert.equal(res.chunks.length, 0); + }); + + it('should send headers on first writeComment', () => { + const writer = new HttpSSEWriter(res as any); + + assert.equal(res.writtenHead, null); + writer.writeComment('ping'); + + assert.ok(res.writtenHead); + assert.equal(res.writtenHead.statusCode, 200); + assert.equal(res.writtenHead.headers['content-type'], 'text/event-stream'); + }); }); diff --git a/tegg/core/agent-runtime/test/MessageConverter.test.ts b/tegg/core/agent-runtime/test/MessageConverter.test.ts index 96c50bda71..60e2fe9454 100644 --- a/tegg/core/agent-runtime/test/MessageConverter.test.ts +++ b/tegg/core/agent-runtime/test/MessageConverter.test.ts @@ -1,183 +1,185 @@ import assert from 'node:assert'; -import type { AgentStreamMessage, AgentStreamMessagePayload } from '@eggjs/tegg-types/agent-runtime'; -import { MessageRole, MessageStatus, AgentObjectType, ContentBlockType } from '@eggjs/tegg-types/agent-runtime'; +import type { AgentMessage, InputMessage, SDKResultMessage } from '@eggjs/tegg-types/agent-runtime'; import { describe, it } from 'vitest'; import { MessageConverter } from '../src/MessageConverter.ts'; describe('test/MessageConverter.test.ts', () => { - describe('toContentBlocks', () => { - it('should return empty array for falsy payload', () => { - const result = MessageConverter.toContentBlocks(undefined as unknown as AgentStreamMessagePayload); - assert.deepStrictEqual(result, []); - }); - - it('should convert string content to a single text block', () => { - const payload: AgentStreamMessagePayload = { content: 'hello world' }; - const result = MessageConverter.toContentBlocks(payload); - assert.equal(result.length, 1); - assert.equal(result[0].type, ContentBlockType.Text); - assert.equal(result[0].text.value, 'hello world'); - assert.deepStrictEqual(result[0].text.annotations, []); - }); - - it('should convert array content parts to text blocks', () => { - const payload: AgentStreamMessagePayload = { - content: [ - { type: 'text', text: 'part1' }, - { type: 'text', text: 'part2' }, - ], - }; - const result = MessageConverter.toContentBlocks(payload); - assert.equal(result.length, 2); - assert.equal(result[0].text.value, 'part1'); - assert.equal(result[1].text.value, 'part2'); + describe('extractUsage', () => { + it('should return undefined when no result messages', () => { + const messages: AgentMessage[] = [ + { type: 'assistant', message: { role: 'assistant', content: [{ type: 'text', text: 'hi' }] } }, + { type: 'user', message: { role: 'user', content: 'hello' } }, + ]; + const usage = MessageConverter.extractUsage(messages); + assert.equal(usage, undefined); }); - it('should filter out non-text content parts', () => { - const payload: AgentStreamMessagePayload = { - content: [ - { type: 'text', text: 'keep' }, - { type: 'image' as 'text', text: 'discard' }, - ], - }; - const result = MessageConverter.toContentBlocks(payload); - assert.equal(result.length, 1); - assert.equal(result[0].text.value, 'keep'); + it('should extract usage from a single result message', () => { + const messages: AgentMessage[] = [ + { type: 'assistant', message: { role: 'assistant', content: [{ type: 'text', text: 'hi' }] } }, + { + type: 'result', + subtype: 'success', + usage: { input_tokens: 10, output_tokens: 5 }, + } as SDKResultMessage, + ]; + const usage = MessageConverter.extractUsage(messages); + assert.ok(usage); + assert.equal(usage.promptTokens, 10); + assert.equal(usage.completionTokens, 5); + assert.equal(usage.totalTokens, 15); }); - it('should return empty array for non-string non-array content', () => { - const payload = { content: 123 } as unknown as AgentStreamMessagePayload; - const result = MessageConverter.toContentBlocks(payload); - assert.deepStrictEqual(result, []); + it('should accumulate usage from multiple result messages', () => { + const messages: AgentMessage[] = [ + { + type: 'result', + subtype: 'success', + usage: { input_tokens: 10, output_tokens: 5 }, + } as SDKResultMessage, + { + type: 'result', + subtype: 'success', + usage: { input_tokens: 20, output_tokens: 8 }, + } as SDKResultMessage, + ]; + const usage = MessageConverter.extractUsage(messages); + assert.ok(usage); + assert.equal(usage.promptTokens, 30); + assert.equal(usage.completionTokens, 13); + assert.equal(usage.totalTokens, 43); }); - }); - describe('toMessageObject', () => { - it('should create a completed assistant message', () => { - const payload: AgentStreamMessagePayload = { content: 'reply' }; - const msg = MessageConverter.toMessageObject(payload, 'run_1'); - - assert.ok(msg.id.startsWith('msg_')); - assert.equal(msg.object, AgentObjectType.ThreadMessage); - assert.equal(msg.runId, 'run_1'); - assert.equal(msg.role, MessageRole.Assistant); - assert.equal(msg.status, MessageStatus.Completed); - assert.equal(typeof msg.createdAt, 'number'); - const content = msg.content; - assert.equal(content.length, 1); - assert.equal(content[0].text.value, 'reply'); + it('should handle result message without usage field', () => { + const messages: AgentMessage[] = [{ type: 'result', subtype: 'success' } as SDKResultMessage]; + const usage = MessageConverter.extractUsage(messages); + assert.equal(usage, undefined); }); - it('should work without runId', () => { - const payload: AgentStreamMessagePayload = { content: 'test' }; - const msg = MessageConverter.toMessageObject(payload); - assert.equal(msg.runId, undefined); + it('should handle empty messages array', () => { + const usage = MessageConverter.extractUsage([]); + assert.equal(usage, undefined); }); - }); - describe('createStreamMessage', () => { - it('should create an in-progress message with empty content', () => { - const msg = MessageConverter.createStreamMessage('msg_abc', 'run_1'); - - assert.equal(msg.id, 'msg_abc'); - assert.equal(msg.object, AgentObjectType.ThreadMessage); - assert.equal(msg.runId, 'run_1'); - assert.equal(msg.role, MessageRole.Assistant); - assert.equal(msg.status, MessageStatus.InProgress); - assert.deepStrictEqual(msg.content, []); - assert.equal(typeof msg.createdAt, 'number'); - }); - }); - - describe('extractFromStreamMessages', () => { - it('should extract messages and accumulate usage', () => { - const messages: AgentStreamMessage[] = [ - { message: { content: 'chunk1' }, usage: { promptTokens: 10, completionTokens: 5 } }, - { message: { content: 'chunk2' }, usage: { promptTokens: 0, completionTokens: 8 } }, + it('should handle partial usage fields (missing output_tokens)', () => { + const messages: AgentMessage[] = [ + { + type: 'result', + subtype: 'success', + usage: { input_tokens: 10 }, + } as SDKResultMessage, ]; - const { output, usage } = MessageConverter.extractFromStreamMessages(messages, 'run_1'); - - assert.equal(output.length, 2); - assert.equal(output[0].content[0].text.value, 'chunk1'); - assert.equal(output[1].content[0].text.value, 'chunk2'); + const usage = MessageConverter.extractUsage(messages); assert.ok(usage); assert.equal(usage.promptTokens, 10); - assert.equal(usage.completionTokens, 13); - assert.equal(usage.totalTokens, 23); - }); - - it('should return undefined usage when no usage info', () => { - const messages: AgentStreamMessage[] = [{ message: { content: 'data' } }]; - const { output, usage } = MessageConverter.extractFromStreamMessages(messages); - assert.equal(output.length, 1); - assert.equal(usage, undefined); + assert.equal(usage.completionTokens, 0); + assert.equal(usage.totalTokens, 10); }); - it('should handle messages without message payload (usage only)', () => { - const messages: AgentStreamMessage[] = [{ usage: { promptTokens: 5, completionTokens: 3 } }]; - const { output, usage } = MessageConverter.extractFromStreamMessages(messages); - assert.equal(output.length, 0); + it('should handle cache-related usage fields', () => { + const messages: AgentMessage[] = [ + { + type: 'result', + subtype: 'success', + usage: { + input_tokens: 100, + output_tokens: 50, + cache_creation_input_tokens: 200, + cache_read_input_tokens: 80, + }, + } as SDKResultMessage, + ]; + const usage = MessageConverter.extractUsage(messages); assert.ok(usage); - assert.equal(usage.totalTokens, 8); - }); - - it('should handle empty message array', () => { - const { output, usage } = MessageConverter.extractFromStreamMessages([]); - assert.equal(output.length, 0); - assert.equal(usage, undefined); + assert.equal(usage.promptTokens, 100); + assert.equal(usage.completionTokens, 50); + assert.equal(usage.totalTokens, 150); }); }); - describe('toInputMessageObjects', () => { - it('should convert user and assistant messages', () => { - const messages = [ - { role: MessageRole.User as MessageRole, content: 'hi' }, - { role: MessageRole.Assistant as MessageRole, content: 'hello' }, + describe('filterForStorage', () => { + it('should filter out stream_event messages', () => { + const messages: AgentMessage[] = [ + { type: 'system', subtype: 'init', session_id: 'sess-1' }, + { type: 'user', message: { role: 'user', content: 'hello' } }, + { type: 'stream_event', event: { type: 'content_block_delta' }, session_id: 'sess-1' }, + { type: 'stream_event', event: { type: 'content_block_delta' }, session_id: 'sess-1' }, + { type: 'assistant', message: { role: 'assistant', content: [{ type: 'text', text: 'hi' }] } }, + { type: 'result', subtype: 'success', usage: { input_tokens: 10, output_tokens: 5 } } as SDKResultMessage, ]; - const result = MessageConverter.toInputMessageObjects(messages, 'thread_1'); - + const result = MessageConverter.filterForStorage(messages); + assert.equal(result.length, 4); + assert.equal(result[0].type, 'system'); + assert.equal(result[1].type, 'user'); + assert.equal(result[2].type, 'assistant'); + assert.equal(result[3].type, 'result'); + }); + + it('should return all messages when no stream_event present', () => { + const messages: AgentMessage[] = [ + { type: 'user', message: { role: 'user', content: 'hello' } }, + { type: 'assistant', message: { role: 'assistant', content: [{ type: 'text', text: 'hi' }] } }, + ]; + const result = MessageConverter.filterForStorage(messages); assert.equal(result.length, 2); - assert.equal(result[0].role, MessageRole.User); - assert.equal(result[0].threadId, 'thread_1'); - assert.equal(result[1].role, MessageRole.Assistant); + }); + + it('should handle empty array', () => { + const result = MessageConverter.filterForStorage([]); + assert.deepStrictEqual(result, []); + }); + }); - const content0 = result[0].content; - assert.equal(content0[0].text.value, 'hi'); + describe('toAgentMessages', () => { + it('should convert user messages to AgentMessage format', () => { + const messages: InputMessage[] = [{ role: 'user', content: 'hello' }]; + const result = MessageConverter.toAgentMessages(messages); + assert.equal(result.length, 1); + assert.equal(result[0].type, 'user'); + assert.deepStrictEqual((result[0] as any).message, { role: 'user', content: 'hello' }); }); it('should filter out system messages', () => { - const messages = [ - { role: MessageRole.System as MessageRole, content: 'you are a bot' }, - { role: MessageRole.User as MessageRole, content: 'hi' }, + const messages: InputMessage[] = [ + { role: 'system', content: 'you are a bot' }, + { role: 'user', content: 'hello' }, ]; - const result = MessageConverter.toInputMessageObjects(messages); + const result = MessageConverter.toAgentMessages(messages); assert.equal(result.length, 1); - assert.equal(result[0].role, MessageRole.User); + assert.equal(result[0].type, 'user'); }); - it('should handle array content parts', () => { - const messages = [ + it('should handle array content', () => { + const messages: InputMessage[] = [ { - role: MessageRole.User as MessageRole, + role: 'user', content: [ - { type: 'text' as const, text: 'part1' }, - { type: 'text' as const, text: 'part2' }, + { type: 'text', text: 'part1' }, + { type: 'text', text: 'part2' }, ], }, ]; - const result = MessageConverter.toInputMessageObjects(messages); - const content = result[0].content; - assert.equal(content.length, 2); - assert.equal(content[0].text.value, 'part1'); - assert.equal(content[1].text.value, 'part2'); + const result = MessageConverter.toAgentMessages(messages); + assert.equal(result.length, 1); + assert.deepStrictEqual((result[0] as any).message.content, [ + { type: 'text', text: 'part1' }, + { type: 'text', text: 'part2' }, + ]); + }); + + it('should handle empty array', () => { + const result = MessageConverter.toAgentMessages([]); + assert.deepStrictEqual(result, []); }); - it('should work without threadId', () => { - const messages = [{ role: MessageRole.User as MessageRole, content: 'hi' }]; - const result = MessageConverter.toInputMessageObjects(messages); - assert.equal(result[0].threadId, undefined); + it('should preserve assistant role messages with correct type', () => { + const messages: InputMessage[] = [{ role: 'assistant', content: 'I said something' }]; + const result = MessageConverter.toAgentMessages(messages); + assert.equal(result.length, 1); + assert.equal(result[0].type, 'assistant'); + assert.deepStrictEqual((result[0] as any).message.role, 'assistant'); }); }); }); diff --git a/tegg/core/agent-runtime/test/OSSAgentStore.test.ts b/tegg/core/agent-runtime/test/OSSAgentStore.test.ts index 5441344fc2..0a00ab1b46 100644 --- a/tegg/core/agent-runtime/test/OSSAgentStore.test.ts +++ b/tegg/core/agent-runtime/test/OSSAgentStore.test.ts @@ -1,9 +1,9 @@ import assert from 'node:assert'; -import { describe, it, beforeEach, vi } from 'vitest'; +import type { AgentMessage } from '@eggjs/tegg-types/agent-runtime'; +import { describe, it, beforeEach } from 'vitest'; -import { AgentNotFoundError } from '../src/index.ts'; -import { OSSAgentStore } from '../src/index.ts'; +import { AgentNotFoundError, OSSAgentStore, reverseMs } from '../src/index.ts'; import { MapStorageClient, MapStorageClientWithoutAppend } from './helpers.ts'; describe('test/OSSAgentStore.test.ts', () => { @@ -62,107 +62,151 @@ describe('test/OSSAgentStore.test.ts', () => { it('should append messages to a thread', async () => { const thread = await store.createThread(); - await store.appendMessages(thread.id, [ - { - id: 'msg_1', - object: 'thread.message', - createdAt: Math.floor(Date.now() / 1000), - role: 'user', - status: 'completed', - content: [{ type: 'text', text: { value: 'Hello', annotations: [] } }], - }, - { - id: 'msg_2', - object: 'thread.message', - createdAt: Math.floor(Date.now() / 1000), - role: 'assistant', - status: 'completed', - content: [{ type: 'text', text: { value: 'Hi!', annotations: [] } }], - }, - ]); + const messages: AgentMessage[] = [ + { type: 'user', message: { role: 'user', content: 'Hello' } }, + { type: 'assistant', message: { role: 'assistant', content: [{ type: 'text', text: 'Hi!' }] } }, + ]; + await store.appendMessages(thread.id, messages); const fetched = await store.getThread(thread.id); assert.equal(fetched.messages.length, 2); - assert.equal(fetched.messages[0].id, 'msg_1'); - assert.equal(fetched.messages[1].id, 'msg_2'); + assert.equal(fetched.messages[0].type, 'user'); + assert.equal(fetched.messages[1].type, 'assistant'); }); it('should append messages incrementally', async () => { const thread = await store.createThread(); + await store.appendMessages(thread.id, [{ type: 'user', message: { role: 'user', content: 'First' } }]); await store.appendMessages(thread.id, [ - { - id: 'msg_1', - object: 'thread.message', - createdAt: Math.floor(Date.now() / 1000), - role: 'user', - status: 'completed', - content: [{ type: 'text', text: { value: 'First', annotations: [] } }], - }, - ]); - await store.appendMessages(thread.id, [ - { - id: 'msg_2', - object: 'thread.message', - createdAt: Math.floor(Date.now() / 1000), - role: 'assistant', - status: 'completed', - content: [{ type: 'text', text: { value: 'Second', annotations: [] } }], - }, + { type: 'assistant', message: { role: 'assistant', content: [{ type: 'text', text: 'Second' }] } }, ]); const fetched = await store.getThread(thread.id); assert.equal(fetched.messages.length, 2); - assert.equal(fetched.messages[0].id, 'msg_1'); - assert.equal(fetched.messages[1].id, 'msg_2'); + assert.equal(fetched.messages[0].type, 'user'); + assert.equal(fetched.messages[1].type, 'assistant'); }); it('should throw AgentNotFoundError when appending to non-existent thread', async () => { await assert.rejects( () => - store.appendMessages('thread_non_existent', [ - { - id: 'msg_1', - object: 'thread.message', - createdAt: Math.floor(Date.now() / 1000), - role: 'user', - status: 'completed', - content: [{ type: 'text', text: { value: 'Hello', annotations: [] } }], - }, - ]), + store.appendMessages('thread_non_existent', [{ type: 'user', message: { role: 'user', content: 'Hello' } }]), (err: unknown) => { assert(err instanceof AgentNotFoundError); return true; }, ); }); + + it('should only return user and assistant messages by default', async () => { + const thread = await store.createThread(); + await store.appendMessages(thread.id, [ + { type: 'system', subtype: 'init', session_id: 'sess-1' }, + { type: 'user', message: { role: 'user', content: 'Hello' } }, + { type: 'assistant', message: { role: 'assistant', content: [{ type: 'text', text: 'Hi!' }] } }, + { type: 'result', subtype: 'success', usage: { input_tokens: 10, output_tokens: 5 } }, + ]); + const fetched = await store.getThread(thread.id); + assert.equal(fetched.messages.length, 2); + assert.equal(fetched.messages[0].type, 'user'); + assert.equal(fetched.messages[1].type, 'assistant'); + }); + + it('should return all message types when includeAllMessages is true', async () => { + const thread = await store.createThread(); + await store.appendMessages(thread.id, [ + { type: 'system', subtype: 'init', session_id: 'sess-1' }, + { type: 'user', message: { role: 'user', content: 'Hello' } }, + { type: 'result', subtype: 'success', usage: { input_tokens: 10, output_tokens: 5 } }, + ]); + const fetched = await store.getThread(thread.id, { includeAllMessages: true }); + assert.equal(fetched.messages.length, 3); + assert.equal(fetched.messages[0].type, 'system'); + assert.equal(fetched.messages[1].type, 'user'); + assert.equal(fetched.messages[2].type, 'result'); + }); }); describe('threads (without append)', () => { it('should fall back to get-concat-put when client has no append', async () => { const fallbackStore = new OSSAgentStore({ client: new MapStorageClientWithoutAppend() }); const thread = await fallbackStore.createThread(); + await fallbackStore.appendMessages(thread.id, [{ type: 'user', message: { role: 'user', content: 'Hello' } }]); await fallbackStore.appendMessages(thread.id, [ - { - id: 'msg_1', - object: 'thread.message', - createdAt: Math.floor(Date.now() / 1000), - role: 'user', - status: 'completed', - content: [{ type: 'text', text: { value: 'Hello', annotations: [] } }], - }, - ]); - await fallbackStore.appendMessages(thread.id, [ - { - id: 'msg_2', - object: 'thread.message', - createdAt: Math.floor(Date.now() / 1000), - role: 'assistant', - status: 'completed', - content: [{ type: 'text', text: { value: 'Hi!', annotations: [] } }], - }, + { type: 'assistant', message: { role: 'assistant', content: [{ type: 'text', text: 'Hi!' }] } }, ]); const fetched = await fallbackStore.getThread(thread.id); assert.equal(fetched.messages.length, 2); - assert.equal(fetched.messages[0].id, 'msg_1'); - assert.equal(fetched.messages[1].id, 'msg_2'); + assert.equal(fetched.messages[0].type, 'user'); + assert.equal(fetched.messages[1].type, 'assistant'); + }); + }); + + describe('thread metadata', () => { + it('should shallow-merge metadata and preserve omitted keys', async () => { + const thread = await store.createThread({ + bizId: 'order_123', + source: 'customer_service', + nested: { old: true }, + }); + + await store.updateThreadMetadata(thread.id, { + source: 'operator_console', + nested: { replacement: true }, + nullable: null, + }); + + assert.deepStrictEqual((await store.getThread(thread.id)).metadata, { + bizId: 'order_123', + source: 'operator_console', + nested: { replacement: true }, + nullable: null, + }); + }); + + it('should reject updating a missing thread', async () => { + await assert.rejects( + () => store.updateThreadMetadata('thread_missing', { bizId: 'order_123' }), + AgentNotFoundError, + ); + }); + + it('should serialize concurrent updates in the same process', async () => { + const client = new MapStorageClient(); + const localStore = new OSSAgentStore({ client }); + const thread = await localStore.createThread({ original: true }); + client.delayPutWhenKeyMatches(/threads\/.*\/meta\.json$/, 20); + + await Promise.all([ + localStore.updateThreadMetadata(thread.id, { first: 1 }), + localStore.updateThreadMetadata(thread.id, { second: 2 }), + ]); + + assert.deepStrictEqual((await localStore.getThread(thread.id)).metadata, { + original: true, + first: 1, + second: 2, + }); + }); + + it('should not clobber metadata when an updateThreadMetadata races a latestRunId write', async () => { + const client = new MapStorageClient(); + const localStore = new OSSAgentStore({ client }); + const thread = await localStore.createThread({ original: true }); + // Both writers do a read-modify-write on the same meta.json; the delay + // forces them to overlap so an unsynchronized pair would clobber. + client.delayPutWhenKeyMatches(/threads\/.*\/meta\.json$/, 20); + + const [, run] = await Promise.all([ + localStore.updateThreadMetadata(thread.id, { merged: 1 }), + // createRun fires recordLatestRunId in the background (latestRunId write). + localStore.createRun([{ role: 'user', content: 'Hi' }], thread.id), + ]); + await localStore.awaitPendingWrites(); + + const stored = await localStore.getThread(thread.id); + // The metadata merge survives... + assert.deepStrictEqual(stored.metadata, { original: true, merged: 1 }); + // ...and so does the latestRunId pointer written by createRun. + assert.equal(await localStore.getLatestRunId(thread.id), run.id); }); }); @@ -224,23 +268,13 @@ describe('test/OSSAgentStore.test.ts', () => { const run = await store.createRun([{ role: 'user', content: 'Hello' }]); await store.updateRun(run.id, { status: 'completed', - output: [ - { - id: 'msg_1', - object: 'thread.message', - createdAt: Math.floor(Date.now() / 1000), - role: 'assistant', - status: 'completed', - content: [{ type: 'text', text: { value: 'World', annotations: [] } }], - }, - ], completedAt: Math.floor(Date.now() / 1000), + usage: { promptTokens: 10, completionTokens: 5, totalTokens: 15 }, }); const fetched = await store.getRun(run.id); assert.equal(fetched.status, 'completed'); - assert(fetched.output); - assert.equal(fetched.output.length, 1); assert(typeof fetched.completedAt === 'number'); + assert.deepStrictEqual(fetched.usage, { promptTokens: 10, completionTokens: 5, totalTokens: 15 }); }); it('should not allow overwriting id or object via updateRun', async () => { @@ -257,21 +291,83 @@ describe('test/OSSAgentStore.test.ts', () => { }); }); + describe('recent run id', () => { + it('should record latestRunId on the thread when createRun has a threadId', async () => { + const thread = await store.createThread(); + const run = await store.createRun([{ role: 'user', content: 'Hello' }], thread.id); + // latestRunId is written in the background; drain before asserting. + await store.awaitPendingWrites(); + assert.equal(await store.getLatestRunId(thread.id), run.id); + }); + + it('should point getLatestRunId at the most recent run', async () => { + const thread = await store.createThread(); + await store.createRun([{ role: 'user', content: 'first' }], thread.id); + const second = await store.createRun([{ role: 'user', content: 'second' }], thread.id); + await store.awaitPendingWrites(); + assert.equal(await store.getLatestRunId(thread.id), second.id); + }); + + it('should preserve existing thread metadata when recording latestRunId', async () => { + const thread = await store.createThread({ origin: 'audit', agentName: 'foo' }); + const run = await store.createRun([{ role: 'user', content: 'Hello' }], thread.id); + await store.awaitPendingWrites(); + const fetched = await store.getThread(thread.id); + assert.deepEqual(fetched.metadata, { origin: 'audit', agentName: 'foo' }); + assert.equal(await store.getLatestRunId(thread.id), run.id); + }); + + it('should return null for a thread that exists but has no run (historical data)', async () => { + const thread = await store.createThread(); + assert.equal(await store.getLatestRunId(thread.id), null); + }); + + it('should not record latestRunId when createRun has no threadId', async () => { + const run = await store.createRun([{ role: 'user', content: 'Hello' }]); + assert(run.id.startsWith('run_')); + // No thread to query; the run simply has no latest-run pointer anywhere. + }); + + it('should not fail run creation when the threadId does not exist', async () => { + // Best-effort pointer write: a missing thread meta is swallowed. + const run = await store.createRun([{ role: 'user', content: 'Hello' }], 'thread_missing'); + assert(run.id.startsWith('run_')); + }); + + it('should throw AgentNotFoundError from getLatestRunId for a non-existent thread', async () => { + await assert.rejects( + () => store.getLatestRunId('thread_non_existent'), + (err: unknown) => { + assert(err instanceof AgentNotFoundError); + assert.equal(err.status, 404); + assert.match(err.message, /Thread thread_non_existent not found/); + return true; + }, + ); + }); + }); + describe('init / destroy', () => { it('should call client init when present', async () => { + let initCalled = false; const client = new MapStorageClient(); - client.init = vi.fn(); + client.init = async () => { + initCalled = true; + }; const s = new OSSAgentStore({ client }); await s.init(); - assert.equal((client.init as ReturnType).mock.calls.length, 1); + assert.equal(initCalled, true); }); it('should call client destroy when present', async () => { + let destroyCalled = false; const client = new MapStorageClient(); - client.destroy = vi.fn(); + client.destroy = async () => { + destroyCalled = true; + }; const s = new OSSAgentStore({ client }); await s.destroy(); - assert.equal((client.destroy as ReturnType).mock.calls.length, 1); + assert.equal(destroyCalled, true); }); it('should not throw when client has no init/destroy', async () => { @@ -322,4 +418,427 @@ describe('test/OSSAgentStore.test.ts', () => { ); }); }); + + describe('thread activity index', () => { + function withFixedNow(ms: number, fn: () => Promise): Promise { + const realNow = Date.now; + Date.now = () => ms; + return Promise.resolve() + .then(fn) + .finally(() => { + Date.now = realNow; + }); + } + + const INDEX_KEY_RE_AGENT = /^agent\/index\/threads-by-updated-date\/\d{4}-\d{2}-\d{2}\/\d{13}_thread_[0-9a-f-]+$/; + const INDEX_KEY_RE_ANY = + /^(?:[^/]+(?:\/[^/]+)*\/)?index\/threads-by-updated-date\/\d{4}-\d{2}-\d{2}\/\d{13}_thread_[0-9a-f-]+$/; + + const T_EARLY = Date.UTC(2025, 10, 13, 8, 0, 0, 0); + const T_MID = Date.UTC(2025, 10, 13, 8, 0, 1, 234); + const T_LATE = Date.UTC(2025, 10, 13, 10, 0, 0, 0); + + it('writes a sidecar activity index next to meta.json on createThread', async () => { + const client = new MapStorageClient(); + const localStore = new OSSAgentStore({ client, prefix: 'agent/' }); + + const thread = await withFixedNow(T_MID, () => localStore.createThread({ user: 'alice' })); + await localStore.awaitPendingWrites(); + + const keys = client.keys(); + assert.ok( + keys.includes(`agent/threads/${thread.id}/meta.json`), + `meta.json not found amongst keys: ${JSON.stringify(keys)}`, + ); + + const indexKeys = client.keysWithPrefix('agent/index/threads-by-updated-date/'); + assert.equal(indexKeys.length, 1, `expected exactly one index entry, got ${JSON.stringify(indexKeys)}`); + const indexKey = indexKeys[0]; + assert.match(indexKey, INDEX_KEY_RE_AGENT); + + const expectedRev = reverseMs(T_MID); + const expectedKey = `agent/index/threads-by-updated-date/2025-11-13/${expectedRev}_${thread.id}`; + assert.equal(indexKey, expectedKey); + }); + + it('the index body carries threadId, createdAt, updatedAt and a snapshot of metadata', async () => { + const client = new MapStorageClient(); + const localStore = new OSSAgentStore({ client, prefix: 'agent/' }); + const meta: Record = { user: 'alice', channel: 'web' }; + + const thread = await withFixedNow(T_MID, () => localStore.createThread(meta)); + await localStore.awaitPendingWrites(); + + const indexKey = client.keysWithPrefix('agent/index/threads-by-updated-date/')[0]; + const raw = await client.get(indexKey); + assert.ok(raw, 'index body should be present after drain'); + const body = JSON.parse(raw) as { + threadId: string; + createdAt: number; + updatedAt: number; + metadata: Record; + }; + + assert.deepStrictEqual(body, { + threadId: thread.id, + createdAt: Math.floor(T_MID / 1000), + updatedAt: Math.floor(T_MID / 1000), + metadata: { user: 'alice', channel: 'web' }, + }); + assert.strictEqual(body.createdAt, thread.createdAt); + assert.strictEqual(body.updatedAt, thread.createdAt); + + meta.user = 'mutated'; + const reread = JSON.parse((await client.get(indexKey))!) as typeof body; + assert.equal(reread.metadata.user, 'alice'); + }); + + it('within a date directory, ASC dictionary order equals time-DESC activity order', async () => { + const client = new MapStorageClient(); + const localStore = new OSSAgentStore({ client, prefix: 'agent/' }); + + const early = await withFixedNow(T_EARLY, () => localStore.createThread({ tag: 'early' })); + const mid = await withFixedNow(T_MID, () => localStore.createThread({ tag: 'mid' })); + const late = await withFixedNow(T_LATE, () => localStore.createThread({ tag: 'late' })); + await localStore.awaitPendingWrites(); + + const indexKeys = client.keysWithPrefix('agent/index/threads-by-updated-date/2025-11-13/'); + assert.equal(indexKeys.length, 3); + + const threadIdsInListOrder = indexKeys.map((k) => { + const fileName = k.slice(k.lastIndexOf('/') + 1); + return fileName.slice(fileName.indexOf('_') + 1); + }); + assert.deepStrictEqual(threadIdsInListOrder, [late.id, mid.id, early.id]); + + const revs = indexKeys.map((k) => { + const fileName = k.slice(k.lastIndexOf('/') + 1); + return fileName.slice(0, fileName.indexOf('_')); + }); + assert.deepStrictEqual(revs, [reverseMs(T_LATE), reverseMs(T_MID), reverseMs(T_EARLY)]); + assert.ok(revs[0] < revs[1] && revs[1] < revs[2], 'revs must be ASCII-ascending'); + }); + + it('threads active on either side of a UTC day boundary land in different date buckets', async () => { + const client = new MapStorageClient(); + const localStore = new OSSAgentStore({ client, prefix: 'agent/' }); + + const lastMsOfNov13Utc = Date.UTC(2025, 10, 13, 23, 59, 59, 999); + const firstMsOfNov14Utc = Date.UTC(2025, 10, 14, 0, 0, 0, 0); + + const earlier = await withFixedNow(lastMsOfNov13Utc, () => localStore.createThread()); + const later = await withFixedNow(firstMsOfNov14Utc, () => localStore.createThread()); + await localStore.awaitPendingWrites(); + + const nov13 = client.keysWithPrefix('agent/index/threads-by-updated-date/2025-11-13/'); + const nov14 = client.keysWithPrefix('agent/index/threads-by-updated-date/2025-11-14/'); + assert.equal(nov13.length, 1, `Nov 13 bucket: ${JSON.stringify(nov13)}`); + assert.equal(nov14.length, 1, `Nov 14 bucket: ${JSON.stringify(nov14)}`); + assert.ok(nov13[0].endsWith(`_${earlier.id}`)); + assert.ok(nov14[0].endsWith(`_${later.id}`)); + }); + + it('createThread does not wait for the index PUT to complete (non-blocking)', async () => { + const client = new MapStorageClient(); + const localStore = new OSSAgentStore({ client, prefix: 'agent/' }); + + const INDEX_DELAY_MS = 120; + client.delayPutWhenKeyMatches(/^agent\/index\/threads-by-updated-date\//, INDEX_DELAY_MS); + + const t0 = Date.now(); + const thread = await localStore.createThread(); + const elapsedMs = Date.now() - t0; + + assert.ok( + elapsedMs < INDEX_DELAY_MS - 50, + `createThread should return before the index PUT settles, but elapsed=${elapsedMs}ms (delay=${INDEX_DELAY_MS}ms)`, + ); + + assert.equal( + client.keysWithPrefix('agent/index/').length, + 0, + 'the slow background PUT should not be visible at the moment createThread returns', + ); + + assert.ok(await client.get(`agent/threads/${thread.id}/meta.json`)); + + await localStore.awaitPendingWrites(); + const indexKeys = client.keysWithPrefix('agent/index/threads-by-updated-date/'); + assert.equal(indexKeys.length, 1); + assert.ok(indexKeys[0].endsWith(`_${thread.id}`)); + }); + + it('createThread succeeds even when the index PUT throws; the failure becomes a single warn line', async () => { + const client = new MapStorageClient(); + client.failPutWhenKeyMatches(/^agent\/index\/threads-by-updated-date\//); + + const warnCalls: unknown[][] = []; + const logger = { + warn(message: string, ...args: unknown[]): void { + warnCalls.push([message, ...args]); + }, + }; + const localStore = new OSSAgentStore({ + client, + prefix: 'agent/', + logger, + }); + + const thread = await localStore.createThread({ trace: 'fail-path' }); + assert.equal(thread.object, 'thread'); + assert.match(thread.id, /^thread_[0-9a-f-]{36}$/); + + await localStore.awaitPendingWrites(); + + const metaRaw = await client.get(`agent/threads/${thread.id}/meta.json`); + assert.ok(metaRaw, 'meta.json should still be present when only the index PUT failed'); + const metaObj = JSON.parse(metaRaw) as { id: string; createdAt: number }; + assert.equal(metaObj.id, thread.id); + assert.equal(metaObj.createdAt, thread.createdAt); + + assert.equal( + client.keysWithPrefix('agent/index/').length, + 0, + 'the simulated index PUT failure means no index entries exist', + ); + + assert.equal( + warnCalls.length, + 1, + `expected one warn call, got ${warnCalls.length}: ${JSON.stringify(warnCalls)}`, + ); + const call = warnCalls[0]; + const formatStr = call[0]; + assert.equal(typeof formatStr, 'string'); + assert.ok( + (formatStr as string).includes('failed to write thread activity index'), + `unexpected warn format string: ${String(formatStr)}`, + ); + assert.equal(call[1], thread.id); + assert.match(call[2] as string, /^agent\/index\/threads-by-updated-date\/\d{4}-\d{2}-\d{2}\/\d{13}_thread_/); + const errArg = call[3]; + assert.ok( + errArg instanceof Error, + `expected the fourth warn-arg to be an Error instance (so the stack is preserved when the logger renders it), got ${typeof errArg}: ${String(errArg)}`, + ); + assert.match((errArg as Error).message, /simulated PUT failure/); + assert.ok( + !(call[0] as string).includes('err=%s'), + 'the format string should no longer carry an err=%s placeholder', + ); + + const fetched = await localStore.getThread(thread.id); + assert.equal(fetched.id, thread.id); + }); + + it('with no logger configured, an index PUT failure falls back to console.warn without throwing', async () => { + const client = new MapStorageClient(); + client.failPutWhenKeyMatches(/^agent\/index\/threads-by-updated-date\//); + const localStore = new OSSAgentStore({ client, prefix: 'agent/' }); + + const realWarn = console.warn; + const captured: unknown[][] = []; + console.warn = ((...args: unknown[]) => { + captured.push(args); + }) as typeof console.warn; + + try { + const thread = await localStore.createThread(); + await localStore.awaitPendingWrites(); + assert.equal(captured.length, 1); + const args = captured[0]; + assert.equal(typeof args[0], 'string'); + assert.ok((args[0] as string).includes('failed to write thread activity index')); + assert.equal(args[1], thread.id); + } finally { + console.warn = realWarn; + } + }); + + it('destroy() drains in-flight index writes before tearing down the underlying client', async () => { + const client = new MapStorageClient(); + const localStore = new OSSAgentStore({ client, prefix: 'agent/' }); + + const INDEX_DELAY_MS = 120; + client.delayPutWhenKeyMatches(/^agent\/index\/threads-by-updated-date\//, INDEX_DELAY_MS); + + let clientDestroyCalledAt: number | null = null; + let indexKeyVisibleAtDestroy = false; + client.destroy = async () => { + clientDestroyCalledAt = Date.now(); + indexKeyVisibleAtDestroy = client.keysWithPrefix('agent/index/').length > 0; + }; + + const createReturnedAt = Date.now(); + await localStore.createThread(); + assert.equal(client.keysWithPrefix('agent/index/').length, 0); + + const beforeDestroy = Date.now(); + await localStore.destroy(); + const elapsedMs = Date.now() - beforeDestroy; + + assert.ok( + elapsedMs >= INDEX_DELAY_MS - 20, + `destroy() should wait for the in-flight index write, elapsedMs=${elapsedMs}ms delay=${INDEX_DELAY_MS}ms`, + ); + + assert(clientDestroyCalledAt !== null, 'client.destroy should have been invoked'); + assert.ok( + indexKeyVisibleAtDestroy, + 'the index entry should have been observable to client.destroy, meaning the drain ran first', + ); + + assert.equal(client.keysWithPrefix('agent/index/threads-by-updated-date/').length, 1); + assert.ok(clientDestroyCalledAt >= createReturnedAt); + }); + + it('destroy() drains a write that arrives during a previous drain iteration', async () => { + const client = new MapStorageClient(); + const localStore = new OSSAgentStore({ client, prefix: 'agent/' }); + + const INDEX_DELAY_MS = 80; + client.delayPutWhenKeyMatches(/^agent\/index\/threads-by-updated-date\//, INDEX_DELAY_MS); + + const thread1 = await localStore.createThread({ ordinal: 1 }); + assert.equal( + client.keysWithPrefix('agent/index/threads-by-updated-date/').length, + 0, + "thread1's index PUT should still be in flight at this instant", + ); + + const destroyP = localStore.destroy(); + const thread2 = await localStore.createThread({ ordinal: 2 }); + await destroyP; + + const indexKeys = client.keysWithPrefix('agent/index/threads-by-updated-date/'); + assert.equal( + indexKeys.length, + 2, + `the drain loop should have captured both the original and the late-arriving index write, got ${JSON.stringify(indexKeys)}`, + ); + const threadIdsInLexOrder = indexKeys + .map((k) => { + const fileName = k.slice(k.lastIndexOf('/') + 1); + return fileName.slice(fileName.indexOf('_') + 1); + }) + .sort(); + const expectedIds = [thread1.id, thread2.id].sort(); + assert.deepStrictEqual(threadIdsInLexOrder, expectedIds); + }); + + it('awaitPendingWrites() is a no-op resolved promise when there are no pending writes', async () => { + const client = new MapStorageClient(); + const localStore = new OSSAgentStore({ client, prefix: 'agent/' }); + + const t0 = Date.now(); + await localStore.awaitPendingWrites(); + const elapsedMs = Date.now() - t0; + assert.ok( + elapsedMs < 50, + `awaitPendingWrites with an empty queue should resolve immediately, took ${elapsedMs}ms`, + ); + + await localStore.awaitPendingWrites(); + await localStore.awaitPendingWrites(); + }); + + it('uses the normalized prefix on the index path, including the empty-prefix case', async () => { + const bareClient = new MapStorageClient(); + const bareStore = new OSSAgentStore({ client: bareClient }); + await withFixedNow(T_MID, () => bareStore.createThread()); + await bareStore.awaitPendingWrites(); + const bareIndex = bareClient.keysWithPrefix('index/threads-by-updated-date/'); + assert.equal(bareIndex.length, 1); + assert.ok(!bareIndex[0].startsWith('/'), `expected no leading "/", got ${bareIndex[0]}`); + assert.match(bareIndex[0], INDEX_KEY_RE_ANY); + + const nestedClient = new MapStorageClient(); + const nestedStore = new OSSAgentStore({ client: nestedClient, prefix: 'a/b' }); + await withFixedNow(T_MID, () => nestedStore.createThread()); + await nestedStore.awaitPendingWrites(); + const nestedIndex = nestedClient.keysWithPrefix('a/b/index/threads-by-updated-date/'); + assert.equal( + nestedIndex.length, + 1, + `expected exactly one nested index entry: ${JSON.stringify(nestedClient.keys())}`, + ); + assert.match(nestedIndex[0], /^a\/b\/index\/threads-by-updated-date\/2025-11-13\/\d{13}_thread_[0-9a-f-]+$/); + + const metaKey = nestedClient.keys().find((k) => k.startsWith('a/b/threads/') && k.endsWith('/meta.json')); + assert.ok(metaKey, 'meta.json key for the nested prefix should exist'); + const expectedThreadId = metaKey.slice('a/b/threads/'.length, -'/meta.json'.length); + assert.ok(nestedIndex[0].endsWith(`_${expectedThreadId}`)); + }); + + it('appendMessages writes a new activity index entry for a reused historical thread', async () => { + const client = new MapStorageClient(); + const localStore = new OSSAgentStore({ client, prefix: 'agent/' }); + + const createdMs = Date.UTC(2025, 10, 13, 8, 0, 0, 0); + const updatedMs = Date.UTC(2025, 10, 14, 9, 0, 0, 0); + const thread = await withFixedNow(createdMs, () => localStore.createThread({ origin: 'audit' })); + await localStore.awaitPendingWrites(); + + const messages: AgentMessage[] = [ + { type: 'user', message: { role: 'user', content: 'hello' } } as unknown as AgentMessage, + { + type: 'assistant', + message: { role: 'assistant', content: [{ type: 'text', text: 'hi' }] }, + } as unknown as AgentMessage, + ]; + await withFixedNow(updatedMs, () => localStore.appendMessages(thread.id, messages)); + await localStore.awaitPendingWrites(); + + const createdDay = client.keysWithPrefix('agent/index/threads-by-updated-date/2025-11-13/'); + const updatedDay = client.keysWithPrefix('agent/index/threads-by-updated-date/2025-11-14/'); + assert.equal(createdDay.length, 1, `created-day index: ${JSON.stringify(createdDay)}`); + assert.equal(updatedDay.length, 1, `updated-day index: ${JSON.stringify(updatedDay)}`); + assert.ok(createdDay[0].endsWith(`_${thread.id}`)); + assert.ok(updatedDay[0].endsWith(`_${thread.id}`)); + + const updatedBody = JSON.parse((await client.get(updatedDay[0]))!) as { + threadId: string; + createdAt: number; + updatedAt: number; + metadata: Record; + }; + assert.deepStrictEqual(updatedBody, { + threadId: thread.id, + createdAt: Math.floor(createdMs / 1000), + updatedAt: Math.floor(updatedMs / 1000), + metadata: { origin: 'audit' }, + }); + }); + + it('createRun and updateRun never touch the activity index', async () => { + const client = new MapStorageClient(); + const localStore = new OSSAgentStore({ client, prefix: 'agent/' }); + + const thread = await localStore.createThread({ origin: 'audit' }); + await localStore.awaitPendingWrites(); + const indexKeysBefore = client.keysWithPrefix('agent/index/').slice(); + const indexBodyBefore = await client.get(indexKeysBefore[0]); + assert.equal(indexKeysBefore.length, 1, 'baseline: exactly one index entry from createThread'); + + const run = await localStore.createRun( + [{ role: 'user', content: 'first turn' }], + thread.id, + { maxIterations: 1 }, + { trace: 't' }, + ); + await localStore.updateRun(run.id, { + status: 'completed', + completedAt: Math.floor(Date.now() / 1000), + }); + + const indexKeysAfter = client.keysWithPrefix('agent/index/').slice(); + assert.deepStrictEqual( + indexKeysAfter.slice().sort(), + indexKeysBefore.slice().sort(), + 'no new index keys should appear from run operations', + ); + const indexBodyAfter = await client.get(indexKeysAfter[0]); + assert.equal(indexBodyAfter, indexBodyBefore, 'the index body should be byte-identical'); + }); + }); }); diff --git a/tegg/core/agent-runtime/test/OSSObjectStorageClient.test.ts b/tegg/core/agent-runtime/test/OSSObjectStorageClient.test.ts index 7a77c04e83..0318feb5d4 100644 --- a/tegg/core/agent-runtime/test/OSSObjectStorageClient.test.ts +++ b/tegg/core/agent-runtime/test/OSSObjectStorageClient.test.ts @@ -1,25 +1,92 @@ import assert from 'node:assert'; import type { OSSObject } from 'oss-client'; -import { describe, it, vi, beforeEach } from 'vitest'; +import { describe, it, beforeEach } from 'vitest'; import { OSSObjectStorageClient } from '../src/OSSObjectStorageClient.ts'; +/** Simple mock function helper for mocha tests. */ +function mockFn() { + const calls: any[][] = []; + let nextResults: Array<{ type: 'resolve' | 'reject'; value: any }> = []; + const fn = (...args: any[]) => { + calls.push(args); + const result = nextResults.shift(); + if (result) { + return result.type === 'resolve' ? Promise.resolve(result.value) : Promise.reject(result.value); + } + return Promise.resolve({}); + }; + fn.mock = { calls }; + fn.mockResolvedValue = (val: any) => { + nextResults = []; + fn.mockResolvedValueOnce(val); + nextResults = nextResults.map(() => ({ type: 'resolve' as const, value: val })); + (fn as any)._defaultResult = { type: 'resolve', value: val }; + return fn; + }; + fn.mockResolvedValueOnce = (val: any) => { + nextResults.push({ type: 'resolve', value: val }); + return fn; + }; + fn.mockRejectedValue = (val: any) => { + (fn as any)._defaultResult = { type: 'reject', value: val }; + return fn; + }; + fn.mockRejectedValueOnce = (val: any) => { + nextResults.push({ type: 'reject', value: val }); + return fn; + }; + + // Override fn to use default result when nextResults is empty + const wrappedFn: any = (...args: any[]) => { + calls.push(args); + const result = nextResults.shift(); + if (result) { + return result.type === 'resolve' ? Promise.resolve(result.value) : Promise.reject(result.value); + } + const def = (wrappedFn as any)._defaultResult; + if (def) { + return def.type === 'resolve' ? Promise.resolve(def.value) : Promise.reject(def.value); + } + return Promise.resolve({}); + }; + wrappedFn.mock = { calls }; + wrappedFn.mockResolvedValue = (val: any) => { + (wrappedFn as any)._defaultResult = { type: 'resolve', value: val }; + return wrappedFn; + }; + wrappedFn.mockResolvedValueOnce = (val: any) => { + nextResults.push({ type: 'resolve', value: val }); + return wrappedFn; + }; + wrappedFn.mockRejectedValue = (val: any) => { + (wrappedFn as any)._defaultResult = { type: 'reject', value: val }; + return wrappedFn; + }; + wrappedFn.mockRejectedValueOnce = (val: any) => { + nextResults.push({ type: 'reject', value: val }); + return wrappedFn; + }; + + return wrappedFn; +} + describe('test/OSSObjectStorageClient.test.ts', () => { let client: OSSObjectStorageClient; let mockOSS: { - put: ReturnType; - get: ReturnType; - append: ReturnType; - head: ReturnType; + put: ReturnType; + get: ReturnType; + append: ReturnType; + head: ReturnType; }; beforeEach(() => { mockOSS = { - put: vi.fn(), - get: vi.fn(), - append: vi.fn(), - head: vi.fn(), + put: mockFn(), + get: mockFn(), + append: mockFn(), + head: mockFn(), }; client = new OSSObjectStorageClient(mockOSS as unknown as OSSObject); }); diff --git a/tegg/core/agent-runtime/test/RunBuilder.test.ts b/tegg/core/agent-runtime/test/RunBuilder.test.ts index 638e17a930..a23f41b0b6 100644 --- a/tegg/core/agent-runtime/test/RunBuilder.test.ts +++ b/tegg/core/agent-runtime/test/RunBuilder.test.ts @@ -1,8 +1,12 @@ import assert from 'node:assert'; -import type { RunRecord, MessageObject } from '@eggjs/tegg-types/agent-runtime'; -import { RunStatus, AgentObjectType, AgentErrorCode } from '@eggjs/tegg-types/agent-runtime'; -import { InvalidRunStateTransitionError } from '@eggjs/tegg-types/agent-runtime'; +import type { RunRecord } from '@eggjs/tegg-types/agent-runtime'; +import { + RunStatus, + AgentObjectType, + AgentErrorCode, + InvalidRunStateTransitionError, +} from '@eggjs/tegg-types/agent-runtime'; import { describe, it } from 'vitest'; import { RunBuilder } from '../src/RunBuilder.ts'; @@ -45,16 +49,6 @@ describe('test/RunBuilder.test.ts', () => { status: RunStatus.Completed, startedAt: 1001, completedAt: 1002, - output: [ - { - id: 'msg_1', - object: 'thread.message', - createdAt: 1001, - role: 'assistant', - status: 'completed', - content: [], - }, - ], usage: { promptTokens: 10, completionTokens: 5, totalTokens: 15 }, metadata: { key: 'value' }, config: { maxIterations: 10 }, @@ -64,7 +58,6 @@ describe('test/RunBuilder.test.ts', () => { assert.equal(snap.status, RunStatus.Completed); assert.equal(snap.startedAt, 1001); assert.equal(snap.completedAt, 1002); - assert.equal(snap.output?.length, 1); assert.deepStrictEqual(snap.usage, { promptTokens: 10, completionTokens: 5, totalTokens: 15 }); assert.deepStrictEqual(snap.metadata, { key: 'value' }); assert.deepStrictEqual(snap.config, { maxIterations: 10 }); @@ -102,15 +95,12 @@ describe('test/RunBuilder.test.ts', () => { }); describe('complete', () => { - it('should transition in_progress → completed with output and usage', () => { + it('should transition in_progress → completed with usage', () => { const rb = RunBuilder.create(makeRunRecord(), 'thread_1'); rb.start(); - const output: MessageObject[] = [ - { id: 'msg_1', object: 'thread.message', createdAt: 1001, role: 'assistant', status: 'completed', content: [] }, - ]; const usage: RunUsage = { promptTokens: 10, completionTokens: 5, totalTokens: 15 }; - const update = rb.complete(output, usage); + const update = rb.complete(usage); assert.equal(update.status, RunStatus.Completed); assert.equal(typeof update.completedAt, 'number'); @@ -119,7 +109,6 @@ describe('test/RunBuilder.test.ts', () => { completionTokens: 5, totalTokens: 15, }); - assert.equal(update.output, output); const snap = rb.snapshot(); assert.equal(snap.status, RunStatus.Completed); @@ -134,7 +123,7 @@ describe('test/RunBuilder.test.ts', () => { const rb = RunBuilder.create(makeRunRecord(), 'thread_1'); rb.start(); - const update = rb.complete([]); + const update = rb.complete(); assert.equal(update.status, RunStatus.Completed); assert.equal(update.usage, undefined); @@ -144,7 +133,7 @@ describe('test/RunBuilder.test.ts', () => { it('should throw for non-in_progress status', () => { const rb = RunBuilder.create(makeRunRecord(), 'thread_1'); - assert.throws(() => rb.complete([]), InvalidRunStateTransitionError); + assert.throws(() => rb.complete(), InvalidRunStateTransitionError); }); }); @@ -171,6 +160,12 @@ describe('test/RunBuilder.test.ts', () => { assert.equal(update.status, RunStatus.Failed); }); + it('should allow failing from cancelling status (cancel watchdog timeout)', () => { + const rb = RunBuilder.create(makeRunRecord({ status: RunStatus.Cancelling }), 'thread_1'); + const update = rb.fail(new Error('commit timeout')); + assert.equal(update.status, RunStatus.Failed); + }); + it('should throw for terminal status', () => { const rb = RunBuilder.create(makeRunRecord({ status: RunStatus.Completed }), 'thread_1'); assert.throws(() => rb.fail(new Error('nope')), InvalidRunStateTransitionError); @@ -237,7 +232,7 @@ describe('test/RunBuilder.test.ts', () => { rb.start(); assert.equal(rb.snapshot().status, RunStatus.InProgress); - rb.complete([], { promptTokens: 1, completionTokens: 2, totalTokens: 3 }); + rb.complete({ promptTokens: 1, completionTokens: 2, totalTokens: 3 }); const snap = rb.snapshot(); assert.equal(snap.status, RunStatus.Completed); assert.ok(snap.startedAt); diff --git a/tegg/core/agent-runtime/test/helpers.ts b/tegg/core/agent-runtime/test/helpers.ts index 7f3b98d5bc..36ce2360ce 100644 --- a/tegg/core/agent-runtime/test/helpers.ts +++ b/tegg/core/agent-runtime/test/helpers.ts @@ -1,12 +1,58 @@ import type { ObjectStorageClient } from '@eggjs/tegg-types/agent-runtime'; +interface PutDelayRule { + pattern: RegExp; + ms: number; +} + +function shouldFailPut(key: string, exact: ReadonlySet, pattern: RegExp | null): boolean { + return exact.has(key) || (pattern !== null && pattern.test(key)); +} + +function findPutDelay(key: string, rules: readonly PutDelayRule[]): PutDelayRule | undefined { + return rules.find((r) => r.pattern.test(key)); +} + +function sleep(ms: number): Promise { + return new Promise((resolve) => { + setTimeout(resolve, ms); + }); +} + /** In-memory ObjectStorageClient backed by a Map — for testing only. */ export class MapStorageClient implements ObjectStorageClient { private readonly store = new Map(); + private readonly putFailureExact = new Set(); + private putFailurePattern: RegExp | null = null; + private readonly putDelays: PutDelayRule[] = []; + init?(): Promise; destroy?(): Promise; + keys(): string[] { + return [...this.store.keys()].sort(); + } + + keysWithPrefix(prefix: string): string[] { + return this.keys().filter((k) => k.startsWith(prefix)); + } + + failPutWhenKeyMatches(pattern: RegExp): void { + this.putFailurePattern = pattern; + } + + delayPutWhenKeyMatches(pattern: RegExp, ms: number): void { + this.putDelays.push({ pattern, ms }); + } + async put(key: string, value: string): Promise { + if (shouldFailPut(key, this.putFailureExact, this.putFailurePattern)) { + throw new Error(`MapStorageClient: simulated PUT failure for ${key}`); + } + const delay = findPutDelay(key, this.putDelays); + if (delay) { + await sleep(delay.ms); + } this.store.set(key, value); } @@ -20,13 +66,40 @@ export class MapStorageClient implements ObjectStorageClient { } } -/** MapStorageClient variant without append — tests the get-concat-put fallback path. */ +/** MapStorageClient variant without `append`. */ export class MapStorageClientWithoutAppend implements ObjectStorageClient { private readonly store = new Map(); + private readonly putFailureExact = new Set(); + private putFailurePattern: RegExp | null = null; + private readonly putDelays: PutDelayRule[] = []; + init?(): Promise; destroy?(): Promise; + keys(): string[] { + return [...this.store.keys()].sort(); + } + + keysWithPrefix(prefix: string): string[] { + return this.keys().filter((k) => k.startsWith(prefix)); + } + + failPutWhenKeyMatches(pattern: RegExp): void { + this.putFailurePattern = pattern; + } + + delayPutWhenKeyMatches(pattern: RegExp, ms: number): void { + this.putDelays.push({ pattern, ms }); + } + async put(key: string, value: string): Promise { + if (shouldFailPut(key, this.putFailureExact, this.putFailurePattern)) { + throw new Error(`MapStorageClientWithoutAppend: simulated PUT failure for ${key}`); + } + const delay = findPutDelay(key, this.putDelays); + if (delay) { + await sleep(delay.ms); + } this.store.set(key, value); } diff --git a/tegg/core/controller-decorator/src/decorator/agent/AgentController.ts b/tegg/core/controller-decorator/src/decorator/agent/AgentController.ts index 98d801f58d..32e90ac2d1 100644 --- a/tegg/core/controller-decorator/src/decorator/agent/AgentController.ts +++ b/tegg/core/controller-decorator/src/decorator/agent/AgentController.ts @@ -14,22 +14,32 @@ import { ControllerInfoUtil } from '../../util/ControllerInfoUtil.ts'; import { HTTPInfoUtil } from '../../util/HTTPInfoUtil.ts'; import { MethodInfoUtil } from '../../util/MethodInfoUtil.ts'; +interface AgentRouteParam { + index: number; + type: 'body' | 'pathParam' | 'query'; + name?: string; +} + interface AgentRouteDefinition { methodName: string; httpMethod: HTTPMethodEnum; path: string; - paramType?: 'body' | 'pathParam'; - paramName?: string; - hasParam: boolean; + params: AgentRouteParam[]; } // Default implementations for unimplemented methods. -// Methods with hasParam=true need function.length === 1 for param validation. +// function.length must match the param count for framework param validation. // Stubs are marked with Symbol.for('AGENT_NOT_IMPLEMENTED') so agent-runtime // can distinguish them from user-defined methods at enhancement time. -function createNotImplemented(methodName: string, hasParam: boolean) { +function createNotImplemented(methodName: string, paramCount: number) { let fn; - if (hasParam) { + if (paramCount >= 2) { + // eslint-disable-next-line @typescript-eslint/no-unused-vars + fn = async function (_a: unknown, _b: unknown) { + throw new Error(`${methodName} not implemented`); + }; + } else if (paramCount === 1) { + // eslint-disable-next-line @typescript-eslint/no-unused-vars fn = async function (_arg: unknown) { throw new Error(`${methodName} not implemented`); }; @@ -47,52 +57,58 @@ const AGENT_ROUTES: AgentRouteDefinition[] = [ methodName: 'createThread', httpMethod: HTTPMethodEnum.POST, path: '/threads', - hasParam: false, + params: [], }, { methodName: 'getThread', httpMethod: HTTPMethodEnum.GET, path: '/threads/:id', - paramType: 'pathParam', - paramName: 'id', - hasParam: true, + params: [{ index: 0, type: 'pathParam', name: 'id' }], + }, + { + methodName: 'getLatestRunId', + httpMethod: HTTPMethodEnum.GET, + path: '/threads/:id/latest-run', + params: [{ index: 0, type: 'pathParam', name: 'id' }], }, { methodName: 'asyncRun', httpMethod: HTTPMethodEnum.POST, path: '/runs', - paramType: 'body', - hasParam: true, + params: [{ index: 0, type: 'body' }], }, { methodName: 'streamRun', httpMethod: HTTPMethodEnum.POST, path: '/runs/stream', - paramType: 'body', - hasParam: true, + params: [{ index: 0, type: 'body' }], + }, + { + methodName: 'getRunStream', + httpMethod: HTTPMethodEnum.GET, + path: '/runs/:id/stream', + params: [ + { index: 0, type: 'pathParam', name: 'id' }, + { index: 1, type: 'query', name: 'lastSeq' }, + ], }, { methodName: 'syncRun', httpMethod: HTTPMethodEnum.POST, path: '/runs/wait', - paramType: 'body', - hasParam: true, + params: [{ index: 0, type: 'body' }], }, { methodName: 'getRun', httpMethod: HTTPMethodEnum.GET, path: '/runs/:id', - paramType: 'pathParam', - paramName: 'id', - hasParam: true, + params: [{ index: 0, type: 'pathParam', name: 'id' }], }, { methodName: 'cancelRun', httpMethod: HTTPMethodEnum.POST, path: '/runs/:id/cancel', - paramType: 'pathParam', - paramName: 'id', - hasParam: true, + params: [{ index: 0, type: 'pathParam', name: 'id' }], }, ]; @@ -119,7 +135,7 @@ export function AgentController(): (constructor: EggProtoImplClass) => void { for (const route of AGENT_ROUTES) { // Inject default implementation if method not defined if (!constructor.prototype[route.methodName]) { - constructor.prototype[route.methodName] = createNotImplemented(route.methodName, route.hasParam); + constructor.prototype[route.methodName] = createNotImplemented(route.methodName, route.params.length); } // Set method controller type @@ -132,11 +148,16 @@ export function AgentController(): (constructor: EggProtoImplClass) => void { HTTPInfoUtil.setHTTPMethodPath(route.path, constructor, route.methodName); // Set parameter metadata - if (route.paramType === 'body') { - HTTPInfoUtil.setHTTPMethodParamType(HTTPParamType.BODY, 0, constructor, route.methodName); - } else if (route.paramType === 'pathParam') { - HTTPInfoUtil.setHTTPMethodParamType(HTTPParamType.PARAM, 0, constructor, route.methodName); - HTTPInfoUtil.setHTTPMethodParamName(route.paramName!, 0, constructor, route.methodName); + for (const param of route.params) { + if (param.type === 'body') { + HTTPInfoUtil.setHTTPMethodParamType(HTTPParamType.BODY, param.index, constructor, route.methodName); + } else if (param.type === 'pathParam') { + HTTPInfoUtil.setHTTPMethodParamType(HTTPParamType.PARAM, param.index, constructor, route.methodName); + HTTPInfoUtil.setHTTPMethodParamName(param.name!, param.index, constructor, route.methodName); + } else if (param.type === 'query') { + HTTPInfoUtil.setHTTPMethodParamType(HTTPParamType.QUERY, param.index, constructor, route.methodName); + HTTPInfoUtil.setHTTPMethodParamName(param.name!, param.index, constructor, route.methodName); + } } } diff --git a/tegg/core/controller-decorator/src/decorator/agent/AgentHandler.ts b/tegg/core/controller-decorator/src/decorator/agent/AgentHandler.ts index 634950674e..abbb1d9532 100644 --- a/tegg/core/controller-decorator/src/decorator/agent/AgentHandler.ts +++ b/tegg/core/controller-decorator/src/decorator/agent/AgentHandler.ts @@ -3,18 +3,31 @@ import type { ThreadObjectWithMessages, CreateRunInput, RunObject, - AgentStreamMessage, + AgentMessage, } from '@eggjs/tegg-types/agent-runtime'; // Interface for AgentController classes. The `execRun` method is required — // the framework uses it to auto-wire thread/run management, store persistence, // SSE streaming, async execution, and cancellation via smart defaults. export interface AgentHandler { - execRun(input: CreateRunInput, signal?: AbortSignal): AsyncGenerator; + execRun(input: CreateRunInput, signal?: AbortSignal): AsyncGenerator; /** Create the AgentStore used to persist threads and runs. */ createStore(): Promise; + /** + * Optional hook to decide whether the executor's underlying session has + * been committed to persistent storage (for example the Claude Code SDK + * jsonl file on disk). The runtime calls this each time a new message is + * yielded from `execRun`; once it returns true, `cancelRun` is allowed to + * abort and persist the thread. + * + * When not implemented, the runtime uses a default heuristic: any message + * with `type !== 'system'` counts as committed (covers the Claude Code SDK + * case where `system/init` is emitted before the session is fully written). + */ + isSessionCommitted?(msg: AgentMessage, history: AgentMessage[]): boolean | Promise; createThread?(): Promise; getThread?(threadId: string): Promise; + getLatestRunId?(threadId: string): Promise<{ threadId: string; runId: string | null }>; asyncRun?(input: CreateRunInput): Promise; streamRun?(input: CreateRunInput): Promise; syncRun?(input: CreateRunInput): Promise; diff --git a/tegg/core/controller-decorator/test/AgentController.test.ts b/tegg/core/controller-decorator/test/AgentController.test.ts index ccc2cb2aef..5d77f3d051 100644 --- a/tegg/core/controller-decorator/test/AgentController.test.ts +++ b/tegg/core/controller-decorator/test/AgentController.test.ts @@ -8,12 +8,13 @@ import { ControllerMetaBuilderFactory, BodyParamMeta, PathParamMeta, + QueryParamMeta, ControllerInfoUtil, MethodInfoUtil, - HTTPInfoUtil, } from '../src/index.ts'; import { HTTPControllerMeta } from '../src/model/index.ts'; -import { AgentFooController } from './fixtures/AgentFooController.js'; +import { HTTPInfoUtil } from '../src/util/HTTPInfoUtil.ts'; +import { AgentFooController } from './fixtures/AgentFooController.ts'; describe('core/controller-decorator/test/AgentController.test.ts', () => { describe('decorator metadata', () => { @@ -36,6 +37,7 @@ describe('core/controller-decorator/test/AgentController.test.ts', () => { const methodRoutes = [ { methodName: 'createThread', httpMethod: HTTPMethodEnum.POST, path: '/threads' }, { methodName: 'getThread', httpMethod: HTTPMethodEnum.GET, path: '/threads/:id' }, + { methodName: 'getLatestRunId', httpMethod: HTTPMethodEnum.GET, path: '/threads/:id/latest-run' }, { methodName: 'asyncRun', httpMethod: HTTPMethodEnum.POST, path: '/runs' }, { methodName: 'streamRun', httpMethod: HTTPMethodEnum.POST, path: '/runs/stream' }, { methodName: 'syncRun', httpMethod: HTTPMethodEnum.POST, path: '/runs/wait' }, @@ -106,7 +108,16 @@ describe('core/controller-decorator/test/AgentController.test.ts', () => { describe('context index', () => { it('should not set contextIndex on any method', () => { - const methods = ['createThread', 'getThread', 'asyncRun', 'streamRun', 'syncRun', 'getRun', 'cancelRun']; + const methods = [ + 'createThread', + 'getThread', + 'getLatestRunId', + 'asyncRun', + 'streamRun', + 'syncRun', + 'getRun', + 'cancelRun', + ]; for (const methodName of methods) { const contextIndex = MethodInfoUtil.getMethodContextIndex(AgentFooController, methodName); assert.strictEqual(contextIndex, undefined, `${methodName} should not have contextIndex`); @@ -128,11 +139,21 @@ describe('core/controller-decorator/test/AgentController.test.ts', () => { }); describe('default implementations', () => { - it('should inject default stubs for all 7 route methods', () => { + it('should inject default stubs for all 9 route methods', () => { // AgentFooController only implements execRun (smart defaults pattern) - // All 7 route methods should have stub defaults that throw + // All 9 route methods should have stub defaults that throw const proto = AgentFooController.prototype as any; - const routeMethods = ['createThread', 'getThread', 'asyncRun', 'streamRun', 'syncRun', 'getRun', 'cancelRun']; + const routeMethods = [ + 'createThread', + 'getThread', + 'getLatestRunId', + 'asyncRun', + 'streamRun', + 'getRunStream', + 'syncRun', + 'getRun', + 'cancelRun', + ]; for (const methodName of routeMethods) { assert(typeof proto[methodName] === 'function', `${methodName} should be a function`); assert.strictEqual( @@ -146,8 +167,10 @@ describe('core/controller-decorator/test/AgentController.test.ts', () => { const stubMethods = [ { name: 'createThread', args: [] }, { name: 'getThread', args: ['thread_1'] }, + { name: 'getLatestRunId', args: ['thread_1'] }, { name: 'asyncRun', args: [{ input: { messages: [] } }] }, { name: 'streamRun', args: [{ input: { messages: [] } }] }, + { name: 'getRunStream', args: ['run_1', '0'] }, { name: 'syncRun', args: [{ input: { messages: [] } }] }, { name: 'getRun', args: ['run_1'] }, { name: 'cancelRun', args: ['run_1'] }, @@ -162,10 +185,10 @@ describe('core/controller-decorator/test/AgentController.test.ts', () => { }); describe('HTTPControllerMetaBuilder integration', () => { - it('should build metadata with 7 HTTPMethodMeta entries', () => { + it('should build metadata with 9 HTTPMethodMeta entries', () => { const meta = ControllerMetaBuilderFactory.build(AgentFooController, ControllerType.HTTP) as HTTPControllerMeta; assert(meta); - assert.strictEqual(meta.methods.length, 7); + assert.strictEqual(meta.methods.length, 9); assert.strictEqual(meta.path, '/api/v1'); }); @@ -182,6 +205,11 @@ describe('core/controller-decorator/test/AgentController.test.ts', () => { assert.strictEqual(getThread.method, HTTPMethodEnum.GET); assert.deepStrictEqual(getThread.paramMap, new Map([[0, new PathParamMeta('id')]])); + const getLatestRunId = meta.methods.find((m) => m.name === 'getLatestRunId')!; + assert.strictEqual(getLatestRunId.path, '/threads/:id/latest-run'); + assert.strictEqual(getLatestRunId.method, HTTPMethodEnum.GET); + assert.deepStrictEqual(getLatestRunId.paramMap, new Map([[0, new PathParamMeta('id')]])); + const asyncRun = meta.methods.find((m) => m.name === 'asyncRun')!; assert.strictEqual(asyncRun.path, '/runs'); assert.strictEqual(asyncRun.method, HTTPMethodEnum.POST); @@ -192,6 +220,17 @@ describe('core/controller-decorator/test/AgentController.test.ts', () => { assert.strictEqual(streamRun.method, HTTPMethodEnum.POST); assert.deepStrictEqual(streamRun.paramMap, new Map([[0, new BodyParamMeta()]])); + const getRunStream = meta.methods.find((m) => m.name === 'getRunStream')!; + assert.strictEqual(getRunStream.path, '/runs/:id/stream'); + assert.strictEqual(getRunStream.method, HTTPMethodEnum.GET); + assert.deepStrictEqual( + getRunStream.paramMap, + new Map([ + [0, new PathParamMeta('id')], + [1, new QueryParamMeta('lastSeq')], + ]), + ); + const syncRun = meta.methods.find((m) => m.name === 'syncRun')!; assert.strictEqual(syncRun.path, '/runs/wait'); assert.strictEqual(syncRun.method, HTTPMethodEnum.POST); diff --git a/tegg/core/controller-decorator/test/fixtures/AgentFooController.ts b/tegg/core/controller-decorator/test/fixtures/AgentFooController.ts index ae7287ee56..35e69bfb72 100644 --- a/tegg/core/controller-decorator/test/fixtures/AgentFooController.ts +++ b/tegg/core/controller-decorator/test/fixtures/AgentFooController.ts @@ -1,4 +1,4 @@ -import type { CreateRunInput, AgentStreamMessage } from '@eggjs/tegg-types/agent-runtime'; +import type { CreateRunInput, AgentMessage } from '@eggjs/tegg-types/agent-runtime'; import { AgentController } from '../../src/decorator/agent/AgentController.ts'; import type { AgentHandler } from '../../src/decorator/agent/AgentHandler.ts'; @@ -10,7 +10,7 @@ export class AgentFooController implements AgentHandler { return new Map(); } - async *execRun(input: CreateRunInput): AsyncGenerator { + async *execRun(input: CreateRunInput): AsyncGenerator { const messages = input.input.messages; yield { type: 'assistant', diff --git a/tegg/core/tegg/src/agent.ts b/tegg/core/tegg/src/agent.ts index 38552c0fd3..b8d986a1a2 100644 --- a/tegg/core/tegg/src/agent.ts +++ b/tegg/core/tegg/src/agent.ts @@ -11,17 +11,24 @@ export type { ThreadRecord, RunRecord, CreateRunInput, - AgentStreamMessage, RunObject, ThreadObject, ThreadObjectWithMessages, - MessageObject, InputMessage, InputContentPart, - MessageContentBlock, - TextContentBlock, - MessageDeltaObject, + TextInputContentPart, + ToolUseInputContentPart, + ToolResultInputContentPart, + GenericInputContentPart, AgentRunConfig, - AgentRunUsage, + GetThreadOptions, RunStatus, + // SDK-aligned message types + AgentMessage, + SDKSystemMessage, + SDKStreamEvent, + SDKUserMessage, + SDKAssistantMessage, + SDKResultMessage, + SDKGenericMessage, } from '@eggjs/agent-runtime'; diff --git a/tegg/core/types/src/agent-runtime/AgentMessage.ts b/tegg/core/types/src/agent-runtime/AgentMessage.ts index 5a0c6c4dc5..f7e24bfba7 100644 --- a/tegg/core/types/src/agent-runtime/AgentMessage.ts +++ b/tegg/core/types/src/agent-runtime/AgentMessage.ts @@ -1,39 +1,96 @@ -// ===== Content block types ===== +// ===== Input content types (for CreateRunInput) ===== -export const ContentBlockType = { - Text: 'text', -} as const; -export type ContentBlockType = (typeof ContentBlockType)[keyof typeof ContentBlockType]; +export interface TextInputContentPart { + type: 'text'; + text: string; +} -// ===== Content types ===== +export interface ToolUseInputContentPart { + type: 'tool_use'; + id: string; + name: string; + input: Record; +} -export interface InputContentPart { - type: typeof ContentBlockType.Text; - text: string; +export interface ToolResultInputContentPart { + type: 'tool_result'; + tool_use_id: string; + content?: string | { type: string; text?: string; [key: string]: unknown }[]; + is_error?: boolean; } -export interface TextContentBlock { - type: typeof ContentBlockType.Text; - text: { value: string; annotations: unknown[] }; +export interface GenericInputContentPart { + type: string; + [key: string]: unknown; } -export type MessageContentBlock = TextContentBlock; +export type InputContentPart = + | TextInputContentPart + | ToolUseInputContentPart + | ToolResultInputContentPart + | GenericInputContentPart; -// ===== Input / Output message types ===== +// ===== Input message (for CreateRunInput) ===== export interface InputMessage { role: string; - content: string | { type: string; text: string }[]; + content: string | InputContentPart[]; metadata?: Record; } -export interface MessageObject { - id: string; - object: string; - createdAt: number; - role: string; - status: string; - content: MessageContentBlock[]; - runId?: string; - threadId?: string; +// ===== AgentMessage — aligned with Claude Agent SDK SDKMessage ===== +// Lightweight subset of SDK message types. The framework only needs to +// discriminate on `type` for a handful of core message kinds; everything +// else passes through as SDKGenericMessage. + +export interface SDKSystemMessage { + type: 'system'; + subtype: string; + session_id?: string; + [key: string]: unknown; +} + +export interface SDKStreamEvent { + type: 'stream_event'; + event: unknown; + session_id?: string; + [key: string]: unknown; } + +export interface SDKUserMessage { + type: 'user'; + message: unknown; + [key: string]: unknown; +} + +export interface SDKAssistantMessage { + type: 'assistant'; + message: unknown; + [key: string]: unknown; +} + +export interface SDKResultMessage { + type: 'result'; + subtype: string; + usage?: { + input_tokens?: number; + output_tokens?: number; + cache_creation_input_tokens?: number; + cache_read_input_tokens?: number; + [key: string]: unknown; + }; + [key: string]: unknown; +} + +export interface SDKGenericMessage { + type: string; + [key: string]: unknown; +} + +export type AgentMessage = + | SDKSystemMessage + | SDKStreamEvent + | SDKUserMessage + | SDKAssistantMessage + | SDKResultMessage + | SDKGenericMessage; diff --git a/tegg/core/types/src/agent-runtime/AgentRuntime.ts b/tegg/core/types/src/agent-runtime/AgentRuntime.ts index ec6f7c19a4..2f4213aaa9 100644 --- a/tegg/core/types/src/agent-runtime/AgentRuntime.ts +++ b/tegg/core/types/src/agent-runtime/AgentRuntime.ts @@ -1,41 +1,5 @@ -import type { InputContentPart, MessageContentBlock } from './AgentMessage.ts'; -import type { AgentRunConfig, InputMessage, MessageObject, RunStatus } from './AgentStore.ts'; - -export { ContentBlockType } from './AgentMessage.ts'; -export type { InputContentPart, MessageContentBlock, TextContentBlock } from './AgentMessage.ts'; - -// ===== Message roles ===== - -export const MessageRole = { - User: 'user', - Assistant: 'assistant', - System: 'system', -} as const; -export type MessageRole = (typeof MessageRole)[keyof typeof MessageRole]; - -// ===== Message statuses ===== - -export const MessageStatus = { - InProgress: 'in_progress', - Incomplete: 'incomplete', - Completed: 'completed', -} as const; -export type MessageStatus = (typeof MessageStatus)[keyof typeof MessageStatus]; - -// ===== SSE events ===== - -export const AgentSSEEvent = { - ThreadRunCreated: 'thread.run.created', - ThreadRunInProgress: 'thread.run.in_progress', - ThreadRunCompleted: 'thread.run.completed', - ThreadRunFailed: 'thread.run.failed', - ThreadRunCancelled: 'thread.run.cancelled', - ThreadMessageCreated: 'thread.message.created', - ThreadMessageDelta: 'thread.message.delta', - ThreadMessageCompleted: 'thread.message.completed', - Done: 'done', -} as const; -export type AgentSSEEvent = (typeof AgentSSEEvent)[keyof typeof AgentSSEEvent]; +import type { AgentMessage } from './AgentMessage.ts'; +import type { AgentRunConfig, RunStatus } from './AgentStore.ts'; // ===== Error codes ===== @@ -54,7 +18,7 @@ export interface ThreadObject { } export interface ThreadObjectWithMessages extends ThreadObject { - messages: MessageObject[]; + messages: AgentMessage[]; } // ===== Run objects ===== @@ -72,7 +36,6 @@ export interface RunObject { failedAt?: number | null; usage?: { promptTokens: number; completionTokens: number; totalTokens: number } | null; metadata?: Record; - output?: MessageObject[]; config?: AgentRunConfig; } @@ -80,37 +43,53 @@ export interface RunObject { export interface CreateRunInput { threadId?: string; + /** + * Populated by AgentRuntime before calling execRun. + * - true: threadId was provided (resume existing conversation) + * - false: no threadId provided, new thread created + */ + isResume?: boolean; input: { - messages: InputMessage[]; + messages: import('./AgentMessage.ts').InputMessage[]; }; config?: AgentRunConfig; + /** + * Metadata for the run. Stored verbatim on the run record, and additionally + * shallow-merged into the thread metadata (`meta.json`): + * - For an auto-created thread, it initializes the thread metadata. + * - For an existing thread, the keys are shallow-merged: new values overwrite + * matching keys, while keys not present are preserved. + * - An omitted or empty object leaves the thread metadata unchanged. + */ metadata?: Record; } -// ===== Message delta ===== +// ===== Thread input ===== -export interface MessageDeltaObject { - id: string; - object: 'thread.message.delta'; - delta: { - content: MessageContentBlock[]; - }; +/** + * Options for {@link AgentRuntime.createThread}. + * + * `metadata` is forwarded verbatim to {@link AgentStore.createThread} so callers + * can persist additional business semantics on the thread record (e.g. the + * resolved agent name, owning sandbox id, trace id). + */ +export interface CreateThreadOptions { + metadata?: Record; } -// ===== Stream message types ===== +// ===== Stream event (TaskEvent-style wrapper) ===== -export interface AgentStreamMessagePayload { - role?: string; - content: string | InputContentPart[]; +export interface StreamEvent { + seq: number; + type: string; + data: unknown; + ts: number; } -export interface AgentRunUsage { - promptTokens?: number; - completionTokens?: number; -} +// ===== Get thread options ===== -export interface AgentStreamMessage { - type?: string; - message?: AgentStreamMessagePayload; - usage?: AgentRunUsage; +export interface GetThreadOptions { + /** When true, return all message types (system, result, stream_event, etc.). + * Defaults to false — only user and assistant messages are returned. */ + includeAllMessages?: boolean; } diff --git a/tegg/core/types/src/agent-runtime/AgentStore.ts b/tegg/core/types/src/agent-runtime/AgentStore.ts index e2337cfe9b..bb9f6bcb3d 100644 --- a/tegg/core/types/src/agent-runtime/AgentStore.ts +++ b/tegg/core/types/src/agent-runtime/AgentStore.ts @@ -1,14 +1,11 @@ -import type { InputMessage, MessageObject } from './AgentMessage.ts'; - -export type { InputMessage, MessageObject } from './AgentMessage.ts'; +import type { AgentMessage, InputMessage } from './AgentMessage.ts'; +import type { GetThreadOptions } from './AgentRuntime.ts'; // ===== Object types ===== export const AgentObjectType = { Thread: 'thread', ThreadRun: 'thread.run', - ThreadMessage: 'thread.message', - ThreadMessageDelta: 'thread.message.delta', } as const; export type AgentObjectType = (typeof AgentObjectType)[keyof typeof AgentObjectType]; @@ -38,13 +35,22 @@ export interface ThreadRecord { id: string; object: typeof AgentObjectType.Thread; /** - * Logically belongs to the thread. In OSSAgentStore the messages are stored - * separately as a JSONL file and assembled on read — callers should treat - * this as a unified view regardless of the underlying storage layout. + * All messages in the thread, stored as SDK-format AgentMessage objects. + * In OSSAgentStore the messages are stored separately as a JSONL file + * and assembled on read — callers should treat this as a unified view + * regardless of the underlying storage layout. */ - messages: MessageObject[]; + messages: AgentMessage[]; metadata: Record; createdAt: number; // Unix seconds + /** + * Id of the most recently created run on this thread, maintained by + * `createRun` as a best-effort pointer so callers can resolve a thread's + * latest run without an external index. Absent on threads that have no + * runs yet, and on threads created before this field existed — a missing + * value must be treated as "no recent run". + */ + latestRunId?: string; } export interface RunRecord { @@ -53,7 +59,6 @@ export interface RunRecord { threadId?: string; status: RunStatus; input: InputMessage[]; - output?: MessageObject[]; lastError?: { code: string; message: string } | null; usage?: { promptTokens: number; completionTokens: number; totalTokens: number } | null; config?: AgentRunConfig; @@ -71,8 +76,13 @@ export interface AgentStore { init?(): Promise; destroy?(): Promise; createThread(metadata?: Record): Promise; - getThread(threadId: string): Promise; - appendMessages(threadId: string, messages: MessageObject[]): Promise; + getThread(threadId: string, options?: GetThreadOptions): Promise; + /** + * Shallow-merge metadata into an existing thread. + * New values overwrite matching keys; omitted keys are preserved. + */ + updateThreadMetadata?(threadId: string, metadata: Record): Promise; + appendMessages(threadId: string, messages: AgentMessage[]): Promise; createRun( input: InputMessage[], threadId?: string, @@ -80,5 +90,12 @@ export interface AgentStore { metadata?: Record, ): Promise; getRun(runId: string): Promise; + /** + * Id of the most recent run created on the thread, or `null` when the + * thread exists but has no recorded run (including threads created before + * run tracking existed). Throws AgentNotFoundError when the thread itself + * does not exist. + */ + getLatestRunId(threadId: string): Promise; updateRun(runId: string, updates: Partial): Promise; } diff --git a/tegg/core/types/src/agent-runtime/errors.ts b/tegg/core/types/src/agent-runtime/errors.ts index 76fbf3265b..61dce9467e 100644 --- a/tegg/core/types/src/agent-runtime/errors.ts +++ b/tegg/core/types/src/agent-runtime/errors.ts @@ -4,7 +4,7 @@ * to set the corresponding HTTP response status code. */ export class AgentNotFoundError extends Error { - status: number = 404; + status = 404; constructor(message: string) { super(message); @@ -12,12 +12,24 @@ export class AgentNotFoundError extends Error { } } +/** + * Error thrown when an agent API request contains invalid input. + */ +export class AgentInvalidRequestError extends Error { + status = 400; + + constructor(message: string) { + super(message); + this.name = 'AgentInvalidRequestError'; + } +} + /** * Error thrown when an operation conflicts with the current state * (e.g., cancelling a completed run). */ export class AgentConflictError extends Error { - status: number = 409; + status = 409; constructor(message: string) { super(message); @@ -30,10 +42,26 @@ export class AgentConflictError extends Error { * (e.g., calling `complete()` on a queued run). */ export class InvalidRunStateTransitionError extends Error { - status: number = 409; + status = 409; constructor(from: string, to: string) { super(`Invalid run state transition: '${from}' -> '${to}'`); this.name = 'InvalidRunStateTransitionError'; } } + +/** + * Error thrown when cancelRun waits for the executor's session to be + * committed to persistent storage (e.g. Claude Code SDK jsonl on disk) + * but the commit never arrives within the configured timeout. The run is + * transitioned to `failed` rather than `cancelled` to reflect that the + * executor never reached a resumable state. + */ +export class AgentTimeoutError extends Error { + status = 408; + + constructor(message: string) { + super(message); + this.name = 'AgentTimeoutError'; + } +} diff --git a/tegg/core/types/test/__snapshots__/index.test.ts.snap b/tegg/core/types/test/__snapshots__/index.test.ts.snap index 440fdf4cac..3a01927c00 100644 --- a/tegg/core/types/test/__snapshots__/index.test.ts.snap +++ b/tegg/core/types/test/__snapshots__/index.test.ts.snap @@ -12,24 +12,13 @@ exports[`should export stable 1`] = ` "AgentErrorCode": { "ExecError": "EXEC_ERROR", }, + "AgentInvalidRequestError": [Function], "AgentNotFoundError": [Function], "AgentObjectType": { "Thread": "thread", - "ThreadMessage": "thread.message", - "ThreadMessageDelta": "thread.message.delta", "ThreadRun": "thread.run", }, - "AgentSSEEvent": { - "Done": "done", - "ThreadMessageCompleted": "thread.message.completed", - "ThreadMessageCreated": "thread.message.created", - "ThreadMessageDelta": "thread.message.delta", - "ThreadRunCancelled": "thread.run.cancelled", - "ThreadRunCompleted": "thread.run.completed", - "ThreadRunCreated": "thread.run.created", - "ThreadRunFailed": "thread.run.failed", - "ThreadRunInProgress": "thread.run.in_progress", - }, + "AgentTimeoutError": [Function], "CONSTRUCTOR_QUALIFIER_META_DATA": Symbol(EggPrototype#constructorQualifier), "CONTROLLER_ACL": Symbol(EggPrototype#controller#acl), "CONTROLLER_AGENT_CONTROLLER": Symbol(EggPrototype#controller#agent#isAgent), @@ -112,9 +101,6 @@ exports[`should export stable 1`] = ` "ZLIB": "ZLIB", }, "ConfigSourceQualifierAttribute": Symbol(Qualifier.ConfigSource), - "ContentBlockType": { - "Text": "text", - }, "ControllerType": { "HEADERS": "HEADERS", "HTTP": "HTTP", @@ -225,16 +211,6 @@ exports[`should export stable 1`] = ` "MODEL_DATA_SOURCE": Symbol(EggPrototype#model#dataSource), "MODEL_DATA_TABLE_NAME": Symbol(EggPrototype#model#tableName), "MODEL_PROTO_IMPL_TYPE": "MODEL_PROTO", - "MessageRole": { - "Assistant": "assistant", - "System": "system", - "User": "user", - }, - "MessageStatus": { - "Completed": "completed", - "InProgress": "in_progress", - "Incomplete": "incomplete", - }, "MethodType": { "HTTP": "HTTP", "MESSAGE": "MESSAGE", diff --git a/tegg/plugin/controller/src/lib/AgentControllerObject.ts b/tegg/plugin/controller/src/lib/AgentControllerObject.ts index 4d72be6df8..a89283615e 100644 --- a/tegg/plugin/controller/src/lib/AgentControllerObject.ts +++ b/tegg/plugin/controller/src/lib/AgentControllerObject.ts @@ -19,11 +19,19 @@ import type { EggLogger } from 'egg'; import { AgentControllerProto } from './AgentControllerProto.ts'; /** Method names that can be delegated to AgentRuntime. */ -type AgentMethodName = 'createThread' | 'getThread' | 'asyncRun' | 'syncRun' | 'getRun' | 'cancelRun'; +type AgentMethodName = + | 'createThread' + | 'getThread' + | 'getLatestRunId' + | 'asyncRun' + | 'syncRun' + | 'getRun' + | 'cancelRun'; const AGENT_METHOD_NAMES: AgentMethodName[] = [ 'createThread', 'getThread', + 'getLatestRunId', 'asyncRun', 'syncRun', 'getRun', @@ -245,6 +253,20 @@ export class AgentControllerObject implements EggObject { return runtime.streamRun(input, writer); }; } + + // getRunStream: always delegate to runtime (no user override needed) + // lastSeq comes from query string as a string, needs parseInt + instance['getRunStream'] = async (runId: string, lastSeq?: string): Promise => { + const runtimeCtx = ContextHandler.getContext(); + if (!runtimeCtx) { + throw new Error('getRunStream must be called within a request context'); + } + const eggCtx = runtimeCtx.get(EGG_CONTEXT); + eggCtx.respond = false; + const writer = new HttpSSEWriter(eggCtx.res); + const seq = parseInt(lastSeq as string, 10) || 0; + return runtime.getRunStream(runId, writer, seq); + }; } static async createObject(