Julius Marmingemigrate to single turn store1d975ad8 months agoHistory
473 lines · 15.4 KB · TypeScript
RawDownload
import type { OrchestrationEvent, OrchestrationReadModel, ThreadId } from "@t3tools/contracts";import {  OrchestrationCheckpointSummary,  OrchestrationMessage,  OrchestrationSession,  OrchestrationThread,} from "@t3tools/contracts";import { Effect, Schema } from "effect"; import { toProjectorDecodeError, type OrchestrationProjectorDecodeError } from "./Errors.ts";import {  MessageSentPayloadSchema,  ProjectCreatedPayload,  ProjectDeletedPayload,  ProjectMetaUpdatedPayload,  ThreadActivityAppendedPayload,  ThreadCreatedPayload,  ThreadDeletedPayload,  ThreadMetaUpdatedPayload,  ThreadRevertedPayload,  ThreadSessionSetPayload,  ThreadTurnDiffCompletedPayload,} from "./Schemas.ts"; type ThreadPatch = Partial<Omit<OrchestrationThread, "id" | "projectId">>;const MAX_THREAD_MESSAGES = 2_000;const MAX_THREAD_CHECKPOINTS = 500; function updateThread(  threads: ReadonlyArray<OrchestrationThread>,  threadId: ThreadId,  patch: ThreadPatch,): OrchestrationThread[] {  return threads.map((thread) => (thread.id === threadId ? { ...thread, ...patch } : thread));} function decodeForEvent<A>(  schema: Schema.Schema<A>,  value: unknown,  eventType: OrchestrationEvent["type"],  field: string,): Effect.Effect<A, OrchestrationProjectorDecodeError> {  return Schema.decodeUnknownEffect(schema)(value).pipe(    Effect.mapError(toProjectorDecodeError(`${eventType}:${field}`)),  ) as Effect.Effect<A, OrchestrationProjectorDecodeError>;} function retainThreadMessagesAfterRevert(  messages: ReadonlyArray<OrchestrationMessage>,  retainedTurnIds: ReadonlySet<string>,  turnCount: number,): ReadonlyArray<OrchestrationMessage> {  const retainedMessageIds = new Set<string>();  for (const message of messages) {    if (message.role === "system") {      retainedMessageIds.add(message.id);      continue;    }    if (message.turnId !== null && retainedTurnIds.has(message.turnId)) {      retainedMessageIds.add(message.id);    }  }   const retainedUserCount = messages.filter(    (message) => message.role === "user" && retainedMessageIds.has(message.id),  ).length;  const missingUserCount = Math.max(0, turnCount - retainedUserCount);  if (missingUserCount > 0) {    const fallbackUserMessages = messages      .filter(        (message) =>          message.role === "user" &&          !retainedMessageIds.has(message.id) &&          (message.turnId === null || retainedTurnIds.has(message.turnId)),      )      .toSorted(        (left, right) =>          left.createdAt.localeCompare(right.createdAt) || left.id.localeCompare(right.id),      )      .slice(0, missingUserCount);    for (const message of fallbackUserMessages) {      retainedMessageIds.add(message.id);    }  }   const retainedAssistantCount = messages.filter(    (message) => message.role === "assistant" && retainedMessageIds.has(message.id),  ).length;  const missingAssistantCount = Math.max(0, turnCount - retainedAssistantCount);  if (missingAssistantCount > 0) {    const fallbackAssistantMessages = messages      .filter(        (message) =>          message.role === "assistant" &&          !retainedMessageIds.has(message.id) &&          (message.turnId === null || retainedTurnIds.has(message.turnId)),      )      .toSorted(        (left, right) =>          left.createdAt.localeCompare(right.createdAt) || left.id.localeCompare(right.id),      )      .slice(0, missingAssistantCount);    for (const message of fallbackAssistantMessages) {      retainedMessageIds.add(message.id);    }  }   return messages.filter((message) => retainedMessageIds.has(message.id));} function retainThreadActivitiesAfterRevert(  activities: ReadonlyArray<OrchestrationThread["activities"][number]>,  retainedTurnIds: ReadonlySet<string>,): ReadonlyArray<OrchestrationThread["activities"][number]> {  return activities.filter(    (activity) => activity.turnId === null || retainedTurnIds.has(activity.turnId),  );} export function createEmptyReadModel(nowIso: string): OrchestrationReadModel {  return {    snapshotSequence: 0,    projects: [],    threads: [],    updatedAt: nowIso,  };} export function projectEvent(  model: OrchestrationReadModel,  event: OrchestrationEvent,): Effect.Effect<OrchestrationReadModel, OrchestrationProjectorDecodeError> {  const nextBase: OrchestrationReadModel = {    ...model,    snapshotSequence: event.sequence,    updatedAt: event.occurredAt,  };   switch (event.type) {    case "project.created":      return decodeForEvent(ProjectCreatedPayload, event.payload, event.type, "payload").pipe(        Effect.map((payload) => {          const existing = nextBase.projects.find((entry) => entry.id === payload.projectId);          const nextProject = {            id: payload.projectId,            title: payload.title,            workspaceRoot: payload.workspaceRoot,            defaultModel: payload.defaultModel,            scripts: payload.scripts,            createdAt: payload.createdAt,            updatedAt: payload.updatedAt,            deletedAt: null,          };           return {            ...nextBase,            projects: existing              ? nextBase.projects.map((entry) =>                  entry.id === payload.projectId ? nextProject : entry,                )              : [...nextBase.projects, nextProject],          };        }),      );     case "project.meta-updated":      return decodeForEvent(ProjectMetaUpdatedPayload, event.payload, event.type, "payload").pipe(        Effect.map((payload) => ({          ...nextBase,          projects: nextBase.projects.map((project) =>            project.id === payload.projectId              ? {                  ...project,                  ...(payload.title !== undefined ? { title: payload.title } : {}),                  ...(payload.workspaceRoot !== undefined                    ? { workspaceRoot: payload.workspaceRoot }                    : {}),                  ...(payload.defaultModel !== undefined                    ? { defaultModel: payload.defaultModel }                    : {}),                  ...(payload.scripts !== undefined ? { scripts: payload.scripts } : {}),                  updatedAt: payload.updatedAt,                }              : project,          ),        })),      );     case "project.deleted":      return decodeForEvent(ProjectDeletedPayload, event.payload, event.type, "payload").pipe(        Effect.map((payload) => ({          ...nextBase,          projects: nextBase.projects.map((project) =>            project.id === payload.projectId              ? {                  ...project,                  deletedAt: payload.deletedAt,                  updatedAt: payload.deletedAt,                }              : project,          ),        })),      );     case "thread.created":      return Effect.gen(function* () {        const payload = yield* decodeForEvent(          ThreadCreatedPayload,          event.payload,          event.type,          "payload",        );        const thread: OrchestrationThread = yield* decodeForEvent(          OrchestrationThread,          {            id: payload.threadId,            projectId: payload.projectId,            title: payload.title,            model: payload.model,            branch: payload.branch,            worktreePath: payload.worktreePath,            latestTurnId: null,            createdAt: payload.createdAt,            updatedAt: payload.updatedAt,            deletedAt: null,            messages: [],            activities: [],            checkpoints: [],            session: null,          },          event.type,          "thread",        );        const existing = nextBase.threads.find((entry) => entry.id === thread.id);        return {          ...nextBase,          threads: existing            ? nextBase.threads.map((entry) => (entry.id === thread.id ? thread : entry))            : [...nextBase.threads, thread],        };      });     case "thread.deleted":      return decodeForEvent(ThreadDeletedPayload, event.payload, event.type, "payload").pipe(        Effect.map((payload) => ({          ...nextBase,          threads: updateThread(nextBase.threads, payload.threadId, {            deletedAt: payload.deletedAt,            updatedAt: payload.deletedAt,          }),        })),      );     case "thread.meta-updated":      return decodeForEvent(ThreadMetaUpdatedPayload, event.payload, event.type, "payload").pipe(        Effect.map((payload) => ({          ...nextBase,          threads: updateThread(nextBase.threads, payload.threadId, {            ...(payload.title !== undefined ? { title: payload.title } : {}),            ...(payload.model !== undefined ? { model: payload.model } : {}),            ...(payload.branch !== undefined ? { branch: payload.branch } : {}),            ...(payload.worktreePath !== undefined ? { worktreePath: payload.worktreePath } : {}),            updatedAt: payload.updatedAt,          }),        })),      );     case "thread.message-sent":      return Effect.gen(function* () {        const payload = yield* decodeForEvent(          MessageSentPayloadSchema,          event.payload,          event.type,          "payload",        );        const thread = nextBase.threads.find((entry) => entry.id === payload.threadId);        if (!thread) {          return nextBase;        }         const message: OrchestrationMessage = yield* decodeForEvent(          OrchestrationMessage,          {            id: payload.messageId,            role: payload.role,            text: payload.text,            ...(payload.attachments !== undefined ? { attachments: payload.attachments } : {}),            turnId: payload.turnId,            streaming: payload.streaming,            createdAt: payload.createdAt,            updatedAt: payload.updatedAt,          },          event.type,          "message",        );         const existingMessage = thread.messages.find((entry) => entry.id === message.id);        const messages = existingMessage          ? thread.messages.map((entry) =>              entry.id === message.id                ? {                    ...entry,                    text: message.streaming                      ? `${entry.text}${message.text}`                      : message.text.length > 0                        ? message.text                        : entry.text,                    streaming: message.streaming,                    updatedAt: message.updatedAt,                    turnId: message.turnId,                    ...(message.attachments !== undefined                      ? { attachments: message.attachments }                      : {}),                  }                : entry,            )          : [...thread.messages, message];        const cappedMessages = messages.slice(-MAX_THREAD_MESSAGES);         return {          ...nextBase,          threads: updateThread(nextBase.threads, payload.threadId, {            messages: cappedMessages,            updatedAt: event.occurredAt,          }),        };      });     case "thread.session-set":      return Effect.gen(function* () {        const payload = yield* decodeForEvent(          ThreadSessionSetPayload,          event.payload,          event.type,          "payload",        );        const thread = nextBase.threads.find((entry) => entry.id === payload.threadId);        if (!thread) {          return nextBase;        }         const session: OrchestrationSession = yield* decodeForEvent(          OrchestrationSession,          payload.session,          event.type,          "session",        );         return {          ...nextBase,          threads: updateThread(nextBase.threads, payload.threadId, {            session,            latestTurnId: session.activeTurnId,            updatedAt: event.occurredAt,          }),        };      });     case "thread.turn-diff-completed":      return Effect.gen(function* () {        const payload = yield* decodeForEvent(          ThreadTurnDiffCompletedPayload,          event.payload,          event.type,          "payload",        );        const thread = nextBase.threads.find((entry) => entry.id === payload.threadId);        if (!thread) {          return nextBase;        }         const checkpoint = yield* decodeForEvent(          OrchestrationCheckpointSummary,          {            turnId: payload.turnId,            checkpointTurnCount: payload.checkpointTurnCount,            checkpointRef: payload.checkpointRef,            status: payload.status,            files: payload.files,            assistantMessageId: payload.assistantMessageId,            completedAt: payload.completedAt,          },          event.type,          "checkpoint",        );         const checkpoints = [          ...thread.checkpoints.filter((entry) => entry.turnId !== checkpoint.turnId),          checkpoint,        ]          .toSorted((left, right) => left.checkpointTurnCount - right.checkpointTurnCount)          .slice(-MAX_THREAD_CHECKPOINTS);         return {          ...nextBase,          threads: updateThread(nextBase.threads, payload.threadId, {            checkpoints,            latestTurnId: payload.turnId,            updatedAt: event.occurredAt,          }),        };      });     case "thread.reverted":      return decodeForEvent(ThreadRevertedPayload, event.payload, event.type, "payload").pipe(        Effect.map((payload) => {          const thread = nextBase.threads.find((entry) => entry.id === payload.threadId);          if (!thread) {            return nextBase;          }           const checkpoints = thread.checkpoints            .filter((entry) => entry.checkpointTurnCount <= payload.turnCount)            .toSorted((left, right) => left.checkpointTurnCount - right.checkpointTurnCount)            .slice(-MAX_THREAD_CHECKPOINTS);          const retainedTurnIds = new Set(checkpoints.map((checkpoint) => checkpoint.turnId));          const messages = retainThreadMessagesAfterRevert(            thread.messages,            retainedTurnIds,            payload.turnCount,          ).slice(-MAX_THREAD_MESSAGES);          const activities = retainThreadActivitiesAfterRevert(thread.activities, retainedTurnIds);           const latestTurnId =            checkpoints.length > 0 ? (checkpoints[checkpoints.length - 1]?.turnId ?? null) : null;           return {            ...nextBase,            threads: updateThread(nextBase.threads, payload.threadId, {              checkpoints,              messages,              activities,              latestTurnId,              updatedAt: event.occurredAt,            }),          };        }),      );     case "thread.activity-appended":      return decodeForEvent(        ThreadActivityAppendedPayload,        event.payload,        event.type,        "payload",      ).pipe(        Effect.map((payload) => {          const thread = nextBase.threads.find((entry) => entry.id === payload.threadId);          if (!thread) {            return nextBase;          }           const activities = [            ...thread.activities.filter((entry) => entry.id !== payload.activity.id),            payload.activity,          ]            .toSorted((left, right) => left.createdAt.localeCompare(right.createdAt))            .slice(-500);           return {            ...nextBase,            threads: updateThread(nextBase.threads, payload.threadId, {              activities,              updatedAt: event.occurredAt,            }),          };        }),      );     default:      return Effect.succeed(nextBase);  }}