From 3bedbdc66966db302e8b3c70cfa58c4fc435ca34 Mon Sep 17 00:00:00 2001 From: pphua Date: Mon, 24 Aug 2026 19:29:38 +0800 Subject: [PATCH] feat(chat): add conversation pinning and rename controls --- db/conversation.go | 109 ++++++++++-- db/conversation_pin_test.go | 69 ++++++++ server/conversations.go | 77 +++++++-- web/src/app/(main)/chat/page.tsx | 280 +++++++++++++++++++++---------- 4 files changed, 418 insertions(+), 117 deletions(-) create mode 100644 db/conversation_pin_test.go diff --git a/db/conversation.go b/db/conversation.go index 9e7e97a..5a7edc5 100644 --- a/db/conversation.go +++ b/db/conversation.go @@ -50,43 +50,98 @@ GROUP BY c.id`) // Conversation is one ChatGPT-style chat thread bound to an agent key. It lives // independent of the pentest exploration graph — see schema.sql §I. type Conversation struct { - ID int64 `json:"id"` - AgentKey string `json:"agent_key"` - Title string `json:"title"` - LLMProfileID *int64 `json:"llm_profile_id,omitempty"` - CreatedAt time.Time `json:"created_at"` - UpdatedAt time.Time `json:"updated_at"` + ID int64 `json:"id"` + AgentKey string `json:"agent_key"` + Title string `json:"title"` + LLMProfileID *int64 `json:"llm_profile_id,omitempty"` + Pinned bool `json:"pinned"` + PinnedAt *time.Time `json:"pinned_at,omitempty"` + CreatedAt time.Time `json:"created_at"` + UpdatedAt time.Time `json:"updated_at"` } -const convCols = `id, agent_key, title, llm_profile_id, created_at, updated_at` +// ConversationPatch updates only the fields whose pointers are non-nil. +type ConversationPatch struct { + Title *string + Pinned *bool +} + +const convCols = `id, agent_key, title, llm_profile_id, pinned_at, created_at, updated_at` func scanConv(row interface{ Scan(...any) error }) (Conversation, error) { var c Conversation - err := row.Scan(&c.ID, &c.AgentKey, &c.Title, &c.LLMProfileID, &c.CreatedAt, &c.UpdatedAt) + var pinnedAt sql.NullTime + err := row.Scan(&c.ID, &c.AgentKey, &c.Title, &c.LLMProfileID, &pinnedAt, &c.CreatedAt, &c.UpdatedAt) + if pinnedAt.Valid { + c.Pinned = true + c.PinnedAt = &pinnedAt.Time + } return c, err } // CreateConversation opens a new chat thread for agentKey with an initial title. // llmProfileID may be nil to use the globally active profile. func (d *DB) CreateConversation(agentKey, title string, llmProfileID *int64) (*Conversation, error) { - c, err := scanConv(d.QueryRow(` -INSERT INTO conversations(agent_key, title, llm_profile_id) VALUES ($1, $2, $3) -RETURNING `+convCols, agentKey, title, llmProfileID)) + tx, err := d.Begin() if err != nil { return nil, err } + defer tx.Rollback() + + // Insert the child without a profile first. The new row is exclusively owned + // by this transaction before it takes a profile lock, matching DeleteProfile's + // child-row -> profile-row protocol. + c, err := scanConv(tx.QueryRow(` +INSERT INTO conversations(agent_key, title, llm_profile_id) VALUES ($1, $2, NULL) +RETURNING `+convCols, agentKey, title)) + if err != nil { + return nil, err + } + if err := lockLLMProfileForReference(tx, llmProfileID); err != nil { + return nil, err + } + if llmProfileID != nil { + c, err = scanConv(tx.QueryRow(`UPDATE conversations SET llm_profile_id=$2 +WHERE id=$1 RETURNING `+convCols, c.ID, llmProfileID)) + if err != nil { + return nil, err + } + } + if err := tx.Commit(); err != nil { + return nil, err + } return &c, nil } // UpdateConversationProfile sets (or clears) the LLM profile override for a conversation. func (d *DB) UpdateConversationProfile(id int64, llmProfileID *int64) error { - _, err := d.Exec(`UPDATE conversations SET llm_profile_id=$2 WHERE id=$1`, id, llmProfileID) - return err + tx, err := d.Begin() + if err != nil { + return err + } + defer tx.Rollback() + + var lockedID int64 + if err := tx.QueryRow(`SELECT id FROM conversations WHERE id=$1 FOR UPDATE`, id).Scan(&lockedID); err != nil { + if err == sql.ErrNoRows { + // Preserve the previous UPDATE semantics: an unknown id is a no-op. + return tx.Commit() + } + return err + } + if err := lockLLMProfileForReference(tx, llmProfileID); err != nil { + return err + } + if _, err := tx.Exec(`UPDATE conversations SET llm_profile_id=$2 WHERE id=$1`, id, llmProfileID); err != nil { + return err + } + return tx.Commit() } // ListConversations returns all threads, most-recently-updated first. func (d *DB) ListConversations() ([]*Conversation, error) { - rows, err := d.Query(`SELECT ` + convCols + ` FROM conversations ORDER BY updated_at DESC`) + rows, err := d.Query(`SELECT ` + convCols + ` FROM conversations +ORDER BY (pinned_at IS NOT NULL) DESC, pinned_at DESC NULLS LAST, updated_at DESC, id DESC`) if err != nil { return nil, err } @@ -114,9 +169,31 @@ func (d *DB) GetConversation(id int64) (*Conversation, error) { return &c, nil } -// RenameConversation sets a thread's title. +// UpdateConversation applies a partial title/pin mutation and returns the updated +// row. Pinning an already-pinned conversation preserves its original pin order. +func (d *DB) UpdateConversation(id int64, patch ConversationPatch) (*Conversation, error) { + c, err := scanConv(d.QueryRow(`UPDATE conversations SET + title = CASE WHEN $2::boolean THEN $3 ELSE title END, + pinned_at = CASE + WHEN $4::boolean IS NULL THEN pinned_at + WHEN $4::boolean THEN COALESCE(pinned_at, now()) + ELSE NULL + END +WHERE id=$1 +RETURNING `+convCols, id, patch.Title != nil, patch.Title, patch.Pinned)) + if err == sql.ErrNoRows { + return nil, nil + } + if err != nil { + return nil, err + } + return &c, nil +} + +// RenameConversation sets a thread's title. Kept for automatic first-message +// titles and compatibility with existing callers. func (d *DB) RenameConversation(id int64, title string) error { - _, err := d.Exec(`UPDATE conversations SET title=$2 WHERE id=$1`, id, title) + _, err := d.UpdateConversation(id, ConversationPatch{Title: &title}) return err } diff --git a/db/conversation_pin_test.go b/db/conversation_pin_test.go new file mode 100644 index 0000000..da7a106 --- /dev/null +++ b/db/conversation_pin_test.go @@ -0,0 +1,69 @@ +package db + +import ( + "fmt" + "testing" + "time" +) + +func TestConversationPinOrderingAndPatch(t *testing.T) { + d, err := Open(testDSN(t)) + if err != nil { + t.Skipf("postgres unavailable (%v) — skipping", err) + } + defer d.Close() + + suffix := time.Now().UnixNano() + first, err := d.CreateConversation("mainagent", fmt.Sprintf("pin-first-%d", suffix), nil) + if err != nil { + t.Fatal(err) + } + second, err := d.CreateConversation("mainagent", fmt.Sprintf("pin-second-%d", suffix), nil) + if err != nil { + _ = d.DeleteConversation(first.ID) + t.Fatal(err) + } + t.Cleanup(func() { + _ = d.DeleteConversation(first.ID) + _ = d.DeleteConversation(second.ID) + }) + + pinned := true + first, err = d.UpdateConversation(first.ID, ConversationPatch{Pinned: &pinned}) + if err != nil || first == nil || !first.Pinned || first.PinnedAt == nil { + t.Fatalf("pin first = %+v, %v", first, err) + } + firstPinnedAt := *first.PinnedAt + time.Sleep(5 * time.Millisecond) + second, err = d.UpdateConversation(second.ID, ConversationPatch{Pinned: &pinned}) + if err != nil || second == nil || !second.Pinned { + t.Fatalf("pin second = %+v, %v", second, err) + } + + newTitle := "renamed while pinned" + first, err = d.UpdateConversation(first.ID, ConversationPatch{Title: &newTitle, Pinned: &pinned}) + if err != nil || first == nil { + t.Fatal(err) + } + if first.Title != newTitle || first.PinnedAt == nil || !first.PinnedAt.Equal(firstPinnedAt) { + t.Fatalf("repeat pin should preserve pin time: %+v (want %v)", first, firstPinnedAt) + } + + conversations, err := d.ListConversations() + if err != nil { + t.Fatal(err) + } + positions := map[int64]int{} + for index, conversation := range conversations { + positions[conversation.ID] = index + } + if positions[second.ID] >= positions[first.ID] { + t.Fatalf("newer pin must sort first: second=%d first=%d", positions[second.ID], positions[first.ID]) + } + + pinned = false + first, err = d.UpdateConversation(first.ID, ConversationPatch{Pinned: &pinned}) + if err != nil || first == nil || first.Pinned || first.PinnedAt != nil { + t.Fatalf("unpin first = %+v, %v", first, err) + } +} diff --git a/server/conversations.go b/server/conversations.go index 7b0c77e..de81e4c 100644 --- a/server/conversations.go +++ b/server/conversations.go @@ -2,6 +2,7 @@ package server import ( "context" + "errors" "fmt" "log" "math" @@ -10,11 +11,33 @@ import ( "strconv" "strings" "time" + "unicode/utf8" + "github.com/Autumn-27/artex/agent" "github.com/Autumn-27/artex/db" "github.com/Autumn-27/artex/intercept" ) +const ( + maxConversationRequestBytes = 64 << 10 + maxConversationAgentKeyRunes = 120 + maxConversationTitleRunes = 200 +) + +func decodeConversationRequest(w http.ResponseWriter, r *http.Request, value any) bool { + r.Body = http.MaxBytesReader(w, r.Body, maxConversationRequestBytes) + if err := decode(r, value); err != nil { + var tooLarge *http.MaxBytesError + if errors.As(err, &tooLarge) { + writeErr(w, http.StatusRequestEntityTooLarge, "请求正文过大") + } else { + writeErr(w, http.StatusBadRequest, err.Error()) + } + return false + } + return true +} + // ---------- conversations (chat page) ---------- // // A conversation is a ChatGPT-style thread bound to an agent key, independent of @@ -45,8 +68,7 @@ func (s *Server) pgCreateConversation(w http.ResponseWriter, r *http.Request) { Title string `json:"title"` LLMProfileID *int64 `json:"llm_profile_id"` } - if err := decode(r, &req); err != nil { - writeErr(w, 400, err.Error()) + if !decodeConversationRequest(w, r, &req) { return } req.AgentKey = strings.TrimSpace(req.AgentKey) @@ -54,6 +76,10 @@ func (s *Server) pgCreateConversation(w http.ResponseWriter, r *http.Request) { writeErr(w, 400, "agent_key 不能为空") return } + if utf8.RuneCountInString(req.AgentKey) > maxConversationAgentKeyRunes { + writeErr(w, 400, fmt.Sprintf("agent_key 最多 %d 个字符", maxConversationAgentKeyRunes)) + return + } a, err := pg.GetAgentByKey(req.AgentKey) if err != nil { writeErr(w, 500, err.Error()) @@ -73,6 +99,10 @@ func (s *Server) pgCreateConversation(w http.ResponseWriter, r *http.Request) { if title == "" { title = "新对话" } + if utf8.RuneCountInString(title) > maxConversationTitleRunes { + writeErr(w, 400, fmt.Sprintf("标题最多 %d 个字符", maxConversationTitleRunes)) + return + } c, err := pg.CreateConversation(req.AgentKey, title, req.LLMProfileID) if err != nil { writeErr(w, 500, err.Error()) @@ -89,8 +119,7 @@ func (s *Server) pgUpdateConversation(w http.ResponseWriter, r *http.Request) { var req struct { LLMProfileID *int64 `json:"llm_profile_id"` // null clears the override } - if err := decode(r, &req); err != nil { - writeErr(w, 400, err.Error()) + if !decodeConversationRequest(w, r, &req) { return } if req.LLMProfileID != nil { @@ -134,21 +163,39 @@ func (s *Server) pgRenameConversation(w http.ResponseWriter, r *http.Request) { if !ok { return } - var req struct{ Title string } - if err := decode(r, &req); err != nil { - writeErr(w, 400, err.Error()) + var req struct { + Title *string `json:"title"` + Pinned *bool `json:"pinned"` + } + if !decodeConversationRequest(w, r, &req) { return } - title := strings.TrimSpace(req.Title) - if title == "" { - writeErr(w, 400, "标题不能为空") + if req.Title == nil && req.Pinned == nil { + writeErr(w, 400, "至少需要提供 title 或 pinned") return } - if err := pg.RenameConversation(c.ID, title); err != nil { + if req.Title != nil { + title := strings.TrimSpace(*req.Title) + if title == "" { + writeErr(w, 400, "标题不能为空") + return + } + if utf8.RuneCountInString(title) > maxConversationTitleRunes { + writeErr(w, 400, fmt.Sprintf("标题最多 %d 个字符", maxConversationTitleRunes)) + return + } + req.Title = &title + } + updated, err := pg.UpdateConversation(c.ID, db.ConversationPatch{Title: req.Title, Pinned: req.Pinned}) + if err != nil { writeErr(w, 500, err.Error()) return } - writeJSON(w, 200, map[string]any{"ok": true}) + if updated == nil { + writeErr(w, 404, "conversation not found") + return + } + writeJSON(w, 200, updated) } func (s *Server) pgDeleteConversation(w http.ResponseWriter, r *http.Request) { @@ -243,7 +290,7 @@ func (s *Server) pgStopConversation(w http.ResponseWriter, r *http.Request) { writeJSON(w, 200, map[string]any{"status": "idle"}) // nothing running return } - cancel() + cancel(agent.AbortChatStoppedByUser) writeJSON(w, 200, map[string]any{"status": "stopping"}) } @@ -331,12 +378,12 @@ func (s *Server) runConversation(c *db.Conversation, msg, busyKey string) { func (s *Server) runConversationSync(c *db.Conversation, msg, busyKey string) { // Per-run cancellable context so a manual stop (pgStopConversation) can abort // just this session. Registered under chatMu so the stop handler can find it. - ctx, cancel := context.WithCancel(intercept.WithConvID(s.ctx, c.ID)) + ctx, cancel := context.WithCancelCause(intercept.WithConvID(s.ctx, c.ID)) s.chatMu.Lock() s.chatCancel[busyKey] = cancel s.chatMu.Unlock() defer func() { - cancel() + cancel(agent.AbortChatTurnFinished) s.chatMu.Lock() delete(s.chatBusy, busyKey) delete(s.chatCancel, busyKey) diff --git a/web/src/app/(main)/chat/page.tsx b/web/src/app/(main)/chat/page.tsx index e72c31b..ded6bce 100644 --- a/web/src/app/(main)/chat/page.tsx +++ b/web/src/app/(main)/chat/page.tsx @@ -1,20 +1,27 @@ "use client"; import * as React from "react"; -import { toast } from "sonner"; -import { Bot, ChevronDownIcon, Loader2Icon, PaperclipIcon, PlusIcon, SendIcon, Square, Trash2Icon, XIcon, ZapIcon } from "lucide-react"; -import { Button } from "@/components/ui/button"; -import { Textarea } from "@/components/ui/textarea"; -import { Badge } from "@/components/ui/badge"; -import { ScrollArea } from "@/components/ui/scroll-area"; import { - Select, - SelectContent, - SelectItem, - SelectTrigger, - SelectValue, -} from "@/components/ui/select"; + ArrowUpIcon, + Bot, + ChevronDownIcon, + Loader2Icon, + MoreHorizontalIcon, + PaperclipIcon, + PencilIcon, + PinIcon, + PinOffIcon, + PlusIcon, + Square, + Trash2Icon, + XIcon, + ZapIcon, +} from "lucide-react"; +import { toast } from "sonner"; + +import { TodoPopover } from "@/components/todo-popover"; +import { Transcript } from "@/components/transcript"; import { AlertDialog, AlertDialogAction, @@ -24,15 +31,21 @@ import { AlertDialogFooter, AlertDialogHeader, AlertDialogTitle, - AlertDialogTrigger, } from "@/components/ui/alert-dialog"; -import { Transcript } from "@/components/transcript"; -import { TodoPopover } from "@/components/todo-popover"; +import { Badge } from "@/components/ui/badge"; +import { Button } from "@/components/ui/button"; import { - Popover, - PopoverContent, - PopoverTrigger, -} from "@/components/ui/popover"; + DropdownMenu, + DropdownMenuContent, + DropdownMenuGroup, + DropdownMenuItem, + DropdownMenuSeparator, + DropdownMenuTrigger, +} from "@/components/ui/dropdown-menu"; +import { Popover, PopoverContent, PopoverTrigger } from "@/components/ui/popover"; +import { ScrollArea } from "@/components/ui/scroll-area"; +import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; +import { Textarea } from "@/components/ui/textarea"; import { api } from "@/lib/api"; import { shouldSubmitOnKey, useChatSendMode } from "@/lib/chat-send-mode"; import type { Activity, Agent, ChatAttachment, Conversation, LLMProfile } from "@/lib/types"; @@ -47,8 +60,8 @@ function fmtBytes(n: number): string { // fmtTokens renders a compact token count (1234 → 1.2k, 2_000_000 → 2M). function fmtTokens(n: number): string { - if (n >= 1_000_000) return (n / 1_000_000).toFixed(n >= 10_000_000 ? 0 : 1) + "M"; - if (n >= 1000) return (n / 1000).toFixed(n >= 10000 ? 0 : 1) + "k"; + if (n >= 1_000_000) return `${(n / 1_000_000).toFixed(n >= 10_000_000 ? 0 : 1)}M`; + if (n >= 1000) return `${(n / 1000).toFixed(n >= 10000 ? 0 : 1)}k`; return String(n); } @@ -67,6 +80,10 @@ function fmtDuration(ms: number): string { // thread stays snappy (only ~a page of rows is in the DOM until you scroll up). const HISTORY_PAGE = 200; +function conversationIsPinned(conversation: Conversation): boolean { + return conversation.pinned ?? Boolean(conversation.pinned_at); +} + // LiveBadge is the small pulsing "实时" chip reused from the task's main-agent // console — shown while a turn is streaming. function LiveBadge() { @@ -115,16 +132,9 @@ function Composer({ // 发送键位由系统设置决定（localStorage），默认 Enter 发送。 const sendMode = useChatSendMode(); function onKeyDown(e: React.KeyboardEvent) { -<<<<<<< Updated upstream - if (e.key === "Enter" && !e.shiftKey) { - e.preventDefault(); - onSend(); - } -======= if (!shouldSubmitOnKey(e, sendMode)) return; e.preventDefault(); onSend(); ->>>>>>> Stashed changes } return (
@@ -153,8 +163,8 @@ function Composer({ ))}
)} -
- {leftSlot} +
+ {leftSlot ?
{leftSlot}
: null} {onPickFiles && ( <> )}