@ai-setting/roy-agent-core 1.0.0

package/src/env/session/storage/memory.ts

@@ -0,0 +1,502 @@
+/**
+ * @fileoverview Memory Session Store
+ * 
+ * In-memory implementation of SessionStore with checkpoint support
+ */
+
+import type {
+  Session,
+  CreateSessionOptions,
+  UpdateSessionOptions,
+  ListSessionsOptions,
+  SessionMessage,
+  SessionCheckpoint,
+  GetMessagesOptions,
+  SessionCheckpointsMeta,
+  SearchMessagesOptions,
+  MessageMatch,
+} from "../types";
+import type { SessionStore } from "../session-store";
+import { SessionMessageConverter } from "../session-message-converter";
+import { randomUUID } from "crypto";
+
+/**
+ * Generate a unique session ID using crypto UUID
+ */
+function generateSessionId(): string {
+  return `session_${randomUUID()}`;
+}
+
+/**
+ * Generate a unique message ID using crypto UUID
+ */
+function generateMessageId(): string {
+  return `msg_${randomUUID()}`;
+}
+
+/**
+ * Memory-based Session Store with checkpoint support
+ */
+export class MemorySessionStore implements SessionStore {
+  private sessions = new Map<string, Session>();
+  private messages = new Map<string, SessionMessage[]>();
+  private checkpoints = new Map<string, SessionCheckpoint[]>();
+  private messageConverter = new SessionMessageConverter();
+
+  // ========================================================================
+  // Session CRUD
+  // ========================================================================
+
+  async create(options: CreateSessionOptions): Promise<Session> {
+    const now = Date.now();
+    const session: Session = {
+      id: options.id ?? generateSessionId(),  // Use custom ID if provided
+      title: options.title ?? "Session",
+      directory: options.directory ?? "/tmp",
+      parentID: options.parentID,
+      createdAt: now,
+      updatedAt: now,
+      messageCount: 0,
+      metadata: {
+        ...options.metadata,
+        messageCount: 0,
+        checkpoints: { checkpoints: [] },
+      },
+    };
+
+    // Check for duplicate ID
+    if (this.sessions.has(session.id)) {
+      throw new Error(`Session with ID '${session.id}' already exists`);
+    }
+
+    this.sessions.set(session.id, session);
+    this.messages.set(session.id, []);
+    this.checkpoints.set(session.id, []);
+
+    return session;
+  }
+
+  async get(id: string): Promise<Session | undefined> {
+    return this.sessions.get(id);
+  }
+
+  async list(options?: ListSessionsOptions): Promise<Session[]> {
+    let sessions = Array.from(this.sessions.values());
+
+    // Filter by metadata
+    if (options?.filter?.metadata) {
+      sessions = sessions.filter((s) => {
+        for (const [key, value] of Object.entries(options.filter!.metadata!)) {
+          if (s.metadata?.[key] !== value) return false;
+        }
+        return true;
+      });
+    }
+
+    // Filter by time range
+    if (options?.filter?.startTime) {
+      sessions = sessions.filter((s) => s.createdAt >= options.filter!.startTime!);
+    }
+    if (options?.filter?.endTime) {
+      sessions = sessions.filter((s) => s.createdAt <= options.filter!.endTime!);
+    }
+
+    // Sort
+    const sortField = options?.sort?.field ?? "updatedAt";
+    const sortOrder = options?.sort?.order ?? "desc";
+    sessions.sort((a, b) => {
+      let aVal: number | string;
+      let bVal: number | string;
+
+      if (sortField === "title") {
+        aVal = a.title;
+        bVal = b.title;
+      } else if (sortField === "createdAt") {
+        aVal = a.createdAt;
+        bVal = b.createdAt;
+      } else {
+        aVal = a.updatedAt;
+        bVal = b.updatedAt;
+      }
+
+      if (aVal < bVal) return sortOrder === "asc" ? -1 : 1;
+      if (aVal > bVal) return sortOrder === "asc" ? 1 : -1;
+      return 0;
+    });
+
+    // Paginate
+    const offset = options?.offset ?? 0;
+    const limit = options?.limit ?? 100;
+    return sessions.slice(offset, offset + limit);
+  }
+
+  async update(id: string, updates: UpdateSessionOptions): Promise<boolean> {
+    const session = this.sessions.get(id);
+    if (!session) return false;
+
+    if (updates.title !== undefined) {
+      session.title = updates.title;
+    }
+    if (updates.metadata !== undefined) {
+      session.metadata = { ...session.metadata, ...updates.metadata };
+    }
+    session.updatedAt = Date.now();
+
+    return true;
+  }
+
+  async delete(id: string): Promise<boolean> {
+    const deleted = this.sessions.delete(id);
+    if (deleted) {
+      this.messages.delete(id);
+      this.checkpoints.delete(id);
+    }
+    return deleted;
+  }
+
+  async getCount(): Promise<number> {
+    return this.sessions.size;
+  }
+
+  // ========================================================================
+  // Messages
+  // ========================================================================
+
+  async addMessage(
+    sessionId: string,
+    message: { role: string; content: string; parts?: unknown[]; metadata?: Record<string, unknown> }
+  ): Promise<string> {
+    const session = this.sessions.get(sessionId);
+    if (!session) {
+      throw new Error(`Session not found: ${sessionId}`);
+    }
+
+    // Create raw message
+    const rawMessage: Partial<SessionMessage> = {
+      sessionID: sessionId,
+      role: message.role as SessionMessage["role"],
+      content: message.content,
+      parts: message.parts as SessionMessage["parts"],
+      metadata: message.metadata as SessionMessage["metadata"],
+    };
+
+    // Apply migration to ensure new format (backward compatibility)
+    const msg = this.messageConverter.migrateMessage(rawMessage);
+    msg.id = generateMessageId();
+    msg.timestamp = Date.now();
+
+    const sessionMessages = this.messages.get(sessionId) ?? [];
+    sessionMessages.push(msg);
+    this.messages.set(sessionId, sessionMessages);
+
+    // Update message count (non-archived only)
+    session.messageCount = sessionMessages.filter((m) => !m.isArchived).length;
+    
+    // Update metadata with messageCount
+    session.metadata = {
+      ...session.metadata,
+      messageCount: session.messageCount
+    };
+    
+    session.updatedAt = Date.now();
+
+    return msg.id;
+  }
+
+  async getMessages(
+    sessionId: string,
+    offset = 0,
+    limit = 50,
+    options?: GetMessagesOptions
+  ): Promise<SessionMessage[]> {
+    let messages = this.messages.get(sessionId) ?? [];
+
+    // Handle archived messages
+    if (!options?.includeArchived) {
+      messages = messages.filter((m) => !m.isArchived);
+    }
+
+    // Handle fromCheckpointId
+    if (options?.fromCheckpointId) {
+      const checkpoint = await this.getCheckpoint(sessionId, options.fromCheckpointId);
+      if (checkpoint) {
+        const checkpointIndex = messages.findIndex(
+          (m) => m.metadata?.checkpointId === options.fromCheckpointId
+        );
+        if (checkpointIndex >= 0) {
+          messages = messages.slice(checkpointIndex);
+        }
+      }
+    }
+
+    // Handle negative offset: count total messages first, then calculate actual offset
+    // offset = -1, limit = 50: latest 50 messages → offset = total - limit
+    // offset = -2, limit = 50: 2nd latest 50 messages → offset = total - 2*limit
+    let actualOffset = offset;
+    if (offset < 0) {
+      const total = messages.length;
+      actualOffset = Math.max(0, total + offset * limit);
+    }
+
+    // Handle reverse
+    if (options?.reverse) {
+      messages = messages.reverse();
+    }
+
+    return messages.slice(actualOffset, actualOffset + limit);
+  }
+
+  async getMessageCount(sessionId: string, includeArchived = false): Promise<number> {
+    const messages = this.messages.get(sessionId) ?? [];
+    if (includeArchived) {
+      return messages.length;
+    }
+    return messages.filter((m) => !m.isArchived).length;
+  }
+
+  async getMessageIndexes(sessionId: string, role: string): Promise<number[]> {
+    const messages = this.messages.get(sessionId) ?? [];
+    const indexes: number[] = [];
+    
+    // Return index based on timestamp ASC order (same as getMessages default)
+    messages
+      .filter(m => !m.isArchived)
+      .forEach((msg, filteredIdx) => {
+        if (msg.role === role) {
+          indexes.push(filteredIdx);
+        }
+      });
+    
+    return indexes;
+  }
+
+  // ========================================================================
+  // Checkpoint Management
+  // ========================================================================
+
+  async saveCheckpoint(sessionId: string, checkpoint: SessionCheckpoint): Promise<void> {
+    const session = this.sessions.get(sessionId);
+    if (!session) {
+      throw new Error(`Session not found: ${sessionId}`);
+    }
+
+    // Save to checkpoints map
+    const cps = this.checkpoints.get(sessionId) ?? [];
+    cps.push(checkpoint);
+    this.checkpoints.set(sessionId, cps);
+
+    // Update session metadata
+    const meta = session.metadata as Record<string, unknown> || {};
+    const checkpointsMeta: SessionCheckpointsMeta = (meta.checkpoints as SessionCheckpointsMeta) || { checkpoints: [] };
+    checkpointsMeta.checkpoints.push({
+      id: checkpoint.id,
+      messageIndex: checkpoint.messageIndex,
+      title: checkpoint.title,
+      createdAt: checkpoint.createdAt,
+      type: checkpoint.type,
+    });
+    checkpointsMeta.latestCheckpointId = checkpoint.id;
+    
+    meta.checkpoints = checkpointsMeta;
+    
+    // Store checkpoint details separately
+    if (!meta.checkpointDetails) {
+      meta.checkpointDetails = {};
+    }
+    (meta.checkpointDetails as Record<string, SessionCheckpoint>)[checkpoint.id] = checkpoint;
+    
+    session.metadata = meta as Session["metadata"];
+    session.updatedAt = Date.now();
+  }
+
+  async getCheckpoints(sessionId: string): Promise<SessionCheckpoint[]> {
+    return this.checkpoints.get(sessionId) ?? [];
+  }
+
+  async getCheckpoint(sessionId: string, checkpointId: string): Promise<SessionCheckpoint | undefined> {
+    const cps = this.checkpoints.get(sessionId) ?? [];
+    return cps.find((cp) => cp.id === checkpointId);
+  }
+
+  async deleteCheckpoint(sessionId: string, checkpointId: string): Promise<boolean> {
+    const session = this.sessions.get(sessionId);
+    if (!session) return false;
+
+    const cps = this.checkpoints.get(sessionId) ?? [];
+    const index = cps.findIndex((cp) => cp.id === checkpointId);
+    if (index === -1) return false;
+
+    // Remove checkpoint
+    cps.splice(index, 1);
+    this.checkpoints.set(sessionId, cps);
+
+    // Restore archived messages
+    const messages = this.messages.get(sessionId) ?? [];
+    for (const msg of messages) {
+      if (msg.checkpointId === checkpointId) {
+        msg.isArchived = false;
+        msg.archivedAt = undefined;
+        msg.checkpointId = undefined;
+      }
+    }
+
+    // Update session metadata
+    const meta = session.metadata as Record<string, unknown> || {};
+    const checkpointsMeta: SessionCheckpointsMeta = (meta.checkpoints as SessionCheckpointsMeta) || { checkpoints: [] };
+    checkpointsMeta.checkpoints = checkpointsMeta.checkpoints.filter((cp) => cp.id !== checkpointId);
+    checkpointsMeta.latestCheckpointId = cps[cps.length - 1]?.id;
+    
+    if (meta.checkpointDetails) {
+      delete (meta.checkpointDetails as Record<string, SessionCheckpoint>)[checkpointId];
+    }
+    
+    session.metadata = meta as Session["metadata"];
+    session.messageCount = messages.filter((m) => !m.isArchived).length;
+    session.updatedAt = Date.now();
+
+    return true;
+  }
+
+  async archiveMessages(sessionId: string, checkpointId: string, beforeIndex: number): Promise<number> {
+    const session = this.sessions.get(sessionId);
+    if (!session) return 0;
+
+    const messages = this.messages.get(sessionId) ?? [];
+    let archivedCount = 0;
+
+    for (let i = 0; i < beforeIndex; i++) {
+      if (!messages[i].isArchived) {
+        messages[i].isArchived = true;
+        messages[i].archivedAt = Date.now();
+        messages[i].checkpointId = checkpointId;
+        archivedCount++;
+      }
+    }
+
+    // Update session message count (non-archived only)
+    session.messageCount = messages.filter((m) => !m.isArchived).length;
+    session.updatedAt = Date.now();
+
+    return archivedCount;
+  }
+
+  async getActiveMessages(sessionId: string, offset = 0, limit = 50): Promise<SessionMessage[]> {
+    const messages = this.messages.get(sessionId) ?? [];
+    const activeMessages = messages.filter((m) => !m.isArchived);
+    return activeMessages.slice(offset, offset + limit);
+  }
+
+  // ========================================================================
+  // Search [new]
+  // ========================================================================
+
+  async searchMessages(options: SearchMessagesOptions): Promise<MessageMatch[]> {
+    const matches: MessageMatch[] = [];
+    const {
+      query,
+      sessionId,
+      limit = 10,
+      beforeTime,
+      afterTime,
+      includeArchived = false,
+    } = options;
+
+    // Determine which sessions to search
+    let sessionsToSearch: Session[];
+    if (sessionId) {
+      const session = await this.get(sessionId);
+      sessionsToSearch = session ? [session] : [];
+    } else {
+      sessionsToSearch = Array.from(this.sessions.values());
+    }
+
+    // Parse query for simple text matching
+    const queryTerms = this.parseQuery(query);
+
+    // Search through each session's messages
+    for (const session of sessionsToSearch) {
+      // Time constraint on session
+      if (beforeTime && session.updatedAt > beforeTime) continue;
+      if (afterTime && session.updatedAt < afterTime) continue;
+
+      let messages = this.messages.get(session.id) ?? [];
+      
+      // Filter archived messages if needed
+      if (!includeArchived) {
+        messages = messages.filter((m) => !m.isArchived);
+      }
+
+      // Search through messages
+      for (const msg of messages) {
+        // Time constraint on message
+        if (beforeTime && msg.timestamp > beforeTime) continue;
+        if (afterTime && msg.timestamp < afterTime) continue;
+
+        // Simple text search
+        const content = msg.content.toLowerCase();
+        let matchesQuery = true;
+
+        for (const term of queryTerms.required) {
+          if (!content.includes(term.toLowerCase())) {
+            matchesQuery = false;
+            break;
+          }
+        }
+
+        for (const term of queryTerms.excluded) {
+          if (content.includes(term.toLowerCase())) {
+            matchesQuery = false;
+            break;
+          }
+        }
+
+        if (matchesQuery) {
+          matches.push({
+            messageId: msg.id,
+            sessionId: session.id,
+            role: msg.role,
+            content: msg.content,
+            timestamp: msg.timestamp,
+          });
+
+          if (matches.length >= limit) {
+            return matches;
+          }
+        }
+      }
+    }
+
+    return matches;
+  }
+
+  /**
+   * Parse simple query into terms
+   */
+  private parseQuery(query: string): { required: string[]; excluded: string[] } {
+    // Remove quoted phrases (simplified)
+    const terms = query.split(/\s+/).filter((t) => !t.startsWith('"') && !t.endsWith('"'));
+    
+    const required: string[] = [];
+    const excluded: string[] = [];
+    
+    for (const term of terms) {
+      if (term.startsWith("-") || term.startsWith("NOT ")) {
+        excluded.push(term.replace(/^-|^NOT /i, ""));
+      } else {
+        required.push(term);
+      }
+    }
+    
+    return { required, excluded };
+  }
+
+  // ========================================================================
+  // Lifecycle
+  // ========================================================================
+
+  async close(): Promise<void> {
+    this.sessions.clear();
+    this.messages.clear();
+    this.checkpoints.clear();
+  }
+}