Merge pull request #24 from wangjk9527/feat/production-password-auth

feat: add conversation branching checkpoints

Author: wangjk9527

apps/api/src/config-api.ts

@@ -40,6 +40,7 @@ import {
   type QueryHistoryRecord,
   type RunRecord,
   type RunEventRecord,
+  type SessionBranchRecord,
   type SessionRecord,
   type UserRecord
 } from "@datafoundry/metadata";
@@ -53,6 +54,14 @@ import { resolveEffectiveRunConfig } from "./run-input.js";
 import { handleCapabilitiesRequest } from "./routes/capabilities.js";
 import type { ConfigApiContext, ConfigApiResponse } from "./routes/types.js";
 import { sessionTitleDto } from "./session-title.js";
+import {
+  createSessionBranch,
+  latestVisibleConversationSummary,
+  listConversationBranchOptions,
+  listVisibleConversationMessages,
+  resolveSessionLineage,
+  type ConversationBranchOption
+} from "./session-branching.js";
 import { readMultipartFiles, readMultipartUpload } from "./upload-parser.js";
 
 const MAX_JSON_BODY_BYTES = 1024 * 1024;
@@ -359,6 +368,22 @@ const handleSessionRequest = async (
     });
     return ok(sessionTitleDto(session));
   }
+  if (action === "branches" && request.method === "POST") {
+    const body = await readJsonBody(request);
+    const runId = stringValue(body.runId ?? body.run_id);
+    if (!runId) {
+      throw new Error("RUN_ID_REQUIRED");
+    }
+    const title = stringValue(body.title);
+    const created = createSessionBranch({
+      activeSessionId: sessionId,
+      metadataStore: context.metadataStore,
+      runId,
+      ...(title ? { title } : {}),
+      userId: context.userId
+    });
+    return ok(sessionBranchCreatedDto(created), 201);
+  }
   if (action !== "conversation") {
     return methodNotAllowed();
   }
@@ -367,16 +392,24 @@ const handleSessionRequest = async (
   }
 
   const session = context.metadataStore.sessions.get({ user_id: context.userId, session_id: sessionId });
+  const lineage = resolveSessionLineage({
+    metadataStore: context.metadataStore,
+    sessionId,
+    userId: context.userId
+  });
   const requestUrl = new URL(request.url ?? "/", "http://127.0.0.1");
   const limit = clampInteger(Number.parseInt(requestUrl.searchParams.get("limit") ?? "", 10), 1, 200, 80);
-  const messages = context.metadataStore.conversationMessages.listRecent({
-    user_id: context.userId,
-    session_id: sessionId,
-    limit
+  const messages = listVisibleConversationMessages({
+    lineage,
+    limit,
+    metadataStore: context.metadataStore,
+    userId: context.userId
   });
-  const latestSummary = context.metadataStore.conversationSummaries.latest({
-    user_id: context.userId,
-    session_id: sessionId
+  const latestSummary = latestVisibleConversationSummary({
+    lineage,
+    metadataStore: context.metadataStore,
+    sessionId,
+    userId: context.userId
   });
   const runIds = [...new Set([
     ...messages.map((message) => message.run_id),
@@ -386,11 +419,19 @@ const handleSessionRequest = async (
     runId,
     events: context.metadataStore.runEvents.listByRun({ user_id: context.userId, run_id: runId })
   }));
+  const visiblePositions = new Map(messages.map((message, index) => [message.id, index + 1]));
   const checkpoints = runCheckpointDtos({
     context,
     messages,
     runEventGroups,
-    runIds
+    runIds,
+    visiblePositions
+  });
+  const branchOptions = listConversationBranchOptions({
+    lineage,
+    metadataStore: context.metadataStore,
+    sessionId,
+    userId: context.userId
   });
   const pendingInteractions = context.metadataStore.interactions.listPendingBySession({
     user_id: context.userId,
@@ -402,10 +443,12 @@ const handleSessionRequest = async (
     title: session.title ?? "",
     titleSource: session.title_source ?? "fallback",
     updatedAt: session.updated_at,
-    messages: messages.map(conversationMessageDto),
+    messages: messages.map((message, index) => conversationMessageDto(message, index + 1)),
     ...(latestSummary ? { summary: conversationSummaryDto(latestSummary) } : {}),
     runEventRefs: runEventGroups.map(({ runId, events }) => runEventRefDto(runId, events)),
     ...(checkpoints.length > 0 ? { checkpoints } : {}),
+    ...(lineage.branch ? { branch: sessionBranchDto(lineage.branch, session) } : {}),
+    ...(branchOptions.length > 0 ? { branches: branchOptions.map(conversationBranchOptionDto) } : {}),
     toolCalls: runEventGroups.flatMap(({ runId, events }) => toolCallPairDtos(runId, events)),
     pendingInteractions: pendingInteractions.map(pendingInteractionDto),
     restorableCustomEvents: runEventGroups.flatMap(({ runId, events }) =>
@@ -1616,14 +1659,14 @@ const handleArtifactRequest = async (
   });
 };
 
-const conversationMessageDto = (message: ConversationMessageRecord): Record<string, unknown> => ({
+const conversationMessageDto = (message: ConversationMessageRecord, visiblePosition = message.position): Record<string, unknown> => ({
   id: message.id,
   runId: message.run_id,
   role: message.role,
   source: message.source,
   ...(message.message_id ? { messageId: message.message_id } : {}),
   contentText: message.content_text,
-  position: message.position,
+  position: visiblePosition,
   createdAt: message.created_at
 });
 
@@ -1637,6 +1680,38 @@ const sessionListDto = (session: SessionRecord): Record<string, unknown> => ({
   lastMessageAt: session.last_message_at ?? session.updated_at
 });
 
+const sessionBranchCreatedDto = (input: { branch: SessionBranchRecord; session: SessionRecord }): Record<string, unknown> => ({
+  ...sessionBranchDto(input.branch, input.session),
+  session: sessionListDto(input.session)
+});
+
+const sessionBranchDto = (
+  branch: SessionBranchRecord,
+  session?: SessionRecord
+): Record<string, unknown> => ({
+  id: branch.id,
+  sessionId: branch.child_session_id,
+  threadId: branch.child_session_id,
+  parentSessionId: branch.parent_session_id,
+  rootSessionId: branch.root_session_id,
+  forkRunId: branch.fork_run_id,
+  forkMessageEndPosition: branch.fork_message_end_position,
+  createdAt: branch.created_at,
+  ...(session?.title ? { title: session.title } : {})
+});
+
+const conversationBranchOptionDto = (branch: ConversationBranchOption): Record<string, unknown> => ({
+  sessionId: branch.sessionId,
+  threadId: branch.sessionId,
+  parentSessionId: branch.parentSessionId,
+  rootSessionId: branch.rootSessionId,
+  forkRunId: branch.forkRunId,
+  forkMessageEndPosition: branch.forkMessageEndPosition,
+  isOriginal: branch.isOriginal,
+  createdAt: branch.createdAt,
+  ...(branch.title ? { title: branch.title } : {})
+});
+
 const queryHistoryDto = (record: QueryHistoryRecord): Record<string, unknown> => ({
   id: record.id,
   sessionId: record.session_id,
@@ -1689,6 +1764,7 @@ const runCheckpointDtos = (input: {
   messages: ConversationMessageRecord[];
   runEventGroups: Array<{ events: RunEventRecord[]; runId: string }>;
   runIds: string[];
+  visiblePositions?: Map<string, number>;
 }): Record<string, unknown>[] => {
   const eventsByRun = new Map(input.runEventGroups.map((group) => [group.runId, group.events]));
   return input.runIds.flatMap((runId) => {
@@ -1703,7 +1779,8 @@ const runCheckpointDtos = (input: {
       runCheckpointDto({
         events: eventsByRun.get(runId) ?? [],
         messages: input.messages.filter((message) => message.run_id === runId),
-        run
+        run,
+        ...(input.visiblePositions ? { visiblePositions: input.visiblePositions } : {})
       })
     ];
   });
@@ -1713,8 +1790,9 @@ const runCheckpointDto = (input: {
   events: RunEventRecord[];
   messages: ConversationMessageRecord[];
   run: RunRecord;
+  visiblePositions?: Map<string, number>;
 }): Record<string, unknown> => {
-  const positions = input.messages.map((message) => message.position);
+  const positions = input.messages.map((message) => input.visiblePositions?.get(message.id) ?? message.position);
   const firstEvent = input.events[0];
   const lastEvent = input.events.at(-1);
   return {
@@ -2670,6 +2748,18 @@ const errorResponse = (error: unknown): ConfigApiResponse => {
   if (message.startsWith("REVISION_CONFLICT")) {
     return fail(409, "REVISION_CONFLICT", message);
   }
+  if (message.startsWith("RUN_NOT_BRANCHABLE")) {
+    return fail(409, "RUN_NOT_BRANCHABLE", message);
+  }
+  if (message.startsWith("RUN_NOT_VISIBLE")) {
+    return fail(404, "RESOURCE_NOT_FOUND", message);
+  }
+  if (message.startsWith("SESSION_BRANCH_CYCLE")) {
+    return fail(409, "REVISION_CONFLICT", message);
+  }
+  if (message.startsWith("SESSION_BRANCH_PARENT_NOT_VISIBLE")) {
+    return fail(404, "RESOURCE_NOT_FOUND", message);
+  }
   if (message.includes("NOT_FOUND") || message.includes("not found")) {
     return fail(404, "RESOURCE_NOT_FOUND", message);
   }

apps/api/src/conversation-memory.ts

@@ -50,8 +50,19 @@ export type ConversationMemoryRunMessages = {
   report: ConversationMemoryWindowReport;
 };
 
+export type ConversationMemoryHistoryProvider = (input: {
+  excludeRunId: string;
+  limit: number;
+  sessionId: string;
+  userId: string;
+}) => {
+  history: ConversationMessageRecord[];
+  summary?: ConversationSummaryRecord;
+};
+
 export type ConversationMemoryServiceInput = {
   compactMemorySource?: ConversationCompactMemorySource | undefined;
+  historyProvider?: ConversationMemoryHistoryProvider | undefined;
   memoryBridge?: ConversationMemoryBridge | undefined;
   policy?: ConversationMemoryWindowPolicy | undefined;
   repository: ConversationMessageRepository;
@@ -136,16 +147,24 @@ export class ConversationMemoryService {
     runInput: RunAgentInput;
   }): ConversationMemoryRunMessages {
     const policy = normalizePolicy(this.input.policy);
-    const history = this.input.repository.listRecent({
+    const provided = this.input.historyProvider?.({
+      excludeRunId: input.runId,
+      limit: policy.historyLoadLimit,
+      sessionId: this.input.sessionId,
+      userId: this.input.userId
+    });
+    const history = provided?.history ?? this.input.repository.listRecent({
       user_id: this.input.userId,
       session_id: this.input.sessionId,
       exclude_run_id: input.runId,
       limit: policy.historyLoadLimit
     });
-    const summary = this.input.summaryRepository?.latest({
-      user_id: this.input.userId,
-      session_id: this.input.sessionId
-    });
+    const summary = provided
+      ? provided.summary
+      : this.input.summaryRepository?.latest({
+        user_id: this.input.userId,
+        session_id: this.input.sessionId
+      });
     const effectiveHistory = summary ? history.filter((record) => record.position > summary.to_position) : history;
     const toolCheckpoints = this.input.runEvents
       ? buildToolCheckpointsForHistory({

apps/api/src/run-identity.ts

@@ -1,6 +1,10 @@
 import type { BaseEvent, RunAgentInput } from "@ag-ui/client";
 import { type MetadataStore, type RunEventWriter, type RunRecord } from "@datafoundry/metadata";
 import { createHash } from "node:crypto";
+import {
+  isRunVisibleInSessionLineage,
+  resolveSessionLineage
+} from "./session-branching.js";
 
 export type ResolveExistingRunInput = {
   existingRun: RunRecord;
@@ -62,7 +66,19 @@ export const validateParentRun = ({
   if (!parentRun) {
     throw new Error(`PARENT_RUN_NOT_FOUND: ${parentRunId}`);
   }
-  if (parentRun.session_id !== sessionId) {
+  if (parentRun.session_id === sessionId) {
+    return;
+  }
+
+  try {
+    const lineage = resolveSessionLineage({ metadataStore, sessionId, userId });
+    if (!isRunVisibleInSessionLineage({ lineage, run: parentRun })) {
+      throw new Error(`PARENT_RUN_SESSION_MISMATCH: ${parentRunId}`);
+    }
+  } catch (error) {
+    if (error instanceof Error && error.message.includes("PARENT_RUN_SESSION_MISMATCH")) {
+      throw error;
+    }
     throw new Error(`PARENT_RUN_SESSION_MISMATCH: ${parentRunId}`);
   }
 };

apps/api/src/run-memory-assembly.ts

@@ -16,6 +16,11 @@ import {
   type ConversationMemoryEventObserver
 } from "./conversation-memory.js";
 import { createMastraConversationSummarizer } from "./conversation-summarizer.js";
+import {
+  latestVisibleConversationSummary,
+  listVisibleConversationMessages,
+  resolveSessionLineage
+} from "./session-branching.js";
 import { LongTermMemoryService } from "./long-term-memory.js";
 import { createMastraLongTermMemoryExtractor } from "./long-term-memory-extractor.js";
 
@@ -56,6 +61,29 @@ export const createRunMemoryAssembly = async (
       : createMastraConversationMemoryBridge({
         memory: input.taskStateRuntime.memory
       }),
+    historyProvider: ({ excludeRunId, limit, sessionId, userId }) => {
+      const lineage = resolveSessionLineage({
+        metadataStore: input.metadataStore,
+        sessionId,
+        userId
+      });
+      const summary = latestVisibleConversationSummary({
+        lineage,
+        metadataStore: input.metadataStore,
+        sessionId,
+        userId
+      });
+      return {
+        history: listVisibleConversationMessages({
+          excludeRunId,
+          lineage,
+          limit,
+          metadataStore: input.metadataStore,
+          userId
+        }),
+        ...(summary ? { summary } : {})
+      };
+    },
     repository: input.metadataStore.conversationMessages,
     runEvents: input.metadataStore.runEvents,
     sessionId: input.sessionId,

apps/api/src/server.ts

@@ -5,6 +5,7 @@ import {
   copilotRuntimeNodeHttpEndpoint
 } from "@copilotkit/runtime";
 import {
+  CONVERSATION_WORKING_MEMORY_CONFIG,
   createTaskStateRuntime,
   createCustomEvent,
   parseAgentMemoryMode,
@@ -631,6 +632,14 @@ class DataFoundryAgUiAgent extends AbstractAgent {
         });
         subscriber.add(() => unregisterCancel());
 
+        if (this.input.conversationMemoryMode === "working-memory-readonly") {
+          await ensureConversationWorkingMemoryThread({
+            resourceId: this.input.user.id,
+            taskStateRuntime: this.input.taskStateRuntime,
+            threadId: sessionId
+          });
+        }
+
         subscription = agentAssembly.mastraAgent.run({
           ...normalizedRunInput,
           runId,
@@ -891,6 +900,7 @@ const resolveRequestAuth = (
       workspaceId: identity.workspace.id
     };
   }
+
   const token = extractAuthToken(request);
   const workspaceId = sanitizeWorkspaceId(headerString(request.headers["x-workspace-id"]));
   const devUser = metadataStore.users.getById({ user_id: DEV_USER.id });
@@ -922,6 +932,26 @@ const resolveRequestAuth = (
   };
 };
 
+const ensureConversationWorkingMemoryThread = async (input: {
+  resourceId: string;
+  taskStateRuntime: TaskStateRuntime;
+  threadId: string;
+}): Promise<void> => {
+  const existing = await input.taskStateRuntime.memory.getThreadById({
+    resourceId: input.resourceId,
+    threadId: input.threadId
+  });
+  if (existing) {
+    return;
+  }
+  await input.taskStateRuntime.memory.createThread({
+    memoryConfig: CONVERSATION_WORKING_MEMORY_CONFIG,
+    resourceId: input.resourceId,
+    saveThread: true,
+    threadId: input.threadId
+  });
+};
+
 const isPasswordAuth = (config: PasswordAuthConfig): boolean => config.mode === "password";
 
 const extractAuthToken = (request: IncomingMessage): string | undefined => {