wehub-resource-sync
68 lines
2.7 KiB
TypeScript
68 lines
2.7 KiB
TypeScript
import { randomUUID } from "node:crypto";
|
|
import {
|
|
EventType,
|
|
type BaseEvent,
|
|
type Message,
|
|
type MessagesSnapshotEvent,
|
|
type ToolCall,
|
|
type ToolCallResultEvent,
|
|
} from "@ag-ui/client";
|
|
import { InMemoryAgentRunner } from "@copilotkit/runtime/v2";
|
|
import { concatMap, Observable, of } from "rxjs";
|
|
|
|
/**
|
|
* AgentCore stores conversation history server-side via AgentCoreMemorySaver /
|
|
* AgentCoreMemorySessionManager. When CopilotKit reconnects to an existing
|
|
* thread (e.g. page refresh), two issues arise that this runner fixes:
|
|
*
|
|
* 1. Unknown threads — CopilotKit may call `connect()` for a thread that has
|
|
* never had a `run()` (first load). The base runner would error; we emit an
|
|
* empty snapshot instead so the UI initialises cleanly.
|
|
* 2. Missing tool-call results — AgentCore's replayed history contains assistant
|
|
* messages with tool calls but no corresponding TOOL_CALL_RESULT events.
|
|
* CopilotKit needs those to reconcile its message state, so we synthesise
|
|
* empty results for each past tool call before the snapshot.
|
|
*/
|
|
export class AgentCoreRunner extends InMemoryAgentRunner {
|
|
private readonly knownThreadIds = new Set<string>();
|
|
|
|
override run(request: Parameters<InMemoryAgentRunner["run"]>[0]) {
|
|
if (request.threadId) this.knownThreadIds.add(request.threadId);
|
|
return super.run(request);
|
|
}
|
|
|
|
override connect(request: Parameters<InMemoryAgentRunner["connect"]>[0]) {
|
|
if (!request.threadId || !this.knownThreadIds.has(request.threadId)) {
|
|
const threadId = request.threadId ?? randomUUID();
|
|
const runId = randomUUID();
|
|
return of(
|
|
{ type: EventType.RUN_STARTED, threadId, runId },
|
|
{ type: EventType.MESSAGES_SNAPSHOT, messages: [] },
|
|
{ type: EventType.RUN_FINISHED, threadId, runId },
|
|
) as Observable<BaseEvent>;
|
|
}
|
|
|
|
return super.connect(request).pipe(
|
|
concatMap((event: BaseEvent) => {
|
|
if (event.type !== EventType.MESSAGES_SNAPSHOT) return of(event);
|
|
const snapshot = event as MessagesSnapshotEvent;
|
|
const replayedResults: ToolCallResultEvent[] =
|
|
snapshot.messages.flatMap((message: Message) => {
|
|
if (message.role !== "assistant" || !message.toolCalls?.length)
|
|
return [];
|
|
return message.toolCalls.map<ToolCallResultEvent>(
|
|
(toolCall: ToolCall) => ({
|
|
type: EventType.TOOL_CALL_RESULT,
|
|
toolCallId: toolCall.id,
|
|
messageId: `${toolCall.id}-result`,
|
|
content: "",
|
|
role: "tool",
|
|
}),
|
|
);
|
|
});
|
|
return of(...replayedResults, snapshot) as Observable<BaseEvent>;
|
|
}),
|
|
);
|
|
}
|
|
}
|