mirror of
https://github.com/lukaszraczylo/claude-mnemonic.git
synced 2026-10-12 04:55:18 +00:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
cce1c62184 | ||
|
|
9629cff306 | ||
|
|
bc336a892c | ||
|
|
949e82b8e6 | ||
|
|
b09d9aafa8 | ||
|
|
03bc296ab4 | ||
|
|
73dacf5f3f | ||
|
|
2084a28548 | ||
|
|
6541d63f68 | ||
|
|
c331899040 | ||
|
|
99850838e8 | ||
|
|
d8a4939201 | ||
|
|
50f078bc6a | ||
|
|
8ff08732c6 | ||
|
|
433d165f17 | ||
|
|
378e7dc489 | ||
|
|
c2c31c8c5d | ||
|
|
b396d05c0f | ||
|
|
66c53deb02 | ||
|
|
372942ec8b | ||
|
|
e3d3f2bb7b | ||
|
|
21bcbfae07 | ||
|
|
03e366da0f | ||
|
|
f8d98f7b6b | ||
|
|
7ce7b80650 | ||
|
|
f23df63059 | ||
|
|
f56037a49a | ||
|
|
d5ccbd641e | ||
|
|
862ff7c3ac | ||
|
|
e7b58312c5 | ||
|
|
e5f4a100ab | ||
|
|
31d67ceb2a | ||
|
|
5c7b623585 | ||
|
|
db528b2544 | ||
|
|
f452591971 | ||
|
|
7fe679f83b | ||
|
|
6c28ecb22a | ||
|
|
7e49113d24 | ||
|
|
e52f3284ba | ||
|
|
f0c35eef0e | ||
|
|
085c2cb3d0 | ||
|
|
cd14ba8915 | ||
|
|
0a0dfc1c2f | ||
|
|
d68a700917 | ||
|
|
c705449467 | ||
|
|
ad8aef92e1 | ||
|
|
268c286859 | ||
|
|
c664ad5fd8 | ||
|
|
0eca679139 | ||
|
|
c14677f1f2 | ||
|
|
ca711a8a7e | ||
|
|
0769feaca5 | ||
|
|
bbd6c3f1ff | ||
|
|
637412e8bb | ||
|
|
a1c1a5cad5 | ||
|
|
3a223da320 | ||
|
|
feb7f251b4 |
@@ -1 +0,0 @@
|
||||
internal/embedding/assets/*.onnx filter=lfs diff=lfs merge=lfs -text
|
||||
@@ -1,114 +0,0 @@
|
||||
# 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
|
||||
@@ -1,113 +0,0 @@
|
||||
# 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.
|
||||
@@ -1,63 +0,0 @@
|
||||
# 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
|
||||
@@ -17,5 +17,4 @@ jobs:
|
||||
with:
|
||||
go-version: ">=1.24"
|
||||
release-workflow: "release.yaml"
|
||||
lfs: true
|
||||
secrets: inherit
|
||||
|
||||
@@ -20,6 +20,4 @@ jobs:
|
||||
uses: lukaszraczylo/shared-actions/.github/workflows/go-pr.yaml@main
|
||||
with:
|
||||
go-version: ">=1.24"
|
||||
lfs: true
|
||||
build-tags: "fts5"
|
||||
secrets: inherit
|
||||
|
||||
+102
-103
@@ -1,121 +1,120 @@
|
||||
name: Release
|
||||
|
||||
on:
|
||||
push:
|
||||
branches:
|
||||
- main
|
||||
paths-ignore:
|
||||
- "**.md"
|
||||
- "docs/**"
|
||||
workflow_dispatch:
|
||||
push:
|
||||
branches:
|
||||
- main
|
||||
paths-ignore:
|
||||
- "**.md"
|
||||
- "docs/**"
|
||||
workflow_dispatch:
|
||||
|
||||
permissions:
|
||||
id-token: write
|
||||
contents: write
|
||||
packages: write
|
||||
id-token: write
|
||||
contents: write
|
||||
packages: write
|
||||
|
||||
jobs:
|
||||
release:
|
||||
uses: lukaszraczylo/shared-actions/.github/workflows/go-release-cgo.yaml@main
|
||||
release:
|
||||
uses: lukaszraczylo/shared-actions/.github/workflows/go-release-cgo.yaml@main
|
||||
with:
|
||||
go-version: ">=1.24"
|
||||
node-enabled: true
|
||||
node-version: "20"
|
||||
node-build-script: "cd ui && npm ci && npm run build"
|
||||
node-cache-dependency-path: "ui/package-lock.json"
|
||||
node-output-path: "ui/dist"
|
||||
node-embed-path: "internal/worker/static"
|
||||
secrets: inherit
|
||||
|
||||
commit-marketplace:
|
||||
needs: release
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
with:
|
||||
go-version: ">=1.24"
|
||||
node-enabled: true
|
||||
node-version: "20"
|
||||
node-build-script: "cd ui && npm ci && npm run build"
|
||||
node-cache-dependency-path: "ui/package-lock.json"
|
||||
node-output-path: "ui/dist"
|
||||
node-embed-path: "internal/worker/static"
|
||||
lfs: true
|
||||
secrets: inherit
|
||||
ref: main
|
||||
|
||||
commit-marketplace:
|
||||
needs: release
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
with:
|
||||
ref: main
|
||||
- name: Download checksums
|
||||
run: |
|
||||
VERSION="${{ needs.release.outputs.version_tag }}"
|
||||
curl -sSL -o checksums.txt \
|
||||
"https://github.com/lukaszraczylo/claude-mnemonic/releases/download/${VERSION}/checksums.txt"
|
||||
|
||||
- name: Download checksums
|
||||
run: |
|
||||
VERSION="${{ needs.release.outputs.version_tag }}"
|
||||
curl -sSL -o checksums.txt \
|
||||
"https://github.com/lukaszraczylo/claude-mnemonic/releases/download/${VERSION}/checksums.txt"
|
||||
- name: Generate marketplace.json with checksums
|
||||
run: |
|
||||
VERSION="${{ needs.release.outputs.version }}"
|
||||
VERSION_TAG="${{ needs.release.outputs.version_tag }}"
|
||||
TODAY=$(date -u +%Y-%m-%d)
|
||||
|
||||
- name: Generate marketplace.json with checksums
|
||||
run: |
|
||||
VERSION="${{ needs.release.outputs.version }}"
|
||||
VERSION_TAG="${{ needs.release.outputs.version_tag }}"
|
||||
TODAY=$(date -u +%Y-%m-%d)
|
||||
# Extract checksums from checksums.txt
|
||||
SHA_DARWIN_ARM64=$(grep "darwin_arm64" checksums.txt | awk '{print $1}')
|
||||
SHA_LINUX_AMD64=$(grep "linux_amd64" checksums.txt | awk '{print $1}')
|
||||
SHA_WINDOWS_AMD64=$(grep "windows_amd64" checksums.txt | awk '{print $1}')
|
||||
|
||||
# Extract checksums from checksums.txt
|
||||
SHA_DARWIN_ARM64=$(grep "darwin_arm64" checksums.txt | awk '{print $1}')
|
||||
SHA_LINUX_AMD64=$(grep "linux_amd64" checksums.txt | awk '{print $1}')
|
||||
SHA_WINDOWS_AMD64=$(grep "windows_amd64" checksums.txt | awk '{print $1}')
|
||||
|
||||
# Generate marketplace.json
|
||||
cat > marketplace.json << EOF
|
||||
{
|
||||
"\$schema": "https://anthropic.com/claude-code/marketplace.schema.json",
|
||||
"name": "claude-mnemonic",
|
||||
"version": "1.0.0",
|
||||
"description": "Persistent memory system for Claude Code - stores observations, session summaries, and user prompts with semantic search",
|
||||
"owner": {
|
||||
"name": "lukaszraczylo",
|
||||
"email": "[email protected]",
|
||||
"url": "https://github.com/lukaszraczylo"
|
||||
},
|
||||
"repository": "https://github.com/lukaszraczylo/claude-mnemonic",
|
||||
"plugins": [
|
||||
{
|
||||
"name": "claude-mnemonic",
|
||||
"description": "Persistent memory system for Claude Code - Go implementation with SQLite and ChromaDB vector search",
|
||||
"version": "${VERSION}",
|
||||
"author": {
|
||||
"name": "lukaszraczylo",
|
||||
"url": "https://github.com/lukaszraczylo"
|
||||
# Generate marketplace.json
|
||||
cat > marketplace.json << EOF
|
||||
{
|
||||
"\$schema": "https://anthropic.com/claude-code/marketplace.schema.json",
|
||||
"name": "claude-mnemonic",
|
||||
"version": "1.0.0",
|
||||
"description": "Persistent memory system for Claude Code - stores observations, session summaries, and user prompts with semantic search",
|
||||
"owner": {
|
||||
"name": "lukaszraczylo",
|
||||
"email": "[email protected]",
|
||||
"url": "https://github.com/lukaszraczylo"
|
||||
},
|
||||
"repository": "https://github.com/lukaszraczylo/claude-mnemonic",
|
||||
"plugins": [
|
||||
{
|
||||
"name": "claude-mnemonic",
|
||||
"description": "Persistent memory system for Claude Code - Go implementation with SQLite and ChromaDB vector search",
|
||||
"version": "${VERSION}",
|
||||
"author": {
|
||||
"name": "lukaszraczylo",
|
||||
"url": "https://github.com/lukaszraczylo"
|
||||
},
|
||||
"category": "productivity",
|
||||
"tags": ["memory", "persistence", "search", "context"],
|
||||
"license": "MIT",
|
||||
"homepage": "https://github.com/lukaszraczylo/claude-mnemonic",
|
||||
"releases": {
|
||||
"latest": "${VERSION}",
|
||||
"versions": {
|
||||
"${VERSION}": {
|
||||
"releaseDate": "${TODAY}",
|
||||
"downloads": {
|
||||
"darwin-arm64": {
|
||||
"url": "https://github.com/lukaszraczylo/claude-mnemonic/releases/download/${VERSION_TAG}/claude-mnemonic_${VERSION}_darwin_arm64.tar.gz",
|
||||
"format": "tar.gz",
|
||||
"sha256": "${SHA_DARWIN_ARM64}"
|
||||
},
|
||||
"category": "productivity",
|
||||
"tags": ["memory", "persistence", "search", "context"],
|
||||
"license": "MIT",
|
||||
"homepage": "https://github.com/lukaszraczylo/claude-mnemonic",
|
||||
"releases": {
|
||||
"latest": "${VERSION}",
|
||||
"versions": {
|
||||
"${VERSION}": {
|
||||
"releaseDate": "${TODAY}",
|
||||
"downloads": {
|
||||
"darwin-arm64": {
|
||||
"url": "https://github.com/lukaszraczylo/claude-mnemonic/releases/download/${VERSION_TAG}/claude-mnemonic_${VERSION}_darwin_arm64.tar.gz",
|
||||
"format": "tar.gz",
|
||||
"sha256": "${SHA_DARWIN_ARM64}"
|
||||
},
|
||||
"linux-amd64": {
|
||||
"url": "https://github.com/lukaszraczylo/claude-mnemonic/releases/download/${VERSION_TAG}/claude-mnemonic_${VERSION}_linux_amd64.tar.gz",
|
||||
"format": "tar.gz",
|
||||
"sha256": "${SHA_LINUX_AMD64}"
|
||||
},
|
||||
"windows-amd64": {
|
||||
"url": "https://github.com/lukaszraczylo/claude-mnemonic/releases/download/${VERSION_TAG}/claude-mnemonic_${VERSION}_windows_amd64.zip",
|
||||
"format": "zip",
|
||||
"sha256": "${SHA_WINDOWS_AMD64}"
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
"linux-amd64": {
|
||||
"url": "https://github.com/lukaszraczylo/claude-mnemonic/releases/download/${VERSION_TAG}/claude-mnemonic_${VERSION}_linux_amd64.tar.gz",
|
||||
"format": "tar.gz",
|
||||
"sha256": "${SHA_LINUX_AMD64}"
|
||||
},
|
||||
"windows-amd64": {
|
||||
"url": "https://github.com/lukaszraczylo/claude-mnemonic/releases/download/${VERSION_TAG}/claude-mnemonic_${VERSION}_windows_amd64.zip",
|
||||
"format": "zip",
|
||||
"sha256": "${SHA_WINDOWS_AMD64}"
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
}
|
||||
EOF
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
EOF
|
||||
|
||||
echo "Generated marketplace.json:"
|
||||
cat marketplace.json
|
||||
echo "Generated marketplace.json:"
|
||||
cat marketplace.json
|
||||
|
||||
- name: Commit and push
|
||||
run: |
|
||||
git config user.name "github-actions[bot]"
|
||||
git config user.email "github-actions[bot]@users.noreply.github.com"
|
||||
git add marketplace.json
|
||||
git diff --staged --quiet || git commit -m "chore: update marketplace for ${{ needs.release.outputs.version_tag }}"
|
||||
git push
|
||||
- name: Commit and push
|
||||
run: |
|
||||
git config user.name "github-actions[bot]"
|
||||
git config user.email "github-actions[bot]@users.noreply.github.com"
|
||||
git add marketplace.json
|
||||
git diff --staged --quiet || git commit -m "chore: update marketplace for ${{ needs.release.outputs.version_tag }}"
|
||||
git push
|
||||
|
||||
@@ -83,10 +83,3 @@ logs/
|
||||
dist/
|
||||
docs/dist
|
||||
.claude-plugin
|
||||
|
||||
# Auto-generated plugin configs (generated by scripts/generate-plugin-config.sh)
|
||||
.claude-plugin/
|
||||
|
||||
# Non-template plugin configs (keep only .tpl files)
|
||||
plugin/.claude-plugin/plugin.json
|
||||
plugin/.claude-plugin/marketplace.json
|
||||
|
||||
@@ -1,35 +0,0 @@
|
||||
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-dirs:
|
||||
- vendor
|
||||
# Exclude some linters from running on test files
|
||||
exclude-rules:
|
||||
- path: _test\.go
|
||||
linters:
|
||||
- errcheck
|
||||
- gosec
|
||||
|
||||
run:
|
||||
timeout: 5m
|
||||
tests: true
|
||||
@@ -146,8 +146,8 @@ install: build stop-worker
|
||||
@# Copy slash commands if they exist
|
||||
@if [ -d "$(PLUGIN_DIR)/commands" ]; then cp -r $(PLUGIN_DIR)/commands/* $(HOME)/.claude/plugins/marketplaces/claude-mnemonic/commands/ 2>/dev/null || true; fi
|
||||
@# Update plugin.json and marketplace.json with current version to prevent stale version directories
|
||||
@sed 's/{{ .Version }}/$(VERSION)/g; s/{{.Version}}/$(VERSION)/g' $(PLUGIN_DIR)/.claude-plugin/plugin.json.tpl > $(HOME)/.claude/plugins/marketplaces/claude-mnemonic/.claude-plugin/plugin.json
|
||||
@sed 's/{{ .Version }}/$(VERSION)/g; s/{{.Version}}/$(VERSION)/g' $(PLUGIN_DIR)/.claude-plugin/marketplace.json.tpl > $(HOME)/.claude/plugins/marketplaces/claude-mnemonic/.claude-plugin/marketplace.json
|
||||
@sed 's/"version": "[^"]*"/"version": "$(VERSION)"/g' $(PLUGIN_DIR)/.claude-plugin/plugin.json > $(HOME)/.claude/plugins/marketplaces/claude-mnemonic/.claude-plugin/plugin.json
|
||||
@sed 's/"version": "[^"]*"/"version": "$(VERSION)"/g' $(PLUGIN_DIR)/.claude-plugin/marketplace.json > $(HOME)/.claude/plugins/marketplaces/claude-mnemonic/.claude-plugin/marketplace.json
|
||||
@echo "Registering plugin with Claude Code..."
|
||||
@./scripts/register-plugin.sh "$(VERSION)"
|
||||
@$(MAKE) start-worker
|
||||
|
||||
@@ -12,21 +12,7 @@ Claude Code forgets everything when your session ends. Claude Mnemonic fixes tha
|
||||
|
||||
It captures what Claude learns during your coding sessions - bug fixes, architecture decisions, patterns that work - and brings that knowledge back in future conversations. No more re-explaining your codebase.
|
||||
|
||||

|
||||
|
||||
## What's New in v0.7
|
||||
|
||||
- **Two-Stage Retrieval** - Cross-encoder reranking for dramatically improved search relevance
|
||||
- **Knowledge Graph** - Automatic relationship detection between observations with visual graph in dashboard
|
||||

|
||||
- **Pattern Detection** - Identifies recurring patterns across sessions and projects
|
||||
- **Importance Scoring** - Time decay and voting system to surface the most valuable memories
|
||||
- **Query Expansion** - Reformulates searches to find semantically related content
|
||||
- **Conflict Detection** - Identifies and resolves contradictory observations
|
||||
- **Observation Lifecycle** - Memories can be superseded when better information arrives
|
||||
|
||||
<details>
|
||||
<summary>Previous: v0.6</summary>
|
||||
## What's New in v0.6
|
||||
|
||||
- **Auto-Updates** - Automatically stays up-to-date with the latest version
|
||||
- **Slash Command: `/restart`** - Restart the worker directly from Claude Code
|
||||
@@ -34,7 +20,6 @@ It captures what Claude learns during your coding sessions - bug fixes, architec
|
||||
- **Async Queue Processing** - Non-blocking observation capture for faster sessions
|
||||
- **Smarter Storage** - Filters out system/agent summaries to keep knowledge relevant
|
||||
- **Improved Reliability** - Better handling of connectivity issues and dead connections
|
||||
</details>
|
||||
|
||||
## Requirements
|
||||
|
||||
@@ -125,39 +110,16 @@ Config file: `~/.claude-mnemonic/settings.json`
|
||||
{
|
||||
"CLAUDE_MNEMONIC_WORKER_PORT": 37777,
|
||||
"CLAUDE_MNEMONIC_CONTEXT_OBSERVATIONS": 100,
|
||||
"CLAUDE_MNEMONIC_CONTEXT_FULL_COUNT": 25,
|
||||
"CLAUDE_MNEMONIC_RERANKING_ENABLED": true
|
||||
"CLAUDE_MNEMONIC_CONTEXT_FULL_COUNT": 25
|
||||
}
|
||||
```
|
||||
|
||||
### Core Settings
|
||||
|
||||
| Variable | Default | What it does |
|
||||
|----------|---------|--------------|
|
||||
| `WORKER_PORT` | `37777` | Dashboard & API port |
|
||||
| `CONTEXT_OBSERVATIONS` | `100` | Max memories per session |
|
||||
| `CONTEXT_FULL_COUNT` | `25` | Full detail memories (rest are condensed) |
|
||||
| `CONTEXT_FULL_COUNT` | `25` | Full detail memories (rest are condensed to title only) |
|
||||
| `CONTEXT_SESSION_COUNT` | `10` | Recent sessions to reference |
|
||||
| `CONTEXT_RELEVANCE_THRESHOLD` | `0.3` | Minimum similarity score (0.0-1.0) for inclusion |
|
||||
| `CONTEXT_MAX_PROMPT_RESULTS` | `10` | Max results per prompt search |
|
||||
|
||||
### Reranking Settings (Two-Stage Retrieval)
|
||||
|
||||
| Variable | Default | What it does |
|
||||
|----------|---------|--------------|
|
||||
| `RERANKING_ENABLED` | `true` | Enable cross-encoder reranking |
|
||||
| `RERANKING_CANDIDATES` | `100` | Candidates to retrieve before reranking |
|
||||
| `RERANKING_RESULTS` | `10` | Final results after reranking |
|
||||
| `RERANKING_ALPHA` | `0.7` | Score blend: alpha×rerank + (1-alpha)×original |
|
||||
| `RERANKING_PURE_MODE` | `false` | Use pure cross-encoder scores only |
|
||||
|
||||
### Embedding Settings
|
||||
|
||||
| Variable | Default | What it does |
|
||||
|----------|---------|--------------|
|
||||
| `EMBEDDING_MODEL` | `bge-v1.5` | Embedding model for semantic search |
|
||||
|
||||
All variables are prefixed with `CLAUDE_MNEMONIC_` in the config file.
|
||||
|
||||
## Project vs Global scope
|
||||
|
||||
@@ -240,11 +202,10 @@ curl -sSL https://raw.githubusercontent.com/lukaszraczylo/claude-mnemonic/main/s
|
||||
|
||||
- **SQLite + FTS5** - Full-text search for exact matches
|
||||
- **sqlite-vec** - Vector database embedded in SQLite
|
||||
- **Two-Stage Retrieval** - Bi-encoder (embedding) + cross-encoder (reranking) for high accuracy
|
||||
- **Local Models** - all-MiniLM-L6-v2 for embeddings, BGE reranker for relevance scoring
|
||||
- **all-MiniLM-L6-v2** - Local embedding model (384 dimensions) via ONNX Runtime
|
||||
- **Go** - Single binary, no external dependencies
|
||||
|
||||
Everything runs locally. No Python. No external vector database. No API calls.
|
||||
Everything runs locally. No Python. No external vector database. No API calls for embeddings.
|
||||
|
||||
## Platform support
|
||||
|
||||
|
||||
@@ -2,7 +2,9 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/url"
|
||||
"os"
|
||||
"strings"
|
||||
@@ -12,41 +14,71 @@ import (
|
||||
|
||||
// Input is the hook input from Claude Code.
|
||||
type Input struct {
|
||||
hooks.BaseInput
|
||||
Source string `json:"source"` // "startup", "resume", "clear", "compact"
|
||||
SessionID string `json:"session_id"`
|
||||
CWD string `json:"cwd"`
|
||||
PermissionMode string `json:"permission_mode"`
|
||||
HookEventName string `json:"hook_event_name"`
|
||||
Source string `json:"source"` // "startup", "resume", "clear", "compact"
|
||||
}
|
||||
|
||||
// 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() {
|
||||
hooks.RunHook("SessionStart", handleSessionStart)
|
||||
}
|
||||
// Skip if this is an internal call (from SDK processor)
|
||||
if os.Getenv("CLAUDE_MNEMONIC_INTERNAL") == "1" {
|
||||
hooks.WriteResponse("SessionStart", true)
|
||||
return
|
||||
}
|
||||
|
||||
// Read input from stdin
|
||||
inputData, err := io.ReadAll(os.Stdin)
|
||||
if err != nil {
|
||||
hooks.WriteError("SessionStart", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
|
||||
var input Input
|
||||
if err := json.Unmarshal(inputData, &input); err != nil {
|
||||
hooks.WriteError("SessionStart", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
|
||||
// Ensure worker is running
|
||||
port, err := hooks.EnsureWorkerRunning()
|
||||
if err != nil {
|
||||
hooks.WriteError("SessionStart", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
|
||||
// Generate unique project ID from CWD (dirname_hash format)
|
||||
project := hooks.ProjectIDWithName(input.CWD)
|
||||
|
||||
func handleSessionStart(ctx *hooks.HookContext, input *Input) (string, error) {
|
||||
// Fetch observations for context injection
|
||||
endpoint := fmt.Sprintf("/api/context/inject?project=%s&cwd=%s",
|
||||
url.QueryEscape(ctx.Project),
|
||||
url.QueryEscape(ctx.CWD))
|
||||
url.QueryEscape(project),
|
||||
url.QueryEscape(input.CWD))
|
||||
|
||||
result, err := hooks.GET(ctx.Port, endpoint)
|
||||
result, err := hooks.GET(port, endpoint)
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "[claude-mnemonic] Warning: context fetch failed: %v\n", err)
|
||||
return "", nil
|
||||
hooks.WriteResponse("SessionStart", true)
|
||||
return
|
||||
}
|
||||
|
||||
// Parse observations from response
|
||||
obsData, ok := result["observations"].([]interface{})
|
||||
if !ok || len(obsData) == 0 {
|
||||
// No observations - just continue normally
|
||||
return "", nil
|
||||
hooks.WriteResponse("SessionStart", true)
|
||||
return
|
||||
}
|
||||
|
||||
// Get full_count from response (how many observations get full detail)
|
||||
@@ -104,7 +136,17 @@ func handleSessionStart(ctx *hooks.HookContext, input *Input) (string, error) {
|
||||
}
|
||||
|
||||
contextBuilder += "</claude-mnemonic-context>\n"
|
||||
return contextBuilder, nil
|
||||
|
||||
// Output context as JSON with additionalContext field
|
||||
response := map[string]interface{}{
|
||||
"continue": true,
|
||||
"hookSpecificOutput": map[string]interface{}{
|
||||
"hookEventName": "SessionStart",
|
||||
"additionalContext": contextBuilder,
|
||||
},
|
||||
}
|
||||
_ = json.NewEncoder(os.Stdout).Encode(response)
|
||||
os.Exit(0)
|
||||
}
|
||||
|
||||
func getString(m map[string]interface{}, key string) string {
|
||||
|
||||
@@ -5,6 +5,7 @@ package main
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"os"
|
||||
@@ -43,21 +44,21 @@ type StatusInput struct {
|
||||
|
||||
// WorkerStats is the response from the worker's /api/stats endpoint.
|
||||
type WorkerStats struct {
|
||||
Uptime string `json:"uptime"`
|
||||
Project string `json:"project,omitempty"`
|
||||
Retrieval 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 {
|
||||
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
|
||||
@@ -71,13 +72,18 @@ const (
|
||||
)
|
||||
|
||||
func main() {
|
||||
hooks.RunStatuslineHook(handleStatusline)
|
||||
}
|
||||
// Read input from stdin
|
||||
inputData, err := io.ReadAll(os.Stdin)
|
||||
if err != nil {
|
||||
// On error, output minimal status
|
||||
fmt.Println(formatOffline())
|
||||
return
|
||||
}
|
||||
|
||||
func handleStatusline(input *StatusInput, port int) string {
|
||||
// Handle error cases (nil input)
|
||||
if input == nil {
|
||||
return formatOffline()
|
||||
var input StatusInput
|
||||
if err := json.Unmarshal(inputData, &input); err != nil {
|
||||
fmt.Println(formatOffline())
|
||||
return
|
||||
}
|
||||
|
||||
// Determine project directory
|
||||
@@ -96,14 +102,16 @@ func handleStatusline(input *StatusInput, port int) string {
|
||||
}
|
||||
|
||||
// Get worker stats
|
||||
stats := getWorkerStats(port, project)
|
||||
stats := getWorkerStats(project)
|
||||
|
||||
// Format and return statusline
|
||||
return formatStatusLine(stats, *input)
|
||||
// Format and output statusline
|
||||
fmt.Println(formatStatusLine(stats, input))
|
||||
}
|
||||
|
||||
// getWorkerStats fetches stats from the worker service.
|
||||
func getWorkerStats(port int, project string) *WorkerStats {
|
||||
func getWorkerStats(project string) *WorkerStats {
|
||||
port := hooks.GetWorkerPort()
|
||||
|
||||
// Build URL with optional project parameter
|
||||
endpoint := fmt.Sprintf("http://127.0.0.1:%d/api/stats", port)
|
||||
if project != "" {
|
||||
@@ -179,45 +187,54 @@ func formatDefault(stats *WorkerStats, useColors bool) string {
|
||||
}
|
||||
|
||||
// Build status parts with clear labels
|
||||
parts := []string{
|
||||
prefix,
|
||||
indicator,
|
||||
}
|
||||
parts := []string{}
|
||||
|
||||
// Add retrieval stats if available
|
||||
if stats.Retrieval.ObservationsServed > 0 {
|
||||
parts = append(parts, fmt.Sprintf("served:%d", stats.Retrieval.ObservationsServed))
|
||||
}
|
||||
// Total memories served to Claude this session
|
||||
parts = append(parts, fmt.Sprintf("served:%d", stats.Retrieval.ObservationsServed))
|
||||
|
||||
// Context injections (memories auto-loaded at session start)
|
||||
if stats.Retrieval.ContextInjections > 0 {
|
||||
parts = append(parts, fmt.Sprintf("injected:%d", stats.Retrieval.ContextInjections))
|
||||
}
|
||||
|
||||
// Semantic searches performed
|
||||
if stats.Retrieval.SearchRequests > 0 {
|
||||
parts = append(parts, fmt.Sprintf("searches:%d", stats.Retrieval.SearchRequests))
|
||||
}
|
||||
|
||||
// Add project-specific observation count if available
|
||||
// Project-specific memory count
|
||||
if stats.ProjectObservations > 0 {
|
||||
parts = append(parts, fmt.Sprintf("project:%d memories", stats.ProjectObservations))
|
||||
if useColors {
|
||||
parts = append(parts, fmt.Sprintf("%sproject:%d memories%s", colorYellow, stats.ProjectObservations, reset))
|
||||
} else {
|
||||
parts = append(parts, fmt.Sprintf("project:%d memories", stats.ProjectObservations))
|
||||
}
|
||||
}
|
||||
|
||||
// Join with separators
|
||||
result := parts[0] + " " + parts[1]
|
||||
if len(parts) > 2 {
|
||||
for i := 2; i < len(parts); i++ {
|
||||
if useColors {
|
||||
result += colorGray + " | " + reset + parts[i]
|
||||
} else {
|
||||
result += " | " + parts[i]
|
||||
}
|
||||
// Processing indicator
|
||||
if stats.IsProcessing || stats.QueueDepth > 0 {
|
||||
if useColors {
|
||||
parts = append(parts, colorYellow+"processing..."+colorReset)
|
||||
} else {
|
||||
parts = append(parts, "processing...")
|
||||
}
|
||||
}
|
||||
|
||||
result := prefix + " " + indicator
|
||||
for i, part := range parts {
|
||||
if i == 0 {
|
||||
result += " " + part
|
||||
} else {
|
||||
result += " | " + part
|
||||
}
|
||||
}
|
||||
|
||||
return result
|
||||
}
|
||||
|
||||
// formatCompact returns a compact status line format.
|
||||
// formatCompact returns a compact status line.
|
||||
func formatCompact(stats *WorkerStats, useColors bool) string {
|
||||
// [m] ● 42/5/3
|
||||
// [m] ● 42/5/3 (28)
|
||||
var prefix, indicator string
|
||||
if useColors {
|
||||
prefix = colorCyan + "[m]" + colorReset
|
||||
@@ -227,16 +244,31 @@ func formatCompact(stats *WorkerStats, useColors bool) string {
|
||||
indicator = "●"
|
||||
}
|
||||
|
||||
return fmt.Sprintf("%s %s %d/%d/%d",
|
||||
result := fmt.Sprintf("%s %s %d/%d/%d",
|
||||
prefix, indicator,
|
||||
stats.Retrieval.ObservationsServed,
|
||||
stats.Retrieval.ContextInjections,
|
||||
stats.Retrieval.SearchRequests)
|
||||
stats.Retrieval.SearchRequests,
|
||||
)
|
||||
|
||||
if stats.ProjectObservations > 0 {
|
||||
result += fmt.Sprintf(" (%d)", stats.ProjectObservations)
|
||||
}
|
||||
|
||||
if stats.IsProcessing || stats.QueueDepth > 0 {
|
||||
if useColors {
|
||||
result += " " + colorYellow + "⚙" + colorReset
|
||||
} else {
|
||||
result += " ⚙"
|
||||
}
|
||||
}
|
||||
|
||||
return result
|
||||
}
|
||||
|
||||
// formatMinimal returns a minimal status line format.
|
||||
// formatMinimal returns a minimal status line.
|
||||
func formatMinimal(stats *WorkerStats, useColors bool) string {
|
||||
// ● 28 memories
|
||||
// ● 42 obs
|
||||
var indicator string
|
||||
if useColors {
|
||||
indicator = colorGreen + "●" + colorReset
|
||||
@@ -244,31 +276,32 @@ func formatMinimal(stats *WorkerStats, useColors bool) string {
|
||||
indicator = "●"
|
||||
}
|
||||
|
||||
result := fmt.Sprintf("%s %d", indicator, stats.Retrieval.ObservationsServed)
|
||||
|
||||
if stats.ProjectObservations > 0 {
|
||||
return fmt.Sprintf("%s %d memories", indicator, stats.ProjectObservations)
|
||||
result += fmt.Sprintf("/%d", stats.ProjectObservations)
|
||||
}
|
||||
|
||||
return fmt.Sprintf("%s mnemonic ready", indicator)
|
||||
return result
|
||||
}
|
||||
|
||||
// formatOffline returns status for when worker is offline.
|
||||
// formatOffline returns the offline status.
|
||||
func formatOffline() string {
|
||||
useColors := os.Getenv("NO_COLOR") == "" && os.Getenv("TERM") != "dumb"
|
||||
return formatOfflineColored(useColors)
|
||||
return formatOfflineColored(true)
|
||||
}
|
||||
|
||||
// formatOfflineColored returns colored offline status.
|
||||
// formatOfflineColored returns the offline status with optional colors.
|
||||
func formatOfflineColored(useColors bool) string {
|
||||
if useColors {
|
||||
return colorGray + "[mnemonic]" + colorReset + " " + colorGray + "○" + colorReset + " offline"
|
||||
return colorCyan + "[mnemonic]" + colorReset + " " + colorGray + "○" + colorReset
|
||||
}
|
||||
return "[mnemonic] ○ offline"
|
||||
return "[mnemonic] ○"
|
||||
}
|
||||
|
||||
// formatStartingColored returns colored starting status.
|
||||
// formatStartingColored returns the starting status with optional colors.
|
||||
func formatStartingColored(useColors bool) string {
|
||||
if useColors {
|
||||
return colorYellow + "[mnemonic]" + colorReset + " " + colorYellow + "○" + colorReset + " starting..."
|
||||
return colorCyan + "[mnemonic]" + colorReset + " " + colorYellow + "◐" + colorReset + " starting"
|
||||
}
|
||||
return "[mnemonic] ○ starting..."
|
||||
return "[mnemonic] ◐ starting"
|
||||
}
|
||||
|
||||
@@ -14,17 +14,17 @@ import (
|
||||
// Input is the hook input from Claude Code.
|
||||
type Input struct {
|
||||
hooks.BaseInput
|
||||
TranscriptPath string `json:"transcript_path"`
|
||||
StopHookActive bool `json:"stop_hook_active"`
|
||||
TranscriptPath string `json:"transcript_path"`
|
||||
}
|
||||
|
||||
// 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).
|
||||
|
||||
@@ -2,7 +2,9 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/url"
|
||||
"os"
|
||||
|
||||
@@ -11,28 +13,60 @@ import (
|
||||
|
||||
// Input is the hook input from Claude Code.
|
||||
type Input struct {
|
||||
hooks.BaseInput
|
||||
Prompt string `json:"prompt"`
|
||||
SessionID string `json:"session_id"`
|
||||
CWD string `json:"cwd"`
|
||||
PermissionMode string `json:"permission_mode"`
|
||||
HookEventName string `json:"hook_event_name"`
|
||||
Prompt string `json:"prompt"`
|
||||
}
|
||||
|
||||
func main() {
|
||||
hooks.RunHook("UserPromptSubmit", handleUserPrompt)
|
||||
}
|
||||
// Skip if this is an internal call (from SDK processor)
|
||||
if os.Getenv("CLAUDE_MNEMONIC_INTERNAL") == "1" {
|
||||
hooks.WriteResponse("UserPromptSubmit", true)
|
||||
return
|
||||
}
|
||||
|
||||
// Read input from stdin
|
||||
inputData, err := io.ReadAll(os.Stdin)
|
||||
if err != nil {
|
||||
hooks.WriteError("UserPromptSubmit", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
|
||||
var input Input
|
||||
if err := json.Unmarshal(inputData, &input); err != nil {
|
||||
hooks.WriteError("UserPromptSubmit", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
|
||||
// Ensure worker is running
|
||||
port, err := hooks.EnsureWorkerRunning()
|
||||
if err != nil {
|
||||
hooks.WriteError("UserPromptSubmit", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
|
||||
// Generate unique project ID from CWD
|
||||
project := hooks.ProjectIDWithName(input.CWD)
|
||||
|
||||
func handleUserPrompt(ctx *hooks.HookContext, input *Input) (string, error) {
|
||||
// Search for relevant observations based on the prompt
|
||||
searchURL := fmt.Sprintf("/api/context/search?project=%s&query=%s&cwd=%s",
|
||||
url.QueryEscape(ctx.Project),
|
||||
url.QueryEscape(project),
|
||||
url.QueryEscape(input.Prompt),
|
||||
url.QueryEscape(ctx.CWD))
|
||||
url.QueryEscape(input.CWD))
|
||||
|
||||
var contextToInject string
|
||||
var observationCount int
|
||||
|
||||
searchResult, _ := hooks.GET(ctx.Port, searchURL)
|
||||
searchResult, _ := hooks.GET(port, searchURL)
|
||||
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)
|
||||
// Limit to top 5 most relevant observations
|
||||
maxObs := 5
|
||||
if len(observations) < maxObs {
|
||||
maxObs = len(observations)
|
||||
}
|
||||
observations = observations[:maxObs]
|
||||
observationCount = len(observations)
|
||||
|
||||
// Build context from search results
|
||||
@@ -74,56 +108,59 @@ func handleUserPrompt(ctx *hooks.HookContext, input *Input) (string, error) {
|
||||
}
|
||||
|
||||
contextBuilder += "</relevant-memory>\n"
|
||||
|
||||
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,
|
||||
result, err := hooks.POST(port, "/api/sessions/init", map[string]interface{}{
|
||||
"claudeSessionId": input.SessionID,
|
||||
"project": project,
|
||||
"prompt": input.Prompt,
|
||||
"matchedObservations": observationCount,
|
||||
})
|
||||
if err != nil {
|
||||
return "", err
|
||||
hooks.WriteError("UserPromptSubmit", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
|
||||
// Check if skipped due to privacy
|
||||
if skipped, ok := result["skipped"].(bool); ok && skipped {
|
||||
fmt.Fprintf(os.Stderr, "[user-prompt] Session skipped (private)\n")
|
||||
return "", nil
|
||||
hooks.WriteResponse("UserPromptSubmit", true)
|
||||
return
|
||||
}
|
||||
|
||||
// Safely extract session ID and prompt number with type checking
|
||||
sessionDbIdRaw, ok := result["sessionDbId"].(float64)
|
||||
if !ok {
|
||||
return "", fmt.Errorf("invalid or missing sessionDbId in response")
|
||||
}
|
||||
sessionID := int64(sessionDbIdRaw)
|
||||
|
||||
promptNumberRaw, ok := result["promptNumber"].(float64)
|
||||
if !ok {
|
||||
return "", fmt.Errorf("invalid or missing promptNumber in response")
|
||||
}
|
||||
promptNumber := int(promptNumberRaw)
|
||||
sessionID := int64(result["sessionDbId"].(float64))
|
||||
promptNumber := int(result["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{}{
|
||||
_, err = hooks.POST(port, fmt.Sprintf("/sessions/%d/init", sessionID), map[string]interface{}{
|
||||
"userPrompt": input.Prompt,
|
||||
"promptNumber": promptNumber,
|
||||
})
|
||||
if err != nil {
|
||||
return "", err
|
||||
hooks.WriteError("UserPromptSubmit", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
|
||||
// Return context if we found relevant observations
|
||||
// Output results - stdout with exit 0 adds context to Claude's prompt
|
||||
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
|
||||
// Output context as JSON with additionalContext field
|
||||
response := map[string]interface{}{
|
||||
"continue": true,
|
||||
"hookSpecificOutput": map[string]interface{}{
|
||||
"hookEventName": "UserPromptSubmit",
|
||||
"additionalContext": contextToInject,
|
||||
},
|
||||
}
|
||||
_ = json.NewEncoder(os.Stdout).Encode(response)
|
||||
os.Exit(0)
|
||||
} else {
|
||||
hooks.WriteResponse("UserPromptSubmit", true)
|
||||
}
|
||||
|
||||
return "", nil
|
||||
}
|
||||
|
||||
+12
-36
@@ -10,14 +10,12 @@ import (
|
||||
"time"
|
||||
|
||||
"github.com/lukaszraczylo/claude-mnemonic/internal/config"
|
||||
"github.com/lukaszraczylo/claude-mnemonic/internal/db/gorm"
|
||||
"github.com/lukaszraczylo/claude-mnemonic/internal/db/sqlite"
|
||||
"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"
|
||||
)
|
||||
@@ -73,25 +71,22 @@ func main() {
|
||||
cancel()
|
||||
}()
|
||||
|
||||
// Initialize database store (migrations run automatically)
|
||||
storeCfg := gorm.Config{
|
||||
// Initialize SQLite store (migrations run automatically)
|
||||
storeCfg := sqlite.StoreConfig{
|
||||
Path: dbPath,
|
||||
MaxConns: cfg.MaxConns,
|
||||
// WALMode is enabled automatically by GORM
|
||||
WALMode: true,
|
||||
}
|
||||
store, err := gorm.NewStore(storeCfg)
|
||||
store, err := sqlite.NewStore(storeCfg)
|
||||
if err != nil {
|
||||
log.Fatal().Err(err).Msg("Failed to initialize database store")
|
||||
log.Fatal().Err(err).Msg("Failed to initialize SQLite store")
|
||||
}
|
||||
defer store.Close()
|
||||
|
||||
// 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)
|
||||
observationStore := sqlite.NewObservationStore(store)
|
||||
summaryStore := sqlite.NewSummaryStore(store)
|
||||
promptStore := sqlite.NewPromptStore(store)
|
||||
|
||||
// Initialize embedding service and vector client
|
||||
var vectorClient *sqlitevec.Client
|
||||
@@ -100,7 +95,7 @@ func main() {
|
||||
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)
|
||||
vectorClient, err = sqlitevec.NewClient(sqlitevec.Config{DB: store.DB()}, embedSvc)
|
||||
if err != nil {
|
||||
log.Warn().Err(err).Msg("Vector client unavailable, vector search disabled")
|
||||
} else {
|
||||
@@ -108,33 +103,14 @@ func main() {
|
||||
}
|
||||
}
|
||||
|
||||
// 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
|
||||
// Note: maintenanceService is nil because it runs in the worker process
|
||||
server := mcp.NewServer(
|
||||
searchMgr,
|
||||
Version,
|
||||
observationStore,
|
||||
patternStore,
|
||||
relationStore,
|
||||
sessionStore,
|
||||
vectorClient,
|
||||
scoreCalculator,
|
||||
recalculator,
|
||||
nil, // maintenanceService - handled by worker
|
||||
)
|
||||
// Create and run MCP server
|
||||
server := mcp.NewServer(searchMgr, Version)
|
||||
log.Info().Str("project", *project).Str("version", Version).Msg("Starting MCP server")
|
||||
|
||||
if err := server.Run(ctx); err != nil {
|
||||
|
||||
Binary file not shown.
|
Before Width: | Height: | Size: 252 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 123 KiB |
+11
-58
@@ -29,21 +29,6 @@
|
||||
subtitle="Stop re-explaining your codebase. Claude Mnemonic captures bug fixes, architecture decisions, and coding patterns - then brings them back exactly when you need them."
|
||||
/>
|
||||
|
||||
<!-- Dashboard Preview -->
|
||||
<section class="py-12 lg:py-16 px-4 sm:px-6 relative">
|
||||
<div class="max-w-6xl mx-auto">
|
||||
<div class="relative rounded-xl overflow-hidden border border-slate-700/50 shadow-2xl shadow-amber-500/5">
|
||||
<div class="absolute inset-0 bg-gradient-to-t from-slate-950 via-transparent to-transparent pointer-events-none z-10"></div>
|
||||
<img
|
||||
src="/claude-mnemonic.jpg"
|
||||
alt="Claude Mnemonic Dashboard"
|
||||
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. View graph stats, vector metrics, storage savings, and performance analytics.</p>
|
||||
</div>
|
||||
</section>
|
||||
|
||||
<!-- Problem Section -->
|
||||
<section class="py-20 lg:py-28 px-4 sm:px-6 relative">
|
||||
<div class="max-w-6xl mx-auto grid lg:grid-cols-2 gap-8 lg:gap-16 items-center">
|
||||
@@ -100,21 +85,6 @@
|
||||
</div>
|
||||
</section>
|
||||
|
||||
<!-- Knowledge Graph Preview -->
|
||||
<section class="py-16 lg:py-20 px-4 sm:px-6 relative">
|
||||
<div class="max-w-5xl mx-auto">
|
||||
<SectionHeader title="See how knowledge connects" subtitle="The knowledge graph reveals relationships between your memories automatically" />
|
||||
<div class="relative rounded-xl overflow-hidden border border-slate-700/50 shadow-2xl shadow-purple-500/5">
|
||||
<img
|
||||
src="/observation-relation-graph.jpg"
|
||||
alt="Knowledge Graph Visualization"
|
||||
class="w-full h-auto"
|
||||
/>
|
||||
</div>
|
||||
<p class="text-center text-slate-500 text-sm mt-4">Click any observation to explore its relationships - see what causes what, what fixes what, and how concepts evolve</p>
|
||||
</div>
|
||||
</section>
|
||||
|
||||
<!-- Before/After Section -->
|
||||
<section class="py-20 lg:py-28 px-4 sm:px-6">
|
||||
<div class="max-w-6xl mx-auto">
|
||||
@@ -304,7 +274,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-3 gap-4 sm:gap-6 text-center">
|
||||
<div class="grid sm:grid-cols-2 lg:grid-cols-4 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,19 +285,11 @@
|
||||
</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">Hybrid vector storage with LEANN-inspired selective embeddings. 60-80% storage reduction.</p>
|
||||
<p class="text-slate-400 text-xs sm:text-sm">Embedded vector database. No external services required.</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 class="text-3xl sm:text-4xl font-bold text-amber-500 mb-2">MiniLM</div>
|
||||
<p class="text-slate-400 text-xs sm:text-sm">Local embeddings via ONNX. "Fix auth" finds "JWT issue" automatically.</p>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
@@ -424,13 +386,10 @@ 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: '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-globe', title: 'Shared best practices', description: 'Security patterns, performance tips, and universal learnings automatically available across all your projects.' },
|
||||
{ icon: 'fas fa-search', title: 'Finds what matters', description: 'Semantic search finds relevant memories even when you don\'t remember the exact words. "That auth thing" just works.' },
|
||||
{ icon: 'fas fa-chart-line', title: 'Live statusline', description: 'Real-time metrics right in Claude Code: memories served, searches performed, and project memory count at a glance.' },
|
||||
{ icon: 'fas fa-lock', title: '100% private', description: 'Your code context never leaves your machine. No telemetry. No cloud sync. Your memories are yours.' },
|
||||
]
|
||||
|
||||
@@ -456,12 +415,8 @@ const installCommands = {
|
||||
const configOptions = [
|
||||
{ name: 'CLAUDE_MNEMONIC_WORKER_PORT', description: 'HTTP port for the worker service (default: 37777)', icon: 'fas fa-network-wired' },
|
||||
{ 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' },
|
||||
{ name: 'CLAUDE_MNEMONIC_CONTEXT_FULL_COUNT', description: 'Observations with full narrative detail, rest are condensed (default: 25)', icon: 'fas fa-expand' },
|
||||
{ name: 'CLAUDE_MNEMONIC_MODEL', description: 'Model for processing observations (default: haiku)', icon: 'fas fa-microchip' },
|
||||
]
|
||||
|
||||
const requiredDeps = [
|
||||
@@ -472,11 +427,9 @@ 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: '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: '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: '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). 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.' },
|
||||
{ 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.' },
|
||||
]
|
||||
</script>
|
||||
|
||||
@@ -7,25 +7,18 @@ 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.5
|
||||
github.com/go-gormigrate/gormigrate/v2 v2.1.5
|
||||
github.com/go-chi/chi/v5 v5.2.3
|
||||
github.com/goccy/go-json v0.10.5
|
||||
github.com/mattn/go-sqlite3 v1.14.33
|
||||
github.com/mattn/go-sqlite3 v1.14.32
|
||||
github.com/rs/zerolog v1.34.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
|
||||
golang.org/x/sync v0.19.0
|
||||
gorm.io/driver/sqlite v1.6.0
|
||||
gorm.io/gorm v1.31.1
|
||||
)
|
||||
|
||||
require (
|
||||
github.com/davecgh/go-spew v1.1.1 // indirect
|
||||
github.com/emirpasic/gods v1.18.1 // indirect
|
||||
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/mitchellh/colorstring v0.0.0-20190213212951-d06e56a500db // indirect
|
||||
@@ -34,7 +27,7 @@ require (
|
||||
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.41.0 // indirect
|
||||
golang.org/x/text v0.33.0 // indirect
|
||||
golang.org/x/sys v0.39.0 // indirect
|
||||
golang.org/x/text v0.32.0 // indirect
|
||||
gopkg.in/yaml.v3 v3.0.1 // indirect
|
||||
)
|
||||
|
||||
@@ -10,17 +10,11 @@ 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.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/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/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/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=
|
||||
@@ -28,8 +22,8 @@ github.com/mattn/go-isatty v0.0.16/go.mod h1:kYGgaQfpe5nmfYZH+SKPsOc2e4SrIfOl2e/
|
||||
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-sqlite3 v1.14.32 h1:JD12Ag3oLy1zQA+BNn74xRgaBbdhbNIDYvQUEuuErjs=
|
||||
github.com/mattn/go-sqlite3 v1.14.32/go.mod h1:Uh1q+B4BYcTPb+yiD3kU8Ct7aC0hY9fxUwlHK0RXw+Y=
|
||||
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=
|
||||
@@ -44,8 +38,6 @@ github.com/rs/zerolog v1.34.0 h1:k43nTLIwcTVQAncfCw4KZ2VY6ukYoZaBPNOE8txlOeY=
|
||||
github.com/rs/zerolog v1.34.0/go.mod h1:bJsvje4Z08ROH4Nhs5iH600c3IkWhwp44iRc54W6wYQ=
|
||||
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=
|
||||
@@ -54,20 +46,14 @@ github.com/sugarme/regexpset v0.0.0-20200920021344-4d4ec8eaf93c h1:pwb4kNSHb4K89
|
||||
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/sync v0.19.0 h1:vV+1eWNmZ5geRlYjzm2adRgW2/mcpevXNg50YZtPCE4=
|
||||
golang.org/x/sync v0.19.0/go.mod h1:9KTHXmSnoGruLpwFjVSX0lNNA75CykiMECbovNTZqGI=
|
||||
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.41.0 h1:Ivj+2Cp/ylzLiEU89QhWblYnOE9zerudt9Ftecq2C6k=
|
||||
golang.org/x/sys v0.41.0/go.mod h1:OgkHotnGiDImocRcuBABYBEXf8A9a87e/uXjp9XT3ks=
|
||||
golang.org/x/text v0.33.0 h1:B3njUFyqtHDUI5jMn1YIr5B0IE2U0qck04r6d4KPAxE=
|
||||
golang.org/x/text v0.33.0/go.mod h1:LuMebE6+rBincTi9+xWTY8TztLzKHc/9C1uBCG27+q8=
|
||||
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=
|
||||
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=
|
||||
gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
|
||||
gorm.io/driver/sqlite v1.6.0 h1:WHRRrIiulaPiPFmDcod6prc4l2VGVWHz80KspNsxSfQ=
|
||||
gorm.io/driver/sqlite v1.6.0/go.mod h1:AO9V1qIQddBESngQUKWL9yoH93HIeA1X6V633rBwyT8=
|
||||
gorm.io/gorm v1.31.1 h1:7CA8FTFz/gRfgqgpeKIBcervUn3xSyPUmr6B2WXJ7kg=
|
||||
gorm.io/gorm v1.31.1/go.mod h1:XyQVbO2k6YkOis7C2437jSit3SsDK72s7n7rsSHd+Gs=
|
||||
|
||||
@@ -1,285 +0,0 @@
|
||||
// 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")
|
||||
}
|
||||
@@ -1,214 +0,0 @@
|
||||
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)
|
||||
}
|
||||
}
|
||||
@@ -1,106 +0,0 @@
|
||||
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
|
||||
}
|
||||
@@ -1,162 +0,0 @@
|
||||
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)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,291 +0,0 @@
|
||||
// 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")
|
||||
}
|
||||
@@ -1,298 +0,0 @@
|
||||
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"])
|
||||
}
|
||||
@@ -1,140 +0,0 @@
|
||||
// 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
|
||||
}
|
||||
}
|
||||
@@ -1,213 +0,0 @@
|
||||
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)
|
||||
}
|
||||
@@ -1,403 +0,0 @@
|
||||
// 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")
|
||||
}
|
||||
@@ -1,398 +0,0 @@
|
||||
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")
|
||||
}
|
||||
+35
-123
@@ -36,42 +36,28 @@ var CriticalConcepts = []string{
|
||||
}
|
||||
|
||||
// Config holds the application configuration.
|
||||
// Field order optimized for memory alignment (fieldalignment).
|
||||
type Config struct {
|
||||
ContextFullField string `json:"context_full_field"`
|
||||
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"`
|
||||
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"`
|
||||
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"`
|
||||
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"`
|
||||
MaintenanceEnabled bool `json:"maintenance_enabled"`
|
||||
RerankingEnabled bool `json:"reranking_enabled"`
|
||||
ContextShowLastSummary bool `json:"context_show_last_summary"`
|
||||
CleanupStaleObservations bool `json:"cleanup_stale_observations"`
|
||||
// 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"`
|
||||
|
||||
// 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"`
|
||||
ContextObsConcepts []string `json:"context_obs_concepts"`
|
||||
}
|
||||
|
||||
var (
|
||||
@@ -97,9 +83,8 @@ 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(), 0700)
|
||||
return os.MkdirAll(DataDir(), 0750)
|
||||
}
|
||||
|
||||
// EnsureSettings creates a default settings file if it doesn't exist.
|
||||
@@ -134,44 +119,22 @@ func EnsureAll() error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// DefaultEmbeddingModel is the default embedding model to use.
|
||||
const DefaultEmbeddingModel = "bge-v1.5"
|
||||
|
||||
// Default returns a Config with default values.
|
||||
func Default() *Config {
|
||||
return &Config{
|
||||
WorkerPort: DefaultWorkerPort,
|
||||
DBPath: DBPath(),
|
||||
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
|
||||
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,
|
||||
ContextShowReadTokens: true,
|
||||
ContextShowWorkTokens: true,
|
||||
ContextFullField: "narrative",
|
||||
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)
|
||||
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
|
||||
WorkerPort: DefaultWorkerPort,
|
||||
DBPath: DBPath(),
|
||||
MaxConns: 4,
|
||||
Model: DefaultModel,
|
||||
ContextObservations: 100,
|
||||
ContextFullCount: 25,
|
||||
ContextSessionCount: 10,
|
||||
ContextShowReadTokens: true,
|
||||
ContextShowWorkTokens: true,
|
||||
ContextFullField: "narrative",
|
||||
ContextShowLastSummary: true,
|
||||
ContextObsTypes: DefaultObservationTypes,
|
||||
ContextObsConcepts: DefaultObservationConcepts,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -203,28 +166,6 @@ func Load() (*Config, error) {
|
||||
if v, ok := settings["CLAUDE_CODE_PATH"].(string); ok {
|
||||
cfg.ClaudeCodePath = v
|
||||
}
|
||||
if v, ok := settings["CLAUDE_MNEMONIC_EMBEDDING_MODEL"].(string); ok && v != "" {
|
||||
cfg.EmbeddingModel = v
|
||||
}
|
||||
// Reranking settings
|
||||
if v, ok := settings["CLAUDE_MNEMONIC_RERANKING_ENABLED"].(bool); ok {
|
||||
cfg.RerankingEnabled = v
|
||||
}
|
||||
if v, ok := settings["CLAUDE_MNEMONIC_RERANKING_CANDIDATES"].(float64); ok && v > 0 {
|
||||
cfg.RerankingCandidates = int(v)
|
||||
}
|
||||
if v, ok := settings["CLAUDE_MNEMONIC_RERANKING_RESULTS"].(float64); ok && v > 0 {
|
||||
cfg.RerankingResults = int(v)
|
||||
}
|
||||
if v, ok := settings["CLAUDE_MNEMONIC_RERANKING_ALPHA"].(float64); ok && v >= 0 && v <= 1 {
|
||||
cfg.RerankingAlpha = v
|
||||
}
|
||||
if v, ok := settings["CLAUDE_MNEMONIC_RERANKING_MIN_IMPROVEMENT"].(float64); ok && v >= 0 {
|
||||
cfg.RerankingMinImprovement = v
|
||||
}
|
||||
if v, ok := settings["CLAUDE_MNEMONIC_RERANKING_PURE_MODE"].(bool); ok {
|
||||
cfg.RerankingPureMode = v
|
||||
}
|
||||
if v, ok := settings["CLAUDE_MNEMONIC_CONTEXT_OBSERVATIONS"].(float64); ok {
|
||||
cfg.ContextObservations = int(v)
|
||||
}
|
||||
@@ -240,35 +181,6 @@ func Load() (*Config, error) {
|
||||
if v, ok := settings["CLAUDE_MNEMONIC_CONTEXT_OBS_CONCEPTS"].(string); ok && v != "" {
|
||||
cfg.ContextObsConcepts = splitTrim(v)
|
||||
}
|
||||
if v, ok := settings["CLAUDE_MNEMONIC_CONTEXT_RELEVANCE_THRESHOLD"].(float64); ok && v >= 0 && v <= 1 {
|
||||
cfg.ContextRelevanceThreshold = v
|
||||
}
|
||||
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)
|
||||
}
|
||||
|
||||
return cfg, nil
|
||||
}
|
||||
|
||||
@@ -121,8 +121,8 @@ func (s *ConfigSuite) TestLoad_TableDriven() {
|
||||
tests := []struct {
|
||||
name string
|
||||
settingsJSON string
|
||||
expectedModel string
|
||||
expectedPort int
|
||||
expectedModel string
|
||||
expectedObsObs int
|
||||
}{
|
||||
{
|
||||
@@ -183,12 +183,12 @@ func (s *ConfigSuite) TestLoad_TableDriven() {
|
||||
s.Require().NoError(err)
|
||||
|
||||
if tt.settingsJSON != "" {
|
||||
writeErr := os.WriteFile(
|
||||
err := os.WriteFile(
|
||||
filepath.Join(tempDir, ".claude-mnemonic", "settings.json"),
|
||||
[]byte(tt.settingsJSON),
|
||||
0600,
|
||||
)
|
||||
s.Require().NoError(writeErr)
|
||||
s.Require().NoError(err)
|
||||
}
|
||||
|
||||
cfg, err := Load()
|
||||
|
||||
@@ -1,331 +0,0 @@
|
||||
//go:build fts5
|
||||
|
||||
// Package gorm provides GORM-based database operations for claude-mnemonic.
|
||||
package gorm
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
"gorm.io/gorm/logger"
|
||||
|
||||
"github.com/lukaszraczylo/claude-mnemonic/pkg/models"
|
||||
)
|
||||
|
||||
// setupBenchStore creates a temporary store for benchmarking.
|
||||
func setupBenchStore(b *testing.B) (*Store, func()) {
|
||||
b.Helper()
|
||||
|
||||
tmpDir, err := os.MkdirTemp("", "gorm_bench_*")
|
||||
if err != nil {
|
||||
b.Fatalf("create temp dir: %v", err)
|
||||
}
|
||||
|
||||
dbPath := filepath.Join(tmpDir, "bench.db")
|
||||
cfg := Config{
|
||||
Path: dbPath,
|
||||
MaxConns: 4,
|
||||
LogLevel: logger.Silent,
|
||||
}
|
||||
|
||||
store, err := NewStore(cfg)
|
||||
if err != nil {
|
||||
os.RemoveAll(tmpDir)
|
||||
b.Fatalf("NewStore failed: %v", err)
|
||||
}
|
||||
|
||||
cleanup := func() {
|
||||
store.Close()
|
||||
os.RemoveAll(tmpDir)
|
||||
}
|
||||
|
||||
return store, cleanup
|
||||
}
|
||||
|
||||
// BenchmarkSessionStore_CreateSDKSession benchmarks session creation (most frequent operation).
|
||||
func BenchmarkSessionStore_CreateSDKSession(b *testing.B) {
|
||||
store, cleanup := setupBenchStore(b)
|
||||
defer cleanup()
|
||||
|
||||
sessionStore := NewSessionStore(store)
|
||||
ctx := context.Background()
|
||||
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
sessionID := fmt.Sprintf("claude-bench-%d", i)
|
||||
_, err := sessionStore.CreateSDKSession(ctx, sessionID, "bench-project", "test prompt")
|
||||
if err != nil {
|
||||
b.Fatalf("CreateSDKSession failed: %v", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// BenchmarkSessionStore_CreateSDKSession_Idempotent benchmarks idempotent session creation (INSERT OR IGNORE).
|
||||
func BenchmarkSessionStore_CreateSDKSession_Idempotent(b *testing.B) {
|
||||
store, cleanup := setupBenchStore(b)
|
||||
defer cleanup()
|
||||
|
||||
sessionStore := NewSessionStore(store)
|
||||
ctx := context.Background()
|
||||
|
||||
// Pre-create session
|
||||
sessionStore.CreateSDKSession(ctx, "claude-bench", "bench-project", "test prompt")
|
||||
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
_, err := sessionStore.CreateSDKSession(ctx, "claude-bench", "bench-project", "updated prompt")
|
||||
if err != nil {
|
||||
b.Fatalf("CreateSDKSession failed: %v", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// BenchmarkObservationStore_StoreObservation benchmarks observation storage (high frequency).
|
||||
func BenchmarkObservationStore_StoreObservation(b *testing.B) {
|
||||
store, cleanup := setupBenchStore(b)
|
||||
defer cleanup()
|
||||
|
||||
sessionStore := NewSessionStore(store)
|
||||
obsStore := NewObservationStore(store, nil, nil, nil)
|
||||
ctx := context.Background()
|
||||
|
||||
// Create session
|
||||
sessionID, _ := sessionStore.CreateSDKSession(ctx, "claude-bench", "bench-project", "")
|
||||
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
obs := &models.ParsedObservation{
|
||||
Type: models.ObsTypeDiscovery,
|
||||
Title: fmt.Sprintf("Observation %d", i),
|
||||
Narrative: "Benchmark observation content",
|
||||
}
|
||||
_, _, err := obsStore.StoreObservation(ctx, "claude-bench", "bench-project", obs, int(sessionID), int64(i+1))
|
||||
if err != nil {
|
||||
b.Fatalf("StoreObservation failed: %v", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// BenchmarkObservationStore_GetRecentObservations benchmarks recent observation retrieval.
|
||||
func BenchmarkObservationStore_GetRecentObservations(b *testing.B) {
|
||||
store, cleanup := setupBenchStore(b)
|
||||
defer cleanup()
|
||||
|
||||
sessionStore := NewSessionStore(store)
|
||||
obsStore := NewObservationStore(store, nil, nil, nil)
|
||||
ctx := context.Background()
|
||||
|
||||
// Create session and observations
|
||||
sessionID, _ := sessionStore.CreateSDKSession(ctx, "claude-bench", "bench-project", "")
|
||||
for i := 0; i < 100; i++ {
|
||||
obs := &models.ParsedObservation{
|
||||
Type: models.ObsTypeDiscovery,
|
||||
Title: fmt.Sprintf("Observation %d", i),
|
||||
Narrative: "Benchmark observation content",
|
||||
}
|
||||
obsStore.StoreObservation(ctx, "claude-bench", "bench-project", obs, int(sessionID), int64(i+1))
|
||||
}
|
||||
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
_, err := obsStore.GetRecentObservations(ctx, "bench-project", 20)
|
||||
if err != nil {
|
||||
b.Fatalf("GetRecentObservations failed: %v", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// BenchmarkObservationStore_SearchObservationsFTS benchmarks FTS5 search (latency-sensitive).
|
||||
func BenchmarkObservationStore_SearchObservationsFTS(b *testing.B) {
|
||||
store, cleanup := setupBenchStore(b)
|
||||
defer cleanup()
|
||||
|
||||
sessionStore := NewSessionStore(store)
|
||||
obsStore := NewObservationStore(store, nil, nil, nil)
|
||||
ctx := context.Background()
|
||||
|
||||
// Create session and observations with searchable content
|
||||
sessionID, _ := sessionStore.CreateSDKSession(ctx, "claude-bench", "bench-project", "")
|
||||
for i := 0; i < 100; i++ {
|
||||
obs := &models.ParsedObservation{
|
||||
Type: models.ObsTypeDiscovery,
|
||||
Title: fmt.Sprintf("Security best practice %d", i),
|
||||
Narrative: "This observation discusses security patterns and authentication mechanisms",
|
||||
}
|
||||
obsStore.StoreObservation(ctx, "claude-bench", "bench-project", obs, int(sessionID), int64(i+1))
|
||||
}
|
||||
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
_, err := obsStore.SearchObservationsFTS(ctx, "security authentication", "bench-project", 10)
|
||||
if err != nil {
|
||||
b.Fatalf("SearchObservationsFTS failed: %v", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// BenchmarkObservationStore_UpdateImportanceScore benchmarks scoring updates.
|
||||
func BenchmarkObservationStore_UpdateImportanceScore(b *testing.B) {
|
||||
store, cleanup := setupBenchStore(b)
|
||||
defer cleanup()
|
||||
|
||||
sessionStore := NewSessionStore(store)
|
||||
obsStore := NewObservationStore(store, nil, nil, nil)
|
||||
ctx := context.Background()
|
||||
|
||||
// Create session and observation
|
||||
sessionID, _ := sessionStore.CreateSDKSession(ctx, "claude-bench", "bench-project", "")
|
||||
obs := &models.ParsedObservation{Type: models.ObsTypeDiscovery, Title: "Test"}
|
||||
obsID, _, _ := obsStore.StoreObservation(ctx, "claude-bench", "bench-project", obs, int(sessionID), 1)
|
||||
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
score := float64(i%10) + 1.0
|
||||
err := obsStore.UpdateImportanceScore(ctx, obsID, score)
|
||||
if err != nil {
|
||||
b.Fatalf("UpdateImportanceScore failed: %v", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// BenchmarkObservationStore_UpdateImportanceScores_Bulk benchmarks bulk scoring updates.
|
||||
func BenchmarkObservationStore_UpdateImportanceScores_Bulk(b *testing.B) {
|
||||
store, cleanup := setupBenchStore(b)
|
||||
defer cleanup()
|
||||
|
||||
sessionStore := NewSessionStore(store)
|
||||
obsStore := NewObservationStore(store, nil, nil, nil)
|
||||
ctx := context.Background()
|
||||
|
||||
// Create session and 100 observations
|
||||
sessionID, _ := sessionStore.CreateSDKSession(ctx, "claude-bench", "bench-project", "")
|
||||
var obsIDs []int64
|
||||
for i := 0; i < 100; i++ {
|
||||
obs := &models.ParsedObservation{Type: models.ObsTypeDiscovery, Title: fmt.Sprintf("Obs %d", i)}
|
||||
obsID, _, _ := obsStore.StoreObservation(ctx, "claude-bench", "bench-project", obs, int(sessionID), int64(i+1))
|
||||
obsIDs = append(obsIDs, obsID)
|
||||
}
|
||||
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
scores := make(map[int64]float64)
|
||||
for _, id := range obsIDs {
|
||||
scores[id] = float64(i%10) + 1.0
|
||||
}
|
||||
err := obsStore.UpdateImportanceScores(ctx, scores)
|
||||
if err != nil {
|
||||
b.Fatalf("UpdateImportanceScores failed: %v", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// BenchmarkPromptStore_SaveUserPromptWithMatches benchmarks prompt storage with matches.
|
||||
func BenchmarkPromptStore_SaveUserPromptWithMatches(b *testing.B) {
|
||||
store, cleanup := setupBenchStore(b)
|
||||
defer cleanup()
|
||||
|
||||
sessionStore := NewSessionStore(store)
|
||||
promptStore := NewPromptStore(store, nil)
|
||||
ctx := context.Background()
|
||||
|
||||
// Create session
|
||||
sessionID, _ := sessionStore.CreateSDKSession(ctx, "claude-bench", "bench-project", "")
|
||||
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
_, err := promptStore.SaveUserPromptWithMatches(ctx, "claude-bench", int(sessionID), fmt.Sprintf("Prompt %d", i), i+1)
|
||||
if err != nil {
|
||||
b.Fatalf("SaveUserPromptWithMatches failed: %v", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// BenchmarkSummaryStore_StoreSummary benchmarks summary storage.
|
||||
func BenchmarkSummaryStore_StoreSummary(b *testing.B) {
|
||||
store, cleanup := setupBenchStore(b)
|
||||
defer cleanup()
|
||||
|
||||
sessionStore := NewSessionStore(store)
|
||||
summaryStore := NewSummaryStore(store)
|
||||
ctx := context.Background()
|
||||
|
||||
// Create session
|
||||
sessionStore.CreateSDKSession(ctx, "claude-bench", "bench-project", "")
|
||||
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
summary := &models.ParsedSummary{
|
||||
Request: fmt.Sprintf("Request %d", i),
|
||||
Investigated: "Investigation details",
|
||||
Learned: "Learning summary",
|
||||
Completed: "Completion status",
|
||||
}
|
||||
_, _, err := summaryStore.StoreSummary(ctx, "claude-bench", "bench-project", summary, i+1, 100)
|
||||
if err != nil {
|
||||
b.Fatalf("StoreSummary failed: %v", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// BenchmarkRelationStore_StoreRelation benchmarks relation storage.
|
||||
func BenchmarkRelationStore_StoreRelation(b *testing.B) {
|
||||
store, cleanup := setupBenchStore(b)
|
||||
defer cleanup()
|
||||
|
||||
sessionStore := NewSessionStore(store)
|
||||
obsStore := NewObservationStore(store, nil, nil, nil)
|
||||
relationStore := NewRelationStore(store)
|
||||
ctx := context.Background()
|
||||
|
||||
// Create session and observations
|
||||
sessionID, _ := sessionStore.CreateSDKSession(ctx, "claude-bench", "bench-project", "")
|
||||
obs1 := &models.ParsedObservation{Type: models.ObsTypeDiscovery, Title: "Source"}
|
||||
obsID1, _, _ := obsStore.StoreObservation(ctx, "claude-bench", "bench-project", obs1, int(sessionID), 1)
|
||||
obs2 := &models.ParsedObservation{Type: models.ObsTypeDiscovery, Title: "Target"}
|
||||
obsID2, _, _ := obsStore.StoreObservation(ctx, "claude-bench", "bench-project", obs2, int(sessionID), 2)
|
||||
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
relation := &models.ObservationRelation{
|
||||
SourceID: obsID1,
|
||||
TargetID: obsID2,
|
||||
RelationType: models.RelationCauses,
|
||||
Confidence: 0.9,
|
||||
DetectionSource: models.DetectionSourceFileOverlap,
|
||||
}
|
||||
_, err := relationStore.StoreRelation(ctx, relation)
|
||||
if err != nil {
|
||||
b.Fatalf("StoreRelation failed: %v", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// BenchmarkPatternStore_StorePattern benchmarks pattern storage.
|
||||
func BenchmarkPatternStore_StorePattern(b *testing.B) {
|
||||
store, cleanup := setupBenchStore(b)
|
||||
defer cleanup()
|
||||
|
||||
patternStore := NewPatternStore(store)
|
||||
ctx := context.Background()
|
||||
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
pattern := &models.Pattern{
|
||||
Name: fmt.Sprintf("Pattern %d", i),
|
||||
Type: models.PatternTypeBug,
|
||||
Description: sql.NullString{String: "Benchmark pattern", Valid: true},
|
||||
Frequency: 1,
|
||||
Confidence: 0.8,
|
||||
Projects: []string{"bench-project"},
|
||||
Status: models.PatternStatusActive,
|
||||
}
|
||||
_, err := patternStore.StorePattern(ctx, pattern)
|
||||
if err != nil {
|
||||
b.Fatalf("StorePattern failed: %v", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,281 +0,0 @@
|
||||
// Package gorm provides GORM-based database operations for claude-mnemonic.
|
||||
package gorm
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"time"
|
||||
|
||||
"gorm.io/gorm"
|
||||
|
||||
"github.com/lukaszraczylo/claude-mnemonic/pkg/models"
|
||||
)
|
||||
|
||||
// SupersededRetentionDays is the number of days to keep superseded observations before deletion.
|
||||
const SupersededRetentionDays = 3
|
||||
|
||||
// ConflictStore provides conflict-related database operations using GORM.
|
||||
type ConflictStore struct {
|
||||
db *gorm.DB
|
||||
}
|
||||
|
||||
// NewConflictStore creates a new conflict store.
|
||||
func NewConflictStore(store *Store) *ConflictStore {
|
||||
return &ConflictStore{
|
||||
db: store.DB,
|
||||
}
|
||||
}
|
||||
|
||||
// StoreConflict stores a new observation conflict.
|
||||
func (s *ConflictStore) StoreConflict(ctx context.Context, conflict *models.ObservationConflict) (int64, error) {
|
||||
dbConflict := &ObservationConflict{
|
||||
NewerObsID: conflict.NewerObsID,
|
||||
OlderObsID: conflict.OlderObsID,
|
||||
ConflictType: conflict.ConflictType,
|
||||
Resolution: conflict.Resolution,
|
||||
DetectedAt: conflict.DetectedAt,
|
||||
DetectedAtEpoch: conflict.DetectedAtEpoch,
|
||||
Resolved: 0,
|
||||
}
|
||||
|
||||
// Convert bool to int
|
||||
if conflict.Resolved {
|
||||
dbConflict.Resolved = 1
|
||||
}
|
||||
|
||||
// Handle nullable fields
|
||||
if conflict.Reason != "" {
|
||||
dbConflict.Reason = sql.NullString{String: conflict.Reason, Valid: true}
|
||||
}
|
||||
if conflict.ResolvedAt != nil && *conflict.ResolvedAt != "" {
|
||||
dbConflict.ResolvedAt = sql.NullString{String: *conflict.ResolvedAt, Valid: true}
|
||||
}
|
||||
|
||||
result := s.db.WithContext(ctx).Create(dbConflict)
|
||||
if result.Error != nil {
|
||||
return 0, result.Error
|
||||
}
|
||||
|
||||
return dbConflict.ID, nil
|
||||
}
|
||||
|
||||
// MarkObservationSuperseded marks an observation as superseded.
|
||||
func (s *ConflictStore) MarkObservationSuperseded(ctx context.Context, obsID int64) error {
|
||||
result := s.db.WithContext(ctx).
|
||||
Model(&Observation{}).
|
||||
Where("id = ?", obsID).
|
||||
Update("is_superseded", 1)
|
||||
|
||||
return result.Error
|
||||
}
|
||||
|
||||
// MarkObservationsSuperseded marks multiple observations as superseded.
|
||||
func (s *ConflictStore) MarkObservationsSuperseded(ctx context.Context, obsIDs []int64) error {
|
||||
if len(obsIDs) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
result := s.db.WithContext(ctx).
|
||||
Model(&Observation{}).
|
||||
Where("id IN ?", obsIDs).
|
||||
Update("is_superseded", 1)
|
||||
|
||||
return result.Error
|
||||
}
|
||||
|
||||
// GetConflictsByObservationID retrieves all conflicts involving an observation.
|
||||
func (s *ConflictStore) GetConflictsByObservationID(ctx context.Context, obsID int64) ([]*models.ObservationConflict, error) {
|
||||
var conflicts []ObservationConflict
|
||||
|
||||
err := s.db.WithContext(ctx).
|
||||
Where("newer_obs_id = ? OR older_obs_id = ?", obsID, obsID).
|
||||
Order("detected_at_epoch DESC").
|
||||
Find(&conflicts).Error
|
||||
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return toModelConflicts(conflicts), nil
|
||||
}
|
||||
|
||||
// GetUnresolvedConflicts retrieves all unresolved conflicts.
|
||||
func (s *ConflictStore) GetUnresolvedConflicts(ctx context.Context, limit int) ([]*models.ObservationConflict, error) {
|
||||
var conflicts []ObservationConflict
|
||||
|
||||
err := s.db.WithContext(ctx).
|
||||
Where("resolved = 0").
|
||||
Order("detected_at_epoch DESC").
|
||||
Limit(limit).
|
||||
Find(&conflicts).Error
|
||||
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return toModelConflicts(conflicts), nil
|
||||
}
|
||||
|
||||
// GetSupersededObservationIDs returns IDs of all observations that have been superseded.
|
||||
func (s *ConflictStore) GetSupersededObservationIDs(ctx context.Context, project string) ([]int64, error) {
|
||||
var ids []int64
|
||||
|
||||
err := s.db.WithContext(ctx).
|
||||
Table("observation_conflicts oc").
|
||||
Select("DISTINCT oc.older_obs_id").
|
||||
Joins("JOIN observations o ON o.id = oc.older_obs_id").
|
||||
Where("oc.resolution = ?", models.ResolutionPreferNewer).
|
||||
Where("o.project = ? OR o.scope = 'global'", project).
|
||||
Pluck("oc.older_obs_id", &ids).Error
|
||||
|
||||
return ids, err
|
||||
}
|
||||
|
||||
// ResolveConflict marks a conflict as resolved.
|
||||
func (s *ConflictStore) ResolveConflict(ctx context.Context, conflictID int64, resolution models.ConflictResolution) error {
|
||||
now := time.Now().Format(time.RFC3339)
|
||||
|
||||
result := s.db.WithContext(ctx).
|
||||
Model(&ObservationConflict{}).
|
||||
Where("id = ?", conflictID).
|
||||
Updates(map[string]interface{}{
|
||||
"resolved": 1,
|
||||
"resolved_at": now,
|
||||
"resolution": resolution,
|
||||
})
|
||||
|
||||
return result.Error
|
||||
}
|
||||
|
||||
// DeleteConflictsByObservationID deletes all conflicts involving an observation.
|
||||
// Called when an observation is deleted.
|
||||
func (s *ConflictStore) DeleteConflictsByObservationID(ctx context.Context, obsID int64) error {
|
||||
result := s.db.WithContext(ctx).
|
||||
Where("newer_obs_id = ? OR older_obs_id = ?", obsID, obsID).
|
||||
Delete(&ObservationConflict{})
|
||||
|
||||
return result.Error
|
||||
}
|
||||
|
||||
// ConflictWithDetails contains a conflict with its observation details.
|
||||
type ConflictWithDetails struct {
|
||||
Conflict *models.ObservationConflict
|
||||
NewerObsTitle string
|
||||
OlderObsTitle string
|
||||
}
|
||||
|
||||
// CleanupSupersededObservations deletes observations that have been superseded for longer than
|
||||
// SupersededRetentionDays. Returns the IDs of deleted observations for downstream cleanup (e.g., vector DB).
|
||||
func (s *ConflictStore) CleanupSupersededObservations(ctx context.Context, project string) ([]int64, error) {
|
||||
// Calculate cutoff time (3 days ago in milliseconds)
|
||||
cutoffEpoch := time.Now().AddDate(0, 0, -SupersededRetentionDays).UnixMilli()
|
||||
|
||||
var toDelete []int64
|
||||
|
||||
// Use a transaction to prevent TOCTOU race condition
|
||||
err := s.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
// Find IDs to delete
|
||||
err := tx.Table("observations o").
|
||||
Select("DISTINCT o.id").
|
||||
Joins("JOIN observation_conflicts oc ON o.id = oc.older_obs_id").
|
||||
Where("o.is_superseded = 1").
|
||||
Where("o.project = ?", project).
|
||||
Where("oc.detected_at_epoch < ?", cutoffEpoch).
|
||||
Pluck("o.id", &toDelete).Error
|
||||
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if len(toDelete) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
// Delete the conflict records first (due to foreign key constraints)
|
||||
for _, obsID := range toDelete {
|
||||
err := tx.Where("newer_obs_id = ? OR older_obs_id = ?", obsID, obsID).
|
||||
Delete(&ObservationConflict{}).Error
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
// Delete the observations
|
||||
return tx.Delete(&Observation{}, toDelete).Error
|
||||
})
|
||||
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return toDelete, nil
|
||||
}
|
||||
|
||||
// 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 {
|
||||
NewerTitle sql.NullString `gorm:"column:newer_title"`
|
||||
OlderTitle sql.NullString `gorm:"column:older_title"`
|
||||
ObservationConflict
|
||||
}
|
||||
|
||||
err := s.db.WithContext(ctx).
|
||||
Table("observation_conflicts oc").
|
||||
Select("oc.*, "+
|
||||
"COALESCE(newer.title, '') as newer_title, "+
|
||||
"COALESCE(older.title, '') as older_title").
|
||||
Joins("JOIN observations newer ON newer.id = oc.newer_obs_id").
|
||||
Joins("JOIN observations older ON older.id = oc.older_obs_id").
|
||||
Where("newer.project = ? OR older.project = ?", project, project).
|
||||
Order("oc.detected_at_epoch DESC").
|
||||
Limit(limit).
|
||||
Scan(&results).Error
|
||||
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
conflicts := make([]*ConflictWithDetails, len(results))
|
||||
for i, r := range results {
|
||||
conflicts[i] = &ConflictWithDetails{
|
||||
Conflict: toModelConflict(&r.ObservationConflict),
|
||||
NewerObsTitle: r.NewerTitle.String,
|
||||
OlderObsTitle: r.OlderTitle.String,
|
||||
}
|
||||
}
|
||||
|
||||
return conflicts, nil
|
||||
}
|
||||
|
||||
// toModelConflict converts a GORM ObservationConflict to a pkg/models ObservationConflict.
|
||||
func toModelConflict(c *ObservationConflict) *models.ObservationConflict {
|
||||
conflict := &models.ObservationConflict{
|
||||
ID: c.ID,
|
||||
NewerObsID: c.NewerObsID,
|
||||
OlderObsID: c.OlderObsID,
|
||||
ConflictType: c.ConflictType,
|
||||
Resolution: c.Resolution,
|
||||
DetectedAt: c.DetectedAt,
|
||||
DetectedAtEpoch: c.DetectedAtEpoch,
|
||||
Resolved: c.Resolved == 1,
|
||||
}
|
||||
|
||||
if c.Reason.Valid {
|
||||
conflict.Reason = c.Reason.String
|
||||
}
|
||||
if c.ResolvedAt.Valid {
|
||||
s := c.ResolvedAt.String
|
||||
conflict.ResolvedAt = &s
|
||||
}
|
||||
|
||||
return conflict
|
||||
}
|
||||
|
||||
// toModelConflicts converts a slice of GORM ObservationConflicts to pkg/models ObservationConflicts.
|
||||
func toModelConflicts(conflicts []ObservationConflict) []*models.ObservationConflict {
|
||||
result := make([]*models.ObservationConflict, len(conflicts))
|
||||
for i, c := range conflicts {
|
||||
result[i] = toModelConflict(&c)
|
||||
}
|
||||
return result
|
||||
}
|
||||
@@ -1,637 +0,0 @@
|
||||
//go:build fts5
|
||||
|
||||
// Package gorm provides GORM-based database operations for claude-mnemonic.
|
||||
package gorm
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"gorm.io/gorm/logger"
|
||||
|
||||
"github.com/lukaszraczylo/claude-mnemonic/pkg/models"
|
||||
)
|
||||
|
||||
// testConflictStore creates a ConflictStore with a temporary database for testing.
|
||||
func testConflictStore(t *testing.T) (*ConflictStore, *Store, func()) {
|
||||
t.Helper()
|
||||
|
||||
tmpDir, err := os.MkdirTemp("", "gorm_conflict_test_*")
|
||||
if err != nil {
|
||||
t.Fatalf("create temp dir: %v", err)
|
||||
}
|
||||
|
||||
dbPath := filepath.Join(tmpDir, "test.db")
|
||||
cfg := Config{
|
||||
Path: dbPath,
|
||||
MaxConns: 4,
|
||||
LogLevel: logger.Silent,
|
||||
}
|
||||
|
||||
store, err := NewStore(cfg)
|
||||
if err != nil {
|
||||
os.RemoveAll(tmpDir)
|
||||
t.Fatalf("NewStore failed: %v", err)
|
||||
}
|
||||
|
||||
conflictStore := NewConflictStore(store)
|
||||
|
||||
cleanup := func() {
|
||||
store.Close()
|
||||
os.RemoveAll(tmpDir)
|
||||
}
|
||||
|
||||
return conflictStore, store, cleanup
|
||||
}
|
||||
|
||||
func TestConflictStore_StoreConflict(t *testing.T) {
|
||||
conflictStore, store, cleanup := testConflictStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Create session for observations
|
||||
sessionStore := NewSessionStore(store)
|
||||
sessionID, err := sessionStore.CreateSDKSession(ctx, "claude-1", "test-project", "")
|
||||
require.NoError(t, err)
|
||||
|
||||
// Create test observations
|
||||
obsStore := NewObservationStore(store, nil, nil, nil)
|
||||
obs1 := &models.ParsedObservation{
|
||||
Type: models.ObsTypeDiscovery,
|
||||
Title: "Newer observation",
|
||||
}
|
||||
obsID1, _, err := obsStore.StoreObservation(ctx, "claude-1", "test-project", obs1, int(sessionID), 1)
|
||||
require.NoError(t, err)
|
||||
|
||||
obs2 := &models.ParsedObservation{
|
||||
Type: models.ObsTypeDiscovery,
|
||||
Title: "Older observation",
|
||||
}
|
||||
obsID2, _, err := obsStore.StoreObservation(ctx, "claude-1", "test-project", obs2, int(sessionID), 2)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Create conflict
|
||||
now := time.Now()
|
||||
conflict := &models.ObservationConflict{
|
||||
NewerObsID: obsID1,
|
||||
OlderObsID: obsID2,
|
||||
ConflictType: models.ConflictContradicts,
|
||||
Resolution: models.ResolutionPreferNewer,
|
||||
Reason: "Newer observation contradicts older one",
|
||||
DetectedAt: now.Format(time.RFC3339),
|
||||
DetectedAtEpoch: now.UnixMilli(),
|
||||
Resolved: false,
|
||||
}
|
||||
|
||||
id, err := conflictStore.StoreConflict(ctx, conflict)
|
||||
require.NoError(t, err)
|
||||
assert.Greater(t, id, int64(0))
|
||||
|
||||
// Verify conflict was stored
|
||||
var count int64
|
||||
store.DB.Model(&ObservationConflict{}).Where("id = ?", id).Count(&count)
|
||||
assert.Equal(t, int64(1), count)
|
||||
}
|
||||
|
||||
func TestConflictStore_MarkObservationSuperseded(t *testing.T) {
|
||||
conflictStore, store, cleanup := testConflictStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Create observation
|
||||
sessionStore := NewSessionStore(store)
|
||||
sessionID, err := sessionStore.CreateSDKSession(ctx, "claude-1", "test-project", "")
|
||||
require.NoError(t, err)
|
||||
|
||||
obsStore := NewObservationStore(store, nil, nil, nil)
|
||||
obs := &models.ParsedObservation{
|
||||
Type: models.ObsTypeDiscovery,
|
||||
Title: "Test observation",
|
||||
}
|
||||
obsID, _, err := obsStore.StoreObservation(ctx, "claude-1", "test-project", obs, int(sessionID), 1)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Mark as superseded
|
||||
err = conflictStore.MarkObservationSuperseded(ctx, obsID)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Verify it's marked
|
||||
var dbObs Observation
|
||||
store.DB.First(&dbObs, obsID)
|
||||
assert.Equal(t, 1, dbObs.IsSuperseded)
|
||||
}
|
||||
|
||||
func TestConflictStore_MarkObservationsSuperseded(t *testing.T) {
|
||||
conflictStore, store, cleanup := testConflictStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
obsIDs []int64
|
||||
setup func() []int64
|
||||
}{
|
||||
{
|
||||
name: "empty list",
|
||||
obsIDs: []int64{},
|
||||
setup: func() []int64 { return []int64{} },
|
||||
},
|
||||
{
|
||||
name: "single observation",
|
||||
setup: func() []int64 {
|
||||
sessionStore := NewSessionStore(store)
|
||||
sessionID, _ := sessionStore.CreateSDKSession(ctx, "claude-1", "test-project", "")
|
||||
obsStore := NewObservationStore(store, nil, nil, nil)
|
||||
obs := &models.ParsedObservation{
|
||||
Type: models.ObsTypeDiscovery,
|
||||
Title: "Test",
|
||||
}
|
||||
id, _, _ := obsStore.StoreObservation(ctx, "claude-1", "test-project", obs, int(sessionID), 1)
|
||||
return []int64{id}
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "multiple observations",
|
||||
setup: func() []int64 {
|
||||
sessionStore := NewSessionStore(store)
|
||||
sessionID, _ := sessionStore.CreateSDKSession(ctx, "claude-1", "test-project", "")
|
||||
obsStore := NewObservationStore(store, nil, nil, nil)
|
||||
var ids []int64
|
||||
for i := 0; i < 3; i++ {
|
||||
obs := &models.ParsedObservation{
|
||||
Type: models.ObsTypeDiscovery,
|
||||
Title: "Test",
|
||||
}
|
||||
id, _, _ := obsStore.StoreObservation(ctx, "claude-1", "test-project", obs, int(sessionID), int64(i+1))
|
||||
ids = append(ids, id)
|
||||
}
|
||||
return ids
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
obsIDs := tt.setup()
|
||||
err := conflictStore.MarkObservationsSuperseded(ctx, obsIDs)
|
||||
require.NoError(t, err)
|
||||
|
||||
if len(obsIDs) > 0 {
|
||||
// Verify all are marked
|
||||
for _, id := range obsIDs {
|
||||
var dbObs Observation
|
||||
store.DB.First(&dbObs, id)
|
||||
assert.Equal(t, 1, dbObs.IsSuperseded)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestConflictStore_GetConflictsByObservationID(t *testing.T) {
|
||||
conflictStore, store, cleanup := testConflictStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Create observations
|
||||
sessionStore := NewSessionStore(store)
|
||||
sessionID, err := sessionStore.CreateSDKSession(ctx, "claude-1", "test-project", "")
|
||||
require.NoError(t, err)
|
||||
|
||||
obsStore := NewObservationStore(store, nil, nil, nil)
|
||||
var obsIDs []int64
|
||||
for i := 0; i < 3; i++ {
|
||||
obs := &models.ParsedObservation{
|
||||
Type: models.ObsTypeDiscovery,
|
||||
Title: "Test",
|
||||
}
|
||||
id, _, err := obsStore.StoreObservation(ctx, "claude-1", "test-project", obs, int(sessionID), int64(i+1))
|
||||
require.NoError(t, err)
|
||||
obsIDs = append(obsIDs, id)
|
||||
}
|
||||
|
||||
// Create conflicts involving observation 2 (index 1)
|
||||
now := time.Now()
|
||||
conflict1 := &models.ObservationConflict{
|
||||
NewerObsID: obsIDs[0],
|
||||
OlderObsID: obsIDs[1],
|
||||
ConflictType: models.ConflictContradicts,
|
||||
Resolution: models.ResolutionPreferNewer,
|
||||
Reason: "reason1",
|
||||
DetectedAt: now.Format(time.RFC3339),
|
||||
DetectedAtEpoch: now.UnixMilli(),
|
||||
}
|
||||
_, err = conflictStore.StoreConflict(ctx, conflict1)
|
||||
require.NoError(t, err)
|
||||
|
||||
conflict2 := &models.ObservationConflict{
|
||||
NewerObsID: obsIDs[1],
|
||||
OlderObsID: obsIDs[2],
|
||||
ConflictType: models.ConflictSuperseded,
|
||||
Resolution: models.ResolutionPreferNewer,
|
||||
Reason: "reason2",
|
||||
DetectedAt: now.Format(time.RFC3339),
|
||||
DetectedAtEpoch: now.UnixMilli(),
|
||||
}
|
||||
_, err = conflictStore.StoreConflict(ctx, conflict2)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Get conflicts for observation 2 (involved in 2 conflicts)
|
||||
conflicts, err := conflictStore.GetConflictsByObservationID(ctx, obsIDs[1])
|
||||
require.NoError(t, err)
|
||||
assert.Len(t, conflicts, 2)
|
||||
}
|
||||
|
||||
func TestConflictStore_GetUnresolvedConflicts(t *testing.T) {
|
||||
conflictStore, store, cleanup := testConflictStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Create observations
|
||||
sessionStore := NewSessionStore(store)
|
||||
sessionID, err := sessionStore.CreateSDKSession(ctx, "claude-1", "test-project", "")
|
||||
require.NoError(t, err)
|
||||
|
||||
obsStore := NewObservationStore(store, nil, nil, nil)
|
||||
obs1 := &models.ParsedObservation{
|
||||
Type: models.ObsTypeDiscovery,
|
||||
Title: "Test1",
|
||||
}
|
||||
obsID1, _, err := obsStore.StoreObservation(ctx, "claude-1", "test-project", obs1, int(sessionID), 1)
|
||||
require.NoError(t, err)
|
||||
|
||||
obs2 := &models.ParsedObservation{
|
||||
Type: models.ObsTypeDiscovery,
|
||||
Title: "Test2",
|
||||
}
|
||||
obsID2, _, err := obsStore.StoreObservation(ctx, "claude-1", "test-project", obs2, int(sessionID), 2)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Create unresolved conflicts
|
||||
now := time.Now()
|
||||
for i := 0; i < 5; i++ {
|
||||
conflict := &models.ObservationConflict{
|
||||
NewerObsID: obsID1,
|
||||
OlderObsID: obsID2,
|
||||
ConflictType: models.ConflictContradicts,
|
||||
Resolution: models.ResolutionPreferNewer,
|
||||
Reason: "reason",
|
||||
DetectedAt: now.Format(time.RFC3339),
|
||||
DetectedAtEpoch: now.UnixMilli(),
|
||||
Resolved: false,
|
||||
}
|
||||
_, err = conflictStore.StoreConflict(ctx, conflict)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
// Create resolved conflict
|
||||
resolvedAt := now.Format(time.RFC3339)
|
||||
resolvedConflict := &models.ObservationConflict{
|
||||
NewerObsID: obsID1,
|
||||
OlderObsID: obsID2,
|
||||
ConflictType: models.ConflictContradicts,
|
||||
Resolution: models.ResolutionPreferNewer,
|
||||
Reason: "reason",
|
||||
DetectedAt: now.Format(time.RFC3339),
|
||||
DetectedAtEpoch: now.UnixMilli(),
|
||||
Resolved: true,
|
||||
ResolvedAt: &resolvedAt,
|
||||
}
|
||||
_, err = conflictStore.StoreConflict(ctx, resolvedConflict)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Get unresolved conflicts with limit
|
||||
conflicts, err := conflictStore.GetUnresolvedConflicts(ctx, 3)
|
||||
require.NoError(t, err)
|
||||
assert.Len(t, conflicts, 3)
|
||||
|
||||
// Verify all are unresolved
|
||||
for _, c := range conflicts {
|
||||
assert.False(t, c.Resolved)
|
||||
}
|
||||
}
|
||||
|
||||
func TestConflictStore_GetSupersededObservationIDs(t *testing.T) {
|
||||
conflictStore, store, cleanup := testConflictStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Create observations
|
||||
sessionStore := NewSessionStore(store)
|
||||
sessionID, err := sessionStore.CreateSDKSession(ctx, "claude-1", "test-project", "")
|
||||
require.NoError(t, err)
|
||||
|
||||
obsStore := NewObservationStore(store, nil, nil, nil)
|
||||
|
||||
// Create newer observations
|
||||
newer1 := &models.ParsedObservation{
|
||||
Type: models.ObsTypeDiscovery,
|
||||
Title: "Newer1",
|
||||
}
|
||||
newerID1, _, err := obsStore.StoreObservation(ctx, "claude-1", "test-project", newer1, int(sessionID), 1)
|
||||
require.NoError(t, err)
|
||||
|
||||
newer2 := &models.ParsedObservation{
|
||||
Type: models.ObsTypeDiscovery,
|
||||
Title: "Newer2",
|
||||
}
|
||||
newerID2, _, err := obsStore.StoreObservation(ctx, "claude-1", "test-project", newer2, int(sessionID), 2)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Create older observations
|
||||
older1 := &models.ParsedObservation{
|
||||
Type: models.ObsTypeDiscovery,
|
||||
Title: "Older1",
|
||||
}
|
||||
olderID1, _, err := obsStore.StoreObservation(ctx, "claude-1", "test-project", older1, int(sessionID), 3)
|
||||
require.NoError(t, err)
|
||||
|
||||
older2 := &models.ParsedObservation{
|
||||
Type: models.ObsTypeDiscovery,
|
||||
Title: "Older2",
|
||||
}
|
||||
olderID2, _, err := obsStore.StoreObservation(ctx, "claude-1", "test-project", older2, int(sessionID), 4)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Mark older observations as superseded
|
||||
err = conflictStore.MarkObservationsSuperseded(ctx, []int64{olderID1, olderID2})
|
||||
require.NoError(t, err)
|
||||
|
||||
// Create conflicts with prefer_newer resolution
|
||||
now := time.Now()
|
||||
conflict1 := &models.ObservationConflict{
|
||||
NewerObsID: newerID1,
|
||||
OlderObsID: olderID1,
|
||||
ConflictType: models.ConflictSuperseded,
|
||||
Resolution: models.ResolutionPreferNewer,
|
||||
Reason: "reason1",
|
||||
DetectedAt: now.Format(time.RFC3339),
|
||||
DetectedAtEpoch: now.UnixMilli(),
|
||||
}
|
||||
_, err = conflictStore.StoreConflict(ctx, conflict1)
|
||||
require.NoError(t, err)
|
||||
|
||||
conflict2 := &models.ObservationConflict{
|
||||
NewerObsID: newerID2,
|
||||
OlderObsID: olderID2,
|
||||
ConflictType: models.ConflictSuperseded,
|
||||
Resolution: models.ResolutionPreferNewer,
|
||||
Reason: "reason2",
|
||||
DetectedAt: now.Format(time.RFC3339),
|
||||
DetectedAtEpoch: now.UnixMilli(),
|
||||
}
|
||||
_, err = conflictStore.StoreConflict(ctx, conflict2)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Get superseded IDs (should return older observation IDs)
|
||||
ids, err := conflictStore.GetSupersededObservationIDs(ctx, "test-project")
|
||||
require.NoError(t, err)
|
||||
assert.Len(t, ids, 2)
|
||||
assert.Contains(t, ids, olderID1)
|
||||
assert.Contains(t, ids, olderID2)
|
||||
}
|
||||
|
||||
func TestConflictStore_ResolveConflict(t *testing.T) {
|
||||
conflictStore, _, cleanup := testConflictStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Create a simple conflict by inserting directly to DB
|
||||
conflict := &ObservationConflict{
|
||||
NewerObsID: 1,
|
||||
OlderObsID: 2,
|
||||
ConflictType: models.ConflictContradicts,
|
||||
Resolution: models.ResolutionManual,
|
||||
DetectedAt: time.Now().Format(time.RFC3339),
|
||||
DetectedAtEpoch: time.Now().UnixMilli(),
|
||||
Resolved: 0,
|
||||
}
|
||||
conflictStore.db.Create(conflict)
|
||||
|
||||
// Resolve conflict
|
||||
err := conflictStore.ResolveConflict(ctx, conflict.ID, models.ResolutionPreferNewer)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Verify resolved
|
||||
var resolved ObservationConflict
|
||||
conflictStore.db.First(&resolved, conflict.ID)
|
||||
assert.Equal(t, 1, resolved.Resolved)
|
||||
assert.True(t, resolved.ResolvedAt.Valid)
|
||||
assert.Equal(t, models.ResolutionPreferNewer, resolved.Resolution)
|
||||
}
|
||||
|
||||
func TestConflictStore_DeleteConflictsByObservationID(t *testing.T) {
|
||||
conflictStore, _, cleanup := testConflictStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Create conflicts directly in DB
|
||||
now := time.Now()
|
||||
conflicts := []ObservationConflict{
|
||||
{
|
||||
NewerObsID: 1,
|
||||
OlderObsID: 2,
|
||||
ConflictType: models.ConflictContradicts,
|
||||
Resolution: models.ResolutionPreferNewer,
|
||||
DetectedAt: now.Format(time.RFC3339),
|
||||
DetectedAtEpoch: now.UnixMilli(),
|
||||
},
|
||||
{
|
||||
NewerObsID: 3,
|
||||
OlderObsID: 1,
|
||||
ConflictType: models.ConflictContradicts,
|
||||
Resolution: models.ResolutionPreferNewer,
|
||||
DetectedAt: now.Format(time.RFC3339),
|
||||
DetectedAtEpoch: now.UnixMilli(),
|
||||
},
|
||||
{
|
||||
NewerObsID: 2,
|
||||
OlderObsID: 3,
|
||||
ConflictType: models.ConflictContradicts,
|
||||
Resolution: models.ResolutionPreferNewer,
|
||||
DetectedAt: now.Format(time.RFC3339),
|
||||
DetectedAtEpoch: now.UnixMilli(),
|
||||
},
|
||||
}
|
||||
for _, c := range conflicts {
|
||||
conflictStore.db.Create(&c)
|
||||
}
|
||||
|
||||
// Delete conflicts for observation 1
|
||||
err := conflictStore.DeleteConflictsByObservationID(ctx, 1)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Verify only conflicts involving 1 are deleted
|
||||
var count int64
|
||||
conflictStore.db.Model(&ObservationConflict{}).
|
||||
Where("newer_obs_id = 1 OR older_obs_id = 1").
|
||||
Count(&count)
|
||||
assert.Equal(t, int64(0), count)
|
||||
|
||||
// Other conflict should still exist
|
||||
conflictStore.db.Model(&ObservationConflict{}).
|
||||
Where("newer_obs_id = 2 AND older_obs_id = 3").
|
||||
Count(&count)
|
||||
assert.Equal(t, int64(1), count)
|
||||
}
|
||||
|
||||
func TestConflictStore_CleanupSupersededObservations(t *testing.T) {
|
||||
conflictStore, store, cleanup := testConflictStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Create observations
|
||||
sessionStore := NewSessionStore(store)
|
||||
sessionID, err := sessionStore.CreateSDKSession(ctx, "claude-1", "test-project", "")
|
||||
require.NoError(t, err)
|
||||
|
||||
obsStore := NewObservationStore(store, nil, nil, nil)
|
||||
|
||||
// Create newer observations
|
||||
newer1 := &models.ParsedObservation{
|
||||
Type: models.ObsTypeDiscovery,
|
||||
Title: "Newer1",
|
||||
}
|
||||
newerID1, _, err := obsStore.StoreObservation(ctx, "claude-1", "test-project", newer1, int(sessionID), 1)
|
||||
require.NoError(t, err)
|
||||
|
||||
newer2 := &models.ParsedObservation{
|
||||
Type: models.ObsTypeDiscovery,
|
||||
Title: "Newer2",
|
||||
}
|
||||
newerID2, _, err := obsStore.StoreObservation(ctx, "claude-1", "test-project", newer2, int(sessionID), 2)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Create older observations
|
||||
older1 := &models.ParsedObservation{
|
||||
Type: models.ObsTypeDiscovery,
|
||||
Title: "OldSuperseded",
|
||||
}
|
||||
oldSupersededID, _, err := obsStore.StoreObservation(ctx, "claude-1", "test-project", older1, int(sessionID), 3)
|
||||
require.NoError(t, err)
|
||||
|
||||
older2 := &models.ParsedObservation{
|
||||
Type: models.ObsTypeDiscovery,
|
||||
Title: "RecentSuperseded",
|
||||
}
|
||||
recentSupersededID, _, err := obsStore.StoreObservation(ctx, "claude-1", "test-project", older2, int(sessionID), 4)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Mark as superseded
|
||||
err = conflictStore.MarkObservationsSuperseded(ctx, []int64{oldSupersededID, recentSupersededID})
|
||||
require.NoError(t, err)
|
||||
|
||||
// Create conflicts
|
||||
oldTime := time.Now().AddDate(0, 0, -SupersededRetentionDays-1)
|
||||
recentTime := time.Now().AddDate(0, 0, -1)
|
||||
|
||||
// Old conflict (should be deleted)
|
||||
oldConflict := &models.ObservationConflict{
|
||||
NewerObsID: newerID1,
|
||||
OlderObsID: oldSupersededID,
|
||||
ConflictType: models.ConflictSuperseded,
|
||||
Resolution: models.ResolutionPreferNewer,
|
||||
Reason: "old",
|
||||
DetectedAt: oldTime.Format(time.RFC3339),
|
||||
DetectedAtEpoch: oldTime.UnixMilli(),
|
||||
}
|
||||
_, err = conflictStore.StoreConflict(ctx, oldConflict)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Recent conflict (should be kept)
|
||||
recentConflict := &models.ObservationConflict{
|
||||
NewerObsID: newerID2,
|
||||
OlderObsID: recentSupersededID,
|
||||
ConflictType: models.ConflictSuperseded,
|
||||
Resolution: models.ResolutionPreferNewer,
|
||||
Reason: "recent",
|
||||
DetectedAt: recentTime.Format(time.RFC3339),
|
||||
DetectedAtEpoch: recentTime.UnixMilli(),
|
||||
}
|
||||
_, err = conflictStore.StoreConflict(ctx, recentConflict)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Cleanup old superseded observations
|
||||
deletedIDs, err := conflictStore.CleanupSupersededObservations(ctx, "test-project")
|
||||
require.NoError(t, err)
|
||||
assert.Len(t, deletedIDs, 1)
|
||||
assert.Contains(t, deletedIDs, oldSupersededID)
|
||||
|
||||
// Verify only old superseded observation was deleted
|
||||
var count int64
|
||||
store.DB.Model(&Observation{}).Count(&count)
|
||||
assert.Equal(t, int64(3), count) // newer1, newer2, recentSuperseded remain
|
||||
|
||||
// Verify old observation was deleted
|
||||
store.DB.Model(&Observation{}).Where("id = ?", oldSupersededID).Count(&count)
|
||||
assert.Equal(t, int64(0), count)
|
||||
}
|
||||
|
||||
func TestConflictStore_GetConflictsWithDetails(t *testing.T) {
|
||||
conflictStore, store, cleanup := testConflictStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Create observations
|
||||
sessionStore := NewSessionStore(store)
|
||||
sessionID, err := sessionStore.CreateSDKSession(ctx, "claude-1", "test-project", "")
|
||||
require.NoError(t, err)
|
||||
|
||||
obsStore := NewObservationStore(store, nil, nil, nil)
|
||||
|
||||
newer := &models.ParsedObservation{
|
||||
Type: models.ObsTypeDiscovery,
|
||||
Title: "Newer observation",
|
||||
}
|
||||
newerID, _, err := obsStore.StoreObservation(ctx, "claude-1", "test-project", newer, int(sessionID), 1)
|
||||
require.NoError(t, err)
|
||||
|
||||
older := &models.ParsedObservation{
|
||||
Type: models.ObsTypeDiscovery,
|
||||
Title: "Older observation",
|
||||
}
|
||||
olderID, _, err := obsStore.StoreObservation(ctx, "claude-1", "test-project", older, int(sessionID), 2)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Create conflict
|
||||
now := time.Now()
|
||||
conflict := &models.ObservationConflict{
|
||||
NewerObsID: newerID,
|
||||
OlderObsID: olderID,
|
||||
ConflictType: models.ConflictContradicts,
|
||||
Resolution: models.ResolutionPreferNewer,
|
||||
Reason: "Test conflict",
|
||||
DetectedAt: now.Format(time.RFC3339),
|
||||
DetectedAtEpoch: now.UnixMilli(),
|
||||
}
|
||||
_, err = conflictStore.StoreConflict(ctx, conflict)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Get conflicts with details
|
||||
conflicts, err := conflictStore.GetConflictsWithDetails(ctx, "test-project", 10)
|
||||
require.NoError(t, err)
|
||||
assert.Len(t, conflicts, 1)
|
||||
|
||||
// Verify conflict details
|
||||
assert.Equal(t, newerID, conflicts[0].Conflict.NewerObsID)
|
||||
assert.Equal(t, olderID, conflicts[0].Conflict.OlderObsID)
|
||||
assert.Equal(t, models.ConflictContradicts, conflicts[0].Conflict.ConflictType)
|
||||
assert.Equal(t, "Test conflict", conflicts[0].Conflict.Reason)
|
||||
assert.Equal(t, "Newer observation", conflicts[0].NewerObsTitle)
|
||||
assert.Equal(t, "Older observation", conflicts[0].OlderObsTitle)
|
||||
}
|
||||
@@ -1,38 +0,0 @@
|
||||
// Package gorm provides a GORM-based database implementation for claude-mnemonic.
|
||||
//
|
||||
// This is a drop-in replacement for internal/db/sqlite with the following benefits:
|
||||
// - 50% code reduction (8,500 → 4,250 lines)
|
||||
// - Type-safe query building
|
||||
// - Automatic statement caching
|
||||
// - Same performance characteristics
|
||||
// - Zero breaking changes
|
||||
//
|
||||
// Status: Production-ready, not yet integrated
|
||||
//
|
||||
// # Integration
|
||||
//
|
||||
// To use this package instead of internal/db/sqlite:
|
||||
//
|
||||
// import "github.com/lukaszraczylo/claude-mnemonic/internal/db/gorm"
|
||||
//
|
||||
// store, err := gorm.NewStore(gorm.Config{
|
||||
// Path: "/path/to/database.db",
|
||||
// MaxConns: 4,
|
||||
// LogLevel: logger.Silent,
|
||||
// })
|
||||
//
|
||||
// See INTEGRATION_GUIDE.md for complete migration instructions.
|
||||
//
|
||||
// # Testing
|
||||
//
|
||||
// All tests require the fts5 build tag:
|
||||
//
|
||||
// go test -tags "fts5" -v ./internal/db/gorm
|
||||
//
|
||||
// # Performance
|
||||
//
|
||||
// See PERFORMANCE.md for detailed benchmark results.
|
||||
package gorm
|
||||
|
||||
// This file exists for package documentation and to prevent deadcode warnings
|
||||
// on an intentionally unused (but complete and tested) implementation.
|
||||
@@ -1,42 +0,0 @@
|
||||
//go:build fts5
|
||||
|
||||
package gorm
|
||||
|
||||
import (
|
||||
"database/sql"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
_ "github.com/mattn/go-sqlite3"
|
||||
)
|
||||
|
||||
// TestFTS5Available verifies FTS5 is available in mattn/go-sqlite3
|
||||
func TestFTS5Available(t *testing.T) {
|
||||
tmpDir, err := os.MkdirTemp("", "fts5_test_*")
|
||||
if err != nil {
|
||||
t.Fatalf("create temp dir: %v", err)
|
||||
}
|
||||
defer os.RemoveAll(tmpDir)
|
||||
|
||||
dbPath := filepath.Join(tmpDir, "test.db")
|
||||
|
||||
// Open with mattn/go-sqlite3 driver
|
||||
db, err := sql.Open("sqlite3", dbPath)
|
||||
if err != nil {
|
||||
t.Fatalf("open database: %v", err)
|
||||
}
|
||||
defer db.Close()
|
||||
|
||||
// Try to create FTS5 virtual table
|
||||
_, err = db.Exec(`
|
||||
CREATE VIRTUAL TABLE test_fts USING fts5(
|
||||
content
|
||||
)
|
||||
`)
|
||||
if err != nil {
|
||||
t.Fatalf("create FTS5 table failed: %v", err)
|
||||
}
|
||||
|
||||
t.Logf("✅ FTS5 is available in mattn/go-sqlite3")
|
||||
}
|
||||
@@ -1,101 +0,0 @@
|
||||
// Package gorm provides GORM-based database operations for claude-mnemonic.
|
||||
package gorm
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"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 {
|
||||
now := time.Now()
|
||||
session := &SDKSession{
|
||||
ClaudeSessionID: sdkSessionID,
|
||||
SDKSessionID: sqlNullString(sdkSessionID),
|
||||
Project: project,
|
||||
Status: "active",
|
||||
StartedAt: now.Format(time.RFC3339),
|
||||
StartedAtEpoch: now.UnixMilli(),
|
||||
PromptCounter: 0,
|
||||
}
|
||||
|
||||
// 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.
|
||||
func sqlNullString(s string) sql.NullString {
|
||||
if s == "" {
|
||||
return sql.NullString{Valid: false}
|
||||
}
|
||||
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 {
|
||||
return parsed
|
||||
}
|
||||
}
|
||||
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),
|
||||
}
|
||||
}
|
||||
@@ -1,343 +0,0 @@
|
||||
//go:build fts5
|
||||
|
||||
// Package gorm provides GORM-based database operations for claude-mnemonic.
|
||||
package gorm
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"gorm.io/gorm/logger"
|
||||
|
||||
"github.com/lukaszraczylo/claude-mnemonic/pkg/models"
|
||||
)
|
||||
|
||||
// TestIntegration_EndToEndWorkflow verifies a complete workflow
|
||||
// simulating real usage of the GORM package.
|
||||
func TestIntegration_EndToEndWorkflow(t *testing.T) {
|
||||
tmpDir, err := os.MkdirTemp("", "gorm_integration_test_*")
|
||||
require.NoError(t, err)
|
||||
defer os.RemoveAll(tmpDir)
|
||||
|
||||
dbPath := filepath.Join(tmpDir, "test.db")
|
||||
cfg := Config{
|
||||
Path: dbPath,
|
||||
MaxConns: 4,
|
||||
LogLevel: logger.Silent,
|
||||
}
|
||||
|
||||
// Step 1: Initialize store
|
||||
store, err := NewStore(cfg)
|
||||
require.NoError(t, err)
|
||||
defer store.Close()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Step 2: Create all store types
|
||||
sessionStore := NewSessionStore(store)
|
||||
summaryStore := NewSummaryStore(store)
|
||||
conflictStore := NewConflictStore(store)
|
||||
relationStore := NewRelationStore(store)
|
||||
patternStore := NewPatternStore(store)
|
||||
|
||||
// Create observation store with dependencies
|
||||
observationStore := NewObservationStore(store, nil, conflictStore, relationStore)
|
||||
promptStore := NewPromptStore(store, nil)
|
||||
|
||||
// Step 3: Create a session
|
||||
sessionID, err := sessionStore.CreateSDKSession(ctx, "claude-test", "test-project", "")
|
||||
require.NoError(t, err)
|
||||
assert.Greater(t, sessionID, int64(0))
|
||||
|
||||
// Step 4: Store observations
|
||||
obs1 := &models.ParsedObservation{
|
||||
Type: models.ObsTypeDiscovery,
|
||||
Title: "Test Discovery",
|
||||
Subtitle: "Testing GORM integration",
|
||||
Facts: []string{"Fact 1", "Fact 2"},
|
||||
Concepts: []string{"testing", "integration"},
|
||||
}
|
||||
|
||||
obsID1, _, err := observationStore.StoreObservation(ctx, "claude-test", "test-project", obs1, int(sessionID), 1)
|
||||
require.NoError(t, err)
|
||||
assert.Greater(t, obsID1, int64(0))
|
||||
|
||||
obs2 := &models.ParsedObservation{
|
||||
Type: models.ObsTypeBugfix,
|
||||
Title: "Test Bugfix",
|
||||
Facts: []string{"Fixed bug"},
|
||||
Concepts: []string{"bugfix"},
|
||||
}
|
||||
|
||||
obsID2, _, err := observationStore.StoreObservation(ctx, "claude-test", "test-project", obs2, int(sessionID), 2)
|
||||
require.NoError(t, err)
|
||||
assert.Greater(t, obsID2, int64(0))
|
||||
|
||||
// Step 5: Create relations
|
||||
now := time.Now()
|
||||
relation := &models.ObservationRelation{
|
||||
SourceID: obsID1,
|
||||
TargetID: obsID2,
|
||||
RelationType: models.RelationCauses,
|
||||
Confidence: 0.8,
|
||||
DetectionSource: models.DetectionSourceFileOverlap,
|
||||
CreatedAt: now.Format(time.RFC3339),
|
||||
CreatedAtEpoch: now.UnixMilli(),
|
||||
}
|
||||
|
||||
relID, err := relationStore.StoreRelation(ctx, relation)
|
||||
require.NoError(t, err)
|
||||
assert.Greater(t, relID, int64(0))
|
||||
|
||||
// Step 6: Update importance scores
|
||||
err = observationStore.UpdateImportanceScore(ctx, obsID1, 5.0)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Step 7: Increment retrieval counts
|
||||
err = observationStore.IncrementRetrievalCount(ctx, []int64{obsID1, obsID2})
|
||||
require.NoError(t, err)
|
||||
|
||||
// Step 8: Create a pattern
|
||||
pattern := &models.Pattern{
|
||||
Name: "Test Pattern",
|
||||
Type: models.PatternTypeBug,
|
||||
Signature: []string{"bug", "fix"},
|
||||
Frequency: 1,
|
||||
Projects: []string{"test-project"},
|
||||
ObservationIDs: []int64{obsID1, obsID2},
|
||||
Status: models.PatternStatusActive,
|
||||
Confidence: 0.75,
|
||||
LastSeenAt: now.Format(time.RFC3339),
|
||||
LastSeenEpoch: now.UnixMilli(),
|
||||
CreatedAt: now.Format(time.RFC3339),
|
||||
CreatedAtEpoch: now.UnixMilli(),
|
||||
}
|
||||
|
||||
patternID, err := patternStore.StorePattern(ctx, pattern)
|
||||
require.NoError(t, err)
|
||||
assert.Greater(t, patternID, int64(0))
|
||||
|
||||
// Step 9: Store a prompt
|
||||
promptID, err := promptStore.SaveUserPromptWithMatches(ctx, "claude-test", 1, "Test prompt", 2)
|
||||
require.NoError(t, err)
|
||||
assert.Greater(t, promptID, int64(0))
|
||||
|
||||
// Step 10: Store a summary
|
||||
summary := &models.ParsedSummary{
|
||||
Request: "Test request",
|
||||
Investigated: "Test investigation",
|
||||
Learned: "Test learning",
|
||||
Completed: "Test completion",
|
||||
NextSteps: "Test next steps",
|
||||
Notes: "Test notes",
|
||||
}
|
||||
|
||||
summaryID, _, err := summaryStore.StoreSummary(ctx, "claude-test", "test-project", summary, 1, 100)
|
||||
require.NoError(t, err)
|
||||
assert.Greater(t, summaryID, int64(0))
|
||||
|
||||
// Step 11: Verify data retrieval
|
||||
retrievedObs, err := observationStore.GetObservationByID(ctx, obsID1)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, retrievedObs)
|
||||
assert.Equal(t, "Test Discovery", retrievedObs.Title.String)
|
||||
assert.Equal(t, 5.0, retrievedObs.ImportanceScore)
|
||||
assert.Equal(t, 1, retrievedObs.RetrievalCount)
|
||||
|
||||
// Step 12: Verify relations
|
||||
relations, err := relationStore.GetRelationsByObservationID(ctx, obsID1)
|
||||
require.NoError(t, err)
|
||||
assert.Len(t, relations, 1)
|
||||
assert.Equal(t, obsID2, relations[0].TargetID)
|
||||
|
||||
// Step 13: Verify pattern
|
||||
retrievedPattern, err := patternStore.GetPatternByID(ctx, patternID)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, retrievedPattern)
|
||||
assert.Equal(t, "Test Pattern", retrievedPattern.Name)
|
||||
|
||||
// Step 14: Verify stats
|
||||
stats, err := observationStore.GetObservationFeedbackStats(ctx, "test-project")
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 2, stats.Total)
|
||||
|
||||
t.Log("✅ End-to-end integration test passed!")
|
||||
}
|
||||
|
||||
// TestIntegration_StoreCompatibility verifies that Store methods work correctly.
|
||||
func TestIntegration_StoreCompatibility(t *testing.T) {
|
||||
tmpDir, err := os.MkdirTemp("", "gorm_store_test_*")
|
||||
require.NoError(t, err)
|
||||
defer os.RemoveAll(tmpDir)
|
||||
|
||||
dbPath := filepath.Join(tmpDir, "test.db")
|
||||
cfg := Config{
|
||||
Path: dbPath,
|
||||
MaxConns: 4,
|
||||
LogLevel: logger.Silent,
|
||||
}
|
||||
|
||||
store, err := NewStore(cfg)
|
||||
require.NoError(t, err)
|
||||
defer store.Close()
|
||||
|
||||
// Verify raw DB access (needed for vector client)
|
||||
rawDB := store.GetRawDB()
|
||||
require.NotNil(t, rawDB)
|
||||
assert.IsType(t, &sql.DB{}, rawDB)
|
||||
|
||||
// Verify GORM DB access
|
||||
gormDB := store.GetDB()
|
||||
require.NotNil(t, gormDB)
|
||||
|
||||
// Verify Close works
|
||||
err = store.Close()
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
// TestIntegration_ConcurrentAccess verifies thread-safe operations.
|
||||
func TestIntegration_ConcurrentAccess(t *testing.T) {
|
||||
tmpDir, err := os.MkdirTemp("", "gorm_concurrent_test_*")
|
||||
require.NoError(t, err)
|
||||
defer os.RemoveAll(tmpDir)
|
||||
|
||||
dbPath := filepath.Join(tmpDir, "test.db")
|
||||
cfg := Config{
|
||||
Path: dbPath,
|
||||
MaxConns: 4,
|
||||
LogLevel: logger.Silent,
|
||||
}
|
||||
|
||||
store, err := NewStore(cfg)
|
||||
require.NoError(t, err)
|
||||
defer store.Close()
|
||||
|
||||
sessionStore := NewSessionStore(store)
|
||||
ctx := context.Background()
|
||||
|
||||
// Create session
|
||||
sessionID, err := sessionStore.CreateSDKSession(ctx, "claude-concurrent", "test-project", "")
|
||||
require.NoError(t, err)
|
||||
|
||||
// Concurrent prompt counter increments
|
||||
done := make(chan bool)
|
||||
numGoroutines := 10
|
||||
|
||||
for i := 0; i < numGoroutines; i++ {
|
||||
go func() {
|
||||
_, err := sessionStore.IncrementPromptCounter(ctx, sessionID)
|
||||
assert.NoError(t, err)
|
||||
done <- true
|
||||
}()
|
||||
}
|
||||
|
||||
// Wait for all goroutines
|
||||
for i := 0; i < numGoroutines; i++ {
|
||||
<-done
|
||||
}
|
||||
|
||||
// Verify final count
|
||||
session, err := sessionStore.GetSessionByID(ctx, sessionID)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, int64(numGoroutines), int64(session.PromptCounter))
|
||||
|
||||
t.Log("✅ Concurrent access test passed!")
|
||||
}
|
||||
|
||||
// TestIntegration_WALMode verifies WAL mode is enabled.
|
||||
func TestIntegration_WALMode(t *testing.T) {
|
||||
tmpDir, err := os.MkdirTemp("", "gorm_wal_test_*")
|
||||
require.NoError(t, err)
|
||||
defer os.RemoveAll(tmpDir)
|
||||
|
||||
dbPath := filepath.Join(tmpDir, "test.db")
|
||||
cfg := Config{
|
||||
Path: dbPath,
|
||||
MaxConns: 4,
|
||||
LogLevel: logger.Silent,
|
||||
}
|
||||
|
||||
store, err := NewStore(cfg)
|
||||
require.NoError(t, err)
|
||||
defer store.Close()
|
||||
|
||||
// Check WAL mode via raw SQL
|
||||
var journalMode string
|
||||
err = store.GetRawDB().QueryRow("PRAGMA journal_mode").Scan(&journalMode)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "wal", journalMode, "WAL mode should be enabled")
|
||||
|
||||
t.Log("✅ WAL mode verification passed!")
|
||||
}
|
||||
|
||||
// TestIntegration_FTS5Search verifies FTS5 functionality.
|
||||
func TestIntegration_FTS5Search(t *testing.T) {
|
||||
tmpDir, err := os.MkdirTemp("", "gorm_fts5_test_*")
|
||||
require.NoError(t, err)
|
||||
defer os.RemoveAll(tmpDir)
|
||||
|
||||
dbPath := filepath.Join(tmpDir, "test.db")
|
||||
cfg := Config{
|
||||
Path: dbPath,
|
||||
MaxConns: 4,
|
||||
LogLevel: logger.Silent,
|
||||
}
|
||||
|
||||
store, err := NewStore(cfg)
|
||||
require.NoError(t, err)
|
||||
defer store.Close()
|
||||
|
||||
ctx := context.Background()
|
||||
sessionStore := NewSessionStore(store)
|
||||
observationStore := NewObservationStore(store, nil, nil, nil)
|
||||
|
||||
// Create session
|
||||
sessionID, _ := sessionStore.CreateSDKSession(ctx, "claude-fts5", "test-project", "")
|
||||
|
||||
// Store observations with searchable text
|
||||
obs1 := &models.ParsedObservation{
|
||||
Type: models.ObsTypeDiscovery,
|
||||
Title: "Database optimization techniques",
|
||||
Subtitle: "Improving query performance",
|
||||
Facts: []string{"Use indexes", "Optimize queries"},
|
||||
Concepts: []string{"performance", "optimization"},
|
||||
}
|
||||
|
||||
obsID1, _, _ := observationStore.StoreObservation(ctx, "claude-fts5", "test-project", obs1, int(sessionID), 1)
|
||||
|
||||
obs2 := &models.ParsedObservation{
|
||||
Type: models.ObsTypeBugfix,
|
||||
Title: "Fixed memory leak",
|
||||
Facts: []string{"Closed connections properly"},
|
||||
Concepts: []string{"bugfix", "memory"},
|
||||
}
|
||||
|
||||
observationStore.StoreObservation(ctx, "claude-fts5", "test-project", obs2, int(sessionID), 2)
|
||||
|
||||
// Give FTS5 triggers time to process
|
||||
time.Sleep(100 * time.Millisecond)
|
||||
|
||||
// Search using FTS5
|
||||
results, err := observationStore.SearchObservationsFTS(ctx, "optimization", "test-project", 10)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Should find the optimization observation
|
||||
assert.NotEmpty(t, results, "FTS5 search should return results")
|
||||
|
||||
found := false
|
||||
for _, obs := range results {
|
||||
if obs.ID == obsID1 {
|
||||
found = true
|
||||
break
|
||||
}
|
||||
}
|
||||
assert.True(t, found, "FTS5 should find the optimization observation")
|
||||
|
||||
t.Log("✅ FTS5 search test passed!")
|
||||
}
|
||||
@@ -1,570 +0,0 @@
|
||||
// Package gorm provides GORM-based database operations for claude-mnemonic.
|
||||
package gorm
|
||||
|
||||
import (
|
||||
"database/sql"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"github.com/go-gormigrate/gormigrate/v2"
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/clause"
|
||||
)
|
||||
|
||||
// runMigrations runs all database migrations using gormigrate.
|
||||
func runMigrations(db *gorm.DB, sqlDB *sql.DB) error {
|
||||
m := gormigrate.New(db, gormigrate.DefaultOptions, []*gormigrate.Migration{
|
||||
// Migration 001: Core tables (SDKSession, Observation, SessionSummary)
|
||||
{
|
||||
ID: "001_core_tables",
|
||||
Migrate: func(tx *gorm.DB) error {
|
||||
// AutoMigrate creates tables with all indexes from struct tags
|
||||
if err := tx.AutoMigrate(&SDKSession{}); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := tx.AutoMigrate(&Observation{}); err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.AutoMigrate(&SessionSummary{})
|
||||
},
|
||||
Rollback: func(tx *gorm.DB) error {
|
||||
return tx.Migrator().DropTable("sdk_sessions", "observations", "session_summaries")
|
||||
},
|
||||
},
|
||||
|
||||
// Migration 002: User prompts table
|
||||
{
|
||||
ID: "002_user_prompts",
|
||||
Migrate: func(tx *gorm.DB) error {
|
||||
return tx.AutoMigrate(&UserPrompt{})
|
||||
},
|
||||
Rollback: func(tx *gorm.DB) error {
|
||||
return tx.Migrator().DropTable("user_prompts")
|
||||
},
|
||||
},
|
||||
|
||||
// Migration 003: FTS5 virtual table for user prompts
|
||||
{
|
||||
ID: "003_user_prompts_fts",
|
||||
Migrate: func(tx *gorm.DB) error {
|
||||
sqls := []string{
|
||||
`CREATE VIRTUAL TABLE IF NOT EXISTS user_prompts_fts USING fts5(
|
||||
prompt_text,
|
||||
content='user_prompts',
|
||||
content_rowid='id'
|
||||
)`,
|
||||
`CREATE TRIGGER IF NOT EXISTS user_prompts_ai AFTER INSERT ON user_prompts BEGIN
|
||||
INSERT INTO user_prompts_fts(rowid, prompt_text)
|
||||
VALUES (new.id, new.prompt_text);
|
||||
END`,
|
||||
`CREATE TRIGGER IF NOT EXISTS user_prompts_ad AFTER DELETE ON user_prompts BEGIN
|
||||
INSERT INTO user_prompts_fts(user_prompts_fts, rowid, prompt_text)
|
||||
VALUES('delete', old.id, old.prompt_text);
|
||||
END`,
|
||||
`CREATE TRIGGER IF NOT EXISTS user_prompts_au AFTER UPDATE ON user_prompts BEGIN
|
||||
INSERT INTO user_prompts_fts(user_prompts_fts, rowid, prompt_text)
|
||||
VALUES('delete', old.id, old.prompt_text);
|
||||
INSERT INTO user_prompts_fts(rowid, prompt_text)
|
||||
VALUES (new.id, new.prompt_text);
|
||||
END`,
|
||||
}
|
||||
for _, s := range sqls {
|
||||
if err := tx.Exec(s).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
},
|
||||
Rollback: func(tx *gorm.DB) error {
|
||||
sqls := []string{
|
||||
"DROP TRIGGER IF EXISTS user_prompts_au",
|
||||
"DROP TRIGGER IF EXISTS user_prompts_ad",
|
||||
"DROP TRIGGER IF EXISTS user_prompts_ai",
|
||||
"DROP TABLE IF EXISTS user_prompts_fts",
|
||||
}
|
||||
for _, s := range sqls {
|
||||
if err := tx.Exec(s).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
},
|
||||
},
|
||||
|
||||
// Migration 004: FTS5 virtual table for observations
|
||||
{
|
||||
ID: "004_observations_fts",
|
||||
Migrate: func(tx *gorm.DB) error {
|
||||
sqls := []string{
|
||||
`CREATE VIRTUAL TABLE IF NOT EXISTS observations_fts USING fts5(
|
||||
title, subtitle, narrative,
|
||||
content='observations',
|
||||
content_rowid='id'
|
||||
)`,
|
||||
`CREATE TRIGGER IF NOT EXISTS observations_ai AFTER INSERT ON observations BEGIN
|
||||
INSERT INTO observations_fts(rowid, title, subtitle, narrative)
|
||||
VALUES (new.id, new.title, new.subtitle, new.narrative);
|
||||
END`,
|
||||
`CREATE TRIGGER IF NOT EXISTS observations_ad AFTER DELETE ON observations BEGIN
|
||||
INSERT INTO observations_fts(observations_fts, rowid, title, subtitle, narrative)
|
||||
VALUES('delete', old.id, old.title, old.subtitle, old.narrative);
|
||||
END`,
|
||||
`CREATE TRIGGER IF NOT EXISTS observations_au AFTER UPDATE ON observations BEGIN
|
||||
INSERT INTO observations_fts(observations_fts, rowid, title, subtitle, narrative)
|
||||
VALUES('delete', old.id, old.title, old.subtitle, old.narrative);
|
||||
INSERT INTO observations_fts(rowid, title, subtitle, narrative)
|
||||
VALUES (new.id, new.title, new.subtitle, new.narrative);
|
||||
END`,
|
||||
}
|
||||
for _, s := range sqls {
|
||||
if err := tx.Exec(s).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
},
|
||||
Rollback: func(tx *gorm.DB) error {
|
||||
sqls := []string{
|
||||
"DROP TRIGGER IF EXISTS observations_au",
|
||||
"DROP TRIGGER IF EXISTS observations_ad",
|
||||
"DROP TRIGGER IF EXISTS observations_ai",
|
||||
"DROP TABLE IF EXISTS observations_fts",
|
||||
}
|
||||
for _, s := range sqls {
|
||||
if err := tx.Exec(s).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
},
|
||||
},
|
||||
|
||||
// Migration 005: FTS5 virtual table for session summaries
|
||||
{
|
||||
ID: "005_session_summaries_fts",
|
||||
Migrate: func(tx *gorm.DB) error {
|
||||
sqls := []string{
|
||||
`CREATE VIRTUAL TABLE IF NOT EXISTS session_summaries_fts USING fts5(
|
||||
request, investigated, learned, completed, next_steps, notes,
|
||||
content='session_summaries',
|
||||
content_rowid='id'
|
||||
)`,
|
||||
`CREATE TRIGGER IF NOT EXISTS session_summaries_ai AFTER INSERT ON session_summaries BEGIN
|
||||
INSERT INTO session_summaries_fts(rowid, request, investigated, learned, completed, next_steps, notes)
|
||||
VALUES (new.id, new.request, new.investigated, new.learned, new.completed, new.next_steps, new.notes);
|
||||
END`,
|
||||
`CREATE TRIGGER IF NOT EXISTS session_summaries_ad AFTER DELETE ON session_summaries BEGIN
|
||||
INSERT INTO session_summaries_fts(session_summaries_fts, rowid, request, investigated, learned, completed, next_steps, notes)
|
||||
VALUES('delete', old.id, old.request, old.investigated, old.learned, old.completed, old.next_steps, old.notes);
|
||||
END`,
|
||||
`CREATE TRIGGER IF NOT EXISTS session_summaries_au AFTER UPDATE ON session_summaries BEGIN
|
||||
INSERT INTO session_summaries_fts(session_summaries_fts, rowid, request, investigated, learned, completed, next_steps, notes)
|
||||
VALUES('delete', old.id, old.request, old.investigated, old.learned, old.completed, old.next_steps, old.notes);
|
||||
INSERT INTO session_summaries_fts(rowid, request, investigated, learned, completed, next_steps, notes)
|
||||
VALUES (new.id, new.request, new.investigated, new.learned, new.completed, new.next_steps, new.notes);
|
||||
END`,
|
||||
}
|
||||
for _, s := range sqls {
|
||||
if err := tx.Exec(s).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
},
|
||||
Rollback: func(tx *gorm.DB) error {
|
||||
sqls := []string{
|
||||
"DROP TRIGGER IF EXISTS session_summaries_au",
|
||||
"DROP TRIGGER IF EXISTS session_summaries_ad",
|
||||
"DROP TRIGGER IF EXISTS session_summaries_ai",
|
||||
"DROP TABLE IF EXISTS session_summaries_fts",
|
||||
}
|
||||
for _, s := range sqls {
|
||||
if err := tx.Exec(s).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
},
|
||||
},
|
||||
|
||||
// Migration 006: sqlite-vec vectors table
|
||||
{
|
||||
ID: "006_sqlite_vec_vectors",
|
||||
Migrate: func(tx *gorm.DB) error {
|
||||
// Note: Uses bge-small-en-v1.5 embeddings (384 dimensions) with model_version
|
||||
sql := `CREATE VIRTUAL TABLE IF NOT EXISTS vectors USING vec0(
|
||||
doc_id TEXT PRIMARY KEY,
|
||||
embedding float[384],
|
||||
sqlite_id INTEGER,
|
||||
doc_type TEXT,
|
||||
field_type TEXT,
|
||||
project TEXT,
|
||||
scope TEXT,
|
||||
model_version TEXT
|
||||
)`
|
||||
return tx.Exec(sql).Error
|
||||
},
|
||||
Rollback: func(tx *gorm.DB) error {
|
||||
return tx.Exec("DROP TABLE IF EXISTS vectors").Error
|
||||
},
|
||||
},
|
||||
|
||||
// Migration 007: Concept weights table with seed data
|
||||
{
|
||||
ID: "007_concept_weights",
|
||||
Migrate: func(tx *gorm.DB) error {
|
||||
if err := tx.AutoMigrate(&ConceptWeight{}); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// Seed default concept weights
|
||||
now := time.Now().Format(time.RFC3339)
|
||||
weights := []ConceptWeight{
|
||||
{Concept: "security", Weight: 0.30, UpdatedAt: now},
|
||||
{Concept: "gotcha", Weight: 0.25, UpdatedAt: now},
|
||||
{Concept: "best-practice", Weight: 0.20, UpdatedAt: now},
|
||||
{Concept: "anti-pattern", Weight: 0.20, UpdatedAt: now},
|
||||
{Concept: "architecture", Weight: 0.15, UpdatedAt: now},
|
||||
{Concept: "performance", Weight: 0.15, UpdatedAt: now},
|
||||
{Concept: "error-handling", Weight: 0.15, UpdatedAt: now},
|
||||
{Concept: "pattern", Weight: 0.10, UpdatedAt: now},
|
||||
{Concept: "testing", Weight: 0.10, UpdatedAt: now},
|
||||
{Concept: "debugging", Weight: 0.10, UpdatedAt: now},
|
||||
{Concept: "workflow", Weight: 0.05, UpdatedAt: now},
|
||||
{Concept: "tooling", Weight: 0.05, UpdatedAt: now},
|
||||
}
|
||||
|
||||
// INSERT OR IGNORE equivalent in GORM
|
||||
return tx.Clauses(clause.OnConflict{DoNothing: true}).Create(&weights).Error
|
||||
},
|
||||
Rollback: func(tx *gorm.DB) error {
|
||||
return tx.Migrator().DropTable("concept_weights")
|
||||
},
|
||||
},
|
||||
|
||||
// Migration 008: Observation conflicts table
|
||||
{
|
||||
ID: "008_observation_conflicts",
|
||||
Migrate: func(tx *gorm.DB) error {
|
||||
return tx.AutoMigrate(&ObservationConflict{})
|
||||
},
|
||||
Rollback: func(tx *gorm.DB) error {
|
||||
return tx.Migrator().DropTable("observation_conflicts")
|
||||
},
|
||||
},
|
||||
|
||||
// Migration 009: Patterns table
|
||||
{
|
||||
ID: "009_patterns",
|
||||
Migrate: func(tx *gorm.DB) error {
|
||||
return tx.AutoMigrate(&Pattern{})
|
||||
},
|
||||
Rollback: func(tx *gorm.DB) error {
|
||||
return tx.Migrator().DropTable("patterns")
|
||||
},
|
||||
},
|
||||
|
||||
// Migration 010: FTS5 virtual table for patterns
|
||||
{
|
||||
ID: "010_patterns_fts",
|
||||
Migrate: func(tx *gorm.DB) error {
|
||||
sqls := []string{
|
||||
`CREATE VIRTUAL TABLE IF NOT EXISTS patterns_fts USING fts5(
|
||||
name, description, recommendation,
|
||||
content='patterns',
|
||||
content_rowid='id'
|
||||
)`,
|
||||
`CREATE TRIGGER IF NOT EXISTS patterns_ai AFTER INSERT ON patterns BEGIN
|
||||
INSERT INTO patterns_fts(rowid, name, description, recommendation)
|
||||
VALUES (new.id, new.name, new.description, new.recommendation);
|
||||
END`,
|
||||
`CREATE TRIGGER IF NOT EXISTS patterns_ad AFTER DELETE ON patterns BEGIN
|
||||
INSERT INTO patterns_fts(patterns_fts, rowid, name, description, recommendation)
|
||||
VALUES('delete', old.id, old.name, old.description, old.recommendation);
|
||||
END`,
|
||||
`CREATE TRIGGER IF NOT EXISTS patterns_au AFTER UPDATE ON patterns BEGIN
|
||||
INSERT INTO patterns_fts(patterns_fts, rowid, name, description, recommendation)
|
||||
VALUES('delete', old.id, old.name, old.description, old.recommendation);
|
||||
INSERT INTO patterns_fts(rowid, name, description, recommendation)
|
||||
VALUES (new.id, new.name, new.description, new.recommendation);
|
||||
END`,
|
||||
}
|
||||
for _, s := range sqls {
|
||||
if err := tx.Exec(s).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
},
|
||||
Rollback: func(tx *gorm.DB) error {
|
||||
sqls := []string{
|
||||
"DROP TRIGGER IF EXISTS patterns_au",
|
||||
"DROP TRIGGER IF EXISTS patterns_ad",
|
||||
"DROP TRIGGER IF EXISTS patterns_ai",
|
||||
"DROP TABLE IF EXISTS patterns_fts",
|
||||
}
|
||||
for _, s := range sqls {
|
||||
if err := tx.Exec(s).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
},
|
||||
},
|
||||
|
||||
// Migration 011: Observation relations table
|
||||
{
|
||||
ID: "011_observation_relations",
|
||||
Migrate: func(tx *gorm.DB) error {
|
||||
return tx.AutoMigrate(&ObservationRelation{})
|
||||
},
|
||||
Rollback: func(tx *gorm.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 {
|
||||
return fmt.Errorf("run gormigrate migrations: %w", err)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
@@ -1,268 +0,0 @@
|
||||
// Package gorm provides GORM-based database operations for claude-mnemonic.
|
||||
package gorm
|
||||
|
||||
import (
|
||||
"database/sql"
|
||||
"time"
|
||||
|
||||
"gorm.io/gorm"
|
||||
|
||||
"github.com/lukaszraczylo/claude-mnemonic/pkg/models"
|
||||
)
|
||||
|
||||
// GORM Models
|
||||
|
||||
// Note: JSON types (JSONStringArray, JSONInt64Map) are imported from pkg/models
|
||||
// and already implement sql.Scanner and driver.Valuer interfaces.
|
||||
|
||||
// SDKSession represents a Claude Code session.
|
||||
type SDKSession struct {
|
||||
ClaudeSessionID string `gorm:"uniqueIndex;not null"`
|
||||
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
|
||||
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" }
|
||||
|
||||
// BeforeCreate hook to ensure timestamps are set.
|
||||
func (s *SDKSession) BeforeCreate(tx *gorm.DB) error {
|
||||
if s.StartedAtEpoch == 0 {
|
||||
s.StartedAtEpoch = time.Now().UnixMilli()
|
||||
}
|
||||
if s.StartedAt == "" {
|
||||
s.StartedAt = time.Now().Format(time.RFC3339)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Observation represents a stored observation (learning).
|
||||
// Field order optimized for memory alignment (fieldalignment).
|
||||
type Observation struct {
|
||||
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
|
||||
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"`
|
||||
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" }
|
||||
|
||||
// BeforeCreate hook to ensure defaults are set.
|
||||
func (o *Observation) BeforeCreate(tx *gorm.DB) error {
|
||||
if o.CreatedAtEpoch == 0 {
|
||||
o.CreatedAtEpoch = time.Now().UnixMilli()
|
||||
}
|
||||
if o.CreatedAt == "" {
|
||||
o.CreatedAt = time.Now().Format(time.RFC3339)
|
||||
}
|
||||
if o.ImportanceScore == 0 {
|
||||
o.ImportanceScore = 1.0
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// SessionSummary represents a session summary.
|
||||
type SessionSummary struct {
|
||||
CreatedAt string `gorm:"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" }
|
||||
|
||||
// BeforeCreate hook to ensure timestamps are set.
|
||||
func (s *SessionSummary) BeforeCreate(tx *gorm.DB) error {
|
||||
if s.CreatedAtEpoch == 0 {
|
||||
s.CreatedAtEpoch = time.Now().UnixMilli()
|
||||
}
|
||||
if s.CreatedAt == "" {
|
||||
s.CreatedAt = time.Now().Format(time.RFC3339)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// UserPrompt represents a user prompt.
|
||||
type UserPrompt struct {
|
||||
ClaudeSessionID string `gorm:"index;not null;uniqueIndex:idx_user_prompts_session_number_unique,priority:1"`
|
||||
PromptText string `gorm:"type:text;not null"`
|
||||
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"`
|
||||
}
|
||||
|
||||
func (UserPrompt) TableName() string { return "user_prompts" }
|
||||
|
||||
// BeforeCreate hook to ensure timestamps are set.
|
||||
func (p *UserPrompt) BeforeCreate(tx *gorm.DB) error {
|
||||
if p.CreatedAtEpoch == 0 {
|
||||
p.CreatedAtEpoch = time.Now().UnixMilli()
|
||||
}
|
||||
if p.CreatedAt == "" {
|
||||
p.CreatedAt = time.Now().Format(time.RFC3339)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// ObservationConflict tracks conflicts between observations.
|
||||
type ObservationConflict struct {
|
||||
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"`
|
||||
DetectedAt string `gorm:"not null"`
|
||||
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" }
|
||||
|
||||
// BeforeCreate hook to ensure timestamps are set.
|
||||
func (c *ObservationConflict) BeforeCreate(tx *gorm.DB) error {
|
||||
if c.DetectedAtEpoch == 0 {
|
||||
c.DetectedAtEpoch = time.Now().UnixMilli()
|
||||
}
|
||||
if c.DetectedAt == "" {
|
||||
c.DetectedAt = time.Now().Format(time.RFC3339)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// 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"`
|
||||
Confidence float64 `gorm:"type:real;default:0.5;index:idx_relations_confidence,sort:desc;not null"`
|
||||
CreatedAtEpoch int64 `gorm:"not null"`
|
||||
}
|
||||
|
||||
func (ObservationRelation) TableName() string { return "observation_relations" }
|
||||
|
||||
// BeforeCreate hook to ensure timestamps are set.
|
||||
func (r *ObservationRelation) BeforeCreate(tx *gorm.DB) error {
|
||||
if r.CreatedAtEpoch == 0 {
|
||||
r.CreatedAtEpoch = time.Now().UnixMilli()
|
||||
}
|
||||
if r.CreatedAt == "" {
|
||||
r.CreatedAt = time.Now().Format(time.RFC3339)
|
||||
}
|
||||
if r.Confidence == 0 {
|
||||
r.Confidence = 0.5
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Pattern represents a detected recurring pattern.
|
||||
type Pattern struct {
|
||||
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"`
|
||||
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"`
|
||||
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"`
|
||||
ID int64 `gorm:"primaryKey;autoIncrement"`
|
||||
LastSeenAtEpoch int64 `gorm:"index:idx_patterns_last_seen,sort:desc;not null"`
|
||||
CreatedAtEpoch int64 `gorm:"not null"`
|
||||
}
|
||||
|
||||
func (Pattern) TableName() string { return "patterns" }
|
||||
|
||||
// BeforeCreate hook to ensure timestamps and defaults are set.
|
||||
func (p *Pattern) BeforeCreate(tx *gorm.DB) error {
|
||||
now := time.Now()
|
||||
if p.CreatedAtEpoch == 0 {
|
||||
p.CreatedAtEpoch = now.UnixMilli()
|
||||
}
|
||||
if p.CreatedAt == "" {
|
||||
p.CreatedAt = now.Format(time.RFC3339)
|
||||
}
|
||||
if p.LastSeenAtEpoch == 0 {
|
||||
p.LastSeenAtEpoch = now.UnixMilli()
|
||||
}
|
||||
if p.LastSeenAt == "" {
|
||||
p.LastSeenAt = now.Format(time.RFC3339)
|
||||
}
|
||||
if p.Confidence == 0 {
|
||||
p.Confidence = 0.5
|
||||
}
|
||||
if p.Frequency == 0 {
|
||||
p.Frequency = 1
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// ConceptWeight stores configurable weights for importance scoring.
|
||||
type ConceptWeight struct {
|
||||
Concept string `gorm:"primaryKey;type:text"`
|
||||
UpdatedAt string `gorm:"not null"`
|
||||
Weight float64 `gorm:"type:real;not null;default:0.1"`
|
||||
}
|
||||
|
||||
func (ConceptWeight) TableName() string { return "concept_weights" }
|
||||
|
||||
// BeforeCreate hook to ensure timestamp is set.
|
||||
func (c *ConceptWeight) BeforeCreate(tx *gorm.DB) error {
|
||||
if c.UpdatedAt == "" {
|
||||
c.UpdatedAt = time.Now().Format(time.RFC3339)
|
||||
}
|
||||
if c.Weight == 0 {
|
||||
c.Weight = 0.1
|
||||
}
|
||||
return nil
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,593 +0,0 @@
|
||||
//go:build fts5
|
||||
|
||||
// Package gorm provides GORM-based database operations for claude-mnemonic.
|
||||
package gorm
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"gorm.io/gorm/logger"
|
||||
|
||||
"github.com/lukaszraczylo/claude-mnemonic/pkg/models"
|
||||
)
|
||||
|
||||
// testObservationStore creates an ObservationStore with a temporary database for testing.
|
||||
func testObservationStore(t *testing.T) (*ObservationStore, *Store, func()) {
|
||||
t.Helper()
|
||||
|
||||
tmpDir, err := os.MkdirTemp("", "gorm_observation_test_*")
|
||||
if err != nil {
|
||||
t.Fatalf("create temp dir: %v", err)
|
||||
}
|
||||
|
||||
dbPath := filepath.Join(tmpDir, "test.db")
|
||||
cfg := Config{
|
||||
Path: dbPath,
|
||||
MaxConns: 4,
|
||||
LogLevel: logger.Silent,
|
||||
}
|
||||
|
||||
store, err := NewStore(cfg)
|
||||
if err != nil {
|
||||
os.RemoveAll(tmpDir)
|
||||
t.Fatalf("NewStore failed: %v", err)
|
||||
}
|
||||
|
||||
observationStore := NewObservationStore(store, nil, nil, nil)
|
||||
|
||||
cleanup := func() {
|
||||
store.Close()
|
||||
os.RemoveAll(tmpDir)
|
||||
}
|
||||
|
||||
return observationStore, store, cleanup
|
||||
}
|
||||
|
||||
func TestObservationStore_StoreObservation(t *testing.T) {
|
||||
observationStore, store, cleanup := testObservationStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Create a session first
|
||||
sessionStore := NewSessionStore(store)
|
||||
_, err := sessionStore.CreateSDKSession(ctx, "claude-1", "test-project", "")
|
||||
require.NoError(t, err)
|
||||
|
||||
// Store an observation
|
||||
observation := &models.ParsedObservation{
|
||||
Type: models.ObsTypeDecision,
|
||||
Title: "User prefers tabs over spaces",
|
||||
Narrative: "Observed in code formatting",
|
||||
Concepts: []string{"coding-style", "preferences"},
|
||||
}
|
||||
|
||||
id, epoch, err := observationStore.StoreObservation(ctx, "claude-1", "test-project", observation, 1, 100)
|
||||
require.NoError(t, err)
|
||||
assert.Greater(t, id, int64(0))
|
||||
assert.Greater(t, epoch, int64(0))
|
||||
}
|
||||
|
||||
func TestObservationStore_StoreObservation_AutoCreateSession(t *testing.T) {
|
||||
observationStore, _, cleanup := testObservationStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Store observation without pre-creating session
|
||||
observation := &models.ParsedObservation{
|
||||
Type: models.ObsTypeDiscovery,
|
||||
Title: "Test auto-create",
|
||||
}
|
||||
|
||||
id, _, err := observationStore.StoreObservation(ctx, "claude-auto", "auto-project", observation, 1, 50)
|
||||
require.NoError(t, err)
|
||||
assert.Greater(t, id, int64(0))
|
||||
}
|
||||
|
||||
func TestObservationStore_StoreObservation_WithScope(t *testing.T) {
|
||||
observationStore, _, cleanup := testObservationStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
tags []string
|
||||
expectedScope models.ObservationScope
|
||||
}{
|
||||
{
|
||||
name: "Global scope - best practice",
|
||||
tags: []string{"best-practice", "testing"},
|
||||
expectedScope: models.ScopeGlobal,
|
||||
},
|
||||
{
|
||||
name: "Global scope - security",
|
||||
tags: []string{"security", "auth"},
|
||||
expectedScope: models.ScopeGlobal,
|
||||
},
|
||||
{
|
||||
name: "Project scope - specific feature",
|
||||
tags: []string{"feature", "implementation"},
|
||||
expectedScope: models.ScopeProject,
|
||||
},
|
||||
{
|
||||
name: "Project scope - no tags",
|
||||
tags: []string{},
|
||||
expectedScope: models.ScopeProject,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
observation := &models.ParsedObservation{
|
||||
Type: models.ObsTypeDiscovery,
|
||||
Title: "Test scope determination",
|
||||
Concepts: tt.tags,
|
||||
}
|
||||
|
||||
id, _, err := observationStore.StoreObservation(ctx, "claude-1", "test-project", observation, 1, 50)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Verify scope was set correctly
|
||||
observations, err := observationStore.GetObservationsByIDs(ctx, []int64{id}, "default", 10)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, observations, 1)
|
||||
assert.Equal(t, tt.expectedScope, observations[0].Scope)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestObservationStore_StoreObservation_AsyncCleanup(t *testing.T) {
|
||||
observationStore, _, cleanup := testObservationStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Track cleanup calls
|
||||
var cleanupMutex sync.Mutex
|
||||
cleanupCalled := false
|
||||
var cleanupIDs []int64
|
||||
|
||||
cleanupFunc := func(ctx context.Context, deletedIDs []int64) {
|
||||
cleanupMutex.Lock()
|
||||
defer cleanupMutex.Unlock()
|
||||
cleanupCalled = true
|
||||
cleanupIDs = deletedIDs
|
||||
}
|
||||
|
||||
observationStore.cleanupFunc = cleanupFunc
|
||||
|
||||
// Store observations beyond the limit (MaxObservationsPerProject = 100)
|
||||
for i := 0; i < 105; i++ {
|
||||
observation := &models.ParsedObservation{
|
||||
Type: models.ObsTypeDiscovery,
|
||||
Title: "Observation",
|
||||
}
|
||||
_, _, err := observationStore.StoreObservation(ctx, "claude-1", "test-project", observation, i, 50)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
// Wait for async cleanup to complete
|
||||
time.Sleep(200 * time.Millisecond)
|
||||
|
||||
// Verify cleanup was called
|
||||
cleanupMutex.Lock()
|
||||
defer cleanupMutex.Unlock()
|
||||
assert.True(t, cleanupCalled, "Cleanup function should have been called")
|
||||
assert.NotEmpty(t, cleanupIDs, "Cleanup should have deleted some observations")
|
||||
}
|
||||
|
||||
func TestObservationStore_GetObservationsByIDs(t *testing.T) {
|
||||
observationStore, _, cleanup := testObservationStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Store multiple observations with different importance scores
|
||||
var ids []int64
|
||||
for i := 1; i <= 3; i++ {
|
||||
observation := &models.ParsedObservation{
|
||||
Type: models.ObsTypeDiscovery,
|
||||
Title: "Test",
|
||||
}
|
||||
id, _, err := observationStore.StoreObservation(ctx, "claude-1", "test-project", observation, i, 10)
|
||||
require.NoError(t, err)
|
||||
ids = append(ids, id)
|
||||
|
||||
// Update importance score directly
|
||||
observationStore.db.Model(&Observation{}).Where("id = ?", id).Update("importance_score", float64(i))
|
||||
time.Sleep(10 * time.Millisecond) // Ensure different timestamps
|
||||
}
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
orderBy string
|
||||
expected []int64
|
||||
}{
|
||||
{
|
||||
name: "Default ordering - importance desc",
|
||||
orderBy: "default",
|
||||
expected: []int64{ids[2], ids[1], ids[0]}, // High to low importance
|
||||
},
|
||||
{
|
||||
name: "Importance ordering",
|
||||
orderBy: "importance",
|
||||
expected: []int64{ids[2], ids[1], ids[0]},
|
||||
},
|
||||
{
|
||||
name: "Date ascending",
|
||||
orderBy: "date_asc",
|
||||
expected: []int64{ids[0], ids[1], ids[2]}, // Oldest to newest
|
||||
},
|
||||
{
|
||||
name: "Date descending",
|
||||
orderBy: "date_desc",
|
||||
expected: []int64{ids[2], ids[1], ids[0]}, // Newest to oldest
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
observations, err := observationStore.GetObservationsByIDs(ctx, ids, tt.orderBy, 10)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, observations, 3)
|
||||
|
||||
// Verify ordering
|
||||
for i, obs := range observations {
|
||||
assert.Equal(t, tt.expected[i], obs.ID, "Position %d should have ID %d", i, tt.expected[i])
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestObservationStore_GetObservationsByIDs_Limit(t *testing.T) {
|
||||
observationStore, _, cleanup := testObservationStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Store multiple observations
|
||||
var ids []int64
|
||||
for i := 1; i <= 5; i++ {
|
||||
observation := &models.ParsedObservation{
|
||||
Type: models.ObsTypeDiscovery,
|
||||
Title: "Test",
|
||||
}
|
||||
id, _, err := observationStore.StoreObservation(ctx, "claude-1", "test-project", observation, i, 10)
|
||||
require.NoError(t, err)
|
||||
ids = append(ids, id)
|
||||
}
|
||||
|
||||
// Get with limit
|
||||
observations, err := observationStore.GetObservationsByIDs(ctx, ids, "default", 3)
|
||||
require.NoError(t, err)
|
||||
assert.Len(t, observations, 3)
|
||||
}
|
||||
|
||||
func TestObservationStore_GetObservationsByIDs_EmptyInput(t *testing.T) {
|
||||
observationStore, _, cleanup := testObservationStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Get with empty IDs
|
||||
observations, err := observationStore.GetObservationsByIDs(ctx, []int64{}, "default", 10)
|
||||
require.NoError(t, err)
|
||||
assert.Nil(t, observations)
|
||||
}
|
||||
|
||||
func TestObservationStore_GetRecentObservations(t *testing.T) {
|
||||
observationStore, _, cleanup := testObservationStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Store project-scoped observations
|
||||
for i := 1; i <= 3; i++ {
|
||||
observation := &models.ParsedObservation{
|
||||
Type: models.ObsTypeDiscovery,
|
||||
Title: "Project A fact",
|
||||
}
|
||||
_, _, err := observationStore.StoreObservation(ctx, "claude-1", "project-a", observation, i, 10)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
// Store global-scoped observation
|
||||
observation := &models.ParsedObservation{
|
||||
Type: models.ObsTypeDiscovery,
|
||||
Title: "Global best practice",
|
||||
Concepts: []string{"best-practice"},
|
||||
}
|
||||
_, _, err := observationStore.StoreObservation(ctx, "claude-2", "project-b", observation, 1, 10)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Store observation for different project
|
||||
observation = &models.ParsedObservation{
|
||||
Type: models.ObsTypeDiscovery,
|
||||
Title: "Project B fact",
|
||||
}
|
||||
_, _, err = observationStore.StoreObservation(ctx, "claude-2", "project-b", observation, 2, 10)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Wait for any async cleanup to complete before querying
|
||||
time.Sleep(100 * time.Millisecond)
|
||||
|
||||
// Get recent observations for project-a (should include project-a + global)
|
||||
observations, err := observationStore.GetRecentObservations(ctx, "project-a", 10)
|
||||
require.NoError(t, err)
|
||||
assert.Len(t, observations, 4) // 3 project-a + 1 global
|
||||
|
||||
// Verify scope filtering
|
||||
projectCount := 0
|
||||
globalCount := 0
|
||||
for _, obs := range observations {
|
||||
if obs.Scope == models.ScopeProject {
|
||||
assert.Equal(t, "project-a", obs.Project)
|
||||
projectCount++
|
||||
} else if obs.Scope == models.ScopeGlobal {
|
||||
globalCount++
|
||||
}
|
||||
}
|
||||
assert.Equal(t, 3, projectCount)
|
||||
assert.Equal(t, 1, globalCount)
|
||||
}
|
||||
|
||||
func TestObservationStore_GetActiveObservations(t *testing.T) {
|
||||
observationStore, _, cleanup := testObservationStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Store active observation
|
||||
activeObs := &models.ParsedObservation{
|
||||
Type: models.ObsTypeDiscovery,
|
||||
Title: "Active observation",
|
||||
}
|
||||
activeID, _, err := observationStore.StoreObservation(ctx, "claude-1", "test-project", activeObs, 1, 10)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Store superseded observation
|
||||
supersededObs := &models.ParsedObservation{
|
||||
Type: models.ObsTypeDiscovery,
|
||||
Title: "Superseded observation",
|
||||
}
|
||||
supersededID, _, err := observationStore.StoreObservation(ctx, "claude-1", "test-project", supersededObs, 2, 10)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Mark as superseded
|
||||
observationStore.db.Model(&Observation{}).Where("id = ?", supersededID).Update("is_superseded", 1)
|
||||
|
||||
// Get active observations (should exclude superseded)
|
||||
observations, err := observationStore.GetActiveObservations(ctx, "test-project", 10)
|
||||
require.NoError(t, err)
|
||||
assert.Len(t, observations, 1)
|
||||
assert.Equal(t, activeID, observations[0].ID)
|
||||
assert.False(t, observations[0].IsSuperseded)
|
||||
}
|
||||
|
||||
func TestObservationStore_GetSupersededObservations(t *testing.T) {
|
||||
observationStore, _, cleanup := testObservationStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Store active observation
|
||||
activeObs := &models.ParsedObservation{
|
||||
Type: models.ObsTypeDiscovery,
|
||||
Title: "Active observation",
|
||||
}
|
||||
_, _, err := observationStore.StoreObservation(ctx, "claude-1", "test-project", activeObs, 1, 10)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Store superseded observation
|
||||
supersededObs := &models.ParsedObservation{
|
||||
Type: models.ObsTypeDiscovery,
|
||||
Title: "Superseded observation",
|
||||
}
|
||||
supersededID, _, err := observationStore.StoreObservation(ctx, "claude-1", "test-project", supersededObs, 2, 10)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Mark as superseded
|
||||
observationStore.db.Model(&Observation{}).Where("id = ?", supersededID).Update("is_superseded", 1)
|
||||
|
||||
// Get superseded observations (should exclude active)
|
||||
observations, err := observationStore.GetSupersededObservations(ctx, "test-project", 10)
|
||||
require.NoError(t, err)
|
||||
assert.Len(t, observations, 1)
|
||||
assert.Equal(t, supersededID, observations[0].ID)
|
||||
assert.True(t, observations[0].IsSuperseded)
|
||||
}
|
||||
|
||||
func TestObservationStore_GetObservationsByProjectStrict(t *testing.T) {
|
||||
observationStore, _, cleanup := testObservationStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Store project-scoped observations
|
||||
for i := 1; i <= 2; i++ {
|
||||
observation := &models.ParsedObservation{
|
||||
Type: models.ObsTypeDiscovery,
|
||||
Title: "Project A fact",
|
||||
}
|
||||
_, _, err := observationStore.StoreObservation(ctx, "claude-1", "project-a", observation, i, 10)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
// Store global-scoped observation
|
||||
observation := &models.ParsedObservation{
|
||||
Type: models.ObsTypeDiscovery,
|
||||
Title: "Global best practice",
|
||||
Concepts: []string{"best-practice"},
|
||||
}
|
||||
_, _, err := observationStore.StoreObservation(ctx, "claude-2", "project-b", observation, 1, 10)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Get strict project observations (should exclude global)
|
||||
observations, err := observationStore.GetObservationsByProjectStrict(ctx, "project-a", 10)
|
||||
require.NoError(t, err)
|
||||
assert.Len(t, observations, 2) // Only project-a observations
|
||||
|
||||
// Verify all are project-scoped
|
||||
for _, obs := range observations {
|
||||
assert.Equal(t, models.ScopeProject, obs.Scope)
|
||||
assert.Equal(t, "project-a", obs.Project)
|
||||
}
|
||||
}
|
||||
|
||||
func TestObservationStore_SearchObservationsFTS(t *testing.T) {
|
||||
observationStore, _, cleanup := testObservationStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Store observations with searchable content
|
||||
observations := []*models.ParsedObservation{
|
||||
{
|
||||
Type: models.ObsTypeDiscovery,
|
||||
Title: "User prefers React for frontend development",
|
||||
Concepts: []string{"frontend", "react"},
|
||||
},
|
||||
{
|
||||
Type: models.ObsTypeDiscovery,
|
||||
Title: "Backend uses Go with chi router",
|
||||
Concepts: []string{"backend", "golang"},
|
||||
},
|
||||
{
|
||||
Type: models.ObsTypeDiscovery,
|
||||
Title: "Database is SQLite with FTS5",
|
||||
Concepts: []string{"database", "sqlite"},
|
||||
},
|
||||
}
|
||||
|
||||
for i, obs := range observations {
|
||||
_, _, err := observationStore.StoreObservation(ctx, "claude-1", "test-project", obs, i+1, 10)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
// Wait for FTS5 triggers to fire
|
||||
time.Sleep(200 * time.Millisecond)
|
||||
|
||||
// Search for "React frontend"
|
||||
results, err := observationStore.SearchObservationsFTS(ctx, "React frontend", "test-project", 10)
|
||||
require.NoError(t, err)
|
||||
assert.NotEmpty(t, results, "Should find observations matching 'React frontend'")
|
||||
|
||||
// Verify results contain relevant observation
|
||||
found := false
|
||||
for _, obs := range results {
|
||||
if obs.Title.String == "User prefers React for frontend development" {
|
||||
found = true
|
||||
break
|
||||
}
|
||||
}
|
||||
assert.True(t, found, "Should find the React observation")
|
||||
}
|
||||
|
||||
func TestObservationStore_CleanupOldObservations(t *testing.T) {
|
||||
observationStore, _, cleanup := testObservationStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Store observations beyond the limit WITHOUT async cleanup
|
||||
// We disable async cleanup by not setting cleanupFunc
|
||||
var allIDs []int64
|
||||
for i := 0; i < 105; i++ {
|
||||
observation := &models.ParsedObservation{
|
||||
Type: models.ObsTypeDiscovery,
|
||||
Title: "Observation",
|
||||
}
|
||||
id, _, err := observationStore.StoreObservation(ctx, "claude-1", "test-project", observation, i, 10)
|
||||
require.NoError(t, err)
|
||||
allIDs = append(allIDs, id)
|
||||
time.Sleep(2 * time.Millisecond) // Ensure different timestamps
|
||||
}
|
||||
|
||||
// Wait for any async cleanups to complete (even though cleanupFunc is nil)
|
||||
time.Sleep(200 * time.Millisecond)
|
||||
|
||||
// Verify we have 105 observations initially (async cleanup should have run but deleted items)
|
||||
initial, err := observationStore.GetRecentObservations(ctx, "test-project", 200)
|
||||
require.NoError(t, err)
|
||||
|
||||
// If async cleanup already happened, we'll have <= 100
|
||||
// Run cleanup manually to ensure cleanup logic works
|
||||
deletedIDs, err := observationStore.CleanupOldObservations(ctx, "test-project")
|
||||
require.NoError(t, err)
|
||||
|
||||
// After cleanup (manual or async), we should have at most 100
|
||||
remaining, err := observationStore.GetRecentObservations(ctx, "test-project", 200)
|
||||
require.NoError(t, err)
|
||||
assert.LessOrEqual(t, len(remaining), 100, "Should have at most 100 observations after cleanup")
|
||||
|
||||
// The number deleted should match how many were over the limit
|
||||
expectedDeleted := len(initial) - len(remaining)
|
||||
assert.Len(t, deletedIDs, expectedDeleted, "Should delete observations beyond limit")
|
||||
}
|
||||
|
||||
func TestObservationStore_DeleteObservations(t *testing.T) {
|
||||
observationStore, _, cleanup := testObservationStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Store multiple observations
|
||||
var ids []int64
|
||||
for i := 1; i <= 5; i++ {
|
||||
observation := &models.ParsedObservation{
|
||||
Type: models.ObsTypeDiscovery,
|
||||
Title: "Test",
|
||||
}
|
||||
id, _, err := observationStore.StoreObservation(ctx, "claude-1", "test-project", observation, i, 10)
|
||||
require.NoError(t, err)
|
||||
ids = append(ids, id)
|
||||
}
|
||||
|
||||
// Delete first 3 observations
|
||||
_, err := observationStore.DeleteObservations(ctx, ids[:3])
|
||||
require.NoError(t, err)
|
||||
|
||||
// Verify only 2 remain
|
||||
remaining, err := observationStore.GetRecentObservations(ctx, "test-project", 10)
|
||||
require.NoError(t, err)
|
||||
assert.Len(t, remaining, 2)
|
||||
|
||||
// Verify deleted observations are gone
|
||||
deleted, err := observationStore.GetObservationsByIDs(ctx, ids[:3], "default", 10)
|
||||
require.NoError(t, err)
|
||||
assert.Empty(t, deleted)
|
||||
}
|
||||
|
||||
// Note: TestObservationStore_MarkObservationsSuperseded is omitted because
|
||||
// MarkObservationsSuperseded is a ConflictStore method (Phase 4), not ObservationStore
|
||||
|
||||
func TestObservationStore_GetAllObservations(t *testing.T) {
|
||||
observationStore, _, cleanup := testObservationStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Store observations across projects
|
||||
_, _, err := observationStore.StoreObservation(ctx, "claude-1", "project-a", &models.ParsedObservation{Type: models.ObsTypeDiscovery, Title: "A1"}, 1, 10)
|
||||
require.NoError(t, err)
|
||||
|
||||
_, _, err = observationStore.StoreObservation(ctx, "claude-2", "project-b", &models.ParsedObservation{Type: models.ObsTypeDiscovery, Title: "B1"}, 1, 10)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Get all observations (for vector rebuild)
|
||||
all, err := observationStore.GetAllObservations(ctx)
|
||||
require.NoError(t, err)
|
||||
assert.Len(t, all, 2)
|
||||
|
||||
// Verify ordering by ID
|
||||
assert.Less(t, all[0].ID, all[1].ID)
|
||||
}
|
||||
@@ -1,420 +0,0 @@
|
||||
// Package gorm provides GORM-based database operations for claude-mnemonic.
|
||||
package gorm
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"time"
|
||||
|
||||
"gorm.io/gorm"
|
||||
|
||||
"github.com/lukaszraczylo/claude-mnemonic/pkg/models"
|
||||
)
|
||||
|
||||
// PatternCleanupFunc is a callback for when patterns are deleted.
|
||||
type PatternCleanupFunc func(ctx context.Context, deletedIDs []int64)
|
||||
|
||||
// PatternStore provides pattern-related database operations using GORM.
|
||||
type PatternStore struct {
|
||||
db *gorm.DB
|
||||
cleanupFunc PatternCleanupFunc
|
||||
}
|
||||
|
||||
// NewPatternStore creates a new pattern store.
|
||||
func NewPatternStore(store *Store) *PatternStore {
|
||||
return &PatternStore{
|
||||
db: store.DB,
|
||||
}
|
||||
}
|
||||
|
||||
// SetCleanupFunc sets the callback for when patterns are deleted.
|
||||
func (s *PatternStore) SetCleanupFunc(fn PatternCleanupFunc) {
|
||||
s.cleanupFunc = fn
|
||||
}
|
||||
|
||||
// StorePattern stores a new pattern.
|
||||
func (s *PatternStore) StorePattern(ctx context.Context, pattern *models.Pattern) (int64, error) {
|
||||
dbPattern := &Pattern{
|
||||
Name: pattern.Name,
|
||||
Type: pattern.Type,
|
||||
Signature: pattern.Signature,
|
||||
Frequency: pattern.Frequency,
|
||||
Projects: pattern.Projects,
|
||||
ObservationIDs: pattern.ObservationIDs,
|
||||
Status: pattern.Status,
|
||||
Confidence: pattern.Confidence,
|
||||
LastSeenAt: pattern.LastSeenAt,
|
||||
LastSeenAtEpoch: pattern.LastSeenEpoch,
|
||||
CreatedAt: pattern.CreatedAt,
|
||||
CreatedAtEpoch: pattern.CreatedAtEpoch,
|
||||
}
|
||||
|
||||
if pattern.Description.Valid {
|
||||
dbPattern.Description = sql.NullString{String: pattern.Description.String, Valid: true}
|
||||
}
|
||||
|
||||
if pattern.Recommendation.Valid {
|
||||
dbPattern.Recommendation = sql.NullString{String: pattern.Recommendation.String, Valid: true}
|
||||
}
|
||||
|
||||
if pattern.MergedIntoID.Valid {
|
||||
dbPattern.MergedIntoID = sql.NullInt64{Int64: pattern.MergedIntoID.Int64, Valid: true}
|
||||
}
|
||||
|
||||
result := s.db.WithContext(ctx).Create(dbPattern)
|
||||
if result.Error != nil {
|
||||
return 0, result.Error
|
||||
}
|
||||
|
||||
return dbPattern.ID, nil
|
||||
}
|
||||
|
||||
// UpdatePattern updates an existing pattern.
|
||||
func (s *PatternStore) UpdatePattern(ctx context.Context, pattern *models.Pattern) error {
|
||||
updates := map[string]interface{}{
|
||||
"name": pattern.Name,
|
||||
"type": pattern.Type,
|
||||
"signature": pattern.Signature,
|
||||
"frequency": pattern.Frequency,
|
||||
"projects": pattern.Projects,
|
||||
"observation_ids": pattern.ObservationIDs,
|
||||
"status": pattern.Status,
|
||||
"confidence": pattern.Confidence,
|
||||
"last_seen_at": pattern.LastSeenAt,
|
||||
"last_seen_at_epoch": pattern.LastSeenEpoch,
|
||||
}
|
||||
|
||||
if pattern.Description.Valid {
|
||||
updates["description"] = pattern.Description.String
|
||||
} else {
|
||||
updates["description"] = nil
|
||||
}
|
||||
|
||||
if pattern.Recommendation.Valid {
|
||||
updates["recommendation"] = pattern.Recommendation.String
|
||||
} else {
|
||||
updates["recommendation"] = nil
|
||||
}
|
||||
|
||||
if pattern.MergedIntoID.Valid {
|
||||
updates["merged_into_id"] = pattern.MergedIntoID.Int64
|
||||
} else {
|
||||
updates["merged_into_id"] = nil
|
||||
}
|
||||
|
||||
result := s.db.WithContext(ctx).
|
||||
Model(&Pattern{}).
|
||||
Where("id = ?", pattern.ID).
|
||||
Updates(updates)
|
||||
|
||||
return result.Error
|
||||
}
|
||||
|
||||
// GetPatternByID retrieves a pattern by ID.
|
||||
func (s *PatternStore) GetPatternByID(ctx context.Context, id int64) (*models.Pattern, error) {
|
||||
var dbPattern Pattern
|
||||
|
||||
err := s.db.WithContext(ctx).First(&dbPattern, id).Error
|
||||
if err == gorm.ErrRecordNotFound {
|
||||
return nil, nil
|
||||
}
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return toModelPattern(&dbPattern), nil
|
||||
}
|
||||
|
||||
// GetPatternByName retrieves a pattern by name.
|
||||
func (s *PatternStore) GetPatternByName(ctx context.Context, name string) (*models.Pattern, error) {
|
||||
var dbPattern Pattern
|
||||
|
||||
err := s.db.WithContext(ctx).
|
||||
Where("name = ? AND status = ?", name, models.PatternStatusActive).
|
||||
First(&dbPattern).Error
|
||||
|
||||
if err == gorm.ErrRecordNotFound {
|
||||
return nil, nil
|
||||
}
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return toModelPattern(&dbPattern), nil
|
||||
}
|
||||
|
||||
// GetActivePatterns retrieves all active patterns.
|
||||
func (s *PatternStore) GetActivePatterns(ctx context.Context, limit int) ([]*models.Pattern, error) {
|
||||
var patterns []Pattern
|
||||
|
||||
err := s.db.WithContext(ctx).
|
||||
Where("status = ?", models.PatternStatusActive).
|
||||
Order("frequency DESC, confidence DESC").
|
||||
Limit(limit).
|
||||
Find(&patterns).Error
|
||||
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return toModelPatterns(patterns), nil
|
||||
}
|
||||
|
||||
// GetPatternsByType retrieves patterns of a specific type.
|
||||
func (s *PatternStore) GetPatternsByType(ctx context.Context, patternType models.PatternType, limit int) ([]*models.Pattern, error) {
|
||||
var patterns []Pattern
|
||||
|
||||
err := s.db.WithContext(ctx).
|
||||
Where("type = ? AND status = ?", patternType, models.PatternStatusActive).
|
||||
Order("frequency DESC, confidence DESC").
|
||||
Limit(limit).
|
||||
Find(&patterns).Error
|
||||
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return toModelPatterns(patterns), nil
|
||||
}
|
||||
|
||||
// GetPatternsByProject retrieves patterns that have been observed in a specific project.
|
||||
// Uses raw SQL since JSON_EACH is complex in GORM.
|
||||
func (s *PatternStore) GetPatternsByProject(ctx context.Context, project string, limit int) ([]*models.Pattern, error) {
|
||||
var patterns []Pattern
|
||||
|
||||
// Use raw SQL for JSON_EACH query
|
||||
query := `
|
||||
SELECT * FROM patterns
|
||||
WHERE status = 'active'
|
||||
AND EXISTS (
|
||||
SELECT 1 FROM json_each(projects)
|
||||
WHERE json_each.value = ?
|
||||
)
|
||||
ORDER BY frequency DESC, confidence DESC
|
||||
LIMIT ?
|
||||
`
|
||||
|
||||
err := s.db.WithContext(ctx).
|
||||
Raw(query, project, limit).
|
||||
Scan(&patterns).Error
|
||||
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return toModelPatterns(patterns), nil
|
||||
}
|
||||
|
||||
// FindMatchingPatterns searches for patterns that match a given signature.
|
||||
// Pattern matching is done in Go code for simplicity.
|
||||
func (s *PatternStore) FindMatchingPatterns(ctx context.Context, signature []string, minScore float64) ([]*models.Pattern, error) {
|
||||
// Get all active patterns
|
||||
patterns, err := s.GetActivePatterns(ctx, 100)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Filter by signature match in Go
|
||||
var matches []*models.Pattern
|
||||
for _, pattern := range patterns {
|
||||
score := models.CalculateMatchScore(signature, pattern.Signature)
|
||||
if score >= minScore {
|
||||
matches = append(matches, pattern)
|
||||
}
|
||||
}
|
||||
|
||||
return matches, nil
|
||||
}
|
||||
|
||||
// MarkPatternDeprecated marks a pattern as deprecated.
|
||||
func (s *PatternStore) MarkPatternDeprecated(ctx context.Context, id int64) error {
|
||||
result := s.db.WithContext(ctx).
|
||||
Model(&Pattern{}).
|
||||
Where("id = ?", id).
|
||||
Update("status", models.PatternStatusDeprecated)
|
||||
|
||||
return result.Error
|
||||
}
|
||||
|
||||
// MergePatterns merges a source pattern into a target pattern.
|
||||
func (s *PatternStore) MergePatterns(ctx context.Context, sourceID, targetID int64) error {
|
||||
// Get both patterns
|
||||
source, err := s.GetPatternByID(ctx, sourceID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
target, err := s.GetPatternByID(ctx, targetID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// Merge source into target
|
||||
target.Frequency += source.Frequency
|
||||
|
||||
// Merge projects (deduplicate)
|
||||
for _, proj := range source.Projects {
|
||||
found := false
|
||||
for _, existing := range target.Projects {
|
||||
if existing == proj {
|
||||
found = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
target.Projects = append(target.Projects, proj)
|
||||
}
|
||||
}
|
||||
|
||||
// Merge observation IDs (deduplicate)
|
||||
for _, obsID := range source.ObservationIDs {
|
||||
found := false
|
||||
for _, existing := range target.ObservationIDs {
|
||||
if existing == obsID {
|
||||
found = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
target.ObservationIDs = append(target.ObservationIDs, obsID)
|
||||
}
|
||||
}
|
||||
|
||||
// Update target
|
||||
if err := s.UpdatePattern(ctx, target); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// Mark source as merged
|
||||
source.Status = models.PatternStatusMerged
|
||||
source.MergedIntoID = sql.NullInt64{Int64: targetID, Valid: true}
|
||||
return s.UpdatePattern(ctx, source)
|
||||
}
|
||||
|
||||
// DeletePattern deletes a pattern by ID.
|
||||
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})
|
||||
}
|
||||
|
||||
return result.Error
|
||||
}
|
||||
|
||||
// SearchPatternsFTS performs full-text search on patterns.
|
||||
// Uses raw SQL for FTS5 query.
|
||||
func (s *PatternStore) SearchPatternsFTS(ctx context.Context, searchQuery string, limit int) ([]*models.Pattern, error) {
|
||||
var patterns []Pattern
|
||||
|
||||
// Use raw SQL for FTS5 MATCH query
|
||||
query := `
|
||||
SELECT p.*
|
||||
FROM patterns p
|
||||
JOIN patterns_fts fts ON p.id = fts.rowid
|
||||
WHERE patterns_fts MATCH ?
|
||||
AND p.status = 'active'
|
||||
ORDER BY rank
|
||||
LIMIT ?
|
||||
`
|
||||
|
||||
err := s.db.WithContext(ctx).
|
||||
Raw(query, searchQuery, limit).
|
||||
Scan(&patterns).Error
|
||||
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return toModelPatterns(patterns), nil
|
||||
}
|
||||
|
||||
// PatternStats contains aggregate statistics about patterns.
|
||||
type PatternStats struct {
|
||||
Total int `json:"total"`
|
||||
Active int `json:"active"`
|
||||
Deprecated int `json:"deprecated"`
|
||||
Merged int `json:"merged"`
|
||||
TotalOccurrences int `json:"total_occurrences"`
|
||||
AvgConfidence float64 `json:"avg_confidence"`
|
||||
Bugs int `json:"bugs"`
|
||||
Refactors int `json:"refactors"`
|
||||
Architectures int `json:"architectures"`
|
||||
AntiPatterns int `json:"anti_patterns"`
|
||||
BestPractices int `json:"best_practices"`
|
||||
}
|
||||
|
||||
// GetPatternStats returns statistics about patterns.
|
||||
// Uses raw SQL for complex aggregate query.
|
||||
func (s *PatternStore) GetPatternStats(ctx context.Context) (*PatternStats, error) {
|
||||
var stats PatternStats
|
||||
|
||||
query := `
|
||||
SELECT
|
||||
COUNT(*) as total,
|
||||
COUNT(CASE WHEN status = 'active' THEN 1 END) as active,
|
||||
COUNT(CASE WHEN status = 'deprecated' THEN 1 END) as deprecated,
|
||||
COUNT(CASE WHEN status = 'merged' THEN 1 END) as merged,
|
||||
COALESCE(SUM(frequency), 0) as total_occurrences,
|
||||
COALESCE(AVG(confidence), 0) as avg_confidence,
|
||||
COUNT(CASE WHEN type = 'bug' THEN 1 END) as bugs,
|
||||
COUNT(CASE WHEN type = 'refactor' THEN 1 END) as refactors,
|
||||
COUNT(CASE WHEN type = 'architecture' THEN 1 END) as architectures,
|
||||
COUNT(CASE WHEN type = 'anti-pattern' THEN 1 END) as anti_patterns,
|
||||
COUNT(CASE WHEN type = 'best-practice' THEN 1 END) as best_practices
|
||||
FROM patterns
|
||||
`
|
||||
|
||||
err := s.db.WithContext(ctx).Raw(query).Scan(&stats).Error
|
||||
return &stats, err
|
||||
}
|
||||
|
||||
// IncrementPatternFrequency atomically increments a pattern's frequency and updates last_seen.
|
||||
func (s *PatternStore) IncrementPatternFrequency(ctx context.Context, id int64, project string, observationID int64) error {
|
||||
now := time.Now()
|
||||
|
||||
// Get current pattern
|
||||
pattern, err := s.GetPatternByID(ctx, id)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// Add occurrence
|
||||
pattern.AddOccurrence(project, observationID)
|
||||
pattern.LastSeenAt = now.Format(time.RFC3339)
|
||||
pattern.LastSeenEpoch = now.UnixMilli()
|
||||
|
||||
return s.UpdatePattern(ctx, pattern)
|
||||
}
|
||||
|
||||
// toModelPattern converts a GORM Pattern to a pkg/models Pattern.
|
||||
func toModelPattern(p *Pattern) *models.Pattern {
|
||||
pattern := &models.Pattern{
|
||||
ID: p.ID,
|
||||
Name: p.Name,
|
||||
Type: p.Type,
|
||||
Description: p.Description,
|
||||
Signature: p.Signature,
|
||||
Recommendation: p.Recommendation,
|
||||
Frequency: p.Frequency,
|
||||
Projects: p.Projects,
|
||||
ObservationIDs: p.ObservationIDs,
|
||||
Status: p.Status,
|
||||
MergedIntoID: p.MergedIntoID,
|
||||
Confidence: p.Confidence,
|
||||
LastSeenAt: p.LastSeenAt,
|
||||
LastSeenEpoch: p.LastSeenAtEpoch,
|
||||
CreatedAt: p.CreatedAt,
|
||||
CreatedAtEpoch: p.CreatedAtEpoch,
|
||||
}
|
||||
|
||||
return pattern
|
||||
}
|
||||
|
||||
// toModelPatterns converts a slice of GORM Patterns to pkg/models Patterns.
|
||||
func toModelPatterns(patterns []Pattern) []*models.Pattern {
|
||||
result := make([]*models.Pattern, len(patterns))
|
||||
for i, p := range patterns {
|
||||
result[i] = toModelPattern(&p)
|
||||
}
|
||||
return result
|
||||
}
|
||||
@@ -1,485 +0,0 @@
|
||||
//go:build fts5
|
||||
|
||||
// Package gorm provides GORM-based database operations for claude-mnemonic.
|
||||
package gorm
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"gorm.io/gorm/logger"
|
||||
|
||||
"github.com/lukaszraczylo/claude-mnemonic/pkg/models"
|
||||
)
|
||||
|
||||
func testPatternStore(t *testing.T) (*PatternStore, *Store, func()) {
|
||||
t.Helper()
|
||||
|
||||
tmpDir, err := os.MkdirTemp("", "gorm_pattern_test_*")
|
||||
if err != nil {
|
||||
t.Fatalf("create temp dir: %v", err)
|
||||
}
|
||||
|
||||
dbPath := filepath.Join(tmpDir, "test.db")
|
||||
cfg := Config{
|
||||
Path: dbPath,
|
||||
MaxConns: 4,
|
||||
LogLevel: logger.Silent,
|
||||
}
|
||||
|
||||
store, err := NewStore(cfg)
|
||||
if err != nil {
|
||||
os.RemoveAll(tmpDir)
|
||||
t.Fatalf("NewStore failed: %v", err)
|
||||
}
|
||||
|
||||
patternStore := NewPatternStore(store)
|
||||
|
||||
cleanup := func() {
|
||||
store.Close()
|
||||
os.RemoveAll(tmpDir)
|
||||
}
|
||||
|
||||
return patternStore, store, cleanup
|
||||
}
|
||||
|
||||
func TestPatternStore_StorePattern(t *testing.T) {
|
||||
patternStore, _, cleanup := testPatternStore(t)
|
||||
defer cleanup()
|
||||
ctx := context.Background()
|
||||
|
||||
now := time.Now()
|
||||
pattern := &models.Pattern{
|
||||
Name: "Test Pattern",
|
||||
Type: models.PatternTypeBug,
|
||||
Description: sql.NullString{String: "Test description", Valid: true},
|
||||
Signature: []string{"bug", "error"},
|
||||
Recommendation: sql.NullString{String: "Fix it", Valid: true},
|
||||
Frequency: 1,
|
||||
Projects: []string{"test-project"},
|
||||
ObservationIDs: []int64{1, 2, 3},
|
||||
Status: models.PatternStatusActive,
|
||||
Confidence: 0.8,
|
||||
LastSeenAt: now.Format(time.RFC3339),
|
||||
LastSeenEpoch: now.UnixMilli(),
|
||||
CreatedAt: now.Format(time.RFC3339),
|
||||
CreatedAtEpoch: now.UnixMilli(),
|
||||
}
|
||||
|
||||
id, err := patternStore.StorePattern(ctx, pattern)
|
||||
require.NoError(t, err)
|
||||
assert.Greater(t, id, int64(0))
|
||||
|
||||
// Verify pattern was stored
|
||||
retrieved, err := patternStore.GetPatternByID(ctx, id)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, pattern.Name, retrieved.Name)
|
||||
assert.Equal(t, pattern.Type, retrieved.Type)
|
||||
assert.Equal(t, pattern.Signature, retrieved.Signature)
|
||||
assert.Equal(t, pattern.Frequency, retrieved.Frequency)
|
||||
assert.Equal(t, pattern.Status, retrieved.Status)
|
||||
assert.Equal(t, pattern.Confidence, retrieved.Confidence)
|
||||
}
|
||||
|
||||
func TestPatternStore_UpdatePattern(t *testing.T) {
|
||||
patternStore, _, cleanup := testPatternStore(t)
|
||||
defer cleanup()
|
||||
ctx := context.Background()
|
||||
|
||||
now := time.Now()
|
||||
pattern := &models.Pattern{
|
||||
Name: "Original",
|
||||
Type: models.PatternTypeBug,
|
||||
Frequency: 1,
|
||||
Status: models.PatternStatusActive,
|
||||
Confidence: 0.5,
|
||||
LastSeenAt: now.Format(time.RFC3339),
|
||||
LastSeenEpoch: now.UnixMilli(),
|
||||
CreatedAt: now.Format(time.RFC3339),
|
||||
CreatedAtEpoch: now.UnixMilli(),
|
||||
}
|
||||
|
||||
id, err := patternStore.StorePattern(ctx, pattern)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Update pattern
|
||||
pattern.ID = id
|
||||
pattern.Name = "Updated"
|
||||
pattern.Frequency = 5
|
||||
pattern.Confidence = 0.9
|
||||
|
||||
err = patternStore.UpdatePattern(ctx, pattern)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Verify update
|
||||
retrieved, err := patternStore.GetPatternByID(ctx, id)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "Updated", retrieved.Name)
|
||||
assert.Equal(t, 5, retrieved.Frequency)
|
||||
assert.Equal(t, 0.9, retrieved.Confidence)
|
||||
}
|
||||
|
||||
func TestPatternStore_GetPatternByName(t *testing.T) {
|
||||
patternStore, _, cleanup := testPatternStore(t)
|
||||
defer cleanup()
|
||||
ctx := context.Background()
|
||||
|
||||
now := time.Now()
|
||||
pattern := &models.Pattern{
|
||||
Name: "Unique Pattern",
|
||||
Type: models.PatternTypeRefactor,
|
||||
Frequency: 1,
|
||||
Status: models.PatternStatusActive,
|
||||
Confidence: 0.7,
|
||||
LastSeenAt: now.Format(time.RFC3339),
|
||||
LastSeenEpoch: now.UnixMilli(),
|
||||
CreatedAt: now.Format(time.RFC3339),
|
||||
CreatedAtEpoch: now.UnixMilli(),
|
||||
}
|
||||
|
||||
_, err := patternStore.StorePattern(ctx, pattern)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Retrieve by name
|
||||
retrieved, err := patternStore.GetPatternByName(ctx, "Unique Pattern")
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, retrieved)
|
||||
assert.Equal(t, "Unique Pattern", retrieved.Name)
|
||||
assert.Equal(t, models.PatternTypeRefactor, retrieved.Type)
|
||||
|
||||
// Non-existent pattern
|
||||
notFound, err := patternStore.GetPatternByName(ctx, "Nonexistent")
|
||||
require.NoError(t, err)
|
||||
assert.Nil(t, notFound)
|
||||
}
|
||||
|
||||
func TestPatternStore_GetActivePatterns(t *testing.T) {
|
||||
patternStore, _, cleanup := testPatternStore(t)
|
||||
defer cleanup()
|
||||
ctx := context.Background()
|
||||
|
||||
now := time.Now()
|
||||
|
||||
// Create active patterns
|
||||
for i := 0; i < 3; i++ {
|
||||
pattern := &models.Pattern{
|
||||
Name: "Pattern " + string(rune('A'+i)),
|
||||
Type: models.PatternTypeBug,
|
||||
Frequency: i + 1, // Different frequencies for sorting
|
||||
Status: models.PatternStatusActive,
|
||||
Confidence: 0.8,
|
||||
LastSeenAt: now.Format(time.RFC3339),
|
||||
LastSeenEpoch: now.UnixMilli(),
|
||||
CreatedAt: now.Format(time.RFC3339),
|
||||
CreatedAtEpoch: now.UnixMilli(),
|
||||
}
|
||||
_, err := patternStore.StorePattern(ctx, pattern)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
// Create deprecated pattern (should not be included)
|
||||
deprecatedPattern := &models.Pattern{
|
||||
Name: "Deprecated Pattern",
|
||||
Type: models.PatternTypeBug,
|
||||
Frequency: 100,
|
||||
Status: models.PatternStatusDeprecated,
|
||||
Confidence: 0.9,
|
||||
LastSeenAt: now.Format(time.RFC3339),
|
||||
LastSeenEpoch: now.UnixMilli(),
|
||||
CreatedAt: now.Format(time.RFC3339),
|
||||
CreatedAtEpoch: now.UnixMilli(),
|
||||
}
|
||||
_, err := patternStore.StorePattern(ctx, deprecatedPattern)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Get active patterns
|
||||
patterns, err := patternStore.GetActivePatterns(ctx, 10)
|
||||
require.NoError(t, err)
|
||||
assert.Len(t, patterns, 3) // Only active patterns
|
||||
|
||||
// Verify sorted by frequency DESC
|
||||
assert.Equal(t, 3, patterns[0].Frequency)
|
||||
assert.Equal(t, 2, patterns[1].Frequency)
|
||||
assert.Equal(t, 1, patterns[2].Frequency)
|
||||
}
|
||||
|
||||
func TestPatternStore_GetPatternsByType(t *testing.T) {
|
||||
patternStore, _, cleanup := testPatternStore(t)
|
||||
defer cleanup()
|
||||
ctx := context.Background()
|
||||
|
||||
now := time.Now()
|
||||
|
||||
// Create patterns of different types
|
||||
bugPattern := &models.Pattern{
|
||||
Name: "Bug Pattern",
|
||||
Type: models.PatternTypeBug,
|
||||
Frequency: 1,
|
||||
Status: models.PatternStatusActive,
|
||||
Confidence: 0.8,
|
||||
LastSeenAt: now.Format(time.RFC3339),
|
||||
LastSeenEpoch: now.UnixMilli(),
|
||||
CreatedAt: now.Format(time.RFC3339),
|
||||
CreatedAtEpoch: now.UnixMilli(),
|
||||
}
|
||||
_, err := patternStore.StorePattern(ctx, bugPattern)
|
||||
require.NoError(t, err)
|
||||
|
||||
refactorPattern := &models.Pattern{
|
||||
Name: "Refactor Pattern",
|
||||
Type: models.PatternTypeRefactor,
|
||||
Frequency: 1,
|
||||
Status: models.PatternStatusActive,
|
||||
Confidence: 0.7,
|
||||
LastSeenAt: now.Format(time.RFC3339),
|
||||
LastSeenEpoch: now.UnixMilli(),
|
||||
CreatedAt: now.Format(time.RFC3339),
|
||||
CreatedAtEpoch: now.UnixMilli(),
|
||||
}
|
||||
_, err = patternStore.StorePattern(ctx, refactorPattern)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Get only bug patterns
|
||||
bugPatterns, err := patternStore.GetPatternsByType(ctx, models.PatternTypeBug, 10)
|
||||
require.NoError(t, err)
|
||||
assert.Len(t, bugPatterns, 1)
|
||||
assert.Equal(t, "Bug Pattern", bugPatterns[0].Name)
|
||||
assert.Equal(t, models.PatternTypeBug, bugPatterns[0].Type)
|
||||
|
||||
// Get only refactor patterns
|
||||
refactorPatterns, err := patternStore.GetPatternsByType(ctx, models.PatternTypeRefactor, 10)
|
||||
require.NoError(t, err)
|
||||
assert.Len(t, refactorPatterns, 1)
|
||||
assert.Equal(t, "Refactor Pattern", refactorPatterns[0].Name)
|
||||
}
|
||||
|
||||
func TestPatternStore_MarkPatternDeprecated(t *testing.T) {
|
||||
patternStore, _, cleanup := testPatternStore(t)
|
||||
defer cleanup()
|
||||
ctx := context.Background()
|
||||
|
||||
now := time.Now()
|
||||
pattern := &models.Pattern{
|
||||
Name: "To Deprecate",
|
||||
Type: models.PatternTypeBug,
|
||||
Frequency: 1,
|
||||
Status: models.PatternStatusActive,
|
||||
Confidence: 0.5,
|
||||
LastSeenAt: now.Format(time.RFC3339),
|
||||
LastSeenEpoch: now.UnixMilli(),
|
||||
CreatedAt: now.Format(time.RFC3339),
|
||||
CreatedAtEpoch: now.UnixMilli(),
|
||||
}
|
||||
|
||||
id, err := patternStore.StorePattern(ctx, pattern)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Mark as deprecated
|
||||
err = patternStore.MarkPatternDeprecated(ctx, id)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Verify status changed
|
||||
retrieved, err := patternStore.GetPatternByID(ctx, id)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, models.PatternStatusDeprecated, retrieved.Status)
|
||||
}
|
||||
|
||||
func TestPatternStore_MergePatterns(t *testing.T) {
|
||||
patternStore, _, cleanup := testPatternStore(t)
|
||||
defer cleanup()
|
||||
ctx := context.Background()
|
||||
|
||||
now := time.Now()
|
||||
|
||||
// Create source pattern
|
||||
source := &models.Pattern{
|
||||
Name: "Source Pattern",
|
||||
Type: models.PatternTypeBug,
|
||||
Frequency: 5,
|
||||
Projects: []string{"project-a", "project-b"},
|
||||
ObservationIDs: []int64{1, 2, 3},
|
||||
Status: models.PatternStatusActive,
|
||||
Confidence: 0.7,
|
||||
LastSeenAt: now.Format(time.RFC3339),
|
||||
LastSeenEpoch: now.UnixMilli(),
|
||||
CreatedAt: now.Format(time.RFC3339),
|
||||
CreatedAtEpoch: now.UnixMilli(),
|
||||
}
|
||||
sourceID, err := patternStore.StorePattern(ctx, source)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Create target pattern
|
||||
target := &models.Pattern{
|
||||
Name: "Target Pattern",
|
||||
Type: models.PatternTypeBug,
|
||||
Frequency: 10,
|
||||
Projects: []string{"project-b", "project-c"},
|
||||
ObservationIDs: []int64{3, 4, 5},
|
||||
Status: models.PatternStatusActive,
|
||||
Confidence: 0.8,
|
||||
LastSeenAt: now.Format(time.RFC3339),
|
||||
LastSeenEpoch: now.UnixMilli(),
|
||||
CreatedAt: now.Format(time.RFC3339),
|
||||
CreatedAtEpoch: now.UnixMilli(),
|
||||
}
|
||||
targetID, err := patternStore.StorePattern(ctx, target)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Merge source into target
|
||||
err = patternStore.MergePatterns(ctx, sourceID, targetID)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Verify target was updated
|
||||
mergedTarget, err := patternStore.GetPatternByID(ctx, targetID)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 15, mergedTarget.Frequency) // 5 + 10
|
||||
assert.ElementsMatch(t, []string{"project-a", "project-b", "project-c"}, mergedTarget.Projects)
|
||||
assert.ElementsMatch(t, []int64{1, 2, 3, 4, 5}, mergedTarget.ObservationIDs)
|
||||
|
||||
// Verify source was marked as merged
|
||||
mergedSource, err := patternStore.GetPatternByID(ctx, sourceID)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, models.PatternStatusMerged, mergedSource.Status)
|
||||
assert.True(t, mergedSource.MergedIntoID.Valid)
|
||||
assert.Equal(t, targetID, mergedSource.MergedIntoID.Int64)
|
||||
}
|
||||
|
||||
func TestPatternStore_DeletePattern(t *testing.T) {
|
||||
patternStore, _, cleanup := testPatternStore(t)
|
||||
defer cleanup()
|
||||
ctx := context.Background()
|
||||
|
||||
now := time.Now()
|
||||
pattern := &models.Pattern{
|
||||
Name: "To Delete",
|
||||
Type: models.PatternTypeBug,
|
||||
Frequency: 1,
|
||||
Status: models.PatternStatusActive,
|
||||
Confidence: 0.5,
|
||||
LastSeenAt: now.Format(time.RFC3339),
|
||||
LastSeenEpoch: now.UnixMilli(),
|
||||
CreatedAt: now.Format(time.RFC3339),
|
||||
CreatedAtEpoch: now.UnixMilli(),
|
||||
}
|
||||
|
||||
id, err := patternStore.StorePattern(ctx, pattern)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Delete pattern
|
||||
err = patternStore.DeletePattern(ctx, id)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Verify deleted
|
||||
deleted, err := patternStore.GetPatternByID(ctx, id)
|
||||
require.NoError(t, err)
|
||||
assert.Nil(t, deleted)
|
||||
}
|
||||
|
||||
func TestPatternStore_IncrementPatternFrequency(t *testing.T) {
|
||||
patternStore, _, cleanup := testPatternStore(t)
|
||||
defer cleanup()
|
||||
ctx := context.Background()
|
||||
|
||||
now := time.Now()
|
||||
pattern := &models.Pattern{
|
||||
Name: "Frequency Test",
|
||||
Type: models.PatternTypeBug,
|
||||
Frequency: 1,
|
||||
Projects: []string{"project-a"},
|
||||
ObservationIDs: []int64{},
|
||||
Status: models.PatternStatusActive,
|
||||
Confidence: 0.7,
|
||||
LastSeenAt: now.Format(time.RFC3339),
|
||||
LastSeenEpoch: now.UnixMilli(),
|
||||
CreatedAt: now.Format(time.RFC3339),
|
||||
CreatedAtEpoch: now.UnixMilli(),
|
||||
}
|
||||
|
||||
id, err := patternStore.StorePattern(ctx, pattern)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Increment frequency with new project and observation
|
||||
err = patternStore.IncrementPatternFrequency(ctx, id, "project-b", 42)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Verify frequency incremented and new data added
|
||||
updated, err := patternStore.GetPatternByID(ctx, id)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 2, updated.Frequency)
|
||||
assert.ElementsMatch(t, []string{"project-a", "project-b"}, updated.Projects)
|
||||
assert.Contains(t, updated.ObservationIDs, int64(42))
|
||||
|
||||
// Last seen should be updated (rough check - within last 5 seconds)
|
||||
updatedTime, _ := time.Parse(time.RFC3339, updated.LastSeenAt)
|
||||
assert.WithinDuration(t, time.Now(), updatedTime, 5*time.Second)
|
||||
}
|
||||
|
||||
func TestPatternStore_GetPatternStats(t *testing.T) {
|
||||
patternStore, _, cleanup := testPatternStore(t)
|
||||
defer cleanup()
|
||||
ctx := context.Background()
|
||||
|
||||
now := time.Now()
|
||||
|
||||
// Create patterns with different statuses and types
|
||||
patterns := []*models.Pattern{
|
||||
{
|
||||
Name: "Bug 1",
|
||||
Type: models.PatternTypeBug,
|
||||
Frequency: 10,
|
||||
Status: models.PatternStatusActive,
|
||||
Confidence: 0.8,
|
||||
LastSeenAt: now.Format(time.RFC3339),
|
||||
LastSeenEpoch: now.UnixMilli(),
|
||||
CreatedAt: now.Format(time.RFC3339),
|
||||
CreatedAtEpoch: now.UnixMilli(),
|
||||
},
|
||||
{
|
||||
Name: "Refactor 1",
|
||||
Type: models.PatternTypeRefactor,
|
||||
Frequency: 5,
|
||||
Status: models.PatternStatusActive,
|
||||
Confidence: 0.7,
|
||||
LastSeenAt: now.Format(time.RFC3339),
|
||||
LastSeenEpoch: now.UnixMilli(),
|
||||
CreatedAt: now.Format(time.RFC3339),
|
||||
CreatedAtEpoch: now.UnixMilli(),
|
||||
},
|
||||
{
|
||||
Name: "Deprecated 1",
|
||||
Type: models.PatternTypeBestPractice,
|
||||
Frequency: 3,
|
||||
Status: models.PatternStatusDeprecated,
|
||||
Confidence: 0.6,
|
||||
LastSeenAt: now.Format(time.RFC3339),
|
||||
LastSeenEpoch: now.UnixMilli(),
|
||||
CreatedAt: now.Format(time.RFC3339),
|
||||
CreatedAtEpoch: now.UnixMilli(),
|
||||
},
|
||||
}
|
||||
|
||||
for _, p := range patterns {
|
||||
_, err := patternStore.StorePattern(ctx, p)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
// Get stats
|
||||
stats, err := patternStore.GetPatternStats(ctx)
|
||||
require.NoError(t, err)
|
||||
|
||||
assert.Equal(t, 3, stats.Total)
|
||||
assert.Equal(t, 2, stats.Active)
|
||||
assert.Equal(t, 1, stats.Deprecated)
|
||||
assert.Equal(t, 0, stats.Merged)
|
||||
assert.Equal(t, 18, stats.TotalOccurrences) // 10 + 5 + 3
|
||||
assert.InDelta(t, 0.7, stats.AvgConfidence, 0.05) // (0.8 + 0.7 + 0.6) / 3
|
||||
assert.Equal(t, 1, stats.Bugs)
|
||||
assert.Equal(t, 1, stats.Refactors)
|
||||
assert.Equal(t, 1, stats.BestPractices)
|
||||
}
|
||||
@@ -1,317 +0,0 @@
|
||||
// Package gorm provides GORM-based database operations for claude-mnemonic.
|
||||
package gorm
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"time"
|
||||
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/clause"
|
||||
|
||||
"github.com/lukaszraczylo/claude-mnemonic/pkg/models"
|
||||
)
|
||||
|
||||
// PromptCleanupFunc is a callback for when prompts are cleaned up.
|
||||
// Receives the IDs of deleted prompts for downstream cleanup (e.g., vector DB).
|
||||
type PromptCleanupFunc func(ctx context.Context, deletedIDs []int64)
|
||||
|
||||
// MaxPromptsGlobal is the hard limit of prompts across all projects.
|
||||
const MaxPromptsGlobal = 500
|
||||
|
||||
// PromptStore provides user prompt-related database operations using GORM.
|
||||
type PromptStore struct {
|
||||
db *gorm.DB
|
||||
cleanupFunc PromptCleanupFunc
|
||||
}
|
||||
|
||||
// NewPromptStore creates a new prompt store.
|
||||
func NewPromptStore(store *Store, cleanupFunc PromptCleanupFunc) *PromptStore {
|
||||
return &PromptStore{
|
||||
db: store.DB,
|
||||
cleanupFunc: cleanupFunc,
|
||||
}
|
||||
}
|
||||
|
||||
// SetCleanupFunc sets the callback for when prompts are deleted during cleanup.
|
||||
func (s *PromptStore) SetCleanupFunc(fn PromptCleanupFunc) {
|
||||
s.cleanupFunc = fn
|
||||
}
|
||||
|
||||
// SaveUserPromptWithMatches saves a user prompt with matched observation count.
|
||||
// Uses INSERT OR IGNORE to be idempotent - duplicate (session, prompt_number) pairs are silently ignored.
|
||||
// This prevents duplicate prompts when the user-prompt hook fires multiple times.
|
||||
func (s *PromptStore) SaveUserPromptWithMatches(ctx context.Context, claudeSessionID string, promptNumber int, promptText string, matchedObservations int) (int64, error) {
|
||||
now := time.Now()
|
||||
|
||||
prompt := &UserPrompt{
|
||||
ClaudeSessionID: claudeSessionID,
|
||||
PromptNumber: promptNumber,
|
||||
PromptText: promptText,
|
||||
MatchedObservations: matchedObservations,
|
||||
CreatedAt: now.Format(time.RFC3339),
|
||||
CreatedAtEpoch: now.UnixMilli(),
|
||||
}
|
||||
|
||||
// INSERT OR IGNORE using OnConflict
|
||||
result := s.db.WithContext(ctx).
|
||||
Clauses(clause.OnConflict{
|
||||
Columns: []clause.Column{{Name: "claude_session_id"}, {Name: "prompt_number"}},
|
||||
DoNothing: true,
|
||||
}).
|
||||
Create(prompt)
|
||||
|
||||
if result.Error != nil {
|
||||
return 0, result.Error
|
||||
}
|
||||
|
||||
// If RowsAffected is 0, the insert was ignored (duplicate) - fetch the existing ID
|
||||
if result.RowsAffected == 0 {
|
||||
var existing UserPrompt
|
||||
err := s.db.Where("claude_session_id = ? AND prompt_number = ?", claudeSessionID, promptNumber).
|
||||
First(&existing).Error
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
// Return existing ID without triggering cleanup (already handled when first inserted)
|
||||
return existing.ID, nil
|
||||
}
|
||||
|
||||
// Cleanup old prompts beyond the global limit (async to not block handler)
|
||||
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)
|
||||
}
|
||||
}()
|
||||
|
||||
return prompt.ID, nil
|
||||
}
|
||||
|
||||
// CleanupOldPrompts deletes prompts beyond the global limit.
|
||||
// Keeps the most recent MaxPromptsGlobal prompts.
|
||||
// Returns the IDs of deleted prompts for downstream cleanup (e.g., vector DB).
|
||||
func (s *PromptStore) CleanupOldPrompts(ctx context.Context) ([]int64, error) {
|
||||
// Use a transaction to prevent TOCTOU race condition
|
||||
var idsToDelete []int64
|
||||
|
||||
err := s.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
// Find IDs to keep (most recent MaxPromptsGlobal)
|
||||
var idsToKeep []int64
|
||||
err := tx.Model(&UserPrompt{}).
|
||||
Order("created_at_epoch DESC").
|
||||
Limit(MaxPromptsGlobal).
|
||||
Pluck("id", &idsToKeep).Error
|
||||
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if len(idsToKeep) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
// Find IDs to delete (all IDs not in the keep list)
|
||||
// This happens in the same transaction to prevent race conditions
|
||||
err = tx.Model(&UserPrompt{}).
|
||||
Where("id NOT IN ?", idsToKeep).
|
||||
Pluck("id", &idsToDelete).Error
|
||||
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if len(idsToDelete) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
// Delete the prompts
|
||||
return tx.Delete(&UserPrompt{}, idsToDelete).Error
|
||||
})
|
||||
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return idsToDelete, nil
|
||||
}
|
||||
|
||||
// GetPromptsByIDs retrieves user prompts by a list of IDs.
|
||||
func (s *PromptStore) GetPromptsByIDs(ctx context.Context, ids []int64, orderBy string, limit int) ([]*models.UserPromptWithSession, error) {
|
||||
if len(ids) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
var results []struct {
|
||||
Project sql.NullString `gorm:"column:project"`
|
||||
SDKSessionID sql.NullString `gorm:"column:sdk_session_id"`
|
||||
UserPrompt
|
||||
}
|
||||
|
||||
query := s.db.WithContext(ctx).
|
||||
Table("user_prompts up").
|
||||
Select("up.id, up.claude_session_id, up.prompt_number, up.prompt_text, "+
|
||||
"COALESCE(up.matched_observations, 0) as matched_observations, "+
|
||||
"up.created_at, up.created_at_epoch, "+
|
||||
"COALESCE(s.project, '') as project, "+
|
||||
"COALESCE(s.sdk_session_id, '') as sdk_session_id").
|
||||
Joins("LEFT JOIN sdk_sessions s ON up.claude_session_id = s.claude_session_id").
|
||||
Where("up.id IN ?", ids)
|
||||
|
||||
// Apply ordering
|
||||
switch orderBy {
|
||||
case "date_asc":
|
||||
query = query.Order("up.created_at_epoch ASC")
|
||||
case "date_desc", "default", "":
|
||||
query = query.Order("up.created_at_epoch DESC")
|
||||
}
|
||||
|
||||
// Apply limit
|
||||
if limit > 0 {
|
||||
query = query.Limit(limit)
|
||||
}
|
||||
|
||||
err := query.Scan(&results).Error
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return toModelUserPromptsWithSession(results), nil
|
||||
}
|
||||
|
||||
// GetAllRecentUserPrompts retrieves recent user prompts across all projects.
|
||||
func (s *PromptStore) GetAllRecentUserPrompts(ctx context.Context, limit int) ([]*models.UserPromptWithSession, error) {
|
||||
var results []struct {
|
||||
Project sql.NullString `gorm:"column:project"`
|
||||
SDKSessionID sql.NullString `gorm:"column:sdk_session_id"`
|
||||
UserPrompt
|
||||
}
|
||||
|
||||
query := s.db.WithContext(ctx).
|
||||
Table("user_prompts up").
|
||||
Select("up.id, up.claude_session_id, up.prompt_number, up.prompt_text, " +
|
||||
"COALESCE(up.matched_observations, 0) as matched_observations, " +
|
||||
"up.created_at, up.created_at_epoch, " +
|
||||
"COALESCE(s.project, '') as project, " +
|
||||
"COALESCE(s.sdk_session_id, '') as sdk_session_id").
|
||||
Joins("LEFT JOIN sdk_sessions s ON up.claude_session_id = s.claude_session_id").
|
||||
Order("up.created_at_epoch DESC").
|
||||
Limit(limit)
|
||||
|
||||
err := query.Scan(&results).Error
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return toModelUserPromptsWithSession(results), nil
|
||||
}
|
||||
|
||||
// GetAllPrompts retrieves all user prompts (for vector rebuild).
|
||||
func (s *PromptStore) GetAllPrompts(ctx context.Context) ([]*models.UserPromptWithSession, error) {
|
||||
var results []struct {
|
||||
Project sql.NullString `gorm:"column:project"`
|
||||
SDKSessionID sql.NullString `gorm:"column:sdk_session_id"`
|
||||
UserPrompt
|
||||
}
|
||||
|
||||
query := s.db.WithContext(ctx).
|
||||
Table("user_prompts up").
|
||||
Select("up.id, up.claude_session_id, up.prompt_number, up.prompt_text, " +
|
||||
"COALESCE(up.matched_observations, 0) as matched_observations, " +
|
||||
"up.created_at, up.created_at_epoch, " +
|
||||
"COALESCE(s.project, '') as project, " +
|
||||
"COALESCE(s.sdk_session_id, '') as sdk_session_id").
|
||||
Joins("LEFT JOIN sdk_sessions s ON up.claude_session_id = s.claude_session_id").
|
||||
Order("up.id")
|
||||
|
||||
err := query.Scan(&results).Error
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return toModelUserPromptsWithSession(results), nil
|
||||
}
|
||||
|
||||
// FindRecentPromptByText finds a recent prompt by exact text match within a time window.
|
||||
// Returns (promptID, promptNumber, found).
|
||||
func (s *PromptStore) FindRecentPromptByText(ctx context.Context, claudeSessionID, 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("claude_session_id = ? AND prompt_text = ? AND created_at_epoch >= ?",
|
||||
claudeSessionID, 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 {
|
||||
Project sql.NullString `gorm:"column:project"`
|
||||
SDKSessionID sql.NullString `gorm:"column:sdk_session_id"`
|
||||
UserPrompt
|
||||
}
|
||||
|
||||
query := s.db.WithContext(ctx).
|
||||
Table("user_prompts up").
|
||||
Select("up.id, up.claude_session_id, up.prompt_number, up.prompt_text, "+
|
||||
"COALESCE(up.matched_observations, 0) as matched_observations, "+
|
||||
"up.created_at, up.created_at_epoch, "+
|
||||
"COALESCE(s.project, '') as project, "+
|
||||
"COALESCE(s.sdk_session_id, '') as sdk_session_id").
|
||||
Joins("LEFT JOIN sdk_sessions s ON up.claude_session_id = s.claude_session_id").
|
||||
Where("s.project = ?", project).
|
||||
Order("up.created_at_epoch DESC").
|
||||
Limit(limit)
|
||||
|
||||
err := query.Scan(&results).Error
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return toModelUserPromptsWithSession(results), nil
|
||||
}
|
||||
|
||||
// toModelUserPromptsWithSession converts query results to pkg/models.UserPromptWithSession.
|
||||
func toModelUserPromptsWithSession(results []struct {
|
||||
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 {
|
||||
project := ""
|
||||
if r.Project.Valid {
|
||||
project = r.Project.String
|
||||
}
|
||||
|
||||
sdkSessionID := ""
|
||||
if r.SDKSessionID.Valid {
|
||||
sdkSessionID = r.SDKSessionID.String
|
||||
}
|
||||
|
||||
prompts[i] = &models.UserPromptWithSession{
|
||||
UserPrompt: models.UserPrompt{
|
||||
ID: r.ID,
|
||||
ClaudeSessionID: r.ClaudeSessionID,
|
||||
PromptNumber: r.PromptNumber,
|
||||
PromptText: r.PromptText,
|
||||
MatchedObservations: r.MatchedObservations,
|
||||
CreatedAt: r.CreatedAt,
|
||||
CreatedAtEpoch: r.CreatedAtEpoch,
|
||||
},
|
||||
Project: project,
|
||||
SDKSessionID: sdkSessionID,
|
||||
}
|
||||
}
|
||||
return prompts
|
||||
}
|
||||
@@ -1,396 +0,0 @@
|
||||
//go:build fts5
|
||||
|
||||
// Package gorm provides GORM-based database operations for claude-mnemonic.
|
||||
package gorm
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"gorm.io/gorm/logger"
|
||||
)
|
||||
|
||||
// testPromptStore creates a PromptStore with a temporary database for testing.
|
||||
func testPromptStore(t *testing.T) (*PromptStore, *Store, func()) {
|
||||
t.Helper()
|
||||
|
||||
tmpDir, err := os.MkdirTemp("", "gorm_prompt_test_*")
|
||||
if err != nil {
|
||||
t.Fatalf("create temp dir: %v", err)
|
||||
}
|
||||
|
||||
dbPath := filepath.Join(tmpDir, "test.db")
|
||||
cfg := Config{
|
||||
Path: dbPath,
|
||||
MaxConns: 4,
|
||||
LogLevel: logger.Silent,
|
||||
}
|
||||
|
||||
store, err := NewStore(cfg)
|
||||
if err != nil {
|
||||
os.RemoveAll(tmpDir)
|
||||
t.Fatalf("NewStore failed: %v", err)
|
||||
}
|
||||
|
||||
promptStore := NewPromptStore(store, nil)
|
||||
|
||||
cleanup := func() {
|
||||
store.Close()
|
||||
os.RemoveAll(tmpDir)
|
||||
}
|
||||
|
||||
return promptStore, store, cleanup
|
||||
}
|
||||
|
||||
func TestPromptStore_SaveUserPromptWithMatches(t *testing.T) {
|
||||
promptStore, store, cleanup := testPromptStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Create a session first
|
||||
sessionStore := NewSessionStore(store)
|
||||
_, err := sessionStore.CreateSDKSession(ctx, "claude-1", "test-project", "")
|
||||
require.NoError(t, err)
|
||||
|
||||
// Save a prompt
|
||||
id, err := promptStore.SaveUserPromptWithMatches(ctx, "claude-1", 1, "What is the codebase structure?", 5)
|
||||
require.NoError(t, err)
|
||||
assert.Greater(t, id, int64(0))
|
||||
}
|
||||
|
||||
func TestPromptStore_SaveUserPromptWithMatches_Idempotency(t *testing.T) {
|
||||
promptStore, _, cleanup := testPromptStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Save the same prompt twice (same claudeSessionID + promptNumber)
|
||||
id1, err := promptStore.SaveUserPromptWithMatches(ctx, "claude-1", 1, "Test prompt", 3)
|
||||
require.NoError(t, err)
|
||||
|
||||
id2, err := promptStore.SaveUserPromptWithMatches(ctx, "claude-1", 1, "Different text", 5)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Should return the same ID (INSERT OR IGNORE)
|
||||
assert.Equal(t, id1, id2, "Duplicate prompts should return same ID")
|
||||
}
|
||||
|
||||
func TestPromptStore_SaveUserPromptWithMatches_AsyncCleanup(t *testing.T) {
|
||||
promptStore, _, cleanup := testPromptStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Track cleanup calls
|
||||
var cleanupMutex sync.Mutex
|
||||
cleanupCalled := false
|
||||
var cleanupIDs []int64
|
||||
|
||||
cleanupFunc := func(ctx context.Context, deletedIDs []int64) {
|
||||
cleanupMutex.Lock()
|
||||
defer cleanupMutex.Unlock()
|
||||
cleanupCalled = true
|
||||
cleanupIDs = deletedIDs
|
||||
}
|
||||
|
||||
promptStore.cleanupFunc = cleanupFunc
|
||||
|
||||
// Save prompts beyond the global limit (MaxPromptsGlobal = 500)
|
||||
// Insert with slower pacing to avoid database lock contention
|
||||
for i := 0; i < 505; i++ {
|
||||
_, err := promptStore.SaveUserPromptWithMatches(ctx, "claude-1", i+1, "Prompt", 1)
|
||||
require.NoError(t, err)
|
||||
if i > 500 {
|
||||
time.Sleep(5 * time.Millisecond) // Slow down after hitting limit
|
||||
}
|
||||
}
|
||||
|
||||
// Wait longer for async cleanup to complete
|
||||
time.Sleep(500 * time.Millisecond)
|
||||
|
||||
// Verify cleanup was called
|
||||
cleanupMutex.Lock()
|
||||
defer cleanupMutex.Unlock()
|
||||
assert.True(t, cleanupCalled, "Cleanup function should have been called")
|
||||
assert.NotEmpty(t, cleanupIDs, "Cleanup should have deleted some prompts")
|
||||
}
|
||||
|
||||
func TestPromptStore_CleanupOldPrompts(t *testing.T) {
|
||||
promptStore, _, cleanup := testPromptStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Save prompts beyond the limit
|
||||
// Async cleanup should fire after each insert beyond 500
|
||||
for i := 0; i < 505; i++ {
|
||||
_, err := promptStore.SaveUserPromptWithMatches(ctx, "claude-1", i+1, "Prompt", 1)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
// Wait for all async cleanups to complete
|
||||
time.Sleep(1 * time.Second)
|
||||
|
||||
// After async cleanup, we should have at most 500 prompts
|
||||
remaining, err := promptStore.GetAllPrompts(ctx)
|
||||
require.NoError(t, err)
|
||||
assert.LessOrEqual(t, len(remaining), MaxPromptsGlobal, "Should have at most %d prompts after async cleanup", MaxPromptsGlobal)
|
||||
}
|
||||
|
||||
func TestPromptStore_GetPromptsByIDs(t *testing.T) {
|
||||
promptStore, _, cleanup := testPromptStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Save multiple prompts
|
||||
var ids []int64
|
||||
for i := 1; i <= 3; i++ {
|
||||
id, err := promptStore.SaveUserPromptWithMatches(ctx, "claude-1", i, "Prompt", i)
|
||||
require.NoError(t, err)
|
||||
ids = append(ids, id)
|
||||
time.Sleep(10 * time.Millisecond) // Ensure different timestamps
|
||||
}
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
orderBy string
|
||||
expected []int64
|
||||
}{
|
||||
{
|
||||
name: "Default ordering - date desc",
|
||||
orderBy: "default",
|
||||
expected: []int64{ids[2], ids[1], ids[0]}, // Newest to oldest
|
||||
},
|
||||
{
|
||||
name: "Date ascending",
|
||||
orderBy: "date_asc",
|
||||
expected: []int64{ids[0], ids[1], ids[2]}, // Oldest to newest
|
||||
},
|
||||
{
|
||||
name: "Date descending",
|
||||
orderBy: "date_desc",
|
||||
expected: []int64{ids[2], ids[1], ids[0]}, // Newest to oldest
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
prompts, err := promptStore.GetPromptsByIDs(ctx, ids, tt.orderBy, 10)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, prompts, 3)
|
||||
|
||||
// Verify ordering
|
||||
for i, prompt := range prompts {
|
||||
assert.Equal(t, tt.expected[i], prompt.ID, "Position %d should have ID %d", i, tt.expected[i])
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestPromptStore_GetPromptsByIDs_Limit(t *testing.T) {
|
||||
promptStore, _, cleanup := testPromptStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Save multiple prompts
|
||||
var ids []int64
|
||||
for i := 1; i <= 5; i++ {
|
||||
id, err := promptStore.SaveUserPromptWithMatches(ctx, "claude-1", i, "Prompt", i)
|
||||
require.NoError(t, err)
|
||||
ids = append(ids, id)
|
||||
}
|
||||
|
||||
// Get with limit
|
||||
prompts, err := promptStore.GetPromptsByIDs(ctx, ids, "default", 3)
|
||||
require.NoError(t, err)
|
||||
assert.Len(t, prompts, 3)
|
||||
}
|
||||
|
||||
func TestPromptStore_GetPromptsByIDs_EmptyInput(t *testing.T) {
|
||||
promptStore, _, cleanup := testPromptStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Get with empty IDs
|
||||
prompts, err := promptStore.GetPromptsByIDs(ctx, []int64{}, "default", 10)
|
||||
require.NoError(t, err)
|
||||
assert.Nil(t, prompts)
|
||||
}
|
||||
|
||||
func TestPromptStore_GetPromptsByIDs_WithSession(t *testing.T) {
|
||||
promptStore, store, cleanup := testPromptStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Create a session
|
||||
sessionStore := NewSessionStore(store)
|
||||
_, err := sessionStore.CreateSDKSession(ctx, "claude-1", "test-project", "")
|
||||
require.NoError(t, err)
|
||||
|
||||
// Save a prompt
|
||||
id, err := promptStore.SaveUserPromptWithMatches(ctx, "claude-1", 1, "Test prompt", 5)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Get with session join
|
||||
prompts, err := promptStore.GetPromptsByIDs(ctx, []int64{id}, "default", 10)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, prompts, 1)
|
||||
|
||||
// Verify session data is populated
|
||||
assert.Equal(t, "test-project", prompts[0].Project)
|
||||
assert.NotEmpty(t, prompts[0].SDKSessionID)
|
||||
}
|
||||
|
||||
func TestPromptStore_GetAllRecentUserPrompts(t *testing.T) {
|
||||
promptStore, _, cleanup := testPromptStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Save prompts across multiple sessions with timestamps
|
||||
for i := 1; i <= 3; i++ {
|
||||
_, err := promptStore.SaveUserPromptWithMatches(ctx, "claude-1", i, "Prompt A", i)
|
||||
require.NoError(t, err)
|
||||
time.Sleep(10 * time.Millisecond) // Ensure different timestamps
|
||||
}
|
||||
|
||||
for i := 1; i <= 2; i++ {
|
||||
_, err := promptStore.SaveUserPromptWithMatches(ctx, "claude-2", i, "Prompt B", i)
|
||||
require.NoError(t, err)
|
||||
time.Sleep(10 * time.Millisecond) // Ensure different timestamps
|
||||
}
|
||||
|
||||
// Get all recent prompts
|
||||
prompts, err := promptStore.GetAllRecentUserPrompts(ctx, 10)
|
||||
require.NoError(t, err)
|
||||
assert.Len(t, prompts, 5)
|
||||
|
||||
// Verify ordering (most recent first) - last inserted should be first
|
||||
assert.Equal(t, "claude-2", prompts[0].ClaudeSessionID)
|
||||
assert.Equal(t, 2, prompts[0].PromptNumber)
|
||||
}
|
||||
|
||||
func TestPromptStore_GetAllPrompts(t *testing.T) {
|
||||
promptStore, _, cleanup := testPromptStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Save prompts
|
||||
for i := 1; i <= 5; i++ {
|
||||
_, err := promptStore.SaveUserPromptWithMatches(ctx, "claude-1", i, "Prompt", i)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
// Wait for any async cleanup to complete (longer wait for race detector)
|
||||
time.Sleep(500 * time.Millisecond)
|
||||
|
||||
// Get all prompts (for vector rebuild)
|
||||
prompts, err := promptStore.GetAllPrompts(ctx)
|
||||
require.NoError(t, err)
|
||||
assert.Len(t, prompts, 5)
|
||||
|
||||
// Verify ordering by ID
|
||||
for i := 0; i < len(prompts)-1; i++ {
|
||||
assert.Less(t, prompts[i].ID, prompts[i+1].ID)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPromptStore_FindRecentPromptByText(t *testing.T) {
|
||||
promptStore, _, cleanup := testPromptStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Save a prompt
|
||||
id, err := promptStore.SaveUserPromptWithMatches(ctx, "claude-1", 1, "What is the architecture?", 3)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Find by exact text match within time window
|
||||
foundID, foundNumber, found := promptStore.FindRecentPromptByText(ctx, "claude-1", "What is the architecture?", 60)
|
||||
assert.True(t, found, "Should find the prompt")
|
||||
assert.Equal(t, id, foundID)
|
||||
assert.Equal(t, 1, foundNumber)
|
||||
|
||||
// Try to find with different text
|
||||
_, _, notFound := promptStore.FindRecentPromptByText(ctx, "claude-1", "Different text", 60)
|
||||
assert.False(t, notFound, "Should not find a different prompt")
|
||||
|
||||
// Try to find outside time window
|
||||
time.Sleep(100 * time.Millisecond)
|
||||
_, _, notFound = promptStore.FindRecentPromptByText(ctx, "claude-1", "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()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Create sessions for different projects
|
||||
sessionStore := NewSessionStore(store)
|
||||
_, err := sessionStore.CreateSDKSession(ctx, "claude-1", "project-a", "")
|
||||
require.NoError(t, err)
|
||||
_, err = sessionStore.CreateSDKSession(ctx, "claude-2", "project-b", "")
|
||||
require.NoError(t, err)
|
||||
|
||||
// Save prompts for project-a
|
||||
for i := 1; i <= 3; i++ {
|
||||
_, err := promptStore.SaveUserPromptWithMatches(ctx, "claude-1", i, "Prompt A", i)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
// Save prompts for project-b
|
||||
for i := 1; i <= 2; i++ {
|
||||
_, err := promptStore.SaveUserPromptWithMatches(ctx, "claude-2", i, "Prompt B", i)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
// Get prompts for project-a
|
||||
prompts, err := promptStore.GetRecentUserPromptsByProject(ctx, "project-a", 10)
|
||||
require.NoError(t, err)
|
||||
assert.Len(t, prompts, 3)
|
||||
|
||||
// Verify all prompts are from project-a
|
||||
for _, prompt := range prompts {
|
||||
assert.Equal(t, "project-a", prompt.Project)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPromptStore_GetRecentUserPromptsByProject_Limit(t *testing.T) {
|
||||
promptStore, store, cleanup := testPromptStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Create session
|
||||
sessionStore := NewSessionStore(store)
|
||||
_, err := sessionStore.CreateSDKSession(ctx, "claude-1", "test-project", "")
|
||||
require.NoError(t, err)
|
||||
|
||||
// Save multiple prompts
|
||||
for i := 1; i <= 10; i++ {
|
||||
_, err := promptStore.SaveUserPromptWithMatches(ctx, "claude-1", i, "Prompt", i)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
// Wait for any async cleanup to complete
|
||||
time.Sleep(100 * time.Millisecond)
|
||||
|
||||
// Get with limit
|
||||
prompts, err := promptStore.GetRecentUserPromptsByProject(ctx, "test-project", 5)
|
||||
require.NoError(t, err)
|
||||
assert.Len(t, prompts, 5)
|
||||
}
|
||||
@@ -1,383 +0,0 @@
|
||||
// Package gorm provides GORM-based database operations for claude-mnemonic.
|
||||
package gorm
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/clause"
|
||||
|
||||
"github.com/lukaszraczylo/claude-mnemonic/pkg/models"
|
||||
)
|
||||
|
||||
// RelationStore provides relation-related database operations using GORM.
|
||||
type RelationStore struct {
|
||||
db *gorm.DB
|
||||
}
|
||||
|
||||
// NewRelationStore creates a new relation store.
|
||||
func NewRelationStore(store *Store) *RelationStore {
|
||||
return &RelationStore{
|
||||
db: store.DB,
|
||||
}
|
||||
}
|
||||
|
||||
// StoreRelation stores a new observation relation.
|
||||
// Uses INSERT OR IGNORE to handle duplicate (source_id, target_id, relation_type) combinations.
|
||||
func (s *RelationStore) StoreRelation(ctx context.Context, relation *models.ObservationRelation) (int64, error) {
|
||||
dbRelation := &ObservationRelation{
|
||||
SourceID: relation.SourceID,
|
||||
TargetID: relation.TargetID,
|
||||
RelationType: relation.RelationType,
|
||||
Confidence: relation.Confidence,
|
||||
DetectionSource: relation.DetectionSource,
|
||||
CreatedAt: relation.CreatedAt,
|
||||
CreatedAtEpoch: relation.CreatedAtEpoch,
|
||||
}
|
||||
|
||||
// Handle nullable fields
|
||||
if relation.Reason != "" {
|
||||
dbRelation.Reason = sql.NullString{String: relation.Reason, Valid: true}
|
||||
}
|
||||
|
||||
// INSERT OR IGNORE using OnConflict
|
||||
result := s.db.WithContext(ctx).
|
||||
Clauses(clause.OnConflict{
|
||||
Columns: []clause.Column{{Name: "source_id"}, {Name: "target_id"}, {Name: "relation_type"}},
|
||||
DoNothing: true,
|
||||
}).
|
||||
Create(dbRelation)
|
||||
|
||||
if result.Error != nil {
|
||||
return 0, result.Error
|
||||
}
|
||||
|
||||
// If RowsAffected is 0, the insert was ignored (duplicate)
|
||||
if result.RowsAffected == 0 {
|
||||
var existing ObservationRelation
|
||||
err := s.db.Where("source_id = ? AND target_id = ? AND relation_type = ?",
|
||||
relation.SourceID, relation.TargetID, relation.RelationType).
|
||||
First(&existing).Error
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return existing.ID, nil
|
||||
}
|
||||
|
||||
return dbRelation.ID, nil
|
||||
}
|
||||
|
||||
// StoreRelations stores multiple relations in a single transaction.
|
||||
func (s *RelationStore) StoreRelations(ctx context.Context, relations []*models.ObservationRelation) error {
|
||||
if len(relations) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
return s.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
for _, rel := range relations {
|
||||
dbRelation := &ObservationRelation{
|
||||
SourceID: rel.SourceID,
|
||||
TargetID: rel.TargetID,
|
||||
RelationType: rel.RelationType,
|
||||
Confidence: rel.Confidence,
|
||||
DetectionSource: rel.DetectionSource,
|
||||
CreatedAt: rel.CreatedAt,
|
||||
CreatedAtEpoch: rel.CreatedAtEpoch,
|
||||
}
|
||||
|
||||
if rel.Reason != "" {
|
||||
dbRelation.Reason = sql.NullString{String: rel.Reason, Valid: true}
|
||||
}
|
||||
|
||||
result := tx.Clauses(clause.OnConflict{
|
||||
Columns: []clause.Column{{Name: "source_id"}, {Name: "target_id"}, {Name: "relation_type"}},
|
||||
DoNothing: true,
|
||||
}).Create(dbRelation)
|
||||
|
||||
if result.Error != nil {
|
||||
return result.Error
|
||||
}
|
||||
}
|
||||
return nil
|
||||
})
|
||||
}
|
||||
|
||||
// GetRelationsByObservationID retrieves all relations involving an observation (as source or target).
|
||||
func (s *RelationStore) GetRelationsByObservationID(ctx context.Context, obsID int64) ([]*models.ObservationRelation, error) {
|
||||
var relations []ObservationRelation
|
||||
|
||||
err := s.db.WithContext(ctx).
|
||||
Where("source_id = ? OR target_id = ?", obsID, obsID).
|
||||
Order("confidence DESC, created_at_epoch DESC").
|
||||
Find(&relations).Error
|
||||
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return toModelRelations(relations), nil
|
||||
}
|
||||
|
||||
// GetOutgoingRelations retrieves relations where the observation is the source.
|
||||
func (s *RelationStore) GetOutgoingRelations(ctx context.Context, obsID int64) ([]*models.ObservationRelation, error) {
|
||||
var relations []ObservationRelation
|
||||
|
||||
err := s.db.WithContext(ctx).
|
||||
Where("source_id = ?", obsID).
|
||||
Order("confidence DESC, created_at_epoch DESC").
|
||||
Find(&relations).Error
|
||||
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return toModelRelations(relations), nil
|
||||
}
|
||||
|
||||
// GetIncomingRelations retrieves relations where the observation is the target.
|
||||
func (s *RelationStore) GetIncomingRelations(ctx context.Context, obsID int64) ([]*models.ObservationRelation, error) {
|
||||
var relations []ObservationRelation
|
||||
|
||||
err := s.db.WithContext(ctx).
|
||||
Where("target_id = ?", obsID).
|
||||
Order("confidence DESC, created_at_epoch DESC").
|
||||
Find(&relations).Error
|
||||
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return toModelRelations(relations), nil
|
||||
}
|
||||
|
||||
// GetRelationsByType retrieves all relations of a specific type.
|
||||
func (s *RelationStore) GetRelationsByType(ctx context.Context, relationType models.RelationType, limit int) ([]*models.ObservationRelation, error) {
|
||||
var relations []ObservationRelation
|
||||
|
||||
err := s.db.WithContext(ctx).
|
||||
Where("relation_type = ?", relationType).
|
||||
Order("confidence DESC, created_at_epoch DESC").
|
||||
Limit(limit).
|
||||
Find(&relations).Error
|
||||
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return toModelRelations(relations), nil
|
||||
}
|
||||
|
||||
// GetRelationsWithDetails retrieves relations with observation titles for display.
|
||||
func (s *RelationStore) GetRelationsWithDetails(ctx context.Context, obsID int64) ([]*models.RelationWithDetails, error) {
|
||||
var results []struct {
|
||||
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).
|
||||
Table("observation_relations r").
|
||||
Select("r.*, "+
|
||||
"COALESCE(src.title, '') as source_title, "+
|
||||
"COALESCE(tgt.title, '') as target_title, "+
|
||||
"src.type as source_type, "+
|
||||
"tgt.type as target_type").
|
||||
Joins("JOIN observations src ON src.id = r.source_id").
|
||||
Joins("JOIN observations tgt ON tgt.id = r.target_id").
|
||||
Where("r.source_id = ? OR r.target_id = ?", obsID, obsID).
|
||||
Order("r.confidence DESC, r.created_at_epoch DESC").
|
||||
Scan(&results).Error
|
||||
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
relations := make([]*models.RelationWithDetails, len(results))
|
||||
for i, r := range results {
|
||||
relations[i] = &models.RelationWithDetails{
|
||||
Relation: toModelRelation(&r.ObservationRelation),
|
||||
SourceTitle: r.SourceTitle.String,
|
||||
TargetTitle: r.TargetTitle.String,
|
||||
SourceType: models.ObservationType(r.SourceType),
|
||||
TargetType: models.ObservationType(r.TargetType),
|
||||
}
|
||||
}
|
||||
|
||||
return relations, nil
|
||||
}
|
||||
|
||||
// GetRelationGraph retrieves a relation graph centered on an observation.
|
||||
// This returns all observations within N hops from the center.
|
||||
func (s *RelationStore) GetRelationGraph(ctx context.Context, centerID int64, maxDepth int) (*models.RelationGraph, error) {
|
||||
// Get all relations involving the center observation
|
||||
relations, err := s.GetRelationsWithDetails(ctx, centerID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
graph := &models.RelationGraph{
|
||||
CenterID: centerID,
|
||||
Relations: relations,
|
||||
}
|
||||
|
||||
// If depth > 1, recursively get relations for connected observations
|
||||
if maxDepth > 1 {
|
||||
visited := map[int64]bool{centerID: true}
|
||||
toVisit := make([]int64, 0)
|
||||
|
||||
// Collect IDs of directly connected observations
|
||||
for _, r := range relations {
|
||||
if !visited[r.Relation.SourceID] {
|
||||
toVisit = append(toVisit, r.Relation.SourceID)
|
||||
visited[r.Relation.SourceID] = true
|
||||
}
|
||||
if !visited[r.Relation.TargetID] {
|
||||
toVisit = append(toVisit, r.Relation.TargetID)
|
||||
visited[r.Relation.TargetID] = true
|
||||
}
|
||||
}
|
||||
|
||||
// Get relations for connected observations (depth - 1)
|
||||
for depth := 1; depth < maxDepth && len(toVisit) > 0; depth++ {
|
||||
nextLevel := make([]int64, 0)
|
||||
for _, obsID := range toVisit {
|
||||
moreRelations, err := s.GetRelationsWithDetails(ctx, obsID)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
for _, r := range moreRelations {
|
||||
// Avoid duplicates
|
||||
exists := false
|
||||
for _, existing := range graph.Relations {
|
||||
if existing.Relation.ID == r.Relation.ID {
|
||||
exists = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !exists {
|
||||
graph.Relations = append(graph.Relations, r)
|
||||
}
|
||||
|
||||
// Queue next level
|
||||
if !visited[r.Relation.SourceID] {
|
||||
nextLevel = append(nextLevel, r.Relation.SourceID)
|
||||
visited[r.Relation.SourceID] = true
|
||||
}
|
||||
if !visited[r.Relation.TargetID] {
|
||||
nextLevel = append(nextLevel, r.Relation.TargetID)
|
||||
visited[r.Relation.TargetID] = true
|
||||
}
|
||||
}
|
||||
}
|
||||
toVisit = nextLevel
|
||||
}
|
||||
}
|
||||
|
||||
return graph, nil
|
||||
}
|
||||
|
||||
// DeleteRelationsByObservationID deletes all relations involving an observation.
|
||||
// Called when an observation is deleted.
|
||||
func (s *RelationStore) DeleteRelationsByObservationID(ctx context.Context, obsID int64) error {
|
||||
result := s.db.WithContext(ctx).
|
||||
Where("source_id = ? OR target_id = ?", obsID, obsID).
|
||||
Delete(&ObservationRelation{})
|
||||
|
||||
return result.Error
|
||||
}
|
||||
|
||||
// GetRelationCount returns the count of relations for an observation.
|
||||
func (s *RelationStore) GetRelationCount(ctx context.Context, obsID int64) (int, error) {
|
||||
var count int64
|
||||
err := s.db.WithContext(ctx).
|
||||
Model(&ObservationRelation{}).
|
||||
Where("source_id = ? OR target_id = ?", obsID, obsID).
|
||||
Count(&count).Error
|
||||
|
||||
return int(count), err
|
||||
}
|
||||
|
||||
// GetTotalRelationCount returns the total count of all relations.
|
||||
func (s *RelationStore) GetTotalRelationCount(ctx context.Context) (int, error) {
|
||||
var count int64
|
||||
err := s.db.WithContext(ctx).
|
||||
Model(&ObservationRelation{}).
|
||||
Count(&count).Error
|
||||
|
||||
return int(count), err
|
||||
}
|
||||
|
||||
// GetHighConfidenceRelations retrieves relations with confidence above threshold.
|
||||
func (s *RelationStore) GetHighConfidenceRelations(ctx context.Context, minConfidence float64, limit int) ([]*models.ObservationRelation, error) {
|
||||
var relations []ObservationRelation
|
||||
|
||||
err := s.db.WithContext(ctx).
|
||||
Where("confidence >= ?", minConfidence).
|
||||
Order("confidence DESC, created_at_epoch DESC").
|
||||
Limit(limit).
|
||||
Find(&relations).Error
|
||||
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return toModelRelations(relations), nil
|
||||
}
|
||||
|
||||
// UpdateRelationConfidence updates the confidence of a relation.
|
||||
func (s *RelationStore) UpdateRelationConfidence(ctx context.Context, relationID int64, newConfidence float64) error {
|
||||
result := s.db.WithContext(ctx).
|
||||
Model(&ObservationRelation{}).
|
||||
Where("id = ?", relationID).
|
||||
Update("confidence", newConfidence)
|
||||
|
||||
return result.Error
|
||||
}
|
||||
|
||||
// GetRelatedObservationIDs returns IDs of observations related to the given one.
|
||||
// This is useful for expanding search results.
|
||||
// Uses CASE expression for bidirectional ID lookup (GORM doesn't support this well, so we use raw SQL).
|
||||
func (s *RelationStore) GetRelatedObservationIDs(ctx context.Context, obsID int64, minConfidence float64) ([]int64, error) {
|
||||
var ids []int64
|
||||
|
||||
err := s.db.WithContext(ctx).
|
||||
Raw("SELECT DISTINCT CASE WHEN source_id = ? THEN target_id ELSE source_id END as related_id "+
|
||||
"FROM observation_relations "+
|
||||
"WHERE (source_id = ? OR target_id = ?) AND confidence >= ?",
|
||||
obsID, obsID, obsID, minConfidence).
|
||||
Pluck("related_id", &ids).Error
|
||||
|
||||
return ids, err
|
||||
}
|
||||
|
||||
// toModelRelation converts a GORM ObservationRelation to a pkg/models ObservationRelation.
|
||||
func toModelRelation(r *ObservationRelation) *models.ObservationRelation {
|
||||
relation := &models.ObservationRelation{
|
||||
ID: r.ID,
|
||||
SourceID: r.SourceID,
|
||||
TargetID: r.TargetID,
|
||||
RelationType: r.RelationType,
|
||||
Confidence: r.Confidence,
|
||||
DetectionSource: r.DetectionSource,
|
||||
CreatedAt: r.CreatedAt,
|
||||
CreatedAtEpoch: r.CreatedAtEpoch,
|
||||
}
|
||||
|
||||
if r.Reason.Valid {
|
||||
relation.Reason = r.Reason.String
|
||||
}
|
||||
|
||||
return relation
|
||||
}
|
||||
|
||||
// toModelRelations converts a slice of GORM ObservationRelations to pkg/models ObservationRelations.
|
||||
func toModelRelations(relations []ObservationRelation) []*models.ObservationRelation {
|
||||
result := make([]*models.ObservationRelation, len(relations))
|
||||
for i, r := range relations {
|
||||
result[i] = toModelRelation(&r)
|
||||
}
|
||||
return result
|
||||
}
|
||||
@@ -1,306 +0,0 @@
|
||||
//go:build fts5
|
||||
|
||||
// Package gorm provides GORM-based database operations for claude-mnemonic.
|
||||
package gorm
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"gorm.io/gorm/logger"
|
||||
|
||||
"github.com/lukaszraczylo/claude-mnemonic/pkg/models"
|
||||
)
|
||||
|
||||
func testRelationStore(t *testing.T) (*RelationStore, *Store, func()) {
|
||||
t.Helper()
|
||||
|
||||
tmpDir, err := os.MkdirTemp("", "gorm_relation_test_*")
|
||||
if err != nil {
|
||||
t.Fatalf("create temp dir: %v", err)
|
||||
}
|
||||
|
||||
dbPath := filepath.Join(tmpDir, "test.db")
|
||||
cfg := Config{
|
||||
Path: dbPath,
|
||||
MaxConns: 4,
|
||||
LogLevel: logger.Silent,
|
||||
}
|
||||
|
||||
store, err := NewStore(cfg)
|
||||
if err != nil {
|
||||
os.RemoveAll(tmpDir)
|
||||
t.Fatalf("NewStore failed: %v", err)
|
||||
}
|
||||
|
||||
relationStore := NewRelationStore(store)
|
||||
|
||||
cleanup := func() {
|
||||
store.Close()
|
||||
os.RemoveAll(tmpDir)
|
||||
}
|
||||
|
||||
return relationStore, store, cleanup
|
||||
}
|
||||
|
||||
func TestRelationStore_StoreRelation(t *testing.T) {
|
||||
relationStore, _, cleanup := testRelationStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
now := time.Now()
|
||||
|
||||
relation := &models.ObservationRelation{
|
||||
SourceID: 1,
|
||||
TargetID: 2,
|
||||
RelationType: models.RelationCauses,
|
||||
Confidence: 0.8,
|
||||
DetectionSource: models.DetectionSourceFileOverlap,
|
||||
Reason: "Test relation",
|
||||
CreatedAt: now.Format(time.RFC3339),
|
||||
CreatedAtEpoch: now.UnixMilli(),
|
||||
}
|
||||
|
||||
id, err := relationStore.StoreRelation(ctx, relation)
|
||||
require.NoError(t, err)
|
||||
assert.Greater(t, id, int64(0))
|
||||
}
|
||||
|
||||
func TestRelationStore_StoreRelation_Idempotency(t *testing.T) {
|
||||
relationStore, _, cleanup := testRelationStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
now := time.Now()
|
||||
|
||||
relation := &models.ObservationRelation{
|
||||
SourceID: 1,
|
||||
TargetID: 2,
|
||||
RelationType: models.RelationCauses,
|
||||
Confidence: 0.8,
|
||||
DetectionSource: models.DetectionSourceFileOverlap,
|
||||
CreatedAt: now.Format(time.RFC3339),
|
||||
CreatedAtEpoch: now.UnixMilli(),
|
||||
}
|
||||
|
||||
id1, err := relationStore.StoreRelation(ctx, relation)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Store again with same source/target/type - should return same ID
|
||||
id2, err := relationStore.StoreRelation(ctx, relation)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, id1, id2)
|
||||
}
|
||||
|
||||
func TestRelationStore_StoreRelations(t *testing.T) {
|
||||
relationStore, _, cleanup := testRelationStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
now := time.Now()
|
||||
|
||||
relations := []*models.ObservationRelation{
|
||||
{
|
||||
SourceID: 1,
|
||||
TargetID: 2,
|
||||
RelationType: models.RelationCauses,
|
||||
Confidence: 0.8,
|
||||
DetectionSource: models.DetectionSourceFileOverlap,
|
||||
CreatedAt: now.Format(time.RFC3339),
|
||||
CreatedAtEpoch: now.UnixMilli(),
|
||||
},
|
||||
{
|
||||
SourceID: 2,
|
||||
TargetID: 3,
|
||||
RelationType: models.RelationFixes,
|
||||
Confidence: 0.9,
|
||||
DetectionSource: models.DetectionSourceTemporalProximity,
|
||||
CreatedAt: now.Format(time.RFC3339),
|
||||
CreatedAtEpoch: now.UnixMilli(),
|
||||
},
|
||||
}
|
||||
|
||||
err := relationStore.StoreRelations(ctx, relations)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Verify both were stored
|
||||
count, err := relationStore.GetTotalRelationCount(ctx)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 2, count)
|
||||
}
|
||||
|
||||
func TestRelationStore_GetRelationsByObservationID(t *testing.T) {
|
||||
relationStore, _, cleanup := testRelationStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
now := time.Now()
|
||||
|
||||
// Create relations involving observation 2
|
||||
relations := []*models.ObservationRelation{
|
||||
{
|
||||
SourceID: 1,
|
||||
TargetID: 2,
|
||||
RelationType: models.RelationCauses,
|
||||
Confidence: 0.8,
|
||||
DetectionSource: models.DetectionSourceFileOverlap,
|
||||
CreatedAt: now.Format(time.RFC3339),
|
||||
CreatedAtEpoch: now.UnixMilli(),
|
||||
},
|
||||
{
|
||||
SourceID: 2,
|
||||
TargetID: 3,
|
||||
RelationType: models.RelationFixes,
|
||||
Confidence: 0.9,
|
||||
DetectionSource: models.DetectionSourceTemporalProximity,
|
||||
CreatedAt: now.Format(time.RFC3339),
|
||||
CreatedAtEpoch: now.UnixMilli(),
|
||||
},
|
||||
}
|
||||
|
||||
err := relationStore.StoreRelations(ctx, relations)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Get relations for observation 2 (involved in both)
|
||||
result, err := relationStore.GetRelationsByObservationID(ctx, 2)
|
||||
require.NoError(t, err)
|
||||
assert.Len(t, result, 2)
|
||||
}
|
||||
|
||||
func TestRelationStore_GetOutgoingAndIncomingRelations(t *testing.T) {
|
||||
relationStore, _, cleanup := testRelationStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
now := time.Now()
|
||||
|
||||
relations := []*models.ObservationRelation{
|
||||
{
|
||||
SourceID: 2,
|
||||
TargetID: 1,
|
||||
RelationType: models.RelationCauses,
|
||||
Confidence: 0.8,
|
||||
DetectionSource: models.DetectionSourceFileOverlap,
|
||||
CreatedAt: now.Format(time.RFC3339),
|
||||
CreatedAtEpoch: now.UnixMilli(),
|
||||
},
|
||||
{
|
||||
SourceID: 3,
|
||||
TargetID: 2,
|
||||
RelationType: models.RelationFixes,
|
||||
Confidence: 0.9,
|
||||
DetectionSource: models.DetectionSourceTemporalProximity,
|
||||
CreatedAt: now.Format(time.RFC3339),
|
||||
CreatedAtEpoch: now.UnixMilli(),
|
||||
},
|
||||
}
|
||||
|
||||
err := relationStore.StoreRelations(ctx, relations)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Observation 2 has 1 outgoing (to 1) and 1 incoming (from 3)
|
||||
outgoing, err := relationStore.GetOutgoingRelations(ctx, 2)
|
||||
require.NoError(t, err)
|
||||
assert.Len(t, outgoing, 1)
|
||||
assert.Equal(t, int64(1), outgoing[0].TargetID)
|
||||
|
||||
incoming, err := relationStore.GetIncomingRelations(ctx, 2)
|
||||
require.NoError(t, err)
|
||||
assert.Len(t, incoming, 1)
|
||||
assert.Equal(t, int64(3), incoming[0].SourceID)
|
||||
}
|
||||
|
||||
func TestRelationStore_GetRelationCount(t *testing.T) {
|
||||
relationStore, _, cleanup := testRelationStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
now := time.Now()
|
||||
|
||||
relations := []*models.ObservationRelation{
|
||||
{
|
||||
SourceID: 1,
|
||||
TargetID: 2,
|
||||
RelationType: models.RelationCauses,
|
||||
Confidence: 0.8,
|
||||
DetectionSource: models.DetectionSourceFileOverlap,
|
||||
CreatedAt: now.Format(time.RFC3339),
|
||||
CreatedAtEpoch: now.UnixMilli(),
|
||||
},
|
||||
{
|
||||
SourceID: 2,
|
||||
TargetID: 3,
|
||||
RelationType: models.RelationFixes,
|
||||
Confidence: 0.9,
|
||||
DetectionSource: models.DetectionSourceTemporalProximity,
|
||||
CreatedAt: now.Format(time.RFC3339),
|
||||
CreatedAtEpoch: now.UnixMilli(),
|
||||
},
|
||||
}
|
||||
|
||||
err := relationStore.StoreRelations(ctx, relations)
|
||||
require.NoError(t, err)
|
||||
|
||||
count, err := relationStore.GetRelationCount(ctx, 2)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 2, count)
|
||||
|
||||
count, err = relationStore.GetRelationCount(ctx, 1)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 1, count)
|
||||
}
|
||||
|
||||
func TestRelationStore_DeleteRelationsByObservationID(t *testing.T) {
|
||||
relationStore, _, cleanup := testRelationStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
now := time.Now()
|
||||
|
||||
relations := []*models.ObservationRelation{
|
||||
{
|
||||
SourceID: 1,
|
||||
TargetID: 2,
|
||||
RelationType: models.RelationCauses,
|
||||
Confidence: 0.8,
|
||||
DetectionSource: models.DetectionSourceFileOverlap,
|
||||
CreatedAt: now.Format(time.RFC3339),
|
||||
CreatedAtEpoch: now.UnixMilli(),
|
||||
},
|
||||
{
|
||||
SourceID: 2,
|
||||
TargetID: 3,
|
||||
RelationType: models.RelationFixes,
|
||||
Confidence: 0.9,
|
||||
DetectionSource: models.DetectionSourceTemporalProximity,
|
||||
CreatedAt: now.Format(time.RFC3339),
|
||||
CreatedAtEpoch: now.UnixMilli(),
|
||||
},
|
||||
{
|
||||
SourceID: 4,
|
||||
TargetID: 5,
|
||||
RelationType: models.RelationRelatesTo,
|
||||
Confidence: 0.7,
|
||||
DetectionSource: models.DetectionSourceConceptOverlap,
|
||||
CreatedAt: now.Format(time.RFC3339),
|
||||
CreatedAtEpoch: now.UnixMilli(),
|
||||
},
|
||||
}
|
||||
|
||||
err := relationStore.StoreRelations(ctx, relations)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Delete relations involving observation 2
|
||||
err = relationStore.DeleteRelationsByObservationID(ctx, 2)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Verify only 1 relation remains (4->5)
|
||||
total, err := relationStore.GetTotalRelationCount(ctx)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 1, total)
|
||||
}
|
||||
@@ -1,303 +0,0 @@
|
||||
// Package gorm provides GORM-based database operations for claude-mnemonic.
|
||||
package gorm
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"gorm.io/gorm"
|
||||
|
||||
"github.com/lukaszraczylo/claude-mnemonic/pkg/models"
|
||||
)
|
||||
|
||||
// UpdateObservationFeedback updates the user feedback for an observation.
|
||||
// Feedback values: -1 (thumbs down), 0 (neutral), 1 (thumbs up).
|
||||
func (s *ObservationStore) UpdateObservationFeedback(ctx context.Context, id int64, feedback int) error {
|
||||
now := time.Now().UnixMilli()
|
||||
|
||||
result := s.db.WithContext(ctx).
|
||||
Model(&Observation{}).
|
||||
Where("id = ?", id).
|
||||
Updates(map[string]interface{}{
|
||||
"user_feedback": feedback,
|
||||
"score_updated_at_epoch": now,
|
||||
})
|
||||
|
||||
return result.Error
|
||||
}
|
||||
|
||||
// IncrementRetrievalCount increments the retrieval counter for the given observation IDs.
|
||||
// This is called when observations are returned in search results.
|
||||
func (s *ObservationStore) IncrementRetrievalCount(ctx context.Context, ids []int64) error {
|
||||
if len(ids) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
now := time.Now().UnixMilli()
|
||||
|
||||
// Use raw SQL for increment expression
|
||||
result := s.db.WithContext(ctx).
|
||||
Exec("UPDATE observations SET retrieval_count = COALESCE(retrieval_count, 0) + 1, last_retrieved_at_epoch = ? WHERE id IN ?",
|
||||
now, ids)
|
||||
|
||||
return result.Error
|
||||
}
|
||||
|
||||
// UpdateImportanceScore updates the importance score for a single observation.
|
||||
func (s *ObservationStore) UpdateImportanceScore(ctx context.Context, id int64, score float64) error {
|
||||
now := time.Now().UnixMilli()
|
||||
|
||||
result := s.db.WithContext(ctx).
|
||||
Model(&Observation{}).
|
||||
Where("id = ?", id).
|
||||
Updates(map[string]interface{}{
|
||||
"importance_score": score,
|
||||
"score_updated_at_epoch": now,
|
||||
})
|
||||
|
||||
return result.Error
|
||||
}
|
||||
|
||||
// UpdateImportanceScores bulk updates importance scores for multiple observations.
|
||||
// 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()
|
||||
|
||||
for id, score := range scores {
|
||||
err := tx.Model(&Observation{}).
|
||||
Where("id = ?", id).
|
||||
Updates(map[string]interface{}{
|
||||
"importance_score": score,
|
||||
"score_updated_at_epoch": now,
|
||||
}).Error
|
||||
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
})
|
||||
}
|
||||
|
||||
// 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) {
|
||||
cutoff := time.Now().Add(-threshold).UnixMilli()
|
||||
|
||||
var observations []Observation
|
||||
|
||||
err := s.db.WithContext(ctx).
|
||||
Where("score_updated_at_epoch IS NULL OR score_updated_at_epoch < ?", cutoff).
|
||||
Order("created_at_epoch DESC").
|
||||
Limit(limit).
|
||||
Find(&observations).Error
|
||||
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return toModelObservations(observations), nil
|
||||
}
|
||||
|
||||
// GetConceptWeights returns all concept weights from the database.
|
||||
func (s *ObservationStore) GetConceptWeights(ctx context.Context) (map[string]float64, error) {
|
||||
var weights []struct {
|
||||
Concept string
|
||||
Weight float64
|
||||
}
|
||||
|
||||
err := s.db.WithContext(ctx).
|
||||
Table("concept_weights").
|
||||
Select("concept, weight").
|
||||
Scan(&weights).Error
|
||||
|
||||
if err != nil {
|
||||
return models.DefaultConceptWeights, nil
|
||||
}
|
||||
|
||||
if len(weights) == 0 {
|
||||
return models.DefaultConceptWeights, nil
|
||||
}
|
||||
|
||||
result := make(map[string]float64, len(weights))
|
||||
for _, w := range weights {
|
||||
result[w.Concept] = w.Weight
|
||||
}
|
||||
|
||||
return result, nil
|
||||
}
|
||||
|
||||
// SetConceptWeights stores concept weights in the database using UPSERT.
|
||||
func (s *ObservationStore) SetConceptWeights(ctx context.Context, weights map[string]float64) error {
|
||||
if len(weights) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
return s.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
for concept, weight := range weights {
|
||||
// UPSERT using raw SQL since GORM's ON CONFLICT is complex for this case
|
||||
err := tx.Exec(`
|
||||
INSERT INTO concept_weights (concept, weight, updated_at)
|
||||
VALUES (?, ?, datetime('now'))
|
||||
ON CONFLICT(concept) DO UPDATE SET weight = excluded.weight, updated_at = excluded.updated_at
|
||||
`, concept, weight).Error
|
||||
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
})
|
||||
}
|
||||
|
||||
// UpdateConceptWeight updates a single concept weight in the database using UPSERT.
|
||||
func (s *ObservationStore) UpdateConceptWeight(ctx context.Context, concept string, weight float64) error {
|
||||
return s.db.WithContext(ctx).Exec(`
|
||||
INSERT INTO concept_weights (concept, weight, updated_at)
|
||||
VALUES (?, ?, datetime('now'))
|
||||
ON CONFLICT(concept) DO UPDATE SET weight = excluded.weight, updated_at = excluded.updated_at
|
||||
`, concept, weight).Error
|
||||
}
|
||||
|
||||
// FeedbackStats contains statistics about observation feedback and scoring.
|
||||
type FeedbackStats struct {
|
||||
Total int `json:"total"`
|
||||
Positive int `json:"positive"`
|
||||
Negative int `json:"negative"`
|
||||
Neutral int `json:"neutral"`
|
||||
AvgScore float64 `json:"avg_score"`
|
||||
AvgRetrieval float64 `json:"avg_retrieval"`
|
||||
}
|
||||
|
||||
// GetObservationFeedbackStats returns statistics about user feedback.
|
||||
func (s *ObservationStore) GetObservationFeedbackStats(ctx context.Context, project string) (*FeedbackStats, error) {
|
||||
var stats FeedbackStats
|
||||
|
||||
query := s.db.WithContext(ctx).
|
||||
Model(&Observation{}).
|
||||
Select(`
|
||||
COUNT(*) as total,
|
||||
COALESCE(SUM(CASE WHEN user_feedback = 1 THEN 1 ELSE 0 END), 0) as positive,
|
||||
COALESCE(SUM(CASE WHEN user_feedback = -1 THEN 1 ELSE 0 END), 0) as negative,
|
||||
COALESCE(SUM(CASE WHEN user_feedback = 0 THEN 1 ELSE 0 END), 0) as neutral,
|
||||
COALESCE(AVG(COALESCE(importance_score, 1.0)), 0) as avg_score,
|
||||
COALESCE(AVG(COALESCE(retrieval_count, 0)), 0) as avg_retrieval
|
||||
`)
|
||||
|
||||
if project != "" {
|
||||
query = query.Where("project = ? OR scope = 'global'", project)
|
||||
}
|
||||
|
||||
err := query.Scan(&stats).Error
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return &stats, nil
|
||||
}
|
||||
|
||||
// GetTopScoringObservations returns the highest-scoring observations.
|
||||
func (s *ObservationStore) GetTopScoringObservations(ctx context.Context, project string, limit int) ([]*models.Observation, error) {
|
||||
var observations []Observation
|
||||
|
||||
query := s.db.WithContext(ctx).
|
||||
Order("COALESCE(importance_score, 1.0) DESC, created_at_epoch DESC").
|
||||
Limit(limit)
|
||||
|
||||
if project != "" {
|
||||
query = query.Where("project = ? OR scope = 'global'", project)
|
||||
}
|
||||
|
||||
err := query.Find(&observations).Error
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return toModelObservations(observations), nil
|
||||
}
|
||||
|
||||
// GetMostRetrievedObservations returns the most frequently retrieved observations.
|
||||
func (s *ObservationStore) GetMostRetrievedObservations(ctx context.Context, project string, limit int) ([]*models.Observation, error) {
|
||||
var observations []Observation
|
||||
|
||||
query := s.db.WithContext(ctx).
|
||||
Where("retrieval_count > 0").
|
||||
Order("retrieval_count DESC, created_at_epoch DESC").
|
||||
Limit(limit)
|
||||
|
||||
if project != "" {
|
||||
query = query.Where("project = ? OR scope = 'global'", project)
|
||||
}
|
||||
|
||||
err := query.Find(&observations).Error
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return toModelObservations(observations), nil
|
||||
}
|
||||
|
||||
// ResetObservationScores resets all observation scores to their default values.
|
||||
// This is useful for testing or when changing the scoring algorithm.
|
||||
func (s *ObservationStore) ResetObservationScores(ctx context.Context) error {
|
||||
// Use Where("1 = 1") to explicitly allow bulk update of all rows
|
||||
result := s.db.WithContext(ctx).
|
||||
Model(&Observation{}).
|
||||
Where("1 = 1").
|
||||
Updates(map[string]interface{}{
|
||||
"importance_score": 1.0,
|
||||
"score_updated_at_epoch": nil,
|
||||
})
|
||||
|
||||
return result.Error
|
||||
}
|
||||
@@ -1,355 +0,0 @@
|
||||
//go:build fts5
|
||||
|
||||
// Package gorm provides GORM-based database operations for claude-mnemonic.
|
||||
package gorm
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/lukaszraczylo/claude-mnemonic/pkg/models"
|
||||
)
|
||||
|
||||
func TestObservationStore_UpdateImportanceScore(t *testing.T) {
|
||||
obsStore, store, cleanup := testObservationStore(t)
|
||||
defer cleanup()
|
||||
ctx := context.Background()
|
||||
|
||||
// Create observation
|
||||
sessionStore := NewSessionStore(store)
|
||||
sessionID, _ := sessionStore.CreateSDKSession(ctx, "claude-1", "test-project", "")
|
||||
obs := &models.ParsedObservation{Type: models.ObsTypeDiscovery, Title: "Test"}
|
||||
obsID, _, _ := obsStore.StoreObservation(ctx, "claude-1", "test-project", obs, int(sessionID), 1)
|
||||
|
||||
// Update score
|
||||
err := obsStore.UpdateImportanceScore(ctx, obsID, 5.0)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Verify
|
||||
var dbObs Observation
|
||||
store.DB.First(&dbObs, obsID)
|
||||
assert.Equal(t, 5.0, dbObs.ImportanceScore)
|
||||
}
|
||||
|
||||
func TestObservationStore_IncrementRetrievalCount(t *testing.T) {
|
||||
obsStore, store, cleanup := testObservationStore(t)
|
||||
defer cleanup()
|
||||
ctx := context.Background()
|
||||
|
||||
sessionStore := NewSessionStore(store)
|
||||
sessionID, _ := sessionStore.CreateSDKSession(ctx, "claude-1", "test-project", "")
|
||||
obs := &models.ParsedObservation{Type: models.ObsTypeDiscovery, Title: "Test"}
|
||||
obsID, _, _ := obsStore.StoreObservation(ctx, "claude-1", "test-project", obs, int(sessionID), 1)
|
||||
|
||||
err := obsStore.IncrementRetrievalCount(ctx, []int64{obsID})
|
||||
require.NoError(t, err)
|
||||
|
||||
var dbObs Observation
|
||||
store.DB.First(&dbObs, obsID)
|
||||
assert.Equal(t, 1, dbObs.RetrievalCount)
|
||||
}
|
||||
|
||||
func TestObservationStore_IncrementRetrievalCount_Multiple(t *testing.T) {
|
||||
obsStore, store, cleanup := testObservationStore(t)
|
||||
defer cleanup()
|
||||
ctx := context.Background()
|
||||
|
||||
sessionStore := NewSessionStore(store)
|
||||
sessionID, _ := sessionStore.CreateSDKSession(ctx, "claude-1", "test-project", "")
|
||||
|
||||
// Create 3 observations
|
||||
ids := make([]int64, 3)
|
||||
for i := 0; i < 3; i++ {
|
||||
obs := &models.ParsedObservation{Type: models.ObsTypeDiscovery, Title: "Test"}
|
||||
obsID, _, _ := obsStore.StoreObservation(ctx, "claude-1", "test-project", obs, int(sessionID), 1)
|
||||
ids[i] = obsID
|
||||
}
|
||||
|
||||
// Increment all
|
||||
err := obsStore.IncrementRetrievalCount(ctx, ids)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Verify all were incremented
|
||||
for _, id := range ids {
|
||||
var dbObs Observation
|
||||
store.DB.First(&dbObs, id)
|
||||
assert.Equal(t, 1, dbObs.RetrievalCount)
|
||||
}
|
||||
|
||||
// Increment again
|
||||
err = obsStore.IncrementRetrievalCount(ctx, ids)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Verify all are now 2
|
||||
for _, id := range ids {
|
||||
var dbObs Observation
|
||||
store.DB.First(&dbObs, id)
|
||||
assert.Equal(t, 2, dbObs.RetrievalCount)
|
||||
}
|
||||
}
|
||||
|
||||
func TestObservationStore_UpdateImportanceScores_Bulk(t *testing.T) {
|
||||
obsStore, store, cleanup := testObservationStore(t)
|
||||
defer cleanup()
|
||||
ctx := context.Background()
|
||||
|
||||
sessionStore := NewSessionStore(store)
|
||||
sessionID, _ := sessionStore.CreateSDKSession(ctx, "claude-1", "test-project", "")
|
||||
|
||||
// Create 3 observations
|
||||
ids := make([]int64, 3)
|
||||
for i := 0; i < 3; i++ {
|
||||
obs := &models.ParsedObservation{Type: models.ObsTypeDiscovery, Title: "Test"}
|
||||
obsID, _, _ := obsStore.StoreObservation(ctx, "claude-1", "test-project", obs, int(sessionID), 1)
|
||||
ids[i] = obsID
|
||||
}
|
||||
|
||||
// Bulk update scores
|
||||
scores := map[int64]float64{
|
||||
ids[0]: 2.5,
|
||||
ids[1]: 3.7,
|
||||
ids[2]: 1.2,
|
||||
}
|
||||
|
||||
err := obsStore.UpdateImportanceScores(ctx, scores)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Verify scores
|
||||
for id, expectedScore := range scores {
|
||||
var dbObs Observation
|
||||
store.DB.First(&dbObs, id)
|
||||
assert.Equal(t, expectedScore, dbObs.ImportanceScore)
|
||||
}
|
||||
}
|
||||
|
||||
func TestObservationStore_UpdateObservationFeedback(t *testing.T) {
|
||||
obsStore, store, cleanup := testObservationStore(t)
|
||||
defer cleanup()
|
||||
ctx := context.Background()
|
||||
|
||||
sessionStore := NewSessionStore(store)
|
||||
sessionID, _ := sessionStore.CreateSDKSession(ctx, "claude-1", "test-project", "")
|
||||
obs := &models.ParsedObservation{Type: models.ObsTypeDiscovery, Title: "Test"}
|
||||
obsID, _, _ := obsStore.StoreObservation(ctx, "claude-1", "test-project", obs, int(sessionID), 1)
|
||||
|
||||
// Set thumbs up
|
||||
err := obsStore.UpdateObservationFeedback(ctx, obsID, 1)
|
||||
require.NoError(t, err)
|
||||
|
||||
var dbObs Observation
|
||||
store.DB.First(&dbObs, obsID)
|
||||
assert.Equal(t, 1, dbObs.UserFeedback)
|
||||
|
||||
// Set thumbs down
|
||||
err = obsStore.UpdateObservationFeedback(ctx, obsID, -1)
|
||||
require.NoError(t, err)
|
||||
|
||||
store.DB.First(&dbObs, obsID)
|
||||
assert.Equal(t, -1, dbObs.UserFeedback)
|
||||
}
|
||||
|
||||
func TestObservationStore_GetObservationFeedbackStats(t *testing.T) {
|
||||
obsStore, store, cleanup := testObservationStore(t)
|
||||
defer cleanup()
|
||||
ctx := context.Background()
|
||||
|
||||
sessionStore := NewSessionStore(store)
|
||||
sessionID, _ := sessionStore.CreateSDKSession(ctx, "claude-1", "test-project", "")
|
||||
|
||||
// Create observations with different feedback
|
||||
obs1 := &models.ParsedObservation{Type: models.ObsTypeDiscovery, Title: "Test1"}
|
||||
obsID1, _, _ := obsStore.StoreObservation(ctx, "claude-1", "test-project", obs1, int(sessionID), 1)
|
||||
obsStore.UpdateObservationFeedback(ctx, obsID1, 1) // thumbs up
|
||||
obsStore.UpdateImportanceScore(ctx, obsID1, 3.0)
|
||||
|
||||
obs2 := &models.ParsedObservation{Type: models.ObsTypeBugfix, Title: "Test2"}
|
||||
obsID2, _, _ := obsStore.StoreObservation(ctx, "claude-1", "test-project", obs2, int(sessionID), 2)
|
||||
obsStore.UpdateObservationFeedback(ctx, obsID2, -1) // thumbs down
|
||||
obsStore.UpdateImportanceScore(ctx, obsID2, 2.0)
|
||||
|
||||
obs3 := &models.ParsedObservation{Type: models.ObsTypeFeature, Title: "Test3"}
|
||||
obsID3, _, _ := obsStore.StoreObservation(ctx, "claude-1", "test-project", obs3, int(sessionID), 3)
|
||||
// neutral (0)
|
||||
obsStore.UpdateImportanceScore(ctx, obsID3, 1.5)
|
||||
obsStore.IncrementRetrievalCount(ctx, []int64{obsID1, obsID2, obsID3})
|
||||
|
||||
// Get stats
|
||||
stats, err := obsStore.GetObservationFeedbackStats(ctx, "test-project")
|
||||
require.NoError(t, err)
|
||||
|
||||
assert.Equal(t, 3, stats.Total)
|
||||
assert.Equal(t, 1, stats.Positive)
|
||||
assert.Equal(t, 1, stats.Negative)
|
||||
assert.Equal(t, 1, stats.Neutral)
|
||||
assert.InDelta(t, 2.166, stats.AvgScore, 0.01) // (3.0 + 2.0 + 1.5) / 3
|
||||
assert.InDelta(t, 1.0, stats.AvgRetrieval, 0.01)
|
||||
}
|
||||
|
||||
func TestObservationStore_GetTopScoringObservations(t *testing.T) {
|
||||
obsStore, store, cleanup := testObservationStore(t)
|
||||
defer cleanup()
|
||||
ctx := context.Background()
|
||||
|
||||
sessionStore := NewSessionStore(store)
|
||||
sessionID, _ := sessionStore.CreateSDKSession(ctx, "claude-1", "test-project", "")
|
||||
|
||||
// Create observations with different scores
|
||||
obs1 := &models.ParsedObservation{Type: models.ObsTypeDiscovery, Title: "High"}
|
||||
obsID1, _, _ := obsStore.StoreObservation(ctx, "claude-1", "test-project", obs1, int(sessionID), 1)
|
||||
obsStore.UpdateImportanceScore(ctx, obsID1, 5.0)
|
||||
|
||||
obs2 := &models.ParsedObservation{Type: models.ObsTypeBugfix, Title: "Low"}
|
||||
obsID2, _, _ := obsStore.StoreObservation(ctx, "claude-1", "test-project", obs2, int(sessionID), 2)
|
||||
obsStore.UpdateImportanceScore(ctx, obsID2, 1.0)
|
||||
|
||||
obs3 := &models.ParsedObservation{Type: models.ObsTypeFeature, Title: "Medium"}
|
||||
obsID3, _, _ := obsStore.StoreObservation(ctx, "claude-1", "test-project", obs3, int(sessionID), 3)
|
||||
obsStore.UpdateImportanceScore(ctx, obsID3, 3.0)
|
||||
|
||||
// Get top 2
|
||||
topObs, err := obsStore.GetTopScoringObservations(ctx, "test-project", 2)
|
||||
require.NoError(t, err)
|
||||
|
||||
require.Len(t, topObs, 2)
|
||||
assert.True(t, topObs[0].Title.Valid)
|
||||
assert.Equal(t, "High", topObs[0].Title.String)
|
||||
assert.Equal(t, 5.0, topObs[0].ImportanceScore)
|
||||
assert.True(t, topObs[1].Title.Valid)
|
||||
assert.Equal(t, "Medium", topObs[1].Title.String)
|
||||
assert.Equal(t, 3.0, topObs[1].ImportanceScore)
|
||||
}
|
||||
|
||||
func TestObservationStore_GetMostRetrievedObservations(t *testing.T) {
|
||||
obsStore, store, cleanup := testObservationStore(t)
|
||||
defer cleanup()
|
||||
ctx := context.Background()
|
||||
|
||||
sessionStore := NewSessionStore(store)
|
||||
sessionID, _ := sessionStore.CreateSDKSession(ctx, "claude-1", "test-project", "")
|
||||
|
||||
// Create observations
|
||||
obs1 := &models.ParsedObservation{Type: models.ObsTypeDiscovery, Title: "Popular"}
|
||||
obsID1, _, _ := obsStore.StoreObservation(ctx, "claude-1", "test-project", obs1, int(sessionID), 1)
|
||||
|
||||
obs2 := &models.ParsedObservation{Type: models.ObsTypeBugfix, Title: "Unpopular"}
|
||||
obsID2, _, _ := obsStore.StoreObservation(ctx, "claude-1", "test-project", obs2, int(sessionID), 2)
|
||||
|
||||
// Increment retrieval counts - each call increments by 1
|
||||
obsStore.IncrementRetrievalCount(ctx, []int64{obsID1}) // increment by 1
|
||||
obsStore.IncrementRetrievalCount(ctx, []int64{obsID1}) // increment by 1 again (total: 2)
|
||||
obsStore.IncrementRetrievalCount(ctx, []int64{obsID1}) // increment by 1 again (total: 3)
|
||||
obsStore.IncrementRetrievalCount(ctx, []int64{obsID2}) // increment by 1 (total: 1)
|
||||
|
||||
// Get most retrieved - should return obsID1 (Popular) with count 3
|
||||
topObs, err := obsStore.GetMostRetrievedObservations(ctx, "test-project", 2)
|
||||
require.NoError(t, err)
|
||||
|
||||
require.Len(t, topObs, 2)
|
||||
assert.True(t, topObs[0].Title.Valid)
|
||||
assert.Equal(t, "Popular", topObs[0].Title.String)
|
||||
assert.Equal(t, 3, topObs[0].RetrievalCount)
|
||||
assert.True(t, topObs[1].Title.Valid)
|
||||
assert.Equal(t, "Unpopular", topObs[1].Title.String)
|
||||
assert.Equal(t, 1, topObs[1].RetrievalCount)
|
||||
}
|
||||
|
||||
func TestObservationStore_SetConceptWeights(t *testing.T) {
|
||||
obsStore, _, cleanup := testObservationStore(t)
|
||||
defer cleanup()
|
||||
ctx := context.Background()
|
||||
|
||||
// Set weights
|
||||
weights := map[string]float64{
|
||||
"security": 2.0,
|
||||
"performance": 1.5,
|
||||
"best-practice": 1.8,
|
||||
}
|
||||
|
||||
err := obsStore.SetConceptWeights(ctx, weights)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Get weights back
|
||||
retrieved, err := obsStore.GetConceptWeights(ctx)
|
||||
require.NoError(t, err)
|
||||
|
||||
assert.Equal(t, 2.0, retrieved["security"])
|
||||
assert.Equal(t, 1.5, retrieved["performance"])
|
||||
assert.Equal(t, 1.8, retrieved["best-practice"])
|
||||
|
||||
// Update weights (UPSERT)
|
||||
weights["security"] = 2.5
|
||||
weights["scalability"] = 1.2
|
||||
|
||||
err = obsStore.SetConceptWeights(ctx, weights)
|
||||
require.NoError(t, err)
|
||||
|
||||
retrieved, err = obsStore.GetConceptWeights(ctx)
|
||||
require.NoError(t, err)
|
||||
|
||||
assert.Equal(t, 2.5, retrieved["security"]) // updated
|
||||
assert.Equal(t, 1.5, retrieved["performance"]) // unchanged
|
||||
assert.Equal(t, 1.2, retrieved["scalability"]) // new
|
||||
assert.Equal(t, 1.8, retrieved["best-practice"]) // unchanged
|
||||
}
|
||||
|
||||
func TestObservationStore_GetObservationsNeedingScoreUpdate(t *testing.T) {
|
||||
obsStore, store, cleanup := testObservationStore(t)
|
||||
defer cleanup()
|
||||
ctx := context.Background()
|
||||
|
||||
sessionStore := NewSessionStore(store)
|
||||
sessionID, _ := sessionStore.CreateSDKSession(ctx, "claude-1", "test-project", "")
|
||||
|
||||
// Create observation with no score update
|
||||
obs1 := &models.ParsedObservation{Type: models.ObsTypeDiscovery, Title: "Needs Update"}
|
||||
obsID1, _, _ := obsStore.StoreObservation(ctx, "claude-1", "test-project", obs1, int(sessionID), 1)
|
||||
|
||||
// Create observation with recent score update
|
||||
obs2 := &models.ParsedObservation{Type: models.ObsTypeBugfix, Title: "Recently Updated"}
|
||||
obsID2, _, _ := obsStore.StoreObservation(ctx, "claude-1", "test-project", obs2, int(sessionID), 2)
|
||||
obsStore.UpdateImportanceScore(ctx, obsID2, 2.0)
|
||||
|
||||
// Get observations needing update (within 1 hour threshold)
|
||||
needsUpdate, err := obsStore.GetObservationsNeedingScoreUpdate(ctx, 1*time.Hour, 10)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Only obs1 should need update (obs2 was just updated)
|
||||
assert.Len(t, needsUpdate, 1)
|
||||
assert.Equal(t, obsID1, needsUpdate[0].ID)
|
||||
assert.True(t, needsUpdate[0].Title.Valid)
|
||||
assert.Equal(t, "Needs Update", needsUpdate[0].Title.String)
|
||||
}
|
||||
|
||||
func TestObservationStore_ResetObservationScores(t *testing.T) {
|
||||
obsStore, store, cleanup := testObservationStore(t)
|
||||
defer cleanup()
|
||||
ctx := context.Background()
|
||||
|
||||
sessionStore := NewSessionStore(store)
|
||||
sessionID, _ := sessionStore.CreateSDKSession(ctx, "claude-1", "test-project", "")
|
||||
|
||||
// Create observations with custom scores
|
||||
obs1 := &models.ParsedObservation{Type: models.ObsTypeDiscovery, Title: "Test1"}
|
||||
obsID1, _, _ := obsStore.StoreObservation(ctx, "claude-1", "test-project", obs1, int(sessionID), 1)
|
||||
obsStore.UpdateImportanceScore(ctx, obsID1, 5.0)
|
||||
|
||||
obs2 := &models.ParsedObservation{Type: models.ObsTypeBugfix, Title: "Test2"}
|
||||
obsID2, _, _ := obsStore.StoreObservation(ctx, "claude-1", "test-project", obs2, int(sessionID), 2)
|
||||
obsStore.UpdateImportanceScore(ctx, obsID2, 3.0)
|
||||
|
||||
// Reset all scores
|
||||
err := obsStore.ResetObservationScores(ctx)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Verify all scores are 1.0
|
||||
var dbObs1, dbObs2 Observation
|
||||
store.DB.First(&dbObs1, obsID1)
|
||||
store.DB.First(&dbObs2, obsID2)
|
||||
|
||||
assert.Equal(t, 1.0, dbObs1.ImportanceScore)
|
||||
assert.Equal(t, 1.0, dbObs2.ImportanceScore)
|
||||
}
|
||||
@@ -1,219 +0,0 @@
|
||||
// Package gorm provides GORM-based database operations for claude-mnemonic.
|
||||
package gorm
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/clause"
|
||||
|
||||
"github.com/lukaszraczylo/claude-mnemonic/pkg/models"
|
||||
)
|
||||
|
||||
// SessionStore provides session-related database operations using GORM.
|
||||
type SessionStore struct {
|
||||
db *gorm.DB
|
||||
}
|
||||
|
||||
// NewSessionStore creates a new session store.
|
||||
func NewSessionStore(store *Store) *SessionStore {
|
||||
return &SessionStore{db: store.DB}
|
||||
}
|
||||
|
||||
// CreateSDKSession creates a new SDK session (idempotent - returns existing ID if exists).
|
||||
// This is the KEY to how claude-mnemonic stays unified across hooks.
|
||||
func (s *SessionStore) CreateSDKSession(ctx context.Context, claudeSessionID, project, userPrompt string) (int64, error) {
|
||||
now := time.Now()
|
||||
|
||||
session := &SDKSession{
|
||||
ClaudeSessionID: claudeSessionID,
|
||||
SDKSessionID: func() sql.NullString {
|
||||
return sql.NullString{String: claudeSessionID, Valid: true}
|
||||
}(),
|
||||
Project: project,
|
||||
UserPrompt: func() sql.NullString {
|
||||
if userPrompt != "" {
|
||||
return sql.NullString{String: userPrompt, Valid: true}
|
||||
}
|
||||
return sql.NullString{Valid: false}
|
||||
}(),
|
||||
Status: "active",
|
||||
StartedAt: now.Format(time.RFC3339),
|
||||
StartedAtEpoch: now.UnixMilli(),
|
||||
}
|
||||
|
||||
// CRITICAL: INSERT OR IGNORE makes this idempotent
|
||||
// Use OnConflict with DoNothing to achieve INSERT OR IGNORE behavior
|
||||
result := s.db.WithContext(ctx).
|
||||
Clauses(clause.OnConflict{
|
||||
Columns: []clause.Column{{Name: "claude_session_id"}},
|
||||
DoNothing: true,
|
||||
}).
|
||||
Create(session)
|
||||
|
||||
if result.Error != nil {
|
||||
return 0, result.Error
|
||||
}
|
||||
|
||||
// Check if insert happened
|
||||
if result.RowsAffected == 0 {
|
||||
// Session exists - UPDATE project and user_prompt if we have non-empty values
|
||||
if project != "" {
|
||||
updates := map[string]interface{}{
|
||||
"project": project,
|
||||
}
|
||||
if userPrompt != "" {
|
||||
updates["user_prompt"] = userPrompt
|
||||
}
|
||||
if err := s.db.WithContext(ctx).
|
||||
Model(&SDKSession{}).
|
||||
Where("claude_session_id = ?", claudeSessionID).
|
||||
Updates(updates).Error; err != nil {
|
||||
return 0, fmt.Errorf("failed to update session: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
// Fetch existing session
|
||||
var existing SDKSession
|
||||
err := s.db.WithContext(ctx).
|
||||
Where("claude_session_id = ?", claudeSessionID).
|
||||
First(&existing).Error
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return existing.ID, nil
|
||||
}
|
||||
|
||||
return session.ID, nil
|
||||
}
|
||||
|
||||
// GetSessionByID retrieves a session by its database ID.
|
||||
func (s *SessionStore) GetSessionByID(ctx context.Context, id int64) (*models.SDKSession, error) {
|
||||
var sess SDKSession
|
||||
err := s.db.WithContext(ctx).First(&sess, id).Error
|
||||
if err == gorm.ErrRecordNotFound {
|
||||
return nil, nil
|
||||
}
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return toModelSDKSession(&sess), nil
|
||||
}
|
||||
|
||||
// FindAnySDKSession finds any session by Claude session ID (any status).
|
||||
func (s *SessionStore) FindAnySDKSession(ctx context.Context, claudeSessionID string) (*models.SDKSession, error) {
|
||||
var sess SDKSession
|
||||
err := s.db.WithContext(ctx).
|
||||
Where("claude_session_id = ?", claudeSessionID).
|
||||
First(&sess).Error
|
||||
if err == gorm.ErrRecordNotFound {
|
||||
return nil, nil
|
||||
}
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return toModelSDKSession(&sess), nil
|
||||
}
|
||||
|
||||
// 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) {
|
||||
// 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
|
||||
}
|
||||
|
||||
return newCounter, nil
|
||||
}
|
||||
|
||||
// GetPromptCounter returns the current prompt counter for a session.
|
||||
func (s *SessionStore) GetPromptCounter(ctx context.Context, id int64) (int, error) {
|
||||
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
|
||||
}
|
||||
|
||||
// GetSessionsToday returns the count of sessions started today.
|
||||
func (s *SessionStore) GetSessionsToday(ctx context.Context) (int, error) {
|
||||
// Get start of today in milliseconds
|
||||
now := time.Now()
|
||||
startOfDay := time.Date(now.Year(), now.Month(), now.Day(), 0, 0, 0, 0, now.Location())
|
||||
startEpoch := startOfDay.UnixMilli()
|
||||
|
||||
var count int64
|
||||
err := s.db.WithContext(ctx).
|
||||
Model(&SDKSession{}).
|
||||
Where("started_at_epoch >= ?", startEpoch).
|
||||
Count(&count).Error
|
||||
|
||||
return int(count), err
|
||||
}
|
||||
|
||||
// GetAllProjects returns all unique project names.
|
||||
func (s *SessionStore) GetAllProjects(ctx context.Context) ([]string, error) {
|
||||
var projects []string
|
||||
err := s.db.WithContext(ctx).
|
||||
Model(&SDKSession{}).
|
||||
Distinct("project").
|
||||
Where("project IS NOT NULL AND project != ''").
|
||||
Order("project ASC").
|
||||
Pluck("project", &projects).Error
|
||||
|
||||
return projects, err
|
||||
}
|
||||
|
||||
// toModelSDKSession converts a GORM SDKSession to pkg/models.SDKSession.
|
||||
func toModelSDKSession(sess *SDKSession) *models.SDKSession {
|
||||
return &models.SDKSession{
|
||||
ID: sess.ID,
|
||||
ClaudeSessionID: sess.ClaudeSessionID,
|
||||
SDKSessionID: sess.SDKSessionID,
|
||||
Project: sess.Project,
|
||||
UserPrompt: sess.UserPrompt,
|
||||
WorkerPort: sess.WorkerPort,
|
||||
PromptCounter: int64(sess.PromptCounter),
|
||||
Status: models.SessionStatus(sess.Status),
|
||||
StartedAt: sess.StartedAt,
|
||||
StartedAtEpoch: sess.StartedAtEpoch,
|
||||
CompletedAt: sess.CompletedAt,
|
||||
CompletedAtEpoch: sess.CompletedAtEpoch,
|
||||
}
|
||||
}
|
||||
@@ -1,259 +0,0 @@
|
||||
//go:build fts5
|
||||
|
||||
// Package gorm provides GORM-based database operations for claude-mnemonic.
|
||||
package gorm
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"gorm.io/gorm/logger"
|
||||
|
||||
"github.com/lukaszraczylo/claude-mnemonic/pkg/models"
|
||||
)
|
||||
|
||||
// testSessionStore creates a SessionStore with a temporary database for testing.
|
||||
func testSessionStore(t *testing.T) (*SessionStore, *Store, func()) {
|
||||
t.Helper()
|
||||
|
||||
tmpDir, err := os.MkdirTemp("", "gorm_session_test_*")
|
||||
if err != nil {
|
||||
t.Fatalf("create temp dir: %v", err)
|
||||
}
|
||||
|
||||
dbPath := filepath.Join(tmpDir, "test.db")
|
||||
cfg := Config{
|
||||
Path: dbPath,
|
||||
MaxConns: 4,
|
||||
LogLevel: logger.Silent,
|
||||
}
|
||||
|
||||
store, err := NewStore(cfg)
|
||||
if err != nil {
|
||||
os.RemoveAll(tmpDir)
|
||||
t.Fatalf("NewStore failed: %v", err)
|
||||
}
|
||||
|
||||
sessionStore := NewSessionStore(store)
|
||||
|
||||
cleanup := func() {
|
||||
store.Close()
|
||||
os.RemoveAll(tmpDir)
|
||||
}
|
||||
|
||||
return sessionStore, store, cleanup
|
||||
}
|
||||
|
||||
func TestSessionStore_CreateSDKSession(t *testing.T) {
|
||||
sessionStore, _, cleanup := testSessionStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Create a new session
|
||||
id, err := sessionStore.CreateSDKSession(ctx, "claude-1", "test-project", "initial prompt")
|
||||
require.NoError(t, err)
|
||||
assert.Greater(t, id, int64(0))
|
||||
|
||||
// Retrieve and verify
|
||||
sess, err := sessionStore.GetSessionByID(ctx, id)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, sess)
|
||||
assert.Equal(t, "claude-1", sess.ClaudeSessionID)
|
||||
assert.Equal(t, "test-project", sess.Project)
|
||||
assert.Equal(t, models.SessionStatusActive, sess.Status)
|
||||
assert.True(t, sess.UserPrompt.Valid)
|
||||
assert.Equal(t, "initial prompt", sess.UserPrompt.String)
|
||||
}
|
||||
|
||||
func TestSessionStore_CreateSDKSession_Idempotent(t *testing.T) {
|
||||
sessionStore, _, cleanup := testSessionStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Create first session
|
||||
id1, err := sessionStore.CreateSDKSession(ctx, "claude-1", "project-a", "prompt 1")
|
||||
require.NoError(t, err)
|
||||
|
||||
// Create again with same claude_session_id but different project
|
||||
id2, err := sessionStore.CreateSDKSession(ctx, "claude-1", "project-b", "prompt 2")
|
||||
require.NoError(t, err)
|
||||
|
||||
// Should return same ID (idempotent)
|
||||
assert.Equal(t, id1, id2)
|
||||
|
||||
// Should have updated project to project-b
|
||||
sess, err := sessionStore.GetSessionByID(ctx, id1)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "project-b", sess.Project)
|
||||
assert.Equal(t, "prompt 2", sess.UserPrompt.String)
|
||||
}
|
||||
|
||||
func TestSessionStore_CreateSDKSession_EmptyPrompt(t *testing.T) {
|
||||
sessionStore, _, cleanup := testSessionStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Create session with empty prompt
|
||||
id, err := sessionStore.CreateSDKSession(ctx, "claude-2", "test-project", "")
|
||||
require.NoError(t, err)
|
||||
assert.Greater(t, id, int64(0))
|
||||
|
||||
// Verify prompt is NULL
|
||||
sess, err := sessionStore.GetSessionByID(ctx, id)
|
||||
require.NoError(t, err)
|
||||
assert.False(t, sess.UserPrompt.Valid)
|
||||
}
|
||||
|
||||
func TestSessionStore_FindAnySDKSession(t *testing.T) {
|
||||
sessionStore, _, cleanup := testSessionStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Create a session
|
||||
_, err := sessionStore.CreateSDKSession(ctx, "claude-1", "test-project", "")
|
||||
require.NoError(t, err)
|
||||
|
||||
// Find it
|
||||
sess, err := sessionStore.FindAnySDKSession(ctx, "claude-1")
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, sess)
|
||||
assert.Equal(t, "claude-1", sess.ClaudeSessionID)
|
||||
|
||||
// Try to find non-existent
|
||||
sess, err = sessionStore.FindAnySDKSession(ctx, "claude-nonexistent")
|
||||
require.NoError(t, err)
|
||||
assert.Nil(t, sess)
|
||||
}
|
||||
|
||||
func TestSessionStore_GetSessionByID_NotFound(t *testing.T) {
|
||||
sessionStore, _, cleanup := testSessionStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Try to get non-existent session
|
||||
sess, err := sessionStore.GetSessionByID(ctx, 99999)
|
||||
require.NoError(t, err)
|
||||
assert.Nil(t, sess)
|
||||
}
|
||||
|
||||
func TestSessionStore_IncrementPromptCounter(t *testing.T) {
|
||||
sessionStore, _, cleanup := testSessionStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Create a session
|
||||
id, err := sessionStore.CreateSDKSession(ctx, "claude-1", "test-project", "")
|
||||
require.NoError(t, err)
|
||||
|
||||
// Initial counter should be 0
|
||||
counter, err := sessionStore.GetPromptCounter(ctx, id)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 0, counter)
|
||||
|
||||
// Increment
|
||||
counter, err = sessionStore.IncrementPromptCounter(ctx, id)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 1, counter)
|
||||
|
||||
// Increment again
|
||||
counter, err = sessionStore.IncrementPromptCounter(ctx, id)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 2, counter)
|
||||
|
||||
// Verify via GetPromptCounter
|
||||
counter, err = sessionStore.GetPromptCounter(ctx, id)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 2, counter)
|
||||
}
|
||||
|
||||
func TestSessionStore_GetSessionsToday(t *testing.T) {
|
||||
sessionStore, _, cleanup := testSessionStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Initially should be 0
|
||||
count, err := sessionStore.GetSessionsToday(ctx)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 0, count)
|
||||
|
||||
// Create some sessions
|
||||
_, err = sessionStore.CreateSDKSession(ctx, "claude-1", "project-1", "")
|
||||
require.NoError(t, err)
|
||||
|
||||
_, err = sessionStore.CreateSDKSession(ctx, "claude-2", "project-2", "")
|
||||
require.NoError(t, err)
|
||||
|
||||
// Should now have 2 sessions today
|
||||
count, err = sessionStore.GetSessionsToday(ctx)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 2, count)
|
||||
}
|
||||
|
||||
func TestSessionStore_GetAllProjects(t *testing.T) {
|
||||
sessionStore, _, cleanup := testSessionStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Initially should be empty
|
||||
projects, err := sessionStore.GetAllProjects(ctx)
|
||||
require.NoError(t, err)
|
||||
assert.Empty(t, projects)
|
||||
|
||||
// Create sessions with different projects
|
||||
_, err = sessionStore.CreateSDKSession(ctx, "claude-1", "project-a", "")
|
||||
require.NoError(t, err)
|
||||
|
||||
_, err = sessionStore.CreateSDKSession(ctx, "claude-2", "project-b", "")
|
||||
require.NoError(t, err)
|
||||
|
||||
_, err = sessionStore.CreateSDKSession(ctx, "claude-3", "project-a", "") // Duplicate project
|
||||
require.NoError(t, err)
|
||||
|
||||
// Should get distinct projects in alphabetical order
|
||||
projects, err = sessionStore.GetAllProjects(ctx)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, []string{"project-a", "project-b"}, projects)
|
||||
}
|
||||
|
||||
func TestSessionStore_SessionFields(t *testing.T) {
|
||||
sessionStore, _, cleanup := testSessionStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Create a session
|
||||
id, err := sessionStore.CreateSDKSession(ctx, "claude-1", "test-project", "test prompt")
|
||||
require.NoError(t, err)
|
||||
|
||||
// Retrieve and verify all fields
|
||||
sess, err := sessionStore.GetSessionByID(ctx, id)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, sess)
|
||||
|
||||
// Verify all fields
|
||||
assert.Equal(t, id, sess.ID)
|
||||
assert.Equal(t, "claude-1", sess.ClaudeSessionID)
|
||||
assert.True(t, sess.SDKSessionID.Valid)
|
||||
assert.Equal(t, "claude-1", sess.SDKSessionID.String) // Should be same as ClaudeSessionID
|
||||
assert.Equal(t, "test-project", sess.Project)
|
||||
assert.True(t, sess.UserPrompt.Valid)
|
||||
assert.Equal(t, "test prompt", sess.UserPrompt.String)
|
||||
assert.Equal(t, int64(0), sess.PromptCounter)
|
||||
assert.Equal(t, models.SessionStatusActive, sess.Status)
|
||||
assert.NotEmpty(t, sess.StartedAt)
|
||||
assert.Greater(t, sess.StartedAtEpoch, int64(0))
|
||||
assert.False(t, sess.CompletedAt.Valid) // Should be NULL
|
||||
assert.False(t, sess.CompletedAtEpoch.Valid) // Should be NULL
|
||||
}
|
||||
@@ -1,8 +0,0 @@
|
||||
//go:build !sqlite_omit_load_extension
|
||||
// +build !sqlite_omit_load_extension
|
||||
|
||||
// Package gorm provides GORM-based database operations for claude-mnemonic.
|
||||
package gorm
|
||||
|
||||
// This file ensures mattn/go-sqlite3 is built with FTS5 and other extensions enabled.
|
||||
// The build tag ensures extensions are not omitted.
|
||||
@@ -1,539 +0,0 @@
|
||||
// Package gorm provides GORM-based database operations for claude-mnemonic.
|
||||
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"
|
||||
)
|
||||
|
||||
// Store represents the GORM database connection with sqlite-vec support.
|
||||
type Store struct {
|
||||
healthCacheTime time.Time
|
||||
DB *gorm.DB
|
||||
sqlDB *sql.DB
|
||||
metrics *PoolMetrics
|
||||
cachedHealth *HealthInfo
|
||||
healthCacheTTL time.Duration
|
||||
healthCacheMu sync.RWMutex
|
||||
}
|
||||
|
||||
// Config holds database configuration.
|
||||
type Config struct {
|
||||
Path string // Path to SQLite database file
|
||||
MaxConns int // Maximum number of open connections (default: 4)
|
||||
LogLevel logger.LogLevel // GORM log level (logger.Silent for production)
|
||||
}
|
||||
|
||||
// NewStore creates a new Store with WAL mode enabled and sqlite-vec registered.
|
||||
// CRITICAL: WAL mode and foreign keys are enabled via pragmas for concurrent reads.
|
||||
func NewStore(cfg Config) (*Store, error) {
|
||||
// 1. Register sqlite-vec extension (must be done before opening database)
|
||||
sqlite_vec.Auto()
|
||||
|
||||
// 2. Build connection string (foreign keys enabled in DSN)
|
||||
// Use sqlite3 driver (mattn/go-sqlite3) which has FTS5 support
|
||||
dsn := cfg.Path + "?_foreign_keys=ON"
|
||||
|
||||
// 3. Open raw database connection with mattn/go-sqlite3 (has FTS5 support)
|
||||
sqlDB, err := sql.Open("sqlite3", dsn)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("open database: %w", err)
|
||||
}
|
||||
|
||||
// 4. Wrap with GORM using existing connection
|
||||
db, err := gorm.Open(sqlite.Dialector{
|
||||
Conn: sqlDB,
|
||||
}, &gorm.Config{
|
||||
Logger: logger.Default.LogMode(cfg.LogLevel),
|
||||
// PrepareStmt enables prepared statement caching for performance
|
||||
PrepareStmt: true,
|
||||
// Disable default timestamp fields (we manage created_at manually)
|
||||
NowFunc: nil,
|
||||
})
|
||||
if err != nil {
|
||||
_ = sqlDB.Close() // Explicitly ignore close error during cleanup
|
||||
return nil, fmt.Errorf("open gorm: %w", err)
|
||||
}
|
||||
|
||||
// 5. Configure connection pool (same settings as current implementation)
|
||||
maxConns := cfg.MaxConns
|
||||
if maxConns <= 0 {
|
||||
maxConns = 4
|
||||
}
|
||||
sqlDB.SetMaxOpenConns(maxConns)
|
||||
sqlDB.SetMaxIdleConns(maxConns)
|
||||
sqlDB.SetConnMaxLifetime(0) // Never expire (SQLite connections are cheap)
|
||||
|
||||
// 6. Verify connection
|
||||
if err := sqlDB.Ping(); err != nil {
|
||||
return nil, fmt.Errorf("ping database: %w", err)
|
||||
}
|
||||
|
||||
store := &Store{
|
||||
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)
|
||||
if err := runMigrations(db, sqlDB); err != nil {
|
||||
return nil, fmt.Errorf("run migrations: %w", err)
|
||||
}
|
||||
|
||||
// 8. CRITICAL: Set WAL mode and other performance pragmas
|
||||
// Use raw sqlDB to avoid GORM transaction issues
|
||||
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)
|
||||
}
|
||||
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()
|
||||
}
|
||||
|
||||
// Ping verifies the database connection is alive.
|
||||
func (s *Store) Ping() error {
|
||||
return s.sqlDB.Ping()
|
||||
}
|
||||
|
||||
// GetRawDB returns the underlying *sql.DB for operations GORM can't handle.
|
||||
// Use this for:
|
||||
// - FTS5 full-text search queries (MATCH operator)
|
||||
// - sqlite-vec vector operations
|
||||
// - Complex raw SQL queries
|
||||
func (s *Store) GetRawDB() *sql.DB {
|
||||
return s.sqlDB
|
||||
}
|
||||
|
||||
// GetDB returns the GORM DB instance for standard queries.
|
||||
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)
|
||||
})
|
||||
}
|
||||
@@ -1,152 +0,0 @@
|
||||
//go:build fts5
|
||||
|
||||
// Package gorm provides GORM-based database operations for claude-mnemonic.
|
||||
package gorm
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
"gorm.io/gorm/logger"
|
||||
)
|
||||
|
||||
func TestNewStore(t *testing.T) {
|
||||
// Create temporary directory for test database
|
||||
tmpDir, err := os.MkdirTemp("", "gorm_test_*")
|
||||
if err != nil {
|
||||
t.Fatalf("create temp dir: %v", err)
|
||||
}
|
||||
defer os.RemoveAll(tmpDir)
|
||||
|
||||
dbPath := filepath.Join(tmpDir, "test.db")
|
||||
|
||||
// Create store with migrations
|
||||
cfg := Config{
|
||||
Path: dbPath,
|
||||
MaxConns: 4,
|
||||
LogLevel: logger.Silent,
|
||||
}
|
||||
|
||||
store, err := NewStore(cfg)
|
||||
if err != nil {
|
||||
t.Fatalf("NewStore failed: %v", err)
|
||||
}
|
||||
defer store.Close()
|
||||
|
||||
// Verify connection works
|
||||
sqlDB := store.GetRawDB()
|
||||
if err := sqlDB.Ping(); err != nil {
|
||||
t.Fatalf("ping failed: %v", err)
|
||||
}
|
||||
|
||||
// Verify WAL mode is enabled
|
||||
var journalMode string
|
||||
err = store.DB.Raw("PRAGMA journal_mode").Scan(&journalMode).Error
|
||||
if err != nil {
|
||||
t.Fatalf("query journal_mode failed: %v", err)
|
||||
}
|
||||
if journalMode != "wal" {
|
||||
t.Errorf("expected WAL mode, got %q", journalMode)
|
||||
}
|
||||
|
||||
// Verify core tables exist
|
||||
tables := []string{
|
||||
"sdk_sessions",
|
||||
"observations",
|
||||
"session_summaries",
|
||||
"user_prompts",
|
||||
"observation_conflicts",
|
||||
"observation_relations",
|
||||
"patterns",
|
||||
"concept_weights",
|
||||
}
|
||||
|
||||
for _, table := range tables {
|
||||
exists := store.DB.Migrator().HasTable(table)
|
||||
if !exists {
|
||||
t.Errorf("table %q does not exist", table)
|
||||
}
|
||||
}
|
||||
|
||||
// Verify FTS5 virtual tables exist (cannot use Migrator().HasTable for virtual tables)
|
||||
ftsTables := []string{
|
||||
"user_prompts_fts",
|
||||
"observations_fts",
|
||||
"session_summaries_fts",
|
||||
"patterns_fts",
|
||||
}
|
||||
|
||||
for _, table := range ftsTables {
|
||||
var count int
|
||||
err := store.DB.Raw("SELECT COUNT(*) FROM sqlite_master WHERE type='table' AND name=?", table).Scan(&count).Error
|
||||
if err != nil {
|
||||
t.Errorf("check FTS table %q failed: %v", table, err)
|
||||
}
|
||||
if count != 1 {
|
||||
t.Errorf("FTS table %q does not exist", table)
|
||||
}
|
||||
}
|
||||
|
||||
// Verify vectors table exists (virtual table)
|
||||
var vectorsCount int
|
||||
err = store.DB.Raw("SELECT COUNT(*) FROM sqlite_master WHERE type='table' AND name='vectors'").Scan(&vectorsCount).Error
|
||||
if err != nil {
|
||||
t.Errorf("check vectors table failed: %v", err)
|
||||
}
|
||||
if vectorsCount != 1 {
|
||||
t.Errorf("vectors table does not exist")
|
||||
}
|
||||
|
||||
// Verify concept_weights seed data exists
|
||||
var conceptCount int64
|
||||
store.DB.Model(&ConceptWeight{}).Count(&conceptCount)
|
||||
if conceptCount != 12 {
|
||||
t.Errorf("expected 12 concept weights, got %d", conceptCount)
|
||||
}
|
||||
|
||||
t.Logf("✅ Phase 1 Foundation: All migrations successful")
|
||||
t.Logf(" - Core tables: %d", len(tables))
|
||||
t.Logf(" - FTS5 tables: %d", len(ftsTables))
|
||||
t.Logf(" - Vector table: 1")
|
||||
t.Logf(" - Seed data: %d concept weights", conceptCount)
|
||||
}
|
||||
|
||||
func TestMigrationIdempotency(t *testing.T) {
|
||||
// Create temporary directory for test database
|
||||
tmpDir, err := os.MkdirTemp("", "gorm_idempotency_*")
|
||||
if err != nil {
|
||||
t.Fatalf("create temp dir: %v", err)
|
||||
}
|
||||
defer os.RemoveAll(tmpDir)
|
||||
|
||||
dbPath := filepath.Join(tmpDir, "test.db")
|
||||
cfg := Config{
|
||||
Path: dbPath,
|
||||
MaxConns: 4,
|
||||
LogLevel: logger.Silent,
|
||||
}
|
||||
|
||||
// Run migrations first time
|
||||
store1, err := NewStore(cfg)
|
||||
if err != nil {
|
||||
t.Fatalf("NewStore (first) failed: %v", err)
|
||||
}
|
||||
store1.Close()
|
||||
|
||||
// Run migrations second time (should be idempotent)
|
||||
store2, err := NewStore(cfg)
|
||||
if err != nil {
|
||||
t.Fatalf("NewStore (second) failed: %v", err)
|
||||
}
|
||||
defer store2.Close()
|
||||
|
||||
// Verify concept_weights seed data is still exactly 12 (INSERT OR IGNORE)
|
||||
var conceptCount int64
|
||||
store2.DB.Model(&ConceptWeight{}).Count(&conceptCount)
|
||||
if conceptCount != 12 {
|
||||
t.Errorf("expected 12 concept weights after second migration, got %d", conceptCount)
|
||||
}
|
||||
|
||||
t.Logf("✅ Migrations are idempotent")
|
||||
}
|
||||
@@ -1,171 +0,0 @@
|
||||
// Package gorm provides GORM-based database operations for claude-mnemonic.
|
||||
package gorm
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"time"
|
||||
|
||||
"gorm.io/gorm"
|
||||
|
||||
"github.com/lukaszraczylo/claude-mnemonic/pkg/models"
|
||||
)
|
||||
|
||||
// SummaryStore provides summary-related database operations using GORM.
|
||||
type SummaryStore struct {
|
||||
db *gorm.DB
|
||||
}
|
||||
|
||||
// NewSummaryStore creates a new summary store.
|
||||
func NewSummaryStore(store *Store) *SummaryStore {
|
||||
return &SummaryStore{db: store.DB}
|
||||
}
|
||||
|
||||
// StoreSummary stores a new session summary.
|
||||
func (s *SummaryStore) StoreSummary(ctx context.Context, sdkSessionID, project string, summary *models.ParsedSummary, promptNumber int, discoveryTokens int64) (int64, int64, error) {
|
||||
now := time.Now()
|
||||
nowEpoch := now.UnixMilli()
|
||||
|
||||
// Ensure session exists (auto-create if missing)
|
||||
if err := EnsureSessionExists(ctx, s.db, sdkSessionID, project); err != nil {
|
||||
return 0, 0, err
|
||||
}
|
||||
|
||||
dbSummary := &SessionSummary{
|
||||
SDKSessionID: sdkSessionID,
|
||||
Project: project,
|
||||
Request: nullString(summary.Request),
|
||||
Investigated: nullString(summary.Investigated),
|
||||
Learned: nullString(summary.Learned),
|
||||
Completed: nullString(summary.Completed),
|
||||
NextSteps: nullString(summary.NextSteps),
|
||||
Notes: nullString(summary.Notes),
|
||||
PromptNumber: func() sql.NullInt64 {
|
||||
if promptNumber > 0 {
|
||||
return sql.NullInt64{Int64: int64(promptNumber), Valid: true}
|
||||
}
|
||||
return sql.NullInt64{Valid: false}
|
||||
}(),
|
||||
DiscoveryTokens: discoveryTokens,
|
||||
CreatedAt: now.Format(time.RFC3339),
|
||||
CreatedAtEpoch: nowEpoch,
|
||||
}
|
||||
|
||||
err := s.db.WithContext(ctx).Create(dbSummary).Error
|
||||
if err != nil {
|
||||
return 0, 0, err
|
||||
}
|
||||
|
||||
return dbSummary.ID, nowEpoch, nil
|
||||
}
|
||||
|
||||
// GetSummariesByIDs retrieves summaries by a list of IDs.
|
||||
func (s *SummaryStore) GetSummariesByIDs(ctx context.Context, ids []int64, orderBy string, limit int) ([]*models.SessionSummary, error) {
|
||||
if len(ids) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
var dbSummaries []SessionSummary
|
||||
query := s.db.WithContext(ctx).Where("id IN ?", ids)
|
||||
|
||||
// Apply ordering
|
||||
switch orderBy {
|
||||
case "date_asc":
|
||||
query = query.Order("created_at_epoch ASC")
|
||||
case "date_desc", "default", "":
|
||||
query = query.Order("created_at_epoch DESC")
|
||||
}
|
||||
|
||||
// Apply limit
|
||||
if limit > 0 {
|
||||
query = query.Limit(limit)
|
||||
}
|
||||
|
||||
err := query.Find(&dbSummaries).Error
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return toModelSessionSummaries(dbSummaries), nil
|
||||
}
|
||||
|
||||
// GetRecentSummaries retrieves recent summaries for a project.
|
||||
func (s *SummaryStore) GetRecentSummaries(ctx context.Context, project string, limit int) ([]*models.SessionSummary, error) {
|
||||
var dbSummaries []SessionSummary
|
||||
err := s.db.WithContext(ctx).
|
||||
Where("project = ?", project).
|
||||
Order("created_at_epoch DESC").
|
||||
Limit(limit).
|
||||
Find(&dbSummaries).Error
|
||||
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return toModelSessionSummaries(dbSummaries), nil
|
||||
}
|
||||
|
||||
// GetAllRecentSummaries retrieves recent summaries across all projects.
|
||||
func (s *SummaryStore) GetAllRecentSummaries(ctx context.Context, limit int) ([]*models.SessionSummary, error) {
|
||||
var dbSummaries []SessionSummary
|
||||
err := s.db.WithContext(ctx).
|
||||
Order("created_at_epoch DESC").
|
||||
Limit(limit).
|
||||
Find(&dbSummaries).Error
|
||||
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return toModelSessionSummaries(dbSummaries), nil
|
||||
}
|
||||
|
||||
// GetAllSummaries retrieves all summaries (for vector rebuild).
|
||||
func (s *SummaryStore) GetAllSummaries(ctx context.Context) ([]*models.SessionSummary, error) {
|
||||
var dbSummaries []SessionSummary
|
||||
err := s.db.WithContext(ctx).
|
||||
Order("id").
|
||||
Find(&dbSummaries).Error
|
||||
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return toModelSessionSummaries(dbSummaries), nil
|
||||
}
|
||||
|
||||
// toModelSessionSummary converts a GORM SessionSummary to pkg/models.SessionSummary.
|
||||
func toModelSessionSummary(s *SessionSummary) *models.SessionSummary {
|
||||
return &models.SessionSummary{
|
||||
ID: s.ID,
|
||||
SDKSessionID: s.SDKSessionID,
|
||||
Project: s.Project,
|
||||
Request: s.Request,
|
||||
Investigated: s.Investigated,
|
||||
Learned: s.Learned,
|
||||
Completed: s.Completed,
|
||||
NextSteps: s.NextSteps,
|
||||
Notes: s.Notes,
|
||||
PromptNumber: s.PromptNumber,
|
||||
DiscoveryTokens: s.DiscoveryTokens,
|
||||
CreatedAt: s.CreatedAt,
|
||||
CreatedAtEpoch: s.CreatedAtEpoch,
|
||||
}
|
||||
}
|
||||
|
||||
// toModelSessionSummaries converts a slice of GORM SessionSummary to pkg/models.SessionSummary.
|
||||
func toModelSessionSummaries(summaries []SessionSummary) []*models.SessionSummary {
|
||||
result := make([]*models.SessionSummary, len(summaries))
|
||||
for i := range summaries {
|
||||
result[i] = toModelSessionSummary(&summaries[i])
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
// nullString converts a string to sql.NullString.
|
||||
func nullString(s string) sql.NullString {
|
||||
if s == "" {
|
||||
return sql.NullString{Valid: false}
|
||||
}
|
||||
return sql.NullString{String: s, Valid: true}
|
||||
}
|
||||
@@ -1,278 +0,0 @@
|
||||
//go:build fts5
|
||||
|
||||
// Package gorm provides GORM-based database operations for claude-mnemonic.
|
||||
package gorm
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"gorm.io/gorm/logger"
|
||||
|
||||
"github.com/lukaszraczylo/claude-mnemonic/pkg/models"
|
||||
)
|
||||
|
||||
// testSummaryStore creates a SummaryStore with a temporary database for testing.
|
||||
func testSummaryStore(t *testing.T) (*SummaryStore, *Store, func()) {
|
||||
t.Helper()
|
||||
|
||||
tmpDir, err := os.MkdirTemp("", "gorm_summary_test_*")
|
||||
if err != nil {
|
||||
t.Fatalf("create temp dir: %v", err)
|
||||
}
|
||||
|
||||
dbPath := filepath.Join(tmpDir, "test.db")
|
||||
cfg := Config{
|
||||
Path: dbPath,
|
||||
MaxConns: 4,
|
||||
LogLevel: logger.Silent,
|
||||
}
|
||||
|
||||
store, err := NewStore(cfg)
|
||||
if err != nil {
|
||||
os.RemoveAll(tmpDir)
|
||||
t.Fatalf("NewStore failed: %v", err)
|
||||
}
|
||||
|
||||
summaryStore := NewSummaryStore(store)
|
||||
|
||||
cleanup := func() {
|
||||
store.Close()
|
||||
os.RemoveAll(tmpDir)
|
||||
}
|
||||
|
||||
return summaryStore, store, cleanup
|
||||
}
|
||||
|
||||
func TestSummaryStore_StoreSummary(t *testing.T) {
|
||||
summaryStore, store, cleanup := testSummaryStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Create a session first
|
||||
sessionStore := NewSessionStore(store)
|
||||
_, err := sessionStore.CreateSDKSession(ctx, "claude-1", "test-project", "")
|
||||
require.NoError(t, err)
|
||||
|
||||
// Store a summary
|
||||
summary := &models.ParsedSummary{
|
||||
Request: "Build a feature",
|
||||
Investigated: "Examined the codebase",
|
||||
Learned: "Discovered patterns",
|
||||
Completed: "Implemented solution",
|
||||
NextSteps: "Write tests",
|
||||
Notes: "Additional notes",
|
||||
}
|
||||
|
||||
id, epoch, err := summaryStore.StoreSummary(ctx, "claude-1", "test-project", summary, 1, 100)
|
||||
require.NoError(t, err)
|
||||
assert.Greater(t, id, int64(0))
|
||||
assert.Greater(t, epoch, int64(0))
|
||||
}
|
||||
|
||||
func TestSummaryStore_StoreSummary_AutoCreateSession(t *testing.T) {
|
||||
summaryStore, _, cleanup := testSummaryStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Store summary without pre-creating session
|
||||
summary := &models.ParsedSummary{
|
||||
Request: "Test auto-create",
|
||||
}
|
||||
|
||||
id, _, err := summaryStore.StoreSummary(ctx, "claude-auto", "auto-project", summary, 1, 50)
|
||||
require.NoError(t, err)
|
||||
assert.Greater(t, id, int64(0))
|
||||
}
|
||||
|
||||
func TestSummaryStore_GetRecentSummaries(t *testing.T) {
|
||||
summaryStore, _, cleanup := testSummaryStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Store multiple summaries
|
||||
for i := 1; i <= 5; i++ {
|
||||
summary := &models.ParsedSummary{
|
||||
Request: "Request " + string(rune('0'+i)),
|
||||
}
|
||||
_, _, err := summaryStore.StoreSummary(ctx, "claude-1", "project-a", summary, i, 10)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
// Store summary for different project
|
||||
summary := &models.ParsedSummary{Request: "Other project"}
|
||||
_, _, err := summaryStore.StoreSummary(ctx, "claude-2", "project-b", summary, 1, 10)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Get recent summaries for project-a
|
||||
summaries, err := summaryStore.GetRecentSummaries(ctx, "project-a", 10)
|
||||
require.NoError(t, err)
|
||||
assert.Len(t, summaries, 5)
|
||||
|
||||
// Verify ordering (most recent first)
|
||||
assert.Equal(t, "project-a", summaries[0].Project)
|
||||
}
|
||||
|
||||
func TestSummaryStore_GetAllRecentSummaries(t *testing.T) {
|
||||
summaryStore, _, cleanup := testSummaryStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Store summaries across projects
|
||||
_, _, err := summaryStore.StoreSummary(ctx, "claude-1", "project-a", &models.ParsedSummary{Request: "A1"}, 1, 10)
|
||||
require.NoError(t, err)
|
||||
|
||||
_, _, err = summaryStore.StoreSummary(ctx, "claude-2", "project-b", &models.ParsedSummary{Request: "B1"}, 1, 10)
|
||||
require.NoError(t, err)
|
||||
|
||||
_, _, err = summaryStore.StoreSummary(ctx, "claude-3", "project-c", &models.ParsedSummary{Request: "C1"}, 1, 10)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Get all recent summaries
|
||||
summaries, err := summaryStore.GetAllRecentSummaries(ctx, 10)
|
||||
require.NoError(t, err)
|
||||
assert.Len(t, summaries, 3)
|
||||
}
|
||||
|
||||
func TestSummaryStore_GetSummariesByIDs(t *testing.T) {
|
||||
summaryStore, _, cleanup := testSummaryStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Store multiple summaries
|
||||
var ids []int64
|
||||
for i := 1; i <= 3; i++ {
|
||||
id, _, err := summaryStore.StoreSummary(ctx, "claude-1", "project-a", &models.ParsedSummary{Request: "Test"}, i, 10)
|
||||
require.NoError(t, err)
|
||||
ids = append(ids, id)
|
||||
}
|
||||
|
||||
// Get by IDs
|
||||
summaries, err := summaryStore.GetSummariesByIDs(ctx, ids, "date_desc", 10)
|
||||
require.NoError(t, err)
|
||||
assert.Len(t, summaries, 3)
|
||||
|
||||
// Get with limit
|
||||
summaries, err = summaryStore.GetSummariesByIDs(ctx, ids, "date_desc", 2)
|
||||
require.NoError(t, err)
|
||||
assert.Len(t, summaries, 2)
|
||||
}
|
||||
|
||||
func TestSummaryStore_GetSummariesByIDs_EmptyInput(t *testing.T) {
|
||||
summaryStore, _, cleanup := testSummaryStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Get with empty IDs
|
||||
summaries, err := summaryStore.GetSummariesByIDs(ctx, []int64{}, "date_desc", 10)
|
||||
require.NoError(t, err)
|
||||
assert.Nil(t, summaries)
|
||||
}
|
||||
|
||||
func TestSummaryStore_SummaryFields(t *testing.T) {
|
||||
summaryStore, _, cleanup := testSummaryStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Store a summary with all fields
|
||||
summary := &models.ParsedSummary{
|
||||
Request: "Full request",
|
||||
Investigated: "Full investigation",
|
||||
Learned: "Full learning",
|
||||
Completed: "Full completion",
|
||||
NextSteps: "Full next steps",
|
||||
Notes: "Full notes",
|
||||
}
|
||||
|
||||
id, epoch, err := summaryStore.StoreSummary(ctx, "claude-1", "test-project", summary, 5, 200)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Retrieve and verify all fields
|
||||
summaries, err := summaryStore.GetSummariesByIDs(ctx, []int64{id}, "date_desc", 1)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, summaries, 1)
|
||||
|
||||
s := summaries[0]
|
||||
assert.Equal(t, id, s.ID)
|
||||
assert.Equal(t, "claude-1", s.SDKSessionID)
|
||||
assert.Equal(t, "test-project", s.Project)
|
||||
assert.True(t, s.Request.Valid)
|
||||
assert.Equal(t, "Full request", s.Request.String)
|
||||
assert.True(t, s.Investigated.Valid)
|
||||
assert.Equal(t, "Full investigation", s.Investigated.String)
|
||||
assert.True(t, s.Learned.Valid)
|
||||
assert.Equal(t, "Full learning", s.Learned.String)
|
||||
assert.True(t, s.Completed.Valid)
|
||||
assert.Equal(t, "Full completion", s.Completed.String)
|
||||
assert.True(t, s.NextSteps.Valid)
|
||||
assert.Equal(t, "Full next steps", s.NextSteps.String)
|
||||
assert.True(t, s.Notes.Valid)
|
||||
assert.Equal(t, "Full notes", s.Notes.String)
|
||||
assert.True(t, s.PromptNumber.Valid)
|
||||
assert.Equal(t, int64(5), s.PromptNumber.Int64)
|
||||
assert.Equal(t, int64(200), s.DiscoveryTokens)
|
||||
assert.NotEmpty(t, s.CreatedAt)
|
||||
assert.Equal(t, epoch, s.CreatedAtEpoch)
|
||||
}
|
||||
|
||||
func TestSummaryStore_EmptySummary(t *testing.T) {
|
||||
summaryStore, _, cleanup := testSummaryStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Store a summary with empty fields
|
||||
summary := &models.ParsedSummary{}
|
||||
|
||||
id, _, err := summaryStore.StoreSummary(ctx, "claude-1", "test-project", summary, 0, 0)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Retrieve and verify NULL fields
|
||||
summaries, err := summaryStore.GetSummariesByIDs(ctx, []int64{id}, "date_desc", 1)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, summaries, 1)
|
||||
|
||||
s := summaries[0]
|
||||
assert.False(t, s.Request.Valid)
|
||||
assert.False(t, s.Investigated.Valid)
|
||||
assert.False(t, s.Learned.Valid)
|
||||
assert.False(t, s.Completed.Valid)
|
||||
assert.False(t, s.NextSteps.Valid)
|
||||
assert.False(t, s.Notes.Valid)
|
||||
assert.False(t, s.PromptNumber.Valid)
|
||||
assert.Equal(t, int64(0), s.DiscoveryTokens)
|
||||
}
|
||||
|
||||
func TestSummaryStore_GetAllSummaries(t *testing.T) {
|
||||
summaryStore, _, cleanup := testSummaryStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Store multiple summaries
|
||||
for i := 1; i <= 5; i++ {
|
||||
_, _, err := summaryStore.StoreSummary(ctx, "claude-1", "project-a", &models.ParsedSummary{Request: "Test"}, i, 10)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
// Get all summaries
|
||||
summaries, err := summaryStore.GetAllSummaries(ctx)
|
||||
require.NoError(t, err)
|
||||
assert.Len(t, summaries, 5)
|
||||
|
||||
// Verify ordering by ID
|
||||
for i := 0; i < len(summaries)-1; i++ {
|
||||
assert.Less(t, summaries[i].ID, summaries[i+1].ID)
|
||||
}
|
||||
}
|
||||
@@ -1,71 +0,0 @@
|
||||
// Package db defines database interfaces for the claude-mnemonic stores.
|
||||
package db
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
"github.com/lukaszraczylo/claude-mnemonic/pkg/models"
|
||||
)
|
||||
|
||||
// ObservationReader defines read operations for observations.
|
||||
type ObservationReader interface {
|
||||
GetObservationByID(ctx context.Context, id int64) (*models.Observation, error)
|
||||
GetObservationsByIDs(ctx context.Context, ids []int64, orderBy string, limit int) ([]*models.Observation, error)
|
||||
GetRecentObservations(ctx context.Context, project string, limit int) ([]*models.Observation, error)
|
||||
GetActiveObservations(ctx context.Context, project string, limit int) ([]*models.Observation, error)
|
||||
GetAllRecentObservations(ctx context.Context, limit int) ([]*models.Observation, error)
|
||||
GetAllObservations(ctx context.Context) ([]*models.Observation, error)
|
||||
SearchObservationsFTS(ctx context.Context, query, project string, limit int) ([]*models.Observation, error)
|
||||
GetObservationCount(ctx context.Context, project string) (int, error)
|
||||
}
|
||||
|
||||
// ObservationWriter defines write operations for observations.
|
||||
type ObservationWriter interface {
|
||||
StoreObservation(ctx context.Context, sdkSessionID, project string, obs *models.ParsedObservation, promptNumber int, discoveryTokens int64) (int64, int64, error)
|
||||
DeleteObservations(ctx context.Context, ids []int64) (int64, error)
|
||||
}
|
||||
|
||||
// ObservationStore combines read and write operations for observations.
|
||||
type ObservationStore interface {
|
||||
ObservationReader
|
||||
ObservationWriter
|
||||
}
|
||||
|
||||
// SummaryReader defines read operations for summaries.
|
||||
type SummaryReader interface {
|
||||
GetSummariesByIDs(ctx context.Context, ids []int64, orderBy string, limit int) ([]*models.SessionSummary, error)
|
||||
GetRecentSummaries(ctx context.Context, project string, limit int) ([]*models.SessionSummary, error)
|
||||
GetAllRecentSummaries(ctx context.Context, limit int) ([]*models.SessionSummary, error)
|
||||
GetAllSummaries(ctx context.Context) ([]*models.SessionSummary, error)
|
||||
}
|
||||
|
||||
// SummaryWriter defines write operations for summaries.
|
||||
type SummaryWriter interface {
|
||||
StoreSummary(ctx context.Context, sdkSessionID, project string, summary *models.ParsedSummary, promptNumber int, discoveryTokens int64) (int64, int64, error)
|
||||
}
|
||||
|
||||
// SummaryStore combines read and write operations for summaries.
|
||||
type SummaryStore interface {
|
||||
SummaryReader
|
||||
SummaryWriter
|
||||
}
|
||||
|
||||
// PromptReader defines read operations for prompts.
|
||||
type PromptReader interface {
|
||||
GetPromptsByIDs(ctx context.Context, ids []int64, orderBy string, limit int) ([]*models.UserPromptWithSession, error)
|
||||
GetAllRecentUserPrompts(ctx context.Context, limit int) ([]*models.UserPromptWithSession, error)
|
||||
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)
|
||||
}
|
||||
|
||||
// PromptWriter defines write operations for prompts.
|
||||
type PromptWriter interface {
|
||||
SaveUserPromptWithMatches(ctx context.Context, claudeSessionID string, promptNumber int, promptText string, matchedObservations int) (int64, error)
|
||||
}
|
||||
|
||||
// PromptStore combines read and write operations for prompts.
|
||||
type PromptStore interface {
|
||||
PromptReader
|
||||
PromptWriter
|
||||
}
|
||||
@@ -0,0 +1,160 @@
|
||||
// Package sqlite provides SQLite database operations for claude-mnemonic.
|
||||
package sqlite
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"time"
|
||||
|
||||
"github.com/lukaszraczylo/claude-mnemonic/pkg/models"
|
||||
)
|
||||
|
||||
// EnsureSessionExists creates a session if it doesn't exist.
|
||||
// This is shared between ObservationStore and SummaryStore to avoid duplication.
|
||||
func EnsureSessionExists(ctx context.Context, store *Store, sdkSessionID, project string) error {
|
||||
const checkQuery = `SELECT id FROM sdk_sessions WHERE sdk_session_id = ?`
|
||||
var id int64
|
||||
err := store.QueryRowContext(ctx, checkQuery, sdkSessionID).Scan(&id)
|
||||
if err == nil {
|
||||
return nil // Session exists
|
||||
}
|
||||
if err != sql.ErrNoRows {
|
||||
return err
|
||||
}
|
||||
|
||||
// Auto-create session
|
||||
now := time.Now()
|
||||
const insertQuery = `
|
||||
INSERT INTO sdk_sessions
|
||||
(claude_session_id, sdk_session_id, project, started_at, started_at_epoch, status)
|
||||
VALUES (?, ?, ?, ?, ?, 'active')
|
||||
`
|
||||
_, err = store.ExecContext(ctx, insertQuery,
|
||||
sdkSessionID, sdkSessionID, project,
|
||||
now.Format(time.RFC3339), now.UnixMilli(),
|
||||
)
|
||||
return err
|
||||
}
|
||||
|
||||
// nullString converts a string to sql.NullString.
|
||||
func nullString(s string) sql.NullString {
|
||||
return sql.NullString{String: s, Valid: s != ""}
|
||||
}
|
||||
|
||||
// nullInt converts an int to sql.NullInt64.
|
||||
func nullInt(i int) sql.NullInt64 {
|
||||
return sql.NullInt64{Int64: int64(i), Valid: i > 0}
|
||||
}
|
||||
|
||||
// repeatPlaceholders generates n comma-prefixed placeholders for SQL IN clauses.
|
||||
// e.g., repeatPlaceholders(2) returns ", ?, ?"
|
||||
func repeatPlaceholders(n int) string {
|
||||
if n <= 0 {
|
||||
return ""
|
||||
}
|
||||
result := ""
|
||||
for i := 0; i < n; i++ {
|
||||
result += ", ?"
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
// int64SliceToInterface converts []int64 to []interface{} for SQL queries.
|
||||
func int64SliceToInterface(ids []int64) []interface{} {
|
||||
args := make([]interface{}, len(ids))
|
||||
for i, id := range ids {
|
||||
args[i] = id
|
||||
}
|
||||
return args
|
||||
}
|
||||
|
||||
// ParseLimitParam parses a limit query parameter with a default value.
|
||||
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 {
|
||||
return parsed
|
||||
}
|
||||
}
|
||||
return defaultLimit
|
||||
}
|
||||
|
||||
// scanSummary scans a single summary from a row scanner.
|
||||
func scanSummary(scanner interface{ Scan(...interface{}) error }) (*models.SessionSummary, error) {
|
||||
var summary models.SessionSummary
|
||||
if err := scanner.Scan(
|
||||
&summary.ID, &summary.SDKSessionID, &summary.Project,
|
||||
&summary.Request, &summary.Investigated, &summary.Learned, &summary.Completed,
|
||||
&summary.NextSteps, &summary.Notes, &summary.PromptNumber, &summary.DiscoveryTokens,
|
||||
&summary.CreatedAt, &summary.CreatedAtEpoch,
|
||||
); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &summary, nil
|
||||
}
|
||||
|
||||
// scanSummaryRows scans multiple summaries from rows.
|
||||
func scanSummaryRows(rows *sql.Rows) ([]*models.SessionSummary, error) {
|
||||
var summaries []*models.SessionSummary
|
||||
for rows.Next() {
|
||||
summary, err := scanSummary(rows)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
summaries = append(summaries, summary)
|
||||
}
|
||||
return summaries, rows.Err()
|
||||
}
|
||||
|
||||
// scanPromptWithSession scans a single prompt with session info from a row scanner.
|
||||
func scanPromptWithSession(scanner interface{ Scan(...interface{}) error }) (*models.UserPromptWithSession, error) {
|
||||
var prompt models.UserPromptWithSession
|
||||
if err := scanner.Scan(
|
||||
&prompt.ID, &prompt.ClaudeSessionID, &prompt.PromptNumber, &prompt.PromptText,
|
||||
&prompt.MatchedObservations, &prompt.CreatedAt, &prompt.CreatedAtEpoch,
|
||||
&prompt.Project, &prompt.SDKSessionID,
|
||||
); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &prompt, nil
|
||||
}
|
||||
|
||||
// scanPromptWithSessionRows scans multiple prompts with session info from rows.
|
||||
func scanPromptWithSessionRows(rows *sql.Rows) ([]*models.UserPromptWithSession, error) {
|
||||
var prompts []*models.UserPromptWithSession
|
||||
for rows.Next() {
|
||||
prompt, err := scanPromptWithSession(rows)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
prompts = append(prompts, prompt)
|
||||
}
|
||||
return prompts, rows.Err()
|
||||
}
|
||||
|
||||
// BuildGetByIDsQuery builds a query for fetching records by IDs with optional ordering and limit.
|
||||
// Returns the query string and args slice.
|
||||
func BuildGetByIDsQuery(baseQuery string, ids []int64, orderBy string, limit int) (string, []interface{}) {
|
||||
// Build query with placeholders
|
||||
// #nosec G202 -- query uses parameterized placeholders, not user input
|
||||
query := baseQuery + ` WHERE id IN (?` + repeatPlaceholders(len(ids)-1) + `)
|
||||
ORDER BY created_at_epoch `
|
||||
|
||||
if orderBy == "date_asc" {
|
||||
query += "ASC"
|
||||
} else {
|
||||
query += "DESC"
|
||||
}
|
||||
|
||||
if limit > 0 {
|
||||
query += " LIMIT ?"
|
||||
}
|
||||
|
||||
args := int64SliceToInterface(ids)
|
||||
if limit > 0 {
|
||||
args = append(args, limit)
|
||||
}
|
||||
|
||||
return query, args
|
||||
}
|
||||
@@ -0,0 +1,254 @@
|
||||
package sqlite
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestNullString(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
input string
|
||||
expected string
|
||||
valid bool
|
||||
}{
|
||||
{"empty_string", "", "", false},
|
||||
{"non_empty_string", "hello", "hello", true},
|
||||
{"whitespace", " ", " ", true},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
result := nullString(tt.input)
|
||||
assert.Equal(t, tt.expected, result.String)
|
||||
assert.Equal(t, tt.valid, result.Valid)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestNullInt(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
input int
|
||||
expected int64
|
||||
valid bool
|
||||
}{
|
||||
{"zero", 0, 0, false},
|
||||
{"positive", 42, 42, true},
|
||||
{"negative", -1, -1, false},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
result := nullInt(tt.input)
|
||||
assert.Equal(t, tt.expected, result.Int64)
|
||||
assert.Equal(t, tt.valid, result.Valid)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestRepeatPlaceholders(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
n int
|
||||
expected string
|
||||
}{
|
||||
{"zero", 0, ""},
|
||||
{"negative", -1, ""},
|
||||
{"one", 1, ", ?"},
|
||||
{"two", 2, ", ?, ?"},
|
||||
{"three", 3, ", ?, ?, ?"},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
result := repeatPlaceholders(tt.n)
|
||||
assert.Equal(t, tt.expected, result)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestInt64SliceToInterface(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
input []int64
|
||||
expected []interface{}
|
||||
}{
|
||||
{"empty", []int64{}, []interface{}{}},
|
||||
{"single", []int64{42}, []interface{}{int64(42)}},
|
||||
{"multiple", []int64{1, 2, 3}, []interface{}{int64(1), int64(2), int64(3)}},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
result := int64SliceToInterface(tt.input)
|
||||
assert.Equal(t, tt.expected, result)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseLimitParam(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
query string
|
||||
defaultLimit int
|
||||
expected int
|
||||
}{
|
||||
{"no_param_uses_default", "", 10, 10},
|
||||
{"valid_limit", "limit=20", 10, 20},
|
||||
{"invalid_limit_uses_default", "limit=abc", 10, 10},
|
||||
{"zero_limit_uses_default", "limit=0", 10, 10},
|
||||
{"negative_limit_uses_default", "limit=-5", 10, 10},
|
||||
{"large_limit", "limit=1000", 10, 1000},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
req := httptest.NewRequest("GET", "/?"+tt.query, nil)
|
||||
result := ParseLimitParam(req, tt.defaultLimit)
|
||||
assert.Equal(t, tt.expected, result)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestScanSummary(t *testing.T) {
|
||||
db, _, cleanup := testDB(t)
|
||||
defer cleanup()
|
||||
createBaseTables(t, db)
|
||||
seedSession(t, db, "claude-123", "sdk-123", "test-project")
|
||||
|
||||
// Insert a test summary
|
||||
_, err := db.Exec(`
|
||||
INSERT INTO session_summaries (sdk_session_id, project, request, investigated, learned, completed, next_steps, notes, prompt_number, discovery_tokens, created_at, created_at_epoch)
|
||||
VALUES ('sdk-123', 'test-project', 'test request', 'test investigated', 'test learned', 'test completed', 'test next steps', 'test notes', 1, 100, '2025-01-01T00:00:00Z', 1704067200000)
|
||||
`)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Query and scan
|
||||
row := db.QueryRow(`
|
||||
SELECT id, sdk_session_id, project, request, investigated, learned, completed, next_steps, notes, prompt_number, discovery_tokens, created_at, created_at_epoch
|
||||
FROM session_summaries WHERE sdk_session_id = ?
|
||||
`, "sdk-123")
|
||||
|
||||
summary, err := scanSummary(row)
|
||||
require.NoError(t, err)
|
||||
assert.NotNil(t, summary)
|
||||
assert.Equal(t, "sdk-123", summary.SDKSessionID)
|
||||
assert.Equal(t, "test-project", summary.Project)
|
||||
assert.Equal(t, "test request", summary.Request.String)
|
||||
assert.Equal(t, "test investigated", summary.Investigated.String)
|
||||
assert.Equal(t, "test learned", summary.Learned.String)
|
||||
assert.Equal(t, "test completed", summary.Completed.String)
|
||||
assert.Equal(t, "test next steps", summary.NextSteps.String)
|
||||
assert.Equal(t, "test notes", summary.Notes.String)
|
||||
}
|
||||
|
||||
func TestScanSummaryRows(t *testing.T) {
|
||||
db, _, cleanup := testDB(t)
|
||||
defer cleanup()
|
||||
createBaseTables(t, db)
|
||||
seedSession(t, db, "claude-123", "sdk-123", "test-project")
|
||||
|
||||
// Insert multiple summaries
|
||||
_, err := db.Exec(`
|
||||
INSERT INTO session_summaries (sdk_session_id, project, request, investigated, learned, completed, next_steps, notes, prompt_number, discovery_tokens, created_at, created_at_epoch)
|
||||
VALUES
|
||||
('sdk-123', 'test-project', 'request 1', '', '', '', '', '', 1, 0, '2025-01-01T00:00:00Z', 1704067200000),
|
||||
('sdk-123', 'test-project', 'request 2', '', '', '', '', '', 2, 0, '2025-01-02T00:00:00Z', 1704153600000)
|
||||
`)
|
||||
require.NoError(t, err)
|
||||
|
||||
rows, err := db.Query(`
|
||||
SELECT id, sdk_session_id, project, request, investigated, learned, completed, next_steps, notes, prompt_number, discovery_tokens, created_at, created_at_epoch
|
||||
FROM session_summaries WHERE sdk_session_id = ? ORDER BY id
|
||||
`, "sdk-123")
|
||||
require.NoError(t, err)
|
||||
defer rows.Close()
|
||||
|
||||
summaries, err := scanSummaryRows(rows)
|
||||
require.NoError(t, err)
|
||||
assert.Len(t, summaries, 2)
|
||||
assert.Equal(t, "request 1", summaries[0].Request.String)
|
||||
assert.Equal(t, "request 2", summaries[1].Request.String)
|
||||
}
|
||||
|
||||
func TestScanPromptWithSession(t *testing.T) {
|
||||
db, _, cleanup := testDB(t)
|
||||
defer cleanup()
|
||||
createBaseTables(t, db)
|
||||
seedSession(t, db, "claude-123", "sdk-123", "test-project")
|
||||
|
||||
// Insert a test prompt
|
||||
_, err := db.Exec(`
|
||||
INSERT INTO user_prompts (claude_session_id, prompt_number, prompt_text, matched_observations, created_at, created_at_epoch)
|
||||
VALUES ('claude-123', 1, 'test prompt', 5, '2025-01-01T00:00:00Z', 1704067200000)
|
||||
`)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Query with session join
|
||||
row := db.QueryRow(`
|
||||
SELECT p.id, p.claude_session_id, p.prompt_number, p.prompt_text, p.matched_observations, p.created_at, p.created_at_epoch, s.project, s.sdk_session_id
|
||||
FROM user_prompts p
|
||||
JOIN sdk_sessions s ON p.claude_session_id = s.claude_session_id
|
||||
WHERE p.claude_session_id = ?
|
||||
`, "claude-123")
|
||||
|
||||
prompt, err := scanPromptWithSession(row)
|
||||
require.NoError(t, err)
|
||||
assert.NotNil(t, prompt)
|
||||
assert.Equal(t, "claude-123", prompt.ClaudeSessionID)
|
||||
assert.Equal(t, 1, prompt.PromptNumber)
|
||||
assert.Equal(t, "test prompt", prompt.PromptText)
|
||||
assert.Equal(t, 5, prompt.MatchedObservations)
|
||||
assert.Equal(t, "test-project", prompt.Project)
|
||||
assert.Equal(t, "sdk-123", prompt.SDKSessionID)
|
||||
}
|
||||
|
||||
func TestScanPromptWithSessionRows(t *testing.T) {
|
||||
db, _, cleanup := testDB(t)
|
||||
defer cleanup()
|
||||
createBaseTables(t, db)
|
||||
seedSession(t, db, "claude-123", "sdk-123", "test-project")
|
||||
|
||||
// Insert multiple prompts
|
||||
_, err := db.Exec(`
|
||||
INSERT INTO user_prompts (claude_session_id, prompt_number, prompt_text, matched_observations, created_at, created_at_epoch)
|
||||
VALUES
|
||||
('claude-123', 1, 'prompt one', 3, '2025-01-01T00:00:00Z', 1704067200000),
|
||||
('claude-123', 2, 'prompt two', 5, '2025-01-02T00:00:00Z', 1704153600000)
|
||||
`)
|
||||
require.NoError(t, err)
|
||||
|
||||
rows, err := db.Query(`
|
||||
SELECT p.id, p.claude_session_id, p.prompt_number, p.prompt_text, p.matched_observations, p.created_at, p.created_at_epoch, s.project, s.sdk_session_id
|
||||
FROM user_prompts p
|
||||
JOIN sdk_sessions s ON p.claude_session_id = s.claude_session_id
|
||||
WHERE p.claude_session_id = ? ORDER BY p.id
|
||||
`, "claude-123")
|
||||
require.NoError(t, err)
|
||||
defer rows.Close()
|
||||
|
||||
prompts, err := scanPromptWithSessionRows(rows)
|
||||
require.NoError(t, err)
|
||||
assert.Len(t, prompts, 2)
|
||||
assert.Equal(t, "prompt one", prompts[0].PromptText)
|
||||
assert.Equal(t, "prompt two", prompts[1].PromptText)
|
||||
}
|
||||
|
||||
func TestParseLimitParam_HTTPRequest(t *testing.T) {
|
||||
// Test with an actual HTTP request
|
||||
handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
limit := ParseLimitParam(r, 25)
|
||||
if limit != 50 {
|
||||
t.Errorf("Expected limit 50, got %d", limit)
|
||||
}
|
||||
})
|
||||
|
||||
req := httptest.NewRequest("GET", "http://example.com/api?limit=50", nil)
|
||||
w := httptest.NewRecorder()
|
||||
handler.ServeHTTP(w, req)
|
||||
}
|
||||
@@ -0,0 +1,376 @@
|
||||
// Package sqlite provides SQLite database operations for claude-mnemonic.
|
||||
package sqlite
|
||||
|
||||
import (
|
||||
"database/sql"
|
||||
"fmt"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Migration represents a database schema migration.
|
||||
type Migration struct {
|
||||
Version int
|
||||
Name string
|
||||
SQL string
|
||||
}
|
||||
|
||||
// Migrations is the list of all database migrations in order.
|
||||
var Migrations = []Migration{
|
||||
{
|
||||
Version: 4,
|
||||
Name: "sdk_agent_architecture",
|
||||
SQL: `
|
||||
-- SDK Sessions (main session tracking)
|
||||
CREATE TABLE IF NOT EXISTS sdk_sessions (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
claude_session_id TEXT UNIQUE NOT NULL,
|
||||
sdk_session_id TEXT UNIQUE,
|
||||
project TEXT NOT NULL,
|
||||
user_prompt TEXT,
|
||||
started_at TEXT NOT NULL,
|
||||
started_at_epoch INTEGER NOT NULL,
|
||||
completed_at TEXT,
|
||||
completed_at_epoch INTEGER,
|
||||
status TEXT CHECK(status IN ('active', 'completed', 'failed')) NOT NULL DEFAULT 'active'
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_sdk_sessions_claude_id ON sdk_sessions(claude_session_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_sdk_sessions_sdk_id ON sdk_sessions(sdk_session_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_sdk_sessions_project ON sdk_sessions(project);
|
||||
CREATE INDEX IF NOT EXISTS idx_sdk_sessions_status ON sdk_sessions(status);
|
||||
CREATE INDEX IF NOT EXISTS idx_sdk_sessions_started ON sdk_sessions(started_at_epoch DESC);
|
||||
|
||||
-- Observations (extracted learnings)
|
||||
CREATE TABLE IF NOT EXISTS observations (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
sdk_session_id TEXT NOT NULL,
|
||||
project TEXT NOT NULL,
|
||||
text TEXT,
|
||||
type TEXT NOT NULL CHECK(type IN ('decision', 'bugfix', 'feature', 'refactor', 'discovery', 'change')),
|
||||
created_at TEXT NOT NULL,
|
||||
created_at_epoch INTEGER NOT NULL,
|
||||
FOREIGN KEY(sdk_session_id) REFERENCES sdk_sessions(sdk_session_id) ON DELETE CASCADE
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_observations_sdk_session ON observations(sdk_session_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_observations_project ON observations(project);
|
||||
CREATE INDEX IF NOT EXISTS idx_observations_type ON observations(type);
|
||||
CREATE INDEX IF NOT EXISTS idx_observations_created ON observations(created_at_epoch DESC);
|
||||
|
||||
-- Session Summaries
|
||||
CREATE TABLE IF NOT EXISTS session_summaries (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
sdk_session_id TEXT NOT NULL,
|
||||
project TEXT NOT NULL,
|
||||
request TEXT,
|
||||
investigated TEXT,
|
||||
learned TEXT,
|
||||
completed TEXT,
|
||||
next_steps TEXT,
|
||||
files_read TEXT,
|
||||
files_edited TEXT,
|
||||
notes TEXT,
|
||||
created_at TEXT NOT NULL,
|
||||
created_at_epoch INTEGER NOT NULL,
|
||||
FOREIGN KEY(sdk_session_id) REFERENCES sdk_sessions(sdk_session_id) ON DELETE CASCADE
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_session_summaries_sdk_session ON session_summaries(sdk_session_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_session_summaries_project ON session_summaries(project);
|
||||
CREATE INDEX IF NOT EXISTS idx_session_summaries_created ON session_summaries(created_at_epoch DESC);
|
||||
`,
|
||||
},
|
||||
{
|
||||
Version: 5,
|
||||
Name: "worker_port_column",
|
||||
SQL: `ALTER TABLE sdk_sessions ADD COLUMN worker_port INTEGER;`,
|
||||
},
|
||||
{
|
||||
Version: 6,
|
||||
Name: "prompt_tracking_columns",
|
||||
SQL: `
|
||||
ALTER TABLE sdk_sessions ADD COLUMN prompt_counter INTEGER DEFAULT 0;
|
||||
ALTER TABLE observations ADD COLUMN prompt_number INTEGER;
|
||||
ALTER TABLE session_summaries ADD COLUMN prompt_number INTEGER;
|
||||
`,
|
||||
},
|
||||
{
|
||||
Version: 8,
|
||||
Name: "observation_hierarchical_fields",
|
||||
SQL: `
|
||||
ALTER TABLE observations ADD COLUMN title TEXT;
|
||||
ALTER TABLE observations ADD COLUMN subtitle TEXT;
|
||||
ALTER TABLE observations ADD COLUMN facts TEXT;
|
||||
ALTER TABLE observations ADD COLUMN narrative TEXT;
|
||||
ALTER TABLE observations ADD COLUMN concepts TEXT;
|
||||
ALTER TABLE observations ADD COLUMN files_read TEXT;
|
||||
ALTER TABLE observations ADD COLUMN files_modified TEXT;
|
||||
`,
|
||||
},
|
||||
{
|
||||
Version: 10,
|
||||
Name: "user_prompts_table",
|
||||
SQL: `
|
||||
-- User prompts table
|
||||
CREATE TABLE IF NOT EXISTS user_prompts (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
claude_session_id TEXT NOT NULL,
|
||||
prompt_number INTEGER NOT NULL,
|
||||
prompt_text TEXT NOT NULL,
|
||||
created_at TEXT NOT NULL,
|
||||
created_at_epoch INTEGER NOT NULL,
|
||||
FOREIGN KEY(claude_session_id) REFERENCES sdk_sessions(claude_session_id) ON DELETE CASCADE
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_user_prompts_claude_session ON user_prompts(claude_session_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_user_prompts_created ON user_prompts(created_at_epoch DESC);
|
||||
CREATE INDEX IF NOT EXISTS idx_user_prompts_prompt_number ON user_prompts(prompt_number);
|
||||
CREATE INDEX IF NOT EXISTS idx_user_prompts_lookup ON user_prompts(claude_session_id, prompt_number);
|
||||
|
||||
-- FTS5 virtual table for user prompts
|
||||
CREATE VIRTUAL TABLE IF NOT EXISTS user_prompts_fts USING fts5(
|
||||
prompt_text,
|
||||
content='user_prompts',
|
||||
content_rowid='id'
|
||||
);
|
||||
|
||||
-- Triggers for FTS5 sync
|
||||
CREATE TRIGGER IF NOT EXISTS user_prompts_ai AFTER INSERT ON user_prompts BEGIN
|
||||
INSERT INTO user_prompts_fts(rowid, prompt_text)
|
||||
VALUES (new.id, new.prompt_text);
|
||||
END;
|
||||
|
||||
CREATE TRIGGER IF NOT EXISTS user_prompts_ad AFTER DELETE ON user_prompts BEGIN
|
||||
INSERT INTO user_prompts_fts(user_prompts_fts, rowid, prompt_text)
|
||||
VALUES('delete', old.id, old.prompt_text);
|
||||
END;
|
||||
|
||||
CREATE TRIGGER IF NOT EXISTS user_prompts_au AFTER UPDATE ON user_prompts BEGIN
|
||||
INSERT INTO user_prompts_fts(user_prompts_fts, rowid, prompt_text)
|
||||
VALUES('delete', old.id, old.prompt_text);
|
||||
INSERT INTO user_prompts_fts(rowid, prompt_text)
|
||||
VALUES (new.id, new.prompt_text);
|
||||
END;
|
||||
`,
|
||||
},
|
||||
{
|
||||
Version: 11,
|
||||
Name: "discovery_tokens_column",
|
||||
SQL: `
|
||||
ALTER TABLE observations ADD COLUMN discovery_tokens INTEGER DEFAULT 0;
|
||||
ALTER TABLE session_summaries ADD COLUMN discovery_tokens INTEGER DEFAULT 0;
|
||||
`,
|
||||
},
|
||||
{
|
||||
Version: 12,
|
||||
Name: "observations_fts",
|
||||
SQL: `
|
||||
-- FTS5 virtual table for observations
|
||||
CREATE VIRTUAL TABLE IF NOT EXISTS observations_fts USING fts5(
|
||||
title, subtitle, narrative,
|
||||
content='observations',
|
||||
content_rowid='id'
|
||||
);
|
||||
|
||||
-- Triggers for FTS5 sync
|
||||
CREATE TRIGGER IF NOT EXISTS observations_ai AFTER INSERT ON observations BEGIN
|
||||
INSERT INTO observations_fts(rowid, title, subtitle, narrative)
|
||||
VALUES (new.id, new.title, new.subtitle, new.narrative);
|
||||
END;
|
||||
|
||||
CREATE TRIGGER IF NOT EXISTS observations_ad AFTER DELETE ON observations BEGIN
|
||||
INSERT INTO observations_fts(observations_fts, rowid, title, subtitle, narrative)
|
||||
VALUES('delete', old.id, old.title, old.subtitle, old.narrative);
|
||||
END;
|
||||
|
||||
CREATE TRIGGER IF NOT EXISTS observations_au AFTER UPDATE ON observations BEGIN
|
||||
INSERT INTO observations_fts(observations_fts, rowid, title, subtitle, narrative)
|
||||
VALUES('delete', old.id, old.title, old.subtitle, old.narrative);
|
||||
INSERT INTO observations_fts(rowid, title, subtitle, narrative)
|
||||
VALUES (new.id, new.title, new.subtitle, new.narrative);
|
||||
END;
|
||||
`,
|
||||
},
|
||||
{
|
||||
Version: 13,
|
||||
Name: "session_summaries_fts",
|
||||
SQL: `
|
||||
-- FTS5 virtual table for session summaries
|
||||
CREATE VIRTUAL TABLE IF NOT EXISTS session_summaries_fts USING fts5(
|
||||
request, investigated, learned, completed, next_steps, notes,
|
||||
content='session_summaries',
|
||||
content_rowid='id'
|
||||
);
|
||||
|
||||
-- Triggers for FTS5 sync
|
||||
CREATE TRIGGER IF NOT EXISTS session_summaries_ai AFTER INSERT ON session_summaries BEGIN
|
||||
INSERT INTO session_summaries_fts(rowid, request, investigated, learned, completed, next_steps, notes)
|
||||
VALUES (new.id, new.request, new.investigated, new.learned, new.completed, new.next_steps, new.notes);
|
||||
END;
|
||||
|
||||
CREATE TRIGGER IF NOT EXISTS session_summaries_ad AFTER DELETE ON session_summaries BEGIN
|
||||
INSERT INTO session_summaries_fts(session_summaries_fts, rowid, request, investigated, learned, completed, next_steps, notes)
|
||||
VALUES('delete', old.id, old.request, old.investigated, old.learned, old.completed, old.next_steps, old.notes);
|
||||
END;
|
||||
|
||||
CREATE TRIGGER IF NOT EXISTS session_summaries_au AFTER UPDATE ON session_summaries BEGIN
|
||||
INSERT INTO session_summaries_fts(session_summaries_fts, rowid, request, investigated, learned, completed, next_steps, notes)
|
||||
VALUES('delete', old.id, old.request, old.investigated, old.learned, old.completed, old.next_steps, old.notes);
|
||||
INSERT INTO session_summaries_fts(rowid, request, investigated, learned, completed, next_steps, notes)
|
||||
VALUES (new.id, new.request, new.investigated, new.learned, new.completed, new.next_steps, new.notes);
|
||||
END;
|
||||
`,
|
||||
},
|
||||
{
|
||||
Version: 14,
|
||||
Name: "observation_scope_column",
|
||||
SQL: `
|
||||
-- Add scope column for project isolation
|
||||
-- 'project' = only visible within same project (default)
|
||||
-- 'global' = visible across all projects (best practices, patterns)
|
||||
ALTER TABLE observations ADD COLUMN scope TEXT DEFAULT 'project' CHECK(scope IN ('project', 'global'));
|
||||
|
||||
-- Index for efficient scope-based queries
|
||||
CREATE INDEX IF NOT EXISTS idx_observations_scope ON observations(scope);
|
||||
CREATE INDEX IF NOT EXISTS idx_observations_project_scope ON observations(project, scope);
|
||||
`,
|
||||
},
|
||||
{
|
||||
Version: 15,
|
||||
Name: "observation_file_mtimes",
|
||||
SQL: `
|
||||
-- Store file modification times at observation creation
|
||||
-- JSON object: {"path": mtime_epoch_ms, ...}
|
||||
-- Used to detect staleness when files change
|
||||
ALTER TABLE observations ADD COLUMN file_mtimes TEXT;
|
||||
`,
|
||||
},
|
||||
{
|
||||
Version: 16,
|
||||
Name: "prompt_matched_observations",
|
||||
SQL: `
|
||||
-- Track how many observations were found relevant for each prompt
|
||||
-- Displayed in dashboard timeline
|
||||
ALTER TABLE user_prompts ADD COLUMN matched_observations INTEGER DEFAULT 0;
|
||||
`,
|
||||
},
|
||||
{
|
||||
Version: 17,
|
||||
Name: "sqlite_vec_vectors",
|
||||
SQL: `
|
||||
-- Vector embeddings table using sqlite-vec
|
||||
-- Each document (narrative, fact, summary field, prompt) gets one vector
|
||||
-- Uses all-MiniLM-L6-v2 embeddings (384 dimensions)
|
||||
CREATE VIRTUAL TABLE IF NOT EXISTS vectors USING vec0(
|
||||
doc_id TEXT PRIMARY KEY,
|
||||
embedding float[384],
|
||||
sqlite_id INTEGER,
|
||||
doc_type TEXT,
|
||||
field_type TEXT,
|
||||
project TEXT,
|
||||
scope TEXT
|
||||
);
|
||||
`,
|
||||
},
|
||||
{
|
||||
Version: 18,
|
||||
Name: "user_prompts_unique_constraint",
|
||||
SQL: `
|
||||
-- Add unique constraint to prevent duplicate prompts
|
||||
-- This fixes a bug where the user-prompt hook could fire multiple times
|
||||
-- creating duplicate prompt records with incrementing numbers
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS idx_user_prompts_session_number_unique
|
||||
ON user_prompts(claude_session_id, prompt_number);
|
||||
`,
|
||||
},
|
||||
}
|
||||
|
||||
// MigrationManager handles database schema migrations.
|
||||
type MigrationManager struct {
|
||||
db *sql.DB
|
||||
}
|
||||
|
||||
// NewMigrationManager creates a new migration manager.
|
||||
func NewMigrationManager(db *sql.DB) *MigrationManager {
|
||||
return &MigrationManager{db: db}
|
||||
}
|
||||
|
||||
// EnsureSchemaVersionsTable creates the schema_versions table if it doesn't exist.
|
||||
func (m *MigrationManager) EnsureSchemaVersionsTable() error {
|
||||
_, err := m.db.Exec(`
|
||||
CREATE TABLE IF NOT EXISTS schema_versions (
|
||||
id INTEGER PRIMARY KEY,
|
||||
version INTEGER UNIQUE NOT NULL,
|
||||
applied_at TEXT NOT NULL
|
||||
)
|
||||
`)
|
||||
return err
|
||||
}
|
||||
|
||||
// GetAppliedVersions returns all applied migration versions.
|
||||
func (m *MigrationManager) GetAppliedVersions() (map[int]bool, error) {
|
||||
rows, err := m.db.Query("SELECT version FROM schema_versions ORDER BY version")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
versions := make(map[int]bool)
|
||||
for rows.Next() {
|
||||
var version int
|
||||
if err := rows.Scan(&version); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
versions[version] = true
|
||||
}
|
||||
return versions, rows.Err()
|
||||
}
|
||||
|
||||
// ApplyMigration applies a single migration.
|
||||
func (m *MigrationManager) ApplyMigration(migration Migration) error {
|
||||
tx, err := m.db.Begin()
|
||||
if err != nil {
|
||||
return fmt.Errorf("begin transaction: %w", err)
|
||||
}
|
||||
defer tx.Rollback()
|
||||
|
||||
// Execute migration SQL
|
||||
if _, err := tx.Exec(migration.SQL); err != nil {
|
||||
return fmt.Errorf("execute migration %d (%s): %w", migration.Version, migration.Name, err)
|
||||
}
|
||||
|
||||
// Record migration
|
||||
_, err = tx.Exec(
|
||||
"INSERT INTO schema_versions (version, applied_at) VALUES (?, ?)",
|
||||
migration.Version, time.Now().Format(time.RFC3339),
|
||||
)
|
||||
if err != nil {
|
||||
return fmt.Errorf("record migration %d: %w", migration.Version, err)
|
||||
}
|
||||
|
||||
return tx.Commit()
|
||||
}
|
||||
|
||||
// RunMigrations applies all pending migrations.
|
||||
func (m *MigrationManager) RunMigrations() error {
|
||||
if err := m.EnsureSchemaVersionsTable(); err != nil {
|
||||
return fmt.Errorf("ensure schema_versions table: %w", err)
|
||||
}
|
||||
|
||||
applied, err := m.GetAppliedVersions()
|
||||
if err != nil {
|
||||
return fmt.Errorf("get applied versions: %w", err)
|
||||
}
|
||||
|
||||
for _, migration := range Migrations {
|
||||
if applied[migration.Version] {
|
||||
continue
|
||||
}
|
||||
|
||||
if err := m.ApplyMigration(migration); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,196 @@
|
||||
package sqlite
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestNewMigrationManager(t *testing.T) {
|
||||
db, _, cleanup := testDB(t)
|
||||
defer cleanup()
|
||||
|
||||
manager := NewMigrationManager(db)
|
||||
require.NotNil(t, manager)
|
||||
assert.Equal(t, db, manager.db)
|
||||
}
|
||||
|
||||
func TestMigrationManager_EnsureSchemaVersionsTable(t *testing.T) {
|
||||
db, _, cleanup := testDB(t)
|
||||
defer cleanup()
|
||||
|
||||
manager := NewMigrationManager(db)
|
||||
|
||||
// Should create table without error
|
||||
err := manager.EnsureSchemaVersionsTable()
|
||||
require.NoError(t, err)
|
||||
|
||||
// Table should exist
|
||||
var count int
|
||||
err = db.QueryRow("SELECT COUNT(*) FROM schema_versions").Scan(&count)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 0, count) // Empty table
|
||||
|
||||
// Calling again should not error (IF NOT EXISTS)
|
||||
err = manager.EnsureSchemaVersionsTable()
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
func TestMigrationManager_GetAppliedVersions_Empty(t *testing.T) {
|
||||
db, _, cleanup := testDB(t)
|
||||
defer cleanup()
|
||||
|
||||
manager := NewMigrationManager(db)
|
||||
err := manager.EnsureSchemaVersionsTable()
|
||||
require.NoError(t, err)
|
||||
|
||||
versions, err := manager.GetAppliedVersions()
|
||||
require.NoError(t, err)
|
||||
assert.Empty(t, versions)
|
||||
}
|
||||
|
||||
func TestMigrationManager_GetAppliedVersions_WithVersions(t *testing.T) {
|
||||
db, _, cleanup := testDB(t)
|
||||
defer cleanup()
|
||||
|
||||
manager := NewMigrationManager(db)
|
||||
err := manager.EnsureSchemaVersionsTable()
|
||||
require.NoError(t, err)
|
||||
|
||||
// Insert some versions
|
||||
_, err = db.Exec("INSERT INTO schema_versions (version, applied_at) VALUES (1, '2025-01-01T00:00:00Z')")
|
||||
require.NoError(t, err)
|
||||
_, err = db.Exec("INSERT INTO schema_versions (version, applied_at) VALUES (2, '2025-01-02T00:00:00Z')")
|
||||
require.NoError(t, err)
|
||||
|
||||
versions, err := manager.GetAppliedVersions()
|
||||
require.NoError(t, err)
|
||||
assert.Len(t, versions, 2)
|
||||
assert.True(t, versions[1])
|
||||
assert.True(t, versions[2])
|
||||
assert.False(t, versions[3]) // Not applied
|
||||
}
|
||||
|
||||
func TestMigrationManager_ApplyMigration(t *testing.T) {
|
||||
db, _, cleanup := testDB(t)
|
||||
defer cleanup()
|
||||
|
||||
manager := NewMigrationManager(db)
|
||||
err := manager.EnsureSchemaVersionsTable()
|
||||
require.NoError(t, err)
|
||||
|
||||
// Apply a simple migration
|
||||
migration := Migration{
|
||||
Version: 100,
|
||||
Name: "test_migration",
|
||||
SQL: "CREATE TABLE test_table (id INTEGER PRIMARY KEY, name TEXT)",
|
||||
}
|
||||
|
||||
err = manager.ApplyMigration(migration)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Verify table was created
|
||||
var count int
|
||||
err = db.QueryRow("SELECT COUNT(*) FROM sqlite_master WHERE type='table' AND name='test_table'").Scan(&count)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 1, count)
|
||||
|
||||
// Verify migration was recorded
|
||||
var version int
|
||||
err = db.QueryRow("SELECT version FROM schema_versions WHERE version = 100").Scan(&version)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 100, version)
|
||||
}
|
||||
|
||||
func TestMigrationManager_ApplyMigration_InvalidSQL(t *testing.T) {
|
||||
db, _, cleanup := testDB(t)
|
||||
defer cleanup()
|
||||
|
||||
manager := NewMigrationManager(db)
|
||||
err := manager.EnsureSchemaVersionsTable()
|
||||
require.NoError(t, err)
|
||||
|
||||
// Try to apply invalid migration
|
||||
migration := Migration{
|
||||
Version: 100,
|
||||
Name: "invalid_migration",
|
||||
SQL: "INVALID SQL SYNTAX",
|
||||
}
|
||||
|
||||
err = manager.ApplyMigration(migration)
|
||||
assert.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "execute migration 100")
|
||||
}
|
||||
|
||||
func TestMigrationManager_RunMigrations_SingleMigration(t *testing.T) {
|
||||
db, _, cleanup := testDB(t)
|
||||
defer cleanup()
|
||||
|
||||
// Create a test migration manager with a subset of migrations
|
||||
manager := NewMigrationManager(db)
|
||||
|
||||
// First ensure schema versions table exists
|
||||
err := manager.EnsureSchemaVersionsTable()
|
||||
require.NoError(t, err)
|
||||
|
||||
// Apply first migration manually
|
||||
err = manager.ApplyMigration(Migrations[0])
|
||||
require.NoError(t, err)
|
||||
|
||||
// Verify the first migration version was recorded
|
||||
versions, err := manager.GetAppliedVersions()
|
||||
require.NoError(t, err)
|
||||
assert.True(t, versions[Migrations[0].Version])
|
||||
}
|
||||
|
||||
func TestMigrationManager_RunMigrations_SkipsApplied(t *testing.T) {
|
||||
db, _, cleanup := testDB(t)
|
||||
defer cleanup()
|
||||
|
||||
manager := NewMigrationManager(db)
|
||||
err := manager.EnsureSchemaVersionsTable()
|
||||
require.NoError(t, err)
|
||||
|
||||
// Mark some migrations as already applied
|
||||
_, err = db.Exec("INSERT INTO schema_versions (version, applied_at) VALUES (4, '2025-01-01T00:00:00Z')")
|
||||
require.NoError(t, err)
|
||||
|
||||
// Get applied versions
|
||||
versions, err := manager.GetAppliedVersions()
|
||||
require.NoError(t, err)
|
||||
assert.True(t, versions[4])
|
||||
}
|
||||
|
||||
func TestMigration_Struct(t *testing.T) {
|
||||
m := Migration{
|
||||
Version: 1,
|
||||
Name: "test",
|
||||
SQL: "SELECT 1",
|
||||
}
|
||||
|
||||
assert.Equal(t, 1, m.Version)
|
||||
assert.Equal(t, "test", m.Name)
|
||||
assert.Equal(t, "SELECT 1", m.SQL)
|
||||
}
|
||||
|
||||
func TestMigrations_List(t *testing.T) {
|
||||
// Verify migrations are ordered correctly
|
||||
assert.NotEmpty(t, Migrations)
|
||||
|
||||
// Verify all migrations have required fields
|
||||
for i, m := range Migrations {
|
||||
assert.Greater(t, m.Version, 0, "Migration %d has invalid version", i)
|
||||
assert.NotEmpty(t, m.Name, "Migration %d has empty name", i)
|
||||
assert.NotEmpty(t, m.SQL, "Migration %d has empty SQL", i)
|
||||
}
|
||||
|
||||
// Verify key migrations exist
|
||||
versionSet := make(map[int]bool)
|
||||
for _, m := range Migrations {
|
||||
versionSet[m.Version] = true
|
||||
}
|
||||
|
||||
assert.True(t, versionSet[4], "Should have sdk_agent_architecture migration")
|
||||
assert.True(t, versionSet[17], "Should have sqlite_vec_vectors migration")
|
||||
}
|
||||
@@ -0,0 +1,468 @@
|
||||
// Package sqlite provides SQLite database operations for claude-mnemonic.
|
||||
package sqlite
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/lukaszraczylo/claude-mnemonic/pkg/models"
|
||||
)
|
||||
|
||||
// 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.
|
||||
type ObservationStore struct {
|
||||
store *Store
|
||||
cleanupFunc CleanupFunc
|
||||
}
|
||||
|
||||
// NewObservationStore creates a new observation store.
|
||||
func NewObservationStore(store *Store) *ObservationStore {
|
||||
return &ObservationStore{store: store}
|
||||
}
|
||||
|
||||
// SetCleanupFunc sets the callback for when observations are deleted during cleanup.
|
||||
func (s *ObservationStore) SetCleanupFunc(fn CleanupFunc) {
|
||||
s.cleanupFunc = fn
|
||||
}
|
||||
|
||||
// StoreObservation stores a new observation.
|
||||
func (s *ObservationStore) StoreObservation(ctx context.Context, sdkSessionID, project string, obs *models.ParsedObservation, promptNumber int, discoveryTokens int64) (int64, int64, error) {
|
||||
now := time.Now()
|
||||
nowEpoch := now.UnixMilli()
|
||||
|
||||
// Ensure session exists (auto-create if missing)
|
||||
if err := s.ensureSessionExists(ctx, sdkSessionID, project); err != nil {
|
||||
return 0, 0, err
|
||||
}
|
||||
|
||||
// Determine scope: use parsed scope if set, otherwise auto-determine from concepts
|
||||
scope := obs.Scope
|
||||
if scope == "" {
|
||||
scope = models.DetermineScope(obs.Concepts)
|
||||
}
|
||||
|
||||
factsJSON, _ := json.Marshal(obs.Facts)
|
||||
conceptsJSON, _ := json.Marshal(obs.Concepts)
|
||||
filesReadJSON, _ := json.Marshal(obs.FilesRead)
|
||||
filesModifiedJSON, _ := json.Marshal(obs.FilesModified)
|
||||
fileMtimesJSON, _ := json.Marshal(obs.FileMtimes)
|
||||
|
||||
const query = `
|
||||
INSERT INTO observations
|
||||
(sdk_session_id, project, scope, type, title, subtitle, facts, narrative, concepts,
|
||||
files_read, files_modified, file_mtimes, prompt_number, discovery_tokens, created_at, created_at_epoch)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
`
|
||||
|
||||
result, err := s.store.ExecContext(ctx, query,
|
||||
sdkSessionID, project, string(scope), string(obs.Type),
|
||||
nullString(obs.Title), nullString(obs.Subtitle),
|
||||
string(factsJSON), nullString(obs.Narrative), string(conceptsJSON),
|
||||
string(filesReadJSON), string(filesModifiedJSON), string(fileMtimesJSON),
|
||||
nullInt(promptNumber), discoveryTokens,
|
||||
now.Format(time.RFC3339), nowEpoch,
|
||||
)
|
||||
if err != nil {
|
||||
return 0, 0, err
|
||||
}
|
||||
|
||||
id, _ := result.LastInsertId()
|
||||
|
||||
// Cleanup 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)
|
||||
}
|
||||
|
||||
return id, nowEpoch, nil
|
||||
}
|
||||
|
||||
// ensureSessionExists creates a session if it doesn't exist.
|
||||
func (s *ObservationStore) ensureSessionExists(ctx context.Context, sdkSessionID, project string) error {
|
||||
return EnsureSessionExists(ctx, s.store, sdkSessionID, project)
|
||||
}
|
||||
|
||||
// GetObservationByID retrieves an observation by ID.
|
||||
func (s *ObservationStore) GetObservationByID(ctx context.Context, id int64) (*models.Observation, error) {
|
||||
const query = `
|
||||
SELECT id, sdk_session_id, project, COALESCE(scope, 'project') as scope, type, title, subtitle, facts, narrative,
|
||||
concepts, files_read, files_modified, file_mtimes, prompt_number, discovery_tokens,
|
||||
created_at, created_at_epoch
|
||||
FROM observations
|
||||
WHERE id = ?
|
||||
`
|
||||
|
||||
obs, err := scanObservation(s.store.QueryRowContext(ctx, query, id))
|
||||
if err == sql.ErrNoRows {
|
||||
return nil, nil
|
||||
}
|
||||
return obs, err
|
||||
}
|
||||
|
||||
// GetObservationsByIDs retrieves observations by a list of IDs.
|
||||
func (s *ObservationStore) GetObservationsByIDs(ctx context.Context, ids []int64, orderBy string, limit int) ([]*models.Observation, error) {
|
||||
if len(ids) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
// Build query with placeholders
|
||||
// #nosec G202 -- query uses parameterized placeholders, not user input
|
||||
query := `
|
||||
SELECT id, sdk_session_id, project, COALESCE(scope, 'project') as scope, type, title, subtitle, facts, narrative,
|
||||
concepts, files_read, files_modified, file_mtimes, prompt_number, discovery_tokens,
|
||||
created_at, created_at_epoch
|
||||
FROM observations
|
||||
WHERE id IN (?` + repeatPlaceholders(len(ids)-1) + `)
|
||||
ORDER BY created_at_epoch `
|
||||
|
||||
if orderBy == "date_asc" {
|
||||
query += "ASC"
|
||||
} else {
|
||||
query += "DESC"
|
||||
}
|
||||
|
||||
if limit > 0 {
|
||||
query += " LIMIT ?"
|
||||
}
|
||||
|
||||
// Convert []int64 to []interface{}
|
||||
args := int64SliceToInterface(ids)
|
||||
if limit > 0 {
|
||||
args = append(args, limit)
|
||||
}
|
||||
|
||||
rows, err := s.store.db.QueryContext(ctx, query, args...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
return scanObservationRows(rows)
|
||||
}
|
||||
|
||||
// GetRecentObservations retrieves recent observations for a project.
|
||||
// This includes project-scoped observations for the specified project AND global observations.
|
||||
func (s *ObservationStore) GetRecentObservations(ctx context.Context, project string, limit int) ([]*models.Observation, error) {
|
||||
const query = `
|
||||
SELECT id, sdk_session_id, project, COALESCE(scope, 'project') as scope, type, title, subtitle, facts, narrative,
|
||||
concepts, files_read, files_modified, file_mtimes, prompt_number, discovery_tokens,
|
||||
created_at, created_at_epoch
|
||||
FROM observations
|
||||
WHERE (project = ? AND (scope IS NULL OR scope = 'project'))
|
||||
OR scope = 'global'
|
||||
ORDER BY created_at_epoch DESC
|
||||
LIMIT ?
|
||||
`
|
||||
|
||||
rows, err := s.store.QueryContext(ctx, query, project, limit)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
return scanObservationRows(rows)
|
||||
}
|
||||
|
||||
// GetObservationsByProjectStrict retrieves observations strictly for a specific project.
|
||||
// Unlike GetRecentObservations, this does NOT include global observations from other projects.
|
||||
// Use this for dashboard filtering where the user expects to see only that project's data.
|
||||
func (s *ObservationStore) GetObservationsByProjectStrict(ctx context.Context, project string, limit int) ([]*models.Observation, error) {
|
||||
const query = `
|
||||
SELECT id, sdk_session_id, project, COALESCE(scope, 'project') as scope, type, title, subtitle, facts, narrative,
|
||||
concepts, files_read, files_modified, file_mtimes, prompt_number, discovery_tokens,
|
||||
created_at, created_at_epoch
|
||||
FROM observations
|
||||
WHERE project = ?
|
||||
ORDER BY created_at_epoch DESC
|
||||
LIMIT ?
|
||||
`
|
||||
|
||||
rows, err := s.store.QueryContext(ctx, query, project, limit)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
return scanObservationRows(rows)
|
||||
}
|
||||
|
||||
// GetObservationCount returns the count of observations for a project (including global).
|
||||
func (s *ObservationStore) GetObservationCount(ctx context.Context, project string) (int, error) {
|
||||
const query = `
|
||||
SELECT COUNT(*) FROM observations
|
||||
WHERE project = ? OR scope = 'global'
|
||||
`
|
||||
var count int
|
||||
err := s.store.QueryRowContext(ctx, query, project).Scan(&count)
|
||||
return count, err
|
||||
}
|
||||
|
||||
// GetAllRecentObservations retrieves recent observations across all projects.
|
||||
func (s *ObservationStore) GetAllRecentObservations(ctx context.Context, limit int) ([]*models.Observation, error) {
|
||||
const query = `
|
||||
SELECT id, sdk_session_id, project, COALESCE(scope, 'project') as scope, type, title, subtitle, facts, narrative,
|
||||
concepts, files_read, files_modified, file_mtimes, prompt_number, discovery_tokens,
|
||||
created_at, created_at_epoch
|
||||
FROM observations
|
||||
ORDER BY created_at_epoch DESC
|
||||
LIMIT ?
|
||||
`
|
||||
|
||||
rows, err := s.store.QueryContext(ctx, query, limit)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
return scanObservationRows(rows)
|
||||
}
|
||||
|
||||
// SearchObservationsFTS performs full-text search on observations.
|
||||
func (s *ObservationStore) SearchObservationsFTS(ctx context.Context, query, project string, limit int) ([]*models.Observation, error) {
|
||||
if limit <= 0 {
|
||||
limit = 10
|
||||
}
|
||||
|
||||
// Extract keywords from the query (words > 3 chars, not common)
|
||||
keywords := extractKeywords(query)
|
||||
if len(keywords) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
// Build FTS5 query: keyword1 OR keyword2 OR keyword3
|
||||
ftsTerms := strings.Join(keywords, " OR ")
|
||||
|
||||
// Use FTS5 to search title, subtitle, and narrative
|
||||
const ftsQuery = `
|
||||
SELECT o.id, o.sdk_session_id, o.project, COALESCE(o.scope, 'project') as scope, o.type,
|
||||
o.title, o.subtitle, o.facts, o.narrative, o.concepts, o.files_read, o.files_modified,
|
||||
o.file_mtimes, o.prompt_number, o.discovery_tokens, o.created_at, o.created_at_epoch
|
||||
FROM observations o
|
||||
JOIN observations_fts fts ON o.id = fts.rowid
|
||||
WHERE observations_fts MATCH ?
|
||||
AND (o.project = ? OR o.scope = 'global')
|
||||
ORDER BY rank
|
||||
LIMIT ?
|
||||
`
|
||||
|
||||
rows, err := s.store.QueryContext(ctx, ftsQuery, ftsTerms, project, limit)
|
||||
if err != nil {
|
||||
// FTS failed, try LIKE fallback
|
||||
return s.searchObservationsLike(ctx, keywords, project, limit)
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
observations, err := scanObservationRows(rows)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// If FTS returned nothing, try LIKE search
|
||||
if len(observations) == 0 {
|
||||
return s.searchObservationsLike(ctx, keywords, project, limit)
|
||||
}
|
||||
|
||||
return observations, nil
|
||||
}
|
||||
|
||||
// searchObservationsLike performs fallback LIKE search on observations.
|
||||
func (s *ObservationStore) searchObservationsLike(ctx context.Context, keywords []string, project string, limit int) ([]*models.Observation, error) {
|
||||
if len(keywords) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
// Build LIKE conditions for each keyword
|
||||
var conditions []string
|
||||
var args []interface{}
|
||||
|
||||
for _, kw := range keywords {
|
||||
pattern := "%" + kw + "%"
|
||||
conditions = append(conditions, "(title LIKE ? OR subtitle LIKE ? OR narrative LIKE ?)")
|
||||
args = append(args, pattern, pattern, pattern)
|
||||
}
|
||||
|
||||
// #nosec G202 -- query uses parameterized placeholders, not user input
|
||||
query := `
|
||||
SELECT id, sdk_session_id, project, COALESCE(scope, 'project') as scope, type,
|
||||
title, subtitle, facts, narrative, concepts, files_read, files_modified,
|
||||
file_mtimes, prompt_number, discovery_tokens, created_at, created_at_epoch
|
||||
FROM observations
|
||||
WHERE (` + strings.Join(conditions, " OR ") + `)
|
||||
AND (project = ? OR scope = 'global')
|
||||
ORDER BY created_at_epoch DESC
|
||||
LIMIT ?
|
||||
`
|
||||
args = append(args, project, limit)
|
||||
|
||||
rows, err := s.store.db.QueryContext(ctx, query, args...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
return scanObservationRows(rows)
|
||||
}
|
||||
|
||||
// extractKeywords extracts significant words from a query.
|
||||
func extractKeywords(query string) []string {
|
||||
// Common words to skip
|
||||
stopWords := map[string]bool{
|
||||
"what": true, "is": true, "the": true, "a": true, "an": true,
|
||||
"how": true, "does": true, "do": true, "can": true, "could": true,
|
||||
"would": true, "should": true, "where": true, "when": true, "why": true,
|
||||
"which": true, "who": true, "this": true, "that": true, "these": true,
|
||||
"those": true, "it": true, "its": true, "for": true, "from": true,
|
||||
"with": true, "about": true, "into": true, "through": true, "during": true,
|
||||
"before": true, "after": true, "above": true, "below": true, "to": true,
|
||||
"of": true, "in": true, "on": true, "at": true, "by": true, "and": true,
|
||||
"or": true, "but": true, "if": true, "then": true, "else": true,
|
||||
"function": true, "method": true, "class": true, "file": true,
|
||||
"code": true, "work": true, "works": true, "working": true,
|
||||
"please": true, "help": true, "me": true, "my": true, "i": true,
|
||||
"tell": true, "show": true, "explain": true, "describe": true,
|
||||
}
|
||||
|
||||
// Split and filter
|
||||
words := strings.FieldsFunc(strings.ToLower(query), func(r rune) bool {
|
||||
return !((r >= 'a' && r <= 'z') || (r >= 'A' && r <= 'Z') || (r >= '0' && r <= '9') || r == '_')
|
||||
})
|
||||
|
||||
var keywords []string
|
||||
seen := make(map[string]bool)
|
||||
|
||||
for _, word := range words {
|
||||
// Skip short words, stop words, and duplicates
|
||||
if len(word) < 4 || stopWords[word] || seen[word] {
|
||||
continue
|
||||
}
|
||||
seen[word] = true
|
||||
keywords = append(keywords, word)
|
||||
}
|
||||
|
||||
return keywords
|
||||
}
|
||||
|
||||
// DeleteObservations deletes multiple observations by ID.
|
||||
func (s *ObservationStore) DeleteObservations(ctx context.Context, ids []int64) (int64, error) {
|
||||
if len(ids) == 0 {
|
||||
return 0, nil
|
||||
}
|
||||
|
||||
query := `DELETE FROM observations WHERE id IN (?` + repeatPlaceholders(len(ids)-1) + `)` // #nosec G202 -- uses parameterized placeholders
|
||||
|
||||
args := make([]interface{}, len(ids))
|
||||
for i, id := range ids {
|
||||
args[i] = id
|
||||
}
|
||||
|
||||
result, err := s.store.db.ExecContext(ctx, query, args...)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return result.RowsAffected()
|
||||
}
|
||||
|
||||
// MaxObservationsPerProject is the hard limit of observations per project.
|
||||
const MaxObservationsPerProject = 100
|
||||
|
||||
// CleanupOldObservations deletes observations beyond the limit for a project.
|
||||
// Keeps the most recent MaxObservationsPerProject observations per project.
|
||||
// Returns the IDs of deleted observations for downstream cleanup (e.g., vector DB).
|
||||
func (s *ObservationStore) CleanupOldObservations(ctx context.Context, project string) ([]int64, error) {
|
||||
// First, find IDs that will be deleted
|
||||
const selectQuery = `
|
||||
SELECT id FROM observations
|
||||
WHERE project = ? AND id NOT IN (
|
||||
SELECT id FROM observations
|
||||
WHERE project = ?
|
||||
ORDER BY created_at_epoch DESC
|
||||
LIMIT ?
|
||||
)
|
||||
`
|
||||
|
||||
rows, err := s.store.QueryContext(ctx, selectQuery, project, project, MaxObservationsPerProject)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
var toDelete []int64
|
||||
for rows.Next() {
|
||||
var id int64
|
||||
if err := rows.Scan(&id); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
toDelete = append(toDelete, id)
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if len(toDelete) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
// Delete the observations
|
||||
const deleteQuery = `
|
||||
DELETE FROM observations
|
||||
WHERE project = ? AND id NOT IN (
|
||||
SELECT id FROM observations
|
||||
WHERE project = ?
|
||||
ORDER BY created_at_epoch DESC
|
||||
LIMIT ?
|
||||
)
|
||||
`
|
||||
|
||||
_, err = s.store.ExecContext(ctx, deleteQuery, project, project, MaxObservationsPerProject)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return toDelete, nil
|
||||
}
|
||||
|
||||
// Helper functions
|
||||
|
||||
// scanObservation scans a single observation from a row scanner.
|
||||
// This reduces code duplication across all observation query methods.
|
||||
func scanObservation(scanner interface{ Scan(...interface{}) error }) (*models.Observation, error) {
|
||||
var obs models.Observation
|
||||
if err := scanner.Scan(
|
||||
&obs.ID, &obs.SDKSessionID, &obs.Project, &obs.Scope, &obs.Type,
|
||||
&obs.Title, &obs.Subtitle, &obs.Facts, &obs.Narrative,
|
||||
&obs.Concepts, &obs.FilesRead, &obs.FilesModified, &obs.FileMtimes,
|
||||
&obs.PromptNumber, &obs.DiscoveryTokens,
|
||||
&obs.CreatedAt, &obs.CreatedAtEpoch,
|
||||
); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &obs, nil
|
||||
}
|
||||
|
||||
// scanObservationRows scans multiple observations from rows.
|
||||
// Caller must close rows after calling this function.
|
||||
func scanObservationRows(rows *sql.Rows) ([]*models.Observation, error) {
|
||||
var observations []*models.Observation
|
||||
for rows.Next() {
|
||||
obs, err := scanObservation(rows)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
observations = append(observations, obs)
|
||||
}
|
||||
return observations, rows.Err()
|
||||
}
|
||||
|
||||
// Note: nullString, nullInt, and repeatPlaceholders are in helpers.go
|
||||
@@ -0,0 +1,947 @@
|
||||
// Package sqlite provides SQLite database operations for claude-mnemonic.
|
||||
package sqlite
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/lukaszraczylo/claude-mnemonic/pkg/models"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"github.com/stretchr/testify/suite"
|
||||
)
|
||||
|
||||
// testObservationStoreBasic creates an ObservationStore with base tables (no FTS5).
|
||||
func testObservationStoreBasic(t *testing.T) (*ObservationStore, *Store, func()) {
|
||||
t.Helper()
|
||||
|
||||
db, _, cleanup := testDB(t)
|
||||
createBaseTables(t, db)
|
||||
|
||||
store := newStoreFromDB(db)
|
||||
obsStore := NewObservationStore(store)
|
||||
|
||||
return obsStore, store, cleanup
|
||||
}
|
||||
|
||||
// testObservationStore creates an ObservationStore with a test database including FTS5.
|
||||
func testObservationStore(t *testing.T) (*ObservationStore, *Store, func()) {
|
||||
t.Helper()
|
||||
|
||||
db, _, cleanup := testDB(t)
|
||||
createAllTables(t, db)
|
||||
|
||||
store := newStoreFromDB(db)
|
||||
obsStore := NewObservationStore(store)
|
||||
|
||||
return obsStore, store, cleanup
|
||||
}
|
||||
|
||||
// ObservationStoreSuite is a test suite for ObservationStore operations.
|
||||
type ObservationStoreSuite struct {
|
||||
suite.Suite
|
||||
obsStore *ObservationStore
|
||||
store *Store
|
||||
cleanup func()
|
||||
}
|
||||
|
||||
func (s *ObservationStoreSuite) SetupTest() {
|
||||
s.obsStore, s.store, s.cleanup = testObservationStoreBasic(s.T())
|
||||
}
|
||||
|
||||
func (s *ObservationStoreSuite) TearDownTest() {
|
||||
if s.cleanup != nil {
|
||||
s.cleanup()
|
||||
}
|
||||
}
|
||||
|
||||
func TestObservationStoreSuite(t *testing.T) {
|
||||
suite.Run(t, new(ObservationStoreSuite))
|
||||
}
|
||||
|
||||
// TestStoreObservation_TableDriven tests observation storage with various scenarios.
|
||||
func (s *ObservationStoreSuite) TestStoreObservation_TableDriven() {
|
||||
ctx := context.Background()
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
sdkSessionID string
|
||||
project string
|
||||
obs *models.ParsedObservation
|
||||
promptNum int
|
||||
tokens int64
|
||||
wantErr bool
|
||||
}{
|
||||
{
|
||||
name: "basic discovery observation",
|
||||
sdkSessionID: "session-basic",
|
||||
project: "project-a",
|
||||
obs: &models.ParsedObservation{
|
||||
Type: models.ObsTypeDiscovery,
|
||||
Title: "Test Title",
|
||||
Subtitle: "Test Subtitle",
|
||||
Narrative: "Test narrative content",
|
||||
Facts: []string{"Fact 1", "Fact 2"},
|
||||
Concepts: []string{"testing", "golang"},
|
||||
},
|
||||
promptNum: 1,
|
||||
tokens: 100,
|
||||
wantErr: false,
|
||||
},
|
||||
{
|
||||
name: "bugfix observation",
|
||||
sdkSessionID: "session-bugfix",
|
||||
project: "project-b",
|
||||
obs: &models.ParsedObservation{
|
||||
Type: models.ObsTypeBugfix,
|
||||
Title: "Fixed null pointer",
|
||||
Narrative: "Fixed null pointer exception in handler",
|
||||
FilesModified: []string{"handler.go"},
|
||||
},
|
||||
promptNum: 2,
|
||||
tokens: 50,
|
||||
wantErr: false,
|
||||
},
|
||||
{
|
||||
name: "global scope observation",
|
||||
sdkSessionID: "session-global",
|
||||
project: "project-c",
|
||||
obs: &models.ParsedObservation{
|
||||
Type: models.ObsTypeDiscovery,
|
||||
Title: "Security best practice",
|
||||
Narrative: "Always validate user input",
|
||||
Concepts: []string{"security", "best-practice"},
|
||||
},
|
||||
promptNum: 1,
|
||||
tokens: 75,
|
||||
wantErr: false,
|
||||
},
|
||||
{
|
||||
name: "observation with files",
|
||||
sdkSessionID: "session-files",
|
||||
project: "project-d",
|
||||
obs: &models.ParsedObservation{
|
||||
Type: models.ObsTypeFeature,
|
||||
Title: "Added authentication",
|
||||
Narrative: "Implemented JWT authentication",
|
||||
FilesRead: []string{"config.go", "auth.go"},
|
||||
FilesModified: []string{"handler.go", "middleware.go"},
|
||||
FileMtimes: map[string]int64{"handler.go": 1234567890, "middleware.go": 1234567891},
|
||||
},
|
||||
promptNum: 3,
|
||||
tokens: 200,
|
||||
wantErr: false,
|
||||
},
|
||||
{
|
||||
name: "minimal observation",
|
||||
sdkSessionID: "session-minimal",
|
||||
project: "project-e",
|
||||
obs: &models.ParsedObservation{
|
||||
Type: models.ObsTypeChange,
|
||||
},
|
||||
promptNum: 0,
|
||||
tokens: 0,
|
||||
wantErr: false,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
s.Run(tt.name, func() {
|
||||
id, epoch, err := s.obsStore.StoreObservation(ctx, tt.sdkSessionID, tt.project, tt.obs, tt.promptNum, tt.tokens)
|
||||
if tt.wantErr {
|
||||
s.Error(err)
|
||||
return
|
||||
}
|
||||
|
||||
s.NoError(err)
|
||||
s.Greater(id, int64(0))
|
||||
s.Greater(epoch, int64(0))
|
||||
|
||||
// Retrieve and verify
|
||||
retrieved, err := s.obsStore.GetObservationByID(ctx, id)
|
||||
s.NoError(err)
|
||||
s.NotNil(retrieved)
|
||||
s.Equal(id, retrieved.ID)
|
||||
s.Equal(tt.project, retrieved.Project)
|
||||
s.Equal(tt.obs.Type, retrieved.Type)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestGetObservationByID_NotFound tests retrieval of non-existent observation.
|
||||
func (s *ObservationStoreSuite) TestGetObservationByID_NotFound() {
|
||||
ctx := context.Background()
|
||||
|
||||
obs, err := s.obsStore.GetObservationByID(ctx, 99999)
|
||||
s.NoError(err)
|
||||
s.Nil(obs)
|
||||
}
|
||||
|
||||
// TestGetRecentObservations_TableDriven tests recent observations retrieval.
|
||||
func (s *ObservationStoreSuite) TestGetRecentObservations_TableDriven() {
|
||||
ctx := context.Background()
|
||||
|
||||
// Create 15 observations
|
||||
for i := 0; i < 15; i++ {
|
||||
obs := &models.ParsedObservation{
|
||||
Type: models.ObsTypeDiscovery,
|
||||
Title: "Observation " + string(rune('A'+i)),
|
||||
}
|
||||
_, _, err := s.obsStore.StoreObservation(ctx, "session-"+string(rune('0'+i)), "project-a", obs, i, 10)
|
||||
s.NoError(err)
|
||||
time.Sleep(time.Millisecond) // Ensure different timestamps
|
||||
}
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
project string
|
||||
limit int
|
||||
wantCount int
|
||||
}{
|
||||
{
|
||||
name: "limit 5",
|
||||
project: "project-a",
|
||||
limit: 5,
|
||||
wantCount: 5,
|
||||
},
|
||||
{
|
||||
name: "limit 10",
|
||||
project: "project-a",
|
||||
limit: 10,
|
||||
wantCount: 10,
|
||||
},
|
||||
{
|
||||
name: "limit higher than count",
|
||||
project: "project-a",
|
||||
limit: 50,
|
||||
wantCount: 15,
|
||||
},
|
||||
{
|
||||
name: "different project (no results)",
|
||||
project: "project-b",
|
||||
limit: 10,
|
||||
wantCount: 0,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
s.Run(tt.name, func() {
|
||||
observations, err := s.obsStore.GetRecentObservations(ctx, tt.project, tt.limit)
|
||||
s.NoError(err)
|
||||
s.Len(observations, tt.wantCount)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestDeleteObservations_TableDriven tests observation deletion.
|
||||
func (s *ObservationStoreSuite) TestDeleteObservations_TableDriven() {
|
||||
ctx := context.Background()
|
||||
|
||||
// Create observations
|
||||
var ids []int64
|
||||
for i := 0; i < 5; i++ {
|
||||
obs := &models.ParsedObservation{
|
||||
Type: models.ObsTypeDiscovery,
|
||||
Title: "To delete " + string(rune('A'+i)),
|
||||
}
|
||||
id, _, err := s.obsStore.StoreObservation(ctx, "session-del", "project-del", obs, i, 10)
|
||||
s.NoError(err)
|
||||
ids = append(ids, id)
|
||||
}
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
toDelete []int64
|
||||
wantDeleted int64
|
||||
wantRemain int
|
||||
}{
|
||||
{
|
||||
name: "delete none",
|
||||
toDelete: []int64{},
|
||||
wantDeleted: 0,
|
||||
wantRemain: 5,
|
||||
},
|
||||
{
|
||||
name: "delete one",
|
||||
toDelete: ids[0:1],
|
||||
wantDeleted: 1,
|
||||
wantRemain: 4,
|
||||
},
|
||||
{
|
||||
name: "delete remaining",
|
||||
toDelete: ids[1:],
|
||||
wantDeleted: 4,
|
||||
wantRemain: 0,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
s.Run(tt.name, func() {
|
||||
deleted, err := s.obsStore.DeleteObservations(ctx, tt.toDelete)
|
||||
s.NoError(err)
|
||||
s.Equal(tt.wantDeleted, deleted)
|
||||
|
||||
remaining, err := s.obsStore.GetAllRecentObservations(ctx, 100)
|
||||
s.NoError(err)
|
||||
s.Len(remaining, tt.wantRemain)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestGetObservationsByIDs tests retrieval by multiple IDs.
|
||||
func (s *ObservationStoreSuite) TestGetObservationsByIDs() {
|
||||
ctx := context.Background()
|
||||
|
||||
// Create observations
|
||||
var ids []int64
|
||||
for i := 0; i < 5; i++ {
|
||||
obs := &models.ParsedObservation{
|
||||
Type: models.ObsTypeDiscovery,
|
||||
Title: "By ID " + string(rune('A'+i)),
|
||||
}
|
||||
id, _, err := s.obsStore.StoreObservation(ctx, "session-byid", "project-byid", obs, i, 10)
|
||||
s.NoError(err)
|
||||
ids = append(ids, id)
|
||||
time.Sleep(time.Millisecond)
|
||||
}
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
queryIDs []int64
|
||||
orderBy string
|
||||
limit int
|
||||
wantCount int
|
||||
}{
|
||||
{
|
||||
name: "empty IDs",
|
||||
queryIDs: []int64{},
|
||||
orderBy: "date_desc",
|
||||
limit: 10,
|
||||
wantCount: 0,
|
||||
},
|
||||
{
|
||||
name: "single ID",
|
||||
queryIDs: ids[0:1],
|
||||
orderBy: "date_desc",
|
||||
limit: 10,
|
||||
wantCount: 1,
|
||||
},
|
||||
{
|
||||
name: "all IDs",
|
||||
queryIDs: ids,
|
||||
orderBy: "date_desc",
|
||||
limit: 10,
|
||||
wantCount: 5,
|
||||
},
|
||||
{
|
||||
name: "with limit less than IDs",
|
||||
queryIDs: ids,
|
||||
orderBy: "date_desc",
|
||||
limit: 3,
|
||||
wantCount: 3,
|
||||
},
|
||||
{
|
||||
name: "ascending order",
|
||||
queryIDs: ids,
|
||||
orderBy: "date_asc",
|
||||
limit: 10,
|
||||
wantCount: 5,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
s.Run(tt.name, func() {
|
||||
observations, err := s.obsStore.GetObservationsByIDs(ctx, tt.queryIDs, tt.orderBy, tt.limit)
|
||||
if tt.wantCount == 0 {
|
||||
s.NoError(err)
|
||||
s.Nil(observations)
|
||||
} else {
|
||||
s.NoError(err)
|
||||
s.Len(observations, tt.wantCount)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestGlobalScope tests global vs project scope.
|
||||
func (s *ObservationStoreSuite) TestGlobalScope() {
|
||||
ctx := context.Background()
|
||||
|
||||
// Create project-scoped observation
|
||||
projectObs := &models.ParsedObservation{
|
||||
Type: models.ObsTypeDiscovery,
|
||||
Title: "Project specific",
|
||||
Concepts: []string{"project-specific"},
|
||||
}
|
||||
_, _, err := s.obsStore.StoreObservation(ctx, "session-scope", "project-a", projectObs, 1, 10)
|
||||
s.NoError(err)
|
||||
|
||||
// Create global-scoped observation (security concept triggers global)
|
||||
globalObs := &models.ParsedObservation{
|
||||
Type: models.ObsTypeDiscovery,
|
||||
Title: "Global security",
|
||||
Concepts: []string{"security"},
|
||||
}
|
||||
_, _, err = s.obsStore.StoreObservation(ctx, "session-scope", "project-a", globalObs, 2, 10)
|
||||
s.NoError(err)
|
||||
|
||||
// Project-a should see both
|
||||
resultsA, err := s.obsStore.GetRecentObservations(ctx, "project-a", 10)
|
||||
s.NoError(err)
|
||||
s.Len(resultsA, 2)
|
||||
|
||||
// Project-b should only see global
|
||||
resultsB, err := s.obsStore.GetRecentObservations(ctx, "project-b", 10)
|
||||
s.NoError(err)
|
||||
s.Len(resultsB, 1)
|
||||
s.Equal("Global security", resultsB[0].Title.String)
|
||||
s.Equal(models.ScopeGlobal, resultsB[0].Scope)
|
||||
}
|
||||
|
||||
// TestSetCleanupFunc tests the cleanup function callback.
|
||||
func (s *ObservationStoreSuite) TestSetCleanupFunc() {
|
||||
ctx := context.Background()
|
||||
|
||||
var calledWith []int64
|
||||
s.obsStore.SetCleanupFunc(func(ctx context.Context, deletedIDs []int64) {
|
||||
calledWith = deletedIDs
|
||||
})
|
||||
|
||||
// Store an observation
|
||||
obs := &models.ParsedObservation{
|
||||
Type: models.ObsTypeDiscovery,
|
||||
Title: "Test cleanup",
|
||||
}
|
||||
_, _, err := s.obsStore.StoreObservation(ctx, "session-cleanup", "project-cleanup", obs, 1, 10)
|
||||
s.NoError(err)
|
||||
|
||||
// Cleanup should not have been called since nothing was deleted
|
||||
s.Empty(calledWith)
|
||||
}
|
||||
|
||||
// TestGetObservationCount tests observation counting.
|
||||
func (s *ObservationStoreSuite) TestGetObservationCount() {
|
||||
ctx := context.Background()
|
||||
|
||||
// Create observations for project-a
|
||||
for i := 0; i < 5; i++ {
|
||||
obs := &models.ParsedObservation{
|
||||
Type: models.ObsTypeDiscovery,
|
||||
}
|
||||
_, _, err := s.obsStore.StoreObservation(ctx, "session-count", "project-a", obs, i, 10)
|
||||
s.NoError(err)
|
||||
}
|
||||
|
||||
// Create global observation
|
||||
globalObs := &models.ParsedObservation{
|
||||
Type: models.ObsTypeDiscovery,
|
||||
Concepts: []string{"security"},
|
||||
}
|
||||
_, _, err := s.obsStore.StoreObservation(ctx, "session-count", "project-a", globalObs, 6, 10)
|
||||
s.NoError(err)
|
||||
|
||||
// Project-a should count 6 (5 project + 1 global)
|
||||
count, err := s.obsStore.GetObservationCount(ctx, "project-a")
|
||||
s.NoError(err)
|
||||
s.Equal(6, count)
|
||||
|
||||
// Project-b should count 1 (only global)
|
||||
count, err = s.obsStore.GetObservationCount(ctx, "project-b")
|
||||
s.NoError(err)
|
||||
s.Equal(1, count)
|
||||
}
|
||||
|
||||
func TestObservationStore_StoreAndRetrieve(t *testing.T) {
|
||||
obsStore, _, cleanup := testObservationStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
obs := &models.ParsedObservation{
|
||||
Type: models.ObsTypeDiscovery,
|
||||
Title: "Test Observation",
|
||||
Subtitle: "A subtitle",
|
||||
Narrative: "This is a test observation about testing",
|
||||
Facts: []string{"Fact 1", "Fact 2"},
|
||||
Concepts: []string{"testing", "golang"},
|
||||
FilesRead: []string{"test.go"},
|
||||
FilesModified: []string{},
|
||||
}
|
||||
|
||||
id, epoch, err := obsStore.StoreObservation(ctx, "session-1", "project-a", obs, 1, 100)
|
||||
require.NoError(t, err)
|
||||
assert.Greater(t, id, int64(0))
|
||||
assert.Greater(t, epoch, int64(0))
|
||||
|
||||
// Retrieve by ID
|
||||
retrieved, err := obsStore.GetObservationByID(ctx, id)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, retrieved)
|
||||
|
||||
assert.Equal(t, id, retrieved.ID)
|
||||
assert.Equal(t, "session-1", retrieved.SDKSessionID)
|
||||
assert.Equal(t, "project-a", retrieved.Project)
|
||||
assert.Equal(t, models.ObsTypeDiscovery, retrieved.Type)
|
||||
assert.Equal(t, "Test Observation", retrieved.Title.String)
|
||||
assert.Equal(t, "A subtitle", retrieved.Subtitle.String)
|
||||
assert.Equal(t, "This is a test observation about testing", retrieved.Narrative.String)
|
||||
assert.Equal(t, []string{"Fact 1", "Fact 2"}, []string(retrieved.Facts))
|
||||
assert.Equal(t, []string{"testing", "golang"}, []string(retrieved.Concepts))
|
||||
}
|
||||
|
||||
func TestObservationStore_GetRecentObservations(t *testing.T) {
|
||||
obsStore, _, cleanup := testObservationStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Create multiple observations
|
||||
for i := 0; i < 10; i++ {
|
||||
obs := &models.ParsedObservation{
|
||||
Type: models.ObsTypeDiscovery,
|
||||
Title: "Observation " + string(rune('A'+i)),
|
||||
Narrative: "Content " + string(rune('A'+i)),
|
||||
Concepts: []string{"test"},
|
||||
}
|
||||
_, _, err := obsStore.StoreObservation(ctx, "session-1", "project-a", obs, i+1, 100)
|
||||
require.NoError(t, err)
|
||||
time.Sleep(time.Millisecond) // Ensure different timestamps
|
||||
}
|
||||
|
||||
// Get recent with limit 5
|
||||
recent, err := obsStore.GetRecentObservations(ctx, "project-a", 5)
|
||||
require.NoError(t, err)
|
||||
assert.Len(t, recent, 5)
|
||||
|
||||
// Get recent with limit 20 (more than exists)
|
||||
recent, err = obsStore.GetRecentObservations(ctx, "project-a", 20)
|
||||
require.NoError(t, err)
|
||||
assert.Len(t, recent, 10)
|
||||
}
|
||||
|
||||
func TestObservationStore_SearchObservationsFTS(t *testing.T) {
|
||||
obsStore, _, cleanup := testObservationStore(t)
|
||||
defer cleanup()
|
||||
|
||||
// FTS5 tables are created by testObservationStore via testutil.CreateAllTables
|
||||
ctx := context.Background()
|
||||
|
||||
// Create observations with different content
|
||||
observations := []struct {
|
||||
title string
|
||||
narrative string
|
||||
}{
|
||||
{"Authentication implementation", "JWT based authentication flow"},
|
||||
{"Database setup", "PostgreSQL configuration and migrations"},
|
||||
{"Caching layer", "Redis caching implementation"},
|
||||
{"User authentication fix", "Fixed authentication bug in login"},
|
||||
{"API endpoints", "REST API implementation details"},
|
||||
}
|
||||
|
||||
for _, o := range observations {
|
||||
obs := &models.ParsedObservation{
|
||||
Type: models.ObsTypeDiscovery,
|
||||
Title: o.title,
|
||||
Narrative: o.narrative,
|
||||
}
|
||||
_, _, err := obsStore.StoreObservation(ctx, "session-1", "project-a", obs, 1, 100)
|
||||
require.NoError(t, err)
|
||||
time.Sleep(time.Millisecond)
|
||||
}
|
||||
|
||||
// Search for authentication - should find 2 observations
|
||||
results, err := obsStore.SearchObservationsFTS(ctx, "authentication", "project-a", 50)
|
||||
require.NoError(t, err)
|
||||
assert.GreaterOrEqual(t, len(results), 2, "should find at least 2 authentication-related observations")
|
||||
|
||||
// Search for database - should find 1 observation
|
||||
results, err = obsStore.SearchObservationsFTS(ctx, "database PostgreSQL", "project-a", 50)
|
||||
require.NoError(t, err)
|
||||
assert.GreaterOrEqual(t, len(results), 1, "should find at least 1 database-related observation")
|
||||
}
|
||||
|
||||
func TestObservationStore_SearchObservationsFTS_LimitRespected(t *testing.T) {
|
||||
obsStore, _, cleanup := testObservationStore(t)
|
||||
defer cleanup()
|
||||
|
||||
// FTS5 tables are created by testObservationStore via testutil.CreateAllTables
|
||||
ctx := context.Background()
|
||||
|
||||
// Create 20 observations with similar content
|
||||
for i := 0; i < 20; i++ {
|
||||
obs := &models.ParsedObservation{
|
||||
Type: models.ObsTypeDiscovery,
|
||||
Title: "Testing observation " + string(rune('A'+i)),
|
||||
Narrative: "This is about testing and quality assurance " + string(rune('A'+i)),
|
||||
}
|
||||
_, _, err := obsStore.StoreObservation(ctx, "session-1", "project-a", obs, 1, 100)
|
||||
require.NoError(t, err)
|
||||
time.Sleep(time.Millisecond)
|
||||
}
|
||||
|
||||
// Search with limit 5
|
||||
results, err := obsStore.SearchObservationsFTS(ctx, "testing quality", "project-a", 5)
|
||||
require.NoError(t, err)
|
||||
assert.LessOrEqual(t, len(results), 5, "should respect limit of 5")
|
||||
|
||||
// Search with limit 15
|
||||
results, err = obsStore.SearchObservationsFTS(ctx, "testing quality", "project-a", 15)
|
||||
require.NoError(t, err)
|
||||
assert.LessOrEqual(t, len(results), 15, "should respect limit of 15")
|
||||
|
||||
// Search with limit 50 (our new default)
|
||||
results, err = obsStore.SearchObservationsFTS(ctx, "testing quality", "project-a", 50)
|
||||
require.NoError(t, err)
|
||||
assert.LessOrEqual(t, len(results), 50, "should respect limit of 50")
|
||||
assert.Equal(t, 20, len(results), "should return all 20 matching observations")
|
||||
}
|
||||
|
||||
func TestObservationStore_GlobalScope(t *testing.T) {
|
||||
obsStore, _, cleanup := testObservationStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Create a project-scoped observation
|
||||
projectObs := &models.ParsedObservation{
|
||||
Type: models.ObsTypeDiscovery,
|
||||
Title: "Project specific code",
|
||||
Narrative: "This is specific to project-a",
|
||||
Concepts: []string{"project-specific"},
|
||||
}
|
||||
_, _, err := obsStore.StoreObservation(ctx, "session-1", "project-a", projectObs, 1, 100)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Create a global-scoped observation (has a globalizable concept)
|
||||
globalObs := &models.ParsedObservation{
|
||||
Type: models.ObsTypeDiscovery,
|
||||
Title: "Security best practice",
|
||||
Narrative: "Always validate user input",
|
||||
Concepts: []string{"security", "best-practice"}, // "security" is in GlobalizableConcepts
|
||||
}
|
||||
_, _, err = obsStore.StoreObservation(ctx, "session-1", "project-a", globalObs, 1, 100)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Get recent for project-a - should see both
|
||||
results, err := obsStore.GetRecentObservations(ctx, "project-a", 10)
|
||||
require.NoError(t, err)
|
||||
assert.Len(t, results, 2)
|
||||
|
||||
// Get recent for project-b - should only see global observation
|
||||
results, err = obsStore.GetRecentObservations(ctx, "project-b", 10)
|
||||
require.NoError(t, err)
|
||||
assert.Len(t, results, 1)
|
||||
assert.Equal(t, "Security best practice", results[0].Title.String)
|
||||
assert.Equal(t, models.ScopeGlobal, results[0].Scope)
|
||||
}
|
||||
|
||||
func TestObservationStore_DeleteObservations(t *testing.T) {
|
||||
obsStore, _, cleanup := testObservationStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Create observations
|
||||
var ids []int64
|
||||
for i := 0; i < 5; i++ {
|
||||
obs := &models.ParsedObservation{
|
||||
Type: models.ObsTypeDiscovery,
|
||||
Title: "Observation " + string(rune('A'+i)),
|
||||
}
|
||||
id, _, err := obsStore.StoreObservation(ctx, "session-1", "project-a", obs, 1, 100)
|
||||
require.NoError(t, err)
|
||||
ids = append(ids, id)
|
||||
}
|
||||
|
||||
// Verify all exist
|
||||
all, err := obsStore.GetRecentObservations(ctx, "project-a", 10)
|
||||
require.NoError(t, err)
|
||||
assert.Len(t, all, 5)
|
||||
|
||||
// Delete first 3
|
||||
deleted, err := obsStore.DeleteObservations(ctx, ids[:3])
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, int64(3), deleted)
|
||||
|
||||
// Verify only 2 remain
|
||||
remaining, err := obsStore.GetRecentObservations(ctx, "project-a", 10)
|
||||
require.NoError(t, err)
|
||||
assert.Len(t, remaining, 2)
|
||||
}
|
||||
|
||||
func TestObservationStore_GetObservationCount(t *testing.T) {
|
||||
obsStore, _, cleanup := testObservationStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Create observations for different projects
|
||||
for i := 0; i < 5; i++ {
|
||||
obs := &models.ParsedObservation{
|
||||
Type: models.ObsTypeDiscovery,
|
||||
Title: "Project A observation " + string(rune('0'+i)),
|
||||
}
|
||||
_, _, err := obsStore.StoreObservation(ctx, "session-1", "project-a", obs, 1, 100)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
for i := 0; i < 3; i++ {
|
||||
obs := &models.ParsedObservation{
|
||||
Type: models.ObsTypeDiscovery,
|
||||
Title: "Project B observation " + string(rune('0'+i)),
|
||||
}
|
||||
_, _, err := obsStore.StoreObservation(ctx, "session-1", "project-b", obs, 1, 100)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
// Create a global observation
|
||||
globalObs := &models.ParsedObservation{
|
||||
Type: models.ObsTypeDiscovery,
|
||||
Title: "Global observation",
|
||||
Concepts: []string{"best-practice"}, // Makes it global
|
||||
}
|
||||
_, _, err := obsStore.StoreObservation(ctx, "session-1", "project-a", globalObs, 1, 100)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Count for project-a includes its own + global
|
||||
count, err := obsStore.GetObservationCount(ctx, "project-a")
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 6, count) // 5 project-a + 1 global
|
||||
|
||||
// Count for project-b includes its own + global
|
||||
count, err = obsStore.GetObservationCount(ctx, "project-b")
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 4, count) // 3 project-b + 1 global
|
||||
}
|
||||
|
||||
func TestObservationStore_CleanupOldObservations(t *testing.T) {
|
||||
obsStore, _, cleanup := testObservationStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Create more observations than the limit (MaxObservationsPerProject = 100)
|
||||
// We'll create a smaller number and verify the logic works
|
||||
for i := 0; i < 10; i++ {
|
||||
obs := &models.ParsedObservation{
|
||||
Type: models.ObsTypeDiscovery,
|
||||
Title: "Observation " + string(rune('A'+i)),
|
||||
}
|
||||
_, _, err := obsStore.StoreObservation(ctx, "session-1", "project-a", obs, i+1, 100)
|
||||
require.NoError(t, err)
|
||||
time.Sleep(time.Millisecond)
|
||||
}
|
||||
|
||||
// Cleanup should return empty since we're under the limit
|
||||
deletedIDs, err := obsStore.CleanupOldObservations(ctx, "project-a")
|
||||
require.NoError(t, err)
|
||||
assert.Empty(t, deletedIDs)
|
||||
|
||||
// All 10 should still exist
|
||||
count, err := obsStore.GetObservationCount(ctx, "project-a")
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 10, count)
|
||||
}
|
||||
|
||||
func TestObservationStore_SetCleanupFunc(t *testing.T) {
|
||||
obsStore, _, cleanup := testObservationStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Track cleanup calls
|
||||
var cleanupCalledWith []int64
|
||||
obsStore.SetCleanupFunc(func(ctx context.Context, deletedIDs []int64) {
|
||||
cleanupCalledWith = deletedIDs
|
||||
})
|
||||
|
||||
// Store an observation (should trigger cleanup, but won't delete anything under limit)
|
||||
obs := &models.ParsedObservation{
|
||||
Type: models.ObsTypeDiscovery,
|
||||
Title: "Test observation",
|
||||
}
|
||||
_, _, err := obsStore.StoreObservation(ctx, "session-1", "project-a", obs, 1, 100)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Cleanup func should not have been called since nothing was deleted
|
||||
assert.Empty(t, cleanupCalledWith)
|
||||
}
|
||||
|
||||
func TestExtractKeywords(t *testing.T) {
|
||||
tests := []struct {
|
||||
query string
|
||||
expected []string
|
||||
}{
|
||||
{
|
||||
query: "What is the authentication flow?",
|
||||
expected: []string{"authentication", "flow"},
|
||||
},
|
||||
{
|
||||
query: "How does the database connection work?",
|
||||
expected: []string{"database", "connection"},
|
||||
},
|
||||
{
|
||||
query: "JWT token validation",
|
||||
expected: []string{"token", "validation"},
|
||||
},
|
||||
{
|
||||
query: "the a an is are", // All stop words
|
||||
expected: nil,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.query, func(t *testing.T) {
|
||||
keywords := extractKeywords(tt.query)
|
||||
for _, exp := range tt.expected {
|
||||
assert.Contains(t, keywords, exp, "should contain keyword: "+exp)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestObservationStore_GetObservationsByProjectStrict(t *testing.T) {
|
||||
obsStore, _, cleanup := testObservationStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Create project-scoped observation for project-a
|
||||
projectObs := &models.ParsedObservation{
|
||||
Type: models.ObsTypeDiscovery,
|
||||
Title: "Project A specific",
|
||||
Narrative: "Only for project-a",
|
||||
Concepts: []string{"local-concept"},
|
||||
}
|
||||
_, _, err := obsStore.StoreObservation(ctx, "session-1", "project-a", projectObs, 1, 100)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Create global observation from project-a
|
||||
globalObs := &models.ParsedObservation{
|
||||
Type: models.ObsTypeDiscovery,
|
||||
Title: "Global security practice",
|
||||
Narrative: "Best practice for all",
|
||||
Concepts: []string{"security", "best-practice"},
|
||||
}
|
||||
_, _, err = obsStore.StoreObservation(ctx, "session-1", "project-a", globalObs, 2, 100)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Create observation for project-b
|
||||
projectBObs := &models.ParsedObservation{
|
||||
Type: models.ObsTypeDiscovery,
|
||||
Title: "Project B specific",
|
||||
Narrative: "Only for project-b",
|
||||
}
|
||||
_, _, err = obsStore.StoreObservation(ctx, "session-1", "project-b", projectBObs, 1, 100)
|
||||
require.NoError(t, err)
|
||||
|
||||
// GetObservationsByProjectStrict for project-a should only return project-a observations
|
||||
// This is different from GetRecentObservations which includes globals from other projects
|
||||
results, err := obsStore.GetObservationsByProjectStrict(ctx, "project-a", 10)
|
||||
require.NoError(t, err)
|
||||
assert.Len(t, results, 2) // Only observations created in project-a
|
||||
|
||||
// Verify both are from project-a
|
||||
for _, obs := range results {
|
||||
assert.Equal(t, "project-a", obs.Project)
|
||||
}
|
||||
|
||||
// GetObservationsByProjectStrict for project-b should only return project-b observations
|
||||
results, err = obsStore.GetObservationsByProjectStrict(ctx, "project-b", 10)
|
||||
require.NoError(t, err)
|
||||
assert.Len(t, results, 1)
|
||||
assert.Equal(t, "Project B specific", results[0].Title.String)
|
||||
}
|
||||
|
||||
func TestObservationStore_SearchObservationsFTS_EmptyQuery(t *testing.T) {
|
||||
obsStore, _, cleanup := testObservationStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Create an observation
|
||||
obs := &models.ParsedObservation{
|
||||
Type: models.ObsTypeDiscovery,
|
||||
Title: "Test observation",
|
||||
Narrative: "Some content here",
|
||||
}
|
||||
_, _, err := obsStore.StoreObservation(ctx, "session-1", "project-a", obs, 1, 100)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Search with only stop words (should return nil)
|
||||
results, err := obsStore.SearchObservationsFTS(ctx, "the a an is are", "project-a", 10)
|
||||
require.NoError(t, err)
|
||||
assert.Nil(t, results)
|
||||
|
||||
// Search with empty query
|
||||
results, err = obsStore.SearchObservationsFTS(ctx, "", "project-a", 10)
|
||||
require.NoError(t, err)
|
||||
assert.Nil(t, results)
|
||||
}
|
||||
|
||||
func TestObservationStore_SearchObservationsFTS_DefaultLimit(t *testing.T) {
|
||||
obsStore, _, cleanup := testObservationStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Create observations
|
||||
for i := 0; i < 15; i++ {
|
||||
obs := &models.ParsedObservation{
|
||||
Type: models.ObsTypeDiscovery,
|
||||
Title: "Authentication test " + string(rune('A'+i)),
|
||||
Narrative: "Auth related content",
|
||||
}
|
||||
_, _, err := obsStore.StoreObservation(ctx, "session-1", "project-a", obs, i+1, 100)
|
||||
require.NoError(t, err)
|
||||
time.Sleep(time.Millisecond)
|
||||
}
|
||||
|
||||
// Search with limit 0 (should default to 10)
|
||||
results, err := obsStore.SearchObservationsFTS(ctx, "authentication", "project-a", 0)
|
||||
require.NoError(t, err)
|
||||
assert.LessOrEqual(t, len(results), 10)
|
||||
|
||||
// Search with negative limit (should default to 10)
|
||||
results, err = obsStore.SearchObservationsFTS(ctx, "authentication", "project-a", -5)
|
||||
require.NoError(t, err)
|
||||
assert.LessOrEqual(t, len(results), 10)
|
||||
}
|
||||
|
||||
func TestObservationStore_GetAllRecentObservations(t *testing.T) {
|
||||
obsStore, _, cleanup := testObservationStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Create observations across different projects
|
||||
projects := []string{"project-a", "project-b", "project-c"}
|
||||
for _, proj := range projects {
|
||||
for i := 0; i < 3; i++ {
|
||||
obs := &models.ParsedObservation{
|
||||
Type: models.ObsTypeDiscovery,
|
||||
Title: proj + " observation " + string(rune('A'+i)),
|
||||
Narrative: "Content for " + proj,
|
||||
}
|
||||
_, _, err := obsStore.StoreObservation(ctx, "session-1", proj, obs, i+1, 100)
|
||||
require.NoError(t, err)
|
||||
time.Sleep(time.Millisecond)
|
||||
}
|
||||
}
|
||||
|
||||
// Get all recent observations
|
||||
results, err := obsStore.GetAllRecentObservations(ctx, 100)
|
||||
require.NoError(t, err)
|
||||
assert.Len(t, results, 9) // 3 projects * 3 observations
|
||||
|
||||
// Verify they are in descending order by epoch
|
||||
for i := 1; i < len(results); i++ {
|
||||
assert.GreaterOrEqual(t, results[i-1].CreatedAtEpoch, results[i].CreatedAtEpoch)
|
||||
}
|
||||
|
||||
// Test with limit
|
||||
results, err = obsStore.GetAllRecentObservations(ctx, 5)
|
||||
require.NoError(t, err)
|
||||
assert.Len(t, results, 5)
|
||||
}
|
||||
@@ -0,0 +1,249 @@
|
||||
// Package sqlite provides SQLite database operations for claude-mnemonic.
|
||||
package sqlite
|
||||
|
||||
import (
|
||||
"context"
|
||||
"time"
|
||||
|
||||
"github.com/lukaszraczylo/claude-mnemonic/pkg/models"
|
||||
)
|
||||
|
||||
// PromptCleanupFunc is a callback for when prompts are cleaned up.
|
||||
// Receives the IDs of deleted prompts for downstream cleanup (e.g., vector DB).
|
||||
type PromptCleanupFunc func(ctx context.Context, deletedIDs []int64)
|
||||
|
||||
// MaxPromptsGlobal is the hard limit of prompts across all projects.
|
||||
const MaxPromptsGlobal = 500
|
||||
|
||||
// PromptStore provides user prompt-related database operations.
|
||||
type PromptStore struct {
|
||||
store *Store
|
||||
cleanupFunc PromptCleanupFunc
|
||||
}
|
||||
|
||||
// NewPromptStore creates a new prompt store.
|
||||
func NewPromptStore(store *Store) *PromptStore {
|
||||
return &PromptStore{store: store}
|
||||
}
|
||||
|
||||
// SetCleanupFunc sets the callback for when prompts are deleted during cleanup.
|
||||
func (s *PromptStore) SetCleanupFunc(fn PromptCleanupFunc) {
|
||||
s.cleanupFunc = fn
|
||||
}
|
||||
|
||||
// SaveUserPromptWithMatches saves a user prompt with matched observation count.
|
||||
// Uses INSERT OR IGNORE to be idempotent - duplicate (session, prompt_number) pairs are silently ignored.
|
||||
// This prevents duplicate prompts when the user-prompt hook fires multiple times.
|
||||
func (s *PromptStore) SaveUserPromptWithMatches(ctx context.Context, claudeSessionID string, promptNumber int, promptText string, matchedObservations int) (int64, error) {
|
||||
now := time.Now()
|
||||
|
||||
// Use INSERT OR IGNORE for idempotency - if (claude_session_id, prompt_number) already exists,
|
||||
// the insert is silently ignored. This handles concurrent/duplicate hook invocations.
|
||||
const query = `
|
||||
INSERT OR IGNORE INTO user_prompts
|
||||
(claude_session_id, prompt_number, prompt_text, matched_observations, created_at, created_at_epoch)
|
||||
VALUES (?, ?, ?, ?, ?, ?)
|
||||
`
|
||||
|
||||
result, err := s.store.ExecContext(ctx, query,
|
||||
claudeSessionID, promptNumber, promptText, matchedObservations,
|
||||
now.Format(time.RFC3339), now.UnixMilli(),
|
||||
)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
|
||||
id, _ := result.LastInsertId()
|
||||
|
||||
// If id is 0, the insert was ignored (duplicate) - fetch the existing ID
|
||||
if id == 0 {
|
||||
const selectQuery = `SELECT id FROM user_prompts WHERE claude_session_id = ? AND prompt_number = ?`
|
||||
row := s.store.QueryRowContext(ctx, selectQuery, claudeSessionID, promptNumber)
|
||||
if err := row.Scan(&id); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
// Return existing ID without triggering cleanup (already handled when first inserted)
|
||||
return id, nil
|
||||
}
|
||||
|
||||
// Cleanup old prompts beyond the global limit (async to not block handler)
|
||||
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)
|
||||
}
|
||||
}()
|
||||
|
||||
return id, nil
|
||||
}
|
||||
|
||||
// CleanupOldPrompts deletes prompts beyond the global limit.
|
||||
// Keeps the most recent MaxPromptsGlobal prompts.
|
||||
// Returns the IDs of deleted prompts for downstream cleanup (e.g., vector DB).
|
||||
func (s *PromptStore) CleanupOldPrompts(ctx context.Context) ([]int64, error) {
|
||||
// First, find IDs that will be deleted
|
||||
const selectQuery = `
|
||||
SELECT id FROM user_prompts
|
||||
WHERE id NOT IN (
|
||||
SELECT id FROM user_prompts
|
||||
ORDER BY created_at_epoch DESC
|
||||
LIMIT ?
|
||||
)
|
||||
`
|
||||
|
||||
rows, err := s.store.QueryContext(ctx, selectQuery, MaxPromptsGlobal)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
var toDelete []int64
|
||||
for rows.Next() {
|
||||
var id int64
|
||||
if err := rows.Scan(&id); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
toDelete = append(toDelete, id)
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if len(toDelete) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
// Delete the prompts
|
||||
const deleteQuery = `
|
||||
DELETE FROM user_prompts
|
||||
WHERE id NOT IN (
|
||||
SELECT id FROM user_prompts
|
||||
ORDER BY created_at_epoch DESC
|
||||
LIMIT ?
|
||||
)
|
||||
`
|
||||
|
||||
_, err = s.store.ExecContext(ctx, deleteQuery, MaxPromptsGlobal)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return toDelete, nil
|
||||
}
|
||||
|
||||
// GetPromptsByIDs retrieves user prompts by a list of IDs.
|
||||
func (s *PromptStore) GetPromptsByIDs(ctx context.Context, ids []int64, orderBy string, limit int) ([]*models.UserPromptWithSession, error) {
|
||||
if len(ids) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
// Build query with placeholders
|
||||
// #nosec G202 -- query uses parameterized placeholders, not user input
|
||||
query := `
|
||||
SELECT up.id, up.claude_session_id, up.prompt_number, up.prompt_text,
|
||||
COALESCE(up.matched_observations, 0) as matched_observations,
|
||||
up.created_at, up.created_at_epoch,
|
||||
COALESCE(s.project, '') as project,
|
||||
COALESCE(s.sdk_session_id, '') as sdk_session_id
|
||||
FROM user_prompts up
|
||||
LEFT JOIN sdk_sessions s ON up.claude_session_id = s.claude_session_id
|
||||
WHERE up.id IN (?` + repeatPlaceholders(len(ids)-1) + `)
|
||||
ORDER BY up.created_at_epoch `
|
||||
|
||||
if orderBy == "date_asc" {
|
||||
query += "ASC"
|
||||
} else {
|
||||
query += "DESC"
|
||||
}
|
||||
|
||||
if limit > 0 {
|
||||
query += " LIMIT ?"
|
||||
}
|
||||
|
||||
args := int64SliceToInterface(ids)
|
||||
if limit > 0 {
|
||||
args = append(args, limit)
|
||||
}
|
||||
|
||||
rows, err := s.store.db.QueryContext(ctx, query, args...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
return scanPromptWithSessionRows(rows)
|
||||
}
|
||||
|
||||
// GetAllRecentUserPrompts retrieves recent user prompts across all sessions.
|
||||
func (s *PromptStore) GetAllRecentUserPrompts(ctx context.Context, limit int) ([]*models.UserPromptWithSession, error) {
|
||||
const query = `
|
||||
SELECT up.id, up.claude_session_id, up.prompt_number, up.prompt_text,
|
||||
COALESCE(up.matched_observations, 0) as matched_observations,
|
||||
up.created_at, up.created_at_epoch,
|
||||
COALESCE(s.project, '') as project,
|
||||
COALESCE(s.sdk_session_id, '') as sdk_session_id
|
||||
FROM user_prompts up
|
||||
LEFT JOIN sdk_sessions s ON up.claude_session_id = s.claude_session_id
|
||||
ORDER BY up.created_at_epoch DESC
|
||||
LIMIT ?
|
||||
`
|
||||
|
||||
rows, err := s.store.QueryContext(ctx, query, limit)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
return scanPromptWithSessionRows(rows)
|
||||
}
|
||||
|
||||
// FindRecentPromptByText finds a prompt with the same text for a session within the last few seconds.
|
||||
// This is used to detect duplicate hook invocations.
|
||||
// Returns (promptID, promptNumber, found).
|
||||
func (s *PromptStore) FindRecentPromptByText(ctx context.Context, claudeSessionID, promptText string, withinSeconds int) (int64, int, bool) {
|
||||
// Look for an existing prompt with the same text within the time window
|
||||
// This catches duplicate hook invocations that happen in quick succession
|
||||
const query = `
|
||||
SELECT id, prompt_number FROM user_prompts
|
||||
WHERE claude_session_id = ? AND prompt_text = ?
|
||||
AND created_at_epoch > ?
|
||||
ORDER BY created_at_epoch DESC
|
||||
LIMIT 1
|
||||
`
|
||||
|
||||
cutoff := time.Now().Add(-time.Duration(withinSeconds) * time.Second).UnixMilli()
|
||||
|
||||
var id int64
|
||||
var promptNumber int
|
||||
err := s.store.QueryRowContext(ctx, query, claudeSessionID, promptText, cutoff).Scan(&id, &promptNumber)
|
||||
if err != nil {
|
||||
return 0, 0, false
|
||||
}
|
||||
return id, 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) {
|
||||
const query = `
|
||||
SELECT up.id, up.claude_session_id, up.prompt_number, up.prompt_text,
|
||||
COALESCE(up.matched_observations, 0) as matched_observations,
|
||||
up.created_at, up.created_at_epoch,
|
||||
COALESCE(s.project, '') as project,
|
||||
COALESCE(s.sdk_session_id, '') as sdk_session_id
|
||||
FROM user_prompts up
|
||||
LEFT JOIN sdk_sessions s ON up.claude_session_id = s.claude_session_id
|
||||
WHERE s.project = ?
|
||||
ORDER BY up.created_at_epoch DESC
|
||||
LIMIT ?
|
||||
`
|
||||
|
||||
rows, err := s.store.QueryContext(ctx, query, project, limit)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
return scanPromptWithSessionRows(rows)
|
||||
}
|
||||
@@ -0,0 +1,289 @@
|
||||
package sqlite
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func testPromptStore(t *testing.T) (*PromptStore, *Store, func()) {
|
||||
t.Helper()
|
||||
|
||||
db, _, cleanup := testDB(t)
|
||||
createAllTables(t, db)
|
||||
|
||||
store := newStoreFromDB(db)
|
||||
promptStore := NewPromptStore(store)
|
||||
|
||||
return promptStore, store, cleanup
|
||||
}
|
||||
|
||||
func TestPromptStore_SaveUserPromptWithMatches(t *testing.T) {
|
||||
promptStore, store, cleanup := testPromptStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Create a session first
|
||||
seedSession(t, storeDB(store), "claude-1", "sdk-1", "test-project")
|
||||
|
||||
// Save a prompt
|
||||
id, err := promptStore.SaveUserPromptWithMatches(ctx, "claude-1", 1, "Help me fix this bug", 5)
|
||||
require.NoError(t, err)
|
||||
assert.Greater(t, id, int64(0))
|
||||
|
||||
// Verify it was saved
|
||||
var count int
|
||||
err = storeDB(store).QueryRow("SELECT COUNT(*) FROM user_prompts WHERE id = ?", id).Scan(&count)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 1, count)
|
||||
}
|
||||
|
||||
func TestPromptStore_GetAllRecentUserPrompts(t *testing.T) {
|
||||
promptStore, store, cleanup := testPromptStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Create a session
|
||||
seedSession(t, storeDB(store), "claude-1", "sdk-1", "test-project")
|
||||
|
||||
// Save multiple prompts
|
||||
for i := 1; i <= 5; i++ {
|
||||
_, err := promptStore.SaveUserPromptWithMatches(ctx, "claude-1", i, "Prompt "+string(rune('A'+i-1)), i)
|
||||
require.NoError(t, err)
|
||||
time.Sleep(time.Millisecond) // Ensure different timestamps
|
||||
}
|
||||
|
||||
// Get recent prompts
|
||||
prompts, err := promptStore.GetAllRecentUserPrompts(ctx, 3)
|
||||
require.NoError(t, err)
|
||||
assert.Len(t, prompts, 3)
|
||||
|
||||
// Should be in descending order (most recent first)
|
||||
assert.Equal(t, 5, prompts[0].PromptNumber)
|
||||
}
|
||||
|
||||
func TestPromptStore_GetRecentUserPromptsByProject(t *testing.T) {
|
||||
promptStore, store, cleanup := testPromptStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Create sessions for different projects
|
||||
seedSession(t, storeDB(store), "claude-1", "sdk-1", "project-a")
|
||||
seedSession(t, storeDB(store), "claude-2", "sdk-2", "project-b")
|
||||
|
||||
// Save prompts for both projects
|
||||
for i := 1; i <= 3; i++ {
|
||||
_, err := promptStore.SaveUserPromptWithMatches(ctx, "claude-1", i, "Project A prompt", 0)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
for i := 1; i <= 2; i++ {
|
||||
_, err := promptStore.SaveUserPromptWithMatches(ctx, "claude-2", i, "Project B prompt", 0)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
// Get prompts for project-a
|
||||
prompts, err := promptStore.GetRecentUserPromptsByProject(ctx, "project-a", 10)
|
||||
require.NoError(t, err)
|
||||
assert.Len(t, prompts, 3)
|
||||
|
||||
// Get prompts for project-b
|
||||
prompts, err = promptStore.GetRecentUserPromptsByProject(ctx, "project-b", 10)
|
||||
require.NoError(t, err)
|
||||
assert.Len(t, prompts, 2)
|
||||
}
|
||||
|
||||
func TestPromptStore_CleanupOldPrompts(t *testing.T) {
|
||||
promptStore, store, cleanup := testPromptStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Create a session
|
||||
seedSession(t, storeDB(store), "claude-1", "sdk-1", "test-project")
|
||||
|
||||
// Save more prompts than the limit
|
||||
// Note: MaxPromptsGlobal is 500, but we'll test with a smaller number
|
||||
// by directly calling CleanupOldPrompts
|
||||
for i := 1; i <= 10; i++ {
|
||||
_, err := storeDB(store).Exec(`
|
||||
INSERT INTO user_prompts (claude_session_id, prompt_number, prompt_text, created_at, created_at_epoch)
|
||||
VALUES (?, ?, ?, datetime('now'), ?)
|
||||
`, "claude-1", i, "Prompt "+string(rune('A'+i-1)), time.Now().UnixMilli()+int64(i))
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
// Verify we have 10 prompts
|
||||
var count int
|
||||
err := storeDB(store).QueryRow("SELECT COUNT(*) FROM user_prompts").Scan(&count)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 10, count)
|
||||
|
||||
// Cleanup should return empty since we're under the limit
|
||||
deletedIDs, err := promptStore.CleanupOldPrompts(ctx)
|
||||
require.NoError(t, err)
|
||||
assert.Empty(t, deletedIDs)
|
||||
}
|
||||
|
||||
func TestPromptStore_SetCleanupFunc(t *testing.T) {
|
||||
promptStore, store, cleanup := testPromptStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Track cleanup calls
|
||||
var cleanupCalledWith []int64
|
||||
promptStore.SetCleanupFunc(func(ctx context.Context, deletedIDs []int64) {
|
||||
cleanupCalledWith = deletedIDs
|
||||
})
|
||||
|
||||
// Create a session
|
||||
seedSession(t, storeDB(store), "claude-1", "sdk-1", "test-project")
|
||||
|
||||
// Save a prompt (should trigger cleanup, but won't delete anything under limit)
|
||||
_, err := promptStore.SaveUserPromptWithMatches(ctx, "claude-1", 1, "Test prompt", 0)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Cleanup func should not have been called since nothing was deleted
|
||||
assert.Empty(t, cleanupCalledWith)
|
||||
}
|
||||
|
||||
func TestPromptStore_GetPromptsByIDs(t *testing.T) {
|
||||
promptStore, store, cleanup := testPromptStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Create a session
|
||||
seedSession(t, storeDB(store), "claude-1", "sdk-1", "test-project")
|
||||
|
||||
// Save some prompts and collect their IDs
|
||||
var ids []int64
|
||||
for i := 1; i <= 5; i++ {
|
||||
id, err := promptStore.SaveUserPromptWithMatches(ctx, "claude-1", i, "Prompt "+string(rune('A'+i-1)), 0)
|
||||
require.NoError(t, err)
|
||||
ids = append(ids, id)
|
||||
time.Sleep(time.Millisecond)
|
||||
}
|
||||
|
||||
// Get specific prompts by ID
|
||||
prompts, err := promptStore.GetPromptsByIDs(ctx, ids[:3], "date_desc", 10)
|
||||
require.NoError(t, err)
|
||||
assert.Len(t, prompts, 3)
|
||||
|
||||
// Test with ascending order
|
||||
prompts, err = promptStore.GetPromptsByIDs(ctx, ids, "date_asc", 2)
|
||||
require.NoError(t, err)
|
||||
assert.Len(t, prompts, 2)
|
||||
assert.Equal(t, 1, prompts[0].PromptNumber)
|
||||
}
|
||||
|
||||
func TestPromptStore_GetPromptsByIDs_EmptyInput(t *testing.T) {
|
||||
promptStore, _, cleanup := testPromptStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Empty IDs should return nil
|
||||
prompts, err := promptStore.GetPromptsByIDs(ctx, []int64{}, "date_desc", 10)
|
||||
require.NoError(t, err)
|
||||
assert.Nil(t, prompts)
|
||||
}
|
||||
|
||||
func TestPromptStore_FindRecentPromptByText(t *testing.T) {
|
||||
promptStore, store, cleanup := testPromptStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Create a session
|
||||
seedSession(t, storeDB(store), "claude-1", "sdk-1", "test-project")
|
||||
|
||||
// Save a prompt
|
||||
_, err := promptStore.SaveUserPromptWithMatches(ctx, "claude-1", 1, "Help me fix this bug in the code", 0)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Find the prompt by text - returns (id, promptNumber, found)
|
||||
id, promptNum, found := promptStore.FindRecentPromptByText(ctx, "claude-1", "Help me fix this bug in the code", 60)
|
||||
assert.True(t, found, "should find the exact prompt text")
|
||||
assert.Greater(t, id, int64(0))
|
||||
assert.Equal(t, 1, promptNum)
|
||||
|
||||
// Try to find non-existent prompt
|
||||
_, _, found = promptStore.FindRecentPromptByText(ctx, "claude-1", "This prompt does not exist", 60)
|
||||
assert.False(t, found, "should not find non-existent prompt")
|
||||
|
||||
// Try with different session
|
||||
_, _, found = promptStore.FindRecentPromptByText(ctx, "claude-2", "Help me fix this bug in the code", 60)
|
||||
assert.False(t, found, "should not find prompt for different session")
|
||||
}
|
||||
|
||||
func TestPromptStore_FindRecentPromptByText_WindowSeconds(t *testing.T) {
|
||||
promptStore, store, cleanup := testPromptStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Create a session
|
||||
seedSession(t, storeDB(store), "claude-1", "sdk-1", "test-project")
|
||||
|
||||
// Save a prompt with an old timestamp
|
||||
oldEpoch := time.Now().Add(-2 * time.Hour).UnixMilli()
|
||||
_, err := storeDB(store).Exec(`
|
||||
INSERT INTO user_prompts (claude_session_id, prompt_number, prompt_text, created_at, created_at_epoch)
|
||||
VALUES (?, ?, ?, datetime('now'), ?)
|
||||
`, "claude-1", 1, "Old prompt text", oldEpoch)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Search within last hour - should not find old prompt
|
||||
_, _, found := promptStore.FindRecentPromptByText(ctx, "claude-1", "Old prompt text", 3600)
|
||||
assert.False(t, found, "should not find prompt outside window")
|
||||
|
||||
// Search within last 3 hours - should find old prompt
|
||||
_, _, found = promptStore.FindRecentPromptByText(ctx, "claude-1", "Old prompt text", 3*3600)
|
||||
assert.True(t, found, "should find prompt within extended window")
|
||||
}
|
||||
|
||||
func TestPromptStore_SaveMultiplePrompts(t *testing.T) {
|
||||
promptStore, store, cleanup := testPromptStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Create sessions
|
||||
seedSession(t, storeDB(store), "claude-1", "sdk-1", "project-x")
|
||||
seedSession(t, storeDB(store), "claude-2", "sdk-2", "project-y")
|
||||
|
||||
tests := []struct {
|
||||
claudeSessionID string
|
||||
promptNum int
|
||||
text string
|
||||
matches int
|
||||
}{
|
||||
{"claude-1", 1, "First prompt", 5},
|
||||
{"claude-1", 2, "Second prompt", 3},
|
||||
{"claude-2", 1, "Third prompt", 0},
|
||||
{"claude-1", 3, "Fourth prompt", 10},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
id, err := promptStore.SaveUserPromptWithMatches(ctx, tt.claudeSessionID, tt.promptNum, tt.text, tt.matches)
|
||||
require.NoError(t, err)
|
||||
assert.Greater(t, id, int64(0))
|
||||
}
|
||||
|
||||
// Verify counts
|
||||
var count int
|
||||
err := storeDB(store).QueryRow("SELECT COUNT(*) FROM user_prompts WHERE claude_session_id = 'claude-1'").Scan(&count)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 3, count)
|
||||
|
||||
err = storeDB(store).QueryRow("SELECT COUNT(*) FROM user_prompts WHERE claude_session_id = 'claude-2'").Scan(&count)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 1, count)
|
||||
}
|
||||
@@ -0,0 +1,184 @@
|
||||
// Package sqlite provides SQLite database operations for claude-mnemonic.
|
||||
package sqlite
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"time"
|
||||
|
||||
"github.com/lukaszraczylo/claude-mnemonic/pkg/models"
|
||||
)
|
||||
|
||||
// SessionStore provides session-related database operations.
|
||||
type SessionStore struct {
|
||||
store *Store
|
||||
}
|
||||
|
||||
// NewSessionStore creates a new session store.
|
||||
func NewSessionStore(store *Store) *SessionStore {
|
||||
return &SessionStore{store: store}
|
||||
}
|
||||
|
||||
// CreateSDKSession creates a new SDK session (idempotent - returns existing ID if exists).
|
||||
// This is the KEY to how claude-mnemonic stays unified across hooks.
|
||||
func (s *SessionStore) CreateSDKSession(ctx context.Context, claudeSessionID, project, userPrompt string) (int64, error) {
|
||||
now := time.Now()
|
||||
|
||||
// CRITICAL: INSERT OR IGNORE makes this idempotent
|
||||
const query = `
|
||||
INSERT OR IGNORE INTO sdk_sessions
|
||||
(claude_session_id, sdk_session_id, project, user_prompt, started_at, started_at_epoch, status)
|
||||
VALUES (?, ?, ?, ?, ?, ?, 'active')
|
||||
`
|
||||
|
||||
result, err := s.store.ExecContext(ctx, query,
|
||||
claudeSessionID, claudeSessionID, project, userPrompt,
|
||||
now.Format(time.RFC3339), now.UnixMilli(),
|
||||
)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
|
||||
// Check if insert happened
|
||||
rowsAffected, _ := result.RowsAffected()
|
||||
if rowsAffected == 0 {
|
||||
// Session exists - UPDATE project and user_prompt if we have non-empty values
|
||||
if project != "" {
|
||||
const updateQuery = `
|
||||
UPDATE sdk_sessions
|
||||
SET project = ?, user_prompt = ?
|
||||
WHERE claude_session_id = ?
|
||||
`
|
||||
_, _ = s.store.ExecContext(ctx, updateQuery, project, userPrompt, claudeSessionID)
|
||||
}
|
||||
|
||||
// Fetch existing ID
|
||||
var id int64
|
||||
const selectQuery = `SELECT id FROM sdk_sessions WHERE claude_session_id = ? LIMIT 1`
|
||||
err := s.store.QueryRowContext(ctx, selectQuery, claudeSessionID).Scan(&id)
|
||||
return id, err
|
||||
}
|
||||
|
||||
id, _ := result.LastInsertId()
|
||||
return id, nil
|
||||
}
|
||||
|
||||
// GetSessionByID retrieves a session by its database ID.
|
||||
func (s *SessionStore) GetSessionByID(ctx context.Context, id int64) (*models.SDKSession, error) {
|
||||
const query = `
|
||||
SELECT id, claude_session_id, sdk_session_id, project, user_prompt,
|
||||
worker_port, prompt_counter, status, started_at, started_at_epoch,
|
||||
completed_at, completed_at_epoch
|
||||
FROM sdk_sessions
|
||||
WHERE id = ?
|
||||
LIMIT 1
|
||||
`
|
||||
|
||||
var sess models.SDKSession
|
||||
err := s.store.QueryRowContext(ctx, query, id).Scan(
|
||||
&sess.ID, &sess.ClaudeSessionID, &sess.SDKSessionID, &sess.Project, &sess.UserPrompt,
|
||||
&sess.WorkerPort, &sess.PromptCounter, &sess.Status, &sess.StartedAt, &sess.StartedAtEpoch,
|
||||
&sess.CompletedAt, &sess.CompletedAtEpoch,
|
||||
)
|
||||
if err == sql.ErrNoRows {
|
||||
return nil, nil
|
||||
}
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &sess, nil
|
||||
}
|
||||
|
||||
// FindAnySDKSession finds any session by Claude session ID (any status).
|
||||
func (s *SessionStore) FindAnySDKSession(ctx context.Context, claudeSessionID string) (*models.SDKSession, error) {
|
||||
const query = `
|
||||
SELECT id, claude_session_id, sdk_session_id, project, user_prompt,
|
||||
worker_port, prompt_counter, status, started_at, started_at_epoch,
|
||||
completed_at, completed_at_epoch
|
||||
FROM sdk_sessions
|
||||
WHERE claude_session_id = ?
|
||||
LIMIT 1
|
||||
`
|
||||
|
||||
var sess models.SDKSession
|
||||
err := s.store.QueryRowContext(ctx, query, claudeSessionID).Scan(
|
||||
&sess.ID, &sess.ClaudeSessionID, &sess.SDKSessionID, &sess.Project, &sess.UserPrompt,
|
||||
&sess.WorkerPort, &sess.PromptCounter, &sess.Status, &sess.StartedAt, &sess.StartedAtEpoch,
|
||||
&sess.CompletedAt, &sess.CompletedAtEpoch,
|
||||
)
|
||||
if err == sql.ErrNoRows {
|
||||
return nil, nil
|
||||
}
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &sess, nil
|
||||
}
|
||||
|
||||
// IncrementPromptCounter increments the prompt counter and returns the new value.
|
||||
func (s *SessionStore) IncrementPromptCounter(ctx context.Context, id int64) (int, error) {
|
||||
const updateQuery = `
|
||||
UPDATE sdk_sessions
|
||||
SET prompt_counter = COALESCE(prompt_counter, 0) + 1
|
||||
WHERE id = ?
|
||||
`
|
||||
if _, err := s.store.ExecContext(ctx, updateQuery, id); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
|
||||
const selectQuery = `SELECT prompt_counter FROM sdk_sessions WHERE id = ?`
|
||||
var counter int
|
||||
err := s.store.QueryRowContext(ctx, selectQuery, id).Scan(&counter)
|
||||
return counter, err
|
||||
}
|
||||
|
||||
// GetPromptCounter returns the current prompt counter for a session.
|
||||
func (s *SessionStore) GetPromptCounter(ctx context.Context, id int64) (int, error) {
|
||||
const query = `SELECT COALESCE(prompt_counter, 0) FROM sdk_sessions WHERE id = ?`
|
||||
var counter int
|
||||
err := s.store.QueryRowContext(ctx, query, id).Scan(&counter)
|
||||
return counter, err
|
||||
}
|
||||
|
||||
// GetSessionsToday returns the count of sessions started today.
|
||||
func (s *SessionStore) GetSessionsToday(ctx context.Context) (int, error) {
|
||||
// Get start of today in milliseconds
|
||||
now := time.Now()
|
||||
startOfDay := time.Date(now.Year(), now.Month(), now.Day(), 0, 0, 0, 0, now.Location())
|
||||
startEpoch := startOfDay.UnixMilli()
|
||||
|
||||
const query = `SELECT COUNT(*) FROM sdk_sessions WHERE started_at_epoch >= ?`
|
||||
|
||||
var count int
|
||||
err := s.store.QueryRowContext(ctx, query, startEpoch).Scan(&count)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return count, nil
|
||||
}
|
||||
|
||||
// GetAllProjects returns all unique project names.
|
||||
func (s *SessionStore) GetAllProjects(ctx context.Context) ([]string, error) {
|
||||
const query = `
|
||||
SELECT DISTINCT project
|
||||
FROM sdk_sessions
|
||||
WHERE project IS NOT NULL AND project != ''
|
||||
ORDER BY project ASC
|
||||
`
|
||||
|
||||
rows, err := s.store.QueryContext(ctx, query)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
var projects []string
|
||||
for rows.Next() {
|
||||
var project string
|
||||
if err := rows.Scan(&project); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
projects = append(projects, project)
|
||||
}
|
||||
return projects, rows.Err()
|
||||
}
|
||||
@@ -0,0 +1,449 @@
|
||||
package sqlite
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/lukaszraczylo/claude-mnemonic/pkg/models"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"github.com/stretchr/testify/suite"
|
||||
)
|
||||
|
||||
func testSessionStore(t *testing.T) (*SessionStore, *Store, func()) {
|
||||
t.Helper()
|
||||
|
||||
db, _, cleanup := testDB(t)
|
||||
createBaseTables(t, db) // Use base tables without FTS5 for session tests
|
||||
|
||||
store := newStoreFromDB(db)
|
||||
sessionStore := NewSessionStore(store)
|
||||
|
||||
return sessionStore, store, cleanup
|
||||
}
|
||||
|
||||
// SessionStoreSuite is a test suite for SessionStore operations.
|
||||
type SessionStoreSuite struct {
|
||||
suite.Suite
|
||||
sessionStore *SessionStore
|
||||
store *Store
|
||||
cleanup func()
|
||||
}
|
||||
|
||||
func (s *SessionStoreSuite) SetupTest() {
|
||||
s.sessionStore, s.store, s.cleanup = testSessionStore(s.T())
|
||||
}
|
||||
|
||||
func (s *SessionStoreSuite) TearDownTest() {
|
||||
if s.cleanup != nil {
|
||||
s.cleanup()
|
||||
}
|
||||
}
|
||||
|
||||
func TestSessionStoreSuite(t *testing.T) {
|
||||
suite.Run(t, new(SessionStoreSuite))
|
||||
}
|
||||
|
||||
// TestCreateSDKSession_TableDriven tests session creation with various scenarios.
|
||||
func (s *SessionStoreSuite) TestCreateSDKSession_TableDriven() {
|
||||
ctx := context.Background()
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
claudeSessionID string
|
||||
project string
|
||||
userPrompt string
|
||||
wantErr bool
|
||||
wantID bool
|
||||
}{
|
||||
{
|
||||
name: "basic session creation",
|
||||
claudeSessionID: "claude-basic",
|
||||
project: "project-a",
|
||||
userPrompt: "hello world",
|
||||
wantErr: false,
|
||||
wantID: true,
|
||||
},
|
||||
{
|
||||
name: "empty user prompt",
|
||||
claudeSessionID: "claude-noprompt",
|
||||
project: "project-b",
|
||||
userPrompt: "",
|
||||
wantErr: false,
|
||||
wantID: true,
|
||||
},
|
||||
{
|
||||
name: "long project name",
|
||||
claudeSessionID: "claude-longproj",
|
||||
project: "/Users/test/Documents/very/long/path/to/some/project/directory",
|
||||
userPrompt: "test",
|
||||
wantErr: false,
|
||||
wantID: true,
|
||||
},
|
||||
{
|
||||
name: "unicode project name",
|
||||
claudeSessionID: "claude-unicode",
|
||||
project: "项目名称-プロジェクト",
|
||||
userPrompt: "测试 テスト",
|
||||
wantErr: false,
|
||||
wantID: true,
|
||||
},
|
||||
{
|
||||
name: "special characters in prompt",
|
||||
claudeSessionID: "claude-special",
|
||||
project: "project-special",
|
||||
userPrompt: "Fix the bug in file.go:123 with \"quotes\" and 'apostrophes'",
|
||||
wantErr: false,
|
||||
wantID: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
s.Run(tt.name, func() {
|
||||
id, err := s.sessionStore.CreateSDKSession(ctx, tt.claudeSessionID, tt.project, tt.userPrompt)
|
||||
if tt.wantErr {
|
||||
s.Error(err)
|
||||
} else {
|
||||
s.NoError(err)
|
||||
if tt.wantID {
|
||||
s.Greater(id, int64(0))
|
||||
}
|
||||
|
||||
// Verify created session
|
||||
sess, err := s.sessionStore.GetSessionByID(ctx, id)
|
||||
s.NoError(err)
|
||||
s.NotNil(sess)
|
||||
s.Equal(tt.claudeSessionID, sess.ClaudeSessionID)
|
||||
s.Equal(tt.project, sess.Project)
|
||||
s.Equal(models.SessionStatusActive, sess.Status)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestIdempotentSession tests that session creation is idempotent.
|
||||
func (s *SessionStoreSuite) TestIdempotentSession() {
|
||||
ctx := context.Background()
|
||||
|
||||
// Create initial session
|
||||
id1, err := s.sessionStore.CreateSDKSession(ctx, "claude-idem", "project-1", "prompt-1")
|
||||
s.NoError(err)
|
||||
s.Greater(id1, int64(0))
|
||||
|
||||
// Create with same claude_session_id - should return same ID
|
||||
id2, err := s.sessionStore.CreateSDKSession(ctx, "claude-idem", "project-2", "prompt-2")
|
||||
s.NoError(err)
|
||||
s.Equal(id1, id2)
|
||||
|
||||
// Verify project was updated
|
||||
sess, err := s.sessionStore.GetSessionByID(ctx, id1)
|
||||
s.NoError(err)
|
||||
s.Equal("project-2", sess.Project)
|
||||
}
|
||||
|
||||
// TestPromptCounterOperations tests prompt counter increment and retrieval.
|
||||
func (s *SessionStoreSuite) TestPromptCounterOperations() {
|
||||
ctx := context.Background()
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
increments int
|
||||
expectedCount int
|
||||
}{
|
||||
{
|
||||
name: "no increments",
|
||||
increments: 0,
|
||||
expectedCount: 0,
|
||||
},
|
||||
{
|
||||
name: "single increment",
|
||||
increments: 1,
|
||||
expectedCount: 1,
|
||||
},
|
||||
{
|
||||
name: "multiple increments",
|
||||
increments: 5,
|
||||
expectedCount: 5,
|
||||
},
|
||||
{
|
||||
name: "many increments",
|
||||
increments: 100,
|
||||
expectedCount: 100,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
s.Run(tt.name, func() {
|
||||
// Create fresh session for each test
|
||||
id, err := s.sessionStore.CreateSDKSession(ctx, "claude-counter-"+tt.name, "project", "")
|
||||
s.NoError(err)
|
||||
|
||||
// Increment specified number of times
|
||||
var lastCount int
|
||||
for i := 0; i < tt.increments; i++ {
|
||||
lastCount, err = s.sessionStore.IncrementPromptCounter(ctx, id)
|
||||
s.NoError(err)
|
||||
}
|
||||
|
||||
// Get final count
|
||||
finalCount, err := s.sessionStore.GetPromptCounter(ctx, id)
|
||||
s.NoError(err)
|
||||
s.Equal(tt.expectedCount, finalCount)
|
||||
|
||||
if tt.increments > 0 {
|
||||
s.Equal(tt.expectedCount, lastCount)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestFindAnySDKSession tests session lookup scenarios.
|
||||
func (s *SessionStoreSuite) TestFindAnySDKSession_Scenarios() {
|
||||
ctx := context.Background()
|
||||
|
||||
// Create test sessions
|
||||
_, err := s.sessionStore.CreateSDKSession(ctx, "session-find-1", "project-a", "")
|
||||
s.NoError(err)
|
||||
_, err = s.sessionStore.CreateSDKSession(ctx, "session-find-2", "project-b", "")
|
||||
s.NoError(err)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
claudeSessionID string
|
||||
wantFound bool
|
||||
wantProject string
|
||||
}{
|
||||
{
|
||||
name: "find existing session 1",
|
||||
claudeSessionID: "session-find-1",
|
||||
wantFound: true,
|
||||
wantProject: "project-a",
|
||||
},
|
||||
{
|
||||
name: "find existing session 2",
|
||||
claudeSessionID: "session-find-2",
|
||||
wantFound: true,
|
||||
wantProject: "project-b",
|
||||
},
|
||||
{
|
||||
name: "find non-existent session",
|
||||
claudeSessionID: "session-nonexistent",
|
||||
wantFound: false,
|
||||
},
|
||||
{
|
||||
name: "find with empty string",
|
||||
claudeSessionID: "",
|
||||
wantFound: false,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
s.Run(tt.name, func() {
|
||||
sess, err := s.sessionStore.FindAnySDKSession(ctx, tt.claudeSessionID)
|
||||
s.NoError(err) // FindAnySDKSession returns nil,nil for not found
|
||||
|
||||
if tt.wantFound {
|
||||
s.NotNil(sess)
|
||||
s.Equal(tt.wantProject, sess.Project)
|
||||
} else {
|
||||
s.Nil(sess)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSessionStore_CreateSDKSession(t *testing.T) {
|
||||
sessionStore, _, cleanup := testSessionStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Create a new session
|
||||
id, err := sessionStore.CreateSDKSession(ctx, "claude-1", "test-project", "initial prompt")
|
||||
require.NoError(t, err)
|
||||
assert.Greater(t, id, int64(0))
|
||||
|
||||
// Retrieve and verify
|
||||
sess, err := sessionStore.GetSessionByID(ctx, id)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, sess)
|
||||
assert.Equal(t, "claude-1", sess.ClaudeSessionID)
|
||||
assert.Equal(t, "test-project", sess.Project)
|
||||
assert.Equal(t, models.SessionStatusActive, sess.Status)
|
||||
}
|
||||
|
||||
func TestSessionStore_CreateSDKSession_Idempotent(t *testing.T) {
|
||||
sessionStore, _, cleanup := testSessionStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Create first session
|
||||
id1, err := sessionStore.CreateSDKSession(ctx, "claude-1", "project-a", "prompt 1")
|
||||
require.NoError(t, err)
|
||||
|
||||
// Create again with same claude_session_id but different project
|
||||
id2, err := sessionStore.CreateSDKSession(ctx, "claude-1", "project-b", "prompt 2")
|
||||
require.NoError(t, err)
|
||||
|
||||
// Should return same ID (idempotent)
|
||||
assert.Equal(t, id1, id2)
|
||||
|
||||
// Should have updated project to project-b
|
||||
sess, err := sessionStore.GetSessionByID(ctx, id1)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "project-b", sess.Project)
|
||||
}
|
||||
|
||||
func TestSessionStore_FindAnySDKSession(t *testing.T) {
|
||||
sessionStore, _, cleanup := testSessionStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Create a session
|
||||
_, err := sessionStore.CreateSDKSession(ctx, "claude-1", "test-project", "")
|
||||
require.NoError(t, err)
|
||||
|
||||
// Find it
|
||||
sess, err := sessionStore.FindAnySDKSession(ctx, "claude-1")
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, sess)
|
||||
assert.Equal(t, "claude-1", sess.ClaudeSessionID)
|
||||
|
||||
// Try to find non-existent
|
||||
sess, err = sessionStore.FindAnySDKSession(ctx, "claude-nonexistent")
|
||||
require.NoError(t, err)
|
||||
assert.Nil(t, sess)
|
||||
}
|
||||
|
||||
func TestSessionStore_IncrementPromptCounter(t *testing.T) {
|
||||
sessionStore, _, cleanup := testSessionStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Create a session
|
||||
id, err := sessionStore.CreateSDKSession(ctx, "claude-1", "test-project", "")
|
||||
require.NoError(t, err)
|
||||
|
||||
// Initial counter should be 0
|
||||
counter, err := sessionStore.GetPromptCounter(ctx, id)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 0, counter)
|
||||
|
||||
// Increment
|
||||
counter, err = sessionStore.IncrementPromptCounter(ctx, id)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 1, counter)
|
||||
|
||||
// Increment again
|
||||
counter, err = sessionStore.IncrementPromptCounter(ctx, id)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 2, counter)
|
||||
|
||||
// Verify via GetPromptCounter
|
||||
counter, err = sessionStore.GetPromptCounter(ctx, id)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 2, counter)
|
||||
}
|
||||
|
||||
func TestSessionStore_GetSessionsToday(t *testing.T) {
|
||||
sessionStore, _, cleanup := testSessionStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Initially no sessions today
|
||||
count, err := sessionStore.GetSessionsToday(ctx)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 0, count)
|
||||
|
||||
// Create some sessions
|
||||
_, err = sessionStore.CreateSDKSession(ctx, "claude-1", "project-a", "")
|
||||
require.NoError(t, err)
|
||||
_, err = sessionStore.CreateSDKSession(ctx, "claude-2", "project-b", "")
|
||||
require.NoError(t, err)
|
||||
_, err = sessionStore.CreateSDKSession(ctx, "claude-3", "project-c", "")
|
||||
require.NoError(t, err)
|
||||
|
||||
// Should have 3 sessions today
|
||||
count, err = sessionStore.GetSessionsToday(ctx)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 3, count)
|
||||
}
|
||||
|
||||
func TestSessionStore_GetAllProjects(t *testing.T) {
|
||||
sessionStore, _, cleanup := testSessionStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Create sessions for different projects
|
||||
_, err := sessionStore.CreateSDKSession(ctx, "claude-1", "alpha-project", "")
|
||||
require.NoError(t, err)
|
||||
_, err = sessionStore.CreateSDKSession(ctx, "claude-2", "beta-project", "")
|
||||
require.NoError(t, err)
|
||||
_, err = sessionStore.CreateSDKSession(ctx, "claude-3", "alpha-project", "") // duplicate
|
||||
require.NoError(t, err)
|
||||
_, err = sessionStore.CreateSDKSession(ctx, "claude-4", "gamma-project", "")
|
||||
require.NoError(t, err)
|
||||
|
||||
// Get all projects
|
||||
projects, err := sessionStore.GetAllProjects(ctx)
|
||||
require.NoError(t, err)
|
||||
assert.Len(t, projects, 3)
|
||||
assert.Contains(t, projects, "alpha-project")
|
||||
assert.Contains(t, projects, "beta-project")
|
||||
assert.Contains(t, projects, "gamma-project")
|
||||
|
||||
// Should be sorted alphabetically
|
||||
assert.Equal(t, "alpha-project", projects[0])
|
||||
}
|
||||
|
||||
func TestSessionStore_GetSessionByID_NotFound(t *testing.T) {
|
||||
sessionStore, _, cleanup := testSessionStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Non-existent ID should return nil, nil (not an error)
|
||||
sess, err := sessionStore.GetSessionByID(ctx, 999)
|
||||
require.NoError(t, err)
|
||||
assert.Nil(t, sess)
|
||||
}
|
||||
|
||||
func TestSessionStore_SessionFields(t *testing.T) {
|
||||
sessionStore, store, cleanup := testSessionStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Create a session with full details
|
||||
id, err := sessionStore.CreateSDKSession(ctx, "claude-full", "full-project", "full user prompt")
|
||||
require.NoError(t, err)
|
||||
|
||||
// Manually update additional fields for testing
|
||||
now := time.Now()
|
||||
_, err = storeDB(store).Exec(`
|
||||
UPDATE sdk_sessions
|
||||
SET worker_port = ?, completed_at = ?, completed_at_epoch = ?, status = 'completed'
|
||||
WHERE id = ?
|
||||
`, 37777, now.Format(time.RFC3339), now.UnixMilli(), id)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Retrieve and verify all fields
|
||||
sess, err := sessionStore.GetSessionByID(ctx, id)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, sess)
|
||||
|
||||
assert.Equal(t, id, sess.ID)
|
||||
assert.Equal(t, "claude-full", sess.ClaudeSessionID)
|
||||
assert.Equal(t, "full-project", sess.Project)
|
||||
assert.Equal(t, models.SessionStatusCompleted, sess.Status)
|
||||
assert.True(t, sess.WorkerPort.Valid)
|
||||
assert.Equal(t, int64(37777), sess.WorkerPort.Int64)
|
||||
assert.True(t, sess.CompletedAt.Valid)
|
||||
assert.True(t, sess.CompletedAtEpoch.Valid)
|
||||
}
|
||||
@@ -0,0 +1,149 @@
|
||||
// Package sqlite provides SQLite database operations for claude-mnemonic.
|
||||
package sqlite
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"fmt"
|
||||
"sync"
|
||||
|
||||
sqlite_vec "github.com/asg017/sqlite-vec-go-bindings/cgo"
|
||||
_ "github.com/mattn/go-sqlite3"
|
||||
)
|
||||
|
||||
// Store provides database operations with connection pooling and prepared statements.
|
||||
type Store struct {
|
||||
db *sql.DB
|
||||
stmtCache map[string]*sql.Stmt
|
||||
stmtMu sync.RWMutex
|
||||
}
|
||||
|
||||
// StoreConfig holds configuration for the database store.
|
||||
type StoreConfig struct {
|
||||
Path string
|
||||
MaxConns int
|
||||
WALMode bool
|
||||
}
|
||||
|
||||
// NewStore creates a new database store with the given configuration.
|
||||
func NewStore(cfg StoreConfig) (*Store, error) {
|
||||
// Register sqlite-vec extension for vector operations
|
||||
sqlite_vec.Auto()
|
||||
|
||||
// Build connection string with pragmas
|
||||
connStr := cfg.Path + "?_journal_mode=WAL&_synchronous=NORMAL&_foreign_keys=ON"
|
||||
|
||||
db, err := sql.Open("sqlite3", connStr)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("open database: %w", err)
|
||||
}
|
||||
|
||||
// Configure connection pool
|
||||
maxConns := cfg.MaxConns
|
||||
if maxConns <= 0 {
|
||||
maxConns = 4
|
||||
}
|
||||
db.SetMaxOpenConns(maxConns)
|
||||
db.SetMaxIdleConns(maxConns)
|
||||
db.SetConnMaxLifetime(0) // Never expire - SQLite connections are cheap
|
||||
|
||||
// Verify connection
|
||||
if err := db.Ping(); err != nil {
|
||||
_ = db.Close()
|
||||
return nil, fmt.Errorf("ping database: %w", err)
|
||||
}
|
||||
|
||||
store := &Store{
|
||||
db: db,
|
||||
stmtCache: make(map[string]*sql.Stmt),
|
||||
}
|
||||
|
||||
// Run migrations
|
||||
mgr := NewMigrationManager(db)
|
||||
if err := mgr.RunMigrations(); err != nil {
|
||||
_ = db.Close()
|
||||
return nil, fmt.Errorf("run migrations: %w", err)
|
||||
}
|
||||
|
||||
return store, nil
|
||||
}
|
||||
|
||||
// Close closes the database connection and all cached statements.
|
||||
func (s *Store) Close() error {
|
||||
s.stmtMu.Lock()
|
||||
defer s.stmtMu.Unlock()
|
||||
|
||||
for _, stmt := range s.stmtCache {
|
||||
_ = stmt.Close()
|
||||
}
|
||||
s.stmtCache = nil
|
||||
|
||||
return s.db.Close()
|
||||
}
|
||||
|
||||
// GetStmt returns a cached prepared statement, creating it if necessary.
|
||||
func (s *Store) GetStmt(query string) (*sql.Stmt, error) {
|
||||
s.stmtMu.RLock()
|
||||
stmt, ok := s.stmtCache[query]
|
||||
s.stmtMu.RUnlock()
|
||||
if ok {
|
||||
return stmt, nil
|
||||
}
|
||||
|
||||
s.stmtMu.Lock()
|
||||
defer s.stmtMu.Unlock()
|
||||
|
||||
// Double-check after acquiring write lock
|
||||
if stmt, ok := s.stmtCache[query]; ok {
|
||||
return stmt, nil
|
||||
}
|
||||
|
||||
stmt, err := s.db.Prepare(query)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
s.stmtCache[query] = stmt
|
||||
return stmt, nil
|
||||
}
|
||||
|
||||
// ExecContext executes a query that doesn't return rows.
|
||||
func (s *Store) ExecContext(ctx context.Context, query string, args ...interface{}) (sql.Result, error) {
|
||||
stmt, err := s.GetStmt(query)
|
||||
if err != nil {
|
||||
// Fall back to direct execution
|
||||
return s.db.ExecContext(ctx, query, args...)
|
||||
}
|
||||
return stmt.ExecContext(ctx, args...)
|
||||
}
|
||||
|
||||
// QueryContext executes a query that returns rows.
|
||||
func (s *Store) QueryContext(ctx context.Context, query string, args ...interface{}) (*sql.Rows, error) {
|
||||
stmt, err := s.GetStmt(query)
|
||||
if err != nil {
|
||||
// Fall back to direct execution
|
||||
return s.db.QueryContext(ctx, query, args...)
|
||||
}
|
||||
return stmt.QueryContext(ctx, args...)
|
||||
}
|
||||
|
||||
// QueryRowContext executes a query that returns a single row.
|
||||
func (s *Store) QueryRowContext(ctx context.Context, query string, args ...interface{}) *sql.Row {
|
||||
stmt, err := s.GetStmt(query)
|
||||
if err != nil {
|
||||
// Fall back to direct execution
|
||||
return s.db.QueryRowContext(ctx, query, args...)
|
||||
}
|
||||
return stmt.QueryRowContext(ctx, args...)
|
||||
}
|
||||
|
||||
// Ping checks if the database connection is alive.
|
||||
func (s *Store) Ping() error {
|
||||
return s.db.Ping()
|
||||
}
|
||||
|
||||
// DB returns the underlying database connection for direct access.
|
||||
// Use this sparingly - prefer the store methods for most operations.
|
||||
func (s *Store) DB() *sql.DB {
|
||||
return s.db
|
||||
}
|
||||
@@ -0,0 +1,529 @@
|
||||
// Package sqlite provides SQLite database operations for claude-mnemonic.
|
||||
package sqlite
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"github.com/stretchr/testify/suite"
|
||||
)
|
||||
|
||||
// StoreSuite is a test suite for Store operations.
|
||||
type StoreSuite struct {
|
||||
suite.Suite
|
||||
db *sql.DB
|
||||
store *Store
|
||||
cleanup func()
|
||||
}
|
||||
|
||||
// SetupTest creates a fresh database before each test.
|
||||
func (s *StoreSuite) SetupTest() {
|
||||
s.db, _, s.cleanup = testDB(s.T())
|
||||
createBaseTables(s.T(), s.db)
|
||||
s.store = newStoreFromDB(s.db)
|
||||
}
|
||||
|
||||
// TearDownTest cleans up after each test.
|
||||
func (s *StoreSuite) TearDownTest() {
|
||||
if s.cleanup != nil {
|
||||
s.cleanup()
|
||||
}
|
||||
}
|
||||
|
||||
func TestStoreSuite(t *testing.T) {
|
||||
suite.Run(t, new(StoreSuite))
|
||||
}
|
||||
|
||||
// TestGetStmt tests prepared statement caching.
|
||||
func (s *StoreSuite) TestGetStmt() {
|
||||
tests := []struct {
|
||||
name string
|
||||
query string
|
||||
wantErr bool
|
||||
}{
|
||||
{
|
||||
name: "valid simple query",
|
||||
query: "SELECT 1",
|
||||
wantErr: false,
|
||||
},
|
||||
{
|
||||
name: "valid query with parameter",
|
||||
query: "SELECT * FROM sdk_sessions WHERE id = ?",
|
||||
wantErr: false,
|
||||
},
|
||||
{
|
||||
name: "invalid query syntax",
|
||||
query: "SELECT * FROM nonexistent_table WHERE",
|
||||
wantErr: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
s.Run(tt.name, func() {
|
||||
stmt, err := s.store.GetStmt(tt.query)
|
||||
if tt.wantErr {
|
||||
s.Error(err)
|
||||
s.Nil(stmt)
|
||||
} else {
|
||||
s.NoError(err)
|
||||
s.NotNil(stmt)
|
||||
|
||||
// Second call should return cached statement
|
||||
stmt2, err := s.store.GetStmt(tt.query)
|
||||
s.NoError(err)
|
||||
s.Same(stmt, stmt2)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestExecContext tests query execution.
|
||||
func (s *StoreSuite) TestExecContext() {
|
||||
ctx := context.Background()
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
query string
|
||||
args []interface{}
|
||||
wantErr bool
|
||||
wantAffected int64
|
||||
}{
|
||||
{
|
||||
name: "insert session",
|
||||
query: `INSERT INTO sdk_sessions (claude_session_id, sdk_session_id, project, started_at, started_at_epoch, status)
|
||||
VALUES (?, ?, ?, datetime('now'), strftime('%s', 'now') * 1000, 'active')`,
|
||||
args: []interface{}{"claude-1", "sdk-1", "test-project"},
|
||||
wantErr: false,
|
||||
wantAffected: 1,
|
||||
},
|
||||
{
|
||||
name: "invalid query",
|
||||
query: "INSERT INTO nonexistent_table VALUES (?)",
|
||||
args: []interface{}{"test"},
|
||||
wantErr: true,
|
||||
wantAffected: 0,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
s.Run(tt.name, func() {
|
||||
result, err := s.store.ExecContext(ctx, tt.query, tt.args...)
|
||||
if tt.wantErr {
|
||||
s.Error(err)
|
||||
} else {
|
||||
s.NoError(err)
|
||||
affected, _ := result.RowsAffected()
|
||||
s.Equal(tt.wantAffected, affected)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestQueryContext tests query execution that returns rows.
|
||||
func (s *StoreSuite) TestQueryContext() {
|
||||
ctx := context.Background()
|
||||
|
||||
// Seed data
|
||||
seedSession(s.T(), s.db, "claude-1", "sdk-1", "project-a")
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
query string
|
||||
args []interface{}
|
||||
wantErr bool
|
||||
wantRows int
|
||||
setupFunc func()
|
||||
assertFunc func(rows *sql.Rows)
|
||||
}{
|
||||
{
|
||||
name: "query existing session",
|
||||
query: "SELECT id, project FROM sdk_sessions WHERE claude_session_id = ?",
|
||||
args: []interface{}{"claude-1"},
|
||||
wantErr: false,
|
||||
wantRows: 1,
|
||||
},
|
||||
{
|
||||
name: "query non-existent session",
|
||||
query: "SELECT id, project FROM sdk_sessions WHERE claude_session_id = ?",
|
||||
args: []interface{}{"nonexistent"},
|
||||
wantErr: false,
|
||||
wantRows: 0,
|
||||
},
|
||||
{
|
||||
name: "query all sessions",
|
||||
query: "SELECT id, project FROM sdk_sessions",
|
||||
args: nil,
|
||||
wantErr: false,
|
||||
wantRows: 1,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
s.Run(tt.name, func() {
|
||||
rows, err := s.store.QueryContext(ctx, tt.query, tt.args...)
|
||||
if tt.wantErr {
|
||||
s.Error(err)
|
||||
return
|
||||
}
|
||||
|
||||
s.NoError(err)
|
||||
defer rows.Close()
|
||||
|
||||
count := 0
|
||||
for rows.Next() {
|
||||
count++
|
||||
}
|
||||
s.Equal(tt.wantRows, count)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestQueryRowContext tests single row query execution.
|
||||
func (s *StoreSuite) TestQueryRowContext() {
|
||||
ctx := context.Background()
|
||||
|
||||
// Seed data
|
||||
seedSession(s.T(), s.db, "claude-1", "sdk-1", "project-a")
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
query string
|
||||
args []interface{}
|
||||
wantErr bool
|
||||
}{
|
||||
{
|
||||
name: "query existing session",
|
||||
query: "SELECT id FROM sdk_sessions WHERE claude_session_id = ?",
|
||||
args: []interface{}{"claude-1"},
|
||||
wantErr: false,
|
||||
},
|
||||
{
|
||||
name: "query non-existent session",
|
||||
query: "SELECT id FROM sdk_sessions WHERE claude_session_id = ?",
|
||||
args: []interface{}{"nonexistent"},
|
||||
wantErr: true, // sql.ErrNoRows
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
s.Run(tt.name, func() {
|
||||
row := s.store.QueryRowContext(ctx, tt.query, tt.args...)
|
||||
var id int64
|
||||
err := row.Scan(&id)
|
||||
if tt.wantErr {
|
||||
s.Error(err)
|
||||
} else {
|
||||
s.NoError(err)
|
||||
s.Greater(id, int64(0))
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestPing tests database connection health check.
|
||||
func (s *StoreSuite) TestPing() {
|
||||
err := s.store.Ping()
|
||||
s.NoError(err)
|
||||
}
|
||||
|
||||
// TestDB tests getting the underlying database connection.
|
||||
func (s *StoreSuite) TestDB() {
|
||||
db := s.store.DB()
|
||||
s.NotNil(db)
|
||||
s.Same(s.db, db)
|
||||
}
|
||||
|
||||
// TestClose tests closing the store.
|
||||
func (s *StoreSuite) TestClose() {
|
||||
// Create a separate store for close test
|
||||
db, _, cleanup := testDB(s.T())
|
||||
defer cleanup()
|
||||
|
||||
store := newStoreFromDB(db)
|
||||
|
||||
// Cache a statement first
|
||||
_, err := store.GetStmt("SELECT 1")
|
||||
s.NoError(err)
|
||||
|
||||
// Close should not error
|
||||
err = store.Close()
|
||||
s.NoError(err)
|
||||
|
||||
// Operations after close should fail
|
||||
err = store.Ping()
|
||||
s.Error(err)
|
||||
}
|
||||
|
||||
// TestConcurrentStmtCache tests concurrent access to statement cache.
|
||||
func (s *StoreSuite) TestConcurrentStmtCache() {
|
||||
ctx := context.Background()
|
||||
queries := []string{
|
||||
"SELECT 1",
|
||||
"SELECT 2",
|
||||
"SELECT id FROM sdk_sessions",
|
||||
"SELECT project FROM sdk_sessions",
|
||||
}
|
||||
|
||||
done := make(chan struct{})
|
||||
for i := 0; i < 10; i++ {
|
||||
go func(i int) {
|
||||
query := queries[i%len(queries)]
|
||||
_, _ = s.store.GetStmt(query)
|
||||
_, _ = s.store.ExecContext(ctx, "SELECT 1")
|
||||
done <- struct{}{}
|
||||
}(i)
|
||||
}
|
||||
|
||||
for i := 0; i < 10; i++ {
|
||||
<-done
|
||||
}
|
||||
}
|
||||
|
||||
// HelpersSuite tests helper functions.
|
||||
type HelpersSuite struct {
|
||||
suite.Suite
|
||||
}
|
||||
|
||||
func TestHelpersSuite(t *testing.T) {
|
||||
suite.Run(t, new(HelpersSuite))
|
||||
}
|
||||
|
||||
func (s *HelpersSuite) TestNullString() {
|
||||
tests := []struct {
|
||||
name string
|
||||
input string
|
||||
wantStr string
|
||||
wantBool bool
|
||||
}{
|
||||
{
|
||||
name: "empty string",
|
||||
input: "",
|
||||
wantStr: "",
|
||||
wantBool: false,
|
||||
},
|
||||
{
|
||||
name: "non-empty string",
|
||||
input: "test",
|
||||
wantStr: "test",
|
||||
wantBool: true,
|
||||
},
|
||||
{
|
||||
name: "whitespace string",
|
||||
input: " ",
|
||||
wantStr: " ",
|
||||
wantBool: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
s.Run(tt.name, func() {
|
||||
result := nullString(tt.input)
|
||||
s.Equal(tt.wantStr, result.String)
|
||||
s.Equal(tt.wantBool, result.Valid)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func (s *HelpersSuite) TestNullInt() {
|
||||
tests := []struct {
|
||||
name string
|
||||
input int
|
||||
wantInt int64
|
||||
wantBool bool
|
||||
}{
|
||||
{
|
||||
name: "zero",
|
||||
input: 0,
|
||||
wantInt: 0,
|
||||
wantBool: false,
|
||||
},
|
||||
{
|
||||
name: "negative",
|
||||
input: -1,
|
||||
wantInt: -1,
|
||||
wantBool: false,
|
||||
},
|
||||
{
|
||||
name: "positive",
|
||||
input: 42,
|
||||
wantInt: 42,
|
||||
wantBool: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
s.Run(tt.name, func() {
|
||||
result := nullInt(tt.input)
|
||||
s.Equal(tt.wantInt, result.Int64)
|
||||
s.Equal(tt.wantBool, result.Valid)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func (s *HelpersSuite) TestRepeatPlaceholders() {
|
||||
tests := []struct {
|
||||
name string
|
||||
input int
|
||||
expected string
|
||||
}{
|
||||
{
|
||||
name: "zero",
|
||||
input: 0,
|
||||
expected: "",
|
||||
},
|
||||
{
|
||||
name: "negative",
|
||||
input: -1,
|
||||
expected: "",
|
||||
},
|
||||
{
|
||||
name: "one",
|
||||
input: 1,
|
||||
expected: ", ?",
|
||||
},
|
||||
{
|
||||
name: "three",
|
||||
input: 3,
|
||||
expected: ", ?, ?, ?",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
s.Run(tt.name, func() {
|
||||
result := repeatPlaceholders(tt.input)
|
||||
s.Equal(tt.expected, result)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func (s *HelpersSuite) TestInt64SliceToInterface() {
|
||||
tests := []struct {
|
||||
name string
|
||||
input []int64
|
||||
expected int
|
||||
}{
|
||||
{
|
||||
name: "empty slice",
|
||||
input: []int64{},
|
||||
expected: 0,
|
||||
},
|
||||
{
|
||||
name: "single element",
|
||||
input: []int64{42},
|
||||
expected: 1,
|
||||
},
|
||||
{
|
||||
name: "multiple elements",
|
||||
input: []int64{1, 2, 3, 4, 5},
|
||||
expected: 5,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
s.Run(tt.name, func() {
|
||||
result := int64SliceToInterface(tt.input)
|
||||
s.Len(result, tt.expected)
|
||||
for i, v := range result {
|
||||
s.Equal(tt.input[i], v)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestBuildGetByIDsQuery tests the shared query builder.
|
||||
func TestBuildGetByIDsQuery(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
baseQuery string
|
||||
ids []int64
|
||||
orderBy string
|
||||
limit int
|
||||
wantQuery string
|
||||
wantArgs int
|
||||
}{
|
||||
{
|
||||
name: "single id, no limit, desc order",
|
||||
baseQuery: "SELECT * FROM test",
|
||||
ids: []int64{1},
|
||||
orderBy: "date_desc",
|
||||
limit: 0,
|
||||
wantQuery: "SELECT * FROM test WHERE id IN (?)\n\t\tORDER BY created_at_epoch DESC",
|
||||
wantArgs: 1,
|
||||
},
|
||||
{
|
||||
name: "multiple ids with limit and asc order",
|
||||
baseQuery: "SELECT * FROM test",
|
||||
ids: []int64{1, 2, 3},
|
||||
orderBy: "date_asc",
|
||||
limit: 10,
|
||||
wantQuery: "SELECT * FROM test WHERE id IN (?, ?, ?)\n\t\tORDER BY created_at_epoch ASC LIMIT ?",
|
||||
wantArgs: 4,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
query, args := BuildGetByIDsQuery(tt.baseQuery, tt.ids, tt.orderBy, tt.limit)
|
||||
assert.Contains(t, query, "WHERE id IN")
|
||||
assert.Len(t, args, tt.wantArgs)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestEnsureSessionExists tests session auto-creation.
|
||||
func TestEnsureSessionExists(t *testing.T) {
|
||||
db, _, cleanup := testDB(t)
|
||||
defer cleanup()
|
||||
createBaseTables(t, db)
|
||||
|
||||
store := newStoreFromDB(db)
|
||||
ctx := context.Background()
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
sdkSessionID string
|
||||
project string
|
||||
setup func()
|
||||
wantErr bool
|
||||
}{
|
||||
{
|
||||
name: "create new session",
|
||||
sdkSessionID: "sdk-new",
|
||||
project: "project-a",
|
||||
wantErr: false,
|
||||
},
|
||||
{
|
||||
name: "session already exists",
|
||||
sdkSessionID: "sdk-existing",
|
||||
project: "project-b",
|
||||
setup: func() {
|
||||
seedSession(t, db, "sdk-existing", "sdk-existing", "project-b")
|
||||
},
|
||||
wantErr: false,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
if tt.setup != nil {
|
||||
tt.setup()
|
||||
}
|
||||
|
||||
err := EnsureSessionExists(ctx, store, tt.sdkSessionID, tt.project)
|
||||
if tt.wantErr {
|
||||
require.Error(t, err)
|
||||
} else {
|
||||
require.NoError(t, err)
|
||||
|
||||
// Verify session exists
|
||||
var id int64
|
||||
err := db.QueryRow("SELECT id FROM sdk_sessions WHERE sdk_session_id = ?", tt.sdkSessionID).Scan(&id)
|
||||
require.NoError(t, err)
|
||||
assert.Greater(t, id, int64(0))
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,118 @@
|
||||
// Package sqlite provides SQLite database operations for claude-mnemonic.
|
||||
package sqlite
|
||||
|
||||
import (
|
||||
"context"
|
||||
"time"
|
||||
|
||||
"github.com/lukaszraczylo/claude-mnemonic/pkg/models"
|
||||
)
|
||||
|
||||
// SummaryStore provides summary-related database operations.
|
||||
type SummaryStore struct {
|
||||
store *Store
|
||||
}
|
||||
|
||||
// NewSummaryStore creates a new summary store.
|
||||
func NewSummaryStore(store *Store) *SummaryStore {
|
||||
return &SummaryStore{store: store}
|
||||
}
|
||||
|
||||
// StoreSummary stores a new session summary.
|
||||
func (s *SummaryStore) StoreSummary(ctx context.Context, sdkSessionID, project string, summary *models.ParsedSummary, promptNumber int, discoveryTokens int64) (int64, int64, error) {
|
||||
now := time.Now()
|
||||
nowEpoch := now.UnixMilli()
|
||||
|
||||
// Ensure session exists (auto-create if missing)
|
||||
if err := s.ensureSessionExists(ctx, sdkSessionID, project); err != nil {
|
||||
return 0, 0, err
|
||||
}
|
||||
|
||||
const query = `
|
||||
INSERT INTO session_summaries
|
||||
(sdk_session_id, project, request, investigated, learned, completed,
|
||||
next_steps, notes, prompt_number, discovery_tokens, created_at, created_at_epoch)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
`
|
||||
|
||||
result, err := s.store.ExecContext(ctx, query,
|
||||
sdkSessionID, project,
|
||||
nullString(summary.Request), nullString(summary.Investigated),
|
||||
nullString(summary.Learned), nullString(summary.Completed),
|
||||
nullString(summary.NextSteps), nullString(summary.Notes),
|
||||
nullInt(promptNumber), discoveryTokens,
|
||||
now.Format(time.RFC3339), nowEpoch,
|
||||
)
|
||||
if err != nil {
|
||||
return 0, 0, err
|
||||
}
|
||||
|
||||
id, _ := result.LastInsertId()
|
||||
return id, nowEpoch, nil
|
||||
}
|
||||
|
||||
// ensureSessionExists creates a session if it doesn't exist.
|
||||
func (s *SummaryStore) ensureSessionExists(ctx context.Context, sdkSessionID, project string) error {
|
||||
return EnsureSessionExists(ctx, s.store, sdkSessionID, project)
|
||||
}
|
||||
|
||||
// GetSummariesByIDs retrieves summaries by a list of IDs.
|
||||
func (s *SummaryStore) GetSummariesByIDs(ctx context.Context, ids []int64, orderBy string, limit int) ([]*models.SessionSummary, error) {
|
||||
if len(ids) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
const baseQuery = `
|
||||
SELECT id, sdk_session_id, project, request, investigated, learned, completed,
|
||||
next_steps, notes, prompt_number, discovery_tokens, created_at, created_at_epoch
|
||||
FROM session_summaries`
|
||||
|
||||
query, args := BuildGetByIDsQuery(baseQuery, ids, orderBy, limit)
|
||||
|
||||
rows, err := s.store.db.QueryContext(ctx, query, args...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
return scanSummaryRows(rows)
|
||||
}
|
||||
|
||||
// GetRecentSummaries retrieves recent summaries for a project.
|
||||
func (s *SummaryStore) GetRecentSummaries(ctx context.Context, project string, limit int) ([]*models.SessionSummary, error) {
|
||||
const query = `
|
||||
SELECT id, sdk_session_id, project, request, investigated, learned, completed,
|
||||
next_steps, notes, prompt_number, discovery_tokens, created_at, created_at_epoch
|
||||
FROM session_summaries
|
||||
WHERE project = ?
|
||||
ORDER BY created_at_epoch DESC
|
||||
LIMIT ?
|
||||
`
|
||||
|
||||
rows, err := s.store.QueryContext(ctx, query, project, limit)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
return scanSummaryRows(rows)
|
||||
}
|
||||
|
||||
// GetAllRecentSummaries retrieves recent summaries across all projects.
|
||||
func (s *SummaryStore) GetAllRecentSummaries(ctx context.Context, limit int) ([]*models.SessionSummary, error) {
|
||||
const query = `
|
||||
SELECT id, sdk_session_id, project, request, investigated, learned, completed,
|
||||
next_steps, notes, prompt_number, discovery_tokens, created_at, created_at_epoch
|
||||
FROM session_summaries
|
||||
ORDER BY created_at_epoch DESC
|
||||
LIMIT ?
|
||||
`
|
||||
|
||||
rows, err := s.store.QueryContext(ctx, query, limit)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
return scanSummaryRows(rows)
|
||||
}
|
||||
@@ -0,0 +1,242 @@
|
||||
package sqlite
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/lukaszraczylo/claude-mnemonic/pkg/models"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func testSummaryStore(t *testing.T) (*SummaryStore, *Store, func()) {
|
||||
t.Helper()
|
||||
|
||||
db, _, cleanup := testDB(t)
|
||||
createAllTables(t, db)
|
||||
|
||||
store := newStoreFromDB(db)
|
||||
summaryStore := NewSummaryStore(store)
|
||||
|
||||
return summaryStore, store, cleanup
|
||||
}
|
||||
|
||||
func TestSummaryStore_StoreSummary(t *testing.T) {
|
||||
summaryStore, store, cleanup := testSummaryStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Create a session first
|
||||
seedSession(t, storeDB(store), "claude-1", "sdk-1", "test-project")
|
||||
|
||||
summary := &models.ParsedSummary{
|
||||
Request: "Add new feature",
|
||||
Investigated: "Looked at existing code",
|
||||
Learned: "Found the pattern to follow",
|
||||
Completed: "Implemented the feature",
|
||||
NextSteps: "Add tests",
|
||||
Notes: "Some additional notes",
|
||||
}
|
||||
|
||||
id, epoch, err := summaryStore.StoreSummary(ctx, "sdk-1", "test-project", summary, 1, 100)
|
||||
require.NoError(t, err)
|
||||
assert.Greater(t, id, int64(0))
|
||||
assert.Greater(t, epoch, int64(0))
|
||||
|
||||
// Verify it was saved
|
||||
var count int
|
||||
err = storeDB(store).QueryRow("SELECT COUNT(*) FROM session_summaries WHERE id = ?", id).Scan(&count)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 1, count)
|
||||
}
|
||||
|
||||
func TestSummaryStore_StoreSummary_AutoCreateSession(t *testing.T) {
|
||||
summaryStore, store, cleanup := testSummaryStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Don't create session beforehand - should be auto-created
|
||||
summary := &models.ParsedSummary{
|
||||
Request: "Test request",
|
||||
}
|
||||
|
||||
id, _, err := summaryStore.StoreSummary(ctx, "auto-session", "test-project", summary, 1, 0)
|
||||
require.NoError(t, err)
|
||||
assert.Greater(t, id, int64(0))
|
||||
|
||||
// Verify session was auto-created
|
||||
var sessionCount int
|
||||
err = storeDB(store).QueryRow("SELECT COUNT(*) FROM sdk_sessions WHERE sdk_session_id = ?", "auto-session").Scan(&sessionCount)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 1, sessionCount)
|
||||
}
|
||||
|
||||
func TestSummaryStore_GetRecentSummaries(t *testing.T) {
|
||||
summaryStore, store, cleanup := testSummaryStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Create a session
|
||||
seedSession(t, storeDB(store), "claude-1", "sdk-1", "test-project")
|
||||
|
||||
// Store multiple summaries
|
||||
for i := 0; i < 5; i++ {
|
||||
summary := &models.ParsedSummary{
|
||||
Request: "Request " + string(rune('A'+i)),
|
||||
}
|
||||
_, _, err := summaryStore.StoreSummary(ctx, "sdk-1", "test-project", summary, i+1, 0)
|
||||
require.NoError(t, err)
|
||||
time.Sleep(time.Millisecond) // Ensure different timestamps
|
||||
}
|
||||
|
||||
// Get recent summaries with limit
|
||||
summaries, err := summaryStore.GetRecentSummaries(ctx, "test-project", 3)
|
||||
require.NoError(t, err)
|
||||
assert.Len(t, summaries, 3)
|
||||
|
||||
// Should be in descending order
|
||||
assert.Equal(t, int64(5), summaries[0].PromptNumber.Int64)
|
||||
}
|
||||
|
||||
func TestSummaryStore_GetAllRecentSummaries(t *testing.T) {
|
||||
summaryStore, store, cleanup := testSummaryStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Create sessions for different projects
|
||||
seedSession(t, storeDB(store), "claude-1", "sdk-1", "project-a")
|
||||
seedSession(t, storeDB(store), "claude-2", "sdk-2", "project-b")
|
||||
|
||||
// Store summaries for both projects
|
||||
for i := 0; i < 3; i++ {
|
||||
summary := &models.ParsedSummary{Request: "Project A request"}
|
||||
_, _, err := summaryStore.StoreSummary(ctx, "sdk-1", "project-a", summary, i+1, 0)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
for i := 0; i < 2; i++ {
|
||||
summary := &models.ParsedSummary{Request: "Project B request"}
|
||||
_, _, err := summaryStore.StoreSummary(ctx, "sdk-2", "project-b", summary, i+1, 0)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
// Get all summaries (should include both projects)
|
||||
summaries, err := summaryStore.GetAllRecentSummaries(ctx, 10)
|
||||
require.NoError(t, err)
|
||||
assert.Len(t, summaries, 5)
|
||||
}
|
||||
|
||||
func TestSummaryStore_GetSummariesByIDs(t *testing.T) {
|
||||
summaryStore, store, cleanup := testSummaryStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Create a session
|
||||
seedSession(t, storeDB(store), "claude-1", "sdk-1", "test-project")
|
||||
|
||||
// Store summaries and collect IDs
|
||||
var ids []int64
|
||||
for i := 0; i < 5; i++ {
|
||||
summary := &models.ParsedSummary{Request: "Request " + string(rune('A'+i))}
|
||||
id, _, err := summaryStore.StoreSummary(ctx, "sdk-1", "test-project", summary, i+1, 0)
|
||||
require.NoError(t, err)
|
||||
ids = append(ids, id)
|
||||
time.Sleep(time.Millisecond)
|
||||
}
|
||||
|
||||
// Get specific summaries by ID
|
||||
summaries, err := summaryStore.GetSummariesByIDs(ctx, ids[:3], "date_desc", 10)
|
||||
require.NoError(t, err)
|
||||
assert.Len(t, summaries, 3)
|
||||
|
||||
// Test with ascending order
|
||||
summaries, err = summaryStore.GetSummariesByIDs(ctx, ids, "date_asc", 2)
|
||||
require.NoError(t, err)
|
||||
assert.Len(t, summaries, 2)
|
||||
assert.Equal(t, int64(1), summaries[0].PromptNumber.Int64)
|
||||
}
|
||||
|
||||
func TestSummaryStore_GetSummariesByIDs_EmptyInput(t *testing.T) {
|
||||
summaryStore, _, cleanup := testSummaryStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Empty IDs should return nil
|
||||
summaries, err := summaryStore.GetSummariesByIDs(ctx, []int64{}, "date_desc", 10)
|
||||
require.NoError(t, err)
|
||||
assert.Nil(t, summaries)
|
||||
}
|
||||
|
||||
func TestSummaryStore_SummaryFields(t *testing.T) {
|
||||
summaryStore, store, cleanup := testSummaryStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Create a session
|
||||
seedSession(t, storeDB(store), "claude-1", "sdk-1", "test-project")
|
||||
|
||||
// Store a summary with all fields
|
||||
summary := &models.ParsedSummary{
|
||||
Request: "Add authentication",
|
||||
Investigated: "Reviewed existing auth code",
|
||||
Learned: "OAuth is preferred",
|
||||
Completed: "Implemented OAuth flow",
|
||||
NextSteps: "Add refresh token support",
|
||||
Notes: "Consider rate limiting",
|
||||
}
|
||||
|
||||
id, _, err := summaryStore.StoreSummary(ctx, "sdk-1", "test-project", summary, 5, 1500)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Retrieve and verify all fields
|
||||
summaries, err := summaryStore.GetSummariesByIDs(ctx, []int64{id}, "date_desc", 1)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, summaries, 1)
|
||||
|
||||
s := summaries[0]
|
||||
assert.Equal(t, id, s.ID)
|
||||
assert.Equal(t, "sdk-1", s.SDKSessionID)
|
||||
assert.Equal(t, "test-project", s.Project)
|
||||
assert.Equal(t, "Add authentication", s.Request.String)
|
||||
assert.Equal(t, "Reviewed existing auth code", s.Investigated.String)
|
||||
assert.Equal(t, "OAuth is preferred", s.Learned.String)
|
||||
assert.Equal(t, "Implemented OAuth flow", s.Completed.String)
|
||||
assert.Equal(t, "Add refresh token support", s.NextSteps.String)
|
||||
assert.Equal(t, "Consider rate limiting", s.Notes.String)
|
||||
assert.Equal(t, int64(5), s.PromptNumber.Int64)
|
||||
assert.Equal(t, int64(1500), s.DiscoveryTokens)
|
||||
}
|
||||
|
||||
func TestSummaryStore_EmptySummary(t *testing.T) {
|
||||
summaryStore, store, cleanup := testSummaryStore(t)
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Create a session
|
||||
seedSession(t, storeDB(store), "claude-1", "sdk-1", "test-project")
|
||||
|
||||
// Store an empty summary
|
||||
summary := &models.ParsedSummary{}
|
||||
|
||||
id, _, err := summaryStore.StoreSummary(ctx, "sdk-1", "test-project", summary, 0, 0)
|
||||
require.NoError(t, err)
|
||||
assert.Greater(t, id, int64(0))
|
||||
|
||||
// Retrieve and verify null fields
|
||||
summaries, err := summaryStore.GetSummariesByIDs(ctx, []int64{id}, "date_desc", 1)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, summaries, 1)
|
||||
|
||||
s := summaries[0]
|
||||
assert.False(t, s.Request.Valid || s.Request.String != "")
|
||||
assert.False(t, s.Investigated.Valid || s.Investigated.String != "")
|
||||
assert.False(t, s.Learned.Valid || s.Learned.String != "")
|
||||
}
|
||||
@@ -0,0 +1,315 @@
|
||||
package sqlite
|
||||
|
||||
import (
|
||||
"database/sql"
|
||||
"os"
|
||||
"testing"
|
||||
|
||||
_ "github.com/mattn/go-sqlite3"
|
||||
)
|
||||
|
||||
// newStoreFromDB creates a Store from an existing database connection for testing.
|
||||
func newStoreFromDB(db *sql.DB) *Store {
|
||||
return &Store{
|
||||
db: db,
|
||||
stmtCache: make(map[string]*sql.Stmt),
|
||||
}
|
||||
}
|
||||
|
||||
// storeDB returns the underlying database connection from a store for testing.
|
||||
func storeDB(s *Store) *sql.DB {
|
||||
return s.db
|
||||
}
|
||||
|
||||
// testDB creates a temporary SQLite database for testing.
|
||||
// Returns the database, path, and a cleanup function.
|
||||
func testDB(t *testing.T) (*sql.DB, string, func()) {
|
||||
t.Helper()
|
||||
|
||||
tmpDir, err := os.MkdirTemp("", "claude-mnemonic-test-*")
|
||||
if err != nil {
|
||||
t.Fatalf("create temp dir: %v", err)
|
||||
}
|
||||
|
||||
dbPath := tmpDir + "/test.db"
|
||||
connStr := dbPath + "?_journal_mode=WAL&_synchronous=NORMAL&_foreign_keys=ON"
|
||||
|
||||
db, err := sql.Open("sqlite3", connStr)
|
||||
if err != nil {
|
||||
_ = os.RemoveAll(tmpDir)
|
||||
t.Fatalf("open database: %v", err)
|
||||
}
|
||||
|
||||
cleanup := func() {
|
||||
_ = db.Close()
|
||||
_ = os.RemoveAll(tmpDir)
|
||||
}
|
||||
|
||||
return db, dbPath, cleanup
|
||||
}
|
||||
|
||||
// createBaseTables creates the base tables without FTS5 for unit testing.
|
||||
func createBaseTables(t *testing.T, db *sql.DB) {
|
||||
t.Helper()
|
||||
|
||||
_, err := db.Exec(`
|
||||
CREATE TABLE IF NOT EXISTS schema_versions (
|
||||
id INTEGER PRIMARY KEY,
|
||||
version INTEGER UNIQUE NOT NULL,
|
||||
applied_at TEXT NOT NULL
|
||||
)
|
||||
`)
|
||||
if err != nil {
|
||||
t.Fatalf("create schema_versions: %v", err)
|
||||
}
|
||||
|
||||
_, err = db.Exec(`
|
||||
CREATE TABLE IF NOT EXISTS sdk_sessions (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
claude_session_id TEXT UNIQUE NOT NULL,
|
||||
sdk_session_id TEXT UNIQUE,
|
||||
project TEXT NOT NULL,
|
||||
user_prompt TEXT,
|
||||
started_at TEXT NOT NULL,
|
||||
started_at_epoch INTEGER NOT NULL,
|
||||
completed_at TEXT,
|
||||
completed_at_epoch INTEGER,
|
||||
status TEXT CHECK(status IN ('active', 'completed', 'failed')) NOT NULL DEFAULT 'active',
|
||||
worker_port INTEGER,
|
||||
prompt_counter INTEGER DEFAULT 0
|
||||
)
|
||||
`)
|
||||
if err != nil {
|
||||
t.Fatalf("create sdk_sessions: %v", err)
|
||||
}
|
||||
|
||||
_, err = db.Exec(`
|
||||
CREATE TABLE IF NOT EXISTS observations (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
sdk_session_id TEXT NOT NULL,
|
||||
project TEXT NOT NULL,
|
||||
text TEXT,
|
||||
type TEXT NOT NULL CHECK(type IN ('decision', 'bugfix', 'feature', 'refactor', 'discovery', 'change')),
|
||||
title TEXT,
|
||||
subtitle TEXT,
|
||||
facts TEXT,
|
||||
narrative TEXT,
|
||||
concepts TEXT,
|
||||
files_read TEXT,
|
||||
files_modified TEXT,
|
||||
file_mtimes TEXT,
|
||||
scope TEXT DEFAULT 'project' CHECK(scope IN ('project', 'global')),
|
||||
prompt_number INTEGER,
|
||||
discovery_tokens INTEGER DEFAULT 0,
|
||||
created_at TEXT NOT NULL,
|
||||
created_at_epoch INTEGER NOT NULL,
|
||||
FOREIGN KEY(sdk_session_id) REFERENCES sdk_sessions(sdk_session_id) ON DELETE CASCADE
|
||||
)
|
||||
`)
|
||||
if err != nil {
|
||||
t.Fatalf("create observations: %v", err)
|
||||
}
|
||||
|
||||
_, err = db.Exec(`
|
||||
CREATE TABLE IF NOT EXISTS session_summaries (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
sdk_session_id TEXT NOT NULL,
|
||||
project TEXT NOT NULL,
|
||||
request TEXT,
|
||||
investigated TEXT,
|
||||
learned TEXT,
|
||||
completed TEXT,
|
||||
next_steps TEXT,
|
||||
files_read TEXT,
|
||||
files_edited TEXT,
|
||||
notes TEXT,
|
||||
prompt_number INTEGER,
|
||||
discovery_tokens INTEGER DEFAULT 0,
|
||||
created_at TEXT NOT NULL,
|
||||
created_at_epoch INTEGER NOT NULL,
|
||||
FOREIGN KEY(sdk_session_id) REFERENCES sdk_sessions(sdk_session_id) ON DELETE CASCADE
|
||||
)
|
||||
`)
|
||||
if err != nil {
|
||||
t.Fatalf("create session_summaries: %v", err)
|
||||
}
|
||||
|
||||
_, err = db.Exec(`
|
||||
CREATE TABLE IF NOT EXISTS user_prompts (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
claude_session_id TEXT NOT NULL,
|
||||
prompt_number INTEGER NOT NULL,
|
||||
prompt_text TEXT NOT NULL,
|
||||
matched_observations INTEGER DEFAULT 0,
|
||||
created_at TEXT NOT NULL,
|
||||
created_at_epoch INTEGER NOT NULL,
|
||||
FOREIGN KEY(claude_session_id) REFERENCES sdk_sessions(claude_session_id) ON DELETE CASCADE
|
||||
)
|
||||
`)
|
||||
if err != nil {
|
||||
t.Fatalf("create user_prompts: %v", err)
|
||||
}
|
||||
|
||||
indexes := []string{
|
||||
`CREATE INDEX IF NOT EXISTS idx_sdk_sessions_claude_id ON sdk_sessions(claude_session_id)`,
|
||||
`CREATE INDEX IF NOT EXISTS idx_sdk_sessions_sdk_id ON sdk_sessions(sdk_session_id)`,
|
||||
`CREATE INDEX IF NOT EXISTS idx_sdk_sessions_project ON sdk_sessions(project)`,
|
||||
`CREATE INDEX IF NOT EXISTS idx_observations_sdk_session ON observations(sdk_session_id)`,
|
||||
`CREATE INDEX IF NOT EXISTS idx_observations_project ON observations(project)`,
|
||||
`CREATE INDEX IF NOT EXISTS idx_observations_scope ON observations(scope)`,
|
||||
`CREATE INDEX IF NOT EXISTS idx_observations_created ON observations(created_at_epoch DESC)`,
|
||||
`CREATE INDEX IF NOT EXISTS idx_session_summaries_sdk_session ON session_summaries(sdk_session_id)`,
|
||||
`CREATE INDEX IF NOT EXISTS idx_session_summaries_project ON session_summaries(project)`,
|
||||
`CREATE INDEX IF NOT EXISTS idx_user_prompts_claude_session ON user_prompts(claude_session_id)`,
|
||||
`CREATE INDEX IF NOT EXISTS idx_user_prompts_created ON user_prompts(created_at_epoch DESC)`,
|
||||
}
|
||||
for _, idx := range indexes {
|
||||
if _, err := db.Exec(idx); err != nil {
|
||||
t.Fatalf("create index: %v", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// seedSession creates a test session in the database.
|
||||
func seedSession(t *testing.T, db *sql.DB, claudeSessionID, sdkSessionID, project string) {
|
||||
t.Helper()
|
||||
|
||||
_, err := db.Exec(`
|
||||
INSERT INTO sdk_sessions (claude_session_id, sdk_session_id, project, started_at, started_at_epoch, status)
|
||||
VALUES (?, ?, ?, datetime('now'), strftime('%s', 'now') * 1000, 'active')
|
||||
`, claudeSessionID, sdkSessionID, project)
|
||||
if err != nil {
|
||||
t.Fatalf("seed session: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// hasFTS5 checks if FTS5 is available in the SQLite build.
|
||||
func hasFTS5(db *sql.DB) bool {
|
||||
_, err := db.Exec("CREATE VIRTUAL TABLE IF NOT EXISTS fts5_test USING fts5(content)")
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
_, _ = db.Exec("DROP TABLE IF EXISTS fts5_test")
|
||||
return true
|
||||
}
|
||||
|
||||
// createFTSTables creates FTS5 virtual tables and triggers for full-text search.
|
||||
func createFTSTables(t *testing.T, db *sql.DB) {
|
||||
t.Helper()
|
||||
|
||||
if !hasFTS5(db) {
|
||||
t.Skip("FTS5 not available in this SQLite build")
|
||||
}
|
||||
|
||||
_, err := db.Exec(`
|
||||
CREATE VIRTUAL TABLE IF NOT EXISTS observations_fts USING fts5(
|
||||
title, subtitle, narrative,
|
||||
content='observations',
|
||||
content_rowid='id'
|
||||
)
|
||||
`)
|
||||
if err != nil {
|
||||
t.Fatalf("create observations_fts: %v", err)
|
||||
}
|
||||
|
||||
_, err = db.Exec(`
|
||||
CREATE TRIGGER IF NOT EXISTS observations_ai AFTER INSERT ON observations BEGIN
|
||||
INSERT INTO observations_fts(rowid, title, subtitle, narrative)
|
||||
VALUES (new.id, new.title, new.subtitle, new.narrative);
|
||||
END
|
||||
`)
|
||||
if err != nil {
|
||||
t.Fatalf("create observations_ai trigger: %v", err)
|
||||
}
|
||||
|
||||
_, err = db.Exec(`
|
||||
CREATE TRIGGER IF NOT EXISTS observations_ad AFTER DELETE ON observations BEGIN
|
||||
INSERT INTO observations_fts(observations_fts, rowid, title, subtitle, narrative)
|
||||
VALUES ('delete', old.id, old.title, old.subtitle, old.narrative);
|
||||
END
|
||||
`)
|
||||
if err != nil {
|
||||
t.Fatalf("create observations_ad trigger: %v", err)
|
||||
}
|
||||
|
||||
_, err = db.Exec(`
|
||||
CREATE TRIGGER IF NOT EXISTS observations_au AFTER UPDATE ON observations BEGIN
|
||||
INSERT INTO observations_fts(observations_fts, rowid, title, subtitle, narrative)
|
||||
VALUES ('delete', old.id, old.title, old.subtitle, old.narrative);
|
||||
INSERT INTO observations_fts(rowid, title, subtitle, narrative)
|
||||
VALUES (new.id, new.title, new.subtitle, new.narrative);
|
||||
END
|
||||
`)
|
||||
if err != nil {
|
||||
t.Fatalf("create observations_au trigger: %v", err)
|
||||
}
|
||||
|
||||
_, err = db.Exec(`
|
||||
CREATE VIRTUAL TABLE IF NOT EXISTS session_summaries_fts USING fts5(
|
||||
request, investigated, learned, completed, next_steps, notes,
|
||||
content='session_summaries',
|
||||
content_rowid='id'
|
||||
)
|
||||
`)
|
||||
if err != nil {
|
||||
t.Fatalf("create session_summaries_fts: %v", err)
|
||||
}
|
||||
|
||||
_, err = db.Exec(`
|
||||
CREATE TRIGGER IF NOT EXISTS summaries_ai AFTER INSERT ON session_summaries BEGIN
|
||||
INSERT INTO session_summaries_fts(rowid, request, investigated, learned, completed, next_steps, notes)
|
||||
VALUES (new.id, new.request, new.investigated, new.learned, new.completed, new.next_steps, new.notes);
|
||||
END
|
||||
`)
|
||||
if err != nil {
|
||||
t.Fatalf("create summaries_ai trigger: %v", err)
|
||||
}
|
||||
|
||||
_, err = db.Exec(`
|
||||
CREATE TRIGGER IF NOT EXISTS summaries_ad AFTER DELETE ON session_summaries BEGIN
|
||||
INSERT INTO session_summaries_fts(session_summaries_fts, rowid, request, investigated, learned, completed, next_steps, notes)
|
||||
VALUES ('delete', old.id, old.request, old.investigated, old.learned, old.completed, old.next_steps, old.notes);
|
||||
END
|
||||
`)
|
||||
if err != nil {
|
||||
t.Fatalf("create summaries_ad trigger: %v", err)
|
||||
}
|
||||
|
||||
_, err = db.Exec(`
|
||||
CREATE VIRTUAL TABLE IF NOT EXISTS user_prompts_fts USING fts5(
|
||||
prompt_text,
|
||||
content='user_prompts',
|
||||
content_rowid='id'
|
||||
)
|
||||
`)
|
||||
if err != nil {
|
||||
t.Fatalf("create user_prompts_fts: %v", err)
|
||||
}
|
||||
|
||||
_, err = db.Exec(`
|
||||
CREATE TRIGGER IF NOT EXISTS prompts_ai AFTER INSERT ON user_prompts BEGIN
|
||||
INSERT INTO user_prompts_fts(rowid, prompt_text)
|
||||
VALUES (new.id, new.prompt_text);
|
||||
END
|
||||
`)
|
||||
if err != nil {
|
||||
t.Fatalf("create prompts_ai trigger: %v", err)
|
||||
}
|
||||
|
||||
_, err = db.Exec(`
|
||||
CREATE TRIGGER IF NOT EXISTS prompts_ad AFTER DELETE ON user_prompts BEGIN
|
||||
INSERT INTO user_prompts_fts(user_prompts_fts, rowid, prompt_text)
|
||||
VALUES ('delete', old.id, old.prompt_text);
|
||||
END
|
||||
`)
|
||||
if err != nil {
|
||||
t.Fatalf("create prompts_ad trigger: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// createAllTables creates all tables including FTS5 for comprehensive testing.
|
||||
func createAllTables(t *testing.T, db *sql.DB) {
|
||||
t.Helper()
|
||||
createBaseTables(t, db)
|
||||
createFTSTables(t, db)
|
||||
}
|
||||
@@ -1 +0,0 @@
|
||||
bge-small-en-v1.5
|
||||
Binary file not shown.
@@ -1,7 +1,21 @@
|
||||
{
|
||||
"version": "1.0",
|
||||
"truncation": null,
|
||||
"padding": null,
|
||||
"truncation": {
|
||||
"direction": "Right",
|
||||
"max_length": 128,
|
||||
"strategy": "LongestFirst",
|
||||
"stride": 0
|
||||
},
|
||||
"padding": {
|
||||
"strategy": {
|
||||
"Fixed": 128
|
||||
},
|
||||
"direction": "Right",
|
||||
"pad_to_multiple_of": null,
|
||||
"pad_id": 0,
|
||||
"pad_type_id": 0,
|
||||
"pad_token": "[PAD]"
|
||||
},
|
||||
"added_tokens": [
|
||||
{
|
||||
"id": 0,
|
||||
|
||||
@@ -1,152 +0,0 @@
|
||||
// Package embedding provides text embedding generation with swappable models.
|
||||
package embedding
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"sync"
|
||||
)
|
||||
|
||||
// PoolingStrategy defines how to pool token embeddings into sentence embeddings.
|
||||
type PoolingStrategy string
|
||||
|
||||
const (
|
||||
// PoolingNone means the model already outputs sentence embeddings directly.
|
||||
PoolingNone PoolingStrategy = "none"
|
||||
// PoolingMean averages all token embeddings (weighted by attention mask).
|
||||
PoolingMean PoolingStrategy = "mean"
|
||||
// PoolingCLS uses only the [CLS] token embedding.
|
||||
PoolingCLS PoolingStrategy = "cls"
|
||||
)
|
||||
|
||||
// ONNXConfig describes ONNX-specific model configuration.
|
||||
// This allows different models to specify their tensor names and pooling needs.
|
||||
type ONNXConfig struct {
|
||||
Pooling PoolingStrategy
|
||||
InputNames []string
|
||||
OutputNames []string
|
||||
HiddenSize int
|
||||
}
|
||||
|
||||
// EmbeddingModel represents a text embedding model.
|
||||
type EmbeddingModel interface {
|
||||
// Name returns the human-readable model name (e.g., "bge-small-en-v1.5").
|
||||
Name() string
|
||||
|
||||
// Version returns a short version string for storage (e.g., "bge-v1.5").
|
||||
Version() string
|
||||
|
||||
// Dimensions returns the embedding vector size.
|
||||
Dimensions() int
|
||||
|
||||
// Embed generates an embedding for a single text.
|
||||
Embed(text string) ([]float32, error)
|
||||
|
||||
// EmbedBatch generates embeddings for multiple texts.
|
||||
EmbedBatch(texts []string) ([][]float32, error)
|
||||
|
||||
// Close releases model resources.
|
||||
Close() error
|
||||
}
|
||||
|
||||
// ONNXConfigurer is an optional interface that models can implement
|
||||
// to expose their ONNX configuration for introspection.
|
||||
type ONNXConfigurer interface {
|
||||
// ONNXConfig returns the model's ONNX configuration.
|
||||
ONNXConfig() ONNXConfig
|
||||
}
|
||||
|
||||
// ModelMetadata describes an embedding model for UI/config.
|
||||
type ModelMetadata struct {
|
||||
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.
|
||||
type ModelFactory func() (EmbeddingModel, error)
|
||||
|
||||
// ModelRegistry provides model lookup by version.
|
||||
type ModelRegistry struct {
|
||||
models map[string]ModelFactory
|
||||
metadata map[string]ModelMetadata
|
||||
defaultModel string
|
||||
mu sync.RWMutex
|
||||
}
|
||||
|
||||
// NewModelRegistry creates a new model registry.
|
||||
func NewModelRegistry() *ModelRegistry {
|
||||
return &ModelRegistry{
|
||||
models: make(map[string]ModelFactory),
|
||||
metadata: make(map[string]ModelMetadata),
|
||||
}
|
||||
}
|
||||
|
||||
// Register adds a model factory to the registry.
|
||||
func (r *ModelRegistry) Register(meta ModelMetadata, factory ModelFactory) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
|
||||
r.models[meta.Version] = factory
|
||||
r.metadata[meta.Version] = meta
|
||||
|
||||
if meta.Default {
|
||||
r.defaultModel = meta.Version
|
||||
}
|
||||
}
|
||||
|
||||
// Get creates a new instance of the model with the given version.
|
||||
func (r *ModelRegistry) Get(version string) (EmbeddingModel, error) {
|
||||
r.mu.RLock()
|
||||
factory, ok := r.models[version]
|
||||
r.mu.RUnlock()
|
||||
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("unknown model version: %s", version)
|
||||
}
|
||||
|
||||
return factory()
|
||||
}
|
||||
|
||||
// Default returns the default model version.
|
||||
func (r *ModelRegistry) Default() string {
|
||||
r.mu.RLock()
|
||||
defer r.mu.RUnlock()
|
||||
return r.defaultModel
|
||||
}
|
||||
|
||||
// List returns metadata for all registered models.
|
||||
func (r *ModelRegistry) List() []ModelMetadata {
|
||||
r.mu.RLock()
|
||||
defer r.mu.RUnlock()
|
||||
|
||||
result := make([]ModelMetadata, 0, len(r.metadata))
|
||||
for _, meta := range r.metadata {
|
||||
result = append(result, meta)
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
// DefaultRegistry is the global model registry with all available models.
|
||||
var DefaultRegistry = NewModelRegistry()
|
||||
|
||||
// RegisterModel adds a model to the default registry.
|
||||
func RegisterModel(meta ModelMetadata, factory ModelFactory) {
|
||||
DefaultRegistry.Register(meta, factory)
|
||||
}
|
||||
|
||||
// GetModel creates a model instance from the default registry.
|
||||
func GetModel(version string) (EmbeddingModel, error) {
|
||||
return DefaultRegistry.Get(version)
|
||||
}
|
||||
|
||||
// GetDefaultModel returns the default model version from the default registry.
|
||||
func GetDefaultModel() string {
|
||||
return DefaultRegistry.Default()
|
||||
}
|
||||
|
||||
// ListModels returns metadata for all models in the default registry.
|
||||
func ListModels() []ModelMetadata {
|
||||
return DefaultRegistry.List()
|
||||
}
|
||||
+58
-279
@@ -1,4 +1,4 @@
|
||||
// Package embedding provides text embedding generation with swappable models.
|
||||
// Package embedding provides text embedding generation using all-MiniLM-L6-v2.
|
||||
package embedding
|
||||
|
||||
import (
|
||||
@@ -15,50 +15,19 @@ import (
|
||||
ort "github.com/yalue/onnxruntime_go"
|
||||
)
|
||||
|
||||
// EmbeddingDim is the dimension of embeddings produced by the current model.
|
||||
// Both all-MiniLM-L6-v2 and bge-small-en-v1.5 produce 384-dimensional embeddings.
|
||||
// EmbeddingDim is the dimension of embeddings produced by all-MiniLM-L6-v2.
|
||||
const EmbeddingDim = 384
|
||||
|
||||
// Model version constants
|
||||
const (
|
||||
// BGEModelVersion is the version string for bge-small-en-v1.5
|
||||
BGEModelVersion = "bge-v1.5"
|
||||
// BGEModelName is the human-readable name for bge-small-en-v1.5
|
||||
BGEModelName = "bge-small-en-v1.5"
|
||||
// DefaultModelVersion is the default model to use
|
||||
DefaultModelVersion = BGEModelVersion
|
||||
)
|
||||
|
||||
// MaxSequenceLength is the maximum token sequence length for the model.
|
||||
const MaxSequenceLength = 512
|
||||
|
||||
// bgeONNXConfig defines the ONNX configuration for BGE models.
|
||||
// BGE outputs last_hidden_state and requires mean pooling.
|
||||
var bgeONNXConfig = ONNXConfig{
|
||||
InputNames: []string{"input_ids", "attention_mask", "token_type_ids"},
|
||||
OutputNames: []string{"last_hidden_state"},
|
||||
Pooling: PoolingMean,
|
||||
HiddenSize: EmbeddingDim,
|
||||
}
|
||||
|
||||
// bgeModel is the ONNX-based embedding model implementation.
|
||||
// Currently supports bge-small-en-v1.5 (previously all-MiniLM-L6-v2).
|
||||
type bgeModel struct {
|
||||
// Service provides thread-safe text embedding generation.
|
||||
type Service struct {
|
||||
tk *tokenizer.Tokenizer
|
||||
session *ort.DynamicAdvancedSession
|
||||
libDir string
|
||||
config ONNXConfig
|
||||
mu sync.Mutex
|
||||
libDir string // temp directory containing extracted libraries
|
||||
}
|
||||
|
||||
// Compile-time check that bgeModel implements EmbeddingModel
|
||||
var _ EmbeddingModel = (*bgeModel)(nil)
|
||||
|
||||
// Compile-time check that bgeModel implements ONNXConfigurer
|
||||
var _ ONNXConfigurer = (*bgeModel)(nil)
|
||||
|
||||
// newBGEModel creates a new BGE embedding model using bundled ONNX runtime and model.
|
||||
func newBGEModel() (EmbeddingModel, error) {
|
||||
// NewService creates a new embedding service using bundled ONNX runtime and model.
|
||||
func NewService() (*Service, error) {
|
||||
// Extract ONNX runtime library to temp directory
|
||||
libDir, err := extractONNXLibrary()
|
||||
if err != nil {
|
||||
@@ -80,41 +49,22 @@ func newBGEModel() (EmbeddingModel, error) {
|
||||
return nil, fmt.Errorf("load tokenizer: %w", err)
|
||||
}
|
||||
|
||||
// Create ONNX session using model-specific configuration
|
||||
config := bgeONNXConfig
|
||||
session, err := ort.NewDynamicAdvancedSessionWithONNXData(modelData, config.InputNames, config.OutputNames, nil)
|
||||
// Create ONNX session with embedded model
|
||||
inputNames := []string{"input_ids", "attention_mask", "token_type_ids"}
|
||||
outputNames := []string{"sentence_embedding"}
|
||||
|
||||
session, err := ort.NewDynamicAdvancedSessionWithONNXData(modelData, inputNames, outputNames, nil)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("create ONNX session: %w", err)
|
||||
}
|
||||
|
||||
return &bgeModel{
|
||||
return &Service{
|
||||
tk: tk,
|
||||
session: session,
|
||||
libDir: libDir,
|
||||
config: config,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// ONNXConfig returns the model's ONNX configuration.
|
||||
func (m *bgeModel) ONNXConfig() ONNXConfig {
|
||||
return m.config
|
||||
}
|
||||
|
||||
// Name returns the human-readable model name.
|
||||
func (m *bgeModel) Name() string {
|
||||
return BGEModelName
|
||||
}
|
||||
|
||||
// Version returns the short version string for storage.
|
||||
func (m *bgeModel) Version() string {
|
||||
return BGEModelVersion
|
||||
}
|
||||
|
||||
// Dimensions returns the embedding vector size.
|
||||
func (m *bgeModel) Dimensions() int {
|
||||
return EmbeddingDim
|
||||
}
|
||||
|
||||
// extractONNXLibrary extracts the embedded ONNX runtime library to a temp directory.
|
||||
// Uses content hash to avoid re-extracting if already present.
|
||||
func extractONNXLibrary() (string, error) {
|
||||
@@ -157,15 +107,15 @@ func extractONNXLibrary() (string, error) {
|
||||
|
||||
// Embed generates an embedding for a single text.
|
||||
// Returns a 384-dimensional float32 vector.
|
||||
func (m *bgeModel) Embed(text string) ([]float32, error) {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
func (s *Service) Embed(text string) ([]float32, error) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
if text == "" {
|
||||
return make([]float32, EmbeddingDim), nil
|
||||
}
|
||||
|
||||
results, err := m.computeBatch([]string{text})
|
||||
results, err := s.computeBatch([]string{text})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -177,13 +127,13 @@ func (m *bgeModel) Embed(text string) ([]float32, error) {
|
||||
|
||||
// EmbedBatch generates embeddings for multiple texts.
|
||||
// Returns slice of 384-dimensional float32 vectors.
|
||||
func (m *bgeModel) EmbedBatch(texts []string) ([][]float32, error) {
|
||||
func (s *Service) EmbedBatch(texts []string) ([][]float32, error) {
|
||||
if len(texts) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
// Filter out empty texts and track indices
|
||||
nonEmpty := make([]string, 0, len(texts))
|
||||
@@ -205,7 +155,7 @@ func (m *bgeModel) EmbedBatch(texts []string) ([][]float32, error) {
|
||||
}
|
||||
|
||||
// Compute embeddings for non-empty texts
|
||||
embeddings, err := m.computeBatch(nonEmpty)
|
||||
embeddings, err := s.computeBatch(nonEmpty)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("compute batch embeddings: %w", err)
|
||||
}
|
||||
@@ -223,7 +173,7 @@ func (m *bgeModel) EmbedBatch(texts []string) ([][]float32, error) {
|
||||
}
|
||||
|
||||
// computeBatch runs inference on a batch of texts. Must be called with lock held.
|
||||
func (m *bgeModel) computeBatch(sentences []string) ([][]float32, error) {
|
||||
func (s *Service) computeBatch(sentences []string) ([][]float32, error) {
|
||||
if len(sentences) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
@@ -234,57 +184,31 @@ func (m *bgeModel) computeBatch(sentences []string) ([][]float32, error) {
|
||||
inputBatch[i] = tokenizer.NewSingleEncodeInput(tokenizer.NewRawInputSequence(sent))
|
||||
}
|
||||
|
||||
encodings, err := m.tk.EncodeBatch(inputBatch, true)
|
||||
encodings, err := s.tk.EncodeBatch(inputBatch, true)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("tokenize: %w", err)
|
||||
}
|
||||
|
||||
batchSize := len(encodings)
|
||||
hiddenSize := m.config.HiddenSize
|
||||
|
||||
// Find max sequence length across all encodings (tokenizer may not pad uniformly)
|
||||
// Also enforce MaxSequenceLength to prevent model errors
|
||||
seqLength := 0
|
||||
for _, enc := range encodings {
|
||||
if len(enc.Ids) > seqLength {
|
||||
seqLength = len(enc.Ids)
|
||||
}
|
||||
}
|
||||
// Truncate to max model sequence length
|
||||
if seqLength > MaxSequenceLength {
|
||||
seqLength = MaxSequenceLength
|
||||
}
|
||||
seqLength := len(encodings[0].Ids)
|
||||
hiddenSize := EmbeddingDim
|
||||
|
||||
inputShape := ort.NewShape(int64(batchSize), int64(seqLength))
|
||||
|
||||
// Create input tensors (pre-filled with zeros for padding)
|
||||
// Create input tensors
|
||||
inputIdsData := make([]int64, batchSize*seqLength)
|
||||
attentionMaskData := make([]int64, batchSize*seqLength)
|
||||
tokenTypeIdsData := make([]int64, batchSize*seqLength)
|
||||
|
||||
for b := 0; b < batchSize; b++ {
|
||||
// Copy actual token data (rest remains 0 as padding)
|
||||
// Truncate to seqLength to handle long inputs
|
||||
copyLen := len(encodings[b].Ids)
|
||||
if copyLen > seqLength {
|
||||
copyLen = seqLength
|
||||
for i, id := range encodings[b].Ids {
|
||||
inputIdsData[b*seqLength+i] = int64(id)
|
||||
}
|
||||
for i := 0; i < copyLen; i++ {
|
||||
inputIdsData[b*seqLength+i] = int64(encodings[b].Ids[i])
|
||||
for i, mask := range encodings[b].AttentionMask {
|
||||
attentionMaskData[b*seqLength+i] = int64(mask)
|
||||
}
|
||||
copyLen = len(encodings[b].AttentionMask)
|
||||
if copyLen > seqLength {
|
||||
copyLen = seqLength
|
||||
}
|
||||
for i := 0; i < copyLen; i++ {
|
||||
attentionMaskData[b*seqLength+i] = int64(encodings[b].AttentionMask[i])
|
||||
}
|
||||
copyLen = len(encodings[b].TypeIds)
|
||||
if copyLen > seqLength {
|
||||
copyLen = seqLength
|
||||
}
|
||||
for i := 0; i < copyLen; i++ {
|
||||
tokenTypeIdsData[b*seqLength+i] = int64(encodings[b].TypeIds[i])
|
||||
for i, typeId := range encodings[b].TypeIds {
|
||||
tokenTypeIdsData[b*seqLength+i] = int64(typeId)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -292,145 +216,65 @@ func (m *bgeModel) computeBatch(sentences []string) ([][]float32, error) {
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("create input_ids tensor: %w", err)
|
||||
}
|
||||
defer func() { _ = inputIdsTensor.Destroy() }()
|
||||
defer inputIdsTensor.Destroy()
|
||||
|
||||
attentionMaskTensor, err := ort.NewTensor(inputShape, attentionMaskData)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("create attention_mask tensor: %w", err)
|
||||
}
|
||||
defer func() { _ = attentionMaskTensor.Destroy() }()
|
||||
defer attentionMaskTensor.Destroy()
|
||||
|
||||
tokenTypeIdsTensor, err := ort.NewTensor(inputShape, tokenTypeIdsData)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("create token_type_ids tensor: %w", err)
|
||||
}
|
||||
defer func() { _ = tokenTypeIdsTensor.Destroy() }()
|
||||
defer tokenTypeIdsTensor.Destroy()
|
||||
|
||||
// Create output tensor based on pooling strategy
|
||||
var outputShape ort.Shape
|
||||
|
||||
switch m.config.Pooling {
|
||||
case PoolingNone:
|
||||
// Direct sentence embedding output: [batch, hidden]
|
||||
outputShape = ort.NewShape(int64(batchSize), int64(hiddenSize))
|
||||
case PoolingMean, PoolingCLS:
|
||||
// Token-level output: [batch, seq_len, hidden]
|
||||
outputShape = ort.NewShape(int64(batchSize), int64(seqLength), int64(hiddenSize))
|
||||
default:
|
||||
outputShape = ort.NewShape(int64(batchSize), int64(hiddenSize))
|
||||
}
|
||||
|
||||
outputTensor, err := ort.NewEmptyTensor[float32](outputShape)
|
||||
sentenceOutputShape := ort.NewShape(int64(batchSize), int64(hiddenSize))
|
||||
sentenceOutputTensor, err := ort.NewEmptyTensor[float32](sentenceOutputShape)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("create output tensor: %w", err)
|
||||
}
|
||||
defer func() { _ = outputTensor.Destroy() }()
|
||||
defer sentenceOutputTensor.Destroy()
|
||||
|
||||
// Run inference
|
||||
inputTensors := []ort.Value{inputIdsTensor, attentionMaskTensor, tokenTypeIdsTensor}
|
||||
outputTensors := []ort.Value{outputTensor}
|
||||
outputTensors := []ort.Value{sentenceOutputTensor}
|
||||
|
||||
if err := m.session.Run(inputTensors, outputTensors); err != nil {
|
||||
if err := s.session.Run(inputTensors, outputTensors); err != nil {
|
||||
return nil, fmt.Errorf("run inference: %w", err)
|
||||
}
|
||||
|
||||
// Extract and pool results based on strategy
|
||||
flatOutput := outputTensor.GetData()
|
||||
|
||||
switch m.config.Pooling {
|
||||
case PoolingNone:
|
||||
// Direct output, no pooling needed
|
||||
expectedSize := batchSize * hiddenSize
|
||||
if len(flatOutput) != expectedSize {
|
||||
return nil, fmt.Errorf("unexpected output size: got %d, expected %d", len(flatOutput), expectedSize)
|
||||
}
|
||||
results := make([][]float32, batchSize)
|
||||
for i := 0; i < batchSize; i++ {
|
||||
start := i * hiddenSize
|
||||
end := start + hiddenSize
|
||||
results[i] = make([]float32, hiddenSize)
|
||||
copy(results[i], flatOutput[start:end])
|
||||
}
|
||||
return results, nil
|
||||
|
||||
case PoolingMean:
|
||||
// Mean pooling over tokens (weighted by attention mask)
|
||||
return meanPooling(flatOutput, attentionMaskData, batchSize, seqLength, hiddenSize), nil
|
||||
|
||||
case PoolingCLS:
|
||||
// CLS token pooling (first token of each sequence)
|
||||
return clsPooling(flatOutput, batchSize, seqLength, hiddenSize), nil
|
||||
|
||||
default:
|
||||
return nil, fmt.Errorf("unknown pooling strategy: %s", m.config.Pooling)
|
||||
// Extract results
|
||||
flatOutput := sentenceOutputTensor.GetData()
|
||||
expectedSize := batchSize * hiddenSize
|
||||
if len(flatOutput) != expectedSize {
|
||||
return nil, fmt.Errorf("unexpected output size: got %d, expected %d", len(flatOutput), expectedSize)
|
||||
}
|
||||
}
|
||||
|
||||
// meanPooling applies mean pooling over token embeddings, weighted by attention mask.
|
||||
// Input shape: [batch, seq_len, hidden], attention mask: [batch, seq_len]
|
||||
// Output shape: [batch, hidden]
|
||||
func meanPooling(embeddings []float32, attentionMask []int64, batchSize, seqLen, hiddenSize int) [][]float32 {
|
||||
results := make([][]float32, batchSize)
|
||||
|
||||
for b := 0; b < batchSize; b++ {
|
||||
result := make([]float32, hiddenSize)
|
||||
var maskSum float32
|
||||
|
||||
// Sum embeddings weighted by attention mask
|
||||
for s := 0; s < seqLen; s++ {
|
||||
maskVal := float32(attentionMask[b*seqLen+s])
|
||||
maskSum += maskVal
|
||||
|
||||
if maskVal > 0 {
|
||||
embOffset := (b*seqLen + s) * hiddenSize
|
||||
for h := 0; h < hiddenSize; h++ {
|
||||
result[h] += embeddings[embOffset+h] * maskVal
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Normalize by mask sum (avoid division by zero)
|
||||
if maskSum > 0 {
|
||||
for h := 0; h < hiddenSize; h++ {
|
||||
result[h] /= maskSum
|
||||
}
|
||||
}
|
||||
|
||||
results[b] = result
|
||||
for i := 0; i < batchSize; i++ {
|
||||
start := i * hiddenSize
|
||||
end := start + hiddenSize
|
||||
results[i] = make([]float32, hiddenSize)
|
||||
copy(results[i], flatOutput[start:end])
|
||||
}
|
||||
|
||||
return results
|
||||
}
|
||||
|
||||
// clsPooling extracts the [CLS] token embedding (first token).
|
||||
// Input shape: [batch, seq_len, hidden]
|
||||
// Output shape: [batch, hidden]
|
||||
func clsPooling(embeddings []float32, batchSize, seqLen, hiddenSize int) [][]float32 {
|
||||
results := make([][]float32, batchSize)
|
||||
|
||||
for b := 0; b < batchSize; b++ {
|
||||
result := make([]float32, hiddenSize)
|
||||
// CLS token is at position 0
|
||||
embOffset := b * seqLen * hiddenSize
|
||||
copy(result, embeddings[embOffset:embOffset+hiddenSize])
|
||||
results[b] = result
|
||||
}
|
||||
|
||||
return results
|
||||
return results, nil
|
||||
}
|
||||
|
||||
// Close releases model resources.
|
||||
func (m *bgeModel) Close() error {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
func (s *Service) Close() error {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
var errs []error
|
||||
|
||||
if m.session != nil {
|
||||
if err := m.session.Destroy(); err != nil {
|
||||
if s.session != nil {
|
||||
if err := s.session.Destroy(); err != nil {
|
||||
errs = append(errs, fmt.Errorf("destroy session: %w", err))
|
||||
}
|
||||
m.session = nil
|
||||
s.session = nil
|
||||
}
|
||||
|
||||
if err := ort.DestroyEnvironment(); err != nil {
|
||||
@@ -438,75 +282,10 @@ func (m *bgeModel) Close() error {
|
||||
}
|
||||
|
||||
// Optionally clean up extracted library (leave for caching)
|
||||
// os.RemoveAll(m.libDir)
|
||||
// os.RemoveAll(s.libDir)
|
||||
|
||||
if len(errs) > 0 {
|
||||
return errs[0]
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Register the BGE model with the default registry at init time
|
||||
func init() {
|
||||
RegisterModel(ModelMetadata{
|
||||
Name: BGEModelName,
|
||||
Version: BGEModelVersion,
|
||||
Dimensions: EmbeddingDim,
|
||||
Description: "High-quality semantic search model",
|
||||
Default: true,
|
||||
}, newBGEModel)
|
||||
}
|
||||
|
||||
// Service provides thread-safe text embedding generation with model abstraction.
|
||||
type Service struct {
|
||||
model EmbeddingModel
|
||||
}
|
||||
|
||||
// NewService creates a new embedding service using the default model.
|
||||
func NewService() (*Service, error) {
|
||||
return NewServiceWithModel(DefaultModelVersion)
|
||||
}
|
||||
|
||||
// NewServiceWithModel creates a new embedding service using the specified model.
|
||||
func NewServiceWithModel(version string) (*Service, error) {
|
||||
if version == "" {
|
||||
version = DefaultModelVersion
|
||||
}
|
||||
|
||||
model, err := GetModel(version)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("get model %s: %w", version, err)
|
||||
}
|
||||
|
||||
return &Service{model: model}, nil
|
||||
}
|
||||
|
||||
// Name returns the human-readable model name.
|
||||
func (s *Service) Name() string {
|
||||
return s.model.Name()
|
||||
}
|
||||
|
||||
// Version returns the short version string for storage.
|
||||
func (s *Service) Version() string {
|
||||
return s.model.Version()
|
||||
}
|
||||
|
||||
// Dimensions returns the embedding vector size.
|
||||
func (s *Service) Dimensions() int {
|
||||
return s.model.Dimensions()
|
||||
}
|
||||
|
||||
// Embed generates an embedding for a single text.
|
||||
func (s *Service) Embed(text string) ([]float32, error) {
|
||||
return s.model.Embed(text)
|
||||
}
|
||||
|
||||
// EmbedBatch generates embeddings for multiple texts.
|
||||
func (s *Service) EmbedBatch(texts []string) ([][]float32, error) {
|
||||
return s.model.EmbedBatch(texts)
|
||||
}
|
||||
|
||||
// Close releases model resources.
|
||||
func (s *Service) Close() error {
|
||||
return s.model.Close()
|
||||
}
|
||||
|
||||
@@ -22,10 +22,8 @@ func TestNewService(t *testing.T) {
|
||||
|
||||
defer svc.Close()
|
||||
|
||||
// Verify the service is properly initialized via public methods
|
||||
assert.NotEmpty(t, svc.Name())
|
||||
assert.NotEmpty(t, svc.Version())
|
||||
assert.Equal(t, EmbeddingDim, svc.Dimensions())
|
||||
assert.NotNil(t, svc.tk)
|
||||
assert.NotNil(t, svc.session)
|
||||
}
|
||||
|
||||
// TestEmbed_SingleText tests embedding a single text.
|
||||
@@ -271,8 +269,8 @@ func TestClose(t *testing.T) {
|
||||
err = svc.Close()
|
||||
require.NoError(t, err)
|
||||
|
||||
// After close, embedding should fail (model resources released)
|
||||
// Note: This behavior is model-specific; some models may still work after close
|
||||
// Session should be nil after close
|
||||
assert.Nil(t, svc.session)
|
||||
}
|
||||
|
||||
// TestEmbedBatch_SingleItem tests batch embedding with single item.
|
||||
|
||||
@@ -1,417 +0,0 @@
|
||||
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)))
|
||||
}
|
||||
@@ -1,423 +0,0 @@
|
||||
// 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"
|
||||
}
|
||||
}
|
||||
@@ -1,290 +0,0 @@
|
||||
// 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)
|
||||
}
|
||||
+34
-2747
File diff suppressed because it is too large
Load Diff
+241
-1946
File diff suppressed because it is too large
Load Diff
@@ -1,499 +0,0 @@
|
||||
// Package pattern provides pattern detection and recognition functionality.
|
||||
package pattern
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"sort"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/lukaszraczylo/claude-mnemonic/internal/db/gorm"
|
||||
"github.com/lukaszraczylo/claude-mnemonic/pkg/models"
|
||||
"github.com/rs/zerolog/log"
|
||||
)
|
||||
|
||||
// DetectorConfig contains configuration for the pattern detector.
|
||||
type DetectorConfig struct {
|
||||
// MinMatchScore is the minimum similarity score to consider a match (0.0-1.0).
|
||||
MinMatchScore float64
|
||||
// MinFrequencyForPattern is the minimum occurrences before creating a pattern.
|
||||
MinFrequencyForPattern int
|
||||
// AnalysisInterval is how often to run background pattern analysis.
|
||||
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.
|
||||
func DefaultConfig() DetectorConfig {
|
||||
return DetectorConfig{
|
||||
MinMatchScore: 0.3, // 30% similarity threshold
|
||||
MinFrequencyForPattern: 2, // At least 2 occurrences to form a pattern
|
||||
AnalysisInterval: 5 * time.Minute,
|
||||
MaxPatternsToTrack: 1000,
|
||||
MaxCandidates: 500, // Prevent unbounded growth
|
||||
}
|
||||
}
|
||||
|
||||
// PatternSyncFunc is a callback for syncing patterns to vector store.
|
||||
type PatternSyncFunc func(pattern *models.Pattern)
|
||||
|
||||
// Detector detects and tracks recurring patterns across observations.
|
||||
type Detector struct {
|
||||
ctx context.Context
|
||||
patternStore *gorm.PatternStore
|
||||
observationStore *gorm.ObservationStore
|
||||
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.
|
||||
func (d *Detector) SetSyncFunc(fn PatternSyncFunc) {
|
||||
d.syncFunc = fn
|
||||
}
|
||||
|
||||
// candidatePattern tracks a potential pattern before it reaches frequency threshold.
|
||||
type candidatePattern struct {
|
||||
patternType models.PatternType
|
||||
title string
|
||||
signature []string
|
||||
observationIDs []int64
|
||||
projects []string
|
||||
lastSeenEpoch int64
|
||||
}
|
||||
|
||||
// NewDetector creates a new pattern detector.
|
||||
func NewDetector(patternStore *gorm.PatternStore, observationStore *gorm.ObservationStore, config DetectorConfig) *Detector {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
return &Detector{
|
||||
config: config,
|
||||
patternStore: patternStore,
|
||||
observationStore: observationStore,
|
||||
candidates: make(map[string]*candidatePattern),
|
||||
ctx: ctx,
|
||||
cancel: cancel,
|
||||
}
|
||||
}
|
||||
|
||||
// Start begins background pattern analysis.
|
||||
func (d *Detector) Start() {
|
||||
d.wg.Add(1)
|
||||
go d.backgroundAnalysis()
|
||||
log.Info().Dur("interval", d.config.AnalysisInterval).Msg("Pattern detector started")
|
||||
}
|
||||
|
||||
// Stop stops background pattern analysis.
|
||||
func (d *Detector) Stop() {
|
||||
d.cancel()
|
||||
d.wg.Wait()
|
||||
log.Info().Msg("Pattern detector stopped")
|
||||
}
|
||||
|
||||
// backgroundAnalysis runs periodic pattern analysis.
|
||||
func (d *Detector) backgroundAnalysis() {
|
||||
defer d.wg.Done()
|
||||
|
||||
ticker := time.NewTicker(d.config.AnalysisInterval)
|
||||
defer ticker.Stop()
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-d.ctx.Done():
|
||||
return
|
||||
case <-ticker.C:
|
||||
if err := d.AnalyzeRecentObservations(d.ctx); err != nil {
|
||||
log.Warn().Err(err).Msg("Background pattern analysis failed")
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// AnalyzeObservation processes a new observation for pattern detection.
|
||||
// This is called synchronously when a new observation is stored.
|
||||
func (d *Detector) AnalyzeObservation(ctx context.Context, obs *models.Observation) (*DetectionResult, error) {
|
||||
result := &DetectionResult{}
|
||||
|
||||
// Extract signature from observation
|
||||
signature := models.ExtractSignature(
|
||||
obs.Concepts,
|
||||
obs.Title.String,
|
||||
obs.Narrative.String,
|
||||
)
|
||||
if len(signature) == 0 {
|
||||
return result, nil // Nothing to detect
|
||||
}
|
||||
|
||||
// Check against existing patterns
|
||||
matches, err := d.patternStore.FindMatchingPatterns(ctx, signature, d.config.MinMatchScore)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if len(matches) > 0 {
|
||||
// Found existing pattern match
|
||||
bestMatch := matches[0]
|
||||
for _, m := range matches[1:] {
|
||||
if models.CalculateMatchScore(signature, m.Signature) > models.CalculateMatchScore(signature, bestMatch.Signature) {
|
||||
bestMatch = m
|
||||
}
|
||||
}
|
||||
|
||||
// Update the pattern with new occurrence
|
||||
if err := d.patternStore.IncrementPatternFrequency(ctx, bestMatch.ID, obs.Project, obs.ID); err != nil {
|
||||
log.Warn().Err(err).Int64("pattern_id", bestMatch.ID).Msg("Failed to update pattern frequency")
|
||||
}
|
||||
|
||||
result.MatchedPattern = bestMatch
|
||||
result.MatchScore = models.CalculateMatchScore(signature, bestMatch.Signature)
|
||||
result.IsNewPattern = false
|
||||
|
||||
log.Debug().
|
||||
Int64("pattern_id", bestMatch.ID).
|
||||
Str("pattern_name", bestMatch.Name).
|
||||
Float64("score", result.MatchScore).
|
||||
Msg("Observation matched existing pattern")
|
||||
|
||||
return result, nil
|
||||
}
|
||||
|
||||
// No existing pattern match - check candidates
|
||||
candidateKey := generateCandidateKey(signature)
|
||||
d.candidatesMu.Lock()
|
||||
defer d.candidatesMu.Unlock()
|
||||
|
||||
if candidate, exists := d.candidates[candidateKey]; exists {
|
||||
// Update existing candidate
|
||||
candidate.observationIDs = append(candidate.observationIDs, obs.ID)
|
||||
|
||||
// Add project if not already present
|
||||
found := false
|
||||
for _, p := range candidate.projects {
|
||||
if p == obs.Project {
|
||||
found = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
candidate.projects = append(candidate.projects, obs.Project)
|
||||
}
|
||||
candidate.lastSeenEpoch = time.Now().UnixMilli()
|
||||
|
||||
// Check if candidate should be promoted to pattern
|
||||
if len(candidate.observationIDs) >= d.config.MinFrequencyForPattern {
|
||||
pattern, err := d.promoteCandidate(ctx, candidateKey, candidate)
|
||||
if err != nil {
|
||||
log.Warn().Err(err).Msg("Failed to promote candidate to pattern")
|
||||
} else {
|
||||
result.MatchedPattern = pattern
|
||||
result.IsNewPattern = true
|
||||
log.Info().
|
||||
Str("pattern_name", pattern.Name).
|
||||
Int("frequency", pattern.Frequency).
|
||||
Msg("New pattern detected and stored")
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// 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,
|
||||
observationIDs: []int64{obs.ID},
|
||||
projects: []string{obs.Project},
|
||||
patternType: patternType,
|
||||
title: obs.Title.String,
|
||||
lastSeenEpoch: time.Now().UnixMilli(),
|
||||
}
|
||||
log.Debug().
|
||||
Str("candidate_key", candidateKey).
|
||||
Strs("signature", signature).
|
||||
Msg("New pattern candidate created")
|
||||
}
|
||||
|
||||
return result, nil
|
||||
}
|
||||
|
||||
// promoteCandidate converts a candidate to a stored pattern.
|
||||
func (d *Detector) promoteCandidate(ctx context.Context, key string, candidate *candidatePattern) (*models.Pattern, error) {
|
||||
// Generate pattern name from signature
|
||||
name := generatePatternName(candidate.patternType, candidate.signature, candidate.title)
|
||||
|
||||
// Create base pattern using NewPattern with first observation
|
||||
firstProject := ""
|
||||
if len(candidate.projects) > 0 {
|
||||
firstProject = candidate.projects[0]
|
||||
}
|
||||
var firstObsID int64
|
||||
if len(candidate.observationIDs) > 0 {
|
||||
firstObsID = candidate.observationIDs[0]
|
||||
}
|
||||
pattern := models.NewPattern(
|
||||
name,
|
||||
candidate.patternType,
|
||||
"Automatically detected pattern from recurring observations",
|
||||
candidate.signature,
|
||||
firstProject,
|
||||
firstObsID,
|
||||
)
|
||||
|
||||
// Add remaining projects and observations
|
||||
for i := 1; i < len(candidate.projects); i++ {
|
||||
pattern.Projects = append(pattern.Projects, candidate.projects[i])
|
||||
}
|
||||
for i := 1; i < len(candidate.observationIDs); i++ {
|
||||
pattern.ObservationIDs = append(pattern.ObservationIDs, candidate.observationIDs[i])
|
||||
}
|
||||
pattern.Frequency = len(candidate.observationIDs)
|
||||
|
||||
id, err := d.patternStore.StorePattern(ctx, pattern)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
pattern.ID = id
|
||||
|
||||
// Sync to vector store if callback is set
|
||||
if d.syncFunc != nil {
|
||||
d.syncFunc(pattern)
|
||||
}
|
||||
|
||||
// Remove from candidates
|
||||
delete(d.candidates, key)
|
||||
|
||||
return pattern, nil
|
||||
}
|
||||
|
||||
// AnalyzeRecentObservations analyzes recent observations for pattern detection.
|
||||
// This is used for background batch analysis.
|
||||
func (d *Detector) AnalyzeRecentObservations(ctx context.Context) error {
|
||||
// Get observations from the last 24 hours that haven't been analyzed
|
||||
observations, err := d.observationStore.GetRecentObservations(ctx, "", 100)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
analyzed := 0
|
||||
patternsFound := 0
|
||||
for _, obs := range observations {
|
||||
result, err := d.AnalyzeObservation(ctx, obs)
|
||||
if err != nil {
|
||||
log.Warn().Err(err).Int64("obs_id", obs.ID).Msg("Failed to analyze observation")
|
||||
continue
|
||||
}
|
||||
analyzed++
|
||||
if result.MatchedPattern != nil {
|
||||
patternsFound++
|
||||
}
|
||||
}
|
||||
|
||||
if analyzed > 0 {
|
||||
log.Info().
|
||||
Int("analyzed", analyzed).
|
||||
Int("patterns_found", patternsFound).
|
||||
Msg("Background pattern analysis completed")
|
||||
}
|
||||
|
||||
// Clean up old candidates (older than 7 days)
|
||||
d.cleanupOldCandidates()
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// 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
|
||||
}
|
||||
|
||||
// DetectionResult contains the result of pattern detection.
|
||||
type DetectionResult struct {
|
||||
MatchedPattern *models.Pattern
|
||||
MatchScore float64
|
||||
IsNewPattern bool
|
||||
}
|
||||
|
||||
// generateCandidateKey creates a unique key for a signature.
|
||||
func generateCandidateKey(signature []string) string {
|
||||
if len(signature) == 0 {
|
||||
return ""
|
||||
}
|
||||
// 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 {
|
||||
b.WriteString(s)
|
||||
b.WriteByte('|')
|
||||
}
|
||||
return b.String()
|
||||
}
|
||||
|
||||
// generatePatternName creates a human-readable name for a pattern.
|
||||
func generatePatternName(patternType models.PatternType, signature []string, title string) string {
|
||||
// Use title if available and meaningful
|
||||
if title != "" && len(title) < 60 {
|
||||
return title
|
||||
}
|
||||
|
||||
// Otherwise generate from type and signature
|
||||
prefix := ""
|
||||
switch patternType {
|
||||
case models.PatternTypeBug:
|
||||
prefix = "Bug Pattern: "
|
||||
case models.PatternTypeRefactor:
|
||||
prefix = "Refactor Pattern: "
|
||||
case models.PatternTypeArchitecture:
|
||||
prefix = "Architecture Pattern: "
|
||||
case models.PatternTypeAntiPattern:
|
||||
prefix = "Anti-Pattern: "
|
||||
case models.PatternTypeBestPractice:
|
||||
prefix = "Best Practice: "
|
||||
}
|
||||
|
||||
// Use first few signature elements
|
||||
if len(signature) > 0 {
|
||||
name := prefix
|
||||
for i, s := range signature {
|
||||
if i >= 3 {
|
||||
break
|
||||
}
|
||||
if i > 0 {
|
||||
name += ", "
|
||||
}
|
||||
name += s
|
||||
}
|
||||
return name
|
||||
}
|
||||
|
||||
return prefix + "Unnamed"
|
||||
}
|
||||
|
||||
// formatPatternInsight creates a human-readable insight from a pattern.
|
||||
func formatPatternInsight(pattern *models.Pattern) string {
|
||||
insight := "I've encountered this pattern " +
|
||||
itoa(pattern.Frequency) + " times"
|
||||
|
||||
if len(pattern.Projects) > 1 {
|
||||
insight += " across " + itoa(len(pattern.Projects)) + " projects"
|
||||
}
|
||||
|
||||
insight += ". "
|
||||
|
||||
if pattern.Recommendation.Valid && pattern.Recommendation.String != "" {
|
||||
insight += "What works: " + pattern.Recommendation.String
|
||||
} else {
|
||||
switch pattern.Type {
|
||||
case models.PatternTypeBug:
|
||||
insight += "This appears to be a recurring bug pattern."
|
||||
case models.PatternTypeAntiPattern:
|
||||
insight += "This is an identified anti-pattern to avoid."
|
||||
case models.PatternTypeBestPractice:
|
||||
insight += "This is a validated best practice."
|
||||
default:
|
||||
insight += "This is a recognized pattern in the codebase."
|
||||
}
|
||||
}
|
||||
|
||||
return insight
|
||||
}
|
||||
|
||||
// itoa converts int to string without importing strconv.
|
||||
func itoa(n int) string {
|
||||
if n == 0 {
|
||||
return "0"
|
||||
}
|
||||
negative := false
|
||||
if n < 0 {
|
||||
negative = true
|
||||
n = -n
|
||||
}
|
||||
var digits []byte
|
||||
for n > 0 {
|
||||
digits = append([]byte{byte('0' + n%10)}, digits...)
|
||||
n /= 10
|
||||
}
|
||||
if negative {
|
||||
digits = append([]byte{'-'}, digits...)
|
||||
}
|
||||
return string(digits)
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,98 +0,0 @@
|
||||
// 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)
|
||||
})
|
||||
}
|
||||
@@ -1,213 +0,0 @@
|
||||
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)
|
||||
}
|
||||
}
|
||||
@@ -1,14 +0,0 @@
|
||||
// Package reranking provides cross-encoder reranking for search results.
|
||||
package reranking
|
||||
|
||||
import (
|
||||
_ "embed"
|
||||
)
|
||||
|
||||
// Cross-encoder model and tokenizer files - embedded for all platforms
|
||||
//
|
||||
//go:embed assets/model.onnx
|
||||
var crossEncoderModelData []byte
|
||||
|
||||
//go:embed assets/tokenizer.json
|
||||
var crossEncoderTokenizerData []byte
|
||||
Binary file not shown.
File diff suppressed because it is too large
Load Diff
@@ -1,57 +0,0 @@
|
||||
{
|
||||
"added_tokens_decoder": {
|
||||
"0": {
|
||||
"content": "[PAD]",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false,
|
||||
"special": true
|
||||
},
|
||||
"100": {
|
||||
"content": "[UNK]",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false,
|
||||
"special": true
|
||||
},
|
||||
"101": {
|
||||
"content": "[CLS]",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false,
|
||||
"special": true
|
||||
},
|
||||
"102": {
|
||||
"content": "[SEP]",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false,
|
||||
"special": true
|
||||
},
|
||||
"103": {
|
||||
"content": "[MASK]",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false,
|
||||
"special": true
|
||||
}
|
||||
},
|
||||
"clean_up_tokenization_spaces": true,
|
||||
"cls_token": "[CLS]",
|
||||
"do_basic_tokenize": true,
|
||||
"do_lower_case": true,
|
||||
"mask_token": "[MASK]",
|
||||
"model_max_length": 512,
|
||||
"never_split": null,
|
||||
"pad_token": "[PAD]",
|
||||
"sep_token": "[SEP]",
|
||||
"strip_accents": null,
|
||||
"tokenize_chinese_chars": true,
|
||||
"tokenizer_class": "BertTokenizer",
|
||||
"unk_token": "[UNK]"
|
||||
}
|
||||
@@ -1,380 +0,0 @@
|
||||
// Package reranking provides cross-encoder reranking for search results.
|
||||
// Uses MS-MARCO MiniLM L6 v2 cross-encoder model for relevance scoring.
|
||||
package reranking
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"fmt"
|
||||
"math"
|
||||
"sort"
|
||||
"sync"
|
||||
|
||||
"github.com/rs/zerolog/log"
|
||||
"github.com/sugarme/tokenizer"
|
||||
"github.com/sugarme/tokenizer/pretrained"
|
||||
ort "github.com/yalue/onnxruntime_go"
|
||||
)
|
||||
|
||||
const (
|
||||
// ModelName is the human-readable name for the cross-encoder model
|
||||
ModelName = "ms-marco-MiniLM-L6-v2"
|
||||
// ModelVersion is the short version string for identification
|
||||
ModelVersion = "msmarco-v2"
|
||||
// MaxSequenceLength is the maximum combined query+document token length
|
||||
MaxSequenceLength = 512
|
||||
// DefaultCandidateLimit is the default number of candidates to rerank
|
||||
DefaultCandidateLimit = 100
|
||||
// DefaultResultLimit is the default number of results to return after reranking
|
||||
DefaultResultLimit = 10
|
||||
)
|
||||
|
||||
// Candidate represents a search result candidate for reranking.
|
||||
type Candidate struct {
|
||||
Metadata map[string]any
|
||||
RerankInfo map[string]float64
|
||||
ID string
|
||||
Content string
|
||||
Score float64
|
||||
}
|
||||
|
||||
// RerankResult represents a reranked search result.
|
||||
type RerankResult struct {
|
||||
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.
|
||||
type Service struct {
|
||||
tk *tokenizer.Tokenizer
|
||||
session *ort.DynamicAdvancedSession
|
||||
mu sync.Mutex
|
||||
|
||||
// Weight for combining scores: combined = alpha*rerank + (1-alpha)*original
|
||||
// Default 0.7 favors cross-encoder score
|
||||
Alpha float64
|
||||
}
|
||||
|
||||
// Config holds configuration for the reranking service.
|
||||
type Config struct {
|
||||
// Alpha is the weight for combining scores (0.0-1.0)
|
||||
// Higher values favor cross-encoder scores, lower values favor bi-encoder scores
|
||||
Alpha float64
|
||||
}
|
||||
|
||||
// DefaultConfig returns sensible defaults for reranking.
|
||||
func DefaultConfig() Config {
|
||||
return Config{
|
||||
Alpha: 0.7, // Favor cross-encoder by default
|
||||
}
|
||||
}
|
||||
|
||||
// NewService creates a new cross-encoder reranking service.
|
||||
// Note: ONNX runtime must be initialized before calling this (via embedding.NewService).
|
||||
func NewService(cfg Config) (*Service, error) {
|
||||
// Load tokenizer from embedded data
|
||||
tk, err := pretrained.FromReader(bytes.NewReader(crossEncoderTokenizerData))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("load cross-encoder tokenizer: %w", err)
|
||||
}
|
||||
|
||||
// Configure tokenizer for sequence classification (pairs)
|
||||
tk.WithTruncation(&tokenizer.TruncationParams{
|
||||
MaxLength: MaxSequenceLength,
|
||||
Strategy: tokenizer.LongestFirst,
|
||||
Stride: 0,
|
||||
})
|
||||
|
||||
// Cross-encoder outputs a single logit for relevance scoring
|
||||
inputNames := []string{"input_ids", "attention_mask", "token_type_ids"}
|
||||
outputNames := []string{"logits"}
|
||||
|
||||
session, err := ort.NewDynamicAdvancedSessionWithONNXData(
|
||||
crossEncoderModelData,
|
||||
inputNames,
|
||||
outputNames,
|
||||
nil,
|
||||
)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("create cross-encoder ONNX session: %w", err)
|
||||
}
|
||||
|
||||
alpha := cfg.Alpha
|
||||
if alpha <= 0 || alpha > 1 {
|
||||
alpha = 0.7
|
||||
}
|
||||
|
||||
return &Service{
|
||||
tk: tk,
|
||||
session: session,
|
||||
Alpha: alpha,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// Rerank reranks candidates using the cross-encoder model.
|
||||
// Takes a query and list of candidates, returns reranked results.
|
||||
func (s *Service) Rerank(query string, candidates []Candidate, limit int) ([]RerankResult, error) {
|
||||
if len(candidates) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
if limit <= 0 {
|
||||
limit = DefaultResultLimit
|
||||
}
|
||||
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
// Score all query-document pairs
|
||||
scores, err := s.scoreAll(query, candidates)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("score candidates: %w", err)
|
||||
}
|
||||
|
||||
// Build results with combined scores
|
||||
results := make([]RerankResult, len(candidates))
|
||||
for i, c := range candidates {
|
||||
// Normalize cross-encoder score to 0-1 range using sigmoid
|
||||
normalizedRerank := sigmoid(scores[i])
|
||||
|
||||
results[i] = RerankResult{
|
||||
ID: c.ID,
|
||||
Content: c.Content,
|
||||
OriginalScore: c.Score,
|
||||
RerankScore: normalizedRerank,
|
||||
CombinedScore: s.Alpha*normalizedRerank + (1-s.Alpha)*c.Score,
|
||||
Metadata: c.Metadata,
|
||||
OriginalRank: i + 1,
|
||||
}
|
||||
}
|
||||
|
||||
// Sort by combined score (descending)
|
||||
sort.Slice(results, func(i, j int) bool {
|
||||
return results[i].CombinedScore > results[j].CombinedScore
|
||||
})
|
||||
|
||||
// Assign rerank positions and calculate improvement
|
||||
for i := range results {
|
||||
results[i].RerankRank = i + 1
|
||||
results[i].RankImprovement = results[i].OriginalRank - results[i].RerankRank
|
||||
}
|
||||
|
||||
// Apply limit
|
||||
if len(results) > limit {
|
||||
results = results[:limit]
|
||||
}
|
||||
|
||||
log.Debug().
|
||||
Int("candidates", len(candidates)).
|
||||
Int("returned", len(results)).
|
||||
Float64("alpha", s.Alpha).
|
||||
Msg("Cross-encoder reranking completed")
|
||||
|
||||
return results, nil
|
||||
}
|
||||
|
||||
// RerankByScore reranks candidates and returns sorted by pure cross-encoder score.
|
||||
// Useful when you want to completely replace bi-encoder ranking.
|
||||
func (s *Service) RerankByScore(query string, candidates []Candidate, limit int) ([]RerankResult, error) {
|
||||
if len(candidates) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
if limit <= 0 {
|
||||
limit = DefaultResultLimit
|
||||
}
|
||||
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
scores, err := s.scoreAll(query, candidates)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("score candidates: %w", err)
|
||||
}
|
||||
|
||||
results := make([]RerankResult, len(candidates))
|
||||
for i, c := range candidates {
|
||||
normalizedRerank := sigmoid(scores[i])
|
||||
results[i] = RerankResult{
|
||||
ID: c.ID,
|
||||
Content: c.Content,
|
||||
OriginalScore: c.Score,
|
||||
RerankScore: normalizedRerank,
|
||||
CombinedScore: normalizedRerank, // Use pure rerank score
|
||||
Metadata: c.Metadata,
|
||||
OriginalRank: i + 1,
|
||||
}
|
||||
}
|
||||
|
||||
// Sort by rerank score only
|
||||
sort.Slice(results, func(i, j int) bool {
|
||||
return results[i].RerankScore > results[j].RerankScore
|
||||
})
|
||||
|
||||
for i := range results {
|
||||
results[i].RerankRank = i + 1
|
||||
results[i].RankImprovement = results[i].OriginalRank - results[i].RerankRank
|
||||
}
|
||||
|
||||
if len(results) > limit {
|
||||
results = results[:limit]
|
||||
}
|
||||
|
||||
return results, nil
|
||||
}
|
||||
|
||||
// scoreAll scores all query-document pairs using the cross-encoder.
|
||||
// Returns raw logits (before sigmoid normalization).
|
||||
func (s *Service) scoreAll(query string, candidates []Candidate) ([]float64, error) {
|
||||
batchSize := len(candidates)
|
||||
|
||||
// Tokenize all query-document pairs
|
||||
pairs := make([]tokenizer.EncodeInput, batchSize)
|
||||
for i, c := range candidates {
|
||||
// Cross-encoder takes query and document as a pair
|
||||
pairs[i] = tokenizer.NewDualEncodeInput(
|
||||
tokenizer.NewRawInputSequence(query),
|
||||
tokenizer.NewRawInputSequence(c.Content),
|
||||
)
|
||||
}
|
||||
|
||||
encodings, err := s.tk.EncodeBatch(pairs, true)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("tokenize pairs: %w", err)
|
||||
}
|
||||
|
||||
// Find max sequence length
|
||||
seqLength := 0
|
||||
for _, enc := range encodings {
|
||||
if len(enc.Ids) > seqLength {
|
||||
seqLength = len(enc.Ids)
|
||||
}
|
||||
}
|
||||
if seqLength > MaxSequenceLength {
|
||||
seqLength = MaxSequenceLength
|
||||
}
|
||||
|
||||
inputShape := ort.NewShape(int64(batchSize), int64(seqLength))
|
||||
|
||||
// Create input tensors
|
||||
inputIdsData := make([]int64, batchSize*seqLength)
|
||||
attentionMaskData := make([]int64, batchSize*seqLength)
|
||||
tokenTypeIdsData := make([]int64, batchSize*seqLength)
|
||||
|
||||
for b := 0; b < batchSize; b++ {
|
||||
copyLen := len(encodings[b].Ids)
|
||||
if copyLen > seqLength {
|
||||
copyLen = seqLength
|
||||
}
|
||||
for i := 0; i < copyLen; i++ {
|
||||
inputIdsData[b*seqLength+i] = int64(encodings[b].Ids[i])
|
||||
}
|
||||
|
||||
copyLen = len(encodings[b].AttentionMask)
|
||||
if copyLen > seqLength {
|
||||
copyLen = seqLength
|
||||
}
|
||||
for i := 0; i < copyLen; i++ {
|
||||
attentionMaskData[b*seqLength+i] = int64(encodings[b].AttentionMask[i])
|
||||
}
|
||||
|
||||
copyLen = len(encodings[b].TypeIds)
|
||||
if copyLen > seqLength {
|
||||
copyLen = seqLength
|
||||
}
|
||||
for i := 0; i < copyLen; i++ {
|
||||
tokenTypeIdsData[b*seqLength+i] = int64(encodings[b].TypeIds[i])
|
||||
}
|
||||
}
|
||||
|
||||
inputIdsTensor, err := ort.NewTensor(inputShape, inputIdsData)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("create input_ids tensor: %w", err)
|
||||
}
|
||||
defer func() { _ = inputIdsTensor.Destroy() }()
|
||||
|
||||
attentionMaskTensor, err := ort.NewTensor(inputShape, attentionMaskData)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("create attention_mask tensor: %w", err)
|
||||
}
|
||||
defer func() { _ = attentionMaskTensor.Destroy() }()
|
||||
|
||||
tokenTypeIdsTensor, err := ort.NewTensor(inputShape, tokenTypeIdsData)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("create token_type_ids tensor: %w", err)
|
||||
}
|
||||
defer func() { _ = tokenTypeIdsTensor.Destroy() }()
|
||||
|
||||
// Cross-encoder outputs [batch, 1] logits
|
||||
outputShape := ort.NewShape(int64(batchSize), 1)
|
||||
outputTensor, err := ort.NewEmptyTensor[float32](outputShape)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("create output tensor: %w", err)
|
||||
}
|
||||
defer func() { _ = outputTensor.Destroy() }()
|
||||
|
||||
// Run inference
|
||||
inputTensors := []ort.Value{inputIdsTensor, attentionMaskTensor, tokenTypeIdsTensor}
|
||||
outputTensors := []ort.Value{outputTensor}
|
||||
|
||||
if err := s.session.Run(inputTensors, outputTensors); err != nil {
|
||||
return nil, fmt.Errorf("run cross-encoder inference: %w", err)
|
||||
}
|
||||
|
||||
// Extract scores
|
||||
flatOutput := outputTensor.GetData()
|
||||
scores := make([]float64, batchSize)
|
||||
for i := 0; i < batchSize; i++ {
|
||||
scores[i] = float64(flatOutput[i])
|
||||
}
|
||||
|
||||
return scores, nil
|
||||
}
|
||||
|
||||
// Score scores a single query-document pair.
|
||||
// Returns the raw cross-encoder logit and normalized score.
|
||||
func (s *Service) Score(query, document string) (rawScore, normalizedScore float64, err error) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
scores, err := s.scoreAll(query, []Candidate{{Content: document}})
|
||||
if err != nil {
|
||||
return 0, 0, err
|
||||
}
|
||||
|
||||
rawScore = scores[0]
|
||||
normalizedScore = sigmoid(rawScore)
|
||||
return rawScore, normalizedScore, nil
|
||||
}
|
||||
|
||||
// Close releases model resources.
|
||||
func (s *Service) Close() error {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
if s.session != nil {
|
||||
if err := s.session.Destroy(); err != nil {
|
||||
return fmt.Errorf("destroy cross-encoder session: %w", err)
|
||||
}
|
||||
s.session = nil
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// sigmoid applies the sigmoid function to normalize scores to 0-1 range.
|
||||
func sigmoid(x float64) float64 {
|
||||
if x > 20 {
|
||||
return 1.0
|
||||
}
|
||||
if x < -20 {
|
||||
return 0.0
|
||||
}
|
||||
return 1.0 / (1.0 + math.Exp(-x))
|
||||
}
|
||||
@@ -1,448 +0,0 @@
|
||||
package reranking
|
||||
|
||||
import (
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
"github.com/lukaszraczylo/claude-mnemonic/internal/embedding"
|
||||
)
|
||||
|
||||
// initONNX initializes the ONNX runtime via the embedding service.
|
||||
// Must be called before creating reranking service.
|
||||
func initONNX(t *testing.T) func() {
|
||||
t.Helper()
|
||||
embSvc, err := embedding.NewService()
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to initialize ONNX via embedding service: %v", err)
|
||||
}
|
||||
return func() {
|
||||
embSvc.Close()
|
||||
}
|
||||
}
|
||||
|
||||
// TestSigmoid tests the sigmoid normalization function.
|
||||
func TestSigmoid(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
input float64
|
||||
wantMin float64
|
||||
wantMax float64
|
||||
}{
|
||||
{"positive large", 10, 0.9999, 1.0},
|
||||
{"positive small", 1, 0.7, 0.8},
|
||||
{"zero", 0, 0.4999, 0.5001},
|
||||
{"negative small", -1, 0.2, 0.3},
|
||||
{"negative large", -10, 0, 0.0001},
|
||||
{"very positive", 25, 0.999999, 1.0},
|
||||
{"very negative", -25, 0, 0.000001},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
got := sigmoid(tt.input)
|
||||
if got < tt.wantMin || got > tt.wantMax {
|
||||
t.Errorf("sigmoid(%v) = %v, want in range [%v, %v]",
|
||||
tt.input, got, tt.wantMin, tt.wantMax)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestDefaultConfig tests the default configuration values.
|
||||
func TestDefaultConfig(t *testing.T) {
|
||||
cfg := DefaultConfig()
|
||||
|
||||
if cfg.Alpha < 0 || cfg.Alpha > 1 {
|
||||
t.Errorf("DefaultConfig().Alpha = %v, want in range [0, 1]", cfg.Alpha)
|
||||
}
|
||||
if cfg.Alpha != 0.7 {
|
||||
t.Errorf("DefaultConfig().Alpha = %v, want 0.7", cfg.Alpha)
|
||||
}
|
||||
}
|
||||
|
||||
// TestNewService tests service creation.
|
||||
func TestNewService(t *testing.T) {
|
||||
cleanup := initONNX(t)
|
||||
defer cleanup()
|
||||
|
||||
cfg := DefaultConfig()
|
||||
svc, err := NewService(cfg)
|
||||
if err != nil {
|
||||
t.Fatalf("NewService() error = %v", err)
|
||||
}
|
||||
defer svc.Close()
|
||||
|
||||
if svc == nil {
|
||||
t.Fatal("NewService() returned nil")
|
||||
}
|
||||
|
||||
if svc.Alpha != cfg.Alpha {
|
||||
t.Errorf("Service.Alpha = %v, want %v", svc.Alpha, cfg.Alpha)
|
||||
}
|
||||
}
|
||||
|
||||
// TestScore tests single pair scoring.
|
||||
func TestScore(t *testing.T) {
|
||||
cleanup := initONNX(t)
|
||||
defer cleanup()
|
||||
|
||||
cfg := DefaultConfig()
|
||||
svc, err := NewService(cfg)
|
||||
if err != nil {
|
||||
t.Fatalf("NewService() error = %v", err)
|
||||
}
|
||||
defer svc.Close()
|
||||
|
||||
query := "What is the capital of France?"
|
||||
relevant := "Paris is the capital and largest city of France."
|
||||
irrelevant := "Dogs are popular pets known for their loyalty."
|
||||
|
||||
// Score relevant document
|
||||
_, relevantNorm, err := svc.Score(query, relevant)
|
||||
if err != nil {
|
||||
t.Fatalf("Score(relevant) error = %v", err)
|
||||
}
|
||||
|
||||
// Score irrelevant document
|
||||
_, irrelevantNorm, err := svc.Score(query, irrelevant)
|
||||
if err != nil {
|
||||
t.Fatalf("Score(irrelevant) error = %v", err)
|
||||
}
|
||||
|
||||
// Relevant document should score higher
|
||||
if relevantNorm <= irrelevantNorm {
|
||||
t.Errorf("Expected relevant (%v) > irrelevant (%v)",
|
||||
relevantNorm, irrelevantNorm)
|
||||
}
|
||||
}
|
||||
|
||||
// TestRerank tests the reranking functionality.
|
||||
func TestRerank(t *testing.T) {
|
||||
cleanup := initONNX(t)
|
||||
defer cleanup()
|
||||
|
||||
cfg := DefaultConfig()
|
||||
svc, err := NewService(cfg)
|
||||
if err != nil {
|
||||
t.Fatalf("NewService() error = %v", err)
|
||||
}
|
||||
defer svc.Close()
|
||||
|
||||
query := "How to handle errors in Go?"
|
||||
candidates := []Candidate{
|
||||
{
|
||||
ID: "1",
|
||||
Content: "Python exception handling with try/except blocks.",
|
||||
Score: 0.8, // High bi-encoder score but irrelevant
|
||||
},
|
||||
{
|
||||
ID: "2",
|
||||
Content: "Go error handling uses explicit return values. Functions return error as the last value.",
|
||||
Score: 0.6, // Lower bi-encoder score but relevant
|
||||
},
|
||||
{
|
||||
ID: "3",
|
||||
Content: "JavaScript uses Promise.catch for async error handling.",
|
||||
Score: 0.7,
|
||||
},
|
||||
}
|
||||
|
||||
results, err := svc.Rerank(query, candidates, 3)
|
||||
if err != nil {
|
||||
t.Fatalf("Rerank() error = %v", err)
|
||||
}
|
||||
|
||||
if len(results) != 3 {
|
||||
t.Fatalf("Rerank() returned %d results, want 3", len(results))
|
||||
}
|
||||
|
||||
// The Go error handling document should rank higher after reranking
|
||||
var goRank int
|
||||
for i, r := range results {
|
||||
if r.ID == "2" {
|
||||
goRank = i + 1
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
if goRank == 0 {
|
||||
t.Error("Go document not found in results")
|
||||
}
|
||||
|
||||
// Verify all results have required fields populated
|
||||
for i, r := range results {
|
||||
if r.ID == "" {
|
||||
t.Errorf("Result %d has empty ID", i)
|
||||
}
|
||||
if r.Content == "" {
|
||||
t.Errorf("Result %d has empty Content", i)
|
||||
}
|
||||
if r.RerankRank != i+1 {
|
||||
t.Errorf("Result %d has RerankRank %d, want %d", i, r.RerankRank, i+1)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestRerankEmpty tests reranking with empty input.
|
||||
func TestRerankEmpty(t *testing.T) {
|
||||
cleanup := initONNX(t)
|
||||
defer cleanup()
|
||||
|
||||
cfg := DefaultConfig()
|
||||
svc, err := NewService(cfg)
|
||||
if err != nil {
|
||||
t.Fatalf("NewService() error = %v", err)
|
||||
}
|
||||
defer svc.Close()
|
||||
|
||||
results, err := svc.Rerank("test query", nil, 10)
|
||||
if err != nil {
|
||||
t.Fatalf("Rerank(nil) error = %v", err)
|
||||
}
|
||||
|
||||
if results != nil {
|
||||
t.Errorf("Rerank(nil) = %v, want nil", results)
|
||||
}
|
||||
|
||||
results, err = svc.Rerank("test query", []Candidate{}, 10)
|
||||
if err != nil {
|
||||
t.Fatalf("Rerank([]) error = %v", err)
|
||||
}
|
||||
|
||||
if results != nil {
|
||||
t.Errorf("Rerank([]) = %v, want nil", results)
|
||||
}
|
||||
}
|
||||
|
||||
// TestRerankLimit tests that limit is respected.
|
||||
func TestRerankLimit(t *testing.T) {
|
||||
cleanup := initONNX(t)
|
||||
defer cleanup()
|
||||
|
||||
cfg := DefaultConfig()
|
||||
svc, err := NewService(cfg)
|
||||
if err != nil {
|
||||
t.Fatalf("NewService() error = %v", err)
|
||||
}
|
||||
defer svc.Close()
|
||||
|
||||
candidates := make([]Candidate, 20)
|
||||
for i := range candidates {
|
||||
candidates[i] = Candidate{
|
||||
ID: string(rune('A' + i)),
|
||||
Content: "Test document content for ranking.",
|
||||
Score: 0.5,
|
||||
}
|
||||
}
|
||||
|
||||
results, err := svc.Rerank("test query", candidates, 5)
|
||||
if err != nil {
|
||||
t.Fatalf("Rerank() error = %v", err)
|
||||
}
|
||||
|
||||
if len(results) != 5 {
|
||||
t.Errorf("Rerank() returned %d results, want 5", len(results))
|
||||
}
|
||||
}
|
||||
|
||||
// TestRerankByScore tests pure cross-encoder ranking.
|
||||
func TestRerankByScore(t *testing.T) {
|
||||
cleanup := initONNX(t)
|
||||
defer cleanup()
|
||||
|
||||
cfg := DefaultConfig()
|
||||
svc, err := NewService(cfg)
|
||||
if err != nil {
|
||||
t.Fatalf("NewService() error = %v", err)
|
||||
}
|
||||
defer svc.Close()
|
||||
|
||||
query := "machine learning algorithms"
|
||||
candidates := []Candidate{
|
||||
{
|
||||
ID: "1",
|
||||
Content: "Cooking recipes for Italian pasta dishes.",
|
||||
Score: 0.9, // High original score
|
||||
},
|
||||
{
|
||||
ID: "2",
|
||||
Content: "Neural networks are a type of machine learning algorithm.",
|
||||
Score: 0.3, // Low original score
|
||||
},
|
||||
}
|
||||
|
||||
results, err := svc.RerankByScore(query, candidates, 2)
|
||||
if err != nil {
|
||||
t.Fatalf("RerankByScore() error = %v", err)
|
||||
}
|
||||
|
||||
// Document 2 should rank first since it's about ML
|
||||
if results[0].ID != "2" {
|
||||
t.Errorf("Expected ML document to rank first, got %v", results[0].ID)
|
||||
}
|
||||
|
||||
// CombinedScore should equal RerankScore when using RerankByScore
|
||||
for _, r := range results {
|
||||
if r.CombinedScore != r.RerankScore {
|
||||
t.Errorf("RerankByScore: CombinedScore (%v) != RerankScore (%v)",
|
||||
r.CombinedScore, r.RerankScore)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestRankImprovement tests that rank improvement is calculated correctly.
|
||||
func TestRankImprovement(t *testing.T) {
|
||||
cleanup := initONNX(t)
|
||||
defer cleanup()
|
||||
|
||||
cfg := DefaultConfig()
|
||||
svc, err := NewService(cfg)
|
||||
if err != nil {
|
||||
t.Fatalf("NewService() error = %v", err)
|
||||
}
|
||||
defer svc.Close()
|
||||
|
||||
// Create candidates where we know the expected reranking
|
||||
candidates := []Candidate{
|
||||
{ID: "A", Content: "Unrelated content about weather forecasting.", Score: 0.9},
|
||||
{ID: "B", Content: "How to fix memory leaks in Go programs.", Score: 0.8},
|
||||
{ID: "C", Content: "More unrelated content about gardening tips.", Score: 0.7},
|
||||
}
|
||||
|
||||
results, err := svc.Rerank("debugging memory issues in Go", candidates, 3)
|
||||
if err != nil {
|
||||
t.Fatalf("Rerank() error = %v", err)
|
||||
}
|
||||
|
||||
for _, r := range results {
|
||||
// RankImprovement = OriginalRank - RerankRank
|
||||
// Positive means moved up, negative means moved down
|
||||
expectedImprovement := r.OriginalRank - r.RerankRank
|
||||
if r.RankImprovement != expectedImprovement {
|
||||
t.Errorf("ID %s: RankImprovement = %d, want %d (orig=%d, new=%d)",
|
||||
r.ID, r.RankImprovement, expectedImprovement,
|
||||
r.OriginalRank, r.RerankRank)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestConcurrentRerank tests concurrent reranking calls.
|
||||
func TestConcurrentRerank(t *testing.T) {
|
||||
cleanup := initONNX(t)
|
||||
defer cleanup()
|
||||
|
||||
cfg := DefaultConfig()
|
||||
svc, err := NewService(cfg)
|
||||
if err != nil {
|
||||
t.Fatalf("NewService() error = %v", err)
|
||||
}
|
||||
defer svc.Close()
|
||||
|
||||
candidates := []Candidate{
|
||||
{ID: "1", Content: "Test document one.", Score: 0.5},
|
||||
{ID: "2", Content: "Test document two.", Score: 0.5},
|
||||
}
|
||||
|
||||
var wg sync.WaitGroup
|
||||
errors := make(chan error, 10)
|
||||
|
||||
for i := 0; i < 10; i++ {
|
||||
wg.Add(1)
|
||||
go func(idx int) {
|
||||
defer wg.Done()
|
||||
_, err := svc.Rerank("concurrent test query", candidates, 2)
|
||||
if err != nil {
|
||||
errors <- err
|
||||
}
|
||||
}(i)
|
||||
}
|
||||
|
||||
wg.Wait()
|
||||
close(errors)
|
||||
|
||||
for err := range errors {
|
||||
t.Errorf("Concurrent Rerank error: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestClose tests service cleanup.
|
||||
func TestClose(t *testing.T) {
|
||||
cleanup := initONNX(t)
|
||||
defer cleanup()
|
||||
|
||||
cfg := DefaultConfig()
|
||||
svc, err := NewService(cfg)
|
||||
if err != nil {
|
||||
t.Fatalf("NewService() error = %v", err)
|
||||
}
|
||||
|
||||
err = svc.Close()
|
||||
if err != nil {
|
||||
t.Errorf("Close() error = %v", err)
|
||||
}
|
||||
|
||||
// Double close should not panic
|
||||
err = svc.Close()
|
||||
if err != nil {
|
||||
t.Errorf("Close() on closed service error = %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestMetadataPreserved tests that metadata is preserved through reranking.
|
||||
func TestMetadataPreserved(t *testing.T) {
|
||||
cleanup := initONNX(t)
|
||||
defer cleanup()
|
||||
|
||||
cfg := DefaultConfig()
|
||||
svc, err := NewService(cfg)
|
||||
if err != nil {
|
||||
t.Fatalf("NewService() error = %v", err)
|
||||
}
|
||||
defer svc.Close()
|
||||
|
||||
candidates := []Candidate{
|
||||
{
|
||||
ID: "1",
|
||||
Content: "Test content.",
|
||||
Score: 0.5,
|
||||
Metadata: map[string]any{"custom": "value1", "num": 42},
|
||||
},
|
||||
{
|
||||
ID: "2",
|
||||
Content: "Another test.",
|
||||
Score: 0.5,
|
||||
Metadata: map[string]any{"custom": "value2"},
|
||||
},
|
||||
}
|
||||
|
||||
results, err := svc.Rerank("query", candidates, 2)
|
||||
if err != nil {
|
||||
t.Fatalf("Rerank() error = %v", err)
|
||||
}
|
||||
|
||||
for _, r := range results {
|
||||
if r.Metadata == nil {
|
||||
t.Errorf("Result %s has nil metadata", r.ID)
|
||||
continue
|
||||
}
|
||||
|
||||
// Find original candidate
|
||||
var original *Candidate
|
||||
for i := range candidates {
|
||||
if candidates[i].ID == r.ID {
|
||||
original = &candidates[i]
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
if original == nil {
|
||||
t.Errorf("Could not find original for result %s", r.ID)
|
||||
continue
|
||||
}
|
||||
|
||||
// Check metadata preserved
|
||||
if original.Metadata["custom"] != r.Metadata["custom"] {
|
||||
t.Errorf("Metadata not preserved for %s", r.ID)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,139 +0,0 @@
|
||||
// Package scoring provides importance score calculation for observations.
|
||||
package scoring
|
||||
|
||||
import (
|
||||
"math"
|
||||
"time"
|
||||
|
||||
"github.com/lukaszraczylo/claude-mnemonic/pkg/models"
|
||||
)
|
||||
|
||||
// Calculator computes importance scores for observations.
|
||||
type Calculator struct {
|
||||
config *models.ScoringConfig
|
||||
}
|
||||
|
||||
// NewCalculator creates a new scoring calculator.
|
||||
// If config is nil, uses the default configuration.
|
||||
func NewCalculator(config *models.ScoringConfig) *Calculator {
|
||||
if config == nil {
|
||||
config = models.DefaultScoringConfig()
|
||||
}
|
||||
return &Calculator{config: config}
|
||||
}
|
||||
|
||||
// Calculate computes the importance score for an observation at the given time.
|
||||
//
|
||||
// The scoring formula:
|
||||
//
|
||||
// FinalScore = (BaseScore × TypeWeight × RecencyDecay) + FeedbackContrib + ConceptContrib + RetrievalContrib
|
||||
//
|
||||
// Where:
|
||||
// - BaseScore = 1.0
|
||||
// - TypeWeight = observation type multiplier (e.g., bugfix=1.3, change=0.9)
|
||||
// - RecencyDecay = 0.5^(age_days / half_life_days) - halves every 7 days by default
|
||||
// - FeedbackContrib = user_feedback × feedback_weight
|
||||
// - 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)
|
||||
|
||||
// 2. Calculate recency decay: 0.5^(age_days / half_life_days)
|
||||
ageDays := now.Sub(time.UnixMilli(obs.CreatedAtEpoch)).Hours() / 24.0
|
||||
if ageDays < 0 {
|
||||
ageDays = 0 // Handle future timestamps gracefully
|
||||
}
|
||||
recencyDecay := math.Pow(0.5, ageDays/c.config.RecencyHalfLifeDays)
|
||||
|
||||
// Core score = 1.0 × type_weight × recency_decay
|
||||
coreScore := 1.0 * typeWeight * recencyDecay
|
||||
|
||||
// 3. User feedback contribution: feedback × weight
|
||||
feedbackContrib := float64(obs.UserFeedback) * c.config.FeedbackWeight
|
||||
|
||||
// 4. Concept boost contribution: sum of matching concept weights × factor
|
||||
conceptBoost := 0.0
|
||||
for _, concept := range obs.Concepts {
|
||||
if weight, ok := c.config.ConceptWeights[concept]; ok {
|
||||
conceptBoost += weight
|
||||
}
|
||||
}
|
||||
conceptContrib := conceptBoost * c.config.ConceptWeight
|
||||
|
||||
// 5. Retrieval boost: log2(count + 1) × 0.1 × weight (diminishing returns)
|
||||
retrievalContrib := 0.0
|
||||
if obs.RetrievalCount > 0 {
|
||||
// log2(count + 1) gives diminishing returns: 1→1, 3→2, 7→3, 15→4, etc.
|
||||
retrievalBoost := math.Log2(float64(obs.RetrievalCount)+1) * 0.1
|
||||
retrievalContrib = retrievalBoost * c.config.RetrievalWeight
|
||||
}
|
||||
|
||||
// Final score with minimum threshold
|
||||
finalScore := coreScore + feedbackContrib + conceptContrib + retrievalContrib
|
||||
if finalScore < c.config.MinScore {
|
||||
finalScore = c.config.MinScore
|
||||
}
|
||||
|
||||
return ScoreComponents{
|
||||
TypeWeight: typeWeight,
|
||||
RecencyDecay: recencyDecay,
|
||||
CoreScore: coreScore,
|
||||
FeedbackContrib: feedbackContrib,
|
||||
ConceptContrib: conceptContrib,
|
||||
RetrievalContrib: retrievalContrib,
|
||||
FinalScore: finalScore,
|
||||
AgeDays: ageDays,
|
||||
}
|
||||
}
|
||||
|
||||
// ScoreComponents contains the breakdown of an importance score calculation.
|
||||
type ScoreComponents struct {
|
||||
TypeWeight float64 `json:"type_weight"`
|
||||
RecencyDecay float64 `json:"recency_decay"`
|
||||
CoreScore float64 `json:"core_score"`
|
||||
FeedbackContrib float64 `json:"feedback_contrib"`
|
||||
ConceptContrib float64 `json:"concept_contrib"`
|
||||
RetrievalContrib float64 `json:"retrieval_contrib"`
|
||||
FinalScore float64 `json:"final_score"`
|
||||
AgeDays float64 `json:"age_days"`
|
||||
}
|
||||
|
||||
// BatchCalculate computes scores for multiple observations.
|
||||
// Returns a map of observation ID to calculated score.
|
||||
func (c *Calculator) BatchCalculate(observations []*models.Observation, now time.Time) map[int64]float64 {
|
||||
scores := make(map[int64]float64, len(observations))
|
||||
for _, obs := range observations {
|
||||
scores[obs.ID] = c.Calculate(obs, now)
|
||||
}
|
||||
return scores
|
||||
}
|
||||
|
||||
// RecalculateThreshold returns the minimum duration before an observation
|
||||
// should have its score recalculated. This prevents excessive recalculation
|
||||
// while ensuring scores stay reasonably fresh.
|
||||
func (c *Calculator) RecalculateThreshold() time.Duration {
|
||||
// Recalculate at most every 6 hours
|
||||
// This balances freshness with performance
|
||||
return 6 * time.Hour
|
||||
}
|
||||
|
||||
// UpdateConfig updates the calculator's scoring configuration.
|
||||
// This allows runtime tuning of scoring parameters.
|
||||
func (c *Calculator) UpdateConfig(config *models.ScoringConfig) {
|
||||
if config != nil {
|
||||
c.config = config
|
||||
}
|
||||
}
|
||||
|
||||
// GetConfig returns the current scoring configuration.
|
||||
func (c *Calculator) GetConfig() *models.ScoringConfig {
|
||||
return c.config
|
||||
}
|
||||
@@ -1,638 +0,0 @@
|
||||
// Package scoring provides importance score calculation for observations.
|
||||
package scoring
|
||||
|
||||
import (
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/lukaszraczylo/claude-mnemonic/pkg/models"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"github.com/stretchr/testify/suite"
|
||||
)
|
||||
|
||||
// CalculatorSuite is a test suite for the Calculator.
|
||||
type CalculatorSuite struct {
|
||||
suite.Suite
|
||||
calc *Calculator
|
||||
config *models.ScoringConfig
|
||||
now time.Time
|
||||
}
|
||||
|
||||
func (s *CalculatorSuite) SetupTest() {
|
||||
s.config = models.DefaultScoringConfig()
|
||||
s.calc = NewCalculator(s.config)
|
||||
s.now = time.Date(2025, 1, 15, 12, 0, 0, 0, time.UTC)
|
||||
}
|
||||
|
||||
func TestCalculatorSuite(t *testing.T) {
|
||||
suite.Run(t, new(CalculatorSuite))
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// GOOD SCENARIOS - Expected normal operations
|
||||
// =============================================================================
|
||||
|
||||
func (s *CalculatorSuite) TestCalculate_GoodScenarios_NewObservation() {
|
||||
// A brand new observation should have score close to type weight
|
||||
obs := &models.Observation{
|
||||
ID: 1,
|
||||
Type: models.ObsTypeBugfix,
|
||||
CreatedAtEpoch: s.now.UnixMilli(),
|
||||
}
|
||||
|
||||
score := s.calc.Calculate(obs, s.now)
|
||||
|
||||
// Expected: 1.0 × 1.3 (bugfix weight) × 1.0 (no decay) = 1.3
|
||||
s.InDelta(1.3, score, 0.01, "new bugfix should score ~1.3")
|
||||
}
|
||||
|
||||
func (s *CalculatorSuite) TestCalculate_GoodScenarios_OneWeekOld() {
|
||||
// One week old observation should have half the recency score
|
||||
obs := &models.Observation{
|
||||
ID: 1,
|
||||
Type: models.ObsTypeDiscovery,
|
||||
CreatedAtEpoch: s.now.Add(-7 * 24 * time.Hour).UnixMilli(),
|
||||
}
|
||||
|
||||
score := s.calc.Calculate(obs, s.now)
|
||||
|
||||
// Expected: 1.0 × 1.1 (discovery) × 0.5 (7 days half-life) = 0.55
|
||||
s.InDelta(0.55, score, 0.05, "7-day old discovery should score ~0.55")
|
||||
}
|
||||
|
||||
func (s *CalculatorSuite) TestCalculate_GoodScenarios_TwoWeeksOld() {
|
||||
// Two weeks old should have 1/4 recency score
|
||||
obs := &models.Observation{
|
||||
ID: 1,
|
||||
Type: models.ObsTypeFeature,
|
||||
CreatedAtEpoch: s.now.Add(-14 * 24 * time.Hour).UnixMilli(),
|
||||
}
|
||||
|
||||
score := s.calc.Calculate(obs, s.now)
|
||||
|
||||
// Expected: 1.0 × 1.2 (feature) × 0.25 (14 days = 2 half-lives) = 0.30
|
||||
s.InDelta(0.30, score, 0.05, "14-day old feature should score ~0.30")
|
||||
}
|
||||
|
||||
func (s *CalculatorSuite) TestCalculate_GoodScenarios_PositiveFeedback() {
|
||||
// Positive feedback should boost score
|
||||
obs := &models.Observation{
|
||||
ID: 1,
|
||||
Type: models.ObsTypeChange,
|
||||
CreatedAtEpoch: s.now.UnixMilli(),
|
||||
UserFeedback: 1, // thumbs up
|
||||
}
|
||||
|
||||
score := s.calc.Calculate(obs, s.now)
|
||||
|
||||
// Expected: (1.0 × 0.9) + 0.30 (feedback) = 1.20
|
||||
s.InDelta(1.20, score, 0.01, "thumbs up should boost score by 0.30")
|
||||
}
|
||||
|
||||
func (s *CalculatorSuite) TestCalculate_GoodScenarios_NegativeFeedback() {
|
||||
// Negative feedback should reduce score
|
||||
obs := &models.Observation{
|
||||
ID: 1,
|
||||
Type: models.ObsTypeChange,
|
||||
CreatedAtEpoch: s.now.UnixMilli(),
|
||||
UserFeedback: -1, // thumbs down
|
||||
}
|
||||
|
||||
score := s.calc.Calculate(obs, s.now)
|
||||
|
||||
// Expected: (1.0 × 0.9) - 0.30 (feedback) = 0.60
|
||||
s.InDelta(0.60, score, 0.01, "thumbs down should reduce score by 0.30")
|
||||
}
|
||||
|
||||
func (s *CalculatorSuite) TestCalculate_GoodScenarios_WithConcepts() {
|
||||
// Observation with valuable concepts should get boost
|
||||
obs := &models.Observation{
|
||||
ID: 1,
|
||||
Type: models.ObsTypeBugfix,
|
||||
CreatedAtEpoch: s.now.UnixMilli(),
|
||||
Concepts: []string{"security", "gotcha"},
|
||||
}
|
||||
|
||||
score := s.calc.Calculate(obs, s.now)
|
||||
|
||||
// Concept boost: (0.30 + 0.25) × 0.20 = 0.11
|
||||
// Expected: 1.3 + 0.11 = 1.41
|
||||
s.InDelta(1.41, score, 0.05, "security+gotcha concepts should boost score")
|
||||
}
|
||||
|
||||
func (s *CalculatorSuite) TestCalculate_GoodScenarios_WithRetrievals() {
|
||||
// Popular observations should get retrieval boost
|
||||
obs := &models.Observation{
|
||||
ID: 1,
|
||||
Type: models.ObsTypeDiscovery,
|
||||
CreatedAtEpoch: s.now.UnixMilli(),
|
||||
RetrievalCount: 7, // log2(8) = 3
|
||||
}
|
||||
|
||||
score := s.calc.Calculate(obs, s.now)
|
||||
|
||||
// Retrieval boost: log2(7+1) × 0.1 × 0.15 = 3 × 0.1 × 0.15 = 0.045
|
||||
// Expected: 1.1 + 0.045 ≈ 1.145
|
||||
s.InDelta(1.145, score, 0.05, "7 retrievals should add small boost")
|
||||
}
|
||||
|
||||
func (s *CalculatorSuite) TestCalculate_GoodScenarios_CombinedFactors() {
|
||||
// Test with all factors combined
|
||||
obs := &models.Observation{
|
||||
ID: 1,
|
||||
Type: models.ObsTypeBugfix,
|
||||
CreatedAtEpoch: s.now.Add(-7 * 24 * time.Hour).UnixMilli(), // 7 days old
|
||||
UserFeedback: 1,
|
||||
Concepts: []string{"security"},
|
||||
RetrievalCount: 3,
|
||||
}
|
||||
|
||||
score := s.calc.Calculate(obs, s.now)
|
||||
|
||||
// Core: 1.0 × 1.3 × 0.5 = 0.65
|
||||
// Feedback: 0.30
|
||||
// Concept: 0.30 × 0.20 = 0.06
|
||||
// Retrieval: log2(4) × 0.1 × 0.15 = 2 × 0.1 × 0.15 = 0.03
|
||||
// Total ≈ 1.04
|
||||
s.InDelta(1.04, score, 0.1, "combined factors should result in ~1.04")
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// WORSE SCENARIOS - Degraded but acceptable operations
|
||||
// =============================================================================
|
||||
|
||||
func (s *CalculatorSuite) TestCalculate_WorseScenarios_VeryOldObservation() {
|
||||
// Very old observation should have low but non-zero score
|
||||
obs := &models.Observation{
|
||||
ID: 1,
|
||||
Type: models.ObsTypeChange,
|
||||
CreatedAtEpoch: s.now.Add(-90 * 24 * time.Hour).UnixMilli(), // 90 days old
|
||||
}
|
||||
|
||||
score := s.calc.Calculate(obs, s.now)
|
||||
|
||||
// 90 days = ~12.86 half-lives → decay ≈ 0.00014
|
||||
// Core: 1.0 × 0.9 × 0.00014 = 0.000126
|
||||
// But minimum score is 0.01
|
||||
s.GreaterOrEqual(score, 0.01, "very old observation should still meet minimum")
|
||||
s.Less(score, 0.1, "very old observation should be low scoring")
|
||||
}
|
||||
|
||||
func (s *CalculatorSuite) TestCalculate_WorseScenarios_NegativeFeedbackOld() {
|
||||
// Old observation with negative feedback should still have minimum score
|
||||
obs := &models.Observation{
|
||||
ID: 1,
|
||||
Type: models.ObsTypeChange,
|
||||
CreatedAtEpoch: s.now.Add(-60 * 24 * time.Hour).UnixMilli(),
|
||||
UserFeedback: -1,
|
||||
}
|
||||
|
||||
score := s.calc.Calculate(obs, s.now)
|
||||
|
||||
s.GreaterOrEqual(score, s.config.MinScore, "should never go below minimum score")
|
||||
}
|
||||
|
||||
func (s *CalculatorSuite) TestCalculate_WorseScenarios_UnknownConcepts() {
|
||||
// Unknown concepts should not affect score negatively
|
||||
obs := &models.Observation{
|
||||
ID: 1,
|
||||
Type: models.ObsTypeDiscovery,
|
||||
CreatedAtEpoch: s.now.UnixMilli(),
|
||||
Concepts: []string{"unknown-concept", "another-unknown"},
|
||||
}
|
||||
|
||||
score := s.calc.Calculate(obs, s.now)
|
||||
|
||||
// Should just be the base score without concept boost
|
||||
s.InDelta(1.1, score, 0.01, "unknown concepts should not affect score")
|
||||
}
|
||||
|
||||
func (s *CalculatorSuite) TestCalculate_WorseScenarios_MixedConcepts() {
|
||||
// Mix of known and unknown concepts
|
||||
obs := &models.Observation{
|
||||
ID: 1,
|
||||
Type: models.ObsTypeDiscovery,
|
||||
CreatedAtEpoch: s.now.UnixMilli(),
|
||||
Concepts: []string{"security", "unknown-concept"},
|
||||
}
|
||||
|
||||
score := s.calc.Calculate(obs, s.now)
|
||||
|
||||
// Only security should contribute
|
||||
// Expected: 1.1 + (0.30 × 0.20) = 1.16
|
||||
s.InDelta(1.16, score, 0.05, "only known concepts should boost score")
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// BAD SCENARIOS - Edge cases and error conditions
|
||||
// =============================================================================
|
||||
|
||||
func (s *CalculatorSuite) TestCalculate_BadScenarios_FutureTimestamp() {
|
||||
// Observation created in the future (clock skew)
|
||||
obs := &models.Observation{
|
||||
ID: 1,
|
||||
Type: models.ObsTypeBugfix,
|
||||
CreatedAtEpoch: s.now.Add(24 * time.Hour).UnixMilli(), // 1 day in future
|
||||
}
|
||||
|
||||
score := s.calc.Calculate(obs, s.now)
|
||||
|
||||
// Should handle gracefully - age should be 0
|
||||
s.InDelta(1.3, score, 0.01, "future timestamp should be treated as now")
|
||||
}
|
||||
|
||||
func (s *CalculatorSuite) TestCalculate_BadScenarios_ZeroEpoch() {
|
||||
// Missing creation timestamp
|
||||
obs := &models.Observation{
|
||||
ID: 1,
|
||||
Type: models.ObsTypeDiscovery,
|
||||
CreatedAtEpoch: 0, // Missing timestamp
|
||||
}
|
||||
|
||||
score := s.calc.Calculate(obs, s.now)
|
||||
|
||||
// This will be treated as very old (1970)
|
||||
s.GreaterOrEqual(score, s.config.MinScore, "should still meet minimum")
|
||||
}
|
||||
|
||||
func (s *CalculatorSuite) TestCalculate_BadScenarios_EmptyObservation() {
|
||||
// Minimal observation with defaults
|
||||
obs := &models.Observation{
|
||||
ID: 1,
|
||||
Type: "", // Empty type
|
||||
CreatedAtEpoch: s.now.UnixMilli(),
|
||||
}
|
||||
|
||||
score := s.calc.Calculate(obs, s.now)
|
||||
|
||||
// Unknown type should default to 1.0 weight
|
||||
s.InDelta(1.0, score, 0.01, "empty type should use default weight 1.0")
|
||||
}
|
||||
|
||||
func (s *CalculatorSuite) TestCalculate_BadScenarios_ExtremeRetrievalCount() {
|
||||
// Very high retrieval count
|
||||
obs := &models.Observation{
|
||||
ID: 1,
|
||||
Type: models.ObsTypeDiscovery,
|
||||
CreatedAtEpoch: s.now.UnixMilli(),
|
||||
RetrievalCount: 1000000, // Extreme value
|
||||
}
|
||||
|
||||
score := s.calc.Calculate(obs, s.now)
|
||||
|
||||
// log2(1000001) ≈ 19.93, so boost = 19.93 × 0.1 × 0.15 ≈ 0.30
|
||||
// Score should be reasonable, not exploding
|
||||
s.Less(score, 2.0, "extreme retrieval count should not explode score")
|
||||
}
|
||||
|
||||
func (s *CalculatorSuite) TestCalculate_BadScenarios_NegativeRetrievalCount() {
|
||||
// Negative retrieval count (should not happen but test defensively)
|
||||
obs := &models.Observation{
|
||||
ID: 1,
|
||||
Type: models.ObsTypeDiscovery,
|
||||
CreatedAtEpoch: s.now.UnixMilli(),
|
||||
RetrievalCount: -5,
|
||||
}
|
||||
|
||||
score := s.calc.Calculate(obs, s.now)
|
||||
|
||||
// Should not panic and should give base score
|
||||
s.InDelta(1.1, score, 0.01, "negative retrieval should be ignored")
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// EDGE CASES - Boundary conditions
|
||||
// =============================================================================
|
||||
|
||||
func (s *CalculatorSuite) TestCalculate_EdgeCases_ExactlyOneHalfLife() {
|
||||
obs := &models.Observation{
|
||||
ID: 1,
|
||||
Type: models.ObsTypeChange, // 0.9 weight
|
||||
CreatedAtEpoch: s.now.Add(-7 * 24 * time.Hour).UnixMilli(),
|
||||
}
|
||||
|
||||
score := s.calc.Calculate(obs, s.now)
|
||||
s.InDelta(0.45, score, 0.01, "exactly 7 days should give 0.5 decay")
|
||||
}
|
||||
|
||||
func (s *CalculatorSuite) TestCalculate_EdgeCases_ExactlyTwoHalfLives() {
|
||||
obs := &models.Observation{
|
||||
ID: 1,
|
||||
Type: models.ObsTypeChange,
|
||||
CreatedAtEpoch: s.now.Add(-14 * 24 * time.Hour).UnixMilli(),
|
||||
}
|
||||
|
||||
score := s.calc.Calculate(obs, s.now)
|
||||
s.InDelta(0.225, score, 0.01, "exactly 14 days should give 0.25 decay")
|
||||
}
|
||||
|
||||
func (s *CalculatorSuite) TestCalculate_EdgeCases_AllTypeWeights() {
|
||||
types := []struct {
|
||||
t models.ObservationType
|
||||
weight float64
|
||||
}{
|
||||
{models.ObsTypeBugfix, 1.3},
|
||||
{models.ObsTypeFeature, 1.2},
|
||||
{models.ObsTypeDiscovery, 1.1},
|
||||
{models.ObsTypeDecision, 1.1},
|
||||
{models.ObsTypeRefactor, 1.0},
|
||||
{models.ObsTypeChange, 0.9},
|
||||
}
|
||||
|
||||
for _, tt := range types {
|
||||
s.Run(string(tt.t), func() {
|
||||
obs := &models.Observation{
|
||||
ID: 1,
|
||||
Type: tt.t,
|
||||
CreatedAtEpoch: s.now.UnixMilli(),
|
||||
}
|
||||
score := s.calc.Calculate(obs, s.now)
|
||||
s.InDelta(tt.weight, score, 0.01)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func (s *CalculatorSuite) TestCalculate_EdgeCases_MinimumScoreEnforced() {
|
||||
// Create worst case scenario
|
||||
obs := &models.Observation{
|
||||
ID: 1,
|
||||
Type: models.ObsTypeChange, // Lowest weight 0.9
|
||||
CreatedAtEpoch: time.Date(2000, 1, 1, 0, 0, 0, 0, time.UTC).UnixMilli(), // Very old
|
||||
UserFeedback: -1, // Negative feedback
|
||||
}
|
||||
|
||||
score := s.calc.Calculate(obs, s.now)
|
||||
|
||||
s.Equal(s.config.MinScore, score, "should be exactly minimum score")
|
||||
}
|
||||
|
||||
func (s *CalculatorSuite) TestCalculate_EdgeCases_AllConceptsMaxWeight() {
|
||||
// Observation with all high-value concepts
|
||||
obs := &models.Observation{
|
||||
ID: 1,
|
||||
Type: models.ObsTypeBugfix,
|
||||
CreatedAtEpoch: s.now.UnixMilli(),
|
||||
Concepts: []string{"security", "gotcha", "best-practice", "anti-pattern"},
|
||||
}
|
||||
|
||||
score := s.calc.Calculate(obs, s.now)
|
||||
|
||||
// security=0.30, gotcha=0.25, best-practice=0.20, anti-pattern=0.20 = 0.95
|
||||
// Concept contrib: 0.95 × 0.20 = 0.19
|
||||
// Total: 1.3 + 0.19 = 1.49
|
||||
s.InDelta(1.49, score, 0.05, "all high-value concepts should boost significantly")
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// CALCULATE COMPONENTS TESTS
|
||||
// =============================================================================
|
||||
|
||||
func (s *CalculatorSuite) TestCalculateComponents_ReturnsAllComponents() {
|
||||
obs := &models.Observation{
|
||||
ID: 1,
|
||||
Type: models.ObsTypeBugfix,
|
||||
CreatedAtEpoch: s.now.Add(-7 * 24 * time.Hour).UnixMilli(),
|
||||
UserFeedback: 1,
|
||||
Concepts: []string{"security"},
|
||||
RetrievalCount: 7,
|
||||
}
|
||||
|
||||
components := s.calc.CalculateComponents(obs, s.now)
|
||||
|
||||
s.InDelta(1.3, components.TypeWeight, 0.01)
|
||||
s.InDelta(0.5, components.RecencyDecay, 0.01)
|
||||
s.InDelta(0.65, components.CoreScore, 0.05)
|
||||
s.InDelta(0.30, components.FeedbackContrib, 0.01)
|
||||
s.InDelta(0.06, components.ConceptContrib, 0.02)
|
||||
s.Greater(components.RetrievalContrib, 0.0)
|
||||
s.InDelta(7.0, components.AgeDays, 0.1)
|
||||
s.Greater(components.FinalScore, 0.0)
|
||||
}
|
||||
|
||||
func (s *CalculatorSuite) TestCalculateComponents_MatchesCalculate() {
|
||||
obs := &models.Observation{
|
||||
ID: 1,
|
||||
Type: models.ObsTypeFeature,
|
||||
CreatedAtEpoch: s.now.Add(-3 * 24 * time.Hour).UnixMilli(),
|
||||
UserFeedback: -1,
|
||||
Concepts: []string{"performance", "architecture"},
|
||||
RetrievalCount: 15,
|
||||
}
|
||||
|
||||
score := s.calc.Calculate(obs, s.now)
|
||||
components := s.calc.CalculateComponents(obs, s.now)
|
||||
|
||||
s.InDelta(score, components.FinalScore, 0.001, "Calculate and CalculateComponents should match")
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// BATCH CALCULATE TESTS
|
||||
// =============================================================================
|
||||
|
||||
func (s *CalculatorSuite) TestBatchCalculate_Empty() {
|
||||
scores := s.calc.BatchCalculate(nil, s.now)
|
||||
s.Empty(scores)
|
||||
|
||||
scores = s.calc.BatchCalculate([]*models.Observation{}, s.now)
|
||||
s.Empty(scores)
|
||||
}
|
||||
|
||||
func (s *CalculatorSuite) TestBatchCalculate_Multiple() {
|
||||
obs := []*models.Observation{
|
||||
{ID: 1, Type: models.ObsTypeBugfix, CreatedAtEpoch: s.now.UnixMilli()},
|
||||
{ID: 2, Type: models.ObsTypeFeature, CreatedAtEpoch: s.now.Add(-7 * 24 * time.Hour).UnixMilli()},
|
||||
{ID: 3, Type: models.ObsTypeChange, CreatedAtEpoch: s.now.Add(-14 * 24 * time.Hour).UnixMilli()},
|
||||
}
|
||||
|
||||
scores := s.calc.BatchCalculate(obs, s.now)
|
||||
|
||||
s.Len(scores, 3)
|
||||
s.Contains(scores, int64(1))
|
||||
s.Contains(scores, int64(2))
|
||||
s.Contains(scores, int64(3))
|
||||
|
||||
s.InDelta(1.3, scores[1], 0.01) // New bugfix
|
||||
s.InDelta(0.6, scores[2], 0.1) // 7-day feature
|
||||
s.InDelta(0.225, scores[3], 0.05) // 14-day change
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// CONFIGURATION TESTS
|
||||
// =============================================================================
|
||||
|
||||
func (s *CalculatorSuite) TestNewCalculator_NilConfig() {
|
||||
calc := NewCalculator(nil)
|
||||
s.NotNil(calc)
|
||||
s.NotNil(calc.config)
|
||||
s.Equal(7.0, calc.config.RecencyHalfLifeDays)
|
||||
}
|
||||
|
||||
func (s *CalculatorSuite) TestUpdateConfig() {
|
||||
newConfig := &models.ScoringConfig{
|
||||
RecencyHalfLifeDays: 14.0, // Changed from 7
|
||||
FeedbackWeight: 0.50,
|
||||
ConceptWeight: 0.10,
|
||||
RetrievalWeight: 0.05,
|
||||
MinScore: 0.001,
|
||||
ConceptWeights: map[string]float64{"test": 0.5},
|
||||
}
|
||||
|
||||
s.calc.UpdateConfig(newConfig)
|
||||
|
||||
obs := &models.Observation{
|
||||
ID: 1,
|
||||
Type: models.ObsTypeChange,
|
||||
CreatedAtEpoch: s.now.Add(-14 * 24 * time.Hour).UnixMilli(),
|
||||
}
|
||||
|
||||
score := s.calc.Calculate(obs, s.now)
|
||||
|
||||
// With 14-day half-life, 14 days = exactly one half-life
|
||||
// Expected: 1.0 × 0.9 × 0.5 = 0.45
|
||||
s.InDelta(0.45, score, 0.01)
|
||||
}
|
||||
|
||||
func (s *CalculatorSuite) TestUpdateConfig_NilIgnored() {
|
||||
originalConfig := s.calc.GetConfig()
|
||||
s.calc.UpdateConfig(nil)
|
||||
s.Equal(originalConfig, s.calc.GetConfig())
|
||||
}
|
||||
|
||||
func (s *CalculatorSuite) TestGetConfig() {
|
||||
config := s.calc.GetConfig()
|
||||
s.NotNil(config)
|
||||
s.Equal(7.0, config.RecencyHalfLifeDays)
|
||||
}
|
||||
|
||||
func (s *CalculatorSuite) TestRecalculateThreshold() {
|
||||
threshold := s.calc.RecalculateThreshold()
|
||||
s.Equal(6*time.Hour, threshold)
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// STANDALONE TESTS (non-suite)
|
||||
// =============================================================================
|
||||
|
||||
func TestNewCalculator_DefaultConfig(t *testing.T) {
|
||||
calc := NewCalculator(nil)
|
||||
require.NotNil(t, calc)
|
||||
assert.Equal(t, 7.0, calc.config.RecencyHalfLifeDays)
|
||||
assert.Equal(t, 0.30, calc.config.FeedbackWeight)
|
||||
assert.Equal(t, 0.01, calc.config.MinScore)
|
||||
}
|
||||
|
||||
func TestCalculator_ConcurrentAccess(t *testing.T) {
|
||||
calc := NewCalculator(nil)
|
||||
now := time.Now()
|
||||
|
||||
// Test that calculator is safe for concurrent reads
|
||||
done := make(chan bool, 10)
|
||||
for i := 0; i < 10; i++ {
|
||||
go func(id int64) {
|
||||
obs := &models.Observation{
|
||||
ID: id,
|
||||
Type: models.ObsTypeBugfix,
|
||||
CreatedAtEpoch: now.UnixMilli(),
|
||||
}
|
||||
score := calc.Calculate(obs, now)
|
||||
assert.Greater(t, score, 0.0)
|
||||
done <- true
|
||||
}(int64(i))
|
||||
}
|
||||
|
||||
for i := 0; i < 10; i++ {
|
||||
<-done
|
||||
}
|
||||
}
|
||||
|
||||
func TestCalculator_DecayPrecision(t *testing.T) {
|
||||
calc := NewCalculator(nil)
|
||||
now := time.Now()
|
||||
|
||||
// Test that decay is mathematically correct
|
||||
testCases := []struct {
|
||||
days int
|
||||
expectedDecay float64
|
||||
}{
|
||||
{0, 1.0},
|
||||
{7, 0.5},
|
||||
{14, 0.25},
|
||||
{21, 0.125},
|
||||
{28, 0.0625},
|
||||
}
|
||||
|
||||
for _, tc := range testCases {
|
||||
t.Run(string(rune('0'+tc.days/7))+"_half_lives", func(t *testing.T) {
|
||||
obs := &models.Observation{
|
||||
ID: 1,
|
||||
Type: models.ObsTypeRefactor, // 1.0 weight
|
||||
CreatedAtEpoch: now.Add(-time.Duration(tc.days) * 24 * time.Hour).UnixMilli(),
|
||||
}
|
||||
components := calc.CalculateComponents(obs, now)
|
||||
assert.InDelta(t, tc.expectedDecay, components.RecencyDecay, 0.001)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestTypeBaseScore_UnknownType(t *testing.T) {
|
||||
score := models.TypeBaseScore("unknown-type")
|
||||
assert.Equal(t, 1.0, score, "unknown type should default to 1.0")
|
||||
}
|
||||
|
||||
func TestTypeBaseScore_AllKnownTypes(t *testing.T) {
|
||||
expected := map[models.ObservationType]float64{
|
||||
models.ObsTypeBugfix: 1.3,
|
||||
models.ObsTypeFeature: 1.2,
|
||||
models.ObsTypeDiscovery: 1.1,
|
||||
models.ObsTypeDecision: 1.1,
|
||||
models.ObsTypeRefactor: 1.0,
|
||||
models.ObsTypeChange: 0.9,
|
||||
}
|
||||
|
||||
for obsType, expectedScore := range expected {
|
||||
t.Run(string(obsType), func(t *testing.T) {
|
||||
score := models.TypeBaseScore(obsType)
|
||||
assert.Equal(t, expectedScore, score)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestCalculator_RetrievalBoostDiminishingReturns(t *testing.T) {
|
||||
calc := NewCalculator(nil)
|
||||
now := time.Now()
|
||||
|
||||
// Test that retrieval boost has diminishing returns
|
||||
// When retrieval count doubles, the boost should NOT double (log2 gives diminishing returns)
|
||||
|
||||
// Collect boosts for different counts
|
||||
boosts := make([]float64, 0)
|
||||
retrievalCounts := []int{1, 3, 7, 15, 31, 63, 127}
|
||||
|
||||
for _, count := range retrievalCounts {
|
||||
obs := &models.Observation{
|
||||
ID: 1,
|
||||
Type: models.ObsTypeRefactor,
|
||||
CreatedAtEpoch: now.UnixMilli(),
|
||||
RetrievalCount: count,
|
||||
}
|
||||
components := calc.CalculateComponents(obs, now)
|
||||
boosts = append(boosts, components.RetrievalContrib)
|
||||
}
|
||||
|
||||
// Verify boost increases but at a decreasing rate
|
||||
for i := 1; i < len(boosts); i++ {
|
||||
// Each boost should be higher than the previous
|
||||
assert.Greater(t, boosts[i], boosts[i-1],
|
||||
"boost should increase with more retrievals")
|
||||
|
||||
// But not proportionally - calculate the ratios
|
||||
if i >= 2 {
|
||||
ratio1 := boosts[i-1] / boosts[i-2]
|
||||
ratio2 := boosts[i] / boosts[i-1]
|
||||
// The growth ratio should be decreasing (diminishing returns)
|
||||
assert.Less(t, ratio2, ratio1+0.01, // Allow small floating point tolerance
|
||||
"growth rate should decrease (diminishing returns)")
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,186 +0,0 @@
|
||||
// Package scoring provides importance score calculation for observations.
|
||||
package scoring
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/rs/zerolog"
|
||||
|
||||
"github.com/lukaszraczylo/claude-mnemonic/internal/db/gorm"
|
||||
"github.com/lukaszraczylo/claude-mnemonic/pkg/models"
|
||||
)
|
||||
|
||||
// ObservationStore defines the interface for observation storage operations needed by the recalculator.
|
||||
type ObservationStore interface {
|
||||
GetObservationsNeedingScoreUpdate(ctx context.Context, threshold time.Duration, limit int) ([]*models.Observation, error)
|
||||
UpdateImportanceScores(ctx context.Context, scores map[int64]float64) error
|
||||
GetConceptWeights(ctx context.Context) (map[string]float64, error)
|
||||
}
|
||||
|
||||
// Recalculator periodically recalculates importance scores for observations.
|
||||
type Recalculator struct {
|
||||
log zerolog.Logger
|
||||
store ObservationStore
|
||||
calculator *Calculator
|
||||
stopCh chan struct{}
|
||||
doneCh chan struct{}
|
||||
interval time.Duration
|
||||
batchSize int
|
||||
mu sync.Mutex
|
||||
running bool
|
||||
}
|
||||
|
||||
// NewRecalculator creates a new background recalculator.
|
||||
func NewRecalculator(store ObservationStore, calc *Calculator, log zerolog.Logger) *Recalculator {
|
||||
return &Recalculator{
|
||||
store: store,
|
||||
calculator: calc,
|
||||
log: log.With().Str("component", "recalculator").Logger(),
|
||||
interval: 1 * time.Hour, // Run every hour
|
||||
batchSize: 500, // Process 500 observations at a time
|
||||
stopCh: make(chan struct{}),
|
||||
doneCh: make(chan struct{}),
|
||||
}
|
||||
}
|
||||
|
||||
// Start begins the background recalculation loop.
|
||||
// This should be called in a goroutine.
|
||||
func (r *Recalculator) Start(ctx context.Context) {
|
||||
r.mu.Lock()
|
||||
if r.running {
|
||||
r.mu.Unlock()
|
||||
return
|
||||
}
|
||||
r.running = true
|
||||
r.mu.Unlock()
|
||||
|
||||
defer func() {
|
||||
r.mu.Lock()
|
||||
r.running = false
|
||||
r.mu.Unlock()
|
||||
close(r.doneCh)
|
||||
}()
|
||||
|
||||
// Initial run
|
||||
r.recalculate(ctx)
|
||||
|
||||
r.mu.Lock()
|
||||
interval := r.interval
|
||||
r.mu.Unlock()
|
||||
|
||||
ticker := time.NewTicker(interval)
|
||||
defer ticker.Stop()
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
r.log.Info().Msg("recalculator shutting down due to context cancellation")
|
||||
return
|
||||
case <-r.stopCh:
|
||||
r.log.Info().Msg("recalculator stopping")
|
||||
return
|
||||
case <-ticker.C:
|
||||
r.recalculate(ctx)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Stop stops the background recalculation loop.
|
||||
func (r *Recalculator) Stop() {
|
||||
r.mu.Lock()
|
||||
if !r.running {
|
||||
r.mu.Unlock()
|
||||
return
|
||||
}
|
||||
r.mu.Unlock()
|
||||
|
||||
close(r.stopCh)
|
||||
<-r.doneCh
|
||||
}
|
||||
|
||||
// recalculate performs a single recalculation batch.
|
||||
func (r *Recalculator) recalculate(ctx context.Context) {
|
||||
now := time.Now()
|
||||
threshold := r.calculator.RecalculateThreshold()
|
||||
|
||||
r.mu.Lock()
|
||||
batchSize := r.batchSize
|
||||
r.mu.Unlock()
|
||||
|
||||
observations, err := r.store.GetObservationsNeedingScoreUpdate(ctx, threshold, batchSize)
|
||||
if err != nil {
|
||||
r.log.Error().Err(err).Msg("failed to get observations for score update")
|
||||
return
|
||||
}
|
||||
|
||||
if len(observations) == 0 {
|
||||
return
|
||||
}
|
||||
|
||||
scores := r.calculator.BatchCalculate(observations, now)
|
||||
|
||||
if err := r.store.UpdateImportanceScores(ctx, scores); err != nil {
|
||||
r.log.Error().Err(err).Msg("failed to update importance scores")
|
||||
return
|
||||
}
|
||||
|
||||
r.log.Info().
|
||||
Int("count", len(scores)).
|
||||
Dur("elapsed", time.Since(now)).
|
||||
Msg("recalculated importance scores")
|
||||
}
|
||||
|
||||
// RecalculateNow triggers an immediate recalculation.
|
||||
// This is useful for testing or when scores need to be updated urgently.
|
||||
func (r *Recalculator) RecalculateNow(ctx context.Context) error {
|
||||
r.recalculate(ctx)
|
||||
return nil
|
||||
}
|
||||
|
||||
// RefreshConceptWeights reloads concept weights from the database.
|
||||
// Call this after updating concept weights to apply changes.
|
||||
func (r *Recalculator) RefreshConceptWeights(ctx context.Context) error {
|
||||
weights, err := r.store.GetConceptWeights(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
config := r.calculator.GetConfig()
|
||||
config.ConceptWeights = weights
|
||||
r.calculator.UpdateConfig(config)
|
||||
|
||||
r.log.Info().Int("count", len(weights)).Msg("refreshed concept weights")
|
||||
return nil
|
||||
}
|
||||
|
||||
// Stats returns statistics about the recalculator.
|
||||
type Stats struct {
|
||||
Running bool `json:"running"`
|
||||
Interval time.Duration `json:"interval"`
|
||||
BatchSize int `json:"batch_size"`
|
||||
HalfLife float64 `json:"half_life_days"`
|
||||
MinScore float64 `json:"min_score"`
|
||||
ConceptsLen int `json:"concepts_count"`
|
||||
}
|
||||
|
||||
// GetStats returns current recalculator statistics.
|
||||
func (r *Recalculator) GetStats() Stats {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
|
||||
config := r.calculator.GetConfig()
|
||||
|
||||
return Stats{
|
||||
Running: r.running,
|
||||
Interval: r.interval,
|
||||
BatchSize: r.batchSize,
|
||||
HalfLife: config.RecencyHalfLifeDays,
|
||||
MinScore: config.MinScore,
|
||||
ConceptsLen: len(config.ConceptWeights),
|
||||
}
|
||||
}
|
||||
|
||||
// Ensure ObservationStore satisfies the interface
|
||||
var _ ObservationStore = (*gorm.ObservationStore)(nil)
|
||||
@@ -1,447 +0,0 @@
|
||||
// Package scoring provides importance score calculation for observations.
|
||||
package scoring
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/rs/zerolog"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/lukaszraczylo/claude-mnemonic/pkg/models"
|
||||
)
|
||||
|
||||
// MockObservationStore is a mock implementation of ObservationStore for testing.
|
||||
type MockObservationStore struct {
|
||||
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 {
|
||||
return &MockObservationStore{
|
||||
observations: []*models.Observation{},
|
||||
scores: make(map[int64]float64),
|
||||
conceptWeights: make(map[string]float64),
|
||||
}
|
||||
}
|
||||
|
||||
func (m *MockObservationStore) GetObservationsNeedingScoreUpdate(ctx context.Context, threshold time.Duration, limit int) ([]*models.Observation, error) {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
|
||||
if m.getErr != nil {
|
||||
return nil, m.getErr
|
||||
}
|
||||
|
||||
// Return observations that haven't been updated within threshold
|
||||
now := time.Now()
|
||||
var result []*models.Observation
|
||||
for _, obs := range m.observations {
|
||||
if !obs.ScoreUpdatedAt.Valid || now.Sub(time.Unix(obs.ScoreUpdatedAt.Int64, 0)) > threshold {
|
||||
result = append(result, obs)
|
||||
if len(result) >= limit {
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func (m *MockObservationStore) UpdateImportanceScores(ctx context.Context, scores map[int64]float64) error {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
|
||||
m.updateScoresCalls++
|
||||
if m.updateErr != nil {
|
||||
return m.updateErr
|
||||
}
|
||||
|
||||
for id, score := range scores {
|
||||
m.scores[id] = score
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (m *MockObservationStore) GetConceptWeights(ctx context.Context) (map[string]float64, error) {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
|
||||
if m.getConceptsErr != nil {
|
||||
return nil, m.getConceptsErr
|
||||
}
|
||||
return m.conceptWeights, nil
|
||||
}
|
||||
|
||||
func (m *MockObservationStore) AddObservation(obs *models.Observation) {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
m.observations = append(m.observations, obs)
|
||||
}
|
||||
|
||||
func (m *MockObservationStore) SetConceptWeights(weights map[string]float64) {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
m.conceptWeights = weights
|
||||
}
|
||||
|
||||
func (m *MockObservationStore) GetScore(id int64) (float64, bool) {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
score, ok := m.scores[id]
|
||||
return score, ok
|
||||
}
|
||||
|
||||
func (m *MockObservationStore) GetUpdateScoresCalls() int {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
return m.updateScoresCalls
|
||||
}
|
||||
|
||||
func (m *MockObservationStore) SetUpdateError(err error) {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
m.updateErr = err
|
||||
}
|
||||
|
||||
func (m *MockObservationStore) SetGetError(err error) {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
m.getErr = err
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// RECALCULATOR TESTS
|
||||
// =============================================================================
|
||||
|
||||
func TestNewRecalculator(t *testing.T) {
|
||||
store := NewMockObservationStore()
|
||||
calc := NewCalculator(nil)
|
||||
log := zerolog.Nop()
|
||||
|
||||
recalc := NewRecalculator(store, calc, log)
|
||||
|
||||
require.NotNil(t, recalc)
|
||||
assert.NotNil(t, recalc.store)
|
||||
assert.NotNil(t, recalc.calculator)
|
||||
assert.Equal(t, 1*time.Hour, recalc.interval)
|
||||
assert.Equal(t, 500, recalc.batchSize)
|
||||
assert.False(t, recalc.running)
|
||||
}
|
||||
|
||||
func TestRecalculator_RecalculateNow(t *testing.T) {
|
||||
store := NewMockObservationStore()
|
||||
calc := NewCalculator(nil)
|
||||
log := zerolog.Nop()
|
||||
recalc := NewRecalculator(store, calc, log)
|
||||
|
||||
// Add observations
|
||||
now := time.Now()
|
||||
store.AddObservation(&models.Observation{
|
||||
ID: 1,
|
||||
Type: models.ObsTypeBugfix,
|
||||
CreatedAtEpoch: now.UnixMilli(),
|
||||
})
|
||||
store.AddObservation(&models.Observation{
|
||||
ID: 2,
|
||||
Type: models.ObsTypeFeature,
|
||||
CreatedAtEpoch: now.Add(-7 * 24 * time.Hour).UnixMilli(),
|
||||
})
|
||||
|
||||
ctx := context.Background()
|
||||
err := recalc.RecalculateNow(ctx)
|
||||
|
||||
require.NoError(t, err)
|
||||
|
||||
// Verify scores were calculated
|
||||
score1, ok := store.GetScore(1)
|
||||
assert.True(t, ok)
|
||||
assert.Greater(t, score1, 0.0)
|
||||
|
||||
score2, ok := store.GetScore(2)
|
||||
assert.True(t, ok)
|
||||
assert.Greater(t, score2, 0.0)
|
||||
assert.Less(t, score2, score1, "older observation should have lower score")
|
||||
}
|
||||
|
||||
func TestRecalculator_RefreshConceptWeights(t *testing.T) {
|
||||
store := NewMockObservationStore()
|
||||
calc := NewCalculator(nil)
|
||||
log := zerolog.Nop()
|
||||
recalc := NewRecalculator(store, calc, log)
|
||||
|
||||
// Set up custom weights in store
|
||||
customWeights := map[string]float64{
|
||||
"security": 0.50,
|
||||
"performance": 0.25,
|
||||
}
|
||||
store.SetConceptWeights(customWeights)
|
||||
|
||||
ctx := context.Background()
|
||||
err := recalc.RefreshConceptWeights(ctx)
|
||||
|
||||
require.NoError(t, err)
|
||||
|
||||
// Verify weights were updated in calculator config
|
||||
config := calc.GetConfig()
|
||||
assert.Equal(t, 0.50, config.ConceptWeights["security"])
|
||||
assert.Equal(t, 0.25, config.ConceptWeights["performance"])
|
||||
}
|
||||
|
||||
func TestRecalculator_RefreshConceptWeights_Error(t *testing.T) {
|
||||
store := NewMockObservationStore()
|
||||
calc := NewCalculator(nil)
|
||||
log := zerolog.Nop()
|
||||
recalc := NewRecalculator(store, calc, log)
|
||||
|
||||
store.getConceptsErr = assert.AnError
|
||||
|
||||
ctx := context.Background()
|
||||
err := recalc.RefreshConceptWeights(ctx)
|
||||
|
||||
require.Error(t, err)
|
||||
}
|
||||
|
||||
func TestRecalculator_GetStats(t *testing.T) {
|
||||
store := NewMockObservationStore()
|
||||
calc := NewCalculator(nil)
|
||||
log := zerolog.Nop()
|
||||
recalc := NewRecalculator(store, calc, log)
|
||||
|
||||
// Set fields directly for testing
|
||||
recalc.interval = 2 * time.Hour
|
||||
recalc.batchSize = 250
|
||||
|
||||
stats := recalc.GetStats()
|
||||
|
||||
assert.False(t, stats.Running)
|
||||
assert.Equal(t, 2*time.Hour, stats.Interval)
|
||||
assert.Equal(t, 250, stats.BatchSize)
|
||||
assert.Equal(t, 7.0, stats.HalfLife)
|
||||
assert.Equal(t, 0.01, stats.MinScore)
|
||||
}
|
||||
|
||||
func TestRecalculator_StartStop(t *testing.T) {
|
||||
store := NewMockObservationStore()
|
||||
calc := NewCalculator(nil)
|
||||
log := zerolog.Nop()
|
||||
recalc := NewRecalculator(store, calc, log)
|
||||
|
||||
// Use a short interval for testing
|
||||
recalc.interval = 50 * time.Millisecond
|
||||
|
||||
// Add an observation
|
||||
store.AddObservation(&models.Observation{
|
||||
ID: 1,
|
||||
Type: models.ObsTypeBugfix,
|
||||
CreatedAtEpoch: time.Now().UnixMilli(),
|
||||
})
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
|
||||
// Start in goroutine
|
||||
go recalc.Start(ctx)
|
||||
|
||||
// Wait for initial run and at least one tick
|
||||
time.Sleep(100 * time.Millisecond)
|
||||
|
||||
// Verify it ran
|
||||
calls := store.GetUpdateScoresCalls()
|
||||
assert.GreaterOrEqual(t, calls, 1, "should have run at least once")
|
||||
|
||||
// Stop via context cancellation
|
||||
cancel()
|
||||
time.Sleep(100 * time.Millisecond)
|
||||
|
||||
// Verify stopped
|
||||
stats := recalc.GetStats()
|
||||
assert.False(t, stats.Running)
|
||||
}
|
||||
|
||||
func TestRecalculator_StartTwice(t *testing.T) {
|
||||
store := NewMockObservationStore()
|
||||
calc := NewCalculator(nil)
|
||||
log := zerolog.Nop()
|
||||
recalc := NewRecalculator(store, calc, log)
|
||||
|
||||
recalc.interval = 1 * time.Hour // Long interval so it doesn't tick during test
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
|
||||
// Start first time
|
||||
go recalc.Start(ctx)
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
|
||||
// Try to start second time (should return immediately)
|
||||
go recalc.Start(ctx)
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
|
||||
// Should still be running (only once)
|
||||
stats := recalc.GetStats()
|
||||
assert.True(t, stats.Running)
|
||||
|
||||
cancel()
|
||||
}
|
||||
|
||||
func TestRecalculator_StopWhenNotRunning(t *testing.T) {
|
||||
store := NewMockObservationStore()
|
||||
calc := NewCalculator(nil)
|
||||
log := zerolog.Nop()
|
||||
recalc := NewRecalculator(store, calc, log)
|
||||
|
||||
// Stop without starting - should not panic
|
||||
recalc.Stop()
|
||||
|
||||
stats := recalc.GetStats()
|
||||
assert.False(t, stats.Running)
|
||||
}
|
||||
|
||||
func TestRecalculator_EmptyStore(t *testing.T) {
|
||||
store := NewMockObservationStore()
|
||||
calc := NewCalculator(nil)
|
||||
log := zerolog.Nop()
|
||||
recalc := NewRecalculator(store, calc, log)
|
||||
|
||||
ctx := context.Background()
|
||||
err := recalc.RecalculateNow(ctx)
|
||||
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 0, store.GetUpdateScoresCalls(), "should not call update with no observations")
|
||||
}
|
||||
|
||||
func TestRecalculator_GetObservationsError(t *testing.T) {
|
||||
store := NewMockObservationStore()
|
||||
calc := NewCalculator(nil)
|
||||
log := zerolog.Nop()
|
||||
recalc := NewRecalculator(store, calc, log)
|
||||
|
||||
store.SetGetError(assert.AnError)
|
||||
|
||||
ctx := context.Background()
|
||||
err := recalc.RecalculateNow(ctx)
|
||||
|
||||
// Should not return error (logs it instead)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 0, store.GetUpdateScoresCalls())
|
||||
}
|
||||
|
||||
func TestRecalculator_UpdateScoresError(t *testing.T) {
|
||||
store := NewMockObservationStore()
|
||||
calc := NewCalculator(nil)
|
||||
log := zerolog.Nop()
|
||||
recalc := NewRecalculator(store, calc, log)
|
||||
|
||||
store.AddObservation(&models.Observation{
|
||||
ID: 1,
|
||||
Type: models.ObsTypeBugfix,
|
||||
CreatedAtEpoch: time.Now().UnixMilli(),
|
||||
})
|
||||
|
||||
store.SetUpdateError(assert.AnError)
|
||||
|
||||
ctx := context.Background()
|
||||
err := recalc.RecalculateNow(ctx)
|
||||
|
||||
// Should not return error (logs it instead)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 1, store.GetUpdateScoresCalls())
|
||||
}
|
||||
|
||||
func TestRecalculator_BatchProcessing(t *testing.T) {
|
||||
store := NewMockObservationStore()
|
||||
calc := NewCalculator(nil)
|
||||
log := zerolog.Nop()
|
||||
recalc := NewRecalculator(store, calc, log)
|
||||
|
||||
// Set small batch size
|
||||
recalc.batchSize = 3
|
||||
|
||||
// Add 5 observations
|
||||
now := time.Now()
|
||||
for i := 1; i <= 5; i++ {
|
||||
store.AddObservation(&models.Observation{
|
||||
ID: int64(i),
|
||||
Type: models.ObsTypeBugfix,
|
||||
CreatedAtEpoch: now.UnixMilli(),
|
||||
})
|
||||
}
|
||||
|
||||
ctx := context.Background()
|
||||
err := recalc.RecalculateNow(ctx)
|
||||
|
||||
require.NoError(t, err)
|
||||
|
||||
// Should only process batch size (3)
|
||||
scores := 0
|
||||
for i := 1; i <= 5; i++ {
|
||||
if _, ok := store.GetScore(int64(i)); ok {
|
||||
scores++
|
||||
}
|
||||
}
|
||||
assert.Equal(t, 3, scores, "should only process batch size observations")
|
||||
}
|
||||
|
||||
func TestRecalculator_ConcurrentAccess(t *testing.T) {
|
||||
store := NewMockObservationStore()
|
||||
calc := NewCalculator(nil)
|
||||
log := zerolog.Nop()
|
||||
recalc := NewRecalculator(store, calc, log)
|
||||
|
||||
// Add observations
|
||||
now := time.Now()
|
||||
for i := 1; i <= 10; i++ {
|
||||
store.AddObservation(&models.Observation{
|
||||
ID: int64(i),
|
||||
Type: models.ObsTypeBugfix,
|
||||
CreatedAtEpoch: now.UnixMilli(),
|
||||
})
|
||||
}
|
||||
|
||||
ctx := context.Background()
|
||||
var wg sync.WaitGroup
|
||||
|
||||
// Run multiple recalculations concurrently
|
||||
for i := 0; i < 5; i++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
_ = recalc.RecalculateNow(ctx)
|
||||
}()
|
||||
}
|
||||
|
||||
wg.Wait()
|
||||
|
||||
// Should complete without race conditions
|
||||
// (use -race flag to verify)
|
||||
assert.GreaterOrEqual(t, store.GetUpdateScoresCalls(), 1)
|
||||
}
|
||||
|
||||
func TestRecalculator_StatsThreadSafe(t *testing.T) {
|
||||
store := NewMockObservationStore()
|
||||
calc := NewCalculator(nil)
|
||||
log := zerolog.Nop()
|
||||
recalc := NewRecalculator(store, calc, log)
|
||||
|
||||
var wg sync.WaitGroup
|
||||
|
||||
// Concurrent reads (use -race flag to verify)
|
||||
for i := 0; i < 10; i++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
_ = recalc.GetStats()
|
||||
}()
|
||||
}
|
||||
|
||||
wg.Wait()
|
||||
}
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user