chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,790 @@
|
||||
import { describe, it, expect, afterEach } from "vitest";
|
||||
import { SqliteAgentRunner } from "..";
|
||||
import {
|
||||
AbstractAgent,
|
||||
BaseEvent,
|
||||
EventType,
|
||||
Message,
|
||||
RunAgentInput,
|
||||
RunErrorEvent,
|
||||
RunFinishedEvent,
|
||||
RunStartedEvent,
|
||||
} from "@ag-ui/client";
|
||||
import { EMPTY, Subscription, firstValueFrom, from } from "rxjs";
|
||||
import { toArray } from "rxjs/operators";
|
||||
|
||||
type RunCallbacks = {
|
||||
onEvent: (event: { event: BaseEvent }) => void;
|
||||
onNewMessage?: (args: { message: Message }) => void;
|
||||
onRunStartedEvent?: () => void;
|
||||
};
|
||||
|
||||
const createdRunners: SqliteAgentRunner[] = [];
|
||||
|
||||
afterEach(() => {
|
||||
while (createdRunners.length > 0) {
|
||||
const runner = createdRunners.pop();
|
||||
runner?.close();
|
||||
}
|
||||
});
|
||||
|
||||
function createRunner(): SqliteAgentRunner {
|
||||
const runner = new SqliteAgentRunner();
|
||||
createdRunners.push(runner);
|
||||
return runner;
|
||||
}
|
||||
|
||||
interface EmitAgentOptions {
|
||||
events?: BaseEvent[];
|
||||
emitDefaultRunStarted?: boolean;
|
||||
includeRunFinished?: boolean;
|
||||
runFinishedEvent?: RunFinishedEvent;
|
||||
afterEvent?: (args: {
|
||||
event: BaseEvent;
|
||||
index: number;
|
||||
}) => void | Promise<void>;
|
||||
}
|
||||
|
||||
class EmitAgent extends AbstractAgent {
|
||||
constructor(private readonly options: EmitAgentOptions = {}) {
|
||||
super();
|
||||
}
|
||||
|
||||
async runAgent(input: RunAgentInput, callbacks: RunCallbacks): Promise<void> {
|
||||
const {
|
||||
emitDefaultRunStarted = true,
|
||||
includeRunFinished = true,
|
||||
runFinishedEvent,
|
||||
afterEvent,
|
||||
} = this.options;
|
||||
const scriptedEvents = this.options.events ?? [];
|
||||
|
||||
let index = 0;
|
||||
const emit = async (event: BaseEvent) => {
|
||||
callbacks.onEvent({ event });
|
||||
if (event.type === EventType.RUN_STARTED) {
|
||||
callbacks.onRunStartedEvent?.();
|
||||
}
|
||||
await afterEvent?.({ event, index });
|
||||
index += 1;
|
||||
};
|
||||
|
||||
if (emitDefaultRunStarted) {
|
||||
const runStarted: RunStartedEvent = {
|
||||
type: EventType.RUN_STARTED,
|
||||
threadId: input.threadId,
|
||||
runId: input.runId,
|
||||
parentRunId: input.parentRunId,
|
||||
};
|
||||
await emit(runStarted);
|
||||
}
|
||||
|
||||
for (const event of scriptedEvents) {
|
||||
await emit(event);
|
||||
}
|
||||
|
||||
const hasRunFinishedEvent =
|
||||
scriptedEvents.some((event) => event.type === EventType.RUN_FINISHED) ||
|
||||
runFinishedEvent?.type === EventType.RUN_FINISHED;
|
||||
|
||||
if (includeRunFinished && !hasRunFinishedEvent) {
|
||||
const finishEvent: RunFinishedEvent = runFinishedEvent ?? {
|
||||
type: EventType.RUN_FINISHED,
|
||||
threadId: input.threadId,
|
||||
runId: input.runId,
|
||||
};
|
||||
await emit(finishEvent);
|
||||
}
|
||||
}
|
||||
|
||||
clone(): AbstractAgent {
|
||||
return new EmitAgent({
|
||||
...this.options,
|
||||
events: this.options.events ? [...this.options.events] : undefined,
|
||||
});
|
||||
}
|
||||
|
||||
protected run(): ReturnType<AbstractAgent["run"]> {
|
||||
return EMPTY;
|
||||
}
|
||||
|
||||
protected connect(): ReturnType<AbstractAgent["connect"]> {
|
||||
return EMPTY;
|
||||
}
|
||||
}
|
||||
|
||||
class ReplayAgent extends AbstractAgent {
|
||||
constructor(
|
||||
private readonly replayEvents: BaseEvent[],
|
||||
threadId: string,
|
||||
) {
|
||||
super({ threadId });
|
||||
}
|
||||
|
||||
async runAgent(): Promise<void> {
|
||||
throw new Error("not used");
|
||||
}
|
||||
|
||||
protected run(): ReturnType<AbstractAgent["run"]> {
|
||||
return EMPTY;
|
||||
}
|
||||
|
||||
protected connect(): ReturnType<AbstractAgent["connect"]> {
|
||||
return from(this.replayEvents);
|
||||
}
|
||||
}
|
||||
|
||||
class RunnerConnectAgent extends AbstractAgent {
|
||||
constructor(
|
||||
private readonly runner: SqliteAgentRunner,
|
||||
threadId: string,
|
||||
) {
|
||||
super({ threadId });
|
||||
}
|
||||
|
||||
async runAgent(): Promise<void> {
|
||||
throw new Error("not used");
|
||||
}
|
||||
|
||||
protected run(): ReturnType<AbstractAgent["run"]> {
|
||||
return EMPTY;
|
||||
}
|
||||
|
||||
protected connect(
|
||||
input: RunAgentInput,
|
||||
): ReturnType<AbstractAgent["connect"]> {
|
||||
return this.runner.connect({ threadId: input.threadId });
|
||||
}
|
||||
}
|
||||
|
||||
type Deferred<T> = {
|
||||
promise: Promise<T>;
|
||||
resolve: (value: T) => void;
|
||||
reject: (reason?: unknown) => void;
|
||||
};
|
||||
|
||||
function createDeferred<T>(): Deferred<T> {
|
||||
let resolve!: (value: T) => void;
|
||||
let reject!: (reason?: unknown) => void;
|
||||
const promise = new Promise<T>((res, rej) => {
|
||||
resolve = res;
|
||||
reject = rej;
|
||||
});
|
||||
return { promise, resolve, reject };
|
||||
}
|
||||
|
||||
async function collectEvents(
|
||||
observable:
|
||||
| ReturnType<SqliteAgentRunner["run"]>
|
||||
| ReturnType<SqliteAgentRunner["connect"]>,
|
||||
) {
|
||||
return firstValueFrom(observable.pipe(toArray()));
|
||||
}
|
||||
|
||||
function createRunInput({
|
||||
threadId,
|
||||
runId,
|
||||
messages,
|
||||
state,
|
||||
parentRunId,
|
||||
}: {
|
||||
threadId: string;
|
||||
runId: string;
|
||||
messages: Message[];
|
||||
state?: Record<string, unknown>;
|
||||
parentRunId?: string | null;
|
||||
}): RunAgentInput {
|
||||
return {
|
||||
threadId,
|
||||
runId,
|
||||
parentRunId: parentRunId ?? undefined,
|
||||
state: state ?? {},
|
||||
messages,
|
||||
tools: [],
|
||||
context: [],
|
||||
forwardedProps: undefined,
|
||||
};
|
||||
}
|
||||
|
||||
function expectRunStartedEvent(event: BaseEvent, expectedMessages: Message[]) {
|
||||
expect(event.type).toBe(EventType.RUN_STARTED);
|
||||
const runStarted = event as RunStartedEvent;
|
||||
expect(runStarted.input?.messages).toEqual(expectedMessages);
|
||||
}
|
||||
|
||||
function createTextMessageEvents({
|
||||
messageId,
|
||||
role = "assistant",
|
||||
content,
|
||||
}: {
|
||||
messageId: string;
|
||||
role?: "assistant" | "developer" | "system" | "user";
|
||||
content: string;
|
||||
}): BaseEvent[] {
|
||||
return [
|
||||
{
|
||||
type: EventType.TEXT_MESSAGE_START,
|
||||
messageId,
|
||||
role,
|
||||
},
|
||||
{
|
||||
type: EventType.TEXT_MESSAGE_CONTENT,
|
||||
messageId,
|
||||
delta: content,
|
||||
},
|
||||
{
|
||||
type: EventType.TEXT_MESSAGE_END,
|
||||
messageId,
|
||||
},
|
||||
] as BaseEvent[];
|
||||
}
|
||||
|
||||
function createToolCallEvents({
|
||||
toolCallId,
|
||||
parentMessageId,
|
||||
toolName,
|
||||
argsJson,
|
||||
resultMessageId,
|
||||
resultContent,
|
||||
}: {
|
||||
toolCallId: string;
|
||||
parentMessageId: string;
|
||||
toolName: string;
|
||||
argsJson: string;
|
||||
resultMessageId: string;
|
||||
resultContent: string;
|
||||
}): BaseEvent[] {
|
||||
return [
|
||||
{
|
||||
type: EventType.TOOL_CALL_START,
|
||||
toolCallId,
|
||||
toolCallName: toolName,
|
||||
parentMessageId,
|
||||
},
|
||||
{
|
||||
type: EventType.TOOL_CALL_ARGS,
|
||||
toolCallId,
|
||||
delta: argsJson,
|
||||
},
|
||||
{
|
||||
type: EventType.TOOL_CALL_END,
|
||||
toolCallId,
|
||||
},
|
||||
{
|
||||
type: EventType.TOOL_CALL_RESULT,
|
||||
toolCallId,
|
||||
messageId: resultMessageId,
|
||||
content: resultContent,
|
||||
role: "tool",
|
||||
},
|
||||
] as BaseEvent[];
|
||||
}
|
||||
|
||||
describe("SqliteAgentRunner e2e", () => {
|
||||
describe("Fresh Replay After Single Run", () => {
|
||||
it("replays sanitized message history on connectAgent", async () => {
|
||||
const runner = createRunner();
|
||||
const threadId = "thread-fresh-replay";
|
||||
const existingMessage: Message = {
|
||||
id: "message-existing",
|
||||
role: "user",
|
||||
content: "Hello there",
|
||||
};
|
||||
|
||||
const runEvents = await collectEvents(
|
||||
runner.run({
|
||||
threadId,
|
||||
agent: new EmitAgent(),
|
||||
input: createRunInput({
|
||||
threadId,
|
||||
runId: "run-0",
|
||||
messages: [existingMessage],
|
||||
}),
|
||||
}),
|
||||
);
|
||||
|
||||
expectRunStartedEvent(runEvents[0], [existingMessage]);
|
||||
expect(runEvents.at(-1)?.type).toBe(EventType.RUN_FINISHED);
|
||||
|
||||
const replayEvents = await collectEvents(runner.connect({ threadId }));
|
||||
const replayAgent = new ReplayAgent(replayEvents, threadId);
|
||||
await replayAgent.connectAgent({ runId: "replay-run" });
|
||||
|
||||
expect(replayAgent.messages).toEqual([existingMessage]);
|
||||
});
|
||||
});
|
||||
|
||||
describe("New Messages on Subsequent Runs", () => {
|
||||
it("merges new message IDs without duplicating history", async () => {
|
||||
const runner = createRunner();
|
||||
const threadId = "thread-subsequent-runs";
|
||||
const existingMessage: Message = {
|
||||
id: "msg-existing",
|
||||
role: "user",
|
||||
content: "First turn",
|
||||
};
|
||||
|
||||
const initialRunEvents = await collectEvents(
|
||||
runner.run({
|
||||
threadId,
|
||||
agent: new EmitAgent(),
|
||||
input: createRunInput({
|
||||
threadId,
|
||||
runId: "run-0",
|
||||
messages: [existingMessage],
|
||||
}),
|
||||
}),
|
||||
);
|
||||
expectRunStartedEvent(initialRunEvents[0], [existingMessage]);
|
||||
|
||||
const newMessage: Message = {
|
||||
id: "msg-new",
|
||||
role: "user",
|
||||
content: "Second turn",
|
||||
};
|
||||
|
||||
const secondRunEvents = await collectEvents(
|
||||
runner.run({
|
||||
threadId,
|
||||
agent: new EmitAgent(),
|
||||
input: createRunInput({
|
||||
threadId,
|
||||
runId: "run-1",
|
||||
messages: [existingMessage, newMessage],
|
||||
}),
|
||||
}),
|
||||
);
|
||||
|
||||
expectRunStartedEvent(secondRunEvents[0], [newMessage]);
|
||||
|
||||
const replayEvents = await collectEvents(runner.connect({ threadId }));
|
||||
const replayAgent = new ReplayAgent(replayEvents, threadId);
|
||||
await replayAgent.connectAgent({ runId: "replay-run" });
|
||||
|
||||
expect(replayAgent.messages).toEqual([existingMessage, newMessage]);
|
||||
expect(
|
||||
new Set(replayAgent.messages.map((message) => message.id)).size,
|
||||
).toBe(replayAgent.messages.length);
|
||||
});
|
||||
});
|
||||
|
||||
describe("Fresh Agent Connection After Prior Runs", () => {
|
||||
it("hydrates a brand-new agent via connect()", async () => {
|
||||
const runner = createRunner();
|
||||
const threadId = "thread-new-agent-connection";
|
||||
const existingMessage: Message = {
|
||||
id: "existing-connection",
|
||||
role: "user",
|
||||
content: "Persist me",
|
||||
};
|
||||
|
||||
const runEvents = await collectEvents(
|
||||
runner.run({
|
||||
threadId,
|
||||
agent: new EmitAgent(),
|
||||
input: createRunInput({
|
||||
threadId,
|
||||
runId: "run-0",
|
||||
messages: [existingMessage],
|
||||
}),
|
||||
}),
|
||||
);
|
||||
expectRunStartedEvent(runEvents[0], [existingMessage]);
|
||||
|
||||
const connectingAgent = new RunnerConnectAgent(runner, threadId);
|
||||
await connectingAgent.connectAgent({ runId: "connect-run" });
|
||||
|
||||
expect(connectingAgent.messages).toEqual([existingMessage]);
|
||||
});
|
||||
});
|
||||
|
||||
describe("Mixed Roles and Tool Results", () => {
|
||||
it("preserves agent-emitted tool events alongside heterogeneous inputs", async () => {
|
||||
const runner = createRunner();
|
||||
const threadId = "thread-mixed-roles";
|
||||
|
||||
const systemMessage: Message = {
|
||||
id: "sys-1",
|
||||
role: "system",
|
||||
content: "Global directive",
|
||||
};
|
||||
const developerMessage: Message = {
|
||||
id: "dev-1",
|
||||
role: "developer",
|
||||
content: "Internal guidance",
|
||||
};
|
||||
const userMessage: Message = {
|
||||
id: "user-1",
|
||||
role: "user",
|
||||
content: "Need the weather",
|
||||
};
|
||||
const baseMessages = [systemMessage, developerMessage, userMessage];
|
||||
|
||||
const assistantMessageId = "assistant-1";
|
||||
const toolCallId = "tool-call-1";
|
||||
const toolMessageId = "tool-msg-1";
|
||||
|
||||
const agentEvents: BaseEvent[] = [
|
||||
...createTextMessageEvents({
|
||||
messageId: assistantMessageId,
|
||||
content: "Calling the weather tool",
|
||||
}),
|
||||
...createToolCallEvents({
|
||||
toolCallId,
|
||||
parentMessageId: assistantMessageId,
|
||||
toolName: "getWeather",
|
||||
argsJson: '{"location":"NYC"}',
|
||||
resultMessageId: toolMessageId,
|
||||
resultContent: '{"temp":72}',
|
||||
}),
|
||||
];
|
||||
|
||||
const runEvents = await collectEvents(
|
||||
runner.run({
|
||||
threadId,
|
||||
agent: new EmitAgent({ events: agentEvents }),
|
||||
input: createRunInput({
|
||||
threadId,
|
||||
runId: "run-0",
|
||||
messages: baseMessages,
|
||||
}),
|
||||
}),
|
||||
);
|
||||
|
||||
expectRunStartedEvent(runEvents[0], baseMessages);
|
||||
expect(
|
||||
runEvents.filter((event) => event.type === EventType.TOOL_CALL_RESULT),
|
||||
).toHaveLength(1);
|
||||
|
||||
const replayEvents = await collectEvents(runner.connect({ threadId }));
|
||||
const replayAgent = new ReplayAgent(replayEvents, threadId);
|
||||
await replayAgent.connectAgent({ runId: "replay-run" });
|
||||
|
||||
expect(replayAgent.messages).toEqual([
|
||||
systemMessage,
|
||||
developerMessage,
|
||||
userMessage,
|
||||
{
|
||||
id: assistantMessageId,
|
||||
role: "assistant",
|
||||
content: "Calling the weather tool",
|
||||
toolCalls: [
|
||||
{
|
||||
id: toolCallId,
|
||||
type: "function",
|
||||
function: {
|
||||
name: "getWeather",
|
||||
arguments: '{"location":"NYC"}',
|
||||
},
|
||||
},
|
||||
],
|
||||
},
|
||||
{
|
||||
id: toolMessageId,
|
||||
role: "tool",
|
||||
content: '{"temp":72}',
|
||||
toolCallId,
|
||||
},
|
||||
]);
|
||||
expect(
|
||||
replayAgent.messages.filter((message) => message.role === "tool"),
|
||||
).toHaveLength(1);
|
||||
});
|
||||
});
|
||||
|
||||
describe("Multiple Consecutive Runs with Agent Output", () => {
|
||||
it("deduplicates input history while emitting each agent message once", async () => {
|
||||
const runner = createRunner();
|
||||
const threadId = "thread-multi-runs";
|
||||
const systemMessage: Message = {
|
||||
id: "system-shared",
|
||||
role: "system",
|
||||
content: "System context",
|
||||
};
|
||||
const userMessages: Message[] = [];
|
||||
|
||||
for (let index = 0; index < 3; index += 1) {
|
||||
const userMessage: Message = {
|
||||
id: `user-${index + 1}`,
|
||||
role: "user",
|
||||
content: `User message ${index + 1}`,
|
||||
};
|
||||
userMessages.push(userMessage);
|
||||
|
||||
const messagesForRun = [systemMessage, ...userMessages];
|
||||
const assistantId = `assistant-${index + 1}`;
|
||||
const toolCallId = `tool-call-${index + 1}`;
|
||||
const toolMessageId = `tool-msg-${index + 1}`;
|
||||
|
||||
const events: BaseEvent[] = [
|
||||
...createTextMessageEvents({
|
||||
messageId: assistantId,
|
||||
content: `Assistant reply ${index + 1}`,
|
||||
}),
|
||||
...createToolCallEvents({
|
||||
toolCallId,
|
||||
parentMessageId: assistantId,
|
||||
toolName: `tool-${index + 1}`,
|
||||
argsJson: `{"step":${index + 1}}`,
|
||||
resultMessageId: toolMessageId,
|
||||
resultContent: `{"ok":${index + 1}}`,
|
||||
}),
|
||||
];
|
||||
|
||||
const runEvents = await collectEvents(
|
||||
runner.run({
|
||||
threadId,
|
||||
agent: new EmitAgent({ events }),
|
||||
input: createRunInput({
|
||||
threadId,
|
||||
runId: `run-${index}`,
|
||||
messages: messagesForRun,
|
||||
}),
|
||||
}),
|
||||
);
|
||||
|
||||
if (index === 0) {
|
||||
expectRunStartedEvent(runEvents[0], messagesForRun);
|
||||
} else {
|
||||
expectRunStartedEvent(runEvents[0], [userMessage]);
|
||||
}
|
||||
expect(runEvents.at(-1)?.type).toBe(EventType.RUN_FINISHED);
|
||||
}
|
||||
|
||||
const replayEvents = await collectEvents(runner.connect({ threadId }));
|
||||
const replayAgent = new ReplayAgent(replayEvents, threadId);
|
||||
await replayAgent.connectAgent({ runId: "replay-final" });
|
||||
|
||||
const finalMessages = replayAgent.messages;
|
||||
expect(new Set(finalMessages.map((message) => message.id)).size).toBe(
|
||||
finalMessages.length,
|
||||
);
|
||||
const roleCounts = finalMessages.reduce<Record<string, number>>(
|
||||
(counts, message) => {
|
||||
counts[message.role] = (counts[message.role] ?? 0) + 1;
|
||||
return counts;
|
||||
},
|
||||
{},
|
||||
);
|
||||
expect(roleCounts.system).toBe(1);
|
||||
expect(roleCounts.user).toBe(3);
|
||||
expect(roleCounts.assistant).toBe(3);
|
||||
expect(roleCounts.tool).toBe(3);
|
||||
});
|
||||
});
|
||||
|
||||
describe("Agent-Provided RUN_STARTED input", () => {
|
||||
it("forwards the agent-specified payload without sanitizing", async () => {
|
||||
const runner = createRunner();
|
||||
const threadId = "thread-custom-run-started";
|
||||
const runId = "run-0";
|
||||
|
||||
const customMessages: Message[] = [
|
||||
{
|
||||
id: "custom-user",
|
||||
role: "user",
|
||||
content: "Pre-sent content",
|
||||
},
|
||||
];
|
||||
const customInput: RunAgentInput = {
|
||||
threadId,
|
||||
runId,
|
||||
parentRunId: undefined,
|
||||
state: { injected: true },
|
||||
messages: customMessages,
|
||||
tools: [],
|
||||
context: [],
|
||||
forwardedProps: { source: "agent" },
|
||||
};
|
||||
const customRunStarted: RunStartedEvent = {
|
||||
type: EventType.RUN_STARTED,
|
||||
threadId,
|
||||
runId,
|
||||
parentRunId: null,
|
||||
input: customInput,
|
||||
};
|
||||
|
||||
const agentEvents: BaseEvent[] = [
|
||||
customRunStarted,
|
||||
...createTextMessageEvents({
|
||||
messageId: "agent-message",
|
||||
content: "Custom start acknowledged",
|
||||
}),
|
||||
];
|
||||
|
||||
const runEvents = await collectEvents(
|
||||
runner.run({
|
||||
threadId,
|
||||
agent: new EmitAgent({
|
||||
events: agentEvents,
|
||||
emitDefaultRunStarted: false,
|
||||
}),
|
||||
input: createRunInput({
|
||||
threadId,
|
||||
runId,
|
||||
messages: [],
|
||||
}),
|
||||
}),
|
||||
);
|
||||
|
||||
expect(runEvents[0]).toEqual(customRunStarted);
|
||||
expect(
|
||||
runEvents.filter((event) => event.type === EventType.RUN_FINISHED),
|
||||
).toHaveLength(1);
|
||||
|
||||
const replayEvents = await collectEvents(runner.connect({ threadId }));
|
||||
const replayAgent = new ReplayAgent(replayEvents, threadId);
|
||||
await replayAgent.connectAgent({ runId: "replay-run" });
|
||||
expect(
|
||||
replayAgent.messages.find((message) => message.id === "custom-user"),
|
||||
).toEqual(customMessages[0]);
|
||||
});
|
||||
});
|
||||
|
||||
describe("Concurrent Connections During Run", () => {
|
||||
it("streams in-flight events to live subscribers and persists final history", async () => {
|
||||
const runner = createRunner();
|
||||
const threadId = "thread-concurrency";
|
||||
const runId = "run-live";
|
||||
const initialMessage: Message = {
|
||||
id: "initial-user",
|
||||
role: "user",
|
||||
content: "Start run",
|
||||
};
|
||||
|
||||
const runStartedSignal = createDeferred<void>();
|
||||
const resumeSignal = createDeferred<void>();
|
||||
|
||||
const agent = new EmitAgent({
|
||||
events: [
|
||||
...createTextMessageEvents({
|
||||
messageId: "assistant-live",
|
||||
content: "Streaming content",
|
||||
}),
|
||||
],
|
||||
afterEvent: async ({ event }) => {
|
||||
if (event.type === EventType.RUN_STARTED) {
|
||||
runStartedSignal.resolve();
|
||||
await resumeSignal.promise;
|
||||
}
|
||||
},
|
||||
});
|
||||
|
||||
const runEvents: BaseEvent[] = [];
|
||||
const run$ = runner.run({
|
||||
threadId,
|
||||
agent,
|
||||
input: createRunInput({
|
||||
threadId,
|
||||
runId,
|
||||
messages: [initialMessage],
|
||||
}),
|
||||
});
|
||||
|
||||
let runSubscription: Subscription;
|
||||
const runCompletion = new Promise<void>((resolve, reject) => {
|
||||
runSubscription = run$.subscribe({
|
||||
next: (event) => runEvents.push(event),
|
||||
error: (error) => {
|
||||
runSubscription.unsubscribe();
|
||||
reject(error);
|
||||
},
|
||||
complete: () => {
|
||||
runSubscription.unsubscribe();
|
||||
resolve();
|
||||
},
|
||||
});
|
||||
});
|
||||
|
||||
await runStartedSignal.promise;
|
||||
|
||||
const liveEvents: BaseEvent[] = [];
|
||||
const connect$ = runner.connect({ threadId });
|
||||
let connectSubscription: Subscription;
|
||||
const connectCompletion = new Promise<void>((resolve, reject) => {
|
||||
connectSubscription = connect$.subscribe({
|
||||
next: (event) => liveEvents.push(event),
|
||||
error: (error) => {
|
||||
connectSubscription.unsubscribe();
|
||||
reject(error);
|
||||
},
|
||||
complete: () => {
|
||||
connectSubscription.unsubscribe();
|
||||
resolve();
|
||||
},
|
||||
});
|
||||
});
|
||||
|
||||
resumeSignal.resolve();
|
||||
|
||||
await Promise.all([runCompletion, connectCompletion]);
|
||||
|
||||
expectRunStartedEvent(runEvents[0], [initialMessage]);
|
||||
expect(runEvents.at(-1)?.type).toBe(EventType.RUN_FINISHED);
|
||||
expect(liveEvents).toEqual(runEvents);
|
||||
|
||||
const persistedEvents = await collectEvents(runner.connect({ threadId }));
|
||||
expect(persistedEvents).toEqual(runEvents);
|
||||
|
||||
const replayAgent = new ReplayAgent(persistedEvents, threadId);
|
||||
await replayAgent.connectAgent({ runId: "replay-run" });
|
||||
expect(replayAgent.messages.map((message) => message.id)).toEqual([
|
||||
initialMessage.id,
|
||||
"assistant-live",
|
||||
]);
|
||||
});
|
||||
});
|
||||
|
||||
describe("Error Handling", () => {
|
||||
it("propagates RUN_ERROR while retaining input history", async () => {
|
||||
const runner = createRunner();
|
||||
const threadId = "thread-run-error";
|
||||
const userMessage: Message = {
|
||||
id: "error-user",
|
||||
role: "user",
|
||||
content: "Trigger error",
|
||||
};
|
||||
|
||||
const runErrorEvent: RunErrorEvent = {
|
||||
type: EventType.RUN_ERROR,
|
||||
message: "Agent failure",
|
||||
};
|
||||
|
||||
const runEvents = await collectEvents(
|
||||
runner.run({
|
||||
threadId,
|
||||
agent: new EmitAgent({
|
||||
events: [runErrorEvent],
|
||||
includeRunFinished: false,
|
||||
}),
|
||||
input: createRunInput({
|
||||
threadId,
|
||||
runId: "run-error",
|
||||
messages: [userMessage],
|
||||
}),
|
||||
}),
|
||||
);
|
||||
|
||||
expectRunStartedEvent(runEvents[0], [userMessage]);
|
||||
expect(runEvents.at(-1)).toEqual(runErrorEvent);
|
||||
|
||||
const replayEvents = await collectEvents(runner.connect({ threadId }));
|
||||
const replayAgent = new ReplayAgent(replayEvents, threadId);
|
||||
const capturedRunErrors: RunErrorEvent[] = [];
|
||||
const result = await replayAgent.connectAgent(
|
||||
{ runId: "replay-run" },
|
||||
{
|
||||
onRunErrorEvent: ({ event }) => {
|
||||
capturedRunErrors.push(event);
|
||||
},
|
||||
},
|
||||
);
|
||||
|
||||
expect(runEvents.at(-1)?.type).toBe(EventType.RUN_ERROR);
|
||||
expect(capturedRunErrors).toHaveLength(1);
|
||||
expect(capturedRunErrors[0]).toMatchObject(runErrorEvent);
|
||||
expect(result.newMessages).toEqual([userMessage]);
|
||||
expect(replayAgent.messages).toEqual([userMessage]);
|
||||
});
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,450 @@
|
||||
import { describe, it, expect, beforeEach, afterEach } from "vitest";
|
||||
import { SqliteAgentRunner } from "..";
|
||||
import {
|
||||
AbstractAgent,
|
||||
BaseEvent,
|
||||
EventType,
|
||||
Message,
|
||||
RunAgentInput,
|
||||
RunFinishedEvent,
|
||||
RunStartedEvent,
|
||||
TextMessageContentEvent,
|
||||
TextMessageEndEvent,
|
||||
TextMessageStartEvent,
|
||||
} from "@ag-ui/client";
|
||||
import { EMPTY, firstValueFrom } from "rxjs";
|
||||
import { toArray } from "rxjs/operators";
|
||||
import Database from "better-sqlite3";
|
||||
import fs from "fs";
|
||||
import os from "os";
|
||||
import path from "path";
|
||||
|
||||
type RunCallbacks = {
|
||||
onEvent: (event: { event: BaseEvent }) => void | Promise<void>;
|
||||
onNewMessage?: (args: { message: Message }) => void | Promise<void>;
|
||||
onRunStartedEvent?: () => void | Promise<void>;
|
||||
};
|
||||
|
||||
class MockAgent extends AbstractAgent {
|
||||
constructor(
|
||||
private readonly events: BaseEvent[] = [],
|
||||
private readonly emitDefaultRunStarted = true,
|
||||
) {
|
||||
super();
|
||||
}
|
||||
|
||||
async runAgent(input: RunAgentInput, callbacks: RunCallbacks): Promise<void> {
|
||||
if (this.emitDefaultRunStarted) {
|
||||
const runStarted: RunStartedEvent = {
|
||||
type: EventType.RUN_STARTED,
|
||||
threadId: input.threadId,
|
||||
runId: input.runId,
|
||||
};
|
||||
await callbacks.onEvent({ event: runStarted });
|
||||
await callbacks.onRunStartedEvent?.();
|
||||
}
|
||||
|
||||
for (const event of this.events) {
|
||||
await callbacks.onEvent({ event });
|
||||
}
|
||||
|
||||
const hasTerminalEvent = this.events.some(
|
||||
(event) =>
|
||||
event.type === EventType.RUN_FINISHED ||
|
||||
event.type === EventType.RUN_ERROR,
|
||||
);
|
||||
|
||||
if (!hasTerminalEvent) {
|
||||
const runFinished: RunFinishedEvent = {
|
||||
type: EventType.RUN_FINISHED,
|
||||
threadId: input.threadId,
|
||||
runId: input.runId,
|
||||
};
|
||||
await callbacks.onEvent({ event: runFinished });
|
||||
}
|
||||
}
|
||||
|
||||
protected run(): ReturnType<AbstractAgent["run"]> {
|
||||
return EMPTY;
|
||||
}
|
||||
|
||||
protected connect(): ReturnType<AbstractAgent["connect"]> {
|
||||
return EMPTY;
|
||||
}
|
||||
|
||||
clone(): AbstractAgent {
|
||||
return new MockAgent(this.events, this.emitDefaultRunStarted);
|
||||
}
|
||||
}
|
||||
|
||||
class StoppableAgent extends AbstractAgent {
|
||||
private shouldStop = false;
|
||||
private eventDelay: number;
|
||||
|
||||
constructor(eventDelay = 5) {
|
||||
super();
|
||||
this.eventDelay = eventDelay;
|
||||
}
|
||||
|
||||
async runAgent(input: RunAgentInput, callbacks: RunCallbacks): Promise<void> {
|
||||
this.shouldStop = false;
|
||||
let counter = 0;
|
||||
|
||||
const runStarted: RunStartedEvent = {
|
||||
type: EventType.RUN_STARTED,
|
||||
threadId: input.threadId,
|
||||
runId: input.runId,
|
||||
};
|
||||
await callbacks.onEvent({ event: runStarted });
|
||||
await callbacks.onRunStartedEvent?.();
|
||||
|
||||
while (!this.shouldStop && counter < 10_000) {
|
||||
await new Promise((resolve) => setTimeout(resolve, this.eventDelay));
|
||||
const event: BaseEvent = {
|
||||
type: EventType.TEXT_MESSAGE_CONTENT,
|
||||
messageId: `sqlite-stop-${counter}`,
|
||||
delta: `chunk-${counter}`,
|
||||
} as TextMessageContentEvent;
|
||||
await callbacks.onEvent({ event });
|
||||
counter += 1;
|
||||
}
|
||||
}
|
||||
|
||||
abortRun(): void {
|
||||
this.shouldStop = true;
|
||||
}
|
||||
|
||||
clone(): AbstractAgent {
|
||||
return new StoppableAgent(this.eventDelay);
|
||||
}
|
||||
}
|
||||
|
||||
class OpenEventsAgent extends AbstractAgent {
|
||||
private shouldStop = false;
|
||||
|
||||
async runAgent(input: RunAgentInput, callbacks: RunCallbacks): Promise<void> {
|
||||
this.shouldStop = false;
|
||||
const messageId = "open-message";
|
||||
const toolCallId = "open-tool";
|
||||
|
||||
await callbacks.onEvent({
|
||||
event: {
|
||||
type: EventType.TEXT_MESSAGE_START,
|
||||
messageId,
|
||||
role: "assistant",
|
||||
} as BaseEvent,
|
||||
});
|
||||
|
||||
await callbacks.onEvent({
|
||||
event: {
|
||||
type: EventType.TEXT_MESSAGE_CONTENT,
|
||||
messageId,
|
||||
delta: "Partial content",
|
||||
} as BaseEvent,
|
||||
});
|
||||
|
||||
await callbacks.onEvent({
|
||||
event: {
|
||||
type: EventType.TOOL_CALL_START,
|
||||
toolCallId,
|
||||
toolCallName: "testTool",
|
||||
parentMessageId: messageId,
|
||||
} as BaseEvent,
|
||||
});
|
||||
|
||||
while (!this.shouldStop) {
|
||||
await new Promise((resolve) => setTimeout(resolve, 5));
|
||||
}
|
||||
}
|
||||
|
||||
abortRun(): void {
|
||||
this.shouldStop = true;
|
||||
}
|
||||
|
||||
clone(): AbstractAgent {
|
||||
return new OpenEventsAgent();
|
||||
}
|
||||
}
|
||||
|
||||
describe("SqliteAgentRunner", () => {
|
||||
let tempDir: string;
|
||||
let dbPath: string;
|
||||
let runner: SqliteAgentRunner;
|
||||
|
||||
beforeEach(() => {
|
||||
tempDir = fs.mkdtempSync(path.join(os.tmpdir(), "sqlite-runner-test-"));
|
||||
dbPath = path.join(tempDir, "test.db");
|
||||
runner = new SqliteAgentRunner({ dbPath });
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
runner.close();
|
||||
|
||||
if (fs.existsSync(dbPath)) fs.unlinkSync(dbPath);
|
||||
if (fs.existsSync(tempDir)) fs.rmdirSync(tempDir);
|
||||
});
|
||||
|
||||
it("emits RUN_STARTED and agent events", async () => {
|
||||
const threadId = "sqlite-basic";
|
||||
const agent = new MockAgent([
|
||||
{
|
||||
type: EventType.TEXT_MESSAGE_START,
|
||||
messageId: "msg-1",
|
||||
role: "assistant",
|
||||
} as TextMessageStartEvent,
|
||||
{
|
||||
type: EventType.TEXT_MESSAGE_CONTENT,
|
||||
messageId: "msg-1",
|
||||
delta: "Hello",
|
||||
} as TextMessageContentEvent,
|
||||
{
|
||||
type: EventType.TEXT_MESSAGE_END,
|
||||
messageId: "msg-1",
|
||||
} as TextMessageEndEvent,
|
||||
{
|
||||
type: EventType.RUN_FINISHED,
|
||||
threadId,
|
||||
runId: "run-1",
|
||||
} as RunFinishedEvent,
|
||||
]);
|
||||
|
||||
const events = await firstValueFrom(
|
||||
runner
|
||||
.run({
|
||||
threadId,
|
||||
agent,
|
||||
input: { threadId, runId: "run-1", messages: [], state: {} },
|
||||
})
|
||||
.pipe(toArray()),
|
||||
);
|
||||
|
||||
expect(events.map((event) => event.type)).toEqual([
|
||||
EventType.RUN_STARTED,
|
||||
EventType.TEXT_MESSAGE_START,
|
||||
EventType.TEXT_MESSAGE_CONTENT,
|
||||
EventType.TEXT_MESSAGE_END,
|
||||
EventType.RUN_FINISHED,
|
||||
]);
|
||||
});
|
||||
|
||||
it("attaches only new messages on subsequent runs", async () => {
|
||||
const threadId = "sqlite-new-messages";
|
||||
const existing: Message = { id: "existing", role: "user", content: "hi" };
|
||||
|
||||
await firstValueFrom(
|
||||
runner
|
||||
.run({
|
||||
threadId,
|
||||
agent: new MockAgent(),
|
||||
input: { threadId, runId: "run-0", messages: [existing], state: {} },
|
||||
})
|
||||
.pipe(toArray()),
|
||||
);
|
||||
|
||||
const newMessage: Message = {
|
||||
id: "new",
|
||||
role: "user",
|
||||
content: "follow up",
|
||||
};
|
||||
|
||||
const secondRun = await firstValueFrom(
|
||||
runner
|
||||
.run({
|
||||
threadId,
|
||||
agent: new MockAgent(),
|
||||
input: {
|
||||
threadId,
|
||||
runId: "run-1",
|
||||
messages: [existing, newMessage],
|
||||
state: { counter: 1 },
|
||||
},
|
||||
})
|
||||
.pipe(toArray()),
|
||||
);
|
||||
|
||||
const runStarted = secondRun[0] as RunStartedEvent;
|
||||
expect(runStarted.input?.messages?.map((m) => m.id)).toEqual(["new"]);
|
||||
|
||||
const db = new Database(dbPath);
|
||||
const rows = db
|
||||
.prepare(
|
||||
"SELECT events FROM agent_runs WHERE thread_id = ? ORDER BY created_at",
|
||||
)
|
||||
.all(threadId) as {
|
||||
events: string;
|
||||
}[];
|
||||
db.close();
|
||||
|
||||
expect(rows).toHaveLength(2);
|
||||
const run1Stored = JSON.parse(rows[0].events) as BaseEvent[];
|
||||
const run2Stored = JSON.parse(rows[1].events) as BaseEvent[];
|
||||
|
||||
const run1Started = run1Stored.find(
|
||||
(event) => event.type === EventType.RUN_STARTED,
|
||||
) as RunStartedEvent;
|
||||
expect(run1Started.input?.messages?.map((m) => m.id)).toEqual(["existing"]);
|
||||
|
||||
const run2Started = run2Stored.find(
|
||||
(event) => event.type === EventType.RUN_STARTED,
|
||||
) as RunStartedEvent;
|
||||
expect(run2Started.input?.messages?.map((m) => m.id)).toEqual(["new"]);
|
||||
});
|
||||
|
||||
it("preserves agent-provided input", async () => {
|
||||
const threadId = "sqlite-agent-input";
|
||||
const providedInput: RunAgentInput = {
|
||||
threadId,
|
||||
runId: "run-keep",
|
||||
messages: [],
|
||||
state: { fromAgent: true },
|
||||
};
|
||||
|
||||
const agent = new MockAgent(
|
||||
[
|
||||
{
|
||||
type: EventType.RUN_STARTED,
|
||||
threadId,
|
||||
runId: "run-keep",
|
||||
input: providedInput,
|
||||
} as RunStartedEvent,
|
||||
{
|
||||
type: EventType.RUN_FINISHED,
|
||||
threadId,
|
||||
runId: "run-keep",
|
||||
} as RunFinishedEvent,
|
||||
],
|
||||
false,
|
||||
);
|
||||
|
||||
const events = await firstValueFrom(
|
||||
runner
|
||||
.run({
|
||||
threadId,
|
||||
agent,
|
||||
input: {
|
||||
threadId,
|
||||
runId: "run-keep",
|
||||
messages: [{ id: "ignored", role: "user", content: "hi" }],
|
||||
state: {},
|
||||
},
|
||||
})
|
||||
.pipe(toArray()),
|
||||
);
|
||||
|
||||
expect(events.map((event) => event.type)).toEqual([
|
||||
EventType.RUN_STARTED,
|
||||
EventType.RUN_FINISHED,
|
||||
]);
|
||||
const runStarted = events[0] as RunStartedEvent;
|
||||
expect(runStarted.input).toBe(providedInput);
|
||||
});
|
||||
|
||||
it("persists events across runner instances", async () => {
|
||||
const threadId = "sqlite-persist";
|
||||
const agent = new MockAgent([
|
||||
{
|
||||
type: EventType.TEXT_MESSAGE_START,
|
||||
messageId: "msg",
|
||||
role: "assistant",
|
||||
} as TextMessageStartEvent,
|
||||
{
|
||||
type: EventType.TEXT_MESSAGE_CONTENT,
|
||||
messageId: "msg",
|
||||
delta: "hi",
|
||||
} as TextMessageContentEvent,
|
||||
{
|
||||
type: EventType.TEXT_MESSAGE_END,
|
||||
messageId: "msg",
|
||||
} as TextMessageEndEvent,
|
||||
{
|
||||
type: EventType.RUN_FINISHED,
|
||||
threadId,
|
||||
runId: "run-1",
|
||||
} as RunFinishedEvent,
|
||||
]);
|
||||
|
||||
await firstValueFrom(
|
||||
runner
|
||||
.run({
|
||||
threadId,
|
||||
agent,
|
||||
input: { threadId, runId: "run-1", messages: [], state: {} },
|
||||
})
|
||||
.pipe(toArray()),
|
||||
);
|
||||
|
||||
const newRunner = new SqliteAgentRunner({ dbPath });
|
||||
try {
|
||||
const replayed = await firstValueFrom(
|
||||
newRunner.connect({ threadId }).pipe(toArray()),
|
||||
);
|
||||
|
||||
expect(replayed[0].type).toBe(EventType.RUN_STARTED);
|
||||
expect(replayed.slice(1).map((event) => event.type)).toEqual([
|
||||
EventType.TEXT_MESSAGE_START,
|
||||
EventType.TEXT_MESSAGE_CONTENT,
|
||||
EventType.TEXT_MESSAGE_END,
|
||||
EventType.RUN_FINISHED,
|
||||
]);
|
||||
} finally {
|
||||
newRunner.close();
|
||||
}
|
||||
});
|
||||
|
||||
it("returns false when stopping a thread that is not running", async () => {
|
||||
await expect(runner.stop({ threadId: "sqlite-missing" })).resolves.toBe(
|
||||
false,
|
||||
);
|
||||
});
|
||||
|
||||
it("stops an active run and completes observables", async () => {
|
||||
const threadId = "sqlite-stop";
|
||||
const agent = new StoppableAgent(2);
|
||||
const input: RunAgentInput = {
|
||||
threadId,
|
||||
runId: "sqlite-stop-run",
|
||||
messages: [],
|
||||
state: {},
|
||||
};
|
||||
|
||||
const run$ = runner.run({ threadId, agent, input });
|
||||
const collected = firstValueFrom(run$.pipe(toArray()));
|
||||
|
||||
await new Promise((resolve) => setTimeout(resolve, 20));
|
||||
expect(await runner.isRunning({ threadId })).toBe(true);
|
||||
|
||||
const stopped = await runner.stop({ threadId });
|
||||
expect(stopped).toBe(true);
|
||||
|
||||
const events = await collected;
|
||||
expect(events.length).toBeGreaterThan(0);
|
||||
expect(events[events.length - 1].type).toBe(EventType.RUN_FINISHED);
|
||||
expect(await runner.isRunning({ threadId })).toBe(false);
|
||||
});
|
||||
|
||||
it("closes open text and tool events when stopping", async () => {
|
||||
const threadId = "sqlite-open-events";
|
||||
const agent = new OpenEventsAgent();
|
||||
const input: RunAgentInput = {
|
||||
threadId,
|
||||
runId: "sqlite-open-run",
|
||||
messages: [],
|
||||
state: {},
|
||||
};
|
||||
|
||||
const run$ = runner.run({ threadId, agent, input });
|
||||
const collected = firstValueFrom(run$.pipe(toArray()));
|
||||
|
||||
await new Promise((resolve) => setTimeout(resolve, 20));
|
||||
await runner.stop({ threadId });
|
||||
|
||||
const events = await collected;
|
||||
const endingTypes = events.slice(-4).map((event) => event.type);
|
||||
expect(endingTypes).toEqual([
|
||||
EventType.TEXT_MESSAGE_END,
|
||||
EventType.TOOL_CALL_END,
|
||||
EventType.TOOL_CALL_RESULT,
|
||||
EventType.RUN_FINISHED,
|
||||
]);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1 @@
|
||||
export * from "./sqlite-runner.js";
|
||||
@@ -0,0 +1,517 @@
|
||||
import {
|
||||
AgentRunner,
|
||||
finalizeRunEvents,
|
||||
type AgentRunnerConnectRequest,
|
||||
type AgentRunnerIsRunningRequest,
|
||||
type AgentRunnerRunRequest,
|
||||
type AgentRunnerStopRequest,
|
||||
} from "@copilotkit/runtime/v2";
|
||||
import { Observable, ReplaySubject } from "rxjs";
|
||||
import {
|
||||
AbstractAgent,
|
||||
BaseEvent,
|
||||
RunAgentInput,
|
||||
EventType,
|
||||
RunStartedEvent,
|
||||
compactEvents,
|
||||
} from "@ag-ui/client";
|
||||
import Database from "better-sqlite3";
|
||||
|
||||
const SCHEMA_VERSION = 1;
|
||||
|
||||
interface AgentRunRecord {
|
||||
id: number;
|
||||
thread_id: string;
|
||||
run_id: string;
|
||||
parent_run_id: string | null;
|
||||
events: BaseEvent[];
|
||||
input: RunAgentInput;
|
||||
created_at: number;
|
||||
version: number;
|
||||
}
|
||||
|
||||
export interface SqliteAgentRunnerOptions {
|
||||
dbPath?: string;
|
||||
}
|
||||
|
||||
interface ActiveConnectionContext {
|
||||
subject: ReplaySubject<BaseEvent>;
|
||||
agent?: AbstractAgent;
|
||||
runSubject?: ReplaySubject<BaseEvent>;
|
||||
currentEvents?: BaseEvent[];
|
||||
stopRequested?: boolean;
|
||||
}
|
||||
|
||||
// Active connections for streaming events and stop support
|
||||
const ACTIVE_CONNECTIONS = new Map<string, ActiveConnectionContext>();
|
||||
|
||||
export class SqliteAgentRunner extends AgentRunner {
|
||||
private db: any;
|
||||
|
||||
constructor(options: SqliteAgentRunnerOptions = {}) {
|
||||
super();
|
||||
const dbPath = options.dbPath ?? ":memory:";
|
||||
|
||||
if (!Database) {
|
||||
throw new Error(
|
||||
"better-sqlite3 is required for SqliteAgentRunner but was not found.\n" +
|
||||
"Please install it in your project:\n" +
|
||||
" npm install better-sqlite3\n" +
|
||||
" or\n" +
|
||||
" pnpm add better-sqlite3\n" +
|
||||
" or\n" +
|
||||
" yarn add better-sqlite3\n\n" +
|
||||
"If you don't need persistence, use InMemoryAgentRunner instead.",
|
||||
);
|
||||
}
|
||||
|
||||
this.db = new Database(dbPath);
|
||||
this.initializeSchema();
|
||||
}
|
||||
|
||||
private initializeSchema(): void {
|
||||
// Create the agent_runs table
|
||||
this.db.exec(`
|
||||
CREATE TABLE IF NOT EXISTS agent_runs (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
thread_id TEXT NOT NULL,
|
||||
run_id TEXT NOT NULL UNIQUE,
|
||||
parent_run_id TEXT,
|
||||
events TEXT NOT NULL,
|
||||
input TEXT NOT NULL,
|
||||
created_at INTEGER NOT NULL,
|
||||
version INTEGER NOT NULL
|
||||
)
|
||||
`);
|
||||
|
||||
// Create run_state table to track active runs
|
||||
this.db.exec(`
|
||||
CREATE TABLE IF NOT EXISTS run_state (
|
||||
thread_id TEXT PRIMARY KEY,
|
||||
is_running INTEGER DEFAULT 0,
|
||||
current_run_id TEXT,
|
||||
updated_at INTEGER NOT NULL
|
||||
)
|
||||
`);
|
||||
|
||||
// Create indexes for efficient queries
|
||||
this.db.exec(`
|
||||
CREATE INDEX IF NOT EXISTS idx_thread_id ON agent_runs(thread_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_parent_run_id ON agent_runs(parent_run_id);
|
||||
`);
|
||||
|
||||
// Create schema version table
|
||||
this.db.exec(`
|
||||
CREATE TABLE IF NOT EXISTS schema_version (
|
||||
version INTEGER PRIMARY KEY,
|
||||
applied_at INTEGER NOT NULL
|
||||
)
|
||||
`);
|
||||
|
||||
// Check and set schema version
|
||||
const currentVersion = this.db
|
||||
.prepare(
|
||||
"SELECT version FROM schema_version ORDER BY version DESC LIMIT 1",
|
||||
)
|
||||
.get() as { version: number } | undefined;
|
||||
|
||||
if (!currentVersion || currentVersion.version < SCHEMA_VERSION) {
|
||||
this.db
|
||||
.prepare(
|
||||
"INSERT OR REPLACE INTO schema_version (version, applied_at) VALUES (?, ?)",
|
||||
)
|
||||
.run(SCHEMA_VERSION, Date.now());
|
||||
}
|
||||
}
|
||||
|
||||
private storeRun(
|
||||
threadId: string,
|
||||
runId: string,
|
||||
events: BaseEvent[],
|
||||
input: RunAgentInput,
|
||||
parentRunId?: string | null,
|
||||
): void {
|
||||
// Compact ONLY the events from this run
|
||||
const compactedEvents = compactEvents(events);
|
||||
|
||||
const stmt = this.db.prepare(`
|
||||
INSERT INTO agent_runs (thread_id, run_id, parent_run_id, events, input, created_at, version)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?)
|
||||
`);
|
||||
|
||||
stmt.run(
|
||||
threadId,
|
||||
runId,
|
||||
parentRunId ?? null,
|
||||
JSON.stringify(compactedEvents), // Store only this run's compacted events
|
||||
JSON.stringify(input),
|
||||
Date.now(),
|
||||
SCHEMA_VERSION,
|
||||
);
|
||||
}
|
||||
|
||||
private getHistoricRuns(threadId: string): AgentRunRecord[] {
|
||||
const stmt = this.db.prepare(`
|
||||
WITH RECURSIVE run_chain AS (
|
||||
-- Base case: find the root runs (those without parent)
|
||||
SELECT * FROM agent_runs
|
||||
WHERE thread_id = ? AND parent_run_id IS NULL
|
||||
|
||||
UNION ALL
|
||||
|
||||
-- Recursive case: find children of current level
|
||||
SELECT ar.* FROM agent_runs ar
|
||||
INNER JOIN run_chain rc ON ar.parent_run_id = rc.run_id
|
||||
WHERE ar.thread_id = ?
|
||||
)
|
||||
SELECT * FROM run_chain
|
||||
ORDER BY created_at ASC
|
||||
`);
|
||||
|
||||
const rows = stmt.all(threadId, threadId) as any[];
|
||||
|
||||
return rows.map((row) => ({
|
||||
id: row.id,
|
||||
thread_id: row.thread_id,
|
||||
run_id: row.run_id,
|
||||
parent_run_id: row.parent_run_id,
|
||||
events: JSON.parse(row.events),
|
||||
input: JSON.parse(row.input),
|
||||
created_at: row.created_at,
|
||||
version: row.version,
|
||||
}));
|
||||
}
|
||||
|
||||
private getLatestRunId(threadId: string): string | null {
|
||||
const stmt = this.db.prepare(`
|
||||
SELECT run_id FROM agent_runs
|
||||
WHERE thread_id = ?
|
||||
ORDER BY created_at DESC
|
||||
LIMIT 1
|
||||
`);
|
||||
|
||||
const result = stmt.get(threadId) as { run_id: string } | undefined;
|
||||
return result?.run_id ?? null;
|
||||
}
|
||||
|
||||
private setRunState(
|
||||
threadId: string,
|
||||
isRunning: boolean,
|
||||
runId?: string,
|
||||
): void {
|
||||
const stmt = this.db.prepare(`
|
||||
INSERT OR REPLACE INTO run_state (thread_id, is_running, current_run_id, updated_at)
|
||||
VALUES (?, ?, ?, ?)
|
||||
`);
|
||||
stmt.run(threadId, isRunning ? 1 : 0, runId ?? null, Date.now());
|
||||
}
|
||||
|
||||
private getRunState(threadId: string): {
|
||||
isRunning: boolean;
|
||||
currentRunId: string | null;
|
||||
} {
|
||||
const stmt = this.db.prepare(`
|
||||
SELECT is_running, current_run_id FROM run_state WHERE thread_id = ?
|
||||
`);
|
||||
const result = stmt.get(threadId) as
|
||||
| { is_running: number; current_run_id: string | null }
|
||||
| undefined;
|
||||
|
||||
return {
|
||||
isRunning: result?.is_running === 1,
|
||||
currentRunId: result?.current_run_id ?? null,
|
||||
};
|
||||
}
|
||||
|
||||
run(request: AgentRunnerRunRequest): Observable<BaseEvent> {
|
||||
// Check if thread is already running in database
|
||||
const runState = this.getRunState(request.threadId);
|
||||
if (runState.isRunning) {
|
||||
throw new Error("Thread already running");
|
||||
}
|
||||
|
||||
// Mark thread as running in database
|
||||
this.setRunState(request.threadId, true, request.input.runId);
|
||||
|
||||
// Track seen message IDs and current run events in memory for this run
|
||||
const seenMessageIds = new Set<string>();
|
||||
const currentRunEvents: BaseEvent[] = [];
|
||||
|
||||
// Get all previously seen message IDs from historic runs
|
||||
const historicRuns = this.getHistoricRuns(request.threadId);
|
||||
const historicMessageIds = new Set<string>();
|
||||
for (const run of historicRuns) {
|
||||
for (const event of run.events) {
|
||||
if ("messageId" in event && typeof event.messageId === "string") {
|
||||
historicMessageIds.add(event.messageId);
|
||||
}
|
||||
if (event.type === EventType.RUN_STARTED) {
|
||||
const runStarted = event as RunStartedEvent;
|
||||
const messages = runStarted.input?.messages ?? [];
|
||||
for (const message of messages) {
|
||||
historicMessageIds.add(message.id);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Get or create subject for this thread's connections
|
||||
const nextSubject = new ReplaySubject<BaseEvent>(Infinity);
|
||||
const prevConnection = ACTIVE_CONNECTIONS.get(request.threadId);
|
||||
const prevSubject = prevConnection?.subject;
|
||||
|
||||
// Create a subject for run() return value
|
||||
const runSubject = new ReplaySubject<BaseEvent>(Infinity);
|
||||
|
||||
// Update the active connection for this thread
|
||||
ACTIVE_CONNECTIONS.set(request.threadId, {
|
||||
subject: nextSubject,
|
||||
agent: request.agent,
|
||||
runSubject,
|
||||
currentEvents: currentRunEvents,
|
||||
stopRequested: false,
|
||||
});
|
||||
|
||||
// Helper function to run the agent and handle errors
|
||||
const runAgent = async () => {
|
||||
// Get parent run ID for chaining
|
||||
const parentRunId = this.getLatestRunId(request.threadId);
|
||||
|
||||
try {
|
||||
await request.agent.runAgent(request.input, {
|
||||
onEvent: ({ event }) => {
|
||||
let processedEvent: BaseEvent = event;
|
||||
if (event.type === EventType.RUN_STARTED) {
|
||||
const runStartedEvent = event as RunStartedEvent;
|
||||
if (!runStartedEvent.input) {
|
||||
const sanitizedMessages = request.input.messages
|
||||
? request.input.messages.filter(
|
||||
(message) => !historicMessageIds.has(message.id),
|
||||
)
|
||||
: undefined;
|
||||
const updatedInput = {
|
||||
...request.input,
|
||||
...(sanitizedMessages !== undefined
|
||||
? { messages: sanitizedMessages }
|
||||
: {}),
|
||||
};
|
||||
processedEvent = {
|
||||
...runStartedEvent,
|
||||
input: updatedInput,
|
||||
} as RunStartedEvent;
|
||||
}
|
||||
}
|
||||
|
||||
runSubject.next(processedEvent); // For run() return - only agent events
|
||||
nextSubject.next(processedEvent); // For connect() / store - all events
|
||||
currentRunEvents.push(processedEvent); // Accumulate for database storage
|
||||
},
|
||||
onNewMessage: ({ message }) => {
|
||||
// Called for each new message
|
||||
if (!seenMessageIds.has(message.id)) {
|
||||
seenMessageIds.add(message.id);
|
||||
}
|
||||
},
|
||||
onRunStartedEvent: () => {
|
||||
// Mark input messages as seen without emitting duplicates
|
||||
if (request.input.messages) {
|
||||
for (const message of request.input.messages) {
|
||||
if (!seenMessageIds.has(message.id)) {
|
||||
seenMessageIds.add(message.id);
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
});
|
||||
|
||||
const connection = ACTIVE_CONNECTIONS.get(request.threadId);
|
||||
const appendedEvents = finalizeRunEvents(currentRunEvents, {
|
||||
stopRequested: connection?.stopRequested ?? false,
|
||||
});
|
||||
for (const event of appendedEvents) {
|
||||
runSubject.next(event);
|
||||
nextSubject.next(event);
|
||||
}
|
||||
|
||||
// Store the run in database
|
||||
this.storeRun(
|
||||
request.threadId,
|
||||
request.input.runId,
|
||||
currentRunEvents,
|
||||
request.input,
|
||||
parentRunId,
|
||||
);
|
||||
|
||||
// Mark run as complete in database
|
||||
this.setRunState(request.threadId, false);
|
||||
|
||||
if (connection) {
|
||||
connection.agent = undefined;
|
||||
connection.runSubject = undefined;
|
||||
connection.currentEvents = undefined;
|
||||
connection.stopRequested = false;
|
||||
}
|
||||
|
||||
// Complete the subjects
|
||||
runSubject.complete();
|
||||
nextSubject.complete();
|
||||
|
||||
ACTIVE_CONNECTIONS.delete(request.threadId);
|
||||
} catch {
|
||||
const connection = ACTIVE_CONNECTIONS.get(request.threadId);
|
||||
const appendedEvents = finalizeRunEvents(currentRunEvents, {
|
||||
stopRequested: connection?.stopRequested ?? false,
|
||||
});
|
||||
for (const event of appendedEvents) {
|
||||
runSubject.next(event);
|
||||
nextSubject.next(event);
|
||||
}
|
||||
|
||||
// Store the run even if it failed (partial events)
|
||||
if (currentRunEvents.length > 0) {
|
||||
this.storeRun(
|
||||
request.threadId,
|
||||
request.input.runId,
|
||||
currentRunEvents,
|
||||
request.input,
|
||||
parentRunId,
|
||||
);
|
||||
}
|
||||
|
||||
// Mark run as complete in database
|
||||
this.setRunState(request.threadId, false);
|
||||
|
||||
if (connection) {
|
||||
connection.agent = undefined;
|
||||
connection.runSubject = undefined;
|
||||
connection.currentEvents = undefined;
|
||||
connection.stopRequested = false;
|
||||
}
|
||||
|
||||
// Don't emit error to the subject, just complete it
|
||||
// This allows subscribers to get events emitted before the error
|
||||
runSubject.complete();
|
||||
nextSubject.complete();
|
||||
|
||||
ACTIVE_CONNECTIONS.delete(request.threadId);
|
||||
}
|
||||
};
|
||||
|
||||
// Bridge previous events if they exist
|
||||
if (prevSubject) {
|
||||
prevSubject.subscribe({
|
||||
next: (e) => nextSubject.next(e),
|
||||
error: (err) => nextSubject.error(err),
|
||||
complete: () => {
|
||||
// Don't complete nextSubject here - it needs to stay open for new events
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
// Start the agent execution immediately (not lazily)
|
||||
runAgent();
|
||||
|
||||
// Return the run subject (only agent events, no injected messages)
|
||||
return runSubject.asObservable();
|
||||
}
|
||||
|
||||
connect(request: AgentRunnerConnectRequest): Observable<BaseEvent> {
|
||||
const connectionSubject = new ReplaySubject<BaseEvent>(Infinity);
|
||||
|
||||
// Load historic runs from database
|
||||
const historicRuns = this.getHistoricRuns(request.threadId);
|
||||
|
||||
// Collect all historic events from database
|
||||
const allHistoricEvents: BaseEvent[] = [];
|
||||
for (const run of historicRuns) {
|
||||
allHistoricEvents.push(...run.events);
|
||||
}
|
||||
|
||||
// Compact all events together before emitting
|
||||
const compactedEvents = compactEvents(allHistoricEvents);
|
||||
|
||||
// Emit compacted events and track message IDs
|
||||
const emittedMessageIds = new Set<string>();
|
||||
for (const event of compactedEvents) {
|
||||
connectionSubject.next(event);
|
||||
if ("messageId" in event && typeof event.messageId === "string") {
|
||||
emittedMessageIds.add(event.messageId);
|
||||
}
|
||||
}
|
||||
|
||||
// Bridge active run to connection if exists
|
||||
const activeConnection = ACTIVE_CONNECTIONS.get(request.threadId);
|
||||
const runState = this.getRunState(request.threadId);
|
||||
|
||||
if (
|
||||
activeConnection &&
|
||||
(runState.isRunning || activeConnection.stopRequested)
|
||||
) {
|
||||
activeConnection.subject.subscribe({
|
||||
next: (event) => {
|
||||
// Skip message events that we've already emitted from historic
|
||||
if (
|
||||
"messageId" in event &&
|
||||
typeof event.messageId === "string" &&
|
||||
emittedMessageIds.has(event.messageId)
|
||||
) {
|
||||
return;
|
||||
}
|
||||
connectionSubject.next(event);
|
||||
},
|
||||
complete: () => connectionSubject.complete(),
|
||||
error: (err) => connectionSubject.error(err),
|
||||
});
|
||||
} else {
|
||||
// No active run, complete after historic events
|
||||
connectionSubject.complete();
|
||||
}
|
||||
|
||||
return connectionSubject.asObservable();
|
||||
}
|
||||
|
||||
isRunning(request: AgentRunnerIsRunningRequest): Promise<boolean> {
|
||||
const runState = this.getRunState(request.threadId);
|
||||
return Promise.resolve(runState.isRunning);
|
||||
}
|
||||
|
||||
stop(request: AgentRunnerStopRequest): Promise<boolean | undefined> {
|
||||
const runState = this.getRunState(request.threadId);
|
||||
if (!runState.isRunning) {
|
||||
return Promise.resolve(false);
|
||||
}
|
||||
|
||||
const connection = ACTIVE_CONNECTIONS.get(request.threadId);
|
||||
const agent = connection?.agent;
|
||||
|
||||
if (!connection || !agent) {
|
||||
return Promise.resolve(false);
|
||||
}
|
||||
|
||||
if (connection.stopRequested) {
|
||||
return Promise.resolve(false);
|
||||
}
|
||||
|
||||
connection.stopRequested = true;
|
||||
this.setRunState(request.threadId, false);
|
||||
|
||||
try {
|
||||
agent.abortRun();
|
||||
return Promise.resolve(true);
|
||||
} catch (error) {
|
||||
console.error("Failed to abort sqlite agent run", error);
|
||||
connection.stopRequested = false;
|
||||
this.setRunState(request.threadId, true);
|
||||
return Promise.resolve(false);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Close the database connection (for cleanup)
|
||||
*/
|
||||
close(): void {
|
||||
if (this.db) {
|
||||
this.db.close();
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user