diff --git a/cmd/chat_tools.go b/cmd/chat_tools.go index 03aedaa5..49c6d737 100644 --- a/cmd/chat_tools.go +++ b/cmd/chat_tools.go @@ -52,6 +52,7 @@ func essentialTools() []tool.Tool { tool.WebSearchTool{}, tool.ToolSearchTool{}, tool.SkillTool{}, + tool.SessionQueryTool{}, tool.AgentTool{}, tool.AskUserQuestionTool{}, tool.TodoWriteTool{}, diff --git a/internal/sessionquery/db.go b/internal/sessionquery/db.go new file mode 100644 index 00000000..b9ff9e18 --- /dev/null +++ b/internal/sessionquery/db.go @@ -0,0 +1,89 @@ +package sessionquery + +import ( + "database/sql" + "fmt" + "os" + "path/filepath" + "sync" + + "github.com/GrayCodeAI/hawk/internal/storage" + _ "modernc.org/sqlite" +) + +const schemaDDL = ` +CREATE TABLE IF NOT EXISTS sessions_meta ( + session_id TEXT PRIMARY KEY, + workspace TEXT NOT NULL, + model TEXT NOT NULL, + provider TEXT NOT NULL, + updated_at INTEGER NOT NULL, + message_count INTEGER NOT NULL, + file_mod_time INTEGER NOT NULL +); + +CREATE INDEX IF NOT EXISTS idx_sessions_workspace ON sessions_meta(workspace); + +CREATE VIRTUAL TABLE IF NOT EXISTS messages_fts USING fts5( + session_id UNINDEXED, + role UNINDEXED, + msg_index UNINDEXED, + content, + tokenize='porter unicode61' +); +` + +// DB wraps the SQLite database connection with connection synchronization. +type DB struct { + mu sync.RWMutex + conn *sql.DB + dbPath string +} + +// OpenDB opens (or creates) the SQLite database at dbPath and initializes the schema. +func OpenDB(dbPath string) (*DB, error) { + if dbPath == "" { + dbPath = filepath.Join(storage.CacheDir(), "session_query.db") + } + + if err := os.MkdirAll(filepath.Dir(dbPath), 0o750); err != nil { + return nil, fmt.Errorf("failed to create directory for session query database: %w", err) + } + + conn, err := sql.Open("sqlite", dbPath+"?_pragma=busy_timeout(5000)&_pragma=journal_mode(WAL)&_pragma=synchronous(NORMAL)") + if err != nil { + return nil, fmt.Errorf("failed to open session query database: %w", err) + } + + // Single writer connection pool + conn.SetMaxOpenConns(1) + + if _, err := conn.Exec(schemaDDL); err != nil { + _ = conn.Close() + return nil, fmt.Errorf("failed to initialize session query schema: %w", err) + } + + return &DB{ + conn: conn, + dbPath: dbPath, + }, nil +} + +// Close closes the underlying SQLite connection. +func (d *DB) Close() error { + d.mu.Lock() + defer d.mu.Unlock() + if d.conn != nil { + err := d.conn.Close() + d.conn = nil + return err + } + return nil +} + +// Conn returns the raw SQL DB connection. +func (d *DB) Conn() *sql.DB { + d.mu.RLock() + defer d.mu.RUnlock() + return d.conn +} diff --git a/internal/sessionquery/indexer.go b/internal/sessionquery/indexer.go new file mode 100644 index 00000000..bd925c7e --- /dev/null +++ b/internal/sessionquery/indexer.go @@ -0,0 +1,265 @@ +package sessionquery + +import ( + "context" + "database/sql" + "fmt" + "os" + "path/filepath" + "strings" + "sync" + "time" + + "github.com/GrayCodeAI/hawk/internal/session" + "github.com/GrayCodeAI/hawk/internal/storage" +) + +// Indexer manages incremental full-text indexing of sessions. +type Indexer struct { + db *DB + sessionsDir string + mu sync.Mutex +} + +// NewIndexer creates a new session indexer. +func NewIndexer(db *DB, sessionsDir string) *Indexer { + if sessionsDir == "" { + sessionsDir = storage.SessionsDir() + } + return &Indexer{ + db: db, + sessionsDir: sessionsDir, + } +} + +// IndexSession indexes or updates a single session if modified. +func (idx *Indexer) IndexSession(ctx context.Context, sessionID string) (bool, error) { + idx.mu.Lock() + defer idx.mu.Unlock() + + conn := idx.db.Conn() + if conn == nil { + return false, fmt.Errorf("database connection closed") + } + + jsonlPath := filepath.Join(idx.sessionsDir, sessionID+".jsonl") + legacyPath := filepath.Join(idx.sessionsDir, sessionID+".json") + + var targetPath string + var info os.FileInfo + var err error + + if info, err = os.Stat(jsonlPath); err == nil && !info.IsDir() { + targetPath = jsonlPath + } else if info, err = os.Stat(legacyPath); err == nil && !info.IsDir() { + targetPath = legacyPath + } else { + // Session file does not exist, remove from index if present + _ = idx.removeSessionLocked(ctx, conn, sessionID) + return false, fmt.Errorf("session %s not found on disk: %w", sessionID, session.ErrNotFound) + } + + modTimeNano := info.ModTime().UnixNano() + + // Check if already indexed with matching modTime + var storedModTime int64 + err = conn.QueryRowContext(ctx, "SELECT file_mod_time FROM sessions_meta WHERE session_id = ?", sessionID).Scan(&storedModTime) + if err == nil && storedModTime == modTimeNano { + // Fresh, no update needed + return false, nil + } + + // Load session + sess, err := session.Load(sessionID) + if err != nil { + return false, fmt.Errorf("failed to load session %s for indexing: %w", sessionID, err) + } + + tx, err := conn.BeginTx(ctx, nil) + if err != nil { + return false, fmt.Errorf("failed to begin transaction: %w", err) + } + defer func() { _ = tx.Rollback() }() + + // Delete old records + if _, err := tx.ExecContext(ctx, "DELETE FROM messages_fts WHERE session_id = ?", sessionID); err != nil { + return false, fmt.Errorf("failed to clear old FTS entries for %s: %w", sessionID, err) + } + + // Insert message records + stmt, err := tx.PrepareContext(ctx, "INSERT INTO messages_fts(session_id, role, msg_index, content) VALUES (?, ?, ?, ?)") + if err != nil { + return false, fmt.Errorf("failed to prepare FTS insert: %w", err) + } + defer func() { _ = stmt.Close() }() + + msgCount := 0 + for i, msg := range sess.Messages { + content := extractMessageContent(&msg) + if strings.TrimSpace(content) == "" { + continue + } + msgCount++ + if _, err := stmt.ExecContext(ctx, sessionID, msg.Role, i, content); err != nil { + return false, fmt.Errorf("failed to insert FTS entry: %w", err) + } + } + + // Upsert sessions_meta + workspace := sess.CWD + if workspace == "" { + workspace = "." + } + + updatedAtUnix := sess.UpdatedAt.Unix() + if updatedAtUnix <= 0 { + updatedAtUnix = time.Now().Unix() + } + + upsertMeta := ` +INSERT INTO sessions_meta(session_id, workspace, model, provider, updated_at, message_count, file_mod_time) +VALUES (?, ?, ?, ?, ?, ?, ?) +ON CONFLICT(session_id) DO UPDATE SET + workspace = excluded.workspace, + model = excluded.model, + provider = excluded.provider, + updated_at = excluded.updated_at, + message_count = excluded.message_count, + file_mod_time = excluded.file_mod_time; +` + if _, err := tx.ExecContext(ctx, upsertMeta, sessionID, workspace, sess.Model, sess.Provider, updatedAtUnix, msgCount, modTimeNano); err != nil { + return false, fmt.Errorf("failed to upsert session meta: %w", err) + } + + if err := tx.Commit(); err != nil { + return false, fmt.Errorf("failed to commit indexing transaction: %w", err) + } + + _ = targetPath // referenced + return true, nil +} + +func (idx *Indexer) removeSessionLocked(ctx context.Context, conn *sql.DB, sessionID string) error { + _, _ = conn.ExecContext(ctx, "DELETE FROM messages_fts WHERE session_id = ?", sessionID) + _, _ = conn.ExecContext(ctx, "DELETE FROM sessions_meta WHERE session_id = ?", sessionID) + return nil +} + +// SyncAll incrementally indexes all session files in sessionsDir and removes deleted ones. +func (idx *Indexer) SyncAll(ctx context.Context) (int, error) { + entries, err := os.ReadDir(idx.sessionsDir) + if err != nil { + if os.IsNotExist(err) { + return 0, nil + } + return 0, fmt.Errorf("failed to read sessions directory: %w", err) + } + + activeIDs := make(map[string]bool) + indexedCount := 0 + + for _, entry := range entries { + if entry.IsDir() { + continue + } + name := entry.Name() + var sessionID string + if strings.HasSuffix(name, ".jsonl") { + sessionID = strings.TrimSuffix(name, ".jsonl") + } else if strings.HasSuffix(name, ".json") { + sessionID = strings.TrimSuffix(name, ".json") + } else { + continue + } + + activeIDs[sessionID] = true + updated, err := idx.IndexSession(ctx, sessionID) + if err != nil { + continue + } + if updated { + indexedCount++ + } + } + + // Clean up stale sessions in database no longer present on disk + conn := idx.db.Conn() + if conn != nil { + rows, err := conn.QueryContext(ctx, "SELECT session_id FROM sessions_meta") + if err == nil { + var staleIDs []string + for rows.Next() { + var id string + if err := rows.Scan(&id); err == nil { + if !activeIDs[id] { + staleIDs = append(staleIDs, id) + } + } + } + _ = rows.Close() + + idx.mu.Lock() + for _, id := range staleIDs { + _ = idx.removeSessionLocked(ctx, conn, id) + } + idx.mu.Unlock() + } + } + + return indexedCount, nil +} + +// RebuildIndex drops all existing index tables and re-indexes all sessions from scratch. +func (idx *Indexer) RebuildIndex(ctx context.Context) error { + idx.mu.Lock() + defer idx.mu.Unlock() + + conn := idx.db.Conn() + if conn == nil { + return fmt.Errorf("database connection closed") + } + + dropSQL := ` +DROP TABLE IF EXISTS messages_fts; +DROP TABLE IF EXISTS sessions_meta; +` + if _, err := conn.ExecContext(ctx, dropSQL); err != nil { + return fmt.Errorf("failed to drop tables: %w", err) + } + + if _, err := conn.ExecContext(ctx, schemaDDL); err != nil { + return fmt.Errorf("failed to recreate schema: %w", err) + } + + idx.mu.Unlock() + _, err := idx.SyncAll(ctx) + idx.mu.Lock() + return err +} + +func extractMessageContent(msg *session.Message) string { + var sb strings.Builder + if msg.Content != "" { + sb.WriteString(msg.Content) + sb.WriteString(" ") + } + for _, p := range msg.ContentParts { + if p.Type == "text" && p.Text != "" { + sb.WriteString(p.Text) + sb.WriteString(" ") + } + } + for _, tc := range msg.ToolUse { + if tc.Name != "" { + sb.WriteString(tc.Name) + sb.WriteString(" ") + } + } + for _, tr := range msg.ToolResults { + if tr.Content != "" { + sb.WriteString(tr.Content) + sb.WriteString(" ") + } + } + return strings.TrimSpace(sb.String()) +} diff --git a/internal/sessionquery/search.go b/internal/sessionquery/search.go new file mode 100644 index 00000000..269b385d --- /dev/null +++ b/internal/sessionquery/search.go @@ -0,0 +1,274 @@ +package sessionquery + +import ( + "context" + "database/sql" + "errors" + "fmt" + "path/filepath" + "strings" + "unicode" + + "github.com/GrayCodeAI/hawk/internal/session" +) + +// ErrUnauthorized indicates the caller is not authorized to access sessions in the requested workspace. +var ErrUnauthorized = errors.New("unauthorized: caller workspace does not have access to the target session") + +// SearchParams specifies filters and parameters for full-text session queries. +type SearchParams struct { + CallerWorkspace string `json:"caller_workspace,omitempty"` // Enforces workspace authorization boundary + Workspace string `json:"workspace,omitempty"` // Filter by workspace path + SessionID string `json:"session_id,omitempty"` // Scope search to a specific session + Query string `json:"query"` // Full-text search term + Roles []string `json:"roles,omitempty"` // e.g. ["user", "assistant"] + Limit int `json:"limit,omitempty"` // Results page limit (default 10, max 50) + Offset int `json:"offset,omitempty"` // Pagination offset + MaxBytes int `json:"max_bytes,omitempty"` // Maximum total bytes of content to return (default 16 KiB) +} + +// SearchMatch represents an individual matched message within a session. +type SearchMatch struct { + SessionID string `json:"session_id"` + Workspace string `json:"workspace"` + Model string `json:"model,omitempty"` + Provider string `json:"provider,omitempty"` + Role string `json:"role"` + MsgIndex int `json:"msg_index"` + Snippet string `json:"snippet"` + Content string `json:"content"` +} + +// SearchResponse represents the paginated search result. +type SearchResponse struct { + Matches []SearchMatch `json:"matches"` + TotalCount int `json:"total_count"` + HasMore bool `json:"has_more"` + Offset int `json:"offset"` + Limit int `json:"limit"` + Query string `json:"query"` +} + +// Search executes an FTS5 search query over the indexed session database. +func (d *DB) Search(ctx context.Context, params SearchParams) (*SearchResponse, error) { + conn := d.Conn() + if conn == nil { + return nil, fmt.Errorf("database connection closed") + } + + rawQuery := strings.TrimSpace(params.Query) + if rawQuery == "" { + return &SearchResponse{ + Matches: []SearchMatch{}, + TotalCount: 0, + HasMore: false, + Offset: params.Offset, + Limit: params.Limit, + Query: rawQuery, + }, nil + } + + sanitizedQuery := sanitizeFTS5Query(rawQuery) + if sanitizedQuery == "" { + return &SearchResponse{ + Matches: []SearchMatch{}, + TotalCount: 0, + HasMore: false, + Offset: params.Offset, + Limit: params.Limit, + Query: rawQuery, + }, nil + } + + limit := params.Limit + if limit <= 0 { + limit = 10 + } else if limit > 50 { + limit = 50 + } + + offset := params.Offset + if offset < 0 { + offset = 0 + } + + maxBytes := params.MaxBytes + if maxBytes <= 0 { + maxBytes = 16 * 1024 // 16 KiB cap + } + + // 1. Authorization check if specific SessionID is requested + if params.SessionID != "" && params.CallerWorkspace != "" { + var sessWorkspace string + err := conn.QueryRowContext(ctx, "SELECT workspace FROM sessions_meta WHERE session_id = ?", params.SessionID).Scan(&sessWorkspace) + if err == sql.ErrNoRows { + // Session not in index yet or doesn't exist + return nil, fmt.Errorf("session %s not found: %w", params.SessionID, session.ErrNotFound) + } else if err != nil { + return nil, fmt.Errorf("failed to query session metadata: %w", err) + } + + if !isWorkspaceAuthorized(params.CallerWorkspace, sessWorkspace) { + return nil, ErrUnauthorized + } + } + + // 2. Build Query + var whereClauses []string + var args []interface{} + + // FTS match clause + whereClauses = append(whereClauses, "messages_fts MATCH ?") + args = append(args, sanitizedQuery) + + if params.SessionID != "" { + whereClauses = append(whereClauses, "f.session_id = ?") + args = append(args, params.SessionID) + } + + if params.Workspace != "" { + whereClauses = append(whereClauses, "(m.workspace = ? OR m.workspace LIKE ?)") + args = append(args, params.Workspace, params.Workspace+"/%") + } else if params.CallerWorkspace != "" && params.SessionID == "" { + // Restrict to caller's workspace + whereClauses = append(whereClauses, "(m.workspace = ? OR m.workspace LIKE ?)") + args = append(args, params.CallerWorkspace, params.CallerWorkspace+"/%") + } + + if len(params.Roles) > 0 { + rolePlaceholders := make([]string, len(params.Roles)) + for i, r := range params.Roles { + rolePlaceholders[i] = "?" + args = append(args, r) + } + whereClauses = append(whereClauses, fmt.Sprintf("f.role IN (%s)", strings.Join(rolePlaceholders, ", "))) + } + + whereSQL := strings.Join(whereClauses, " AND ") + + // Total count query + // #nosec G201 -- SQL where clauses use fixed parameterized placeholders with bind arguments + countSQL := fmt.Sprintf(` +SELECT COUNT(*) +FROM messages_fts f +JOIN sessions_meta m ON f.session_id = m.session_id +WHERE %s +`, whereSQL) + + var totalCount int + if err := conn.QueryRowContext(ctx, countSQL, args...).Scan(&totalCount); err != nil { + return nil, fmt.Errorf("search count query failed: %w", err) + } + + // Select query with snippet + // #nosec G201 -- SQL where clauses use fixed parameterized placeholders with bind arguments + selectSQL := fmt.Sprintf(` +SELECT f.session_id, m.workspace, m.model, m.provider, f.role, f.msg_index, f.content, + snippet(messages_fts, 3, '', '', '...', 24) AS match_snippet +FROM messages_fts f +JOIN sessions_meta m ON f.session_id = m.session_id +WHERE %s +ORDER BY f.rank +LIMIT ? OFFSET ? +`, whereSQL) + + selectArgs := append(args, limit+1, offset) // fetch limit+1 to check hasMore + + rows, err := conn.QueryContext(ctx, selectSQL, selectArgs...) + if err != nil { + return nil, fmt.Errorf("search query failed: %w", err) + } + defer func() { _ = rows.Close() }() + + var matches []SearchMatch + currentBytes := 0 + hasMore := false + + for rows.Next() { + if len(matches) >= limit { + hasMore = true + break + } + + var m SearchMatch + var rawSnippet sql.NullString + if err := rows.Scan(&m.SessionID, &m.Workspace, &m.Model, &m.Provider, &m.Role, &m.MsgIndex, &m.Content, &rawSnippet); err != nil { + return nil, fmt.Errorf("failed to scan search match: %w", err) + } + + if rawSnippet.Valid { + m.Snippet = rawSnippet.String + } + + // Apply SecretDetector redaction at the read boundary + m.Snippet = session.RedactSecrets(m.Snippet) + m.Content = session.RedactSecrets(m.Content) + + // Byte bounding check + matchSize := len(m.Snippet) + len(m.Content) + len(m.SessionID) + 64 + if currentBytes+matchSize > maxBytes && len(matches) > 0 { + hasMore = true + break + } + + currentBytes += matchSize + matches = append(matches, m) + } + + if err := rows.Err(); err != nil { + return nil, fmt.Errorf("row iteration error: %w", err) + } + + return &SearchResponse{ + Matches: matches, + TotalCount: totalCount, + HasMore: hasMore || (offset+len(matches) < totalCount), + Offset: offset, + Limit: limit, + Query: rawQuery, + }, nil +} + +func isWorkspaceAuthorized(callerWorkspace, sessionWorkspace string) bool { + if callerWorkspace == "" || callerWorkspace == "." { + return true + } + cClean := filepath.Clean(callerWorkspace) + sClean := filepath.Clean(sessionWorkspace) + if cClean == sClean { + return true + } + // Check if session is inside caller's workspace or vice versa + rel, err := filepath.Rel(cClean, sClean) + if err == nil && !strings.HasPrefix(rel, "..") { + return true + } + return false +} + +// sanitizeFTS5Query cleans raw query text into safe FTS5 terms with prefix support. +func sanitizeFTS5Query(raw string) string { + terms := strings.FieldsFunc(raw, func(r rune) bool { + return unicode.IsSpace(r) || r == '"' || r == '\'' || r == '(' || r == ')' || r == '*' || r == ':' + }) + + if len(terms) == 0 { + return "" + } + + var sanitized []string + for _, term := range terms { + t := strings.TrimSpace(term) + if t == "" || strings.EqualFold(t, "AND") || strings.EqualFold(t, "OR") || strings.EqualFold(t, "NOT") { + continue + } + // Prefix search for each term + sanitized = append(sanitized, fmt.Sprintf("%q*", t)) + } + + if len(sanitized) == 0 { + return "" + } + + return strings.Join(sanitized, " ") +} diff --git a/internal/sessionquery/service.go b/internal/sessionquery/service.go new file mode 100644 index 00000000..c4e6b28f --- /dev/null +++ b/internal/sessionquery/service.go @@ -0,0 +1,89 @@ +package sessionquery + +import ( + "context" + "fmt" + "sync" + "time" + + "github.com/GrayCodeAI/hawk/internal/storage" +) + +// Service encapsulates database indexing and search capabilities. +type Service struct { + db *DB + indexer *Indexer + mu sync.RWMutex +} + +var ( + defaultServiceInstance *Service + defaultServiceOnce sync.Once + defaultServiceErr error +) + +// DefaultService returns the process-wide session query service. +func DefaultService() (*Service, error) { + defaultServiceOnce.Do(func() { + dbPath := storage.CacheDir() + "/session_query.db" + svc, err := NewService(dbPath, storage.SessionsDir()) + if err != nil { + defaultServiceErr = err + return + } + defaultServiceInstance = svc + }) + if defaultServiceErr != nil { + return nil, defaultServiceErr + } + return defaultServiceInstance, nil +} + +// NewService creates a new SessionQuery service with the specified paths. +func NewService(dbPath, sessionsDir string) (*Service, error) { + db, err := OpenDB(dbPath) + if err != nil { + return nil, fmt.Errorf("failed to open sessionquery DB: %w", err) + } + indexer := NewIndexer(db, sessionsDir) + return &Service{ + db: db, + indexer: indexer, + }, nil +} + +// Close closes the underlying resources. +func (s *Service) Close() error { + s.mu.Lock() + defer s.mu.Unlock() + if s.db != nil { + return s.db.Close() + } + return nil +} + +// Indexer returns the underlying session indexer. +func (s *Service) Indexer() *Indexer { + return s.indexer +} + +// DB returns the underlying DB instance. +func (s *Service) DB() *DB { + return s.db +} + +// Search runs a full-text search against the indexed sessions. +// Automatically ensures the index is synchronized first. +func (s *Service) Search(ctx context.Context, params SearchParams) (*SearchResponse, error) { + // Sync index before searching with a short timeout + syncCtx, cancel := context.WithTimeout(ctx, 3*time.Second) + _, _ = s.indexer.SyncAll(syncCtx) + cancel() + + return s.db.Search(ctx, params) +} + +// RebuildIndex drops and rebuilds the full-text search index from scratch. +func (s *Service) RebuildIndex(ctx context.Context) error { + return s.indexer.RebuildIndex(ctx) +} diff --git a/internal/sessionquery/sessionquery_test.go b/internal/sessionquery/sessionquery_test.go new file mode 100644 index 00000000..7d4bcd2c --- /dev/null +++ b/internal/sessionquery/sessionquery_test.go @@ -0,0 +1,286 @@ +package sessionquery + +import ( + "context" + "os" + "path/filepath" + "strings" + "testing" + "time" + + "github.com/GrayCodeAI/hawk/internal/session" +) + +func createTestSessionFile(t *testing.T, dir, sessionID, workspace string, messages []session.Message) { + t.Helper() + sess := &session.Session{ + ID: sessionID, + Model: "gpt-4o", + Provider: "openai", + CWD: workspace, + Messages: messages, + CreatedAt: time.Now(), + UpdatedAt: time.Now(), + } + + // We set test session dir environment so session.Save and session.Load work + t.Setenv("HAWK_STATE_DIR", dir) + sessDir := filepath.Join(dir, "sessions") + _ = os.MkdirAll(sessDir, 0o755) + + if err := session.Save(sess); err != nil { + t.Fatalf("failed to save test session %s: %v", sessionID, err) + } +} + +func setupTestService(t *testing.T) (*Service, string, string) { + t.Helper() + tmpDir := t.TempDir() + dbPath := filepath.Join(tmpDir, "test_session_query.db") + sessDir := filepath.Join(tmpDir, "sessions") + _ = os.MkdirAll(sessDir, 0o755) + t.Setenv("HAWK_STATE_DIR", tmpDir) + + svc, err := NewService(dbPath, sessDir) + if err != nil { + t.Fatalf("failed to create sessionquery service: %v", err) + } + t.Cleanup(func() { + _ = svc.Close() + }) + + return svc, tmpDir, sessDir +} + +func TestIndexingAndTokenization(t *testing.T) { + svc, tmpDir, _ := setupTestService(t) + ctx := context.Background() + + createTestSessionFile(t, tmpDir, "sess-1", "/projects/backend", []session.Message{ + {Role: "user", Content: "How do we deploy the microservice to Kubernetes cluster?"}, + {Role: "assistant", Content: "You can use kubectl apply with deployment.yaml."}, + }) + + createTestSessionFile(t, tmpDir, "sess-2", "/projects/frontend", []session.Message{ + {Role: "user", Content: "The React button component is failing to render."}, + {Role: "assistant", Content: "Check your CSS imports and Vite config."}, + }) + + // 1. Search for "Kubernetes" + res, err := svc.Search(ctx, SearchParams{ + Query: "Kubernetes", + }) + if err != nil { + t.Fatalf("Search failed: %v", err) + } + if len(res.Matches) != 1 { + t.Fatalf("expected 1 match for 'Kubernetes', got %d", len(res.Matches)) + } + if res.Matches[0].SessionID != "sess-1" { + t.Fatalf("expected match from sess-1, got %s", res.Matches[0].SessionID) + } + + // 2. Stemming test: searching "deploying" should match "deploy" (Porter stemmer) + resStem, err := svc.Search(ctx, SearchParams{ + Query: "deploying", + }) + if err != nil { + t.Fatalf("Search with stemming failed: %v", err) + } + if len(resStem.Matches) != 1 { + t.Fatalf("expected stem match for 'deploying', got %d", len(resStem.Matches)) + } + if resStem.Matches[0].SessionID != "sess-1" { + t.Fatalf("expected match from sess-1, got %s", resStem.Matches[0].SessionID) + } +} + +func TestWorkspaceAuthorization(t *testing.T) { + svc, tmpDir, _ := setupTestService(t) + ctx := context.Background() + + createTestSessionFile(t, tmpDir, "sess-auth-1", "/workspaces/team-a", []session.Message{ + {Role: "user", Content: "Configuring Redis cache cluster for team A."}, + }) + + createTestSessionFile(t, tmpDir, "sess-auth-2", "/workspaces/team-b", []session.Message{ + {Role: "user", Content: "Configuring Redis cache cluster for team B."}, + }) + + // Sync index + _, _ = svc.Indexer().SyncAll(ctx) + + // Caller in /workspaces/team-a searching without sessionID + resA, err := svc.Search(ctx, SearchParams{ + CallerWorkspace: "/workspaces/team-a", + Query: "Redis", + }) + if err != nil { + t.Fatalf("Search team-a failed: %v", err) + } + if len(resA.Matches) != 1 || resA.Matches[0].SessionID != "sess-auth-1" { + t.Fatalf("expected only team-a session, got %#v", resA.Matches) + } + + // Caller in /workspaces/team-a attempting to target team-b's sessionID -> ErrUnauthorized + _, err = svc.Search(ctx, SearchParams{ + CallerWorkspace: "/workspaces/team-a", + SessionID: "sess-auth-2", + Query: "Redis", + }) + if err == nil { + t.Fatal("expected ErrUnauthorized when querying unauthorized session, got nil") + } + if !errorsIs(err, ErrUnauthorized) { + t.Fatalf("expected ErrUnauthorized, got %v", err) + } +} + +func TestSecretRedaction(t *testing.T) { + svc, tmpDir, _ := setupTestService(t) + ctx := context.Background() + + secretKey := "sk-ant-api03-abcdef12345678901234567890abcdef12345678" + createTestSessionFile(t, tmpDir, "sess-secret", "/workspaces/app", []session.Message{ + {Role: "user", Content: "Connecting to provider using secret key " + secretKey}, + }) + + res, err := svc.Search(ctx, SearchParams{ + Query: "secret key", + }) + if err != nil { + t.Fatalf("Search failed: %v", err) + } + if len(res.Matches) == 0 { + t.Fatal("expected match for 'secret key'") + } + + match := res.Matches[0] + if strings.Contains(match.Content, secretKey) { + t.Errorf("secret was not redacted in Content: %s", match.Content) + } + if strings.Contains(match.Snippet, secretKey) { + t.Errorf("secret was not redacted in Snippet: %s", match.Snippet) + } + if !strings.Contains(match.Content, "[REDACTED") { + t.Errorf("expected [REDACTED] in Content: %s", match.Content) + } +} + +func TestPagingAndByteBounding(t *testing.T) { + svc, tmpDir, _ := setupTestService(t) + ctx := context.Background() + + var msgs []session.Message + for i := 0; i < 20; i++ { + msgs = append(msgs, session.Message{ + Role: "user", + Content: "Database migration step and schema updates for transaction processing.", + }) + } + createTestSessionFile(t, tmpDir, "sess-paged", "/workspaces/app", msgs) + + // Fetch page 1 (limit 5) + p1, err := svc.Search(ctx, SearchParams{ + Query: "transaction processing", + Limit: 5, + Offset: 0, + }) + if err != nil { + t.Fatalf("page 1 search failed: %v", err) + } + if len(p1.Matches) != 5 { + t.Fatalf("expected 5 matches on page 1, got %d", len(p1.Matches)) + } + if !p1.HasMore { + t.Fatal("expected HasMore=true on page 1") + } + + // Fetch page 2 (limit 5, offset 5) + p2, err := svc.Search(ctx, SearchParams{ + Query: "transaction processing", + Limit: 5, + Offset: 5, + }) + if err != nil { + t.Fatalf("page 2 search failed: %v", err) + } + if len(p2.Matches) != 5 { + t.Fatalf("expected 5 matches on page 2, got %d", len(p2.Matches)) + } + if p2.Matches[0].MsgIndex != p1.Matches[4].MsgIndex+1 { + t.Errorf("expected sequential message indexing across pages") + } +} + +func TestIndexRebuildFromScratch(t *testing.T) { + svc, tmpDir, _ := setupTestService(t) + ctx := context.Background() + + createTestSessionFile(t, tmpDir, "sess-rebuild", "/workspaces/app", []session.Message{ + {Role: "user", Content: "Initial indexing test data."}, + }) + + _, err := svc.Indexer().SyncAll(ctx) + if err != nil { + t.Fatalf("SyncAll failed: %v", err) + } + + // Rebuild index + if err := svc.RebuildIndex(ctx); err != nil { + t.Fatalf("RebuildIndex failed: %v", err) + } + + res, err := svc.Search(ctx, SearchParams{Query: "indexing"}) + if err != nil { + t.Fatalf("Search after rebuild failed: %v", err) + } + if len(res.Matches) != 1 || res.Matches[0].SessionID != "sess-rebuild" { + t.Fatalf("expected sess-rebuild after index rebuild, got %#v", res.Matches) + } +} + +func TestIncrementalTailAfterSessionRewrite(t *testing.T) { + svc, tmpDir, _ := setupTestService(t) + ctx := context.Background() + + createTestSessionFile(t, tmpDir, "sess-rewrite", "/workspaces/app", []session.Message{ + {Role: "user", Content: "Version 1 message about architecture."}, + }) + + _, _ = svc.Indexer().SyncAll(ctx) + + // Verify v1 is found + res1, _ := svc.Search(ctx, SearchParams{Query: "architecture"}) + if len(res1.Matches) != 1 { + t.Fatalf("expected 1 match for architecture in v1") + } + + // Rewrite session with new message + time.Sleep(10 * time.Millisecond) // Ensure modTime changes + createTestSessionFile(t, tmpDir, "sess-rewrite", "/workspaces/app", []session.Message{ + {Role: "user", Content: "Version 2 message about telemetry and distributed tracing."}, + }) + + // Search for "telemetry" -> should trigger incremental sync and match immediately + res2, err := svc.Search(ctx, SearchParams{Query: "telemetry"}) + if err != nil { + t.Fatalf("Search v2 failed: %v", err) + } + if len(res2.Matches) != 1 || res2.Matches[0].SessionID != "sess-rewrite" { + t.Fatalf("expected match for rewritten telemetry message, got %#v", res2.Matches) + } + + // Old term should no longer match + resOld, _ := svc.Search(ctx, SearchParams{Query: "architecture"}) + if len(resOld.Matches) != 0 { + t.Fatalf("expected 0 matches for old content after rewrite, got %d", len(resOld.Matches)) + } +} + +func errorsIs(err, target error) bool { + if err == target { + return true + } + return strings.Contains(err.Error(), target.Error()) +} diff --git a/internal/tool/session_query.go b/internal/tool/session_query.go new file mode 100644 index 00000000..f0185894 --- /dev/null +++ b/internal/tool/session_query.go @@ -0,0 +1,123 @@ +package tool + +import ( + "context" + "encoding/json" + "fmt" + "os" + "strings" + + "github.com/GrayCodeAI/hawk/internal/sessionquery" +) + +// SessionQueryTool enables the agent to search past and current conversation sessions via FTS5. +type SessionQueryTool struct { + Service *sessionquery.Service +} + +func (SessionQueryTool) Name() string { return "SessionQuery" } +func (SessionQueryTool) Aliases() []string { return []string{"session_query", "session-query"} } +func (SessionQueryTool) Description() string { + return "Search conversation history across past and current sessions using SQLite full-text search." +} + +func (SessionQueryTool) Parameters() map[string]interface{} { + return map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "query": map[string]interface{}{ + "type": "string", + "description": "Full-text search query to find relevant past conversations", + }, + "session_id": map[string]interface{}{ + "type": "string", + "description": "Optional specific session ID to filter search within", + }, + "workspace": map[string]interface{}{ + "type": "string", + "description": "Optional workspace directory path to filter sessions within", + }, + "limit": map[string]interface{}{ + "type": "integer", + "description": "Maximum number of search results to return (default 10, max 50)", + }, + "offset": map[string]interface{}{ + "type": "integer", + "description": "Pagination offset for scrolling through results (default 0)", + }, + }, + "required": []string{"query"}, + } +} + +func (t SessionQueryTool) Execute(ctx context.Context, input json.RawMessage) (string, error) { + var p struct { + Query string `json:"query"` + SessionID string `json:"session_id"` + Workspace string `json:"workspace"` + Limit int `json:"limit"` + Offset int `json:"offset"` + } + if len(input) > 0 { + if err := json.Unmarshal(input, &p); err != nil { + return "", fmt.Errorf("invalid parameters: %w", err) + } + } + + if strings.TrimSpace(p.Query) == "" { + return "", fmt.Errorf("query parameter is required") + } + + svc := t.Service + if svc == nil { + var err error + svc, err = sessionquery.DefaultService() + if err != nil { + return "", fmt.Errorf("failed to initialize session query service: %w", err) + } + } + + cwd, _ := os.Getwd() + callerWs := cwd + if p.Workspace != "" { + callerWs = p.Workspace + } + + res, err := svc.Search(ctx, sessionquery.SearchParams{ + CallerWorkspace: callerWs, + Workspace: p.Workspace, + SessionID: p.SessionID, + Query: p.Query, + Limit: p.Limit, + Offset: p.Offset, + }) + if err != nil { + return "", err + } + + if len(res.Matches) == 0 { + return fmt.Sprintf("No matching messages found for query: %q", p.Query), nil + } + + var sb strings.Builder + sb.WriteString(fmt.Sprintf("Found %d match(es) for query %q (showing %d–%d):\n\n", + res.TotalCount, p.Query, res.Offset+1, res.Offset+len(res.Matches))) + + for _, m := range res.Matches { + sb.WriteString(fmt.Sprintf("### Session `%s` (Message #%d, Role: %s)\n", m.SessionID, m.MsgIndex, m.Role)) + if m.Workspace != "" { + sb.WriteString(fmt.Sprintf("**Workspace:** `%s`\n", m.Workspace)) + } + if m.Snippet != "" { + sb.WriteString(fmt.Sprintf("**Match:** %s\n\n", m.Snippet)) + } else { + sb.WriteString(fmt.Sprintf("**Content:** %s\n\n", m.Content)) + } + } + + if res.HasMore { + sb.WriteString(fmt.Sprintf("*(More results available. Use offset=%d to fetch the next page)*\n", res.Offset+len(res.Matches))) + } + + return strings.TrimRight(sb.String(), "\n"), nil +} diff --git a/internal/tool/session_query_test.go b/internal/tool/session_query_test.go new file mode 100644 index 00000000..e3e0696d --- /dev/null +++ b/internal/tool/session_query_test.go @@ -0,0 +1,84 @@ +package tool + +import ( + "context" + "encoding/json" + "os" + "path/filepath" + "strings" + "testing" + "time" + + "github.com/GrayCodeAI/hawk/internal/session" + "github.com/GrayCodeAI/hawk/internal/sessionquery" +) + +func TestSessionQueryTool_Execute(t *testing.T) { + tmpDir := t.TempDir() + t.Setenv("HAWK_STATE_DIR", tmpDir) + sessDir := filepath.Join(tmpDir, "sessions") + _ = os.MkdirAll(sessDir, 0o755) + + dbPath := filepath.Join(tmpDir, "tool_session_query.db") + svc, err := sessionquery.NewService(dbPath, sessDir) + if err != nil { + t.Fatalf("failed to create sessionquery service: %v", err) + } + defer func() { _ = svc.Close() }() + + sess := &session.Session{ + ID: "tool-sess-1", + Model: "claude-3-7-sonnet", + Provider: "anthropic", + CWD: tmpDir, + Messages: []session.Message{ + {Role: "user", Content: "How to configure OpenTelemetry collector with Prometheus?"}, + {Role: "assistant", Content: "Configure otel-collector.yaml with prometheus exporter."}, + }, + CreatedAt: time.Now(), + UpdatedAt: time.Now(), + } + if err := session.Save(sess); err != nil { + t.Fatalf("failed to save session: %v", err) + } + + qTool := SessionQueryTool{Service: svc} + ctx := context.Background() + + // 1. Search for OpenTelemetry + in, _ := json.Marshal(map[string]interface{}{ + "query": "OpenTelemetry", + "workspace": tmpDir, + }) + out, err := qTool.Execute(ctx, in) + if err != nil { + t.Fatalf("Execute failed: %v", err) + } + if !strings.Contains(out, "tool-sess-1") { + t.Fatalf("output %q should contain session ID tool-sess-1", out) + } + if !strings.Contains(out, "OpenTelemetry") { + t.Fatalf("output %q should contain OpenTelemetry", out) + } + + // 2. Empty query returns error + inEmpty, _ := json.Marshal(map[string]interface{}{ + "query": "", + }) + _, err = qTool.Execute(ctx, inEmpty) + if err == nil { + t.Fatal("expected error for empty query, got nil") + } + + // 3. No matches returns friendly response + inNoMatch, _ := json.Marshal(map[string]interface{}{ + "query": "nonexistenttermxyz123", + }) + outNoMatch, err := qTool.Execute(ctx, inNoMatch) + if err != nil { + t.Fatalf("Execute no match failed: %v", err) + } + if !strings.Contains(outNoMatch, "No matching messages found") { + t.Fatalf("expected 'No matching messages found', got %s", outNoMatch) + } +}