From facb1dd08d9fe85465259441c337b52e9b2aaec7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9Cloccen=E2=80=9D?= Date: Sun, 26 Apr 2026 03:43:54 +0800 Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E5=A4=8D=E4=B8=BB=E5=8A=A8=E5=8F=91?= =?UTF-8?q?=E9=80=81=E7=BC=BA=E5=A4=B1=E4=B8=8A=E4=B8=8B=E6=96=87=E4=BB=A4?= =?UTF-8?q?=E7=89=8C?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- ilink/context_store.go | 108 ++++++++++++++++++++++++++++++++++++ ilink/context_store_test.go | 56 +++++++++++++++++++ ilink/monitor.go | 3 + messaging/handler.go | 9 ++- messaging/media.go | 7 +++ messaging/sender.go | 7 +++ 6 files changed, 187 insertions(+), 3 deletions(-) create mode 100644 ilink/context_store.go create mode 100644 ilink/context_store_test.go diff --git a/ilink/context_store.go b/ilink/context_store.go new file mode 100644 index 0000000..bc3bed0 --- /dev/null +++ b/ilink/context_store.go @@ -0,0 +1,108 @@ +package ilink + +import ( + "encoding/json" + "fmt" + "os" + "path/filepath" + "sync" +) + +var contextStoreMu sync.Mutex + +type contextTokenData struct { + Tokens map[string]string `json:"tokens"` +} + +func contextTokenPath(botID string) (string, error) { + dir, err := AccountsDir() + if err != nil { + return "", err + } + return filepath.Join(dir, NormalizeAccountID(botID)+".contexts.json"), nil +} + +// SaveContextToken stores the latest iLink context token for a user. +func SaveContextToken(botID, userID, token string) error { + if botID == "" || userID == "" || token == "" { + return nil + } + + contextStoreMu.Lock() + defer contextStoreMu.Unlock() + + path, err := contextTokenPath(botID) + if err != nil { + return err + } + + data := contextTokenData{Tokens: map[string]string{}} + if raw, err := os.ReadFile(path); err == nil { + _ = json.Unmarshal(raw, &data) + } + if data.Tokens == nil { + data.Tokens = map[string]string{} + } + data.Tokens[userID] = token + + if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil { + return fmt.Errorf("create context token dir: %w", err) + } + + raw, err := json.MarshalIndent(data, "", " ") + if err != nil { + return fmt.Errorf("marshal context tokens: %w", err) + } + if err := os.WriteFile(path, raw, 0o600); err != nil { + return fmt.Errorf("write context tokens: %w", err) + } + return nil +} + +// LoadContextToken returns the latest cached iLink context token for a user. +func LoadContextToken(botID, userID string) (string, error) { + if botID == "" || userID == "" { + return "", nil + } + + contextStoreMu.Lock() + defer contextStoreMu.Unlock() + + path, err := contextTokenPath(botID) + if err != nil { + return "", err + } + + raw, err := os.ReadFile(path) + if err != nil { + if os.IsNotExist(err) { + return "", nil + } + return "", fmt.Errorf("read context tokens: %w", err) + } + + var data contextTokenData + if err := json.Unmarshal(raw, &data); err != nil { + return "", fmt.Errorf("parse context tokens: %w", err) + } + return data.Tokens[userID], nil +} + +// ClearContextTokens removes cached iLink context tokens for a bot account. +func ClearContextTokens(botID string) error { + if botID == "" { + return nil + } + + contextStoreMu.Lock() + defer contextStoreMu.Unlock() + + path, err := contextTokenPath(botID) + if err != nil { + return err + } + if err := os.Remove(path); err != nil && !os.IsNotExist(err) { + return fmt.Errorf("remove context tokens: %w", err) + } + return nil +} diff --git a/ilink/context_store_test.go b/ilink/context_store_test.go new file mode 100644 index 0000000..6f62ed7 --- /dev/null +++ b/ilink/context_store_test.go @@ -0,0 +1,56 @@ +package ilink + +import ( + "os" + "path/filepath" + "testing" +) + +func TestContextTokenStoreRoundTrip(t *testing.T) { + home := t.TempDir() + t.Setenv("HOME", home) + + botID := "bot@example" + userID := "user@im.wechat" + token := "context-token" + + if err := SaveContextToken(botID, userID, token); err != nil { + t.Fatalf("SaveContextToken() error = %v", err) + } + + got, err := LoadContextToken(botID, userID) + if err != nil { + t.Fatalf("LoadContextToken() error = %v", err) + } + if got != token { + t.Fatalf("LoadContextToken() = %q, want %q", got, token) + } + + path := filepath.Join(home, ".weclaw", "accounts", "bot-example.contexts.json") + if _, err := os.Stat(path); err != nil { + t.Fatalf("context token file was not written: %v", err) + } +} + +func TestClearContextTokens(t *testing.T) { + home := t.TempDir() + t.Setenv("HOME", home) + + botID := "bot@example" + userID := "user@im.wechat" + + if err := SaveContextToken(botID, userID, "context-token"); err != nil { + t.Fatalf("SaveContextToken() error = %v", err) + } + if err := ClearContextTokens(botID); err != nil { + t.Fatalf("ClearContextTokens() error = %v", err) + } + + got, err := LoadContextToken(botID, userID) + if err != nil { + t.Fatalf("LoadContextToken() error = %v", err) + } + if got != "" { + t.Fatalf("LoadContextToken() = %q, want empty", got) + } +} diff --git a/ilink/monitor.go b/ilink/monitor.go index a8c360c..231fd47 100644 --- a/ilink/monitor.go +++ b/ilink/monitor.go @@ -93,6 +93,9 @@ func (m *Monitor) Run(ctx context.Context) error { log.Printf("[monitor] session expired, resetting sync buf") m.getUpdatesBuf = "" m.saveBuf() + if err := ClearContextTokens(m.client.BotID()); err != nil { + log.Printf("[monitor] failed to clear context tokens: %v", err) + } } else { // Sync buf already empty but still getting session expired: // the bot token itself has expired. The user needs to re-login. diff --git a/messaging/handler.go b/messaging/handler.go index dcfee34..42ea830 100644 --- a/messaging/handler.go +++ b/messaging/handler.go @@ -39,9 +39,9 @@ type Handler struct { customAliases map[string]string // custom alias -> agent name (from config) factory AgentFactory saveDefault SaveDefaultFunc - contextTokens sync.Map // map[userID]contextToken - saveDir string // directory to save images/files to - seenMsgs sync.Map // map[int64]time.Time — dedup by message_id + contextTokens sync.Map // map[userID]contextToken + saveDir string // directory to save images/files to + seenMsgs sync.Map // map[int64]time.Time — dedup by message_id } // NewHandler creates a new message handler. @@ -299,6 +299,9 @@ func (h *Handler) HandleMessage(ctx context.Context, client *ilink.Client, msg i // Store context token for this user h.contextTokens.Store(msg.FromUserID, msg.ContextToken) + if err := ilink.SaveContextToken(client.BotID(), msg.FromUserID, msg.ContextToken); err != nil { + log.Printf("[handler] failed to save context token for %s: %v", msg.FromUserID, err) + } // Generate a clientID for this reply (used to correlate typing → finish) clientID := NewClientID() diff --git a/messaging/media.go b/messaging/media.go index 94ff753..48147e6 100644 --- a/messaging/media.go +++ b/messaging/media.go @@ -56,6 +56,13 @@ func sendMediaData(ctx context.Context, client *ilink.Client, toUserID, fileName if fileName == "" { fileName = "file" } + if contextToken == "" { + token, err := ilink.LoadContextToken(client.BotID(), toUserID) + if err != nil { + return fmt.Errorf("load context token: %w", err) + } + contextToken = token + } cdnMediaType, itemType := classifyMedia(contentType, source) diff --git a/messaging/sender.go b/messaging/sender.go index cd439b2..229d34e 100644 --- a/messaging/sender.go +++ b/messaging/sender.go @@ -41,6 +41,13 @@ func SendTextReply(ctx context.Context, client *ilink.Client, toUserID, text, co if clientID == "" { clientID = NewClientID() } + if contextToken == "" { + token, err := ilink.LoadContextToken(client.BotID(), toUserID) + if err != nil { + return fmt.Errorf("load context token: %w", err) + } + contextToken = token + } // Convert markdown to plain text for WeChat display plainText := MarkdownToPlainText(text)