Enha: do not remove tag

2 files changed, 136 insertions, 131 deletions

diff --git a/bot.go b/bot.go
index fc878d5..3af0ac9 100644
--- a/bot.go
+++ b/bot.go
@@ -71,9 +71,9 @@ var (
 
 // parseKnownToTag extracts known_to list from content using configured tag.
 // Returns cleaned content and list of character names.
-func parseKnownToTag(content string) (string, []string) {
+func parseKnownToTag(content string) []string {
 	if cfg == nil || !cfg.CharSpecificContextEnabled {
-		return content, nil
+		return nil
 	}
 	tag := cfg.CharSpecificContextTag
 	if tag == "" {
@@ -84,17 +84,15 @@ func parseKnownToTag(content string) (string, []string) {
 	re := regexp.MustCompile(pattern)
 	matches := re.FindAllStringSubmatch(content, -1)
 	if len(matches) == 0 {
-		return content, nil
+		return nil
 	}
 	// There may be multiple tags; we combine all.
 	var knownTo []string
-	cleaned := content
 	for _, match := range matches {
 		if len(match) < 2 {
 			continue
 		}
 		// Remove the entire matched tag from content
-		cleaned = strings.Replace(cleaned, match[0], "", 1)
 		list := strings.TrimSpace(match[1])
 		if list == "" {
 			continue
@@ -108,7 +106,7 @@ func parseKnownToTag(content string) (string, []string) {
 		}
 	}
 	// Also remove any leftover trailing "__" that might be orphaned? Not needed.
-	return strings.TrimSpace(cleaned), knownTo
+	return knownTo
 }
 
 // processMessageTag processes a message for known_to tag and sets KnownTo field.
@@ -120,11 +118,8 @@ func processMessageTag(msg models.RoleMsg) models.RoleMsg {
 	}
 	// If KnownTo already set, assume tag already processed (content cleaned).
 	// However, we still check for new tags (maybe added later).
-	cleaned, knownTo := parseKnownToTag(msg.Content)
+	knownTo := parseKnownToTag(msg.Content)
 	logger.Info("processing tags", "msg", msg.Content, "known_to", knownTo)
-	if cleaned != msg.Content {
-		msg.Content = cleaned
-	}
 	// If tag found, replace KnownTo with new list (merge with existing?)
 	// For simplicity, if knownTo is not nil, replace.
 	if knownTo != nil {
@@ -789,10 +784,17 @@ out:
 	if resume {
 		chatBody.Messages[len(chatBody.Messages)-1].Content += respText.String()
 		// lastM.Content = lastM.Content + respText.String()
+		// Process the updated message to check for known_to tags in resumed response
+		updatedMsg := chatBody.Messages[len(chatBody.Messages)-1]
+		processedMsg := processMessageTag(updatedMsg)
+		chatBody.Messages[len(chatBody.Messages)-1] = processedMsg
 	} else {
-		chatBody.Messages = append(chatBody.Messages, models.RoleMsg{
+		newMsg := models.RoleMsg{
 			Role: botPersona, Content: respText.String(),
-		})
+		}
+		// Process the new message to check for known_to tags in LLM response
+		newMsg = processMessageTag(newMsg)
+		chatBody.Messages = append(chatBody.Messages, newMsg)
 	}
 	logger.Debug("chatRound: before cleanChatBody", "messages_before_clean", len(chatBody.Messages))
 	for i, msg := range chatBody.Messages {
diff --git a/bot_test.go b/bot_test.go
index 7496175..a745dde 100644
--- a/bot_test.go
+++ b/bot_test.go
@@ -118,34 +118,34 @@ func TestConsolidateConsecutiveAssistantMessages(t *testing.T) {
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			result := consolidateConsecutiveAssistantMessages(tt.input)
-			
+
 			if len(result) != len(tt.expected) {
 				t.Errorf("Expected %d messages, got %d", len(tt.expected), len(result))
 				t.Logf("Result: %+v", result)
 				t.Logf("Expected: %+v", tt.expected)
 				return
 			}
-			
+
 			for i, expectedMsg := range tt.expected {
 				if i >= len(result) {
 					t.Errorf("Result has fewer messages than expected at index %d", i)
 					continue
 				}
-				
+
 				actualMsg := result[i]
 				if actualMsg.Role != expectedMsg.Role {
 					t.Errorf("Message %d: expected role '%s', got '%s'", i, expectedMsg.Role, actualMsg.Role)
 				}
-				
+
 				if actualMsg.Content != expectedMsg.Content {
 					t.Errorf("Message %d: expected content '%s', got '%s'", i, expectedMsg.Content, actualMsg.Content)
 				}
-				
+
 				if actualMsg.ToolCallID != expectedMsg.ToolCallID {
 					t.Errorf("Message %d: expected ToolCallID '%s', got '%s'", i, expectedMsg.ToolCallID, actualMsg.ToolCallID)
 				}
 			}
-			
+
 			// Additional check: ensure no messages were lost
 			if !reflect.DeepEqual(result, tt.expected) {
 				t.Errorf("Result does not match expected:\nResult:   %+v\nExpected: %+v", result, tt.expected)
@@ -290,92 +290,92 @@ func TestConvertJSONToMapStringString(t *testing.T) {
 
 func TestParseKnownToTag(t *testing.T) {
 	tests := []struct {
-		name           string
-		content        string
-		enabled        bool
-		tag            string
-		wantCleaned    string
-		wantKnownTo    []string
+		name        string
+		content     string
+		enabled     bool
+		tag         string
+		wantCleaned string
+		wantKnownTo []string
 	}{
 		{
-			name:           "feature disabled returns original",
-			content:        "Hello __known_to_chars__Alice__",
-			enabled:        false,
-			tag:            "__known_to_chars__",
-			wantCleaned:    "Hello __known_to_chars__Alice__",
-			wantKnownTo:    nil,
+			name:        "feature disabled returns original",
+			content:     "Hello __known_to_chars__Alice__",
+			enabled:     false,
+			tag:         "__known_to_chars__",
+			wantCleaned: "Hello __known_to_chars__Alice__",
+			wantKnownTo: nil,
 		},
 		{
-			name:           "no tag returns original",
-			content:        "Hello Alice",
-			enabled:        true,
-			tag:            "__known_to_chars__",
-			wantCleaned:    "Hello Alice",
-			wantKnownTo:    nil,
+			name:        "no tag returns original",
+			content:     "Hello Alice",
+			enabled:     true,
+			tag:         "__known_to_chars__",
+			wantCleaned: "Hello Alice",
+			wantKnownTo: nil,
 		},
 		{
-			name:           "single tag with one char",
-			content:        "Hello __known_to_chars__Alice__",
-			enabled:        true,
-			tag:            "__known_to_chars__",
-			wantCleaned:    "Hello",
-			wantKnownTo:    []string{"Alice"},
+			name:        "single tag with one char",
+			content:     "Hello __known_to_chars__Alice__",
+			enabled:     true,
+			tag:         "__known_to_chars__",
+			wantCleaned: "Hello",
+			wantKnownTo: []string{"Alice"},
 		},
 		{
-			name:           "single tag with two chars",
-			content:        "Secret __known_to_chars__Alice,Bob__ message",
-			enabled:        true,
-			tag:            "__known_to_chars__",
-			wantCleaned:    "Secret  message",
-			wantKnownTo:    []string{"Alice", "Bob"},
+			name:        "single tag with two chars",
+			content:     "Secret __known_to_chars__Alice,Bob__ message",
+			enabled:     true,
+			tag:         "__known_to_chars__",
+			wantCleaned: "Secret  message",
+			wantKnownTo: []string{"Alice", "Bob"},
 		},
 		{
-			name:           "tag at beginning",
-			content:        "__known_to_chars__Alice__ Hello",
-			enabled:        true,
-			tag:            "__known_to_chars__",
-			wantCleaned:    "Hello",
-			wantKnownTo:    []string{"Alice"},
+			name:        "tag at beginning",
+			content:     "__known_to_chars__Alice__ Hello",
+			enabled:     true,
+			tag:         "__known_to_chars__",
+			wantCleaned: "Hello",
+			wantKnownTo: []string{"Alice"},
 		},
 		{
-			name:           "tag at end",
-			content:        "Hello __known_to_chars__Alice__",
-			enabled:        true,
-			tag:            "__known_to_chars__",
-			wantCleaned:    "Hello",
-			wantKnownTo:    []string{"Alice"},
+			name:        "tag at end",
+			content:     "Hello __known_to_chars__Alice__",
+			enabled:     true,
+			tag:         "__known_to_chars__",
+			wantCleaned: "Hello",
+			wantKnownTo: []string{"Alice"},
 		},
 		{
-			name:           "multiple tags",
-			content:        "First __known_to_chars__Alice__ then __known_to_chars__Bob__",
-			enabled:        true,
-			tag:            "__known_to_chars__",
-			wantCleaned:    "First  then",
-			wantKnownTo:    []string{"Alice", "Bob"},
+			name:        "multiple tags",
+			content:     "First __known_to_chars__Alice__ then __known_to_chars__Bob__",
+			enabled:     true,
+			tag:         "__known_to_chars__",
+			wantCleaned: "First  then",
+			wantKnownTo: []string{"Alice", "Bob"},
 		},
 		{
-			name:           "custom tag",
-			content:        "Secret __secret__Alice,Bob__ message",
-			enabled:        true,
-			tag:            "__secret__",
-			wantCleaned:    "Secret  message",
-			wantKnownTo:    []string{"Alice", "Bob"},
+			name:        "custom tag",
+			content:     "Secret __secret__Alice,Bob__ message",
+			enabled:     true,
+			tag:         "__secret__",
+			wantCleaned: "Secret  message",
+			wantKnownTo: []string{"Alice", "Bob"},
 		},
 		{
-			name:           "empty list",
-			content:        "Secret __known_to_chars____",
-			enabled:        true,
-			tag:            "__known_to_chars__",
-			wantCleaned:    "Secret",
-			wantKnownTo:    nil,
+			name:        "empty list",
+			content:     "Secret __known_to_chars____",
+			enabled:     true,
+			tag:         "__known_to_chars__",
+			wantCleaned: "Secret",
+			wantKnownTo: nil,
 		},
 		{
-			name:           "whitespace around commas",
-			content:        "__known_to_chars__ Alice , Bob , Carl __",
-			enabled:        true,
-			tag:            "__known_to_chars__",
-			wantCleaned:    "",
-			wantKnownTo:    []string{"Alice", "Bob", "Carl"},
+			name:        "whitespace around commas",
+			content:     "__known_to_chars__ Alice , Bob , Carl __",
+			enabled:     true,
+			tag:         "__known_to_chars__",
+			wantCleaned: "",
+			wantKnownTo: []string{"Alice", "Bob", "Carl"},
 		},
 	}
 
@@ -387,13 +387,7 @@ func TestParseKnownToTag(t *testing.T) {
 				CharSpecificContextTag:     tt.tag,
 			}
 			cfg = testCfg
-
-			cleaned, knownTo := parseKnownToTag(tt.content)
-
-			if cleaned != tt.wantCleaned {
-				t.Errorf("parseKnownToTag() cleaned = %q, want %q", cleaned, tt.wantCleaned)
-			}
-
+			knownTo := parseKnownToTag(tt.content)
 			if len(knownTo) != len(tt.wantKnownTo) {
 				t.Errorf("parseKnownToTag() knownTo length = %v, want %v", len(knownTo), len(tt.wantKnownTo))
 				t.Logf("got: %v", knownTo)
@@ -411,11 +405,11 @@ func TestParseKnownToTag(t *testing.T) {
 
 func TestProcessMessageTag(t *testing.T) {
 	tests := []struct {
-		name     string
-		msg      models.RoleMsg
-		enabled  bool
-		tag      string
-		wantMsg  models.RoleMsg
+		name    string
+		msg     models.RoleMsg
+		enabled bool
+		tag     string
+		wantMsg models.RoleMsg
 	}{
 		{
 			name: "feature disabled returns unchanged",
@@ -488,6 +482,21 @@ func TestProcessMessageTag(t *testing.T) {
 				KnownTo: []string{"Bob", "Alice"},
 			},
 		},
+		{
+			name: "example from real use",
+			msg: models.RoleMsg{
+				Role:    "Alice",
+				Content: "I'll start with a simple one! The word is 'banana'. (ooc: __known_to_chars__Bob__)",
+				KnownTo: []string{"Alice"}, // from previous processing
+			},
+			enabled: true,
+			tag:     "__known_to_chars__",
+			wantMsg: models.RoleMsg{
+				Role:    "Alice",
+				Content: "I'll start with a simple one! The word is 'banana'. (ooc: __known_to_chars__Bob__)",
+				KnownTo: []string{"Bob", "Alice"},
+			},
+		},
 	}
 
 	for _, tt := range tests {
@@ -497,13 +506,7 @@ func TestProcessMessageTag(t *testing.T) {
 				CharSpecificContextTag:     tt.tag,
 			}
 			cfg = testCfg
-
 			got := processMessageTag(tt.msg)
-
-			if got.Content != tt.wantMsg.Content {
-				t.Errorf("processMessageTag() content = %q, want %q", got.Content, tt.wantMsg.Content)
-			}
-
 			if len(got.KnownTo) != len(tt.wantMsg.KnownTo) {
 				t.Errorf("processMessageTag() KnownTo length = %v, want %v", len(got.KnownTo), len(tt.wantMsg.KnownTo))
 				t.Logf("got: %v", got.KnownTo)
@@ -530,7 +533,7 @@ func TestProcessMessageTag(t *testing.T) {
 func TestFilterMessagesForCharacter(t *testing.T) {
 	messages := []models.RoleMsg{
 		{Role: "system", Content: "System message", KnownTo: nil}, // visible to all
-		{Role: "Alice", Content: "Hello everyone", KnownTo: nil}, // visible to all
+		{Role: "Alice", Content: "Hello everyone", KnownTo: nil},  // visible to all
 		{Role: "Alice", Content: "Secret for Bob", KnownTo: []string{"Alice", "Bob"}},
 		{Role: "Bob", Content: "Reply to Alice", KnownTo: []string{"Alice", "Bob"}},
 		{Role: "Alice", Content: "Private to Carl", KnownTo: []string{"Alice", "Carl"}},
@@ -538,46 +541,46 @@ func TestFilterMessagesForCharacter(t *testing.T) {
 	}
 
 	tests := []struct {
-		name      string
-		enabled   bool
-		character string
+		name        string
+		enabled     bool
+		character   string
 		wantIndices []int // indices from original messages that should be included
 	}{
 		{
-			name:      "feature disabled returns all",
-			enabled:   false,
-			character: "Alice",
-			wantIndices: []int{0,1,2,3,4,5},
+			name:        "feature disabled returns all",
+			enabled:     false,
+			character:   "Alice",
+			wantIndices: []int{0, 1, 2, 3, 4, 5},
 		},
 		{
-			name:      "character empty returns all",
-			enabled:   true,
-			character: "",
-			wantIndices: []int{0,1,2,3,4,5},
+			name:        "character empty returns all",
+			enabled:     true,
+			character:   "",
+			wantIndices: []int{0, 1, 2, 3, 4, 5},
 		},
 		{
-			name:      "Alice sees all including Carl-private",
-			enabled:   true,
-			character: "Alice",
-			wantIndices: []int{0,1,2,3,4,5},
+			name:        "Alice sees all including Carl-private",
+			enabled:     true,
+			character:   "Alice",
+			wantIndices: []int{0, 1, 2, 3, 4, 5},
 		},
 		{
-			name:      "Bob sees Alice-Bob secrets and all public",
-			enabled:   true,
-			character: "Bob",
-			wantIndices: []int{0,1,2,3,5},
+			name:        "Bob sees Alice-Bob secrets and all public",
+			enabled:     true,
+			character:   "Bob",
+			wantIndices: []int{0, 1, 2, 3, 5},
 		},
 		{
-			name:      "Carl sees Alice-Carl secret and public",
-			enabled:   true,
-			character: "Carl",
-			wantIndices: []int{0,1,4,5},
+			name:        "Carl sees Alice-Carl secret and public",
+			enabled:     true,
+			character:   "Carl",
+			wantIndices: []int{0, 1, 4, 5},
 		},
 		{
-			name:      "David sees only public messages",
-			enabled:   true,
-			character: "David",
-			wantIndices: []int{0,1,5},
+			name:        "David sees only public messages",
+			enabled:     true,
+			character:   "David",
+			wantIndices: []int{0, 1, 5},
 		},
 	}
 
@@ -604,4 +607,4 @@ func TestFilterMessagesForCharacter(t *testing.T) {
 			}
 		})
 	}
-}
\ No newline at end of file
+}