diff --git a/internal/database/database.go b/internal/database/database.go index 8de22a54..4a354294 100644 --- a/internal/database/database.go +++ b/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) diff --git a/internal/database/robot_session.go b/internal/database/robot_session.go new file mode 100644 index 00000000..b7631260 --- /dev/null +++ b/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 +}