Merge 95393d9b68 into 412705783d

pkg/seahorse/short_retrieval.go

@@ -44,6 +44,7 @@ type GrepInput struct {
 	Pattern          string     `json:"pattern"`
 	Scope            string     `json:"scope,omitempty"` // "both" (default), "summary", or "message"
 	Role             string     `json:"role,omitempty"`  // "user", "assistant", or "" (all)
+	ConversationID   int64      `json:"conversationId,omitempty"`
 	AllConversations bool       `json:"allConversations,omitempty"`
 	Since            *time.Time `json:"since,omitempty"`
 	Before           *time.Time `json:"before,omitempty"`
@@ -84,8 +85,9 @@ type GrepMessageResult struct {
 
 // ExpandMessagesResult contains expanded messages.
 type ExpandMessagesResult struct {
-	Messages   []Message `json:"messages"`
+	Messages           []Message `json:"messages"`
-	TokenCount int       `json:"tokenCount"`
+	RejectedMessageIDs []int64   `json:"rejectedMessageIds,omitempty"`
+	TokenCount         int       `json:"tokenCount"`
 }
 
 // Grep searches summaries and messages for matching content.
@@ -120,6 +122,7 @@ func (r *RetrievalEngine) Grep(ctx context.Context, input GrepInput) (*GrepResul
 		Pattern:          input.Pattern,
 		Mode:             mode,
 		Role:             input.Role,
+		ConversationID:   input.ConversationID,
 		AllConversations: input.AllConversations,
 		Since:            since,
 		Before:           input.Before,
@@ -193,8 +196,43 @@ func (r *RetrievalEngine) Grep(ctx context.Context, input GrepInput) (*GrepResul
 	return result, nil
 }
 
+// ConversationIDForSession returns the storage conversation ID for a session key.
+func (r *RetrievalEngine) ConversationIDForSession(ctx context.Context, sessionKey string) (int64, bool, error) {
+	if strings.TrimSpace(sessionKey) == "" {
+		return 0, false, nil
+	}
+	conv, err := r.store.GetConversationBySessionKey(ctx, sessionKey)
+	if err != nil {
+		return 0, false, err
+	}
+	if conv == nil {
+		return 0, false, nil
+	}
+	return conv.ConversationID, true, nil
+}
+
 // ExpandMessages retrieves full message content by IDs.
 func (r *RetrievalEngine) ExpandMessages(ctx context.Context, messageIDs []int64) (*ExpandMessagesResult, error) {
+	return r.expandMessages(ctx, messageIDs, 0, true)
+}
+
+// ExpandMessagesScoped retrieves full message content by IDs, restricted to a conversation
+// unless allConversations is true.
+func (r *RetrievalEngine) ExpandMessagesScoped(
+	ctx context.Context,
+	messageIDs []int64,
+	conversationID int64,
+	allConversations bool,
+) (*ExpandMessagesResult, error) {
+	return r.expandMessages(ctx, messageIDs, conversationID, allConversations)
+}
+
+func (r *RetrievalEngine) expandMessages(
+	ctx context.Context,
+	messageIDs []int64,
+	conversationID int64,
+	allConversations bool,
+) (*ExpandMessagesResult, error) {
 	result := &ExpandMessagesResult{
 		Messages: make([]Message, 0, len(messageIDs)),
 	}
@@ -202,6 +240,11 @@ func (r *RetrievalEngine) ExpandMessages(ctx context.Context, messageIDs []int64
 	for _, msgID := range messageIDs {
 		msg, err := r.store.GetMessageByID(ctx, msgID)
 		if err != nil {
+			result.RejectedMessageIDs = append(result.RejectedMessageIDs, msgID)
+			continue
+		}
+		if !allConversations && (conversationID <= 0 || msg.ConversationID != conversationID) {
+			result.RejectedMessageIDs = append(result.RejectedMessageIDs, msgID)
 			continue
 		}
 		result.Messages = append(result.Messages, *msg)

pkg/seahorse/tool_expand.go

@@ -28,6 +28,7 @@ Use when short_grep returns messages and you need complete content (not just sni
 
 Parameters:
 - message_ids (required): Array of message ID strings (from short_grep results)
+- all_conversations: Expand IDs from any conversation (default: current conversation only)
 
 Returns message with:
 - content: Full text content
@@ -40,9 +41,12 @@ Returns message with:
 Notes:
 - tool_result content is not returned (can be large). Re-run the tool if you need the result.
 - Media files are stored on disk at mediaUri path, use bash to access.
+- By default, IDs outside the current conversation are rejected and reported in rejectedMessageIds.
+- If short_grep used all_conversations: true, pass all_conversations: true to expand those IDs.
 
 Example:
-  {"message_ids": ["10", "25"]}`
+  {"message_ids": ["10", "25"]}
+  {"message_ids": ["10", "25"], "all_conversations": true}`
 }
 
 func (t *ExpandTool) Parameters() map[string]any {
@@ -54,6 +58,10 @@ func (t *ExpandTool) Parameters() map[string]any {
 				"items":       map[string]any{"type": "string"},
 				"description": "Message IDs to expand (from short_grep results, e.g., [\"10\", \"25\"])",
 			},
+			"all_conversations": map[string]any{
+				"type":        "boolean",
+				"description": "Expand IDs across all conversations (default: current conversation only)",
+			},
 		},
 		"required": []string{"message_ids"},
 	}
@@ -82,7 +90,23 @@ func (t *ExpandTool) Execute(ctx context.Context, args map[string]any) *tools.To
 		}
 	}
 
-	result, err := t.engine.ExpandMessages(ctx, messageIDs)
+	allConversations, _ := args["all_conversations"].(bool)
+	var conversationID int64
+	if !allConversations {
+		var found bool
+		var err error
+		conversationID, found, err = t.engine.ConversationIDForSession(ctx, tools.ToolSessionKey(ctx))
+		if err != nil {
+			return tools.ErrorResult("Expand failed: resolve current conversation: " + err.Error())
+		}
+		if !found {
+			return tools.ErrorResult(
+				"Expand failed: no current conversation found for this session. Use all_conversations: true to expand across conversations.",
+			)
+		}
+	}
+
+	result, err := t.engine.ExpandMessagesScoped(ctx, messageIDs, conversationID, allConversations)
 	if err != nil {
 		return tools.ErrorResult("Expand failed: " + err.Error())
 	}
@@ -120,9 +144,10 @@ func (t *ExpandTool) Execute(ctx context.Context, args map[string]any) *tools.To
 	}
 
 	output := map[string]any{
-		"success":    true,
+		"success":            true,
-		"tokenCount": result.TokenCount,
+		"tokenCount":         result.TokenCount,
-		"messages":   messages,
+		"messages":           messages,
+		"rejectedMessageIds": result.RejectedMessageIDs,
 	}
 	data, _ := json.Marshal(output)
 	return tools.NewToolResult(string(data))

pkg/seahorse/tool_expand_test.go

@@ -5,6 +5,8 @@ import (
 	"encoding/json"
 	"fmt"
 	"testing"
+
+	"github.com/sipeed/picoclaw/pkg/tools"
 )
 
 func TestExpandToolByMessageIDs(t *testing.T) {
@@ -19,7 +21,8 @@ func TestExpandToolByMessageIDs(t *testing.T) {
 	tool := NewExpandTool(re)
 
 	result := tool.Execute(ctx, map[string]any{
-		"message_ids": []any{fmt.Sprintf("%d", msg1.ID), fmt.Sprintf("%d", msg2.ID)},
+		"message_ids":       []any{fmt.Sprintf("%d", msg1.ID), fmt.Sprintf("%d", msg2.ID)},
+		"all_conversations": true,
 	})
 
 	if result.IsError {
@@ -76,7 +79,8 @@ func TestExpandToolWithParts(t *testing.T) {
 	tool := NewExpandTool(re)
 
 	result := tool.Execute(ctx, map[string]any{
-		"message_ids": []any{fmt.Sprintf("%d", msg.ID)},
+		"message_ids":       []any{fmt.Sprintf("%d", msg.ID)},
+		"all_conversations": true,
 	})
 
 	if result.IsError {
@@ -134,3 +138,108 @@ func TestExpandToolWithParts(t *testing.T) {
 		t.Error("missing tool_result part")
 	}
 }
+
+func TestExpandToolScopesToCurrentSession(t *testing.T) {
+	s := openTestStore(t)
+	ctx := context.Background()
+	current, _ := s.GetOrCreateConversation(ctx, "session:current")
+	other, _ := s.GetOrCreateConversation(ctx, "session:other")
+	currentMsg, _ := s.AddMessage(ctx, current.ConversationID, "user", "current message", 5)
+	otherMsg, _ := s.AddMessage(ctx, other.ConversationID, "user", "other message", 5)
+
+	tool := NewExpandTool(&RetrievalEngine{store: s})
+	toolCtx := tools.WithToolSessionContext(ctx, "agent", "session:current", nil)
+	result := tool.Execute(toolCtx, map[string]any{
+		"message_ids": []any{
+			float64(currentMsg.ID),
+			float64(otherMsg.ID),
+		},
+	})
+	if result.IsError {
+		t.Fatalf("Execute returned error: %s", result.ContentForLLM())
+	}
+
+	var output struct {
+		Messages []struct {
+			Content string `json:"content"`
+		} `json:"messages"`
+		RejectedMessageIDs []int64 `json:"rejectedMessageIds"`
+	}
+	if err := json.Unmarshal([]byte(result.ContentForLLM()), &output); err != nil {
+		t.Fatalf("unmarshal result: %v", err)
+	}
+	if len(output.Messages) != 1 {
+		t.Fatalf("messages = %d, want 1: %#v", len(output.Messages), output.Messages)
+	}
+	if output.Messages[0].Content != "current message" {
+		t.Fatalf("content = %q, want current message", output.Messages[0].Content)
+	}
+	if len(output.RejectedMessageIDs) != 1 || output.RejectedMessageIDs[0] != otherMsg.ID {
+		t.Fatalf("rejectedMessageIds = %#v, want [%d]", output.RejectedMessageIDs, otherMsg.ID)
+	}
+}
+
+func TestExpandToolCanExpandAllConversations(t *testing.T) {
+	s := openTestStore(t)
+	ctx := context.Background()
+	current, _ := s.GetOrCreateConversation(ctx, "session:current")
+	other, _ := s.GetOrCreateConversation(ctx, "session:other")
+	currentMsg, _ := s.AddMessage(ctx, current.ConversationID, "user", "current message", 5)
+	otherMsg, _ := s.AddMessage(ctx, other.ConversationID, "user", "other message", 5)
+
+	tool := NewExpandTool(&RetrievalEngine{store: s})
+	toolCtx := tools.WithToolSessionContext(ctx, "agent", "session:current", nil)
+	result := tool.Execute(toolCtx, map[string]any{
+		"message_ids": []any{
+			float64(currentMsg.ID),
+			float64(otherMsg.ID),
+		},
+		"all_conversations": true,
+	})
+	if result.IsError {
+		t.Fatalf("Execute returned error: %s", result.ContentForLLM())
+	}
+
+	var output struct {
+		Messages []struct {
+			Content string `json:"content"`
+		} `json:"messages"`
+		RejectedMessageIDs []int64 `json:"rejectedMessageIds"`
+	}
+	if err := json.Unmarshal([]byte(result.ContentForLLM()), &output); err != nil {
+		t.Fatalf("unmarshal result: %v", err)
+	}
+	if len(output.Messages) != 2 {
+		t.Fatalf("messages = %d, want 2: %#v", len(output.Messages), output.Messages)
+	}
+	if len(output.RejectedMessageIDs) != 0 {
+		t.Fatalf("rejectedMessageIds = %#v, want none", output.RejectedMessageIDs)
+	}
+}
+
+func TestExpandToolUnknownSessionErrors(t *testing.T) {
+	s := openTestStore(t)
+	ctx := context.Background()
+	conv, _ := s.GetOrCreateConversation(ctx, "session:current")
+	msg, _ := s.AddMessage(ctx, conv.ConversationID, "user", "current message", 5)
+
+	tool := NewExpandTool(&RetrievalEngine{store: s})
+	toolCtx := tools.WithToolSessionContext(ctx, "agent", "session:missing", nil)
+	result := tool.Execute(toolCtx, map[string]any{
+		"message_ids": []any{float64(msg.ID)},
+	})
+	if !result.IsError {
+		t.Fatal("expected error for unknown current session")
+	}
+}
+
+func TestExpandToolSupportsAllConversationsParameter(t *testing.T) {
+	s := openTestStore(t)
+	tool := NewExpandTool(&RetrievalEngine{store: s})
+	params := tool.Parameters()
+	props := params["properties"].(map[string]any)
+
+	if _, ok := props["all_conversations"]; !ok {
+		t.Error("Parameters missing 'all_conversations' field")
+	}
+}

pkg/seahorse/tool_grep.go

@@ -131,6 +131,21 @@ func (t *GrepTool) Execute(ctx context.Context, args map[string]any) *tools.Tool
 	if allConv, ok := args["all_conversations"].(bool); ok {
 		input.AllConversations = allConv
 	}
+	if !input.AllConversations {
+		conversationID, found, err := t.engine.ConversationIDForSession(ctx, tools.ToolSessionKey(ctx))
+		if err != nil {
+			return tools.ErrorResult("Grep failed: resolve current conversation: " + err.Error())
+		}
+		if !found {
+			return grepJSONResult(&GrepResult{
+				Success:   true,
+				Summaries: make([]GrepSummaryResult, 0),
+				Messages:  make([]GrepMessageResult, 0),
+				Hint:      "No current conversation found for this session. Use all_conversations: true to search across conversations.",
+			})
+		}
+		input.ConversationID = conversationID
+	}
 	if limit, ok := args["limit"].(float64); ok {
 		input.Limit = int(limit)
 	}
@@ -155,7 +170,10 @@ func (t *GrepTool) Execute(ctx context.Context, args map[string]any) *tools.Tool
 		return tools.ErrorResult("Grep failed: " + err.Error())
 	}
 
-	// Build response
+	return grepJSONResult(result)
+}
+
+func grepJSONResult(result *GrepResult) *tools.ToolResult {
 	output := map[string]any{
 		"success":   result.Success,
 		"summaries": result.Summaries,

pkg/seahorse/tool_grep_test.go

@@ -2,7 +2,10 @@ package seahorse
 
 import (
 	"context"
+	"encoding/json"
 	"testing"
+
+	"github.com/sipeed/picoclaw/pkg/tools"
 )
 
 func TestGrepSearchSummaries(t *testing.T) {
@@ -70,3 +73,111 @@ func TestGrepToolSupportsAllConversations(t *testing.T) {
 		t.Error("Parameters missing 'all_conversations' field")
 	}
 }
+
+func TestGrepToolScopesToCurrentSessionByDefault(t *testing.T) {
+	s := openTestStore(t)
+	ctx := context.Background()
+	current, _ := s.GetOrCreateConversation(ctx, "session:current")
+	other, _ := s.GetOrCreateConversation(ctx, "session:other")
+	s.AddMessage(ctx, current.ConversationID, "user", "shared needle from current topic", 5)
+	s.AddMessage(ctx, other.ConversationID, "user", "shared needle from other topic", 5)
+
+	tool := NewGrepTool(&RetrievalEngine{store: s})
+	toolCtx := tools.WithToolSessionContext(ctx, "agent", "session:current", nil)
+	result := tool.Execute(toolCtx, map[string]any{"pattern": "needle"})
+	if result.IsError {
+		t.Fatalf("Execute returned error: %s", result.ContentForLLM())
+	}
+
+	var output struct {
+		Messages []GrepMessageResult `json:"messages"`
+	}
+	if err := json.Unmarshal([]byte(result.ContentForLLM()), &output); err != nil {
+		t.Fatalf("unmarshal result: %v", err)
+	}
+	if len(output.Messages) != 1 {
+		t.Fatalf("messages = %d, want 1: %#v", len(output.Messages), output.Messages)
+	}
+	if output.Messages[0].ConversationID != current.ConversationID {
+		t.Fatalf("conversation id = %d, want %d", output.Messages[0].ConversationID, current.ConversationID)
+	}
+}
+
+func TestGrepToolCanSearchAllConversations(t *testing.T) {
+	s := openTestStore(t)
+	ctx := context.Background()
+	current, _ := s.GetOrCreateConversation(ctx, "session:current")
+	other, _ := s.GetOrCreateConversation(ctx, "session:other")
+	s.AddMessage(ctx, current.ConversationID, "user", "shared needle from current topic", 5)
+	s.AddMessage(ctx, other.ConversationID, "user", "shared needle from other topic", 5)
+
+	tool := NewGrepTool(&RetrievalEngine{store: s})
+	toolCtx := tools.WithToolSessionContext(ctx, "agent", "session:current", nil)
+	result := tool.Execute(toolCtx, map[string]any{"pattern": "needle", "all_conversations": true})
+	if result.IsError {
+		t.Fatalf("Execute returned error: %s", result.ContentForLLM())
+	}
+
+	var output struct {
+		Messages []GrepMessageResult `json:"messages"`
+	}
+	if err := json.Unmarshal([]byte(result.ContentForLLM()), &output); err != nil {
+		t.Fatalf("unmarshal result: %v", err)
+	}
+	if len(output.Messages) != 2 {
+		t.Fatalf("messages = %d, want 2: %#v", len(output.Messages), output.Messages)
+	}
+}
+
+func TestGrepToolUnknownSessionDoesNotSearchAllConversations(t *testing.T) {
+	s := openTestStore(t)
+	ctx := context.Background()
+	current, _ := s.GetOrCreateConversation(ctx, "session:current")
+	other, _ := s.GetOrCreateConversation(ctx, "session:other")
+	s.AddMessage(ctx, current.ConversationID, "user", "shared needle from current topic", 5)
+	s.AddMessage(ctx, other.ConversationID, "user", "shared needle from other topic", 5)
+
+	tool := NewGrepTool(&RetrievalEngine{store: s})
+	toolCtx := tools.WithToolSessionContext(ctx, "agent", "session:missing", nil)
+	result := tool.Execute(toolCtx, map[string]any{"pattern": "needle"})
+	if result.IsError {
+		t.Fatalf("Execute returned error: %s", result.ContentForLLM())
+	}
+
+	var output struct {
+		Messages []GrepMessageResult `json:"messages"`
+		Hint     string              `json:"hint"`
+	}
+	if err := json.Unmarshal([]byte(result.ContentForLLM()), &output); err != nil {
+		t.Fatalf("unmarshal result: %v", err)
+	}
+	if len(output.Messages) != 0 {
+		t.Fatalf("messages = %d, want 0: %#v", len(output.Messages), output.Messages)
+	}
+	if output.Hint == "" {
+		t.Fatal("expected hint for missing current conversation")
+	}
+}
+
+func TestGrepToolEmptySessionDoesNotSearchAllConversations(t *testing.T) {
+	s := openTestStore(t)
+	ctx := context.Background()
+	conv, _ := s.GetOrCreateConversation(ctx, "session:current")
+	s.AddMessage(ctx, conv.ConversationID, "user", "shared needle from current topic", 5)
+
+	tool := NewGrepTool(&RetrievalEngine{store: s})
+	result := tool.Execute(ctx, map[string]any{"pattern": "needle"})
+	if result.IsError {
+		t.Fatalf("Execute returned error: %s", result.ContentForLLM())
+	}
+
+	var output struct {
+		Messages []GrepMessageResult `json:"messages"`
+	}
+	if err := json.Unmarshal([]byte(result.ContentForLLM()), &output); err != nil {
+		t.Fatalf("unmarshal result: %v", err)
+	}
+	if len(output.Messages) != 0 {
+		t.Fatalf("messages = %d, want 0: %#v", len(output.Messages), output.Messages)
+	}
+}