package database import ( "database/sql" "fmt" "strings" "time" "github.com/google/uuid" ) // VulnerabilityRetest identifies the independent snapshot used for a retest. type VulnerabilityRetest struct { ConversationID string `json:"conversation_id"` SourceConversationID string `json:"source_conversation_id"` VulnerabilityID string `json:"vulnerability_id"` SourceProjectID string `json:"source_project_id"` PendingPrompt string `json:"pending_prompt,omitempty"` } func (db *DB) GetVulnerabilityRetest(conversationID string) (*VulnerabilityRetest, error) { r := &VulnerabilityRetest{ConversationID: conversationID} err := db.QueryRow(`SELECT source_conversation_id, vulnerability_id, source_project_id, pending_prompt FROM vulnerability_retests WHERE conversation_id = ?`, conversationID).Scan(&r.SourceConversationID, &r.VulnerabilityID, &r.SourceProjectID, &r.PendingPrompt) if err == sql.ErrNoRows { return nil, nil } return r, err } // ResolveConversationPlaceholders replaces identifiers throughout the input, // including follow-up turns. Unknown placeholders remain untouched. func (db *DB) ResolveConversationPlaceholders(conversationID, input string) (string, error) { r, err := db.GetVulnerabilityRetest(conversationID) if err != nil { return "", err } replacements := []string{"{{conversation_id}}", conversationID} if r != nil { replacements = append(replacements, "{{source_conversation_id}}", r.SourceConversationID, "{{vulnerability_id}}", r.VulnerabilityID) } return strings.NewReplacer(replacements...).Replace(input), nil } // copiedTime re-binds a scanned timestamp as time.Time so copied rows use the same created_at format AddMessage writes, which the raw string ORDER BY depends on. func copiedTime(value sql.NullString) interface{} { if !value.Valid || strings.TrimSpace(value.String) == "" { return nil } for _, layout := range []string{time.RFC3339Nano, time.RFC3339, "2006-01-02 15:04:05.999999999-07:00", "2006-01-02 15:04:05"} { if parsed, err := time.Parse(layout, value.String); err == nil { return parsed } } return value.String } // CreateVulnerabilityRetest snapshots history and model trace in one transaction. // Each call creates a sibling; nothing is written to the source conversation. func (db *DB) CreateVulnerabilityRetest(v *Vulnerability, meta ConversationCreateMeta) (*Conversation, error) { tx, err := db.Begin() if err != nil { return nil, err } defer tx.Rollback() id, now := uuid.NewString(), time.Now() title := "复测:" + v.Title result, err := tx.Exec(`INSERT INTO conversations (id, title, project_id, role_name, agent_mode, last_react_input, last_react_output, created_at, updated_at) SELECT ?, ?, project_id, role_name, agent_mode, last_react_input, last_react_output, ?, ? FROM conversations WHERE id = ?`, id, title, now, now, v.ConversationID) if err != nil { return nil, err } n, err := result.RowsAffected() if err != nil || n != 1 { return nil, fmt.Errorf("原始对话不存在") } conv := &Conversation{ID: id, Title: title, CreatedAt: now, UpdatedAt: now} if err = tx.QueryRow(`SELECT COALESCE(project_id,''), role_name, agent_mode FROM conversations WHERE id=?`, id).Scan(&conv.ProjectID, &conv.RoleName, &conv.AgentMode); err != nil { return nil, err } // Copy visible messages using new IDs, preserving chronological order. Execution // records belong to the source and must not be shared with a mutable branch. rows, err := tx.Query(`SELECT role, content, reasoning_content, created_at, updated_at FROM messages WHERE conversation_id=? ORDER BY created_at, rowid`, v.ConversationID) if err != nil { return nil, err } type snapshotMessage struct { role, content string reasoning, created, updated sql.NullString } var messages []snapshotMessage for rows.Next() { var m snapshotMessage if err = rows.Scan(&m.role, &m.content, &m.reasoning, &m.created, &m.updated); err != nil { rows.Close() return nil, err } messages = append(messages, m) } err = rows.Err() rows.Close() if err != nil { return nil, err } for _, m := range messages { if _, err = tx.Exec(`INSERT INTO messages (id, conversation_id, role, content, reasoning_content, created_at, updated_at) VALUES (?,?,?,?,?,?,?)`, uuid.NewString(), id, m.role, m.content, m.reasoning, copiedTime(m.created), copiedTime(m.updated)); err != nil { return nil, err } } if _, err = tx.Exec(`INSERT INTO vulnerability_retests (conversation_id, source_conversation_id, vulnerability_id, source_project_id) VALUES (?,?,?,?)`, id, v.ConversationID, v.ID, conv.ProjectID); err != nil { return nil, err } // Preserve the source's approval policy without copying live interrupts. // HITL schema is installed by the agent handler rather than by NewDB. var hasHITL int if err = tx.QueryRow(`SELECT COUNT(*) FROM sqlite_master WHERE type='table' AND name='hitl_conversation_configs'`).Scan(&hasHITL); err != nil { return nil, err } if hasHITL != 0 { if _, err = tx.Exec(`INSERT INTO hitl_conversation_configs (conversation_id, enabled, mode, reviewer, sensitive_tools, timeout_seconds, updated_at) SELECT ?, enabled, mode, reviewer, sensitive_tools, timeout_seconds, ? FROM hitl_conversation_configs WHERE conversation_id=?`, id, now, v.ConversationID); err != nil { return nil, err } } if err = tx.Commit(); err != nil { return nil, err } meta.ProjectID, meta.RoleName, meta.AgentMode = conv.ProjectID, conv.RoleName, conv.AgentMode notifyConversationCreated(conv, meta) return conv, nil } // UpdateVulnerabilityRetestResult appends evidence to retest_log; retest_notes stays operator-owned. func (db *DB) UpdateVulnerabilityRetestResult(id, status, notes, conversationID string) error { entry := fmt.Sprintf("\n\n[复测 %s | 对话 %s | 状态 %s]\n%s", time.Now().Format(time.RFC3339), conversationID, status, notes) result, err := db.Exec(`UPDATE vulnerabilities SET status=?, retest_log=COALESCE(retest_log,'') || ?, updated_at=? WHERE id=?`, status, entry, time.Now(), id) if err != nil { return err } n, err := result.RowsAffected() if err != nil { return err } if n != 1 { return fmt.Errorf("漏洞不存在") } if v, err := db.GetVulnerability(id); err == nil { db.refreshAssetRiskCacheForConversationsBestEffort(v.ConversationID) } return nil } // A pending prompt survives a lost response, navigation or browser refresh. func (db *DB) SetVulnerabilityRetestPrompt(conversationID, prompt string) error { _, err := db.Exec(`UPDATE vulnerability_retests SET pending_prompt=? WHERE conversation_id=?`, prompt, conversationID) return err }