Add files via upload

internal/database/database.go

@@ -203,6 +203,16 @@ func (db *DB) initTables() error {
 		UNIQUE(conversation_id, group_id)
 	);`
 
+	// 机器人会话绑定表（用于跨重启保持「平台+租户+用户」到 conversation 的映射）
+	createRobotUserSessionsTable := `
+	CREATE TABLE IF NOT EXISTS robot_user_sessions (
+		session_key TEXT PRIMARY KEY,
+		conversation_id TEXT NOT NULL,
+		role_name TEXT NOT NULL DEFAULT '默认',
+		updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
+		FOREIGN KEY (conversation_id) REFERENCES conversations(id) ON DELETE CASCADE
+	);`
+
 	// 创建漏洞表
 	createVulnerabilitiesTable := `
 	CREATE TABLE IF NOT EXISTS vulnerabilities (
@@ -409,6 +419,7 @@ func (db *DB) initTables() error {
 	CREATE INDEX IF NOT EXISTS idx_knowledge_retrieval_logs_created_at ON knowledge_retrieval_logs(created_at);
 	CREATE INDEX IF NOT EXISTS idx_conversation_group_mappings_conversation ON conversation_group_mappings(conversation_id);
 	CREATE INDEX IF NOT EXISTS idx_conversation_group_mappings_group ON conversation_group_mappings(group_id);
+	CREATE INDEX IF NOT EXISTS idx_robot_user_sessions_updated_at ON robot_user_sessions(updated_at);
 	CREATE INDEX IF NOT EXISTS idx_conversations_pinned ON conversations(pinned);
 	CREATE INDEX IF NOT EXISTS idx_vulnerabilities_conversation_id ON vulnerabilities(conversation_id);
 	CREATE INDEX IF NOT EXISTS idx_vulnerabilities_conversation_tag ON vulnerabilities(conversation_tag);
@@ -479,6 +490,9 @@ func (db *DB) initTables() error {
 	if _, err := db.Exec(createConversationGroupMappingsTable); err != nil {
 		return fmt.Errorf("创建conversation_group_mappings表失败: %w", err)
 	}
+	if _, err := db.Exec(createRobotUserSessionsTable); err != nil {
+		return fmt.Errorf("创建robot_user_sessions表失败: %w", err)
+	}
 
 	if _, err := db.Exec(createVulnerabilitiesTable); err != nil {
 		return fmt.Errorf("创建vulnerabilities表失败: %w", err)

internal/database/robot_session.go

@@ -0,0 +1,84 @@
+package database
+
+import (
+	"database/sql"
+	"fmt"
+	"strings"
+	"time"
+)
+
+// RobotSessionBinding 机器人会话绑定信息。
+type RobotSessionBinding struct {
+	SessionKey     string
+	ConversationID string
+	RoleName       string
+	UpdatedAt      time.Time
+}
+
+// GetRobotSessionBinding 按 session_key 获取机器人会话绑定。
+func (db *DB) GetRobotSessionBinding(sessionKey string) (*RobotSessionBinding, error) {
+	sessionKey = strings.TrimSpace(sessionKey)
+	if sessionKey == "" {
+		return nil, nil
+	}
+	var b RobotSessionBinding
+	var updatedAt string
+	err := db.QueryRow(
+		"SELECT session_key, conversation_id, role_name, updated_at FROM robot_user_sessions WHERE session_key = ?",
+		sessionKey,
+	).Scan(&b.SessionKey, &b.ConversationID, &b.RoleName, &updatedAt)
+	if err != nil {
+		if err == sql.ErrNoRows {
+			return nil, nil
+		}
+		return nil, fmt.Errorf("查询机器人会话绑定失败: %w", err)
+	}
+	if t, e := time.Parse("2006-01-02 15:04:05.999999999-07:00", updatedAt); e == nil {
+		b.UpdatedAt = t
+	} else if t, e := time.Parse("2006-01-02 15:04:05", updatedAt); e == nil {
+		b.UpdatedAt = t
+	} else {
+		b.UpdatedAt, _ = time.Parse(time.RFC3339, updatedAt)
+	}
+	if strings.TrimSpace(b.RoleName) == "" {
+		b.RoleName = "默认"
+	}
+	return &b, nil
+}
+
+// UpsertRobotSessionBinding 写入或更新机器人会话绑定（包含角色）。
+func (db *DB) UpsertRobotSessionBinding(sessionKey, conversationID, roleName string) error {
+	sessionKey = strings.TrimSpace(sessionKey)
+	conversationID = strings.TrimSpace(conversationID)
+	roleName = strings.TrimSpace(roleName)
+	if sessionKey == "" || conversationID == "" {
+		return nil
+	}
+	if roleName == "" {
+		roleName = "默认"
+	}
+	_, err := db.Exec(`
+		INSERT INTO robot_user_sessions (session_key, conversation_id, role_name, updated_at)
+		VALUES (?, ?, ?, ?)
+		ON CONFLICT(session_key) DO UPDATE SET
+			conversation_id = excluded.conversation_id,
+			role_name = excluded.role_name,
+			updated_at = excluded.updated_at
+	`, sessionKey, conversationID, roleName, time.Now())
+	if err != nil {
+		return fmt.Errorf("写入机器人会话绑定失败: %w", err)
+	}
+	return nil
+}
+
+// DeleteRobotSessionBinding 删除机器人会话绑定。
+func (db *DB) DeleteRobotSessionBinding(sessionKey string) error {
+	sessionKey = strings.TrimSpace(sessionKey)
+	if sessionKey == "" {
+		return nil
+	}
+	if _, err := db.Exec("DELETE FROM robot_user_sessions WHERE session_key = ?", sessionKey); err != nil {
+		return fmt.Errorf("删除机器人会话绑定失败: %w", err)
+	}
+	return nil
+}