MailUI4Agents/gateway/internal/repo/platform_sessions.go

package repo

import (
	"context"
	"database/sql"
	"errors"
	"strings"
	"time"

	"github.com/agentmail/gateway/internal/db"
	"github.com/google/uuid"
)

// ---------- 平台会话镜像 ----------
//
// Agent 平台自己也在开会话：有些经由邮件驱动，有些是人直接在平台界面上开的。
// 写信时想续谈某条会话，得先知道那个工作区下有哪些会话可续 —— 而 Gateway
// 只看得见邮件驱动的那部分。
//
// **由插件在心跳里上报，Gateway 不反向拉取。**
// 当前架构是单向的（Agent 持密钥主动连 Gateway，Gateway 从不外呼）；
// 让 Gateway 去调平台接口需要它保存各平台的地址与凭证，那是另一套信任模型。
// 代价是插件没运行时同步不了 —— 但插件没运行时邮件本来也投不进去。

// PlatformSession 是插件上报的一条平台侧会话。
type PlatformSession struct {
	PlatformID string     `json:"platform_id"`
	Workspace  string     `json:"workspace"`
	Slug       string     `json:"slug,omitempty"`
	Title      string     `json:"title,omitempty"`
	MailDriven bool       `json:"mail_driven"`
	UpdatedAt  *time.Time `json:"updated_at,omitempty"`
}

// maxPlatformSessions 限制单次上报的会话数。
//
// 一个长期运行的平台可以累积上千条会话，而候选列表上千项对人没有意义。
// 插件按最近活跃排序后上报前 N 条即可。
const maxPlatformSessions = 200

// ReplacePlatformSessions 整表替换某 Agent 的平台会话镜像。
//
// 整表替换而非增量合并：镜像是平台当前状态的快照。增量合并会让已经删掉的
// 平台会话永远留在候选列表里，而那正是「选了却送不到」的来源
// —— session 位是三态语义，指向一条不存在的会话会直接 404。
func ReplacePlatformSessions(ctx context.Context, agentName string, list []PlatformSession) error {
	agentName = strings.TrimSpace(agentName)
	if agentName == "" {
		return nil
	}
	if len(list) > maxPlatformSessions {
		list = list[:maxPlatformSessions]
	}

	tx, err := db.DB.BeginTx(ctx, nil)
	if err != nil {
		return err
	}
	defer tx.Rollback()

	if _, err := tx.ExecContext(ctx,
		`DELETE FROM agent_platform_sessions WHERE agent_name = $1`, agentName); err != nil {
		return err
	}

	seen := map[string]bool{}
	for _, ps := range list {
		id := strings.TrimSpace(ps.PlatformID)
		if id == "" || seen[id] {
			continue
		}
		seen[id] = true
		driven := 0
		if ps.MailDriven {
			driven = 1
		}
		if _, err := tx.ExecContext(ctx, `
			INSERT INTO agent_platform_sessions
			    (agent_name, platform_id, workspace, slug, title, mail_driven, updated_at, reported_at)
			VALUES ($1, $2, $3, $4, $5, $6, $7, NOW())
		`, agentName, id, strings.TrimSpace(ps.Workspace), strings.TrimSpace(ps.Slug),
			strings.TrimSpace(ps.Title), driven, ps.UpdatedAt); err != nil {
			return err
		}
	}
	return tx.Commit()
}

// SessionCandidate 是「续谈某条会话」的一个候选项。
type SessionCandidate struct {
	// Alias 是填进 session 位的值 —— 候选项的实际用途就是它
	Alias string `json:"alias"`
	// Title 给人看，用来分辨两条别名相似的会话在谈什么
	Title string `json:"title,omitempty"`
	// Source 说明这条候选从哪来：
	//   mail     本侧邮件线索（可直接送达）
	//   platform 平台侧会话镜像（本侧还没有对应线索）
	Source string `json:"source"`
	// Unread 仅 mail 来源有意义
	Unread int `json:"unread,omitempty"`
}

// SuggestSessionCandidates 汇总某 name@path 下可续谈的会话。
//
// 两个来源合并：
//  1. 本侧邮件线索（sessions.workspace 匹配，或历史数据里靠 mails 反推）
//  2. 平台会话镜像里带 slug 的那些
//
// 本侧优先：邮件线索是「这个别名一定送得到」的保证，而镜像只是平台的说法。
// 同名时保留本侧那条，并把镜像的标题补上去（镜像通常有更新的标题）。
func SuggestSessionCandidates(ctx context.Context, forUser, peerName, path string) ([]SessionCandidate, error) {
	out := []SessionCandidate{}
	seen := map[string]int{} // alias -> out 下标

	// ---- 来源 1：本侧邮件线索 ----
	//
	// sessions.workspace 是权威来源。它是新加的列，历史会话为空串，
	// 因此保留 mails 反推作为兜底：`s.workspace = $2 OR (s.workspace = '' AND <mails 反推>)`。
	// 反推只看 to_workspace —— Agent 回信时 from_workspace 存的是 Agent 名而非路径，
	// 拿它比路径永远匹配不上。
	rows, err := db.DB.QueryContext(ctx, `
		SELECT s.session_alias,
		       COALESCE(s.subject, ''),
		       COALESCE(s.platform_id, ''),
		       (SELECT COUNT(*) FROM mails u
		         WHERE u.session_id = s.session_id AND u.status = 'unread')
		FROM sessions s
		WHERE s.session_alias IS NOT NULL AND s.session_alias <> ''
		  AND s.status <> 'archived'
		  AND EXISTS (
		      SELECT 1 FROM mails m
		      WHERE m.session_id = s.session_id
		        AND (m.to_name = $1 OR m.from_name = $1 OR `+db.CCHas("m.cc_list", 1)+`)
		  )
		  AND ($2 = ''
		       OR s.workspace = $2
		       OR (s.workspace = '' AND EXISTS (
		           SELECT 1 FROM mails w
		           WHERE w.session_id = s.session_id
		             AND COALESCE(w.to_workspace,'') = $2
		       )))
		  AND ($3 = '' OR s.owner_user_id = (SELECT user_id FROM users WHERE username = $3)
		       OR EXISTS (
		           SELECT 1 FROM mails mm
		           WHERE mm.session_id = s.session_id
		             AND (mm.from_name = $3 OR mm.to_name = $3
		                  OR `+db.CCHas("mm.cc_list", 3)+`)
		       ))
		ORDER BY s.updated_at DESC
	`, peerName, path, forUser)
	if err != nil {
		return out, err
	}
	defer rows.Close()

	for rows.Next() {
		var alias, title, pid string
		var unread int
		if err := rows.Scan(&alias, &title, &pid, &unread); err != nil {
			return out, err
		}
		if alias == "" {
			continue
		}
		seen[alias] = len(out)
		out = append(out, SessionCandidate{
			Alias: alias, Title: title, Source: "mail", Unread: unread,
		})
		if pid != "" {
			seen["pid:"+pid] = len(out) - 1
		}
	}
	if err := rows.Err(); err != nil {
		return out, err
	}

	// ---- 来源 2：平台会话镜像 ----
	prows, err := db.DB.QueryContext(ctx, `
		SELECT slug, title, platform_id
		FROM agent_platform_sessions
		WHERE agent_name = $1
		  AND slug <> ''
		  AND ($2 = '' OR workspace = $2)
		-- 不用 NULLS LAST：它要 SQLite 3.30+，而驱动自带的版本不由我们控制。
		-- COALESCE 在两个方言里都成立，语义也更直接：没有 updated_at 就用上报时间。
		ORDER BY COALESCE(updated_at, reported_at) DESC
	`, peerName, path)
	if err != nil {
		// 镜像查不到不该让整个补全失败：本侧线索已经够用了
		return out, nil
	}
	defer prows.Close()

	for prows.Next() {
		var slug, title, pid string
		if err := prows.Scan(&slug, &title, &pid); err != nil {
			break
		}
		if slug == "" {
			continue
		}
		// 已被接管的平台会话不再单独列：选它也会落进已有的那条本侧线索，
		// 但候选列表出现两次会让人以为有两条不同的会话（项目定位 x2 的场景）。
		if pid != "" {
			if _, dup := seen["pid:"+pid]; dup {
				continue
			}
		}
		if i, ok := seen[slug]; ok {
			// 本侧已有同名线索：保留 mail 来源（它保证送得到），
			// 但补上镜像的标题 —— 平台侧标题通常比会话建立时的主题更贴切
			if out[i].Title == "" && title != "" {
				out[i].Title = title
			}
			continue
		}
		seen[slug] = len(out)
		out = append(out, SessionCandidate{Alias: slug, Title: title, Source: "platform"})
	}
	return out, nil
}

// SetSessionWorkspace 记下会话所属的工作目录。
//
// 只在为空时写入：会话的工作区在建立时就定下了，之后不该被一封发往
// 别处的邮件改掉 —— 那会让这条会话在候选列表里凭空换一个工作区。
func SetSessionWorkspace(ctx context.Context, sessionID interface{ String() string }, workspace string) error {
	ws := strings.TrimSpace(workspace)
	if ws == "" {
		return nil
	}
	_, err := db.DB.ExecContext(ctx,
		`UPDATE sessions SET workspace = $1 WHERE session_id = $2 AND workspace = ''`,
		ws, sessionID.String())
	return err
}

// ---------- 接管平台会话 ----------
//
// TUI 与邮箱是同一个 Agent 的**两个入口**，不是两套隔离的世界。
// 人在平台界面上开的会话，应该也能被邮件投进去 —— 补全早就把它们列为候选，
// 缺的只是投递侧这一跳。
//
// 「接管」= 在本侧建一条会话并把 platform_id 记上。之后：
//   - 这条会话在 sessions 表里有正式身份（可寻址、有预算、能归档）
//   - 插件收到投递事件时看到 platform_id，就去 resume 那条平台会话
//     而不是新建一条
//
// 一条平台会话只能被接管一次：第二次投递复用第一次建的本侧会话，
// 否则同一条 TUI 对话会在邮箱里裂成多条互不相干的线索。

// FindPlatformSession 按 (agent, slug, workspace) 找一条平台会话镜像。
//
// workspace 为空表示不限（地址省略 path 位时）。返回 platform_id 与它的
// 真实 workspace —— 后者是权威的：**会话的 cwd 在它创建时就定了**，
// 地址里的 path 位若与之不同，以会话为准。人是从候选列表里选的，
// 他要的是「那条会话」而不是「那个目录」。
func FindPlatformSession(ctx context.Context, agentName, slug, workspace string) (platformID, realWorkspace, title string, err error) {
	agentName = strings.TrimSpace(agentName)
	slug = strings.TrimSpace(slug)
	if agentName == "" || slug == "" {
		return "", "", "", ErrSessionNotFound
	}
	ws := strings.TrimSpace(workspace)
	err = db.DB.QueryRowContext(ctx, `
		SELECT platform_id, workspace, title
		FROM agent_platform_sessions
		WHERE agent_name = $1 AND slug = $2
		  AND ($3 = '' OR workspace = $3)
		ORDER BY COALESCE(updated_at, reported_at) DESC
		LIMIT 1
	`, agentName, slug, ws).Scan(&platformID, &realWorkspace, &title)
	if errors.Is(err, sql.ErrNoRows) {
		return "", "", "", ErrSessionNotFound
	}
	return platformID, realWorkspace, title, err
}

// FindSessionByPlatformID 找出已经接管了某条平台会话的本侧会话。
//
// 返回 ErrSessionNotFound 表示还没被接管。归档的也算 —— 让归档过的会话
// 重新被接管会造出第二条本侧会话，同一条 TUI 对话在邮箱里就裂成两截。
// 需要恢复的话人应该去取消归档。
func FindSessionByPlatformID(ctx context.Context, agentName, platformID string) (uuid.UUID, error) {
	var id uuid.UUID
	err := db.DB.QueryRowContext(ctx, `
		SELECT s.session_id
		FROM sessions s
		WHERE s.platform_id = $1
		  AND EXISTS (
		      SELECT 1 FROM mails m
		      WHERE m.session_id = s.session_id
		        AND (m.to_name = $2 OR m.from_name = $2 OR `+db.CCHas("m.cc_list", 2)+`)
		  )
		ORDER BY s.updated_at DESC
		LIMIT 1
	`, platformID, agentName).Scan(&id)
	if errors.Is(err, sql.ErrNoRows) {
		return uuid.Nil, ErrSessionNotFound
	}
	return id, err
}

// AdoptPlatformSession 接管一条平台会话：建本侧会话并绑定 platform_id。
//
// alias 用平台自己的 slug —— 「别名复用平台命名」是既定决策，而且人在补全里
// 看到的就是那个 slug，投递后别名换成别的会让他找不到自己刚发的信。
//
// workspace 用平台会话的真实 cwd 而不是地址里的 path 位，理由见
// FindPlatformSession 的注释。
func AdoptPlatformSession(ctx context.Context, agentName, platformID, slug, workspace, subject string) (uuid.UUID, error) {
	// slug 可能与本侧某条无关会话撞名（别名全局唯一）。撞了就加后缀 ——
	// EnsureSessionAlias 已有这套逻辑，这里先建后命名即可。
	id, err := CreateSession(ctx, nil, agentName, subject, workspace)
	if err != nil {
		return uuid.Nil, err
	}
	if _, err := db.DB.ExecContext(ctx,
		`UPDATE sessions SET platform_id = $1 WHERE session_id = $2`,
		platformID, id); err != nil {
		return uuid.Nil, err
	}
	// 别名尽量用 slug；撞名时 EnsureSessionAlias 自动加后缀
	_, _ = EnsureSessionAlias(ctx, id, slug)
	return id, nil
}

// PlatformIDOf 读一条本侧会话绑定的平台会话 id（空 = 不是接管来的）。
//
// 投递时要把它放进 SSE 事件：插件据此决定 resume 还是新建。
//
// 只在「已知这条会话只有一个参与方」时用它。有抄送时必须用
// PlatformSessionFor 拿到归属方 —— 理由见那个函数。
func PlatformIDOf(ctx context.Context, sessionID uuid.UUID) string {
	pid, _ := PlatformSessionFor(ctx, sessionID)
	return pid
}

// PlatformSessionFor 返回一条本侧会话绑定的平台会话 id **及其归属 Agent**。
//
// # 为什么归属方是必须的
//
// `platform_id` 是**会话级**的一个值，而一封邮件可以有多个参与方。
// 把它无差别推给所有人，收到的一方会拿它去自己的磁盘上找会话文件 ——
// 那个 id 属于别的平台。
//
// 生产实测：会话 `16845133` 接管了 pi 的会话 `01a05a5e-…`，而那封邮件抄送了
// `dsh@/home/program/agentmail.new`。DSH 收到同一个 platform_session_id，
// 在 `~/.dsh/sessions/` 里查不到（那是 `/root/.pi/agent/sessions/` 下的文件），
// 于是走进「平台侧会话已删」那条防线抛错。那道防线本身是对的（N-8：
// 不能退回新建，否则人在界面上看不到这封邮件带来的对话），它拦下的却是
// 「别人的会话」—— 邮件因此静默消失，而插件侧的日志走的是不进 journalctl
// 的通道，连线索都没有。
//
// 归属方以镜像（`agent_platform_sessions.agent_name`，Agent 自己上报的）为准；
// 镜像整表替换，平台侧删了会话那行就没了，此时退回 `sessions.from_agent` ——
// `AdoptPlatformSession` 建会话时把归属 Agent 写在那里，是可靠的第二来源。
func PlatformSessionFor(ctx context.Context, sessionID uuid.UUID) (platformID, owner string) {
	var pid, fromAgent string
	var mirrored *string
	if err := db.DB.QueryRowContext(ctx, `
		SELECT COALESCE(s.platform_id, ''), COALESCE(s.from_agent, ''), aps.agent_name
		FROM sessions s
		LEFT JOIN agent_platform_sessions aps
		       ON aps.platform_id = s.platform_id AND COALESCE(s.platform_id, '') <> ''
		WHERE s.session_id = $1`, sessionID).Scan(&pid, &fromAgent, &mirrored); err != nil {
		return "", ""
	}
	if pid == "" {
		return "", ""
	}
	if mirrored != nil && *mirrored != "" {
		return pid, *mirrored
	}
	return pid, fromAgent
}