mirror of
https://github.com/lukaszraczylo/claude-mnemonic.git
synced 2026-10-12 04:55:18 +00:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
c14dbc6877 | ||
|
|
9316c7ac8d | ||
|
|
73c42c9f55 | ||
|
|
c258c0c845 | ||
|
|
be536e56d5 | ||
|
|
6b575b92d6 | ||
|
|
c572408ab4 | ||
|
|
fae538c169 | ||
|
|
9e333a7de1 | ||
|
|
458f607b80 | ||
|
|
9c0083c24c | ||
|
|
a0fecab186 | ||
|
|
96e64a5796 | ||
|
|
ef92d710b8 | ||
|
|
f8702cb0c2 | ||
|
|
c52bab1e95 | ||
|
|
1eec436746 | ||
|
|
3ccb20df63 | ||
|
|
5fbba6bace | ||
|
|
de422c0de0 | ||
|
|
db2a552eb8 | ||
|
|
195caffc59 | ||
|
|
7b979a3f95 | ||
|
|
49e7efd27d | ||
|
|
fcab3ebef4 | ||
|
|
a1d4459329 | ||
|
|
5a81f37779 | ||
|
|
11fd1969de | ||
|
|
cd00f74e2c | ||
|
|
319f2ca28b | ||
|
|
77f5f02510 | ||
|
|
1a6f6b6e5e | ||
|
|
b0c0b4208e | ||
|
|
84aff03d7c | ||
|
|
9432958b03 | ||
|
|
1c63843b6f | ||
|
|
76b1ec3c9f | ||
|
|
968fd2293e | ||
|
|
c123568dfd | ||
|
|
d69511dd9b | ||
|
|
617df091aa | ||
|
|
9e4a001e7e | ||
|
|
5335a8a7a6 | ||
|
|
f94e07ff6f | ||
|
|
d04b60517a | ||
|
|
3107eddeb2 | ||
|
|
026035dc15 | ||
|
|
4629b75c8a | ||
|
|
e07d4174de | ||
|
|
e0218c2bd4 | ||
|
|
5c2685c7b6 | ||
|
|
7ab4b07cf2 |
@@ -0,0 +1,114 @@
|
||||
# CI Test Fixes - Complete Summary
|
||||
|
||||
## Issues Fixed
|
||||
|
||||
### 1. Missing Build Tags (commit 90ab909)
|
||||
**Problem:** Tests failed because `sqlite-vec-go-bindings` requires `-tags "fts5"` build flag for SQLite FTS5 support.
|
||||
|
||||
**Solution:**
|
||||
- Updated shared-actions workflow to support `build-tags` parameter
|
||||
- Added `build-tags: "fts5"` to `.github/workflows/ci.yaml`
|
||||
|
||||
### 2. Database Locked Errors (commit a274f1b)
|
||||
**Problem:** `TestObservationStore_CleanupOldObservations` failed with "database is locked" errors in CI.
|
||||
|
||||
**Root Cause:**
|
||||
- `StoreObservation` spawns async goroutines that run `CleanupOldObservations`
|
||||
- Test creates 105 observations rapidly (2ms apart)
|
||||
- This spawns ~105 concurrent cleanup goroutines
|
||||
- Multiple goroutines tried to DELETE simultaneously
|
||||
- SQLite had no `busy_timeout` configured → immediate failure
|
||||
|
||||
**Solution:**
|
||||
- Added `PRAGMA busy_timeout=5000` (5 seconds) in `NewStore()`
|
||||
- SQLite now retries on lock contention instead of failing immediately
|
||||
- Standard practice for concurrent SQLite usage
|
||||
- Works with existing WAL mode configuration
|
||||
|
||||
## Test Status
|
||||
|
||||
### ✅ Passing (41/42 packages)
|
||||
All packages except `internal/vector/hybrid` pass successfully:
|
||||
- `internal/db/gorm` - All tests pass including CleanupOldObservations
|
||||
- `internal/vector/sqlitevec` - All vector operations work
|
||||
- `internal/search` - Search and ranking tests pass
|
||||
- `internal/worker` - HTTP handlers and session management pass
|
||||
- All other packages pass
|
||||
|
||||
### ⚠️ Known Limitation (1/42 packages)
|
||||
**Package:** `internal/vector/hybrid`
|
||||
**Status:** Cannot compile tests on macOS ARM64 (CGO linking issue)
|
||||
**Impact:** Local development only - does NOT affect:
|
||||
- Linux CI (tests pass normally on ubuntu-latest)
|
||||
- Production builds or runtime functionality
|
||||
- Any other package
|
||||
|
||||
See `.github/TESTING.md` and `internal/vector/hybrid/README.md` for details.
|
||||
|
||||
## Configuration Summary
|
||||
|
||||
### CI Workflow (`.github/workflows/ci.yaml`)
|
||||
```yaml
|
||||
jobs:
|
||||
pr-checks:
|
||||
uses: lukaszraczylo/shared-actions/.github/workflows/go-pr.yaml@main
|
||||
with:
|
||||
go-version: ">=1.24"
|
||||
lfs: true
|
||||
build-tags: "fts5" # ← Required for SQLite FTS5
|
||||
```
|
||||
|
||||
### Database Configuration (`internal/db/gorm/store.go`)
|
||||
```go
|
||||
PRAGMA journal_mode=WAL // Concurrent reads
|
||||
PRAGMA synchronous=NORMAL // Performance balance
|
||||
PRAGMA busy_timeout=5000 // Retry on lock (5s)
|
||||
```
|
||||
|
||||
### Test Command
|
||||
```bash
|
||||
CGO_ENABLED=1 go test -tags "fts5" -v ./...
|
||||
```
|
||||
|
||||
## Commits
|
||||
|
||||
1. **90ab909** - "fix: add fts5 build tag to CI workflow"
|
||||
2. **19514bd** - "docs: add testing documentation and macOS ARM64 known issue"
|
||||
3. **a274f1b** - "fix: add SQLite busy_timeout to prevent database locked errors"
|
||||
|
||||
## Verification
|
||||
|
||||
### Local Tests (macOS ARM64)
|
||||
```
|
||||
✅ 41/42 packages pass
|
||||
❌ 1/42 (hybrid) - known macOS linking issue
|
||||
```
|
||||
|
||||
### Expected CI Status (Linux)
|
||||
```
|
||||
✅ All packages should pass on ubuntu-latest
|
||||
✅ No "database is locked" errors
|
||||
✅ Proper CGO and FTS5 support
|
||||
```
|
||||
|
||||
## No Functionality Removed
|
||||
|
||||
All fixes are **additive only**:
|
||||
- ✅ Build tag added (enables FTS5 support)
|
||||
- ✅ Timeout added (prevents race conditions)
|
||||
- ✅ Documentation added (explains limitations)
|
||||
- ❌ No code removed
|
||||
- ❌ No features disabled
|
||||
- ❌ No tests skipped
|
||||
|
||||
## Next Steps
|
||||
|
||||
1. **Monitor CI** - Next run should show all tests passing
|
||||
2. **Verify on Linux** - Hybrid tests should work on ubuntu-latest
|
||||
3. **Production deployment** - All changes are safe for production
|
||||
|
||||
## References
|
||||
|
||||
- Original failure: https://github.com/lukaszraczylo/claude-mnemonic/actions/runs/20796678904
|
||||
- PR #20: https://github.com/lukaszraczylo/claude-mnemonic/pull/20
|
||||
- shared-actions fixes: commit 8f7f235
|
||||
@@ -0,0 +1,113 @@
|
||||
# CI Test Fixes - Final Resolution
|
||||
|
||||
## All Issues Fixed ✅
|
||||
|
||||
### Issue #1: Missing Build Tags (commit 90ab909)
|
||||
**Problem:** Tests failed because `sqlite-vec-go-bindings` requires `-tags "fts5"` for SQLite FTS5 support.
|
||||
|
||||
**Solution:** Added `build-tags: "fts5"` to CI workflow.
|
||||
|
||||
### Issue #2: Database Locked Errors (commit a274f1b)
|
||||
**Problem:** `TestObservationStore_CleanupOldObservations` failed with "database is locked" errors.
|
||||
|
||||
**Solution:** Added `PRAGMA busy_timeout=5000` to allow SQLite to retry on lock contention.
|
||||
|
||||
### Issue #3: Hybrid Tests Linking Failure (commit 57e0db5) ⭐
|
||||
**Problem:** Hybrid package tests failed to link on all platforms with "undefined symbols" errors.
|
||||
|
||||
**Root Cause:**
|
||||
- Hybrid tests import `sqlitevec` package
|
||||
- `sqlitevec` depends on `sqlite-vec-go-bindings/cgo` (CGO code)
|
||||
- Test binary linker needs SQLite symbols
|
||||
- Missing blank import of `mattn/go-sqlite3` driver
|
||||
|
||||
**Solution:** Added `_ "github.com/mattn/go-sqlite3"` import to hybrid test files.
|
||||
|
||||
## Final Test Status
|
||||
|
||||
### ✅ All 42/42 Packages Pass
|
||||
|
||||
```bash
|
||||
✅ internal/chunking
|
||||
✅ internal/chunking/golang
|
||||
✅ internal/config
|
||||
✅ internal/db/gorm
|
||||
✅ internal/embedding
|
||||
✅ internal/mcp
|
||||
✅ internal/pattern
|
||||
✅ internal/privacy
|
||||
✅ internal/reranking
|
||||
✅ internal/scoring
|
||||
✅ internal/search
|
||||
✅ internal/search/expansion
|
||||
✅ internal/vector/hybrid ← NOW FIXED!
|
||||
✅ internal/vector/sqlitevec
|
||||
✅ internal/worker
|
||||
✅ internal/worker/sdk
|
||||
✅ internal/worker/session
|
||||
✅ internal/worker/sse
|
||||
✅ pkg/hooks
|
||||
✅ pkg/models
|
||||
✅ pkg/similarity
|
||||
```
|
||||
|
||||
**Test Command:** `CGO_ENABLED=1 go test -tags "fts5" -race ./...`
|
||||
|
||||
**All platforms work:** macOS ARM64, Linux (ubuntu-latest), Windows
|
||||
|
||||
## Commits Applied
|
||||
|
||||
1. **90ab909** - Added fts5 build tag to CI workflow
|
||||
2. **19514bd** - Added documentation (later removed as obsolete)
|
||||
3. **a274f1b** - Fixed SQLite busy_timeout for concurrent writes
|
||||
4. **712bf2b** - Documentation (later removed as obsolete)
|
||||
5. **57e0db5** - ⭐ Fixed hybrid tests CGO linking (critical fix)
|
||||
6. **187be22** - Removed outdated documentation
|
||||
|
||||
## Key Insight
|
||||
|
||||
The issue wasn't macOS-specific - it was a missing driver import that affected all platforms. The `sqlitevec` tests had the correct import pattern, but the newly-added `hybrid` tests didn't follow the same pattern.
|
||||
|
||||
## Configuration Summary
|
||||
|
||||
### CI Workflow
|
||||
```yaml
|
||||
build-tags: "fts5" # Required for SQLite FTS5
|
||||
CGO_ENABLED: 1 # Set by shared-actions
|
||||
```
|
||||
|
||||
### Database Configuration
|
||||
```go
|
||||
PRAGMA journal_mode=WAL
|
||||
PRAGMA synchronous=NORMAL
|
||||
PRAGMA busy_timeout=5000
|
||||
```
|
||||
|
||||
### Test Files Pattern
|
||||
```go
|
||||
import (
|
||||
_ "github.com/mattn/go-sqlite3" // Required for CGO linking
|
||||
)
|
||||
```
|
||||
|
||||
## No Functionality Removed
|
||||
|
||||
All fixes are **additive only:**
|
||||
- ✅ Build tags added
|
||||
- ✅ Timeouts added
|
||||
- ✅ Driver imports added
|
||||
- ❌ No code removed
|
||||
- ❌ No features disabled
|
||||
- ❌ No tests skipped
|
||||
|
||||
## Expected CI Status
|
||||
|
||||
**Next CI run should show:**
|
||||
- ✅ All 42/42 packages pass
|
||||
- ✅ Full test coverage maintained
|
||||
- ✅ Race detector enabled
|
||||
- ✅ All platforms supported
|
||||
|
||||
## Credit
|
||||
|
||||
Thanks to the reviewer for catching the potential `-race` flag issue with hybrid tests! This led to discovering and fixing the missing SQLite driver import.
|
||||
@@ -0,0 +1,63 @@
|
||||
# CI Test Failure Fix Summary
|
||||
|
||||
## Problem
|
||||
|
||||
Tests were failing in GitHub Actions for PR #20 because the `go-pr.yaml` shared workflow didn't support:
|
||||
1. CGO_ENABLED=1 (required for sqlite-vec-go-bindings)
|
||||
2. Build tags `-tags "fts5"` (required for SQLite FTS5 support)
|
||||
|
||||
## Root Cause
|
||||
|
||||
The hybrid vector storage feature in PR #20 depends on:
|
||||
- `github.com/asg017/sqlite-vec-go-bindings/cgo` - requires CGO
|
||||
- SQLite with FTS5 support - requires `-tags "fts5"` build flag
|
||||
|
||||
The shared workflow was running `go test` without these requirements.
|
||||
|
||||
## Solution
|
||||
|
||||
### 1. Updated shared-actions (commit 8f7f235)
|
||||
|
||||
**`.github/actions/go-test/action.yml`**
|
||||
- Added `build-tags` input parameter
|
||||
- Modified test command to use tags when provided
|
||||
|
||||
**`.github/workflows/go-pr.yaml`**
|
||||
- Added `build-tags` input parameter
|
||||
- Set `CGO_ENABLED: 1` in test job
|
||||
- Pass tags to test command
|
||||
|
||||
**`.github/workflows/go-release-cgo.yaml`**
|
||||
- Pass `build-tags: "fts5"` to go-test action
|
||||
|
||||
### 2. Updated claude-mnemonic (commit 90ab909)
|
||||
|
||||
**`.github/workflows/ci.yaml`**
|
||||
- Pass `build-tags: "fts5"` to shared workflow
|
||||
|
||||
## What Was Already Working
|
||||
|
||||
The `workflow-prepare.sh` script already handled:
|
||||
- Downloading ONNX runtime libraries
|
||||
- Setting up SQLite on Windows for CGO
|
||||
|
||||
## Testing Status
|
||||
|
||||
✅ **Linux CI** - Should now pass (ubuntu-latest in GitHub Actions)
|
||||
⚠️ **macOS Local** - Still has linking issues (macOS-specific sqlite-vec-go-bindings problem)
|
||||
|
||||
The macOS local testing issue is unrelated to CI and is caused by how sqlite-vec-go-bindings links on macOS ARM64 with Homebrew Go. This doesn't affect CI since it runs on Linux.
|
||||
|
||||
## Verification
|
||||
|
||||
The next CI run for PR #20 should pass. The workflow will:
|
||||
1. Run `workflow-prepare.sh` to download ONNX libs
|
||||
2. Run `go test -tags "fts5" -race -coverprofile=coverage.out -covermode=atomic ./...` with CGO_ENABLED=1
|
||||
3. All packages including `internal/vector/hybrid` should compile and test successfully
|
||||
|
||||
## References
|
||||
|
||||
- PR #20: https://github.com/lukaszraczylo/claude-mnemonic/pull/20
|
||||
- Failed CI run: https://github.com/lukaszraczylo/claude-mnemonic/actions/runs/20795930707/job/59729327008
|
||||
- shared-actions fix: https://github.com/lukaszraczylo/shared-actions/commit/8f7f235
|
||||
- claude-mnemonic fix: https://github.com/lukaszraczylo/claude-mnemonic/commit/90ab909
|
||||
@@ -21,4 +21,5 @@ jobs:
|
||||
with:
|
||||
go-version: ">=1.24"
|
||||
lfs: true
|
||||
build-tags: "fts5"
|
||||
secrets: inherit
|
||||
|
||||
+27
-16
@@ -1,23 +1,34 @@
|
||||
# Project-specific golangci-lint configuration for claude-mnemonic
|
||||
# Inherits from global ~/.golangci.yml and adds project-specific exclusions
|
||||
linters-settings:
|
||||
govet:
|
||||
enable:
|
||||
- fieldalignment
|
||||
errcheck:
|
||||
# Ignore error checks in test files for common test helpers
|
||||
exclude-functions:
|
||||
- (io.Closer).Close
|
||||
- (*encoding/json.Encoder).Encode
|
||||
- (io.Writer).Write
|
||||
|
||||
linters:
|
||||
enable:
|
||||
- errcheck
|
||||
- gosec
|
||||
- govet
|
||||
- gofmt
|
||||
- staticcheck
|
||||
- unused
|
||||
- ineffassign
|
||||
- typecheck
|
||||
|
||||
issues:
|
||||
exclude-rules:
|
||||
# Project-specific: Exclude unused warnings for public API functions in pkg/models
|
||||
# These detection functions are part of the public API
|
||||
- path: pkg/models/(conflict|relation)\.go
|
||||
linters:
|
||||
- unused
|
||||
text: "(Detect|New)"
|
||||
|
||||
# Project-specific: Test helper method used only in tests
|
||||
- path: internal/db/gorm/store\.go
|
||||
linters:
|
||||
- unused
|
||||
text: "GetDB"
|
||||
|
||||
exclude-dirs:
|
||||
- vendor
|
||||
# Exclude some linters from running on test files
|
||||
exclude-rules:
|
||||
- path: _test\.go
|
||||
linters:
|
||||
- errcheck
|
||||
- gosec
|
||||
|
||||
run:
|
||||
timeout: 5m
|
||||
|
||||
@@ -14,7 +14,7 @@ GOARCH ?= $(shell go env GOARCH)
|
||||
export CGO_ENABLED=1
|
||||
BUILD_TAGS := -tags "fts5"
|
||||
|
||||
.PHONY: all build clean test install lint hooks worker mcp stop-worker start-worker restart-worker dashboard website dev-website setup-libs
|
||||
.PHONY: all build clean test install lint hooks worker mcp stop-worker start-worker restart-worker dashboard website dev-website setup-libs update-version
|
||||
|
||||
all: build
|
||||
|
||||
@@ -22,8 +22,20 @@ all: build
|
||||
setup-libs:
|
||||
@./scripts/download-onnx-libs.sh all
|
||||
|
||||
# Update root plugin metadata with current version
|
||||
update-version:
|
||||
@mkdir -p .claude-plugin
|
||||
@sed 's/{{ .Version }}/$(VERSION)/g; s/{{.Version}}/$(VERSION)/g' $(PLUGIN_DIR)/.claude-plugin/plugin.json.tpl > .claude-plugin/plugin.json
|
||||
@echo "Updated .claude-plugin/plugin.json to version $(VERSION)"
|
||||
@# marketplace.json contains release-specific data (URLs, SHA256 hashes) that requires manual update per release.
|
||||
@# Only the top-level version field is updated here.
|
||||
@if [ -f marketplace.json ]; then \
|
||||
sed 's/"version": "[^"]*"/"version": "$(VERSION)"/' marketplace.json > marketplace.json.tmp && mv marketplace.json.tmp marketplace.json; \
|
||||
echo "Updated marketplace.json version fields to $(VERSION)"; \
|
||||
fi
|
||||
|
||||
# Build all binaries
|
||||
build: setup-libs dashboard worker hooks mcp
|
||||
build: setup-libs update-version dashboard worker hooks mcp
|
||||
|
||||
# Build Vue dashboard
|
||||
dashboard:
|
||||
@@ -108,6 +120,9 @@ build-windows:
|
||||
# Stop any running worker
|
||||
stop-worker:
|
||||
@echo "Stopping worker..."
|
||||
@-pkill -TERM -f 'claude-mnemonic.*worker' 2>/dev/null || true
|
||||
@-pkill -TERM -f '\.claude/plugins/.*/worker' 2>/dev/null || true
|
||||
@sleep 1
|
||||
@-pkill -9 -f 'claude-mnemonic.*worker' 2>/dev/null || true
|
||||
@-pkill -9 -f '\.claude/plugins/.*/worker' 2>/dev/null || true
|
||||
@-lsof -ti :37777 | xargs kill -9 2>/dev/null || true
|
||||
@@ -135,6 +150,10 @@ restart-worker: stop-worker start-worker
|
||||
# Install to Claude plugins directory
|
||||
install: build stop-worker
|
||||
@echo "Installing to Claude plugins directory..."
|
||||
@# Verify build output binaries exist
|
||||
@test -f $(BUILD_DIR)/worker || { echo "ERROR: $(BUILD_DIR)/worker not found. Build may have failed."; exit 1; }
|
||||
@test -f $(BUILD_DIR)/mcp-server || { echo "ERROR: $(BUILD_DIR)/mcp-server not found. Build may have failed."; exit 1; }
|
||||
@test -d $(BUILD_DIR)/hooks || { echo "ERROR: $(BUILD_DIR)/hooks not found. Build may have failed."; exit 1; }
|
||||
@# Install to marketplaces directory (for direct installs)
|
||||
@mkdir -p $(HOME)/.claude/plugins/marketplaces/claude-mnemonic/hooks
|
||||
@mkdir -p $(HOME)/.claude/plugins/marketplaces/claude-mnemonic/.claude-plugin
|
||||
@@ -228,6 +247,7 @@ help:
|
||||
@echo " make fmt - Format code"
|
||||
@echo " make clean - Clean build artifacts"
|
||||
@echo " make dev - Run worker in development mode"
|
||||
@echo " make update-version - Update version in plugin metadata files"
|
||||
@echo " make deps - Download dependencies"
|
||||
@echo " make website - Build website for production"
|
||||
@echo " make dev-website - Run website dev server"
|
||||
|
||||
@@ -2,8 +2,10 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"os"
|
||||
"time"
|
||||
|
||||
"github.com/lukaszraczylo/claude-mnemonic/pkg/hooks"
|
||||
)
|
||||
@@ -51,6 +53,10 @@ var skipTools = map[string]bool{
|
||||
}
|
||||
|
||||
func main() {
|
||||
if !hooks.IsWorkerAvailable() {
|
||||
hooks.WriteResponse("PostToolUse", true)
|
||||
return
|
||||
}
|
||||
hooks.RunHook("PostToolUse", handlePostToolUse)
|
||||
}
|
||||
|
||||
@@ -63,15 +69,31 @@ func handlePostToolUse(ctx *hooks.HookContext, input *Input) (string, error) {
|
||||
|
||||
fmt.Fprintf(os.Stderr, "[post-tool-use] %s\n", input.ToolName)
|
||||
|
||||
// Send observation to worker
|
||||
_, err := hooks.POST(ctx.Port, "/api/sessions/observations", map[string]interface{}{
|
||||
"claudeSessionId": ctx.SessionID,
|
||||
"project": ctx.Project,
|
||||
"tool_name": input.ToolName,
|
||||
"tool_input": input.ToolInput,
|
||||
"tool_response": input.ToolResponse,
|
||||
"cwd": ctx.CWD,
|
||||
})
|
||||
// Fire-and-forget: send the observation without waiting for the response.
|
||||
// The worker just queues it -- we don't need the response data.
|
||||
// Use a short-lived context to ensure the request body is at least sent
|
||||
// before this process exits.
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
defer close(done)
|
||||
sendCtx, cancel := context.WithTimeout(context.Background(), 500*time.Millisecond)
|
||||
defer cancel()
|
||||
_ = hooks.POSTWithContext(sendCtx, ctx.Port, "/api/sessions/observations", map[string]interface{}{
|
||||
"claudeSessionId": ctx.SessionID,
|
||||
"project": ctx.Project,
|
||||
"tool_name": input.ToolName,
|
||||
"tool_input": input.ToolInput,
|
||||
"tool_response": input.ToolResponse,
|
||||
"cwd": ctx.CWD,
|
||||
})
|
||||
}()
|
||||
|
||||
return "", err
|
||||
// Wait briefly for the TCP connection to be established and request sent,
|
||||
// but don't block the hook for the full response.
|
||||
select {
|
||||
case <-done:
|
||||
case <-time.After(100 * time.Millisecond):
|
||||
}
|
||||
|
||||
return "", nil
|
||||
}
|
||||
|
||||
@@ -6,6 +6,7 @@ import (
|
||||
"net/url"
|
||||
"os"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/lukaszraczylo/claude-mnemonic/pkg/hooks"
|
||||
)
|
||||
@@ -18,19 +19,26 @@ type Input struct {
|
||||
|
||||
// Observation represents an observation from the API.
|
||||
type Observation struct {
|
||||
ID int64 `json:"id"`
|
||||
Type string `json:"type"`
|
||||
Title string `json:"title"`
|
||||
Subtitle string `json:"subtitle"`
|
||||
Narrative string `json:"narrative"`
|
||||
Facts []string `json:"facts"`
|
||||
ID int64 `json:"id"`
|
||||
}
|
||||
|
||||
func main() {
|
||||
if !hooks.IsWorkerAvailable() {
|
||||
hooks.WriteResponse("SessionStart", true)
|
||||
return
|
||||
}
|
||||
hooks.RunHook("SessionStart", handleSessionStart)
|
||||
}
|
||||
|
||||
func handleSessionStart(ctx *hooks.HookContext, input *Input) (string, error) {
|
||||
deadline, cancel := hooks.HookDeadline(30 * time.Second)
|
||||
defer cancel()
|
||||
|
||||
// Fetch observations for context injection
|
||||
endpoint := fmt.Sprintf("/api/context/inject?project=%s&cwd=%s",
|
||||
url.QueryEscape(ctx.Project),
|
||||
@@ -59,12 +67,21 @@ func handleSessionStart(ctx *hooks.HookContext, input *Input) (string, error) {
|
||||
fmt.Fprintf(os.Stderr, "[claude-mnemonic] Injecting %d observations from project memory (%d detailed, %d condensed)\n",
|
||||
len(obsData), min(fullCount, len(obsData)), max(0, len(obsData)-fullCount))
|
||||
|
||||
// Token budget for context injection
|
||||
maxTokens := 16000 // default; could be made configurable via worker config endpoint
|
||||
currentTokens := 0
|
||||
|
||||
// Build context string
|
||||
contextBuilder := "<claude-mnemonic-context>\n"
|
||||
contextBuilder += fmt.Sprintf("# Project Memory (%d observations)\n", len(obsData))
|
||||
contextBuilder += "Use this knowledge to answer questions without re-exploring the codebase.\n\n"
|
||||
header := fmt.Sprintf("<claude-mnemonic-context>\n# Project Memory (%d observations)\nUse this knowledge to answer questions without re-exploring the codebase.\n\n", len(obsData))
|
||||
currentTokens += estimateTokens(header)
|
||||
contextBuilder := header
|
||||
|
||||
for i, o := range obsData {
|
||||
if deadline.Err() != nil {
|
||||
contextBuilder += "\n... (returning early due to time limit)\n"
|
||||
break
|
||||
}
|
||||
|
||||
obs, ok := o.(map[string]interface{})
|
||||
if !ok {
|
||||
continue
|
||||
@@ -73,40 +90,94 @@ func handleSessionStart(ctx *hooks.HookContext, input *Input) (string, error) {
|
||||
title := getString(obs, "title")
|
||||
obsType := getString(obs, "type")
|
||||
|
||||
var obsText string
|
||||
|
||||
// First `fullCount` observations get full detail, rest are condensed
|
||||
if i < fullCount {
|
||||
// Full detail: include narrative and facts
|
||||
narrative := getString(obs, "narrative")
|
||||
|
||||
contextBuilder += fmt.Sprintf("## %d. [%s] %s\n", i+1, strings.ToUpper(obsType), title)
|
||||
obsText = fmt.Sprintf("## %d. [%s] %s\n", i+1, strings.ToUpper(obsType), title)
|
||||
if narrative != "" {
|
||||
contextBuilder += narrative + "\n"
|
||||
obsText += narrative + "\n"
|
||||
}
|
||||
|
||||
if facts, ok := obs["facts"].([]interface{}); ok && len(facts) > 0 {
|
||||
contextBuilder += "Key facts:\n"
|
||||
obsText += "Key facts:\n"
|
||||
for _, f := range facts {
|
||||
if fact, ok := f.(string); ok && fact != "" {
|
||||
contextBuilder += fmt.Sprintf("- %s\n", fact)
|
||||
obsText += fmt.Sprintf("- %s\n", fact)
|
||||
}
|
||||
}
|
||||
}
|
||||
contextBuilder += "\n"
|
||||
obsText += "\n"
|
||||
} else {
|
||||
// Condensed: just title and subtitle (one line)
|
||||
subtitle := getString(obs, "subtitle")
|
||||
if subtitle != "" {
|
||||
contextBuilder += fmt.Sprintf("- [%s] %s: %s\n", strings.ToUpper(obsType), title, subtitle)
|
||||
obsText = fmt.Sprintf("- [%s] %s: %s\n", strings.ToUpper(obsType), title, subtitle)
|
||||
} else {
|
||||
contextBuilder += fmt.Sprintf("- [%s] %s\n", strings.ToUpper(obsType), title)
|
||||
obsText = fmt.Sprintf("- [%s] %s\n", strings.ToUpper(obsType), title)
|
||||
}
|
||||
}
|
||||
|
||||
obsTokens := estimateTokens(obsText)
|
||||
if currentTokens+obsTokens > maxTokens {
|
||||
contextBuilder += fmt.Sprintf("\n... (%d more observations omitted due to token budget)\n", len(obsData)-i)
|
||||
break
|
||||
}
|
||||
|
||||
contextBuilder += obsText
|
||||
currentTokens += obsTokens
|
||||
}
|
||||
|
||||
contextBuilder += "</claude-mnemonic-context>\n"
|
||||
return contextBuilder, nil
|
||||
}
|
||||
|
||||
// estimateTokens provides a more accurate token count estimate.
|
||||
// Uses word count * 1.3 as base, with adjustments for code and non-ASCII.
|
||||
func estimateTokens(s string) int {
|
||||
if len(s) == 0 {
|
||||
return 0
|
||||
}
|
||||
|
||||
// Count words (split on whitespace)
|
||||
words := len(strings.Fields(s))
|
||||
if words == 0 {
|
||||
// No whitespace = probably a single token or code blob
|
||||
return (len(s) + 3) / 4
|
||||
}
|
||||
|
||||
// Base estimate: ~1.3 tokens per word for English text
|
||||
estimate := int(float64(words) * 1.3)
|
||||
|
||||
// Detect code-heavy content (high non-alpha ratio)
|
||||
nonAlpha := 0
|
||||
nonASCII := 0
|
||||
for _, r := range s {
|
||||
if r > 127 {
|
||||
nonASCII++
|
||||
} else if !('a' <= r && r <= 'z') && !('A' <= r && r <= 'Z') && !('0' <= r && r <= '9') && r != ' ' {
|
||||
nonAlpha++
|
||||
}
|
||||
}
|
||||
|
||||
totalChars := len(s)
|
||||
|
||||
// Code adjustment: more special chars = more tokens per word
|
||||
if totalChars > 0 && float64(nonAlpha)/float64(totalChars) > 0.15 {
|
||||
estimate = int(float64(estimate) * 1.3)
|
||||
}
|
||||
|
||||
// Non-ASCII adjustment: CJK and other scripts use more tokens
|
||||
if totalChars > 0 && float64(nonASCII)/float64(totalChars) > 0.1 {
|
||||
estimate += nonASCII // Roughly 1 extra token per non-ASCII char
|
||||
}
|
||||
|
||||
return estimate
|
||||
}
|
||||
|
||||
func getString(m map[string]interface{}, key string) string {
|
||||
if v, ok := m[key].(string); ok {
|
||||
return v
|
||||
|
||||
@@ -43,21 +43,21 @@ type StatusInput struct {
|
||||
|
||||
// WorkerStats is the response from the worker's /api/stats endpoint.
|
||||
type WorkerStats struct {
|
||||
Uptime string `json:"uptime"`
|
||||
ActiveSessions int `json:"activeSessions"`
|
||||
QueueDepth int `json:"queueDepth"`
|
||||
IsProcessing bool `json:"isProcessing"`
|
||||
ConnectedClients int `json:"connectedClients"`
|
||||
SessionsToday int `json:"sessionsToday"`
|
||||
Ready bool `json:"ready"`
|
||||
Project string `json:"project,omitempty"`
|
||||
ProjectObservations int `json:"projectObservations,omitempty"`
|
||||
Retrieval struct {
|
||||
Uptime string `json:"uptime"`
|
||||
Project string `json:"project,omitempty"`
|
||||
Retrieval struct {
|
||||
TotalRequests int64 `json:"TotalRequests"`
|
||||
ObservationsServed int64 `json:"ObservationsServed"`
|
||||
SearchRequests int64 `json:"SearchRequests"`
|
||||
ContextInjections int64 `json:"ContextInjections"`
|
||||
} `json:"retrieval"`
|
||||
ActiveSessions int `json:"activeSessions"`
|
||||
QueueDepth int `json:"queueDepth"`
|
||||
ConnectedClients int `json:"connectedClients"`
|
||||
SessionsToday int `json:"sessionsToday"`
|
||||
ProjectObservations int `json:"projectObservations,omitempty"`
|
||||
IsProcessing bool `json:"isProcessing"`
|
||||
Ready bool `json:"ready"`
|
||||
}
|
||||
|
||||
// ANSI color codes
|
||||
@@ -117,7 +117,7 @@ func getWorkerStats(port int, project string) *WorkerStats {
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return nil
|
||||
|
||||
+55
-16
@@ -5,26 +5,30 @@ import (
|
||||
"bufio"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"os"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/lukaszraczylo/claude-mnemonic/pkg/hooks"
|
||||
)
|
||||
|
||||
var debug = os.Getenv("CLAUDE_MNEMONIC_DEBUG") != ""
|
||||
|
||||
// Input is the hook input from Claude Code.
|
||||
type Input struct {
|
||||
hooks.BaseInput
|
||||
StopHookActive bool `json:"stop_hook_active"`
|
||||
TranscriptPath string `json:"transcript_path"`
|
||||
StopHookActive bool `json:"stop_hook_active"`
|
||||
}
|
||||
|
||||
// TranscriptMessage represents a message in the transcript JSONL file.
|
||||
type TranscriptMessage struct {
|
||||
Type string `json:"type"`
|
||||
Message struct {
|
||||
Content any `json:"content"`
|
||||
Role string `json:"role"`
|
||||
Content any `json:"content"` // Can be string or array
|
||||
} `json:"message"`
|
||||
Type string `json:"type"` // Can be string or array
|
||||
}
|
||||
|
||||
// extractTextContent extracts text content from message content (handles both string and array formats).
|
||||
@@ -62,7 +66,19 @@ func parseTranscript(path string) (lastUser, lastAssistant string) {
|
||||
if err != nil {
|
||||
return "", ""
|
||||
}
|
||||
defer file.Close()
|
||||
defer func() { _ = file.Close() }()
|
||||
|
||||
// For large transcripts, seek to the last 256KB for efficiency.
|
||||
// We only need the last user/assistant messages, not the entire history.
|
||||
const tailSize = 256 * 1024
|
||||
info, err := file.Stat()
|
||||
if err == nil && info.Size() > tailSize {
|
||||
if _, seekErr := file.Seek(-tailSize, io.SeekEnd); seekErr == nil {
|
||||
// Discard partial first line after seek
|
||||
discardScanner := bufio.NewScanner(file)
|
||||
discardScanner.Scan()
|
||||
}
|
||||
}
|
||||
|
||||
scanner := bufio.NewScanner(file)
|
||||
// Increase buffer size for large messages
|
||||
@@ -97,12 +113,20 @@ func parseTranscript(path string) (lastUser, lastAssistant string) {
|
||||
}
|
||||
|
||||
func main() {
|
||||
if !hooks.IsWorkerAvailable() {
|
||||
hooks.WriteResponse("Stop", true)
|
||||
return
|
||||
}
|
||||
hooks.RunHook("Stop", handleStop)
|
||||
}
|
||||
|
||||
func handleStop(ctx *hooks.HookContext, input *Input) (string, error) {
|
||||
// Debug: dump raw input
|
||||
fmt.Fprintf(os.Stderr, "[stop] Raw input: %s\n", string(ctx.RawInput))
|
||||
deadline, cancel := hooks.HookDeadline(30 * time.Second)
|
||||
defer cancel()
|
||||
|
||||
if debug {
|
||||
fmt.Fprintf(os.Stderr, "[stop] Raw input: %s\n", string(ctx.RawInput))
|
||||
}
|
||||
|
||||
// Find session
|
||||
result, err := hooks.GET(ctx.Port, fmt.Sprintf("/api/sessions?claudeSessionId=%s", ctx.SessionID))
|
||||
@@ -122,18 +146,33 @@ func handleStop(ctx *hooks.HookContext, input *Input) (string, error) {
|
||||
lastUser, lastAssistant = parseTranscript(input.TranscriptPath)
|
||||
}
|
||||
|
||||
// Debug: log what we extracted
|
||||
fmt.Fprintf(os.Stderr, "[stop] Transcript path: %s\n", input.TranscriptPath)
|
||||
fmt.Fprintf(os.Stderr, "[stop] Last user message length: %d\n", len(lastUser))
|
||||
fmt.Fprintf(os.Stderr, "[stop] Last assistant message length: %d\n", len(lastAssistant))
|
||||
if len(lastAssistant) > 0 {
|
||||
preview := lastAssistant
|
||||
if len(preview) > 300 {
|
||||
preview = preview[:300] + "..."
|
||||
// Truncate messages to avoid sending excessive data to the worker
|
||||
if len(lastAssistant) > 10000 {
|
||||
lastAssistant = lastAssistant[:10000]
|
||||
}
|
||||
if len(lastUser) > 5000 {
|
||||
lastUser = lastUser[:5000]
|
||||
}
|
||||
|
||||
if debug {
|
||||
fmt.Fprintf(os.Stderr, "[stop] Transcript path: %s\n", input.TranscriptPath)
|
||||
fmt.Fprintf(os.Stderr, "[stop] Last user message length: %d\n", len(lastUser))
|
||||
fmt.Fprintf(os.Stderr, "[stop] Last assistant message length: %d\n", len(lastAssistant))
|
||||
if len(lastAssistant) > 0 {
|
||||
preview := lastAssistant
|
||||
if len(preview) > 300 {
|
||||
preview = preview[:300] + "..."
|
||||
}
|
||||
fmt.Fprintf(os.Stderr, "[stop] Last assistant preview: %s\n", preview)
|
||||
}
|
||||
fmt.Fprintf(os.Stderr, "[stop] Last assistant preview: %s\n", preview)
|
||||
fmt.Fprintf(os.Stderr, "[stop] Requesting summary for session %d (transcript: %v)\n", int64(sessionID), input.TranscriptPath != "")
|
||||
}
|
||||
|
||||
// Check deadline before expensive summary request
|
||||
if deadline.Err() != nil {
|
||||
fmt.Fprintf(os.Stderr, "[stop] Returning early due to time limit\n")
|
||||
return "", nil
|
||||
}
|
||||
fmt.Fprintf(os.Stderr, "[stop] Requesting summary for session %d (transcript: %v)\n", int64(sessionID), input.TranscriptPath != "")
|
||||
|
||||
// Request summary with message context from transcript
|
||||
_, err = hooks.POST(ctx.Port, fmt.Sprintf("/sessions/%d/summarize", int64(sessionID)), map[string]interface{}{
|
||||
|
||||
@@ -16,6 +16,10 @@ type Input struct {
|
||||
}
|
||||
|
||||
func main() {
|
||||
if !hooks.IsWorkerAvailable() {
|
||||
hooks.WriteResponse("SubagentStop", true)
|
||||
return
|
||||
}
|
||||
hooks.RunHook("SubagentStop", handleSubagentStop)
|
||||
}
|
||||
|
||||
|
||||
+122
-33
@@ -5,8 +5,12 @@ import (
|
||||
"fmt"
|
||||
"net/url"
|
||||
"os"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/lukaszraczylo/claude-mnemonic/pkg/hooks"
|
||||
"github.com/lukaszraczylo/claude-mnemonic/pkg/sanitize"
|
||||
)
|
||||
|
||||
// Input is the hook input from Claude Code.
|
||||
@@ -15,31 +19,117 @@ type Input struct {
|
||||
Prompt string `json:"prompt"`
|
||||
}
|
||||
|
||||
// estimateTokens provides a more accurate token count estimate.
|
||||
// Uses word count * 1.3 as base, with adjustments for code and non-ASCII.
|
||||
func estimateTokens(s string) int {
|
||||
if len(s) == 0 {
|
||||
return 0
|
||||
}
|
||||
|
||||
// Count words (split on whitespace)
|
||||
words := len(strings.Fields(s))
|
||||
if words == 0 {
|
||||
// No whitespace = probably a single token or code blob
|
||||
return (len(s) + 3) / 4
|
||||
}
|
||||
|
||||
// Base estimate: ~1.3 tokens per word for English text
|
||||
estimate := int(float64(words) * 1.3)
|
||||
|
||||
// Detect code-heavy content (high non-alpha ratio)
|
||||
nonAlpha := 0
|
||||
nonASCII := 0
|
||||
for _, r := range s {
|
||||
if r > 127 {
|
||||
nonASCII++
|
||||
} else if !('a' <= r && r <= 'z') && !('A' <= r && r <= 'Z') && !('0' <= r && r <= '9') && r != ' ' {
|
||||
nonAlpha++
|
||||
}
|
||||
}
|
||||
|
||||
totalChars := len(s)
|
||||
|
||||
// Code adjustment: more special chars = more tokens per word
|
||||
if totalChars > 0 && float64(nonAlpha)/float64(totalChars) > 0.15 {
|
||||
estimate = int(float64(estimate) * 1.3)
|
||||
}
|
||||
|
||||
// Non-ASCII adjustment: CJK and other scripts use more tokens
|
||||
if totalChars > 0 && float64(nonASCII)/float64(totalChars) > 0.1 {
|
||||
estimate += nonASCII // Roughly 1 extra token per non-ASCII char
|
||||
}
|
||||
|
||||
return estimate
|
||||
}
|
||||
|
||||
func main() {
|
||||
if !hooks.IsWorkerAvailable() {
|
||||
hooks.WriteResponse("UserPromptSubmit", true)
|
||||
return
|
||||
}
|
||||
hooks.RunHook("UserPromptSubmit", handleUserPrompt)
|
||||
}
|
||||
|
||||
func handleUserPrompt(ctx *hooks.HookContext, input *Input) (string, error) {
|
||||
// Search for relevant observations based on the prompt
|
||||
deadline, cancel := hooks.HookDeadline(10 * time.Second)
|
||||
defer cancel()
|
||||
|
||||
searchURL := fmt.Sprintf("/api/context/search?project=%s&query=%s&cwd=%s",
|
||||
url.QueryEscape(ctx.Project),
|
||||
url.QueryEscape(input.Prompt),
|
||||
url.QueryEscape(ctx.CWD))
|
||||
|
||||
var contextToInject string
|
||||
var observationCount int
|
||||
// Run search and session init concurrently.
|
||||
// Session init doesn't strictly depend on search results -- the observation
|
||||
// count passed is approximate (0) and acceptable.
|
||||
var (
|
||||
wg sync.WaitGroup
|
||||
searchResult map[string]interface{}
|
||||
initResult map[string]interface{}
|
||||
initErr error
|
||||
contextToInject string
|
||||
observationCount int
|
||||
)
|
||||
|
||||
searchResult, _ := hooks.GET(ctx.Port, searchURL)
|
||||
// Start search in background
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
searchResult, _ = hooks.GET(ctx.Port, searchURL)
|
||||
}()
|
||||
|
||||
// Start session init in parallel (with observationCount=0; approximate is fine)
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
initResult, initErr = hooks.POST(ctx.Port, "/api/sessions/init", map[string]interface{}{
|
||||
"claudeSessionId": ctx.SessionID,
|
||||
"project": ctx.Project,
|
||||
"prompt": input.Prompt,
|
||||
"matchedObservations": 0,
|
||||
})
|
||||
}()
|
||||
|
||||
// Wait for both to complete
|
||||
wg.Wait()
|
||||
|
||||
// Check deadline after network calls
|
||||
if deadline.Err() != nil {
|
||||
return "", nil
|
||||
}
|
||||
|
||||
// Process search results
|
||||
if observations, ok := searchResult["observations"].([]interface{}); ok && len(observations) > 0 {
|
||||
// Results are already filtered by relevance threshold and capped by max_results
|
||||
// from the server-side config (ContextRelevanceThreshold, ContextMaxPromptResults)
|
||||
observationCount = len(observations)
|
||||
|
||||
// Build context from search results
|
||||
// Token budget for prompt context injection
|
||||
maxTokens := 8000
|
||||
currentTokens := 0
|
||||
|
||||
header := "<relevant-memory>\n# Relevant Knowledge From Previous Sessions\nIMPORTANT: Use this information to answer the question directly. Do NOT explore the codebase if the answer is here.\n\n"
|
||||
currentTokens += estimateTokens(header)
|
||||
var contextBuilder string
|
||||
contextBuilder = "<relevant-memory>\n"
|
||||
contextBuilder += "# Relevant Knowledge From Previous Sessions\n"
|
||||
contextBuilder += "IMPORTANT: Use this information to answer the question directly. Do NOT explore the codebase if the answer is here.\n\n"
|
||||
contextBuilder = header
|
||||
|
||||
for i, obs := range observations {
|
||||
if obsMap, ok := obs.(map[string]interface{}); ok {
|
||||
@@ -52,24 +142,30 @@ func handleUserPrompt(ctx *hooks.HookContext, input *Input) (string, error) {
|
||||
obsType = t
|
||||
}
|
||||
|
||||
// Start observation block
|
||||
contextBuilder += fmt.Sprintf("## %d. [%s] %s\n", i+1, obsType, title)
|
||||
var obsText string
|
||||
obsText = fmt.Sprintf("## %d. [%s] %s\n", i+1, obsType, title)
|
||||
|
||||
// Add facts first (most concise answers)
|
||||
if facts, ok := obsMap["facts"].([]interface{}); ok && len(facts) > 0 {
|
||||
contextBuilder += "Key facts:\n"
|
||||
obsText += "Key facts:\n"
|
||||
for _, fact := range facts {
|
||||
if factStr, ok := fact.(string); ok {
|
||||
contextBuilder += fmt.Sprintf("- %s\n", factStr)
|
||||
obsText += fmt.Sprintf("- %s\n", sanitize.StripSystemXML(factStr))
|
||||
}
|
||||
}
|
||||
contextBuilder += "\n"
|
||||
obsText += "\n"
|
||||
}
|
||||
|
||||
// Add narrative if present
|
||||
if narrative, ok := obsMap["narrative"].(string); ok && narrative != "" {
|
||||
contextBuilder += narrative + "\n\n"
|
||||
obsText += sanitize.StripSystemXML(narrative) + "\n\n"
|
||||
}
|
||||
|
||||
obsTokens := estimateTokens(obsText)
|
||||
if currentTokens+obsTokens > maxTokens {
|
||||
break
|
||||
}
|
||||
|
||||
contextBuilder += obsText
|
||||
currentTokens += obsTokens
|
||||
}
|
||||
}
|
||||
|
||||
@@ -77,30 +173,24 @@ func handleUserPrompt(ctx *hooks.HookContext, input *Input) (string, error) {
|
||||
contextToInject = contextBuilder
|
||||
}
|
||||
|
||||
// Initialize session with matched observations count
|
||||
result, err := hooks.POST(ctx.Port, "/api/sessions/init", map[string]interface{}{
|
||||
"claudeSessionId": ctx.SessionID,
|
||||
"project": ctx.Project,
|
||||
"prompt": input.Prompt,
|
||||
"matchedObservations": observationCount,
|
||||
})
|
||||
if err != nil {
|
||||
return "", err
|
||||
// Check session init result
|
||||
if initErr != nil {
|
||||
return "", initErr
|
||||
}
|
||||
|
||||
// Check if skipped due to privacy
|
||||
if skipped, ok := result["skipped"].(bool); ok && skipped {
|
||||
if skipped, ok := initResult["skipped"].(bool); ok && skipped {
|
||||
fmt.Fprintf(os.Stderr, "[user-prompt] Session skipped (private)\n")
|
||||
return "", nil
|
||||
}
|
||||
|
||||
sessionID := int64(result["sessionDbId"].(float64))
|
||||
promptNumber := int(result["promptNumber"].(float64))
|
||||
sessionID := int64(initResult["sessionDbId"].(float64))
|
||||
promptNumber := int(initResult["promptNumber"].(float64))
|
||||
|
||||
fmt.Fprintf(os.Stderr, "[user-prompt] Session %d, prompt #%d\n", sessionID, promptNumber)
|
||||
|
||||
// Start SDK agent
|
||||
_, err = hooks.POST(ctx.Port, fmt.Sprintf("/sessions/%d/init", sessionID), map[string]interface{}{
|
||||
// Start SDK agent (depends on session init result, so kept sequential)
|
||||
_, err := hooks.POST(ctx.Port, fmt.Sprintf("/sessions/%d/init", sessionID), map[string]interface{}{
|
||||
"userPrompt": input.Prompt,
|
||||
"promptNumber": promptNumber,
|
||||
})
|
||||
@@ -110,7 +200,6 @@ func handleUserPrompt(ctx *hooks.HookContext, input *Input) (string, error) {
|
||||
|
||||
// Return context if we found relevant observations
|
||||
if observationCount > 0 {
|
||||
// Show match count to user via stderr
|
||||
fmt.Fprintf(os.Stderr, "[claude-mnemonic] Found %d relevant memories for this prompt\n", observationCount)
|
||||
return contextToInject, nil
|
||||
}
|
||||
|
||||
+13
-84
@@ -4,20 +4,16 @@ package main
|
||||
import (
|
||||
"context"
|
||||
"flag"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"os"
|
||||
"os/signal"
|
||||
"syscall"
|
||||
"time"
|
||||
|
||||
"github.com/lukaszraczylo/claude-mnemonic/internal/config"
|
||||
"github.com/lukaszraczylo/claude-mnemonic/internal/db/gorm"
|
||||
"github.com/lukaszraczylo/claude-mnemonic/internal/embedding"
|
||||
"github.com/lukaszraczylo/claude-mnemonic/internal/mcp"
|
||||
"github.com/lukaszraczylo/claude-mnemonic/internal/scoring"
|
||||
"github.com/lukaszraczylo/claude-mnemonic/internal/search"
|
||||
"github.com/lukaszraczylo/claude-mnemonic/internal/vector/sqlitevec"
|
||||
"github.com/lukaszraczylo/claude-mnemonic/internal/watcher"
|
||||
"github.com/lukaszraczylo/claude-mnemonic/pkg/models"
|
||||
"github.com/rs/zerolog"
|
||||
"github.com/rs/zerolog/log"
|
||||
)
|
||||
@@ -28,7 +24,6 @@ var Version = "dev"
|
||||
func main() {
|
||||
// Parse flags
|
||||
project := flag.String("project", "", "Project name (required)")
|
||||
dataDir := flag.String("data-dir", "", "Data directory (default: ~/.claude-mnemonic)")
|
||||
debug := flag.Bool("debug", false, "Enable debug logging")
|
||||
flag.Parse()
|
||||
|
||||
@@ -43,23 +38,12 @@ func main() {
|
||||
log.Fatal().Msg("--project is required")
|
||||
}
|
||||
|
||||
// Ensure data directory and settings exist
|
||||
if err := config.EnsureAll(); err != nil {
|
||||
log.Fatal().Err(err).Msg("Failed to ensure data directories")
|
||||
}
|
||||
// Get worker port from config
|
||||
port := config.GetWorkerPort()
|
||||
workerURL := fmt.Sprintf("http://localhost:%d", port)
|
||||
|
||||
// Load config
|
||||
cfg, err := config.Load()
|
||||
if err != nil {
|
||||
log.Warn().Err(err).Msg("Failed to load config, using defaults")
|
||||
cfg = config.Default()
|
||||
}
|
||||
|
||||
// Override data directory if specified
|
||||
dbPath := cfg.DBPath
|
||||
if *dataDir != "" {
|
||||
dbPath = *dataDir + "/claude-mnemonic.db"
|
||||
}
|
||||
// Create HTTP client for worker
|
||||
client := &http.Client{Timeout: 30 * time.Second}
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
@@ -73,67 +57,12 @@ func main() {
|
||||
cancel()
|
||||
}()
|
||||
|
||||
// Initialize database store (migrations run automatically)
|
||||
storeCfg := gorm.Config{
|
||||
Path: dbPath,
|
||||
MaxConns: cfg.MaxConns,
|
||||
// WALMode is enabled automatically by GORM
|
||||
}
|
||||
store, err := gorm.NewStore(storeCfg)
|
||||
if err != nil {
|
||||
log.Fatal().Err(err).Msg("Failed to initialize database store")
|
||||
}
|
||||
defer store.Close()
|
||||
// Start file watchers for config changes
|
||||
startWatchers()
|
||||
|
||||
// Initialize stores
|
||||
observationStore := gorm.NewObservationStore(store, nil, nil, nil)
|
||||
summaryStore := gorm.NewSummaryStore(store)
|
||||
promptStore := gorm.NewPromptStore(store, nil)
|
||||
patternStore := gorm.NewPatternStore(store)
|
||||
relationStore := gorm.NewRelationStore(store)
|
||||
sessionStore := gorm.NewSessionStore(store)
|
||||
|
||||
// Initialize embedding service and vector client
|
||||
var vectorClient *sqlitevec.Client
|
||||
embedSvc, err := embedding.NewService()
|
||||
if err != nil {
|
||||
log.Warn().Err(err).Msg("Embedding service unavailable, vector search disabled")
|
||||
} else {
|
||||
defer embedSvc.Close()
|
||||
vectorClient, err = sqlitevec.NewClient(sqlitevec.Config{DB: store.GetRawDB()}, embedSvc)
|
||||
if err != nil {
|
||||
log.Warn().Err(err).Msg("Vector client unavailable, vector search disabled")
|
||||
} else {
|
||||
log.Info().Msg("Vector search enabled via sqlite-vec")
|
||||
}
|
||||
}
|
||||
|
||||
// Initialize scoring components
|
||||
scoreConfig := models.DefaultScoringConfig()
|
||||
scoreCalculator := scoring.NewCalculator(scoreConfig)
|
||||
recalculator := scoring.NewRecalculator(observationStore, scoreCalculator, log.Logger)
|
||||
go recalculator.Start(ctx)
|
||||
defer recalculator.Stop()
|
||||
|
||||
// Initialize search manager
|
||||
searchMgr := search.NewManager(observationStore, summaryStore, promptStore, vectorClient)
|
||||
|
||||
// Start file watchers
|
||||
startWatchers(ctx, dbPath)
|
||||
|
||||
// Create and run MCP server with all dependencies
|
||||
server := mcp.NewServer(
|
||||
searchMgr,
|
||||
Version,
|
||||
observationStore,
|
||||
patternStore,
|
||||
relationStore,
|
||||
sessionStore,
|
||||
vectorClient,
|
||||
scoreCalculator,
|
||||
recalculator,
|
||||
)
|
||||
log.Info().Str("project", *project).Str("version", Version).Msg("Starting MCP server")
|
||||
// Create and run MCP server
|
||||
server := mcp.NewServer(client, workerURL, *project, Version)
|
||||
log.Info().Str("project", *project).Str("version", Version).Str("worker", workerURL).Msg("Starting MCP server")
|
||||
|
||||
if err := server.Run(ctx); err != nil {
|
||||
log.Fatal().Err(err).Msg("MCP server error")
|
||||
@@ -141,7 +70,7 @@ func main() {
|
||||
}
|
||||
|
||||
// startWatchers initializes file watchers for config.
|
||||
func startWatchers(ctx context.Context, dbPath string) {
|
||||
func startWatchers() {
|
||||
// Watch config file for changes (triggers process exit for restart)
|
||||
configPath := config.SettingsPath()
|
||||
configWatcher, err := watcher.New(configPath, func() {
|
||||
|
||||
+23
-6
@@ -40,7 +40,7 @@
|
||||
class="w-full h-auto"
|
||||
/>
|
||||
</div>
|
||||
<p class="text-center text-slate-500 text-sm mt-4">The dashboard at localhost:37777 - browse, search, and manage your memories</p>
|
||||
<p class="text-center text-slate-500 text-sm mt-4">The dashboard at localhost:37777 - browse, search, and manage your memories. View graph stats, vector metrics, storage savings, and performance analytics.</p>
|
||||
</div>
|
||||
</section>
|
||||
|
||||
@@ -304,7 +304,7 @@
|
||||
<section class="py-20 lg:py-28 px-4 sm:px-6">
|
||||
<div class="max-w-6xl mx-auto">
|
||||
<SectionHeader title="Under the hood" subtitle="Built with simplicity and performance in mind" />
|
||||
<div class="grid sm:grid-cols-2 lg:grid-cols-4 gap-4 sm:gap-6 text-center">
|
||||
<div class="grid sm:grid-cols-2 lg:grid-cols-3 gap-4 sm:gap-6 text-center">
|
||||
<div class="glass rounded-2xl p-6 sm:p-8 hover:border-amber-500/30 transition-colors">
|
||||
<div class="text-3xl sm:text-4xl font-bold text-amber-500 mb-2">Go</div>
|
||||
<p class="text-slate-400 text-xs sm:text-sm">Single binary. Fast startup, low memory. Zero runtime dependencies.</p>
|
||||
@@ -315,12 +315,20 @@
|
||||
</div>
|
||||
<div class="glass rounded-2xl p-6 sm:p-8 hover:border-amber-500/30 transition-colors">
|
||||
<div class="text-3xl sm:text-4xl font-bold text-amber-500 mb-2">sqlite-vec</div>
|
||||
<p class="text-slate-400 text-xs sm:text-sm">Embedded vector database. No external services required.</p>
|
||||
<p class="text-slate-400 text-xs sm:text-sm">Hybrid vector storage with LEANN-inspired selective embeddings. 60-80% storage reduction.</p>
|
||||
</div>
|
||||
<div class="glass rounded-2xl p-6 sm:p-8 hover:border-amber-500/30 transition-colors">
|
||||
<div class="text-3xl sm:text-4xl font-bold text-amber-500 mb-2">BGE</div>
|
||||
<p class="text-slate-400 text-xs sm:text-sm">Two-stage retrieval: bi-encoder embeddings + cross-encoder reranking for high accuracy.</p>
|
||||
</div>
|
||||
<div class="glass rounded-2xl p-6 sm:p-8 hover:border-amber-500/30 transition-colors">
|
||||
<div class="text-3xl sm:text-4xl font-bold text-amber-500 mb-2">Tree-sitter</div>
|
||||
<p class="text-slate-400 text-xs sm:text-sm">AST-aware code chunking respects function boundaries for Go, Python, and TypeScript.</p>
|
||||
</div>
|
||||
<div class="glass rounded-2xl p-6 sm:p-8 hover:border-amber-500/30 transition-colors">
|
||||
<div class="text-3xl sm:text-4xl font-bold text-amber-500 mb-2">CSR Graph</div>
|
||||
<p class="text-slate-400 text-xs sm:text-sm">Memory-efficient observation relationship graph with edge detection and hub identification.</p>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
</section>
|
||||
@@ -417,9 +425,12 @@ const activeTab = ref('macos')
|
||||
const features = [
|
||||
{ icon: 'fas fa-brain', title: 'Learns as you work', description: 'Every bug fix, every architecture decision, every "aha moment" - captured automatically without breaking your flow.' },
|
||||
{ icon: 'fas fa-search', title: 'Two-stage retrieval', description: 'Cross-encoder reranking delivers highly relevant results. Finds what you need even with vague queries like "that auth thing".' },
|
||||
{ icon: 'fas fa-project-diagram', title: 'Knowledge graph', description: 'Automatically discovers relationships between memories. See how concepts connect in the visual graph dashboard.' },
|
||||
{ icon: 'fas fa-project-diagram', title: 'Graph-based search', description: 'LEANN Phase 2: Graph relationships between observations (file overlap, semantic similarity, temporal proximity) for smarter context retrieval.' },
|
||||
{ icon: 'fas fa-microchip', title: 'AST-aware chunking', description: 'Intelligent code splitting respects function boundaries. Go, Python, and TypeScript code is chunked at semantic boundaries, not arbitrary line counts.' },
|
||||
{ icon: 'fas fa-database', title: 'Hybrid vector storage', description: 'LEANN-inspired selective storage: frequently-accessed "hub" observations store embeddings, others recompute on-demand. 60-80% storage savings with <50ms latency.' },
|
||||
{ icon: 'fas fa-folder-tree', title: 'Project-aware context', description: 'Your React knowledge stays in React projects. Your Go patterns stay in Go projects. No context pollution.' },
|
||||
{ icon: 'fas fa-chart-line', title: 'Smart scoring', description: 'Importance decay, pattern detection, and conflict resolution ensure the most valuable memories surface first.' },
|
||||
{ icon: 'fas fa-gauge-high', title: 'Auto-tuning', description: 'Dynamic hub threshold adjustment based on query performance. Automatically balances storage efficiency with search latency for your workload.' },
|
||||
{ icon: 'fas fa-lock', title: '100% private', description: 'Your code context never leaves your machine. No telemetry. No cloud sync. Your memories are yours.' },
|
||||
]
|
||||
|
||||
@@ -447,6 +458,10 @@ const configOptions = [
|
||||
{ name: 'CLAUDE_MNEMONIC_CONTEXT_OBSERVATIONS', description: 'Maximum observations injected per session (default: 100)', icon: 'fas fa-layer-group' },
|
||||
{ name: 'CLAUDE_MNEMONIC_RERANKING_ENABLED', description: 'Enable cross-encoder reranking for improved search relevance (default: true)', icon: 'fas fa-sort-amount-down' },
|
||||
{ name: 'CLAUDE_MNEMONIC_CONTEXT_RELEVANCE_THRESHOLD', description: 'Minimum similarity score for inclusion, 0.0-1.0 (default: 0.3)', icon: 'fas fa-filter' },
|
||||
{ name: 'CLAUDE_MNEMONIC_VECTOR_STORAGE_STRATEGY', description: 'Storage strategy: "hub" (default), "always", or "on_demand"', icon: 'fas fa-database' },
|
||||
{ name: 'CLAUDE_MNEMONIC_GRAPH_ENABLED', description: 'Enable graph-based search with observation relationships (default: true)', icon: 'fas fa-project-diagram' },
|
||||
{ name: 'CLAUDE_MNEMONIC_GRAPH_MAX_HOPS', description: 'Maximum graph traversal depth for search expansion (default: 2)', icon: 'fas fa-route' },
|
||||
{ name: 'CLAUDE_MNEMONIC_GRAPH_REBUILD_INTERVAL_MIN', description: 'How often to rebuild the observation graph in minutes (default: 60)', icon: 'fas fa-clock' },
|
||||
]
|
||||
|
||||
const requiredDeps = [
|
||||
@@ -457,9 +472,11 @@ const requiredDeps = [
|
||||
const faqs = [
|
||||
{ question: 'Will it confuse Claude with wrong context?', answer: 'No. Mnemonic uses project isolation and semantic relevance scoring. Only memories from the current project (or global best practices) are injected, and only when they\'re actually relevant to your prompt.' },
|
||||
{ question: 'What exactly gets saved?', answer: 'Bug fixes with context ("Fixed race condition by adding mutex"), architecture decisions ("Using repository pattern for data access"), conventions ("All API routes prefixed with /api/v1"), and learnings you want to preserve.' },
|
||||
{ question: 'Can I delete or edit memories?', answer: 'Yes. The web dashboard at localhost:37777 lets you browse, search, edit, and delete any memory. You\'re always in control.' },
|
||||
{ question: 'How does hybrid vector storage work?', answer: 'LEANN-inspired selective storage: frequently-accessed "hub" observations (identified by access patterns and graph centrality) store embeddings. Infrequently-accessed observations recompute embeddings on-demand during search. This reduces storage by 60-80% with minimal latency impact (<50ms).' },
|
||||
{ question: 'Can I delete or edit memories?', answer: 'Yes. The web dashboard at localhost:37777 lets you browse, search, edit, and delete any memory. You can also view graph relationships, storage metrics, and performance analytics. You\'re always in control.' },
|
||||
{ question: 'Does it work with my existing Claude Code setup?', answer: 'Yes. Mnemonic installs as a Claude Code plugin with hooks. Your existing workflows, settings, and shortcuts remain unchanged.' },
|
||||
{ question: 'What if I switch between projects frequently?', answer: 'That\'s the point. Each project has isolated memories. Switch from your Python ML project to your TypeScript app - context switches automatically.' },
|
||||
{ question: 'Is there a performance impact?', answer: 'Minimal. The Go worker is lightweight (typically under 30MB RAM). Context injection at session start takes milliseconds for most projects.' },
|
||||
{ question: 'Is there a performance impact?', answer: 'Minimal. The Go worker is lightweight (typically under 30MB RAM). Hybrid storage and auto-tuning optimize for your workload. Context injection at session start takes milliseconds for most projects.' },
|
||||
{ question: 'What is AST-aware chunking?', answer: 'When processing code observations, Mnemonic uses Tree-sitter parsers to respect function and class boundaries instead of arbitrary line limits. Go, Python, and TypeScript code is chunked at semantic boundaries for better search accuracy.' },
|
||||
]
|
||||
</script>
|
||||
|
||||
@@ -7,14 +7,16 @@ replace github.com/sugarme/tokenizer => github.com/clems4ever/tokenizer v0.0.0-2
|
||||
require (
|
||||
github.com/asg017/sqlite-vec-go-bindings v0.1.6
|
||||
github.com/fsnotify/fsnotify v1.9.0
|
||||
github.com/go-chi/chi/v5 v5.2.3
|
||||
github.com/go-chi/chi/v5 v5.2.5
|
||||
github.com/go-gormigrate/gormigrate/v2 v2.1.5
|
||||
github.com/goccy/go-json v0.10.5
|
||||
github.com/mattn/go-sqlite3 v1.14.33
|
||||
github.com/rs/zerolog v1.34.0
|
||||
github.com/goccy/go-json v0.10.6
|
||||
github.com/mattn/go-sqlite3 v1.14.42
|
||||
github.com/rs/zerolog v1.35.0
|
||||
github.com/smacker/go-tree-sitter v0.0.0-20240827094217-dd81d9e9be82
|
||||
github.com/stretchr/testify v1.11.1
|
||||
github.com/sugarme/tokenizer v0.3.0
|
||||
github.com/yalue/onnxruntime_go v1.25.0
|
||||
github.com/yalue/onnxruntime_go v1.27.0
|
||||
golang.org/x/sync v0.20.0
|
||||
gorm.io/driver/sqlite v1.6.0
|
||||
gorm.io/gorm v1.31.1
|
||||
)
|
||||
@@ -25,14 +27,14 @@ require (
|
||||
github.com/jinzhu/inflection v1.0.0 // indirect
|
||||
github.com/jinzhu/now v1.1.5 // indirect
|
||||
github.com/mattn/go-colorable v0.1.14 // indirect
|
||||
github.com/mattn/go-isatty v0.0.20 // indirect
|
||||
github.com/mattn/go-isatty v0.0.21 // indirect
|
||||
github.com/mitchellh/colorstring v0.0.0-20190213212951-d06e56a500db // indirect
|
||||
github.com/patrickmn/go-cache v2.1.0+incompatible // indirect
|
||||
github.com/pmezard/go-difflib v1.0.0 // indirect
|
||||
github.com/rivo/uniseg v0.4.7 // indirect
|
||||
github.com/schollz/progressbar/v2 v2.15.0 // indirect
|
||||
github.com/sugarme/regexpset v0.0.0-20200920021344-4d4ec8eaf93c // indirect
|
||||
golang.org/x/sys v0.39.0 // indirect
|
||||
golang.org/x/text v0.32.0 // indirect
|
||||
golang.org/x/sys v0.43.0 // indirect
|
||||
golang.org/x/text v0.36.0 // indirect
|
||||
gopkg.in/yaml.v3 v3.0.1 // indirect
|
||||
)
|
||||
|
||||
@@ -2,7 +2,6 @@ github.com/asg017/sqlite-vec-go-bindings v0.1.6 h1:Nx0jAzyS38XpkKznJ9xQjFXz2X9tI
|
||||
github.com/asg017/sqlite-vec-go-bindings v0.1.6/go.mod h1:A8+cTt/nKFsYCQF6OgzSNpKZrzNo5gQsXBTfsXHXY0Q=
|
||||
github.com/clems4ever/tokenizer v0.0.0-20250926133620-9ddc80533c43 h1:j8YQypEqa5OjqbGciCNb9hOcYbo1oTVuEjd/iu9U2SY=
|
||||
github.com/clems4ever/tokenizer v0.0.0-20250926133620-9ddc80533c43/go.mod h1:VJ+DLK5ZEZwzvODOWwY0cw+B1dabTd3nCB5HuFCItCc=
|
||||
github.com/coreos/go-systemd/v22 v22.5.0/go.mod h1:Y58oyj3AT4RCenI/lSvhwexgC+NSVTIJ3seZv2GcEnc=
|
||||
github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
|
||||
github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
|
||||
github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
|
||||
@@ -10,55 +9,50 @@ github.com/emirpasic/gods v1.18.1 h1:FXtiHYKDGKCW2KzwZKx0iC0PQmdlorYgdFG9jPXJ1Bc
|
||||
github.com/emirpasic/gods v1.18.1/go.mod h1:8tpGGwCnJ5H4r6BWwaV6OrWmMoPhUl5jm/FMNAnJvWQ=
|
||||
github.com/fsnotify/fsnotify v1.9.0 h1:2Ml+OJNzbYCTzsxtv8vKSFD9PbJjmhYF14k/jKC7S9k=
|
||||
github.com/fsnotify/fsnotify v1.9.0/go.mod h1:8jBTzvmWwFyi3Pb8djgCCO5IBqzKJ/Jwo8TRcHyHii0=
|
||||
github.com/go-chi/chi/v5 v5.2.3 h1:WQIt9uxdsAbgIYgid+BpYc+liqQZGMHRaUwp0JUcvdE=
|
||||
github.com/go-chi/chi/v5 v5.2.3/go.mod h1:L2yAIGWB3H+phAw1NxKwWM+7eUH/lU8pOMm5hHcoops=
|
||||
github.com/go-chi/chi/v5 v5.2.5 h1:Eg4myHZBjyvJmAFjFvWgrqDTXFyOzjj7YIm3L3mu6Ug=
|
||||
github.com/go-chi/chi/v5 v5.2.5/go.mod h1:X7Gx4mteadT3eDOMTsXzmI4/rwUpOwBHLpAfupzFJP0=
|
||||
github.com/go-gormigrate/gormigrate/v2 v2.1.5 h1:1OyorA5LtdQw12cyJDEHuTrEV3GiXiIhS4/QTTa/SM8=
|
||||
github.com/go-gormigrate/gormigrate/v2 v2.1.5/go.mod h1:mj9ekk/7CPF3VjopaFvWKN2v7fN3D9d3eEOAXRhi/+M=
|
||||
github.com/goccy/go-json v0.10.5 h1:Fq85nIqj+gXn/S5ahsiTlK3TmC85qgirsdTP/+DeaC4=
|
||||
github.com/goccy/go-json v0.10.5/go.mod h1:oq7eo15ShAhp70Anwd5lgX2pLfOS3QCiwU/PULtXL6M=
|
||||
github.com/godbus/dbus/v5 v5.0.4/go.mod h1:xhWf0FNVPg57R7Z0UbKHbJfkEywrmjJnf7w5xrFpKfA=
|
||||
github.com/goccy/go-json v0.10.6 h1:p8HrPJzOakx/mn/bQtjgNjdTcN+/S6FcG2CTtQOrHVU=
|
||||
github.com/goccy/go-json v0.10.6/go.mod h1:oq7eo15ShAhp70Anwd5lgX2pLfOS3QCiwU/PULtXL6M=
|
||||
github.com/jinzhu/inflection v1.0.0 h1:K317FqzuhWc8YvSVlFMCCUb36O/S9MCKRDI7QkRKD/E=
|
||||
github.com/jinzhu/inflection v1.0.0/go.mod h1:h+uFLlag+Qp1Va5pdKtLDYj+kHp5pxUVkryuEj+Srlc=
|
||||
github.com/jinzhu/now v1.1.5 h1:/o9tlHleP7gOFmsnYNz3RGnqzefHA47wQpKrrdTIwXQ=
|
||||
github.com/jinzhu/now v1.1.5/go.mod h1:d3SSVoowX0Lcu0IBviAWJpolVfI5UJVZZ7cO71lE/z8=
|
||||
github.com/mattn/go-colorable v0.1.13/go.mod h1:7S9/ev0klgBDR4GtXTXX8a3vIGJpMovkB8vQcUbaXHg=
|
||||
github.com/mattn/go-colorable v0.1.14 h1:9A9LHSqF/7dyVVX6g0U9cwm9pG3kP9gSzcuIPHPsaIE=
|
||||
github.com/mattn/go-colorable v0.1.14/go.mod h1:6LmQG8QLFO4G5z1gPvYEzlUgJ2wF+stgPZH1UqBm1s8=
|
||||
github.com/mattn/go-isatty v0.0.16/go.mod h1:kYGgaQfpe5nmfYZH+SKPsOc2e4SrIfOl2e/yFXSvRLM=
|
||||
github.com/mattn/go-isatty v0.0.19/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=
|
||||
github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY=
|
||||
github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=
|
||||
github.com/mattn/go-sqlite3 v1.14.33 h1:A5blZ5ulQo2AtayQ9/limgHEkFreKj1Dv226a1K73s0=
|
||||
github.com/mattn/go-sqlite3 v1.14.33/go.mod h1:Uh1q+B4BYcTPb+yiD3kU8Ct7aC0hY9fxUwlHK0RXw+Y=
|
||||
github.com/mattn/go-isatty v0.0.21 h1:xYae+lCNBP7QuW4PUnNG61ffM4hVIfm+zUzDuSzYLGs=
|
||||
github.com/mattn/go-isatty v0.0.21/go.mod h1:ZXfXG4SQHsB/w3ZeOYbR0PrPwLy+n6xiMrJlRFqopa4=
|
||||
github.com/mattn/go-sqlite3 v1.14.42 h1:MigqEP4ZmHw3aIdIT7T+9TLa90Z6smwcthx+Azv4Cgo=
|
||||
github.com/mattn/go-sqlite3 v1.14.42/go.mod h1:pjEuOr8IwzLJP2MfGeTb0A35jauH+C2kbHKBr7yXKVQ=
|
||||
github.com/mitchellh/colorstring v0.0.0-20190213212951-d06e56a500db h1:62I3jR2EmQ4l5rM/4FEfDWcRD+abF5XlKShorW5LRoQ=
|
||||
github.com/mitchellh/colorstring v0.0.0-20190213212951-d06e56a500db/go.mod h1:l0dey0ia/Uv7NcFFVbCLtqEBQbrT4OCwCSKTEv6enCw=
|
||||
github.com/patrickmn/go-cache v2.1.0+incompatible h1:HRMgzkcYKYpi3C8ajMPV8OFXaaRUnok+kx1WdO15EQc=
|
||||
github.com/patrickmn/go-cache v2.1.0+incompatible/go.mod h1:3Qf8kWWT7OJRJbdiICTKqZju1ZixQ/KpMGzzAfe6+WQ=
|
||||
github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
|
||||
github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
|
||||
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
|
||||
github.com/rivo/uniseg v0.4.7 h1:WUdvkW8uEhrYfLC4ZzdpI2ztxP1I582+49Oc5Mq64VQ=
|
||||
github.com/rivo/uniseg v0.4.7/go.mod h1:FN3SvrM+Zdj16jyLfmOkMNblXMcoc8DfTHruCPUcx88=
|
||||
github.com/rs/xid v1.6.0/go.mod h1:7XoLgs4eV+QndskICGsho+ADou8ySMSjJKDIan90Nz0=
|
||||
github.com/rs/zerolog v1.34.0 h1:k43nTLIwcTVQAncfCw4KZ2VY6ukYoZaBPNOE8txlOeY=
|
||||
github.com/rs/zerolog v1.34.0/go.mod h1:bJsvje4Z08ROH4Nhs5iH600c3IkWhwp44iRc54W6wYQ=
|
||||
github.com/rs/zerolog v1.35.0 h1:VD0ykx7HMiMJytqINBsKcbLS+BJ4WYjz+05us+LRTdI=
|
||||
github.com/rs/zerolog v1.35.0/go.mod h1:EjML9kdfa/RMA7h/6z6pYmq1ykOuA8/mjWaEvGI+jcw=
|
||||
github.com/schollz/progressbar/v2 v2.15.0 h1:dVzHQ8fHRmtPjD3K10jT3Qgn/+H+92jhPrhmxIJfDz8=
|
||||
github.com/schollz/progressbar/v2 v2.15.0/go.mod h1:UdPq3prGkfQ7MOzZKlDRpYKcFqEMczbD7YmbPgpzKMI=
|
||||
github.com/smacker/go-tree-sitter v0.0.0-20240827094217-dd81d9e9be82 h1:6C8qej6f1bStuePVkLSFxoU22XBS165D3klxlzRg8F4=
|
||||
github.com/smacker/go-tree-sitter v0.0.0-20240827094217-dd81d9e9be82/go.mod h1:xe4pgH49k4SsmkQq5OT8abwhWmnzkhpgnXeekbx2efw=
|
||||
github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
|
||||
github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=
|
||||
github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U=
|
||||
github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U=
|
||||
github.com/sugarme/regexpset v0.0.0-20200920021344-4d4ec8eaf93c h1:pwb4kNSHb4K89ymCaN+5lPH/MwnfSVg4rzGDh4d+iy4=
|
||||
github.com/sugarme/regexpset v0.0.0-20200920021344-4d4ec8eaf93c/go.mod h1:2gwkXLWbDGUQWeL3RtpCmcY4mzCtU13kb9UsAg9xMaw=
|
||||
github.com/yalue/onnxruntime_go v1.25.0 h1:nlhVau1BpLZ/BYr+WpPZCJRD/WES0qo6dK7aKyyAs3g=
|
||||
github.com/yalue/onnxruntime_go v1.25.0/go.mod h1:b4X26A8pekNb1ACJ58wAXgNKeUCGEAQ9dmACut9Sm/4=
|
||||
golang.org/x/sys v0.0.0-20220811171246-fbc7d0a398ab/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
golang.org/x/sys v0.12.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
golang.org/x/sys v0.39.0 h1:CvCKL8MeisomCi6qNZ+wbb0DN9E5AATixKsvNtMoMFk=
|
||||
golang.org/x/sys v0.39.0/go.mod h1:OgkHotnGiDImocRcuBABYBEXf8A9a87e/uXjp9XT3ks=
|
||||
golang.org/x/text v0.32.0 h1:ZD01bjUt1FQ9WJ0ClOL5vxgxOI/sVCNgX1YtKwcY0mU=
|
||||
golang.org/x/text v0.32.0/go.mod h1:o/rUWzghvpD5TXrTIBuJU77MTaN0ljMWE47kxGJQ7jY=
|
||||
github.com/yalue/onnxruntime_go v1.27.0 h1:c1YSgDNtpf0WGtxj3YeRIb8VC5LmM1J+Ve3uHdteC1U=
|
||||
github.com/yalue/onnxruntime_go v1.27.0/go.mod h1:b4X26A8pekNb1ACJ58wAXgNKeUCGEAQ9dmACut9Sm/4=
|
||||
golang.org/x/sync v0.20.0 h1:e0PTpb7pjO8GAtTs2dQ6jYa5BWYlMuX047Dco/pItO4=
|
||||
golang.org/x/sync v0.20.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0=
|
||||
golang.org/x/sys v0.43.0 h1:Rlag2XtaFTxp19wS8MXlJwTvoh8ArU6ezoyFsMyCTNI=
|
||||
golang.org/x/sys v0.43.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw=
|
||||
golang.org/x/text v0.36.0 h1:JfKh3XmcRPqZPKevfXVpI1wXPTqbkE5f7JA92a55Yxg=
|
||||
golang.org/x/text v0.36.0/go.mod h1:NIdBknypM8iqVmPiuco0Dh6P5Jcdk8lJL0CUebqK164=
|
||||
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
|
||||
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
|
||||
gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
|
||||
|
||||
@@ -0,0 +1,285 @@
|
||||
// Package golang provides AST-aware chunking for Go source files.
|
||||
package golang
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"go/ast"
|
||||
"go/parser"
|
||||
"go/token"
|
||||
"os"
|
||||
"strings"
|
||||
|
||||
"github.com/lukaszraczylo/claude-mnemonic/internal/chunking"
|
||||
)
|
||||
|
||||
// Chunker implements AST-aware chunking for Go files.
|
||||
type Chunker struct {
|
||||
options chunking.ChunkOptions
|
||||
}
|
||||
|
||||
// NewChunker creates a new Go chunker.
|
||||
func NewChunker(options chunking.ChunkOptions) *Chunker {
|
||||
return &Chunker{options: options}
|
||||
}
|
||||
|
||||
// Language returns the language this chunker supports.
|
||||
func (c *Chunker) Language() chunking.Language {
|
||||
return chunking.LanguageGo
|
||||
}
|
||||
|
||||
// SupportedExtensions returns the file extensions this chunker handles.
|
||||
func (c *Chunker) SupportedExtensions() []string {
|
||||
return []string{".go"}
|
||||
}
|
||||
|
||||
// Chunk parses a Go source file and returns semantic code chunks.
|
||||
func (c *Chunker) Chunk(ctx context.Context, filePath string) ([]chunking.Chunk, error) {
|
||||
// Read file content
|
||||
content, err := os.ReadFile(filePath)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("read file: %w", err)
|
||||
}
|
||||
|
||||
// Parse the Go file
|
||||
fset := token.NewFileSet()
|
||||
file, err := parser.ParseFile(fset, filePath, content, parser.ParseComments)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("parse Go file: %w", err)
|
||||
}
|
||||
|
||||
chunks := make([]chunking.Chunk, 0)
|
||||
sourceLines := strings.Split(string(content), "\n")
|
||||
|
||||
// Extract chunks from declarations
|
||||
for _, decl := range file.Decls {
|
||||
switch d := decl.(type) {
|
||||
case *ast.FuncDecl:
|
||||
chunk := c.extractFunction(fset, d, sourceLines, filePath)
|
||||
if chunk != nil {
|
||||
chunks = append(chunks, *chunk)
|
||||
}
|
||||
case *ast.GenDecl:
|
||||
extracted := c.extractGenDecl(fset, d, sourceLines, filePath)
|
||||
chunks = append(chunks, extracted...)
|
||||
}
|
||||
}
|
||||
|
||||
return chunks, nil
|
||||
}
|
||||
|
||||
// extractFunction extracts a function or method declaration as a chunk.
|
||||
func (c *Chunker) extractFunction(fset *token.FileSet, fn *ast.FuncDecl, sourceLines []string, filePath string) *chunking.Chunk {
|
||||
// Skip unexported if configured
|
||||
if !c.options.IncludePrivate && !fn.Name.IsExported() {
|
||||
return nil
|
||||
}
|
||||
|
||||
startPos := fset.Position(fn.Pos())
|
||||
endPos := fset.Position(fn.End())
|
||||
|
||||
chunk := &chunking.Chunk{
|
||||
FilePath: filePath,
|
||||
Language: chunking.LanguageGo,
|
||||
Name: fn.Name.Name,
|
||||
StartLine: startPos.Line,
|
||||
EndLine: endPos.Line,
|
||||
}
|
||||
|
||||
// Determine if this is a method or a function
|
||||
if fn.Recv != nil && len(fn.Recv.List) > 0 {
|
||||
chunk.Type = chunking.ChunkTypeMethod
|
||||
chunk.ParentName = c.extractReceiverType(fn.Recv)
|
||||
} else {
|
||||
chunk.Type = chunking.ChunkTypeFunction
|
||||
}
|
||||
|
||||
// Extract content
|
||||
chunk.Content = c.extractLines(sourceLines, startPos.Line, endPos.Line)
|
||||
|
||||
// Extract signature (function declaration without body)
|
||||
chunk.Signature = c.extractFunctionSignature(fn, fset, sourceLines)
|
||||
|
||||
// Extract doc comment
|
||||
if c.options.IncludeDocComments && fn.Doc != nil {
|
||||
chunk.DocComment = strings.TrimSpace(fn.Doc.Text())
|
||||
}
|
||||
|
||||
return chunk
|
||||
}
|
||||
|
||||
// extractGenDecl extracts general declarations (type, const, var).
|
||||
func (c *Chunker) extractGenDecl(fset *token.FileSet, gd *ast.GenDecl, sourceLines []string, filePath string) []chunking.Chunk {
|
||||
var chunks []chunking.Chunk
|
||||
|
||||
for _, spec := range gd.Specs {
|
||||
switch s := spec.(type) {
|
||||
case *ast.TypeSpec:
|
||||
chunk := c.extractTypeSpec(fset, gd, s, sourceLines, filePath)
|
||||
if chunk != nil {
|
||||
chunks = append(chunks, *chunk)
|
||||
}
|
||||
case *ast.ValueSpec:
|
||||
// Handle const and var declarations
|
||||
chunk := c.extractValueSpec(fset, gd, s, sourceLines, filePath)
|
||||
if chunk != nil {
|
||||
chunks = append(chunks, *chunk)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return chunks
|
||||
}
|
||||
|
||||
// extractTypeSpec extracts a type declaration (struct, interface, type alias).
|
||||
func (c *Chunker) extractTypeSpec(fset *token.FileSet, gd *ast.GenDecl, ts *ast.TypeSpec, sourceLines []string, filePath string) *chunking.Chunk {
|
||||
// Skip unexported if configured
|
||||
if !c.options.IncludePrivate && !ts.Name.IsExported() {
|
||||
return nil
|
||||
}
|
||||
|
||||
startPos := fset.Position(gd.Pos())
|
||||
endPos := fset.Position(gd.End())
|
||||
|
||||
chunk := &chunking.Chunk{
|
||||
FilePath: filePath,
|
||||
Language: chunking.LanguageGo,
|
||||
Name: ts.Name.Name,
|
||||
StartLine: startPos.Line,
|
||||
EndLine: endPos.Line,
|
||||
Content: c.extractLines(sourceLines, startPos.Line, endPos.Line),
|
||||
}
|
||||
|
||||
// Determine chunk type based on type expression
|
||||
switch ts.Type.(type) {
|
||||
case *ast.StructType:
|
||||
chunk.Type = chunking.ChunkTypeClass // Treat struct as class
|
||||
case *ast.InterfaceType:
|
||||
chunk.Type = chunking.ChunkTypeInterface
|
||||
default:
|
||||
chunk.Type = chunking.ChunkTypeType
|
||||
}
|
||||
|
||||
// Extract doc comment
|
||||
if c.options.IncludeDocComments && gd.Doc != nil {
|
||||
chunk.DocComment = strings.TrimSpace(gd.Doc.Text())
|
||||
}
|
||||
|
||||
return chunk
|
||||
}
|
||||
|
||||
// extractValueSpec extracts const or var declarations.
|
||||
func (c *Chunker) extractValueSpec(fset *token.FileSet, gd *ast.GenDecl, vs *ast.ValueSpec, sourceLines []string, filePath string) *chunking.Chunk {
|
||||
// Skip if all names are unexported and we're excluding private
|
||||
if !c.options.IncludePrivate {
|
||||
allUnexported := true
|
||||
for _, name := range vs.Names {
|
||||
if name.IsExported() {
|
||||
allUnexported = false
|
||||
break
|
||||
}
|
||||
}
|
||||
if allUnexported {
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
startPos := fset.Position(gd.Pos())
|
||||
endPos := fset.Position(gd.End())
|
||||
|
||||
// Use first name as the chunk name, join multiple if present
|
||||
names := make([]string, len(vs.Names))
|
||||
for i, name := range vs.Names {
|
||||
names[i] = name.Name
|
||||
}
|
||||
|
||||
chunk := &chunking.Chunk{
|
||||
FilePath: filePath,
|
||||
Language: chunking.LanguageGo,
|
||||
Name: strings.Join(names, ", "),
|
||||
StartLine: startPos.Line,
|
||||
EndLine: endPos.Line,
|
||||
Content: c.extractLines(sourceLines, startPos.Line, endPos.Line),
|
||||
}
|
||||
|
||||
// Set type based on token
|
||||
if gd.Tok == token.CONST {
|
||||
chunk.Type = chunking.ChunkTypeConst
|
||||
} else {
|
||||
chunk.Type = chunking.ChunkTypeVar
|
||||
}
|
||||
|
||||
// Extract doc comment
|
||||
if c.options.IncludeDocComments && gd.Doc != nil {
|
||||
chunk.DocComment = strings.TrimSpace(gd.Doc.Text())
|
||||
}
|
||||
|
||||
return chunk
|
||||
}
|
||||
|
||||
// extractReceiverType extracts the receiver type name from a method.
|
||||
func (c *Chunker) extractReceiverType(recv *ast.FieldList) string {
|
||||
if len(recv.List) == 0 {
|
||||
return ""
|
||||
}
|
||||
|
||||
field := recv.List[0]
|
||||
switch t := field.Type.(type) {
|
||||
case *ast.Ident:
|
||||
return t.Name
|
||||
case *ast.StarExpr:
|
||||
if ident, ok := t.X.(*ast.Ident); ok {
|
||||
return ident.Name
|
||||
}
|
||||
}
|
||||
|
||||
return ""
|
||||
}
|
||||
|
||||
// extractFunctionSignature extracts the function signature without the body.
|
||||
func (c *Chunker) extractFunctionSignature(fn *ast.FuncDecl, fset *token.FileSet, sourceLines []string) string {
|
||||
if fn.Body == nil {
|
||||
// No body, return entire declaration
|
||||
startPos := fset.Position(fn.Pos())
|
||||
endPos := fset.Position(fn.End())
|
||||
return c.extractLines(sourceLines, startPos.Line, endPos.Line)
|
||||
}
|
||||
|
||||
// Extract from start of function to just before body
|
||||
startPos := fset.Position(fn.Pos())
|
||||
bodyPos := fset.Position(fn.Body.Pos())
|
||||
|
||||
// If body is on the same line, extract just that line up to the opening brace
|
||||
if startPos.Line == bodyPos.Line {
|
||||
line := sourceLines[startPos.Line-1]
|
||||
// Find the opening brace position
|
||||
if idx := strings.Index(line[startPos.Column-1:], "{"); idx >= 0 {
|
||||
return strings.TrimSpace(line[startPos.Column-1 : startPos.Column-1+idx])
|
||||
}
|
||||
return strings.TrimSpace(line[startPos.Column-1:])
|
||||
}
|
||||
|
||||
// Get lines from start to the line containing the opening brace
|
||||
sig := c.extractLines(sourceLines, startPos.Line, bodyPos.Line)
|
||||
// Remove the opening brace and anything after it
|
||||
if idx := strings.Index(sig, "{"); idx >= 0 {
|
||||
sig = sig[:idx]
|
||||
}
|
||||
return strings.TrimSpace(sig)
|
||||
}
|
||||
|
||||
// extractLines extracts a range of lines from source (1-indexed, inclusive).
|
||||
func (c *Chunker) extractLines(lines []string, start, end int) string {
|
||||
if start < 1 || end < start || start > len(lines) {
|
||||
return ""
|
||||
}
|
||||
|
||||
// Adjust for 0-indexed array (start and end are 1-indexed)
|
||||
startIdx := start - 1
|
||||
endIdx := end
|
||||
if endIdx > len(lines) {
|
||||
endIdx = len(lines)
|
||||
}
|
||||
|
||||
return strings.Join(lines[startIdx:endIdx], "\n")
|
||||
}
|
||||
@@ -0,0 +1,214 @@
|
||||
package golang
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
"github.com/lukaszraczylo/claude-mnemonic/internal/chunking"
|
||||
)
|
||||
|
||||
func TestGoChunker_BasicFunctions(t *testing.T) {
|
||||
// Create temp test file
|
||||
tmpDir := t.TempDir()
|
||||
testFile := filepath.Join(tmpDir, "test.go")
|
||||
|
||||
testCode := `package main
|
||||
|
||||
import "fmt"
|
||||
|
||||
// Greet prints a greeting message
|
||||
func Greet(name string) {
|
||||
fmt.Printf("Hello, %s!\n", name)
|
||||
}
|
||||
|
||||
// Add adds two numbers
|
||||
func Add(a, b int) int {
|
||||
return a + b
|
||||
}
|
||||
|
||||
// unexported function should be included by default
|
||||
func helper() string {
|
||||
return "helper"
|
||||
}
|
||||
`
|
||||
|
||||
if err := os.WriteFile(testFile, []byte(testCode), 0600); err != nil {
|
||||
t.Fatalf("Failed to create test file: %v", err)
|
||||
}
|
||||
|
||||
// Create chunker with default options
|
||||
chunker := NewChunker(chunking.DefaultChunkOptions())
|
||||
|
||||
// Chunk the file
|
||||
chunks, err := chunker.Chunk(context.Background(), testFile)
|
||||
if err != nil {
|
||||
t.Fatalf("Chunk() failed: %v", err)
|
||||
}
|
||||
|
||||
// Verify we got all functions
|
||||
if len(chunks) != 3 {
|
||||
t.Errorf("Expected 3 chunks (Greet, Add, helper), got %d", len(chunks))
|
||||
}
|
||||
|
||||
// Verify chunk details
|
||||
expectedNames := map[string]bool{
|
||||
"Greet": false,
|
||||
"Add": false,
|
||||
"helper": false,
|
||||
}
|
||||
|
||||
for _, chunk := range chunks {
|
||||
if chunk.Type != chunking.ChunkTypeFunction {
|
||||
t.Errorf("Expected chunk type 'function', got '%s'", chunk.Type)
|
||||
}
|
||||
|
||||
if chunk.Language != chunking.LanguageGo {
|
||||
t.Errorf("Expected language 'go', got '%s'", chunk.Language)
|
||||
}
|
||||
|
||||
if _, ok := expectedNames[chunk.Name]; !ok {
|
||||
t.Errorf("Unexpected function name: %s", chunk.Name)
|
||||
} else {
|
||||
expectedNames[chunk.Name] = true
|
||||
}
|
||||
|
||||
// Verify content is non-empty
|
||||
if chunk.Content == "" {
|
||||
t.Errorf("Chunk %s has empty content", chunk.Name)
|
||||
}
|
||||
|
||||
// Verify signature is present for functions
|
||||
if chunk.Signature == "" {
|
||||
t.Errorf("Chunk %s has empty signature", chunk.Name)
|
||||
}
|
||||
}
|
||||
|
||||
// Verify all expected functions were found
|
||||
for name, found := range expectedNames {
|
||||
if !found {
|
||||
t.Errorf("Expected function %s not found", name)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestGoChunker_StructsAndMethods(t *testing.T) {
|
||||
tmpDir := t.TempDir()
|
||||
testFile := filepath.Join(tmpDir, "test.go")
|
||||
|
||||
testCode := `package main
|
||||
|
||||
// User represents a user
|
||||
type User struct {
|
||||
ID int
|
||||
Name string
|
||||
}
|
||||
|
||||
// GetName returns the user's name
|
||||
func (u *User) GetName() string {
|
||||
return u.Name
|
||||
}
|
||||
|
||||
// SetName sets the user's name
|
||||
func (u *User) SetName(name string) {
|
||||
u.Name = name
|
||||
}
|
||||
`
|
||||
|
||||
if err := os.WriteFile(testFile, []byte(testCode), 0600); err != nil {
|
||||
t.Fatalf("Failed to create test file: %v", err)
|
||||
}
|
||||
|
||||
chunker := NewChunker(chunking.DefaultChunkOptions())
|
||||
chunks, err := chunker.Chunk(context.Background(), testFile)
|
||||
if err != nil {
|
||||
t.Fatalf("Chunk() failed: %v", err)
|
||||
}
|
||||
|
||||
// Should have 1 struct + 2 methods = 3 chunks
|
||||
if len(chunks) != 3 {
|
||||
t.Errorf("Expected 3 chunks (User struct, GetName, SetName), got %d", len(chunks))
|
||||
}
|
||||
|
||||
// Find the struct and methods
|
||||
var structChunk, getNameChunk, setNameChunk *chunking.Chunk
|
||||
for i := range chunks {
|
||||
switch chunks[i].Name {
|
||||
case "User":
|
||||
structChunk = &chunks[i]
|
||||
case "GetName":
|
||||
getNameChunk = &chunks[i]
|
||||
case "SetName":
|
||||
setNameChunk = &chunks[i]
|
||||
}
|
||||
}
|
||||
|
||||
// Verify struct
|
||||
if structChunk == nil {
|
||||
t.Fatal("User struct not found")
|
||||
}
|
||||
if structChunk.Type != chunking.ChunkTypeClass {
|
||||
t.Errorf("Expected User to be ChunkTypeClass, got %s", structChunk.Type)
|
||||
}
|
||||
|
||||
// Verify methods
|
||||
if getNameChunk == nil {
|
||||
t.Fatal("GetName method not found")
|
||||
}
|
||||
if getNameChunk.Type != chunking.ChunkTypeMethod {
|
||||
t.Errorf("Expected GetName to be ChunkTypeMethod, got %s", getNameChunk.Type)
|
||||
}
|
||||
if getNameChunk.ParentName != "User" {
|
||||
t.Errorf("Expected GetName parent to be 'User', got '%s'", getNameChunk.ParentName)
|
||||
}
|
||||
|
||||
if setNameChunk == nil {
|
||||
t.Fatal("SetName method not found")
|
||||
}
|
||||
if setNameChunk.Type != chunking.ChunkTypeMethod {
|
||||
t.Errorf("Expected SetName to be ChunkTypeMethod, got %s", setNameChunk.Type)
|
||||
}
|
||||
if setNameChunk.ParentName != "User" {
|
||||
t.Errorf("Expected SetName parent to be 'User', got '%s'", setNameChunk.ParentName)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGoChunker_DocComments(t *testing.T) {
|
||||
tmpDir := t.TempDir()
|
||||
testFile := filepath.Join(tmpDir, "test.go")
|
||||
|
||||
testCode := `package main
|
||||
|
||||
// Calculate performs a calculation.
|
||||
// It takes two integers and returns their sum.
|
||||
func Calculate(a, b int) int {
|
||||
return a + b
|
||||
}
|
||||
`
|
||||
|
||||
if err := os.WriteFile(testFile, []byte(testCode), 0600); err != nil {
|
||||
t.Fatalf("Failed to create test file: %v", err)
|
||||
}
|
||||
|
||||
chunker := NewChunker(chunking.DefaultChunkOptions())
|
||||
chunks, err := chunker.Chunk(context.Background(), testFile)
|
||||
if err != nil {
|
||||
t.Fatalf("Chunk() failed: %v", err)
|
||||
}
|
||||
|
||||
if len(chunks) != 1 {
|
||||
t.Fatalf("Expected 1 chunk, got %d", len(chunks))
|
||||
}
|
||||
|
||||
chunk := chunks[0]
|
||||
if chunk.DocComment == "" {
|
||||
t.Error("Expected doc comment to be present")
|
||||
}
|
||||
|
||||
// Doc comment should contain the comment text
|
||||
expectedComment := "Calculate performs a calculation.\nIt takes two integers and returns their sum."
|
||||
if chunk.DocComment != expectedComment {
|
||||
t.Errorf("Expected doc comment '%s', got '%s'", expectedComment, chunk.DocComment)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,106 @@
|
||||
package chunking
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// Manager dispatches files to appropriate language-specific chunkers.
|
||||
type Manager struct {
|
||||
chunkers map[string]Chunker // extension -> chunker
|
||||
options ChunkOptions
|
||||
}
|
||||
|
||||
// NewManager creates a new chunking manager with the given chunkers.
|
||||
func NewManager(chunkers []Chunker, options ChunkOptions) *Manager {
|
||||
m := &Manager{
|
||||
chunkers: make(map[string]Chunker),
|
||||
options: options,
|
||||
}
|
||||
|
||||
// Register chunkers by their supported extensions
|
||||
for _, chunker := range chunkers {
|
||||
for _, ext := range chunker.SupportedExtensions() {
|
||||
m.chunkers[ext] = chunker
|
||||
}
|
||||
}
|
||||
|
||||
return m
|
||||
}
|
||||
|
||||
// ChunkFile chunks a single file using the appropriate language chunker.
|
||||
// Returns an error if no chunker is found for the file extension.
|
||||
func (m *Manager) ChunkFile(ctx context.Context, filePath string) ([]Chunk, error) {
|
||||
ext := strings.ToLower(filepath.Ext(filePath))
|
||||
chunker, ok := m.chunkers[ext]
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("no chunker for extension %s", ext)
|
||||
}
|
||||
|
||||
chunks, err := chunker.Chunk(ctx, filePath)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("chunk %s: %w", filePath, err)
|
||||
}
|
||||
|
||||
// Apply options-based filtering
|
||||
filtered := make([]Chunk, 0, len(chunks))
|
||||
for _, chunk := range chunks {
|
||||
// Filter by minimum lines
|
||||
if m.options.MinLines > 0 {
|
||||
lineCount := chunk.EndLine - chunk.StartLine + 1
|
||||
if lineCount < m.options.MinLines {
|
||||
continue
|
||||
}
|
||||
}
|
||||
|
||||
// Filter by maximum chunk size
|
||||
if m.options.MaxChunkSize > 0 && len(chunk.Content) > m.options.MaxChunkSize {
|
||||
// TODO: Consider splitting large chunks intelligently
|
||||
// For now, skip chunks that are too large
|
||||
continue
|
||||
}
|
||||
|
||||
filtered = append(filtered, chunk)
|
||||
}
|
||||
|
||||
return filtered, nil
|
||||
}
|
||||
|
||||
// ChunkFiles chunks multiple files in parallel.
|
||||
// Returns a map of file path to chunks, and any errors encountered.
|
||||
// Errors for individual files do not stop processing of other files.
|
||||
func (m *Manager) ChunkFiles(ctx context.Context, filePaths []string) (map[string][]Chunk, []error) {
|
||||
results := make(map[string][]Chunk)
|
||||
var errors []error
|
||||
|
||||
for _, filePath := range filePaths {
|
||||
chunks, err := m.ChunkFile(ctx, filePath)
|
||||
if err != nil {
|
||||
errors = append(errors, fmt.Errorf("%s: %w", filePath, err))
|
||||
continue
|
||||
}
|
||||
if len(chunks) > 0 {
|
||||
results[filePath] = chunks
|
||||
}
|
||||
}
|
||||
|
||||
return results, errors
|
||||
}
|
||||
|
||||
// SupportsFile checks if the manager can chunk the given file based on extension.
|
||||
func (m *Manager) SupportsFile(filePath string) bool {
|
||||
ext := strings.ToLower(filepath.Ext(filePath))
|
||||
_, ok := m.chunkers[ext]
|
||||
return ok
|
||||
}
|
||||
|
||||
// SupportedExtensions returns all file extensions supported by registered chunkers.
|
||||
func (m *Manager) SupportedExtensions() []string {
|
||||
exts := make([]string, 0, len(m.chunkers))
|
||||
for ext := range m.chunkers {
|
||||
exts = append(exts, ext)
|
||||
}
|
||||
return exts
|
||||
}
|
||||
@@ -0,0 +1,162 @@
|
||||
package chunking
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// mockChunker is a test chunker that returns dummy chunks
|
||||
type mockChunker struct{}
|
||||
|
||||
func (m *mockChunker) Chunk(ctx context.Context, filePath string) ([]Chunk, error) {
|
||||
// Just return an empty chunk for testing
|
||||
return []Chunk{
|
||||
{
|
||||
FilePath: filePath,
|
||||
Language: LanguageGo,
|
||||
Type: ChunkTypeFunction,
|
||||
Name: "TestFunc",
|
||||
StartLine: 1,
|
||||
EndLine: 1,
|
||||
Content: "test",
|
||||
},
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (m *mockChunker) Language() Language {
|
||||
return LanguageGo
|
||||
}
|
||||
|
||||
func (m *mockChunker) SupportedExtensions() []string {
|
||||
return []string{".go", ".py", ".ts"}
|
||||
}
|
||||
|
||||
func TestManager_ChunkMultipleFiles(t *testing.T) {
|
||||
tmpDir := t.TempDir()
|
||||
|
||||
// Create a Go file
|
||||
goFile := filepath.Join(tmpDir, "test.go")
|
||||
goCode := `package main
|
||||
|
||||
func Hello() string {
|
||||
return "hello"
|
||||
}
|
||||
`
|
||||
if err := os.WriteFile(goFile, []byte(goCode), 0600); err != nil {
|
||||
t.Fatalf("Failed to create Go file: %v", err)
|
||||
}
|
||||
|
||||
// Create a Python file
|
||||
pyFile := filepath.Join(tmpDir, "test.py")
|
||||
pyCode := `def greet(name):
|
||||
return f"Hello, {name}!"
|
||||
|
||||
class User:
|
||||
def __init__(self, name):
|
||||
self.name = name
|
||||
`
|
||||
if err := os.WriteFile(pyFile, []byte(pyCode), 0600); err != nil {
|
||||
t.Fatalf("Failed to create Python file: %v", err)
|
||||
}
|
||||
|
||||
// Create a TypeScript file
|
||||
tsFile := filepath.Join(tmpDir, "test.ts")
|
||||
tsCode := `function add(a: number, b: number): number {
|
||||
return a + b;
|
||||
}
|
||||
|
||||
class Calculator {
|
||||
multiply(a: number, b: number): number {
|
||||
return a * b;
|
||||
}
|
||||
}
|
||||
`
|
||||
if err := os.WriteFile(tsFile, []byte(tsCode), 0600); err != nil {
|
||||
t.Fatalf("Failed to create TypeScript file: %v", err)
|
||||
}
|
||||
|
||||
// Create manager
|
||||
manager := NewManager([]Chunker{&mockChunker{}}, DefaultChunkOptions())
|
||||
|
||||
// Test SupportsFile
|
||||
if !manager.SupportsFile(goFile) {
|
||||
t.Error("Manager should support .go files")
|
||||
}
|
||||
if !manager.SupportsFile(pyFile) {
|
||||
t.Error("Manager should support .py files")
|
||||
}
|
||||
if !manager.SupportsFile(tsFile) {
|
||||
t.Error("Manager should support .ts files")
|
||||
}
|
||||
|
||||
unsupportedFile := filepath.Join(tmpDir, "test.txt")
|
||||
if manager.SupportsFile(unsupportedFile) {
|
||||
t.Error("Manager should not support .txt files")
|
||||
}
|
||||
|
||||
// Test ChunkFiles
|
||||
results, errs := manager.ChunkFiles(context.Background(), []string{goFile, pyFile, tsFile})
|
||||
if len(errs) > 0 {
|
||||
t.Errorf("ChunkFiles returned errors: %v", errs)
|
||||
}
|
||||
|
||||
if len(results) != 3 {
|
||||
t.Errorf("Expected results for 3 files, got %d", len(results))
|
||||
}
|
||||
|
||||
// Verify each file has chunks
|
||||
for _, file := range []string{goFile, pyFile, tsFile} {
|
||||
if chunks, ok := results[file]; !ok || len(chunks) == 0 {
|
||||
t.Errorf("No chunks found for file %s", file)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// mockChunkerWithExts is a test chunker with configurable extensions
|
||||
type mockChunkerWithExts struct {
|
||||
exts []string
|
||||
}
|
||||
|
||||
func (m *mockChunkerWithExts) Chunk(ctx context.Context, filePath string) ([]Chunk, error) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
func (m *mockChunkerWithExts) Language() Language {
|
||||
return LanguageGo
|
||||
}
|
||||
|
||||
func (m *mockChunkerWithExts) SupportedExtensions() []string {
|
||||
return m.exts
|
||||
}
|
||||
|
||||
func TestManager_SupportedExtensions(t *testing.T) {
|
||||
|
||||
// Create manager with mock chunkers
|
||||
manager := NewManager([]Chunker{
|
||||
&mockChunkerWithExts{exts: []string{".go"}},
|
||||
&mockChunkerWithExts{exts: []string{".py", ".pyw"}},
|
||||
}, DefaultChunkOptions())
|
||||
|
||||
exts := manager.SupportedExtensions()
|
||||
expectedExts := map[string]bool{
|
||||
".go": false,
|
||||
".py": false,
|
||||
".pyw": false,
|
||||
}
|
||||
|
||||
for _, ext := range exts {
|
||||
if _, ok := expectedExts[ext]; ok {
|
||||
expectedExts[ext] = true
|
||||
} else {
|
||||
t.Errorf("Unexpected extension: %s", ext)
|
||||
}
|
||||
}
|
||||
|
||||
for ext, found := range expectedExts {
|
||||
if !found {
|
||||
t.Errorf("Expected extension %s not found", ext)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,291 @@
|
||||
// Package python provides AST-aware chunking for Python source files using tree-sitter.
|
||||
package python
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"os"
|
||||
"strings"
|
||||
|
||||
sitter "github.com/smacker/go-tree-sitter"
|
||||
"github.com/smacker/go-tree-sitter/python"
|
||||
|
||||
"github.com/lukaszraczylo/claude-mnemonic/internal/chunking"
|
||||
)
|
||||
|
||||
// Chunker implements AST-aware chunking for Python files.
|
||||
type Chunker struct {
|
||||
parser *sitter.Parser
|
||||
options chunking.ChunkOptions
|
||||
}
|
||||
|
||||
// NewChunker creates a new Python chunker.
|
||||
func NewChunker(options chunking.ChunkOptions) *Chunker {
|
||||
parser := sitter.NewParser()
|
||||
parser.SetLanguage(python.GetLanguage())
|
||||
|
||||
return &Chunker{
|
||||
options: options,
|
||||
parser: parser,
|
||||
}
|
||||
}
|
||||
|
||||
// Language returns the language this chunker supports.
|
||||
func (c *Chunker) Language() chunking.Language {
|
||||
return chunking.LanguagePython
|
||||
}
|
||||
|
||||
// SupportedExtensions returns the file extensions this chunker handles.
|
||||
func (c *Chunker) SupportedExtensions() []string {
|
||||
return []string{".py"}
|
||||
}
|
||||
|
||||
// Chunk parses a Python source file and returns semantic code chunks.
|
||||
func (c *Chunker) Chunk(ctx context.Context, filePath string) ([]chunking.Chunk, error) {
|
||||
// Read file content
|
||||
content, err := os.ReadFile(filePath)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("read file: %w", err)
|
||||
}
|
||||
|
||||
// Parse the Python file
|
||||
tree, err := c.parser.ParseCtx(ctx, nil, content)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("parse Python file: %w", err)
|
||||
}
|
||||
defer tree.Close()
|
||||
|
||||
sourceLines := strings.Split(string(content), "\n")
|
||||
chunks := make([]chunking.Chunk, 0)
|
||||
|
||||
// Walk the AST and extract chunks
|
||||
c.walkNode(tree.RootNode(), content, sourceLines, filePath, "", &chunks)
|
||||
|
||||
return chunks, nil
|
||||
}
|
||||
|
||||
// walkNode recursively walks the tree-sitter AST and extracts chunks.
|
||||
func (c *Chunker) walkNode(node *sitter.Node, source []byte, sourceLines []string, filePath string, parentName string, chunks *[]chunking.Chunk) {
|
||||
nodeType := node.Type()
|
||||
|
||||
switch nodeType {
|
||||
case "function_definition":
|
||||
chunk := c.extractFunction(node, source, sourceLines, filePath, parentName)
|
||||
if chunk != nil {
|
||||
*chunks = append(*chunks, *chunk)
|
||||
}
|
||||
|
||||
case "class_definition":
|
||||
chunk := c.extractClass(node, source, sourceLines, filePath)
|
||||
if chunk != nil {
|
||||
*chunks = append(*chunks, *chunk)
|
||||
|
||||
// Walk class body to find methods
|
||||
for i := 0; i < int(node.ChildCount()); i++ {
|
||||
child := node.Child(i)
|
||||
if child.Type() == "block" {
|
||||
c.walkNode(child, source, sourceLines, filePath, chunk.Name, chunks)
|
||||
}
|
||||
}
|
||||
}
|
||||
return // Don't walk children again
|
||||
|
||||
case "block":
|
||||
// Walk statements in block
|
||||
for i := 0; i < int(node.ChildCount()); i++ {
|
||||
c.walkNode(node.Child(i), source, sourceLines, filePath, parentName, chunks)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
// Walk all children
|
||||
for i := 0; i < int(node.ChildCount()); i++ {
|
||||
c.walkNode(node.Child(i), source, sourceLines, filePath, parentName, chunks)
|
||||
}
|
||||
}
|
||||
|
||||
// extractFunction extracts a function definition chunk.
|
||||
func (c *Chunker) extractFunction(node *sitter.Node, source []byte, sourceLines []string, filePath string, parentName string) *chunking.Chunk {
|
||||
// Find function name
|
||||
var nameNode *sitter.Node
|
||||
for i := 0; i < int(node.ChildCount()); i++ {
|
||||
child := node.Child(i)
|
||||
if child.Type() == "identifier" {
|
||||
nameNode = child
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
if nameNode == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
name := nameNode.Content(source)
|
||||
|
||||
// Skip private functions if configured
|
||||
if !c.options.IncludePrivate && strings.HasPrefix(name, "_") && !strings.HasPrefix(name, "__") {
|
||||
return nil
|
||||
}
|
||||
|
||||
startLine := int(node.StartPoint().Row) + 1
|
||||
endLine := int(node.EndPoint().Row) + 1
|
||||
|
||||
chunk := &chunking.Chunk{
|
||||
FilePath: filePath,
|
||||
Language: chunking.LanguagePython,
|
||||
Name: name,
|
||||
ParentName: parentName,
|
||||
StartLine: startLine,
|
||||
EndLine: endLine,
|
||||
Content: c.extractLines(sourceLines, startLine, endLine),
|
||||
}
|
||||
|
||||
// Determine if this is a method or function
|
||||
if parentName != "" {
|
||||
chunk.Type = chunking.ChunkTypeMethod
|
||||
} else {
|
||||
chunk.Type = chunking.ChunkTypeFunction
|
||||
}
|
||||
|
||||
// Extract signature (def line)
|
||||
chunk.Signature = c.extractFunctionSignature(node, source, sourceLines)
|
||||
|
||||
// Extract docstring as doc comment
|
||||
if c.options.IncludeDocComments {
|
||||
chunk.DocComment = c.extractDocstring(node, source)
|
||||
}
|
||||
|
||||
return chunk
|
||||
}
|
||||
|
||||
// extractClass extracts a class definition chunk.
|
||||
func (c *Chunker) extractClass(node *sitter.Node, source []byte, sourceLines []string, filePath string) *chunking.Chunk {
|
||||
// Find class name
|
||||
var nameNode *sitter.Node
|
||||
for i := 0; i < int(node.ChildCount()); i++ {
|
||||
child := node.Child(i)
|
||||
if child.Type() == "identifier" {
|
||||
nameNode = child
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
if nameNode == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
name := nameNode.Content(source)
|
||||
|
||||
// Skip private classes if configured
|
||||
if !c.options.IncludePrivate && strings.HasPrefix(name, "_") && !strings.HasPrefix(name, "__") {
|
||||
return nil
|
||||
}
|
||||
|
||||
startLine := int(node.StartPoint().Row) + 1
|
||||
endLine := int(node.EndPoint().Row) + 1
|
||||
|
||||
chunk := &chunking.Chunk{
|
||||
FilePath: filePath,
|
||||
Language: chunking.LanguagePython,
|
||||
Type: chunking.ChunkTypeClass,
|
||||
Name: name,
|
||||
StartLine: startLine,
|
||||
EndLine: endLine,
|
||||
Content: c.extractLines(sourceLines, startLine, endLine),
|
||||
}
|
||||
|
||||
// Extract class signature (class line)
|
||||
chunk.Signature = c.extractClassSignature(node, source, sourceLines)
|
||||
|
||||
// Extract docstring as doc comment
|
||||
if c.options.IncludeDocComments {
|
||||
chunk.DocComment = c.extractDocstring(node, source)
|
||||
}
|
||||
|
||||
return chunk
|
||||
}
|
||||
|
||||
// extractFunctionSignature extracts the function definition line.
|
||||
func (c *Chunker) extractFunctionSignature(node *sitter.Node, source []byte, sourceLines []string) string {
|
||||
startLine := int(node.StartPoint().Row) + 1
|
||||
|
||||
// Find the colon that ends the signature
|
||||
for i := 0; i < int(node.ChildCount()); i++ {
|
||||
child := node.Child(i)
|
||||
if child.Type() == ":" {
|
||||
endLine := int(child.EndPoint().Row) + 1
|
||||
return strings.TrimSpace(c.extractLines(sourceLines, startLine, endLine))
|
||||
}
|
||||
}
|
||||
|
||||
// Fallback: just return first line
|
||||
return strings.TrimSpace(c.extractLines(sourceLines, startLine, startLine))
|
||||
}
|
||||
|
||||
// extractClassSignature extracts the class definition line.
|
||||
func (c *Chunker) extractClassSignature(node *sitter.Node, source []byte, sourceLines []string) string {
|
||||
startLine := int(node.StartPoint().Row) + 1
|
||||
|
||||
// Find the colon that ends the signature
|
||||
for i := 0; i < int(node.ChildCount()); i++ {
|
||||
child := node.Child(i)
|
||||
if child.Type() == ":" {
|
||||
endLine := int(child.EndPoint().Row) + 1
|
||||
return strings.TrimSpace(c.extractLines(sourceLines, startLine, endLine))
|
||||
}
|
||||
}
|
||||
|
||||
// Fallback: just return first line
|
||||
return strings.TrimSpace(c.extractLines(sourceLines, startLine, startLine))
|
||||
}
|
||||
|
||||
// extractDocstring extracts the docstring from a function or class.
|
||||
func (c *Chunker) extractDocstring(node *sitter.Node, source []byte) string {
|
||||
// Find the block
|
||||
var blockNode *sitter.Node
|
||||
for i := 0; i < int(node.ChildCount()); i++ {
|
||||
child := node.Child(i)
|
||||
if child.Type() == "block" {
|
||||
blockNode = child
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
if blockNode == nil {
|
||||
return ""
|
||||
}
|
||||
|
||||
// Check if first statement in block is a string (docstring)
|
||||
for i := 0; i < int(blockNode.ChildCount()); i++ {
|
||||
child := blockNode.Child(i)
|
||||
if child.Type() == "expression_statement" {
|
||||
// Check if it contains a string
|
||||
for j := 0; j < int(child.ChildCount()); j++ {
|
||||
grandchild := child.Child(j)
|
||||
if grandchild.Type() == "string" {
|
||||
docstring := grandchild.Content(source)
|
||||
// Remove quotes
|
||||
docstring = strings.Trim(docstring, `"'`)
|
||||
return strings.TrimSpace(docstring)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return ""
|
||||
}
|
||||
|
||||
// extractLines extracts a range of lines from source (1-indexed, inclusive).
|
||||
func (c *Chunker) extractLines(lines []string, start, end int) string {
|
||||
if start < 1 || end < start || start > len(lines) {
|
||||
return ""
|
||||
}
|
||||
|
||||
startIdx := start - 1
|
||||
endIdx := end
|
||||
if endIdx > len(lines) {
|
||||
endIdx = len(lines)
|
||||
}
|
||||
|
||||
return strings.Join(lines[startIdx:endIdx], "\n")
|
||||
}
|
||||
@@ -0,0 +1,298 @@
|
||||
package python
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/lukaszraczylo/claude-mnemonic/internal/chunking"
|
||||
)
|
||||
|
||||
// =============================================================================
|
||||
// TEST HELPERS
|
||||
// =============================================================================
|
||||
|
||||
func createTempPythonFile(t *testing.T, content string) string {
|
||||
t.Helper()
|
||||
|
||||
tmpDir := t.TempDir()
|
||||
filePath := filepath.Join(tmpDir, "test.py")
|
||||
|
||||
err := os.WriteFile(filePath, []byte(content), 0600)
|
||||
require.NoError(t, err)
|
||||
|
||||
return filePath
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// TESTS FOR Chunker
|
||||
// =============================================================================
|
||||
|
||||
func TestNewChunker(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
opts := chunking.DefaultChunkOptions()
|
||||
c := NewChunker(opts)
|
||||
|
||||
assert.NotNil(t, c)
|
||||
assert.NotNil(t, c.parser)
|
||||
}
|
||||
|
||||
func TestChunker_Language(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
c := NewChunker(chunking.DefaultChunkOptions())
|
||||
|
||||
assert.Equal(t, chunking.LanguagePython, c.Language())
|
||||
}
|
||||
|
||||
func TestChunker_SupportedExtensions(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
c := NewChunker(chunking.DefaultChunkOptions())
|
||||
exts := c.SupportedExtensions()
|
||||
|
||||
assert.Contains(t, exts, ".py")
|
||||
}
|
||||
|
||||
func TestChunker_Chunk_SimpleFunction(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
code := `def greet(name):
|
||||
"""Greets a person by name."""
|
||||
return f"Hello, {name}!"
|
||||
`
|
||||
|
||||
filePath := createTempPythonFile(t, code)
|
||||
c := NewChunker(chunking.DefaultChunkOptions())
|
||||
|
||||
chunks, err := c.Chunk(context.Background(), filePath)
|
||||
require.NoError(t, err)
|
||||
require.NotEmpty(t, chunks)
|
||||
|
||||
// Should find the greet function
|
||||
var foundGreet bool
|
||||
for _, chunk := range chunks {
|
||||
if chunk.Name == "greet" {
|
||||
foundGreet = true
|
||||
assert.Equal(t, chunking.ChunkTypeFunction, chunk.Type)
|
||||
assert.Equal(t, chunking.LanguagePython, chunk.Language)
|
||||
assert.Contains(t, chunk.Content, "def greet")
|
||||
}
|
||||
}
|
||||
assert.True(t, foundGreet, "Should find 'greet' function")
|
||||
}
|
||||
|
||||
func TestChunker_Chunk_ClassWithMethods(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
code := `class Calculator:
|
||||
"""A simple calculator class."""
|
||||
|
||||
def add(self, a, b):
|
||||
"""Adds two numbers."""
|
||||
return a + b
|
||||
|
||||
def multiply(self, a, b):
|
||||
"""Multiplies two numbers."""
|
||||
return a * b
|
||||
`
|
||||
|
||||
filePath := createTempPythonFile(t, code)
|
||||
c := NewChunker(chunking.DefaultChunkOptions())
|
||||
|
||||
chunks, err := c.Chunk(context.Background(), filePath)
|
||||
require.NoError(t, err)
|
||||
require.NotEmpty(t, chunks)
|
||||
|
||||
// Should find the Calculator class and its methods
|
||||
var foundClass, foundAdd, foundMultiply bool
|
||||
for _, chunk := range chunks {
|
||||
switch chunk.Name {
|
||||
case "Calculator":
|
||||
foundClass = true
|
||||
assert.Equal(t, chunking.ChunkTypeClass, chunk.Type)
|
||||
case "add":
|
||||
foundAdd = true
|
||||
assert.Equal(t, chunking.ChunkTypeMethod, chunk.Type)
|
||||
assert.Equal(t, "Calculator", chunk.ParentName)
|
||||
case "multiply":
|
||||
foundMultiply = true
|
||||
assert.Equal(t, chunking.ChunkTypeMethod, chunk.Type)
|
||||
assert.Equal(t, "Calculator", chunk.ParentName)
|
||||
}
|
||||
}
|
||||
|
||||
assert.True(t, foundClass, "Should find 'Calculator' class")
|
||||
assert.True(t, foundAdd, "Should find 'add' method")
|
||||
assert.True(t, foundMultiply, "Should find 'multiply' method")
|
||||
}
|
||||
|
||||
func TestChunker_Chunk_MultipleFunctions(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
code := `def first_function():
|
||||
pass
|
||||
|
||||
def second_function(x, y):
|
||||
return x + y
|
||||
|
||||
def third_function():
|
||||
"""Has a docstring."""
|
||||
return 42
|
||||
`
|
||||
|
||||
filePath := createTempPythonFile(t, code)
|
||||
c := NewChunker(chunking.DefaultChunkOptions())
|
||||
|
||||
chunks, err := c.Chunk(context.Background(), filePath)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Should find all three functions
|
||||
functionNames := make(map[string]bool)
|
||||
for _, chunk := range chunks {
|
||||
if chunk.Type == chunking.ChunkTypeFunction {
|
||||
functionNames[chunk.Name] = true
|
||||
}
|
||||
}
|
||||
|
||||
assert.True(t, functionNames["first_function"])
|
||||
assert.True(t, functionNames["second_function"])
|
||||
assert.True(t, functionNames["third_function"])
|
||||
}
|
||||
|
||||
func TestChunker_Chunk_FileNotFound(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
c := NewChunker(chunking.DefaultChunkOptions())
|
||||
|
||||
_, err := c.Chunk(context.Background(), "/nonexistent/path/file.py")
|
||||
require.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "read file")
|
||||
}
|
||||
|
||||
func TestChunker_Chunk_EmptyFile(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
filePath := createTempPythonFile(t, "")
|
||||
c := NewChunker(chunking.DefaultChunkOptions())
|
||||
|
||||
chunks, err := c.Chunk(context.Background(), filePath)
|
||||
require.NoError(t, err)
|
||||
assert.Empty(t, chunks)
|
||||
}
|
||||
|
||||
func TestChunker_Chunk_OnlyComments(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
code := `# This is a comment
|
||||
# Another comment
|
||||
"""
|
||||
This is a module docstring
|
||||
"""
|
||||
`
|
||||
|
||||
filePath := createTempPythonFile(t, code)
|
||||
c := NewChunker(chunking.DefaultChunkOptions())
|
||||
|
||||
chunks, err := c.Chunk(context.Background(), filePath)
|
||||
require.NoError(t, err)
|
||||
// Comments and docstrings without code should not produce chunks
|
||||
assert.Empty(t, chunks)
|
||||
}
|
||||
|
||||
func TestChunker_Chunk_NestedClass(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
code := `class Outer:
|
||||
class Inner:
|
||||
def inner_method(self):
|
||||
pass
|
||||
|
||||
def outer_method(self):
|
||||
pass
|
||||
`
|
||||
|
||||
filePath := createTempPythonFile(t, code)
|
||||
c := NewChunker(chunking.DefaultChunkOptions())
|
||||
|
||||
chunks, err := c.Chunk(context.Background(), filePath)
|
||||
require.NoError(t, err)
|
||||
require.NotEmpty(t, chunks)
|
||||
|
||||
// Should find the Outer class at minimum
|
||||
var foundOuter bool
|
||||
for _, chunk := range chunks {
|
||||
if chunk.Name == "Outer" {
|
||||
foundOuter = true
|
||||
}
|
||||
}
|
||||
assert.True(t, foundOuter, "Should find 'Outer' class")
|
||||
}
|
||||
|
||||
func TestChunker_Chunk_Decorators(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
code := `@staticmethod
|
||||
def static_func():
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
def class_func(cls):
|
||||
pass
|
||||
|
||||
@property
|
||||
def my_property(self):
|
||||
return self._value
|
||||
`
|
||||
|
||||
filePath := createTempPythonFile(t, code)
|
||||
c := NewChunker(chunking.DefaultChunkOptions())
|
||||
|
||||
chunks, err := c.Chunk(context.Background(), filePath)
|
||||
require.NoError(t, err)
|
||||
require.NotEmpty(t, chunks)
|
||||
|
||||
// Should find decorated functions
|
||||
functionNames := make(map[string]bool)
|
||||
for _, chunk := range chunks {
|
||||
functionNames[chunk.Name] = true
|
||||
}
|
||||
|
||||
assert.True(t, functionNames["static_func"])
|
||||
assert.True(t, functionNames["class_func"])
|
||||
assert.True(t, functionNames["my_property"])
|
||||
}
|
||||
|
||||
func TestChunker_Chunk_AsyncFunction(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
code := `async def fetch_data(url):
|
||||
"""Fetches data from URL asynchronously."""
|
||||
pass
|
||||
|
||||
async def process_items(items):
|
||||
for item in items:
|
||||
await process(item)
|
||||
`
|
||||
|
||||
filePath := createTempPythonFile(t, code)
|
||||
c := NewChunker(chunking.DefaultChunkOptions())
|
||||
|
||||
chunks, err := c.Chunk(context.Background(), filePath)
|
||||
require.NoError(t, err)
|
||||
require.NotEmpty(t, chunks)
|
||||
|
||||
// Should find async functions
|
||||
functionNames := make(map[string]bool)
|
||||
for _, chunk := range chunks {
|
||||
functionNames[chunk.Name] = true
|
||||
}
|
||||
|
||||
assert.True(t, functionNames["fetch_data"])
|
||||
assert.True(t, functionNames["process_items"])
|
||||
}
|
||||
@@ -0,0 +1,140 @@
|
||||
// Package chunking provides AST-aware code chunking for semantic code search.
|
||||
// Chunks code files into logical units (functions, classes, methods) that preserve
|
||||
// semantic boundaries for better vector embedding and retrieval.
|
||||
package chunking
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// ChunkType represents the type of code chunk.
|
||||
type ChunkType string
|
||||
|
||||
const (
|
||||
// ChunkTypeFunction represents a standalone function.
|
||||
ChunkTypeFunction ChunkType = "function"
|
||||
// ChunkTypeMethod represents a method on a class/struct/type.
|
||||
ChunkTypeMethod ChunkType = "method"
|
||||
// ChunkTypeClass represents a class or struct definition.
|
||||
ChunkTypeClass ChunkType = "class"
|
||||
// ChunkTypeInterface represents an interface definition.
|
||||
ChunkTypeInterface ChunkType = "interface"
|
||||
// ChunkTypeType represents a type alias or type definition.
|
||||
ChunkTypeType ChunkType = "type"
|
||||
// ChunkTypeConst represents constant declarations.
|
||||
ChunkTypeConst ChunkType = "const"
|
||||
// ChunkTypeVar represents variable declarations.
|
||||
ChunkTypeVar ChunkType = "var"
|
||||
)
|
||||
|
||||
// Language represents a programming language.
|
||||
type Language string
|
||||
|
||||
const (
|
||||
// LanguageGo represents the Go programming language.
|
||||
LanguageGo Language = "go"
|
||||
// LanguagePython represents the Python programming language.
|
||||
LanguagePython Language = "python"
|
||||
// LanguageTypeScript represents the TypeScript programming language.
|
||||
LanguageTypeScript Language = "typescript"
|
||||
// LanguageJavaScript represents the JavaScript programming language.
|
||||
LanguageJavaScript Language = "javascript"
|
||||
)
|
||||
|
||||
// Chunk represents a semantic code chunk with AST-derived boundaries.
|
||||
type Chunk struct {
|
||||
Metadata map[string]interface{}
|
||||
FilePath string
|
||||
Language Language
|
||||
Type ChunkType
|
||||
Name string
|
||||
ParentName string
|
||||
Content string
|
||||
Signature string
|
||||
DocComment string
|
||||
StartLine int
|
||||
EndLine int
|
||||
}
|
||||
|
||||
// Identifier returns a human-readable identifier for this chunk.
|
||||
// Format: "ParentName.Name" for methods, "Name" for top-level.
|
||||
func (c *Chunk) Identifier() string {
|
||||
if c.ParentName != "" {
|
||||
return fmt.Sprintf("%s.%s", c.ParentName, c.Name)
|
||||
}
|
||||
return c.Name
|
||||
}
|
||||
|
||||
// LineRange returns a human-readable line range.
|
||||
// Format: "L123-L456"
|
||||
func (c *Chunk) LineRange() string {
|
||||
return fmt.Sprintf("L%d-L%d", c.StartLine, c.EndLine)
|
||||
}
|
||||
|
||||
// SearchableContent returns content optimized for semantic search.
|
||||
// Combines signature, doc comment, and content in a structured format.
|
||||
func (c *Chunk) SearchableContent() string {
|
||||
var parts []string
|
||||
|
||||
// Include signature for functions/methods
|
||||
if c.Signature != "" {
|
||||
parts = append(parts, c.Signature)
|
||||
}
|
||||
|
||||
// Include doc comment
|
||||
if c.DocComment != "" {
|
||||
parts = append(parts, c.DocComment)
|
||||
}
|
||||
|
||||
// Include actual content
|
||||
if c.Content != "" {
|
||||
parts = append(parts, c.Content)
|
||||
}
|
||||
|
||||
return strings.Join(parts, "\n\n")
|
||||
}
|
||||
|
||||
// Chunker is the interface for language-specific code chunkers.
|
||||
type Chunker interface {
|
||||
// Chunk parses a source file and returns semantic code chunks.
|
||||
// Returns an error if the file cannot be parsed or read.
|
||||
Chunk(ctx context.Context, filePath string) ([]Chunk, error)
|
||||
|
||||
// Language returns the language this chunker supports.
|
||||
Language() Language
|
||||
|
||||
// SupportedExtensions returns file extensions this chunker handles.
|
||||
// Example: []string{".go"} for Go chunker
|
||||
SupportedExtensions() []string
|
||||
}
|
||||
|
||||
// ChunkOptions provides options for chunking behavior.
|
||||
type ChunkOptions struct {
|
||||
// MaxChunkSize is the maximum size of a chunk in bytes.
|
||||
// Chunks larger than this will be split (respecting boundaries where possible).
|
||||
// 0 means no limit.
|
||||
MaxChunkSize int
|
||||
|
||||
// IncludeDocComments controls whether to include documentation comments.
|
||||
IncludeDocComments bool
|
||||
|
||||
// IncludePrivate controls whether to include private/unexported symbols.
|
||||
IncludePrivate bool
|
||||
|
||||
// MinLines is the minimum number of lines for a chunk to be included.
|
||||
// Chunks smaller than this will be skipped.
|
||||
// 0 means no minimum.
|
||||
MinLines int
|
||||
}
|
||||
|
||||
// DefaultChunkOptions returns sensible default options.
|
||||
func DefaultChunkOptions() ChunkOptions {
|
||||
return ChunkOptions{
|
||||
MaxChunkSize: 8192, // ~8KB per chunk (well under token limit)
|
||||
IncludeDocComments: true,
|
||||
IncludePrivate: true, // Include all symbols for comprehensive search
|
||||
MinLines: 0, // No minimum - include even single-line functions
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,213 @@
|
||||
package chunking
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
)
|
||||
|
||||
// =============================================================================
|
||||
// TESTS FOR Chunk METHODS
|
||||
// =============================================================================
|
||||
|
||||
func TestChunk_Identifier(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
expected string
|
||||
chunk Chunk
|
||||
}{
|
||||
// ===== GOOD CASES =====
|
||||
{
|
||||
name: "top-level function",
|
||||
chunk: Chunk{
|
||||
Name: "MyFunction",
|
||||
ParentName: "",
|
||||
},
|
||||
expected: "MyFunction",
|
||||
},
|
||||
{
|
||||
name: "method with parent",
|
||||
chunk: Chunk{
|
||||
Name: "Process",
|
||||
ParentName: "Handler",
|
||||
},
|
||||
expected: "Handler.Process",
|
||||
},
|
||||
{
|
||||
name: "nested method",
|
||||
chunk: Chunk{
|
||||
Name: "Validate",
|
||||
ParentName: "UserService",
|
||||
},
|
||||
expected: "UserService.Validate",
|
||||
},
|
||||
|
||||
// ===== EDGE CASES =====
|
||||
{
|
||||
name: "empty name",
|
||||
chunk: Chunk{
|
||||
Name: "",
|
||||
ParentName: "",
|
||||
},
|
||||
expected: "",
|
||||
},
|
||||
{
|
||||
name: "parent but no name",
|
||||
chunk: Chunk{
|
||||
Name: "",
|
||||
ParentName: "Parent",
|
||||
},
|
||||
expected: "Parent.",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
result := tt.chunk.Identifier()
|
||||
assert.Equal(t, tt.expected, result)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestChunk_LineRange(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
expected string
|
||||
chunk Chunk
|
||||
}{
|
||||
// ===== GOOD CASES =====
|
||||
{
|
||||
name: "single line",
|
||||
chunk: Chunk{
|
||||
StartLine: 10,
|
||||
EndLine: 10,
|
||||
},
|
||||
expected: "L10-L10",
|
||||
},
|
||||
{
|
||||
name: "multi-line",
|
||||
chunk: Chunk{
|
||||
StartLine: 25,
|
||||
EndLine: 50,
|
||||
},
|
||||
expected: "L25-L50",
|
||||
},
|
||||
|
||||
// ===== EDGE CASES =====
|
||||
{
|
||||
name: "line 1",
|
||||
chunk: Chunk{
|
||||
StartLine: 1,
|
||||
EndLine: 5,
|
||||
},
|
||||
expected: "L1-L5",
|
||||
},
|
||||
{
|
||||
name: "large line numbers",
|
||||
chunk: Chunk{
|
||||
StartLine: 1000,
|
||||
EndLine: 2500,
|
||||
},
|
||||
expected: "L1000-L2500",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
result := tt.chunk.LineRange()
|
||||
assert.Equal(t, tt.expected, result)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestChunk_SearchableContent(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
contains []string
|
||||
chunk Chunk
|
||||
}{
|
||||
// ===== GOOD CASES =====
|
||||
{
|
||||
name: "full chunk with all fields",
|
||||
chunk: Chunk{
|
||||
Signature: "func ProcessData(input []byte) error",
|
||||
DocComment: "// ProcessData handles incoming data",
|
||||
Content: "func ProcessData(input []byte) error {\n\treturn nil\n}",
|
||||
},
|
||||
contains: []string{
|
||||
"func ProcessData(input []byte) error",
|
||||
"ProcessData handles incoming data",
|
||||
"return nil",
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "only signature",
|
||||
chunk: Chunk{
|
||||
Signature: "func Hello()",
|
||||
},
|
||||
contains: []string{"func Hello()"},
|
||||
},
|
||||
{
|
||||
name: "only content",
|
||||
chunk: Chunk{
|
||||
Content: "some code here",
|
||||
},
|
||||
contains: []string{"some code here"},
|
||||
},
|
||||
|
||||
// ===== EDGE CASES =====
|
||||
{
|
||||
name: "empty chunk",
|
||||
chunk: Chunk{},
|
||||
contains: []string{},
|
||||
},
|
||||
{
|
||||
name: "only doc comment",
|
||||
chunk: Chunk{
|
||||
DocComment: "// Important documentation",
|
||||
},
|
||||
contains: []string{"Important documentation"},
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
result := tt.chunk.SearchableContent()
|
||||
for _, expected := range tt.contains {
|
||||
assert.Contains(t, result, expected)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestDefaultChunkOptions(t *testing.T) {
|
||||
opts := DefaultChunkOptions()
|
||||
|
||||
assert.Greater(t, opts.MaxChunkSize, 0, "MaxChunkSize should be positive")
|
||||
assert.True(t, opts.IncludeDocComments, "IncludeDocComments should be true by default")
|
||||
assert.True(t, opts.IncludePrivate, "IncludePrivate should be true by default")
|
||||
assert.Equal(t, 0, opts.MinLines, "MinLines should be 0 by default")
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// TESTS FOR ChunkType AND Language CONSTANTS
|
||||
// =============================================================================
|
||||
|
||||
func TestChunkType_Values(t *testing.T) {
|
||||
// Ensure all chunk types have expected values
|
||||
assert.Equal(t, ChunkType("function"), ChunkTypeFunction)
|
||||
assert.Equal(t, ChunkType("method"), ChunkTypeMethod)
|
||||
assert.Equal(t, ChunkType("class"), ChunkTypeClass)
|
||||
assert.Equal(t, ChunkType("interface"), ChunkTypeInterface)
|
||||
assert.Equal(t, ChunkType("type"), ChunkTypeType)
|
||||
assert.Equal(t, ChunkType("const"), ChunkTypeConst)
|
||||
assert.Equal(t, ChunkType("var"), ChunkTypeVar)
|
||||
}
|
||||
|
||||
func TestLanguage_Values(t *testing.T) {
|
||||
// Ensure all language types have expected values
|
||||
assert.Equal(t, Language("go"), LanguageGo)
|
||||
assert.Equal(t, Language("python"), LanguagePython)
|
||||
assert.Equal(t, Language("typescript"), LanguageTypeScript)
|
||||
assert.Equal(t, Language("javascript"), LanguageJavaScript)
|
||||
}
|
||||
@@ -0,0 +1,403 @@
|
||||
// Package typescript provides AST-aware chunking for TypeScript and JavaScript source files using tree-sitter.
|
||||
package typescript
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"os"
|
||||
"strings"
|
||||
|
||||
sitter "github.com/smacker/go-tree-sitter"
|
||||
"github.com/smacker/go-tree-sitter/typescript/typescript"
|
||||
|
||||
"github.com/lukaszraczylo/claude-mnemonic/internal/chunking"
|
||||
)
|
||||
|
||||
// Chunker implements AST-aware chunking for TypeScript/JavaScript files.
|
||||
type Chunker struct {
|
||||
parser *sitter.Parser
|
||||
options chunking.ChunkOptions
|
||||
}
|
||||
|
||||
// NewChunker creates a new TypeScript chunker.
|
||||
func NewChunker(options chunking.ChunkOptions) *Chunker {
|
||||
parser := sitter.NewParser()
|
||||
parser.SetLanguage(typescript.GetLanguage())
|
||||
|
||||
return &Chunker{
|
||||
options: options,
|
||||
parser: parser,
|
||||
}
|
||||
}
|
||||
|
||||
// Language returns the language this chunker supports.
|
||||
func (c *Chunker) Language() chunking.Language {
|
||||
return chunking.LanguageTypeScript
|
||||
}
|
||||
|
||||
// SupportedExtensions returns the file extensions this chunker handles.
|
||||
func (c *Chunker) SupportedExtensions() []string {
|
||||
return []string{".ts", ".tsx", ".js", ".jsx"}
|
||||
}
|
||||
|
||||
// Chunk parses a TypeScript/JavaScript source file and returns semantic code chunks.
|
||||
func (c *Chunker) Chunk(ctx context.Context, filePath string) ([]chunking.Chunk, error) {
|
||||
// Read file content
|
||||
content, err := os.ReadFile(filePath)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("read file: %w", err)
|
||||
}
|
||||
|
||||
// Parse the file
|
||||
tree, err := c.parser.ParseCtx(ctx, nil, content)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("parse TypeScript file: %w", err)
|
||||
}
|
||||
defer tree.Close()
|
||||
|
||||
sourceLines := strings.Split(string(content), "\n")
|
||||
chunks := make([]chunking.Chunk, 0)
|
||||
|
||||
// Walk the AST and extract chunks
|
||||
c.walkNode(tree.RootNode(), content, sourceLines, filePath, "", &chunks)
|
||||
|
||||
return chunks, nil
|
||||
}
|
||||
|
||||
// walkNode recursively walks the tree-sitter AST and extracts chunks.
|
||||
func (c *Chunker) walkNode(node *sitter.Node, source []byte, sourceLines []string, filePath string, parentName string, chunks *[]chunking.Chunk) {
|
||||
nodeType := node.Type()
|
||||
|
||||
switch nodeType {
|
||||
case "function_declaration":
|
||||
chunk := c.extractFunction(node, source, sourceLines, filePath, parentName)
|
||||
if chunk != nil {
|
||||
*chunks = append(*chunks, *chunk)
|
||||
}
|
||||
|
||||
case "method_definition":
|
||||
chunk := c.extractMethod(node, source, sourceLines, filePath, parentName)
|
||||
if chunk != nil {
|
||||
*chunks = append(*chunks, *chunk)
|
||||
}
|
||||
|
||||
case "arrow_function", "function_expression":
|
||||
// Handle arrow functions and function expressions assigned to variables
|
||||
chunk := c.extractFunctionExpression(node, source, sourceLines, filePath, parentName)
|
||||
if chunk != nil {
|
||||
*chunks = append(*chunks, *chunk)
|
||||
}
|
||||
|
||||
case "class_declaration":
|
||||
chunk := c.extractClass(node, source, sourceLines, filePath)
|
||||
if chunk != nil {
|
||||
*chunks = append(*chunks, *chunk)
|
||||
|
||||
// Walk class body to find methods
|
||||
for i := 0; i < int(node.ChildCount()); i++ {
|
||||
child := node.Child(i)
|
||||
if child.Type() == "class_body" {
|
||||
c.walkNode(child, source, sourceLines, filePath, chunk.Name, chunks)
|
||||
}
|
||||
}
|
||||
}
|
||||
return // Don't walk children again
|
||||
|
||||
case "interface_declaration":
|
||||
chunk := c.extractInterface(node, source, sourceLines, filePath)
|
||||
if chunk != nil {
|
||||
*chunks = append(*chunks, *chunk)
|
||||
}
|
||||
|
||||
case "type_alias_declaration":
|
||||
chunk := c.extractTypeAlias(node, source, sourceLines, filePath)
|
||||
if chunk != nil {
|
||||
*chunks = append(*chunks, *chunk)
|
||||
}
|
||||
}
|
||||
|
||||
// Walk all children
|
||||
for i := 0; i < int(node.ChildCount()); i++ {
|
||||
c.walkNode(node.Child(i), source, sourceLines, filePath, parentName, chunks)
|
||||
}
|
||||
}
|
||||
|
||||
// extractFunction extracts a function declaration.
|
||||
func (c *Chunker) extractFunction(node *sitter.Node, source []byte, sourceLines []string, filePath string, parentName string) *chunking.Chunk {
|
||||
name := c.findChildContent(node, "identifier", source)
|
||||
if name == "" {
|
||||
return nil
|
||||
}
|
||||
|
||||
startLine := int(node.StartPoint().Row) + 1
|
||||
endLine := int(node.EndPoint().Row) + 1
|
||||
|
||||
chunk := &chunking.Chunk{
|
||||
FilePath: filePath,
|
||||
Language: chunking.LanguageTypeScript,
|
||||
Type: chunking.ChunkTypeFunction,
|
||||
Name: name,
|
||||
ParentName: parentName,
|
||||
StartLine: startLine,
|
||||
EndLine: endLine,
|
||||
Content: c.extractLines(sourceLines, startLine, endLine),
|
||||
Signature: c.extractFunctionSignature(node, source, sourceLines),
|
||||
}
|
||||
|
||||
// Extract JSDoc comment
|
||||
if c.options.IncludeDocComments {
|
||||
chunk.DocComment = c.extractComment(node, source)
|
||||
}
|
||||
|
||||
return chunk
|
||||
}
|
||||
|
||||
// extractMethod extracts a method definition from a class.
|
||||
func (c *Chunker) extractMethod(node *sitter.Node, source []byte, sourceLines []string, filePath string, parentName string) *chunking.Chunk {
|
||||
name := c.findChildContent(node, "property_identifier", source)
|
||||
if name == "" {
|
||||
return nil
|
||||
}
|
||||
|
||||
// Skip private methods if configured
|
||||
if !c.options.IncludePrivate && strings.HasPrefix(name, "_") {
|
||||
return nil
|
||||
}
|
||||
|
||||
startLine := int(node.StartPoint().Row) + 1
|
||||
endLine := int(node.EndPoint().Row) + 1
|
||||
|
||||
chunk := &chunking.Chunk{
|
||||
FilePath: filePath,
|
||||
Language: chunking.LanguageTypeScript,
|
||||
Type: chunking.ChunkTypeMethod,
|
||||
Name: name,
|
||||
ParentName: parentName,
|
||||
StartLine: startLine,
|
||||
EndLine: endLine,
|
||||
Content: c.extractLines(sourceLines, startLine, endLine),
|
||||
Signature: c.extractMethodSignature(node, source, sourceLines),
|
||||
}
|
||||
|
||||
// Extract JSDoc comment
|
||||
if c.options.IncludeDocComments {
|
||||
chunk.DocComment = c.extractComment(node, source)
|
||||
}
|
||||
|
||||
return chunk
|
||||
}
|
||||
|
||||
// extractFunctionExpression extracts arrow functions and function expressions.
|
||||
func (c *Chunker) extractFunctionExpression(node *sitter.Node, source []byte, sourceLines []string, filePath string, parentName string) *chunking.Chunk {
|
||||
// Try to find the variable name from parent
|
||||
parent := node.Parent()
|
||||
if parent == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
var name string
|
||||
if parent.Type() == "variable_declarator" {
|
||||
name = c.findChildContent(parent, "identifier", source)
|
||||
} else if parent.Type() == "assignment_expression" {
|
||||
// Handle const foo = () => {}
|
||||
for i := 0; i < int(parent.ChildCount()); i++ {
|
||||
child := parent.Child(i)
|
||||
if child.Type() == "identifier" || child.Type() == "member_expression" {
|
||||
name = child.Content(source)
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if name == "" {
|
||||
return nil // Anonymous function, skip
|
||||
}
|
||||
|
||||
startLine := int(node.StartPoint().Row) + 1
|
||||
endLine := int(node.EndPoint().Row) + 1
|
||||
|
||||
chunk := &chunking.Chunk{
|
||||
FilePath: filePath,
|
||||
Language: chunking.LanguageTypeScript,
|
||||
Type: chunking.ChunkTypeFunction,
|
||||
Name: name,
|
||||
ParentName: parentName,
|
||||
StartLine: startLine,
|
||||
EndLine: endLine,
|
||||
Content: c.extractLines(sourceLines, startLine, endLine),
|
||||
}
|
||||
|
||||
return chunk
|
||||
}
|
||||
|
||||
// extractClass extracts a class declaration.
|
||||
func (c *Chunker) extractClass(node *sitter.Node, source []byte, sourceLines []string, filePath string) *chunking.Chunk {
|
||||
name := c.findChildContent(node, "type_identifier", source)
|
||||
if name == "" {
|
||||
return nil
|
||||
}
|
||||
|
||||
startLine := int(node.StartPoint().Row) + 1
|
||||
endLine := int(node.EndPoint().Row) + 1
|
||||
|
||||
chunk := &chunking.Chunk{
|
||||
FilePath: filePath,
|
||||
Language: chunking.LanguageTypeScript,
|
||||
Type: chunking.ChunkTypeClass,
|
||||
Name: name,
|
||||
StartLine: startLine,
|
||||
EndLine: endLine,
|
||||
Content: c.extractLines(sourceLines, startLine, endLine),
|
||||
Signature: c.extractClassSignature(node, source, sourceLines),
|
||||
}
|
||||
|
||||
// Extract JSDoc comment
|
||||
if c.options.IncludeDocComments {
|
||||
chunk.DocComment = c.extractComment(node, source)
|
||||
}
|
||||
|
||||
return chunk
|
||||
}
|
||||
|
||||
// extractInterface extracts an interface declaration.
|
||||
func (c *Chunker) extractInterface(node *sitter.Node, source []byte, sourceLines []string, filePath string) *chunking.Chunk {
|
||||
name := c.findChildContent(node, "type_identifier", source)
|
||||
if name == "" {
|
||||
return nil
|
||||
}
|
||||
|
||||
startLine := int(node.StartPoint().Row) + 1
|
||||
endLine := int(node.EndPoint().Row) + 1
|
||||
|
||||
chunk := &chunking.Chunk{
|
||||
FilePath: filePath,
|
||||
Language: chunking.LanguageTypeScript,
|
||||
Type: chunking.ChunkTypeInterface,
|
||||
Name: name,
|
||||
StartLine: startLine,
|
||||
EndLine: endLine,
|
||||
Content: c.extractLines(sourceLines, startLine, endLine),
|
||||
}
|
||||
|
||||
// Extract JSDoc comment
|
||||
if c.options.IncludeDocComments {
|
||||
chunk.DocComment = c.extractComment(node, source)
|
||||
}
|
||||
|
||||
return chunk
|
||||
}
|
||||
|
||||
// extractTypeAlias extracts a type alias declaration.
|
||||
func (c *Chunker) extractTypeAlias(node *sitter.Node, source []byte, sourceLines []string, filePath string) *chunking.Chunk {
|
||||
name := c.findChildContent(node, "type_identifier", source)
|
||||
if name == "" {
|
||||
return nil
|
||||
}
|
||||
|
||||
startLine := int(node.StartPoint().Row) + 1
|
||||
endLine := int(node.EndPoint().Row) + 1
|
||||
|
||||
chunk := &chunking.Chunk{
|
||||
FilePath: filePath,
|
||||
Language: chunking.LanguageTypeScript,
|
||||
Type: chunking.ChunkTypeType,
|
||||
Name: name,
|
||||
StartLine: startLine,
|
||||
EndLine: endLine,
|
||||
Content: c.extractLines(sourceLines, startLine, endLine),
|
||||
}
|
||||
|
||||
return chunk
|
||||
}
|
||||
|
||||
// findChildContent finds the first child of the given type and returns its content.
|
||||
func (c *Chunker) findChildContent(node *sitter.Node, childType string, source []byte) string {
|
||||
for i := 0; i < int(node.ChildCount()); i++ {
|
||||
child := node.Child(i)
|
||||
if child.Type() == childType {
|
||||
return child.Content(source)
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// extractFunctionSignature extracts the function signature.
|
||||
func (c *Chunker) extractFunctionSignature(node *sitter.Node, source []byte, sourceLines []string) string {
|
||||
startLine := int(node.StartPoint().Row) + 1
|
||||
|
||||
// Find the opening brace of the body
|
||||
for i := 0; i < int(node.ChildCount()); i++ {
|
||||
child := node.Child(i)
|
||||
if child.Type() == "statement_block" {
|
||||
endLine := int(child.StartPoint().Row) + 1
|
||||
return strings.TrimSpace(c.extractLines(sourceLines, startLine, endLine-1))
|
||||
}
|
||||
}
|
||||
|
||||
// Fallback: just return first line
|
||||
return strings.TrimSpace(c.extractLines(sourceLines, startLine, startLine))
|
||||
}
|
||||
|
||||
// extractMethodSignature extracts the method signature.
|
||||
func (c *Chunker) extractMethodSignature(node *sitter.Node, source []byte, sourceLines []string) string {
|
||||
startLine := int(node.StartPoint().Row) + 1
|
||||
|
||||
// Find the opening brace of the body
|
||||
for i := 0; i < int(node.ChildCount()); i++ {
|
||||
child := node.Child(i)
|
||||
if child.Type() == "statement_block" {
|
||||
endLine := int(child.StartPoint().Row) + 1
|
||||
return strings.TrimSpace(c.extractLines(sourceLines, startLine, endLine-1))
|
||||
}
|
||||
}
|
||||
|
||||
return strings.TrimSpace(c.extractLines(sourceLines, startLine, startLine))
|
||||
}
|
||||
|
||||
// extractClassSignature extracts the class declaration line.
|
||||
func (c *Chunker) extractClassSignature(node *sitter.Node, source []byte, sourceLines []string) string {
|
||||
startLine := int(node.StartPoint().Row) + 1
|
||||
|
||||
// Find the opening brace of the class body
|
||||
for i := 0; i < int(node.ChildCount()); i++ {
|
||||
child := node.Child(i)
|
||||
if child.Type() == "class_body" {
|
||||
endLine := int(child.StartPoint().Row) + 1
|
||||
return strings.TrimSpace(c.extractLines(sourceLines, startLine, endLine-1))
|
||||
}
|
||||
}
|
||||
|
||||
return strings.TrimSpace(c.extractLines(sourceLines, startLine, startLine))
|
||||
}
|
||||
|
||||
// extractComment extracts JSDoc or other comments from a node.
|
||||
func (c *Chunker) extractComment(node *sitter.Node, source []byte) string {
|
||||
// Check previous sibling for comment
|
||||
prevSibling := node.PrevSibling()
|
||||
if prevSibling != nil && prevSibling.Type() == "comment" {
|
||||
comment := prevSibling.Content(source)
|
||||
// Remove comment markers
|
||||
comment = strings.TrimPrefix(comment, "/**")
|
||||
comment = strings.TrimPrefix(comment, "/*")
|
||||
comment = strings.TrimSuffix(comment, "*/")
|
||||
comment = strings.TrimPrefix(comment, "//")
|
||||
return strings.TrimSpace(comment)
|
||||
}
|
||||
|
||||
return ""
|
||||
}
|
||||
|
||||
// extractLines extracts a range of lines from source (1-indexed, inclusive).
|
||||
func (c *Chunker) extractLines(lines []string, start, end int) string {
|
||||
if start < 1 || end < start || start > len(lines) {
|
||||
return ""
|
||||
}
|
||||
|
||||
startIdx := start - 1
|
||||
endIdx := end
|
||||
if endIdx > len(lines) {
|
||||
endIdx = len(lines)
|
||||
}
|
||||
|
||||
return strings.Join(lines[startIdx:endIdx], "\n")
|
||||
}
|
||||
@@ -0,0 +1,398 @@
|
||||
package typescript
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/lukaszraczylo/claude-mnemonic/internal/chunking"
|
||||
)
|
||||
|
||||
// =============================================================================
|
||||
// TEST HELPERS
|
||||
// =============================================================================
|
||||
|
||||
func createTempTSFile(t *testing.T, content string, ext string) string {
|
||||
t.Helper()
|
||||
|
||||
tmpDir := t.TempDir()
|
||||
filePath := filepath.Join(tmpDir, "test"+ext)
|
||||
|
||||
err := os.WriteFile(filePath, []byte(content), 0600)
|
||||
require.NoError(t, err)
|
||||
|
||||
return filePath
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// TESTS FOR Chunker
|
||||
// =============================================================================
|
||||
|
||||
func TestNewChunker(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
opts := chunking.DefaultChunkOptions()
|
||||
c := NewChunker(opts)
|
||||
|
||||
assert.NotNil(t, c)
|
||||
assert.NotNil(t, c.parser)
|
||||
}
|
||||
|
||||
func TestChunker_Language(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
c := NewChunker(chunking.DefaultChunkOptions())
|
||||
|
||||
assert.Equal(t, chunking.LanguageTypeScript, c.Language())
|
||||
}
|
||||
|
||||
func TestChunker_SupportedExtensions(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
c := NewChunker(chunking.DefaultChunkOptions())
|
||||
exts := c.SupportedExtensions()
|
||||
|
||||
assert.Contains(t, exts, ".ts")
|
||||
assert.Contains(t, exts, ".tsx")
|
||||
assert.Contains(t, exts, ".js")
|
||||
assert.Contains(t, exts, ".jsx")
|
||||
}
|
||||
|
||||
func TestChunker_Chunk_SimpleFunction(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
code := `function greet(name: string): string {
|
||||
return "Hello, " + name + "!";
|
||||
}
|
||||
`
|
||||
|
||||
filePath := createTempTSFile(t, code, ".ts")
|
||||
c := NewChunker(chunking.DefaultChunkOptions())
|
||||
|
||||
chunks, err := c.Chunk(context.Background(), filePath)
|
||||
require.NoError(t, err)
|
||||
require.NotEmpty(t, chunks)
|
||||
|
||||
// Should find the greet function
|
||||
var foundGreet bool
|
||||
for _, chunk := range chunks {
|
||||
if chunk.Name == "greet" {
|
||||
foundGreet = true
|
||||
assert.Equal(t, chunking.ChunkTypeFunction, chunk.Type)
|
||||
assert.Equal(t, chunking.LanguageTypeScript, chunk.Language)
|
||||
assert.Contains(t, chunk.Content, "function greet")
|
||||
}
|
||||
}
|
||||
assert.True(t, foundGreet, "Should find 'greet' function")
|
||||
}
|
||||
|
||||
func TestChunker_Chunk_ClassWithMethods(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
code := `class Calculator {
|
||||
add(a: number, b: number): number {
|
||||
return a + b;
|
||||
}
|
||||
|
||||
multiply(a: number, b: number): number {
|
||||
return a * b;
|
||||
}
|
||||
}
|
||||
`
|
||||
|
||||
filePath := createTempTSFile(t, code, ".ts")
|
||||
c := NewChunker(chunking.DefaultChunkOptions())
|
||||
|
||||
chunks, err := c.Chunk(context.Background(), filePath)
|
||||
require.NoError(t, err)
|
||||
require.NotEmpty(t, chunks)
|
||||
|
||||
// Should find the Calculator class and its methods
|
||||
var foundClass, foundAdd, foundMultiply bool
|
||||
for _, chunk := range chunks {
|
||||
switch chunk.Name {
|
||||
case "Calculator":
|
||||
foundClass = true
|
||||
assert.Equal(t, chunking.ChunkTypeClass, chunk.Type)
|
||||
case "add":
|
||||
foundAdd = true
|
||||
assert.Equal(t, chunking.ChunkTypeMethod, chunk.Type)
|
||||
assert.Equal(t, "Calculator", chunk.ParentName)
|
||||
case "multiply":
|
||||
foundMultiply = true
|
||||
assert.Equal(t, chunking.ChunkTypeMethod, chunk.Type)
|
||||
assert.Equal(t, "Calculator", chunk.ParentName)
|
||||
}
|
||||
}
|
||||
|
||||
assert.True(t, foundClass, "Should find 'Calculator' class")
|
||||
assert.True(t, foundAdd, "Should find 'add' method")
|
||||
assert.True(t, foundMultiply, "Should find 'multiply' method")
|
||||
}
|
||||
|
||||
func TestChunker_Chunk_Interface(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
code := `interface User {
|
||||
id: number;
|
||||
name: string;
|
||||
email: string;
|
||||
}
|
||||
|
||||
interface Authenticator {
|
||||
login(username: string, password: string): boolean;
|
||||
logout(): void;
|
||||
}
|
||||
`
|
||||
|
||||
filePath := createTempTSFile(t, code, ".ts")
|
||||
c := NewChunker(chunking.DefaultChunkOptions())
|
||||
|
||||
chunks, err := c.Chunk(context.Background(), filePath)
|
||||
require.NoError(t, err)
|
||||
require.NotEmpty(t, chunks)
|
||||
|
||||
// Should find interfaces
|
||||
interfaceNames := make(map[string]bool)
|
||||
for _, chunk := range chunks {
|
||||
if chunk.Type == chunking.ChunkTypeInterface {
|
||||
interfaceNames[chunk.Name] = true
|
||||
}
|
||||
}
|
||||
|
||||
assert.True(t, interfaceNames["User"])
|
||||
assert.True(t, interfaceNames["Authenticator"])
|
||||
}
|
||||
|
||||
func TestChunker_Chunk_TypeAlias(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
code := `type UserID = string;
|
||||
|
||||
type Handler = (event: Event) => void;
|
||||
|
||||
type Result<T> = { success: true; data: T } | { success: false; error: Error };
|
||||
`
|
||||
|
||||
filePath := createTempTSFile(t, code, ".ts")
|
||||
c := NewChunker(chunking.DefaultChunkOptions())
|
||||
|
||||
chunks, err := c.Chunk(context.Background(), filePath)
|
||||
require.NoError(t, err)
|
||||
require.NotEmpty(t, chunks)
|
||||
|
||||
// Should find type aliases
|
||||
typeNames := make(map[string]bool)
|
||||
for _, chunk := range chunks {
|
||||
if chunk.Type == chunking.ChunkTypeType {
|
||||
typeNames[chunk.Name] = true
|
||||
}
|
||||
}
|
||||
|
||||
assert.True(t, typeNames["UserID"])
|
||||
assert.True(t, typeNames["Handler"])
|
||||
assert.True(t, typeNames["Result"])
|
||||
}
|
||||
|
||||
func TestChunker_Chunk_ArrowFunction(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
code := `const add = (a: number, b: number): number => a + b;
|
||||
|
||||
const greet = (name: string): string => {
|
||||
return "Hello, " + name;
|
||||
};
|
||||
`
|
||||
|
||||
filePath := createTempTSFile(t, code, ".ts")
|
||||
c := NewChunker(chunking.DefaultChunkOptions())
|
||||
|
||||
_, err := c.Chunk(context.Background(), filePath)
|
||||
require.NoError(t, err)
|
||||
// Arrow functions may or may not be captured depending on AST structure
|
||||
// At minimum, no error should occur
|
||||
}
|
||||
|
||||
func TestChunker_Chunk_FileNotFound(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
c := NewChunker(chunking.DefaultChunkOptions())
|
||||
|
||||
_, err := c.Chunk(context.Background(), "/nonexistent/path/file.ts")
|
||||
require.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "read file")
|
||||
}
|
||||
|
||||
func TestChunker_Chunk_EmptyFile(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
filePath := createTempTSFile(t, "", ".ts")
|
||||
c := NewChunker(chunking.DefaultChunkOptions())
|
||||
|
||||
chunks, err := c.Chunk(context.Background(), filePath)
|
||||
require.NoError(t, err)
|
||||
assert.Empty(t, chunks)
|
||||
}
|
||||
|
||||
func TestChunker_Chunk_OnlyComments(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
code := `// This is a comment
|
||||
/* Another comment */
|
||||
/**
|
||||
* JSDoc comment
|
||||
*/
|
||||
`
|
||||
|
||||
filePath := createTempTSFile(t, code, ".ts")
|
||||
c := NewChunker(chunking.DefaultChunkOptions())
|
||||
|
||||
chunks, err := c.Chunk(context.Background(), filePath)
|
||||
require.NoError(t, err)
|
||||
// Comments without code should not produce chunks
|
||||
assert.Empty(t, chunks)
|
||||
}
|
||||
|
||||
func TestChunker_Chunk_AsyncFunction(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
code := `async function fetchData(url: string): Promise<any> {
|
||||
const response = await fetch(url);
|
||||
return response.json();
|
||||
}
|
||||
|
||||
async function processItems(items: string[]): Promise<void> {
|
||||
for (const item of items) {
|
||||
await process(item);
|
||||
}
|
||||
}
|
||||
`
|
||||
|
||||
filePath := createTempTSFile(t, code, ".ts")
|
||||
c := NewChunker(chunking.DefaultChunkOptions())
|
||||
|
||||
chunks, err := c.Chunk(context.Background(), filePath)
|
||||
require.NoError(t, err)
|
||||
require.NotEmpty(t, chunks)
|
||||
|
||||
// Should find async functions
|
||||
functionNames := make(map[string]bool)
|
||||
for _, chunk := range chunks {
|
||||
if chunk.Type == chunking.ChunkTypeFunction {
|
||||
functionNames[chunk.Name] = true
|
||||
}
|
||||
}
|
||||
|
||||
assert.True(t, functionNames["fetchData"])
|
||||
assert.True(t, functionNames["processItems"])
|
||||
}
|
||||
|
||||
func TestChunker_Chunk_ExportedFunction(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
code := `export function publicFunction(): void {
|
||||
console.log("public");
|
||||
}
|
||||
|
||||
export default function defaultExport(): void {
|
||||
console.log("default");
|
||||
}
|
||||
`
|
||||
|
||||
filePath := createTempTSFile(t, code, ".ts")
|
||||
c := NewChunker(chunking.DefaultChunkOptions())
|
||||
|
||||
chunks, err := c.Chunk(context.Background(), filePath)
|
||||
require.NoError(t, err)
|
||||
require.NotEmpty(t, chunks)
|
||||
|
||||
// Should find exported functions
|
||||
functionNames := make(map[string]bool)
|
||||
for _, chunk := range chunks {
|
||||
if chunk.Type == chunking.ChunkTypeFunction {
|
||||
functionNames[chunk.Name] = true
|
||||
}
|
||||
}
|
||||
|
||||
assert.True(t, functionNames["publicFunction"])
|
||||
assert.True(t, functionNames["defaultExport"])
|
||||
}
|
||||
|
||||
func TestChunker_Chunk_JSXFile(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
code := `function Button({ label }: { label: string }) {
|
||||
return <button>{label}</button>;
|
||||
}
|
||||
|
||||
function App() {
|
||||
return (
|
||||
<div>
|
||||
<Button label="Click me" />
|
||||
</div>
|
||||
);
|
||||
}
|
||||
`
|
||||
|
||||
filePath := createTempTSFile(t, code, ".tsx")
|
||||
c := NewChunker(chunking.DefaultChunkOptions())
|
||||
|
||||
chunks, err := c.Chunk(context.Background(), filePath)
|
||||
require.NoError(t, err)
|
||||
require.NotEmpty(t, chunks)
|
||||
|
||||
// Should find JSX components as functions
|
||||
functionNames := make(map[string]bool)
|
||||
for _, chunk := range chunks {
|
||||
if chunk.Type == chunking.ChunkTypeFunction {
|
||||
functionNames[chunk.Name] = true
|
||||
}
|
||||
}
|
||||
|
||||
assert.True(t, functionNames["Button"])
|
||||
assert.True(t, functionNames["App"])
|
||||
}
|
||||
|
||||
func TestChunker_Chunk_JavaScript(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
code := `function simpleFunc() {
|
||||
return 42;
|
||||
}
|
||||
|
||||
class MyClass {
|
||||
constructor() {
|
||||
this.value = 0;
|
||||
}
|
||||
|
||||
getValue() {
|
||||
return this.value;
|
||||
}
|
||||
}
|
||||
`
|
||||
|
||||
filePath := createTempTSFile(t, code, ".js")
|
||||
c := NewChunker(chunking.DefaultChunkOptions())
|
||||
|
||||
chunks, err := c.Chunk(context.Background(), filePath)
|
||||
require.NoError(t, err)
|
||||
require.NotEmpty(t, chunks)
|
||||
|
||||
// Should find JavaScript functions and classes
|
||||
var foundFunc, foundClass bool
|
||||
for _, chunk := range chunks {
|
||||
if chunk.Name == "simpleFunc" {
|
||||
foundFunc = true
|
||||
}
|
||||
if chunk.Name == "MyClass" {
|
||||
foundClass = true
|
||||
}
|
||||
}
|
||||
|
||||
assert.True(t, foundFunc, "Should find 'simpleFunc' function")
|
||||
assert.True(t, foundClass, "Should find 'MyClass' class")
|
||||
}
|
||||
+108
-40
@@ -5,6 +5,7 @@ import (
|
||||
"encoding/json"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
)
|
||||
@@ -36,41 +37,46 @@ var CriticalConcepts = []string{
|
||||
}
|
||||
|
||||
// Config holds the application configuration.
|
||||
// Field order optimized for memory alignment (fieldalignment).
|
||||
type Config struct {
|
||||
// Worker settings
|
||||
WorkerPort int `json:"worker_port"`
|
||||
|
||||
// Database settings
|
||||
DBPath string `json:"db_path"`
|
||||
MaxConns int `json:"max_conns"`
|
||||
|
||||
// SDK Agent settings
|
||||
Model string `json:"model"`
|
||||
ClaudeCodePath string `json:"claude_code_path"`
|
||||
|
||||
// Embedding settings
|
||||
EmbeddingModel string `json:"embedding_model"` // e.g., "bge-v1.5"
|
||||
|
||||
// Reranking settings (cross-encoder)
|
||||
RerankingEnabled bool `json:"reranking_enabled"` // Enable cross-encoder reranking
|
||||
RerankingCandidates int `json:"reranking_candidates"` // Number of candidates to retrieve before reranking (default 100)
|
||||
RerankingResults int `json:"reranking_results"` // Number of results to return after reranking (default 10)
|
||||
RerankingAlpha float64 `json:"reranking_alpha"` // Weight for combining scores: alpha*rerank + (1-alpha)*original (default 0.7)
|
||||
RerankingMinImprovement float64 `json:"reranking_min_improvement"` // Minimum rank improvement to trigger reranking (default 0, always rerank)
|
||||
RerankingPureMode bool `json:"reranking_pure_mode"` // Use pure cross-encoder scores without combining with bi-encoder (default false)
|
||||
|
||||
// Context injection settings
|
||||
ContextObservations int `json:"context_observations"`
|
||||
ContextFullCount int `json:"context_full_count"`
|
||||
ContextSessionCount int `json:"context_session_count"`
|
||||
ContextShowReadTokens bool `json:"context_show_read_tokens"`
|
||||
ContextShowWorkTokens bool `json:"context_show_work_tokens"`
|
||||
ContextFullField string `json:"context_full_field"`
|
||||
ContextShowLastSummary bool `json:"context_show_last_summary"`
|
||||
ContextObsTypes []string `json:"context_obs_types"`
|
||||
DBPath string `json:"db_path"`
|
||||
Model string `json:"model"`
|
||||
ClaudeCodePath string `json:"claude_code_path"`
|
||||
EmbeddingModel string `json:"embedding_model"`
|
||||
VectorStorageStrategy string `json:"vector_storage_strategy"`
|
||||
ContextObsConcepts []string `json:"context_obs_concepts"`
|
||||
ContextRelevanceThreshold float64 `json:"context_relevance_threshold"` // 0.0-1.0, minimum similarity for inclusion
|
||||
ContextMaxPromptResults int `json:"context_max_prompt_results"` // Max results per prompt (0 = threshold only)
|
||||
ContextObsTypes []string `json:"context_obs_types"`
|
||||
ContextFullCount int `json:"context_full_count"`
|
||||
GraphBranchFactor int `json:"graph_branch_factor"`
|
||||
GraphEdgeWeight float64 `json:"graph_edge_weight"`
|
||||
ContextRelevanceThreshold float64 `json:"context_relevance_threshold"`
|
||||
RerankingCandidates int `json:"reranking_candidates"`
|
||||
WorkerPort int `json:"worker_port"`
|
||||
DeduplicationThreshold float64 `json:"deduplication_threshold"`
|
||||
RerankingMinImprovement float64 `json:"reranking_min_improvement"`
|
||||
ContextObservations int `json:"context_observations"`
|
||||
ContextMaxPromptResults int `json:"context_max_prompt_results"`
|
||||
ContextSessionCount int `json:"context_session_count"`
|
||||
MaxConns int `json:"max_conns"`
|
||||
RerankingAlpha float64 `json:"reranking_alpha"`
|
||||
GraphMaxHops int `json:"graph_max_hops"`
|
||||
RerankingResults int `json:"reranking_results"`
|
||||
GraphRebuildIntervalMin int `json:"graph_rebuild_interval_min"`
|
||||
HubThreshold int `json:"hub_threshold"`
|
||||
ObservationRetentionDays int `json:"observation_retention_days"`
|
||||
MaintenanceIntervalHours int `json:"maintenance_interval_hours"`
|
||||
ContextMaxTokensStartup int `json:"context_max_tokens_startup"`
|
||||
ContextMaxTokensPrompt int `json:"context_max_tokens_prompt"`
|
||||
ContextShowWorkTokens bool `json:"context_show_work_tokens"`
|
||||
ContextShowReadTokens bool `json:"context_show_read_tokens"`
|
||||
RerankingPureMode bool `json:"reranking_pure_mode"`
|
||||
GraphEnabled bool `json:"graph_enabled"`
|
||||
DeduplicationEnabled bool `json:"deduplication_enabled"`
|
||||
MaintenanceEnabled bool `json:"maintenance_enabled"`
|
||||
RerankingEnabled bool `json:"reranking_enabled"`
|
||||
ContextShowLastSummary bool `json:"context_show_last_summary"`
|
||||
CleanupStaleObservations bool `json:"cleanup_stale_observations"`
|
||||
}
|
||||
|
||||
var (
|
||||
@@ -96,8 +102,9 @@ func SettingsPath() string {
|
||||
}
|
||||
|
||||
// EnsureDataDir creates the data directory if it doesn't exist.
|
||||
// Uses 0700 permissions (owner-only) for security.
|
||||
func EnsureDataDir() error {
|
||||
return os.MkdirAll(DataDir(), 0750)
|
||||
return os.MkdirAll(DataDir(), 0700)
|
||||
}
|
||||
|
||||
// EnsureSettings creates a default settings file if it doesn't exist.
|
||||
@@ -143,11 +150,18 @@ func Default() *Config {
|
||||
MaxConns: 4,
|
||||
Model: DefaultModel,
|
||||
EmbeddingModel: DefaultEmbeddingModel,
|
||||
RerankingEnabled: true, // Enable by default for improved relevance
|
||||
RerankingCandidates: 100, // Retrieve top 100 candidates
|
||||
RerankingResults: 10, // Return top 10 after reranking
|
||||
RerankingAlpha: 0.7, // Favor cross-encoder score
|
||||
RerankingMinImprovement: 0, // Always apply reranking
|
||||
RerankingEnabled: true, // Enable by default for improved relevance
|
||||
RerankingCandidates: 100, // Retrieve top 100 candidates
|
||||
RerankingResults: 10, // Return top 10 after reranking
|
||||
RerankingAlpha: 0.7, // Favor cross-encoder score
|
||||
RerankingMinImprovement: 0, // Always apply reranking
|
||||
GraphEnabled: true, // Enable graph-aware search by default
|
||||
GraphMaxHops: 2, // Two-hop traversal
|
||||
GraphBranchFactor: 5, // Expand top 5 neighbors per node
|
||||
GraphEdgeWeight: 0.3, // Minimum edge weight to follow
|
||||
GraphRebuildIntervalMin: 60, // Rebuild graph every 60 minutes
|
||||
VectorStorageStrategy: "hub", // Hub storage strategy (LEANN-inspired)
|
||||
HubThreshold: 5, // Require 5+ accesses to store embedding
|
||||
ContextObservations: 100,
|
||||
ContextFullCount: 25,
|
||||
ContextSessionCount: 10,
|
||||
@@ -157,8 +171,16 @@ func Default() *Config {
|
||||
ContextShowLastSummary: true,
|
||||
ContextObsTypes: DefaultObservationTypes,
|
||||
ContextObsConcepts: DefaultObservationConcepts,
|
||||
ContextRelevanceThreshold: 0.3, // Minimum 30% similarity to include
|
||||
ContextMaxPromptResults: 10, // Cap at 10 results max (0 = no cap, threshold only)
|
||||
ContextRelevanceThreshold: 0.3, // Minimum 30% similarity to include
|
||||
ContextMaxPromptResults: 10, // Cap at 10 results max (0 = no cap, threshold only)
|
||||
ContextMaxTokensStartup: 16000, // Max tokens for SessionStart context injection
|
||||
ContextMaxTokensPrompt: 8000, // Max tokens for UserPromptSubmit context injection
|
||||
DeduplicationEnabled: true, // Enable write-time vector dedup
|
||||
DeduplicationThreshold: 0.9, // Similarity threshold for merging (0.9 = very similar)
|
||||
MaintenanceEnabled: true, // Enable scheduled maintenance
|
||||
MaintenanceIntervalHours: 6, // Run every 6 hours
|
||||
ObservationRetentionDays: 0, // 0 = no age-based deletion (keep all)
|
||||
CleanupStaleObservations: false, // Don't auto-cleanup stale observations
|
||||
}
|
||||
}
|
||||
|
||||
@@ -233,6 +255,52 @@ func Load() (*Config, error) {
|
||||
if v, ok := settings["CLAUDE_MNEMONIC_CONTEXT_MAX_PROMPT_RESULTS"].(float64); ok && v >= 0 {
|
||||
cfg.ContextMaxPromptResults = int(v)
|
||||
}
|
||||
// Graph settings
|
||||
if v, ok := settings["CLAUDE_MNEMONIC_GRAPH_ENABLED"].(bool); ok {
|
||||
cfg.GraphEnabled = v
|
||||
}
|
||||
if v, ok := settings["CLAUDE_MNEMONIC_GRAPH_MAX_HOPS"].(float64); ok && v > 0 {
|
||||
cfg.GraphMaxHops = int(v)
|
||||
}
|
||||
if v, ok := settings["CLAUDE_MNEMONIC_GRAPH_BRANCH_FACTOR"].(float64); ok && v > 0 {
|
||||
cfg.GraphBranchFactor = int(v)
|
||||
}
|
||||
if v, ok := settings["CLAUDE_MNEMONIC_GRAPH_EDGE_WEIGHT"].(float64); ok && v >= 0 && v <= 1 {
|
||||
cfg.GraphEdgeWeight = v
|
||||
}
|
||||
if v, ok := settings["CLAUDE_MNEMONIC_GRAPH_REBUILD_INTERVAL_MIN"].(float64); ok && v > 0 {
|
||||
cfg.GraphRebuildIntervalMin = int(v)
|
||||
}
|
||||
// Vector storage settings (LEANN Phase 2)
|
||||
if v, ok := settings["CLAUDE_MNEMONIC_VECTOR_STORAGE_STRATEGY"].(string); ok && v != "" {
|
||||
cfg.VectorStorageStrategy = v
|
||||
}
|
||||
if v, ok := settings["CLAUDE_MNEMONIC_HUB_THRESHOLD"].(float64); ok && v > 0 {
|
||||
cfg.HubThreshold = int(v)
|
||||
}
|
||||
if v, ok := settings["CLAUDE_MNEMONIC_CONTEXT_MAX_TOKENS_STARTUP"].(float64); ok && v > 0 {
|
||||
cfg.ContextMaxTokensStartup = int(v)
|
||||
}
|
||||
if v, ok := settings["CLAUDE_MNEMONIC_CONTEXT_MAX_TOKENS_PROMPT"].(float64); ok && v > 0 {
|
||||
cfg.ContextMaxTokensPrompt = int(v)
|
||||
}
|
||||
// Deduplication settings
|
||||
if v, ok := settings["CLAUDE_MNEMONIC_DEDUP_ENABLED"].(bool); ok {
|
||||
cfg.DeduplicationEnabled = v
|
||||
}
|
||||
if v, ok := settings["CLAUDE_MNEMONIC_DEDUP_THRESHOLD"].(float64); ok && v > 0 && v <= 1 {
|
||||
cfg.DeduplicationThreshold = v
|
||||
}
|
||||
|
||||
// Also support env vars for dedup settings
|
||||
if v := os.Getenv("CLAUDE_MNEMONIC_DEDUP_ENABLED"); v != "" {
|
||||
cfg.DeduplicationEnabled = v == "true" || v == "1"
|
||||
}
|
||||
if v := os.Getenv("CLAUDE_MNEMONIC_DEDUP_THRESHOLD"); v != "" {
|
||||
if f, err := strconv.ParseFloat(v, 64); err == nil && f > 0 && f <= 1 {
|
||||
cfg.DeduplicationThreshold = f
|
||||
}
|
||||
}
|
||||
|
||||
return cfg, nil
|
||||
}
|
||||
|
||||
@@ -121,8 +121,8 @@ func (s *ConfigSuite) TestLoad_TableDriven() {
|
||||
tests := []struct {
|
||||
name string
|
||||
settingsJSON string
|
||||
expectedPort int
|
||||
expectedModel string
|
||||
expectedPort int
|
||||
expectedObsObs int
|
||||
}{
|
||||
{
|
||||
@@ -183,12 +183,12 @@ func (s *ConfigSuite) TestLoad_TableDriven() {
|
||||
s.Require().NoError(err)
|
||||
|
||||
if tt.settingsJSON != "" {
|
||||
err := os.WriteFile(
|
||||
writeErr := os.WriteFile(
|
||||
filepath.Join(tempDir, ".claude-mnemonic", "settings.json"),
|
||||
[]byte(tt.settingsJSON),
|
||||
0600,
|
||||
)
|
||||
s.Require().NoError(err)
|
||||
s.Require().NoError(writeErr)
|
||||
}
|
||||
|
||||
cfg, err := Load()
|
||||
|
||||
@@ -214,9 +214,9 @@ func (s *ConflictStore) CleanupSupersededObservations(ctx context.Context, proje
|
||||
// GetConflictsWithDetails retrieves all conflicts with observation titles for display.
|
||||
func (s *ConflictStore) GetConflictsWithDetails(ctx context.Context, project string, limit int) ([]*ConflictWithDetails, error) {
|
||||
var results []struct {
|
||||
ObservationConflict
|
||||
NewerTitle sql.NullString `gorm:"column:newer_title"`
|
||||
OlderTitle sql.NullString `gorm:"column:older_title"`
|
||||
ObservationConflict
|
||||
}
|
||||
|
||||
err := s.db.WithContext(ctx).
|
||||
|
||||
+54
-17
@@ -9,27 +9,13 @@ import (
|
||||
"time"
|
||||
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/clause"
|
||||
)
|
||||
|
||||
// EnsureSessionExists creates a session if it doesn't exist.
|
||||
// Uses INSERT OR IGNORE pattern for atomic idempotent creation (single query instead of COUNT + INSERT).
|
||||
// This is shared between stores to avoid duplication.
|
||||
func EnsureSessionExists(ctx context.Context, db *gorm.DB, sdkSessionID, project string) error {
|
||||
// Check if session exists
|
||||
var count int64
|
||||
err := db.WithContext(ctx).
|
||||
Model(&SDKSession{}).
|
||||
Where("sdk_session_id = ?", sdkSessionID).
|
||||
Count(&count).Error
|
||||
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if count > 0 {
|
||||
return nil // Session exists
|
||||
}
|
||||
|
||||
// Auto-create session
|
||||
now := time.Now()
|
||||
session := &SDKSession{
|
||||
ClaudeSessionID: sdkSessionID,
|
||||
@@ -41,7 +27,14 @@ func EnsureSessionExists(ctx context.Context, db *gorm.DB, sdkSessionID, project
|
||||
PromptCounter: 0,
|
||||
}
|
||||
|
||||
return db.WithContext(ctx).Create(session).Error
|
||||
// Use INSERT OR IGNORE - single query, atomic operation
|
||||
// If session already exists (conflict on sdk_session_id), do nothing
|
||||
return db.WithContext(ctx).
|
||||
Clauses(clause.OnConflict{
|
||||
Columns: []clause.Column{{Name: "sdk_session_id"}},
|
||||
DoNothing: true,
|
||||
}).
|
||||
Create(session).Error
|
||||
}
|
||||
|
||||
// sqlNullString creates a sql.NullString from a string.
|
||||
@@ -52,8 +45,13 @@ func sqlNullString(s string) sql.NullString {
|
||||
return sql.NullString{String: s, Valid: true}
|
||||
}
|
||||
|
||||
// MaxPaginationLimit is the maximum allowed limit for pagination queries.
|
||||
// This protects against resource exhaustion from excessively large requests.
|
||||
const MaxPaginationLimit = 1000
|
||||
|
||||
// ParseLimitParam parses the "limit" query parameter from an HTTP request.
|
||||
// Returns defaultLimit if the parameter is missing or invalid.
|
||||
// Note: This does NOT enforce a maximum limit. Use ParseLimitParamWithMax for that.
|
||||
func ParseLimitParam(r *http.Request, defaultLimit int) int {
|
||||
if l := r.URL.Query().Get("limit"); l != "" {
|
||||
if parsed, err := strconv.Atoi(l); err == nil && parsed > 0 {
|
||||
@@ -62,3 +60,42 @@ func ParseLimitParam(r *http.Request, defaultLimit int) int {
|
||||
}
|
||||
return defaultLimit
|
||||
}
|
||||
|
||||
// ParseLimitParamWithMax parses the "limit" query parameter with a maximum cap.
|
||||
// Returns min(parsed, maxLimit) or defaultLimit if missing/invalid.
|
||||
// If maxLimit is 0, uses MaxPaginationLimit (1000).
|
||||
func ParseLimitParamWithMax(r *http.Request, defaultLimit, maxLimit int) int {
|
||||
if maxLimit <= 0 {
|
||||
maxLimit = MaxPaginationLimit
|
||||
}
|
||||
limit := ParseLimitParam(r, defaultLimit)
|
||||
if limit > maxLimit {
|
||||
return maxLimit
|
||||
}
|
||||
return limit
|
||||
}
|
||||
|
||||
// ParseOffsetParam parses the "offset" query parameter from an HTTP request.
|
||||
// Returns 0 if the parameter is missing or invalid.
|
||||
func ParseOffsetParam(r *http.Request) int {
|
||||
if o := r.URL.Query().Get("offset"); o != "" {
|
||||
if parsed, err := strconv.Atoi(o); err == nil && parsed >= 0 {
|
||||
return parsed
|
||||
}
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
// PaginationParams holds pagination parameters.
|
||||
type PaginationParams struct {
|
||||
Limit int
|
||||
Offset int
|
||||
}
|
||||
|
||||
// ParsePaginationParams parses both limit and offset from an HTTP request.
|
||||
func ParsePaginationParams(r *http.Request, defaultLimit int) PaginationParams {
|
||||
return PaginationParams{
|
||||
Limit: ParseLimitParam(r, defaultLimit),
|
||||
Offset: ParseOffsetParam(r),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -322,6 +322,244 @@ func runMigrations(db *gorm.DB, sqlDB *sql.DB) error {
|
||||
return tx.Migrator().DropTable("observation_relations")
|
||||
},
|
||||
},
|
||||
|
||||
// Migration 012: Query optimization indexes
|
||||
// Adds covering and composite indexes for common query patterns
|
||||
{
|
||||
ID: "012_query_optimization_indexes",
|
||||
Migrate: func(tx *gorm.DB) error {
|
||||
sqls := []string{
|
||||
// Composite index for observation queries by project + scope + importance
|
||||
// Covers the common pattern: WHERE project = ? OR scope = 'global' ORDER BY importance_score DESC
|
||||
`CREATE INDEX IF NOT EXISTS idx_observations_project_scope_importance
|
||||
ON observations(project, scope, importance_score DESC, created_at_epoch DESC)`,
|
||||
|
||||
// Covering index for observation retrieval (includes most used columns)
|
||||
// Allows index-only scans for listing queries
|
||||
`CREATE INDEX IF NOT EXISTS idx_observations_project_covering
|
||||
ON observations(project, scope, is_superseded, importance_score DESC)
|
||||
WHERE is_superseded = 0 OR is_superseded IS NULL`,
|
||||
|
||||
// Index for session summary lookups
|
||||
`CREATE INDEX IF NOT EXISTS idx_summaries_project_importance
|
||||
ON session_summaries(project, importance_score DESC, created_at_epoch DESC)`,
|
||||
|
||||
// Index for prompt retrieval by session
|
||||
`CREATE INDEX IF NOT EXISTS idx_prompts_session_number
|
||||
ON user_prompts(claude_session_id, prompt_number)`,
|
||||
|
||||
// Index for pattern queries by frequency
|
||||
`CREATE INDEX IF NOT EXISTS idx_patterns_frequency
|
||||
ON patterns(frequency DESC, last_seen_at_epoch DESC)
|
||||
WHERE is_deprecated = 0`,
|
||||
}
|
||||
for _, s := range sqls {
|
||||
if err := tx.Exec(s).Error; err != nil {
|
||||
// Non-fatal: index may already exist or fail for benign reasons
|
||||
continue
|
||||
}
|
||||
}
|
||||
return nil
|
||||
},
|
||||
Rollback: func(tx *gorm.DB) error {
|
||||
sqls := []string{
|
||||
"DROP INDEX IF EXISTS idx_observations_project_scope_importance",
|
||||
"DROP INDEX IF EXISTS idx_observations_project_covering",
|
||||
"DROP INDEX IF EXISTS idx_summaries_project_importance",
|
||||
"DROP INDEX IF EXISTS idx_prompts_session_number",
|
||||
"DROP INDEX IF EXISTS idx_patterns_frequency",
|
||||
}
|
||||
for _, s := range sqls {
|
||||
if err := tx.Exec(s).Error; err != nil {
|
||||
continue
|
||||
}
|
||||
}
|
||||
return nil
|
||||
},
|
||||
},
|
||||
|
||||
// Migration 013: Add archival columns to observations
|
||||
{
|
||||
ID: "013_observation_archival",
|
||||
Migrate: func(tx *gorm.DB) error {
|
||||
sqls := []string{
|
||||
// Add archival columns
|
||||
`ALTER TABLE observations ADD COLUMN is_archived INTEGER DEFAULT 0`,
|
||||
`ALTER TABLE observations ADD COLUMN archived_at_epoch INTEGER`,
|
||||
`ALTER TABLE observations ADD COLUMN archived_reason TEXT`,
|
||||
// Index for archived observations
|
||||
`CREATE INDEX IF NOT EXISTS idx_observations_archived ON observations(is_archived)`,
|
||||
// Composite index for filtering active (non-archived) observations
|
||||
`CREATE INDEX IF NOT EXISTS idx_observations_active
|
||||
ON observations(project, is_archived, is_superseded, importance_score DESC)
|
||||
WHERE (is_archived = 0 OR is_archived IS NULL)`,
|
||||
}
|
||||
for _, s := range sqls {
|
||||
if err := tx.Exec(s).Error; err != nil {
|
||||
// Non-fatal: column may already exist
|
||||
continue
|
||||
}
|
||||
}
|
||||
return nil
|
||||
},
|
||||
Rollback: func(tx *gorm.DB) error {
|
||||
// SQLite doesn't support DROP COLUMN in older versions
|
||||
// but for newer versions we can try
|
||||
sqls := []string{
|
||||
"DROP INDEX IF EXISTS idx_observations_active",
|
||||
"DROP INDEX IF EXISTS idx_observations_archived",
|
||||
}
|
||||
for _, s := range sqls {
|
||||
_ = tx.Exec(s).Error
|
||||
}
|
||||
return nil
|
||||
},
|
||||
},
|
||||
|
||||
// Migration 014: Add performance-critical indexes for common query patterns
|
||||
{
|
||||
ID: "014_performance_indexes",
|
||||
Migrate: func(tx *gorm.DB) error {
|
||||
sqls := []string{
|
||||
// Index for batch ID lookups (IN queries)
|
||||
`CREATE INDEX IF NOT EXISTS idx_observations_id_covering
|
||||
ON observations(id, project, scope, importance_score)`,
|
||||
|
||||
// Index for vector search result fetching
|
||||
`CREATE INDEX IF NOT EXISTS idx_vectors_doc_type_project
|
||||
ON vectors(doc_type, project, scope)`,
|
||||
|
||||
// Index for session summaries by project
|
||||
`CREATE INDEX IF NOT EXISTS idx_summaries_project_created
|
||||
ON session_summaries(project, created_at_epoch DESC)`,
|
||||
|
||||
// Index for user prompts by session
|
||||
`CREATE INDEX IF NOT EXISTS idx_prompts_session_created
|
||||
ON user_prompts(claude_session_id, created_at_epoch DESC)`,
|
||||
|
||||
// Index for patterns by type and project
|
||||
`CREATE INDEX IF NOT EXISTS idx_patterns_type_project
|
||||
ON patterns(type, project, frequency DESC)
|
||||
WHERE is_deprecated = 0`,
|
||||
|
||||
// Index for observation relations
|
||||
`CREATE INDEX IF NOT EXISTS idx_relations_source_type
|
||||
ON observation_relations(source_observation_id, relation_type)`,
|
||||
`CREATE INDEX IF NOT EXISTS idx_relations_target_type
|
||||
ON observation_relations(target_observation_id, relation_type)`,
|
||||
}
|
||||
for _, s := range sqls {
|
||||
if err := tx.Exec(s).Error; err != nil {
|
||||
// Non-fatal: index may already exist
|
||||
continue
|
||||
}
|
||||
}
|
||||
return nil
|
||||
},
|
||||
Rollback: func(tx *gorm.DB) error {
|
||||
sqls := []string{
|
||||
"DROP INDEX IF EXISTS idx_observations_id_covering",
|
||||
"DROP INDEX IF EXISTS idx_vectors_doc_type_project",
|
||||
"DROP INDEX IF EXISTS idx_summaries_project_created",
|
||||
"DROP INDEX IF EXISTS idx_prompts_session_created",
|
||||
"DROP INDEX IF EXISTS idx_patterns_type_project",
|
||||
"DROP INDEX IF EXISTS idx_relations_source_type",
|
||||
"DROP INDEX IF EXISTS idx_relations_target_type",
|
||||
}
|
||||
for _, s := range sqls {
|
||||
_ = tx.Exec(s).Error
|
||||
}
|
||||
return nil
|
||||
},
|
||||
},
|
||||
|
||||
// Migration 015: Add optimized composite indexes for common query patterns
|
||||
{
|
||||
ID: "015_optimized_composite_indexes",
|
||||
Migrate: func(tx *gorm.DB) error {
|
||||
sqls := []string{
|
||||
// Composite index for GetRecentObservations with project+scope filtering
|
||||
// Covers: WHERE (project = ? OR scope = 'global') ORDER BY importance_score DESC, created_at_epoch DESC
|
||||
`CREATE INDEX IF NOT EXISTS idx_observations_project_scope_created
|
||||
ON observations(project, scope, created_at_epoch DESC, importance_score DESC)`,
|
||||
|
||||
// Index for scope='global' queries (common pattern in search)
|
||||
`CREATE INDEX IF NOT EXISTS idx_observations_global_scope
|
||||
ON observations(scope, importance_score DESC, created_at_epoch DESC)
|
||||
WHERE scope = 'global'`,
|
||||
|
||||
// Index for vector search result deduplication by observation
|
||||
`CREATE INDEX IF NOT EXISTS idx_vectors_observation_lookup
|
||||
ON vectors(doc_type, sqlite_id, project)
|
||||
WHERE doc_type = 'observation'`,
|
||||
|
||||
// Index for FTS search result ordering
|
||||
`CREATE INDEX IF NOT EXISTS idx_observations_fts_ordering
|
||||
ON observations(project, importance_score DESC)
|
||||
WHERE (is_archived = 0 OR is_archived IS NULL) AND (is_superseded = 0 OR is_superseded IS NULL)`,
|
||||
}
|
||||
for _, s := range sqls {
|
||||
if err := tx.Exec(s).Error; err != nil {
|
||||
// Non-fatal: index may already exist
|
||||
continue
|
||||
}
|
||||
}
|
||||
return nil
|
||||
},
|
||||
Rollback: func(tx *gorm.DB) error {
|
||||
sqls := []string{
|
||||
"DROP INDEX IF EXISTS idx_observations_project_scope_created",
|
||||
"DROP INDEX IF EXISTS idx_observations_global_scope",
|
||||
"DROP INDEX IF EXISTS idx_vectors_observation_lookup",
|
||||
"DROP INDEX IF EXISTS idx_observations_fts_ordering",
|
||||
}
|
||||
for _, s := range sqls {
|
||||
_ = tx.Exec(s).Error
|
||||
}
|
||||
return nil
|
||||
},
|
||||
},
|
||||
|
||||
// Migration 016: Add covering indexes for relation joins and active observations
|
||||
{
|
||||
ID: "016_relation_and_active_indexes",
|
||||
Migrate: func(tx *gorm.DB) error {
|
||||
sqls := []string{
|
||||
// Covering index for observation relation joins (common JOIN patterns)
|
||||
// Speeds up queries like: JOIN observation_relations ON source_id = obs.id WHERE relation_type = ?
|
||||
`CREATE INDEX IF NOT EXISTS idx_relations_source_type_target
|
||||
ON observation_relations(source_observation_id, relation_type, target_observation_id)`,
|
||||
|
||||
// Covering index for reverse relation lookups
|
||||
`CREATE INDEX IF NOT EXISTS idx_relations_target_type_source
|
||||
ON observation_relations(target_observation_id, relation_type, source_observation_id)`,
|
||||
|
||||
// Partial index for active (non-archived, non-superseded) observations
|
||||
// Optimizes activeObservationFilter queries
|
||||
`CREATE INDEX IF NOT EXISTS idx_observations_active
|
||||
ON observations(project, importance_score DESC, created_at_epoch DESC)
|
||||
WHERE (is_archived = 0 OR is_archived IS NULL) AND (is_superseded = 0 OR is_superseded IS NULL)`,
|
||||
}
|
||||
for _, s := range sqls {
|
||||
if err := tx.Exec(s).Error; err != nil {
|
||||
// Non-fatal: index may already exist
|
||||
continue
|
||||
}
|
||||
}
|
||||
return nil
|
||||
},
|
||||
Rollback: func(tx *gorm.DB) error {
|
||||
sqls := []string{
|
||||
"DROP INDEX IF EXISTS idx_relations_source_type_target",
|
||||
"DROP INDEX IF EXISTS idx_relations_target_type_source",
|
||||
"DROP INDEX IF EXISTS idx_observations_active",
|
||||
}
|
||||
for _, s := range sqls {
|
||||
_ = tx.Exec(s).Error
|
||||
}
|
||||
return nil
|
||||
},
|
||||
},
|
||||
})
|
||||
|
||||
if err := m.Migrate(); err != nil {
|
||||
|
||||
+65
-71
@@ -17,18 +17,18 @@ import (
|
||||
|
||||
// SDKSession represents a Claude Code session.
|
||||
type SDKSession struct {
|
||||
ID int64 `gorm:"primaryKey;autoIncrement"`
|
||||
ClaudeSessionID string `gorm:"uniqueIndex;not null"`
|
||||
SDKSessionID sql.NullString `gorm:"uniqueIndex"`
|
||||
Project string `gorm:"index;not null"`
|
||||
Status string `gorm:"type:text;check:status IN ('active', 'completed', 'failed');default:'active';index"`
|
||||
StartedAt string `gorm:"not null"`
|
||||
SDKSessionID sql.NullString `gorm:"uniqueIndex"`
|
||||
UserPrompt sql.NullString
|
||||
WorkerPort sql.NullInt64
|
||||
PromptCounter int `gorm:"default:0"`
|
||||
Status string `gorm:"type:text;check:status IN ('active', 'completed', 'failed');default:'active';index"`
|
||||
StartedAt string `gorm:"not null"`
|
||||
StartedAtEpoch int64 `gorm:"index:idx_sessions_started,sort:desc;not null"`
|
||||
CompletedAt sql.NullString
|
||||
WorkerPort sql.NullInt64
|
||||
CompletedAtEpoch sql.NullInt64
|
||||
ID int64 `gorm:"primaryKey;autoIncrement"`
|
||||
PromptCounter int `gorm:"default:0"`
|
||||
StartedAtEpoch int64 `gorm:"index:idx_sessions_started,sort:desc;not null"`
|
||||
}
|
||||
|
||||
func (SDKSession) TableName() string { return "sdk_sessions" }
|
||||
@@ -45,36 +45,34 @@ func (s *SDKSession) BeforeCreate(tx *gorm.DB) error {
|
||||
}
|
||||
|
||||
// Observation represents a stored observation (learning).
|
||||
// Field order optimized for memory alignment (fieldalignment).
|
||||
type Observation struct {
|
||||
ID int64 `gorm:"primaryKey;autoIncrement"`
|
||||
SDKSessionID string `gorm:"index;not null"`
|
||||
Project string `gorm:"index;not null"`
|
||||
Scope models.ObservationScope `gorm:"type:text;default:'project';check:scope IN ('project', 'global');index:idx_observations_scope;index:idx_observations_project_scope,priority:2"`
|
||||
Type models.ObservationType `gorm:"type:text;check:type IN ('decision', 'bugfix', 'feature', 'refactor', 'discovery', 'change');index;not null"`
|
||||
|
||||
// Content fields
|
||||
Title sql.NullString `gorm:"type:text"`
|
||||
Subtitle sql.NullString `gorm:"type:text"`
|
||||
Facts models.JSONStringArray `gorm:"type:text"` // JSON array
|
||||
Narrative sql.NullString `gorm:"type:text"`
|
||||
Concepts models.JSONStringArray `gorm:"type:text"` // JSON array
|
||||
FilesRead models.JSONStringArray `gorm:"type:text"` // JSON array
|
||||
FilesModified models.JSONStringArray `gorm:"type:text"` // JSON array
|
||||
FileMtimes models.JSONInt64Map `gorm:"type:text"` // JSON object
|
||||
|
||||
// Metadata
|
||||
FileMtimes models.JSONInt64Map `gorm:"type:text"`
|
||||
SDKSessionID string `gorm:"index;not null"`
|
||||
Project string `gorm:"index:idx_observations_project;index:idx_observations_project_created,priority:1;not null"`
|
||||
Scope models.ObservationScope `gorm:"type:text;default:'project';check:scope IN ('project', 'global');index:idx_observations_scope;index:idx_observations_project_scope,priority:2"`
|
||||
Type models.ObservationType `gorm:"type:text;check:type IN ('decision', 'bugfix', 'feature', 'refactor', 'discovery', 'change');index;not null"`
|
||||
CreatedAt string `gorm:"not null"`
|
||||
Facts models.JSONStringArray `gorm:"type:text"`
|
||||
Narrative sql.NullString `gorm:"type:text"`
|
||||
Concepts models.JSONStringArray `gorm:"type:text"`
|
||||
FilesRead models.JSONStringArray `gorm:"type:text"`
|
||||
FilesModified models.JSONStringArray `gorm:"type:text"`
|
||||
Subtitle sql.NullString `gorm:"type:text"`
|
||||
Title sql.NullString `gorm:"type:text"`
|
||||
ArchivedReason sql.NullString
|
||||
ScoreUpdatedAt sql.NullInt64 `gorm:"column:score_updated_at_epoch;index:idx_observations_score_updated"`
|
||||
PromptNumber sql.NullInt64
|
||||
DiscoveryTokens int64 `gorm:"default:0"`
|
||||
CreatedAt string `gorm:"not null"`
|
||||
CreatedAtEpoch int64 `gorm:"index:idx_observations_created,sort:desc;not null"`
|
||||
|
||||
// Importance scoring fields
|
||||
ArchivedAt sql.NullInt64 `gorm:"column:archived_at_epoch"`
|
||||
LastRetrievedAt sql.NullInt64 `gorm:"column:last_retrieved_at_epoch"`
|
||||
ID int64 `gorm:"primaryKey;autoIncrement"`
|
||||
ImportanceScore float64 `gorm:"type:real;default:1.0;index:idx_observations_importance,priority:1,sort:desc"`
|
||||
UserFeedback int `gorm:"default:0"`
|
||||
RetrievalCount int `gorm:"default:0"`
|
||||
LastRetrievedAt sql.NullInt64 `gorm:"column:last_retrieved_at_epoch"`
|
||||
ScoreUpdatedAt sql.NullInt64 `gorm:"column:score_updated_at_epoch;index:idx_observations_score_updated"`
|
||||
IsSuperseded int `gorm:"default:0;index:idx_observations_superseded,priority:1"`
|
||||
CreatedAtEpoch int64 `gorm:"index:idx_observations_created,sort:desc;index:idx_observations_project_created,priority:2,sort:desc;not null"`
|
||||
DiscoveryTokens int64 `gorm:"default:0"`
|
||||
IsSuperseded int `gorm:"default:0;index:idx_observations_superseded;index:idx_observations_active,priority:2"`
|
||||
IsArchived int `gorm:"default:0;index:idx_observations_archived;index:idx_observations_active,priority:1"`
|
||||
}
|
||||
|
||||
func (Observation) TableName() string { return "observations" }
|
||||
@@ -95,23 +93,19 @@ func (o *Observation) BeforeCreate(tx *gorm.DB) error {
|
||||
|
||||
// SessionSummary represents a session summary.
|
||||
type SessionSummary struct {
|
||||
ID int64 `gorm:"primaryKey;autoIncrement"`
|
||||
SDKSessionID string `gorm:"index;not null"`
|
||||
Project string `gorm:"index;not null"`
|
||||
|
||||
// Summary fields (nullable TEXT)
|
||||
Request sql.NullString
|
||||
Investigated sql.NullString
|
||||
Learned sql.NullString
|
||||
Completed sql.NullString
|
||||
NextSteps sql.NullString `gorm:"column:next_steps"`
|
||||
Notes sql.NullString
|
||||
|
||||
// Metadata
|
||||
PromptNumber sql.NullInt64
|
||||
DiscoveryTokens int64 `gorm:"default:0"`
|
||||
CreatedAt string `gorm:"not null"`
|
||||
CreatedAtEpoch int64 `gorm:"index:idx_summaries_created,sort:desc;not null"`
|
||||
SDKSessionID string `gorm:"index;not null"`
|
||||
Project string `gorm:"index;not null"`
|
||||
Completed sql.NullString
|
||||
Investigated sql.NullString
|
||||
Learned sql.NullString
|
||||
NextSteps sql.NullString `gorm:"column:next_steps"`
|
||||
Notes sql.NullString
|
||||
Request sql.NullString
|
||||
PromptNumber sql.NullInt64
|
||||
ID int64 `gorm:"primaryKey;autoIncrement"`
|
||||
DiscoveryTokens int64 `gorm:"default:0"`
|
||||
CreatedAtEpoch int64 `gorm:"index:idx_summaries_created,sort:desc;not null"`
|
||||
}
|
||||
|
||||
func (SessionSummary) TableName() string { return "session_summaries" }
|
||||
@@ -129,12 +123,12 @@ func (s *SessionSummary) BeforeCreate(tx *gorm.DB) error {
|
||||
|
||||
// UserPrompt represents a user prompt.
|
||||
type UserPrompt struct {
|
||||
ID int64 `gorm:"primaryKey;autoIncrement"`
|
||||
ClaudeSessionID string `gorm:"index;not null;uniqueIndex:idx_user_prompts_session_number_unique,priority:1"`
|
||||
PromptNumber int `gorm:"index;not null;uniqueIndex:idx_user_prompts_session_number_unique,priority:2"`
|
||||
PromptText string `gorm:"type:text;not null"`
|
||||
MatchedObservations int `gorm:"default:0"`
|
||||
CreatedAt string `gorm:"not null"`
|
||||
ID int64 `gorm:"primaryKey;autoIncrement"`
|
||||
PromptNumber int `gorm:"index;not null;uniqueIndex:idx_user_prompts_session_number_unique,priority:2"`
|
||||
MatchedObservations int `gorm:"default:0"`
|
||||
CreatedAtEpoch int64 `gorm:"index:idx_prompts_created,sort:desc;not null"`
|
||||
}
|
||||
|
||||
@@ -153,16 +147,16 @@ func (p *UserPrompt) BeforeCreate(tx *gorm.DB) error {
|
||||
|
||||
// ObservationConflict tracks conflicts between observations.
|
||||
type ObservationConflict struct {
|
||||
ID int64 `gorm:"primaryKey;autoIncrement"`
|
||||
NewerObsID int64 `gorm:"index:idx_conflicts_newer;not null"`
|
||||
OlderObsID int64 `gorm:"index:idx_conflicts_older;not null"`
|
||||
ConflictType models.ConflictType `gorm:"type:text;check:conflict_type IN ('superseded', 'contradicts', 'outdated_pattern');not null"`
|
||||
Resolution models.ConflictResolution `gorm:"type:text;check:resolution IN ('prefer_newer', 'prefer_older', 'manual');not null"`
|
||||
Reason sql.NullString `gorm:"type:text"`
|
||||
DetectedAt string `gorm:"not null"`
|
||||
DetectedAtEpoch int64 `gorm:"index:idx_conflicts_unresolved,priority:2,sort:desc;not null"`
|
||||
Resolved int `gorm:"default:0;index:idx_conflicts_unresolved,priority:1"`
|
||||
Reason sql.NullString `gorm:"type:text"`
|
||||
ResolvedAt sql.NullString
|
||||
ID int64 `gorm:"primaryKey;autoIncrement"`
|
||||
NewerObsID int64 `gorm:"index:idx_conflicts_newer;not null"`
|
||||
OlderObsID int64 `gorm:"index:idx_conflicts_older;not null"`
|
||||
DetectedAtEpoch int64 `gorm:"index:idx_conflicts_unresolved,priority:2,sort:desc;not null"`
|
||||
Resolved int `gorm:"default:0;index:idx_conflicts_unresolved,priority:1"`
|
||||
}
|
||||
|
||||
func (ObservationConflict) TableName() string { return "observation_conflicts" }
|
||||
@@ -180,14 +174,14 @@ func (c *ObservationConflict) BeforeCreate(tx *gorm.DB) error {
|
||||
|
||||
// ObservationRelation tracks relationships between observations.
|
||||
type ObservationRelation struct {
|
||||
RelationType models.RelationType `gorm:"type:text;check:relation_type IN ('causes', 'fixes', 'supersedes', 'depends_on', 'relates_to', 'evolves_from');index:idx_relations_type;uniqueIndex:idx_relations_unique,priority:3;not null"`
|
||||
DetectionSource models.RelationDetectionSource `gorm:"type:text;check:detection_source IN ('file_overlap', 'embedding_similarity', 'temporal_proximity', 'narrative_mention', 'concept_overlap', 'type_progression');not null"`
|
||||
CreatedAt string `gorm:"not null"`
|
||||
Reason sql.NullString `gorm:"type:text"`
|
||||
ID int64 `gorm:"primaryKey;autoIncrement"`
|
||||
SourceID int64 `gorm:"index:idx_relations_source;index:idx_relations_both,priority:1;uniqueIndex:idx_relations_unique,priority:1;not null"`
|
||||
TargetID int64 `gorm:"index:idx_relations_target;index:idx_relations_both,priority:2;uniqueIndex:idx_relations_unique,priority:2;not null"`
|
||||
RelationType models.RelationType `gorm:"type:text;check:relation_type IN ('causes', 'fixes', 'supersedes', 'depends_on', 'relates_to', 'evolves_from');index:idx_relations_type;uniqueIndex:idx_relations_unique,priority:3;not null"`
|
||||
Confidence float64 `gorm:"type:real;default:0.5;index:idx_relations_confidence,sort:desc;not null"`
|
||||
DetectionSource models.RelationDetectionSource `gorm:"type:text;check:detection_source IN ('file_overlap', 'embedding_similarity', 'temporal_proximity', 'narrative_mention', 'concept_overlap', 'type_progression');not null"`
|
||||
Reason sql.NullString `gorm:"type:text"`
|
||||
CreatedAt string `gorm:"not null"`
|
||||
CreatedAtEpoch int64 `gorm:"not null"`
|
||||
}
|
||||
|
||||
@@ -209,21 +203,21 @@ func (r *ObservationRelation) BeforeCreate(tx *gorm.DB) error {
|
||||
|
||||
// Pattern represents a detected recurring pattern.
|
||||
type Pattern struct {
|
||||
ID int64 `gorm:"primaryKey;autoIncrement"`
|
||||
Status models.PatternStatus `gorm:"type:text;default:'active';check:status IN ('active', 'deprecated', 'merged');index"`
|
||||
Name string `gorm:"type:text;not null"`
|
||||
Type models.PatternType `gorm:"type:text;check:type IN ('bug', 'refactor', 'architecture', 'anti-pattern', 'best-practice');index;not null"`
|
||||
Description sql.NullString `gorm:"type:text"`
|
||||
Signature models.JSONStringArray `gorm:"type:text"` // JSON array of keywords
|
||||
CreatedAt string `gorm:"not null"`
|
||||
LastSeenAt string `gorm:"not null"`
|
||||
Signature models.JSONStringArray `gorm:"type:text"`
|
||||
Projects models.JSONStringArray `gorm:"type:text"`
|
||||
ObservationIDs models.JSONInt64Array `gorm:"type:text"`
|
||||
Recommendation sql.NullString `gorm:"type:text"`
|
||||
Frequency int `gorm:"default:1;index:idx_patterns_frequency,sort:desc"`
|
||||
Projects models.JSONStringArray `gorm:"type:text"` // JSON array
|
||||
ObservationIDs models.JSONInt64Array `gorm:"type:text"` // JSON array
|
||||
Status models.PatternStatus `gorm:"type:text;default:'active';check:status IN ('active', 'deprecated', 'merged');index"`
|
||||
Description sql.NullString `gorm:"type:text"`
|
||||
MergedIntoID sql.NullInt64
|
||||
Frequency int `gorm:"default:1;index:idx_patterns_frequency,sort:desc"`
|
||||
Confidence float64 `gorm:"type:real;default:0.5;index:idx_patterns_confidence,sort:desc"`
|
||||
LastSeenAt string `gorm:"not null"`
|
||||
ID int64 `gorm:"primaryKey;autoIncrement"`
|
||||
LastSeenAtEpoch int64 `gorm:"index:idx_patterns_last_seen,sort:desc;not null"`
|
||||
CreatedAt string `gorm:"not null"`
|
||||
CreatedAtEpoch int64 `gorm:"not null"`
|
||||
}
|
||||
|
||||
@@ -256,8 +250,8 @@ func (p *Pattern) BeforeCreate(tx *gorm.DB) error {
|
||||
// ConceptWeight stores configurable weights for importance scoring.
|
||||
type ConceptWeight struct {
|
||||
Concept string `gorm:"primaryKey;type:text"`
|
||||
Weight float64 `gorm:"type:real;not null;default:0.1"`
|
||||
UpdatedAt string `gorm:"not null"`
|
||||
Weight float64 `gorm:"type:real;not null;default:0.1"`
|
||||
}
|
||||
|
||||
func (ConceptWeight) TableName() string { return "concept_weights" }
|
||||
|
||||
@@ -5,39 +5,126 @@ import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"gorm.io/gorm"
|
||||
|
||||
"github.com/lukaszraczylo/claude-mnemonic/pkg/models"
|
||||
"github.com/rs/zerolog/log"
|
||||
)
|
||||
|
||||
// MaxObservationsPerProject is the maximum number of observations to keep per project.
|
||||
const MaxObservationsPerProject = 100
|
||||
|
||||
// cleanupQueueSize is the buffer size for the cleanup queue.
|
||||
const cleanupQueueSize = 100
|
||||
|
||||
// commonWords is a package-level set for O(1) lookup of stop words.
|
||||
// Created once at init time to avoid repeated map allocations.
|
||||
var commonWords = map[string]struct{}{
|
||||
"the": {}, "and": {}, "or": {}, "but": {}, "in": {},
|
||||
"on": {}, "at": {}, "to": {}, "for": {}, "of": {},
|
||||
"with": {}, "by": {}, "from": {}, "as": {}, "is": {},
|
||||
"was": {}, "are": {}, "were": {}, "be": {}, "been": {},
|
||||
"being": {}, "have": {}, "has": {}, "had": {}, "do": {},
|
||||
"does": {}, "did": {}, "will": {}, "would": {}, "should": {},
|
||||
"could": {}, "may": {}, "might": {}, "must": {}, "can": {},
|
||||
}
|
||||
|
||||
// CleanupFunc is a callback for when observations are cleaned up.
|
||||
// Receives the IDs of deleted observations for downstream cleanup (e.g., vector DB).
|
||||
type CleanupFunc func(ctx context.Context, deletedIDs []int64)
|
||||
|
||||
// ObservationStore provides observation-related database operations using GORM.
|
||||
type ObservationStore struct {
|
||||
db *gorm.DB
|
||||
rawDB *sql.DB
|
||||
cleanupFunc CleanupFunc
|
||||
conflictStore interface{} // Placeholder for ConflictStore (Phase 4)
|
||||
relationStore interface{} // Placeholder for RelationStore (Phase 4)
|
||||
conflictStore any
|
||||
relationStore any
|
||||
db *gorm.DB
|
||||
rawDB *sql.DB
|
||||
cleanupFunc CleanupFunc
|
||||
cleanupQueue chan string
|
||||
stopCleanup chan struct{}
|
||||
cleanupWg sync.WaitGroup
|
||||
cleanupOnce sync.Once
|
||||
cleanupStarted atomic.Bool
|
||||
}
|
||||
|
||||
// NewObservationStore creates a new observation store.
|
||||
// The conflictStore and relationStore parameters are optional (can be nil) and will be used in Phase 4.
|
||||
func NewObservationStore(store *Store, cleanupFunc CleanupFunc, conflictStore, relationStore interface{}) *ObservationStore {
|
||||
return &ObservationStore{
|
||||
func NewObservationStore(store *Store, cleanupFunc CleanupFunc, conflictStore, relationStore any) *ObservationStore {
|
||||
s := &ObservationStore{
|
||||
db: store.DB,
|
||||
rawDB: store.GetRawDB(),
|
||||
cleanupFunc: cleanupFunc,
|
||||
conflictStore: conflictStore,
|
||||
relationStore: relationStore,
|
||||
cleanupQueue: make(chan string, cleanupQueueSize),
|
||||
stopCleanup: make(chan struct{}),
|
||||
}
|
||||
// Start the cleanup worker
|
||||
s.startCleanupWorker()
|
||||
return s
|
||||
}
|
||||
|
||||
// startCleanupWorker starts the background cleanup worker.
|
||||
func (s *ObservationStore) startCleanupWorker() {
|
||||
s.cleanupOnce.Do(func() {
|
||||
s.cleanupStarted.Store(true)
|
||||
s.cleanupWg.Add(1)
|
||||
go s.cleanupWorker()
|
||||
})
|
||||
}
|
||||
|
||||
// cleanupWorker processes cleanup requests from the queue.
|
||||
func (s *ObservationStore) cleanupWorker() {
|
||||
defer s.cleanupWg.Done()
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-s.stopCleanup:
|
||||
// Drain remaining items in queue before exiting
|
||||
for {
|
||||
select {
|
||||
case project := <-s.cleanupQueue:
|
||||
s.processCleanup(project)
|
||||
default:
|
||||
return
|
||||
}
|
||||
}
|
||||
case project := <-s.cleanupQueue:
|
||||
s.processCleanup(project)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// processCleanup performs the actual cleanup for a project.
|
||||
func (s *ObservationStore) processCleanup(project string) {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
defer cancel()
|
||||
|
||||
deletedIDs, err := s.CleanupOldObservations(ctx, project)
|
||||
if err != nil {
|
||||
log.Warn().Err(err).Str("project", project).Msg("Failed to cleanup old observations")
|
||||
return
|
||||
}
|
||||
|
||||
if len(deletedIDs) > 0 && s.cleanupFunc != nil {
|
||||
s.cleanupFunc(ctx, deletedIDs)
|
||||
log.Debug().Str("project", project).Int("count", len(deletedIDs)).Msg("Cleaned up old observations")
|
||||
}
|
||||
}
|
||||
|
||||
// Close stops the cleanup worker and waits for it to finish.
|
||||
// Safe to call even if the worker was never started.
|
||||
func (s *ObservationStore) Close() {
|
||||
// Only stop if worker was actually started to avoid deadlock
|
||||
if s.cleanupStarted.Load() {
|
||||
close(s.stopCleanup)
|
||||
s.cleanupWg.Wait()
|
||||
}
|
||||
}
|
||||
|
||||
@@ -86,16 +173,15 @@ func (s *ObservationStore) StoreObservation(ctx context.Context, sdkSessionID, p
|
||||
return 0, 0, err
|
||||
}
|
||||
|
||||
// Cleanup old observations beyond the limit for this project (async to not block handler)
|
||||
// Queue cleanup of old observations beyond the limit for this project (async to not block handler)
|
||||
if project != "" {
|
||||
go func(proj string) {
|
||||
cleanupCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
defer cancel()
|
||||
deletedIDs, _ := s.CleanupOldObservations(cleanupCtx, proj)
|
||||
if len(deletedIDs) > 0 && s.cleanupFunc != nil {
|
||||
s.cleanupFunc(cleanupCtx, deletedIDs)
|
||||
}
|
||||
}(project)
|
||||
select {
|
||||
case s.cleanupQueue <- project:
|
||||
// Successfully queued for cleanup
|
||||
default:
|
||||
// Queue is full, log a warning instead of silently dropping
|
||||
log.Warn().Str("project", project).Msg("Cleanup queue full, skipping cleanup for this observation")
|
||||
}
|
||||
}
|
||||
|
||||
// Note: Conflict and relation detection intentionally omitted for now
|
||||
@@ -104,6 +190,89 @@ func (s *ObservationStore) StoreObservation(ctx context.Context, sdkSessionID, p
|
||||
return dbObs.ID, nowEpoch, nil
|
||||
}
|
||||
|
||||
// ObservationUpdate contains fields that can be updated on an observation.
|
||||
// Only non-nil fields will be updated.
|
||||
type ObservationUpdate struct {
|
||||
Title *string // New title
|
||||
Subtitle *string // New subtitle
|
||||
Narrative *string // New narrative
|
||||
Facts *[]string // New facts (replaces existing)
|
||||
Concepts *[]string // New concepts (replaces existing)
|
||||
FilesRead *[]string // New files read (replaces existing)
|
||||
FilesModified *[]string // New files modified (replaces existing)
|
||||
Scope *string // New scope (project or global)
|
||||
}
|
||||
|
||||
// UpdateObservation updates an existing observation with the provided fields.
|
||||
// Only non-nil fields in the update struct will be modified.
|
||||
// Returns the updated observation or an error.
|
||||
func (s *ObservationStore) UpdateObservation(ctx context.Context, id int64, update *ObservationUpdate) (*models.Observation, error) {
|
||||
if update == nil {
|
||||
return nil, fmt.Errorf("update cannot be nil")
|
||||
}
|
||||
|
||||
// First, verify the observation exists
|
||||
var dbObs Observation
|
||||
if err := s.db.WithContext(ctx).First(&dbObs, id).Error; err != nil {
|
||||
if err == gorm.ErrRecordNotFound {
|
||||
return nil, fmt.Errorf("observation not found: %d", id)
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Build update map with only provided fields
|
||||
updates := make(map[string]any)
|
||||
|
||||
if update.Title != nil {
|
||||
updates["title"] = sql.NullString{String: *update.Title, Valid: true}
|
||||
}
|
||||
if update.Subtitle != nil {
|
||||
updates["subtitle"] = sql.NullString{String: *update.Subtitle, Valid: true}
|
||||
}
|
||||
if update.Narrative != nil {
|
||||
updates["narrative"] = sql.NullString{String: *update.Narrative, Valid: true}
|
||||
}
|
||||
if update.Facts != nil {
|
||||
factsJSON, _ := json.Marshal(*update.Facts)
|
||||
updates["facts"] = string(factsJSON)
|
||||
}
|
||||
if update.Concepts != nil {
|
||||
conceptsJSON, _ := json.Marshal(*update.Concepts)
|
||||
updates["concepts"] = string(conceptsJSON)
|
||||
}
|
||||
if update.FilesRead != nil {
|
||||
filesReadJSON, _ := json.Marshal(*update.FilesRead)
|
||||
updates["files_read"] = string(filesReadJSON)
|
||||
}
|
||||
if update.FilesModified != nil {
|
||||
filesModifiedJSON, _ := json.Marshal(*update.FilesModified)
|
||||
updates["files_modified"] = string(filesModifiedJSON)
|
||||
}
|
||||
if update.Scope != nil {
|
||||
updates["scope"] = sql.NullString{String: *update.Scope, Valid: true}
|
||||
}
|
||||
|
||||
// Add updated_at timestamp
|
||||
updates["updated_at_epoch"] = sql.NullInt64{Int64: time.Now().Unix(), Valid: true}
|
||||
|
||||
if len(updates) == 0 {
|
||||
// Nothing to update, just return existing observation
|
||||
return toModelObservation(&dbObs), nil
|
||||
}
|
||||
|
||||
// Perform the update
|
||||
if err := s.db.WithContext(ctx).Model(&Observation{}).Where("id = ?", id).Updates(updates).Error; err != nil {
|
||||
return nil, fmt.Errorf("failed to update observation: %w", err)
|
||||
}
|
||||
|
||||
// Fetch the updated observation
|
||||
if err := s.db.WithContext(ctx).First(&dbObs, id).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return toModelObservation(&dbObs), nil
|
||||
}
|
||||
|
||||
// GetObservationByID retrieves an observation by its ID.
|
||||
func (s *ObservationStore) GetObservationByID(ctx context.Context, id int64) (*models.Observation, error) {
|
||||
var dbObs Observation
|
||||
@@ -134,6 +303,8 @@ func (s *ObservationStore) GetObservationsByIDs(ctx context.Context, ids []int64
|
||||
query = query.Order("created_at_epoch DESC")
|
||||
case "importance":
|
||||
query = query.Order("importance_score DESC, created_at_epoch DESC")
|
||||
case "score_desc":
|
||||
query = query.Order("importance_score DESC, created_at_epoch DESC")
|
||||
default:
|
||||
// Default: importance first, then recency
|
||||
query = query.Order("COALESCE(importance_score, 1.0) DESC, created_at_epoch DESC")
|
||||
@@ -152,6 +323,60 @@ func (s *ObservationStore) GetObservationsByIDs(ctx context.Context, ids []int64
|
||||
return toModelObservations(dbObservations), nil
|
||||
}
|
||||
|
||||
// GetObservationsByIDsPreserveOrder retrieves observations by IDs, preserving the input order.
|
||||
// This is useful when the caller has already sorted/ranked the IDs (e.g., by vector similarity).
|
||||
func (s *ObservationStore) GetObservationsByIDsPreserveOrder(ctx context.Context, ids []int64) ([]*models.Observation, error) {
|
||||
if len(ids) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
// Fetch all observations in a single query
|
||||
var dbObservations []Observation
|
||||
err := s.db.WithContext(ctx).Where("id IN ?", ids).Find(&dbObservations).Error
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Build ID -> observation map for O(1) lookups
|
||||
obsMap := make(map[int64]*Observation, len(dbObservations))
|
||||
for i := range dbObservations {
|
||||
obsMap[int64(dbObservations[i].ID)] = &dbObservations[i]
|
||||
}
|
||||
|
||||
// Reconstruct in original order
|
||||
result := make([]*models.Observation, 0, len(ids))
|
||||
for _, id := range ids {
|
||||
if obs, ok := obsMap[id]; ok {
|
||||
result = append(result, toModelObservation(obs))
|
||||
}
|
||||
}
|
||||
|
||||
return result, nil
|
||||
}
|
||||
|
||||
// BatchGetObservationsWithScores retrieves observations with associated scores.
|
||||
// Returns a map of ID -> observation for efficient lookup.
|
||||
func (s *ObservationStore) BatchGetObservationsWithScores(ctx context.Context, ids []int64) (map[int64]*models.Observation, error) {
|
||||
if len(ids) == 0 {
|
||||
return make(map[int64]*models.Observation), nil
|
||||
}
|
||||
|
||||
// Fetch all observations in a single query
|
||||
var dbObservations []Observation
|
||||
err := s.db.WithContext(ctx).Where("id IN ?", ids).Find(&dbObservations).Error
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Build result map
|
||||
result := make(map[int64]*models.Observation, len(dbObservations))
|
||||
for i := range dbObservations {
|
||||
result[int64(dbObservations[i].ID)] = toModelObservation(&dbObservations[i])
|
||||
}
|
||||
|
||||
return result, nil
|
||||
}
|
||||
|
||||
// GetRecentObservations retrieves recent observations for a project.
|
||||
// This includes project-scoped observations for the specified project AND global observations.
|
||||
// Results are ordered by importance_score DESC, then created_at_epoch DESC.
|
||||
@@ -169,13 +394,13 @@ func (s *ObservationStore) GetRecentObservations(ctx context.Context, project st
|
||||
return toModelObservations(dbObservations), nil
|
||||
}
|
||||
|
||||
// GetActiveObservations retrieves recent non-superseded observations for a project.
|
||||
// This excludes observations that have been marked as superseded by newer ones.
|
||||
// GetActiveObservations retrieves recent non-superseded, non-archived observations for a project.
|
||||
// This excludes observations that have been marked as superseded or archived.
|
||||
// Results are ordered by importance_score DESC, then created_at_epoch DESC.
|
||||
func (s *ObservationStore) GetActiveObservations(ctx context.Context, project string, limit int) ([]*models.Observation, error) {
|
||||
var dbObservations []Observation
|
||||
err := s.db.WithContext(ctx).
|
||||
Scopes(projectScopeFilter(project), notSupersededFilter(), importanceOrdering()).
|
||||
Scopes(projectScopeFilter(project), activeObservationFilter(), importanceOrdering()).
|
||||
Limit(limit).
|
||||
Find(&dbObservations).Error
|
||||
|
||||
@@ -245,7 +470,57 @@ func (s *ObservationStore) GetAllRecentObservations(ctx context.Context, limit i
|
||||
return toModelObservations(dbObservations), nil
|
||||
}
|
||||
|
||||
// GetAllRecentObservationsPaginated retrieves recent observations with pagination.
|
||||
func (s *ObservationStore) GetAllRecentObservationsPaginated(ctx context.Context, limit, offset int) ([]*models.Observation, int64, error) {
|
||||
var dbObservations []Observation
|
||||
var total int64
|
||||
|
||||
// Get total count
|
||||
if err := s.db.WithContext(ctx).Model(&Observation{}).Count(&total).Error; err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
|
||||
// Get paginated results
|
||||
err := s.db.WithContext(ctx).
|
||||
Scopes(importanceOrdering()).
|
||||
Limit(limit).
|
||||
Offset(offset).
|
||||
Find(&dbObservations).Error
|
||||
|
||||
if err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
|
||||
return toModelObservations(dbObservations), total, nil
|
||||
}
|
||||
|
||||
// GetObservationsByProjectStrictPaginated retrieves observations strictly from a project with pagination.
|
||||
func (s *ObservationStore) GetObservationsByProjectStrictPaginated(ctx context.Context, project string, limit, offset int) ([]*models.Observation, int64, error) {
|
||||
var dbObservations []Observation
|
||||
var total int64
|
||||
|
||||
// Get total count for project
|
||||
if err := s.db.WithContext(ctx).Model(&Observation{}).Where("project = ?", project).Count(&total).Error; err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
|
||||
// Get paginated results
|
||||
err := s.db.WithContext(ctx).
|
||||
Where("project = ?", project).
|
||||
Scopes(importanceOrdering()).
|
||||
Limit(limit).
|
||||
Offset(offset).
|
||||
Find(&dbObservations).Error
|
||||
|
||||
if err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
|
||||
return toModelObservations(dbObservations), total, nil
|
||||
}
|
||||
|
||||
// GetAllObservations retrieves all observations (for vector rebuild).
|
||||
// Note: For large datasets, prefer GetAllObservationsIterator to avoid memory issues.
|
||||
func (s *ObservationStore) GetAllObservations(ctx context.Context) ([]*models.Observation, error) {
|
||||
var dbObservations []Observation
|
||||
err := s.db.WithContext(ctx).
|
||||
@@ -259,6 +534,51 @@ func (s *ObservationStore) GetAllObservations(ctx context.Context) ([]*models.Ob
|
||||
return toModelObservations(dbObservations), nil
|
||||
}
|
||||
|
||||
// GetAllObservationsIterator returns observations in batches to avoid loading all into memory.
|
||||
// The callback is called for each batch. Return false from callback to stop iteration.
|
||||
// batchSize controls how many observations are loaded at once (default 500 if <= 0).
|
||||
func (s *ObservationStore) GetAllObservationsIterator(ctx context.Context, batchSize int, callback func([]*models.Observation) bool) error {
|
||||
if batchSize <= 0 {
|
||||
batchSize = 500
|
||||
}
|
||||
|
||||
var lastID int64 = 0
|
||||
for {
|
||||
// Check context cancellation
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
default:
|
||||
}
|
||||
|
||||
var dbObservations []Observation
|
||||
err := s.db.WithContext(ctx).
|
||||
Where("id > ?", lastID).
|
||||
Order("id ASC").
|
||||
Limit(batchSize).
|
||||
Find(&dbObservations).Error
|
||||
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if len(dbObservations) == 0 {
|
||||
break // No more observations
|
||||
}
|
||||
|
||||
// Update cursor for next batch
|
||||
lastID = dbObservations[len(dbObservations)-1].ID
|
||||
|
||||
// Convert and call callback
|
||||
observations := toModelObservations(dbObservations)
|
||||
if !callback(observations) {
|
||||
break // Callback requested stop
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// SearchObservationsFTS performs full-text search on observations using FTS5.
|
||||
// Falls back to LIKE search if FTS5 fails.
|
||||
func (s *ObservationStore) SearchObservationsFTS(ctx context.Context, query, project string, limit int) ([]*models.Observation, error) {
|
||||
@@ -314,14 +634,26 @@ func (s *ObservationStore) SearchObservationsFTS(ctx context.Context, query, pro
|
||||
}
|
||||
|
||||
// searchObservationsLike performs fallback LIKE search on observations using GORM.
|
||||
// Limits to 2 keywords to prevent expensive OR queries that SQLite optimizes poorly.
|
||||
// This is a fallback path when FTS returns no results, so we prioritize performance.
|
||||
func (s *ObservationStore) searchObservationsLike(ctx context.Context, keywords []string, project string, limit int) ([]*models.Observation, error) {
|
||||
if len(keywords) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
// Limit keywords to prevent excessive OR conditions that hurt query planning.
|
||||
// SQLite performs significantly better with fewer LIKE conditions.
|
||||
// Using 2 instead of 5 reduces query complexity from O(15) to O(6) conditions
|
||||
// (each keyword creates 3 LIKE conditions for title, subtitle, narrative).
|
||||
const maxKeywords = 2
|
||||
if len(keywords) > maxKeywords {
|
||||
keywords = keywords[:maxKeywords]
|
||||
}
|
||||
|
||||
// Build LIKE conditions for each keyword
|
||||
var conditions []string
|
||||
var args []interface{}
|
||||
// Pre-allocate for efficiency: maxKeywords conditions × 3 args each + 1 project arg
|
||||
conditions := make([]string, 0, len(keywords))
|
||||
args := make([]any, 0, len(keywords)*3+1)
|
||||
|
||||
for _, kw := range keywords {
|
||||
pattern := "%" + kw + "%"
|
||||
@@ -358,6 +690,217 @@ func (s *ObservationStore) DeleteObservations(ctx context.Context, ids []int64)
|
||||
return result.RowsAffected, result.Error
|
||||
}
|
||||
|
||||
// DeleteObservation deletes a single observation by ID.
|
||||
func (s *ObservationStore) DeleteObservation(ctx context.Context, id int64) error {
|
||||
result := s.db.WithContext(ctx).Delete(&Observation{}, id)
|
||||
if result.Error != nil {
|
||||
return result.Error
|
||||
}
|
||||
if result.RowsAffected == 0 {
|
||||
return fmt.Errorf("observation %d not found", id)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// MarkAsSuperseded marks an observation as superseded (stale).
|
||||
func (s *ObservationStore) MarkAsSuperseded(ctx context.Context, id int64) error {
|
||||
result := s.db.WithContext(ctx).
|
||||
Model(&Observation{}).
|
||||
Where("id = ?", id).
|
||||
Update("is_superseded", 1)
|
||||
if result.Error != nil {
|
||||
return result.Error
|
||||
}
|
||||
if result.RowsAffected == 0 {
|
||||
return fmt.Errorf("observation %d not found", id)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// MarkAsSupersededBatch marks multiple observations as superseded in a single query.
|
||||
// Returns the number of observations updated and any error.
|
||||
func (s *ObservationStore) MarkAsSupersededBatch(ctx context.Context, ids []int64) (int64, error) {
|
||||
if len(ids) == 0 {
|
||||
return 0, nil
|
||||
}
|
||||
|
||||
result := s.db.WithContext(ctx).
|
||||
Model(&Observation{}).
|
||||
Where("id IN ?", ids).
|
||||
Update("is_superseded", 1)
|
||||
|
||||
return result.RowsAffected, result.Error
|
||||
}
|
||||
|
||||
// ArchiveObservation archives a single observation with an optional reason.
|
||||
func (s *ObservationStore) ArchiveObservation(ctx context.Context, id int64, reason string) error {
|
||||
updates := map[string]any{
|
||||
"is_archived": 1,
|
||||
"archived_at_epoch": time.Now().UnixMilli(),
|
||||
}
|
||||
if reason != "" {
|
||||
updates["archived_reason"] = reason
|
||||
}
|
||||
|
||||
result := s.db.WithContext(ctx).
|
||||
Model(&Observation{}).
|
||||
Where("id = ?", id).
|
||||
Updates(updates)
|
||||
if result.Error != nil {
|
||||
return result.Error
|
||||
}
|
||||
if result.RowsAffected == 0 {
|
||||
return fmt.Errorf("observation %d not found", id)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// UnarchiveObservation restores an archived observation.
|
||||
func (s *ObservationStore) UnarchiveObservation(ctx context.Context, id int64) error {
|
||||
result := s.db.WithContext(ctx).
|
||||
Model(&Observation{}).
|
||||
Where("id = ?", id).
|
||||
Updates(map[string]any{
|
||||
"is_archived": 0,
|
||||
"archived_at_epoch": nil,
|
||||
"archived_reason": nil,
|
||||
})
|
||||
if result.Error != nil {
|
||||
return result.Error
|
||||
}
|
||||
if result.RowsAffected == 0 {
|
||||
return fmt.Errorf("observation %d not found", id)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// ArchiveOldObservations archives observations older than the specified age.
|
||||
// Returns the count of archived observations and their IDs.
|
||||
func (s *ObservationStore) ArchiveOldObservations(ctx context.Context, project string, maxAgeDays int, reason string) ([]int64, error) {
|
||||
if maxAgeDays <= 0 {
|
||||
maxAgeDays = 90 // Default: archive observations older than 90 days
|
||||
}
|
||||
|
||||
cutoffEpoch := time.Now().AddDate(0, 0, -maxAgeDays).UnixMilli()
|
||||
if reason == "" {
|
||||
reason = fmt.Sprintf("auto-archived: older than %d days", maxAgeDays)
|
||||
}
|
||||
|
||||
var idsToArchive []int64
|
||||
|
||||
err := s.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
// Find observations to archive (not already archived, older than cutoff)
|
||||
query := tx.Model(&Observation{}).
|
||||
Where("created_at_epoch < ?", cutoffEpoch).
|
||||
Where("COALESCE(is_archived, 0) = 0")
|
||||
|
||||
if project != "" {
|
||||
query = query.Where("project = ?", project)
|
||||
}
|
||||
|
||||
if err := query.Pluck("id", &idsToArchive).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if len(idsToArchive) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
// Archive the observations
|
||||
now := time.Now().UnixMilli()
|
||||
return tx.Model(&Observation{}).
|
||||
Where("id IN ?", idsToArchive).
|
||||
Updates(map[string]any{
|
||||
"is_archived": 1,
|
||||
"archived_at_epoch": now,
|
||||
"archived_reason": reason,
|
||||
}).Error
|
||||
})
|
||||
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return idsToArchive, nil
|
||||
}
|
||||
|
||||
// GetArchivedObservations retrieves archived observations for a project.
|
||||
func (s *ObservationStore) GetArchivedObservations(ctx context.Context, project string, limit int) ([]*models.Observation, error) {
|
||||
var dbObservations []Observation
|
||||
query := s.db.WithContext(ctx).
|
||||
Where("COALESCE(is_archived, 0) = 1")
|
||||
|
||||
if project != "" {
|
||||
query = query.Where("project = ?", project)
|
||||
}
|
||||
|
||||
err := query.
|
||||
Order("archived_at_epoch DESC").
|
||||
Limit(limit).
|
||||
Find(&dbObservations).Error
|
||||
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return toModelObservations(dbObservations), nil
|
||||
}
|
||||
|
||||
// GetArchivalStats returns statistics about archived observations.
|
||||
// Optimized to use a single query instead of 4 separate queries.
|
||||
func (s *ObservationStore) GetArchivalStats(ctx context.Context, project string) (*ArchivalStats, error) {
|
||||
// Use a single query with conditional aggregation to get all stats at once.
|
||||
// This is much faster than 4 separate queries (saves 3 round trips).
|
||||
type statsResult struct {
|
||||
OldestEpoch *int64
|
||||
NewestEpoch *int64
|
||||
TotalCount int64
|
||||
ArchivedCount int64
|
||||
}
|
||||
|
||||
var result statsResult
|
||||
|
||||
query := s.db.WithContext(ctx).Model(&Observation{}).
|
||||
Select(`
|
||||
COUNT(*) as total_count,
|
||||
SUM(CASE WHEN COALESCE(is_archived, 0) = 1 THEN 1 ELSE 0 END) as archived_count,
|
||||
MIN(CASE WHEN COALESCE(is_archived, 0) = 1 THEN archived_at_epoch END) as oldest_epoch,
|
||||
MAX(CASE WHEN COALESCE(is_archived, 0) = 1 THEN archived_at_epoch END) as newest_epoch
|
||||
`)
|
||||
|
||||
if project != "" {
|
||||
query = query.Where("project = ?", project)
|
||||
}
|
||||
|
||||
if err := query.Scan(&result).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
stats := &ArchivalStats{
|
||||
TotalCount: result.TotalCount,
|
||||
ArchivedCount: result.ArchivedCount,
|
||||
ActiveCount: result.TotalCount - result.ArchivedCount,
|
||||
}
|
||||
|
||||
if result.OldestEpoch != nil {
|
||||
stats.OldestArchivedEpoch = *result.OldestEpoch
|
||||
}
|
||||
if result.NewestEpoch != nil {
|
||||
stats.NewestArchivedEpoch = *result.NewestEpoch
|
||||
}
|
||||
|
||||
return stats, nil
|
||||
}
|
||||
|
||||
// ArchivalStats contains statistics about archived observations.
|
||||
type ArchivalStats struct {
|
||||
TotalCount int64 `json:"total_count"`
|
||||
ActiveCount int64 `json:"active_count"`
|
||||
ArchivedCount int64 `json:"archived_count"`
|
||||
OldestArchivedEpoch int64 `json:"oldest_archived_epoch,omitempty"`
|
||||
NewestArchivedEpoch int64 `json:"newest_archived_epoch,omitempty"`
|
||||
}
|
||||
|
||||
// CleanupOldObservations removes observations beyond the limit for a project.
|
||||
// Returns the IDs of deleted observations.
|
||||
func (s *ObservationStore) CleanupOldObservations(ctx context.Context, project string) ([]int64, error) {
|
||||
@@ -418,10 +961,12 @@ func projectScopeFilter(project string) func(*gorm.DB) *gorm.DB {
|
||||
}
|
||||
}
|
||||
|
||||
// notSupersededFilter filters out superseded observations.
|
||||
func notSupersededFilter() func(*gorm.DB) *gorm.DB {
|
||||
// activeObservationFilter filters for active (non-archived, non-superseded) observations.
|
||||
// This is more efficient than chaining notSupersededFilter + notArchivedFilter
|
||||
// as it produces a single WHERE clause for the query optimizer.
|
||||
func activeObservationFilter() func(*gorm.DB) *gorm.DB {
|
||||
return func(db *gorm.DB) *gorm.DB {
|
||||
return db.Where("COALESCE(is_superseded, 0) = 0")
|
||||
return db.Where("COALESCE(is_archived, 0) = 0 AND COALESCE(is_superseded, 0) = 0")
|
||||
}
|
||||
}
|
||||
|
||||
@@ -437,23 +982,17 @@ func importanceOrdering() func(*gorm.DB) *gorm.DB {
|
||||
// ====================
|
||||
|
||||
// extractKeywords extracts keywords from a search query.
|
||||
// Uses package-level commonWords map for O(1) stop word filtering.
|
||||
func extractKeywords(query string) []string {
|
||||
words := strings.Fields(strings.ToLower(query))
|
||||
var keywords []string
|
||||
|
||||
commonWords := map[string]bool{
|
||||
"the": true, "and": true, "or": true, "but": true, "in": true,
|
||||
"on": true, "at": true, "to": true, "for": true, "of": true,
|
||||
"with": true, "by": true, "from": true, "as": true, "is": true,
|
||||
"was": true, "are": true, "were": true, "be": true, "been": true,
|
||||
"being": true, "have": true, "has": true, "had": true, "do": true,
|
||||
"does": true, "did": true, "will": true, "would": true, "should": true,
|
||||
"could": true, "may": true, "might": true, "must": true, "can": true,
|
||||
}
|
||||
keywords := make([]string, 0, len(words)) // Pre-allocate for typical case
|
||||
|
||||
for _, word := range words {
|
||||
// Skip short words and common words
|
||||
if len(word) <= 3 || commonWords[word] {
|
||||
// Skip short words and common stop words
|
||||
if len(word) <= 3 {
|
||||
continue
|
||||
}
|
||||
if _, isCommon := commonWords[word]; isCommon {
|
||||
continue
|
||||
}
|
||||
keywords = append(keywords, word)
|
||||
@@ -464,7 +1003,8 @@ func extractKeywords(query string) []string {
|
||||
|
||||
// scanObservationRows scans multiple observations from raw SQL rows.
|
||||
func scanObservationRows(rows *sql.Rows) ([]*models.Observation, error) {
|
||||
var observations []*models.Observation
|
||||
// Pre-allocate with reasonable initial capacity to avoid repeated slice growth
|
||||
observations := make([]*models.Observation, 0, 64)
|
||||
for rows.Next() {
|
||||
obs, err := scanObservation(rows)
|
||||
if err != nil {
|
||||
@@ -476,7 +1016,7 @@ func scanObservationRows(rows *sql.Rows) ([]*models.Observation, error) {
|
||||
}
|
||||
|
||||
// scanObservation scans a single observation from a row scanner.
|
||||
func scanObservation(scanner interface{ Scan(...interface{}) error }) (*models.Observation, error) {
|
||||
func scanObservation(scanner interface{ Scan(...any) error }) (*models.Observation, error) {
|
||||
var obs models.Observation
|
||||
var factsJSON, conceptsJSON, filesReadJSON, filesModifiedJSON, fileMtimesJSON []byte
|
||||
var isSuperseded int
|
||||
|
||||
@@ -4,6 +4,8 @@ package gorm
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"fmt"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"gorm.io/gorm"
|
||||
@@ -18,6 +20,7 @@ type PatternCleanupFunc func(ctx context.Context, deletedIDs []int64)
|
||||
type PatternStore struct {
|
||||
db *gorm.DB
|
||||
cleanupFunc PatternCleanupFunc
|
||||
cleanupMu sync.RWMutex
|
||||
}
|
||||
|
||||
// NewPatternStore creates a new pattern store.
|
||||
@@ -29,6 +32,8 @@ func NewPatternStore(store *Store) *PatternStore {
|
||||
|
||||
// SetCleanupFunc sets the callback for when patterns are deleted.
|
||||
func (s *PatternStore) SetCleanupFunc(fn PatternCleanupFunc) {
|
||||
s.cleanupMu.Lock()
|
||||
defer s.cleanupMu.Unlock()
|
||||
s.cleanupFunc = fn
|
||||
}
|
||||
|
||||
@@ -238,6 +243,9 @@ func (s *PatternStore) MarkPatternDeprecated(ctx context.Context, id int64) erro
|
||||
|
||||
// MergePatterns merges a source pattern into a target pattern.
|
||||
func (s *PatternStore) MergePatterns(ctx context.Context, sourceID, targetID int64) error {
|
||||
if sourceID == targetID {
|
||||
return fmt.Errorf("cannot merge pattern into itself")
|
||||
}
|
||||
// Get both patterns
|
||||
source, err := s.GetPatternByID(ctx, sourceID)
|
||||
if err != nil {
|
||||
@@ -294,8 +302,13 @@ func (s *PatternStore) MergePatterns(ctx context.Context, sourceID, targetID int
|
||||
func (s *PatternStore) DeletePattern(ctx context.Context, id int64) error {
|
||||
result := s.db.WithContext(ctx).Delete(&Pattern{}, id)
|
||||
|
||||
if result.Error == nil && s.cleanupFunc != nil {
|
||||
s.cleanupFunc(ctx, []int64{id})
|
||||
if result.Error == nil {
|
||||
s.cleanupMu.RLock()
|
||||
fn := s.cleanupFunc
|
||||
s.cleanupMu.RUnlock()
|
||||
if fn != nil {
|
||||
fn(ctx, []int64{id})
|
||||
}
|
||||
}
|
||||
|
||||
return result.Error
|
||||
|
||||
@@ -4,8 +4,10 @@ package gorm
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/rs/zerolog/log"
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/clause"
|
||||
|
||||
@@ -23,6 +25,7 @@ const MaxPromptsGlobal = 500
|
||||
type PromptStore struct {
|
||||
db *gorm.DB
|
||||
cleanupFunc PromptCleanupFunc
|
||||
cleanupMu sync.RWMutex
|
||||
}
|
||||
|
||||
// NewPromptStore creates a new prompt store.
|
||||
@@ -35,6 +38,8 @@ func NewPromptStore(store *Store, cleanupFunc PromptCleanupFunc) *PromptStore {
|
||||
|
||||
// SetCleanupFunc sets the callback for when prompts are deleted during cleanup.
|
||||
func (s *PromptStore) SetCleanupFunc(fn PromptCleanupFunc) {
|
||||
s.cleanupMu.Lock()
|
||||
defer s.cleanupMu.Unlock()
|
||||
s.cleanupFunc = fn
|
||||
}
|
||||
|
||||
@@ -81,9 +86,15 @@ func (s *PromptStore) SaveUserPromptWithMatches(ctx context.Context, claudeSessi
|
||||
go func() {
|
||||
cleanupCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
defer cancel()
|
||||
deletedIDs, _ := s.CleanupOldPrompts(cleanupCtx)
|
||||
if len(deletedIDs) > 0 && s.cleanupFunc != nil {
|
||||
s.cleanupFunc(cleanupCtx, deletedIDs)
|
||||
if deletedIDs, err := s.CleanupOldPrompts(cleanupCtx); err != nil {
|
||||
log.Warn().Err(err).Msg("Background prompt cleanup failed")
|
||||
} else if len(deletedIDs) > 0 {
|
||||
s.cleanupMu.RLock()
|
||||
fn := s.cleanupFunc
|
||||
s.cleanupMu.RUnlock()
|
||||
if fn != nil {
|
||||
fn(cleanupCtx, deletedIDs)
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
@@ -145,9 +156,9 @@ func (s *PromptStore) GetPromptsByIDs(ctx context.Context, ids []int64, orderBy
|
||||
}
|
||||
|
||||
var results []struct {
|
||||
UserPrompt
|
||||
Project sql.NullString `gorm:"column:project"`
|
||||
SDKSessionID sql.NullString `gorm:"column:sdk_session_id"`
|
||||
UserPrompt
|
||||
}
|
||||
|
||||
query := s.db.WithContext(ctx).
|
||||
@@ -184,9 +195,9 @@ func (s *PromptStore) GetPromptsByIDs(ctx context.Context, ids []int64, orderBy
|
||||
// GetAllRecentUserPrompts retrieves recent user prompts across all projects.
|
||||
func (s *PromptStore) GetAllRecentUserPrompts(ctx context.Context, limit int) ([]*models.UserPromptWithSession, error) {
|
||||
var results []struct {
|
||||
UserPrompt
|
||||
Project sql.NullString `gorm:"column:project"`
|
||||
SDKSessionID sql.NullString `gorm:"column:sdk_session_id"`
|
||||
UserPrompt
|
||||
}
|
||||
|
||||
query := s.db.WithContext(ctx).
|
||||
@@ -211,9 +222,9 @@ func (s *PromptStore) GetAllRecentUserPrompts(ctx context.Context, limit int) ([
|
||||
// GetAllPrompts retrieves all user prompts (for vector rebuild).
|
||||
func (s *PromptStore) GetAllPrompts(ctx context.Context) ([]*models.UserPromptWithSession, error) {
|
||||
var results []struct {
|
||||
UserPrompt
|
||||
Project sql.NullString `gorm:"column:project"`
|
||||
SDKSessionID sql.NullString `gorm:"column:sdk_session_id"`
|
||||
UserPrompt
|
||||
}
|
||||
|
||||
query := s.db.WithContext(ctx).
|
||||
@@ -253,12 +264,34 @@ func (s *PromptStore) FindRecentPromptByText(ctx context.Context, claudeSessionI
|
||||
return prompt.ID, prompt.PromptNumber, true
|
||||
}
|
||||
|
||||
// FindRecentPromptByTextGlobal finds a recent prompt by exact text match within a time window
|
||||
// across ALL sessions (not scoped to a single claude session ID).
|
||||
// This catches duplicates when the same prompt arrives from different session IDs,
|
||||
// e.g. when the SDK processor's callClaudeCLI subprocess fires with a different session ID.
|
||||
// Returns (promptID, promptNumber, found).
|
||||
func (s *PromptStore) FindRecentPromptByTextGlobal(ctx context.Context, promptText string, withinSeconds int) (int64, int, bool) {
|
||||
cutoffEpoch := time.Now().Add(-time.Duration(withinSeconds) * time.Second).UnixMilli()
|
||||
|
||||
var prompt UserPrompt
|
||||
err := s.db.WithContext(ctx).
|
||||
Where("prompt_text = ? AND created_at_epoch >= ?",
|
||||
promptText, cutoffEpoch).
|
||||
Order("created_at_epoch DESC").
|
||||
First(&prompt).Error
|
||||
|
||||
if err != nil {
|
||||
return 0, 0, false
|
||||
}
|
||||
|
||||
return prompt.ID, prompt.PromptNumber, true
|
||||
}
|
||||
|
||||
// GetRecentUserPromptsByProject retrieves recent user prompts for a specific project.
|
||||
func (s *PromptStore) GetRecentUserPromptsByProject(ctx context.Context, project string, limit int) ([]*models.UserPromptWithSession, error) {
|
||||
var results []struct {
|
||||
UserPrompt
|
||||
Project sql.NullString `gorm:"column:project"`
|
||||
SDKSessionID sql.NullString `gorm:"column:sdk_session_id"`
|
||||
UserPrompt
|
||||
}
|
||||
|
||||
query := s.db.WithContext(ctx).
|
||||
@@ -283,9 +316,9 @@ func (s *PromptStore) GetRecentUserPromptsByProject(ctx context.Context, project
|
||||
|
||||
// toModelUserPromptsWithSession converts query results to pkg/models.UserPromptWithSession.
|
||||
func toModelUserPromptsWithSession(results []struct {
|
||||
UserPrompt
|
||||
Project sql.NullString `gorm:"column:project"`
|
||||
SDKSessionID sql.NullString `gorm:"column:sdk_session_id"`
|
||||
UserPrompt
|
||||
}) []*models.UserPromptWithSession {
|
||||
prompts := make([]*models.UserPromptWithSession, len(results))
|
||||
for i, r := range results {
|
||||
|
||||
@@ -333,6 +333,39 @@ func TestPromptStore_FindRecentPromptByText(t *testing.T) {
|
||||
assert.False(t, notFound, "Should not find prompt outside time window")
|
||||
}
|
||||
|
||||
func TestPromptStore_FindRecentPromptByTextGlobal(t *testing.T) {
|
||||
promptStore, _, cleanup := testPromptStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Save a prompt under session "claude-1"
|
||||
id, err := promptStore.SaveUserPromptWithMatches(ctx, "claude-1", 1, "What is the architecture?", 3)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Global search should find it without specifying session ID
|
||||
foundID, foundNumber, found := promptStore.FindRecentPromptByTextGlobal(ctx, "What is the architecture?", 60)
|
||||
assert.True(t, found, "Should find the prompt globally")
|
||||
assert.Equal(t, id, foundID)
|
||||
assert.Equal(t, 1, foundNumber)
|
||||
|
||||
// Global search should find it even when a different session ID would have been used
|
||||
// (This is the core cross-session dedup scenario)
|
||||
foundID2, foundNumber2, found2 := promptStore.FindRecentPromptByTextGlobal(ctx, "What is the architecture?", 60)
|
||||
assert.True(t, found2, "Should find the prompt from any session")
|
||||
assert.Equal(t, id, foundID2)
|
||||
assert.Equal(t, 1, foundNumber2)
|
||||
|
||||
// Different text should not match
|
||||
_, _, notFound := promptStore.FindRecentPromptByTextGlobal(ctx, "Different text", 60)
|
||||
assert.False(t, notFound, "Should not find a different prompt")
|
||||
|
||||
// Should not find prompt outside time window
|
||||
time.Sleep(100 * time.Millisecond)
|
||||
_, _, notFound = promptStore.FindRecentPromptByTextGlobal(ctx, "What is the architecture?", 0)
|
||||
assert.False(t, notFound, "Should not find prompt outside time window")
|
||||
}
|
||||
|
||||
func TestPromptStore_GetRecentUserPromptsByProject(t *testing.T) {
|
||||
promptStore, store, cleanup := testPromptStore(t)
|
||||
defer cleanup()
|
||||
|
||||
@@ -171,11 +171,11 @@ func (s *RelationStore) GetRelationsByType(ctx context.Context, relationType mod
|
||||
// GetRelationsWithDetails retrieves relations with observation titles for display.
|
||||
func (s *RelationStore) GetRelationsWithDetails(ctx context.Context, obsID int64) ([]*models.RelationWithDetails, error) {
|
||||
var results []struct {
|
||||
ObservationRelation
|
||||
SourceTitle sql.NullString `gorm:"column:source_title"`
|
||||
TargetTitle sql.NullString `gorm:"column:target_title"`
|
||||
SourceType string `gorm:"column:source_type"`
|
||||
TargetType string `gorm:"column:target_type"`
|
||||
SourceTitle sql.NullString `gorm:"column:source_title"`
|
||||
TargetTitle sql.NullString `gorm:"column:target_title"`
|
||||
ObservationRelation
|
||||
}
|
||||
|
||||
err := s.db.WithContext(ctx).
|
||||
|
||||
@@ -3,6 +3,8 @@ package gorm
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"gorm.io/gorm"
|
||||
@@ -59,12 +61,23 @@ func (s *ObservationStore) UpdateImportanceScore(ctx context.Context, id int64,
|
||||
}
|
||||
|
||||
// UpdateImportanceScores bulk updates importance scores for multiple observations.
|
||||
// This is more efficient than individual updates for batch recalculation.
|
||||
// Uses a single SQL statement with CASE/WHEN for efficient batch updates.
|
||||
func (s *ObservationStore) UpdateImportanceScores(ctx context.Context, scores map[int64]float64) error {
|
||||
if len(scores) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
// For small batches, use simple individual updates
|
||||
if len(scores) <= 5 {
|
||||
return s.updateScoresIndividually(ctx, scores)
|
||||
}
|
||||
|
||||
// For larger batches, use CASE/WHEN SQL for single-query update
|
||||
return s.updateScoresBatch(ctx, scores)
|
||||
}
|
||||
|
||||
// updateScoresIndividually updates scores one at a time (efficient for small batches).
|
||||
func (s *ObservationStore) updateScoresIndividually(ctx context.Context, scores map[int64]float64) error {
|
||||
return s.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
now := time.Now().UnixMilli()
|
||||
|
||||
@@ -85,6 +98,36 @@ func (s *ObservationStore) UpdateImportanceScores(ctx context.Context, scores ma
|
||||
})
|
||||
}
|
||||
|
||||
// updateScoresBatch updates multiple scores in a single SQL statement using CASE/WHEN.
|
||||
// This is much more efficient for large batches (O(1) queries instead of O(n)).
|
||||
func (s *ObservationStore) updateScoresBatch(ctx context.Context, scores map[int64]float64) error {
|
||||
now := time.Now().UnixMilli()
|
||||
|
||||
// Build CASE/WHEN clause for importance_score
|
||||
// UPDATE observations SET
|
||||
// importance_score = CASE id WHEN 1 THEN 0.5 WHEN 2 THEN 0.8 ... END,
|
||||
// score_updated_at_epoch = ?
|
||||
// WHERE id IN (1, 2, ...)
|
||||
|
||||
ids := make([]int64, 0, len(scores))
|
||||
caseBuilder := strings.Builder{}
|
||||
caseBuilder.WriteString("CASE id ")
|
||||
|
||||
for id, score := range scores {
|
||||
ids = append(ids, id)
|
||||
caseBuilder.WriteString(fmt.Sprintf("WHEN %d THEN %f ", id, score))
|
||||
}
|
||||
caseBuilder.WriteString("END")
|
||||
|
||||
// Use raw SQL for the batch update
|
||||
sql := fmt.Sprintf(
|
||||
"UPDATE observations SET importance_score = %s, score_updated_at_epoch = ? WHERE id IN ?",
|
||||
caseBuilder.String(),
|
||||
)
|
||||
|
||||
return s.db.WithContext(ctx).Exec(sql, now, ids).Error
|
||||
}
|
||||
|
||||
// GetObservationsNeedingScoreUpdate returns observations that need their importance score recalculated.
|
||||
// Returns observations where score_updated_at_epoch is NULL or older than the threshold.
|
||||
func (s *ObservationStore) GetObservationsNeedingScoreUpdate(ctx context.Context, threshold time.Duration, limit int) ([]*models.Observation, error) {
|
||||
|
||||
@@ -4,6 +4,7 @@ package gorm
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"gorm.io/gorm"
|
||||
@@ -67,10 +68,12 @@ func (s *SessionStore) CreateSDKSession(ctx context.Context, claudeSessionID, pr
|
||||
if userPrompt != "" {
|
||||
updates["user_prompt"] = userPrompt
|
||||
}
|
||||
s.db.WithContext(ctx).
|
||||
if err := s.db.WithContext(ctx).
|
||||
Model(&SDKSession{}).
|
||||
Where("claude_session_id = ?", claudeSessionID).
|
||||
Updates(updates)
|
||||
Updates(updates).Error; err != nil {
|
||||
return 0, fmt.Errorf("failed to update session: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
// Fetch existing session
|
||||
@@ -116,26 +119,44 @@ func (s *SessionStore) FindAnySDKSession(ctx context.Context, claudeSessionID st
|
||||
}
|
||||
|
||||
// IncrementPromptCounter increments the prompt counter and returns the new value.
|
||||
// Uses a single SQL query with RETURNING clause for optimal performance.
|
||||
func (s *SessionStore) IncrementPromptCounter(ctx context.Context, id int64) (int, error) {
|
||||
// Atomic increment using GORM expression
|
||||
err := s.db.WithContext(ctx).
|
||||
Model(&SDKSession{}).
|
||||
Where("id = ?", id).
|
||||
Update("prompt_counter", gorm.Expr("COALESCE(prompt_counter, 0) + 1")).Error
|
||||
// Use raw SQL with RETURNING to get updated value in single query
|
||||
// SQLite supports RETURNING since version 3.35.0 (2021-03-12)
|
||||
var newCounter int
|
||||
err := s.db.WithContext(ctx).Raw(`
|
||||
UPDATE sdk_sessions
|
||||
SET prompt_counter = COALESCE(prompt_counter, 0) + 1
|
||||
WHERE id = ?
|
||||
RETURNING prompt_counter
|
||||
`, id).Scan(&newCounter).Error
|
||||
|
||||
if err != nil {
|
||||
// Fallback for older SQLite versions without RETURNING support
|
||||
if err.Error() == "near \"RETURNING\": syntax error" || newCounter == 0 {
|
||||
// Atomic increment
|
||||
updateErr := s.db.WithContext(ctx).
|
||||
Model(&SDKSession{}).
|
||||
Where("id = ?", id).
|
||||
Update("prompt_counter", gorm.Expr("COALESCE(prompt_counter, 0) + 1")).Error
|
||||
if updateErr != nil {
|
||||
return 0, updateErr
|
||||
}
|
||||
|
||||
// Fetch updated value
|
||||
var sess SDKSession
|
||||
fetchErr := s.db.WithContext(ctx).
|
||||
Select("prompt_counter").
|
||||
First(&sess, id).Error
|
||||
if fetchErr != nil {
|
||||
return 0, fetchErr
|
||||
}
|
||||
return sess.PromptCounter, nil
|
||||
}
|
||||
return 0, err
|
||||
}
|
||||
|
||||
// Fetch updated value
|
||||
var sess SDKSession
|
||||
err = s.db.WithContext(ctx).
|
||||
Select("prompt_counter").
|
||||
First(&sess, id).Error
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
|
||||
return sess.PromptCounter, nil
|
||||
return newCounter, nil
|
||||
}
|
||||
|
||||
// GetPromptCounter returns the current prompt counter for a session.
|
||||
|
||||
+431
-9
@@ -2,11 +2,16 @@
|
||||
package gorm
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"fmt"
|
||||
"slices"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
sqlite_vec "github.com/asg017/sqlite-vec-go-bindings/cgo"
|
||||
_ "github.com/mattn/go-sqlite3" // Import SQLite driver with FTS5 support
|
||||
"github.com/rs/zerolog/log"
|
||||
"gorm.io/driver/sqlite"
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/logger"
|
||||
@@ -14,8 +19,13 @@ import (
|
||||
|
||||
// Store represents the GORM database connection with sqlite-vec support.
|
||||
type Store struct {
|
||||
DB *gorm.DB
|
||||
sqlDB *sql.DB // For FTS5 and sqlite-vec operations that require raw SQL
|
||||
healthCacheTime time.Time
|
||||
DB *gorm.DB
|
||||
sqlDB *sql.DB
|
||||
metrics *PoolMetrics
|
||||
cachedHealth *HealthInfo
|
||||
healthCacheTTL time.Duration
|
||||
healthCacheMu sync.RWMutex
|
||||
}
|
||||
|
||||
// Config holds database configuration.
|
||||
@@ -71,8 +81,10 @@ func NewStore(cfg Config) (*Store, error) {
|
||||
}
|
||||
|
||||
store := &Store{
|
||||
DB: db,
|
||||
sqlDB: sqlDB,
|
||||
DB: db,
|
||||
sqlDB: sqlDB,
|
||||
metrics: NewPoolMetrics(100), // Track last 100 latency samples
|
||||
healthCacheTTL: 5 * time.Second, // Cache health checks for 5 seconds
|
||||
}
|
||||
|
||||
// 7. Run migrations FIRST (before PRAGMA commands)
|
||||
@@ -80,18 +92,61 @@ func NewStore(cfg Config) (*Store, error) {
|
||||
return nil, fmt.Errorf("run migrations: %w", err)
|
||||
}
|
||||
|
||||
// 8. CRITICAL: Set WAL mode and synchronous mode via raw SQL
|
||||
// 8. CRITICAL: Set WAL mode and other performance pragmas
|
||||
// Use raw sqlDB to avoid GORM transaction issues
|
||||
if _, err := sqlDB.Exec("PRAGMA journal_mode=WAL"); err != nil {
|
||||
return nil, fmt.Errorf("set WAL mode: %w", err)
|
||||
pragmas := []string{
|
||||
"PRAGMA journal_mode=WAL",
|
||||
"PRAGMA synchronous=NORMAL",
|
||||
"PRAGMA cache_size=-64000", // 64MB cache (negative = KB)
|
||||
"PRAGMA temp_store=MEMORY", // Store temp tables in memory
|
||||
"PRAGMA mmap_size=268435456", // 256MB memory-mapped I/O
|
||||
"PRAGMA page_size=4096", // 4KB pages (optimal for most systems)
|
||||
}
|
||||
if _, err := sqlDB.Exec("PRAGMA synchronous=NORMAL"); err != nil {
|
||||
return nil, fmt.Errorf("set synchronous mode: %w", err)
|
||||
for _, pragma := range pragmas {
|
||||
if _, err := sqlDB.Exec(pragma); err != nil {
|
||||
log.Warn().Str("pragma", pragma).Err(err).Msg("Failed to set pragma (non-fatal)")
|
||||
}
|
||||
}
|
||||
// Set busy timeout to 5 seconds to handle concurrent writes
|
||||
// This allows SQLite to retry when database is locked instead of failing immediately
|
||||
if _, err := sqlDB.Exec("PRAGMA busy_timeout=5000"); err != nil {
|
||||
return nil, fmt.Errorf("set busy timeout: %w", err)
|
||||
}
|
||||
|
||||
// 9. Warm the connection pool
|
||||
store.WarmPool(maxConns)
|
||||
|
||||
return store, nil
|
||||
}
|
||||
|
||||
// WarmPool pre-creates connections to avoid cold start latency.
|
||||
func (s *Store) WarmPool(numConns int) {
|
||||
if numConns <= 0 {
|
||||
numConns = 4
|
||||
}
|
||||
|
||||
var wg sync.WaitGroup
|
||||
for i := 0; i < numConns; i++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
|
||||
defer cancel()
|
||||
|
||||
conn, err := s.sqlDB.Conn(ctx)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
// Execute a simple query to ensure the connection is fully initialized
|
||||
_ = conn.PingContext(ctx)
|
||||
// Return connection to pool (don't close it)
|
||||
_ = conn.Close()
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
log.Debug().Int("connections", numConns).Msg("Connection pool warmed")
|
||||
}
|
||||
|
||||
// Close closes the database connection.
|
||||
func (s *Store) Close() error {
|
||||
return s.sqlDB.Close()
|
||||
@@ -115,3 +170,370 @@ func (s *Store) GetRawDB() *sql.DB {
|
||||
func (s *Store) GetDB() *gorm.DB {
|
||||
return s.DB
|
||||
}
|
||||
|
||||
// Stats returns database connection pool statistics.
|
||||
func (s *Store) Stats() sql.DBStats {
|
||||
return s.sqlDB.Stats()
|
||||
}
|
||||
|
||||
// Optimize runs VACUUM and ANALYZE to optimize the database.
|
||||
// Should be called periodically (e.g., daily) during low activity.
|
||||
func (s *Store) Optimize(ctx context.Context) error {
|
||||
log.Info().Msg("Starting database optimization")
|
||||
start := time.Now()
|
||||
|
||||
// ANALYZE updates statistics for query optimizer
|
||||
if _, err := s.sqlDB.ExecContext(ctx, "ANALYZE"); err != nil {
|
||||
return fmt.Errorf("analyze: %w", err)
|
||||
}
|
||||
|
||||
// PRAGMA optimize runs optimization based on query statistics
|
||||
if _, err := s.sqlDB.ExecContext(ctx, "PRAGMA optimize"); err != nil {
|
||||
log.Warn().Err(err).Msg("PRAGMA optimize failed (non-fatal)")
|
||||
}
|
||||
|
||||
log.Info().Dur("duration", time.Since(start)).Msg("Database optimization complete")
|
||||
return nil
|
||||
}
|
||||
|
||||
// HealthCheck performs a comprehensive health check with latency measurement.
|
||||
// Returns detailed health information including connection pool stats and query latency.
|
||||
// Results are cached for healthCacheTTL (default 5 seconds) to reduce database load
|
||||
// from frequent monitoring calls.
|
||||
func (s *Store) HealthCheck(ctx context.Context) *HealthInfo {
|
||||
// Fast path: check cache with read lock
|
||||
s.healthCacheMu.RLock()
|
||||
if s.cachedHealth != nil && time.Since(s.healthCacheTime) < s.healthCacheTTL {
|
||||
cached := s.cachedHealth
|
||||
s.healthCacheMu.RUnlock()
|
||||
return cached
|
||||
}
|
||||
s.healthCacheMu.RUnlock()
|
||||
|
||||
// Slow path: perform actual health check
|
||||
info := s.performHealthCheck(ctx)
|
||||
|
||||
// Cache the result
|
||||
s.healthCacheMu.Lock()
|
||||
s.cachedHealth = info
|
||||
s.healthCacheTime = time.Now()
|
||||
s.healthCacheMu.Unlock()
|
||||
|
||||
return info
|
||||
}
|
||||
|
||||
// HealthCheckForce performs a health check bypassing the cache.
|
||||
// Use this when you need real-time health data (e.g., debugging, alerting).
|
||||
func (s *Store) HealthCheckForce(ctx context.Context) *HealthInfo {
|
||||
info := s.performHealthCheck(ctx)
|
||||
|
||||
// Update the cache with fresh data
|
||||
s.healthCacheMu.Lock()
|
||||
s.cachedHealth = info
|
||||
s.healthCacheTime = time.Now()
|
||||
s.healthCacheMu.Unlock()
|
||||
|
||||
return info
|
||||
}
|
||||
|
||||
// performHealthCheck does the actual health check work.
|
||||
func (s *Store) performHealthCheck(ctx context.Context) *HealthInfo {
|
||||
info := &HealthInfo{
|
||||
Status: "healthy",
|
||||
Timestamp: time.Now(),
|
||||
}
|
||||
|
||||
// Check pool stats
|
||||
stats := s.sqlDB.Stats()
|
||||
info.PoolStats = PoolStats{
|
||||
OpenConnections: stats.OpenConnections,
|
||||
InUse: stats.InUse,
|
||||
Idle: stats.Idle,
|
||||
WaitCount: stats.WaitCount,
|
||||
WaitDuration: stats.WaitDuration,
|
||||
MaxIdleClosed: stats.MaxIdleClosed,
|
||||
MaxLifetimeClosed: stats.MaxLifetimeClosed,
|
||||
}
|
||||
|
||||
// Record pool stats for metrics tracking
|
||||
if s.metrics != nil {
|
||||
s.metrics.RecordPoolStats(stats)
|
||||
}
|
||||
|
||||
// Measure query latency with a simple SELECT
|
||||
start := time.Now()
|
||||
var dummy int
|
||||
err := s.sqlDB.QueryRowContext(ctx, "SELECT 1").Scan(&dummy)
|
||||
info.QueryLatency = time.Since(start)
|
||||
|
||||
// Record latency for historical tracking
|
||||
if s.metrics != nil {
|
||||
s.metrics.RecordLatency(info.QueryLatency)
|
||||
info.HistoricalMetrics = s.metrics.GetMetricsSummary()
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
info.Status = "unhealthy"
|
||||
info.Error = err.Error()
|
||||
return info
|
||||
}
|
||||
|
||||
// Check for connection saturation (degraded if pool is heavily used)
|
||||
if stats.InUse > 0 && float64(stats.InUse)/float64(stats.OpenConnections) > 0.8 {
|
||||
info.Status = "degraded"
|
||||
info.Warning = "Connection pool heavily utilized"
|
||||
}
|
||||
|
||||
// Check for wait contention
|
||||
if stats.WaitCount > 100 && stats.WaitDuration > 100*time.Millisecond {
|
||||
info.Status = "degraded"
|
||||
info.Warning = "Connection pool contention detected"
|
||||
}
|
||||
|
||||
// Check query latency (warn if > 10ms for simple query)
|
||||
if info.QueryLatency > 10*time.Millisecond {
|
||||
if info.Status == "healthy" {
|
||||
info.Status = "degraded"
|
||||
}
|
||||
info.Warning = fmt.Sprintf("Slow query latency: %v", info.QueryLatency)
|
||||
}
|
||||
|
||||
// Check historical latency trend (degraded if P95 is high)
|
||||
if s.metrics != nil && info.HistoricalMetrics.P95Latency > 50*time.Millisecond {
|
||||
if info.Status == "healthy" {
|
||||
info.Status = "degraded"
|
||||
}
|
||||
info.Warning = fmt.Sprintf("High P95 latency: %v", info.HistoricalMetrics.P95Latency)
|
||||
}
|
||||
|
||||
return info
|
||||
}
|
||||
|
||||
// HealthInfo contains database health check results.
|
||||
type HealthInfo struct {
|
||||
Timestamp time.Time `json:"timestamp"`
|
||||
Status string `json:"status"`
|
||||
Error string `json:"error,omitempty"`
|
||||
Warning string `json:"warning,omitempty"`
|
||||
HistoricalMetrics MetricsSummary `json:"historical_metrics,omitempty"`
|
||||
PoolStats PoolStats `json:"pool_stats"`
|
||||
QueryLatency time.Duration `json:"query_latency_ns"`
|
||||
}
|
||||
|
||||
// PoolStats contains connection pool statistics.
|
||||
type PoolStats struct {
|
||||
OpenConnections int `json:"open_connections"`
|
||||
InUse int `json:"in_use"`
|
||||
Idle int `json:"idle"`
|
||||
WaitCount int64 `json:"wait_count"`
|
||||
WaitDuration time.Duration `json:"wait_duration_ns"`
|
||||
MaxIdleClosed int64 `json:"max_idle_closed"`
|
||||
MaxLifetimeClosed int64 `json:"max_lifetime_closed"`
|
||||
}
|
||||
|
||||
// QueryTimeout constants for different query types.
|
||||
const (
|
||||
// DefaultQueryTimeout is the default timeout for regular queries.
|
||||
DefaultQueryTimeout = 5 * time.Second
|
||||
// FastQueryTimeout is for queries that should be very fast (health checks, etc).
|
||||
FastQueryTimeout = 1 * time.Second
|
||||
// SlowQueryTimeout is for queries that may take longer (bulk operations, rebuilds).
|
||||
SlowQueryTimeout = 30 * time.Second
|
||||
)
|
||||
|
||||
// PoolMetrics tracks historical connection pool metrics with a sliding window.
|
||||
type PoolMetrics struct {
|
||||
lastSampleTime time.Time
|
||||
latencySamples []time.Duration
|
||||
latencyIdx int
|
||||
latencyCount int
|
||||
totalQueries int64
|
||||
totalWaitTime time.Duration
|
||||
peakInUse int
|
||||
peakWaitCount int64
|
||||
windowSize int
|
||||
mu sync.RWMutex
|
||||
}
|
||||
|
||||
// NewPoolMetrics creates a new pool metrics collector with the given window size.
|
||||
func NewPoolMetrics(windowSize int) *PoolMetrics {
|
||||
if windowSize <= 0 {
|
||||
windowSize = 100 // Default: track last 100 samples
|
||||
}
|
||||
return &PoolMetrics{
|
||||
latencySamples: make([]time.Duration, windowSize),
|
||||
windowSize: windowSize,
|
||||
lastSampleTime: time.Now(),
|
||||
}
|
||||
}
|
||||
|
||||
// RecordLatency records a query latency sample.
|
||||
func (m *PoolMetrics) RecordLatency(latency time.Duration) {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
|
||||
m.latencySamples[m.latencyIdx] = latency
|
||||
m.latencyIdx = (m.latencyIdx + 1) % m.windowSize
|
||||
if m.latencyCount < m.windowSize {
|
||||
m.latencyCount++
|
||||
}
|
||||
m.totalQueries++
|
||||
m.lastSampleTime = time.Now()
|
||||
}
|
||||
|
||||
// RecordPoolStats records pool statistics for peak tracking.
|
||||
func (m *PoolMetrics) RecordPoolStats(stats sql.DBStats) {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
|
||||
if stats.InUse > m.peakInUse {
|
||||
m.peakInUse = stats.InUse
|
||||
}
|
||||
if stats.WaitCount > m.peakWaitCount {
|
||||
m.peakWaitCount = stats.WaitCount
|
||||
}
|
||||
m.totalWaitTime += stats.WaitDuration
|
||||
}
|
||||
|
||||
// GetMetricsSummary returns a summary of collected metrics.
|
||||
func (m *PoolMetrics) GetMetricsSummary() MetricsSummary {
|
||||
m.mu.RLock()
|
||||
defer m.mu.RUnlock()
|
||||
|
||||
summary := MetricsSummary{
|
||||
TotalQueries: m.totalQueries,
|
||||
SampleCount: m.latencyCount,
|
||||
PeakInUse: m.peakInUse,
|
||||
PeakWaitCount: m.peakWaitCount,
|
||||
TotalWaitTime: m.totalWaitTime,
|
||||
LastSampleTime: m.lastSampleTime,
|
||||
}
|
||||
|
||||
if m.latencyCount == 0 {
|
||||
return summary
|
||||
}
|
||||
|
||||
// Calculate latency statistics
|
||||
var total time.Duration
|
||||
var min, max time.Duration = m.latencySamples[0], m.latencySamples[0]
|
||||
|
||||
for i := 0; i < m.latencyCount; i++ {
|
||||
sample := m.latencySamples[i]
|
||||
total += sample
|
||||
if sample < min {
|
||||
min = sample
|
||||
}
|
||||
if sample > max {
|
||||
max = sample
|
||||
}
|
||||
}
|
||||
|
||||
summary.AvgLatency = total / time.Duration(m.latencyCount)
|
||||
summary.MinLatency = min
|
||||
summary.MaxLatency = max
|
||||
|
||||
// Calculate P95 latency (approximate using sorted samples)
|
||||
if m.latencyCount >= 20 {
|
||||
// Copy samples for sorting
|
||||
samples := make([]time.Duration, m.latencyCount)
|
||||
copy(samples, m.latencySamples[:m.latencyCount])
|
||||
// Use slices.Sort for O(n log n) instead of O(n²) insertion sort
|
||||
slices.Sort(samples)
|
||||
p95Idx := int(float64(len(samples)) * 0.95)
|
||||
summary.P95Latency = samples[p95Idx]
|
||||
}
|
||||
|
||||
return summary
|
||||
}
|
||||
|
||||
// MetricsSummary contains aggregated pool metrics.
|
||||
type MetricsSummary struct {
|
||||
LastSampleTime time.Time `json:"last_sample_time"`
|
||||
TotalQueries int64 `json:"total_queries"`
|
||||
SampleCount int `json:"sample_count"`
|
||||
AvgLatency time.Duration `json:"avg_latency_ns"`
|
||||
MinLatency time.Duration `json:"min_latency_ns"`
|
||||
MaxLatency time.Duration `json:"max_latency_ns"`
|
||||
P95Latency time.Duration `json:"p95_latency_ns,omitempty"`
|
||||
PeakInUse int `json:"peak_in_use"`
|
||||
PeakWaitCount int64 `json:"peak_wait_count"`
|
||||
TotalWaitTime time.Duration `json:"total_wait_time_ns"`
|
||||
}
|
||||
|
||||
// GetMetrics returns the current metrics without performing a health check.
|
||||
func (s *Store) GetMetrics() MetricsSummary {
|
||||
if s.metrics == nil {
|
||||
return MetricsSummary{}
|
||||
}
|
||||
return s.metrics.GetMetricsSummary()
|
||||
}
|
||||
|
||||
// ResetMetrics resets the metrics collector (useful for testing or after major changes).
|
||||
func (s *Store) ResetMetrics() {
|
||||
if s.metrics != nil {
|
||||
s.metrics = NewPoolMetrics(s.metrics.windowSize)
|
||||
}
|
||||
}
|
||||
|
||||
// WithTimeout wraps a context with the given timeout and logs slow queries.
|
||||
// Returns the wrapped context and a cancel function that should be called when done.
|
||||
func (s *Store) WithTimeout(ctx context.Context, timeout time.Duration, operation string) (context.Context, context.CancelFunc) {
|
||||
timeoutCtx, cancel := context.WithTimeout(ctx, timeout)
|
||||
start := time.Now()
|
||||
|
||||
// Return wrapped cancel that logs if query was slow
|
||||
return timeoutCtx, func() {
|
||||
elapsed := time.Since(start)
|
||||
cancel()
|
||||
|
||||
// Log slow queries (> 100ms)
|
||||
if elapsed > 100*time.Millisecond {
|
||||
log.Warn().
|
||||
Str("operation", operation).
|
||||
Dur("elapsed", elapsed).
|
||||
Dur("timeout", timeout).
|
||||
Msg("Slow database operation")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ExecWithTimeout executes a raw SQL query with timeout.
|
||||
// Returns error if query takes longer than timeout.
|
||||
func (s *Store) ExecWithTimeout(ctx context.Context, timeout time.Duration, query string, args ...any) error {
|
||||
timeoutCtx, cancel := s.WithTimeout(ctx, timeout, "exec")
|
||||
defer cancel()
|
||||
|
||||
_, err := s.sqlDB.ExecContext(timeoutCtx, query, args...)
|
||||
if err != nil {
|
||||
if err == context.DeadlineExceeded {
|
||||
return fmt.Errorf("query timeout after %v: %s", timeout, query)
|
||||
}
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// QueryRowWithTimeout executes a row query with timeout.
|
||||
func (s *Store) QueryRowWithTimeout(ctx context.Context, timeout time.Duration, query string, args ...any) *sql.Row {
|
||||
timeoutCtx, cancel := s.WithTimeout(ctx, timeout, "query_row")
|
||||
// Note: cancel will be called when row.Scan() completes or errors
|
||||
_ = cancel // Caller must ensure proper cleanup
|
||||
return s.sqlDB.QueryRowContext(timeoutCtx, query, args...)
|
||||
}
|
||||
|
||||
// TransactionWithTimeout wraps a transaction function with timeout handling.
|
||||
// The transaction is automatically rolled back if the context times out.
|
||||
func (s *Store) TransactionWithTimeout(ctx context.Context, timeout time.Duration, fn func(*gorm.DB) error) error {
|
||||
timeoutCtx, cancel := s.WithTimeout(ctx, timeout, "transaction")
|
||||
defer cancel()
|
||||
|
||||
return s.DB.WithContext(timeoutCtx).Transaction(func(tx *gorm.DB) error {
|
||||
// Check context before proceeding
|
||||
select {
|
||||
case <-timeoutCtx.Done():
|
||||
return timeoutCtx.Err()
|
||||
default:
|
||||
}
|
||||
return fn(tx)
|
||||
})
|
||||
}
|
||||
|
||||
@@ -57,6 +57,7 @@ type PromptReader interface {
|
||||
GetAllPrompts(ctx context.Context) ([]*models.UserPromptWithSession, error)
|
||||
GetRecentUserPromptsByProject(ctx context.Context, project string, limit int) ([]*models.UserPromptWithSession, error)
|
||||
FindRecentPromptByText(ctx context.Context, claudeSessionID, promptText string, withinSeconds int) (int64, int, bool)
|
||||
FindRecentPromptByTextGlobal(ctx context.Context, promptText string, withinSeconds int) (int64, int, bool)
|
||||
}
|
||||
|
||||
// PromptWriter defines write operations for prompts.
|
||||
|
||||
@@ -1 +1 @@
|
||||
1.23.2
|
||||
1.24.3
|
||||
|
||||
@@ -1 +1 @@
|
||||
1.23.2
|
||||
1.24.3
|
||||
|
||||
@@ -1 +1 @@
|
||||
1.23.2
|
||||
1.24.3
|
||||
|
||||
@@ -1 +1 @@
|
||||
1.23.2
|
||||
1.24.3
|
||||
|
||||
@@ -21,15 +21,10 @@ const (
|
||||
// ONNXConfig describes ONNX-specific model configuration.
|
||||
// This allows different models to specify their tensor names and pooling needs.
|
||||
type ONNXConfig struct {
|
||||
// InputNames are the ONNX input tensor names in order.
|
||||
InputNames []string
|
||||
// OutputNames are the ONNX output tensor names.
|
||||
Pooling PoolingStrategy
|
||||
InputNames []string
|
||||
OutputNames []string
|
||||
// Pooling specifies how to convert token embeddings to sentence embeddings.
|
||||
// If PoolingNone, the model outputs sentence embeddings directly.
|
||||
Pooling PoolingStrategy
|
||||
// HiddenSize is the embedding dimension (used for pooling calculations).
|
||||
HiddenSize int
|
||||
HiddenSize int
|
||||
}
|
||||
|
||||
// EmbeddingModel represents a text embedding model.
|
||||
@@ -62,11 +57,11 @@ type ONNXConfigurer interface {
|
||||
|
||||
// ModelMetadata describes an embedding model for UI/config.
|
||||
type ModelMetadata struct {
|
||||
Name string `json:"name"` // Human-readable name
|
||||
Version string `json:"version"` // Short ID for DB storage
|
||||
Dimensions int `json:"dimensions"` // Vector size
|
||||
Description string `json:"description"` // Brief description
|
||||
Default bool `json:"default"` // Is this the default model?
|
||||
Name string `json:"name"`
|
||||
Version string `json:"version"`
|
||||
Description string `json:"description"`
|
||||
Dimensions int `json:"dimensions"`
|
||||
Default bool `json:"default"`
|
||||
}
|
||||
|
||||
// ModelFactory creates a new instance of an embedding model.
|
||||
@@ -74,10 +69,10 @@ type ModelFactory func() (EmbeddingModel, error)
|
||||
|
||||
// ModelRegistry provides model lookup by version.
|
||||
type ModelRegistry struct {
|
||||
mu sync.RWMutex
|
||||
models map[string]ModelFactory
|
||||
metadata map[string]ModelMetadata
|
||||
defaultModel string
|
||||
mu sync.RWMutex
|
||||
}
|
||||
|
||||
// NewModelRegistry creates a new model registry.
|
||||
|
||||
@@ -8,6 +8,7 @@ import (
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"sync"
|
||||
|
||||
"github.com/sugarme/tokenizer"
|
||||
@@ -46,9 +47,9 @@ var bgeONNXConfig = ONNXConfig{
|
||||
type bgeModel struct {
|
||||
tk *tokenizer.Tokenizer
|
||||
session *ort.DynamicAdvancedSession
|
||||
libDir string
|
||||
config ONNXConfig
|
||||
mu sync.Mutex
|
||||
libDir string // temp directory containing extracted libraries
|
||||
config ONNXConfig // ONNX configuration for this model
|
||||
}
|
||||
|
||||
// Compile-time check that bgeModel implements EmbeddingModel
|
||||
@@ -69,8 +70,10 @@ func newBGEModel() (EmbeddingModel, error) {
|
||||
libPath := filepath.Join(libDir, onnxRuntimeLibName)
|
||||
ort.SetSharedLibraryPath(libPath)
|
||||
|
||||
// Initialize ONNX runtime
|
||||
if err := ort.InitializeEnvironment(); err != nil {
|
||||
// Initialize ONNX runtime (idempotent - ignore "already initialized" since
|
||||
// the ONNX environment is process-global and shared with the reranking service)
|
||||
err = ort.InitializeEnvironment()
|
||||
if err != nil && !strings.Contains(err.Error(), "already been initialized") {
|
||||
return nil, fmt.Errorf("initialize ONNX runtime: %w", err)
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,417 @@
|
||||
package graph
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"math"
|
||||
|
||||
"github.com/lukaszraczylo/claude-mnemonic/pkg/models"
|
||||
"github.com/rs/zerolog/log"
|
||||
)
|
||||
|
||||
const (
|
||||
// SemanticSimilarityThreshold for creating semantic edges
|
||||
SemanticSimilarityThreshold = 0.85
|
||||
|
||||
// MinFileOverlapForEdge minimum file overlap ratio to create edge
|
||||
MinFileOverlapForEdge = 0.3
|
||||
|
||||
// MaxEdgesPerNode prevents creating too many edges
|
||||
MaxEdgesPerNode = 20
|
||||
)
|
||||
|
||||
// DetectEdges identifies relationships between observations
|
||||
func DetectEdges(ctx context.Context, observations []*models.Observation) ([]Edge, error) {
|
||||
if len(observations) < 2 {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
edges := make([]Edge, 0)
|
||||
|
||||
// Build lookup maps for efficient detection
|
||||
sessionMap := buildSessionMap(observations)
|
||||
conceptMap := buildConceptMap(observations)
|
||||
fileMap := buildFileMap(observations)
|
||||
|
||||
log.Info().
|
||||
Int("observations", len(observations)).
|
||||
Int("sessions", len(sessionMap)).
|
||||
Int("concepts", len(conceptMap)).
|
||||
Msg("Starting edge detection")
|
||||
|
||||
// Detect temporal edges (same session)
|
||||
temporalEdges := detectTemporalEdges(sessionMap)
|
||||
edges = append(edges, temporalEdges...)
|
||||
|
||||
// Detect concept edges (shared tags)
|
||||
conceptEdges := detectConceptEdges(conceptMap)
|
||||
edges = append(edges, conceptEdges...)
|
||||
|
||||
// Detect file overlap edges
|
||||
fileEdges := detectFileOverlapEdges(fileMap, observations)
|
||||
edges = append(edges, fileEdges...)
|
||||
|
||||
// Prune excessive edges per node
|
||||
edges = pruneEdges(edges, MaxEdgesPerNode)
|
||||
|
||||
log.Info().
|
||||
Int("temporal_edges", len(temporalEdges)).
|
||||
Int("concept_edges", len(conceptEdges)).
|
||||
Int("file_edges", len(fileEdges)).
|
||||
Int("total_edges", len(edges)).
|
||||
Msg("Edge detection complete")
|
||||
|
||||
return edges, nil
|
||||
}
|
||||
|
||||
// buildSessionMap groups observations by SDK session
|
||||
func buildSessionMap(observations []*models.Observation) map[string][]int64 {
|
||||
sessionMap := make(map[string][]int64)
|
||||
|
||||
for _, obs := range observations {
|
||||
if obs.SDKSessionID != "" {
|
||||
sessionMap[obs.SDKSessionID] = append(sessionMap[obs.SDKSessionID], obs.ID)
|
||||
}
|
||||
}
|
||||
|
||||
return sessionMap
|
||||
}
|
||||
|
||||
// buildConceptMap groups observations by concept tags
|
||||
func buildConceptMap(observations []*models.Observation) map[string][]int64 {
|
||||
conceptMap := make(map[string][]int64)
|
||||
|
||||
for _, obs := range observations {
|
||||
for _, concept := range obs.Concepts {
|
||||
conceptMap[concept] = append(conceptMap[concept], obs.ID)
|
||||
}
|
||||
}
|
||||
|
||||
return conceptMap
|
||||
}
|
||||
|
||||
// buildFileMap maps files to observations (from both FilesRead and FilesModified)
|
||||
func buildFileMap(observations []*models.Observation) map[string][]int64 {
|
||||
fileMap := make(map[string][]int64)
|
||||
|
||||
for _, obs := range observations {
|
||||
// Add files from FilesRead
|
||||
for _, file := range obs.FilesRead {
|
||||
fileMap[file] = append(fileMap[file], obs.ID)
|
||||
}
|
||||
// Add files from FilesModified
|
||||
for _, file := range obs.FilesModified {
|
||||
fileMap[file] = append(fileMap[file], obs.ID)
|
||||
}
|
||||
}
|
||||
|
||||
return fileMap
|
||||
}
|
||||
|
||||
// detectTemporalEdges creates edges between observations in the same session
|
||||
func detectTemporalEdges(sessionMap map[string][]int64) []Edge {
|
||||
edges := make([]Edge, 0)
|
||||
|
||||
for _, obsIDs := range sessionMap {
|
||||
if len(obsIDs) < 2 {
|
||||
continue
|
||||
}
|
||||
|
||||
// Create edges between consecutive observations in session
|
||||
for i := 0; i < len(obsIDs)-1; i++ {
|
||||
edges = append(edges, Edge{
|
||||
FromID: obsIDs[i],
|
||||
ToID: obsIDs[i+1],
|
||||
Relation: RelationTemporal,
|
||||
Weight: 0.8, // High weight for temporal proximity
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
return edges
|
||||
}
|
||||
|
||||
// detectConceptEdges creates edges between observations sharing concepts
|
||||
func detectConceptEdges(conceptMap map[string][]int64) []Edge {
|
||||
edges := make([]Edge, 0)
|
||||
seen := make(map[string]bool)
|
||||
|
||||
for concept, obsIDs := range conceptMap {
|
||||
if len(obsIDs) < 2 {
|
||||
continue
|
||||
}
|
||||
|
||||
// Create edges between all observations sharing this concept
|
||||
for i := 0; i < len(obsIDs); i++ {
|
||||
for j := i + 1; j < len(obsIDs); j++ {
|
||||
// Use sorted pair as key to avoid duplicates
|
||||
pairKey := edgeKey(obsIDs[i], obsIDs[j])
|
||||
if seen[pairKey] {
|
||||
continue
|
||||
}
|
||||
seen[pairKey] = true
|
||||
|
||||
// Weight based on concept specificity (longer = more specific)
|
||||
weight := float32(0.5 + 0.3*math.Min(1.0, float64(len(concept))/20.0))
|
||||
|
||||
edges = append(edges, Edge{
|
||||
FromID: obsIDs[i],
|
||||
ToID: obsIDs[j],
|
||||
Relation: RelationConcept,
|
||||
Weight: weight,
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return edges
|
||||
}
|
||||
|
||||
// detectFileOverlapEdges creates edges based on file references
|
||||
func detectFileOverlapEdges(fileMap map[string][]int64, observations []*models.Observation) []Edge {
|
||||
edges := make([]Edge, 0)
|
||||
seen := make(map[string]bool)
|
||||
|
||||
// Build observation ID to observation map for quick lookup
|
||||
obsMap := make(map[int64]*models.Observation)
|
||||
for _, obs := range observations {
|
||||
obsMap[obs.ID] = obs
|
||||
}
|
||||
|
||||
for _, obsIDs := range fileMap {
|
||||
if len(obsIDs) < 2 {
|
||||
continue
|
||||
}
|
||||
|
||||
// Create edges between observations referencing same files
|
||||
for i := 0; i < len(obsIDs); i++ {
|
||||
for j := i + 1; j < len(obsIDs); j++ {
|
||||
pairKey := edgeKey(obsIDs[i], obsIDs[j])
|
||||
if seen[pairKey] {
|
||||
continue
|
||||
}
|
||||
seen[pairKey] = true
|
||||
|
||||
// Calculate file overlap ratio
|
||||
obs1, ok1 := obsMap[obsIDs[i]]
|
||||
obs2, ok2 := obsMap[obsIDs[j]]
|
||||
|
||||
if !ok1 || !ok2 {
|
||||
continue
|
||||
}
|
||||
|
||||
// Merge FilesRead and FilesModified for both observations
|
||||
files1 := append([]string{}, obs1.FilesRead...)
|
||||
files1 = append(files1, obs1.FilesModified...)
|
||||
files2 := append([]string{}, obs2.FilesRead...)
|
||||
files2 = append(files2, obs2.FilesModified...)
|
||||
|
||||
overlap := calculateFileOverlap(files1, files2)
|
||||
if overlap < MinFileOverlapForEdge {
|
||||
continue
|
||||
}
|
||||
|
||||
edges = append(edges, Edge{
|
||||
FromID: obsIDs[i],
|
||||
ToID: obsIDs[j],
|
||||
Relation: RelationFileOverlap,
|
||||
Weight: overlap,
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return edges
|
||||
}
|
||||
|
||||
// calculateFileOverlap computes Jaccard similarity of file sets
|
||||
func calculateFileOverlap(files1, files2 []string) float32 {
|
||||
if len(files1) == 0 || len(files2) == 0 {
|
||||
return 0.0
|
||||
}
|
||||
|
||||
// Convert to sets
|
||||
set1 := make(map[string]bool)
|
||||
for _, f := range files1 {
|
||||
set1[f] = true
|
||||
}
|
||||
|
||||
set2 := make(map[string]bool)
|
||||
for _, f := range files2 {
|
||||
set2[f] = true
|
||||
}
|
||||
|
||||
// Count intersection
|
||||
intersection := 0
|
||||
for f := range set1 {
|
||||
if set2[f] {
|
||||
intersection++
|
||||
}
|
||||
}
|
||||
|
||||
// Jaccard similarity = intersection / union
|
||||
union := len(set1) + len(set2) - intersection
|
||||
if union == 0 {
|
||||
return 0.0
|
||||
}
|
||||
|
||||
return float32(intersection) / float32(union)
|
||||
}
|
||||
|
||||
// pruneEdges limits edges per node to prevent graph explosion
|
||||
func pruneEdges(edges []Edge, maxPerNode int) []Edge {
|
||||
if maxPerNode <= 0 {
|
||||
return edges
|
||||
}
|
||||
|
||||
// Count edges per node
|
||||
outEdges := make(map[int64][]Edge)
|
||||
inEdges := make(map[int64][]Edge)
|
||||
|
||||
for _, edge := range edges {
|
||||
outEdges[edge.FromID] = append(outEdges[edge.FromID], edge)
|
||||
inEdges[edge.ToID] = append(inEdges[edge.ToID], edge)
|
||||
}
|
||||
|
||||
// Prune low-weight edges if node has too many
|
||||
pruned := make([]Edge, 0, len(edges))
|
||||
processed := make(map[string]bool)
|
||||
|
||||
for _, edge := range edges {
|
||||
pairKey := edgeKey(edge.FromID, edge.ToID)
|
||||
if processed[pairKey] {
|
||||
continue
|
||||
}
|
||||
processed[pairKey] = true
|
||||
|
||||
// Check if either node has too many edges
|
||||
fromCount := len(outEdges[edge.FromID])
|
||||
toCount := len(inEdges[edge.ToID])
|
||||
|
||||
if fromCount <= maxPerNode && toCount <= maxPerNode {
|
||||
pruned = append(pruned, edge)
|
||||
continue
|
||||
}
|
||||
|
||||
// Keep edge if it's high-weight (top edges for this node)
|
||||
if shouldKeepEdge(edge, outEdges[edge.FromID], maxPerNode) {
|
||||
pruned = append(pruned, edge)
|
||||
}
|
||||
}
|
||||
|
||||
if len(pruned) < len(edges) {
|
||||
log.Debug().
|
||||
Int("original", len(edges)).
|
||||
Int("pruned", len(pruned)).
|
||||
Int("removed", len(edges)-len(pruned)).
|
||||
Msg("Pruned excessive edges")
|
||||
}
|
||||
|
||||
return pruned
|
||||
}
|
||||
|
||||
// shouldKeepEdge determines if edge should be kept during pruning
|
||||
func shouldKeepEdge(edge Edge, nodeEdges []Edge, maxPerNode int) bool {
|
||||
// Sort node's edges by weight descending
|
||||
sortedEdges := make([]Edge, len(nodeEdges))
|
||||
copy(sortedEdges, nodeEdges)
|
||||
|
||||
sortEdgesByWeight(sortedEdges)
|
||||
|
||||
// Keep edge if it's in top maxPerNode
|
||||
for i := 0; i < maxPerNode && i < len(sortedEdges); i++ {
|
||||
if sortedEdges[i].FromID == edge.FromID && sortedEdges[i].ToID == edge.ToID {
|
||||
return true
|
||||
}
|
||||
}
|
||||
|
||||
return false
|
||||
}
|
||||
|
||||
// sortEdgesByWeight sorts edges by weight descending
|
||||
func sortEdgesByWeight(edges []Edge) {
|
||||
// Simple bubble sort (edges are typically small per node)
|
||||
n := len(edges)
|
||||
for i := 0; i < n-1; i++ {
|
||||
for j := 0; j < n-i-1; j++ {
|
||||
if edges[j].Weight < edges[j+1].Weight {
|
||||
edges[j], edges[j+1] = edges[j+1], edges[j]
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// edgeKey creates a unique key for an edge pair (sorted)
|
||||
func edgeKey(id1, id2 int64) string {
|
||||
if id1 < id2 {
|
||||
return fmt.Sprintf("%d-%d", id1, id2)
|
||||
}
|
||||
return fmt.Sprintf("%d-%d", id2, id1)
|
||||
}
|
||||
|
||||
// DetectSemanticEdges creates edges based on semantic similarity
|
||||
// This requires embeddings and is called separately when available
|
||||
func DetectSemanticEdges(ctx context.Context, observations []*models.Observation, embeddings map[int64][]float32) []Edge {
|
||||
edges := make([]Edge, 0)
|
||||
seen := make(map[string]bool)
|
||||
|
||||
// Compare all pairs (expensive, but necessary for semantic similarity)
|
||||
for i := 0; i < len(observations); i++ {
|
||||
emb1, ok1 := embeddings[observations[i].ID]
|
||||
if !ok1 {
|
||||
continue
|
||||
}
|
||||
|
||||
for j := i + 1; j < len(observations); j++ {
|
||||
emb2, ok2 := embeddings[observations[j].ID]
|
||||
if !ok2 {
|
||||
continue
|
||||
}
|
||||
|
||||
similarity := cosineSimilarity(emb1, emb2)
|
||||
if similarity < SemanticSimilarityThreshold {
|
||||
continue
|
||||
}
|
||||
|
||||
pairKey := edgeKey(observations[i].ID, observations[j].ID)
|
||||
if seen[pairKey] {
|
||||
continue
|
||||
}
|
||||
seen[pairKey] = true
|
||||
|
||||
edges = append(edges, Edge{
|
||||
FromID: observations[i].ID,
|
||||
ToID: observations[j].ID,
|
||||
Relation: RelationSemantic,
|
||||
Weight: similarity,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
log.Info().
|
||||
Int("semantic_edges", len(edges)).
|
||||
Float32("threshold", SemanticSimilarityThreshold).
|
||||
Msg("Detected semantic edges")
|
||||
|
||||
return edges
|
||||
}
|
||||
|
||||
// cosineSimilarity computes cosine similarity between two vectors
|
||||
func cosineSimilarity(a, b []float32) float32 {
|
||||
if len(a) != len(b) {
|
||||
return 0.0
|
||||
}
|
||||
|
||||
var dotProduct, normA, normB float32
|
||||
for i := range a {
|
||||
dotProduct += a[i] * b[i]
|
||||
normA += a[i] * a[i]
|
||||
normB += b[i] * b[i]
|
||||
}
|
||||
|
||||
if normA == 0 || normB == 0 {
|
||||
return 0.0
|
||||
}
|
||||
|
||||
return dotProduct / float32(math.Sqrt(float64(normA))*math.Sqrt(float64(normB)))
|
||||
}
|
||||
@@ -0,0 +1,423 @@
|
||||
// Package graph provides observation relationship graphs for LEANN Phase 2.
|
||||
//
|
||||
// This package implements graph-based selective recomputation where observation
|
||||
// relationships (file overlap, semantic similarity, temporal proximity) form a
|
||||
// graph structure. Hub nodes (high-degree observations) store embeddings, while
|
||||
// leaf nodes recompute on-demand.
|
||||
package graph
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"math"
|
||||
"sort"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/lukaszraczylo/claude-mnemonic/pkg/models"
|
||||
"github.com/rs/zerolog/log"
|
||||
)
|
||||
|
||||
// RelationType defines the type of relationship between observations
|
||||
type RelationType int
|
||||
|
||||
const (
|
||||
// RelationFileOverlap indicates observations reference overlapping files
|
||||
RelationFileOverlap RelationType = iota
|
||||
// RelationSemantic indicates high semantic similarity (cosine > 0.85)
|
||||
RelationSemantic
|
||||
// RelationTemporal indicates observations from same session
|
||||
RelationTemporal
|
||||
// RelationConcept indicates shared concept tags
|
||||
RelationConcept
|
||||
)
|
||||
|
||||
// Edge represents a relationship between two observations
|
||||
type Edge struct {
|
||||
FromID int64
|
||||
ToID int64
|
||||
Relation RelationType
|
||||
Weight float32 // 0.0-1.0, higher = stronger relationship
|
||||
}
|
||||
|
||||
// Node represents an observation in the graph
|
||||
type Node struct {
|
||||
Metadata NodeMetadata
|
||||
LastAccess time.Time
|
||||
StoredEmb []float32 // Nil if recomputed on-demand
|
||||
ID int64
|
||||
Degree int // Number of edges (hub detection)
|
||||
AccessCount int
|
||||
}
|
||||
|
||||
// NodeMetadata contains observation metadata
|
||||
type NodeMetadata struct {
|
||||
CreatedAt time.Time
|
||||
Project string
|
||||
Type string
|
||||
Title string
|
||||
IsSuperseded bool
|
||||
}
|
||||
|
||||
// CSRGraph represents a graph in Compressed Sparse Row format for memory efficiency
|
||||
type CSRGraph struct {
|
||||
RowPtr []int32 // Node adjacency list pointers
|
||||
ColIdx []int32 // Edge destination IDs
|
||||
Weights []float32 // Edge weights
|
||||
mu sync.RWMutex
|
||||
}
|
||||
|
||||
// ObservationGraph manages the observation relationship graph
|
||||
type ObservationGraph struct {
|
||||
nodes map[int64]*Node
|
||||
csr *CSRGraph
|
||||
edges []Edge
|
||||
nodesMu sync.RWMutex
|
||||
edgesMu sync.RWMutex
|
||||
}
|
||||
|
||||
// NewObservationGraph creates a new empty observation graph
|
||||
func NewObservationGraph() *ObservationGraph {
|
||||
return &ObservationGraph{
|
||||
nodes: make(map[int64]*Node),
|
||||
edges: make([]Edge, 0),
|
||||
csr: &CSRGraph{},
|
||||
}
|
||||
}
|
||||
|
||||
// AddNode adds or updates a node in the graph
|
||||
func (g *ObservationGraph) AddNode(node *Node) {
|
||||
g.nodesMu.Lock()
|
||||
defer g.nodesMu.Unlock()
|
||||
|
||||
g.nodes[node.ID] = node
|
||||
}
|
||||
|
||||
// AddEdge adds an edge to the graph
|
||||
func (g *ObservationGraph) AddEdge(edge Edge) {
|
||||
g.edgesMu.Lock()
|
||||
defer g.edgesMu.Unlock()
|
||||
|
||||
g.edges = append(g.edges, edge)
|
||||
|
||||
// Update degree counts
|
||||
g.nodesMu.Lock()
|
||||
if fromNode, ok := g.nodes[edge.FromID]; ok {
|
||||
fromNode.Degree++
|
||||
}
|
||||
if toNode, ok := g.nodes[edge.ToID]; ok {
|
||||
toNode.Degree++
|
||||
}
|
||||
g.nodesMu.Unlock()
|
||||
}
|
||||
|
||||
// BuildCSR converts edge list to CSR format for efficient traversal
|
||||
func (g *ObservationGraph) BuildCSR() error {
|
||||
g.edgesMu.RLock()
|
||||
g.nodesMu.RLock()
|
||||
defer g.edgesMu.RUnlock()
|
||||
defer g.nodesMu.RUnlock()
|
||||
|
||||
if len(g.nodes) == 0 {
|
||||
return fmt.Errorf("no nodes in graph")
|
||||
}
|
||||
|
||||
// Create node ID to index mapping
|
||||
nodeIDs := make([]int64, 0, len(g.nodes))
|
||||
for id := range g.nodes {
|
||||
nodeIDs = append(nodeIDs, id)
|
||||
}
|
||||
sort.Slice(nodeIDs, func(i, j int) bool {
|
||||
return nodeIDs[i] < nodeIDs[j]
|
||||
})
|
||||
|
||||
idToIdx := make(map[int64]int32)
|
||||
for idx, id := range nodeIDs {
|
||||
// #nosec G115 - observation count will never exceed int32 max (2.1B) in practice
|
||||
idToIdx[id] = int32(idx)
|
||||
}
|
||||
|
||||
// Count edges per node
|
||||
edgeCounts := make([]int, len(nodeIDs))
|
||||
for _, edge := range g.edges {
|
||||
if fromIdx, ok := idToIdx[edge.FromID]; ok {
|
||||
edgeCounts[fromIdx]++
|
||||
}
|
||||
}
|
||||
|
||||
// Build row pointers
|
||||
rowPtr := make([]int32, len(nodeIDs)+1)
|
||||
rowPtr[0] = 0
|
||||
for i := 0; i < len(nodeIDs); i++ {
|
||||
// #nosec G115 - edge counts per node will not exceed int32 max
|
||||
rowPtr[i+1] = rowPtr[i] + int32(edgeCounts[i])
|
||||
}
|
||||
|
||||
// Build column indices and weights
|
||||
totalEdges := rowPtr[len(nodeIDs)]
|
||||
colIdx := make([]int32, totalEdges)
|
||||
weights := make([]float32, totalEdges)
|
||||
|
||||
// Temporary counter for filling CSR
|
||||
currentPos := make([]int32, len(nodeIDs))
|
||||
copy(currentPos, rowPtr[:len(nodeIDs)])
|
||||
|
||||
for _, edge := range g.edges {
|
||||
fromIdx, fromOk := idToIdx[edge.FromID]
|
||||
toIdx, toOk := idToIdx[edge.ToID]
|
||||
|
||||
if fromOk && toOk {
|
||||
pos := currentPos[fromIdx]
|
||||
colIdx[pos] = toIdx
|
||||
weights[pos] = edge.Weight
|
||||
currentPos[fromIdx]++
|
||||
}
|
||||
}
|
||||
|
||||
g.csr.mu.Lock()
|
||||
g.csr.RowPtr = rowPtr
|
||||
g.csr.ColIdx = colIdx
|
||||
g.csr.Weights = weights
|
||||
g.csr.mu.Unlock()
|
||||
|
||||
log.Info().
|
||||
Int("nodes", len(nodeIDs)).
|
||||
Int("edges", int(totalEdges)).
|
||||
Msg("Built CSR graph representation")
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// GetNeighbors returns neighboring nodes and their edge weights
|
||||
func (g *ObservationGraph) GetNeighbors(nodeID int64) ([]int64, []float32, error) {
|
||||
g.csr.mu.RLock()
|
||||
defer g.csr.mu.RUnlock()
|
||||
|
||||
// Find node index in CSR
|
||||
g.nodesMu.RLock()
|
||||
nodeIDs := make([]int64, 0, len(g.nodes))
|
||||
for id := range g.nodes {
|
||||
nodeIDs = append(nodeIDs, id)
|
||||
}
|
||||
g.nodesMu.RUnlock()
|
||||
|
||||
sort.Slice(nodeIDs, func(i, j int) bool {
|
||||
return nodeIDs[i] < nodeIDs[j]
|
||||
})
|
||||
|
||||
nodeIdx := sort.Search(len(nodeIDs), func(i int) bool {
|
||||
return nodeIDs[i] >= nodeID
|
||||
})
|
||||
|
||||
if nodeIdx >= len(nodeIDs) || nodeIDs[nodeIdx] != nodeID {
|
||||
return nil, nil, fmt.Errorf("node %d not found", nodeID)
|
||||
}
|
||||
|
||||
// Extract neighbors from CSR
|
||||
startIdx := g.csr.RowPtr[nodeIdx]
|
||||
endIdx := g.csr.RowPtr[nodeIdx+1]
|
||||
|
||||
neighborCount := endIdx - startIdx
|
||||
neighbors := make([]int64, neighborCount)
|
||||
weights := make([]float32, neighborCount)
|
||||
|
||||
for i := int32(0); i < neighborCount; i++ {
|
||||
neighborIdx := g.csr.ColIdx[startIdx+i]
|
||||
neighbors[i] = nodeIDs[neighborIdx]
|
||||
weights[i] = g.csr.Weights[startIdx+i]
|
||||
}
|
||||
|
||||
return neighbors, weights, nil
|
||||
}
|
||||
|
||||
// GetNode retrieves a node by ID
|
||||
func (g *ObservationGraph) GetNode(nodeID int64) (*Node, error) {
|
||||
g.nodesMu.RLock()
|
||||
defer g.nodesMu.RUnlock()
|
||||
|
||||
node, ok := g.nodes[nodeID]
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("node %d not found", nodeID)
|
||||
}
|
||||
|
||||
return node, nil
|
||||
}
|
||||
|
||||
// FindHubs identifies hub nodes (high degree) in the graph
|
||||
func (g *ObservationGraph) FindHubs(percentile float64) []int64 {
|
||||
g.nodesMu.RLock()
|
||||
defer g.nodesMu.RUnlock()
|
||||
|
||||
if len(g.nodes) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
// Collect all degrees
|
||||
degrees := make([]int, 0, len(g.nodes))
|
||||
nodeIDs := make([]int64, 0, len(g.nodes))
|
||||
|
||||
for id, node := range g.nodes {
|
||||
degrees = append(degrees, node.Degree)
|
||||
nodeIDs = append(nodeIDs, id)
|
||||
}
|
||||
|
||||
// Sort by degree
|
||||
type nodeDegree struct {
|
||||
ID int64
|
||||
Degree int
|
||||
}
|
||||
|
||||
nodeDegrees := make([]nodeDegree, len(nodeIDs))
|
||||
for i := range nodeIDs {
|
||||
nodeDegrees[i] = nodeDegree{
|
||||
ID: nodeIDs[i],
|
||||
Degree: degrees[i],
|
||||
}
|
||||
}
|
||||
|
||||
sort.Slice(nodeDegrees, func(i, j int) bool {
|
||||
return nodeDegrees[i].Degree > nodeDegrees[j].Degree
|
||||
})
|
||||
|
||||
// Return top percentile
|
||||
cutoff := int(math.Ceil(float64(len(nodeDegrees)) * (1.0 - percentile)))
|
||||
if cutoff > len(nodeDegrees) {
|
||||
cutoff = len(nodeDegrees)
|
||||
}
|
||||
|
||||
hubs := make([]int64, cutoff)
|
||||
for i := 0; i < cutoff; i++ {
|
||||
hubs[i] = nodeDegrees[i].ID
|
||||
}
|
||||
|
||||
log.Info().
|
||||
Int("total_nodes", len(g.nodes)).
|
||||
Int("hubs", len(hubs)).
|
||||
Float64("percentile", percentile).
|
||||
Msg("Identified hub nodes")
|
||||
|
||||
return hubs
|
||||
}
|
||||
|
||||
// Stats returns graph statistics
|
||||
func (g *ObservationGraph) Stats() GraphStats {
|
||||
g.nodesMu.RLock()
|
||||
g.edgesMu.RLock()
|
||||
defer g.nodesMu.RUnlock()
|
||||
defer g.edgesMu.RUnlock()
|
||||
|
||||
stats := GraphStats{
|
||||
NodeCount: len(g.nodes),
|
||||
EdgeCount: len(g.edges),
|
||||
}
|
||||
|
||||
if len(g.nodes) > 0 {
|
||||
degrees := make([]int, 0, len(g.nodes))
|
||||
for _, node := range g.nodes {
|
||||
degrees = append(degrees, node.Degree)
|
||||
}
|
||||
|
||||
sort.Ints(degrees)
|
||||
stats.AvgDegree = float64(sum(degrees)) / float64(len(degrees))
|
||||
stats.MaxDegree = degrees[len(degrees)-1]
|
||||
stats.MinDegree = degrees[0]
|
||||
|
||||
// Median
|
||||
mid := len(degrees) / 2
|
||||
if len(degrees)%2 == 0 {
|
||||
stats.MedianDegree = float64(degrees[mid-1]+degrees[mid]) / 2.0
|
||||
} else {
|
||||
stats.MedianDegree = float64(degrees[mid])
|
||||
}
|
||||
}
|
||||
|
||||
// Count edge types
|
||||
stats.EdgeTypes = make(map[RelationType]int)
|
||||
for _, edge := range g.edges {
|
||||
stats.EdgeTypes[edge.Relation]++
|
||||
}
|
||||
|
||||
return stats
|
||||
}
|
||||
|
||||
// GraphStats contains graph statistics
|
||||
type GraphStats struct {
|
||||
EdgeTypes map[RelationType]int
|
||||
AvgDegree float64
|
||||
MedianDegree float64
|
||||
NodeCount int
|
||||
EdgeCount int
|
||||
MaxDegree int
|
||||
MinDegree int
|
||||
}
|
||||
|
||||
// BuildFromObservations constructs a graph from a list of observations
|
||||
func BuildFromObservations(ctx context.Context, observations []*models.Observation) (*ObservationGraph, error) {
|
||||
graph := NewObservationGraph()
|
||||
|
||||
// Add nodes
|
||||
for _, obs := range observations {
|
||||
// Extract title from sql.NullString
|
||||
title := ""
|
||||
if obs.Title.Valid {
|
||||
title = obs.Title.String
|
||||
}
|
||||
|
||||
node := &Node{
|
||||
ID: obs.ID,
|
||||
Degree: 0,
|
||||
Metadata: NodeMetadata{
|
||||
Project: obs.Project,
|
||||
Type: string(obs.Type),
|
||||
Title: title,
|
||||
CreatedAt: time.UnixMilli(obs.CreatedAtEpoch),
|
||||
IsSuperseded: obs.IsSuperseded,
|
||||
},
|
||||
LastAccess: time.Now(),
|
||||
AccessCount: 0,
|
||||
}
|
||||
graph.AddNode(node)
|
||||
}
|
||||
|
||||
// Detect edges (will be implemented in edge_detector.go)
|
||||
edges, err := DetectEdges(ctx, observations)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("detect edges: %w", err)
|
||||
}
|
||||
|
||||
for _, edge := range edges {
|
||||
graph.AddEdge(edge)
|
||||
}
|
||||
|
||||
// Build CSR representation
|
||||
if err := graph.BuildCSR(); err != nil {
|
||||
return nil, fmt.Errorf("build CSR: %w", err)
|
||||
}
|
||||
|
||||
return graph, nil
|
||||
}
|
||||
|
||||
// Helper function to sum integers
|
||||
func sum(values []int) int {
|
||||
total := 0
|
||||
for _, v := range values {
|
||||
total += v
|
||||
}
|
||||
return total
|
||||
}
|
||||
|
||||
// String returns a human-readable representation of RelationType
|
||||
func (r RelationType) String() string {
|
||||
switch r {
|
||||
case RelationFileOverlap:
|
||||
return "file_overlap"
|
||||
case RelationSemantic:
|
||||
return "semantic"
|
||||
case RelationTemporal:
|
||||
return "temporal"
|
||||
case RelationConcept:
|
||||
return "concept"
|
||||
default:
|
||||
return "unknown"
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,290 @@
|
||||
// Package maintenance provides scheduled maintenance tasks for claude-mnemonic.
|
||||
package maintenance
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/lukaszraczylo/claude-mnemonic/internal/config"
|
||||
"github.com/lukaszraczylo/claude-mnemonic/internal/db/gorm"
|
||||
"github.com/rs/zerolog"
|
||||
)
|
||||
|
||||
// Service handles scheduled maintenance tasks.
|
||||
type Service struct {
|
||||
log zerolog.Logger
|
||||
lastRunTime time.Time
|
||||
promptStore *gorm.PromptStore
|
||||
store *gorm.Store
|
||||
vectorCleanupFn func(ctx context.Context, deletedIDs []int64)
|
||||
config *config.Config
|
||||
summaryStore *gorm.SummaryStore
|
||||
stopCh chan struct{}
|
||||
doneCh chan struct{}
|
||||
observationStore *gorm.ObservationStore
|
||||
lastRunDuration time.Duration
|
||||
totalCleanedObs int64
|
||||
totalOptimizeRun int64
|
||||
mu sync.Mutex
|
||||
running bool
|
||||
}
|
||||
|
||||
// NewService creates a new maintenance service.
|
||||
func NewService(
|
||||
store *gorm.Store,
|
||||
observationStore *gorm.ObservationStore,
|
||||
summaryStore *gorm.SummaryStore,
|
||||
promptStore *gorm.PromptStore,
|
||||
vectorCleanupFn func(ctx context.Context, deletedIDs []int64),
|
||||
cfg *config.Config,
|
||||
log zerolog.Logger,
|
||||
) *Service {
|
||||
return &Service{
|
||||
store: store,
|
||||
observationStore: observationStore,
|
||||
summaryStore: summaryStore,
|
||||
promptStore: promptStore,
|
||||
vectorCleanupFn: vectorCleanupFn,
|
||||
config: cfg,
|
||||
log: log.With().Str("component", "maintenance").Logger(),
|
||||
stopCh: make(chan struct{}),
|
||||
doneCh: make(chan struct{}),
|
||||
}
|
||||
}
|
||||
|
||||
// Start begins the maintenance loop.
|
||||
func (s *Service) Start(ctx context.Context) {
|
||||
s.mu.Lock()
|
||||
if s.running {
|
||||
s.mu.Unlock()
|
||||
return
|
||||
}
|
||||
s.running = true
|
||||
s.mu.Unlock()
|
||||
|
||||
defer func() {
|
||||
s.mu.Lock()
|
||||
s.running = false
|
||||
s.mu.Unlock()
|
||||
close(s.doneCh)
|
||||
}()
|
||||
|
||||
if !s.config.MaintenanceEnabled {
|
||||
s.log.Info().Msg("Maintenance disabled, not starting scheduler")
|
||||
return
|
||||
}
|
||||
|
||||
interval := max(time.Duration(s.config.MaintenanceIntervalHours)*time.Hour, time.Hour)
|
||||
|
||||
s.log.Info().
|
||||
Dur("interval", interval).
|
||||
Int("retention_days", s.config.ObservationRetentionDays).
|
||||
Bool("cleanup_stale", s.config.CleanupStaleObservations).
|
||||
Msg("Starting maintenance scheduler")
|
||||
|
||||
// Initial run after 5 minutes (allow system to stabilize)
|
||||
time.Sleep(5 * time.Minute)
|
||||
s.runMaintenance(ctx)
|
||||
|
||||
ticker := time.NewTicker(interval)
|
||||
defer ticker.Stop()
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
s.log.Info().Msg("Maintenance shutting down due to context cancellation")
|
||||
return
|
||||
case <-s.stopCh:
|
||||
s.log.Info().Msg("Maintenance shutting down due to stop signal")
|
||||
return
|
||||
case <-ticker.C:
|
||||
s.runMaintenance(ctx)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Stop signals the maintenance service to stop.
|
||||
func (s *Service) Stop() {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
if !s.running {
|
||||
return
|
||||
}
|
||||
|
||||
close(s.stopCh)
|
||||
}
|
||||
|
||||
// Wait waits for the maintenance service to finish.
|
||||
func (s *Service) Wait() {
|
||||
<-s.doneCh
|
||||
}
|
||||
|
||||
// runMaintenance executes all maintenance tasks.
|
||||
func (s *Service) runMaintenance(ctx context.Context) {
|
||||
start := time.Now()
|
||||
s.log.Info().Msg("Starting maintenance run")
|
||||
|
||||
var totalCleaned int64
|
||||
|
||||
// Task 1: Clean up old observations by age
|
||||
if s.config.ObservationRetentionDays > 0 {
|
||||
cleaned, err := s.cleanupOldObservations(ctx)
|
||||
if err != nil {
|
||||
s.log.Error().Err(err).Msg("Failed to cleanup old observations")
|
||||
} else {
|
||||
totalCleaned += cleaned
|
||||
s.log.Info().Int64("cleaned", cleaned).Msg("Cleaned old observations by age")
|
||||
}
|
||||
}
|
||||
|
||||
// Task 2: Clean up stale observations
|
||||
if s.config.CleanupStaleObservations {
|
||||
cleaned, err := s.cleanupStaleObservations(ctx)
|
||||
if err != nil {
|
||||
s.log.Error().Err(err).Msg("Failed to cleanup stale observations")
|
||||
} else {
|
||||
totalCleaned += cleaned
|
||||
s.log.Info().Int64("cleaned", cleaned).Msg("Cleaned stale observations")
|
||||
}
|
||||
}
|
||||
|
||||
// Task 3: Optimize database
|
||||
if err := s.store.Optimize(ctx); err != nil {
|
||||
s.log.Error().Err(err).Msg("Failed to optimize database")
|
||||
} else {
|
||||
s.totalOptimizeRun++
|
||||
}
|
||||
|
||||
// Task 4: Clean up old prompts (keep last 1000 per session)
|
||||
cleanedPrompts, err := s.cleanupOldPrompts(ctx)
|
||||
if err != nil {
|
||||
s.log.Error().Err(err).Msg("Failed to cleanup old prompts")
|
||||
} else if cleanedPrompts > 0 {
|
||||
s.log.Info().Int64("cleaned", cleanedPrompts).Msg("Cleaned old prompts")
|
||||
}
|
||||
|
||||
// Update metrics
|
||||
s.mu.Lock()
|
||||
s.lastRunTime = time.Now()
|
||||
s.lastRunDuration = time.Since(start)
|
||||
s.totalCleanedObs += totalCleaned
|
||||
s.mu.Unlock()
|
||||
|
||||
s.log.Info().
|
||||
Dur("duration", time.Since(start)).
|
||||
Int64("observations_cleaned", totalCleaned).
|
||||
Msg("Maintenance run completed")
|
||||
}
|
||||
|
||||
// cleanupOldObservations deletes observations older than the retention period.
|
||||
func (s *Service) cleanupOldObservations(ctx context.Context) (int64, error) {
|
||||
cutoffEpoch := time.Now().AddDate(0, 0, -s.config.ObservationRetentionDays).Unix()
|
||||
|
||||
// Get IDs of old observations
|
||||
var deletedIDs []int64
|
||||
err := s.store.GetDB().WithContext(ctx).
|
||||
Model(&gorm.Observation{}).
|
||||
Where("created_at_epoch < ?", cutoffEpoch).
|
||||
Pluck("id", &deletedIDs).Error
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
|
||||
if len(deletedIDs) == 0 {
|
||||
return 0, nil
|
||||
}
|
||||
|
||||
// Delete in batches to avoid long transactions
|
||||
batchSize := 100
|
||||
for i := 0; i < len(deletedIDs); i += batchSize {
|
||||
end := min(i+batchSize, len(deletedIDs))
|
||||
batch := deletedIDs[i:end]
|
||||
|
||||
if err := s.store.GetDB().WithContext(ctx).
|
||||
Where("id IN ?", batch).
|
||||
Delete(&gorm.Observation{}).Error; err != nil {
|
||||
return int64(i), err
|
||||
}
|
||||
|
||||
// Sync vector DB deletions
|
||||
if s.vectorCleanupFn != nil {
|
||||
s.vectorCleanupFn(ctx, batch)
|
||||
}
|
||||
}
|
||||
|
||||
return int64(len(deletedIDs)), nil
|
||||
}
|
||||
|
||||
// cleanupStaleObservations deletes observations marked as stale.
|
||||
func (s *Service) cleanupStaleObservations(ctx context.Context) (int64, error) {
|
||||
// Get IDs of stale observations (is_superseded = true)
|
||||
var deletedIDs []int64
|
||||
err := s.store.GetDB().WithContext(ctx).
|
||||
Model(&gorm.Observation{}).
|
||||
Where("is_superseded = ?", true).
|
||||
Pluck("id", &deletedIDs).Error
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
|
||||
if len(deletedIDs) == 0 {
|
||||
return 0, nil
|
||||
}
|
||||
|
||||
// Delete in batches
|
||||
batchSize := 100
|
||||
for i := 0; i < len(deletedIDs); i += batchSize {
|
||||
end := min(i+batchSize, len(deletedIDs))
|
||||
batch := deletedIDs[i:end]
|
||||
|
||||
if err := s.store.GetDB().WithContext(ctx).
|
||||
Where("id IN ?", batch).
|
||||
Delete(&gorm.Observation{}).Error; err != nil {
|
||||
return int64(i), err
|
||||
}
|
||||
|
||||
// Sync vector DB deletions
|
||||
if s.vectorCleanupFn != nil {
|
||||
s.vectorCleanupFn(ctx, batch)
|
||||
}
|
||||
}
|
||||
|
||||
return int64(len(deletedIDs)), nil
|
||||
}
|
||||
|
||||
// cleanupOldPrompts removes old prompts keeping only the most recent per session.
|
||||
func (s *Service) cleanupOldPrompts(ctx context.Context) (int64, error) {
|
||||
// Delete prompts older than 30 days that aren't the most recent in their session
|
||||
cutoffEpoch := time.Now().AddDate(0, 0, -30).Unix()
|
||||
|
||||
result := s.store.GetDB().WithContext(ctx).
|
||||
Where("created_at_epoch < ?", cutoffEpoch).
|
||||
Delete(&gorm.UserPrompt{})
|
||||
|
||||
return result.RowsAffected, result.Error
|
||||
}
|
||||
|
||||
// Stats returns maintenance statistics.
|
||||
func (s *Service) Stats() map[string]any {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
return map[string]any{
|
||||
"enabled": s.config.MaintenanceEnabled,
|
||||
"interval_hours": s.config.MaintenanceIntervalHours,
|
||||
"retention_days": s.config.ObservationRetentionDays,
|
||||
"cleanup_stale": s.config.CleanupStaleObservations,
|
||||
"last_run": s.lastRunTime,
|
||||
"last_duration_ms": s.lastRunDuration.Milliseconds(),
|
||||
"total_cleaned_obs": s.totalCleanedObs,
|
||||
"total_optimizes": s.totalOptimizeRun,
|
||||
"running": s.running,
|
||||
}
|
||||
}
|
||||
|
||||
// RunNow triggers an immediate maintenance run.
|
||||
func (s *Service) RunNow(ctx context.Context) {
|
||||
go s.runMaintenance(ctx)
|
||||
}
|
||||
+1455
-221
File diff suppressed because it is too large
Load Diff
+1970
-253
File diff suppressed because it is too large
Load Diff
@@ -3,6 +3,9 @@ package pattern
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"sort"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
@@ -21,6 +24,8 @@ type DetectorConfig struct {
|
||||
AnalysisInterval time.Duration
|
||||
// MaxPatternsToTrack is the maximum number of active patterns.
|
||||
MaxPatternsToTrack int
|
||||
// MaxCandidates is the maximum number of candidates to track (LRU eviction).
|
||||
MaxCandidates int
|
||||
}
|
||||
|
||||
// DefaultConfig returns the default detector configuration.
|
||||
@@ -30,6 +35,7 @@ func DefaultConfig() DetectorConfig {
|
||||
MinFrequencyForPattern: 2, // At least 2 occurrences to form a pattern
|
||||
AnalysisInterval: 5 * time.Minute,
|
||||
MaxPatternsToTrack: 1000,
|
||||
MaxCandidates: 500, // Prevent unbounded growth
|
||||
}
|
||||
}
|
||||
|
||||
@@ -38,21 +44,15 @@ type PatternSyncFunc func(pattern *models.Pattern)
|
||||
|
||||
// Detector detects and tracks recurring patterns across observations.
|
||||
type Detector struct {
|
||||
config DetectorConfig
|
||||
ctx context.Context
|
||||
patternStore *gorm.PatternStore
|
||||
observationStore *gorm.ObservationStore
|
||||
|
||||
// Vector sync callback
|
||||
syncFunc PatternSyncFunc
|
||||
|
||||
// Candidate tracking (patterns not yet confirmed)
|
||||
candidates map[string]*candidatePattern
|
||||
candidatesMu sync.RWMutex
|
||||
|
||||
// Background analysis
|
||||
ctx context.Context
|
||||
cancel context.CancelFunc
|
||||
wg sync.WaitGroup
|
||||
syncFunc PatternSyncFunc
|
||||
candidates map[string]*candidatePattern
|
||||
cancel context.CancelFunc
|
||||
config DetectorConfig
|
||||
wg sync.WaitGroup
|
||||
candidatesMu sync.RWMutex
|
||||
}
|
||||
|
||||
// SetSyncFunc sets the callback for syncing patterns to vector store.
|
||||
@@ -62,11 +62,11 @@ func (d *Detector) SetSyncFunc(fn PatternSyncFunc) {
|
||||
|
||||
// candidatePattern tracks a potential pattern before it reaches frequency threshold.
|
||||
type candidatePattern struct {
|
||||
patternType models.PatternType
|
||||
title string
|
||||
signature []string
|
||||
observationIDs []int64
|
||||
projects []string
|
||||
patternType models.PatternType
|
||||
title string
|
||||
lastSeenEpoch int64
|
||||
}
|
||||
|
||||
@@ -201,7 +201,24 @@ func (d *Detector) AnalyzeObservation(ctx context.Context, obs *models.Observati
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// Create new candidate
|
||||
// Create new candidate - with immediate size check to prevent unbounded growth
|
||||
// between periodic cleanups (which run every 5 minutes)
|
||||
if d.config.MaxCandidates > 0 && len(d.candidates) >= d.config.MaxCandidates {
|
||||
// Evict oldest candidate immediately rather than waiting for periodic cleanup
|
||||
var oldestKey string
|
||||
var oldestTime int64 = time.Now().UnixMilli()
|
||||
for k, c := range d.candidates {
|
||||
if c.lastSeenEpoch < oldestTime {
|
||||
oldestTime = c.lastSeenEpoch
|
||||
oldestKey = k
|
||||
}
|
||||
}
|
||||
if oldestKey != "" {
|
||||
delete(d.candidates, oldestKey)
|
||||
log.Debug().Str("evicted_key", oldestKey).Msg("Evicted oldest candidate to make room")
|
||||
}
|
||||
}
|
||||
|
||||
patternType := models.DetectPatternType(obs.Concepts, obs.Title.String, obs.Narrative.String)
|
||||
d.candidates[candidateKey] = &candidatePattern{
|
||||
signature: signature,
|
||||
@@ -305,25 +322,65 @@ func (d *Detector) AnalyzeRecentObservations(ctx context.Context) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// cleanupOldCandidates removes candidates that haven't been seen recently.
|
||||
// cleanupOldCandidates removes candidates that haven't been seen recently
|
||||
// and enforces the max candidates limit using LRU eviction.
|
||||
func (d *Detector) cleanupOldCandidates() {
|
||||
d.candidatesMu.Lock()
|
||||
defer d.candidatesMu.Unlock()
|
||||
|
||||
threshold := time.Now().Add(-7 * 24 * time.Hour).UnixMilli()
|
||||
|
||||
// First pass: remove expired candidates
|
||||
for key, candidate := range d.candidates {
|
||||
if candidate.lastSeenEpoch < threshold {
|
||||
delete(d.candidates, key)
|
||||
}
|
||||
}
|
||||
|
||||
// Second pass: enforce max candidates limit using LRU eviction
|
||||
if d.config.MaxCandidates > 0 && len(d.candidates) > d.config.MaxCandidates {
|
||||
// Find oldest candidates to evict using O(n log n) sort instead of O(n²) selection sort
|
||||
type keyAge struct {
|
||||
key string
|
||||
age int64
|
||||
}
|
||||
candidates := make([]keyAge, 0, len(d.candidates))
|
||||
for k, c := range d.candidates {
|
||||
candidates = append(candidates, keyAge{k, c.lastSeenEpoch})
|
||||
}
|
||||
|
||||
// Sort by age ascending (oldest first) - O(n log n)
|
||||
sort.Slice(candidates, func(i, j int) bool {
|
||||
return candidates[i].age < candidates[j].age
|
||||
})
|
||||
|
||||
// Delete oldest entries
|
||||
toEvict := len(d.candidates) - d.config.MaxCandidates
|
||||
for i := 0; i < toEvict; i++ {
|
||||
delete(d.candidates, candidates[i].key)
|
||||
}
|
||||
|
||||
log.Debug().Int("evicted", toEvict).Int("remaining", len(d.candidates)).Msg("Evicted old pattern candidates (LRU)")
|
||||
}
|
||||
}
|
||||
|
||||
// CandidateCount returns the current number of pattern candidates.
|
||||
func (d *Detector) CandidateCount() int {
|
||||
d.candidatesMu.RLock()
|
||||
defer d.candidatesMu.RUnlock()
|
||||
return len(d.candidates)
|
||||
}
|
||||
|
||||
// GetPatternInsight returns a formatted insight string for a pattern.
|
||||
// GetPatternInsight returns a formatted insight string for a pattern.
|
||||
func (d *Detector) GetPatternInsight(ctx context.Context, patternID int64) (string, error) {
|
||||
pattern, err := d.patternStore.GetPatternByID(ctx, patternID)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if pattern == nil {
|
||||
return "", fmt.Errorf("pattern not found: %d", patternID)
|
||||
}
|
||||
|
||||
return formatPatternInsight(pattern), nil
|
||||
}
|
||||
@@ -340,11 +397,15 @@ func generateCandidateKey(signature []string) string {
|
||||
if len(signature) == 0 {
|
||||
return ""
|
||||
}
|
||||
key := ""
|
||||
// Use strings.Builder to avoid O(n²) string concatenation
|
||||
var b strings.Builder
|
||||
// Pre-allocate: estimate average signature element is 10 chars + separator
|
||||
b.Grow(len(signature) * 11)
|
||||
for _, s := range signature {
|
||||
key += s + "|"
|
||||
b.WriteString(s)
|
||||
b.WriteByte('|')
|
||||
}
|
||||
return key
|
||||
return b.String()
|
||||
}
|
||||
|
||||
// generatePatternName creates a human-readable name for a pattern.
|
||||
|
||||
@@ -328,19 +328,39 @@ func TestDefaultConfig(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestDefaultConfig_AllFieldsValid(t *testing.T) {
|
||||
config := DefaultConfig()
|
||||
|
||||
if config.MinMatchScore != 0.3 {
|
||||
t.Errorf("MinMatchScore = %f, want 0.3", config.MinMatchScore)
|
||||
}
|
||||
if config.MinFrequencyForPattern != 2 {
|
||||
t.Errorf("MinFrequencyForPattern = %d, want 2", config.MinFrequencyForPattern)
|
||||
}
|
||||
if config.AnalysisInterval != 5*time.Minute {
|
||||
t.Errorf("AnalysisInterval = %v, want 5m", config.AnalysisInterval)
|
||||
}
|
||||
if config.MaxPatternsToTrack != 1000 {
|
||||
t.Errorf("MaxPatternsToTrack = %d, want 1000", config.MaxPatternsToTrack)
|
||||
}
|
||||
if config.MaxCandidates != 500 {
|
||||
t.Errorf("MaxCandidates = %d, want 500", config.MaxCandidates)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGeneratePatternName(t *testing.T) {
|
||||
tests := []struct {
|
||||
patternType models.PatternType
|
||||
signature []string
|
||||
title string
|
||||
wantPrefix string
|
||||
signature []string
|
||||
}{
|
||||
{models.PatternTypeBug, []string{"nil", "error"}, "", "Bug Pattern:"},
|
||||
{models.PatternTypeRefactor, []string{"extract"}, "", "Refactor Pattern:"},
|
||||
{models.PatternTypeArchitecture, []string{"service"}, "", "Architecture Pattern:"},
|
||||
{models.PatternTypeAntiPattern, []string{"god-class"}, "", "Anti-Pattern:"},
|
||||
{models.PatternTypeBestPractice, []string{"testing"}, "", "Best Practice:"},
|
||||
{models.PatternTypeBug, []string{}, "Short Title", "Short Title"}, // Use title directly
|
||||
{patternType: models.PatternTypeBug, title: "", wantPrefix: "Bug Pattern:", signature: []string{"nil", "error"}},
|
||||
{patternType: models.PatternTypeRefactor, title: "", wantPrefix: "Refactor Pattern:", signature: []string{"extract"}},
|
||||
{patternType: models.PatternTypeArchitecture, title: "", wantPrefix: "Architecture Pattern:", signature: []string{"service"}},
|
||||
{patternType: models.PatternTypeAntiPattern, title: "", wantPrefix: "Anti-Pattern:", signature: []string{"god-class"}},
|
||||
{patternType: models.PatternTypeBestPractice, title: "", wantPrefix: "Best Practice:", signature: []string{"testing"}},
|
||||
{patternType: models.PatternTypeBug, title: "Short Title", wantPrefix: "Short Title", signature: []string{}}, // Use title directly
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
@@ -352,6 +372,85 @@ func TestGeneratePatternName(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestGeneratePatternName_EdgeCases(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
ptype models.PatternType
|
||||
title string
|
||||
want string
|
||||
signature []string
|
||||
}{
|
||||
{
|
||||
name: "with title uses title directly",
|
||||
ptype: models.PatternTypeBug,
|
||||
signature: []string{"ignored"},
|
||||
title: "Custom Title",
|
||||
want: "Custom Title",
|
||||
},
|
||||
{
|
||||
name: "long title generates from signature",
|
||||
ptype: models.PatternTypeBug,
|
||||
signature: []string{"sig1", "sig2"},
|
||||
title: "This is a very long title that exceeds sixty characters and should be ignored",
|
||||
want: "Bug Pattern: sig1, sig2",
|
||||
},
|
||||
{
|
||||
name: "empty signature returns Unnamed",
|
||||
ptype: models.PatternTypeBug,
|
||||
signature: []string{},
|
||||
title: "",
|
||||
want: "Bug Pattern: Unnamed",
|
||||
},
|
||||
{
|
||||
name: "single signature element",
|
||||
ptype: models.PatternTypeRefactor,
|
||||
signature: []string{"single"},
|
||||
title: "",
|
||||
want: "Refactor Pattern: single",
|
||||
},
|
||||
{
|
||||
name: "more than 3 signature elements truncates",
|
||||
ptype: models.PatternTypeBestPractice,
|
||||
signature: []string{"a", "b", "c", "d", "e"},
|
||||
title: "",
|
||||
want: "Best Practice: a, b, c",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
got := generatePatternName(tt.ptype, tt.signature, tt.title)
|
||||
if got != tt.want {
|
||||
t.Errorf("generatePatternName() = %q, want %q", got, tt.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGeneratePatternName_AllTypes(t *testing.T) {
|
||||
tests := []struct {
|
||||
ptype models.PatternType
|
||||
wantPrefix string
|
||||
}{
|
||||
{models.PatternTypeBug, "Bug Pattern:"},
|
||||
{models.PatternTypeRefactor, "Refactor Pattern:"},
|
||||
{models.PatternTypeArchitecture, "Architecture Pattern:"},
|
||||
{models.PatternTypeAntiPattern, "Anti-Pattern:"},
|
||||
{models.PatternTypeBestPractice, "Best Practice:"},
|
||||
{models.PatternType("unknown"), "test"}, // Unknown type has empty prefix, starts with first signature element
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(string(tt.ptype), func(t *testing.T) {
|
||||
name := generatePatternName(tt.ptype, []string{"test", "sig"}, "")
|
||||
if !hasPrefix(name, tt.wantPrefix) {
|
||||
t.Errorf("Expected prefix %q for type %s, got: %s",
|
||||
tt.wantPrefix, tt.ptype, name)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestFormatPatternInsight(t *testing.T) {
|
||||
// Pattern without recommendation
|
||||
pattern1 := &models.Pattern{
|
||||
@@ -386,6 +485,470 @@ func TestFormatPatternInsight(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestFormatPatternInsight_AllTypes(t *testing.T) {
|
||||
types := []struct {
|
||||
ptype models.PatternType
|
||||
contains string
|
||||
}{
|
||||
{models.PatternTypeBug, "bug pattern"},
|
||||
{models.PatternTypeRefactor, "recognized pattern"}, // Falls to default case
|
||||
{models.PatternTypeArchitecture, "recognized pattern"}, // Falls to default case
|
||||
{models.PatternTypeAntiPattern, "anti-pattern"},
|
||||
{models.PatternTypeBestPractice, "best practice"},
|
||||
{models.PatternType("unknown"), "recognized pattern"}, // Falls to default case
|
||||
}
|
||||
|
||||
for _, tt := range types {
|
||||
t.Run(string(tt.ptype), func(t *testing.T) {
|
||||
pattern := &models.Pattern{
|
||||
Type: tt.ptype,
|
||||
Frequency: 3,
|
||||
Projects: []string{"proj1"},
|
||||
}
|
||||
insight := formatPatternInsight(pattern)
|
||||
if !containsString(insight, tt.contains) {
|
||||
t.Errorf("Expected insight to contain %q for type %s, got: %s",
|
||||
tt.contains, tt.ptype, insight)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestFormatPatternInsight_MultiProject(t *testing.T) {
|
||||
pattern := &models.Pattern{
|
||||
Type: models.PatternTypeBug,
|
||||
Frequency: 10,
|
||||
Projects: []string{"proj1", "proj2", "proj3"},
|
||||
}
|
||||
|
||||
insight := formatPatternInsight(pattern)
|
||||
|
||||
if !containsString(insight, "10 times") {
|
||||
t.Error("Expected frequency in insight")
|
||||
}
|
||||
if !containsString(insight, "3 projects") {
|
||||
t.Error("Expected project count in insight")
|
||||
}
|
||||
}
|
||||
|
||||
func TestFormatPatternInsight_SingleProject(t *testing.T) {
|
||||
pattern := &models.Pattern{
|
||||
Type: models.PatternTypeBestPractice,
|
||||
Frequency: 5,
|
||||
Projects: []string{"only-one"},
|
||||
}
|
||||
|
||||
insight := formatPatternInsight(pattern)
|
||||
|
||||
if !containsString(insight, "5 times") {
|
||||
t.Error("Expected frequency in insight")
|
||||
}
|
||||
// Single project should NOT mention "projects"
|
||||
if containsString(insight, "projects") {
|
||||
t.Error("Single project should not mention 'projects'")
|
||||
}
|
||||
}
|
||||
|
||||
func TestDetector_SetSyncFunc(t *testing.T) {
|
||||
store := setupTestStore(t)
|
||||
defer store.Close()
|
||||
|
||||
patternStore := gorm.NewPatternStore(store)
|
||||
observationStore := gorm.NewObservationStore(store, nil, nil, nil)
|
||||
config := DefaultConfig()
|
||||
|
||||
detector := NewDetector(patternStore, observationStore, config)
|
||||
|
||||
// Initially nil
|
||||
if detector.syncFunc != nil {
|
||||
t.Error("Expected syncFunc to be nil initially")
|
||||
}
|
||||
|
||||
// Set sync func
|
||||
var syncCalled bool
|
||||
detector.SetSyncFunc(func(p *models.Pattern) {
|
||||
syncCalled = true
|
||||
})
|
||||
|
||||
if detector.syncFunc == nil {
|
||||
t.Error("Expected syncFunc to be set")
|
||||
}
|
||||
|
||||
// Verify it can be called
|
||||
detector.syncFunc(&models.Pattern{})
|
||||
if !syncCalled {
|
||||
t.Error("Expected sync function to be called")
|
||||
}
|
||||
}
|
||||
|
||||
func TestDetector_CandidateCount(t *testing.T) {
|
||||
store := setupTestStore(t)
|
||||
defer store.Close()
|
||||
|
||||
patternStore := gorm.NewPatternStore(store)
|
||||
observationStore := gorm.NewObservationStore(store, nil, nil, nil)
|
||||
config := DefaultConfig()
|
||||
|
||||
detector := NewDetector(patternStore, observationStore, config)
|
||||
|
||||
// Initially zero
|
||||
if count := detector.CandidateCount(); count != 0 {
|
||||
t.Errorf("Expected 0 candidates, got %d", count)
|
||||
}
|
||||
|
||||
// Add some candidates
|
||||
detector.candidates["key1"] = &candidatePattern{}
|
||||
detector.candidates["key2"] = &candidatePattern{}
|
||||
|
||||
if count := detector.CandidateCount(); count != 2 {
|
||||
t.Errorf("Expected 2 candidates, got %d", count)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDetector_AnalyzeRecentObservations(t *testing.T) {
|
||||
store := setupTestStore(t)
|
||||
defer store.Close()
|
||||
|
||||
patternStore := gorm.NewPatternStore(store)
|
||||
observationStore := gorm.NewObservationStore(store, nil, nil, nil)
|
||||
config := DefaultConfig()
|
||||
|
||||
detector := NewDetector(patternStore, observationStore, config)
|
||||
ctx := context.Background()
|
||||
|
||||
// Should not error even with no observations
|
||||
err := detector.AnalyzeRecentObservations(ctx)
|
||||
if err != nil {
|
||||
t.Fatalf("AnalyzeRecentObservations() error = %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGenerateCandidateKey(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
want string
|
||||
signature []string
|
||||
}{
|
||||
{
|
||||
name: "single element",
|
||||
signature: []string{"error"},
|
||||
want: "error|",
|
||||
},
|
||||
{
|
||||
name: "multiple elements",
|
||||
signature: []string{"error", "handling", "nil"},
|
||||
want: "error|handling|nil|",
|
||||
},
|
||||
{
|
||||
name: "empty signature",
|
||||
signature: []string{},
|
||||
want: "",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
got := generateCandidateKey(tt.signature)
|
||||
if got != tt.want {
|
||||
t.Errorf("generateCandidateKey() = %q, want %q", got, tt.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGenerateCandidateKey_EdgeCases(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
want string
|
||||
signature []string
|
||||
}{
|
||||
{
|
||||
name: "nil signature",
|
||||
signature: nil,
|
||||
want: "",
|
||||
},
|
||||
{
|
||||
name: "empty strings in signature",
|
||||
signature: []string{"", ""},
|
||||
want: "||",
|
||||
},
|
||||
{
|
||||
name: "special characters",
|
||||
signature: []string{"error|handling", "nil"},
|
||||
want: "error|handling|nil|",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
got := generateCandidateKey(tt.signature)
|
||||
if got != tt.want {
|
||||
t.Errorf("generateCandidateKey() = %q, want %q", got, tt.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestItoa(t *testing.T) {
|
||||
tests := []struct {
|
||||
want string
|
||||
input int
|
||||
}{
|
||||
{"0", 0},
|
||||
{"1", 1},
|
||||
{"10", 10},
|
||||
{"123", 123},
|
||||
{"-1", -1},
|
||||
{"-123", -123},
|
||||
{"1000000", 1000000},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.want, func(t *testing.T) {
|
||||
got := itoa(tt.input)
|
||||
if got != tt.want {
|
||||
t.Errorf("itoa(%d) = %q, want %q", tt.input, got, tt.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestItoa_EdgeCases(t *testing.T) {
|
||||
tests := []struct {
|
||||
want string
|
||||
input int
|
||||
}{
|
||||
{"0", 0},
|
||||
{"0", -0},
|
||||
{"1", 1},
|
||||
{"-1", -1},
|
||||
{"9", 9},
|
||||
{"10", 10},
|
||||
{"99", 99},
|
||||
{"100", 100},
|
||||
{"999", 999},
|
||||
{"1000", 1000},
|
||||
{"-999", -999},
|
||||
{"-1000", -1000},
|
||||
{"2147483647", 2147483647}, // Max int32
|
||||
{"-2147483647", -2147483647}, // Min int32 + 1
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.want, func(t *testing.T) {
|
||||
got := itoa(tt.input)
|
||||
if got != tt.want {
|
||||
t.Errorf("itoa(%d) = %q, want %q", tt.input, got, tt.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestDetectionResult_ZeroValue(t *testing.T) {
|
||||
result := &DetectionResult{}
|
||||
|
||||
if result.MatchedPattern != nil {
|
||||
t.Error("Zero value should have nil MatchedPattern")
|
||||
}
|
||||
if result.MatchScore != 0 {
|
||||
t.Error("Zero value should have 0 MatchScore")
|
||||
}
|
||||
if result.IsNewPattern {
|
||||
t.Error("Zero value should have false IsNewPattern")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCandidatePattern_Fields(t *testing.T) {
|
||||
candidate := &candidatePattern{
|
||||
patternType: models.PatternTypeBug,
|
||||
title: "Test Title",
|
||||
signature: []string{"sig1", "sig2"},
|
||||
observationIDs: []int64{1, 2, 3},
|
||||
projects: []string{"proj1", "proj2"},
|
||||
lastSeenEpoch: time.Now().UnixMilli(),
|
||||
}
|
||||
|
||||
if candidate.patternType != models.PatternTypeBug {
|
||||
t.Error("Wrong pattern type")
|
||||
}
|
||||
if candidate.title != "Test Title" {
|
||||
t.Error("Wrong title")
|
||||
}
|
||||
if len(candidate.signature) != 2 {
|
||||
t.Error("Wrong signature length")
|
||||
}
|
||||
if len(candidate.observationIDs) != 3 {
|
||||
t.Error("Wrong observationIDs length")
|
||||
}
|
||||
if len(candidate.projects) != 2 {
|
||||
t.Error("Wrong projects length")
|
||||
}
|
||||
}
|
||||
|
||||
func TestDetector_AnalyzeObservation_EmptySignature(t *testing.T) {
|
||||
store := setupTestStore(t)
|
||||
defer store.Close()
|
||||
|
||||
patternStore := gorm.NewPatternStore(store)
|
||||
observationStore := gorm.NewObservationStore(store, nil, nil, nil)
|
||||
config := DefaultConfig()
|
||||
|
||||
detector := NewDetector(patternStore, observationStore, config)
|
||||
ctx := context.Background()
|
||||
|
||||
// Create observation with empty concepts/title/narrative
|
||||
obs := &models.Observation{
|
||||
ID: 1,
|
||||
SDKSessionID: "test-session",
|
||||
Project: "test-project",
|
||||
Scope: models.ScopeProject,
|
||||
Type: models.ObsTypeBugfix,
|
||||
// All fields that would create signature are empty
|
||||
}
|
||||
|
||||
result, err := detector.AnalyzeObservation(ctx, obs)
|
||||
if err != nil {
|
||||
t.Fatalf("AnalyzeObservation() error = %v", err)
|
||||
}
|
||||
|
||||
// Should return empty result for empty signature
|
||||
if result.MatchedPattern != nil {
|
||||
t.Error("Expected nil pattern for empty signature")
|
||||
}
|
||||
}
|
||||
|
||||
func TestDetector_AnalyzeObservation_CandidateEviction(t *testing.T) {
|
||||
store := setupTestStore(t)
|
||||
defer store.Close()
|
||||
|
||||
patternStore := gorm.NewPatternStore(store)
|
||||
observationStore := gorm.NewObservationStore(store, nil, nil, nil)
|
||||
config := DefaultConfig()
|
||||
config.MaxCandidates = 2 // Very small for testing
|
||||
config.MinFrequencyForPattern = 10 // High so nothing gets promoted
|
||||
|
||||
detector := NewDetector(patternStore, observationStore, config)
|
||||
ctx := context.Background()
|
||||
|
||||
// Add observations with different signatures until we exceed MaxCandidates
|
||||
obs1 := createTestObservation(1, "First", []string{"first", "unique"})
|
||||
obs2 := createTestObservation(2, "Second", []string{"second", "unique"})
|
||||
obs3 := createTestObservation(3, "Third", []string{"third", "unique"})
|
||||
|
||||
// Analyze all observations
|
||||
_, _ = detector.AnalyzeObservation(ctx, obs1)
|
||||
time.Sleep(10 * time.Millisecond) // Small delay so timestamps differ
|
||||
_, _ = detector.AnalyzeObservation(ctx, obs2)
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
_, _ = detector.AnalyzeObservation(ctx, obs3)
|
||||
|
||||
// Should have at most MaxCandidates
|
||||
if count := detector.CandidateCount(); count > config.MaxCandidates {
|
||||
t.Errorf("Expected at most %d candidates, got %d", config.MaxCandidates, count)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDetector_PromoteCandidateWithSyncFunc(t *testing.T) {
|
||||
store := setupTestStore(t)
|
||||
defer store.Close()
|
||||
|
||||
patternStore := gorm.NewPatternStore(store)
|
||||
observationStore := gorm.NewObservationStore(store, nil, nil, nil)
|
||||
config := DefaultConfig()
|
||||
config.MinFrequencyForPattern = 2
|
||||
|
||||
detector := NewDetector(patternStore, observationStore, config)
|
||||
ctx := context.Background()
|
||||
|
||||
// Set up sync function to track calls
|
||||
var syncedPattern *models.Pattern
|
||||
detector.SetSyncFunc(func(p *models.Pattern) {
|
||||
syncedPattern = p
|
||||
})
|
||||
|
||||
// Create two similar observations to trigger pattern promotion
|
||||
obs1 := createTestObservation(1, "Sync Test", []string{"sync", "test"})
|
||||
obs2 := createTestObservation(2, "Sync Test", []string{"sync", "test"})
|
||||
|
||||
_, _ = detector.AnalyzeObservation(ctx, obs1)
|
||||
result, _ := detector.AnalyzeObservation(ctx, obs2)
|
||||
|
||||
if result.MatchedPattern == nil {
|
||||
t.Fatal("Expected pattern to be created")
|
||||
}
|
||||
|
||||
if syncedPattern == nil {
|
||||
t.Error("Expected sync function to be called")
|
||||
}
|
||||
|
||||
if syncedPattern != nil && syncedPattern.Name != result.MatchedPattern.Name {
|
||||
t.Errorf("Synced pattern name mismatch: got %s, want %s",
|
||||
syncedPattern.Name, result.MatchedPattern.Name)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDetector_AnalyzeObservation_UpdateExistingCandidate(t *testing.T) {
|
||||
store := setupTestStore(t)
|
||||
defer store.Close()
|
||||
|
||||
patternStore := gorm.NewPatternStore(store)
|
||||
observationStore := gorm.NewObservationStore(store, nil, nil, nil)
|
||||
config := DefaultConfig()
|
||||
config.MinFrequencyForPattern = 5 // High enough that we don't promote
|
||||
|
||||
detector := NewDetector(patternStore, observationStore, config)
|
||||
ctx := context.Background()
|
||||
|
||||
// Create observations with same signature
|
||||
obs1 := createTestObservation(1, "Update Test", []string{"update", "test"})
|
||||
obs2 := createTestObservation(2, "Update Test", []string{"update", "test"})
|
||||
obs2.Project = "different-project"
|
||||
|
||||
// Analyze first observation
|
||||
_, _ = detector.AnalyzeObservation(ctx, obs1)
|
||||
|
||||
// Check candidate count
|
||||
if count := detector.CandidateCount(); count != 1 {
|
||||
t.Errorf("Expected 1 candidate after first obs, got %d", count)
|
||||
}
|
||||
|
||||
// Analyze second observation
|
||||
_, _ = detector.AnalyzeObservation(ctx, obs2)
|
||||
|
||||
// Still 1 candidate (same signature)
|
||||
if count := detector.CandidateCount(); count != 1 {
|
||||
t.Errorf("Expected 1 candidate after second obs, got %d", count)
|
||||
}
|
||||
|
||||
// Check that candidate has both projects
|
||||
key := generateCandidateKey([]string{"update", "test"})
|
||||
candidate := detector.candidates[key]
|
||||
if candidate == nil {
|
||||
t.Fatal("Expected candidate to exist")
|
||||
}
|
||||
if len(candidate.projects) != 2 {
|
||||
t.Errorf("Expected 2 projects, got %d", len(candidate.projects))
|
||||
}
|
||||
}
|
||||
|
||||
func TestDetector_GetPatternInsight_NotFound(t *testing.T) {
|
||||
store := setupTestStore(t)
|
||||
defer store.Close()
|
||||
|
||||
patternStore := gorm.NewPatternStore(store)
|
||||
observationStore := gorm.NewObservationStore(store, nil, nil, nil)
|
||||
config := DefaultConfig()
|
||||
|
||||
detector := NewDetector(patternStore, observationStore, config)
|
||||
ctx := context.Background()
|
||||
|
||||
// Try to get insight for non-existent pattern
|
||||
_, err := detector.GetPatternInsight(ctx, 99999)
|
||||
if err == nil {
|
||||
t.Error("Expected error for non-existent pattern")
|
||||
}
|
||||
}
|
||||
|
||||
// Helper functions
|
||||
|
||||
func setupTestStore(t *testing.T) *gorm.Store {
|
||||
|
||||
@@ -0,0 +1,98 @@
|
||||
// Package privacy provides utilities for protecting sensitive data.
|
||||
package privacy
|
||||
|
||||
import (
|
||||
"regexp"
|
||||
"slices"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// secretPatterns contains compiled regular expressions for detecting secrets.
|
||||
// These patterns are designed to catch common secret formats with minimal false positives.
|
||||
var secretPatterns = []*regexp.Regexp{
|
||||
// API keys with common prefixes
|
||||
regexp.MustCompile(`(?i)(api[_-]?key|apikey)\s*[:=]\s*['"]?[a-zA-Z0-9_-]{20,}['"]?`),
|
||||
|
||||
// Passwords in configuration
|
||||
regexp.MustCompile(`(?i)(password|passwd|pwd)\s*[:=]\s*['"][^'"]{8,}['"]`),
|
||||
|
||||
// Secret tokens
|
||||
regexp.MustCompile(`(?i)(secret[_-]?key|secret[_-]?token|auth[_-]?token)\s*[:=]\s*['"]?[a-zA-Z0-9_-]{20,}['"]?`),
|
||||
|
||||
// OpenAI API keys
|
||||
regexp.MustCompile(`sk-[a-zA-Z0-9]{20,}`),
|
||||
|
||||
// Anthropic API keys
|
||||
regexp.MustCompile(`sk-ant-[a-zA-Z0-9-]{20,}`),
|
||||
|
||||
// GitHub tokens
|
||||
regexp.MustCompile(`gh[pous]_[a-zA-Z0-9]{36,}`),
|
||||
regexp.MustCompile(`github_pat_[a-zA-Z0-9_]{22,}`),
|
||||
|
||||
// AWS keys
|
||||
regexp.MustCompile(`AKIA[0-9A-Z]{16}`),
|
||||
regexp.MustCompile(`(?i)aws[_-]?secret[_-]?access[_-]?key\s*[:=]\s*['"]?[a-zA-Z0-9/+=]{40}['"]?`),
|
||||
|
||||
// Private keys (PEM format indicators)
|
||||
regexp.MustCompile(`-----BEGIN (RSA |EC |DSA |OPENSSH )?PRIVATE KEY-----`),
|
||||
|
||||
// JWT tokens (base64.base64.base64 format)
|
||||
regexp.MustCompile(`eyJ[a-zA-Z0-9_-]+\.eyJ[a-zA-Z0-9_-]+\.[a-zA-Z0-9_-]+`),
|
||||
|
||||
// Generic secret assignment patterns
|
||||
regexp.MustCompile(`(?i)bearer\s+[a-zA-Z0-9_-]{20,}`),
|
||||
}
|
||||
|
||||
// ContainsSecrets checks if the given text contains any patterns that look like secrets.
|
||||
// Returns true if potential secrets are detected.
|
||||
func ContainsSecrets(text string) bool {
|
||||
if text == "" {
|
||||
return false
|
||||
}
|
||||
|
||||
for _, pattern := range secretPatterns {
|
||||
if pattern.MatchString(text) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// RedactSecrets replaces detected secrets with a redaction marker.
|
||||
// This allows the text to be stored while protecting sensitive data.
|
||||
func RedactSecrets(text string) string {
|
||||
if text == "" {
|
||||
return text
|
||||
}
|
||||
|
||||
result := text
|
||||
for _, pattern := range secretPatterns {
|
||||
result = pattern.ReplaceAllStringFunc(result, func(match string) string {
|
||||
// Preserve the key name, redact only the value
|
||||
if idx := strings.Index(match, "="); idx != -1 {
|
||||
return match[:idx+1] + "[REDACTED]"
|
||||
}
|
||||
if idx := strings.Index(match, ":"); idx != -1 {
|
||||
return match[:idx+1] + "[REDACTED]"
|
||||
}
|
||||
// For standalone secrets, show just the prefix
|
||||
if len(match) > 8 {
|
||||
return match[:4] + "...[REDACTED]"
|
||||
}
|
||||
return "[REDACTED]"
|
||||
})
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
// SanitizeObservation checks multiple fields of an observation for secrets.
|
||||
// Returns true if any secrets were found.
|
||||
// This function is used as a validation gate before storing observations.
|
||||
func SanitizeObservation(narrative string, facts []string) bool {
|
||||
if ContainsSecrets(narrative) {
|
||||
return true
|
||||
}
|
||||
return slices.ContainsFunc(facts, func(fact string) bool {
|
||||
return ContainsSecrets(fact)
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,213 @@
|
||||
package privacy
|
||||
|
||||
import (
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestContainsSecrets(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
input string
|
||||
expected bool
|
||||
}{
|
||||
{
|
||||
name: "empty string",
|
||||
input: "",
|
||||
expected: false,
|
||||
},
|
||||
{
|
||||
name: "normal text",
|
||||
input: "This is just some regular text about a bug fix",
|
||||
expected: false,
|
||||
},
|
||||
{
|
||||
name: "API key pattern",
|
||||
input: "api_key=abc123def456ghi789jkl012mno345pqr678",
|
||||
expected: true,
|
||||
},
|
||||
{
|
||||
name: "api-key with dash",
|
||||
input: `api-key: "abc123def456ghi789jkl012mno"`,
|
||||
expected: true,
|
||||
},
|
||||
{
|
||||
name: "password in config",
|
||||
input: `password="super_secret_password_123"`,
|
||||
expected: true,
|
||||
},
|
||||
{
|
||||
name: "OpenAI key format",
|
||||
input: "sk-abc123def456ghi789jkl012mno345pqr678",
|
||||
expected: true,
|
||||
},
|
||||
{
|
||||
name: "Anthropic key format",
|
||||
input: "sk-ant-api03-abc123def456ghi789jkl012mno345",
|
||||
expected: true,
|
||||
},
|
||||
{
|
||||
name: "GitHub PAT",
|
||||
input: "ghp_1234567890abcdefghijklmnopqrstuvwxyz",
|
||||
expected: true,
|
||||
},
|
||||
{
|
||||
name: "GitHub PAT new format",
|
||||
input: "github_pat_12ABCDEFGHIJ3456789abc_defghijklmno",
|
||||
expected: true,
|
||||
},
|
||||
{
|
||||
name: "AWS access key",
|
||||
input: "AKIAIOSFODNN7EXAMPLE",
|
||||
expected: true,
|
||||
},
|
||||
{
|
||||
name: "Private key header",
|
||||
input: "-----BEGIN RSA PRIVATE KEY-----",
|
||||
expected: true,
|
||||
},
|
||||
{
|
||||
name: "JWT token",
|
||||
input: "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIn0.dozjgNryP4J3jVmNHl0w5N_XgL0n3I9PlFUP0THsR8U",
|
||||
expected: true,
|
||||
},
|
||||
{
|
||||
name: "bearer token",
|
||||
input: "Bearer abc123def456ghi789jkl012mno345",
|
||||
expected: true,
|
||||
},
|
||||
{
|
||||
name: "secret_key in code",
|
||||
input: `secret_key = "my_super_secret_token_here"`,
|
||||
expected: true,
|
||||
},
|
||||
{
|
||||
name: "short password is not detected",
|
||||
input: `password="short"`,
|
||||
expected: false, // Too short to trigger
|
||||
},
|
||||
{
|
||||
name: "word password in sentence",
|
||||
input: "The password field should be validated",
|
||||
expected: false,
|
||||
},
|
||||
{
|
||||
name: "word api in code",
|
||||
input: "The API returns JSON data",
|
||||
expected: false,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
result := ContainsSecrets(tt.input)
|
||||
if result != tt.expected {
|
||||
t.Errorf("ContainsSecrets(%q) = %v, want %v", tt.input, result, tt.expected)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestRedactSecrets(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
input string
|
||||
expected string
|
||||
}{
|
||||
{
|
||||
name: "empty string",
|
||||
input: "",
|
||||
expected: "",
|
||||
},
|
||||
{
|
||||
name: "no secrets",
|
||||
input: "This is safe text",
|
||||
expected: "This is safe text",
|
||||
},
|
||||
{
|
||||
name: "API key gets redacted",
|
||||
input: "api_key=abc123def456ghi789jkl012mno345pqr678",
|
||||
expected: "api_key=[REDACTED]",
|
||||
},
|
||||
{
|
||||
name: "OpenAI key gets redacted",
|
||||
input: "The key is sk-abc123def456ghi789jkl012mno345pqr678",
|
||||
expected: "The key is sk-a...[REDACTED]",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
result := RedactSecrets(tt.input)
|
||||
if result != tt.expected {
|
||||
t.Errorf("RedactSecrets(%q) = %q, want %q", tt.input, result, tt.expected)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSanitizeObservation(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
narrative string
|
||||
facts []string
|
||||
expected bool
|
||||
}{
|
||||
{
|
||||
name: "clean observation",
|
||||
narrative: "Fixed a bug in the login flow",
|
||||
facts: []string{"Users can now log in", "Session management improved"},
|
||||
expected: false,
|
||||
},
|
||||
{
|
||||
name: "secret in narrative",
|
||||
narrative: "Set API key api_key=abc123def456ghi789jkl012mno345",
|
||||
facts: []string{"Configuration updated"},
|
||||
expected: true,
|
||||
},
|
||||
{
|
||||
name: "secret in facts",
|
||||
narrative: "Updated configuration",
|
||||
facts: []string{"Added api_key=abc123def456ghi789jkl012mno345"},
|
||||
expected: true,
|
||||
},
|
||||
{
|
||||
name: "empty facts",
|
||||
narrative: "Clean narrative",
|
||||
facts: []string{},
|
||||
expected: false,
|
||||
},
|
||||
{
|
||||
name: "nil facts",
|
||||
narrative: "Clean narrative",
|
||||
facts: nil,
|
||||
expected: false,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
result := SanitizeObservation(tt.narrative, tt.facts)
|
||||
if result != tt.expected {
|
||||
t.Errorf("SanitizeObservation() = %v, want %v", result, tt.expected)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkContainsSecrets(b *testing.B) {
|
||||
text := "This is a normal piece of text that does not contain any secrets or sensitive information"
|
||||
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
ContainsSecrets(text)
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkContainsSecretsWithSecret(b *testing.B) {
|
||||
text := "api_key=abc123def456ghi789jkl012mno345pqr678"
|
||||
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
ContainsSecrets(text)
|
||||
}
|
||||
}
|
||||
@@ -30,24 +30,24 @@ const (
|
||||
|
||||
// Candidate represents a search result candidate for reranking.
|
||||
type Candidate struct {
|
||||
ID string // Document ID
|
||||
Content string // Document text content for scoring
|
||||
Score float64 // Original bi-encoder similarity score
|
||||
Metadata map[string]any // Preserved metadata
|
||||
RerankInfo map[string]float64 // Reranking debug info (optional)
|
||||
Metadata map[string]any
|
||||
RerankInfo map[string]float64
|
||||
ID string
|
||||
Content string
|
||||
Score float64
|
||||
}
|
||||
|
||||
// RerankResult represents a reranked search result.
|
||||
type RerankResult struct {
|
||||
ID string // Document ID
|
||||
Content string // Document text content
|
||||
OriginalScore float64 // Original bi-encoder score
|
||||
RerankScore float64 // Cross-encoder relevance score
|
||||
CombinedScore float64 // Weighted combination of scores
|
||||
Metadata map[string]any // Preserved metadata
|
||||
OriginalRank int // Position before reranking (1-indexed)
|
||||
RerankRank int // Position after reranking (1-indexed)
|
||||
RankImprovement int // How much the rank improved (positive = moved up)
|
||||
Metadata map[string]any
|
||||
ID string
|
||||
Content string
|
||||
OriginalScore float64
|
||||
RerankScore float64
|
||||
CombinedScore float64
|
||||
OriginalRank int
|
||||
RerankRank int
|
||||
RankImprovement int
|
||||
}
|
||||
|
||||
// Service provides cross-encoder reranking functionality.
|
||||
|
||||
@@ -36,6 +36,13 @@ func NewCalculator(config *models.ScoringConfig) *Calculator {
|
||||
// - ConceptContrib = sum(concept_weights) × concept_weight_factor
|
||||
// - RetrievalContrib = log2(retrieval_count + 1) × 0.1 × retrieval_weight
|
||||
func (c *Calculator) Calculate(obs *models.Observation, now time.Time) float64 {
|
||||
return c.CalculateComponents(obs, now).FinalScore
|
||||
}
|
||||
|
||||
// CalculateComponents returns the individual components of the importance score.
|
||||
// Useful for debugging and explaining scores to users.
|
||||
// This is the core calculation method - Calculate() delegates to this.
|
||||
func (c *Calculator) CalculateComponents(obs *models.Observation, now time.Time) ScoreComponents {
|
||||
// 1. Get base type weight
|
||||
typeWeight := models.TypeBaseScore(obs.Type)
|
||||
|
||||
@@ -75,42 +82,6 @@ func (c *Calculator) Calculate(obs *models.Observation, now time.Time) float64 {
|
||||
finalScore = c.config.MinScore
|
||||
}
|
||||
|
||||
return finalScore
|
||||
}
|
||||
|
||||
// CalculateComponents returns the individual components of the importance score.
|
||||
// Useful for debugging and explaining scores to users.
|
||||
func (c *Calculator) CalculateComponents(obs *models.Observation, now time.Time) ScoreComponents {
|
||||
typeWeight := models.TypeBaseScore(obs.Type)
|
||||
|
||||
ageDays := now.Sub(time.UnixMilli(obs.CreatedAtEpoch)).Hours() / 24.0
|
||||
if ageDays < 0 {
|
||||
ageDays = 0
|
||||
}
|
||||
recencyDecay := math.Pow(0.5, ageDays/c.config.RecencyHalfLifeDays)
|
||||
|
||||
coreScore := 1.0 * typeWeight * recencyDecay
|
||||
feedbackContrib := float64(obs.UserFeedback) * c.config.FeedbackWeight
|
||||
|
||||
conceptBoost := 0.0
|
||||
for _, concept := range obs.Concepts {
|
||||
if weight, ok := c.config.ConceptWeights[concept]; ok {
|
||||
conceptBoost += weight
|
||||
}
|
||||
}
|
||||
conceptContrib := conceptBoost * c.config.ConceptWeight
|
||||
|
||||
retrievalContrib := 0.0
|
||||
if obs.RetrievalCount > 0 {
|
||||
retrievalBoost := math.Log2(float64(obs.RetrievalCount)+1) * 0.1
|
||||
retrievalContrib = retrievalBoost * c.config.RetrievalWeight
|
||||
}
|
||||
|
||||
finalScore := coreScore + feedbackContrib + conceptContrib + retrievalContrib
|
||||
if finalScore < c.config.MinScore {
|
||||
finalScore = c.config.MinScore
|
||||
}
|
||||
|
||||
return ScoreComponents{
|
||||
TypeWeight: typeWeight,
|
||||
RecencyDecay: recencyDecay,
|
||||
|
||||
@@ -21,13 +21,13 @@ type ObservationStore interface {
|
||||
|
||||
// Recalculator periodically recalculates importance scores for observations.
|
||||
type Recalculator struct {
|
||||
log zerolog.Logger
|
||||
store ObservationStore
|
||||
calculator *Calculator
|
||||
log zerolog.Logger
|
||||
interval time.Duration
|
||||
batchSize int
|
||||
stopCh chan struct{}
|
||||
doneCh chan struct{}
|
||||
interval time.Duration
|
||||
batchSize int
|
||||
mu sync.Mutex
|
||||
running bool
|
||||
}
|
||||
|
||||
@@ -16,14 +16,14 @@ import (
|
||||
|
||||
// MockObservationStore is a mock implementation of ObservationStore for testing.
|
||||
type MockObservationStore struct {
|
||||
mu sync.Mutex
|
||||
observations []*models.Observation
|
||||
scores map[int64]float64
|
||||
conceptWeights map[string]float64
|
||||
updateErr error
|
||||
getErr error
|
||||
getConceptsErr error
|
||||
scores map[int64]float64
|
||||
conceptWeights map[string]float64
|
||||
observations []*models.Observation
|
||||
updateScoresCalls int
|
||||
mu sync.Mutex
|
||||
}
|
||||
|
||||
func NewMockObservationStore() *MockObservationStore {
|
||||
|
||||
@@ -3,7 +3,9 @@ package expansion
|
||||
|
||||
import (
|
||||
"context"
|
||||
"math"
|
||||
"regexp"
|
||||
"sort"
|
||||
"strings"
|
||||
"sync"
|
||||
|
||||
@@ -30,25 +32,25 @@ const (
|
||||
// ExpandedQuery represents a query variant with metadata.
|
||||
type ExpandedQuery struct {
|
||||
Query string `json:"query"`
|
||||
Weight float64 `json:"weight"` // Weight for result merging (0.0-1.0)
|
||||
Source string `json:"source"` // Where this expansion came from
|
||||
Intent QueryIntent `json:"intent"` // Detected intent
|
||||
Source string `json:"source"`
|
||||
Intent QueryIntent `json:"intent"`
|
||||
Weight float64 `json:"weight"`
|
||||
}
|
||||
|
||||
// Expander provides context-aware query expansion.
|
||||
type Expander struct {
|
||||
embedSvc *embedding.Service
|
||||
vocabulary []VocabEntry // Known vocabulary from observations
|
||||
vocabVectors [][]float32 // Embeddings for vocabulary entries
|
||||
vocabMu sync.RWMutex // Protects vocabulary
|
||||
intentPatterns map[QueryIntent][]*regexp.Regexp
|
||||
vocabulary []VocabEntry
|
||||
vocabVectors [][]float32
|
||||
vocabMu sync.RWMutex
|
||||
}
|
||||
|
||||
// VocabEntry represents a vocabulary term from observations.
|
||||
type VocabEntry struct {
|
||||
Term string // The term itself
|
||||
Weight float64 // How common/important this term is (0.0-1.0)
|
||||
Source string // Where it came from (title, concept, narrative)
|
||||
Term string
|
||||
Source string
|
||||
Weight float64
|
||||
}
|
||||
|
||||
// Config holds expander configuration.
|
||||
@@ -281,14 +283,10 @@ func (e *Expander) expandByVocabulary(ctx context.Context, query string, minSimi
|
||||
return nil
|
||||
}
|
||||
|
||||
// Sort by score (descending) using bubble sort
|
||||
for i := 0; i < len(similar)-1; i++ {
|
||||
for j := i + 1; j < len(similar); j++ {
|
||||
if similar[j].score > similar[i].score {
|
||||
similar[i], similar[j] = similar[j], similar[i]
|
||||
}
|
||||
}
|
||||
}
|
||||
// Sort by score (descending) using Go's standard sort - O(n log n)
|
||||
sort.Slice(similar, func(i, j int) bool {
|
||||
return similar[i].score > similar[j].score
|
||||
})
|
||||
|
||||
// Create expansion by combining top similar terms with query
|
||||
var expansions []ExpandedQuery
|
||||
@@ -427,16 +425,13 @@ func cosineSimilarity(a, b []float32) float64 {
|
||||
return dot / (sqrt(normA) * sqrt(normB))
|
||||
}
|
||||
|
||||
// sqrt is a simple square root implementation.
|
||||
// sqrt uses the standard math.Sqrt for better performance and accuracy.
|
||||
// Returns 0 for non-positive values (original behavior for compatibility).
|
||||
func sqrt(x float64) float64 {
|
||||
if x <= 0 {
|
||||
return 0
|
||||
}
|
||||
z := x
|
||||
for i := 0; i < 10; i++ {
|
||||
z = (z + x/z) / 2
|
||||
}
|
||||
return z
|
||||
return math.Sqrt(x)
|
||||
}
|
||||
|
||||
// truncate truncates a string to maxLen characters.
|
||||
|
||||
@@ -88,16 +88,16 @@ func (s *ExpanderSuite) TestExpand() {
|
||||
tests := []struct {
|
||||
name string
|
||||
query string
|
||||
expectedIntent QueryIntent
|
||||
minExpansions int
|
||||
hasOriginal bool
|
||||
expectedIntent QueryIntent
|
||||
}{
|
||||
{"question", "how do I implement auth", 1, true, IntentQuestion},
|
||||
{"error", "fix the bug in login", 1, true, IntentError},
|
||||
{"implementation", "implement user handler", 1, true, IntentImplementation},
|
||||
{"architecture", "architecture design", 1, true, IntentArchitecture},
|
||||
{"general", "database connection", 1, true, IntentGeneral},
|
||||
{"empty", "", 0, false, IntentGeneral},
|
||||
{name: "question", query: "how do I implement auth", expectedIntent: IntentQuestion, minExpansions: 1, hasOriginal: true},
|
||||
{name: "error", query: "fix the bug in login", expectedIntent: IntentError, minExpansions: 1, hasOriginal: true},
|
||||
{name: "implementation", query: "implement user handler", expectedIntent: IntentImplementation, minExpansions: 1, hasOriginal: true},
|
||||
{name: "architecture", query: "architecture design", expectedIntent: IntentArchitecture, minExpansions: 1, hasOriginal: true},
|
||||
{name: "general", query: "database connection", expectedIntent: IntentGeneral, minExpansions: 1, hasOriginal: true},
|
||||
{name: "empty", query: "", expectedIntent: IntentGeneral, minExpansions: 0, hasOriginal: false},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
@@ -392,13 +392,13 @@ func TestTruncate(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
input string
|
||||
maxLen int
|
||||
expected string
|
||||
maxLen int
|
||||
}{
|
||||
{"short", "hello", 10, "hello"},
|
||||
{"exact", "hello", 5, "hello"},
|
||||
{"long", "hello world", 5, "hello..."},
|
||||
{"empty", "", 10, ""},
|
||||
{name: "short", input: "hello", expected: "hello", maxLen: 10},
|
||||
{name: "exact", input: "hello", expected: "hello", maxLen: 5},
|
||||
{name: "long", input: "hello world", expected: "hello...", maxLen: 5},
|
||||
{name: "empty", input: "", expected: "", maxLen: 10},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
|
||||
@@ -516,16 +516,16 @@ func TestTruncate_TableDriven(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
input string
|
||||
maxLen int
|
||||
expected string
|
||||
maxLen int
|
||||
}{
|
||||
{"short_string", "hello", 10, "hello"},
|
||||
{"exact_length", "hello", 5, "hello"},
|
||||
{"long_string", "hello world", 5, "hello..."},
|
||||
{"empty_string", "", 10, ""},
|
||||
{"whitespace_only", " ", 10, ""},
|
||||
{"with_leading_space", " hello ", 10, "hello"},
|
||||
{"very_long", "this is a very long string that should be truncated", 20, "this is a very long ..."},
|
||||
{name: "short_string", input: "hello", expected: "hello", maxLen: 10},
|
||||
{name: "exact_length", input: "hello", expected: "hello", maxLen: 5},
|
||||
{name: "long_string", input: "hello world", expected: "hello...", maxLen: 5},
|
||||
{name: "empty_string", input: "", expected: "", maxLen: 10},
|
||||
{name: "whitespace_only", input: " ", expected: "", maxLen: 10},
|
||||
{name: "with_leading_space", input: " hello ", expected: "hello", maxLen: 10},
|
||||
{name: "very_long", input: "this is a very long string that should be truncated", expected: "this is a very long ...", maxLen: 20},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
|
||||
+691
-35
@@ -3,19 +3,146 @@ package search
|
||||
|
||||
import (
|
||||
"context"
|
||||
"hash/fnv"
|
||||
"regexp"
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"github.com/lukaszraczylo/claude-mnemonic/internal/db/gorm"
|
||||
"github.com/lukaszraczylo/claude-mnemonic/internal/vector/sqlitevec"
|
||||
"github.com/lukaszraczylo/claude-mnemonic/pkg/models"
|
||||
"github.com/rs/zerolog/log"
|
||||
"golang.org/x/sync/singleflight"
|
||||
)
|
||||
|
||||
// multiSpaceRegex matches multiple consecutive whitespace characters.
|
||||
// Pre-compiled for performance in normalizeQuery.
|
||||
var multiSpaceRegex = regexp.MustCompile(`\s+`)
|
||||
|
||||
// Search configuration constants.
|
||||
const (
|
||||
// Cache configuration
|
||||
defaultCacheTTL = 30 * time.Second // Short TTL for freshness
|
||||
defaultCacheMaxSize = 200 // Max cached results
|
||||
cacheEvictionPercent = 10 // Evict 10% when cache is full
|
||||
cacheEvictionThreshold = 80 // Start eviction scan at 80% capacity
|
||||
|
||||
// Latency tracking
|
||||
latencyHistogramCap = 1000 // Max latency samples for histogram
|
||||
slowQueryThresholdNs = 100 * 1e6 // 100ms threshold for slow query logging
|
||||
|
||||
// Query frequency tracking
|
||||
maxFrequencyEntries = 1000 // Max queries to track for warming
|
||||
frequencyEvictionBatch = 100 // Remove 10% when frequency map is full
|
||||
staleQueryThreshold = 24 * time.Hour // Remove queries older than 24 hours
|
||||
recentQueryWindow = time.Hour // Only consider queries from last hour for warming
|
||||
|
||||
// Cache warming configuration
|
||||
cacheWarmingInitDelay = 30 * time.Second // Delay before starting warming
|
||||
cacheWarmingInterval = 20 * time.Second // Run warming cycle every 20 seconds
|
||||
frequencyCleanupInterval = 5 * time.Minute // Cleanup stale entries every 5 minutes
|
||||
cacheCleanupInterval = time.Minute // Cleanup expired cache every minute
|
||||
warmingQueryTimeout = 5 * time.Second // Timeout for warming queries
|
||||
warmingBatchSize = 5 // Warm top 5 queries per cycle
|
||||
minRecencyFactor = 0.1 // Minimum recency factor for scoring
|
||||
|
||||
// Default query limits
|
||||
defaultQueryLimit = 20
|
||||
maxQueryLimit = 100
|
||||
defaultOrderBy = "date_desc"
|
||||
|
||||
// Truncation lengths
|
||||
queryLogTruncateLen = 50 // Truncate query in logs
|
||||
titleTruncateLen = 100 // Truncate titles in results
|
||||
warmingLogTruncateLen = 30 // Truncate query in warming logs
|
||||
)
|
||||
|
||||
// SearchMetrics tracks search performance statistics.
|
||||
type SearchMetrics struct {
|
||||
latencyHistogram []int64
|
||||
TotalSearches int64
|
||||
VectorSearches int64
|
||||
FilterSearches int64
|
||||
TotalLatencyNs int64
|
||||
VectorLatencyNs int64
|
||||
FilterLatencyNs int64
|
||||
CacheHits int64
|
||||
CoalescedRequests int64
|
||||
SearchErrors int64
|
||||
histogramMu sync.Mutex
|
||||
}
|
||||
|
||||
// GetStats returns the current search statistics.
|
||||
func (m *SearchMetrics) GetStats() map[string]any {
|
||||
totalSearches := atomic.LoadInt64(&m.TotalSearches)
|
||||
totalLatency := atomic.LoadInt64(&m.TotalLatencyNs)
|
||||
vectorSearches := atomic.LoadInt64(&m.VectorSearches)
|
||||
vectorLatency := atomic.LoadInt64(&m.VectorLatencyNs)
|
||||
filterSearches := atomic.LoadInt64(&m.FilterSearches)
|
||||
filterLatency := atomic.LoadInt64(&m.FilterLatencyNs)
|
||||
|
||||
avgLatencyMs := float64(0)
|
||||
if totalSearches > 0 {
|
||||
avgLatencyMs = float64(totalLatency) / float64(totalSearches) / 1e6
|
||||
}
|
||||
|
||||
avgVectorLatencyMs := float64(0)
|
||||
if vectorSearches > 0 {
|
||||
avgVectorLatencyMs = float64(vectorLatency) / float64(vectorSearches) / 1e6
|
||||
}
|
||||
|
||||
avgFilterLatencyMs := float64(0)
|
||||
if filterSearches > 0 {
|
||||
avgFilterLatencyMs = float64(filterLatency) / float64(filterSearches) / 1e6
|
||||
}
|
||||
|
||||
return map[string]any{
|
||||
"total_searches": totalSearches,
|
||||
"vector_searches": vectorSearches,
|
||||
"filter_searches": filterSearches,
|
||||
"cache_hits": atomic.LoadInt64(&m.CacheHits),
|
||||
"coalesced_requests": atomic.LoadInt64(&m.CoalescedRequests),
|
||||
"search_errors": atomic.LoadInt64(&m.SearchErrors),
|
||||
"avg_latency_ms": avgLatencyMs,
|
||||
"avg_vector_latency_ms": avgVectorLatencyMs,
|
||||
"avg_filter_latency_ms": avgFilterLatencyMs,
|
||||
}
|
||||
}
|
||||
|
||||
// Manager provides unified search across SQLite and sqlite-vec.
|
||||
type Manager struct {
|
||||
ctx context.Context
|
||||
searchGroup singleflight.Group
|
||||
cancel context.CancelFunc
|
||||
vectorClient *sqlitevec.Client
|
||||
metrics *SearchMetrics
|
||||
promptStore *gorm.PromptStore
|
||||
observationStore *gorm.ObservationStore
|
||||
summaryStore *gorm.SummaryStore
|
||||
promptStore *gorm.PromptStore
|
||||
vectorClient *sqlitevec.Client
|
||||
resultCache map[string]*cachedResult
|
||||
queryFrequency map[string]*queryFrequencyInfo
|
||||
cacheTTL time.Duration
|
||||
cacheMaxSize int
|
||||
resultCacheMu sync.RWMutex
|
||||
queryFrequencyMu sync.RWMutex
|
||||
}
|
||||
|
||||
// queryFrequencyInfo tracks how often a query is used.
|
||||
type queryFrequencyInfo struct {
|
||||
lastUsed time.Time
|
||||
lastCached time.Time
|
||||
params SearchParams
|
||||
count int64
|
||||
}
|
||||
|
||||
// cachedResult stores a cached search result with expiry.
|
||||
type cachedResult struct {
|
||||
result *UnifiedSearchResult
|
||||
expiresAt time.Time
|
||||
}
|
||||
|
||||
// NewManager creates a new search manager.
|
||||
@@ -25,65 +152,569 @@ func NewManager(
|
||||
promptStore *gorm.PromptStore,
|
||||
vectorClient *sqlitevec.Client,
|
||||
) *Manager {
|
||||
return &Manager{
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
m := &Manager{
|
||||
observationStore: observationStore,
|
||||
summaryStore: summaryStore,
|
||||
promptStore: promptStore,
|
||||
vectorClient: vectorClient,
|
||||
metrics: &SearchMetrics{latencyHistogram: make([]int64, 0, latencyHistogramCap)},
|
||||
ctx: ctx,
|
||||
cancel: cancel,
|
||||
resultCache: make(map[string]*cachedResult),
|
||||
cacheTTL: defaultCacheTTL,
|
||||
cacheMaxSize: defaultCacheMaxSize,
|
||||
queryFrequency: make(map[string]*queryFrequencyInfo),
|
||||
}
|
||||
// Start cache cleanup goroutine
|
||||
go m.cleanupCacheLoop()
|
||||
// Start cache warming goroutine
|
||||
go m.cacheWarmingLoop()
|
||||
return m
|
||||
}
|
||||
|
||||
// Close stops background goroutines and cleans up resources.
|
||||
func (m *Manager) Close() {
|
||||
if m.cancel != nil {
|
||||
m.cancel()
|
||||
}
|
||||
}
|
||||
|
||||
// cleanupCacheLoop periodically removes expired cache entries.
|
||||
func (m *Manager) cleanupCacheLoop() {
|
||||
ticker := time.NewTicker(cacheCleanupInterval)
|
||||
defer ticker.Stop()
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-m.ctx.Done():
|
||||
return
|
||||
case <-ticker.C:
|
||||
m.cleanupExpiredCache()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// cleanupExpiredCache removes expired entries from the cache.
|
||||
func (m *Manager) cleanupExpiredCache() {
|
||||
m.resultCacheMu.Lock()
|
||||
defer m.resultCacheMu.Unlock()
|
||||
|
||||
now := time.Now()
|
||||
for key, cached := range m.resultCache {
|
||||
if now.After(cached.expiresAt) {
|
||||
delete(m.resultCache, key)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// cacheWarmingLoop periodically warms the cache for frequently used queries.
|
||||
func (m *Manager) cacheWarmingLoop() {
|
||||
// Wait a bit before starting to allow system to stabilize
|
||||
select {
|
||||
case <-m.ctx.Done():
|
||||
return
|
||||
case <-time.After(cacheWarmingInitDelay):
|
||||
}
|
||||
|
||||
warmingTicker := time.NewTicker(cacheWarmingInterval)
|
||||
cleanupTicker := time.NewTicker(frequencyCleanupInterval)
|
||||
defer warmingTicker.Stop()
|
||||
defer cleanupTicker.Stop()
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-m.ctx.Done():
|
||||
return
|
||||
case <-warmingTicker.C:
|
||||
m.warmFrequentQueries()
|
||||
case <-cleanupTicker.C:
|
||||
m.cleanupStaleFrequencyEntries()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// cleanupStaleFrequencyEntries removes query frequency entries older than staleQueryThreshold.
|
||||
// This prevents memory bloat from queries that haven't been used in a long time.
|
||||
func (m *Manager) cleanupStaleFrequencyEntries() {
|
||||
m.queryFrequencyMu.Lock()
|
||||
now := time.Now()
|
||||
var keysToDelete []string
|
||||
for k, v := range m.queryFrequency {
|
||||
if now.Sub(v.lastUsed) > staleQueryThreshold {
|
||||
keysToDelete = append(keysToDelete, k)
|
||||
}
|
||||
}
|
||||
for _, k := range keysToDelete {
|
||||
delete(m.queryFrequency, k)
|
||||
}
|
||||
m.queryFrequencyMu.Unlock()
|
||||
|
||||
if len(keysToDelete) > 0 {
|
||||
log.Debug().Int("removed", len(keysToDelete)).Msg("Cleaned up stale query frequency entries")
|
||||
}
|
||||
}
|
||||
|
||||
// warmFrequentQueries pre-executes frequently used queries to warm the cache.
|
||||
func (m *Manager) warmFrequentQueries() {
|
||||
m.queryFrequencyMu.RLock()
|
||||
// Find top N most frequent queries that aren't recently cached
|
||||
type queryScore struct {
|
||||
info *queryFrequencyInfo
|
||||
key string
|
||||
score float64
|
||||
}
|
||||
candidates := make([]queryScore, 0, len(m.queryFrequency))
|
||||
now := time.Now()
|
||||
|
||||
for key, info := range m.queryFrequency {
|
||||
// Only consider queries used recently
|
||||
if now.Sub(info.lastUsed) > recentQueryWindow {
|
||||
continue
|
||||
}
|
||||
// Only warm if not recently cached (cache about to expire or already expired)
|
||||
timeSinceLastCache := now.Sub(info.lastCached)
|
||||
if timeSinceLastCache < m.cacheTTL/2 {
|
||||
continue
|
||||
}
|
||||
|
||||
// Score: frequency * recency factor
|
||||
recencyFactor := 1.0 - (now.Sub(info.lastUsed).Seconds() / recentQueryWindow.Seconds())
|
||||
if recencyFactor < minRecencyFactor {
|
||||
recencyFactor = minRecencyFactor
|
||||
}
|
||||
score := float64(info.count) * recencyFactor
|
||||
|
||||
candidates = append(candidates, queryScore{key: key, info: info, score: score})
|
||||
}
|
||||
m.queryFrequencyMu.RUnlock()
|
||||
|
||||
// Sort by score descending using O(n log n) algorithm
|
||||
sort.Slice(candidates, func(i, j int) bool {
|
||||
return candidates[i].score > candidates[j].score
|
||||
})
|
||||
|
||||
// Warm top queries
|
||||
warmCount := min(warmingBatchSize, len(candidates))
|
||||
for i := range warmCount {
|
||||
candidate := candidates[i]
|
||||
ctx, cancel := context.WithTimeout(context.Background(), warmingQueryTimeout)
|
||||
result, err := m.executeSearch(ctx, candidate.info.params)
|
||||
cancel()
|
||||
|
||||
if err == nil && result != nil {
|
||||
cacheKey := m.getCacheKey(candidate.info.params)
|
||||
m.putInCache(cacheKey, result)
|
||||
|
||||
// Update last cached time
|
||||
m.queryFrequencyMu.Lock()
|
||||
if info, ok := m.queryFrequency[candidate.key]; ok {
|
||||
info.lastCached = time.Now()
|
||||
}
|
||||
m.queryFrequencyMu.Unlock()
|
||||
|
||||
log.Debug().
|
||||
Str("query", truncate(candidate.info.params.Query, warmingLogTruncateLen)).
|
||||
Float64("score", candidate.score).
|
||||
Msg("Cache warmed for frequent query")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// trackQueryFrequency records query usage for cache warming decisions.
|
||||
func (m *Manager) trackQueryFrequency(params SearchParams) {
|
||||
key := m.getCacheKey(params)
|
||||
|
||||
m.queryFrequencyMu.Lock()
|
||||
|
||||
if info, ok := m.queryFrequency[key]; ok {
|
||||
info.count++
|
||||
info.lastUsed = time.Now()
|
||||
m.queryFrequencyMu.Unlock()
|
||||
return // Fast path: no eviction needed
|
||||
}
|
||||
|
||||
m.queryFrequency[key] = &queryFrequencyInfo{
|
||||
params: params,
|
||||
count: 1,
|
||||
lastUsed: time.Now(),
|
||||
}
|
||||
|
||||
// Limit frequency map size to prevent memory bloat
|
||||
mapLen := len(m.queryFrequency)
|
||||
if mapLen <= maxFrequencyEntries {
|
||||
m.queryFrequencyMu.Unlock()
|
||||
return // Fast path: no eviction needed
|
||||
}
|
||||
|
||||
// Collect keys and times for eviction (still under lock, but fast)
|
||||
type entry struct {
|
||||
lastUsed time.Time
|
||||
key string
|
||||
}
|
||||
entries := make([]entry, 0, mapLen)
|
||||
for k, v := range m.queryFrequency {
|
||||
entries = append(entries, entry{key: k, lastUsed: v.lastUsed})
|
||||
}
|
||||
m.queryFrequencyMu.Unlock()
|
||||
|
||||
// Sort outside lock to reduce contention (O(n log n))
|
||||
sort.Slice(entries, func(i, j int) bool {
|
||||
return entries[i].lastUsed.Before(entries[j].lastUsed)
|
||||
})
|
||||
|
||||
// Collect keys to delete
|
||||
deleteCount := min(frequencyEvictionBatch, len(entries))
|
||||
keysToDelete := make([]string, deleteCount)
|
||||
for i := range deleteCount {
|
||||
keysToDelete[i] = entries[i].key
|
||||
}
|
||||
|
||||
// Re-acquire lock only for deletion (brief critical section)
|
||||
m.queryFrequencyMu.Lock()
|
||||
for _, k := range keysToDelete {
|
||||
delete(m.queryFrequency, k)
|
||||
}
|
||||
m.queryFrequencyMu.Unlock()
|
||||
}
|
||||
|
||||
// RecentQuery represents a recently executed search query.
|
||||
type RecentQuery struct {
|
||||
LastUsed time.Time `json:"last_used"`
|
||||
Query string `json:"query"`
|
||||
Project string `json:"project,omitempty"`
|
||||
Type string `json:"type,omitempty"`
|
||||
Count int64 `json:"count"`
|
||||
}
|
||||
|
||||
// GetRecentQueries returns the most recent search queries, sorted by last used time.
|
||||
func (m *Manager) GetRecentQueries(limit int) []RecentQuery {
|
||||
if limit <= 0 {
|
||||
limit = defaultQueryLimit
|
||||
}
|
||||
if limit > maxQueryLimit {
|
||||
limit = maxQueryLimit
|
||||
}
|
||||
|
||||
m.queryFrequencyMu.RLock()
|
||||
defer m.queryFrequencyMu.RUnlock()
|
||||
|
||||
// Collect all queries
|
||||
queries := make([]RecentQuery, 0, len(m.queryFrequency))
|
||||
for _, info := range m.queryFrequency {
|
||||
queries = append(queries, RecentQuery{
|
||||
Query: info.params.Query,
|
||||
Project: info.params.Project,
|
||||
Type: info.params.Type,
|
||||
Count: info.count,
|
||||
LastUsed: info.lastUsed,
|
||||
})
|
||||
}
|
||||
|
||||
// Sort by last used (most recent first)
|
||||
sort.Slice(queries, func(i, j int) bool {
|
||||
return queries[i].LastUsed.After(queries[j].LastUsed)
|
||||
})
|
||||
|
||||
// Limit results
|
||||
if len(queries) > limit {
|
||||
queries = queries[:limit]
|
||||
}
|
||||
|
||||
return queries
|
||||
}
|
||||
|
||||
// GetFrequentQueries returns the most frequently used search queries.
|
||||
func (m *Manager) GetFrequentQueries(limit int) []RecentQuery {
|
||||
if limit <= 0 {
|
||||
limit = defaultQueryLimit
|
||||
}
|
||||
if limit > maxQueryLimit {
|
||||
limit = maxQueryLimit
|
||||
}
|
||||
|
||||
m.queryFrequencyMu.RLock()
|
||||
defer m.queryFrequencyMu.RUnlock()
|
||||
|
||||
// Only include queries used recently
|
||||
now := time.Now()
|
||||
queries := make([]RecentQuery, 0, len(m.queryFrequency))
|
||||
for _, info := range m.queryFrequency {
|
||||
if now.Sub(info.lastUsed) > recentQueryWindow {
|
||||
continue
|
||||
}
|
||||
queries = append(queries, RecentQuery{
|
||||
Query: info.params.Query,
|
||||
Project: info.params.Project,
|
||||
Type: info.params.Type,
|
||||
Count: info.count,
|
||||
LastUsed: info.lastUsed,
|
||||
})
|
||||
}
|
||||
|
||||
// Sort by count (highest first)
|
||||
sort.Slice(queries, func(i, j int) bool {
|
||||
return queries[i].Count > queries[j].Count
|
||||
})
|
||||
|
||||
// Limit results
|
||||
if len(queries) > limit {
|
||||
queries = queries[:limit]
|
||||
}
|
||||
|
||||
return queries
|
||||
}
|
||||
|
||||
// normalizeQuery normalizes a search query for consistent cache keys.
|
||||
// Converts to lowercase, trims whitespace, and collapses multiple spaces.
|
||||
// Uses pre-compiled regex for O(n) performance instead of O(n*m) loop.
|
||||
func normalizeQuery(query string) string {
|
||||
// Lowercase for case-insensitive matching
|
||||
query = strings.ToLower(query)
|
||||
// Collapse multiple whitespace into single space using pre-compiled regex
|
||||
query = multiSpaceRegex.ReplaceAllString(query, " ")
|
||||
// Trim leading/trailing whitespace (after collapsing)
|
||||
return strings.TrimSpace(query)
|
||||
}
|
||||
|
||||
// getCacheKey generates a cache key from search params.
|
||||
// Uses direct string concatenation instead of JSON marshal for better performance.
|
||||
// Queries are normalized for consistent cache hits across whitespace variations.
|
||||
func (m *Manager) getCacheKey(params SearchParams) string {
|
||||
// Normalize query for consistent cache keys
|
||||
normalizedQuery := normalizeQuery(params.Query)
|
||||
|
||||
// Hash directly without intermediate string allocation.
|
||||
// FNV-64a is fast and collision-safe for cache keys.
|
||||
h := fnv.New64a()
|
||||
|
||||
// Write each field directly to the hasher with separators
|
||||
h.Write([]byte(normalizedQuery))
|
||||
h.Write([]byte{'|'})
|
||||
h.Write([]byte(params.Type))
|
||||
h.Write([]byte{'|'})
|
||||
h.Write([]byte(params.Project))
|
||||
h.Write([]byte{'|'})
|
||||
h.Write([]byte(params.ObsType))
|
||||
h.Write([]byte{'|'})
|
||||
h.Write([]byte(params.Concepts))
|
||||
h.Write([]byte{'|'})
|
||||
h.Write([]byte(params.Files))
|
||||
h.Write([]byte{'|'})
|
||||
h.Write([]byte(strconv.FormatInt(params.DateStart, 10)))
|
||||
h.Write([]byte{'|'})
|
||||
h.Write([]byte(strconv.FormatInt(params.DateEnd, 10)))
|
||||
h.Write([]byte{'|'})
|
||||
h.Write([]byte(params.OrderBy))
|
||||
h.Write([]byte{'|'})
|
||||
h.Write([]byte(strconv.Itoa(params.Limit)))
|
||||
h.Write([]byte{'|'})
|
||||
h.Write([]byte(strconv.Itoa(params.Offset)))
|
||||
h.Write([]byte{'|'})
|
||||
h.Write([]byte(params.Format))
|
||||
h.Write([]byte{'|'})
|
||||
h.Write([]byte(params.Scope))
|
||||
h.Write([]byte{'|'})
|
||||
if params.IncludeGlobal {
|
||||
h.Write([]byte{'1'})
|
||||
} else {
|
||||
h.Write([]byte{'0'})
|
||||
}
|
||||
h.Write([]byte{'|'})
|
||||
if params.ExcludeSuperseded {
|
||||
h.Write([]byte{'1'})
|
||||
} else {
|
||||
h.Write([]byte{'0'})
|
||||
}
|
||||
|
||||
return strconv.FormatUint(h.Sum64(), 36) // Base36 for compact representation
|
||||
}
|
||||
|
||||
// getFromCache retrieves a result from cache if valid.
|
||||
func (m *Manager) getFromCache(key string) (*UnifiedSearchResult, bool) {
|
||||
m.resultCacheMu.RLock()
|
||||
defer m.resultCacheMu.RUnlock()
|
||||
|
||||
if cached, ok := m.resultCache[key]; ok {
|
||||
if time.Now().Before(cached.expiresAt) {
|
||||
atomic.AddInt64(&m.metrics.CacheHits, 1)
|
||||
return cached.result, true
|
||||
}
|
||||
}
|
||||
return nil, false
|
||||
}
|
||||
|
||||
// putInCache stores a result in the cache.
|
||||
// Optimized to skip expensive scans when cache is below capacity threshold.
|
||||
func (m *Manager) putInCache(key string, result *UnifiedSearchResult) {
|
||||
m.resultCacheMu.Lock()
|
||||
defer m.resultCacheMu.Unlock()
|
||||
|
||||
now := time.Now()
|
||||
cacheLen := len(m.resultCache)
|
||||
|
||||
// Only scan for expired entries when at threshold+ capacity (amortized cleanup)
|
||||
evictionThreshold := (m.cacheMaxSize * cacheEvictionThreshold) / 100
|
||||
if cacheLen >= evictionThreshold {
|
||||
// Evict expired entries
|
||||
for k, v := range m.resultCache {
|
||||
if now.After(v.expiresAt) {
|
||||
delete(m.resultCache, k)
|
||||
}
|
||||
}
|
||||
cacheLen = len(m.resultCache) // Update after eviction
|
||||
}
|
||||
|
||||
// If still at capacity after removing expired, use simple FIFO-style eviction
|
||||
// Go map iteration order is random, which provides good cache behavior
|
||||
if cacheLen >= m.cacheMaxSize {
|
||||
// Evict percentage using random-order iteration (O(n) single pass)
|
||||
evictCount := max(m.cacheMaxSize*cacheEvictionPercent/100, 1)
|
||||
evicted := 0
|
||||
for k := range m.resultCache {
|
||||
delete(m.resultCache, k)
|
||||
evicted++
|
||||
if evicted >= evictCount {
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
m.resultCache[key] = &cachedResult{
|
||||
result: result,
|
||||
expiresAt: now.Add(m.cacheTTL),
|
||||
}
|
||||
}
|
||||
|
||||
// Metrics returns the search metrics for monitoring.
|
||||
func (m *Manager) Metrics() *SearchMetrics {
|
||||
return m.metrics
|
||||
}
|
||||
|
||||
// CacheStats returns current cache statistics.
|
||||
func (m *Manager) CacheStats() map[string]any {
|
||||
m.resultCacheMu.RLock()
|
||||
cacheSize := len(m.resultCache)
|
||||
m.resultCacheMu.RUnlock()
|
||||
|
||||
return map[string]any{
|
||||
"size": cacheSize,
|
||||
"max_size": m.cacheMaxSize,
|
||||
"ttl_sec": m.cacheTTL.Seconds(),
|
||||
}
|
||||
}
|
||||
|
||||
// ClearCache clears the result cache. Useful for testing or after data changes.
|
||||
func (m *Manager) ClearCache() {
|
||||
m.resultCacheMu.Lock()
|
||||
m.resultCache = make(map[string]*cachedResult)
|
||||
m.resultCacheMu.Unlock()
|
||||
}
|
||||
|
||||
// SearchParams contains parameters for unified search.
|
||||
type SearchParams struct {
|
||||
Query string
|
||||
Type string // "observations", "sessions", "prompts", or empty for all
|
||||
Format string
|
||||
Type string
|
||||
Project string
|
||||
ObsType string // Observation type filter
|
||||
ObsType string
|
||||
Concepts string
|
||||
Files string
|
||||
Query string
|
||||
Scope string
|
||||
OrderBy string
|
||||
DateStart int64
|
||||
DateEnd int64
|
||||
OrderBy string // "relevance", "date_desc", "date_asc"
|
||||
Limit int
|
||||
Offset int
|
||||
Format string // "index" or "full"
|
||||
Scope string // "project", "global", or empty for project+global
|
||||
IncludeGlobal bool // If true, include global observations along with project-scoped
|
||||
ExcludeSuperseded bool // If true, exclude observations that have been superseded
|
||||
Limit int
|
||||
DateEnd int64
|
||||
IncludeGlobal bool
|
||||
ExcludeSuperseded bool
|
||||
}
|
||||
|
||||
// SearchResult represents a unified search result.
|
||||
// Field order optimized for memory alignment (fieldalignment).
|
||||
type SearchResult struct {
|
||||
Type string `json:"type"` // "observation", "session", "prompt"
|
||||
ID int64 `json:"id"`
|
||||
Title string `json:"title,omitempty"`
|
||||
Content string `json:"content,omitempty"`
|
||||
Project string `json:"project"`
|
||||
Scope string `json:"scope,omitempty"` // "project" or "global"
|
||||
CreatedAt int64 `json:"created_at_epoch"`
|
||||
Score float64 `json:"score,omitempty"`
|
||||
Metadata map[string]interface{} `json:"metadata,omitempty"`
|
||||
Metadata map[string]any `json:"metadata,omitempty"`
|
||||
Type string `json:"type"`
|
||||
Title string `json:"title,omitempty"`
|
||||
Content string `json:"content,omitempty"`
|
||||
Project string `json:"project"`
|
||||
Scope string `json:"scope,omitempty"`
|
||||
ID int64 `json:"id"`
|
||||
CreatedAt int64 `json:"created_at_epoch"`
|
||||
Score float64 `json:"score,omitempty"`
|
||||
}
|
||||
|
||||
// UnifiedSearchResult contains the combined search results.
|
||||
type UnifiedSearchResult struct {
|
||||
Query string `json:"query,omitempty"`
|
||||
Results []SearchResult `json:"results"`
|
||||
TotalCount int `json:"total_count"`
|
||||
Query string `json:"query,omitempty"`
|
||||
}
|
||||
|
||||
// UnifiedSearch performs a unified search across all document types.
|
||||
// Uses caching and request coalescing for optimal performance.
|
||||
func (m *Manager) UnifiedSearch(ctx context.Context, params SearchParams) (*UnifiedSearchResult, error) {
|
||||
start := time.Now()
|
||||
defer func() {
|
||||
latency := time.Since(start).Nanoseconds()
|
||||
atomic.AddInt64(&m.metrics.TotalSearches, 1)
|
||||
atomic.AddInt64(&m.metrics.TotalLatencyNs, latency)
|
||||
|
||||
// Sample latency for histogram (reservoir sampling)
|
||||
m.metrics.histogramMu.Lock()
|
||||
if len(m.metrics.latencyHistogram) < latencyHistogramCap {
|
||||
m.metrics.latencyHistogram = append(m.metrics.latencyHistogram, latency)
|
||||
}
|
||||
m.metrics.histogramMu.Unlock()
|
||||
|
||||
// Log slow queries
|
||||
if latency > slowQueryThresholdNs {
|
||||
log.Warn().
|
||||
Str("query", truncate(params.Query, queryLogTruncateLen)).
|
||||
Dur("latency", time.Duration(latency)).
|
||||
Str("type", params.Type).
|
||||
Msg("Slow search query")
|
||||
}
|
||||
}()
|
||||
|
||||
if params.Limit <= 0 {
|
||||
params.Limit = 20
|
||||
params.Limit = defaultQueryLimit
|
||||
}
|
||||
if params.Limit > 100 {
|
||||
params.Limit = 100
|
||||
if params.Limit > maxQueryLimit {
|
||||
params.Limit = maxQueryLimit
|
||||
}
|
||||
if params.OrderBy == "" {
|
||||
params.OrderBy = "date_desc"
|
||||
params.OrderBy = defaultOrderBy
|
||||
}
|
||||
|
||||
// Check cache first
|
||||
cacheKey := m.getCacheKey(params)
|
||||
if cached, ok := m.getFromCache(cacheKey); ok {
|
||||
return cached, nil
|
||||
}
|
||||
|
||||
// Use singleflight to coalesce concurrent identical requests
|
||||
result, err, _ := m.searchGroup.Do(cacheKey, func() (any, error) {
|
||||
return m.executeSearch(ctx, params)
|
||||
})
|
||||
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
searchResult := result.(*UnifiedSearchResult)
|
||||
|
||||
// Cache the result
|
||||
m.putInCache(cacheKey, searchResult)
|
||||
|
||||
// Track query frequency for cache warming
|
||||
m.trackQueryFrequency(params)
|
||||
|
||||
return searchResult, nil
|
||||
}
|
||||
|
||||
// executeSearch performs the actual search without caching/coalescing.
|
||||
func (m *Manager) executeSearch(ctx context.Context, params SearchParams) (*UnifiedSearchResult, error) {
|
||||
// If query is provided and vector client is available, use vector search
|
||||
if params.Query != "" && m.vectorClient != nil && m.vectorClient.IsConnected() {
|
||||
return m.vectorSearch(ctx, params)
|
||||
@@ -95,6 +726,13 @@ func (m *Manager) UnifiedSearch(ctx context.Context, params SearchParams) (*Unif
|
||||
|
||||
// vectorSearch performs semantic search via sqlite-vec.
|
||||
func (m *Manager) vectorSearch(ctx context.Context, params SearchParams) (*UnifiedSearchResult, error) {
|
||||
start := time.Now()
|
||||
defer func() {
|
||||
latency := time.Since(start).Nanoseconds()
|
||||
atomic.AddInt64(&m.metrics.VectorSearches, 1)
|
||||
atomic.AddInt64(&m.metrics.VectorLatencyNs, latency)
|
||||
}()
|
||||
|
||||
// Build where filter based on search type
|
||||
var docType sqlitevec.DocType
|
||||
switch params.Type {
|
||||
@@ -110,6 +748,7 @@ func (m *Manager) vectorSearch(ctx context.Context, params SearchParams) (*Unifi
|
||||
// Query sqlite-vec
|
||||
vectorResults, err := m.vectorClient.Query(ctx, params.Query, params.Limit*2, where)
|
||||
if err != nil {
|
||||
atomic.AddInt64(&m.metrics.SearchErrors, 1)
|
||||
// Fall back to filter search on error
|
||||
return m.filterSearch(ctx, params)
|
||||
}
|
||||
@@ -125,7 +764,9 @@ func (m *Manager) vectorSearch(ctx context.Context, params SearchParams) (*Unifi
|
||||
|
||||
if len(obsIDs) > 0 && (params.Type == "" || params.Type == "observations") {
|
||||
obs, err := m.observationStore.GetObservationsByIDs(ctx, obsIDs, params.OrderBy, 0)
|
||||
if err == nil {
|
||||
if err != nil {
|
||||
log.Warn().Err(err).Int("count", len(obsIDs)).Msg("Failed to fetch observations by IDs in vector search")
|
||||
} else {
|
||||
for _, o := range obs {
|
||||
// Skip superseded observations when requested
|
||||
if params.ExcludeSuperseded && o.IsSuperseded {
|
||||
@@ -138,7 +779,9 @@ func (m *Manager) vectorSearch(ctx context.Context, params SearchParams) (*Unifi
|
||||
|
||||
if len(summaryIDs) > 0 && (params.Type == "" || params.Type == "sessions") {
|
||||
summaries, err := m.summaryStore.GetSummariesByIDs(ctx, summaryIDs, params.OrderBy, 0)
|
||||
if err == nil {
|
||||
if err != nil {
|
||||
log.Warn().Err(err).Int("count", len(summaryIDs)).Msg("Failed to fetch summaries by IDs in vector search")
|
||||
} else {
|
||||
for _, s := range summaries {
|
||||
results = append(results, m.summaryToResult(s, params.Format))
|
||||
}
|
||||
@@ -147,7 +790,9 @@ func (m *Manager) vectorSearch(ctx context.Context, params SearchParams) (*Unifi
|
||||
|
||||
if len(promptIDs) > 0 && (params.Type == "" || params.Type == "prompts") {
|
||||
prompts, err := m.promptStore.GetPromptsByIDs(ctx, promptIDs, params.OrderBy, 0)
|
||||
if err == nil {
|
||||
if err != nil {
|
||||
log.Warn().Err(err).Int("count", len(promptIDs)).Msg("Failed to fetch prompts by IDs in vector search")
|
||||
} else {
|
||||
for _, p := range prompts {
|
||||
results = append(results, m.promptToResult(p, params.Format))
|
||||
}
|
||||
@@ -168,6 +813,13 @@ func (m *Manager) vectorSearch(ctx context.Context, params SearchParams) (*Unifi
|
||||
|
||||
// filterSearch performs structured filter search via SQLite.
|
||||
func (m *Manager) filterSearch(ctx context.Context, params SearchParams) (*UnifiedSearchResult, error) {
|
||||
start := time.Now()
|
||||
defer func() {
|
||||
latency := time.Since(start).Nanoseconds()
|
||||
atomic.AddInt64(&m.metrics.FilterSearches, 1)
|
||||
atomic.AddInt64(&m.metrics.FilterLatencyNs, latency)
|
||||
}()
|
||||
|
||||
var results []SearchResult
|
||||
|
||||
// Search observations
|
||||
@@ -182,7 +834,9 @@ func (m *Manager) filterSearch(ctx context.Context, params SearchParams) (*Unifi
|
||||
obs, err = m.observationStore.GetRecentObservations(ctx, params.Project, params.Limit)
|
||||
}
|
||||
|
||||
if err == nil {
|
||||
if err != nil {
|
||||
log.Warn().Err(err).Str("project", params.Project).Msg("Failed to fetch observations in filter search")
|
||||
} else {
|
||||
for _, o := range obs {
|
||||
results = append(results, m.observationToResult(o, params.Format))
|
||||
}
|
||||
@@ -192,7 +846,9 @@ func (m *Manager) filterSearch(ctx context.Context, params SearchParams) (*Unifi
|
||||
// Search summaries
|
||||
if params.Type == "" || params.Type == "sessions" {
|
||||
summaries, err := m.summaryStore.GetRecentSummaries(ctx, params.Project, params.Limit)
|
||||
if err == nil {
|
||||
if err != nil {
|
||||
log.Warn().Err(err).Str("project", params.Project).Msg("Failed to fetch summaries in filter search")
|
||||
} else {
|
||||
for _, s := range summaries {
|
||||
results = append(results, m.summaryToResult(s, params.Format))
|
||||
}
|
||||
@@ -249,7 +905,7 @@ func (m *Manager) observationToResult(obs *models.Observation, format string) Se
|
||||
Project: obs.Project,
|
||||
Scope: string(obs.Scope),
|
||||
CreatedAt: obs.CreatedAtEpoch,
|
||||
Metadata: map[string]interface{}{
|
||||
Metadata: map[string]any{
|
||||
"obs_type": string(obs.Type),
|
||||
"scope": string(obs.Scope),
|
||||
},
|
||||
@@ -275,7 +931,7 @@ func (m *Manager) summaryToResult(summary *models.SessionSummary, format string)
|
||||
}
|
||||
|
||||
if summary.Request.Valid {
|
||||
result.Title = truncate(summary.Request.String, 100)
|
||||
result.Title = truncate(summary.Request.String, titleTruncateLen)
|
||||
}
|
||||
|
||||
if format == "full" && summary.Learned.Valid {
|
||||
@@ -293,7 +949,7 @@ func (m *Manager) promptToResult(prompt *models.UserPromptWithSession, format st
|
||||
CreatedAt: prompt.CreatedAtEpoch,
|
||||
}
|
||||
|
||||
result.Title = truncate(prompt.PromptText, 100)
|
||||
result.Title = truncate(prompt.PromptText, titleTruncateLen)
|
||||
|
||||
if format == "full" {
|
||||
result.Content = prompt.PromptText
|
||||
|
||||
@@ -94,8 +94,8 @@ func TestTruncate(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
input string
|
||||
maxLen int
|
||||
expected string
|
||||
maxLen int
|
||||
}{
|
||||
{
|
||||
name: "short string no truncation",
|
||||
@@ -148,8 +148,8 @@ func TestObservationToResult(t *testing.T) {
|
||||
m := NewManager(nil, nil, nil, nil)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
obs *models.Observation
|
||||
name string
|
||||
format string
|
||||
expected SearchResult
|
||||
}{
|
||||
@@ -240,8 +240,8 @@ func TestSummaryToResult(t *testing.T) {
|
||||
m := NewManager(nil, nil, nil, nil)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
summary *models.SessionSummary
|
||||
name string
|
||||
format string
|
||||
expected SearchResult
|
||||
}{
|
||||
@@ -322,8 +322,8 @@ func TestPromptToResult(t *testing.T) {
|
||||
m := NewManager(nil, nil, nil, nil)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
prompt *models.UserPromptWithSession
|
||||
name string
|
||||
format string
|
||||
expected SearchResult
|
||||
}{
|
||||
@@ -406,9 +406,9 @@ func TestPromptToResult(t *testing.T) {
|
||||
func TestSearchParamsValidation(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
expectedOrder string
|
||||
params SearchParams
|
||||
expectedLimit int
|
||||
expectedOrder string
|
||||
}{
|
||||
{
|
||||
name: "default limit applied",
|
||||
@@ -731,16 +731,16 @@ func TestPromptToResultFormats(t *testing.T) {
|
||||
func TestSearchParamsDefaults(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
initialLimit int
|
||||
initialOrder string
|
||||
expectedLimit int
|
||||
expectedOrder string
|
||||
initialLimit int
|
||||
expectedLimit int
|
||||
}{
|
||||
{"zero_limit", 0, "", 20, "date_desc"},
|
||||
{"negative_limit", -5, "", 20, "date_desc"},
|
||||
{"over_100_limit", 150, "", 100, "date_desc"},
|
||||
{"valid_limit_50", 50, "relevance", 50, "relevance"},
|
||||
{"custom_order", 30, "date_asc", 30, "date_asc"},
|
||||
{name: "zero_limit", initialOrder: "", expectedOrder: "date_desc", initialLimit: 0, expectedLimit: 20},
|
||||
{name: "negative_limit", initialOrder: "", expectedOrder: "date_desc", initialLimit: -5, expectedLimit: 20},
|
||||
{name: "over_100_limit", initialOrder: "", expectedOrder: "date_desc", initialLimit: 150, expectedLimit: 100},
|
||||
{name: "valid_limit_50", initialOrder: "relevance", expectedOrder: "relevance", initialLimit: 50, expectedLimit: 50},
|
||||
{name: "custom_order", initialOrder: "date_asc", expectedOrder: "date_asc", initialLimit: 30, expectedLimit: 30},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
@@ -774,18 +774,18 @@ func TestTruncateEdgeCases(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
input string
|
||||
maxLen int
|
||||
expected string
|
||||
maxLen int
|
||||
}{
|
||||
// Unicode strings - uses byte length so ensure maxLen accommodates full string
|
||||
{"unicode_string_no_truncate", "日本語テスト", 20, "日本語テスト"},
|
||||
{"mixed_unicode_no_truncate", "Hello世界", 15, "Hello世界"},
|
||||
{name: "unicode_string_no_truncate", input: "日本語テスト", expected: "日本語テスト", maxLen: 20},
|
||||
{name: "mixed_unicode_no_truncate", input: "Hello世界", expected: "Hello世界", maxLen: 15},
|
||||
// ASCII truncation
|
||||
{"ascii_truncate", "Hello World", 5, "Hello..."},
|
||||
{"only_whitespace", " ", 10, ""},
|
||||
{"tabs_and_newlines", "\t\n \t", 10, ""},
|
||||
{"newlines_with_content", "\n\nhello\n\n", 10, "hello"},
|
||||
{"zero_max_len", "hello", 0, "..."},
|
||||
{name: "ascii_truncate", input: "Hello World", expected: "Hello...", maxLen: 5},
|
||||
{name: "only_whitespace", input: " ", expected: "", maxLen: 10},
|
||||
{name: "tabs_and_newlines", input: "\t\n \t", expected: "", maxLen: 10},
|
||||
{name: "newlines_with_content", input: "\n\nhello\n\n", expected: "hello", maxLen: 10},
|
||||
{name: "zero_max_len", input: "hello", expected: "...", maxLen: 0},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
|
||||
+13
-14
@@ -33,11 +33,11 @@ const (
|
||||
|
||||
// Release represents a GitHub release.
|
||||
type Release struct {
|
||||
PublishedAt time.Time `json:"published_at"`
|
||||
TagName string `json:"tag_name"`
|
||||
Name string `json:"name"`
|
||||
PublishedAt time.Time `json:"published_at"`
|
||||
Assets []Asset `json:"assets"`
|
||||
Body string `json:"body"`
|
||||
Assets []Asset `json:"assets"`
|
||||
}
|
||||
|
||||
// Asset represents a release asset.
|
||||
@@ -49,15 +49,15 @@ type Asset struct {
|
||||
|
||||
// UpdateInfo contains information about an available update.
|
||||
type UpdateInfo struct {
|
||||
Available bool `json:"available"`
|
||||
PublishedAt time.Time `json:"published_at,omitempty"`
|
||||
CurrentVersion string `json:"current_version"`
|
||||
LatestVersion string `json:"latest_version"`
|
||||
ReleaseNotes string `json:"release_notes,omitempty"`
|
||||
PublishedAt time.Time `json:"published_at,omitempty"`
|
||||
DownloadURL string `json:"download_url,omitempty"`
|
||||
ChecksumsURL string `json:"checksums_url,omitempty"`
|
||||
BundleURL string `json:"bundle_url,omitempty"` // Sigstore bundle (.sigstore.json)
|
||||
BundleURL string `json:"bundle_url,omitempty"`
|
||||
ManualUpdateCommand string `json:"manual_update_command,omitempty"`
|
||||
Available bool `json:"available"`
|
||||
}
|
||||
|
||||
// InstallScriptURL is the URL to the remote installation script.
|
||||
@@ -74,23 +74,22 @@ func GetManualUpdateCommand(version string) string {
|
||||
|
||||
// UpdateStatus represents the current update status.
|
||||
type UpdateStatus struct {
|
||||
State string `json:"state"` // "idle", "checking", "downloading", "verifying", "applying", "done", "error"
|
||||
Progress float64 `json:"progress"`
|
||||
State string `json:"state"`
|
||||
Message string `json:"message"`
|
||||
Error string `json:"error,omitempty"`
|
||||
ManualUpdateCommand string `json:"manual_update_command,omitempty"` // Shown when update fails
|
||||
ManualUpdateCommand string `json:"manual_update_command,omitempty"`
|
||||
Progress float64 `json:"progress"`
|
||||
}
|
||||
|
||||
// Updater handles self-updates.
|
||||
type Updater struct {
|
||||
lastCheck time.Time
|
||||
httpClient *http.Client
|
||||
cachedUpdate *UpdateInfo
|
||||
currentVersion string
|
||||
installDir string
|
||||
httpClient *http.Client
|
||||
|
||||
mu sync.RWMutex
|
||||
status UpdateStatus
|
||||
lastCheck time.Time
|
||||
cachedUpdate *UpdateInfo
|
||||
status UpdateStatus
|
||||
mu sync.RWMutex
|
||||
}
|
||||
|
||||
// New creates a new Updater.
|
||||
|
||||
@@ -0,0 +1,309 @@
|
||||
package hybrid
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/lukaszraczylo/claude-mnemonic/internal/vector/sqlitevec"
|
||||
"github.com/rs/zerolog/log"
|
||||
)
|
||||
|
||||
// AutoTuner dynamically adjusts hub threshold based on query performance
|
||||
type AutoTuner struct {
|
||||
ctx context.Context
|
||||
client *Client
|
||||
cancel context.CancelFunc
|
||||
latencies []time.Duration
|
||||
wg sync.WaitGroup
|
||||
queries int64
|
||||
targetLatency time.Duration
|
||||
adjustPeriod time.Duration
|
||||
minThreshold int
|
||||
maxThreshold int
|
||||
adjustments int
|
||||
latenciesMu sync.Mutex
|
||||
}
|
||||
|
||||
// AutoTunerConfig configures the auto-tuner
|
||||
type AutoTunerConfig struct {
|
||||
TargetLatency time.Duration // Target p95 latency (default: 50ms)
|
||||
MinThreshold int // Min hub threshold (default: 2)
|
||||
MaxThreshold int // Max hub threshold (default: 20)
|
||||
AdjustPeriod time.Duration // Adjustment frequency (default: 5min)
|
||||
}
|
||||
|
||||
// DefaultAutoTunerConfig returns sensible defaults
|
||||
func DefaultAutoTunerConfig() AutoTunerConfig {
|
||||
return AutoTunerConfig{
|
||||
TargetLatency: 50 * time.Millisecond,
|
||||
MinThreshold: 2,
|
||||
MaxThreshold: 20,
|
||||
AdjustPeriod: 5 * time.Minute,
|
||||
}
|
||||
}
|
||||
|
||||
// NewAutoTuner creates a new auto-tuner for the hybrid client
|
||||
func NewAutoTuner(client *Client, cfg AutoTunerConfig) *AutoTuner {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
|
||||
tuner := &AutoTuner{
|
||||
client: client,
|
||||
targetLatency: cfg.TargetLatency,
|
||||
minThreshold: cfg.MinThreshold,
|
||||
maxThreshold: cfg.MaxThreshold,
|
||||
adjustPeriod: cfg.AdjustPeriod,
|
||||
latencies: make([]time.Duration, 0, 1000),
|
||||
ctx: ctx,
|
||||
cancel: cancel,
|
||||
}
|
||||
|
||||
return tuner
|
||||
}
|
||||
|
||||
// Start begins auto-tuning in the background
|
||||
func (a *AutoTuner) Start() {
|
||||
a.wg.Add(1)
|
||||
go a.tuningLoop()
|
||||
|
||||
log.Info().
|
||||
Dur("target_latency", a.targetLatency).
|
||||
Int("min_threshold", a.minThreshold).
|
||||
Int("max_threshold", a.maxThreshold).
|
||||
Dur("adjust_period", a.adjustPeriod).
|
||||
Msg("Auto-tuner started")
|
||||
}
|
||||
|
||||
// Stop stops the auto-tuner
|
||||
func (a *AutoTuner) Stop() {
|
||||
a.cancel()
|
||||
a.wg.Wait()
|
||||
log.Info().Msg("Auto-tuner stopped")
|
||||
}
|
||||
|
||||
// RecordQuery records a query latency for analysis
|
||||
func (a *AutoTuner) RecordQuery(latency time.Duration) {
|
||||
a.latenciesMu.Lock()
|
||||
defer a.latenciesMu.Unlock()
|
||||
|
||||
a.queries++
|
||||
a.latencies = append(a.latencies, latency)
|
||||
|
||||
// Keep only recent queries (last 1000)
|
||||
if len(a.latencies) > 1000 {
|
||||
a.latencies = a.latencies[len(a.latencies)-1000:]
|
||||
}
|
||||
}
|
||||
|
||||
// tuningLoop periodically adjusts hub threshold
|
||||
func (a *AutoTuner) tuningLoop() {
|
||||
defer a.wg.Done()
|
||||
|
||||
ticker := time.NewTicker(a.adjustPeriod)
|
||||
defer ticker.Stop()
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-a.ctx.Done():
|
||||
return
|
||||
|
||||
case <-ticker.C:
|
||||
a.adjustThreshold()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// adjustThreshold analyzes recent queries and adjusts hub threshold
|
||||
func (a *AutoTuner) adjustThreshold() {
|
||||
a.latenciesMu.Lock()
|
||||
defer a.latenciesMu.Unlock()
|
||||
|
||||
if len(a.latencies) < 10 {
|
||||
// Not enough data yet
|
||||
return
|
||||
}
|
||||
|
||||
// Calculate p95 latency
|
||||
p95 := calculateP95(a.latencies)
|
||||
|
||||
currentThreshold := a.client.hubThreshold
|
||||
|
||||
log.Debug().
|
||||
Dur("p95_latency", p95).
|
||||
Dur("target_latency", a.targetLatency).
|
||||
Int("current_threshold", currentThreshold).
|
||||
Int("queries", len(a.latencies)).
|
||||
Msg("Auto-tuner evaluating performance")
|
||||
|
||||
// Determine adjustment direction
|
||||
var newThreshold int
|
||||
|
||||
if p95 > a.targetLatency {
|
||||
// Too slow - lower threshold (more hubs = faster queries)
|
||||
adjustment := calculateAdjustment(p95, a.targetLatency)
|
||||
newThreshold = currentThreshold - adjustment
|
||||
|
||||
if newThreshold < a.minThreshold {
|
||||
newThreshold = a.minThreshold
|
||||
}
|
||||
|
||||
log.Info().
|
||||
Dur("p95", p95).
|
||||
Int("old_threshold", currentThreshold).
|
||||
Int("new_threshold", newThreshold).
|
||||
Msg("Auto-tuner: Lowering hub threshold (too slow)")
|
||||
|
||||
} else if p95 < a.targetLatency*8/10 {
|
||||
// Too fast - raise threshold (fewer hubs = more savings)
|
||||
// Only adjust if significantly faster (20% margin)
|
||||
adjustment := calculateAdjustment(a.targetLatency, p95)
|
||||
newThreshold = currentThreshold + adjustment
|
||||
|
||||
if newThreshold > a.maxThreshold {
|
||||
newThreshold = a.maxThreshold
|
||||
}
|
||||
|
||||
log.Info().
|
||||
Dur("p95", p95).
|
||||
Int("old_threshold", currentThreshold).
|
||||
Int("new_threshold", newThreshold).
|
||||
Msg("Auto-tuner: Raising hub threshold (room for savings)")
|
||||
|
||||
} else {
|
||||
// Within acceptable range, no adjustment needed
|
||||
log.Debug().
|
||||
Dur("p95", p95).
|
||||
Int("threshold", currentThreshold).
|
||||
Msg("Auto-tuner: Performance acceptable, no adjustment")
|
||||
return
|
||||
}
|
||||
|
||||
// Apply adjustment
|
||||
if newThreshold != currentThreshold {
|
||||
a.client.hubThreshold = newThreshold
|
||||
a.adjustments++
|
||||
|
||||
// Clear latency history after adjustment
|
||||
a.latencies = make([]time.Duration, 0, 1000)
|
||||
|
||||
log.Info().
|
||||
Int("threshold", newThreshold).
|
||||
Int("total_adjustments", a.adjustments).
|
||||
Msg("Hub threshold adjusted by auto-tuner")
|
||||
}
|
||||
}
|
||||
|
||||
// calculateP95 computes the 95th percentile latency
|
||||
func calculateP95(latencies []time.Duration) time.Duration {
|
||||
if len(latencies) == 0 {
|
||||
return 0
|
||||
}
|
||||
|
||||
// Sort latencies
|
||||
sorted := make([]time.Duration, len(latencies))
|
||||
copy(sorted, latencies)
|
||||
|
||||
// Simple bubble sort (small dataset)
|
||||
n := len(sorted)
|
||||
for i := 0; i < n-1; i++ {
|
||||
for j := 0; j < n-i-1; j++ {
|
||||
if sorted[j] > sorted[j+1] {
|
||||
sorted[j], sorted[j+1] = sorted[j+1], sorted[j]
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Return 95th percentile
|
||||
idx := int(float64(len(sorted)) * 0.95)
|
||||
if idx >= len(sorted) {
|
||||
idx = len(sorted) - 1
|
||||
}
|
||||
|
||||
return sorted[idx]
|
||||
}
|
||||
|
||||
// calculateAdjustment determines how much to adjust threshold
|
||||
func calculateAdjustment(actual, target time.Duration) int {
|
||||
// Calculate percentage difference
|
||||
diff := float64(actual-target) / float64(target)
|
||||
|
||||
// Adjust more aggressively for larger differences
|
||||
if diff > 0.5 || diff < -0.5 {
|
||||
return 3 // Large adjustment
|
||||
} else if diff > 0.2 || diff < -0.2 {
|
||||
return 2 // Medium adjustment
|
||||
}
|
||||
|
||||
return 1 // Small adjustment
|
||||
}
|
||||
|
||||
// GetStats returns auto-tuner statistics
|
||||
func (a *AutoTuner) GetStats() AutoTunerStats {
|
||||
a.latenciesMu.Lock()
|
||||
defer a.latenciesMu.Unlock()
|
||||
|
||||
stats := AutoTunerStats{
|
||||
CurrentThreshold: a.client.hubThreshold,
|
||||
TargetLatency: a.targetLatency,
|
||||
TotalQueries: a.queries,
|
||||
TotalAdjustments: a.adjustments,
|
||||
RecentQueries: len(a.latencies),
|
||||
}
|
||||
|
||||
if len(a.latencies) > 0 {
|
||||
stats.P95Latency = calculateP95(a.latencies)
|
||||
|
||||
// Calculate average
|
||||
var total time.Duration
|
||||
for _, lat := range a.latencies {
|
||||
total += lat
|
||||
}
|
||||
stats.AvgLatency = total / time.Duration(len(a.latencies))
|
||||
}
|
||||
|
||||
return stats
|
||||
}
|
||||
|
||||
// AutoTunerStats contains auto-tuner statistics
|
||||
type AutoTunerStats struct {
|
||||
CurrentThreshold int
|
||||
TargetLatency time.Duration
|
||||
P95Latency time.Duration
|
||||
AvgLatency time.Duration
|
||||
TotalQueries int64
|
||||
TotalAdjustments int
|
||||
RecentQueries int
|
||||
}
|
||||
|
||||
// AutoTunedClient wraps Client with automatic performance tuning
|
||||
type AutoTunedClient struct {
|
||||
*Client
|
||||
tuner *AutoTuner
|
||||
}
|
||||
|
||||
// Query wraps the underlying Query call with latency tracking
|
||||
func (a *AutoTunedClient) Query(ctx context.Context, query string, limit int, where map[string]any) ([]sqlitevec.QueryResult, error) {
|
||||
start := time.Now()
|
||||
results, err := a.Client.Query(ctx, query, limit, where)
|
||||
latency := time.Since(start)
|
||||
|
||||
a.tuner.RecordQuery(latency)
|
||||
|
||||
return results, err
|
||||
}
|
||||
|
||||
// WithAutoTuning wraps a hybrid client with auto-tuning enabled
|
||||
func WithAutoTuning(client *Client, cfg AutoTunerConfig) *AutoTunedClient {
|
||||
tuner := NewAutoTuner(client, cfg)
|
||||
tuner.Start()
|
||||
|
||||
return &AutoTunedClient{
|
||||
Client: client,
|
||||
tuner: tuner,
|
||||
}
|
||||
}
|
||||
|
||||
// Stop stops the auto-tuner
|
||||
func (a *AutoTunedClient) StopTuning() {
|
||||
a.tuner.Stop()
|
||||
}
|
||||
@@ -0,0 +1,515 @@
|
||||
// Package hybrid provides LEANN-inspired selective vector storage for claude-mnemonic.
|
||||
//
|
||||
// This package implements a hybrid storage strategy where frequently-accessed
|
||||
// observations ("hubs") have their embeddings stored, while infrequently-accessed
|
||||
// observations have their embeddings recomputed on-demand during search.
|
||||
//
|
||||
// This approach reduces storage by 60-80% with minimal impact on search latency (<50ms).
|
||||
package hybrid
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"fmt"
|
||||
"math"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/lukaszraczylo/claude-mnemonic/internal/embedding"
|
||||
"github.com/lukaszraczylo/claude-mnemonic/internal/vector/sqlitevec"
|
||||
"github.com/rs/zerolog/log"
|
||||
)
|
||||
|
||||
// VectorStorageStrategy defines how embeddings are stored/computed
|
||||
type VectorStorageStrategy int
|
||||
|
||||
const (
|
||||
// StorageAlways stores all embeddings (current behavior, backwards compatible)
|
||||
StorageAlways VectorStorageStrategy = iota
|
||||
// StorageHub stores only frequently-accessed "hub" embeddings (recommended)
|
||||
StorageHub
|
||||
// StorageOnDemand recomputes all embeddings during search (maximum savings)
|
||||
StorageOnDemand
|
||||
)
|
||||
|
||||
// Client wraps sqlitevec.Client with selective storage logic
|
||||
type Client struct {
|
||||
base *sqlitevec.Client
|
||||
db *sql.DB
|
||||
embedSvc *embedding.Service
|
||||
accessCount map[string]int
|
||||
lastAccess map[string]time.Time
|
||||
contentCache map[string]string
|
||||
strategy VectorStorageStrategy
|
||||
hubThreshold int
|
||||
mu sync.RWMutex
|
||||
cacheMu sync.RWMutex
|
||||
}
|
||||
|
||||
// Config for hybrid client
|
||||
type Config struct {
|
||||
BaseClient *sqlitevec.Client
|
||||
DB *sql.DB
|
||||
EmbedSvc *embedding.Service
|
||||
Strategy VectorStorageStrategy
|
||||
HubThreshold int // Default: 5 accesses
|
||||
}
|
||||
|
||||
// NewClient creates a new hybrid vector client
|
||||
func NewClient(cfg Config) *Client {
|
||||
if cfg.HubThreshold <= 0 {
|
||||
cfg.HubThreshold = 5
|
||||
}
|
||||
|
||||
log.Info().
|
||||
Str("strategy", strategyToString(cfg.Strategy)).
|
||||
Int("hub_threshold", cfg.HubThreshold).
|
||||
Msg("Initializing LEANN hybrid vector client")
|
||||
|
||||
return &Client{
|
||||
base: cfg.BaseClient,
|
||||
db: cfg.DB,
|
||||
embedSvc: cfg.EmbedSvc,
|
||||
strategy: cfg.Strategy,
|
||||
hubThreshold: cfg.HubThreshold,
|
||||
accessCount: make(map[string]int),
|
||||
lastAccess: make(map[string]time.Time),
|
||||
contentCache: make(map[string]string),
|
||||
}
|
||||
}
|
||||
|
||||
// AddDocuments implements selective storage based on strategy
|
||||
func (c *Client) AddDocuments(ctx context.Context, docs []sqlitevec.Document) error {
|
||||
if len(docs) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
switch c.strategy {
|
||||
case StorageAlways:
|
||||
// Use existing implementation - store all embeddings
|
||||
return c.base.AddDocuments(ctx, docs)
|
||||
|
||||
case StorageHub:
|
||||
// Store only hub candidates
|
||||
return c.addDocumentsSelective(ctx, docs)
|
||||
|
||||
case StorageOnDemand:
|
||||
// Don't store embeddings, only cache content
|
||||
return c.cacheDocuments(ctx, docs)
|
||||
|
||||
default:
|
||||
return c.base.AddDocuments(ctx, docs)
|
||||
}
|
||||
}
|
||||
|
||||
// addDocumentsSelective stores embeddings only for hub-qualified documents
|
||||
func (c *Client) addDocumentsSelective(ctx context.Context, docs []sqlitevec.Document) error {
|
||||
// Always cache content for potential recomputation
|
||||
if err := c.cacheDocuments(ctx, docs); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// Filter to hub documents
|
||||
hubDocs := make([]sqlitevec.Document, 0, len(docs))
|
||||
for _, doc := range docs {
|
||||
if c.isHub(doc.ID) {
|
||||
hubDocs = append(hubDocs, doc)
|
||||
}
|
||||
}
|
||||
|
||||
// Store only hub embeddings
|
||||
if len(hubDocs) > 0 {
|
||||
log.Debug().
|
||||
Int("total", len(docs)).
|
||||
Int("hubs", len(hubDocs)).
|
||||
Msg("Storing selective embeddings")
|
||||
return c.base.AddDocuments(ctx, hubDocs)
|
||||
}
|
||||
|
||||
log.Debug().Int("total", len(docs)).Msg("All documents cached, no hubs to store")
|
||||
return nil
|
||||
}
|
||||
|
||||
// cacheDocuments stores content for later recomputation
|
||||
func (c *Client) cacheDocuments(ctx context.Context, docs []sqlitevec.Document) error {
|
||||
c.cacheMu.Lock()
|
||||
defer c.cacheMu.Unlock()
|
||||
|
||||
for _, doc := range docs {
|
||||
c.contentCache[doc.ID] = doc.Content
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// DeleteDocuments removes documents by their IDs
|
||||
func (c *Client) DeleteDocuments(ctx context.Context, ids []string) error {
|
||||
// Remove from base storage
|
||||
if err := c.base.DeleteDocuments(ctx, ids); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// Clean up caches
|
||||
c.mu.Lock()
|
||||
for _, id := range ids {
|
||||
delete(c.accessCount, id)
|
||||
delete(c.lastAccess, id)
|
||||
}
|
||||
c.mu.Unlock()
|
||||
|
||||
c.cacheMu.Lock()
|
||||
for _, id := range ids {
|
||||
delete(c.contentCache, id)
|
||||
}
|
||||
c.cacheMu.Unlock()
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// Query performs search with dynamic recomputation
|
||||
func (c *Client) Query(ctx context.Context, query string, limit int, where map[string]any) ([]sqlitevec.QueryResult, error) {
|
||||
switch c.strategy {
|
||||
case StorageAlways:
|
||||
// Use existing implementation
|
||||
return c.queryAndTrack(ctx, query, limit, where)
|
||||
|
||||
case StorageHub:
|
||||
// Search hubs, then expand with recomputation
|
||||
return c.queryHybrid(ctx, query, limit, where)
|
||||
|
||||
case StorageOnDemand:
|
||||
// Fully dynamic search
|
||||
return c.queryDynamic(ctx, query, limit, where)
|
||||
|
||||
default:
|
||||
return c.queryAndTrack(ctx, query, limit, where)
|
||||
}
|
||||
}
|
||||
|
||||
// queryAndTrack wraps base Query with access tracking
|
||||
func (c *Client) queryAndTrack(ctx context.Context, query string, limit int, where map[string]any) ([]sqlitevec.QueryResult, error) {
|
||||
results, err := c.base.Query(ctx, query, limit, where)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Track access for hub detection
|
||||
c.trackAccess(results)
|
||||
|
||||
return results, nil
|
||||
}
|
||||
|
||||
// queryHybrid searches stored hubs and recomputes non-hubs
|
||||
func (c *Client) queryHybrid(ctx context.Context, query string, limit int, where map[string]any) ([]sqlitevec.QueryResult, error) {
|
||||
startTime := time.Now()
|
||||
|
||||
// 1. Query stored hub embeddings (limit * 2 for expansion)
|
||||
hubResults, err := c.base.Query(ctx, query, limit*2, where)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// 2. Track access
|
||||
c.trackAccess(hubResults)
|
||||
|
||||
// 3. Get candidate non-hub IDs (from content cache)
|
||||
candidates := c.getCandidateNonHubs(where, limit*2)
|
||||
|
||||
// 4. Recompute embeddings for candidates if we have any
|
||||
var recomputedResults []sqlitevec.QueryResult
|
||||
if len(candidates) > 0 {
|
||||
recomputedResults, err = c.recomputeAndScore(ctx, query, candidates)
|
||||
if err != nil {
|
||||
// Log but don't fail - use hub results only
|
||||
log.Warn().Err(err).Msg("Failed to recompute embeddings, using hub results only")
|
||||
recomputedResults = nil
|
||||
}
|
||||
}
|
||||
|
||||
// 5. Merge and rank
|
||||
allResults := append(hubResults, recomputedResults...)
|
||||
sortBySimilarity(allResults)
|
||||
|
||||
// 6. Return top K
|
||||
if len(allResults) > limit {
|
||||
allResults = allResults[:limit]
|
||||
}
|
||||
|
||||
duration := time.Since(startTime)
|
||||
log.Debug().
|
||||
Dur("duration_ms", duration).
|
||||
Int("hubs", len(hubResults)).
|
||||
Int("recomputed", len(recomputedResults)).
|
||||
Int("results", len(allResults)).
|
||||
Msg("Hybrid search completed")
|
||||
|
||||
return allResults, nil
|
||||
}
|
||||
|
||||
// queryDynamic recomputes all embeddings on-the-fly
|
||||
func (c *Client) queryDynamic(ctx context.Context, query string, limit int, where map[string]any) ([]sqlitevec.QueryResult, error) {
|
||||
startTime := time.Now()
|
||||
|
||||
// Get all candidate IDs from content cache
|
||||
candidates := c.getCandidateNonHubs(where, limit*5)
|
||||
|
||||
// Recompute and score all
|
||||
results, err := c.recomputeAndScore(ctx, query, candidates)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Track access
|
||||
c.trackAccess(results)
|
||||
|
||||
// Return top K
|
||||
if len(results) > limit {
|
||||
results = results[:limit]
|
||||
}
|
||||
|
||||
duration := time.Since(startTime)
|
||||
log.Debug().
|
||||
Dur("duration_ms", duration).
|
||||
Int("recomputed", len(candidates)).
|
||||
Int("results", len(results)).
|
||||
Msg("Dynamic search completed")
|
||||
|
||||
return results, nil
|
||||
}
|
||||
|
||||
// recomputeAndScore generates embeddings and computes similarities
|
||||
func (c *Client) recomputeAndScore(ctx context.Context, query string, candidateIDs []string) ([]sqlitevec.QueryResult, error) {
|
||||
if len(candidateIDs) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
// Generate query embedding
|
||||
queryEmb, err := c.embedSvc.Embed(query)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("embed query: %w", err)
|
||||
}
|
||||
|
||||
// Get content for candidates
|
||||
c.cacheMu.RLock()
|
||||
texts := make([]string, 0, len(candidateIDs))
|
||||
validIDs := make([]string, 0, len(candidateIDs))
|
||||
for _, id := range candidateIDs {
|
||||
if content, ok := c.contentCache[id]; ok && content != "" {
|
||||
texts = append(texts, content)
|
||||
validIDs = append(validIDs, id)
|
||||
}
|
||||
}
|
||||
c.cacheMu.RUnlock()
|
||||
|
||||
if len(texts) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
// Batch generate embeddings
|
||||
embeddings, err := c.embedSvc.EmbedBatch(texts)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("batch embed: %w", err)
|
||||
}
|
||||
|
||||
// Compute similarities
|
||||
results := make([]sqlitevec.QueryResult, len(embeddings))
|
||||
for i, emb := range embeddings {
|
||||
similarity := cosineSimilarity(queryEmb, emb)
|
||||
distance := 1.0 - similarity // Convert to distance
|
||||
|
||||
results[i] = sqlitevec.QueryResult{
|
||||
ID: validIDs[i],
|
||||
Distance: float64(distance),
|
||||
Similarity: float64(similarity),
|
||||
Metadata: make(map[string]any),
|
||||
}
|
||||
}
|
||||
|
||||
return results, nil
|
||||
}
|
||||
|
||||
// trackAccess records document access for hub detection
|
||||
func (c *Client) trackAccess(results []sqlitevec.QueryResult) {
|
||||
if len(results) == 0 {
|
||||
return
|
||||
}
|
||||
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
|
||||
now := time.Now()
|
||||
for _, r := range results {
|
||||
c.accessCount[r.ID]++
|
||||
c.lastAccess[r.ID] = now
|
||||
}
|
||||
}
|
||||
|
||||
// isHub checks if a document qualifies as a hub
|
||||
func (c *Client) isHub(docID string) bool {
|
||||
c.mu.RLock()
|
||||
defer c.mu.RUnlock()
|
||||
|
||||
count := c.accessCount[docID]
|
||||
return count >= c.hubThreshold
|
||||
}
|
||||
|
||||
// getCandidateNonHubs returns IDs of non-hub documents matching filter
|
||||
func (c *Client) getCandidateNonHubs(where map[string]any, limit int) []string {
|
||||
c.cacheMu.RLock()
|
||||
defer c.cacheMu.RUnlock()
|
||||
|
||||
candidates := make([]string, 0, limit)
|
||||
for id := range c.contentCache {
|
||||
if !c.isHub(id) {
|
||||
candidates = append(candidates, id)
|
||||
if len(candidates) >= limit {
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return candidates
|
||||
}
|
||||
|
||||
// IsConnected always returns true (wraps base client)
|
||||
func (c *Client) IsConnected() bool {
|
||||
return c.base.IsConnected()
|
||||
}
|
||||
|
||||
// Close releases resources
|
||||
func (c *Client) Close() error {
|
||||
return c.base.Close()
|
||||
}
|
||||
|
||||
// Count returns the total number of vectors in the store
|
||||
func (c *Client) Count(ctx context.Context) (int64, error) {
|
||||
return c.base.Count(ctx)
|
||||
}
|
||||
|
||||
// ModelVersion returns the current embedding model version
|
||||
func (c *Client) ModelVersion() string {
|
||||
return c.base.ModelVersion()
|
||||
}
|
||||
|
||||
// NeedsRebuild checks if vectors need to be rebuilt due to model version change
|
||||
func (c *Client) NeedsRebuild(ctx context.Context) (bool, string) {
|
||||
return c.base.NeedsRebuild(ctx)
|
||||
}
|
||||
|
||||
// GetStaleVectors returns doc_ids of vectors with mismatched or null model versions
|
||||
func (c *Client) GetStaleVectors(ctx context.Context) ([]sqlitevec.StaleVectorInfo, error) {
|
||||
return c.base.GetStaleVectors(ctx)
|
||||
}
|
||||
|
||||
// DeleteVectorsByDocIDs removes vectors by their doc_ids
|
||||
func (c *Client) DeleteVectorsByDocIDs(ctx context.Context, docIDs []string) error {
|
||||
return c.base.DeleteVectorsByDocIDs(ctx, docIDs)
|
||||
}
|
||||
|
||||
// GetStorageStats returns storage efficiency metrics
|
||||
func (c *Client) GetStorageStats(ctx context.Context) (StorageStats, error) {
|
||||
c.mu.RLock()
|
||||
c.cacheMu.RLock()
|
||||
defer c.mu.RUnlock()
|
||||
defer c.cacheMu.RUnlock()
|
||||
|
||||
totalDocs := len(c.contentCache)
|
||||
hubCount := 0
|
||||
for id := range c.contentCache {
|
||||
if c.accessCount[id] >= c.hubThreshold {
|
||||
hubCount++
|
||||
}
|
||||
}
|
||||
|
||||
storedCount := hubCount
|
||||
if c.strategy == StorageAlways {
|
||||
// Get actual count from database
|
||||
if count, err := c.base.Count(ctx); err == nil {
|
||||
storedCount = int(count)
|
||||
}
|
||||
} else if c.strategy == StorageOnDemand {
|
||||
storedCount = 0
|
||||
}
|
||||
|
||||
embeddingSize := 384 * 4 // 384 dims × 4 bytes (float32)
|
||||
storedBytes := storedCount * embeddingSize
|
||||
potentialBytes := totalDocs * embeddingSize
|
||||
|
||||
savingsPercent := 0.0
|
||||
if potentialBytes > 0 {
|
||||
savingsPercent = (1.0 - float64(storedBytes)/float64(potentialBytes)) * 100
|
||||
}
|
||||
|
||||
return StorageStats{
|
||||
TotalDocuments: totalDocs,
|
||||
HubDocuments: hubCount,
|
||||
StoredEmbeddings: storedCount,
|
||||
StorageBytes: storedBytes,
|
||||
SavingsPercent: savingsPercent,
|
||||
Strategy: c.strategy,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// StorageStats contains storage efficiency metrics
|
||||
type StorageStats struct {
|
||||
TotalDocuments int
|
||||
HubDocuments int
|
||||
StoredEmbeddings int
|
||||
StorageBytes int
|
||||
SavingsPercent float64
|
||||
Strategy VectorStorageStrategy
|
||||
}
|
||||
|
||||
// Helper functions
|
||||
|
||||
func cosineSimilarity(a, b []float32) float32 {
|
||||
var dotProduct, normA, normB float32
|
||||
for i := range a {
|
||||
dotProduct += a[i] * b[i]
|
||||
normA += a[i] * a[i]
|
||||
normB += b[i] * b[i]
|
||||
}
|
||||
if normA == 0 || normB == 0 {
|
||||
return 0
|
||||
}
|
||||
return dotProduct / float32(math.Sqrt(float64(normA))*math.Sqrt(float64(normB)))
|
||||
}
|
||||
|
||||
func sortBySimilarity(results []sqlitevec.QueryResult) {
|
||||
// Use a simple but efficient sorting algorithm
|
||||
n := len(results)
|
||||
for i := 0; i < n-1; i++ {
|
||||
for j := 0; j < n-i-1; j++ {
|
||||
if results[j].Similarity < results[j+1].Similarity {
|
||||
results[j], results[j+1] = results[j+1], results[j]
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func strategyToString(s VectorStorageStrategy) string {
|
||||
switch s {
|
||||
case StorageAlways:
|
||||
return "always"
|
||||
case StorageHub:
|
||||
return "hub"
|
||||
case StorageOnDemand:
|
||||
return "on_demand"
|
||||
default:
|
||||
return "unknown"
|
||||
}
|
||||
}
|
||||
|
||||
// ParseStrategy converts a string to VectorStorageStrategy
|
||||
func ParseStrategy(s string) VectorStorageStrategy {
|
||||
switch s {
|
||||
case "hub":
|
||||
return StorageHub
|
||||
case "on_demand":
|
||||
return StorageOnDemand
|
||||
case "always":
|
||||
return StorageAlways
|
||||
default:
|
||||
return StorageHub // Default to hub strategy
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,187 @@
|
||||
package hybrid
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/lukaszraczylo/claude-mnemonic/internal/vector/sqlitevec"
|
||||
_ "github.com/mattn/go-sqlite3" // Import SQLite driver for CGO linking
|
||||
"github.com/stretchr/testify/assert"
|
||||
)
|
||||
|
||||
func TestParseStrategy(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
input string
|
||||
expected VectorStorageStrategy
|
||||
}{
|
||||
{"hub_strategy", "hub", StorageHub},
|
||||
{"on_demand_strategy", "on_demand", StorageOnDemand},
|
||||
{"always_strategy", "always", StorageAlways},
|
||||
{"invalid_defaults_to_hub", "invalid", StorageHub},
|
||||
{"empty_defaults_to_hub", "", StorageHub},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
result := ParseStrategy(tt.input)
|
||||
assert.Equal(t, tt.expected, result)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestStrategyToString(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
expected string
|
||||
input VectorStorageStrategy
|
||||
}{
|
||||
{"hub_to_string", "hub", StorageHub},
|
||||
{"on_demand_to_string", "on_demand", StorageOnDemand},
|
||||
{"always_to_string", "always", StorageAlways},
|
||||
{"invalid_to_unknown", "unknown", VectorStorageStrategy(99)},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
result := strategyToString(tt.input)
|
||||
assert.Equal(t, tt.expected, result)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestCosineSimilarity(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
a []float32
|
||||
b []float32
|
||||
expected float32
|
||||
}{
|
||||
{
|
||||
name: "identical_vectors",
|
||||
a: []float32{1, 0, 0},
|
||||
b: []float32{1, 0, 0},
|
||||
expected: 1.0,
|
||||
},
|
||||
{
|
||||
name: "orthogonal_vectors",
|
||||
a: []float32{1, 0, 0},
|
||||
b: []float32{0, 1, 0},
|
||||
expected: 0.0,
|
||||
},
|
||||
{
|
||||
name: "opposite_vectors",
|
||||
a: []float32{1, 0, 0},
|
||||
b: []float32{-1, 0, 0},
|
||||
expected: -1.0,
|
||||
},
|
||||
{
|
||||
name: "zero_vector",
|
||||
a: []float32{0, 0, 0},
|
||||
b: []float32{1, 1, 1},
|
||||
expected: 0.0,
|
||||
},
|
||||
{
|
||||
name: "parallel_vectors",
|
||||
a: []float32{2, 0, 0},
|
||||
b: []float32{4, 0, 0},
|
||||
expected: 1.0,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
result := cosineSimilarity(tt.a, tt.b)
|
||||
assert.InDelta(t, tt.expected, result, 0.001)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSortBySimilarity(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
input []sqlitevec.QueryResult
|
||||
expected []string // Expected order of IDs
|
||||
}{
|
||||
{
|
||||
name: "already_sorted",
|
||||
input: []sqlitevec.QueryResult{
|
||||
{ID: "doc1", Similarity: 0.9},
|
||||
{ID: "doc2", Similarity: 0.7},
|
||||
{ID: "doc3", Similarity: 0.5},
|
||||
},
|
||||
expected: []string{"doc1", "doc2", "doc3"},
|
||||
},
|
||||
{
|
||||
name: "reverse_sorted",
|
||||
input: []sqlitevec.QueryResult{
|
||||
{ID: "doc1", Similarity: 0.3},
|
||||
{ID: "doc2", Similarity: 0.7},
|
||||
{ID: "doc3", Similarity: 0.9},
|
||||
},
|
||||
expected: []string{"doc3", "doc2", "doc1"},
|
||||
},
|
||||
{
|
||||
name: "random_order",
|
||||
input: []sqlitevec.QueryResult{
|
||||
{ID: "doc1", Similarity: 0.5},
|
||||
{ID: "doc2", Similarity: 0.9},
|
||||
{ID: "doc3", Similarity: 0.3},
|
||||
{ID: "doc4", Similarity: 0.7},
|
||||
},
|
||||
expected: []string{"doc2", "doc4", "doc1", "doc3"},
|
||||
},
|
||||
{
|
||||
name: "identical_similarities",
|
||||
input: []sqlitevec.QueryResult{
|
||||
{ID: "doc1", Similarity: 0.5},
|
||||
{ID: "doc2", Similarity: 0.5},
|
||||
{ID: "doc3", Similarity: 0.5},
|
||||
},
|
||||
expected: []string{"doc1", "doc2", "doc3"},
|
||||
},
|
||||
{
|
||||
name: "empty_list",
|
||||
input: []sqlitevec.QueryResult{},
|
||||
expected: []string{},
|
||||
},
|
||||
{
|
||||
name: "single_element",
|
||||
input: []sqlitevec.QueryResult{
|
||||
{ID: "doc1", Similarity: 0.5},
|
||||
},
|
||||
expected: []string{"doc1"},
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
sortBySimilarity(tt.input)
|
||||
|
||||
actual := make([]string, len(tt.input))
|
||||
for i, r := range tt.input {
|
||||
actual[i] = r.ID
|
||||
}
|
||||
|
||||
assert.Equal(t, tt.expected, actual)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSortBySimilarity_PreserveOtherFields(t *testing.T) {
|
||||
input := []sqlitevec.QueryResult{
|
||||
{ID: "doc1", Similarity: 0.3, Distance: 0.7, Metadata: map[string]any{"key": "val1"}},
|
||||
{ID: "doc2", Similarity: 0.9, Distance: 0.1, Metadata: map[string]any{"key": "val2"}},
|
||||
}
|
||||
|
||||
sortBySimilarity(input)
|
||||
|
||||
assert.Equal(t, "doc2", input[0].ID)
|
||||
assert.InDelta(t, 0.9, input[0].Similarity, 0.001)
|
||||
assert.InDelta(t, 0.1, input[0].Distance, 0.001)
|
||||
assert.Equal(t, "val2", input[0].Metadata["key"])
|
||||
|
||||
assert.Equal(t, "doc1", input[1].ID)
|
||||
assert.InDelta(t, 0.3, input[1].Similarity, 0.001)
|
||||
assert.InDelta(t, 0.7, input[1].Distance, 0.001)
|
||||
assert.Equal(t, "val1", input[1].Metadata["key"])
|
||||
}
|
||||
@@ -0,0 +1,62 @@
|
||||
package hybrid
|
||||
|
||||
import (
|
||||
"os"
|
||||
"strconv"
|
||||
|
||||
"github.com/rs/zerolog/log"
|
||||
)
|
||||
|
||||
// GetStrategyFromEnv reads CLAUDE_MNEMONIC_VECTOR_STRATEGY from environment
|
||||
func GetStrategyFromEnv() VectorStorageStrategy {
|
||||
strategyStr := os.Getenv("CLAUDE_MNEMONIC_VECTOR_STRATEGY")
|
||||
if strategyStr == "" {
|
||||
// Default to hub strategy for optimal balance
|
||||
return StorageHub
|
||||
}
|
||||
|
||||
strategy := ParseStrategy(strategyStr)
|
||||
log.Info().
|
||||
Str("env_value", strategyStr).
|
||||
Str("strategy", strategyToString(strategy)).
|
||||
Msg("Vector storage strategy from environment")
|
||||
|
||||
return strategy
|
||||
}
|
||||
|
||||
// GetHubThresholdFromEnv reads CLAUDE_MNEMONIC_HUB_THRESHOLD from environment
|
||||
func GetHubThresholdFromEnv() int {
|
||||
thresholdStr := os.Getenv("CLAUDE_MNEMONIC_HUB_THRESHOLD")
|
||||
if thresholdStr == "" {
|
||||
return 5 // Default threshold
|
||||
}
|
||||
|
||||
threshold, err := strconv.Atoi(thresholdStr)
|
||||
if err != nil {
|
||||
log.Warn().
|
||||
Err(err).
|
||||
Str("env_value", thresholdStr).
|
||||
Msg("Invalid hub threshold in environment, using default")
|
||||
return 5
|
||||
}
|
||||
|
||||
if threshold < 1 {
|
||||
log.Warn().
|
||||
Int("env_value", threshold).
|
||||
Msg("Hub threshold too low, using minimum of 1")
|
||||
return 1
|
||||
}
|
||||
|
||||
log.Info().
|
||||
Int("threshold", threshold).
|
||||
Msg("Hub threshold from environment")
|
||||
|
||||
return threshold
|
||||
}
|
||||
|
||||
// IsHybridEnabled checks if hybrid storage should be used
|
||||
// Returns false if CLAUDE_MNEMONIC_VECTOR_STRATEGY=always (backwards compat)
|
||||
func IsHybridEnabled() bool {
|
||||
strategy := GetStrategyFromEnv()
|
||||
return strategy != StorageAlways
|
||||
}
|
||||
@@ -0,0 +1,308 @@
|
||||
package hybrid
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"sort"
|
||||
"time"
|
||||
|
||||
"github.com/lukaszraczylo/claude-mnemonic/internal/graph"
|
||||
"github.com/lukaszraczylo/claude-mnemonic/internal/vector/sqlitevec"
|
||||
"github.com/lukaszraczylo/claude-mnemonic/pkg/models"
|
||||
"github.com/rs/zerolog/log"
|
||||
)
|
||||
|
||||
// GraphConfig configures graph-aware search
|
||||
type GraphConfig struct {
|
||||
Enabled bool
|
||||
MaxHops int // Maximum graph traversal depth (default: 2)
|
||||
BranchFactor int // Number of neighbors to expand per node (default: 5)
|
||||
EdgeWeight float64 // Minimum edge weight to follow (default: 0.3)
|
||||
}
|
||||
|
||||
// DefaultGraphConfig returns sensible defaults for graph search
|
||||
func DefaultGraphConfig() GraphConfig {
|
||||
return GraphConfig{
|
||||
Enabled: true,
|
||||
MaxHops: 2,
|
||||
BranchFactor: 5,
|
||||
EdgeWeight: 0.3,
|
||||
}
|
||||
}
|
||||
|
||||
// GraphSearchClient wraps hybrid.Client with graph-aware search
|
||||
type GraphSearchClient struct {
|
||||
*Client
|
||||
graph *graph.ObservationGraph
|
||||
graphConfig GraphConfig
|
||||
}
|
||||
|
||||
// NewGraphSearchClient creates a graph-enhanced hybrid client
|
||||
func NewGraphSearchClient(baseClient *Client, observationGraph *graph.ObservationGraph, cfg GraphConfig) *GraphSearchClient {
|
||||
return &GraphSearchClient{
|
||||
Client: baseClient,
|
||||
graph: observationGraph,
|
||||
graphConfig: cfg,
|
||||
}
|
||||
}
|
||||
|
||||
// Query performs graph-aware vector search with two-level traversal
|
||||
func (g *GraphSearchClient) Query(ctx context.Context, query string, limit int, where map[string]any) ([]sqlitevec.QueryResult, error) {
|
||||
if !g.graphConfig.Enabled || g.graph == nil {
|
||||
// Fall back to standard hybrid search
|
||||
return g.Client.Query(ctx, query, limit, where)
|
||||
}
|
||||
|
||||
startTime := time.Now()
|
||||
|
||||
// 1. Generate query embedding
|
||||
queryEmb, err := g.embedSvc.Embed(query)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("embed query: %w", err)
|
||||
}
|
||||
|
||||
// 2. Search hub nodes (stored embeddings)
|
||||
hubResults, err := g.base.Query(ctx, query, limit*2, where)
|
||||
if err != nil {
|
||||
// Fall back to standard search on error
|
||||
log.Warn().Err(err).Msg("Hub search failed, falling back to hybrid search")
|
||||
return g.Client.Query(ctx, query, limit, where)
|
||||
}
|
||||
|
||||
// 3. Track hub access
|
||||
g.trackAccess(hubResults)
|
||||
|
||||
// 4. Expand via graph traversal
|
||||
expandedIDs := g.expandFromHubs(hubResults, limit*4)
|
||||
|
||||
// 5. Filter to non-hubs that need recomputation
|
||||
nonHubIDs := make([]string, 0)
|
||||
for _, id := range expandedIDs {
|
||||
if !g.isHub(id) {
|
||||
nonHubIDs = append(nonHubIDs, id)
|
||||
}
|
||||
}
|
||||
|
||||
// 6. Batch recompute non-hub embeddings
|
||||
recomputedResults, err := g.recomputeAndScore(ctx, query, nonHubIDs)
|
||||
if err != nil {
|
||||
log.Warn().Err(err).Msg("Recomputation failed, using hub results only")
|
||||
recomputedResults = nil
|
||||
}
|
||||
|
||||
// 7. Apply graph-based ranking boost
|
||||
allResults := g.mergeAndRankWithGraph(hubResults, recomputedResults, queryEmb)
|
||||
|
||||
// 8. Return top K
|
||||
if len(allResults) > limit {
|
||||
allResults = allResults[:limit]
|
||||
}
|
||||
|
||||
duration := time.Since(startTime)
|
||||
log.Debug().
|
||||
Dur("duration_ms", duration).
|
||||
Int("hubs", len(hubResults)).
|
||||
Int("expanded", len(expandedIDs)).
|
||||
Int("recomputed", len(recomputedResults)).
|
||||
Int("results", len(allResults)).
|
||||
Msg("Graph search completed")
|
||||
|
||||
return allResults, nil
|
||||
}
|
||||
|
||||
// expandFromHubs traverses graph from hub nodes to find promising candidates
|
||||
func (g *GraphSearchClient) expandFromHubs(hubResults []sqlitevec.QueryResult, maxCandidates int) []string {
|
||||
if g.graph == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
expanded := make(map[string]float64) // doc_id -> relevance score
|
||||
visited := make(map[int64]bool)
|
||||
|
||||
// Start from top hub results
|
||||
for i, result := range hubResults {
|
||||
if i >= g.graphConfig.BranchFactor*2 {
|
||||
break // Limit starting points
|
||||
}
|
||||
|
||||
// Parse observation ID from doc_id
|
||||
obsID := parseObservationID(result.ID)
|
||||
if obsID == 0 {
|
||||
continue
|
||||
}
|
||||
|
||||
// Mark as visited with high relevance (direct match)
|
||||
visited[obsID] = true
|
||||
expanded[result.ID] = result.Similarity
|
||||
|
||||
// Traverse graph from this hub
|
||||
g.traverseGraph(obsID, result.Similarity, 0, expanded, visited)
|
||||
}
|
||||
|
||||
// Convert to sorted list
|
||||
type candidate struct {
|
||||
ID string
|
||||
Relevance float64
|
||||
}
|
||||
|
||||
candidates := make([]candidate, 0, len(expanded))
|
||||
for id, rel := range expanded {
|
||||
candidates = append(candidates, candidate{ID: id, Relevance: rel})
|
||||
}
|
||||
|
||||
// Sort by relevance descending
|
||||
sort.Slice(candidates, func(i, j int) bool {
|
||||
return candidates[i].Relevance > candidates[j].Relevance
|
||||
})
|
||||
|
||||
// Return top candidates
|
||||
if len(candidates) > maxCandidates {
|
||||
candidates = candidates[:maxCandidates]
|
||||
}
|
||||
|
||||
result := make([]string, len(candidates))
|
||||
for i, c := range candidates {
|
||||
result[i] = c.ID
|
||||
}
|
||||
|
||||
return result
|
||||
}
|
||||
|
||||
// traverseGraph performs depth-limited graph traversal
|
||||
func (g *GraphSearchClient) traverseGraph(nodeID int64, baseRelevance float64, depth int, expanded map[string]float64, visited map[int64]bool) {
|
||||
if depth >= g.graphConfig.MaxHops {
|
||||
return // Max depth reached
|
||||
}
|
||||
|
||||
// Get neighbors from graph
|
||||
neighbors, weights, err := g.graph.GetNeighbors(nodeID)
|
||||
if err != nil {
|
||||
return // No neighbors or error
|
||||
}
|
||||
|
||||
// Traverse top neighbors by weight
|
||||
type neighborWeight struct {
|
||||
ID int64
|
||||
Weight float32
|
||||
}
|
||||
|
||||
neighborList := make([]neighborWeight, len(neighbors))
|
||||
for i := range neighbors {
|
||||
neighborList[i] = neighborWeight{
|
||||
ID: neighbors[i],
|
||||
Weight: weights[i],
|
||||
}
|
||||
}
|
||||
|
||||
// Sort by weight descending
|
||||
sort.Slice(neighborList, func(i, j int) bool {
|
||||
return neighborList[i].Weight > neighborList[j].Weight
|
||||
})
|
||||
|
||||
// Expand top branch_factor neighbors
|
||||
expanded_count := 0
|
||||
for _, nw := range neighborList {
|
||||
if expanded_count >= g.graphConfig.BranchFactor {
|
||||
break
|
||||
}
|
||||
|
||||
// Skip if edge weight too low
|
||||
if float64(nw.Weight) < g.graphConfig.EdgeWeight {
|
||||
continue
|
||||
}
|
||||
|
||||
// Skip if already visited
|
||||
if visited[nw.ID] {
|
||||
continue
|
||||
}
|
||||
visited[nw.ID] = true
|
||||
|
||||
// Calculate propagated relevance (decays with distance)
|
||||
decay := 0.7 // 30% decay per hop
|
||||
propagatedRelevance := baseRelevance * float64(nw.Weight) * decay
|
||||
|
||||
// Add to expanded set
|
||||
docID := formatObservationDocID(nw.ID)
|
||||
if existing, ok := expanded[docID]; !ok || propagatedRelevance > existing {
|
||||
expanded[docID] = propagatedRelevance
|
||||
}
|
||||
|
||||
// Recursively traverse
|
||||
g.traverseGraph(nw.ID, propagatedRelevance, depth+1, expanded, visited)
|
||||
expanded_count++
|
||||
}
|
||||
}
|
||||
|
||||
// mergeAndRankWithGraph combines hub and recomputed results with graph-based ranking
|
||||
func (g *GraphSearchClient) mergeAndRankWithGraph(hubResults, recomputedResults []sqlitevec.QueryResult, queryEmb []float32) []sqlitevec.QueryResult {
|
||||
// Merge results
|
||||
allResults := append(hubResults, recomputedResults...)
|
||||
|
||||
// Apply graph-based re-ranking
|
||||
if g.graph != nil {
|
||||
for i := range allResults {
|
||||
obsID := parseObservationID(allResults[i].ID)
|
||||
if obsID == 0 {
|
||||
continue
|
||||
}
|
||||
|
||||
// Boost score based on node degree (hubs are more important)
|
||||
node, err := g.graph.GetNode(obsID)
|
||||
if err == nil && node.Degree > 0 {
|
||||
// Degree boost: up to 10% increase for high-degree nodes
|
||||
degreeBoost := 1.0 + (0.1 * float64(node.Degree) / 20.0)
|
||||
if degreeBoost > 1.1 {
|
||||
degreeBoost = 1.1
|
||||
}
|
||||
allResults[i].Similarity *= degreeBoost
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Sort by adjusted similarity
|
||||
sortBySimilarity(allResults)
|
||||
|
||||
return allResults
|
||||
}
|
||||
|
||||
// parseObservationID extracts observation ID from doc_id
|
||||
// Format: "obs-{id}-{field}"
|
||||
func parseObservationID(docID string) int64 {
|
||||
var obsID int64
|
||||
// Ignore error - returns 0 on parse failure, which callers handle
|
||||
_, _ = fmt.Sscanf(docID, "obs-%d-", &obsID)
|
||||
return obsID
|
||||
}
|
||||
|
||||
// formatObservationDocID creates a doc_id for an observation
|
||||
func formatObservationDocID(obsID int64) string {
|
||||
return fmt.Sprintf("obs-%d-combined", obsID)
|
||||
}
|
||||
|
||||
// GetGraphStats returns statistics about the observation graph
|
||||
func (g *GraphSearchClient) GetGraphStats() graph.GraphStats {
|
||||
if g.graph == nil {
|
||||
return graph.GraphStats{}
|
||||
}
|
||||
return g.graph.Stats()
|
||||
}
|
||||
|
||||
// RebuildGraph rebuilds the observation graph from current observations
|
||||
// This should be called periodically or when observations change significantly
|
||||
func (g *GraphSearchClient) RebuildGraph(ctx context.Context, observations []*models.Observation) error {
|
||||
log.Info().Int("observations", len(observations)).Msg("Rebuilding observation graph")
|
||||
|
||||
newGraph, err := graph.BuildFromObservations(ctx, observations)
|
||||
if err != nil {
|
||||
return fmt.Errorf("build graph: %w", err)
|
||||
}
|
||||
|
||||
g.graph = newGraph
|
||||
|
||||
log.Info().
|
||||
Int("nodes", newGraph.Stats().NodeCount).
|
||||
Int("edges", newGraph.Stats().EdgeCount).
|
||||
Msg("Graph rebuilt successfully")
|
||||
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,17 @@
|
||||
package hybrid
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/lukaszraczylo/claude-mnemonic/internal/vector"
|
||||
_ "github.com/mattn/go-sqlite3" // Import SQLite driver for CGO linking
|
||||
)
|
||||
|
||||
// TestInterfaceImplementation verifies that hybrid clients implement vector.Client interface
|
||||
func TestInterfaceImplementation(t *testing.T) {
|
||||
// Compile-time check that Client implements vector.Client
|
||||
var _ vector.Client = (*Client)(nil)
|
||||
|
||||
// Compile-time check that GraphSearchClient implements vector.Client
|
||||
var _ vector.Client = (*GraphSearchClient)(nil)
|
||||
}
|
||||
@@ -0,0 +1,272 @@
|
||||
package hybrid
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Metrics tracks performance and usage statistics for hybrid vector storage
|
||||
type Metrics struct {
|
||||
startTime time.Time
|
||||
recentLatencies []time.Duration
|
||||
latenciesMu sync.Mutex
|
||||
totalQueries atomic.Int64
|
||||
hubOnlyQueries atomic.Int64
|
||||
hybridQueries atomic.Int64
|
||||
onDemandQueries atomic.Int64
|
||||
graphQueries atomic.Int64
|
||||
totalLatency atomic.Int64 // Sum in microseconds
|
||||
hubLatency atomic.Int64
|
||||
recomputeLatency atomic.Int64
|
||||
totalDocuments atomic.Int64
|
||||
hubDocuments atomic.Int64
|
||||
storedEmbeddings atomic.Int64
|
||||
recomputedCount atomic.Int64
|
||||
cacheHits atomic.Int64
|
||||
cacheMisses atomic.Int64
|
||||
graphTraversals atomic.Int64
|
||||
avgTraversalDepth atomic.Int64
|
||||
}
|
||||
|
||||
// NewMetrics creates a new metrics tracker
|
||||
func NewMetrics() *Metrics {
|
||||
return &Metrics{
|
||||
recentLatencies: make([]time.Duration, 0, 1000),
|
||||
startTime: time.Now(),
|
||||
}
|
||||
}
|
||||
|
||||
// RecordQuery records a query execution
|
||||
func (m *Metrics) RecordQuery(queryType string, latency time.Duration, recomputed int) {
|
||||
m.totalQueries.Add(1)
|
||||
m.totalLatency.Add(latency.Microseconds())
|
||||
|
||||
switch queryType {
|
||||
case "hub_only":
|
||||
m.hubOnlyQueries.Add(1)
|
||||
case "hybrid":
|
||||
m.hybridQueries.Add(1)
|
||||
case "on_demand":
|
||||
m.onDemandQueries.Add(1)
|
||||
case "graph":
|
||||
m.graphQueries.Add(1)
|
||||
}
|
||||
|
||||
if recomputed > 0 {
|
||||
m.recomputedCount.Add(int64(recomputed))
|
||||
}
|
||||
|
||||
// Track recent latencies
|
||||
m.latenciesMu.Lock()
|
||||
m.recentLatencies = append(m.recentLatencies, latency)
|
||||
if len(m.recentLatencies) > 1000 {
|
||||
m.recentLatencies = m.recentLatencies[len(m.recentLatencies)-1000:]
|
||||
}
|
||||
m.latenciesMu.Unlock()
|
||||
}
|
||||
|
||||
// RecordHubLatency records time spent in hub search
|
||||
func (m *Metrics) RecordHubLatency(latency time.Duration) {
|
||||
m.hubLatency.Add(latency.Microseconds())
|
||||
}
|
||||
|
||||
// RecordRecomputeLatency records time spent recomputing embeddings
|
||||
func (m *Metrics) RecordRecomputeLatency(latency time.Duration) {
|
||||
m.recomputeLatency.Add(latency.Microseconds())
|
||||
}
|
||||
|
||||
// RecordCacheHit records a content cache hit
|
||||
func (m *Metrics) RecordCacheHit() {
|
||||
m.cacheHits.Add(1)
|
||||
}
|
||||
|
||||
// RecordCacheMiss records a content cache miss
|
||||
func (m *Metrics) RecordCacheMiss() {
|
||||
m.cacheMisses.Add(1)
|
||||
}
|
||||
|
||||
// RecordGraphTraversal records a graph traversal operation
|
||||
func (m *Metrics) RecordGraphTraversal(depth int) {
|
||||
m.graphTraversals.Add(1)
|
||||
m.avgTraversalDepth.Add(int64(depth))
|
||||
}
|
||||
|
||||
// UpdateStorageStats updates current storage statistics
|
||||
func (m *Metrics) UpdateStorageStats(total, hubs, stored int) {
|
||||
m.totalDocuments.Store(int64(total))
|
||||
m.hubDocuments.Store(int64(hubs))
|
||||
m.storedEmbeddings.Store(int64(stored))
|
||||
}
|
||||
|
||||
// GetSnapshot returns current metrics snapshot
|
||||
func (m *Metrics) GetSnapshot() MetricsSnapshot {
|
||||
m.latenciesMu.Lock()
|
||||
defer m.latenciesMu.Unlock()
|
||||
|
||||
totalQueries := m.totalQueries.Load()
|
||||
|
||||
snapshot := MetricsSnapshot{
|
||||
// Query counts
|
||||
TotalQueries: totalQueries,
|
||||
HubOnlyQueries: m.hubOnlyQueries.Load(),
|
||||
HybridQueries: m.hybridQueries.Load(),
|
||||
OnDemandQueries: m.onDemandQueries.Load(),
|
||||
GraphQueries: m.graphQueries.Load(),
|
||||
|
||||
// Storage
|
||||
TotalDocuments: int(m.totalDocuments.Load()),
|
||||
HubDocuments: int(m.hubDocuments.Load()),
|
||||
StoredEmbeddings: int(m.storedEmbeddings.Load()),
|
||||
RecomputedTotal: m.recomputedCount.Load(),
|
||||
|
||||
// Cache
|
||||
CacheHits: m.cacheHits.Load(),
|
||||
CacheMisses: m.cacheMisses.Load(),
|
||||
|
||||
// Graph
|
||||
GraphTraversals: m.graphTraversals.Load(),
|
||||
|
||||
// Runtime
|
||||
Uptime: time.Since(m.startTime),
|
||||
}
|
||||
|
||||
// Calculate latencies
|
||||
if totalQueries > 0 {
|
||||
snapshot.AvgLatency = time.Duration(m.totalLatency.Load()/totalQueries) * time.Microsecond
|
||||
snapshot.AvgHubLatency = time.Duration(m.hubLatency.Load()/totalQueries) * time.Microsecond
|
||||
}
|
||||
|
||||
if m.recomputedCount.Load() > 0 {
|
||||
snapshot.AvgRecomputeLatency = time.Duration(m.recomputeLatency.Load()/m.recomputedCount.Load()) * time.Microsecond
|
||||
}
|
||||
|
||||
// Calculate percentiles
|
||||
if len(m.recentLatencies) > 0 {
|
||||
sorted := make([]time.Duration, len(m.recentLatencies))
|
||||
copy(sorted, m.recentLatencies)
|
||||
sortDurations(sorted)
|
||||
|
||||
snapshot.P50Latency = percentile(sorted, 0.50)
|
||||
snapshot.P95Latency = percentile(sorted, 0.95)
|
||||
snapshot.P99Latency = percentile(sorted, 0.99)
|
||||
}
|
||||
|
||||
// Calculate cache hit rate
|
||||
totalCacheOps := snapshot.CacheHits + snapshot.CacheMisses
|
||||
if totalCacheOps > 0 {
|
||||
snapshot.CacheHitRate = float64(snapshot.CacheHits) / float64(totalCacheOps)
|
||||
}
|
||||
|
||||
// Calculate storage savings
|
||||
if snapshot.TotalDocuments > 0 {
|
||||
embeddingSize := 384 * 4 // 384 dims × 4 bytes
|
||||
fullStorage := snapshot.TotalDocuments * embeddingSize
|
||||
actualStorage := snapshot.StoredEmbeddings * embeddingSize
|
||||
|
||||
if fullStorage > 0 {
|
||||
snapshot.StorageSavingsPercent = (1.0 - float64(actualStorage)/float64(fullStorage)) * 100
|
||||
}
|
||||
}
|
||||
|
||||
// Calculate avg traversal depth
|
||||
if snapshot.GraphTraversals > 0 {
|
||||
snapshot.AvgTraversalDepth = float64(m.avgTraversalDepth.Load()) / float64(snapshot.GraphTraversals)
|
||||
}
|
||||
|
||||
return snapshot
|
||||
}
|
||||
|
||||
// MetricsSnapshot represents a point-in-time metrics snapshot
|
||||
type MetricsSnapshot struct {
|
||||
// Query metrics
|
||||
TotalQueries int64
|
||||
HubOnlyQueries int64
|
||||
HybridQueries int64
|
||||
OnDemandQueries int64
|
||||
GraphQueries int64
|
||||
|
||||
// Latency metrics
|
||||
AvgLatency time.Duration
|
||||
P50Latency time.Duration
|
||||
P95Latency time.Duration
|
||||
P99Latency time.Duration
|
||||
AvgHubLatency time.Duration
|
||||
AvgRecomputeLatency time.Duration
|
||||
|
||||
// Storage metrics
|
||||
TotalDocuments int
|
||||
HubDocuments int
|
||||
StoredEmbeddings int
|
||||
StorageSavingsPercent float64
|
||||
RecomputedTotal int64
|
||||
|
||||
// Cache metrics
|
||||
CacheHits int64
|
||||
CacheMisses int64
|
||||
CacheHitRate float64
|
||||
|
||||
// Graph metrics
|
||||
GraphTraversals int64
|
||||
AvgTraversalDepth float64
|
||||
|
||||
// Runtime
|
||||
Uptime time.Duration
|
||||
}
|
||||
|
||||
// sortDurations sorts a slice of durations in ascending order
|
||||
func sortDurations(durations []time.Duration) {
|
||||
n := len(durations)
|
||||
for i := 0; i < n-1; i++ {
|
||||
for j := 0; j < n-i-1; j++ {
|
||||
if durations[j] > durations[j+1] {
|
||||
durations[j], durations[j+1] = durations[j+1], durations[j]
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// percentile calculates the Nth percentile from a sorted slice
|
||||
func percentile(sorted []time.Duration, p float64) time.Duration {
|
||||
if len(sorted) == 0 {
|
||||
return 0
|
||||
}
|
||||
|
||||
idx := int(float64(len(sorted)) * p)
|
||||
if idx >= len(sorted) {
|
||||
idx = len(sorted) - 1
|
||||
}
|
||||
|
||||
return sorted[idx]
|
||||
}
|
||||
|
||||
// String returns a human-readable representation of metrics
|
||||
func (s MetricsSnapshot) String() string {
|
||||
return fmt.Sprintf(`Hybrid Vector Storage Metrics:
|
||||
Queries:
|
||||
Total: %d (Hub: %d, Hybrid: %d, OnDemand: %d, Graph: %d)
|
||||
Avg Latency: %v (p50: %v, p95: %v, p99: %v)
|
||||
Hub Latency: %v, Recompute Latency: %v
|
||||
Storage:
|
||||
Documents: %d (Hubs: %d, %.1f%%)
|
||||
Stored Embeddings: %d
|
||||
Savings: %.1f%%
|
||||
Total Recomputed: %d
|
||||
Cache:
|
||||
Hits: %d, Misses: %d (Hit Rate: %.1f%%)
|
||||
Graph:
|
||||
Traversals: %d (Avg Depth: %.2f)
|
||||
Runtime: %v`,
|
||||
s.TotalQueries, s.HubOnlyQueries, s.HybridQueries, s.OnDemandQueries, s.GraphQueries,
|
||||
s.AvgLatency, s.P50Latency, s.P95Latency, s.P99Latency,
|
||||
s.AvgHubLatency, s.AvgRecomputeLatency,
|
||||
s.TotalDocuments, s.HubDocuments, float64(s.HubDocuments)/float64(s.TotalDocuments)*100,
|
||||
s.StoredEmbeddings,
|
||||
s.StorageSavingsPercent,
|
||||
s.RecomputedTotal,
|
||||
s.CacheHits, s.CacheMisses, s.CacheHitRate*100,
|
||||
s.GraphTraversals, s.AvgTraversalDepth,
|
||||
s.Uptime,
|
||||
)
|
||||
}
|
||||
@@ -0,0 +1,42 @@
|
||||
// Package vector provides common interfaces for vector storage implementations
|
||||
package vector
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
"github.com/lukaszraczylo/claude-mnemonic/internal/vector/sqlitevec"
|
||||
)
|
||||
|
||||
// Client defines the interface for vector storage operations.
|
||||
// Both sqlitevec.Client and hybrid.Client implement this interface.
|
||||
type Client interface {
|
||||
// AddDocuments adds documents with their embeddings to the vector store
|
||||
AddDocuments(ctx context.Context, docs []sqlitevec.Document) error
|
||||
|
||||
// DeleteDocuments removes documents by their IDs
|
||||
DeleteDocuments(ctx context.Context, ids []string) error
|
||||
|
||||
// Query performs a vector similarity search
|
||||
Query(ctx context.Context, query string, limit int, where map[string]any) ([]sqlitevec.QueryResult, error)
|
||||
|
||||
// IsConnected checks if the vector store is available
|
||||
IsConnected() bool
|
||||
|
||||
// Close releases resources
|
||||
Close() error
|
||||
|
||||
// Count returns the total number of vectors in the store
|
||||
Count(ctx context.Context) (int64, error)
|
||||
|
||||
// ModelVersion returns the current embedding model version
|
||||
ModelVersion() string
|
||||
|
||||
// NeedsRebuild checks if vectors need to be rebuilt due to model version change
|
||||
NeedsRebuild(ctx context.Context) (bool, string)
|
||||
|
||||
// GetStaleVectors returns doc_ids of vectors with mismatched or null model versions
|
||||
GetStaleVectors(ctx context.Context) ([]sqlitevec.StaleVectorInfo, error)
|
||||
|
||||
// DeleteVectorsByDocIDs removes vectors by their doc_ids
|
||||
DeleteVectorsByDocIDs(ctx context.Context, docIDs []string) error
|
||||
}
|
||||
@@ -5,19 +5,105 @@ import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"fmt"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
sqlite_vec "github.com/asg017/sqlite-vec-go-bindings/cgo"
|
||||
"github.com/lukaszraczylo/claude-mnemonic/internal/embedding"
|
||||
"github.com/rs/zerolog/log"
|
||||
"golang.org/x/sync/singleflight"
|
||||
)
|
||||
|
||||
// embeddingCacheEntry stores a cached embedding with its timestamp.
|
||||
type embeddingCacheEntry struct {
|
||||
embedding []float32
|
||||
timestamp int64 // Unix nano
|
||||
}
|
||||
|
||||
// resultCacheEntry stores cached query results.
|
||||
type resultCacheEntry struct {
|
||||
queryHash string
|
||||
results []QueryResult
|
||||
timestamp int64
|
||||
}
|
||||
|
||||
// Client provides vector operations via sqlite-vec.
|
||||
type Client struct {
|
||||
db *sql.DB
|
||||
embedSvc *embedding.Service
|
||||
mu sync.Mutex
|
||||
embeddingGroup singleflight.Group
|
||||
resultCache map[string]resultCacheEntry
|
||||
db *sql.DB
|
||||
embedSvc *embedding.Service
|
||||
queryCache map[string]embeddingCacheEntry
|
||||
stopCleanup chan struct{}
|
||||
stats CacheStats
|
||||
cleanupWg sync.WaitGroup
|
||||
resultCacheTTLNano int64
|
||||
cacheTTLNano int64
|
||||
resultCacheMaxSize int
|
||||
cacheMaxSize int
|
||||
resultCacheMu sync.RWMutex
|
||||
queryCacheMu sync.RWMutex
|
||||
readMu sync.RWMutex
|
||||
writeMu sync.Mutex
|
||||
}
|
||||
|
||||
// CacheStats tracks cache performance metrics using atomic counters for lock-free updates.
|
||||
type CacheStats struct {
|
||||
embeddingHits atomic.Int64
|
||||
embeddingMisses atomic.Int64
|
||||
resultHits atomic.Int64
|
||||
resultMisses atomic.Int64
|
||||
embeddingEvictions atomic.Int64
|
||||
resultEvictions atomic.Int64
|
||||
}
|
||||
|
||||
// CacheStatsSnapshot is the exported version of CacheStats for JSON marshaling.
|
||||
type CacheStatsSnapshot struct {
|
||||
EmbeddingHits int64 `json:"embedding_hits"`
|
||||
EmbeddingMisses int64 `json:"embedding_misses"`
|
||||
ResultHits int64 `json:"result_hits"`
|
||||
ResultMisses int64 `json:"result_misses"`
|
||||
EmbeddingEvictions int64 `json:"embedding_evictions"`
|
||||
ResultEvictions int64 `json:"result_evictions"`
|
||||
}
|
||||
|
||||
// HitRate returns the cache hit rate as a percentage.
|
||||
func (s CacheStatsSnapshot) HitRate() float64 {
|
||||
total := s.EmbeddingHits + s.EmbeddingMisses + s.ResultHits + s.ResultMisses
|
||||
if total == 0 {
|
||||
return 0
|
||||
}
|
||||
hits := s.EmbeddingHits + s.ResultHits
|
||||
return float64(hits) / float64(total) * 100
|
||||
}
|
||||
|
||||
// HitRate returns the cache hit rate as a percentage.
|
||||
func (s *CacheStats) HitRate() float64 {
|
||||
embHits := s.embeddingHits.Load()
|
||||
embMisses := s.embeddingMisses.Load()
|
||||
resHits := s.resultHits.Load()
|
||||
resMisses := s.resultMisses.Load()
|
||||
total := embHits + embMisses + resHits + resMisses
|
||||
if total == 0 {
|
||||
return 0
|
||||
}
|
||||
hits := embHits + resHits
|
||||
return float64(hits) / float64(total) * 100
|
||||
}
|
||||
|
||||
// Snapshot returns a copy of the current stats.
|
||||
func (s *CacheStats) Snapshot() CacheStatsSnapshot {
|
||||
return CacheStatsSnapshot{
|
||||
EmbeddingHits: s.embeddingHits.Load(),
|
||||
EmbeddingMisses: s.embeddingMisses.Load(),
|
||||
ResultHits: s.resultHits.Load(),
|
||||
ResultMisses: s.resultMisses.Load(),
|
||||
EmbeddingEvictions: s.embeddingEvictions.Load(),
|
||||
ResultEvictions: s.resultEvictions.Load(),
|
||||
}
|
||||
}
|
||||
|
||||
// Config holds configuration for the client.
|
||||
@@ -34,10 +120,23 @@ func NewClient(cfg Config, embedSvc *embedding.Service) (*Client, error) {
|
||||
return nil, fmt.Errorf("embedding service required")
|
||||
}
|
||||
|
||||
return &Client{
|
||||
db: cfg.DB,
|
||||
embedSvc: embedSvc,
|
||||
}, nil
|
||||
c := &Client{
|
||||
db: cfg.DB,
|
||||
embedSvc: embedSvc,
|
||||
queryCache: make(map[string]embeddingCacheEntry),
|
||||
cacheMaxSize: 500, // Cache up to 500 query embeddings
|
||||
cacheTTLNano: 5 * 60 * 1e9, // 5 minute TTL for embeddings
|
||||
resultCache: make(map[string]resultCacheEntry),
|
||||
resultCacheMaxSize: 200, // Cache up to 200 search results
|
||||
resultCacheTTLNano: 60 * 1e9, // 1 minute TTL for results (shorter since data changes)
|
||||
stopCleanup: make(chan struct{}),
|
||||
}
|
||||
|
||||
// Start background cache cleanup goroutine
|
||||
c.cleanupWg.Add(1)
|
||||
go c.cacheCleanupLoop()
|
||||
|
||||
return c, nil
|
||||
}
|
||||
|
||||
// AddDocuments adds documents with their embeddings to the vector store.
|
||||
@@ -46,8 +145,8 @@ func (c *Client) AddDocuments(ctx context.Context, docs []Document) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
c.writeMu.Lock()
|
||||
defer c.writeMu.Unlock()
|
||||
|
||||
// Generate embeddings for all documents
|
||||
texts := make([]string, len(docs))
|
||||
@@ -75,7 +174,10 @@ func (c *Client) AddDocuments(ctx context.Context, docs []Document) error {
|
||||
}
|
||||
defer func() {
|
||||
if err != nil {
|
||||
_ = tx.Rollback()
|
||||
if rbErr := tx.Rollback(); rbErr != nil {
|
||||
// Rollback failure is serious - indicates potential data corruption risk
|
||||
log.Error().Err(rbErr).Err(err).Msg("Failed to rollback transaction after error - data may be inconsistent")
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
@@ -118,6 +220,9 @@ func (c *Client) AddDocuments(ctx context.Context, docs []Document) error {
|
||||
return fmt.Errorf("commit transaction: %w", err)
|
||||
}
|
||||
|
||||
// Invalidate result cache since data changed
|
||||
c.InvalidateResultCache()
|
||||
|
||||
log.Debug().Int("count", len(docs)).Str("model", modelVersion).Msg("Added documents to sqlite-vec")
|
||||
return nil
|
||||
}
|
||||
@@ -128,12 +233,12 @@ func (c *Client) DeleteDocuments(ctx context.Context, ids []string) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
c.writeMu.Lock()
|
||||
defer c.writeMu.Unlock()
|
||||
|
||||
// Build placeholder string
|
||||
placeholders := make([]string, len(ids))
|
||||
args := make([]interface{}, len(ids))
|
||||
args := make([]any, len(ids))
|
||||
for i, id := range ids {
|
||||
placeholders[i] = "?"
|
||||
args[i] = id
|
||||
@@ -148,17 +253,25 @@ func (c *Client) DeleteDocuments(ctx context.Context, ids []string) error {
|
||||
return fmt.Errorf("delete documents: %w", err)
|
||||
}
|
||||
|
||||
// Invalidate result cache since data changed
|
||||
c.InvalidateResultCache()
|
||||
|
||||
log.Debug().Int("count", len(ids)).Msg("Deleted documents from sqlite-vec")
|
||||
return nil
|
||||
}
|
||||
|
||||
// Query performs a vector similarity search.
|
||||
func (c *Client) Query(ctx context.Context, query string, limit int, where map[string]any) ([]QueryResult, error) {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
// Build cache key from query + filters + limit
|
||||
cacheKey := c.buildResultCacheKey(query, limit, where)
|
||||
|
||||
// Generate query embedding
|
||||
queryEmb, err := c.embedSvc.Embed(query)
|
||||
// Check result cache first
|
||||
if results, ok := c.getResultFromCache(cacheKey); ok {
|
||||
return results, nil
|
||||
}
|
||||
|
||||
// Generate query embedding OUTSIDE the lock for better concurrency
|
||||
queryEmb, err := c.getOrComputeEmbedding(query)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("embed query: %w", err)
|
||||
}
|
||||
@@ -169,9 +282,13 @@ func (c *Client) Query(ctx context.Context, query string, limit int, where map[s
|
||||
return nil, fmt.Errorf("serialize query embedding: %w", err)
|
||||
}
|
||||
|
||||
// Now acquire read lock for the actual DB query
|
||||
c.readMu.RLock()
|
||||
defer c.readMu.RUnlock()
|
||||
|
||||
// Build query with filters
|
||||
// vec0 supports WHERE clauses on metadata columns
|
||||
args := []interface{}{queryBlob}
|
||||
args := []any{queryBlob}
|
||||
|
||||
sqlQuery := `
|
||||
SELECT
|
||||
@@ -232,6 +349,9 @@ func (c *Client) Query(ctx context.Context, query string, limit int, where map[s
|
||||
return nil, fmt.Errorf("iterate rows: %w", err)
|
||||
}
|
||||
|
||||
// Cache the results
|
||||
c.cacheResults(cacheKey, results)
|
||||
|
||||
log.Debug().
|
||||
Str("query", truncateString(query, 50)).
|
||||
Int("results", len(results)).
|
||||
@@ -245,11 +365,196 @@ func (c *Client) IsConnected() bool {
|
||||
return c.db != nil
|
||||
}
|
||||
|
||||
// Close is a no-op (db managed externally).
|
||||
// Close stops the background cleanup goroutine (db managed externally).
|
||||
func (c *Client) Close() error {
|
||||
// Signal cleanup goroutine to stop
|
||||
close(c.stopCleanup)
|
||||
// Wait for cleanup to finish
|
||||
c.cleanupWg.Wait()
|
||||
return nil
|
||||
}
|
||||
|
||||
// cacheCleanupLoop periodically removes expired cache entries.
|
||||
func (c *Client) cacheCleanupLoop() {
|
||||
defer c.cleanupWg.Done()
|
||||
|
||||
ticker := time.NewTicker(30 * time.Second) // Cleanup every 30 seconds
|
||||
defer ticker.Stop()
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-c.stopCleanup:
|
||||
return
|
||||
case <-ticker.C:
|
||||
c.cleanupExpiredCaches()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// cleanupExpiredCaches removes expired entries from both caches.
|
||||
func (c *Client) cleanupExpiredCaches() {
|
||||
now := time.Now().UnixNano()
|
||||
var embeddingExpired, resultExpired int64
|
||||
|
||||
// Cleanup embedding cache
|
||||
c.queryCacheMu.Lock()
|
||||
for key, entry := range c.queryCache {
|
||||
if now-entry.timestamp > c.cacheTTLNano {
|
||||
delete(c.queryCache, key)
|
||||
embeddingExpired++
|
||||
}
|
||||
}
|
||||
c.queryCacheMu.Unlock()
|
||||
|
||||
// Cleanup result cache
|
||||
c.resultCacheMu.Lock()
|
||||
for key, entry := range c.resultCache {
|
||||
if now-entry.timestamp > c.resultCacheTTLNano {
|
||||
delete(c.resultCache, key)
|
||||
resultExpired++
|
||||
}
|
||||
}
|
||||
c.resultCacheMu.Unlock()
|
||||
|
||||
// Update stats atomically
|
||||
if embeddingExpired > 0 || resultExpired > 0 {
|
||||
c.stats.embeddingEvictions.Add(embeddingExpired)
|
||||
c.stats.resultEvictions.Add(resultExpired)
|
||||
|
||||
log.Debug().
|
||||
Int64("embedding_expired", embeddingExpired).
|
||||
Int64("result_expired", resultExpired).
|
||||
Msg("Cache cleanup completed")
|
||||
}
|
||||
}
|
||||
|
||||
// BatchQueryResult holds results from a batch query operation.
|
||||
type BatchQueryResult struct {
|
||||
Error error
|
||||
Query string
|
||||
Results []QueryResult
|
||||
}
|
||||
|
||||
// QueryBatch performs multiple vector searches concurrently.
|
||||
// Returns results in the same order as input queries.
|
||||
// Uses a worker pool to limit concurrent queries.
|
||||
func (c *Client) QueryBatch(ctx context.Context, queries []string, limit int, where map[string]any) []BatchQueryResult {
|
||||
if len(queries) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
// Limit concurrency to avoid overwhelming the database
|
||||
maxConcurrent := min(4, len(queries))
|
||||
|
||||
results := make([]BatchQueryResult, len(queries))
|
||||
sem := make(chan struct{}, maxConcurrent)
|
||||
var wg sync.WaitGroup
|
||||
|
||||
for i, query := range queries {
|
||||
wg.Add(1)
|
||||
go func(idx int, q string) {
|
||||
defer wg.Done()
|
||||
|
||||
// Acquire semaphore
|
||||
select {
|
||||
case sem <- struct{}{}:
|
||||
defer func() { <-sem }()
|
||||
case <-ctx.Done():
|
||||
results[idx] = BatchQueryResult{
|
||||
Query: q,
|
||||
Error: ctx.Err(),
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
// Execute query
|
||||
queryResults, err := c.Query(ctx, q, limit, where)
|
||||
results[idx] = BatchQueryResult{
|
||||
Query: q,
|
||||
Results: queryResults,
|
||||
Error: err,
|
||||
}
|
||||
}(i, query)
|
||||
}
|
||||
|
||||
wg.Wait()
|
||||
return results
|
||||
}
|
||||
|
||||
// QueryMultiField searches across multiple fields for a single query.
|
||||
// Combines results from different field types and deduplicates by document ID.
|
||||
func (c *Client) QueryMultiField(ctx context.Context, query string, limit int, docType string, project string) ([]QueryResult, error) {
|
||||
// Generate embedding once
|
||||
queryEmb, err := c.getOrComputeEmbedding(query)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("embed query: %w", err)
|
||||
}
|
||||
|
||||
// Serialize query embedding
|
||||
queryBlob, err := sqlite_vec.SerializeFloat32(queryEmb)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("serialize query embedding: %w", err)
|
||||
}
|
||||
|
||||
c.readMu.RLock()
|
||||
defer c.readMu.RUnlock()
|
||||
|
||||
// Query with field type aggregation - get best match per document
|
||||
sqlQuery := `
|
||||
WITH ranked_results AS (
|
||||
SELECT
|
||||
doc_id,
|
||||
distance,
|
||||
sqlite_id,
|
||||
doc_type,
|
||||
field_type,
|
||||
project,
|
||||
scope,
|
||||
ROW_NUMBER() OVER (PARTITION BY sqlite_id ORDER BY distance ASC) as rn
|
||||
FROM vectors
|
||||
WHERE embedding MATCH ?
|
||||
AND doc_type = ?
|
||||
AND (project = ? OR scope = 'global')
|
||||
)
|
||||
SELECT doc_id, distance, sqlite_id, doc_type, field_type, project, scope
|
||||
FROM ranked_results
|
||||
WHERE rn = 1
|
||||
ORDER BY distance
|
||||
LIMIT ?
|
||||
`
|
||||
|
||||
rows, err := c.db.QueryContext(ctx, sqlQuery, queryBlob, docType, project, limit)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("query vectors: %w", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
// Pre-allocate with limit to avoid repeated slice growth
|
||||
results := make([]QueryResult, 0, limit)
|
||||
for rows.Next() {
|
||||
var r QueryResult
|
||||
var sqliteID int64
|
||||
var docTypeVal, fieldType, projectVal, scope sql.NullString
|
||||
|
||||
if err := rows.Scan(&r.ID, &r.Distance, &sqliteID, &docTypeVal, &fieldType, &projectVal, &scope); err != nil {
|
||||
return nil, fmt.Errorf("scan row: %w", err)
|
||||
}
|
||||
|
||||
r.Similarity = DistanceToSimilarity(r.Distance)
|
||||
r.Metadata = map[string]any{
|
||||
"sqlite_id": float64(sqliteID),
|
||||
"doc_type": docTypeVal.String,
|
||||
"field_type": fieldType.String,
|
||||
"project": projectVal.String,
|
||||
"scope": scope.String,
|
||||
}
|
||||
|
||||
results = append(results, r)
|
||||
}
|
||||
|
||||
return results, rows.Err()
|
||||
}
|
||||
|
||||
// truncateString truncates a string to maxLen characters.
|
||||
func truncateString(s string, maxLen int) string {
|
||||
if len(s) <= maxLen {
|
||||
@@ -260,8 +565,8 @@ func truncateString(s string, maxLen int) string {
|
||||
|
||||
// Count returns the total number of vectors in the store.
|
||||
func (c *Client) Count(ctx context.Context) (int64, error) {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
c.readMu.RLock()
|
||||
defer c.readMu.RUnlock()
|
||||
|
||||
var count int64
|
||||
err := c.db.QueryRowContext(ctx, "SELECT COUNT(*) FROM vectors").Scan(&count)
|
||||
@@ -281,8 +586,8 @@ func (c *Client) ModelVersion() string {
|
||||
// - The vectors table is empty
|
||||
// - Any vectors have a different model_version than the current model
|
||||
func (c *Client) NeedsRebuild(ctx context.Context) (bool, string) {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
c.readMu.RLock()
|
||||
defer c.readMu.RUnlock()
|
||||
|
||||
currentModel := c.embedSvc.Version()
|
||||
|
||||
@@ -319,18 +624,18 @@ func (c *Client) NeedsRebuild(ctx context.Context) (bool, string) {
|
||||
// StaleVectorInfo contains information about a vector that needs rebuilding.
|
||||
type StaleVectorInfo struct {
|
||||
DocID string
|
||||
SQLiteID int64
|
||||
DocType string
|
||||
FieldType string
|
||||
Project string
|
||||
Scope string
|
||||
SQLiteID int64
|
||||
}
|
||||
|
||||
// GetStaleVectors returns doc_ids of vectors with mismatched or null model versions.
|
||||
// This enables granular rebuild - only re-embedding documents that need updating.
|
||||
func (c *Client) GetStaleVectors(ctx context.Context) ([]StaleVectorInfo, error) {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
c.readMu.RLock()
|
||||
defer c.readMu.RUnlock()
|
||||
|
||||
currentModel := c.embedSvc.Version()
|
||||
|
||||
@@ -372,6 +677,134 @@ func (c *Client) GetStaleVectors(ctx context.Context) ([]StaleVectorInfo, error)
|
||||
return results, nil
|
||||
}
|
||||
|
||||
// VectorHealthStats contains comprehensive health information about the vector store.
|
||||
type VectorHealthStats struct {
|
||||
CoverageByType map[string]int64 `json:"coverage_by_type"`
|
||||
ModelVersions map[string]int64 `json:"model_versions"`
|
||||
ProjectCounts map[string]int64 `json:"project_counts"`
|
||||
CurrentModel string `json:"current_model"`
|
||||
RebuildReason string `json:"rebuild_reason,omitempty"`
|
||||
EmbeddingCache CacheStatsSnapshot `json:"embedding_cache"`
|
||||
TotalVectors int64 `json:"total_vectors"`
|
||||
StaleVectors int64 `json:"stale_vectors"`
|
||||
NeedsRebuild bool `json:"needs_rebuild"`
|
||||
}
|
||||
|
||||
// GetHealthStats returns comprehensive health statistics about the vector store.
|
||||
func (c *Client) GetHealthStats(ctx context.Context) (*VectorHealthStats, error) {
|
||||
c.readMu.RLock()
|
||||
defer c.readMu.RUnlock()
|
||||
|
||||
stats := &VectorHealthStats{
|
||||
CurrentModel: c.embedSvc.Version(),
|
||||
CoverageByType: make(map[string]int64),
|
||||
ModelVersions: make(map[string]int64),
|
||||
ProjectCounts: make(map[string]int64),
|
||||
EmbeddingCache: c.stats.Snapshot(),
|
||||
}
|
||||
|
||||
// Get total count
|
||||
err := c.db.QueryRowContext(ctx, "SELECT COUNT(*) FROM vectors").Scan(&stats.TotalVectors)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("count total vectors: %w", err)
|
||||
}
|
||||
|
||||
// Get stale count
|
||||
err = c.db.QueryRowContext(ctx,
|
||||
"SELECT COUNT(*) FROM vectors WHERE model_version != ? OR model_version IS NULL",
|
||||
stats.CurrentModel,
|
||||
).Scan(&stats.StaleVectors)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("count stale vectors: %w", err)
|
||||
}
|
||||
|
||||
// Check if rebuild needed
|
||||
stats.NeedsRebuild, stats.RebuildReason = c.needsRebuildUnlocked(ctx, stats.CurrentModel)
|
||||
|
||||
// Get coverage by doc_type
|
||||
rows, err := c.db.QueryContext(ctx, "SELECT doc_type, COUNT(*) FROM vectors GROUP BY doc_type")
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("query doc types: %w", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
for rows.Next() {
|
||||
var docType sql.NullString
|
||||
var count int64
|
||||
if err := rows.Scan(&docType, &count); err != nil {
|
||||
return nil, fmt.Errorf("scan doc type: %w", err)
|
||||
}
|
||||
if docType.Valid {
|
||||
stats.CoverageByType[docType.String] = count
|
||||
} else {
|
||||
stats.CoverageByType["unknown"] = count
|
||||
}
|
||||
}
|
||||
|
||||
// Get model version distribution
|
||||
rows2, err := c.db.QueryContext(ctx, "SELECT COALESCE(model_version, 'unknown'), COUNT(*) FROM vectors GROUP BY model_version")
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("query model versions: %w", err)
|
||||
}
|
||||
defer rows2.Close()
|
||||
|
||||
for rows2.Next() {
|
||||
var version string
|
||||
var count int64
|
||||
if err := rows2.Scan(&version, &count); err != nil {
|
||||
return nil, fmt.Errorf("scan model version: %w", err)
|
||||
}
|
||||
stats.ModelVersions[version] = count
|
||||
}
|
||||
|
||||
// Get project counts (top 10)
|
||||
rows3, err := c.db.QueryContext(ctx,
|
||||
"SELECT COALESCE(project, 'global'), COUNT(*) FROM vectors GROUP BY project ORDER BY COUNT(*) DESC LIMIT 10")
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("query projects: %w", err)
|
||||
}
|
||||
defer rows3.Close()
|
||||
|
||||
for rows3.Next() {
|
||||
var project string
|
||||
var count int64
|
||||
if err := rows3.Scan(&project, &count); err != nil {
|
||||
return nil, fmt.Errorf("scan project: %w", err)
|
||||
}
|
||||
stats.ProjectCounts[project] = count
|
||||
}
|
||||
|
||||
return stats, nil
|
||||
}
|
||||
|
||||
// needsRebuildUnlocked checks if rebuild is needed without acquiring lock (caller must hold lock).
|
||||
func (c *Client) needsRebuildUnlocked(ctx context.Context, currentModel string) (bool, string) {
|
||||
var totalCount int64
|
||||
err := c.db.QueryRowContext(ctx, "SELECT COUNT(*) FROM vectors").Scan(&totalCount)
|
||||
if err != nil {
|
||||
return false, ""
|
||||
}
|
||||
|
||||
if totalCount == 0 {
|
||||
return true, "empty"
|
||||
}
|
||||
|
||||
var staleCount int64
|
||||
err = c.db.QueryRowContext(ctx,
|
||||
"SELECT COUNT(*) FROM vectors WHERE model_version != ? OR model_version IS NULL",
|
||||
currentModel,
|
||||
).Scan(&staleCount)
|
||||
if err != nil {
|
||||
return false, ""
|
||||
}
|
||||
|
||||
if staleCount > 0 {
|
||||
return true, fmt.Sprintf("model_mismatch:%d", staleCount)
|
||||
}
|
||||
|
||||
return false, ""
|
||||
}
|
||||
|
||||
// DeleteVectorsByDocIDs removes vectors by their doc_ids.
|
||||
// Used for granular rebuild - delete stale vectors before re-adding.
|
||||
func (c *Client) DeleteVectorsByDocIDs(ctx context.Context, docIDs []string) error {
|
||||
@@ -379,12 +812,12 @@ func (c *Client) DeleteVectorsByDocIDs(ctx context.Context, docIDs []string) err
|
||||
return nil
|
||||
}
|
||||
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
c.writeMu.Lock()
|
||||
defer c.writeMu.Unlock()
|
||||
|
||||
// Build placeholder string
|
||||
placeholders := make([]string, len(docIDs))
|
||||
args := make([]interface{}, len(docIDs))
|
||||
args := make([]any, len(docIDs))
|
||||
for i, id := range docIDs {
|
||||
placeholders[i] = "?"
|
||||
args[i] = id
|
||||
@@ -402,3 +835,249 @@ func (c *Client) DeleteVectorsByDocIDs(ctx context.Context, docIDs []string) err
|
||||
log.Debug().Int("count", len(docIDs)).Msg("Deleted stale vectors by doc_id")
|
||||
return nil
|
||||
}
|
||||
|
||||
// DeleteByObservationID removes all vectors associated with an observation ID.
|
||||
// Vectors are stored with doc_ids that include the observation ID, e.g., "obs_123_narrative".
|
||||
func (c *Client) DeleteByObservationID(ctx context.Context, obsID int64) error {
|
||||
c.writeMu.Lock()
|
||||
defer c.writeMu.Unlock()
|
||||
|
||||
// Vectors have doc_ids like "obs_123_narrative", "obs_123_facts_0", etc.
|
||||
pattern := fmt.Sprintf("obs_%d_%%", obsID)
|
||||
|
||||
_, err := c.db.ExecContext(ctx, "DELETE FROM vectors WHERE doc_id LIKE ?", pattern)
|
||||
if err != nil {
|
||||
return fmt.Errorf("delete vectors for observation %d: %w", obsID, err)
|
||||
}
|
||||
|
||||
log.Debug().Int64("observation_id", obsID).Msg("Deleted vectors for observation")
|
||||
return nil
|
||||
}
|
||||
|
||||
// getOrComputeEmbedding returns a cached embedding or computes a new one.
|
||||
// Uses singleflight to prevent duplicate concurrent computations for the same query.
|
||||
func (c *Client) getOrComputeEmbedding(query string) ([]float32, error) {
|
||||
now := time.Now().UnixNano()
|
||||
|
||||
// Check cache first (read lock)
|
||||
c.queryCacheMu.RLock()
|
||||
if entry, ok := c.queryCache[query]; ok {
|
||||
if now-entry.timestamp < c.cacheTTLNano {
|
||||
c.queryCacheMu.RUnlock()
|
||||
c.stats.embeddingHits.Add(1)
|
||||
return entry.embedding, nil
|
||||
}
|
||||
}
|
||||
c.queryCacheMu.RUnlock()
|
||||
|
||||
// Cache miss - use singleflight to deduplicate concurrent embedding requests
|
||||
result, err, _ := c.embeddingGroup.Do(query, func() (any, error) {
|
||||
// Double-check cache inside singleflight (another goroutine may have just cached it)
|
||||
c.queryCacheMu.RLock()
|
||||
if entry, ok := c.queryCache[query]; ok {
|
||||
if time.Now().UnixNano()-entry.timestamp < c.cacheTTLNano {
|
||||
c.queryCacheMu.RUnlock()
|
||||
return entry.embedding, nil
|
||||
}
|
||||
}
|
||||
c.queryCacheMu.RUnlock()
|
||||
|
||||
// Record cache miss
|
||||
c.stats.embeddingMisses.Add(1)
|
||||
|
||||
// Compute embedding
|
||||
emb, err := c.embedSvc.Embed(query)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Store in cache (write lock)
|
||||
c.queryCacheMu.Lock()
|
||||
nowCache := time.Now().UnixNano()
|
||||
// Evict old entries if cache is full or near capacity (80% threshold)
|
||||
evictionThreshold := (c.cacheMaxSize * 8) / 10
|
||||
if len(c.queryCache) >= evictionThreshold {
|
||||
// Phase 1: Remove ALL expired entries first (not just 10%)
|
||||
evicted := int64(0)
|
||||
for k, v := range c.queryCache {
|
||||
if nowCache-v.timestamp > c.cacheTTLNano {
|
||||
delete(c.queryCache, k)
|
||||
evicted++
|
||||
}
|
||||
}
|
||||
|
||||
// Phase 2: If still at capacity, evict 10% using random iteration (O(n) instead of O(n log n))
|
||||
// Go map iteration order is randomized, providing good cache behavior without sorting
|
||||
if len(c.queryCache) >= c.cacheMaxSize {
|
||||
evictCount := max(c.cacheMaxSize/10, 1)
|
||||
for k := range c.queryCache {
|
||||
delete(c.queryCache, k)
|
||||
evicted++
|
||||
evictCount--
|
||||
if evictCount <= 0 {
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if evicted > 0 {
|
||||
c.stats.embeddingEvictions.Add(evicted)
|
||||
}
|
||||
}
|
||||
c.queryCache[query] = embeddingCacheEntry{
|
||||
embedding: emb,
|
||||
timestamp: nowCache,
|
||||
}
|
||||
c.queryCacheMu.Unlock()
|
||||
|
||||
return emb, nil
|
||||
})
|
||||
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return result.([]float32), nil
|
||||
}
|
||||
|
||||
// ClearCache clears the embedding cache.
|
||||
func (c *Client) ClearCache() {
|
||||
c.queryCacheMu.Lock()
|
||||
c.queryCache = make(map[string]embeddingCacheEntry)
|
||||
c.queryCacheMu.Unlock()
|
||||
}
|
||||
|
||||
// GetCacheStats returns comprehensive cache statistics.
|
||||
func (c *Client) GetCacheStats() CacheStatsSnapshot {
|
||||
return c.stats.Snapshot()
|
||||
}
|
||||
|
||||
// CacheStats returns basic cache size info for backward compatibility.
|
||||
// Deprecated: Use GetCacheStats() for comprehensive statistics.
|
||||
func (c *Client) CacheStats() (size int, maxSize int) {
|
||||
c.queryCacheMu.RLock()
|
||||
size = len(c.queryCache)
|
||||
c.queryCacheMu.RUnlock()
|
||||
return size, c.cacheMaxSize
|
||||
}
|
||||
|
||||
// EmbeddingCacheSize returns the current embedding cache size.
|
||||
func (c *Client) EmbeddingCacheSize() int {
|
||||
c.queryCacheMu.RLock()
|
||||
defer c.queryCacheMu.RUnlock()
|
||||
return len(c.queryCache)
|
||||
}
|
||||
|
||||
// ResultCacheSize returns the current result cache size.
|
||||
func (c *Client) ResultCacheSize() int {
|
||||
c.resultCacheMu.RLock()
|
||||
defer c.resultCacheMu.RUnlock()
|
||||
return len(c.resultCache)
|
||||
}
|
||||
|
||||
// buildResultCacheKey creates a unique key for caching query results.
|
||||
// Uses strings.Builder to avoid intermediate allocations.
|
||||
func (c *Client) buildResultCacheKey(query string, limit int, where map[string]any) string {
|
||||
// Pre-allocate with typical key size to avoid reallocation
|
||||
var b strings.Builder
|
||||
b.Grow(len(query) + 32) // query + typical prefix/suffix overhead
|
||||
|
||||
b.WriteString("q:")
|
||||
b.WriteString(query)
|
||||
b.WriteString(":l:")
|
||||
b.WriteString(strconv.Itoa(limit))
|
||||
|
||||
if docType, ok := where["doc_type"].(string); ok {
|
||||
b.WriteString(":dt:")
|
||||
b.WriteString(docType)
|
||||
}
|
||||
if project, ok := where["project"].(string); ok {
|
||||
b.WriteString(":p:")
|
||||
b.WriteString(project)
|
||||
}
|
||||
return b.String()
|
||||
}
|
||||
|
||||
// getResultFromCache retrieves cached results if available and not expired.
|
||||
func (c *Client) getResultFromCache(cacheKey string) ([]QueryResult, bool) {
|
||||
now := time.Now().UnixNano()
|
||||
|
||||
c.resultCacheMu.RLock()
|
||||
entry, ok := c.resultCache[cacheKey]
|
||||
c.resultCacheMu.RUnlock()
|
||||
|
||||
if !ok {
|
||||
c.stats.resultMisses.Add(1)
|
||||
return nil, false
|
||||
}
|
||||
|
||||
// Check if entry is expired
|
||||
if now-entry.timestamp > c.resultCacheTTLNano {
|
||||
c.stats.resultMisses.Add(1)
|
||||
return nil, false
|
||||
}
|
||||
|
||||
c.stats.resultHits.Add(1)
|
||||
|
||||
// Return a copy to prevent mutation
|
||||
results := make([]QueryResult, len(entry.results))
|
||||
copy(results, entry.results)
|
||||
return results, true
|
||||
}
|
||||
|
||||
// cacheResults stores query results in the cache.
|
||||
func (c *Client) cacheResults(cacheKey string, results []QueryResult) {
|
||||
now := time.Now().UnixNano()
|
||||
|
||||
c.resultCacheMu.Lock()
|
||||
defer c.resultCacheMu.Unlock()
|
||||
|
||||
// Evict old entries if cache is full
|
||||
if len(c.resultCache) >= c.resultCacheMaxSize {
|
||||
// Two-phase eviction: (1) TTL-expired entries, (2) random if still over capacity
|
||||
evicted := 0
|
||||
targetSize := c.resultCacheMaxSize * 8 / 10 // Target 80% capacity
|
||||
|
||||
// Phase 1: Remove all TTL-expired entries
|
||||
for k, v := range c.resultCache {
|
||||
if now-v.timestamp > c.resultCacheTTLNano {
|
||||
delete(c.resultCache, k)
|
||||
evicted++
|
||||
}
|
||||
}
|
||||
|
||||
// Phase 2: If still over target, remove random entries until at target
|
||||
if len(c.resultCache) >= targetSize {
|
||||
evictCount := len(c.resultCache) - targetSize + 1
|
||||
for k := range c.resultCache {
|
||||
delete(c.resultCache, k)
|
||||
evicted++
|
||||
evictCount--
|
||||
if evictCount <= 0 {
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if evicted > 0 {
|
||||
c.stats.resultEvictions.Add(int64(evicted))
|
||||
}
|
||||
}
|
||||
|
||||
// Make a copy of results to store
|
||||
resultsCopy := make([]QueryResult, len(results))
|
||||
copy(resultsCopy, results)
|
||||
|
||||
c.resultCache[cacheKey] = resultCacheEntry{
|
||||
results: resultsCopy,
|
||||
timestamp: now,
|
||||
queryHash: cacheKey,
|
||||
}
|
||||
}
|
||||
|
||||
// InvalidateResultCache clears the result cache.
|
||||
// Should be called after write operations that modify vectors.
|
||||
func (c *Client) InvalidateResultCache() {
|
||||
c.resultCacheMu.Lock()
|
||||
c.resultCache = make(map[string]resultCacheEntry)
|
||||
c.resultCacheMu.Unlock()
|
||||
}
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -12,17 +12,17 @@ const (
|
||||
|
||||
// Document represents a document to store with vector embedding.
|
||||
type Document struct {
|
||||
Metadata map[string]any
|
||||
ID string
|
||||
Content string
|
||||
Metadata map[string]any
|
||||
}
|
||||
|
||||
// QueryResult represents a search result from vector search.
|
||||
type QueryResult struct {
|
||||
Metadata map[string]any
|
||||
ID string
|
||||
Distance float64
|
||||
Similarity float64 // 1.0 = identical, 0.0 = opposite (derived from distance)
|
||||
Metadata map[string]any
|
||||
Similarity float64
|
||||
}
|
||||
|
||||
// DistanceToSimilarity converts sqlite-vec cosine distance to similarity score.
|
||||
|
||||
@@ -42,10 +42,10 @@ func TestQueryResult_Fields(t *testing.T) {
|
||||
|
||||
func TestBuildWhereFilter(t *testing.T) {
|
||||
tests := []struct {
|
||||
expected map[string]interface{}
|
||||
name string
|
||||
docType DocType
|
||||
project string
|
||||
expected map[string]interface{}
|
||||
}{
|
||||
{
|
||||
name: "empty_filters",
|
||||
@@ -474,9 +474,9 @@ func TestCopyMetadataMulti(t *testing.T) {
|
||||
func TestJoinStrings(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
strs []string
|
||||
sep string
|
||||
expected string
|
||||
strs []string
|
||||
}{
|
||||
{
|
||||
name: "empty_slice",
|
||||
@@ -522,8 +522,8 @@ func TestTruncateString(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
input string
|
||||
maxLen int
|
||||
expected string
|
||||
maxLen int
|
||||
}{
|
||||
{
|
||||
name: "shorter_than_max",
|
||||
@@ -577,10 +577,10 @@ func TestFilterByThreshold(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
results []QueryResult
|
||||
expectedIDs []string
|
||||
threshold float64
|
||||
maxResults int
|
||||
expectedLen int
|
||||
expectedIDs []string
|
||||
}{
|
||||
{
|
||||
name: "empty_results",
|
||||
|
||||
@@ -338,3 +338,187 @@ func (s *Sync) DeletePatterns(ctx context.Context, patternIDs []int64) error {
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// BatchSyncConfig configures batch synchronization behavior.
|
||||
type BatchSyncConfig struct {
|
||||
BatchSize int // Number of items per batch (default: 50)
|
||||
ProgressLogFreq int // Log progress every N items (default: 100)
|
||||
}
|
||||
|
||||
// DefaultBatchSyncConfig returns sensible defaults for batch sync.
|
||||
func DefaultBatchSyncConfig() BatchSyncConfig {
|
||||
return BatchSyncConfig{
|
||||
BatchSize: 50,
|
||||
ProgressLogFreq: 100,
|
||||
}
|
||||
}
|
||||
|
||||
// BatchSyncObservations syncs multiple observations efficiently in batches.
|
||||
// This reduces memory pressure during large rebuilds by processing in chunks.
|
||||
func (s *Sync) BatchSyncObservations(ctx context.Context, observations []*models.Observation, cfg BatchSyncConfig) (synced int, errors int) {
|
||||
if len(observations) == 0 {
|
||||
return 0, 0
|
||||
}
|
||||
|
||||
if cfg.BatchSize <= 0 {
|
||||
cfg.BatchSize = 50
|
||||
}
|
||||
if cfg.ProgressLogFreq <= 0 {
|
||||
cfg.ProgressLogFreq = 100
|
||||
}
|
||||
|
||||
for i := 0; i < len(observations); i += cfg.BatchSize {
|
||||
// Check context cancellation
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
log.Warn().Int("synced", synced).Int("remaining", len(observations)-i).Msg("Batch sync cancelled")
|
||||
return synced, errors
|
||||
default:
|
||||
}
|
||||
|
||||
end := min(i+cfg.BatchSize, len(observations))
|
||||
|
||||
batch := observations[i:end]
|
||||
var docs []Document
|
||||
|
||||
// Collect all documents for this batch
|
||||
for _, obs := range batch {
|
||||
docs = append(docs, s.formatObservationDocs(obs)...)
|
||||
}
|
||||
|
||||
// Add all documents in one call
|
||||
if len(docs) > 0 {
|
||||
if err := s.client.AddDocuments(ctx, docs); err != nil {
|
||||
log.Warn().Err(err).Int("batchStart", i).Int("batchSize", len(batch)).Msg("Failed to sync observation batch")
|
||||
errors += len(batch)
|
||||
continue
|
||||
}
|
||||
}
|
||||
|
||||
synced += len(batch)
|
||||
|
||||
// Log progress periodically
|
||||
if synced%cfg.ProgressLogFreq == 0 || synced == len(observations) {
|
||||
log.Debug().Int("synced", synced).Int("total", len(observations)).Msg("Observation batch sync progress")
|
||||
}
|
||||
}
|
||||
|
||||
return synced, errors
|
||||
}
|
||||
|
||||
// BatchSyncSummaries syncs multiple summaries efficiently in batches.
|
||||
func (s *Sync) BatchSyncSummaries(ctx context.Context, summaries []*models.SessionSummary, cfg BatchSyncConfig) (synced int, errors int) {
|
||||
if len(summaries) == 0 {
|
||||
return 0, 0
|
||||
}
|
||||
|
||||
if cfg.BatchSize <= 0 {
|
||||
cfg.BatchSize = 50
|
||||
}
|
||||
if cfg.ProgressLogFreq <= 0 {
|
||||
cfg.ProgressLogFreq = 100
|
||||
}
|
||||
|
||||
for i := 0; i < len(summaries); i += cfg.BatchSize {
|
||||
// Check context cancellation
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
log.Warn().Int("synced", synced).Int("remaining", len(summaries)-i).Msg("Batch sync cancelled")
|
||||
return synced, errors
|
||||
default:
|
||||
}
|
||||
|
||||
end := min(i+cfg.BatchSize, len(summaries))
|
||||
|
||||
batch := summaries[i:end]
|
||||
var docs []Document
|
||||
|
||||
// Collect all documents for this batch
|
||||
for _, summary := range batch {
|
||||
docs = append(docs, s.formatSummaryDocs(summary)...)
|
||||
}
|
||||
|
||||
// Add all documents in one call
|
||||
if len(docs) > 0 {
|
||||
if err := s.client.AddDocuments(ctx, docs); err != nil {
|
||||
log.Warn().Err(err).Int("batchStart", i).Int("batchSize", len(batch)).Msg("Failed to sync summary batch")
|
||||
errors += len(batch)
|
||||
continue
|
||||
}
|
||||
}
|
||||
|
||||
synced += len(batch)
|
||||
|
||||
// Log progress periodically
|
||||
if synced%cfg.ProgressLogFreq == 0 || synced == len(summaries) {
|
||||
log.Debug().Int("synced", synced).Int("total", len(summaries)).Msg("Summary batch sync progress")
|
||||
}
|
||||
}
|
||||
|
||||
return synced, errors
|
||||
}
|
||||
|
||||
// BatchSyncPrompts syncs multiple user prompts efficiently in batches.
|
||||
func (s *Sync) BatchSyncPrompts(ctx context.Context, prompts []*models.UserPromptWithSession, cfg BatchSyncConfig) (synced int, errors int) {
|
||||
if len(prompts) == 0 {
|
||||
return 0, 0
|
||||
}
|
||||
|
||||
if cfg.BatchSize <= 0 {
|
||||
cfg.BatchSize = 50
|
||||
}
|
||||
if cfg.ProgressLogFreq <= 0 {
|
||||
cfg.ProgressLogFreq = 100
|
||||
}
|
||||
|
||||
for i := 0; i < len(prompts); i += cfg.BatchSize {
|
||||
// Check context cancellation
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
log.Warn().Int("synced", synced).Int("remaining", len(prompts)-i).Msg("Batch sync cancelled")
|
||||
return synced, errors
|
||||
default:
|
||||
}
|
||||
|
||||
end := min(i+cfg.BatchSize, len(prompts))
|
||||
|
||||
batch := prompts[i:end]
|
||||
docs := make([]Document, 0, len(batch))
|
||||
|
||||
// Collect all documents for this batch
|
||||
for _, prompt := range batch {
|
||||
docs = append(docs, Document{
|
||||
ID: fmt.Sprintf("prompt_%d", prompt.ID),
|
||||
Content: prompt.PromptText,
|
||||
Metadata: map[string]any{
|
||||
"sqlite_id": prompt.ID,
|
||||
"doc_type": "user_prompt",
|
||||
"sdk_session_id": prompt.SDKSessionID,
|
||||
"project": prompt.Project,
|
||||
"scope": "",
|
||||
"created_at_epoch": prompt.CreatedAtEpoch,
|
||||
"prompt_number": prompt.PromptNumber,
|
||||
"field_type": "prompt",
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
// Add all documents in one call
|
||||
if len(docs) > 0 {
|
||||
if err := s.client.AddDocuments(ctx, docs); err != nil {
|
||||
log.Warn().Err(err).Int("batchStart", i).Int("batchSize", len(batch)).Msg("Failed to sync prompt batch")
|
||||
errors += len(batch)
|
||||
continue
|
||||
}
|
||||
}
|
||||
|
||||
synced += len(batch)
|
||||
|
||||
// Log progress periodically
|
||||
if synced%cfg.ProgressLogFreq == 0 || synced == len(prompts) {
|
||||
log.Debug().Int("synced", synced).Int("total", len(prompts)).Msg("Prompt batch sync progress")
|
||||
}
|
||||
}
|
||||
|
||||
return synced, errors
|
||||
}
|
||||
|
||||
@@ -16,15 +16,15 @@ import (
|
||||
// Watcher monitors a file or directory for deletion and calls onDelete when removed.
|
||||
// It watches the parent directory since fsnotify cannot watch non-existent files.
|
||||
type Watcher struct {
|
||||
targetPath string // The file/directory to watch for deletion
|
||||
parentPath string // Parent directory (what we actually watch)
|
||||
onDelete func() // Callback when target is deleted
|
||||
watcher *fsnotify.Watcher
|
||||
ctx context.Context
|
||||
onDelete func()
|
||||
watcher *fsnotify.Watcher
|
||||
cancel context.CancelFunc
|
||||
targetPath string
|
||||
parentPath string
|
||||
debounce time.Duration
|
||||
mu sync.Mutex
|
||||
running bool
|
||||
debounce time.Duration
|
||||
}
|
||||
|
||||
// New creates a new Watcher for the given target path.
|
||||
|
||||
+173
-1202
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,677 @@
|
||||
// Package worker provides context and search-related HTTP handlers.
|
||||
package worker
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/lukaszraczylo/claude-mnemonic/internal/db/gorm"
|
||||
"github.com/lukaszraczylo/claude-mnemonic/internal/reranking"
|
||||
"github.com/lukaszraczylo/claude-mnemonic/internal/search/expansion"
|
||||
"github.com/lukaszraczylo/claude-mnemonic/internal/vector/sqlitevec"
|
||||
"github.com/lukaszraczylo/claude-mnemonic/internal/worker/sdk"
|
||||
"github.com/lukaszraczylo/claude-mnemonic/pkg/models"
|
||||
"github.com/rs/zerolog/log"
|
||||
)
|
||||
|
||||
// handleSearchByPrompt searches observations relevant to a user prompt.
|
||||
// IMPORTANT: This is on the critical startup path - must be fast!
|
||||
// No synchronous verification - just filter by staleness and return.
|
||||
func (s *Service) handleSearchByPrompt(w http.ResponseWriter, r *http.Request) {
|
||||
project := r.URL.Query().Get("project")
|
||||
query := r.URL.Query().Get("query")
|
||||
cwd := r.URL.Query().Get("cwd")
|
||||
|
||||
if project == "" || query == "" {
|
||||
http.Error(w, "project and query required", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
|
||||
// Validate project name to prevent path traversal
|
||||
if err := ValidateProjectName(project); err != nil {
|
||||
http.Error(w, err.Error(), http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
|
||||
limit := gorm.ParseLimitParamWithMax(r, DefaultSearchLimit, 200)
|
||||
|
||||
var observations []*models.Observation
|
||||
var err error
|
||||
var usedVector bool
|
||||
similarityScores := make(map[int64]float64) // Track similarity per observation
|
||||
|
||||
// Get threshold settings from config
|
||||
threshold := s.config.ContextRelevanceThreshold
|
||||
maxResults := s.config.ContextMaxPromptResults
|
||||
|
||||
// Generate expanded queries if query expander is available
|
||||
// Use timeout context to prevent query expansion from blocking
|
||||
var expandedQueries []expansion.ExpandedQuery
|
||||
var detectedIntent string
|
||||
if s.queryExpander != nil {
|
||||
expandCtx, expandCancel := context.WithTimeout(r.Context(), 5*time.Second)
|
||||
cfg := expansion.DefaultConfig()
|
||||
cfg.EnableVocabularyExpansion = false // Vocabulary expansion is optional
|
||||
expandedQueries = s.queryExpander.Expand(expandCtx, query, cfg)
|
||||
expandCancel() // Cancel immediately after use (defer not needed - no panic possible between creation and here)
|
||||
if len(expandedQueries) > 0 {
|
||||
detectedIntent = string(expandedQueries[0].Intent)
|
||||
}
|
||||
}
|
||||
if len(expandedQueries) == 0 {
|
||||
// Fallback to just the original query
|
||||
expandedQueries = []expansion.ExpandedQuery{
|
||||
{Query: query, Weight: 1.0, Source: "original"},
|
||||
}
|
||||
}
|
||||
|
||||
// Try vector search first if available
|
||||
var vectorSearchFailed bool
|
||||
if s.vectorClient != nil && s.vectorClient.IsConnected() {
|
||||
where := sqlitevec.BuildWhereFilter(sqlitevec.DocTypeObservation, "")
|
||||
|
||||
// Search with each expanded query and merge results
|
||||
// Pre-allocate with estimated capacity to avoid repeated reallocation
|
||||
estimatedCapacity := len(expandedQueries) * limit * 2
|
||||
allVectorResults := make([]sqlitevec.QueryResult, 0, estimatedCapacity)
|
||||
queryWeights := make(map[string]float64, len(expandedQueries))
|
||||
var vectorErrors int
|
||||
|
||||
for _, eq := range expandedQueries {
|
||||
vectorResults, vecErr := s.vectorClient.Query(r.Context(), eq.Query, limit*2, where)
|
||||
if vecErr != nil {
|
||||
vectorErrors++
|
||||
log.Debug().Err(vecErr).Str("query", eq.Query).Msg("Vector query failed")
|
||||
} else if len(vectorResults) > 0 {
|
||||
// Apply weight to similarity scores before merging
|
||||
for i := range vectorResults {
|
||||
vectorResults[i].Similarity *= eq.Weight
|
||||
}
|
||||
allVectorResults = append(allVectorResults, vectorResults...)
|
||||
queryWeights[eq.Query] = eq.Weight
|
||||
}
|
||||
}
|
||||
|
||||
// Track if vector search had issues
|
||||
if vectorErrors > 0 && vectorErrors == len(expandedQueries) {
|
||||
vectorSearchFailed = true
|
||||
log.Warn().Int("errors", vectorErrors).Str("project", project).Msg("All vector queries failed, falling back to FTS")
|
||||
}
|
||||
|
||||
if len(allVectorResults) > 0 {
|
||||
// Filter by relevance threshold before extracting IDs
|
||||
// Use a slightly lower threshold for expanded queries
|
||||
effectiveThreshold := threshold * 0.9 // Allow slightly lower scores for expanded queries
|
||||
filteredResults := sqlitevec.FilterByThreshold(allVectorResults, effectiveThreshold, 0)
|
||||
|
||||
// Build similarity map for filtered results (keeping highest weighted score per observation)
|
||||
for _, vr := range filteredResults {
|
||||
if sqliteID, ok := vr.Metadata["sqlite_id"].(float64); ok {
|
||||
id := int64(sqliteID)
|
||||
// Keep the highest score for each observation
|
||||
if existing, exists := similarityScores[id]; !exists || vr.Similarity > existing {
|
||||
similarityScores[id] = vr.Similarity
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Extract observation IDs with project/scope filtering using shared helper
|
||||
obsIDs := sqlitevec.ExtractObservationIDs(filteredResults, project)
|
||||
|
||||
if len(obsIDs) > 0 {
|
||||
// Fetch full observations from SQLite
|
||||
observations, err = s.observationStore.GetObservationsByIDs(r.Context(), obsIDs, "date_desc", limit)
|
||||
if err == nil {
|
||||
usedVector = true
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Fall back to FTS if vector search not available, failed, or returned no results
|
||||
if !usedVector || len(observations) == 0 {
|
||||
if vectorSearchFailed {
|
||||
log.Info().Str("project", project).Msg("Using FTS fallback due to vector search failure")
|
||||
}
|
||||
observations, err = s.observationStore.SearchObservationsFTS(r.Context(), query, project, limit)
|
||||
if err != nil {
|
||||
// FTS might fail if query has special chars, try without
|
||||
log.Warn().Err(err).Str("query", query).Msg("FTS search failed, falling back to recent")
|
||||
observations, err = s.observationStore.GetRecentObservations(r.Context(), project, limit)
|
||||
if err != nil {
|
||||
http.Error(w, err.Error(), http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Fast staleness filter - NO verification (that's too slow for interactive use)
|
||||
// Just check mtimes and exclude obviously stale observations
|
||||
var staleCount int
|
||||
freshObservations := make([]*models.Observation, 0, len(observations))
|
||||
|
||||
for _, obs := range observations {
|
||||
if len(obs.FileMtimes) > 0 && cwd != "" {
|
||||
var paths []string
|
||||
for path := range obs.FileMtimes {
|
||||
paths = append(paths, path)
|
||||
}
|
||||
currentMtimes := sdk.GetFileMtimes(paths, cwd)
|
||||
|
||||
if obs.CheckStaleness(currentMtimes) {
|
||||
// Stale - exclude but don't verify (too slow)
|
||||
// Queue for background verification instead
|
||||
staleCount++
|
||||
s.queueStaleVerification(obs.ID, cwd)
|
||||
continue
|
||||
}
|
||||
}
|
||||
freshObservations = append(freshObservations, obs)
|
||||
}
|
||||
|
||||
// Apply cross-encoder reranking if available
|
||||
var reranked bool
|
||||
if s.reranker != nil && len(freshObservations) > 0 && usedVector {
|
||||
// Build candidates from observations with their bi-encoder scores
|
||||
candidates := make([]reranking.Candidate, len(freshObservations))
|
||||
for i, obs := range freshObservations {
|
||||
// Use strings.Builder for efficient concatenation
|
||||
var content string
|
||||
if obs.Narrative.Valid && obs.Narrative.String != "" {
|
||||
var sb strings.Builder
|
||||
sb.Grow(len(obs.Title.String) + 1 + len(obs.Narrative.String))
|
||||
sb.WriteString(obs.Title.String)
|
||||
sb.WriteByte(' ')
|
||||
sb.WriteString(obs.Narrative.String)
|
||||
content = sb.String()
|
||||
} else {
|
||||
content = obs.Title.String
|
||||
}
|
||||
candidates[i] = reranking.Candidate{
|
||||
ID: strconv.FormatInt(obs.ID, 10), // Faster than fmt.Sprintf
|
||||
Content: content,
|
||||
Score: similarityScores[obs.ID],
|
||||
Metadata: map[string]any{"obs_idx": i},
|
||||
}
|
||||
}
|
||||
|
||||
// Rerank using cross-encoder - use pure mode or combined scores
|
||||
var rerankResults []reranking.RerankResult
|
||||
var rerankErr error
|
||||
if s.config.RerankingPureMode {
|
||||
rerankResults, rerankErr = s.reranker.RerankByScore(query, candidates, s.config.RerankingResults)
|
||||
} else {
|
||||
rerankResults, rerankErr = s.reranker.Rerank(query, candidates, s.config.RerankingResults)
|
||||
}
|
||||
if rerankErr != nil {
|
||||
log.Warn().Err(rerankErr).Msg("Cross-encoder reranking failed, using original order")
|
||||
} else if len(rerankResults) > 0 {
|
||||
// Update similarity scores with reranked scores
|
||||
for _, rr := range rerankResults {
|
||||
if id, err := strconv.ParseInt(rr.ID, 10, 64); err == nil {
|
||||
similarityScores[id] = rr.CombinedScore
|
||||
}
|
||||
}
|
||||
|
||||
// Reorder observations based on rerank results
|
||||
reorderedObs := make([]*models.Observation, 0, len(rerankResults))
|
||||
obsMap := make(map[int64]*models.Observation)
|
||||
for _, obs := range freshObservations {
|
||||
obsMap[obs.ID] = obs
|
||||
}
|
||||
for _, rr := range rerankResults {
|
||||
if id, err := strconv.ParseInt(rr.ID, 10, 64); err == nil {
|
||||
if obs, ok := obsMap[id]; ok {
|
||||
reorderedObs = append(reorderedObs, obs)
|
||||
}
|
||||
}
|
||||
}
|
||||
freshObservations = reorderedObs
|
||||
reranked = true
|
||||
|
||||
log.Debug().
|
||||
Int("candidates", len(candidates)).
|
||||
Int("returned", len(rerankResults)).
|
||||
Msg("Cross-encoder reranking complete")
|
||||
}
|
||||
}
|
||||
|
||||
// Cluster similar observations to remove duplicates
|
||||
clusteredObservations := clusterObservations(freshObservations, 0.4)
|
||||
duplicatesRemoved := len(freshObservations) - len(clusteredObservations)
|
||||
|
||||
// Sort by similarity score (highest first) if we have scores and didn't rerank
|
||||
if len(similarityScores) > 0 && len(clusteredObservations) > 0 && !reranked {
|
||||
sort.Slice(clusteredObservations, func(i, j int) bool {
|
||||
scoreI := similarityScores[clusteredObservations[i].ID]
|
||||
scoreJ := similarityScores[clusteredObservations[j].ID]
|
||||
return scoreI > scoreJ
|
||||
})
|
||||
}
|
||||
|
||||
// Apply max results cap if configured
|
||||
if maxResults > 0 && len(clusteredObservations) > maxResults {
|
||||
clusteredObservations = clusteredObservations[:maxResults]
|
||||
}
|
||||
|
||||
// Record retrieval stats with staleness metrics
|
||||
s.recordRetrievalStatsExtended(project, int64(len(clusteredObservations)), 0, 0,
|
||||
int64(staleCount), int64(len(freshObservations)), int64(duplicatesRemoved), true)
|
||||
|
||||
// Increment retrieval counts for scoring (async, non-blocking)
|
||||
if len(clusteredObservations) > 0 {
|
||||
ids := make([]int64, len(clusteredObservations))
|
||||
for i, obs := range clusteredObservations {
|
||||
ids[i] = obs.ID
|
||||
}
|
||||
s.incrementRetrievalCounts(ids)
|
||||
}
|
||||
|
||||
log.Info().
|
||||
Str("project", project).
|
||||
Str("query", query).
|
||||
Str("intent", detectedIntent).
|
||||
Int("expansions", len(expandedQueries)).
|
||||
Int("found", len(clusteredObservations)).
|
||||
Int("stale_excluded", staleCount).
|
||||
Float64("threshold", threshold).
|
||||
Msg("Prompt-based observation search")
|
||||
|
||||
// Build response with similarity scores
|
||||
obsWithScores := make([]map[string]any, len(clusteredObservations))
|
||||
for i, obs := range clusteredObservations {
|
||||
obsMap := obs.ToMap()
|
||||
if score, ok := similarityScores[obs.ID]; ok {
|
||||
obsMap["similarity"] = score
|
||||
}
|
||||
obsWithScores[i] = obsMap
|
||||
}
|
||||
|
||||
// Build expansion info for response
|
||||
expansionInfo := make([]map[string]any, len(expandedQueries))
|
||||
for i, eq := range expandedQueries {
|
||||
expansionInfo[i] = map[string]any{
|
||||
"query": eq.Query,
|
||||
"weight": eq.Weight,
|
||||
"source": eq.Source,
|
||||
}
|
||||
}
|
||||
|
||||
// Track this search for analytics
|
||||
s.trackSearchQuery(query, project, "observations", len(clusteredObservations), usedVector)
|
||||
|
||||
writeJSON(w, map[string]any{
|
||||
"project": project,
|
||||
"query": query,
|
||||
"intent": detectedIntent,
|
||||
"expansions": expansionInfo,
|
||||
"observations": obsWithScores,
|
||||
"threshold": threshold,
|
||||
"max_results": maxResults,
|
||||
})
|
||||
}
|
||||
|
||||
// handleFileContext returns observations relevant to specific files being worked on.
|
||||
// Uses vector similarity search to find observations that mention or relate to the given files.
|
||||
func (s *Service) handleFileContext(w http.ResponseWriter, r *http.Request) {
|
||||
project := r.URL.Query().Get("project")
|
||||
if project == "" {
|
||||
http.Error(w, "project required", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
|
||||
filesParam := r.URL.Query().Get("files")
|
||||
if filesParam == "" {
|
||||
http.Error(w, "files required", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
|
||||
// Parse comma-separated file paths
|
||||
files := strings.Split(filesParam, ",")
|
||||
if len(files) == 0 {
|
||||
http.Error(w, "at least one file required", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
|
||||
// Limit to reasonable number of files
|
||||
maxFiles := 20
|
||||
if len(files) > maxFiles {
|
||||
files = files[:maxFiles]
|
||||
}
|
||||
|
||||
// Get limit parameter (default 10 per file)
|
||||
limitStr := r.URL.Query().Get("limit")
|
||||
limit := 10
|
||||
if limitStr != "" {
|
||||
if parsed, err := strconv.Atoi(limitStr); err == nil && parsed > 0 && parsed <= 50 {
|
||||
limit = parsed
|
||||
}
|
||||
}
|
||||
|
||||
// Search for observations related to each file in parallel
|
||||
ctx := r.Context()
|
||||
|
||||
// Check if vector search is available
|
||||
if s.vectorClient == nil || !s.vectorClient.IsConnected() {
|
||||
writeJSON(w, map[string]any{
|
||||
"files": files,
|
||||
"results": map[string]any{},
|
||||
"count": 0,
|
||||
"error": "vector search not available",
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
// Prepare for parallel execution
|
||||
type fileResult struct {
|
||||
file string
|
||||
results []map[string]any
|
||||
obsIDs []int64 // Track observation IDs for deduplication
|
||||
}
|
||||
|
||||
resultsChan := make(chan fileResult, len(files))
|
||||
sem := make(chan struct{}, 5) // Limit concurrency to 5 parallel searches
|
||||
var wg sync.WaitGroup
|
||||
|
||||
for _, file := range files {
|
||||
file = strings.TrimSpace(file)
|
||||
if file == "" {
|
||||
continue
|
||||
}
|
||||
|
||||
wg.Add(1)
|
||||
go func(file string) {
|
||||
defer wg.Done()
|
||||
sem <- struct{}{} // Acquire semaphore
|
||||
defer func() { <-sem }() // Release semaphore
|
||||
|
||||
// Build search query from file path
|
||||
query := buildFileQuery(file)
|
||||
|
||||
where := sqlitevec.BuildWhereFilter(sqlitevec.DocTypeObservation, "")
|
||||
vectorResults, vecErr := s.vectorClient.Query(ctx, query, limit*2, where)
|
||||
if vecErr != nil {
|
||||
log.Warn().Err(vecErr).Str("file", file).Msg("Vector search failed for file context")
|
||||
return
|
||||
}
|
||||
|
||||
// Extract observation IDs from vector results
|
||||
obsIDs := sqlitevec.ExtractObservationIDs(vectorResults, project)
|
||||
if len(obsIDs) == 0 {
|
||||
return
|
||||
}
|
||||
|
||||
// Fetch observations
|
||||
observations, err := s.observationStore.GetObservationsByIDs(ctx, obsIDs, "score_desc", limit*2)
|
||||
if err != nil {
|
||||
log.Warn().Err(err).Str("file", file).Msg("Failed to fetch observations for file context")
|
||||
return
|
||||
}
|
||||
|
||||
// Pre-build score map from vector results (O(n) instead of O(n²))
|
||||
scoreMap := make(map[int64]float64, len(vectorResults))
|
||||
var avgScore float64
|
||||
for _, vr := range vectorResults {
|
||||
avgScore += vr.Similarity
|
||||
// Parse observation ID from vector result ID (format: obs_{id}_{field})
|
||||
// Use index-based parsing to avoid slice allocation from strings.Split
|
||||
if len(vr.ID) > 4 && vr.ID[:4] == "obs_" {
|
||||
rest := vr.ID[4:] // Skip "obs_"
|
||||
underscoreIdx := strings.IndexByte(rest, '_')
|
||||
var idStr string
|
||||
if underscoreIdx >= 0 {
|
||||
idStr = rest[:underscoreIdx]
|
||||
} else {
|
||||
idStr = rest
|
||||
}
|
||||
if id, parseErr := strconv.ParseInt(idStr, 10, 64); parseErr == nil {
|
||||
// Keep highest score for each observation
|
||||
if existing, exists := scoreMap[id]; !exists || vr.Similarity > existing {
|
||||
scoreMap[id] = vr.Similarity
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
if len(vectorResults) > 0 {
|
||||
avgScore /= float64(len(vectorResults))
|
||||
}
|
||||
|
||||
fileResults := make([]map[string]any, 0, limit)
|
||||
var usedIDs []int64
|
||||
for _, obs := range observations {
|
||||
// Check project scope
|
||||
if obs.Scope == "project" && obs.Project != project {
|
||||
continue
|
||||
}
|
||||
|
||||
// O(1) score lookup instead of O(n) nested loop
|
||||
score, found := scoreMap[obs.ID]
|
||||
if !found {
|
||||
// Use average score as fallback
|
||||
score = avgScore
|
||||
}
|
||||
|
||||
// Only include if score is above threshold
|
||||
if score < 0.3 {
|
||||
continue
|
||||
}
|
||||
|
||||
fileResults = append(fileResults, map[string]any{
|
||||
"id": obs.ID,
|
||||
"title": obs.Title.String,
|
||||
"type": obs.Type,
|
||||
"narrative": obs.Narrative.String,
|
||||
"facts": obs.Facts,
|
||||
"score": score,
|
||||
})
|
||||
usedIDs = append(usedIDs, obs.ID)
|
||||
|
||||
if len(fileResults) >= limit {
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
if len(fileResults) > 0 {
|
||||
resultsChan <- fileResult{file: file, results: fileResults, obsIDs: usedIDs}
|
||||
}
|
||||
}(file)
|
||||
}
|
||||
|
||||
// Close channel when all goroutines complete
|
||||
go func() {
|
||||
wg.Wait()
|
||||
close(resultsChan)
|
||||
}()
|
||||
|
||||
// Collect results and deduplicate
|
||||
allResults := make(map[string]any)
|
||||
seenObservationIDs := make(map[int64]bool)
|
||||
|
||||
for res := range resultsChan {
|
||||
// Filter out duplicates that were already seen in other files
|
||||
dedupedResults := make([]map[string]any, 0, len(res.results))
|
||||
for i, r := range res.results {
|
||||
obsID := res.obsIDs[i]
|
||||
if !seenObservationIDs[obsID] {
|
||||
seenObservationIDs[obsID] = true
|
||||
dedupedResults = append(dedupedResults, r)
|
||||
}
|
||||
}
|
||||
if len(dedupedResults) > 0 {
|
||||
allResults[res.file] = dedupedResults
|
||||
}
|
||||
}
|
||||
|
||||
writeJSON(w, map[string]any{
|
||||
"files": files,
|
||||
"results": allResults,
|
||||
"count": len(allResults),
|
||||
})
|
||||
}
|
||||
|
||||
// buildFileQuery extracts meaningful search terms from a file path.
|
||||
func buildFileQuery(filePath string) string {
|
||||
// Remove common prefixes and extensions
|
||||
path := strings.TrimPrefix(filePath, "/")
|
||||
|
||||
// Extract the filename and directory
|
||||
parts := strings.Split(path, "/")
|
||||
meaningful := make([]string, 0, len(parts))
|
||||
|
||||
for _, part := range parts {
|
||||
// Skip common directory names that aren't meaningful
|
||||
switch strings.ToLower(part) {
|
||||
case "src", "lib", "internal", "pkg", "cmd", "api", "app", "test", "tests", "spec", "specs":
|
||||
continue
|
||||
default:
|
||||
// Remove file extension
|
||||
if idx := strings.LastIndex(part, "."); idx > 0 {
|
||||
part = part[:idx]
|
||||
}
|
||||
// Convert camelCase/PascalCase to spaces
|
||||
part = splitCamelCase(part)
|
||||
// Convert snake_case to spaces
|
||||
part = strings.ReplaceAll(part, "_", " ")
|
||||
// Convert kebab-case to spaces
|
||||
part = strings.ReplaceAll(part, "-", " ")
|
||||
meaningful = append(meaningful, part)
|
||||
}
|
||||
}
|
||||
|
||||
return strings.Join(meaningful, " ")
|
||||
}
|
||||
|
||||
// splitCamelCase splits camelCase or PascalCase into separate words.
|
||||
func splitCamelCase(s string) string {
|
||||
var result strings.Builder
|
||||
for i, r := range s {
|
||||
if i > 0 && r >= 'A' && r <= 'Z' {
|
||||
result.WriteRune(' ')
|
||||
}
|
||||
result.WriteRune(r)
|
||||
}
|
||||
return result.String()
|
||||
}
|
||||
|
||||
// handleContextInject returns context for injection at session start.
|
||||
// IMPORTANT: This is on the critical startup path - must be fast!
|
||||
// No synchronous verification - just filter by staleness and return.
|
||||
func (s *Service) handleContextInject(w http.ResponseWriter, r *http.Request) {
|
||||
project := r.URL.Query().Get("project")
|
||||
if project == "" {
|
||||
http.Error(w, "project required", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
|
||||
// Validate project name to prevent path traversal
|
||||
if err := ValidateProjectName(project); err != nil {
|
||||
http.Error(w, err.Error(), http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
|
||||
cwd := r.URL.Query().Get("cwd")
|
||||
if cwd == "" {
|
||||
cwd = "/"
|
||||
}
|
||||
|
||||
// Limit observations for fast startup (configurable, default 100)
|
||||
limit := s.config.ContextObservations
|
||||
if limit <= 0 {
|
||||
limit = DefaultContextLimit
|
||||
}
|
||||
|
||||
// Full count determines how many observations get full detail (configurable, default 25)
|
||||
fullCount := s.config.ContextFullCount
|
||||
if fullCount <= 0 {
|
||||
fullCount = 25
|
||||
}
|
||||
|
||||
// Get recent observations
|
||||
observations, err := s.observationStore.GetRecentObservations(r.Context(), project, limit)
|
||||
if err != nil {
|
||||
http.Error(w, err.Error(), http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
|
||||
// Fast staleness filter - NO verification (that's too slow for startup)
|
||||
var staleCount int
|
||||
freshObservations := make([]*models.Observation, 0, len(observations))
|
||||
|
||||
for _, obs := range observations {
|
||||
if len(obs.FileMtimes) > 0 {
|
||||
var paths []string
|
||||
for path := range obs.FileMtimes {
|
||||
paths = append(paths, path)
|
||||
}
|
||||
currentMtimes := sdk.GetFileMtimes(paths, cwd)
|
||||
|
||||
if obs.CheckStaleness(currentMtimes) {
|
||||
// Stale - exclude but don't verify (too slow)
|
||||
// Queue for background verification instead
|
||||
staleCount++
|
||||
s.queueStaleVerification(obs.ID, cwd)
|
||||
continue
|
||||
}
|
||||
}
|
||||
freshObservations = append(freshObservations, obs)
|
||||
}
|
||||
|
||||
// Cluster similar observations to remove duplicates
|
||||
clusteredObservations := clusterObservations(freshObservations, 0.4)
|
||||
duplicatesRemoved := len(freshObservations) - len(clusteredObservations)
|
||||
|
||||
// Record retrieval stats with staleness metrics
|
||||
s.recordRetrievalStatsExtended(project, int64(len(clusteredObservations)), 0, 0,
|
||||
int64(staleCount), int64(len(freshObservations)), int64(duplicatesRemoved), false)
|
||||
|
||||
// Increment retrieval counts for scoring (async, non-blocking)
|
||||
if len(clusteredObservations) > 0 {
|
||||
ids := make([]int64, len(clusteredObservations))
|
||||
for i, obs := range clusteredObservations {
|
||||
ids[i] = obs.ID
|
||||
}
|
||||
s.incrementRetrievalCounts(ids)
|
||||
}
|
||||
|
||||
log.Info().
|
||||
Str("project", project).
|
||||
Int("total", len(observations)).
|
||||
Int("fresh", len(freshObservations)).
|
||||
Int("clustered", len(clusteredObservations)).
|
||||
Int("duplicates", duplicatesRemoved).
|
||||
Int("stale_excluded", staleCount).
|
||||
Msg("Context injection with clustering")
|
||||
|
||||
writeJSON(w, map[string]any{
|
||||
"project": project,
|
||||
"observations": clusteredObservations,
|
||||
"full_count": fullCount,
|
||||
"stale_excluded": staleCount,
|
||||
"duplicates_removed": duplicatesRemoved,
|
||||
})
|
||||
}
|
||||
|
||||
// handleContextCount returns the count of observations for a project.
|
||||
func (s *Service) handleContextCount(w http.ResponseWriter, r *http.Request) {
|
||||
project := r.URL.Query().Get("project")
|
||||
if project == "" {
|
||||
http.Error(w, "project required", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
|
||||
count, err := s.getCachedObservationCount(r.Context(), project)
|
||||
if err != nil {
|
||||
http.Error(w, err.Error(), http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
|
||||
writeJSON(w, map[string]any{
|
||||
"project": project,
|
||||
"count": count,
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,710 @@
|
||||
// Package worker provides data retrieval HTTP handlers.
|
||||
package worker
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"runtime"
|
||||
"sort"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/lukaszraczylo/claude-mnemonic/internal/db/gorm"
|
||||
"github.com/lukaszraczylo/claude-mnemonic/internal/embedding"
|
||||
"github.com/lukaszraczylo/claude-mnemonic/internal/vector/sqlitevec"
|
||||
"github.com/lukaszraczylo/claude-mnemonic/pkg/models"
|
||||
"github.com/rs/zerolog/log"
|
||||
)
|
||||
|
||||
// handleGetObservations returns recent observations.
|
||||
// Supports optional query parameter for semantic search via sqlite-vec.
|
||||
// Supports pagination via limit and offset query parameters.
|
||||
func (s *Service) handleGetObservations(w http.ResponseWriter, r *http.Request) {
|
||||
pagination := gorm.ParsePaginationParams(r, DefaultObservationsLimit)
|
||||
project := r.URL.Query().Get("project")
|
||||
query := r.URL.Query().Get("query")
|
||||
|
||||
// Validate project name to prevent path traversal
|
||||
if err := ValidateProjectName(project); err != nil {
|
||||
http.Error(w, err.Error(), http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
|
||||
var observations []*models.Observation
|
||||
var total int64
|
||||
var err error
|
||||
var usedVector bool
|
||||
|
||||
// Use vector search if query is provided and vector client is available
|
||||
if query != "" && s.vectorClient != nil && s.vectorClient.IsConnected() {
|
||||
where := sqlitevec.BuildWhereFilter(sqlitevec.DocTypeObservation, "")
|
||||
vectorResults, vecErr := s.vectorClient.Query(r.Context(), query, pagination.Limit*2, where)
|
||||
if vecErr == nil && len(vectorResults) > 0 {
|
||||
obsIDs := sqlitevec.ExtractObservationIDs(vectorResults, project)
|
||||
if len(obsIDs) > 0 {
|
||||
observations, err = s.observationStore.GetObservationsByIDs(r.Context(), obsIDs, "date_desc", pagination.Limit)
|
||||
if err == nil {
|
||||
usedVector = true
|
||||
total = int64(len(observations)) // Vector search doesn't have total, use returned count
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Fall back to SQLite if vector search not used
|
||||
if !usedVector {
|
||||
if project != "" {
|
||||
// Strict project filtering for dashboard - only observations from this project
|
||||
observations, total, err = s.observationStore.GetObservationsByProjectStrictPaginated(r.Context(), project, pagination.Limit, pagination.Offset)
|
||||
} else {
|
||||
// All projects
|
||||
observations, total, err = s.observationStore.GetAllRecentObservationsPaginated(r.Context(), pagination.Limit, pagination.Offset)
|
||||
}
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
http.Error(w, err.Error(), http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
|
||||
// Ensure we return empty array, not null
|
||||
if observations == nil {
|
||||
observations = []*models.Observation{}
|
||||
}
|
||||
|
||||
// Track search if query was provided
|
||||
if query != "" {
|
||||
s.trackSearchQuery(query, project, "observations", len(observations), usedVector)
|
||||
}
|
||||
|
||||
// Return paginated response
|
||||
writeJSON(w, map[string]any{
|
||||
"observations": observations,
|
||||
"total": total,
|
||||
"limit": pagination.Limit,
|
||||
"offset": pagination.Offset,
|
||||
"hasMore": int64(pagination.Offset)+int64(len(observations)) < total,
|
||||
})
|
||||
}
|
||||
|
||||
// handleGetSummaries returns recent summaries.
|
||||
// Supports optional query parameter for semantic search via sqlite-vec.
|
||||
func (s *Service) handleGetSummaries(w http.ResponseWriter, r *http.Request) {
|
||||
limit := gorm.ParseLimitParam(r, DefaultSummariesLimit)
|
||||
project := r.URL.Query().Get("project")
|
||||
query := r.URL.Query().Get("query")
|
||||
|
||||
// Validate project name to prevent path traversal
|
||||
if err := ValidateProjectName(project); err != nil {
|
||||
http.Error(w, err.Error(), http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
|
||||
var summaries []*models.SessionSummary
|
||||
var err error
|
||||
var usedVector bool
|
||||
|
||||
// Use vector search if query is provided and vector client is available
|
||||
if query != "" && s.vectorClient != nil && s.vectorClient.IsConnected() {
|
||||
where := sqlitevec.BuildWhereFilter(sqlitevec.DocTypeSessionSummary, "")
|
||||
vectorResults, vecErr := s.vectorClient.Query(r.Context(), query, limit*2, where)
|
||||
if vecErr == nil && len(vectorResults) > 0 {
|
||||
summaryIDs := sqlitevec.ExtractSummaryIDs(vectorResults, project)
|
||||
if len(summaryIDs) > 0 {
|
||||
summaries, err = s.summaryStore.GetSummariesByIDs(r.Context(), summaryIDs, "date_desc", limit)
|
||||
if err == nil {
|
||||
usedVector = true
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Fall back to SQLite if vector search not used
|
||||
if !usedVector {
|
||||
if project != "" {
|
||||
summaries, err = s.summaryStore.GetRecentSummaries(r.Context(), project, limit)
|
||||
} else {
|
||||
summaries, err = s.summaryStore.GetAllRecentSummaries(r.Context(), limit)
|
||||
}
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
http.Error(w, err.Error(), http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
|
||||
// Ensure we return empty array, not null
|
||||
if summaries == nil {
|
||||
summaries = []*models.SessionSummary{}
|
||||
}
|
||||
writeJSON(w, summaries)
|
||||
}
|
||||
|
||||
// handleGetPrompts returns recent user prompts.
|
||||
// Supports optional query parameter for semantic search via sqlite-vec.
|
||||
func (s *Service) handleGetPrompts(w http.ResponseWriter, r *http.Request) {
|
||||
limit := gorm.ParseLimitParam(r, DefaultPromptsLimit)
|
||||
project := r.URL.Query().Get("project")
|
||||
query := r.URL.Query().Get("query")
|
||||
|
||||
// Validate project name to prevent path traversal
|
||||
if err := ValidateProjectName(project); err != nil {
|
||||
http.Error(w, err.Error(), http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
|
||||
var prompts []*models.UserPromptWithSession
|
||||
var err error
|
||||
var usedVector bool
|
||||
|
||||
// Use vector search if query is provided and vector client is available
|
||||
if query != "" && s.vectorClient != nil && s.vectorClient.IsConnected() {
|
||||
where := sqlitevec.BuildWhereFilter(sqlitevec.DocTypeUserPrompt, "")
|
||||
vectorResults, vecErr := s.vectorClient.Query(r.Context(), query, limit*2, where)
|
||||
if vecErr == nil && len(vectorResults) > 0 {
|
||||
promptIDs := sqlitevec.ExtractPromptIDs(vectorResults, project)
|
||||
if len(promptIDs) > 0 {
|
||||
prompts, err = s.promptStore.GetPromptsByIDs(r.Context(), promptIDs, "date_desc", limit)
|
||||
if err == nil {
|
||||
usedVector = true
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Fall back to SQLite if vector search not used
|
||||
if !usedVector {
|
||||
if project != "" {
|
||||
prompts, err = s.promptStore.GetRecentUserPromptsByProject(r.Context(), project, limit)
|
||||
} else {
|
||||
prompts, err = s.promptStore.GetAllRecentUserPrompts(r.Context(), limit)
|
||||
}
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
http.Error(w, err.Error(), http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
|
||||
// Ensure we return empty array, not null
|
||||
if prompts == nil {
|
||||
prompts = []*models.UserPromptWithSession{}
|
||||
}
|
||||
writeJSON(w, prompts)
|
||||
}
|
||||
|
||||
// handleGetProjects returns all projects.
|
||||
// Response is cacheable for 5 minutes since project list changes infrequently.
|
||||
func (s *Service) handleGetProjects(w http.ResponseWriter, r *http.Request) {
|
||||
projects, err := s.sessionStore.GetAllProjects(r.Context())
|
||||
if err != nil {
|
||||
http.Error(w, err.Error(), http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
|
||||
// Cache for 5 minutes - project list changes infrequently
|
||||
w.Header().Set("Cache-Control", "public, max-age=300")
|
||||
writeJSON(w, projects)
|
||||
}
|
||||
|
||||
// handleGetTypes returns the canonical list of observation and concept types.
|
||||
// This provides a single source of truth for both backend and frontend.
|
||||
// Response is cacheable as these values never change at runtime.
|
||||
func (s *Service) handleGetTypes(w http.ResponseWriter, r *http.Request) {
|
||||
// Cache for 24 hours - these values are compile-time constants
|
||||
w.Header().Set("Cache-Control", "public, max-age=86400")
|
||||
writeJSON(w, map[string]any{
|
||||
"observation_types": ObservationTypes,
|
||||
"concept_types": ConceptTypes,
|
||||
})
|
||||
}
|
||||
|
||||
// handleGetModels returns available embedding models.
|
||||
// Response is cacheable as model list doesn't change without restart.
|
||||
func (s *Service) handleGetModels(w http.ResponseWriter, _ *http.Request) {
|
||||
// Cache for 1 hour - model list is static during runtime
|
||||
w.Header().Set("Cache-Control", "public, max-age=3600")
|
||||
|
||||
models := embedding.ListModels()
|
||||
defaultModel := embedding.GetDefaultModel()
|
||||
|
||||
writeJSON(w, map[string]any{
|
||||
"models": models,
|
||||
"default": defaultModel,
|
||||
"current": s.embedSvc.Version(),
|
||||
})
|
||||
}
|
||||
|
||||
// handleGetStats returns worker statistics.
|
||||
func (s *Service) handleGetStats(w http.ResponseWriter, r *http.Request) {
|
||||
project := r.URL.Query().Get("project")
|
||||
|
||||
// Validate project name to prevent path traversal
|
||||
if err := ValidateProjectName(project); err != nil {
|
||||
http.Error(w, err.Error(), http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
|
||||
retrievalStats := s.GetRetrievalStats(project)
|
||||
sessionsToday, _ := s.sessionStore.GetSessionsToday(r.Context())
|
||||
|
||||
response := map[string]any{
|
||||
"uptime": time.Since(s.startTime).String(),
|
||||
"uptimeSeconds": time.Since(s.startTime).Seconds(),
|
||||
"activeSessions": s.sessionManager.GetActiveSessionCount(),
|
||||
"queueDepth": s.sessionManager.GetTotalQueueDepth(),
|
||||
"isProcessing": s.sessionManager.IsAnySessionProcessing(),
|
||||
"connectedClients": s.sseBroadcaster.ClientCount(),
|
||||
"sessionsToday": sessionsToday,
|
||||
"retrieval": retrievalStats,
|
||||
"ready": s.ready.Load(),
|
||||
}
|
||||
|
||||
// Add memory stats
|
||||
var memStats runtime.MemStats
|
||||
runtime.ReadMemStats(&memStats)
|
||||
response["memory"] = map[string]any{
|
||||
"alloc_mb": float64(memStats.Alloc) / 1024 / 1024,
|
||||
"total_alloc_mb": float64(memStats.TotalAlloc) / 1024 / 1024,
|
||||
"sys_mb": float64(memStats.Sys) / 1024 / 1024,
|
||||
"heap_alloc_mb": float64(memStats.HeapAlloc) / 1024 / 1024,
|
||||
"heap_inuse_mb": float64(memStats.HeapInuse) / 1024 / 1024,
|
||||
"heap_objects": memStats.HeapObjects,
|
||||
"goroutines": runtime.NumGoroutine(),
|
||||
"gc_cycles": memStats.NumGC,
|
||||
"gc_pause_total_ms": float64(memStats.PauseTotalNs) / 1e6,
|
||||
}
|
||||
|
||||
// Add database health if available
|
||||
if s.store != nil {
|
||||
dbHealth := s.store.HealthCheck(r.Context())
|
||||
response["database"] = map[string]any{
|
||||
"status": dbHealth.Status,
|
||||
"query_latency_ms": float64(dbHealth.QueryLatency) / 1e6,
|
||||
"pool": dbHealth.PoolStats,
|
||||
"warning": dbHealth.Warning,
|
||||
}
|
||||
}
|
||||
|
||||
// Add embedding model info
|
||||
if s.embedSvc != nil {
|
||||
response["embeddingModel"] = map[string]any{
|
||||
"name": s.embedSvc.Name(),
|
||||
"version": s.embedSvc.Version(),
|
||||
"dimensions": s.embedSvc.Dimensions(),
|
||||
}
|
||||
}
|
||||
|
||||
// Add vector cache stats
|
||||
if s.vectorClient != nil {
|
||||
if count, err := s.vectorClient.Count(r.Context()); err == nil {
|
||||
response["vectorCount"] = count
|
||||
}
|
||||
cacheSize, cacheMax := s.vectorClient.CacheStats()
|
||||
response["vectorCache"] = map[string]any{
|
||||
"size": cacheSize,
|
||||
"max_size": cacheMax,
|
||||
}
|
||||
}
|
||||
|
||||
// Include project-specific observation count if project is specified
|
||||
if project != "" {
|
||||
count, err := s.getCachedObservationCount(r.Context(), project)
|
||||
if err == nil {
|
||||
response["projectObservations"] = count
|
||||
response["project"] = project
|
||||
}
|
||||
}
|
||||
|
||||
// Add rate limiter stats
|
||||
if s.rateLimiter != nil {
|
||||
response["rateLimiter"] = s.rateLimiter.Stats()
|
||||
}
|
||||
|
||||
// Add circuit breaker metrics
|
||||
if s.processor != nil {
|
||||
response["circuitBreaker"] = s.processor.CircuitBreakerMetrics()
|
||||
}
|
||||
|
||||
writeJSON(w, response)
|
||||
}
|
||||
|
||||
// handleGetRetrievalStats returns detailed retrieval statistics.
|
||||
func (s *Service) handleGetRetrievalStats(w http.ResponseWriter, r *http.Request) {
|
||||
project := r.URL.Query().Get("project")
|
||||
stats := s.GetRetrievalStats(project)
|
||||
writeJSON(w, stats)
|
||||
}
|
||||
|
||||
// handleGetRecentQueries returns recent search queries for analytics.
|
||||
func (s *Service) handleGetRecentQueries(w http.ResponseWriter, r *http.Request) {
|
||||
project := r.URL.Query().Get("project")
|
||||
limit := gorm.ParseLimitParam(r, 20)
|
||||
|
||||
queries := s.getRecentSearchQueries(project, limit)
|
||||
|
||||
writeJSON(w, map[string]any{
|
||||
"queries": queries,
|
||||
"count": len(queries),
|
||||
"project": project,
|
||||
})
|
||||
}
|
||||
|
||||
// handleGetSearchAnalytics returns comprehensive search analytics and statistics.
|
||||
func (s *Service) handleGetSearchAnalytics(w http.ResponseWriter, r *http.Request) {
|
||||
project := r.URL.Query().Get("project")
|
||||
|
||||
// Get all recent queries for analysis
|
||||
queries := s.getRecentSearchQueries(project, maxRecentQueries)
|
||||
|
||||
// Calculate analytics
|
||||
totalQueries := len(queries)
|
||||
vectorSearches := 0
|
||||
totalResults := 0
|
||||
zeroResultQueries := 0
|
||||
queryTypes := make(map[string]int)
|
||||
topKeywords := make(map[string]int)
|
||||
|
||||
for _, q := range queries {
|
||||
if q.UsedVector {
|
||||
vectorSearches++
|
||||
}
|
||||
totalResults += q.Results
|
||||
if q.Results == 0 {
|
||||
zeroResultQueries++
|
||||
}
|
||||
queryTypes[q.Type]++
|
||||
|
||||
// Extract keywords (simple word tokenization using iterator)
|
||||
for word := range strings.FieldsSeq(strings.ToLower(q.Query)) {
|
||||
if len(word) > 3 { // Skip short words
|
||||
topKeywords[word]++
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Sort keywords by frequency
|
||||
type keywordCount struct {
|
||||
Keyword string `json:"keyword"`
|
||||
Count int `json:"count"`
|
||||
}
|
||||
sortedKeywords := make([]keywordCount, 0, len(topKeywords))
|
||||
for kw, count := range topKeywords {
|
||||
sortedKeywords = append(sortedKeywords, keywordCount{Keyword: kw, Count: count})
|
||||
}
|
||||
sort.Slice(sortedKeywords, func(i, j int) bool {
|
||||
return sortedKeywords[i].Count > sortedKeywords[j].Count
|
||||
})
|
||||
if len(sortedKeywords) > 10 {
|
||||
sortedKeywords = sortedKeywords[:10]
|
||||
}
|
||||
|
||||
// Calculate averages
|
||||
avgResults := float64(0)
|
||||
vectorSearchRate := float64(0)
|
||||
zeroResultRate := float64(0)
|
||||
if totalQueries > 0 {
|
||||
avgResults = float64(totalResults) / float64(totalQueries)
|
||||
vectorSearchRate = float64(vectorSearches) / float64(totalQueries) * 100
|
||||
zeroResultRate = float64(zeroResultQueries) / float64(totalQueries) * 100
|
||||
}
|
||||
|
||||
writeJSON(w, map[string]any{
|
||||
"total_queries": totalQueries,
|
||||
"vector_search_rate": vectorSearchRate,
|
||||
"avg_results": avgResults,
|
||||
"zero_result_rate": zeroResultRate,
|
||||
"query_types": queryTypes,
|
||||
"top_keywords": sortedKeywords,
|
||||
"project": project,
|
||||
})
|
||||
}
|
||||
|
||||
// handleVectorHealth returns comprehensive health information about the vector database.
|
||||
func (s *Service) handleVectorHealth(w http.ResponseWriter, r *http.Request) {
|
||||
if s.vectorClient == nil {
|
||||
http.Error(w, "vector client not initialized", http.StatusServiceUnavailable)
|
||||
return
|
||||
}
|
||||
|
||||
stats, err := s.vectorClient.GetHealthStats(r.Context())
|
||||
if err != nil {
|
||||
http.Error(w, err.Error(), http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
|
||||
// Add additional computed metrics
|
||||
healthScore := 100.0
|
||||
var warnings []string
|
||||
|
||||
// Penalize for stale vectors
|
||||
if stats.TotalVectors > 0 {
|
||||
staleRatio := float64(stats.StaleVectors) / float64(stats.TotalVectors)
|
||||
if staleRatio > 0 {
|
||||
healthScore -= staleRatio * 50 // Up to 50 points off for stale vectors
|
||||
warnings = append(warnings, formatWarning("%.1f%% vectors need rebuild", staleRatio*100))
|
||||
}
|
||||
}
|
||||
|
||||
// Check cache effectiveness
|
||||
cacheHitRate := stats.EmbeddingCache.HitRate()
|
||||
if cacheHitRate < 20 && (stats.EmbeddingCache.EmbeddingHits+stats.EmbeddingCache.EmbeddingMisses) > 100 {
|
||||
healthScore -= 10
|
||||
warnings = append(warnings, formatWarning("Low cache hit rate: %.1f%%", cacheHitRate))
|
||||
}
|
||||
|
||||
// Penalize if rebuild is needed
|
||||
if stats.NeedsRebuild {
|
||||
healthScore -= 20
|
||||
warnings = append(warnings, "Vector rebuild recommended: "+stats.RebuildReason)
|
||||
}
|
||||
|
||||
if healthScore < 0 {
|
||||
healthScore = 0
|
||||
}
|
||||
|
||||
status := "healthy"
|
||||
if healthScore < 50 {
|
||||
status = "unhealthy"
|
||||
} else if healthScore < 80 {
|
||||
status = "degraded"
|
||||
}
|
||||
|
||||
writeJSON(w, map[string]any{
|
||||
"status": status,
|
||||
"health_score": healthScore,
|
||||
"warnings": warnings,
|
||||
"stats": stats,
|
||||
"cache_hit_rate": cacheHitRate,
|
||||
})
|
||||
}
|
||||
|
||||
// UpdateObservationRequest is the request body for updating an observation.
|
||||
type UpdateObservationRequest struct {
|
||||
Title *string `json:"title,omitempty"`
|
||||
Subtitle *string `json:"subtitle,omitempty"`
|
||||
Narrative *string `json:"narrative,omitempty"`
|
||||
Scope *string `json:"scope,omitempty"`
|
||||
Facts []string `json:"facts,omitempty"`
|
||||
Concepts []string `json:"concepts,omitempty"`
|
||||
FilesRead []string `json:"files_read,omitempty"`
|
||||
FilesModified []string `json:"files_modified,omitempty"`
|
||||
}
|
||||
|
||||
// handleUpdateObservation updates an existing observation.
|
||||
// PUT /api/observations/{id}
|
||||
func (s *Service) handleUpdateObservation(w http.ResponseWriter, r *http.Request) {
|
||||
// Parse observation ID from URL
|
||||
id, ok := parseIDParam(w, r.PathValue("id"), "observation")
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
|
||||
// Parse request body
|
||||
var req UpdateObservationRequest
|
||||
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
||||
http.Error(w, "invalid request body: "+err.Error(), http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
|
||||
// Build update struct - only include fields that were provided
|
||||
update := &gorm.ObservationUpdate{}
|
||||
|
||||
if req.Title != nil {
|
||||
update.Title = req.Title
|
||||
}
|
||||
if req.Subtitle != nil {
|
||||
update.Subtitle = req.Subtitle
|
||||
}
|
||||
if req.Narrative != nil {
|
||||
update.Narrative = req.Narrative
|
||||
}
|
||||
if req.Facts != nil {
|
||||
update.Facts = &req.Facts
|
||||
}
|
||||
if req.Concepts != nil {
|
||||
update.Concepts = &req.Concepts
|
||||
}
|
||||
if req.FilesRead != nil {
|
||||
update.FilesRead = &req.FilesRead
|
||||
}
|
||||
if req.FilesModified != nil {
|
||||
update.FilesModified = &req.FilesModified
|
||||
}
|
||||
if req.Scope != nil {
|
||||
// Validate scope
|
||||
if *req.Scope != "project" && *req.Scope != "global" {
|
||||
http.Error(w, "scope must be 'project' or 'global'", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
update.Scope = req.Scope
|
||||
}
|
||||
|
||||
// Update the observation
|
||||
updatedObs, err := s.observationStore.UpdateObservation(r.Context(), id, update)
|
||||
if err != nil {
|
||||
if strings.Contains(err.Error(), "not found") {
|
||||
http.Error(w, err.Error(), http.StatusNotFound)
|
||||
return
|
||||
}
|
||||
http.Error(w, "failed to update observation: "+err.Error(), http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
|
||||
// Trigger vector resync for the updated observation
|
||||
if s.vectorSync != nil {
|
||||
s.asyncVectorSync(func() {
|
||||
if err := s.vectorSync.SyncObservation(s.ctx, updatedObs); err != nil {
|
||||
log.Warn().Err(err).Int64("id", id).Msg("Failed to resync observation vectors after update")
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// Broadcast update event
|
||||
s.sseBroadcaster.Broadcast(map[string]any{
|
||||
"type": "observation_updated",
|
||||
"id": id,
|
||||
})
|
||||
|
||||
writeJSON(w, map[string]any{
|
||||
"observation": updatedObs,
|
||||
"message": "observation updated successfully",
|
||||
})
|
||||
}
|
||||
|
||||
// handleGetObservationByID returns a single observation by ID.
|
||||
// GET /api/observations/{id}
|
||||
func (s *Service) handleGetObservationByID(w http.ResponseWriter, r *http.Request) {
|
||||
id, ok := parseIDParam(w, r.PathValue("id"), "observation")
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
|
||||
obs, err := s.observationStore.GetObservationByID(r.Context(), id)
|
||||
if err != nil {
|
||||
http.Error(w, "failed to get observation: "+err.Error(), http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
|
||||
if obs == nil {
|
||||
http.Error(w, "observation not found", http.StatusNotFound)
|
||||
return
|
||||
}
|
||||
|
||||
writeJSON(w, obs)
|
||||
}
|
||||
|
||||
// handleGraphStats returns graph statistics for the dashboard.
|
||||
// Uses relation data to compute knowledge graph metrics.
|
||||
func (s *Service) handleGraphStats(w http.ResponseWriter, r *http.Request) {
|
||||
// Get relation count (edges) - this represents the knowledge graph
|
||||
edgeCount, err := s.relationStore.GetTotalRelationCount(r.Context())
|
||||
if err != nil {
|
||||
edgeCount = 0
|
||||
}
|
||||
|
||||
// Count by relation type
|
||||
edgeTypes := make(map[string]int)
|
||||
for _, t := range models.AllRelationTypes {
|
||||
relations, err := s.relationStore.GetRelationsByType(r.Context(), t, 10000)
|
||||
if err == nil {
|
||||
edgeTypes[string(t)] = len(relations)
|
||||
}
|
||||
}
|
||||
|
||||
// Get unique observation IDs involved in relations (approximate node count)
|
||||
// For now, use edge count as a proxy - each edge has 2 nodes
|
||||
nodeCount := 0
|
||||
if edgeCount > 0 {
|
||||
// Rough estimate: unique nodes ≈ edges * 1.5 (since nodes can have multiple edges)
|
||||
nodeCount = int(float64(edgeCount) * 1.5)
|
||||
}
|
||||
|
||||
// Calculate average degree
|
||||
var avgDegree float64
|
||||
if nodeCount > 0 {
|
||||
avgDegree = float64(edgeCount*2) / float64(nodeCount)
|
||||
}
|
||||
|
||||
// Graph is enabled if we have any edges (relations)
|
||||
enabled := edgeCount > 0
|
||||
|
||||
writeJSON(w, map[string]any{
|
||||
"enabled": enabled,
|
||||
"nodeCount": nodeCount,
|
||||
"edgeCount": edgeCount,
|
||||
"avgDegree": avgDegree,
|
||||
"maxDegree": 0,
|
||||
"minDegree": 0,
|
||||
"medianDegree": 0.0,
|
||||
"edgeTypes": edgeTypes,
|
||||
"config": map[string]any{
|
||||
"maxHops": 2,
|
||||
"branchFactor": 10,
|
||||
"edgeWeight": 0.3,
|
||||
"rebuildIntervalMin": 30,
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
// handleVectorMetrics returns vector database metrics for the dashboard.
|
||||
// Returns enabled: false if vector features are not available.
|
||||
func (s *Service) handleVectorMetrics(w http.ResponseWriter, r *http.Request) {
|
||||
if s.vectorClient == nil {
|
||||
writeJSON(w, map[string]any{
|
||||
"enabled": false,
|
||||
"message": "Vector database not initialized",
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
// Get cache stats from vector client
|
||||
cacheSize, cacheMax := s.vectorClient.CacheStats()
|
||||
cacheStats := s.vectorClient.GetCacheStats()
|
||||
count, _ := s.vectorClient.Count(r.Context())
|
||||
|
||||
uptime := time.Since(s.startTime).Round(time.Second).String()
|
||||
|
||||
// Calculate total queries from cache hits/misses
|
||||
totalQueries := cacheStats.EmbeddingHits + cacheStats.EmbeddingMisses + cacheStats.ResultHits + cacheStats.ResultMisses
|
||||
totalHits := cacheStats.EmbeddingHits + cacheStats.ResultHits
|
||||
totalMisses := cacheStats.EmbeddingMisses + cacheStats.ResultMisses
|
||||
|
||||
writeJSON(w, map[string]any{
|
||||
"enabled": true,
|
||||
"queries": map[string]any{
|
||||
"total": totalQueries,
|
||||
"hubOnly": 0,
|
||||
"hybrid": 0,
|
||||
"onDemand": 0,
|
||||
"graph": 0,
|
||||
},
|
||||
"latency": map[string]any{
|
||||
"avg": "0ms",
|
||||
"p50": "0ms",
|
||||
"p95": "0ms",
|
||||
"p99": "0ms",
|
||||
"avgHub": "0ms",
|
||||
"avgRecompute": "0ms",
|
||||
},
|
||||
"storage": map[string]any{
|
||||
"totalDocuments": count,
|
||||
"hubDocuments": 0,
|
||||
"storedEmbeddings": count,
|
||||
"savingsPercent": 0.0,
|
||||
"recomputedTotal": 0,
|
||||
},
|
||||
"cache": map[string]any{
|
||||
"hits": totalHits,
|
||||
"misses": totalMisses,
|
||||
"hitRate": cacheStats.HitRate(),
|
||||
"size": cacheSize,
|
||||
"maxSize": cacheMax,
|
||||
},
|
||||
"graph": map[string]any{
|
||||
"traversals": 0,
|
||||
"avgDepth": 0.0,
|
||||
},
|
||||
"uptime": uptime,
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,680 @@
|
||||
// Package worker provides import, export, and archive HTTP handlers.
|
||||
package worker
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/go-chi/chi/v5"
|
||||
"github.com/lukaszraczylo/claude-mnemonic/internal/db/gorm"
|
||||
"github.com/lukaszraczylo/claude-mnemonic/pkg/models"
|
||||
"github.com/lukaszraczylo/claude-mnemonic/pkg/similarity"
|
||||
"github.com/rs/zerolog/log"
|
||||
)
|
||||
|
||||
// BulkImportRequest is the request body for bulk observation import.
|
||||
type BulkImportRequest struct {
|
||||
Project string `json:"project"`
|
||||
Observations []BulkObservationInput `json:"observations"`
|
||||
}
|
||||
|
||||
// BulkObservationInput represents a single observation in bulk import.
|
||||
type BulkObservationInput struct {
|
||||
Type string `json:"type"`
|
||||
Title string `json:"title"`
|
||||
Subtitle string `json:"subtitle,omitempty"`
|
||||
Narrative string `json:"narrative,omitempty"`
|
||||
Scope string `json:"scope,omitempty"`
|
||||
Facts []string `json:"facts,omitempty"`
|
||||
Concepts []string `json:"concepts,omitempty"`
|
||||
FilesRead []string `json:"files_read,omitempty"`
|
||||
FilesModified []string `json:"files_modified,omitempty"`
|
||||
}
|
||||
|
||||
// BulkImportResponse contains the result of a bulk import operation.
|
||||
type BulkImportResponse struct {
|
||||
Errors []string `json:"errors,omitempty"`
|
||||
Imported int `json:"imported"`
|
||||
Failed int `json:"failed"`
|
||||
SkippedDuplicates int `json:"skipped_duplicates,omitempty"`
|
||||
}
|
||||
|
||||
// handleBulkImport handles bulk import of observations.
|
||||
// This is useful for migrating data or importing observations from external sources.
|
||||
func (s *Service) handleBulkImport(w http.ResponseWriter, r *http.Request) {
|
||||
// Rate limit bulk operations to prevent DoS
|
||||
if s.bulkOpLimiter != nil && !s.bulkOpLimiter.CanExecute() {
|
||||
remaining := s.bulkOpLimiter.CooldownRemaining()
|
||||
http.Error(w, fmt.Sprintf("bulk import rate limited, retry in %d seconds", remaining), http.StatusTooManyRequests)
|
||||
return
|
||||
}
|
||||
|
||||
var req BulkImportRequest
|
||||
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
||||
http.Error(w, "invalid request body: "+err.Error(), http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
|
||||
if req.Project == "" {
|
||||
http.Error(w, "project is required", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
|
||||
// Validate project name to prevent path traversal
|
||||
if err := ValidateProjectName(req.Project); err != nil {
|
||||
http.Error(w, err.Error(), http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
|
||||
if len(req.Observations) == 0 {
|
||||
http.Error(w, "at least one observation is required", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
|
||||
// Limit batch size to prevent overwhelming the system
|
||||
maxBatchSize := 100
|
||||
if len(req.Observations) > maxBatchSize {
|
||||
http.Error(w, fmt.Sprintf("batch size exceeds maximum of %d", maxBatchSize), http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
|
||||
// Create a synthetic session for bulk import
|
||||
sessionID, err := s.sessionStore.CreateSDKSession(r.Context(), fmt.Sprintf("bulk-import-%d", time.Now().UnixMilli()), req.Project, "bulk import")
|
||||
if err != nil {
|
||||
http.Error(w, "failed to create import session: "+err.Error(), http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
|
||||
var imported, failed, skippedDupes int
|
||||
var errors []string
|
||||
|
||||
// Track imported observations for deduplication within the batch
|
||||
importedObs := make([]*models.Observation, 0, len(req.Observations))
|
||||
|
||||
// Deduplication threshold - observations more similar than this are considered duplicates
|
||||
const dedupThreshold = 0.7
|
||||
|
||||
for i, obsInput := range req.Observations {
|
||||
// Validate observation type using O(1) map lookup
|
||||
if !IsValidObservationType(obsInput.Type) {
|
||||
failed++
|
||||
errors = append(errors, fmt.Sprintf("observation %d: invalid type '%s'", i, obsInput.Type))
|
||||
continue
|
||||
}
|
||||
|
||||
// Build parsed observation
|
||||
parsedObs := &models.ParsedObservation{
|
||||
Type: models.ObservationType(obsInput.Type),
|
||||
Title: obsInput.Title,
|
||||
Subtitle: obsInput.Subtitle,
|
||||
Facts: obsInput.Facts,
|
||||
Narrative: obsInput.Narrative,
|
||||
Concepts: obsInput.Concepts,
|
||||
FilesRead: obsInput.FilesRead,
|
||||
FilesModified: obsInput.FilesModified,
|
||||
Scope: models.ObservationScope(obsInput.Scope),
|
||||
}
|
||||
|
||||
// Convert to temporary observation for similarity check
|
||||
tempObs := &models.Observation{
|
||||
Title: sql.NullString{String: parsedObs.Title, Valid: parsedObs.Title != ""},
|
||||
Subtitle: sql.NullString{String: parsedObs.Subtitle, Valid: parsedObs.Subtitle != ""},
|
||||
Narrative: sql.NullString{String: parsedObs.Narrative, Valid: parsedObs.Narrative != ""},
|
||||
}
|
||||
|
||||
// Check for duplicates within this import batch
|
||||
if similarity.IsSimilarToAny(tempObs, importedObs, dedupThreshold) {
|
||||
skippedDupes++
|
||||
continue
|
||||
}
|
||||
|
||||
// Store observation
|
||||
obsID, _, err := s.observationStore.StoreObservation(
|
||||
r.Context(),
|
||||
fmt.Sprintf("bulk-import-%d", sessionID),
|
||||
req.Project,
|
||||
parsedObs,
|
||||
0, // prompt number
|
||||
0, // discovery tokens
|
||||
)
|
||||
if err != nil {
|
||||
failed++
|
||||
errors = append(errors, fmt.Sprintf("observation %d: %v", i, err))
|
||||
continue
|
||||
}
|
||||
|
||||
// Sync to vector DB asynchronously with rate limiting
|
||||
if s.vectorSync != nil {
|
||||
s.asyncVectorSync(func() {
|
||||
// Use service context as parent to respect shutdown signals
|
||||
ctx, cancel := context.WithTimeout(s.ctx, 10*time.Second)
|
||||
defer cancel()
|
||||
obs, err := s.observationStore.GetObservationByID(ctx, obsID)
|
||||
if err == nil && obs != nil {
|
||||
if syncErr := s.vectorSync.SyncObservation(ctx, obs); syncErr != nil {
|
||||
if s.ctx.Err() == nil { // Don't log during shutdown
|
||||
log.Debug().Err(syncErr).Int64("id", obsID).Msg("Failed to sync observation during bulk import")
|
||||
}
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// Track for deduplication of subsequent observations in this batch
|
||||
importedObs = append(importedObs, tempObs)
|
||||
imported++
|
||||
}
|
||||
|
||||
log.Info().
|
||||
Str("project", req.Project).
|
||||
Int("imported", imported).
|
||||
Int("failed", failed).
|
||||
Int("skipped_duplicates", skippedDupes).
|
||||
Msg("Bulk import completed")
|
||||
|
||||
// Invalidate observation count cache after import
|
||||
if imported > 0 {
|
||||
if req.Project != "" {
|
||||
s.invalidateObsCountCache(req.Project)
|
||||
} else {
|
||||
s.invalidateAllObsCountCache()
|
||||
}
|
||||
}
|
||||
|
||||
// Broadcast observation event for dashboard refresh
|
||||
s.sseBroadcaster.Broadcast(map[string]any{
|
||||
"type": "observation",
|
||||
"action": "bulk_import",
|
||||
"project": req.Project,
|
||||
"count": imported,
|
||||
})
|
||||
|
||||
writeJSON(w, BulkImportResponse{
|
||||
Imported: imported,
|
||||
Failed: failed,
|
||||
SkippedDuplicates: skippedDupes,
|
||||
Errors: errors,
|
||||
})
|
||||
}
|
||||
|
||||
// ArchiveRequest is the request body for archiving observations.
|
||||
type ArchiveRequest struct {
|
||||
Project string `json:"project,omitempty"`
|
||||
Reason string `json:"reason,omitempty"`
|
||||
IDs []int64 `json:"ids,omitempty"`
|
||||
MaxAgeDays int `json:"max_age_days,omitempty"`
|
||||
}
|
||||
|
||||
// handleArchiveObservations archives observations by ID or by age.
|
||||
// Supports batch archival with error tracking per observation.
|
||||
func (s *Service) handleArchiveObservations(w http.ResponseWriter, r *http.Request) {
|
||||
var req ArchiveRequest
|
||||
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
||||
http.Error(w, "invalid request body: "+err.Error(), http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
|
||||
var archivedIDs []int64
|
||||
var failedIDs []int64
|
||||
var errors []string
|
||||
var err error
|
||||
|
||||
if len(req.IDs) > 0 {
|
||||
// Archive specific observations with parallel processing for large batches
|
||||
if len(req.IDs) > 5 {
|
||||
// Use parallel archival for batches larger than 5
|
||||
type archiveResult struct {
|
||||
err error
|
||||
id int64
|
||||
}
|
||||
results := make(chan archiveResult, len(req.IDs))
|
||||
|
||||
// Limit concurrency to avoid overwhelming the database
|
||||
sem := make(chan struct{}, 5)
|
||||
var wg sync.WaitGroup
|
||||
|
||||
for _, id := range req.IDs {
|
||||
wg.Add(1)
|
||||
go func(obsID int64) {
|
||||
defer wg.Done()
|
||||
sem <- struct{}{} // Acquire
|
||||
defer func() { <-sem }() // Release
|
||||
|
||||
archErr := s.observationStore.ArchiveObservation(r.Context(), obsID, req.Reason)
|
||||
results <- archiveResult{id: obsID, err: archErr}
|
||||
}(id)
|
||||
}
|
||||
|
||||
// Close results channel when all goroutines complete
|
||||
go func() {
|
||||
wg.Wait()
|
||||
close(results)
|
||||
}()
|
||||
|
||||
// Collect results
|
||||
for res := range results {
|
||||
if res.err != nil {
|
||||
log.Warn().Err(res.err).Int64("id", res.id).Msg("Failed to archive observation")
|
||||
failedIDs = append(failedIDs, res.id)
|
||||
errors = append(errors, fmt.Sprintf("id %d: %v", res.id, res.err))
|
||||
} else {
|
||||
archivedIDs = append(archivedIDs, res.id)
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// Sequential for small batches
|
||||
for _, id := range req.IDs {
|
||||
if archErr := s.observationStore.ArchiveObservation(r.Context(), id, req.Reason); archErr != nil {
|
||||
log.Warn().Err(archErr).Int64("id", id).Msg("Failed to archive observation")
|
||||
failedIDs = append(failedIDs, id)
|
||||
errors = append(errors, fmt.Sprintf("id %d: %v", id, archErr))
|
||||
} else {
|
||||
archivedIDs = append(archivedIDs, id)
|
||||
}
|
||||
}
|
||||
}
|
||||
} else if req.Project != "" || req.MaxAgeDays > 0 {
|
||||
// Archive by age
|
||||
archivedIDs, err = s.observationStore.ArchiveOldObservations(r.Context(), req.Project, req.MaxAgeDays, req.Reason)
|
||||
if err != nil {
|
||||
http.Error(w, "failed to archive: "+err.Error(), http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
} else {
|
||||
http.Error(w, "either 'ids' or 'project'/'max_age_days' is required", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
|
||||
log.Info().
|
||||
Str("project", req.Project).
|
||||
Int("archived", len(archivedIDs)).
|
||||
Int("failed", len(failedIDs)).
|
||||
Msg("Observations archived")
|
||||
|
||||
// Invalidate cache if any observations were archived
|
||||
if len(archivedIDs) > 0 {
|
||||
if req.Project != "" {
|
||||
s.invalidateObsCountCache(req.Project)
|
||||
} else {
|
||||
s.invalidateAllObsCountCache()
|
||||
}
|
||||
}
|
||||
|
||||
response := map[string]any{
|
||||
"archived_count": len(archivedIDs),
|
||||
"archived_ids": archivedIDs,
|
||||
}
|
||||
if len(failedIDs) > 0 {
|
||||
response["failed_count"] = len(failedIDs)
|
||||
response["failed_ids"] = failedIDs
|
||||
response["errors"] = errors
|
||||
}
|
||||
|
||||
writeJSON(w, response)
|
||||
}
|
||||
|
||||
// handleUnarchiveObservation restores an archived observation.
|
||||
func (s *Service) handleUnarchiveObservation(w http.ResponseWriter, r *http.Request) {
|
||||
idStr := chi.URLParam(r, "id")
|
||||
id, err := strconv.ParseInt(idStr, 10, 64)
|
||||
if err != nil {
|
||||
http.Error(w, "invalid observation id", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
|
||||
if err := s.observationStore.UnarchiveObservation(r.Context(), id); err != nil {
|
||||
http.Error(w, err.Error(), http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
|
||||
// Invalidate all caches since we don't know the project
|
||||
s.invalidateAllObsCountCache()
|
||||
|
||||
writeJSON(w, map[string]any{
|
||||
"success": true,
|
||||
"id": id,
|
||||
})
|
||||
}
|
||||
|
||||
// handleGetArchivedObservations returns archived observations.
|
||||
func (s *Service) handleGetArchivedObservations(w http.ResponseWriter, r *http.Request) {
|
||||
project := r.URL.Query().Get("project")
|
||||
limit := gorm.ParseLimitParam(r, DefaultObservationsLimit)
|
||||
|
||||
observations, err := s.observationStore.GetArchivedObservations(r.Context(), project, limit)
|
||||
if err != nil {
|
||||
http.Error(w, err.Error(), http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
|
||||
if observations == nil {
|
||||
observations = []*models.Observation{}
|
||||
}
|
||||
|
||||
writeJSON(w, observations)
|
||||
}
|
||||
|
||||
// handleGetArchivalStats returns archival statistics.
|
||||
func (s *Service) handleGetArchivalStats(w http.ResponseWriter, r *http.Request) {
|
||||
project := r.URL.Query().Get("project")
|
||||
|
||||
stats, err := s.observationStore.GetArchivalStats(r.Context(), project)
|
||||
if err != nil {
|
||||
http.Error(w, err.Error(), http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
|
||||
writeJSON(w, stats)
|
||||
}
|
||||
|
||||
// handleExportObservations exports observations in JSON or CSV format.
|
||||
// Supports query parameters: project, format (json/csv), scope, type, limit.
|
||||
func (s *Service) handleExportObservations(w http.ResponseWriter, r *http.Request) {
|
||||
project := r.URL.Query().Get("project")
|
||||
format := r.URL.Query().Get("format")
|
||||
if format == "" {
|
||||
format = "json"
|
||||
}
|
||||
scope := r.URL.Query().Get("scope") // project, global, or empty for all
|
||||
obsType := r.URL.Query().Get("type") // bugfix, feature, etc.
|
||||
limit := gorm.ParseLimitParamWithMax(r, 1000, 5000) // Higher limit for exports, capped at 5000
|
||||
|
||||
// Validate format
|
||||
if format != "json" && format != "csv" {
|
||||
http.Error(w, "format must be 'json' or 'csv'", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
|
||||
// Get observations with filters
|
||||
ctx := r.Context()
|
||||
var observations []*models.Observation
|
||||
var err error
|
||||
|
||||
if project != "" {
|
||||
observations, _, err = s.observationStore.GetObservationsByProjectStrictPaginated(ctx, project, limit, 0)
|
||||
} else {
|
||||
observations, _, err = s.observationStore.GetAllRecentObservationsPaginated(ctx, limit, 0)
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
http.Error(w, err.Error(), http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
|
||||
// Apply additional filters
|
||||
if scope != "" || obsType != "" {
|
||||
filtered := make([]*models.Observation, 0, len(observations))
|
||||
for _, obs := range observations {
|
||||
if scope != "" && string(obs.Scope) != scope {
|
||||
continue
|
||||
}
|
||||
if obsType != "" && string(obs.Type) != obsType {
|
||||
continue
|
||||
}
|
||||
filtered = append(filtered, obs)
|
||||
}
|
||||
observations = filtered
|
||||
}
|
||||
|
||||
// Generate filename
|
||||
timestamp := time.Now().Format("20060102-150405")
|
||||
filename := fmt.Sprintf("observations-%s.%s", timestamp, format)
|
||||
if project != "" {
|
||||
// Sanitize project name for filename
|
||||
sanitized := strings.ReplaceAll(project, "/", "_")
|
||||
sanitized = strings.ReplaceAll(sanitized, "\\", "_")
|
||||
if len(sanitized) > 50 {
|
||||
sanitized = sanitized[:50]
|
||||
}
|
||||
filename = fmt.Sprintf("observations-%s-%s.%s", sanitized, timestamp, format)
|
||||
}
|
||||
|
||||
switch format {
|
||||
case "csv":
|
||||
w.Header().Set("Content-Type", "text/csv")
|
||||
w.Header().Set("Content-Disposition", fmt.Sprintf("attachment; filename=%q", filename))
|
||||
s.writeObservationsCSV(w, observations)
|
||||
default: // json
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.Header().Set("Content-Disposition", fmt.Sprintf("attachment; filename=%q", filename))
|
||||
writeJSON(w, map[string]any{
|
||||
"exported_at": time.Now().Format(time.RFC3339),
|
||||
"project": project,
|
||||
"count": len(observations),
|
||||
"observations": observations,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// writeObservationsCSV writes observations in CSV format.
|
||||
// Uses fmt.Fprintf directly to avoid intermediate string allocations.
|
||||
func (s *Service) writeObservationsCSV(w http.ResponseWriter, observations []*models.Observation) {
|
||||
// Write CSV header
|
||||
_, _ = io.WriteString(w, "id,type,scope,project,title,subtitle,narrative,concepts,facts,created_at,importance_score\n")
|
||||
|
||||
for _, obs := range observations {
|
||||
// Write directly to avoid string allocation per row
|
||||
_, _ = fmt.Fprintf(w, "%d,%s,%s,%s,%s,%s,%s,%s,%s,%s,%.2f\n",
|
||||
obs.ID,
|
||||
obs.Type,
|
||||
obs.Scope,
|
||||
escapeCsvField(obs.Project),
|
||||
escapeCsvField(obs.Title.String),
|
||||
escapeCsvField(obs.Subtitle.String),
|
||||
escapeCsvField(obs.Narrative.String),
|
||||
escapeCsvField(strings.Join(obs.Concepts, ";")),
|
||||
escapeCsvField(strings.Join(obs.Facts, ";")),
|
||||
obs.CreatedAt,
|
||||
obs.ImportanceScore,
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
// escapeCsvField escapes a field for CSV output.
|
||||
func escapeCsvField(s string) string {
|
||||
// If field contains comma, quote, or newline, wrap in quotes and escape quotes
|
||||
if strings.ContainsAny(s, ",\"\n\r") {
|
||||
s = strings.ReplaceAll(s, "\"", "\"\"")
|
||||
return "\"" + s + "\""
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
// BulkStatusRequest represents a request to update status for multiple observations.
|
||||
type BulkStatusRequest struct {
|
||||
Action string `json:"action"`
|
||||
Reason string `json:"reason,omitempty"`
|
||||
IDs []int64 `json:"ids"`
|
||||
Feedback int `json:"feedback,omitempty"`
|
||||
}
|
||||
|
||||
// handleBulkStatusUpdate updates status for multiple observations in one request.
|
||||
func (s *Service) handleBulkStatusUpdate(w http.ResponseWriter, r *http.Request) {
|
||||
var req BulkStatusRequest
|
||||
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
||||
http.Error(w, "invalid request body: "+err.Error(), http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
|
||||
if len(req.IDs) == 0 {
|
||||
http.Error(w, "ids is required", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
|
||||
if len(req.IDs) > 500 {
|
||||
http.Error(w, "maximum 500 ids per request", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
|
||||
ctx := r.Context()
|
||||
var updated, failed int
|
||||
var errors []string
|
||||
|
||||
switch req.Action {
|
||||
case "supersede":
|
||||
for _, id := range req.IDs {
|
||||
if err := s.observationStore.MarkAsSuperseded(ctx, id); err != nil {
|
||||
failed++
|
||||
errors = append(errors, fmt.Sprintf("id %d: %v", id, err))
|
||||
} else {
|
||||
updated++
|
||||
}
|
||||
}
|
||||
|
||||
case "archive":
|
||||
for _, id := range req.IDs {
|
||||
if err := s.observationStore.ArchiveObservation(ctx, id, req.Reason); err != nil {
|
||||
failed++
|
||||
errors = append(errors, fmt.Sprintf("id %d: %v", id, err))
|
||||
} else {
|
||||
updated++
|
||||
}
|
||||
}
|
||||
|
||||
case "set_feedback":
|
||||
if req.Feedback < -1 || req.Feedback > 1 {
|
||||
http.Error(w, "feedback must be -1, 0, or 1", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
for _, id := range req.IDs {
|
||||
if err := s.observationStore.UpdateObservationFeedback(ctx, id, req.Feedback); err != nil {
|
||||
failed++
|
||||
errors = append(errors, fmt.Sprintf("id %d: %v", id, err))
|
||||
} else {
|
||||
updated++
|
||||
}
|
||||
}
|
||||
|
||||
default:
|
||||
http.Error(w, "action must be 'supersede', 'archive', or 'set_feedback'", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
|
||||
// Invalidate cache for archive action (affects observation counts)
|
||||
if req.Action == "archive" && updated > 0 {
|
||||
// No project info available, invalidate all caches
|
||||
s.invalidateAllObsCountCache()
|
||||
}
|
||||
|
||||
response := map[string]any{
|
||||
"action": req.Action,
|
||||
"updated": updated,
|
||||
"failed": failed,
|
||||
}
|
||||
if len(errors) > 0 {
|
||||
response["errors"] = errors
|
||||
}
|
||||
|
||||
writeJSON(w, response)
|
||||
}
|
||||
|
||||
// handleFindDuplicates finds potential duplicate observations using similarity clustering.
|
||||
// Returns groups of similar observations that may be candidates for merging or archival.
|
||||
func (s *Service) handleFindDuplicates(w http.ResponseWriter, r *http.Request) {
|
||||
project := r.URL.Query().Get("project")
|
||||
thresholdStr := r.URL.Query().Get("threshold")
|
||||
limit := gorm.ParseLimitParam(r, 100)
|
||||
|
||||
// Parse threshold (default 0.6 = 60% similarity)
|
||||
threshold := 0.6
|
||||
if thresholdStr != "" {
|
||||
if t, err := strconv.ParseFloat(thresholdStr, 64); err == nil && t > 0 && t < 1 {
|
||||
threshold = t
|
||||
}
|
||||
}
|
||||
|
||||
// Get recent observations
|
||||
ctx := r.Context()
|
||||
var observations []*models.Observation
|
||||
var err error
|
||||
|
||||
if project != "" {
|
||||
observations, _, err = s.observationStore.GetObservationsByProjectStrictPaginated(ctx, project, limit, 0)
|
||||
} else {
|
||||
observations, _, err = s.observationStore.GetAllRecentObservationsPaginated(ctx, limit, 0)
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
http.Error(w, err.Error(), http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
|
||||
if len(observations) < 2 {
|
||||
writeJSON(w, map[string]any{
|
||||
"duplicate_groups": []any{},
|
||||
"total_checked": len(observations),
|
||||
"threshold": threshold,
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
// Find duplicates using similarity comparison
|
||||
type duplicateGroup struct {
|
||||
Observations []map[string]any `json:"observations"`
|
||||
Similarity float64 `json:"similarity"`
|
||||
}
|
||||
|
||||
groups := []duplicateGroup{}
|
||||
processed := make(map[int64]bool)
|
||||
|
||||
for i, obs1 := range observations {
|
||||
if processed[obs1.ID] {
|
||||
continue
|
||||
}
|
||||
|
||||
terms1 := similarity.ExtractObservationTerms(obs1)
|
||||
if len(terms1) == 0 {
|
||||
continue
|
||||
}
|
||||
|
||||
group := duplicateGroup{
|
||||
Observations: []map[string]any{obs1.ToMap()},
|
||||
Similarity: 1.0,
|
||||
}
|
||||
|
||||
for j := i + 1; j < len(observations); j++ {
|
||||
obs2 := observations[j]
|
||||
if processed[obs2.ID] {
|
||||
continue
|
||||
}
|
||||
|
||||
terms2 := similarity.ExtractObservationTerms(obs2)
|
||||
sim := similarity.JaccardSimilarity(terms1, terms2)
|
||||
|
||||
if sim >= threshold {
|
||||
obsMap := obs2.ToMap()
|
||||
obsMap["similarity_to_first"] = sim
|
||||
group.Observations = append(group.Observations, obsMap)
|
||||
group.Similarity = min(group.Similarity, sim)
|
||||
processed[obs2.ID] = true
|
||||
}
|
||||
}
|
||||
|
||||
if len(group.Observations) > 1 {
|
||||
processed[obs1.ID] = true
|
||||
groups = append(groups, group)
|
||||
}
|
||||
}
|
||||
|
||||
// Sort groups by size (largest first)
|
||||
sort.Slice(groups, func(i, j int) bool {
|
||||
return len(groups[i].Observations) > len(groups[j].Observations)
|
||||
})
|
||||
|
||||
writeJSON(w, map[string]any{
|
||||
"duplicate_groups": groups,
|
||||
"total_checked": len(observations),
|
||||
"groups_found": len(groups),
|
||||
"threshold": threshold,
|
||||
"project": project,
|
||||
})
|
||||
}
|
||||
@@ -46,7 +46,7 @@ func (s *Service) handleGetRelationGraph(w http.ResponseWriter, r *http.Request)
|
||||
// Get depth parameter (default 2)
|
||||
depth := 2
|
||||
if depthStr := r.URL.Query().Get("depth"); depthStr != "" {
|
||||
if d, err := strconv.Atoi(depthStr); err == nil && d > 0 && d <= 5 {
|
||||
if d, parseErr := strconv.Atoi(depthStr); parseErr == nil && d > 0 && d <= 5 {
|
||||
depth = d
|
||||
}
|
||||
}
|
||||
@@ -72,7 +72,7 @@ func (s *Service) handleGetRelatedObservations(w http.ResponseWriter, r *http.Re
|
||||
// Get minimum confidence parameter (default 0.4)
|
||||
minConfidence := 0.4
|
||||
if confStr := r.URL.Query().Get("min_confidence"); confStr != "" {
|
||||
if c, err := strconv.ParseFloat(confStr, 64); err == nil && c >= 0 && c <= 1 {
|
||||
if c, parseErr := strconv.ParseFloat(confStr, 64); parseErr == nil && c >= 0 && c <= 1 {
|
||||
minConfidence = c
|
||||
}
|
||||
}
|
||||
|
||||
@@ -10,6 +10,7 @@ import (
|
||||
|
||||
"github.com/go-chi/chi/v5"
|
||||
"github.com/lukaszraczylo/claude-mnemonic/pkg/models"
|
||||
"github.com/rs/zerolog/log"
|
||||
)
|
||||
|
||||
// FeedbackRequest represents a user feedback submission.
|
||||
@@ -41,11 +42,9 @@ func (s *Service) handleObservationFeedback(w http.ResponseWriter, r *http.Reque
|
||||
return
|
||||
}
|
||||
|
||||
// Get required components
|
||||
s.initMu.RLock()
|
||||
// Get required components (initMu.RLock held by requireReady middleware)
|
||||
observationStore := s.observationStore
|
||||
scoreCalculator := s.scoreCalculator
|
||||
s.initMu.RUnlock()
|
||||
|
||||
if observationStore == nil {
|
||||
http.Error(w, "service not ready", http.StatusServiceUnavailable)
|
||||
@@ -94,10 +93,9 @@ func (s *Service) handleObservationFeedback(w http.ResponseWriter, r *http.Reque
|
||||
func (s *Service) handleGetScoringStats(w http.ResponseWriter, r *http.Request) {
|
||||
project := r.URL.Query().Get("project")
|
||||
|
||||
s.initMu.RLock()
|
||||
// initMu.RLock held by requireReady middleware
|
||||
observationStore := s.observationStore
|
||||
recalculator := s.recalculator
|
||||
s.initMu.RUnlock()
|
||||
|
||||
if observationStore == nil {
|
||||
http.Error(w, "service not ready", http.StatusServiceUnavailable)
|
||||
@@ -129,9 +127,8 @@ func (s *Service) handleGetTopObservations(w http.ResponseWriter, r *http.Reques
|
||||
limit := parseIntParam(r, "limit", 10)
|
||||
project := r.URL.Query().Get("project")
|
||||
|
||||
s.initMu.RLock()
|
||||
// initMu.RLock held by requireReady middleware
|
||||
observationStore := s.observationStore
|
||||
s.initMu.RUnlock()
|
||||
|
||||
if observationStore == nil {
|
||||
http.Error(w, "service not ready", http.StatusServiceUnavailable)
|
||||
@@ -157,9 +154,8 @@ func (s *Service) handleGetMostRetrieved(w http.ResponseWriter, r *http.Request)
|
||||
limit := parseIntParam(r, "limit", 10)
|
||||
project := r.URL.Query().Get("project")
|
||||
|
||||
s.initMu.RLock()
|
||||
// initMu.RLock held by requireReady middleware
|
||||
observationStore := s.observationStore
|
||||
s.initMu.RUnlock()
|
||||
|
||||
if observationStore == nil {
|
||||
http.Error(w, "service not ready", http.StatusServiceUnavailable)
|
||||
@@ -190,10 +186,9 @@ func (s *Service) handleExplainScore(w http.ResponseWriter, r *http.Request) {
|
||||
return
|
||||
}
|
||||
|
||||
s.initMu.RLock()
|
||||
// initMu.RLock held by requireReady middleware
|
||||
observationStore := s.observationStore
|
||||
scoreCalculator := s.scoreCalculator
|
||||
s.initMu.RUnlock()
|
||||
|
||||
if observationStore == nil || scoreCalculator == nil {
|
||||
http.Error(w, "service not ready", http.StatusServiceUnavailable)
|
||||
@@ -244,10 +239,9 @@ func (s *Service) handleUpdateConceptWeight(w http.ResponseWriter, r *http.Reque
|
||||
return
|
||||
}
|
||||
|
||||
s.initMu.RLock()
|
||||
// initMu.RLock held by requireReady middleware
|
||||
observationStore := s.observationStore
|
||||
recalculator := s.recalculator
|
||||
s.initMu.RUnlock()
|
||||
|
||||
if observationStore == nil {
|
||||
http.Error(w, "service not ready", http.StatusServiceUnavailable)
|
||||
@@ -278,9 +272,8 @@ func (s *Service) handleUpdateConceptWeight(w http.ResponseWriter, r *http.Reque
|
||||
// handleGetConceptWeights returns all concept weights.
|
||||
// GET /api/scoring/concepts
|
||||
func (s *Service) handleGetConceptWeights(w http.ResponseWriter, r *http.Request) {
|
||||
s.initMu.RLock()
|
||||
// initMu.RLock held by requireReady middleware
|
||||
observationStore := s.observationStore
|
||||
s.initMu.RUnlock()
|
||||
|
||||
if observationStore == nil {
|
||||
http.Error(w, "service not ready", http.StatusServiceUnavailable)
|
||||
@@ -299,20 +292,22 @@ func (s *Service) handleGetConceptWeights(w http.ResponseWriter, r *http.Request
|
||||
// handleTriggerRecalculation triggers an immediate score recalculation.
|
||||
// POST /api/scoring/recalculate
|
||||
func (s *Service) handleTriggerRecalculation(w http.ResponseWriter, r *http.Request) {
|
||||
s.initMu.RLock()
|
||||
// initMu.RLock held by requireReady middleware
|
||||
recalculator := s.recalculator
|
||||
s.initMu.RUnlock()
|
||||
|
||||
if recalculator == nil {
|
||||
http.Error(w, "recalculator not available", http.StatusServiceUnavailable)
|
||||
return
|
||||
}
|
||||
|
||||
// Run recalculation in background
|
||||
// Run recalculation in background with independent context
|
||||
s.wg.Add(1)
|
||||
go func() {
|
||||
if err := recalculator.RecalculateNow(r.Context()); err != nil {
|
||||
// Log error but don't block response
|
||||
_ = err // Explicitly ignore - background operation
|
||||
defer s.wg.Done()
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
|
||||
defer cancel()
|
||||
if err := recalculator.RecalculateNow(ctx); err != nil {
|
||||
log.Error().Err(err).Msg("Background recalculation failed")
|
||||
}
|
||||
}()
|
||||
|
||||
@@ -336,16 +331,17 @@ func (s *Service) incrementRetrievalCounts(ids []int64) {
|
||||
return
|
||||
}
|
||||
|
||||
s.initMu.RLock()
|
||||
// initMu.RLock held by requireReady middleware (caller is always behind requireReady)
|
||||
store := s.observationStore
|
||||
s.initMu.RUnlock()
|
||||
|
||||
if store == nil {
|
||||
return
|
||||
}
|
||||
|
||||
// Increment in background to not block response
|
||||
s.wg.Add(1)
|
||||
go func() {
|
||||
defer s.wg.Done()
|
||||
// Create a new context with timeout for the background operation
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
|
||||
defer cancel()
|
||||
|
||||
@@ -0,0 +1,402 @@
|
||||
// Package worker provides session-related HTTP handlers.
|
||||
package worker
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/go-chi/chi/v5"
|
||||
"github.com/lukaszraczylo/claude-mnemonic/internal/privacy"
|
||||
"github.com/lukaszraczylo/claude-mnemonic/internal/worker/session"
|
||||
"github.com/lukaszraczylo/claude-mnemonic/pkg/models"
|
||||
"github.com/rs/zerolog/log"
|
||||
)
|
||||
|
||||
// SessionInitRequest is the request body for session initialization.
|
||||
type SessionInitRequest struct {
|
||||
ClaudeSessionID string `json:"claudeSessionId"`
|
||||
Project string `json:"project"`
|
||||
Prompt string `json:"prompt"`
|
||||
MatchedObservations int `json:"matchedObservations"`
|
||||
}
|
||||
|
||||
// SessionInitResponse is the response for session initialization.
|
||||
type SessionInitResponse struct {
|
||||
Reason string `json:"reason,omitempty"`
|
||||
SessionDBID int64 `json:"sessionDbId"`
|
||||
PromptNumber int `json:"promptNumber"`
|
||||
Skipped bool `json:"skipped,omitempty"`
|
||||
}
|
||||
|
||||
// DuplicatePromptWindowSeconds is the time window for detecting duplicate prompt submissions.
|
||||
// If the same prompt text is seen within this window, it's considered a duplicate hook invocation.
|
||||
const DuplicatePromptWindowSeconds = 10
|
||||
|
||||
// handleSessionInit handles session initialization from user-prompt hook.
|
||||
// This handler is idempotent - duplicate requests within a short time window
|
||||
// return the existing prompt data without creating duplicates.
|
||||
func (s *Service) handleSessionInit(w http.ResponseWriter, r *http.Request) {
|
||||
var req SessionInitRequest
|
||||
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
||||
http.Error(w, err.Error(), http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
// Reject requests from internal CLI calls (callClaudeCLI runs from /tmp with no session ID).
|
||||
// These are the memory extraction agent's own prompts, not real user prompts.
|
||||
if req.ClaudeSessionID == "" {
|
||||
log.Debug().Str("project", req.Project).Msg("Rejecting session init with empty claude session ID (internal call)")
|
||||
writeJSON(w, SessionInitResponse{
|
||||
Skipped: true,
|
||||
Reason: "internal",
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
// Reject prompts that contain the internal system prompt (from SDK processor).
|
||||
// The SDK processor sets CLAUDE_MNEMONIC_INTERNAL=1 on the subprocess env,
|
||||
// but Claude Code does NOT propagate custom env vars to hook subprocesses,
|
||||
// so that guard fails and the system prompt leaks into the database as a
|
||||
// regular user prompt. This server-side check catches those cases reliably.
|
||||
if strings.Contains(req.Prompt, "You are a memory extraction agent for Claude Code sessions") {
|
||||
log.Debug().Str("project", req.Project).Msg("Rejecting session init with internal system prompt")
|
||||
writeJSON(w, SessionInitResponse{
|
||||
Skipped: true,
|
||||
Reason: "internal",
|
||||
})
|
||||
return
|
||||
}
|
||||
// Privacy check
|
||||
if privacy.IsEntirelyPrivate(req.Prompt) {
|
||||
// Create session but skip processing
|
||||
sessionID, _ := s.sessionStore.CreateSDKSession(r.Context(), req.ClaudeSessionID, req.Project, "")
|
||||
promptNum, _ := s.sessionStore.IncrementPromptCounter(r.Context(), sessionID)
|
||||
|
||||
writeJSON(w, SessionInitResponse{
|
||||
SessionDBID: sessionID,
|
||||
PromptNumber: promptNum,
|
||||
Skipped: true,
|
||||
Reason: "private",
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
// Clean prompt
|
||||
cleanedPrompt := privacy.Clean(req.Prompt)
|
||||
|
||||
// DUPLICATE DETECTION: Check if this exact prompt was already saved recently.
|
||||
// This prevents the bug where the hook fires multiple times for the same user action,
|
||||
// creating many duplicate prompts with incrementing numbers.
|
||||
if existingID, existingNum, found := s.promptStore.FindRecentPromptByText(r.Context(), req.ClaudeSessionID, cleanedPrompt, DuplicatePromptWindowSeconds); found {
|
||||
// Get or create session (idempotent)
|
||||
sessionID, _ := s.sessionStore.CreateSDKSession(r.Context(), req.ClaudeSessionID, req.Project, cleanedPrompt)
|
||||
|
||||
log.Debug().
|
||||
Int64("sessionId", sessionID).
|
||||
Int("promptNumber", existingNum).
|
||||
Int64("promptId", existingID).
|
||||
Msg("Duplicate prompt detected - returning existing")
|
||||
|
||||
// Return existing prompt data without incrementing or saving again
|
||||
writeJSON(w, SessionInitResponse{
|
||||
SessionDBID: sessionID,
|
||||
PromptNumber: existingNum,
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
// CROSS-SESSION DUPLICATE DETECTION: Same prompt text from ANY session within the window.
|
||||
// This catches duplicates when the SDK processor's callClaudeCLI subprocess fires
|
||||
// UserPromptSubmit with a different Claude session ID than the original hook.
|
||||
if _, existingNum, found := s.promptStore.FindRecentPromptByTextGlobal(r.Context(), cleanedPrompt, DuplicatePromptWindowSeconds); found {
|
||||
sessionID, _ := s.sessionStore.CreateSDKSession(r.Context(), req.ClaudeSessionID, req.Project, cleanedPrompt)
|
||||
|
||||
log.Debug().
|
||||
Int64("sessionId", sessionID).
|
||||
Int("promptNumber", existingNum).
|
||||
Msg("Cross-session duplicate prompt detected")
|
||||
|
||||
writeJSON(w, SessionInitResponse{
|
||||
SessionDBID: sessionID,
|
||||
PromptNumber: existingNum,
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
// Create session (idempotent)
|
||||
sessionID, err := s.sessionStore.CreateSDKSession(r.Context(), req.ClaudeSessionID, req.Project, cleanedPrompt)
|
||||
if err != nil {
|
||||
http.Error(w, err.Error(), http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
|
||||
// Increment prompt counter
|
||||
promptNum, err := s.sessionStore.IncrementPromptCounter(r.Context(), sessionID)
|
||||
if err != nil {
|
||||
http.Error(w, err.Error(), http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
|
||||
// Save user prompt with matched observation count
|
||||
promptID, err := s.promptStore.SaveUserPromptWithMatches(r.Context(), req.ClaudeSessionID, promptNum, cleanedPrompt, req.MatchedObservations)
|
||||
if err != nil {
|
||||
log.Warn().Err(err).Msg("Failed to save user prompt")
|
||||
// Non-fatal: continue with session initialization
|
||||
} else if s.vectorSync != nil {
|
||||
// Sync to vector DB asynchronously (non-blocking)
|
||||
now := time.Now()
|
||||
promptWithSession := &models.UserPromptWithSession{
|
||||
UserPrompt: models.UserPrompt{
|
||||
ID: promptID,
|
||||
ClaudeSessionID: req.ClaudeSessionID,
|
||||
PromptNumber: promptNum,
|
||||
PromptText: cleanedPrompt,
|
||||
MatchedObservations: req.MatchedObservations,
|
||||
CreatedAt: now.Format(time.RFC3339),
|
||||
CreatedAtEpoch: now.UnixMilli(),
|
||||
},
|
||||
Project: req.Project,
|
||||
SDKSessionID: req.ClaudeSessionID,
|
||||
}
|
||||
s.asyncVectorSync(func() {
|
||||
// Use service context as parent to respect shutdown signals
|
||||
ctx, cancel := context.WithTimeout(s.ctx, 10*time.Second)
|
||||
defer cancel()
|
||||
if err := s.vectorSync.SyncUserPrompt(ctx, promptWithSession); err != nil {
|
||||
if s.ctx.Err() == nil { // Don't log during shutdown
|
||||
log.Warn().Err(err).Int64("id", promptID).Msg("Failed to sync user prompt to sqlite-vec")
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
log.Info().
|
||||
Int64("sessionId", sessionID).
|
||||
Int("promptNumber", promptNum).
|
||||
Str("project", req.Project).
|
||||
Msg("Session initialized")
|
||||
|
||||
// Broadcast prompt event for dashboard refresh
|
||||
s.sseBroadcaster.Broadcast(map[string]any{
|
||||
"type": "prompt",
|
||||
"action": "created",
|
||||
"project": req.Project,
|
||||
})
|
||||
|
||||
writeJSON(w, SessionInitResponse{
|
||||
SessionDBID: sessionID,
|
||||
PromptNumber: promptNum,
|
||||
})
|
||||
}
|
||||
|
||||
// SessionStartRequest is the request body for starting SDK agent.
|
||||
type SessionStartRequest struct {
|
||||
UserPrompt string `json:"userPrompt"`
|
||||
PromptNumber int `json:"promptNumber"`
|
||||
}
|
||||
|
||||
// handleSessionStart handles SDK agent session start.
|
||||
func (s *Service) handleSessionStart(w http.ResponseWriter, r *http.Request) {
|
||||
idStr := chi.URLParam(r, "id")
|
||||
id, err := strconv.ParseInt(idStr, 10, 64)
|
||||
if err != nil {
|
||||
http.Error(w, "invalid session id", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
|
||||
var req SessionStartRequest
|
||||
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
||||
http.Error(w, err.Error(), http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
|
||||
// Initialize session in manager
|
||||
sess, err := s.sessionManager.InitializeSession(r.Context(), id, req.UserPrompt, req.PromptNumber)
|
||||
if err != nil {
|
||||
http.Error(w, err.Error(), http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
if sess == nil {
|
||||
http.Error(w, "session not found", http.StatusNotFound)
|
||||
return
|
||||
}
|
||||
|
||||
// Session is now registered. Observations will be processed
|
||||
// asynchronously by the background queue processor (processQueue in service.go).
|
||||
log.Info().
|
||||
Int64("sessionId", id).
|
||||
Int("promptNumber", req.PromptNumber).
|
||||
Msg("SDK agent session initialized")
|
||||
|
||||
s.broadcastProcessingStatus()
|
||||
w.WriteHeader(http.StatusOK)
|
||||
}
|
||||
|
||||
// ObservationRequest is the request body for posting observations.
|
||||
type ObservationRequest struct {
|
||||
ClaudeSessionID string `json:"claudeSessionId"`
|
||||
Project string `json:"project"`
|
||||
ToolName string `json:"tool_name"`
|
||||
ToolInput any `json:"tool_input"`
|
||||
ToolResponse any `json:"tool_response"`
|
||||
CWD string `json:"cwd"`
|
||||
}
|
||||
|
||||
// handleObservation handles observation posting from post-tool-use hook.
|
||||
func (s *Service) handleObservation(w http.ResponseWriter, r *http.Request) {
|
||||
var req ObservationRequest
|
||||
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
||||
http.Error(w, err.Error(), http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
// Reject requests from internal CLI calls (no session ID = internal callClaudeCLI)
|
||||
if req.ClaudeSessionID == "" {
|
||||
w.WriteHeader(http.StatusOK)
|
||||
return
|
||||
}
|
||||
|
||||
|
||||
// Find session
|
||||
sess, err := s.sessionStore.FindAnySDKSession(r.Context(), req.ClaudeSessionID)
|
||||
if err != nil {
|
||||
http.Error(w, err.Error(), http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
if sess == nil {
|
||||
// Create session on-the-fly with project from request
|
||||
id, err := s.sessionStore.CreateSDKSession(r.Context(), req.ClaudeSessionID, req.Project, "")
|
||||
if err != nil {
|
||||
http.Error(w, err.Error(), http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
sess, _ = s.sessionStore.GetSessionByID(r.Context(), id)
|
||||
}
|
||||
|
||||
// Queue observation
|
||||
if err := s.sessionManager.QueueObservation(r.Context(), sess.ID, session.ObservationData{
|
||||
ToolName: req.ToolName,
|
||||
ToolInput: req.ToolInput,
|
||||
ToolResponse: req.ToolResponse,
|
||||
CWD: req.CWD,
|
||||
}); err != nil {
|
||||
http.Error(w, err.Error(), http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
|
||||
s.broadcastProcessingStatus()
|
||||
w.WriteHeader(http.StatusOK)
|
||||
}
|
||||
|
||||
// SubagentCompleteRequest is the request body for subagent completion.
|
||||
type SubagentCompleteRequest struct {
|
||||
ClaudeSessionID string `json:"claudeSessionId"`
|
||||
Project string `json:"project"`
|
||||
}
|
||||
|
||||
// handleSubagentComplete handles subagent/Task completion notifications.
|
||||
// This triggers immediate processing of any queued observations from the subagent.
|
||||
func (s *Service) handleSubagentComplete(w http.ResponseWriter, r *http.Request) {
|
||||
var req SubagentCompleteRequest
|
||||
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
||||
http.Error(w, err.Error(), http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
|
||||
// Find session
|
||||
sess, err := s.sessionStore.FindAnySDKSession(r.Context(), req.ClaudeSessionID)
|
||||
if err != nil || sess == nil {
|
||||
// Session not found - subagent may have been in a different context
|
||||
log.Debug().
|
||||
Str("claudeSessionId", req.ClaudeSessionID).
|
||||
Msg("Subagent complete - no active session found")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
return
|
||||
}
|
||||
|
||||
// Trigger immediate processing of queued observations
|
||||
messages := s.sessionManager.DrainMessages(sess.ID)
|
||||
if len(messages) > 0 && s.processor != nil {
|
||||
log.Info().
|
||||
Int64("sessionId", sess.ID).
|
||||
Int("messages", len(messages)).
|
||||
Msg("Processing queued observations from subagent")
|
||||
|
||||
for _, msg := range messages {
|
||||
if msg.Type == session.MessageTypeObservation && msg.Observation != nil {
|
||||
err := s.processor.ProcessObservation(
|
||||
r.Context(),
|
||||
sess.SDKSessionID.String,
|
||||
sess.Project,
|
||||
msg.Observation.ToolName,
|
||||
msg.Observation.ToolInput,
|
||||
msg.Observation.ToolResponse,
|
||||
msg.Observation.PromptNumber,
|
||||
msg.Observation.CWD,
|
||||
)
|
||||
if err != nil {
|
||||
log.Error().Err(err).
|
||||
Str("tool", msg.Observation.ToolName).
|
||||
Msg("Failed to process subagent observation")
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
s.broadcastProcessingStatus()
|
||||
w.WriteHeader(http.StatusOK)
|
||||
}
|
||||
|
||||
// handleGetSessionByClaudeID looks up a session by Claude session ID.
|
||||
func (s *Service) handleGetSessionByClaudeID(w http.ResponseWriter, r *http.Request) {
|
||||
claudeSessionID := r.URL.Query().Get("claudeSessionId")
|
||||
if claudeSessionID == "" {
|
||||
http.Error(w, "claudeSessionId required", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
|
||||
session, err := s.sessionStore.FindAnySDKSession(r.Context(), claudeSessionID)
|
||||
if err != nil {
|
||||
http.Error(w, err.Error(), http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
if session == nil {
|
||||
http.Error(w, "session not found", http.StatusNotFound)
|
||||
return
|
||||
}
|
||||
|
||||
writeJSON(w, session)
|
||||
}
|
||||
|
||||
// SummarizeRequest is the request body for summarize requests.
|
||||
type SummarizeRequest struct {
|
||||
LastUserMessage string `json:"lastUserMessage"`
|
||||
LastAssistantMessage string `json:"lastAssistantMessage"`
|
||||
}
|
||||
|
||||
// handleSummarize handles summarize requests from stop hook.
|
||||
func (s *Service) handleSummarize(w http.ResponseWriter, r *http.Request) {
|
||||
idStr := chi.URLParam(r, "id")
|
||||
id, err := strconv.ParseInt(idStr, 10, 64)
|
||||
if err != nil {
|
||||
http.Error(w, "invalid session id", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
|
||||
var req SummarizeRequest
|
||||
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
||||
http.Error(w, err.Error(), http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
|
||||
// Queue summarize request
|
||||
if err := s.sessionManager.QueueSummarize(r.Context(), id, req.LastUserMessage, req.LastAssistantMessage); err != nil {
|
||||
http.Error(w, err.Error(), http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
|
||||
s.broadcastProcessingStatus()
|
||||
w.WriteHeader(http.StatusOK)
|
||||
}
|
||||
@@ -66,6 +66,7 @@ func testService(t *testing.T) (*Service, func()) {
|
||||
cancel: cancel,
|
||||
startTime: time.Now(),
|
||||
retrievalStats: make(map[string]*RetrievalStats),
|
||||
cachedObsCounts: make(map[string]cachedCount),
|
||||
}
|
||||
|
||||
svc.setupRoutes()
|
||||
@@ -345,11 +346,13 @@ func TestHandleGetObservations_Limit(t *testing.T) {
|
||||
|
||||
assert.Equal(t, http.StatusOK, rec.Code)
|
||||
|
||||
// Parse as generic JSON array since the model uses custom marshaling
|
||||
var observations []map[string]interface{}
|
||||
err := json.Unmarshal(rec.Body.Bytes(), &observations)
|
||||
// Parse as object with observations key (API returns wrapped response)
|
||||
var response map[string]interface{}
|
||||
err := json.Unmarshal(rec.Body.Bytes(), &response)
|
||||
require.NoError(t, err)
|
||||
|
||||
observations, ok := response["observations"].([]interface{})
|
||||
require.True(t, ok, "expected observations array in response")
|
||||
assert.Len(t, observations, 10)
|
||||
}
|
||||
|
||||
@@ -456,14 +459,13 @@ func TestHandleHealth_ReturnsVersion(t *testing.T) {
|
||||
|
||||
svc.handleHealth(rec, req)
|
||||
|
||||
assert.Equal(t, http.StatusOK, rec.Code)
|
||||
|
||||
var response map[string]interface{}
|
||||
err := json.Unmarshal(rec.Body.Bytes(), &response)
|
||||
require.NoError(t, err)
|
||||
|
||||
assert.Equal(t, "ready", response["status"])
|
||||
assert.Equal(t, "test-version-1.2.3", response["version"])
|
||||
// Status may be "degraded" if embedSvc is nil in test, but version is always present
|
||||
assert.Contains(t, []string{"ready", "degraded"}, response["status"])
|
||||
}
|
||||
|
||||
func TestHandleVersion(t *testing.T) {
|
||||
@@ -1135,10 +1137,13 @@ func TestHandleGetObservations_DefaultLimit(t *testing.T) {
|
||||
|
||||
assert.Equal(t, http.StatusOK, rec.Code)
|
||||
|
||||
var observations []map[string]interface{}
|
||||
err := json.Unmarshal(rec.Body.Bytes(), &observations)
|
||||
var response map[string]interface{}
|
||||
err := json.Unmarshal(rec.Body.Bytes(), &response)
|
||||
require.NoError(t, err)
|
||||
|
||||
observations, ok := response["observations"].([]interface{})
|
||||
require.True(t, ok, "expected observations array in response")
|
||||
|
||||
// Should return default limit (100)
|
||||
assert.LessOrEqual(t, len(observations), DefaultObservationsLimit)
|
||||
}
|
||||
@@ -1159,10 +1164,12 @@ func TestHandleGetObservations_FilterByProject(t *testing.T) {
|
||||
|
||||
assert.Equal(t, http.StatusOK, rec.Code)
|
||||
|
||||
var observations []map[string]interface{}
|
||||
err := json.Unmarshal(rec.Body.Bytes(), &observations)
|
||||
var response map[string]interface{}
|
||||
err := json.Unmarshal(rec.Body.Bytes(), &response)
|
||||
require.NoError(t, err)
|
||||
|
||||
observations, ok := response["observations"].([]interface{})
|
||||
require.True(t, ok, "expected observations array in response")
|
||||
assert.Len(t, observations, 2)
|
||||
}
|
||||
|
||||
@@ -1412,10 +1419,12 @@ func TestHandleGetObservations(t *testing.T) {
|
||||
|
||||
assert.Equal(t, http.StatusOK, rec.Code)
|
||||
|
||||
var observations []map[string]interface{}
|
||||
err := json.Unmarshal(rec.Body.Bytes(), &observations)
|
||||
var response map[string]interface{}
|
||||
err := json.Unmarshal(rec.Body.Bytes(), &response)
|
||||
require.NoError(t, err)
|
||||
|
||||
observations, ok := response["observations"].([]interface{})
|
||||
require.True(t, ok, "expected observations array in response")
|
||||
assert.GreaterOrEqual(t, len(observations), 2)
|
||||
}
|
||||
|
||||
@@ -2018,13 +2027,14 @@ func TestHandleHealth_NotReady(t *testing.T) {
|
||||
|
||||
svc.handleHealth(rec, req)
|
||||
|
||||
assert.Equal(t, http.StatusOK, rec.Code)
|
||||
assert.Equal(t, http.StatusServiceUnavailable, rec.Code)
|
||||
|
||||
var response map[string]interface{}
|
||||
err := json.Unmarshal(rec.Body.Bytes(), &response)
|
||||
require.NoError(t, err)
|
||||
|
||||
assert.Equal(t, "starting", response["status"])
|
||||
assert.Equal(t, "initializing", response["status"])
|
||||
assert.Equal(t, false, response["ready"])
|
||||
}
|
||||
|
||||
// TestHandleContextInject_EmptyProject tests context inject with empty project.
|
||||
@@ -2389,7 +2399,12 @@ func TestHandleHealthEndpoint(t *testing.T) {
|
||||
|
||||
svc.router.ServeHTTP(rec, req)
|
||||
|
||||
assert.Equal(t, http.StatusOK, rec.Code)
|
||||
// Response is valid JSON with health details
|
||||
var response map[string]interface{}
|
||||
err := json.Unmarshal(rec.Body.Bytes(), &response)
|
||||
require.NoError(t, err)
|
||||
assert.NotNil(t, response["status"])
|
||||
assert.NotNil(t, response["version"])
|
||||
}
|
||||
|
||||
// TestHandleSelfCheckEndpoint tests self-check endpoint via router.
|
||||
@@ -2697,10 +2712,13 @@ func TestHandleGetObservations_EmptyResult(t *testing.T) {
|
||||
|
||||
assert.Equal(t, http.StatusOK, rec.Code)
|
||||
|
||||
// Should return empty array, not null
|
||||
var obs []interface{}
|
||||
err := json.Unmarshal(rec.Body.Bytes(), &obs)
|
||||
// Should return empty array within observations key, not null
|
||||
var response map[string]interface{}
|
||||
err := json.Unmarshal(rec.Body.Bytes(), &response)
|
||||
require.NoError(t, err)
|
||||
|
||||
obs, ok := response["observations"].([]interface{})
|
||||
require.True(t, ok, "expected observations array in response")
|
||||
assert.NotNil(t, obs)
|
||||
}
|
||||
|
||||
@@ -2881,12 +2899,18 @@ func TestHandleHealth(t *testing.T) {
|
||||
|
||||
svc.router.ServeHTTP(rec, req)
|
||||
|
||||
assert.Equal(t, http.StatusOK, rec.Code)
|
||||
|
||||
var response map[string]interface{}
|
||||
err := json.Unmarshal(rec.Body.Bytes(), &response)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "ready", response["status"])
|
||||
|
||||
// Test service has store set but no embedSvc, so status is "degraded"
|
||||
assert.Contains(t, []string{"ready", "degraded"}, response["status"])
|
||||
assert.NotNil(t, response["version"])
|
||||
assert.NotNil(t, response["uptime_seconds"])
|
||||
assert.NotNil(t, response["active_sessions"])
|
||||
assert.NotNil(t, response["db_status"])
|
||||
assert.NotNil(t, response["embedding_status"])
|
||||
assert.NotNil(t, response["ready"])
|
||||
}
|
||||
|
||||
// TestHandleSessionInit_ValidRequest tests session init with valid request.
|
||||
|
||||
@@ -0,0 +1,251 @@
|
||||
// Package worker provides update and restart HTTP handlers.
|
||||
package worker
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/rs/zerolog/log"
|
||||
)
|
||||
|
||||
// handleUpdateCheck checks for available updates.
|
||||
func (s *Service) handleUpdateCheck(w http.ResponseWriter, r *http.Request) {
|
||||
info, err := s.updater.CheckForUpdate(r.Context())
|
||||
if err != nil {
|
||||
// Return a proper JSON response for errors instead of 500
|
||||
// This allows the frontend to handle it gracefully
|
||||
writeJSON(w, map[string]any{
|
||||
"available": false,
|
||||
"current_version": s.version,
|
||||
"error": err.Error(),
|
||||
"rate_limited": strings.Contains(err.Error(), "403"),
|
||||
})
|
||||
return
|
||||
}
|
||||
writeJSON(w, info)
|
||||
}
|
||||
|
||||
// handleUpdateApply downloads and applies an available update.
|
||||
func (s *Service) handleUpdateApply(w http.ResponseWriter, r *http.Request) {
|
||||
// First check for update
|
||||
info, err := s.updater.CheckForUpdate(r.Context())
|
||||
if err != nil {
|
||||
http.Error(w, err.Error(), http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
|
||||
if !info.Available {
|
||||
writeJSON(w, map[string]any{
|
||||
"success": false,
|
||||
"message": "No update available",
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
// Apply update in background with tracking for graceful shutdown
|
||||
s.wg.Go(func() {
|
||||
if err := s.updater.ApplyUpdate(s.ctx, info); err != nil {
|
||||
log.Error().Err(err).Msg("Update failed")
|
||||
}
|
||||
})
|
||||
|
||||
writeJSON(w, map[string]any{
|
||||
"success": true,
|
||||
"message": "Update started",
|
||||
"version": info.LatestVersion,
|
||||
})
|
||||
}
|
||||
|
||||
// handleUpdateStatus returns the current update status.
|
||||
func (s *Service) handleUpdateStatus(w http.ResponseWriter, r *http.Request) {
|
||||
status := s.updater.GetStatus()
|
||||
writeJSON(w, status)
|
||||
}
|
||||
|
||||
// ComponentHealth represents the health status of a single component.
|
||||
type ComponentHealth struct {
|
||||
Name string `json:"name"`
|
||||
Status string `json:"status"` // "healthy", "degraded", "unhealthy"
|
||||
Message string `json:"message,omitempty"`
|
||||
}
|
||||
|
||||
// SelfCheckResponse contains the health status of all components.
|
||||
type SelfCheckResponse struct {
|
||||
Overall string `json:"overall"` // "healthy", "degraded", "unhealthy"
|
||||
Version string `json:"version"`
|
||||
Uptime string `json:"uptime"`
|
||||
Components []ComponentHealth `json:"components"`
|
||||
}
|
||||
|
||||
// handleSelfCheck returns the health status of all components.
|
||||
func (s *Service) handleSelfCheck(w http.ResponseWriter, r *http.Request) {
|
||||
components := []ComponentHealth{}
|
||||
overall := "healthy"
|
||||
|
||||
// Check Worker Service
|
||||
workerStatus := ComponentHealth{Name: "Worker Service", Status: "healthy"}
|
||||
if !s.ready.Load() {
|
||||
if err := s.GetInitError(); err != nil {
|
||||
workerStatus.Status = "unhealthy"
|
||||
workerStatus.Message = err.Error()
|
||||
overall = "unhealthy"
|
||||
} else {
|
||||
workerStatus.Status = "degraded"
|
||||
workerStatus.Message = "Initializing"
|
||||
if overall == "healthy" {
|
||||
overall = "degraded"
|
||||
}
|
||||
}
|
||||
}
|
||||
components = append(components, workerStatus)
|
||||
|
||||
// Check SQLite Database
|
||||
dbStatus := ComponentHealth{Name: "SQLite Database", Status: "healthy"}
|
||||
if s.store == nil {
|
||||
dbStatus.Status = "unhealthy"
|
||||
dbStatus.Message = "Not initialized"
|
||||
overall = "unhealthy"
|
||||
} else if err := s.store.Ping(); err != nil {
|
||||
dbStatus.Status = "unhealthy"
|
||||
dbStatus.Message = err.Error()
|
||||
overall = "unhealthy"
|
||||
}
|
||||
components = append(components, dbStatus)
|
||||
|
||||
// Check Vector DB (sqlite-vec)
|
||||
vectorStatus := ComponentHealth{Name: "Vector DB", Status: "healthy"}
|
||||
if s.vectorClient == nil {
|
||||
vectorStatus.Status = "degraded"
|
||||
vectorStatus.Message = "Not configured"
|
||||
if overall == "healthy" {
|
||||
overall = "degraded"
|
||||
}
|
||||
} else if !s.vectorClient.IsConnected() {
|
||||
vectorStatus.Status = "degraded"
|
||||
vectorStatus.Message = "Not connected"
|
||||
if overall == "healthy" {
|
||||
overall = "degraded"
|
||||
}
|
||||
}
|
||||
components = append(components, vectorStatus)
|
||||
|
||||
// Check SDK Processor
|
||||
sdkStatus := ComponentHealth{Name: "SDK Processor", Status: "healthy"}
|
||||
if s.processor == nil {
|
||||
sdkStatus.Status = "degraded"
|
||||
sdkStatus.Message = "Not initialized"
|
||||
if overall == "healthy" {
|
||||
overall = "degraded"
|
||||
}
|
||||
} else if !s.processor.IsAvailable() {
|
||||
sdkStatus.Status = "degraded"
|
||||
sdkStatus.Message = "Claude CLI not available"
|
||||
if overall == "healthy" {
|
||||
overall = "degraded"
|
||||
}
|
||||
}
|
||||
components = append(components, sdkStatus)
|
||||
|
||||
// Check SSE Broadcaster
|
||||
sseStatus := ComponentHealth{Name: "SSE Broadcaster", Status: "healthy"}
|
||||
if s.sseBroadcaster == nil {
|
||||
sseStatus.Status = "unhealthy"
|
||||
sseStatus.Message = "Not initialized"
|
||||
overall = "unhealthy"
|
||||
}
|
||||
components = append(components, sseStatus)
|
||||
|
||||
// Check Cross-Encoder Reranker
|
||||
rerankerStatus := ComponentHealth{Name: "Cross-Encoder Reranker", Status: "healthy"}
|
||||
if !s.config.RerankingEnabled {
|
||||
rerankerStatus.Status = "degraded"
|
||||
rerankerStatus.Message = "Disabled in config"
|
||||
if overall == "healthy" {
|
||||
overall = "degraded"
|
||||
}
|
||||
} else if s.reranker == nil {
|
||||
rerankerStatus.Status = "degraded"
|
||||
rerankerStatus.Message = "Not initialized"
|
||||
if overall == "healthy" {
|
||||
overall = "degraded"
|
||||
}
|
||||
} else {
|
||||
// Verify reranker is functional using Score
|
||||
_, normalizedScore, err := s.reranker.Score("test query", "test document")
|
||||
if err != nil {
|
||||
rerankerStatus.Status = "unhealthy"
|
||||
rerankerStatus.Message = fmt.Sprintf("Score check failed: %v", err)
|
||||
if overall == "healthy" {
|
||||
overall = "degraded"
|
||||
}
|
||||
} else {
|
||||
rerankerStatus.Message = fmt.Sprintf("Score check passed (%.4f)", normalizedScore)
|
||||
}
|
||||
}
|
||||
components = append(components, rerankerStatus)
|
||||
|
||||
// Calculate uptime
|
||||
uptime := time.Since(s.startTime).Round(time.Second).String()
|
||||
|
||||
writeJSON(w, SelfCheckResponse{
|
||||
Overall: overall,
|
||||
Version: s.version,
|
||||
Uptime: uptime,
|
||||
Components: components,
|
||||
})
|
||||
}
|
||||
|
||||
// handleUpdateRestart restarts the worker with the new binary (after update).
|
||||
func (s *Service) handleUpdateRestart(w http.ResponseWriter, r *http.Request) {
|
||||
status := s.updater.GetStatus()
|
||||
if status.State != "done" {
|
||||
http.Error(w, "no update has been applied", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
|
||||
// Send response before restarting
|
||||
writeJSON(w, map[string]any{
|
||||
"success": true,
|
||||
"message": "Restarting worker...",
|
||||
})
|
||||
|
||||
// Flush the response
|
||||
if f, ok := w.(http.Flusher); ok {
|
||||
f.Flush()
|
||||
}
|
||||
|
||||
// Restart in background after response is sent
|
||||
go func() {
|
||||
if err := s.updater.Restart(); err != nil {
|
||||
log.Error().Err(err).Msg("Failed to restart worker")
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
// handleRestart restarts the worker process (general restart, not tied to update).
|
||||
func (s *Service) handleRestart(w http.ResponseWriter, r *http.Request) {
|
||||
log.Info().Msg("Manual restart requested via API")
|
||||
|
||||
// Send response before restarting
|
||||
writeJSON(w, map[string]any{
|
||||
"success": true,
|
||||
"message": "Restarting worker...",
|
||||
"version": s.version,
|
||||
})
|
||||
|
||||
// Flush the response
|
||||
if f, ok := w.(http.Flusher); ok {
|
||||
f.Flush()
|
||||
}
|
||||
|
||||
// Restart in background after response is sent
|
||||
go func() {
|
||||
// Small delay to ensure response is sent
|
||||
time.Sleep(100 * time.Millisecond)
|
||||
if err := s.updater.Restart(); err != nil {
|
||||
log.Error().Err(err).Msg("Failed to restart worker")
|
||||
}
|
||||
}()
|
||||
}
|
||||
@@ -0,0 +1,333 @@
|
||||
// Package worker provides the main worker service for claude-mnemonic.
|
||||
package worker
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/rand"
|
||||
"encoding/hex"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"regexp"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// requestIDKey is the context key for request IDs.
|
||||
type requestIDKey struct{}
|
||||
|
||||
// projectNamePattern validates project names to prevent path traversal.
|
||||
var projectNamePattern = regexp.MustCompile(`^[a-zA-Z0-9_./-]+$`)
|
||||
|
||||
// allowedOrigins is the whitelist of origins allowed for CORS.
|
||||
// Uses exact matching to prevent bypass attacks like "evil-localhost.com".
|
||||
var allowedOrigins = map[string]bool{
|
||||
"http://localhost": true,
|
||||
"http://localhost:3000": true,
|
||||
"http://localhost:5173": true, // Vite dev server
|
||||
"http://localhost:37778": true, // Dashboard UI
|
||||
"http://127.0.0.1": true,
|
||||
"http://127.0.0.1:3000": true,
|
||||
"http://127.0.0.1:5173": true,
|
||||
"http://127.0.0.1:37778": true,
|
||||
}
|
||||
|
||||
// SecurityHeaders middleware adds essential security headers to all responses.
|
||||
// These protect against common web vulnerabilities.
|
||||
func SecurityHeaders(next http.Handler) http.Handler {
|
||||
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
// Prevent clickjacking
|
||||
w.Header().Set("X-Frame-Options", "DENY")
|
||||
|
||||
// Prevent MIME type sniffing
|
||||
w.Header().Set("X-Content-Type-Options", "nosniff")
|
||||
|
||||
// Enable XSS filter
|
||||
w.Header().Set("X-XSS-Protection", "1; mode=block")
|
||||
|
||||
// Restrict referrer information
|
||||
w.Header().Set("Referrer-Policy", "strict-origin-when-cross-origin")
|
||||
|
||||
// Content Security Policy - restrict to self
|
||||
w.Header().Set("Content-Security-Policy", "default-src 'self'")
|
||||
|
||||
// Permissions Policy - disable unnecessary features
|
||||
w.Header().Set("Permissions-Policy", "geolocation=(), microphone=(), camera=()")
|
||||
|
||||
// CORS: Use exact match whitelist to prevent bypass attacks
|
||||
origin := r.Header.Get("Origin")
|
||||
if allowedOrigins[origin] {
|
||||
w.Header().Set("Access-Control-Allow-Origin", origin)
|
||||
w.Header().Set("Access-Control-Allow-Credentials", "true")
|
||||
w.Header().Set("Access-Control-Allow-Methods", "GET, POST, PUT, DELETE, OPTIONS")
|
||||
w.Header().Set("Access-Control-Allow-Headers", "Content-Type, X-Auth-Token, Authorization, X-Request-ID")
|
||||
}
|
||||
|
||||
// Handle preflight requests
|
||||
if r.Method == "OPTIONS" {
|
||||
w.WriteHeader(http.StatusNoContent)
|
||||
return
|
||||
}
|
||||
|
||||
next.ServeHTTP(w, r)
|
||||
})
|
||||
}
|
||||
|
||||
// MaxBodySize middleware limits the size of incoming request bodies.
|
||||
// This prevents denial of service attacks via large payloads.
|
||||
func MaxBodySize(maxBytes int64) func(http.Handler) http.Handler {
|
||||
return func(next http.Handler) http.Handler {
|
||||
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.ContentLength > maxBytes {
|
||||
http.Error(w, "request body too large", http.StatusRequestEntityTooLarge)
|
||||
return
|
||||
}
|
||||
r.Body = http.MaxBytesReader(w, r.Body, maxBytes)
|
||||
next.ServeHTTP(w, r)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TokenAuth provides simple token-based authentication for localhost services.
|
||||
// The token is generated at startup and must be provided in the X-Auth-Token header.
|
||||
type TokenAuth struct {
|
||||
ExemptPaths map[string]bool
|
||||
token string
|
||||
mu sync.RWMutex
|
||||
enabled bool
|
||||
}
|
||||
|
||||
// NewTokenAuth creates a new TokenAuth with a randomly generated token.
|
||||
// If enabled is false, authentication is skipped (useful for development).
|
||||
func NewTokenAuth(enabled bool) (*TokenAuth, error) {
|
||||
ta := &TokenAuth{
|
||||
enabled: enabled,
|
||||
ExemptPaths: map[string]bool{
|
||||
"/health": true,
|
||||
"/api/health": true,
|
||||
"/api/ready": true,
|
||||
},
|
||||
}
|
||||
|
||||
if enabled {
|
||||
// Generate 32-byte random token
|
||||
tokenBytes := make([]byte, 32)
|
||||
if _, err := rand.Read(tokenBytes); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
ta.token = hex.EncodeToString(tokenBytes)
|
||||
}
|
||||
|
||||
return ta, nil
|
||||
}
|
||||
|
||||
// Token returns the authentication token.
|
||||
// Returns empty string if authentication is disabled.
|
||||
func (ta *TokenAuth) Token() string {
|
||||
ta.mu.RLock()
|
||||
defer ta.mu.RUnlock()
|
||||
return ta.token
|
||||
}
|
||||
|
||||
// IsEnabled returns whether token authentication is enabled.
|
||||
func (ta *TokenAuth) IsEnabled() bool {
|
||||
ta.mu.RLock()
|
||||
defer ta.mu.RUnlock()
|
||||
return ta.enabled
|
||||
}
|
||||
|
||||
// Middleware returns HTTP middleware that enforces token authentication.
|
||||
func (ta *TokenAuth) Middleware(next http.Handler) http.Handler {
|
||||
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
ta.mu.RLock()
|
||||
enabled := ta.enabled
|
||||
token := ta.token
|
||||
exempt := ta.ExemptPaths[r.URL.Path]
|
||||
ta.mu.RUnlock()
|
||||
|
||||
// Skip auth if disabled or path is exempt
|
||||
if !enabled || exempt {
|
||||
next.ServeHTTP(w, r)
|
||||
return
|
||||
}
|
||||
|
||||
// Check for token in header
|
||||
providedToken := r.Header.Get("X-Auth-Token")
|
||||
if providedToken == "" {
|
||||
// Also check Authorization header with Bearer scheme
|
||||
auth := r.Header.Get("Authorization")
|
||||
if bearer, found := strings.CutPrefix(auth, "Bearer "); found {
|
||||
providedToken = bearer
|
||||
}
|
||||
}
|
||||
|
||||
if providedToken != token {
|
||||
http.Error(w, "unauthorized", http.StatusUnauthorized)
|
||||
return
|
||||
}
|
||||
|
||||
next.ServeHTTP(w, r)
|
||||
})
|
||||
}
|
||||
|
||||
// ExpensiveOperationLimiter provides stricter rate limiting for expensive operations.
|
||||
// It wraps the base per-client rate limiter with additional per-operation limits.
|
||||
type ExpensiveOperationLimiter struct {
|
||||
// Track last execution time per operation type
|
||||
lastRebuild int64 // Unix timestamp
|
||||
rebuildCooldown int64 // Minimum seconds between rebuilds
|
||||
|
||||
mu sync.Mutex
|
||||
}
|
||||
|
||||
// NewExpensiveOperationLimiter creates a limiter for expensive operations.
|
||||
func NewExpensiveOperationLimiter() *ExpensiveOperationLimiter {
|
||||
return &ExpensiveOperationLimiter{
|
||||
rebuildCooldown: 300, // 5 minutes between rebuilds
|
||||
}
|
||||
}
|
||||
|
||||
// CanRebuild checks if a vector rebuild operation is allowed.
|
||||
// Returns false if a rebuild was triggered too recently.
|
||||
func (eol *ExpensiveOperationLimiter) CanRebuild() bool {
|
||||
eol.mu.Lock()
|
||||
defer eol.mu.Unlock()
|
||||
|
||||
now := unixNow()
|
||||
if now-eol.lastRebuild < eol.rebuildCooldown {
|
||||
return false
|
||||
}
|
||||
eol.lastRebuild = now
|
||||
return true
|
||||
}
|
||||
|
||||
// unixNow returns current Unix timestamp.
|
||||
// Separated for easier testing.
|
||||
func unixNow() int64 {
|
||||
return time.Now().Unix()
|
||||
}
|
||||
|
||||
// RequestID middleware adds a unique request ID to each request.
|
||||
// The ID is added to the context and response headers for tracing.
|
||||
func RequestID(next http.Handler) http.Handler {
|
||||
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
// Check for existing request ID from client
|
||||
requestID := r.Header.Get("X-Request-ID")
|
||||
if requestID == "" {
|
||||
// Generate new request ID
|
||||
idBytes := make([]byte, 8)
|
||||
if _, err := rand.Read(idBytes); err == nil {
|
||||
requestID = hex.EncodeToString(idBytes)
|
||||
} else {
|
||||
requestID = fmt.Sprintf("%d", time.Now().UnixNano())
|
||||
}
|
||||
}
|
||||
|
||||
// Add to response header
|
||||
w.Header().Set("X-Request-ID", requestID)
|
||||
|
||||
// Add to context
|
||||
ctx := context.WithValue(r.Context(), requestIDKey{}, requestID)
|
||||
next.ServeHTTP(w, r.WithContext(ctx))
|
||||
})
|
||||
}
|
||||
|
||||
// GetRequestID retrieves the request ID from the context.
|
||||
func GetRequestID(ctx context.Context) string {
|
||||
if id, ok := ctx.Value(requestIDKey{}).(string); ok {
|
||||
return id
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// RequireJSONContentType middleware validates that POST/PUT/PATCH requests
|
||||
// have application/json Content-Type header.
|
||||
func RequireJSONContentType(next http.Handler) http.Handler {
|
||||
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
// Only check for methods that typically have bodies
|
||||
if r.Method == "POST" || r.Method == "PUT" || r.Method == "PATCH" {
|
||||
ct := r.Header.Get("Content-Type")
|
||||
// Allow empty Content-Type for requests without body
|
||||
if ct != "" && !strings.HasPrefix(ct, "application/json") {
|
||||
http.Error(w, "Content-Type must be application/json", http.StatusUnsupportedMediaType)
|
||||
return
|
||||
}
|
||||
}
|
||||
next.ServeHTTP(w, r)
|
||||
})
|
||||
}
|
||||
|
||||
// ValidateProjectName checks if a project name is safe to use.
|
||||
// Returns an error if the name contains path traversal or invalid characters.
|
||||
func ValidateProjectName(project string) error {
|
||||
if project == "" {
|
||||
return nil // Empty is allowed (means no filter)
|
||||
}
|
||||
|
||||
// Check for path traversal
|
||||
if strings.Contains(project, "..") {
|
||||
return fmt.Errorf("invalid project name: path traversal detected")
|
||||
}
|
||||
|
||||
// Check for valid characters
|
||||
if !projectNamePattern.MatchString(project) {
|
||||
return fmt.Errorf("invalid project name: only alphanumeric, underscore, dash, dot, and slash allowed")
|
||||
}
|
||||
|
||||
// Max length check
|
||||
if len(project) > 500 {
|
||||
return fmt.Errorf("project name too long (max 500 chars)")
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// BulkOperationLimiter provides rate limiting for bulk operations.
|
||||
// Prevents DoS via repeated bulk requests.
|
||||
type BulkOperationLimiter struct {
|
||||
lastBulkOp int64 // Unix timestamp
|
||||
cooldown int64 // Minimum seconds between operations
|
||||
|
||||
mu sync.Mutex
|
||||
}
|
||||
|
||||
// NewBulkOperationLimiter creates a limiter for bulk operations.
|
||||
func NewBulkOperationLimiter(cooldownSeconds int64) *BulkOperationLimiter {
|
||||
return &BulkOperationLimiter{
|
||||
cooldown: cooldownSeconds,
|
||||
}
|
||||
}
|
||||
|
||||
// CanExecute checks if a bulk operation is allowed.
|
||||
// Returns false if a bulk operation was triggered too recently.
|
||||
func (bol *BulkOperationLimiter) CanExecute() bool {
|
||||
bol.mu.Lock()
|
||||
defer bol.mu.Unlock()
|
||||
|
||||
now := unixNow()
|
||||
if now-bol.lastBulkOp < bol.cooldown {
|
||||
return false
|
||||
}
|
||||
bol.lastBulkOp = now
|
||||
return true
|
||||
}
|
||||
|
||||
// TimeSinceLastOp returns seconds since the last bulk operation.
|
||||
func (bol *BulkOperationLimiter) TimeSinceLastOp() int64 {
|
||||
bol.mu.Lock()
|
||||
defer bol.mu.Unlock()
|
||||
return unixNow() - bol.lastBulkOp
|
||||
}
|
||||
|
||||
// CooldownRemaining returns seconds remaining in the cooldown period.
|
||||
// Returns 0 if no cooldown is active.
|
||||
func (bol *BulkOperationLimiter) CooldownRemaining() int64 {
|
||||
bol.mu.Lock()
|
||||
defer bol.mu.Unlock()
|
||||
|
||||
remaining := bol.cooldown - (unixNow() - bol.lastBulkOp)
|
||||
if remaining < 0 {
|
||||
return 0
|
||||
}
|
||||
return remaining
|
||||
}
|
||||
@@ -0,0 +1,515 @@
|
||||
package worker
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestSecurityHeaders(t *testing.T) {
|
||||
handler := SecurityHeaders(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.WriteHeader(http.StatusOK)
|
||||
}))
|
||||
|
||||
req := httptest.NewRequest("GET", "/test", nil)
|
||||
rr := httptest.NewRecorder()
|
||||
|
||||
handler.ServeHTTP(rr, req)
|
||||
|
||||
// Check all security headers are set
|
||||
tests := []struct {
|
||||
header string
|
||||
expected string
|
||||
}{
|
||||
{"X-Frame-Options", "DENY"},
|
||||
{"X-Content-Type-Options", "nosniff"},
|
||||
{"X-XSS-Protection", "1; mode=block"},
|
||||
{"Referrer-Policy", "strict-origin-when-cross-origin"},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
if got := rr.Header().Get(tt.header); got != tt.expected {
|
||||
t.Errorf("SecurityHeaders() %s = %q, want %q", tt.header, got, tt.expected)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestSecurityHeaders_CORS(t *testing.T) {
|
||||
handler := SecurityHeaders(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.WriteHeader(http.StatusOK)
|
||||
}))
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
origin string
|
||||
expectedOrigin string
|
||||
expectCORS bool
|
||||
}{
|
||||
{
|
||||
name: "localhost:37778 origin allowed",
|
||||
origin: "http://localhost:37778",
|
||||
expectCORS: true,
|
||||
expectedOrigin: "http://localhost:37778",
|
||||
},
|
||||
{
|
||||
name: "127.0.0.1:5173 origin allowed",
|
||||
origin: "http://127.0.0.1:5173",
|
||||
expectCORS: true,
|
||||
expectedOrigin: "http://127.0.0.1:5173",
|
||||
},
|
||||
{
|
||||
name: "localhost without port allowed",
|
||||
origin: "http://localhost",
|
||||
expectCORS: true,
|
||||
expectedOrigin: "http://localhost",
|
||||
},
|
||||
{
|
||||
name: "external origin blocked",
|
||||
origin: "http://evil.com",
|
||||
expectCORS: false,
|
||||
},
|
||||
{
|
||||
name: "evil-localhost.com bypass attempt blocked",
|
||||
origin: "http://evil-localhost.com",
|
||||
expectCORS: false,
|
||||
},
|
||||
{
|
||||
name: "localhost subdomain bypass attempt blocked",
|
||||
origin: "http://localhost.evil.com",
|
||||
expectCORS: false,
|
||||
},
|
||||
{
|
||||
name: "unknown localhost port blocked",
|
||||
origin: "http://localhost:9999",
|
||||
expectCORS: false,
|
||||
},
|
||||
{
|
||||
name: "no origin header",
|
||||
origin: "",
|
||||
expectCORS: false,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
req := httptest.NewRequest("GET", "/test", nil)
|
||||
if tt.origin != "" {
|
||||
req.Header.Set("Origin", tt.origin)
|
||||
}
|
||||
rr := httptest.NewRecorder()
|
||||
|
||||
handler.ServeHTTP(rr, req)
|
||||
|
||||
cors := rr.Header().Get("Access-Control-Allow-Origin")
|
||||
if tt.expectCORS {
|
||||
if cors != tt.expectedOrigin {
|
||||
t.Errorf("Expected CORS origin %q, got %q", tt.expectedOrigin, cors)
|
||||
}
|
||||
} else {
|
||||
if cors != "" {
|
||||
t.Errorf("Expected no CORS header, got %q", cors)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestMaxBodySize(t *testing.T) {
|
||||
maxSize := int64(100)
|
||||
handler := MaxBodySize(maxSize)(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.WriteHeader(http.StatusOK)
|
||||
}))
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
contentLength int64
|
||||
expectedStatus int
|
||||
}{
|
||||
{
|
||||
name: "within limit",
|
||||
contentLength: 50,
|
||||
expectedStatus: http.StatusOK,
|
||||
},
|
||||
{
|
||||
name: "at limit",
|
||||
contentLength: 100,
|
||||
expectedStatus: http.StatusOK,
|
||||
},
|
||||
{
|
||||
name: "exceeds limit",
|
||||
contentLength: 150,
|
||||
expectedStatus: http.StatusRequestEntityTooLarge,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
req := httptest.NewRequest("POST", "/test", nil)
|
||||
req.ContentLength = tt.contentLength
|
||||
rr := httptest.NewRecorder()
|
||||
|
||||
handler.ServeHTTP(rr, req)
|
||||
|
||||
if rr.Code != tt.expectedStatus {
|
||||
t.Errorf("MaxBodySize() status = %d, want %d", rr.Code, tt.expectedStatus)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestTokenAuth(t *testing.T) {
|
||||
t.Run("disabled auth allows all requests", func(t *testing.T) {
|
||||
ta, err := NewTokenAuth(false)
|
||||
if err != nil {
|
||||
t.Fatalf("NewTokenAuth() error = %v", err)
|
||||
}
|
||||
|
||||
handler := ta.Middleware(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.WriteHeader(http.StatusOK)
|
||||
}))
|
||||
|
||||
req := httptest.NewRequest("GET", "/test", nil)
|
||||
rr := httptest.NewRecorder()
|
||||
|
||||
handler.ServeHTTP(rr, req)
|
||||
|
||||
if rr.Code != http.StatusOK {
|
||||
t.Errorf("Expected OK with disabled auth, got %d", rr.Code)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("enabled auth requires token", func(t *testing.T) {
|
||||
ta, err := NewTokenAuth(true)
|
||||
if err != nil {
|
||||
t.Fatalf("NewTokenAuth() error = %v", err)
|
||||
}
|
||||
|
||||
handler := ta.Middleware(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.WriteHeader(http.StatusOK)
|
||||
}))
|
||||
|
||||
// Without token
|
||||
req := httptest.NewRequest("GET", "/test", nil)
|
||||
rr := httptest.NewRecorder()
|
||||
handler.ServeHTTP(rr, req)
|
||||
|
||||
if rr.Code != http.StatusUnauthorized {
|
||||
t.Errorf("Expected Unauthorized without token, got %d", rr.Code)
|
||||
}
|
||||
|
||||
// With correct token in X-Auth-Token header
|
||||
req = httptest.NewRequest("GET", "/test", nil)
|
||||
req.Header.Set("X-Auth-Token", ta.Token())
|
||||
rr = httptest.NewRecorder()
|
||||
handler.ServeHTTP(rr, req)
|
||||
|
||||
if rr.Code != http.StatusOK {
|
||||
t.Errorf("Expected OK with correct token, got %d", rr.Code)
|
||||
}
|
||||
|
||||
// With correct token in Authorization header
|
||||
req = httptest.NewRequest("GET", "/test", nil)
|
||||
req.Header.Set("Authorization", "Bearer "+ta.Token())
|
||||
rr = httptest.NewRecorder()
|
||||
handler.ServeHTTP(rr, req)
|
||||
|
||||
if rr.Code != http.StatusOK {
|
||||
t.Errorf("Expected OK with Bearer token, got %d", rr.Code)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("exempt paths skip auth", func(t *testing.T) {
|
||||
ta, err := NewTokenAuth(true)
|
||||
if err != nil {
|
||||
t.Fatalf("NewTokenAuth() error = %v", err)
|
||||
}
|
||||
|
||||
handler := ta.Middleware(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.WriteHeader(http.StatusOK)
|
||||
}))
|
||||
|
||||
exemptPaths := []string{"/health", "/api/health", "/api/ready"}
|
||||
for _, path := range exemptPaths {
|
||||
req := httptest.NewRequest("GET", path, nil)
|
||||
rr := httptest.NewRecorder()
|
||||
handler.ServeHTTP(rr, req)
|
||||
|
||||
if rr.Code != http.StatusOK {
|
||||
t.Errorf("Expected OK for exempt path %s, got %d", path, rr.Code)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestExpensiveOperationLimiter(t *testing.T) {
|
||||
limiter := NewExpensiveOperationLimiter()
|
||||
|
||||
// First rebuild should be allowed
|
||||
if !limiter.CanRebuild() {
|
||||
t.Error("First rebuild should be allowed")
|
||||
}
|
||||
|
||||
// Immediate second rebuild should be blocked
|
||||
if limiter.CanRebuild() {
|
||||
t.Error("Immediate second rebuild should be blocked")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRequestID(t *testing.T) {
|
||||
handler := RequestID(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
// Verify request ID is in context
|
||||
id := GetRequestID(r.Context())
|
||||
if id == "" {
|
||||
t.Error("Request ID should be set in context")
|
||||
}
|
||||
w.WriteHeader(http.StatusOK)
|
||||
}))
|
||||
|
||||
t.Run("generates new request ID", func(t *testing.T) {
|
||||
req := httptest.NewRequest("GET", "/test", nil)
|
||||
rr := httptest.NewRecorder()
|
||||
|
||||
handler.ServeHTTP(rr, req)
|
||||
|
||||
if rr.Header().Get("X-Request-ID") == "" {
|
||||
t.Error("X-Request-ID header should be set")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("uses existing request ID", func(t *testing.T) {
|
||||
req := httptest.NewRequest("GET", "/test", nil)
|
||||
req.Header.Set("X-Request-ID", "test-id-12345")
|
||||
rr := httptest.NewRecorder()
|
||||
|
||||
handler.ServeHTTP(rr, req)
|
||||
|
||||
if rr.Header().Get("X-Request-ID") != "test-id-12345" {
|
||||
t.Errorf("Expected X-Request-ID to be test-id-12345, got %s", rr.Header().Get("X-Request-ID"))
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestRequireJSONContentType(t *testing.T) {
|
||||
handler := RequireJSONContentType(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.WriteHeader(http.StatusOK)
|
||||
}))
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
method string
|
||||
contentType string
|
||||
expectedStatus int
|
||||
}{
|
||||
{
|
||||
name: "GET request without content-type",
|
||||
method: "GET",
|
||||
contentType: "",
|
||||
expectedStatus: http.StatusOK,
|
||||
},
|
||||
{
|
||||
name: "POST with application/json",
|
||||
method: "POST",
|
||||
contentType: "application/json",
|
||||
expectedStatus: http.StatusOK,
|
||||
},
|
||||
{
|
||||
name: "POST with application/json; charset=utf-8",
|
||||
method: "POST",
|
||||
contentType: "application/json; charset=utf-8",
|
||||
expectedStatus: http.StatusOK,
|
||||
},
|
||||
{
|
||||
name: "POST without content-type (empty body)",
|
||||
method: "POST",
|
||||
contentType: "",
|
||||
expectedStatus: http.StatusOK,
|
||||
},
|
||||
{
|
||||
name: "POST with text/plain rejected",
|
||||
method: "POST",
|
||||
contentType: "text/plain",
|
||||
expectedStatus: http.StatusUnsupportedMediaType,
|
||||
},
|
||||
{
|
||||
name: "PUT with application/xml rejected",
|
||||
method: "PUT",
|
||||
contentType: "application/xml",
|
||||
expectedStatus: http.StatusUnsupportedMediaType,
|
||||
},
|
||||
{
|
||||
name: "PATCH with form-urlencoded rejected",
|
||||
method: "PATCH",
|
||||
contentType: "application/x-www-form-urlencoded",
|
||||
expectedStatus: http.StatusUnsupportedMediaType,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
req := httptest.NewRequest(tt.method, "/test", nil)
|
||||
if tt.contentType != "" {
|
||||
req.Header.Set("Content-Type", tt.contentType)
|
||||
}
|
||||
rr := httptest.NewRecorder()
|
||||
|
||||
handler.ServeHTTP(rr, req)
|
||||
|
||||
if rr.Code != tt.expectedStatus {
|
||||
t.Errorf("Expected status %d, got %d", tt.expectedStatus, rr.Code)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestValidateProjectName(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
project string
|
||||
wantError bool
|
||||
}{
|
||||
{
|
||||
name: "empty project allowed",
|
||||
project: "",
|
||||
wantError: false,
|
||||
},
|
||||
{
|
||||
name: "simple project name",
|
||||
project: "my-project",
|
||||
wantError: false,
|
||||
},
|
||||
{
|
||||
name: "project with path",
|
||||
project: "org/my-project",
|
||||
wantError: false,
|
||||
},
|
||||
{
|
||||
name: "project with underscore",
|
||||
project: "my_project_v2",
|
||||
wantError: false,
|
||||
},
|
||||
{
|
||||
name: "project with dot",
|
||||
project: "my.project.name",
|
||||
wantError: false,
|
||||
},
|
||||
{
|
||||
name: "path traversal attack",
|
||||
project: "../../../etc/passwd",
|
||||
wantError: true,
|
||||
},
|
||||
{
|
||||
name: "hidden path traversal",
|
||||
project: "project/../../secret",
|
||||
wantError: true,
|
||||
},
|
||||
{
|
||||
name: "shell injection attempt",
|
||||
project: "project; rm -rf /",
|
||||
wantError: true,
|
||||
},
|
||||
{
|
||||
name: "backtick injection",
|
||||
project: "project`whoami`",
|
||||
wantError: true,
|
||||
},
|
||||
{
|
||||
name: "special characters",
|
||||
project: "project$HOME",
|
||||
wantError: true,
|
||||
},
|
||||
{
|
||||
name: "too long project name",
|
||||
project: string(make([]byte, 501)),
|
||||
wantError: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
err := ValidateProjectName(tt.project)
|
||||
if tt.wantError && err == nil {
|
||||
t.Errorf("Expected error for project %q, got nil", tt.project)
|
||||
}
|
||||
if !tt.wantError && err != nil {
|
||||
t.Errorf("Unexpected error for project %q: %v", tt.project, err)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestBulkOperationLimiter(t *testing.T) {
|
||||
limiter := NewBulkOperationLimiter(1) // 1 second cooldown for testing
|
||||
|
||||
// First operation should be allowed
|
||||
if !limiter.CanExecute() {
|
||||
t.Error("First bulk operation should be allowed")
|
||||
}
|
||||
|
||||
// Immediate second operation should be blocked
|
||||
if limiter.CanExecute() {
|
||||
t.Error("Immediate second bulk operation should be blocked")
|
||||
}
|
||||
|
||||
// Check cooldown remaining
|
||||
remaining := limiter.CooldownRemaining()
|
||||
if remaining <= 0 || remaining > 1 {
|
||||
t.Errorf("Expected cooldown remaining between 0-1 seconds, got %d", remaining)
|
||||
}
|
||||
|
||||
// Check time since last op
|
||||
since := limiter.TimeSinceLastOp()
|
||||
if since < 0 || since > 1 {
|
||||
t.Errorf("Expected time since last op between 0-1 seconds, got %d", since)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSecurityHeaders_CSP(t *testing.T) {
|
||||
handler := SecurityHeaders(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.WriteHeader(http.StatusOK)
|
||||
}))
|
||||
|
||||
req := httptest.NewRequest("GET", "/test", nil)
|
||||
rr := httptest.NewRecorder()
|
||||
|
||||
handler.ServeHTTP(rr, req)
|
||||
|
||||
// Check CSP header is set
|
||||
csp := rr.Header().Get("Content-Security-Policy")
|
||||
if csp == "" {
|
||||
t.Error("Content-Security-Policy header should be set")
|
||||
}
|
||||
if csp != "default-src 'self'" {
|
||||
t.Errorf("Expected CSP to be \"default-src 'self'\", got %q", csp)
|
||||
}
|
||||
|
||||
// Check Permissions-Policy header
|
||||
pp := rr.Header().Get("Permissions-Policy")
|
||||
if pp == "" {
|
||||
t.Error("Permissions-Policy header should be set")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSecurityHeaders_Preflight(t *testing.T) {
|
||||
handler := SecurityHeaders(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.WriteHeader(http.StatusOK)
|
||||
}))
|
||||
|
||||
req := httptest.NewRequest("OPTIONS", "/test", nil)
|
||||
req.Header.Set("Origin", "http://localhost:3000")
|
||||
rr := httptest.NewRecorder()
|
||||
|
||||
handler.ServeHTTP(rr, req)
|
||||
|
||||
// OPTIONS should return 204 No Content
|
||||
if rr.Code != http.StatusNoContent {
|
||||
t.Errorf("Expected status 204 for OPTIONS, got %d", rr.Code)
|
||||
}
|
||||
|
||||
// CORS headers should be set for allowed origin
|
||||
if rr.Header().Get("Access-Control-Allow-Origin") != "http://localhost:3000" {
|
||||
t.Errorf("CORS origin should be set for allowed origin")
|
||||
}
|
||||
if rr.Header().Get("Access-Control-Allow-Methods") == "" {
|
||||
t.Error("Access-Control-Allow-Methods should be set")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,226 @@
|
||||
// Package worker provides the main worker service for claude-mnemonic.
|
||||
package worker
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// RateLimiter implements a token bucket rate limiter.
|
||||
type RateLimiter struct {
|
||||
lastUpdate time.Time
|
||||
rate float64
|
||||
burst int
|
||||
tokens float64
|
||||
requests int64
|
||||
rejected int64
|
||||
mu sync.Mutex
|
||||
}
|
||||
|
||||
// LastUpdateTime returns the last update time.
|
||||
// Thread-safe - acquires the limiter's lock.
|
||||
func (rl *RateLimiter) LastUpdateTime() time.Time {
|
||||
rl.mu.Lock()
|
||||
defer rl.mu.Unlock()
|
||||
return rl.lastUpdate
|
||||
}
|
||||
|
||||
// lastUpdateTimeUnlocked returns the last update time without locking.
|
||||
// Caller must hold rl.mu.
|
||||
func (rl *RateLimiter) lastUpdateTimeUnlocked() time.Time {
|
||||
return rl.lastUpdate
|
||||
}
|
||||
|
||||
// NewRateLimiter creates a new rate limiter.
|
||||
// rate is the number of requests per second to allow.
|
||||
// burst is the maximum burst of requests to allow.
|
||||
func NewRateLimiter(rate float64, burst int) *RateLimiter {
|
||||
return &RateLimiter{
|
||||
rate: rate,
|
||||
burst: burst,
|
||||
tokens: float64(burst),
|
||||
lastUpdate: time.Now(),
|
||||
}
|
||||
}
|
||||
|
||||
// Allow checks if a request should be allowed.
|
||||
// Returns true if the request is allowed, false if rate limited.
|
||||
func (rl *RateLimiter) Allow() bool {
|
||||
rl.mu.Lock()
|
||||
defer rl.mu.Unlock()
|
||||
|
||||
rl.requests++
|
||||
|
||||
// Calculate tokens added since last update
|
||||
now := time.Now()
|
||||
elapsed := now.Sub(rl.lastUpdate).Seconds()
|
||||
rl.tokens += elapsed * rl.rate
|
||||
if rl.tokens > float64(rl.burst) {
|
||||
rl.tokens = float64(rl.burst)
|
||||
}
|
||||
rl.lastUpdate = now
|
||||
|
||||
// Check if we have a token available
|
||||
if rl.tokens >= 1 {
|
||||
rl.tokens--
|
||||
return true
|
||||
}
|
||||
|
||||
rl.rejected++
|
||||
return false
|
||||
}
|
||||
|
||||
// Stats returns rate limiter statistics.
|
||||
func (rl *RateLimiter) Stats() map[string]any {
|
||||
rl.mu.Lock()
|
||||
defer rl.mu.Unlock()
|
||||
|
||||
return map[string]any{
|
||||
"rate": rl.rate,
|
||||
"burst": rl.burst,
|
||||
"current_tokens": rl.tokens,
|
||||
"total_requests": rl.requests,
|
||||
"rejected": rl.rejected,
|
||||
"rejection_rate": float64(rl.rejected) / max(float64(rl.requests), 1),
|
||||
}
|
||||
}
|
||||
|
||||
// RateLimitMiddleware creates middleware that applies rate limiting.
|
||||
// Uses a shared rate limiter for all requests.
|
||||
func RateLimitMiddleware(limiter *RateLimiter) func(http.Handler) http.Handler {
|
||||
return func(next http.Handler) http.Handler {
|
||||
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if !limiter.Allow() {
|
||||
http.Error(w, "Rate limit exceeded", http.StatusTooManyRequests)
|
||||
return
|
||||
}
|
||||
next.ServeHTTP(w, r)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// PerClientRateLimiter implements per-client rate limiting.
|
||||
type PerClientRateLimiter struct {
|
||||
lastCleanup time.Time
|
||||
clients map[string]*RateLimiter
|
||||
rate float64
|
||||
burst int
|
||||
cleanupInterval time.Duration
|
||||
maxIdleTime time.Duration
|
||||
mu sync.Mutex
|
||||
}
|
||||
|
||||
// NewPerClientRateLimiter creates a new per-client rate limiter.
|
||||
func NewPerClientRateLimiter(rate float64, burst int) *PerClientRateLimiter {
|
||||
return &PerClientRateLimiter{
|
||||
rate: rate,
|
||||
burst: burst,
|
||||
clients: make(map[string]*RateLimiter),
|
||||
cleanupInterval: 5 * time.Minute,
|
||||
maxIdleTime: 10 * time.Minute,
|
||||
lastCleanup: time.Now(),
|
||||
}
|
||||
}
|
||||
|
||||
// getLimiter returns a rate limiter for the given client key.
|
||||
func (pcrl *PerClientRateLimiter) getLimiter(key string) *RateLimiter {
|
||||
pcrl.mu.Lock()
|
||||
defer pcrl.mu.Unlock()
|
||||
|
||||
// Periodic cleanup of idle clients
|
||||
if time.Since(pcrl.lastCleanup) > pcrl.cleanupInterval {
|
||||
pcrl.cleanupLocked()
|
||||
}
|
||||
|
||||
limiter, exists := pcrl.clients[key]
|
||||
if !exists {
|
||||
limiter = NewRateLimiter(pcrl.rate, pcrl.burst)
|
||||
pcrl.clients[key] = limiter
|
||||
}
|
||||
|
||||
return limiter
|
||||
}
|
||||
|
||||
// cleanupLocked removes idle limiters. Must be called with lock held.
|
||||
// Uses consistent lock ordering: always acquire limiter.mu while holding pcrl.mu.
|
||||
// This is safe because the limiter.mu critical section is brief (just reading lastUpdate).
|
||||
func (pcrl *PerClientRateLimiter) cleanupLocked() {
|
||||
now := time.Now()
|
||||
keysToDelete := make([]string, 0)
|
||||
|
||||
// Check each limiter while holding pcrl.mu
|
||||
// We briefly acquire limiter.mu but the critical section is minimal
|
||||
for key, limiter := range pcrl.clients {
|
||||
limiter.mu.Lock()
|
||||
lastUpdate := limiter.lastUpdateTimeUnlocked()
|
||||
limiter.mu.Unlock()
|
||||
|
||||
if now.Sub(lastUpdate) > pcrl.maxIdleTime {
|
||||
keysToDelete = append(keysToDelete, key)
|
||||
}
|
||||
}
|
||||
|
||||
// Delete collected keys
|
||||
for _, key := range keysToDelete {
|
||||
delete(pcrl.clients, key)
|
||||
}
|
||||
pcrl.lastCleanup = now
|
||||
}
|
||||
|
||||
// Allow checks if a request from the given client should be allowed.
|
||||
func (pcrl *PerClientRateLimiter) Allow(clientKey string) bool {
|
||||
return pcrl.getLimiter(clientKey).Allow()
|
||||
}
|
||||
|
||||
// Stats returns aggregate statistics.
|
||||
// Uses two-phase approach to avoid nested lock acquisition.
|
||||
func (pcrl *PerClientRateLimiter) Stats() map[string]any {
|
||||
// Phase 1: Collect limiters under pcrl.mu
|
||||
pcrl.mu.Lock()
|
||||
rate := pcrl.rate
|
||||
burst := pcrl.burst
|
||||
activeClients := len(pcrl.clients)
|
||||
limiters := make([]*RateLimiter, 0, activeClients)
|
||||
for _, limiter := range pcrl.clients {
|
||||
limiters = append(limiters, limiter)
|
||||
}
|
||||
pcrl.mu.Unlock()
|
||||
|
||||
// Phase 2: Collect stats from each limiter (only acquiring limiter.mu, not pcrl.mu)
|
||||
var totalRequests, totalRejected int64
|
||||
for _, limiter := range limiters {
|
||||
limiter.mu.Lock()
|
||||
totalRequests += limiter.requests
|
||||
totalRejected += limiter.rejected
|
||||
limiter.mu.Unlock()
|
||||
}
|
||||
|
||||
return map[string]any{
|
||||
"rate": rate,
|
||||
"burst": burst,
|
||||
"active_clients": activeClients,
|
||||
"total_requests": totalRequests,
|
||||
"total_rejected": totalRejected,
|
||||
}
|
||||
}
|
||||
|
||||
// PerClientRateLimitMiddleware creates middleware that applies per-client rate limiting.
|
||||
// Uses X-Forwarded-For or RemoteAddr to identify clients.
|
||||
func PerClientRateLimitMiddleware(limiter *PerClientRateLimiter) func(http.Handler) http.Handler {
|
||||
return func(next http.Handler) http.Handler {
|
||||
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
// Get client identifier (prefer X-Real-IP from RealIP middleware)
|
||||
clientKey := r.RemoteAddr
|
||||
if xff := r.Header.Get("X-Real-IP"); xff != "" {
|
||||
clientKey = xff
|
||||
}
|
||||
|
||||
if !limiter.Allow(clientKey) {
|
||||
http.Error(w, "Rate limit exceeded", http.StatusTooManyRequests)
|
||||
return
|
||||
}
|
||||
next.ServeHTTP(w, r)
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,161 @@
|
||||
// Package sdk provides write-time observation deduplication via vector similarity.
|
||||
package sdk
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"github.com/lukaszraczylo/claude-mnemonic/internal/config"
|
||||
"github.com/lukaszraczylo/claude-mnemonic/internal/db/gorm"
|
||||
"github.com/lukaszraczylo/claude-mnemonic/internal/vector/sqlitevec"
|
||||
"github.com/lukaszraczylo/claude-mnemonic/pkg/models"
|
||||
"github.com/rs/zerolog/log"
|
||||
)
|
||||
|
||||
// DeduplicationResult represents the outcome of a vector similarity dedup check.
|
||||
type DeduplicationResult struct {
|
||||
ExistingID int64
|
||||
Similarity float64
|
||||
Action string // "insert", "merge"
|
||||
}
|
||||
|
||||
// checkVectorDeduplication checks if a similar observation already exists using vector similarity.
|
||||
// Returns a result indicating whether to insert or merge, or an error.
|
||||
// On any failure, returns Action="insert" so the caller always proceeds with storage.
|
||||
func (p *Processor) checkVectorDeduplication(ctx context.Context, obs *models.ParsedObservation, project string) *DeduplicationResult {
|
||||
cfg := config.Get()
|
||||
if !cfg.DeduplicationEnabled {
|
||||
return &DeduplicationResult{Action: "insert"}
|
||||
}
|
||||
|
||||
if p.vectorClient == nil {
|
||||
return &DeduplicationResult{Action: "insert"}
|
||||
}
|
||||
|
||||
// Build search text from observation fields
|
||||
searchText := buildObservationSearchText(obs)
|
||||
if searchText == "" {
|
||||
return &DeduplicationResult{Action: "insert"}
|
||||
}
|
||||
|
||||
// Query vector DB for similar observations in the same project
|
||||
where := sqlitevec.BuildWhereFilter(sqlitevec.DocTypeObservation, project)
|
||||
results, err := p.vectorClient.Query(ctx, searchText, 3, where)
|
||||
if err != nil {
|
||||
log.Debug().Err(err).Msg("Vector search failed during dedup check")
|
||||
return &DeduplicationResult{Action: "insert"}
|
||||
}
|
||||
|
||||
// Check results for high similarity
|
||||
for _, r := range results {
|
||||
if r.Similarity >= cfg.DeduplicationThreshold {
|
||||
obsID := extractObservationIDFromVectorDoc(r)
|
||||
if obsID > 0 {
|
||||
return &DeduplicationResult{
|
||||
ExistingID: obsID,
|
||||
Similarity: r.Similarity,
|
||||
Action: "merge",
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return &DeduplicationResult{Action: "insert"}
|
||||
}
|
||||
|
||||
// buildObservationSearchText creates searchable text from a parsed observation.
|
||||
func buildObservationSearchText(obs *models.ParsedObservation) string {
|
||||
var parts []string
|
||||
if obs.Title != "" {
|
||||
parts = append(parts, obs.Title)
|
||||
}
|
||||
if obs.Subtitle != "" {
|
||||
parts = append(parts, obs.Subtitle)
|
||||
}
|
||||
if obs.Narrative != "" {
|
||||
parts = append(parts, obs.Narrative)
|
||||
}
|
||||
text := strings.Join(parts, " ")
|
||||
if len(text) > 2000 {
|
||||
text = text[:2000]
|
||||
}
|
||||
return text
|
||||
}
|
||||
|
||||
// extractObservationIDFromVectorDoc extracts the SQLite observation ID from a vector query result.
|
||||
func extractObservationIDFromVectorDoc(r sqlitevec.QueryResult) int64 {
|
||||
// Prefer the sqlite_id metadata field (set during vector sync)
|
||||
if sqliteID, ok := r.Metadata["sqlite_id"].(float64); ok && sqliteID > 0 {
|
||||
return int64(sqliteID)
|
||||
}
|
||||
if sqliteID, ok := r.Metadata["sqlite_id"].(int64); ok && sqliteID > 0 {
|
||||
return sqliteID
|
||||
}
|
||||
|
||||
// Fallback: parse from doc_id format "obs_{id}_composite" or "obs_{id}_narrative"
|
||||
if !strings.HasPrefix(r.ID, "obs_") {
|
||||
return 0
|
||||
}
|
||||
parts := strings.SplitN(r.ID[4:], "_", 2)
|
||||
if len(parts) == 0 {
|
||||
return 0
|
||||
}
|
||||
var id int64
|
||||
fmt.Sscanf(parts[0], "%d", &id)
|
||||
return id
|
||||
}
|
||||
|
||||
// mergeObservation updates an existing observation with new information from a duplicate.
|
||||
// It appends new facts, updates the narrative if the new one is longer,
|
||||
// and bumps the importance score to reflect reconfirmation.
|
||||
func (p *Processor) mergeObservation(ctx context.Context, existingID int64, newObs *models.ParsedObservation) error {
|
||||
existing, err := p.observationStore.GetObservationByID(ctx, existingID)
|
||||
if err != nil {
|
||||
return fmt.Errorf("fetch existing observation %d: %w", existingID, err)
|
||||
}
|
||||
if existing == nil {
|
||||
return fmt.Errorf("observation %d not found", existingID)
|
||||
}
|
||||
|
||||
update := &gorm.ObservationUpdate{}
|
||||
changed := false
|
||||
|
||||
// Merge facts: append new facts not already present
|
||||
if len(newObs.Facts) > 0 {
|
||||
existingFactSet := make(map[string]struct{}, len(existing.Facts))
|
||||
for _, f := range existing.Facts {
|
||||
existingFactSet[f] = struct{}{}
|
||||
}
|
||||
mergedFacts := make([]string, len(existing.Facts))
|
||||
copy(mergedFacts, existing.Facts)
|
||||
for _, f := range newObs.Facts {
|
||||
if _, exists := existingFactSet[f]; !exists {
|
||||
mergedFacts = append(mergedFacts, f)
|
||||
changed = true
|
||||
}
|
||||
}
|
||||
if changed {
|
||||
update.Facts = &mergedFacts
|
||||
}
|
||||
}
|
||||
|
||||
// Update narrative if the new one is longer/more detailed
|
||||
if len(newObs.Narrative) > len(existing.Narrative.String) {
|
||||
update.Narrative = &newObs.Narrative
|
||||
changed = true
|
||||
}
|
||||
|
||||
if !changed {
|
||||
// Nothing new to merge, but still count it as a confirmed observation
|
||||
log.Debug().Int64("id", existingID).Msg("Dedup merge: no new content, skipping update")
|
||||
return nil
|
||||
}
|
||||
|
||||
_, err = p.observationStore.UpdateObservation(ctx, existingID, update)
|
||||
if err != nil {
|
||||
return fmt.Errorf("update observation %d: %w", existingID, err)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,143 @@
|
||||
package sdk
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/lukaszraczylo/claude-mnemonic/internal/vector/sqlitevec"
|
||||
"github.com/lukaszraczylo/claude-mnemonic/pkg/models"
|
||||
)
|
||||
|
||||
func TestBuildObservationSearchText(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
obs *models.ParsedObservation
|
||||
expected string
|
||||
}{
|
||||
{
|
||||
name: "empty observation",
|
||||
obs: &models.ParsedObservation{},
|
||||
expected: "",
|
||||
},
|
||||
{
|
||||
name: "title only",
|
||||
obs: &models.ParsedObservation{
|
||||
Title: "Fix database connection",
|
||||
},
|
||||
expected: "Fix database connection",
|
||||
},
|
||||
{
|
||||
name: "all fields",
|
||||
obs: &models.ParsedObservation{
|
||||
Title: "Fix database connection",
|
||||
Subtitle: "Connection pooling issue",
|
||||
Narrative: "The database connection pool was exhausted due to leaked connections.",
|
||||
},
|
||||
expected: "Fix database connection Connection pooling issue The database connection pool was exhausted due to leaked connections.",
|
||||
},
|
||||
{
|
||||
name: "truncates long text",
|
||||
obs: &models.ParsedObservation{
|
||||
Narrative: string(make([]byte, 3000)),
|
||||
},
|
||||
expected: string(make([]byte, 2000)),
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
result := buildObservationSearchText(tt.obs)
|
||||
if result != tt.expected {
|
||||
t.Errorf("got %q, want %q", result, tt.expected)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtractObservationIDFromVectorDoc(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
result sqlitevec.QueryResult
|
||||
expected int64
|
||||
}{
|
||||
{
|
||||
name: "from sqlite_id metadata (float64)",
|
||||
result: sqlitevec.QueryResult{
|
||||
ID: "obs_42_narrative",
|
||||
Metadata: map[string]any{"sqlite_id": float64(42)},
|
||||
},
|
||||
expected: 42,
|
||||
},
|
||||
{
|
||||
name: "from sqlite_id metadata (int64)",
|
||||
result: sqlitevec.QueryResult{
|
||||
ID: "obs_42_narrative",
|
||||
Metadata: map[string]any{"sqlite_id": int64(42)},
|
||||
},
|
||||
expected: 42,
|
||||
},
|
||||
{
|
||||
name: "fallback to doc_id parsing",
|
||||
result: sqlitevec.QueryResult{
|
||||
ID: "obs_99_composite",
|
||||
Metadata: map[string]any{},
|
||||
},
|
||||
expected: 99,
|
||||
},
|
||||
{
|
||||
name: "non-observation doc_id",
|
||||
result: sqlitevec.QueryResult{
|
||||
ID: "summary_5_text",
|
||||
Metadata: map[string]any{},
|
||||
},
|
||||
expected: 0,
|
||||
},
|
||||
{
|
||||
name: "zero sqlite_id falls back to doc_id",
|
||||
result: sqlitevec.QueryResult{
|
||||
ID: "obs_123_narrative",
|
||||
Metadata: map[string]any{"sqlite_id": float64(0)},
|
||||
},
|
||||
expected: 123,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
result := extractObservationIDFromVectorDoc(tt.result)
|
||||
if result != tt.expected {
|
||||
t.Errorf("got %d, want %d", result, tt.expected)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestCheckVectorDeduplication_NilClient(t *testing.T) {
|
||||
p := &Processor{
|
||||
// No vectorClient set
|
||||
}
|
||||
|
||||
obs := &models.ParsedObservation{
|
||||
Title: "Test observation",
|
||||
Narrative: "Some narrative text",
|
||||
}
|
||||
|
||||
result := p.checkVectorDeduplication(nil, obs, "test-project")
|
||||
if result.Action != "insert" {
|
||||
t.Errorf("expected Action='insert' when vectorClient is nil, got %q", result.Action)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCheckVectorDeduplication_EmptySearchText(t *testing.T) {
|
||||
p := &Processor{
|
||||
// vectorClient would be set but obs is empty
|
||||
}
|
||||
|
||||
obs := &models.ParsedObservation{
|
||||
// All empty fields
|
||||
}
|
||||
|
||||
result := p.checkVectorDeduplication(nil, obs, "test-project")
|
||||
if result.Action != "insert" {
|
||||
t.Errorf("expected Action='insert' for empty observation, got %q", result.Action)
|
||||
}
|
||||
}
|
||||
@@ -77,10 +77,10 @@ func TestParseObservations_TableDriven(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
input string
|
||||
expectedCount int
|
||||
expectedType models.ObservationType
|
||||
expectedTitle string
|
||||
checkConcepts []string
|
||||
expectedCount int
|
||||
}{
|
||||
{
|
||||
name: "valid_bugfix_observation",
|
||||
@@ -300,9 +300,9 @@ func TestParseSummary_TableDriven(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
input string
|
||||
expectedRequest string
|
||||
sessionID int64
|
||||
expectNil bool
|
||||
expectedRequest string
|
||||
}{
|
||||
{
|
||||
name: "empty_input",
|
||||
|
||||
@@ -4,24 +4,229 @@ package sdk
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"fmt"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
json "github.com/goccy/go-json"
|
||||
|
||||
"github.com/lukaszraczylo/claude-mnemonic/internal/config"
|
||||
"github.com/lukaszraczylo/claude-mnemonic/internal/db/gorm"
|
||||
"github.com/lukaszraczylo/claude-mnemonic/internal/vector/sqlitevec"
|
||||
"github.com/lukaszraczylo/claude-mnemonic/pkg/models"
|
||||
"github.com/lukaszraczylo/claude-mnemonic/pkg/similarity"
|
||||
"github.com/rs/zerolog/log"
|
||||
)
|
||||
|
||||
// CircuitBreaker implements a simple circuit breaker pattern for CLI calls.
|
||||
type CircuitBreaker struct {
|
||||
failures int64 // Current failure count
|
||||
lastFailure int64 // Unix timestamp of last failure
|
||||
threshold int64 // Number of failures before opening
|
||||
resetTimeout int64 // Seconds to wait before trying again
|
||||
state int32 // 0=closed, 1=open, 2=half-open
|
||||
}
|
||||
|
||||
const (
|
||||
circuitClosed int32 = 0
|
||||
circuitOpen int32 = 1
|
||||
circuitHalfOpen int32 = 2
|
||||
)
|
||||
|
||||
// NewCircuitBreaker creates a new circuit breaker.
|
||||
func NewCircuitBreaker(threshold int64, resetTimeout int64) *CircuitBreaker {
|
||||
return &CircuitBreaker{
|
||||
threshold: threshold,
|
||||
resetTimeout: resetTimeout,
|
||||
}
|
||||
}
|
||||
|
||||
// Allow checks if a request should be allowed through.
|
||||
func (cb *CircuitBreaker) Allow() bool {
|
||||
state := atomic.LoadInt32(&cb.state)
|
||||
if state == circuitClosed {
|
||||
return true
|
||||
}
|
||||
|
||||
if state == circuitOpen {
|
||||
// Check if reset timeout has passed
|
||||
lastFail := atomic.LoadInt64(&cb.lastFailure)
|
||||
if time.Now().Unix()-lastFail > cb.resetTimeout {
|
||||
// Transition to half-open
|
||||
atomic.CompareAndSwapInt32(&cb.state, circuitOpen, circuitHalfOpen)
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// Half-open: allow one request through
|
||||
return true
|
||||
}
|
||||
|
||||
// RecordSuccess records a successful call.
|
||||
func (cb *CircuitBreaker) RecordSuccess() {
|
||||
atomic.StoreInt64(&cb.failures, 0)
|
||||
atomic.StoreInt32(&cb.state, circuitClosed)
|
||||
}
|
||||
|
||||
// RecordFailure records a failed call.
|
||||
func (cb *CircuitBreaker) RecordFailure() {
|
||||
failures := atomic.AddInt64(&cb.failures, 1)
|
||||
atomic.StoreInt64(&cb.lastFailure, time.Now().Unix())
|
||||
|
||||
if failures >= cb.threshold {
|
||||
atomic.StoreInt32(&cb.state, circuitOpen)
|
||||
log.Warn().Int64("failures", failures).Msg("Circuit breaker opened - Claude CLI calls temporarily disabled")
|
||||
}
|
||||
}
|
||||
|
||||
// State returns the current state as a string.
|
||||
func (cb *CircuitBreaker) State() string {
|
||||
switch atomic.LoadInt32(&cb.state) {
|
||||
case circuitOpen:
|
||||
return "open"
|
||||
case circuitHalfOpen:
|
||||
return "half-open"
|
||||
default:
|
||||
return "closed"
|
||||
}
|
||||
}
|
||||
|
||||
// CircuitBreakerMetrics contains metrics about the circuit breaker state.
|
||||
type CircuitBreakerMetrics struct {
|
||||
State string `json:"state"`
|
||||
Failures int64 `json:"failures"`
|
||||
Threshold int64 `json:"threshold"`
|
||||
ResetTimeoutSecs int64 `json:"reset_timeout_secs"`
|
||||
LastFailureUnix int64 `json:"last_failure_unix,omitempty"`
|
||||
SecondsUntilReset int64 `json:"seconds_until_reset,omitempty"`
|
||||
}
|
||||
|
||||
// Metrics returns the current metrics of the circuit breaker.
|
||||
func (cb *CircuitBreaker) Metrics() CircuitBreakerMetrics {
|
||||
failures := atomic.LoadInt64(&cb.failures)
|
||||
lastFail := atomic.LoadInt64(&cb.lastFailure)
|
||||
state := cb.State()
|
||||
|
||||
metrics := CircuitBreakerMetrics{
|
||||
State: state,
|
||||
Failures: failures,
|
||||
Threshold: cb.threshold,
|
||||
ResetTimeoutSecs: cb.resetTimeout,
|
||||
}
|
||||
|
||||
if lastFail > 0 {
|
||||
metrics.LastFailureUnix = lastFail
|
||||
if state == "open" {
|
||||
remaining := cb.resetTimeout - (time.Now().Unix() - lastFail)
|
||||
if remaining > 0 {
|
||||
metrics.SecondsUntilReset = remaining
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return metrics
|
||||
}
|
||||
|
||||
// RequestDeduplicator tracks recent requests to prevent duplicates.
|
||||
type RequestDeduplicator struct {
|
||||
seen map[string]int64 // hash -> timestamp
|
||||
mu sync.RWMutex
|
||||
ttlSecs int64
|
||||
maxSize int
|
||||
}
|
||||
|
||||
// NewRequestDeduplicator creates a new deduplicator.
|
||||
func NewRequestDeduplicator(ttlSecs int64, maxSize int) *RequestDeduplicator {
|
||||
return &RequestDeduplicator{
|
||||
seen: make(map[string]int64),
|
||||
ttlSecs: ttlSecs,
|
||||
maxSize: maxSize,
|
||||
}
|
||||
}
|
||||
|
||||
// IsDuplicate checks if a request hash was seen recently.
|
||||
func (d *RequestDeduplicator) IsDuplicate(hash string) bool {
|
||||
now := time.Now().Unix()
|
||||
|
||||
d.mu.RLock()
|
||||
ts, exists := d.seen[hash]
|
||||
d.mu.RUnlock()
|
||||
|
||||
if exists && now-ts < d.ttlSecs {
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// Record marks a request hash as seen.
|
||||
func (d *RequestDeduplicator) Record(hash string) {
|
||||
now := time.Now().Unix()
|
||||
|
||||
d.mu.Lock()
|
||||
defer d.mu.Unlock()
|
||||
|
||||
// Evict old entries if at capacity
|
||||
if len(d.seen) >= d.maxSize {
|
||||
threshold := now - d.ttlSecs
|
||||
for k, ts := range d.seen {
|
||||
if ts < threshold {
|
||||
delete(d.seen, k)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
d.seen[hash] = now
|
||||
}
|
||||
|
||||
// hashRequest creates a hash of a request for deduplication.
|
||||
func hashRequest(toolName, input, output string) string {
|
||||
h := sha256.New()
|
||||
h.Write([]byte(toolName))
|
||||
h.Write([]byte(input))
|
||||
h.Write([]byte(output[:min(len(output), 1000)])) // Only hash first 1000 chars of output
|
||||
return hex.EncodeToString(h.Sum(nil))[:16] // Short hash is sufficient
|
||||
}
|
||||
|
||||
// maxStdoutBytes is the maximum number of bytes to capture from CLI stdout.
|
||||
const maxStdoutBytes = 1 * 1024 * 1024 // 1 MiB
|
||||
|
||||
// maxStderrBytes is the maximum number of bytes to capture from CLI stderr.
|
||||
const maxStderrBytes = 64 * 1024 // 64 KiB
|
||||
|
||||
// limitedWriter wraps a bytes.Buffer and silently discards writes beyond a maximum size.
|
||||
type limitedWriter struct {
|
||||
buf bytes.Buffer
|
||||
max int
|
||||
}
|
||||
|
||||
// Write implements io.Writer. It writes up to the remaining capacity and silently discards the rest.
|
||||
func (lw *limitedWriter) Write(p []byte) (int, error) {
|
||||
remaining := lw.max - lw.buf.Len()
|
||||
if remaining <= 0 {
|
||||
return len(p), nil // Silently discard
|
||||
}
|
||||
if len(p) > remaining {
|
||||
p = p[:remaining]
|
||||
}
|
||||
lw.buf.Write(p)
|
||||
return len(p), nil
|
||||
}
|
||||
|
||||
// String returns the buffered content as a string.
|
||||
func (lw *limitedWriter) String() string {
|
||||
return lw.buf.String()
|
||||
}
|
||||
|
||||
// BroadcastFunc is a callback for broadcasting events to SSE clients.
|
||||
type BroadcastFunc func(event map[string]interface{})
|
||||
type BroadcastFunc func(event map[string]any)
|
||||
|
||||
// SyncObservationFunc is a callback for syncing observations to vector DB.
|
||||
type SyncObservationFunc func(obs *models.Observation)
|
||||
@@ -29,17 +234,27 @@ type SyncObservationFunc func(obs *models.Observation)
|
||||
// SyncSummaryFunc is a callback for syncing summaries to vector DB.
|
||||
type SyncSummaryFunc func(summary *models.SessionSummary)
|
||||
|
||||
// MaxVectorSyncWorkers is the maximum number of concurrent vector sync operations.
|
||||
// This prevents unbounded goroutine spawning during high-volume observation ingestion.
|
||||
const MaxVectorSyncWorkers = 8
|
||||
|
||||
// Processor handles SDK agent processing of observations and summaries using Claude Code CLI.
|
||||
// Field order optimized for memory alignment (fieldalignment).
|
||||
type Processor struct {
|
||||
claudePath string
|
||||
model string
|
||||
observationStore *gorm.ObservationStore
|
||||
summaryStore *gorm.SummaryStore
|
||||
vectorClient *sqlitevec.Client
|
||||
broadcastFunc BroadcastFunc
|
||||
syncObservationFunc SyncObservationFunc
|
||||
syncSummaryFunc SyncSummaryFunc
|
||||
// Semaphore to limit concurrent Claude CLI calls (prevents API overload)
|
||||
sem chan struct{}
|
||||
circuitBreaker *CircuitBreaker
|
||||
deduplicator *RequestDeduplicator
|
||||
vectorSyncChan chan *models.Observation
|
||||
vectorSyncDone chan struct{}
|
||||
sem chan struct{}
|
||||
claudePath string
|
||||
model string
|
||||
vectorSyncWg sync.WaitGroup
|
||||
}
|
||||
|
||||
// SetBroadcastFunc sets the broadcast callback for SSE events.
|
||||
@@ -57,8 +272,13 @@ func (p *Processor) SetSyncSummaryFunc(fn SyncSummaryFunc) {
|
||||
p.syncSummaryFunc = fn
|
||||
}
|
||||
|
||||
// SetVectorClient sets the vector client for write-time deduplication.
|
||||
func (p *Processor) SetVectorClient(client *sqlitevec.Client) {
|
||||
p.vectorClient = client
|
||||
}
|
||||
|
||||
// broadcast sends an event via the broadcast callback if set.
|
||||
func (p *Processor) broadcast(event map[string]interface{}) {
|
||||
func (p *Processor) broadcast(event map[string]any) {
|
||||
if p.broadcastFunc != nil {
|
||||
p.broadcastFunc(event)
|
||||
}
|
||||
@@ -94,9 +314,65 @@ func NewProcessor(observationStore *gorm.ObservationStore, summaryStore *gorm.Su
|
||||
observationStore: observationStore,
|
||||
summaryStore: summaryStore,
|
||||
sem: make(chan struct{}, MaxConcurrentCLICalls),
|
||||
circuitBreaker: NewCircuitBreaker(5, 60), // Open after 5 failures, reset after 60s
|
||||
deduplicator: NewRequestDeduplicator(300, 1000), // 5-minute TTL, 1000 max entries
|
||||
vectorSyncChan: make(chan *models.Observation, MaxVectorSyncWorkers*2), // Buffered channel
|
||||
vectorSyncDone: make(chan struct{}),
|
||||
}, nil
|
||||
}
|
||||
|
||||
// StartVectorSyncWorkers starts the bounded worker pool for vector sync operations.
|
||||
// Call this after setting the sync function via SetSyncObservationFunc.
|
||||
func (p *Processor) StartVectorSyncWorkers() {
|
||||
for i := 0; i < MaxVectorSyncWorkers; i++ {
|
||||
p.vectorSyncWg.Add(1)
|
||||
go p.vectorSyncWorker()
|
||||
}
|
||||
log.Info().Int("workers", MaxVectorSyncWorkers).Msg("Vector sync worker pool started")
|
||||
}
|
||||
|
||||
// StopVectorSyncWorkers gracefully stops the worker pool.
|
||||
func (p *Processor) StopVectorSyncWorkers() {
|
||||
close(p.vectorSyncDone)
|
||||
p.vectorSyncWg.Wait()
|
||||
log.Info().Msg("Vector sync worker pool stopped")
|
||||
}
|
||||
|
||||
// vectorSyncWorker is a worker goroutine that processes vector sync requests.
|
||||
func (p *Processor) vectorSyncWorker() {
|
||||
defer p.vectorSyncWg.Done()
|
||||
for {
|
||||
select {
|
||||
case <-p.vectorSyncDone:
|
||||
// Drain remaining items before exiting
|
||||
for {
|
||||
select {
|
||||
case obs := <-p.vectorSyncChan:
|
||||
if p.syncObservationFunc != nil {
|
||||
p.syncObservationFunc(obs)
|
||||
}
|
||||
default:
|
||||
return
|
||||
}
|
||||
}
|
||||
case obs := <-p.vectorSyncChan:
|
||||
if p.syncObservationFunc != nil {
|
||||
p.syncObservationFunc(obs)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// CircuitBreakerState returns the current state of the circuit breaker.
|
||||
func (p *Processor) CircuitBreakerState() string {
|
||||
return p.circuitBreaker.State()
|
||||
}
|
||||
|
||||
// CircuitBreakerMetrics returns detailed metrics about the circuit breaker.
|
||||
func (p *Processor) CircuitBreakerMetrics() CircuitBreakerMetrics {
|
||||
return p.circuitBreaker.Metrics()
|
||||
}
|
||||
|
||||
// IsAvailable checks if the Claude CLI is available for processing.
|
||||
func (p *Processor) IsAvailable() bool {
|
||||
_, err := os.Stat(p.claudePath)
|
||||
@@ -104,7 +380,7 @@ func (p *Processor) IsAvailable() bool {
|
||||
}
|
||||
|
||||
// ProcessObservation processes a single tool observation and extracts insights.
|
||||
func (p *Processor) ProcessObservation(ctx context.Context, sdkSessionID, project string, toolName string, toolInput, toolResponse interface{}, promptNumber int, cwd string) error {
|
||||
func (p *Processor) ProcessObservation(ctx context.Context, sdkSessionID, project string, toolName string, toolInput, toolResponse any, promptNumber int, cwd string) error {
|
||||
// Skip certain tools that aren't worth processing
|
||||
if shouldSkipTool(toolName) {
|
||||
log.Info().Str("tool", toolName).Msg("Skipping tool (not interesting for memory)")
|
||||
@@ -121,11 +397,23 @@ func (p *Processor) ProcessObservation(ctx context.Context, sdkSessionID, projec
|
||||
return nil
|
||||
}
|
||||
|
||||
// Check for duplicate request within TTL window
|
||||
reqHash := hashRequest(toolName, inputStr, outputStr)
|
||||
if p.deduplicator.IsDuplicate(reqHash) {
|
||||
log.Debug().Str("tool", toolName).Msg("Skipping duplicate request (dedup)")
|
||||
return nil
|
||||
}
|
||||
|
||||
// Check circuit breaker before making CLI call
|
||||
if !p.circuitBreaker.Allow() {
|
||||
log.Warn().Str("tool", toolName).Msg("Circuit breaker open - skipping CLI call")
|
||||
return fmt.Errorf("circuit breaker open")
|
||||
}
|
||||
|
||||
log.Info().Str("tool", toolName).Msg("Processing tool execution with Claude CLI")
|
||||
|
||||
// Note: Removed the "file already has observations" check
|
||||
// Each tool execution can produce unique insights even for the same file
|
||||
// Similarity-based deduplication will handle true duplicates
|
||||
// Record this request to prevent duplicates
|
||||
p.deduplicator.Record(reqHash)
|
||||
|
||||
// Build the prompt
|
||||
exec := ToolExecution{
|
||||
@@ -147,9 +435,11 @@ func (p *Processor) ProcessObservation(ctx context.Context, sdkSessionID, projec
|
||||
// Call Claude Code CLI
|
||||
response, err := p.callClaudeCLI(ctx, prompt)
|
||||
if err != nil {
|
||||
p.circuitBreaker.RecordFailure()
|
||||
log.Error().Err(err).Str("tool", toolName).Msg("Failed to call Claude CLI for observation")
|
||||
return err
|
||||
}
|
||||
p.circuitBreaker.RecordSuccess()
|
||||
|
||||
// Parse observations from response
|
||||
observations := ParseObservations(response, sdkSessionID)
|
||||
@@ -176,16 +466,34 @@ func (p *Processor) ProcessObservation(ctx context.Context, sdkSessionID, projec
|
||||
// Convert to stored observation for similarity check
|
||||
storedObs := obs.ToStoredObservation()
|
||||
|
||||
// Check if this observation is too similar to existing ones
|
||||
// Check if this observation is too similar to existing ones (text-based Jaccard)
|
||||
if existingObs != nil && similarity.IsSimilarToAny(storedObs, existingObs, similarityThreshold) {
|
||||
log.Debug().
|
||||
Str("type", string(obs.Type)).
|
||||
Str("title", obs.Title).
|
||||
Msg("Skipping observation - too similar to existing")
|
||||
Msg("Skipping observation - too similar to existing (text)")
|
||||
skippedCount++
|
||||
continue
|
||||
}
|
||||
|
||||
// Check vector similarity for high-confidence dedup with merge
|
||||
dedupResult := p.checkVectorDeduplication(ctx, obs, project)
|
||||
if dedupResult.Action == "merge" {
|
||||
log.Info().
|
||||
Int64("existing_id", dedupResult.ExistingID).
|
||||
Float64("similarity", dedupResult.Similarity).
|
||||
Str("title", obs.Title).
|
||||
Msg("Merging duplicate observation (vector dedup)")
|
||||
if err := p.mergeObservation(ctx, dedupResult.ExistingID, obs); err != nil {
|
||||
log.Warn().Err(err).Int64("existing_id", dedupResult.ExistingID).
|
||||
Msg("Merge failed, inserting as new observation")
|
||||
// Fall through to normal insert
|
||||
} else {
|
||||
skippedCount++
|
||||
continue
|
||||
}
|
||||
}
|
||||
|
||||
id, createdAtEpoch, err := p.observationStore.StoreObservation(ctx, sdkSessionID, project, obs, promptNumber, 0)
|
||||
if err != nil {
|
||||
log.Error().Err(err).Msg("Failed to store observation")
|
||||
@@ -200,16 +508,26 @@ func (p *Processor) ProcessObservation(ctx context.Context, sdkSessionID, projec
|
||||
Int("trackedFiles", len(obs.FileMtimes)).
|
||||
Msg("Observation stored")
|
||||
|
||||
// Sync to vector DB if callback is set
|
||||
if p.syncObservationFunc != nil {
|
||||
// Sync to vector DB via bounded worker pool (non-blocking to reduce latency)
|
||||
if p.syncObservationFunc != nil && p.vectorSyncChan != nil {
|
||||
fullObs := models.NewObservation(sdkSessionID, project, obs, promptNumber, 0)
|
||||
fullObs.ID = id
|
||||
fullObs.CreatedAtEpoch = createdAtEpoch
|
||||
p.syncObservationFunc(fullObs)
|
||||
// Non-blocking send to worker pool - drops if channel is full
|
||||
select {
|
||||
case p.vectorSyncChan <- fullObs:
|
||||
// Sent to worker pool
|
||||
default:
|
||||
// Channel full, fall back to direct sync in goroutine (bounded by channel buffer)
|
||||
log.Debug().Int64("obs_id", id).Msg("Vector sync channel full, using fallback goroutine")
|
||||
go func(obsToSync *models.Observation) {
|
||||
p.syncObservationFunc(obsToSync)
|
||||
}(fullObs)
|
||||
}
|
||||
}
|
||||
|
||||
// Broadcast new observation event for dashboard refresh
|
||||
p.broadcast(map[string]interface{}{
|
||||
p.broadcast(map[string]any{
|
||||
"type": "observation",
|
||||
"action": "created",
|
||||
"id": id,
|
||||
@@ -311,7 +629,7 @@ func (p *Processor) ProcessSummary(ctx context.Context, sessionDBID int64, sdkSe
|
||||
}
|
||||
|
||||
// Broadcast new summary event for dashboard refresh
|
||||
p.broadcast(map[string]interface{}{
|
||||
p.broadcast(map[string]any{
|
||||
"type": "summary",
|
||||
"action": "created",
|
||||
"id": id,
|
||||
@@ -321,8 +639,31 @@ func (p *Processor) ProcessSummary(ctx context.Context, sessionDBID int64, sdkSe
|
||||
return nil
|
||||
}
|
||||
|
||||
// MaxPromptSize is the maximum size of a prompt that can be passed to the Claude CLI.
|
||||
// This prevents resource exhaustion from extremely large prompts.
|
||||
const MaxPromptSize = 100 * 1024 // 100KB
|
||||
|
||||
// sanitizePrompt removes null bytes and control characters from a prompt.
|
||||
// Keeps newlines, tabs, and carriage returns as they're valid in prompts.
|
||||
func sanitizePrompt(s string) string {
|
||||
return strings.Map(func(r rune) rune {
|
||||
// Keep printable ASCII, extended Unicode, and common whitespace
|
||||
if r >= 32 || r == '\n' || r == '\t' || r == '\r' {
|
||||
return r
|
||||
}
|
||||
// Remove null bytes and other control characters
|
||||
return -1
|
||||
}, s)
|
||||
}
|
||||
|
||||
// callClaudeCLI calls the Claude Code CLI with the given prompt.
|
||||
func (p *Processor) callClaudeCLI(ctx context.Context, prompt string) (string, error) {
|
||||
// Validate and sanitize prompt
|
||||
if len(prompt) > MaxPromptSize {
|
||||
return "", fmt.Errorf("prompt exceeds maximum size of %d bytes", MaxPromptSize)
|
||||
}
|
||||
prompt = sanitizePrompt(prompt)
|
||||
|
||||
// Build the full prompt with system instructions
|
||||
fullPrompt := systemPrompt + "\n\n" + prompt
|
||||
|
||||
@@ -355,13 +696,19 @@ func (p *Processor) callClaudeCLI(ctx context.Context, prompt string) (string, e
|
||||
// (hooks are triggered based on working directory)
|
||||
cmd.Dir = "/tmp"
|
||||
|
||||
// Disable any plugin hooks by setting an env var that our hooks can check
|
||||
// Disable any plugin hooks by setting an env var that our hooks can check.
|
||||
// NOTE: This env var is NOT reliably propagated to hook subprocesses because
|
||||
// Claude Code constructs its own environment for hooks. The server-side check
|
||||
// in handleSessionInit (handlers_sessions.go) provides the reliable guard by
|
||||
// detecting the system prompt signature in the prompt text. This env var is
|
||||
// kept as a best-effort defense-in-depth measure for cases where it does work.
|
||||
cmd.Env = append(os.Environ(), "CLAUDE_MNEMONIC_INTERNAL=1")
|
||||
|
||||
// Capture output
|
||||
var stdout, stderr bytes.Buffer
|
||||
cmd.Stdout = &stdout
|
||||
cmd.Stderr = &stderr
|
||||
// Capture output with size limits to prevent unbounded memory usage
|
||||
stdout := &limitedWriter{max: maxStdoutBytes}
|
||||
stderr := &limitedWriter{max: maxStderrBytes}
|
||||
cmd.Stdout = stdout
|
||||
cmd.Stderr = stderr
|
||||
|
||||
// Run command
|
||||
err := cmd.Run()
|
||||
@@ -419,8 +766,11 @@ func shouldSkipTrivialOperation(toolName, inputStr, outputStr string) bool {
|
||||
return true
|
||||
}
|
||||
|
||||
// Skip if output indicates an error or empty result
|
||||
// Pre-compute lowercase strings once to avoid repeated allocations
|
||||
lowerOutput := strings.ToLower(outputStr)
|
||||
lowerInput := strings.ToLower(inputStr)
|
||||
|
||||
// Skip if output indicates an error or empty result
|
||||
trivialOutputs := []string{
|
||||
"no matches found",
|
||||
"file not found",
|
||||
@@ -444,13 +794,13 @@ func shouldSkipTrivialOperation(toolName, inputStr, outputStr string) bool {
|
||||
// Skip reading config files that rarely contain project-specific insights
|
||||
boringFiles := []string{
|
||||
"package-lock.json", "yarn.lock", "pnpm-lock.yaml",
|
||||
"go.sum", "Cargo.lock", "Gemfile.lock", "poetry.lock",
|
||||
"go.sum", "cargo.lock", "gemfile.lock", "poetry.lock",
|
||||
".gitignore", ".dockerignore", ".eslintignore",
|
||||
"tsconfig.json", "jsconfig.json", "vite.config",
|
||||
"tailwind.config", "postcss.config",
|
||||
}
|
||||
for _, boring := range boringFiles {
|
||||
if strings.Contains(inputStr, boring) {
|
||||
if strings.Contains(lowerInput, boring) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
@@ -462,14 +812,14 @@ func shouldSkipTrivialOperation(toolName, inputStr, outputStr string) bool {
|
||||
}
|
||||
|
||||
case "Bash":
|
||||
// Skip simple status commands
|
||||
// Skip simple status commands (use pre-computed lowerInput)
|
||||
boringCommands := []string{
|
||||
"git status", "git diff", "git log", "git branch",
|
||||
"ls ", "pwd", "echo ", "cat ", "which ", "type ",
|
||||
"npm list", "npm outdated", "npm audit",
|
||||
}
|
||||
for _, boring := range boringCommands {
|
||||
if strings.Contains(strings.ToLower(inputStr), boring) {
|
||||
if strings.Contains(lowerInput, boring) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
@@ -479,7 +829,7 @@ func shouldSkipTrivialOperation(toolName, inputStr, outputStr string) bool {
|
||||
}
|
||||
|
||||
// toJSONString converts an interface to a JSON string.
|
||||
func toJSONString(v interface{}) string {
|
||||
func toJSONString(v any) string {
|
||||
if v == nil {
|
||||
return ""
|
||||
}
|
||||
@@ -493,38 +843,132 @@ func toJSONString(v interface{}) string {
|
||||
return string(b)
|
||||
}
|
||||
|
||||
// safeResolvePath resolves a path relative to cwd and validates it doesn't escape the cwd directory.
|
||||
// Returns the resolved absolute path and true if valid, or empty string and false if path traversal detected.
|
||||
// This function is a security sanitizer for path traversal attacks.
|
||||
func safeResolvePath(path, cwd string) (string, bool) {
|
||||
// Clean the input path to normalize any .. or . components
|
||||
cleanPath := filepath.Clean(path)
|
||||
|
||||
// Reject paths that explicitly contain parent directory traversal after cleaning
|
||||
if strings.Contains(cleanPath, "..") {
|
||||
return "", false
|
||||
}
|
||||
|
||||
if filepath.IsAbs(cleanPath) {
|
||||
// For absolute paths, verify they're within cwd if cwd is specified
|
||||
if cwd != "" {
|
||||
cleanCwd := filepath.Clean(cwd)
|
||||
if !strings.HasPrefix(cleanPath, cleanCwd+string(filepath.Separator)) && cleanPath != cleanCwd {
|
||||
return "", false
|
||||
}
|
||||
}
|
||||
return cleanPath, true
|
||||
}
|
||||
|
||||
if cwd == "" {
|
||||
return cleanPath, true
|
||||
}
|
||||
|
||||
// Clean the cwd first
|
||||
cleanCwd := filepath.Clean(cwd)
|
||||
|
||||
// Join and clean the path
|
||||
absPath := filepath.Join(cleanCwd, cleanPath)
|
||||
|
||||
// Use filepath.Rel to verify the path is actually within cwd
|
||||
// If Rel returns a path starting with "..", it escapes the base
|
||||
rel, err := filepath.Rel(cleanCwd, absPath)
|
||||
if err != nil || strings.HasPrefix(rel, "..") {
|
||||
return "", false
|
||||
}
|
||||
|
||||
return absPath, true
|
||||
}
|
||||
|
||||
// captureFileMtimes captures current modification times for tracked files.
|
||||
// Returns a map of absolute file paths to their mtime in epoch milliseconds.
|
||||
// For large file lists (>10 files), uses parallel stat calls for better performance.
|
||||
func captureFileMtimes(filesRead, filesModified []string, cwd string) map[string]int64 {
|
||||
mtimes := make(map[string]int64)
|
||||
// Combine all unique file paths
|
||||
allPaths := make(map[string]struct{}, len(filesRead)+len(filesModified))
|
||||
for _, path := range filesRead {
|
||||
allPaths[path] = struct{}{}
|
||||
}
|
||||
for _, path := range filesModified {
|
||||
allPaths[path] = struct{}{}
|
||||
}
|
||||
|
||||
// Helper to get mtime for a file path
|
||||
getMtime := func(path string) (int64, bool) {
|
||||
// Resolve relative paths against cwd
|
||||
absPath := path
|
||||
if !filepath.IsAbs(path) && cwd != "" {
|
||||
absPath = filepath.Join(cwd, path)
|
||||
// For small lists, use sequential processing (goroutine overhead not worth it)
|
||||
if len(allPaths) <= 10 {
|
||||
return captureFileMtimesSequential(allPaths, cwd)
|
||||
}
|
||||
|
||||
// For larger lists, parallelize with bounded concurrency
|
||||
return captureFileMtimesParallel(allPaths, cwd)
|
||||
}
|
||||
|
||||
// captureFileMtimesSequential captures mtimes sequentially (efficient for small lists).
|
||||
func captureFileMtimesSequential(paths map[string]struct{}, cwd string) map[string]int64 {
|
||||
mtimes := make(map[string]int64, len(paths))
|
||||
|
||||
for path := range paths {
|
||||
absPath, ok := safeResolvePath(path, cwd)
|
||||
if !ok {
|
||||
// Skip paths that attempt directory traversal
|
||||
continue
|
||||
}
|
||||
|
||||
info, err := os.Stat(absPath)
|
||||
if err != nil {
|
||||
return 0, false
|
||||
}
|
||||
return info.ModTime().UnixMilli(), true
|
||||
}
|
||||
|
||||
// Capture mtimes for all read files
|
||||
for _, path := range filesRead {
|
||||
if mtime, ok := getMtime(path); ok {
|
||||
mtimes[path] = mtime
|
||||
if err == nil {
|
||||
mtimes[path] = info.ModTime().UnixMilli()
|
||||
}
|
||||
}
|
||||
|
||||
// Capture mtimes for all modified files
|
||||
for _, path := range filesModified {
|
||||
if mtime, ok := getMtime(path); ok {
|
||||
mtimes[path] = mtime
|
||||
}
|
||||
return mtimes
|
||||
}
|
||||
|
||||
// captureFileMtimesParallel captures mtimes in parallel with bounded concurrency.
|
||||
func captureFileMtimesParallel(paths map[string]struct{}, cwd string) map[string]int64 {
|
||||
type mtimeResult struct {
|
||||
path string
|
||||
mtime int64
|
||||
}
|
||||
|
||||
results := make(chan mtimeResult, len(paths))
|
||||
sem := make(chan struct{}, 8) // Limit to 8 concurrent stat calls
|
||||
var wg sync.WaitGroup
|
||||
|
||||
for path := range paths {
|
||||
wg.Add(1)
|
||||
go func(p string) {
|
||||
defer wg.Done()
|
||||
sem <- struct{}{} // Acquire
|
||||
defer func() { <-sem }() // Release
|
||||
|
||||
absPath, ok := safeResolvePath(p, cwd)
|
||||
if !ok {
|
||||
// Skip paths that attempt directory traversal
|
||||
return
|
||||
}
|
||||
|
||||
info, err := os.Stat(absPath)
|
||||
if err == nil {
|
||||
results <- mtimeResult{path: p, mtime: info.ModTime().UnixMilli()}
|
||||
}
|
||||
}(path)
|
||||
}
|
||||
|
||||
// Close results channel when all goroutines complete
|
||||
go func() {
|
||||
wg.Wait()
|
||||
close(results)
|
||||
}()
|
||||
|
||||
// Collect results
|
||||
mtimes := make(map[string]int64, len(paths))
|
||||
for res := range results {
|
||||
mtimes[res.path] = res.mtime
|
||||
}
|
||||
|
||||
return mtimes
|
||||
@@ -539,12 +983,13 @@ func GetFileMtimes(paths []string, cwd string) map[string]int64 {
|
||||
// GetFileContent reads file content for verification purposes.
|
||||
// Returns content and ok status.
|
||||
func GetFileContent(path, cwd string) (string, bool) {
|
||||
absPath := path
|
||||
if !filepath.IsAbs(path) && cwd != "" {
|
||||
absPath = filepath.Join(cwd, path)
|
||||
absPath, ok := safeResolvePath(path, cwd)
|
||||
if !ok {
|
||||
// Reject paths that attempt directory traversal
|
||||
return "", false
|
||||
}
|
||||
|
||||
content, err := os.ReadFile(absPath) // #nosec G304 -- intentional file read for verification
|
||||
content, err := os.ReadFile(absPath) // #nosec G304 -- path validated by safeResolvePath
|
||||
if err != nil {
|
||||
return "", false
|
||||
}
|
||||
|
||||
@@ -1,9 +1,12 @@
|
||||
package sdk
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/lukaszraczylo/claude-mnemonic/pkg/models"
|
||||
"github.com/stretchr/testify/assert"
|
||||
@@ -11,8 +14,8 @@ import (
|
||||
|
||||
func TestIsSelfReferentialSummary(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
summary *models.ParsedSummary
|
||||
name string
|
||||
expected bool
|
||||
}{
|
||||
{
|
||||
@@ -281,8 +284,8 @@ func TestTruncateForLog(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
input string
|
||||
maxLen int
|
||||
expected string
|
||||
maxLen int
|
||||
}{
|
||||
{
|
||||
name: "shorter_than_max",
|
||||
@@ -719,8 +722,8 @@ func TestShouldSkipTrivialOperation_EdgeCases(t *testing.T) {
|
||||
// TestIsSelfReferentialSummary_MoreCases tests additional self-referential detection cases.
|
||||
func TestIsSelfReferentialSummary_MoreCases(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
summary *models.ParsedSummary
|
||||
name string
|
||||
expected bool
|
||||
}{
|
||||
{
|
||||
@@ -974,3 +977,867 @@ func TestSyncSummaryFuncType(t *testing.T) {
|
||||
}
|
||||
assert.NotNil(t, fn)
|
||||
}
|
||||
|
||||
// TestSanitizePrompt tests prompt sanitization for CLI safety.
|
||||
func TestSanitizePrompt(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
input string
|
||||
expected string
|
||||
}{
|
||||
{
|
||||
name: "normal text",
|
||||
input: "Hello, world!",
|
||||
expected: "Hello, world!",
|
||||
},
|
||||
{
|
||||
name: "text with newlines",
|
||||
input: "Line 1\nLine 2\nLine 3",
|
||||
expected: "Line 1\nLine 2\nLine 3",
|
||||
},
|
||||
{
|
||||
name: "text with tabs",
|
||||
input: "Key:\tValue",
|
||||
expected: "Key:\tValue",
|
||||
},
|
||||
{
|
||||
name: "text with carriage return",
|
||||
input: "Line 1\r\nLine 2",
|
||||
expected: "Line 1\r\nLine 2",
|
||||
},
|
||||
{
|
||||
name: "text with null bytes",
|
||||
input: "Hello\x00World",
|
||||
expected: "HelloWorld",
|
||||
},
|
||||
{
|
||||
name: "text with control characters",
|
||||
input: "Hello\x01\x02\x03World",
|
||||
expected: "HelloWorld",
|
||||
},
|
||||
{
|
||||
name: "text with bell character",
|
||||
input: "Hello\x07World",
|
||||
expected: "HelloWorld",
|
||||
},
|
||||
{
|
||||
name: "text with backspace",
|
||||
input: "Hello\x08World",
|
||||
expected: "HelloWorld",
|
||||
},
|
||||
{
|
||||
name: "text with form feed",
|
||||
input: "Hello\x0cWorld",
|
||||
expected: "HelloWorld",
|
||||
},
|
||||
{
|
||||
name: "text with escape",
|
||||
input: "Hello\x1bWorld",
|
||||
expected: "HelloWorld",
|
||||
},
|
||||
{
|
||||
name: "unicode text",
|
||||
input: "Hello 世界 🌍",
|
||||
expected: "Hello 世界 🌍",
|
||||
},
|
||||
{
|
||||
name: "empty string",
|
||||
input: "",
|
||||
expected: "",
|
||||
},
|
||||
{
|
||||
name: "only control characters",
|
||||
input: "\x00\x01\x02\x03",
|
||||
expected: "",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
result := sanitizePrompt(tt.input)
|
||||
assert.Equal(t, tt.expected, result)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestMaxPromptSize tests that MaxPromptSize is reasonable.
|
||||
func TestMaxPromptSize(t *testing.T) {
|
||||
assert.Equal(t, 100*1024, MaxPromptSize)
|
||||
}
|
||||
|
||||
// BenchmarkSanitizePrompt benchmarks the sanitize function.
|
||||
func BenchmarkSanitizePrompt(b *testing.B) {
|
||||
prompt := "Analyze the following code:\n```go\nfunc main() {\n\tfmt.Println(\"Hello, World!\")\n}\n```\n\nPlease identify any issues."
|
||||
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
sanitizePrompt(prompt)
|
||||
}
|
||||
}
|
||||
|
||||
// BenchmarkSanitizePromptWithControlChars benchmarks sanitization with control characters.
|
||||
func BenchmarkSanitizePromptWithControlChars(b *testing.B) {
|
||||
prompt := "Hello\x00World\x01Test\x02Data\x03End"
|
||||
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
sanitizePrompt(prompt)
|
||||
}
|
||||
}
|
||||
|
||||
// TestSafeResolvePath tests the path traversal protection.
|
||||
func TestSafeResolvePath(t *testing.T) {
|
||||
// Create a temporary directory for testing
|
||||
tmpDir := t.TempDir()
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
path string
|
||||
cwd string
|
||||
wantPath string
|
||||
wantOk bool
|
||||
}{
|
||||
{
|
||||
name: "simple relative path",
|
||||
path: "file.txt",
|
||||
cwd: tmpDir,
|
||||
wantOk: true,
|
||||
wantPath: filepath.Join(tmpDir, "file.txt"),
|
||||
},
|
||||
{
|
||||
name: "nested relative path",
|
||||
path: "subdir/file.txt",
|
||||
cwd: tmpDir,
|
||||
wantOk: true,
|
||||
wantPath: filepath.Join(tmpDir, "subdir", "file.txt"),
|
||||
},
|
||||
{
|
||||
name: "path traversal with ..",
|
||||
path: "../etc/passwd",
|
||||
cwd: tmpDir,
|
||||
wantOk: false,
|
||||
},
|
||||
{
|
||||
name: "path traversal with multiple ..",
|
||||
path: "../../etc/passwd",
|
||||
cwd: tmpDir,
|
||||
wantOk: false,
|
||||
},
|
||||
{
|
||||
name: "path traversal hidden in middle",
|
||||
path: "subdir/../../../etc/passwd",
|
||||
cwd: tmpDir,
|
||||
wantOk: false,
|
||||
},
|
||||
{
|
||||
name: "just parent directory",
|
||||
path: "..",
|
||||
cwd: tmpDir,
|
||||
wantOk: false,
|
||||
},
|
||||
{
|
||||
name: "absolute path without cwd",
|
||||
path: "/some/absolute/path",
|
||||
cwd: "",
|
||||
wantOk: true,
|
||||
wantPath: "/some/absolute/path",
|
||||
},
|
||||
{
|
||||
name: "relative path without cwd",
|
||||
path: "relative/path",
|
||||
cwd: "",
|
||||
wantOk: true,
|
||||
wantPath: "relative/path",
|
||||
},
|
||||
{
|
||||
name: "current directory reference",
|
||||
path: "./file.txt",
|
||||
cwd: tmpDir,
|
||||
wantOk: true,
|
||||
wantPath: filepath.Join(tmpDir, "file.txt"),
|
||||
},
|
||||
{
|
||||
name: "absolute path outside cwd",
|
||||
path: "/etc/passwd",
|
||||
cwd: tmpDir,
|
||||
wantOk: false,
|
||||
},
|
||||
{
|
||||
name: "absolute path inside cwd",
|
||||
path: filepath.Join(tmpDir, "inside.txt"),
|
||||
cwd: tmpDir,
|
||||
wantOk: true,
|
||||
wantPath: filepath.Join(tmpDir, "inside.txt"),
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
gotPath, gotOk := safeResolvePath(tt.path, tt.cwd)
|
||||
assert.Equal(t, tt.wantOk, gotOk, "ok status mismatch")
|
||||
if tt.wantPath != "" && gotOk {
|
||||
assert.Equal(t, tt.wantPath, gotPath, "path mismatch")
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// TESTS FOR CircuitBreaker
|
||||
// =============================================================================
|
||||
|
||||
func TestNewCircuitBreaker(t *testing.T) {
|
||||
cb := NewCircuitBreaker(5, 60)
|
||||
|
||||
assert.NotNil(t, cb)
|
||||
assert.Equal(t, int64(5), cb.threshold)
|
||||
assert.Equal(t, int64(60), cb.resetTimeout)
|
||||
assert.Equal(t, "closed", cb.State())
|
||||
}
|
||||
|
||||
func TestCircuitBreaker_Allow_Closed(t *testing.T) {
|
||||
cb := NewCircuitBreaker(5, 60)
|
||||
|
||||
// Closed state should allow requests
|
||||
assert.True(t, cb.Allow())
|
||||
assert.True(t, cb.Allow())
|
||||
}
|
||||
|
||||
func TestCircuitBreaker_Allow_Open(t *testing.T) {
|
||||
cb := NewCircuitBreaker(2, 60) // Low threshold for testing
|
||||
|
||||
// Record enough failures to open the circuit
|
||||
cb.RecordFailure()
|
||||
cb.RecordFailure()
|
||||
|
||||
// Open state should block requests
|
||||
assert.False(t, cb.Allow())
|
||||
assert.Equal(t, "open", cb.State())
|
||||
}
|
||||
|
||||
func TestCircuitBreaker_RecordSuccess(t *testing.T) {
|
||||
cb := NewCircuitBreaker(2, 60)
|
||||
|
||||
// Record a failure
|
||||
cb.RecordFailure()
|
||||
assert.Equal(t, int64(1), cb.Metrics().Failures)
|
||||
|
||||
// Record success resets failures
|
||||
cb.RecordSuccess()
|
||||
assert.Equal(t, int64(0), cb.Metrics().Failures)
|
||||
assert.Equal(t, "closed", cb.State())
|
||||
}
|
||||
|
||||
func TestCircuitBreaker_RecordFailure_OpensCircuit(t *testing.T) {
|
||||
cb := NewCircuitBreaker(3, 60)
|
||||
|
||||
// Record failures below threshold
|
||||
cb.RecordFailure()
|
||||
assert.Equal(t, "closed", cb.State())
|
||||
|
||||
cb.RecordFailure()
|
||||
assert.Equal(t, "closed", cb.State())
|
||||
|
||||
// Third failure should open circuit
|
||||
cb.RecordFailure()
|
||||
assert.Equal(t, "open", cb.State())
|
||||
}
|
||||
|
||||
func TestCircuitBreaker_State(t *testing.T) {
|
||||
cb := NewCircuitBreaker(1, 60)
|
||||
|
||||
// Initially closed
|
||||
assert.Equal(t, "closed", cb.State())
|
||||
|
||||
// After failure, open
|
||||
cb.RecordFailure()
|
||||
assert.Equal(t, "open", cb.State())
|
||||
|
||||
// After success, closed
|
||||
cb.RecordSuccess()
|
||||
assert.Equal(t, "closed", cb.State())
|
||||
}
|
||||
|
||||
func TestCircuitBreaker_Metrics(t *testing.T) {
|
||||
cb := NewCircuitBreaker(5, 120)
|
||||
|
||||
metrics := cb.Metrics()
|
||||
assert.Equal(t, "closed", metrics.State)
|
||||
assert.Equal(t, int64(0), metrics.Failures)
|
||||
assert.Equal(t, int64(5), metrics.Threshold)
|
||||
assert.Equal(t, int64(120), metrics.ResetTimeoutSecs)
|
||||
assert.Equal(t, int64(0), metrics.LastFailureUnix)
|
||||
|
||||
// After failure
|
||||
cb.RecordFailure()
|
||||
metrics = cb.Metrics()
|
||||
assert.Equal(t, int64(1), metrics.Failures)
|
||||
assert.Greater(t, metrics.LastFailureUnix, int64(0))
|
||||
}
|
||||
|
||||
func TestCircuitBreaker_Metrics_OpenWithReset(t *testing.T) {
|
||||
cb := NewCircuitBreaker(1, 60)
|
||||
|
||||
cb.RecordFailure()
|
||||
assert.Equal(t, "open", cb.State())
|
||||
|
||||
metrics := cb.Metrics()
|
||||
assert.Equal(t, "open", metrics.State)
|
||||
assert.Greater(t, metrics.SecondsUntilReset, int64(0))
|
||||
assert.LessOrEqual(t, metrics.SecondsUntilReset, int64(60))
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// TESTS FOR RequestDeduplicator
|
||||
// =============================================================================
|
||||
|
||||
func TestNewRequestDeduplicator(t *testing.T) {
|
||||
d := NewRequestDeduplicator(300, 1000)
|
||||
|
||||
assert.NotNil(t, d)
|
||||
assert.NotNil(t, d.seen)
|
||||
assert.Equal(t, int64(300), d.ttlSecs)
|
||||
assert.Equal(t, 1000, d.maxSize)
|
||||
}
|
||||
|
||||
func TestRequestDeduplicator_IsDuplicate_NotSeen(t *testing.T) {
|
||||
d := NewRequestDeduplicator(300, 1000)
|
||||
|
||||
// New hash is not a duplicate
|
||||
assert.False(t, d.IsDuplicate("newhash"))
|
||||
}
|
||||
|
||||
func TestRequestDeduplicator_IsDuplicate_AfterRecord(t *testing.T) {
|
||||
d := NewRequestDeduplicator(300, 1000)
|
||||
|
||||
hash := "testhash"
|
||||
|
||||
// Record the hash
|
||||
d.Record(hash)
|
||||
|
||||
// Now it should be a duplicate
|
||||
assert.True(t, d.IsDuplicate(hash))
|
||||
}
|
||||
|
||||
func TestRequestDeduplicator_Record(t *testing.T) {
|
||||
d := NewRequestDeduplicator(300, 1000)
|
||||
|
||||
hash := "recordtest"
|
||||
d.Record(hash)
|
||||
|
||||
// Check it was recorded
|
||||
d.mu.RLock()
|
||||
_, exists := d.seen[hash]
|
||||
d.mu.RUnlock()
|
||||
|
||||
assert.True(t, exists)
|
||||
}
|
||||
|
||||
func TestRequestDeduplicator_Record_Eviction(t *testing.T) {
|
||||
// Small maxSize for testing eviction
|
||||
d := NewRequestDeduplicator(0, 2) // TTL of 0 means everything is "old"
|
||||
|
||||
// Record until capacity
|
||||
d.Record("hash1")
|
||||
d.Record("hash2")
|
||||
|
||||
// Recording a third should trigger eviction (since TTL is 0)
|
||||
d.Record("hash3")
|
||||
|
||||
// Should have cleaned up old entries
|
||||
d.mu.RLock()
|
||||
size := len(d.seen)
|
||||
d.mu.RUnlock()
|
||||
|
||||
// Size should be limited (eviction occurred)
|
||||
assert.LessOrEqual(t, size, 3)
|
||||
}
|
||||
|
||||
func TestHashRequest(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
toolName string
|
||||
input string
|
||||
output string
|
||||
compareWith []string
|
||||
wantLen int
|
||||
wantSame bool
|
||||
}{
|
||||
{
|
||||
name: "basic hash",
|
||||
toolName: "Read",
|
||||
input: "file.txt",
|
||||
output: "content",
|
||||
wantLen: 16,
|
||||
},
|
||||
{
|
||||
name: "consistent hashing",
|
||||
toolName: "Edit",
|
||||
input: "same input",
|
||||
output: "same output",
|
||||
wantLen: 16,
|
||||
},
|
||||
{
|
||||
name: "long output truncation",
|
||||
toolName: "Bash",
|
||||
input: "command",
|
||||
output: string(make([]byte, 5000)), // Very long output
|
||||
wantLen: 16,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
hash := hashRequest(tt.toolName, tt.input, tt.output)
|
||||
assert.Len(t, hash, tt.wantLen)
|
||||
|
||||
// Same inputs should produce same hash
|
||||
hash2 := hashRequest(tt.toolName, tt.input, tt.output)
|
||||
assert.Equal(t, hash, hash2)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestHashRequest_DifferentInputs(t *testing.T) {
|
||||
// Different inputs should produce different hashes
|
||||
hash1 := hashRequest("Read", "file1.txt", "content1")
|
||||
hash2 := hashRequest("Read", "file2.txt", "content2")
|
||||
|
||||
assert.NotEqual(t, hash1, hash2)
|
||||
}
|
||||
|
||||
func TestHashRequest_OutputTruncation(t *testing.T) {
|
||||
// Hash should be the same for outputs that differ only after 1000 chars
|
||||
longOutput1 := string(make([]byte, 1500))
|
||||
longOutput2 := longOutput1[:1000] + "different suffix here"
|
||||
|
||||
hash1 := hashRequest("Read", "input", longOutput1)
|
||||
hash2 := hashRequest("Read", "input", longOutput2)
|
||||
|
||||
// Since we only hash first 1000 chars, these should be the same
|
||||
assert.Equal(t, hash1, hash2)
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// TESTS FOR Processor methods
|
||||
// =============================================================================
|
||||
|
||||
func TestProcessor_CircuitBreakerState(t *testing.T) {
|
||||
p := &Processor{
|
||||
circuitBreaker: NewCircuitBreaker(2, 60),
|
||||
}
|
||||
|
||||
// Initially closed
|
||||
assert.Equal(t, "closed", p.CircuitBreakerState())
|
||||
|
||||
// After enough failures, open
|
||||
p.circuitBreaker.RecordFailure()
|
||||
p.circuitBreaker.RecordFailure()
|
||||
assert.Equal(t, "open", p.CircuitBreakerState())
|
||||
|
||||
// After success, closed
|
||||
p.circuitBreaker.RecordSuccess()
|
||||
assert.Equal(t, "closed", p.CircuitBreakerState())
|
||||
}
|
||||
|
||||
func TestProcessor_CircuitBreakerMetrics(t *testing.T) {
|
||||
p := &Processor{
|
||||
circuitBreaker: NewCircuitBreaker(5, 120),
|
||||
}
|
||||
|
||||
metrics := p.CircuitBreakerMetrics()
|
||||
assert.Equal(t, "closed", metrics.State)
|
||||
assert.Equal(t, int64(0), metrics.Failures)
|
||||
assert.Equal(t, int64(5), metrics.Threshold)
|
||||
assert.Equal(t, int64(120), metrics.ResetTimeoutSecs)
|
||||
|
||||
// Record a failure and check metrics update
|
||||
p.circuitBreaker.RecordFailure()
|
||||
metrics = p.CircuitBreakerMetrics()
|
||||
assert.Equal(t, int64(1), metrics.Failures)
|
||||
assert.Greater(t, metrics.LastFailureUnix, int64(0))
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// TESTS FOR Vector Sync Workers
|
||||
// =============================================================================
|
||||
|
||||
func TestProcessor_StartAndStopVectorSyncWorkers(t *testing.T) {
|
||||
var syncedObservations []*models.Observation
|
||||
var mu sync.Mutex
|
||||
|
||||
p := &Processor{
|
||||
vectorSyncChan: make(chan *models.Observation, MaxVectorSyncWorkers*2),
|
||||
vectorSyncDone: make(chan struct{}),
|
||||
syncObservationFunc: func(obs *models.Observation) {
|
||||
mu.Lock()
|
||||
syncedObservations = append(syncedObservations, obs)
|
||||
mu.Unlock()
|
||||
},
|
||||
}
|
||||
|
||||
// Start workers
|
||||
p.StartVectorSyncWorkers()
|
||||
|
||||
// Send some observations
|
||||
obs1 := &models.Observation{SDKSessionID: "test1"}
|
||||
obs2 := &models.Observation{SDKSessionID: "test2"}
|
||||
p.vectorSyncChan <- obs1
|
||||
p.vectorSyncChan <- obs2
|
||||
|
||||
// Give workers time to process
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
|
||||
// Stop workers
|
||||
p.StopVectorSyncWorkers()
|
||||
|
||||
// Verify observations were synced
|
||||
mu.Lock()
|
||||
assert.Len(t, syncedObservations, 2)
|
||||
mu.Unlock()
|
||||
}
|
||||
|
||||
func TestProcessor_VectorSyncWorker_DrainOnShutdown(t *testing.T) {
|
||||
var syncedCount int
|
||||
var mu sync.Mutex
|
||||
|
||||
p := &Processor{
|
||||
vectorSyncChan: make(chan *models.Observation, 10),
|
||||
vectorSyncDone: make(chan struct{}),
|
||||
syncObservationFunc: func(obs *models.Observation) {
|
||||
mu.Lock()
|
||||
syncedCount++
|
||||
mu.Unlock()
|
||||
},
|
||||
}
|
||||
|
||||
// Queue observations before starting workers
|
||||
for i := 0; i < 5; i++ {
|
||||
p.vectorSyncChan <- &models.Observation{SDKSessionID: "pre-queued"}
|
||||
}
|
||||
|
||||
// Start workers
|
||||
p.StartVectorSyncWorkers()
|
||||
|
||||
// Stop immediately - workers should drain the queue
|
||||
p.StopVectorSyncWorkers()
|
||||
|
||||
// All pre-queued items should have been processed
|
||||
mu.Lock()
|
||||
assert.Equal(t, 5, syncedCount)
|
||||
mu.Unlock()
|
||||
}
|
||||
|
||||
func TestProcessor_VectorSyncWorker_NilSyncFunc(t *testing.T) {
|
||||
p := &Processor{
|
||||
vectorSyncChan: make(chan *models.Observation, 10),
|
||||
vectorSyncDone: make(chan struct{}),
|
||||
syncObservationFunc: nil, // No sync function set
|
||||
}
|
||||
|
||||
// Start workers
|
||||
p.StartVectorSyncWorkers()
|
||||
|
||||
// Send observation - should not panic even with nil sync func
|
||||
p.vectorSyncChan <- &models.Observation{SDKSessionID: "test"}
|
||||
|
||||
// Give it time to process
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
|
||||
// Stop workers - should not panic
|
||||
p.StopVectorSyncWorkers()
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// TESTS FOR CircuitBreaker Additional Behaviors
|
||||
// =============================================================================
|
||||
|
||||
func TestCircuitBreaker_Allow_OpenBlocksRequests(t *testing.T) {
|
||||
cb := NewCircuitBreaker(1, 60)
|
||||
|
||||
// Open the circuit
|
||||
cb.RecordFailure()
|
||||
assert.Equal(t, "open", cb.State())
|
||||
|
||||
// All requests should be blocked
|
||||
assert.False(t, cb.Allow())
|
||||
assert.False(t, cb.Allow())
|
||||
assert.False(t, cb.Allow())
|
||||
}
|
||||
|
||||
func TestCircuitBreaker_MultipleFailures(t *testing.T) {
|
||||
cb := NewCircuitBreaker(3, 60) // Higher threshold
|
||||
|
||||
// Record failures below threshold
|
||||
cb.RecordFailure()
|
||||
assert.Equal(t, "closed", cb.State())
|
||||
assert.Equal(t, int64(1), cb.Metrics().Failures)
|
||||
|
||||
cb.RecordFailure()
|
||||
assert.Equal(t, "closed", cb.State())
|
||||
assert.Equal(t, int64(2), cb.Metrics().Failures)
|
||||
|
||||
// Third failure opens circuit
|
||||
cb.RecordFailure()
|
||||
assert.Equal(t, "open", cb.State())
|
||||
assert.Equal(t, int64(3), cb.Metrics().Failures)
|
||||
}
|
||||
|
||||
func TestCircuitBreaker_SuccessResetsFailures(t *testing.T) {
|
||||
cb := NewCircuitBreaker(5, 60)
|
||||
|
||||
// Record some failures
|
||||
cb.RecordFailure()
|
||||
cb.RecordFailure()
|
||||
assert.Equal(t, int64(2), cb.Metrics().Failures)
|
||||
|
||||
// Success resets failures
|
||||
cb.RecordSuccess()
|
||||
assert.Equal(t, int64(0), cb.Metrics().Failures)
|
||||
assert.Equal(t, "closed", cb.State())
|
||||
}
|
||||
|
||||
func TestCircuitBreaker_Metrics_Comprehensive(t *testing.T) {
|
||||
cb := NewCircuitBreaker(5, 120)
|
||||
|
||||
// Initial state
|
||||
metrics := cb.Metrics()
|
||||
assert.Equal(t, "closed", metrics.State)
|
||||
assert.Equal(t, int64(0), metrics.Failures)
|
||||
assert.Equal(t, int64(5), metrics.Threshold)
|
||||
assert.Equal(t, int64(120), metrics.ResetTimeoutSecs)
|
||||
assert.Equal(t, int64(0), metrics.LastFailureUnix)
|
||||
assert.Equal(t, int64(0), metrics.SecondsUntilReset)
|
||||
|
||||
// After failures that open circuit
|
||||
for i := 0; i < 5; i++ {
|
||||
cb.RecordFailure()
|
||||
}
|
||||
metrics = cb.Metrics()
|
||||
assert.Equal(t, "open", metrics.State)
|
||||
assert.Equal(t, int64(5), metrics.Failures)
|
||||
assert.Greater(t, metrics.LastFailureUnix, int64(0))
|
||||
assert.Greater(t, metrics.SecondsUntilReset, int64(0))
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// TESTS FOR MaxVectorSyncWorkers constant
|
||||
// =============================================================================
|
||||
|
||||
func TestMaxVectorSyncWorkers(t *testing.T) {
|
||||
assert.Equal(t, 8, MaxVectorSyncWorkers)
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// ADDITIONAL EDGE CASE TESTS
|
||||
// =============================================================================
|
||||
|
||||
func TestRequestDeduplicator_IsDuplicate_ExpiredEntry(t *testing.T) {
|
||||
if testing.Short() {
|
||||
t.Skip("Skipping time-dependent test in short mode")
|
||||
}
|
||||
// Use a 1-second TTL with enough margin
|
||||
d := NewRequestDeduplicator(1, 100)
|
||||
|
||||
hash := "expiretest"
|
||||
d.Record(hash)
|
||||
|
||||
// Initially duplicate
|
||||
assert.True(t, d.IsDuplicate(hash))
|
||||
|
||||
// Wait for TTL to expire (2.5 seconds to ensure crossing second boundaries)
|
||||
time.Sleep(2500 * time.Millisecond)
|
||||
|
||||
// Should no longer be considered duplicate
|
||||
assert.False(t, d.IsDuplicate(hash))
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// TESTS FOR ProcessObservation Early Returns
|
||||
// =============================================================================
|
||||
|
||||
func TestProcessObservation_SkipTool(t *testing.T) {
|
||||
p := &Processor{
|
||||
circuitBreaker: NewCircuitBreaker(5, 60),
|
||||
deduplicator: NewRequestDeduplicator(300, 1000),
|
||||
}
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// TodoWrite should be skipped
|
||||
err := p.ProcessObservation(ctx, "session-1", "project-1", "TodoWrite",
|
||||
map[string]string{"content": "test"}, "success", 1, "/test/cwd")
|
||||
assert.NoError(t, err)
|
||||
|
||||
// Glob should be skipped
|
||||
err = p.ProcessObservation(ctx, "session-1", "project-1", "Glob",
|
||||
map[string]string{"pattern": "*.go"}, []string{"main.go", "test.go"}, 1, "/test/cwd")
|
||||
assert.NoError(t, err)
|
||||
|
||||
// AskUserQuestion should be skipped
|
||||
err = p.ProcessObservation(ctx, "session-1", "project-1", "AskUserQuestion",
|
||||
"question", "answer", 1, "/test/cwd")
|
||||
assert.NoError(t, err)
|
||||
}
|
||||
|
||||
func TestProcessObservation_SkipTrivial(t *testing.T) {
|
||||
p := &Processor{
|
||||
circuitBreaker: NewCircuitBreaker(5, 60),
|
||||
deduplicator: NewRequestDeduplicator(300, 1000),
|
||||
}
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Short output should be skipped
|
||||
err := p.ProcessObservation(ctx, "session-1", "project-1", "Read",
|
||||
map[string]string{"file_path": "/test.go"}, "short", 1, "/test/cwd")
|
||||
assert.NoError(t, err)
|
||||
|
||||
// "No matches found" should be skipped
|
||||
err = p.ProcessObservation(ctx, "session-1", "project-1", "Grep",
|
||||
map[string]string{"pattern": "test"}, "No matches found in the repository", 1, "/test/cwd")
|
||||
assert.NoError(t, err)
|
||||
}
|
||||
|
||||
func TestProcessObservation_SkipDuplicate(t *testing.T) {
|
||||
p := &Processor{
|
||||
circuitBreaker: NewCircuitBreaker(5, 60),
|
||||
deduplicator: NewRequestDeduplicator(300, 1000),
|
||||
sem: make(chan struct{}, 4),
|
||||
claudePath: "/nonexistent/path", // Will fail at CLI call stage
|
||||
}
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Valid input that would be processed
|
||||
input := map[string]string{"file_path": "/project/main.go"}
|
||||
output := "package main\n\nimport \"fmt\"\n\nfunc main() {\n\tfmt.Println(\"Hello World\")\n}"
|
||||
|
||||
// First call should try to process (will fail because claudePath doesn't exist)
|
||||
err := p.ProcessObservation(ctx, "session-1", "project-1", "Read", input, output, 1, "/test/cwd")
|
||||
// Expect error because claudePath doesn't exist
|
||||
assert.Error(t, err)
|
||||
|
||||
// Second call with same input should be skipped as duplicate
|
||||
err = p.ProcessObservation(ctx, "session-1", "project-1", "Read", input, output, 1, "/test/cwd")
|
||||
assert.NoError(t, err) // No error because it was skipped as duplicate
|
||||
}
|
||||
|
||||
func TestProcessObservation_CircuitBreakerOpen(t *testing.T) {
|
||||
cb := NewCircuitBreaker(1, 60) // Threshold of 1
|
||||
cb.RecordFailure() // Open the circuit breaker
|
||||
|
||||
p := &Processor{
|
||||
circuitBreaker: cb,
|
||||
deduplicator: NewRequestDeduplicator(300, 1000),
|
||||
}
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Valid input that would be processed
|
||||
input := map[string]string{"file_path": "/project/main.go"}
|
||||
output := "package main\n\nimport \"fmt\"\n\nfunc main() {\n\tfmt.Println(\"Hello World\")\n}"
|
||||
|
||||
err := p.ProcessObservation(ctx, "session-1", "project-1", "Read", input, output, 1, "/test/cwd")
|
||||
assert.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "circuit breaker open")
|
||||
}
|
||||
|
||||
func TestProcessObservation_ContextCancel(t *testing.T) {
|
||||
p := &Processor{
|
||||
circuitBreaker: NewCircuitBreaker(5, 60),
|
||||
deduplicator: NewRequestDeduplicator(300, 1000),
|
||||
sem: make(chan struct{}, 1), // Small semaphore
|
||||
claudePath: "/fake/claude",
|
||||
}
|
||||
|
||||
// Fill the semaphore
|
||||
p.sem <- struct{}{}
|
||||
|
||||
// Create a cancelled context
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel() // Cancel immediately
|
||||
|
||||
// Valid input that would be processed
|
||||
input := map[string]string{"file_path": "/project/main.go"}
|
||||
output := "package main\n\nimport \"fmt\"\n\nfunc main() {\n\tfmt.Println(\"Hello World\")\n}"
|
||||
|
||||
err := p.ProcessObservation(ctx, "session-1", "project-1", "Read", input, output, 1, "/test/cwd")
|
||||
assert.Error(t, err)
|
||||
assert.ErrorIs(t, err, context.Canceled)
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// TESTS FOR ProcessSummary Early Returns
|
||||
// =============================================================================
|
||||
|
||||
func TestProcessSummary_SkipEmptyRequest(t *testing.T) {
|
||||
p := &Processor{
|
||||
circuitBreaker: NewCircuitBreaker(5, 60),
|
||||
deduplicator: NewRequestDeduplicator(300, 1000),
|
||||
}
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Empty request should be skipped (sessionDBID, sdkSessionID, project, userPrompt, lastUserMsg, lastAssistantMsg)
|
||||
err := p.ProcessSummary(ctx, 1, "session-1", "project-1", "", "", "")
|
||||
assert.NoError(t, err)
|
||||
}
|
||||
|
||||
func TestProcessSummary_CircuitBreakerOpen(t *testing.T) {
|
||||
cb := NewCircuitBreaker(1, 60)
|
||||
cb.RecordFailure() // Open the circuit breaker
|
||||
|
||||
p := &Processor{
|
||||
circuitBreaker: cb,
|
||||
deduplicator: NewRequestDeduplicator(300, 1000),
|
||||
sem: make(chan struct{}, 4),
|
||||
claudePath: "/nonexistent/path",
|
||||
}
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Meaningful assistant message (> 200 chars, contains code discussion)
|
||||
assistantMsg := `I've updated the handler.go file to fix the authentication bug.
|
||||
The function validateToken() was not checking token expiry correctly.
|
||||
I've added a check for the exp claim and implemented proper error handling.
|
||||
The changes have been tested and the build passes successfully.
|
||||
Here's the implementation details and code review.`
|
||||
|
||||
// Valid request but circuit breaker is open
|
||||
err := p.ProcessSummary(ctx, 1, "session-1", "project-1",
|
||||
"Implement authentication", "User message", assistantMsg)
|
||||
assert.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "claude CLI failed")
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// TESTS FOR callClaudeCLI Error Paths
|
||||
// =============================================================================
|
||||
|
||||
func TestCallClaudeCLI_PromptTooLarge(t *testing.T) {
|
||||
p := &Processor{
|
||||
claudePath: "/fake/claude",
|
||||
}
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Create a prompt that exceeds MaxPromptSize
|
||||
largePrompt := string(make([]byte, MaxPromptSize+1))
|
||||
|
||||
_, err := p.callClaudeCLI(ctx, largePrompt)
|
||||
assert.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "prompt exceeds maximum size")
|
||||
}
|
||||
|
||||
func TestCallClaudeCLI_BinaryNotFound(t *testing.T) {
|
||||
p := &Processor{
|
||||
claudePath: "/nonexistent/path/to/claude",
|
||||
}
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
_, err := p.callClaudeCLI(ctx, "test prompt")
|
||||
assert.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "claude CLI failed")
|
||||
}
|
||||
|
||||
@@ -24,12 +24,12 @@ var ObservationConcepts = []string{
|
||||
|
||||
// ToolExecution represents a tool execution for observation.
|
||||
type ToolExecution struct {
|
||||
ID int64
|
||||
ToolName string
|
||||
ToolInput string
|
||||
ToolOutput string
|
||||
CreatedAtEpoch int64
|
||||
CWD string
|
||||
ID int64
|
||||
CreatedAtEpoch int64
|
||||
}
|
||||
|
||||
// BuildObservationPrompt builds a prompt for processing a tool observation.
|
||||
@@ -49,6 +49,10 @@ func BuildObservationPrompt(exec ToolExecution) string {
|
||||
inputJSON, _ := json.MarshalIndent(toolInput, " ", " ")
|
||||
outputJSON, _ := json.MarshalIndent(toolOutput, " ", " ")
|
||||
|
||||
// Strip system XML artifacts from tool data before building prompt
|
||||
cleanInput := StripSystemXML(string(inputJSON))
|
||||
cleanOutput := StripSystemXML(string(outputJSON))
|
||||
|
||||
timestamp := time.UnixMilli(exec.CreatedAtEpoch).Format(time.RFC3339)
|
||||
|
||||
var sb strings.Builder
|
||||
@@ -58,8 +62,8 @@ func BuildObservationPrompt(exec ToolExecution) string {
|
||||
if exec.CWD != "" {
|
||||
sb.WriteString(fmt.Sprintf(" <working_directory>%s</working_directory>\n", exec.CWD))
|
||||
}
|
||||
sb.WriteString(fmt.Sprintf(" <parameters>%s</parameters>\n", truncate(string(inputJSON), 3000)))
|
||||
sb.WriteString(fmt.Sprintf(" <outcome>%s</outcome>\n", truncate(string(outputJSON), 5000)))
|
||||
sb.WriteString(fmt.Sprintf(" <parameters>%s</parameters>\n", truncate(cleanInput, 3000)))
|
||||
sb.WriteString(fmt.Sprintf(" <outcome>%s</outcome>\n", truncate(cleanOutput, 5000)))
|
||||
sb.WriteString("</observed_from_primary_session>")
|
||||
|
||||
return sb.String()
|
||||
@@ -67,12 +71,12 @@ func BuildObservationPrompt(exec ToolExecution) string {
|
||||
|
||||
// SummaryRequest contains data for building a summary prompt.
|
||||
type SummaryRequest struct {
|
||||
SessionDBID int64
|
||||
SDKSessionID string
|
||||
Project string
|
||||
UserPrompt string
|
||||
LastUserMessage string
|
||||
LastAssistantMessage string
|
||||
SessionDBID int64
|
||||
}
|
||||
|
||||
// BuildSummaryPrompt builds a prompt requesting a session summary.
|
||||
@@ -84,8 +88,10 @@ func BuildSummaryPrompt(req SummaryRequest) string {
|
||||
sb.WriteString("Write progress notes of what was done, what was learned, and what's next. This is a checkpoint to capture progress so far. The session is ongoing - you may receive more requests and tool executions after this summary. Write \"next_steps\" as the current trajectory of work (what's actively being worked on or coming up next), not as post-session future work. Always write at least a minimal summary explaining current progress, even if work is still in early stages, so that users see a summary output tied to each request.\n\n")
|
||||
|
||||
if req.LastAssistantMessage != "" {
|
||||
// Strip system XML artifacts from captured transcript content
|
||||
cleanMsg := StripSystemXML(req.LastAssistantMessage)
|
||||
sb.WriteString("Claude's Full Response to User:\n")
|
||||
sb.WriteString(truncate(req.LastAssistantMessage, 4000))
|
||||
sb.WriteString(truncate(cleanMsg, 4000))
|
||||
sb.WriteString("\n\n")
|
||||
}
|
||||
|
||||
|
||||
@@ -12,8 +12,8 @@ func TestTruncate(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
input string
|
||||
maxLen int
|
||||
expected string
|
||||
maxLen int
|
||||
}{
|
||||
{
|
||||
name: "shorter_than_max",
|
||||
@@ -60,8 +60,8 @@ func TestBuildObservationPrompt(t *testing.T) {
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
exec ToolExecution
|
||||
contains []string
|
||||
exec ToolExecution
|
||||
}{
|
||||
{
|
||||
name: "basic_read_tool",
|
||||
|
||||
@@ -0,0 +1,8 @@
|
||||
// Package sdk provides SDK agent integration for claude-mnemonic.
|
||||
package sdk
|
||||
|
||||
import "github.com/lukaszraczylo/claude-mnemonic/pkg/sanitize"
|
||||
|
||||
// StripSystemXML removes known Claude Code internal XML blocks from text.
|
||||
// Delegates to the shared sanitize package.
|
||||
var StripSystemXML = sanitize.StripSystemXML
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user