mirror of
https://github.com/lukaszraczylo/claude-mnemonic.git
synced 2026-10-12 04:55:18 +00:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
f2886bbedb | ||
|
|
9b6e5f7648 | ||
|
|
055103cb3b | ||
|
|
67de0f0c84 | ||
|
|
eba3136fff | ||
|
|
17f7bffef5 | ||
|
|
0d75b867e6 | ||
|
|
a8e9d0f256 | ||
|
|
8857cc9c72 | ||
|
|
096490a2d3 | ||
|
|
eaf7c4ddca | ||
|
|
2d1f52cdf4 | ||
|
|
9d7959f98f | ||
|
|
daa9959961 | ||
|
|
f8e5b08de7 | ||
|
|
ad718f9292 | ||
|
|
1cfe6936de | ||
|
|
2b2bc47fe5 | ||
|
|
88d316b3ec | ||
|
|
4cef2eb792 | ||
|
|
c3cad756f9 | ||
|
|
c3fdfed512 | ||
|
|
7cb0dfbc13 | ||
|
|
9e5bcc64a8 | ||
|
|
0bbba6c078 | ||
|
|
da1d96dab4 | ||
|
|
443f557000 | ||
|
|
affdd4fab0 | ||
|
|
41b774f973 | ||
|
|
ee285441e4 | ||
|
|
2658694839 | ||
|
|
e71d34912e | ||
|
|
9c832f9ac0 | ||
|
|
7d55a6a493 | ||
|
|
8995781628 | ||
|
|
10ab2a6d41 | ||
|
|
052fef0af8 | ||
|
|
3384bc1d5f | ||
|
|
3caf83b7ea | ||
|
|
93b19d5402 | ||
|
|
6a819918f6 | ||
|
|
98ab1cd3e9 | ||
|
|
55ba09e0c6 | ||
|
|
a369a421aa | ||
|
|
76f31e97ae | ||
|
|
132666a190 | ||
|
|
3a8e3667b3 | ||
|
|
3ea3904875 | ||
|
|
8a8814968e | ||
|
|
6b8ecdd2ec | ||
|
|
aae542565e | ||
|
|
1fcb8722c4 | ||
|
|
d64a384e92 | ||
|
|
bd04da2b69 | ||
|
|
5191bc474a | ||
|
|
1b5c974473 | ||
|
|
b38ecee4da | ||
|
|
406b539c84 | ||
|
|
a23b8c4936 | ||
|
|
98e9a2d59c | ||
|
|
af19b6a798 | ||
|
|
11ee9a4f1f | ||
|
|
70a6eb2e52 | ||
|
|
86ee0e28ed | ||
|
|
1b5697b316 | ||
|
|
a8a8c532c4 | ||
|
|
e057f121eb | ||
|
|
7d1998e975 | ||
|
|
af8e6f44cd | ||
|
|
cf4161d459 | ||
|
|
910aabda22 | ||
|
|
6115ea8a44 | ||
|
|
93446d1437 | ||
|
|
1d1807da9e |
@@ -14,7 +14,6 @@
|
||||
"mcpServers": {
|
||||
"claude-mnemonic": {
|
||||
"command": "${CLAUDE_PLUGIN_ROOT}/mcp-server",
|
||||
"args": ["--project", "${CLAUDE_PROJECT}"],
|
||||
"env": {}
|
||||
}
|
||||
},
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
name: Autoupdate go.mod and go.sum
|
||||
name: Update dependencies
|
||||
|
||||
on:
|
||||
workflow_dispatch:
|
||||
@@ -6,16 +6,16 @@ on:
|
||||
- cron: "0 3 * * *"
|
||||
|
||||
permissions:
|
||||
contents: write
|
||||
actions: write
|
||||
contents: write
|
||||
issues: write
|
||||
pull-requests: write
|
||||
security-events: write
|
||||
|
||||
jobs:
|
||||
autoupdate:
|
||||
uses: lukaszraczylo/shared-actions/.github/workflows/go-autoupdate.yaml@main
|
||||
update:
|
||||
uses: lukaszraczylo/shared-actions/.github/workflows/renovate-autoupdate.yaml@main
|
||||
with:
|
||||
go-version: ">=1.24"
|
||||
release-workflow: "release.yaml"
|
||||
lfs: false
|
||||
secrets: inherit
|
||||
node-working-directory: ui
|
||||
require-node-tests: false
|
||||
release-workflow: release.yaml
|
||||
|
||||
@@ -32,7 +32,7 @@ jobs:
|
||||
needs: release
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v5
|
||||
with:
|
||||
ref: main
|
||||
|
||||
|
||||
@@ -22,7 +22,7 @@ jobs:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- name: Checkout
|
||||
uses: actions/checkout@v4
|
||||
uses: actions/checkout@v5
|
||||
|
||||
- name: Setup Node.js
|
||||
uses: actions/setup-node@v4
|
||||
|
||||
@@ -95,3 +95,6 @@ docs/dist
|
||||
plugin/.claude-plugin/plugin.json
|
||||
plugin/.claude-plugin/marketplace.json
|
||||
user-prompt
|
||||
|
||||
# Compressed ONNX runtime libs are produced by scripts/download-onnx-libs.sh
|
||||
internal/embedding/assets/lib/**/*.gz
|
||||
|
||||
+23
-21
@@ -1,34 +1,36 @@
|
||||
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
|
||||
version: "2"
|
||||
|
||||
linters:
|
||||
enable:
|
||||
- errcheck
|
||||
- gosec
|
||||
- govet
|
||||
- gofmt
|
||||
- staticcheck
|
||||
- unused
|
||||
- ineffassign
|
||||
- typecheck
|
||||
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
|
||||
exclusions:
|
||||
paths:
|
||||
- vendor
|
||||
# Exclude some linters from running on test files
|
||||
rules:
|
||||
- path: _test\.go
|
||||
linters:
|
||||
- errcheck
|
||||
- gosec
|
||||
|
||||
issues:
|
||||
exclude-dirs:
|
||||
- vendor
|
||||
# Exclude some linters from running on test files
|
||||
exclude-rules:
|
||||
- path: _test\.go
|
||||
linters:
|
||||
- errcheck
|
||||
- gosec
|
||||
formatters:
|
||||
enable:
|
||||
- gofmt
|
||||
|
||||
run:
|
||||
timeout: 5m
|
||||
|
||||
@@ -172,13 +172,19 @@ Example: A bug fix in your auth module stays local. "Always validate JWT server-
|
||||
|
||||
## MCP Tools
|
||||
|
||||
These search tools are available via MCP:
|
||||
Four tools are exposed via MCP:
|
||||
|
||||
- `search` - semantic search across all memories
|
||||
- `timeline` - browse by time
|
||||
- `decisions` - find architecture decisions
|
||||
- `changes` - find code modifications
|
||||
- `how_it_works` - system understanding queries
|
||||
- `search` - semantic search across all memories. Accepts `obs_type` (e.g.
|
||||
`decision`, `code_change`, `architecture`) plus filter params (`concepts`,
|
||||
`files`, `type`) - these replace the old per-type shortcut tools.
|
||||
- `timeline` - browse observations around a point in time.
|
||||
- `observation` - manage individual observations. Set `action` to one of:
|
||||
`get`, `edit`, `delete`, `supersede`, `boost`, `merge`, `related`, `similar`,
|
||||
`quality`, `relationships`, `scoring`, `tag`, `by_tag`, `batch_tag`.
|
||||
- `memory_admin` - administration and analytics. Set `action` to one of:
|
||||
`stats`, `health`, `maintenance_stats`, `run_maintenance`, `importance`,
|
||||
`search_patterns`, `explain_ranking`, `temporal_trends`, `data_quality`,
|
||||
`export`, `suggest_consolidations`, `patterns`.
|
||||
|
||||
## Slash Commands
|
||||
|
||||
|
||||
@@ -3,9 +3,9 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"os"
|
||||
"time"
|
||||
@@ -13,6 +13,9 @@ import (
|
||||
"github.com/lukaszraczylo/claude-mnemonic/pkg/hooks"
|
||||
)
|
||||
|
||||
// statsTimeout bounds the worker stats request.
|
||||
const statsTimeout = 100 * time.Millisecond
|
||||
|
||||
// StatusInput is the JSON input from Claude Code's statusline feature.
|
||||
type StatusInput struct {
|
||||
HookEventName string `json:"hook_event_name"`
|
||||
@@ -105,26 +108,19 @@ func handleStatusline(input *StatusInput, port int) string {
|
||||
// getWorkerStats fetches stats from the worker service.
|
||||
func getWorkerStats(port int, project string) *WorkerStats {
|
||||
// Build URL with optional project parameter
|
||||
endpoint := fmt.Sprintf("http://127.0.0.1:%d/api/stats", port)
|
||||
endpoint := "/api/stats"
|
||||
if project != "" {
|
||||
endpoint += "?project=" + url.QueryEscape(project)
|
||||
}
|
||||
|
||||
// Create HTTP client with short timeout (statusline must be fast)
|
||||
client := &http.Client{Timeout: 100 * time.Millisecond}
|
||||
|
||||
resp, err := client.Get(endpoint)
|
||||
// Short timeout: statusline must be fast
|
||||
body, err := hooks.GETBody(context.Background(), statsTimeout, port, endpoint)
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return nil
|
||||
}
|
||||
|
||||
var stats WorkerStats
|
||||
if err := json.NewDecoder(resp.Body).Decode(&stats); err != nil {
|
||||
if err := json.Unmarshal(body, &stats); err != nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
|
||||
+14
-3
@@ -8,12 +8,14 @@ import (
|
||||
"net/http"
|
||||
"os"
|
||||
"os/signal"
|
||||
"strings"
|
||||
"syscall"
|
||||
"time"
|
||||
|
||||
"github.com/lukaszraczylo/claude-mnemonic/internal/config"
|
||||
"github.com/lukaszraczylo/claude-mnemonic/internal/mcp"
|
||||
"github.com/lukaszraczylo/claude-mnemonic/internal/watcher"
|
||||
"github.com/lukaszraczylo/claude-mnemonic/pkg/hooks"
|
||||
"github.com/lukaszraczylo/oss-telemetry"
|
||||
"github.com/rs/zerolog"
|
||||
"github.com/rs/zerolog/log"
|
||||
@@ -22,9 +24,18 @@ import (
|
||||
// Version is set at build time via ldflags.
|
||||
var Version = "dev"
|
||||
|
||||
// defaultProject derives the project ID the same way the hooks do.
|
||||
func defaultProject() string {
|
||||
dir := os.Getenv("CLAUDE_PROJECT_DIR")
|
||||
if dir == "" {
|
||||
dir, _ = os.Getwd()
|
||||
}
|
||||
return hooks.ProjectIDWithName(dir)
|
||||
}
|
||||
|
||||
func main() {
|
||||
// Parse flags
|
||||
project := flag.String("project", "", "Project name (required)")
|
||||
project := flag.String("project", "", "Project ID (default: derived from CLAUDE_PROJECT_DIR or cwd)")
|
||||
debug := flag.Bool("debug", false, "Enable debug logging")
|
||||
flag.Parse()
|
||||
|
||||
@@ -35,8 +46,8 @@ func main() {
|
||||
}
|
||||
log.Logger = log.Output(zerolog.ConsoleWriter{Out: os.Stderr, NoColor: true})
|
||||
|
||||
if *project == "" {
|
||||
log.Fatal().Msg("--project is required")
|
||||
if *project == "" || strings.Contains(*project, "${") {
|
||||
*project = defaultProject()
|
||||
}
|
||||
|
||||
// Get worker port from config
|
||||
|
||||
@@ -1,3 +1,8 @@
|
||||
---
|
||||
description: Restart the Claude Mnemonic worker process
|
||||
allowed-tools: Bash(curl:*)
|
||||
---
|
||||
|
||||
# Restart Claude Mnemonic Worker
|
||||
|
||||
Restart the claude-mnemonic worker process. Use this command when experiencing issues with the memory system.
|
||||
|
||||
@@ -0,0 +1,60 @@
|
||||
# Experiment journal
|
||||
|
||||
## Campaign B: binary size
|
||||
|
||||
Metric: sum of bytes of worker, mcp and 6 hook binaries (release flags, `-s -w -trimpath`, `-tags fts5`). Lower is better.
|
||||
Threshold: at least 3% drop, tests green. Stop: 5 consecutive no-improve trials or 15 trials.
|
||||
Mutate: Go code and build flags. Frozen: public API, DB schema, plugin manifest.
|
||||
|
||||
| id | mutation | total bytes | delta vs base | kept? |
|
||||
|----|----------|-------------|---------------|-------|
|
||||
| 01 | baseline | 102698416 | - | base |
|
||||
| 02 | pkg/hooks: raw HTTP/1.0 client replaces net/http | 87603952 | -14.7% | KEEP |
|
||||
| 03 | statusline: shared client, drop net/http | 84531056 | -3.5% vs 02 | KEEP |
|
||||
| 04 | embed gzipped ONNX runtime libs, gunzip at first use | 57216880 | -32.3% vs 03 | KEEP |
|
||||
| 05 | hooks built with CGO_ENABLED=0 | 57216880 | 0% | revert (not applied) |
|
||||
| - | statically dedupe or compress tokenizer.json | - | projected ~1.2-1.6% | not run, under 3% threshold |
|
||||
| - | drop net/http from mcp | - | none possible | not run, oss-telemetry links net/http |
|
||||
|
||||
Campaign B stopped on diminishing returns: 102698416 -> 57216880 bytes (-44.3%).
|
||||
|
||||
## Campaign A: hook latency (stub worker, 200 runs, p50)
|
||||
|
||||
Harness: hook binary against a local stub worker (100-observation context), same machine, alternating runs. Lower is better.
|
||||
|
||||
| hook | before (ee28544) | after trial 02 | delta |
|
||||
|------|------------------|----------------|-------|
|
||||
| session-start | 4.6-5.6 ms | 3.3 ms | -30% or better |
|
||||
| user-prompt | 4.3-4.7 ms | 3.5-3.7 ms | -15% to -22% |
|
||||
|
||||
Hooks now sit near the process-start floor. No further trials run.
|
||||
|
||||
## Campaign D: reliability
|
||||
|
||||
Metric: failing cases in the extraction fault test (`TestWriteGunzipped`). Lower is better.
|
||||
|
||||
| id | mutation | failing cases | kept? |
|
||||
|----|----------|---------------|-------|
|
||||
| 01 | baseline (truncated archive leaves partial lib.so) | 1 | base |
|
||||
| 02 | extract via temp file and rename | 0 | KEEP |
|
||||
|
||||
## Campaign C: search accuracy
|
||||
|
||||
Harness: `scripts/eval-search` seeds 30 labelled observations, starts a real worker (fresh HOME, model cache copied), waits for the vector rebuild, then runs 30 paraphrased queries against `/api/context/search`. Metric: MRR@10 (higher is better). Threshold: at least +0.02. Runs are deterministic (repeat runs identical).
|
||||
|
||||
| id | mutation | MRR@10 | recall@5 | kept? |
|
||||
|----|----------|--------|----------|-------|
|
||||
| 01 | baseline | 0.7944 | 0.8667 | base |
|
||||
| 02 | L2-normalise embeddings (similarity was negative for every vector hit, so search fell back to FTS) | 0.8944 | 0.9667 | KEEP |
|
||||
| 03 | reranking alpha 0.7 -> 0.5 | 0.9111 | 0.9667 | revert (+0.017 < 0.02) |
|
||||
| 04 | reranking alpha 0.7 -> 0.9 | 0.9111 | 0.9667 | revert (+0.017 < 0.02) |
|
||||
| 05 | reranker off | 0.8889 | 0.9667 | revert |
|
||||
| 06 | reranker pure mode | 0.8714 | 0.9667 | revert |
|
||||
| 07 | CLS pooling instead of mean | 0.8944 | 0.9667 | revert |
|
||||
|
||||
Trial 02 changes stored vectors. `BGEStorageVersion` is bumped to `bge-v1.5-l2`, so existing installs rebuild vectors on next start.
|
||||
Caveat: 30 queries is a small set, so borderline gains such as alpha 0.5 are not distinguishable from noise.
|
||||
|
||||
Notes:
|
||||
- Trial 04 needs `scripts/download-onnx-libs.sh` to produce `.gz` files; existing raw libs trigger a re-download.
|
||||
- Extracted library verified byte-identical to the raw dylib (`cmp`).
|
||||
@@ -1,41 +1,38 @@
|
||||
module github.com/lukaszraczylo/claude-mnemonic
|
||||
|
||||
go 1.25.1
|
||||
go 1.26.0
|
||||
|
||||
replace github.com/sugarme/tokenizer => github.com/clems4ever/tokenizer v0.0.0-20250926133620-9ddc80533c43
|
||||
|
||||
require (
|
||||
github.com/asg017/sqlite-vec-go-bindings v0.1.6
|
||||
github.com/fsnotify/fsnotify v1.10.1
|
||||
github.com/go-chi/chi/v5 v5.3.0
|
||||
github.com/go-gormigrate/gormigrate/v2 v2.1.6
|
||||
github.com/goccy/go-json v0.10.6
|
||||
github.com/lukaszraczylo/oss-telemetry v0.2.3
|
||||
github.com/mattn/go-sqlite3 v1.14.44
|
||||
github.com/go-chi/chi/v5 v5.3.2
|
||||
github.com/go-gormigrate/gormigrate/v2 v2.1.7
|
||||
github.com/goccy/go-json v0.11.2
|
||||
github.com/lukaszraczylo/oss-telemetry v0.2.5
|
||||
github.com/mattn/go-sqlite3 v1.14.52
|
||||
github.com/rs/zerolog v1.35.1
|
||||
github.com/smacker/go-tree-sitter v0.0.0-20240827094217-dd81d9e9be82
|
||||
github.com/stretchr/testify v1.11.1
|
||||
github.com/stretchr/testify v1.12.1
|
||||
github.com/sugarme/tokenizer v0.3.0
|
||||
github.com/yalue/onnxruntime_go v1.30.1
|
||||
golang.org/x/sync v0.20.0
|
||||
github.com/yalue/onnxruntime_go v1.36.0
|
||||
golang.org/x/sync v0.23.0
|
||||
gorm.io/driver/sqlite v1.6.0
|
||||
gorm.io/gorm v1.31.1
|
||||
gorm.io/gorm v1.31.2
|
||||
)
|
||||
|
||||
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.15 // indirect
|
||||
github.com/mattn/go-isatty v0.0.22 // indirect
|
||||
github.com/mattn/go-isatty v0.0.24 // indirect
|
||||
github.com/mitchellh/colorstring v0.0.0-20190213212951-d06e56a500db // indirect
|
||||
github.com/patrickmn/go-cache v2.1.0+incompatible // indirect
|
||||
github.com/pmezard/go-difflib v1.0.0 // indirect
|
||||
github.com/rivo/uniseg v0.4.7 // indirect
|
||||
github.com/schollz/progressbar/v2 v2.15.0 // indirect
|
||||
github.com/sugarme/regexpset v0.0.0-20200920021344-4d4ec8eaf93c // indirect
|
||||
golang.org/x/sys v0.45.0 // indirect
|
||||
golang.org/x/text v0.37.0 // indirect
|
||||
gopkg.in/yaml.v3 v3.0.1 // indirect
|
||||
go.yaml.in/yaml/v3 v3.0.5 // indirect
|
||||
golang.org/x/sys v0.48.0 // indirect
|
||||
golang.org/x/text v0.42.0 // indirect
|
||||
)
|
||||
|
||||
@@ -3,35 +3,33 @@ github.com/asg017/sqlite-vec-go-bindings v0.1.6/go.mod h1:A8+cTt/nKFsYCQF6OgzSNp
|
||||
github.com/clems4ever/tokenizer v0.0.0-20250926133620-9ddc80533c43 h1:j8YQypEqa5OjqbGciCNb9hOcYbo1oTVuEjd/iu9U2SY=
|
||||
github.com/clems4ever/tokenizer v0.0.0-20250926133620-9ddc80533c43/go.mod h1:VJ+DLK5ZEZwzvODOWwY0cw+B1dabTd3nCB5HuFCItCc=
|
||||
github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
|
||||
github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
|
||||
github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
|
||||
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.10.1 h1:b0/UzAf9yR5rhf3RPm9gf3ehBPpf0oZKIjtpKrx59Ho=
|
||||
github.com/fsnotify/fsnotify v1.10.1/go.mod h1:TLheqan6HD6GBK6PrDWyDPBaEV8LspOxvPSjC+bVfgo=
|
||||
github.com/go-chi/chi/v5 v5.3.0 h1:halUjDxhshgXHMrao5bB8eNBXo/rnzwr8m5m36glehM=
|
||||
github.com/go-chi/chi/v5 v5.3.0/go.mod h1:R+tYY2hNuVUUjxoPtqUdgBqevM9s9njzkTLutVsOCto=
|
||||
github.com/go-gormigrate/gormigrate/v2 v2.1.6 h1:VtX+l1Stj2v5RGubVQk0LS/8EPGXR+ldcOyCmlmKoyg=
|
||||
github.com/go-gormigrate/gormigrate/v2 v2.1.6/go.mod h1:PZpedQc4tWaxn6kvXicwhinh3L0seLpMc5ReKRX5id4=
|
||||
github.com/goccy/go-json v0.10.6 h1:p8HrPJzOakx/mn/bQtjgNjdTcN+/S6FcG2CTtQOrHVU=
|
||||
github.com/goccy/go-json v0.10.6/go.mod h1:oq7eo15ShAhp70Anwd5lgX2pLfOS3QCiwU/PULtXL6M=
|
||||
github.com/go-chi/chi/v5 v5.3.2 h1:5YQkICvTCSZ25hoRsyJazN0scjzKGiu4VAUc7H1o1nY=
|
||||
github.com/go-chi/chi/v5 v5.3.2/go.mod h1:R+tYY2hNuVUUjxoPtqUdgBqevM9s9njzkTLutVsOCto=
|
||||
github.com/go-gormigrate/gormigrate/v2 v2.1.7 h1:PdT4jVPbRb4R+0Ey2R0yJOdctVf4Whiq1Qi4necaZdg=
|
||||
github.com/go-gormigrate/gormigrate/v2 v2.1.7/go.mod h1:3ouXglTuPrKF5+7cQyVGfvAXTU4vLMaYh9+EPl03uog=
|
||||
github.com/goccy/go-json v0.11.2 h1:jdZv93Tt4ioR8yW1CoNsvSxrcZlCXAUU1aZXN7gpXUA=
|
||||
github.com/goccy/go-json v0.11.2/go.mod h1:3NdmfEkZlB7YI5UFw/qdFKq8XN1aiWR0YyRPWZNQltY=
|
||||
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/lukaszraczylo/oss-telemetry v0.2.3 h1:xoDtBqeZGmXj7IteiE1M5WMuzeoqag58qEleI0Cf2Ms=
|
||||
github.com/lukaszraczylo/oss-telemetry v0.2.3/go.mod h1:+Cn78qZo8rc3T9eZt0v3oICYRdd75wORtSidc8lNjDQ=
|
||||
github.com/lukaszraczylo/oss-telemetry v0.2.5 h1:W6Qgwp7IY4yBRO361MmnEAmM59motdQpbAFAx8v6fJM=
|
||||
github.com/lukaszraczylo/oss-telemetry v0.2.5/go.mod h1:+Cn78qZo8rc3T9eZt0v3oICYRdd75wORtSidc8lNjDQ=
|
||||
github.com/mattn/go-colorable v0.1.15 h1:+u9SLTRGnXv73cEsnsmoZBom+dMU88B2M0aDcWy0/jY=
|
||||
github.com/mattn/go-colorable v0.1.15/go.mod h1:6LmQG8QLFO4G5z1gPvYEzlUgJ2wF+stgPZH1UqBm1s8=
|
||||
github.com/mattn/go-isatty v0.0.22 h1:j8l17JJ9i6VGPUFUYoTUKPSgKe/83EYU2zBC7YNKMw4=
|
||||
github.com/mattn/go-isatty v0.0.22/go.mod h1:ZXfXG4SQHsB/w3ZeOYbR0PrPwLy+n6xiMrJlRFqopa4=
|
||||
github.com/mattn/go-sqlite3 v1.14.44 h1:3VSe+xafpbzsLbdr2AWlAZk9yRHiBhTBakioXaCKTF8=
|
||||
github.com/mattn/go-sqlite3 v1.14.44/go.mod h1:pjEuOr8IwzLJP2MfGeTb0A35jauH+C2kbHKBr7yXKVQ=
|
||||
github.com/mattn/go-isatty v0.0.24 h1:tGZZoVgT/KiqK1c8ocVLeDS8BSWMRd47J3Lbz7vsReI=
|
||||
github.com/mattn/go-isatty v0.0.24/go.mod h1:nMCL3Zebbrt45jsMDgnfIwz6ydEQApk5oEI3HqDio6A=
|
||||
github.com/mattn/go-sqlite3 v1.14.52 h1:wVbm2Qnf4OXkqhBTSPuCRZDRnxfbVrrmiCEroVdog8U=
|
||||
github.com/mattn/go-sqlite3 v1.14.52/go.mod h1:6JTjA44L93a0QCyJef5YvlPoKXntQPjzWv5gtm9sB6w=
|
||||
github.com/mitchellh/colorstring v0.0.0-20190213212951-d06e56a500db h1:62I3jR2EmQ4l5rM/4FEfDWcRD+abF5XlKShorW5LRoQ=
|
||||
github.com/mitchellh/colorstring v0.0.0-20190213212951-d06e56a500db/go.mod h1:l0dey0ia/Uv7NcFFVbCLtqEBQbrT4OCwCSKTEv6enCw=
|
||||
github.com/patrickmn/go-cache v2.1.0+incompatible h1:HRMgzkcYKYpi3C8ajMPV8OFXaaRUnok+kx1WdO15EQc=
|
||||
github.com/patrickmn/go-cache v2.1.0+incompatible/go.mod h1:3Qf8kWWT7OJRJbdiICTKqZju1ZixQ/KpMGzzAfe6+WQ=
|
||||
github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
|
||||
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
|
||||
github.com/rivo/uniseg v0.4.7 h1:WUdvkW8uEhrYfLC4ZzdpI2ztxP1I582+49Oc5Mq64VQ=
|
||||
github.com/rivo/uniseg v0.4.7/go.mod h1:FN3SvrM+Zdj16jyLfmOkMNblXMcoc8DfTHruCPUcx88=
|
||||
@@ -39,27 +37,23 @@ github.com/rs/zerolog v1.35.1 h1:m7xQeoiLIiV0BCEY4Hs+j2NG4Gp2o2KPKmhnnLiazKI=
|
||||
github.com/rs/zerolog v1.35.1/go.mod h1:EjML9kdfa/RMA7h/6z6pYmq1ykOuA8/mjWaEvGI+jcw=
|
||||
github.com/schollz/progressbar/v2 v2.15.0 h1:dVzHQ8fHRmtPjD3K10jT3Qgn/+H+92jhPrhmxIJfDz8=
|
||||
github.com/schollz/progressbar/v2 v2.15.0/go.mod h1:UdPq3prGkfQ7MOzZKlDRpYKcFqEMczbD7YmbPgpzKMI=
|
||||
github.com/smacker/go-tree-sitter v0.0.0-20240827094217-dd81d9e9be82 h1:6C8qej6f1bStuePVkLSFxoU22XBS165D3klxlzRg8F4=
|
||||
github.com/smacker/go-tree-sitter v0.0.0-20240827094217-dd81d9e9be82/go.mod h1:xe4pgH49k4SsmkQq5OT8abwhWmnzkhpgnXeekbx2efw=
|
||||
github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
|
||||
github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=
|
||||
github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U=
|
||||
github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U=
|
||||
github.com/stretchr/testify v1.12.1 h1:EuwCh5fleGS7H32xRwO3wRGT7DxrDhLAT6FF8MpWDWE=
|
||||
github.com/stretchr/testify v1.12.1/go.mod h1:MDEgiDPPsNp5cuIrHPPCyornHKgEVbtFUmoNlxoYthg=
|
||||
github.com/sugarme/regexpset v0.0.0-20200920021344-4d4ec8eaf93c h1:pwb4kNSHb4K89ymCaN+5lPH/MwnfSVg4rzGDh4d+iy4=
|
||||
github.com/sugarme/regexpset v0.0.0-20200920021344-4d4ec8eaf93c/go.mod h1:2gwkXLWbDGUQWeL3RtpCmcY4mzCtU13kb9UsAg9xMaw=
|
||||
github.com/yalue/onnxruntime_go v1.30.1 h1:NaEng5lWbsHZ/8X1dtaw1mIj7eV1ozyjbFo//g0ktl4=
|
||||
github.com/yalue/onnxruntime_go v1.30.1/go.mod h1:b4X26A8pekNb1ACJ58wAXgNKeUCGEAQ9dmACut9Sm/4=
|
||||
golang.org/x/sync v0.20.0 h1:e0PTpb7pjO8GAtTs2dQ6jYa5BWYlMuX047Dco/pItO4=
|
||||
golang.org/x/sync v0.20.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0=
|
||||
golang.org/x/sys v0.45.0 h1:dO4czNzziLiiXplLQgBCEpCvXQ3dnkn0SdaZSYdQ+FY=
|
||||
golang.org/x/sys v0.45.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw=
|
||||
golang.org/x/text v0.37.0 h1:Cqjiwd9eSg8e0QAkyCaQTNHFIIzWtidPahFWR83rTrc=
|
||||
golang.org/x/text v0.37.0/go.mod h1:a5sjxXGs9hsn/AJVwuElvCAo9v8QYLzvavO5z2PiM38=
|
||||
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=
|
||||
github.com/yalue/onnxruntime_go v1.36.0 h1:iH1Q++DcsyT9sWtN26KYimESlI5hhXpKaChHDS44oV4=
|
||||
github.com/yalue/onnxruntime_go v1.36.0/go.mod h1:b4X26A8pekNb1ACJ58wAXgNKeUCGEAQ9dmACut9Sm/4=
|
||||
go.yaml.in/yaml/v3 v3.0.5 h1:N6y/pJk8buWs9NY5ERU2HSMfm+IuD/OtfdAnq6kESPw=
|
||||
go.yaml.in/yaml/v3 v3.0.5/go.mod h1:HVTZu1O7/Vkt2N+BFy8Zza+lnLsABggaTM2ZpNIGuKg=
|
||||
golang.org/x/sync v0.23.0 h1:KameEIfc1IkluZyXWLn39Wd4tURc6GbCiISGiZm2bQk=
|
||||
golang.org/x/sync v0.23.0/go.mod h1:sUUOizhqBxiL6pEWpqNLUiaJn1ShEbZ6BBqskPbjZm0=
|
||||
golang.org/x/sys v0.48.0 h1:bbX/i/6MgT9BVLM9RT1thmxL04yeTAhbEz4SyadbXoo=
|
||||
golang.org/x/sys v0.48.0/go.mod h1:hNLxWAXmnKAxqDtdwIYC4bM9oQPEecfsnNMuSxOs3og=
|
||||
golang.org/x/text v0.42.0 h1:JbOZXgfeCPU9gacVtYliJqOhD+zhrEqK4LfdpmlUZqI=
|
||||
golang.org/x/text v0.42.0/go.mod h1:ojzP1Z+2QtioaF8DTtO8K5q7JWVVYwZKenzujK0Zd0E=
|
||||
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=
|
||||
gorm.io/gorm v1.31.2 h1:3o8FXNo9v9S858gil+3LlZA1LkCOzgb4g5BL64FgaCo=
|
||||
gorm.io/gorm v1.31.2/go.mod h1:XyQVbO2k6YkOis7C2437jSit3SsDK72s7n7rsSHd+Gs=
|
||||
|
||||
+5
-5
@@ -6,7 +6,7 @@
|
||||
"hooks": [
|
||||
{
|
||||
"type": "command",
|
||||
"command": "${CLAUDE_PLUGIN_ROOT}/hooks/session-start",
|
||||
"command": "\"${CLAUDE_PLUGIN_ROOT}/hooks/session-start\"",
|
||||
"timeout": 30
|
||||
}
|
||||
]
|
||||
@@ -17,7 +17,7 @@
|
||||
"hooks": [
|
||||
{
|
||||
"type": "command",
|
||||
"command": "${CLAUDE_PLUGIN_ROOT}/hooks/user-prompt",
|
||||
"command": "\"${CLAUDE_PLUGIN_ROOT}/hooks/user-prompt\"",
|
||||
"timeout": 10
|
||||
}
|
||||
]
|
||||
@@ -29,7 +29,7 @@
|
||||
"hooks": [
|
||||
{
|
||||
"type": "command",
|
||||
"command": "${CLAUDE_PLUGIN_ROOT}/hooks/post-tool-use",
|
||||
"command": "\"${CLAUDE_PLUGIN_ROOT}/hooks/post-tool-use\"",
|
||||
"timeout": 10
|
||||
}
|
||||
]
|
||||
@@ -40,7 +40,7 @@
|
||||
"hooks": [
|
||||
{
|
||||
"type": "command",
|
||||
"command": "${CLAUDE_PLUGIN_ROOT}/hooks/subagent-stop",
|
||||
"command": "\"${CLAUDE_PLUGIN_ROOT}/hooks/subagent-stop\"",
|
||||
"timeout": 10
|
||||
}
|
||||
]
|
||||
@@ -51,7 +51,7 @@
|
||||
"hooks": [
|
||||
{
|
||||
"type": "command",
|
||||
"command": "${CLAUDE_PLUGIN_ROOT}/hooks/stop",
|
||||
"command": "\"${CLAUDE_PLUGIN_ROOT}/hooks/stop\"",
|
||||
"timeout": 30
|
||||
}
|
||||
]
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
@@ -252,9 +252,6 @@ func (s *ObservationStore) UpdateObservation(ctx context.Context, id int64, upda
|
||||
updates["scope"] = sql.NullString{String: *update.Scope, Valid: true}
|
||||
}
|
||||
|
||||
// Add updated_at timestamp
|
||||
updates["updated_at_epoch"] = sql.NullInt64{Int64: time.Now().Unix(), Valid: true}
|
||||
|
||||
if len(updates) == 0 {
|
||||
// Nothing to update, just return existing observation
|
||||
return toModelObservation(&dbObs), nil
|
||||
@@ -618,7 +615,7 @@ func (s *ObservationStore) SearchObservationsFTS(ctx context.Context, query, pro
|
||||
// FTS failed, try LIKE fallback
|
||||
return s.searchObservationsLike(ctx, keywords, project, limit)
|
||||
}
|
||||
defer rows.Close()
|
||||
defer func() { _ = rows.Close() }()
|
||||
|
||||
observations, err := scanObservationRows(rows)
|
||||
if err != nil {
|
||||
|
||||
@@ -500,15 +500,13 @@ func TestObservationStore_CleanupOldObservations(t *testing.T) {
|
||||
|
||||
// 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)
|
||||
_, _, 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
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,94 @@
|
||||
package gorm
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/lukaszraczylo/claude-mnemonic/pkg/models"
|
||||
)
|
||||
|
||||
func storeForUpdate(t *testing.T) (*ObservationStore, int64, func()) {
|
||||
t.Helper()
|
||||
obsStore, store, cleanup := testObservationStore(t)
|
||||
ctx := context.Background()
|
||||
sessionID, err := NewSessionStore(store).CreateSDKSession(ctx, "claude-1", "test-project", "")
|
||||
require.NoError(t, err)
|
||||
id, _, err := obsStore.StoreObservation(ctx, "claude-1", "test-project", &models.ParsedObservation{
|
||||
Type: models.ObsTypeDiscovery, Title: "Original title", Subtitle: "Original subtitle", Narrative: "Original narrative",
|
||||
Facts: []string{"fact one"}, Concepts: []string{"cache"}, FilesRead: []string{"a.go"},
|
||||
}, int(sessionID), 1)
|
||||
require.NoError(t, err)
|
||||
return obsStore, id, cleanup
|
||||
}
|
||||
|
||||
func ptr[T any](v T) *T { return &v }
|
||||
|
||||
func TestUpdateObservation_ChangesOnlyTheGivenFields(t *testing.T) {
|
||||
obsStore, id, cleanup := storeForUpdate(t)
|
||||
defer cleanup()
|
||||
ctx := context.Background()
|
||||
|
||||
got, err := obsStore.UpdateObservation(ctx, id, &ObservationUpdate{Title: ptr("New title"), Narrative: ptr("New narrative")})
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "New title", got.Title.String)
|
||||
assert.Equal(t, "New narrative", got.Narrative.String)
|
||||
assert.Equal(t, "Original subtitle", got.Subtitle.String, "a field that was not given stays")
|
||||
assert.Equal(t, []string{"fact one"}, []string(got.Facts))
|
||||
|
||||
stored, err := obsStore.GetObservationByID(ctx, id)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "New title", stored.Title.String, "and it is stored, not only returned")
|
||||
}
|
||||
|
||||
func TestUpdateObservation_ReplacesListsAndScope(t *testing.T) {
|
||||
obsStore, id, cleanup := storeForUpdate(t)
|
||||
defer cleanup()
|
||||
ctx := context.Background()
|
||||
|
||||
got, err := obsStore.UpdateObservation(ctx, id, &ObservationUpdate{
|
||||
Facts: &[]string{"a", "b"}, Concepts: &[]string{}, FilesRead: &[]string{"x.go"}, FilesModified: &[]string{"y.go"}, Scope: ptr("global"),
|
||||
})
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, []string{"a", "b"}, []string(got.Facts))
|
||||
assert.Empty(t, got.Concepts, "an empty list clears the field")
|
||||
assert.Equal(t, []string{"x.go"}, []string(got.FilesRead))
|
||||
assert.Equal(t, []string{"y.go"}, []string(got.FilesModified))
|
||||
assert.Equal(t, models.ScopeGlobal, got.Scope)
|
||||
}
|
||||
|
||||
func TestUpdateObservation_NothingToChangeReturnsTheObservation(t *testing.T) {
|
||||
obsStore, id, cleanup := storeForUpdate(t)
|
||||
defer cleanup()
|
||||
|
||||
got, err := obsStore.UpdateObservation(context.Background(), id, &ObservationUpdate{})
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "Original title", got.Title.String)
|
||||
}
|
||||
|
||||
func TestUpdateObservation_Refusals(t *testing.T) {
|
||||
obsStore, id, cleanup := storeForUpdate(t)
|
||||
defer cleanup()
|
||||
ctx := context.Background()
|
||||
|
||||
_, err := obsStore.UpdateObservation(ctx, id, nil)
|
||||
assert.Error(t, err, "a nil update")
|
||||
_, err = obsStore.UpdateObservation(ctx, 999999, &ObservationUpdate{Title: ptr("x")})
|
||||
require.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "not found")
|
||||
}
|
||||
|
||||
func TestUpdateObservation_KeepsTheSearchIndexInStep(t *testing.T) {
|
||||
obsStore, id, cleanup := storeForUpdate(t)
|
||||
defer cleanup()
|
||||
ctx := context.Background()
|
||||
|
||||
_, err := obsStore.UpdateObservation(ctx, id, &ObservationUpdate{Title: ptr("Zebrafish enclosure")})
|
||||
require.NoError(t, err)
|
||||
var n int64
|
||||
require.NoError(t, obsStore.db.WithContext(ctx).Raw(
|
||||
`SELECT COUNT(*) FROM observations_fts WHERE observations_fts MATCH 'zebrafish'`).Scan(&n).Error)
|
||||
assert.EqualValues(t, 1, n, "the full-text index sees the new title")
|
||||
}
|
||||
@@ -1 +1 @@
|
||||
1.26.0
|
||||
1.30.0
|
||||
|
||||
@@ -1 +1 @@
|
||||
1.26.0
|
||||
1.30.0
|
||||
|
||||
@@ -1 +1 @@
|
||||
1.26.0
|
||||
1.30.0
|
||||
|
||||
@@ -1 +1 @@
|
||||
1.26.0
|
||||
1.30.0
|
||||
|
||||
@@ -6,7 +6,7 @@ import (
|
||||
_ "embed"
|
||||
)
|
||||
|
||||
//go:embed assets/lib/darwin-arm64/libonnxruntime.dylib
|
||||
//go:embed assets/lib/darwin-arm64/libonnxruntime.dylib.gz
|
||||
var onnxRuntimeLib []byte
|
||||
|
||||
const onnxRuntimeLibName = "libonnxruntime.dylib"
|
||||
|
||||
@@ -6,10 +6,10 @@ import (
|
||||
_ "embed"
|
||||
)
|
||||
|
||||
//go:embed assets/lib/linux-amd64/libonnxruntime.so
|
||||
//go:embed assets/lib/linux-amd64/libonnxruntime.so.gz
|
||||
var onnxRuntimeLib []byte
|
||||
|
||||
//go:embed assets/lib/linux-amd64/libonnxruntime_providers_shared.so
|
||||
//go:embed assets/lib/linux-amd64/libonnxruntime_providers_shared.so.gz
|
||||
var onnxRuntimeProvidersLib []byte
|
||||
|
||||
const onnxRuntimeLibName = "libonnxruntime.so"
|
||||
|
||||
@@ -6,10 +6,10 @@ import (
|
||||
_ "embed"
|
||||
)
|
||||
|
||||
//go:embed assets/lib/linux-arm64/libonnxruntime.so
|
||||
//go:embed assets/lib/linux-arm64/libonnxruntime.so.gz
|
||||
var onnxRuntimeLib []byte
|
||||
|
||||
//go:embed assets/lib/linux-arm64/libonnxruntime_providers_shared.so
|
||||
//go:embed assets/lib/linux-arm64/libonnxruntime_providers_shared.so.gz
|
||||
var onnxRuntimeProvidersLib []byte
|
||||
|
||||
const onnxRuntimeLibName = "libonnxruntime.so"
|
||||
|
||||
@@ -6,7 +6,7 @@ import (
|
||||
_ "embed"
|
||||
)
|
||||
|
||||
//go:embed assets/lib/windows-amd64/onnxruntime.dll
|
||||
//go:embed assets/lib/windows-amd64/onnxruntime.dll.gz
|
||||
var onnxRuntimeLib []byte
|
||||
|
||||
const onnxRuntimeLibName = "onnxruntime.dll"
|
||||
|
||||
@@ -138,7 +138,25 @@ func (r *ModelRegistry) List() []ModelMetadata {
|
||||
}
|
||||
|
||||
// DefaultRegistry is the global model registry with all available models.
|
||||
var DefaultRegistry = NewModelRegistry()
|
||||
// It is constructed eagerly at package load via newDefaultRegistry, which
|
||||
// registers the built-in BGE model — replacing the previous init()-based
|
||||
// mutate-after-init pattern while preserving identical behaviour.
|
||||
var DefaultRegistry = newDefaultRegistry()
|
||||
|
||||
// newDefaultRegistry builds a model registry pre-populated with the built-in
|
||||
// models (currently BGE). Registering at construction time keeps DefaultRegistry
|
||||
// ready-to-use without a separate init() mutating it after package load.
|
||||
func newDefaultRegistry() *ModelRegistry {
|
||||
r := NewModelRegistry()
|
||||
r.Register(ModelMetadata{
|
||||
Name: BGEModelName,
|
||||
Version: BGEModelVersion,
|
||||
Dimensions: EmbeddingDim,
|
||||
Description: "High-quality semantic search model",
|
||||
Default: true,
|
||||
}, newBGEModel)
|
||||
return r
|
||||
}
|
||||
|
||||
// RegisterModel adds a model to the default registry.
|
||||
func RegisterModel(meta ModelMetadata, factory ModelFactory) {
|
||||
|
||||
@@ -3,10 +3,13 @@ package embedding
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"compress/gzip"
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"fmt"
|
||||
"io"
|
||||
"math"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
@@ -24,6 +27,9 @@ const EmbeddingDim = 384
|
||||
|
||||
// Model version constants
|
||||
const (
|
||||
// BGEStorageVersion tags stored vectors; bump it to force a rebuild when vectors change.
|
||||
BGEStorageVersion = "bge-v1.5-l2"
|
||||
|
||||
// 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
|
||||
@@ -126,7 +132,7 @@ func (m *bgeModel) Name() string {
|
||||
|
||||
// Version returns the short version string for storage.
|
||||
func (m *bgeModel) Version() string {
|
||||
return BGEModelVersion
|
||||
return BGEStorageVersion
|
||||
}
|
||||
|
||||
// Dimensions returns the embedding vector size.
|
||||
@@ -134,7 +140,7 @@ func (m *bgeModel) Dimensions() int {
|
||||
return EmbeddingDim
|
||||
}
|
||||
|
||||
// extractONNXLibrary extracts the embedded ONNX runtime library to a temp directory.
|
||||
// extractONNXLibrary gunzips the embedded ONNX runtime library to a temp directory.
|
||||
// Uses content hash to avoid re-extracting if already present.
|
||||
func extractONNXLibrary() (string, error) {
|
||||
// Create a hash of the library content for cache key
|
||||
@@ -157,16 +163,14 @@ func extractONNXLibrary() (string, error) {
|
||||
}
|
||||
|
||||
// Write main library
|
||||
// #nosec G306 -- Shared library needs executable permission (0755) for dynamic linker
|
||||
if err := os.WriteFile(libPath, onnxRuntimeLib, 0755); err != nil {
|
||||
if err := writeGunzipped(libPath, onnxRuntimeLib); err != nil {
|
||||
return "", fmt.Errorf("write library: %w", err)
|
||||
}
|
||||
|
||||
// Write providers library if present (Linux only)
|
||||
if len(onnxRuntimeProvidersLib) > 0 && onnxRuntimeProvidersLibName != "" {
|
||||
providersPath := filepath.Join(cacheDir, onnxRuntimeProvidersLibName)
|
||||
// #nosec G306 -- Shared library needs executable permission (0755) for dynamic linker
|
||||
if err := os.WriteFile(providersPath, onnxRuntimeProvidersLib, 0755); err != nil {
|
||||
if err := writeGunzipped(providersPath, onnxRuntimeProvidersLib); err != nil {
|
||||
return "", fmt.Errorf("write providers library: %w", err)
|
||||
}
|
||||
}
|
||||
@@ -512,12 +516,28 @@ func meanPooling(embeddings []float32, attentionMask []int64, batchSize, seqLen,
|
||||
}
|
||||
}
|
||||
|
||||
results[b] = result
|
||||
results[b] = l2Normalize(result)
|
||||
}
|
||||
|
||||
return results
|
||||
}
|
||||
|
||||
// l2Normalize scales v to unit length in place; sqlite-vec distances assume unit vectors.
|
||||
func l2Normalize(v []float32) []float32 {
|
||||
var sum float64
|
||||
for _, x := range v {
|
||||
sum += float64(x) * float64(x)
|
||||
}
|
||||
if sum == 0 {
|
||||
return v
|
||||
}
|
||||
inv := float32(1 / math.Sqrt(sum))
|
||||
for i := range v {
|
||||
v[i] *= inv
|
||||
}
|
||||
return v
|
||||
}
|
||||
|
||||
// clsPooling extracts the [CLS] token embedding (first token).
|
||||
// Input shape: [batch, seq_len, hidden]
|
||||
// Output shape: [batch, hidden]
|
||||
@@ -529,7 +549,7 @@ func clsPooling(embeddings []float32, batchSize, seqLen, hiddenSize int) [][]flo
|
||||
// CLS token is at position 0
|
||||
embOffset := b * seqLen * hiddenSize
|
||||
copy(result, embeddings[embOffset:embOffset+hiddenSize])
|
||||
results[b] = result
|
||||
results[b] = l2Normalize(result)
|
||||
}
|
||||
|
||||
return results
|
||||
@@ -562,17 +582,6 @@ func (m *bgeModel) Close() error {
|
||||
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
|
||||
@@ -637,3 +646,38 @@ func (s *Service) EmbedBatchWithContext(ctx context.Context, texts []string) ([]
|
||||
func (s *Service) Close() error {
|
||||
return s.model.Close()
|
||||
}
|
||||
|
||||
// writeGunzipped decompresses gz into path with the executable bit the dynamic linker needs.
|
||||
// It writes a temp file and renames it, so a crash never leaves a truncated library at path.
|
||||
func writeGunzipped(path string, gz []byte) (err error) {
|
||||
zr, err := gzip.NewReader(bytes.NewReader(gz))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer func() { _ = zr.Close() }()
|
||||
|
||||
f, err := os.CreateTemp(filepath.Dir(path), filepath.Base(path)+".tmp*")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
tmp := f.Name()
|
||||
defer func() {
|
||||
if err != nil {
|
||||
_ = os.Remove(tmp)
|
||||
}
|
||||
}()
|
||||
|
||||
// #nosec G110 -- archive is embedded at build time, not external input
|
||||
if _, err = io.Copy(f, zr); err != nil {
|
||||
_ = f.Close()
|
||||
return err
|
||||
}
|
||||
if err = f.Close(); err != nil {
|
||||
return err
|
||||
}
|
||||
// #nosec G302 -- shared library needs the executable bit for the dynamic linker
|
||||
if err = os.Chmod(tmp, 0o755); err != nil {
|
||||
return err
|
||||
}
|
||||
return os.Rename(tmp, path)
|
||||
}
|
||||
|
||||
@@ -1,9 +1,13 @@
|
||||
package embedding
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"compress/gzip"
|
||||
"context"
|
||||
"errors"
|
||||
"math"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
@@ -331,8 +335,8 @@ func holdMutex(m *bgeModel) (ready <-chan struct{}, release func()) {
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
m.mu.Lock()
|
||||
close(ch) // signal: lock acquired
|
||||
<-done // wait for release signal
|
||||
close(ch) // signal: lock acquired
|
||||
<-done // wait for release signal
|
||||
m.mu.Unlock()
|
||||
}()
|
||||
return ch, func() { close(done) }
|
||||
@@ -475,3 +479,57 @@ func cosineSimilarity(a, b []float32) float64 {
|
||||
|
||||
return dotProduct / (math.Sqrt(normA) * math.Sqrt(normB))
|
||||
}
|
||||
|
||||
func gzipBytes(t *testing.T, data []byte) []byte {
|
||||
t.Helper()
|
||||
var buf bytes.Buffer
|
||||
zw := gzip.NewWriter(&buf)
|
||||
_, err := zw.Write(data)
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, zw.Close())
|
||||
return buf.Bytes()
|
||||
}
|
||||
|
||||
func TestWriteGunzipped(t *testing.T) {
|
||||
payload := bytes.Repeat([]byte("onnx"), 50_000)
|
||||
good := gzipBytes(t, payload)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
gz []byte
|
||||
wantErr bool
|
||||
wantFile bool
|
||||
}{
|
||||
{"valid archive", good, false, true},
|
||||
{"truncated archive leaves no file", good[:len(good)/2], true, false},
|
||||
{"not gzip leaves no file", []byte("plain text"), true, false},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
path := filepath.Join(dir, "lib.so")
|
||||
|
||||
err := writeGunzipped(path, tt.gz)
|
||||
|
||||
if tt.wantErr {
|
||||
require.Error(t, err)
|
||||
} else {
|
||||
require.NoError(t, err)
|
||||
}
|
||||
_, statErr := os.Stat(path)
|
||||
assert.Equal(t, tt.wantFile, statErr == nil)
|
||||
entries, _ := os.ReadDir(dir)
|
||||
if !tt.wantFile {
|
||||
assert.Empty(t, entries, "no partial or temp files may remain")
|
||||
}
|
||||
if tt.wantFile {
|
||||
got, err := os.ReadFile(path)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, payload, got)
|
||||
info, err := os.Stat(path)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, os.FileMode(0o755), info.Mode().Perm())
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -363,24 +363,24 @@ func TestRunNow_RetentionDays_DeletesExpiredObservations(t *testing.T) {
|
||||
|
||||
// Insert an observation and back-date it to 2 days ago.
|
||||
obs := &gormdb.Observation{
|
||||
SDKSessionID: "old-session",
|
||||
Project: "proj",
|
||||
Type: models.ObsTypeDiscovery,
|
||||
CreatedAt: "2000-01-01T00:00:00Z",
|
||||
CreatedAtEpoch: time.Now().AddDate(0, 0, -2).Unix(),
|
||||
Scope: models.ScopeProject,
|
||||
SDKSessionID: "old-session",
|
||||
Project: "proj",
|
||||
Type: models.ObsTypeDiscovery,
|
||||
CreatedAt: "2000-01-01T00:00:00Z",
|
||||
CreatedAtEpoch: time.Now().AddDate(0, 0, -2).Unix(),
|
||||
Scope: models.ScopeProject,
|
||||
ImportanceScore: 1.0,
|
||||
}
|
||||
require.NoError(t, store.GetDB().WithContext(ctx).Create(obs).Error)
|
||||
|
||||
// Insert a recent observation (should survive).
|
||||
recentObs := &gormdb.Observation{
|
||||
SDKSessionID: "new-session",
|
||||
Project: "proj",
|
||||
Type: models.ObsTypeDiscovery,
|
||||
CreatedAt: time.Now().Format(time.RFC3339),
|
||||
CreatedAtEpoch: time.Now().Unix(),
|
||||
Scope: models.ScopeProject,
|
||||
SDKSessionID: "new-session",
|
||||
Project: "proj",
|
||||
Type: models.ObsTypeDiscovery,
|
||||
CreatedAt: time.Now().Format(time.RFC3339),
|
||||
CreatedAtEpoch: time.Now().Unix(),
|
||||
Scope: models.ScopeProject,
|
||||
ImportanceScore: 1.0,
|
||||
}
|
||||
require.NoError(t, store.GetDB().WithContext(ctx).Create(recentObs).Error)
|
||||
@@ -434,12 +434,12 @@ func TestRunNow_RetentionDays_VectorCleanupCalled(t *testing.T) {
|
||||
|
||||
// Insert an expired observation directly.
|
||||
obs := &gormdb.Observation{
|
||||
SDKSessionID: "session-x",
|
||||
Project: "proj",
|
||||
Type: models.ObsTypeDiscovery,
|
||||
CreatedAt: "2000-01-01T00:00:00Z",
|
||||
CreatedAtEpoch: time.Now().AddDate(0, 0, -2).Unix(),
|
||||
Scope: models.ScopeProject,
|
||||
SDKSessionID: "session-x",
|
||||
Project: "proj",
|
||||
Type: models.ObsTypeDiscovery,
|
||||
CreatedAt: "2000-01-01T00:00:00Z",
|
||||
CreatedAtEpoch: time.Now().AddDate(0, 0, -2).Unix(),
|
||||
Scope: models.ScopeProject,
|
||||
ImportanceScore: 1.0,
|
||||
}
|
||||
require.NoError(t, store.GetDB().WithContext(ctx).Create(obs).Error)
|
||||
@@ -670,12 +670,12 @@ func TestRunNow_RetentionDays_BatchDeletion_MoreThan100Rows(t *testing.T) {
|
||||
// Insert 150 expired observations (forces 2 batches of 100).
|
||||
for i := 0; i < 150; i++ {
|
||||
obs := &gormdb.Observation{
|
||||
SDKSessionID: "session-old",
|
||||
Project: "proj",
|
||||
Type: models.ObsTypeDiscovery,
|
||||
CreatedAt: "2000-01-01T00:00:00Z",
|
||||
CreatedAtEpoch: time.Now().AddDate(0, 0, -2).Unix(),
|
||||
Scope: models.ScopeProject,
|
||||
SDKSessionID: "session-old",
|
||||
Project: "proj",
|
||||
Type: models.ObsTypeDiscovery,
|
||||
CreatedAt: "2000-01-01T00:00:00Z",
|
||||
CreatedAtEpoch: time.Now().AddDate(0, 0, -2).Unix(),
|
||||
Scope: models.ScopeProject,
|
||||
ImportanceScore: 1.0,
|
||||
}
|
||||
require.NoError(t, store.GetDB().WithContext(ctx).Create(obs).Error)
|
||||
@@ -706,14 +706,14 @@ func TestRunNow_CleanupStale_BatchDeletion_MoreThan100Rows(t *testing.T) {
|
||||
// Insert 120 superseded observations.
|
||||
for i := 0; i < 120; i++ {
|
||||
obs := &gormdb.Observation{
|
||||
SDKSessionID: "session-stale",
|
||||
Project: "proj",
|
||||
Type: models.ObsTypeDiscovery,
|
||||
CreatedAt: time.Now().Format(time.RFC3339),
|
||||
CreatedAtEpoch: time.Now().Unix(),
|
||||
Scope: models.ScopeProject,
|
||||
SDKSessionID: "session-stale",
|
||||
Project: "proj",
|
||||
Type: models.ObsTypeDiscovery,
|
||||
CreatedAt: time.Now().Format(time.RFC3339),
|
||||
CreatedAtEpoch: time.Now().Unix(),
|
||||
Scope: models.ScopeProject,
|
||||
ImportanceScore: 1.0,
|
||||
IsSuperseded: 1,
|
||||
IsSuperseded: 1,
|
||||
}
|
||||
require.NoError(t, store.GetDB().WithContext(ctx).Create(obs).Error)
|
||||
}
|
||||
|
||||
@@ -0,0 +1,38 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
)
|
||||
|
||||
func TestIdleTimeoutFromEnv(t *testing.T) {
|
||||
// The idle shutdown is off unless asked for; anything that is not a positive duration leaves it off.
|
||||
for value, want := range map[string]time.Duration{
|
||||
"": 0,
|
||||
"0": 0,
|
||||
"abc": 0,
|
||||
"-5m": 0,
|
||||
"45": 0, // a bare number is not a duration
|
||||
"90m": 90 * time.Minute,
|
||||
"2h": 2 * time.Hour,
|
||||
"30s": 30 * time.Second,
|
||||
" 2h ": 0, // not trimmed: the operator wrote it, so it is wrong rather than guessed
|
||||
} {
|
||||
t.Run(value, func(t *testing.T) {
|
||||
t.Setenv(IdleTimeoutEnv, value)
|
||||
assert.Equal(t, want, idleTimeoutFromEnv())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestNewServer_HasNoIdleTimeoutByDefault(t *testing.T) {
|
||||
t.Setenv(IdleTimeoutEnv, "")
|
||||
assert.Zero(t, NewServer(nil, "http://localhost:1", "p", "v").idleTimeout)
|
||||
}
|
||||
|
||||
func TestNewServer_TakesTheIdleTimeoutFromTheEnvironment(t *testing.T) {
|
||||
t.Setenv(IdleTimeoutEnv, "3h")
|
||||
assert.Equal(t, 3*time.Hour, NewServer(nil, "http://localhost:1", "p", "v").idleTimeout)
|
||||
}
|
||||
@@ -0,0 +1,97 @@
|
||||
//go:build unix
|
||||
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"os"
|
||||
"sync"
|
||||
"syscall"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// blockingPipe returns the two ends of a pipe in blocking mode, the way an inherited stdin is: a Read on it cannot be
|
||||
// woken by closing the file (the standard pipe from os.Pipe is pollable and can, which hid this).
|
||||
func blockingPipe(t *testing.T) (r, w *os.File) {
|
||||
t.Helper()
|
||||
var fds [2]int
|
||||
require.NoError(t, syscall.Pipe(fds[:]))
|
||||
r = os.NewFile(uintptr(fds[0]), "stdin-like")
|
||||
w = os.NewFile(uintptr(fds[1]), "stdin-writer")
|
||||
t.Cleanup(func() { _ = w.Close(); _ = r.Close() })
|
||||
return r, w
|
||||
}
|
||||
|
||||
type lockedBuffer struct {
|
||||
buf bytes.Buffer
|
||||
mu sync.Mutex
|
||||
}
|
||||
|
||||
func (b *lockedBuffer) Write(p []byte) (int, error) {
|
||||
b.mu.Lock()
|
||||
defer b.mu.Unlock()
|
||||
return b.buf.Write(p)
|
||||
}
|
||||
|
||||
func (b *lockedBuffer) String() string {
|
||||
b.mu.Lock()
|
||||
defer b.mu.Unlock()
|
||||
return b.buf.String()
|
||||
}
|
||||
|
||||
// The idle timeout used to close stdin and hope the read loop ended. On a real (blocking) stdin the read stays blocked,
|
||||
// so the server logged "shutting down" and then lived on with a closed stdin; the next message from the client woke the
|
||||
// read, which failed with "file already closed" and killed the connection mid-call (seen in Claude Desktop after the
|
||||
// connector sat unused for more than 30 minutes). An idle timeout that is on must end Run for real.
|
||||
func TestRun_IdleTimeoutEndsTheServerCleanly(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
stdin, _ := blockingPipe(t)
|
||||
server := &Server{stdin: stdin, stdout: &lockedBuffer{}, idleTimeout: 50 * time.Millisecond, monitorInterval: 10 * time.Millisecond}
|
||||
|
||||
done := make(chan error, 1)
|
||||
go func() { done <- server.Run(context.Background()) }()
|
||||
|
||||
select {
|
||||
case err := <-done:
|
||||
require.NoError(t, err, "an idle shutdown is a clean exit, not an error")
|
||||
case <-time.After(3 * time.Second):
|
||||
t.Fatal("Run did not return after the idle timeout: stdin was closed but the blocked read never ended")
|
||||
}
|
||||
}
|
||||
|
||||
// By default there is no idle timeout: a stdio server lives as long as its client, and a client that talks to it after a
|
||||
// long pause must get an answer.
|
||||
func TestRun_WithoutAnIdleTimeoutTheServerKeepsAnsweringAfterALongPause(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
stdin, stdinWriter := blockingPipe(t)
|
||||
stdout := &lockedBuffer{}
|
||||
server := &Server{stdin: stdin, stdout: stdout, monitorInterval: 5 * time.Millisecond}
|
||||
|
||||
done := make(chan error, 1)
|
||||
go func() { done <- server.Run(context.Background()) }()
|
||||
|
||||
time.Sleep(300 * time.Millisecond) // many monitor ticks with no traffic
|
||||
select {
|
||||
case err := <-done:
|
||||
t.Fatalf("Run ended during a pause with no traffic: %v", err)
|
||||
default:
|
||||
}
|
||||
|
||||
_, err := stdinWriter.WriteString(`{"jsonrpc":"2.0","id":7,"method":"tools/list"}` + "\n")
|
||||
require.NoError(t, err)
|
||||
|
||||
assert.Eventually(t, func() bool { return bytes.Contains([]byte(stdout.String()), []byte(`"id":7`)) }, 3*time.Second, 10*time.Millisecond,
|
||||
"the request after the pause is answered")
|
||||
select {
|
||||
case err := <-done:
|
||||
t.Fatalf("Run ended after answering: %v", err)
|
||||
default:
|
||||
}
|
||||
}
|
||||
+143
-467
@@ -30,6 +30,11 @@ type Server struct {
|
||||
version string
|
||||
writeMu sync.Mutex
|
||||
lastActivity atomic.Int64
|
||||
// idleTimeout ends the server after this long without a message; 0 (the default) never does. A stdio server lives
|
||||
// as long as its client: ending it while the client is still connected makes the client's next call fail.
|
||||
idleTimeout time.Duration
|
||||
// monitorInterval is how often the parent and idle checks run (30s when 0).
|
||||
monitorInterval time.Duration
|
||||
}
|
||||
|
||||
// NewServer creates a new MCP server that proxies to the worker HTTP API.
|
||||
@@ -41,9 +46,24 @@ func NewServer(client *http.Client, workerURL, project, version string) *Server
|
||||
version: version,
|
||||
stdin: os.Stdin,
|
||||
stdout: os.Stdout,
|
||||
|
||||
idleTimeout: idleTimeoutFromEnv(),
|
||||
}
|
||||
}
|
||||
|
||||
// IdleTimeoutEnv names the environment variable that turns the idle shutdown on, as a Go duration such as "2h".
|
||||
// It is off by default.
|
||||
const IdleTimeoutEnv = "CLAUDE_MNEMONIC_MCP_IDLE_TIMEOUT"
|
||||
|
||||
// idleTimeoutFromEnv reads IdleTimeoutEnv; an unset, invalid or non-positive value means no idle shutdown.
|
||||
func idleTimeoutFromEnv() time.Duration {
|
||||
d, err := time.ParseDuration(os.Getenv(IdleTimeoutEnv))
|
||||
if err != nil || d <= 0 {
|
||||
return 0
|
||||
}
|
||||
return d
|
||||
}
|
||||
|
||||
// Request represents a JSON-RPC request.
|
||||
type Request struct {
|
||||
JSONRPC string `json:"jsonrpc"`
|
||||
@@ -82,6 +102,13 @@ type Tool struct {
|
||||
|
||||
// Run starts the MCP server loop.
|
||||
func (s *Server) Run(ctx context.Context) error {
|
||||
// An idle shutdown cancels this context. Closing stdin does not do it: on a real (blocking) stdin the read stays
|
||||
// blocked, the server lived on with a closed stdin, and the client's next message woke the read with "file already
|
||||
// closed", which killed the connection mid-call.
|
||||
ctx, cancel := context.WithCancel(ctx)
|
||||
defer cancel()
|
||||
var idleStopped atomic.Bool
|
||||
|
||||
scanner := bufio.NewScanner(s.stdin)
|
||||
buf := make([]byte, 0, 64*1024)
|
||||
scanner.Buffer(buf, 1024*1024) // 1MB max message size
|
||||
@@ -100,11 +127,15 @@ func (s *Server) Run(ctx context.Context) error {
|
||||
scanErr <- scanner.Err()
|
||||
}()
|
||||
|
||||
// Monitor parent process liveness and idle timeout.
|
||||
// If the parent dies (ppid changes) or no messages arrive for 30 minutes, shut down.
|
||||
// Monitor parent process liveness and, when one is set, the idle timeout.
|
||||
// If the parent dies (ppid changes) or no messages arrive for idleTimeout, shut down.
|
||||
parentPID := os.Getppid()
|
||||
go func() {
|
||||
ticker := time.NewTicker(30 * time.Second)
|
||||
interval := s.monitorInterval
|
||||
if interval <= 0 {
|
||||
interval = 30 * time.Second
|
||||
}
|
||||
ticker := time.NewTicker(interval)
|
||||
defer ticker.Stop()
|
||||
for {
|
||||
select {
|
||||
@@ -118,11 +149,10 @@ func (s *Server) Run(ctx context.Context) error {
|
||||
}
|
||||
return
|
||||
}
|
||||
if time.Since(time.Unix(s.lastActivity.Load(), 0)) > 30*time.Minute {
|
||||
log.Info().Msg("MCP server idle timeout (30m), shutting down")
|
||||
if closer, ok := s.stdin.(io.Closer); ok {
|
||||
_ = closer.Close()
|
||||
}
|
||||
if s.idleTimeout > 0 && time.Since(time.Unix(s.lastActivity.Load(), 0)) > s.idleTimeout {
|
||||
log.Info().Dur("idle", s.idleTimeout).Msg("MCP server idle timeout, shutting down")
|
||||
idleStopped.Store(true)
|
||||
cancel()
|
||||
return
|
||||
}
|
||||
}
|
||||
@@ -140,6 +170,9 @@ func (s *Server) Run(ctx context.Context) error {
|
||||
case <-ctx.Done():
|
||||
// Drain in-flight requests before returning.
|
||||
wg.Wait()
|
||||
if idleStopped.Load() {
|
||||
return nil // an idle shutdown the operator asked for is a clean exit
|
||||
}
|
||||
return ctx.Err()
|
||||
case line, ok := <-lines:
|
||||
if !ok {
|
||||
@@ -290,475 +323,65 @@ func (s *Server) handleToolsList(req *Request) *Response {
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "decisions",
|
||||
Description: "Semantic shortcut for finding architectural, design, and implementation decisions.",
|
||||
Name: "observation",
|
||||
Description: "Manage individual observations. Set 'action' to choose the operation: get, edit, delete, supersede, boost, merge, related, similar, quality, relationships, scoring, tag, by_tag, batch_tag.",
|
||||
InputSchema: map[string]any{
|
||||
"type": "object",
|
||||
"required": []string{"query"},
|
||||
"required": []string{"action"},
|
||||
"properties": map[string]any{
|
||||
"query": map[string]any{"type": "string", "description": "Natural language query for finding decisions"},
|
||||
"dateStart": map[string]any{"type": []string{"string", "number"}},
|
||||
"dateEnd": map[string]any{"type": []string{"string", "number"}},
|
||||
"limit": map[string]any{"type": "number", "default": 20, "minimum": 1, "maximum": 100},
|
||||
"format": map[string]any{"type": "string", "enum": []string{"index", "full"}, "default": "index"},
|
||||
"action": map[string]any{"type": "string", "enum": []string{"get", "edit", "delete", "supersede", "boost", "merge", "related", "similar", "quality", "relationships", "scoring", "tag", "by_tag", "batch_tag"}, "description": "Operation to perform"},
|
||||
"id": map[string]any{"type": "number", "description": "Observation ID (get, edit, quality, relationships, scoring, related)"},
|
||||
"ids": map[string]any{"type": "array", "items": map[string]any{"type": "number"}, "description": "Observation IDs (delete, supersede, boost)"},
|
||||
"title": map[string]any{"type": "string", "description": "edit: new title"},
|
||||
"subtitle": map[string]any{"type": "string", "description": "edit: new subtitle"},
|
||||
"narrative": map[string]any{"type": "string", "description": "edit: new narrative"},
|
||||
"facts": map[string]any{"type": "array", "items": map[string]any{"type": "string"}, "description": "edit: new facts"},
|
||||
"concepts": map[string]any{"type": "array", "items": map[string]any{"type": "string"}, "description": "edit: new concept tags"},
|
||||
"files_read": map[string]any{"type": "array", "items": map[string]any{"type": "string"}, "description": "edit: new files read"},
|
||||
"files_modified": map[string]any{"type": "array", "items": map[string]any{"type": "string"}, "description": "edit: new files modified"},
|
||||
"scope": map[string]any{"type": "string", "enum": []string{"project", "global"}, "description": "edit: new scope"},
|
||||
"query": map[string]any{"type": "string", "description": "similar: text to find similar observations for"},
|
||||
"project": map[string]any{"type": "string", "description": "Filter by project (similar, by_tag, batch_tag)"},
|
||||
"min_confidence": map[string]any{"type": "number", "minimum": 0.0, "maximum": 1.0, "description": "related: minimum confidence (default 0.5)"},
|
||||
"min_similarity": map[string]any{"type": "number", "minimum": 0.0, "maximum": 1.0, "description": "similar: minimum similarity (default 0.7)"},
|
||||
"limit": map[string]any{"type": "number", "description": "Max results (related, similar, by_tag, batch_tag)"},
|
||||
"delete_vectors": map[string]any{"type": "boolean", "description": "delete: also delete vectors (default true)"},
|
||||
"boost": map[string]any{"type": "number", "minimum": -1.0, "maximum": 1.0, "description": "boost: amount; merge: target boost"},
|
||||
"source_id": map[string]any{"type": "number", "description": "merge: source observation ID (superseded)"},
|
||||
"target_id": map[string]any{"type": "number", "description": "merge: target observation ID (kept)"},
|
||||
"max_depth": map[string]any{"type": "number", "minimum": 1, "maximum": 5, "description": "relationships: traversal hops (default 2)"},
|
||||
"tag": map[string]any{"type": "string", "description": "by_tag: tag/concept to search for"},
|
||||
"tags": map[string]any{"type": "array", "items": map[string]any{"type": "string"}, "description": "tag/batch_tag: tags to apply"},
|
||||
"mode": map[string]any{"type": "string", "enum": []string{"add", "remove", "set"}, "description": "tag: add/remove/set (default add)"},
|
||||
"pattern": map[string]any{"type": "string", "description": "batch_tag: search pattern to match"},
|
||||
"dry_run": map[string]any{"type": "boolean", "description": "batch_tag: preview only (default true)"},
|
||||
"max_matches": map[string]any{"type": "number", "description": "batch_tag: max observations to tag"},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "changes",
|
||||
Description: "Semantic shortcut for finding code changes, refactorings, and modifications.",
|
||||
Name: "memory_admin",
|
||||
Description: "Memory system administration and analytics. Set 'action': stats, health, maintenance_stats, run_maintenance, importance, search_patterns, explain_ranking, temporal_trends, data_quality, export, suggest_consolidations, patterns.",
|
||||
InputSchema: map[string]any{
|
||||
"type": "object",
|
||||
"required": []string{"query"},
|
||||
"required": []string{"action"},
|
||||
"properties": map[string]any{
|
||||
"query": map[string]any{"type": "string", "description": "Natural language query for finding changes"},
|
||||
"dateStart": map[string]any{"type": []string{"string", "number"}},
|
||||
"dateEnd": map[string]any{"type": []string{"string", "number"}},
|
||||
"limit": map[string]any{"type": "number", "default": 20, "minimum": 1, "maximum": 100},
|
||||
"format": map[string]any{"type": "string", "enum": []string{"index", "full"}, "default": "index"},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "how_it_works",
|
||||
Description: "Semantic shortcut for understanding system architecture, design patterns, and implementation details.",
|
||||
InputSchema: map[string]any{
|
||||
"type": "object",
|
||||
"required": []string{"query"},
|
||||
"properties": map[string]any{
|
||||
"query": map[string]any{"type": "string", "description": "Natural language query for understanding how something works"},
|
||||
"dateStart": map[string]any{"type": []string{"string", "number"}},
|
||||
"dateEnd": map[string]any{"type": []string{"string", "number"}},
|
||||
"limit": map[string]any{"type": "number", "default": 20, "minimum": 1, "maximum": 100},
|
||||
"format": map[string]any{"type": "string", "enum": []string{"index", "full"}, "default": "index"},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "find_by_concept",
|
||||
Description: "Find observations tagged with specific concepts.",
|
||||
InputSchema: map[string]any{
|
||||
"type": "object",
|
||||
"required": []string{"concepts"},
|
||||
"properties": map[string]any{
|
||||
"concepts": map[string]any{"type": "string", "description": "Concept tag(s) to filter by"},
|
||||
"type": map[string]any{"type": "string"},
|
||||
"files": map[string]any{"type": "string"},
|
||||
"project": map[string]any{"type": "string"},
|
||||
"dateStart": map[string]any{"type": []string{"string", "number"}},
|
||||
"dateEnd": map[string]any{"type": []string{"string", "number"}},
|
||||
"orderBy": map[string]any{"type": "string", "enum": []string{"date_desc", "date_asc"}, "default": "date_desc"},
|
||||
"limit": map[string]any{"type": "number", "default": 20},
|
||||
"offset": map[string]any{"type": "number", "default": 0},
|
||||
"format": map[string]any{"type": "string", "enum": []string{"index", "full"}, "default": "index"},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "find_by_file",
|
||||
Description: "Find observations related to specific file paths.",
|
||||
InputSchema: map[string]any{
|
||||
"type": "object",
|
||||
"required": []string{"files"},
|
||||
"properties": map[string]any{
|
||||
"files": map[string]any{"type": "string", "description": "File path(s) to filter by"},
|
||||
"type": map[string]any{"type": "string"},
|
||||
"concepts": map[string]any{"type": "string"},
|
||||
"project": map[string]any{"type": "string"},
|
||||
"dateStart": map[string]any{"type": []string{"string", "number"}},
|
||||
"dateEnd": map[string]any{"type": []string{"string", "number"}},
|
||||
"orderBy": map[string]any{"type": "string", "enum": []string{"date_desc", "date_asc"}, "default": "date_desc"},
|
||||
"limit": map[string]any{"type": "number", "default": 20},
|
||||
"offset": map[string]any{"type": "number", "default": 0},
|
||||
"format": map[string]any{"type": "string", "enum": []string{"index", "full"}, "default": "index"},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "find_by_type",
|
||||
Description: "Find observations of specific types.",
|
||||
InputSchema: map[string]any{
|
||||
"type": "object",
|
||||
"required": []string{"type"},
|
||||
"properties": map[string]any{
|
||||
"type": map[string]any{"type": "string", "description": "Observation type(s) to filter by"},
|
||||
"concepts": map[string]any{"type": "string"},
|
||||
"files": map[string]any{"type": "string"},
|
||||
"project": map[string]any{"type": "string"},
|
||||
"dateStart": map[string]any{"type": []string{"string", "number"}},
|
||||
"dateEnd": map[string]any{"type": []string{"string", "number"}},
|
||||
"orderBy": map[string]any{"type": "string", "enum": []string{"date_desc", "date_asc"}, "default": "date_desc"},
|
||||
"limit": map[string]any{"type": "number", "default": 20},
|
||||
"offset": map[string]any{"type": "number", "default": 0},
|
||||
"format": map[string]any{"type": "string", "enum": []string{"index", "full"}, "default": "index"},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "get_recent_context",
|
||||
Description: "Get recent session context for timeline display.",
|
||||
InputSchema: map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"project": map[string]any{"type": "string"},
|
||||
"type": map[string]any{"type": "string"},
|
||||
"concepts": map[string]any{"type": "string"},
|
||||
"files": map[string]any{"type": "string"},
|
||||
"dateStart": map[string]any{"type": []string{"string", "number"}},
|
||||
"dateEnd": map[string]any{"type": []string{"string", "number"}},
|
||||
"limit": map[string]any{"type": "number", "default": 30, "minimum": 1, "maximum": 100},
|
||||
"format": map[string]any{"type": "string", "enum": []string{"index", "full"}, "default": "index"},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "get_context_timeline",
|
||||
Description: "Get timeline of observations around a specific observation ID.",
|
||||
InputSchema: map[string]any{
|
||||
"type": "object",
|
||||
"required": []string{"anchor_id"},
|
||||
"properties": map[string]any{
|
||||
"anchor_id": map[string]any{"type": "number", "description": "Observation ID to use as anchor point"},
|
||||
"before": map[string]any{"type": "number", "default": 10, "minimum": 0, "maximum": 100},
|
||||
"after": map[string]any{"type": "number", "default": 10, "minimum": 0, "maximum": 100},
|
||||
"project": map[string]any{"type": "string"},
|
||||
"type": map[string]any{"type": "string"},
|
||||
"concepts": map[string]any{"type": "string"},
|
||||
"files": map[string]any{"type": "string"},
|
||||
"format": map[string]any{"type": "string", "enum": []string{"index", "full"}, "default": "index"},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "get_timeline_by_query",
|
||||
Description: "Combined search + timeline tool. First searches for observations matching the query, then returns timeline around the best match.",
|
||||
InputSchema: map[string]any{
|
||||
"type": "object",
|
||||
"required": []string{"query"},
|
||||
"properties": map[string]any{
|
||||
"query": map[string]any{"type": "string", "description": "Natural language query to find anchor observation"},
|
||||
"before": map[string]any{"type": "number", "default": 10, "minimum": 0, "maximum": 100},
|
||||
"after": map[string]any{"type": "number", "default": 10, "minimum": 0, "maximum": 100},
|
||||
"project": map[string]any{"type": "string"},
|
||||
"type": map[string]any{"type": "string"},
|
||||
"concepts": map[string]any{"type": "string"},
|
||||
"files": map[string]any{"type": "string"},
|
||||
"dateStart": map[string]any{"type": []string{"string", "number"}},
|
||||
"dateEnd": map[string]any{"type": []string{"string", "number"}},
|
||||
"format": map[string]any{"type": "string", "enum": []string{"index", "full"}, "default": "index"},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "find_related_observations",
|
||||
Description: "Find observations related to a given observation ID filtered by confidence threshold. Returns related observations sorted by confidence score. Useful for discovering relevant context.",
|
||||
InputSchema: map[string]any{
|
||||
"type": "object",
|
||||
"required": []string{"id"},
|
||||
"properties": map[string]any{
|
||||
"id": map[string]any{"type": "number", "description": "Observation ID"},
|
||||
"min_confidence": map[string]any{"type": "number", "default": 0.5, "minimum": 0.0, "maximum": 1.0, "description": "Minimum confidence threshold"},
|
||||
"limit": map[string]any{"type": "number", "default": 20, "minimum": 1, "maximum": 100},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "find_similar_observations",
|
||||
Description: "Find observations semantically similar to a query or observation. Uses vector similarity search to find related content. Useful for detecting duplicates before creating new observations.",
|
||||
InputSchema: map[string]any{
|
||||
"type": "object",
|
||||
"required": []string{"query"},
|
||||
"properties": map[string]any{
|
||||
"query": map[string]any{"type": "string", "description": "Text to find similar observations for"},
|
||||
"project": map[string]any{"type": "string", "description": "Filter by project name"},
|
||||
"min_similarity": map[string]any{"type": "number", "default": 0.7, "minimum": 0.0, "maximum": 1.0, "description": "Minimum similarity threshold (0-1)"},
|
||||
"limit": map[string]any{"type": "number", "default": 10, "minimum": 1, "maximum": 50},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "get_patterns",
|
||||
Description: "Get detected patterns from observations. Patterns represent recurring themes, workflows, or practices discovered across observations.",
|
||||
InputSchema: map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"type": map[string]any{"type": "string", "enum": []string{"workflow", "preference", "best_practice", "anti_pattern", "tooling"}, "description": "Filter by pattern type"},
|
||||
"project": map[string]any{"type": "string", "description": "Filter by project"},
|
||||
"query": map[string]any{"type": "string", "description": "Search patterns by name/description"},
|
||||
"limit": map[string]any{"type": "number", "default": 20, "minimum": 1, "maximum": 100},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "get_memory_stats",
|
||||
Description: "Get statistics about the memory system including observation counts, vector stats, pattern counts, and search metrics.",
|
||||
InputSchema: map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{},
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "bulk_delete_observations",
|
||||
Description: "Delete multiple observations by their IDs. Returns count of successfully deleted observations.",
|
||||
InputSchema: map[string]any{
|
||||
"type": "object",
|
||||
"required": []string{"ids"},
|
||||
"properties": map[string]any{
|
||||
"ids": map[string]any{"type": "array", "items": map[string]any{"type": "number"}, "description": "Array of observation IDs to delete"},
|
||||
"delete_vectors": map[string]any{"type": "boolean", "default": true, "description": "Also delete associated vectors"},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "bulk_mark_superseded",
|
||||
Description: "Mark multiple observations as superseded (stale). Useful for cleanup without permanent deletion.",
|
||||
InputSchema: map[string]any{
|
||||
"type": "object",
|
||||
"required": []string{"ids"},
|
||||
"properties": map[string]any{
|
||||
"ids": map[string]any{"type": "array", "items": map[string]any{"type": "number"}, "description": "Array of observation IDs to mark as superseded"},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "bulk_boost_observations",
|
||||
Description: "Boost or reduce the importance score of multiple observations. Positive values increase importance, negative decrease.",
|
||||
InputSchema: map[string]any{
|
||||
"type": "object",
|
||||
"required": []string{"ids", "boost"},
|
||||
"properties": map[string]any{
|
||||
"ids": map[string]any{"type": "array", "items": map[string]any{"type": "number"}, "description": "Array of observation IDs to boost"},
|
||||
"boost": map[string]any{"type": "number", "minimum": -1.0, "maximum": 1.0, "description": "Boost amount (-1.0 to 1.0)"},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "trigger_maintenance",
|
||||
Description: "Trigger an immediate database maintenance run: optimize/checkpoint the database, clean up old prompts, apply any configured observation retention/stale cleanup, and recalculate importance scores.",
|
||||
InputSchema: map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{},
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "get_maintenance_stats",
|
||||
Description: "Get statistics about the maintenance system including last run time, cleanup counts, and configuration.",
|
||||
InputSchema: map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{},
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "merge_observations",
|
||||
Description: "Merge two observations into one. The target observation is kept and boosted, the source is marked as superseded. Useful for deduplication without data loss.",
|
||||
InputSchema: map[string]any{
|
||||
"type": "object",
|
||||
"required": []string{"source_id", "target_id"},
|
||||
"properties": map[string]any{
|
||||
"source_id": map[string]any{"type": "number", "description": "ID of the observation to merge FROM (will be superseded)"},
|
||||
"target_id": map[string]any{"type": "number", "description": "ID of the observation to merge INTO (will be kept and boosted)"},
|
||||
"boost": map[string]any{"type": "number", "default": 0.1, "minimum": 0, "maximum": 0.5, "description": "Score boost for the target observation (default 0.1)"},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "get_observation",
|
||||
Description: "Get a single observation by its ID. Returns full observation details including all metadata.",
|
||||
InputSchema: map[string]any{
|
||||
"type": "object",
|
||||
"required": []string{"id"},
|
||||
"properties": map[string]any{
|
||||
"id": map[string]any{"type": "number", "description": "Observation ID to retrieve"},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "edit_observation",
|
||||
Description: "Edit an existing observation. Only provided fields will be updated, others remain unchanged. Useful for correcting errors, adding details, or updating scope.",
|
||||
InputSchema: map[string]any{
|
||||
"type": "object",
|
||||
"required": []string{"id"},
|
||||
"properties": map[string]any{
|
||||
"id": map[string]any{"type": "number", "description": "Observation ID to edit"},
|
||||
"title": map[string]any{"type": "string", "description": "New title (optional)"},
|
||||
"subtitle": map[string]any{"type": "string", "description": "New subtitle (optional)"},
|
||||
"narrative": map[string]any{"type": "string", "description": "New narrative text (optional)"},
|
||||
"facts": map[string]any{"type": "array", "items": map[string]any{"type": "string"}, "description": "New facts array (optional)"},
|
||||
"concepts": map[string]any{"type": "array", "items": map[string]any{"type": "string"}, "description": "New concept tags (optional)"},
|
||||
"files_read": map[string]any{"type": "array", "items": map[string]any{"type": "string"}, "description": "New files read list (optional)"},
|
||||
"files_modified": map[string]any{"type": "array", "items": map[string]any{"type": "string"}, "description": "New files modified list (optional)"},
|
||||
"scope": map[string]any{"type": "string", "enum": []string{"project", "global"}, "description": "New scope (optional)"},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "get_observation_quality",
|
||||
Description: "Get quality metrics for an observation. Returns completeness score, usage stats, and improvement suggestions.",
|
||||
InputSchema: map[string]any{
|
||||
"type": "object",
|
||||
"required": []string{"id"},
|
||||
"properties": map[string]any{
|
||||
"id": map[string]any{"type": "number", "description": "Observation ID to analyze"},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "suggest_consolidations",
|
||||
Description: "Find observations that could be merged or consolidated. Returns groups of similar observations with merge recommendations.",
|
||||
InputSchema: map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"project": map[string]any{"type": "string", "description": "Filter by project"},
|
||||
"min_similarity": map[string]any{"type": "number", "default": 0.8, "minimum": 0.5, "maximum": 1.0, "description": "Minimum similarity threshold for grouping"},
|
||||
"limit": map[string]any{"type": "number", "default": 10, "minimum": 1, "maximum": 50, "description": "Maximum groups to return"},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "tag_observation",
|
||||
Description: "Add or remove concept tags from an observation. Tags help with organization and filtering. Use mode 'add' to add new tags, 'remove' to remove specific tags, or 'set' to replace all tags.",
|
||||
InputSchema: map[string]any{
|
||||
"type": "object",
|
||||
"required": []string{"id", "tags"},
|
||||
"properties": map[string]any{
|
||||
"id": map[string]any{"type": "number", "description": "Observation ID to tag"},
|
||||
"tags": map[string]any{"type": "array", "items": map[string]any{"type": "string"}, "description": "Tags to add, remove, or set"},
|
||||
"mode": map[string]any{"type": "string", "enum": []string{"add", "remove", "set"}, "default": "add", "description": "Operation mode: 'add' appends tags, 'remove' removes specific tags, 'set' replaces all tags"},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "get_observations_by_tag",
|
||||
Description: "Find all observations that have a specific concept tag. Useful for browsing by category.",
|
||||
InputSchema: map[string]any{
|
||||
"type": "object",
|
||||
"required": []string{"tag"},
|
||||
"properties": map[string]any{
|
||||
"tag": map[string]any{"type": "string", "description": "Tag/concept to search for"},
|
||||
"project": map[string]any{"type": "string", "description": "Filter by project (optional)"},
|
||||
"limit": map[string]any{"type": "number", "default": 50, "minimum": 1, "maximum": 200, "description": "Maximum observations to return"},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "get_temporal_trends",
|
||||
Description: "Analyze observation creation patterns over time. Returns daily counts, peak activity times, and trend insights. Useful for understanding work patterns.",
|
||||
InputSchema: map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"project": map[string]any{"type": "string", "description": "Filter by project (optional)"},
|
||||
"days": map[string]any{"type": "number", "default": 30, "minimum": 1, "maximum": 365, "description": "Number of days to analyze"},
|
||||
"group_by": map[string]any{"type": "string", "enum": []string{"day", "week", "hour_of_day"}, "default": "day", "description": "How to group the data"},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "get_data_quality_report",
|
||||
Description: "Get a comprehensive quality assessment of observations. Shows completeness distribution, common issues, and improvement suggestions.",
|
||||
InputSchema: map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"project": map[string]any{"type": "string", "description": "Filter by project (optional)"},
|
||||
"limit": map[string]any{"type": "number", "default": 100, "minimum": 10, "maximum": 500, "description": "Number of observations to analyze"},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "batch_tag_by_pattern",
|
||||
Description: "Apply tags to observations matching a pattern. Useful for retroactive organization and categorization.",
|
||||
InputSchema: map[string]any{
|
||||
"type": "object",
|
||||
"required": []string{"pattern", "tags"},
|
||||
"properties": map[string]any{
|
||||
"pattern": map[string]any{"type": "string", "description": "Search pattern to match (searches title, narrative, facts)"},
|
||||
"tags": map[string]any{"type": "array", "items": map[string]any{"type": "string"}, "description": "Tags to add to matching observations"},
|
||||
"project": map[string]any{"type": "string", "description": "Filter by project (optional)"},
|
||||
"dry_run": map[string]any{"type": "boolean", "default": true, "description": "If true, only preview matches without applying tags"},
|
||||
"max_matches": map[string]any{"type": "number", "default": 100, "minimum": 1, "maximum": 500, "description": "Maximum observations to tag"},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "explain_search_ranking",
|
||||
Description: "Debug search results by showing score breakdown for top matches. Explains why each observation ranked where it did.",
|
||||
InputSchema: map[string]any{
|
||||
"type": "object",
|
||||
"required": []string{"query"},
|
||||
"properties": map[string]any{
|
||||
"query": map[string]any{"type": "string", "description": "Search query to analyze"},
|
||||
"project": map[string]any{"type": "string", "description": "Project context for search"},
|
||||
"top_n": map[string]any{"type": "number", "default": 5, "minimum": 1, "maximum": 20, "description": "Number of top results to explain"},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "export_observations",
|
||||
Description: "Export observations in various formats for backup or analysis.",
|
||||
InputSchema: map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"format": map[string]any{"type": "string", "enum": []string{"json", "jsonl", "markdown"}, "default": "json", "description": "Export format"},
|
||||
"project": map[string]any{"type": "string", "description": "Filter by project (optional)"},
|
||||
"limit": map[string]any{"type": "number", "default": 100, "minimum": 1, "maximum": 1000, "description": "Maximum observations to export"},
|
||||
"date_start": map[string]any{"type": "number", "description": "Filter by creation date (epoch milliseconds)"},
|
||||
"date_end": map[string]any{"type": "number", "description": "Filter by creation date (epoch milliseconds)"},
|
||||
"obs_type": map[string]any{"type": "string", "description": "Filter by observation type"},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "check_system_health",
|
||||
Description: "Comprehensive system health check. Returns status of all subsystems (database, vectors, cache, search) with actionable diagnostics.",
|
||||
InputSchema: map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{},
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "analyze_search_patterns",
|
||||
Description: "Analyze search query patterns to identify common searches, missed queries, and optimization opportunities.",
|
||||
InputSchema: map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"days": map[string]any{"type": "number", "default": 7, "minimum": 1, "maximum": 30, "description": "Number of days to analyze"},
|
||||
"top_n": map[string]any{"type": "number", "default": 10, "minimum": 1, "maximum": 50, "description": "Number of top patterns to return"},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "get_observation_relationships",
|
||||
Description: "Get relationship graph for an observation. Shows how observations relate to each other (depends_on, extends, conflicts_with, supersedes). Useful for understanding dependencies and context.",
|
||||
InputSchema: map[string]any{
|
||||
"type": "object",
|
||||
"required": []string{"id"},
|
||||
"properties": map[string]any{
|
||||
"id": map[string]any{"type": "number", "description": "Observation ID to analyze relationships for"},
|
||||
"max_depth": map[string]any{"type": "number", "default": 2, "minimum": 1, "maximum": 5, "description": "How many hops to traverse (1=direct, 2=neighbors of neighbors)"},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "get_observation_scoring_breakdown",
|
||||
Description: "Get detailed scoring breakdown for an observation. Shows how importance scores are calculated including type weight, recency decay, feedback contribution, concept boost, and retrieval frequency. Useful for understanding why observations are ranked the way they are.",
|
||||
InputSchema: map[string]any{
|
||||
"type": "object",
|
||||
"required": []string{"id"},
|
||||
"properties": map[string]any{
|
||||
"id": map[string]any{"type": "number", "description": "Observation ID to get scoring breakdown for"},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "analyze_observation_importance",
|
||||
Description: "Analyze observation importance patterns in a project. Returns statistics on feedback distribution, top-scoring observations, most-retrieved observations, and concept weights. Useful for understanding what makes observations valuable.",
|
||||
InputSchema: map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"project": map[string]any{"type": "string", "description": "Project to analyze (optional, analyzes all if omitted)"},
|
||||
"include_top_scored": map[string]any{"type": "boolean", "default": true, "description": "Include top-scoring observations"},
|
||||
"include_most_retrieved": map[string]any{"type": "boolean", "default": true, "description": "Include most-retrieved observations"},
|
||||
"include_concept_weights": map[string]any{"type": "boolean", "default": true, "description": "Include concept weight analysis"},
|
||||
"limit": map[string]any{"type": "number", "default": 10, "minimum": 1, "maximum": 50, "description": "Number of top observations to include"},
|
||||
"action": map[string]any{"type": "string", "enum": []string{"stats", "health", "maintenance_stats", "run_maintenance", "importance", "search_patterns", "explain_ranking", "temporal_trends", "data_quality", "export", "suggest_consolidations", "patterns"}, "description": "Operation to perform"},
|
||||
"project": map[string]any{"type": "string", "description": "Filter by project (importance, temporal_trends, data_quality, export, suggest_consolidations, patterns, explain_ranking)"},
|
||||
"query": map[string]any{"type": "string", "description": "explain_ranking: query to analyze; patterns: search by name/description"},
|
||||
"top_n": map[string]any{"type": "number", "description": "Top results (explain_ranking, search_patterns)"},
|
||||
"days": map[string]any{"type": "number", "description": "Days to analyze (search_patterns, temporal_trends)"},
|
||||
"group_by": map[string]any{"type": "string", "enum": []string{"day", "week", "hour_of_day"}, "description": "temporal_trends: grouping (default day)"},
|
||||
"limit": map[string]any{"type": "number", "description": "Max results (importance, data_quality, export, suggest_consolidations, patterns)"},
|
||||
"min_similarity": map[string]any{"type": "number", "minimum": 0.5, "maximum": 1.0, "description": "suggest_consolidations: grouping threshold (default 0.8)"},
|
||||
"format": map[string]any{"type": "string", "enum": []string{"json", "jsonl", "markdown"}, "description": "export: output format (default json)"},
|
||||
"date_start": map[string]any{"type": "number", "description": "export: filter start (epoch ms)"},
|
||||
"date_end": map[string]any{"type": "number", "description": "export: filter end (epoch ms)"},
|
||||
"obs_type": map[string]any{"type": "string", "description": "export: filter by observation type"},
|
||||
"type": map[string]any{"type": "string", "enum": []string{"workflow", "preference", "best_practice", "anti_pattern", "tooling"}, "description": "patterns: filter by pattern type"},
|
||||
"include_top_scored": map[string]any{"type": "boolean", "description": "importance: include top-scoring (default true)"},
|
||||
"include_most_retrieved": map[string]any{"type": "boolean", "description": "importance: include most-retrieved (default true)"},
|
||||
"include_concept_weights": map[string]any{"type": "boolean", "description": "importance: include concept weights (default true)"},
|
||||
},
|
||||
},
|
||||
},
|
||||
@@ -833,6 +456,55 @@ type searchArgs struct {
|
||||
Offset int `json:"offset"`
|
||||
}
|
||||
|
||||
// observationActions maps the 'observation' tool's action values to the internal tool name that implements them.
|
||||
var observationActions = map[string]string{
|
||||
"get": "get_observation",
|
||||
"edit": "edit_observation",
|
||||
"delete": "bulk_delete_observations",
|
||||
"supersede": "bulk_mark_superseded",
|
||||
"boost": "bulk_boost_observations",
|
||||
"merge": "merge_observations",
|
||||
"related": "find_related_observations",
|
||||
"similar": "find_similar_observations",
|
||||
"quality": "get_observation_quality",
|
||||
"relationships": "get_observation_relationships",
|
||||
"scoring": "get_observation_scoring_breakdown",
|
||||
"tag": "tag_observation",
|
||||
"by_tag": "get_observations_by_tag",
|
||||
"batch_tag": "batch_tag_by_pattern",
|
||||
}
|
||||
|
||||
// adminActions maps the 'memory_admin' tool's action values to the internal tool name that implements them.
|
||||
var adminActions = map[string]string{
|
||||
"stats": "get_memory_stats",
|
||||
"health": "check_system_health",
|
||||
"maintenance_stats": "get_maintenance_stats",
|
||||
"run_maintenance": "trigger_maintenance",
|
||||
"importance": "analyze_observation_importance",
|
||||
"search_patterns": "analyze_search_patterns",
|
||||
"explain_ranking": "explain_search_ranking",
|
||||
"temporal_trends": "get_temporal_trends",
|
||||
"data_quality": "get_data_quality_report",
|
||||
"export": "export_observations",
|
||||
"suggest_consolidations": "suggest_consolidations",
|
||||
"patterns": "get_patterns",
|
||||
}
|
||||
|
||||
// dispatchAction routes a multiplexed tool call to its underlying implementation by reading the "action" field.
|
||||
func (s *Server) dispatchAction(ctx context.Context, tool string, actions map[string]string, args json.RawMessage) (string, error) {
|
||||
var a struct {
|
||||
Action string `json:"action"`
|
||||
}
|
||||
if err := json.Unmarshal(args, &a); err != nil {
|
||||
return "", fmt.Errorf("%s: invalid arguments: %w", tool, err)
|
||||
}
|
||||
target, ok := actions[a.Action]
|
||||
if !ok {
|
||||
return "", fmt.Errorf("%s: unknown action %q", tool, a.Action)
|
||||
}
|
||||
return s.callTool(ctx, target, args)
|
||||
}
|
||||
|
||||
// callTool dispatches to the appropriate tool handler by proxying to the worker HTTP API.
|
||||
func (s *Server) callTool(ctx context.Context, name string, args json.RawMessage) (string, error) {
|
||||
// Parse common search params used by many tools
|
||||
@@ -849,6 +521,10 @@ func (s *Server) callTool(ctx context.Context, name string, args json.RawMessage
|
||||
// --- Search-based tools: proxy to GET /api/context/search ---
|
||||
case "search":
|
||||
return s.handleSearchProxy(ctx, sa)
|
||||
case "observation":
|
||||
return s.dispatchAction(ctx, "observation", observationActions, args)
|
||||
case "memory_admin":
|
||||
return s.dispatchAction(ctx, "memory_admin", adminActions, args)
|
||||
case "decisions":
|
||||
sa.ObsType = "decision"
|
||||
return s.handleSearchProxy(ctx, sa)
|
||||
|
||||
+59
-44
@@ -434,12 +434,7 @@ func TestHandleToolsList(t *testing.T) {
|
||||
toolNames[tool.Name] = true
|
||||
}
|
||||
|
||||
expectedTools := []string{
|
||||
"search", "timeline", "decisions", "changes",
|
||||
"how_it_works", "find_by_concept", "find_by_file",
|
||||
"find_by_type", "get_recent_context", "get_context_timeline",
|
||||
"get_timeline_by_query",
|
||||
}
|
||||
expectedTools := []string{"search", "timeline", "observation", "memory_admin"}
|
||||
|
||||
for _, name := range expectedTools {
|
||||
assert.True(t, toolNames[name], "expected tool %s to be present", name)
|
||||
@@ -1707,43 +1702,10 @@ func TestCallTool_ToolNameRecognition(t *testing.T) {
|
||||
|
||||
// Verify all expected tools are registered
|
||||
expectedTools := map[string]bool{
|
||||
"search": true,
|
||||
"timeline": true,
|
||||
"decisions": true,
|
||||
"changes": true,
|
||||
"how_it_works": true,
|
||||
"find_by_concept": true,
|
||||
"find_by_file": true,
|
||||
"find_by_type": true,
|
||||
"get_recent_context": true,
|
||||
"get_context_timeline": true,
|
||||
"get_timeline_by_query": true,
|
||||
"find_related_observations": true,
|
||||
"find_similar_observations": true,
|
||||
"get_patterns": true,
|
||||
"get_memory_stats": true,
|
||||
"bulk_delete_observations": true,
|
||||
"bulk_mark_superseded": true,
|
||||
"bulk_boost_observations": true,
|
||||
"trigger_maintenance": true,
|
||||
"get_maintenance_stats": true,
|
||||
"merge_observations": true,
|
||||
"get_observation": true,
|
||||
"edit_observation": true,
|
||||
"get_observation_quality": true,
|
||||
"suggest_consolidations": true,
|
||||
"tag_observation": true,
|
||||
"get_observations_by_tag": true,
|
||||
"get_temporal_trends": true,
|
||||
"get_data_quality_report": true,
|
||||
"batch_tag_by_pattern": true,
|
||||
"explain_search_ranking": true,
|
||||
"export_observations": true,
|
||||
"check_system_health": true,
|
||||
"analyze_search_patterns": true,
|
||||
"get_observation_relationships": true,
|
||||
"get_observation_scoring_breakdown": true,
|
||||
"analyze_observation_importance": true,
|
||||
"search": true,
|
||||
"timeline": true,
|
||||
"observation": true,
|
||||
"memory_admin": true,
|
||||
}
|
||||
|
||||
foundTools := make(map[string]bool)
|
||||
@@ -3219,7 +3181,7 @@ func TestRun_SemaphoreDoesNotBlockMainLoop(t *testing.T) {
|
||||
_, _ = io.WriteString(w, `{}`)
|
||||
}))
|
||||
defer func() {
|
||||
handlerCancel() // unblock all handlers first
|
||||
handlerCancel() // unblock all handlers first
|
||||
ts.CloseClientConnections()
|
||||
ts.Close()
|
||||
}()
|
||||
@@ -3418,3 +3380,56 @@ func TestHandleToolsCall_ErrorUsesIsError(t *testing.T) {
|
||||
require.True(t, ok)
|
||||
assert.Contains(t, errText, "unknown tool: nonexistent_tool")
|
||||
}
|
||||
|
||||
// TestDispatchAction verifies the multiplexed tool action routing.
|
||||
// Only action-routing error cases are asserted: a nil HTTP client means valid
|
||||
// actions fail later at the HTTP layer, but unknown/invalid actions short-circuit
|
||||
// before any HTTP call.
|
||||
func TestDispatchAction(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
server := NewServer(nil, "", "", "1.0.0")
|
||||
ctx := context.Background()
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
tool string
|
||||
args string
|
||||
wantErr string
|
||||
}{
|
||||
{
|
||||
name: "observation unknown action",
|
||||
tool: "observation",
|
||||
args: `{"action":"bogus"}`,
|
||||
wantErr: "unknown action",
|
||||
},
|
||||
{
|
||||
name: "memory_admin unknown action",
|
||||
tool: "memory_admin",
|
||||
args: `{"action":"bogus"}`,
|
||||
wantErr: "unknown action",
|
||||
},
|
||||
{
|
||||
name: "observation no action field",
|
||||
tool: "observation",
|
||||
args: `{}`,
|
||||
wantErr: "unknown action",
|
||||
},
|
||||
{
|
||||
name: "observation invalid action JSON",
|
||||
tool: "observation",
|
||||
args: `{"action": 123}`,
|
||||
wantErr: "invalid arguments",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
_, err := server.callTool(ctx, tt.tool, json.RawMessage(tt.args))
|
||||
require.Error(t, err)
|
||||
assert.Contains(t, err.Error(), tt.wantErr)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
+10
-10
@@ -85,23 +85,23 @@ func EnsureModel(assetName, url, expectedSHA256 string) (string, error) {
|
||||
return modelPath, nil
|
||||
}
|
||||
log.Warn().Str("model", assetName).Msg("Cached model corrupted, re-downloading")
|
||||
os.Remove(modelPath)
|
||||
_ = os.Remove(modelPath) // best-effort cleanup
|
||||
}
|
||||
} else if len(currentVersion) > 0 {
|
||||
log.Info().Str("old", string(currentVersion)).Str("new", ModelVersion).Msg("Model version updated, clearing cache")
|
||||
os.RemoveAll(dir)
|
||||
_ = os.RemoveAll(dir) // best-effort cleanup
|
||||
if err := os.MkdirAll(dir, 0700); err != nil {
|
||||
return "", fmt.Errorf("create model cache dir: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
if err := downloadWithRetries(url, modelPath); err != nil {
|
||||
os.Remove(modelPath)
|
||||
_ = os.Remove(modelPath) // best-effort cleanup
|
||||
return "", fmt.Errorf("download %s: %w", assetName, err)
|
||||
}
|
||||
|
||||
if valid, err := verifyChecksum(modelPath, expectedSHA256); err != nil || !valid {
|
||||
os.Remove(modelPath)
|
||||
_ = os.Remove(modelPath) // best-effort cleanup
|
||||
return "", fmt.Errorf("downloaded model %s failed checksum verification — the file at %s may have changed", assetName, url)
|
||||
}
|
||||
|
||||
@@ -137,7 +137,7 @@ func downloadFile(url, dest string) error {
|
||||
if err != nil {
|
||||
return fmt.Errorf("HTTP GET: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
defer func() { _ = resp.Body.Close() }() // best-effort cleanup
|
||||
|
||||
if resp.StatusCode < 200 || resp.StatusCode >= 400 {
|
||||
return fmt.Errorf("HTTP %d", resp.StatusCode)
|
||||
@@ -148,7 +148,7 @@ func downloadFile(url, dest string) error {
|
||||
if err != nil {
|
||||
return fmt.Errorf("create temp file: %w", err)
|
||||
}
|
||||
cleanup := func() { f.Close(); os.Remove(tmp) }
|
||||
cleanup := func() { _ = f.Close(); _ = os.Remove(tmp) } // best-effort cleanup
|
||||
|
||||
name := filepath.Base(dest)
|
||||
reader := &progressReader{
|
||||
@@ -166,10 +166,10 @@ func downloadFile(url, dest string) error {
|
||||
cleanup()
|
||||
return fmt.Errorf("sync file: %w", err)
|
||||
}
|
||||
f.Close()
|
||||
_ = f.Close() // best-effort cleanup; data already flushed via Sync above
|
||||
|
||||
if err := os.Rename(tmp, dest); err != nil {
|
||||
os.Remove(tmp)
|
||||
_ = os.Remove(tmp) // best-effort cleanup
|
||||
return fmt.Errorf("commit file: %w", err)
|
||||
}
|
||||
|
||||
@@ -221,7 +221,7 @@ func verifyChecksum(path, expectedHex string) (bool, error) {
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
defer f.Close()
|
||||
defer func() { _ = f.Close() }() // best-effort cleanup
|
||||
|
||||
h := sha256.New()
|
||||
if _, err := io.Copy(h, f); err != nil {
|
||||
@@ -249,7 +249,7 @@ func CleanStale() error {
|
||||
tmpPath := filepath.Join(dir, e.Name())
|
||||
if info, err := e.Info(); err == nil {
|
||||
if time.Since(info.ModTime()) > time.Hour {
|
||||
os.Remove(tmpPath)
|
||||
_ = os.Remove(tmpPath) // best-effort cleanup
|
||||
log.Debug().Str("file", tmpPath).Msg("Cleaned stale temp file")
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,235 @@
|
||||
// Package models provides runtime model download and caching.
|
||||
package models
|
||||
|
||||
import (
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// sha256Hex returns the hex-encoded SHA-256 of b.
|
||||
func sha256Hex(t *testing.T, b []byte) string {
|
||||
t.Helper()
|
||||
h := sha256.Sum256(b)
|
||||
return hex.EncodeToString(h[:])
|
||||
}
|
||||
|
||||
// isolateCacheDir points ModelsDir() at a fresh temp directory by overriding
|
||||
// the OS cache-dir resolution env vars (HOME and XDG_CACHE_HOME, which
|
||||
// os.UserCacheDir honors on Linux and macOS). It also clears
|
||||
// CLAUDE_MNEMONIC_MODEL_DIR so EnsureModel takes the download/cache path
|
||||
// rather than the local-dir shortcut. It returns the resulting models dir.
|
||||
func isolateCacheDir(t *testing.T) string {
|
||||
t.Helper()
|
||||
home := t.TempDir()
|
||||
t.Setenv("HOME", home)
|
||||
t.Setenv("XDG_CACHE_HOME", filepath.Join(home, ".cache"))
|
||||
t.Setenv("CLAUDE_MNEMONIC_MODEL_DIR", "")
|
||||
|
||||
dir := ModelsDir()
|
||||
require.True(t, strings.HasPrefix(dir, home),
|
||||
"ModelsDir() %q must live under temp HOME %q for hermetic test", dir, home)
|
||||
return dir
|
||||
}
|
||||
|
||||
// fileServer returns an httptest server that serves body for any request and
|
||||
// records how many requests it received.
|
||||
func fileServer(t *testing.T, body []byte) (*httptest.Server, *int32) {
|
||||
t.Helper()
|
||||
var hits int32
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
atomic.AddInt32(&hits, 1)
|
||||
_, _ = w.Write(body)
|
||||
}))
|
||||
t.Cleanup(srv.Close)
|
||||
return srv, &hits
|
||||
}
|
||||
|
||||
// TestEnsureModel_Download_CachesValidModel verifies the happy path: a model is
|
||||
// downloaded, passes checksum verification, and the version file is written.
|
||||
func TestEnsureModel_Download_CachesValidModel(t *testing.T) {
|
||||
dir := isolateCacheDir(t)
|
||||
body := []byte("model-bytes-happy-path")
|
||||
srv, hits := fileServer(t, body)
|
||||
|
||||
const asset = "test-model.onnx"
|
||||
path, err := EnsureModel(asset, srv.URL, sha256Hex(t, body))
|
||||
require.NoError(t, err)
|
||||
|
||||
assert.Equal(t, filepath.Join(dir, asset), path)
|
||||
got, err := os.ReadFile(path)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, body, got, "cached file content must match downloaded body")
|
||||
|
||||
ver, err := os.ReadFile(filepath.Join(dir, ".model_version"))
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, ModelVersion, string(ver), "version file must record current ModelVersion")
|
||||
|
||||
assert.Equal(t, int32(1), atomic.LoadInt32(hits), "exactly one download on first fetch")
|
||||
}
|
||||
|
||||
// TestEnsureModel_ChecksumMismatch_FailsAndRemoves verifies that a downloaded
|
||||
// file whose content does not match the expected SHA-256 is rejected with a
|
||||
// verification error and is not left behind in the cache.
|
||||
func TestEnsureModel_ChecksumMismatch_FailsAndRemoves(t *testing.T) {
|
||||
dir := isolateCacheDir(t)
|
||||
body := []byte("the-actual-downloaded-bytes")
|
||||
srv, _ := fileServer(t, body)
|
||||
|
||||
const asset = "mismatch-model.onnx"
|
||||
wrongSum := sha256Hex(t, []byte("a-different-expected-payload"))
|
||||
|
||||
path, err := EnsureModel(asset, srv.URL, wrongSum)
|
||||
require.Error(t, err)
|
||||
assert.Empty(t, path)
|
||||
assert.Contains(t, err.Error(), "checksum verification")
|
||||
|
||||
_, statErr := os.Stat(filepath.Join(dir, asset))
|
||||
assert.True(t, os.IsNotExist(statErr),
|
||||
"corrupt download must be removed from cache, got statErr=%v", statErr)
|
||||
}
|
||||
|
||||
// TestEnsureModel_VersionBump_InvalidatesCache verifies that when the on-disk
|
||||
// .model_version differs from ModelVersion, the cache directory is cleared and
|
||||
// the model is re-downloaded fresh.
|
||||
func TestEnsureModel_VersionBump_InvalidatesCache(t *testing.T) {
|
||||
dir := isolateCacheDir(t)
|
||||
require.NoError(t, os.MkdirAll(dir, 0700))
|
||||
|
||||
// Seed a stale cache: an old version file, a leftover model file, and an
|
||||
// unrelated sidecar file. Only the version-bump branch wipes the whole dir
|
||||
// (os.RemoveAll), so the sidecar's removal uniquely proves invalidation ran
|
||||
// — a plain checksum-driven re-download would leave the sidecar in place.
|
||||
const asset = "versioned-model.onnx"
|
||||
staleModel := filepath.Join(dir, asset)
|
||||
staleSidecar := filepath.Join(dir, "leftover-sidecar.bin")
|
||||
require.NoError(t, os.WriteFile(staleModel, []byte("stale-garbage"), 0600))
|
||||
require.NoError(t, os.WriteFile(staleSidecar, []byte("unrelated-leftover"), 0600))
|
||||
require.NoError(t, os.WriteFile(filepath.Join(dir, ".model_version"), []byte("OLD-VERSION"), 0600))
|
||||
|
||||
body := []byte("freshly-downloaded-model-bytes")
|
||||
srv, hits := fileServer(t, body)
|
||||
|
||||
path, err := EnsureModel(asset, srv.URL, sha256Hex(t, body))
|
||||
require.NoError(t, err)
|
||||
|
||||
got, err := os.ReadFile(path)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, body, got, "stale cached file must be replaced by fresh download")
|
||||
|
||||
ver, err := os.ReadFile(filepath.Join(dir, ".model_version"))
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, ModelVersion, string(ver), "version file must be rewritten to current ModelVersion")
|
||||
|
||||
assert.Equal(t, int32(1), atomic.LoadInt32(hits),
|
||||
"version bump must force exactly one re-download, not serve the stale file")
|
||||
|
||||
_, sidecarErr := os.Stat(staleSidecar)
|
||||
assert.True(t, os.IsNotExist(sidecarErr),
|
||||
"version bump must wipe the whole cache dir, removing unrelated sidecar files")
|
||||
}
|
||||
|
||||
// TestEnsureModel_SameVersionValidCache_NoDownload verifies that a matching
|
||||
// version file plus a checksum-valid cached file short-circuits the download.
|
||||
func TestEnsureModel_SameVersionValidCache_NoDownload(t *testing.T) {
|
||||
dir := isolateCacheDir(t)
|
||||
require.NoError(t, os.MkdirAll(dir, 0700))
|
||||
|
||||
const asset = "cached-model.onnx"
|
||||
body := []byte("already-cached-valid-bytes")
|
||||
require.NoError(t, os.WriteFile(filepath.Join(dir, asset), body, 0600))
|
||||
require.NoError(t, os.WriteFile(filepath.Join(dir, ".model_version"), []byte(ModelVersion), 0600))
|
||||
|
||||
// Server would record a hit if EnsureModel wrongly re-downloaded.
|
||||
srv, hits := fileServer(t, []byte("should-not-be-fetched"))
|
||||
|
||||
path, err := EnsureModel(asset, srv.URL, sha256Hex(t, body))
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, filepath.Join(dir, asset), path)
|
||||
assert.Equal(t, int32(0), atomic.LoadInt32(hits), "valid same-version cache must not trigger a download")
|
||||
}
|
||||
|
||||
// TestEnsureModel_RetryExhaustion_ReturnsError verifies that when every download
|
||||
// attempt fails (server always returns HTTP 500), EnsureModel exhausts its
|
||||
// retries and returns an error reporting the attempt count, leaving no model
|
||||
// file behind.
|
||||
//
|
||||
// Note: downloadWithRetries sleeps with exponential backoff between attempts,
|
||||
// so this test takes ~6s by design; the backoff is production behavior and is
|
||||
// intentionally not stubbed.
|
||||
func TestEnsureModel_RetryExhaustion_ReturnsError(t *testing.T) {
|
||||
if testing.Short() {
|
||||
t.Skip("skipping retry-exhaustion test (real backoff sleeps) in -short mode")
|
||||
}
|
||||
|
||||
dir := isolateCacheDir(t)
|
||||
|
||||
var hits int32
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
atomic.AddInt32(&hits, 1)
|
||||
w.WriteHeader(http.StatusInternalServerError)
|
||||
}))
|
||||
t.Cleanup(srv.Close)
|
||||
|
||||
const asset = "unreachable-model.onnx"
|
||||
path, err := EnsureModel(asset, srv.URL, sha256Hex(t, []byte("never-served")))
|
||||
require.Error(t, err)
|
||||
assert.Empty(t, path)
|
||||
assert.Contains(t, err.Error(), "attempts", "error must report retry exhaustion")
|
||||
|
||||
assert.Equal(t, int32(maxRetries), atomic.LoadInt32(&hits),
|
||||
"server must be hit exactly maxRetries times")
|
||||
|
||||
_, statErr := os.Stat(filepath.Join(dir, asset))
|
||||
assert.True(t, os.IsNotExist(statErr), "failed download must not leave a model file")
|
||||
}
|
||||
|
||||
// TestEnsureModel_LocalDir verifies the CLAUDE_MNEMONIC_MODEL_DIR shortcut:
|
||||
// a present, checksum-valid file is returned directly, and a checksum mismatch
|
||||
// or missing file is reported without any network access.
|
||||
func TestEnsureModel_LocalDir(t *testing.T) {
|
||||
body := []byte("local-dir-model-bytes")
|
||||
validSum := sha256Hex(t, body)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
writeFile bool
|
||||
expectSum string
|
||||
wantErr string
|
||||
}{
|
||||
{name: "valid_local_model", writeFile: true, expectSum: validSum, wantErr: ""},
|
||||
{name: "checksum_mismatch", writeFile: true, expectSum: sha256Hex(t, []byte("other")), wantErr: "checksum mismatch"},
|
||||
{name: "missing_file", writeFile: false, expectSum: validSum, wantErr: "not found"},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
localDir := t.TempDir()
|
||||
t.Setenv("CLAUDE_MNEMONIC_MODEL_DIR", localDir)
|
||||
|
||||
const asset = "local-model.onnx"
|
||||
if tt.writeFile {
|
||||
require.NoError(t, os.WriteFile(filepath.Join(localDir, asset), body, 0600))
|
||||
}
|
||||
|
||||
path, err := EnsureModel(asset, "http://127.0.0.1:0/should-not-be-used", tt.expectSum)
|
||||
if tt.wantErr != "" {
|
||||
require.Error(t, err)
|
||||
assert.Contains(t, err.Error(), tt.wantErr)
|
||||
assert.Empty(t, path)
|
||||
return
|
||||
}
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, filepath.Join(localDir, asset), path)
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -34,6 +34,10 @@ const (
|
||||
CrossEncoderAssetURL = "https://huggingface.co/cross-encoder/ms-marco-MiniLM-L6-v2/resolve/main/onnx/model.onnx"
|
||||
// CrossEncoderModelSHA256 is the expected SHA-256 of the cross-encoder model file.
|
||||
CrossEncoderModelSHA256 = "5d3e70fd0c9ff14b9b5169a51e957b7a9c74897afd0a35ce4bd318150c1d4d4a"
|
||||
|
||||
// DefaultAlpha is the default score-combining weight, favoring the
|
||||
// cross-encoder: combined = alpha*rerank + (1-alpha)*original.
|
||||
DefaultAlpha = 0.7
|
||||
)
|
||||
|
||||
// Candidate represents a search result candidate for reranking.
|
||||
@@ -79,7 +83,7 @@ type Config struct {
|
||||
// DefaultConfig returns sensible defaults for reranking.
|
||||
func DefaultConfig() Config {
|
||||
return Config{
|
||||
Alpha: 0.7, // Favor cross-encoder by default
|
||||
Alpha: DefaultAlpha, // Favor cross-encoder by default
|
||||
}
|
||||
}
|
||||
|
||||
@@ -121,7 +125,7 @@ func NewService(cfg Config) (*Service, error) {
|
||||
|
||||
alpha := cfg.Alpha
|
||||
if alpha <= 0 || alpha > 1 {
|
||||
alpha = 0.7
|
||||
alpha = DefaultAlpha
|
||||
}
|
||||
|
||||
return &Service{
|
||||
|
||||
@@ -8,6 +8,11 @@ import (
|
||||
"github.com/lukaszraczylo/claude-mnemonic/pkg/models"
|
||||
)
|
||||
|
||||
// retrievalLogFactor scales the diminishing-returns retrieval boost:
|
||||
// log2(retrieval_count + 1) is multiplied by this factor before applying
|
||||
// the configured retrieval weight.
|
||||
const retrievalLogFactor = 0.1
|
||||
|
||||
// Calculator computes importance scores for observations.
|
||||
type Calculator struct {
|
||||
config *models.ScoringConfig
|
||||
@@ -72,7 +77,7 @@ func (c *Calculator) CalculateComponents(obs *models.Observation, now time.Time)
|
||||
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
|
||||
retrievalBoost := math.Log2(float64(obs.RetrievalCount)+1) * retrievalLogFactor
|
||||
retrievalContrib = retrievalBoost * c.config.RetrievalWeight
|
||||
}
|
||||
|
||||
|
||||
@@ -254,12 +254,16 @@ func (u *Updater) ApplyUpdate(ctx context.Context, info *UpdateInfo) error {
|
||||
u.setStatus("verifying", 0.2, "Verifying signature...")
|
||||
|
||||
if info.ChecksumsURL != "" && info.BundleURL != "" {
|
||||
// Fail-closed: a signature bundle is published for this release, so a
|
||||
// verification failure (including cosign being unavailable) MUST abort
|
||||
// the update before any binary is replaced. Silently proceeding would
|
||||
// allow an unverified binary to overwrite the running one.
|
||||
if err := u.verifySigstoreBundle(ctx, checksumsPath, bundlePath); err != nil {
|
||||
// Log warning but continue - signature verification is optional if cosign isn't installed
|
||||
log.Warn().Err(err).Msg("Signature verification failed or skipped")
|
||||
} else {
|
||||
log.Info().Msg("Sigstore signature verification passed")
|
||||
u.setError(err)
|
||||
return fmt.Errorf("release signature verification failed; aborting update for safety, please update manually with: %s: %w",
|
||||
GetManualUpdateCommand("v"+info.LatestVersion), err)
|
||||
}
|
||||
log.Info().Msg("Sigstore signature verification passed")
|
||||
}
|
||||
|
||||
// Step 3: Download the archive
|
||||
|
||||
@@ -182,9 +182,9 @@ func TestGetManualUpdateCommand(t *testing.T) {
|
||||
wantNotContains []string
|
||||
}{
|
||||
{
|
||||
name: "empty_version_returns_latest",
|
||||
version: "",
|
||||
wantContains: []string{"curl -sSL", InstallScriptURL, "| bash"},
|
||||
name: "empty_version_returns_latest",
|
||||
version: "",
|
||||
wantContains: []string{"curl -sSL", InstallScriptURL, "| bash"},
|
||||
wantNotContains: []string{"bash -s --"},
|
||||
},
|
||||
{
|
||||
@@ -312,9 +312,9 @@ func makeTarGzFile(t *testing.T, dir, content string) (path string, checksum str
|
||||
|
||||
body := []byte(content)
|
||||
hdr := &tar.Header{
|
||||
Name: "worker",
|
||||
Mode: 0755,
|
||||
Size: int64(len(body)),
|
||||
Name: "worker",
|
||||
Mode: 0755,
|
||||
Size: int64(len(body)),
|
||||
Typeflag: tar.TypeReg,
|
||||
}
|
||||
require.NoError(t, tw.WriteHeader(hdr))
|
||||
@@ -455,9 +455,9 @@ func makeTarGzArchive(t *testing.T, files map[string]string) []byte {
|
||||
for name, content := range files {
|
||||
body := []byte(content)
|
||||
hdr := &tar.Header{
|
||||
Name: name,
|
||||
Mode: 0755,
|
||||
Size: int64(len(body)),
|
||||
Name: name,
|
||||
Mode: 0755,
|
||||
Size: int64(len(body)),
|
||||
Typeflag: tar.TypeReg,
|
||||
}
|
||||
require.NoError(t, tw.WriteHeader(hdr))
|
||||
|
||||
+124
-107
@@ -5,6 +5,7 @@ import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"fmt"
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
@@ -139,9 +140,22 @@ func NewClient(cfg Config, embedSvc *embedding.Service) (*Client, error) {
|
||||
return c, nil
|
||||
}
|
||||
|
||||
// AddDocuments adds documents with their embeddings to the vector store.
|
||||
// AddDocuments embeds the documents and inserts them. A document whose ID already exists is an error (the vec0
|
||||
// table does not replace on insert); use ReplaceDocuments to write the documents of a note that may have some.
|
||||
func (c *Client) AddDocuments(ctx context.Context, docs []Document) error {
|
||||
if len(docs) == 0 {
|
||||
return c.writeDocuments(ctx, docs, "", 0)
|
||||
}
|
||||
|
||||
// ReplaceDocuments writes the documents of one note (observation, summary, ...) in place of the ones it had: the old
|
||||
// documents are deleted and the new ones inserted in one transaction, so an edit that removes a fact leaves no stale
|
||||
// document behind and a failure leaves the old ones untouched. docs may be empty, which only deletes.
|
||||
func (c *Client) ReplaceDocuments(ctx context.Context, docType string, sqliteID int64, docs []Document) error {
|
||||
return c.writeDocuments(ctx, docs, docType, sqliteID)
|
||||
}
|
||||
|
||||
// writeDocuments inserts docs; when docType is not empty it first deletes every document of that note.
|
||||
func (c *Client) writeDocuments(ctx context.Context, docs []Document, docType string, sqliteID int64) (err error) {
|
||||
if len(docs) == 0 && docType == "" {
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -154,9 +168,11 @@ func (c *Client) AddDocuments(ctx context.Context, docs []Document) error {
|
||||
// Compute embeddings OUTSIDE the write lock for better concurrency.
|
||||
// Embedding is ONNX inference (slow, mutex-protected internally) — holding
|
||||
// writeMu during inference blocks all concurrent writes and reads.
|
||||
embeddings, err := c.embedSvc.EmbedBatchWithContext(ctx, texts)
|
||||
if err != nil {
|
||||
return fmt.Errorf("generate embeddings: %w", err)
|
||||
var embeddings [][]float32
|
||||
if len(texts) > 0 {
|
||||
if embeddings, err = c.embedSvc.EmbedBatchWithContext(ctx, texts); err != nil {
|
||||
return fmt.Errorf("generate embeddings: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
// Acquire write lock for DB operations only
|
||||
@@ -165,7 +181,7 @@ func (c *Client) AddDocuments(ctx context.Context, docs []Document) error {
|
||||
|
||||
// Insert into vectors table with model version tracking
|
||||
const insertQuery = `
|
||||
INSERT OR REPLACE INTO vectors (doc_id, embedding, sqlite_id, doc_type, field_type, project, scope, model_version)
|
||||
INSERT INTO vectors (doc_id, embedding, sqlite_id, doc_type, field_type, project, scope, model_version)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?)
|
||||
`
|
||||
|
||||
@@ -176,6 +192,8 @@ func (c *Client) AddDocuments(ctx context.Context, docs []Document) error {
|
||||
if err != nil {
|
||||
return fmt.Errorf("begin transaction: %w", err)
|
||||
}
|
||||
// err is the named result, so every failed return below reaches this and the write lock is released. (A shadowed
|
||||
// err in the loop once left the transaction open, and the database locked, after a failed insert.)
|
||||
defer func() {
|
||||
if err != nil {
|
||||
if rbErr := tx.Rollback(); rbErr != nil {
|
||||
@@ -185,6 +203,12 @@ func (c *Client) AddDocuments(ctx context.Context, docs []Document) error {
|
||||
}
|
||||
}()
|
||||
|
||||
if docType != "" {
|
||||
if _, err = tx.ExecContext(ctx, `DELETE FROM vectors WHERE doc_type = ? AND sqlite_id = ?`, docType, sqliteID); err != nil {
|
||||
return fmt.Errorf("delete old documents: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
stmt, err := tx.PrepareContext(ctx, insertQuery)
|
||||
if err != nil {
|
||||
return fmt.Errorf("prepare statement: %w", err)
|
||||
@@ -193,29 +217,28 @@ func (c *Client) AddDocuments(ctx context.Context, docs []Document) error {
|
||||
|
||||
for i, doc := range docs {
|
||||
// Serialize embedding to blob format
|
||||
embBlob, err := sqlite_vec.SerializeFloat32(embeddings[i])
|
||||
if err != nil {
|
||||
var embBlob []byte
|
||||
if embBlob, err = sqlite_vec.SerializeFloat32(embeddings[i]); err != nil {
|
||||
return fmt.Errorf("serialize embedding for %s: %w", doc.ID, err)
|
||||
}
|
||||
|
||||
// Extract metadata
|
||||
sqliteID, _ := doc.Metadata["sqlite_id"].(int64)
|
||||
docType, _ := doc.Metadata["doc_type"].(string)
|
||||
docSQLiteID, _ := doc.Metadata["sqlite_id"].(int64)
|
||||
docDocType, _ := doc.Metadata["doc_type"].(string)
|
||||
fieldType, _ := doc.Metadata["field_type"].(string)
|
||||
project, _ := doc.Metadata["project"].(string)
|
||||
scope, _ := doc.Metadata["scope"].(string)
|
||||
|
||||
_, err = stmt.ExecContext(ctx,
|
||||
if _, err = stmt.ExecContext(ctx,
|
||||
doc.ID,
|
||||
embBlob,
|
||||
sqliteID,
|
||||
docType,
|
||||
docSQLiteID,
|
||||
docDocType,
|
||||
fieldType,
|
||||
project,
|
||||
scope,
|
||||
modelVersion,
|
||||
)
|
||||
if err != nil {
|
||||
); err != nil {
|
||||
return fmt.Errorf("insert document %s: %w", doc.ID, err)
|
||||
}
|
||||
}
|
||||
@@ -292,67 +315,13 @@ func (c *Client) Query(ctx context.Context, query string, limit int, where map[s
|
||||
}
|
||||
defer c.readMu.RUnlock()
|
||||
|
||||
// Build query with filters
|
||||
// vec0 supports WHERE clauses on metadata columns
|
||||
args := []any{queryBlob}
|
||||
|
||||
sqlQuery := `
|
||||
SELECT
|
||||
doc_id,
|
||||
distance,
|
||||
sqlite_id,
|
||||
doc_type,
|
||||
field_type,
|
||||
project,
|
||||
scope
|
||||
FROM vectors
|
||||
WHERE embedding MATCH ?
|
||||
`
|
||||
|
||||
// Add filters - these work with vec0 metadata columns
|
||||
if docType, ok := where["doc_type"].(string); ok && docType != "" {
|
||||
sqlQuery += " AND doc_type = ?"
|
||||
args = append(args, docType)
|
||||
}
|
||||
if project, ok := where["project"].(string); ok && project != "" {
|
||||
// Include project-specific OR global scope
|
||||
sqlQuery += " AND (project = ? OR scope = 'global')"
|
||||
args = append(args, project)
|
||||
}
|
||||
|
||||
sqlQuery += " ORDER BY distance LIMIT ?"
|
||||
args = append(args, limit)
|
||||
|
||||
rows, err := c.db.QueryContext(ctx, sqlQuery, args...)
|
||||
// vec0 answers equality filters on metadata columns, but not an OR between them, so a project filter
|
||||
// (the project's own documents, plus global ones) is asked as two searches and merged.
|
||||
docType, _ := where["doc_type"].(string)
|
||||
project, _ := where["project"].(string)
|
||||
results, err := c.nearest(ctx, queryBlob, docType, project, limit)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("query vectors: %w", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
var results []QueryResult
|
||||
for rows.Next() {
|
||||
var r QueryResult
|
||||
var sqliteID int64
|
||||
var docType, fieldType, project, scope sql.NullString
|
||||
|
||||
if err := rows.Scan(&r.ID, &r.Distance, &sqliteID, &docType, &fieldType, &project, &scope); err != nil {
|
||||
return nil, fmt.Errorf("scan row: %w", err)
|
||||
}
|
||||
|
||||
r.Similarity = DistanceToSimilarity(r.Distance)
|
||||
r.Metadata = map[string]any{
|
||||
"sqlite_id": float64(sqliteID), // Keep as float64 for compatibility
|
||||
"doc_type": docType.String,
|
||||
"field_type": fieldType.String,
|
||||
"project": project.String,
|
||||
"scope": scope.String,
|
||||
}
|
||||
|
||||
results = append(results, r)
|
||||
}
|
||||
|
||||
if err = rows.Err(); err != nil {
|
||||
return nil, fmt.Errorf("iterate rows: %w", err)
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Cache the results
|
||||
@@ -507,60 +476,108 @@ func (c *Client) QueryMultiField(ctx context.Context, query string, limit int, d
|
||||
}
|
||||
defer c.readMu.RUnlock()
|
||||
|
||||
// Query with field type aggregation - get best match per document
|
||||
sqlQuery := `
|
||||
WITH ranked_results AS (
|
||||
SELECT
|
||||
doc_id,
|
||||
distance,
|
||||
sqlite_id,
|
||||
doc_type,
|
||||
field_type,
|
||||
project,
|
||||
scope,
|
||||
ROW_NUMBER() OVER (PARTITION BY sqlite_id ORDER BY distance ASC) as rn
|
||||
FROM vectors
|
||||
WHERE embedding MATCH ?
|
||||
AND doc_type = ?
|
||||
AND (project = ? OR scope = 'global')
|
||||
)
|
||||
SELECT doc_id, distance, sqlite_id, doc_type, field_type, project, scope
|
||||
FROM ranked_results
|
||||
WHERE rn = 1
|
||||
ORDER BY distance
|
||||
LIMIT ?
|
||||
`
|
||||
// Each observation has several vector documents (one per field). Ask for enough of them that `limit`
|
||||
// distinct observations are likely to be among them, then keep the best match per observation.
|
||||
found, err := c.nearest(ctx, queryBlob, docType, project, limit*multiFieldFetchFactor)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
results := make([]QueryResult, 0, limit)
|
||||
seen := make(map[float64]bool, limit)
|
||||
for _, r := range found { // best first
|
||||
id, _ := r.Metadata["sqlite_id"].(float64)
|
||||
if seen[id] {
|
||||
continue
|
||||
}
|
||||
seen[id] = true
|
||||
results = append(results, r)
|
||||
if len(results) == limit {
|
||||
break
|
||||
}
|
||||
}
|
||||
return results, nil
|
||||
}
|
||||
|
||||
rows, err := c.db.QueryContext(ctx, sqlQuery, queryBlob, docType, project, limit)
|
||||
// multiFieldFetchFactor is how many vector documents QueryMultiField asks for per wanted observation.
|
||||
const multiFieldFetchFactor = 6
|
||||
|
||||
// nearest runs the nearest-neighbour search for a document type and, when project is not empty, for that
|
||||
// project's documents and the global ones. The caller holds the read lock. Results are best first, at most limit.
|
||||
func (c *Client) nearest(ctx context.Context, queryBlob []byte, docType, project string, limit int) ([]QueryResult, error) {
|
||||
base := " WHERE embedding MATCH ?"
|
||||
baseArgs := []any{queryBlob}
|
||||
if docType != "" {
|
||||
base += " AND doc_type = ?"
|
||||
baseArgs = append(baseArgs, docType)
|
||||
}
|
||||
if project == "" {
|
||||
return c.knn(ctx, base, baseArgs, limit)
|
||||
}
|
||||
own, err := c.knn(ctx, base+" AND project = ?", append(append([]any{}, baseArgs...), project), limit)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
global, err := c.knn(ctx, base+" AND scope = 'global'", baseArgs, limit)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return mergeNearest(limit, own, global), nil
|
||||
}
|
||||
|
||||
// knn runs one vec0 nearest-neighbour query. where holds the MATCH and filters vec0 can answer on their own.
|
||||
func (c *Client) knn(ctx context.Context, where string, args []any, limit int) ([]QueryResult, error) {
|
||||
rows, err := c.db.QueryContext(ctx,
|
||||
"SELECT doc_id, distance, sqlite_id, doc_type, field_type, project, scope FROM vectors"+where+" ORDER BY distance LIMIT ?",
|
||||
append(append([]any{}, args...), limit)...)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("query vectors: %w", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
// Pre-allocate with limit to avoid repeated slice growth
|
||||
results := make([]QueryResult, 0, limit)
|
||||
var results []QueryResult
|
||||
for rows.Next() {
|
||||
var r QueryResult
|
||||
var sqliteID int64
|
||||
var docTypeVal, fieldType, projectVal, scope sql.NullString
|
||||
var docType, fieldType, project, scope sql.NullString
|
||||
|
||||
if err := rows.Scan(&r.ID, &r.Distance, &sqliteID, &docTypeVal, &fieldType, &projectVal, &scope); err != nil {
|
||||
if err := rows.Scan(&r.ID, &r.Distance, &sqliteID, &docType, &fieldType, &project, &scope); err != nil {
|
||||
return nil, fmt.Errorf("scan row: %w", err)
|
||||
}
|
||||
|
||||
r.Similarity = DistanceToSimilarity(r.Distance)
|
||||
r.Metadata = map[string]any{
|
||||
"sqlite_id": float64(sqliteID),
|
||||
"doc_type": docTypeVal.String,
|
||||
"sqlite_id": float64(sqliteID), // Keep as float64 for compatibility
|
||||
"doc_type": docType.String,
|
||||
"field_type": fieldType.String,
|
||||
"project": projectVal.String,
|
||||
"project": project.String,
|
||||
"scope": scope.String,
|
||||
}
|
||||
|
||||
results = append(results, r)
|
||||
}
|
||||
if err = rows.Err(); err != nil {
|
||||
return nil, fmt.Errorf("iterate rows: %w", err)
|
||||
}
|
||||
return results, nil
|
||||
}
|
||||
|
||||
return results, rows.Err()
|
||||
// mergeNearest combines result lists into one, best first, without repeating a document, at most limit long.
|
||||
func mergeNearest(limit int, lists ...[]QueryResult) []QueryResult {
|
||||
var all []QueryResult
|
||||
seen := map[string]bool{}
|
||||
for _, list := range lists {
|
||||
for _, r := range list {
|
||||
if !seen[r.ID] {
|
||||
seen[r.ID] = true
|
||||
all = append(all, r)
|
||||
}
|
||||
}
|
||||
}
|
||||
sort.SliceStable(all, func(i, j int) bool { return all[i].Distance < all[j].Distance })
|
||||
if len(all) > limit {
|
||||
all = all[:limit]
|
||||
}
|
||||
return all
|
||||
}
|
||||
|
||||
// acquireRLockWithContext acquires a read lock on mu, respecting ctx cancellation.
|
||||
|
||||
@@ -405,6 +405,50 @@ func TestClient_Query_WithProjectFilter(t *testing.T) {
|
||||
results, err := client.Query(context.Background(), "authentication security", 10, nil)
|
||||
require.NoError(t, err)
|
||||
assert.NotEmpty(t, results, "Should find some results")
|
||||
assert.Len(t, results, 3)
|
||||
|
||||
ids := func(rs []QueryResult) []string {
|
||||
var out []string
|
||||
for _, r := range rs {
|
||||
out = append(out, r.ID)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// With a project filter: the project's own documents and the global ones, not another project's.
|
||||
results, err = client.Query(context.Background(), "authentication security", 10, BuildWhereFilter(DocTypeObservation, "project-a"))
|
||||
require.NoError(t, err, "a project filter must be answerable (vec0 rejects an OR between metadata filters)")
|
||||
assert.ElementsMatch(t, []string{"obs-1", "obs-3"}, ids(results), "project-a's note and the global one, not project-b's own")
|
||||
|
||||
results, err = client.Query(context.Background(), "authentication security", 10, BuildWhereFilter(DocTypeObservation, "project-b"))
|
||||
require.NoError(t, err)
|
||||
assert.ElementsMatch(t, []string{"obs-2", "obs-3"}, ids(results))
|
||||
|
||||
// Best first, and the limit applies to the merged list.
|
||||
results, err = client.Query(context.Background(), "authentication security", 1, BuildWhereFilter(DocTypeObservation, "project-a"))
|
||||
require.NoError(t, err)
|
||||
require.Len(t, results, 1)
|
||||
all, err := client.Query(context.Background(), "authentication security", 10, BuildWhereFilter(DocTypeObservation, "project-a"))
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, all[0].ID, results[0].ID, "the one result is the best of the merged list")
|
||||
assert.LessOrEqual(t, all[0].Distance, all[1].Distance)
|
||||
|
||||
// A project nobody wrote to still sees the global documents.
|
||||
results, err = client.Query(context.Background(), "authentication security", 10, BuildWhereFilter(DocTypeObservation, "project-z"))
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, []string{"obs-3"}, ids(results))
|
||||
}
|
||||
|
||||
func TestMergeNearest(t *testing.T) {
|
||||
r := func(id string, d float64) QueryResult { return QueryResult{ID: id, Distance: d} }
|
||||
got := mergeNearest(3, []QueryResult{r("a", 0.5), r("b", 0.9)}, []QueryResult{r("b", 0.9), r("c", 0.1), r("d", 0.7)})
|
||||
var ids []string
|
||||
for _, x := range got {
|
||||
ids = append(ids, x.ID)
|
||||
}
|
||||
assert.Equal(t, []string{"c", "a", "d"}, ids, "best first, a document found twice counts once, cut to the limit")
|
||||
assert.Empty(t, mergeNearest(5))
|
||||
assert.Len(t, mergeNearest(5, []QueryResult{r("a", 1)}, nil), 1)
|
||||
}
|
||||
|
||||
func TestClient_IsConnected(t *testing.T) {
|
||||
@@ -1105,8 +1149,6 @@ func TestClient_QueryBatch_WithContextCancellation(t *testing.T) {
|
||||
// =============================================================================
|
||||
|
||||
func TestClient_QueryMultiField_Basic(t *testing.T) {
|
||||
t.Skip("QueryMultiField SQL query needs 'k' parameter fix for sqlite-vec")
|
||||
|
||||
db, dbCleanup := testDB(t)
|
||||
defer dbCleanup()
|
||||
|
||||
@@ -1173,8 +1215,6 @@ func TestClient_QueryMultiField_Basic(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestClient_QueryMultiField_WithGlobalScope(t *testing.T) {
|
||||
t.Skip("QueryMultiField SQL query needs 'k' parameter fix for sqlite-vec")
|
||||
|
||||
db, dbCleanup := testDB(t)
|
||||
defer dbCleanup()
|
||||
|
||||
@@ -1217,7 +1257,37 @@ func TestClient_QueryMultiField_WithGlobalScope(t *testing.T) {
|
||||
require.NoError(t, err)
|
||||
|
||||
// Should include both project-scoped (matching project) and global
|
||||
assert.NotEmpty(t, results)
|
||||
require.Len(t, results, 2)
|
||||
assert.ElementsMatch(t, []float64{1, 2}, []float64{
|
||||
results[0].Metadata["sqlite_id"].(float64), results[1].Metadata["sqlite_id"].(float64),
|
||||
})
|
||||
}
|
||||
|
||||
// The query the de-duplication of new observations makes: the same project, a few results, and a note that
|
||||
// says nearly the same thing must come back with a high similarity.
|
||||
func TestClient_Query_ProjectFilterFindsANearDuplicate(t *testing.T) {
|
||||
db, dbCleanup := testDB(t)
|
||||
defer dbCleanup()
|
||||
embedSvc, embedCleanup := testEmbeddingService(t)
|
||||
defer embedCleanup()
|
||||
client, err := NewClient(Config{DB: db}, embedSvc)
|
||||
require.NoError(t, err)
|
||||
|
||||
text := "The shipping rate cache lives for sixty minutes and is stored in Redis."
|
||||
require.NoError(t, client.AddDocuments(context.Background(), []Document{
|
||||
{ID: "obs-1", Content: text, Metadata: map[string]any{"sqlite_id": int64(1), "doc_type": "observation", "project": "shop", "scope": "project"}},
|
||||
{ID: "obs-2", Content: text, Metadata: map[string]any{"sqlite_id": int64(2), "doc_type": "observation", "project": "elsewhere", "scope": "project"}},
|
||||
{ID: "obs-3", Content: "Parcel labels are printed as PDF.", Metadata: map[string]any{"sqlite_id": int64(3), "doc_type": "observation", "project": "shop", "scope": "project"}},
|
||||
}))
|
||||
|
||||
results, err := client.Query(context.Background(), text, 3, BuildWhereFilter(DocTypeObservation, "shop"))
|
||||
require.NoError(t, err)
|
||||
require.NotEmpty(t, results)
|
||||
assert.Equal(t, "obs-1", results[0].ID, "the same text in the same project is the best match")
|
||||
assert.GreaterOrEqual(t, results[0].Similarity, 0.99)
|
||||
for _, r := range results {
|
||||
assert.NotEqual(t, "obs-2", r.ID, "another project's note is not a duplicate")
|
||||
}
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
@@ -2152,3 +2222,67 @@ func TestGetOrComputeEmbedding_ContextCancelDuringSingleflight(t *testing.T) {
|
||||
// Wait for first call to finish
|
||||
<-firstDone
|
||||
}
|
||||
|
||||
func TestClient_AddDocuments_FailedBatchIsRolledBackAndDoesNotLockTheDatabase(t *testing.T) {
|
||||
db, dbCleanup := testDB(t)
|
||||
defer dbCleanup()
|
||||
embedSvc, embedCleanup := testEmbeddingService(t)
|
||||
defer embedCleanup()
|
||||
client, err := NewClient(Config{DB: db}, embedSvc)
|
||||
require.NoError(t, err)
|
||||
|
||||
doc := func(id string, sqliteID int64) Document {
|
||||
return Document{ID: id, Content: "text of " + id, Metadata: map[string]any{
|
||||
"sqlite_id": sqliteID, "doc_type": "observation", "field_type": "narrative", "project": "alpha", "scope": "project",
|
||||
}}
|
||||
}
|
||||
require.NoError(t, client.AddDocuments(context.Background(), []Document{doc("obs_1_narrative", 1)}))
|
||||
|
||||
// The second document of the batch exists already, so the batch fails after the first was inserted.
|
||||
err = client.AddDocuments(context.Background(), []Document{doc("obs_2_narrative", 2), doc("obs_1_narrative", 1)})
|
||||
require.Error(t, err)
|
||||
|
||||
var n int
|
||||
require.NoError(t, db.QueryRow("SELECT COUNT(*) FROM vectors WHERE doc_id = 'obs_2_narrative'").Scan(&n))
|
||||
assert.Zero(t, n, "nothing of the failed batch is kept")
|
||||
|
||||
// With the transaction left open, this would wait for the lock until the context ends.
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
require.NoError(t, client.AddDocuments(ctx, []Document{doc("obs_3_narrative", 3)}), "a write after a failed one goes through")
|
||||
}
|
||||
|
||||
func TestClient_ReplaceDocuments_ReplacesOnlyTheDocumentsOfThatNote(t *testing.T) {
|
||||
db, dbCleanup := testDB(t)
|
||||
defer dbCleanup()
|
||||
embedSvc, embedCleanup := testEmbeddingService(t)
|
||||
defer embedCleanup()
|
||||
client, err := NewClient(Config{DB: db}, embedSvc)
|
||||
require.NoError(t, err)
|
||||
ctx := context.Background()
|
||||
|
||||
doc := func(id string, sqliteID int64, docType, scope string) Document {
|
||||
return Document{ID: id, Content: "text of " + id, Metadata: map[string]any{
|
||||
"sqlite_id": sqliteID, "doc_type": docType, "field_type": "narrative", "project": "alpha", "scope": scope,
|
||||
}}
|
||||
}
|
||||
require.NoError(t, client.AddDocuments(ctx, []Document{
|
||||
doc("obs_1_narrative", 1, "observation", "global"), doc("obs_1_fact_0", 1, "observation", "global"),
|
||||
doc("obs_2_narrative", 2, "observation", "global"),
|
||||
doc("summary_1_request", 1, "session_summary", ""), // the same sqlite id, another kind of document
|
||||
}))
|
||||
|
||||
require.NoError(t, client.ReplaceDocuments(ctx, "observation", 1, []Document{doc("obs_1_narrative", 1, "observation", "project")}))
|
||||
|
||||
count := func(where string, args ...any) (n int) {
|
||||
require.NoError(t, db.QueryRow("SELECT COUNT(*) FROM vectors WHERE "+where, args...).Scan(&n))
|
||||
return n
|
||||
}
|
||||
assert.Equal(t, 1, count("doc_type = 'observation' AND sqlite_id = 1"), "the fact document is gone")
|
||||
assert.Equal(t, 1, count("doc_id = 'obs_1_narrative' AND scope = 'project'"), "and the narrative carries the new scope")
|
||||
assert.Equal(t, 1, count("doc_id = 'obs_2_narrative' AND scope = 'global'"), "another note is untouched")
|
||||
assert.Equal(t, 1, count("doc_id = 'summary_1_request'"), "so is a summary with the same id")
|
||||
|
||||
require.NoError(t, client.ReplaceDocuments(ctx, "observation", 1, nil), "no documents left: only the old ones are deleted")
|
||||
assert.Zero(t, count("doc_type = 'observation' AND sqlite_id = 1"))
|
||||
}
|
||||
|
||||
@@ -22,11 +22,8 @@ func NewSync(client *Client) *Sync {
|
||||
// SyncObservation syncs a single observation to the vector store.
|
||||
func (s *Sync) SyncObservation(ctx context.Context, obs *models.Observation) error {
|
||||
docs := s.formatObservationDocs(obs)
|
||||
if len(docs) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
if err := s.client.AddDocuments(ctx, docs); err != nil {
|
||||
// Replace, not add: a note that is synced again (after an edit, a re-scope or a re-import) already has documents.
|
||||
if err := s.client.ReplaceDocuments(ctx, "observation", obs.ID, docs); err != nil {
|
||||
return fmt.Errorf("add observation docs: %w", err)
|
||||
}
|
||||
|
||||
@@ -102,11 +99,7 @@ func (s *Sync) formatObservationDocs(obs *models.Observation) []Document {
|
||||
// SyncSummary syncs a single session summary to the vector store.
|
||||
func (s *Sync) SyncSummary(ctx context.Context, summary *models.SessionSummary) error {
|
||||
docs := s.formatSummaryDocs(summary)
|
||||
if len(docs) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
if err := s.client.AddDocuments(ctx, docs); err != nil {
|
||||
if err := s.client.ReplaceDocuments(ctx, "session_summary", summary.ID, docs); err != nil {
|
||||
return fmt.Errorf("add summary docs: %w", err)
|
||||
}
|
||||
|
||||
|
||||
@@ -346,3 +346,57 @@ func TestSync_FormatObservationDocs_EmptyScope(t *testing.T) {
|
||||
assert.Len(t, docs, 1)
|
||||
assert.Equal(t, "project", docs[0].Metadata["scope"])
|
||||
}
|
||||
|
||||
func vectorCount(t *testing.T, c *Client, where string, args ...any) int {
|
||||
t.Helper()
|
||||
var n int
|
||||
require.NoError(t, c.db.QueryRow("SELECT COUNT(*) FROM vectors WHERE "+where, args...).Scan(&n))
|
||||
return n
|
||||
}
|
||||
|
||||
func TestSync_SyncObservation_AgainReplacesTheDocumentsOfTheNote(t *testing.T) {
|
||||
client, cleanup := testClient(t)
|
||||
defer cleanup()
|
||||
sync := NewSync(client)
|
||||
ctx := context.Background()
|
||||
|
||||
obs := &models.Observation{
|
||||
ID: 7, SDKSessionID: "s", Project: "alpha", Type: models.ObsTypeDiscovery, Scope: models.ScopeGlobal,
|
||||
Narrative: sql.NullString{String: "The queue keeps one consumer per tenant.", Valid: true},
|
||||
Facts: []string{"one consumer per tenant", "retries are capped", "dead letters go to a second queue"},
|
||||
}
|
||||
other := &models.Observation{
|
||||
ID: 8, SDKSessionID: "s", Project: "alpha", Type: models.ObsTypeDiscovery, Scope: models.ScopeProject,
|
||||
Narrative: sql.NullString{String: "Releases are cut on Thursdays.", Valid: true},
|
||||
}
|
||||
require.NoError(t, sync.SyncObservation(ctx, obs))
|
||||
require.NoError(t, sync.SyncObservation(ctx, other))
|
||||
require.Equal(t, 4, vectorCount(t, client, "doc_type = 'observation' AND sqlite_id = ?", 7))
|
||||
|
||||
// The note was edited: one fact fewer, and it is a project note now.
|
||||
obs.Facts = obs.Facts[:2]
|
||||
obs.Scope = models.ScopeProject
|
||||
require.NoError(t, sync.SyncObservation(ctx, obs), "syncing a note that already has documents is not an error")
|
||||
|
||||
assert.Equal(t, 3, vectorCount(t, client, "doc_type = 'observation' AND sqlite_id = ?", 7), "narrative and two facts; the dropped fact's document is gone")
|
||||
assert.Equal(t, 0, vectorCount(t, client, "doc_id = ?", "obs_7_fact_2"))
|
||||
assert.Equal(t, 3, vectorCount(t, client, "sqlite_id = ? AND scope = 'project'", 7), "the new scope is what the vectors carry")
|
||||
assert.Equal(t, 1, vectorCount(t, client, "sqlite_id = ?", 8), "other notes are not touched")
|
||||
}
|
||||
|
||||
func TestSync_SyncSummary_AgainReplacesTheDocumentsOfTheSummary(t *testing.T) {
|
||||
client, cleanup := testClient(t)
|
||||
defer cleanup()
|
||||
sync := NewSync(client)
|
||||
ctx := context.Background()
|
||||
|
||||
summary := &models.SessionSummary{
|
||||
ID: 3, SDKSessionID: "s", Project: "alpha",
|
||||
Request: sql.NullString{String: "Add the queue", Valid: true}, Learned: sql.NullString{String: "One consumer per tenant", Valid: true},
|
||||
}
|
||||
require.NoError(t, sync.SyncSummary(ctx, summary))
|
||||
summary.Learned = sql.NullString{}
|
||||
require.NoError(t, sync.SyncSummary(ctx, summary))
|
||||
|
||||
assert.Equal(t, 1, vectorCount(t, client, "doc_type = 'session_summary' AND sqlite_id = ?", 3), "the field that was emptied has no document any more")
|
||||
}
|
||||
|
||||
@@ -258,7 +258,6 @@ func (s *Service) handleObservation(w http.ResponseWriter, r *http.Request) {
|
||||
return
|
||||
}
|
||||
|
||||
|
||||
// Find session
|
||||
sess, err := s.sessionStore.FindAnySDKSession(r.Context(), req.ClaudeSessionID)
|
||||
if err != nil {
|
||||
|
||||
@@ -218,9 +218,8 @@ func (s *Service) handleUpdateRestart(w http.ResponseWriter, r *http.Request) {
|
||||
|
||||
// Restart in background after response is sent
|
||||
go func() {
|
||||
if err := s.updater.Restart(); err != nil {
|
||||
log.Error().Err(err).Msg("Failed to restart worker")
|
||||
}
|
||||
// Restart only returns when the new worker failed to start.
|
||||
log.Error().Err(s.updater.Restart()).Msg("Failed to restart worker")
|
||||
}()
|
||||
}
|
||||
|
||||
@@ -244,8 +243,7 @@ func (s *Service) handleRestart(w http.ResponseWriter, r *http.Request) {
|
||||
go func() {
|
||||
// Small delay to ensure response is sent
|
||||
time.Sleep(100 * time.Millisecond)
|
||||
if err := s.updater.Restart(); err != nil {
|
||||
log.Error().Err(err).Msg("Failed to restart worker")
|
||||
}
|
||||
// Restart only returns when the new worker failed to start.
|
||||
log.Error().Err(s.updater.Restart()).Msg("Failed to restart worker")
|
||||
}()
|
||||
}
|
||||
|
||||
@@ -0,0 +1,59 @@
|
||||
package worker
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func doUpdateRequest(t *testing.T, svc *Service, method, target string, body any) *httptest.ResponseRecorder {
|
||||
t.Helper()
|
||||
var buf bytes.Buffer
|
||||
if body != nil {
|
||||
require.NoError(t, json.NewEncoder(&buf).Encode(body))
|
||||
}
|
||||
rec := httptest.NewRecorder()
|
||||
svc.router.ServeHTTP(rec, httptest.NewRequest(method, target, &buf))
|
||||
return rec
|
||||
}
|
||||
|
||||
func TestHandleUpdateObservation_EditsAndReturnsTheObservation(t *testing.T) {
|
||||
svc, cleanup := testService(t)
|
||||
defer cleanup()
|
||||
id := createTestObservation(t, svc.observationStore, "proj_aaaaaa", "Old title", "Old narrative", []string{"cache"})
|
||||
|
||||
rec := doUpdateRequest(t, svc, http.MethodPut, fmt.Sprintf("/api/observations/%d", id),
|
||||
map[string]any{"title": "New title", "facts": []string{"one", "two"}})
|
||||
require.Equal(t, http.StatusOK, rec.Code, rec.Body.String())
|
||||
var resp struct {
|
||||
Observation struct {
|
||||
Title string `json:"title"`
|
||||
Narrative string `json:"narrative"`
|
||||
Facts []string `json:"facts"`
|
||||
} `json:"observation"`
|
||||
}
|
||||
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp))
|
||||
assert.Equal(t, "New title", resp.Observation.Title)
|
||||
assert.Equal(t, "Old narrative", resp.Observation.Narrative)
|
||||
assert.Equal(t, []string{"one", "two"}, resp.Observation.Facts)
|
||||
|
||||
rec = doUpdateRequest(t, svc, http.MethodGet, fmt.Sprintf("/api/observations/%d", id), nil)
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
assert.Contains(t, rec.Body.String(), "New title", "the edit is stored")
|
||||
}
|
||||
|
||||
func TestHandleUpdateObservation_Refusals(t *testing.T) {
|
||||
svc, cleanup := testService(t)
|
||||
defer cleanup()
|
||||
id := createTestObservation(t, svc.observationStore, "proj_aaaaaa", "A title", "A narrative", nil)
|
||||
|
||||
assert.Equal(t, http.StatusNotFound, doUpdateRequest(t, svc, http.MethodPut, "/api/observations/999999", map[string]any{"title": "x"}).Code)
|
||||
assert.Equal(t, http.StatusBadRequest, doUpdateRequest(t, svc, http.MethodPut, fmt.Sprintf("/api/observations/%d", id), map[string]any{"scope": "galaxy"}).Code)
|
||||
assert.Equal(t, http.StatusBadRequest, doUpdateRequest(t, svc, http.MethodPut, "/api/observations/abc", map[string]any{"title": "x"}).Code)
|
||||
}
|
||||
@@ -4,6 +4,7 @@ package worker
|
||||
import (
|
||||
"context"
|
||||
"crypto/rand"
|
||||
"crypto/subtle"
|
||||
"encoding/hex"
|
||||
"fmt"
|
||||
"net/http"
|
||||
@@ -121,6 +122,22 @@ func NewTokenAuth(enabled bool) (*TokenAuth, error) {
|
||||
return ta, nil
|
||||
}
|
||||
|
||||
// NewTokenAuthWithToken creates a TokenAuth using an explicitly supplied token.
|
||||
// It is used to enforce a caller-provided token (e.g. from an environment
|
||||
// variable) so the value is known to clients, unlike the random token produced
|
||||
// by NewTokenAuth. Authentication is enabled only when token is non-empty.
|
||||
func NewTokenAuthWithToken(token string) *TokenAuth {
|
||||
return &TokenAuth{
|
||||
enabled: token != "",
|
||||
token: token,
|
||||
ExemptPaths: map[string]bool{
|
||||
"/health": true,
|
||||
"/api/health": true,
|
||||
"/api/ready": true,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// Token returns the authentication token.
|
||||
// Returns empty string if authentication is disabled.
|
||||
func (ta *TokenAuth) Token() string {
|
||||
@@ -161,7 +178,8 @@ func (ta *TokenAuth) Middleware(next http.Handler) http.Handler {
|
||||
}
|
||||
}
|
||||
|
||||
if providedToken != token {
|
||||
// Constant-time comparison to avoid leaking the token via timing.
|
||||
if subtle.ConstantTimeCompare([]byte(providedToken), []byte(token)) != 1 {
|
||||
http.Error(w, "unauthorized", http.StatusUnauthorized)
|
||||
return
|
||||
}
|
||||
|
||||
@@ -2,6 +2,7 @@
|
||||
package worker
|
||||
|
||||
import (
|
||||
"net"
|
||||
"net/http"
|
||||
"sync"
|
||||
"time"
|
||||
@@ -206,14 +207,13 @@ func (pcrl *PerClientRateLimiter) Stats() map[string]any {
|
||||
}
|
||||
|
||||
// PerClientRateLimitMiddleware creates middleware that applies per-client rate limiting.
|
||||
// Uses X-Forwarded-For or RemoteAddr to identify clients.
|
||||
// Identifies clients by the peer IP of the connection; forwarding headers are spoofable and ignored.
|
||||
func PerClientRateLimitMiddleware(limiter *PerClientRateLimiter) func(http.Handler) http.Handler {
|
||||
return func(next http.Handler) http.Handler {
|
||||
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
// Get client identifier (prefer X-Real-IP from RealIP middleware)
|
||||
clientKey := r.RemoteAddr
|
||||
if xff := r.Header.Get("X-Real-IP"); xff != "" {
|
||||
clientKey = xff
|
||||
if host, _, err := net.SplitHostPort(clientKey); err == nil {
|
||||
clientKey = host
|
||||
}
|
||||
|
||||
if !limiter.Allow(clientKey) {
|
||||
|
||||
@@ -0,0 +1,35 @@
|
||||
package worker
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestPerClientRateLimitMiddleware_KeysOnPeerIP(t *testing.T) {
|
||||
h := PerClientRateLimitMiddleware(NewPerClientRateLimiter(0.001, 1))(
|
||||
http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { w.WriteHeader(http.StatusOK) }))
|
||||
|
||||
do := func(remote, realIP string) int {
|
||||
req := httptest.NewRequest(http.MethodGet, "/", nil)
|
||||
req.RemoteAddr = remote
|
||||
if realIP != "" {
|
||||
req.Header.Set("X-Real-IP", realIP)
|
||||
}
|
||||
rec := httptest.NewRecorder()
|
||||
h.ServeHTTP(rec, req)
|
||||
return rec.Code
|
||||
}
|
||||
|
||||
if got := do("10.0.0.1:1111", "1.1.1.1"); got != http.StatusOK {
|
||||
t.Fatalf("first request: got %d", got)
|
||||
}
|
||||
// Same peer, new port and forged header: must share the exhausted bucket.
|
||||
if got := do("10.0.0.1:2222", "2.2.2.2"); got != http.StatusTooManyRequests {
|
||||
t.Fatalf("spoofed header/new port bypassed limit: got %d", got)
|
||||
}
|
||||
// Different peer is unaffected.
|
||||
if got := do("10.0.0.2:1111", ""); got != http.StatusOK {
|
||||
t.Fatalf("other peer: got %d", got)
|
||||
}
|
||||
}
|
||||
@@ -51,7 +51,7 @@ func (p *Processor) checkVectorDeduplication(ctx context.Context, obs *models.Pa
|
||||
for _, r := range results {
|
||||
if r.Similarity >= cfg.DeduplicationThreshold {
|
||||
obsID := extractObservationIDFromVectorDoc(r)
|
||||
if obsID > 0 {
|
||||
if obsID > 0 && p.canMergeInto(ctx, obsID) {
|
||||
return &DeduplicationResult{
|
||||
ExistingID: obsID,
|
||||
Similarity: r.Similarity,
|
||||
@@ -64,6 +64,20 @@ func (p *Processor) checkVectorDeduplication(ctx context.Context, obs *models.Pa
|
||||
return &DeduplicationResult{Action: "insert"}
|
||||
}
|
||||
|
||||
// canMergeInto says whether a new observation may be merged into the existing one. A note a person superseded
|
||||
// is hidden from sessions, so what was merged into it would be hidden as well: such a note is not a target,
|
||||
// and neither is one that is gone.
|
||||
func (p *Processor) canMergeInto(ctx context.Context, id int64) bool {
|
||||
if p.observationStore == nil {
|
||||
return true
|
||||
}
|
||||
existing, err := p.observationStore.GetObservationByID(ctx, id)
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
return existing != nil && !existing.IsSuperseded
|
||||
}
|
||||
|
||||
// buildObservationSearchText creates searchable text from a parsed observation.
|
||||
func buildObservationSearchText(obs *models.ParsedObservation) string {
|
||||
var parts []string
|
||||
|
||||
@@ -1,8 +1,10 @@
|
||||
package sdk
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/lukaszraczylo/claude-mnemonic/internal/db/gorm"
|
||||
"github.com/lukaszraczylo/claude-mnemonic/internal/vector/sqlitevec"
|
||||
"github.com/lukaszraczylo/claude-mnemonic/pkg/models"
|
||||
)
|
||||
@@ -121,7 +123,7 @@ func TestCheckVectorDeduplication_NilClient(t *testing.T) {
|
||||
Narrative: "Some narrative text",
|
||||
}
|
||||
|
||||
result := p.checkVectorDeduplication(nil, obs, "test-project")
|
||||
result := p.checkVectorDeduplication(context.Background(), obs, "test-project")
|
||||
if result.Action != "insert" {
|
||||
t.Errorf("expected Action='insert' when vectorClient is nil, got %q", result.Action)
|
||||
}
|
||||
@@ -136,8 +138,40 @@ func TestCheckVectorDeduplication_EmptySearchText(t *testing.T) {
|
||||
// All empty fields
|
||||
}
|
||||
|
||||
result := p.checkVectorDeduplication(nil, obs, "test-project")
|
||||
result := p.checkVectorDeduplication(context.Background(), obs, "test-project")
|
||||
if result.Action != "insert" {
|
||||
t.Errorf("expected Action='insert' for empty observation, got %q", result.Action)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCanMergeInto(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
dir := t.TempDir()
|
||||
store, err := gorm.NewStore(gorm.Config{Path: dir + "/test.db", MaxConns: 1})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer store.Close()
|
||||
obsStore := gorm.NewObservationStore(store, nil, nil, nil)
|
||||
id, _, err := obsStore.StoreObservation(ctx, "s1", "proj", &models.ParsedObservation{Type: models.ObsTypeDiscovery, Title: "A note"}, 1, 1)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
p := &Processor{observationStore: obsStore}
|
||||
|
||||
if !p.canMergeInto(ctx, id) {
|
||||
t.Error("a live observation is a merge target")
|
||||
}
|
||||
if p.canMergeInto(ctx, id+1000) {
|
||||
t.Error("an observation that does not exist is not")
|
||||
}
|
||||
if err := store.DB.Exec(`UPDATE observations SET is_superseded = 1 WHERE id = ?`, id).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if p.canMergeInto(ctx, id) {
|
||||
t.Error("a superseded observation is hidden from sessions, so nothing is merged into it")
|
||||
}
|
||||
if !(&Processor{}).canMergeInto(ctx, id) {
|
||||
t.Error("without an observation store there is nothing to check against")
|
||||
}
|
||||
}
|
||||
|
||||
+16
-11
@@ -156,7 +156,6 @@ type Service struct {
|
||||
updater *update.Updater
|
||||
rateLimiter *PerClientRateLimiter
|
||||
expensiveOpLimiter *ExpensiveOperationLimiter
|
||||
contextCache sync.Map
|
||||
version string
|
||||
recentQueriesBuf [maxRecentQueries]RecentSearchQuery
|
||||
wg sync.WaitGroup
|
||||
@@ -197,13 +196,6 @@ type staleVerifyRequest struct {
|
||||
observationID int64
|
||||
}
|
||||
|
||||
// contextCacheEntry caches clustering results for context injection.
|
||||
type contextCacheEntry struct {
|
||||
timestamp time.Time
|
||||
observations []*models.Observation
|
||||
obsCount int
|
||||
}
|
||||
|
||||
// RecentSearchQuery tracks a search query for analytics.
|
||||
type RecentSearchQuery struct {
|
||||
Timestamp time.Time `json:"timestamp"`
|
||||
@@ -1221,7 +1213,6 @@ func (s *Service) setupMiddleware() {
|
||||
|
||||
s.router.Use(middleware.Logger)
|
||||
s.router.Use(middleware.Recoverer)
|
||||
s.router.Use(middleware.RealIP)
|
||||
|
||||
// Add security headers (X-Frame-Options, X-Content-Type-Options, CSP, etc.)
|
||||
s.router.Use(SecurityHeaders)
|
||||
@@ -1235,11 +1226,24 @@ func (s *Service) setupMiddleware() {
|
||||
// Add gzip compression for responses >1KB (reduces bandwidth ~70% for JSON)
|
||||
s.router.Use(middleware.Compress(5)) // Level 5 = good balance of speed vs compression
|
||||
|
||||
// Apply per-client rate limiting (after RealIP so we get the real client IP)
|
||||
// Apply per-client rate limiting (keyed on the socket peer address; forwarding headers are untrusted)
|
||||
if s.rateLimiter != nil {
|
||||
s.router.Use(PerClientRateLimitMiddleware(s.rateLimiter))
|
||||
}
|
||||
|
||||
// Token authentication: opt-in and default-preserving.
|
||||
// When CLAUDE_MNEMONIC_AUTH_TOKEN is set, every request (except exempt
|
||||
// health/ready paths) must present that token via X-Auth-Token or
|
||||
// "Authorization: Bearer". When the env var is unset (the default), no
|
||||
// token is configured, auth stays disabled, and the server starts as
|
||||
// before. The worker binds loopback only, so this is defense-in-depth for
|
||||
// shared/multi-user hosts rather than a hard requirement.
|
||||
if authToken := os.Getenv("CLAUDE_MNEMONIC_AUTH_TOKEN"); authToken != "" {
|
||||
ta := NewTokenAuthWithToken(authToken)
|
||||
s.router.Use(ta.Middleware)
|
||||
log.Info().Msg("Token authentication enabled (CLAUDE_MNEMONIC_AUTH_TOKEN set)")
|
||||
}
|
||||
|
||||
// Note: Timeout middleware is applied per-route, not globally,
|
||||
// to avoid killing SSE connections which need to stay open indefinitely
|
||||
}
|
||||
@@ -1588,7 +1592,8 @@ func (s *Service) Start() error {
|
||||
port := config.GetWorkerPort()
|
||||
|
||||
s.server = &http.Server{
|
||||
Addr: fmt.Sprintf(":%d", port),
|
||||
// Bind loopback only so the worker is never exposed on external interfaces.
|
||||
Addr: fmt.Sprintf("127.0.0.1:%d", port),
|
||||
Handler: s.router,
|
||||
ReadHeaderTimeout: 10 * time.Second,
|
||||
ReadTimeout: 30 * time.Second,
|
||||
|
||||
+14
-18
@@ -5,33 +5,28 @@ import (
|
||||
"io/fs"
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"github.com/rs/zerolog/log"
|
||||
"sync"
|
||||
)
|
||||
|
||||
//go:embed static/*
|
||||
var staticFS embed.FS
|
||||
|
||||
// staticSubFS is the static subdirectory filesystem
|
||||
var staticSubFS fs.FS
|
||||
|
||||
// staticInitErr stores any error from static filesystem initialization
|
||||
var staticInitErr error
|
||||
|
||||
func init() {
|
||||
staticSubFS, staticInitErr = fs.Sub(staticFS, "static")
|
||||
if staticInitErr != nil {
|
||||
log.Warn().Err(staticInitErr).Msg("Static filesystem initialization failed - dashboard will be unavailable")
|
||||
}
|
||||
}
|
||||
// staticFiles lazily resolves the embedded "static" subdirectory exactly once.
|
||||
// Using sync.OnceValues avoids an init() and package-level mutable state while
|
||||
// preserving the previous behavior: the dashboard is served when resolution
|
||||
// succeeds and reports a clear error otherwise.
|
||||
var staticFiles = sync.OnceValues(func() (fs.FS, error) {
|
||||
return fs.Sub(staticFS, "static")
|
||||
})
|
||||
|
||||
// serveIndex serves the index.html file for the root path
|
||||
func serveIndex(w http.ResponseWriter, r *http.Request) {
|
||||
if staticInitErr != nil {
|
||||
sub, err := staticFiles()
|
||||
if err != nil {
|
||||
http.Error(w, "Dashboard unavailable: static files not initialized", http.StatusServiceUnavailable)
|
||||
return
|
||||
}
|
||||
content, err := fs.ReadFile(staticSubFS, "index.html")
|
||||
content, err := fs.ReadFile(sub, "index.html")
|
||||
if err != nil {
|
||||
http.Error(w, "Dashboard not found", http.StatusNotFound)
|
||||
return
|
||||
@@ -46,14 +41,15 @@ func serveIndex(w http.ResponseWriter, r *http.Request) {
|
||||
|
||||
// serveAssets serves static assets from the embedded filesystem
|
||||
func serveAssets(w http.ResponseWriter, r *http.Request) {
|
||||
if staticInitErr != nil {
|
||||
sub, err := staticFiles()
|
||||
if err != nil {
|
||||
http.Error(w, "Assets unavailable: static files not initialized", http.StatusServiceUnavailable)
|
||||
return
|
||||
}
|
||||
// Strip the /assets/ prefix and serve the file
|
||||
path := strings.TrimPrefix(r.URL.Path, "/")
|
||||
|
||||
content, err := fs.ReadFile(staticSubFS, path)
|
||||
content, err := fs.ReadFile(sub, path)
|
||||
if err != nil {
|
||||
http.Error(w, "Asset not found", http.StatusNotFound)
|
||||
return
|
||||
|
||||
+10
-10
@@ -13,7 +13,7 @@
|
||||
{
|
||||
"name": "claude-mnemonic",
|
||||
"description": "Persistent memory system for Claude Code - Go implementation with SQLite and ChromaDB vector search",
|
||||
"version": "0.13.17",
|
||||
"version": "0.21.95",
|
||||
"author": {
|
||||
"name": "lukaszraczylo",
|
||||
"url": "https://github.com/lukaszraczylo"
|
||||
@@ -23,25 +23,25 @@
|
||||
"license": "MIT",
|
||||
"homepage": "https://github.com/lukaszraczylo/claude-mnemonic",
|
||||
"releases": {
|
||||
"latest": "0.13.17",
|
||||
"latest": "0.21.95",
|
||||
"versions": {
|
||||
"0.13.17": {
|
||||
"releaseDate": "2026-05-31",
|
||||
"0.21.95": {
|
||||
"releaseDate": "2026-10-04",
|
||||
"downloads": {
|
||||
"darwin-arm64": {
|
||||
"url": "https://github.com/lukaszraczylo/claude-mnemonic/releases/download/v0.13.17/claude-mnemonic_0.13.17_darwin_arm64.tar.gz",
|
||||
"url": "https://github.com/lukaszraczylo/claude-mnemonic/releases/download/v0.21.95/claude-mnemonic_0.21.95_darwin_arm64.tar.gz",
|
||||
"format": "tar.gz",
|
||||
"sha256": "efcdcdf0ba20a03b31bb501dba6e659ffe4a55f47ce82b3ccfbc8a47c56d1476"
|
||||
"sha256": "3224be650ead03a23078f7033909ecaedf89c88d665b3aa83c70b2eb952722b9"
|
||||
},
|
||||
"linux-amd64": {
|
||||
"url": "https://github.com/lukaszraczylo/claude-mnemonic/releases/download/v0.13.17/claude-mnemonic_0.13.17_linux_amd64.tar.gz",
|
||||
"url": "https://github.com/lukaszraczylo/claude-mnemonic/releases/download/v0.21.95/claude-mnemonic_0.21.95_linux_amd64.tar.gz",
|
||||
"format": "tar.gz",
|
||||
"sha256": "1aab142d32aaf84f148ec42629959d6038d740187b8270afda426d28acc1f707"
|
||||
"sha256": "4af6927b7857b03143e6eb817da7f44c39dd2b5674a7d1b7f505db7c3dd25b2f"
|
||||
},
|
||||
"windows-amd64": {
|
||||
"url": "https://github.com/lukaszraczylo/claude-mnemonic/releases/download/v0.13.17/claude-mnemonic_0.13.17_windows_amd64.zip",
|
||||
"url": "https://github.com/lukaszraczylo/claude-mnemonic/releases/download/v0.21.95/claude-mnemonic_0.21.95_windows_amd64.zip",
|
||||
"format": "zip",
|
||||
"sha256": "e03b34be90a29109f45c375eecccebfb3d1662cc3d4814aa86386d147f6128bd"
|
||||
"sha256": "d779920957eea6b21a344556b5a36ec2a1a6c596e095575e73110c5cb997fc1a"
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,94 @@
|
||||
package hooks
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
// maxResponseBytes caps a worker response so a misbehaving peer cannot exhaust memory.
|
||||
const maxResponseBytes = 32 << 20
|
||||
|
||||
// Minimal HTTP/1.0 client over a raw TCP socket. Linking net/http adds ~3 MB to
|
||||
// every hook binary; hooks only speak plain JSON to the local worker. HTTP/1.0
|
||||
// makes the server skip chunked encoding and close the connection, so the body
|
||||
// is simply everything after the header block.
|
||||
type httpResponse struct {
|
||||
StatusText string
|
||||
Body []byte
|
||||
Status int
|
||||
}
|
||||
|
||||
func httpDo(ctx context.Context, timeout time.Duration, method string, port int, path string, body []byte) (*httpResponse, error) {
|
||||
ctx, cancel := context.WithTimeout(ctx, timeout)
|
||||
defer cancel()
|
||||
|
||||
addr := net.JoinHostPort("127.0.0.1", strconv.Itoa(port))
|
||||
var d net.Dialer
|
||||
conn, err := d.DialContext(ctx, "tcp", addr)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer func() { _ = conn.Close() }()
|
||||
|
||||
if dl, ok := ctx.Deadline(); ok {
|
||||
_ = conn.SetDeadline(dl)
|
||||
}
|
||||
stop := context.AfterFunc(ctx, func() { _ = conn.Close() })
|
||||
defer stop()
|
||||
|
||||
var req bytes.Buffer
|
||||
fmt.Fprintf(&req, "%s %s HTTP/1.0\r\nHost: %s\r\n", method, path, addr)
|
||||
if body != nil {
|
||||
fmt.Fprintf(&req, "Content-Type: application/json\r\nContent-Length: %d\r\n", len(body))
|
||||
}
|
||||
req.WriteString("\r\n")
|
||||
req.Write(body)
|
||||
if _, err := conn.Write(req.Bytes()); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
raw, err := io.ReadAll(io.LimitReader(bufio.NewReader(conn), maxResponseBytes))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return parseHTTPResponse(raw)
|
||||
}
|
||||
|
||||
func parseHTTPResponse(raw []byte) (*httpResponse, error) {
|
||||
head, body, found := bytes.Cut(raw, []byte("\r\n\r\n"))
|
||||
if !found {
|
||||
return nil, errors.New("malformed http response")
|
||||
}
|
||||
statusLine, _, _ := strings.Cut(string(head), "\r\n")
|
||||
_, rest, ok := strings.Cut(statusLine, " ")
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("malformed status line %q", statusLine)
|
||||
}
|
||||
codeStr, _, _ := strings.Cut(rest, " ")
|
||||
code, err := strconv.Atoi(codeStr)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("malformed status line %q", statusLine)
|
||||
}
|
||||
return &httpResponse{Status: code, StatusText: rest, Body: body}, nil
|
||||
}
|
||||
|
||||
// GETBody fetches path from the worker within timeout and returns the raw body.
|
||||
// Statuses of 400 and above are returned as errors.
|
||||
func GETBody(ctx context.Context, timeout time.Duration, port int, path string) ([]byte, error) {
|
||||
resp, err := httpDo(ctx, timeout, "GET", port, path, nil)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if resp.Status >= 400 {
|
||||
return nil, fmt.Errorf("request failed: %s", resp.StatusText)
|
||||
}
|
||||
return resp.Body, nil
|
||||
}
|
||||
+39
-137
@@ -2,12 +2,10 @@
|
||||
package hooks
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net"
|
||||
"net/http"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
@@ -28,6 +26,9 @@ const (
|
||||
// HealthCheckTimeout is the timeout for health checks.
|
||||
HealthCheckTimeout = 2 * time.Second
|
||||
|
||||
// hookTimeout bounds a single hook->worker request.
|
||||
hookTimeout = 10 * time.Second
|
||||
|
||||
// StartupTimeout is the timeout for worker startup.
|
||||
StartupTimeout = 10 * time.Second
|
||||
|
||||
@@ -52,26 +53,6 @@ var (
|
||||
// circuitBreakerMu protects lastStartupFailure.
|
||||
circuitBreakerMu sync.Mutex
|
||||
lastStartupFailure time.Time
|
||||
|
||||
// hookClient is a shared HTTP client for hook->worker requests.
|
||||
// DisableKeepAlives prevents TIME_WAIT connection leaks since each hook
|
||||
// is a separate OS process that exits quickly.
|
||||
hookClient = &http.Client{
|
||||
Timeout: 10 * time.Second,
|
||||
Transport: &http.Transport{
|
||||
DisableKeepAlives: true,
|
||||
MaxIdleConns: 1,
|
||||
},
|
||||
}
|
||||
|
||||
// healthClient is a shared HTTP client for health/version checks.
|
||||
healthClient = &http.Client{
|
||||
Timeout: HealthCheckTimeout,
|
||||
Transport: &http.Transport{
|
||||
DisableKeepAlives: true,
|
||||
MaxIdleConns: 1,
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
// IsWorkerAvailable performs a fast check without network calls.
|
||||
@@ -110,22 +91,21 @@ func GetWorkerPort() int {
|
||||
// Parses the JSON health response to check the "ready" field when available.
|
||||
// Falls back to HTTP status code check for backwards compatibility.
|
||||
func IsWorkerRunning(port int) bool {
|
||||
resp, err := healthClient.Get(fmt.Sprintf("http://127.0.0.1:%d/api/health", port))
|
||||
resp, err := httpDo(context.Background(), HealthCheckTimeout, "GET", port, "/api/health", nil)
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
|
||||
// Try to parse JSON response for structured health check
|
||||
var health struct {
|
||||
Ready bool `json:"ready"`
|
||||
}
|
||||
if err := json.NewDecoder(resp.Body).Decode(&health); err == nil {
|
||||
if err := json.Unmarshal(resp.Body, &health); err == nil {
|
||||
return health.Ready
|
||||
}
|
||||
|
||||
// Fallback: treat HTTP 200 as healthy (backwards compatibility)
|
||||
return resp.StatusCode == http.StatusOK
|
||||
return resp.Status == 200
|
||||
}
|
||||
|
||||
// workerCachePath returns the path to the worker cache file.
|
||||
@@ -396,18 +376,17 @@ func updateCacheFromPort(port int) {
|
||||
|
||||
// GetWorkerVersion gets the version of the running worker.
|
||||
func GetWorkerVersion(port int) string {
|
||||
resp, err := healthClient.Get(fmt.Sprintf("http://127.0.0.1:%d/api/version", port))
|
||||
resp, err := httpDo(context.Background(), HealthCheckTimeout, "GET", port, "/api/version", nil)
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
if resp.Status != 200 {
|
||||
return ""
|
||||
}
|
||||
|
||||
var result map[string]string
|
||||
if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
|
||||
if err := json.Unmarshal(resp.Body, &result); err != nil {
|
||||
return ""
|
||||
}
|
||||
|
||||
@@ -510,58 +489,46 @@ func findWorkerBinary() string {
|
||||
return ""
|
||||
}
|
||||
|
||||
// POST sends a POST request to the worker.
|
||||
func POST(port int, path string, body interface{}) (map[string]interface{}, error) {
|
||||
jsonBody, err := json.Marshal(body)
|
||||
// request performs one worker call and returns the decoded JSON object.
|
||||
// A non-JSON body yields (nil, nil) when lenient, else the decode error.
|
||||
func request(ctx context.Context, method string, port int, path string, body interface{}, lenient bool) (map[string]interface{}, error) {
|
||||
var payload []byte
|
||||
if body != nil {
|
||||
var err error
|
||||
if payload, err = json.Marshal(body); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
|
||||
resp, err := httpDo(ctx, hookTimeout, method, port, path, payload)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
resp, err := hookClient.Post(
|
||||
fmt.Sprintf("http://127.0.0.1:%d%s", port, path),
|
||||
"application/json",
|
||||
bytes.NewReader(jsonBody),
|
||||
)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
|
||||
if resp.StatusCode >= 400 {
|
||||
return nil, fmt.Errorf("request failed: %s", resp.Status)
|
||||
if resp.Status >= 400 {
|
||||
return nil, fmt.Errorf("request failed: %s", resp.StatusText)
|
||||
}
|
||||
|
||||
var result map[string]interface{}
|
||||
if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
|
||||
// Not all endpoints return JSON
|
||||
return nil, nil
|
||||
if err := json.Unmarshal(resp.Body, &result); err != nil {
|
||||
if lenient {
|
||||
// Not all endpoints return JSON
|
||||
return nil, nil
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return result, nil
|
||||
}
|
||||
|
||||
// POST sends a POST request to the worker.
|
||||
func POST(port int, path string, body interface{}) (map[string]interface{}, error) {
|
||||
return request(context.Background(), "POST", port, path, body, true)
|
||||
}
|
||||
|
||||
// POSTWithContext sends a POST request using the provided context.
|
||||
// Used for fire-and-forget calls where we want to control the timeout externally.
|
||||
func POSTWithContext(ctx context.Context, port int, path string, body interface{}) error {
|
||||
jsonBody, err := json.Marshal(body)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost,
|
||||
fmt.Sprintf("http://127.0.0.1:%d%s", port, path),
|
||||
bytes.NewReader(jsonBody))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
|
||||
resp, err := hookClient.Do(req)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
return nil
|
||||
_, err := request(ctx, "POST", port, path, body, true)
|
||||
return err
|
||||
}
|
||||
|
||||
// POSTWithContextResult sends a POST request using the provided context and
|
||||
@@ -570,56 +537,12 @@ func POSTWithContext(ctx context.Context, port int, path string, body interface{
|
||||
// aborts at the hook deadline instead of blocking for the full client timeout.
|
||||
// A non-JSON body is returned as (nil, nil), matching POST's behavior.
|
||||
func POSTWithContextResult(ctx context.Context, port int, path string, body interface{}) (map[string]interface{}, error) {
|
||||
jsonBody, err := json.Marshal(body)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost,
|
||||
fmt.Sprintf("http://127.0.0.1:%d%s", port, path),
|
||||
bytes.NewReader(jsonBody))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
|
||||
resp, err := hookClient.Do(req)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
|
||||
if resp.StatusCode >= 400 {
|
||||
return nil, fmt.Errorf("request failed: %s", resp.Status)
|
||||
}
|
||||
|
||||
var result map[string]interface{}
|
||||
if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
|
||||
// Not all endpoints return JSON
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
return result, nil
|
||||
return request(ctx, "POST", port, path, body, true)
|
||||
}
|
||||
|
||||
// GET sends a GET request to the worker.
|
||||
func GET(port int, path string) (map[string]interface{}, error) {
|
||||
resp, err := hookClient.Get(fmt.Sprintf("http://127.0.0.1:%d%s", port, path))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
|
||||
if resp.StatusCode >= 400 {
|
||||
return nil, fmt.Errorf("request failed: %s", resp.Status)
|
||||
}
|
||||
|
||||
var result map[string]interface{}
|
||||
if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return result, nil
|
||||
return request(context.Background(), "GET", port, path, nil, false)
|
||||
}
|
||||
|
||||
// GETWithContext sends a GET request using the provided context and decodes the
|
||||
@@ -627,28 +550,7 @@ func GET(port int, path string) (map[string]interface{}, error) {
|
||||
// Used on the prompt critical path so a wedged worker aborts at the hook
|
||||
// deadline instead of blocking for the full client timeout.
|
||||
func GETWithContext(ctx context.Context, port int, path string) (map[string]interface{}, error) {
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet,
|
||||
fmt.Sprintf("http://127.0.0.1:%d%s", port, path), nil)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
resp, err := hookClient.Do(req)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
|
||||
if resp.StatusCode >= 400 {
|
||||
return nil, fmt.Errorf("request failed: %s", resp.Status)
|
||||
}
|
||||
|
||||
var result map[string]interface{}
|
||||
if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return result, nil
|
||||
return request(ctx, "GET", port, path, nil, false)
|
||||
}
|
||||
|
||||
// versionsCompatible checks if two versions are compatible for dev builds.
|
||||
|
||||
+98
-15
@@ -5,8 +5,12 @@ import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strconv"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
@@ -1481,22 +1485,101 @@ func TestSleepCtx_Cancelled(t *testing.T) {
|
||||
assert.Less(t, elapsed, 500*time.Millisecond)
|
||||
}
|
||||
|
||||
// TestHookClients_DisableKeepAlives asserts shared clients disable keep-alives
|
||||
// to prevent TIME_WAIT connection leaks in short-lived hook processes.
|
||||
func TestHookClients_DisableKeepAlives(t *testing.T) {
|
||||
hTransport, ok := hookClient.Transport.(*http.Transport)
|
||||
require.True(t, ok, "hookClient.Transport should be *http.Transport")
|
||||
assert.True(t, hTransport.DisableKeepAlives, "hookClient must disable keep-alives")
|
||||
assert.Equal(t, 1, hTransport.MaxIdleConns)
|
||||
// TestHTTPDo_Behaviour covers the raw HTTP/1.0 client against a real server.
|
||||
func TestHTTPDo_Behaviour(t *testing.T) {
|
||||
big := strings.Repeat("x", 200_000)
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
switch r.URL.Path {
|
||||
case "/echo":
|
||||
b, _ := io.ReadAll(r.Body)
|
||||
assert.Equal(t, "application/json", r.Header.Get("Content-Type"))
|
||||
_, _ = w.Write(b)
|
||||
case "/big":
|
||||
_, _ = w.Write([]byte(`{"v":"` + big + `"}`))
|
||||
case "/missing":
|
||||
w.WriteHeader(http.StatusNotFound)
|
||||
case "/empty":
|
||||
w.WriteHeader(http.StatusNoContent)
|
||||
}
|
||||
}))
|
||||
defer srv.Close()
|
||||
port, err := strconv.Atoi(srv.URL[strings.LastIndex(srv.URL, ":")+1:])
|
||||
require.NoError(t, err)
|
||||
|
||||
hcTransport, ok := healthClient.Transport.(*http.Transport)
|
||||
require.True(t, ok, "healthClient.Transport should be *http.Transport")
|
||||
assert.True(t, hcTransport.DisableKeepAlives, "healthClient must disable keep-alives")
|
||||
assert.Equal(t, 1, hcTransport.MaxIdleConns)
|
||||
tests := []struct {
|
||||
name string
|
||||
method string
|
||||
path string
|
||||
body []byte
|
||||
wantStatus int
|
||||
wantLen int
|
||||
}{
|
||||
{"post echoes body", "POST", "/echo", []byte(`{"a":1}`), 200, 7},
|
||||
{"large body over server buffer", "GET", "/big", nil, 200, len(big) + 8},
|
||||
{"not found", "GET", "/missing", nil, 404, 0},
|
||||
{"no content", "GET", "/empty", nil, 204, 0},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
resp, err := httpDo(context.Background(), time.Second, tt.method, port, tt.path, tt.body)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, tt.wantStatus, resp.Status)
|
||||
assert.Len(t, resp.Body, tt.wantLen)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestHookClient_Timeout verifies hookClient timeout is set.
|
||||
func TestHookClient_Timeout(t *testing.T) {
|
||||
assert.Equal(t, 10*time.Second, hookClient.Timeout)
|
||||
assert.Equal(t, HealthCheckTimeout, healthClient.Timeout)
|
||||
// TestHTTPDo_Failures covers connection, timeout, cancel and malformed replies.
|
||||
func TestHTTPDo_Failures(t *testing.T) {
|
||||
t.Run("connection refused", func(t *testing.T) {
|
||||
l, err := net.Listen("tcp", "127.0.0.1:0")
|
||||
require.NoError(t, err)
|
||||
port := l.Addr().(*net.TCPAddr).Port
|
||||
require.NoError(t, l.Close())
|
||||
_, err = httpDo(context.Background(), time.Second, "GET", port, "/", nil)
|
||||
assert.Error(t, err)
|
||||
})
|
||||
|
||||
t.Run("wedged server times out", func(t *testing.T) {
|
||||
l, err := net.Listen("tcp", "127.0.0.1:0")
|
||||
require.NoError(t, err)
|
||||
defer func() { _ = l.Close() }()
|
||||
go func() {
|
||||
c, err := l.Accept()
|
||||
if err == nil {
|
||||
defer func() { _ = c.Close() }()
|
||||
time.Sleep(2 * time.Second)
|
||||
}
|
||||
}()
|
||||
start := time.Now()
|
||||
_, err = httpDo(context.Background(), 200*time.Millisecond, "GET", l.Addr().(*net.TCPAddr).Port, "/", nil)
|
||||
assert.Error(t, err)
|
||||
assert.Less(t, time.Since(start), time.Second)
|
||||
})
|
||||
|
||||
t.Run("cancelled context aborts", func(t *testing.T) {
|
||||
l, err := net.Listen("tcp", "127.0.0.1:0")
|
||||
require.NoError(t, err)
|
||||
defer func() { _ = l.Close() }()
|
||||
go func() {
|
||||
c, err := l.Accept()
|
||||
if err == nil {
|
||||
defer func() { _ = c.Close() }()
|
||||
time.Sleep(2 * time.Second)
|
||||
}
|
||||
}()
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
|
||||
defer cancel()
|
||||
start := time.Now()
|
||||
_, err = httpDo(ctx, 5*time.Second, "GET", l.Addr().(*net.TCPAddr).Port, "/", nil)
|
||||
assert.Error(t, err)
|
||||
assert.Less(t, time.Since(start), time.Second)
|
||||
})
|
||||
|
||||
t.Run("malformed response", func(t *testing.T) {
|
||||
for _, raw := range []string{"", "garbage", "HTTP/1.0\r\n\r\n", "HTTP/1.0 abc OK\r\n\r\n"} {
|
||||
_, err := parseHTTPResponse([]byte(raw))
|
||||
assert.Error(t, err, "input %q", raw)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
@@ -5,6 +5,7 @@ package sanitize
|
||||
import (
|
||||
"regexp"
|
||||
"strings"
|
||||
"sync"
|
||||
)
|
||||
|
||||
// systemXMLTags lists Claude Code internal XML tags that should be stripped
|
||||
@@ -30,16 +31,17 @@ var systemXMLTags = []string{
|
||||
"antml_function_calls",
|
||||
}
|
||||
|
||||
// systemXMLRegexps are compiled regexps for each tag, built once at init.
|
||||
var systemXMLRegexps []*regexp.Regexp
|
||||
|
||||
func init() {
|
||||
systemXMLRegexps = make([]*regexp.Regexp, len(systemXMLTags))
|
||||
// systemXMLRegexps lazily compiles a regexp for each tag exactly once on first
|
||||
// use. The result is immutable after compilation, so sync.OnceValue safely
|
||||
// shares it across callers without re-compiling at package init.
|
||||
var systemXMLRegexps = sync.OnceValue(func() []*regexp.Regexp {
|
||||
res := make([]*regexp.Regexp, len(systemXMLTags))
|
||||
for i, tag := range systemXMLTags {
|
||||
// Match opening tag (with optional attributes), content (including newlines), and closing tag
|
||||
systemXMLRegexps[i] = regexp.MustCompile(`(?s)<` + regexp.QuoteMeta(tag) + `[^>]*>.*?</` + regexp.QuoteMeta(tag) + `>`)
|
||||
res[i] = regexp.MustCompile(`(?s)<` + regexp.QuoteMeta(tag) + `[^>]*>.*?</` + regexp.QuoteMeta(tag) + `>`)
|
||||
}
|
||||
}
|
||||
return res
|
||||
})
|
||||
|
||||
// StripSystemXML removes known Claude Code internal XML blocks from text.
|
||||
// This prevents system artifacts like <task-notification>, <system-reminder>,
|
||||
@@ -50,7 +52,7 @@ func StripSystemXML(s string) string {
|
||||
return s
|
||||
}
|
||||
|
||||
for _, re := range systemXMLRegexps {
|
||||
for _, re := range systemXMLRegexps() {
|
||||
s = re.ReplaceAllString(s, "")
|
||||
}
|
||||
|
||||
|
||||
@@ -23,8 +23,8 @@ func TestStripSystemXML(t *testing.T) {
|
||||
expected: "",
|
||||
},
|
||||
{
|
||||
name: "task-notification block",
|
||||
input: "Before\n<task-notification>\n<task-id>abc123</task-id>\n<status>completed</status>\n<summary>Agent completed</summary>\n<result>Some result text</result>\n</task-notification>\nAfter",
|
||||
name: "task-notification block",
|
||||
input: "Before\n<task-notification>\n<task-id>abc123</task-id>\n<status>completed</status>\n<summary>Agent completed</summary>\n<result>Some result text</result>\n</task-notification>\nAfter",
|
||||
expected: "Before\n\nAfter",
|
||||
},
|
||||
{
|
||||
@@ -78,8 +78,8 @@ func TestStripSystemXML(t *testing.T) {
|
||||
expected: "Before\n\nAfter",
|
||||
},
|
||||
{
|
||||
name: "real-world task notification",
|
||||
input: "Here is the analysis:\n<task-notification>\n<task-id>a077fd04aed547ce1</task-id>\n<tool-use-id>toolu_01Kw2GwsArQQR1F9aV3VfzhP</tool-use-id>\n<status>completed</status>\n<summary>Agent \"Analyze agency-agents repo structure\" completed</summary>\n<result>Now I have all the information I need. Let me compile a comprehensive report.\n\n## Agency-Agents Directory Exploration Report\n\n### 1. Agent File Count and Organization\n\n**Total Agent Files (excluding strategy docs):** 61 agents across 9 categories</result>\n</task-notification>\nThe repo has 61 agents.",
|
||||
name: "real-world task notification",
|
||||
input: "Here is the analysis:\n<task-notification>\n<task-id>a077fd04aed547ce1</task-id>\n<tool-use-id>toolu_01Kw2GwsArQQR1F9aV3VfzhP</tool-use-id>\n<status>completed</status>\n<summary>Agent \"Analyze agency-agents repo structure\" completed</summary>\n<result>Now I have all the information I need. Let me compile a comprehensive report.\n\n## Agency-Agents Directory Exploration Report\n\n### 1. Agent File Count and Organization\n\n**Total Agent Files (excluding strategy docs):** 61 agents across 9 categories</result>\n</task-notification>\nThe repo has 61 agents.",
|
||||
expected: "Here is the analysis:\n\nThe repo has 61 agents.",
|
||||
},
|
||||
}
|
||||
|
||||
@@ -10,7 +10,6 @@
|
||||
"mcpServers": {
|
||||
"claude-mnemonic": {
|
||||
"command": "${CLAUDE_PLUGIN_ROOT}/mcp-server",
|
||||
"args": ["--project", "${CLAUDE_PROJECT}"],
|
||||
"env": {}
|
||||
}
|
||||
},
|
||||
|
||||
@@ -1,3 +1,8 @@
|
||||
---
|
||||
description: Restart the Claude Mnemonic worker process
|
||||
allowed-tools: Bash(curl:*)
|
||||
---
|
||||
|
||||
# Restart Claude Mnemonic Worker
|
||||
|
||||
Restart the claude-mnemonic worker process. Use this command when experiencing issues with the memory system.
|
||||
|
||||
@@ -6,7 +6,7 @@
|
||||
"hooks": [
|
||||
{
|
||||
"type": "command",
|
||||
"command": "${CLAUDE_PLUGIN_ROOT}/hooks/session-start",
|
||||
"command": "\"${CLAUDE_PLUGIN_ROOT}/hooks/session-start\"",
|
||||
"timeout": 30
|
||||
}
|
||||
]
|
||||
@@ -17,7 +17,7 @@
|
||||
"hooks": [
|
||||
{
|
||||
"type": "command",
|
||||
"command": "${CLAUDE_PLUGIN_ROOT}/hooks/user-prompt",
|
||||
"command": "\"${CLAUDE_PLUGIN_ROOT}/hooks/user-prompt\"",
|
||||
"timeout": 10
|
||||
}
|
||||
]
|
||||
@@ -29,7 +29,7 @@
|
||||
"hooks": [
|
||||
{
|
||||
"type": "command",
|
||||
"command": "${CLAUDE_PLUGIN_ROOT}/hooks/post-tool-use",
|
||||
"command": "\"${CLAUDE_PLUGIN_ROOT}/hooks/post-tool-use\"",
|
||||
"timeout": 10
|
||||
}
|
||||
]
|
||||
@@ -40,7 +40,7 @@
|
||||
"hooks": [
|
||||
{
|
||||
"type": "command",
|
||||
"command": "${CLAUDE_PLUGIN_ROOT}/hooks/subagent-stop",
|
||||
"command": "\"${CLAUDE_PLUGIN_ROOT}/hooks/subagent-stop\"",
|
||||
"timeout": 10
|
||||
}
|
||||
]
|
||||
@@ -51,7 +51,7 @@
|
||||
"hooks": [
|
||||
{
|
||||
"type": "command",
|
||||
"command": "${CLAUDE_PLUGIN_ROOT}/hooks/stop",
|
||||
"command": "\"${CLAUDE_PLUGIN_ROOT}/hooks/stop\"",
|
||||
"timeout": 30
|
||||
}
|
||||
]
|
||||
|
||||
@@ -6,7 +6,7 @@
|
||||
|
||||
set -e
|
||||
|
||||
ONNX_VERSION="1.26.0"
|
||||
ONNX_VERSION="1.30.0"
|
||||
ASSETS_DIR="internal/embedding/assets/lib"
|
||||
PLATFORM="${1:-all}"
|
||||
FORCE_DOWNLOAD=false
|
||||
@@ -51,9 +51,22 @@ get_installed_version() {
|
||||
fi
|
||||
}
|
||||
|
||||
# Gzip the libraries: Go embeds the .gz and gunzips once at first use, which
|
||||
# shrinks the worker binary by roughly 25 MB.
|
||||
compress_libs() {
|
||||
local plat="$1"
|
||||
local f
|
||||
for f in "${ASSETS_DIR}/${plat}"/*.dylib "${ASSETS_DIR}/${plat}"/*.so "${ASSETS_DIR}/${plat}"/*.dll; do
|
||||
if [ -f "$f" ]; then
|
||||
gzip -9 -n -f "$f"
|
||||
fi
|
||||
done
|
||||
}
|
||||
|
||||
# Write version file after successful download
|
||||
write_version_file() {
|
||||
local plat="$1"
|
||||
compress_libs "$plat"
|
||||
echo "${ONNX_VERSION}" > "${ASSETS_DIR}/${plat}/.version"
|
||||
}
|
||||
|
||||
@@ -110,9 +123,9 @@ download_windows_amd64() {
|
||||
lib_exists() {
|
||||
local plat="$1"
|
||||
case "$plat" in
|
||||
darwin-*) [ -f "${ASSETS_DIR}/${plat}/libonnxruntime.dylib" ] ;;
|
||||
linux-*) [ -f "${ASSETS_DIR}/${plat}/libonnxruntime.so" ] ;;
|
||||
windows-*) [ -f "${ASSETS_DIR}/${plat}/onnxruntime.dll" ] ;;
|
||||
darwin-*) [ -f "${ASSETS_DIR}/${plat}/libonnxruntime.dylib.gz" ] ;;
|
||||
linux-*) [ -f "${ASSETS_DIR}/${plat}/libonnxruntime.so.gz" ] ;;
|
||||
windows-*) [ -f "${ASSETS_DIR}/${plat}/onnxruntime.dll.gz" ] ;;
|
||||
*) return 1 ;;
|
||||
esac
|
||||
}
|
||||
|
||||
@@ -0,0 +1,276 @@
|
||||
{
|
||||
"docs": [
|
||||
{
|
||||
"type": "bugfix",
|
||||
"title": "Fixed race in session cache",
|
||||
"narrative": "Concurrent map writes in the session cache caused panics; guarded with a RWMutex."
|
||||
},
|
||||
{
|
||||
"type": "decision",
|
||||
"title": "Chose SQLite WAL mode",
|
||||
"narrative": "WAL lets readers proceed while a writer commits, so hooks can query during writes."
|
||||
},
|
||||
{
|
||||
"type": "feature",
|
||||
"title": "Added Prometheus metrics endpoint",
|
||||
"narrative": "Exposes request latency histograms and error counters at /metrics for scraping."
|
||||
},
|
||||
{
|
||||
"type": "bugfix",
|
||||
"title": "Fixed memory leak in websocket hub",
|
||||
"narrative": "Goroutines were never released when clients disconnected; added context cancellation."
|
||||
},
|
||||
{
|
||||
"type": "refactor",
|
||||
"title": "Extracted retry helper with backoff",
|
||||
"narrative": "Exponential backoff with jitter now lives in one function used by all HTTP callers."
|
||||
},
|
||||
{
|
||||
"type": "decision",
|
||||
"title": "Use argon2id for password hashing",
|
||||
"narrative": "bcrypt cost was too low for new hardware; argon2id with 64MB memory chosen."
|
||||
},
|
||||
{
|
||||
"type": "feature",
|
||||
"title": "Dark mode toggle in dashboard",
|
||||
"narrative": "Theme preference persists in localStorage and follows prefers-color-scheme by default."
|
||||
},
|
||||
{
|
||||
"type": "discovery",
|
||||
"title": "Kubernetes readiness probe timing",
|
||||
"narrative": "Pods were killed during warmup because the readiness probe fired before the model loaded."
|
||||
},
|
||||
{
|
||||
"type": "bugfix",
|
||||
"title": "Fixed off-by-one in pagination",
|
||||
"narrative": "The last page skipped one row when total count was an exact multiple of page size."
|
||||
},
|
||||
{
|
||||
"type": "decision",
|
||||
"title": "Adopted table-driven tests",
|
||||
"narrative": "Subtests with named cases make failures easier to locate and new cases cheap to add."
|
||||
},
|
||||
{
|
||||
"type": "feature",
|
||||
"title": "Rate limiting middleware",
|
||||
"narrative": "Token bucket per client IP returns 429 with a Retry-After header when exceeded."
|
||||
},
|
||||
{
|
||||
"type": "bugfix",
|
||||
"title": "Fixed timezone bug in scheduler",
|
||||
"narrative": "Cron jobs fired an hour late after daylight saving because times were stored in local zone."
|
||||
},
|
||||
{
|
||||
"type": "discovery",
|
||||
"title": "Slow query caused by missing index",
|
||||
"narrative": "EXPLAIN showed a sequential scan on the events table; a composite index cut it from 900ms to 4ms."
|
||||
},
|
||||
{
|
||||
"type": "refactor",
|
||||
"title": "Split monolithic handler file",
|
||||
"narrative": "The 3000 line handlers file was divided by domain: sessions, observations, search."
|
||||
},
|
||||
{
|
||||
"type": "decision",
|
||||
"title": "Embeddings computed locally with ONNX",
|
||||
"narrative": "Avoids sending user text to third parties and works offline."
|
||||
},
|
||||
{
|
||||
"type": "feature",
|
||||
"title": "Export observations as JSON",
|
||||
"narrative": "New CLI command dumps all stored observations to a portable JSON file for backup."
|
||||
},
|
||||
{
|
||||
"type": "bugfix",
|
||||
"title": "Fixed CORS preflight failure",
|
||||
"narrative": "OPTIONS requests were rejected by auth middleware; moved CORS handling before authentication."
|
||||
},
|
||||
{
|
||||
"type": "discovery",
|
||||
"title": "Docker image size dominated by build cache",
|
||||
"narrative": "Multi-stage build with a distroless runtime reduced the image from 1.2GB to 80MB."
|
||||
},
|
||||
{
|
||||
"type": "decision",
|
||||
"title": "Structured logging with zerolog",
|
||||
"narrative": "JSON logs with fields are easier to filter in aggregation than free-form strings."
|
||||
},
|
||||
{
|
||||
"type": "feature",
|
||||
"title": "Full text search with FTS5",
|
||||
"narrative": "Keyword queries use SQLite FTS5 with porter stemming alongside vector similarity."
|
||||
},
|
||||
{
|
||||
"type": "bugfix",
|
||||
"title": "Fixed panic on nil config",
|
||||
"narrative": "Missing settings file returned a nil pointer; defaults are now loaded instead."
|
||||
},
|
||||
{
|
||||
"type": "discovery",
|
||||
"title": "GC pauses from large slices",
|
||||
"narrative": "Preallocating slice capacity removed repeated growth and lowered p99 latency."
|
||||
},
|
||||
{
|
||||
"type": "decision",
|
||||
"title": "Graceful shutdown on SIGTERM",
|
||||
"narrative": "Server drains in-flight requests for 30 seconds before exiting so deploys drop nothing."
|
||||
},
|
||||
{
|
||||
"type": "refactor",
|
||||
"title": "Replaced global state with dependency injection",
|
||||
"narrative": "Services receive their stores through constructors, which makes tests independent."
|
||||
},
|
||||
{
|
||||
"type": "feature",
|
||||
"title": "OAuth login with GitHub",
|
||||
"narrative": "Users authenticate via the GitHub authorization code flow and receive a session cookie."
|
||||
},
|
||||
{
|
||||
"type": "bugfix",
|
||||
"title": "Fixed flaky integration test",
|
||||
"narrative": "The test relied on sleep; replaced with polling on a readiness condition."
|
||||
},
|
||||
{
|
||||
"type": "discovery",
|
||||
"title": "Git hooks slowed commits",
|
||||
"narrative": "Pre-commit ran the full test suite; limited it to changed packages."
|
||||
},
|
||||
{
|
||||
"type": "decision",
|
||||
"title": "Pinned dependency versions",
|
||||
"narrative": "Lockfile committed and CI uses frozen installs to keep builds reproducible."
|
||||
},
|
||||
{
|
||||
"type": "feature",
|
||||
"title": "Webhook retries with dead letter queue",
|
||||
"narrative": "Failed deliveries retry five times and then land in a dead letter table for inspection."
|
||||
},
|
||||
{
|
||||
"type": "discovery",
|
||||
"title": "Unicode normalization broke search",
|
||||
"narrative": "Accented characters did not match until queries and content were NFC normalized."
|
||||
}
|
||||
],
|
||||
"queries": [
|
||||
{
|
||||
"doc": 0,
|
||||
"query": "why did concurrent writes to the cache crash the program"
|
||||
},
|
||||
{
|
||||
"doc": 1,
|
||||
"query": "how can readers keep working while the database is being written"
|
||||
},
|
||||
{
|
||||
"doc": 2,
|
||||
"query": "where do I scrape request latency numbers from"
|
||||
},
|
||||
{
|
||||
"doc": 3,
|
||||
"query": "goroutines not freed after clients drop connection"
|
||||
},
|
||||
{
|
||||
"doc": 4,
|
||||
"query": "shared function for repeating failed requests with growing delays"
|
||||
},
|
||||
{
|
||||
"doc": 5,
|
||||
"query": "which algorithm protects stored passwords"
|
||||
},
|
||||
{
|
||||
"doc": 6,
|
||||
"query": "can the UI switch to a dark theme"
|
||||
},
|
||||
{
|
||||
"doc": 7,
|
||||
"query": "containers restarted while the model was still loading"
|
||||
},
|
||||
{
|
||||
"doc": 8,
|
||||
"query": "last page of results missing an item"
|
||||
},
|
||||
{
|
||||
"doc": 9,
|
||||
"query": "why do our tests use named cases"
|
||||
},
|
||||
{
|
||||
"doc": 10,
|
||||
"query": "how are abusive clients throttled"
|
||||
},
|
||||
{
|
||||
"doc": 11,
|
||||
"query": "scheduled jobs running late after clocks changed"
|
||||
},
|
||||
{
|
||||
"doc": 12,
|
||||
"query": "query got fast after adding an index"
|
||||
},
|
||||
{
|
||||
"doc": 13,
|
||||
"query": "large source file was broken up by domain"
|
||||
},
|
||||
{
|
||||
"doc": 14,
|
||||
"query": "do we send text to an external embedding API"
|
||||
},
|
||||
{
|
||||
"doc": 15,
|
||||
"query": "how to back up everything I stored"
|
||||
},
|
||||
{
|
||||
"doc": 16,
|
||||
"query": "browser blocked cross-origin request before login check"
|
||||
},
|
||||
{
|
||||
"doc": 17,
|
||||
"query": "how did we make the container image smaller"
|
||||
},
|
||||
{
|
||||
"doc": 18,
|
||||
"query": "why are logs emitted as JSON"
|
||||
},
|
||||
{
|
||||
"doc": 19,
|
||||
"query": "keyword matching with stemming"
|
||||
},
|
||||
{
|
||||
"doc": 20,
|
||||
"query": "crash when settings file is absent"
|
||||
},
|
||||
{
|
||||
"doc": 21,
|
||||
"query": "reduce latency spikes from memory allocation"
|
||||
},
|
||||
{
|
||||
"doc": 22,
|
||||
"query": "what happens to requests during deployment shutdown"
|
||||
},
|
||||
{
|
||||
"doc": 23,
|
||||
"query": "stop using package level variables for services"
|
||||
},
|
||||
{
|
||||
"doc": 24,
|
||||
"query": "sign in with a third party identity provider"
|
||||
},
|
||||
{
|
||||
"doc": 25,
|
||||
"query": "test that fails randomly because of timing"
|
||||
},
|
||||
{
|
||||
"doc": 26,
|
||||
"query": "commit takes too long because of checks"
|
||||
},
|
||||
{
|
||||
"doc": 27,
|
||||
"query": "make builds reproducible across machines"
|
||||
},
|
||||
{
|
||||
"doc": 28,
|
||||
"query": "what happens when a callback delivery fails repeatedly"
|
||||
},
|
||||
{
|
||||
"doc": 29,
|
||||
"query": "accented letters not found when searching"
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,157 @@
|
||||
// Command eval-search measures search quality against a labelled corpus.
|
||||
//
|
||||
// eval-search seed <dbdir> writes the corpus into <dbdir>/claude-mnemonic.db
|
||||
// eval-search run <port> queries a running worker and prints MRR@10 and recall@5
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
_ "embed"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"time"
|
||||
|
||||
"github.com/lukaszraczylo/claude-mnemonic/internal/db/gorm"
|
||||
"github.com/lukaszraczylo/claude-mnemonic/pkg/models"
|
||||
)
|
||||
|
||||
const (
|
||||
project = "evalproj"
|
||||
sessionID = "eval-session"
|
||||
dbFile = "claude-mnemonic.db"
|
||||
recallAt = 5
|
||||
mrrAt = 10
|
||||
)
|
||||
|
||||
//go:embed corpus.json
|
||||
var corpusJSON []byte
|
||||
|
||||
type corpus struct {
|
||||
Docs []struct {
|
||||
Type string `json:"type"`
|
||||
Title string `json:"title"`
|
||||
Narrative string `json:"narrative"`
|
||||
} `json:"docs"`
|
||||
Queries []struct {
|
||||
Query string `json:"query"`
|
||||
Doc int `json:"doc"`
|
||||
} `json:"queries"`
|
||||
}
|
||||
|
||||
func main() {
|
||||
if len(os.Args) != 3 {
|
||||
fmt.Fprintln(os.Stderr, "usage: eval-search seed <dbdir> | run <port>")
|
||||
os.Exit(2)
|
||||
}
|
||||
var c corpus
|
||||
if err := json.Unmarshal(corpusJSON, &c); err != nil {
|
||||
fatal(err)
|
||||
}
|
||||
var err error
|
||||
switch os.Args[1] {
|
||||
case "seed":
|
||||
err = seed(c, os.Args[2])
|
||||
case "run":
|
||||
err = run(c, os.Args[2])
|
||||
default:
|
||||
err = fmt.Errorf("unknown mode %q", os.Args[1])
|
||||
}
|
||||
if err != nil {
|
||||
fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
func fatal(err error) {
|
||||
fmt.Fprintln(os.Stderr, "error:", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
|
||||
// seed stores docs in order, so doc i gets observation id i+1.
|
||||
func seed(c corpus, dir string) error {
|
||||
if err := os.MkdirAll(dir, 0o750); err != nil {
|
||||
return err
|
||||
}
|
||||
store, err := gorm.NewStore(gorm.Config{Path: filepath.Join(dir, dbFile), MaxConns: 1})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer func() { _ = store.Close() }()
|
||||
obsStore := gorm.NewObservationStore(store, nil, nil, nil)
|
||||
for i, d := range c.Docs {
|
||||
_, _, err := obsStore.StoreObservation(context.Background(), sessionID, project, &models.ParsedObservation{
|
||||
Type: models.ObservationType(d.Type),
|
||||
Title: d.Title,
|
||||
Narrative: d.Narrative,
|
||||
Concepts: []string{"eval"},
|
||||
}, i+1, 0)
|
||||
if err != nil {
|
||||
return fmt.Errorf("store doc %d: %w", i, err)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
type obsRef struct {
|
||||
ID int64 `json:"id"`
|
||||
}
|
||||
|
||||
func printRanks(query string, want int64, got []obsRef) {
|
||||
pos := -1
|
||||
for i, o := range got {
|
||||
if o.ID == want {
|
||||
pos = i + 1
|
||||
break
|
||||
}
|
||||
}
|
||||
fmt.Printf("rank=%d results=%d %q\n", pos, len(got), query)
|
||||
}
|
||||
|
||||
func run(c corpus, port string) error {
|
||||
client := &http.Client{Timeout: 20 * time.Second}
|
||||
var mrr, recall, results float64
|
||||
for _, q := range c.Queries {
|
||||
u := fmt.Sprintf("http://127.0.0.1:%s/api/context/search?project=%s&query=%s", port, project, url.QueryEscape(q.Query))
|
||||
resp, err := client.Get(u) // #nosec G704 -- dev tool, loopback URL from a numeric port
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
body, err := io.ReadAll(resp.Body)
|
||||
_ = resp.Body.Close()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
var out struct {
|
||||
Observations []obsRef `json:"observations"`
|
||||
}
|
||||
if err := json.Unmarshal(body, &out); err != nil {
|
||||
return fmt.Errorf("decode %q: %w", q.Query, err)
|
||||
}
|
||||
results += float64(len(out.Observations))
|
||||
want := int64(q.Doc + 1)
|
||||
if os.Getenv("EVAL_VERBOSE") != "" {
|
||||
printRanks(q.Query, want, out.Observations)
|
||||
}
|
||||
for rank, o := range out.Observations {
|
||||
if o.ID != want {
|
||||
continue
|
||||
}
|
||||
if rank < mrrAt {
|
||||
mrr += 1 / float64(rank+1)
|
||||
}
|
||||
if rank < recallAt {
|
||||
recall++
|
||||
}
|
||||
break
|
||||
}
|
||||
}
|
||||
n := float64(len(c.Queries))
|
||||
fmt.Printf("queries=%d MRR@%d=%s recall@%d=%s avg_results=%s\n", len(c.Queries), mrrAt, strconv.FormatFloat(mrr/n, 'f', 4, 64),
|
||||
recallAt, strconv.FormatFloat(recall/n, 'f', 4, 64), strconv.FormatFloat(results/n, 'f', 2, 64))
|
||||
return nil
|
||||
}
|
||||
Generated
+660
-355
File diff suppressed because it is too large
Load Diff
+11
-11
@@ -10,19 +10,19 @@
|
||||
"type-check": "vue-tsc --noEmit"
|
||||
},
|
||||
"dependencies": {
|
||||
"vis-data": "^8.0.4",
|
||||
"vis-network": "^10.1.0",
|
||||
"vue": "^3.5.34"
|
||||
"vis-data": "^8.0.5",
|
||||
"vis-network": "^10.1.2",
|
||||
"vue": "^3.5.43"
|
||||
},
|
||||
"devDependencies": {
|
||||
"@fortawesome/fontawesome-free": "^7.2.0",
|
||||
"@tailwindcss/postcss": "^4.3.0",
|
||||
"@types/node": "^25.9.1",
|
||||
"@vitejs/plugin-vue": "^6.0.7",
|
||||
"postcss": "^8.5.15",
|
||||
"tailwindcss": "^4.3.0",
|
||||
"@fortawesome/fontawesome-free": "^7.3.1",
|
||||
"@tailwindcss/postcss": "^4.3.3",
|
||||
"@types/node": "^26.0.0",
|
||||
"@vitejs/plugin-vue": "^6.0.9",
|
||||
"postcss": "^8.5.28",
|
||||
"tailwindcss": "^4.3.3",
|
||||
"typescript": "~6.0.3",
|
||||
"vite": "^8.0.14",
|
||||
"vue-tsc": "^3.3.1"
|
||||
"vite": "^8.3.0",
|
||||
"vue-tsc": "^3.3.11"
|
||||
}
|
||||
}
|
||||
|
||||
+11
-11
@@ -10,19 +10,19 @@
|
||||
"type-check": "vue-tsc --noEmit"
|
||||
},
|
||||
"dependencies": {
|
||||
"vis-data": "^8.0.4",
|
||||
"vis-network": "^10.1.0",
|
||||
"vue": "^3.5.34"
|
||||
"vis-data": "^8.0.5",
|
||||
"vis-network": "^10.1.2",
|
||||
"vue": "^3.5.43"
|
||||
},
|
||||
"devDependencies": {
|
||||
"@fortawesome/fontawesome-free": "^7.2.0",
|
||||
"@tailwindcss/postcss": "^4.3.0",
|
||||
"@types/node": "^25.9.1",
|
||||
"@vitejs/plugin-vue": "^6.0.7",
|
||||
"postcss": "^8.5.15",
|
||||
"tailwindcss": "^4.3.0",
|
||||
"@fortawesome/fontawesome-free": "^7.3.1",
|
||||
"@tailwindcss/postcss": "^4.3.3",
|
||||
"@types/node": "^26.0.0",
|
||||
"@vitejs/plugin-vue": "^6.0.9",
|
||||
"postcss": "^8.5.28",
|
||||
"tailwindcss": "^4.3.3",
|
||||
"typescript": "~6.0.3",
|
||||
"vite": "^8.0.14",
|
||||
"vue-tsc": "^3.3.1"
|
||||
"vite": "^8.3.0",
|
||||
"vue-tsc": "^3.3.11"
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user