laimiaoai-agent-new/src/sessions/mysql_store.ts

import { randomUUID } from "node:crypto";
import { createPool, type Pool, type RowDataPacket } from "mysql2/promise";

import type { MysqlConfig } from "../config.js";
import type { Logger } from "../logger.js";
import type { Identity } from "../types.js";
import type { SessionMessage, SessionMeta, SessionStoreLike } from "./store.js";

const SESSION_DDL = `
CREATE TABLE IF NOT EXISTS agent_session (
  session_id VARCHAR(64) NOT NULL,
  uid VARCHAR(64) NOT NULL,
  team_id VARCHAR(64) NOT NULL DEFAULT '',
  title VARCHAR(200) NOT NULL DEFAULT '',
  type INT NOT NULL DEFAULT 0,
  last_mode VARCHAR(16) NOT NULL DEFAULT 'agent',
  message_count INT NOT NULL DEFAULT 0,
  add_time INT NOT NULL DEFAULT 0,
  update_time INT NOT NULL DEFAULT 0,
  PRIMARY KEY (session_id),
  KEY idx_uid_time (uid, update_time),
  KEY idx_team_time (team_id, update_time)
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4
`;

const MESSAGE_DDL = `
CREATE TABLE IF NOT EXISTS agent_session_message (
  id VARCHAR(64) NOT NULL,
  session_id VARCHAR(64) NOT NULL,
  role VARCHAR(16) NOT NULL,
  mode VARCHAR(16) NOT NULL DEFAULT 'agent',
  content MEDIUMTEXT,
  thinking MEDIUMTEXT,
  media MEDIUMTEXT,
  status VARCHAR(16) NOT NULL DEFAULT 'ok',
  error VARCHAR(2000) NOT NULL DEFAULT '',
  add_time INT NOT NULL DEFAULT 0,
  PRIMARY KEY (id),
  KEY idx_session_time (session_id, add_time)
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4
`;

interface SessionRow extends RowDataPacket {
  session_id: string;
  uid: string;
  team_id: string;
  title: string;
  type: number;
  last_mode: string;
  message_count: number;
  add_time: number;
  update_time: number;
}

interface MessageRow extends RowDataPacket {
  id: string;
  session_id: string;
  role: "user" | "assistant";
  mode: "agent" | "image" | "video";
  content: string | null;
  thinking: string | null;
  media: string | null;
  status: "ok" | "error" | "loading";
  error: string | null;
  add_time: number;
}

function toMeta(row: SessionRow): SessionMeta {
  return {
    session_id: row.session_id,
    uid: row.uid,
    team_id: row.team_id,
    title: row.title,
    type: row.type,
    last_mode: row.last_mode,
    message_count: row.message_count,
    add_time: row.add_time,
    update_time: row.update_time,
  };
}

function toMessage(row: MessageRow): SessionMessage {
  let media: SessionMessage["media"] = [];
  if (row.media) {
    try {
      const parsed = JSON.parse(row.media);
      if (Array.isArray(parsed)) media = parsed;
    } catch {
      media = [];
    }
  }
  const message: SessionMessage = {
    id: row.id,
    session_id: row.session_id,
    role: row.role,
    mode: row.mode,
    content: row.content ?? "",
    media,
    status: row.status,
    add_time: row.add_time,
  };
  if (row.thinking) message.thinking = row.thinking;
  if (row.error) message.error = row.error;
  return message;
}

export class MySqlSessionStore implements SessionStoreLike {
  private readonly pool: Pool;

  constructor(
    config: MysqlConfig,
    private readonly logger: Logger,
  ) {
    this.pool = createPool({
      host: config.host,
      port: config.port,
      user: config.user,
      password: config.password,
      database: config.database,
      connectionLimit: config.connectionLimit,
      connectTimeout: config.connectTimeoutMs,
      charset: "utf8mb4",
      supportBigNumbers: true,
      bigNumberStrings: true,
    });
  }

  async init(): Promise<void> {
    await this.pool.query(SESSION_DDL);
    await this.pool.query(MESSAGE_DDL);
    this.logger.info("session tables ready");
  }

  async close(): Promise<void> {
    await this.pool.end();
  }

  async create(identity: Identity, input: { title?: string; type?: number } = {}): Promise<SessionMeta> {
    const now = Math.floor(Date.now() / 1000);
    const meta: SessionMeta = {
      session_id: randomUUID(),
      uid: identity.uid,
      team_id: identity.team_id || "",
      title: input.title?.trim() ?? "",
      type: input.type ?? 0,
      last_mode: "agent",
      message_count: 0,
      add_time: now,
      update_time: now,
    };
    await this.pool.execute(
      `INSERT INTO agent_session
        (session_id, uid, team_id, title, type, last_mode, message_count, add_time, update_time)
       VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)`,
      [
        meta.session_id,
        meta.uid,
        meta.team_id,
        meta.title,
        meta.type,
        meta.last_mode,
        meta.message_count,
        meta.add_time,
        meta.update_time,
      ],
    );
    return meta;
  }

  async get(sessionId: string): Promise<SessionMeta | undefined> {
    const [rows] = await this.pool.execute<SessionRow[]>(
      "SELECT * FROM agent_session WHERE session_id = ? LIMIT 1",
      [sessionId],
    );
    return rows[0] ? toMeta(rows[0]) : undefined;
  }

  async getOwned(identity: Identity, sessionId: string): Promise<SessionMeta | undefined> {
    const session = await this.get(sessionId);
    return session && session.uid === identity.uid ? session : undefined;
  }

  async list(
    filter: { uid?: string; team_id?: string; limit?: number; offset?: number } = {},
  ): Promise<{ sessions: SessionMeta[]; total: number }> {
    const where: string[] = [];
    const params: string[] = [];
    if (filter.uid) {
      where.push("uid = ?");
      params.push(filter.uid);
    }
    if (filter.team_id) {
      where.push("team_id = ?");
      params.push(filter.team_id);
    }
    const whereSql = where.length ? `WHERE ${where.join(" AND ")}` : "";
    const limit = Math.max(1, Math.min(filter.limit ?? 50, 200));
    const offset = Math.max(0, filter.offset ?? 0);

    const [countRows] = await this.pool.execute<(RowDataPacket & { total: number })[]>(
      `SELECT COUNT(*) AS total FROM agent_session ${whereSql}`,
      params,
    );
    const [rows] = await this.pool.execute<SessionRow[]>(
      `SELECT * FROM agent_session ${whereSql}
       ORDER BY update_time DESC, session_id DESC
       LIMIT ${limit} OFFSET ${offset}`,
      params,
    );
    return { sessions: rows.map(toMeta), total: Number(countRows[0]?.total ?? 0) };
  }

  async append(
    sessionId: string,
    message: Omit<SessionMessage, "id" | "session_id" | "add_time">,
    addTime?: number,
  ): Promise<SessionMessage> {
    const full: SessionMessage = {
      ...message,
      id: randomUUID(),
      session_id: sessionId,
      add_time: addTime ?? Math.floor(Date.now() / 1000),
    };

    const connection = await this.pool.getConnection();
    try {
      await connection.beginTransaction();
      await connection.execute(
        `INSERT INTO agent_session_message
          (id, session_id, role, mode, content, thinking, media, status, error, add_time)
         VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`,
        [
          full.id,
          sessionId,
          full.role,
          full.mode,
          full.content ?? "",
          full.thinking ?? null,
          JSON.stringify(full.media ?? []),
          full.status,
          full.error ?? "",
          full.add_time,
        ],
      );
      await connection.execute(
        `UPDATE agent_session
           SET message_count = message_count + 1,
               update_time = ?,
               last_mode = ?,
               title = CASE WHEN title = '' AND ? = 'user' THEN LEFT(?, 30) ELSE title END
         WHERE session_id = ?`,
        [full.add_time, full.mode, full.role, full.content ?? "", sessionId],
      );
      await connection.commit();
    } catch (error) {
      await connection.rollback();
      throw error;
    } finally {
      connection.release();
    }
    return full;
  }

  async update(
    sessionId: string,
    messageId: string,
    patch: Partial<Pick<SessionMessage, "mode" | "content" | "thinking" | "media" | "status" | "error">>,
  ): Promise<void> {
    const connection = await this.pool.getConnection();
    try {
      await connection.beginTransaction();
      await connection.execute(
        `UPDATE agent_session_message
           SET mode = COALESCE(?, mode),
               content = COALESCE(?, content),
               thinking = COALESCE(?, thinking),
               media = COALESCE(?, media),
               status = COALESCE(?, status),
               error = COALESCE(?, error)
         WHERE id = ? AND session_id = ?`,
        [
          patch.mode ?? null,
          patch.content ?? null,
          patch.thinking ?? null,
          patch.media ? JSON.stringify(patch.media) : null,
          patch.status ?? null,
          patch.error ?? null,
          messageId,
          sessionId,
        ],
      );
      if (patch.mode || patch.status) {
        await connection.execute(
          `UPDATE agent_session
             SET update_time = ?, last_mode = COALESCE(?, last_mode)
           WHERE session_id = ?`,
          [Math.floor(Date.now() / 1000), patch.mode ?? null, sessionId],
        );
      }
      await connection.commit();
    } catch (error) {
      await connection.rollback();
      throw error;
    } finally {
      connection.release();
    }
  }

  async upsertSession(meta: SessionMeta): Promise<void> {
    await this.pool.execute(
      `INSERT INTO agent_session
        (session_id, uid, team_id, title, type, last_mode, message_count, add_time, update_time)
       VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
       ON DUPLICATE KEY UPDATE
         uid = VALUES(uid), team_id = VALUES(team_id), title = VALUES(title), type = VALUES(type),
         last_mode = VALUES(last_mode), message_count = VALUES(message_count),
         add_time = VALUES(add_time), update_time = VALUES(update_time)`,
      [
        meta.session_id,
        meta.uid,
        meta.team_id || "",
        meta.title,
        meta.type,
        meta.last_mode,
        meta.message_count,
        meta.add_time,
        meta.update_time,
      ],
    );
  }

  async insertMessage(message: SessionMessage): Promise<void> {
    await this.pool.execute(
      `INSERT IGNORE INTO agent_session_message
        (id, session_id, role, mode, content, thinking, media, status, error, add_time)
       VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`,
      [
        message.id,
        message.session_id,
        message.role,
        message.mode,
        message.content ?? "",
        message.thinking ?? null,
        JSON.stringify(message.media ?? []),
        message.status,
        message.error ?? "",
        message.add_time,
      ],
    );
  }

  async messages(sessionId: string, limit = 200): Promise<SessionMessage[]> {
    const size = Math.max(1, Math.min(limit, 1000));
    const [rows] = await this.pool.execute<MessageRow[]>(
      `SELECT * FROM agent_session_message
        WHERE session_id = ?
        ORDER BY add_time DESC, id DESC
        LIMIT ${size}`,
      [sessionId],
    );
    return rows.reverse().map(toMessage);
  }
}