popagent

Latest change 7a8a4e2ba0b7062e5ea31eb88cd856b14c3cfbed - Add workspace chat bulk deletion by AkurAI Build

import { DEFAULT_WORKSPACE_ID, type ChatMessage, type Session, type SessionSummary } from "./api-types";
import { memory, storage } from "./storage";
import { retryableInit } from "./retryable-init";

export const RESOURCE_ID = "popagent-user";
const NEW_CHAT_TITLE = "New chat";

function titleFrom(messages: ChatMessage[]): string {
  const firstUser = messages.find((message) => message.role === "user");
  const text = firstUser?.parts.find(
    (part): part is { type: "text"; text: string } =>
      typeof part === "object" && part !== null && "type" in part && "text" in part &&
      part.type === "text" && typeof part.text === "string",
  )?.text.trim();
  if (!text) return NEW_CHAT_TITLE;
  return text.length > 48 ? `${text.slice(0, 47).trimEnd()}…` : text;
}

export class SessionStore {
  readonly storage = storage;
  readonly memory = memory;
  private readonly initializeOnce = retryableInit(() => this.initialize());

  init(): Promise<void> {
    return this.initializeOnce();
  }

  private async initialize() {
    await this.storage.init();
    await this.storage.db.none(`
      CREATE TABLE IF NOT EXISTS popagent_sessions (
        id TEXT PRIMARY KEY,
        title TEXT NOT NULL,
        model TEXT NOT NULL,
        workspace_id TEXT NOT NULL DEFAULT '${DEFAULT_WORKSPACE_ID}',
        messages JSONB NOT NULL DEFAULT '[]'::jsonb,
        revision INTEGER NOT NULL DEFAULT 0,
        title_locked BOOLEAN NOT NULL DEFAULT FALSE,
        archived_at TIMESTAMPTZ,
        created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
        updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
      );
      ALTER TABLE popagent_sessions ADD COLUMN IF NOT EXISTS workspace_id TEXT NOT NULL DEFAULT '${DEFAULT_WORKSPACE_ID}';
      ALTER TABLE popagent_sessions ADD COLUMN IF NOT EXISTS title_locked BOOLEAN NOT NULL DEFAULT FALSE;
      ALTER TABLE popagent_sessions ADD COLUMN IF NOT EXISTS archived_at TIMESTAMPTZ;
      ALTER TABLE popagent_sessions ADD COLUMN IF NOT EXISTS revision INTEGER NOT NULL DEFAULT 0;
      CREATE INDEX IF NOT EXISTS popagent_sessions_workspace_archive_updated_idx
        ON popagent_sessions (workspace_id, archived_at, updated_at DESC);
    `);
  }

  async list(archived = false, workspaceId?: string): Promise<SessionSummary[]> {
    await this.init();
    return this.storage.db.any<SessionSummary>(`
      SELECT id, title, model, workspace_id AS "workspaceId", revision,
             archived_at::text AS "archivedAt",
             created_at::text AS "createdAt", updated_at::text AS "updatedAt"
      FROM popagent_sessions
      WHERE archived_at IS ${archived ? "NOT NULL" : "NULL"}
        AND ($1::text IS NULL OR workspace_id = $1)
      ORDER BY updated_at DESC
      LIMIT 200
    `, [workspaceId ?? null]);
  }

  async listForWorkspace(workspaceId: string): Promise<SessionSummary[]> {
    await this.init();
    return this.storage.db.any<SessionSummary>(`
      SELECT id, title, model, workspace_id AS "workspaceId", revision,
             archived_at::text AS "archivedAt",
             created_at::text AS "createdAt", updated_at::text AS "updatedAt"
      FROM popagent_sessions
      WHERE workspace_id = $1
      ORDER BY updated_at DESC
    `, [workspaceId]);
  }

  async create(model: string, workspaceId = DEFAULT_WORKSPACE_ID): Promise<Session> {
    await this.init();
    const id = crypto.randomUUID();
    const row = await this.storage.db.one<SessionSummary>(`
      INSERT INTO popagent_sessions (id, title, model, workspace_id)
      VALUES ($1, $2, $3, $4)
      RETURNING id, title, model, workspace_id AS "workspaceId", revision,
                archived_at::text AS "archivedAt",
                created_at::text AS "createdAt", updated_at::text AS "updatedAt"
    `, [id, NEW_CHAT_TITLE, model, workspaceId]);
    await this.memory.saveThread({
      thread: {
        id,
        resourceId: RESOURCE_ID,
        title: NEW_CHAT_TITLE,
        metadata: { model, workspaceId },
        createdAt: new Date(row.createdAt),
        updatedAt: new Date(row.updatedAt),
      },
    });
    return { ...row, messages: [] };
  }

  async get(id: string): Promise<Session | undefined> {
    await this.init();
    const row = await this.storage.db.oneOrNone<Session>(`
      SELECT id, title, model, workspace_id AS "workspaceId", revision, messages,
             archived_at::text AS "archivedAt",
             created_at::text AS "createdAt", updated_at::text AS "updatedAt"
      FROM popagent_sessions
      WHERE id = $1
    `, [id]);
    return row ?? undefined;
  }

  async save(
    id: string,
    input: { model: string; messages: ChatMessage[]; expectedRevision: number },
  ): Promise<Session | undefined | "conflict"> {
    await this.init();
    const generatedTitle = titleFrom(input.messages);
    const row = await this.storage.db.oneOrNone<SessionSummary>(`
      UPDATE popagent_sessions
      SET title = CASE WHEN title_locked THEN title ELSE $1 END,
          model = $2,
          messages = $3::jsonb,
          revision = revision + 1,
          updated_at = NOW()
      WHERE id = $4 AND revision = $5
      RETURNING id, title, model, workspace_id AS "workspaceId", revision,
                archived_at::text AS "archivedAt",
                created_at::text AS "createdAt", updated_at::text AS "updatedAt"
    `, [generatedTitle, input.model, JSON.stringify(input.messages), id, input.expectedRevision]);
    if (!row) return (await this.get(id)) ? "conflict" : undefined;
    await this.memory.updateThread({
      id,
      title: row.title,
      metadata: { model: input.model, workspaceId: row.workspaceId },
    });
    return { ...row, messages: input.messages };
  }

  async rename(id: string, title: string): Promise<Session | undefined> {
    await this.init();
    const cleanTitle = title.trim().slice(0, 120);
    if (!cleanTitle) return undefined;
    const row = await this.storage.db.oneOrNone<Session>(`
      UPDATE popagent_sessions
      SET title = $2, title_locked = TRUE, updated_at = NOW()
      WHERE id = $1
      RETURNING id, title, model, workspace_id AS "workspaceId", revision, messages,
                archived_at::text AS "archivedAt",
                created_at::text AS "createdAt", updated_at::text AS "updatedAt"
    `, [id, cleanTitle]);
    if (row) await this.memory.updateThread({ id, title: cleanTitle });
    return row ?? undefined;
  }

  async setArchived(id: string, archived: boolean): Promise<Session | undefined> {
    await this.init();
    const row = await this.storage.db.oneOrNone<Session>(`
      UPDATE popagent_sessions
      SET archived_at = ${archived ? "NOW()" : "NULL"}, updated_at = NOW()
      WHERE id = $1
      RETURNING id, title, model, workspace_id AS "workspaceId", revision, messages,
                archived_at::text AS "archivedAt",
                created_at::text AS "createdAt", updated_at::text AS "updatedAt"
    `, [id]);
    return row ?? undefined;
  }

  async delete(id: string): Promise<boolean> {
    await this.init();
    const exists = await this.storage.db.oneOrNone<{ id: string }>(
      "SELECT id FROM popagent_sessions WHERE id = $1",
      [id],
    );
    if (!exists) return false;
    await this.memory.deleteThread(id);
    await this.storage.db.tx(async (db) => {
      await db.none("DELETE FROM mastra_thread_state WHERE \"threadId\" = $1", [id]);
      await db.none("DELETE FROM popagent_sessions WHERE id = $1", [id]);
    });
    return true;
  }
}