diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index de6f3431..4045a313 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -24,7 +24,7 @@ jobs: run: go build -o gortex ./cmd/gortex/ - name: Test - run: go test -race -timeout=20m -coverprofile=coverage.out ./... + run: go test -race -timeout=30m -coverprofile=coverage.out ./... - name: Upload coverage if: matrix.os == 'ubuntu-latest' && matrix.go-version == '1.26' diff --git a/cmd/gortex/testdata/agent-render/codex.txt b/cmd/gortex/testdata/agent-render/codex.txt index 81f0d66a..c6ba5140 100644 --- a/cmd/gortex/testdata/agent-render/codex.txt +++ b/cmd/gortex/testdata/agent-render/codex.txt @@ -5,7 +5,7 @@ direct_only_tool_namespaces = ['mcp__gortex', 'gortex'] [hooks] [[hooks.PostToolUse]] -matcher = '^(Bash|apply_patch)$' +matcher = '^(Bash|apply_patch|(mcp__gortex__|gortex__)(explore|search|read|relations|trace|analyze))$' [[hooks.PostToolUse.hooks]] command = 'gortex hook --agent=codex --mode=enrich' @@ -14,20 +14,11 @@ timeout = 5 type = 'command' [[hooks.PreToolUse]] -matcher = '^Bash$' +matcher = '.*' [[hooks.PreToolUse.hooks]] command = 'gortex hook --agent=codex --mode=enrich' -statusMessage = 'Loading Gortex Bash guidance...' -timeout = 5 -type = 'command' - -[[hooks.PreToolUse]] -matcher = '^mcp__gortex__read$' - -[[hooks.PreToolUse.hooks]] -command = 'gortex hook --agent=codex --mode=enrich' -statusMessage = 'Loading Gortex read guidance...' +statusMessage = 'Loading Gortex tool guidance...' timeout = 5 type = 'command' @@ -40,6 +31,13 @@ statusMessage = 'Loading Gortex graph orientation...' timeout = 5 type = 'command' +[[hooks.Stop]] +[[hooks.Stop.hooks]] +command = 'gortex hook --agent=codex --mode=enrich' +statusMessage = 'Checking Gortex evidence authority...' +timeout = 5 +type = 'command' + [[hooks.UserPromptSubmit]] [[hooks.UserPromptSubmit.hooks]] command = 'gortex hook --agent=codex --mode=enrich' @@ -607,7 +605,7 @@ direct_only_tool_namespaces = ['mcp__gortex', 'gortex'] [hooks] [[hooks.PostToolUse]] -matcher = '^(Bash|apply_patch)$' +matcher = '^(Bash|apply_patch|(mcp__gortex__|gortex__)(explore|search|read|relations|trace|analyze))$' [[hooks.PostToolUse.hooks]] command = 'gortex hook --agent=codex --mode=enrich' @@ -616,20 +614,11 @@ timeout = 5 type = 'command' [[hooks.PreToolUse]] -matcher = '^Bash$' +matcher = '.*' [[hooks.PreToolUse.hooks]] command = 'gortex hook --agent=codex --mode=enrich' -statusMessage = 'Loading Gortex Bash guidance...' -timeout = 5 -type = 'command' - -[[hooks.PreToolUse]] -matcher = '^mcp__gortex__read$' - -[[hooks.PreToolUse.hooks]] -command = 'gortex hook --agent=codex --mode=enrich' -statusMessage = 'Loading Gortex read guidance...' +statusMessage = 'Loading Gortex tool guidance...' timeout = 5 type = 'command' @@ -642,6 +631,13 @@ statusMessage = 'Loading Gortex graph orientation...' timeout = 5 type = 'command' +[[hooks.Stop]] +[[hooks.Stop.hooks]] +command = 'gortex hook --agent=codex --mode=enrich' +statusMessage = 'Checking Gortex evidence authority...' +timeout = 5 +type = 'command' + [[hooks.UserPromptSubmit]] [[hooks.UserPromptSubmit.hooks]] command = 'gortex hook --agent=codex --mode=enrich' diff --git a/docs/agents.md b/docs/agents.md index b406e7cb..4a6e98f1 100644 --- a/docs/agents.md +++ b/docs/agents.md @@ -376,6 +376,7 @@ Current Codex hook coverage: | Surface | Coverage | | ------- | -------- | | `SessionStart` | Matches `startup|resume|clear|compact` and emits graph-tools orientation for new, resumed, cleared, and compacted sessions. | +| Hook-observable local-tool `PreToolUse` terminal gate | Uses one match-all installer-owned hook so an enforceable localization `answer_ready` contract blocks every local tool Codex routes through `PreToolUse`, including Bash, `apply_patch`, MCP tools, and other local function tools. Hosted tools such as `WebSearch`, and specialized paths that opt out of the default hook path, are outside this host boundary. Without a terminal marker it is a local no-op; advisory completion still permits non-navigation contract operations. | | Bash `PreToolUse` | Advises by default; `deny` hard-blocks graph-detectable fallback reads/searches; `rewrite` converts only an unambiguous indexed `cat ` into the exact public `gortex call read` mirror. Compound or ambiguous commands remain advisory. A shell command that would *rewrite* indexed source — `sed -i` / `perl -pi`, a `>` / `>>` redirect, `tee`, or an inline interpreter script opening the path for writing — gets the same redirect Edit and Write get, escalating to a hard block under `GORTEX_HOOK_BLOCK_EDIT`. Writing a path the daemon does not know is a new file and passes through. | | Gortex MCP `read` `PreToolUse` | Advises broad file/editing-context reads by default; `deny` blocks them; `rewrite` preserves the request and adds `options.compress_bodies=true`. Selector-driven reads with no explicit operation are covered. | | Bash `PostToolUse` | Adds graph context for grep/search, source reads, and conservative file-list shapes: `find -name`, `fd`, `ls`, `tree -fi`, and `git ls-files`; bounded `sed`/`awk` reads get file graph context. Execution-capable or ambiguous forms are no-ops. | diff --git a/internal/agents/codex/adapter.go b/internal/agents/codex/adapter.go index 456222d4..8b9d3252 100644 --- a/internal/agents/codex/adapter.go +++ b/internal/agents/codex/adapter.go @@ -62,11 +62,20 @@ const ( v060CodexSessionStartMessage = "IMPORTANT: Prefer Gortex MCP tools (search_symbols, get_callers, get_file_summary, edit_file) over Read/Grep/Glob/Edit." v060CodexSessionStartCommand = "printf '%s\\n' '" + v060CodexSessionStartMessage + "'" v060CodexSessionStartWindowsCommand = "powershell -NoProfile -Command \"Write-Output '" + v060CodexSessionStartMessage + "'\"" - codexPreToolUseMatcher = "^Bash$" - codexMCPReadPreToolUseMatcher = "^mcp__gortex__read$" - codexPostToolUseMatcher = "^(Bash|apply_patch)$" - codexHookTimeoutSeconds = 5 - codexHookModeEnvVar = "GORTEX_CODEX_HOOK_MODE" + // Codex matchers are regular expressions. A match-all PreToolUse hook is + // required because terminal localization covers every local tool routed + // through Codex's hook path; hosted and specialized opt-out tools remain + // outside the host's hook boundary. Without a marker the handler is a + // strict local no-op. + codexPreToolUseMatcher = ".*" + // Retained as migration fingerprints in tests: upsertCodexHookSet removes + // both split predecessors by managed command identity before installing the + // singleton match-all hook. + codexLegacyBashPreToolUseMatcher = "^Bash$" + codexLegacyMCPNavigationPreToolUseMatcher = "^(mcp__gortex__|gortex__)(explore|search|read|relations|trace|analyze)$" + codexPostToolUseMatcher = "^(Bash|apply_patch|(mcp__gortex__|gortex__)(explore|search|read|relations|trace|analyze))$" + codexHookTimeoutSeconds = 5 + codexHookModeEnvVar = "GORTEX_CODEX_HOOK_MODE" // Codex merges its home instructions file into every session ahead of // the repo's own AGENTS.md, preferring the override name when present. codexGlobalInstructionsFile = "AGENTS.md" @@ -550,11 +559,77 @@ func upsertSessionStartHook(root map[string]any, env agents.Env, opts agents.App } func upsertPreToolUseHook(root map[string]any, env agents.Env, opts agents.ApplyOpts) bool { - desired := []map[string]any{ - codexPreToolUseHookEntry(env), - codexMCPReadPreToolUseHookEntry(env), + // Codex permits several handlers in one matcher group. Normalize at handler + // granularity before replacing legacy Gortex matchers: this preserves every + // co-located user handler while collapsing duplicate managed invocations. + splitChanged := splitMixedCodexPreToolUseGroups(root) + desired := []map[string]any{codexPreToolUseHookEntry(env)} + upsertChanged := upsertCodexHookSet(root, "PreToolUse", codexHookEntryIsGortexPreToolUse, desired, opts) + return splitChanged || upsertChanged +} + +func splitMixedCodexPreToolUseGroups(root map[string]any) bool { + hooks, ok := root["hooks"].(map[string]any) + if !ok { + return false + } + entries, ok := codexHookList(hooks["PreToolUse"]) + if !ok { + return false + } + + normalized := make([]any, 0, len(entries)) + changed := false + for _, entry := range entries { + group, ok := entry.(map[string]any) + if !ok { + normalized = append(normalized, entry) + continue + } + handlers, ok := codexHookList(group["hooks"]) + if !ok { + normalized = append(normalized, entry) + continue + } + managed := 0 + kept := make([]any, 0, len(handlers)) + for _, handler := range handlers { + fields, ok := handler.(map[string]any) + if ok { + command, _ := fields["command"].(string) + if codexCommandInvokesCodexHook(command) { + managed++ + continue + } + } + kept = append(kept, handler) + } + switch { + case managed == 0: + normalized = append(normalized, entry) + case managed == 1 && len(handlers) == 1: + // Leave a singleton managed group for upsertCodexHookSet to + // validate or replace. A current singleton remains idempotent. + normalized = append(normalized, entry) + default: + changed = true + if len(kept) == 0 { + continue + } + preserved := make(map[string]any, len(group)) + for key, value := range group { + preserved[key] = value + } + preserved["hooks"] = kept + normalized = append(normalized, preserved) + } + } + if !changed { + return false } - return upsertCodexHookSet(root, "PreToolUse", codexHookEntryIsGortexPreToolUse, desired, opts) + hooks["PreToolUse"] = normalized + root["hooks"] = hooks + return true } func upsertPostToolUseHook(root map[string]any, env agents.Env, opts agents.ApplyOpts) bool { @@ -565,6 +640,10 @@ func upsertUserPromptSubmitHook(root map[string]any, env agents.Env, opts agents return upsertCodexHookSet(root, "UserPromptSubmit", codexHookEntryIsGortexUserPromptSubmit, []map[string]any{codexUserPromptSubmitHookEntry(env)}, opts) } +func upsertStopHook(root map[string]any, env agents.Env, opts agents.ApplyOpts) bool { + return upsertCodexHookSet(root, "Stop", codexHookEntryIsGortexStop, []map[string]any{codexStopHookEntry(env)}, opts) +} + // InstallHooksOnly refreshes the Codex lifecycle hooks in configPath without // touching MCP server entries, AGENTS.md, or any other Codex adapter surface. func InstallHooksOnly(w io.Writer, configPath string, env agents.Env, opts agents.ApplyOpts) (agents.FileAction, error) { @@ -585,7 +664,8 @@ func upsertCodexHooks(root map[string]any, env agents.Env, opts agents.ApplyOpts preChanged := upsertPreToolUseHook(root, env, opts) postChanged := upsertPostToolUseHook(root, env, opts) promptChanged := upsertUserPromptSubmitHook(root, env, opts) - return sessionChanged || preChanged || postChanged || promptChanged + stopChanged := upsertStopHook(root, env, opts) + return sessionChanged || preChanged || postChanged || promptChanged || stopChanged } func upsertCodexHookSet(root map[string]any, event string, isGortex func(any) bool, desired []map[string]any, opts agents.ApplyOpts) bool { @@ -734,6 +814,10 @@ func codexHookEntryIsGortexUserPromptSubmit(entry any) bool { return codexHookEntryInvokesCodexHook(entry) } +func codexHookEntryIsGortexStop(entry any) bool { + return codexHookEntryInvokesCodexHook(entry) +} + func codexHookEntryInvokesCodexHook(entry any) bool { group, ok := entry.(map[string]any) if !ok { @@ -789,21 +873,7 @@ func codexPreToolUseHookEntry(env agents.Env) map[string]any { "type": "command", "command": codexPreToolUseCommand(env), "timeout": codexHookTimeoutSeconds, - "statusMessage": "Loading Gortex Bash guidance...", - }, - }, - } -} - -func codexMCPReadPreToolUseHookEntry(env agents.Env) map[string]any { - return map[string]any{ - "matcher": codexMCPReadPreToolUseMatcher, - "hooks": []any{ - map[string]any{ - "type": "command", - "command": codexPreToolUseCommand(env), - "timeout": codexHookTimeoutSeconds, - "statusMessage": "Loading Gortex read guidance...", + "statusMessage": "Loading Gortex tool guidance...", }, }, } @@ -841,6 +911,22 @@ func codexUserPromptSubmitHookEntry(env agents.Env) map[string]any { } } +// codexStopHookEntry has no matcher: Stop applies to every final response. +// Older Codex hosts that omit last_assistant_message remain fail-open in the +// shared Stop handler. +func codexStopHookEntry(env agents.Env) map[string]any { + return map[string]any{ + "hooks": []any{ + map[string]any{ + "type": "command", + "command": codexHookCommand(env), + "timeout": codexHookTimeoutSeconds, + "statusMessage": "Checking Gortex evidence authority...", + }, + }, + } +} + func codexPreToolUseCommand(env agents.Env) string { return codexHookCommand(env) } diff --git a/internal/agents/codex/adapter_test.go b/internal/agents/codex/adapter_test.go index b952c093..72f57ee6 100644 --- a/internal/agents/codex/adapter_test.go +++ b/internal/agents/codex/adapter_test.go @@ -3,6 +3,7 @@ package codex import ( "os" "path/filepath" + "regexp" "strings" "testing" @@ -359,26 +360,26 @@ func TestCodexInstallsPreToolUseHook(t *testing.T) { cfg := readCodexConfig(t, env) entries := preToolUseEntries(t, cfg) - if len(entries) != 2 { - t.Fatalf("PreToolUse entries=%d want Bash+MCP read entries: %#v", len(entries), entries) + if len(entries) != 1 { + t.Fatalf("PreToolUse entries=%d want one match-all entry: %#v", len(entries), entries) } assertGortexPreToolUseHooks(t, cfg) - bashHandler := requireHookEntry(t, cfg, "PreToolUse", codexPreToolUseMatcher, testCodexHookCommand) - mcpHandler := requireHookEntry(t, cfg, "PreToolUse", codexMCPReadPreToolUseMatcher, testCodexHookCommand) - for name, handler := range map[string]map[string]any{"Bash": bashHandler, "MCP read": mcpHandler} { - if handler["type"] != "command" { - t.Errorf("%s hook type=%v want command", name, handler["type"]) - } - if handler["timeout"] != int64(codexHookTimeoutSeconds) { - t.Errorf("%s timeout=%v want %d", name, handler["timeout"], codexHookTimeoutSeconds) - } + handler := requireHookEntry(t, cfg, "PreToolUse", codexPreToolUseMatcher, testCodexHookCommand) + if handler["type"] != "command" { + t.Errorf("hook type=%v want command", handler["type"]) } - if bashHandler["statusMessage"] != "Loading Gortex Bash guidance..." { - t.Errorf("Bash statusMessage=%v", bashHandler["statusMessage"]) + if handler["timeout"] != int64(codexHookTimeoutSeconds) { + t.Errorf("timeout=%v want %d", handler["timeout"], codexHookTimeoutSeconds) } - if mcpHandler["statusMessage"] != "Loading Gortex read guidance..." { - t.Errorf("MCP read statusMessage=%v", mcpHandler["statusMessage"]) + if handler["statusMessage"] != "Loading Gortex tool guidance..." { + t.Errorf("statusMessage=%v", handler["statusMessage"]) + } + matcher := regexp.MustCompile(codexPreToolUseMatcher) + for _, tool := range []string{"Bash", "apply_patch", "Read", "WebSearch", "mcp__gortex__search", "mcp__gortex__change"} { + if !matcher.MatchString(tool) { + t.Errorf("match-all PreToolUse matcher missed %q", tool) + } } } @@ -450,8 +451,8 @@ func TestCodexHookModeIsOptInAndMigratesInPlace(t *testing.T) { if hasHookCommand(t, cfg, "PreToolUse", "/tmp/test-gortex hook --agent=codex --mode=deny") { t.Fatalf("stale deny hook survived posture migration: %#v", preToolUseEntries(t, cfg)) } - if count := gortexPreToolUseHookCount(t, cfg); count != 2 { - t.Fatalf("posture migration duplicated PreToolUse hooks: %d", count) + if count := gortexPreToolUseHookCount(t, cfg); count != 1 { + t.Fatalf("posture migration duplicated match-all PreToolUse hook: %d", count) } } @@ -748,6 +749,13 @@ command = "` + testCodexHookCommand + `" [[hooks.PreToolUse]] matcher = "^mcp__gortex__(read_file|get_editing_context)$" +[[hooks.PreToolUse.hooks]] +type = "command" +command = "` + testCodexHookCommand + `" + +[[hooks.PreToolUse]] +matcher = "` + codexLegacyMCPNavigationPreToolUseMatcher + `" + [[hooks.PreToolUse.hooks]] type = "command" command = "` + testCodexHookCommand + `" @@ -766,11 +774,16 @@ command = "` + testCodexHookCommand + `" if !hasSessionStartCommand(t, cfg, testCodexHookCommand) { t.Fatalf("managed SessionStart hook missing after static-command migration: %#v", sessionStartEntries(t, cfg)) } - if count := hookMatcherCommandCount(t, cfg, "PreToolUse", codexMCPReadPreToolUseMatcher, testCodexHookCommand); count != 1 { - t.Fatalf("compact read matcher count=%d want 1: %#v", count, preToolUseEntries(t, cfg)) + if count := hookMatcherCommandCount(t, cfg, "PreToolUse", codexPreToolUseMatcher, testCodexHookCommand); count != 1 { + t.Fatalf("match-all matcher count=%d want 1: %#v", count, preToolUseEntries(t, cfg)) } - if count := hookMatcherCommandCount(t, cfg, "PreToolUse", v060CodexMCPReadPreToolUseMatcher, testCodexHookCommand); count != 0 { - t.Fatalf("v0.60.0 read matcher survived upgrade: %#v", preToolUseEntries(t, cfg)) + for _, stale := range []string{codexLegacyBashPreToolUseMatcher, codexLegacyMCPNavigationPreToolUseMatcher, v060CodexMCPReadPreToolUseMatcher} { + if count := hookMatcherCommandCount(t, cfg, "PreToolUse", stale, testCodexHookCommand); count != 0 { + t.Fatalf("stale matcher %q survived upgrade: %#v", stale, preToolUseEntries(t, cfg)) + } + } + if count := gortexPreToolUseHookCount(t, cfg); count != 1 { + t.Fatalf("split hooks were not collapsed: got %d entries %#v", count, preToolUseEntries(t, cfg)) } } @@ -839,8 +852,8 @@ statusMessage = "User PostToolUse" t.Fatalf("Gortex SessionStart hooks=%d want 1", count) } preEntries := preToolUseEntries(t, cfg) - if len(preEntries) != 3 { - t.Fatalf("PreToolUse entries=%d want user+Bash+MCP read entries: %#v", len(preEntries), preEntries) + if len(preEntries) != 2 { + t.Fatalf("PreToolUse entries=%d want user+match-all entries: %#v", len(preEntries), preEntries) } if !hasHookCommand(t, cfg, "PreToolUse", "echo user-pretooluse") { t.Fatalf("user PreToolUse hook was not preserved: %#v", preEntries) @@ -898,8 +911,8 @@ statusMessage = "Old Gortex MCP Read PreToolUse" cfg := readCodexConfig(t, env) preEntries := preToolUseEntries(t, cfg) - if len(preEntries) != 3 { - t.Fatalf("PreToolUse entries=%d want user+Bash+MCP read entries: %#v", len(preEntries), preEntries) + if len(preEntries) != 2 { + t.Fatalf("PreToolUse entries=%d want user+match-all entries: %#v", len(preEntries), preEntries) } if !hasHookCommand(t, cfg, "PreToolUse", "echo user-pretooluse") { t.Fatalf("Force removed user PreToolUse hook: %#v", preEntries) @@ -1038,14 +1051,11 @@ func gortexPreToolUseHookCount(t *testing.T, cfg map[string]any) int { func assertGortexPreToolUseHooks(t *testing.T, cfg map[string]any) { t.Helper() - if count := gortexPreToolUseHookCount(t, cfg); count != 2 { - t.Fatalf("Gortex PreToolUse hooks=%d want Bash+MCP read hooks: %#v", count, preToolUseEntries(t, cfg)) + if count := gortexPreToolUseHookCount(t, cfg); count != 1 { + t.Fatalf("Gortex PreToolUse hooks=%d want one match-all hook: %#v", count, preToolUseEntries(t, cfg)) } if count := hookMatcherCommandCount(t, cfg, "PreToolUse", codexPreToolUseMatcher, testCodexHookCommand); count != 1 { - t.Fatalf("Bash PreToolUse hook count=%d want 1: %#v", count, preToolUseEntries(t, cfg)) - } - if count := hookMatcherCommandCount(t, cfg, "PreToolUse", codexMCPReadPreToolUseMatcher, testCodexHookCommand); count != 1 { - t.Fatalf("MCP read PreToolUse hook count=%d want 1: %#v", count, preToolUseEntries(t, cfg)) + t.Fatalf("match-all PreToolUse hook count=%d want 1: %#v", count, preToolUseEntries(t, cfg)) } } diff --git a/internal/agents/codex/inspect.go b/internal/agents/codex/inspect.go index 9726fed3..8d5429cc 100644 --- a/internal/agents/codex/inspect.go +++ b/internal/agents/codex/inspect.go @@ -35,7 +35,7 @@ type InstallState struct { // HookEvents are the lifecycle events this adapter installs, in the order // they matter for adoption: SessionStart is what states the Gortex rule for // the session, so its absence explains a silent integration on its own. -var HookEvents = []string{"SessionStart", "UserPromptSubmit", "PreToolUse", "PostToolUse"} +var HookEvents = []string{"SessionStart", "UserPromptSubmit", "PreToolUse", "PostToolUse", "Stop"} // TrustRemedy is how a user approves hooks Codex is skipping. const TrustRemedy = "run `/hooks` inside Codex, review the gortex entries, and trust them" @@ -96,6 +96,7 @@ func countGortexHooks(root map[string]any, out map[string]int) { "UserPromptSubmit": codexHookEntryIsGortexUserPromptSubmit, "PreToolUse": codexHookEntryIsGortexPreToolUse, "PostToolUse": codexHookEntryIsGortexPostToolUse, + "Stop": codexHookEntryIsGortexStop, } for event, isGortex := range recognisers { entries, ok := codexHookList(hooks[event]) diff --git a/internal/agents/codex/localization_hook_test.go b/internal/agents/codex/localization_hook_test.go new file mode 100644 index 00000000..8fd1731a --- /dev/null +++ b/internal/agents/codex/localization_hook_test.go @@ -0,0 +1,37 @@ +package codex + +import ( + "regexp" + "testing" + + "github.com/zzet/gortex/internal/agents" +) + +func TestCodexLocalizationLifecycleMatchersCoverEveryFacadeAndNamespace(t *testing.T) { + env := codexGlobalEnv(t) + if _, err := New().Apply(env, agents.ApplyOpts{}); err != nil { + t.Fatalf("apply: %v", err) + } + cfg := readCodexConfig(t, env) + requireHookEntry(t, cfg, "PreToolUse", codexPreToolUseMatcher, testCodexHookCommand) + requireHookEntry(t, cfg, "PostToolUse", codexPostToolUseMatcher, testCodexHookCommand) + + preMatcher := regexp.MustCompile(codexPreToolUseMatcher) + postMatcher := regexp.MustCompile(codexPostToolUseMatcher) + for _, prefix := range []string{"mcp__gortex__", "gortex__"} { + for _, operation := range []string{"explore", "search", "read", "relations", "trace", "analyze"} { + tool := prefix + operation + if !preMatcher.MatchString(tool) { + t.Errorf("PreToolUse matcher does not cover %q", tool) + } + if !postMatcher.MatchString(tool) { + t.Errorf("PostToolUse matcher does not cover %q", tool) + } + } + } + + state := Inspect(env.Home) + if state.Hooks["PreToolUse"] != 1 || state.Hooks["PostToolUse"] != 1 { + t.Fatalf("doctor did not recognize localization hooks: %+v", state.Hooks) + } +} diff --git a/internal/agents/codex/pretooluse_migration_test.go b/internal/agents/codex/pretooluse_migration_test.go new file mode 100644 index 00000000..f8d8fe20 --- /dev/null +++ b/internal/agents/codex/pretooluse_migration_test.go @@ -0,0 +1,151 @@ +package codex + +import ( + "os" + "testing" + + "github.com/zzet/gortex/internal/agents" +) + +func TestCodexPreToolUseMigrationPreservesCoLocatedUserHandlers(t *testing.T) { + managed := ` +[[hooks.PreToolUse.hooks]] +type = "command" +command = "` + testCodexHookCommand + `" +statusMessage = "Old Gortex PreToolUse" +` + user := ` +[[hooks.PreToolUse.hooks]] +type = "command" +command = "echo colocated-user-pretooluse" +statusMessage = "User PreToolUse" +` + for _, tt := range []struct { + name string + handlers string + force bool + }{ + {name: "managed first", handlers: managed + user}, + {name: "user first", handlers: user + managed}, + {name: "force", handlers: managed + user, force: true}, + } { + t.Run(tt.name, func(t *testing.T) { + env := codexGlobalEnv(t) + path := codexConfigPath(env) + seed := `[[hooks.PreToolUse]] +matcher = "^Bash$" +` + tt.handlers + if err := os.WriteFile(path, []byte(seed), 0o644); err != nil { + t.Fatalf("seed config: %v", err) + } + + if _, err := InstallHooksOnly(env.Stderr, path, env, agents.ApplyOpts{Force: tt.force}); err != nil { + t.Fatalf("migrate mixed PreToolUse group: %v", err) + } + cfg := readCodexConfig(t, env) + if count := hookMatcherCommandCount(t, cfg, "PreToolUse", codexLegacyBashPreToolUseMatcher, "echo colocated-user-pretooluse"); count != 1 { + t.Fatalf("co-located user handler count=%d want 1: %#v", count, preToolUseEntries(t, cfg)) + } + if count := hookMatcherCommandCount(t, cfg, "PreToolUse", codexLegacyBashPreToolUseMatcher, testCodexHookCommand); count != 0 { + t.Fatalf("legacy Gortex handler survived mixed-group migration: %#v", preToolUseEntries(t, cfg)) + } + assertGortexPreToolUseHooks(t, cfg) + + if _, err := InstallHooksOnly(env.Stderr, path, env, agents.ApplyOpts{}); err != nil { + t.Fatalf("reapply hooks: %v", err) + } + cfg = readCodexConfig(t, env) + if count := hookMatcherCommandCount(t, cfg, "PreToolUse", codexLegacyBashPreToolUseMatcher, "echo colocated-user-pretooluse"); count != 1 { + t.Fatalf("reapply changed co-located user handler count=%d: %#v", count, preToolUseEntries(t, cfg)) + } + assertGortexPreToolUseHooks(t, cfg) + }) + } +} + +func TestCodexPreToolUseMigrationCollapsesCoLocatedManagedHandlers(t *testing.T) { + const staleCommand = "/tmp/stale-gortex hook --agent=codex --mode=enrich" + current := ` +[[hooks.PreToolUse.hooks]] +type = "command" +command = "` + testCodexHookCommand + `" +statusMessage = "Current Gortex PreToolUse" +` + stale := ` +[[hooks.PreToolUse.hooks]] +type = "command" +command = "` + staleCommand + `" +statusMessage = "Stale Gortex PreToolUse" +` + for _, tt := range []struct { + name string + handlers string + force bool + }{ + {name: "duplicate current", handlers: current + current}, + {name: "current then stale", handlers: current + stale}, + {name: "stale then current", handlers: stale + current}, + {name: "force duplicate current", handlers: current + current, force: true}, + } { + t.Run(tt.name, func(t *testing.T) { + env := codexGlobalEnv(t) + path := codexConfigPath(env) + seed := `[[hooks.PreToolUse]] +matcher = ".*" +` + tt.handlers + if err := os.WriteFile(path, []byte(seed), 0o644); err != nil { + t.Fatalf("seed config: %v", err) + } + + assertNormalized := func(stage string, cfg map[string]any) { + t.Helper() + if count := codexPreToolUseManagedHandlerCount(t, cfg); count != 1 { + t.Fatalf("%s managed handler invocations=%d want 1: %#v", stage, count, preToolUseEntries(t, cfg)) + } + if count := hookMatcherCommandCount(t, cfg, "PreToolUse", codexPreToolUseMatcher, testCodexHookCommand); count != 1 { + t.Fatalf("%s current match-all handler count=%d want 1: %#v", stage, count, preToolUseEntries(t, cfg)) + } + if hasHookCommand(t, cfg, "PreToolUse", staleCommand) { + t.Fatalf("%s stale managed handler survived: %#v", stage, preToolUseEntries(t, cfg)) + } + assertGortexPreToolUseHooks(t, cfg) + } + + if _, err := InstallHooksOnly(env.Stderr, path, env, agents.ApplyOpts{Force: tt.force}); err != nil { + t.Fatalf("collapse co-located managed handlers: %v", err) + } + assertNormalized("initial migration", readCodexConfig(t, env)) + + if _, err := InstallHooksOnly(env.Stderr, path, env, agents.ApplyOpts{}); err != nil { + t.Fatalf("reapply hooks: %v", err) + } + assertNormalized("idempotent reapply", readCodexConfig(t, env)) + }) + } +} + +func codexPreToolUseManagedHandlerCount(t *testing.T, cfg map[string]any) int { + t.Helper() + count := 0 + for _, entry := range preToolUseEntries(t, cfg) { + group, ok := entry.(map[string]any) + if !ok { + continue + } + handlers, ok := codexHookList(group["hooks"]) + if !ok { + continue + } + for _, handler := range handlers { + fields, ok := handler.(map[string]any) + if !ok { + continue + } + command, _ := fields["command"].(string) + if codexCommandInvokesCodexHook(command) { + count++ + } + } + } + return count +} diff --git a/internal/agents/codex/remove.go b/internal/agents/codex/remove.go index 0d4717bc..0621040f 100644 --- a/internal/agents/codex/remove.go +++ b/internal/agents/codex/remove.go @@ -238,6 +238,7 @@ func removeCodexHooks(root map[string]any) bool { "UserPromptSubmit": codexHookEntryIsGortexUserPromptSubmit, "PreToolUse": codexHookEntryIsGortexPreToolUse, "PostToolUse": codexHookEntryIsGortexPostToolUse, + "Stop": codexHookEntryIsGortexStop, } changed := false for event, isGortex := range recognisers { diff --git a/internal/agents/codex/stop_hook_test.go b/internal/agents/codex/stop_hook_test.go new file mode 100644 index 00000000..5d2a4ad6 --- /dev/null +++ b/internal/agents/codex/stop_hook_test.go @@ -0,0 +1,60 @@ +package codex + +import ( + "path/filepath" + "testing" + + "github.com/zzet/gortex/internal/agents" +) + +func TestCodexInstallsAndInspectsStopHook(t *testing.T) { + env := codexGlobalEnv(t) + env.InstallGlobalInstructions = true + if _, err := New().Apply(env, agents.ApplyOpts{}); err != nil { + t.Fatalf("apply: %v", err) + } + assertCodexStopHook(t, readCodexConfig(t, env)) + + state := Inspect(env.Home) + if got := state.Hooks["Stop"]; got != 1 { + t.Fatalf("doctor counted Stop hooks=%d want 1: %+v", got, state.Hooks) + } +} + +func TestCodexInstallHooksOnlyIncludesStop(t *testing.T) { + env := codexGlobalEnv(t) + path := filepath.Join(env.Home, ".codex", "config.toml") + if _, err := InstallHooksOnly(env.Stderr, path, env, agents.ApplyOpts{}); err != nil { + t.Fatalf("install hooks only: %v", err) + } + assertCodexStopHook(t, readCodexConfig(t, env)) +} + +func assertCodexStopHook(t *testing.T, cfg map[string]any) { + t.Helper() + entries := hookEntries(t, cfg, "Stop") + if len(entries) != 1 { + t.Fatalf("Stop entries=%d want 1: %#v", len(entries), entries) + } + entry, ok := entries[0].(map[string]any) + if !ok { + t.Fatalf("Stop entry has unexpected shape: %#v", entries[0]) + } + if _, hasMatcher := entry["matcher"]; hasMatcher { + t.Fatalf("Stop entry should carry no matcher: %#v", entry) + } + if !codexHookEntryIsGortexStop(entry) { + t.Fatalf("installed Stop entry is not recognized as Gortex: %#v", entry) + } + handlers, ok := codexHookList(entry["hooks"]) + if !ok || len(handlers) != 1 { + t.Fatalf("Stop handlers=%#v", entry["hooks"]) + } + handler := handlers[0].(map[string]any) + if handler["command"] != testCodexHookCommand { + t.Errorf("command=%v want %q", handler["command"], testCodexHookCommand) + } + if handler["timeout"] != int64(codexHookTimeoutSeconds) { + t.Errorf("timeout=%v want %d", handler["timeout"], codexHookTimeoutSeconds) + } +} diff --git a/internal/daemon/overlay.go b/internal/daemon/overlay.go index 22776c9d..f7e2bce0 100644 --- a/internal/daemon/overlay.go +++ b/internal/daemon/overlay.go @@ -2,6 +2,7 @@ package daemon import ( "errors" + "fmt" "os" "regexp" "sort" @@ -161,6 +162,19 @@ var ErrSessionNotFound = errors.New("overlay session not found") // surface as wrong-line errors that look like graph bugs. var ErrOverlayDrift = errors.New("overlay base SHA mismatch — re-read and resubmit") +// ErrOverlaySnapshotTooLarge reports that a bounded overlay snapshot could not +// be materialized within its caller-supplied file or byte envelope. Resource is +// either "files" or "bytes" and Limit is the corresponding hard cap. Bounded +// snapshot APIs return no workspace or file slice with this error. +type ErrOverlaySnapshotTooLarge struct { + Resource string + Limit int +} + +func (e *ErrOverlaySnapshotTooLarge) Error() string { + return fmt.Sprintf("overlay snapshot too large: %s exceed limit %d", e.Resource, e.Limit) +} + // ErrBranchNotFound is returned by branch tools when the named // branch does not exist on the session. The MCP surface translates // this into a structured tool error with the branch name embedded. @@ -369,30 +383,112 @@ func (m *OverlayManager) SnapshotFor(sessionID string) (workspace string, files if m == nil { return "", nil, ErrSessionNotFound } - // Promoted to write lock so we can refresh LastUsed alongside - // the snapshot copy: every tool-call view-build flows through - // here, and that activity must reset the idle timer. Without - // this, a session that only queries (no further Push) would - // trip the TTL while in active use. The cost is one extra - // mutex promotion per overlay-active tool call — negligible - // against the parse work the view builder is about to do. + // Ordinary MCP sessions are not overlay sessions. Prove absence under a + // shared lock so their request preparation never queues for the manager-wide + // writer. A Register racing after this check belongs to the next request; + // this read-lock instant is the current request's empty snapshot boundary. + m.mu.RLock() + _, exists := m.sessions[sessionID] + m.mu.RUnlock() + if !exists { + return "", nil, ErrSessionNotFound + } + + // Registered sessions still count reads as activity. Re-check after lock + // promotion because Drop may win the gap, then copy the active branch while + // protected. Sorting happens after unlock: the returned slice owns its + // headers and no longer needs to serialize unrelated sessions. m.mu.Lock() - defer m.mu.Unlock() sess, ok := m.sessions[sessionID] if !ok { + m.mu.Unlock() return "", nil, ErrSessionNotFound } sess.LastUsed = time.Now() + workspace = sess.WorkspaceID br := sess.activeBranch() if br == nil { - return sess.WorkspaceID, nil, nil + m.mu.Unlock() + return workspace, nil, nil } out := make([]OverlayFile, 0, len(br.files)) for _, f := range br.files { out = append(out, f) } + m.mu.Unlock() + + sort.Slice(out, func(i, j int) bool { return out[i].Path < out[j].Path }) + return workspace, out, nil +} + +// SnapshotForBounded is SnapshotFor with a hard pre-materialization envelope. +// Both limits must be positive. File count and Path+Content+BaseSHA bytes are +// checked while holding the manager lock and before allocating the result +// slice. LastUsed is bumped even when the envelope rejects the snapshot. +func (m *OverlayManager) SnapshotForBounded( + sessionID string, + maxFiles, maxBytes int, +) (workspace string, files []OverlayFile, err error) { + if maxFiles <= 0 || maxBytes <= 0 { + return "", nil, fmt.Errorf("overlay snapshot limits must be positive: files=%d bytes=%d", maxFiles, maxBytes) + } + if m == nil { + return "", nil, ErrSessionNotFound + } + + m.mu.RLock() + _, exists := m.sessions[sessionID] + m.mu.RUnlock() + if !exists { + return "", nil, ErrSessionNotFound + } + + m.mu.Lock() + sess, ok := m.sessions[sessionID] + if !ok { + m.mu.Unlock() + return "", nil, ErrSessionNotFound + } + sess.LastUsed = time.Now() + workspace = sess.WorkspaceID + br := sess.activeBranch() + if br == nil { + m.mu.Unlock() + return workspace, nil, nil + } + out, snapshotErr := copyOverlayFilesBounded(br.files, maxFiles, maxBytes) + m.mu.Unlock() + if snapshotErr != nil { + return "", nil, snapshotErr + } + sort.Slice(out, func(i, j int) bool { return out[i].Path < out[j].Path }) - return sess.WorkspaceID, out, nil + return workspace, out, nil +} + +// copyOverlayFilesBounded validates a locked branch map before allocating or +// copying its snapshot. Callers must hold the manager lock for the full call. +func copyOverlayFilesBounded( + files map[string]OverlayFile, + maxFiles, maxBytes int, +) ([]OverlayFile, error) { + if len(files) > maxFiles { + return nil, &ErrOverlaySnapshotTooLarge{Resource: "files", Limit: maxFiles} + } + totalBytes := 0 + for _, file := range files { + for _, size := range [...]int{len(file.Path), len(file.Content), len(file.BaseSHA)} { + if size > maxBytes-totalBytes { + return nil, &ErrOverlaySnapshotTooLarge{Resource: "bytes", Limit: maxBytes} + } + totalBytes += size + } + } + out := make([]OverlayFile, 0, len(files)) + for _, file := range files { + out = append(out, file) + } + return out, nil } // Push attaches one overlay file to a session's active branch. @@ -970,3 +1066,45 @@ func (m *OverlayManager) FilesForBranch(sessionID, branchName string) ([]Overlay sort.Slice(out, func(i, j int) bool { return out[i].Path < out[j].Path }) return out, nil } + +// FilesForBranchBounded snapshots a named branch under the same bounded +// materialization contract as SnapshotForBounded. +func (m *OverlayManager) FilesForBranchBounded( + sessionID, branchName string, + maxFiles, maxBytes int, +) ([]OverlayFile, error) { + if maxFiles <= 0 || maxBytes <= 0 { + return nil, fmt.Errorf("overlay snapshot limits must be positive: files=%d bytes=%d", maxFiles, maxBytes) + } + if m == nil { + return nil, ErrSessionNotFound + } + + m.mu.RLock() + _, exists := m.sessions[sessionID] + m.mu.RUnlock() + if !exists { + return nil, ErrSessionNotFound + } + + m.mu.Lock() + sess, ok := m.sessions[sessionID] + if !ok { + m.mu.Unlock() + return nil, ErrSessionNotFound + } + br, ok := sess.branches[branchName] + if !ok { + m.mu.Unlock() + return nil, ErrBranchNotFound + } + sess.LastUsed = time.Now() + out, snapshotErr := copyOverlayFilesBounded(br.files, maxFiles, maxBytes) + m.mu.Unlock() + if snapshotErr != nil { + return nil, snapshotErr + } + + sort.Slice(out, func(i, j int) bool { return out[i].Path < out[j].Path }) + return out, nil +} diff --git a/internal/daemon/overlay_bounded_test.go b/internal/daemon/overlay_bounded_test.go new file mode 100644 index 00000000..db509f8c --- /dev/null +++ b/internal/daemon/overlay_bounded_test.go @@ -0,0 +1,253 @@ +package daemon + +import ( + "fmt" + "strings" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +const ( + testOverlaySnapshotMaxFiles = 257 + testOverlaySnapshotMaxBytes = 4_259_840 +) + +func TestOverlayManagerSnapshotForBoundedFileBoundary(t *testing.T) { + m := NewOverlayManager(time.Minute) + id := m.Register("ws") + for i := 0; i < testOverlaySnapshotMaxFiles; i++ { + require.NoError(t, m.Push(id, OverlayFile{ + Path: fmt.Sprintf("file-%03d.go", i), + Content: "package p", + }, nil)) + } + + workspace, files, err := m.SnapshotForBounded( + id, testOverlaySnapshotMaxFiles, testOverlaySnapshotMaxBytes, + ) + require.NoError(t, err) + require.Equal(t, "ws", workspace) + require.Len(t, files, testOverlaySnapshotMaxFiles) + require.Equal(t, "file-000.go", files[0].Path) + require.Equal(t, "file-256.go", files[len(files)-1].Path) + + require.NoError(t, m.Push(id, OverlayFile{Path: "file-257.go", Content: "package p"}, nil)) + workspace, files, err = m.SnapshotForBounded( + id, testOverlaySnapshotMaxFiles, testOverlaySnapshotMaxBytes, + ) + var tooLarge *ErrOverlaySnapshotTooLarge + require.ErrorAs(t, err, &tooLarge) + require.Equal(t, "files", tooLarge.Resource) + require.Equal(t, testOverlaySnapshotMaxFiles, tooLarge.Limit) + require.Empty(t, workspace, "rejected snapshots must not return partial metadata") + require.Nil(t, files, "rejected snapshots must not return a partial slice") +} + +func TestOverlayManagerSnapshotForBoundedCountsAllStringFields(t *testing.T) { + t.Run("exact combined boundary", func(t *testing.T) { + m := NewOverlayManager(time.Minute) + id := m.Register("ws") + require.NoError(t, m.Push(id, OverlayFile{ + Path: "path", + Content: "abc", + BaseSHA: "sha", + }, nil)) + + _, files, err := m.SnapshotForBounded(id, 1, 10) + require.NoError(t, err) + require.Len(t, files, 1) + + require.NoError(t, m.Push(id, OverlayFile{ + Path: "path", + Content: "abcd", + BaseSHA: "sha", + }, nil)) + workspace, files, err := m.SnapshotForBounded(id, 1, 10) + var tooLarge *ErrOverlaySnapshotTooLarge + require.ErrorAs(t, err, &tooLarge) + require.Equal(t, "bytes", tooLarge.Resource) + require.Empty(t, workspace) + require.Nil(t, files) + }) + + t.Run("path exact and plus one", func(t *testing.T) { + m := NewOverlayManager(time.Minute) + id := m.Register("ws") + require.NoError(t, m.Push(id, OverlayFile{Path: "1234"}, nil)) + _, files, err := m.SnapshotForBounded(id, 1, 4) + require.NoError(t, err) + require.Len(t, files, 1) + + require.NoError(t, m.Push(id, OverlayFile{Path: "12345"}, nil)) + _, files, err = m.SnapshotForBounded(id, 2, 4) + var tooLarge *ErrOverlaySnapshotTooLarge + require.ErrorAs(t, err, &tooLarge) + require.Equal(t, "bytes", tooLarge.Resource) + require.Nil(t, files) + }) + + t.Run("base sha exact and plus one", func(t *testing.T) { + m := NewOverlayManager(time.Minute) + id := m.Register("ws") + require.NoError(t, m.Push(id, OverlayFile{Path: "p", BaseSHA: "1234"}, nil)) + _, files, err := m.SnapshotForBounded(id, 1, 5) + require.NoError(t, err) + require.Len(t, files, 1) + + require.NoError(t, m.Push(id, OverlayFile{Path: "p", BaseSHA: "12345"}, nil)) + _, files, err = m.SnapshotForBounded(id, 1, 5) + var tooLarge *ErrOverlaySnapshotTooLarge + require.ErrorAs(t, err, &tooLarge) + require.Equal(t, "bytes", tooLarge.Resource) + require.Nil(t, files) + }) +} + +func TestOverlayManagerBoundedSnapshotsRequirePositiveLimits(t *testing.T) { + m := NewOverlayManager(time.Minute) + id := m.Register("ws") + + for _, limits := range [][2]int{{0, 1}, {1, 0}, {-1, 1}, {1, -1}} { + _, files, err := m.SnapshotForBounded(id, limits[0], limits[1]) + require.Error(t, err) + require.Nil(t, files) + + files, err = m.FilesForBranchBounded(id, MainBranchName, limits[0], limits[1]) + require.Error(t, err) + require.Nil(t, files) + } +} + +func TestOverlayManagerBoundedRejectionBumpsLastUsed(t *testing.T) { + m := NewOverlayManager(time.Minute) + id := m.Register("ws") + require.NoError(t, m.Push(id, OverlayFile{Path: "a.go", Content: "a"}, nil)) + + old := time.Unix(1, 0) + m.mu.Lock() + m.sessions[id].LastUsed = old + m.mu.Unlock() + + _, _, err := m.SnapshotForBounded(id, 1, 1) + var tooLarge *ErrOverlaySnapshotTooLarge + require.ErrorAs(t, err, &tooLarge) + status, statusErr := m.StatusFor(id) + require.NoError(t, statusErr) + require.True(t, status.LastUsed.After(old)) + + m.mu.Lock() + m.sessions[id].LastUsed = old + m.mu.Unlock() + files, err := m.FilesForBranchBounded(id, MainBranchName, 1, 1) + require.ErrorAs(t, err, &tooLarge) + require.Nil(t, files) + status, statusErr = m.StatusFor(id) + require.NoError(t, statusErr) + require.True(t, status.LastUsed.After(old)) +} + +func TestOverlayManagerBoundedMissingUsesSharedLock(t *testing.T) { + m := NewOverlayManager(time.Minute) + m.mu.RLock() + defer m.mu.RUnlock() + + tests := map[string]func() error{ + "active": func() error { + _, _, err := m.SnapshotForBounded("ordinary-mcp-session", 1, 1) + return err + }, + "named branch": func() error { + _, err := m.FilesForBranchBounded("ordinary-mcp-session", MainBranchName, 1, 1) + return err + }, + } + for name, call := range tests { + done := make(chan error, 1) + go func() { done <- call() }() + select { + case err := <-done: + require.ErrorIs(t, err, ErrSessionNotFound, name) + case <-time.After(time.Second): + t.Fatalf("%s snapshot waited for an exclusive manager lock", name) + } + } +} + +func TestOverlayManagerSnapshotForBoundedConcurrentCoherent(t *testing.T) { + m := NewOverlayManager(time.Minute) + id := m.Register("ws") + require.NoError(t, m.Push(id, OverlayFile{Path: "state.go", Content: "main"}, nil)) + _, err := m.Fork(id, ForkOptions{Name: "alternate"}) + require.NoError(t, err) + require.NoError(t, m.PushToBranch(id, "alternate", OverlayFile{ + Path: "state.go", Content: "alternate", + }, nil)) + + const iterations = 500 + start := make(chan struct{}) + errs := make(chan error, 2) + var wg sync.WaitGroup + wg.Add(2) + go func() { + defer wg.Done() + <-start + for i := 0; i < iterations; i++ { + branch := MainBranchName + if i%2 == 0 { + branch = "alternate" + } + if switchErr := m.SwitchBranch(id, branch); switchErr != nil { + errs <- switchErr + return + } + } + }() + go func() { + defer wg.Done() + <-start + for i := 0; i < iterations; i++ { + workspace, files, snapshotErr := m.SnapshotForBounded(id, 1, 64) + if snapshotErr != nil { + errs <- snapshotErr + return + } + if workspace != "ws" || len(files) != 1 || files[0].Path != "state.go" || + (files[0].Content != "main" && files[0].Content != "alternate") { + errs <- fmt.Errorf("incoherent snapshot: workspace=%q files=%+v", workspace, files) + return + } + } + }() + close(start) + wg.Wait() + close(errs) + for err := range errs { + require.NoError(t, err) + } +} + +func TestOverlayManagerFilesForBranchBoundedSortsAndDoesNotAlias(t *testing.T) { + m := NewOverlayManager(time.Minute) + id := m.Register("ws") + _, err := m.Fork(id, ForkOptions{Name: "candidate"}) + require.NoError(t, err) + require.NoError(t, m.PushToBranch(id, "candidate", OverlayFile{Path: "b.go", Content: "b"}, nil)) + require.NoError(t, m.PushToBranch(id, "candidate", OverlayFile{Path: "a.go", Content: "a"}, nil)) + + files, err := m.FilesForBranchBounded(id, "candidate", 2, 10) + require.NoError(t, err) + require.Equal(t, []string{"a.go", "b.go"}, []string{files[0].Path, files[1].Path}) + files[0].Content = strings.Repeat("x", 3) + + again, err := m.FilesForBranchBounded(id, "candidate", 2, 10) + require.NoError(t, err) + require.Equal(t, "a", again[0].Content) + + rejected, err := m.FilesForBranchBounded(id, "candidate", 1, 10) + var tooLarge *ErrOverlaySnapshotTooLarge + require.ErrorAs(t, err, &tooLarge) + require.Nil(t, rejected) +} diff --git a/internal/daemon/overlay_test.go b/internal/daemon/overlay_test.go index bb0b6ca5..05ddd41d 100644 --- a/internal/daemon/overlay_test.go +++ b/internal/daemon/overlay_test.go @@ -72,6 +72,25 @@ func TestOverlayManager_HasAndFileCount(t *testing.T) { require.Zero(t, m.FileCount(id)) } +func TestOverlayManager_SnapshotForMissingUsesSharedLock(t *testing.T) { + m := NewOverlayManager(time.Minute) + m.mu.RLock() + defer m.mu.RUnlock() + + done := make(chan error, 1) + go func() { + _, _, err := m.SnapshotFor("ordinary-mcp-session") + done <- err + }() + + select { + case err := <-done: + require.ErrorIs(t, err, ErrSessionNotFound) + case <-time.After(time.Second): + t.Fatal("missing-session snapshot waited for an exclusive manager lock") + } +} + // TestOverlayManager_DriftCheck verifies that Push surfaces a drift // error when the supplied BaseSHA disagrees with the on-disk SHA // reported by the callback. Without drift detection two clients diff --git a/internal/graph/bounded_adjacency.go b/internal/graph/bounded_adjacency.go new file mode 100644 index 00000000..8ccd8e85 --- /dev/null +++ b/internal/graph/bounded_adjacency.go @@ -0,0 +1,354 @@ +package graph + +import ( + "context" + "sort" +) + +const ( + // MaxBoundedAdjacencyKeys bounds both endpoint and source-site batches. + MaxBoundedAdjacencyKeys = 256 + // MaxBoundedAdjacencyKinds bounds the raw requested kind set. + MaxBoundedAdjacencyKinds = 16 + // MaxBoundedAdjacencyRowsPerKey bounds complete identities retained for one key. + MaxBoundedAdjacencyRowsPerKey = 256 + // MaxBoundedAdjacencyTotalIdentities bounds matching identities inspected + // across one projection, including each limit+1 truncation sentinel. + MaxBoundedAdjacencyTotalIdentities = 16_384 + // MaxBoundedAdjacencyInspectedEdges bounds raw adjacency slots per physical + // reader layer. An OverlaidView makes at most one bounded base call and one + // separately capped local scan; matching base and local identities are + // recharged against MaxBoundedAdjacencyTotalIdentities before return. + MaxBoundedAdjacencyInspectedEdges = 16_384 +) + +// EdgeSourceSite identifies one source line. FilePath deliberately is not part of the +// key: source-literal attribution historically keys call sites by graph source +// identity and line, while each returned EdgeIdentity retains its provenance. +type EdgeSourceSite struct { + From string + Line int +} + +// BoundedEdgeIdentityProjection is a metadata-free adjacency projection keyed +// by source (outgoing) or target (incoming) ID. A truncated key has no retained +// identities; callers must fail closed for that key. +type BoundedEdgeIdentityProjection struct { + ByEndpoint map[string][]EdgeIdentity + Truncated map[string]bool +} + +// BoundedSiteEdgeIdentityProjection is the exact-source-site sibling. +type BoundedSiteEdgeIdentityProjection struct { + BySite map[EdgeSourceSite][]EdgeIdentity + Truncated map[EdgeSourceSite]bool +} + +// BoundedOutgoingEdgeIdentityReader is an optional metadata-free projection. +type BoundedOutgoingEdgeIdentityReader interface { + FindOutgoingEdgeIdentitiesBounded(context.Context, []string, []EdgeKind, int) (BoundedEdgeIdentityProjection, error) +} + +// BoundedIncomingEdgeIdentityReader is the inbound sibling. +type BoundedIncomingEdgeIdentityReader interface { + FindIncomingEdgeIdentitiesBounded(context.Context, []string, []EdgeKind, int) (BoundedEdgeIdentityProjection, error) +} + +// BoundedOutgoingSiteEdgeIdentityReader projects exact source-line adjacency. +type BoundedOutgoingSiteEdgeIdentityReader interface { + FindOutgoingSiteEdgeIdentitiesBounded(context.Context, []EdgeSourceSite, []EdgeKind, int) (BoundedSiteEdgeIdentityProjection, error) +} + +var ( + _ BoundedOutgoingEdgeIdentityReader = (*Graph)(nil) + _ BoundedIncomingEdgeIdentityReader = (*Graph)(nil) + _ BoundedOutgoingSiteEdgeIdentityReader = (*Graph)(nil) +) + +func emptyBoundedEdgeIdentityProjection() BoundedEdgeIdentityProjection { + return BoundedEdgeIdentityProjection{ + ByEndpoint: make(map[string][]EdgeIdentity), + Truncated: make(map[string]bool), + } +} + +func emptyBoundedSiteEdgeIdentityProjection() BoundedSiteEdgeIdentityProjection { + return BoundedSiteEdgeIdentityProjection{ + BySite: make(map[EdgeSourceSite][]EdgeIdentity), + Truncated: make(map[EdgeSourceSite]bool), + } +} + +func validateBoundedAdjacencyLimit(limit int) error { + if limit < 1 || limit > MaxBoundedAdjacencyRowsPerKey { + return &BoundedLocalizationLimitError{ + Resource: "adjacency rows per key", + Limit: MaxBoundedAdjacencyRowsPerKey, + } + } + return nil +} + +func canonicalBoundedAdjacencyKinds(ctx context.Context, kinds []EdgeKind) (map[EdgeKind]struct{}, error) { + if len(kinds) > MaxBoundedAdjacencyKinds { + return nil, &BoundedLocalizationLimitError{ + Resource: "adjacency kinds", + Limit: MaxBoundedAdjacencyKinds, + } + } + out := make(map[EdgeKind]struct{}, len(kinds)) + for index, kind := range kinds { + if index&127 == 0 { + if err := ctx.Err(); err != nil { + return nil, err + } + } + if kind != "" { + out[kind] = struct{}{} + } + } + return out, nil +} + +func canonicalBoundedAdjacencyEndpoints(ctx context.Context, ids []string) ([]string, error) { + if len(ids) > MaxBoundedAdjacencyKeys { + return nil, &BoundedLocalizationLimitError{ + Resource: "adjacency endpoint keys", + Limit: MaxBoundedAdjacencyKeys, + } + } + seen := make(map[string]struct{}, len(ids)) + out := make([]string, 0, len(ids)) + for index, id := range ids { + if index&127 == 0 { + if err := ctx.Err(); err != nil { + return nil, err + } + } + if id == "" { + continue + } + if _, duplicate := seen[id]; duplicate { + continue + } + seen[id] = struct{}{} + out = append(out, id) + } + sort.Strings(out) + return out, nil +} + +func canonicalBoundedAdjacencySites(ctx context.Context, sites []EdgeSourceSite) ([]EdgeSourceSite, error) { + if len(sites) > MaxBoundedAdjacencyKeys { + return nil, &BoundedLocalizationLimitError{ + Resource: "adjacency source-site keys", + Limit: MaxBoundedAdjacencyKeys, + } + } + seen := make(map[EdgeSourceSite]struct{}, len(sites)) + out := make([]EdgeSourceSite, 0, len(sites)) + for index, site := range sites { + if index&127 == 0 { + if err := ctx.Err(); err != nil { + return nil, err + } + } + if site.From == "" { + continue + } + if _, duplicate := seen[site]; duplicate { + continue + } + seen[site] = struct{}{} + out = append(out, site) + } + sort.Slice(out, func(i, j int) bool { + if out[i].From != out[j].From { + return out[i].From < out[j].From + } + return out[i].Line < out[j].Line + }) + return out, nil +} + +func sortEdgeIdentities(identities []EdgeIdentity) { + sort.Slice(identities, func(i, j int) bool { + if identities[i].From != identities[j].From { + return identities[i].From < identities[j].From + } + if identities[i].To != identities[j].To { + return identities[i].To < identities[j].To + } + if identities[i].Kind != identities[j].Kind { + return identities[i].Kind < identities[j].Kind + } + if identities[i].FilePath != identities[j].FilePath { + return identities[i].FilePath < identities[j].FilePath + } + return identities[i].Line < identities[j].Line + }) +} + +type boundedAdjacencyBudget struct { + inspected int + relevant int +} + +func (budget *boundedAdjacencyBudget) inspect() error { + if budget.inspected == MaxBoundedAdjacencyInspectedEdges { + return &BoundedLocalizationLimitError{ + Resource: "adjacency inspected edge rows", + Limit: MaxBoundedAdjacencyInspectedEdges, + } + } + budget.inspected++ + return nil +} + +func (budget *boundedAdjacencyBudget) retain() error { + if budget.relevant == MaxBoundedAdjacencyTotalIdentities { + return &BoundedLocalizationLimitError{ + Resource: "adjacency matching identities", + Limit: MaxBoundedAdjacencyTotalIdentities, + } + } + budget.relevant++ + return nil +} + +func scanBoundedEdgeIdentities( + ctx context.Context, + edges []*Edge, + kinds map[EdgeKind]struct{}, + limit int, + site *EdgeSourceSite, + budget *boundedAdjacencyBudget, + accept func(EdgeIdentity) bool, +) ([]EdgeIdentity, bool, error) { + var identities []EdgeIdentity + var seen map[EdgeIdentity]struct{} + for index, edge := range edges { + if index&127 == 0 { + if err := ctx.Err(); err != nil { + return nil, false, err + } + } + if err := budget.inspect(); err != nil { + return nil, false, err + } + if edge == nil { + continue + } + if site != nil && (edge.From != site.From || edge.Line != site.Line) { + continue + } + if _, requested := kinds[edge.Kind]; !requested { + continue + } + identity := EdgeIdentityFor(edge) + if accept != nil && !accept(identity) { + continue + } + if seen == nil { + identities = make([]EdgeIdentity, 0, 1) + seen = make(map[EdgeIdentity]struct{}, 1) + } + if _, duplicate := seen[identity]; duplicate { + continue + } + seen[identity] = struct{}{} + if err := budget.retain(); err != nil { + return nil, false, err + } + if len(identities) == limit { + return nil, true, nil + } + identities = append(identities, identity) + } + sortEdgeIdentities(identities) + return identities, false, nil +} + +type boundedSiteScanState struct { + site EdgeSourceSite + identities []EdgeIdentity + seen map[EdgeIdentity]struct{} + truncated bool +} + +// scanBoundedSiteEdgeIdentities scans one source adjacency exactly once and +// dispatches matching rows to every requested line for that source. Raw work is +// therefore charged per physical adjacency slot, not once per requested site. +func scanBoundedSiteEdgeIdentities( + ctx context.Context, + edges []*Edge, + sites []EdgeSourceSite, + kinds map[EdgeKind]struct{}, + limit int, + budget *boundedAdjacencyBudget, + accept func(EdgeIdentity) bool, +) (map[EdgeSourceSite][]EdgeIdentity, map[EdgeSourceSite]bool, error) { + states := make(map[int]*boundedSiteScanState, len(sites)) + for _, site := range sites { + states[site.Line] = &boundedSiteScanState{site: site} + } + active := len(states) + for index, edge := range edges { + if active == 0 { + break + } + if index&127 == 0 { + if err := ctx.Err(); err != nil { + return nil, nil, err + } + } + if err := budget.inspect(); err != nil { + return nil, nil, err + } + if edge == nil { + continue + } + state := states[edge.Line] + if state == nil || state.truncated { + continue + } + if _, requested := kinds[edge.Kind]; !requested { + continue + } + identity := EdgeIdentityFor(edge) + if accept != nil && !accept(identity) { + continue + } + if state.seen == nil { + state.identities = make([]EdgeIdentity, 0, 1) + state.seen = make(map[EdgeIdentity]struct{}, 1) + } + if _, duplicate := state.seen[identity]; duplicate { + continue + } + state.seen[identity] = struct{}{} + if err := budget.retain(); err != nil { + return nil, nil, err + } + if len(state.identities) == limit { + state.identities = nil + state.truncated = true + active-- + continue + } + state.identities = append(state.identities, identity) + } + bySite := make(map[EdgeSourceSite][]EdgeIdentity) + truncated := make(map[EdgeSourceSite]bool) + for _, state := range states { + if state.truncated { + truncated[state.site] = true + continue + } + if len(state.identities) > 0 { + sortEdgeIdentities(state.identities) + bySite[state.site] = state.identities + } + } + return bySite, truncated, nil +} diff --git a/internal/graph/bounded_adjacency_graph.go b/internal/graph/bounded_adjacency_graph.go new file mode 100644 index 00000000..a109bc8e --- /dev/null +++ b/internal/graph/bounded_adjacency_graph.go @@ -0,0 +1,172 @@ +package graph + +import "context" + +// FindOutgoingEdgeIdentitiesBounded projects metadata-free outgoing adjacency. +// Requested kinds are filtered before the per-source limit. Raw adjacency and +// matching identities share request-wide hard budgets; any budget or context +// failure discards the whole projection. +func (g *Graph) FindOutgoingEdgeIdentitiesBounded( + ctx context.Context, + sourceIDs []string, + kinds []EdgeKind, + limit int, +) (BoundedEdgeIdentityProjection, error) { + if ctx == nil { + ctx = context.Background() + } + if err := ctx.Err(); err != nil { + return BoundedEdgeIdentityProjection{}, err + } + if err := validateBoundedAdjacencyLimit(limit); err != nil { + return BoundedEdgeIdentityProjection{}, err + } + kindSet, err := canonicalBoundedAdjacencyKinds(ctx, kinds) + if err != nil { + return BoundedEdgeIdentityProjection{}, err + } + ids, err := canonicalBoundedAdjacencyEndpoints(ctx, sourceIDs) + if err != nil { + return BoundedEdgeIdentityProjection{}, err + } + projection := emptyBoundedEdgeIdentityProjection() + if g == nil || len(ids) == 0 || len(kindSet) == 0 { + return projection, nil + } + budget := &boundedAdjacencyBudget{} + for _, sourceID := range ids { + if err := ctx.Err(); err != nil { + return BoundedEdgeIdentityProjection{}, err + } + shard := g.shardFor(sourceID) + shard.mu.RLock() + identities, truncated, scanErr := scanBoundedEdgeIdentities( + ctx, shard.outEdges[sourceID], kindSet, limit, nil, budget, nil, + ) + shard.mu.RUnlock() + if scanErr != nil { + return BoundedEdgeIdentityProjection{}, scanErr + } + if truncated { + projection.Truncated[sourceID] = true + continue + } + if len(identities) > 0 { + projection.ByEndpoint[sourceID] = identities + } + } + return projection, nil +} + +// FindIncomingEdgeIdentitiesBounded is the inbound sibling, keyed by target. +func (g *Graph) FindIncomingEdgeIdentitiesBounded( + ctx context.Context, + targetIDs []string, + kinds []EdgeKind, + limit int, +) (BoundedEdgeIdentityProjection, error) { + if ctx == nil { + ctx = context.Background() + } + if err := ctx.Err(); err != nil { + return BoundedEdgeIdentityProjection{}, err + } + if err := validateBoundedAdjacencyLimit(limit); err != nil { + return BoundedEdgeIdentityProjection{}, err + } + kindSet, err := canonicalBoundedAdjacencyKinds(ctx, kinds) + if err != nil { + return BoundedEdgeIdentityProjection{}, err + } + ids, err := canonicalBoundedAdjacencyEndpoints(ctx, targetIDs) + if err != nil { + return BoundedEdgeIdentityProjection{}, err + } + projection := emptyBoundedEdgeIdentityProjection() + if g == nil || len(ids) == 0 || len(kindSet) == 0 { + return projection, nil + } + budget := &boundedAdjacencyBudget{} + for _, targetID := range ids { + if err := ctx.Err(); err != nil { + return BoundedEdgeIdentityProjection{}, err + } + shard := g.shardFor(targetID) + shard.mu.RLock() + identities, truncated, scanErr := scanBoundedEdgeIdentities( + ctx, shard.inEdges[targetID], kindSet, limit, nil, budget, nil, + ) + shard.mu.RUnlock() + if scanErr != nil { + return BoundedEdgeIdentityProjection{}, scanErr + } + if truncated { + projection.Truncated[targetID] = true + continue + } + if len(identities) > 0 { + projection.ByEndpoint[targetID] = identities + } + } + return projection, nil +} + +// FindOutgoingSiteEdgeIdentitiesBounded projects outgoing identities at exact +// {source,line} sites. Edge FilePath remains output provenance and is never a +// site predicate. +func (g *Graph) FindOutgoingSiteEdgeIdentitiesBounded( + ctx context.Context, + sites []EdgeSourceSite, + kinds []EdgeKind, + limit int, +) (BoundedSiteEdgeIdentityProjection, error) { + if ctx == nil { + ctx = context.Background() + } + if err := ctx.Err(); err != nil { + return BoundedSiteEdgeIdentityProjection{}, err + } + if err := validateBoundedAdjacencyLimit(limit); err != nil { + return BoundedSiteEdgeIdentityProjection{}, err + } + kindSet, err := canonicalBoundedAdjacencyKinds(ctx, kinds) + if err != nil { + return BoundedSiteEdgeIdentityProjection{}, err + } + canonical, err := canonicalBoundedAdjacencySites(ctx, sites) + if err != nil { + return BoundedSiteEdgeIdentityProjection{}, err + } + projection := emptyBoundedSiteEdgeIdentityProjection() + if g == nil || len(canonical) == 0 || len(kindSet) == 0 { + return projection, nil + } + budget := &boundedAdjacencyBudget{} + for start := 0; start < len(canonical); { + if err := ctx.Err(); err != nil { + return BoundedSiteEdgeIdentityProjection{}, err + } + end := start + 1 + for end < len(canonical) && canonical[end].From == canonical[start].From { + end++ + } + from := canonical[start].From + shard := g.shardFor(from) + shard.mu.RLock() + bySite, truncated, scanErr := scanBoundedSiteEdgeIdentities( + ctx, shard.outEdges[from], canonical[start:end], kindSet, limit, budget, nil, + ) + shard.mu.RUnlock() + if scanErr != nil { + return BoundedSiteEdgeIdentityProjection{}, scanErr + } + for site, identities := range bySite { + projection.BySite[site] = identities + } + for site := range truncated { + projection.Truncated[site] = true + } + start = end + } + return projection, nil +} diff --git a/internal/graph/bounded_adjacency_overlay.go b/internal/graph/bounded_adjacency_overlay.go new file mode 100644 index 00000000..8d781700 --- /dev/null +++ b/internal/graph/bounded_adjacency_overlay.go @@ -0,0 +1,316 @@ +package graph + +import "context" + +var ( + _ BoundedOutgoingEdgeIdentityReader = (*OverlaidView)(nil) + _ BoundedIncomingEdgeIdentityReader = (*OverlaidView)(nil) +) + +func overlayAdjacencyCompensatedLimit(layer *OverlayLayer, limit int) int { + if layer == nil { + return limit + } + // One shadow identity may hide arbitrarily many location-distinct edges. + // The only honest bounded compensation is the full per-key envelope; the + // caller-visible limit is re-applied after overlay filtering. + return MaxBoundedAdjacencyRowsPerKey +} + +func (v *OverlaidView) overlayOwnsIdentity(id string) bool { + return v != nil && v.layer != nil && + (v.nodeBelongsToOverlay(id) || v.layer.ownsNodeIdentity(id)) +} + +func (v *OverlaidView) overlayTargetVisible(id string) bool { + return !v.overlayOwnsIdentity(id) || v.layer.nodeByID[id] != nil +} + +func appendBoundedBaseIdentities( + ctx context.Context, + identities []EdgeIdentity, + limit int, + budget *boundedAdjacencyBudget, + accept func(EdgeIdentity) bool, +) ([]EdgeIdentity, bool, error) { + var out []EdgeIdentity + seen := make(map[EdgeIdentity]struct{}) + for index, identity := range identities { + if index&127 == 0 { + if err := ctx.Err(); err != nil { + return nil, false, err + } + } + if accept != nil && !accept(identity) { + continue + } + if _, duplicate := seen[identity]; duplicate { + continue + } + seen[identity] = struct{}{} + if err := budget.retain(); err != nil { + return nil, false, err + } + if len(out) == limit { + return nil, true, nil + } + out = append(out, identity) + } + sortEdgeIdentities(out) + return out, false, nil +} + +func mergeBoundedIdentitySlices(left, right []EdgeIdentity, limit int) ([]EdgeIdentity, bool) { + out := make([]EdgeIdentity, 0, len(left)+len(right)) + seen := make(map[EdgeIdentity]struct{}, len(left)+len(right)) + for _, identities := range [][]EdgeIdentity{left, right} { + for _, identity := range identities { + if _, duplicate := seen[identity]; duplicate { + continue + } + seen[identity] = struct{}{} + if len(out) == limit { + return nil, true + } + out = append(out, identity) + } + } + sortEdgeIdentities(out) + return out, false +} + +// FindOutgoingEdgeIdentitiesBounded preserves GetOutEdges replacement and +// tombstone semantics without materializing Edge payloads. Overlay-owned +// sources use only current overlay adjacency; durable rows targeting removed +// overlay identities are filtered before the caller-visible cap. +func (v *OverlaidView) FindOutgoingEdgeIdentitiesBounded( + ctx context.Context, + sourceIDs []string, + kinds []EdgeKind, + limit int, +) (BoundedEdgeIdentityProjection, error) { + if ctx == nil { + ctx = context.Background() + } + if err := ctx.Err(); err != nil { + return BoundedEdgeIdentityProjection{}, err + } + if err := validateBoundedAdjacencyLimit(limit); err != nil { + return BoundedEdgeIdentityProjection{}, err + } + kindSet, err := canonicalBoundedAdjacencyKinds(ctx, kinds) + if err != nil { + return BoundedEdgeIdentityProjection{}, err + } + ids, err := canonicalBoundedAdjacencyEndpoints(ctx, sourceIDs) + if err != nil { + return BoundedEdgeIdentityProjection{}, err + } + projection := emptyBoundedEdgeIdentityProjection() + if v == nil || len(ids) == 0 || len(kindSet) == 0 { + return projection, nil + } + if v.layer == nil { + if v.base == nil { + return projection, nil + } + bounded, ok := v.base.(BoundedOutgoingEdgeIdentityReader) + if !ok { + return BoundedEdgeIdentityProjection{}, ErrBoundedLocalizationUnavailable + } + return bounded.FindOutgoingEdgeIdentitiesBounded(ctx, ids, kinds, limit) + } + + baseIDs := make([]string, 0, len(ids)) + overlayIDs := make(map[string]bool, len(ids)) + for _, id := range ids { + if v.overlayOwnsIdentity(id) { + overlayIDs[id] = v.layer.nodeByID[id] != nil + continue + } + baseIDs = append(baseIDs, id) + } + baseProjection := emptyBoundedEdgeIdentityProjection() + if len(baseIDs) > 0 && v.base != nil { + bounded, ok := v.base.(BoundedOutgoingEdgeIdentityReader) + if !ok { + return BoundedEdgeIdentityProjection{}, ErrBoundedLocalizationUnavailable + } + baseProjection, err = bounded.FindOutgoingEdgeIdentitiesBounded( + ctx, baseIDs, kinds, overlayAdjacencyCompensatedLimit(v.layer, limit), + ) + if err != nil { + return BoundedEdgeIdentityProjection{}, err + } + } + + budget := &boundedAdjacencyBudget{} + for _, sourceID := range ids { + if err := ctx.Err(); err != nil { + return BoundedEdgeIdentityProjection{}, err + } + if current, owned := overlayIDs[sourceID]; owned { + if !current { + continue + } + identities, truncated, scanErr := scanBoundedEdgeIdentities( + ctx, v.layer.outEdges[sourceID], kindSet, limit, nil, budget, + func(identity EdgeIdentity) bool { return v.overlayTargetVisible(identity.To) }, + ) + if scanErr != nil { + return BoundedEdgeIdentityProjection{}, scanErr + } + if truncated { + projection.Truncated[sourceID] = true + continue + } + if len(identities) > 0 { + projection.ByEndpoint[sourceID] = identities + } + continue + } + if baseProjection.Truncated[sourceID] { + projection.Truncated[sourceID] = true + continue + } + identities, truncated, mergeErr := appendBoundedBaseIdentities( + ctx, baseProjection.ByEndpoint[sourceID], limit, budget, + func(identity EdgeIdentity) bool { + return identity.From == sourceID && v.overlayTargetVisible(identity.To) && kindRequested(kindSet, identity.Kind) + }, + ) + if mergeErr != nil { + return BoundedEdgeIdentityProjection{}, mergeErr + } + if truncated { + projection.Truncated[sourceID] = true + continue + } + if len(identities) > 0 { + projection.ByEndpoint[sourceID] = identities + } + } + return projection, nil +} + +// FindIncomingEdgeIdentitiesBounded preserves GetInEdges semantics: durable +// edges from overlay-owned sources are replaced, then current overlay edges are +// merged. A removed target yields an empty complete key; same-ID replacement +// retains durable callers from untouched sources. +func (v *OverlaidView) FindIncomingEdgeIdentitiesBounded( + ctx context.Context, + targetIDs []string, + kinds []EdgeKind, + limit int, +) (BoundedEdgeIdentityProjection, error) { + if ctx == nil { + ctx = context.Background() + } + if err := ctx.Err(); err != nil { + return BoundedEdgeIdentityProjection{}, err + } + if err := validateBoundedAdjacencyLimit(limit); err != nil { + return BoundedEdgeIdentityProjection{}, err + } + kindSet, err := canonicalBoundedAdjacencyKinds(ctx, kinds) + if err != nil { + return BoundedEdgeIdentityProjection{}, err + } + ids, err := canonicalBoundedAdjacencyEndpoints(ctx, targetIDs) + if err != nil { + return BoundedEdgeIdentityProjection{}, err + } + projection := emptyBoundedEdgeIdentityProjection() + if v == nil || len(ids) == 0 || len(kindSet) == 0 { + return projection, nil + } + if v.layer == nil { + if v.base == nil { + return projection, nil + } + bounded, ok := v.base.(BoundedIncomingEdgeIdentityReader) + if !ok { + return BoundedEdgeIdentityProjection{}, ErrBoundedLocalizationUnavailable + } + return bounded.FindIncomingEdgeIdentitiesBounded(ctx, ids, kinds, limit) + } + + baseIDs := make([]string, 0, len(ids)) + removedTargets := make(map[string]bool) + for _, id := range ids { + if v.overlayOwnsIdentity(id) && v.layer.nodeByID[id] == nil { + removedTargets[id] = true + continue + } + baseIDs = append(baseIDs, id) + } + baseProjection := emptyBoundedEdgeIdentityProjection() + if len(baseIDs) > 0 && v.base != nil { + bounded, ok := v.base.(BoundedIncomingEdgeIdentityReader) + if !ok { + return BoundedEdgeIdentityProjection{}, ErrBoundedLocalizationUnavailable + } + baseProjection, err = bounded.FindIncomingEdgeIdentitiesBounded( + ctx, baseIDs, kinds, overlayAdjacencyCompensatedLimit(v.layer, limit), + ) + if err != nil { + return BoundedEdgeIdentityProjection{}, err + } + } + + budget := &boundedAdjacencyBudget{} + for _, targetID := range ids { + if err := ctx.Err(); err != nil { + return BoundedEdgeIdentityProjection{}, err + } + if removedTargets[targetID] { + continue + } + if baseProjection.Truncated[targetID] { + projection.Truncated[targetID] = true + continue + } + baseIdentities, baseTruncated, mergeErr := appendBoundedBaseIdentities( + ctx, baseProjection.ByEndpoint[targetID], limit, budget, + func(identity EdgeIdentity) bool { + return identity.To == targetID && !v.overlayOwnsIdentity(identity.From) && + v.overlayTargetVisible(identity.To) && kindRequested(kindSet, identity.Kind) + }, + ) + if mergeErr != nil { + return BoundedEdgeIdentityProjection{}, mergeErr + } + if baseTruncated { + projection.Truncated[targetID] = true + continue + } + overlayIdentities, overlayTruncated, scanErr := scanBoundedEdgeIdentities( + ctx, v.layer.inEdges[targetID], kindSet, limit, nil, budget, + func(identity EdgeIdentity) bool { + return identity.To == targetID && v.overlayTargetVisible(identity.To) && + v.overlayOwnsIdentity(identity.From) && v.layer.nodeByID[identity.From] != nil + }, + ) + if scanErr != nil { + return BoundedEdgeIdentityProjection{}, scanErr + } + if overlayTruncated { + projection.Truncated[targetID] = true + continue + } + identities, truncated := mergeBoundedIdentitySlices(baseIdentities, overlayIdentities, limit) + if truncated { + projection.Truncated[targetID] = true + continue + } + if len(identities) > 0 { + projection.ByEndpoint[targetID] = identities + } + } + return projection, nil +} + +func kindRequested(kinds map[EdgeKind]struct{}, kind EdgeKind) bool { + _, ok := kinds[kind] + return ok +} diff --git a/internal/graph/bounded_adjacency_overlay_site.go b/internal/graph/bounded_adjacency_overlay_site.go new file mode 100644 index 00000000..00029942 --- /dev/null +++ b/internal/graph/bounded_adjacency_overlay_site.go @@ -0,0 +1,126 @@ +package graph + +import "context" + +var _ BoundedOutgoingSiteEdgeIdentityReader = (*OverlaidView)(nil) + +// FindOutgoingSiteEdgeIdentitiesBounded is the exact {source,line} overlay +// projection. Source ownership is decided once per source, so multiple sites +// from the same overlay file share one raw adjacency scan. +func (v *OverlaidView) FindOutgoingSiteEdgeIdentitiesBounded( + ctx context.Context, + sites []EdgeSourceSite, + kinds []EdgeKind, + limit int, +) (BoundedSiteEdgeIdentityProjection, error) { + if ctx == nil { + ctx = context.Background() + } + if err := ctx.Err(); err != nil { + return BoundedSiteEdgeIdentityProjection{}, err + } + if err := validateBoundedAdjacencyLimit(limit); err != nil { + return BoundedSiteEdgeIdentityProjection{}, err + } + kindSet, err := canonicalBoundedAdjacencyKinds(ctx, kinds) + if err != nil { + return BoundedSiteEdgeIdentityProjection{}, err + } + canonical, err := canonicalBoundedAdjacencySites(ctx, sites) + if err != nil { + return BoundedSiteEdgeIdentityProjection{}, err + } + projection := emptyBoundedSiteEdgeIdentityProjection() + if v == nil || len(canonical) == 0 || len(kindSet) == 0 { + return projection, nil + } + if v.layer == nil { + if v.base == nil { + return projection, nil + } + bounded, ok := v.base.(BoundedOutgoingSiteEdgeIdentityReader) + if !ok { + return BoundedSiteEdgeIdentityProjection{}, ErrBoundedLocalizationUnavailable + } + return bounded.FindOutgoingSiteEdgeIdentitiesBounded(ctx, canonical, kinds, limit) + } + + baseSites := make([]EdgeSourceSite, 0, len(canonical)) + overlayCurrent := make(map[string]bool) + for _, site := range canonical { + if v.overlayOwnsIdentity(site.From) { + overlayCurrent[site.From] = v.layer.nodeByID[site.From] != nil + continue + } + baseSites = append(baseSites, site) + } + baseProjection := emptyBoundedSiteEdgeIdentityProjection() + if len(baseSites) > 0 && v.base != nil { + bounded, ok := v.base.(BoundedOutgoingSiteEdgeIdentityReader) + if !ok { + return BoundedSiteEdgeIdentityProjection{}, ErrBoundedLocalizationUnavailable + } + baseProjection, err = bounded.FindOutgoingSiteEdgeIdentitiesBounded( + ctx, baseSites, kinds, overlayAdjacencyCompensatedLimit(v.layer, limit), + ) + if err != nil { + return BoundedSiteEdgeIdentityProjection{}, err + } + } + + budget := &boundedAdjacencyBudget{} + for start := 0; start < len(canonical); { + if err := ctx.Err(); err != nil { + return BoundedSiteEdgeIdentityProjection{}, err + } + end := start + 1 + for end < len(canonical) && canonical[end].From == canonical[start].From { + end++ + } + from := canonical[start].From + if current, owned := overlayCurrent[from]; owned { + if current { + bySite, truncated, scanErr := scanBoundedSiteEdgeIdentities( + ctx, v.layer.outEdges[from], canonical[start:end], kindSet, limit, budget, + func(identity EdgeIdentity) bool { return v.overlayTargetVisible(identity.To) }, + ) + if scanErr != nil { + return BoundedSiteEdgeIdentityProjection{}, scanErr + } + for site, identities := range bySite { + projection.BySite[site] = identities + } + for site := range truncated { + projection.Truncated[site] = true + } + } + start = end + continue + } + for _, site := range canonical[start:end] { + if baseProjection.Truncated[site] { + projection.Truncated[site] = true + continue + } + identities, truncated, mergeErr := appendBoundedBaseIdentities( + ctx, baseProjection.BySite[site], limit, budget, + func(identity EdgeIdentity) bool { + return identity.From == site.From && identity.Line == site.Line && + v.overlayTargetVisible(identity.To) && kindRequested(kindSet, identity.Kind) + }, + ) + if mergeErr != nil { + return BoundedSiteEdgeIdentityProjection{}, mergeErr + } + if truncated { + projection.Truncated[site] = true + continue + } + if len(identities) > 0 { + projection.BySite[site] = identities + } + } + start = end + } + return projection, nil +} diff --git a/internal/graph/bounded_adjacency_overlay_test.go b/internal/graph/bounded_adjacency_overlay_test.go new file mode 100644 index 00000000..8e00f82a --- /dev/null +++ b/internal/graph/bounded_adjacency_overlay_test.go @@ -0,0 +1,414 @@ +package graph + +import ( + "context" + "errors" + "fmt" + "reflect" + "testing" +) + +func requireSingleBoundedIdentity(t *testing.T, identities []EdgeIdentity, want EdgeIdentity) { + t.Helper() + if !reflect.DeepEqual(identities, []EdgeIdentity{want}) { + t.Fatalf("identities = %#v, want %#v", identities, []EdgeIdentity{want}) + } +} + +func TestOverlaidViewBoundedAdjacencyReplacementAndTombstoneParity(t *testing.T) { + const ( + sourceFile = "repo/source.go" + targetFile = "repo/target.go" + sourceID = sourceFile + "::source" + targetID = targetFile + "::target" + otherID = "repo/other.go::caller" + ) + baseEdge := &Edge{From: sourceID, To: targetID, Kind: EdgeCalls, FilePath: sourceFile, Line: 10} + otherEdge := &Edge{From: otherID, To: targetID, Kind: EdgeCalls, FilePath: "repo/other.go", Line: 20} + base := New() + base.AddEdge(baseEdge) + base.AddEdge(otherEdge) + + t.Run("source tombstone suppresses base and stray overlay edges", func(t *testing.T) { + layer := NewOverlayLayer() + layer.MarkFile(sourceFile, false) + layer.MarkRemoved("source", sourceID) + layer.AddEdge(&Edge{From: sourceID, To: targetID, Kind: EdgeCalls, Line: 99}) + view := NewOverlaidView(base, layer) + out, err := view.FindOutgoingEdgeIdentitiesBounded(context.Background(), []string{sourceID}, []EdgeKind{EdgeCalls}, 1) + if err != nil || len(out.ByEndpoint) != 0 || len(out.Truncated) != 0 { + t.Fatalf("source tombstone outgoing = %#v, %v", out, err) + } + in, err := view.FindIncomingEdgeIdentitiesBounded(context.Background(), []string{targetID}, []EdgeKind{EdgeCalls}, 1) + if err != nil { + t.Fatalf("source tombstone incoming: %v", err) + } + requireSingleBoundedIdentity(t, in.ByEndpoint[targetID], EdgeIdentityFor(otherEdge)) + }) + + t.Run("same ID source replacement uses only overlay adjacency", func(t *testing.T) { + layer := NewOverlayLayer() + layer.MarkFile(sourceFile, false) + layer.MarkRemoved("source", sourceID) + layer.AddNode(sourceFile, &Node{ID: sourceID, Name: "source", Kind: KindMethod, FilePath: sourceFile}) + replacementEdge := &Edge{From: sourceID, To: targetID, Kind: EdgeCalls, FilePath: sourceFile, Line: 30} + layer.AddEdge(replacementEdge) + view := NewOverlaidView(base, layer) + out, err := view.FindOutgoingEdgeIdentitiesBounded(context.Background(), []string{sourceID}, []EdgeKind{EdgeCalls}, 1) + if err != nil { + t.Fatalf("source replacement outgoing: %v", err) + } + requireSingleBoundedIdentity(t, out.ByEndpoint[sourceID], EdgeIdentityFor(replacementEdge)) + in, err := view.FindIncomingEdgeIdentitiesBounded(context.Background(), []string{targetID}, []EdgeKind{EdgeCalls}, 2) + if err != nil { + t.Fatalf("source replacement incoming: %v", err) + } + want := []EdgeIdentity{EdgeIdentityFor(replacementEdge), EdgeIdentityFor(otherEdge)} + sortEdgeIdentities(want) + if !reflect.DeepEqual(in.ByEndpoint[targetID], want) { + t.Fatalf("source replacement incoming = %#v, want %#v", in, want) + } + }) + + t.Run("target tombstone and replacement", func(t *testing.T) { + tombstone := NewOverlayLayer() + tombstone.MarkFile(targetFile, false) + tombstone.MarkRemoved("target", targetID) + view := NewOverlaidView(base, tombstone) + out, err := view.FindOutgoingEdgeIdentitiesBounded(context.Background(), []string{sourceID}, []EdgeKind{EdgeCalls}, 1) + if err != nil || len(out.ByEndpoint) != 0 { + t.Fatalf("target tombstone outgoing = %#v, %v", out, err) + } + in, err := view.FindIncomingEdgeIdentitiesBounded(context.Background(), []string{targetID}, []EdgeKind{EdgeCalls}, 2) + if err != nil || len(in.ByEndpoint) != 0 { + t.Fatalf("target tombstone incoming = %#v, %v", in, err) + } + + replacement := NewOverlayLayer() + replacement.MarkFile(targetFile, false) + replacement.MarkRemoved("target", targetID) + replacement.AddNode(targetFile, &Node{ID: targetID, Name: "target", Kind: KindType, FilePath: targetFile}) + view = NewOverlaidView(base, replacement) + out, err = view.FindOutgoingEdgeIdentitiesBounded(context.Background(), []string{sourceID}, []EdgeKind{EdgeCalls}, 1) + if err != nil { + t.Fatalf("target replacement outgoing: %v", err) + } + requireSingleBoundedIdentity(t, out.ByEndpoint[sourceID], EdgeIdentityFor(baseEdge)) + in, err = view.FindIncomingEdgeIdentitiesBounded(context.Background(), []string{targetID}, []EdgeKind{EdgeCalls}, 2) + if err != nil || len(in.ByEndpoint[targetID]) != 2 { + t.Fatalf("target replacement incoming = %#v, %v", in, err) + } + }) +} + +func TestOverlaidViewBoundedIncomingRejectsStrayAndTombstonedOverlaySources(t *testing.T) { + const target = "repo/target.go::target" + for _, test := range []struct { + name string + source string + configure func(*OverlayLayer, string) + wantSource bool + }{ + {name: "untouched stray", source: "repo/untouched.go::source", configure: func(layer *OverlayLayer, source string) {}}, + {name: "standard tombstone", source: "repo/source.go::source", configure: func(layer *OverlayLayer, source string) { + layer.MarkFile("repo/source.go", false) + layer.MarkRemoved("source", source) + }}, + {name: "standard current", source: "repo/source.go::source", configure: func(layer *OverlayLayer, source string) { + layer.MarkFile("repo/source.go", false) + layer.AddNode("repo/source.go", &Node{ID: source, Name: "source", Kind: KindMethod, FilePath: "repo/source.go"}) + }, wantSource: true}, + {name: "detached tombstone", source: "legacy-source", configure: func(layer *OverlayLayer, source string) { + layer.MarkRemoved("legacy-source", source) + }}, + {name: "detached current", source: "legacy-source", configure: func(layer *OverlayLayer, source string) { + layer.MarkRemoved("legacy-source", source) + layer.AddNode("overlay.go", &Node{ID: source, Name: "legacy-source", Kind: KindMethod, FilePath: "overlay.go"}) + }, wantSource: true}, + } { + t.Run(test.name, func(t *testing.T) { + layer := NewOverlayLayer() + test.configure(layer, test.source) + edge := &Edge{From: test.source, To: target, Kind: EdgeCalls, FilePath: "overlay.go", Line: 5} + layer.AddEdge(edge) + projection, err := NewOverlaidView(New(), layer).FindIncomingEdgeIdentitiesBounded( + context.Background(), []string{target}, []EdgeKind{EdgeCalls}, 1, + ) + if err != nil { + t.Fatalf("incoming projection: %v", err) + } + if test.wantSource { + requireSingleBoundedIdentity(t, projection.ByEndpoint[target], EdgeIdentityFor(edge)) + } else if len(projection.ByEndpoint) != 0 || len(projection.Truncated) != 0 { + t.Fatalf("stray/tombstoned source authenticated: %#v", projection) + } + }) + } +} + +func TestOverlaidViewBoundedAdjacencyDetachedReplacementAndTombstoneParity(t *testing.T) { + const ( + source = "legacy-source" + target = "legacy-target" + ) + baseEdge := &Edge{From: source, To: target, Kind: EdgeCalls, FilePath: "legacy.go", Line: 1} + base := New() + base.AddEdge(baseEdge) + + sourceTombstone := NewOverlayLayer() + sourceTombstone.MarkRemoved("legacy-source", source) + sourceTombstone.AddEdge(&Edge{From: source, To: target, Kind: EdgeCalls, Line: 99}) + out, err := NewOverlaidView(base, sourceTombstone).FindOutgoingEdgeIdentitiesBounded( + context.Background(), []string{source}, []EdgeKind{EdgeCalls}, 1, + ) + if err != nil || len(out.ByEndpoint) != 0 || len(out.Truncated) != 0 { + t.Fatalf("detached source tombstone = %#v, %v", out, err) + } + + sourceReplacement := NewOverlayLayer() + sourceReplacement.MarkRemoved("legacy-source", source) + sourceReplacement.AddNode("overlay.go", &Node{ID: source, Name: "legacy-source", Kind: KindMethod, FilePath: "overlay.go"}) + replacementEdge := &Edge{From: source, To: target, Kind: EdgeCalls, FilePath: "overlay.go", Line: 2} + sourceReplacement.AddEdge(replacementEdge) + out, err = NewOverlaidView(base, sourceReplacement).FindOutgoingEdgeIdentitiesBounded( + context.Background(), []string{source}, []EdgeKind{EdgeCalls}, 1, + ) + if err != nil { + t.Fatalf("detached source replacement: %v", err) + } + requireSingleBoundedIdentity(t, out.ByEndpoint[source], EdgeIdentityFor(replacementEdge)) + + targetTombstone := NewOverlayLayer() + targetTombstone.MarkRemoved("legacy-target", target) + view := NewOverlaidView(base, targetTombstone) + out, err = view.FindOutgoingEdgeIdentitiesBounded(context.Background(), []string{source}, []EdgeKind{EdgeCalls}, 1) + if err != nil || len(out.ByEndpoint) != 0 { + t.Fatalf("detached target tombstone outgoing = %#v, %v", out, err) + } + in, err := view.FindIncomingEdgeIdentitiesBounded(context.Background(), []string{target}, []EdgeKind{EdgeCalls}, 1) + if err != nil || len(in.ByEndpoint) != 0 || len(in.Truncated) != 0 { + t.Fatalf("detached target tombstone incoming = %#v, %v", in, err) + } + + targetReplacement := NewOverlayLayer() + targetReplacement.MarkRemoved("legacy-target", target) + targetReplacement.AddNode("overlay.go", &Node{ID: target, Name: "legacy-target", Kind: KindType, FilePath: "overlay.go"}) + view = NewOverlaidView(base, targetReplacement) + out, err = view.FindOutgoingEdgeIdentitiesBounded(context.Background(), []string{source}, []EdgeKind{EdgeCalls}, 1) + if err != nil { + t.Fatalf("detached target replacement outgoing: %v", err) + } + requireSingleBoundedIdentity(t, out.ByEndpoint[source], EdgeIdentityFor(baseEdge)) + in, err = view.FindIncomingEdgeIdentitiesBounded(context.Background(), []string{target}, []EdgeKind{EdgeCalls}, 1) + if err != nil { + t.Fatalf("detached target replacement incoming: %v", err) + } + requireSingleBoundedIdentity(t, in.ByEndpoint[target], EdgeIdentityFor(baseEdge)) +} + +func TestOverlaidViewBoundedSiteReplacementAndTombstoneParity(t *testing.T) { + for _, test := range []struct { + name string + source string + sourceFile string + target string + targetFile string + }{ + {name: "standard", source: "repo/source.go::source", sourceFile: "repo/source.go", target: "repo/target.go::target", targetFile: "repo/target.go"}, + {name: "detached", source: "legacy-source", target: "legacy-target"}, + } { + t.Run(test.name, func(t *testing.T) { + baseEdge := &Edge{From: test.source, To: test.target, Kind: EdgeCalls, FilePath: "base.go", Line: 10} + base := New() + base.AddEdge(baseEdge) + site := EdgeSourceSite{From: test.source, Line: 10} + markRemoved := func(layer *OverlayLayer, file, name, id string) { + if file != "" { + layer.MarkFile(file, false) + } + layer.MarkRemoved(name, id) + } + addCurrent := func(layer *OverlayLayer, file, name, id string, kind NodeKind) { + if file == "" { + file = "overlay.go" + } + layer.AddNode(file, &Node{ID: id, Name: name, Kind: kind, FilePath: file}) + } + + sourceTombstone := NewOverlayLayer() + markRemoved(sourceTombstone, test.sourceFile, "source", test.source) + sourceTombstone.AddEdge(&Edge{From: test.source, To: test.target, Kind: EdgeCalls, FilePath: "stray.go", Line: 10}) + projection, err := NewOverlaidView(base, sourceTombstone).FindOutgoingSiteEdgeIdentitiesBounded( + context.Background(), []EdgeSourceSite{site}, []EdgeKind{EdgeCalls}, 1, + ) + if err != nil || len(projection.BySite) != 0 || len(projection.Truncated) != 0 { + t.Fatalf("source tombstone site = %#v, %v", projection, err) + } + + sourceReplacement := NewOverlayLayer() + markRemoved(sourceReplacement, test.sourceFile, "source", test.source) + addCurrent(sourceReplacement, test.sourceFile, "source", test.source, KindMethod) + replacementEdge := &Edge{From: test.source, To: test.target, Kind: EdgeCalls, FilePath: "overlay.go", Line: 10} + sourceReplacement.AddEdge(replacementEdge) + projection, err = NewOverlaidView(base, sourceReplacement).FindOutgoingSiteEdgeIdentitiesBounded( + context.Background(), []EdgeSourceSite{site}, []EdgeKind{EdgeCalls}, 1, + ) + if err != nil { + t.Fatalf("source replacement site: %v", err) + } + requireSingleBoundedIdentity(t, projection.BySite[site], EdgeIdentityFor(replacementEdge)) + + targetTombstone := NewOverlayLayer() + markRemoved(targetTombstone, test.targetFile, "target", test.target) + projection, err = NewOverlaidView(base, targetTombstone).FindOutgoingSiteEdgeIdentitiesBounded( + context.Background(), []EdgeSourceSite{site}, []EdgeKind{EdgeCalls}, 1, + ) + if err != nil || len(projection.BySite) != 0 || len(projection.Truncated) != 0 { + t.Fatalf("target tombstone site = %#v, %v", projection, err) + } + + targetReplacement := NewOverlayLayer() + markRemoved(targetReplacement, test.targetFile, "target", test.target) + addCurrent(targetReplacement, test.targetFile, "target", test.target, KindType) + projection, err = NewOverlaidView(base, targetReplacement).FindOutgoingSiteEdgeIdentitiesBounded( + context.Background(), []EdgeSourceSite{site}, []EdgeKind{EdgeCalls}, 1, + ) + if err != nil { + t.Fatalf("target replacement site: %v", err) + } + requireSingleBoundedIdentity(t, projection.BySite[site], EdgeIdentityFor(baseEdge)) + }) + } +} + +func TestOverlaidViewBoundedAdjacencyCompensatesMultipleHiddenIdentities(t *testing.T) { + for _, test := range []struct { + name string + hiddenID string + hiddenFile string + }{ + {name: "standard", hiddenID: "repo/hidden.go::hidden", hiddenFile: "repo/hidden.go"}, + {name: "detached", hiddenID: "legacy-hidden"}, + } { + t.Run(test.name, func(t *testing.T) { + const source = "repo/source.go::source" + visible := &Edge{From: source, To: "visible-target", Kind: EdgeCalls, FilePath: "visible.go", Line: 10} + base := New() + base.AddEdge(&Edge{From: source, To: test.hiddenID, Kind: EdgeCalls, FilePath: "hidden-a.go", Line: 10}) + base.AddEdge(&Edge{From: source, To: test.hiddenID, Kind: EdgeCalls, FilePath: "hidden-b.go", Line: 10}) + base.AddEdge(visible) + layer := NewOverlayLayer() + if test.hiddenFile != "" { + layer.MarkFile(test.hiddenFile, false) + } + layer.MarkRemoved("hidden", test.hiddenID) + view := NewOverlaidView(base, layer) + out, err := view.FindOutgoingEdgeIdentitiesBounded(context.Background(), []string{source}, []EdgeKind{EdgeCalls}, 1) + if err != nil || out.Truncated[source] { + t.Fatalf("compensated outgoing = %#v, %v", out, err) + } + requireSingleBoundedIdentity(t, out.ByEndpoint[source], EdgeIdentityFor(visible)) + site := EdgeSourceSite{From: source, Line: 10} + bySite, err := view.FindOutgoingSiteEdgeIdentitiesBounded(context.Background(), []EdgeSourceSite{site}, []EdgeKind{EdgeCalls}, 1) + if err != nil || bySite.Truncated[site] { + t.Fatalf("compensated site = %#v, %v", bySite, err) + } + requireSingleBoundedIdentity(t, bySite.BySite[site], EdgeIdentityFor(visible)) + + const target = "repo/target.go::target" + visibleSource := &Edge{From: "repo/visible.go::source", To: target, Kind: EdgeCalls, FilePath: "visible.go", Line: 20} + incomingBase := New() + incomingBase.AddEdge(&Edge{From: test.hiddenID, To: target, Kind: EdgeCalls, FilePath: "hidden-a.go", Line: 20}) + incomingBase.AddEdge(&Edge{From: test.hiddenID, To: target, Kind: EdgeCalls, FilePath: "hidden-b.go", Line: 21}) + incomingBase.AddEdge(visibleSource) + incoming, err := NewOverlaidView(incomingBase, layer).FindIncomingEdgeIdentitiesBounded( + context.Background(), []string{target}, []EdgeKind{EdgeCalls}, 1, + ) + if err != nil || incoming.Truncated[target] { + t.Fatalf("compensated incoming = %#v, %v", incoming, err) + } + requireSingleBoundedIdentity(t, incoming.ByEndpoint[target], EdgeIdentityFor(visibleSource)) + }) + } +} + +type boundedAdjacencyUnsupportedReader struct{ Reader } + +func TestOverlaidViewBoundedAdjacencyFailsClosedWithoutBaseCapability(t *testing.T) { + base := boundedAdjacencyUnsupportedReader{Reader: New()} + layer := NewOverlayLayer() + const overlaySource = "overlay.go::source" + layer.MarkFile("overlay.go", false) + layer.AddNode("overlay.go", &Node{ID: overlaySource, Name: "source", Kind: KindMethod, FilePath: "overlay.go"}) + layer.AddEdge(&Edge{From: overlaySource, To: "target", Kind: EdgeCalls}) + projection, err := NewOverlaidView(base, layer).FindOutgoingEdgeIdentitiesBounded( + context.Background(), []string{overlaySource, "durable-source"}, []EdgeKind{EdgeCalls}, 1, + ) + if !errors.Is(err, ErrBoundedLocalizationUnavailable) || len(projection.ByEndpoint) != 0 || len(projection.Truncated) != 0 { + t.Fatalf("unsupported base leaked partial overlay evidence: %#v, %v", projection, err) + } +} + +var errBoundedAdjacencyBase = errors.New("bounded adjacency base failure") + +type boundedAdjacencyErrorReader struct{ Reader } + +func (boundedAdjacencyErrorReader) FindOutgoingEdgeIdentitiesBounded(context.Context, []string, []EdgeKind, int) (BoundedEdgeIdentityProjection, error) { + return BoundedEdgeIdentityProjection{ByEndpoint: map[string][]EdgeIdentity{"prefilled": {{From: "stale"}}}}, errBoundedAdjacencyBase +} + +func (boundedAdjacencyErrorReader) FindIncomingEdgeIdentitiesBounded(context.Context, []string, []EdgeKind, int) (BoundedEdgeIdentityProjection, error) { + return BoundedEdgeIdentityProjection{ByEndpoint: map[string][]EdgeIdentity{"prefilled": {{From: "stale"}}}}, errBoundedAdjacencyBase +} + +func (boundedAdjacencyErrorReader) FindOutgoingSiteEdgeIdentitiesBounded(context.Context, []EdgeSourceSite, []EdgeKind, int) (BoundedSiteEdgeIdentityProjection, error) { + return BoundedSiteEdgeIdentityProjection{BySite: map[EdgeSourceSite][]EdgeIdentity{{From: "stale", Line: 1}: {{From: "stale"}}}}, errBoundedAdjacencyBase +} + +func TestOverlaidViewBoundedAdjacencyDiscardsBasePartialOnError(t *testing.T) { + base := boundedAdjacencyErrorReader{Reader: New()} + layer := NewOverlayLayer() + const overlaySource = "overlay.go::source" + layer.MarkFile("overlay.go", false) + layer.AddNode("overlay.go", &Node{ID: overlaySource, Name: "source", Kind: KindMethod, FilePath: "overlay.go"}) + layer.AddEdge(&Edge{From: overlaySource, To: "target", Kind: EdgeCalls, Line: 1}) + view := NewOverlaidView(base, layer) + + out, err := view.FindOutgoingEdgeIdentitiesBounded(context.Background(), []string{overlaySource, "base-source"}, []EdgeKind{EdgeCalls}, 1) + if !errors.Is(err, errBoundedAdjacencyBase) || len(out.ByEndpoint) != 0 || len(out.Truncated) != 0 { + t.Fatalf("outgoing base error leaked partial: %#v, %v", out, err) + } + in, err := view.FindIncomingEdgeIdentitiesBounded(context.Background(), []string{"target", "base-target"}, []EdgeKind{EdgeCalls}, 1) + if !errors.Is(err, errBoundedAdjacencyBase) || len(in.ByEndpoint) != 0 || len(in.Truncated) != 0 { + t.Fatalf("incoming base error leaked partial: %#v, %v", in, err) + } + site := EdgeSourceSite{From: overlaySource, Line: 1} + bySite, err := view.FindOutgoingSiteEdgeIdentitiesBounded(context.Background(), []EdgeSourceSite{site, {From: "base-source", Line: 1}}, []EdgeKind{EdgeCalls}, 1) + if !errors.Is(err, errBoundedAdjacencyBase) || len(bySite.BySite) != 0 || len(bySite.Truncated) != 0 { + t.Fatalf("site base error leaked partial: %#v, %v", bySite, err) + } +} + +func TestOverlaidViewBoundedAdjacencyInspectionAndCancellationAreFailClosed(t *testing.T) { + const source = "overlay.go::source" + layer := NewOverlayLayer() + layer.MarkFile("overlay.go", false) + layer.AddNode("overlay.go", &Node{ID: source, Name: "source", Kind: KindMethod, FilePath: "overlay.go"}) + for index := 0; index < MaxBoundedAdjacencyInspectedEdges; index++ { + layer.AddEdge(&Edge{From: source, To: fmt.Sprintf("noise-%05d", index), Kind: EdgeReferences, Line: index}) + } + layer.AddEdge(&Edge{From: source, To: "wanted", Kind: EdgeCalls, Line: MaxBoundedAdjacencyInspectedEdges}) + view := NewOverlaidView(New(), layer) + projection, err := view.FindOutgoingEdgeIdentitiesBounded(context.Background(), []string{source}, []EdgeKind{EdgeCalls}, 1) + var limitErr *BoundedLocalizationLimitError + if !errors.As(err, &limitErr) || limitErr.Resource != "adjacency inspected edge rows" || + len(projection.ByEndpoint) != 0 || len(projection.Truncated) != 0 { + t.Fatalf("overlay inspection overflow = %#v, %v", projection, err) + } + + ctx := &cancelAdjacencyAfterChecksContext{Context: context.Background(), remaining: 5} + projection, err = view.FindOutgoingEdgeIdentitiesBounded(ctx, []string{source}, []EdgeKind{EdgeCalls}, 1) + if !errors.Is(err, context.Canceled) || len(projection.ByEndpoint) != 0 || len(projection.Truncated) != 0 { + t.Fatalf("canceled overlay projection leaked partial rows: %#v, %v", projection, err) + } +} diff --git a/internal/graph/bounded_adjacency_test.go b/internal/graph/bounded_adjacency_test.go new file mode 100644 index 00000000..43e4d80b --- /dev/null +++ b/internal/graph/bounded_adjacency_test.go @@ -0,0 +1,232 @@ +package graph + +import ( + "context" + "errors" + "fmt" + "math" + "reflect" + "testing" +) + +func TestGraphBoundedAdjacencyFiltersKindsAndSitesBeforeLimit(t *testing.T) { + g := New() + const source = "repo/source.go::source" + for index := 0; index < 300; index++ { + g.AddEdge(&Edge{ + From: source, To: fmt.Sprintf("noise-%03d", index), Kind: EdgeReferences, + FilePath: "repo/source.go", Line: 1_000 + index, + Meta: map[string]any{"payload": "must not cross projection"}, + }) + } + first := &Edge{From: source, To: "target", Kind: EdgeCalls, FilePath: "repo/source.go", Line: 10} + second := &Edge{From: source, To: "other", Kind: EdgeCalls, FilePath: "generated.go", Line: 20} + third := &Edge{From: "repo/other.go::caller", To: "target", Kind: EdgeCalls, FilePath: "repo/other.go", Line: 30} + g.AddEdge(first) + g.AddEdge(second) + g.AddEdge(third) + + sources := []string{source, source, ""} + kinds := []EdgeKind{EdgeCalls, EdgeCalls, ""} + sourcesBefore := append([]string(nil), sources...) + kindsBefore := append([]EdgeKind(nil), kinds...) + outgoing, err := g.FindOutgoingEdgeIdentitiesBounded(context.Background(), sources, kinds, 2) + if err != nil { + t.Fatalf("outgoing projection: %v", err) + } + if !reflect.DeepEqual(sources, sourcesBefore) || !reflect.DeepEqual(kinds, kindsBefore) { + t.Fatalf("caller inputs mutated: sources=%#v kinds=%#v", sources, kinds) + } + wantOutgoing := []EdgeIdentity{EdgeIdentityFor(second), EdgeIdentityFor(first)} + sortEdgeIdentities(wantOutgoing) + if !reflect.DeepEqual(outgoing.ByEndpoint[source], wantOutgoing) || outgoing.Truncated[source] { + t.Fatalf("outgoing = %#v, want %#v complete", outgoing, wantOutgoing) + } + + incoming, err := g.FindIncomingEdgeIdentitiesBounded( + context.Background(), []string{"target"}, []EdgeKind{EdgeCalls}, 2, + ) + if err != nil { + t.Fatalf("incoming projection: %v", err) + } + wantIncoming := []EdgeIdentity{EdgeIdentityFor(first), EdgeIdentityFor(third)} + sortEdgeIdentities(wantIncoming) + if !reflect.DeepEqual(incoming.ByEndpoint["target"], wantIncoming) || incoming.Truncated["target"] { + t.Fatalf("incoming = %#v, want %#v complete", incoming, wantIncoming) + } + + site := EdgeSourceSite{From: source, Line: 10} + sites := []EdgeSourceSite{site, site} + sitesBefore := append([]EdgeSourceSite(nil), sites...) + siteProjection, err := g.FindOutgoingSiteEdgeIdentitiesBounded( + context.Background(), sites, []EdgeKind{EdgeCalls}, 1, + ) + if err != nil { + t.Fatalf("site projection: %v", err) + } + if !reflect.DeepEqual(sites, sitesBefore) { + t.Fatalf("caller sites mutated: %#v", sites) + } + if got := siteProjection.BySite[site]; !reflect.DeepEqual(got, []EdgeIdentity{EdgeIdentityFor(first)}) || siteProjection.Truncated[site] { + t.Fatalf("site = %#v, want exact source+line identity", siteProjection) + } +} + +func TestGraphBoundedAdjacencyRejectsInvalidRawEnvelopes(t *testing.T) { + g := New() + tooManyKeys := make([]string, MaxBoundedAdjacencyKeys+1) + for index := range tooManyKeys { + tooManyKeys[index] = "duplicate" + } + tooManyKinds := make([]EdgeKind, MaxBoundedAdjacencyKinds+1) + for index := range tooManyKinds { + tooManyKinds[index] = EdgeCalls + } + for _, test := range []struct { + name string + run func() error + }{ + {name: "raw keys", run: func() error { + _, err := g.FindOutgoingEdgeIdentitiesBounded(context.Background(), tooManyKeys, []EdgeKind{EdgeCalls}, 1) + return err + }}, + {name: "raw kinds", run: func() error { + _, err := g.FindIncomingEdgeIdentitiesBounded(context.Background(), []string{"target"}, tooManyKinds, 1) + return err + }}, + {name: "zero limit on empty input", run: func() error { + _, err := g.FindOutgoingSiteEdgeIdentitiesBounded(context.Background(), nil, nil, 0) + return err + }}, + {name: "max int limit", run: func() error { + _, err := g.FindOutgoingEdgeIdentitiesBounded(context.Background(), nil, nil, math.MaxInt) + return err + }}, + } { + t.Run(test.name, func(t *testing.T) { + var limitErr *BoundedLocalizationLimitError + if err := test.run(); !errors.As(err, &limitErr) { + t.Fatalf("error = %v, want typed bounded limit", err) + } + }) + } +} + +func TestGraphBoundedAdjacencyPerKeyAndInspectionSentinels(t *testing.T) { + g := New() + const source = "dense-source" + for index := 0; index <= MaxBoundedAdjacencyRowsPerKey; index++ { + g.AddEdge(&Edge{From: source, To: fmt.Sprintf("target-%03d", index), Kind: EdgeCalls, Line: index}) + } + projection, err := g.FindOutgoingEdgeIdentitiesBounded( + context.Background(), []string{source}, []EdgeKind{EdgeCalls}, MaxBoundedAdjacencyRowsPerKey, + ) + if err != nil { + t.Fatalf("per-key sentinel: %v", err) + } + if !projection.Truncated[source] || len(projection.ByEndpoint[source]) != 0 { + t.Fatalf("per-key sentinel leaked partial rows: %#v", projection) + } + + flood := New() + for index := 0; index < MaxBoundedAdjacencyInspectedEdges; index++ { + flood.AddEdge(&Edge{From: source, To: fmt.Sprintf("noise-%05d", index), Kind: EdgeReferences, Line: index}) + } + flood.AddEdge(&Edge{From: source, To: "wanted", Kind: EdgeCalls, Line: MaxBoundedAdjacencyInspectedEdges}) + projection, err = flood.FindOutgoingEdgeIdentitiesBounded( + context.Background(), []string{source}, []EdgeKind{EdgeCalls}, 1, + ) + var limitErr *BoundedLocalizationLimitError + if !errors.As(err, &limitErr) || limitErr.Resource != "adjacency inspected edge rows" || len(projection.ByEndpoint) != 0 { + t.Fatalf("inspection overflow = %#v, %v", projection, err) + } +} + +type cancelAdjacencyAfterChecksContext struct { + context.Context + remaining int + checks int +} + +func (ctx *cancelAdjacencyAfterChecksContext) Err() error { + ctx.checks++ + if ctx.remaining == 0 { + return context.Canceled + } + ctx.remaining-- + return nil +} + +func TestGraphBoundedAdjacencyMidScanCancellationReturnsNoPartial(t *testing.T) { + g := New() + const source = "cancel-source" + for index := 0; index < 512; index++ { + g.AddEdge(&Edge{From: source, To: fmt.Sprintf("target-%03d", index), Kind: EdgeCalls, Line: index}) + } + + // The early checks cover entry, canonicalization, and entering the source. + // This threshold cancels at a later scan checkpoint, after rows have been + // accumulated that must not escape. + ctx := &cancelAdjacencyAfterChecksContext{Context: context.Background(), remaining: 5} + projection, err := g.FindOutgoingEdgeIdentitiesBounded(ctx, []string{source}, []EdgeKind{EdgeCalls}, MaxBoundedAdjacencyRowsPerKey) + if !errors.Is(err, context.Canceled) { + t.Fatalf("error = %v, want context cancellation", err) + } + if len(projection.ByEndpoint) != 0 || len(projection.Truncated) != 0 { + t.Fatalf("canceled projection leaked partial rows: %#v", projection) + } + if ctx.checks < 6 { + t.Fatalf("cancellation checks = %d, want scan progress", ctx.checks) + } +} + +func TestGraphBoundedAdjacencySiteScanChargesEachSourceRowOnce(t *testing.T) { + g := New() + const source = "many-sites" + sites := make([]EdgeSourceSite, 0, MaxBoundedAdjacencyKeys) + for line := 0; line < MaxBoundedAdjacencyKeys; line++ { + g.AddEdge(&Edge{From: source, To: fmt.Sprintf("target-%03d", line), Kind: EdgeCalls, Line: line}) + sites = append(sites, EdgeSourceSite{From: source, Line: line}) + } + for index := MaxBoundedAdjacencyKeys; index < MaxBoundedAdjacencyInspectedEdges; index++ { + g.AddEdge(&Edge{From: source, To: fmt.Sprintf("noise-%05d", index), Kind: EdgeReferences, Line: index}) + } + projection, err := g.FindOutgoingSiteEdgeIdentitiesBounded( + context.Background(), sites, []EdgeKind{EdgeCalls}, 1, + ) + if err != nil { + t.Fatalf("same-source site projection: %v", err) + } + if len(projection.BySite) != MaxBoundedAdjacencyKeys || len(projection.Truncated) != 0 { + t.Fatalf("same-source sites = %d complete/%d truncated, want %d/0", + len(projection.BySite), len(projection.Truncated), MaxBoundedAdjacencyKeys) + } +} + +func TestGraphBoundedAdjacencyTotalIdentityCap(t *testing.T) { + g := New() + sources := make([]string, 0, 65) + for sourceIndex := 0; sourceIndex < 65; sourceIndex++ { + source := fmt.Sprintf("source-%02d", sourceIndex) + sources = append(sources, source) + for edgeIndex := 0; edgeIndex < MaxBoundedAdjacencyRowsPerKey; edgeIndex++ { + g.AddEdge(&Edge{ + From: source, To: fmt.Sprintf("target-%02d-%03d", sourceIndex, edgeIndex), + Kind: EdgeCalls, Line: edgeIndex, + }) + } + } + projection, err := g.FindOutgoingEdgeIdentitiesBounded( + context.Background(), sources[:64], []EdgeKind{EdgeCalls}, MaxBoundedAdjacencyRowsPerKey, + ) + if err != nil || len(projection.ByEndpoint) != 64 { + t.Fatalf("exact aggregate cap = %d endpoints, %v", len(projection.ByEndpoint), err) + } + projection, err = g.FindOutgoingEdgeIdentitiesBounded( + context.Background(), sources, []EdgeKind{EdgeCalls}, MaxBoundedAdjacencyRowsPerKey, + ) + var limitErr *BoundedLocalizationLimitError + if !errors.As(err, &limitErr) || limitErr.Resource != "adjacency inspected edge rows" || len(projection.ByEndpoint) != 0 { + t.Fatalf("aggregate overflow = %#v, %v", projection, err) + } +} diff --git a/internal/graph/bounded_edge_existence.go b/internal/graph/bounded_edge_existence.go new file mode 100644 index 00000000..d75c7247 --- /dev/null +++ b/internal/graph/bounded_edge_existence.go @@ -0,0 +1,274 @@ +package graph + +import ( + "context" + "sort" +) + +// MaxBoundedEdgeExistencePredicates is the hard request envelope for exact +// edge-existence projections. It is intentionally sized to the qualified-leaf +// localization frontier: ten candidate declarations times at most thirty-two +// admitted owner declarations. +const MaxBoundedEdgeExistencePredicates = 320 + +// TypedEdgeEndpoint identifies an exact directed edge independent of source +// location and payload. It is comparable and can be used directly as a map key. +type TypedEdgeEndpoint struct { + From string + To string + Kind EdgeKind +} + +// BoundedEdgeExistenceReader answers a finite set of exact edge predicates +// without materializing adjacency or Edge payloads. Implementations must either +// return the complete set of existing predicates or an error; partial evidence +// is never valid. limit applies to unique, non-empty predicates and may not +// exceed MaxBoundedEdgeExistencePredicates. +type BoundedEdgeExistenceReader interface { + FindExistingEdgeEndpoints(context.Context, []TypedEdgeEndpoint, int) (map[TypedEdgeEndpoint]struct{}, error) +} + +var ( + _ BoundedEdgeExistenceReader = (*Graph)(nil) + _ BoundedEdgeExistenceReader = (*OverlaidView)(nil) +) + +type typedEdgeTarget struct { + to string + kind EdgeKind +} + +func canonicalBoundedEdgeEndpoints( + ctx context.Context, + endpoints []TypedEdgeEndpoint, + limit int, +) ([]TypedEdgeEndpoint, error) { + if ctx == nil { + ctx = context.Background() + } + if err := ctx.Err(); err != nil { + return nil, err + } + if limit < 1 || limit > MaxBoundedEdgeExistencePredicates { + return nil, &BoundedLocalizationLimitError{ + Resource: "edge-existence predicates", + Limit: MaxBoundedEdgeExistencePredicates, + } + } + if len(endpoints) == 0 { + return nil, nil + } + capacity := len(endpoints) + if capacity > limit+1 { + capacity = limit + 1 + } + seen := make(map[TypedEdgeEndpoint]struct{}, capacity) + out := make([]TypedEdgeEndpoint, 0, capacity) + for index, endpoint := range endpoints { + if index&127 == 0 { + if err := ctx.Err(); err != nil { + return nil, err + } + } + if endpoint.From == "" || endpoint.To == "" || endpoint.Kind == "" { + continue + } + if _, duplicate := seen[endpoint]; duplicate { + continue + } + seen[endpoint] = struct{}{} + if len(seen) > limit { + return nil, &BoundedLocalizationLimitError{ + Resource: "edge-existence predicates", + Limit: limit, + } + } + out = append(out, endpoint) + } + sort.Slice(out, func(i, j int) bool { + if out[i].From != out[j].From { + return out[i].From < out[j].From + } + if out[i].To != out[j].To { + return out[i].To < out[j].To + } + return out[i].Kind < out[j].Kind + }) + return out, nil +} + +// FindExistingEdgeEndpoints is the in-memory reference implementation. It +// groups predicates by source, takes one shard read lock per source, and stops +// scanning that adjacency bucket as soon as every requested predicate is found. +func (g *Graph) FindExistingEdgeEndpoints( + ctx context.Context, + endpoints []TypedEdgeEndpoint, + limit int, +) (map[TypedEdgeEndpoint]struct{}, error) { + if ctx == nil { + ctx = context.Background() + } + keys, err := canonicalBoundedEdgeEndpoints(ctx, endpoints, limit) + if err != nil { + return nil, err + } + found := make(map[TypedEdgeEndpoint]struct{}, len(keys)) + if g == nil || len(keys) == 0 { + return found, nil + } + for start := 0; start < len(keys); { + if err := ctx.Err(); err != nil { + return nil, err + } + end := start + 1 + for end < len(keys) && keys[end].From == keys[start].From { + end++ + } + from := keys[start].From + wanted := make(map[typedEdgeTarget]TypedEdgeEndpoint, end-start) + for _, endpoint := range keys[start:end] { + wanted[typedEdgeTarget{to: endpoint.To, kind: endpoint.Kind}] = endpoint + } + shard := g.shardFor(from) + shard.mu.RLock() + for index, edge := range shard.outEdges[from] { + if index&127 == 0 { + if err := ctx.Err(); err != nil { + shard.mu.RUnlock() + return nil, err + } + } + if edge == nil { + continue + } + target := typedEdgeTarget{to: edge.To, kind: edge.Kind} + endpoint, ok := wanted[target] + if !ok { + continue + } + found[endpoint] = struct{}{} + delete(wanted, target) + if len(wanted) == 0 { + break + } + } + shard.mu.RUnlock() + start = end + } + return found, nil +} + +// FindExistingEdgeEndpoints composes exact base evidence with one immutable +// overlay. A source owned by the overlay uses only current overlay edges. For an +// untouched source, base evidence survives unless the exact target identity is +// tombstoned by the overlay; a same-ID replacement preserves that base edge. +func (v *OverlaidView) FindExistingEdgeEndpoints( + ctx context.Context, + endpoints []TypedEdgeEndpoint, + limit int, +) (map[TypedEdgeEndpoint]struct{}, error) { + if ctx == nil { + ctx = context.Background() + } + keys, err := canonicalBoundedEdgeEndpoints(ctx, endpoints, limit) + if err != nil { + return nil, err + } + found := make(map[TypedEdgeEndpoint]struct{}, len(keys)) + if v == nil || len(keys) == 0 { + return found, nil + } + if v.layer == nil { + if v.base == nil { + return found, nil + } + bounded, ok := v.base.(BoundedEdgeExistenceReader) + if !ok { + return nil, ErrBoundedLocalizationUnavailable + } + return bounded.FindExistingEdgeEndpoints(ctx, keys, limit) + } + + baseKeys := make([]TypedEdgeEndpoint, 0, len(keys)) + for start := 0; start < len(keys); { + if err := ctx.Err(); err != nil { + return nil, err + } + end := start + 1 + for end < len(keys) && keys[end].From == keys[start].From { + end++ + } + from := keys[start].From + sourceOwned := v.nodeBelongsToOverlay(from) || v.layer.ownsNodeIdentity(from) + if sourceOwned { + // A covered or detached tombstone owns the identity but carries no + // current adjacency. Never authenticate a stray staged edge for a + // source declaration the overlay removed. + if v.layer.nodeByID[from] == nil { + start = end + continue + } + wanted := make(map[typedEdgeTarget]TypedEdgeEndpoint, end-start) + for _, endpoint := range keys[start:end] { + targetOwned := v.nodeBelongsToOverlay(endpoint.To) || v.layer.ownsNodeIdentity(endpoint.To) + if targetOwned && v.layer.nodeByID[endpoint.To] == nil { + continue + } + wanted[typedEdgeTarget{to: endpoint.To, kind: endpoint.Kind}] = endpoint + } + if len(wanted) == 0 { + start = end + continue + } + for index, edge := range v.layer.outEdges[from] { + if index&127 == 0 { + if err := ctx.Err(); err != nil { + return nil, err + } + } + if edge == nil { + continue + } + target := typedEdgeTarget{to: edge.To, kind: edge.Kind} + endpoint, ok := wanted[target] + if !ok { + continue + } + found[endpoint] = struct{}{} + delete(wanted, target) + if len(wanted) == 0 { + break + } + } + start = end + continue + } + for _, endpoint := range keys[start:end] { + targetOwned := v.nodeBelongsToOverlay(endpoint.To) || v.layer.ownsNodeIdentity(endpoint.To) + if targetOwned && v.layer.nodeByID[endpoint.To] == nil { + continue + } + baseKeys = append(baseKeys, endpoint) + } + start = end + } + + if len(baseKeys) == 0 || v.base == nil { + return found, nil + } + bounded, ok := v.base.(BoundedEdgeExistenceReader) + if !ok { + return nil, ErrBoundedLocalizationUnavailable + } + baseFound, err := bounded.FindExistingEdgeEndpoints(ctx, baseKeys, limit) + if err != nil { + return nil, err + } + if err := ctx.Err(); err != nil { + return nil, err + } + for endpoint := range baseFound { + found[endpoint] = struct{}{} + } + return found, nil +} diff --git a/internal/graph/bounded_edge_existence_test.go b/internal/graph/bounded_edge_existence_test.go new file mode 100644 index 00000000..76708b15 --- /dev/null +++ b/internal/graph/bounded_edge_existence_test.go @@ -0,0 +1,248 @@ +package graph + +import ( + "context" + "errors" + "fmt" + "reflect" + "testing" +) + +func TestGraphFindExistingEdgeEndpointsBoundsKindsAndCancellation(t *testing.T) { + g := New() + g.AddEdge(&Edge{From: "source", To: "owner", Kind: EdgeMemberOf, FilePath: "a.go", Line: 1}) + g.AddEdge(&Edge{From: "source", To: "owner", Kind: EdgeMemberOf, FilePath: "a.go", Line: 2}) + g.AddEdge(&Edge{From: "source", To: "owner", Kind: EdgeCalls}) + wanted := TypedEdgeEndpoint{From: "source", To: "owner", Kind: EdgeMemberOf} + missing := TypedEdgeEndpoint{From: "source", To: "other", Kind: EdgeMemberOf} + input := []TypedEdgeEndpoint{wanted, wanted, missing, {}, {From: "source", To: "owner", Kind: EdgeCalls}} + before := append([]TypedEdgeEndpoint(nil), input...) + found, err := g.FindExistingEdgeEndpoints(context.Background(), input, 4) + if err != nil { + t.Fatalf("bounded edge existence: %v", err) + } + if !reflect.DeepEqual(input, before) { + t.Fatalf("caller predicates mutated: got %#v want %#v", input, before) + } + if _, ok := found[wanted]; !ok { + t.Fatalf("missing exact member_of predicate: %#v", found) + } + if _, ok := found[missing]; ok { + t.Fatalf("reported missing predicate: %#v", found) + } + if _, ok := found[TypedEdgeEndpoint{From: "source", To: "owner", Kind: EdgeCalls}]; !ok { + t.Fatalf("kind-specific predicate missing: %#v", found) + } + + boundary := make([]TypedEdgeEndpoint, 0, MaxBoundedEdgeExistencePredicates+1) + for index := 0; index <= MaxBoundedEdgeExistencePredicates; index++ { + boundary = append(boundary, TypedEdgeEndpoint{ + From: "source", To: fmt.Sprintf("target-%03d", index), Kind: EdgeCalls, + }) + } + if result, err := g.FindExistingEdgeEndpoints( + context.Background(), boundary[:MaxBoundedEdgeExistencePredicates], MaxBoundedEdgeExistencePredicates, + ); err != nil || len(result) != 0 { + t.Fatalf("exact boundary = %#v, %v", result, err) + } + result, err := g.FindExistingEdgeEndpoints(context.Background(), boundary, MaxBoundedEdgeExistencePredicates) + var limitErr *BoundedLocalizationLimitError + if !errors.As(err, &limitErr) || len(result) != 0 || limitErr.Limit != MaxBoundedEdgeExistencePredicates { + t.Fatalf("over boundary = %#v, %v", result, err) + } + result, err = g.FindExistingEdgeEndpoints(context.Background(), []TypedEdgeEndpoint{wanted}, MaxBoundedEdgeExistencePredicates+1) + if !errors.As(err, &limitErr) || len(result) != 0 { + t.Fatalf("oversized declared limit = %#v, %v", result, err) + } + for _, invalidLimit := range []int{0, MaxBoundedEdgeExistencePredicates + 1} { + result, err = g.FindExistingEdgeEndpoints(context.Background(), nil, invalidLimit) + if !errors.As(err, &limitErr) || len(result) != 0 { + t.Fatalf("empty input with invalid limit %d = %#v, %v", invalidLimit, result, err) + } + } + + for index := 0; index < 512; index++ { + g.AddEdge(&Edge{From: "hot", To: fmt.Sprintf("noise-%03d", index), Kind: EdgeCalls}) + } + ctx := &cancelAfterLocalizationChecksContext{Context: context.Background(), remaining: 5, done: make(chan struct{})} + result, err = g.FindExistingEdgeEndpoints(ctx, []TypedEdgeEndpoint{{From: "hot", To: "absent", Kind: EdgeMemberOf}}, 1) + if !errors.Is(err, context.Canceled) || len(result) != 0 { + t.Fatalf("cancelled scan returned partial evidence: %#v, %v", result, err) + } +} + +func TestOverlaidViewFindExistingEdgeEndpointsReplacementAndTombstone(t *testing.T) { + const ( + sourceFile = "repo/source.go" + targetFile = "repo/target.go" + sourceID = sourceFile + "::source" + targetID = targetFile + "::target" + ) + key := TypedEdgeEndpoint{From: sourceID, To: targetID, Kind: EdgeMemberOf} + base := New() + base.AddEdge(&Edge{From: sourceID, To: targetID, Kind: EdgeMemberOf}) + assertFound := func(t *testing.T, layer *OverlayLayer, want bool) { + t.Helper() + found, err := NewOverlaidView(base, layer).FindExistingEdgeEndpoints( + context.Background(), []TypedEdgeEndpoint{key}, 1, + ) + if err != nil { + t.Fatalf("overlay edge existence: %v", err) + } + _, got := found[key] + if got != want { + t.Fatalf("found = %v (%#v), want %v", got, found, want) + } + } + + t.Run("source tombstone rejects stray edge", func(t *testing.T) { + layer := NewOverlayLayer() + layer.MarkFile(sourceFile, false) + layer.MarkRemoved("source", sourceID) + layer.AddEdge(&Edge{From: sourceID, To: targetID, Kind: EdgeMemberOf}) + assertFound(t, layer, false) + }) + t.Run("same ID source replacement uses overlay edge", func(t *testing.T) { + layer := NewOverlayLayer() + layer.MarkFile(sourceFile, false) + layer.MarkRemoved("source", sourceID) + layer.AddNode(sourceFile, &Node{ID: sourceID, Name: "source", Kind: KindMethod, FilePath: sourceFile}) + layer.AddEdge(&Edge{From: sourceID, To: targetID, Kind: EdgeMemberOf}) + assertFound(t, layer, true) + }) + t.Run("target tombstone suppresses base edge", func(t *testing.T) { + layer := NewOverlayLayer() + layer.MarkFile(targetFile, false) + layer.MarkRemoved("target", targetID) + assertFound(t, layer, false) + }) + t.Run("same ID target replacement preserves base edge", func(t *testing.T) { + layer := NewOverlayLayer() + layer.MarkFile(targetFile, false) + layer.MarkRemoved("target", targetID) + layer.AddNode(targetFile, &Node{ID: targetID, Name: "target", Kind: KindType, FilePath: targetFile}) + assertFound(t, layer, true) + }) + t.Run("overlay source cannot point to tombstoned target", func(t *testing.T) { + layer := NewOverlayLayer() + layer.MarkFile(sourceFile, false) + layer.MarkFile(targetFile, false) + layer.AddNode(sourceFile, &Node{ID: sourceID, Name: "source", Kind: KindMethod, FilePath: sourceFile}) + layer.MarkRemoved("target", targetID) + layer.AddEdge(&Edge{From: sourceID, To: targetID, Kind: EdgeMemberOf}) + assertFound(t, layer, false) + }) + t.Run("overlay source can point to same ID target replacement", func(t *testing.T) { + layer := NewOverlayLayer() + layer.MarkFile(sourceFile, false) + layer.MarkFile(targetFile, false) + layer.AddNode(sourceFile, &Node{ID: sourceID, Name: "source", Kind: KindMethod, FilePath: sourceFile}) + layer.MarkRemoved("target", targetID) + layer.AddNode(targetFile, &Node{ID: targetID, Name: "target", Kind: KindType, FilePath: targetFile}) + layer.AddEdge(&Edge{From: sourceID, To: targetID, Kind: EdgeMemberOf}) + assertFound(t, layer, true) + }) +} + +func TestOverlaidViewFindExistingEdgeEndpointsDetachedIdentityParity(t *testing.T) { + key := TypedEdgeEndpoint{From: "legacy-source", To: "legacy-target", Kind: EdgeMemberOf} + base := New() + base.AddEdge(&Edge{From: key.From, To: key.To, Kind: key.Kind}) + + tombstone := NewOverlayLayer() + tombstone.MarkRemoved("legacy-source", key.From) + tombstone.AddEdge(&Edge{From: key.From, To: key.To, Kind: key.Kind}) + found, err := NewOverlaidView(base, tombstone).FindExistingEdgeEndpoints(context.Background(), []TypedEdgeEndpoint{key}, 1) + if err != nil || len(found) != 0 { + t.Fatalf("detached source tombstone authenticated edge: %#v, %v", found, err) + } + + replacement := NewOverlayLayer() + replacement.MarkRemoved("legacy-source", key.From) + replacement.AddNode("overlay.go", &Node{ID: key.From, Name: "legacy-source", Kind: KindMethod, FilePath: "overlay.go"}) + replacement.AddEdge(&Edge{From: key.From, To: key.To, Kind: key.Kind}) + found, err = NewOverlaidView(base, replacement).FindExistingEdgeEndpoints(context.Background(), []TypedEdgeEndpoint{key}, 1) + if err != nil { + t.Fatalf("detached replacement: %v", err) + } + if _, ok := found[key]; !ok { + t.Fatalf("detached same-ID replacement missing: %#v", found) + } + + targetKey := TypedEdgeEndpoint{From: "stable-source", To: key.To, Kind: key.Kind} + base.AddEdge(&Edge{From: targetKey.From, To: targetKey.To, Kind: targetKey.Kind}) + targetTombstone := NewOverlayLayer() + targetTombstone.MarkRemoved("legacy-target", targetKey.To) + found, err = NewOverlaidView(base, targetTombstone).FindExistingEdgeEndpoints( + context.Background(), []TypedEdgeEndpoint{targetKey}, 1, + ) + if err != nil || len(found) != 0 { + t.Fatalf("detached target tombstone authenticated base edge: %#v, %v", found, err) + } + + targetReplacement := NewOverlayLayer() + targetReplacement.MarkRemoved("legacy-target", targetKey.To) + targetReplacement.AddNode("overlay.go", &Node{ + ID: targetKey.To, Name: "legacy-target", Kind: KindType, FilePath: "overlay.go", + }) + found, err = NewOverlaidView(base, targetReplacement).FindExistingEdgeEndpoints( + context.Background(), []TypedEdgeEndpoint{targetKey}, 1, + ) + if err != nil { + t.Fatalf("detached target replacement: %v", err) + } + if _, ok := found[targetKey]; !ok { + t.Fatalf("detached same-ID target replacement hid base edge: %#v", found) + } + + overlaySourceTombstone := NewOverlayLayer() + overlaySourceTombstone.AddNode("overlay.go", &Node{ + ID: targetKey.From, Name: "stable-source", Kind: KindMethod, FilePath: "overlay.go", + }) + overlaySourceTombstone.MarkRemoved("legacy-target", targetKey.To) + overlaySourceTombstone.AddEdge(&Edge{From: targetKey.From, To: targetKey.To, Kind: targetKey.Kind}) + found, err = NewOverlaidView(base, overlaySourceTombstone).FindExistingEdgeEndpoints( + context.Background(), []TypedEdgeEndpoint{targetKey}, 1, + ) + if err != nil || len(found) != 0 { + t.Fatalf("overlay source authenticated detached target tombstone: %#v, %v", found, err) + } + + overlaySourceReplacement := NewOverlayLayer() + overlaySourceReplacement.AddNode("overlay.go", &Node{ + ID: targetKey.From, Name: "stable-source", Kind: KindMethod, FilePath: "overlay.go", + }) + overlaySourceReplacement.MarkRemoved("legacy-target", targetKey.To) + overlaySourceReplacement.AddNode("overlay.go", &Node{ + ID: targetKey.To, Name: "legacy-target", Kind: KindType, FilePath: "overlay.go", + }) + overlaySourceReplacement.AddEdge(&Edge{From: targetKey.From, To: targetKey.To, Kind: targetKey.Kind}) + found, err = NewOverlaidView(base, overlaySourceReplacement).FindExistingEdgeEndpoints( + context.Background(), []TypedEdgeEndpoint{targetKey}, 1, + ) + if err != nil { + t.Fatalf("overlay source detached target replacement: %v", err) + } + if _, ok := found[targetKey]; !ok { + t.Fatalf("overlay source missed detached target replacement: %#v", found) + } +} + +type edgeExistenceUnsupportedReader struct{ Reader } + +func TestOverlaidViewFindExistingEdgeEndpointsFailsClosedWithoutCapability(t *testing.T) { + base := edgeExistenceUnsupportedReader{Reader: New()} + baseKey := TypedEdgeEndpoint{From: "source", To: "target", Kind: EdgeCalls} + overlayKey := TypedEdgeEndpoint{From: "overlay.go::source", To: "target", Kind: EdgeCalls} + layer := NewOverlayLayer() + layer.AddNode("overlay.go", &Node{ + ID: overlayKey.From, Name: "source", Kind: KindMethod, FilePath: "overlay.go", + }) + layer.AddEdge(&Edge{From: overlayKey.From, To: overlayKey.To, Kind: overlayKey.Kind}) + found, err := NewOverlaidView(base, layer).FindExistingEdgeEndpoints( + context.Background(), []TypedEdgeEndpoint{overlayKey, baseKey}, 2, + ) + if !errors.Is(err, ErrBoundedLocalizationUnavailable) || len(found) != 0 { + t.Fatalf("unsupported base leaked partial overlay evidence: %#v, %v", found, err) + } +} diff --git a/internal/graph/bounded_incoming_sources.go b/internal/graph/bounded_incoming_sources.go new file mode 100644 index 00000000..16b071d7 --- /dev/null +++ b/internal/graph/bounded_incoming_sources.go @@ -0,0 +1,367 @@ +package graph + +import ( + "context" + "sort" +) + +// BoundedIncomingSourceProjection is a metadata-free incoming-adjacency page. +// Sources contains distinct source IDs for each requested target. Truncated is +// set per target when more than limit distinct sources of the requested kind +// exist; callers must not interpret that target's partial Sources as complete. +type BoundedIncomingSourceProjection struct { + Sources map[string][]string + Truncated map[string]bool +} + +// BoundedIncomingSourceReader projects distinct incoming source identities for +// one edge kind under a per-target cap. It is optional so bounded callers can +// fail closed instead of falling back to full incoming-edge materialization. +type BoundedIncomingSourceReader interface { + FindIncomingSourcesBounded(context.Context, []string, EdgeKind, int) (BoundedIncomingSourceProjection, error) +} + +type contextNodesByIDsReader interface { + GetNodesByIDsContext(context.Context, []string) (map[string]*Node, error) +} + +var ( + _ BoundedIncomingSourceReader = (*Graph)(nil) + _ BoundedIncomingSourceReader = (*OverlaidView)(nil) + + maxBoundedIncomingSourceLimit = int(^uint(0) >> 1) +) + +func (l *OverlayLayer) ownsNodeIdentity(id string) bool { + if l == nil || id == "" { + return false + } + return l.nodeByID[id] != nil || l.removedByID[id] +} + +// GetNodesByIDsContext is the cancellable exact-refetch sibling used by +// bounded request paths. It preserves overlay ownership for both ordinary and +// detached legacy identities, never mutates the caller's ID slice, and delegates +// the durable partition to a contextual base when available. +func (v *OverlaidView) GetNodesByIDsContext(ctx context.Context, ids []string) (map[string]*Node, error) { + if len(ids) == 0 { + return nil, nil + } + if ctx == nil { + ctx = context.Background() + } + if err := ctx.Err(); err != nil { + return nil, err + } + out := make(map[string]*Node, len(ids)) + baseIDs := make([]string, 0, len(ids)) + seen := make(map[string]struct{}, len(ids)) + for index, id := range ids { + if index&127 == 0 { + if err := ctx.Err(); err != nil { + return nil, err + } + } + if id == "" { + continue + } + if _, duplicate := seen[id]; duplicate { + continue + } + seen[id] = struct{}{} + if v.layer != nil && (v.nodeBelongsToOverlay(id) || v.layer.ownsNodeIdentity(id)) { + if node := v.layer.nodeByID[id]; node != nil { + out[id] = node + } + continue + } + baseIDs = append(baseIDs, id) + } + if len(baseIDs) == 0 || v.base == nil { + return out, nil + } + var ( + base map[string]*Node + err error + ) + if contextual, ok := v.base.(contextNodesByIDsReader); ok { + base, err = contextual.GetNodesByIDsContext(ctx, baseIDs) + } else { + base = v.base.GetNodesByIDs(baseIDs) + } + if err != nil { + return nil, err + } + if err := ctx.Err(); err != nil { + return nil, err + } + for id, node := range base { + if node != nil { + out[id] = node + } + } + return out, nil +} + +func boundedIncomingTargetIDs(ids []string) []string { + seen := make(map[string]struct{}, len(ids)) + out := make([]string, 0, len(ids)) + for _, id := range ids { + if id == "" { + continue + } + if _, duplicate := seen[id]; duplicate { + continue + } + seen[id] = struct{}{} + out = append(out, id) + } + return out +} + +// FindIncomingSourcesBounded is the in-memory reference implementation. It +// scans immutable edge pointers under one shard read lock per target, retains +// only distinct source identities, and stops at the limit+1 sentinel. +func (g *Graph) FindIncomingSourcesBounded( + ctx context.Context, + targetIDs []string, + kind EdgeKind, + limit int, +) (BoundedIncomingSourceProjection, error) { + projection := BoundedIncomingSourceProjection{ + Sources: make(map[string][]string), + Truncated: make(map[string]bool), + } + if g == nil { + return projection, nil + } + if ctx == nil { + ctx = context.Background() + } + ids := boundedIncomingTargetIDs(targetIDs) + if limit <= 0 { + for _, id := range ids { + projection.Truncated[id] = true + } + return projection, nil + } + if limit >= maxBoundedIncomingSourceLimit { + return BoundedIncomingSourceProjection{}, &BoundedLocalizationLimitError{ + Resource: "incoming-source sentinel", + Limit: maxBoundedIncomingSourceLimit - 1, + } + } + for _, id := range ids { + if err := ctx.Err(); err != nil { + return BoundedIncomingSourceProjection{}, err + } + shard := g.shardFor(id) + shard.mu.RLock() + seen := make(map[string]struct{}, limit+1) + for index, edge := range shard.inEdges[id] { + if index&127 == 0 { + if err := ctx.Err(); err != nil { + shard.mu.RUnlock() + return BoundedIncomingSourceProjection{}, err + } + } + if edge == nil || edge.Kind != kind || edge.From == "" { + continue + } + seen[edge.From] = struct{}{} + if len(seen) > limit { + projection.Truncated[id] = true + break + } + } + shard.mu.RUnlock() + if projection.Truncated[id] { + continue + } + sources := make([]string, 0, len(seen)) + for sourceID := range seen { + sources = append(sources, sourceID) + } + sort.Strings(sources) + if len(sources) > 0 { + projection.Sources[id] = sources + } + } + return projection, nil +} + +// FindIncomingSourcesBounded composes a bounded base projection with the +// request overlay. Base sources from overlay-owned files are hidden; current +// overlay edges replace them. Shadow compensation is bounded by the same hard +// detached-identity envelope used by exact-name localization. +func (v *OverlaidView) FindIncomingSourcesBounded( + ctx context.Context, + targetIDs []string, + kind EdgeKind, + limit int, +) (BoundedIncomingSourceProjection, error) { + projection := BoundedIncomingSourceProjection{ + Sources: make(map[string][]string), + Truncated: make(map[string]bool), + } + if v == nil { + return projection, nil + } + if ctx == nil { + ctx = context.Background() + } + ids := boundedIncomingTargetIDs(targetIDs) + if limit <= 0 { + for _, id := range ids { + projection.Truncated[id] = true + } + return projection, nil + } + if err := ctx.Err(); err != nil { + return BoundedIncomingSourceProjection{}, err + } + if limit >= maxBoundedIncomingSourceLimit { + return BoundedIncomingSourceProjection{}, &BoundedLocalizationLimitError{ + Resource: "overlay incoming-source sentinel", + Limit: maxBoundedIncomingSourceLimit - 1, + } + } + if v.layer == nil { + if v.base == nil { + return projection, nil + } + bounded, ok := v.base.(BoundedIncomingSourceReader) + if !ok { + return BoundedIncomingSourceProjection{}, ErrBoundedLocalizationUnavailable + } + return bounded.FindIncomingSourcesBounded(ctx, ids, kind, limit) + } + + shadowIDs := make(map[string]struct{}) + standardShadows := 0 + detachedShadows := 0 + addShadow := func(id string) error { + if id == "" { + return nil + } + if _, duplicate := shadowIDs[id]; duplicate { + return nil + } + shadowIDs[id] = struct{}{} + if filePath := IDFile(id); filePath != "" && v.layer.HasFile(filePath) { + standardShadows++ + if standardShadows > overlayExactNameInspectionLimit { + return &BoundedLocalizationLimitError{ + Resource: "overlay incoming-source standard shadows", + Limit: overlayExactNameInspectionLimit, + } + } + return nil + } + detachedShadows++ + if detachedShadows > overlayDetachedShadowLimit { + return &BoundedLocalizationLimitError{ + Resource: "overlay incoming-source detached shadows", + Limit: overlayDetachedShadowLimit, + } + } + return nil + } + inspectedShadows := 0 + for id := range v.layer.removedByID { + if inspectedShadows&127 == 0 { + if err := ctx.Err(); err != nil { + return BoundedIncomingSourceProjection{}, err + } + } + inspectedShadows++ + if err := addShadow(id); err != nil { + return BoundedIncomingSourceProjection{}, err + } + } + for id := range v.layer.nodeByID { + if inspectedShadows&127 == 0 { + if err := ctx.Err(); err != nil { + return BoundedIncomingSourceProjection{}, err + } + } + inspectedShadows++ + if err := addShadow(id); err != nil { + return BoundedIncomingSourceProjection{}, err + } + } + + baseProjection := BoundedIncomingSourceProjection{ + Sources: make(map[string][]string), + Truncated: make(map[string]bool), + } + if v.base != nil { + bounded, ok := v.base.(BoundedIncomingSourceReader) + if !ok { + return BoundedIncomingSourceProjection{}, ErrBoundedLocalizationUnavailable + } + var err error + if limit > maxBoundedIncomingSourceLimit-len(shadowIDs) { + return BoundedIncomingSourceProjection{}, &BoundedLocalizationLimitError{ + Resource: "overlay incoming-source compensation", + Limit: maxBoundedIncomingSourceLimit, + } + } + baseProjection, err = bounded.FindIncomingSourcesBounded(ctx, ids, kind, limit+len(shadowIDs)) + if err != nil { + return BoundedIncomingSourceProjection{}, err + } + } + + for _, targetID := range ids { + if err := ctx.Err(); err != nil { + return BoundedIncomingSourceProjection{}, err + } + targetRemoved := v.layer.removedByID[targetID] + if (targetRemoved || v.layer.HasFile(IDFile(targetID))) && v.layer.nodeByID[targetID] == nil { + continue + } + seen := make(map[string]struct{}, limit+1) + for _, sourceID := range baseProjection.Sources[targetID] { + if _, shadowed := shadowIDs[sourceID]; shadowed || v.layer.HasFile(IDFile(sourceID)) { + continue + } + seen[sourceID] = struct{}{} + if len(seen) > limit { + projection.Truncated[targetID] = true + break + } + } + if baseProjection.Truncated[targetID] || projection.Truncated[targetID] { + projection.Truncated[targetID] = true + continue + } + for index, edge := range v.layer.inEdges[targetID] { + if index&127 == 0 { + if err := ctx.Err(); err != nil { + return BoundedIncomingSourceProjection{}, err + } + } + if edge == nil || edge.Kind != kind || edge.From == "" { + continue + } + seen[edge.From] = struct{}{} + if len(seen) > limit { + projection.Truncated[targetID] = true + break + } + } + if projection.Truncated[targetID] { + continue + } + sources := make([]string, 0, len(seen)) + for sourceID := range seen { + sources = append(sources, sourceID) + } + sort.Strings(sources) + if len(sources) > 0 { + projection.Sources[targetID] = sources + } + } + return projection, nil +} diff --git a/internal/graph/bounded_incoming_sources_test.go b/internal/graph/bounded_incoming_sources_test.go new file mode 100644 index 00000000..c07fb3f8 --- /dev/null +++ b/internal/graph/bounded_incoming_sources_test.go @@ -0,0 +1,354 @@ +package graph + +import ( + "context" + "errors" + "fmt" + "reflect" + "sort" + "testing" +) + +type recordingBoundedIncomingReader struct { + Reader + bounded BoundedIncomingSourceReader + calls int + limits []int +} + +func (reader *recordingBoundedIncomingReader) FindIncomingSourcesBounded( + ctx context.Context, + targetIDs []string, + kind EdgeKind, + limit int, +) (BoundedIncomingSourceProjection, error) { + reader.calls++ + reader.limits = append(reader.limits, limit) + return reader.bounded.FindIncomingSourcesBounded(ctx, targetIDs, kind, limit) +} + +func TestGraphFindIncomingSourcesBoundedCountsDistinctRelevantSources(t *testing.T) { + memory := New() + const targetID = "repo/target.go::target" + for index := 0; index < 8; index++ { + sourceID := fmt.Sprintf("repo/source-%02d.go::source", index) + for line := 1; line <= 3; line++ { + memory.AddEdge(&Edge{From: sourceID, To: targetID, Kind: EdgeCalls, Line: line}) + } + memory.AddEdge(&Edge{From: fmt.Sprintf("repo/noise-%02d.go::noise", index), To: targetID, Kind: EdgeReferences}) + } + + page, err := memory.FindIncomingSourcesBounded(context.Background(), []string{targetID}, EdgeCalls, 8) + if err != nil { + t.Fatalf("bounded incoming projection: %v", err) + } + if page.Truncated[targetID] || len(page.Sources[targetID]) != 8 { + t.Fatalf("projection = %#v, want eight distinct CALLS sources", page) + } + if !sort.StringsAreSorted(page.Sources[targetID]) { + t.Fatalf("sources are not deterministic: %v", page.Sources[targetID]) + } + + memory.AddEdge(&Edge{From: "repo/source-08.go::source", To: targetID, Kind: EdgeCalls}) + page, err = memory.FindIncomingSourcesBounded(context.Background(), []string{targetID}, EdgeCalls, 8) + if err != nil { + t.Fatalf("saturated incoming projection: %v", err) + } + if !page.Truncated[targetID] || len(page.Sources[targetID]) != 0 { + t.Fatalf("saturated projection exposed a partial source set: %#v", page) + } +} + +func TestGraphFindIncomingSourcesBoundedRejectsImpossibleLimitAndCancellation(t *testing.T) { + memory := New() + memory.AddEdge(&Edge{From: "source", To: "target", Kind: EdgeCalls}) + maxInt := int(^uint(0) >> 1) + if page, err := memory.FindIncomingSourcesBounded(context.Background(), []string{"target"}, EdgeCalls, maxInt); err == nil || len(page.Sources) != 0 { + t.Fatalf("max-int projection = %#v, %v; want empty typed limit error", page, err) + } else { + var limitErr *BoundedLocalizationLimitError + if !errors.As(err, &limitErr) { + t.Fatalf("max-int error = %T %v, want BoundedLocalizationLimitError", err, err) + } + } + ctx, cancel := context.WithCancel(context.Background()) + cancel() + if page, err := memory.FindIncomingSourcesBounded(ctx, []string{"target"}, EdgeCalls, 8); !errors.Is(err, context.Canceled) || len(page.Sources) != 0 { + t.Fatalf("cancelled projection = %#v, %v", page, err) + } +} + +func TestOverlaidViewFindIncomingSourcesBoundedReappliesLimitAfterCompensation(t *testing.T) { + base := New() + const targetID = "repo/target.go::target" + for index := 0; index < 9; index++ { + base.AddEdge(&Edge{ + From: fmt.Sprintf("repo/visible-%02d.go::source", index), To: targetID, Kind: EdgeCalls, + }) + } + layer := NewOverlayLayer() + layer.MarkFile("repo/unrelated.go", false) + for index := 0; index < 300; index++ { + layer.AddNode("repo/unrelated.go", &Node{ + ID: fmt.Sprintf("repo/unrelated.go::shadow-%03d", index), Name: "shadow", + Kind: KindFunction, FilePath: "repo/unrelated.go", + }) + } + recording := &recordingBoundedIncomingReader{Reader: base, bounded: base} + page, err := NewOverlaidView(recording, layer).FindIncomingSourcesBounded( + context.Background(), []string{targetID}, EdgeCalls, 8, + ) + if err != nil { + t.Fatalf("compensated overlay projection: %v", err) + } + if !page.Truncated[targetID] || len(page.Sources[targetID]) != 0 { + t.Fatalf("compensation bypassed caller cap: %#v", page) + } + if !reflect.DeepEqual(recording.limits, []int{308}) { + t.Fatalf("base compensation limit = %v, want 308", recording.limits) + } +} + +func TestOverlaidViewFindIncomingSourcesBoundedSeparatesStandardAndDetachedShadows(t *testing.T) { + base := New() + const ( + targetID = "repo/target.go::target" + filePath = "repo/edited.go" + ) + layer := NewOverlayLayer() + layer.MarkFile(filePath, false) + for index := 0; index < 300; index++ { + id := fmt.Sprintf("%s::source-%03d", filePath, index) + base.AddNode(&Node{ID: id, Name: "source", Kind: KindFunction, FilePath: filePath}) + base.AddEdge(&Edge{From: id, To: targetID, Kind: EdgeCalls}) + layer.AddNode(filePath, &Node{ID: id, Name: "source", Kind: KindFunction, FilePath: filePath}) + } + currentID := filePath + "::source-000" + layer.AddEdge(&Edge{From: currentID, To: targetID, Kind: EdgeCalls}) + recording := &recordingBoundedIncomingReader{Reader: base, bounded: base} + page, err := NewOverlaidView(recording, layer).FindIncomingSourcesBounded( + context.Background(), []string{targetID}, EdgeCalls, 8, + ) + if err != nil { + t.Fatalf("ordinary overlay shadows above detached cap: %v", err) + } + if page.Truncated[targetID] || !reflect.DeepEqual(page.Sources[targetID], []string{currentID}) { + t.Fatalf("overlay replacement sources = %#v, want only current %q", page, currentID) + } + if recording.calls != 1 || len(recording.limits) != 1 || recording.limits[0] != 308 { + t.Fatalf("base calls/limits = %d/%v, want one exact compensation call at 308", recording.calls, recording.limits) + } + + for _, test := range []struct { + name string + shadows int + wantError bool + }{ + {name: "exact detached cap", shadows: overlayDetachedShadowLimit}, + {name: "above detached cap", shadows: overlayDetachedShadowLimit + 1, wantError: true}, + } { + t.Run(test.name, func(t *testing.T) { + empty := New() + counted := &recordingBoundedIncomingReader{Reader: empty, bounded: empty} + detached := NewOverlayLayer() + for index := 0; index < test.shadows; index++ { + detached.MarkRemoved("source", fmt.Sprintf("legacy-source-%03d", index)) + } + got, gotErr := NewOverlaidView(counted, detached).FindIncomingSourcesBounded( + context.Background(), []string{targetID}, EdgeCalls, 8, + ) + if test.wantError { + var limitErr *BoundedLocalizationLimitError + if !errors.As(gotErr, &limitErr) || + limitErr.Resource != "overlay incoming-source detached shadows" || + limitErr.Limit != overlayDetachedShadowLimit || counted.calls != 0 || + len(got.Sources) != 0 || len(got.Truncated) != 0 { + t.Fatalf("overflow = %#v, %v, base calls %d; want exact pre-base detached limit failure", got, gotErr, counted.calls) + } + return + } + if gotErr != nil || counted.calls != 1 || counted.limits[0] != 8+overlayDetachedShadowLimit { + t.Fatalf("exact-cap projection = %#v, %v, calls/limits %d/%v", got, gotErr, counted.calls, counted.limits) + } + }) + } +} + +func TestOverlaidViewFindIncomingSourcesBoundedHonorsDetachedSourceAndTargetState(t *testing.T) { + base := New() + const ( + targetID = "legacy-target" + staleID = "legacy-stale-source" + freshID = "legacy-fresh-source" + ) + base.AddNode(&Node{ID: targetID, Name: "target", Kind: KindFunction}) + base.AddNode(&Node{ID: staleID, Name: "stale", Kind: KindFunction}) + base.AddEdge(&Edge{From: staleID, To: targetID, Kind: EdgeCalls}) + + tombstoneSource := NewOverlayLayer() + tombstoneSource.MarkRemoved("stale", staleID) + page, err := NewOverlaidView(base, tombstoneSource).FindIncomingSourcesBounded( + context.Background(), []string{targetID}, EdgeCalls, 8, + ) + if err != nil || len(page.Sources[targetID]) != 0 { + t.Fatalf("detached source tombstone leaked stale caller: %#v, %v", page, err) + } + + tombstoneTarget := NewOverlayLayer() + tombstoneTarget.MarkRemoved("target", targetID) + page, err = NewOverlaidView(base, tombstoneTarget).FindIncomingSourcesBounded( + context.Background(), []string{targetID}, EdgeCalls, 8, + ) + if err != nil || len(page.Sources[targetID]) != 0 { + t.Fatalf("detached target tombstone leaked base adjacency: %#v, %v", page, err) + } + + replacement := NewOverlayLayer() + replacement.MarkRemoved("target", targetID) + replacement.AddNode("repo/replacement.go", &Node{ID: targetID, Name: "target", Kind: KindFunction, FilePath: "repo/replacement.go"}) + replacement.AddNode("repo/caller.go", &Node{ID: freshID, Name: "fresh", Kind: KindFunction, FilePath: "repo/caller.go"}) + replacement.AddEdge(&Edge{From: freshID, To: targetID, Kind: EdgeCalls}) + page, err = NewOverlaidView(base, replacement).FindIncomingSourcesBounded( + context.Background(), []string{targetID}, EdgeCalls, 8, + ) + if err != nil || !reflect.DeepEqual(page.Sources[targetID], []string{freshID, staleID}) { + t.Fatalf("detached target replacement callers = %#v, %v; want current overlay plus unaffected base caller", page, err) + } +} + +func TestOverlaidViewGetNodesByIDsContextPreservesInputAndDetachedMasks(t *testing.T) { + base := New() + base.AddNode(&Node{ID: "base", Name: "base", Kind: KindFunction}) + base.AddNode(&Node{ID: "removed", Name: "removed", Kind: KindFunction}) + layer := NewOverlayLayer() + layer.MarkRemoved("removed", "removed") + layer.AddNode("repo/overlay.go", &Node{ID: "overlay", Name: "overlay", Kind: KindFunction, FilePath: "repo/overlay.go"}) + view := NewOverlaidView(base, layer) + ids := []string{"overlay", "base", "removed", "base"} + before := append([]string(nil), ids...) + nodes, err := view.GetNodesByIDsContext(context.Background(), ids) + if err != nil { + t.Fatalf("contextual overlay refetch: %v", err) + } + if !reflect.DeepEqual(ids, before) { + t.Fatalf("caller IDs mutated: got %v want %v", ids, before) + } + if nodes["overlay"] == nil || nodes["base"] == nil || nodes["removed"] != nil { + t.Fatalf("contextual overlay refetch = %#v", nodes) + } + ctx, cancel := context.WithCancel(context.Background()) + cancel() + if cancelled, err := view.GetNodesByIDsContext(ctx, ids); !errors.Is(err, context.Canceled) || len(cancelled) != 0 { + t.Fatalf("cancelled contextual refetch = %#v, %v", cancelled, err) + } +} + +func TestBoundedIncomingSourcesCancelsDuringGraphAndOverlayInspection(t *testing.T) { + memory := New() + for index := 0; index < 256; index++ { + memory.AddEdge(&Edge{ + From: fmt.Sprintf("source-%03d", index), To: "target", Kind: EdgeCalls, + }) + } + ctx := &cancelAfterLocalizationChecksContext{ + Context: context.Background(), remaining: 3, done: make(chan struct{}), + } + page, err := memory.FindIncomingSourcesBounded(ctx, []string{"target"}, EdgeCalls, 512) + if !errors.Is(err, context.Canceled) || len(page.Sources) != 0 || len(page.Truncated) != 0 { + t.Fatalf("mid-graph cancellation = %#v, %v", page, err) + } + + base := New() + recording := &recordingBoundedIncomingReader{Reader: base, bounded: base} + layer := NewOverlayLayer() + layer.MarkFile("repo/edited.go", false) + for index := 0; index < 512; index++ { + layer.MarkRemoved("source", fmt.Sprintf("repo/edited.go::source-%03d", index)) + } + ctx = &cancelAfterLocalizationChecksContext{ + Context: context.Background(), remaining: 3, done: make(chan struct{}), + } + page, err = NewOverlaidView(recording, layer).FindIncomingSourcesBounded( + ctx, []string{"target"}, EdgeCalls, 8, + ) + if !errors.Is(err, context.Canceled) || len(page.Sources) != 0 || len(page.Truncated) != 0 || recording.calls != 0 { + t.Fatalf("mid-overlay cancellation = %#v, %v, base calls=%d", page, err, recording.calls) + } +} + +func TestOverlaidViewFindIncomingSourcesBoundedDedupesShadowsAndGuardsCompensation(t *testing.T) { + base := New() + recording := &recordingBoundedIncomingReader{Reader: base, bounded: base} + layer := NewOverlayLayer() + layer.MarkRemoved("first", "legacy-source") + layer.MarkRemoved("second", "legacy-source") + layer.AddNode("repo/overlay.go", &Node{ + ID: "legacy-source", Name: "replacement", Kind: KindFunction, FilePath: "repo/overlay.go", + }) + page, err := NewOverlaidView(recording, layer).FindIncomingSourcesBounded( + context.Background(), []string{"target"}, EdgeCalls, 8, + ) + if err != nil || len(page.Sources) != 0 || recording.calls != 1 || !reflect.DeepEqual(recording.limits, []int{9}) { + t.Fatalf("deduped shadow projection = %#v, %v, calls/limits=%d/%v", page, err, recording.calls, recording.limits) + } + + overflowBase := New() + overflowRecording := &recordingBoundedIncomingReader{Reader: overflowBase, bounded: overflowBase} + overflowLayer := NewOverlayLayer() + overflowLayer.MarkRemoved("first", "legacy-one") + overflowLayer.MarkRemoved("second", "legacy-two") + maxInt := int(^uint(0) >> 1) + page, err = NewOverlaidView(overflowRecording, overflowLayer).FindIncomingSourcesBounded( + context.Background(), []string{"target"}, EdgeCalls, maxInt-1, + ) + var limitErr *BoundedLocalizationLimitError + if !errors.As(err, &limitErr) || limitErr.Resource != "overlay incoming-source compensation" || + overflowRecording.calls != 0 || len(page.Sources) != 0 || len(page.Truncated) != 0 { + t.Fatalf("compensation overflow = %#v, %v, base calls=%d", page, err, overflowRecording.calls) + } +} + +type contextualExactNodeSpy struct { + Reader + nodes map[string]*Node + err error + calls int + ctx context.Context +} + +func (spy *contextualExactNodeSpy) GetNodesByIDs([]string) map[string]*Node { + panic("non-contextual exact refetch must not be used") +} + +func (spy *contextualExactNodeSpy) GetNodesByIDsContext(ctx context.Context, ids []string) (map[string]*Node, error) { + spy.calls++ + spy.ctx = ctx + out := make(map[string]*Node, len(ids)) + for _, id := range ids { + if node := spy.nodes[id]; node != nil { + out[id] = node + } + } + return out, spy.err +} + +func TestOverlaidViewGetNodesByIDsContextDelegatesContextAndDropsPartialErrors(t *testing.T) { + type contextKey string + const key contextKey = "request" + requestCtx := context.WithValue(context.Background(), key, "bounded") + spy := &contextualExactNodeSpy{ + nodes: map[string]*Node{"base": {ID: "base", Name: "base", Kind: KindFunction}}, + } + view := NewOverlaidView(spy, NewOverlayLayer()) + nodes, err := view.GetNodesByIDsContext(requestCtx, []string{"base"}) + if err != nil || nodes["base"] == nil || spy.calls != 1 || spy.ctx.Value(key) != "bounded" { + t.Fatalf("contextual delegation = %#v, %v, calls=%d, ctx=%v", nodes, err, spy.calls, spy.ctx) + } + + spy.err = errors.New("base refetch failed") + nodes, err = view.GetNodesByIDsContext(requestCtx, []string{"base"}) + if err == nil || len(nodes) != 0 { + t.Fatalf("base error returned partial exact nodes: %#v, %v", nodes, err) + } +} + diff --git a/internal/graph/localization_projection.go b/internal/graph/localization_projection.go new file mode 100644 index 00000000..a7e2aa76 --- /dev/null +++ b/internal/graph/localization_projection.go @@ -0,0 +1,565 @@ +package graph + +import ( + "context" + "errors" + "fmt" + "sort" +) + +// LocalizationNodeScope is the storage-layer form of a request scope. It is +// intentionally independent of query.QueryOptions so graph does not import the +// query package. Empty fields mean no corresponding restriction. +type LocalizationNodeScope struct { + WorkspaceID string + ProjectID string + RepoAllow map[string]bool + // Kinds is an optional declaration-kind pushdown. It keeps homonym counts + // meaningful for a caller that can only localize selected node classes. + Kinds map[NodeKind]bool + // ExcludeKinds is the complementary pushdown for open-ended declaration + // sets. It lets callers exclude structural/body-internal kinds without a + // brittle allow-list that would hide future definition kinds. + ExcludeKinds map[NodeKind]bool + // ExcludeTests preserves localization's production-anchor gate. SQLite + // evaluates the legacy is_test metadata while keyset-paging bounded rows. + ExcludeTests bool + // ExcludeFiles applies caller-owned whole-file exclusions before LIMIT. + // SQLite evaluates it between bounded keyset pages, avoiding a variable-sized + // NOT IN clause and SQLite variable limits. + ExcludeFiles map[string]bool + + // excludeFile is a request-local immutable view over an overlay's covered + // paths. Keeping it separate avoids cloning every overlay file into + // ExcludeFiles for each exact-name lookup. It composes through + // withFileExcluder and is deliberately unexported: callers own ExcludeFiles. + excludeFile func(string) bool +} + +// ExcludesFile applies caller and immutable overlay exclusions consistently in +// the in-memory graph, SQLite keyset projections, and overlay composition. +func (s LocalizationNodeScope) ExcludesFile(path string) bool { + return s.ExcludeFiles[path] || s.excludeFile != nil && s.excludeFile(path) +} + +func (s LocalizationNodeScope) withFileExcluder(exclude func(string) bool) LocalizationNodeScope { + if exclude == nil { + return s + } + if s.excludeFile == nil { + s.excludeFile = exclude + return s + } + prior := s.excludeFile + s.excludeFile = func(path string) bool { return prior(path) || exclude(path) } + return s +} + +// Allows applies the same effective workspace/project and repository rules as +// query.QueryOptions.ScopeAllows. Keep this predicate in lock-step with that +// method: SQLite projections use the same rules in SQL while the in-memory and +// overlay readers use this implementation. +func (s LocalizationNodeScope) Allows(n *Node) bool { + if n == nil { + return false + } + if len(s.Kinds) > 0 && !s.Kinds[n.Kind] { + return false + } + if s.ExcludeKinds[n.Kind] { + return false + } + if s.ExcludeTests { + if isTest, _ := n.Meta["is_test"].(bool); isTest { + return false + } + } + path := n.FilePath + if path == "" { + path = IDFile(n.ID) + } + if s.ExcludesFile(path) { + return false + } + if s.WorkspaceID != "" { + workspace := n.WorkspaceID + if workspace == "" { + workspace = n.RepoPrefix + } + if workspace != s.WorkspaceID { + return false + } + if s.ProjectID != "" { + project := n.ProjectID + if project == "" { + project = n.RepoPrefix + } + if project != s.ProjectID { + return false + } + } + } + // Empty RepoPrefix nodes are global synthetic externals and remain visible + // under a repository narrow, matching QueryOptions.ScopeAllows. + return len(s.RepoAllow) == 0 || n.RepoPrefix == "" || s.RepoAllow[n.RepoPrefix] +} + +// BoundedNodeProjection is a deterministic, request-bounded node page. Total +// is the number of admitted rows observed up to limit+1, not a corpus-wide +// count. Truncated means that sentinel was reached. This saturation contract is +// sufficient for ambiguity decisions while bounding both memory and row work. +type BoundedNodeProjection struct { + Nodes []*Node + Total int + Truncated bool +} + +// BoundedExactNameReader is the localization exact-name projection. It is an +// optional capability instead of part of Reader so callers cannot silently +// fall back to the legacy unbounded, full-row FindNodesByName path. +type BoundedExactNameReader interface { + FindNodesByNameBounded(context.Context, string, LocalizationNodeScope, int) (BoundedNodeProjection, error) +} + +// BoundedFileNodeReader projects one file's scoped identity, kind, and source +// location rows without transferring retrieval payloads. It is optional rather +// than part of Reader so localization callers can fail closed instead of +// silently falling back to the legacy full-row GetFileNodes path. +type BoundedFileNodeReader interface { + FindFileNodesBounded(context.Context, string, LocalizationNodeScope, int) (BoundedNodeProjection, error) +} + +var ( + _ BoundedExactNameReader = (*Graph)(nil) + _ BoundedFileNodeReader = (*Graph)(nil) +) + +// FindNodesByNameBounded reads at most limit+1 matching pointers while still +// counting every scoped homonym. The bounded sorted insertion keeps memory +// proportional to the response cap even for names shared by tens of thousands +// of declarations. This legacy in-memory backend takes one read lock per shard, +// so concurrent mutation may produce a weak cross-shard snapshot; each returned +// page still obeys its cap and Total is never smaller than len(Nodes). SQLite, +// the release backend, provides a single read-transaction snapshot. +func (g *Graph) FindNodesByNameBounded( + ctx context.Context, + name string, + scope LocalizationNodeScope, + limit int, +) (BoundedNodeProjection, error) { + if g == nil || name == "" || limit <= 0 { + return BoundedNodeProjection{}, nil + } + if ctx == nil { + ctx = context.Background() + } + if err := ctx.Err(); err != nil { + return BoundedNodeProjection{}, err + } + + pageSize := limit + 1 + kept := make([]*Node, 0, pageSize) + total := 0 + for _, shard := range g.shards { + if err := ctx.Err(); err != nil { + return BoundedNodeProjection{}, err + } + shard.mu.RLock() + for i, node := range shard.byName[name] { + if i&127 == 0 { + if err := ctx.Err(); err != nil { + shard.mu.RUnlock() + return BoundedNodeProjection{}, err + } + } + if !scope.Allows(node) { + continue + } + total++ + kept = insertBoundedLocalizationNode(kept, node, pageSize) + } + shard.mu.RUnlock() + } + + truncated := len(kept) > limit || total > limit + if len(kept) > limit { + kept = kept[:limit] + } + if total > pageSize { + total = pageSize + } + return BoundedNodeProjection{Nodes: kept, Total: total, Truncated: truncated}, nil +} + +// FindFileNodesBounded is the in-memory reference implementation of the +// lightweight file projection. It retains only limit+1 pointers while scanning +// the file buckets, applies scope before the cap, and returns a deterministic +// ID-ordered page. As with FindNodesByNameBounded, concurrent writes can yield a +// weak cross-shard snapshot; the page invariants remain intact. +func (g *Graph) FindFileNodesBounded( + ctx context.Context, + filePath string, + scope LocalizationNodeScope, + limit int, +) (BoundedNodeProjection, error) { + if g == nil || filePath == "" || limit <= 0 { + return BoundedNodeProjection{}, nil + } + if ctx == nil { + ctx = context.Background() + } + if err := ctx.Err(); err != nil { + return BoundedNodeProjection{}, err + } + + pageSize := limit + 1 + kept := make([]*Node, 0, pageSize) + total := 0 + for _, shard := range g.shards { + if err := ctx.Err(); err != nil { + return BoundedNodeProjection{}, err + } + shard.mu.RLock() + for index, node := range shard.byFile[filePath] { + if index&127 == 0 { + if err := ctx.Err(); err != nil { + shard.mu.RUnlock() + return BoundedNodeProjection{}, err + } + } + if !scope.Allows(node) { + continue + } + total++ + kept = insertBoundedLocalizationNode(kept, node, pageSize) + } + shard.mu.RUnlock() + } + + truncated := len(kept) > limit || total > limit + if len(kept) > limit { + kept = kept[:limit] + } + kept = localizationNodeSummaries(kept) + if total > pageSize { + total = pageSize + } + return BoundedNodeProjection{Nodes: kept, Total: total, Truncated: truncated}, nil +} + +func localizationNodeSummary(node *Node) *Node { + if node == nil { + return nil + } + return &Node{ + ID: node.ID, Kind: node.Kind, Name: node.Name, QualName: node.QualName, + FilePath: node.FilePath, StartLine: node.StartLine, EndLine: node.EndLine, + StartColumn: node.StartColumn, EndColumn: node.EndColumn, Language: node.Language, + RepoPrefix: node.RepoPrefix, WorkspaceID: node.WorkspaceID, ProjectID: node.ProjectID, + } +} + +// localizationNodeSummaries copies only the retained response page onto a fresh, +// exact-capacity backing array. Callers first rank raw immutable graph pointers, +// so discarded candidates never allocate summary objects, and a saturated raw +// sentinel cannot remain reachable behind the returned slice length. +func localizationNodeSummaries(nodes []*Node) []*Node { + summaries := make([]*Node, len(nodes)) + for index, node := range nodes { + summaries[index] = localizationNodeSummary(node) + } + return summaries +} + +func insertBoundedLocalizationNode(nodes []*Node, node *Node, limit int) []*Node { + if node == nil || limit <= 0 { + return nodes + } + position := sort.Search(len(nodes), func(i int) bool { return nodes[i].ID >= node.ID }) + if position < len(nodes) && nodes[position].ID == node.ID { + return nodes + } + if position >= limit { + return nodes + } + nodes = append(nodes, nil) + copy(nodes[position+1:], nodes[position:]) + nodes[position] = node + if len(nodes) > limit { + nodes = nodes[:limit] + } + return nodes +} + +const ( + overlayExactNameInspectionLimit = 4096 + overlayDetachedShadowLimit = 256 +) + +// BoundedLocalizationLimitError reports that a bounded projection cannot +// preserve its completeness contract within a hard work or allocation limit. +// Callers must not use a partial result when this error is returned. +type BoundedLocalizationLimitError struct { + Resource string + Limit int +} + +func (e *BoundedLocalizationLimitError) Error() string { + return fmt.Sprintf("bounded localization %s exceeds limit %d", e.Resource, e.Limit) +} + +var ( + _ BoundedExactNameReader = (*OverlaidView)(nil) + _ BoundedFileNodeReader = (*OverlaidView)(nil) + // ErrBoundedLocalizationUnavailable lets MCP localization fail closed when + // a third-party Reader has not implemented the bounded projection. Falling + // back to FindNodesByName would silently restore the unbounded allocation. + ErrBoundedLocalizationUnavailable = errors.New("bounded localization projection unavailable") +) + +// FindNodesByNameBounded merges a bounded base page with the request overlay. +// It inflates the base cap only by the exact set of base identities the layer +// can shadow, so filtering a replacement/tombstone cannot leave a short page. +func (v *OverlaidView) FindNodesByNameBounded( + ctx context.Context, + name string, + scope LocalizationNodeScope, + limit int, +) (BoundedNodeProjection, error) { + if v == nil || name == "" || limit <= 0 { + return BoundedNodeProjection{}, nil + } + if ctx == nil { + ctx = context.Background() + } + if err := ctx.Err(); err != nil { + return BoundedNodeProjection{}, err + } + + var baseReader BoundedExactNameReader + if v.base != nil { + var ok bool + baseReader, ok = v.base.(BoundedExactNameReader) + if !ok { + return BoundedNodeProjection{}, ErrBoundedLocalizationUnavailable + } + } + + removedCount, overlayNodeCount := 0, 0 + if v.layer != nil { + removedCount = len(v.layer.nameRemoved[name]) + overlayNodeCount = len(v.layer.nodesByName[name]) + } + if removedCount > overlayExactNameInspectionLimit || + overlayNodeCount > overlayExactNameInspectionLimit-removedCount { + return BoundedNodeProjection{}, &BoundedLocalizationLimitError{ + Resource: "overlay exact-name entries", + Limit: overlayExactNameInspectionLimit, + } + } + + const maxInt = int(^uint(0) >> 1) + if limit == maxInt { + return BoundedNodeProjection{}, &BoundedLocalizationLimitError{ + Resource: "exact-name page capacity", + Limit: maxInt - 1, + } + } + pageSize := limit + 1 + kept := make([]*Node, 0, pageSize) + overlayCount := 0 + + // Exclude whole-file replacements and tombstones in the base reader without + // cloning a potentially repository-sized entries map into every request. + baseScope := scope + if v.layer != nil { + baseScope = baseScope.withFileExcluder(v.layer.HasFile) + } + + var detachedShadowIDs map[string]struct{} + shadowCapacity := removedCount + overlayNodeCount + if shadowCapacity > overlayDetachedShadowLimit { + shadowCapacity = overlayDetachedShadowLimit + } + addDetachedShadow := func(id, path string) error { + if baseReader == nil || id == "" || baseScope.ExcludesFile(path) { + return nil + } + if _, exists := detachedShadowIDs[id]; exists { + return nil + } + if len(detachedShadowIDs) >= overlayDetachedShadowLimit { + return &BoundedLocalizationLimitError{ + Resource: "detached overlay shadow identities", + Limit: overlayDetachedShadowLimit, + } + } + if detachedShadowIDs == nil { + detachedShadowIDs = make(map[string]struct{}, shadowCapacity) + } + detachedShadowIDs[id] = struct{}{} + return nil + } + + inspected := 0 + checkInspection := func() error { + inspected++ + if inspected&127 == 0 { + return ctx.Err() + } + return nil + } + if v.layer != nil { + for id := range v.layer.nameRemoved[name] { + if err := checkInspection(); err != nil { + return BoundedNodeProjection{}, err + } + if err := addDetachedShadow(id, IDFile(id)); err != nil { + return BoundedNodeProjection{}, err + } + } + for _, node := range v.layer.nodesByName[name] { + if err := checkInspection(); err != nil { + return BoundedNodeProjection{}, err + } + if node == nil { + continue + } + path := node.FilePath + if path == "" { + path = IDFile(node.ID) + } + if err := addDetachedShadow(node.ID, path); err != nil { + return BoundedNodeProjection{}, err + } + if !scope.Allows(node) { + continue + } + overlayCount++ + kept = insertBoundedLocalizationNode(kept, node, pageSize) + } + } + if err := ctx.Err(); err != nil { + return BoundedNodeProjection{}, err + } + + if baseReader == nil { + truncated := overlayCount > limit + if len(kept) > limit { + kept = kept[:limit] + } + total := overlayCount + if total > pageSize { + total = pageSize + } + kept = kept[:len(kept):len(kept)] + return BoundedNodeProjection{Nodes: kept, Total: total, Truncated: truncated}, nil + } + if limit > maxInt-len(detachedShadowIDs) { + return BoundedNodeProjection{}, &BoundedLocalizationLimitError{ + Resource: "exact-name shadow compensation", + Limit: maxInt - limit, + } + } + baseLimit := limit + len(detachedShadowIDs) + basePage, err := baseReader.FindNodesByNameBounded(ctx, name, baseScope, baseLimit) + if err != nil { + return BoundedNodeProjection{}, err + } + + visible := overlayCount + for index, node := range basePage.Nodes { + if index&127 == 0 { + if err := ctx.Err(); err != nil { + return BoundedNodeProjection{}, err + } + } + if node == nil { + continue + } + if _, shadowed := detachedShadowIDs[node.ID]; shadowed { + continue + } + path := node.FilePath + if path == "" { + path = IDFile(node.ID) + } + if baseScope.ExcludesFile(path) { + continue + } + visible++ + kept = insertBoundedLocalizationNode(kept, node, pageSize) + } + // Saturation of the shadow-inflated base page proves at least limit+1 + // visible rows: no more than len(detachedShadowIDs) returned identities can + // vanish, while whole-file shadows were excluded inside the base reader. + if basePage.Truncated && visible <= limit { + visible = limit + 1 + } + if visible > pageSize { + visible = pageSize + } + truncated := visible > limit + if len(kept) > limit { + kept = kept[:limit] + } + kept = kept[:len(kept):len(kept)] + return BoundedNodeProjection{Nodes: kept, Total: visible, Truncated: truncated}, nil +} + +// FindFileNodesBounded preserves file-overlay replacement semantics: an +// overlaid file is answered exclusively from its request-local replacement (or +// as empty for a tombstone), while an untouched file delegates to the bounded +// base capability. It never merges stale base declarations into an overlay. +func (v *OverlaidView) FindFileNodesBounded( + ctx context.Context, + filePath string, + scope LocalizationNodeScope, + limit int, +) (BoundedNodeProjection, error) { + if v == nil || filePath == "" || limit <= 0 { + return BoundedNodeProjection{}, nil + } + if ctx == nil { + ctx = context.Background() + } + if err := ctx.Err(); err != nil { + return BoundedNodeProjection{}, err + } + + if v.layer != nil && v.layer.HasFile(filePath) { + pageSize := limit + 1 + kept := make([]*Node, 0, pageSize) + total := 0 + for index, node := range v.layer.nodesForFileReadOnly(filePath) { + if index&127 == 0 { + if err := ctx.Err(); err != nil { + return BoundedNodeProjection{}, err + } + } + if !scope.Allows(node) { + continue + } + total++ + kept = insertBoundedLocalizationNode(kept, node, pageSize) + } + truncated := total > limit + if len(kept) > limit { + kept = kept[:limit] + } + kept = localizationNodeSummaries(kept) + if total > pageSize { + total = pageSize + } + return BoundedNodeProjection{Nodes: kept, Total: total, Truncated: truncated}, nil + } + + if v.base == nil { + return BoundedNodeProjection{}, nil + } + baseReader, ok := v.base.(BoundedFileNodeReader) + if !ok { + return BoundedNodeProjection{}, ErrBoundedLocalizationUnavailable + } + return baseReader.FindFileNodesBounded(ctx, filePath, scope, limit) +} diff --git a/internal/graph/localization_projection_test.go b/internal/graph/localization_projection_test.go new file mode 100644 index 00000000..5eb3c0e7 --- /dev/null +++ b/internal/graph/localization_projection_test.go @@ -0,0 +1,658 @@ +package graph + +import ( + "context" + "errors" + "fmt" + "sort" + "sync" + "testing" +) + +func TestFindNodesByNameBoundedCapsAndCancels(t *testing.T) { + graph := New() + for index := 0; index < 128; index++ { + graph.AddNode(&Node{ + ID: fmt.Sprintf("repo/file-%03d.go::handle", index), + Name: "handle", + Kind: KindFunction, + FilePath: fmt.Sprintf("repo/file-%03d.go", index), + }) + } + + page, err := graph.FindNodesByNameBounded( + context.Background(), "handle", + LocalizationNodeScope{Kinds: map[NodeKind]bool{KindFunction: true}}, + 8, + ) + if err != nil { + t.Fatalf("bounded lookup: %v", err) + } + if page.Total != 9 || !page.Truncated || len(page.Nodes) != 8 { + t.Fatalf("page = %#v, want threshold total 9, truncated, cap 8", page) + } + if !sort.SliceIsSorted(page.Nodes, func(i, j int) bool { return page.Nodes[i].ID < page.Nodes[j].ID }) { + t.Fatalf("nodes are not deterministic: %#v", page.Nodes) + } + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + if cancelled, err := graph.FindNodesByNameBounded(ctx, "handle", LocalizationNodeScope{}, 8); err == nil || len(cancelled.Nodes) != 0 { + t.Fatalf("cancelled lookup = %#v, %v; want empty error result", cancelled, err) + } +} + +func TestFindNodesByNameBoundedWeakSnapshotKeepsPageInvariants(t *testing.T) { + graph := New() + var writer sync.WaitGroup + writer.Add(1) + go func() { + defer writer.Done() + for index := 0; index < 2_000; index++ { + graph.AddNode(&Node{ + ID: fmt.Sprintf("repo/live-%04d.go::handle", index), + Name: "handle", + Kind: KindFunction, + FilePath: fmt.Sprintf("repo/live-%04d.go", index), + }) + } + }() + + for attempt := 0; attempt < 100; attempt++ { + page, err := graph.FindNodesByNameBounded(context.Background(), "handle", LocalizationNodeScope{}, 16) + if err != nil { + t.Fatalf("bounded lookup: %v", err) + } + if len(page.Nodes) > 16 || page.Total < len(page.Nodes) { + t.Fatalf("weak snapshot violated bounds: %#v", page) + } + for index := 1; index < len(page.Nodes); index++ { + if page.Nodes[index-1].ID >= page.Nodes[index].ID { + t.Fatalf("weak snapshot is not sorted/deduplicated: %#v", page.Nodes) + } + } + } + writer.Wait() +} + +func TestOverlaidViewFindNodesByNameBoundedRefillsAfterFileTombstone(t *testing.T) { + base := New() + for index := 0; index < 32; index++ { + base.AddNode(&Node{ + ID: fmt.Sprintf("repo/a-hidden.go::handle:%02d", index), Name: "handle", + Kind: KindFunction, FilePath: "repo/a-hidden.go", + }) + } + visible := &Node{ + ID: "repo/z-visible.go::handle", Name: "handle", Kind: KindFunction, + FilePath: "repo/z-visible.go", + } + base.AddNode(visible) + + layer := NewOverlayLayer() + layer.MarkFile("repo/a-hidden.go", true) + view := NewOverlaidView(base, layer) + + page, err := view.FindNodesByNameBounded(context.Background(), "handle", LocalizationNodeScope{}, 8) + if err != nil { + t.Fatalf("bounded overlay lookup: %v", err) + } + if page.Total != 1 || page.Truncated || len(page.Nodes) != 1 || page.Nodes[0].ID != visible.ID { + t.Fatalf("page = %#v, want the later visible homonym after file tombstone", page) + } +} + +type recordingBoundedExactNameReader struct { + Reader + bounded BoundedExactNameReader + limits []int +} + +func (reader *recordingBoundedExactNameReader) FindNodesByNameBounded( + ctx context.Context, + name string, + scope LocalizationNodeScope, + limit int, +) (BoundedNodeProjection, error) { + reader.limits = append(reader.limits, limit) + return reader.bounded.FindNodesByNameBounded(ctx, name, scope, limit) +} + +func TestOverlaidViewFindNodesByNameBoundedDoesNotInflateWholeFileShadows(t *testing.T) { + base := New() + visible := &Node{ + ID: "repo/visible.go::handle", Name: "handle", Kind: KindFunction, + FilePath: "repo/visible.go", + } + base.AddNode(visible) + recording := &recordingBoundedExactNameReader{Reader: base, bounded: base} + + layer := NewOverlayLayer() + layer.MarkFile("repo/generated.go", false) + for index := 0; index < overlayExactNameInspectionLimit; index++ { + layer.MarkRemoved("handle", fmt.Sprintf("repo/generated.go::handle:%04d", index)) + } + view := NewOverlaidView(recording, layer) + + page, err := view.FindNodesByNameBounded(context.Background(), "handle", LocalizationNodeScope{}, 8) + if err != nil { + t.Fatalf("bounded overlay lookup: %v", err) + } + if len(recording.limits) != 1 || recording.limits[0] != 8 { + t.Fatalf("base limits = %v, want unchanged request limit 8", recording.limits) + } + if page.Total != 1 || page.Truncated || len(page.Nodes) != 1 || page.Nodes[0].ID != visible.ID { + t.Fatalf("page = %#v, want the one visible base homonym", page) + } +} + +func TestOverlaidViewFindNodesByNameBoundedFailsClosedAboveInspectionLimit(t *testing.T) { + base := New() + recording := &recordingBoundedExactNameReader{Reader: base, bounded: base} + layer := NewOverlayLayer() + layer.MarkFile("repo/generated.go", false) + for index := 0; index <= overlayExactNameInspectionLimit; index++ { + layer.MarkRemoved("handle", fmt.Sprintf("repo/generated.go::handle:%04d", index)) + } + + page, err := NewOverlaidView(recording, layer).FindNodesByNameBounded( + context.Background(), "handle", LocalizationNodeScope{}, 8, + ) + var limitErr *BoundedLocalizationLimitError + if !errors.As(err, &limitErr) || limitErr.Resource != "overlay exact-name entries" || + limitErr.Limit != overlayExactNameInspectionLimit { + t.Fatalf("error = %v, want typed inspection limit error", err) + } + if len(page.Nodes) != 0 || page.Total != 0 || page.Truncated { + t.Fatalf("overflow returned a partial page: %#v", page) + } + if len(recording.limits) != 0 { + t.Fatalf("overflow called base with limits %v", recording.limits) + } +} + +func TestOverlaidViewFindNodesByNameBoundedAllowsDetachedShadowLimit(t *testing.T) { + base := New() + layer := NewOverlayLayer() + for index := 0; index < overlayDetachedShadowLimit; index++ { + stale := &Node{ + ID: fmt.Sprintf("repo/a-old-%03d.go::handle", index), Name: "handle", + Kind: KindFunction, FilePath: fmt.Sprintf("repo/a-old-%03d.go", index), + } + base.AddNode(stale) + layer.MarkRemoved(stale.Name, stale.ID) + } + visible := &Node{ + ID: "repo/z-visible.go::handle", Name: "handle", Kind: KindFunction, + FilePath: "repo/z-visible.go", + } + base.AddNode(visible) + recording := &recordingBoundedExactNameReader{Reader: base, bounded: base} + + page, err := NewOverlaidView(recording, layer).FindNodesByNameBounded( + context.Background(), "handle", LocalizationNodeScope{}, 8, + ) + if err != nil { + t.Fatalf("bounded overlay lookup: %v", err) + } + if len(recording.limits) != 1 || recording.limits[0] != 8+overlayDetachedShadowLimit { + t.Fatalf("base limits = %v, want exact detached-shadow compensation", recording.limits) + } + if page.Total != 1 || page.Truncated || len(page.Nodes) != 1 || page.Nodes[0].ID != visible.ID { + t.Fatalf("page = %#v, want only the visible base homonym", page) + } + if cap(page.Nodes) != len(page.Nodes) { + t.Fatalf("page capacity = %d, want exact returned length %d", cap(page.Nodes), len(page.Nodes)) + } +} + +func TestOverlaidViewFindNodesByNameBoundedFailsClosedAboveDetachedShadowLimit(t *testing.T) { + base := New() + recording := &recordingBoundedExactNameReader{Reader: base, bounded: base} + layer := NewOverlayLayer() + for index := 0; index <= overlayDetachedShadowLimit; index++ { + layer.MarkRemoved("handle", fmt.Sprintf("repo/old-%03d.go::handle", index)) + } + + page, err := NewOverlaidView(recording, layer).FindNodesByNameBounded( + context.Background(), "handle", LocalizationNodeScope{}, 8, + ) + var limitErr *BoundedLocalizationLimitError + if !errors.As(err, &limitErr) || limitErr.Resource != "detached overlay shadow identities" || + limitErr.Limit != overlayDetachedShadowLimit { + t.Fatalf("error = %v, want typed detached-shadow limit error", err) + } + if len(page.Nodes) != 0 || page.Total != 0 || page.Truncated { + t.Fatalf("overflow returned a partial page: %#v", page) + } + if len(recording.limits) != 0 { + t.Fatalf("overflow called base with limits %v", recording.limits) + } +} + +func TestOverlaidViewFindNodesByNameBoundedCancelsDuringOverlayInspection(t *testing.T) { + base := New() + recording := &recordingBoundedExactNameReader{Reader: base, bounded: base} + layer := NewOverlayLayer() + layer.MarkFile("repo/generated.go", false) + for index := 0; index < 512; index++ { + layer.MarkRemoved("handle", fmt.Sprintf("repo/generated.go::handle:%04d", index)) + } + ctx := &cancelAfterLocalizationChecksContext{ + Context: context.Background(), remaining: 3, done: make(chan struct{}), + } + + page, err := NewOverlaidView(recording, layer).FindNodesByNameBounded( + ctx, "handle", LocalizationNodeScope{}, 8, + ) + if err != context.Canceled { + t.Fatalf("error = %v, want context cancellation during overlay inspection", err) + } + if len(page.Nodes) != 0 || page.Total != 0 || page.Truncated { + t.Fatalf("cancelled inspection returned a partial page: %#v", page) + } + if len(recording.limits) != 0 { + t.Fatalf("cancelled inspection called base with limits %v", recording.limits) + } +} + +func TestOverlaidViewFindNodesByNameBoundedComposesExclusionsWithoutMutatingCaller(t *testing.T) { + base := New() + for _, filePath := range []string{ + "repo/caller-hidden.go", "repo/inner-hidden.go", "repo/outer-hidden.go", "repo/visible.go", + } { + base.AddNode(&Node{ + ID: filePath + "::handle", Name: "handle", Kind: KindFunction, FilePath: filePath, + }) + } + innerLayer := NewOverlayLayer() + innerLayer.MarkFile("repo/inner-hidden.go", true) + outerLayer := NewOverlayLayer() + outerLayer.MarkFile("repo/outer-hidden.go", true) + view := NewOverlaidView(NewOverlaidView(base, innerLayer), outerLayer) + callerFiles := map[string]bool{"repo/caller-hidden.go": true} + + page, err := view.FindNodesByNameBounded( + context.Background(), "handle", LocalizationNodeScope{ExcludeFiles: callerFiles}, 8, + ) + if err != nil { + t.Fatalf("bounded nested-overlay lookup: %v", err) + } + if page.Total != 1 || page.Truncated || len(page.Nodes) != 1 || page.Nodes[0].FilePath != "repo/visible.go" { + t.Fatalf("page = %#v, want only the file outside all exclusion layers", page) + } + if len(callerFiles) != 1 || !callerFiles["repo/caller-hidden.go"] { + t.Fatalf("caller exclusion map was mutated: %#v", callerFiles) + } + for _, filePath := range []string{"repo/inner-hidden.go", "repo/outer-hidden.go"} { + if _, added := callerFiles[filePath]; added { + t.Fatalf("overlay path %q leaked into caller exclusion map: %#v", filePath, callerFiles) + } + } +} + +func TestOverlaidViewFindNodesByNameBoundedAllocationDoesNotScaleWithCoveredFiles(t *testing.T) { + base := New() + base.AddNode(&Node{ + ID: "repo/visible.go::handle", Name: "handle", Kind: KindFunction, + FilePath: "repo/visible.go", + }) + layer := NewOverlayLayer() + for index := 0; index < overlayExactNameInspectionLimit; index++ { + layer.MarkFile(fmt.Sprintf("repo/covered-%04d.go", index), true) + } + view := NewOverlaidView(base, layer) + + result := testing.Benchmark(func(b *testing.B) { + for iteration := 0; iteration < b.N; iteration++ { + page, err := view.FindNodesByNameBounded( + context.Background(), "handle", LocalizationNodeScope{}, 8, + ) + if err != nil || len(page.Nodes) != 1 { + b.Fatalf("bounded overlay lookup = %#v, %v", page, err) + } + } + }) + if bytes := result.AllocedBytesPerOp(); bytes > 16<<10 { + t.Fatalf("allocated %d bytes/op with many covered files, want request-bounded allocation", bytes) + } +} + +func TestOverlaidViewFindNodesByNameBoundedHonorsDetachedRemoval(t *testing.T) { + base := New() + stale := &Node{ + ID: "repo/old.go::handle", Name: "handle", Kind: KindFunction, + FilePath: "repo/old.go", + } + base.AddNode(stale) + recording := &recordingBoundedExactNameReader{Reader: base, bounded: base} + + layer := NewOverlayLayer() + // MarkRemoved is sufficient in the legacy overlay contract even when the + // layer was assembled without MarkFile. The bounded path must keep parity. + layer.MarkRemoved(stale.Name, stale.ID) + view := NewOverlaidView(recording, layer) + + page, err := view.FindNodesByNameBounded(context.Background(), "handle", LocalizationNodeScope{}, 8) + if err != nil { + t.Fatalf("bounded overlay lookup: %v", err) + } + if len(recording.limits) != 1 || recording.limits[0] != 9 { + t.Fatalf("base limits = %v, want one detached-shadow refill slot", recording.limits) + } + if page.Total != 0 || len(page.Nodes) != 0 || page.Truncated { + t.Fatalf("detached removal leaked stale base node: %#v", page) + } +} + +func TestFindFileNodesBoundedCapsScopeKindsAndCancels(t *testing.T) { + graph := New() + const filePath = "repo/dense.go" + for index := 0; index < 32; index++ { + graph.AddNode(&Node{ + ID: fmt.Sprintf("repo/dense.go::fn-%03d", index), Name: fmt.Sprintf("fn%d", index), + Kind: KindFunction, FilePath: filePath, RepoPrefix: "repo", + WorkspaceID: "workspace", ProjectID: "project", + Meta: map[string]any{"doc": "must not escape the summary projection"}, + }) + graph.AddNode(&Node{ + ID: fmt.Sprintf("repo/dense.go::value-%03d", index), Name: fmt.Sprintf("value%d", index), + Kind: KindVariable, FilePath: filePath, RepoPrefix: "repo", + WorkspaceID: "workspace", ProjectID: "project", + }) + graph.AddNode(&Node{ + ID: fmt.Sprintf("foreign/dense.go::fn-%03d", index), Name: fmt.Sprintf("foreign%d", index), + Kind: KindFunction, FilePath: filePath, RepoPrefix: "foreign", + WorkspaceID: "foreign", ProjectID: "foreign", + }) + } + + page, err := graph.FindFileNodesBounded( + context.Background(), filePath, + LocalizationNodeScope{ + WorkspaceID: "workspace", ProjectID: "project", + RepoAllow: map[string]bool{"repo": true}, + Kinds: map[NodeKind]bool{KindFunction: true}, + }, + 8, + ) + if err != nil { + t.Fatalf("bounded file lookup: %v", err) + } + if page.Total != 9 || !page.Truncated || len(page.Nodes) != 8 { + t.Fatalf("page = %#v, want threshold total 9, truncated, cap 8", page) + } + if !sort.SliceIsSorted(page.Nodes, func(i, j int) bool { return page.Nodes[i].ID < page.Nodes[j].ID }) { + t.Fatalf("nodes are not deterministic: %#v", page.Nodes) + } + for _, node := range page.Nodes { + if node.Kind != KindFunction || node.RepoPrefix != "repo" || node.WorkspaceID != "workspace" { + t.Fatalf("scope or kind was applied after cap: %#v", node) + } + if node.Meta != nil { + t.Fatalf("in-memory file summary retained metadata: %#v", node.Meta) + } + } + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + if cancelled, err := graph.FindFileNodesBounded(ctx, filePath, LocalizationNodeScope{}, 8); err == nil || len(cancelled.Nodes) != 0 { + t.Fatalf("cancelled file lookup = %#v, %v; want empty error result", cancelled, err) + } +} + +func TestFindFileNodesBoundedFiltersTestsBeforeCap(t *testing.T) { + graph := New() + const filePath = "repo/dense_test.go" + for index := 0; index < 32; index++ { + graph.AddNode(&Node{ + ID: fmt.Sprintf("repo/dense_test.go::a-test-%03d", index), Name: "test", + Kind: KindFunction, FilePath: filePath, Meta: map[string]any{"is_test": true}, + }) + } + for index := 0; index < 10; index++ { + graph.AddNode(&Node{ + ID: fmt.Sprintf("repo/dense_test.go::z-prod-%03d", index), Name: "prod", + Kind: KindFunction, FilePath: filePath, + }) + } + + page, err := graph.FindFileNodesBounded( + context.Background(), filePath, LocalizationNodeScope{ExcludeTests: true}, 8, + ) + if err != nil { + t.Fatalf("bounded file lookup: %v", err) + } + if page.Total != 9 || !page.Truncated || len(page.Nodes) != 8 { + t.Fatalf("page = %#v, want production sentinel after test rows", page) + } + for _, node := range page.Nodes { + if node.Name != "prod" { + t.Fatalf("test node consumed the production cap: %#v", node) + } + } +} + +func TestFindFileNodesBoundedExcludesKindsBeforeCap(t *testing.T) { + graph := New() + const filePath = "repo/generated.go" + for index := 0; index < 32; index++ { + graph.AddNode(&Node{ + ID: fmt.Sprintf("repo/generated.go::a-param-%03d", index), Name: "arg", + Kind: KindParam, FilePath: filePath, + }) + } + for index := 0; index < 10; index++ { + graph.AddNode(&Node{ + ID: fmt.Sprintf("repo/generated.go::z-function-%03d", index), Name: "function", + Kind: KindFunction, FilePath: filePath, + }) + } + + page, err := graph.FindFileNodesBounded( + context.Background(), filePath, + LocalizationNodeScope{ExcludeKinds: map[NodeKind]bool{KindParam: true}}, 8, + ) + if err != nil { + t.Fatalf("bounded file lookup: %v", err) + } + if page.Total != 9 || !page.Truncated || len(page.Nodes) != 8 { + t.Fatalf("page = %#v, want definition sentinel behind excluded params", page) + } + for _, node := range page.Nodes { + if node.Kind != KindFunction { + t.Fatalf("excluded kind consumed the cap: %#v", node) + } + } +} + +func TestOverlaidViewFindFileNodesBoundedReplacesAndTombstones(t *testing.T) { + base := New() + const filePath = "repo/handler.go" + base.AddNode(&Node{ID: filePath + "::old", Name: "old", Kind: KindFunction, FilePath: filePath}) + + layer := NewOverlayLayer() + layer.MarkFile(filePath, false) + layer.AddNode(filePath, &Node{ + ID: filePath + "::replacement", Name: "replacement", Kind: KindFunction, FilePath: filePath, + Meta: map[string]any{"doc": "must not escape the overlay summary projection"}, + }) + layer.AddNode(filePath, &Node{ID: filePath + "::ignored", Name: "ignored", Kind: KindVariable, FilePath: filePath}) + view := NewOverlaidView(base, layer) + + page, err := view.FindFileNodesBounded( + context.Background(), filePath, + LocalizationNodeScope{ExcludeKinds: map[NodeKind]bool{KindVariable: true}}, 8, + ) + if err != nil { + t.Fatalf("bounded overlay file lookup: %v", err) + } + if page.Total != 1 || page.Truncated || len(page.Nodes) != 1 || page.Nodes[0].Name != "replacement" { + t.Fatalf("overlay page = %#v, want replacement only", page) + } + if page.Nodes[0].Meta != nil { + t.Fatalf("overlay file summary retained metadata: %#v", page.Nodes[0].Meta) + } + + tombstone := NewOverlayLayer() + tombstone.MarkFile(filePath, true) + deleted, err := NewOverlaidView(base, tombstone).FindFileNodesBounded( + context.Background(), filePath, LocalizationNodeScope{}, 8, + ) + if err != nil { + t.Fatalf("bounded tombstone lookup: %v", err) + } + if deleted.Total != 0 || deleted.Truncated || len(deleted.Nodes) != 0 { + t.Fatalf("tombstone leaked base nodes: %#v", deleted) + } +} + +func TestFindFileNodesBoundedAllocatesOnlyRetainedSummaries(t *testing.T) { + const ( + filePath = "repo/generated.go" + nodeCount = 2_048 + limit = 8 + ) + memory := New() + layer := NewOverlayLayer() + layer.MarkFile(filePath, false) + for index := 0; index < nodeCount; index++ { + node := &Node{ + ID: fmt.Sprintf("%s::declaration-%04d", filePath, index), + Name: fmt.Sprintf("declaration%d", index), Kind: KindFunction, + FilePath: filePath, Meta: map[string]any{"doc": "large retrieval payload"}, + } + memory.AddNode(node) + layer.AddNode(filePath, node) + } + + readers := []struct { + name string + reader BoundedFileNodeReader + }{ + {name: "graph", reader: memory}, + {name: "overlay", reader: NewOverlaidView(nil, layer)}, + } + for _, test := range readers { + t.Run(test.name, func(t *testing.T) { + var page BoundedNodeProjection + allocations := testing.AllocsPerRun(3, func() { + var err error + page, err = test.reader.FindFileNodesBounded( + context.Background(), filePath, LocalizationNodeScope{}, limit, + ) + if err != nil { + panic(err) + } + }) + if allocations > 64 { + t.Fatalf("allocations = %.0f, want response-bounded summary allocation", allocations) + } + if len(page.Nodes) != limit || page.Total != limit+1 || !page.Truncated { + t.Fatalf("page = %#v, want bounded saturated projection", page) + } + if cap(page.Nodes) != len(page.Nodes) { + t.Fatalf("page capacity = %d, want exact returned length %d", cap(page.Nodes), len(page.Nodes)) + } + for _, node := range page.Nodes[:cap(page.Nodes)] { + if node.Meta != nil { + t.Fatalf("retained or reslice-reachable summary hydrated metadata: %#v", node.Meta) + } + } + }) + } +} + +type cancelAfterLocalizationChecksContext struct { + context.Context + remaining int + done chan struct{} +} + +func (ctx *cancelAfterLocalizationChecksContext) Done() <-chan struct{} { + return ctx.done +} + +func (ctx *cancelAfterLocalizationChecksContext) Err() error { + if ctx.remaining > 0 { + ctx.remaining-- + if ctx.remaining > 0 { + return nil + } + } + select { + case <-ctx.done: + default: + close(ctx.done) + } + return context.Canceled +} + +func TestOverlaidViewFindFileNodesBoundedCancelsDuringLargeScan(t *testing.T) { + const filePath = "repo/generated.go" + layer := NewOverlayLayer() + layer.MarkFile(filePath, false) + for index := 0; index < 1_024; index++ { + layer.AddNode(filePath, &Node{ + ID: fmt.Sprintf("%s::declaration-%04d", filePath, index), + Name: "declaration", Kind: KindFunction, FilePath: filePath, + }) + } + ctx := &cancelAfterLocalizationChecksContext{ + Context: context.Background(), remaining: 3, done: make(chan struct{}), + } + page, err := NewOverlaidView(nil, layer).FindFileNodesBounded( + ctx, filePath, LocalizationNodeScope{}, 8, + ) + if err != context.Canceled { + t.Fatalf("error = %v, want context cancellation during scan", err) + } + if len(page.Nodes) != 0 || page.Total != 0 { + t.Fatalf("cancelled scan returned partial page: %#v", page) + } +} + +func TestOverlayLayerAddNodeReplacesDuplicateIdentity(t *testing.T) { + const ( + filePath = "repo/handler.go" + nodeID = filePath + "::handler" + ) + layer := NewOverlayLayer() + layer.MarkFile(filePath, false) + layer.AddNode(filePath, &Node{ + ID: nodeID, Name: "old", QualName: "repo.Old", Kind: KindFunction, FilePath: filePath, + }) + layer.AddNode(filePath, &Node{ + ID: nodeID, Name: "handler", QualName: "repo.Handler", Kind: KindFunction, + FilePath: filePath, StartLine: 2, + }) + final := &Node{ + ID: nodeID, Name: "handler", QualName: "repo.Handler", Kind: KindFunction, + FilePath: filePath, StartLine: 3, + } + layer.AddNode(filePath, final) + + if nodes := layer.nodesForFile(filePath); len(nodes) != 1 || nodes[0] != final { + t.Fatalf("file nodes = %#v, want one final replacement", nodes) + } + if old := layer.NodesByName("old"); len(old) != 0 { + t.Fatalf("old name index retained replacement: %#v", old) + } + if current := layer.NodesByName("handler"); len(current) != 1 || current[0] != final { + t.Fatalf("name index = %#v, want final replacement", current) + } + view := NewOverlaidView(nil, layer) + if view.GetNode(nodeID) != final || view.GetNodeByQualName("repo.Handler") != final { + t.Fatal("identity or qualified-name index did not retain final replacement") + } + if stale := view.GetNodeByQualName("repo.Old"); stale != nil { + t.Fatalf("stale qualified name retained: %#v", stale) + } + page, err := view.FindFileNodesBounded(context.Background(), filePath, LocalizationNodeScope{}, 1) + if err != nil { + t.Fatalf("bounded file lookup: %v", err) + } + if page.Total != 1 || page.Truncated || len(page.Nodes) != 1 || page.Nodes[0].StartLine != 3 { + t.Fatalf("duplicate identity inflated bounded projection: %#v", page) + } +} diff --git a/internal/graph/overlay.go b/internal/graph/overlay.go index f24ff98d..7254f095 100644 --- a/internal/graph/overlay.go +++ b/internal/graph/overlay.go @@ -1,6 +1,7 @@ package graph import ( + "context" "sort" "strings" "sync" @@ -59,6 +60,9 @@ type OverlayLayer struct { // filter base hits whose enclosing file is overlaid but whose // id disappeared from the overlay's node list. nameRemoved map[string]map[string]bool + // removedByID is the immutable identity-side index for bounded point and + // adjacency projections. It avoids rescanning every name bucket per ID. + removedByID map[string]bool } // overlayFileEntry carries one file's overlay state inside the @@ -83,6 +87,7 @@ func NewOverlayLayer() *OverlayLayer { nodesByName: make(map[string][]*Node), nodesByQual: make(map[string]*Node), nameRemoved: make(map[string]map[string]bool), + removedByID: make(map[string]bool), } } @@ -111,6 +116,17 @@ func (l *OverlayLayer) AddNode(graphPath string, n *Node) { // Tombstone: silently drop. Caller bug — but cheap to absorb. return } + if old := l.nodeByID[n.ID]; old != nil { + for index, candidate := range entry.Nodes { + if candidate == nil || candidate.ID != n.ID { + continue + } + entry.Nodes[index] = n + l.nodeByID[n.ID] = n + l.replaceNodeIndexes(old, n) + return + } + } entry.Nodes = append(entry.Nodes, n) l.nodeByID[n.ID] = n if n.Name != "" { @@ -121,6 +137,50 @@ func (l *OverlayLayer) AddNode(graphPath string, n *Node) { } } +func (l *OverlayLayer) replaceNodeIndexes(old, replacement *Node) { + if old.Name == replacement.Name { + if old.Name != "" { + replaced := false + for index, candidate := range l.nodesByName[old.Name] { + if candidate != nil && candidate.ID == replacement.ID { + l.nodesByName[old.Name][index] = replacement + replaced = true + } + } + if !replaced { + l.nodesByName[old.Name] = append(l.nodesByName[old.Name], replacement) + } + } + } else { + if old.Name != "" { + bucket := l.nodesByName[old.Name] + filtered := bucket[:0] + for _, candidate := range bucket { + if candidate == nil || candidate.ID != replacement.ID { + filtered = append(filtered, candidate) + } + } + if len(filtered) == 0 { + delete(l.nodesByName, old.Name) + } else { + l.nodesByName[old.Name] = filtered + } + } + if replacement.Name != "" { + l.nodesByName[replacement.Name] = append(l.nodesByName[replacement.Name], replacement) + } + } + + if old.QualName != "" && old.QualName != replacement.QualName { + if current := l.nodesByQual[old.QualName]; current != nil && current.ID == replacement.ID { + delete(l.nodesByQual, old.QualName) + } + } + if replacement.QualName != "" { + l.nodesByQual[replacement.QualName] = replacement + } +} + // AddEdge attaches one resolved overlay edge. The local-resolver // pass at layer construction is expected to have rewritten any // `unresolved::*` placeholders to point at concrete (overlay or @@ -150,6 +210,7 @@ func (l *OverlayLayer) MarkRemoved(baseName, baseID string) { l.nameRemoved[baseName] = set } set[baseID] = true + l.removedByID[baseID] = true } // HasFile reports whether the overlay covers a particular graph path @@ -257,6 +318,21 @@ func (l *OverlayLayer) nodesForFile(graphPath string) []*Node { return out } +// nodesForFileReadOnly exposes the immutable layer's backing slice to internal +// bounded readers. Overlay construction is complete before an OverlaidView is +// published, so scanning this slice cannot race a writer and avoids an O(N) +// snapshot allocation before cancellation can be observed. +func (l *OverlayLayer) nodesForFileReadOnly(graphPath string) []*Node { + if l == nil { + return nil + } + entry := l.entries[graphPath] + if entry == nil || entry.Deleted { + return nil + } + return entry.Nodes +} + // OverlaidView composes an immutable base Reader with a per-session // overlay layer. Every read path consults the layer first for paths // the overlay covers; falls through to base otherwise. The base is @@ -337,37 +413,8 @@ func (v *OverlaidView) GetNode(id string) *Node { // fans out as a single batched lookup against the base store. Missing // IDs are simply absent from the returned map. func (v *OverlaidView) GetNodesByIDs(ids []string) map[string]*Node { - if len(ids) == 0 { - return nil - } - out := make(map[string]*Node, len(ids)) - baseIDs := ids[:0:0] // fresh backing array — never aliases caller's slice - for _, id := range ids { - if id == "" { - continue - } - if _, dup := out[id]; dup { - continue - } - if v.layer != nil && v.nodeBelongsToOverlay(id) { - if n := v.layer.nodeByID[id]; n != nil { - out[id] = n - } - // Overlay tombstone — ID is hidden, do not fall back to base. - continue - } - // Track for the single base round-trip; reserve a slot in `out` - // only after the batched lookup returns. - baseIDs = append(baseIDs, id) - } - if len(baseIDs) > 0 && v.base != nil { - for id, n := range v.base.GetNodesByIDs(baseIDs) { - if n != nil { - out[id] = n - } - } - } - return out + nodes, _ := v.GetNodesByIDsContext(context.Background(), ids) + return nodes } // GetNodeByQualName: overlay first, then base. Base hits are filtered diff --git a/internal/graph/store_sqlite/bounded_adjacency.go b/internal/graph/store_sqlite/bounded_adjacency.go new file mode 100644 index 00000000..ef15ab00 --- /dev/null +++ b/internal/graph/store_sqlite/bounded_adjacency.go @@ -0,0 +1,377 @@ +package store_sqlite + +import ( + "context" + "database/sql" + "fmt" + "sort" + "strings" + + "github.com/zzet/gortex/internal/graph" +) + +var ( + _ graph.BoundedOutgoingEdgeIdentityReader = (*Store)(nil) + _ graph.BoundedIncomingEdgeIdentityReader = (*Store)(nil) + _ graph.BoundedOutgoingSiteEdgeIdentityReader = (*Store)(nil) +) + +func canonicalSQLiteBoundedAdjacencyKinds(ctx context.Context, kinds []graph.EdgeKind) ([]graph.EdgeKind, error) { + if len(kinds) > graph.MaxBoundedAdjacencyKinds { + return nil, &graph.BoundedLocalizationLimitError{ + Resource: "SQLite adjacency kinds", + Limit: graph.MaxBoundedAdjacencyKinds, + } + } + seen := make(map[graph.EdgeKind]struct{}, len(kinds)) + out := make([]graph.EdgeKind, 0, len(kinds)) + for index, kind := range kinds { + if index&127 == 0 { + if err := ctx.Err(); err != nil { + return nil, err + } + } + if kind == "" { + continue + } + if _, duplicate := seen[kind]; duplicate { + continue + } + seen[kind] = struct{}{} + out = append(out, kind) + } + sort.Slice(out, func(i, j int) bool { return out[i] < out[j] }) + return out, nil +} + +func canonicalSQLiteBoundedAdjacencyEndpoints(ctx context.Context, ids []string) ([]string, error) { + if len(ids) > graph.MaxBoundedAdjacencyKeys { + return nil, &graph.BoundedLocalizationLimitError{ + Resource: "SQLite adjacency endpoint keys", + Limit: graph.MaxBoundedAdjacencyKeys, + } + } + seen := make(map[string]struct{}, len(ids)) + out := make([]string, 0, len(ids)) + for index, id := range ids { + if index&127 == 0 { + if err := ctx.Err(); err != nil { + return nil, err + } + } + if id == "" { + continue + } + if _, duplicate := seen[id]; duplicate { + continue + } + seen[id] = struct{}{} + out = append(out, id) + } + sort.Strings(out) + return out, nil +} + +func canonicalSQLiteBoundedAdjacencySites(ctx context.Context, sites []graph.EdgeSourceSite) ([]graph.EdgeSourceSite, error) { + if len(sites) > graph.MaxBoundedAdjacencyKeys { + return nil, &graph.BoundedLocalizationLimitError{ + Resource: "SQLite adjacency source-site keys", + Limit: graph.MaxBoundedAdjacencyKeys, + } + } + seen := make(map[graph.EdgeSourceSite]struct{}, len(sites)) + out := make([]graph.EdgeSourceSite, 0, len(sites)) + for index, site := range sites { + if index&127 == 0 { + if err := ctx.Err(); err != nil { + return nil, err + } + } + if site.From == "" { + continue + } + if _, duplicate := seen[site]; duplicate { + continue + } + seen[site] = struct{}{} + out = append(out, site) + } + sort.Slice(out, func(i, j int) bool { + if out[i].From != out[j].From { + return out[i].From < out[j].From + } + return out[i].Line < out[j].Line + }) + return out, nil +} + +func validateSQLiteBoundedAdjacencyLimit(limit int) error { + if limit < 1 || limit > graph.MaxBoundedAdjacencyRowsPerKey { + return &graph.BoundedLocalizationLimitError{ + Resource: "SQLite adjacency rows per key", + Limit: graph.MaxBoundedAdjacencyRowsPerKey, + } + } + return nil +} + +func boundedAdjacencyKindList(count int) string { + if count <= 0 { + return "" + } + return strings.TrimRight(strings.Repeat("?,", count), ",") +} + +func boundedOutgoingAdjacencySQL(kindCount int) string { + return `SELECT from_id, to_id, kind, file_path, line +FROM edges INDEXED BY edges_by_from +WHERE from_id = ? AND kind IN (` + boundedAdjacencyKindList(kindCount) + `) +LIMIT ?` +} + +func boundedIncomingAdjacencySQL(kindCount int) string { + return `SELECT from_id, to_id, kind, file_path, line +FROM edges INDEXED BY edges_by_to +WHERE to_id = ? AND kind IN (` + boundedAdjacencyKindList(kindCount) + `) +LIMIT ?` +} + +func boundedOutgoingSiteAdjacencySQL(kindCount int) string { + return `SELECT from_id, to_id, kind, file_path, line +FROM edges INDEXED BY edges_by_from_line_kind +WHERE from_id = ? AND line = ? AND kind IN (` + boundedAdjacencyKindList(kindCount) + `) +LIMIT ?` +} + +func boundedAdjacencyArgs(prefix []any, kinds []graph.EdgeKind, limit int) []any { + args := make([]any, 0, len(prefix)+len(kinds)+1) + args = append(args, prefix...) + for _, kind := range kinds { + args = append(args, string(kind)) + } + args = append(args, limit+1) + return args +} + +func scanSQLiteBoundedAdjacencyRows( + ctx context.Context, + rows *sql.Rows, + limit int, + total *int, +) ([]graph.EdgeIdentity, bool, error) { + var identities []graph.EdgeIdentity + var seen map[graph.EdgeIdentity]struct{} + rowIndex := 0 + for rows.Next() { + if rowIndex&127 == 0 { + if err := ctx.Err(); err != nil { + _ = rows.Close() + return nil, false, err + } + } + rowIndex++ + var identity graph.EdgeIdentity + var kind string + if err := rows.Scan(&identity.From, &identity.To, &kind, &identity.FilePath, &identity.Line); err != nil { + _ = rows.Close() + return nil, false, err + } + identity.Kind = graph.EdgeKind(kind) + if _, duplicate := seen[identity]; duplicate { + continue + } + if seen == nil { + identities = make([]graph.EdgeIdentity, 0, 1) + seen = make(map[graph.EdgeIdentity]struct{}, 1) + } + seen[identity] = struct{}{} + if *total == graph.MaxBoundedAdjacencyTotalIdentities { + _ = rows.Close() + return nil, false, &graph.BoundedLocalizationLimitError{ + Resource: "SQLite adjacency matching identities", + Limit: graph.MaxBoundedAdjacencyTotalIdentities, + } + } + *total++ + if len(identities) == limit { + if err := rows.Close(); err != nil { + return nil, false, err + } + return nil, true, nil + } + identities = append(identities, identity) + } + if err := rows.Err(); err != nil { + _ = rows.Close() + return nil, false, fmt.Errorf("bounded adjacency query: %w", err) + } + if err := rows.Close(); err != nil { + return nil, false, err + } + sort.Slice(identities, func(i, j int) bool { + if identities[i].From != identities[j].From { + return identities[i].From < identities[j].From + } + if identities[i].To != identities[j].To { + return identities[i].To < identities[j].To + } + if identities[i].Kind != identities[j].Kind { + return identities[i].Kind < identities[j].Kind + } + if identities[i].FilePath != identities[j].FilePath { + return identities[i].FilePath < identities[j].FilePath + } + return identities[i].Line < identities[j].Line + }) + return identities, false, nil +} + +func (s *Store) FindOutgoingEdgeIdentitiesBounded( + ctx context.Context, + sourceIDs []string, + kinds []graph.EdgeKind, + limit int, +) (graph.BoundedEdgeIdentityProjection, error) { + return s.findEndpointEdgeIdentitiesBounded(ctx, sourceIDs, kinds, limit, true) +} + +func (s *Store) FindIncomingEdgeIdentitiesBounded( + ctx context.Context, + targetIDs []string, + kinds []graph.EdgeKind, + limit int, +) (graph.BoundedEdgeIdentityProjection, error) { + return s.findEndpointEdgeIdentitiesBounded(ctx, targetIDs, kinds, limit, false) +} + +func (s *Store) findEndpointEdgeIdentitiesBounded( + ctx context.Context, + endpointIDs []string, + kinds []graph.EdgeKind, + limit int, + outgoing bool, +) (graph.BoundedEdgeIdentityProjection, error) { + if ctx == nil { + ctx = context.Background() + } + if err := ctx.Err(); err != nil { + return graph.BoundedEdgeIdentityProjection{}, err + } + if err := validateSQLiteBoundedAdjacencyLimit(limit); err != nil { + return graph.BoundedEdgeIdentityProjection{}, err + } + kindSet, err := canonicalSQLiteBoundedAdjacencyKinds(ctx, kinds) + if err != nil { + return graph.BoundedEdgeIdentityProjection{}, err + } + ids, err := canonicalSQLiteBoundedAdjacencyEndpoints(ctx, endpointIDs) + if err != nil { + return graph.BoundedEdgeIdentityProjection{}, err + } + projection := graph.BoundedEdgeIdentityProjection{ + ByEndpoint: make(map[string][]graph.EdgeIdentity), + Truncated: make(map[string]bool), + } + if len(ids) == 0 || len(kindSet) == 0 { + return projection, nil + } + tx, err := s.db.BeginTx(ctx, &sql.TxOptions{ReadOnly: true}) + if err != nil { + return graph.BoundedEdgeIdentityProjection{}, err + } + defer func() { _ = tx.Rollback() }() + query := boundedIncomingAdjacencySQL(len(kindSet)) + if outgoing { + query = boundedOutgoingAdjacencySQL(len(kindSet)) + } + total := 0 + for _, id := range ids { + if err := ctx.Err(); err != nil { + return graph.BoundedEdgeIdentityProjection{}, err + } + rows, queryErr := tx.QueryContext(ctx, query, boundedAdjacencyArgs([]any{id}, kindSet, limit)...) + if queryErr != nil { + return graph.BoundedEdgeIdentityProjection{}, queryErr + } + identities, truncated, scanErr := scanSQLiteBoundedAdjacencyRows(ctx, rows, limit, &total) + if scanErr != nil { + return graph.BoundedEdgeIdentityProjection{}, scanErr + } + if truncated { + projection.Truncated[id] = true + continue + } + if len(identities) > 0 { + projection.ByEndpoint[id] = identities + } + } + if err := tx.Commit(); err != nil { + return graph.BoundedEdgeIdentityProjection{}, err + } + return projection, nil +} + +func (s *Store) FindOutgoingSiteEdgeIdentitiesBounded( + ctx context.Context, + sites []graph.EdgeSourceSite, + kinds []graph.EdgeKind, + limit int, +) (graph.BoundedSiteEdgeIdentityProjection, error) { + if ctx == nil { + ctx = context.Background() + } + if err := ctx.Err(); err != nil { + return graph.BoundedSiteEdgeIdentityProjection{}, err + } + if err := validateSQLiteBoundedAdjacencyLimit(limit); err != nil { + return graph.BoundedSiteEdgeIdentityProjection{}, err + } + kindSet, err := canonicalSQLiteBoundedAdjacencyKinds(ctx, kinds) + if err != nil { + return graph.BoundedSiteEdgeIdentityProjection{}, err + } + canonical, err := canonicalSQLiteBoundedAdjacencySites(ctx, sites) + if err != nil { + return graph.BoundedSiteEdgeIdentityProjection{}, err + } + projection := graph.BoundedSiteEdgeIdentityProjection{ + BySite: make(map[graph.EdgeSourceSite][]graph.EdgeIdentity), + Truncated: make(map[graph.EdgeSourceSite]bool), + } + if len(canonical) == 0 || len(kindSet) == 0 { + return projection, nil + } + tx, err := s.db.BeginTx(ctx, &sql.TxOptions{ReadOnly: true}) + if err != nil { + return graph.BoundedSiteEdgeIdentityProjection{}, err + } + defer func() { _ = tx.Rollback() }() + query := boundedOutgoingSiteAdjacencySQL(len(kindSet)) + total := 0 + for _, site := range canonical { + if err := ctx.Err(); err != nil { + return graph.BoundedSiteEdgeIdentityProjection{}, err + } + rows, queryErr := tx.QueryContext( + ctx, query, boundedAdjacencyArgs([]any{site.From, site.Line}, kindSet, limit)..., + ) + if queryErr != nil { + return graph.BoundedSiteEdgeIdentityProjection{}, queryErr + } + identities, truncated, scanErr := scanSQLiteBoundedAdjacencyRows(ctx, rows, limit, &total) + if scanErr != nil { + return graph.BoundedSiteEdgeIdentityProjection{}, scanErr + } + if truncated { + projection.Truncated[site] = true + continue + } + if len(identities) > 0 { + projection.BySite[site] = identities + } + } + if err := tx.Commit(); err != nil { + return graph.BoundedSiteEdgeIdentityProjection{}, err + } + return projection, nil +} diff --git a/internal/graph/store_sqlite/bounded_adjacency_test.go b/internal/graph/store_sqlite/bounded_adjacency_test.go new file mode 100644 index 00000000..c2e839eb --- /dev/null +++ b/internal/graph/store_sqlite/bounded_adjacency_test.go @@ -0,0 +1,363 @@ +package store_sqlite + +import ( + "context" + "errors" + "fmt" + "math" + "path/filepath" + "reflect" + "sort" + "strings" + "testing" + + "github.com/zzet/gortex/internal/graph" +) + +func openBoundedAdjacencyTestStore(t *testing.T) *Store { + t.Helper() + store, err := Open(filepath.Join(t.TempDir(), "graph.sqlite")) + if err != nil { + t.Fatalf("open store: %v", err) + } + t.Cleanup(func() { _ = store.Close() }) + return store +} + +func TestSQLiteBoundedAdjacencyFiltersBeforeLimitAndReenters(t *testing.T) { + store := openBoundedAdjacencyTestStore(t) + const source = "repo/source.go::source" + edges := make([]*graph.Edge, 0, 304) + for index := 0; index < 300; index++ { + edges = append(edges, &graph.Edge{ + From: source, To: fmt.Sprintf("noise-%03d", index), Kind: graph.EdgeReferences, + FilePath: "repo/source.go", Line: 1_000 + index, + Meta: map[string]any{"payload": strings.Repeat("x", 1024)}, + }) + } + first := &graph.Edge{From: source, To: "target", Kind: graph.EdgeCalls, FilePath: "repo/source.go", Line: 10} + second := &graph.Edge{From: source, To: "other", Kind: graph.EdgeCalls, FilePath: "generated.go", Line: 20} + third := &graph.Edge{From: "repo/other.go::caller", To: "target", Kind: graph.EdgeCalls, FilePath: "repo/other.go", Line: 30} + edges = append(edges, first, second, third) + store.AddBatch(nil, edges) + + sources := []string{source, source, ""} + kinds := []graph.EdgeKind{graph.EdgeCalls, graph.EdgeCalls, ""} + sourcesBefore := append([]string(nil), sources...) + kindsBefore := append([]graph.EdgeKind(nil), kinds...) + outgoing, err := store.FindOutgoingEdgeIdentitiesBounded(context.Background(), sources, kinds, 2) + if err != nil { + t.Fatalf("outgoing projection: %v", err) + } + if !reflect.DeepEqual(sources, sourcesBefore) || !reflect.DeepEqual(kinds, kindsBefore) { + t.Fatalf("caller inputs mutated: sources=%#v kinds=%#v", sources, kinds) + } + wantOutgoing := []graph.EdgeIdentity{graph.EdgeIdentityFor(first), graph.EdgeIdentityFor(second)} + sort.Slice(wantOutgoing, func(i, j int) bool { return wantOutgoing[i].To < wantOutgoing[j].To }) + if !reflect.DeepEqual(outgoing.ByEndpoint[source], wantOutgoing) || outgoing.Truncated[source] { + t.Fatalf("outgoing = %#v, want %#v complete", outgoing, wantOutgoing) + } + + incoming, err := store.FindIncomingEdgeIdentitiesBounded(context.Background(), []string{"target"}, []graph.EdgeKind{graph.EdgeCalls}, 2) + if err != nil { + t.Fatalf("incoming projection: %v", err) + } + wantIncoming := []graph.EdgeIdentity{graph.EdgeIdentityFor(first), graph.EdgeIdentityFor(third)} + sort.Slice(wantIncoming, func(i, j int) bool { return wantIncoming[i].From < wantIncoming[j].From }) + if !reflect.DeepEqual(incoming.ByEndpoint["target"], wantIncoming) || incoming.Truncated["target"] { + t.Fatalf("incoming = %#v, want %#v complete", incoming, wantIncoming) + } + + site := graph.EdgeSourceSite{From: source, Line: 10} + sites := []graph.EdgeSourceSite{site, site} + sitesBefore := append([]graph.EdgeSourceSite(nil), sites...) + siteProjection, err := store.FindOutgoingSiteEdgeIdentitiesBounded(context.Background(), sites, []graph.EdgeKind{graph.EdgeCalls}, 1) + if err != nil { + t.Fatalf("site projection: %v", err) + } + if !reflect.DeepEqual(sites, sitesBefore) { + t.Fatalf("caller sites mutated: %#v", sites) + } + if got := siteProjection.BySite[site]; !reflect.DeepEqual(got, []graph.EdgeIdentity{graph.EdgeIdentityFor(first)}) || siteProjection.Truncated[site] { + t.Fatalf("site = %#v, want exact source+line identity", siteProjection) + } + + // Every cursor and the read-only transaction must be closed before return. + newEdge := &graph.Edge{From: "new-source", To: "new-target", Kind: graph.EdgeCalls} + store.AddEdge(newEdge) + reentered, err := store.FindOutgoingEdgeIdentitiesBounded(context.Background(), []string{newEdge.From}, []graph.EdgeKind{graph.EdgeCalls}, 1) + if err != nil || len(reentered.ByEndpoint[newEdge.From]) != 1 { + t.Fatalf("re-entered projection = %#v, %v", reentered, err) + } +} + +func TestSQLiteBoundedAdjacencyHardEnvelopesAndAggregateCap(t *testing.T) { + store := openBoundedAdjacencyTestStore(t) + tooManyKeys := make([]string, graph.MaxBoundedAdjacencyKeys+1) + tooManyKinds := make([]graph.EdgeKind, graph.MaxBoundedAdjacencyKinds+1) + for index := range tooManyKeys { + tooManyKeys[index] = "duplicate" + } + for index := range tooManyKinds { + tooManyKinds[index] = graph.EdgeCalls + } + for _, run := range []func() error{ + func() error { + _, err := store.FindOutgoingEdgeIdentitiesBounded(context.Background(), tooManyKeys, []graph.EdgeKind{graph.EdgeCalls}, 1) + return err + }, + func() error { + _, err := store.FindIncomingEdgeIdentitiesBounded(context.Background(), []string{"target"}, tooManyKinds, 1) + return err + }, + func() error { + _, err := store.FindOutgoingSiteEdgeIdentitiesBounded(context.Background(), nil, nil, 0) + return err + }, + func() error { + _, err := store.FindOutgoingEdgeIdentitiesBounded(context.Background(), nil, nil, math.MaxInt) + return err + }, + } { + var limitErr *graph.BoundedLocalizationLimitError + if err := run(); !errors.As(err, &limitErr) { + t.Fatalf("error = %v, want typed bounded limit", err) + } + } + + edges := make([]*graph.Edge, 0, 65*graph.MaxBoundedAdjacencyRowsPerKey) + sources := make([]string, 0, 65) + for sourceIndex := 0; sourceIndex < 65; sourceIndex++ { + source := fmt.Sprintf("source-%02d", sourceIndex) + sources = append(sources, source) + for edgeIndex := 0; edgeIndex < graph.MaxBoundedAdjacencyRowsPerKey; edgeIndex++ { + edges = append(edges, &graph.Edge{ + From: source, To: fmt.Sprintf("target-%02d-%03d", sourceIndex, edgeIndex), + Kind: graph.EdgeCalls, FilePath: source + ".go", Line: edgeIndex, + }) + } + } + store.AddBatch(nil, edges) + exact, err := store.FindOutgoingEdgeIdentitiesBounded(context.Background(), sources[:64], []graph.EdgeKind{graph.EdgeCalls}, graph.MaxBoundedAdjacencyRowsPerKey) + if err != nil || len(exact.ByEndpoint) != 64 { + t.Fatalf("exact aggregate cap = %d endpoints, %v", len(exact.ByEndpoint), err) + } + over, err := store.FindOutgoingEdgeIdentitiesBounded(context.Background(), sources, []graph.EdgeKind{graph.EdgeCalls}, graph.MaxBoundedAdjacencyRowsPerKey) + var limitErr *graph.BoundedLocalizationLimitError + if !errors.As(err, &limitErr) || limitErr.Resource != "SQLite adjacency matching identities" || len(over.ByEndpoint) != 0 || len(over.Truncated) != 0 { + t.Fatalf("aggregate overflow = %#v, %v", over, err) + } +} + +func TestSQLiteBoundedAdjacencyPerKeySentinels(t *testing.T) { + store := openBoundedAdjacencyTestStore(t) + edges := make([]*graph.Edge, 0, 4*graph.MaxBoundedAdjacencyRowsPerKey+2) + for index := 0; index < graph.MaxBoundedAdjacencyRowsPerKey; index++ { + edges = append(edges, &graph.Edge{From: "exact-source", To: fmt.Sprintf("exact-%03d", index), Kind: graph.EdgeCalls, Line: 10}) + } + for index := 0; index <= graph.MaxBoundedAdjacencyRowsPerKey; index++ { + edges = append(edges, &graph.Edge{From: "overflow-source", To: fmt.Sprintf("overflow-%03d", index), Kind: graph.EdgeCalls, Line: 20}) + } + for index := 0; index < graph.MaxBoundedAdjacencyRowsPerKey; index++ { + edges = append(edges, &graph.Edge{From: fmt.Sprintf("exact-caller-%03d", index), To: "exact-target", Kind: graph.EdgeCalls}) + } + for index := 0; index <= graph.MaxBoundedAdjacencyRowsPerKey; index++ { + edges = append(edges, &graph.Edge{From: fmt.Sprintf("overflow-caller-%03d", index), To: "overflow-target", Kind: graph.EdgeCalls}) + } + store.AddBatch(nil, edges) + + projection, err := store.FindOutgoingEdgeIdentitiesBounded(context.Background(), []string{"exact-source"}, []graph.EdgeKind{graph.EdgeCalls}, graph.MaxBoundedAdjacencyRowsPerKey) + if err != nil || len(projection.ByEndpoint["exact-source"]) != graph.MaxBoundedAdjacencyRowsPerKey || projection.Truncated["exact-source"] { + t.Fatalf("exact endpoint boundary = %#v, %v", projection, err) + } + projection, err = store.FindOutgoingEdgeIdentitiesBounded(context.Background(), []string{"overflow-source"}, []graph.EdgeKind{graph.EdgeCalls}, graph.MaxBoundedAdjacencyRowsPerKey) + if err != nil || !projection.Truncated["overflow-source"] || len(projection.ByEndpoint["overflow-source"]) != 0 { + t.Fatalf("endpoint sentinel leaked rows = %#v, %v", projection, err) + } + + exactSite := graph.EdgeSourceSite{From: "exact-source", Line: 10} + siteProjection, err := store.FindOutgoingSiteEdgeIdentitiesBounded(context.Background(), []graph.EdgeSourceSite{exactSite}, []graph.EdgeKind{graph.EdgeCalls}, graph.MaxBoundedAdjacencyRowsPerKey) + if err != nil || len(siteProjection.BySite[exactSite]) != graph.MaxBoundedAdjacencyRowsPerKey || siteProjection.Truncated[exactSite] { + t.Fatalf("exact site boundary = %#v, %v", siteProjection, err) + } + overflowSite := graph.EdgeSourceSite{From: "overflow-source", Line: 20} + siteProjection, err = store.FindOutgoingSiteEdgeIdentitiesBounded(context.Background(), []graph.EdgeSourceSite{overflowSite}, []graph.EdgeKind{graph.EdgeCalls}, graph.MaxBoundedAdjacencyRowsPerKey) + if err != nil || !siteProjection.Truncated[overflowSite] || len(siteProjection.BySite[overflowSite]) != 0 { + t.Fatalf("site sentinel leaked rows = %#v, %v", siteProjection, err) + } + + projection, err = store.FindIncomingEdgeIdentitiesBounded(context.Background(), []string{"exact-target"}, []graph.EdgeKind{graph.EdgeCalls}, graph.MaxBoundedAdjacencyRowsPerKey) + if err != nil || len(projection.ByEndpoint["exact-target"]) != graph.MaxBoundedAdjacencyRowsPerKey || projection.Truncated["exact-target"] { + t.Fatalf("exact incoming boundary = %#v, %v", projection, err) + } + projection, err = store.FindIncomingEdgeIdentitiesBounded(context.Background(), []string{"overflow-target"}, []graph.EdgeKind{graph.EdgeCalls}, graph.MaxBoundedAdjacencyRowsPerKey) + if err != nil || !projection.Truncated["overflow-target"] || len(projection.ByEndpoint["overflow-target"]) != 0 { + t.Fatalf("incoming sentinel leaked rows = %#v, %v", projection, err) + } +} + +type cancelSQLiteAdjacencyAfterChecksContext struct { + context.Context + remaining int +} + +func (ctx *cancelSQLiteAdjacencyAfterChecksContext) Err() error { + if ctx.remaining == 0 { + return context.Canceled + } + ctx.remaining-- + return nil +} + +func TestSQLiteBoundedAdjacencyMidScanCancellationReturnsNoPartial(t *testing.T) { + store := openBoundedAdjacencyTestStore(t) + const source = "cancel-source" + edges := make([]*graph.Edge, 0, graph.MaxBoundedAdjacencyRowsPerKey) + for index := 0; index < graph.MaxBoundedAdjacencyRowsPerKey; index++ { + edges = append(edges, &graph.Edge{From: source, To: fmt.Sprintf("target-%03d", index), Kind: graph.EdgeCalls, Line: index}) + } + store.AddBatch(nil, edges) + rows, err := store.db.Query(boundedOutgoingAdjacencySQL(1), source, string(graph.EdgeCalls), graph.MaxBoundedAdjacencyRowsPerKey+1) + if err != nil { + t.Fatalf("open rows: %v", err) + } + total := 0 + ctx := &cancelSQLiteAdjacencyAfterChecksContext{Context: context.Background(), remaining: 1} + identities, truncated, err := scanSQLiteBoundedAdjacencyRows(ctx, rows, graph.MaxBoundedAdjacencyRowsPerKey, &total) + if !errors.Is(err, context.Canceled) || len(identities) != 0 || truncated { + t.Fatalf("canceled scan leaked partial identities: len=%d truncated=%v err=%v", len(identities), truncated, err) + } + if err := rows.Err(); err != nil && !errors.Is(err, context.Canceled) { + t.Fatalf("rows: %v", err) + } + store.AddEdge(&graph.Edge{From: "after-cancel", To: "target", Kind: graph.EdgeCalls}) +} + +func TestSQLiteBoundedAdjacencyPlansUsePredicateIndexes(t *testing.T) { + store := openBoundedAdjacencyTestStore(t) + for _, test := range []struct { + name string + query string + args []any + index string + constraints string + }{ + {name: "outgoing", query: boundedOutgoingAdjacencySQL(1), args: []any{"source", string(graph.EdgeCalls), 2}, index: "EDGES_BY_FROM", constraints: "(FROM_ID=? AND KIND=?)"}, + {name: "incoming", query: boundedIncomingAdjacencySQL(1), args: []any{"target", string(graph.EdgeCalls), 2}, index: "EDGES_BY_TO", constraints: "(TO_ID=? AND KIND=?)"}, + {name: "site", query: boundedOutgoingSiteAdjacencySQL(1), args: []any{"source", 10, string(graph.EdgeCalls), 2}, index: "EDGES_BY_FROM_LINE_KIND", constraints: "(FROM_ID=? AND LINE=? AND KIND=?)"}, + } { + t.Run(test.name, func(t *testing.T) { + rows, err := store.db.Query("EXPLAIN QUERY PLAN "+test.query, test.args...) + if err != nil { + t.Fatalf("explain: %v", err) + } + defer rows.Close() + var details []string + for rows.Next() { + var id, parent, unused int + var detail string + if err := rows.Scan(&id, &parent, &unused, &detail); err != nil { + t.Fatalf("scan plan: %v", err) + } + details = append(details, detail) + } + if err := rows.Err(); err != nil { + t.Fatalf("plan rows: %v", err) + } + plan := strings.ToUpper(strings.Join(details, "\n")) + if strings.Contains(plan, "SCAN EDGES") || strings.Contains(plan, "ORDER BY") || strings.Contains(plan, "TEMP B-TREE") { + t.Fatalf("bounded adjacency plan scans/sorts:\n%s", plan) + } + if !strings.Contains(plan, "USING INDEX "+test.index) || !strings.Contains(plan, test.constraints) { + t.Fatalf("bounded adjacency plan missed %s %s:\n%s", test.index, test.constraints, plan) + } + }) + } +} + +func TestSQLiteEdgesByFromLineIndexesCoexistAndOpenIdempotently(t *testing.T) { + path := filepath.Join(t.TempDir(), "graph.sqlite") + store, err := Open(path) + if err != nil { + t.Fatalf("open store: %v", err) + } + if _, err := store.writerDB.Exec(`DROP INDEX edges_by_from_line_kind`); err != nil { + t.Fatalf("drop bounded-site index: %v", err) + } + if _, err := store.writerDB.Exec(`CREATE INDEX IF NOT EXISTS edges_by_from_line ON edges(from_id, line)`); err != nil { + t.Fatalf("create historical index: %v", err) + } + if err := store.Close(); err != nil { + t.Fatalf("close store without bounded-site index: %v", err) + } + + for reopen := 0; reopen < 2; reopen++ { + store, err = Open(path) + if err != nil { + t.Fatalf("reopen %d: %v", reopen, err) + } + rows, err := store.db.Query(`PRAGMA index_info('edges_by_from_line')`) + if err != nil { + t.Fatalf("index info: %v", err) + } + var columns []string + for rows.Next() { + var sequence, columnID int + var name string + if err := rows.Scan(&sequence, &columnID, &name); err != nil { + t.Fatalf("scan index info: %v", err) + } + columns = append(columns, name) + } + if err := rows.Err(); err != nil { + t.Fatalf("index info rows: %v", err) + } + _ = rows.Close() + if !reflect.DeepEqual(columns, []string{"from_id", "line"}) { + t.Fatalf("reopen %d legacy columns = %#v", reopen, columns) + } + rows, err = store.db.Query(`PRAGMA index_info('edges_by_from_line_kind')`) + if err != nil { + t.Fatalf("bounded-site index info: %v", err) + } + columns = columns[:0] + for rows.Next() { + var sequence, columnID int + var name string + if err := rows.Scan(&sequence, &columnID, &name); err != nil { + t.Fatalf("scan bounded-site index info: %v", err) + } + columns = append(columns, name) + } + if err := rows.Err(); err != nil { + t.Fatalf("bounded-site index rows: %v", err) + } + _ = rows.Close() + if !reflect.DeepEqual(columns, []string{"from_id", "line", "kind"}) { + t.Fatalf("reopen %d bounded-site columns = %#v", reopen, columns) + } + var schemaBefore, schemaAfter int + if err := store.writerDB.QueryRow(`PRAGMA schema_version`).Scan(&schemaBefore); err != nil { + t.Fatalf("schema version before ensure: %v", err) + } + if _, err := store.writerDB.Exec(edgesByFromLineKindIndexDDL); err != nil { + t.Fatalf("idempotent bounded-site index ensure: %v", err) + } + if err := store.writerDB.QueryRow(`PRAGMA schema_version`).Scan(&schemaAfter); err != nil { + t.Fatalf("schema version after ensure: %v", err) + } + if schemaAfter != schemaBefore { + t.Fatalf("correct index shape rebuilt: schema version %d -> %d", schemaBefore, schemaAfter) + } + var siblingCount int + if err := store.db.QueryRow(`SELECT COUNT(*) FROM sqlite_master WHERE type='index' AND name LIKE 'edges_by_from_line%'`).Scan(&siblingCount); err != nil { + t.Fatalf("count site indexes: %v", err) + } + if siblingCount != 2 { + t.Fatalf("reopen %d site index count = %d, want 2", reopen, siblingCount) + } + if err := store.Close(); err != nil { + t.Fatalf("close reopen %d: %v", reopen, err) + } + } +} diff --git a/internal/graph/store_sqlite/bounded_edge_existence.go b/internal/graph/store_sqlite/bounded_edge_existence.go new file mode 100644 index 00000000..1a8f42f0 --- /dev/null +++ b/internal/graph/store_sqlite/bounded_edge_existence.go @@ -0,0 +1,161 @@ +package store_sqlite + +import ( + "context" + "database/sql" + "fmt" + "sort" + "strings" + + "github.com/zzet/gortex/internal/graph" +) + +var _ graph.BoundedEdgeExistenceReader = (*Store)(nil) + +func canonicalBoundedEdgeExistenceKeys( + ctx context.Context, + endpoints []graph.TypedEdgeEndpoint, + limit int, +) ([]graph.TypedEdgeEndpoint, error) { + if ctx == nil { + ctx = context.Background() + } + if err := ctx.Err(); err != nil { + return nil, err + } + if limit < 1 || limit > graph.MaxBoundedEdgeExistencePredicates { + return nil, &graph.BoundedLocalizationLimitError{ + Resource: "SQLite edge-existence predicates", + Limit: graph.MaxBoundedEdgeExistencePredicates, + } + } + if len(endpoints) == 0 { + return nil, nil + } + capacity := len(endpoints) + if capacity > limit+1 { + capacity = limit + 1 + } + seen := make(map[graph.TypedEdgeEndpoint]struct{}, capacity) + keys := make([]graph.TypedEdgeEndpoint, 0, capacity) + for index, endpoint := range endpoints { + if index&127 == 0 { + if err := ctx.Err(); err != nil { + return nil, err + } + } + if endpoint.From == "" || endpoint.To == "" || endpoint.Kind == "" { + continue + } + if _, duplicate := seen[endpoint]; duplicate { + continue + } + seen[endpoint] = struct{}{} + if len(seen) > limit { + return nil, &graph.BoundedLocalizationLimitError{ + Resource: "SQLite edge-existence predicates", + Limit: limit, + } + } + keys = append(keys, endpoint) + } + sort.Slice(keys, func(i, j int) bool { + if keys[i].From != keys[j].From { + return keys[i].From < keys[j].From + } + if keys[i].To != keys[j].To { + return keys[i].To < keys[j].To + } + return keys[i].Kind < keys[j].Kind + }) + return keys, nil +} + +func boundedEdgeExistenceQuery(rows int) string { + var query strings.Builder + query.Grow(192 + rows*8) + query.WriteString("WITH wanted(from_id, to_id, kind) AS (VALUES ") + for index := 0; index < rows; index++ { + if index > 0 { + query.WriteByte(',') + } + query.WriteString("(?,?,?)") + } + query.WriteString(`) +SELECT wanted.from_id, wanted.to_id, wanted.kind +FROM wanted +WHERE EXISTS ( + SELECT 1 + FROM edges + WHERE edges.from_id = wanted.from_id + AND edges.to_id = wanted.to_id + AND edges.kind = wanted.kind + LIMIT 1 +)`) + return query.String() +} + +// FindExistingEdgeEndpoints projects only exact typed endpoint identities. A +// short read-only transaction freezes the chunked lookup; each cursor is fully +// closed before the next query, and correlated EXISTS stops at the first +// physical edge row for each requested predicate. +func (s *Store) FindExistingEdgeEndpoints( + ctx context.Context, + endpoints []graph.TypedEdgeEndpoint, + limit int, +) (map[graph.TypedEdgeEndpoint]struct{}, error) { + if ctx == nil { + ctx = context.Background() + } + keys, err := canonicalBoundedEdgeExistenceKeys(ctx, endpoints, limit) + if err != nil { + return nil, err + } + found := make(map[graph.TypedEdgeEndpoint]struct{}, len(keys)) + if len(keys) == 0 { + return found, nil + } + tx, err := s.db.BeginTx(ctx, &sql.TxOptions{ReadOnly: true}) + if err != nil { + return nil, err + } + defer func() { _ = tx.Rollback() }() + + for start := 0; start < len(keys); start += lookupChunkSize { + if err := ctx.Err(); err != nil { + return nil, err + } + end := start + lookupChunkSize + if end > len(keys) { + end = len(keys) + } + chunk := keys[start:end] + args := make([]any, 0, len(chunk)*3) + for _, endpoint := range chunk { + args = append(args, endpoint.From, endpoint.To, string(endpoint.Kind)) + } + rows, queryErr := tx.QueryContext(ctx, boundedEdgeExistenceQuery(len(chunk)), args...) + if queryErr != nil { + return nil, queryErr + } + for rows.Next() { + var from, to, kind string + if scanErr := rows.Scan(&from, &to, &kind); scanErr != nil { + _ = rows.Close() + return nil, scanErr + } + found[graph.TypedEdgeEndpoint{From: from, To: to, Kind: graph.EdgeKind(kind)}] = struct{}{} + } + if rowsErr := rows.Err(); rowsErr != nil { + _ = rows.Close() + return nil, fmt.Errorf("bounded edge-existence query: %w", rowsErr) + } + if closeErr := rows.Close(); closeErr != nil { + return nil, closeErr + } + } + if err := tx.Commit(); err != nil { + return nil, err + } + return found, nil +} diff --git a/internal/graph/store_sqlite/bounded_edge_existence_test.go b/internal/graph/store_sqlite/bounded_edge_existence_test.go new file mode 100644 index 00000000..4173a219 --- /dev/null +++ b/internal/graph/store_sqlite/bounded_edge_existence_test.go @@ -0,0 +1,149 @@ +package store_sqlite + +import ( + "context" + "errors" + "fmt" + "path/filepath" + "reflect" + "strings" + "testing" + + "github.com/zzet/gortex/internal/graph" +) + +func TestFindExistingEdgeEndpointsProjectsDistinctPredicatesAndReenters(t *testing.T) { + store, err := Open(filepath.Join(t.TempDir(), "graph.sqlite")) + if err != nil { + t.Fatalf("open store: %v", err) + } + t.Cleanup(func() { _ = store.Close() }) + + const ( + from = "repo/source.go::source" + owner = "repo/owner.go::Owner" + ) + edges := make([]*graph.Edge, 0, 2_002) + for line := 1; line <= 2_000; line++ { + edges = append(edges, &graph.Edge{ + From: from, To: owner, Kind: graph.EdgeMemberOf, + FilePath: "repo/source.go", Line: line, + Meta: map[string]any{"payload": strings.Repeat("x", 1024)}, + }) + } + edges = append(edges, + &graph.Edge{From: from, To: owner, Kind: graph.EdgeCalls}, + &graph.Edge{From: from, To: "repo/other.go::Other", Kind: graph.EdgeMemberOf}, + ) + store.AddBatch(nil, edges) + + member := graph.TypedEdgeEndpoint{From: from, To: owner, Kind: graph.EdgeMemberOf} + calls := graph.TypedEdgeEndpoint{From: from, To: owner, Kind: graph.EdgeCalls} + missing := graph.TypedEdgeEndpoint{From: from, To: "missing", Kind: graph.EdgeMemberOf} + input := []graph.TypedEdgeEndpoint{member, member, calls, missing, {}} + before := append([]graph.TypedEdgeEndpoint(nil), input...) + found, err := store.FindExistingEdgeEndpoints(context.Background(), input, 4) + if err != nil { + t.Fatalf("bounded edge existence: %v", err) + } + if !reflect.DeepEqual(input, before) { + t.Fatalf("caller predicates mutated: got %#v want %#v", input, before) + } + if len(found) != 2 { + t.Fatalf("found = %#v, want exactly member_of and calls", found) + } + if _, ok := found[member]; !ok { + t.Fatalf("duplicate-heavy exact endpoint missing: %#v", found) + } + if _, ok := found[calls]; !ok { + t.Fatalf("kind-specific endpoint missing: %#v", found) + } + + // The read-only transaction and cursor must be closed before return. + newKey := graph.TypedEdgeEndpoint{From: "new-source", To: "new-target", Kind: graph.EdgeCalls} + store.AddEdge(&graph.Edge{From: newKey.From, To: newKey.To, Kind: newKey.Kind}) + found, err = store.FindExistingEdgeEndpoints(context.Background(), []graph.TypedEdgeEndpoint{newKey}, 1) + if err != nil { + t.Fatalf("re-entered projection: %v", err) + } + if _, ok := found[newKey]; !ok { + t.Fatalf("re-entered projection missed newly committed edge: %#v", found) + } +} + +func TestFindExistingEdgeEndpointsEnforcesHardCapAndCancellation(t *testing.T) { + store, err := Open(filepath.Join(t.TempDir(), "graph.sqlite")) + if err != nil { + t.Fatalf("open store: %v", err) + } + t.Cleanup(func() { _ = store.Close() }) + + keys := make([]graph.TypedEdgeEndpoint, 0, graph.MaxBoundedEdgeExistencePredicates+1) + for index := 0; index <= graph.MaxBoundedEdgeExistencePredicates; index++ { + keys = append(keys, graph.TypedEdgeEndpoint{ + From: "source", To: fmt.Sprintf("target-%03d", index), Kind: graph.EdgeCalls, + }) + } + if found, err := store.FindExistingEdgeEndpoints( + context.Background(), keys[:graph.MaxBoundedEdgeExistencePredicates], graph.MaxBoundedEdgeExistencePredicates, + ); err != nil || len(found) != 0 { + t.Fatalf("exact boundary = %#v, %v", found, err) + } + found, err := store.FindExistingEdgeEndpoints(context.Background(), keys, graph.MaxBoundedEdgeExistencePredicates) + var limitErr *graph.BoundedLocalizationLimitError + if !errors.As(err, &limitErr) || len(found) != 0 || limitErr.Limit != graph.MaxBoundedEdgeExistencePredicates { + t.Fatalf("over boundary = %#v, %v", found, err) + } + for _, invalidLimit := range []int{0, graph.MaxBoundedEdgeExistencePredicates + 1} { + found, err = store.FindExistingEdgeEndpoints(context.Background(), nil, invalidLimit) + limitErr = nil + if !errors.As(err, &limitErr) || len(found) != 0 || limitErr.Limit != graph.MaxBoundedEdgeExistencePredicates { + t.Fatalf("empty input with invalid limit %d = %#v, %v", invalidLimit, found, err) + } + } + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + found, err = store.FindExistingEdgeEndpoints(ctx, keys[:1], 1) + if !errors.Is(err, context.Canceled) || len(found) != 0 { + t.Fatalf("cancelled projection returned partial evidence: %#v, %v", found, err) + } +} + +func TestFindExistingEdgeEndpointsPlanUsesExactIndexWithoutSorter(t *testing.T) { + store, err := Open(filepath.Join(t.TempDir(), "graph.sqlite")) + if err != nil { + t.Fatalf("open store: %v", err) + } + t.Cleanup(func() { _ = store.Close() }) + store.AddEdge(&graph.Edge{From: "source", To: "target", Kind: graph.EdgeMemberOf}) + + rows, err := store.db.Query( + "EXPLAIN QUERY PLAN "+boundedEdgeExistenceQuery(1), + "source", "target", string(graph.EdgeMemberOf), + ) + if err != nil { + t.Fatalf("explain edge existence: %v", err) + } + defer rows.Close() + var details []string + for rows.Next() { + var id, parent, unused int + var detail string + if err := rows.Scan(&id, &parent, &unused, &detail); err != nil { + t.Fatalf("scan query plan: %v", err) + } + details = append(details, detail) + } + if err := rows.Err(); err != nil { + t.Fatalf("query plan rows: %v", err) + } + plan := strings.ToUpper(strings.Join(details, "\n")) + if strings.Contains(plan, "SCAN EDGES") || strings.Contains(plan, "ORDER BY") { + t.Fatalf("bounded exact endpoint query scans/sorts edges:\n%s", plan) + } + if !strings.Contains(plan, "USING COVERING INDEX SQLITE_AUTOINDEX_EDGES_1") || + !strings.Contains(plan, "(FROM_ID=? AND TO_ID=? AND KIND=?)") { + t.Fatalf("bounded exact endpoint query did not use the exact covering edge key:\n%s", plan) + } +} diff --git a/internal/graph/store_sqlite/bounded_incoming_sources.go b/internal/graph/store_sqlite/bounded_incoming_sources.go new file mode 100644 index 00000000..bdc82eb6 --- /dev/null +++ b/internal/graph/store_sqlite/bounded_incoming_sources.go @@ -0,0 +1,105 @@ +package store_sqlite + +import ( + "context" + "database/sql" + "fmt" + "sort" + + "github.com/zzet/gortex/internal/graph" +) + +var _ graph.BoundedIncomingSourceReader = (*Store)(nil) + +const findIncomingSourcesBoundedSQL = `SELECT DISTINCT from_id + FROM edges + WHERE to_id = ? AND kind = ? AND from_id <> '' + LIMIT ?` + +// FindIncomingSourcesBounded projects distinct incoming source identities for +// one edge kind. One read-only transaction freezes the batch snapshot; each +// target cursor is closed before the next query so callers can safely perform +// later store reads without cursor re-entry. +func (s *Store) FindIncomingSourcesBounded( + ctx context.Context, + targetIDs []string, + kind graph.EdgeKind, + limit int, +) (graph.BoundedIncomingSourceProjection, error) { + projection := graph.BoundedIncomingSourceProjection{ + Sources: make(map[string][]string), + Truncated: make(map[string]bool), + } + ids := dedupeNonEmpty(targetIDs) + if len(ids) == 0 { + return projection, nil + } + if ctx == nil { + ctx = context.Background() + } + if err := ctx.Err(); err != nil { + return graph.BoundedIncomingSourceProjection{}, err + } + if limit <= 0 { + for _, id := range ids { + projection.Truncated[id] = true + } + return projection, nil + } + maxInt := int(^uint(0) >> 1) + if limit >= maxInt { + return graph.BoundedIncomingSourceProjection{}, &graph.BoundedLocalizationLimitError{ + Resource: "SQLite incoming-source sentinel", + Limit: maxInt - 1, + } + } + + tx, err := s.db.BeginTx(ctx, &sql.TxOptions{ReadOnly: true}) + if err != nil { + return graph.BoundedIncomingSourceProjection{}, err + } + defer func() { _ = tx.Rollback() }() + + for _, targetID := range ids { + if err := ctx.Err(); err != nil { + return graph.BoundedIncomingSourceProjection{}, err + } + rows, queryErr := tx.QueryContext(ctx, findIncomingSourcesBoundedSQL, targetID, kind, limit+1) + if queryErr != nil { + return graph.BoundedIncomingSourceProjection{}, queryErr + } + sources := make([]string, 0, limit) + truncated := false + for rows.Next() { + var sourceID string + if scanErr := rows.Scan(&sourceID); scanErr != nil { + _ = rows.Close() + return graph.BoundedIncomingSourceProjection{}, scanErr + } + if len(sources) >= limit { + truncated = true + break + } + sources = append(sources, sourceID) + } + if rowsErr := rows.Err(); rowsErr != nil { + _ = rows.Close() + return graph.BoundedIncomingSourceProjection{}, fmt.Errorf("bounded incoming-source query: %w", rowsErr) + } + if closeErr := rows.Close(); closeErr != nil { + return graph.BoundedIncomingSourceProjection{}, closeErr + } + if truncated { + projection.Truncated[targetID] = true + continue + } + if len(sources) > 0 { + sort.Strings(sources) + projection.Sources[targetID] = sources + } + } + if err := tx.Commit(); err != nil { + return graph.BoundedIncomingSourceProjection{}, err + } + return projection, nil +} diff --git a/internal/graph/store_sqlite/bounded_incoming_sources_test.go b/internal/graph/store_sqlite/bounded_incoming_sources_test.go new file mode 100644 index 00000000..c998c645 --- /dev/null +++ b/internal/graph/store_sqlite/bounded_incoming_sources_test.go @@ -0,0 +1,178 @@ +package store_sqlite + +import ( + "context" + "errors" + "fmt" + "path/filepath" + "reflect" + "strings" + "testing" + + "github.com/zzet/gortex/internal/graph" +) + +func TestFindIncomingSourcesBoundedCountsDistinctRelevantRowsAndReenters(t *testing.T) { + store, err := Open(filepath.Join(t.TempDir(), "graph.sqlite")) + if err != nil { + t.Fatalf("open store: %v", err) + } + t.Cleanup(func() { _ = store.Close() }) + const targetID = "repo/target.go::target" + edges := make([]*graph.Edge, 0, 32) + for index := 0; index < 8; index++ { + sourceID := fmt.Sprintf("repo/source-%02d.go::source", index) + for line := 1; line <= 3; line++ { + edges = append(edges, &graph.Edge{From: sourceID, To: targetID, Kind: graph.EdgeCalls, Line: line}) + } + edges = append(edges, &graph.Edge{From: fmt.Sprintf("repo/noise-%02d.go::noise", index), To: targetID, Kind: graph.EdgeReferences}) + } + store.AddBatch(nil, edges) + + page, err := store.FindIncomingSourcesBounded(context.Background(), []string{targetID}, graph.EdgeCalls, 8) + if err != nil { + t.Fatalf("bounded incoming projection: %v", err) + } + if page.Truncated[targetID] || len(page.Sources[targetID]) != 8 { + t.Fatalf("projection = %#v, want eight distinct CALLS sources", page) + } + + // The projection must close its cursor and read transaction before return; + // this write and second read would otherwise expose store re-entry trouble. + store.AddEdge(&graph.Edge{From: "repo/source-08.go::source", To: targetID, Kind: graph.EdgeCalls}) + page, err = store.FindIncomingSourcesBounded(context.Background(), []string{targetID, "missing"}, graph.EdgeCalls, 8) + if err != nil { + t.Fatalf("re-entered saturated projection: %v", err) + } + if !page.Truncated[targetID] || len(page.Sources[targetID]) != 0 || page.Truncated["missing"] { + t.Fatalf("re-entered projection = %#v", page) + } +} + +func TestFindIncomingSourcesBoundedDuplicateHeavySentinel(t *testing.T) { + store, err := Open(filepath.Join(t.TempDir(), "graph.sqlite")) + if err != nil { + t.Fatalf("open store: %v", err) + } + t.Cleanup(func() { _ = store.Close() }) + const targetID = "target" + edges := make([]*graph.Edge, 0, 2_001) + for line := 1; line <= 2_000; line++ { + edges = append(edges, &graph.Edge{From: "one-source", To: targetID, Kind: graph.EdgeCalls, Line: line}) + } + store.AddBatch(nil, edges) + page, err := store.FindIncomingSourcesBounded(context.Background(), []string{targetID}, graph.EdgeCalls, 1) + if err != nil || page.Truncated[targetID] || len(page.Sources[targetID]) != 1 { + t.Fatalf("duplicate-heavy projection = %#v, %v; duplicates consumed sentinel", page, err) + } + store.AddEdge(&graph.Edge{From: "second-source", To: targetID, Kind: graph.EdgeCalls}) + page, err = store.FindIncomingSourcesBounded(context.Background(), []string{targetID}, graph.EdgeCalls, 1) + if err != nil || !page.Truncated[targetID] || len(page.Sources[targetID]) != 0 { + t.Fatalf("second distinct source did not saturate: %#v, %v", page, err) + } +} + +func TestFindIncomingSourcesBoundedRejectsImpossibleLimitAndCancellation(t *testing.T) { + store, err := Open(filepath.Join(t.TempDir(), "graph.sqlite")) + if err != nil { + t.Fatalf("open store: %v", err) + } + t.Cleanup(func() { _ = store.Close() }) + store.AddEdge(&graph.Edge{From: "source", To: "target", Kind: graph.EdgeCalls}) + + maxInt := int(^uint(0) >> 1) + page, err := store.FindIncomingSourcesBounded(context.Background(), []string{"target"}, graph.EdgeCalls, maxInt) + var limitErr *graph.BoundedLocalizationLimitError + if !errors.As(err, &limitErr) || len(page.Sources) != 0 { + t.Fatalf("max-int projection = %#v, %v; want typed empty failure", page, err) + } + ctx, cancel := context.WithCancel(context.Background()) + cancel() + page, err = store.FindIncomingSourcesBounded(ctx, []string{"target"}, graph.EdgeCalls, 8) + if !errors.Is(err, context.Canceled) || len(page.Sources) != 0 { + t.Fatalf("cancelled projection = %#v, %v", page, err) + } + + layer := graph.NewOverlayLayer() + layer.AddNode("repo/overlay.go", &graph.Node{ID: "overlay", Name: "overlay", Kind: graph.KindFunction, FilePath: "repo/overlay.go"}) + ids := []string{"overlay", "target"} + if nodes, err := graph.NewOverlaidView(store, layer).GetNodesByIDsContext(ctx, ids); !errors.Is(err, context.Canceled) || len(nodes) != 0 { + t.Fatalf("cancelled SQLite-backed overlay refetch = %#v, %v", nodes, err) + } +} + +type contextualSQLiteNodeSpy struct { + *Store + contextCalls int + legacyCalls int +} + +func (spy *contextualSQLiteNodeSpy) GetNodesByIDs(ids []string) map[string]*graph.Node { + spy.legacyCalls++ + return spy.Store.GetNodesByIDs(ids) +} + +func (spy *contextualSQLiteNodeSpy) GetNodesByIDsContext(context.Context, []string) (map[string]*graph.Node, error) { + spy.contextCalls++ + return map[string]*graph.Node{"base": {ID: "base", Name: "base", Kind: graph.KindFunction}}, context.Canceled +} + +func TestOverlaidViewContextualExactRefetchUsesSQLiteContextPathAndFailsClosed(t *testing.T) { + store, err := Open(filepath.Join(t.TempDir(), "graph.sqlite")) + if err != nil { + t.Fatalf("open store: %v", err) + } + t.Cleanup(func() { _ = store.Close() }) + store.AddNode(&graph.Node{ID: "base", Name: "base", Kind: graph.KindFunction}) + spy := &contextualSQLiteNodeSpy{Store: store} + layer := graph.NewOverlayLayer() + layer.AddNode("repo/overlay.go", &graph.Node{ + ID: "repo/overlay.go::overlay", Name: "overlay", Kind: graph.KindFunction, FilePath: "repo/overlay.go", + }) + ids := []string{"repo/overlay.go::overlay", "base"} + before := append([]string(nil), ids...) + nodes, err := graph.NewOverlaidView(spy, layer).GetNodesByIDsContext(context.Background(), ids) + if !errors.Is(err, context.Canceled) || len(nodes) != 0 { + t.Fatalf("contextual SQLite error returned partial overlay nodes: %#v, %v", nodes, err) + } + if spy.contextCalls != 1 || spy.legacyCalls != 0 || !reflect.DeepEqual(ids, before) { + t.Fatalf("dispatch/caller slice = contextual:%d legacy:%d ids:%v want:%v", spy.contextCalls, spy.legacyCalls, ids, before) + } +} + +func TestFindIncomingSourcesBoundedPlanAvoidsOrderBySorter(t *testing.T) { + store, err := Open(filepath.Join(t.TempDir(), "graph.sqlite")) + if err != nil { + t.Fatalf("open store: %v", err) + } + t.Cleanup(func() { _ = store.Close() }) + store.AddEdge(&graph.Edge{From: "source", To: "target", Kind: graph.EdgeCalls}) + + rows, err := store.db.Query( + "EXPLAIN QUERY PLAN "+findIncomingSourcesBoundedSQL, + "target", graph.EdgeCalls, 17, + ) + if err != nil { + t.Fatalf("explain incoming projection: %v", err) + } + defer rows.Close() + var details []string + for rows.Next() { + var id, parent, unused int + var detail string + if err := rows.Scan(&id, &parent, &unused, &detail); err != nil { + t.Fatalf("scan query plan: %v", err) + } + details = append(details, detail) + } + if err := rows.Err(); err != nil { + t.Fatalf("query plan rows: %v", err) + } + plan := strings.ToUpper(strings.Join(details, "\n")) + if strings.Contains(plan, "ORDER BY") { + t.Fatalf("bounded incoming query materializes an ORDER BY sorter:\n%s", plan) + } + if !strings.Contains(plan, "EDGES_BY_TO") { + t.Fatalf("bounded incoming query did not use target/kind index:\n%s", plan) + } +} diff --git a/internal/graph/store_sqlite/bulk_load.go b/internal/graph/store_sqlite/bulk_load.go index 0d6e0339..2fe28d3f 100644 --- a/internal/graph/store_sqlite/bulk_load.go +++ b/internal/graph/store_sqlite/bulk_load.go @@ -82,6 +82,11 @@ const coldFTSMergePages = 64 // // Dropping/recreating these is a runtime operation on identical DDL — it is // NOT a schema change, so it does not touch the persisted schema version. +const ( + edgesByFromLineIndexDDL = `CREATE INDEX IF NOT EXISTS edges_by_from_line ON edges(from_id, line)` + edgesByFromLineKindIndexDDL = `CREATE INDEX IF NOT EXISTS edges_by_from_line_kind ON edges(from_id, line, kind)` +) + var bulkDroppableIndexes = []bulkDroppableIndex{ {"nodes_by_name", `CREATE INDEX IF NOT EXISTS nodes_by_name ON nodes(name)`}, {"nodes_by_kind", `CREATE INDEX IF NOT EXISTS nodes_by_kind ON nodes(kind)`}, @@ -113,7 +118,13 @@ var bulkDroppableIndexes = []bulkDroppableIndex{ // into tens of milliseconds, and the cross-package guard alone paid ~980s // of a 28-repo cold index that way. With the index the full candidate // path measured ~30 → ~116k sites/s on a production store copy. - {"edges_by_from_line", `CREATE INDEX IF NOT EXISTS edges_by_from_line ON edges(from_id, line)`}, + // Preserve this exact two-column shape: ordered outgoing projections rely + // on the rowid/primary-key suffix after line and otherwise need a temp sort. + {"edges_by_from_line", edgesByFromLineIndexDDL}, + // Bounded exact-site adjacency also constrains kind before LIMIT. Keep that + // predicate-shaped index separate so it cannot perturb the legacy ordered + // outgoing plan; both participate in the same bulk drop/rebuild lifecycle. + {"edges_by_from_line_kind", edgesByFromLineKindIndexDDL}, {"edges_by_to", `CREATE INDEX IF NOT EXISTS edges_by_to ON edges(to_id, kind)`}, {"edges_by_kind", `CREATE INDEX IF NOT EXISTS edges_by_kind ON edges(kind)`}, // Exact changed-file frontiers (watcher and partial indexing) must not diff --git a/internal/graph/store_sqlite/localization_projection.go b/internal/graph/store_sqlite/localization_projection.go new file mode 100644 index 00000000..f6bccb6d --- /dev/null +++ b/internal/graph/store_sqlite/localization_projection.go @@ -0,0 +1,313 @@ +package store_sqlite + +import ( + "context" + "database/sql" + "sort" + "strings" + + "github.com/zzet/gortex/internal/graph" +) + +var ( + _ graph.BoundedExactNameReader = (*Store)(nil) + _ graph.BoundedFileNodeReader = (*Store)(nil) +) + +// lookupLocalizationNodeCols carries identity, kind, location, and scope plus +// the opaque metadata blob needed to preserve the exact anchor's is_test gate. +// Unlike lookupNodeCols it excludes docs, signatures, retrieval text, and every +// other promoted payload. Raw metadata work is capped to 256 rows per closed +// keyset page; pages continue only until limit+1 eligible declarations prove +// saturation. Promoting is_test later would remove the remaining skipped-row +// blob decoding without changing this projection contract. +var lookupLocalizationNodeCols = lookupNodeSummaryCols + `, meta` + +// FindNodesByNameBounded performs the exact-name localization lookup inside +// one short read-only transaction. Indexed scope predicates and keyset pages +// keep every cursor and allocation bounded. is_test remains in legacy metadata, +// so each 256-row page is decoded and closed before the next; scanning stops as +// soon as limit+1 production declarations prove ambiguity. +func (s *Store) FindNodesByNameBounded( + ctx context.Context, + name string, + scope graph.LocalizationNodeScope, + limit int, +) (graph.BoundedNodeProjection, error) { + if s == nil || s.db == nil || name == "" || limit <= 0 { + return graph.BoundedNodeProjection{}, nil + } + if ctx == nil { + ctx = context.Background() + } + if err := ctx.Err(); err != nil { + return graph.BoundedNodeProjection{}, err + } + + predicate, args := localizationNodePredicate(name, scope) + tx, err := s.db.BeginTx(ctx, &sql.TxOptions{ReadOnly: true}) + if err != nil { + return graph.BoundedNodeProjection{}, err + } + defer func() { _ = tx.Rollback() }() + + const rawPageSize = 256 + sentinel := limit + 1 + nodes := make([]*graph.Node, 0, sentinel) + lastID := "" + for len(nodes) < sentinel { + if err := ctx.Err(); err != nil { + return graph.BoundedNodeProjection{}, err + } + pageArgs := append(append([]any(nil), args...), lastID, rawPageSize) + rows, queryErr := tx.QueryContext( + ctx, + `SELECT `+lookupLocalizationNodeCols+` FROM nodes WHERE `+predicate+` AND id > ? ORDER BY id LIMIT ?`, + pageArgs..., + ) + if queryErr != nil { + return graph.BoundedNodeProjection{}, queryErr + } + rawRows := 0 + for rows.Next() { + node, scanErr := scanLocalizationNode(rows) + if scanErr != nil { + _ = rows.Close() + return graph.BoundedNodeProjection{}, scanErr + } + rawRows++ + lastID = node.ID + if !scope.Allows(node) { + continue + } + nodes = append(nodes, node) + if len(nodes) == sentinel { + break + } + } + if err := rows.Err(); err != nil { + _ = rows.Close() + return graph.BoundedNodeProjection{}, err + } + if err := rows.Close(); err != nil { + return graph.BoundedNodeProjection{}, err + } + if len(nodes) == sentinel || rawRows < rawPageSize { + break + } + } + if err := tx.Commit(); err != nil { + return graph.BoundedNodeProjection{}, err + } + + total := len(nodes) + truncated := total > limit + if len(nodes) > limit { + nodes = nodes[:limit] + } + return graph.BoundedNodeProjection{Nodes: nodes, Total: total, Truncated: truncated}, nil +} + +// FindFileNodesBounded reads one file through the identity/location projection. +// Scope and kind predicates are pushed before the sentinel cap. The normal +// declaration path transfers no metadata column at all; callers that explicitly +// exclude tests pay bounded RawBytes decoding because is_test still lives in the +// legacy metadata blob. Every keyset cursor is closed before the next page. +func (s *Store) FindFileNodesBounded( + ctx context.Context, + filePath string, + scope graph.LocalizationNodeScope, + limit int, +) (graph.BoundedNodeProjection, error) { + if s == nil || s.db == nil || filePath == "" || limit <= 0 { + return graph.BoundedNodeProjection{}, nil + } + if ctx == nil { + ctx = context.Background() + } + if err := ctx.Err(); err != nil { + return graph.BoundedNodeProjection{}, err + } + if scope.ExcludesFile(filePath) { + return graph.BoundedNodeProjection{}, nil + } + + predicate, args := localizationFileNodePredicate(filePath, scope) + tx, err := s.db.BeginTx(ctx, &sql.TxOptions{ReadOnly: true}) + if err != nil { + return graph.BoundedNodeProjection{}, err + } + defer func() { _ = tx.Rollback() }() + + columns := lookupNodeSummaryCols + if scope.ExcludeTests { + columns = lookupLocalizationNodeCols + } + const rawPageSize = 256 + sentinel := limit + 1 + nodes := make([]*graph.Node, 0, sentinel) + lastID := "" + for len(nodes) < sentinel { + if err := ctx.Err(); err != nil { + return graph.BoundedNodeProjection{}, err + } + pageArgs := append(append([]any(nil), args...), lastID, rawPageSize) + rows, queryErr := tx.QueryContext( + ctx, + `SELECT `+columns+` FROM nodes WHERE `+predicate+` AND id > ? ORDER BY id LIMIT ?`, + pageArgs..., + ) + if queryErr != nil { + return graph.BoundedNodeProjection{}, queryErr + } + rawRows := 0 + for rows.Next() { + var node *graph.Node + var scanErr error + if scope.ExcludeTests { + node, scanErr = scanLocalizationNode(rows) + } else { + node, scanErr = scanNodeSummary(rows) + } + if scanErr != nil { + _ = rows.Close() + return graph.BoundedNodeProjection{}, scanErr + } + rawRows++ + lastID = node.ID + if !scope.Allows(node) { + continue + } + nodes = append(nodes, node) + if len(nodes) == sentinel { + break + } + } + if err := rows.Err(); err != nil { + _ = rows.Close() + return graph.BoundedNodeProjection{}, err + } + if err := rows.Close(); err != nil { + return graph.BoundedNodeProjection{}, err + } + if len(nodes) == sentinel || rawRows < rawPageSize { + break + } + } + if err := tx.Commit(); err != nil { + return graph.BoundedNodeProjection{}, err + } + + total := len(nodes) + truncated := total > limit + if len(nodes) > limit { + nodes = nodes[:limit] + } + // Match the in-memory projection's ownership boundary: callers cannot + // reslice a bounded page into the sentinel slot. + nodes = nodes[:len(nodes):len(nodes)] + return graph.BoundedNodeProjection{Nodes: nodes, Total: total, Truncated: truncated}, nil +} + +func localizationNodePredicate(name string, scope graph.LocalizationNodeScope) (string, []any) { + clauses, args := localizationScopePredicate(scope) + clauses = append([]string{`name = ?`}, clauses...) + args = append([]any{name}, args...) + return strings.Join(clauses, ` AND `), args +} + +func localizationFileNodePredicate(filePath string, scope graph.LocalizationNodeScope) (string, []any) { + clauses, args := localizationScopePredicate(scope) + clauses = append([]string{`file_path = ?`}, clauses...) + args = append([]any{filePath}, args...) + return strings.Join(clauses, ` AND `), args +} + +func localizationScopePredicate(scope graph.LocalizationNodeScope) ([]string, []any) { + var clauses []string + var args []any + if len(scope.Kinds) > 0 { + kinds := make([]string, 0, len(scope.Kinds)) + for kind, allowed := range scope.Kinds { + if allowed { + kinds = append(kinds, string(kind)) + } + } + sort.Strings(kinds) + if len(kinds) == 0 { + clauses = append(clauses, `0`) + } else { + clauses = append(clauses, `kind IN (`+inPlaceholders(len(kinds))+`)`) + for _, kind := range kinds { + args = append(args, kind) + } + } + } + if len(scope.ExcludeKinds) > 0 { + kinds := make([]string, 0, len(scope.ExcludeKinds)) + for kind, excluded := range scope.ExcludeKinds { + if excluded { + kinds = append(kinds, string(kind)) + } + } + sort.Strings(kinds) + if len(kinds) > 0 { + clauses = append(clauses, `kind NOT IN (`+inPlaceholders(len(kinds))+`)`) + for _, kind := range kinds { + args = append(args, kind) + } + } + } + if scope.WorkspaceID != "" { + clauses = append(clauses, `COALESCE(NULLIF(workspace_id, ''), repo_prefix) = ?`) + args = append(args, scope.WorkspaceID) + if scope.ProjectID != "" { + clauses = append(clauses, `COALESCE(NULLIF(project_id, ''), repo_prefix) = ?`) + args = append(args, scope.ProjectID) + } + } + if len(scope.RepoAllow) > 0 { + repos := make([]string, 0, len(scope.RepoAllow)) + for repo, allowed := range scope.RepoAllow { + if allowed { + repos = append(repos, repo) + } + } + sort.Strings(repos) + if len(repos) == 0 { + clauses = append(clauses, `repo_prefix = ''`) + } else { + clauses = append(clauses, `(repo_prefix = '' OR repo_prefix IN (`+inPlaceholders(len(repos))+`))`) + for _, repo := range repos { + args = append(args, repo) + } + } + } + return clauses, args +} + +func scanLocalizationNode(scanner interface{ Scan(...any) error }) (*graph.Node, error) { + var ( + node graph.Node + metaBlob sql.RawBytes + ) + if err := scanner.Scan( + &node.ID, &node.Kind, &node.Name, &node.QualName, &node.FilePath, + &node.StartLine, &node.EndLine, &node.StartColumn, &node.EndColumn, + &node.Language, &node.RepoPrefix, &node.WorkspaceID, &node.ProjectID, + &metaBlob, + ); err != nil { + return nil, err + } + if len(metaBlob) == 0 { + return &node, nil + } + meta, err := decodeMeta([]byte(metaBlob)) + if err != nil { + return nil, err + } + if isTest, _ := meta["is_test"].(bool); isTest { + node.Meta = map[string]any{"is_test": true} + } + return &node, nil +} diff --git a/internal/graph/store_sqlite/localization_projection_test.go b/internal/graph/store_sqlite/localization_projection_test.go new file mode 100644 index 00000000..13ef4b24 --- /dev/null +++ b/internal/graph/store_sqlite/localization_projection_test.go @@ -0,0 +1,473 @@ +package store_sqlite + +import ( + "context" + "fmt" + "path/filepath" + "strings" + "testing" + + "github.com/zzet/gortex/internal/graph" +) + +func TestFindNodesByNameBoundedCapsTenThousandHomonyms(t *testing.T) { + store, err := Open(filepath.Join(t.TempDir(), "graph.sqlite")) + if err != nil { + t.Fatalf("open store: %v", err) + } + t.Cleanup(func() { _ = store.Close() }) + + const homonyms = 10_000 + nodes := make([]*graph.Node, 0, homonyms+2) + for index := 0; index < homonyms; index++ { + nodes = append(nodes, &graph.Node{ + ID: fmt.Sprintf("repo/src/handler-%05d.go::handle", index), + Name: "handle", + Kind: graph.KindFunction, + FilePath: fmt.Sprintf("repo/src/handler-%05d.go", index), + RepoPrefix: "repo", + WorkspaceID: "workspace", + ProjectID: "project", + Meta: map[string]any{ + "doc": "payload that the localization projection must not hydrate", + "search_doc": "large search-only payload", + "custom_metadata": stringsRepeatForProjectionTest("payload", 16), + }, + }) + } + // These rows prove every scope predicate is applied before COUNT and LIMIT. + nodes = append(nodes, + &graph.Node{ID: "foreign/src/f.go::handle", Name: "handle", Kind: graph.KindFunction, FilePath: "foreign/src/f.go", RepoPrefix: "foreign", WorkspaceID: "foreign"}, + &graph.Node{ID: "repo/src/value.go::handle", Name: "handle", Kind: graph.KindVariable, FilePath: "repo/src/value.go", RepoPrefix: "repo", WorkspaceID: "workspace", ProjectID: "project"}, + ) + store.BeginBulkLoad() + store.AddBatch(nodes, nil) + if err := store.FlushBulk(); err != nil { + t.Fatalf("flush: %v", err) + } + + page, err := store.FindNodesByNameBounded( + context.Background(), + "handle", + graph.LocalizationNodeScope{ + WorkspaceID: "workspace", + ProjectID: "project", + RepoAllow: map[string]bool{"repo": true}, + Kinds: map[graph.NodeKind]bool{graph.KindFunction: true}, + }, + 8, + ) + if err != nil { + t.Fatalf("bounded lookup: %v", err) + } + if page.Total != 9 { + t.Fatalf("total = %d, want threshold total 9", page.Total) + } + if !page.Truncated { + t.Fatal("truncated = false, want LIMIT+1 sentinel saturation") + } + if len(page.Nodes) != 8 { + t.Fatalf("nodes = %d, want hard cap 8", len(page.Nodes)) + } + for index, node := range page.Nodes { + want := fmt.Sprintf("repo/src/handler-%05d.go::handle", index) + if node.ID != want { + t.Fatalf("node[%d] = %q, want deterministic %q", index, node.ID, want) + } + if _, hydrated := node.Meta["doc"]; hydrated { + t.Fatalf("node[%d] hydrated promoted doc payload", index) + } + if _, hydrated := node.Meta["custom_metadata"]; hydrated { + t.Fatalf("node[%d] retained unrelated metadata", index) + } + } +} + +func TestFindNodesByNameBoundedFindsProductionRowsBehindTests(t *testing.T) { + store, err := Open(filepath.Join(t.TempDir(), "graph.sqlite")) + if err != nil { + t.Fatalf("open store: %v", err) + } + t.Cleanup(func() { _ = store.Close() }) + + nodes := make([]*graph.Node, 0, 42) + for index := 0; index < 32; index++ { + nodes = append(nodes, &graph.Node{ + ID: fmt.Sprintf("a-tests/%02d.go::handle", index), Name: "handle", + Kind: graph.KindFunction, FilePath: fmt.Sprintf("a-tests/%02d.go", index), + Meta: map[string]any{"is_test": true}, + }) + } + for index := 0; index < 10; index++ { + nodes = append(nodes, &graph.Node{ + ID: fmt.Sprintf("z-prod/%02d.go::handle", index), Name: "handle", + Kind: graph.KindFunction, FilePath: fmt.Sprintf("z-prod/%02d.go", index), + }) + } + store.BeginBulkLoad() + store.AddBatch(nodes, nil) + if err := store.FlushBulk(); err != nil { + t.Fatalf("flush: %v", err) + } + + page, err := store.FindNodesByNameBounded( + context.Background(), "handle", graph.LocalizationNodeScope{ExcludeTests: true}, 8, + ) + if err != nil { + t.Fatalf("bounded lookup: %v", err) + } + if len(page.Nodes) != 8 || page.Total != 9 || !page.Truncated { + t.Fatalf("page = %#v, want eight production nodes and saturation sentinel", page) + } + for _, node := range page.Nodes { + if len(node.FilePath) < len("z-prod/") || node.FilePath[:len("z-prod/")] != "z-prod/" { + t.Fatalf("test declaration leaked into production page: %#v", node) + } + } +} + +func TestFindNodesByNameBoundedPreservesTestClassification(t *testing.T) { + store, err := Open(filepath.Join(t.TempDir(), "graph.sqlite")) + if err != nil { + t.Fatalf("open store: %v", err) + } + t.Cleanup(func() { _ = store.Close() }) + store.AddNode(&graph.Node{ + ID: "repo/handler_test.go::handle", Name: "handle", Kind: graph.KindFunction, + FilePath: "repo/handler_test.go", Meta: map[string]any{"is_test": true, "doc": "omit me"}, + }) + + page, err := store.FindNodesByNameBounded(context.Background(), "handle", graph.LocalizationNodeScope{}, 1) + if err != nil { + t.Fatalf("bounded lookup: %v", err) + } + if len(page.Nodes) != 1 { + t.Fatalf("nodes = %d, want 1", len(page.Nodes)) + } + if isTest, _ := page.Nodes[0].Meta["is_test"].(bool); !isTest { + t.Fatalf("projected meta = %#v, want is_test classification", page.Nodes[0].Meta) + } + if _, hydrated := page.Nodes[0].Meta["doc"]; hydrated { + t.Fatal("bounded projection hydrated doc") + } +} + +func TestFindNodesByNameBoundedHonorsCancellation(t *testing.T) { + store, err := Open(filepath.Join(t.TempDir(), "graph.sqlite")) + if err != nil { + t.Fatalf("open store: %v", err) + } + t.Cleanup(func() { _ = store.Close() }) + store.AddNode(&graph.Node{ID: "repo/f.go::handle", Name: "handle", Kind: graph.KindFunction, FilePath: "repo/f.go"}) + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + page, err := store.FindNodesByNameBounded(ctx, "handle", graph.LocalizationNodeScope{}, 8) + if err == nil { + t.Fatalf("error = nil with cancelled context; page = %#v", page) + } + if len(page.Nodes) != 0 || page.Total != 0 { + t.Fatalf("cancelled lookup returned partial page %#v", page) + } +} + +func TestOverlaidViewFindNodesByNameBoundedAppliesSQLiteScopeWithoutMutatingCaller(t *testing.T) { + store, err := Open(filepath.Join(t.TempDir(), "graph.sqlite")) + if err != nil { + t.Fatalf("open store: %v", err) + } + t.Cleanup(func() { _ = store.Close() }) + + nodes := make([]*graph.Node, 0, 302) + for index := 0; index < 300; index++ { + nodes = append(nodes, &graph.Node{ + ID: fmt.Sprintf("repo/a-overlay-hidden.go::handle:%03d", index), Name: "handle", + Kind: graph.KindFunction, FilePath: "repo/a-overlay-hidden.go", + }) + } + nodes = append(nodes, + &graph.Node{ + ID: "repo/b-caller-hidden.go::handle", Name: "handle", Kind: graph.KindFunction, + FilePath: "repo/b-caller-hidden.go", + }, + &graph.Node{ + ID: "repo/z-visible.go::handle", Name: "handle", Kind: graph.KindFunction, + FilePath: "repo/z-visible.go", + }, + ) + store.BeginBulkLoad() + store.AddBatch(nodes, nil) + if err := store.FlushBulk(); err != nil { + t.Fatalf("flush: %v", err) + } + + layer := graph.NewOverlayLayer() + layer.MarkFile("repo/a-overlay-hidden.go", true) + callerFiles := map[string]bool{"repo/b-caller-hidden.go": true} + page, err := graph.NewOverlaidView(store, layer).FindNodesByNameBounded( + context.Background(), "handle", + graph.LocalizationNodeScope{ExcludeFiles: callerFiles}, 8, + ) + if err != nil { + t.Fatalf("bounded SQLite overlay lookup: %v", err) + } + if page.Total != 1 || page.Truncated || len(page.Nodes) != 1 || + page.Nodes[0].FilePath != "repo/z-visible.go" { + t.Fatalf("page = %#v, want only the visible row behind excluded keyset pages", page) + } + if len(callerFiles) != 1 || !callerFiles["repo/b-caller-hidden.go"] { + t.Fatalf("caller exclusion map was mutated: %#v", callerFiles) + } + if _, added := callerFiles["repo/a-overlay-hidden.go"]; added { + t.Fatalf("overlay path leaked into caller exclusion map: %#v", callerFiles) + } +} + +func TestFindFileNodesBoundedCapsScopesKindsAndDropsPayloads(t *testing.T) { + store, err := Open(filepath.Join(t.TempDir(), "graph.sqlite")) + if err != nil { + t.Fatalf("open store: %v", err) + } + t.Cleanup(func() { _ = store.Close() }) + + const filePath = "shared/dense.go" + nodes := make([]*graph.Node, 0, 192) + for index := 0; index < 64; index++ { + nodes = append(nodes, + &graph.Node{ + ID: fmt.Sprintf("repo/dense.go::fn-%03d", index), Name: fmt.Sprintf("fn%d", index), + Kind: graph.KindFunction, FilePath: filePath, RepoPrefix: "repo", + WorkspaceID: "workspace", ProjectID: "project", + Meta: map[string]any{ + "signature": "func with promoted payload", + "doc": stringsRepeatForProjectionTest("documentation", 32), + "custom_metadata": stringsRepeatForProjectionTest("metadata", 32), + }, + }, + &graph.Node{ + ID: fmt.Sprintf("repo/dense.go::value-%03d", index), Name: fmt.Sprintf("value%d", index), + Kind: graph.KindVariable, FilePath: filePath, RepoPrefix: "repo", + WorkspaceID: "workspace", ProjectID: "project", + }, + &graph.Node{ + ID: fmt.Sprintf("foreign/dense.go::fn-%03d", index), Name: fmt.Sprintf("foreign%d", index), + Kind: graph.KindFunction, FilePath: filePath, RepoPrefix: "foreign", + WorkspaceID: "foreign", ProjectID: "foreign", + }, + ) + } + store.BeginBulkLoad() + store.AddBatch(nodes, nil) + if err := store.FlushBulk(); err != nil { + t.Fatalf("flush: %v", err) + } + + page, err := store.FindFileNodesBounded( + context.Background(), filePath, + graph.LocalizationNodeScope{ + WorkspaceID: "workspace", ProjectID: "project", + RepoAllow: map[string]bool{"repo": true}, + Kinds: map[graph.NodeKind]bool{graph.KindFunction: true}, + }, + 8, + ) + if err != nil { + t.Fatalf("bounded file lookup: %v", err) + } + if page.Total != 9 || !page.Truncated || len(page.Nodes) != 8 { + t.Fatalf("page = %#v, want threshold total 9, truncated, cap 8", page) + } + if cap(page.Nodes) != len(page.Nodes) { + t.Fatalf("page capacity = %d, want exact returned length %d", cap(page.Nodes), len(page.Nodes)) + } + for index, node := range page.Nodes { + want := fmt.Sprintf("repo/dense.go::fn-%03d", index) + if node.ID != want { + t.Fatalf("node[%d] = %q, want deterministic %q", index, node.ID, want) + } + if node.Kind != graph.KindFunction || node.RepoPrefix != "repo" || node.WorkspaceID != "workspace" { + t.Fatalf("scope or kind was applied after cap: %#v", node) + } + if node.Meta != nil { + t.Fatalf("summary node[%d] hydrated metadata: %#v", index, node.Meta) + } + } + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + if cancelled, err := store.FindFileNodesBounded(ctx, filePath, graph.LocalizationNodeScope{}, 8); err == nil || len(cancelled.Nodes) != 0 { + t.Fatalf("cancelled file lookup = %#v, %v; want empty error result", cancelled, err) + } +} + +func TestFindFileNodesBoundedFindsProductionRowsBehindTests(t *testing.T) { + store, err := Open(filepath.Join(t.TempDir(), "graph.sqlite")) + if err != nil { + t.Fatalf("open store: %v", err) + } + t.Cleanup(func() { _ = store.Close() }) + + const filePath = "repo/dense.go" + nodes := make([]*graph.Node, 0, 310) + for index := 0; index < 300; index++ { + nodes = append(nodes, &graph.Node{ + ID: fmt.Sprintf("repo/dense.go::a-test-%03d", index), Name: "test", + Kind: graph.KindFunction, FilePath: filePath, Meta: map[string]any{"is_test": true}, + }) + } + for index := 0; index < 10; index++ { + nodes = append(nodes, &graph.Node{ + ID: fmt.Sprintf("repo/dense.go::z-prod-%03d", index), Name: "prod", + Kind: graph.KindFunction, FilePath: filePath, + }) + } + store.BeginBulkLoad() + store.AddBatch(nodes, nil) + if err := store.FlushBulk(); err != nil { + t.Fatalf("flush: %v", err) + } + + page, err := store.FindFileNodesBounded( + context.Background(), filePath, graph.LocalizationNodeScope{ExcludeTests: true}, 8, + ) + if err != nil { + t.Fatalf("bounded file lookup: %v", err) + } + if page.Total != 9 || !page.Truncated || len(page.Nodes) != 8 { + t.Fatalf("page = %#v, want production sentinel after test rows", page) + } + for _, node := range page.Nodes { + if node.Name != "prod" { + t.Fatalf("test declaration consumed production cap: %#v", node) + } + } +} + +func TestFindFileNodesBoundedFindsDefinitionsBehindExcludedKinds(t *testing.T) { + store, err := Open(filepath.Join(t.TempDir(), "graph.sqlite")) + if err != nil { + t.Fatalf("open store: %v", err) + } + t.Cleanup(func() { _ = store.Close() }) + + const filePath = "repo/generated.go" + nodes := make([]*graph.Node, 0, 310) + for index := 0; index < 300; index++ { + nodes = append(nodes, &graph.Node{ + ID: fmt.Sprintf("repo/generated.go::a-param-%03d", index), Name: "arg", + Kind: graph.KindParam, FilePath: filePath, + }) + } + for index := 0; index < 10; index++ { + nodes = append(nodes, &graph.Node{ + ID: fmt.Sprintf("repo/generated.go::z-function-%03d", index), Name: "function", + Kind: graph.KindFunction, FilePath: filePath, + }) + } + store.BeginBulkLoad() + store.AddBatch(nodes, nil) + if err := store.FlushBulk(); err != nil { + t.Fatalf("flush: %v", err) + } + + page, err := store.FindFileNodesBounded( + context.Background(), filePath, + graph.LocalizationNodeScope{ExcludeKinds: map[graph.NodeKind]bool{graph.KindParam: true}}, 8, + ) + if err != nil { + t.Fatalf("bounded file lookup: %v", err) + } + if page.Total != 9 || !page.Truncated || len(page.Nodes) != 8 { + t.Fatalf("page = %#v, want definition sentinel behind excluded params", page) + } + for _, node := range page.Nodes { + if node.Kind != graph.KindFunction { + t.Fatalf("excluded kind consumed the cap: %#v", node) + } + } +} + +func TestFindFileNodesBoundedDoesNotTransferMetadataWhenTestsIncluded(t *testing.T) { + store, err := Open(filepath.Join(t.TempDir(), "graph.sqlite")) + if err != nil { + t.Fatalf("open store: %v", err) + } + t.Cleanup(func() { _ = store.Close() }) + + const nodeID = "repo/file.go::handler" + store.AddNode(&graph.Node{ + ID: nodeID, Name: "handler", Kind: graph.KindFunction, FilePath: "repo/file.go", + Meta: map[string]any{"signature": "func handler()", "doc": "promoted payload"}, + }) + if _, err := store.writerDB.Exec(`UPDATE nodes SET meta = ? WHERE id = ?`, []byte("not-valid-metadata"), nodeID); err != nil { + t.Fatalf("corrupt metadata fixture: %v", err) + } + + page, err := store.FindFileNodesBounded( + context.Background(), "repo/file.go", graph.LocalizationNodeScope{}, 8, + ) + if err != nil { + t.Fatalf("summary lookup decoded metadata it must not select: %v", err) + } + if len(page.Nodes) != 1 || page.Nodes[0].ID != nodeID { + t.Fatalf("page = %#v, want one summary node", page) + } + if page.Nodes[0].Meta != nil { + t.Fatalf("summary hydrated metadata: %#v", page.Nodes[0].Meta) + } +} + +func TestFindFileNodesBoundedPlanUsesFileIndexWithoutSorter(t *testing.T) { + store, err := Open(filepath.Join(t.TempDir(), "graph.sqlite")) + if err != nil { + t.Fatalf("open store: %v", err) + } + t.Cleanup(func() { _ = store.Close() }) + store.AddNode(&graph.Node{ + ID: "repo/file.go::handler", Name: "handler", Kind: graph.KindFunction, + FilePath: "repo/file.go", RepoPrefix: "repo", WorkspaceID: "workspace", ProjectID: "project", + }) + + predicate, args := localizationFileNodePredicate("repo/file.go", graph.LocalizationNodeScope{ + WorkspaceID: "workspace", ProjectID: "project", + RepoAllow: map[string]bool{"repo": true}, + Kinds: map[graph.NodeKind]bool{graph.KindFunction: true, graph.KindMethod: true}, + ExcludeKinds: map[graph.NodeKind]bool{graph.KindParam: true, graph.KindLocal: true}, + }) + args = append(args, "", 257) + rows, err := store.db.Query( + `EXPLAIN QUERY PLAN SELECT `+lookupNodeSummaryCols+` FROM nodes WHERE `+predicate+` AND id > ? ORDER BY id LIMIT ?`, + args..., + ) + if err != nil { + t.Fatalf("explain bounded file query: %v", err) + } + defer rows.Close() + var details []string + for rows.Next() { + var id, parent, unused int + var detail string + if err := rows.Scan(&id, &parent, &unused, &detail); err != nil { + t.Fatalf("scan query plan: %v", err) + } + details = append(details, detail) + } + if err := rows.Err(); err != nil { + t.Fatalf("query plan rows: %v", err) + } + plan := strings.Join(details, "\n") + if !strings.Contains(plan, "nodes_by_file") { + t.Fatalf("query plan does not use nodes_by_file:\n%s", plan) + } + if strings.Contains(strings.ToUpper(plan), "TEMP B-TREE") { + t.Fatalf("query plan sorts bounded file rows:\n%s", plan) + } +} + +func stringsRepeatForProjectionTest(value string, count int) string { + out := "" + for index := 0; index < count; index++ { + out += value + } + return out +} diff --git a/internal/graph/store_sqlite/planner_stats.go b/internal/graph/store_sqlite/planner_stats.go index 4bfbe399..3fab5a2b 100644 --- a/internal/graph/store_sqlite/planner_stats.go +++ b/internal/graph/store_sqlite/planner_stats.go @@ -31,6 +31,7 @@ const plannerStatsAnalysisLimit = 1000 const plannerStatsIndexQuery = ` WITH critical(name) AS (VALUES ('edges_by_from_line'), + ('edges_by_from_line_kind'), ('edges_by_kind'), ('nodes_by_file'), ('nodes_by_kind'), @@ -110,14 +111,21 @@ func quoteSQLiteIdentifier(name string) string { func healPlannerStats(db *sql.DB) { var hasTable bool // sqlite_stat1 does not exist until the first ANALYZE, so probe the catalog - // before the table. A sidecar-only stat row is not enough: nodes_by_file is - // the sentinel for the graph-index refresh that protects the hottest plan. + // before the table. Warm stores with edges must have both the original hot + // file stat and the bounded exact-site stat; Open may have just created the + // latter index on a pre-upgrade database. Empty edge tables emit no sibling + // stat and need no edge-plan repair. if err := db.QueryRow(`SELECT EXISTS(SELECT 1 FROM sqlite_master WHERE type = 'table' AND name = 'sqlite_stat1')`).Scan(&hasTable); err != nil { return } if hasTable { var populated bool - if err := db.QueryRow(`SELECT EXISTS(SELECT 1 FROM sqlite_stat1 WHERE idx = 'nodes_by_file')`).Scan(&populated); err == nil && populated { + if err := db.QueryRow(` +SELECT EXISTS(SELECT 1 FROM sqlite_stat1 WHERE idx = 'nodes_by_file') + AND ( + NOT EXISTS(SELECT 1 FROM edges LIMIT 1) + OR EXISTS(SELECT 1 FROM sqlite_stat1 WHERE idx = 'edges_by_from_line_kind') + )`).Scan(&populated); err == nil && populated { return } } diff --git a/internal/graph/store_sqlite/planner_stats_checkpoint_test.go b/internal/graph/store_sqlite/planner_stats_checkpoint_test.go index 6af02736..ad219fb8 100644 --- a/internal/graph/store_sqlite/planner_stats_checkpoint_test.go +++ b/internal/graph/store_sqlite/planner_stats_checkpoint_test.go @@ -73,6 +73,7 @@ FROM seq`) } wantIndexes := []string{ "edges_by_from_line", + "edges_by_from_line_kind", "edges_by_kind", "nodes_by_file", "nodes_by_kind", @@ -99,8 +100,8 @@ FROM seq`) plan = explainPlannerQueryPlan(t, store.db, `SELECT to_id FROM edges WHERE from_id = ? AND line = ? AND kind = ?`, "hub", 1, "calls") - if !strings.Contains(plan, "edges_by_from_line") { - t.Fatalf("exact-site query missed edges_by_from_line after stats refresh:\n%s", plan) + if !strings.Contains(plan, "edges_by_from_line_kind") { + t.Fatalf("exact-site query missed edges_by_from_line_kind after stats refresh:\n%s", plan) } // These indexes have no competing left-prefix path. They remain selected diff --git a/internal/graph/store_sqlite/planner_stats_test.go b/internal/graph/store_sqlite/planner_stats_test.go index 537c13df..1cdefbab 100644 --- a/internal/graph/store_sqlite/planner_stats_test.go +++ b/internal/graph/store_sqlite/planner_stats_test.go @@ -1,6 +1,7 @@ package store_sqlite import ( + "context" "fmt" "path/filepath" "testing" @@ -113,3 +114,53 @@ func TestOpenHealsPlannerStats(t *testing.T) { t.Fatal("open did not heal sqlite_stat1 for a populated store") } } + +func TestOpenHealsMissingBoundedSitePlannerStat(t *testing.T) { + path := filepath.Join(t.TempDir(), "stats_site_heal.sqlite") + s, err := Open(path) + if err != nil { + t.Fatalf("open: %v", err) + } + seedPlannerStatsNodes(s) + s.AddEdge(&graph.Edge{ + From: "pkg/a.go::sym00", + To: "pkg/a.go::sym01", + Kind: graph.EdgeCalls, + FilePath: "pkg/a.go", + Line: 7, + }) + s.writeMu.Lock() + err = s.refreshPlannerStatsLocked(context.Background()) + s.writeMu.Unlock() + if err != nil { + t.Fatalf("seed graph planner stats: %v", err) + } + if _, err := s.writerDB.Exec(`DROP INDEX edges_by_from_line_kind`); err != nil { + t.Fatalf("drop bounded-site index: %v", err) + } + var warmState bool + if err := s.db.QueryRow(` +SELECT EXISTS(SELECT 1 FROM sqlite_stat1 WHERE idx = 'nodes_by_file') + AND NOT EXISTS(SELECT 1 FROM sqlite_stat1 WHERE idx = 'edges_by_from_line_kind')`).Scan(&warmState); err != nil { + t.Fatalf("probe pre-upgrade planner stats: %v", err) + } + if !warmState { + t.Fatal("fixture did not retain old graph stats while removing the new sibling stat") + } + if err := s.Close(); err != nil { + t.Fatalf("close pre-upgrade store: %v", err) + } + + reopened, err := Open(path) + if err != nil { + t.Fatalf("reopen: %v", err) + } + defer reopened.Close() + var healed bool + if err := reopened.db.QueryRow(`SELECT EXISTS(SELECT 1 FROM sqlite_stat1 WHERE idx = 'edges_by_from_line_kind')`).Scan(&healed); err != nil { + t.Fatalf("probe healed bounded-site stat: %v", err) + } + if !healed { + t.Fatal("open recreated the bounded-site index without healing its planner stat") + } +} diff --git a/internal/hooks/codex.go b/internal/hooks/codex.go index b23d3e7a..d031b026 100644 --- a/internal/hooks/codex.go +++ b/internal/hooks/codex.go @@ -66,6 +66,9 @@ func runCodex(data []byte, port int, selected ...CodexMode) { var peek struct { HookEventName string `json:"hook_event_name"` ToolName string `json:"tool_name"` + SessionID string `json:"session_id"` + PromptID string `json:"prompt_id"` + AgentID string `json:"agent_id"` CWD string `json:"cwd"` // Codex sends the active model slug on every hook event. Claude Code // does not, which is why its hint has to be recovered from the @@ -83,7 +86,28 @@ func runCodex(data []byte, port int, selected ...CodexMode) { setHookCWD(peek.CWD) defer setHookCWD("") + // Codex has specialized Bash and Gortex-read handlers, so enforce the + // shared terminal contract before dispatch. Decode only the identity and + // tool name here: Codex permits tool_input to be any JSON value, and a + // scalar or array must not bypass an enforceable terminal marker merely + // because later tool-specific handlers expect an arguments object. + if peek.HookEventName == "PreToolUse" && enforceLocalizationTerminalPreToolUse(HookInput{ + HookEventName: peek.HookEventName, + ToolName: peek.ToolName, + SessionID: peek.SessionID, + PromptID: peek.PromptID, + AgentID: peek.AgentID, + CWD: peek.CWD, + }, time.Now()) { + return + } + switch { + case peek.HookEventName == "Stop": + // Codex uses the same stop_hook_active and last_assistant_message + // fields as the shared Stop handler. Unsupported/older hosts omit the + // message and therefore fail open in runPostTask. + runPostTask(data, port) case peek.HookEventName == "SessionStart": runSessionStart(data, port) case peek.HookEventName == "PreToolUse" && peek.ToolName == "Bash": @@ -95,9 +119,9 @@ func runCodex(data []byte, port int, selected ...CodexMode) { default: runPreToolUse(data, port, ModeEnrich) } - case peek.HookEventName == "PreToolUse" && codexMCPReadPreToolUseTool(peek.ToolName): - runCodexMCPReadPreToolUse(data, mode) - case peek.HookEventName == "PostToolUse" && (peek.ToolName == "Bash" || peek.ToolName == "apply_patch"): + case peek.HookEventName == "PreToolUse" && codexLocalizationPreToolUseTool(peek.ToolName): + runCodexLocalizationPreToolUse(data, mode) + case peek.HookEventName == "PostToolUse" && (peek.ToolName == "Bash" || peek.ToolName == "apply_patch" || localizationNavigationTool(peek.ToolName)): runCodexPostToolUse(data, port, mode) case peek.HookEventName == "UserPromptSubmit": // Re-surface graph symbols relevant to the prompt on every turn. @@ -289,47 +313,69 @@ func codexMCPReadPreToolUseTool(toolName string) bool { } } -func runCodexMCPReadPreToolUse(data []byte, mode CodexMode) { +func codexLocalizationPreToolUseTool(toolName string) bool { + return codexMCPReadPreToolUseTool(toolName) || localizationNavigationTool(toolName) +} + +func runCodexLocalizationPreToolUse(data []byte, mode CodexMode) { started := time.Now() var input HookInput if err := json.Unmarshal(data, &input); err != nil { return } - if input.HookEventName != "PreToolUse" || !codexMCPReadPreToolUseTool(input.ToolName) { + if input.HookEventName != "PreToolUse" || !codexLocalizationPreToolUseTool(input.ToolName) { return } + + updatedInput := map[string]any(nil) + if turn, ok := currentLocalizationTurnState(input.SessionID, input.PromptID, input.AgentID, input.CWD); ok { + if authToken, ready := snapshotLocalizationToolUseWithAuth(input, turn.Identity); ready { + updatedInput = localizationPreToolUpdatedInput(input, authToken, turn.ProblemStatement) + } + } + daemonUp := daemonReachableFn() emitted := false defer func() { logHookEffectiveness("PreToolUse", emitted, daemonUp, 0, time.Since(started)) }() - // Daemon outage: nudging, rewriting, or (in deny mode) blocking a Gortex - // read is moot when the daemon cannot serve it — the call is about to - // fail on transport, and a compress-bodies deny on top of that failure - // would read as enforcement of a tool that does not answer (#486). - if !daemonUp { - return + ctx := "" + if daemonUp && codexMCPReadPreToolUseTool(input.ToolName) { + ctx = gortexReadNudge(input.ToolName, input.ToolInput) } - - ctx := gortexReadNudge(input.ToolName, input.ToolInput) - if ctx == "" { + if ctx == "" && updatedInput == nil { return } - hso := &HookSpecificOutput{HookEventName: "PreToolUse", AdditionalContext: ctx} - switch mode { - case CodexModeDeny: - hso.AdditionalContext = "" - hso.PermissionDecision = "deny" - hso.PermissionDecisionReason = ctx - case CodexModeRewrite: - hso.PermissionDecision = "allow" - hso.UpdatedInput = rewrittenGortexReadInput(input.ToolName, input.ToolInput) + hso := &HookSpecificOutput{ + HookEventName: "PreToolUse", + AdditionalContext: ctx, + UpdatedInput: updatedInput, + } + if ctx != "" { + switch mode { + case CodexModeDeny: + hso.AdditionalContext = "" + hso.PermissionDecision = "deny" + hso.PermissionDecisionReason = ctx + case CodexModeRewrite: + hso.PermissionDecision = "allow" + rewritten := rewrittenGortexReadInput(input.ToolName, input.ToolInput) + for key, value := range updatedInput { + rewritten[key] = value + } + hso.UpdatedInput = rewritten + } } emitted = true emitPreToolUse(HookOutput{HookSpecificOutput: hso}) } +// Kept for existing callers that exercise the read-only Codex path directly. +func runCodexMCPReadPreToolUse(data []byte, mode CodexMode) { + runCodexLocalizationPreToolUse(data, mode) +} + func rewrittenGortexReadInput(toolName string, input map[string]any) map[string]any { out := cloneStringAnyMap(input) switch strings.TrimSpace(toolName) { @@ -429,6 +475,12 @@ func runCodexPostToolUse(data []byte, port int, mode CodexMode) { if input.HookEventName != "PostToolUse" { return } + if localizationNavigationTool(input.ToolName) { + if terminal, observed := observeLocalizationTerminal(data); observed { + _ = emitLocalizationTerminalContext(terminal) + } + return + } if input.ToolName == "apply_patch" { daemonUp := daemonReachableFn() ctx := "" diff --git a/internal/hooks/codex_localization_alias_test.go b/internal/hooks/codex_localization_alias_test.go new file mode 100644 index 00000000..6277d98c --- /dev/null +++ b/internal/hooks/codex_localization_alias_test.go @@ -0,0 +1,17 @@ +package hooks + +import "testing" + +func TestCodexLocalizationLifecycleRecognizesBothMCPNamespaces(t *testing.T) { + for _, prefix := range []string{gortexMCPToolPrefix, gortexCodexMCPToolPrefix} { + for _, operation := range []string{"explore", "search", "read", "relations", "trace", "analyze"} { + tool := prefix + operation + if !localizationNavigationTool(tool) { + t.Errorf("localization navigation does not recognize %q", tool) + } + if !codexLocalizationPreToolUseTool(tool) { + t.Errorf("Codex PreToolUse does not recognize %q", tool) + } + } + } +} diff --git a/internal/hooks/codex_localization_lifecycle_test.go b/internal/hooks/codex_localization_lifecycle_test.go new file mode 100644 index 00000000..5671a4cb --- /dev/null +++ b/internal/hooks/codex_localization_lifecycle_test.go @@ -0,0 +1,94 @@ +package hooks + +import ( + "encoding/json" + "strings" + "testing" + "time" + + "github.com/zzet/gortex/internal/localizationauth" +) + +func TestCodexLocalizationLifecycleArmsClaimCheckFromReceipt(t *testing.T) { + configureLocalizationTerminalTestHome(t) + oldProbe := userPromptProbe + userPromptProbe = func(string, time.Duration) ([]grepSymbolHit, error) { return nil, nil } + t.Cleanup(func() { userPromptProbe = oldProbe }) + + sessionID := "codex-lifecycle" + promptID := "prompt-1" + cwd := t.TempDir() + prompt := codexLocalizationPayload(t, map[string]any{ + "hook_event_name": "UserPromptSubmit", "session_id": sessionID, + "prompt_id": promptID, "cwd": cwd, "prompt": "Locate the writer implementation", + }) + _ = captureStdout(t, func() { runCodex(prompt, 0) }) + + pre := codexLocalizationPayload(t, map[string]any{ + "hook_event_name": "PreToolUse", "tool_name": gortexMCPToolPrefix + "explore", + "tool_use_id": "call-1", "session_id": sessionID, "prompt_id": promptID, "cwd": cwd, + "tool_input": map[string]any{"operation": "localize", "task": "Locate the writer implementation"}, + }) + preOut := captureStdout(t, func() { runCodex(pre, 0) }) + var hookOutput HookOutput + if err := json.Unmarshal([]byte(preOut), &hookOutput); err != nil { + t.Fatalf("decode Codex PreToolUse output %q: %v", preOut, err) + } + if hookOutput.HookSpecificOutput == nil || hookOutput.HookSpecificOutput.UpdatedInput == nil { + t.Fatalf("Codex explore did not receive authenticated updated input: %#v", hookOutput) + } + token, _ := hookOutput.HookSpecificOutput.UpdatedInput[localizationauth.ArgumentKey].(string) + if token == "" { + t.Fatalf("Codex explore updated input omitted terminal auth: %#v", hookOutput.HookSpecificOutput.UpdatedInput) + } + identity, ok := currentLocalizationTurn(sessionID, promptID, "", cwd) + if !ok { + t.Fatal("Codex UserPromptSubmit did not begin a localization turn") + } + if hasLocalizationTerminal(identity) { + t.Fatal("PreToolUse armed terminal state before an authenticated MCP receipt") + } + + primary := []string{"repo/a.go::Writer.write"} + if !localizationauth.Publish(token, localizationauth.Receipt{ + FinalResponse: "Use Writer.write.", PrimaryIDs: primary, EvidenceIDs: primary, + ContractVersion: localizationTerminalContractV2, Enforceable: true, + }) { + t.Fatal("MCP answer_ready receipt was not published") + } + post := codexLocalizationPayload(t, map[string]any{ + "hook_event_name": "PostToolUse", "tool_name": gortexMCPToolPrefix + "explore", + "tool_use_id": "call-1", "session_id": sessionID, "prompt_id": promptID, "cwd": cwd, + "tool_input": hookOutput.HookSpecificOutput.UpdatedInput, + "tool_response": map[string]any{"content": []any{}}, + }) + postOut := captureStdout(t, func() { runCodex(post, 0) }) + if !strings.Contains(postOut, "Localization for this task is complete") { + t.Fatalf("Codex PostToolUse did not consume the receipt: %q", postOut) + } + if !hasLocalizationTerminal(identity) { + t.Fatal("Codex PostToolUse did not arm terminal state") + } + + stop := codexLocalizationPayload(t, map[string]any{ + "hook_event_name": "Stop", "session_id": sessionID, "prompt_id": promptID, + "cwd": cwd, "last_assistant_message": "SYMBOLS:\n- flush", + }) + stopOut := captureStdout(t, func() { runCodex(stop, 0) }) + var decision HookOutput + if err := json.Unmarshal([]byte(stopOut), &decision); err != nil { + t.Fatalf("decode Codex Stop output %q: %v", stopOut, err) + } + if decision.Decision != "block" || !strings.Contains(decision.Reason, "claim_check") { + t.Fatalf("wrong final claim was not blocked from authenticated lifecycle: %#v", decision) + } +} + +func codexLocalizationPayload(t *testing.T, value map[string]any) []byte { + t.Helper() + data, err := json.Marshal(value) + if err != nil { + t.Fatal(err) + } + return data +} diff --git a/internal/hooks/codex_stop_test.go b/internal/hooks/codex_stop_test.go new file mode 100644 index 00000000..d608212c --- /dev/null +++ b/internal/hooks/codex_stop_test.go @@ -0,0 +1,56 @@ +package hooks + +import ( + "encoding/json" + "strings" + "testing" +) + +func TestRunCodexStopClaimCheckUsesFinalMessageAndStopsAfterRetry(t *testing.T) { + input := claimCheckTestInput(t, "SYMBOLS:\n- flush") + data, err := json.Marshal(input) + if err != nil { + t.Fatal(err) + } + out := captureStdout(t, func() { runCodex(data, 0) }) + var payload HookOutput + if err := json.Unmarshal([]byte(out), &payload); err != nil { + t.Fatalf("decode claim-check output %q: %v", out, err) + } + if payload.Decision != "block" || !strings.Contains(payload.Reason, "claim_check") { + t.Fatalf("Codex Stop was not blocked with a claim check: %#v", payload) + } + + input.StopHookActive = true + data, _ = json.Marshal(input) + if retry := captureStdout(t, func() { runCodex(data, 0) }); retry != "" { + t.Fatalf("recursive Codex Stop should be silent: %q", retry) + } +} + +func TestRunCodexStopClaimCheckReadsPrimaryHeading(t *testing.T) { + input := claimCheckTestInput(t, "# PRIMARY\nFabricated") + data, err := json.Marshal(input) + if err != nil { + t.Fatal(err) + } + out := captureStdout(t, func() { runCodex(data, 0) }) + var payload HookOutput + if err := json.Unmarshal([]byte(out), &payload); err != nil { + t.Fatalf("decode heading claim-check output %q: %v", out, err) + } + if payload.Decision != "block" || !strings.Contains(payload.Reason, "claim_check") { + t.Fatalf("Codex Stop did not enforce PRIMARY heading claims: %#v", payload) + } +} + +func TestRunCodexStopFailsOpenWithoutFinalMessage(t *testing.T) { + input := claimCheckTestInput(t, "") + data, err := json.Marshal(input) + if err != nil { + t.Fatal(err) + } + if out := captureStdout(t, func() { runCodex(data, 0) }); out != "" { + t.Fatalf("Codex Stop without last_assistant_message must fail open: %q", out) + } +} diff --git a/internal/hooks/codex_test.go b/internal/hooks/codex_test.go index c2ab22f2..2e5eeb67 100644 --- a/internal/hooks/codex_test.go +++ b/internal/hooks/codex_test.go @@ -68,6 +68,110 @@ func TestRunCodexIgnoresNonBash(t *testing.T) { } } +func TestRunCodexPreToolUseTerminalGateCoversHookObservableLocalTools(t *testing.T) { + configureLocalizationTerminalTestHome(t) + identity := beginTestLocalizationTurn(t, "codex-terminal-families", "prompt-1", t.TempDir()) + const finalResponse = "Use the retained localization evidence." + if !markLocalizationTerminalWithStrength(identity, localizationTerminalContractV2, false, finalResponse) { + t.Fatal("mark enforceable terminal localization") + } + + tests := []struct { + name string + tool string + mode CodexMode + }{ + {name: "apply patch mutation", tool: "apply_patch"}, + {name: "local plan update", tool: "update_plan"}, + {name: "subagent", tool: "spawn_agent"}, + {name: "image read", tool: "view_image"}, + {name: "bash deny posture", tool: "Bash", mode: CodexModeDeny}, + {name: "bash rewrite posture", tool: "Bash", mode: CodexModeRewrite}, + {name: "gortex navigation", tool: gortexMCPToolPrefix + "search"}, + {name: "gortex change contract", tool: gortexMCPToolPrefix + "change"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + data := preToolPayload(t, tt.tool, "tool-1", identity, map[string]any{}) + out := captureHookStdout(t, func() { runCodex(data, 0, tt.mode) }) + hso := decodeHookOutput(t, out).HookSpecificOutput + if hso == nil || hso.PermissionDecision != "deny" { + t.Fatalf("terminal %s output=%q want deny", tt.tool, out) + } + if !strings.HasPrefix(hso.PermissionDecisionReason, localizationTerminalDenyReason) { + t.Fatalf("terminal reason=%q want prefix %q", hso.PermissionDecisionReason, localizationTerminalDenyReason) + } + if !strings.Contains(hso.PermissionDecisionReason, finalResponse) { + t.Fatalf("terminal reason omitted retained response: %q", hso.PermissionDecisionReason) + } + }) + } +} + +func TestRunCodexPreToolUseTerminalGateAcceptsAnyToolInputShape(t *testing.T) { + configureLocalizationTerminalTestHome(t) + identity := beginTestLocalizationTurn(t, "codex-terminal-input-shapes", "prompt-1", t.TempDir()) + if !markLocalizationTerminalWithStrength(identity, localizationTerminalContractV2, false, "Use retained evidence.") { + t.Fatal("mark enforceable terminal localization") + } + + for _, tt := range []struct { + name string + input any + }{ + {name: "null", input: nil}, + {name: "scalar", input: "freeform input"}, + {name: "array", input: []any{"one", 2}}, + } { + t.Run(tt.name, func(t *testing.T) { + data := mustJSON(t, map[string]any{ + "hook_event_name": "PreToolUse", + "tool_name": "custom_local_tool", + "tool_input": tt.input, + "session_id": identity.SessionID, + "prompt_id": identity.PromptID, + "agent_id": identity.AgentID, + "cwd": identity.CWD, + }) + out := captureHookStdout(t, func() { runCodex(data, 0) }) + hso := decodeHookOutput(t, out).HookSpecificOutput + if hso == nil || hso.PermissionDecision != "deny" { + t.Fatalf("terminal input %T output=%q want deny", tt.input, out) + } + }) + } +} + +func TestRunCodexPreToolUseAdvisoryTerminalPreservesContractOperations(t *testing.T) { + configureLocalizationTerminalTestHome(t) + identity := beginTestLocalizationTurn(t, "codex-advisory-contract", "prompt-1", t.TempDir()) + if !markLocalizationTerminalWithStrength(identity, localizationTerminalContractV2, true, "Advisory localization.") { + t.Fatal("mark advisory terminal localization") + } + + for _, tool := range []string{gortexMCPToolPrefix + "change", gortexCodexMCPToolPrefix + "change"} { + t.Run(tool, func(t *testing.T) { + data := preToolPayload(t, tool, "tool-1", identity, map[string]any{"operation": "impact"}) + if out := captureHookStdout(t, func() { runCodex(data, 0) }); out != "" { + t.Fatalf("advisory terminal blocked non-navigation contract operation %s: %q", tool, out) + } + }) + } +} + +func TestRunCodexPreToolUseWithoutTerminalPreservesBuiltins(t *testing.T) { + configureLocalizationTerminalTestHome(t) + identity := beginTestLocalizationTurn(t, "codex-no-terminal", "prompt-1", t.TempDir()) + for _, tool := range []string{"apply_patch", "update_plan", "view_image", gortexMCPToolPrefix + "change"} { + t.Run(tool, func(t *testing.T) { + data := preToolPayload(t, tool, "tool-1", identity, map[string]any{}) + if out := captureHookStdout(t, func() { runCodex(data, 0) }); out != "" { + t.Fatalf("non-terminal operation %s emitted %q", tool, out) + } + }) + } +} + func TestRunCodexPreToolUseBashSoftAdditionalContext(t *testing.T) { oldProbe := grepProbe grepProbe = func(string, time.Duration) ([]grepSymbolHit, error) { diff --git a/internal/hooks/gortex_read.go b/internal/hooks/gortex_read.go index 01dfd823..0b69e930 100644 --- a/internal/hooks/gortex_read.go +++ b/internal/hooks/gortex_read.go @@ -133,7 +133,8 @@ func gortexReadAdvisory(toolName, path string) string { // mcp__plugin_gortex_gortex__. The advisory should display only the operation. func shortGortexToolName(toolName string) string { toolName = strings.TrimPrefix(toolName, gortexMCPToolPrefix) - return strings.TrimPrefix(toolName, gortexPluginMCPToolPrefix) + toolName = strings.TrimPrefix(toolName, gortexPluginMCPToolPrefix) + return strings.TrimPrefix(toolName, gortexCodexMCPToolPrefix) } // hasReadSizeCap reports whether the read already bounds its output via a diff --git a/internal/hooks/localization_advised_call_test.go b/internal/hooks/localization_advised_call_test.go index b8dbe20d..41bec97d 100644 --- a/internal/hooks/localization_advised_call_test.go +++ b/internal/hooks/localization_advised_call_test.go @@ -4,6 +4,8 @@ import ( "encoding/json" "strings" "testing" + + "github.com/zzet/gortex/internal/localizationauth" ) // A refusal must never be the answer to advice this same hook just gave. When @@ -17,7 +19,7 @@ func TestAdvisoryMarkerAnswersTheHostToolsItWouldOtherwiseRedirect(t *testing.T) t.Run(tool, func(t *testing.T) { configureLocalizationTerminalTestHome(t) identity := beginTestLocalizationTurn(t, t.Name(), "prompt", t.TempDir()) - if !markLocalizationTerminalReceipt(identity, localizationTerminalContractV2, false, answer) { + if !markLocalizationTerminalReceipt(identity, localizationauth.Receipt{FinalResponse: answer, ContractVersion: localizationTerminalContractV2}) { t.Fatal("advisory marker was not written") } input := map[string]any{"file_path": "storage/disk.go", "pattern": "Load"} @@ -56,7 +58,7 @@ func TestAdvisoryMarkerAnswersTheHostToolsItWouldOtherwiseRedirect(t *testing.T) func TestAdvisoryMarkerStillPassesThroughUnrelatedTools(t *testing.T) { configureLocalizationTerminalTestHome(t) identity := beginTestLocalizationTurn(t, t.Name(), "prompt", t.TempDir()) - if !markLocalizationTerminalReceipt(identity, localizationTerminalContractV2, false, "answer") { + if !markLocalizationTerminalReceipt(identity, localizationauth.Receipt{FinalResponse: "answer", ContractVersion: localizationTerminalContractV2}) { t.Fatal("advisory marker was not written") } for _, tool := range []string{"WebSearch", "Write"} { diff --git a/internal/hooks/localization_claim_check.go b/internal/hooks/localization_claim_check.go new file mode 100644 index 00000000..a8700ea7 --- /dev/null +++ b/internal/hooks/localization_claim_check.go @@ -0,0 +1,1307 @@ +package hooks + +import ( + "fmt" + "strings" + "unicode" + "unicode/utf8" +) + +const ( + localizationClaimCheckMaxChars = 600 + localizationClaimCheckMaxMessageBytes = 64 << 10 + localizationClaimCheckMaxClaims = 32 + localizationClaimCheckMaxTokens = 2048 + localizationClaimCheckMaxTokenBytes = 256 + localizationRejectedClaim = "__gortex_invalid_claim_input__" +) + +// localizationClaimCheck returns a single bounded correction only when a Stop +// payload exposes the final assistant message and that message makes an +// explicit code-shaped claim outside the authenticated evidence digest. +func localizationClaimCheck(input PostTaskInput) string { + if input.LastAssistantMessage == "" || input.StopHookActive { + return "" + } + claims := []string(nil) + if len(input.LastAssistantMessage) > localizationClaimCheckMaxMessageBytes { + claims = []string{localizationRejectedClaim} + } else { + message := strings.TrimSpace(input.LastAssistantMessage) + if message == "" { + return "" + } + claims = localizationExplicitSymbolClaims(message) + } + if len(claims) == 0 { + return "" + } + identity, ok := currentLocalizationTurn(input.SessionID, input.PromptID, input.AgentID, input.CWD) + if !ok { + return "" + } + marker, consumed := consumeLocalizationClaimCheck(identity, claims) + if !consumed { + return "" + } + prompt := "[Gortex claim_check] Your answer names one or more symbols outside the authenticated evidence. Cite only these PRIMARY IDs, or explicitly confirm that none fits: " + strings.Join(marker.PrimaryIDs, ", ") + ". Do not retrieve more evidence." + return boundedLocalizationClaimCheck(prompt) +} + +// localizationExplicitSymbolClaims preserves the slice-only helper contract +// while making every resource-limit violation an unmatchable claim. +func localizationExplicitSymbolClaims(message string) []string { + claims, _, valid := localizationBoundedSymbolClaims(message) + if !valid { + return []string{localizationRejectedClaim} + } + return claims +} + +func localizationBoundedSymbolClaims(message string) ([]string, bool, bool) { + if len(message) > localizationClaimCheckMaxMessageBytes { + return nil, false, false + } + budget := newLocalizationClaimBudget() + explicitNone := false + inSymbols := false + symbolsSawContent := false + lines := strings.Split(message, "\n") + var fence localizationMarkdownFenceState + var symbolsContainer localizationMarkdownContainerIdentity + + for index, line := range lines { + markdown := localizationParseMarkdownContainer(line) + if !markdown.valid { + return nil, false, false + } + if fence.open { + contained, sameContainer := localizationParseMarkdownContinuation(line, fence.container) + if sameContainer { + if marker, ok := localizationMarkdownFenceMarker(contained.content); ok && + !contained.codeIndented && marker.character == fence.character && + marker.length >= fence.length && marker.closing { + fence = localizationMarkdownFenceState{} + continue + } + // A heading-looking line inside a code fence is code, not a heading. + // Scan it so a fabricated qualified identity cannot hide behind '#'. + if !localizationScanUnstructuredClaimBody(contained.content, budget) { + return nil, false, false + } + continue + } + // A list or quote container ended before its fence closed. CommonMark + // ends the fenced block with that container; process this same line + // again as ordinary answer content instead of dropping it. + fence = localizationMarkdownFenceState{} + } + + if inSymbols && symbolsContainer.depth > 0 { + scoped, sameContainer := localizationParseMarkdownContinuationWithChildren(line, symbolsContainer) + if sameContainer { + if !scoped.valid { + return nil, false, false + } + markdown = scoped + } else { + inSymbols = false + symbolsSawContent = false + symbolsContainer = localizationMarkdownContainerIdentity{} + } + } + if marker, ok := localizationMarkdownFenceMarker(markdown.content); ok && !markdown.codeIndented { + marker.container = markdown.identity() + fence = marker + continue + } + + trimmed := strings.TrimSpace(markdown.content) + headingBody, heading := localizationMarkdownHeadingBodyAt(lines, index, markdown) + _, atxHeading := localizationMarkdownATXHeadingBody(markdown.content) + if heading { + if role, body, roleLine := localizationClaimRoleLine(headingBody, true); roleLine { + inSymbols = role == localizationClaimRoleOpen + symbolsSawContent = false + symbolsContainer = localizationMarkdownContainerIdentity{} + if inSymbols { + symbolsContainer = markdown.identity() + } + if body == "" { + continue + } + if inSymbols { + material, none, valid := localizationAddStructuredClaimLine(body, budget) + if !valid { + return nil, false, false + } + symbolsSawContent = material + explicitNone = explicitNone || none + } else if !localizationScanHeadingClaimBody(body, budget) { + return nil, false, false + } + continue + } + if atxHeading { + inSymbols = false + symbolsSawContent = false + symbolsContainer = localizationMarkdownContainerIdentity{} + if !localizationScanHeadingClaimBody(headingBody, budget) { + return nil, false, false + } + continue + } + } + if localizationMarkdownSetextUnderlineContent(markdown.content) { + continue + } + if role, body, roleLine := localizationClaimRoleLine(trimmed, false); roleLine { + inSymbols = role == localizationClaimRoleOpen + symbolsSawContent = false + symbolsContainer = localizationMarkdownContainerIdentity{} + if inSymbols { + symbolsContainer = markdown.identity() + } + if body == "" { + continue + } + if inSymbols { + material, none, valid := localizationAddStructuredClaimLine(body, budget) + if !valid { + return nil, false, false + } + symbolsSawContent = material + explicitNone = explicitNone || none + } else if !localizationScanUnstructuredClaimBody(body, budget) { + return nil, false, false + } + continue + } + if inSymbols { + if trimmed == "" { + if symbolsSawContent { + inSymbols = false + symbolsContainer = localizationMarkdownContainerIdentity{} + } + continue + } + material, none, valid := localizationAddStructuredClaimLine(trimmed, budget) + if !valid { + return nil, false, false + } + if material { + symbolsSawContent = true + explicitNone = explicitNone || none + continue + } + inSymbols = false + symbolsContainer = localizationMarkdownContainerIdentity{} + } + if heading { + if !localizationScanHeadingClaimBody(headingBody, budget) { + return nil, false, false + } + continue + } + body, inspect := localizationUnstructuredClaimLine(markdown.content) + if inspect && !localizationScanUnstructuredClaimBody(body, budget) { + return nil, false, false + } + } + return budget.claims, explicitNone, true +} + +func localizationAddStructuredClaimLine(line string, budget *localizationClaimBudget) (bool, bool, bool) { + token, rest, none, material := localizationStructuredClaimLine(line) + if !material { + return false, false, true + } + if token != "" && !budget.countToken(token) { + return false, false, false + } + if !none && token != "" && !budget.addClaim(localizationStructuredSymbolClaim(token), false) { + return false, false, false + } + if rest != "" && !localizationScanUnstructuredClaimBody(rest, budget) { + return false, false, false + } + return true, none, true +} + +func localizationClaimRoleLine(line string, heading bool) (localizationClaimRoleMode, string, bool) { + line = strings.TrimSpace(line) + if colon := strings.IndexByte(line, ':'); colon >= 0 && + (colon+1 >= len(line) || line[colon+1] != ':') { + if role := localizationClaimRole(line[:colon]); role != localizationClaimRoleNone { + return role, strings.TrimSpace(line[colon+1:]), true + } + } + if heading { + if role := localizationClaimRole(strings.TrimSuffix(line, ":")); role != localizationClaimRoleNone { + return role, "", true + } + } + return localizationClaimRoleNone, "", false +} + +type localizationClaimBudget struct { + claims []string + seen map[string]struct{} + tokens int +} + +func newLocalizationClaimBudget() *localizationClaimBudget { + return &localizationClaimBudget{ + claims: make([]string, 0, 4), + seen: make(map[string]struct{}, 4), + } +} + +func (budget *localizationClaimBudget) countToken(token string) bool { + budget.tokens++ + return budget.tokens <= localizationClaimCheckMaxTokens && len(token) <= localizationClaimCheckMaxTokenBytes +} + +func (budget *localizationClaimBudget) addClaim(claim string, requireCodeShape bool) bool { + if requireCodeShape && (len(claim) < 2 || !localizationCodeShapedClaim(claim)) { + return true + } + if claim == "" { + return true + } + if _, duplicate := budget.seen[claim]; duplicate { + return true + } + if len(budget.claims) >= localizationClaimCheckMaxClaims { + return false + } + budget.seen[claim] = struct{}{} + budget.claims = append(budget.claims, claim) + return true +} + +func localizationScanUnstructuredClaimBody(body string, budget *localizationClaimBudget) bool { + return localizationScanClaimBody(body, budget, false) +} + +func localizationScanHeadingClaimBody(body string, budget *localizationClaimBudget) bool { + return localizationScanClaimBody(body, budget, true) +} + +func localizationScanClaimBody(body string, budget *localizationClaimBudget, heading bool) bool { + tokenStart := -1 + skipUntil := -1 + consume := func(start, end int) bool { + token := body[start:end] + if !budget.countToken(token) { + return false + } + // '_' and '$' are identifier characters, not wrappers. Trimming them + // silently changed _private/$foo into another identity (or no claim). + claim := strings.Trim(token, ".:#\\/-") + if localizationContextualFileClaim(body, start, end, claim) { + return true + } + explicitSyntax := localizationExplicitInlineClaim(body, start, end, claim) + if heading { + qualified := strings.ContainsAny(claim, ".:#") && localizationCodeShapedClaim(claim) || + localizationBackslashQualifiedClaim(claim) + if !explicitSyntax && !qualified { + return true + } + return budget.addClaim(claim, false) + } + return budget.addClaim(claim, !explicitSyntax) + } + for index, r := range body { + if index < skipUntil { + continue + } + if r == '(' { + receiverStart := index + if tokenStart >= 0 { + prefix := body[tokenStart:index] + if strings.HasSuffix(prefix, ".") || strings.HasSuffix(prefix, "::") { + receiverStart = tokenStart + } + } + if end, claim, ok := localizationGoReceiverClaimAt(body, receiverStart, index); ok { + if tokenStart >= 0 && tokenStart < receiverStart && !consume(tokenStart, receiverStart) { + return false + } + tokenStart = -1 + if !budget.countToken(body[receiverStart:end]) || !budget.addClaim(claim, false) { + return false + } + skipUntil = end + continue + } + } + if localizationClaimTokenRune(r) { + if tokenStart < 0 { + tokenStart = index + } + continue + } + if tokenStart >= 0 { + if !consume(tokenStart, index) { + return false + } + tokenStart = -1 + } + } + return tokenStart < 0 || consume(tokenStart, len(body)) +} + +func localizationGoReceiverClaimAt(body string, start, receiverOpen int) (int, string, bool) { + if start < 0 || start > receiverOpen || receiverOpen >= len(body) || body[receiverOpen] != '(' { + return 0, "", false + } + prefix := "" + fileQualified := false + if start < receiverOpen { + rawPrefix := body[start:receiverOpen] + switch { + case strings.HasSuffix(rawPrefix, "::"): + file := strings.TrimSuffix(rawPrefix, "::") + if !localizationFileQualifiedClaim(file + "::receiver") { + return 0, "", false + } + prefix = file + "::" + fileQualified = true + case strings.HasSuffix(rawPrefix, "."): + qualifier := strings.TrimSuffix(rawPrefix, ".") + if !localizationGoPackageQualifier(qualifier) { + return 0, "", false + } + prefix = qualifier + "." + default: + return 0, "", false + } + } + index := receiverOpen + 1 + if index < len(body) && body[index] == '*' { + index++ + } + receiverStart := index + for index < len(body) { + r, width := utf8.DecodeRuneInString(body[index:]) + if !unicode.IsLetter(r) && !unicode.IsDigit(r) && r != '_' && r != '$' && r != '.' { + break + } + index += width + } + if receiverStart == index || index+2 >= len(body) || body[index] != ')' || body[index+1] != '.' { + return 0, "", false + } + methodStart := index + 2 + index = methodStart + for index < len(body) { + r, width := utf8.DecodeRuneInString(body[index:]) + if !unicode.IsLetter(r) && !unicode.IsDigit(r) && r != '_' && r != '$' { + break + } + index += width + } + if methodStart == index { + return 0, "", false + } + receiver := localizationNormalizeGoPointerIdentity(body[receiverOpen:index]) + if !localizationInlineIdentitySyntax(receiver) { + return 0, "", false + } + claim := prefix + receiver + if !fileQualified && !localizationInlineIdentitySyntax(claim) { + return 0, "", false + } + end := index + if strings.HasPrefix(body[end:], "()") { + end += 2 + } + return end, claim, true +} + +func localizationGoPackageQualifier(value string) bool { + parts := strings.Split(value, ".") + if len(parts) == 0 { + return false + } + for _, part := range parts { + if !localizationExplicitIdentityRow(part, part) { + return false + } + } + return true +} + +// localizationExplicitInlineClaim admits a simple lower-case leaf only when +// the answer marks it as code or a call. Ordinary prose words remain outside +// claim checking; qualified and file-qualified identities continue through +// localizationCodeShapedClaim instead. +func localizationContextualFileClaim(body string, start, _ int, claim string) bool { + if !localizationAmbiguousFileExtension(claim) { + return false + } + const contextBytes = 64 + beforeStart := start - contextBytes + if beforeStart < 0 { + beforeStart = 0 + } + before := body[beforeStart:start] + if boundary := strings.LastIndexAny(before, ".;!?\n\r"); boundary >= 0 { + before = before[boundary+1:] + } + return localizationFileContextBefore(localizationContextWords(before)) +} + +func localizationContextWords(value string) []string { + return strings.FieldsFunc(strings.ToLower(value), func(r rune) bool { + return !unicode.IsLetter(r) && !unicode.IsDigit(r) && r != '_' + }) +} + +func localizationFileContextBefore(words []string) bool { + if len(words) == 0 { + return false + } + if localizationFileContextWord(words[len(words)-1]) { + return true + } + if len(words) < 2 { + return false + } + switch words[len(words)-1] { + case "is", "at", "named", "called": + return localizationFileContextWord(words[len(words)-2]) + default: + return false + } +} + +func localizationFileContextWord(word string) bool { + switch word { + case "file", "path", "source", "header", "document", "manifest": + return true + default: + return false + } +} + +func localizationExplicitInlineClaim(body string, start, end int, claim string) bool { + if !localizationInlineIdentitySyntax(claim) { + return false + } + callEnd := end + callShaped := strings.HasPrefix(body[end:], "()") + if callShaped { + callEnd += 2 + } + return callShaped || localizationInlineCodeDelimited(body, start, callEnd) +} + +func localizationInlineIdentitySyntax(claim string) bool { + if localizationExplicitIdentityRow(claim, claim) { + return true + } + if claim == "" || strings.ContainsAny(claim, "/\\") { + return false + } + parts := strings.FieldsFunc(claim, func(r rune) bool { + return strings.ContainsRune(".:#", r) + }) + if len(parts) < 2 { + return false + } + for _, part := range parts { + if part == "" || !localizationExplicitIdentityRow(part, part) { + return false + } + } + return true +} + +func localizationBackslashQualifiedClaim(claim string) bool { + parts := strings.Split(claim, "\\") + if len(parts) < 2 { + return false + } + for _, part := range parts { + if !localizationExplicitIdentityRow(part, part) { + return false + } + } + return true +} + +func localizationInlineCodeDelimited(body string, start, end int) bool { + opening := 0 + for index := start - 1; index >= 0 && body[index] == '`'; index-- { + opening++ + } + if opening < 1 || opening > 3 { + return false + } + closing := 0 + for index := end; index < len(body) && body[index] == '`'; index++ { + closing++ + } + return closing == opening +} + +func localizationUnstructuredClaimLine(line string) (string, bool) { + line = strings.TrimSpace(line) + if line == "" { + return "", false + } + if _, heading := localizationMarkdownATXHeadingBody(line); heading { + return "", false + } + if localizationClaimRole(line) != localizationClaimRoleNone { + return "", false + } + if colon := strings.IndexByte(line, ':'); colon >= 0 && + (colon+1 >= len(line) || line[colon+1] != ':') && localizationClaimRoleLabel(line[:colon]) { + line = strings.TrimSpace(line[colon+1:]) + return line, line != "" + } + return line, true +} + +const localizationMarkdownContainerMaxDepth = 8 + +type localizationMarkdownContainerKind uint8 + +const ( + localizationMarkdownQuoteContainer localizationMarkdownContainerKind = iota + 1 + localizationMarkdownListContainer +) + +type localizationMarkdownContainerFrame struct { + kind localizationMarkdownContainerKind + continuationColumns uint8 +} + +type localizationMarkdownContainerIdentity struct { + depth uint8 + frames [localizationMarkdownContainerMaxDepth]localizationMarkdownContainerFrame +} + +type localizationMarkdownFenceState struct { + open bool + character byte + length int + closing bool + container localizationMarkdownContainerIdentity +} + +func localizationMarkdownFenceMarker(line string) (localizationMarkdownFenceState, bool) { + if len(line) < 3 || (line[0] != '`' && line[0] != '~') { + return localizationMarkdownFenceState{}, false + } + character := line[0] + length := 0 + for length < len(line) && line[length] == character { + length++ + } + if length < 3 { + return localizationMarkdownFenceState{}, false + } + return localizationMarkdownFenceState{ + open: true, + character: character, + length: length, + closing: strings.TrimSpace(line[length:]) == "", + }, true +} + +type localizationMarkdownContainer struct { + content string + path localizationMarkdownContainerIdentity + codeIndented bool + valid bool +} + +func (container localizationMarkdownContainer) identity() localizationMarkdownContainerIdentity { + return container.path +} + +func (container *localizationMarkdownContainer) push(frame localizationMarkdownContainerFrame) bool { + if int(container.path.depth) >= len(container.path.frames) { + return false + } + container.path.frames[container.path.depth] = frame + container.path.depth++ + return true +} + +func localizationParseMarkdownContainer(line string) localizationMarkdownContainer { + container := localizationMarkdownContainer{ + content: strings.TrimRight(line, " \t\r"), + valid: true, + } + return localizationParseMarkdownChildContainers(container, 0) +} + +func localizationParseMarkdownChildContainers( + container localizationMarkdownContainer, + column int, +) localizationMarkdownContainer { + for container.content != "" { + content, indent, codeIndented := localizationMarkdownIndentWidthAt(container.content, column) + column += indent + container.content = strings.TrimRight(content, " \t\r") + if codeIndented { + container.codeIndented = true + break + } + line := container.content + if line == "" { + break + } + if line[0] == '>' && (len(line) == 1 || line[1] == ' ' || line[1] == '\t') { + if !container.push(localizationMarkdownContainerFrame{kind: localizationMarkdownQuoteContainer}) { + container.valid = false + return container + } + var padding int + container.content, padding = localizationMarkdownMarkerRemainderAt(line[1:], column+1) + column += 1 + padding + continue + } + frame, remainder, markerColumns, listItem := localizationMarkdownListMarker(line, column) + if !listItem { + break + } + frame.continuationColumns += uint8(indent) + if !container.push(frame) { + container.valid = false + return container + } + column += markerColumns + container.content = strings.TrimRight(remainder, " \t\r") + } + return container +} + +func localizationParseMarkdownContinuation( + line string, + identity localizationMarkdownContainerIdentity, +) (localizationMarkdownContainer, bool) { + return localizationParseMarkdownContinuationMode(line, identity, false) +} + +func localizationParseMarkdownContinuationWithChildren( + line string, + identity localizationMarkdownContainerIdentity, +) (localizationMarkdownContainer, bool) { + return localizationParseMarkdownContinuationMode(line, identity, true) +} + +func localizationParseMarkdownContinuationMode( + line string, + identity localizationMarkdownContainerIdentity, + parseChildren bool, +) (localizationMarkdownContainer, bool) { + container := localizationMarkdownContainer{ + content: strings.TrimRight(line, " \t\r"), + path: identity, + valid: true, + } + column := 0 + for index := 0; index < int(identity.depth); index++ { + frame := identity.frames[index] + if strings.TrimSpace(container.content) == "" { + for remaining := index; remaining < int(identity.depth); remaining++ { + if identity.frames[remaining].kind == localizationMarkdownQuoteContainer { + return localizationMarkdownContainer{}, false + } + } + container.content = "" + return container, true + } + switch frame.kind { + case localizationMarkdownQuoteContainer: + content, indent, codeIndented := localizationMarkdownIndentWidthAt(container.content, column) + column += indent + if codeIndented || content == "" || content[0] != '>' || + len(content) > 1 && content[1] != ' ' && content[1] != '\t' { + return localizationMarkdownContainer{}, false + } + var padding int + container.content, padding = localizationMarkdownMarkerRemainderAt(content[1:], column+1) + column += 1 + padding + case localizationMarkdownListContainer: + required := int(frame.continuationColumns) + content, indent, ok := localizationMarkdownConsumeIndentAt( + container.content, required, column, + ) + if !ok { + return localizationMarkdownContainer{}, false + } + if indent > required { + content = strings.Repeat(" ", indent-required) + content + indent = required + } + column += indent + container.content = strings.TrimRight(content, " \t\r") + default: + return localizationMarkdownContainer{}, false + } + } + if parseChildren { + return localizationParseMarkdownChildContainers(container, column), true + } + content, _, codeIndented := localizationMarkdownIndentWidthAt(container.content, column) + container.content = strings.TrimRight(content, " \t\r") + container.codeIndented = codeIndented + return container, true +} + +func localizationMarkdownListMarker( + line string, + column int, +) (localizationMarkdownContainerFrame, string, int, bool) { + markerEnd := 0 + if len(line) >= 2 && strings.ContainsRune("-*+", rune(line[0])) { + markerEnd = 1 + } else { + for markerEnd < len(line) && markerEnd < 9 && line[markerEnd] >= '0' && line[markerEnd] <= '9' { + markerEnd++ + } + if markerEnd == 0 || markerEnd >= len(line) || + (line[markerEnd] != '.' && line[markerEnd] != ')') { + return localizationMarkdownContainerFrame{}, "", 0, false + } + markerEnd++ + } + if markerEnd >= len(line) || line[markerEnd] != ' ' && line[markerEnd] != '\t' { + return localizationMarkdownContainerFrame{}, "", 0, false + } + + paddingEnd := markerEnd + paddingColumns := 0 + for paddingEnd < len(line) && (line[paddingEnd] == ' ' || line[paddingEnd] == '\t') { + if line[paddingEnd] == ' ' { + paddingColumns++ + } else { + paddingColumns += 4 - (column+markerEnd+paddingColumns)%4 + } + paddingEnd++ + } + if paddingColumns > 4 { + paddingEnd = markerEnd + 1 + paddingColumns = 1 + if line[markerEnd] == '\t' { + paddingColumns = 4 - (column+markerEnd)%4 + } + } + markerColumns := markerEnd + paddingColumns + return localizationMarkdownContainerFrame{ + kind: localizationMarkdownListContainer, + continuationColumns: uint8(markerColumns), + }, line[paddingEnd:], markerColumns, true +} + +func localizationMarkdownIndentWidthAt(line string, column int) (string, int, bool) { + columns := 0 + for index := 0; index < len(line); index++ { + switch line[index] { + case ' ': + columns++ + case '\t': + columns += 4 - (column+columns)%4 + default: + return line[index:], columns, false + } + if columns >= 4 { + return line[index+1:], columns, true + } + } + return "", columns, columns >= 4 +} + +func localizationMarkdownConsumeIndentAt(line string, required, column int) (string, int, bool) { + columns := 0 + index := 0 + for index < len(line) && columns < required { + switch line[index] { + case ' ': + columns++ + case '\t': + columns += 4 - (column+columns)%4 + default: + return "", 0, false + } + index++ + } + if columns < required { + return "", 0, false + } + return line[index:], columns, true +} + +func localizationMarkdownMarkerRemainderAt(line string, column int) (string, int) { + padding := 0 + if line != "" { + switch line[0] { + case ' ': + padding = 1 + line = line[1:] + case '\t': + // CommonMark permits one optional padding column after a quote + // marker. Preserve the tab's remaining virtual columns so they + // still contribute to code indentation after the container prefix. + width := 4 - column%4 + line = line[1:] + if width > 1 { + line = strings.Repeat(" ", width-1) + line + } + padding = 1 + } + } + return strings.TrimRight(line, " \t\r"), padding +} + +func localizationMarkdownATXHeadingBody(line string) (string, bool) { + count := 0 + for count < len(line) && line[count] == '#' { + count++ + } + if count == 0 || count > 6 || count < len(line) && line[count] != ' ' && line[count] != '\t' { + return "", false + } + body := strings.TrimSpace(line[count:]) + closing := len(body) + for closing > 0 && body[closing-1] == '#' { + closing-- + } + if closing < len(body) && (closing == 0 || body[closing-1] == ' ' || body[closing-1] == '\t') { + body = strings.TrimSpace(body[:closing]) + } + return body, true +} + +func localizationMarkdownSetextUnderlineContent(content string) bool { + if content == "" || (content[0] != '=' && content[0] != '-') { + return false + } + for index := 1; index < len(content); index++ { + if content[index] != content[0] { + return false + } + } + return true +} + +func localizationMarkdownHeadingBodyAt( + lines []string, + index int, + current localizationMarkdownContainer, +) (string, bool) { + if index < 0 || index >= len(lines) || current.codeIndented { + return "", false + } + if body, ok := localizationMarkdownATXHeadingBody(current.content); ok { + return body, true + } + if current.content == "" || index+1 >= len(lines) || + localizationMarkdownSetextUnderlineContent(current.content) { + return "", false + } + next, sameContainer := localizationParseMarkdownContinuation(lines[index+1], current.identity()) + if !sameContainer || next.codeIndented || !localizationMarkdownSetextUnderlineContent(next.content) { + return "", false + } + return strings.TrimSpace(current.content), true +} + +type localizationClaimRoleMode uint8 + +const ( + localizationClaimRoleNone localizationClaimRoleMode = iota + localizationClaimRoleOpen + localizationClaimRoleClose +) + +func localizationClaimRole(value string) localizationClaimRoleMode { + value = strings.ToLower(strings.TrimSpace(value)) + value = strings.ReplaceAll(value, "-", "_") + switch value { + case "primary", "supporting", "symbol", "symbols": + return localizationClaimRoleOpen + case "evidence", "file", "files", "implementation", "implementation_details", + "details", "answer", "result", "results": + return localizationClaimRoleClose + default: + return localizationClaimRoleNone + } +} + +func localizationClaimRoleLabel(value string) bool { + return localizationClaimRole(value) != localizationClaimRoleNone +} + +func localizationClaimTokenRune(r rune) bool { + return unicode.IsLetter(r) || unicode.IsDigit(r) || strings.ContainsRune("_.$:#\\/-", r) +} + +func localizationStructuredClaimLine(line string) (token, rest string, explicitNone, material bool) { + line, _ = localizationTrimClaimListMarker(line) + if line == "" { + return "", "", false, false + } + if localizationExplicitNoneClaim(line) { + return localizationFirstClaimToken(line), "", true, true + } + token = localizationFirstClaimToken(line) + claim := localizationStructuredSymbolClaim(token) + rest = strings.TrimSpace(strings.TrimPrefix(line, token)) + if rest == "" && localizationExplicitIdentityRow(token, claim) { + return token, "", false, true + } + if localizationCodeShapedClaim(claim) { + return token, rest, false, true + } + // A prose row inside SYMBOLS is not permission to accept its first bare + // word as an identity. Scan the complete row for qualified claims instead. + return "", line, false, true +} + +func localizationExplicitIdentityRow(token, claim string) bool { + if token == "" || claim == "" || localizationStructuredSymbolClaim(token) != claim { + return false + } + for index, r := range claim { + if index == 0 { + if !unicode.IsLetter(r) && r != '_' && r != '$' { + return false + } + continue + } + if !unicode.IsLetter(r) && !unicode.IsDigit(r) && r != '_' && r != '$' { + return false + } + } + return true +} + +func localizationTrimClaimListMarker(line string) (string, bool) { + line = strings.TrimSpace(line) + if line == "" { + return "", false + } + for _, marker := range []string{"-", "*", "•"} { + if strings.HasPrefix(line, marker) { + return strings.TrimSpace(strings.TrimPrefix(line, marker)), true + } + } + index := 0 + for index < len(line) && line[index] >= '0' && line[index] <= '9' { + index++ + } + if index > 0 && index+1 < len(line) && line[index] == '.' && (line[index+1] == ' ' || line[index+1] == '\t') { + return strings.TrimSpace(line[index+1:]), true + } + return line, false +} + +func localizationFirstClaimToken(value string) string { + value = strings.TrimSpace(value) + if index := strings.IndexFunc(value, unicode.IsSpace); index >= 0 { + return value[:index] + } + return value +} + +func localizationExplicitNoneClaim(value string) bool { + value = strings.ToLower(strings.Trim(strings.TrimSpace(value), "`._:;,-")) + switch value { + case "none", "none fits", "none fit", "none of these", "none of the above", + "no symbol fits", "no symbols fit", "no listed symbol fits", "no listed symbols fit": + return true + default: + return false + } +} + +func localizationStructuredSymbolClaim(value string) string { + claim := strings.Trim(localizationFirstClaimToken(value), "` .:#\\/-,;") + for strings.HasSuffix(claim, "()") { + claim = strings.TrimSuffix(claim, "()") + } + if strings.HasPrefix(claim, "(*") { + if close := strings.Index(claim, ")."); close > 2 { + claim = claim[2:close] + claim[close+1:] + } + } else if strings.HasPrefix(claim, "(") { + if close := strings.Index(claim, ")."); close > 1 { + claim = claim[1:close] + claim[close+1:] + } + } + // '_' and '$' remain identity bytes at either edge. + return strings.Trim(claim, "` .:#\\/-,;") +} + +// localizationLooksLikeFileToken rejects stable file basenames and extensions +// from prose claim extraction. File-qualified identities such as foo.c::flush +// are deliberately exempt: their suffix is an explicit symbol identity. +func localizationLooksLikeFileToken(value string) bool { + if value == "" || strings.Contains(value, "::") || strings.Contains(value, "#") { + return false + } + if strings.ContainsAny(value, "/\\") { + return true + } + if localizationKnownFileBasename(value) { + return true + } + dot := strings.LastIndexByte(value, '.') + return dot > 0 && dot+1 < len(value) && localizationKnownFileExtension(value[dot+1:]) +} + +func localizationKnownFileBasename(value string) bool { + name := strings.ToLower(strings.Trim(value, "`.,;:")) + for _, stem := range []string{"dockerfile.", "containerfile.", "makefile."} { + if strings.HasPrefix(name, stem) && len(name) > len(stem) { + return true + } + } + switch name { + case "readme", "license", "licence", "copying", "notice", "changelog", "changes", "authors", + "makefile", "dockerfile", "containerfile", "rakefile", "gemfile", "procfile", "cmakelists.txt": + return true + default: + return false + } +} + +func localizationAmbiguousFileExtension(value string) bool { + dot := strings.LastIndexByte(value, '.') + if dot <= 0 || dot+2 != len(value) { + return false + } + switch strings.ToLower(value[dot+1:]) { + case "m", "r", "s", "v", "d": + return true + default: + return false + } +} + +func localizationKnownFileExtension(extension string) bool { + switch strings.ToLower(extension) { + // Native and systems languages. One-letter m/r/s/v/d remain ambiguous + // unless the surrounding answer explicitly identifies a file or path. + case "c", "h", "cc", "cp", "cpp", "cxx", "hh", "hpp", "hxx", "mm", + "go", "rs", "zig", "odin", "hare", "carbon", "asm", + "sol", "move", "cairo", "nr", "noir", "tact", "bal": + return true + // Scripting, web, JVM, .NET, functional, and shell. + case "py", "pyi", "pyx", "rb", "php", "pl", "pm", "raku", "lua", "tcl", + "js", "jsx", "mjs", "cjs", "ts", "tsx", "mts", "cts", "coffee", + "java", "kt", "kts", "scala", "groovy", "clj", "cljs", "cljc", "edn", + "cs", "fs", "fsx", "vb", "swift", "dart", "ex", "exs", "erl", "hrl", + "hs", "lhs", "ml", "mli", "mll", "elm", "gleam", "res", "re", "rei", + "sh", "bash", "zsh", "fish", "ps1", "bat", "cmd", "ahk": + return true + // UI, schemas, queries, config, and manifests. + case "html", "htm", "css", "scss", "sass", "less", "vue", "svelte", "astro", + "sql", "graphql", "gql", "proto", "thrift", "capnp", "prisma", "wit", + "json", "jsonc", "json5", "yaml", "yml", "toml", "xml", "xsd", "xsl", "xslt", + "ini", "cfg", "conf", "properties", "env", "hcl", "tf", "tfvars", "nix", + "mod", "sum", "work", "lock", "ipynb": + return true + // Documents, templates, build files, and data assets. + case "md", "markdown", "mdx", "rst", "txt", "adoc", "asciidoc", "org", "tex", + "razor", "cshtml", "jsp", "ejs", "hbs", "twig", "erb", "liquid", "pug", + "blade", "tmpl", "tpl", "gotmpl", "mustache", "cmake", "gradle", "bazel", "bzl", + "make", "mk", "ninja", "csv", "tsv", "parquet", "avro", "pdf", "doc", "docx", + "ppt", "pptx", "xls", "xlsx": + return true + default: + return false + } +} + +func localizationCodeShapedClaim(claim string) bool { + if localizationLooksLikeFileToken(claim) { + return false + } + if strings.Contains(claim, "/") && !strings.Contains(claim, "::") && !strings.Contains(claim, "#") { + return false + } + if strings.Contains(claim, ":") && !strings.Contains(claim, "::") { + return false + } + + parts := strings.FieldsFunc(claim, func(r rune) bool { + return strings.ContainsRune(".:#\\/", r) + }) + if len(parts) == 0 { + return false + } + for _, part := range parts { + if part == "" { + continue + } + for index, r := range part { + if index == 0 { + if !unicode.IsLetter(r) && r != '_' && r != '$' { + return false + } + continue + } + if !unicode.IsLetter(r) && !unicode.IsDigit(r) && r != '_' && r != '$' { + return false + } + } + } + if strings.ContainsAny(claim, "_#$\\") || strings.Contains(claim, "::") { + return true + } + if strings.Contains(claim, ".") { + // Reject prose abbreviations such as e.g. while retaining ordinary + // qualified identities such as writer.flush. + for _, part := range parts { + if utf8.RuneCountInString(part) > 1 { + return true + } + } + return false + } + for index, r := range claim { + if index > 0 && unicode.IsUpper(r) { + return true + } + } + return false +} + +func localizationClaimMatchesEvidence(claim, evidenceID string) bool { + claim = localizationNormalizeGoPointerIdentity(strings.TrimSpace(claim)) + evidenceID = strings.TrimSpace(evidenceID) + if claim == "" || claim == localizationRejectedClaim || evidenceID == "" { + return false + } + if claim == evidenceID { + return true + } + + claimFile, claimSymbol, fileQualified := localizationFileQualifiedIdentity(claim) + evidenceFile, evidenceSymbol := localizationEvidenceIdentity(evidenceID) + if fileQualified { + // Graph paths are case-sensitive identities even when a host filesystem + // happens to fold case. Never let a differently-cased file authenticate. + return claimFile == evidenceFile && localizationQualifiedSymbolMatches(claimSymbol, evidenceSymbol) + } + return localizationQualifiedSymbolMatches(claim, evidenceSymbol) +} + +func localizationQualifiedSymbolMatches(claim, evidence string) bool { + claim = localizationNormalizeGoPointerIdentity(strings.TrimSpace(claim)) + evidence = localizationNormalizeGoPointerIdentity(strings.TrimSpace(evidence)) + if claim == "" || evidence == "" { + return false + } + if claim == evidence { + return true + } + if !localizationQualifiedSymbolClaim(claim) { + return claim == localizationSymbolLeaf(evidence) + } + if !strings.HasSuffix(evidence, claim) { + return false + } + prefix := strings.TrimSuffix(evidence, claim) + // The claim's internal notation remains literal. Its containing identity + // may use another language-appropriate separator, but must end at a real + // boundary rather than merely sharing a textual suffix. + for _, boundary := range []string{"::", ".", "#", "\\"} { + if strings.HasSuffix(prefix, boundary) { + return true + } + } + return false +} + +func localizationQualifiedSymbolClaim(identity string) bool { + return strings.Contains(identity, "::") || strings.ContainsAny(identity, ".#\\") +} + +func localizationSymbolLeaf(identity string) string { + last := -1 + width := 1 + for _, separator := range []string{"::", ".", "#", "\\"} { + if index := strings.LastIndex(identity, separator); index > last { + last = index + width = len(separator) + } + } + if last >= 0 && last+width < len(identity) { + return identity[last+width:] + } + return identity +} + +func localizationNormalizeGoPointerIdentity(identity string) string { + identity = strings.TrimSpace(identity) + if strings.HasPrefix(identity, "(*") { + if close := strings.Index(identity, ")."); close > 2 { + return identity[2:close] + identity[close+1:] + } + } else if strings.HasPrefix(identity, "(") { + if close := strings.Index(identity, ")."); close > 1 { + return identity[1:close] + identity[close+1:] + } + } + return identity +} + +func localizationFileQualifiedIdentity(identity string) (file, symbol string, ok bool) { + separator := strings.Index(identity, "::") + if separator <= 0 || !localizationFileQualifiedClaim(identity) { + return "", identity, false + } + return identity[:separator], identity[separator+2:], true +} + +func localizationEvidenceIdentity(evidenceID string) (file, symbol string) { + if separator := strings.Index(evidenceID, "::"); separator >= 0 && separator+2 < len(evidenceID) { + return evidenceID[:separator], evidenceID[separator+2:] + } + return "", evidenceID +} + +func localizationFileQualifiedClaim(claim string) bool { + separator := strings.Index(claim, "::") + if separator <= 0 { + return false + } + prefix := claim[:separator] + if strings.Contains(prefix, "/") || (len(prefix) >= 2 && prefix[1] == ':') { + return true + } + return localizationLooksLikeSourcePath(prefix) +} + +func localizationLooksLikeSourcePath(value string) bool { + value = strings.ToLower(value) + for _, extension := range []string{".go", ".py", ".js", ".ts", ".tsx", ".rs", ".java", ".php", ".rb", ".swift", ".dart", ".cs"} { + if strings.HasSuffix(value, extension) { + return true + } + } + return false +} + +func boundedLocalizationClaimCheck(prompt string) string { + if len(prompt) <= localizationClaimCheckMaxChars { + return prompt + } + const suffix = "… Do not retrieve more evidence." + cut := localizationClaimCheckMaxChars - len(suffix) + for cut > 0 && !utf8.RuneStart(prompt[cut]) { + cut-- + } + return fmt.Sprintf("%s%s", prompt[:cut], suffix) +} diff --git a/internal/hooks/localization_claim_check_followup_test.go b/internal/hooks/localization_claim_check_followup_test.go new file mode 100644 index 00000000..1c381bb6 --- /dev/null +++ b/internal/hooks/localization_claim_check_followup_test.go @@ -0,0 +1,144 @@ +package hooks + +import ( + "os" + "path/filepath" + "strings" + "syscall" + "testing" + + "github.com/gofrs/flock" +) + +func TestLocalizationClaimCheckUnionsStructuredAndProseClaims(t *testing.T) { + tests := []struct { + name string + message string + }{ + { + name: "authenticated structured plus fabricated prose", + message: "SYMBOLS:\n- Writer.write\n\nThe fix is Fabricated.flush.", + }, + { + name: "none fits plus fabricated prose", + message: "SYMBOLS:\n- none fits\n\nThe fix is Fabricated.flush.", + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + input := claimCheckTestInput(t, test.message) + if got := localizationClaimCheck(input); got == "" { + t.Fatal("fabricated prose outside SYMBOLS was not challenged") + } + }) + } +} + +func TestLocalizationClaimIdentityPreservesCaseAndSeparators(t *testing.T) { + tests := []struct { + name string + claim string + evidenceID string + }{ + {name: "file path case", claim: "repo/A.go::Writer.write", evidenceID: "repo/a.go::Writer.write"}, + {name: "hash does not match dot", claim: "Writer#write", evidenceID: "repo/x.go::Writer.write"}, + {name: "dot does not match hash", claim: "Writer.write", evidenceID: "repo/x.py::Writer#write"}, + {name: "double colon does not match dot", claim: "Writer::write", evidenceID: "repo/x.go::Writer.write"}, + {name: "dot does not match double colon", claim: "Writer.write", evidenceID: "repo/x.rs::Writer::write"}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + if localizationClaimMatchesEvidence(test.claim, test.evidenceID) { + t.Fatalf("claim %q cross-authenticated as %q", test.claim, test.evidenceID) + } + }) + } + for _, claim := range []string{"write", "Writer.write", "(*Writer).write"} { + if !localizationClaimMatchesEvidence(claim, "repo/a.go::Writer.write") { + t.Errorf("documented claim notation %q did not authenticate", claim) + } + } +} + +func TestLocalizationClaimParserKeepsColonEndedMaterialClaims(t *testing.T) { + for _, message := range []string{ + "- Fabricated.flush:", + "The culprit is Fabricated.flush:", + } { + claims, _, valid := localizationBoundedSymbolClaims(message) + if !valid { + t.Fatalf("message %q was rejected", message) + } + found := false + for _, claim := range claims { + if claim == "Fabricated.flush" { + found = true + break + } + } + if !found { + t.Fatalf("colon-ended material claim was ignored: message=%q claims=%v", message, claims) + } + } +} + +func TestLocalizationClaimLockReleaseRetriesEINTR(t *testing.T) { + path := t.TempDir() + "/claim-check.lock" + lock := flock.New(path) + locked, err := lock.TryLock() + if err != nil || !locked { + t.Fatalf("initial lock failed: locked=%v err=%v", locked, err) + } + attempts := 0 + if !releaseLocalizationClaimLockWith(lock, func(candidate *flock.Flock) error { + attempts++ + if attempts < localizationClaimLockReleaseAttempts { + return syscall.EINTR + } + return candidate.Close() + }) { + t.Fatal("EINTR retries did not release the lock") + } + if attempts != localizationClaimLockReleaseAttempts { + t.Fatalf("release attempts = %d, want %d", attempts, localizationClaimLockReleaseAttempts) + } + + challenger := flock.New(path) + locked, err = challenger.TryLock() + if err != nil || !locked { + t.Fatalf("successful retry retained ownership: locked=%v err=%v", locked, err) + } + if err := challenger.Close(); err != nil { + t.Fatalf("challenger release failed: %v", err) + } +} + +func TestLocalizationClaimLockReleaseExhaustionIsObservable(t *testing.T) { + logPath := filepath.Join(t.TempDir(), "effectiveness.jsonl") + t.Setenv("GORTEX_HOOK_EFFECTIVENESS_LOG", logPath) + lock := flock.New(filepath.Join(t.TempDir(), "claim-check.lock")) + locked, err := lock.TryLock() + if err != nil || !locked { + t.Fatalf("initial lock failed: locked=%v err=%v", locked, err) + } + attempts := 0 + if releaseLocalizationClaimLockObserved(lock, func(*flock.Flock) error { + attempts++ + return syscall.EINTR + }) { + t.Fatal("persistent EINTR was reported as a successful release") + } + if attempts != localizationClaimLockReleaseAttempts { + t.Fatalf("release attempts = %d, want %d", attempts, localizationClaimLockReleaseAttempts) + } + if err := lock.Close(); err != nil { + t.Fatalf("test cleanup release failed: %v", err) + } + data, err := os.ReadFile(logPath) + if err != nil { + t.Fatalf("read release telemetry: %v", err) + } + if !strings.Contains(string(data), `"event":"LocalizationTerminal.claim_lock_release_failed"`) { + t.Fatalf("release failure telemetry missing: %s", data) + } +} diff --git a/internal/hooks/localization_claim_check_hardening_test.go b/internal/hooks/localization_claim_check_hardening_test.go new file mode 100644 index 00000000..07e9b34f --- /dev/null +++ b/internal/hooks/localization_claim_check_hardening_test.go @@ -0,0 +1,139 @@ +package hooks + +import ( + "fmt" + "strings" + "sync" + "testing" +) + +func TestLocalizationClaimCheckPreservesQualifiedIdentity(t *testing.T) { + tests := []struct { + name string + claim string + evidenceID string + }{ + {name: "owner", claim: "Other.write", evidenceID: "repo/a.go::Writer.write"}, + {name: "file", claim: "repo/z.go::Writer.write", evidenceID: "repo/a.go::Writer.write"}, + {name: "rust owner", claim: "module::Other::write", evidenceID: "repo/x.rs::module::Writer::write"}, + {name: "php namespace", claim: `Other\Writer::write`, evidenceID: `repo/y.php::App\Writer::write`}, + {name: "hash owner", claim: "Other#write", evidenceID: "repo/z.py::Writer#write"}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + input := claimCheckTestInputWithEvidence(t, "SYMBOLS:\n- "+test.claim, []string{test.evidenceID}, []string{test.evidenceID}) + if got := localizationClaimCheck(input); got == "" { + t.Fatalf("qualified homonym %q authenticated as %q", test.claim, test.evidenceID) + } + }) + } +} + +func TestLocalizationClaimCheckAcceptsQualifiedLanguageNotation(t *testing.T) { + tests := []struct { + claim string + evidenceID string + }{ + {claim: "module::Writer::write", evidenceID: "repo/x.rs::module::Writer::write"}, + {claim: `App\Writer::write`, evidenceID: `repo/y.php::App\Writer::write`}, + {claim: "Writer#write", evidenceID: "repo/z.py::Writer#write"}, + } + for _, test := range tests { + t.Run(test.claim, func(t *testing.T) { + input := claimCheckTestInputWithEvidence(t, "SYMBOLS:\n- "+test.claim, []string{test.evidenceID}, []string{test.evidenceID}) + if got := localizationClaimCheck(input); got != "" { + t.Fatalf("authenticated language-qualified claim was challenged: %q", got) + } + }) + } +} + +func TestLocalizationClaimParserBounds(t *testing.T) { + prefix := "SYMBOLS:\n- Writer.write" + exactMessage := prefix + strings.Repeat(" ", localizationClaimCheckMaxMessageBytes-len(prefix)) + claims, _, valid := localizationBoundedSymbolClaims(exactMessage) + if !valid || len(claims) != 1 { + t.Fatalf("message at byte boundary rejected: valid=%v claims=%v", valid, claims) + } + if _, _, valid := localizationBoundedSymbolClaims(exactMessage + " "); valid { + t.Fatal("message above byte boundary was accepted") + } + + var claimLines strings.Builder + claimLines.WriteString("SYMBOLS:\n") + for index := 0; index < localizationClaimCheckMaxClaims; index++ { + fmt.Fprintf(&claimLines, "- Owner%d.write\n", index) + } + claims, _, valid = localizationBoundedSymbolClaims(claimLines.String()) + if !valid || len(claims) != localizationClaimCheckMaxClaims { + t.Fatalf("claim-count boundary rejected: valid=%v claims=%d", valid, len(claims)) + } + claimLines.WriteString("- Overflow.write\n") + if _, _, valid := localizationBoundedSymbolClaims(claimLines.String()); valid { + t.Fatal("claim count above boundary was accepted") + } + + exactToken := strings.Repeat("A", localizationClaimCheckMaxTokenBytes) + if _, _, valid := localizationBoundedSymbolClaims("SYMBOLS:\n- " + exactToken); !valid { + t.Fatal("claim token at byte boundary was rejected") + } + if _, _, valid := localizationBoundedSymbolClaims("SYMBOLS:\n- " + exactToken + "A"); valid { + t.Fatal("claim token above byte boundary was accepted") + } + + exactTokens := strings.Repeat("word ", localizationClaimCheckMaxTokens) + if _, _, valid := localizationBoundedSymbolClaims(exactTokens); !valid { + t.Fatal("unstructured token count at boundary was rejected") + } + if _, _, valid := localizationBoundedSymbolClaims(exactTokens + "word"); valid { + t.Fatal("unstructured token count above boundary was accepted") + } +} + +func TestLocalizationClaimCheckFailsClosedOnOversizedMessage(t *testing.T) { + input := claimCheckTestInput(t, strings.Repeat(" ", localizationClaimCheckMaxMessageBytes+1)) + if got := localizationClaimCheck(input); got == "" { + t.Fatal("oversized untrusted final message was not challenged") + } +} + +func TestLocalizationClaimCheckLockKeepsTerminalMarkerVisible(t *testing.T) { + input := claimCheckTestInput(t, "SYMBOLS:\n- Fabricated.flush") + identity, ok := currentLocalizationTurn(input.SessionID, input.PromptID, input.AgentID, input.CWD) + if !ok { + t.Fatal("localization turn missing") + } + release, ok := acquireLocalizationClaimCheck(localizationTerminalMarkerPath(identity)) + if !ok { + t.Fatal("claim-check lock was not acquired") + } + + const readers = 64 + missing := make(chan struct{}, readers) + var wait sync.WaitGroup + for index := 0; index < readers; index++ { + wait.Add(1) + go func() { + defer wait.Done() + if !hasLocalizationTerminal(identity) { + missing <- struct{}{} + } + }() + } + wait.Wait() + if len(missing) != 0 { + release() + t.Fatalf("%d concurrent readers observed a missing terminal marker", len(missing)) + } + if got := localizationClaimCheck(input); got != "" { + release() + t.Fatalf("contending claim check did not fail open: %q", got) + } + release() + if !hasLocalizationTerminal(identity) { + t.Fatal("terminal marker disappeared after releasing claim-check lock") + } + if got := localizationClaimCheck(input); got == "" { + t.Fatal("claim check was not available after releasing lock") + } +} diff --git a/internal/hooks/localization_claim_check_lock_test.go b/internal/hooks/localization_claim_check_lock_test.go new file mode 100644 index 00000000..d749517d --- /dev/null +++ b/internal/hooks/localization_claim_check_lock_test.go @@ -0,0 +1,46 @@ +package hooks + +import ( + "os" + "strings" + "testing" + "time" +) + +func TestLocalizationClaimCheckPersistentSidecarDoesNotWedge(t *testing.T) { + input := claimCheckTestInput(t, "SYMBOLS:\n- Fabricated.flush") + identity, ok := currentLocalizationTurn(input.SessionID, input.PromptID, input.AgentID, input.CWD) + if !ok { + t.Fatal("localization turn missing") + } + lockPath := localizationTerminalMarkerPath(identity) + ".claim-check" + if err := os.WriteFile(lockPath, []byte("stale"), 0o600); err != nil { + t.Fatal(err) + } + stale := time.Now().Add(-24 * time.Hour) + if err := os.Chtimes(lockPath, stale, stale); err != nil { + t.Fatal(err) + } + if got := localizationClaimCheck(input); got == "" { + t.Fatal("persistent stale lock sidecar wedged claim checking") + } +} + +func TestLocalizationClaimCheckConcurrentPreToolUseStaysDenied(t *testing.T) { + input := claimCheckTestInput(t, "SYMBOLS:\n- Fabricated.flush") + identity, ok := currentLocalizationTurn(input.SessionID, input.PromptID, input.AgentID, input.CWD) + if !ok { + t.Fatal("localization turn missing") + } + release, ok := acquireLocalizationClaimCheck(localizationTerminalMarkerPath(identity)) + if !ok { + t.Fatal("claim-check lock was not acquired") + } + defer release() + + payload := preToolPayload(t, "Read", "concurrent-read", identity, map[string]any{"file_path": "repo/a.go"}) + output := captureHookStdout(t, func() { runPreToolUse(payload, 0, ModeDeny) }) + if !strings.Contains(output, localizationTerminalDenyReason) || !strings.Contains(output, `"permissionDecision":"deny"`) { + t.Fatalf("concurrent PreToolUse bypassed terminal marker: %q", output) + } +} diff --git a/internal/hooks/localization_claim_check_once_test.go b/internal/hooks/localization_claim_check_once_test.go new file mode 100644 index 00000000..14c3b1d0 --- /dev/null +++ b/internal/hooks/localization_claim_check_once_test.go @@ -0,0 +1,47 @@ +package hooks + +import ( + "testing" + + "github.com/zzet/gortex/internal/localizationauth" +) + +func TestLocalizationClaimCheckSkipsAdvisoryAuthority(t *testing.T) { + configureLocalizationTerminalTestHome(t) + identity := beginTestLocalizationTurn(t, t.Name(), "prompt", t.TempDir()) + ids := []string{"repo/a.go::Writer.write"} + if !markLocalizationTerminalReceipt(identity, localizationauth.Receipt{ + FinalResponse: "Possible answer", PrimaryIDs: ids, EvidenceIDs: ids, + ContractVersion: localizationTerminalContractV2, Enforceable: false, + }) { + t.Fatal("advisory terminal marker was not written") + } + input := PostTaskInput{ + HookEventName: "Stop", SessionID: identity.SessionID, PromptID: identity.PromptID, + AgentID: identity.AgentID, CWD: identity.CWD, LastAssistantMessage: "SYMBOLS:\n- flush", + } + if got := localizationClaimCheck(input); got != "" { + t.Fatalf("advisory evidence must fail open, got %q", got) + } +} + +func TestLocalizationClaimCheckPersistsOneShotConsumption(t *testing.T) { + input := claimCheckTestInput(t, "SYMBOLS:\n- flush") + if first := localizationClaimCheck(input); first == "" { + t.Fatal("first unsupported claim was not challenged") + } + identity, ok := currentLocalizationTurn(input.SessionID, input.PromptID, input.AgentID, input.CWD) + if !ok { + t.Fatal("localization turn disappeared") + } + marker, ok := localizationTerminalMarkerFor(identity) + if !ok || !marker.ClaimCheckConsumed { + t.Fatalf("claim-check consumption was not persisted: %#v ok=%v", marker, ok) + } + // Hosts are not uniformly reliable about stop_hook_active. Persisted state, + // rather than that callback hint alone, makes a false retry acceptable. + input.StopHookActive = false + if retry := localizationClaimCheck(input); retry != "" { + t.Fatalf("persisted claim check fired twice: %q", retry) + } +} diff --git a/internal/hooks/localization_claim_check_roles_test.go b/internal/hooks/localization_claim_check_roles_test.go new file mode 100644 index 00000000..f96ab2cf --- /dev/null +++ b/internal/hooks/localization_claim_check_roles_test.go @@ -0,0 +1,115 @@ +package hooks + +import "testing" + +func TestLocalizationUnstructuredClaimsInspectHeadingBodiesWithoutClaimingRoleLabels(t *testing.T) { + tests := []struct { + name string + message string + want []string + }{ + {name: "ATX qualified identity", message: "# Writer.write", want: []string{"Writer.write"}}, + {name: "Setext qualified identity", message: "Writer.write\n---", want: []string{"Writer.write"}}, + {name: "quoted ATX identity", message: "> ## Writer.write", want: []string{"Writer.write"}}, + {name: "explicit simple heading", message: "# `Writer`", want: []string{"Writer"}}, + {name: "pointer receiver heading", message: "# `(*Server).handle`", want: []string{"Server.handle"}}, + {name: "bare receiver heading", message: "# (*Server).handle", want: []string{"Server.handle"}}, + {name: "receiver call heading", message: "# (*Server).handle()", want: []string{"Server.handle"}}, + {name: "package receiver heading", message: "# pkg.(*Server).handle", want: []string{"pkg.Server.handle"}}, + {name: "file receiver heading", message: "# server.go::(*Server).handle", want: []string{"server.go::Server.handle"}}, + {name: "PHP namespace heading", message: `# Vendor\Class`, want: []string{`Vendor\Class`}}, + {name: "double-colon ATX is identity", message: "# FILES::danger", want: []string{"FILES::danger"}}, + {name: "double-colon row is identity", message: "EVIDENCE::danger", want: []string{"EVIDENCE::danger"}}, + {name: "plain simple heading", message: "# Architecture"}, + {name: "API prose heading", message: "# API Design"}, + {name: "CommonMark prose heading", message: "# CommonMark behavior"}, + {name: "HTTP Setext prose heading", message: "HTTP Server\n---"}, + {name: "opening role heading", message: "### PRIMARY"}, + {name: "closing role heading", message: "### Implementation_details:"}, + {name: "empty inline roles", message: "PRIMARY:\nSUPPORTING:\nEvidence:\nImplementation_details:"}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + claims, _, valid := localizationBoundedSymbolClaims(test.message) + if !valid { + t.Fatal("message was rejected") + } + assertLocalizationClaims(t, claims, test.want) + }) + } +} + +func TestLocalizationClaimsRespectStructuredRoleState(t *testing.T) { + tests := []struct { + name string + message string + want []string + }{ + {name: "inline primary", message: "PRIMARY: Writer", want: []string{"Writer"}}, + {name: "primary heading", message: "# PRIMARY\nWriter", want: []string{"Writer"}}, + {name: "supporting label", message: "SUPPORTING:\nWriter", want: []string{"Writer"}}, + {name: "symbol heading", message: "# SYMBOL\nWriter", want: []string{"Writer"}}, + {name: "symbols label", message: "SYMBOLS:\nWriter", want: []string{"Writer"}}, + {name: "Setext primary heading", message: "PRIMARY\n===\nWriter", want: []string{"Writer"}}, + {name: "closing label", message: "PRIMARY:\nWriter\nFILES:\nReader", want: []string{"Writer"}}, + {name: "closing heading", message: "PRIMARY:\nWriter\n# EVIDENCE\nReader", want: []string{"Writer"}}, + {name: "non-role heading closes", message: "PRIMARY:\nWriter\n# Architecture\nReader", want: []string{"Writer"}}, + {name: "blank before content", message: "PRIMARY:\n\nWriter", want: []string{"Writer"}}, + {name: "blank after content closes", message: "PRIMARY:\nWriter\n\nReader", want: []string{"Writer"}}, + {name: "structured row wins over Setext", message: "PRIMARY:\nflush\n---", want: []string{"flush"}}, + {name: "cross-container underline cannot suppress row", message: "- PRIMARY:\n flush\n---", want: []string{"flush"}}, + {name: "same-container Setext role", message: "- PRIMARY\n ---\n Writer", want: []string{"Writer"}}, + {name: "sibling underline is not Setext role", message: "- PRIMARY\n- ---\n Writer"}, + {name: "nested closing label", message: "- PRIMARY:\n - FILES:\n Reader"}, + {name: "nested closing heading", message: "- PRIMARY:\n - # DETAILS\n Reader"}, + {name: "nested opening label", message: "- PRIMARY:\n - SUPPORTING:\n Writer", want: []string{"Writer"}}, + {name: "absolute-column list tab padding", message: " -\tPRIMARY:\n flush", want: []string{"flush"}}, + {name: "receiver prose", message: "The implementation calls (*Server).handle before returning.", want: []string{"Server.handle"}}, + {name: "receiver in fence", message: "```text\n(*Server).handle\n```", want: []string{"Server.handle"}}, + {name: "prose", message: "The implementation calls Fabricated.flush before returning.", want: []string{"Fabricated.flush"}}, + {name: "bare leaf", message: "The answer relies on flush."}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + claims, _, valid := localizationBoundedSymbolClaims(test.message) + if !valid { + t.Fatal("message was rejected") + } + assertLocalizationClaims(t, claims, test.want) + }) + } +} + +func TestLocalizationClaimsRejectExcessiveContainerDepth(t *testing.T) { + message := "> > > > > > > > > PRIMARY:\nWriter" + claims, _, valid := localizationBoundedSymbolClaims(message) + if valid { + t.Fatalf("claims = %v, want bounded parser rejection", claims) + } +} + +func TestLocalizationGoReceiverClaimPreservesPackageQualifier(t *testing.T) { + claims, _, valid := localizationBoundedSymbolClaims("# pkg.(*Server).handle") + if !valid { + t.Fatal("message was rejected") + } + assertLocalizationClaims(t, claims, []string{"pkg.Server.handle"}) + if localizationClaimMatchesEvidence(claims[0], "repo/a.go::other.Server.handle") { + t.Fatal("package-qualified receiver authenticated a different package") + } + if !localizationClaimMatchesEvidence(claims[0], "repo/a.go::pkg.Server.handle") { + t.Fatal("package-qualified receiver did not authenticate its exact evidence") + } + + claims, _, valid = localizationBoundedSymbolClaims("# server.go::(*Server).handle") + if !valid { + t.Fatal("file-qualified receiver message was rejected") + } + assertLocalizationClaims(t, claims, []string{"server.go::Server.handle"}) + if localizationClaimMatchesEvidence(claims[0], "other.go::Server.handle") { + t.Fatal("file-qualified receiver authenticated a different file") + } + if !localizationClaimMatchesEvidence(claims[0], "server.go::Server.handle") { + t.Fatal("file-qualified receiver did not authenticate its exact evidence") + } +} diff --git a/internal/hooks/localization_claim_check_unicode_test.go b/internal/hooks/localization_claim_check_unicode_test.go new file mode 100644 index 00000000..57884325 --- /dev/null +++ b/internal/hooks/localization_claim_check_unicode_test.go @@ -0,0 +1,24 @@ +package hooks + +import ( + "strings" + "testing" + "unicode/utf8" +) + +func TestBoundedLocalizationClaimCheckPreservesUTF8AndByteLimit(t *testing.T) { + // Place a multibyte rune across the nominal byte cut. The hook reason must + // remain valid UTF-8 because encoding/json otherwise replaces the partial + // rune and can expand the supposedly bounded response. + prompt := strings.Repeat("a", localizationClaimCheckMaxChars-35) + "界" + strings.Repeat("b", 100) + got := boundedLocalizationClaimCheck(prompt) + if !utf8.ValidString(got) { + t.Fatalf("bounded claim check is invalid UTF-8: %q", got) + } + if len(got) > localizationClaimCheckMaxChars { + t.Fatalf("bounded claim check bytes=%d want <=%d", len(got), localizationClaimCheckMaxChars) + } + if !strings.HasSuffix(got, "Do not retrieve more evidence.") { + t.Fatalf("bounded claim check lost terminal instruction: %q", got) + } +} diff --git a/internal/hooks/localization_claim_parser_followup_test.go b/internal/hooks/localization_claim_parser_followup_test.go new file mode 100644 index 00000000..34295854 --- /dev/null +++ b/internal/hooks/localization_claim_parser_followup_test.go @@ -0,0 +1,395 @@ +package hooks + +import ( + "os" + "path/filepath" + "strings" + "syscall" + "testing" + + "github.com/gofrs/flock" +) + +func TestLocalizationClaimParserPairsSetextWithinOneContainer(t *testing.T) { + tests := []struct { + name string + message string + want []string + }{ + {name: "top level heading", message: "Fabricated.flush\n---", want: []string{"Fabricated.flush"}}, + {name: "quoted heading", message: "> Fabricated.flush\n> ===", want: []string{"Fabricated.flush"}}, + {name: "structured row before thematic break", message: "SYMBOLS:\n- Fabricated.flush\n---", want: []string{"Fabricated.flush"}}, + {name: "list row before thematic break", message: "- Fabricated.flush\n---", want: []string{"Fabricated.flush"}}, + {name: "different quote container", message: "Fabricated.flush\n> ---", want: []string{"Fabricated.flush"}}, + {name: "new list cannot underline bare claim", message: "FabricatedThing\n- ---", want: []string{"FabricatedThing"}}, + {name: "new quote cannot underline bare claim", message: "FabricatedThing\n> ---", want: []string{"FabricatedThing"}}, + {name: "atx heading", message: "## Fabricated.flush", want: []string{"Fabricated.flush"}}, + {name: "fenced heading-like code", message: "```text\n# Fabricated.flush\n```", want: []string{"Fabricated.flush"}}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + claims, _, valid := localizationBoundedSymbolClaims(test.message) + if !valid { + t.Fatal("message was rejected") + } + assertLocalizationClaims(t, claims, test.want) + }) + } +} + +func TestLocalizationClaimParserPreservesIdentifierEdgeCharacters(t *testing.T) { + identities := []string{"_private", "$foo", "_Writer.write", "Writer.write_", "$Writer#write"} + for _, identity := range identities { + t.Run(identity, func(t *testing.T) { + claims, _, valid := localizationBoundedSymbolClaims("SYMBOLS:\n- `" + identity + "`()") + if !valid { + t.Fatal("message was rejected") + } + assertLocalizationClaims(t, claims, []string{identity}) + if !localizationClaimMatchesEvidence(identity, "repo/a.go::"+identity) { + t.Fatalf("exact identity %q did not authenticate", identity) + } + trimmed := strings.Trim(identity, "_$") + if trimmed != identity && localizationClaimMatchesEvidence(identity, "repo/a.go::"+trimmed) { + t.Fatalf("identity %q authenticated after edge characters were removed", identity) + } + }) + } +} + +func TestLocalizationClaimParserRequiresExplicitSyntaxForSimpleProseNames(t *testing.T) { + tests := []struct { + name string + message string + want []string + }{ + {name: "plain prose leaf", message: "The answer relies on flush."}, + {name: "inline code leaf", message: "The answer relies on `flush`.", want: []string{"flush"}}, + {name: "inline code call", message: "The answer relies on `flush()`.", want: []string{"flush"}}, + {name: "call expression", message: "The answer relies on flush().", want: []string{"flush"}}, + {name: "unmatched inline delimiter", message: "The answer relies on `flush."}, + {name: "file-qualified leaf", message: "The answer relies on repo/a.go::flush.", want: []string{"repo/a.go::flush"}}, + {name: "benign parenthetical prose", message: "The answer is (ordinary)."}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + claims, _, valid := localizationBoundedSymbolClaims(test.message) + if !valid { + t.Fatal("message was rejected") + } + assertLocalizationClaims(t, claims, test.want) + }) + } +} + +func TestLocalizationClaimParserDistinguishesFilesFromSymbols(t *testing.T) { + tests := []struct { + name string + message string + want []string + }{ + {name: "c source", message: "See foo.c."}, + {name: "markdown document", message: "See README.md."}, + {name: "sql schema", message: "See schema.sql."}, + {name: "vue component", message: "See Component.vue."}, + {name: "slash path", message: "See internal/parser/file.custom."}, + {name: "qualified symbol", message: "The owner is writer.flush.", want: []string{"writer.flush"}}, + {name: "file-qualified symbol", message: "The owner is foo.c::flush.", want: []string{"foo.c::flush"}}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + claims, _, valid := localizationBoundedSymbolClaims(test.message) + if !valid { + t.Fatal("message was rejected") + } + assertLocalizationClaims(t, claims, test.want) + }) + } +} + +func TestLocalizationClaimLockFailureTelemetryDoesNotClaimContext(t *testing.T) { + logPath := filepath.Join(t.TempDir(), "effectiveness.jsonl") + t.Setenv("GORTEX_HOOK_EFFECTIVENESS_LOG", logPath) + lock := flock.New(filepath.Join(t.TempDir(), "claim-check.lock")) + locked, err := lock.TryLock() + if err != nil || !locked { + t.Fatalf("initial lock failed: locked=%v err=%v", locked, err) + } + if releaseLocalizationClaimLockObserved(lock, func(*flock.Flock) error { return syscall.EINTR }) { + t.Fatal("persistent EINTR was reported as a successful release") + } + if err := lock.Close(); err != nil { + t.Fatalf("test cleanup release failed: %v", err) + } + data, err := os.ReadFile(logPath) + if err != nil { + t.Fatalf("read release telemetry: %v", err) + } + record := string(data) + if !strings.Contains(record, `"event":"LocalizationTerminal.claim_lock_release_failed"`) || + !strings.Contains(record, `"emitted_context":false`) { + t.Fatalf("release failure telemetry has wrong context semantics: %s", record) + } +} + +func TestLocalizationClaimParserRespectsIndentedCode(t *testing.T) { + tests := []struct { + name string + message string + want []string + }{ + {name: "zero-space Setext", message: "Fabricated.flush\n---", want: []string{"Fabricated.flush"}}, + {name: "three-space Setext", message: " Fabricated.flush\n ---", want: []string{"Fabricated.flush"}}, + {name: "four-space code", message: " Fabricated.flush\n ---", want: []string{"Fabricated.flush"}}, + {name: "tab-indented code", message: "\tFabricated.flush\n\t---", want: []string{"Fabricated.flush"}}, + {name: "top-level claim and code underline", message: "Fabricated.flush\n ---", want: []string{"Fabricated.flush"}}, + {name: "quoted Setext", message: "> Fabricated.flush\n> ---", want: []string{"Fabricated.flush"}}, + {name: "different quote containers", message: "> Fabricated.flush\n>> ---", want: []string{"Fabricated.flush"}}, + {name: "structured claim before thematic break", message: "SYMBOLS:\n- Fabricated.flush\n---", want: []string{"Fabricated.flush"}}, + {name: "list claim before thematic break", message: "- Fabricated.flush\n---", want: []string{"Fabricated.flush"}}, + {name: "atx heading", message: "# Fabricated.flush", want: []string{"Fabricated.flush"}}, + {name: "fenced heading-like code", message: "```text\n# Fabricated.flush\n```", want: []string{"Fabricated.flush"}}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + claims, _, valid := localizationBoundedSymbolClaims(test.message) + if !valid { + t.Fatal("message was rejected") + } + assertLocalizationClaims(t, claims, test.want) + }) + } +} + +func TestLocalizationClaimParserKeepsAmbiguousExtensionSymbols(t *testing.T) { + for _, identity := range []string{"obj.m", "module.r", "node.s", "pkg.v", "value.d"} { + t.Run(identity+" symbol", func(t *testing.T) { + claims, _, valid := localizationBoundedSymbolClaims("The owner is " + identity + ".") + if !valid { + t.Fatal("message was rejected") + } + assertLocalizationClaims(t, claims, []string{identity}) + }) + t.Run(identity+" file context", func(t *testing.T) { + claims, _, valid := localizationBoundedSymbolClaims("The source file is " + identity + ".") + if !valid { + t.Fatal("message was rejected") + } + assertLocalizationClaims(t, claims, nil) + }) + t.Run(identity+" structured", func(t *testing.T) { + claims, _, valid := localizationBoundedSymbolClaims("SYMBOLS:\n- " + identity) + if !valid { + t.Fatal("message was rejected") + } + assertLocalizationClaims(t, claims, []string{identity}) + }) + } + claims, _, valid := localizationBoundedSymbolClaims("See file foo.c.") + if !valid { + t.Fatal("C file message was rejected") + } + assertLocalizationClaims(t, claims, nil) +} + +func TestLocalizationStructuredClaimsNormalizeSentencePunctuation(t *testing.T) { + for _, test := range []struct { + row string + want string + }{ + {row: "flush,", want: "flush"}, + {row: "Writer;", want: "Writer"}, + {row: "_private.", want: "_private"}, + {row: "$foo:", want: "$foo"}, + {row: "`flush`,", want: "flush"}, + {row: "flush(),", want: "flush"}, + } { + t.Run(test.row, func(t *testing.T) { + claims, _, valid := localizationBoundedSymbolClaims("SYMBOLS:\n- " + test.row) + if !valid { + t.Fatal("message was rejected") + } + assertLocalizationClaims(t, claims, []string{test.want}) + }) + } + claims, _, valid := localizationBoundedSymbolClaims("SYMBOLS:\n- Writer performs the work.") + if !valid { + t.Fatal("prose row was rejected") + } + assertLocalizationClaims(t, claims, nil) +} + +func TestLocalizationClaimParserRecognizesStableFileVocabulary(t *testing.T) { + files := []string{ + "README", "LICENSE", "CHANGELOG", "Dockerfile.ci", "Makefile.am", "CMakeLists.txt", + "go.mod", "go.sum", "go.work", "Cargo.lock", "notebook.ipynb", "contract.sol", + "module.move", "program.cairo", "circuit.noir", "contract.tact", "service.bal", + "settings.toml", "page.vue", "template.gotmpl", "document.adoc", + } + for _, file := range files { + t.Run(file, func(t *testing.T) { + claims, _, valid := localizationBoundedSymbolClaims("See " + file + ".") + if !valid { + t.Fatal("message was rejected") + } + assertLocalizationClaims(t, claims, nil) + }) + } + for _, test := range []struct { + message string + want string + }{ + {message: "The owner is writer.flush.", want: "writer.flush"}, + {message: "The owner is obj.m.", want: "obj.m"}, + {message: "The owner is file.ext::symbol.", want: "file.ext::symbol"}, + } { + claims, _, valid := localizationBoundedSymbolClaims(test.message) + if !valid { + t.Fatalf("message %q was rejected", test.message) + } + assertLocalizationClaims(t, claims, []string{test.want}) + } +} + +func TestLocalizationClaimParserKeepsIndentedPseudoFenceOpen(t *testing.T) { + message := "```text\ninside\n ```\n# Fabricated.flush\n```" + claims, _, valid := localizationBoundedSymbolClaims(message) + if !valid { + t.Fatal("message was rejected") + } + assertLocalizationClaims(t, claims, []string{"Fabricated.flush"}) +} + +func TestLocalizationClaimParserKeepsFenceWithinOpeningContainer(t *testing.T) { + tests := []struct { + name string + message string + want []string + }{ + { + name: "bullet continuation closes", + message: "PRIMARY:\n- ```text\n harmless\n ```\n- Writer", + want: []string{"Writer"}, + }, + { + name: "ordered continuation closes", + message: "PRIMARY:\n1. ```text\n harmless\n ```\n2. Writer", + want: []string{"Writer"}, + }, + { + name: "nested continuation closes", + message: "PRIMARY:\n- 1. ```text\n harmless\n ```\n- 2. Writer", + want: []string{"Writer"}, + }, + { + name: "quote then list continuation closes", + message: "> PRIMARY:\n> - ```text\n> harmless\n> ```\n> - Writer", + want: []string{"Writer"}, + }, + { + name: "list then quote continuation closes", + message: "- PRIMARY:\n > ```text\n > harmless\n > ```\n > Writer", + want: []string{"Writer"}, + }, + { + name: "sibling item cannot close", + message: "PRIMARY:\n- ```text\n harmless\n- ```\n Writer\n ```", + }, + { + name: "list exit reprocesses line", + message: "- ```text\n harmless\nPRIMARY:\nWriter", + want: []string{"Writer"}, + }, + { + name: "quote exit reprocesses line", + message: "> ```text\n> harmless\nPRIMARY:\nWriter", + want: []string{"Writer"}, + }, + { + name: "nested fence opener captures child path", + message: "- PRIMARY:\n - ```text\n Writer\n ```\n - Reader", + want: []string{"Reader"}, + }, + { + name: "list tab overshoot stays code indented", + message: "PRIMARY:\n- ```text\n\t ```\n Writer\n ```", + }, + { + name: "quote tab overshoot stays code indented", + message: "PRIMARY:\n> ```text\n>\t ```\n> Writer\n> ```", + }, + { + name: "three-space closer indentation is legal", + message: "PRIMARY:\n- ```text\n harmless\n ```\n- Writer", + want: []string{"Writer"}, + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + claims, _, valid := localizationBoundedSymbolClaims(test.message) + if !valid { + t.Fatal("message was rejected") + } + assertLocalizationClaims(t, claims, test.want) + }) + } +} + +func TestLocalizationAmbiguousFileContextIsPrefixAndClauseBounded(t *testing.T) { + tests := []struct { + name string + message string + want []string + }{ + {name: "direct file prefix", message: "Use file foo.m."}, + {name: "colon path prefix", message: "Use path: foo.r."}, + {name: "file copula prefix", message: "The source file is node.s."}, + {name: "unsafe suffix adjacency", message: "Call obj.m file handler.", want: []string{"obj.m"}}, + {name: "prior clause does not leak", message: "Inspect the source file; the owner is pkg.v.", want: []string{"pkg.v"}}, + {name: "prior sentence does not leak", message: "Inspect the path. The owner is value.d.", want: []string{"value.d"}}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + claims, _, valid := localizationBoundedSymbolClaims(test.message) + if !valid { + t.Fatal("message was rejected") + } + assertLocalizationClaims(t, claims, test.want) + }) + } +} + +func TestLocalizationExplicitSyntaxOverridesFileExtension(t *testing.T) { + tests := []struct { + name string + message string + want []string + }{ + {name: "plain file", message: "See foo.go."}, + {name: "inline dotted identity", message: "The owner is `writer.go`.", want: []string{"writer.go"}}, + {name: "dotted call", message: "The owner calls writer.go().", want: []string{"writer.go"}}, + {name: "inline config identity", message: "The owner is `config.json`.", want: []string{"config.json"}}, + {name: "file-qualified remains claim", message: "The owner is foo.go::writer.", want: []string{"foo.go::writer"}}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + claims, _, valid := localizationBoundedSymbolClaims(test.message) + if !valid { + t.Fatal("message was rejected") + } + assertLocalizationClaims(t, claims, test.want) + }) + } +} + +func assertLocalizationClaims(t *testing.T, got, want []string) { + t.Helper() + if len(got) != len(want) { + t.Fatalf("claims = %v, want %v", got, want) + } + for index := range want { + if got[index] != want[index] { + t.Fatalf("claims = %v, want %v", got, want) + } + } +} diff --git a/internal/hooks/localization_claim_parser_regression_test.go b/internal/hooks/localization_claim_parser_regression_test.go new file mode 100644 index 00000000..e16566d7 --- /dev/null +++ b/internal/hooks/localization_claim_parser_regression_test.go @@ -0,0 +1,108 @@ +package hooks + +import "testing" + +func TestLocalizationClaimParserIgnoresBenignPunctuation(t *testing.T) { + messages := []string{ + "Release v1.2.3 took 3.14 seconds at 12:30 from 127.0.0.1.", + "For example, e.g. and i.e. are ordinary prose.", + } + for _, message := range messages { + claims, _, valid := localizationBoundedSymbolClaims(message) + if !valid { + t.Fatalf("benign message was rejected: %q", message) + } + if len(claims) != 0 { + t.Fatalf("benign punctuation became symbol claims: message=%q claims=%v", message, claims) + } + } +} + +func TestLocalizationClaimParserHonorsMarkdownStructure(t *testing.T) { + tests := []struct { + name string + message string + want string + }{ + {name: "atx heading", message: "# Fabricated.flush", want: "Fabricated.flush"}, + {name: "setext heading", message: "Fabricated.flush\n---", want: "Fabricated.flush"}, + {name: "container atx heading", message: "> # Fabricated.flush", want: "Fabricated.flush"}, + {name: "container setext heading", message: "> Fabricated.flush\n> ===", want: "Fabricated.flush"}, + {name: "backtick fenced heading-like code", message: "```go\n# Fabricated.flush\n```", want: "Fabricated.flush"}, + {name: "tilde fenced heading-like code", message: "~~~text\n## Fabricated.flush\n~~~~", want: "Fabricated.flush"}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + claims, _, valid := localizationBoundedSymbolClaims(test.message) + if !valid { + t.Fatal("message was rejected") + } + if test.want == "" { + if len(claims) != 0 { + t.Fatalf("heading became a claim: %v", claims) + } + return + } + if len(claims) != 1 || claims[0] != test.want { + t.Fatalf("claims = %v, want %q", claims, test.want) + } + }) + } +} + +func TestLocalizationStructuredClaimsRequireIdentityRows(t *testing.T) { + tests := []struct { + name string + message string + want string + }{ + {name: "bare identity row", message: "SYMBOLS:\n- Writer", want: "Writer"}, + {name: "backticked identity row", message: "SYMBOLS:\n- `Writer`", want: "Writer"}, + {name: "bare prose row", message: "SYMBOLS:\n- Writer performs the work", want: ""}, + {name: "qualified claim in prose remainder", message: "SYMBOLS:\n- This row names Fabricated.flush after context", want: "Fabricated.flush"}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + claims, _, valid := localizationBoundedSymbolClaims(test.message) + if !valid { + t.Fatal("message was rejected") + } + if test.want == "" { + if len(claims) != 0 { + t.Fatalf("prose row produced a bare claim: %v", claims) + } + return + } + if len(claims) != 1 || claims[0] != test.want { + t.Fatalf("claims = %v, want %q", claims, test.want) + } + }) + } +} + +func TestLocalizationQualifiedSuffixUsesLiteralBoundaries(t *testing.T) { + for _, test := range []struct { + claim string + evidence string + }{ + {claim: "Writer::write", evidence: "pkg::Outer::Writer::write"}, + {claim: "Writer#write", evidence: "pkg::Outer#Writer#write"}, + {claim: "Writer#write", evidence: "pkg::Outer::Writer#write"}, + } { + if !localizationQualifiedSymbolMatches(test.claim, test.evidence) { + t.Errorf("claim %q did not match qualified evidence %q", test.claim, test.evidence) + } + } + for _, test := range []struct { + claim string + evidence string + }{ + {claim: "Writer::write", evidence: "pkg::NotWriter::write"}, + {claim: "Writer#write", evidence: "pkg::writer#write"}, + {claim: "Writer#write", evidence: "pkg::NotWriter#write"}, + } { + if localizationQualifiedSymbolMatches(test.claim, test.evidence) { + t.Errorf("claim %q crossed a non-boundary in %q", test.claim, test.evidence) + } + } +} diff --git a/internal/hooks/localization_terminal.go b/internal/hooks/localization_terminal.go index f66452b3..b599c769 100644 --- a/internal/hooks/localization_terminal.go +++ b/internal/hooks/localization_terminal.go @@ -5,13 +5,17 @@ import ( "crypto/sha256" "encoding/hex" "encoding/json" + "errors" "fmt" "os" "path/filepath" "sort" "strings" + "syscall" "time" + "github.com/gofrs/flock" + "github.com/zzet/gortex/internal/localizationauth" ) @@ -29,6 +33,7 @@ const ( localizationTerminalDenyReason = "[Gortex] Localization for this task is complete, so this tool call is blocked. Answer now from the retained evidence below, naming what you rely on; if it does not fit the request, say so in your answer." localizationAdvisoryDenyReason = "[Gortex] Localization for this task is complete, so this additional Gortex navigation call was not run. Answer now from the retained evidence below, naming what you rely on; if it does not fit the request, say so in your answer." gortexPluginMCPToolPrefix = "mcp__plugin_gortex_gortex__" + gortexCodexMCPToolPrefix = "gortex__" localizationHostMetaKey = "gortex/localization" ) @@ -106,7 +111,10 @@ type localizationTerminalMarker struct { // FinalResponse rides the marker so a blocked call can hand back the answer // itself. A refusal that only says "you are done" gives the caller nothing // to act on and it tries the next tool; the answer ends the turn. - FinalResponse string `json:"final_response,omitempty"` + FinalResponse string `json:"final_response,omitempty"` + PrimaryIDs []string `json:"primary_ids,omitempty"` + EvidenceIDs []string `json:"evidence_ids,omitempty"` + ClaimCheckConsumed bool `json:"claim_check_consumed,omitempty"` } type localizationTerminalHookInput struct { @@ -182,7 +190,7 @@ func observeLocalizationTerminal(data []byte) (localizationTerminalHookInput, bo // PreToolUse snapshot. Visible tool JSON is never an authority on this path. if authToken != "" { if receipt, authenticated := localizationauth.Consume(authToken); authenticated { - if !markLocalizationTerminalReceipt(identity, receipt.ContractVersion, receipt.Enforceable, receipt.FinalResponse) { + if !markLocalizationTerminalReceipt(identity, receipt) { return localizationTerminalHookInput{}, false } input.TerminalReceipt = receipt @@ -197,7 +205,10 @@ func observeLocalizationTerminal(data []byte) (localizationTerminalHookInput, bo if !ok || !answerReadyLocalizationTerminalContract(contract) { return localizationTerminalHookInput{}, false } - if !markLocalizationTerminalReceipt(identity, contract.Completion.ContractVersion, contract.Completion.Enforceable, contract.Completion.FinalResponse) { + if !markLocalizationTerminalReceipt(identity, localizationauth.Receipt{ + FinalResponse: contract.Completion.FinalResponse, ContractVersion: contract.Completion.ContractVersion, + Enforceable: contract.Completion.Enforceable, + }) { return localizationTerminalHookInput{}, false } input.TerminalReceipt = localizationauth.Receipt{ @@ -348,7 +359,8 @@ func localizationNavigationTool(tool string) bool { } func isGortexMCPToolName(tool string) bool { - return strings.HasPrefix(tool, gortexMCPToolPrefix) || strings.HasPrefix(tool, gortexPluginMCPToolPrefix) + return strings.HasPrefix(tool, gortexMCPToolPrefix) || strings.HasPrefix(tool, gortexPluginMCPToolPrefix) || + strings.HasPrefix(tool, gortexCodexMCPToolPrefix) } func preToolUsePolicyTool(tool string) bool { @@ -641,14 +653,24 @@ func markLocalizationTerminal(identity localizationTerminalIdentity, contractVer return markLocalizationTerminalWithStrength(identity, contractVersion, false, "") } -func markLocalizationTerminalReceipt( - identity localizationTerminalIdentity, contractVersion int, enforceable bool, finalResponse string, -) bool { - return markLocalizationTerminalWithStrength(identity, contractVersion, !enforceable, finalResponse) +func markLocalizationTerminalReceipt(identity localizationTerminalIdentity, receipt localizationauth.Receipt) bool { + primaryIDs, evidenceIDs, ok := localizationauth.NormalizeEvidenceIDs(receipt.PrimaryIDs, receipt.EvidenceIDs) + if !ok { + return false + } + return markLocalizationTerminalWithEvidence(identity, receipt.ContractVersion, !receipt.Enforceable, + receipt.FinalResponse, primaryIDs, evidenceIDs) } func markLocalizationTerminalWithStrength( identity localizationTerminalIdentity, contractVersion int, advisory bool, finalResponse string, +) bool { + return markLocalizationTerminalWithEvidence(identity, contractVersion, advisory, finalResponse, nil, nil) +} + +func markLocalizationTerminalWithEvidence( + identity localizationTerminalIdentity, contractVersion int, advisory bool, finalResponse string, + primaryIDs, evidenceIDs []string, ) bool { if identity.SessionID == "" || identity.CWD == "" || identity.TurnToken == "" || contractVersion < localizationTerminalContractV2 { @@ -661,6 +683,8 @@ func markLocalizationTerminalWithStrength( ObservedUnixNano: time.Now().UnixNano(), Advisory: advisory, FinalResponse: strings.TrimSpace(finalResponse), + PrimaryIDs: append([]string(nil), primaryIDs...), + EvidenceIDs: append([]string(nil), evidenceIDs...), } path := localizationTerminalMarkerPath(identity) if advisory { @@ -679,6 +703,103 @@ func localizationTerminalMarkerFor(identity localizationTerminalIdentity) (local return localizationTerminalMarker{}, false } +// consumeLocalizationClaimCheck serializes one enforceable terminal marker's +// challenge without ever removing the marker from its canonical path. Keeping +// that path readable closes the window in which a concurrent PreToolUse could +// bypass terminal enforcement. Host-retried Stop hooks still fail open after +// the first successful challenge, even when the host omits stop_hook_active. +func consumeLocalizationClaimCheck(identity localizationTerminalIdentity, claims []string) (localizationTerminalMarker, bool) { + if len(claims) == 0 { + return localizationTerminalMarker{}, false + } + path := localizationTerminalMarkerPath(identity) + release, ok := acquireLocalizationClaimCheck(path) + if !ok { + return localizationTerminalMarker{}, false + } + defer release() + + marker, ok := readLocalizationTerminalMarker(path, identity) + if !ok || marker.Advisory || marker.ClaimCheckConsumed || len(marker.PrimaryIDs) == 0 || len(marker.EvidenceIDs) == 0 { + return localizationTerminalMarker{}, false + } + allAuthenticated := true + for _, claim := range claims { + matched := false + for _, evidenceID := range marker.EvidenceIDs { + if localizationClaimMatchesEvidence(claim, evidenceID) { + matched = true + break + } + } + if !matched { + allAuthenticated = false + break + } + } + if allAuthenticated { + return localizationTerminalMarker{}, false + } + marker.ClaimCheckConsumed = true + if !writeLocalizationState(path, marker) { + return localizationTerminalMarker{}, false + } + return marker, true +} + +func acquireLocalizationClaimCheck(path string) (func(), bool) { + // Kernel-backed advisory locking makes the sidecar reusable: a process + // crash releases ownership automatically, so a stale path can never wedge + // future claim checks. Do not remove the sidecar on unlock; another process + // may already have opened the same inode while waiting to contend. + lock := flock.New(path + ".claim-check") + locked, err := lock.TryLock() + if err != nil || !locked { + return nil, false + } + return func() { releaseLocalizationClaimLock(lock) }, true +} + +const localizationClaimLockReleaseAttempts = 3 + +func releaseLocalizationClaimLock(lock *flock.Flock) { + releaseLocalizationClaimLockObserved(lock, func(candidate *flock.Flock) error { + // Close is gofrs/flock's documented release API: it unlocks and closes + // the underlying descriptor while leaving the reusable sidecar in place. + return candidate.Close() + }) +} + +func releaseLocalizationClaimLockObserved(lock *flock.Flock, release func(*flock.Flock) error) bool { + started := time.Now() + if releaseLocalizationClaimLockWith(lock, release) { + return true + } + // Lock-release telemetry describes an internal failure; no user-visible + // hook context was emitted by this path. + localizationTerminalTelemetry("claim_lock_release_failed", false, started) + return false +} + +func releaseLocalizationClaimLockWith(lock *flock.Flock, release func(*flock.Flock) error) bool { + if lock == nil { + return true + } + if release == nil { + return false + } + for attempt := 0; attempt < localizationClaimLockReleaseAttempts; attempt++ { + err := release(lock) + if err == nil { + return true + } + if !errors.Is(err, syscall.EINTR) { + return false + } + } + return false +} + func readLocalizationTerminalMarker(path string, identity localizationTerminalIdentity) (localizationTerminalMarker, bool) { var marker localizationTerminalMarker if !readLocalizationState(path, &marker) || !freshLocalizationTimestamp(path, marker.ObservedUnixNano) || diff --git a/internal/hooks/localization_terminal_receipt_test.go b/internal/hooks/localization_terminal_receipt_test.go index 7ac8bf8d..147a2568 100644 --- a/internal/hooks/localization_terminal_receipt_test.go +++ b/internal/hooks/localization_terminal_receipt_test.go @@ -3,6 +3,7 @@ package hooks import ( "encoding/json" "os" + "slices" "strings" "sync" "sync/atomic" @@ -174,6 +175,26 @@ func TestLocalizationReceiptMarkerStrengthControlsPreToolUse(t *testing.T) { } } +func TestLocalizationTerminalMarkerRetainsAuthenticatedEvidenceIDs(t *testing.T) { + configureLocalizationTerminalTestHome(t) + identity := beginTestLocalizationTurn(t, t.Name(), "prompt", t.TempDir()) + primary := []string{"repo/a.go::A", "repo/b.go::B"} + evidence := []string{"repo/a.go::A", "repo/b.go::B", "repo/c.go::C"} + if !markLocalizationTerminalReceipt(identity, localizationauth.Receipt{ + FinalResponse: "answer", PrimaryIDs: primary, EvidenceIDs: evidence, + ContractVersion: localizationTerminalContractV2, Enforceable: true, + }) { + t.Fatal("marker was not written") + } + marker, ok := localizationTerminalMarkerFor(identity) + if !ok { + t.Fatal("marker was not readable") + } + if !slices.Equal(marker.PrimaryIDs, primary) || !slices.Equal(marker.EvidenceIDs, evidence) { + t.Fatalf("marker evidence authority = primary %#v evidence %#v", marker.PrimaryIDs, marker.EvidenceIDs) + } +} + func TestLocalizationTerminalMarkerStrengthIsMonotonic(t *testing.T) { configureLocalizationTerminalTestHome(t) identity := beginTestLocalizationTurn(t, "monotonic-strength", "prompt", t.TempDir()) @@ -251,7 +272,7 @@ func TestLocalizationAdvisoryMarkerRotatesAndDelayedPostCannotPoisonNewTurn(t *t if _, observed := observeLocalizationTerminal(pendingPost); observed { t.Fatal("pending pre-rotation advisory receipt armed the next turn") } - if !markLocalizationTerminalReceipt(oldTurn, localizationTerminalContractV2, false, "") { + if !markLocalizationTerminalReceipt(oldTurn, localizationauth.Receipt{ContractVersion: localizationTerminalContractV2}) { t.Fatal("simulate delayed advisory marker write") } if _, marked := localizationTerminalMarkerFor(newTurn); marked { diff --git a/internal/hooks/posttask.go b/internal/hooks/posttask.go index 0cd07f30..d985403b 100644 --- a/internal/hooks/posttask.go +++ b/internal/hooks/posttask.go @@ -24,13 +24,14 @@ const postTaskScopePhrase = "uncommitted (staged + unstaged) changes to tracked // Stop hook asked the agent to continue) — we must skip in that case to // avoid recursion. type PostTaskInput struct { - HookEventName string `json:"hook_event_name"` - SessionID string `json:"session_id"` - PromptID string `json:"prompt_id"` - AgentID string `json:"agent_id"` - TranscriptPath string `json:"transcript_path"` - CWD string `json:"cwd"` - StopHookActive bool `json:"stop_hook_active"` + HookEventName string `json:"hook_event_name"` + SessionID string `json:"session_id"` + PromptID string `json:"prompt_id"` + AgentID string `json:"agent_id"` + TranscriptPath string `json:"transcript_path"` + CWD string `json:"cwd"` + LastAssistantMessage string `json:"last_assistant_message"` + StopHookActive bool `json:"stop_hook_active"` } // runPostTask handles a Stop hook invocation with the raw stdin bytes. @@ -58,6 +59,15 @@ func runPostTask(data []byte, port int) { if input.StopHookActive { return } + if claimCheck := localizationClaimCheck(input); claimCheck != "" { + output := HookOutput{Decision: "block", Reason: claimCheck} + out, err := json.Marshal(output) + if err == nil { + emitted = true + fmt.Print(string(out)) + } + return + } briefing := buildPostTaskBriefing(port, sessionScope{ SessionID: input.SessionID, diff --git a/internal/hooks/posttask_test.go b/internal/hooks/posttask_test.go index 576338f7..0e582e92 100644 --- a/internal/hooks/posttask_test.go +++ b/internal/hooks/posttask_test.go @@ -7,6 +7,7 @@ import ( "testing" "github.com/zzet/gortex/internal/daemon" + "github.com/zzet/gortex/internal/localizationauth" ) func TestRunPostTask_RejectsWrongEvent(t *testing.T) { @@ -424,6 +425,127 @@ func TestRunPostTask_OwnsAll_SkipsImpactCall(t *testing.T) { } } +func claimCheckTestInput(t *testing.T, message string) PostTaskInput { + t.Helper() + primary := []string{"repo/a.go::Writer.write", "repo/b.go::Reader.read", "repo/c.go::Store.load", "repo/d.go::Cache.get", "repo/e.go::Index.find"} + evidence := append(append([]string(nil), primary...), "repo/f.go::Helper.close") + return claimCheckTestInputWithEvidence(t, message, primary, evidence) +} + +func claimCheckTestInputWithEvidence(t *testing.T, message string, primary, evidence []string) PostTaskInput { + t.Helper() + configureLocalizationTerminalTestHome(t) + identity := beginTestLocalizationTurn(t, t.Name(), "prompt", t.TempDir()) + if !markLocalizationTerminalReceipt(identity, localizationauth.Receipt{ + FinalResponse: "answer", PrimaryIDs: primary, EvidenceIDs: evidence, + ContractVersion: localizationTerminalContractV2, Enforceable: true, + }) { + t.Fatal("terminal marker was not written") + } + return PostTaskInput{ + HookEventName: "Stop", SessionID: identity.SessionID, PromptID: identity.PromptID, + AgentID: identity.AgentID, CWD: identity.CWD, LastAssistantMessage: message, + } +} + +func TestLocalizationClaimCheckAcceptsAuthenticatedClaim(t *testing.T) { + input := claimCheckTestInput(t, "SYMBOLS:\n- write") + if got := localizationClaimCheck(input); got != "" { + t.Fatalf("matching evidence was challenged: %q", got) + } + input.LastAssistantMessage = "The implementation is in repo/a.go." + if got := localizationClaimCheck(input); got != "" { + t.Fatalf("a file path in prose was treated as a symbol claim: %q", got) + } +} + +func TestLocalizationClaimCheckNormalizesCommonMethodNotation(t *testing.T) { + input := claimCheckTestInput(t, "SYMBOLS:\n- Writer.write()") + if got := localizationClaimCheck(input); got != "" { + t.Fatalf("qualified method call was challenged: %q", got) + } + input.LastAssistantMessage = "SYMBOLS:\n- (*Writer).write" + if got := localizationClaimCheck(input); got != "" { + t.Fatalf("pointer-receiver method was challenged: %q", got) + } +} + +func TestLocalizationClaimCheckRequiresEveryMaterialClaim(t *testing.T) { + input := claimCheckTestInput(t, "SYMBOLS:\n- Writer.write\n- Fabricated.flush") + if got := localizationClaimCheck(input); got == "" { + t.Fatal("a fabricated claim beside an authenticated claim was not challenged") + } +} + +func TestLocalizationClaimCheckAcceptsExplicitNoneFitsOnly(t *testing.T) { + input := claimCheckTestInput(t, "SYMBOLS:\n- none fits") + if got := localizationClaimCheck(input); got != "" { + t.Fatalf("explicit none-fits response was challenged: %q", got) + } + input.LastAssistantMessage = "SYMBOLS:\n- none fits\n- flush" + if got := localizationClaimCheck(input); got == "" { + t.Fatal("an unsupported claim hidden beside none-fits was not challenged") + } +} + +func TestLocalizationClaimCheckChallengesWrongBareSymbolWithinBound(t *testing.T) { + input := claimCheckTestInput(t, "SYMBOLS:\n- flush") + got := localizationClaimCheck(input) + if got == "" { + t.Fatal("wrong structured bare symbol was not challenged") + } + if len(got) > localizationClaimCheckMaxChars { + t.Fatalf("claim check is %d chars, max %d", len(got), localizationClaimCheckMaxChars) + } + for _, want := range []string{"repo/a.go::Writer.write", "repo/e.go::Index.find", "explicitly confirm", "Do not retrieve"} { + if !strings.Contains(got, want) { + t.Errorf("claim check missing %q: %s", want, got) + } + } + if strings.Contains(got, "repo/f.go::Helper.close") { + t.Fatalf("claim check exposed non-PRIMARY evidence: %s", got) + } +} + +func TestLocalizationClaimCheckFailsOpenWithoutMessageAuthorityOrOnRetry(t *testing.T) { + input := claimCheckTestInput(t, "SYMBOLS:\n- wrong") + input.StopHookActive = true + if got := localizationClaimCheck(input); got != "" { + t.Fatalf("retry was challenged twice: %q", got) + } + input.StopHookActive = false + input.LastAssistantMessage = "" + if got := localizationClaimCheck(input); got != "" { + t.Fatalf("missing final message did not fail open: %q", got) + } + input.LastAssistantMessage = "SYMBOLS:\n- wrong" + input.SessionID = "unsupported-host-without-authority" + if got := localizationClaimCheck(input); got != "" { + t.Fatalf("missing terminal authority did not fail open: %q", got) + } +} + +func TestRunPostTaskClaimCheckBlocksOnceWithoutDaemonRetrieval(t *testing.T) { + input := claimCheckTestInput(t, "SYMBOLS:\n- flush") + data, err := json.Marshal(input) + if err != nil { + t.Fatal(err) + } + out := captureStdout(t, func() { runPostTask(data, 0) }) + var payload HookOutput + if err := json.Unmarshal([]byte(out), &payload); err != nil { + t.Fatalf("decode claim-check output %q: %v", out, err) + } + if payload.Decision != "block" || !strings.Contains(payload.Reason, "claim_check") { + t.Fatalf("Stop was not blocked with a claim check: %#v", payload) + } + input.StopHookActive = true + data, _ = json.Marshal(input) + if retry := captureStdout(t, func() { runPostTask(data, 0) }); retry != "" { + t.Fatalf("second Stop response was not accepted: %q", retry) + } +} + func TestDispatch_RoutesStop(t *testing.T) { srv := newFakeServer(map[string]string{ "detect_changes": `{"changed_files":["a.go"],"changed_symbols":[{"id":"a.go::A","name":"A","kind":"function"}],"risk":"LOW","summary":"1"}`, diff --git a/internal/hooks/pretooluse.go b/internal/hooks/pretooluse.go index 249c142d..756f95a3 100644 --- a/internal/hooks/pretooluse.go +++ b/internal/hooks/pretooluse.go @@ -90,40 +90,11 @@ func runPreToolUse(data []byte, gortexPort int, mode Mode) { // Terminal enforcement is deliberately the first policy branch. It is a // local marker lookup, so it neither waits for the daemon nor gets bypassed // by permissive permission modes. A new user prompt clears the marker. + if enforceLocalizationTerminalPreToolUse(input, started) { + return + } terminalTurn, terminalTurnReady := currentLocalizationTurnState(input.SessionID, input.PromptID, input.AgentID, input.CWD) terminalIdentity := terminalTurn.Identity - if terminalTurnReady { - if marker, marked := localizationTerminalMarkerFor(terminalIdentity); marked { - reason := "" - switch { - case !marker.Advisory: - reason = localizationTerminalDenyReason - case localizationNavigationTool(input.ToolName): - reason = localizationAdvisoryDenyReason - case localizationRedirectedHostTool(input.ToolName): - // Left to the access policy this deny becomes "call a Gortex - // graph tool instead", and the branch above then refuses that - // call. Prescribing a step we will not honour spends the - // caller's turn and teaches it nothing, so answer here instead. - reason = localizationAdvisoryDenyReason - } - if reason != "" { - // Hand the answer back with the refusal. A bare "you are done" - // leaves the caller with nothing to act on, so it reaches for the - // next tool and the turn budget drains one denial at a time. - if answer := strings.TrimSpace(marker.FinalResponse); answer != "" { - reason += "\n\n" + answer - } - emitPreToolUse(HookOutput{HookSpecificOutput: &HookSpecificOutput{ - HookEventName: "PreToolUse", - PermissionDecision: "deny", - PermissionDecisionReason: reason, - }}) - localizationTerminalTelemetry("denied", true, started) - return - } - } - } localizationAuthToken := "" if terminalTurnReady { // Correlate the current turn with this exact tool invocation. The nonce is @@ -266,6 +237,46 @@ func runPreToolUse(data []byte, gortexPort int, mode Mode) { emitPreToolUse(output) } +// enforceLocalizationTerminalPreToolUse applies only the local terminal +// contract and reports whether it emitted a deny. Keeping this seam separate +// lets hosts with specialized per-tool behavior (notably Codex) enforce the +// same all-tool terminal policy before dispatching to those handlers. +func enforceLocalizationTerminalPreToolUse(input HookInput, started time.Time) bool { + terminalTurn, ready := currentLocalizationTurnState(input.SessionID, input.PromptID, input.AgentID, input.CWD) + if !ready { + return false + } + marker, marked := localizationTerminalMarkerFor(terminalTurn.Identity) + if !marked { + return false + } + + reason := "" + switch { + case !marker.Advisory: + reason = localizationTerminalDenyReason + case localizationNavigationTool(input.ToolName): + reason = localizationAdvisoryDenyReason + case localizationRedirectedHostTool(input.ToolName): + // Left to the access policy this deny becomes "call a Gortex graph + // tool instead", and the navigation branch then refuses that call. + reason = localizationAdvisoryDenyReason + } + if reason == "" { + return false + } + if answer := strings.TrimSpace(marker.FinalResponse); answer != "" { + reason += "\n\n" + answer + } + emitPreToolUse(HookOutput{HookSpecificOutput: &HookSpecificOutput{ + HookEventName: "PreToolUse", + PermissionDecision: "deny", + PermissionDecisionReason: reason, + }}) + localizationTerminalTelemetry("denied", true, started) + return true +} + func hookAlternationSegmentCount(input HookInput) int { var pattern string switch input.ToolName { diff --git a/internal/hooks/telemetry.go b/internal/hooks/telemetry.go index e55222fa..eaff2e01 100644 --- a/internal/hooks/telemetry.go +++ b/internal/hooks/telemetry.go @@ -85,19 +85,20 @@ func SetAgent(name string) { func ActiveAgent() string { return activeAgent } var hookEffectivenessEvents = map[string]bool{ - "PostToolUse": true, - "PreToolUse": true, - "SessionStart": true, - "UserPromptSubmit": true, - "PreCompact": true, - "PostCompact": true, - "Stop": true, - "SubagentStart": true, - "SubagentStop": true, - "LocalizationTerminal.observed": true, - "LocalizationTerminal.denied": true, - "LocalizationTerminal.cleared_prompt": true, - "LocalizationTerminal.cleared_session": true, + "PostToolUse": true, + "PreToolUse": true, + "SessionStart": true, + "UserPromptSubmit": true, + "PreCompact": true, + "PostCompact": true, + "Stop": true, + "SubagentStart": true, + "SubagentStop": true, + "LocalizationTerminal.observed": true, + "LocalizationTerminal.denied": true, + "LocalizationTerminal.claim_lock_release_failed": true, + "LocalizationTerminal.cleared_prompt": true, + "LocalizationTerminal.cleared_session": true, } // hookDecisionsPath returns the telemetry file path. Respects GORTEX_HOOK_LOG diff --git a/internal/localizationauth/evidence.go b/internal/localizationauth/evidence.go new file mode 100644 index 00000000..39f56429 --- /dev/null +++ b/internal/localizationauth/evidence.go @@ -0,0 +1,115 @@ +package localizationauth + +import ( + "encoding/json" + "reflect" + "strings" + "unicode" + "unicode/utf8" +) + +const ( + maxPrimaryIDs = 5 + maxEvidenceIDs = 16 + maxEvidenceIDSize = 1024 + maxEvidenceIDsBytes = 8 << 10 +) + +// NormalizeEvidenceIDs validates and clones the bounded evidence authority +// carried by one terminal receipt. Primary IDs must be a unique subset of the +// complete digest IDs. Empty slices preserve compatibility with older receipts. +func NormalizeEvidenceIDs(primaryIDs, evidenceIDs []string) ([]string, []string, bool) { + if len(primaryIDs) == 0 && len(evidenceIDs) == 0 { + return nil, nil, true + } + if len(primaryIDs) > maxPrimaryIDs || len(evidenceIDs) > maxEvidenceIDs { + return nil, nil, false + } + normalize := func(values []string) ([]string, map[string]struct{}, int, bool) { + out := make([]string, 0, len(values)) + seen := make(map[string]struct{}, len(values)) + total := 0 + for _, raw := range values { + id := strings.TrimSpace(raw) + if id == "" || len(id) > maxEvidenceIDSize || !validEvidenceIDText(id) { + return nil, nil, 0, false + } + if _, duplicate := seen[id]; duplicate { + return nil, nil, 0, false + } + total += len(id) + if total > maxEvidenceIDsBytes { + return nil, nil, 0, false + } + seen[id] = struct{}{} + out = append(out, id) + } + return out, seen, total, true + } + evidence, evidenceSet, evidenceBytes, ok := normalize(evidenceIDs) + if !ok { + return nil, nil, false + } + primary, _, primaryBytes, ok := normalize(primaryIDs) + if !ok || primaryBytes+evidenceBytes > maxEvidenceIDsBytes { + return nil, nil, false + } + for _, id := range primary { + if _, present := evidenceSet[id]; !present { + return nil, nil, false + } + } + return primary, evidence, true +} + +// validEvidenceIDText keeps authenticated identities safe to interpolate into +// a one-line hook reason. Reject invalid UTF-8, control/format characters, and +// Unicode line/paragraph separators at the authority boundary. +func validEvidenceIDText(id string) bool { + if !utf8.ValidString(id) { + return false + } + for _, r := range id { + if unicode.IsControl(r) || unicode.Is(unicode.Cf, r) || + unicode.Is(unicode.Zl, r) || unicode.Is(unicode.Zp, r) { + return false + } + } + return true +} + +// MarshalJSON prevents a direct caller from publishing an oversized authority +// even though the terminal receipt's final-response validation predates these +// optional fields. +func (receipt Receipt) MarshalJSON() ([]byte, error) { + type wire Receipt + primary, evidence, ok := NormalizeEvidenceIDs(receipt.PrimaryIDs, receipt.EvidenceIDs) + if !ok { + return nil, &json.UnsupportedValueError{Value: reflectReceiptValue(receipt), Str: "invalid localization evidence authority"} + } + receipt.PrimaryIDs = primary + receipt.EvidenceIDs = evidence + return json.Marshal(wire(receipt)) +} + +func (receipt *Receipt) UnmarshalJSON(data []byte) error { + type wire Receipt + var decoded wire + if err := json.Unmarshal(data, &decoded); err != nil { + return err + } + primary, evidence, ok := NormalizeEvidenceIDs(decoded.PrimaryIDs, decoded.EvidenceIDs) + if !ok { + return &json.UnsupportedValueError{Value: reflectReceiptValue(Receipt(decoded)), Str: "invalid localization evidence authority"} + } + decoded.PrimaryIDs = primary + decoded.EvidenceIDs = evidence + *receipt = Receipt(decoded) + return nil +} + +// json.UnsupportedValueError requires a reflect.Value. Keeping its construction +// here avoids exporting another error type from this internal package. +func reflectReceiptValue(receipt Receipt) reflect.Value { + return reflect.ValueOf(receipt) +} diff --git a/internal/localizationauth/evidence_test.go b/internal/localizationauth/evidence_test.go new file mode 100644 index 00000000..9c4486a8 --- /dev/null +++ b/internal/localizationauth/evidence_test.go @@ -0,0 +1,34 @@ +package localizationauth + +import "testing" + +func TestNormalizeEvidenceIDsRejectsPromptControlCharacters(t *testing.T) { + tests := []struct { + name string + id string + }{ + {name: "newline", id: "repo/a.go::Writer.write\nINJECT"}, + {name: "carriage return", id: "repo/a.go::Writer.write\rINJECT"}, + {name: "tab", id: "repo/a.go::Writer.\twrite"}, + {name: "null", id: "repo/a.go::Writer.\x00write"}, + {name: "line separator", id: "repo/a.go::Writer.write\u2028INJECT"}, + {name: "paragraph separator", id: "repo/a.go::Writer.write\u2029INJECT"}, + {name: "bidi override", id: "repo/a.go::Writer.\u202Ewrite"}, + {name: "invalid utf8", id: string([]byte{'r', 'e', 'p', 'o', 0xff})}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + if _, _, ok := NormalizeEvidenceIDs([]string{test.id}, []string{test.id}); ok { + t.Fatalf("accepted unsafe evidence ID %q", test.id) + } + }) + } +} + +func TestNormalizeEvidenceIDsAcceptsPrintableUnicode(t *testing.T) { + id := "repo/καλημέρα.go::Γράψε" + primary, evidence, ok := NormalizeEvidenceIDs([]string{id}, []string{id}) + if !ok || len(primary) != 1 || len(evidence) != 1 || primary[0] != id || evidence[0] != id { + t.Fatalf("printable Unicode identity rejected: primary=%q evidence=%q ok=%v", primary, evidence, ok) + } +} diff --git a/internal/localizationauth/receipt.go b/internal/localizationauth/receipt.go index cb954e47..224db706 100644 --- a/internal/localizationauth/receipt.go +++ b/internal/localizationauth/receipt.go @@ -35,9 +35,11 @@ const ( // answer_ready contract. PostToolUse treats this server-owned record, rather // than its visible tool_response, as the terminal authority. type Receipt struct { - FinalResponse string `json:"final_response"` - ContractVersion int `json:"contract_version"` - Enforceable bool `json:"enforceable"` + FinalResponse string `json:"final_response"` + PrimaryIDs []string `json:"primary_ids,omitempty"` + EvidenceIDs []string `json:"evidence_ids,omitempty"` + ContractVersion int `json:"contract_version"` + Enforceable bool `json:"enforceable"` } type receiptEnvelope struct { diff --git a/internal/localizationauth/receipt_test.go b/internal/localizationauth/receipt_test.go index b933829c..d8705177 100644 --- a/internal/localizationauth/receipt_test.go +++ b/internal/localizationauth/receipt_test.go @@ -3,6 +3,7 @@ package localizationauth import ( "encoding/json" "os" + "reflect" "sync" "sync/atomic" "testing" @@ -51,7 +52,7 @@ func TestReceiptRoundTripPreservesFinalResponseVerbatim(t *testing.T) { if !ok { t.Fatal("Consume failed") } - if got != want { + if !reflect.DeepEqual(got, want) { t.Fatalf("receipt mismatch\n got: %#v\nwant: %#v", got, want) } if _, ok := Consume(token); ok { diff --git a/internal/mcp/ast_target_file_index_test.go b/internal/mcp/ast_target_file_index_test.go new file mode 100644 index 00000000..0359a749 --- /dev/null +++ b/internal/mcp/ast_target_file_index_test.go @@ -0,0 +1,248 @@ +package mcp + +import ( + "context" + "fmt" + "testing" + + "github.com/zzet/gortex/internal/astquery" + "github.com/zzet/gortex/internal/graph" +) + +type astTargetBoundedStore struct { + graph.Store + calls []boundedFileCall + read func(context.Context, string, graph.LocalizationNodeScope, int) (graph.BoundedNodeProjection, error) +} + +func (store *astTargetBoundedStore) AllNodes() []*graph.Node { + panic("AST enclosing lookup must not call AllNodes") +} + +func (store *astTargetBoundedStore) FindFileNodesBounded( + ctx context.Context, + path string, + scope graph.LocalizationNodeScope, + limit int, +) (graph.BoundedNodeProjection, error) { + store.calls = append(store.calls, boundedFileCall{path: path, scope: scope, limit: limit}) + return store.read(ctx, path, scope, limit) +} + +type astTargetUnboundedStore struct{ graph.Store } + +func (store *astTargetUnboundedStore) AllNodes() []*graph.Node { + panic("missing bounded capability must fail closed without AllNodes") +} + +func TestEnrichASTMatchesPreservesOrderAndDeduplicates(t *testing.T) { + probe := &astTargetBoundedStore{Store: graph.New()} + probe.read = func(_ context.Context, _ string, _ graph.LocalizationNodeScope, _ int) (graph.BoundedNodeProjection, error) { + return graph.BoundedNodeProjection{}, nil + } + server := &Server{graph: probe} + matches := []astquery.Match{ + {File: "repo/b.ts", Line: 1}, + {File: "repo/a.go", Line: 1}, + {File: "repo/b.ts", Line: 2}, + {File: "", Line: 1}, + {File: "repo/c.py", Line: 1}, + } + + server.enrichASTMatchesContext(context.Background(), matches) + + want := []string{"repo/b.ts", "repo/a.go", "repo/c.py"} + if len(probe.calls) != len(want) { + t.Fatalf("bounded calls = %#v, want %v", probe.calls, want) + } + for index, path := range want { + if probe.calls[index].path != path || probe.calls[index].limit != localizationFileNodeLimit { + t.Fatalf("bounded call %d = %#v, want path=%q limit=%d", index, probe.calls[index], path, localizationFileNodeLimit) + } + } +} + +func TestEnrichASTMatchesSharesRequestBudget(t *testing.T) { + probe := &astTargetBoundedStore{Store: graph.New()} + probe.read = func(_ context.Context, path string, _ graph.LocalizationNodeScope, limit int) (graph.BoundedNodeProjection, error) { + return graph.BoundedNodeProjection{ + Nodes: []*graph.Node{{ + ID: path + "::owner", Name: "owner", Kind: graph.KindFunction, + FilePath: path, StartLine: 1, EndLine: 2, + }}, + Total: limit, + }, nil + } + server := &Server{graph: probe} + ctx := withLocalizationFileRequestBudget(context.Background()) + first := []astquery.Match{{File: "repo/a.go", Line: 1}, {File: "repo/b.go", Line: 1}, {File: "repo/c.go", Line: 1}} + second := []astquery.Match{{File: "repo/d.go", Line: 1}, {File: "repo/e.go", Line: 1}} + + server.enrichASTMatchesContext(ctx, first) + server.enrichASTMatchesContext(ctx, second) + + if len(probe.calls) != localizationFileRequestLimit/localizationFileNodeLimit { + t.Fatalf("bounded calls = %d, want shared request cap of %d", len(probe.calls), localizationFileRequestLimit/localizationFileNodeLimit) + } + if probe.calls[3].path != "repo/d.go" || second[0].SymbolID != "repo/d.go::owner" { + t.Fatalf("last budgeted match = %#v / %#v, want repo/d.go owner", probe.calls[3], second[0]) + } + if second[1].SymbolID != "" || second[1].SymbolName != "" { + t.Fatalf("budget-exhausted match = %#v, want blank fail-closed symbol", second[1]) + } +} + +func TestEnrichASTMatchesFailsClosedWithoutFullScan(t *testing.T) { + t.Run("saturated", func(t *testing.T) { + const path = "repo/dense.go" + probe := &astTargetBoundedStore{Store: graph.New()} + probe.read = func(_ context.Context, _ string, _ graph.LocalizationNodeScope, limit int) (graph.BoundedNodeProjection, error) { + return graph.BoundedNodeProjection{ + Nodes: []*graph.Node{{ + ID: path + "::wrong", Name: "wrong", Kind: graph.KindFunction, + FilePath: path, StartLine: 1, EndLine: 100, + }}, + Total: limit, Truncated: true, + }, nil + } + server := &Server{graph: probe} + matches := []astquery.Match{{File: path, Line: 50, SymbolID: "stale", SymbolName: "stale"}} + server.enrichASTMatchesContext(context.Background(), matches) + if matches[0].SymbolID != "" || matches[0].SymbolName != "" { + t.Fatalf("saturated match = %#v, want fail closed", matches[0]) + } + }) + + t.Run("missing capability", func(t *testing.T) { + const path = "repo/unsupported.go" + server := &Server{graph: &astTargetUnboundedStore{Store: graph.New()}} + matches := []astquery.Match{{File: path, Line: 1, SymbolID: "stale", SymbolName: "stale"}} + server.enrichASTMatchesContext(context.Background(), matches) + if matches[0].SymbolID != "" || matches[0].SymbolName != "" { + t.Fatalf("unsupported match = %#v, want fail closed", matches[0]) + } + }) +} + +func TestEnrichASTMatchesReadsBaseInsteadOfOverlay(t *testing.T) { + const path = "repo/service.go" + baseOwner := &graph.Node{ + ID: path + "::base", Name: "base", Kind: graph.KindFunction, + FilePath: path, StartLine: 1, EndLine: 20, + } + overlayOwner := &graph.Node{ + ID: path + "::overlay", Name: "overlay", Kind: graph.KindFunction, + FilePath: path, StartLine: 1, EndLine: 20, + } + probe := &astTargetBoundedStore{Store: graph.New()} + probe.read = func(_ context.Context, gotPath string, _ graph.LocalizationNodeScope, limit int) (graph.BoundedNodeProjection, error) { + if gotPath != path || limit != localizationFileNodeLimit { + return graph.BoundedNodeProjection{}, fmt.Errorf("unexpected base read %q limit %d", gotPath, limit) + } + return graph.BoundedNodeProjection{Nodes: []*graph.Node{baseOwner}, Total: 1}, nil + } + layer := graph.NewOverlayLayer() + layer.MarkFile(path, false) + layer.AddNode(path, overlayOwner) + ctx := WithOverlayView(context.Background(), graph.NewOverlaidView(probe, layer)) + server := &Server{graph: probe} + matches := []astquery.Match{{File: path, Line: 10}} + + server.enrichASTMatchesContext(ctx, matches) + + if matches[0].SymbolID != baseOwner.ID || matches[0].SymbolName != baseOwner.Name { + t.Fatalf("AST match owner = %#v, want durable base owner (%q, %q)", matches[0], baseOwner.ID, baseOwner.Name) + } + if len(probe.calls) != 1 { + t.Fatalf("base projection calls = %d, want one", len(probe.calls)) + } +} + +func TestEnrichASTMatchesCapsRefundedEmptyFileReads(t *testing.T) { + probe := &astTargetBoundedStore{Store: graph.New()} + probe.read = func(_ context.Context, _ string, _ graph.LocalizationNodeScope, _ int) (graph.BoundedNodeProjection, error) { + // Empty pages refund their entire node reservation. The independent + // file-call cap must still stop projection after 64 transactions. + return graph.BoundedNodeProjection{}, nil + } + server := &Server{graph: probe} + matches := make([]astquery.Match, astPostMatchFileLimit+2) + for index := range matches { + matches[index] = astquery.Match{File: fmt.Sprintf("repo/%03d.go", len(matches)-index), Line: 1} + } + + server.enrichASTMatchesContext(context.Background(), matches) + + if len(probe.calls) != astPostMatchFileLimit { + t.Fatalf("bounded calls = %d, want hard file cap %d despite refunds", len(probe.calls), astPostMatchFileLimit) + } + for index := range probe.calls { + if probe.calls[index].path != matches[index].File { + t.Fatalf("priority %d call = %#v, want path %q", index, probe.calls[index], matches[index].File) + } + } + for index := range matches { + if matches[index].SymbolID != "" || matches[index].SymbolName != "" { + t.Fatalf("empty-page match %d = %#v, want blank symbol", index, matches[index]) + } + } +} + +func TestPrepareReviewSASTMatchesAppliesPriorityBeforeLimit(t *testing.T) { + probe := &astTargetBoundedStore{Store: graph.New()} + probe.read = func(_ context.Context, path string, _ graph.LocalizationNodeScope, _ int) (graph.BoundedNodeProjection, error) { + return graph.BoundedNodeProjection{ + Nodes: []*graph.Node{{ + ID: path + "::owner", Name: "owner", Kind: graph.KindFunction, + FilePath: path, StartLine: 1, EndLine: 2, + }}, Total: 1, + }, nil + } + server := &Server{graph: probe} + matches := make([]astquery.Match, 0, astPostMatchFileLimit+2) + for index := 0; index <= astPostMatchFileLimit; index++ { + matches = append(matches, astquery.Match{ + Detector: "z-low", Severity: "low", + File: fmt.Sprintf("repo/%03d-low.go", index), Line: 1, + }) + } + const highPath = "repo/high-priority.go" + matches = append(matches, astquery.Match{ + Detector: "a-high", Severity: "critical", File: highPath, Line: 1, + }) + + server.prepareReviewSASTMatchesContext(context.Background(), matches, 1, true) + + if matches[0].File != highPath { + t.Fatalf("first stable match = %q, want high-priority %q", matches[0].File, highPath) + } + if len(probe.calls) != 1 || probe.calls[0].path != highPath { + t.Fatalf("bounded calls = %#v, want one high-priority projection", probe.calls) + } + if matches[0].SymbolID != highPath+"::owner" { + t.Fatalf("high-priority symbol = %q, want owner", matches[0].SymbolID) + } + for index := 1; index < len(matches); index++ { + if matches[index].SymbolID != "" || matches[index].SymbolName != "" { + t.Fatalf("over-limit match %d = %#v, want blank symbol", index, matches[index]) + } + } +} + +func TestPrepareReviewSASTMatchesKindsOnlySkipsProjection(t *testing.T) { + probe := &astTargetBoundedStore{Store: graph.New()} + probe.read = func(_ context.Context, _ string, _ graph.LocalizationNodeScope, _ int) (graph.BoundedNodeProjection, error) { + t.Fatal("kinds_only must not read enclosing-symbol projections") + return graph.BoundedNodeProjection{}, nil + } + server := &Server{graph: probe} + matches := []astquery.Match{{ + Detector: "review", Severity: "high", File: "repo/review.go", Line: 1, + }} + + server.prepareReviewSASTMatchesContext(context.Background(), matches, 1, false) + + if len(probe.calls) != 0 { + t.Fatalf("bounded calls = %#v, want none for kinds_only", probe.calls) + } +} diff --git a/internal/mcp/change_contract.go b/internal/mcp/change_contract.go index 8dcde9dd..abb2fea4 100644 --- a/internal/mcp/change_contract.go +++ b/internal/mcp/change_contract.go @@ -257,7 +257,14 @@ func (s *Server) lowerEditSource(ctx context.Context, req mcp.CallToolRequest) ( } // Edited ranges that did not add/remove a symbol still touch their // enclosing symbol — lower the edit's ranges so a body change counts. - for _, h := range s.lowerWorkspaceEditRanges(edit) { + loweredEdit := s.lowerWorkspaceEditRanges(ctx, edit) + if err := ctx.Err(); err != nil { + return nil, err + } + if len(loweredEdit.saturated) > 0 { + return nil, saturatedRangeLoweringError(loweredEdit.saturated) + } + for _, h := range loweredEdit.hits { ids = append(ids, h.ID) } ids = dedupeStrings(ids) @@ -322,10 +329,10 @@ func (s *Server) workspaceEditVerificationFiles(edit lsp.WorkspaceEdit) ([]verif } // lowerWorkspaceEditRanges maps each TextEdit's range to its enclosing symbols. -func (s *Server) lowerWorkspaceEditRanges(edit lsp.WorkspaceEdit) []rangeSymbolHit { +func (s *Server) lowerWorkspaceEditRanges(ctx context.Context, edit lsp.WorkspaceEdit) loweredRanges { fileEdits, err := s.groupEditByFile(edit) if err != nil { - return nil + return loweredRanges{} } var specs []rangeSpec for _, fe := range fileEdits { @@ -342,8 +349,13 @@ func (s *Server) lowerWorkspaceEditRanges(edit lsp.WorkspaceEdit) []rangeSymbolH }) } } - hits, _ := s.lowerRanges(specs) - return hits + return s.lowerRangesDetailedContext(ctx, specs) +} + +func saturatedRangeLoweringError(paths []string) error { + paths = append([]string(nil), paths...) + sort.Strings(paths) + return fmt.Errorf("bounded range ownership is unavailable for: %s", strings.Join(dedupeStrings(paths), ", ")) } func (s *Server) lowerRangeSource(ctx context.Context, req mcp.CallToolRequest) (*prediction, error) { @@ -351,7 +363,14 @@ func (s *Server) lowerRangeSource(ctx context.Context, req mcp.CallToolRequest) if err != nil { return nil, err } - hits, _ := s.lowerRanges(specs) + lowered := s.lowerRangesDetailedContext(ctx, specs) + if err := ctx.Err(); err != nil { + return nil, err + } + if len(lowered.saturated) > 0 { + return nil, saturatedRangeLoweringError(lowered.saturated) + } + hits := lowered.hits ids := make([]string, 0, len(hits)) changed := make([]changedSymbolRef, 0, len(hits)) files := make([]string, 0, len(hits)) @@ -875,6 +894,9 @@ func (s *Server) handleChangeContract(ctx context.Context, req mcp.CallToolReque } p, err := s.lowerChange(ctx, req) if err != nil { + if ctxErr := requestContextError(ctx, err); ctxErr != nil { + return nil, ctxErr + } return mcp.NewToolResultError(err.Error()), nil } if req.GetBool("ack", false) { diff --git a/internal/mcp/combo_apply.go b/internal/mcp/combo_apply.go index 3dccc3e1..5f1d9be9 100644 --- a/internal/mcp/combo_apply.go +++ b/internal/mcp/combo_apply.go @@ -69,19 +69,12 @@ func applyRerankBoostsTimed(s *Server, nodes []*graph.Node, query string, rerank return result, prepare, signals } -// recordLastSearchFromNodes stores the query + top-limit IDs on the session -// so a subsequent get_symbol_source / get_editing_context can credit this -// search. Capped at limit to avoid crediting results the agent never saw. -func recordLastSearchFromNodes(sess *sessionState, query string, nodes []*graph.Node, limit int) { - if sess == nil || len(nodes) == 0 { +// recordLastSearchFromNodes stores exactly one returned search page so later +// symbol/file consumption can be attributed without graph reads. Empty pages +// intentionally clear the prior page. +func recordLastSearchFromNodes(sess *sessionState, query string, nodes []*graph.Node) { + if sess == nil { return } - if limit <= 0 || limit > len(nodes) { - limit = len(nodes) - } - ids := make([]string, 0, limit) - for i := 0; i < limit; i++ { - ids = append(ids, nodes[i].ID) - } - sess.recordLastSearch(query, ids) + sess.recordLastSearchPage(query, nodes) } diff --git a/internal/mcp/enclosing.go b/internal/mcp/enclosing.go index e1996c26..ed5af403 100644 --- a/internal/mcp/enclosing.go +++ b/internal/mcp/enclosing.go @@ -6,6 +6,7 @@ import ( "github.com/zzet/gortex/internal/astquery" "github.com/zzet/gortex/internal/graph" + "github.com/zzet/gortex/internal/query" ) // This file owns enclosing-scope resolution shared across the search @@ -15,44 +16,82 @@ import ( // to answer "which symbol contains this?" -- they share this code so // the answer stays consistent. -// buildFileSymbolIndex returns one fileSymbolIndex per Target's graph -// path. Building all indexes up-front (instead of lazily on first -// match) is fine because the cost is one graph pass per file's -// symbol list, and the alternative — locking inside the worker pool -// hot path — is worse for parallel runs. -func (s *Server) buildFileSymbolIndex(targets []astquery.Target) map[string]*fileSymbolIndex { - if s.graph == nil { +const astPostMatchFileLimit = 64 + +// astPostMatchSymbolLookupContext builds enclosing-symbol indexes only for the +// first 64 distinct files that survived a caller's stable ordering and result +// limit. AST targets come from the durable graph, so enrichment deliberately +// reads s.graph even when the request carries an editor overlay. +func (s *Server) astPostMatchSymbolLookupContext( + ctx context.Context, + count int, + pathAt func(int) string, +) astquery.SymbolLookup { + if s == nil || s.graph == nil || count <= 0 || pathAt == nil || ctx.Err() != nil { return nil } - wanted := make(map[string]struct{}, len(targets)) - for _, t := range targets { - wanted[t.GraphPath] = struct{}{} - } - out := make(map[string]*fileSymbolIndex, len(wanted)) - for _, n := range s.graph.AllNodes() { - if _, ok := wanted[n.FilePath]; !ok { + paths := make([]string, 0, astPostMatchFileLimit) + admitted := make(map[string]struct{}, astPostMatchFileLimit) + for index := 0; index < count && len(paths) < astPostMatchFileLimit; index++ { + path := pathAt(index) + if path == "" { continue } - // Functions, methods, closures, and macros are meaningful - // "enclosing scope" candidates. Macro nodes are declaration-backed and - // already carry exact source ranges; token-tree contents are never parsed - // as declarations here. KindType (struct/class) is included too so a - // class-body match still gets a symbol identity. - switch n.Kind { - case graph.KindFunction, graph.KindMethod, graph.KindClosure, graph.KindMacro, - graph.KindType, graph.KindInterface: - idx := out[n.FilePath] - if idx == nil { - idx = &fileSymbolIndex{} - out[n.FilePath] = idx - } - idx.add(n) + if _, duplicate := admitted[path]; duplicate { + continue } + admitted[path] = struct{}{} + paths = append(paths, path) } - for _, idx := range out { - idx.finalise() + if len(paths) == 0 { + return nil + } + indexes := s.buildFileSymbolIndexForOrderedPathsScopedReaderContext( + ctx, s.graph, paths, query.QueryOptions{}, + ) + return func(path string, line int) (string, string) { + if _, ok := admitted[path]; !ok { + return "", "" + } + idx := indexes[path] + if idx == nil || idx.saturated { + return "", "" + } + return idx.find(line) + } +} + +func (s *Server) enrichASTMatchesContext(ctx context.Context, matches []astquery.Match) { + lookup := s.astPostMatchSymbolLookupContext(ctx, len(matches), func(index int) string { + return matches[index].File + }) + for index := range matches { + matches[index].SymbolID = "" + matches[index].SymbolName = "" + if lookup != nil { + matches[index].SymbolID, matches[index].SymbolName = lookup(matches[index].File, matches[index].Line) + } + } +} + +func (s *Server) enrichASTSymbolIDsContext( + ctx context.Context, + count int, + pathAt func(int) string, + lineAt func(int) int, + set func(int, string), +) { + if pathAt == nil || lineAt == nil || set == nil { + return + } + lookup := s.astPostMatchSymbolLookupContext(ctx, count, pathAt) + for index := 0; index < count; index++ { + id := "" + if lookup != nil { + id, _ = lookup(pathAt(index), lineAt(index)) + } + set(index, id) } - return out } // fileSymbolIndex is the per-file lookup used by the SymbolLookup @@ -60,7 +99,9 @@ func (s *Server) buildFileSymbolIndex(targets []astquery.Target) map[string]*fil // width so `find` returns the deepest enclosing scope (a closure // inside a method beats the method itself). type fileSymbolIndex struct { - syms []*graph.Node + syms []*graph.Node + fileNode *graph.Node + saturated bool } func (i *fileSymbolIndex) add(n *graph.Node) { i.syms = append(i.syms, n) } @@ -87,7 +128,7 @@ func (i *fileSymbolIndex) finalise() { // 1-based; graph nodes store the same convention. syms is sorted by // StartLine ascending, so the scan can stop once StartLine passes line. func (i *fileSymbolIndex) smallestEnclosing(line int) *graph.Node { - if i == nil { + if i == nil || i.saturated { return nil } var best *graph.Node @@ -203,65 +244,145 @@ func enclosingName(n *graph.Node, g graph.Reader) (id, name string) { return "", "" } -// buildFileSymbolIndexForPaths builds one fileSymbolIndex per file -// path in `paths`. It is the plain-path sibling of -// buildFileSymbolIndex (which keys off astquery.Target values) -- -// search_text works from trigram match paths, not AST targets, and -// needs the same enclosing-scope lookup. +const ( + localizationFileNodeLimit = 1_024 + localizationFileRequestLimit = 4_096 +) + +var localizationFileIndexKinds = []graph.NodeKind{ + graph.KindFile, graph.KindFunction, graph.KindMethod, graph.KindClosure, + graph.KindMacro, graph.KindType, graph.KindInterface, +} + +// buildFileSymbolIndexForPaths builds one bounded fileSymbolIndex per file +// path. Compatibility callers without a request scope still use the same typed +// bounded projection; there is deliberately no GetFileNodes fallback. func (s *Server) buildFileSymbolIndexForPaths(paths map[string]struct{}) map[string]*fileSymbolIndex { return s.buildFileSymbolIndexForPathsContext(context.Background(), paths) } -type contextFileNodeReader interface { - GetFileNodesContext(context.Context, string) []*graph.Node +func (s *Server) buildFileSymbolIndexForPathsContext(ctx context.Context, paths map[string]struct{}) map[string]*fileSymbolIndex { + return s.buildFileSymbolIndexForPathsScopedContext(ctx, paths, query.QueryOptions{}) } -func (s *Server) buildFileSymbolIndexForPathsContext(ctx context.Context, paths map[string]struct{}) map[string]*fileSymbolIndex { +func (s *Server) buildFileSymbolIndexForPathsScopedContext( + ctx context.Context, + paths map[string]struct{}, + opts query.QueryOptions, +) map[string]*fileSymbolIndex { ordered := make([]string, 0, len(paths)) for path := range paths { ordered = append(ordered, path) } sort.Strings(ordered) - return s.buildFileSymbolIndexForOrderedPathsContext(ctx, ordered) + return s.buildFileSymbolIndexForOrderedPathsScopedContext(ctx, ordered, opts) } -// buildFileSymbolIndexForOrderedPathsContext preserves caller priority while -// keeping file-node lookup bounded by ctx. Source-literal mapping uses this to -// query authoritative match paths before compatibility aliases. -func (s *Server) buildFileSymbolIndexForOrderedPathsContext(ctx context.Context, paths []string) map[string]*fileSymbolIndex { - if s.graph == nil || len(paths) == 0 || ctx.Err() != nil { +// buildFileSymbolIndexForOrderedPathsScopedContext preserves caller priority, +// applies request/session scope before each storage cap, and shares one strict +// node budget across the request. Saturated and unavailable paths retain an +// explicit marker so an exact lookup cannot fall through to a compatibility +// alias and misattribute an omitted narrower declaration. +func (s *Server) buildFileSymbolIndexForOrderedPathsScopedContext( + ctx context.Context, + paths []string, + opts query.QueryOptions, +) map[string]*fileSymbolIndex { + return s.buildFileSymbolIndexForOrderedPathsScopedReaderContext( + ctx, s.readerFor(ctx), paths, opts, + ) +} + +func (s *Server) buildFileSymbolIndexForOrderedPathsScopedReaderContext( + ctx context.Context, + reader graph.Reader, + paths []string, + opts query.QueryOptions, +) map[string]*fileSymbolIndex { + if len(paths) == 0 { return nil } + bounded, ok := reader.(graph.BoundedFileNodeReader) + if reader == nil || !ok || ctx.Err() != nil { + return saturatedFileSymbolIndexes(paths) + } + + scope := s.localizationNodeScopeWithTests(ctx, opts, false, localizationFileIndexKinds...) out := make(map[string]*fileSymbolIndex, len(paths)) - contextReader, hasContextReader := s.graph.(contextFileNodeReader) + budget := localizationFileBudgetFor(ctx) for _, path := range paths { - if ctx.Err() != nil { - break + if path == "" { + continue } - var nodes []*graph.Node - if hasContextReader { - nodes = contextReader.GetFileNodesContext(ctx, path) - } else { - nodes = s.graph.GetFileNodes(path) + if _, duplicate := out[path]; duplicate { + continue } if ctx.Err() != nil { - break + return saturateMissingFileSymbolIndexes(out, paths) + } + limit := budget.reserve(localizationFileNodeLimit) + if limit <= 0 { + out[path] = &fileSymbolIndex{saturated: true} + continue + } + page, err := bounded.FindFileNodesBounded(ctx, path, scope, limit) + if err != nil || ctx.Err() != nil { + return saturateMissingFileSymbolIndexes(out, paths) } - for _, n := range nodes { - switch n.Kind { - case graph.KindFunction, graph.KindMethod, graph.KindClosure, graph.KindMacro, - graph.KindType, graph.KindInterface: - idx := out[n.FilePath] - if idx == nil { - idx = &fileSymbolIndex{} - out[n.FilePath] = idx + consumed := page.Total + if len(page.Nodes) > consumed { + consumed = len(page.Nodes) + } + budget.finish(limit, consumed) + if page.Truncated { + out[path] = &fileSymbolIndex{saturated: true} + continue + } + + idx := &fileSymbolIndex{} + for _, node := range page.Nodes { + if node == nil { + continue + } + if node.Kind == graph.KindFile { + if idx.fileNode == nil || node.ID < idx.fileNode.ID { + idx.fileNode = node } - idx.add(n) + continue } + idx.add(node) + } + if len(idx.syms) == 0 && idx.fileNode == nil { + continue } - } - for _, idx := range out { idx.finalise() + out[path] = idx + for _, node := range page.Nodes { + if node != nil && node.FilePath != "" { + if _, exists := out[node.FilePath]; !exists { + out[node.FilePath] = idx + } + } + } + } + return out +} + +func saturatedFileSymbolIndexes(paths []string) map[string]*fileSymbolIndex { + return saturateMissingFileSymbolIndexes(make(map[string]*fileSymbolIndex, len(paths)), paths) +} + +func saturateMissingFileSymbolIndexes( + out map[string]*fileSymbolIndex, + paths []string, +) map[string]*fileSymbolIndex { + for _, path := range paths { + if path == "" { + continue + } + if _, complete := out[path]; !complete { + out[path] = &fileSymbolIndex{saturated: true} + } } return out } diff --git a/internal/mcp/explore_artifact_intent.go b/internal/mcp/explore_artifact_intent.go index 367fe233..bd801b84 100644 --- a/internal/mcp/explore_artifact_intent.go +++ b/internal/mcp/explore_artifact_intent.go @@ -389,7 +389,8 @@ func (s *Server) gatherExploreArtifactLane(ctx context.Context, intent exploreAr } else { matches = s.indexer.GrepText(probe, exploreArtifactTextHitLimit) } - for _, match := range s.enrichTextMatches(matches) { + enriched, _ := s.enrichTextMatchesContext(ctx, matches, scope) + for _, match := range enriched { hit := byPath[strings.ToLower(strings.ReplaceAll(match.Path, "\\", "/"))] if hit == nil || !exploreArtifactFile(hit.path) { continue diff --git a/internal/mcp/explore_bare_literal.go b/internal/mcp/explore_bare_literal.go new file mode 100644 index 00000000..0dc890ed --- /dev/null +++ b/internal/mcp/explore_bare_literal.go @@ -0,0 +1,427 @@ +package mcp + +import ( + "regexp" + "sort" + "strings" + "unicode" + "unicode/utf8" + + "github.com/zzet/gortex/internal/search/rerank" +) + +const ( + exploreBareLiteralMaxTaskBytes = 16 << 10 + exploreBareLiteralMaxTerms = 5 + exploreBareLiteralMinRunes = 4 + exploreBareLiteralMaxRunes = 128 + exploreBareDiagnosticMinRunes = 12 + exploreBareDiagnosticMinWords = 3 + exploreBareDiagnosticMaxWords = 6 + + exploreDistinctiveBareTokenMaxTerms = 5 + exploreDistinctiveBareTokenMinRunes = 4 + exploreDistinctiveBareTokenSegmentMinRunes = 6 + exploreDistinctiveBareTokenScanCap = 128 +) + +var exploreBareLiteralWordRE = regexp.MustCompile(`[\pL\pN_][\pL\pN_'-]*`) + +var ( + exploreDistinctiveBareTokenRE = regexp.MustCompile(`[A-Za-z_][A-Za-z0-9_.]{3,}`) + exploreCamelTransitionRE = regexp.MustCompile(`[a-z][A-Z]`) +) + +// exploreDistinctiveBareTokens mines the identifier-shaped tokens a caller +// would grep for verbatim: camel-case, underscore, dotted, or digit-bearing +// tokens, plus long plain words. Issue prose names its subject far more often +// than it quotes it, and these tokens feed the same bounded source-literal +// recall as quoted terms — retrieval only, never authored claims. A dotted +// token also contributes its final segment, since a member cited as +// `owner.method` is greppable only by the method name. +func exploreDistinctiveBareTokens(task, repoPrefix string) []string { + task = exploreBareLiteralBoundedTask(task) + if strings.TrimSpace(task) == "" { + return nil + } + repoBase := strings.ToLower(strings.TrimSpace(repoPrefix)) + if cut := strings.LastIndexByte(repoBase, '-'); cut > 0 { + if digits := repoBase[cut+1:]; digits != "" && strings.TrimLeft(digits, "0123456789") == "" { + repoBase = repoBase[:cut] + } + } + type minedToken struct { + term string + structured bool + } + mined := make([]minedToken, 0, exploreDistinctiveBareTokenScanCap) + seen := make(map[string]struct{}, exploreDistinctiveBareTokenScanCap) + appendToken := func(term string, structured bool) { + term = strings.Trim(term, "-_.:") + if utf8.RuneCountInString(term) < exploreDistinctiveBareTokenMinRunes { + return + } + key := strings.ToLower(term) + if key == repoBase { + return + } + if _, stop := assistStopWords[key]; stop { + return + } + if _, generic := exploreTerminalGenericTerms[key]; generic { + return + } + if _, duplicate := seen[key]; duplicate { + return + } + seen[key] = struct{}{} + mined = append(mined, minedToken{term: term, structured: structured}) + } + for _, raw := range exploreDistinctiveBareTokenRE.FindAllString(task, exploreDistinctiveBareTokenScanCap) { + raw = strings.TrimRight(raw, ".") + // Only structurally identifier-shaped tokens qualify. Long plain words + // ("connection", "completion") reach gold files too, but they are hub + // vocabulary: their grep pages map to dozens of owners and the rows + // they admit displace real evidence under the bounded page. + structured := exploreCamelTransitionRE.MatchString(raw) || + strings.Contains(raw, "_") || strings.Contains(raw, ".") || + strings.ContainsAny(raw, "0123456789") + if !structured { + continue + } + appendToken(raw, structured) + if dot := strings.LastIndexByte(raw, '.'); dot >= 0 { + segment := raw[dot+1:] + if utf8.RuneCountInString(segment) >= exploreDistinctiveBareTokenSegmentMinRunes { + appendToken(segment, exploreCamelTransitionRE.MatchString(segment) || strings.Contains(segment, "_")) + } + } + } + sort.SliceStable(mined, func(i, j int) bool { + if mined[i].structured != mined[j].structured { + return mined[i].structured + } + return utf8.RuneCountInString(mined[i].term) > utf8.RuneCountInString(mined[j].term) + }) + out := make([]string, 0, exploreDistinctiveBareTokenMaxTerms) + for _, token := range mined { + if len(out) == exploreDistinctiveBareTokenMaxTerms { + break + } + out = append(out, token.term) + } + return out +} + +// exploreBareLiteralRecallTerms mines only high-signal source values the +// requester wrote without quotes. It is intentionally separate from quoted +// recall: these inferred terms aid retrieval but never become authored claims +// that can hold localization open. +func exploreBareLiteralRecallTerms(task string) []string { + task = exploreBareLiteralBoundedTask(task) + if strings.TrimSpace(task) == "" { + return nil + } + + lines := exploreBareLiteralProseLines(task) + anchors := exploreSyntacticAnchors(task) + diagnostics := make([]string, 0, exploreBareLiteralMaxTerms) + structured := make([]string, 0, exploreBareLiteralMaxTerms) + plain := make([]string, 0, exploreBareLiteralMaxTerms) + for _, line := range lines { + if phrase := exploreBareDiagnosticPhrase(line); phrase != "" { + diagnostics = append(diagnostics, phrase) + } + for _, raw := range exploreUnquotedCodeTokens(line) { + term := strings.Trim(raw, "-_.:()[]{}<>,;'\"") + if !exploreBareAllCapsTerm(term) || exploreBareLiteralRepresentedByAnchor(term, anchors) { + continue + } + if strings.Contains(term, "-") { + if !exploreBareStructuredCue(line) { + continue + } + structured = append(structured, term) + continue + } + if strings.Contains(term, "_") { + structured = append(structured, term) + continue + } + plain = append(plain, term) + } + } + + seen := make(map[string]struct{}, exploreBareLiteralMaxTerms) + out := make([]string, 0, exploreBareLiteralMaxTerms) + for _, lane := range [3][]string{diagnostics, structured, plain} { + for _, term := range lane { + term = strings.TrimSpace(term) + key := strings.ToLower(term) + if term == "" { + continue + } + if _, duplicate := seen[key]; duplicate { + continue + } + seen[key] = struct{}{} + out = append(out, term) + if len(out) == exploreBareLiteralMaxTerms { + return out + } + } + } + return out +} + +func exploreBareLiteralBoundedTask(task string) string { + if len(task) <= exploreBareLiteralMaxTaskBytes { + return task + } + task = task[:exploreBareLiteralMaxTaskBytes] + for task != "" && !utf8.ValidString(task) { + task = task[:len(task)-1] + } + return task +} + +// exploreBareLiteralProseLines drops code blocks and masks explicit literals. +// Those inputs already have their own bounded lanes and must not consume this +// lane a second time. +func exploreBareLiteralProseLines(task string) []string { + lines := make([]string, 0, strings.Count(task, "\n")+1) + fenced := false + for _, raw := range strings.Split(task, "\n") { + trimmed := strings.TrimSpace(strings.TrimRight(raw, "\r")) + if strings.HasPrefix(trimmed, "```") || strings.HasPrefix(trimmed, "~~~") { + fenced = !fenced + continue + } + if fenced || exploreIndentedCodeLine(raw) { + continue + } + if line := strings.TrimSpace(exploreBareMaskQuotedSpans(raw)); line != "" { + lines = append(lines, line) + } + } + return lines +} + +func exploreBareMaskQuotedSpans(text string) string { + var masked strings.Builder + masked.Grow(len(text)) + var quote rune + escaped := false + for _, r := range text { + if quote != 0 { + if escaped { + escaped = false + masked.WriteRune(' ') + continue + } + if r == '\\' && quote == '"' { + escaped = true + masked.WriteRune(' ') + continue + } + if r == quote { + quote = 0 + } + masked.WriteRune(' ') + continue + } + if r == '"' || r == '`' { + quote = r + masked.WriteRune(' ') + continue + } + masked.WriteRune(r) + } + return masked.String() +} + +func exploreBareDiagnosticPhrase(line string) string { + lower := strings.ToLower(line) + start := -1 + for _, cue := range []string{ + "fails with", "error:", "message:", "reports", "returns", "raises", "throws", "emits", "prints", "says", "logs", + } { + if at := exploreBareCueIndex(lower, cue); at >= 0 { + candidateStart := at + len(cue) + if start < 0 || candidateStart < start { + start = candidateStart + } + } + } + if start < 0 || start >= len(line) { + return "" + } + tail := strings.TrimLeft(line[start:], " \t:-—") + if tail == "" || strings.Contains(strings.ToLower(tail), "://") || strings.Contains(strings.ToLower(tail), "www.") { + return "" + } + if cut := strings.IndexAny(tail, ".;!?"); cut >= 0 { + tail = tail[:cut] + } + wordIndexes := exploreBareLiteralWordRE.FindAllStringIndex(tail, exploreBareDiagnosticMaxWords+1) + if len(wordIndexes) < exploreBareDiagnosticMinWords { + return "" + } + endWord := len(wordIndexes) + if endWord > exploreBareDiagnosticMaxWords { + endWord = exploreBareDiagnosticMaxWords + } + for index := 2; index < endWord; index++ { + word := strings.ToLower(tail[wordIndexes[index][0]:wordIndexes[index][1]]) + switch word { + case "after", "before", "because", "when", "while": + endWord = index + } + } + if endWord < exploreBareDiagnosticMinWords { + return "" + } + phrase := strings.TrimSpace(tail[:wordIndexes[endWord-1][1]]) + runes := utf8.RuneCountInString(phrase) + if runes < exploreBareDiagnosticMinRunes || runes > exploreBareLiteralMaxRunes { + return "" + } + meaningful := 0 + for _, index := range wordIndexes[:endWord] { + word := strings.ToLower(tail[index[0]:index[1]]) + if _, stop := assistStopWords[word]; stop { + continue + } + if _, generic := exploreTerminalGenericTerms[exploreTerminalTermRoot(word)]; generic { + continue + } + meaningful++ + } + if meaningful < 2 { + return "" + } + return phrase +} + +func exploreBareCueIndex(lower, cue string) int { + for offset := 0; offset < len(lower); { + at := strings.Index(lower[offset:], cue) + if at < 0 { + return -1 + } + at += offset + beforeOK := at == 0 || !exploreBareWordByte(lower[at-1]) + after := at + len(cue) + afterOK := strings.HasSuffix(cue, ":") || after == len(lower) || !exploreBareWordByte(lower[after]) + if beforeOK && afterOK { + return at + } + offset = at + 1 + } + return -1 +} + +func exploreBareWordByte(b byte) bool { + return b == '_' || b >= 'a' && b <= 'z' || b >= 'A' && b <= 'Z' || b >= '0' && b <= '9' +} + +func exploreBareAllCapsTerm(term string) bool { + runes := utf8.RuneCountInString(term) + if runes < exploreBareLiteralMinRunes || runes > exploreBareLiteralMaxRunes || + strings.Contains(term, "://") || strings.HasPrefix(strings.ToLower(term), "www.") { + return false + } + letters := 0 + hex := true + for _, r := range term { + switch { + case unicode.IsLetter(r): + letters++ + if unicode.IsLower(r) { + return false + } + if !strings.ContainsRune("ABCDEF", r) { + hex = false + } + case unicode.IsDigit(r): + case r == '_' || r == '-': + hex = false + default: + return false + } + } + if letters < 2 || hex && runes >= 7 { + return false + } + compact := strings.ToLower(strings.NewReplacer("_", "", "-", "").Replace(term)) + if exploreSyntacticAnchorNoise(compact) { + return false + } + lower := strings.ToLower(term) + if _, stop := assistStopWords[lower]; stop { + return false + } + if _, generic := exploreTerminalGenericTerms[exploreTerminalTermRoot(lower)]; generic { + return false + } + return true +} + +func exploreBareStructuredCue(line string) bool { + lower := strings.ToLower(line) + for _, cue := range []string{"config", "configuration", "flag", "header", "key", "option", "property", "value"} { + if exploreBareCueIndex(lower, cue) >= 0 { + return true + } + } + return false +} + +func exploreBareLiteralRepresentedByAnchor(term string, anchors []exploreSyntacticAnchor) bool { + candidate, ok := newExploreSyntacticAnchor(term) + if !ok { + return false + } + for _, anchor := range anchors { + if exploreSyntacticAnchorEquivalent(anchor, candidate) { + return true + } + } + return false +} + +func exploreHasExplicitCandidateTarget(query string, candidates []*rerank.Candidate) bool { + if _, hasPath := exploreQueryPathAnchors(query); hasPath { + return true + } + for _, candidate := range candidates { + if candidate != nil && candidate.Node != nil && + exploreLocalizationExplicitAnchor(query, candidate.Node) { + return true + } + } + return false +} + +// exploreLiteralEvidenceEligible is the shared admission and seating gate for +// inferred source literals. Keeping both decisions behind one predicate makes +// anchored tasks a strict no-op and prevents a retrieved bare term from later +// taking a PRIMARY seat under different rules. +func exploreLiteralEvidenceEligible( + query string, + candidates []*rerank.Candidate, + protectedSyntacticAnchors map[int]string, +) bool { + return len(protectedSyntacticAnchors) == 0 && + !exploreHasExplicitCandidateTarget(query, candidates) +} + +func exploreBareLiteralLaneEligible( + task, query string, + conceptTask, artifactReady bool, + candidates []*rerank.Candidate, + protectedSyntacticAnchors map[int]string, +) bool { + return conceptTask && !artifactReady && len(exploreQuotedRecallTerms(task)) == 0 && + exploreLiteralEvidenceEligible(query, candidates, protectedSyntacticAnchors) +} diff --git a/internal/mcp/explore_bare_literal_test.go b/internal/mcp/explore_bare_literal_test.go new file mode 100644 index 00000000..8f50e8b8 --- /dev/null +++ b/internal/mcp/explore_bare_literal_test.go @@ -0,0 +1,97 @@ +package mcp + +import ( + "strings" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/zzet/gortex/internal/graph" + "github.com/zzet/gortex/internal/search/rerank" +) + +func TestExploreBareLiteralRecallTermsKeepsDiagnosticTail(t *testing.T) { + terms := exploreBareLiteralRecallTerms("The daemon reports unable to seal chunk manifest after the retry starts.") + require.Equal(t, []string{"unable to seal chunk manifest"}, terms) +} + +func TestExploreBareLiteralRecallTermsRequiresThreeWordDiagnosticTail(t *testing.T) { + require.Empty(t, exploreBareLiteralRecallTerms("The daemon reports manifest missing.")) + require.Equal(t, []string{"durable manifest missing"}, + exploreBareLiteralRecallTerms("The daemon reports durable manifest missing.")) +} + +func TestExploreBareLiteralRecallTermsAdmitsAllCapsAndCuedStructuredValues(t *testing.T) { + terms := exploreBareLiteralRecallTerms("Run ZRANGEBYSCORE again with the header X-CACHE-MODE value.") + require.Equal(t, []string{"X-CACHE-MODE", "ZRANGEBYSCORE"}, terms) +} + +func TestExploreBareLiteralRecallTermsExcludesExplicitAndFencedLiterals(t *testing.T) { + task := "The daemon logs \"UNABLE TO SEAL CHUNK MANIFEST\" for `CACHE_MISS`.\n\n" + + "```text\nERROR: DURABLE SNAPSHOT MISSING\n```\n" + + " TRACE BUFFER OVERFLOW\n" + require.Empty(t, exploreBareLiteralRecallTerms(task)) +} + +func TestExploreBareLiteralRecallTermsRejectsNoise(t *testing.T) { + task := "The service reports https://EXAMPLE.COM before retrying version V1-2-3 at ABCDEF1234. VERSION remains unchanged." + require.Empty(t, exploreBareLiteralRecallTerms(task)) +} + +func TestExploreBareLiteralRecallTermsDoesNotDuplicateSyntacticLane(t *testing.T) { + terms := exploreBareLiteralRecallTerms("Retry CACHE_KEY through ZRANGEBYSCORE when the command stalls.") + require.Equal(t, []string{"ZRANGEBYSCORE"}, terms) +} + +func TestExploreBareLiteralRecallTermsDeduplicatesAndCaps(t *testing.T) { + terms := exploreBareLiteralRecallTerms("ALPHA alpha BRAVO CHARLIE DELTA ECHO FOXTROT GOLF ALPHA") + require.Equal(t, []string{"ALPHA", "BRAVO", "CHARLIE", "DELTA", "ECHO"}, terms) +} + +func TestExploreBareLiteralRecallTermsBoundsTaskScan(t *testing.T) { + task := strings.Repeat("ordinary prose ", exploreBareLiteralMaxTaskBytes/8) + " ZRANGEBYSCORE" + require.Empty(t, exploreBareLiteralRecallTerms(task)) +} + +func TestExploreBareLiteralRecallTermsDoesNotCreateQuotedClaims(t *testing.T) { + task := "The daemon reports unable to seal chunk manifest before retrying ZRANGEBYSCORE." + require.NotEmpty(t, exploreBareLiteralRecallTerms(task)) + require.Empty(t, exploreQuotedRecallTerms(task)) + require.Empty(t, exploreQuotedRecallClaimTerms(task)) +} + +func TestExploreBareLiteralLaneEligibleOnlyForUnanchoredConcepts(t *testing.T) { + unanchoredTask := "Recovery loses the registry entry during rollback" + require.True(t, exploreBareLiteralLaneEligible( + unanchoredTask, unanchoredTask, true, false, nil, nil, + )) + + explicit := []*rerank.Candidate{{Node: &graph.Node{ + ID: "demo/registry.go::FlushRegistry", Name: "FlushRegistry", + Kind: graph.KindFunction, FilePath: "demo/registry.go", + }}} + cases := []struct { + name string + task string + query string + concept bool + artifactReady bool + candidates []*rerank.Candidate + protected map[int]string + }{ + {name: "quoted", task: `recovery logs "registry entry missing"`, query: unanchoredTask, concept: true}, + {name: "path", task: unanchoredTask, query: "fix src/cache.go during recovery", concept: true}, + {name: "ranked explicit call", task: unanchoredTask, query: "FlushRegistry() loses entries", concept: true, candidates: explicit}, + {name: "protected syntax", task: unanchoredTask, query: unanchoredTask, concept: true, protected: map[int]string{0: "node"}}, + {name: "artifact ready", task: unanchoredTask, query: unanchoredTask, concept: true, artifactReady: true}, + {name: "non concept", task: unanchoredTask, query: unanchoredTask}, + } + for _, test := range cases { + t.Run(test.name, func(t *testing.T) { + require.False(t, exploreBareLiteralLaneEligible( + test.task, test.query, test.concept, test.artifactReady, + test.candidates, test.protected, + )) + }) + } +} diff --git a/internal/mcp/explore_bounded_owner_recovery_test.go b/internal/mcp/explore_bounded_owner_recovery_test.go new file mode 100644 index 00000000..a30b8e28 --- /dev/null +++ b/internal/mcp/explore_bounded_owner_recovery_test.go @@ -0,0 +1,192 @@ +package mcp + +import ( + "context" + "fmt" + "testing" + "time" + + "github.com/zzet/gortex/internal/graph" +) + +type legacyOnlyExploreReader struct{ graph.Reader } + +func TestExploreCausalChangeOwnerHonorsExactFileCapScopeAndOverlay(t *testing.T) { + build := func(rows int) (*graph.Graph, *graph.Node, *graph.Node) { + memory := graph.New() + const path = "src/Builder.php" + leaf := divergentDefaultTestNode(path+"::Builder.build", graph.KindMethod, "build", path, "function build(): Owner") + owner := divergentDefaultTestNode(path+"::Owner", graph.KindType, "Owner", path, "class Owner") + nodes := []*graph.Node{leaf, owner} + for index := len(nodes); index < rows; index++ { + nodes = append(nodes, divergentDefaultTestNode( + fmt.Sprintf("%s::noise-%03d", path, index), graph.KindVariable, "noise", path, "", + )) + } + memory.AddBatch(nodes, nil) + return memory, leaf, owner + } + + memory, leaf, owner := build(exploreCausalChangeFileNodeCap) + got := exploreCausalChangeOwner(context.Background(), "change Owner builder", leaf, memory, graph.LocalizationNodeScope{}) + if got == nil || got.ID != owner.ID { + t.Fatalf("exact-cap owner = %#v, want %s", got, owner.ID) + } + if got := exploreCausalChangeOwner( + context.Background(), "change Owner builder", leaf, memory, + graph.LocalizationNodeScope{WorkspaceID: "other"}, + ); got != nil { + t.Fatalf("out-of-scope owner was promoted: %#v", got) + } + layer := graph.NewOverlayLayer() + layer.MarkFile(leaf.FilePath, true) + if got := exploreCausalChangeOwner( + context.Background(), "change Owner builder", leaf, + graph.NewOverlaidView(memory, layer), graph.LocalizationNodeScope{}, + ); got != nil { + t.Fatalf("overlay tombstone reintroduced durable owner: %#v", got) + } + + overflow, overflowLeaf, _ := build(exploreCausalChangeFileNodeCap + 1) + if got := exploreCausalChangeOwner( + context.Background(), "change Owner builder", overflowLeaf, overflow, graph.LocalizationNodeScope{}, + ); got != nil { + t.Fatalf("97-row causal file produced partial owner evidence: %#v", got) + } + if got := exploreCausalChangeOwner( + context.Background(), "change Owner builder", leaf, + legacyOnlyExploreReader{Reader: memory}, graph.LocalizationNodeScope{}, + ); got != nil { + t.Fatalf("reader without bounded file capability used a legacy fallback: %#v", got) + } +} + +func buildDivergentRankedCallableFiles(rowsPerFile []int) (*graph.Graph, []exploreTarget) { + memory := graph.New() + targets := make([]exploreTarget, 0, len(rowsPerFile)) + for fileIndex, rows := range rowsPerFile { + path := fmt.Sprintf("src/Owner%d.php", fileIndex) + ownerName := fmt.Sprintf("Owner%d", fileIndex) + owner := divergentDefaultTestNode(path+"::"+ownerName, graph.KindType, ownerName, path, "class "+ownerName) + callable := divergentDefaultTestNode(path+"::"+ownerName+".permissionHandler", graph.KindMethod, "permissionHandler", path, "function permissionHandler()") + constructor := divergentDefaultTestNode(path+"::"+ownerName+".__construct", graph.KindMethod, "__construct", path, "function __construct($filePermission = null)") + nodes := []*graph.Node{owner, callable, constructor} + for index := len(nodes); index < rows; index++ { + nodes = append(nodes, divergentDefaultTestNode( + fmt.Sprintf("%s::noise-%03d", path, index), graph.KindVariable, "noise", path, "", + )) + } + memory.AddBatch(nodes, nil) + targets = append(targets, exploreTarget{node: callable}) + } + return memory, targets +} + +func TestExploreDivergentDefaultRankedFilesHonorPerFileAndTotalCaps(t *testing.T) { + terms := exploreTerminalTerms("permission handler") + for _, test := range []struct { + name string + rows []int + wantBases int + wantComplete bool + }{ + {name: "64 row file", rows: []int{64}, wantBases: 1, wantComplete: true}, + {name: "65 row file", rows: []int{65}}, + {name: "96 rows total", rows: []int{32, 32, 32}, wantBases: 3, wantComplete: true}, + {name: "97 rows total", rows: []int{33, 32, 32}}, + } { + t.Run(test.name, func(t *testing.T) { + memory, targets := buildDivergentRankedCallableFiles(test.rows) + bases, complete := exploreDivergentDefaultBasesFromRankedCallables( + context.Background(), terms, targets, memory, graph.LocalizationNodeScope{}, time.Now().Add(30*time.Second), + ) + if complete != test.wantComplete || len(bases) != test.wantBases { + t.Fatalf("bases/complete = %d/%v, want %d/%v", len(bases), complete, test.wantBases, test.wantComplete) + } + }) + } +} + +func divergentProjectionBase(index int) exploreDivergentDefaultBase { + path := fmt.Sprintf("src/Base%d.php", index) + ownerName := fmt.Sprintf("Base%d", index) + return exploreDivergentDefaultBase{ + constructor: divergentDefaultTestNode(path+"::"+ownerName+".__construct", graph.KindMethod, "__construct", path, "function __construct($filePermission = null)"), + owner: divergentDefaultTestNode(path+"::"+ownerName, graph.KindType, ownerName, path, "class "+ownerName), + } +} + +func addDivergentProjectionEndpoints(memory *graph.Graph, base exploreDivergentDefaultBase, calls, extenders int, prefix string) { + for index := 0; index < calls; index++ { + id := fmt.Sprintf("src/%s-call-%02d.php::Child.__construct", prefix, index) + memory.AddNode(divergentDefaultTestNode(id, graph.KindMethod, "__construct", fmt.Sprintf("src/%s-call-%02d.php", prefix, index), "function __construct($filePermission = 0644)")) + memory.AddEdge(&graph.Edge{From: id, To: base.constructor.ID, Kind: graph.EdgeCalls}) + } + for index := 0; index < extenders; index++ { + id := fmt.Sprintf("src/%s-type-%02d.php::Child", prefix, index) + memory.AddNode(divergentDefaultTestNode(id, graph.KindType, "Child", fmt.Sprintf("src/%s-type-%02d.php", prefix, index), "class Child")) + memory.AddEdge(&graph.Edge{From: id, To: base.owner.ID, Kind: graph.EdgeExtends}) + } +} + +func TestProjectExploreDivergentDefaultOwnersHonorsEdgeAndEndpointCaps(t *testing.T) { + base := divergentProjectionBase(0) + memory := graph.New() + addDivergentProjectionEndpoints(memory, base, exploreDefaultOwnerEdgeCap, exploreDefaultOwnerEdgeCap, "exact") + projection, complete := projectExploreDivergentDefaultOwners(context.Background(), memory, []exploreDivergentDefaultBase{base}) + if !complete || len(projection.nodes) != exploreDefaultOwnerEndpointCap { + t.Fatalf("exact endpoint cap = %d/%v, want %d/true", len(projection.nodes), complete, exploreDefaultOwnerEndpointCap) + } + + overEdge := graph.New() + addDivergentProjectionEndpoints(overEdge, base, exploreDefaultOwnerEdgeCap+1, 0, "edge") + if projection, complete := projectExploreDivergentDefaultOwners(context.Background(), overEdge, []exploreDivergentDefaultBase{base}); complete || len(projection.nodes) != 0 { + t.Fatalf("17th caller returned partial projection: %#v, %v", projection, complete) + } + + second := divergentProjectionBase(1) + overUnion := graph.New() + addDivergentProjectionEndpoints(overUnion, base, exploreDefaultOwnerEdgeCap, exploreDefaultOwnerEdgeCap, "union") + addDivergentProjectionEndpoints(overUnion, second, 1, 0, "extra") + if projection, complete := projectExploreDivergentDefaultOwners(context.Background(), overUnion, []exploreDivergentDefaultBase{base, second}); complete || len(projection.nodes) != 0 { + t.Fatalf("33rd endpoint returned partial projection: %#v, %v", projection, complete) + } + + missing := graph.New() + missing.AddEdge(&graph.Edge{From: "missing-node", To: base.constructor.ID, Kind: graph.EdgeCalls}) + if projection, complete := projectExploreDivergentDefaultOwners(context.Background(), missing, []exploreDivergentDefaultBase{base}); complete || len(projection.nodes) != 0 { + t.Fatalf("missing exact refetch returned partial projection: %#v, %v", projection, complete) + } +} + +func TestExploreDivergentDefaultOwnerUsesRequestScopeAndOverlayView(t *testing.T) { + task := `StreamHandler::write reports permission denied; find the divergent filePermission default and owning type` + fixture := monologDivergentDefaultFixture() + targets := fixture.targets[:1] + + got := promoteExploreDivergentDefaultOwner( + context.Background(), task, targets, fixture.store, + graph.LocalizationNodeScope{WorkspaceID: "other"}, 3, divergentDefaultTestSource, 30*time.Second, + ) + if len(got) != 1 || got[0].node.ID != fixture.write.ID { + t.Fatalf("out-of-scope fallback promoted owner: %#v", got) + } + + layer := graph.NewOverlayLayer() + layer.MarkFile(fixture.baseType.FilePath, true) + got = promoteExploreDivergentDefaultOwner( + context.Background(), task, targets, graph.NewOverlaidView(fixture.store, layer), + graph.LocalizationNodeScope{}, 3, divergentDefaultTestSource, 30*time.Second, + ) + if len(got) != 1 || got[0].node.ID != fixture.write.ID { + t.Fatalf("overlay-hidden base owner was reintroduced: %#v", got) + } + + got = promoteExploreDivergentDefaultOwner( + context.Background(), task, targets, legacyOnlyExploreReader{Reader: fixture.store}, + graph.LocalizationNodeScope{}, 3, divergentDefaultTestSource, 30*time.Second, + ) + if len(got) != 1 || got[0].node.ID != fixture.write.ID { + t.Fatalf("reader without bounded capability used legacy full-row fallback: %#v", got) + } +} diff --git a/internal/mcp/explore_bounded_recovery.go b/internal/mcp/explore_bounded_recovery.go new file mode 100644 index 00000000..5b2d73e9 --- /dev/null +++ b/internal/mcp/explore_bounded_recovery.go @@ -0,0 +1,219 @@ +package mcp + +import ( + "context" + "sort" + + "github.com/zzet/gortex/internal/graph" +) + +type exploreContextNodesReader interface { + GetNodesByIDsContext(context.Context, []string) (map[string]*graph.Node, error) +} + +func exploreBoundedNodeIDs(ids []string, limit int) ([]string, bool) { + if limit < 0 { + return nil, false + } + seen := make(map[string]struct{}, min(len(ids), limit)) + out := make([]string, 0, min(len(ids), limit)) + for _, id := range ids { + if id == "" { + continue + } + if _, duplicate := seen[id]; duplicate { + continue + } + if len(out) >= limit { + return nil, false + } + seen[id] = struct{}{} + out = append(out, id) + } + return out, true +} + +// exploreNodesByIDsBounded exact-refetches only an already-bounded identity +// cohort. SQLite receives the request context; overlay and in-memory readers +// retain their exact-ID batch path with cancellation checked on both sides. +func exploreNodesByIDsBounded( + ctx context.Context, + reader graph.Reader, + ids []string, + limit int, +) (map[string]*graph.Node, bool) { + if reader == nil { + return nil, false + } + if ctx == nil { + ctx = context.Background() + } + boundedIDs, ok := exploreBoundedNodeIDs(ids, limit) + if !ok || ctx.Err() != nil { + return nil, false + } + if len(boundedIDs) == 0 { + return map[string]*graph.Node{}, true + } + var ( + nodes map[string]*graph.Node + err error + ) + if contextual, supported := reader.(exploreContextNodesReader); supported { + nodes, err = contextual.GetNodesByIDsContext(ctx, boundedIDs) + } else { + nodes = reader.GetNodesByIDs(boundedIDs) + } + if err != nil || ctx.Err() != nil || len(nodes) != len(boundedIDs) { + return nil, false + } + for _, id := range boundedIDs { + if nodes[id] == nil { + return nil, false + } + } + return nodes, true +} + +func exploreIncomingSourcesBounded( + ctx context.Context, + reader graph.Reader, + targetIDs []string, + kind graph.EdgeKind, + perTargetLimit int, +) (graph.BoundedIncomingSourceProjection, bool) { + bounded, ok := reader.(graph.BoundedIncomingSourceReader) + if !ok || perTargetLimit <= 0 { + return graph.BoundedIncomingSourceProjection{}, false + } + if ctx == nil { + ctx = context.Background() + } + if ctx.Err() != nil { + return graph.BoundedIncomingSourceProjection{}, false + } + projection, err := bounded.FindIncomingSourcesBounded(ctx, targetIDs, kind, perTargetLimit) + if err != nil || ctx.Err() != nil { + return graph.BoundedIncomingSourceProjection{}, false + } + return projection, true +} + +type exploreBoundedAdjacencyDirection uint8 + +const ( + exploreBoundedOutgoing exploreBoundedAdjacencyDirection = iota + exploreBoundedIncoming +) + +// exploreBoundedEdgeIdentities reads one complete metadata-free adjacency +// cohort. Optional capabilities fail closed: operational errors, sticky request +// cancellation, truncated keys, foreign endpoints, and malformed identities +// discard the whole projection rather than exposing a usable prefix. +func exploreBoundedEdgeIdentities( + ctx context.Context, + reader graph.Reader, + endpointIDs []string, + kinds []graph.EdgeKind, + perEndpointLimit int, + direction exploreBoundedAdjacencyDirection, +) (map[string][]graph.EdgeIdentity, bool) { + if reader == nil || perEndpointLimit <= 0 { + return nil, false + } + if ctx == nil { + ctx = context.Background() + } + if ctx.Err() != nil { + return nil, false + } + + requested := make(map[string]struct{}, len(endpointIDs)) + for _, id := range endpointIDs { + if id == "" { + return nil, false + } + requested[id] = struct{}{} + } + allowedKinds := make(map[graph.EdgeKind]struct{}, len(kinds)) + for _, kind := range kinds { + allowedKinds[kind] = struct{}{} + } + if len(allowedKinds) == 0 { + return nil, false + } + + var ( + projection graph.BoundedEdgeIdentityProjection + err error + ) + switch direction { + case exploreBoundedOutgoing: + bounded, ok := reader.(graph.BoundedOutgoingEdgeIdentityReader) + if !ok { + return nil, false + } + projection, err = bounded.FindOutgoingEdgeIdentitiesBounded(ctx, endpointIDs, kinds, perEndpointLimit) + case exploreBoundedIncoming: + bounded, ok := reader.(graph.BoundedIncomingEdgeIdentityReader) + if !ok { + return nil, false + } + projection, err = bounded.FindIncomingEdgeIdentitiesBounded(ctx, endpointIDs, kinds, perEndpointLimit) + default: + return nil, false + } + if err != nil || ctx.Err() != nil { + return nil, false + } + for endpoint, truncated := range projection.Truncated { + if _, ok := requested[endpoint]; !ok || truncated { + return nil, false + } + } + ordered := make(map[string][]graph.EdgeIdentity, len(projection.ByEndpoint)) + for endpoint, identities := range projection.ByEndpoint { + if _, ok := requested[endpoint]; !ok || len(identities) > perEndpointLimit { + return nil, false + } + identities = append([]graph.EdgeIdentity(nil), identities...) + sort.Slice(identities, func(i, j int) bool { + left, right := identities[i], identities[j] + if left.From != right.From { + return left.From < right.From + } + if left.To != right.To { + return left.To < right.To + } + if left.Kind != right.Kind { + return left.Kind < right.Kind + } + if left.FilePath != right.FilePath { + return left.FilePath < right.FilePath + } + return left.Line < right.Line + }) + seen := make(map[graph.EdgeIdentity]struct{}, len(identities)) + for _, identity := range identities { + if _, ok := allowedKinds[identity.Kind]; !ok { + return nil, false + } + if direction == exploreBoundedOutgoing { + if identity.From != endpoint || identity.To == "" { + return nil, false + } + } else if identity.To != endpoint || identity.From == "" { + return nil, false + } + if _, duplicate := seen[identity]; duplicate { + return nil, false + } + seen[identity] = struct{}{} + } + ordered[endpoint] = identities + } + if ctx.Err() != nil { + return nil, false + } + return ordered, true +} diff --git a/internal/mcp/explore_causal_change.go b/internal/mcp/explore_causal_change.go index 91ed2dbe..4b24bc81 100644 --- a/internal/mcp/explore_causal_change.go +++ b/internal/mcp/explore_causal_change.go @@ -1,6 +1,7 @@ package mcp import ( + "context" "sort" "strings" @@ -45,9 +46,11 @@ type exploreCausalChangeHydrator func(*graph.Node) ([]*graph.Node, bool) // never scans a repository, reads at most one source body, and replaces only // unprotected retrieval-tail rows. func promoteExploreCausalChangeTargets( + ctx context.Context, task string, targets []exploreTarget, - store graph.Store, + reader graph.Reader, + scope graph.LocalizationNodeScope, maxSymbols int, readSource func(*graph.Node) string, hydrateBridges ...exploreCausalChangeHydrator, @@ -59,7 +62,7 @@ func promoteExploreCausalChangeTargets( if len(hydrateBridges) > 0 { hydrateBridge = hydrateBridges[0] } - candidate, ok := selectExploreCausalChangeTargetWithConsumers(task, targets, store) + candidate, ok := selectExploreCausalChangeTargetWithConsumers(ctx, task, targets, reader) if !ok { return targets } @@ -96,7 +99,7 @@ func promoteExploreCausalChangeTargets( var owner exploreTarget ownerPresent := false if candidate.crossFile { - if node := exploreCausalChangeOwner(task, candidate.node, store); node != nil { + if node := exploreCausalChangeOwner(ctx, task, candidate.node, reader, scope); node != nil { owner, ownerPresent = exploreTargetByID(targets, node.ID) owner.node = node owner.causalChangeOwner = true @@ -271,12 +274,13 @@ func selectExploreCausalContinuation(task string, bridge exploreTarget, hinted * } func selectExploreCausalChangeTargetWithConsumers( + ctx context.Context, task string, targets []exploreTarget, - store graph.Store, + reader graph.Reader, ) (exploreCausalChangeCandidate, bool) { direct, directOK := selectExploreCausalChangeTarget(task, targets) - consumer, consumerOK := selectExploreCausalConsumerTarget(task, targets, store) + consumer, consumerOK := selectExploreCausalConsumerTarget(ctx, task, targets, reader) switch { case !consumerOK: return direct, directOK @@ -295,11 +299,15 @@ func selectExploreCausalChangeTargetWithConsumers( // it prevents a generic high-fanout caller walk from becoming another broad // search channel. func selectExploreCausalConsumerTarget( + ctx context.Context, task string, targets []exploreTarget, - store graph.Store, + reader graph.Reader, ) (exploreCausalChangeCandidate, bool) { - if store == nil || len(targets) == 0 { + if ctx == nil { + ctx = context.Background() + } + if reader == nil || len(targets) == 0 || ctx.Err() != nil { return exploreCausalChangeCandidate{}, false } terms := exploreTerminalTerms(shapeExploreQuery(task)) @@ -360,20 +368,18 @@ func selectExploreCausalConsumerTarget( for index, node := range frontier { ids[index] = node.ID } - incoming := store.GetInEdgesByNodeIDs(ids) + incoming, complete := exploreIncomingSourcesBounded( + ctx, reader, ids, graph.EdgeCalls, exploreCausalConsumerNodeFanoutCap, + ) + if !complete { + return exploreCausalChangeCandidate{}, false + } callerSet := make(map[string]struct{}, exploreCausalConsumerFrontierCap) for _, id := range ids { - local := make(map[string]struct{}, exploreCausalConsumerNodeFanoutCap+1) - for _, edge := range incoming[id] { - if edge == nil || edge.Kind != graph.EdgeCalls || edge.From == "" { - continue - } - local[edge.From] = struct{}{} - } - if len(local) > exploreCausalConsumerNodeFanoutCap { + if incoming.Truncated[id] { continue } - for callerID := range local { + for _, callerID := range incoming.Sources[id] { callerSet[callerID] = struct{}{} } } @@ -385,7 +391,12 @@ func selectExploreCausalConsumerTarget( if len(callerIDs) > exploreCausalConsumerFrontierCap { callerIDs = callerIDs[:exploreCausalConsumerFrontierCap] } - nodes := store.GetNodesByIDs(callerIDs) + nodes, complete := exploreNodesByIDsBounded( + ctx, reader, callerIDs, exploreCausalConsumerFrontierCap, + ) + if !complete { + return exploreCausalChangeCandidate{}, false + } next := make([]*graph.Node, 0, len(callerIDs)) for _, callerID := range callerIDs { if _, exists := expanded[callerID]; exists { @@ -559,18 +570,45 @@ func exploreCausalChangeLess(left, right exploreCausalChangeCandidate) bool { return exploreDraftNodeKey(left.node) < exploreDraftNodeKey(right.node) } -func exploreCausalChangeOwner(task string, leaf *graph.Node, store graph.Store) *graph.Node { - if leaf == nil || leaf.FilePath == "" || store == nil { +func exploreCausalChangeOwner( + ctx context.Context, + task string, + leaf *graph.Node, + reader graph.Reader, + scope graph.LocalizationNodeScope, +) *graph.Node { + if leaf == nil || leaf.FilePath == "" || reader == nil { return nil } - fileNodes := store.GetFileNodesByPaths([]string{leaf.FilePath})[leaf.FilePath] - if len(fileNodes) == 0 || len(fileNodes) > exploreCausalChangeFileNodeCap { + if ctx == nil { + ctx = context.Background() + } + page, complete := boundedLocalizationFileNodes( + ctx, + reader, + localizationFileBudgetFor(ctx), + leaf.FilePath, + graph.LocalizationNodeScope{}, + exploreCausalChangeFileNodeCap, + ) + if !complete || len(page.Nodes) == 0 { + return nil + } + typeIDs := make([]string, 0, len(page.Nodes)) + for _, summary := range page.Nodes { + if summary != nil && summary.Kind == graph.KindType && summary.ID != "" { + typeIDs = append(typeIDs, summary.ID) + } + } + full, complete := exploreNodesByIDsBounded(ctx, reader, typeIDs, exploreCausalChangeFileNodeCap) + if !complete { return nil } - types := make([]*graph.Node, 0, 8) - for _, node := range fileNodes { + types := make([]*graph.Node, 0, len(typeIDs)) + for _, id := range typeIDs { + node := full[id] if node == nil || node.Kind != graph.KindType || exploreDraftIsTestNode(node) || - !exploreNodesShareExactScope(leaf, node) { + !scope.Allows(node) || !exploreNodesShareExactScope(leaf, node) { continue } types = append(types, node) @@ -645,7 +683,7 @@ func exploreCausalChangeAdmissionProtected(task string, targets []exploreTarget, } return target.causalChangeBridge || target.causalChangeLeaf || target.causalChangeOwner || target.divergentDefaultOwner || target.divergentDefaultType || - target.conceptImplementation || target.conceptComplement || + target.conceptImplementation || target.conceptComplement || target.sourceRange || target.exactContent || target.exactContentAmbiguous || target.sourceLiteral || target.typedAnchorProjection || exploreLocalizationExplicitAnchor(task, target.node) diff --git a/internal/mcp/explore_causal_change_test.go b/internal/mcp/explore_causal_change_test.go index b615d4ac..722f473b 100644 --- a/internal/mcp/explore_causal_change_test.go +++ b/internal/mcp/explore_causal_change_test.go @@ -1,6 +1,7 @@ package mcp import ( + "context" "slices" "testing" @@ -19,6 +20,34 @@ func causalChangeTestNode(id, name, file string, kind graph.NodeKind, signature } } +func promoteExploreCausalChangeTargetsForTest( + task string, + targets []exploreTarget, + reader graph.Reader, + maxSymbols int, + readSource func(*graph.Node) string, + hydrateBridges ...exploreCausalChangeHydrator, +) []exploreTarget { + return promoteExploreCausalChangeTargets( + context.Background(), task, targets, reader, graph.LocalizationNodeScope{}, + maxSymbols, readSource, hydrateBridges..., + ) +} + +func exploreCausalChangeOwnerForTest(task string, leaf *graph.Node, reader graph.Reader) *graph.Node { + return exploreCausalChangeOwner( + context.Background(), task, leaf, reader, graph.LocalizationNodeScope{}, + ) +} + +func selectExploreCausalConsumerTargetForTest( + task string, + targets []exploreTarget, + reader graph.Reader, +) (exploreCausalChangeCandidate, bool) { + return selectExploreCausalConsumerTarget(context.Background(), task, targets, reader) +} + func causalChangeTargetIDs(targets []exploreTarget) []string { ids := make([]string, 0, len(targets)) for _, target := range targets { @@ -48,7 +77,7 @@ func TestPromoteExploreCausalChangeTargetsReservesDelegatedLeaf(t *testing.T) { {node: causalChangeTestNode("repo/router/errors.go::invalidNodeDiagnostic", "invalidNodeDiagnostic", "router/errors.go", graph.KindFunction, "func invalidNodeDiagnostic()")}, } task := "Request execution panics during case insensitive path fallback instead of returning not found" - got := promoteExploreCausalChangeTargets(task, targets, nil, len(targets), func(node *graph.Node) string { + got := promoteExploreCausalChangeTargetsForTest(task, targets, nil, len(targets), func(node *graph.Node) string { if node.ID == leaf.ID { return "func resolveCaseInsensitivePathRec(path string) []byte { return nil }" } @@ -106,7 +135,7 @@ func TestPromoteExploreCausalChangeTargetsPrefersGraphCallerOverKeywordOnlyTail( {node: keywordOnly}, } task := "Multiline replace duplicates output when replacement captures span several lines" - got := promoteExploreCausalChangeTargets(task, targets, g, len(targets), func(node *graph.Node) string { + got := promoteExploreCausalChangeTargetsForTest(task, targets, g, len(targets), func(node *graph.Node) string { if node.ID == leaf.ID { return "fn apply_replacements(&mut self) -> io::Result<()> { Ok(()) }" } @@ -144,7 +173,7 @@ func TestExploreCausalChangeOwnerPrefersUniquelyReturnedStateType(t *testing.T) g.AddNode(builder) g.AddNode(state) - got := exploreCausalChangeOwner("parallel roots leak rule state between traversals", leaf, g) + got := exploreCausalChangeOwnerForTest("parallel roots leak rule state between traversals", leaf, g) if got == nil || got.ID != state.ID { t.Fatalf("change owner = %#v, want uniquely returned state type %q", got, state.ID) } @@ -193,7 +222,7 @@ func TestPromoteExploreCausalChangeTargetsCarriesCrossFileBridgeToContinuationAn readCalls, hydrateCalls := 0, 0 task := "The replace with captures at path should replace all matches while preserving captures in their surrounding context" - got := promoteExploreCausalChangeTargets(task, []exploreTarget{{ + got := promoteExploreCausalChangeTargetsForTest(task, []exploreTarget{{ node: seed, source: "replace_with_captures_at replaces captures", callers: []*graph.Node{bridge}, }}, store, 4, func(node *graph.Node) string { readCalls++ @@ -257,7 +286,7 @@ func TestPromoteExploreCausalChangeTargetsUsesNestedDelegationHint(t *testing.T) causalCallees: []exploreCausalNeighbor{{node: bridge, hop: 1}, {node: continuation, hop: 2}}, }} task := "The redirect fixed path fallback should follow case insensitive path lookup recursively without cleaning away valid request segments" - got := promoteExploreCausalChangeTargets(task, targets, nil, 3, func(node *graph.Node) string { + got := promoteExploreCausalChangeTargetsForTest(task, targets, nil, 3, func(node *graph.Node) string { if node.ID != bridge.ID { t.Fatalf("source read for %q, want hinted bridge %q", node.ID, bridge.ID) } @@ -303,7 +332,7 @@ func TestPromoteExploreCausalChangeTargetsContinuesInheritedCaller(t *testing.T) "combineRecords", "HipChatHandler.combineRecords", "src/Handler/HipChatHandler.php", graph.KindMethod, ) task := "The Hip Chat handler should handle each batch by combining records before checking handling and writing messages to the API" - got := promoteExploreCausalChangeTargets(task, []exploreTarget{{ + got := promoteExploreCausalChangeTargetsForTest(task, []exploreTarget{{ node: seed, source: "isHandling checks every record", callers: []*graph.Node{genericCaller, bridge}, }}, nil, 3, func(node *graph.Node) string { if node.ID != bridge.ID { @@ -357,7 +386,7 @@ func TestPromoteExploreCausalChangeTargetsContinuationFailsClosed(t *testing.T) } for _, test := range tests { t.Run(test.name, func(t *testing.T) { - got := promoteExploreCausalChangeTargets(task, []exploreTarget{{ + got := promoteExploreCausalChangeTargetsForTest(task, []exploreTarget{{ node: seed, source: "resolveRequestPath dispatches request work", callers: []*graph.Node{bridge}, }}, nil, 3, func(*graph.Node) string { return "func dispatchRequest() { dispatchPrimaryWork(); dispatchBackupWork() }" diff --git a/internal/mcp/explore_causal_consumer_test.go b/internal/mcp/explore_causal_consumer_test.go index e814ddb5..f038785e 100644 --- a/internal/mcp/explore_causal_consumer_test.go +++ b/internal/mcp/explore_causal_consumer_test.go @@ -1,6 +1,7 @@ package mcp import ( + "fmt" "testing" "github.com/zzet/gortex/internal/graph" @@ -29,7 +30,7 @@ func TestSelectExploreCausalConsumerTargetBridgesRetainedComponent(t *testing.T) } store.AddEdge(&graph.Edge{From: consumer.ID, To: callee.ID, Kind: graph.EdgeCalls}) - candidate, ok := selectExploreCausalConsumerTarget( + candidate, ok := selectExploreCausalConsumerTargetForTest( "panic while using --replace with multiline captures in context", []exploreTarget{{node: seed, callees: []*graph.Node{callee}}, {node: printerPeer}}, store, @@ -68,7 +69,7 @@ func TestSelectExploreCausalConsumerTargetCollapsesSameFileWrappers(t *testing.T store.AddEdge(&graph.Edge{From: wrapper2.ID, To: wrapper1.ID, Kind: graph.EdgeCalls}) store.AddEdge(&graph.Edge{From: consumer.ID, To: wrapper2.ID, Kind: graph.EdgeCalls}) - candidate, ok := selectExploreCausalConsumerTarget( + candidate, ok := selectExploreCausalConsumerTargetForTest( "hidden files whitelisted by an ancestor .ignore fail for explicit current directory", []exploreTarget{{node: seed}}, store, @@ -78,6 +79,60 @@ func TestSelectExploreCausalConsumerTargetCollapsesSameFileWrappers(t *testing.T } } +func TestSelectExploreCausalConsumerTargetSkipsOnlySaturatedFrontierTarget(t *testing.T) { + build := func(firstFanout int, includeFallback bool) (*graph.Graph, []exploreTarget, *graph.Node, *graph.Node) { + store := graph.New() + firstSeed := &graph.Node{ + ID: "crates/matcher/src/first.rs::Matcher.first", Name: "first", + Kind: graph.KindMethod, FilePath: "crates/matcher/src/first.rs", + } + secondSeed := &graph.Node{ + ID: "crates/matcher/src/second.rs::Matcher.second", Name: "second", + Kind: graph.KindMethod, FilePath: "crates/matcher/src/second.rs", + } + peer := &graph.Node{ + ID: "crates/printer/src/standard.rs::StandardBuilder", Name: "StandardBuilder", + Kind: graph.KindType, FilePath: "crates/printer/src/standard.rs", + } + firstConsumer := &graph.Node{ + ID: "crates/printer/src/first.rs::replace_captures_context", Name: "replace_captures_context", + Kind: graph.KindFunction, FilePath: "crates/printer/src/first.rs", + } + fallback := &graph.Node{ + ID: "crates/printer/src/fallback.rs::replace_captures_context", Name: "replace_captures_context", + Kind: graph.KindFunction, FilePath: "crates/printer/src/fallback.rs", + } + for _, node := range []*graph.Node{firstSeed, secondSeed, peer, firstConsumer, fallback} { + store.AddNode(node) + } + store.AddEdge(&graph.Edge{From: firstConsumer.ID, To: firstSeed.ID, Kind: graph.EdgeCalls}) + for index := 1; index < firstFanout; index++ { + noise := &graph.Node{ + ID: fmt.Sprintf("crates/noise/src/noise-%02d.rs::unrelated", index), Name: "unrelated", + Kind: graph.KindFunction, FilePath: fmt.Sprintf("crates/noise/src/noise-%02d.rs", index), + } + store.AddNode(noise) + store.AddEdge(&graph.Edge{From: noise.ID, To: firstSeed.ID, Kind: graph.EdgeCalls}) + } + if includeFallback { + store.AddEdge(&graph.Edge{From: fallback.ID, To: secondSeed.ID, Kind: graph.EdgeCalls}) + } + return store, []exploreTarget{{node: firstSeed}, {node: secondSeed}, {node: peer}}, firstConsumer, fallback + } + + exact, targets, firstConsumer, _ := build(exploreCausalConsumerNodeFanoutCap, false) + candidate, ok := selectExploreCausalConsumerTargetForTest("replace captures in context", targets, exact) + if !ok || candidate.node == nil || candidate.node.ID != firstConsumer.ID { + t.Fatalf("exact eight-source frontier was not admitted: ok=%v candidate=%#v", ok, candidate) + } + + saturated, targets, _, fallback := build(exploreCausalConsumerNodeFanoutCap+1, true) + candidate, ok = selectExploreCausalConsumerTargetForTest("replace captures in context", targets, saturated) + if !ok || candidate.node == nil || candidate.node.ID != fallback.ID || candidate.parentID != targets[1].node.ID { + t.Fatalf("ninth source poisoned unrelated frontier target: ok=%v candidate=%#v", ok, candidate) + } +} + func TestSelectExploreCausalConsumerTargetRejectsAmbiguousFiles(t *testing.T) { store := graph.New() seed := &graph.Node{ @@ -106,7 +161,7 @@ func TestSelectExploreCausalConsumerTargetRejectsAmbiguousFiles(t *testing.T) { store.AddEdge(&graph.Edge{From: first.ID, To: seed.ID, Kind: graph.EdgeCalls}) store.AddEdge(&graph.Edge{From: second.ID, To: seed.ID, Kind: graph.EdgeCalls}) - if candidate, ok := selectExploreCausalConsumerTarget( + if candidate, ok := selectExploreCausalConsumerTargetForTest( "replace output is duplicated", []exploreTarget{{node: seed}, {node: printerPeer}, {node: otherPeer}}, store, diff --git a/internal/mcp/explore_distinctive_bare_tokens_test.go b/internal/mcp/explore_distinctive_bare_tokens_test.go new file mode 100644 index 00000000..4af85110 --- /dev/null +++ b/internal/mcp/explore_distinctive_bare_tokens_test.go @@ -0,0 +1,52 @@ +package mcp + +import ( + "reflect" + "testing" +) + +func TestExploreDistinctiveBareTokensMinesIdentifierShapes(t *testing.T) { + task := "Localize bug: gin-4372 checkout rejects websocket Accept after it calls " + + "WriteHeaderNow before Hijack. In responseWriter.Hijack, w.Written is true " + + "when size is 0 and maxRedirects stays unset." + tokens := exploreDistinctiveBareTokens(task, "gin-4372") + want := map[string]bool{ + "WriteHeaderNow": true, + "responseWriter.Hijack": true, + "maxRedirects": true, + } + got := make(map[string]bool, len(tokens)) + for _, token := range tokens { + got[token] = true + } + for token := range want { + if !got[token] { + t.Fatalf("expected %q among mined tokens: %v", token, tokens) + } + } + if got["gin-4372"] || got["checkout"] || got["websocket"] { + t.Fatalf("repo name or plain prose leaked into mined tokens: %v", tokens) + } + if len(tokens) > exploreDistinctiveBareTokenMaxTerms { + t.Fatalf("token cap exceeded: %v", tokens) + } +} + +func TestExploreDistinctiveBareTokensEmitsDottedFinalSegment(t *testing.T) { + tokens := exploreDistinctiveBareTokens( + "crash inside c.AbortWithStatus during recovery middleware handling", "gin") + got := make(map[string]bool, len(tokens)) + for _, token := range tokens { + got[token] = true + } + if !got["AbortWithStatus"] { + t.Fatalf("dotted citation must contribute its final segment: %v", tokens) + } +} + +func TestExploreDistinctiveBareTokensEmptyForPlainProse(t *testing.T) { + tokens := exploreDistinctiveBareTokens("the report says the build is slow and fails often", "demo") + if !reflect.DeepEqual(tokens, []string{}) && tokens != nil { + t.Fatalf("plain prose must mine nothing: %v", tokens) + } +} diff --git a/internal/mcp/explore_divergent_default_owner.go b/internal/mcp/explore_divergent_default_owner.go index 6140a5cc..a8dddd8c 100644 --- a/internal/mcp/explore_divergent_default_owner.go +++ b/internal/mcp/explore_divergent_default_owner.go @@ -1,6 +1,7 @@ package mcp import ( + "context" "sort" "strconv" "strings" @@ -81,8 +82,17 @@ type exploreDivergentDefaultProjection struct { // cannot depend on machine load — see the note on // exploreDefaultOwnerFallbackSLO. Left variadic so the fifteen call sites that // never reach the fallback stay untouched. -func promoteExploreDivergentDefaultOwner(task string, targets []exploreTarget, store graph.Store, maxSymbols int, readSource func(*graph.Node) string, fallbackSLO ...time.Duration) []exploreTarget { - if store == nil || readSource == nil || len(targets) == 0 || !exploreQueryIsConceptTask(task) { +func promoteExploreDivergentDefaultOwner( + ctx context.Context, + task string, + targets []exploreTarget, + reader graph.Reader, + scope graph.LocalizationNodeScope, + maxSymbols int, + readSource func(*graph.Node) string, + fallbackSLO ...time.Duration, +) []exploreTarget { + if reader == nil || readSource == nil || len(targets) == 0 || !exploreQueryIsConceptTask(task) { return targets } taskTerms := exploreTerminalTerms(task) @@ -96,12 +106,14 @@ func promoteExploreDivergentDefaultOwner(task string, targets []exploreTarget, s slo = fallbackSLO[0] } deadline := time.Now().Add(slo) - bases, ok = exploreDivergentDefaultBasesFromRankedCallables(taskTerms, targets, store, deadline) + bases, ok = exploreDivergentDefaultBasesFromRankedCallables( + ctx, taskTerms, targets, reader, scope, deadline, + ) if !ok || len(bases) == 0 { return targets } } - projection, ok := projectExploreDivergentDefaultOwners(store, bases) + projection, ok := projectExploreDivergentDefaultOwners(ctx, reader, bases) if !ok { return targets } @@ -148,10 +160,14 @@ func placeExploreDivergentDefaultOwner(task string, targets []exploreTarget, mat maxSymbols = len(targets) } missing := 0 - if exploreTargetIndex(targets, match.constructor.ID) < 0 && exploreTargetIndex(targets, match.baseCtor.ID) < 0 { + constructorIndex := exploreTargetIndex(targets, match.constructor.ID) + baseConstructorIndex := exploreTargetIndex(targets, match.baseCtor.ID) + if constructorIndex < 0 && (baseConstructorIndex < 0 || targets[baseConstructorIndex].sourceRange) { missing++ } - if exploreTargetIndex(targets, match.owner.ID) < 0 && exploreTargetIndex(targets, match.baseOwner.ID) < 0 { + ownerIndex := exploreTargetIndex(targets, match.owner.ID) + baseOwnerIndex := exploreTargetIndex(targets, match.baseOwner.ID) + if ownerIndex < 0 && (baseOwnerIndex < 0 || targets[baseOwnerIndex].sourceRange) { missing++ } overflow := len(targets) + missing - maxSymbols @@ -181,8 +197,10 @@ func placeExploreDivergentDefaultOwner(task string, targets []exploreTarget, mat return } if index := exploreTargetIndex(replaced, baseID); index >= 0 { - replaced[index] = candidate - return + if !replaced[index].sourceRange { + replaced[index] = candidate + return + } } replaced = append(replaced, candidate) } @@ -201,7 +219,7 @@ func exploreDivergentDefaultAdmissionProtected(task string, target exploreTarget return true } return target.divergentDefaultOwner || target.divergentDefaultType || target.conceptImplementation || - target.exactContent || target.exactContentAmbiguous || target.sourceLiteral || + target.sourceRange || target.exactContent || target.exactContentAmbiguous || target.sourceLiteral || exploreLocalizationExplicitAnchor(task, target.node) } @@ -512,16 +530,32 @@ type exploreRankedCallableOwner struct { // exploreDivergentDefaultBasesFromRankedCallables is the strictly bounded // recovery path for a common ranking gap: the consuming method is present but // its constructor/type pair is not. It follows only the callable's exact -// enclosing owner, hydrates those owners in one batch, and reads only their -// indexed file-node sets in one batch. It never scans a repository or source -// tree, and ambiguous/oversized projections are rejected. -func exploreDivergentDefaultBasesFromRankedCallables(taskTerms map[string]struct{}, targets []exploreTarget, store graph.Store, deadline time.Time) ([]exploreDivergentDefaultBase, bool) { - if store == nil || len(taskTerms) == 0 || time.Now().After(deadline) { +// enclosing owner, hydrates those owners in one batch, and reads each distinct +// owner file through a capped metadata-free projection. It never scans a +// repository or source tree, and ambiguous/oversized projections are rejected. +func exploreDivergentDefaultBasesFromRankedCallables( + ctx context.Context, + taskTerms map[string]struct{}, + targets []exploreTarget, + reader graph.Reader, + scope graph.LocalizationNodeScope, + deadline time.Time, +) ([]exploreDivergentDefaultBase, bool) { + if ctx == nil { + ctx = context.Background() + } + if reader == nil || len(taskTerms) == 0 || ctx.Err() != nil || time.Now().After(deadline) { return nil, false } + readCtx, cancel := context.WithDeadline(ctx, deadline) + defer cancel() + owners := make([]exploreRankedCallableOwner, 0, exploreDefaultOwnerFileCap) seenOwners := make(map[string]struct{}, exploreDefaultOwnerFileCap) for _, target := range targets { + if readCtx.Err() != nil { + return nil, false + } callable := target.node if callable == nil || (callable.Kind != graph.KindMethod && callable.Kind != graph.KindFunction) || exploreConstructorNode(callable) || exploreDraftIsTestNode(callable) { @@ -553,8 +587,8 @@ func exploreDivergentDefaultBasesFromRankedCallables(taskTerms map[string]struct for _, candidate := range owners { ownerIDs = append(ownerIDs, candidate.ownerID) } - hydrated := store.GetNodesByIDs(ownerIDs) - if time.Now().After(deadline) || len(hydrated) != len(ownerIDs) { + hydrated, complete := exploreNodesByIDsBounded(readCtx, reader, ownerIDs, exploreDefaultOwnerFileCap) + if !complete { return nil, false } filePaths := make([]string, 0, len(owners)) @@ -576,27 +610,67 @@ func exploreDivergentDefaultBasesFromRankedCallables(taskTerms map[string]struct return nil, len(coherent) == 0 } - fileNodes := store.GetFileNodesByPaths(filePaths) - if time.Now().After(deadline) { - return nil, false - } - total := 0 + budget := localizationFileBudgetFor(ctx) + fileCandidates := make(map[string][]string, len(filePaths)) + constructorIDs := make([]string, 0, exploreDefaultOwnerTotalCap) + seenConstructors := make(map[string]struct{}, exploreDefaultOwnerTotalCap) + totalRows := 0 for _, path := range filePaths { - nodes := fileNodes[path] - if len(nodes) > exploreDefaultOwnerFileNodeCap { + if readCtx.Err() != nil { return nil, false } - total += len(nodes) - if total > exploreDefaultOwnerTotalCap { + remaining := exploreDefaultOwnerTotalCap - totalRows + pageLimit := exploreDefaultOwnerFileNodeCap + if remaining < pageLimit { + pageLimit = remaining + } + if pageLimit <= 0 { + // A one-row sentinel distinguishes an empty remaining file from a + // total-cap overflow without opening an unbounded scan. + pageLimit = 1 + } + page, complete := boundedLocalizationFileNodes( + readCtx, reader, budget, path, graph.LocalizationNodeScope{}, pageLimit, + ) + if !complete { return nil, false } + rows := page.Total + if len(page.Nodes) > rows { + rows = len(page.Nodes) + } + if rows > exploreDefaultOwnerFileNodeCap || totalRows+rows > exploreDefaultOwnerTotalCap { + return nil, false + } + totalRows += rows + for _, summary := range page.Nodes { + if !exploreConstructorNode(summary) || summary.ID == "" { + continue + } + if _, duplicate := seenConstructors[summary.ID]; duplicate { + continue + } + seenConstructors[summary.ID] = struct{}{} + constructorIDs = append(constructorIDs, summary.ID) + fileCandidates[path] = append(fileCandidates[path], summary.ID) + } } + constructors, complete := exploreNodesByIDsBounded( + readCtx, reader, constructorIDs, exploreDefaultOwnerTotalCap, + ) + if !complete { + return nil, false + } bases := make([]exploreDivergentDefaultBase, 0, len(coherent)) for _, candidate := range coherent { + if !scope.Allows(candidate.owner) { + continue + } var viable []exploreDivergentDefaultBase - for _, node := range fileNodes[candidate.owner.FilePath] { - if !exploreConstructorNode(node) || exploreDraftIsTestNode(node) { + for _, id := range fileCandidates[candidate.owner.FilePath] { + node := constructors[id] + if !exploreConstructorNode(node) || exploreDraftIsTestNode(node) || !scope.Allows(node) { continue } ownerID, _ := graph.EnclosingFromID(node.ID, node.Kind) @@ -629,7 +703,7 @@ func exploreDivergentDefaultBasesFromRankedCallables(taskTerms map[string]struct } } } - return bases, true + return bases, readCtx.Err() == nil } func exploreRankedCallableOwnerCoherent(callable, owner *graph.Node) bool { @@ -660,46 +734,59 @@ func exploreNodesShareExactScope(left, right *graph.Node) bool { // preserves the relationship kind and full node metadata needed by this proof. // Any oversized or partially hydrated projection is rejected rather than // interpreted as evidence that no competing child exists. -func projectExploreDivergentDefaultOwners(store graph.Store, bases []exploreDivergentDefaultBase) (exploreDivergentDefaultProjection, bool) { - if store == nil || len(bases) == 0 || len(bases) > exploreDefaultOwnerBaseCap { +func projectExploreDivergentDefaultOwners( + ctx context.Context, + reader graph.Reader, + bases []exploreDivergentDefaultBase, +) (exploreDivergentDefaultProjection, bool) { + if ctx == nil { + ctx = context.Background() + } + if reader == nil || len(bases) == 0 || len(bases) > exploreDefaultOwnerBaseCap || ctx.Err() != nil { return exploreDivergentDefaultProjection{}, false } - seedIDs := make([]string, 0, len(bases)*2) + constructorIDs := make([]string, 0, len(bases)) + ownerIDs := make([]string, 0, len(bases)) for _, base := range bases { - seedIDs = append(seedIDs, base.constructor.ID, base.owner.ID) + if base.constructor == nil || base.constructor.ID == "" || base.owner == nil || base.owner.ID == "" { + return exploreDivergentDefaultProjection{}, false + } + constructorIDs = append(constructorIDs, base.constructor.ID) + ownerIDs = append(ownerIDs, base.owner.ID) + } + callers, complete := exploreIncomingSourcesBounded( + ctx, reader, constructorIDs, graph.EdgeCalls, exploreDefaultOwnerEdgeCap, + ) + if !complete { + return exploreDivergentDefaultProjection{}, false + } + extenders, complete := exploreIncomingSourcesBounded( + ctx, reader, ownerIDs, graph.EdgeExtends, exploreDefaultOwnerEdgeCap, + ) + if !complete { + return exploreDivergentDefaultProjection{}, false } - incoming := store.GetInEdgesByNodeIDs(seedIDs) projection := exploreDivergentDefaultProjection{ callers: make(map[string][]string, len(bases)), extenders: make(map[string][]string, len(bases)), } - endpointIDs := make(map[string]struct{}) - appendEndpoint := func(targetID string, edge *graph.Edge, kind graph.EdgeKind, destinations map[string][]string) bool { - if edge == nil || edge.Kind != kind || edge.To != targetID || edge.From == "" { - return true + endpointIDs := make(map[string]struct{}, exploreDefaultOwnerEndpointCap) + for _, base := range bases { + constructorID := base.constructor.ID + ownerID := base.owner.ID + if callers.Truncated[constructorID] || extenders.Truncated[ownerID] { + return exploreDivergentDefaultProjection{}, false } - for _, existing := range destinations[targetID] { - if existing == edge.From { - return true - } + projection.callers[constructorID] = append([]string(nil), callers.Sources[constructorID]...) + projection.extenders[ownerID] = append([]string(nil), extenders.Sources[ownerID]...) + for _, id := range projection.callers[constructorID] { + endpointIDs[id] = struct{}{} } - destinations[targetID] = append(destinations[targetID], edge.From) - if len(destinations[targetID]) > exploreDefaultOwnerEdgeCap { - return false + for _, id := range projection.extenders[ownerID] { + endpointIDs[id] = struct{}{} } - endpointIDs[edge.From] = struct{}{} - return len(endpointIDs) <= exploreDefaultOwnerEndpointCap - } - for _, base := range bases { - for _, edge := range incoming[base.constructor.ID] { - if !appendEndpoint(base.constructor.ID, edge, graph.EdgeCalls, projection.callers) { - return exploreDivergentDefaultProjection{}, false - } - } - for _, edge := range incoming[base.owner.ID] { - if !appendEndpoint(base.owner.ID, edge, graph.EdgeExtends, projection.extenders) { - return exploreDivergentDefaultProjection{}, false - } + if len(endpointIDs) > exploreDefaultOwnerEndpointCap { + return exploreDivergentDefaultProjection{}, false } } ids := make([]string, 0, len(endpointIDs)) @@ -707,8 +794,10 @@ func projectExploreDivergentDefaultOwners(store graph.Store, bases []exploreDive ids = append(ids, id) } sort.Strings(ids) - projection.nodes = store.GetNodesByIDs(ids) - if len(projection.nodes) != len(ids) { + projection.nodes, complete = exploreNodesByIDsBounded( + ctx, reader, ids, exploreDefaultOwnerEndpointCap, + ) + if !complete { return exploreDivergentDefaultProjection{}, false } return projection, true diff --git a/internal/mcp/explore_divergent_default_owner_test.go b/internal/mcp/explore_divergent_default_owner_test.go index ba06f023..ec775f7b 100644 --- a/internal/mcp/explore_divergent_default_owner_test.go +++ b/internal/mcp/explore_divergent_default_owner_test.go @@ -40,21 +40,63 @@ type countingDivergentDefaultStore struct { } func (s *countingDivergentDefaultStore) GetNodesByIDs(ids []string) map[string]*graph.Node { + s.nodeBatchCalls++ + return s.Store.GetNodesByIDs(ids) +} + +func (s *countingDivergentDefaultStore) GetNodesByIDsContext(ctx context.Context, ids []string) (map[string]*graph.Node, error) { s.nodeBatchCalls++ if s.nodeBatchDelay > 0 { - time.Sleep(s.nodeBatchDelay) + timer := time.NewTimer(s.nodeBatchDelay) + defer timer.Stop() + select { + case <-ctx.Done(): + return nil, ctx.Err() + case <-timer.C: + } } - return s.Store.GetNodesByIDs(ids) + if contextual, ok := s.Store.(interface { + GetNodesByIDsContext(context.Context, []string) (map[string]*graph.Node, error) + }); ok { + return contextual.GetNodesByIDsContext(ctx, ids) + } + return s.Store.GetNodesByIDs(ids), ctx.Err() +} + +func (s *countingDivergentDefaultStore) GetFileNodesByPaths([]string) map[string][]*graph.Node { + panic("legacy full file-node batch must not be used") } -func (s *countingDivergentDefaultStore) GetFileNodesByPaths(paths []string) map[string][]*graph.Node { +func (s *countingDivergentDefaultStore) FindFileNodesBounded( + ctx context.Context, + path string, + scope graph.LocalizationNodeScope, + limit int, +) (graph.BoundedNodeProjection, error) { s.fileBatchCalls++ - return s.Store.GetFileNodesByPaths(paths) + bounded, ok := s.Store.(graph.BoundedFileNodeReader) + if !ok { + return graph.BoundedNodeProjection{}, graph.ErrBoundedLocalizationUnavailable + } + return bounded.FindFileNodesBounded(ctx, path, scope, limit) +} + +func (s *countingDivergentDefaultStore) GetInEdgesByNodeIDs([]string) map[string][]*graph.Edge { + panic("legacy full incoming-edge batch must not be used") } -func (s *countingDivergentDefaultStore) GetInEdgesByNodeIDs(ids []string) map[string][]*graph.Edge { +func (s *countingDivergentDefaultStore) FindIncomingSourcesBounded( + ctx context.Context, + targetIDs []string, + kind graph.EdgeKind, + limit int, +) (graph.BoundedIncomingSourceProjection, error) { s.inBatchCalls++ - return s.Store.GetInEdgesByNodeIDs(ids) + bounded, ok := s.Store.(graph.BoundedIncomingSourceReader) + if !ok { + return graph.BoundedIncomingSourceProjection{}, graph.ErrBoundedLocalizationUnavailable + } + return bounded.FindIncomingSourcesBounded(ctx, targetIDs, kind, limit) } func divergentDefaultTestSource(node *graph.Node) string { @@ -105,6 +147,19 @@ func monologDivergentDefaultFixture() divergentDefaultFixture { } } +func promoteExploreDivergentDefaultOwnerForTest( + task string, + targets []exploreTarget, + reader graph.Reader, + maxSymbols int, + readSource func(*graph.Node) string, + fallbackSLO ...time.Duration, +) []exploreTarget { + return promoteExploreDivergentDefaultOwner( + context.Background(), task, targets, reader, graph.LocalizationNodeScope{}, maxSymbols, readSource, fallbackSLO..., + ) +} + func TestDivergentDefaultOwnerPromotesFromStoreProjectionWithoutEviction(t *testing.T) { task := `StreamHandler::write throws UnexpectedValueException "could not be opened: Permission denied" after upgrade, likely due to chmod and file permission handling` fixture := monologDivergentDefaultFixture() @@ -115,7 +170,7 @@ func TestDivergentDefaultOwnerPromotesFromStoreProjectionWithoutEviction(t *test fixture.targets = append(fixture.targets, exploreTarget{node: protected, conceptImplementation: true}) originalCount := len(fixture.targets) reads := 0 - promoted := promoteExploreDivergentDefaultOwner(task, fixture.targets, fixture.store, len(fixture.targets), func(node *graph.Node) string { + promoted := promoteExploreDivergentDefaultOwnerForTest(task, fixture.targets, fixture.store, len(fixture.targets), func(node *graph.Node) string { reads++ require.Equal(t, fixture.childCtor.ID, node.ID) return monologForwardingConstructorSource @@ -160,7 +215,7 @@ func TestDivergentDefaultOwnerRecoversBasePairFromRankedCallableOwner(t *testing fixture.targets = fixture.targets[:1] counted := &countingDivergentDefaultStore{Store: fixture.store} reads := 0 - promoted := promoteExploreDivergentDefaultOwner(task, fixture.targets, counted, 3, func(node *graph.Node) string { + promoted := promoteExploreDivergentDefaultOwnerForTest(task, fixture.targets, counted, 3, func(node *graph.Node) string { reads++ require.Equal(t, fixture.childCtor.ID, node.ID) return monologForwardingConstructorSource @@ -169,9 +224,9 @@ func TestDivergentDefaultOwnerRecoversBasePairFromRankedCallableOwner(t *testing require.Equal(t, []string{fixture.childCtor.ID, fixture.childType.ID, fixture.write.ID}, []string{ promoted[0].node.ID, promoted[1].node.ID, promoted[2].node.ID, }) - require.Equal(t, 2, counted.nodeBatchCalls, "owner discovery and endpoint hydration must each be batched") + require.Equal(t, 3, counted.nodeBatchCalls, "owner, constructor, and endpoint hydration must each be bounded batches") require.Equal(t, 1, counted.fileBatchCalls) - require.Equal(t, 1, counted.inBatchCalls) + require.Equal(t, 2, counted.inBatchCalls, "CALLS and EXTENDS must use separate bounded projections") require.Equal(t, 1, reads) } @@ -182,7 +237,7 @@ func TestDivergentDefaultOwnerAdmissionAtCapacityEvictsOnlyUnprotectedTail(t *te distractorTwo := divergentDefaultTestNode("src/Noise.php::Noise.two", graph.KindMethod, "two", "src/Noise.php", "function two()") fixture.targets = []exploreTarget{{node: fixture.write}, {node: distractorOne}, {node: distractorTwo}} - promoted := promoteExploreDivergentDefaultOwner(task, fixture.targets, fixture.store, len(fixture.targets), divergentDefaultTestSource) + promoted := promoteExploreDivergentDefaultOwnerForTest(task, fixture.targets, fixture.store, len(fixture.targets), divergentDefaultTestSource) require.Len(t, promoted, len(fixture.targets)) require.Equal(t, []string{fixture.childCtor.ID, fixture.childType.ID, fixture.write.ID}, []string{ promoted[0].node.ID, promoted[1].node.ID, promoted[2].node.ID, @@ -191,7 +246,7 @@ func TestDivergentDefaultOwnerAdmissionAtCapacityEvictsOnlyUnprotectedTail(t *te protectedLiteral := exploreTarget{node: distractorOne, sourceLiteral: true} protectedImplementation := exploreTarget{node: distractorTwo, conceptImplementation: true} fixture.targets = []exploreTarget{{node: fixture.write}, protectedLiteral, protectedImplementation} - unchanged := promoteExploreDivergentDefaultOwner(task, fixture.targets, fixture.store, len(fixture.targets), divergentDefaultTestSource) + unchanged := promoteExploreDivergentDefaultOwnerForTest(task, fixture.targets, fixture.store, len(fixture.targets), divergentDefaultTestSource) require.Equal(t, []string{fixture.write.ID, distractorOne.ID, distractorTwo.ID}, []string{ unchanged[0].node.ID, unchanged[1].node.ID, unchanged[2].node.ID, }) @@ -199,18 +254,94 @@ func TestDivergentDefaultOwnerAdmissionAtCapacityEvictsOnlyUnprotectedTail(t *te require.True(t, unchanged[2].conceptImplementation) } +func TestDivergentDefaultOwnerPreservesSourceRangeBaseSeats(t *testing.T) { + task := `StreamHandler::write reports "could not be opened: Permission denied" after a rotating handler applies chmod` + fixture := monologDivergentDefaultFixture() + tail := divergentDefaultTestNode("src/Noise.php::Noise.tail", graph.KindMethod, "tail", "src/Noise.php", "function tail()") + match := exploreDivergentDefaultOwner{ + constructor: fixture.childCtor, + owner: fixture.childType, + baseCtor: fixture.baseCtor, + baseOwner: fixture.baseType, + consumerID: fixture.write.ID, + } + constructor := exploreTarget{node: fixture.childCtor, divergentDefaultOwner: true} + owner := exploreTarget{node: fixture.childType, divergentDefaultType: true} + + for _, test := range []struct { + name string + baseCtor bool + }{ + {name: "constructor", baseCtor: true}, + {name: "owner"}, + } { + t.Run(test.name, func(t *testing.T) { + targets := []exploreTarget{{node: fixture.write}, {node: fixture.baseType}, {node: fixture.baseCtor}, {node: tail}} + protectedID := fixture.baseType.ID + if test.baseCtor { + targets[2].sourceRange = true + protectedID = fixture.baseCtor.ID + } else { + targets[1].sourceRange = true + } + + got, ok := placeExploreDivergentDefaultOwner(task, targets, match, constructor, owner, len(targets)) + require.True(t, ok) + require.Len(t, got, len(targets)) + protectedIndex := exploreTargetIndex(got, protectedID) + require.NotEqual(t, -1, protectedIndex, "exact source-range owner was replaced") + require.True(t, got[protectedIndex].sourceRange) + require.NotEqual(t, -1, exploreTargetIndex(got, fixture.childCtor.ID)) + require.NotEqual(t, -1, exploreTargetIndex(got, fixture.childType.ID)) + require.Equal(t, -1, exploreTargetIndex(got, tail.ID), "ordinary tail should fund the proven child pair") + }) + } +} + +func TestDivergentDefaultOwnerRejectsPromotionWhenSourceRangeBaseAndTailAreProtected(t *testing.T) { + task := `StreamHandler::write reports "could not be opened: Permission denied" after a rotating handler applies chmod` + fixture := monologDivergentDefaultFixture() + tail := divergentDefaultTestNode("src/Noise.php::Noise.tail", graph.KindMethod, "tail", "src/Noise.php", "function tail()") + targets := []exploreTarget{ + {node: fixture.write}, + {node: fixture.baseType}, + {node: fixture.baseCtor, sourceRange: true}, + {node: tail, sourceLiteral: true}, + } + match := exploreDivergentDefaultOwner{ + constructor: fixture.childCtor, + owner: fixture.childType, + baseCtor: fixture.baseCtor, + baseOwner: fixture.baseType, + consumerID: fixture.write.ID, + } + + got, ok := placeExploreDivergentDefaultOwner( + task, + targets, + match, + exploreTarget{node: fixture.childCtor, divergentDefaultOwner: true}, + exploreTarget{node: fixture.childType, divergentDefaultType: true}, + len(targets), + ) + require.False(t, ok) + require.Nil(t, got) + require.True(t, targets[2].sourceRange, "failed promotion mutated the cited base") + require.True(t, targets[3].sourceLiteral, "failed promotion mutated the protected tail") +} + func TestDivergentDefaultOwnerCallableFallbackFailsClosed(t *testing.T) { task := `StreamHandler::write reports "could not be opened: Permission denied" after a rotating handler applies chmod; find the divergent filePermission default and owning type` t.Run("no output capacity", func(t *testing.T) { fixture := monologDivergentDefaultFixture() fixture.targets = fixture.targets[:1] - got := promoteExploreDivergentDefaultOwner(task, fixture.targets, fixture.store, 2, divergentDefaultTestSource) + got := promoteExploreDivergentDefaultOwnerForTest(task, fixture.targets, fixture.store, 2, divergentDefaultTestSource) require.Equal(t, []exploreTarget{fixture.targets[0]}, got) }) t.Run("irrelevant ranked callable", func(t *testing.T) { fixture := monologDivergentDefaultFixture() fixture.targets = fixture.targets[:1] - got := promoteExploreDivergentDefaultOwner("Investigate queue retry scheduling and backoff policy", fixture.targets, fixture.store, 3, divergentDefaultTestSource) + got := promoteExploreDivergentDefaultOwnerForTest("Investigate queue retry scheduling and backoff policy", fixture.targets, fixture.store, 3, divergentDefaultTestSource) require.Equal(t, fixture.write.ID, got[0].node.ID) require.Len(t, got, 1) }) @@ -225,7 +356,7 @@ func TestDivergentDefaultOwnerCallableFallbackFailsClosed(t *testing.T) { `function StreamHandler($filePermission = null)`, ) fixture.store.AddNode(alternate) - got := promoteExploreDivergentDefaultOwner(task, fixture.targets, fixture.store, 3, divergentDefaultTestSource) + got := promoteExploreDivergentDefaultOwnerForTest(task, fixture.targets, fixture.store, 3, divergentDefaultTestSource) require.Equal(t, fixture.write.ID, got[0].node.ID) require.Len(t, got, 1) }) @@ -241,15 +372,15 @@ func TestDivergentDefaultOwnerCallableFallbackFailsClosed(t *testing.T) { "function helper()", )) } - got := promoteExploreDivergentDefaultOwner(task, fixture.targets, fixture.store, 3, divergentDefaultTestSource) + got := promoteExploreDivergentDefaultOwnerForTest(task, fixture.targets, fixture.store, 3, divergentDefaultTestSource) require.Equal(t, fixture.write.ID, got[0].node.ID) require.Len(t, got, 1) }) t.Run("owner hydration exceeds admission budget", func(t *testing.T) { fixture := monologDivergentDefaultFixture() fixture.targets = fixture.targets[:1] - counted := &countingDivergentDefaultStore{Store: fixture.store, nodeBatchDelay: exploreDefaultOwnerFallbackSLO + time.Millisecond} - got := promoteExploreDivergentDefaultOwner(task, fixture.targets, counted, 3, divergentDefaultTestSource) + counted := &countingDivergentDefaultStore{Store: fixture.store, nodeBatchDelay: 2 * exploreDefaultOwnerFallbackSLO} + got := promoteExploreDivergentDefaultOwnerForTest(task, fixture.targets, counted, 3, divergentDefaultTestSource) require.Equal(t, fixture.write.ID, got[0].node.ID) require.Len(t, got, 1) require.Equal(t, 1, counted.nodeBatchCalls) @@ -284,7 +415,7 @@ func TestDivergentDefaultOwnerFailsClosedOnMissingOrInvalidStoreEvidence(t *test t.Run(test.name, func(t *testing.T) { fixture := monologDivergentDefaultFixture() test.mutate(fixture) - got := promoteExploreDivergentDefaultOwner(task, fixture.targets, fixture.store, len(fixture.targets), divergentDefaultTestSource) + got := promoteExploreDivergentDefaultOwnerForTest(task, fixture.targets, fixture.store, len(fixture.targets), divergentDefaultTestSource) require.Equal(t, fixture.write.ID, got[0].node.ID) require.False(t, got[0].divergentDefaultOwner) }) @@ -301,7 +432,7 @@ func TestDivergentDefaultOwnerRejectsAmbiguousOrOversizedInboundProjection(t *te {From: secondCtor.ID, To: fixture.baseCtor.ID, Kind: graph.EdgeCalls}, {From: secondType.ID, To: fixture.baseType.ID, Kind: graph.EdgeExtends}, }) - got := promoteExploreDivergentDefaultOwner(task, fixture.targets, fixture.store, len(fixture.targets), divergentDefaultTestSource) + got := promoteExploreDivergentDefaultOwnerForTest(task, fixture.targets, fixture.store, len(fixture.targets), divergentDefaultTestSource) require.Equal(t, fixture.write.ID, got[0].node.ID) }) t.Run("edge cap", func(t *testing.T) { @@ -314,7 +445,7 @@ func TestDivergentDefaultOwnerRejectsAmbiguousOrOversizedInboundProjection(t *te edges = append(edges, &graph.Edge{From: node.ID, To: fixture.baseCtor.ID, Kind: graph.EdgeCalls}) } fixture.store.AddBatch(nodes, edges) - got := promoteExploreDivergentDefaultOwner(task, fixture.targets, fixture.store, len(fixture.targets), divergentDefaultTestSource) + got := promoteExploreDivergentDefaultOwnerForTest(task, fixture.targets, fixture.store, len(fixture.targets), divergentDefaultTestSource) require.Equal(t, fixture.write.ID, got[0].node.ID) }) } @@ -367,7 +498,7 @@ func TestDivergentDefaultOwnerRequiresExecutableForwardingIntoProvenBaseCall(t * } for _, source := range negative { fixture := monologDivergentDefaultFixture() - got := promoteExploreDivergentDefaultOwner(task, fixture.targets, fixture.store, len(fixture.targets), func(*graph.Node) string { return source }) + got := promoteExploreDivergentDefaultOwnerForTest(task, fixture.targets, fixture.store, len(fixture.targets), func(*graph.Node) string { return source }) require.Equal(t, fixture.write.ID, got[0].node.ID, "non-executable forwarding was accepted: %q", source) } } @@ -406,7 +537,7 @@ func TestHandleExplorePromotesDivergentDefaultOwnerFromSQLitePHPIndex(t *testing require.True(t, hasFixtureEdge(store.GetInEdges(serializableContract.ID), statusType.ID, serializableContract.ID, graph.EdgeImplements), "PHP index must resolve enum conformance") task := `StreamHandler::write reports "could not be opened: Permission denied" after a rotating handler applies chmod; find the divergent filePermission default and owning type` - direct := promoteExploreDivergentDefaultOwner(task, []exploreTarget{{node: write}, {node: baseType}, {node: baseCtor}}, store, 3, func(node *graph.Node) string { + direct := promoteExploreDivergentDefaultOwnerForTest(task, []exploreTarget{{node: write}, {node: baseType}, {node: baseCtor}}, store, 3, func(node *graph.Node) string { return server.manifestSymbolSource(context.Background(), node) }, pinnedDivergentDefaultFallbackSLO) require.Equal(t, childCtor.ID, direct[0].node.ID, "real store projection did not promote child constructor") @@ -415,7 +546,7 @@ func TestHandleExplorePromotesDivergentDefaultOwnerFromSQLitePHPIndex(t *testing // ranked-callable projection runs and its batched SQLite reads must finish // inside the slice. This assertion is about WHICH owner is promoted, not // about how fast the machine is, so the slice is pinned. - fallback := promoteExploreDivergentDefaultOwner(task, []exploreTarget{{node: write}}, store, 3, func(node *graph.Node) string { + fallback := promoteExploreDivergentDefaultOwnerForTest(task, []exploreTarget{{node: write}}, store, 3, func(node *graph.Node) string { return server.manifestSymbolSource(context.Background(), node) }, pinnedDivergentDefaultFallbackSLO) require.Equal(t, childCtor.ID, fallback[0].node.ID, "real store callable-owner fallback did not promote child constructor") @@ -577,7 +708,7 @@ func BenchmarkPromoteExploreDivergentDefaultOwner24(b *testing.B) { b.ReportAllocs() b.ResetTimer() for index := 0; index < b.N; index++ { - if got := promoteExploreDivergentDefaultOwner(task, fixture.targets, fixture.store, len(fixture.targets), divergentDefaultTestSource); !got[0].divergentDefaultOwner { + if got := promoteExploreDivergentDefaultOwnerForTest(task, fixture.targets, fixture.store, len(fixture.targets), divergentDefaultTestSource); !got[0].divergentDefaultOwner { b.Fatal("expected promotion") } } @@ -594,7 +725,7 @@ func BenchmarkPromoteExploreDivergentDefaultOwnerFromCallable24(b *testing.B) { b.ReportAllocs() b.ResetTimer() for index := 0; index < b.N; index++ { - if got := promoteExploreDivergentDefaultOwner(task, fixture.targets, fixture.store, 26, divergentDefaultTestSource); !got[0].divergentDefaultOwner { + if got := promoteExploreDivergentDefaultOwnerForTest(task, fixture.targets, fixture.store, 26, divergentDefaultTestSource); !got[0].divergentDefaultOwner { b.Fatal("expected promotion") } } diff --git a/internal/mcp/explore_exact_name_anchor.go b/internal/mcp/explore_exact_name_anchor.go index 48db85dc..e463dafb 100644 --- a/internal/mcp/explore_exact_name_anchor.go +++ b/internal/mcp/explore_exact_name_anchor.go @@ -17,9 +17,15 @@ const ( exploreExactNameAnchorMaxChars = 64 // A name this widely shared describes a convention (handle, execute, run), // not the task's subject, so it cannot anchor on its own. - exploreExactNameAnchorMaxShared = 8 - exploreExactNameAnchorMaxNodes = 10 - exploreExactNameAnchorOwnerScan = 32 + exploreExactNameAnchorMaxShared = 8 + exploreExactNameAnchorMaxNodes = 10 + exploreExactNameAnchorOwnerScan = 32 + exploreExactNameAnchorOwnerRawScan = 4 * exploreExactNameAnchorOwnerScan + // Case folding is a miss-only recovery lane. Four tokens, three alternate + // indexed spellings each, and four ranked files are hard request-wide caps. + exploreExactNameAnchorCaseFoldMaxTokens = 4 + exploreExactNameAnchorCaseVariantMax = 4 + exploreExactNameAnchorCaseFoldMaxFiles = 4 ) // exploreTaskAnchors is the anchor lane the retrieval side uses: the @@ -42,9 +48,9 @@ func (s *Server) exploreTaskAnchors( } // exploreExactNameAnchors resolves plain task tokens against the graph name -// index. Cost is fixed by construction: at most exploreExactNameAnchorMaxTokens -// name lookups per task, each one sharded-map hit, with no per-candidate store -// query and no source hydration. +// index. Cost is fixed by construction: sixteen exact lookups plus at most +// three spelling variants for four misses and four ranked-file reads. The +// recovery lane never scans the corpus or hydrates source. func (s *Server) exploreExactNameAnchors( ctx context.Context, task string, @@ -53,20 +59,24 @@ func (s *Server) exploreExactNameAnchors( scope query.QueryOptions, slots int, ) []exploreSyntacticAnchor { - if s == nil || s.graph == nil || slots <= 0 || ctx.Err() != nil { + if s == nil || slots <= 0 || ctx.Err() != nil { + return nil + } + reader := s.readerFor(ctx) + if reader == nil { return nil } tokens := exploreExactNameAnchorTokens(task) if len(tokens) == 0 { return nil } - pooledFiles := exploreRankedPoolFiles(ordinary) + lookup := newExploreExactNameAnchorLookup(ctx, reader, ordinary) out := make([]exploreSyntacticAnchor, 0, slots) for _, token := range tokens { if ctx.Err() != nil { break } - nodes := s.exploreExactNameAnchorNodes(ctx, token, scope, pooledFiles) + nodes := s.exploreExactNameAnchorNodes(ctx, token, scope, lookup) if len(nodes) == 0 { continue } @@ -94,45 +104,304 @@ func exploreExactNameAnchorDuplicate(anchor exploreSyntacticAnchor, groups ...[] return false } -// exploreExactNameAnchorNodes returns the localizable declarations whose name is -// the task token verbatim. Case matters: `Mount` and `mount` are different -// symbols, and a case-insensitive match would readmit exactly the prose noise -// this lane exists to exclude. +type exploreExactNameFileNodes struct { + nodes []*graph.Node + complete bool +} + +type exploreExactNameAnchorLookup struct { + reader graph.Reader + fileBudget *localizationFileRequestBudget + rankedNodes []*graph.Node + rankedFiles []string + pooledFiles map[string]struct{} + fallbackTokens int + fileNodes map[string]exploreExactNameFileNodes +} + +func newExploreExactNameAnchorLookup( + ctx context.Context, + reader graph.Reader, + ordinary []*rerank.Candidate, +) *exploreExactNameAnchorLookup { + lookup := &exploreExactNameAnchorLookup{ + reader: reader, + fileBudget: localizationFileBudgetFor(ctx), + pooledFiles: exploreRankedPoolFiles(ordinary), + fileNodes: make(map[string]exploreExactNameFileNodes, exploreExactNameAnchorCaseFoldMaxFiles), + } + seenFiles := make(map[string]struct{}, exploreExactNameAnchorCaseFoldMaxFiles) + for _, candidate := range ordinary { + if candidate == nil || candidate.Node == nil { + continue + } + lookup.rankedNodes = append(lookup.rankedNodes, candidate.Node) + path := candidate.Node.FilePath + if path == "" || len(lookup.rankedFiles) == exploreExactNameAnchorCaseFoldMaxFiles { + continue + } + if _, duplicate := seenFiles[path]; duplicate { + continue + } + seenFiles[path] = struct{}{} + lookup.rankedFiles = append(lookup.rankedFiles, path) + } + return lookup +} + +func (lookup *exploreExactNameAnchorLookup) boundedFileNodes( + ctx context.Context, + path string, + scope graph.LocalizationNodeScope, +) ([]*graph.Node, bool) { + if lookup == nil { + return nil, false + } + if cached, ok := lookup.fileNodes[path]; ok { + return cached.nodes, cached.complete + } + page, complete := boundedLocalizationFileNodes( + ctx, lookup.reader, lookup.fileBudget, path, scope, localizationFileNodeLimit, + ) + cached := exploreExactNameFileNodes{complete: complete} + if complete { + cached.nodes = page.Nodes + } + lookup.fileNodes[path] = cached + return cached.nodes, cached.complete +} + +// exploreExactNameCaseVariants returns a small deterministic set of spellings +// that exact name indexes commonly contain. It never attempts the combinatorial +// interior-case search needed to invent arbitrary camelCase. +func exploreExactNameCaseVariants(token string) []string { + if token == "" { + return nil + } + variants := make([]string, 0, exploreExactNameAnchorCaseVariantMax) + seen := make(map[string]struct{}, exploreExactNameAnchorCaseVariantMax) + add := func(value string) { + if value == "" || len(variants) == exploreExactNameAnchorCaseVariantMax { + return + } + if _, exists := seen[value]; exists { + return + } + seen[value] = struct{}{} + variants = append(variants, value) + } + add(token) + runes := []rune(token) + if len(runes) == 0 { + return variants + } + lowerFirst := append([]rune(nil), runes...) + lowerFirst[0] = unicode.ToLower(lowerFirst[0]) + add(string(lowerFirst)) + upperFirst := append([]rune(nil), runes...) + upperFirst[0] = unicode.ToUpper(upperFirst[0]) + add(string(upperFirst)) + + lower := strings.ToLower(token) + upper := strings.ToUpper(token) + if token == lower || token == upper || token == string(upperFirst) { + add(lower) + title := []rune(lower) + if len(title) > 0 { + title[0] = unicode.ToUpper(title[0]) + add(string(title)) + } + add(upper) + } + return variants +} + +// caseFoldedMatches performs bounded miss recovery. Exact indexed spelling +// variants come first. Arbitrary interior-case matches are accepted only when +// the independent ordinary ranking already selected the node or its file. +func (lookup *exploreExactNameAnchorLookup) caseFoldedMatches( + ctx context.Context, + token string, + scope graph.LocalizationNodeScope, + eligible func(*graph.Node) bool, +) ([]*graph.Node, int) { + if lookup == nil || lookup.reader == nil || ctx.Err() != nil || + lookup.fallbackTokens == exploreExactNameAnchorCaseFoldMaxTokens { + return nil, 0 + } + lookup.fallbackTokens++ + matches := make([]*graph.Node, 0, exploreExactNameAnchorMaxNodes) + seen := make(map[string]struct{}, exploreExactNameAnchorMaxNodes) + add := func(node *graph.Node) { + if node == nil || node.ID == "" || len(matches) == exploreExactNameAnchorMaxNodes || + !strings.EqualFold(node.Name, token) || eligible == nil || !eligible(node) { + return + } + if _, duplicate := seen[node.ID]; duplicate { + return + } + seen[node.ID] = struct{}{} + matches = append(matches, node) + } + + variants := exploreExactNameCaseVariants(token) + shared := 0 + for _, variant := range variants[1:] { + if ctx.Err() != nil { + return matches, shared + } + page, ok := boundedLocalizationExactName( + ctx, + lookup.reader, + variant, + scope, + exploreExactNameAnchorMaxNodes, + ) + if !ok { + return matches, shared + } + shared += page.Total + for _, node := range page.Nodes { + add(node) + } + } + if len(matches) > 0 { + return matches, shared + } + + for _, node := range lookup.rankedNodes { + add(node) + } + if len(matches) > 0 { + return matches, len(matches) + } + + for _, path := range lookup.rankedFiles { + if ctx.Err() != nil { + return nil, 0 + } + nodes, complete := lookup.boundedFileNodes(ctx, path, scope) + if !complete { + return nil, 0 + } + for _, node := range nodes { + add(node) + } + if len(matches) > 0 { + break + } + } + return matches, len(matches) +} + +// exploreExactNameAnchorNodes returns localizable declarations whose indexed +// name matches the task token. An exact case bucket is authoritative only when +// it contains a declaration eligible for this request's repo/session scope; +// foreign homonyms cannot suppress bounded case-fold recovery. func (s *Server) exploreExactNameAnchorNodes( ctx context.Context, token string, scope query.QueryOptions, - pooledFiles map[string]struct{}, + lookup *exploreExactNameAnchorLookup, ) []*graph.Node { - matches := s.graph.FindNodesByName(token) - if len(matches) == 0 { + eligibleNode := func(node *graph.Node) bool { + return node != nil && exploreSyntacticAnchorEligibleNode(node) && + scope.ScopeAllows(node) && s.nodeInSessionScope(ctx, node) + } + nodeScope := s.localizationNodeScope( + ctx, scope, graph.KindFunction, graph.KindMethod, graph.KindType, graph.KindMacro, + ) + selectMatches := func(matches []*graph.Node, exact bool, shared int) ([]*graph.Node, bool) { + eligibleShared := 0 + eligible := make([]*graph.Node, 0, exploreExactNameAnchorMaxNodes) + pooled := make([]*graph.Node, 0, exploreExactNameAnchorMaxNodes) + for _, node := range matches { + if node == nil || (exact && node.Name != token) || (!exact && !strings.EqualFold(node.Name, token)) || + !eligibleNode(node) { + continue + } + eligibleShared++ + if len(eligible) < exploreExactNameAnchorMaxNodes { + eligible = append(eligible, node) + } + if _, ranked := lookup.pooledFiles[node.FilePath]; ranked && len(pooled) < exploreExactNameAnchorMaxNodes { + pooled = append(pooled, node) + } + } + if shared > exploreExactNameAnchorMaxShared { + // The ranked pool is independent evidence that one of the homonyms is + // the one the task means; without it an ambiguous name stays unanchored. + if len(pooled) > 0 { + return pooled, true + } + return lookup.preferredFileMatches(ctx, token, exact, nodeScope, eligibleNode), true + } + return eligible, eligibleShared > 0 + } + + page, ok := boundedLocalizationExactName( + ctx, + lookup.reader, + token, + nodeScope, + exploreExactNameAnchorMaxNodes, + ) + if ok { + if exact, found := selectMatches(page.Nodes, true, page.Total); found { + return exact + } + } + matches, shared := lookup.caseFoldedMatches(ctx, token, nodeScope, eligibleNode) + selected, _ := selectMatches(matches, false, shared) + return selected +} + +// preferredFileMatches preserves the independent ranked-file evidence used to +// disambiguate a highly shared name even when its declaration sorts beyond the +// bounded exact-name page. rankedFiles is request-capped at four and fileNodes +// is shared with case-fold recovery, so this cannot grow into a corpus scan. +func (lookup *exploreExactNameAnchorLookup) preferredFileMatches( + ctx context.Context, + token string, + exact bool, + scope graph.LocalizationNodeScope, + eligible func(*graph.Node) bool, +) []*graph.Node { + if lookup == nil || lookup.reader == nil || eligible == nil { return nil } - shared := 0 - eligible := make([]*graph.Node, 0, exploreExactNameAnchorMaxNodes) - pooled := make([]*graph.Node, 0, exploreExactNameAnchorMaxNodes) - for _, node := range matches { - if node == nil || node.Name != token { - continue + matches := make([]*graph.Node, 0, exploreExactNameAnchorMaxNodes) + seen := make(map[string]struct{}, exploreExactNameAnchorMaxNodes) + for _, path := range lookup.rankedFiles { + if ctx.Err() != nil { + return nil } - shared++ - if !exploreSyntacticAnchorEligibleNode(node) || !scope.ScopeAllows(node) || - !s.nodeInSessionScope(ctx, node) { - continue + if len(matches) == exploreExactNameAnchorMaxNodes { + break } - if len(eligible) < exploreExactNameAnchorMaxNodes { - eligible = append(eligible, node) + nodes, complete := lookup.boundedFileNodes(ctx, path, scope) + if !complete { + return nil } - if _, ranked := pooledFiles[node.FilePath]; ranked && len(pooled) < exploreExactNameAnchorMaxNodes { - pooled = append(pooled, node) + for _, node := range nodes { + nameMatches := node != nil && strings.EqualFold(node.Name, token) + if exact { + nameMatches = node != nil && node.Name == token + } + if !nameMatches || !eligible(node) { + continue + } + if _, duplicate := seen[node.ID]; duplicate { + continue + } + seen[node.ID] = struct{}{} + matches = append(matches, node) + if len(matches) == exploreExactNameAnchorMaxNodes { + break + } } } - if shared > exploreExactNameAnchorMaxShared { - // The ranked pool is independent evidence that one of the homonyms is the - // one the task means; without it an ambiguous name stays unanchored. - return pooled - } - return eligible + return matches } func exploreRankedPoolFiles(candidates []*rerank.Candidate) map[string]struct{} { @@ -161,10 +430,11 @@ func exploreExactNameAnchorTokens(task string) []string { if !exploreExactNameAnchorToken(token) { continue } - if _, duplicate := seen[token]; duplicate { + key := strings.ToLower(token) + if _, duplicate := seen[key]; duplicate { continue } - seen[token] = struct{}{} + seen[key] = struct{}{} out = append(out, token) if len(out) == exploreExactNameAnchorMaxTokens { return out @@ -254,6 +524,7 @@ func exploreIdentifierSegment(segment string) bool { func (s *Server) exploreExactAnchorCandidate( ctx context.Context, anchor exploreSyntacticAnchor, + ordinary []*rerank.Candidate, scope query.QueryOptions, usedIDs, usedFiles map[string]struct{}, ) *rerank.Candidate { @@ -266,7 +537,7 @@ func (s *Server) exploreExactAnchorCandidate( return got } } - return s.exploreExactQualifiedAnchorCandidate(ctx, anchor, scope, usedIDs, usedFiles) + return s.exploreExactQualifiedAnchorCandidate(ctx, anchor, ordinary, scope, usedIDs, usedFiles) } // exploreQualifiedAnchorOwnerCandidate resolves the declaring type named by a @@ -281,7 +552,11 @@ func (s *Server) exploreQualifiedAnchorOwnerCandidate( scope query.QueryOptions, usedIDs map[string]struct{}, ) *rerank.Candidate { - if s == nil || s.graph == nil || member == nil || anchor.qualifiedName == "" || ctx.Err() != nil { + if s == nil || member == nil || anchor.qualifiedName == "" || ctx.Err() != nil { + return nil + } + reader := s.readerFor(ctx) + if reader == nil { return nil } dot := strings.LastIndexByte(anchor.qualifiedName, '.') @@ -295,13 +570,23 @@ func (s *Server) exploreQualifiedAnchorOwnerCandidate( if owner == "" || owner == member.Name { return nil } + page, ok := boundedLocalizationExactName( + ctx, + reader, + owner, + s.localizationNodeScope(ctx, scope, graph.KindType, graph.KindInterface), + exploreExactNameAnchorOwnerRawScan, + ) + if !ok { + return nil + } var best *graph.Node - scanned := 0 - for _, node := range s.graph.FindNodesByName(owner) { - if scanned == exploreExactNameAnchorOwnerScan { + rawScanned, eligibleScanned := 0, 0 + for _, node := range page.Nodes { + if rawScanned == exploreExactNameAnchorOwnerRawScan { break } - scanned++ + rawScanned++ if node == nil || node.Name != owner || (node.Kind != graph.KindType && node.Kind != graph.KindInterface) { continue @@ -309,6 +594,10 @@ func (s *Server) exploreQualifiedAnchorOwnerCandidate( if !scope.ScopeAllows(node) || !s.nodeInSessionScope(ctx, node) { continue } + if eligibleScanned == exploreExactNameAnchorOwnerScan { + break + } + eligibleScanned++ if _, used := usedIDs[node.ID]; used { continue } diff --git a/internal/mcp/explore_exact_name_anchor_test.go b/internal/mcp/explore_exact_name_anchor_test.go index 375113cb..32b83a9c 100644 --- a/internal/mcp/explore_exact_name_anchor_test.go +++ b/internal/mcp/explore_exact_name_anchor_test.go @@ -3,10 +3,12 @@ package mcp import ( "context" "fmt" + "path/filepath" "strings" "testing" "github.com/zzet/gortex/internal/graph" + "github.com/zzet/gortex/internal/graph/store_sqlite" "github.com/zzet/gortex/internal/query" "github.com/zzet/gortex/internal/search/rerank" ) @@ -35,7 +37,7 @@ func TestExploreTaskAnchorsAdmitLowercaseProseSymbolName(t *testing.T) { t.Fatalf("anchor compact = %q, want vprintf", anchors[0].compact) } got := server.exploreExactAnchorCandidate( - context.Background(), anchors[0], query.QueryOptions{}, + context.Background(), anchors[0], nil, query.QueryOptions{}, map[string]struct{}{}, map[string]struct{}{}, ) if got == nil || got.Node == nil || got.Node.ID != vprintf.ID { @@ -65,14 +67,19 @@ func TestReserveExploreSyntacticAnchorCandidatesKeepsGraphResolvedAnchor(t *test } func TestExploreTaskAnchorsSkipAmbiguousPlainName(t *testing.T) { - nodes := make([]*graph.Node, 0, 9) - for index := 0; index < 9; index++ { + nodes := make([]*graph.Node, 0, 33) + for index := 0; index < 32; index++ { nodes = append(nodes, &graph.Node{ - ID: fmt.Sprintf("pkg/h%d.go::handle", index), + ID: fmt.Sprintf("pkg/a-%02d.go::handle", index), Name: "handle", Kind: graph.KindFunction, - FilePath: fmt.Sprintf("pkg/h%d.go", index), + FilePath: fmt.Sprintf("pkg/a-%02d.go", index), }) } + rankedHandle := &graph.Node{ + ID: "pkg/z-ranked.go::handle", Name: "handle", Kind: graph.KindFunction, + FilePath: "pkg/z-ranked.go", + } + nodes = append(nodes, rankedHandle) server := exactNameAnchorServer(t, nodes...) task := "requests handle poorly under sustained load" if anchors := server.exploreTaskAnchors(context.Background(), task, nil, query.QueryOptions{}); len(anchors) != 0 { @@ -80,14 +87,15 @@ func TestExploreTaskAnchorsSkipAmbiguousPlainName(t *testing.T) { } pool := []*rerank.Candidate{{Node: &graph.Node{ - ID: "pkg/h3.go::serve", Name: "serve", Kind: graph.KindFunction, FilePath: "pkg/h3.go", + ID: "pkg/z-ranked.go::serve", Name: "serve", Kind: graph.KindFunction, + FilePath: rankedHandle.FilePath, }}} anchors := server.exploreTaskAnchors(context.Background(), task, pool, query.QueryOptions{}) if len(anchors) != 1 || anchors[0].compact != "handle" { t.Fatalf("anchors = %#v, want the ranked-pool file to disambiguate handle", anchors) } - if len(anchors[0].exactNodes) != 1 || anchors[0].exactNodes[0].FilePath != "pkg/h3.go" { - t.Fatalf("exact nodes = %#v, want only the ranked-pool declaration", anchors[0].exactNodes) + if len(anchors[0].exactNodes) != 1 || anchors[0].exactNodes[0].ID != rankedHandle.ID { + t.Fatalf("exact nodes = %#v, want ranked declaration %q beyond the bounded page", anchors[0].exactNodes, rankedHandle.ID) } } @@ -125,6 +133,13 @@ func TestExploreTaskAnchorsKeepCodeShapedAnchorsAhead(t *testing.T) { } } +func TestExploreExactNameAnchorTokensDeduplicateCaseFoldedSpellings(t *testing.T) { + got := exploreExactNameAnchorTokens("Mount mount MOUNT latch") + if fmt.Sprint(got) != fmt.Sprint([]string{"Mount", "latch"}) { + t.Fatalf("tokens = %#v, want first spelling plus latch", got) + } +} + func TestExploreExactNameAnchorTokensAreBoundedAndDeduplicated(t *testing.T) { words := []string{ "alpha", "bravo", "delta", "gamma", "kappa", "sigma", "theta", "omega", @@ -186,7 +201,7 @@ func TestExploreDottedQualifiedMentionResolvesMemberAndOwner(t *testing.T) { t.Fatalf("anchors = %#v, want the dotted qualified member", anchors) } got := server.exploreExactAnchorCandidate( - context.Background(), anchors[0], query.QueryOptions{}, + context.Background(), anchors[0], nil, query.QueryOptions{}, map[string]struct{}{}, map[string]struct{}{}, ) if got == nil || got.Node == nil || got.Node.ID != member.ID { @@ -199,3 +214,377 @@ func TestExploreDottedQualifiedMentionResolvesMemberAndOwner(t *testing.T) { t.Fatalf("owner candidate = %#v, want the owner declared beside the member", ownerCandidate) } } + +func forEachExactNameAnchorStore(t *testing.T, test func(*testing.T, graph.Store)) { + t.Helper() + t.Run("memory", func(t *testing.T) { + test(t, graph.New()) + }) + t.Run("sqlite", func(t *testing.T) { + store, err := store_sqlite.Open(filepath.Join(t.TempDir(), "graph.sqlite")) + if err != nil { + t.Fatalf("open SQLite store: %v", err) + } + t.Cleanup(func() { + if err := store.Close(); err != nil { + t.Errorf("close SQLite store: %v", err) + } + }) + test(t, store) + }) +} + +func TestExploreTaskAnchorsRecoverCaseFoldedExactNameAcrossStores(t *testing.T) { + forEachExactNameAnchorStore(t, func(t *testing.T, store graph.Store) { + mount := &graph.Node{ + ID: "pkg/router.go::Mount", Name: "Mount", + Kind: graph.KindFunction, FilePath: "pkg/router.go", StartLine: 12, + } + store.AddNode(mount) + server := &Server{graph: store} + anchors := server.exploreTaskAnchors( + context.Background(), "mount", nil, query.QueryOptions{}, + ) + if len(anchors) != 1 || anchors[0].compact != "mount" { + t.Fatalf("anchors = %#v, want lowercase prose to recover Mount", anchors) + } + if len(anchors[0].exactNodes) != 1 || anchors[0].exactNodes[0].ID != mount.ID { + t.Fatalf("exact nodes = %#v, want %q", anchors[0].exactNodes, mount.ID) + } + }) +} + +func TestExploreTaskAnchorsKeepExactCaseBucketAuthoritativeAcrossStores(t *testing.T) { + forEachExactNameAnchorStore(t, func(t *testing.T, store graph.Store) { + lower := &graph.Node{ + ID: "pkg/lower.go::mount", Name: "mount", + Kind: graph.KindFunction, FilePath: "pkg/lower.go", StartLine: 8, + } + upper := &graph.Node{ + ID: "pkg/upper.go::Mount", Name: "Mount", + Kind: graph.KindFunction, FilePath: "pkg/upper.go", StartLine: 9, + } + store.AddNode(lower) + store.AddNode(upper) + server := &Server{graph: store} + anchors := server.exploreTaskAnchors( + context.Background(), "mount", nil, query.QueryOptions{}, + ) + if len(anchors) != 1 || len(anchors[0].exactNodes) != 1 { + t.Fatalf("anchors = %#v, want one exact-case declaration", anchors) + } + if got := anchors[0].exactNodes[0].ID; got != lower.ID { + t.Fatalf("exact node = %q, want authoritative exact-case node %q", got, lower.ID) + } + }) +} + +func TestExploreTaskAnchorsCaseFoldedFallbackSkipsOutOfScopeVariant(t *testing.T) { + foreign := &graph.Node{ + ID: "foreign/router.go::Mount", Name: "Mount", Kind: graph.KindFunction, + FilePath: "foreign/router.go", RepoPrefix: "foreign", WorkspaceID: "foreign", StartLine: 4, + } + local := &graph.Node{ + ID: "local/router.go::MOUNT", Name: "MOUNT", Kind: graph.KindFunction, + FilePath: "local/router.go", RepoPrefix: "local", WorkspaceID: "local", StartLine: 8, + } + ranked := &graph.Node{ + ID: "local/router.go::serve", Name: "serve", Kind: graph.KindFunction, + FilePath: "local/router.go", RepoPrefix: "local", WorkspaceID: "local", StartLine: 2, + } + server := exactNameAnchorServer(t, foreign, local, ranked) + anchors := server.exploreTaskAnchors( + context.Background(), "mount", []*rerank.Candidate{{Node: ranked}}, + query.QueryOptions{WorkspaceID: "local", RepoAllow: map[string]bool{"local": true}}, + ) + if len(anchors) != 1 || len(anchors[0].exactNodes) != 1 || anchors[0].exactNodes[0].ID != local.ID { + t.Fatalf("anchors = %#v, want scoped fallback %q", anchors, local.ID) + } +} + +func TestExploreTaskAnchorsCaseFoldedFallbackHonorsSessionScope(t *testing.T) { + foreign := &graph.Node{ + ID: "foreign/router.go::Mount", Name: "Mount", Kind: graph.KindFunction, + FilePath: "foreign/router.go", WorkspaceID: "foreign", StartLine: 4, + } + local := &graph.Node{ + ID: "local/router.go::MOUNT", Name: "MOUNT", Kind: graph.KindFunction, + FilePath: "local/router.go", WorkspaceID: "local", StartLine: 8, + } + ranked := &graph.Node{ + ID: "local/router.go::serve", Name: "serve", Kind: graph.KindFunction, + FilePath: "local/router.go", WorkspaceID: "local", StartLine: 2, + } + server := exactNameAnchorServer(t, foreign, local, ranked) + server.session = &sessionState{scopeResolved: true, scopeBound: true, scopeWorkspaceID: "local"} + anchors := server.exploreTaskAnchors( + context.Background(), "mount", []*rerank.Candidate{{Node: ranked}}, query.QueryOptions{}, + ) + if len(anchors) != 1 || len(anchors[0].exactNodes) != 1 || anchors[0].exactNodes[0].ID != local.ID { + t.Fatalf("anchors = %#v, want session-scoped fallback %q", anchors, local.ID) + } +} + +func TestExploreTaskAnchorsPushSessionScopeBeforeBoundedLookup(t *testing.T) { + nodes := make([]*graph.Node, 0, 162) + for index := 0; index < 160; index++ { + nodes = append(nodes, &graph.Node{ + ID: fmt.Sprintf("a-foreign/%03d.go::MOUNT", index), Name: "MOUNT", + Kind: graph.KindFunction, FilePath: fmt.Sprintf("a-foreign/%03d.go", index), + RepoPrefix: "foreign", WorkspaceID: "foreign", + }) + } + local := &graph.Node{ + ID: "z-local/router.go::MOUNT", Name: "MOUNT", Kind: graph.KindFunction, + FilePath: "z-local/router.go", RepoPrefix: "local", WorkspaceID: "local", + } + ranked := &graph.Node{ + ID: "z-local/router.go::serve", Name: "serve", Kind: graph.KindFunction, + FilePath: local.FilePath, RepoPrefix: "local", WorkspaceID: "local", + } + nodes = append(nodes, local, ranked) + server := exactNameAnchorServer(t, nodes...) + server.session = &sessionState{ + scopeResolved: true, scopeBound: true, scopeWorkspaceID: "local", + scopeRepoAllow: map[string]bool{"local": true}, + } + anchors := server.exploreTaskAnchors( + context.Background(), "mount", []*rerank.Candidate{{Node: ranked}}, query.QueryOptions{}, + ) + if len(anchors) != 1 || len(anchors[0].exactNodes) != 1 || anchors[0].exactNodes[0].ID != local.ID { + t.Fatalf("anchors = %#v, want session-local declaration %q beyond foreign homonyms", anchors, local.ID) + } +} + +func TestExploreTaskAnchorsRecoverInteriorCaseOnlyFromRankedFile(t *testing.T) { + target := &graph.Node{ + ID: "pkg/render.go::buildContent", Name: "buildContent", + Kind: graph.KindFunction, FilePath: "pkg/render.go", StartLine: 31, + } + ranked := &graph.Node{ + ID: "pkg/render.go::render", Name: "render", + Kind: graph.KindFunction, FilePath: "pkg/render.go", StartLine: 7, + } + server := exactNameAnchorServer(t, target, ranked) + ordinary := []*rerank.Candidate{{Node: ranked}} + anchors := server.exploreTaskAnchors( + context.Background(), "buildcontent", ordinary, query.QueryOptions{}, + ) + if len(anchors) != 1 || len(anchors[0].exactNodes) != 1 || anchors[0].exactNodes[0].ID != target.ID { + t.Fatalf("anchors = %#v, want ranked-file evidence for %q", anchors, target.ID) + } + if anchors := server.exploreTaskAnchors( + context.Background(), "buildcontent", nil, query.QueryOptions{}, + ); len(anchors) != 0 { + t.Fatalf("anchors = %#v, want no global interior-case guess without ranking evidence", anchors) + } +} + +func TestExploreExactNameCaseVariantsAreBoundedAndUnicodeSafe(t *testing.T) { + for _, test := range []struct { + token string + want []string + }{ + {token: "mount", want: []string{"mount", "Mount", "MOUNT"}}, + {token: "über", want: []string{"über", "Über", "ÜBER"}}, + } { + got := exploreExactNameCaseVariants(test.token) + if len(got) > exploreExactNameAnchorCaseVariantMax { + t.Fatalf("variants(%q) = %#v, exceeds cap %d", test.token, got, exploreExactNameAnchorCaseVariantMax) + } + if fmt.Sprint(got) != fmt.Sprint(test.want) { + t.Fatalf("variants(%q) = %#v, want %#v", test.token, got, test.want) + } + } +} + +type exactNameAnchorCountingReader struct { + graph.Reader + nameLookups []string + fileLookups []string + fileLimits []int + fileScopes []graph.LocalizationNodeScope + fileProjection func(context.Context, string, graph.LocalizationNodeScope, int) (graph.BoundedNodeProjection, error) +} + +func (reader *exactNameAnchorCountingReader) FindNodesByName(string) []*graph.Node { + panic("localization must not use the legacy unbounded exact-name reader") +} + +func (reader *exactNameAnchorCountingReader) FindNodesByNameBounded( + ctx context.Context, + name string, + scope graph.LocalizationNodeScope, + limit int, +) (graph.BoundedNodeProjection, error) { + reader.nameLookups = append(reader.nameLookups, name) + bounded, ok := reader.Reader.(graph.BoundedExactNameReader) + if !ok { + return graph.BoundedNodeProjection{}, graph.ErrBoundedLocalizationUnavailable + } + return bounded.FindNodesByNameBounded(ctx, name, scope, limit) +} + +func (*exactNameAnchorCountingReader) GetFileNodes(string) []*graph.Node { + panic("ranked-file localization must not use full-row GetFileNodes") +} + +func (reader *exactNameAnchorCountingReader) FindFileNodesBounded( + ctx context.Context, + path string, + scope graph.LocalizationNodeScope, + limit int, +) (graph.BoundedNodeProjection, error) { + reader.fileLookups = append(reader.fileLookups, path) + reader.fileLimits = append(reader.fileLimits, limit) + reader.fileScopes = append(reader.fileScopes, scope) + if reader.fileProjection != nil { + return reader.fileProjection(ctx, path, scope, limit) + } + bounded, ok := reader.Reader.(graph.BoundedFileNodeReader) + if !ok { + return graph.BoundedNodeProjection{}, graph.ErrBoundedLocalizationUnavailable + } + return bounded.FindFileNodesBounded(ctx, path, scope, limit) +} + +func (*exactNameAnchorCountingReader) FindNodesByNameContaining(string, int) []*graph.Node { + panic("case-folded exact-name recovery must not scan names") +} + +func (*exactNameAnchorCountingReader) AllNodes() []*graph.Node { + panic("case-folded exact-name recovery must not scan the corpus") +} + +func TestExploreTaskAnchorsCaseFoldedRecoveryHasRequestWideBudgets(t *testing.T) { + base := graph.New() + ordinary := make([]*rerank.Candidate, 0, 6) + for index := 0; index < 6; index++ { + node := &graph.Node{ + ID: fmt.Sprintf("pkg/f%d.go::candidate%d", index, index), + Name: fmt.Sprintf("candidate%d", index), Kind: graph.KindFunction, + FilePath: fmt.Sprintf("pkg/f%d.go", index), StartLine: index + 1, + } + base.AddNode(node) + ordinary = append(ordinary, &rerank.Candidate{Node: node}) + } + reader := &exactNameAnchorCountingReader{Reader: base} + view := graph.NewOverlaidView(reader, graph.NewOverlayLayer()) + ctx := WithOverlayView(context.Background(), view) + server := &Server{graph: base} + words := []string{ + "alpha", "bravo", "delta", "gamma", "kappa", "sigma", "theta", "omega", + "cargo", "flare", "gauge", "hover", "index", "joint", "knurl", "latch", + "mount", "nudge", "orbit", "plumb", + } + if anchors := server.exploreTaskAnchors( + ctx, strings.Join(words, " "), ordinary, query.QueryOptions{}, + ); len(anchors) != 0 { + t.Fatalf("anchors = %#v, want no matches", anchors) + } + maxNameLookups := exploreExactNameAnchorMaxTokens + + exploreExactNameAnchorCaseFoldMaxTokens*(exploreExactNameAnchorCaseVariantMax-1) + if len(reader.nameLookups) <= exploreExactNameAnchorMaxTokens || len(reader.nameLookups) > maxNameLookups { + t.Fatalf("name lookups = %d (%#v), want fallback use bounded to (%d, %d]", + len(reader.nameLookups), reader.nameLookups, exploreExactNameAnchorMaxTokens, maxNameLookups) + } + if len(reader.fileLookups) != exploreExactNameAnchorCaseFoldMaxFiles { + t.Fatalf("file lookups = %d (%#v), want request-wide cap %d", + len(reader.fileLookups), reader.fileLookups, exploreExactNameAnchorCaseFoldMaxFiles) + } + totalLimit := 0 + for index, limit := range reader.fileLimits { + if limit <= 0 || limit > localizationFileNodeLimit { + t.Fatalf("file limit[%d] = %d, want (0, %d]", index, limit, localizationFileNodeLimit) + } + totalLimit += limit + scope := reader.fileScopes[index] + if !scope.ExcludeTests || len(scope.Kinds) != 4 || + !scope.Kinds[graph.KindFunction] || !scope.Kinds[graph.KindMethod] || + !scope.Kinds[graph.KindType] || !scope.Kinds[graph.KindMacro] { + t.Fatalf("file scope[%d] = %#v, want production function/method/type/macro declarations", index, scope) + } + } + if totalLimit > localizationFileRequestLimit { + t.Fatalf("aggregate file limit = %d, want <= %d", totalLimit, localizationFileRequestLimit) + } +} + +type exactNameOnlyBoundedReader struct { + graph.Reader +} + +func (reader *exactNameOnlyBoundedReader) FindNodesByNameBounded( + ctx context.Context, + name string, + scope graph.LocalizationNodeScope, + limit int, +) (graph.BoundedNodeProjection, error) { + return reader.Reader.(graph.BoundedExactNameReader).FindNodesByNameBounded(ctx, name, scope, limit) +} + +func TestExploreTaskAnchorsRankedFileProjectionFailsClosed(t *testing.T) { + target := &graph.Node{ + ID: "pkg/render.go::buildContent", Name: "buildContent", + Kind: graph.KindFunction, FilePath: "pkg/render.go", StartLine: 31, + } + ranked := &graph.Node{ + ID: "pkg/render.go::render", Name: "render", + Kind: graph.KindFunction, FilePath: "pkg/render.go", StartLine: 7, + } + tests := []struct { + name string + readerBase func(*graph.Graph) graph.Reader + projection func(context.CancelFunc) func(context.Context, string, graph.LocalizationNodeScope, int) (graph.BoundedNodeProjection, error) + }{ + { + name: "truncated", + readerBase: func(base *graph.Graph) graph.Reader { return base }, + projection: func(context.CancelFunc) func(context.Context, string, graph.LocalizationNodeScope, int) (graph.BoundedNodeProjection, error) { + return func(context.Context, string, graph.LocalizationNodeScope, int) (graph.BoundedNodeProjection, error) { + return graph.BoundedNodeProjection{Nodes: []*graph.Node{target}, Total: 2, Truncated: true}, nil + } + }, + }, + { + name: "unsupported", + readerBase: func(base *graph.Graph) graph.Reader { + return &exactNameOnlyBoundedReader{Reader: base} + }, + }, + { + name: "canceled", + readerBase: func(base *graph.Graph) graph.Reader { return base }, + projection: func(cancel context.CancelFunc) func(context.Context, string, graph.LocalizationNodeScope, int) (graph.BoundedNodeProjection, error) { + return func(context.Context, string, graph.LocalizationNodeScope, int) (graph.BoundedNodeProjection, error) { + cancel() + return graph.BoundedNodeProjection{}, context.Canceled + } + }, + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + base := graph.New() + base.AddNode(target) + base.AddNode(ranked) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + reader := &exactNameAnchorCountingReader{Reader: test.readerBase(base)} + if test.projection != nil { + reader.fileProjection = test.projection(cancel) + } + ctx = WithOverlayView(ctx, graph.NewOverlaidView(reader, graph.NewOverlayLayer())) + server := &Server{graph: base} + anchors := server.exploreTaskAnchors( + ctx, "buildcontent", []*rerank.Candidate{{Node: ranked}}, query.QueryOptions{}, + ) + if len(anchors) != 0 { + t.Fatalf("anchors = %#v, want incomplete ranked-file lane to fail closed", anchors) + } + if len(reader.fileLookups) != 1 { + t.Fatalf("file lookups = %#v, want one bounded attempt", reader.fileLookups) + } + }) + } +} diff --git a/internal/mcp/explore_fenced_mining_test.go b/internal/mcp/explore_fenced_mining_test.go index 0e196714..38fe19f3 100644 --- a/internal/mcp/explore_fenced_mining_test.go +++ b/internal/mcp/explore_fenced_mining_test.go @@ -50,7 +50,7 @@ func TestExploreTaskAnchorsResolveExactNameOnlyInsideFence(t *testing.T) { require.Len(t, anchors[0].exactNodes, 1) got := server.exploreExactAnchorCandidate( - context.Background(), anchors[0], query.QueryOptions{}, + context.Background(), anchors[0], nil, query.QueryOptions{}, map[string]struct{}{}, map[string]struct{}{}, ) require.NotNil(t, got) @@ -105,12 +105,14 @@ func TestExploreQuotedRecallTermsKeepProseTermsAheadOfMinedLiterals(t *testing.T "debug: reader ignored the pending manifest entry\n" + "trace: writer flushed the pending batch to disk\n" + "info: registry rebuilt from the durable snapshot\n" + + "notice: registry verified after snapshot recovery\n" + "```\n" terms := exploreQuotedRecallTerms(task) require.Len(t, terms, exploreQuotedRecallMaxMinedTerms) require.Equal(t, "tenant registry cleared", terms[0], "prose literals keep the leading slots: %#v", terms) require.Contains(t, terms, "tenant registry rollback discarded entry") - require.NotContains(t, terms, "registry rebuilt from the durable snapshot", + require.Contains(t, terms, "registry rebuilt from the durable snapshot") + require.NotContains(t, terms, "registry verified after snapshot recovery", "mined literals stop at the total cap: %#v", terms) } diff --git a/internal/mcp/explore_implementation_bounded_fault_test.go b/internal/mcp/explore_implementation_bounded_fault_test.go new file mode 100644 index 00000000..d6ab9171 --- /dev/null +++ b/internal/mcp/explore_implementation_bounded_fault_test.go @@ -0,0 +1,486 @@ +package mcp + +import ( + "context" + "errors" + "testing" + + "github.com/zzet/gortex/internal/graph" + "github.com/zzet/gortex/internal/query" +) + +type exploreImplementationAdjacencyRequest struct { + ids []string + kinds []graph.EdgeKind + limit int +} + +type exploreImplementationFaultReader struct { + graph.Reader + outgoing func(context.Context, []string, []graph.EdgeKind, int) (graph.BoundedEdgeIdentityProjection, error) + incoming func(context.Context, []string, []graph.EdgeKind, int) (graph.BoundedEdgeIdentityProjection, error) + nodes func(context.Context, []string) (map[string]*graph.Node, error) + + outgoingRequests []exploreImplementationAdjacencyRequest + incomingRequests []exploreImplementationAdjacencyRequest + nodeRequests [][]string +} + +func (r *exploreImplementationFaultReader) GetOutEdges(string) []*graph.Edge { + panic("legacy GetOutEdges must not be called") +} + +func (r *exploreImplementationFaultReader) GetInEdges(string) []*graph.Edge { + panic("legacy GetInEdges must not be called") +} + +func (r *exploreImplementationFaultReader) FindOutgoingEdgeIdentitiesBounded( + ctx context.Context, + ids []string, + kinds []graph.EdgeKind, + limit int, +) (graph.BoundedEdgeIdentityProjection, error) { + r.outgoingRequests = append(r.outgoingRequests, exploreImplementationAdjacencyRequest{ + ids: append([]string(nil), ids...), kinds: append([]graph.EdgeKind(nil), kinds...), limit: limit, + }) + if r.outgoing != nil { + return r.outgoing(ctx, ids, kinds, limit) + } + return r.Reader.(graph.BoundedOutgoingEdgeIdentityReader).FindOutgoingEdgeIdentitiesBounded(ctx, ids, kinds, limit) +} + +func (r *exploreImplementationFaultReader) FindIncomingEdgeIdentitiesBounded( + ctx context.Context, + ids []string, + kinds []graph.EdgeKind, + limit int, +) (graph.BoundedEdgeIdentityProjection, error) { + r.incomingRequests = append(r.incomingRequests, exploreImplementationAdjacencyRequest{ + ids: append([]string(nil), ids...), kinds: append([]graph.EdgeKind(nil), kinds...), limit: limit, + }) + if r.incoming != nil { + return r.incoming(ctx, ids, kinds, limit) + } + return r.Reader.(graph.BoundedIncomingEdgeIdentityReader).FindIncomingEdgeIdentitiesBounded(ctx, ids, kinds, limit) +} + +func (r *exploreImplementationFaultReader) GetNodesByIDsContext( + ctx context.Context, + ids []string, +) (map[string]*graph.Node, error) { + r.nodeRequests = append(r.nodeRequests, append([]string(nil), ids...)) + if r.nodes != nil { + return r.nodes(ctx, append([]string(nil), ids...)) + } + if err := ctx.Err(); err != nil { + return nil, err + } + return r.GetNodesByIDs(ids), nil +} + +type exploreImplementationLegacyOnlyReader struct{ graph.Reader } + +func (r *exploreImplementationLegacyOnlyReader) GetOutEdges(string) []*graph.Edge { + panic("legacy GetOutEdges must not be called") +} + +func (r *exploreImplementationLegacyOnlyReader) GetInEdges(string) []*graph.Edge { + panic("legacy GetInEdges must not be called") +} + +type exploreImplementationCancelAfterChecksContext struct { + context.Context + remaining int +} + +func (ctx *exploreImplementationCancelAfterChecksContext) Err() error { + ctx.remaining-- + if ctx.remaining <= 0 { + return context.Canceled + } + return nil +} + +func exploreImplementationTestServer(reader graph.Reader) *Server { + return &Server{engine: query.NewEngine(graph.New()).WithReader(reader)} +} + +func TestImplementationRecoveryUsesExactBoundedCapabilities(t *testing.T) { + _, base := newStorageFixtureServer(t) + reader := &exploreImplementationFaultReader{Reader: base} + server := exploreImplementationTestServer(reader) + seed := exploreTarget{node: base.GetNode("repo/storage/storage.go::Storage.load"), score: 1} + if got := server.expandImplementationTargets(context.Background(), []exploreTarget{seed}, graph.LocalizationNodeScope{}); len(got) < 2 { + t.Fatalf("bounded implementation recovery returned %d targets", len(got)) + } + if len(reader.outgoingRequests) != 1 || reader.outgoingRequests[0].limit != exploreImplOwnerRelationLimit || + len(reader.outgoingRequests[0].kinds) != 1 || reader.outgoingRequests[0].kinds[0] != graph.EdgeMemberOf { + t.Fatalf("owner request = %#v", reader.outgoingRequests) + } + foundImplementors, foundMembers := false, false + for _, request := range reader.incomingRequests { + switch request.limit { + case exploreImplImplementorLimit: + foundImplementors = len(request.kinds) == 2 && request.kinds[0] == graph.EdgeImplements && request.kinds[1] == graph.EdgeExtends + case exploreImplMemberRefinementLimit: + foundMembers = len(request.kinds) == 1 && request.kinds[0] == graph.EdgeMemberOf + } + } + if !foundImplementors || !foundMembers { + t.Fatalf("incoming requests = %#v", reader.incomingRequests) + } + wantNodeRequests := [][]string{ + {"repo/storage/storage.go::Storage"}, + {"repo/storage/cloud.go::CloudStorage", "repo/storage/disk.go::DiskStorage"}, + {"repo/storage/cloud.go::CloudStorage.load"}, + {"repo/storage/disk.go::DiskStorage.load"}, + } + if len(reader.nodeRequests) != len(wantNodeRequests) { + t.Fatalf("node requests = %#v, want %#v", reader.nodeRequests, wantNodeRequests) + } + for index := range wantNodeRequests { + if len(reader.nodeRequests[index]) != len(wantNodeRequests[index]) { + t.Fatalf("node request %d = %#v, want %#v", index, reader.nodeRequests[index], wantNodeRequests[index]) + } + for item := range wantNodeRequests[index] { + if reader.nodeRequests[index][item] != wantNodeRequests[index][item] { + t.Fatalf("node request %d = %#v, want %#v", index, reader.nodeRequests[index], wantNodeRequests[index]) + } + } + } + + foldServer, foldBase := newOwnerFoldFixture(t) + _ = foldServer + foldReader := &exploreImplementationFaultReader{Reader: foldBase} + members := []exploreTarget{ + {node: foldBase.GetNode("repo/mw/persist.ts::PersistOptions.serialize"), score: 1}, + {node: foldBase.GetNode("repo/mw/persist.ts::PersistOptions.hydrate"), score: .9}, + } + if got := exploreImplementationTestServer(foldReader).foldMemberOwners(context.Background(), members, graph.LocalizationNodeScope{}); len(got) != 3 { + t.Fatalf("bounded owner fold returned %d targets", len(got)) + } + if len(foldReader.outgoingRequests) != 1 || foldReader.outgoingRequests[0].limit != exploreOwnerFoldRelationLimit || + len(foldReader.outgoingRequests[0].ids) != 2 { + t.Fatalf("fold request = %#v", foldReader.outgoingRequests) + } + if len(foldReader.nodeRequests) != 1 || len(foldReader.nodeRequests[0]) != 1 || + foldReader.nodeRequests[0][0] != "repo/mw/persist.ts::PersistOptions" { + t.Fatalf("fold node requests = %#v", foldReader.nodeRequests) + } +} + +func TestImplementationRecoveryHasNoLegacyAdjacencyFallback(t *testing.T) { + _, base := newStorageFixtureServer(t) + reader := &exploreImplementationLegacyOnlyReader{Reader: base} + server := exploreImplementationTestServer(reader) + method := exploreTarget{node: base.GetNode("repo/storage/storage.go::Storage.load"), score: 1} + if got := server.expandImplementationTargets(context.Background(), []exploreTarget{method}, graph.LocalizationNodeScope{}); len(got) != 1 { + t.Fatalf("unsupported expansion returned %d targets", len(got)) + } + if !exploreImplementationAnswerBlocked( + context.Background(), "find the concrete implementation", []exploreTarget{method}, reader, graph.LocalizationNodeScope{}, + ) { + t.Fatal("unsupported terminality lookup must conservatively block") + } + + _, foldBase := newOwnerFoldFixture(t) + foldReader := &exploreImplementationLegacyOnlyReader{Reader: foldBase} + members := []exploreTarget{ + {node: foldBase.GetNode("repo/mw/persist.ts::PersistOptions.serialize")}, + {node: foldBase.GetNode("repo/mw/persist.ts::PersistOptions.hydrate")}, + } + if got := exploreImplementationTestServer(foldReader).foldMemberOwners(context.Background(), members, graph.LocalizationNodeScope{}); len(got) != len(members) { + t.Fatalf("unsupported owner fold returned %d targets", len(got)) + } +} + +func TestImplementationRecoveryDiscardsFaultedOrCanceledProjections(t *testing.T) { + _, base := newStorageFixtureServer(t) + method := exploreTarget{node: base.GetNode("repo/storage/storage.go::Storage.load"), score: 1} + partialOwner := graph.BoundedEdgeIdentityProjection{ByEndpoint: map[string][]graph.EdgeIdentity{ + method.node.ID: {{From: method.node.ID, To: "repo/storage/storage.go::Storage", Kind: graph.EdgeMemberOf}}, + }} + + t.Run("partial error", func(t *testing.T) { + reader := &exploreImplementationFaultReader{Reader: base} + reader.outgoing = func(context.Context, []string, []graph.EdgeKind, int) (graph.BoundedEdgeIdentityProjection, error) { + return partialOwner, errors.New("projection failed") + } + server := exploreImplementationTestServer(reader) + if got := server.expandImplementationTargets(context.Background(), []exploreTarget{method}, graph.LocalizationNodeScope{}); len(got) != 1 { + t.Fatalf("partial-error expansion returned %d targets", len(got)) + } + if !exploreImplementationAnswerBlocked(context.Background(), "find implementation", []exploreTarget{method}, reader, graph.LocalizationNodeScope{}) { + t.Fatal("partial-error terminality lookup must block") + } + }) + + t.Run("sticky cancellation", func(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + reader := &exploreImplementationFaultReader{Reader: base} + reader.outgoing = func(context.Context, []string, []graph.EdgeKind, int) (graph.BoundedEdgeIdentityProjection, error) { + cancel() + return partialOwner, nil + } + if !exploreImplementationAnswerBlocked(ctx, "find implementation", []exploreTarget{method}, reader, graph.LocalizationNodeScope{}) { + t.Fatal("sticky cancellation must block terminality") + } + }) + + t.Run("missing hydration", func(t *testing.T) { + reader := &exploreImplementationFaultReader{Reader: base} + reader.nodes = func(context.Context, []string) (map[string]*graph.Node, error) { + return map[string]*graph.Node{}, nil + } + server := exploreImplementationTestServer(reader) + if got := server.expandImplementationTargets(context.Background(), []exploreTarget{method}, graph.LocalizationNodeScope{}); len(got) != 1 { + t.Fatalf("missing-hydration expansion returned %d targets", len(got)) + } + }) + + t.Run("malformed endpoint", func(t *testing.T) { + reader := &exploreImplementationFaultReader{Reader: base} + reader.outgoing = func(context.Context, []string, []graph.EdgeKind, int) (graph.BoundedEdgeIdentityProjection, error) { + return graph.BoundedEdgeIdentityProjection{ByEndpoint: map[string][]graph.EdgeIdentity{ + method.node.ID: {{From: "foreign", To: "repo/storage/storage.go::Storage", Kind: graph.EdgeMemberOf}}, + }}, nil + } + if !exploreImplementationAnswerBlocked(context.Background(), "find implementation", []exploreTarget{method}, reader, graph.LocalizationNodeScope{}) { + t.Fatal("malformed endpoint must block terminality") + } + }) +} + +func TestImplementationAnswerBlockedChecksCancellationAfterClassification(t *testing.T) { + base := graph.New() + concrete := &graph.Node{ID: "repo/impl.go::Concrete", Kind: graph.KindType, Name: "Concrete", FilePath: "repo/impl.go"} + base.AddNode(concrete) + reader := &exploreImplementationFaultReader{Reader: base} + ctx := &exploreImplementationCancelAfterChecksContext{Context: context.Background(), remaining: 2} + if !exploreImplementationAnswerBlocked( + ctx, "find implementation", []exploreTarget{{node: concrete}}, reader, graph.LocalizationNodeScope{}, + ) { + t.Fatal("cancellation after concrete classification must block terminality") + } +} + +func TestExploreBoundedEdgeIdentitiesRejectsMalformedProjection(t *testing.T) { + endpoint := "source" + valid := graph.EdgeIdentity{From: endpoint, To: "owner", Kind: graph.EdgeMemberOf} + for _, test := range []struct { + name string + limit int + projection graph.BoundedEdgeIdentityProjection + }{ + {name: "foreign key", projection: graph.BoundedEdgeIdentityProjection{ByEndpoint: map[string][]graph.EdgeIdentity{"foreign": {valid}}}}, + {name: "wrong kind", projection: graph.BoundedEdgeIdentityProjection{ByEndpoint: map[string][]graph.EdgeIdentity{endpoint: {{From: endpoint, To: "owner", Kind: graph.EdgeReads}}}}}, + {name: "wrong source", projection: graph.BoundedEdgeIdentityProjection{ByEndpoint: map[string][]graph.EdgeIdentity{endpoint: {{From: "foreign", To: "owner", Kind: graph.EdgeMemberOf}}}}}, + {name: "empty target", projection: graph.BoundedEdgeIdentityProjection{ByEndpoint: map[string][]graph.EdgeIdentity{endpoint: {{From: endpoint, Kind: graph.EdgeMemberOf}}}}}, + {name: "cardinality", projection: graph.BoundedEdgeIdentityProjection{ByEndpoint: map[string][]graph.EdgeIdentity{endpoint: {valid, {From: endpoint, To: "owner-2", Kind: graph.EdgeMemberOf}}}}}, + {name: "duplicate within limit", limit: 2, projection: graph.BoundedEdgeIdentityProjection{ByEndpoint: map[string][]graph.EdgeIdentity{endpoint: {valid, valid}}}}, + {name: "truncated", projection: graph.BoundedEdgeIdentityProjection{Truncated: map[string]bool{endpoint: true}}}, + } { + t.Run(test.name, func(t *testing.T) { + reader := &exploreImplementationFaultReader{Reader: graph.New()} + reader.outgoing = func(context.Context, []string, []graph.EdgeKind, int) (graph.BoundedEdgeIdentityProjection, error) { + return test.projection, nil + } + limit := test.limit + if limit == 0 { + limit = 1 + } + if rows, ok := exploreBoundedEdgeIdentities( + context.Background(), reader, []string{endpoint}, []graph.EdgeKind{graph.EdgeMemberOf}, limit, exploreBoundedOutgoing, + ); ok || rows != nil { + t.Fatalf("malformed projection accepted: ok=%v rows=%v", ok, rows) + } + }) + } + + t.Run("incoming wrong target", func(t *testing.T) { + reader := &exploreImplementationFaultReader{Reader: graph.New()} + reader.incoming = func(context.Context, []string, []graph.EdgeKind, int) (graph.BoundedEdgeIdentityProjection, error) { + return graph.BoundedEdgeIdentityProjection{ByEndpoint: map[string][]graph.EdgeIdentity{ + "target": {{From: "implementation", To: "foreign", Kind: graph.EdgeImplements}}, + }}, nil + } + if rows, ok := exploreBoundedEdgeIdentities( + context.Background(), reader, []string{"target"}, []graph.EdgeKind{graph.EdgeImplements}, 1, exploreBoundedIncoming, + ); ok || rows != nil { + t.Fatalf("incoming wrong-target projection accepted: ok=%v rows=%v", ok, rows) + } + }) +} + +func TestImplementationRecoveryRejectsIncomingAndHydrationFaultsAtomically(t *testing.T) { + _, base := newStorageFixtureServer(t) + interfaceSeed := exploreTarget{node: base.GetNode("repo/storage/storage.go::Storage"), score: 1} + + t.Run("incoming partial error", func(t *testing.T) { + reader := &exploreImplementationFaultReader{Reader: base} + reader.incoming = func(ctx context.Context, ids []string, kinds []graph.EdgeKind, limit int) (graph.BoundedEdgeIdentityProjection, error) { + projection, err := base.FindIncomingEdgeIdentitiesBounded(ctx, ids, kinds, limit) + if err != nil { + return graph.BoundedEdgeIdentityProjection{}, err + } + return projection, errors.New("after partial projection") + } + got := exploreImplementationTestServer(reader).expandImplementationTargets( + context.Background(), []exploreTarget{interfaceSeed}, graph.LocalizationNodeScope{}, + ) + if len(got) != 1 { + t.Fatalf("incoming partial error returned %d targets", len(got)) + } + }) + + for _, test := range []struct { + name string + lookup func(context.Context, []string) (map[string]*graph.Node, error) + }{ + {name: "full map plus error", lookup: func(_ context.Context, ids []string) (map[string]*graph.Node, error) { + return base.GetNodesByIDs(ids), errors.New("hydration failed") + }}, + {name: "full map plus cancellation", lookup: func(ctx context.Context, ids []string) (map[string]*graph.Node, error) { + if cancel, ok := ctx.Value(exploreImplementationCancelKey{}).(context.CancelFunc); ok { + cancel() + } + return base.GetNodesByIDs(ids), nil + }}, + } { + t.Run(test.name, func(t *testing.T) { + reader := &exploreImplementationFaultReader{Reader: base, nodes: test.lookup} + ctx := context.Background() + if test.name == "full map plus cancellation" { + cancelCtx, cancel := context.WithCancel(ctx) + ctx = context.WithValue(cancelCtx, exploreImplementationCancelKey{}, context.CancelFunc(cancel)) + } + got := exploreImplementationTestServer(reader).expandImplementationTargets( + ctx, []exploreTarget{interfaceSeed}, graph.LocalizationNodeScope{}, + ) + if len(got) != 1 { + t.Fatalf("faulted hydration returned %d targets", len(got)) + } + }) + } + + t.Run("owner fold partial error", func(t *testing.T) { + _, foldBase := newOwnerFoldFixture(t) + reader := &exploreImplementationFaultReader{Reader: foldBase} + reader.outgoing = func(ctx context.Context, ids []string, kinds []graph.EdgeKind, limit int) (graph.BoundedEdgeIdentityProjection, error) { + projection, err := foldBase.FindOutgoingEdgeIdentitiesBounded(ctx, ids, kinds, limit) + if err != nil { + return graph.BoundedEdgeIdentityProjection{}, err + } + return projection, errors.New("after partial projection") + } + members := []exploreTarget{ + {node: foldBase.GetNode("repo/mw/persist.ts::PersistOptions.serialize")}, + {node: foldBase.GetNode("repo/mw/persist.ts::PersistOptions.hydrate")}, + } + if got := exploreImplementationTestServer(reader).foldMemberOwners(context.Background(), members, graph.LocalizationNodeScope{}); len(got) != len(members) { + t.Fatalf("faulted owner fold returned %d targets", len(got)) + } + }) +} + +type exploreImplementationCancelKey struct{} + +func TestExpandImplementationTargetsStageFaultSemantics(t *testing.T) { + for _, test := range []struct { + name string + stageLimit int + mode string + wantOriginal bool + }{ + {name: "implementor partial error", stageLimit: exploreImplImplementorLimit, mode: "error", wantOriginal: true}, + {name: "implementor malformed", stageLimit: exploreImplImplementorLimit, mode: "malformed", wantOriginal: true}, + {name: "implementor canceled", stageLimit: exploreImplImplementorLimit, mode: "cancel", wantOriginal: true}, + {name: "member partial error falls back", stageLimit: exploreImplMemberRefinementLimit, mode: "error"}, + {name: "member malformed falls back", stageLimit: exploreImplMemberRefinementLimit, mode: "malformed"}, + {name: "member truncated falls back", stageLimit: exploreImplMemberRefinementLimit, mode: "truncated"}, + {name: "member canceled aborts lane", stageLimit: exploreImplMemberRefinementLimit, mode: "cancel", wantOriginal: true}, + } { + t.Run(test.name, func(t *testing.T) { + _, base := newStorageFixtureServer(t) + ctx, cancel := context.WithCancel(context.Background()) + reader := &exploreImplementationFaultReader{Reader: base} + reader.incoming = func(callCtx context.Context, ids []string, kinds []graph.EdgeKind, limit int) (graph.BoundedEdgeIdentityProjection, error) { + projection, err := base.FindIncomingEdgeIdentitiesBounded(callCtx, ids, kinds, limit) + if err != nil || limit != test.stageLimit { + return projection, err + } + switch test.mode { + case "error": + return projection, errors.New("stage failed after projection") + case "cancel": + cancel() + return projection, nil + case "malformed": + return graph.BoundedEdgeIdentityProjection{ByEndpoint: map[string][]graph.EdgeIdentity{ + ids[0]: {{From: "foreign", To: ids[0], Kind: kinds[0]}}, + }}, nil + case "truncated": + return graph.BoundedEdgeIdentityProjection{Truncated: map[string]bool{ids[0]: true}}, nil + default: + return projection, nil + } + } + seed := exploreTarget{node: base.GetNode("repo/storage/storage.go::Storage.load"), score: 1} + got := exploreImplementationTestServer(reader).expandImplementationTargets( + ctx, []exploreTarget{seed}, graph.LocalizationNodeScope{}, + ) + if test.wantOriginal { + if len(got) != 1 || got[0].node.ID != seed.node.ID { + t.Fatalf("faulted stage returned partial targets: %#v", got) + } + return + } + if len(got) < 2 { + t.Fatalf("member fault lost concrete-type fallback: %#v", got) + } + for _, target := range got[1:] { + if target.node == nil || target.node.Kind != graph.KindType { + t.Fatalf("member fault admitted non-type refinement: %#v", target.node) + } + } + }) + } +} + +func TestFoldMemberOwnersRelationFaultsAreAtomic(t *testing.T) { + for _, mode := range []string{"error", "cancel", "malformed", "truncated"} { + t.Run(mode, func(t *testing.T) { + _, base := newOwnerFoldFixture(t) + ctx, cancel := context.WithCancel(context.Background()) + reader := &exploreImplementationFaultReader{Reader: base} + reader.outgoing = func(callCtx context.Context, ids []string, kinds []graph.EdgeKind, limit int) (graph.BoundedEdgeIdentityProjection, error) { + projection, err := base.FindOutgoingEdgeIdentitiesBounded(callCtx, ids, kinds, limit) + if err != nil { + return projection, err + } + switch mode { + case "error": + return projection, errors.New("fold projection failed") + case "cancel": + cancel() + return projection, nil + case "malformed": + return graph.BoundedEdgeIdentityProjection{ByEndpoint: map[string][]graph.EdgeIdentity{ + ids[0]: {{From: "foreign", To: "owner", Kind: graph.EdgeMemberOf}}, + }}, nil + case "truncated": + return graph.BoundedEdgeIdentityProjection{Truncated: map[string]bool{ids[0]: true}}, nil + } + return projection, nil + } + members := []exploreTarget{ + {node: base.GetNode("repo/mw/persist.ts::PersistOptions.serialize")}, + {node: base.GetNode("repo/mw/persist.ts::PersistOptions.hydrate")}, + } + got := exploreImplementationTestServer(reader).foldMemberOwners(ctx, members, graph.LocalizationNodeScope{}) + if len(got) != len(members) || got[0].node.ID != members[0].node.ID { + t.Fatalf("faulted fold returned partial mutation: %#v", got) + } + }) + } +} diff --git a/internal/mcp/explore_implementation_file_reserve_test.go b/internal/mcp/explore_implementation_file_reserve_test.go new file mode 100644 index 00000000..553e5ad2 --- /dev/null +++ b/internal/mcp/explore_implementation_file_reserve_test.go @@ -0,0 +1,45 @@ +package mcp + +import ( + "context" + "testing" + + "github.com/zzet/gortex/internal/graph" + "github.com/zzet/gortex/internal/query" +) + +func TestExpandImplementationTargetsReusesAdmittedNewFile(t *testing.T) { + g := graph.New() + owner := &graph.Node{ID: "reserve-owner", Kind: graph.KindInterface, Name: "Runner", FilePath: "repo/api.go"} + implementations := []*graph.Node{ + {ID: "reserve-impl-a", Kind: graph.KindType, Name: "ImplA", FilePath: "repo/impl-a.go"}, + {ID: "reserve-impl-b", Kind: graph.KindType, Name: "ImplB", FilePath: "repo/impl-b.go"}, + {ID: "reserve-impl-c", Kind: graph.KindType, Name: "ImplC", FilePath: "repo/impl-c.go"}, + {ID: "reserve-impl-d", Kind: graph.KindType, Name: "ImplD", FilePath: "repo/impl-a.go"}, + } + nodes := []*graph.Node{owner} + edges := make([]*graph.Edge, 0, len(implementations)) + for _, implementation := range implementations { + nodes = append(nodes, implementation) + edges = append(edges, &graph.Edge{From: implementation.ID, To: owner.ID, Kind: graph.EdgeImplements}) + } + g.AddBatch(nodes, edges) + server := &Server{engine: query.NewEngine(g)} + expanded := server.expandImplementationTargets( + context.Background(), []exploreTarget{{node: owner, score: 1}}, graph.LocalizationNodeScope{}, + ) + if len(expanded) != 1+len(implementations) { + t.Fatalf("expanded targets=%d, want %d", len(expanded), 1+len(implementations)) + } + got := make(map[string]bool, len(expanded)) + for _, target := range expanded { + if target.node != nil { + got[target.node.ID] = true + } + } + for _, implementation := range implementations { + if !got[implementation.ID] { + t.Fatalf("implementation sharing an admitted file was dropped: %s (%v)", implementation.ID, got) + } + } +} diff --git a/internal/mcp/explore_implementation_intent.go b/internal/mcp/explore_implementation_intent.go index d5e93537..b1168e1e 100644 --- a/internal/mcp/explore_implementation_intent.go +++ b/internal/mcp/explore_implementation_intent.go @@ -33,6 +33,10 @@ const ( // exploreImplFileReserve is the number of DISTINCT implementation files // the expansion reserves in the final candidate set. exploreImplFileReserve = 3 + + exploreImplOwnerRelationLimit = 8 + exploreImplImplementorLimit = 16 + exploreImplMemberRefinementLimit = 256 ) var exploreImplementationIntentTerms = []string{ @@ -59,28 +63,54 @@ func exploreImplementationIntent(task string) bool { return false } -// exploreAbstractSeed reports whether a ranked node is an abstract +// exploreBoundedAbstractSeed reports whether a ranked node is an abstract // declaration an implementation-intent query must expand: an interface, or -// an interface member (owner resolved through member_of). -func exploreAbstractSeed(getOut func(string) []*graph.Edge, getNode func(string) *graph.Node, n *graph.Node) (owner *graph.Node, abstract bool) { +// an interface member whose owner is resolved through a complete bounded +// member_of projection. +func exploreBoundedAbstractSeed( + ctx context.Context, + reader graph.Reader, + declarationScope graph.LocalizationNodeScope, + n *graph.Node, +) (owner *graph.Node, abstract, complete bool) { if n == nil { - return nil, false + return nil, false, true } if n.Kind == graph.KindInterface { - return n, true + return n, true, true } if n.Kind != graph.KindMethod && n.Kind != graph.KindFunction { - return nil, false + return nil, false, true + } + relations, complete := exploreBoundedEdgeIdentities( + ctx, reader, []string{n.ID}, []graph.EdgeKind{graph.EdgeMemberOf}, + exploreImplOwnerRelationLimit, exploreBoundedOutgoing, + ) + if !complete { + return nil, false, false + } + edges := relations[n.ID] + ownerIDs := make([]string, 0, len(edges)) + for _, edge := range edges { + ownerIDs = append(ownerIDs, edge.To) + } + owners, complete := exploreNodesByIDsBounded(ctx, reader, ownerIDs, exploreImplOwnerRelationLimit) + if !complete { + return nil, false, false } - for _, e := range getOut(n.ID) { - if e == nil || e.Kind != graph.EdgeMemberOf { + for _, edge := range edges { + owner := owners[edge.To] + if owner == nil || owner.ID != edge.To { + return nil, false, false + } + if !declarationScope.Allows(owner) { continue } - if ownerNode := getNode(e.To); ownerNode != nil && ownerNode.Kind == graph.KindInterface { - return ownerNode, true + if owner.Kind == graph.KindInterface { + return owner, true, ctx.Err() == nil } } - return nil, false + return nil, false, ctx.Err() == nil } // expandImplementationTargets inserts concrete implementors after abstract @@ -89,59 +119,114 @@ func exploreAbstractSeed(getOut func(string) []*graph.Edge, getNode func(string) // The abstract seed is never evicted, admitted implementors cover at most // exploreImplFileReserve distinct new files, and every walk is one bounded // relation hop — no transitive traversal. -func (s *Server) expandImplementationTargets(ctx context.Context, targets []exploreTarget) []exploreTarget { +func (s *Server) expandImplementationTargets( + ctx context.Context, + targets []exploreTarget, + declarationScope graph.LocalizationNodeScope, +) []exploreTarget { eng := s.engineFor(ctx) if eng == nil || len(targets) == 0 { return targets } + reader := eng.Reader() + if reader == nil || ctx.Err() != nil { + return targets + } present := make(map[string]struct{}, len(targets)) presentFiles := make(map[string]struct{}, len(targets)) - for _, t := range targets { - if t.node != nil { - present[t.node.ID] = struct{}{} - presentFiles[t.node.FilePath] = struct{}{} + for _, target := range targets { + if target.node != nil { + present[target.node.ID] = struct{}{} + presentFiles[target.node.FilePath] = struct{}{} } } newFiles := make(map[string]struct{}, exploreImplFileReserve) - admit := func(n *graph.Node) bool { - if n == nil || n.ID == "" { + admit := func(node *graph.Node) bool { + if node == nil || node.ID == "" || !declarationScope.Allows(node) || ctx.Err() != nil { return false } - if _, dup := present[n.ID]; dup { + if _, duplicate := present[node.ID]; duplicate { return false } - if _, known := presentFiles[n.FilePath]; !known { - if len(newFiles) >= exploreImplFileReserve { + _, presentFile := presentFiles[node.FilePath] + _, admittedFile := newFiles[node.FilePath] + if !presentFile && !admittedFile { + if len(newFiles) >= exploreImplFileReserve || ctx.Err() != nil { return false } - newFiles[n.FilePath] = struct{}{} + newFiles[node.FilePath] = struct{}{} } - present[n.ID] = struct{}{} + if ctx.Err() != nil { + return false + } + present[node.ID] = struct{}{} return true } - implementorsOf := func(ownerID string) []*graph.Node { - var out []*graph.Node - for _, e := range eng.GetInEdges(ownerID) { - if e == nil || (e.Kind != graph.EdgeImplements && e.Kind != graph.EdgeExtends) { + implementorsOf := func(ownerID string) ([]*graph.Node, bool) { + relations, complete := exploreBoundedEdgeIdentities( + ctx, reader, []string{ownerID}, + []graph.EdgeKind{graph.EdgeImplements, graph.EdgeExtends}, + exploreImplImplementorLimit, exploreBoundedIncoming, + ) + if !complete { + return nil, false + } + edges := relations[ownerID] + implementorIDs := make([]string, 0, len(edges)) + for _, edge := range edges { + implementorIDs = append(implementorIDs, edge.From) + } + nodes, complete := exploreNodesByIDsBounded(ctx, reader, implementorIDs, exploreImplImplementorLimit) + if !complete { + return nil, false + } + out := make([]*graph.Node, 0, min(exploreImplPerSeed, len(edges))) + for _, edge := range edges { + implementation := nodes[edge.From] + if implementation == nil || implementation.ID != edge.From { + return nil, false + } + if !declarationScope.Allows(implementation) { continue } - if impl := eng.GetSymbol(e.From); impl != nil { - out = append(out, impl) - if len(out) == exploreImplPerSeed { - break - } + if implementation.Kind == graph.KindInterface { + continue + } + out = append(out, implementation) + if len(out) == exploreImplPerSeed { + break } } - return out + return out, ctx.Err() == nil } memberOf := func(typeID, name string) *graph.Node { - for _, e := range eng.GetInEdges(typeID) { - if e == nil || e.Kind != graph.EdgeMemberOf { + relations, complete := exploreBoundedEdgeIdentities( + ctx, reader, []string{typeID}, []graph.EdgeKind{graph.EdgeMemberOf}, + exploreImplMemberRefinementLimit, exploreBoundedIncoming, + ) + if !complete { + return nil + } + edges := relations[typeID] + memberIDs := make([]string, 0, len(edges)) + for _, edge := range edges { + memberIDs = append(memberIDs, edge.From) + } + nodes, complete := exploreNodesByIDsBounded(ctx, reader, memberIDs, exploreImplMemberRefinementLimit) + if !complete { + return nil + } + for _, edge := range edges { + member := nodes[edge.From] + if member == nil || member.ID != edge.From { + return nil + } + if !declarationScope.Allows(member) { continue } - if member := eng.GetSymbol(e.From); member != nil && member.Name == name { + if member.Name == name { return member } } @@ -149,56 +234,87 @@ func (s *Server) expandImplementationTargets(ctx context.Context, targets []expl } expanded := make([]exploreTarget, 0, len(targets)+exploreImplPerSeed) - for index, t := range targets { - expanded = append(expanded, t) - if index >= exploreImplSeedScan || t.node == nil { + for index, target := range targets { + expanded = append(expanded, target) + if index >= exploreImplSeedScan || target.node == nil { continue } - owner, abstract := exploreAbstractSeed(eng.GetOutEdges, eng.GetSymbol, t.node) + owner, abstract, complete := exploreBoundedAbstractSeed(ctx, reader, declarationScope, target.node) + if !complete { + return targets + } if !abstract { continue } + implementations, complete := implementorsOf(owner.ID) + if !complete { + return targets + } memberName := "" - if owner.ID != t.node.ID { - memberName = t.node.Name + if owner.ID != target.node.ID { + memberName = target.node.Name } - for _, impl := range implementorsOf(owner.ID) { - concrete := impl + for _, implementation := range implementations { + concrete := implementation if memberName != "" { - if member := memberOf(impl.ID, memberName); member != nil { + if member := memberOf(implementation.ID, memberName); member != nil { concrete = member } + if ctx.Err() != nil { + return targets + } } if !admit(concrete) { continue } expanded = append(expanded, exploreTarget{ node: concrete, - score: t.score * 0.98, + score: target.score * 0.98, }) } } + if ctx.Err() != nil { + return targets + } return expanded } // exploreImplementationAnswerBlocked refuses answer_ready when an // implementation-intent query's visible head holds only abstract // declarations: the concrete code the query asks for is not in evidence. -func exploreImplementationAnswerBlocked(task string, targets []exploreTarget, getOut func(string) []*graph.Edge, getNode func(string) *graph.Node) bool { +func exploreImplementationAnswerBlocked( + ctx context.Context, + task string, + targets []exploreTarget, + reader graph.Reader, + declarationScope graph.LocalizationNodeScope, +) bool { if !exploreImplementationIntent(task) { return false } + if reader == nil || ctx.Err() != nil { + return len(targets) > 0 + } scanned := 0 - for _, t := range targets { - if t.node == nil { + for _, target := range targets { + if target.node == nil { continue } if scanned++; scanned > exploreImplSeedScan { break } - if _, abstract := exploreAbstractSeed(getOut, getNode, t.node); !abstract { + _, abstract, complete := exploreBoundedAbstractSeed( + ctx, reader, declarationScope, target.node, + ) + if !complete || ctx.Err() != nil { + return true + } + if !abstract { return false } } + if ctx.Err() != nil { + return scanned > 0 + } return scanned > 0 } diff --git a/internal/mcp/explore_implementation_intent_test.go b/internal/mcp/explore_implementation_intent_test.go index f2957522..a58334aa 100644 --- a/internal/mcp/explore_implementation_intent_test.go +++ b/internal/mcp/explore_implementation_intent_test.go @@ -2,6 +2,7 @@ package mcp import ( "context" + "fmt" "testing" "go.uber.org/zap" @@ -57,7 +58,7 @@ func TestExpandImplementationTargets(t *testing.T) { s, g := newStorageFixtureServer(t) seed := exploreTarget{node: g.GetNode("repo/storage/storage.go::Storage.load"), score: 1.0} - expanded := s.expandImplementationTargets(context.Background(), []exploreTarget{seed}) + expanded := s.expandImplementationTargets(context.Background(), []exploreTarget{seed}, graph.LocalizationNodeScope{}) if expanded[0].node.ID != "repo/storage/storage.go::Storage.load" { t.Fatal("abstract seed must stay at its rank") } @@ -70,7 +71,7 @@ func TestExpandImplementationTargets(t *testing.T) { } ifaceSeed := exploreTarget{node: g.GetNode("repo/storage/storage.go::Storage"), score: 1.0} - expanded = s.expandImplementationTargets(context.Background(), []exploreTarget{ifaceSeed}) + expanded = s.expandImplementationTargets(context.Background(), []exploreTarget{ifaceSeed}, graph.LocalizationNodeScope{}) got = map[string]bool{} for _, tgt := range expanded[1:] { got[tgt.node.ID] = true @@ -91,14 +92,252 @@ func TestImplementationAnswerBlockedOnAbstractOnlyHead(t *testing.T) { {node: g.GetNode("repo/storage/storage.go::Storage.load")}, {node: g.GetNode("repo/storage/storage.go::Storage")}, } - if !exploreImplementationAnswerBlocked(task, abstractOnly, eng.GetOutEdges, eng.GetSymbol) { + if !exploreImplementationAnswerBlocked(context.Background(), task, abstractOnly, eng.Reader(), graph.LocalizationNodeScope{}) { t.Fatal("abstract-only head must block answer_ready for implementation intent") } withConcrete := append(abstractOnly, exploreTarget{node: g.GetNode("repo/storage/disk.go::DiskStorage.load")}) - if exploreImplementationAnswerBlocked(task, withConcrete, eng.GetOutEdges, eng.GetSymbol) { + if exploreImplementationAnswerBlocked(context.Background(), task, withConcrete, eng.Reader(), graph.LocalizationNodeScope{}) { t.Fatal("a concrete implementor in the head must unblock answer_ready") } - if exploreImplementationAnswerBlocked("where is Storage.load declared", abstractOnly, eng.GetOutEdges, eng.GetSymbol) { + if exploreImplementationAnswerBlocked(context.Background(), "where is Storage.load declared", abstractOnly, eng.Reader(), graph.LocalizationNodeScope{}) { t.Fatal("non-intent queries must be untouched") } } + +func TestExpandImplementationTargetsBoundedRelationCaps(t *testing.T) { + t.Run("abstract owner exact and plus one", func(t *testing.T) { + for _, test := range []struct { + name string + extra int + wantExpand bool + }{ + {name: "exact", extra: exploreImplOwnerRelationLimit - 1, wantExpand: true}, + {name: "plus one", extra: exploreImplOwnerRelationLimit, wantExpand: false}, + } { + t.Run(test.name, func(t *testing.T) { + s, g := newStorageFixtureServer(t) + for index := 0; index < test.extra; index++ { + id := fmt.Sprintf("repo/storage/storage.go::Owner%02d", index) + g.AddNode(&graph.Node{ID: id, Kind: graph.KindType, Name: fmt.Sprintf("Owner%02d", index), FilePath: "repo/storage/storage.go"}) + g.AddEdge(&graph.Edge{From: "repo/storage/storage.go::Storage.load", To: id, Kind: graph.EdgeMemberOf}) + } + for index := 0; index < 32; index++ { + g.AddEdge(&graph.Edge{From: "repo/storage/storage.go::Storage.load", To: fmt.Sprintf("irrelevant-owner-%02d", index), Kind: graph.EdgeReads, Line: index + 1}) + } + seed := exploreTarget{node: g.GetNode("repo/storage/storage.go::Storage.load"), score: 1} + expanded := s.expandImplementationTargets(context.Background(), []exploreTarget{seed}, graph.LocalizationNodeScope{}) + if got := len(expanded) > 1; got != test.wantExpand { + t.Fatalf("expanded=%v at %d owner relations, want %v", got, test.extra+1, test.wantExpand) + } + }) + } + }) + + t.Run("implementors exact and plus one", func(t *testing.T) { + for _, test := range []struct { + name string + extra int + wantExpand bool + }{ + {name: "exact", extra: exploreImplImplementorLimit - 2, wantExpand: true}, + {name: "plus one", extra: exploreImplImplementorLimit - 1, wantExpand: false}, + } { + t.Run(test.name, func(t *testing.T) { + s, g := newStorageFixtureServer(t) + for index := 0; index < test.extra; index++ { + id := fmt.Sprintf("repo/storage/impl.go::Impl%02d", index) + g.AddNode(&graph.Node{ID: id, Kind: graph.KindType, Name: fmt.Sprintf("Impl%02d", index), FilePath: "repo/storage/impl.go"}) + g.AddEdge(&graph.Edge{From: id, To: "repo/storage/storage.go::Storage", Kind: graph.EdgeImplements}) + } + for index := 0; index < 32; index++ { + g.AddEdge(&graph.Edge{From: fmt.Sprintf("irrelevant-impl-%02d", index), To: "repo/storage/storage.go::Storage", Kind: graph.EdgeReads, Line: index + 1}) + } + seed := exploreTarget{node: g.GetNode("repo/storage/storage.go::Storage"), score: 1} + expanded := s.expandImplementationTargets(context.Background(), []exploreTarget{seed}, graph.LocalizationNodeScope{}) + if got := len(expanded) > 1; got != test.wantExpand { + t.Fatalf("expanded=%v at %d implementors, want %v", got, test.extra+2, test.wantExpand) + } + }) + } + }) + + t.Run("member refinement exact and plus one", func(t *testing.T) { + for _, test := range []struct { + name string + members int + wantMember bool + }{ + {name: "exact", members: exploreImplMemberRefinementLimit, wantMember: true}, + {name: "plus one", members: exploreImplMemberRefinementLimit + 1, wantMember: false}, + } { + t.Run(test.name, func(t *testing.T) { + g := graph.New() + ownerID := "repo/api/api.go::API" + seedID := "repo/api/api.go::API.run" + implID := "repo/impl/impl.go::Concrete" + g.AddBatch([]*graph.Node{ + {ID: ownerID, Kind: graph.KindInterface, Name: "API", FilePath: "repo/api/api.go"}, + {ID: seedID, Kind: graph.KindMethod, Name: "run", FilePath: "repo/api/api.go"}, + {ID: implID, Kind: graph.KindType, Name: "Concrete", FilePath: "repo/impl/impl.go"}, + }, []*graph.Edge{ + {From: seedID, To: ownerID, Kind: graph.EdgeMemberOf}, + {From: implID, To: ownerID, Kind: graph.EdgeImplements}, + }) + memberID := "repo/impl/impl.go::Concrete.000-run" + for index := 0; index < test.members; index++ { + id := fmt.Sprintf("repo/impl/impl.go::Concrete.member%03d", index) + name := fmt.Sprintf("member%03d", index) + if index == 0 { + id, name = memberID, "run" + } + g.AddNode(&graph.Node{ID: id, Kind: graph.KindMethod, Name: name, FilePath: "repo/impl/impl.go"}) + g.AddEdge(&graph.Edge{From: id, To: implID, Kind: graph.EdgeMemberOf}) + } + for index := 0; index < 32; index++ { + g.AddEdge(&graph.Edge{From: fmt.Sprintf("irrelevant-member-%02d", index), To: implID, Kind: graph.EdgeReads, Line: index + 1}) + } + s := NewServer(query.NewEngine(g), g, nil, nil, zap.NewNop(), nil) + seed := exploreTarget{node: g.GetNode(seedID), score: 1} + expanded := s.expandImplementationTargets(context.Background(), []exploreTarget{seed}, graph.LocalizationNodeScope{}) + if len(expanded) != 2 { + t.Fatalf("expanded length=%d, want 2", len(expanded)) + } + gotMember := expanded[1].node.ID == memberID + if gotMember != test.wantMember { + t.Fatalf("refined member=%v at %d rows, want %v (node=%s)", gotMember, test.members, test.wantMember, expanded[1].node.ID) + } + }) + } + }) +} + +func TestExpandImplementationTargetsAppliesScopeBeforeAdmissionQuotas(t *testing.T) { + const ( + workspace = "allowed-workspace" + project = "allowed-project" + repo = "allowed-repo" + ) + scope := graph.LocalizationNodeScope{ + WorkspaceID: workspace, + ProjectID: project, + RepoAllow: map[string]bool{repo: true}, + ExcludeTests: true, + } + g := graph.New() + owner := &graph.Node{ + ID: "repo/api/api.go::API", Kind: graph.KindInterface, Name: "API", FilePath: "repo/api/api.go", + WorkspaceID: workspace, ProjectID: project, RepoPrefix: repo, + } + g.AddNode(owner) + blocked := []*graph.Node{ + {ID: "repo/00-workspace/impl.go::Impl", Kind: graph.KindType, Name: "BlockedWorkspace", FilePath: "repo/00-workspace/impl.go", WorkspaceID: "other", ProjectID: project, RepoPrefix: repo}, + {ID: "repo/01-project/impl.go::Impl", Kind: graph.KindType, Name: "BlockedProject", FilePath: "repo/01-project/impl.go", WorkspaceID: workspace, ProjectID: "other", RepoPrefix: repo}, + {ID: "repo/02-repo/impl.go::Impl", Kind: graph.KindType, Name: "BlockedRepo", FilePath: "repo/02-repo/impl.go", WorkspaceID: workspace, ProjectID: project, RepoPrefix: "other-repo"}, + {ID: "repo/03-test/impl.go::Impl", Kind: graph.KindType, Name: "BlockedTest", FilePath: "repo/03-test/impl.go", WorkspaceID: workspace, ProjectID: project, RepoPrefix: repo, Meta: map[string]any{"is_test": true}}, + } + for _, node := range blocked { + g.AddNode(node) + g.AddEdge(&graph.Edge{From: node.ID, To: owner.ID, Kind: graph.EdgeImplements}) + } + allowed := []*graph.Node{ + {ID: "repo/10-allowed-a/impl.go::Impl", Kind: graph.KindType, Name: "AllowedA", FilePath: "repo/10-allowed-a/impl.go", WorkspaceID: workspace, ProjectID: project, RepoPrefix: repo}, + {ID: "repo/11-allowed-b/impl.go::Impl", Kind: graph.KindType, Name: "AllowedB", FilePath: "repo/11-allowed-b/impl.go", WorkspaceID: workspace, ProjectID: project, RepoPrefix: repo}, + } + for _, node := range allowed { + g.AddNode(node) + g.AddEdge(&graph.Edge{From: node.ID, To: owner.ID, Kind: graph.EdgeImplements}) + } + s := NewServer(query.NewEngine(g), g, nil, nil, zap.NewNop(), nil) + expanded := s.expandImplementationTargets(context.Background(), []exploreTarget{{node: owner, score: 1}}, scope) + got := map[string]bool{} + for _, target := range expanded { + if target.node != nil { + got[target.node.ID] = true + } + } + for _, node := range allowed { + if !got[node.ID] { + t.Fatalf("scope-allowed implementation %s was charged behind blocked rows: %v", node.ID, got) + } + } + for _, node := range blocked { + if got[node.ID] { + t.Fatalf("out-of-scope implementation was admitted: %s", node.ID) + } + } +} + +func TestImplementationRecoveryScopesHydratedOwnersAndMembers(t *testing.T) { + const ( + workspace = "workspace" + project = "project" + repo = "repo" + ) + scope := graph.LocalizationNodeScope{ + WorkspaceID: workspace, ProjectID: project, + RepoAllow: map[string]bool{repo: true}, ExcludeTests: true, + } + allowedNode := func(id string, kind graph.NodeKind, name string) *graph.Node { + return &graph.Node{ID: id, Kind: kind, Name: name, FilePath: graph.IDFile(id), WorkspaceID: workspace, ProjectID: project, RepoPrefix: repo} + } + + t.Run("owner classification", func(t *testing.T) { + g := graph.New() + seed := allowedNode("repo/api.go::API.run", graph.KindMethod, "run") + allowedOwner := allowedNode("repo/api.go::ZAllowedAPI", graph.KindInterface, "AllowedAPI") + implementation := allowedNode("repo/impl.go::Implementation", graph.KindType, "Implementation") + blockedOwners := []*graph.Node{ + {ID: "repo/api.go::AWorkspace", Kind: graph.KindInterface, Name: "AWorkspace", FilePath: "repo/api.go", WorkspaceID: "other", ProjectID: project, RepoPrefix: repo}, + {ID: "repo/api.go::BProject", Kind: graph.KindInterface, Name: "BProject", FilePath: "repo/api.go", WorkspaceID: workspace, ProjectID: "other", RepoPrefix: repo}, + {ID: "repo/api.go::CRepo", Kind: graph.KindInterface, Name: "CRepo", FilePath: "repo/api.go", WorkspaceID: workspace, ProjectID: project, RepoPrefix: "other"}, + {ID: "repo/api.go::DTest", Kind: graph.KindInterface, Name: "DTest", FilePath: "repo/api.go", WorkspaceID: workspace, ProjectID: project, RepoPrefix: repo, Meta: map[string]any{"is_test": true}}, + } + g.AddBatch([]*graph.Node{seed, allowedOwner, implementation}, []*graph.Edge{{From: implementation.ID, To: allowedOwner.ID, Kind: graph.EdgeImplements}}) + for _, owner := range blockedOwners { + g.AddNode(owner) + g.AddEdge(&graph.Edge{From: seed.ID, To: owner.ID, Kind: graph.EdgeMemberOf}) + } + g.AddEdge(&graph.Edge{From: seed.ID, To: allowedOwner.ID, Kind: graph.EdgeMemberOf}) + server := NewServer(query.NewEngine(g), g, nil, nil, zap.NewNop(), nil) + expanded := server.expandImplementationTargets(context.Background(), []exploreTarget{{node: seed, score: 1}}, scope) + if len(expanded) != 2 || expanded[1].node.ID != implementation.ID { + t.Fatalf("scope-filtered owner expansion = %#v", expanded) + } + if !exploreImplementationAnswerBlocked(context.Background(), "find implementation", []exploreTarget{{node: seed}}, g, scope) { + t.Fatal("scope-allowed interface owner must classify the method as abstract") + } + disjoint := scope + disjoint.WorkspaceID = "other-workspace" + if exploreImplementationAnswerBlocked(context.Background(), "find implementation", []exploreTarget{{node: seed}}, g, disjoint) { + t.Fatal("only out-of-scope interface owners must not classify the method as abstract") + } + }) + + t.Run("same-name member refinement", func(t *testing.T) { + g := graph.New() + owner := allowedNode("repo/api.go::API", graph.KindInterface, "API") + seed := allowedNode("repo/api.go::API.run", graph.KindMethod, "run") + implementation := allowedNode("repo/impl.go::Implementation", graph.KindType, "Implementation") + allowedMember := allowedNode("repo/impl.go::Implementation.zrun", graph.KindMethod, "run") + blockedMembers := []*graph.Node{ + {ID: "repo/impl.go::Implementation.arun", Kind: graph.KindMethod, Name: "run", FilePath: "repo/impl.go", WorkspaceID: "other", ProjectID: project, RepoPrefix: repo}, + {ID: "repo/impl.go::Implementation.brun", Kind: graph.KindMethod, Name: "run", FilePath: "repo/impl.go", WorkspaceID: workspace, ProjectID: "other", RepoPrefix: repo}, + {ID: "repo/impl.go::Implementation.crun", Kind: graph.KindMethod, Name: "run", FilePath: "repo/impl.go", WorkspaceID: workspace, ProjectID: project, RepoPrefix: "other"}, + {ID: "repo/impl.go::Implementation.drun", Kind: graph.KindMethod, Name: "run", FilePath: "repo/impl.go", WorkspaceID: workspace, ProjectID: project, RepoPrefix: repo, Meta: map[string]any{"is_test": true}}, + } + g.AddBatch([]*graph.Node{owner, seed, implementation, allowedMember}, []*graph.Edge{ + {From: seed.ID, To: owner.ID, Kind: graph.EdgeMemberOf}, + {From: implementation.ID, To: owner.ID, Kind: graph.EdgeImplements}, + {From: allowedMember.ID, To: implementation.ID, Kind: graph.EdgeMemberOf}, + }) + for _, member := range blockedMembers { + g.AddNode(member) + g.AddEdge(&graph.Edge{From: member.ID, To: implementation.ID, Kind: graph.EdgeMemberOf}) + } + server := NewServer(query.NewEngine(g), g, nil, nil, zap.NewNop(), nil) + expanded := server.expandImplementationTargets(context.Background(), []exploreTarget{{node: seed, score: 1}}, scope) + if len(expanded) != 2 || expanded[1].node.ID != allowedMember.ID { + t.Fatalf("scope-filtered member refinement = %#v", expanded) + } + }) +} diff --git a/internal/mcp/explore_implementation_order_test.go b/internal/mcp/explore_implementation_order_test.go new file mode 100644 index 00000000..8f91ae83 --- /dev/null +++ b/internal/mcp/explore_implementation_order_test.go @@ -0,0 +1,61 @@ +package mcp + +import ( + "context" + "testing" + + "github.com/zzet/gortex/internal/graph" +) + +func TestExpandImplementationTargetsCanonicalizesBackendOrder(t *testing.T) { + owner := &graph.Node{ID: "order-owner", Kind: graph.KindInterface, Name: "Runner", FilePath: "repo/api.go"} + implementations := []*graph.Node{ + {ID: "order-a", Kind: graph.KindType, Name: "A", FilePath: "repo/impl.go"}, + {ID: "order-b", Kind: graph.KindType, Name: "B", FilePath: "repo/impl.go"}, + {ID: "order-c", Kind: graph.KindType, Name: "C", FilePath: "repo/impl.go"}, + {ID: "order-d", Kind: graph.KindType, Name: "D", FilePath: "repo/impl.go"}, + {ID: "order-e", Kind: graph.KindType, Name: "E", FilePath: "repo/impl.go"}, + } + base := graph.New() + nodes := []*graph.Node{owner} + nodes = append(nodes, implementations...) + base.AddBatch(nodes, nil) + identity := func(index int) graph.EdgeIdentity { + kind := graph.EdgeImplements + if index%2 == 0 { + kind = graph.EdgeExtends + } + return graph.EdgeIdentity{From: implementations[index].ID, To: owner.ID, Kind: kind} + } + orders := [][]graph.EdgeIdentity{ + {identity(4), identity(3), identity(2), identity(1), identity(0)}, + {identity(2), identity(0), identity(4), identity(1), identity(3)}, + } + want := []string{"order-a", "order-b", "order-c", "order-d"} + for orderIndex, backendOrder := range orders { + backendOrder := append([]graph.EdgeIdentity(nil), backendOrder...) + original := append([]graph.EdgeIdentity(nil), backendOrder...) + reader := &exploreImplementationFaultReader{Reader: base} + reader.incoming = func(context.Context, []string, []graph.EdgeKind, int) (graph.BoundedEdgeIdentityProjection, error) { + return graph.BoundedEdgeIdentityProjection{ + ByEndpoint: map[string][]graph.EdgeIdentity{owner.ID: backendOrder}, + }, nil + } + expanded := exploreImplementationTestServer(reader).expandImplementationTargets( + context.Background(), []exploreTarget{{node: owner, score: 1}}, graph.LocalizationNodeScope{}, + ) + if len(expanded) != 1+len(want) { + t.Fatalf("order %d expanded=%d, want %d", orderIndex, len(expanded), 1+len(want)) + } + for index, id := range want { + if expanded[index+1].node == nil || expanded[index+1].node.ID != id { + t.Fatalf("order %d candidate %d=%#v, want %s", orderIndex, index, expanded[index+1].node, id) + } + } + for index := range original { + if backendOrder[index] != original[index] { + t.Fatalf("order %d reader slice mutated: %#v, want %#v", orderIndex, backendOrder, original) + } + } + } +} diff --git a/internal/mcp/explore_implementation_overlay_incoming_test.go b/internal/mcp/explore_implementation_overlay_incoming_test.go new file mode 100644 index 00000000..832904cd --- /dev/null +++ b/internal/mcp/explore_implementation_overlay_incoming_test.go @@ -0,0 +1,92 @@ +package mcp + +import ( + "context" + "testing" + + "github.com/zzet/gortex/internal/graph" + "github.com/zzet/gortex/internal/query" +) + +func exploreImplementationOverlayExpandedID(ctx context.Context, server *Server, seed *graph.Node) string { + expanded := server.expandImplementationTargets( + ctx, []exploreTarget{{node: seed, score: 1}}, graph.LocalizationNodeScope{}, + ) + if len(expanded) != 2 || expanded[1].node == nil { + return "" + } + return expanded[1].node.ID +} + +func TestImplementationExpansionOverlayIncomingStages(t *testing.T) { + fixture := newExploreImplementationStoreFixture() + base := graph.New() + addExploreImplementationFixtureToStore(base, fixture) + server := &Server{engine: query.NewEngine(base)} + withLayer := func(layer *graph.OverlayLayer) context.Context { + return WithOverlayView(context.Background(), graph.NewOverlaidView(base, layer)) + } + + t.Run("implementor source", func(t *testing.T) { + tombstone := graph.NewOverlayLayer() + tombstone.MarkRemoved(fixture.impl.Name, fixture.impl.ID) + if got := exploreImplementationOverlayExpandedID(withLayer(tombstone), server, fixture.seed); got != "" { + t.Fatalf("tombstoned implementor source produced %q", got) + } + + replacement := graph.NewOverlayLayer() + replacement.MarkRemoved(fixture.impl.Name, fixture.impl.ID) + replacement.AddNode(fixture.impl.FilePath, fixture.impl) + replacement.AddEdge(&graph.Edge{From: fixture.impl.ID, To: fixture.owner.ID, Kind: graph.EdgeImplements}) + if got := exploreImplementationOverlayExpandedID(withLayer(replacement), server, fixture.seed); got != fixture.member.ID { + t.Fatalf("same-ID implementor source replacement produced %q", got) + } + }) + + t.Run("implementor target", func(t *testing.T) { + tombstone := graph.NewOverlayLayer() + tombstone.MarkRemoved(fixture.owner.Name, fixture.owner.ID) + if got := exploreImplementationOverlayExpandedID(withLayer(tombstone), server, fixture.owner); got != "" { + t.Fatalf("tombstoned implementor target produced %q", got) + } + + replacement := graph.NewOverlayLayer() + replacement.MarkRemoved(fixture.owner.Name, fixture.owner.ID) + replacement.AddNode(fixture.owner.FilePath, fixture.owner) + if got := exploreImplementationOverlayExpandedID(withLayer(replacement), server, fixture.owner); got != fixture.impl.ID { + t.Fatalf("same-ID implementor target replacement produced %q", got) + } + }) + + t.Run("member source", func(t *testing.T) { + tombstone := graph.NewOverlayLayer() + tombstone.MarkRemoved(fixture.member.Name, fixture.member.ID) + if got := exploreImplementationOverlayExpandedID(withLayer(tombstone), server, fixture.seed); got != fixture.impl.ID { + t.Fatalf("tombstoned member source produced %q, want type fallback", got) + } + + replacement := graph.NewOverlayLayer() + replacement.MarkRemoved(fixture.member.Name, fixture.member.ID) + replacement.AddNode(fixture.member.FilePath, fixture.member) + replacement.AddEdge(&graph.Edge{From: fixture.member.ID, To: fixture.impl.ID, Kind: graph.EdgeMemberOf}) + if got := exploreImplementationOverlayExpandedID(withLayer(replacement), server, fixture.seed); got != fixture.member.ID { + t.Fatalf("same-ID member source replacement produced %q", got) + } + }) + + t.Run("member target", func(t *testing.T) { + tombstone := graph.NewOverlayLayer() + tombstone.MarkRemoved(fixture.impl.Name, fixture.impl.ID) + if got := exploreImplementationOverlayExpandedID(withLayer(tombstone), server, fixture.seed); got != "" { + t.Fatalf("tombstoned member target produced %q", got) + } + + replacement := graph.NewOverlayLayer() + replacement.MarkRemoved(fixture.impl.Name, fixture.impl.ID) + replacement.AddNode(fixture.impl.FilePath, fixture.impl) + replacement.AddEdge(&graph.Edge{From: fixture.impl.ID, To: fixture.owner.ID, Kind: graph.EdgeImplements}) + if got := exploreImplementationOverlayExpandedID(withLayer(replacement), server, fixture.seed); got != fixture.member.ID { + t.Fatalf("same-ID member target replacement produced %q", got) + } + }) +} diff --git a/internal/mcp/explore_implementation_overlay_test.go b/internal/mcp/explore_implementation_overlay_test.go new file mode 100644 index 00000000..dd7ad6d6 --- /dev/null +++ b/internal/mcp/explore_implementation_overlay_test.go @@ -0,0 +1,97 @@ +package mcp + +import ( + "context" + "path/filepath" + "testing" + + "github.com/zzet/gortex/internal/graph" + "github.com/zzet/gortex/internal/graph/store_sqlite" + "github.com/zzet/gortex/internal/query" +) + +type exploreImplementationStoreFixture struct { + owner *graph.Node + seed *graph.Node + impl *graph.Node + member *graph.Node + foldOwner *graph.Node + foldA *graph.Node + foldB *graph.Node + nodes []*graph.Node + edges []*graph.Edge +} + +func newExploreImplementationStoreFixture() exploreImplementationStoreFixture { + fixture := exploreImplementationStoreFixture{ + owner: &graph.Node{ID: "implementation-owner", Kind: graph.KindInterface, Name: "Runner", FilePath: "repo/api/runner.go"}, + seed: &graph.Node{ID: "implementation-seed", Kind: graph.KindMethod, Name: "run", FilePath: "repo/api/method.go"}, + impl: &graph.Node{ID: "implementation-type", Kind: graph.KindType, Name: "ConcreteRunner", FilePath: "repo/impl/runner.go"}, + member: &graph.Node{ID: "implementation-member", Kind: graph.KindMethod, Name: "run", FilePath: "repo/impl/runner.go"}, + foldOwner: &graph.Node{ID: "fold-owner", Kind: graph.KindType, Name: "Options", FilePath: "repo/options/options.go"}, + foldA: &graph.Node{ID: "fold-member-a", Kind: graph.KindMethod, Name: "encode", FilePath: "repo/options/encode.go"}, + foldB: &graph.Node{ID: "fold-member-b", Kind: graph.KindMethod, Name: "decode", FilePath: "repo/options/decode.go"}, + } + fixture.nodes = []*graph.Node{ + fixture.owner, fixture.seed, fixture.impl, fixture.member, + fixture.foldOwner, fixture.foldA, fixture.foldB, + } + fixture.edges = []*graph.Edge{ + {From: fixture.seed.ID, To: fixture.owner.ID, Kind: graph.EdgeMemberOf}, + {From: fixture.impl.ID, To: fixture.owner.ID, Kind: graph.EdgeImplements}, + {From: fixture.member.ID, To: fixture.impl.ID, Kind: graph.EdgeMemberOf}, + {From: fixture.foldA.ID, To: fixture.foldOwner.ID, Kind: graph.EdgeMemberOf}, + {From: fixture.foldB.ID, To: fixture.foldOwner.ID, Kind: graph.EdgeMemberOf}, + } + return fixture +} + +func addExploreImplementationFixtureToStore(store graph.Store, fixture exploreImplementationStoreFixture) { + store.AddBatch(fixture.nodes, fixture.edges) +} + +func requireExploreImplementationStoreParity(t *testing.T, store graph.Store) { + t.Helper() + fixture := newExploreImplementationStoreFixture() + addExploreImplementationFixtureToStore(store, fixture) + server := &Server{engine: query.NewEngine(store)} + + expanded := server.expandImplementationTargets( + context.Background(), []exploreTarget{{node: fixture.seed, score: 1}}, graph.LocalizationNodeScope{}, + ) + if len(expanded) != 2 || expanded[0].node.ID != fixture.seed.ID || expanded[1].node.ID != fixture.member.ID { + t.Fatalf("implementation expansion = %#v", expanded) + } + if !exploreImplementationAnswerBlocked( + context.Background(), "find concrete implementation", []exploreTarget{{node: fixture.seed}}, + store, graph.LocalizationNodeScope{}, + ) { + t.Fatal("abstract-only implementation head did not block answer_ready") + } + + folded := server.foldMemberOwners(context.Background(), []exploreTarget{ + {node: fixture.foldA, score: 1}, + {node: fixture.foldB, score: .9}, + }, graph.LocalizationNodeScope{}) + if len(folded) != 3 || folded[0].node.ID != fixture.foldOwner.ID { + t.Fatalf("owner fold = %#v", folded) + } +} + +func TestImplementationRecoveryGraphSQLiteParity(t *testing.T) { + t.Run("graph", func(t *testing.T) { + requireExploreImplementationStoreParity(t, graph.New()) + }) + t.Run("sqlite", func(t *testing.T) { + store, err := store_sqlite.Open(filepath.Join(t.TempDir(), "implementation.sqlite")) + if err != nil { + t.Fatal(err) + } + defer func() { + if err := store.Close(); err != nil { + t.Errorf("close sqlite store: %v", err) + } + }() + requireExploreImplementationStoreParity(t, store) + }) +} diff --git a/internal/mcp/explore_implementation_overlay_view_test.go b/internal/mcp/explore_implementation_overlay_view_test.go new file mode 100644 index 00000000..419fe6b3 --- /dev/null +++ b/internal/mcp/explore_implementation_overlay_view_test.go @@ -0,0 +1,90 @@ +package mcp + +import ( + "context" + "testing" + + "github.com/zzet/gortex/internal/graph" + "github.com/zzet/gortex/internal/query" +) + +func exploreImplementationOverlayExpands( + ctx context.Context, + server *Server, + fixture exploreImplementationStoreFixture, +) bool { + expanded := server.expandImplementationTargets( + ctx, []exploreTarget{{node: fixture.seed, score: 1}}, graph.LocalizationNodeScope{}, + ) + return len(expanded) == 2 && expanded[1].node != nil && expanded[1].node.ID == fixture.member.ID +} + +func TestImplementationExpansionUsesEngineOverlayReader(t *testing.T) { + fixture := newExploreImplementationStoreFixture() + routeStore := graph.New() + addExploreImplementationFixtureToStore(routeStore, fixture) + server := &Server{graph: graph.New(), engine: query.NewEngine(routeStore)} + if !exploreImplementationOverlayExpands(context.Background(), server, fixture) { + t.Fatal("base engine reader lost its independently configured graph") + } + + tests := []struct { + name string + want bool + make func() *graph.OverlayLayer + }{ + { + name: "detached source tombstone", + make: func() *graph.OverlayLayer { + layer := graph.NewOverlayLayer() + layer.MarkRemoved(fixture.seed.Name, fixture.seed.ID) + return layer + }, + }, + { + name: "detached source same-ID replacement", + want: true, + make: func() *graph.OverlayLayer { + layer := graph.NewOverlayLayer() + layer.MarkRemoved(fixture.seed.Name, fixture.seed.ID) + layer.AddNode(fixture.seed.FilePath, fixture.seed) + layer.AddEdge(&graph.Edge{From: fixture.seed.ID, To: fixture.owner.ID, Kind: graph.EdgeMemberOf}) + return layer + }, + }, + { + name: "detached target tombstone", + make: func() *graph.OverlayLayer { + layer := graph.NewOverlayLayer() + layer.MarkRemoved(fixture.owner.Name, fixture.owner.ID) + return layer + }, + }, + { + name: "detached target same-ID replacement", + want: true, + make: func() *graph.OverlayLayer { + layer := graph.NewOverlayLayer() + layer.MarkRemoved(fixture.owner.Name, fixture.owner.ID) + layer.AddNode(fixture.owner.FilePath, fixture.owner) + return layer + }, + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + view := graph.NewOverlaidView(routeStore, test.make()) + ctx := WithOverlayView(context.Background(), view) + eng := server.engineFor(ctx) + if eng == nil || eng.Reader() != view { + t.Fatal("engineFor did not bind the request-local overlay reader") + } + if got := exploreImplementationOverlayExpands(ctx, server, fixture); got != test.want { + t.Fatalf("expanded=%v, want %v", got, test.want) + } + }) + } + if !exploreImplementationOverlayExpands(context.Background(), server, fixture) { + t.Fatal("overlay requests mutated the base engine reader") + } +} diff --git a/internal/mcp/explore_owner_folding.go b/internal/mcp/explore_owner_folding.go index 193be184..9507d8ff 100644 --- a/internal/mcp/explore_owner_folding.go +++ b/internal/mcp/explore_owner_folding.go @@ -22,61 +22,110 @@ import ( // members, and Q2's expansion exists to surface exactly those. const ( - exploreOwnerFoldScan = 8 - exploreOwnerFoldMax = 2 + exploreOwnerFoldScan = 8 + exploreOwnerFoldMax = 2 + exploreOwnerFoldRelationLimit = 8 ) // foldMemberOwners promotes an owner type above its members when at least // two of the scanned top candidates belong to it. The owner node is pulled // from the candidate list when it is already present (its later occurrence // is removed), or fetched by one bounded member_of hop otherwise. -func (s *Server) foldMemberOwners(ctx context.Context, targets []exploreTarget) []exploreTarget { +func (s *Server) foldMemberOwners( + ctx context.Context, + targets []exploreTarget, + declarationScope graph.LocalizationNodeScope, +) []exploreTarget { eng := s.engineFor(ctx) if eng == nil || len(targets) < 2 { return targets } - ownerOf := func(n *graph.Node) *graph.Node { - if n == nil || (n.Kind != graph.KindMethod && n.Kind != graph.KindFunction && n.Kind != graph.KindField) { - return nil + reader := eng.Reader() + if reader == nil || ctx.Err() != nil { + return targets + } + originalTargets := targets + memberIDs := make([]string, 0, min(len(targets), exploreOwnerFoldScan)) + for index, target := range targets { + if index >= exploreOwnerFoldScan { + break + } + if target.node == nil || (target.node.Kind != graph.KindMethod && + target.node.Kind != graph.KindFunction && target.node.Kind != graph.KindField) { + continue + } + memberIDs = append(memberIDs, target.node.ID) + } + if len(memberIDs) == 0 { + return targets + } + relations, complete := exploreBoundedEdgeIdentities( + ctx, reader, memberIDs, []graph.EdgeKind{graph.EdgeMemberOf}, + exploreOwnerFoldRelationLimit, exploreBoundedOutgoing, + ) + if !complete { + return originalTargets + } + ownerIDs := make([]string, 0, len(memberIDs)*exploreOwnerFoldRelationLimit) + for _, memberID := range memberIDs { + for _, edge := range relations[memberID] { + ownerIDs = append(ownerIDs, edge.To) } - for _, e := range eng.GetOutEdges(n.ID) { - if e == nil || e.Kind != graph.EdgeMemberOf { + } + owners, complete := exploreNodesByIDsBounded( + ctx, reader, ownerIDs, exploreOwnerFoldScan*exploreOwnerFoldRelationLimit, + ) + if !complete { + return originalTargets + } + ownerOf := func(memberID string) (*graph.Node, bool) { + for _, edge := range relations[memberID] { + owner := owners[edge.To] + if owner == nil || owner.ID != edge.To { + return nil, false + } + if !declarationScope.Allows(owner) { continue } - if owner := eng.GetSymbol(e.To); owner != nil && - (owner.Kind == graph.KindType || owner.Kind == graph.KindInterface) { - return owner + if owner.Kind == graph.KindType || owner.Kind == graph.KindInterface { + return owner, true } } - return nil + return nil, true } type ownerGroup struct { - owner *graph.Node - firstMember int - members int + owner *graph.Node + firstMemberID string + members int } groups := map[string]*ownerGroup{} order := make([]string, 0, exploreOwnerFoldScan) rankOf := map[string]int{} - for index, t := range targets { - if t.node != nil { - rankOf[t.node.ID] = index + for index, target := range targets { + if target.node != nil { + rankOf[target.node.ID] = index } - if index >= exploreOwnerFoldScan || t.node == nil { + if index >= exploreOwnerFoldScan || target.node == nil { continue } - owner := ownerOf(t.node) + owner, complete := ownerOf(target.node.ID) + if !complete { + return originalTargets + } if owner == nil { continue } - g, ok := groups[owner.ID] + group, ok := groups[owner.ID] if !ok { - g = &ownerGroup{owner: owner, firstMember: index} - groups[owner.ID] = g + group = &ownerGroup{owner: owner, firstMemberID: target.node.ID} + groups[owner.ID] = group order = append(order, owner.ID) } - g.members++ + group.members++ + } + if ctx.Err() != nil { + return originalTargets } folded := 0 @@ -84,11 +133,15 @@ func (s *Server) foldMemberOwners(ctx context.Context, targets []exploreTarget) if folded >= exploreOwnerFoldMax { break } - g := groups[ownerID] - if g.members < 2 { + group := groups[ownerID] + if group.members < 2 { continue } - if existing, present := rankOf[ownerID]; present && existing <= g.firstMember { + firstMember, present := rankOf[group.firstMemberID] + if !present { + return originalTargets + } + if existing, present := rankOf[ownerID]; present && existing <= firstMember { continue // the owner already leads its members } // Remove a lower-ranked occurrence of the owner, then insert it @@ -96,18 +149,18 @@ func (s *Server) foldMemberOwners(ctx context.Context, targets []exploreTarget) kept := make([]exploreTarget, 0, len(targets)+1) var ownerTarget exploreTarget found := false - for _, t := range targets { - if t.node != nil && t.node.ID == ownerID { - ownerTarget = t + for _, target := range targets { + if target.node != nil && target.node.ID == ownerID { + ownerTarget = target found = true continue } - kept = append(kept, t) + kept = append(kept, target) } if !found { - ownerTarget = exploreTarget{node: g.owner, foldedOwner: true} + ownerTarget = exploreTarget{node: group.owner, foldedOwner: true} } - insertAt := g.firstMember + insertAt := firstMember if insertAt > len(kept) { insertAt = len(kept) } @@ -117,12 +170,15 @@ func (s *Server) foldMemberOwners(ctx context.Context, targets []exploreTarget) targets = append(kept[:insertAt:insertAt], append([]exploreTarget{ownerTarget}, kept[insertAt:]...)...) folded++ rankOf = map[string]int{} - for index, t := range targets { - if t.node != nil { - rankOf[t.node.ID] = index + for index, target := range targets { + if target.node != nil { + rankOf[target.node.ID] = index } } } + if ctx.Err() != nil { + return originalTargets + } return targets } @@ -272,7 +328,7 @@ func exploreFoldedTargetMandatory(target exploreTarget, reserved map[string]stru } return target.sourceLiteral || target.sourceLiteralCallee || target.exactContent || target.causalChangeBridge || target.causalChangeLeaf || target.causalChangeOwner || - target.conceptImplementation || target.conceptComplement || target.syntacticAnchor || target.typedAnchorProjection || + target.conceptImplementation || target.conceptComplement || target.syntacticAnchor || target.sourceRange || target.typedAnchorProjection || target.divergentDefaultOwner || target.divergentDefaultType } diff --git a/internal/mcp/explore_owner_folding_overlay_view_test.go b/internal/mcp/explore_owner_folding_overlay_view_test.go new file mode 100644 index 00000000..85371b52 --- /dev/null +++ b/internal/mcp/explore_owner_folding_overlay_view_test.go @@ -0,0 +1,91 @@ +package mcp + +import ( + "context" + "testing" + + "github.com/zzet/gortex/internal/graph" + "github.com/zzet/gortex/internal/query" +) + +func exploreOwnerFoldingOverlayFolds( + ctx context.Context, + server *Server, + fixture exploreImplementationStoreFixture, +) bool { + folded := server.foldMemberOwners(ctx, []exploreTarget{ + {node: fixture.foldA, score: 1}, + {node: fixture.foldB, score: .9}, + }, graph.LocalizationNodeScope{}) + return len(folded) == 3 && folded[0].node != nil && folded[0].node.ID == fixture.foldOwner.ID +} + +func TestOwnerFoldingUsesEngineOverlayReader(t *testing.T) { + fixture := newExploreImplementationStoreFixture() + routeStore := graph.New() + addExploreImplementationFixtureToStore(routeStore, fixture) + server := &Server{graph: graph.New(), engine: query.NewEngine(routeStore)} + if !exploreOwnerFoldingOverlayFolds(context.Background(), server, fixture) { + t.Fatal("base engine reader lost its independently configured graph") + } + + tests := []struct { + name string + want bool + make func() *graph.OverlayLayer + }{ + { + name: "detached source tombstone", + make: func() *graph.OverlayLayer { + layer := graph.NewOverlayLayer() + layer.MarkRemoved(fixture.foldA.Name, fixture.foldA.ID) + return layer + }, + }, + { + name: "detached source same-ID replacement", + want: true, + make: func() *graph.OverlayLayer { + layer := graph.NewOverlayLayer() + layer.MarkRemoved(fixture.foldA.Name, fixture.foldA.ID) + layer.AddNode(fixture.foldA.FilePath, fixture.foldA) + layer.AddEdge(&graph.Edge{From: fixture.foldA.ID, To: fixture.foldOwner.ID, Kind: graph.EdgeMemberOf}) + return layer + }, + }, + { + name: "detached target tombstone", + make: func() *graph.OverlayLayer { + layer := graph.NewOverlayLayer() + layer.MarkRemoved(fixture.foldOwner.Name, fixture.foldOwner.ID) + return layer + }, + }, + { + name: "detached target same-ID replacement", + want: true, + make: func() *graph.OverlayLayer { + layer := graph.NewOverlayLayer() + layer.MarkRemoved(fixture.foldOwner.Name, fixture.foldOwner.ID) + layer.AddNode(fixture.foldOwner.FilePath, fixture.foldOwner) + return layer + }, + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + view := graph.NewOverlaidView(routeStore, test.make()) + ctx := WithOverlayView(context.Background(), view) + eng := server.engineFor(ctx) + if eng == nil || eng.Reader() != view { + t.Fatal("engineFor did not bind the request-local overlay reader") + } + if got := exploreOwnerFoldingOverlayFolds(ctx, server, fixture); got != test.want { + t.Fatalf("folded=%v, want %v", got, test.want) + } + }) + } + if !exploreOwnerFoldingOverlayFolds(context.Background(), server, fixture) { + t.Fatal("overlay requests mutated the base engine reader") + } +} diff --git a/internal/mcp/explore_owner_folding_test.go b/internal/mcp/explore_owner_folding_test.go index c4effdc9..4d3ed575 100644 --- a/internal/mcp/explore_owner_folding_test.go +++ b/internal/mcp/explore_owner_folding_test.go @@ -2,6 +2,7 @@ package mcp import ( "context" + "fmt" "testing" "go.uber.org/zap" @@ -35,7 +36,7 @@ func TestFoldMemberOwnersPromotesSharedOwner(t *testing.T) { {node: g.GetNode("repo/mw/other.ts::unrelated"), score: 0.9}, {node: g.GetNode("repo/mw/persist.ts::PersistOptions.hydrate"), score: 0.8}, } - folded := s.foldMemberOwners(context.Background(), targets) + folded := s.foldMemberOwners(context.Background(), targets, graph.LocalizationNodeScope{}) if folded[0].node.ID != "repo/mw/persist.ts::PersistOptions" { t.Fatalf("owner not promoted ahead of first member: head=%s", folded[0].node.ID) } @@ -48,7 +49,7 @@ func TestFoldMemberOwnersPromotesSharedOwner(t *testing.T) { {node: g.GetNode("repo/mw/persist.ts::PersistOptions.serialize"), score: 1.0}, {node: g.GetNode("repo/mw/other.ts::unrelated"), score: 0.9}, } - unfolded := s.foldMemberOwners(context.Background(), single) + unfolded := s.foldMemberOwners(context.Background(), single, graph.LocalizationNodeScope{}) if unfolded[0].node.ID != "repo/mw/persist.ts::PersistOptions.serialize" { t.Fatal("a lone member must not trigger folding") } @@ -59,12 +60,132 @@ func TestFoldMemberOwnersPromotesSharedOwner(t *testing.T) { {node: g.GetNode("repo/mw/persist.ts::PersistOptions.serialize"), score: 0.9}, {node: g.GetNode("repo/mw/persist.ts::PersistOptions.hydrate"), score: 0.8}, } - same := s.foldMemberOwners(context.Background(), led) + same := s.foldMemberOwners(context.Background(), led, graph.LocalizationNodeScope{}) if len(same) != 3 || same[0].node.ID != "repo/mw/persist.ts::PersistOptions" { t.Fatalf("leading owner must be untouched, got %d entries head=%s", len(same), same[0].node.ID) } } +func TestFoldMemberOwnersPromotesTwoGroupsAtCurrentMemberRanks(t *testing.T) { + g := graph.New() + ownerA := &graph.Node{ID: "repo/groups.go::OwnerA", Kind: graph.KindType, Name: "OwnerA", FilePath: "repo/groups.go"} + ownerB := &graph.Node{ID: "repo/groups.go::OwnerB", Kind: graph.KindType, Name: "OwnerB", FilePath: "repo/groups.go"} + memberA1 := &graph.Node{ID: "repo/groups.go::OwnerA.a", Kind: graph.KindMethod, Name: "a", FilePath: "repo/groups.go"} + memberA2 := &graph.Node{ID: "repo/groups.go::OwnerA.b", Kind: graph.KindMethod, Name: "b", FilePath: "repo/groups.go"} + memberB1 := &graph.Node{ID: "repo/groups.go::OwnerB.a", Kind: graph.KindMethod, Name: "a", FilePath: "repo/groups.go"} + memberB2 := &graph.Node{ID: "repo/groups.go::OwnerB.b", Kind: graph.KindMethod, Name: "b", FilePath: "repo/groups.go"} + unrelatedX := &graph.Node{ID: "repo/groups.go::x", Kind: graph.KindFunction, Name: "x", FilePath: "repo/groups.go"} + unrelatedY := &graph.Node{ID: "repo/groups.go::y", Kind: graph.KindFunction, Name: "y", FilePath: "repo/groups.go"} + g.AddBatch([]*graph.Node{ownerA, ownerB, memberA1, memberA2, memberB1, memberB2, unrelatedX, unrelatedY}, []*graph.Edge{ + {From: memberA1.ID, To: ownerA.ID, Kind: graph.EdgeMemberOf}, + {From: memberA2.ID, To: ownerA.ID, Kind: graph.EdgeMemberOf}, + {From: memberB1.ID, To: ownerB.ID, Kind: graph.EdgeMemberOf}, + {From: memberB2.ID, To: ownerB.ID, Kind: graph.EdgeMemberOf}, + }) + server := NewServer(query.NewEngine(g), g, nil, nil, zap.NewNop(), nil) + targets := []exploreTarget{ + {node: memberA1, score: 1}, + {node: unrelatedX, score: .9}, + {node: memberA2, score: .8}, + {node: memberB1, score: .7}, + {node: unrelatedY, score: .6}, + {node: memberB2, score: .5}, + } + folded := server.foldMemberOwners(context.Background(), targets, graph.LocalizationNodeScope{}) + want := []string{ownerA.ID, memberA1.ID, unrelatedX.ID, memberA2.ID, ownerB.ID, memberB1.ID, unrelatedY.ID, memberB2.ID} + if len(folded) != len(want) { + t.Fatalf("folded len = %d, want %d: %#v", len(folded), len(want), folded) + } + for index, id := range want { + if folded[index].node == nil || folded[index].node.ID != id { + t.Fatalf("folded[%d] = %#v, want %q; full=%#v", index, folded[index], id, folded) + } + } +} + +func TestFoldMemberOwnersBoundedRelationCapAndKindFilter(t *testing.T) { + for _, test := range []struct { + name string + extra int + wantFold bool + }{ + {name: "exact", extra: exploreOwnerFoldRelationLimit - 1, wantFold: true}, + {name: "plus one", extra: exploreOwnerFoldRelationLimit, wantFold: false}, + } { + t.Run(test.name, func(t *testing.T) { + s, g := newOwnerFoldFixture(t) + members := []string{ + "repo/mw/persist.ts::PersistOptions.serialize", + "repo/mw/persist.ts::PersistOptions.hydrate", + } + for index := 0; index < test.extra; index++ { + ownerID := fmt.Sprintf("repo/mw/persist.ts::OtherOwner%02d", index) + g.AddNode(&graph.Node{ID: ownerID, Kind: graph.KindType, Name: fmt.Sprintf("OtherOwner%02d", index), FilePath: "repo/mw/persist.ts"}) + for _, memberID := range members { + g.AddEdge(&graph.Edge{From: memberID, To: ownerID, Kind: graph.EdgeMemberOf}) + } + } + // Unrequested kinds must be filtered before the per-member cap. + for index := 0; index < 32; index++ { + for _, memberID := range members { + g.AddEdge(&graph.Edge{From: memberID, To: fmt.Sprintf("irrelevant-%02d", index), Kind: graph.EdgeReads, Line: index + 1}) + } + } + targets := []exploreTarget{ + {node: g.GetNode(members[0]), score: 1}, + {node: g.GetNode(members[1]), score: .9}, + } + folded := s.foldMemberOwners(context.Background(), targets, graph.LocalizationNodeScope{}) + gotFold := len(folded) == len(targets)+1 && folded[0].node != nil && folded[0].foldedOwner + if gotFold != test.wantFold { + t.Fatalf("folded=%v at %d relevant owner rows, want %v: %#v", gotFold, test.extra+1, test.wantFold, folded) + } + }) + } +} + +func TestFoldMemberOwnersScopesHydratedOwners(t *testing.T) { + const ( + workspace = "workspace" + project = "project" + repo = "repo" + ) + scope := graph.LocalizationNodeScope{ + WorkspaceID: workspace, ProjectID: project, + RepoAllow: map[string]bool{repo: true}, ExcludeTests: true, + } + g := graph.New() + memberA := &graph.Node{ID: "repo/owner.go::Owner.a", Kind: graph.KindMethod, Name: "a", FilePath: "repo/owner.go", WorkspaceID: workspace, ProjectID: project, RepoPrefix: repo} + memberB := &graph.Node{ID: "repo/owner.go::Owner.b", Kind: graph.KindMethod, Name: "b", FilePath: "repo/owner.go", WorkspaceID: workspace, ProjectID: project, RepoPrefix: repo} + allowedOwner := &graph.Node{ID: "repo/owner.go::ZOwner", Kind: graph.KindType, Name: "ZOwner", FilePath: "repo/owner.go", WorkspaceID: workspace, ProjectID: project, RepoPrefix: repo} + blockedOwners := []*graph.Node{ + {ID: "repo/owner.go::AWorkspace", Kind: graph.KindType, Name: "AWorkspace", FilePath: "repo/owner.go", WorkspaceID: "other", ProjectID: project, RepoPrefix: repo}, + {ID: "repo/owner.go::BProject", Kind: graph.KindType, Name: "BProject", FilePath: "repo/owner.go", WorkspaceID: workspace, ProjectID: "other", RepoPrefix: repo}, + {ID: "repo/owner.go::CRepo", Kind: graph.KindType, Name: "CRepo", FilePath: "repo/owner.go", WorkspaceID: workspace, ProjectID: project, RepoPrefix: "other"}, + {ID: "repo/owner.go::DTest", Kind: graph.KindType, Name: "DTest", FilePath: "repo/owner.go", WorkspaceID: workspace, ProjectID: project, RepoPrefix: repo, Meta: map[string]any{"is_test": true}}, + } + g.AddBatch([]*graph.Node{memberA, memberB, allowedOwner}, nil) + for _, owner := range append(blockedOwners, allowedOwner) { + if owner != allowedOwner { + g.AddNode(owner) + } + g.AddEdge(&graph.Edge{From: memberA.ID, To: owner.ID, Kind: graph.EdgeMemberOf}) + g.AddEdge(&graph.Edge{From: memberB.ID, To: owner.ID, Kind: graph.EdgeMemberOf}) + } + server := NewServer(query.NewEngine(g), g, nil, nil, zap.NewNop(), nil) + targets := []exploreTarget{{node: memberA, score: 1}, {node: memberB, score: .9}} + folded := server.foldMemberOwners(context.Background(), targets, scope) + if len(folded) != 3 || folded[0].node.ID != allowedOwner.ID { + t.Fatalf("scope-filtered owner fold = %#v", folded) + } + disjoint := scope + disjoint.WorkspaceID = "other-workspace" + unfolded := server.foldMemberOwners(context.Background(), targets, disjoint) + if len(unfolded) != len(targets) || unfolded[0].node.ID != memberA.ID { + t.Fatalf("out-of-scope owners changed fold: %#v", unfolded) + } +} + func TestLimitExploreFoldedTargetsPreservesReservationsAndCap(t *testing.T) { s, g := newOwnerFoldFixture(t) targets := []exploreTarget{ @@ -72,7 +193,7 @@ func TestLimitExploreFoldedTargetsPreservesReservationsAndCap(t *testing.T) { {node: g.GetNode("repo/mw/other.ts::unrelated"), score: 0.9}, {node: g.GetNode("repo/mw/persist.ts::PersistOptions.hydrate"), score: 0.8, sourceLiteral: true}, } - folded := s.foldMemberOwners(context.Background(), targets) + folded := s.foldMemberOwners(context.Background(), targets, graph.LocalizationNodeScope{}) bounded := limitExploreFoldedTargets("", folded, len(targets), map[string]struct{}{ targets[0].node.ID: {}, targets[2].node.ID: {}, @@ -102,7 +223,7 @@ func TestLimitExploreFoldedTargetsDropsSyntheticOwnerWhenEveryDirectTargetIsRese for _, target := range targets { reserved[target.node.ID] = struct{}{} } - folded := s.foldMemberOwners(context.Background(), targets) + folded := s.foldMemberOwners(context.Background(), targets, graph.LocalizationNodeScope{}) if len(folded) != len(targets)+1 || !folded[0].foldedOwner { t.Fatalf("fixture did not insert a tagged synthetic owner: %#v", folded) } @@ -129,7 +250,7 @@ func TestLimitExploreFoldedTargetsNeverEvictsDirectWindowForSyntheticOwner(t *te {node: g.GetNode("repo/mw/other.ts::unrelated"), score: 0.9}, {node: g.GetNode("repo/mw/persist.ts::PersistOptions.hydrate"), score: 0.8}, } - folded := s.foldMemberOwners(context.Background(), direct) + folded := s.foldMemberOwners(context.Background(), direct, graph.LocalizationNodeScope{}) if len(folded) != len(direct)+1 || !folded[0].foldedOwner { t.Fatalf("fixture did not insert a tagged synthetic owner: %#v", folded) } @@ -161,7 +282,7 @@ func TestLimitExploreFoldedTargetsRetainsDraftSelectedTypeOwnerAtCap(t *testing. {node: g.GetNode("repo/mw/other.ts::unrelated"), score: 0.9}, {node: g.GetNode("repo/mw/persist.ts::PersistOptions.hydrate"), score: 0.8}, } - folded := s.foldMemberOwners(context.Background(), direct) + folded := s.foldMemberOwners(context.Background(), direct, graph.LocalizationNodeScope{}) bounded := limitExploreFoldedTargets( "Find the persist middleware options type that owns serialize and hydrate behavior.", folded, @@ -246,7 +367,7 @@ func TestLimitExploreFoldedTargetsGenericTypeWordsDoNotRetainUnrelatedOwner(t *t {node: g.GetNode("repo/mw/other.ts::unrelated"), score: 0.9}, {node: g.GetNode("repo/mw/persist.ts::PersistOptions.hydrate"), score: 0.8}, } - folded := s.foldMemberOwners(context.Background(), direct) + folded := s.foldMemberOwners(context.Background(), direct, graph.LocalizationNodeScope{}) bounded := limitExploreFoldedTargets( "Find the type config options that own this behavior.", folded, diff --git a/internal/mcp/explore_qualified_leaf.go b/internal/mcp/explore_qualified_leaf.go new file mode 100644 index 00000000..b4584ee9 --- /dev/null +++ b/internal/mcp/explore_qualified_leaf.go @@ -0,0 +1,287 @@ +package mcp + +import ( + "context" + "strings" + + "github.com/zzet/gortex/internal/graph" + "github.com/zzet/gortex/internal/query" + "github.com/zzet/gortex/internal/search/rerank" +) + +const exploreQualifiedLeafEndpointLimit = exploreSyntacticAnchorFetch * exploreExactNameAnchorOwnerScan + +func exploreQualifiedAnchorParts(qualifiedName string) (owner, member string, ok bool) { + dot := strings.LastIndexByte(qualifiedName, '.') + if dot <= 0 || dot == len(qualifiedName)-1 { + return "", "", false + } + owner = exploreQualifiedIdentifierLeaf(qualifiedName[:dot]) + member = exploreQualifiedIdentifierLeaf(qualifiedName[dot+1:]) + return owner, member, owner != "" && member != "" +} + +func exploreQualifiedLeafMatchesNode(node *graph.Node, leaf string) bool { + if node == nil || leaf == "" { + return false + } + for _, identifier := range []string{node.Name, node.QualName, node.ID} { + if strings.EqualFold(exploreQualifiedIdentifierLeaf(identifier), leaf) { + return true + } + } + return false +} + +type exploreQualifiedOwnerFile struct { + path string + ownerIDs map[string]struct{} +} + +// exploreRankedQualifiedOwnerFiles intersects exact owner declarations with +// the ordinary result order. Thus a common owner name cannot authorize a new +// file: independent retrieval must already have ranked that file. +func (s *Server) exploreRankedQualifiedOwnerFiles( + ctx context.Context, + reader graph.Reader, + owner string, + ordinary []*rerank.Candidate, + scope query.QueryOptions, +) []exploreQualifiedOwnerFile { + page, ok := boundedLocalizationExactName( + ctx, + reader, + owner, + s.localizationNodeScope(ctx, scope, graph.KindType, graph.KindInterface), + exploreExactNameAnchorOwnerRawScan, + ) + if !ok { + return nil + } + ownersByFile := make(map[string]map[string]struct{}) + rawScanned, eligibleScanned := 0, 0 + for _, node := range page.Nodes { + if rawScanned == exploreExactNameAnchorOwnerRawScan { + break + } + rawScanned++ + if node == nil || node.Name != owner || node.FilePath == "" || + (node.Kind != graph.KindType && node.Kind != graph.KindInterface) || + !scope.ScopeAllows(node) || !s.nodeInSessionScope(ctx, node) { + continue + } + if eligibleScanned == exploreExactNameAnchorOwnerScan { + break + } + eligibleScanned++ + ownerIDs := ownersByFile[node.FilePath] + if ownerIDs == nil { + ownerIDs = make(map[string]struct{}) + ownersByFile[node.FilePath] = ownerIDs + } + ownerIDs[node.ID] = struct{}{} + } + if len(ownersByFile) == 0 { + return nil + } + files := make([]exploreQualifiedOwnerFile, 0, exploreQualifiedLeafMaxFiles) + seen := make(map[string]struct{}, exploreQualifiedLeafMaxFiles) + for _, candidate := range ordinary { + if candidate == nil || candidate.Node == nil { + continue + } + path := candidate.Node.FilePath + ownerIDs, ownerFile := ownersByFile[path] + if !ownerFile { + continue + } + if _, duplicate := seen[path]; duplicate { + continue + } + seen[path] = struct{}{} + files = append(files, exploreQualifiedOwnerFile{path: path, ownerIDs: ownerIDs}) + if len(files) == exploreQualifiedLeafMaxFiles { + break + } + } + return files +} + +type exploreQualifiedLeafOwnerStatus uint8 + +const ( + exploreQualifiedLeafOwnerUnknown exploreQualifiedLeafOwnerStatus = iota + exploreQualifiedLeafOwnerMatch + exploreQualifiedLeafOwnerMismatch +) + +type exploreQualifiedLeafMatch struct { + node *graph.Node + ownerIDs map[string]struct{} + proven bool +} + +func (s *Server) exploreQualifiedLeafCandidate( + ctx context.Context, + reader graph.Reader, + anchor exploreSyntacticAnchor, + ordinary []*rerank.Candidate, + scope query.QueryOptions, + usedIDs, usedFiles map[string]struct{}, +) *rerank.Candidate { + owner, member, ok := exploreQualifiedAnchorParts(anchor.qualifiedName) + if !ok || ctx.Err() != nil { + return nil + } + files := s.exploreRankedQualifiedOwnerFiles(ctx, reader, owner, ordinary, scope) + if len(files) == 0 { + return nil + } + nodeScope := s.localizationNodeScope( + ctx, scope, graph.KindFunction, graph.KindMethod, graph.KindType, graph.KindMacro, + ) + budget := localizationFileBudgetFor(ctx) + matches := make([]exploreQualifiedLeafMatch, 0, exploreSyntacticAnchorFetch) + seen := make(map[string]struct{}, exploreSyntacticAnchorFetch) + for _, file := range files { + if ctx.Err() != nil { + return nil + } + if len(matches) == exploreSyntacticAnchorFetch { + break + } + page, complete := boundedLocalizationFileNodes( + ctx, reader, budget, file.path, nodeScope, localizationFileNodeLimit, + ) + if !complete { + return nil + } + for _, node := range page.Nodes { + if node == nil || !exploreQualifiedLeafMatchesNode(node, member) || + !exploreSyntacticAnchorEligibleNode(node) || !scope.ScopeAllows(node) || + !s.nodeInSessionScope(ctx, node) { + continue + } + if _, duplicate := seen[node.ID]; duplicate { + continue + } + seen[node.ID] = struct{}{} + ownerStatus := exploreQualifiedLeafIdentityOwnerStatus(node, owner, member) + if ownerStatus == exploreQualifiedLeafOwnerMismatch { + continue + } + matches = append(matches, exploreQualifiedLeafMatch{ + node: node, ownerIDs: file.ownerIDs, + proven: ownerStatus == exploreQualifiedLeafOwnerMatch, + }) + if len(matches) == exploreSyntacticAnchorFetch { + break + } + } + } + if len(matches) == 0 { + return nil + } + + endpoints := make([]graph.TypedEdgeEndpoint, 0, len(matches)*exploreExactNameAnchorOwnerScan) + for _, match := range matches { + if match.proven { + continue + } + for ownerID := range match.ownerIDs { + endpoints = append(endpoints, graph.TypedEdgeEndpoint{ + From: match.node.ID, To: ownerID, Kind: graph.EdgeMemberOf, + }) + if len(endpoints) > exploreQualifiedLeafEndpointLimit { + return nil + } + } + } + if len(endpoints) > 0 { + bounded, ok := reader.(graph.BoundedEdgeExistenceReader) + if !ok { + return nil + } + existing, err := bounded.FindExistingEdgeEndpoints( + ctx, endpoints, exploreQualifiedLeafEndpointLimit, + ) + if err != nil || ctx.Err() != nil { + return nil + } + for index := range matches { + if matches[index].proven { + continue + } + for ownerID := range matches[index].ownerIDs { + key := graph.TypedEdgeEndpoint{ + From: matches[index].node.ID, To: ownerID, Kind: graph.EdgeMemberOf, + } + if _, exactOwner := existing[key]; exactOwner { + matches[index].proven = true + break + } + } + } + } + + var provenFallback *graph.Node + for _, match := range matches { + if !match.proven { + continue + } + if _, used := usedIDs[match.node.ID]; used { + continue + } + if provenFallback == nil { + provenFallback = match.node + } + if _, repeatedFile := usedFiles[match.node.FilePath]; !repeatedFile { + return &rerank.Candidate{Node: match.node, VectorRank: -1} + } + } + if provenFallback != nil { + return &rerank.Candidate{Node: provenFallback, VectorRank: -1} + } + // Parsers that do not retain enclosing-owner metadata still recover one + // globally unique leaf. Any second leaf makes ownership ambiguous and the + // qualified fallback deliberately declines. + if len(matches) != 1 { + return nil + } + if _, used := usedIDs[matches[0].node.ID]; used { + return nil + } + return &rerank.Candidate{Node: matches[0].node, VectorRank: -1} +} + +func exploreQualifiedLeafIdentityOwnerStatus(node *graph.Node, owner, member string) exploreQualifiedLeafOwnerStatus { + if node == nil || owner == "" || member == "" { + return exploreQualifiedLeafOwnerUnknown + } + status := exploreQualifiedLeafOwnerUnknown + identities := []string{node.QualName, node.Name, node.ID} + for index, identity := range identities { + identity = strings.TrimSpace(identity) + if index == len(identities)-1 { + // Graph IDs carry a file/repo prefix before the final ::. Only the + // symbol suffix can describe an enclosing declaration. + if cut := strings.LastIndex(identity, "::"); cut >= 0 { + identity = identity[cut+2:] + } + } + identity = strings.ReplaceAll(identity, "::", ".") + separator := strings.LastIndexByte(identity, '.') + if separator <= 0 || !strings.EqualFold(exploreQualifiedIdentifierLeaf(identity[separator+1:]), member) { + continue + } + ownerPart := exploreQualifiedIdentifierLeaf(identity[:separator]) + if ownerPart == "" { + continue + } + if strings.EqualFold(ownerPart, owner) { + return exploreQualifiedLeafOwnerMatch + } + status = exploreQualifiedLeafOwnerMismatch + } + return status +} diff --git a/internal/mcp/explore_qualified_leaf_edge_bounds_test.go b/internal/mcp/explore_qualified_leaf_edge_bounds_test.go new file mode 100644 index 00000000..c3fe50a1 --- /dev/null +++ b/internal/mcp/explore_qualified_leaf_edge_bounds_test.go @@ -0,0 +1,226 @@ +package mcp + +import ( + "context" + "errors" + "fmt" + "testing" + + "github.com/zzet/gortex/internal/graph" + "github.com/zzet/gortex/internal/query" + "github.com/zzet/gortex/internal/search/rerank" +) + +type qualifiedLeafEdgeProjectionReader struct { + *exactNameAnchorCountingReader + edgeCalls int + endpoints []graph.TypedEdgeEndpoint + edgeLimit int + projection func(context.Context, []graph.TypedEdgeEndpoint, int) (map[graph.TypedEdgeEndpoint]struct{}, error) +} + +func (reader *qualifiedLeafEdgeProjectionReader) FindExistingEdgeEndpoints( + ctx context.Context, + endpoints []graph.TypedEdgeEndpoint, + limit int, +) (map[graph.TypedEdgeEndpoint]struct{}, error) { + reader.edgeCalls++ + reader.endpoints = append([]graph.TypedEdgeEndpoint(nil), endpoints...) + reader.edgeLimit = limit + if reader.projection != nil { + return reader.projection(ctx, endpoints, limit) + } + return map[graph.TypedEdgeEndpoint]struct{}{}, nil +} + +func (*qualifiedLeafEdgeProjectionReader) GetOutEdgesByNodeIDs([]string) map[string][]*graph.Edge { + panic("qualified-leaf proof must not materialize legacy outgoing adjacency") +} + +type qualifiedLeafLegacyPanicReader struct { + *exactNameAnchorCountingReader +} + +func (*qualifiedLeafLegacyPanicReader) GetOutEdgesByNodeIDs([]string) map[string][]*graph.Edge { + panic("qualified-leaf proof must fail closed without the bounded edge capability") +} + +func TestExploreQualifiedLeafUsesBoundedExactEdgeProof(t *testing.T) { + anchor, owner, _ := qualifiedLeafFixture() + requested := &graph.Node{ + ID: owner.FilePath + "::buildContent#requested", Name: "buildContent", + Kind: graph.KindMethod, FilePath: owner.FilePath, StartLine: 40, + } + ambiguous := &graph.Node{ + ID: owner.FilePath + "::buildContent#ambiguous", Name: "buildContent", + Kind: graph.KindMethod, FilePath: owner.FilePath, StartLine: 50, + } + base := graph.New() + for _, node := range []*graph.Node{owner, requested, ambiguous} { + base.AddNode(node) + } + proof := graph.TypedEdgeEndpoint{From: requested.ID, To: owner.ID, Kind: graph.EdgeMemberOf} + reader := &qualifiedLeafEdgeProjectionReader{ + exactNameAnchorCountingReader: &exactNameAnchorCountingReader{Reader: base}, + projection: func(_ context.Context, endpoints []graph.TypedEdgeEndpoint, limit int) (map[graph.TypedEdgeEndpoint]struct{}, error) { + return map[graph.TypedEdgeEndpoint]struct{}{proof: {}}, nil + }, + } + server := &Server{graph: base} + got := server.exploreQualifiedLeafCandidate( + context.Background(), reader, anchor, []*rerank.Candidate{{Node: owner}}, + query.QueryOptions{}, map[string]struct{}{}, map[string]struct{}{}, + ) + if got == nil || got.Node == nil || got.Node.ID != requested.ID { + t.Fatalf("qualified leaf = %#v, want exact member_of proof for %q", got, requested.ID) + } + if reader.edgeCalls != 1 || reader.edgeLimit != exploreQualifiedLeafEndpointLimit { + t.Fatalf("edge projection calls/limit = %d/%d, want 1/%d", reader.edgeCalls, reader.edgeLimit, exploreQualifiedLeafEndpointLimit) + } + if len(reader.endpoints) != 2 { + t.Fatalf("edge predicates = %#v, want the two exact member_of alternatives", reader.endpoints) + } + seen := make(map[graph.TypedEdgeEndpoint]struct{}, len(reader.endpoints)) + for _, endpoint := range reader.endpoints { + seen[endpoint] = struct{}{} + if endpoint.Kind != graph.EdgeMemberOf || endpoint.To != owner.ID { + t.Fatalf("unexpected edge predicate: %#v", endpoint) + } + } + if _, ok := seen[proof]; !ok { + t.Fatalf("edge predicates = %#v, missing requested proof %#v", reader.endpoints, proof) + } +} + +func TestExploreQualifiedLeafEdgeProofFailsClosed(t *testing.T) { + anchor, owner, _ := qualifiedLeafFixture() + bare := &graph.Node{ + ID: owner.FilePath + "::buildContent", Name: "buildContent", + Kind: graph.KindMethod, FilePath: owner.FilePath, StartLine: 40, + } + base := graph.New() + base.AddNode(owner) + base.AddNode(bare) + server := &Server{graph: base} + ordinary := []*rerank.Candidate{{Node: owner}} + + t.Run("unsupported", func(t *testing.T) { + reader := &qualifiedLeafLegacyPanicReader{ + exactNameAnchorCountingReader: &exactNameAnchorCountingReader{Reader: base}, + } + if got := server.exploreQualifiedLeafCandidate( + context.Background(), reader, anchor, ordinary, query.QueryOptions{}, + map[string]struct{}{}, map[string]struct{}{}, + ); got != nil { + t.Fatalf("qualified leaf = %#v, want missing bounded capability to fail closed", got) + } + }) + + for _, test := range []struct { + name string + err error + }{ + {name: "error", err: errors.New("projection failed")}, + {name: "canceled", err: context.Canceled}, + } { + t.Run(test.name, func(t *testing.T) { + reader := &qualifiedLeafEdgeProjectionReader{ + exactNameAnchorCountingReader: &exactNameAnchorCountingReader{Reader: base}, + projection: func(context.Context, []graph.TypedEdgeEndpoint, int) (map[graph.TypedEdgeEndpoint]struct{}, error) { + return map[graph.TypedEdgeEndpoint]struct{}{ + {From: bare.ID, To: owner.ID, Kind: graph.EdgeMemberOf}: {}, + }, test.err + }, + } + if got := server.exploreQualifiedLeafCandidate( + context.Background(), reader, anchor, ordinary, query.QueryOptions{}, + map[string]struct{}{}, map[string]struct{}{}, + ); got != nil { + t.Fatalf("qualified leaf = %#v, want %v to discard partial proof", got, test.err) + } + if reader.edgeCalls != 1 { + t.Fatalf("edge projection calls = %d, want one bounded attempt", reader.edgeCalls) + } + }) + } + + t.Run("request canceled after projection", func(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + reader := &qualifiedLeafEdgeProjectionReader{ + exactNameAnchorCountingReader: &exactNameAnchorCountingReader{Reader: base}, + projection: func(context.Context, []graph.TypedEdgeEndpoint, int) (map[graph.TypedEdgeEndpoint]struct{}, error) { + cancel() + return map[graph.TypedEdgeEndpoint]struct{}{ + {From: bare.ID, To: owner.ID, Kind: graph.EdgeMemberOf}: {}, + }, nil + }, + } + if got := server.exploreQualifiedLeafCandidate( + ctx, reader, anchor, ordinary, query.QueryOptions{}, + map[string]struct{}{}, map[string]struct{}{}, + ); got != nil { + t.Fatalf("qualified leaf = %#v, want post-projection cancellation to discard proof", got) + } + }) +} + +func TestExploreQualifiedLeafEdgeProofUsesExactPredicateCap(t *testing.T) { + if exploreQualifiedLeafEndpointLimit != graph.MaxBoundedEdgeExistencePredicates { + t.Fatalf("MCP/store predicate caps drifted: %d != %d", exploreQualifiedLeafEndpointLimit, graph.MaxBoundedEdgeExistencePredicates) + } + anchor, _, _ := qualifiedLeafFixture() + const path = "src/HipChatHandler.php" + base := graph.New() + owners := make(map[string]struct{}, exploreExactNameAnchorOwnerScan) + for index := 0; index < exploreExactNameAnchorOwnerScan; index++ { + node := &graph.Node{ + ID: fmt.Sprintf("%s::HipChatHandler#%02d", path, index), Name: "HipChatHandler", + Kind: graph.KindType, FilePath: path, StartLine: index + 1, + } + base.AddNode(node) + owners[node.ID] = struct{}{} + } + leaves := make(map[string]struct{}, exploreSyntacticAnchorFetch) + for index := 0; index < exploreSyntacticAnchorFetch; index++ { + node := &graph.Node{ + ID: fmt.Sprintf("%s::leaf#%02d", path, index), Name: "buildContent", + Kind: graph.KindMethod, FilePath: path, StartLine: 100 + index, + } + base.AddNode(node) + leaves[node.ID] = struct{}{} + } + ranked := &graph.Node{ID: path + "::ranked", Name: "ranked", Kind: graph.KindFunction, FilePath: path} + base.AddNode(ranked) + reader := &qualifiedLeafEdgeProjectionReader{ + exactNameAnchorCountingReader: &exactNameAnchorCountingReader{Reader: base}, + } + server := &Server{graph: base} + if got := server.exploreQualifiedLeafCandidate( + context.Background(), reader, anchor, []*rerank.Candidate{{Node: ranked}}, + query.QueryOptions{}, map[string]struct{}{}, map[string]struct{}{}, + ); got != nil { + t.Fatalf("qualified leaf = %#v, want ten ambiguous unproven leaves", got) + } + if reader.edgeCalls != 1 || reader.edgeLimit != exploreQualifiedLeafEndpointLimit || + len(reader.endpoints) != exploreQualifiedLeafEndpointLimit { + t.Fatalf("edge calls/limit/predicates = %d/%d/%d, want 1/%d/%d", + reader.edgeCalls, reader.edgeLimit, len(reader.endpoints), + exploreQualifiedLeafEndpointLimit, exploreQualifiedLeafEndpointLimit) + } + seen := make(map[graph.TypedEdgeEndpoint]struct{}, len(reader.endpoints)) + for _, endpoint := range reader.endpoints { + if endpoint.Kind != graph.EdgeMemberOf { + t.Fatalf("edge predicate kind = %q, want member_of", endpoint.Kind) + } + if _, ok := leaves[endpoint.From]; !ok { + t.Fatalf("edge predicate source = %q, want one of the ten leaves", endpoint.From) + } + if _, ok := owners[endpoint.To]; !ok { + t.Fatalf("edge predicate target = %q, want one of the 32 owners", endpoint.To) + } + if _, duplicate := seen[endpoint]; duplicate { + t.Fatalf("duplicate edge predicate: %#v", endpoint) + } + seen[endpoint] = struct{}{} + } +} diff --git a/internal/mcp/explore_qualified_leaf_identifier.go b/internal/mcp/explore_qualified_leaf_identifier.go new file mode 100644 index 00000000..07a184ce --- /dev/null +++ b/internal/mcp/explore_qualified_leaf_identifier.go @@ -0,0 +1,14 @@ +package mcp + +import "strings" + +func exploreQualifiedIdentifierLeaf(identifier string) string { + identifier = strings.TrimSpace(identifier) + if separator := strings.LastIndex(identifier, "::"); separator >= 0 { + identifier = identifier[separator+2:] + } + if separator := strings.LastIndexAny(identifier, ".\\/"); separator >= 0 { + identifier = identifier[separator+1:] + } + return identifier +} diff --git a/internal/mcp/explore_qualified_leaf_test.go b/internal/mcp/explore_qualified_leaf_test.go new file mode 100644 index 00000000..4c906515 --- /dev/null +++ b/internal/mcp/explore_qualified_leaf_test.go @@ -0,0 +1,282 @@ +package mcp + +import ( + "context" + "fmt" + "testing" + + "github.com/zzet/gortex/internal/graph" + "github.com/zzet/gortex/internal/query" + "github.com/zzet/gortex/internal/search/rerank" +) + +func qualifiedLeafFixture() (exploreSyntacticAnchor, *graph.Node, *graph.Node) { + anchor, _ := newExploreSyntacticAnchor("HipChatHandler::buildContent") + owner := &graph.Node{ + ID: "src/HipChatHandler.php::HipChatHandler", Name: "HipChatHandler", + Kind: graph.KindType, FilePath: "src/HipChatHandler.php", StartLine: 10, + } + member := &graph.Node{ + ID: "src/HipChatHandler.php::HipChatHandler.buildcontent", + Name: "HipChatHandler::buildcontent", QualName: "HipChatHandler.buildcontent", + Kind: graph.KindMethod, FilePath: "src/HipChatHandler.php", StartLine: 40, + } + return anchor, owner, member +} + +func TestExploreQualifiedLeafRecoversWithinRankedOwnerFileAcrossStores(t *testing.T) { + forEachExactNameAnchorStore(t, func(t *testing.T, store graph.Store) { + anchor, owner, member := qualifiedLeafFixture() + store.AddNode(owner) + store.AddNode(member) + server := &Server{graph: store} + ordinary := []*rerank.Candidate{{Node: owner}} + got := server.exploreExactQualifiedAnchorCandidate( + context.Background(), anchor, ordinary, query.QueryOptions{}, + map[string]struct{}{}, map[string]struct{}{}, + ) + if got == nil || got.Node == nil || got.Node.ID != member.ID { + t.Fatalf("qualified leaf = %#v, want %q", got, member.ID) + } + }) +} + +func TestExploreQualifiedLeafRejectsSiblingOwnerCollision(t *testing.T) { + anchor, owner, _ := qualifiedLeafFixture() + other := &graph.Node{ + ID: "src/HipChatHandler.php::Other.buildContent", Name: "buildContent", + QualName: "Other.buildContent", Kind: graph.KindMethod, + FilePath: owner.FilePath, StartLine: 30, + } + server := exactNameAnchorServer(t, owner, other) + got := server.exploreQualifiedLeafCandidate( + context.Background(), server.graph, anchor, []*rerank.Candidate{{Node: owner}}, + query.QueryOptions{}, map[string]struct{}{}, map[string]struct{}{}, + ) + if got != nil { + t.Fatalf("qualified leaf = %#v, want explicit other-owner member rejected", got) + } +} + +func TestExploreQualifiedLeafChoosesQualifiedOwnerAmongSameFileCollisions(t *testing.T) { + anchor, owner, member := qualifiedLeafFixture() + other := &graph.Node{ + ID: "src/HipChatHandler.php::Other.buildContent", Name: "buildContent", + QualName: "Other.buildContent", Kind: graph.KindMethod, + FilePath: owner.FilePath, StartLine: 30, + } + server := exactNameAnchorServer(t, owner, other, member) + got := server.exploreQualifiedLeafCandidate( + context.Background(), server.graph, anchor, []*rerank.Candidate{{Node: owner}}, + query.QueryOptions{}, map[string]struct{}{}, map[string]struct{}{}, + ) + if got == nil || got.Node == nil || got.Node.ID != member.ID { + t.Fatalf("qualified leaf = %#v, want requested owner member %q", got, member.ID) + } +} + +func TestExploreQualifiedLeafUsesMemberOfToDisambiguateBareLeaves(t *testing.T) { + anchor, owner, _ := qualifiedLeafFixture() + requested := &graph.Node{ + ID: owner.FilePath + "::buildContent#requested", Name: "buildContent", + Kind: graph.KindMethod, FilePath: owner.FilePath, StartLine: 40, + } + ambiguous := &graph.Node{ + ID: owner.FilePath + "::buildContent#ambiguous", Name: "buildContent", + Kind: graph.KindMethod, FilePath: owner.FilePath, StartLine: 50, + } + store := graph.New() + store.AddNode(owner) + store.AddNode(requested) + store.AddNode(ambiguous) + store.AddEdge(&graph.Edge{From: requested.ID, To: owner.ID, Kind: graph.EdgeMemberOf}) + server := &Server{graph: store} + got := server.exploreQualifiedLeafCandidate( + context.Background(), store, anchor, []*rerank.Candidate{{Node: owner}}, + query.QueryOptions{}, map[string]struct{}{}, map[string]struct{}{}, + ) + if got == nil || got.Node == nil || got.Node.ID != requested.ID { + t.Fatalf("qualified leaf = %#v, want member_of-proven leaf %q", got, requested.ID) + } +} + +func TestExploreQualifiedLeafAllowsUniqueUnqualifiedLeaf(t *testing.T) { + anchor, owner, _ := qualifiedLeafFixture() + bare := &graph.Node{ + ID: owner.FilePath + "::buildContent", Name: "buildContent", + Kind: graph.KindMethod, FilePath: owner.FilePath, StartLine: 40, + } + server := exactNameAnchorServer(t, owner, bare) + got := server.exploreQualifiedLeafCandidate( + context.Background(), server.graph, anchor, []*rerank.Candidate{{Node: owner}}, + query.QueryOptions{}, map[string]struct{}{}, map[string]struct{}{}, + ) + if got == nil || got.Node == nil || got.Node.ID != bare.ID { + t.Fatalf("qualified leaf = %#v, want unique bare member %q", got, bare.ID) + } +} + +func TestExploreQualifiedLeafRejectsAmbiguousUnqualifiedLeavesAcrossOwnerFiles(t *testing.T) { + anchor, firstOwner, _ := qualifiedLeafFixture() + secondOwner := &graph.Node{ + ID: "src/Legacy/HipChatHandler.php::HipChatHandler", Name: "HipChatHandler", + Kind: graph.KindType, FilePath: "src/Legacy/HipChatHandler.php", StartLine: 8, + } + first := &graph.Node{ID: firstOwner.FilePath + "::buildContent", Name: "buildContent", Kind: graph.KindMethod, FilePath: firstOwner.FilePath, StartLine: 40} + second := &graph.Node{ID: secondOwner.FilePath + "::buildContent", Name: "buildContent", Kind: graph.KindMethod, FilePath: secondOwner.FilePath, StartLine: 50} + server := exactNameAnchorServer(t, firstOwner, secondOwner, first, second) + ordinary := []*rerank.Candidate{{Node: firstOwner}, {Node: secondOwner}} + got := server.exploreQualifiedLeafCandidate( + context.Background(), server.graph, anchor, ordinary, + query.QueryOptions{}, map[string]struct{}{}, map[string]struct{}{}, + ) + if got != nil { + t.Fatalf("qualified leaf = %#v, want ambiguous bare leaves rejected", got) + } +} + +func TestExploreQualifiedLeafRequiresRankedOwnerFile(t *testing.T) { + anchor, owner, member := qualifiedLeafFixture() + other := &graph.Node{ + ID: "src/Other.php::Other", Name: "Other", + Kind: graph.KindType, FilePath: "src/Other.php", StartLine: 7, + } + server := exactNameAnchorServer(t, owner, member, other) + for name, ordinary := range map[string][]*rerank.Candidate{ + "no ranking evidence": nil, + "different ranked file": {{Node: other}}, + } { + t.Run(name, func(t *testing.T) { + got := server.exploreExactQualifiedAnchorCandidate( + context.Background(), anchor, ordinary, query.QueryOptions{}, + map[string]struct{}{}, map[string]struct{}{}, + ) + if got != nil { + t.Fatalf("qualified leaf = %#v, want no unranked file recovery", got) + } + }) + } +} + +func TestExploreQualifiedLeafUsesOverlayOwnerFile(t *testing.T) { + anchor, owner, member := qualifiedLeafFixture() + base := graph.New() + base.AddNode(&graph.Node{ + ID: owner.ID, Name: owner.Name, Kind: owner.Kind, + FilePath: owner.FilePath, StartLine: owner.StartLine, + }) + layer := graph.NewOverlayLayer() + layer.MarkFile(owner.FilePath, false) + layer.AddNode(owner.FilePath, owner) + layer.AddNode(member.FilePath, member) + ctx := WithOverlayView(context.Background(), graph.NewOverlaidView(base, layer)) + server := &Server{graph: base} + got := server.exploreExactQualifiedAnchorCandidate( + ctx, anchor, []*rerank.Candidate{{Node: owner}}, query.QueryOptions{}, + map[string]struct{}{}, map[string]struct{}{}, + ) + if got == nil || got.Node == nil || got.Node.ID != member.ID { + t.Fatalf("overlay qualified leaf = %#v, want %q", got, member.ID) + } +} + +func TestExploreQualifiedLeafHasRankedFileBudgetAndNoCorpusScan(t *testing.T) { + base := graph.New() + ordinary := make([]*rerank.Candidate, 0, 6) + for index := 0; index < 6; index++ { + path := fmt.Sprintf("src/owner%d.go", index) + owner := &graph.Node{ + ID: path + "::Owner", Name: "Owner", + Kind: graph.KindType, FilePath: path, StartLine: 5, + } + decoy := &graph.Node{ + ID: path + "::decoy", Name: "decoy", + Kind: graph.KindFunction, FilePath: path, StartLine: 20, + } + base.AddNode(owner) + base.AddNode(decoy) + ordinary = append(ordinary, &rerank.Candidate{Node: decoy}) + } + reader := &exactNameAnchorCountingReader{Reader: base} + ctx := WithOverlayView( + context.Background(), graph.NewOverlaidView(reader, graph.NewOverlayLayer()), + ) + server := &Server{graph: base} + anchor, _ := newExploreSyntacticAnchor("Owner::missingLeaf") + got := server.exploreExactQualifiedAnchorCandidate( + ctx, anchor, ordinary, query.QueryOptions{}, + map[string]struct{}{}, map[string]struct{}{}, + ) + if got != nil { + t.Fatalf("qualified leaf = %#v, want no match", got) + } + if len(reader.nameLookups) != 3 { + t.Fatalf("name lookups = %d (%#v), want member, qualified name, and owner", len(reader.nameLookups), reader.nameLookups) + } + if len(reader.fileLookups) != exploreQualifiedLeafMaxFiles { + t.Fatalf("file lookups = %d (%#v), want cap %d", len(reader.fileLookups), reader.fileLookups, exploreQualifiedLeafMaxFiles) + } + totalLimit := 0 + for index, limit := range reader.fileLimits { + if limit <= 0 || limit > localizationFileNodeLimit { + t.Fatalf("file limit[%d] = %d, want (0, %d]", index, limit, localizationFileNodeLimit) + } + totalLimit += limit + scope := reader.fileScopes[index] + if !scope.ExcludeTests || len(scope.Kinds) != 4 || + !scope.Kinds[graph.KindFunction] || !scope.Kinds[graph.KindMethod] || + !scope.Kinds[graph.KindType] || !scope.Kinds[graph.KindMacro] { + t.Fatalf("file scope[%d] = %#v, want production function/method/type/macro declarations", index, scope) + } + } + if totalLimit > localizationFileRequestLimit { + t.Fatalf("aggregate file limit = %d, want <= %d", totalLimit, localizationFileRequestLimit) + } +} + +func TestExploreQualifiedLeafProjectionFailsClosed(t *testing.T) { + anchor, owner, member := qualifiedLeafFixture() + tests := []struct { + name string + readerBase func(*graph.Graph) graph.Reader + projection func(context.Context, string, graph.LocalizationNodeScope, int) (graph.BoundedNodeProjection, error) + }{ + { + name: "truncated", + readerBase: func(base *graph.Graph) graph.Reader { return base }, + projection: func(context.Context, string, graph.LocalizationNodeScope, int) (graph.BoundedNodeProjection, error) { + return graph.BoundedNodeProjection{Nodes: []*graph.Node{member}, Total: 2, Truncated: true}, nil + }, + }, + { + name: "unsupported", + readerBase: func(base *graph.Graph) graph.Reader { + return &exactNameOnlyBoundedReader{Reader: base} + }, + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + base := graph.New() + base.AddNode(owner) + base.AddNode(member) + reader := &exactNameAnchorCountingReader{ + Reader: test.readerBase(base), fileProjection: test.projection, + } + ctx := WithOverlayView( + context.Background(), graph.NewOverlaidView(reader, graph.NewOverlayLayer()), + ) + server := &Server{graph: base} + got := server.exploreExactQualifiedAnchorCandidate( + ctx, anchor, []*rerank.Candidate{{Node: owner}}, query.QueryOptions{}, + map[string]struct{}{}, map[string]struct{}{}, + ) + if got != nil { + t.Fatalf("qualified leaf = %#v, want incomplete file projection to fail closed", got) + } + if len(reader.fileLookups) != 1 { + t.Fatalf("file lookups = %#v, want one bounded attempt", reader.fileLookups) + } + }) + } +} diff --git a/internal/mcp/explore_source_literal.go b/internal/mcp/explore_source_literal.go index 05977eb0..3b5e3f88 100644 --- a/internal/mcp/explore_source_literal.go +++ b/internal/mcp/explore_source_literal.go @@ -13,15 +13,18 @@ import ( "github.com/zzet/gortex/internal/graph" "github.com/zzet/gortex/internal/query" "github.com/zzet/gortex/internal/search/trigram" + "github.com/zzet/gortex/internal/testpath" ) const ( exploreSourceLiteralRecallMaxHits = 24 exploreSourceLiteralRecallMaxFiles = 0 exploreSourceLiteralRecallBudget = 75 * time.Millisecond - exploreSourceLiteralRecallMaxTerms = 2 + exploreSourceLiteralRecallMaxTerms = 4 exploreSourceLiteralRecallMaxOwnersPerTerm = 3 - exploreSourceLiteralRecallMaxFilesPerTerm = 2 + exploreSourceLiteralRecallMaxFilesPerTerm = 3 + exploreSourceLiteralCallEdgesPerSite = 8 + exploreSourceLiteralCalleeHydrationLimit = exploreSourceLiteralRecallMaxHits * exploreSourceLiteralCallEdgesPerSite ) // sourceLiteralRecallBudget is the wall-clock slice one anchor of the bounded @@ -41,6 +44,12 @@ type exploreSourceLiteralHit struct { anchor int ambiguous bool callee bool + // The exact source-match coordinate is retained request-locally so the + // localization envelope may offer one bounded live source window without + // repeating the content search. Raw match text is intentionally discarded. + matchPath string + matchLine int + literal string } type exploreSourceLiteralDiagnostic struct { @@ -65,6 +74,7 @@ type exploreSourceLiteralSearch struct { backend string owned bool lookupRepoPrefix string + err error } // explorePreferredSourceLiteral picks one deterministic source-search key. @@ -111,6 +121,42 @@ func exploreCompactSourceLiteral(term string, runeCount int) bool { return true } +// appendExploreSourceLiteralProseTerms fills the remaining bounded term slots +// with non-compact literals, longest first: a longer quoted value carries more +// information and greps with less noise. Compact codes keep the slots they +// already reserved. +func appendExploreSourceLiteralProseTerms(attemptTerms, terms []string) []string { + if len(attemptTerms) >= exploreSourceLiteralRecallMaxTerms { + return attemptTerms + } + seen := make(map[string]struct{}, len(attemptTerms)+len(terms)) + for _, term := range attemptTerms { + seen[strings.ToLower(term)] = struct{}{} + } + prose := make([]string, 0, len(terms)) + for _, term := range terms { + if exploreCompactSourceLiteral(term, utf8.RuneCountInString(term)) { + continue + } + key := strings.ToLower(term) + if _, duplicate := seen[key]; duplicate { + continue + } + seen[key] = struct{}{} + prose = append(prose, term) + } + sort.SliceStable(prose, func(i, j int) bool { + return utf8.RuneCountInString(prose[i]) > utf8.RuneCountInString(prose[j]) + }) + for _, term := range prose { + if len(attemptTerms) >= exploreSourceLiteralRecallMaxTerms { + break + } + attemptTerms = append(attemptTerms, term) + } + return attemptTerms +} + func exploreSourceLiteralFallback(terms []string, primary string) string { best := "" bestLen := 0 @@ -155,6 +201,23 @@ func retainExploreSourceLiteralOwners(recall exploreSourceLiteralRecall) (hits [ seenOwners := make(map[string]struct{}, min(len(recall.hits), exploreSourceLiteralRecallMaxOwnersPerTerm)) seenFiles := make(map[string]struct{}, exploreSourceLiteralRecallMaxFilesPerTerm) selected := make([]bool, len(recall.hits)) + order := make([]int, len(recall.hits)) + for index := range recall.hits { + order[index] = index + } + // A graph-resolved direct callee is more actionable than its enclosing + // source owner. Prefer it inside the same fixed owner/file caps; ambiguity + // remains attached to the hit and therefore cannot become terminal proof. + sort.SliceStable(order, func(i, j int) bool { + left, right := recall.hits[order[i]], recall.hits[order[j]] + if left.callee != right.callee { + return left.callee + } + if left.rank != right.rank { + return left.rank < right.rank + } + return left.nodeID < right.nodeID + }) add := func(index int) { hit := recall.hits[index] if _, duplicate := seenOwners[hit.nodeID]; duplicate { @@ -169,10 +232,11 @@ func retainExploreSourceLiteralOwners(recall exploreSourceLiteralRecall) (hits [ } // First pass: one declaration from each distinct file. - for index, hit := range recall.hits { + for _, index := range order { if len(hits) >= exploreSourceLiteralRecallMaxOwnersPerTerm || len(seenFiles) >= exploreSourceLiteralRecallMaxFilesPerTerm { break } + hit := recall.hits[index] file := recall.ownerFiles[hit.nodeID] if file == "" { continue @@ -183,13 +247,14 @@ func retainExploreSourceLiteralOwners(recall exploreSourceLiteralRecall) (hits [ add(index) } // Second pass: fill remaining owner slots from already-admitted files. - for index, hit := range recall.hits { + for _, index := range order { if len(hits) >= exploreSourceLiteralRecallMaxOwnersPerTerm { break } if selected[index] { continue } + hit := recall.hits[index] file := recall.ownerFiles[hit.nodeID] if file != "" { if _, admitted := seenFiles[file]; !admitted && len(seenFiles) >= exploreSourceLiteralRecallMaxFilesPerTerm { @@ -216,10 +281,12 @@ func retainExploreSourceLiteralOwners(recall exploreSourceLiteralRecall) (hits [ } // gatherExploreSourceLiteralRecall reuses the bounded raw-text path behind -// search(operation:"text") only when content_fts could not produce an exact -// symbol candidates. It searches one repository and at most two compact -// literals, maps 1-based hits to their smallest enclosing declarations, and -// returns a file-diverse owner set for the caller's existing batch hydration. +// search(operation:"text") for every admitted task literal whose value is not +// already visible in a ranked declaration. Compact codes keep their reserved +// slots, then the longest remaining literals fill the bounded term budget. It +// searches one repository, maps 1-based hits to their smallest enclosing +// declarations, and returns a file-diverse owner set for the caller's +// existing batch hydration. func (s *Server) gatherExploreSourceLiteralRecall( ctx context.Context, terms []string, @@ -230,18 +297,18 @@ func (s *Server) gatherExploreSourceLiteralRecall( return exploreSourceLiteralRecall{} } attemptTerms := exploreCompactSourceLiteralTerms(terms) - if len(attemptTerms) == 0 { - if primary := explorePreferredSourceLiteral(terms); primary != "" { - attemptTerms = append(attemptTerms, primary) - } - } + attemptTerms = appendExploreSourceLiteralProseTerms(attemptTerms, terms) if len(attemptTerms) == 0 { return exploreSourceLiteralRecall{} } recallBudget := s.sourceLiteralRecallBudget() started := time.Now() - boundedCtx, cancelBounded := context.WithTimeout(ctx, 2*recallBudget) + // The end-to-end wall grows with the admitted term count so a page that + // quotes several distinct literals can ground each of them, while a page + // with one literal keeps the original two-slice bound. + wallSlices := time.Duration(max(2, len(attemptTerms))) + boundedCtx, cancelBounded := context.WithTimeout(ctx, wallSlices*recallBudget) defer cancelBounded() type literalAttempt struct { term string @@ -266,7 +333,10 @@ func (s *Server) gatherExploreSourceLiteralRecall( // fixed 150ms request budget without silently halving grep time again. searchCtx, cancelSearch := context.WithTimeout(attemptCtx, searchBudget) result.search = s.searchExploreSourceLiteral(searchCtx, term, repoPrefix, scope, exploreSourceLiteralRecallMaxHits) - result.searchErr = searchCtx.Err() + result.searchErr = result.search.err + if result.searchErr == nil { + result.searchErr = searchCtx.Err() + } cancelSearch() if ctx.Err() != nil { return result @@ -309,7 +379,6 @@ func (s *Server) gatherExploreSourceLiteralRecall( }) for _, hit := range retained { hit.anchor = anchor - hit.ambiguous = result.recall.ambiguous aggregate.hits = append(aggregate.hits, hit) if file := result.recall.ownerFiles[hit.nodeID]; file != "" { aggregate.ownerFiles[hit.nodeID] = file @@ -344,12 +413,9 @@ func (s *Server) gatherExploreSourceLiteralRecall( if _, ok := attempted[strings.ToLower(term)]; ok { continue } - reason := "not_compact" - if exploreCompactSourceLiteral(term, utf8.RuneCountInString(term)) { - reason = "term_cap" - if boundedCtx.Err() != nil { - reason = "request_deadline" - } + reason := "term_cap" + if boundedCtx.Err() != nil { + reason = "request_deadline" } aggregate.diagnostics = append(aggregate.diagnostics, exploreSourceLiteralDiagnostic{ literal: term, @@ -458,6 +524,10 @@ func (s *Server) mapExploreSourceLiteralMatchesContext( if s == nil || len(matches) == 0 || ctx.Err() != nil { return exploreSourceLiteralRecall{} } + reader := s.readerFor(ctx) + if reader == nil { + return exploreSourceLiteralRecall{} + } saturated := len(matches) >= exploreSourceLiteralRecallMaxHits if len(matches) > exploreSourceLiteralRecallMaxHits { matches = matches[:exploreSourceLiteralRecallMaxHits] @@ -490,8 +560,6 @@ func (s *Server) mapExploreSourceLiteralMatchesContext( } } } - sort.Strings(exactPaths) - sort.Strings(aliasPaths) orderedPaths := make([]string, 0, len(exactPaths)+len(aliasPaths)) orderedPaths = append(orderedPaths, exactPaths...) for _, alias := range aliasPaths { @@ -499,7 +567,7 @@ func (s *Server) mapExploreSourceLiteralMatchesContext( orderedPaths = append(orderedPaths, alias) } } - indexes := s.buildFileSymbolIndexForOrderedPathsContext(ctx, orderedPaths) + indexes := s.buildFileSymbolIndexForOrderedPathsScopedContext(ctx, orderedPaths, scope) type mappedLiteralOwner struct { owner *graph.Node match trigram.Match @@ -507,11 +575,8 @@ func (s *Server) mapExploreSourceLiteralMatchesContext( callName string } mapped := make([]mappedLiteralOwner, 0, len(matches)) - ownerIDs := make([]string, 0, len(matches)) - seenOwners := make(map[string]struct{}, len(matches)) - calleeIDs := make([]string, 0, len(matches)) - seenCallees := make(map[string]struct{}, len(matches)) - ownerEdges := make(map[string][]*graph.Edge) + ownerSites := make([]graph.EdgeSourceSite, 0, len(matches)) + seenSites := make(map[graph.EdgeSourceSite]struct{}, len(matches)) // First map each exact line to its smallest enclosing declaration. When // the literal is syntactically inside one call on that line, retain the @@ -536,36 +601,61 @@ func (s *Server) mapExploreSourceLiteralMatchesContext( callName, _ := exploreSourceLiteralCallName(match.Text, term) mapped = append(mapped, mappedLiteralOwner{owner: owner, match: match, rank: rank, callName: callName}) if callName != "" { - if _, duplicate := seenOwners[owner.ID]; duplicate { + site := graph.EdgeSourceSite{From: owner.ID, Line: match.Line} + if _, duplicate := seenSites[site]; duplicate { continue } - seenOwners[owner.ID] = struct{}{} - ownerIDs = append(ownerIDs, owner.ID) + seenSites[site] = struct{}{} + ownerSites = append(ownerSites, site) } } - if len(ownerIDs) > 0 { - ownerEdges = s.graph.GetOutEdgesByNodeIDs(ownerIDs) - for _, item := range mapped { - if item.callName == "" { - continue - } - for _, edge := range ownerEdges[item.owner.ID] { - if edge == nil || edge.Kind != graph.EdgeCalls || edge.Line != item.match.Line || edge.To == "" { - continue + siteEdges := make(map[graph.EdgeSourceSite][]graph.EdgeIdentity) + calleeNodes := map[string]*graph.Node{} + if len(ownerSites) > 0 { + bounded, supported := reader.(graph.BoundedOutgoingSiteEdgeIdentityReader) + if supported { + projection, err := bounded.FindOutgoingSiteEdgeIdentitiesBounded( + ctx, ownerSites, []graph.EdgeKind{graph.EdgeCalls}, exploreSourceLiteralCallEdgesPerSite, + ) + complete := err == nil && ctx.Err() == nil + calleeIDs := make([]string, 0, exploreSourceLiteralCalleeHydrationLimit) + seenCallees := make(map[string]struct{}, exploreSourceLiteralCalleeHydrationLimit) + if complete { + for _, site := range ownerSites { + if projection.Truncated[site] { + complete = false + break + } + for _, identity := range projection.BySite[site] { + if identity.To == "" { + continue + } + if _, duplicate := seenCallees[identity.To]; duplicate { + continue + } + if len(calleeIDs) == exploreSourceLiteralCalleeHydrationLimit { + complete = false + break + } + seenCallees[identity.To] = struct{}{} + calleeIDs = append(calleeIDs, identity.To) + } + if !complete { + break + } } - if _, duplicate := seenCallees[edge.To]; duplicate { - continue + } + if complete { + if hydrated, hydratedComplete := exploreNodesByIDsBounded( + ctx, reader, calleeIDs, exploreSourceLiteralCalleeHydrationLimit, + ); hydratedComplete { + siteEdges = projection.BySite + calleeNodes = hydrated } - seenCallees[edge.To] = struct{}{} - calleeIDs = append(calleeIDs, edge.To) } } } - calleeNodes := map[string]*graph.Node{} - if len(calleeIDs) > 0 { - calleeNodes = s.graph.GetNodesByIDs(calleeIDs) - } seen := make(map[string]int, len(mapped)) hits := make([]exploreSourceLiteralHit, 0, len(matches)) @@ -575,11 +665,9 @@ func (s *Server) mapExploreSourceLiteralMatchesContext( calleeResolved := false if item.callName != "" { resolved := make(map[string]*graph.Node) - for _, edge := range ownerEdges[item.owner.ID] { - if edge == nil || edge.Kind != graph.EdgeCalls || edge.Line != item.match.Line { - continue - } - callee := calleeNodes[edge.To] + site := graph.EdgeSourceSite{From: item.owner.ID, Line: item.match.Line} + for _, identity := range siteEdges[site] { + callee := calleeNodes[identity.To] if !exploreSourceLiteralLocalCallee(item.owner, callee, item.callName, scope) { continue } @@ -599,16 +687,73 @@ func (s *Server) mapExploreSourceLiteralMatchesContext( continue } seen[node.ID] = len(hits) - hits = append(hits, exploreSourceLiteralHit{nodeID: node.ID, rank: item.rank, callee: calleeResolved}) + hits = append(hits, exploreSourceLiteralHit{ + nodeID: node.ID, rank: item.rank, callee: calleeResolved, + matchPath: item.match.Path, matchLine: item.match.Line, literal: term, + }) ownerFiles[node.ID] = node.FilePath } + // Ambiguity is a property of the term's evidence, not of having company: + // a complete search that maps to a handful of distinct owners has settled + // each of them — a literal registered in one file and used in another is + // two pieces of evidence, not mutual noise. Only a truncated search, or a + // literal so common it overflows the owner cap, leaves owners unproven. + ambiguous := saturated || len(hits) > exploreSourceLiteralRecallMaxOwnersPerTerm + for index := range hits { + hits[index].ambiguous = ambiguous + } return exploreSourceLiteralRecall{ hits: hits, - ambiguous: saturated || len(hits) > 1, + ambiguous: ambiguous, ownerFiles: ownerFiles, } } +// projectExploreSourceLiteralConstructionAlignment proves only complete, +// current-view instantiation adjacency for the already-bounded literal callee +// cohort. The full per-key projection cap is intentional: a truncated page may +// consist entirely of durable edges hidden by an overlay, so truncation can +// never be treated as a positive witness. +func projectExploreSourceLiteralConstructionAlignment( + ctx context.Context, + reader graph.Reader, + ids []string, +) map[string]bool { + if ctx == nil { + ctx = context.Background() + } + if reader == nil || ctx.Err() != nil { + return nil + } + boundedIDs, complete := exploreBoundedNodeIDs(ids, exploreSourceLiteralRecallMaxHits) + if !complete || len(boundedIDs) == 0 { + return nil + } + bounded, supported := reader.(graph.BoundedOutgoingEdgeIdentityReader) + if !supported { + return nil + } + projection, err := bounded.FindOutgoingEdgeIdentitiesBounded( + ctx, + boundedIDs, + []graph.EdgeKind{graph.EdgeInstantiates}, + graph.MaxBoundedAdjacencyRowsPerKey, + ) + if err != nil || ctx.Err() != nil { + return nil + } + aligned := make(map[string]bool, len(boundedIDs)) + for _, id := range boundedIDs { + if projection.Truncated[id] { + continue + } + if len(projection.ByEndpoint[id]) > 0 { + aligned[id] = true + } + } + return aligned +} + type exploreSourceLiteralSpan struct { start int end int @@ -745,6 +890,9 @@ func exploreSourceLiteralLocalCallee(owner, callee *graph.Node, callName string, if callee.Kind != graph.KindFunction && callee.Kind != graph.KindMethod { return false } + if scope.ExcludeTests && exploreDraftIsTestNode(callee) { + return false + } if callee.RepoPrefix != owner.RepoPrefix { return false } @@ -781,70 +929,204 @@ func exploreSourceLiteralUnprefixedPath(path, repoPrefix string) string { return strings.TrimPrefix(path, marker) } -// searchExploreSourceLiteral mirrors search_text's literal backend while -// deliberately refusing an unscoped multi-repository fan-out. The caller's -// session locality supplies repoPrefix in normal operation. maxHits is the -// caller's own recall cap — every caller is responsible for a bound the rest -// of its response budget can absorb. -func (s *Server) searchExploreSourceLiteral( +func (s *Server) filterExploreSourceLiteralMatches(ctx context.Context, matches []trigram.Match) []trigram.Match { + if OverlayViewFromContext(ctx) == nil || len(matches) == 0 { + return matches + } + filtered := matches[:0:0] + for _, match := range matches { + if !s.overlayShadowsGraphPath(ctx, match.Path, "") { + filtered = append(filtered, match) + } + } + return filtered +} + +func finalizeExploreSourceLiteralSearch( + result exploreSourceLiteralSearch, + overlayScan exploreSourceLiteralOverlayScan, + covered map[string]struct{}, + maxHits int, + overlayPresent bool, + overlayIncomplete bool, + err error, +) exploreSourceLiteralSearch { + if maxHits <= 0 || maxHits > exploreSourceLiteralOverlayMaxHits { + maxHits = exploreSourceLiteralOverlayMaxHits + } + if len(covered) == 0 && len(overlayScan.matches) == 0 { + if len(result.matches) > maxHits { + result.matches = result.matches[:maxHits] + result.incomplete = true + } + result.incomplete = result.incomplete || overlayIncomplete || overlayScan.incomplete || err != nil + result.err = err + if overlayPresent { + result.backend += "+overlay" + } + return result + } + merged, mergeIncomplete := mergeExploreSourceLiteralMatches( + overlayScan.matches, result.matches, covered, maxHits, + ) + result.matches = merged + result.incomplete = result.incomplete || overlayIncomplete || overlayScan.incomplete || mergeIncomplete || err != nil + result.err = err + if overlayPresent { + result.backend += "+overlay" + result.owned = result.owned || len(overlayScan.matches) > 0 + } + return result +} + +type exploreSourceLiteralOverlayScanner func( + context.Context, string, []exploreSourceLiteralOverlayFile, int, +) (exploreSourceLiteralOverlayScan, error) + +type exploreSourceLiteralDurableSearcher func( + context.Context, string, string, int, +) exploreSourceLiteralSearch + +func runExploreSourceLiteralLanes( ctx context.Context, term string, repoPrefix string, - scope query.QueryOptions, maxHits int, + overlayFiles []exploreSourceLiteralOverlayFile, + covered map[string]struct{}, + overlayIncomplete bool, + coveredOverflow bool, + scanOverlay exploreSourceLiteralOverlayScanner, + searchDurable exploreSourceLiteralDurableSearcher, ) exploreSourceLiteralSearch { - if s.multiIndexer != nil { - if repoPrefix == "" { - haveScopedPrefix := false - for prefix, allowed := range scope.RepoAllow { - if !allowed { - continue - } - prefix = strings.TrimSuffix(strings.TrimSpace(prefix), "/") - if haveScopedPrefix && repoPrefix != prefix { - return exploreSourceLiteralSearch{backend: "multi-ambiguous-scope"} - } - repoPrefix = prefix - haveScopedPrefix = true + if maxHits <= 0 || maxHits > exploreSourceLiteralOverlayMaxHits { + maxHits = exploreSourceLiteralOverlayMaxHits + } + if coveredOverflow { + return exploreSourceLiteralSearch{ + backend: "overlay-covered-cap", lookupRepoPrefix: repoPrefix, incomplete: true, owned: true, + } + } + if len(overlayFiles) == 0 && len(covered) == 0 { + result := searchDurable(ctx, term, repoPrefix, maxHits+1) + return finalizeExploreSourceLiteralSearch( + result, exploreSourceLiteralOverlayScan{}, covered, maxHits, false, + overlayIncomplete, ctx.Err(), + ) + } + + overlayScan, scanErr := scanOverlay(ctx, term, overlayFiles, maxHits) + if scanErr != nil { + return finalizeExploreSourceLiteralSearch( + exploreSourceLiteralSearch{backend: "none", lookupRepoPrefix: repoPrefix}, + overlayScan, covered, maxHits, true, overlayIncomplete, scanErr, + ) + } + if err := ctx.Err(); err != nil { + return finalizeExploreSourceLiteralSearch( + exploreSourceLiteralSearch{backend: "none", lookupRepoPrefix: repoPrefix}, + overlayScan, covered, maxHits, true, overlayIncomplete, err, + ) + } + if len(overlayScan.matches) >= maxHits { + allProduction := true + for _, match := range overlayScan.matches[:maxHits] { + if testpath.IsTestFile(match.Path) { + allProduction = false + break } } - result := s.multiIndexer.GrepLiteralForRepoBounded( - ctx, repoPrefix, term, - maxHits, - exploreSourceLiteralRecallMaxFiles, + if allProduction { + overlayScan.incomplete = true // Durable existence is deliberately not probed. + return finalizeExploreSourceLiteralSearch( + exploreSourceLiteralSearch{backend: "none", lookupRepoPrefix: repoPrefix}, + overlayScan, covered, maxHits, true, overlayIncomplete, nil, + ) + } + } + + result := searchDurable(ctx, term, repoPrefix, maxHits+len(covered)+1) + return finalizeExploreSourceLiteralSearch( + result, overlayScan, covered, maxHits, true, overlayIncomplete, ctx.Err(), + ) +} + +func (s *Server) searchExploreSourceLiteralDurable( + ctx context.Context, + term string, + repoPrefix string, + durableLimit int, +) exploreSourceLiteralSearch { + result := exploreSourceLiteralSearch{backend: "none", lookupRepoPrefix: repoPrefix} + if s.multiIndexer != nil { + durable := s.multiIndexer.GrepLiteralForRepoBounded( + ctx, repoPrefix, term, durableLimit, exploreSourceLiteralRecallMaxFiles, ) - if result.Owned { + if durable.Owned { return exploreSourceLiteralSearch{ - matches: result.Matches, - incomplete: result.Incomplete, + matches: s.filterExploreSourceLiteralMatches(ctx, durable.Matches), + incomplete: durable.Incomplete || len(durable.Matches) >= durableLimit, backend: "multi", owned: true, - lookupRepoPrefix: result.RepoPrefix, + lookupRepoPrefix: durable.RepoPrefix, } } - // Once MultiIndexer owns any repository, an unresolved prefix is an - // ownership failure rather than permission to scan the base indexer. - // Falling through here can leak matches from a different repository. - if result.Configured { - return exploreSourceLiteralSearch{ - backend: "multi-unresolved", - lookupRepoPrefix: repoPrefix, - } + if durable.Configured { + return exploreSourceLiteralSearch{backend: "multi-unresolved", lookupRepoPrefix: repoPrefix} } } if s.indexer != nil { matches, incomplete := s.indexer.GrepLiteralBounded( - ctx, term, - maxHits, - exploreSourceLiteralRecallMaxFiles, + ctx, term, durableLimit, exploreSourceLiteralRecallMaxFiles, ) return exploreSourceLiteralSearch{ - matches: matches, - incomplete: incomplete, + matches: s.filterExploreSourceLiteralMatches(ctx, matches), + incomplete: incomplete || len(matches) >= durableLimit, backend: "direct", owned: true, lookupRepoPrefix: s.indexer.RepoPrefix(), } } - return exploreSourceLiteralSearch{backend: "none", lookupRepoPrefix: repoPrefix} + return result +} + +// searchExploreSourceLiteral mirrors search_text's literal backend while +// deliberately refusing an unscoped multi-repository fan-out. The caller's +// session locality supplies repoPrefix in normal operation. maxHits is the +// caller's own recall cap — every caller is responsible for a bound the rest +// of its response budget can absorb. +func (s *Server) searchExploreSourceLiteral( + ctx context.Context, + term string, + repoPrefix string, + scope query.QueryOptions, + maxHits int, +) exploreSourceLiteralSearch { + if err := ctx.Err(); err != nil { + return exploreSourceLiteralSearch{backend: "none", lookupRepoPrefix: repoPrefix, err: err} + } + if maxHits <= 0 || maxHits > exploreSourceLiteralOverlayMaxHits { + maxHits = exploreSourceLiteralOverlayMaxHits + } + var scopeOK bool + scope, scopeOK = s.effectiveExploreSourceLiteralScope(ctx, scope) + if !scopeOK { + return exploreSourceLiteralSearch{backend: "session-scope-mismatch", lookupRepoPrefix: repoPrefix} + } + resolvedRepoPrefix, prefixOK := s.resolveExploreSourceLiteralRepoPrefix(repoPrefix, scope) + if !prefixOK { + return exploreSourceLiteralSearch{backend: "multi-ambiguous-scope", lookupRepoPrefix: repoPrefix} + } + repoPrefix = resolvedRepoPrefix + overlayFiles, covered, overlayIncomplete, coveredOverflow, overlayErr := s.snapshotExploreSourceLiteralOverlays(ctx, scope, repoPrefix) + if overlayErr != nil { + return exploreSourceLiteralSearch{ + backend: "overlay", lookupRepoPrefix: repoPrefix, incomplete: true, owned: true, err: overlayErr, + } + } + return runExploreSourceLiteralLanes( + ctx, term, repoPrefix, maxHits, overlayFiles, covered, + overlayIncomplete, coveredOverflow, + scanExploreSourceLiteralOverlays, s.searchExploreSourceLiteralDurable, + ) } diff --git a/internal/mcp/explore_source_literal_adjacency_test.go b/internal/mcp/explore_source_literal_adjacency_test.go new file mode 100644 index 00000000..53d00575 --- /dev/null +++ b/internal/mcp/explore_source_literal_adjacency_test.go @@ -0,0 +1,524 @@ +package mcp + +import ( + "context" + "errors" + "fmt" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/zzet/gortex/internal/graph" + "github.com/zzet/gortex/internal/query" + "github.com/zzet/gortex/internal/search/trigram" +) + +type exploreSourceLiteralUnsupportedSiteStore struct { + graph.Store +} + +func (s *exploreSourceLiteralUnsupportedSiteStore) FindFileNodesBounded( + ctx context.Context, + path string, + scope graph.LocalizationNodeScope, + limit int, +) (graph.BoundedNodeProjection, error) { + return s.Store.(graph.BoundedFileNodeReader).FindFileNodesBounded(ctx, path, scope, limit) +} + +func (*exploreSourceLiteralUnsupportedSiteStore) GetOutEdgesByNodeIDs([]string) map[string][]*graph.Edge { + panic("legacy outgoing adjacency must not be used") +} + +func sourceLiteralCallFixture( + path string, + owner *graph.Node, + line int, + count int, +) ([]*graph.Node, []*graph.Edge, *graph.Node) { + nodes := make([]*graph.Node, 0, count+1) + nodes = append(nodes, owner) + edges := make([]*graph.Edge, 0, count) + var wanted *graph.Node + for index := 0; index < count; index++ { + name := fmt.Sprintf("Other%d", index) + if index == 0 { + name = "Resolve" + } + callee := sourceLiteralNode( + fmt.Sprintf("%s::callee-%d", path, index), name, path, + graph.KindMethod, 100+index*2, 101+index*2, + ) + if index == 0 { + wanted = callee + } + nodes = append(nodes, callee) + edges = append(edges, &graph.Edge{ + From: owner.ID, To: callee.ID, Kind: graph.EdgeCalls, + FilePath: path, Line: line, + }) + } + return nodes, edges, wanted +} + +func TestMapExploreSourceLiteralMatchesCallsiteExactEightAndNinthSentinel(t *testing.T) { + for _, test := range []struct { + name string + callEdges int + promoted bool + }{ + {name: "exact eight is complete", callEdges: exploreSourceLiteralCallEdgesPerSite, promoted: true}, + {name: "ninth fails closed", callEdges: exploreSourceLiteralCallEdgesPerSite + 1}, + } { + t.Run(test.name, func(t *testing.T) { + const path = "demo/src/calls.cs" + owner := sourceLiteralNode(path+"::owner", "owner", path, graph.KindMethod, 1, 20) + nodes, edges, wanted := sourceLiteralCallFixture(path, owner, 3, test.callEdges) + for index := 0; index < 64; index++ { + edges = append(edges, &graph.Edge{ + From: owner.ID, To: fmt.Sprintf("irrelevant-%d", index), + Kind: graph.EdgeReferences, FilePath: path, Line: 3, + }) + } + server, counting := newExploreSourceLiteralGraphServer(t, nodes, edges) + recall := server.mapExploreSourceLiteralMatches("needle", []trigram.Match{{ + Path: path, Line: 3, Text: `Resolve("needle");`, + }}, query.QueryOptions{RepoAllow: map[string]bool{"demo": true}}) + + wantID := owner.ID + if test.promoted { + wantID = wanted.ID + } + requireSourceLiteralHitIdentity(t, recall.hits, exploreSourceLiteralHit{ + nodeID: wantID, rank: 0, callee: test.promoted, + }) + require.Equal(t, 1, counting.outSiteBatchCalls) + require.Zero(t, counting.outEdgeBatchCalls) + if test.promoted { + require.Len(t, counting.lastNodeIDs, exploreSourceLiteralCallEdgesPerSite) + } else { + require.Zero(t, counting.nodeLookupBatches) + } + }) + } +} + +func TestMapExploreSourceLiteralMatchesDropsWholeCalleeLaneOnOneTruncatedSite(t *testing.T) { + const path = "demo/src/atomic.cs" + firstOwner := sourceLiteralNode(path+"::first", "first", path, graph.KindMethod, 1, 10) + secondOwner := sourceLiteralNode(path+"::second", "second", path, graph.KindMethod, 20, 30) + firstNodes, firstEdges, _ := sourceLiteralCallFixture(path, firstOwner, 3, 1) + secondNodes, secondEdges, _ := sourceLiteralCallFixture(path, secondOwner, 23, exploreSourceLiteralCallEdgesPerSite+1) + nodes := append(firstNodes, secondNodes...) + edges := append(firstEdges, secondEdges...) + server, counting := newExploreSourceLiteralGraphServer(t, nodes, edges) + + recall := server.mapExploreSourceLiteralMatches("needle", []trigram.Match{ + {Path: path, Line: 3, Text: `Resolve("needle");`}, + {Path: path, Line: 23, Text: `Resolve("needle");`}, + }, query.QueryOptions{RepoAllow: map[string]bool{"demo": true}}) + + requireSourceLiteralHitIdentity(t, recall.hits, + exploreSourceLiteralHit{nodeID: firstOwner.ID, rank: 0}, + exploreSourceLiteralHit{nodeID: secondOwner.ID, rank: 1}, + ) + require.Zero(t, counting.nodeLookupBatches, "a truncated peer must discard every partial callee proof") +} + +func TestMapExploreSourceLiteralMatchesDropsWholeCalleeLaneOnMissingHydration(t *testing.T) { + const path = "demo/src/hydration.cs" + firstOwner := sourceLiteralNode(path+"::first", "first", path, graph.KindMethod, 1, 10) + secondOwner := sourceLiteralNode(path+"::second", "second", path, graph.KindMethod, 20, 30) + firstNodes, firstEdges, _ := sourceLiteralCallFixture(path, firstOwner, 3, 1) + secondNodes, secondEdges, _ := sourceLiteralCallFixture(path, secondOwner, 23, 1) + server, counting := newExploreSourceLiteralGraphServer(t, append(firstNodes, secondNodes...), append(firstEdges, secondEdges...)) + counting.nodeLookup = func(ctx context.Context, ids []string) (map[string]*graph.Node, error) { + nodes, err := counting.Store.(exploreContextNodesReader).GetNodesByIDsContext(ctx, ids) + delete(nodes, ids[len(ids)-1]) + return nodes, err + } + + recall := server.mapExploreSourceLiteralMatches("needle", []trigram.Match{ + {Path: path, Line: 3, Text: `Resolve("needle");`}, + {Path: path, Line: 23, Text: `Resolve("needle");`}, + }, query.QueryOptions{RepoAllow: map[string]bool{"demo": true}}) + + requireSourceLiteralHitIdentity(t, recall.hits, + exploreSourceLiteralHit{nodeID: firstOwner.ID, rank: 0}, + exploreSourceLiteralHit{nodeID: secondOwner.ID, rank: 1}, + ) + require.Equal(t, 1, counting.nodeLookupBatches) +} + +func TestMapExploreSourceLiteralMatchesCapsCalleeHydrationAt192(t *testing.T) { + const path = "demo/src/many-sites.cs" + owner := sourceLiteralNode(path+"::owner", "owner", path, graph.KindMethod, 1, 40) + nodes := []*graph.Node{owner} + edges := make([]*graph.Edge, 0, exploreSourceLiteralCalleeHydrationLimit) + matches := make([]trigram.Match, 0, exploreSourceLiteralRecallMaxHits) + for site := 0; site < exploreSourceLiteralRecallMaxHits; site++ { + line := site + 2 + matches = append(matches, trigram.Match{Path: path, Line: line, Text: `Resolve("needle");`}) + for index := 0; index < exploreSourceLiteralCallEdgesPerSite; index++ { + callee := sourceLiteralNode( + fmt.Sprintf("%s::callee-%d-%d", path, site, index), "Resolve", path, + graph.KindMethod, 100+site*20+index, 100+site*20+index, + ) + nodes = append(nodes, callee) + edges = append(edges, &graph.Edge{ + From: owner.ID, To: callee.ID, Kind: graph.EdgeCalls, + FilePath: path, Line: line, + }) + } + } + server, counting := newExploreSourceLiteralGraphServer(t, nodes, edges) + recall := server.mapExploreSourceLiteralMatches( + "needle", matches, query.QueryOptions{RepoAllow: map[string]bool{"demo": true}}, + ) + + // A recall that filled its hit cap is saturated: the single mapped owner + // stays ambiguous because the search cannot prove it is alone. + requireSourceLiteralHitIdentity(t, recall.hits, exploreSourceLiteralHit{nodeID: owner.ID, rank: 0, ambiguous: true}) + require.Equal(t, 1, counting.outSiteBatchCalls) + require.Len(t, counting.lastSites, exploreSourceLiteralRecallMaxHits) + require.Equal(t, 1, counting.nodeLookupBatches) + require.Len(t, counting.lastNodeIDs, exploreSourceLiteralCalleeHydrationLimit) +} + +func TestMapExploreSourceLiteralMatchesFailsClosedWithoutSiteCapability(t *testing.T) { + const path = "demo/src/unsupported.cs" + owner := sourceLiteralNode(path+"::owner", "owner", path, graph.KindMethod, 1, 10) + callee := sourceLiteralNode(path+"::Resolve", "Resolve", path, graph.KindMethod, 20, 30) + server := newExploreSourceLiteralServer(t, []*graph.Node{owner, callee}) + server.graph = &exploreSourceLiteralUnsupportedSiteStore{Store: server.graph} + + recall := server.mapExploreSourceLiteralMatches("needle", []trigram.Match{{ + Path: path, Line: 3, Text: `Resolve("needle");`, + }}, query.QueryOptions{RepoAllow: map[string]bool{"demo": true}}) + requireSourceLiteralHitIdentity(t, recall.hits, exploreSourceLiteralHit{nodeID: owner.ID, rank: 0}) +} + +func TestMapDiscoveredExploreSourceLiteralMatchesPropagatesStickyAdjacencyCancellation(t *testing.T) { + const path = "demo/src/cancel.cs" + owner := sourceLiteralNode(path+"::owner", "owner", path, graph.KindMethod, 1, 10) + callee := sourceLiteralNode(path+"::Resolve", "Resolve", path, graph.KindMethod, 20, 30) + server, counting := newExploreSourceLiteralGraphServer(t, []*graph.Node{owner, callee}, nil) + ctx, cancel := context.WithCancel(context.Background()) + counting.siteLookup = func(_ context.Context, sites []graph.EdgeSourceSite, _ []graph.EdgeKind, _ int) (graph.BoundedSiteEdgeIdentityProjection, error) { + cancel() + return graph.BoundedSiteEdgeIdentityProjection{BySite: map[graph.EdgeSourceSite][]graph.EdgeIdentity{ + sites[0]: {{From: owner.ID, To: callee.ID, Kind: graph.EdgeCalls, FilePath: path, Line: 3}}, + }}, nil + } + + recall, err := server.mapDiscoveredExploreSourceLiteralMatches(ctx, "needle", exploreSourceLiteralSearch{ + matches: []trigram.Match{{Path: path, Line: 3, Text: `Resolve("needle");`}}, + }, query.QueryOptions{RepoAllow: map[string]bool{"demo": true}}, nil) + require.ErrorIs(t, err, context.Canceled) + requireSourceLiteralHitIdentity(t, recall.hits, exploreSourceLiteralHit{nodeID: owner.ID, rank: 0}) + require.Zero(t, counting.nodeLookupBatches) +} + +func TestProjectExploreSourceLiteralConstructionAlignmentBoundsAndFailures(t *testing.T) { + const ( + path = "demo/src/construction.cs" + sourceID = path + "::source" + targetID = path + "::target" + ) + source := sourceLiteralNode(sourceID, "source", path, graph.KindMethod, 1, 10) + target := sourceLiteralNode(targetID, "target", path, graph.KindType, 20, 30) + edge := graph.EdgeIdentity{From: sourceID, To: targetID, Kind: graph.EdgeInstantiates, FilePath: path, Line: 3} + + t.Run("complete positive", func(t *testing.T) { + server, counting := newExploreSourceLiteralGraphServer(t, []*graph.Node{source, target}, []*graph.Edge{{ + From: sourceID, To: targetID, Kind: graph.EdgeInstantiates, FilePath: path, Line: 3, + }}) + aligned := projectExploreSourceLiteralConstructionAlignment(context.Background(), server.graph, []string{sourceID}) + require.Equal(t, map[string]bool{sourceID: true}, aligned) + require.Equal(t, 1, counting.outEndpointBatchCalls) + require.Equal(t, []string{sourceID}, counting.lastEndpointIDs) + require.Equal(t, []graph.EdgeKind{graph.EdgeInstantiates}, counting.lastEndpointKinds) + require.Equal(t, graph.MaxBoundedAdjacencyRowsPerKey, counting.lastEndpointLimit) + require.Zero(t, counting.outEdgeBatchCalls) + }) + + for _, test := range []struct { + name string + count int + aligned bool + }{ + {name: "graph exact 256 is complete", count: graph.MaxBoundedAdjacencyRowsPerKey, aligned: true}, + {name: "graph 257th sentinel fails closed", count: graph.MaxBoundedAdjacencyRowsPerKey + 1}, + } { + t.Run(test.name, func(t *testing.T) { + base := graph.New() + edges := make([]*graph.Edge, 0, test.count) + for line := 1; line <= test.count; line++ { + edges = append(edges, &graph.Edge{ + From: sourceID, To: targetID, Kind: graph.EdgeInstantiates, + FilePath: path, Line: line, + }) + } + base.AddBatch([]*graph.Node{source, target}, edges) + aligned := projectExploreSourceLiteralConstructionAlignment( + context.Background(), base, []string{sourceID}, + ) + require.Equal(t, test.aligned, aligned[sourceID]) + }) + } + + t.Run("zero", func(t *testing.T) { + server, counting := newExploreSourceLiteralGraphServer(t, []*graph.Node{source, target}, nil) + require.Empty(t, projectExploreSourceLiteralConstructionAlignment(context.Background(), server.graph, []string{sourceID})) + require.Equal(t, 1, counting.outEndpointBatchCalls) + }) + + for _, test := range []struct { + name string + projection graph.BoundedEdgeIdentityProjection + err error + }{ + {name: "truncated is not positive", projection: graph.BoundedEdgeIdentityProjection{ + ByEndpoint: map[string][]graph.EdgeIdentity{sourceID: {edge}}, + Truncated: map[string]bool{sourceID: true}, + }}, + {name: "partial backend error", projection: graph.BoundedEdgeIdentityProjection{ + ByEndpoint: map[string][]graph.EdgeIdentity{sourceID: {edge}}, + }, err: errors.New("partial backend failure")}, + } { + t.Run(test.name, func(t *testing.T) { + server, counting := newExploreSourceLiteralGraphServer(t, []*graph.Node{source, target}, nil) + counting.endpointLookup = func(context.Context, []string, []graph.EdgeKind, int) (graph.BoundedEdgeIdentityProjection, error) { + return test.projection, test.err + } + require.Empty(t, projectExploreSourceLiteralConstructionAlignment(context.Background(), server.graph, []string{sourceID})) + }) + } + + t.Run("sticky cancellation", func(t *testing.T) { + server, counting := newExploreSourceLiteralGraphServer(t, []*graph.Node{source, target}, nil) + ctx, cancel := context.WithCancel(context.Background()) + counting.endpointLookup = func(context.Context, []string, []graph.EdgeKind, int) (graph.BoundedEdgeIdentityProjection, error) { + cancel() + return graph.BoundedEdgeIdentityProjection{ByEndpoint: map[string][]graph.EdgeIdentity{sourceID: {edge}}}, nil + } + require.Empty(t, projectExploreSourceLiteralConstructionAlignment(ctx, server.graph, []string{sourceID})) + }) + + t.Run("unsupported has no legacy fallback", func(t *testing.T) { + base := graph.New() + require.Empty(t, projectExploreSourceLiteralConstructionAlignment( + context.Background(), &exploreSourceLiteralUnsupportedSiteStore{Store: base}, []string{sourceID}, + )) + }) +} + +func TestProjectExploreSourceLiteralConstructionAlignmentOverlayParity(t *testing.T) { + const ( + sourceFile = "demo/source.cs" + targetFile = "demo/target.cs" + sourceID = sourceFile + "::source" + targetID = targetFile + "::target" + ) + base := graph.New() + base.AddNode(&graph.Node{ID: sourceID, Name: "source", Kind: graph.KindMethod, FilePath: sourceFile}) + base.AddNode(&graph.Node{ID: targetID, Name: "target", Kind: graph.KindType, FilePath: targetFile}) + base.AddEdge(&graph.Edge{From: sourceID, To: targetID, Kind: graph.EdgeInstantiates, FilePath: sourceFile, Line: 4}) + + require.True(t, projectExploreSourceLiteralConstructionAlignment( + context.Background(), base, []string{sourceID}, + )[sourceID]) + + t.Run("source tombstone and replacement", func(t *testing.T) { + tombstone := graph.NewOverlayLayer() + tombstone.MarkFile(sourceFile, false) + tombstone.MarkRemoved("source", sourceID) + tombstone.AddEdge(&graph.Edge{From: sourceID, To: targetID, Kind: graph.EdgeInstantiates, Line: 99}) + require.Empty(t, projectExploreSourceLiteralConstructionAlignment( + context.Background(), graph.NewOverlaidView(base, tombstone), []string{sourceID}, + )) + + replacement := graph.NewOverlayLayer() + replacement.MarkFile(sourceFile, false) + replacement.MarkRemoved("source", sourceID) + replacement.AddNode(sourceFile, &graph.Node{ID: sourceID, Name: "source", Kind: graph.KindMethod, FilePath: sourceFile}) + replacement.AddEdge(&graph.Edge{From: sourceID, To: targetID, Kind: graph.EdgeInstantiates, FilePath: sourceFile, Line: 8}) + require.True(t, projectExploreSourceLiteralConstructionAlignment( + context.Background(), graph.NewOverlaidView(base, replacement), []string{sourceID}, + )[sourceID]) + }) + + t.Run("target tombstone and replacement", func(t *testing.T) { + tombstone := graph.NewOverlayLayer() + tombstone.MarkFile(targetFile, false) + tombstone.MarkRemoved("target", targetID) + require.Empty(t, projectExploreSourceLiteralConstructionAlignment( + context.Background(), graph.NewOverlaidView(base, tombstone), []string{sourceID}, + )) + + replacement := graph.NewOverlayLayer() + replacement.MarkFile(targetFile, false) + replacement.MarkRemoved("target", targetID) + replacement.AddNode(targetFile, &graph.Node{ID: targetID, Name: "target", Kind: graph.KindType, FilePath: targetFile}) + require.True(t, projectExploreSourceLiteralConstructionAlignment( + context.Background(), graph.NewOverlaidView(base, replacement), []string{sourceID}, + )[sourceID]) + }) +} + +func TestProjectExploreSourceLiteralConstructionAlignmentRejectsOverlayHiddenTruncation(t *testing.T) { + const ( + sourceFile = "demo/source.cs" + targetFile = "demo/target.cs" + sourceID = sourceFile + "::source" + targetID = targetFile + "::target" + ) + base := graph.New() + base.AddNode(&graph.Node{ID: sourceID, Name: "source", Kind: graph.KindMethod, FilePath: sourceFile}) + base.AddNode(&graph.Node{ID: targetID, Name: "target", Kind: graph.KindType, FilePath: targetFile}) + for line := 1; line <= graph.MaxBoundedAdjacencyRowsPerKey+1; line++ { + base.AddEdge(&graph.Edge{ + From: sourceID, To: targetID, Kind: graph.EdgeInstantiates, + FilePath: sourceFile, Line: line, + }) + } + layer := graph.NewOverlayLayer() + layer.MarkFile(targetFile, false) + layer.MarkRemoved("target", targetID) + + require.Empty(t, projectExploreSourceLiteralConstructionAlignment( + context.Background(), graph.NewOverlaidView(base, layer), []string{sourceID}, + ), "a base truncation hidden by the overlay is indeterminate, never positive") +} + +func TestMapExploreSourceLiteralMatchesDedupesProjectedCalleeIdentities(t *testing.T) { + const path = "demo/src/duplicate.cs" + owner := sourceLiteralNode(path+"::owner", "owner", path, graph.KindMethod, 1, 10) + callee := sourceLiteralNode(path+"::Resolve", "Resolve", path, graph.KindMethod, 20, 30) + server, counting := newExploreSourceLiteralGraphServer(t, []*graph.Node{owner, callee}, nil) + counting.siteLookup = func(_ context.Context, sites []graph.EdgeSourceSite, _ []graph.EdgeKind, _ int) (graph.BoundedSiteEdgeIdentityProjection, error) { + identity := graph.EdgeIdentity{From: owner.ID, To: callee.ID, Kind: graph.EdgeCalls, FilePath: path, Line: 3} + return graph.BoundedSiteEdgeIdentityProjection{BySite: map[graph.EdgeSourceSite][]graph.EdgeIdentity{ + sites[0]: {identity, identity}, + }}, nil + } + + recall := server.mapExploreSourceLiteralMatches("needle", []trigram.Match{{ + Path: path, Line: 3, Text: `Resolve("needle");`, + }}, query.QueryOptions{RepoAllow: map[string]bool{"demo": true}}) + requireSourceLiteralHitIdentity(t, recall.hits, exploreSourceLiteralHit{nodeID: callee.ID, rank: 0, callee: true}) + require.Equal(t, []string{callee.ID}, counting.lastNodeIDs) + require.Zero(t, counting.outEdgeBatchCalls) +} + +func TestMapExploreSourceLiteralMatchesOverlayAdjacencyParity(t *testing.T) { + const ( + sourceFile = "demo/source.cs" + targetFile = "demo/target.cs" + sourceID = sourceFile + "::owner" + targetID = targetFile + "::Resolve" + ) + newFixture := func(t *testing.T, layer *graph.OverlayLayer) (*Server, *graph.OverlaidView) { + t.Helper() + owner := sourceLiteralNode(sourceID, "owner", sourceFile, graph.KindMethod, 1, 10) + callee := sourceLiteralNode(targetID, "Resolve", targetFile, graph.KindMethod, 20, 30) + server, counting := newExploreSourceLiteralGraphServer(t, []*graph.Node{owner, callee}, []*graph.Edge{{ + From: sourceID, To: targetID, Kind: graph.EdgeCalls, FilePath: sourceFile, Line: 3, + }}) + return server, graph.NewOverlaidView(counting.Store, layer) + } + match := []trigram.Match{{Path: sourceFile, Line: 3, Text: `Resolve("needle");`}} + scope := query.QueryOptions{RepoAllow: map[string]bool{"demo": true}} + mapWithView := func(server *Server, view *graph.OverlaidView) exploreSourceLiteralRecall { + return server.mapExploreSourceLiteralMatchesContext( + WithOverlayView(context.Background(), view), "needle", match, scope, + ) + } + + t.Run("source tombstone and replacement", func(t *testing.T) { + tombstone := graph.NewOverlayLayer() + tombstone.MarkFile(sourceFile, false) + tombstone.MarkRemoved("owner", sourceID) + tombstone.AddEdge(&graph.Edge{From: sourceID, To: targetID, Kind: graph.EdgeCalls, Line: 3}) + server, view := newFixture(t, tombstone) + require.Empty(t, mapWithView(server, view).hits) + + replacement := graph.NewOverlayLayer() + replacement.MarkFile(sourceFile, false) + replacement.MarkRemoved("owner", sourceID) + replacement.AddNode(sourceFile, sourceLiteralNode(sourceID, "owner", sourceFile, graph.KindMethod, 1, 10)) + replacement.AddEdge(&graph.Edge{ + From: sourceID, To: targetID, Kind: graph.EdgeCalls, FilePath: sourceFile, Line: 3, + }) + server, view = newFixture(t, replacement) + recall := mapWithView(server, view) + requireSourceLiteralHitIdentity(t, recall.hits, exploreSourceLiteralHit{nodeID: targetID, rank: 0, callee: true}) + }) + + t.Run("target tombstone and replacement", func(t *testing.T) { + tombstone := graph.NewOverlayLayer() + tombstone.MarkFile(targetFile, false) + tombstone.MarkRemoved("Resolve", targetID) + server, view := newFixture(t, tombstone) + recall := mapWithView(server, view) + requireSourceLiteralHitIdentity(t, recall.hits, exploreSourceLiteralHit{nodeID: sourceID, rank: 0}) + + replacement := graph.NewOverlayLayer() + replacement.MarkFile(targetFile, false) + replacement.MarkRemoved("Resolve", targetID) + replacement.AddNode(targetFile, sourceLiteralNode(targetID, "Resolve", targetFile, graph.KindMethod, 20, 30)) + server, view = newFixture(t, replacement) + recall = mapWithView(server, view) + requireSourceLiteralHitIdentity(t, recall.hits, exploreSourceLiteralHit{nodeID: targetID, rank: 0, callee: true}) + }) +} + +func TestMapExploreSourceLiteralMatchesFiltersCalleeScopeAndTests(t *testing.T) { + for _, test := range []struct { + name string + configure func(*graph.Node, *graph.Node, *query.QueryOptions) + }{ + {name: "workspace", configure: func(owner, callee *graph.Node, scope *query.QueryOptions) { + owner.WorkspaceID = "demo" + callee.WorkspaceID = "other" + scope.WorkspaceID = "demo" + }}, + {name: "test", configure: func(_ *graph.Node, callee *graph.Node, scope *query.QueryOptions) { + callee.Meta = map[string]any{"is_test": true} + scope.ExcludeTests = true + }}, + } { + t.Run(test.name, func(t *testing.T) { + const path = "demo/src/scope.cs" + owner := sourceLiteralNode(path+"::owner", "owner", path, graph.KindMethod, 1, 10) + callee := sourceLiteralNode(path+"::Resolve", "Resolve", path, graph.KindMethod, 20, 30) + scope := query.QueryOptions{RepoAllow: map[string]bool{"demo": true}} + test.configure(owner, callee, &scope) + server, _ := newExploreSourceLiteralGraphServer(t, []*graph.Node{owner, callee}, []*graph.Edge{{ + From: owner.ID, To: callee.ID, Kind: graph.EdgeCalls, FilePath: path, Line: 3, + }}) + recall := server.mapExploreSourceLiteralMatches("needle", []trigram.Match{{ + Path: path, Line: 3, Text: `Resolve("needle");`, + }}, scope) + requireSourceLiteralHitIdentity(t, recall.hits, exploreSourceLiteralHit{nodeID: owner.ID, rank: 0}) + }) + } +} + +func TestMapExploreSourceLiteralMatchesDiscardsPartialSiteError(t *testing.T) { + const path = "demo/src/error.cs" + owner := sourceLiteralNode(path+"::owner", "owner", path, graph.KindMethod, 1, 10) + callee := sourceLiteralNode(path+"::Resolve", "Resolve", path, graph.KindMethod, 20, 30) + server, counting := newExploreSourceLiteralGraphServer(t, []*graph.Node{owner, callee}, nil) + counting.siteLookup = func(_ context.Context, sites []graph.EdgeSourceSite, _ []graph.EdgeKind, _ int) (graph.BoundedSiteEdgeIdentityProjection, error) { + return graph.BoundedSiteEdgeIdentityProjection{BySite: map[graph.EdgeSourceSite][]graph.EdgeIdentity{ + sites[0]: {{From: owner.ID, To: callee.ID, Kind: graph.EdgeCalls, FilePath: path, Line: 3}}, + }}, errors.New("backend failed after a partial page") + } + + recall := server.mapExploreSourceLiteralMatches("needle", []trigram.Match{{ + Path: path, Line: 3, Text: `Resolve("needle");`, + }}, query.QueryOptions{RepoAllow: map[string]bool{"demo": true}}) + requireSourceLiteralHitIdentity(t, recall.hits, exploreSourceLiteralHit{nodeID: owner.ID, rank: 0}) + require.Zero(t, counting.nodeLookupBatches) +} diff --git a/internal/mcp/explore_source_literal_overlay.go b/internal/mcp/explore_source_literal_overlay.go new file mode 100644 index 00000000..18b92852 --- /dev/null +++ b/internal/mcp/explore_source_literal_overlay.go @@ -0,0 +1,411 @@ +package mcp + +import ( + "bufio" + "context" + "fmt" + "sort" + "strconv" + "strings" + "time" + "unicode" + + "github.com/zzet/gortex/internal/query" + "github.com/zzet/gortex/internal/search/trigram" + "github.com/zzet/gortex/internal/testpath" +) + +const ( + exploreSourceLiteralOverlayMaxHits = 24 + exploreSourceLiteralOverlayMaxBytes = 4 << 20 + exploreSourceLiteralOverlayMaxLineBytes = 64 << 10 + exploreSourceLiteralOverlayMaxCoveredFiles = 256 + exploreSourceLiteralOverlayBudget = 75 * time.Millisecond +) + +// exploreSourceLiteralOverlayFile is the request-local form of an editor +// overlay. Path is already in canonical graph-path form and eligible records +// whether the request scope admits the file. Every path, including tombstones +// and out-of-scope replacements, still shadows its durable counterpart. +type exploreSourceLiteralOverlayFile struct { + path string + content string + deleted bool + eligible bool +} + +type exploreSourceLiteralOverlayScan struct { + matches []trigram.Match + covered map[string]struct{} + incomplete bool +} + +// effectiveExploreSourceLiteralScope intersects the caller's requested scope +// with the session's immutable workspace/repository boundary. Caller filters +// may narrow the session, never widen it. +func (s *Server) effectiveExploreSourceLiteralScope( + ctx context.Context, + requested query.QueryOptions, +) (query.QueryOptions, bool) { + merged := s.localizationNodeScopeWithTests(ctx, requested, requested.ExcludeTests) + if strings.HasPrefix(merged.WorkspaceID, unresolvedWorkspacePrefix) { + return query.QueryOptions{}, false + } + effective := requested + effective.WorkspaceID = merged.WorkspaceID + effective.ProjectID = merged.ProjectID + effective.RepoAllow = merged.RepoAllow + effective.ExcludeTests = merged.ExcludeTests + return effective, true +} + +func (s *Server) resolveExploreSourceLiteralRepoPrefix( + repoPrefix string, + scope query.QueryOptions, +) (string, bool) { + repoPrefix = strings.TrimSuffix(strings.TrimSpace(repoPrefix), "/") + if s != nil && s.multiIndexer == nil && s.indexer != nil { + actual := strings.TrimSuffix(strings.TrimSpace(s.indexer.RepoPrefix()), "/") + if actual != "" && repoPrefix != "" && repoPrefix != actual { + return "", false + } + if actual != "" && len(scope.RepoAllow) > 0 && !scope.RepoAllow[actual] { + return "", false + } + return actual, true + } + if repoPrefix != "" { + if len(scope.RepoAllow) > 0 && !scope.RepoAllow[repoPrefix] { + return "", false + } + return repoPrefix, true + } + for candidate, allowed := range scope.RepoAllow { + if !allowed { + continue + } + candidate = strings.TrimSuffix(strings.TrimSpace(candidate), "/") + if repoPrefix != "" && candidate != repoPrefix { + return "", false + } + repoPrefix = candidate + } + if repoPrefix != "" { + return repoPrefix, true + } + if s != nil && s.multiIndexer != nil { + prefixes := s.multiIndexer.RepoPrefixes() + if len(prefixes) == 1 { + return prefixes[0], true + } + if len(prefixes) > 1 { + return "", false + } + } + if s != nil && s.indexer != nil { + return s.indexer.RepoPrefix(), true + } + return "", true +} + +// snapshotExploreSourceLiteralOverlays consumes the exact immutable editor +// cohort pinned by request middleware, resolves each buffer to its graph path, +// and computes scope eligibility without parsing an overlay graph. A context +// that was not prepared intentionally has no overlay lane. +func (s *Server) snapshotExploreSourceLiteralOverlays( + ctx context.Context, + scope query.QueryOptions, + repoPrefix string, +) ([]exploreSourceLiteralOverlayFile, map[string]struct{}, bool, bool, error) { + covered := make(map[string]struct{}) + if s == nil || ctx == nil { + return nil, covered, false, false, nil + } + snapshot, ok := overlayRequestSnapshotFromContext(ctx) + if !ok || snapshot == nil { + if OverlayViewFromContext(ctx) != nil { + return nil, covered, true, false, fmt.Errorf("overlay view has no pinned request snapshot") + } + return nil, covered, false, false, nil + } + if snapshot.sessionID != SessionIDFromContext(ctx) { + return nil, covered, false, false, fmt.Errorf("overlay request snapshot belongs to session %q, not %q", snapshot.sessionID, SessionIDFromContext(ctx)) + } + if !snapshot.canonical { + return nil, covered, true, false, fmt.Errorf("overlay request snapshot is not canonical") + } + fileCapacity := len(snapshot.files) + if fileCapacity > exploreSourceLiteralOverlayMaxCoveredFiles { + fileCapacity = exploreSourceLiteralOverlayMaxCoveredFiles + } + files := make([]exploreSourceLiteralOverlayFile, 0, fileCapacity) + incomplete := false + recordCount := 0 + for _, overlay := range snapshot.files { + if err := ctx.Err(); err != nil { + return nil, covered, incomplete, false, err + } + absPath, resolveErr := s.resolveOverlayAbsPath(overlay.Path) + if resolveErr != nil || absPath == "" { + incomplete = true + continue + } + owner := s.pickIndexerForPath(absPath) + if owner == nil { + incomplete = true + continue + } + if repoPrefix != "" && owner.RepoPrefix() != repoPrefix { + continue + } + graphPath := canonicalExploreSourceLiteralPath(overlay.Path) + if graphPath == "" { + incomplete = true + continue + } + recordCount++ + if recordCount > exploreSourceLiteralOverlayMaxCoveredFiles { + // The request cohort is already canonical, so this sentinel counts + // distinct effective files rather than raw aliases. + return nil, covered, true, true, nil + } + covered[graphPath] = struct{}{} + eligible := repoNarrowAdmits(scope.RepoAllow, owner.RepoPrefix()) + if eligible && scope.WorkspaceID != "" { + eligible = owner.WorkspaceID() == scope.WorkspaceID + } + if eligible && scope.ProjectID != "" { + eligible = scope.WorkspaceID != "" && owner.ProjectID() == scope.ProjectID + } + if eligible && scope.ExcludeTests { + eligible = !testpath.IsTestFile(graphPath) + } + files = append(files, exploreSourceLiteralOverlayFile{ + path: graphPath, content: overlay.Content, deleted: overlay.Deleted, eligible: eligible, + }) + } + return files, covered, incomplete, false, nil +} + +// scanExploreSourceLiteralOverlays searches editor buffers without building an +// overlay graph. It keeps one representative line per file, uses a hard +// request-local byte/time envelope, and reads one sentinel hit beyond the +// caller's limit so truncation is explicit. +func scanExploreSourceLiteralOverlays( + ctx context.Context, + term string, + files []exploreSourceLiteralOverlayFile, + maxHits int, +) (exploreSourceLiteralOverlayScan, error) { + return scanExploreSourceLiteralOverlaysWithClock(ctx, term, files, maxHits, time.Now) +} + +func scanExploreSourceLiteralOverlaysWithClock( + ctx context.Context, + term string, + files []exploreSourceLiteralOverlayFile, + maxHits int, + now func() time.Time, +) (exploreSourceLiteralOverlayScan, error) { + result := exploreSourceLiteralOverlayScan{ + covered: make(map[string]struct{}, len(files)), + } + if err := ctx.Err(); err != nil { + return result, err + } + term = strings.TrimSpace(term) + if term == "" || len(files) == 0 { + return result, nil + } + + hitLimit := exploreSourceLiteralOverlayMaxHits + if maxHits > 0 && maxHits < hitLimit { + hitLimit = maxHits + } + ordered := make([]exploreSourceLiteralOverlayFile, 0, len(files)) + for _, file := range files { + file.path = canonicalExploreSourceLiteralPath(file.path) + if file.path == "" { + continue + } + result.covered[file.path] = struct{}{} + if file.eligible && !file.deleted { + ordered = append(ordered, file) + } + } + sort.SliceStable(ordered, func(i, j int) bool { + return exploreSourceLiteralPathLess(ordered[i].path, ordered[j].path) + }) + + if now == nil { + now = time.Now + } + deadline := now().Add(exploreSourceLiteralOverlayBudget) + inputBytes := 0 + for _, file := range ordered { + if err := ctx.Err(); err != nil { + return result, err + } + if !now().Before(deadline) { + result.incomplete = true + break + } + scanner := bufio.NewScanner(strings.NewReader(file.content)) + scanner.Buffer(make([]byte, 4096), exploreSourceLiteralOverlayMaxLineBytes) + line := 0 + for scanner.Scan() { + line++ + if err := ctx.Err(); err != nil { + return result, err + } + if !now().Before(deadline) { + result.incomplete = true + break + } + lineBytes := scanner.Bytes() + if len(lineBytes) > exploreSourceLiteralOverlayMaxBytes-inputBytes { + result.incomplete = true + break + } + inputBytes += len(lineBytes) + text := scanner.Text() + if !exploreOverlayLiteralHasBoundary(text, term) { + continue + } + result.matches = append(result.matches, trigram.Match{ + Path: file.path, + Line: line, + Text: strings.Clone(text), + }) + break + } + if scanner.Err() != nil { + result.incomplete = true + } + if result.incomplete && !now().Before(deadline) { + break + } + if len(result.matches) > hitLimit { + result.incomplete = true + result.matches = result.matches[:hitLimit] + break + } + } + return result, nil +} + +// mergeExploreSourceLiteralMatches masks durable rows for every covered +// overlay path, gives overlay rows precedence, then deduplicates and sorts the +// combined cohort before applying the result cap. +func mergeExploreSourceLiteralMatches( + overlay []trigram.Match, + durable []trigram.Match, + covered map[string]struct{}, + maxHits int, +) (matches []trigram.Match, incomplete bool) { + limit := exploreSourceLiteralOverlayMaxHits + if maxHits > 0 && maxHits < limit { + limit = maxHits + } + type rankedMatch struct { + match trigram.Match + overlay bool + } + ranked := make([]rankedMatch, 0, len(overlay)+len(durable)) + seen := make(map[string]struct{}, len(overlay)+len(durable)) + appendUnique := func(match trigram.Match, fromOverlay bool) { + match.Path = canonicalExploreSourceLiteralPath(match.Path) + if match.Path == "" { + return + } + key := match.Path + "\x00" + strconv.Itoa(match.Line) + if _, ok := seen[key]; ok { + return + } + seen[key] = struct{}{} + match.Text = strings.Clone(match.Text) + ranked = append(ranked, rankedMatch{match: match, overlay: fromOverlay}) + } + for _, match := range overlay { + appendUnique(match, true) + } + for _, match := range durable { + canonical := canonicalExploreSourceLiteralPath(match.Path) + if _, shadowed := covered[canonical]; shadowed { + continue + } + appendUnique(match, false) + } + sort.SliceStable(ranked, func(i, j int) bool { + left, right := ranked[i], ranked[j] + leftTest, rightTest := testpath.IsTestFile(left.match.Path), testpath.IsTestFile(right.match.Path) + if leftTest != rightTest { + return !leftTest + } + if left.overlay != right.overlay { + return left.overlay + } + if left.match.Path != right.match.Path { + return left.match.Path < right.match.Path + } + if left.match.Line != right.match.Line { + return left.match.Line < right.match.Line + } + return left.match.Text < right.match.Text + }) + matches = make([]trigram.Match, 0, len(ranked)) + for _, rankedMatch := range ranked { + matches = append(matches, rankedMatch.match) + } + if len(matches) > limit { + matches = matches[:limit] + incomplete = true + } + return matches, incomplete +} + +func exploreSourceLiteralPathLess(left, right string) bool { + leftTest, rightTest := testpath.IsTestFile(left), testpath.IsTestFile(right) + if leftTest != rightTest { + return !leftTest + } + return left < right +} + +func canonicalExploreSourceLiteralPath(candidate string) string { + return canonicalOverlayGraphPath(candidate) +} + +func exploreOverlayLiteralHasBoundary(text, term string) bool { + textRunes := []rune(text) + termRunes := []rune(strings.TrimSpace(term)) + if len(termRunes) == 0 || len(textRunes) < len(termRunes) { + return false + } + for start := 0; start+len(termRunes) <= len(textRunes); start++ { + matched := true + for offset := range termRunes { + if textRunes[start+offset] != termRunes[offset] { + matched = false + break + } + } + if !matched { + continue + } + end := start + len(termRunes) + leftOK := !exploreOverlayIdentifierRune(termRunes[0]) || start == 0 || + !exploreOverlayIdentifierRune(textRunes[start-1]) + rightOK := !exploreOverlayIdentifierRune(termRunes[len(termRunes)-1]) || end == len(textRunes) || + !exploreOverlayIdentifierRune(textRunes[end]) + if leftOK && rightOK { + return true + } + } + return false +} + +func exploreOverlayIdentifierRune(r rune) bool { + return r == '_' || unicode.IsLetter(r) || unicode.IsDigit(r) || + unicode.IsMark(r) || unicode.Is(unicode.Pc, r) +} diff --git a/internal/mcp/explore_source_literal_overlay_integration_test.go b/internal/mcp/explore_source_literal_overlay_integration_test.go new file mode 100644 index 00000000..c63b9437 --- /dev/null +++ b/internal/mcp/explore_source_literal_overlay_integration_test.go @@ -0,0 +1,475 @@ +package mcp + +import ( + "context" + "fmt" + "os" + "path/filepath" + "strings" + "testing" + "time" + + "go.uber.org/zap" + + "github.com/zzet/gortex/internal/config" + "github.com/zzet/gortex/internal/daemon" + "github.com/zzet/gortex/internal/graph" + "github.com/zzet/gortex/internal/indexer" + "github.com/zzet/gortex/internal/parser" + "github.com/zzet/gortex/internal/parser/languages" + "github.com/zzet/gortex/internal/query" +) + +func newExploreSourceLiteralOverlayTestServer(t *testing.T, diskFiles map[string]string) (*Server, string) { + t.Helper() + root := t.TempDir() + mtimes := make(map[string]int64, len(diskFiles)) + for rel, content := range diskFiles { + abs := filepath.Join(root, filepath.FromSlash(rel)) + if err := os.MkdirAll(filepath.Dir(abs), 0o755); err != nil { + t.Fatalf("mkdir %s: %v", rel, err) + } + if err := os.WriteFile(abs, []byte(content), 0o644); err != nil { + t.Fatalf("write %s: %v", rel, err) + } + mtimes[rel] = 1 + } + store := graph.New() + reg := parser.NewRegistry() + languages.RegisterAll(reg) + idx := indexer.New(store, reg, config.IndexConfig{}, zap.NewNop()) + idx.SetRootPath(root) + idx.SetRepoPrefix("repo") + idx.SetWorkspaceID("workspace") + idx.SetProjectID("project") + idx.SetFileMtimes(mtimes) + engine := query.NewEngine(store) + server := NewServer(engine, store, idx, nil, zap.NewNop(), nil) + manager := daemon.NewOverlayManager(time.Hour) + server.SetOverlayManager(manager) + const sessionID = "literal-overlay-test" + if err := manager.RegisterWithID(sessionID, "workspace"); err != nil { + t.Fatalf("register overlay: %v", err) + } + return server, sessionID +} + +func pushExploreSourceLiteralOverlay(t *testing.T, server *Server, sessionID, path, content string, deleted bool) { + t.Helper() + if err := server.overlays.Push(sessionID, daemon.OverlayFile{ + Path: path, Content: content, Deleted: deleted, + }, nil); err != nil { + t.Fatalf("push %s: %v", path, err) + } +} + +func prepareExploreSourceLiteralOverlayContext(t *testing.T, server *Server, sessionID string) context.Context { + t.Helper() + ctx, _, err := server.prepareOverlayRequest(WithSessionID(context.Background(), sessionID)) + if err != nil { + t.Fatalf("prepare overlay request: %v", err) + } + return ctx +} + +func TestSearchExploreSourceLiteralFindsOverlayOnlyAndMasksDurable(t *testing.T) { + server, sessionID := newExploreSourceLiteralOverlayTestServer(t, map[string]string{ + "replace.go": "package repo\nvar _ = \"needle\"\n", + "delete.go": "package repo\nvar _ = \"needle\"\n", + "keep.go": "package repo\nvar _ = \"needle\"\n", + }) + pushExploreSourceLiteralOverlay(t, server, sessionID, "replace.go", "package repo\nvar _ = \"other\"\n", false) + pushExploreSourceLiteralOverlay(t, server, sessionID, "delete.go", "", true) + pushExploreSourceLiteralOverlay(t, server, sessionID, "new.go", "package repo\nvar _ = \"needle\"\n", false) + + result := server.searchExploreSourceLiteral( + prepareExploreSourceLiteralOverlayContext(t, server, sessionID), "needle", "repo", + query.QueryOptions{WorkspaceID: "workspace", ProjectID: "project", RepoAllow: map[string]bool{"repo": true}}, + 24, + ) + if result.err != nil { + t.Fatalf("search: %v", result.err) + } + paths := make(map[string]bool, len(result.matches)) + for _, match := range result.matches { + paths[match.Path] = true + } + if !paths["new.go"] || !paths["keep.go"] { + t.Fatalf("overlay-only/durable match missing: %#v", result.matches) + } + if paths["replace.go"] || paths["delete.go"] { + t.Fatalf("covered durable match leaked: %#v", result.matches) + } +} + +func TestSearchExploreSourceLiteralUsesPinnedCohortAfterPush(t *testing.T) { + server, sessionID := newExploreSourceLiteralOverlayTestServer(t, nil) + pushExploreSourceLiteralOverlay(t, server, sessionID, "new.go", + "package repo\nvar _ = \"needle\"\n", false) + ctx := prepareExploreSourceLiteralOverlayContext(t, server, sessionID) + pushExploreSourceLiteralOverlay(t, server, sessionID, "new.go", + "package repo\nvar _ = \"other\"\n", false) + + result := server.searchExploreSourceLiteral(ctx, "needle", "repo", + query.QueryOptions{RepoAllow: map[string]bool{"repo": true}}, 24) + if result.err != nil { + t.Fatalf("search: %v", result.err) + } + if len(result.matches) != 1 || result.matches[0].Path != "new.go" { + t.Fatalf("pinned cohort changed after push: %#v", result.matches) + } +} + +func TestSearchExploreSourceLiteralOutOfScopeReplacementStillMasksDurable(t *testing.T) { + for _, test := range []struct { + name string + path string + scope query.QueryOptions + }{ + {name: "project", path: "replace.go", scope: query.QueryOptions{ + WorkspaceID: "workspace", ProjectID: "other-project", RepoAllow: map[string]bool{"repo": true}, + }}, + {name: "test", path: "replace_test.go", scope: query.QueryOptions{ + WorkspaceID: "workspace", RepoAllow: map[string]bool{"repo": true}, ExcludeTests: true, + }}, + } { + t.Run(test.name, func(t *testing.T) { + server, sessionID := newExploreSourceLiteralOverlayTestServer(t, map[string]string{ + test.path: "package repo\nvar _ = \"needle\"\n", + }) + pushExploreSourceLiteralOverlay(t, server, sessionID, test.path, + "package repo\nvar _ = \"needle\"\n", false) + ctx := prepareExploreSourceLiteralOverlayContext(t, server, sessionID) + files, covered, incomplete, overflow, err := server.snapshotExploreSourceLiteralOverlays(ctx, test.scope, "repo") + if err != nil || incomplete || overflow { + t.Fatalf("snapshot err/incomplete/overflow = %v/%t/%t", err, incomplete, overflow) + } + if _, ok := covered[test.path]; !ok { + t.Fatalf("covered paths = %#v, want %q", covered, test.path) + } + if len(files) != 1 || files[0].eligible { + t.Fatalf("overlay files = %#v, want one ineligible replacement", files) + } + result := server.searchExploreSourceLiteral(ctx, "needle", "repo", test.scope, 24) + if result.err != nil { + t.Fatalf("search: %v", result.err) + } + if len(result.matches) != 0 { + t.Fatalf("out-of-scope replacement leaked or durable was not masked: %#v", result.matches) + } + }) + } +} + +func TestSearchExploreSourceLiteralCompensatesCoveredLeadingDurableRows(t *testing.T) { + disk := make(map[string]string) + for index := 0; index < 50; index++ { + disk[fmt.Sprintf("a%02d.go", index)] = "package repo\nvar _ = \"needle\"\n" + } + disk["z-visible.go"] = "package repo\nvar _ = \"needle\"\n" + server, sessionID := newExploreSourceLiteralOverlayTestServer(t, disk) + for index := 0; index < 50; index++ { + pushExploreSourceLiteralOverlay(t, server, sessionID, fmt.Sprintf("a%02d.go", index), "package repo\n", false) + } + + result := server.searchExploreSourceLiteral( + prepareExploreSourceLiteralOverlayContext(t, server, sessionID), "needle", "repo", + query.QueryOptions{WorkspaceID: "workspace", RepoAllow: map[string]bool{"repo": true}}, + 1, + ) + if result.err != nil { + t.Fatalf("search: %v", result.err) + } + if len(result.matches) != 1 || result.matches[0].Path != "z-visible.go" { + t.Fatalf("compensated matches = %#v", result.matches) + } +} + +func newExploreSourceLiteralMultiRepoTestServer(t *testing.T) (*Server, string) { + t.Helper() + store := graph.New() + reg := parser.NewRegistry() + languages.RegisterAll(reg) + base := indexer.New(store, reg, config.IndexConfig{}, zap.NewNop()) + managerPath := filepath.Join(t.TempDir(), "missing-global.yaml") + configManager, err := config.NewConfigManager(managerPath) + if err != nil { + t.Fatalf("config manager: %v", err) + } + multi := indexer.NewMultiIndexer(store, reg, base.Search(), configManager, zap.NewNop()) + for _, prefix := range []string{"target", "other"} { + root := t.TempDir() + if err := os.WriteFile(filepath.Join(root, "disk.go"), []byte("package repo\nvar _ = \"needle\"\n"), 0o644); err != nil { + t.Fatalf("write %s: %v", prefix, err) + } + if _, err := multi.TrackRepo(config.RepoEntry{ + Name: prefix, Path: root, Workspace: "workspace", Project: "project", + }); err != nil { + t.Fatalf("track %s: %v", prefix, err) + } + } + engine := query.NewEngine(store) + server := NewServer(engine, store, nil, nil, zap.NewNop(), nil, MultiRepoOptions{ + MultiIndexer: multi, ConfigManager: configManager, + }) + manager := daemon.NewOverlayManager(time.Hour) + server.SetOverlayManager(manager) + const sessionID = "literal-overlay-multi-test" + if err := manager.RegisterWithID(sessionID, "workspace"); err != nil { + t.Fatalf("register overlay: %v", err) + } + return server, sessionID +} + +func TestSearchExploreSourceLiteralIsolatesTargetRepoOverlays(t *testing.T) { + server, sessionID := newExploreSourceLiteralMultiRepoTestServer(t) + for index := 0; index < 30; index++ { + pushExploreSourceLiteralOverlay(t, server, sessionID, + fmt.Sprintf("other/overlay-%02d.go", index), "package other\nvar _ = \"needle\"\n", false) + } + + result := server.searchExploreSourceLiteral( + prepareExploreSourceLiteralOverlayContext(t, server, sessionID), "needle", "target", query.QueryOptions{}, 1, + ) + if result.err != nil { + t.Fatalf("search: %v", result.err) + } + if result.incomplete { + t.Fatal("unrelated overlay volume marked target search incomplete") + } + if len(result.matches) != 1 || result.matches[0].Path != "target/disk.go" { + t.Fatalf("cross-repo overlay leaked or target starved: %#v", result.matches) + } +} + +func TestOverlayRequestCanonicalizesMultiRepoAliasesAcrossConsumers(t *testing.T) { + server, sessionID := newExploreSourceLiteralMultiRepoTestServer(t) + target := server.multiIndexer.GetIndexer("target") + if target == nil { + t.Fatal("target indexer missing") + } + absPath := filepath.Join(target.RootPath(), "disk.go") + const content = "package repo\nvar _ = \"overlay_alias\"\n" + for _, alias := range []string{absPath, "target/disk.go"} { + pushExploreSourceLiteralOverlay(t, server, sessionID, alias, content, false) + } + + ctx := prepareExploreSourceLiteralOverlayContext(t, server, sessionID) + snapshot, ok := overlayRequestSnapshotFromContext(ctx) + if !ok || !snapshot.canonical || len(snapshot.files) != 1 { + t.Fatalf("canonical snapshot = %#v", snapshot) + } + if snapshot.files[0].Path != "target/disk.go" { + t.Fatalf("canonical path = %q", snapshot.files[0].Path) + } + if got, found := server.overlayContentFor(ctx, absPath); !found || got != content { + t.Fatalf("overlay content = %q, %v", got, found) + } + result := server.searchExploreSourceLiteral( + ctx, "overlay_alias", "target", query.QueryOptions{RepoAllow: map[string]bool{"target": true}}, 1, + ) + if result.err != nil || len(result.matches) != 1 || result.matches[0].Path != "target/disk.go" { + t.Fatalf("literal result = %#v, err=%v", result.matches, result.err) + } +} + +func TestSearchExploreSourceLiteralSessionProjectDoesNotHideSiblingProject(t *testing.T) { + server, sessionID := newExploreSourceLiteralMultiRepoTestServer(t) + home := server.multiIndexer.GetIndexer("target") + sibling := server.multiIndexer.GetIndexer("other") + if home == nil || sibling == nil { + t.Fatal("multi-repo test indexers missing") + } + home.SetProjectID("project-a") + sibling.SetProjectID("project-b") + pushExploreSourceLiteralOverlay(t, server, sessionID, "other/new.go", + "package other\nvar _ = \"needle\"\n", false) + baseCtx := WithSessionCWD(WithSessionID(context.Background(), sessionID), home.RootPath()) + ctx, _, err := server.prepareOverlayRequest(baseCtx) + if err != nil { + t.Fatalf("prepare: %v", err) + } + + effective, ok := server.effectiveExploreSourceLiteralScope(ctx, query.QueryOptions{}) + if !ok || effective.ProjectID != "" || effective.WorkspaceID != "workspace" { + t.Fatalf("effective sibling-project scope = %#v, %v", effective, ok) + } + result := server.searchExploreSourceLiteral(ctx, "needle", "other", query.QueryOptions{}, 1) + if result.err != nil { + t.Fatalf("search: %v", result.err) + } + if len(result.matches) != 1 || !strings.HasPrefix(result.matches[0].Path, "other/") { + t.Fatalf("session home project hid sibling-project result: %#v", result.matches) + } + mismatch := server.searchExploreSourceLiteral(ctx, "needle", "other", query.QueryOptions{ + RepoAllow: map[string]bool{"target": true}, + }, 1) + if mismatch.owned || len(mismatch.matches) != 0 || mismatch.backend != "multi-ambiguous-scope" { + t.Fatalf("disjoint request narrow escaped scope: %#v", mismatch) + } +} + +func TestResolveExploreSourceLiteralRepoPrefixRejectsScopeMismatch(t *testing.T) { + server, _ := newExploreSourceLiteralOverlayTestServer(t, nil) + if prefix, ok := server.resolveExploreSourceLiteralRepoPrefix("repo", query.QueryOptions{ + RepoAllow: map[string]bool{"other": true}, + }); ok || prefix != "" { + t.Fatalf("mismatched explicit prefix = %q, %v", prefix, ok) + } +} + +func TestSearchExploreSourceLiteralRejectsUnresolvedOverlayOwnership(t *testing.T) { + server, sessionID := newExploreSourceLiteralMultiRepoTestServer(t) + pushExploreSourceLiteralOverlay(t, server, sessionID, "ambiguous.go", "needle\n", false) + _, _, err := server.prepareOverlayRequest(WithSessionID(context.Background(), sessionID)) + if err == nil { + t.Fatal("unresolved overlay ownership did not fail request preparation") + } +} + +func TestSearchExploreSourceLiteralFailsClosedAboveCoveredPathCap(t *testing.T) { + server, sessionID := newExploreSourceLiteralOverlayTestServer(t, map[string]string{ + "visible.go": "package repo\nvar _ = \"needle\"\n", + }) + for index := 0; index <= exploreSourceLiteralOverlayMaxCoveredFiles; index++ { + pushExploreSourceLiteralOverlay(t, server, sessionID, + fmt.Sprintf("covered-%03d.go", index), "package repo\n", false) + } + result := server.searchExploreSourceLiteral( + prepareExploreSourceLiteralOverlayContext(t, server, sessionID), "needle", "repo", + query.QueryOptions{RepoAllow: map[string]bool{"repo": true}}, 24, + ) + if result.err != nil { + t.Fatalf("search: %v", result.err) + } + if !result.incomplete || len(result.matches) != 0 { + t.Fatalf("covered-cap result: matches=%#v incomplete=%v", result.matches, result.incomplete) + } + if result.backend != "overlay-covered-cap" { + t.Fatalf("backend = %q", result.backend) + } +} + +func TestSearchExploreSourceLiteralAllowsExactlyCoveredPathCap(t *testing.T) { + server, sessionID := newExploreSourceLiteralOverlayTestServer(t, map[string]string{ + "visible.go": "package repo\nvar _ = \"needle\"\n", + }) + for index := 0; index < exploreSourceLiteralOverlayMaxCoveredFiles; index++ { + pushExploreSourceLiteralOverlay(t, server, sessionID, + fmt.Sprintf("covered-%03d.go", index), "package repo\n", false) + } + result := server.searchExploreSourceLiteral( + prepareExploreSourceLiteralOverlayContext(t, server, sessionID), "needle", "repo", + query.QueryOptions{RepoAllow: map[string]bool{"repo": true}}, 1, + ) + if result.err != nil { + t.Fatalf("search: %v", result.err) + } + if result.incomplete || len(result.matches) != 1 || result.matches[0].Path != "visible.go" { + t.Fatalf("exact-cap result: matches=%#v incomplete=%v", result.matches, result.incomplete) + } +} + +func TestSearchExploreSourceLiteralRejectsViewWithoutPinnedSnapshot(t *testing.T) { + server, _ := newExploreSourceLiteralOverlayTestServer(t, map[string]string{ + "visible.go": "package repo\nvar _ = \"needle\"\n", + }) + ctx := WithOverlayView(context.Background(), graph.NewOverlaidView(server.graph, graph.NewOverlayLayer())) + result := server.searchExploreSourceLiteral(ctx, "needle", "repo", query.QueryOptions{}, 1) + if result.err == nil || len(result.matches) != 0 || !result.incomplete { + t.Fatalf("missing snapshot result: matches=%#v incomplete=%v err=%v", result.matches, result.incomplete, result.err) + } +} + +func TestSearchExploreSourceLiteralUsesCanonicalAliasCohort(t *testing.T) { + server, sessionID := newExploreSourceLiteralOverlayTestServer(t, map[string]string{ + "alias.go": "package repo\nvar _ = \"needle\"\n", + }) + const overlayContent = "package repo\nvar _ = \"other\"\n" + pushExploreSourceLiteralOverlay(t, server, sessionID, "./alias.go", overlayContent, false) + pushExploreSourceLiteralOverlay(t, server, sessionID, "alias.go", overlayContent, false) + ctx := prepareExploreSourceLiteralOverlayContext(t, server, sessionID) + snapshot, ok := overlayRequestSnapshotFromContext(ctx) + if !ok || !snapshot.canonical || len(snapshot.files) != 1 { + t.Fatalf("canonical snapshot = %#v", snapshot) + } + result := server.searchExploreSourceLiteral( + ctx, "needle", "repo", query.QueryOptions{RepoAllow: map[string]bool{"repo": true}}, 1, + ) + if result.err != nil { + t.Fatalf("search: %v", result.err) + } + if len(result.matches) != 0 { + t.Fatalf("canonical alias did not mask durable content: %#v", result.matches) + } +} + +func TestSnapshotExploreSourceLiteralOverlaysCollapsesCanonicalAliases(t *testing.T) { + server, sessionID := newExploreSourceLiteralOverlayTestServer(t, nil) + for index := 0; index <= exploreSourceLiteralOverlayMaxCoveredFiles; index++ { + pushExploreSourceLiteralOverlay(t, server, sessionID, + strings.Repeat("./", index+1)+"alias.go", "package repo\n", false) + } + ctx := prepareExploreSourceLiteralOverlayContext(t, server, sessionID) + files, covered, incomplete, overflow, err := server.snapshotExploreSourceLiteralOverlays( + ctx, query.QueryOptions{RepoAllow: map[string]bool{"repo": true}}, "repo", + ) + if err != nil { + t.Fatalf("snapshot: %v", err) + } + if incomplete || overflow || len(files) != 1 || len(covered) != 1 { + t.Fatalf("alias cohort: files=%d covered=%d incomplete=%v overflow=%v", len(files), len(covered), incomplete, overflow) + } +} + +func TestResolveExploreSourceLiteralRepoPrefixRejectsDirectMismatch(t *testing.T) { + server, _ := newExploreSourceLiteralOverlayTestServer(t, nil) + if prefix, ok := server.resolveExploreSourceLiteralRepoPrefix("other", query.QueryOptions{}); ok || prefix != "" { + t.Fatalf("mismatched direct prefix = %q, %v", prefix, ok) + } +} + +func TestSearchExploreSourceLiteralNormalizesHitLimits(t *testing.T) { + diskFiles := make(map[string]string, 30) + for i := 0; i < 30; i++ { + diskFiles[fmt.Sprintf("file-%02d.go", i)] = `const value = "needle"` + } + server, sessionID := newExploreSourceLiteralOverlayTestServer(t, diskFiles) + ctx := prepareExploreSourceLiteralOverlayContext(t, server, sessionID) + scope := query.QueryOptions{RepoAllow: map[string]bool{"repo": true}} + + for _, maxHits := range []int{0, -1, 10_000} { + t.Run(fmt.Sprintf("max_hits_%d", maxHits), func(t *testing.T) { + got := server.searchExploreSourceLiteral(ctx, "needle", "repo", scope, maxHits) + if got.err != nil { + t.Fatalf("search: %v", got.err) + } + if len(got.matches) != exploreSourceLiteralOverlayMaxHits || !got.incomplete { + t.Fatalf("matches/incomplete = %d/%t, want %d/true", len(got.matches), got.incomplete, exploreSourceLiteralOverlayMaxHits) + } + }) + } +} + +func TestSearchExploreSourceLiteralAppliesScopeAndCancellation(t *testing.T) { + server, sessionID := newExploreSourceLiteralOverlayTestServer(t, nil) + pushExploreSourceLiteralOverlay(t, server, sessionID, "new.go", "package repo\nvar _ = \"needle\"\n", false) + ctx := prepareExploreSourceLiteralOverlayContext(t, server, sessionID) + + outOfScope := server.searchExploreSourceLiteral(ctx, "needle", "repo", query.QueryOptions{ + WorkspaceID: "other", RepoAllow: map[string]bool{"repo": true}, + }, 24) + if outOfScope.err != nil { + t.Fatalf("scoped search: %v", outOfScope.err) + } + if len(outOfScope.matches) != 0 { + t.Fatalf("out-of-scope overlay leaked: %#v", outOfScope.matches) + } + + cancelled, cancel := context.WithCancel(ctx) + cancel() + result := server.searchExploreSourceLiteral(cancelled, "needle", "repo", query.QueryOptions{}, 24) + if result.err == nil { + t.Fatal("cancelled search did not return an error") + } +} diff --git a/internal/mcp/explore_source_literal_overlay_test.go b/internal/mcp/explore_source_literal_overlay_test.go new file mode 100644 index 00000000..fcd230b0 --- /dev/null +++ b/internal/mcp/explore_source_literal_overlay_test.go @@ -0,0 +1,384 @@ +package mcp + +import ( + "context" + "errors" + "strings" + "testing" + "time" + + "github.com/zzet/gortex/internal/search/trigram" +) + +func TestScanExploreSourceLiteralOverlaysDeterministicScopeAndBoundaries(t *testing.T) { + files := []exploreSourceLiteralOverlayFile{ + {path: `repo\\z.go`, content: "package z\nvar _ = \"café\"\n", eligible: true}, + {path: "repo/a.go", content: "var _ = \"caféteria\"\nvar x = `café`\n", eligible: true}, + {path: "repo/case.go", content: "var _ = `CAFÉ`\n", eligible: true}, + {path: "repo/deleted.go", content: "café", deleted: true, eligible: true}, + {path: "outside/no.go", content: "café", eligible: false}, + } + + got, err := scanExploreSourceLiteralOverlays(context.Background(), "café", files, 24) + if err != nil { + t.Fatalf("scan: %v", err) + } + if got.incomplete { + t.Fatal("bounded scan unexpectedly incomplete") + } + if len(got.matches) != 2 { + t.Fatalf("matches = %#v, want 2", got.matches) + } + if got.matches[0].Path != "repo/a.go" || got.matches[0].Line != 2 { + t.Fatalf("first match = %#v", got.matches[0]) + } + if got.matches[1].Path != "repo/z.go" || got.matches[1].Line != 2 { + t.Fatalf("second match = %#v", got.matches[1]) + } + for _, covered := range []string{"repo/a.go", "repo/case.go", "repo/z.go", "repo/deleted.go", "outside/no.go"} { + if _, ok := got.covered[covered]; !ok { + t.Errorf("covered path %q missing: %#v", covered, got.covered) + } + } + + for _, text := range []string{"cafe\u0301", "cafe‿value", "decafeteria"} { + if exploreOverlayLiteralHasBoundary(text, "cafe") { + t.Errorf("%q accepted as an identifier-bounded cafe", text) + } + } + if !exploreOverlayLiteralHasBoundary("call(cafe)", "cafe") { + t.Fatal("punctuation-delimited literal not accepted") + } +} + +func TestScanExploreSourceLiteralOverlaysHitSentinel(t *testing.T) { + makeFiles := func(count int) []exploreSourceLiteralOverlayFile { + files := make([]exploreSourceLiteralOverlayFile, 0, count) + for index := 0; index < count; index++ { + files = append(files, exploreSourceLiteralOverlayFile{ + path: strings.Repeat("a", index+1) + ".go", + content: "needle\n", + eligible: true, + }) + } + return files + } + + exact, err := scanExploreSourceLiteralOverlays(context.Background(), "needle", makeFiles(24), 24) + if err != nil { + t.Fatalf("exact scan: %v", err) + } + if len(exact.matches) != 24 || exact.incomplete { + t.Fatalf("exact cap: matches=%d incomplete=%v, want 24/false", len(exact.matches), exact.incomplete) + } + + over, err := scanExploreSourceLiteralOverlays(context.Background(), "needle", makeFiles(25), 24) + if err != nil { + t.Fatalf("sentinel scan: %v", err) + } + if len(over.matches) != 24 || !over.incomplete { + t.Fatalf("sentinel: matches=%d incomplete=%v, want 24/true", len(over.matches), over.incomplete) + } +} + +func TestScanExploreSourceLiteralOverlaysCapsAndCancellation(t *testing.T) { + tooLong := strings.Repeat("x", exploreSourceLiteralOverlayMaxLineBytes+1) + " needle" + got, err := scanExploreSourceLiteralOverlays(context.Background(), "needle", []exploreSourceLiteralOverlayFile{ + {path: "repo/long.go", content: tooLong, eligible: true}, + }, 24) + if err != nil { + t.Fatalf("long-line scan: %v", err) + } + if !got.incomplete || len(got.matches) != 0 { + t.Fatalf("long line: matches=%#v incomplete=%v", got.matches, got.incomplete) + } + + got, err = scanExploreSourceLiteralOverlays(context.Background(), "needle", []exploreSourceLiteralOverlayFile{ + {path: "repo/huge.go", content: strings.Repeat("x", exploreSourceLiteralOverlayMaxBytes+1), eligible: true}, + }, 24) + if err != nil { + t.Fatalf("byte-cap scan: %v", err) + } + if !got.incomplete || len(got.matches) != 0 { + t.Fatalf("byte cap: matches=%#v incomplete=%v", got.matches, got.incomplete) + } + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + _, err = scanExploreSourceLiteralOverlays(ctx, "needle", []exploreSourceLiteralOverlayFile{ + {path: "repo/a.go", content: "needle", eligible: true}, + }, 24) + if !errors.Is(err, context.Canceled) { + t.Fatalf("cancellation error = %v, want context.Canceled", err) + } +} + +func TestScanExploreSourceLiteralOverlaysChargesOnlyInspectedBytes(t *testing.T) { + largeTail := strings.Repeat("x\n", exploreSourceLiteralOverlayMaxBytes/2+1) + early, err := scanExploreSourceLiteralOverlays(context.Background(), "needle", []exploreSourceLiteralOverlayFile{ + {path: "repo/early.go", content: "needle\n" + largeTail, eligible: true}, + }, 24) + if err != nil { + t.Fatalf("early scan: %v", err) + } + if early.incomplete || len(early.matches) != 1 { + t.Fatalf("early hit: matches=%#v incomplete=%v", early.matches, early.incomplete) + } + + line := strings.Repeat("x", exploreSourceLiteralOverlayMaxLineBytes-1) + "\n" + lateContent := strings.Repeat(line, exploreSourceLiteralOverlayMaxBytes/(exploreSourceLiteralOverlayMaxLineBytes-1)+1) + "needle\n" + late, err := scanExploreSourceLiteralOverlays(context.Background(), "needle", []exploreSourceLiteralOverlayFile{ + {path: "repo/late.go", content: lateContent, eligible: true}, + }, 24) + if err != nil { + t.Fatalf("late scan: %v", err) + } + if !late.incomplete || len(late.matches) != 0 { + t.Fatalf("late hit: matches=%#v incomplete=%v", late.matches, late.incomplete) + } +} + +func TestScanExploreSourceLiteralOverlaysDeadline(t *testing.T) { + started := time.Unix(1, 0) + calls := 0 + now := func() time.Time { + calls++ + if calls == 1 { + return started + } + return started.Add(exploreSourceLiteralOverlayBudget) + } + got, err := scanExploreSourceLiteralOverlaysWithClock( + context.Background(), "needle", + []exploreSourceLiteralOverlayFile{{path: "repo/a.go", content: "needle", eligible: true}}, + 24, now, + ) + if err != nil { + t.Fatalf("scan: %v", err) + } + if !got.incomplete || len(got.matches) != 0 { + t.Fatalf("deadline: matches=%#v incomplete=%v", got.matches, got.incomplete) + } +} + +func TestScanExploreSourceLiteralOverlaysPrefersProductionBeforeTestCap(t *testing.T) { + files := make([]exploreSourceLiteralOverlayFile, 0, exploreSourceLiteralOverlayMaxHits+2) + for i := 0; i <= exploreSourceLiteralOverlayMaxHits; i++ { + files = append(files, exploreSourceLiteralOverlayFile{ + path: "repo/aaa/tests/test_" + strings.Repeat("a", i+1) + ".py", + content: `label = "needle"`, + eligible: true, + }) + } + files = append(files, exploreSourceLiteralOverlayFile{ + path: "repo/zzz/app.py", content: `label = "needle"`, eligible: true, + }) + + got, err := scanExploreSourceLiteralOverlays(context.Background(), "needle", files, exploreSourceLiteralOverlayMaxHits) + if err != nil { + t.Fatalf("scan: %v", err) + } + if len(got.matches) != exploreSourceLiteralOverlayMaxHits || !got.incomplete { + t.Fatalf("matches/incomplete = %d/%t, want %d/true", len(got.matches), got.incomplete, exploreSourceLiteralOverlayMaxHits) + } + if got.matches[0].Path != "repo/zzz/app.py" { + t.Fatalf("first match = %q, want production path", got.matches[0].Path) + } +} + +func TestMergeExploreSourceLiteralMatchesPreservesProductionPriority(t *testing.T) { + got, incomplete := mergeExploreSourceLiteralMatches(nil, []trigram.Match{ + {Path: "repo/aaa/tests/test_alpha.py", Line: 1, Text: "test"}, + {Path: "repo/zzz/app.py", Line: 2, Text: "production"}, + }, nil, exploreSourceLiteralOverlayMaxHits) + if incomplete || len(got) != 2 { + t.Fatalf("matches/incomplete = %d/%t, want 2/false", len(got), incomplete) + } + if got[0].Path != "repo/zzz/app.py" { + t.Fatalf("first match = %q, want production path", got[0].Path) + } +} + +func TestFinalizeExploreSourceLiteralSearchPreservesDurableFastPath(t *testing.T) { + durable := []trigram.Match{ + {Path: "repo/production.py", Line: 2, Text: "production"}, + {Path: "repo/tests/test_alpha.py", Line: 1, Text: "test"}, + {Path: "repo/extra.py", Line: 3, Text: "sentinel"}, + } + got := finalizeExploreSourceLiteralSearch( + exploreSourceLiteralSearch{matches: durable, backend: "direct", owned: true}, + exploreSourceLiteralOverlayScan{}, nil, 2, false, false, nil, + ) + if len(got.matches) != 2 || !got.incomplete { + t.Fatalf("matches/incomplete = %d/%t, want 2/true", len(got.matches), got.incomplete) + } + if got.matches[0].Path != durable[0].Path || got.matches[1].Path != durable[1].Path { + t.Fatalf("durable order changed: %#v", got.matches) + } + if &got.matches[0] != &durable[0] { + t.Fatal("durable fast path allocated a replacement backing slice") + } +} + +func TestFinalizeExploreSourceLiteralSearchRetainsOnlySafePartialEvidence(t *testing.T) { + deadlineErr := context.DeadlineExceeded + got := finalizeExploreSourceLiteralSearch( + exploreSourceLiteralSearch{ + matches: []trigram.Match{ + {Path: "repo/covered.go", Line: 1, Text: "stale durable"}, + {Path: "repo/safe.go", Line: 2, Text: "safe durable"}, + }, + backend: "direct", owned: true, + }, + exploreSourceLiteralOverlayScan{ + matches: []trigram.Match{{Path: "repo/live.go", Line: 3, Text: "partial overlay"}}, + incomplete: true, + }, + map[string]struct{}{"repo/covered.go": {}}, + exploreSourceLiteralOverlayMaxHits, + true, + false, + deadlineErr, + ) + if !errors.Is(got.err, deadlineErr) || !got.incomplete { + t.Fatalf("err/incomplete = %v/%t, want deadline/true", got.err, got.incomplete) + } + if len(got.matches) != 2 { + t.Fatalf("matches = %#v, want safe durable plus partial overlay", got.matches) + } + for _, match := range got.matches { + if match.Path == "repo/covered.go" { + t.Fatalf("covered durable row leaked: %#v", got.matches) + } + } + if got.matches[0].Path != "repo/live.go" || got.matches[1].Path != "repo/safe.go" { + t.Fatalf("matches = %#v, want live overlay then safe durable", got.matches) + } +} + +func TestRunExploreSourceLiteralLanesOrdersAndBoundsWork(t *testing.T) { + productionOverlay := trigram.Match{Path: "overlay.go", Line: 1, Text: "needle"} + testOverlay := trigram.Match{Path: "a_test.go", Line: 1, Text: "needle"} + productionDurable := trigram.Match{Path: "durable.go", Line: 1, Text: "needle"} + + t.Run("overlay scans before durable and wins within production", func(t *testing.T) { + order := make([]string, 0, 2) + result := runExploreSourceLiteralLanes( + context.Background(), "needle", "repo", 2, + []exploreSourceLiteralOverlayFile{{path: "overlay.go"}}, map[string]struct{}{"overlay.go": {}}, false, false, + func(context.Context, string, []exploreSourceLiteralOverlayFile, int) (exploreSourceLiteralOverlayScan, error) { + order = append(order, "overlay") + return exploreSourceLiteralOverlayScan{matches: []trigram.Match{productionOverlay}}, nil + }, + func(context.Context, string, string, int) exploreSourceLiteralSearch { + order = append(order, "durable") + return exploreSourceLiteralSearch{matches: []trigram.Match{productionDurable}, backend: "direct", owned: true} + }, + ) + if strings.Join(order, ",") != "overlay,durable" { + t.Fatalf("lane order = %v", order) + } + if len(result.matches) != 2 || result.matches[0].Path != "overlay.go" || result.matches[1].Path != "durable.go" { + t.Fatalf("matches = %#v", result.matches) + } + }) + + t.Run("durable production outranks overlay test", func(t *testing.T) { + result := runExploreSourceLiteralLanes( + context.Background(), "needle", "repo", 2, + []exploreSourceLiteralOverlayFile{{path: "a_test.go"}}, map[string]struct{}{"a_test.go": {}}, false, false, + func(context.Context, string, []exploreSourceLiteralOverlayFile, int) (exploreSourceLiteralOverlayScan, error) { + return exploreSourceLiteralOverlayScan{matches: []trigram.Match{testOverlay}}, nil + }, + func(context.Context, string, string, int) exploreSourceLiteralSearch { + return exploreSourceLiteralSearch{matches: []trigram.Match{productionDurable}, backend: "direct", owned: true} + }, + ) + if len(result.matches) != 2 || result.matches[0].Path != "durable.go" || result.matches[1].Path != "a_test.go" { + t.Fatalf("matches = %#v", result.matches) + } + }) + + t.Run("full production overlay skips durable", func(t *testing.T) { + durableCalls := 0 + result := runExploreSourceLiteralLanes( + context.Background(), "needle", "repo", 1, + []exploreSourceLiteralOverlayFile{{path: "overlay.go"}}, map[string]struct{}{"overlay.go": {}}, false, false, + func(context.Context, string, []exploreSourceLiteralOverlayFile, int) (exploreSourceLiteralOverlayScan, error) { + return exploreSourceLiteralOverlayScan{matches: []trigram.Match{productionOverlay}}, nil + }, + func(context.Context, string, string, int) exploreSourceLiteralSearch { + durableCalls++ + return exploreSourceLiteralSearch{} + }, + ) + if durableCalls != 0 || !result.incomplete || len(result.matches) != 1 || result.matches[0].Path != "overlay.go" { + t.Fatalf("durable calls=%d incomplete=%v matches=%#v", durableCalls, result.incomplete, result.matches) + } + }) + + t.Run("no overlay preserves durable order and skips scanner", func(t *testing.T) { + scanCalls := 0 + durable := []trigram.Match{{Path: "z.go", Line: 1}, {Path: "a.go", Line: 1}} + result := runExploreSourceLiteralLanes( + context.Background(), "needle", "repo", 2, nil, nil, false, false, + func(context.Context, string, []exploreSourceLiteralOverlayFile, int) (exploreSourceLiteralOverlayScan, error) { + scanCalls++ + return exploreSourceLiteralOverlayScan{}, nil + }, + func(context.Context, string, string, int) exploreSourceLiteralSearch { + return exploreSourceLiteralSearch{matches: durable, backend: "direct", owned: true} + }, + ) + if scanCalls != 0 || result.matches[0].Path != "z.go" || result.matches[1].Path != "a.go" { + t.Fatalf("scan calls=%d matches=%#v", scanCalls, result.matches) + } + }) + + t.Run("scan error returns safe partial without durable", func(t *testing.T) { + durableCalls := 0 + result := runExploreSourceLiteralLanes( + context.Background(), "needle", "repo", 2, + []exploreSourceLiteralOverlayFile{{path: "overlay.go"}}, map[string]struct{}{"shadowed.go": {}}, false, false, + func(context.Context, string, []exploreSourceLiteralOverlayFile, int) (exploreSourceLiteralOverlayScan, error) { + return exploreSourceLiteralOverlayScan{matches: []trigram.Match{productionOverlay}, incomplete: true}, context.DeadlineExceeded + }, + func(context.Context, string, string, int) exploreSourceLiteralSearch { + durableCalls++ + return exploreSourceLiteralSearch{} + }, + ) + if durableCalls != 0 || result.err != context.DeadlineExceeded || !result.incomplete || len(result.matches) != 1 { + t.Fatalf("durable calls=%d result=%#v", durableCalls, result) + } + }) +} + +func TestMergeExploreSourceLiteralMatchesOverlayWinsAndSortsBeforeLimit(t *testing.T) { + overlay := []trigram.Match{ + {Path: `repo\\same.go`, Line: 9, Text: "overlay replacement"}, + {Path: "repo/z.go", Line: 1, Text: "overlay z"}, + } + durable := []trigram.Match{ + {Path: "repo/same.go", Line: 9, Text: "stale same-line text"}, + {Path: "repo/same.go", Line: 1, Text: "stale covered path"}, + {Path: "repo/a.go", Line: 2, Text: "durable a"}, + {Path: "repo/a.go", Line: 2, Text: "durable duplicate"}, + } + matches, incomplete := mergeExploreSourceLiteralMatches(overlay, durable, map[string]struct{}{ + "repo/same.go": {}, + }, 24) + if incomplete || len(matches) != 3 { + t.Fatalf("full merge = %#v incomplete=%v", matches, incomplete) + } + if matches[0].Text != "overlay replacement" || matches[1].Path != "repo/z.go" || matches[2].Path != "repo/a.go" { + t.Fatalf("merge order/precedence = %#v", matches) + } + + capped, incomplete := mergeExploreSourceLiteralMatches(overlay, durable, map[string]struct{}{ + "repo/same.go": {}, + }, 2) + if !incomplete || len(capped) != 2 || capped[0].Text != "overlay replacement" || capped[1].Path != "repo/z.go" { + t.Fatalf("capped merge = %#v incomplete=%v", capped, incomplete) + } +} diff --git a/internal/mcp/explore_source_literal_test.go b/internal/mcp/explore_source_literal_test.go index 406594a3..abf0c9ed 100644 --- a/internal/mcp/explore_source_literal_test.go +++ b/internal/mcp/explore_source_literal_test.go @@ -28,9 +28,21 @@ import ( type exploreSourceLiteralCountingStore struct { graph.Store - allNodesCalls int - outEdgeBatchCalls int - nodeLookupBatches int + allNodesCalls int + outEdgeBatchCalls int + outSiteBatchCalls int + outEndpointBatchCalls int + nodeLookupBatches int + lastSites []graph.EdgeSourceSite + lastSiteKinds []graph.EdgeKind + lastSiteLimit int + lastEndpointIDs []string + lastEndpointKinds []graph.EdgeKind + lastEndpointLimit int + lastNodeIDs []string + siteLookup func(context.Context, []graph.EdgeSourceSite, []graph.EdgeKind, int) (graph.BoundedSiteEdgeIdentityProjection, error) + endpointLookup func(context.Context, []string, []graph.EdgeKind, int) (graph.BoundedEdgeIdentityProjection, error) + nodeLookup func(context.Context, []string) (map[string]*graph.Node, error) } func (s *exploreSourceLiteralCountingStore) AllNodes() []*graph.Node { @@ -45,21 +57,82 @@ func (s *exploreSourceLiteralCountingStore) GetOutEdgesByNodeIDs(ids []string) m func (s *exploreSourceLiteralCountingStore) GetNodesByIDs(ids []string) map[string]*graph.Node { s.nodeLookupBatches++ + s.lastNodeIDs = append([]string(nil), ids...) return s.Store.GetNodesByIDs(ids) } +func (s *exploreSourceLiteralCountingStore) GetNodesByIDsContext( + ctx context.Context, + ids []string, +) (map[string]*graph.Node, error) { + s.nodeLookupBatches++ + s.lastNodeIDs = append([]string(nil), ids...) + if s.nodeLookup != nil { + return s.nodeLookup(ctx, ids) + } + return s.Store.(exploreContextNodesReader).GetNodesByIDsContext(ctx, ids) +} + +func (s *exploreSourceLiteralCountingStore) FindOutgoingSiteEdgeIdentitiesBounded( + ctx context.Context, + sites []graph.EdgeSourceSite, + kinds []graph.EdgeKind, + limit int, +) (graph.BoundedSiteEdgeIdentityProjection, error) { + s.outSiteBatchCalls++ + s.lastSites = append([]graph.EdgeSourceSite(nil), sites...) + s.lastSiteKinds = append([]graph.EdgeKind(nil), kinds...) + s.lastSiteLimit = limit + if s.siteLookup != nil { + return s.siteLookup(ctx, sites, kinds, limit) + } + return s.Store.(graph.BoundedOutgoingSiteEdgeIdentityReader). + FindOutgoingSiteEdgeIdentitiesBounded(ctx, sites, kinds, limit) +} + +func (s *exploreSourceLiteralCountingStore) FindOutgoingEdgeIdentitiesBounded( + ctx context.Context, + ids []string, + kinds []graph.EdgeKind, + limit int, +) (graph.BoundedEdgeIdentityProjection, error) { + s.outEndpointBatchCalls++ + s.lastEndpointIDs = append([]string(nil), ids...) + s.lastEndpointKinds = append([]graph.EdgeKind(nil), kinds...) + s.lastEndpointLimit = limit + if s.endpointLookup != nil { + return s.endpointLookup(ctx, ids, kinds, limit) + } + return s.Store.(graph.BoundedOutgoingEdgeIdentityReader). + FindOutgoingEdgeIdentitiesBounded(ctx, ids, kinds, limit) +} + +func (s *exploreSourceLiteralCountingStore) FindFileNodesBounded( + ctx context.Context, + path string, + scope graph.LocalizationNodeScope, + limit int, +) (graph.BoundedNodeProjection, error) { + return s.Store.(graph.BoundedFileNodeReader).FindFileNodesBounded(ctx, path, scope, limit) +} + type exploreSourceLiteralBlockingStore struct { graph.Store started chan struct{} } -func (s *exploreSourceLiteralBlockingStore) GetFileNodesContext(ctx context.Context, _ string) []*graph.Node { +func (s *exploreSourceLiteralBlockingStore) FindFileNodesBounded( + ctx context.Context, + _ string, + _ graph.LocalizationNodeScope, + _ int, +) (graph.BoundedNodeProjection, error) { select { case s.started <- struct{}{}: default: } <-ctx.Done() - return nil + return graph.BoundedNodeProjection{}, ctx.Err() } type exploreSourceLiteralOrderedStore struct { @@ -69,13 +142,26 @@ type exploreSourceLiteralOrderedStore struct { calls []string } -func (s *exploreSourceLiteralOrderedStore) GetFileNodesContext(ctx context.Context, path string) []*graph.Node { +func (s *exploreSourceLiteralOrderedStore) FindFileNodesBounded( + ctx context.Context, + path string, + _ graph.LocalizationNodeScope, + limit int, +) (graph.BoundedNodeProjection, error) { s.calls = append(s.calls, path) if path == s.blockPath { <-ctx.Done() - return nil + return graph.BoundedNodeProjection{}, ctx.Err() + } + nodes := s.nodesByPath[path] + if len(nodes) <= limit { + return graph.BoundedNodeProjection{Nodes: nodes, Total: len(nodes)}, nil } - return s.nodesByPath[path] + return graph.BoundedNodeProjection{ + Nodes: nodes[:limit], + Total: limit + 1, + Truncated: true, + }, nil } func (s *exploreSourceLiteralOrderedStore) GetOutEdgesByNodeIDs([]string) map[string][]*graph.Edge { @@ -128,6 +214,21 @@ func sourceLiteralNode(id, name, path string, kind graph.NodeKind, start, end in } } +func requireSourceLiteralHitIdentity(t testing.TB, hits []exploreSourceLiteralHit, expected ...exploreSourceLiteralHit) { + t.Helper() + require.Len(t, hits, len(expected)) + for index := range expected { + require.Equal(t, expected[index].nodeID, hits[index].nodeID) + require.Equal(t, expected[index].rank, hits[index].rank) + require.Equal(t, expected[index].anchor, hits[index].anchor) + require.Equal(t, expected[index].ambiguous, hits[index].ambiguous) + require.Equal(t, expected[index].callee, hits[index].callee) + require.NotEmpty(t, hits[index].matchPath) + require.Positive(t, hits[index].matchLine) + require.NotEmpty(t, hits[index].literal) + } +} + func TestExploreSourceLiteralCallNameAcrossLanguages(t *testing.T) { tests := []struct { name string @@ -179,10 +280,14 @@ func TestMapExploreSourceLiteralMatchesPromotesUniqueDirectCalleeAcrossLanguages Path: path, Line: 3, Text: test.line, }}, query.QueryOptions{RepoAllow: map[string]bool{"demo": true}}) - require.Equal(t, []exploreSourceLiteralHit{{nodeID: callee.ID, rank: 0, callee: true}}, recall.hits) + requireSourceLiteralHitIdentity(t, recall.hits, exploreSourceLiteralHit{nodeID: callee.ID, rank: 0, callee: true}) require.Equal(t, callee.FilePath, recall.ownerFiles[callee.ID]) require.Zero(t, counting.allNodesCalls, "callee promotion must remain batch- and file-bounded") - require.Equal(t, 1, counting.outEdgeBatchCalls) + require.Zero(t, counting.outEdgeBatchCalls, "legacy adjacency must not be queried") + require.Equal(t, 1, counting.outSiteBatchCalls) + require.Equal(t, []graph.EdgeSourceSite{{From: owner.ID, Line: 3}}, counting.lastSites) + require.Equal(t, []graph.EdgeKind{graph.EdgeCalls}, counting.lastSiteKinds) + require.Equal(t, exploreSourceLiteralCallEdgesPerSite, counting.lastSiteLimit) require.Equal(t, 1, counting.nodeLookupBatches) }) } @@ -202,7 +307,13 @@ func TestMapExploreSourceLiteralMatchesDoesNotPromoteAmbiguousCallee(t *testing. Path: path, Line: 3, Text: `RegisterDefaultFormatter("ku");`, }}, query.QueryOptions{RepoAllow: map[string]bool{"demo": true}}) - require.Equal(t, []exploreSourceLiteralHit{{nodeID: owner.ID, rank: 0}}, recall.hits) + requireSourceLiteralHitIdentity(t, recall.hits, exploreSourceLiteralHit{nodeID: owner.ID, rank: 0}) + target := exploreTarget{ + node: owner, source: "void configure() {}", exactContent: true, + sourceLiteral: true, sourceLiteralCallee: recall.hits[0].callee, + } + require.Equal(t, localizationProvenanceContentLiteral, localizationTargetProvenance(localizationCompletion{}, target)) + require.False(t, localizationStrongSourceLiteralCallee(target), "ambiguous call edges must remain advisory") } func TestMapExploreSourceLiteralMatchesDoesNotPromoteAssignment(t *testing.T) { @@ -217,9 +328,15 @@ func TestMapExploreSourceLiteralMatchesDoesNotPromoteAssignment(t *testing.T) { Path: path, Line: 3, Text: `const locale = "ku";`, }}, query.QueryOptions{RepoAllow: map[string]bool{"demo": true}}) - require.Equal(t, []exploreSourceLiteralHit{{nodeID: owner.ID, rank: 0}}, recall.hits, - "an unrelated same-line edge must not turn an assignment into a callsite") - require.Zero(t, counting.outEdgeBatchCalls, "non-call literal hits must not query graph adjacency") + requireSourceLiteralHitIdentity(t, recall.hits, exploreSourceLiteralHit{nodeID: owner.ID, rank: 0}) + target := exploreTarget{ + node: owner, source: "void configure() {}", exactContent: true, + sourceLiteral: true, sourceLiteralCallee: recall.hits[0].callee, + } + require.Equal(t, localizationProvenanceContentLiteral, localizationTargetProvenance(localizationCompletion{}, target)) + require.False(t, localizationStrongSourceLiteralCallee(target), "assignment literals must remain advisory") + require.Zero(t, counting.outEdgeBatchCalls, "non-call literal hits must not query legacy graph adjacency") + require.Zero(t, counting.outSiteBatchCalls, "non-call literal hits must not query graph adjacency") require.Zero(t, counting.nodeLookupBatches, "non-call literal hits must not query callee nodes") } @@ -228,6 +345,20 @@ func TestExploreSourceLiteralLocalCalleeRejectsCrossRepositoryTarget(t *testing. callee := sourceLiteralNode("other/registry.cs::register", "RegisterDefaultFormatter", "other/registry.cs", graph.KindMethod, 7, 9) callee.RepoPrefix = "other" require.False(t, exploreSourceLiteralLocalCallee(owner, callee, "RegisterDefaultFormatter", query.QueryOptions{})) + + server, _ := newExploreSourceLiteralGraphServer(t, []*graph.Node{owner, callee}, []*graph.Edge{{ + From: owner.ID, To: callee.ID, Kind: graph.EdgeCalls, FilePath: owner.FilePath, Line: 3, + }}) + recall := server.mapExploreSourceLiteralMatches("ku", []trigram.Match{{ + Path: owner.FilePath, Line: 3, Text: `RegisterDefaultFormatter("ku");`, + }}, query.QueryOptions{RepoAllow: map[string]bool{"demo": true}}) + requireSourceLiteralHitIdentity(t, recall.hits, exploreSourceLiteralHit{nodeID: owner.ID, rank: 0}) + target := exploreTarget{ + node: owner, source: "void configure() {}", exactContent: true, + sourceLiteral: true, sourceLiteralCallee: recall.hits[0].callee, + } + require.Equal(t, localizationProvenanceContentLiteral, localizationTargetProvenance(localizationCompletion{}, target)) + require.False(t, localizationStrongSourceLiteralCallee(target), "cross-repository edges must remain advisory") } func TestSourceLiteralCalleeRemainsAuthorizedForRefinement(t *testing.T) { @@ -259,7 +390,7 @@ func TestMapExploreSourceLiteralMatchesFindsCSharpConstructor(t *testing.T) { Path: path, Line: 24, Text: `Register("ku", new CentralKurdishFormatter());`, }}, query.QueryOptions{RepoAllow: map[string]bool{"demo": true}}) - require.Equal(t, []exploreSourceLiteralHit{{nodeID: constructor.ID, rank: 0}}, recall.hits) + requireSourceLiteralHitIdentity(t, recall.hits, exploreSourceLiteralHit{nodeID: constructor.ID, rank: 0}) require.False(t, recall.ambiguous) } @@ -276,7 +407,7 @@ func TestMapExploreSourceLiteralMatchesFallsBackToSingleRepoUnprefixedPath(t *te Path: matchPath, Line: 24, Text: `RegisterDefaultFormatter("ku");`, }}, query.QueryOptions{RepoAllow: map[string]bool{"humanizer-1059": true}}) - require.Equal(t, []exploreSourceLiteralHit{{nodeID: constructor.ID, rank: 0}}, recall.hits) + requireSourceLiteralHitIdentity(t, recall.hits, exploreSourceLiteralHit{nodeID: constructor.ID, rank: 0}) require.False(t, recall.ambiguous) } @@ -293,10 +424,10 @@ func TestMapExploreSourceLiteralMatchesPrefersExactPathOverAlias(t *testing.T) { Path: matchPath, Line: 24, Text: `RegisterDefaultFormatter("ku");`, }}, query.QueryOptions{RepoAllow: map[string]bool{"humanizer-1059": true}}) - require.Equal(t, []exploreSourceLiteralHit{{nodeID: exact.ID, rank: 0}}, recall.hits) + requireSourceLiteralHitIdentity(t, recall.hits, exploreSourceLiteralHit{nodeID: exact.ID, rank: 0}) } -func TestMapExploreSourceLiteralMatchesQueriesExactPathsBeforeAliases(t *testing.T) { +func TestMapExploreSourceLiteralMatchesPreservesFirstSeenExactPathsBeforeAliases(t *testing.T) { exactA := sourceLiteralNode("demo/src/a.cs::Register", "RegisterA", "demo/src/a.cs", graph.KindMethod, 1, 5) exactB := sourceLiteralNode("demo/src/b.cs::Register", "RegisterB", "demo/src/b.cs", graph.KindMethod, 1, 5) store := &exploreSourceLiteralOrderedStore{ @@ -315,12 +446,13 @@ func TestMapExploreSourceLiteralMatchesQueriesExactPathsBeforeAliases(t *testing {Path: exactA.FilePath, Line: 3, Text: `Register("ku")`}, }, query.QueryOptions{RepoAllow: map[string]bool{"demo": true}}) - require.Equal(t, []string{"demo/src/a.cs", "demo/src/b.cs", "src/a.cs"}, store.calls) - require.Equal(t, []exploreSourceLiteralHit{ - {nodeID: exactB.ID, rank: 0}, - {nodeID: exactA.ID, rank: 1}, - }, recall.hits) - require.True(t, recall.ambiguous) + require.Equal(t, []string{"demo/src/b.cs", "demo/src/a.cs", "src/b.cs", "src/a.cs"}, store.calls) + requireSourceLiteralHitIdentity(t, recall.hits, + exploreSourceLiteralHit{nodeID: exactB.ID, rank: 0}, + exploreSourceLiteralHit{nodeID: exactA.ID, rank: 1}, + ) + // Two owners from a complete search are two settled sites, not noise. + require.False(t, recall.ambiguous) } func TestMapExploreSourceLiteralMatchesChoosesSmallestEnclosingSymbol(t *testing.T) { @@ -334,7 +466,7 @@ func TestMapExploreSourceLiteralMatchesChoosesSmallestEnclosingSymbol(t *testing Path: path, Line: 23, Text: `register("ku")`, }}, query.QueryOptions{RepoAllow: map[string]bool{"demo": true}}) - require.Equal(t, []exploreSourceLiteralHit{{nodeID: closure.ID, rank: 0}}, recall.hits) + requireSourceLiteralHitIdentity(t, recall.hits, exploreSourceLiteralHit{nodeID: closure.ID, rank: 0}) } func TestMapExploreSourceLiteralMatchesKeepsCommonLiteralNonTerminal(t *testing.T) { @@ -348,7 +480,9 @@ func TestMapExploreSourceLiteralMatchesKeepsCommonLiteralNonTerminal(t *testing. }, query.QueryOptions{RepoAllow: map[string]bool{"demo": true}}) require.Len(t, recall.hits, 2) - require.True(t, recall.ambiguous) + // A complete two-owner recall is settled evidence for both sites; the + // ambiguity mark is reserved for saturated or owner-cap-overflowing terms. + require.False(t, recall.ambiguous) targets := []exploreTarget{ {node: left, exactContent: true, exactContentAmbiguous: true}, {node: right, exactContent: true, exactContentAmbiguous: true}, @@ -399,7 +533,7 @@ func TestMapDiscoveredExploreSourceLiteralMatchesPreservesHitsAfterDiscoveryDead ) require.NoError(t, mappingErr) - require.Equal(t, []exploreSourceLiteralHit{{nodeID: constructor.ID, rank: 0}}, recall.hits) + requireSourceLiteralHitIdentity(t, recall.hits, exploreSourceLiteralHit{nodeID: constructor.ID, rank: 0}) require.True(t, recall.ambiguous, "deadline-truncated discovery must remain non-terminal") } @@ -554,11 +688,37 @@ func TestRetainExploreSourceLiteralOwnersDiversifiesFilesWithinCaps(t *testing.T hits, files, reason := retainExploreSourceLiteralOwners(recall) - require.Equal(t, []string{"first-a", "second", "first-b"}, []string{ + require.Equal(t, []string{"first-a", "second", "third"}, []string{ hits[0].nodeID, hits[1].nodeID, hits[2].nodeID, }) require.Equal(t, exploreSourceLiteralRecallMaxFilesPerTerm, files) - require.Equal(t, "file_cap", reason) + require.Equal(t, "owner_cap", reason) +} + +func TestRetainExploreSourceLiteralOwnersPrefersResolvedCalleeWithinFileDiverseCaps(t *testing.T) { + recall := exploreSourceLiteralRecall{ + hits: []exploreSourceLiteralHit{ + {nodeID: "owner-a", rank: 0}, + {nodeID: "owner-b", rank: 1}, + {nodeID: "other-file", rank: 2}, + {nodeID: "resolved-callee", rank: 3, callee: true}, + }, + ownerFiles: map[string]string{ + "owner-a": "src/first.go", + "owner-b": "src/first.go", + "other-file": "src/second.go", + "resolved-callee": "src/first.go", + }, + } + + hits, files, reason := retainExploreSourceLiteralOwners(recall) + + require.Equal(t, []string{"resolved-callee", "other-file", "owner-a"}, []string{ + hits[0].nodeID, hits[1].nodeID, hits[2].nodeID, + }) + require.True(t, hits[0].callee) + require.Equal(t, 2, files) + require.Equal(t, "owner_cap", reason) } func TestGatherExploreSourceLiteralRecallAggregatesCompactAnchorsAcrossLanguages(t *testing.T) { @@ -652,7 +812,8 @@ func TestGatherExploreSourceLiteralRecallAggregatesCompactAnchorsAcrossLanguages require.NotNil(t, secondCandidate) require.Equal(t, float64(1), firstCandidate.Signals[exploreSourceLiteralCoverageSignal]) require.Equal(t, float64(2), secondCandidate.Signals[exploreSourceLiteralCoverageSignal]) - require.Positive(t, firstCandidate.Signals[exploreContentRecallAmbiguousSignal]) + require.Zero(t, firstCandidate.Signals[exploreContentRecallAmbiguousSignal], + "a complete two-owner recall settles both sites") require.Zero(t, secondCandidate.Signals[exploreContentRecallAmbiguousSignal]) }) } @@ -743,7 +904,8 @@ func TestGatherExploreSourceLiteralRecallKeepsMultiAnchorOwnerUnderNearCapCompet targetCandidate := candidateByID(sourceCandidates, target.ID) require.NotNil(t, targetCandidate) require.Equal(t, float64(2), targetCandidate.Signals[exploreSourceLiteralCoverageSignal]) - require.Positive(t, targetCandidate.Signals[exploreContentRecallAmbiguousSignal]) + require.Zero(t, targetCandidate.Signals[exploreContentRecallAmbiguousSignal], + "a complete two-owner recall settles the multi-anchor owner") ordinary := []*rerank.Candidate{ sourcePreservationCandidate("semantic-0", 0, 0), sourcePreservationCandidate("semantic-1", 1, 0), @@ -758,18 +920,18 @@ func TestGatherExploreSourceLiteralRecallKeepsMultiAnchorOwnerUnderNearCapCompet } func TestGatherExploreSourceLiteralRecallRecordsTermCapDiagnostic(t *testing.T) { - // "cc" must be reported as term_cap, not as a deadline the runner caused. + // "ee" must be reported as term_cap, not as a deadline the runner caused. server := pinExploreSourceLiteralRecallBudget(&Server{logger: zap.NewNop()}) recall := server.gatherExploreSourceLiteralRecall( - context.Background(), []string{"aa", "bb", "cc"}, "demo", query.QueryOptions{}, + context.Background(), []string{"aa", "bb", "cc", "dd", "ee"}, "demo", query.QueryOptions{}, ) - require.Len(t, recall.diagnostics, 3) + require.Len(t, recall.diagnostics, 5) byLiteral := make(map[string]exploreSourceLiteralDiagnostic, len(recall.diagnostics)) for _, diagnostic := range recall.diagnostics { byLiteral[diagnostic.literal] = diagnostic } - require.Equal(t, "term_cap", byLiteral["cc"].reason) + require.Equal(t, "term_cap", byLiteral["ee"].reason) } func TestGatherExploreSourceLiteralRecallMapsParsedCSharpConstructor(t *testing.T) { @@ -832,6 +994,11 @@ func TestGatherExploreSourceLiteralRecallMapsParsedCSharpConstructor(t *testing. require.NotEmpty(t, envelope.Evidence) require.Equal(t, "RegisterDefaultFormatter", envelope.Evidence[0].Name, "invoked source evidence must lead the final localization envelope") require.Equal(t, localizationProvenanceSourceLiteralCallee, envelope.Evidence[0].Provenance) + require.NotNil(t, envelope.SourceWindow) + require.Equal(t, rel, envelope.SourceWindow.Path) + require.Equal(t, 3, envelope.SourceWindow.MatchLine) + require.Equal(t, envelope.Evidence[0].ID, envelope.SourceWindow.AnchorSymbol) + require.Contains(t, envelope.SourceWindow.Content, `RegisterDefaultFormatter("ku")`) require.Equal(t, localizationStateAnswerReady, envelope.Completion.State) require.True(t, envelope.Terminal) require.True(t, envelope.Completion.Enforceable) @@ -956,10 +1123,17 @@ func TestExploreCompactLiteralIgnoresTestMetadataAndPrefersSpecificProductionCal } } require.NotEqual(t, -1, specificRank, "construction-aligned source evidence must survive final packing") - require.Equal(t, localizationStateNeedsRefinement, envelope.Completion.State) - require.False(t, envelope.Terminal) - require.False(t, envelope.Completion.Enforceable, "ambiguous production literal sites remain advisory") - require.Contains(t, envelope.Completion.AllowedSymbols, specificID) + // Two settled production registration sites are the complete answer, not + // grounds for another recovery turn: the page terminal-claims with both + // sites presented. + require.Equal(t, localizationStateAnswerReady, envelope.Completion.State) + require.True(t, envelope.Terminal) + evidenceIDs := make([]string, 0, len(envelope.Evidence)) + for _, evidence := range envelope.Evidence { + evidenceIDs = append(evidenceIDs, evidence.ID) + } + require.Contains(t, evidenceIDs, specificID) + require.Contains(t, evidenceIDs, genericID, "both settled production sites must be presented") } func TestGatherExploreSourceLiteralRecallBoundsMappingByRequestDeadline(t *testing.T) { diff --git a/internal/mcp/explore_source_range_scope_test.go b/internal/mcp/explore_source_range_scope_test.go new file mode 100644 index 00000000..53336d82 --- /dev/null +++ b/internal/mcp/explore_source_range_scope_test.go @@ -0,0 +1,172 @@ +package mcp + +import ( + "context" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/zzet/gortex/internal/graph" + "github.com/zzet/gortex/internal/query" + "github.com/zzet/gortex/internal/search/rerank" +) + +func TestExploreSourceRangeScopedGraphPathAliases(t *testing.T) { + const ( + rawPath = "monolog/src/Monolog/Handler/FingersCrossedHandler.php" + resolvedPath = "monolog-1800/monolog/src/Monolog/Handler/FingersCrossedHandler.php" + scopedPath = "monolog-1800/src/Monolog/Handler/FingersCrossedHandler.php" + ) + scope := query.QueryOptions{RepoAllow: map[string]bool{"monolog-1800": true}} + got := exploreSourceRangeScopedGraphPathAliases(rawPath, resolvedPath, "monolog-1800", scope) + require.NotEmpty(t, got) + require.Equal(t, resolvedPath, got[0], "a resolved exact path must remain authoritative") + require.Contains(t, got, scopedPath) + require.LessOrEqual(t, len(got), exploreSourceRangeMaxAliases) + + resolvedIndex := &fileSymbolIndex{} + scopedIndex := &fileSymbolIndex{} + require.Same(t, resolvedIndex, exploreSourceRangeIndex(map[string]*fileSymbolIndex{ + resolvedPath: resolvedIndex, + scopedPath: scopedIndex, + }, got), "a synthesized scoped path must not displace an indexed resolved path") + + got = exploreSourceRangeScopedGraphPathAliases( + "spdlog/include/spdlog/details/os-inl.h", + "", + "", + query.QueryOptions{RepoAllow: map[string]bool{"spdlog-3488": true}}, + ) + require.NotEmpty(t, got) + require.Equal(t, "spdlog-3488/include/spdlog/details/os-inl.h", got[0]) + + require.Nil(t, exploreSourceRangeScopedGraphPathAliases( + rawPath, + "", + "", + query.QueryOptions{RepoAllow: map[string]bool{ + "monolog-1800": true, + "monolog-1900": true, + }}, + )) + require.Nil(t, exploreSourceRangeScopedGraphPathAliases( + "other/src/x.go", "other-1900/src/x.go", "other-1900", scope, + ), "a path resolved to a foreign repo must not suffix-fallback into the allowed repo") + unknownAliases := exploreSourceRangeScopedGraphPathAliases("other/src/x.go", "", "", scope) + require.Equal(t, []string{"monolog-1800/other/src/x.go"}, unknownAliases) + require.Nil(t, exploreSourceRangeIndex(map[string]*fileSymbolIndex{ + "monolog-1800/src/x.go": {}, + }, unknownAliases), "an unrecognized label must not be shortened into the allowed repo") + unknownResolved := exploreSourceRangeScopedGraphPathAliases( + "other/src/x.go", "monolog-1800/other/src/x.go", "monolog-1800", scope, + ) + require.Equal(t, []string{"monolog-1800/other/src/x.go"}, unknownResolved) + require.Nil(t, exploreSourceRangeIndex(map[string]*fileSymbolIndex{ + "monolog-1800/src/x.go": {}, + }, unknownResolved), "a resolved unknown label must retain only its exact spelling") + + for _, unsafe := range []string{`C:\other\src\x.go`, `\\server\share\x.go`, "../src/x.go", "other/../src/x.go", "monolog/../src/x.go"} { + require.Nil(t, exploreSourceRangeScopedGraphPathAliases(unsafe, "", "", scope), unsafe) + } + for _, traversal := range []string{"other/../src/x.go", "monolog/../src/x.go"} { + require.Nil(t, exploreSourceRangeScopedGraphPathAliases( + traversal, "monolog-1800/src/x.go", "monolog-1800", scope, + ), traversal) + } + require.Equal(t, []string{"monolog-1800/src/x.go"}, exploreSourceRangeScopedGraphPathAliases( + `C:\other\src\x.go`, "monolog-1800/src/x.go", "monolog-1800", scope, + ), "an absolute citation may retain only its resolved exact spelling") +} + +func TestPromoteExploreSourceRangeCandidatesUsesSingleScopedRepoPrefix(t *testing.T) { + const file = "monolog-1800/src/Monolog/Handler/FingersCrossedHandler.php" + owner := &graph.Node{ + ID: file + "::FingersCrossedHandler", + Name: "FingersCrossedHandler", + QualName: "Monolog.Handler.FingersCrossedHandler", + Kind: graph.KindType, + FilePath: file, + RepoPrefix: "monolog-1800", + StartLine: 1, + EndLine: 240, + } + flush := &graph.Node{ + ID: file + "::FingersCrossedHandler.flushBuffer", + Name: "flushBuffer", + QualName: "FingersCrossedHandler.flushBuffer", + Kind: graph.KindMethod, + FilePath: file, + RepoPrefix: "monolog-1800", + StartLine: 180, + EndLine: 195, + } + g := graph.New() + g.AddBatch([]*graph.Node{owner, flush}, nil) + // The request engine can be bound to a different reader than Server.graph. + // Source-range recovery must follow the same reader as the ranked request. + s := &Server{graph: graph.New(), engine: query.NewEngine(g)} + reader := s.engineFor(context.Background()).Reader() + ordinary := []*rerank.Candidate{{Node: &graph.Node{ + ID: "monolog-1800/src/Utils.php::pcreLastErrorMessage", + Name: "pcreLastErrorMessage", Kind: graph.KindMethod, + FilePath: "monolog-1800/src/Utils.php", RepoPrefix: "monolog-1800", + }}} + scope := query.QueryOptions{RepoAllow: map[string]bool{"monolog-1800": true}} + + got := s.promoteExploreSourceRangeCandidates( + context.Background(), + "Investigate monolog/src/Monolog/Handler/FingersCrossedHandler.php Lines 185–187.", + ordinary, + reader, + scope, + ) + require.Len(t, got, 2) + require.Equal(t, flush.ID, got[0].Node.ID) + require.Equal(t, float64(1), got[0].Signals[exploreSourceRangeSignal]) + require.Same(t, ordinary[0], got[1]) + + ambiguous := s.promoteExploreSourceRangeCandidates( + context.Background(), + "Investigate monolog/src/Monolog/Handler/FingersCrossedHandler.php Lines 185–187.", + ordinary, + reader, + query.QueryOptions{RepoAllow: map[string]bool{ + "monolog-1800": true, + "monolog-1900": true, + }}, + ) + require.Len(t, ambiguous, 1) + require.Same(t, ordinary[0], ambiguous[0]) + + outOfScope := s.promoteExploreSourceRangeCandidates( + context.Background(), + "Investigate monolog/src/Monolog/Handler/FingersCrossedHandler.php Lines 185–187.", + ordinary, + reader, + query.QueryOptions{RepoAllow: map[string]bool{"monolog-1900": true}}, + ) + require.Len(t, outOfScope, 1) + require.Same(t, ordinary[0], outOfScope[0]) + + canceledCtx, cancel := context.WithCancel(context.Background()) + cancel() + canceled := s.promoteExploreSourceRangeCandidates( + canceledCtx, + "Investigate monolog/src/Monolog/Handler/FingersCrossedHandler.php Lines 185–187.", + ordinary, + reader, + scope, + ) + require.Len(t, canceled, 1) + require.Same(t, ordinary[0], canceled[0]) + + noReader := s.promoteExploreSourceRangeCandidates( + context.Background(), + "Investigate monolog/src/Monolog/Handler/FingersCrossedHandler.php Lines 185–187.", + ordinary, + nil, + scope, + ) + require.Len(t, noReader, 1) + require.Same(t, ordinary[0], noReader[0]) +} diff --git a/internal/mcp/explore_syntactic_anchor.go b/internal/mcp/explore_syntactic_anchor.go index c84acf78..6a85818a 100644 --- a/internal/mcp/explore_syntactic_anchor.go +++ b/internal/mcp/explore_syntactic_anchor.go @@ -19,6 +19,9 @@ const ( exploreSyntacticAnchorMaxTerms = 3 exploreSyntacticAnchorFetch = 10 exploreSyntacticAnchorCompetitionFetch = 32 + // A qualified leaf may inspect only the first four ranked files that + // contain an exact owner declaration. + exploreQualifiedLeafMaxFiles = 4 ) // exploreSyntacticAnchor is a bounded, implementation-shaped clue copied @@ -244,8 +247,8 @@ const ( ) var ( - exploreSourceRangeLineRE = regexp.MustCompile(`(?i)\blines?\s+([0-9]{1,8})(?:\s+(?:to|-)\s+([0-9]{1,8}))?`) - exploreSourceRangeInlineRE = regexp.MustCompile(`^\s*:([0-9]{1,8})(?:-([0-9]{1,8}))?`) + exploreSourceRangeLineRE = regexp.MustCompile(`(?i)\blines?\s+([0-9]{1,8})(?:(?:\s+to\s+|\s*[-–—]\s*)([0-9]{1,8}))?`) + exploreSourceRangeInlineRE = regexp.MustCompile(`^\s*:([0-9]{1,8})(?:[-–—]([0-9]{1,8}))?`) ) // exploreSourceRangeSpecs pairs an explicit source path with the line citation @@ -318,7 +321,7 @@ func exploreSourceRangeLines(tail string) (int, int, bool) { // line falls inside one, the enclosing named method/function remains the useful // localization symbol (for example FingersCrossedHandler::flushBuffer). func exploreSourceRangeDefinitions(index *fileSymbolIndex, start, end int) []*graph.Node { - if index == nil { + if index == nil || index.saturated { return nil } if end < start { @@ -385,6 +388,104 @@ func exploreSourceRangeGraphPathAliases(graphPath string) []string { return out } +// exploreSourceRangeScopedGraphPathAliases maps a task's logical repo label to +// the sole request-scoped graph prefix without consulting the filesystem. Replay +// workspaces commonly index monolog/src/... as monolog-1800/src/.... A resolved +// exact spelling remains authoritative; the scoped spelling is only a bounded +// fallback. Multiple-repo scopes and paths resolved to a foreign repo never use +// that fallback. +func exploreSourceRangeScopedGraphPathAliases( + rawPath, resolvedGraphPath, resolvedRepoPrefix string, + scope query.QueryOptions, +) []string { + resolved := exploreSourceRangeGraphPathAliases(resolvedGraphPath) + if resolvedRepoPrefix != "" && len(scope.RepoAllow) > 0 && + !repoNarrowAdmits(scope.RepoAllow, resolvedRepoPrefix) { + return nil + } + repoPrefix := exploreSourceLiteralSingleRepoPrefix(scope) + normalizedRaw := strings.ReplaceAll(strings.TrimSpace(rawPath), "\\", "/") + if exploreSourceRangeTraversalPath(normalizedRaw) { + return nil + } + if repoPrefix == "" { + return resolved + } + if exploreSourceRangeAbsolutePath(normalizedRaw) { + if len(resolved) == 0 { + return nil + } + return resolved[:1] + } + raw := filepath.ToSlash(filepath.Clean(normalizedRaw)) + raw = strings.TrimPrefix(raw, "./") + if raw == "" || raw == "." || raw == ".." || strings.HasPrefix(raw, "../") { + return resolved + } + + relative := raw + if strings.HasPrefix(relative, repoPrefix+"/") { + relative = strings.TrimPrefix(relative, repoPrefix+"/") + } else if slash := strings.IndexByte(relative, '/'); slash > 0 && + exploreSourceRangeLogicalRepoLabel(repoPrefix, relative[:slash]) { + relative = relative[slash+1:] + } + scopedPath := repoPrefix + "/" + relative + out := make([]string, 0, exploreSourceRangeMaxAliases) + seen := make(map[string]struct{}, exploreSourceRangeMaxAliases) + add := func(path string) { + if path == "" || len(out) == exploreSourceRangeMaxAliases { + return + } + if _, duplicate := seen[path]; duplicate { + return + } + seen[path] = struct{}{} + out = append(out, path) + } + if len(resolved) > 0 { + add(resolved[0]) + } + add(scopedPath) + return out +} + +func exploreSourceRangeAbsolutePath(path string) bool { + if path == "" { + return false + } + if filepath.IsAbs(filepath.FromSlash(path)) || strings.HasPrefix(path, "//") { + return true + } + return len(path) >= 2 && path[1] == ':' && + ((path[0] >= 'a' && path[0] <= 'z') || (path[0] >= 'A' && path[0] <= 'Z')) +} + +func exploreSourceRangeTraversalPath(path string) bool { + for _, component := range strings.Split(path, "/") { + if component == ".." { + return true + } + } + return false +} + +func exploreSourceRangeLogicalRepoLabel(repoPrefix, label string) bool { + if strings.EqualFold(repoPrefix, label) { + return true + } + if len(repoPrefix) <= len(label)+1 || repoPrefix[len(label)] != '-' || + !strings.EqualFold(repoPrefix[:len(label)], label) { + return false + } + for _, r := range repoPrefix[len(label)+1:] { + if r < '0' || r > '9' { + return false + } + } + return true +} + // exploreSourceRangeIndex chooses the exact indexed path when present. If it is // absent, exactly one suffix alias may resolve; multiple suffix hits are // deliberately rejected so an explicit citation never promotes an ambiguous @@ -419,10 +520,11 @@ func (s *Server) promoteExploreSourceRangeCandidates( ctx context.Context, task string, ordinary []*rerank.Candidate, + reader graph.Reader, scope query.QueryOptions, ) []*rerank.Candidate { specs := exploreSourceRangeSpecs(task) - if s == nil || s.graph == nil || len(specs) == 0 || ctx.Err() != nil { + if s == nil || reader == nil || len(specs) == 0 || ctx.Err() != nil { return ordinary } type resolvedRange struct { @@ -434,10 +536,17 @@ func (s *Server) promoteExploreSourceRangeCandidates( seenPaths := make(map[string]struct{}, len(specs)*exploreSourceRangeMaxAliases) for _, spec := range specs { absPath, relPath, err := s.resolveFilePath(spec.File) - if err != nil { - continue + resolvedGraphPath := "" + resolvedRepoPrefix := "" + if err == nil { + resolvedGraphPath = s.resolveOverlayGraphPath(relPath, absPath) + if s.multiIndexer != nil { + resolvedRepoPrefix = matchedRepoPrefix(s.multiIndexer, relPath) + } } - graphPaths := exploreSourceRangeGraphPathAliases(s.resolveOverlayGraphPath(relPath, absPath)) + graphPaths := exploreSourceRangeScopedGraphPathAliases( + spec.File, resolvedGraphPath, resolvedRepoPrefix, scope, + ) if len(graphPaths) == 0 { continue } @@ -452,10 +561,16 @@ func (s *Server) promoteExploreSourceRangeCandidates( if len(resolved) == 0 { return ordinary } - indexes := s.buildFileSymbolIndexForOrderedPathsContext(ctx, orderedPaths) + indexes := s.buildFileSymbolIndexForOrderedPathsScopedReaderContext(ctx, reader, orderedPaths, scope) + if ctx.Err() != nil { + return ordinary + } exactNodes := make([]*graph.Node, 0, len(resolved)) seenNodes := make(map[string]struct{}, len(resolved)) for _, item := range resolved { + if ctx.Err() != nil { + return ordinary + } index := exploreSourceRangeIndex(indexes, item.graphPaths) for _, node := range exploreSourceRangeDefinitions(index, item.spec.StartLine, item.spec.EndLine) { if node == nil || !exploreLocalizableKind(node.Kind) || !scope.ScopeAllows(node) || @@ -475,7 +590,7 @@ func (s *Server) promoteExploreSourceRangeCandidates( break } } - if len(exactNodes) == 0 { + if ctx.Err() != nil || len(exactNodes) == 0 { return ordinary } ordinaryByID := make(map[string]*rerank.Candidate, len(ordinary)) @@ -506,6 +621,9 @@ func (s *Server) promoteExploreSourceRangeCandidates( } out = append(out, candidate) } + if ctx.Err() != nil { + return ordinary + } return out } @@ -799,16 +917,21 @@ func exploreSyntacticAnchorReusesProtected( // exploreExactQualifiedAnchorCandidate resolves the parser's exact member name // before the bounded lexical lane. Graph nodes commonly store only the terminal // method Name while the ID tail carries Owner.member, so query both exact forms. -// This lookup is restricted to explicit Owner::member task syntax and still -// applies session, query, kind, and diversity filters through the ordinary -// anchor selector. +// If those exact point lookups cannot select a qualified node, one final lane +// may inspect up to four already-ranked files that contain an exact owner +// declaration. It never performs a substring or corpus-wide scan. func (s *Server) exploreExactQualifiedAnchorCandidate( ctx context.Context, anchor exploreSyntacticAnchor, + ordinary []*rerank.Candidate, scope query.QueryOptions, usedIDs, usedFiles map[string]struct{}, ) *rerank.Candidate { - if s == nil || s.graph == nil || anchor.qualifiedName == "" || ctx.Err() != nil { + if s == nil || anchor.qualifiedName == "" || ctx.Err() != nil { + return nil + } + reader := s.readerFor(ctx) + if reader == nil { return nil } names := []string{anchor.qualifiedName} @@ -818,7 +941,21 @@ func (s *Server) exploreExactQualifiedAnchorCandidate( candidates := make([]*rerank.Candidate, 0, exploreSyntacticAnchorFetch) seen := make(map[string]struct{}, exploreSyntacticAnchorFetch) for _, name := range names { - for _, node := range s.graph.FindNodesByName(name) { + remaining := exploreSyntacticAnchorFetch - len(candidates) + if remaining <= 0 { + break + } + page, ok := boundedLocalizationExactName( + ctx, + reader, + name, + s.localizationNodeScope(ctx, scope, graph.KindFunction, graph.KindMethod, graph.KindType, graph.KindInterface, graph.KindMacro), + remaining, + ) + if !ok { + return nil + } + for _, node := range page.Nodes { if node == nil || !s.nodeInSessionScope(ctx, node) { continue } @@ -835,7 +972,10 @@ func (s *Server) exploreExactQualifiedAnchorCandidate( break } } - return exploreSyntacticAnchorCandidate(anchor, candidates, scope, usedIDs, usedFiles) + if candidate := exploreSyntacticAnchorCandidate(anchor, candidates, scope, usedIDs, usedFiles); candidate != nil { + return candidate + } + return s.exploreQualifiedLeafCandidate(ctx, reader, anchor, ordinary, scope, usedIDs, usedFiles) } // gatherExploreSyntacticAnchorCandidates performs a tiny lexical retrieval for @@ -879,13 +1019,14 @@ func exploreAnchorPoolCandidate( return candidate, false } -func (s *Server) gatherExploreSyntacticAnchorCandidates( +func (s *Server) gatherExploreSyntacticAnchorCandidatesCollecting( ctx context.Context, task string, ordinary []*rerank.Candidate, eng *query.Engine, scope query.QueryOptions, rctx *rerank.Context, + collector *localizationSourceWindowHitCollector, ) ([]*rerank.Candidate, map[int]string) { if s == nil || s.graph == nil || eng == nil || ctx.Err() != nil { return nil, nil @@ -927,7 +1068,7 @@ func (s *Server) gatherExploreSyntacticAnchorCandidates( protected[index] = reused continue } - if exactCandidate := s.exploreExactAnchorCandidate(ctx, anchor, scope, usedIDs, usedFiles); exactCandidate != nil { + if exactCandidate := s.exploreExactAnchorCandidate(ctx, anchor, ordinary, scope, usedIDs, usedFiles); exactCandidate != nil { protectedCandidate, pooled := exploreAnchorPoolCandidate(ordinary, exactCandidate) if !pooled { additions = append(additions, protectedCandidate) @@ -1003,7 +1144,7 @@ func (s *Server) gatherExploreSyntacticAnchorCandidates( for _, hit := range recall.hits { ids = append(ids, hit.nodeID) } - nodes := s.graph.GetNodesByIDs(ids) + nodes := s.readerFor(ctx).GetNodesByIDs(ids) for localIndex, anchorIndex := range missed { var fallback *graph.Node var selected *graph.Node @@ -1035,6 +1176,12 @@ func (s *Server) gatherExploreSyntacticAnchorCandidates( if selected == nil { continue } + for _, hit := range recall.hits { + if hit.anchor == localIndex && hit.nodeID == selected.ID && hit.rank == selectedRank { + collector.add(hit) + break + } + } sourceRank := 1.0 if selectedRank > 0 { sourceRank = 1 / float64(selectedRank+1) diff --git a/internal/mcp/explore_task_terminal.go b/internal/mcp/explore_task_terminal.go new file mode 100644 index 00000000..1b285008 --- /dev/null +++ b/internal/mcp/explore_task_terminal.go @@ -0,0 +1,206 @@ +package mcp + +import ( + "context" + "encoding/json" + "fmt" + "strings" + + "github.com/zzet/gortex/internal/graph" +) + +const ( + exploreTaskOutlineBudgetShare = 4 + exploreTaskOutlineHeading = "## File outlines" + exploreTaskMinimumOutlineTokens = 64 + exploreTaskSectionSeparatorTokens = 2 +) + +type exploreTaskOutlineProvider func([]exploreTarget) *localizationPageOutline + +// newExploreTaskPageOutlineProvider returns the same bounded declaration index +// used by localize pages. It captures the reader selected for this request, so +// session overlays cannot fall back to the server's base engine. Enumeration is +// deferred until the renderer proves a useful outline can fit. +func newExploreTaskPageOutlineProvider( + ctx context.Context, + reader graph.Reader, + task string, + scope graph.LocalizationNodeScope, +) exploreTaskOutlineProvider { + if reader == nil { + return nil + } + terms := exploreTerminalTerms(task) + return func(targets []exploreTarget) *localizationPageOutline { + declarations := newLocalizationFileDeclarationCache(ctx, reader, scope) + provider := localizationPageOutlineProvider(nil, targets, terms, declarations.outlineDefinitions) + if provider == nil { + return nil + } + return provider() + } +} + +// renderExploreTask appends task-mode terminal guidance and, when the remaining +// budget can carry a useful index, lazily loads bounded file outlines. Ordinary +// task mode deliberately retains no authorization state. +func (s *Server) renderExploreTask( + task string, + targets []exploreTarget, + budget int, + outlineProvider exploreTaskOutlineProvider, +) string { + completion := renderExploreTaskCompletion() + completionTokens := estimateTokens(completion) + fullBaseBudget := max(budget-completionTokens-exploreTaskSectionSeparatorTokens, 0) + renderWithoutOutline := func() string { + base := s.renderExplore(task, targets, fullBaseBudget) + return joinExploreTaskSections(base, "", completion) + } + if outlineProvider == nil || budget <= completionTokens+exploreTaskMinimumOutlineTokens { + return renderWithoutOutline() + } + + // Reserve source-packing space for an outline, never candidate space. + // renderExplore always keeps every ranked location and signature even when + // those mandatory rows exceed its approximate source budget. + outlineReserve := min( + budget/exploreTaskOutlineBudgetShare, + budget-completionTokens-exploreTaskSectionSeparatorTokens, + ) + if outlineReserve < exploreTaskMinimumOutlineTokens { + return renderWithoutOutline() + } + baseBudget := max(fullBaseBudget-outlineReserve, 0) + base := s.renderExplore(task, targets, baseBudget) + withoutOutline := joinExploreTaskSections(base, "", completion) + remaining := budget - estimateTokens(withoutOutline) + if remaining < exploreTaskMinimumOutlineTokens { + return renderWithoutOutline() + } + + outlineBudget := min(remaining, outlineReserve) + outlines := renderExploreTaskOutlines(outlineProvider(targets), outlineBudget) + withOutline := joinExploreTaskSections(base, outlines, completion) + if outlines == "" || estimateTokens(withOutline) > budget { + // The optional index did not use its allowance. Re-render with the full + // source budget so a nil or over-budget outline cannot make task mode + // poorer than the same request with outlines disabled. + return renderWithoutOutline() + } + return withOutline +} + +func joinExploreTaskSections(base, outlines, completion string) string { + var b strings.Builder + b.Grow(len(base) + len(outlines) + len(completion) + 3) + b.WriteString(base) + if base != "" && !strings.HasSuffix(base, "\n") { + b.WriteByte('\n') + } + if outlines != "" { + b.WriteByte('\n') + b.WriteString(outlines) + } + if completion != "" { + b.WriteByte('\n') + b.WriteString(completion) + } + return b.String() +} + +func renderExploreTaskCompletion() string { + completion := localizationCompletion{ + State: localizationStateLocalized, + Scope: "task", + RequiredAction: "continue_task", + Instruction: "Continue the requested diagnosis or implementation from the ranked evidence and file outlines above. Editing, navigation, building, and testing remain available.", + FinalResponse: "Use the ranked evidence and file outlines above to continue the requested task.", + AllowedToolCalls: 0, + ContractVersion: localizationTerminalContractV2, + } + contract := localizationContractFor(completion) + encoded, err := json.MarshalIndent(contract, "", " ") + if err != nil { + return "" + } + return "## Completion\n```json\n" + string(encoded) + "\n```\n" +} + +// renderExploreTaskOutlines applies the shared rank-aware relief policy until +// the declaration index fits its fixed allowance. Ranked rows and source bodies +// are never displaced, and the original cached outline is never mutated. +func renderExploreTaskOutlines(page *localizationPageOutline, tokenBudget int) string { + if page == nil || page.empty() || tokenBudget <= 0 { + return "" + } + page = page.clone() + for !page.empty() { + text := formatExploreTaskOutlines(page) + if estimateTokens(text) <= tokenBudget { + return text + } + before := exploreTaskOutlineWeight(page) + page.relieve() + if exploreTaskOutlineWeight(page) >= before { + return "" + } + } + return "" +} + +func formatExploreTaskOutlines(page *localizationPageOutline) string { + if page == nil || page.empty() { + return "" + } + var b strings.Builder + b.WriteString(exploreTaskOutlineHeading + "\n") + appendOutline := func(outline *localizationFileOutline) { + if outline == nil { + return + } + if outline.Truncated { + fmt.Fprintf(&b, "\n### %s — at least %d declaration(s)", outline.File, outline.Declared) + } else { + fmt.Fprintf(&b, "\n### %s — %d declaration(s)", outline.File, outline.Declared) + } + if outline.Elided > 0 { + if outline.Truncated { + fmt.Fprintf(&b, ", at least %d elided", outline.Elided) + } else { + fmt.Fprintf(&b, ", %d elided", outline.Elided) + } + } + b.WriteByte('\n') + for _, row := range outline.Rows { + name := truncateOneLine(row.Name, localizationMaxNameRunes) + if row.Kind != "" { + fmt.Fprintf(&b, "- %d: %s [%s]\n", row.Line, name, row.Kind) + } else { + fmt.Fprintf(&b, "- %d: %s\n", row.Line, name) + } + } + } + appendOutline(page.Leading) + for _, outline := range page.Others { + appendOutline(outline) + } + return b.String() +} + +func exploreTaskOutlineWeight(page *localizationPageOutline) int { + if page == nil { + return 0 + } + weight := 0 + if page.Leading != nil { + weight += 1 + len(page.Leading.Rows) + } + for _, outline := range page.Others { + if outline != nil { + weight += 1 + len(outline.Rows) + } + } + return weight +} diff --git a/internal/mcp/explore_task_terminal_test.go b/internal/mcp/explore_task_terminal_test.go new file mode 100644 index 00000000..2ffc8f25 --- /dev/null +++ b/internal/mcp/explore_task_terminal_test.go @@ -0,0 +1,298 @@ +package mcp + +import ( + "context" + "fmt" + "strings" + "testing" + + "github.com/zzet/gortex/internal/graph" +) + +func TestRenderExploreTaskAddsCompletionWithinBudget(t *testing.T) { + targets := exploreTestTargets() + got := (&Server{}).renderExploreTask("retry backoff", targets, 1600, nil) + + for _, want := range []string{ + "EXPLORE — retry backoff", + "## Completion", + `"state": "localized"`, + `"scope": "task"`, + `"required_action": "continue_task"`, + `"allowed_tool_calls": 0`, + `"terminal": false`, + "Editing, navigation, building, and testing remain available.", + } { + if !strings.Contains(got, want) { + t.Fatalf("task completion missing %q:\n%s", want, got) + } + } + if strings.Contains(got, localizationAnswerReadyInstruction) || strings.Contains(got, `"state": "answer_ready"`) { + t.Fatalf("ordinary task mode incorrectly terminalized navigation:\n%s", got) + } + if used := estimateTokens(got); used > 1600 { + t.Fatalf("task response used %d tokens, budget 1600", used) + } +} + +func TestRenderExploreTaskCompletionLeavesLocalizeTerminalContractUnchanged(t *testing.T) { + contract := localizationContractFor(newLocalizationCompletion(true, "")) + if !contract.Terminal || contract.Completion.State != localizationStateAnswerReady || + contract.Completion.RequiredAction != "respond" || contract.Completion.Instruction != localizationAnswerReadyInstruction { + t.Fatalf("localize terminal contract changed: %#v", contract) + } +} + +func TestRenderExploreTaskNilOutlineDoesNotDisplaceRankedTargets(t *testing.T) { + targets := make([]exploreTarget, 0, exploreDefaultMaxSymbols) + for index := 0; index < exploreDefaultMaxSymbols; index++ { + name := fmt.Sprintf("Candidate%02d", index) + targets = append(targets, exploreTarget{node: &graph.Node{ + ID: "candidate.go::" + name, + Name: name, + Kind: graph.KindFunction, + FilePath: "candidate.go", + StartLine: index + 1, + }}) + } + withoutProvider := (&Server{}).renderExploreTask("retry policy", targets, exploreDefaultBudgetTokens, nil) + calls := 0 + withNilOutline := (&Server{}).renderExploreTask("retry policy", targets, exploreDefaultBudgetTokens, func(actual []exploreTarget) *localizationPageOutline { + calls++ + if len(actual) != len(targets) { + t.Fatalf("outline provider received %d targets, want %d", len(actual), len(targets)) + } + return nil + }) + + if calls != 1 { + t.Fatalf("outline provider called %d times, want 1", calls) + } + if withNilOutline != withoutProvider { + t.Fatalf("nil optional outline changed ranked task response:\n--- without provider ---\n%s\n--- with provider ---\n%s", withoutProvider, withNilOutline) + } + for _, target := range targets { + if !strings.Contains(withNilOutline, target.node.ID) { + t.Fatalf("ranked target %q was displaced:\n%s", target.node.ID, withNilOutline) + } + } +} + +func TestRenderExploreTaskLazilyAddsBoundedTopFileOutline(t *testing.T) { + targets := exploreTestTargets() + nodes := []*graph.Node{ + targets[0].node, + targets[1].node, + {ID: "retry.go::RetryPolicy", Name: "RetryPolicy", QualName: "RetryPolicy", Kind: graph.KindType, FilePath: "retry.go", StartLine: 2}, + } + reads := 0 + provider := exploreTaskOutlineProvider(func(actualTargets []exploreTarget) *localizationPageOutline { + outline := localizationPageOutlineProvider(nil, actualTargets, exploreTerminalTerms("retry policy"), func(file string) []*graph.Node { + reads++ + if file != "retry.go" { + t.Fatalf("enumerated unexpected file %q", file) + } + return nodes + }) + return outline() + }) + got := (&Server{}).renderExploreTask("retry policy", targets, exploreMaxBudgetTokens, provider) + + if reads != 1 { + t.Fatalf("file declarations read %d times, want 1", reads) + } + for _, want := range []string{ + exploreTaskOutlineHeading, + "### retry.go — 3 declaration(s)", + "- 2: RetryPolicy [t]", + "- 6: Backoff [f]", + "- 11: DoWithRetry [f]", + } { + if !strings.Contains(got, want) { + t.Fatalf("task outline missing %q:\n%s", want, got) + } + } + if used := estimateTokens(got); used > exploreMaxBudgetTokens { + t.Fatalf("task response used %d tokens, budget %d", used, exploreMaxBudgetTokens) + } +} + +func TestRenderExploreTaskDoesNotLoadOutlineWithoutUsefulResidual(t *testing.T) { + calls := 0 + provider := exploreTaskOutlineProvider(func([]exploreTarget) *localizationPageOutline { + calls++ + return nil + }) + budget := estimateTokens(renderExploreTaskCompletion()) + exploreTaskMinimumOutlineTokens - 1 + _ = (&Server{}).renderExploreTask("retry policy", exploreTestTargets(), budget, provider) + if calls != 0 { + t.Fatalf("outline provider called %d times without useful residual", calls) + } +} + +func TestRenderExploreTaskHonorsClampedBudgets(t *testing.T) { + for _, budget := range []int{exploreMinBudgetTokens, exploreDefaultBudgetTokens, exploreMaxBudgetTokens} { + t.Run(strings.Repeat("b", budget/1000), func(t *testing.T) { + got := (&Server{}).renderExploreTask("retry policy", exploreTestTargets(), budget, nil) + if used := estimateTokens(got); used > budget { + t.Fatalf("task response used %d tokens, budget %d", used, budget) + } + }) + } +} + +func TestExploreTaskOutlineProviderUsesSelectedReader(t *testing.T) { + targets := exploreTestTargets() + selected := &localizationDeclarationSpyReader{ + files: map[string][]*graph.Node{"retry.go": {targets[0].node}}, + fileCalls: make(map[string]int), + } + provider := newExploreTaskPageOutlineProvider( + context.Background(), selected, "retry policy", + graph.LocalizationNodeScope{RepoAllow: map[string]bool{"repo": true}}, + ) + if provider == nil || provider(targets) == nil { + t.Fatal("selected reader did not produce an outline") + } + if selected.fileCalls["retry.go"] != 1 { + t.Fatalf("selected reader called %d times, want 1", selected.fileCalls["retry.go"]) + } + if len(selected.calls) != 1 || selected.calls[0].limit != localizationFileNodeLimit || + !selected.calls[0].scope.RepoAllow["repo"] || !selected.calls[0].scope.ExcludeKinds[graph.KindParam] { + t.Fatalf("selected reader call = %#v, want 1024-node request scope plus declaration exclusions", selected.calls) + } +} + +func TestExploreTaskOutlineProviderPreservesBoundedDeclarationCounts(t *testing.T) { + const file = "generated.go" + for _, test := range []struct { + name string + declared int + wantRows int + wantDeclared int + wantTruncated bool + wantRendered string + }{ + { + name: "complete retained top file", declared: 140, wantRows: 140, + wantDeclared: 140, wantRendered: "140 declaration(s)", + }, + { + name: "saturated top file lower bound", declared: localizationFileNodeLimit + 1, + wantRows: localizationFileNodeLimit, wantDeclared: localizationFileNodeLimit + 1, + wantTruncated: true, wantRendered: "at least 1025 declaration(s), at least 1 elided", + }, + } { + t.Run(test.name, func(t *testing.T) { + nodes := make([]*graph.Node, 0, test.declared) + for index := 0; index < test.declared; index++ { + name := fmt.Sprintf("Declaration%04d", index) + nodes = append(nodes, &graph.Node{ + ID: file + "::" + name, Name: name, Kind: graph.KindFunction, + FilePath: file, StartLine: index + 1, + }) + } + selected := &localizationDeclarationSpyReader{ + files: map[string][]*graph.Node{file: nodes}, fileCalls: make(map[string]int), + } + provider := newExploreTaskPageOutlineProvider( + context.Background(), selected, "generated declarations", graph.LocalizationNodeScope{}, + ) + page := provider([]exploreTarget{{node: nodes[0]}}) + + if page == nil || page.Leading == nil { + t.Fatal("task provider did not produce a leading outline") + } + if page.Leading.Declared != test.wantDeclared || page.Leading.Truncated != test.wantTruncated || + len(page.Leading.Rows) != test.wantRows { + t.Fatalf("outline = %#v, want rows=%d declared=%d truncated=%v", + page.Leading, test.wantRows, test.wantDeclared, test.wantTruncated) + } + wantElided := test.wantDeclared - test.wantRows + if page.Leading.Elided != wantElided { + t.Fatalf("elided = %d, want %d", page.Leading.Elided, wantElided) + } + if selected.fileCalls[file] != 1 || len(selected.calls) != 1 || + selected.calls[0].limit != localizationFileNodeLimit { + t.Fatalf("selected reader calls = %#v, want one 1024-node read", selected.calls) + } + rendered := formatExploreTaskOutlines(page) + if !strings.Contains(rendered, test.wantRendered) { + t.Fatalf("task outline did not render the truthful count %q:\n%s", test.wantRendered, rendered) + } + if !test.wantTruncated && strings.Contains(rendered, "at least") { + t.Fatalf("complete task outline presented a lower bound:\n%s", rendered) + } + }) + } +} + +func TestExploreTaskAndStructuredOutlinesShareRankAwareFetchPlan(t *testing.T) { + files := make(map[string][]*graph.Node, localizationOutlineFileCap) + targets := make([]exploreTarget, 0, localizationOutlineFileCap) + for rank := 0; rank < localizationOutlineFileCap; rank++ { + file := fmt.Sprintf("repo/file-%d.go", rank) + ranked := &graph.Node{ + ID: file + "::ranked", Name: "ranked", Kind: graph.KindFunction, + FilePath: file, StartLine: 1, + } + files[file] = []*graph.Node{ + ranked, + {ID: file + "::sibling", Name: "sibling", Kind: graph.KindFunction, FilePath: file, StartLine: 2}, + } + targets = append(targets, exploreTarget{node: ranked}) + } + newReader := func() *localizationDeclarationSpyReader { + return &localizationDeclarationSpyReader{files: files, fileCalls: make(map[string]int)} + } + + taskReader := newReader() + taskProvider := newExploreTaskPageOutlineProvider( + context.Background(), taskReader, "ranked declarations", graph.LocalizationNodeScope{}, + ) + taskPage := taskProvider(targets) + + structuredReader := newReader() + structuredCache := newLocalizationFileDeclarationCache( + context.Background(), structuredReader, graph.LocalizationNodeScope{}, + ) + structuredProvider := boundedLocalizationPageOutlineProvider( + nil, targets, exploreTerminalTerms("ranked declarations"), structuredCache.outlineDefinitions, + ) + structuredPage := structuredProvider() + if taskPage == nil || structuredPage == nil { + t.Fatalf("task/structured pages = %#v / %#v, want both", taskPage, structuredPage) + } + if len(taskReader.calls) != localizationOutlineFileCap || len(structuredReader.calls) != localizationOutlineFileCap { + t.Fatalf("task/structured calls = %d/%d, want %d each", + len(taskReader.calls), len(structuredReader.calls), localizationOutlineFileCap) + } + for rank := 0; rank < localizationOutlineFileCap; rank++ { + want := localizationOutlineFileFetchLimit(rank) + taskCall, structuredCall := taskReader.calls[rank], structuredReader.calls[rank] + if taskCall.file != structuredCall.file || taskCall.limit != want || structuredCall.limit != want { + t.Fatalf("rank %d task/structured calls = %#v / %#v, want identical limit %d", + rank, taskCall, structuredCall, want) + } + } +} + +func TestRenderExploreTaskOutlineShrinksCloneOnly(t *testing.T) { + rows := make([]localizationOutlineRow, 0, localizationOutlineRowCap) + for index := 0; index < localizationOutlineRowCap; index++ { + rows = append(rows, localizationOutlineRow{Name: strings.Repeat("LongDeclaration", 8), Line: index + 1, Kind: "f"}) + } + outline := &localizationFileOutline{ + File: "wide.go", Declared: len(rows), Rows: rows, all: rows, + } + page := &localizationPageOutline{Leading: outline} + originalRows := len(page.Leading.Rows) + got := renderExploreTaskOutlines(page, 64) + + if got != "" && estimateTokens(got) > 64 { + t.Fatalf("outline exceeded allowance: %d tokens", estimateTokens(got)) + } + if len(page.Leading.Rows) != originalRows { + t.Fatalf("cached outline mutated: got %d rows want %d", len(page.Leading.Rows), originalRows) + } +} diff --git a/internal/mcp/explore_typed_anchor_adjacency_bounds_test.go b/internal/mcp/explore_typed_anchor_adjacency_bounds_test.go new file mode 100644 index 00000000..475fae71 --- /dev/null +++ b/internal/mcp/explore_typed_anchor_adjacency_bounds_test.go @@ -0,0 +1,659 @@ +package mcp + +import ( + "context" + "errors" + "fmt" + "testing" + + "github.com/zzet/gortex/internal/graph" +) + +func exploreTypedAnchorFixtureProjects(fixture exploreTypedAnchorFixture) bool { + // Boundary cases isolate the intended stage from the fixture's unrelated + // same-name wrong-type route. + fixture.reader.in["rs-owner"] = removeExploreTypedAnchorTestEdge( + fixture.reader.in["rs-owner"], "rs-wrong-consumer", "rs-owner", graph.EdgeMemberOf, + ) + fixture.reader.out["rs-wrong-consumer"] = removeExploreTypedAnchorTestEdge( + fixture.reader.out["rs-wrong-consumer"], "rs-wrong-consumer", "rs-owner", graph.EdgeMemberOf, + ) + got := projectExploreTypedAnchorCandidates( + context.Background(), fixture.task, fixture.candidates, fixture.reader, + exploreTypedAnchorTestScope(), len(fixture.candidates), fixture.protected, "", + ) + return exploreTypedAnchorCandidateByID(got, fixture.member.ID) != nil +} + +func addExploreTypedAnchorIrrelevantEdges( + fixture *exploreTypedAnchorFixture, + from string, + toPrefix string, + kind graph.EdgeKind, +) { + for index := 0; index < 300; index++ { + fixture.reader.addEdge(from, fmt.Sprintf("%s-%03d", toPrefix, index), kind) + } +} + +func addExploreTypedAnchorIrrelevantIncomingEdges( + fixture *exploreTypedAnchorFixture, + fromPrefix string, + to string, + kind graph.EdgeKind, +) { + for index := 0; index < 300; index++ { + fixture.reader.addEdge(fmt.Sprintf("%s-%03d", fromPrefix, index), to, kind) + } +} + +func TestProjectExploreTypedAnchorCandidatesExactStageLimits(t *testing.T) { + t.Run("field relations 8 and 9", func(t *testing.T) { + for _, test := range []struct { + name string + total int + want bool + }{ + {name: "exact", total: exploreTypedAnchorFieldRelationLimit, want: true}, + {name: "sentinel", total: exploreTypedAnchorFieldRelationLimit + 1}, + } { + t.Run(test.name, func(t *testing.T) { + fixture := newExploreTypedAnchorFixture( + "rs", "--replace causes duplicate output", "replacer", "Replacer", "replace", "replace_all", true, + ) + for index := 2; index < test.total; index++ { + node := exploreTypedAnchorNode(fmt.Sprintf("field-type-%02d", index), fmt.Sprintf("Other%d", index), "types.rs", "rs", graph.KindType) + fixture.reader.nodes[node.ID] = node + fixture.reader.addEdge(fixture.field.ID, node.ID, graph.EdgeTypedAs) + } + addExploreTypedAnchorIrrelevantEdges(&fixture, fixture.field.ID, "field-noise", graph.EdgeCalls) + if got := exploreTypedAnchorFixtureProjects(fixture); got != test.want { + t.Fatalf("projection = %v, want %v at %d field relations", got, test.want, test.total) + } + }) + } + }) + + t.Run("direct consumers 12 and 13", func(t *testing.T) { + for _, test := range []struct { + name string + total int + want bool + }{ + {name: "exact", total: exploreTypedAnchorConsumerLimit, want: true}, + {name: "sentinel", total: exploreTypedAnchorConsumerLimit + 1}, + } { + t.Run(test.name, func(t *testing.T) { + fixture := newExploreTypedAnchorFixture( + "rs", "--replace causes duplicate output", "replacer", "Replacer", "replace", "replace_all", false, + ) + for index := 1; index < test.total; index++ { + node := exploreTypedAnchorNode(fmt.Sprintf("direct-consumer-%02d", index), "noise", "sink.rs", "rs", graph.KindMethod) + fixture.reader.nodes[node.ID] = node + fixture.reader.addEdge(node.ID, "rs-owner", graph.EdgeMemberOf) + fixture.reader.addEdge(node.ID, fixture.field.ID, graph.EdgeReads) + } + addExploreTypedAnchorIrrelevantIncomingEdges(&fixture, "irrelevant-reader", fixture.field.ID, graph.EdgeCalls) + if got := exploreTypedAnchorFixtureProjects(fixture); got != test.want { + t.Fatalf("projection = %v, want %v at %d direct consumers", got, test.want, test.total) + } + }) + } + }) + + t.Run("owner members 48 and 49", func(t *testing.T) { + for _, test := range []struct { + name string + total int + want bool + }{ + {name: "exact", total: exploreTypedAnchorOwnerMemberLimit, want: true}, + {name: "sentinel", total: exploreTypedAnchorOwnerMemberLimit + 1}, + } { + t.Run(test.name, func(t *testing.T) { + fixture := newExploreTypedAnchorFixture( + "rs", "--replace causes duplicate output", "replacer", "Replacer", "replace", "replace_all", false, + ) + for index := 2; index < test.total; index++ { + node := exploreTypedAnchorNode(fmt.Sprintf("owner-member-%02d", index), "noise", "sink.rs", "rs", graph.KindMethod) + fixture.reader.nodes[node.ID] = node + fixture.reader.addEdge(node.ID, "rs-owner", graph.EdgeMemberOf) + } + addExploreTypedAnchorIrrelevantIncomingEdges(&fixture, "irrelevant-owner-reader", "rs-owner", graph.EdgeReads) + if got := exploreTypedAnchorFixtureProjects(fixture); got != test.want { + t.Fatalf("projection = %v, want %v at %d owner members", got, test.want, test.total) + } + }) + } + }) + + t.Run("consumer calls 64 and 65", func(t *testing.T) { + for _, test := range []struct { + name string + total int + want bool + }{ + {name: "exact", total: exploreTypedAnchorCallLimit, want: true}, + {name: "sentinel", total: exploreTypedAnchorCallLimit + 1}, + } { + t.Run(test.name, func(t *testing.T) { + fixture := newExploreTypedAnchorFixture( + "rs", "--replace causes duplicate output", "replacer", "Replacer", "replace", "replace_all", false, + ) + for index := 1; index < test.total; index++ { + node := exploreTypedAnchorNode(fmt.Sprintf("called-member-%02d", index), "noise", "codec.rs", "rs", graph.KindMethod) + fixture.reader.nodes[node.ID] = node + fixture.reader.addEdge(fixture.consumer.ID, node.ID, graph.EdgeCalls) + fixture.reader.addEdge(node.ID, "rs-target-type", graph.EdgeMemberOf) + } + addExploreTypedAnchorIrrelevantEdges(&fixture, fixture.consumer.ID, "call-noise", graph.EdgeReads) + if got := exploreTypedAnchorFixtureProjects(fixture); got != test.want { + t.Fatalf("projection = %v, want %v at %d calls", got, test.want, test.total) + } + }) + } + }) + + t.Run("member owners 4 and 5", func(t *testing.T) { + for _, test := range []struct { + name string + total int + want bool + }{ + {name: "exact", total: exploreTypedAnchorMemberOwnerLimit, want: true}, + {name: "sentinel", total: exploreTypedAnchorMemberOwnerLimit + 1}, + } { + t.Run(test.name, func(t *testing.T) { + fixture := newExploreTypedAnchorFixture( + "rs", "--replace causes duplicate output", "replacer", "Replacer", "replace", "replace_all", false, + ) + for index := 1; index < test.total; index++ { + node := exploreTypedAnchorNode(fmt.Sprintf("member-owner-%02d", index), fmt.Sprintf("Other%d", index), "types.rs", "rs", graph.KindType) + fixture.reader.nodes[node.ID] = node + fixture.reader.addEdge(fixture.member.ID, node.ID, graph.EdgeMemberOf) + } + addExploreTypedAnchorIrrelevantEdges(&fixture, fixture.member.ID, "member-noise", graph.EdgeCalls) + if got := exploreTypedAnchorFixtureProjects(fixture); got != test.want { + t.Fatalf("projection = %v, want %v at %d member owners", got, test.want, test.total) + } + }) + } + }) +} + +func TestProjectExploreTypedAnchorCandidatesAggregateCallLimit(t *testing.T) { + for _, test := range []struct { + name string + altCalls int + want bool + }{ + {name: "exact 64", altCalls: 32, want: true}, + {name: "sentinel 65", altCalls: 33}, + } { + t.Run(test.name, func(t *testing.T) { + fixture := newExploreTypedAnchorFixture( + "rs", "--replace causes duplicate output", "replacer", "Replacer", "replace", "replace_all", false, + ) + alt := exploreTypedAnchorNode("aggregate-consumer", "noise", "sink.rs", "rs", graph.KindMethod) + fixture.reader.nodes[alt.ID] = alt + fixture.reader.addEdge(alt.ID, "rs-owner", graph.EdgeMemberOf) + fixture.reader.addEdge(alt.ID, fixture.field.ID, graph.EdgeReads) + + addCalls := func(from, prefix string, count int) { + for index := 0; index < count; index++ { + node := exploreTypedAnchorNode(fmt.Sprintf("%s-%02d", prefix, index), "noise", "codec.rs", "rs", graph.KindMethod) + fixture.reader.nodes[node.ID] = node + fixture.reader.addEdge(from, node.ID, graph.EdgeCalls) + fixture.reader.addEdge(node.ID, "rs-target-type", graph.EdgeMemberOf) + } + } + addCalls(fixture.consumer.ID, "aggregate-primary", 31) + addCalls(alt.ID, "aggregate-alt", test.altCalls) + if got := exploreTypedAnchorFixtureProjects(fixture); got != test.want { + t.Fatalf("projection = %v, want %v with %d aggregate calls", got, test.want, 32+test.altCalls) + } + }) + } +} + +type exploreTypedAnchorUnsupportedAdjacencyReader struct { + delegate *exploreTypedAnchorTestReader + legacyCalls int +} + +func (r *exploreTypedAnchorUnsupportedAdjacencyReader) GetNodesByIDs(ids []string) map[string]*graph.Node { + return r.delegate.GetNodesByIDs(ids) +} + +func (r *exploreTypedAnchorUnsupportedAdjacencyReader) GetInEdgesByNodeIDs([]string) map[string][]*graph.Edge { + r.legacyCalls++ + panic("typed-anchor projection fell back to legacy incoming adjacency") +} + +func (r *exploreTypedAnchorUnsupportedAdjacencyReader) GetOutEdgesByNodeIDs([]string) map[string][]*graph.Edge { + r.legacyCalls++ + panic("typed-anchor projection fell back to legacy outgoing adjacency") +} + +func requireExploreTypedAnchorOriginalCandidates( + t *testing.T, + ctx context.Context, + fixture exploreTypedAnchorFixture, + reader exploreTypedAnchorBatchReader, +) { + t.Helper() + got := projectExploreTypedAnchorCandidates( + ctx, fixture.task, fixture.candidates, reader, + exploreTypedAnchorTestScope(), len(fixture.candidates), fixture.protected, "", + ) + if len(got) != len(fixture.candidates) { + t.Fatalf("candidate count = %d, want original %d", len(got), len(fixture.candidates)) + } + for index := range fixture.candidates { + if got[index] != fixture.candidates[index] { + t.Fatalf("candidate[%d] changed on incomplete projection: got %#v want %#v", index, got[index], fixture.candidates[index]) + } + } + if exploreTypedAnchorCandidateByID(got, fixture.consumer.ID) != nil || + exploreTypedAnchorCandidateByID(got, fixture.member.ID) != nil { + t.Fatal("incomplete projection retained a partial consumer/member pair") + } +} + +func TestProjectExploreTypedAnchorCandidatesFailsClosedOnAdjacencyFaults(t *testing.T) { + t.Run("unsupported capability never falls back", func(t *testing.T) { + fixture := newExploreTypedAnchorFixture( + "rs", "--replace causes duplicate output", "replacer", "Replacer", "replace", "replace_all", false, + ) + reader := &exploreTypedAnchorUnsupportedAdjacencyReader{delegate: fixture.reader} + requireExploreTypedAnchorOriginalCandidates(t, context.Background(), fixture, reader) + if reader.legacyCalls != 0 { + t.Fatalf("legacy adjacency calls = %d, want 0", reader.legacyCalls) + } + }) + + stages := []struct { + name string + direction string + ordinal int + }{ + {name: "field outgoing", direction: "out", ordinal: 1}, + {name: "field direct incoming", direction: "in", ordinal: 1}, + {name: "owner incoming", direction: "in", ordinal: 2}, + {name: "calls outgoing", direction: "out", ordinal: 2}, + {name: "member outgoing", direction: "out", ordinal: 3}, + } + faults := []struct { + name string + sticky bool + }{ + {name: "partial error"}, + {name: "sticky cancellation", sticky: true}, + } + for _, stage := range stages { + stage := stage + t.Run(stage.name, func(t *testing.T) { + for _, fault := range faults { + fault := fault + t.Run(fault.name, func(t *testing.T) { + fixture := newExploreTypedAnchorFixture( + "rs", "--replace causes duplicate output", "replacer", "Replacer", "replace", "replace_all", false, + ) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + incomingCalls, outgoingCalls := 0, 0 + injected := false + inject := func(projection graph.BoundedEdgeIdentityProjection) (graph.BoundedEdgeIdentityProjection, error) { + injected = true + if fault.sticky { + cancel() + return projection, nil + } + return projection, errors.New("injected adjacency failure") + } + fixture.reader.incomingProjection = func( + callCtx context.Context, + ids []string, + kinds []graph.EdgeKind, + limit int, + ) (graph.BoundedEdgeIdentityProjection, error) { + incomingCalls++ + projection, err := exploreTypedAnchorTestEdgeProjection(callCtx, ids, kinds, limit, fixture.reader.in) + if err != nil || stage.direction != "in" || incomingCalls != stage.ordinal { + return projection, err + } + return inject(projection) + } + fixture.reader.outgoingProjection = func( + callCtx context.Context, + ids []string, + kinds []graph.EdgeKind, + limit int, + ) (graph.BoundedEdgeIdentityProjection, error) { + outgoingCalls++ + projection, err := exploreTypedAnchorTestEdgeProjection(callCtx, ids, kinds, limit, fixture.reader.out) + if err != nil || stage.direction != "out" || outgoingCalls != stage.ordinal { + return projection, err + } + return inject(projection) + } + + requireExploreTypedAnchorOriginalCandidates(t, ctx, fixture, fixture.reader) + if !injected { + t.Fatal("fault stage was not reached") + } + if fault.sticky && !errors.Is(ctx.Err(), context.Canceled) { + t.Fatalf("request context error = %v, want context.Canceled", ctx.Err()) + } + }) + } + }) + } +} + +func TestProjectExploreTypedAnchorCandidatesRejectsInvalidAdjacencyShapes(t *testing.T) { + stages := []struct { + name string + direction string + ordinal int + }{ + {name: "field outgoing", direction: "out", ordinal: 1}, + {name: "field direct incoming", direction: "in", ordinal: 1}, + {name: "owner incoming", direction: "in", ordinal: 2}, + {name: "calls outgoing", direction: "out", ordinal: 2}, + {name: "member outgoing", direction: "out", ordinal: 3}, + } + shapes := []struct { + name string + empty bool + }{ + {name: "empty endpoint", empty: true}, + {name: "mis-keyed endpoint"}, + } + for _, stage := range stages { + stage := stage + t.Run(stage.name, func(t *testing.T) { + for _, shape := range shapes { + shape := shape + t.Run(shape.name, func(t *testing.T) { + fixture := newExploreTypedAnchorFixture( + "rs", "--replace causes duplicate output", "replacer", "Replacer", "replace", "replace_all", false, + ) + incomingCalls, outgoingCalls := 0, 0 + injected := false + mutate := func( + ids []string, + projection graph.BoundedEdgeIdentityProjection, + ) graph.BoundedEdgeIdentityProjection { + identities := append([]graph.EdgeIdentity(nil), projection.ByEndpoint[ids[0]]...) + if len(identities) == 0 { + t.Fatal("fixture produced no identity at injected stage") + } + if stage.direction == "out" { + if shape.empty { + identities[0].From = "" + } else { + identities[0].From = "mis-keyed-source" + } + } else if shape.empty { + identities[0].To = "" + } else { + identities[0].To = "mis-keyed-target" + } + projection.ByEndpoint[ids[0]] = identities + injected = true + return projection + } + fixture.reader.incomingProjection = func( + callCtx context.Context, + ids []string, + kinds []graph.EdgeKind, + limit int, + ) (graph.BoundedEdgeIdentityProjection, error) { + incomingCalls++ + projection, err := exploreTypedAnchorTestEdgeProjection(callCtx, ids, kinds, limit, fixture.reader.in) + if err == nil && stage.direction == "in" && incomingCalls == stage.ordinal { + projection = mutate(ids, projection) + } + return projection, err + } + fixture.reader.outgoingProjection = func( + callCtx context.Context, + ids []string, + kinds []graph.EdgeKind, + limit int, + ) (graph.BoundedEdgeIdentityProjection, error) { + outgoingCalls++ + projection, err := exploreTypedAnchorTestEdgeProjection(callCtx, ids, kinds, limit, fixture.reader.out) + if err == nil && stage.direction == "out" && outgoingCalls == stage.ordinal { + projection = mutate(ids, projection) + } + return projection, err + } + requireExploreTypedAnchorOriginalCandidates(t, context.Background(), fixture, fixture.reader) + if !injected { + t.Fatal("invalid-shape stage was not reached") + } + }) + } + }) + } + + t.Run("limit plus one without truncation", func(t *testing.T) { + fixture := newExploreTypedAnchorFixture( + "rs", "--replace causes duplicate output", "replacer", "Replacer", "replace", "replace_all", false, + ) + injected := false + fixture.reader.outgoingProjection = func( + _ context.Context, + ids []string, + kinds []graph.EdgeKind, + limit int, + ) (graph.BoundedEdgeIdentityProjection, error) { + identities := make([]graph.EdgeIdentity, limit+1) + for index := range identities { + identities[index] = graph.EdgeIdentity{ + From: ids[0], To: fmt.Sprintf("dishonest-target-%02d", index), Kind: kinds[0], + } + } + injected = true + return graph.BoundedEdgeIdentityProjection{ + ByEndpoint: map[string][]graph.EdgeIdentity{ids[0]: identities}, + Truncated: map[string]bool{}, + }, nil + } + requireExploreTypedAnchorOriginalCandidates(t, context.Background(), fixture, fixture.reader) + if !injected { + t.Fatal("dishonest capability was not invoked") + } + }) +} + +func TestProjectExploreTypedAnchorCandidatesFailsClosedOnMissingHydration(t *testing.T) { + for _, test := range []struct { + name string + missingID string + wantBatches int + }{ + {name: "field endpoint", missingID: "rs-owner", wantBatches: 1}, + {name: "consumer", missingID: "rs-consumer", wantBatches: 2}, + {name: "member", missingID: "rs-member", wantBatches: 3}, + } { + t.Run(test.name, func(t *testing.T) { + fixture := newExploreTypedAnchorFixture( + "rs", "--replace causes duplicate output", "replacer", "Replacer", "replace", "replace_all", false, + ) + delete(fixture.reader.nodes, test.missingID) + requireExploreTypedAnchorOriginalCandidates(t, context.Background(), fixture, fixture.reader) + if fixture.reader.nodeBatches != test.wantBatches { + t.Fatalf("node batches = %d, want failure at batch %d", fixture.reader.nodeBatches, test.wantBatches) + } + }) + } +} + +type exploreTypedAnchorFinalGateContext struct { + context.Context + armed bool + checks int + cancelAfter int +} + +func (ctx *exploreTypedAnchorFinalGateContext) Err() error { + if err := ctx.Context.Err(); err != nil { + return err + } + if !ctx.armed { + return nil + } + ctx.checks++ + if ctx.checks >= ctx.cancelAfter { + return context.Canceled + } + return nil +} + +func TestProjectExploreTypedAnchorCandidatesRechecksCancellationBeforeReserve(t *testing.T) { + fixture := newExploreTypedAnchorFixture( + "rs", "--replace causes duplicate output", "replacer", "Replacer", "replace", "replace_all", false, + ) + gateCtx := &exploreTypedAnchorFinalGateContext{ + Context: context.Background(), + cancelAfter: 3, + } + fixture.reader.nodeProjection = func(_ context.Context, ids []string) (map[string]*graph.Node, error) { + nodes := make(map[string]*graph.Node, len(ids)) + for _, id := range ids { + if node := fixture.reader.nodes[id]; node != nil { + nodes[id] = node + } + } + if fixture.reader.nodeBatches == 3 { + gateCtx.armed = true + } + return nodes, nil + } + + requireExploreTypedAnchorOriginalCandidates(t, gateCtx, fixture, fixture.reader) + if fixture.reader.nodeBatches != 3 { + t.Fatalf("node batches = %d, want final hydration batch 3", fixture.reader.nodeBatches) + } + if err := gateCtx.Err(); !errors.Is(err, context.Canceled) { + t.Fatalf("final gate context error = %v, want context.Canceled", err) + } +} + +func TestProjectExploreTypedAnchorCandidatesFailsClosedOnNodeRefetchFaults(t *testing.T) { + for _, batch := range []int{1, 2, 3} { + batch := batch + t.Run(fmt.Sprintf("batch %d", batch), func(t *testing.T) { + for _, fault := range []struct { + name string + sticky bool + }{ + {name: "partial error"}, + {name: "sticky cancellation", sticky: true}, + } { + fault := fault + t.Run(fault.name, func(t *testing.T) { + fixture := newExploreTypedAnchorFixture( + "rs", "--replace causes duplicate output", "replacer", "Replacer", "replace", "replace_all", false, + ) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + injected := false + fixture.reader.nodeProjection = func(_ context.Context, ids []string) (map[string]*graph.Node, error) { + nodes := make(map[string]*graph.Node, len(ids)) + for _, id := range ids { + if node := fixture.reader.nodes[id]; node != nil { + nodes[id] = node + } + } + if fixture.reader.nodeBatches != batch { + return nodes, nil + } + injected = true + if fault.sticky { + cancel() + return nodes, nil + } + return nodes, errors.New("injected node refetch failure") + } + + requireExploreTypedAnchorOriginalCandidates(t, ctx, fixture, fixture.reader) + if !injected { + t.Fatalf("node fault batch %d was not reached", batch) + } + if fault.sticky && !errors.Is(ctx.Err(), context.Canceled) { + t.Fatalf("request context error = %v, want context.Canceled", ctx.Err()) + } + }) + } + }) + } +} + +func TestProjectExploreTypedAnchorCandidatesFiltersScopeAndTests(t *testing.T) { + for _, test := range []struct { + name string + mutate func(exploreTypedAnchorFixture) + }{ + { + name: "test consumer", + mutate: func(fixture exploreTypedAnchorFixture) { + fixture.consumer.Meta = map[string]any{"is_test": true} + }, + }, + { + name: "test member", + mutate: func(fixture exploreTypedAnchorFixture) { + fixture.member.Meta = map[string]any{"is_test": true} + }, + }, + { + name: "owner outside repository scope", + mutate: func(fixture exploreTypedAnchorFixture) { + fixture.reader.nodes["rs-owner"].RepoPrefix = "other" + }, + }, + { + name: "consumer outside workspace scope", + mutate: func(fixture exploreTypedAnchorFixture) { + fixture.consumer.WorkspaceID = "other" + }, + }, + { + name: "member outside project scope", + mutate: func(fixture exploreTypedAnchorFixture) { + fixture.member.ProjectID = "other" + }, + }, + { + name: "declared type outside repository scope", + mutate: func(fixture exploreTypedAnchorFixture) { + fixture.reader.nodes["rs-target-type"].RepoPrefix = "other" + }, + }, + } { + t.Run(test.name, func(t *testing.T) { + fixture := newExploreTypedAnchorFixture( + "rs", "--replace causes duplicate output", "replacer", "Replacer", "replace", "replace_all", false, + ) + test.mutate(fixture) + requireExploreTypedAnchorOriginalCandidates(t, context.Background(), fixture, fixture.reader) + }) + } + + t.Run("out-of-scope peer owner is ignored", func(t *testing.T) { + fixture := newExploreTypedAnchorFixture( + "rs", "--replace causes duplicate output", "replacer", "Replacer", "replace", "replace_all", false, + ) + peer := exploreTypedAnchorNode("rs-peer-owner", "PeerSink", "peer.rs", "rs", graph.KindType) + peer.RepoPrefix = "other" + fixture.reader.nodes[peer.ID] = peer + fixture.reader.addEdge(fixture.field.ID, peer.ID, graph.EdgeMemberOf) + if !exploreTypedAnchorFixtureProjects(fixture) { + t.Fatal("out-of-scope peer owner made the in-scope proof ambiguous") + } + }) +} diff --git a/internal/mcp/explore_typed_anchor_overlay_test.go b/internal/mcp/explore_typed_anchor_overlay_test.go new file mode 100644 index 00000000..c53d0f41 --- /dev/null +++ b/internal/mcp/explore_typed_anchor_overlay_test.go @@ -0,0 +1,170 @@ +package mcp + +import ( + "context" + "path/filepath" + "testing" + + "github.com/zzet/gortex/internal/graph" + "github.com/zzet/gortex/internal/graph/store_sqlite" + "github.com/zzet/gortex/internal/query" +) + +func addExploreTypedAnchorFixtureToStore(store graph.Store, fixture exploreTypedAnchorFixture) { + nodes := make([]*graph.Node, 0, len(fixture.reader.nodes)) + for _, node := range fixture.reader.nodes { + nodes = append(nodes, node) + } + edges := make([]*graph.Edge, 0) + for _, outgoing := range fixture.reader.out { + edges = append(edges, outgoing...) + } + store.AddBatch(nodes, edges) +} + +func exploreTypedAnchorReaderProjects( + ctx context.Context, + fixture exploreTypedAnchorFixture, + reader exploreTypedAnchorBatchReader, +) bool { + got := projectExploreTypedAnchorCandidates( + ctx, fixture.task, fixture.candidates, reader, + exploreTypedAnchorTestScope(), len(fixture.candidates), fixture.protected, "", + ) + return exploreTypedAnchorCandidateByID(got, fixture.member.ID) != nil +} + +func requireExploreTypedAnchorStoreParity(t *testing.T, store graph.Store) { + t.Helper() + fixture := newExploreTypedAnchorFixture( + "rs", "--replace causes duplicate output", "replacer", "Replacer", "replace", "replace_all", false, + ) + addExploreTypedAnchorFixtureToStore(store, fixture) + if !exploreTypedAnchorReaderProjects(context.Background(), fixture, store) { + t.Fatal("bounded store reader did not produce the complete typed-anchor proof") + } + + t.Run("target member tombstone and replacement", func(t *testing.T) { + tombstone := graph.NewOverlayLayer() + tombstone.MarkRemoved(fixture.member.Name, fixture.member.ID) + if exploreTypedAnchorReaderProjects( + context.Background(), fixture, graph.NewOverlaidView(store, tombstone), + ) { + t.Fatal("overlay-tombstoned typed member retained a stale base proof") + } + + targetType := fixture.reader.nodes["rs-target-type"] + replacement := graph.NewOverlayLayer() + replacement.MarkFile(fixture.member.FilePath, false) + replacement.AddNode(fixture.member.FilePath, fixture.member) + replacement.AddNode(targetType.FilePath, targetType) + replacement.AddEdge(&graph.Edge{ + From: fixture.member.ID, To: targetType.ID, Kind: graph.EdgeMemberOf, + FilePath: fixture.member.FilePath, + }) + if !exploreTypedAnchorReaderProjects( + context.Background(), fixture, graph.NewOverlaidView(store, replacement), + ) { + t.Fatal("same-identity member replacement did not restore the current-view proof") + } + }) + + t.Run("source consumer tombstone and replacement", func(t *testing.T) { + tombstone := graph.NewOverlayLayer() + tombstone.MarkRemoved(fixture.consumer.Name, fixture.consumer.ID) + if exploreTypedAnchorReaderProjects( + context.Background(), fixture, graph.NewOverlaidView(store, tombstone), + ) { + t.Fatal("overlay-tombstoned consumer retained a stale base proof") + } + + owner := fixture.reader.nodes["rs-owner"] + replacement := graph.NewOverlayLayer() + replacement.MarkFile(fixture.consumer.FilePath, false) + replacement.AddNode(fixture.field.FilePath, fixture.field) + replacement.AddNode(owner.FilePath, owner) + replacement.AddNode(fixture.consumer.FilePath, fixture.consumer) + for _, edge := range []*graph.Edge{ + {From: fixture.field.ID, To: owner.ID, Kind: graph.EdgeMemberOf, FilePath: fixture.field.FilePath}, + {From: fixture.consumer.ID, To: owner.ID, Kind: graph.EdgeMemberOf, FilePath: fixture.consumer.FilePath}, + {From: fixture.consumer.ID, To: fixture.field.ID, Kind: graph.EdgeReads, FilePath: fixture.consumer.FilePath}, + {From: fixture.consumer.ID, To: fixture.member.ID, Kind: graph.EdgeCalls, FilePath: fixture.consumer.FilePath}, + } { + replacement.AddEdge(edge) + } + if !exploreTypedAnchorReaderProjects( + context.Background(), fixture, graph.NewOverlaidView(store, replacement), + ) { + t.Fatal("same-identity consumer replacement did not restore the current-view proof") + } + }) +} + +func TestProjectExploreTypedAnchorCandidatesGraphSQLiteOverlayParity(t *testing.T) { + t.Run("graph", func(t *testing.T) { + requireExploreTypedAnchorStoreParity(t, graph.New()) + }) + t.Run("sqlite", func(t *testing.T) { + store, err := store_sqlite.Open(filepath.Join(t.TempDir(), "typed-anchor.sqlite")) + if err != nil { + t.Fatal(err) + } + defer func() { + if err := store.Close(); err != nil { + t.Errorf("close sqlite store: %v", err) + } + }() + requireExploreTypedAnchorStoreParity(t, store) + }) +} + +func TestProjectExploreTypedAnchorCandidatesUsesEngineOverlayReader(t *testing.T) { + fixture := newExploreTypedAnchorFixture( + "rs", "--replace causes duplicate output", "replacer", "Replacer", "replace", "replace_all", false, + ) + routeStore := graph.New() + addExploreTypedAnchorFixtureToStore(routeStore, fixture) + server := &Server{graph: graph.New(), engine: query.NewEngine(routeStore)} + baseEngine := server.engineFor(context.Background()) + if baseEngine == nil || !exploreTypedAnchorReaderProjects( + context.Background(), fixture, baseEngine.Reader(), + ) { + t.Fatal("non-overlay engine reader lost its independently configured graph") + } + if exploreTypedAnchorReaderProjects(context.Background(), fixture, server.readerFor(context.Background())) { + t.Fatal("empty server graph unexpectedly contains the engine's typed route") + } + + tombstone := graph.NewOverlayLayer() + tombstone.MarkFile(fixture.member.FilePath, false) + tombstone.MarkRemoved(fixture.member.Name, fixture.member.ID) + tombstoneView := graph.NewOverlaidView(routeStore, tombstone) + tombstoneCtx := WithOverlayView(context.Background(), tombstoneView) + tombstoneEngine := server.engineFor(tombstoneCtx) + if tombstoneEngine == nil || tombstoneEngine.Reader() != tombstoneView { + t.Fatal("engineFor did not bind the request-local overlay reader") + } + if exploreTypedAnchorReaderProjects(tombstoneCtx, fixture, tombstoneEngine.Reader()) { + t.Fatal("request-local engine reader bypassed the overlay tombstone") + } + + targetType := fixture.reader.nodes["rs-target-type"] + replacement := graph.NewOverlayLayer() + replacement.MarkFile(fixture.member.FilePath, false) + replacement.AddNode(fixture.member.FilePath, fixture.member) + replacement.AddNode(targetType.FilePath, targetType) + replacement.AddEdge(&graph.Edge{ + From: fixture.member.ID, To: targetType.ID, Kind: graph.EdgeMemberOf, + FilePath: fixture.member.FilePath, + }) + replacementView := graph.NewOverlaidView(routeStore, replacement) + replacementCtx := WithOverlayView(context.Background(), replacementView) + if !exploreTypedAnchorReaderProjects( + replacementCtx, fixture, server.engineFor(replacementCtx).Reader(), + ) { + t.Fatal("engine overlay reader did not restore the same-identity replacement") + } + if !exploreTypedAnchorReaderProjects(context.Background(), fixture, server.engine.Reader()) { + t.Fatal("overlay requests mutated the server's base engine reader") + } +} diff --git a/internal/mcp/explore_typed_anchor_projection.go b/internal/mcp/explore_typed_anchor_projection.go index ec159df2..681b6724 100644 --- a/internal/mcp/explore_typed_anchor_projection.go +++ b/internal/mcp/explore_typed_anchor_projection.go @@ -20,14 +20,164 @@ const ( exploreTypedAnchorProjectionSignal = "typed_anchor_projection" ) -// exploreTypedAnchorBatchReader deliberately exposes only batched graph reads. -// The SQLite implementation turns each stage into one bounded IN query, while -// query overlays preserve request-local edits. A point-reader interface here -// previously amplified one projection into dozens of SQLite round trips. +// exploreTypedAnchorBatchReader deliberately exposes only exact batched node +// hydration. Every adjacency stage below requires the optional metadata-free +// bounded capabilities; unsupported readers fail the whole projection closed +// rather than falling back to legacy full-row edge reads. type exploreTypedAnchorBatchReader interface { GetNodesByIDs([]string) map[string]*graph.Node - GetInEdgesByNodeIDs([]string) map[string][]*graph.Edge - GetOutEdgesByNodeIDs([]string) map[string][]*graph.Edge +} + +func exploreTypedAnchorNodesByIDsBounded( + ctx context.Context, + reader exploreTypedAnchorBatchReader, + ids []string, + limit int, +) (map[string]*graph.Node, bool) { + if ctx == nil { + ctx = context.Background() + } + boundedIDs, complete := exploreBoundedNodeIDs(ids, limit) + if reader == nil || !complete || ctx.Err() != nil { + return nil, false + } + if len(boundedIDs) == 0 { + return map[string]*graph.Node{}, true + } + var ( + nodes map[string]*graph.Node + err error + ) + if contextual, supported := reader.(exploreContextNodesReader); supported { + nodes, err = contextual.GetNodesByIDsContext(ctx, boundedIDs) + } else { + nodes = reader.GetNodesByIDs(boundedIDs) + } + if err != nil || ctx.Err() != nil || len(nodes) != len(boundedIDs) { + return nil, false + } + for _, id := range boundedIDs { + if nodes[id] == nil { + return nil, false + } + } + if ctx.Err() != nil { + return nil, false + } + return nodes, true +} + +func exploreTypedAnchorOutgoingIdentitiesBounded( + ctx context.Context, + reader exploreTypedAnchorBatchReader, + ids []string, + kinds []graph.EdgeKind, + limit int, +) (map[string][]graph.EdgeIdentity, bool) { + if ctx == nil { + ctx = context.Background() + } + bounded, supported := reader.(graph.BoundedOutgoingEdgeIdentityReader) + if !supported || ctx.Err() != nil { + return nil, false + } + ids = exploreTypedAnchorSortedIDs(ids) + projection, err := bounded.FindOutgoingEdgeIdentitiesBounded(ctx, ids, kinds, limit) + if err != nil || ctx.Err() != nil { + return nil, false + } + if !exploreTypedAnchorValidIdentityProjection(projection, ids, kinds, limit, true) || ctx.Err() != nil { + return nil, false + } + return projection.ByEndpoint, true +} + +func exploreTypedAnchorIncomingIdentitiesBounded( + ctx context.Context, + reader exploreTypedAnchorBatchReader, + ids []string, + kinds []graph.EdgeKind, + limit int, +) (map[string][]graph.EdgeIdentity, bool) { + if ctx == nil { + ctx = context.Background() + } + bounded, supported := reader.(graph.BoundedIncomingEdgeIdentityReader) + if !supported || ctx.Err() != nil { + return nil, false + } + ids = exploreTypedAnchorSortedIDs(ids) + projection, err := bounded.FindIncomingEdgeIdentitiesBounded(ctx, ids, kinds, limit) + if err != nil || ctx.Err() != nil { + return nil, false + } + if !exploreTypedAnchorValidIdentityProjection(projection, ids, kinds, limit, false) || ctx.Err() != nil { + return nil, false + } + return projection.ByEndpoint, true +} + +func exploreTypedAnchorValidIdentityProjection( + projection graph.BoundedEdgeIdentityProjection, + ids []string, + kinds []graph.EdgeKind, + limit int, + outgoing bool, +) bool { + if limit < 1 || len(kinds) == 0 || len(projection.ByEndpoint) > len(ids) || len(projection.Truncated) > len(ids) { + return false + } + requested := make(map[string]struct{}, len(ids)) + for _, id := range ids { + requested[id] = struct{}{} + } + allowed := make(map[graph.EdgeKind]struct{}, len(kinds)) + for _, kind := range kinds { + allowed[kind] = struct{}{} + } + for id, truncated := range projection.Truncated { + if _, ok := requested[id]; !ok || truncated { + return false + } + } + for id, identities := range projection.ByEndpoint { + if _, ok := requested[id]; !ok || len(identities) > limit || projection.Truncated[id] { + return false + } + seen := make(map[graph.EdgeIdentity]struct{}, len(identities)) + for _, identity := range identities { + if _, ok := allowed[identity.Kind]; !ok || identity.From == "" || identity.To == "" { + return false + } + if outgoing && identity.From != id || !outgoing && identity.To != id { + return false + } + if _, duplicate := seen[identity]; duplicate { + return false + } + seen[identity] = struct{}{} + } + } + return true +} + +func exploreTypedAnchorSortIdentities(identities []graph.EdgeIdentity) { + sort.SliceStable(identities, func(i, j int) bool { + left, right := identities[i], identities[j] + if left.Kind != right.Kind { + return left.Kind < right.Kind + } + if left.From != right.From { + return left.From < right.From + } + if left.To != right.To { + return left.To < right.To + } + if left.FilePath != right.FilePath { + return left.FilePath < right.FilePath + } + return left.Line < right.Line + }) } type exploreTypedAnchorField struct { @@ -91,7 +241,7 @@ func projectExploreTypedAnchorCandidates( return candidates } projection, ok := findExploreTypedAnchorProjection(ctx, task, candidates, reader, scope, protectedAnchors) - if !ok { + if !ok || ctx.Err() != nil { return candidates } reserved := exploreTypedAnchorReservedCandidateIDs(candidates, protectedAnchors, protectedImplementationID) @@ -167,27 +317,37 @@ func hydrateExploreTypedAnchorFields( for _, field := range fields { fieldIDs = append(fieldIDs, field.field.ID) } - out := reader.GetOutEdgesByNodeIDs(exploreTypedAnchorSortedIDs(fieldIDs)) - if ctx.Err() != nil { + out, complete := exploreTypedAnchorOutgoingIdentitiesBounded( + ctx, + reader, + fieldIDs, + []graph.EdgeKind{graph.EdgeMemberOf, graph.EdgeTypedAs}, + exploreTypedAnchorFieldRelationLimit, + ) + if !complete { return false } - relations := make(map[string][]*graph.Edge, len(fields)) - endpointIDs := make([]string, 0, len(fields)*2) + relations := make(map[string][]graph.EdgeIdentity, len(fields)) + endpointIDs := make([]string, 0, len(fields)*exploreTypedAnchorFieldRelationLimit) for _, field := range fields { - edges, overflow := exploreTypedAnchorBoundedEdges( - out[field.field.ID], exploreTypedAnchorFieldRelationLimit, graph.EdgeMemberOf, graph.EdgeTypedAs, - ) - if overflow { - return false - } + edges := append([]graph.EdgeIdentity(nil), out[field.field.ID]...) + exploreTypedAnchorSortIdentities(edges) relations[field.field.ID] = edges for _, edge := range edges { + if edge.From != field.field.ID || edge.To == "" { + return false + } endpointIDs = append(endpointIDs, edge.To) } } - nodes := reader.GetNodesByIDs(exploreTypedAnchorSortedIDs(endpointIDs)) - if ctx.Err() != nil { + nodes, complete := exploreTypedAnchorNodesByIDsBounded( + ctx, + reader, + endpointIDs, + exploreTypedAnchorFieldLimit*exploreTypedAnchorFieldRelationLimit, + ) + if !complete { return false } @@ -195,13 +355,8 @@ func hydrateExploreTypedAnchorFields( for _, field := range fields { owners := make(map[string]*graph.Node) field.typedIDs = make(map[string]struct{}) - complete := true for _, edge := range relations[field.field.ID] { node := nodes[edge.To] - if node == nil { - complete = false - break - } if !exploreNodeWithinQueryScope(node, scope) { continue } @@ -216,7 +371,7 @@ func hydrateExploreTypedAnchorFields( } } } - if !complete || len(owners) != 1 || (len(field.typedIDs) == 0 && field.canonicalType == "") { + if len(owners) != 1 || (len(field.typedIDs) == 0 && field.canonicalType == "") { continue } for _, owner := range owners { @@ -239,18 +394,36 @@ func hydrateExploreTypedAnchorConsumers( ) ([]*exploreTypedAnchorConsumer, bool) { // hydrateExploreTypedAnchorFields filters in place but cannot resize its // caller's slice. Ignore any cleared or unhydrated seed here. - queryIDs := make([]string, 0, len(fields)*2) + fieldIDs := make([]string, 0, len(fields)) + ownerIDs := make([]string, 0, len(fields)) for _, field := range fields { if field == nil || field.owner == nil { continue } - queryIDs = append(queryIDs, field.field.ID, field.owner.ID) + fieldIDs = append(fieldIDs, field.field.ID) + ownerIDs = append(ownerIDs, field.owner.ID) } - if len(queryIDs) == 0 { + if len(fieldIDs) == 0 { return nil, false } - in := reader.GetInEdgesByNodeIDs(exploreTypedAnchorSortedIDs(queryIDs)) - if ctx.Err() != nil { + fieldIn, complete := exploreTypedAnchorIncomingIdentitiesBounded( + ctx, + reader, + fieldIDs, + []graph.EdgeKind{graph.EdgeReads, graph.EdgeWrites, graph.EdgeAccessesField}, + exploreTypedAnchorConsumerLimit, + ) + if !complete { + return nil, false + } + ownerIn, complete := exploreTypedAnchorIncomingIdentitiesBounded( + ctx, + reader, + ownerIDs, + []graph.EdgeKind{graph.EdgeMemberOf}, + exploreTypedAnchorOwnerMemberLimit, + ) + if !complete { return nil, false } @@ -259,29 +432,26 @@ func hydrateExploreTypedAnchorConsumers( direct bool } rawByField := make(map[string]map[string]rawConsumer, len(fields)) - consumerIDs := make([]string, 0, len(fields)*exploreTypedAnchorConsumerLimit) + consumerIDs := make([]string, 0, len(fields)*exploreTypedAnchorOwnerMemberLimit) for _, field := range fields { if field == nil || field.owner == nil { continue } - direct, overflow := exploreTypedAnchorBoundedEdges( - in[field.field.ID], exploreTypedAnchorConsumerLimit, - graph.EdgeReads, graph.EdgeWrites, graph.EdgeAccessesField, - ) - if overflow { - return nil, false - } - members, overflow := exploreTypedAnchorBoundedEdges( - in[field.owner.ID], exploreTypedAnchorOwnerMemberLimit, graph.EdgeMemberOf, - ) - if overflow { - return nil, false - } + direct := append([]graph.EdgeIdentity(nil), fieldIn[field.field.ID]...) + members := append([]graph.EdgeIdentity(nil), ownerIn[field.owner.ID]...) + exploreTypedAnchorSortIdentities(direct) + exploreTypedAnchorSortIdentities(members) byID := make(map[string]rawConsumer, len(direct)+len(members)) for _, edge := range members { + if edge.From == "" || edge.To != field.owner.ID { + return nil, false + } byID[edge.From] = rawConsumer{id: edge.From} } for _, edge := range direct { + if edge.From == "" || edge.To != field.field.ID { + return nil, false + } consumer, member := byID[edge.From] if !member { // A field read outside the owning type does not prove the @@ -296,8 +466,13 @@ func hydrateExploreTypedAnchorConsumers( consumerIDs = append(consumerIDs, id) } } - nodes := reader.GetNodesByIDs(exploreTypedAnchorSortedIDs(consumerIDs)) - if ctx.Err() != nil { + nodes, complete := exploreTypedAnchorNodesByIDsBounded( + ctx, + reader, + consumerIDs, + exploreTypedAnchorFieldLimit*exploreTypedAnchorOwnerMemberLimit, + ) + if !complete { return nil, false } @@ -353,16 +528,28 @@ func hydrateExploreTypedAnchorCalls( consumerIDs = append(consumerIDs, consumer.node.ID) consumerByID[consumer.node.ID] = append(consumerByID[consumer.node.ID], consumer) } - out := reader.GetOutEdgesByNodeIDs(exploreTypedAnchorSortedIDs(consumerIDs)) - if ctx.Err() != nil { + out, complete := exploreTypedAnchorOutgoingIdentitiesBounded( + ctx, + reader, + consumerIDs, + []graph.EdgeKind{graph.EdgeCalls}, + exploreTypedAnchorCallLimit, + ) + if !complete { return nil, false } calls := make([]exploreTypedAnchorCall, 0, exploreTypedAnchorCallLimit) seenCalls := make(map[string]struct{}) for _, consumerID := range exploreTypedAnchorSortedIDs(consumerIDs) { - edges, overflow := exploreTypedAnchorBoundedEdges(out[consumerID], exploreTypedAnchorCallLimit, graph.EdgeCalls) - if overflow || len(calls)+len(edges)*len(consumerByID[consumerID]) > exploreTypedAnchorCallLimit { + edges := append([]graph.EdgeIdentity(nil), out[consumerID]...) + exploreTypedAnchorSortIdentities(edges) + for _, edge := range edges { + if edge.From != consumerID || edge.To == "" { + return nil, false + } + } + if len(calls)+len(edges)*len(consumerByID[consumerID]) > exploreTypedAnchorCallLimit { return nil, false } for _, consumer := range consumerByID[consumerID] { @@ -384,27 +571,37 @@ func hydrateExploreTypedAnchorCalls( for _, call := range calls { memberIDs = append(memberIDs, call.memberID) } - memberOut := reader.GetOutEdgesByNodeIDs(exploreTypedAnchorSortedIDs(memberIDs)) - if ctx.Err() != nil { + memberOut, complete := exploreTypedAnchorOutgoingIdentitiesBounded( + ctx, + reader, + memberIDs, + []graph.EdgeKind{graph.EdgeMemberOf}, + exploreTypedAnchorMemberOwnerLimit, + ) + if !complete { return nil, false } - memberRelations := make(map[string][]*graph.Edge) - ownerIDs := make([]string, 0, len(memberIDs)) + memberRelations := make(map[string][]graph.EdgeIdentity) + ownerIDs := make([]string, 0, len(memberIDs)*exploreTypedAnchorMemberOwnerLimit) for _, memberID := range exploreTypedAnchorSortedIDs(memberIDs) { - edges, overflow := exploreTypedAnchorBoundedEdges( - memberOut[memberID], exploreTypedAnchorMemberOwnerLimit, graph.EdgeMemberOf, - ) - if overflow { - return nil, false - } + edges := append([]graph.EdgeIdentity(nil), memberOut[memberID]...) + exploreTypedAnchorSortIdentities(edges) memberRelations[memberID] = edges for _, edge := range edges { + if edge.From != memberID || edge.To == "" { + return nil, false + } ownerIDs = append(ownerIDs, edge.To) } } hydrationIDs := append(exploreTypedAnchorSortedIDs(memberIDs), ownerIDs...) - nodes := reader.GetNodesByIDs(exploreTypedAnchorSortedIDs(hydrationIDs)) - if ctx.Err() != nil { + nodes, complete := exploreTypedAnchorNodesByIDsBounded( + ctx, + reader, + hydrationIDs, + exploreTypedAnchorCallLimit*(exploreTypedAnchorMemberOwnerLimit+1), + ) + if !complete { return nil, false } @@ -443,7 +640,7 @@ func hydrateExploreTypedAnchorCalls( func exploreTypedAnchorMemberTypeProof( field *exploreTypedAnchorField, - edges []*graph.Edge, + edges []graph.EdgeIdentity, nodes map[string]*graph.Node, scope query.QueryOptions, ) (exact, matches, complete bool) { @@ -604,48 +801,6 @@ func exploreTypedAnchorCanonicalTypeMatches(fieldIdentity string, owner *graph.N return fieldIdentity == exploreTypedAnchorCanonicalType(owner.QualName) } -func exploreTypedAnchorBoundedEdges(edges []*graph.Edge, limit int, kinds ...graph.EdgeKind) ([]*graph.Edge, bool) { - if limit < 1 || len(edges) == 0 || len(kinds) == 0 { - return nil, false - } - allowed := make(map[graph.EdgeKind]struct{}, len(kinds)) - for _, kind := range kinds { - allowed[kind] = struct{}{} - } - bounded := make([]*graph.Edge, 0, min(limit+1, len(edges))) - for _, edge := range edges { - if edge == nil { - continue - } - if _, ok := allowed[edge.Kind]; ok { - bounded = append(bounded, edge) - } - } - sort.SliceStable(bounded, func(i, j int) bool { - left, right := bounded[i], bounded[j] - if left.Kind != right.Kind { - return left.Kind < right.Kind - } - if left.From != right.From { - return left.From < right.From - } - if left.To != right.To { - return left.To < right.To - } - if left.FilePath != right.FilePath { - return left.FilePath < right.FilePath - } - if left.Line != right.Line { - return left.Line < right.Line - } - return left.Origin < right.Origin - }) - if len(bounded) > limit { - return nil, true - } - return bounded, false -} - func exploreTypedAnchorSortedIDs(ids []string) []string { if len(ids) == 0 { return nil diff --git a/internal/mcp/explore_typed_anchor_projection_test.go b/internal/mcp/explore_typed_anchor_projection_test.go index e323bf57..fac30664 100644 --- a/internal/mcp/explore_typed_anchor_projection_test.go +++ b/internal/mcp/explore_typed_anchor_projection_test.go @@ -11,14 +11,27 @@ import ( "github.com/zzet/gortex/internal/search/rerank" ) +type exploreTypedAnchorAdjacencyRequest struct { + ids []string + kinds []graph.EdgeKind + limit int +} + type exploreTypedAnchorTestReader struct { - nodes map[string]*graph.Node - in map[string][]*graph.Edge - out map[string][]*graph.Edge - delay time.Duration - nodeBatches int - inBatches int - outBatches int + nodes map[string]*graph.Node + in map[string][]*graph.Edge + out map[string][]*graph.Edge + delay time.Duration + + nodeBatches int + legacyInBatches int + legacyOutBatches int + nodeRequests [][]string + incomingRequests []exploreTypedAnchorAdjacencyRequest + outgoingRequests []exploreTypedAnchorAdjacencyRequest + nodeProjection func(context.Context, []string) (map[string]*graph.Node, error) + incomingProjection func(context.Context, []string, []graph.EdgeKind, int) (graph.BoundedEdgeIdentityProjection, error) + outgoingProjection func(context.Context, []string, []graph.EdgeKind, int) (graph.BoundedEdgeIdentityProjection, error) } func newExploreTypedAnchorTestReader(nodes ...*graph.Node) *exploreTypedAnchorTestReader { @@ -40,35 +53,119 @@ func (r *exploreTypedAnchorTestReader) addEdge(from, to string, kind graph.EdgeK } func (r *exploreTypedAnchorTestReader) GetNodesByIDs(ids []string) map[string]*graph.Node { + nodes, _ := r.GetNodesByIDsContext(context.Background(), ids) + return nodes +} + +func (r *exploreTypedAnchorTestReader) GetNodesByIDsContext( + ctx context.Context, + ids []string, +) (map[string]*graph.Node, error) { r.nodeBatches++ + r.nodeRequests = append(r.nodeRequests, append([]string(nil), ids...)) time.Sleep(r.delay) + if r.nodeProjection != nil { + return r.nodeProjection(ctx, ids) + } result := make(map[string]*graph.Node, len(ids)) for _, id := range ids { if node := r.nodes[id]; node != nil { result[id] = node } } - return result + return result, nil } -func (r *exploreTypedAnchorTestReader) GetInEdgesByNodeIDs(ids []string) map[string][]*graph.Edge { - r.inBatches++ +func (r *exploreTypedAnchorTestReader) GetInEdgesByNodeIDs([]string) map[string][]*graph.Edge { + r.legacyInBatches++ + panic("typed-anchor projection used legacy incoming adjacency") +} + +func (r *exploreTypedAnchorTestReader) GetOutEdgesByNodeIDs([]string) map[string][]*graph.Edge { + r.legacyOutBatches++ + panic("typed-anchor projection used legacy outgoing adjacency") +} + +func (r *exploreTypedAnchorTestReader) FindIncomingEdgeIdentitiesBounded( + ctx context.Context, + ids []string, + kinds []graph.EdgeKind, + limit int, +) (graph.BoundedEdgeIdentityProjection, error) { + r.incomingRequests = append(r.incomingRequests, exploreTypedAnchorAdjacencyRequest{ + ids: append([]string(nil), ids...), kinds: append([]graph.EdgeKind(nil), kinds...), limit: limit, + }) time.Sleep(r.delay) - result := make(map[string][]*graph.Edge, len(ids)) - for _, id := range ids { - result[id] = r.in[id] + if r.incomingProjection != nil { + return r.incomingProjection(ctx, ids, kinds, limit) } - return result + return exploreTypedAnchorTestEdgeProjection(ctx, ids, kinds, limit, r.in) } -func (r *exploreTypedAnchorTestReader) GetOutEdgesByNodeIDs(ids []string) map[string][]*graph.Edge { - r.outBatches++ +func (r *exploreTypedAnchorTestReader) FindOutgoingEdgeIdentitiesBounded( + ctx context.Context, + ids []string, + kinds []graph.EdgeKind, + limit int, +) (graph.BoundedEdgeIdentityProjection, error) { + r.outgoingRequests = append(r.outgoingRequests, exploreTypedAnchorAdjacencyRequest{ + ids: append([]string(nil), ids...), kinds: append([]graph.EdgeKind(nil), kinds...), limit: limit, + }) time.Sleep(r.delay) - result := make(map[string][]*graph.Edge, len(ids)) + if r.outgoingProjection != nil { + return r.outgoingProjection(ctx, ids, kinds, limit) + } + return exploreTypedAnchorTestEdgeProjection(ctx, ids, kinds, limit, r.out) +} + +func exploreTypedAnchorTestEdgeProjection( + ctx context.Context, + ids []string, + kinds []graph.EdgeKind, + limit int, + edgesByID map[string][]*graph.Edge, +) (graph.BoundedEdgeIdentityProjection, error) { + projection := graph.BoundedEdgeIdentityProjection{ + ByEndpoint: make(map[string][]graph.EdgeIdentity), + Truncated: make(map[string]bool), + } + allowed := make(map[graph.EdgeKind]struct{}, len(kinds)) + for _, kind := range kinds { + allowed[kind] = struct{}{} + } for _, id := range ids { - result[id] = r.out[id] + if err := ctx.Err(); err != nil { + return graph.BoundedEdgeIdentityProjection{}, err + } + seen := make(map[graph.EdgeIdentity]struct{}) + identities := make([]graph.EdgeIdentity, 0, limit) + for _, edge := range edgesByID[id] { + if edge == nil { + continue + } + if _, ok := allowed[edge.Kind]; !ok { + continue + } + identity := graph.EdgeIdentity{ + From: edge.From, To: edge.To, Kind: edge.Kind, + FilePath: edge.FilePath, Line: edge.Line, + } + if _, duplicate := seen[identity]; duplicate { + continue + } + seen[identity] = struct{}{} + if len(identities) == limit { + projection.Truncated[id] = true + identities = nil + break + } + identities = append(identities, identity) + } + if len(identities) > 0 { + projection.ByEndpoint[id] = identities + } } - return result + return projection, nil } func removeExploreTypedAnchorTestEdge(edges []*graph.Edge, from, to string, kind graph.EdgeKind) []*graph.Edge { @@ -323,7 +420,7 @@ func TestProjectExploreTypedAnchorCandidatesUsesFixedBatchPipelineDespiteDelay(t fixture := newExploreTypedAnchorFixture( "rs", "--replace causes duplicate output", "replacer", "Replacer", "replace", "replace_all", false, ) - fixture.reader.delay = 3 * time.Millisecond // seven stages exceed the removed 10 ms cutoff + fixture.reader.delay = 3 * time.Millisecond // eight stages exceed the removed 10 ms cutoff got := projectExploreTypedAnchorCandidates( context.Background(), fixture.task, fixture.candidates, fixture.reader, exploreTypedAnchorTestScope(), len(fixture.candidates), fixture.protected, "", @@ -331,9 +428,13 @@ func TestProjectExploreTypedAnchorCandidatesUsesFixedBatchPipelineDespiteDelay(t if exploreTypedAnchorCandidateByID(got, fixture.member.ID) == nil { t.Fatal("scheduler/database delay changed a structurally complete projection") } - if fixture.reader.inBatches != 1 || fixture.reader.outBatches != 3 || fixture.reader.nodeBatches != 3 { - t.Fatalf("batch pipeline = in:%d out:%d nodes:%d, want 1/3/3", - fixture.reader.inBatches, fixture.reader.outBatches, fixture.reader.nodeBatches) + if len(fixture.reader.incomingRequests) != 2 || len(fixture.reader.outgoingRequests) != 3 || fixture.reader.nodeBatches != 3 { + t.Fatalf("batch pipeline = in:%d out:%d nodes:%d, want 2/3/3", + len(fixture.reader.incomingRequests), len(fixture.reader.outgoingRequests), fixture.reader.nodeBatches) + } + if fixture.reader.legacyInBatches != 0 || fixture.reader.legacyOutBatches != 0 { + t.Fatalf("legacy adjacency calls = in:%d out:%d, want 0/0", + fixture.reader.legacyInBatches, fixture.reader.legacyOutBatches) } } @@ -496,6 +597,9 @@ func BenchmarkProjectExploreTypedAnchorCandidates(b *testing.B) { b.ReportAllocs() b.ResetTimer() for i := 0; i < b.N; i++ { + fixture.reader.nodeRequests = fixture.reader.nodeRequests[:0] + fixture.reader.incomingRequests = fixture.reader.incomingRequests[:0] + fixture.reader.outgoingRequests = fixture.reader.outgoingRequests[:0] got := projectExploreTypedAnchorCandidates( context.Background(), fixture.task, fixture.candidates, fixture.reader, scope, len(fixture.candidates), fixture.protected, "", diff --git a/internal/mcp/facade_tools.go b/internal/mcp/facade_tools.go index 3c0fa995..2e49d2f7 100644 --- a/internal/mcp/facade_tools.go +++ b/internal/mcp/facade_tools.go @@ -655,8 +655,18 @@ func publishLocalizationAuthReceipt(token string, result *mcpgo.CallToolResult) return } completion := host.Contract.Completion + evidenceIDs := make([]string, 0, localizationReplayEvidenceLimit) + if host.Evidence != nil { + for _, row := range host.Evidence.Evidence { + if id := strings.TrimSpace(row.ID); id != "" { + evidenceIDs = append(evidenceIDs, id) + } + } + } localizationauth.Publish(token, localizationauth.Receipt{ FinalResponse: completion.FinalResponse, + PrimaryIDs: append([]string(nil), host.PrimaryIDs...), + EvidenceIDs: evidenceIDs, ContractVersion: completion.ContractVersion, Enforceable: completion.Enforceable, }) @@ -703,18 +713,20 @@ func captureLocalizationSearchSymbols(ctx context.Context, nodes []*graph.Node) } // captureLocalizationSearchText promotes only graph-backed identities from the -// typed search.text page. A file-level literal may have no enclosing SymbolID; -// in that case the matching path and line are resolved back to the narrowest -// in-scope graph declaration, or finally to the graph's file node. Rendered MCP -// Content is never parsed as evidence. -func (s *Server) captureLocalizationSearchText(ctx context.Context, matches []enrichedTextMatch) { +// typed search.text page. It reuses the exact bounded file indexes built for +// enrichment, so a permitted-evidence capture never repeats a file scan. +func (s *Server) captureLocalizationSearchText( + ctx context.Context, + matches []enrichedTextMatch, + indexes map[string]*fileSymbolIndex, +) { capture, _ := ctx.Value(localizationPermittedEvidenceCaptureKey{}).(*localizationPermittedEvidenceCapture) if capture == nil { return } rows := make([]localizationDigestRow, 0, len(matches)) for _, match := range matches { - node, provenance := s.localizationTextMatchNode(ctx, match) + node, provenance := s.localizationTextMatchNode(ctx, match, indexes) if row, ok := localizationDigestRowFromNode(node, provenance); ok { if match.Line > 0 { row.Line = match.Line @@ -725,12 +737,22 @@ func (s *Server) captureLocalizationSearchText(ctx context.Context, matches []en captureLocalizationRows(ctx, rows) } -func (s *Server) localizationTextMatchNode(ctx context.Context, match enrichedTextMatch) (*graph.Node, string) { - if s == nil || s.graph == nil { +func (s *Server) localizationTextMatchNode( + ctx context.Context, + match enrichedTextMatch, + indexes map[string]*fileSymbolIndex, +) (*graph.Node, string) { + if s == nil { return nil, "" } + reader := s.readerFor(ctx) + if reader == nil { + return nil, "" + } + // Preserve the typed SymbolID path exactly: an already-enriched hit needs + // one identity lookup, not another file projection. if id := strings.TrimSpace(match.SymbolID); id != "" { - if node := s.graph.GetNode(id); node != nil && s.nodeInSessionScope(ctx, node) { + if node := reader.GetNode(id); node != nil && s.nodeInSessionScope(ctx, node) { return node, "permitted_search_text" } return nil, "" @@ -739,45 +761,15 @@ func (s *Server) localizationTextMatchNode(ctx context.Context, match enrichedTe if path == "" { return nil, "" } - var owner *graph.Node - var fileNode *graph.Node - ownerSpan := int(^uint(0) >> 1) - for _, node := range s.graph.GetFileNodes(path) { - if node == nil || !s.nodeInSessionScope(ctx, node) { - continue - } - if node.Kind == graph.KindFile { - if fileNode == nil || node.ID < fileNode.ID { - fileNode = node - } - continue - } - if match.Line <= 0 || node.StartLine <= 0 || node.StartLine > match.Line || !exploreLocalizableKind(node.Kind) { - continue - } - end := node.EndLine - if end <= 0 { - end = node.StartLine - } - if end < match.Line { - continue - } - span := end - node.StartLine - if owner == nil || span < ownerSpan || (span == ownerSpan && node.ID < owner.ID) { - owner = node - ownerSpan = span - } + index := fileSymbolIndexForPath(indexes, path) + if index == nil || index.saturated { + return nil, "" } - if owner != nil { + if owner := index.smallestEnclosing(match.Line); owner != nil && s.nodeInSessionScope(ctx, owner) { return owner, "permitted_search_text_owner" } - if fileNode == nil { - if node := s.graph.GetNode(path); node != nil && node.Kind == graph.KindFile && s.nodeInSessionScope(ctx, node) { - fileNode = node - } - } - if fileNode != nil { - return fileNode, "permitted_search_text_file" + if index.fileNode != nil && s.nodeInSessionScope(ctx, index.fileNode) { + return index.fileNode, "permitted_search_text_file" } return nil, "" } @@ -1102,15 +1094,16 @@ func (s *Server) invokeFacadeSpec(ctx context.Context, req mcpgo.CallToolRequest outcome = facadeOutcomeInvalidArgument return invalid, nil } - if OverlayViewFromContext(ctx) == nil && !facadeLegacyManagesOwnOverlay(spec.Legacy) { - view, viewErr := s.buildOverlayViewForCtx(ctx) + if !facadeLegacyManagesOwnOverlay(spec.Legacy) { + var viewErr error + ctx, _, viewErr = s.prepareOverlayRequest(ctx) if viewErr != nil { + if ctxErr := requestContextError(ctx, viewErr); ctxErr != nil { + return nil, ctxErr + } outcome = facadeOutcomeToolError return mcpgo.NewToolResultError(viewErr.Error()), nil } - if view != nil { - ctx = WithOverlayView(ctx, view) - } } forwarded := req forwarded.Params.Name = spec.Legacy diff --git a/internal/mcp/force_inject.go b/internal/mcp/force_inject.go index 29bfc2f0..0b381708 100644 --- a/internal/mcp/force_inject.go +++ b/internal/mcp/force_inject.go @@ -15,24 +15,14 @@ import ( // Cheap-gated on a fresh search so a file open with no recent search // pays nothing. func (s *Server) creditFileConsumption(ctx context.Context, filePath string) { - if s == nil || s.combo == nil || filePath == "" { + if s == nil || s.combo == nil || ctx == nil || ctx.Err() != nil || filePath == "" { return } sess := s.sessionFor(ctx) - if sess == nil || !sess.hasFreshSearch() { + if sess == nil || ctx.Err() != nil { return } - nodes := s.readerFor(ctx).GetFileNodes(filePath) - if len(nodes) == 0 { - return - } - ids := make([]string, 0, len(nodes)) - for _, n := range nodes { - if n != nil && n.ID != "" { - ids = append(ids, n.ID) - } - } - query, matched := sess.attributedConsumptionBatch(ids) + query, matched := sess.attributedFileConsumption(filePath) if query != "" && len(matched) > 0 { s.combo.RecordBatch(query, matched) } diff --git a/internal/mcp/force_inject_consumption_test.go b/internal/mcp/force_inject_consumption_test.go new file mode 100644 index 00000000..3e777076 --- /dev/null +++ b/internal/mcp/force_inject_consumption_test.go @@ -0,0 +1,65 @@ +package mcp + +import ( + "context" + "testing" + + "github.com/zzet/gortex/internal/graph" +) + +type panicFileNodeStore struct { + graph.Store +} + +func (panicFileNodeStore) GetFileNodes(string) []*graph.Node { + panic("creditFileConsumption must not read file nodes") +} + +func TestCreditFileConsumptionUsesOnlyReturnedPageState(t *testing.T) { + session := &sessionState{} + session.recordLastSearchPage("needle", []*graph.Node{ + {ID: "repo/a.go::A", FilePath: "repo/a.go"}, + {ID: "repo/a.go::B", FilePath: "repo/a.go"}, + {ID: "repo/b.go::C", FilePath: "repo/b.go"}, + }) + server := &Server{ + session: session, + combo: newComboManager("", "", ModeHuman), + graph: panicFileNodeStore{}, + } + + server.creditFileConsumption(context.Background(), " ./repo/a.go ") + server.creditFileConsumption(context.Background(), "repo/a.go") + + session.mu.Lock() + defer session.mu.Unlock() + _, gotA := session.lastSearch.consumed["repo/a.go::A"] + _, gotB := session.lastSearch.consumed["repo/a.go::B"] + _, gotOther := session.lastSearch.consumed["repo/b.go::C"] + if !gotA || !gotB { + t.Fatalf("matching page IDs were not credited once: %#v", session.lastSearch.consumed) + } + if gotOther { + t.Fatalf("other-file ID was credited: %#v", session.lastSearch.consumed) + } +} + +func TestCreditFileConsumptionCanceledContextIsNoOp(t *testing.T) { + session := &sessionState{} + session.recordLastSearchPage("needle", []*graph.Node{{ID: "repo/a.go::A", FilePath: "repo/a.go"}}) + server := &Server{ + session: session, + combo: newComboManager("", "", ModeHuman), + graph: panicFileNodeStore{}, + } + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + server.creditFileConsumption(ctx, "repo/a.go") + + session.mu.Lock() + defer session.mu.Unlock() + if len(session.lastSearch.consumed) != 0 { + t.Fatalf("canceled credit mutated state: %#v", session.lastSearch.consumed) + } +} diff --git a/internal/mcp/localization_body_mentions.go b/internal/mcp/localization_body_mentions.go new file mode 100644 index 00000000..4de92198 --- /dev/null +++ b/internal/mcp/localization_body_mentions.go @@ -0,0 +1,378 @@ +package mcp + +import ( + "sort" + "strings" + "unicode" + "unicode/utf8" + + "github.com/zzet/gortex/internal/graph" +) + +const ( + localizationBodyMentionCap = 8 + localizationBodyMentionSourceCap = 12 + localizationBodyMentionFileCap = 10 + localizationBodyMentionDeclarationCap = 256 + localizationBodyMentionCheckCap = 2048 + localizationBodyMentionTokenCap = 2048 + localizationBodyMentionTokenByteCap = 256 + localizationBodyMentionTotalSourceByteCap = 64 << 10 + localizationProvenanceBodyMention = "body_mention" + // A declaration the raw task names as a whole word is the caller's own + // subject, not a scan inference; it carries its own provenance and is not + // confined to supporting placement. + localizationProvenanceTaskMention = "task_mention" +) + +type localizationBodyMentionCandidate struct { + node *graph.Node + taskNamed bool + sameFile bool + overlap int + longest int + direct bool + ownerRank int + ownerIndex int +} + +type localizationBodyMentionSource struct { + owner localizationEvidence + ownerIndex int + file string + tokens map[string]struct{} +} + +// promoteLocalizationBodyMentions converts declarations already visible inside +// packed source into citeable SUPPORTING rows. Every identity comes from the +// request-local graph declaration cache; source text alone can never fabricate +// a row. Source bytes, files, declarations, comparisons, candidates, and output +// are all independently capped. A truncated scan remains supporting-only and +// therefore cannot establish or extend terminal authority. +func promoteLocalizationBodyMentions( + task string, + envelope localizationExploreEnvelope, + cache *localizationFileDeclarationCache, + maxBytes int, + digest *localizationEvidenceDigest, +) (localizationExploreEnvelope, *localizationEvidenceDigest) { + if cache == nil || cache.reader == nil || len(envelope.Evidence) == 0 { + return envelope, digest + } + seenIDs := make(map[string]struct{}, len(envelope.Evidence)) + for _, row := range envelope.Evidence { + if row.ID != "" { + seenIDs[row.ID] = struct{}{} + } + } + sources := localizationBodyMentionSources(envelope) + if len(sources) == 0 { + return envelope, digest + } + files := localizationBodyMentionFiles(envelope) + declarations := make(map[string][]*graph.Node, len(files)) + for _, file := range files { + declarations[file] = cache.boundedDefinitions(file, localizationBodyMentionDeclarationCap) + } + + taskTerms := exploreTerminalTerms(task) + candidates := make([]localizationBodyMentionCandidate, 0, localizationBodyMentionCap) + checks := 0 +scan: + for _, source := range sources { + for pass := 0; pass < 2; pass++ { + for _, file := range files { + sameFile := file == source.file + if sameFile != (pass == 0) { + continue + } + for _, node := range declarations[file] { + if !localizationBodyMentionNodeEligible(node) { + continue + } + if _, visible := seenIDs[node.ID]; visible { + continue + } + if checks >= localizationBodyMentionCheckCap { + break scan + } + checks++ + name, ok := localizationBodyMentionIdentifier(node.Name) + if !ok { + continue + } + _, mentioned := source.tokens[name] + taskNamed := localizationTaskNamesIdentifier( + task, strings.TrimSpace(node.Name), localizationTaskNameMinRunes) + if !mentioned && !taskNamed { + continue + } + overlap, longest := exploreDraftTermOverlap(taskTerms, node) + candidate := localizationBodyMentionCandidate{ + node: node, + taskNamed: taskNamed, + sameFile: sameFile, + overlap: overlap, + longest: longest, + direct: localizationBodyMentionDirect(source.owner, node.ID), + ownerRank: source.owner.Rank, + ownerIndex: source.ownerIndex, + } + candidates = localizationBodyMentionKeepBest(candidates, candidate) + } + } + } + } + + admitted := 0 + for _, mention := range candidates { + if admitted == localizationBodyMentionCap { + break + } + node := mention.node + row := localizationEvidence{ + Rank: len(envelope.Evidence) + 1, + ID: node.ID, + Name: compactLocalizationField(node.Name, localizationMaxNameRunes), + Kind: string(node.Kind), + File: nodeDisplayPath(node), + Line: node.StartLine, + Provenance: localizationProvenanceBodyMention, + + supportingOnly: true, + } + if mention.taskNamed { + row.Provenance = localizationProvenanceTaskMention + row.supportingOnly = false + } + candidate := envelope + // Files, Symbols, and Evidence are positional arrays. Repeated file + // paths are intentional: row N must align across all three arrays. + candidate.Files = append(append([]string(nil), envelope.Files...), row.File) + candidate.Symbols = append(append([]string(nil), envelope.Symbols...), row.ID) + candidate.Evidence = append(append([]localizationEvidence(nil), envelope.Evidence...), row) + + candidateDigest := newLocalizationEvidenceDigestForTask(task, candidate) + if !localizationBodyMentionDigestContains(candidateDigest, row.ID) { + continue + } + contract := localizationContractReconciledWithDigest(candidate.Completion, candidateDigest) + candidate.Completion = contract.Completion + candidate.Terminal = contract.Terminal + if !localizationEnvelopeFits(candidate, maxBytes) { + continue + } + envelope, digest = candidate, candidateDigest + seenIDs[row.ID] = struct{}{} + admitted++ + } + return envelope, digest +} + +func localizationBodyMentionSources(envelope localizationExploreEnvelope) []localizationBodyMentionSource { + sources := make([]localizationBodyMentionSource, 0, min(len(envelope.Evidence)+1, localizationBodyMentionSourceCap)) + remainingBytes := localizationBodyMentionTotalSourceByteCap + appendSource := func(owner localizationEvidence, ownerIndex int, file, content string) { + if len(sources) >= localizationBodyMentionSourceCap || remainingBytes <= 0 || strings.TrimSpace(content) == "" { + return + } + content = localizationBodyMentionBoundedSource(content, remainingBytes) + tokens := localizationBodyMentionIdentifiers(content) + if len(tokens) == 0 { + return + } + sources = append(sources, localizationBodyMentionSource{ + owner: owner, ownerIndex: ownerIndex, file: strings.TrimSpace(file), tokens: tokens, + }) + remainingBytes -= len(content) + } + // The elected source window is exact request evidence and would otherwise + // be invisible to this pass. Scan it first so ordinary bodies cannot consume + // the bounded source allowance before the window-only sibling is considered. + if window := envelope.SourceWindow; window != nil { + for index, owner := range envelope.Evidence { + if owner.ID != window.AnchorSymbol { + continue + } + appendSource(owner, index, window.Path, window.Content) + break + } + } + for index, owner := range envelope.Evidence { + appendSource(owner, index, owner.File, owner.Source) + } + return sources +} + +func localizationBodyMentionFiles(envelope localizationExploreEnvelope) []string { + files := make([]string, 0, localizationBodyMentionFileCap) + seen := make(map[string]struct{}, localizationBodyMentionFileCap) + appendFile := func(file string) { + file = strings.TrimSpace(file) + if file == "" || len(files) >= localizationBodyMentionFileCap { + return + } + if _, duplicate := seen[file]; duplicate { + return + } + seen[file] = struct{}{} + files = append(files, file) + } + if envelope.SourceWindow != nil { + appendFile(envelope.SourceWindow.Path) + } + for _, row := range envelope.Evidence { + appendFile(row.File) + } + return files +} + +func localizationBodyMentionBoundedSource(source string, limit int) string { + if limit <= 0 { + return "" + } + if len(source) <= limit { + return source + } + cut := limit + for cut > 0 && !utf8.ValidString(source[:cut]) { + cut-- + } + return source[:cut] +} + +func localizationBodyMentionIdentifiers(source string) map[string]struct{} { + identifiers := make(map[string]struct{}, min(localizationBodyMentionTokenCap, 64)) + start := -1 + appendIdentifier := func(end int) { + if start < 0 || end <= start || len(identifiers) >= localizationBodyMentionTokenCap { + start = -1 + return + } + raw := source[start:end] + start = -1 + if len(raw) > localizationBodyMentionTokenByteCap { + return + } + identifiers[strings.ToLower(raw)] = struct{}{} + } + for index, r := range source { + if unicode.IsLetter(r) || unicode.IsDigit(r) || r == '_' || r == '$' { + if start < 0 { + start = index + } + continue + } + appendIdentifier(index) + if len(identifiers) >= localizationBodyMentionTokenCap { + break + } + } + if len(identifiers) < localizationBodyMentionTokenCap { + appendIdentifier(len(source)) + } + return identifiers +} + +func localizationBodyMentionIdentifier(name string) (string, bool) { + name = strings.TrimSpace(name) + if name == "" || len(name) > localizationBodyMentionTokenByteCap { + return "", false + } + for _, r := range name { + if !unicode.IsLetter(r) && !unicode.IsDigit(r) && r != '_' && r != '$' { + return "", false + } + } + return strings.ToLower(name), true +} + +func localizationBodyMentionKeepBest( + candidates []localizationBodyMentionCandidate, + candidate localizationBodyMentionCandidate, +) []localizationBodyMentionCandidate { + for index := range candidates { + if candidates[index].node.ID != candidate.node.ID { + continue + } + if localizationBodyMentionLess(candidate, candidates[index]) { + candidates[index] = candidate + } + sort.SliceStable(candidates, func(first, second int) bool { + return localizationBodyMentionLess(candidates[first], candidates[second]) + }) + return candidates + } + candidates = append(candidates, candidate) + sort.SliceStable(candidates, func(first, second int) bool { + return localizationBodyMentionLess(candidates[first], candidates[second]) + }) + if len(candidates) > localizationBodyMentionCap { + candidates = candidates[:localizationBodyMentionCap] + } + return candidates +} + +func localizationBodyMentionNodeEligible(node *graph.Node) bool { + return node != nil && node.ID != "" && strings.TrimSpace(node.Name) != "" && + nodeDisplayPath(node) != "" && node.StartLine > 0 && + exploreLocalizableKind(node.Kind) && !isNonDefinitionNode(node.Kind) +} + +func localizationBodyMentionDirect(owner localizationEvidence, id string) bool { + for _, neighbor := range owner.Callers { + if neighbor == id { + return true + } + } + for _, neighbor := range owner.Callees { + if neighbor == id { + return true + } + } + return false +} + +func localizationBodyMentionLess(left, right localizationBodyMentionCandidate) bool { + if left.taskNamed != right.taskNamed { + return left.taskNamed + } + if left.sameFile != right.sameFile { + return left.sameFile + } + if left.overlap != right.overlap { + return left.overlap > right.overlap + } + if left.longest != right.longest { + return left.longest > right.longest + } + if left.direct != right.direct { + return left.direct + } + if left.ownerRank != right.ownerRank { + return left.ownerRank < right.ownerRank + } + if left.ownerIndex != right.ownerIndex { + return left.ownerIndex < right.ownerIndex + } + if left.node.StartLine != right.node.StartLine { + return left.node.StartLine < right.node.StartLine + } + return left.node.ID < right.node.ID +} + +func localizationBodyMentionDigestContains(digest *localizationEvidenceDigest, id string) bool { + if digest == nil || id == "" { + return false + } + for _, row := range digest.Evidence { + if row.ID != id { + continue + } + if row.Provenance == localizationProvenanceBodyMention || + row.Provenance == localizationProvenanceTaskMention { + return true + } + } + return false +} diff --git a/internal/mcp/localization_body_mentions_test.go b/internal/mcp/localization_body_mentions_test.go new file mode 100644 index 00000000..23943a07 --- /dev/null +++ b/internal/mcp/localization_body_mentions_test.go @@ -0,0 +1,316 @@ +package mcp + +import ( + "encoding/json" + "fmt" + "strings" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/zzet/gortex/internal/graph" +) + +func localizationBodyMentionTestNode(file, name string, line int) *graph.Node { + return &graph.Node{ + ID: file + "::" + name, + Name: name, + Kind: graph.KindFunction, + FilePath: file, + StartLine: line, + } +} + +func TestPromoteLocalizationBodyMentionsUsesRealPageDeclarations(t *testing.T) { + sibling := localizationBodyMentionTestNode("src/a.go", "sibling", 22) + fragment := localizationBodyMentionTestNode("src/a.go", "Foo", 30) + external := localizationBodyMentionTestNode("src/b.go", "External", 7) + reader := &localizationDeclarationSpyReader{ + files: map[string][]*graph.Node{ + "src/a.go": {sibling, fragment}, + "src/b.go": {external}, + }, + fileCalls: make(map[string]int), + } + envelope := localizationExploreEnvelope{ + Completion: newLocalizationCompletion(true, ""), + Files: []string{"src/a.go", "src/b.go"}, + Symbols: []string{"src/a.go::owner", "src/b.go::page"}, + Evidence: []localizationEvidence{ + { + Rank: 1, ID: "src/a.go::owner", Name: "owner", Kind: string(graph.KindFunction), + File: "src/a.go", Line: 1, Source: "func owner() { sibling(); External(); foobar() }", + Callees: []string{external.ID}, + }, + {Rank: 2, ID: "src/b.go::page", Name: "page", Kind: string(graph.KindFunction), File: "src/b.go", Line: 1}, + }, + } + digest := newLocalizationEvidenceDigestForTask("fix sibling and External", envelope) + + packed, packedDigest := promoteLocalizationBodyMentions( + "fix sibling and External", envelope, newLocalizationDeclarationTestCache(reader), 1<<20, digest, + ) + + require.Equal(t, []string{sibling.ID, external.ID}, []string{packed.Evidence[2].ID, packed.Evidence[3].ID}) + require.Equal(t, localizationProvenanceTaskMention, packed.Evidence[2].Provenance) + require.Equal(t, localizationProvenanceTaskMention, packed.Evidence[3].Provenance) + require.Empty(t, packed.Evidence[2].Source) + require.NotContains(t, packed.Symbols, fragment.ID, "identifier fragments must not promote declarations") + require.Equal(t, map[string]int{"src/a.go": 1, "src/b.go": 1}, reader.fileCalls) + require.Len(t, packed.Files, len(packed.Evidence)) + require.Len(t, packed.Symbols, len(packed.Evidence)) + require.True(t, localizationBodyMentionDigestContains(packedDigest, sibling.ID)) + require.True(t, localizationBodyMentionDigestContains(packedDigest, external.ID)) + + seatedMentions := 0 + for _, row := range localizationFinalResponseRows("fix sibling and External", nil, packedDigest.Evidence) { + if row.row.Provenance == localizationProvenanceTaskMention { + require.True(t, row.primary, "a body-derived declaration the task names seats as primary") + seatedMentions++ + } + } + require.Equal(t, 2, seatedMentions) +} + +func TestPromoteLocalizationBodyMentionsIsCappedAndCached(t *testing.T) { + declarations := make([]*graph.Node, 0, localizationBodyMentionCap+4) + var source strings.Builder + for index := 0; index < localizationBodyMentionCap+4; index++ { + name := fmt.Sprintf("Call%d", index) + declarations = append(declarations, localizationBodyMentionTestNode("src/a.go", name, index+10)) + fmt.Fprintf(&source, "%s(); ", name) + } + reader := &localizationDeclarationSpyReader{ + files: map[string][]*graph.Node{"src/a.go": declarations}, + fileCalls: make(map[string]int), + } + envelope := localizationExploreEnvelope{ + Completion: newLocalizationCompletion(true, ""), + Files: []string{"src/a.go"}, + Symbols: []string{"src/a.go::owner"}, + Evidence: []localizationEvidence{{ + Rank: 1, ID: "src/a.go::owner", Name: "owner", Kind: string(graph.KindFunction), + File: "src/a.go", Line: 1, Source: source.String(), + }}, + } + + packed, _ := promoteLocalizationBodyMentions( + "trace the calls", envelope, newLocalizationDeclarationTestCache(reader), 1<<20, + newLocalizationEvidenceDigestForTask("trace the calls", envelope), + ) + + require.Len(t, packed.Evidence, 1+localizationBodyMentionCap) + require.Equal(t, 1, reader.fileCalls["src/a.go"]) + for index, row := range packed.Evidence { + require.Equal(t, index+1, row.Rank) + } +} + +func TestPromoteLocalizationBodyMentionsUsesPackedSourceWindow(t *testing.T) { + sibling := localizationBodyMentionTestNode("src/a.go", "sibling", 22) + reader := &localizationDeclarationSpyReader{ + files: map[string][]*graph.Node{"src/a.go": {sibling}}, + fileCalls: make(map[string]int), + } + envelope := localizationExploreEnvelope{ + Completion: newLocalizationCompletion(true, ""), + Files: []string{"src/a.go"}, + Symbols: []string{"src/a.go::owner"}, + Evidence: []localizationEvidence{{ + Rank: 1, ID: "src/a.go::owner", Name: "owner", Kind: string(graph.KindFunction), + File: "src/a.go", Line: 1, + }}, + SourceWindow: &localizationSourceWindow{ + Path: "src/a.go", StartLine: 20, EndLine: 24, MatchLine: 21, + Content: "func call() { sibling() }", AnchorSymbol: "src/a.go::owner", + }, + } + + packed, packedDigest := promoteLocalizationBodyMentions( + "find sibling", envelope, newLocalizationDeclarationTestCache(reader), 1<<20, + newLocalizationEvidenceDigestForTask("find sibling", envelope), + ) + + require.Contains(t, packed.Symbols, sibling.ID) + require.True(t, localizationBodyMentionDigestContains(packedDigest, sibling.ID)) + for _, row := range localizationFinalResponseRows("find sibling", nil, packedDigest.Evidence) { + if row.row.ID == sibling.ID { + require.True(t, row.primary, "a window-derived row the task names seats as primary") + } + } +} + +func TestPromoteLocalizationBodyMentionsBoundsDeclarationsPerFile(t *testing.T) { + inside := localizationBodyMentionTestNode("src/generated.go", "Inside", 1) + outside := localizationBodyMentionTestNode("src/generated.go", "Outside", localizationBodyMentionDeclarationCap+1) + declarations := make([]*graph.Node, 0, localizationBodyMentionDeclarationCap+1) + declarations = append(declarations, inside) + for index := 1; index < localizationBodyMentionDeclarationCap; index++ { + declarations = append(declarations, localizationBodyMentionTestNode("src/generated.go", fmt.Sprintf("Filler%d", index), index+1)) + } + declarations = append(declarations, outside) + reader := &localizationDeclarationSpyReader{ + files: map[string][]*graph.Node{"src/generated.go": declarations}, + fileCalls: make(map[string]int), + } + envelope := localizationExploreEnvelope{ + Completion: newLocalizationCompletion(true, ""), + Files: []string{"src/generated.go"}, + Symbols: []string{"src/generated.go::owner"}, + Evidence: []localizationEvidence{{ + Rank: 1, ID: "src/generated.go::owner", Name: "owner", Kind: string(graph.KindFunction), + File: "src/generated.go", Line: 1, Source: "Inside(); Outside()", + }}, + } + + packed, _ := promoteLocalizationBodyMentions( + "find Inside and Outside", envelope, newLocalizationDeclarationTestCache(reader), 1<<20, + newLocalizationEvidenceDigestForTask("find Inside and Outside", envelope), + ) + + require.Contains(t, packed.Symbols, inside.ID) + require.NotContains(t, packed.Symbols, outside.ID) + require.Equal(t, 1, reader.fileCalls["src/generated.go"]) + require.Len(t, reader.calls, 1) + require.Equal(t, localizationBodyMentionDeclarationCap, reader.calls[0].limit) +} + +func TestPromoteLocalizationBodyMentionsHonorsEnvelopeBudget(t *testing.T) { + sibling := localizationBodyMentionTestNode("src/a.go", "sibling", 22) + reader := &localizationDeclarationSpyReader{ + files: map[string][]*graph.Node{"src/a.go": {sibling}}, + fileCalls: make(map[string]int), + } + envelope := localizationExploreEnvelope{ + Completion: newLocalizationCompletion(true, ""), + Files: []string{"src/a.go"}, + Symbols: []string{"src/a.go::owner"}, + Evidence: []localizationEvidence{{ + Rank: 1, ID: "src/a.go::owner", Name: "owner", Kind: string(graph.KindFunction), + File: "src/a.go", Line: 1, Source: "func owner() { sibling() }", + }}, + } + body, err := json.Marshal(envelope) + require.NoError(t, err) + digest := newLocalizationEvidenceDigestForTask("find sibling", envelope) + + packed, packedDigest := promoteLocalizationBodyMentions( + "find sibling", envelope, newLocalizationDeclarationTestCache(reader), len(body), digest, + ) + + require.Equal(t, envelope, packed) + require.Same(t, digest, packedDigest) +} + +func TestPromoteLocalizationBodyMentionsBoundsTheTenFileProjection(t *testing.T) { + reader := &localizationDeclarationSpyReader{files: make(map[string][]*graph.Node)} + envelope := localizationExploreEnvelope{Completion: newLocalizationCompletion(true, "")} + for index := 0; index < localizationBodyMentionFileCap; index++ { + file := fmt.Sprintf("src/file-%02d.go", index) + ownerID := file + "::owner" + reader.files[file] = []*graph.Node{ + localizationBodyMentionTestNode(file, "Mentioned", index+2), + } + envelope.Files = append(envelope.Files, file) + envelope.Symbols = append(envelope.Symbols, ownerID) + envelope.Evidence = append(envelope.Evidence, localizationEvidence{ + Rank: index + 1, ID: ownerID, Name: "owner", Kind: string(graph.KindFunction), + File: file, Line: 1, Source: "func owner() { Mentioned() }", + }) + } + + promoteLocalizationBodyMentions( + "find Mentioned", envelope, newLocalizationDeclarationTestCache(reader), 1<<20, + newLocalizationEvidenceDigestForTask("find Mentioned", envelope), + ) + + require.Len(t, reader.calls, localizationBodyMentionFileCap) + consumed := 0 + for _, call := range reader.calls { + require.LessOrEqual(t, call.limit, localizationBodyMentionDeclarationCap) + consumed += call.limit + } + require.LessOrEqual(t, consumed, localizationFileRequestLimit) +} + +func TestLocalizationDigestReservesBodyMentionTail(t *testing.T) { + evidence := make([]localizationEvidence, 0, localizationReplayEvidenceLimit+1) + for index := 0; index < localizationReplayEvidenceLimit; index++ { + evidence = append(evidence, localizationEvidence{ + Rank: index + 1, ID: fmt.Sprintf("src/a.go::ordinary%d", index), + Name: fmt.Sprintf("ordinary%d", index), Kind: string(graph.KindFunction), File: "src/a.go", Line: index + 1, + }) + } + bodyID := "src/a.go::mentioned" + evidence = append(evidence, localizationEvidence{ + Rank: len(evidence) + 1, ID: bodyID, Name: "mentioned", Kind: string(graph.KindFunction), + File: "src/a.go", Line: 99, Provenance: localizationProvenanceBodyMention, + }) + + digest := newLocalizationEvidenceDigestForTask("find mentioned", localizationExploreEnvelope{Evidence: evidence}) + + require.Len(t, digest.Evidence, localizationReplayEvidenceLimit) + require.True(t, localizationBodyMentionDigestContains(digest, bodyID)) + require.NotContains(t, digest.Symbols, "src/a.go::ordinary15") +} + +func TestShedLocalizationDigestSupportingOnlyBeforeOrdinaryRows(t *testing.T) { + rows := []localizationDigestRow{ + {ID: "ordinary", Provenance: "ranked"}, + {ID: "mentioned", Provenance: localizationProvenanceBodyMention, supportingOnly: true}, + {ID: "adjacent", Provenance: localizationProvenanceDirectAdjacency, supportingOnly: true}, + {ID: "protected", Provenance: localizationProvenanceDirectAdjacency, supportingOnly: true, authorizationPriority: true}, + {ID: "tail", Provenance: "ranked"}, + } + + retained, removed := shedLocalizationDigestSupportingOnly(rows) + require.True(t, removed) + require.Equal(t, []string{"ordinary", "mentioned", "protected", "tail"}, []string{retained[0].ID, retained[1].ID, retained[2].ID, retained[3].ID}) + + retained, removed = shedLocalizationDigestSupportingOnly(retained) + require.True(t, removed) + require.Equal(t, []string{"ordinary", "protected", "tail"}, []string{retained[0].ID, retained[1].ID, retained[2].ID}) + + retained, removed = shedLocalizationDigestSupportingOnly(retained) + require.False(t, removed) + require.Equal(t, []string{"ordinary", "protected", "tail"}, []string{retained[0].ID, retained[1].ID, retained[2].ID}) +} + +func TestPromoteLocalizationTaskNamedDeclarationWithoutBodyMention(t *testing.T) { + target := localizationBodyMentionTestNode("src/a.go", "RotateJournal", 40) + reader := &localizationDeclarationSpyReader{ + files: map[string][]*graph.Node{"src/a.go": {target}}, + fileCalls: make(map[string]int), + } + envelope := localizationExploreEnvelope{ + Completion: newLocalizationCompletion(true, ""), + Files: []string{"src/a.go"}, + Symbols: []string{"src/a.go::owner"}, + Evidence: []localizationEvidence{{ + Rank: 1, ID: "src/a.go::owner", Name: "owner", Kind: string(graph.KindFunction), + File: "src/a.go", Line: 1, Source: "func owner() { unrelated() }", + }}, + } + task := "RotateJournal drops entries during shutdown" + + packed, packedDigest := promoteLocalizationBodyMentions( + task, envelope, newLocalizationDeclarationTestCache(reader), 1<<20, + newLocalizationEvidenceDigestForTask(task, envelope), + ) + + require.Contains(t, packed.Symbols, target.ID) + var provenance string + for _, row := range packed.Evidence { + if row.ID == target.ID { + provenance = row.Provenance + } + } + require.Equal(t, localizationProvenanceTaskMention, provenance) + seated := false + for _, row := range localizationFinalResponseRows(task, nil, packedDigest.Evidence) { + if row.row.ID == target.ID { + seated = row.primary + } + } + require.True(t, seated, "a task-named declaration seats as primary without a body mention") +} diff --git a/internal/mcp/localization_digest.go b/internal/mcp/localization_digest.go index 8c41b1c0..f01b9397 100644 --- a/internal/mcp/localization_digest.go +++ b/internal/mcp/localization_digest.go @@ -3,7 +3,10 @@ package mcp import ( "encoding/json" "fmt" + "sort" "strings" + "unicode" + "unicode/utf8" mcpgo "github.com/mark3labs/mcp-go/mcp" ) @@ -49,6 +52,10 @@ type localizationEvidenceDigest struct { // session that ends holding nothing is the one outcome with no recovery, so // every state carries a page — labelled for what it is. provisionalResponse string + // primaryIDs is the exact bounded PRIMARY projection used to render the + // final response. It stays out of retained digest JSON to avoid counting the + // same identities twice, but is copied into authenticated host authority. + primaryIDs []string } type localizationDigestRow struct { @@ -63,6 +70,17 @@ type localizationDigestRow struct { Callers []string `json:"callers,omitempty"` Callees []string `json:"callees,omitempty"` Provenance string `json:"provenance,omitempty"` + + // Kept only in the request-local/session-retained projection. They separate + // truthful provenance from presentation authority without spending response + // bytes on internal policy state. primaryCohortOrder freezes the ranked + // PRIMARY projection before supplemental evidence is materialized. + literalPrimaryEligible bool + taskCitedPrimaryEligible bool + primaryCohortOrder int + supportingOnly bool + leadingFileDepth bool + authorizationPriority bool } // newLocalizationEvidenceDigestForTask retains only concrete ranked evidence @@ -77,39 +95,75 @@ func newLocalizationEvidenceDigestForTask(task string, envelope localizationExpl priorityIDs := localizationDigestPriorityIDs(envelope.Completion, envelope.Evidence) seen := make(map[string]struct{}, localizationReplayEvidenceLimit) - appendRows := func(priority bool) { - for _, row := range envelope.Evidence { - if len(digest.Evidence) >= localizationReplayEvidenceLimit { - return - } - if row.ID == "" || row.File == "" { - continue - } - _, prioritized := priorityIDs[row.ID] - if prioritized != priority { - continue - } - if _, exists := seen[row.ID]; exists { - continue + appendRow := func(row localizationEvidence) bool { + if row.ID == "" || row.File == "" || len(digest.Evidence) >= localizationReplayEvidenceLimit { + return false + } + if _, exists := seen[row.ID]; exists { + return false + } + seen[row.ID] = struct{}{} + _, authorizationPriority := priorityIDs[row.ID] + digest.Evidence = append(digest.Evidence, localizationDigestRow{ + Rank: row.Rank, + ID: row.ID, + Name: row.Name, + QualName: row.QualName, + Kind: row.Kind, + File: row.File, + Line: row.Line, + Signature: row.Signature, + Callers: append([]string(nil), row.Callers...), + Callees: append([]string(nil), row.Callees...), + Provenance: row.Provenance, + + literalPrimaryEligible: row.literalPrimaryEligible, + taskCitedPrimaryEligible: row.taskCitedPrimaryEligible, + primaryCohortOrder: row.primaryCohortOrder, + supportingOnly: row.supportingOnly || localizationSupportingOnlyProvenance(row.Provenance), + leadingFileDepth: row.leadingFileDepth, + authorizationPriority: authorizationPriority, + }) + return true + } + // Proof/authorization identities and the one task-cited supplemental + // PRIMARY remain the immutable prefix. The latter is presentation evidence, + // not authorization, so it is reserved without entering priorityIDs. + for _, row := range envelope.Evidence { + _, prioritized := priorityIDs[row.ID] + if prioritized || localizationEvidenceTaskCitedPrimary(row) { + appendRow(row) + } + } + // Reserve tail capacity for real body-derived declarations before ordinary + // non-priority rows fill the replay window. Visible envelope rows are never + // removed; only the weakest retained digest tail yields. + bodyIDs := make(map[string]struct{}, localizationBodyMentionCap) + for _, row := range envelope.Evidence { + if row.Provenance == localizationProvenanceBodyMention && row.ID != "" { + if _, already := seen[row.ID]; !already { + bodyIDs[row.ID] = struct{}{} } - seen[row.ID] = struct{}{} - digest.Evidence = append(digest.Evidence, localizationDigestRow{ - Rank: row.Rank, - ID: row.ID, - Name: row.Name, - QualName: row.QualName, - Kind: row.Kind, - File: row.File, - Line: row.Line, - Signature: row.Signature, - Callers: append([]string(nil), row.Callers...), - Callees: append([]string(nil), row.Callees...), - Provenance: row.Provenance, - }) } } - appendRows(true) - appendRows(false) + bodyReserve := min(len(bodyIDs), localizationReplayEvidenceLimit-len(digest.Evidence)) + ordinaryLimit := localizationReplayEvidenceLimit - bodyReserve + for _, row := range envelope.Evidence { + if len(digest.Evidence) >= ordinaryLimit { + break + } + if localizationSupportingOnlyProvenance(row.Provenance) { + continue + } + if _, prioritized := priorityIDs[row.ID]; !prioritized && !localizationEvidenceTaskCitedPrimary(row) { + appendRow(row) + } + } + for _, row := range envelope.Evidence { + if localizationSupportingOnlyProvenance(row.Provenance) { + appendRow(row) + } + } for { rebuildLocalizationDigestSkeleton(digest) @@ -121,6 +175,10 @@ func newLocalizationEvidenceDigestForTask(task string, envelope localizationExpl if len(digest.Evidence) == 0 { return digest } + if retained, removed := shedLocalizationDigestSupportingOnly(digest.Evidence); removed { + digest.Evidence = retained + continue + } last := len(digest.Evidence) - 1 if shedLocalizationDigestOptionalFields(digest.Evidence) { continue @@ -206,6 +264,32 @@ func mergeLocalizationDigestStrings(primary, supplementary []string) []string { return merged } +// localizationProvenanceMergeStrength orders provenance labels for merge +// retention. A literal observation is a claim a later graph hop cannot +// substitute: once a row has been seen as a content or source literal, a +// re-observation through adjacency must not erase that mark. +func localizationProvenanceMergeStrength(provenance string) int { + switch provenance { + case localizationProvenanceSourceLiteralCallee: + return 6 + case localizationProvenanceContentLiteral: + return 5 + case localizationProvenanceDivergentDefault, + localizationProvenanceImplementationTarget, + localizationProvenanceTypedAnchorProjection, + localizationProvenanceDivergentDefaultType, + localizationProvenancePermittedReadSource, + localizationProvenanceImplementationRoute: + return 4 + case localizationProvenanceDirectAdjacency, "direct_caller", "direct_callee": + return 2 + case "": + return 0 + default: + return 3 + } +} + // mergeLocalizationDigestRowEvidence preserves the first-ranked identity while // filling metadata and unioning bounded graph evidence from a later observation. // File disagreement is deliberately not repaired here: one ID cannot authorize @@ -229,9 +313,21 @@ func mergeLocalizationDigestRowEvidence(primary, supplementary localizationDiges if primary.Signature == "" { primary.Signature = supplementary.Signature } - if primary.Provenance == "" { + if localizationProvenanceMergeStrength(supplementary.Provenance) > + localizationProvenanceMergeStrength(primary.Provenance) { primary.Provenance = supplementary.Provenance } + primary.literalPrimaryEligible = primary.literalPrimaryEligible || supplementary.literalPrimaryEligible + primary.taskCitedPrimaryEligible = primary.taskCitedPrimaryEligible || supplementary.taskCitedPrimaryEligible + primary.authorizationPriority = primary.authorizationPriority || supplementary.authorizationPriority + // An identity that was independently ranked remains cohort-eligible when a + // later supplemental observation of the same row is merged into it. + primary.supportingOnly = primary.supportingOnly && supplementary.supportingOnly + primary.leadingFileDepth = primary.leadingFileDepth || supplementary.leadingFileDepth + if primary.primaryCohortOrder == 0 || + (supplementary.primaryCohortOrder > 0 && supplementary.primaryCohortOrder < primary.primaryCohortOrder) { + primary.primaryCohortOrder = supplementary.primaryCohortOrder + } primary.Callers = mergeLocalizationDigestStrings(primary.Callers, supplementary.Callers) primary.Callees = mergeLocalizationDigestStrings(primary.Callees, supplementary.Callees) return primary @@ -294,6 +390,10 @@ func mergeLocalizationEvidenceDigest(current []localizationDigestRow, retained * if len(digest.Evidence) == 0 { return digest } + if retained, removed := shedLocalizationDigestSupportingOnly(digest.Evidence); removed { + digest.Evidence = retained + continue + } last := len(digest.Evidence) - 1 if shedLocalizationDigestOptionalFields(digest.Evidence) { continue @@ -385,6 +485,84 @@ func localizationTaskAwareRetainedRows(task string, current []localizationDigest return ordered } +func localizationIdentifierRune(r rune) bool { + return unicode.IsLetter(r) || unicode.IsDigit(r) || r == '_' || r == '$' +} + +func localizationTaskExactIdentifierOffset(task, value string) int { + value = strings.TrimSpace(value) + if value == "" || len(task) < len(value) { + return -1 + } + quoted := false + for _, quote := range []string{"`", "'", "\""} { + if strings.Contains(task, quote+value+quote) { + quoted = true + break + } + } + if !localizationRecoveryConcreteIdentifier(value) && + !exploreQueryHasCallAnchor(task, value) && !quoted { + return -1 + } + for offset := 0; offset <= len(task)-len(value); { + relative := strings.Index(task[offset:], value) + if relative < 0 { + return -1 + } + start := offset + relative + end := start + len(value) + leftBoundary := start == 0 + if !leftBoundary { + r, _ := utf8.DecodeLastRuneInString(task[:start]) + leftBoundary = !localizationIdentifierRune(r) + } + rightBoundary := end == len(task) + if !rightBoundary { + r, _ := utf8.DecodeRuneInString(task[end:]) + rightBoundary = !localizationIdentifierRune(r) + } + if leftBoundary && rightBoundary { + return start + } + offset = start + 1 + } + return -1 +} + +func localizationTaskCitesExactIdentifier(task, value string) bool { + return localizationTaskExactIdentifierOffset(task, value) >= 0 +} + +func localizationEvidenceTaskCitationOffset(task string, row localizationEvidence) int { + best := -1 + for _, value := range []string{row.ID, row.Name, row.QualName} { + offset := localizationTaskExactIdentifierOffset(task, value) + if offset >= 0 && (best < 0 || offset < best) { + best = offset + } + } + return best +} + +func localizationEvidenceTaskCited(task string, row localizationEvidence) bool { + return localizationEvidenceTaskCitationOffset(task, row) >= 0 +} + +func localizationEvidenceTaskCitedPrimary(row localizationEvidence) bool { + return row.taskCitedPrimaryEligible && row.primaryCohortOrder > 0 && + row.Provenance == localizationProvenanceDirectAdjacency +} + +func localizationDigestRowIdentifierTaskCited(task string, row localizationDigestRow) bool { + for _, value := range []string{row.ID, row.Name, row.QualName} { + if localizationTaskCitesExactIdentifier(task, value) { + return true + } + } + return false +} + func localizationDigestRowTaskCited(task string, row localizationDigestRow) bool { for _, value := range []string{row.ID, row.Name, row.QualName, row.File, row.Signature} { if localizationTaskCitesConcreteEvidence(task, value) { @@ -471,6 +649,22 @@ func shedLocalizationDigestOptionalFields(rows []localizationDigestRow) bool { return false } +// Supplemental rows yield before richer metadata or independently ranked +// evidence under retained-state pressure, regardless of where a later merge +// placed them. An identity required by the live completion is protected even +// when its only visible observation has supplemental provenance. +func shedLocalizationDigestSupportingOnly(rows []localizationDigestRow) ([]localizationDigestRow, bool) { + for index := len(rows) - 1; index >= 0; index-- { + if !localizationFinalResponseSupportingOnly(rows[index]) || rows[index].authorizationPriority { + continue + } + copy(rows[index:], rows[index+1:]) + rows[len(rows)-1] = localizationDigestRow{} + return rows[:len(rows)-1], true + } + return rows, false +} + func rebuildLocalizationDigestSkeleton(digest *localizationEvidenceDigest) { digest.Files = digest.Files[:0] digest.Symbols = digest.Symbols[:0] @@ -494,6 +688,21 @@ const ( // usually within the top five. Slots below stay SUPPORTING so the page // keeps an explicit confidence order. localizationFinalResponsePrimaryLimit = 5 + // Authenticated literal rows may pre-seat ahead of the ranked cohort. The + // second seat is granted only for a second proven file, so a literal + // registered in one place and consumed in another presents both sites + // while same-file corroboration still competes through ordinary ranking. + localizationFinalResponseLiteralReserve = 2 + // A row is task-named only when the raw task text carries its identifier + // as a whole word, case-sensitively. The rune floors keep short prose + // words from naming whatever row happens to share them: five for the + // row's own name, six for an owner segment. + localizationTaskNameMinRunes = 5 + localizationTaskOwnerMinRunes = 6 + // Task-named rows never displace a ranked seat. They may extend the + // primary block past its ranked width by at most this many extra seats: + // a proven ranked conversion is never traded for a hoped-for one. + localizationTaskAlignedExtraSeats = 2 // localizationDigestShrinkFloorRows is the smallest answer the envelope // shed loop may shrink the digest to under extreme budget pressure. It is // deliberately narrower than the primary block: the primary width is what @@ -521,8 +730,10 @@ type localizationFinalResponseTaskScore struct { callable bool } -func localizationFinalResponsePrimaryProvenance(provenance string) bool { - switch provenance { +func localizationFinalResponsePrimaryProvenance(row localizationDigestRow) bool { + switch row.Provenance { + case localizationProvenanceContentLiteral: + return row.literalPrimaryEligible case localizationProvenanceSourceLiteralCallee, localizationProvenanceDivergentDefault, localizationProvenanceImplementationTarget, @@ -534,10 +745,19 @@ func localizationFinalResponsePrimaryProvenance(provenance string) bool { } } +func localizationFinalResponseSupportingOnly(row localizationDigestRow) bool { + if row.taskCitedPrimaryEligible && row.primaryCohortOrder > 0 && + row.Provenance == localizationProvenanceDirectAdjacency { + return false + } + return row.supportingOnly || localizationSupportingOnlyProvenance(row.Provenance) +} + func localizationFinalResponseSupportingProvenance(provenance string) bool { switch provenance { case localizationProvenanceDivergentDefaultType, localizationProvenanceImplementationRoute, + localizationProvenanceDirectAdjacency, "direct_caller", "direct_callee": return true default: @@ -601,6 +821,82 @@ func localizationFinalResponseBareName(row localizationDigestRow) string { return name } +func localizationIdentifierByte(b byte) bool { + return b == '_' || b == '$' || + ('0' <= b && b <= '9') || ('a' <= b && b <= 'z') || ('A' <= b && b <= 'Z') +} + +// localizationTaskNamesIdentifier reports whether the raw task text contains +// identifier as a whole word, case-sensitively. Substring hits inside a longer +// identifier do not count: a task about "starting" does not name "start". +func localizationTaskNamesIdentifier(task, identifier string, minRunes int) bool { + if identifier == "" || task == "" || utf8.RuneCountInString(identifier) < minRunes { + return false + } + for offset := 0; offset+len(identifier) <= len(task); { + index := strings.Index(task[offset:], identifier) + if index < 0 { + return false + } + start := offset + index + end := start + len(identifier) + if (start == 0 || !localizationIdentifierByte(task[start-1])) && + (end == len(task) || !localizationIdentifierByte(task[end])) { + return true + } + offset = start + 1 + } + return false +} + +// localizationTaskNamesRow reports whether the task names the row itself or +// the type the row is declared on. Only the immediate declaring owner counts: +// qualified names carry namespace and package segments, and the repository's +// own name sits in almost every namespace chain and almost every task, so a +// wider match would name the whole page. +func localizationTaskNamesRow(task string, row localizationDigestRow) bool { + name := localizationFinalResponseBareName(row) + if localizationTaskNamesIdentifier(task, name, localizationTaskNameMinRunes) { + return true + } + segments := strings.FieldsFunc(strings.TrimSpace(row.QualName), func(r rune) bool { + return r == '.' || r == ':' + }) + if len(segments) < 2 { + return false + } + owner := segments[len(segments)-2] + if owner == name || strings.EqualFold(owner, localizationDigestRowRepoIdentifier(row)) { + return false + } + return localizationTaskNamesIdentifier(task, owner, localizationTaskOwnerMinRunes) +} + +// localizationDigestRowRepoIdentifier extracts the repository name that +// prefixes the row's node ID, without any trailing checkout discriminator. +func localizationDigestRowRepoIdentifier(row localizationDigestRow) string { + id := strings.TrimSpace(row.ID) + slash := strings.IndexByte(id, '/') + if slash <= 0 { + return "" + } + repo := id[:slash] + if dash := strings.LastIndexByte(repo, '-'); dash > 0 { + digits := repo[dash+1:] + numeric := digits != "" + for _, r := range digits { + if r < '0' || r > '9' { + numeric = false + break + } + } + if numeric { + repo = repo[:dash] + } + } + return repo +} + // localizationFinalResponseAnchorRank orders the provenance flags that can name // the file the page is actually about. A divergent default owner is the // strongest such claim, a causal literal callee next, then a proven @@ -740,13 +1036,61 @@ func localizationFinalResponseRows(task string, current, rows []localizationDige for _, row := range rows { rowsByID[strings.TrimSpace(row.ID)] = row } - // Fresh rows already lead the merged evidence order. Primary status is earned - // from proof, task alignment, or owner coherence below; arrival time alone - // must not displace stronger retained evidence. + // Fresh rows already lead the merged evidence order. Once the initial ranked + // cohort has been frozen, later supplemental projections cannot reshuffle it: + // replay and permitted-read merges recover the exact original order first. + appendPrimary := func(row localizationDigestRow) bool { + if localizationFinalResponseSupportingOnly(row) { + return false + } + return appendRow(&primaries, localizationFinalResponsePrimaryLimit, row) + } + for order := 1; order <= localizationFinalResponsePrimaryLimit; order++ { + for _, row := range rows { + if row.primaryCohortOrder == order { + appendPrimary(row) + break + } + } + } + + // Authenticated literal rows reserve bounded seats before ordinary + // task/rank selection. The second seat must bring a second proven file: + // two owners of the same registration site add nothing over one. + literalReserveUsed := 0 + literalReserveFiles := make(map[string]struct{}, localizationFinalResponseLiteralReserve) + claimReserve := func(row localizationDigestRow) { + literalReserveUsed++ + if file := strings.TrimSpace(row.File); file != "" { + literalReserveFiles[file] = struct{}{} + } + } + for _, row := range primaries { + if row.Provenance == localizationProvenanceSourceLiteralCallee && row.literalPrimaryEligible { + claimReserve(row) + } + } for _, row := range rows { - if localizationFinalResponsePrimaryProvenance(row.Provenance) { - appendRow(&primaries, localizationFinalResponsePrimaryLimit, row) + if !localizationFinalResponsePrimaryProvenance(row) { + continue + } + literal := row.Provenance == localizationProvenanceContentLiteral || + row.Provenance == localizationProvenanceSourceLiteralCallee + if literal { + if literalReserveUsed >= localizationFinalResponseLiteralReserve { + continue + } + if literalReserveUsed > 0 { + if _, sameFile := literalReserveFiles[strings.TrimSpace(row.File)]; sameFile { + continue + } + } + if appendPrimary(row) { + claimReserve(row) + } + continue } + appendPrimary(row) } taskTerms := exploreTerminalTerms(task) @@ -779,7 +1123,7 @@ func localizationFinalResponseRows(task string, current, rows []localizationDige } } if bestSameOwner >= 0 { - appendRow(&primaries, localizationFinalResponsePrimaryLimit, rows[bestSameOwner]) + appendPrimary(rows[bestSameOwner]) } } @@ -799,20 +1143,68 @@ func localizationFinalResponseRows(task string, current, rows []localizationDige } } if bestTaskMatch >= 0 { - appendRow(&primaries, localizationFinalResponsePrimaryLimit, rows[bestTaskMatch]) + appendPrimary(rows[bestTaskMatch]) } } + for _, row := range rows { - appendRow(&primaries, localizationFinalResponsePrimaryLimit, row) + // Further authenticated literal callees may still win through ordinary + // rank; only the pre-seated reserve above is capped at one. + appendPrimary(row) + } + + // After every ranked seat is filled, remaining task-named rows extend the + // block by bounded extra seats. The admission lane must not decide + // presentation authority: a declaration the task names is what the caller + // asked about, whether it arrived through ranked retrieval, leading-file + // completeness, or a body mention. Naming means the raw task carries the + // row's identifier as a whole word, case-sensitively — prose substrings + // do not extend the answer. Adjacency neighbors stay supporting: sharing + // an identifier with the task does not make a graph hop the subject of + // the report. + if strings.TrimSpace(task) != "" { + extendedLimit := localizationFinalResponsePrimaryLimit + localizationTaskAlignedExtraSeats + type taskAlignedCandidate struct { + index int + score localizationFinalResponseTaskScore + } + aligned := make([]taskAlignedCandidate, 0, len(rows)) + for index, row := range rows { + if _, exists := selected[strings.TrimSpace(row.ID)]; exists { + continue + } + if row.Provenance == localizationProvenanceDirectAdjacency { + continue + } + // Only the leading-file-depth writer of supportingOnly may be + // overridden by task naming; a supplemental row stays supporting + // no matter how well it matches the task. + if localizationFinalResponseSupportingOnly(row) && !row.leadingFileDepth { + continue + } + if !localizationTaskNamesRow(task, row) { + continue + } + aligned = append(aligned, taskAlignedCandidate{ + index: index, score: scoreLocalizationFinalResponseTask(taskTerms, row), + }) + } + sort.SliceStable(aligned, func(left, right int) bool { + return localizationFinalResponseBetterTaskScore(aligned[left].score, aligned[right].score) + }) + for _, candidate := range aligned { + appendRow(&primaries, extendedLimit, rows[candidate.index]) + } } + supportingLimit := localizationReplayEvidenceLimit - len(primaries) for _, row := range rows { if localizationFinalResponseDirectRelation(row, primaries) { - appendRow(&supporting, localizationFinalResponseSupportingLimit, row) + appendRow(&supporting, supportingLimit, row) } } for _, row := range rows { - appendRow(&supporting, localizationFinalResponseSupportingLimit, row) + appendRow(&supporting, supportingLimit, row) } presented := make([]localizationFinalResponseRow, 0, len(primaries)+len(supporting)) @@ -825,6 +1217,61 @@ func localizationFinalResponseRows(task string, current, rows []localizationDige return presented } +func localizationFinalResponsePrimaryIDs(task string, current, rows []localizationDigestRow) []string { + presented := localizationFinalResponseRows(task, current, rows) + ids := make([]string, 0, localizationFinalResponsePrimaryLimit) + for _, item := range presented { + if !item.primary { + break + } + if id := strings.TrimSpace(item.row.ID); id != "" { + ids = append(ids, id) + } + } + return ids +} + +// freezeLocalizationPrimaryCohort records the exact initial PRIMARY projection +// immediately before supplemental evidence is materialized. The order is +// request-local policy state: it is retained across digest rebuilds and reads, +// but never serialized into the response payload. +func freezeLocalizationPrimaryCohort( + task string, + envelope *localizationExploreEnvelope, + digest *localizationEvidenceDigest, +) { + if envelope == nil || digest == nil { + return + } + primaryIDs := append([]string(nil), digest.primaryIDs...) + if len(primaryIDs) == 0 { + primaryIDs = localizationFinalResponsePrimaryIDs(task, nil, digest.Evidence) + } + if len(primaryIDs) > localizationFinalResponsePrimaryLimit { + primaryIDs = primaryIDs[:localizationFinalResponsePrimaryLimit] + } + orders := make(map[string]int, len(primaryIDs)) + for index, id := range primaryIDs { + if id = strings.TrimSpace(id); id != "" { + orders[id] = index + 1 + } + } + for index := range envelope.Evidence { + row := &envelope.Evidence[index] + if order := orders[strings.TrimSpace(row.ID)]; order > 0 && !row.supportingOnly && + !localizationSupportingOnlyProvenance(row.Provenance) { + row.primaryCohortOrder = order + } + } + for index := range digest.Evidence { + row := &digest.Evidence[index] + if order := orders[strings.TrimSpace(row.ID)]; order > 0 && !localizationFinalResponseSupportingOnly(*row) { + row.primaryCohortOrder = order + } + } + refreshLocalizationDigestResponses(digest, task, nil) +} + func renderLocalizationFinalResponse(rows []localizationDigestRow) string { return renderLocalizationFinalResponseForTask("", nil, rows) } @@ -948,6 +1395,7 @@ func refreshLocalizationDigestResponses(digest *localizationEvidenceDigest, task if digest == nil { return } + digest.primaryIDs = localizationFinalResponsePrimaryIDs(task, current, digest.Evidence) digest.finalResponse = renderLocalizationFinalResponseForTask(task, current, digest.Evidence) digest.provisionalResponse = renderLocalizationProvisionalResponseForTask(task, current, digest.Evidence) } @@ -1241,6 +1689,7 @@ type localizationHostEnvelope struct { Version int `json:"version"` FallbackFormat string `json:"fallback_format"` Evidence *localizationEvidenceDigest `json:"evidence"` + PrimaryIDs []string `json:"primary_ids,omitempty"` Contract localizationTerminalContract `json:"contract"` } @@ -1279,10 +1728,15 @@ func attachLocalizationHostEnvelope(result *mcpgo.CallToolResult, completion loc if result.Meta.AdditionalFields == nil { result.Meta.AdditionalFields = make(map[string]any) } + primaryIDs := []string(nil) + if digest != nil { + primaryIDs = append(primaryIDs, digest.primaryIDs...) + } result.Meta.AdditionalFields[localizationHostMetaKey] = localizationHostEnvelope{ Version: 1, FallbackFormat: "{file}:{line} — {id} ({signature})", Evidence: digest, + PrimaryIDs: primaryIDs, Contract: contract, } return result diff --git a/internal/mcp/localization_digest_test.go b/internal/mcp/localization_digest_test.go index 6ba37602..341dde39 100644 --- a/internal/mcp/localization_digest_test.go +++ b/internal/mcp/localization_digest_test.go @@ -209,6 +209,122 @@ func TestDigestLifecycleAndLegacyFallback(t *testing.T) { requireLocalizationTerminalReplay(t, blocked, "search", "symbols") } +func TestOnlyAuthenticatedSourceLiteralCalleeReservesOnePrimarySeat(t *testing.T) { + rows := []localizationDigestRow{ + {ID: "repo/d.go::Semantic", Name: "Semantic", File: "repo/d.go"}, + {ID: "repo/e.go::RankedTwo", Name: "RankedTwo", File: "repo/e.go"}, + {ID: "repo/f.go::RankedThree", Name: "RankedThree", File: "repo/f.go"}, + {ID: "repo/g.go::RankedFour", Name: "RankedFour", File: "repo/g.go"}, + {ID: "repo/h.go::RankedFive", Name: "RankedFive", File: "repo/h.go"}, + {ID: "repo/b.go::FirstCallee", Name: "FirstCallee", File: "repo/b.go", Provenance: localizationProvenanceSourceLiteralCallee, literalPrimaryEligible: true}, + {ID: "repo/c.go::SecondCallee", Name: "SecondCallee", File: "repo/c.go", Provenance: localizationProvenanceSourceLiteralCallee}, + {ID: "repo/a.go::Content", Name: "Content", File: "repo/a.go", Provenance: localizationProvenanceContentLiteral}, + } + + presented := localizationFinalResponseRows("find the semantic handler", nil, rows) + reservedLiteralPrimaries := 0 + contentLiteralPrimary := false + semanticPrimary := false + for _, row := range presented { + if !row.primary { + continue + } + if row.row.Provenance == localizationProvenanceSourceLiteralCallee { + reservedLiteralPrimaries++ + } + if row.row.Provenance == localizationProvenanceContentLiteral { + contentLiteralPrimary = true + } + if row.row.ID == "repo/d.go::Semantic" { + semanticPrimary = true + } + } + if reservedLiteralPrimaries != localizationFinalResponseLiteralReserve { + t.Fatalf("reserved literal PRIMARY rows = %d, want %d", reservedLiteralPrimaries, localizationFinalResponseLiteralReserve) + } + if contentLiteralPrimary { + t.Fatal("generic content literal gained PRIMARY authority from the literal alone") + } + if !semanticPrimary { + t.Fatal("literal reserve consumed the semantic PRIMARY opportunity") + } +} + +func TestFrozenPrimaryCohortSurvivesSupplementalRows(t *testing.T) { + cases := []struct { + name string + task string + gold string + }{ + {name: "aiohttp anchored import", task: "import aiohttp fails when zstd is not installed", gold: "aiohttp/compression_utils.py::ZSTDDecompressor.__init__"}, + {name: "fmt anchored print", task: "ostream print fails with constant conditional warning", gold: "include/fmt/ostream.h::print"}, + {name: "vue anchored hydrate", task: "skip detached async component hydration", gold: "packages/runtime-core/src/apiAsyncComponent.ts::performHydrate"}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + envelope := localizationExploreEnvelope{Evidence: []localizationEvidence{ + {ID: tc.gold, Name: "gold", File: strings.Split(tc.gold, "::")[0]}, + {ID: "repo/second.go::Second", Name: "Second", File: "repo/second.go"}, + {ID: "repo/third.go::Third", Name: "Third", File: "repo/third.go"}, + {ID: "repo/fourth.go::Fourth", Name: "Fourth", File: "repo/fourth.go"}, + {ID: "repo/fifth.go::Fifth", Name: "Fifth", File: "repo/fifth.go"}, + }} + digest := newLocalizationEvidenceDigestForTask(tc.task, envelope) + before := append([]string(nil), digest.primaryIDs...) + freezeLocalizationPrimaryCohort(tc.task, &envelope, digest) + for index := 0; index < 4; index++ { + envelope.Evidence = append(envelope.Evidence, localizationEvidence{ + ID: fmt.Sprintf("repo/supplement%02d.go::HydratePrint", index), + Name: "HydratePrint", File: fmt.Sprintf("repo/supplement%02d.go", index), + Provenance: localizationProvenanceDirectAdjacency, supportingOnly: true, + }) + } + afterDigest := newLocalizationEvidenceDigestForTask(tc.task, envelope) + after := localizationFinalResponsePrimaryIDs(tc.task, nil, afterDigest.Evidence) + if !reflect.DeepEqual(after, before) { + t.Fatalf("PRIMARY cohort changed after supplemental rows:\nbefore=%v\nafter=%v", before, after) + } + if len(after) != localizationFinalResponsePrimaryLimit || after[0] != tc.gold { + t.Fatalf("PRIMARY cohort lost anchored candidate or count: %v", after) + } + }) + } +} + +func TestSupplementalRowsStaySupportingAcrossEverySelectorLoop(t *testing.T) { + rows := []localizationDigestRow{ + {ID: "repo/base.go::First", Name: "First", File: "repo/base.go", primaryCohortOrder: 1}, + {ID: "repo/base.go::Second", Name: "Second", File: "repo/base.go", primaryCohortOrder: 2}, + {ID: "repo/base.go::Third", Name: "Third", File: "repo/base.go", primaryCohortOrder: 3}, + {ID: "repo/base.go::Fourth", Name: "Fourth", File: "repo/base.go", primaryCohortOrder: 4}, + {ID: "repo/supp.go::ExactTaskMatch", Name: "ExactTaskMatch", Kind: "method", File: "repo/supp.go", supportingOnly: true}, + {ID: "repo/body.go::Body", Name: "Body", File: "repo/body.go", Provenance: localizationProvenanceBodyMention, supportingOnly: true}, + {ID: "repo/adj.go::Adjacent", Name: "Adjacent", File: "repo/adj.go", Provenance: localizationProvenanceDirectAdjacency, supportingOnly: true}, + } + current := []localizationDigestRow{{ID: "repo/supp.go::Current", Name: "Current", QualName: "Owner.Current", Kind: "method", File: "repo/supp.go"}} + first := localizationFinalResponseRows("ExactTaskMatch", current, rows) + second := localizationFinalResponseRows("ExactTaskMatch", current, rows) + if !reflect.DeepEqual(first, second) { + t.Fatal("selector output is not deterministic") + } + for _, item := range first { + if item.primary && localizationFinalResponseSupportingOnly(item.row) { + t.Fatalf("supplemental row became PRIMARY: %+v", item.row) + } + } +} + +func TestContentLiteralProvenanceDoesNotBecomeStrongTerminalProof(t *testing.T) { + rows := map[string]localizationDigestRow{ + "repo/a.go::First": { + ID: "repo/a.go::First", File: "repo/a.go", Provenance: localizationProvenanceContentLiteral, + }, + } + if localizationDigestStrongProofRetained(rows, "repo/a.go::First") { + t.Fatal("ordinary content evidence became enforceable terminal proof") + } +} + func TestDigestByteCapShedsOptionalMetadataBeforeEvidenceTail(t *testing.T) { envelope := localizationExploreEnvelope{} for i := 0; i < 400; i++ { @@ -254,6 +370,85 @@ func TestDigestByteCapShedsOptionalMetadataBeforeEvidenceTail(t *testing.T) { } } +func TestDigestPressureShedsSupplementalRowsBeforeRankedEvidence(t *testing.T) { + const ( + protectedID = "repo/protected.go::Protected" + bodyID = "repo/body.go::BodyMention" + adjacentID = "repo/adjacent.go::Adjacent" + ordinary = 13 + ) + largeSignature := strings.Repeat("large optional signature ", 120) + completion := newLocalizationRefinementCompletion(protectedID) + evidence := []localizationEvidence{{ + Rank: 1, ID: protectedID, Name: "Protected", Kind: "function", + File: "repo/protected.go", Signature: largeSignature, + Provenance: localizationProvenanceDirectAdjacency, supportingOnly: true, + }} + for index := 0; index < ordinary; index++ { + evidence = append(evidence, localizationEvidence{ + Rank: len(evidence) + 1, + ID: fmt.Sprintf("repo/ordinary%02d.go::Ordinary%02d", index, index), + Name: fmt.Sprintf("Ordinary%02d", index), Kind: "function", + File: fmt.Sprintf("repo/ordinary%02d.go", index), Signature: largeSignature, + }) + } + evidence = append(evidence, + localizationEvidence{Rank: len(evidence) + 1, ID: bodyID, Name: "BodyMention", Kind: "function", File: "repo/body.go", Signature: largeSignature, Provenance: localizationProvenanceBodyMention, supportingOnly: true}, + localizationEvidence{Rank: len(evidence) + 2, ID: adjacentID, Name: "Adjacent", Kind: "function", File: "repo/adjacent.go", Signature: largeSignature, Provenance: localizationProvenanceDirectAdjacency, supportingOnly: true}, + ) + + assertRetained := func(t *testing.T, digest *localizationEvidenceDigest) { + t.Helper() + ids := make(map[string]localizationDigestRow, len(digest.Evidence)) + for index, row := range digest.Evidence { + ids[row.ID] = row + if row.Rank != index+1 || digest.Files[index] != row.File || digest.Symbols[index] != row.ID { + t.Fatalf("unaligned retained row %d: %#v", index+1, row) + } + } + if _, exists := ids[bodyID]; exists { + t.Fatal("body supplement survived retained-state pressure") + } + if _, exists := ids[adjacentID]; exists { + t.Fatal("direct adjacency supplement survived retained-state pressure") + } + protected, exists := ids[protectedID] + if !exists || !protected.authorizationPriority { + t.Fatalf("authorized supplemental identity was shed or lost priority: %#v", protected) + } + for index := 0; index < ordinary; index++ { + id := fmt.Sprintf("repo/ordinary%02d.go::Ordinary%02d", index, index) + if _, exists := ids[id]; !exists { + t.Fatalf("independently ranked row %q was shed before supplements", id) + } + } + encoded, err := json.Marshal(digest) + if err != nil || len(encoded) > localizationDigestMaxBytes || len(digest.finalResponse) > localizationFinalResponseMaxBytes { + t.Fatalf("unbounded digest bytes=%d final=%d err=%v", len(encoded), len(digest.finalResponse), err) + } + } + + t.Run("initial digest", func(t *testing.T) { + digest := newLocalizationEvidenceDigestForTask("find Protected", localizationExploreEnvelope{ + Completion: completion, + Evidence: evidence, + }) + assertRetained(t, digest) + }) + t.Run("post-read merge", func(t *testing.T) { + rows := make([]localizationDigestRow, 0, len(evidence)) + for _, row := range evidence { + rows = append(rows, localizationDigestRow{ + Rank: row.Rank, ID: row.ID, Name: row.Name, Kind: row.Kind, + File: row.File, Signature: row.Signature, Provenance: row.Provenance, + supportingOnly: row.supportingOnly, + authorizationPriority: row.ID == protectedID, + }) + } + assertRetained(t, mergeLocalizationEvidenceDigest(nil, &localizationEvidenceDigest{Evidence: rows})) + }) +} + func authorizationAwareDigestFixture() (localizationExploreEnvelope, []string) { allowed := []string{ "repo/storage/level.go::StorageLevel.Normalize", @@ -832,8 +1027,8 @@ func TestTaskAwareDigestMergeDoesNotCohortQualifiedFunctions(t *testing.T) { func TestDigestByteCapRetainsSingleMandatoryRowAfterSheddingOptionalFields(t *testing.T) { envelope := localizationExploreEnvelope{Evidence: []localizationEvidence{{ - Rank: 1, - ID: "repo/registry.go::Registry.Configure", + Rank: 1, + ID: "repo/registry.go::Registry.Configure", // Each optional field scales with the retention cap so every shed // step is still forced: after callers, callees, and the signature go, // the qual-name alone still busts the cap and must go too. diff --git a/internal/mcp/localization_direct_adjacency.go b/internal/mcp/localization_direct_adjacency.go new file mode 100644 index 00000000..f79213eb --- /dev/null +++ b/internal/mcp/localization_direct_adjacency.go @@ -0,0 +1,418 @@ +package mcp + +import ( + "sort" + "strings" + + "github.com/zzet/gortex/internal/graph" +) + +const ( + localizationDirectAdjacencyCap = 8 + localizationDirectAdjacencyLookupCap = localizationReplayEvidenceLimit * 6 + localizationProvenanceDirectAdjacency = "direct_adjacency" +) + +type localizationDirectAdjacencyCandidate struct { + node *graph.Node + ownerFile string + ownerIndex int + relationOrder int + direction int + matched int + longest int + production bool + callable bool + sameFile bool + taskCited bool + taskCitationOffset int +} + +func localizationDirectAdjacencyNodeTaskCitationOffset(task string, node *graph.Node) int { + if node == nil { + return -1 + } + return localizationEvidenceTaskCitationOffset(task, localizationEvidence{ + ID: node.ID, Name: node.Name, QualName: node.QualName, + }) +} + +// promoteLocalizationDirectAdjacency promotes only graph-authenticated nodes +// already named by the bounded caller/callee IDs in the retained result. It +// performs one bounded batch lookup and never traverses another adjacency list. +func promoteLocalizationDirectAdjacency( + task string, + envelope localizationExploreEnvelope, + reader graph.Reader, + maxEvidence int, + maxBytes int, + digest *localizationEvidenceDigest, +) (localizationExploreEnvelope, *localizationEvidenceDigest) { + if reader == nil || len(envelope.Evidence) == 0 { + return envelope, digest + } + + seenIDs := make(map[string]struct{}, len(envelope.Evidence)+localizationDirectAdjacencyLookupCap) + for _, row := range envelope.Evidence { + if id := strings.TrimSpace(row.ID); id != "" { + seenIDs[id] = struct{}{} + } + } + + type relation struct { + id string + ownerFile string + ownerIndex int + relationOrder int + direction int + } + relations := make([]relation, 0, localizationDirectAdjacencyLookupCap) + lookupIDs := make([]string, 0, localizationDirectAdjacencyLookupCap) + lookupSeen := make(map[string]struct{}, localizationDirectAdjacencyLookupCap) + appendRelations := func(ids []string, owner localizationEvidence, ownerIndex, direction int) { + if len(ids) > localizationMaxNeighborIDs { + ids = ids[:localizationMaxNeighborIDs] + } + for relationOrder, rawID := range ids { + // Relation contexts are independently bounded, but an identity may + // appear beside several owners. Retain those contexts so ranking can + // select its best same-file/direction proof while the store lookup + // remains one deduplicated batch. + if len(relations) >= localizationDirectAdjacencyLookupCap { + return + } + id := strings.TrimSpace(rawID) + if id == "" { + continue + } + if _, exists := seenIDs[id]; exists { + continue + } + if _, exists := lookupSeen[id]; !exists { + lookupSeen[id] = struct{}{} + lookupIDs = append(lookupIDs, id) + } + relations = append(relations, relation{ + id: id, + ownerFile: strings.TrimSpace(owner.File), + ownerIndex: ownerIndex, + relationOrder: relationOrder, + direction: direction, + }) + } + } + for ownerIndex, owner := range envelope.Evidence { + // Callees are usually the implementation named by a caller-shaped task; + // callers remain equally eligible and win whenever task alignment is better. + appendRelations(owner.Callees, owner, ownerIndex, 0) + appendRelations(owner.Callers, owner, ownerIndex, 1) + if len(relations) >= localizationDirectAdjacencyLookupCap { + break + } + } + if len(lookupIDs) == 0 { + return envelope, digest + } + + nodes := reader.GetNodesByIDs(lookupIDs) + taskTerms := exploreTerminalTerms(task) + candidates := make([]localizationDirectAdjacencyCandidate, 0, len(relations)) + for _, relation := range relations { + node := nodes[relation.id] + if node == nil || strings.TrimSpace(node.ID) != relation.id || !localizationBodyMentionNodeEligible(node) { + continue + } + matched, longest := exploreDraftTermOverlap(taskTerms, node) + kind := strings.ToLower(strings.TrimSpace(string(node.Kind))) + taskCitationOffset := localizationDirectAdjacencyNodeTaskCitationOffset(task, node) + candidates = append(candidates, localizationDirectAdjacencyCandidate{ + node: node, + ownerFile: relation.ownerFile, + ownerIndex: relation.ownerIndex, + relationOrder: relation.relationOrder, + direction: relation.direction, + matched: matched, + longest: longest, + production: !exploreDraftIsTestNode(node), + callable: kind == "function" || kind == "method", + sameFile: relation.ownerFile != "" && nodeDisplayPath(node) == relation.ownerFile, + taskCited: taskCitationOffset >= 0, + taskCitationOffset: taskCitationOffset, + }) + } + if len(candidates) == 0 { + return envelope, digest + } + + sort.SliceStable(candidates, func(first, second int) bool { + left, right := candidates[first], candidates[second] + if left.matched != right.matched { + return left.matched > right.matched + } + if left.production != right.production { + return left.production + } + if left.callable != right.callable { + return left.callable + } + if left.sameFile != right.sameFile { + return left.sameFile + } + if left.longest != right.longest { + return left.longest > right.longest + } + if left.ownerIndex != right.ownerIndex { + return left.ownerIndex < right.ownerIndex + } + if left.direction != right.direction { + return left.direction < right.direction + } + if left.relationOrder != right.relationOrder { + return left.relationOrder < right.relationOrder + } + return left.node.ID < right.node.ID + }) + + // PRIMARY authority follows the request's first exact identifier citation, + // while the evidence list keeps the established graph/task-score order. + taskCitedPrimaryID := "" + taskCitedPrimaryOffset := -1 + for _, candidate := range candidates { + if candidate.taskCitationOffset >= 0 && + (taskCitedPrimaryOffset < 0 || candidate.taskCitationOffset < taskCitedPrimaryOffset) { + taskCitedPrimaryID = strings.TrimSpace(candidate.node.ID) + taskCitedPrimaryOffset = candidate.taskCitationOffset + } + } + + promoted := 0 + taskCitedPrimaryPending := taskCitedPrimaryID != "" + for _, candidate := range candidates { + if promoted >= localizationDirectAdjacencyCap { + break + } + node := candidate.node + id := strings.TrimSpace(node.ID) + if taskCitedPrimaryPending && id != taskCitedPrimaryID && + promoted == localizationDirectAdjacencyCap-1 { + continue + } + if id == taskCitedPrimaryID { + taskCitedPrimaryPending = false + } + if _, exists := seenIDs[id]; exists { + continue + } + row := localizationEvidence{ + Rank: len(envelope.Evidence) + 1, + ID: id, + Name: compactLocalizationField(node.Name, localizationMaxNameRunes), + QualName: compactLocalizationField(node.QualName, localizationMaxNameRunes), + Kind: string(node.Kind), + File: nodeDisplayPath(node), + Line: node.StartLine, + EndLine: node.EndLine, + Provenance: localizationProvenanceDirectAdjacency, + taskCitedPrimaryEligible: candidate.taskCited && id == taskCitedPrimaryID, + } + if row.File == "" { + continue + } + + candidateEnvelope, admitted := localizationDirectAdjacencyEnvelopeWithRow(task, envelope, digest, row, maxEvidence) + if !admitted { + continue + } + candidateDigest := newLocalizationEvidenceDigestForTask(task, candidateEnvelope) + if row.taskCitedPrimaryEligible { + candidateEnvelope, candidateDigest = promoteLocalizationTaskCitedAdjacencyPrimary( + task, candidateEnvelope, candidateDigest, row.ID, + ) + } + if !localizationDirectAdjacencyDigestContains(candidateDigest, row.ID) { + continue + } + contract := localizationContractReconciledWithDigest(candidateEnvelope.Completion, candidateDigest) + candidateEnvelope.Completion = contract.Completion + candidateEnvelope.Terminal = contract.Terminal + if !localizationEnvelopeFits(candidateEnvelope, maxBytes) { + continue + } + envelope, digest = candidateEnvelope, candidateDigest + seenIDs[id] = struct{}{} + promoted++ + } + return envelope, digest +} + +func localizationDirectAdjacencyEnvelopeWithRow( + task string, + envelope localizationExploreEnvelope, + digest *localizationEvidenceDigest, + row localizationEvidence, + maxEvidence int, +) (localizationExploreEnvelope, bool) { + if maxEvidence <= 0 { + return envelope, false + } + evidence := append([]localizationEvidence(nil), envelope.Evidence...) + if len(evidence) >= maxEvidence { + retainedDigest := digest + if retainedDigest == nil { + retainedDigest = newLocalizationEvidenceDigestForTask(task, envelope) + } + // Contract priorities include exact/allowed identities and every + // implementation/proof dependency of an authorized refinement route. + // Supporting adjacency must never narrow that live authorization merely + // because only five of its rows fit the model-facing PRIMARY block. + protectedIDs := localizationDigestPriorityIDs(envelope.Completion, evidence) + if retainedDigest != nil { + for _, presented := range localizationFinalResponseRows(task, nil, retainedDigest.Evidence) { + if presented.primary { + protectedIDs[strings.TrimSpace(presented.row.ID)] = struct{}{} + } + } + } + replace := -1 + for index := len(evidence) - 1; index >= 0; index-- { + id := strings.TrimSpace(evidence[index].ID) + if _, protected := protectedIDs[id]; protected { + continue + } + if localizationEvidenceTaskCited(task, evidence[index]) { + continue + } + if localizationSupportingOnlyProvenance(evidence[index].Provenance) { + continue + } + replace = index + break + } + if replace < 0 && row.taskCitedPrimaryEligible { + victimID, victimOrder := localizationTaskCitedAdjacencyPrimaryVictim(task, retainedDigest, row.ID) + for index := range evidence { + if evidence[index].ID == victimID { + replace = index + row.primaryCohortOrder = victimOrder + break + } + } + } + if replace < 0 { + return envelope, false + } + evidence = append(evidence[:replace], evidence[replace+1:]...) + } + if len(evidence) >= maxEvidence { + return envelope, false + } + evidence = append(evidence, row) + for index := range evidence { + evidence[index].Rank = index + 1 + } + + candidate := envelope + candidate.Evidence = evidence + candidate.Files = make([]string, 0, len(evidence)) + candidate.Symbols = make([]string, 0, len(evidence)) + for _, retained := range evidence { + candidate.Files = append(candidate.Files, retained.File) + candidate.Symbols = append(candidate.Symbols, retained.ID) + } + return candidate, true +} + +func localizationTaskCitedAdjacencyPrimaryVictim( + task string, + digest *localizationEvidenceDigest, + candidateID string, +) (string, int) { + if digest == nil { + return "", 0 + } + for _, row := range digest.Evidence { + if row.ID != candidateID && row.taskCitedPrimaryEligible && row.primaryCohortOrder > 0 { + return "", 0 + } + } + for order := localizationFinalResponsePrimaryLimit; order >= 2; order-- { + for _, row := range digest.Evidence { + if row.primaryCohortOrder != order || row.ID == candidateID || + row.authorizationPriority || localizationFinalResponsePrimaryProvenance(row) || + localizationDigestRowIdentifierTaskCited(task, row) { + continue + } + return row.ID, order + } + } + return "", 0 +} + +func promoteLocalizationTaskCitedAdjacencyPrimary( + task string, + envelope localizationExploreEnvelope, + digest *localizationEvidenceDigest, + candidateID string, +) (localizationExploreEnvelope, *localizationEvidenceDigest) { + candidateID = strings.TrimSpace(candidateID) + if digest == nil || candidateID == "" { + return envelope, digest + } + for _, row := range digest.Evidence { + if row.ID == candidateID && row.taskCitedPrimaryEligible && row.primaryCohortOrder > 0 { + return envelope, digest + } + } + victimID, victimOrder := localizationTaskCitedAdjacencyPrimaryVictim(task, digest, candidateID) + if victimID == "" { + return envelope, digest + } + + candidateEnvelopeIndex, victimEnvelopeIndex := -1, -1 + for index := range envelope.Evidence { + switch envelope.Evidence[index].ID { + case candidateID: + candidateEnvelopeIndex = index + case victimID: + victimEnvelopeIndex = index + } + } + candidateDigestIndex, victimDigestIndex := -1, -1 + for index := range digest.Evidence { + switch digest.Evidence[index].ID { + case candidateID: + candidateDigestIndex = index + case victimID: + victimDigestIndex = index + } + } + if candidateEnvelopeIndex < 0 || victimEnvelopeIndex < 0 || + candidateDigestIndex < 0 || victimDigestIndex < 0 { + return envelope, digest + } + + envelope.Evidence[victimEnvelopeIndex].primaryCohortOrder = 0 + envelope.Evidence[candidateEnvelopeIndex].taskCitedPrimaryEligible = true + envelope.Evidence[candidateEnvelopeIndex].primaryCohortOrder = victimOrder + digest.Evidence[victimDigestIndex].primaryCohortOrder = 0 + digest.Evidence[candidateDigestIndex].taskCitedPrimaryEligible = true + digest.Evidence[candidateDigestIndex].primaryCohortOrder = victimOrder + refreshLocalizationDigestResponses(digest, task, nil) + return envelope, digest +} + +func localizationDirectAdjacencyDigestContains(digest *localizationEvidenceDigest, id string) bool { + if digest == nil { + return false + } + for _, row := range digest.Evidence { + if row.ID == id && row.Provenance == localizationProvenanceDirectAdjacency { + return true + } + } + return false +} + +func localizationSupportingOnlyProvenance(provenance string) bool { + return provenance == localizationProvenanceBodyMention || provenance == localizationProvenanceDirectAdjacency +} diff --git a/internal/mcp/localization_direct_adjacency_test.go b/internal/mcp/localization_direct_adjacency_test.go new file mode 100644 index 00000000..8ace9d0c --- /dev/null +++ b/internal/mcp/localization_direct_adjacency_test.go @@ -0,0 +1,676 @@ +package mcp + +import ( + "fmt" + "reflect" + "strings" + "testing" + + "github.com/zzet/gortex/internal/graph" + "github.com/zzet/gortex/internal/search/rerank" +) + +type localizationDirectAdjacencyReader struct { + graph.Reader + nodes map[string]*graph.Node + batches [][]string +} + +func (reader *localizationDirectAdjacencyReader) GetNodesByIDs(ids []string) map[string]*graph.Node { + reader.batches = append(reader.batches, append([]string(nil), ids...)) + result := make(map[string]*graph.Node, len(ids)) + for _, id := range ids { + if node, exists := reader.nodes[id]; exists { + result[id] = node + } + } + return result +} + +func TestPromoteLocalizationDirectAdjacencyAuthenticatesBothDirections(t *testing.T) { + reader := &localizationDirectAdjacencyReader{nodes: map[string]*graph.Node{ + "callee": nodeForDirectAdjacency("callee", "normalizeException", "src/formatter.php", 41), + "caller": nodeForDirectAdjacency("caller", "normalize", "src/formatter.php", 20), + "forged": nodeForDirectAdjacency("different-id", "forged", "src/forged.php", 1), + }} + envelope := localizationExploreEnvelope{Evidence: []localizationEvidence{{ + Rank: 1, ID: "owner", Name: "JsonFormatter", Kind: "class", File: "src/formatter.php", Line: 10, + Callers: []string{"caller", "missing"}, Callees: []string{"callee", "forged"}, + }}} + digest := newLocalizationEvidenceDigestForTask("normalize exception", envelope) + + promoted, _ := promoteLocalizationDirectAdjacency("normalize exception", envelope, reader, 3, 1<<20, digest) + + if len(reader.batches) != 1 { + t.Fatalf("GetNodesByIDs calls = %d, want exactly one", len(reader.batches)) + } + assertLocalizationEvidenceIDs(t, promoted.Evidence, []string{"owner", "callee", "caller"}) + for _, row := range promoted.Evidence[1:] { + if row.Provenance != localizationProvenanceDirectAdjacency { + t.Fatalf("provenance for %q = %q", row.ID, row.Provenance) + } + } +} + +func TestPromoteLocalizationDirectAdjacencyIsBoundedAndDeterministic(t *testing.T) { + nodes := make(map[string]*graph.Node) + evidence := make([]localizationEvidence, 0, 4) + for owner := 0; owner < 4; owner++ { + row := localizationEvidence{Rank: owner + 1, ID: fmt.Sprintf("owner-%d", owner), Name: "owner", Kind: "function", File: "src/owner.go", Line: owner + 1} + for related := 0; related < localizationMaxNeighborIDs; related++ { + id := fmt.Sprintf("target-%02d", owner*localizationMaxNeighborIDs+related) + row.Callees = append(row.Callees, id) + nodes[id] = nodeForDirectAdjacency(id, id, "src/target.go", owner*10+related+1) + } + evidence = append(evidence, row) + } + envelope := localizationExploreEnvelope{Evidence: evidence} + firstReader := &localizationDirectAdjacencyReader{nodes: nodes} + secondReader := &localizationDirectAdjacencyReader{nodes: nodes} + + first, _ := promoteLocalizationDirectAdjacency("target", envelope, firstReader, len(evidence)+localizationDirectAdjacencyCap, 1<<20, newLocalizationEvidenceDigestForTask("target", envelope)) + second, _ := promoteLocalizationDirectAdjacency("target", envelope, secondReader, len(evidence)+localizationDirectAdjacencyCap, 1<<20, newLocalizationEvidenceDigestForTask("target", envelope)) + + if got := len(first.Evidence) - len(evidence); got != localizationDirectAdjacencyCap { + t.Fatalf("promoted rows = %d, want cap %d", got, localizationDirectAdjacencyCap) + } + if !reflect.DeepEqual(localizationEvidenceIDs(first.Evidence), localizationEvidenceIDs(second.Evidence)) { + t.Fatalf("non-deterministic order: %v != %v", localizationEvidenceIDs(first.Evidence), localizationEvidenceIDs(second.Evidence)) + } + if got := len(firstReader.batches[0]); got != 4*localizationMaxNeighborIDs { + t.Fatalf("lookup IDs = %d, want bounded %d", got, 4*localizationMaxNeighborIDs) + } +} + +func TestPromoteLocalizationDirectAdjacencyNeverEscalatesAuthority(t *testing.T) { + reader := &localizationDirectAdjacencyReader{nodes: map[string]*graph.Node{ + "target": nodeForDirectAdjacency("target", "criticalTarget", "src/target.go", 7), + }} + envelope := localizationExploreEnvelope{Evidence: []localizationEvidence{{ + Rank: 1, ID: "owner", Name: "owner", Kind: "function", File: "src/owner.go", Line: 1, Callees: []string{"target"}, + }}} + digest := newLocalizationEvidenceDigestForTask("critical target", envelope) + promoted, promotedDigest := promoteLocalizationDirectAdjacency("critical target", envelope, reader, 2, 1<<20, digest) + if len(promoted.Evidence) != 2 { + t.Fatalf("evidence = %d, want 2", len(promoted.Evidence)) + } + for _, row := range localizationFinalResponseRows("critical target", nil, promotedDigest.Evidence) { + if row.row.ID == "target" && row.primary { + t.Fatal("direct adjacency evidence became primary authority") + } + } +} + +func TestPromoteLocalizationDirectAdjacencyPreservesExactAndRoutePriorities(t *testing.T) { + t.Run("exact symbol", func(t *testing.T) { + evidence := make([]localizationEvidence, 0, 7) + for index := 0; index < localizationFinalResponsePrimaryLimit; index++ { + evidence = append(evidence, localizationEvidence{ + Rank: index + 1, ID: fmt.Sprintf("primary-%d", index), Name: "primary", Kind: "function", File: "src/primary.go", Line: index + 1, + }) + } + evidence[0].Callees = []string{"target"} + evidence = append(evidence, + localizationEvidence{Rank: 6, ID: "weak", Name: "weak", Kind: "function", File: "src/weak.go", Line: 1}, + localizationEvidence{Rank: 7, ID: "exact", Name: "exact", Kind: "function", File: "src/exact.go", Line: 1}, + ) + envelope := localizationExploreEnvelope{ + Completion: newLocalizationExactReadCompletion("exact", false), + Evidence: evidence, + } + reader := &localizationDirectAdjacencyReader{nodes: map[string]*graph.Node{ + "target": nodeForDirectAdjacency("target", "target", "src/target.go", 7), + }} + + promoted, _ := promoteLocalizationDirectAdjacency( + "target", envelope, reader, len(evidence), 1<<20, + newLocalizationEvidenceDigestForTask("target", envelope), + ) + + assertLocalizationEvidenceIDs(t, promoted.Evidence, []string{ + "primary-0", "primary-1", "primary-2", "primary-3", "primary-4", "exact", "target", + }) + if promoted.Completion.ExactSymbol != "exact" { + t.Fatalf("exact symbol = %q, want preserved exact", promoted.Completion.ExactSymbol) + } + }) + + t.Run("allowed symbols and route proofs", func(t *testing.T) { + allowed := make([]string, localizationRefinementAllowedSymbolCap) + evidence := make([]localizationEvidence, 0, len(allowed)+3) + for index := range allowed { + allowed[index] = fmt.Sprintf("allowed-%d", index) + evidence = append(evidence, localizationEvidence{ + Rank: index + 1, ID: allowed[index], Name: allowed[index], Kind: "function", File: "src/allowed.go", Line: index + 1, + }) + } + evidence[0].Callees = []string{"target"} + evidence[0].Provenance = localizationProvenanceImplementationRoute + evidence[1].Provenance = localizationProvenanceImplementationTarget + evidence = append(evidence, + localizationEvidence{Rank: 9, ID: "weak", Name: "weak", Kind: "function", File: "src/weak.go", Line: 1}, + localizationEvidence{Rank: 10, ID: "implementation", Name: "implementation", Kind: "function", File: "src/implementation.go", Line: 1, Provenance: localizationProvenanceImplementationTarget}, + localizationEvidence{Rank: 11, ID: "proof", Name: "proof", Kind: "function", File: "src/proof.go", Line: 1, Provenance: localizationProvenanceImplementationRoute}, + ) + completion := newLocalizationRefinementCompletionForSymbols(allowed[0], allowed) + completion.refinementRoutes = map[string]localizationRefinementRoute{ + allowed[0]: {implementationSymbol: "implementation"}, + allowed[1]: {proofSymbol: "proof"}, + } + envelope := localizationExploreEnvelope{Completion: completion, Evidence: evidence} + reader := &localizationDirectAdjacencyReader{nodes: map[string]*graph.Node{ + "target": nodeForDirectAdjacency("target", "target", "src/target.go", 7), + }} + + promoted, _ := promoteLocalizationDirectAdjacency( + "target", envelope, reader, len(evidence), 1<<20, + newLocalizationEvidenceDigestForTask("target", envelope), + ) + + want := append(append([]string(nil), allowed...), "implementation", "proof", "target") + assertLocalizationEvidenceIDs(t, promoted.Evidence, want) + if !reflect.DeepEqual(promoted.Completion.AllowedSymbols, allowed) { + t.Fatalf("allowed symbols = %v, want %v", promoted.Completion.AllowedSymbols, allowed) + } + if route := promoted.Completion.refinementRoutes[allowed[0]]; route.implementationSymbol != "implementation" { + t.Fatalf("implementation route = %+v, want implementation proof retained", route) + } + if route := promoted.Completion.refinementRoutes[allowed[1]]; route.proofSymbol != "proof" { + t.Fatalf("proof route = %+v, want proof retained", route) + } + }) +} + +func TestPromoteLocalizationDirectAdjacencyUsesBestDuplicateRelationContext(t *testing.T) { + reader := &localizationDirectAdjacencyReader{nodes: map[string]*graph.Node{ + "shared": nodeForDirectAdjacency("shared", "sharedTarget", "src/shared.go", 7), + "competitor": nodeForDirectAdjacency("competitor", "competitorTarget", "src/competitor.go", 8), + }} + allowed := []string{"owner-0", "owner-1", "owner-2"} + envelope := localizationExploreEnvelope{ + Completion: newLocalizationRefinementCompletionForSymbols(allowed[0], allowed), + Evidence: []localizationEvidence{ + {Rank: 1, ID: allowed[0], Name: "owner0", Kind: "function", File: "src/other.go", Line: 1, Callees: []string{"shared"}}, + {Rank: 2, ID: allowed[1], Name: "owner1", Kind: "function", File: "src/shared.go", Line: 1, Callees: []string{"shared"}}, + {Rank: 3, ID: allowed[2], Name: "owner2", Kind: "function", File: "src/competitor.go", Line: 1, Callees: []string{"competitor"}}, + }, + } + + promoted, _ := promoteLocalizationDirectAdjacency( + "target", envelope, reader, len(envelope.Evidence)+1, 1<<20, + newLocalizationEvidenceDigestForTask("target", envelope), + ) + + assertLocalizationEvidenceIDs(t, promoted.Evidence, []string{"owner-0", "owner-1", "owner-2", "shared"}) + if len(reader.batches) != 1 || !reflect.DeepEqual(reader.batches[0], []string{"shared", "competitor"}) { + t.Fatalf("deduplicated batch = %v, want [shared competitor]", reader.batches) + } +} + +func TestLocalizationDirectAdjacencyReplacementRenumbersEvidence(t *testing.T) { + evidence := make([]localizationEvidence, 0, 8) + for index := 0; index < localizationFinalResponsePrimaryLimit; index++ { + evidence = append(evidence, localizationEvidence{ + Rank: index + 1, ID: fmt.Sprintf("primary-%d", index), Name: "primary", Kind: "function", File: "src/primary.go", Line: index + 1, + }) + } + evidence = append(evidence, + localizationEvidence{Rank: 6, ID: "replaceable", Name: "replaceable", Kind: "function", File: "src/weak.go", Line: 1}, + localizationEvidence{Rank: 7, ID: "body", Name: "body", Kind: "function", File: "src/body.go", Line: 1, Provenance: localizationProvenanceBodyMention}, + localizationEvidence{Rank: 8, ID: "adjacent", Name: "adjacent", Kind: "function", File: "src/adjacent.go", Line: 1, Provenance: localizationProvenanceDirectAdjacency}, + ) + envelope := localizationExploreEnvelope{Evidence: evidence} + candidate, admitted := localizationDirectAdjacencyEnvelopeWithRow( + "target", envelope, newLocalizationEvidenceDigestForTask("target", envelope), + localizationEvidence{ID: "target", Name: "target", Kind: "function", File: "src/target.go", Line: 1, Provenance: localizationProvenanceDirectAdjacency}, + len(evidence), + ) + if !admitted { + t.Fatal("replacement was not admitted") + } + for index, row := range candidate.Evidence { + if row.Rank != index+1 { + t.Fatalf("evidence[%d] rank = %d, want %d", index, row.Rank, index+1) + } + } + assertLocalizationEvidenceIDs(t, candidate.Evidence, []string{ + "primary-0", "primary-1", "primary-2", "primary-3", "primary-4", "body", "adjacent", "target", + }) +} + +func TestPromoteLocalizationDirectAdjacencyPromotesOneTaskCitedPrimary(t *testing.T) { + task := "hydrateOnVisible waits before performHydrate checks connectivity; secondExplicit is supporting" + primaryNames := []string{ + "getLeavingNodesForType", "hydrate", "moveTeleport", "isTeleportDisabled", "hydrateOnVisible", + } + evidence := make([]localizationEvidence, 0, localizationReplayEvidenceLimit) + for index, name := range primaryNames { + evidence = append(evidence, localizationEvidence{ + Rank: index + 1, ID: fmt.Sprintf("primary-%d", index), Name: name, + Kind: "function", File: "src/primary.ts", Line: index + 1, + primaryCohortOrder: index + 1, + }) + } + evidence[0].Callees = []string{"perform", "extra-0", "extra-1"} + evidence[1].Callees = []string{"extra-2", "extra-3", "extra-4"} + evidence[2].Callees = []string{"extra-5", "extra-6"} + for len(evidence) < localizationReplayEvidenceLimit { + index := len(evidence) + evidence = append(evidence, localizationEvidence{ + Rank: index + 1, ID: fmt.Sprintf("weak-%d", index), Name: "ordinaryCandidate", + Kind: "function", File: "src/weak.ts", Line: index + 1, + }) + } + envelope := localizationExploreEnvelope{Evidence: evidence} + reader := &localizationDirectAdjacencyReader{nodes: map[string]*graph.Node{ + "perform": nodeForDirectAdjacency("perform", "performHydrate", "src/apiAsyncComponent.ts", 73), + "extra-0": nodeForDirectAdjacency("extra-0", "secondExplicit", "src/extra.ts", 80), + "extra-1": nodeForDirectAdjacency("extra-1", "ordinaryNeighbor1", "src/extra.ts", 81), + "extra-2": nodeForDirectAdjacency("extra-2", "ordinaryNeighbor2", "src/extra.ts", 82), + "extra-3": nodeForDirectAdjacency("extra-3", "ordinaryNeighbor3", "src/extra.ts", 83), + "extra-4": nodeForDirectAdjacency("extra-4", "ordinaryNeighbor4", "src/extra.ts", 84), + "extra-5": nodeForDirectAdjacency("extra-5", "ordinaryNeighbor5", "src/extra.ts", 85), + "extra-6": nodeForDirectAdjacency("extra-6", "ordinaryNeighbor6", "src/extra.ts", 86), + }} + + promoted, digest := promoteLocalizationDirectAdjacency( + task, envelope, reader, len(evidence), 1<<20, + newLocalizationEvidenceDigestForTask(task, envelope), + ) + + if len(reader.batches) != 1 { + t.Fatalf("GetNodesByIDs calls = %d, want exactly one", len(reader.batches)) + } + if len(promoted.Evidence) != localizationReplayEvidenceLimit { + t.Fatalf("evidence rows = %d, want %d", len(promoted.Evidence), localizationReplayEvidenceLimit) + } + wantPrimary := []string{"primary-0", "primary-1", "primary-2", "perform", "primary-4"} + if got := localizationFinalResponsePrimaryIDs(task, nil, digest.Evidence); !reflect.DeepEqual(got, wantPrimary) { + t.Fatalf("primary IDs = %v, want %v", got, wantPrimary) + } + rebuilt := newLocalizationEvidenceDigestForTask(task, promoted) + if got := localizationFinalResponsePrimaryIDs(task, nil, rebuilt.Evidence); !reflect.DeepEqual(got, wantPrimary) { + t.Fatalf("rebuilt primary IDs = %v, want %v", got, wantPrimary) + } + rows := cloneLocalizationDigestRows(rebuilt.Evidence) + for { + var removed bool + rows, removed = shedLocalizationDigestSupportingOnly(rows) + if !removed { + break + } + } + if got := localizationFinalResponsePrimaryIDs(task, nil, rows); !reflect.DeepEqual(got, wantPrimary) { + t.Fatalf("primary IDs after supporting-row pressure = %v, want %v", got, wantPrimary) + } +} + +func TestPromoteLocalizationDirectAdjacencyReplacesWeakPrimaryOnTightPage(t *testing.T) { + task := "hydrateOnVisible then performHydrate" + evidence := make([]localizationEvidence, 0, localizationFinalResponsePrimaryLimit) + for index, name := range []string{"weakZero", "weakOne", "weakTwo", "weakThree", "hydrateOnVisible"} { + evidence = append(evidence, localizationEvidence{ + Rank: index + 1, ID: fmt.Sprintf("primary-%d", index), Name: name, + Kind: "function", File: "src/primary.ts", Line: index + 1, + primaryCohortOrder: index + 1, + }) + } + evidence[0].Callees = []string{"perform"} + envelope := localizationExploreEnvelope{Evidence: evidence} + reader := &localizationDirectAdjacencyReader{nodes: map[string]*graph.Node{ + "perform": nodeForDirectAdjacency("perform", "performHydrate", "src/apiAsyncComponent.ts", 73), + }} + + promoted, digest := promoteLocalizationDirectAdjacency( + task, envelope, reader, len(evidence), 1<<20, + newLocalizationEvidenceDigestForTask(task, envelope), + ) + + assertLocalizationEvidenceIDs(t, promoted.Evidence, []string{"primary-0", "primary-1", "primary-2", "primary-4", "perform"}) + wantPrimary := []string{"primary-0", "primary-1", "primary-2", "perform", "primary-4"} + if got := localizationFinalResponsePrimaryIDs(task, nil, digest.Evidence); !reflect.DeepEqual(got, wantPrimary) { + t.Fatalf("tight-page primary IDs = %v, want %v", got, wantPrimary) + } +} + +func TestPromoteLocalizationDirectAdjacencyDoesNotReplaceProtectedPrimary(t *testing.T) { + task := "KeepZero KeepOne KeepTwo KeepThree KeepFour performHydrate" + primaryNames := []string{"KeepZero", "KeepOne", "KeepTwo", "KeepThree", "KeepFour"} + evidence := make([]localizationEvidence, 0, localizationFinalResponsePrimaryLimit+1) + for index, name := range primaryNames { + evidence = append(evidence, localizationEvidence{ + Rank: index + 1, ID: fmt.Sprintf("primary-%d", index), Name: name, + Kind: "function", File: "src/primary.ts", Line: index + 1, + primaryCohortOrder: index + 1, + }) + } + evidence[0].Callees = []string{"perform"} + evidence = append(evidence, localizationEvidence{ + Rank: 6, ID: "weak", Name: "ordinaryCandidate", Kind: "function", File: "src/weak.ts", Line: 1, + }) + envelope := localizationExploreEnvelope{Evidence: evidence} + reader := &localizationDirectAdjacencyReader{nodes: map[string]*graph.Node{ + "perform": nodeForDirectAdjacency("perform", "performHydrate", "src/apiAsyncComponent.ts", 73), + }} + + promoted, digest := promoteLocalizationDirectAdjacency( + task, envelope, reader, len(evidence), 1<<20, + newLocalizationEvidenceDigestForTask(task, envelope), + ) + + wantPrimary := []string{"primary-0", "primary-1", "primary-2", "primary-3", "primary-4"} + if got := localizationFinalResponsePrimaryIDs(task, nil, digest.Evidence); !reflect.DeepEqual(got, wantPrimary) { + t.Fatalf("primary IDs = %v, want protected cohort %v", got, wantPrimary) + } + wantEvidence := append(append([]string(nil), wantPrimary...), "perform") + if got := localizationEvidenceIDs(promoted.Evidence); !reflect.DeepEqual(got, wantEvidence) { + t.Fatalf("evidence IDs = %v, want task-cited adjacency retained as supporting %v", got, wantEvidence) + } + for _, row := range localizationFinalResponseRows(task, nil, digest.Evidence) { + if row.row.ID == "perform" && row.primary { + t.Fatal("task-cited adjacency displaced an all-protected primary cohort") + } + } +} + +func TestPromoteLocalizationDirectAdjacencyPreservesTaskCitedInitialEvidence(t *testing.T) { + task := "change FMT_FORMAT_AS for custom allocator support" + evidence := make([]localizationEvidence, 0, localizationReplayEvidenceLimit) + nodes := make(map[string]*graph.Node, localizationDirectAdjacencyCap) + for index := 0; index < localizationFinalResponsePrimaryLimit; index++ { + row := localizationEvidence{ + Rank: index + 1, ID: fmt.Sprintf("primary-%d", index), Name: "primaryCandidate", + Kind: "function", File: "include/fmt/base.h", Line: index + 1, + primaryCohortOrder: index + 1, + } + if index < 3 { + for related := 0; related < localizationMaxNeighborIDs; related++ { + id := fmt.Sprintf("adjacent-%d-%d", index, related) + row.Callees = append(row.Callees, id) + nodes[id] = nodeForDirectAdjacency(id, "unrelatedNeighbor", "include/fmt/other.h", index*10+related+1) + } + } + evidence = append(evidence, row) + } + evidence = append(evidence, localizationEvidence{ + Rank: 6, ID: "macro", Name: "FMT_FORMAT_AS", Kind: "macro", File: "include/fmt/format.h", Line: 4009, + }) + for len(evidence) < localizationReplayEvidenceLimit { + index := len(evidence) + evidence = append(evidence, localizationEvidence{ + Rank: index + 1, ID: fmt.Sprintf("weak-%d", index), Name: "ordinaryCandidate", + Kind: "function", File: "include/fmt/base.h", Line: index + 1, + }) + } + envelope := localizationExploreEnvelope{Evidence: evidence} + reader := &localizationDirectAdjacencyReader{nodes: nodes} + + promoted, _ := promoteLocalizationDirectAdjacency( + task, envelope, reader, len(evidence), 1<<20, + newLocalizationEvidenceDigestForTask(task, envelope), + ) + + if got := len(promoted.Evidence); got != localizationReplayEvidenceLimit { + t.Fatalf("evidence rows = %d, want %d", got, localizationReplayEvidenceLimit) + } + promotedCount, macroPresent := 0, false + for _, row := range promoted.Evidence { + if row.Provenance == localizationProvenanceDirectAdjacency { + promotedCount++ + } + if row.ID == "macro" { + macroPresent = true + } + } + if promotedCount != localizationDirectAdjacencyCap { + t.Fatalf("promoted adjacency rows = %d, want %d", promotedCount, localizationDirectAdjacencyCap) + } + if !macroPresent { + t.Fatal("task-cited FMT_FORMAT_AS evidence was evicted") + } +} + +func TestPromoteLocalizationDirectAdjacencyRejectsOverBudgetRow(t *testing.T) { + reader := &localizationDirectAdjacencyReader{nodes: map[string]*graph.Node{ + "target": nodeForDirectAdjacency("target", "target", "src/target.go", 7), + }} + envelope := localizationExploreEnvelope{Evidence: []localizationEvidence{{ + Rank: 1, ID: "owner", Name: "owner", Kind: "function", File: "src/owner.go", Line: 1, Callees: []string{"target"}, + }}} + digest := newLocalizationEvidenceDigestForTask("target", envelope) + promoted, _ := promoteLocalizationDirectAdjacency("target", envelope, reader, 2, 1, digest) + assertLocalizationEvidenceIDs(t, promoted.Evidence, []string{"owner"}) +} + +func TestPromoteLocalizationDirectAdjacencyReservesTaskCitedSlot(t *testing.T) { + task := "alpha beta gamma calls performHydrate" + primaryNames := []string{"firstPrimary", "secondPrimary", "thirdPrimary", "weakPrimary", "fifthPrimary"} + evidence := make([]localizationEvidence, 0, localizationReplayEvidenceLimit) + for index, name := range primaryNames { + evidence = append(evidence, localizationEvidence{ + Rank: index + 1, ID: fmt.Sprintf("primary-%d", index), Name: name, + Kind: "function", File: "src/primary.ts", Line: index + 1, + primaryCohortOrder: index + 1, + }) + } + reader := &localizationDirectAdjacencyReader{nodes: make(map[string]*graph.Node)} + for index := 0; index < localizationDirectAdjacencyCap; index++ { + id := fmt.Sprintf("ordinary-%d", index) + ownerIndex := index / localizationMaxNeighborIDs + evidence[ownerIndex].Callees = append(evidence[ownerIndex].Callees, id) + reader.nodes[id] = nodeForDirectAdjacency(id, fmt.Sprintf("ordinaryNeighbor%d", index), "src/ordinary.ts", 20+index) + reader.nodes[id].Meta = map[string]any{"signature": "alpha beta gamma"} + } + evidence[2].Callees = append(evidence[2].Callees, "perform") + reader.nodes["perform"] = nodeForDirectAdjacency("perform", "performHydrate", "src/apiAsyncComponent.ts", 73) + for len(evidence) < localizationReplayEvidenceLimit { + index := len(evidence) + evidence = append(evidence, localizationEvidence{ + Rank: index + 1, ID: fmt.Sprintf("weak-%d", index), Name: "weakCandidate", + Kind: "function", File: "src/weak.ts", Line: index + 1, + }) + } + envelope := localizationExploreEnvelope{Evidence: evidence} + + promoted, digest := promoteLocalizationDirectAdjacency( + task, envelope, reader, len(evidence), 1<<20, + newLocalizationEvidenceDigestForTask(task, envelope), + ) + + if got := promoted.Evidence[len(promoted.Evidence)-1].ID; got != "perform" { + t.Fatalf("last admitted adjacency = %q, want reserved task-cited perform; evidence=%v", got, localizationEvidenceIDs(promoted.Evidence)) + } + wantPrimary := []string{"primary-0", "primary-1", "primary-2", "primary-3", "perform"} + if got := localizationFinalResponsePrimaryIDs(task, nil, digest.Evidence); !reflect.DeepEqual(got, wantPrimary) { + t.Fatalf("primary IDs = %v, want %v", got, wantPrimary) + } +} + +func TestPromoteLocalizationDirectAdjacencyRequiresConcreteIdentifierCitation(t *testing.T) { + tests := []struct { + name string + task string + candidate *graph.Node + }{ + { + name: "identifier suffix", task: "performHydrateLater handles the request", + candidate: nodeForDirectAdjacency("candidate", "performHydrate", "src/candidate.ts", 7), + }, + { + name: "ordinary prose word", task: "the target handles the request", + candidate: nodeForDirectAdjacency("candidate", "target", "src/candidate.ts", 7), + }, + { + name: "file and signature only", task: "inspect src/candidate.ts and func performHydrate()", + candidate: &graph.Node{ + ID: "candidate", Name: "unrelatedConcrete", Kind: graph.KindFunction, + FilePath: "src/candidate.ts", StartLine: 7, + Meta: map[string]any{"signature": "func performHydrate()"}, + }, + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + evidence := make([]localizationEvidence, 0, localizationFinalResponsePrimaryLimit) + for index := 0; index < localizationFinalResponsePrimaryLimit; index++ { + evidence = append(evidence, localizationEvidence{ + Rank: index + 1, ID: fmt.Sprintf("primary-%d", index), Name: fmt.Sprintf("primaryName%d", index), + Kind: "function", File: "src/primary.ts", Line: index + 1, + primaryCohortOrder: index + 1, + }) + } + evidence[0].Callees = []string{"candidate"} + envelope := localizationExploreEnvelope{Evidence: evidence} + reader := &localizationDirectAdjacencyReader{nodes: map[string]*graph.Node{"candidate": test.candidate}} + promoted, digest := promoteLocalizationDirectAdjacency( + test.task, envelope, reader, len(evidence), 1<<20, + newLocalizationEvidenceDigestForTask(test.task, envelope), + ) + want := []string{"primary-0", "primary-1", "primary-2", "primary-3", "primary-4"} + if got := localizationEvidenceIDs(promoted.Evidence); !reflect.DeepEqual(got, want) { + t.Fatalf("evidence IDs = %v, want %v", got, want) + } + if got := localizationFinalResponsePrimaryIDs(test.task, nil, digest.Evidence); !reflect.DeepEqual(got, want) { + t.Fatalf("primary IDs = %v, want %v", got, want) + } + }) + } +} + +func TestTaskCitedAdjacencyPrimarySurvivesDigestPressureAndMerge(t *testing.T) { + task := "performHydrate handles the request" + evidence := []localizationEvidence{ + {Rank: 1, ID: "primary-0", Name: "firstPrimary", Kind: "function", File: "src/primary.ts", Line: 1, primaryCohortOrder: 1}, + {Rank: 2, ID: "primary-1", Name: "secondPrimary", Kind: "function", File: "src/primary.ts", Line: 2, primaryCohortOrder: 2}, + {Rank: 3, ID: "primary-2", Name: "thirdPrimary", Kind: "function", File: "src/primary.ts", Line: 3, primaryCohortOrder: 3}, + {Rank: 4, ID: "primary-4", Name: "fifthPrimary", Kind: "function", File: "src/primary.ts", Line: 5, primaryCohortOrder: 5}, + } + for index := 0; index < localizationReplayEvidenceLimit-len(evidence)-1; index++ { + evidence = append(evidence, localizationEvidence{ + Rank: len(evidence) + 1, + ID: fmt.Sprintf("ordinary-%02d-%s", index, strings.Repeat("x", 700)), + Name: "ordinaryCandidate", Kind: "function", + File: fmt.Sprintf("src/%s/%02d.ts", strings.Repeat("y", 700), index), Line: index + 10, + }) + } + evidence = append(evidence, localizationEvidence{ + Rank: len(evidence) + 1, ID: "perform", Name: "performHydrate", Kind: "function", + File: "src/apiAsyncComponent.ts", Line: 73, + Provenance: localizationProvenanceDirectAdjacency, + taskCitedPrimaryEligible: true, primaryCohortOrder: 4, + }) + digest := newLocalizationEvidenceDigestForTask(task, localizationExploreEnvelope{Evidence: evidence}) + wantPrimary := []string{"primary-0", "primary-1", "primary-2", "perform", "primary-4"} + if len(digest.Evidence) >= len(evidence) { + t.Fatalf("byte pressure retained %d rows, want fewer than %d", len(digest.Evidence), len(evidence)) + } + if len(digest.Evidence) == 0 || digest.Evidence[0].ID != "perform" { + t.Fatalf("digest first row = %#v, want task-cited primary perform", digest.Evidence) + } + if !reflect.DeepEqual(digest.primaryIDs, wantPrimary) { + t.Fatalf("digest primary IDs = %v, want %v", digest.primaryIDs, wantPrimary) + } + merged := mergeLocalizationEvidenceDigestForTask(task, nil, digest) + if got := localizationFinalResponsePrimaryIDs(task, nil, merged.Evidence); !reflect.DeepEqual(got, wantPrimary) { + t.Fatalf("merged primary IDs = %v, want %v", got, wantPrimary) + } +} + +func TestSourceRangeEvidenceSurvivesFreezeAndDirectAdjacency(t *testing.T) { + task := "Investigate head() using src/Handler.php lines 185-187" + head := exploreTarget{node: &graph.Node{ + ID: "head", Name: "head", Kind: graph.KindFunction, + FilePath: "src/Entry.php", StartLine: 10, + }} + rangeCandidate := &rerank.Candidate{ + Node: &graph.Node{ + ID: "flush", Name: "flushBuffer", Kind: graph.KindMethod, + FilePath: "src/Handler.php", StartLine: 181, EndLine: 195, + }, + Signals: map[string]float64{exploreSourceRangeSignal: 1}, + } + rangeTarget := exploreTargetFromCandidate(rangeCandidate, "", false, false) + if !rangeTarget.sourceRange { + t.Fatal("source-range candidate signal was not projected onto the target") + } + targets := []exploreTarget{head} + for index := 0; index < localizationReplayEvidenceLimit-2; index++ { + targets = append(targets, exploreTarget{node: &graph.Node{ + ID: fmt.Sprintf("weak-%d", index), Name: "ordinaryCandidate", + Kind: graph.KindFunction, FilePath: "src/Weak.php", StartLine: index + 20, + }}) + } + targets = append(targets, rangeTarget) + draft := make([]exploreDraftEntry, 0, len(targets)) + for _, target := range targets { + draft = append(draft, exploreDraftEntry{node: target.node}) + } + selected := localizationEvidenceTargetsFromDraft(task, "", targets, draft) + if len(selected) != len(targets) { + t.Fatalf("selected rows = %d, want %d", len(selected), len(targets)) + } + if selected[0].node.ID != head.node.ID || selected[1].node.ID != rangeTarget.node.ID { + t.Fatalf("selected head = [%s, %s], want explicit head then source range", selected[0].node.ID, selected[1].node.ID) + } + + evidence := make([]localizationEvidence, 0, len(selected)) + for index, target := range selected { + evidence = append(evidence, localizationEvidence{ + Rank: index + 1, ID: target.node.ID, Name: target.node.Name, + Kind: string(target.node.Kind), File: target.node.FilePath, + Line: target.node.StartLine, EndLine: target.node.EndLine, + }) + } + relationIDs := make([]string, 0, localizationDirectAdjacencyCap) + reader := &localizationDirectAdjacencyReader{nodes: make(map[string]*graph.Node)} + for index := 0; index < localizationDirectAdjacencyCap; index++ { + id := fmt.Sprintf("adjacent-%d", index) + relationIDs = append(relationIDs, id) + reader.nodes[id] = nodeForDirectAdjacency(id, "adjacent", "src/Adjacent.php", 100+index) + } + evidence[0].Callees = relationIDs + envelope := localizationExploreEnvelope{Evidence: evidence} + digest := newLocalizationEvidenceDigestForTask(task, envelope) + freezeLocalizationPrimaryCohort(task, &envelope, digest) + + promoted, promotedDigest := promoteLocalizationDirectAdjacency( + task, envelope, reader, localizationReplayEvidenceLimit, 1<<20, digest, + ) + if len(promoted.Evidence) != localizationReplayEvidenceLimit { + t.Fatalf("evidence rows = %d, want %d", len(promoted.Evidence), localizationReplayEvidenceLimit) + } + if got := localizationFinalResponsePrimaryIDs(task, nil, promotedDigest.Evidence); !containsString(got, rangeTarget.node.ID) { + t.Fatalf("primary IDs = %v, missing source-range owner %q", got, rangeTarget.node.ID) + } + if !containsLocalizationEvidenceID(promoted.Evidence, rangeTarget.node.ID) { + t.Fatalf("promoted evidence lost source-range owner %q", rangeTarget.node.ID) + } +} + +func containsLocalizationEvidenceID(rows []localizationEvidence, id string) bool { + for _, row := range rows { + if row.ID == id { + return true + } + } + return false +} + +func nodeForDirectAdjacency(id, name, file string, line int) *graph.Node { + return &graph.Node{ID: id, Name: name, QualName: name, Kind: graph.KindFunction, FilePath: file, StartLine: line, EndLine: line + 1} +} + +func localizationEvidenceIDs(rows []localizationEvidence) []string { + ids := make([]string, 0, len(rows)) + for _, row := range rows { + ids = append(ids, row.ID) + } + return ids +} + +func assertLocalizationEvidenceIDs(t *testing.T, rows []localizationEvidence, want []string) { + t.Helper() + if got := localizationEvidenceIDs(rows); !reflect.DeepEqual(got, want) { + t.Fatalf("evidence IDs = %v, want %v", got, want) + } +} diff --git a/internal/mcp/localization_evidence_policy.go b/internal/mcp/localization_evidence_policy.go index 5a25e0c2..cc7d4773 100644 --- a/internal/mcp/localization_evidence_policy.go +++ b/internal/mcp/localization_evidence_policy.go @@ -10,6 +10,10 @@ import ( // The latter is a ranking decision; the former requires one of these bounded, // production-proven evidence shapes and must survive final response packing. const ( + // localizationProvenanceContentLiteral marks a bounded content/source hit + // that is useful answer evidence but is not, by itself, strong enough to + // make terminal enforcement safe. + localizationProvenanceContentLiteral = "content_literal" localizationProvenanceSourceLiteralCallee = "source_literal_callee" localizationProvenanceDivergentDefault = "divergent_default_owner" localizationProvenanceDivergentDefaultType = "divergent_default_type" @@ -395,6 +399,13 @@ func localizationTargetProvenance(completion localizationCompletion, target expl if localizationStrongSourceLiteralCallee(target) { return localizationProvenanceSourceLiteralCallee } + // Provenance describes how this authenticated graph identity was found; + // PRIMARY seating is a separate presentation decision. Explicit symbol or + // path anchors must not erase an exact content/source observation merely + // because that observation is not allowed to reserve an extra PRIMARY seat. + if target.sourceLiteral || target.exactContent { + return localizationProvenanceContentLiteral + } if target.typedAnchorProjection { return localizationProvenanceTypedAnchorProjection } diff --git a/internal/mcp/localization_evidence_policy_test.go b/internal/mcp/localization_evidence_policy_test.go index 31ce9820..31ba49fe 100644 --- a/internal/mcp/localization_evidence_policy_test.go +++ b/internal/mcp/localization_evidence_policy_test.go @@ -39,6 +39,150 @@ func requireLocalizationHostContractMatchesVisible( return host } +func TestLocalizationTargetProvenanceSeparatesLiteralObservationFromPrimarySeating(t *testing.T) { + node := &graph.Node{ID: "repo/service.go::Handle", Name: "Handle", FilePath: "repo/service.go"} + weak := exploreTarget{node: node, sourceLiteral: true} + if got := localizationTargetProvenance(localizationCompletion{}, weak); got != localizationProvenanceContentLiteral { + t.Fatalf("anchored literal provenance = %q, want %q", got, localizationProvenanceContentLiteral) + } + if localizationFinalResponsePrimaryProvenance(localizationDigestRow{Provenance: localizationProvenanceContentLiteral}) { + t.Fatal("anchored literal provenance must not reserve a PRIMARY seat") + } + + weak.literalPrimaryEligible = true + if got := localizationTargetProvenance(localizationCompletion{}, weak); got != localizationProvenanceContentLiteral { + t.Fatalf("unanchored literal provenance = %q, want %q", got, localizationProvenanceContentLiteral) + } + if !localizationFinalResponsePrimaryProvenance(localizationDigestRow{ + Provenance: localizationProvenanceContentLiteral, literalPrimaryEligible: true, + }) { + t.Fatal("eligible unanchored literal must retain its bounded PRIMARY reservation") + } + proof := localizationStrongEvidenceForCompletion( + newLocalizationCompletion(true, ""), []exploreTarget{weak}, + ) + if proof.provenance != "" { + t.Fatalf("unanchored content seating became strong terminal proof: %#v", proof) + } +} + +func TestLocalizationEligibleContentLiteralReachesPrimaryThroughPackedEnvelope(t *testing.T) { + const literalID = "repo/literal.go::LiteralOwner" + ordinary := []struct { + id string + name string + file string + }{ + {"repo/a.go::Alpha", "Alpha", "repo/a.go"}, + {"repo/b.go::Bravo", "Bravo", "repo/b.go"}, + {"repo/c.go::Charlie", "Charlie", "repo/c.go"}, + {"repo/d.go::Delta", "Delta", "repo/d.go"}, + {"repo/e.go::Echo", "Echo", "repo/e.go"}, + } + buildTargets := func(eligible bool) []exploreTarget { + targets := make([]exploreTarget, 0, len(ordinary)+1) + for _, item := range ordinary { + targets = append(targets, exploreTarget{node: &graph.Node{ + ID: item.id, Name: item.name, Kind: graph.KindFunction, FilePath: item.file, + }}) + } + targets = append(targets, exploreTarget{ + node: &graph.Node{ + ID: literalID, Name: "LiteralOwner", Kind: graph.KindFunction, FilePath: "repo/literal.go", + }, + exactContent: true, sourceLiteral: true, syntacticAnchor: !eligible, + literalPrimaryEligible: eligible, + }) + return targets + } + + for _, test := range []struct { + name string + eligible bool + }{ + {name: "authenticated unanchored literal reserves one primary seat", eligible: true}, + {name: "anchored literal observation remains rank only", eligible: false}, + } { + t.Run(test.name, func(t *testing.T) { + result, _, digest, _ := buildLocalizationExploreResultForTaskFinalized( + newLocalizationCompletion(true, ""), "", buildTargets(test.eligible), + exploreDefaultBudgetTokens, + ) + require.NotNil(t, digest) + require.LessOrEqual(t, len(digest.primaryIDs), localizationFinalResponsePrimaryLimit) + + body, ok := singleTextContent(result) + require.True(t, ok) + var envelope localizationExploreEnvelope + require.NoError(t, json.Unmarshal([]byte(body), &envelope)) + found := false + for _, row := range envelope.Evidence { + if row.ID == literalID { + found = true + require.Equal(t, localizationProvenanceContentLiteral, row.Provenance) + } + } + require.True(t, found, "literal evidence must survive production packing") + retainedEligible := false + for _, row := range digest.Evidence { + if row.ID == literalID { + retainedEligible = row.literalPrimaryEligible + } + } + require.Equal(t, test.eligible, retainedEligible) + require.Contains(t, digest.primaryIDs, literalID, "ordinary ranking may still seat either literal") + if test.eligible { + require.Equal(t, literalID, digest.primaryIDs[0], "the bounded literal reserve must pre-seat the eligible row") + } else { + require.NotEqual(t, literalID, digest.primaryIDs[0], "an ineligible literal must not claim the reserve") + } + }) + } +} + +func TestLocalizationLiteralProvenanceSurvivesFinalProjection(t *testing.T) { + tests := []struct { + name string + target exploreTarget + provenance string + enforceable bool + }{ + { + name: "authenticated content observation remains advisory", + target: exploreTarget{ + node: &graph.Node{ID: "repo/service.go::Handle", Name: "Handle", Kind: graph.KindFunction, FilePath: "repo/service.go"}, + source: "func Handle() {}", exactContent: true, + }, + provenance: localizationProvenanceContentLiteral, + }, + { + name: "unique graph-resolved literal callee remains strong", + target: exploreTarget{ + node: &graph.Node{ID: "repo/registry.go::Register", Name: "Register", Kind: graph.KindFunction, FilePath: "repo/registry.go"}, + source: "func Register(value string) {}", exactContent: true, + sourceLiteral: true, sourceLiteralCallee: true, + }, + provenance: localizationProvenanceSourceLiteralCallee, + enforceable: true, + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + result, _, _, completion := buildLocalizationExploreResultForTaskFinalized( + newLocalizationCompletion(true, ""), `locate the value "wire-key"`, + []exploreTarget{test.target}, exploreDefaultBudgetTokens, + ) + body, ok := singleTextContent(result) + require.True(t, ok) + var envelope localizationExploreEnvelope + require.NoError(t, json.Unmarshal([]byte(body), &envelope)) + require.Len(t, envelope.Evidence, 1) + require.Equal(t, test.provenance, envelope.Evidence[0].Provenance) + require.Equal(t, test.enforceable, completion.Enforceable) + }) + } +} + func TestLocalizationEvidencePolicyHoldsWeakOwnerForOneBoundedCallFromSQLiteCSharpIndex(t *testing.T) { root := t.TempDir() rel := "src/Humanizer/Localisation/Localiser.cs" @@ -113,7 +257,7 @@ func TestLocalizationEvidencePolicyHoldsWeakOwnerForOneBoundedCallFromSQLiteCSha require.Equal(t, localizationStateNeedsRecovery, envelope.Completion.State) require.False(t, envelope.Completion.Enforceable) require.Len(t, envelope.Evidence, 1) - require.Empty(t, envelope.Evidence[0].Provenance) + require.Equal(t, localizationProvenanceContentLiteral, envelope.Evidence[0].Provenance) wire, err := json.Marshal(result) require.NoError(t, err) diff --git a/internal/mcp/localization_explicit_anchor_test.go b/internal/mcp/localization_explicit_anchor_test.go index 03d0e2be..c350e43c 100644 --- a/internal/mcp/localization_explicit_anchor_test.go +++ b/internal/mcp/localization_explicit_anchor_test.go @@ -106,7 +106,7 @@ func TestExploreExactQualifiedAnchorCandidateFindsParserName(t *testing.T) { }) server := &Server{graph: g} got := server.exploreExactQualifiedAnchorCandidate( - context.Background(), anchors[1], query.QueryOptions{}, + context.Background(), anchors[1], nil, query.QueryOptions{}, map[string]struct{}{}, map[string]struct{}{}, ) if got == nil || got.Node == nil { @@ -134,6 +134,39 @@ func TestExploreSourceRangeSpecsPairPathsWithFollowingLines(t *testing.T) { } } +func TestExploreSourceRangeSpecsAcceptCompactUnicodeBounds(t *testing.T) { + tests := []struct { + name string + task string + }{ + {name: "compact ASCII", task: "src/Handler.php lines 185-187"}, + {name: "en dash", task: "src/Handler.php lines 185–187"}, + {name: "em dash", task: "src/Handler.php lines 185—187"}, + {name: "inline en dash", task: "src/Handler.php:185–187"}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + specs := exploreSourceRangeSpecs(test.task) + if len(specs) != 1 || specs[0].StartLine != 185 || specs[0].EndLine != 187 { + t.Fatalf("source range specs = %#v, want 185-187", specs) + } + }) + } +} + +func TestExploreSourceRangeDefinitionsCoverEveryDeclarationInRange(t *testing.T) { + idx := &fileSymbolIndex{} + first := &graph.Node{ID: "handler.php::before", Name: "before", Kind: graph.KindMethod, StartLine: 250, EndLine: 255} + second := &graph.Node{ID: "handler.php::testPassthruOnClose", Name: "testPassthruOnClose", Kind: graph.KindMethod, StartLine: 256, EndLine: 270} + idx.add(first) + idx.add(second) + idx.finalise() + got := exploreSourceRangeDefinitions(idx, 253, 265) + if len(got) != 2 || got[0].ID != first.ID || got[1].ID != second.ID { + t.Fatalf("source range definitions = %#v, want both declarations in cited range", got) + } +} + func TestExploreSourceRangeDefinitionsPreferNamedOwnerOverClosure(t *testing.T) { idx := &fileSymbolIndex{} idx.add(&graph.Node{ID: "handler.php::flushBuffer", Name: "flushBuffer", Kind: graph.KindMethod, StartLine: 170, EndLine: 200}) @@ -151,7 +184,7 @@ func TestPromoteExploreSourceRangeCandidatesMapsToEnclosingMethod(t *testing.T) stop := store.GetNode(navFindMethod(t, store, "Stop")) task := fmt.Sprintf("svc.go\nLines %d to %d", start.StartLine, start.EndLine) ordinary := []*rerank.Candidate{{Node: stop, TextRank: 0, VectorRank: -1}} - got := server.promoteExploreSourceRangeCandidates(context.Background(), task, ordinary, query.QueryOptions{}) + got := server.promoteExploreSourceRangeCandidates(context.Background(), task, ordinary, store, query.QueryOptions{}) if len(got) != 2 || got[0].Node.ID != start.ID || got[1].Node.ID != stop.ID { t.Fatalf("promoted candidates = %#v, want cited Start then ordinary Stop", got) } @@ -166,7 +199,7 @@ func TestPromoteExploreSourceRangeCandidatesStripsIsolatedRepoLabel(t *testing.T stop := store.GetNode(navFindMethod(t, store, "Stop")) task := fmt.Sprintf("monolog/svc.go\nLines %d to %d", start.StartLine, start.EndLine) ordinary := []*rerank.Candidate{{Node: stop, TextRank: 0, VectorRank: -1}} - got := server.promoteExploreSourceRangeCandidates(context.Background(), task, ordinary, query.QueryOptions{}) + got := server.promoteExploreSourceRangeCandidates(context.Background(), task, ordinary, store, query.QueryOptions{}) if len(got) != 2 || got[0].Node.ID != start.ID || got[1].Node.ID != stop.ID { t.Fatalf("promoted candidates = %#v, want repo-prefixed citation to resolve indexed svc.go", got) } diff --git a/internal/mcp/localization_file_budget.go b/internal/mcp/localization_file_budget.go new file mode 100644 index 00000000..83a3ec93 --- /dev/null +++ b/internal/mcp/localization_file_budget.go @@ -0,0 +1,80 @@ +package mcp + +import ( + "context" + "sync" +) + +type localizationFileRequestBudgetContextKey struct{} + +// localizationFileRequestBudget is shared by every bounded file-summary +// lookup descended from one MCP request context. Reservations make parallel +// builders conservative: their combined in-flight and consumed allowance can +// never exceed the request cap. +type localizationFileRequestBudget struct { + mu sync.Mutex + remaining int +} + +func withLocalizationFileRequestBudget(ctx context.Context) context.Context { + if ctx == nil { + return nil + } + if budget, _ := ctx.Value(localizationFileRequestBudgetContextKey{}).(*localizationFileRequestBudget); budget != nil { + return ctx + } + return context.WithValue(ctx, localizationFileRequestBudgetContextKey{}, &localizationFileRequestBudget{ + remaining: localizationFileRequestLimit, + }) +} + +// localizationFileBudgetFor returns the request-owned counter when middleware +// installed one. Direct compatibility callers receive a private counter, so +// they remain bounded without sharing state across unrelated calls. +func localizationFileBudgetFor(ctx context.Context) *localizationFileRequestBudget { + if ctx != nil { + if budget, _ := ctx.Value(localizationFileRequestBudgetContextKey{}).(*localizationFileRequestBudget); budget != nil { + return budget + } + } + return &localizationFileRequestBudget{remaining: localizationFileRequestLimit} +} + +func (budget *localizationFileRequestBudget) reserve(limit int) int { + if budget == nil || limit <= 0 { + return 0 + } + budget.mu.Lock() + defer budget.mu.Unlock() + if budget.remaining <= 0 { + return 0 + } + if limit > budget.remaining { + limit = budget.remaining + } + budget.remaining -= limit + return limit +} + +// finish refunds the portion of a successful reservation that storage proved +// it did not consume. Callers deliberately do not invoke this after errors or +// cancellation: the amount inspected is then uncertain, so the full +// reservation stays charged and subsequent attribution fails closed. +func (budget *localizationFileRequestBudget) finish(reserved, consumed int) { + if budget == nil || reserved <= 0 { + return + } + if consumed < 0 { + consumed = 0 + } + if consumed > reserved { + consumed = reserved + } + refund := reserved - consumed + if refund == 0 { + return + } + budget.mu.Lock() + budget.remaining += refund + budget.mu.Unlock() +} diff --git a/internal/mcp/localization_file_budget_test.go b/internal/mcp/localization_file_budget_test.go new file mode 100644 index 00000000..10eb8105 --- /dev/null +++ b/internal/mcp/localization_file_budget_test.go @@ -0,0 +1,295 @@ +package mcp + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "path/filepath" + "strings" + "sync" + "testing" + + mcpgo "github.com/mark3labs/mcp-go/mcp" + + "github.com/zzet/gortex/internal/graph" + "github.com/zzet/gortex/internal/query" +) + +type concurrentBoundedFileStoreProbe struct { + graph.Store + + mu sync.Mutex + calls []boundedFileCall + read func(context.Context, string, graph.LocalizationNodeScope, int) (graph.BoundedNodeProjection, error) +} + +func (store *concurrentBoundedFileStoreProbe) FindFileNodesBounded( + ctx context.Context, + path string, + scope graph.LocalizationNodeScope, + limit int, +) (graph.BoundedNodeProjection, error) { + store.mu.Lock() + store.calls = append(store.calls, boundedFileCall{path: path, scope: scope, limit: limit}) + store.mu.Unlock() + return store.read(ctx, path, scope, limit) +} + +func (store *concurrentBoundedFileStoreProbe) snapshotCalls() []boundedFileCall { + store.mu.Lock() + defer store.mu.Unlock() + return append([]boundedFileCall(nil), store.calls...) +} + +func budgetTestNode(path string) *graph.Node { + return &graph.Node{ + ID: path + "::owner", Name: "owner", Kind: graph.KindFunction, + FilePath: path, StartLine: 1, EndLine: 1, + } +} + +func TestWrappedToolSharesAndRenewsLocalizationFileBudget(t *testing.T) { + probe := &concurrentBoundedFileStoreProbe{Store: graph.New()} + probe.read = func(_ context.Context, path string, _ graph.LocalizationNodeScope, limit int) (graph.BoundedNodeProjection, error) { + return graph.BoundedNodeProjection{Nodes: []*graph.Node{budgetTestNode(path)}, Total: limit}, nil + } + server := fastPathTestServer(t) + server.graph = probe + + handler := func(ctx context.Context, _ mcpgo.CallToolRequest) (*mcpgo.CallToolResult, error) { + first := server.buildFileSymbolIndexForOrderedPathsScopedContext( + ctx, []string{"repo/a.go", "repo/b.go"}, query.QueryOptions{}, + ) + if first["repo/a.go"] == nil || first["repo/b.go"] == nil { + return nil, errors.New("first builder invocation did not complete") + } + // Facade and overlay preparation may defensively seed again. The helper + // must preserve the existing pointer instead of resetting its allowance. + ctx = withLocalizationFileRequestBudget(ctx) + second := server.buildFileSymbolIndexForOrderedPathsScopedContext( + ctx, []string{"repo/c.go", "repo/d.go", "repo/e.go"}, query.QueryOptions{}, + ) + if index := second["repo/e.go"]; index == nil || !index.saturated { + return nil, fmt.Errorf("shared budget did not saturate final path: %#v", index) + } + return mcpgo.NewToolResultText("ok"), nil + } + + callWrapped(t, server, handler, "search_text") + if calls := probe.snapshotCalls(); len(calls) != 4 { + t.Fatalf("first wrapped request made %d storage calls, want 4", len(calls)) + } + callWrapped(t, server, handler, "search_text") + calls := probe.snapshotCalls() + if len(calls) != 8 { + t.Fatalf("second wrapped request inherited prior budget: calls=%d, want 8", len(calls)) + } + for _, call := range calls { + if call.limit != localizationFileNodeLimit { + t.Fatalf("wrapped call limit = %d, want %d", call.limit, localizationFileNodeLimit) + } + } +} + +func TestLocalizationFileBudgetSerializesParallelReservations(t *testing.T) { + probe := &concurrentBoundedFileStoreProbe{Store: graph.New()} + probe.read = func(_ context.Context, path string, _ graph.LocalizationNodeScope, limit int) (graph.BoundedNodeProjection, error) { + return graph.BoundedNodeProjection{Nodes: []*graph.Node{budgetTestNode(path)}, Total: limit}, nil + } + server := &Server{graph: probe} + ctx := withLocalizationFileRequestBudget(context.Background()) + + const workers = 8 + results := make([]map[string]*fileSymbolIndex, workers) + start := make(chan struct{}) + var wg sync.WaitGroup + for worker := 0; worker < workers; worker++ { + worker := worker + wg.Add(1) + go func() { + defer wg.Done() + <-start + path := fmt.Sprintf("repo/worker-%d.go", worker) + results[worker] = server.buildFileSymbolIndexForOrderedPathsScopedContext( + ctx, []string{path}, query.QueryOptions{}, + ) + }() + } + close(start) + wg.Wait() + + calls := probe.snapshotCalls() + consumed := 0 + for _, call := range calls { + consumed += call.limit + } + if len(calls) != localizationFileRequestLimit/localizationFileNodeLimit || consumed != localizationFileRequestLimit { + t.Fatalf("parallel reservations: calls=%d consumed=%d, want calls=%d consumed=%d", + len(calls), consumed, localizationFileRequestLimit/localizationFileNodeLimit, localizationFileRequestLimit) + } + saturated := 0 + for worker, indexes := range results { + path := fmt.Sprintf("repo/worker-%d.go", worker) + if index := indexes[path]; index != nil && index.saturated { + saturated++ + } + } + if saturated != workers-len(calls) { + t.Fatalf("parallel saturated paths = %d, want %d", saturated, workers-len(calls)) + } +} + +func TestLocalizationFileBudgetRefundsSuccessfulShortPages(t *testing.T) { + probe := &concurrentBoundedFileStoreProbe{Store: graph.New()} + probe.read = func(_ context.Context, path string, _ graph.LocalizationNodeScope, limit int) (graph.BoundedNodeProjection, error) { + total := limit + if path == "repo/small.go" { + total = 1 + } + return graph.BoundedNodeProjection{Nodes: []*graph.Node{budgetTestNode(path)}, Total: total}, nil + } + server := &Server{graph: probe} + ctx := withLocalizationFileRequestBudget(context.Background()) + + server.buildFileSymbolIndexForOrderedPathsScopedContext( + ctx, []string{"repo/small.go"}, query.QueryOptions{}, + ) + paths := []string{"repo/a.go", "repo/b.go", "repo/c.go", "repo/d.go", "repo/e.go"} + indexes := server.buildFileSymbolIndexForOrderedPathsScopedContext(ctx, paths, query.QueryOptions{}) + + calls := probe.snapshotCalls() + wantLimits := []int{1024, 1024, 1024, 1024, 1023} + if len(calls) != len(wantLimits) { + t.Fatalf("calls after short-page refund = %d, want %d", len(calls), len(wantLimits)) + } + for index, want := range wantLimits { + if calls[index].limit != want { + t.Fatalf("call %d limit = %d, want %d", index, calls[index].limit, want) + } + } + if index := indexes["repo/e.go"]; index == nil || !index.saturated { + t.Fatalf("post-refund exhausted path = %#v, want saturated", index) + } +} + +func TestLocalizationFileBudgetChargesUncertainErrorReservation(t *testing.T) { + probe := &concurrentBoundedFileStoreProbe{Store: graph.New()} + probe.read = func(_ context.Context, path string, _ graph.LocalizationNodeScope, limit int) (graph.BoundedNodeProjection, error) { + switch path { + case "repo/ok.go": + return graph.BoundedNodeProjection{Nodes: []*graph.Node{budgetTestNode(path)}, Total: 1}, nil + case "repo/boom.go": + return graph.BoundedNodeProjection{}, errors.New("read failed") + default: + return graph.BoundedNodeProjection{Nodes: []*graph.Node{budgetTestNode(path)}, Total: limit}, nil + } + } + server := &Server{graph: probe} + ctx := withLocalizationFileRequestBudget(context.Background()) + + first := server.buildFileSymbolIndexForOrderedPathsScopedContext( + ctx, []string{"repo/ok.go", "repo/boom.go", "repo/later.go"}, query.QueryOptions{}, + ) + if first["repo/ok.go"] == nil || first["repo/ok.go"].saturated { + t.Fatalf("completed path was not preserved: %#v", first["repo/ok.go"]) + } + for _, path := range []string{"repo/boom.go", "repo/later.go"} { + if index := first[path]; index == nil || !index.saturated { + t.Fatalf("error path %q = %#v, want saturated", path, index) + } + } + + secondPaths := []string{"repo/a.go", "repo/b.go", "repo/c.go", "repo/d.go"} + second := server.buildFileSymbolIndexForOrderedPathsScopedContext(ctx, secondPaths, query.QueryOptions{}) + calls := probe.snapshotCalls() + wantLimits := []int{1024, 1024, 1024, 1024, 1023} + if len(calls) != len(wantLimits) { + t.Fatalf("calls after uncertain error = %d, want %d", len(calls), len(wantLimits)) + } + for index, want := range wantLimits { + if calls[index].limit != want { + t.Fatalf("call %d limit = %d, want %d", index, calls[index].limit, want) + } + } + if index := second["repo/d.go"]; index == nil || !index.saturated { + t.Fatalf("error reservation was refunded: final path=%#v", index) + } +} + +func TestRangeLoweringSharesExhaustedRequestFileBudget(t *testing.T) { + server, graphStore := setupNavServer(t) + start := graphStore.GetNode(navFindMethod(t, graphStore, "Start")) + if start == nil { + t.Fatal("Start node is missing") + } + + req := mcpgo.CallToolRequest{} + req.Params.Arguments = map[string]any{ + "source": "ranges", + "path": "svc.go", + "start_line": float64(start.StartLine), + "end_line": float64(start.EndLine), + } + exhaustedContext := func() context.Context { + ctx := withLocalizationFileRequestBudget(context.Background()) + if got := localizationFileBudgetFor(ctx).reserve(localizationFileRequestLimit); got != localizationFileRequestLimit { + t.Fatalf("reserved %d, want %d", got, localizationFileRequestLimit) + } + return ctx + } + + t.Run("symbols for ranges", func(t *testing.T) { + result, err := server.handleSymbolsForRanges(exhaustedContext(), req) + if err != nil || result == nil || result.IsError { + t.Fatalf("handleSymbolsForRanges = (%#v, %v)", result, err) + } + var payload map[string]any + text := result.Content[0].(mcpgo.TextContent).Text + if err := json.Unmarshal([]byte(text), &payload); err != nil { + t.Fatal(err) + } + if symbols, _ := payload["symbols"].([]any); len(symbols) != 0 { + t.Fatalf("exhausted request budget admitted range symbols: %#v", symbols) + } + if unresolved, _ := payload["unresolved_files"].([]any); len(unresolved) != 1 || unresolved[0] != "svc.go" { + t.Fatalf("saturated range unresolved files = %#v, want [svc.go]", unresolved) + } + }) + + t.Run("change contract range source", func(t *testing.T) { + prediction, err := server.lowerRangeSource(exhaustedContext(), req) + if err == nil || prediction != nil || !strings.Contains(err.Error(), "bounded range ownership is unavailable for: svc.go") { + t.Fatalf("change-contract saturation = (%#v, %v)", prediction, err) + } + }) + + t.Run("change contract workspace edit", func(t *testing.T) { + editReq := mcpgo.CallToolRequest{} + editReq.Params.Arguments = map[string]any{ + "source": "edit", + "workspace_edit": buildSingleFileEdit( + filepath.Join(server.indexer.RootPath(), "svc.go"), + "package svc\n\nfunc Changed() {}\n", + ), + } + prediction, err := server.lowerEditSource(exhaustedContext(), editReq) + if err == nil || prediction != nil || !strings.Contains(err.Error(), "bounded range ownership is unavailable for: svc.go") { + t.Fatalf("workspace-edit saturation = (%#v, %v)", prediction, err) + } + }) + + t.Run("ordinary unresolved remains tolerated", func(t *testing.T) { + missingReq := mcpgo.CallToolRequest{} + missingReq.Params.Arguments = map[string]any{ + "source": "ranges", "path": "missing.go", "start_line": float64(1), + } + prediction, err := server.lowerRangeSource(withLocalizationFileRequestBudget(context.Background()), missingReq) + if err != nil || prediction == nil { + t.Fatalf("ordinary unresolved range = (%#v, %v), want tolerated empty prediction", prediction, err) + } + if len(prediction.changedIDs) != 0 || len(prediction.changed) != 0 { + t.Fatalf("ordinary unresolved range changed prediction: %#v", prediction) + } + }) +} diff --git a/internal/mcp/localization_file_declarations.go b/internal/mcp/localization_file_declarations.go new file mode 100644 index 00000000..3eae49fc --- /dev/null +++ b/internal/mcp/localization_file_declarations.go @@ -0,0 +1,169 @@ +package mcp + +import ( + "context" + "strings" + + "github.com/zzet/gortex/internal/graph" +) + +// localizationFileDeclarations carries a bounded declaration page together with +// the count observed while scanning the file. DeclaredKnown distinguishes an +// exact count from a saturation lower bound. Truncated means Nodes is only a +// prefix and Declared is therefore an "at least" count. +type localizationFileDeclarations struct { + Nodes []*graph.Node + Declared int + DeclaredKnown bool + Truncated bool +} + +var localizationDeclarationExcludedKinds = []graph.NodeKind{ + graph.KindFile, graph.KindImport, graph.KindLocal, graph.KindParam, + graph.KindClosure, graph.KindGenericParam, graph.KindBuiltin, +} + +// localizationFileDeclarationCache keeps one request's bounded lightweight +// declaration pages. It deliberately depends on BoundedFileNodeReader: a store +// without that capability fails closed instead of falling back to full-row +// GetFileNodes decoding. +type localizationFileDeclarationCache struct { + ctx context.Context + reader graph.Reader + bounded graph.BoundedFileNodeReader + scope graph.LocalizationNodeScope + budget *localizationFileRequestBudget + byFile map[string]localizationFileDeclarations + definitionLimit int +} + +func newLocalizationFileDeclarationCache( + ctx context.Context, + reader graph.Reader, + scope graph.LocalizationNodeScope, +) *localizationFileDeclarationCache { + return newBoundedLocalizationFileDeclarationCache(ctx, reader, scope, localizationFileNodeLimit) +} + +func newBoundedLocalizationFileDeclarationCache( + ctx context.Context, + reader graph.Reader, + scope graph.LocalizationNodeScope, + definitionLimit int, +) *localizationFileDeclarationCache { + if ctx == nil { + ctx = context.Background() + } + bounded, _ := reader.(graph.BoundedFileNodeReader) + return &localizationFileDeclarationCache{ + ctx: ctx, + reader: reader, + bounded: bounded, + scope: localizationDeclarationScope(scope), + budget: localizationFileBudgetFor(ctx), + byFile: make(map[string]localizationFileDeclarations), + definitionLimit: min(max(definitionLimit, 0), localizationFileNodeLimit), + } +} + +func localizationDeclarationScope(scope graph.LocalizationNodeScope) graph.LocalizationNodeScope { + excluded := make(map[graph.NodeKind]bool, len(scope.ExcludeKinds)+len(localizationDeclarationExcludedKinds)) + for kind, omit := range scope.ExcludeKinds { + if omit { + excluded[kind] = true + } + } + for _, kind := range localizationDeclarationExcludedKinds { + excluded[kind] = true + } + scope.ExcludeKinds = excluded + return scope +} + +func (cache *localizationFileDeclarationCache) definitions(file string) localizationFileDeclarations { + if cache == nil { + return localizationFileDeclarations{} + } + return cache.definitionsAtLimit(file, cache.definitionLimit) +} + +// outlineDefinitions gives the first two distinct page files the full bounded +// declaration projection and spends only a shallow page on every later file. +// The first read wins: a lower-ranked cached page is never silently upgraded by +// a later consumer, so one file cannot be charged twice against the request. +func (cache *localizationFileDeclarationCache) outlineDefinitions(file string, rank int) localizationFileDeclarations { + if cache == nil { + return localizationFileDeclarations{} + } + limit := localizationOutlineFileFetchLimit(rank) + if cache.definitionLimit > 0 { + limit = min(limit, cache.definitionLimit) + } + return cache.definitionsAtLimit(file, limit) +} + +func (cache *localizationFileDeclarationCache) definitionsAtLimit(file string, limit int) localizationFileDeclarations { + file = strings.TrimSpace(file) + if cache == nil || file == "" { + return localizationFileDeclarations{} + } + if declarations, cached := cache.byFile[file]; cached { + return declarations + } + declarations := cache.readDefinitions(file, limit) + cache.byFile[file] = declarations + return declarations +} + +// boundedDefinitions avoids retaining a generated file's whole declaration set +// for supporting body evidence. A cached outline page can be sliced safely; an +// uncached file gets its own bounded projection and does not poison the outline +// cache with a shallower page. +func (cache *localizationFileDeclarationCache) boundedDefinitions(file string, limit int) []*graph.Node { + file = strings.TrimSpace(file) + if cache == nil || file == "" || limit <= 0 { + return nil + } + if declarations, cached := cache.byFile[file]; cached { + return declarations.Nodes[:min(len(declarations.Nodes), limit)] + } + return cache.readDefinitions(file, limit).Nodes +} + +func (cache *localizationFileDeclarationCache) readDefinitions(file string, limit int) localizationFileDeclarations { + incomplete := localizationFileDeclarations{Truncated: true} + if cache == nil || cache.bounded == nil || cache.ctx == nil || cache.ctx.Err() != nil || file == "" { + return incomplete + } + if limit <= 0 || limit > localizationFileNodeLimit { + limit = localizationFileNodeLimit + } + reserved := cache.budget.reserve(limit) + if reserved <= 0 { + return incomplete + } + page, err := cache.bounded.FindFileNodesBounded(cache.ctx, file, cache.scope, reserved) + if err != nil || cache.ctx.Err() != nil { + // The amount inspected is uncertain, so the reservation remains charged. + return incomplete + } + consumed := max(page.Total, len(page.Nodes)) + cache.budget.finish(reserved, consumed) + + nodes := page.Nodes + truncated := page.Truncated + if len(nodes) > reserved { + nodes = nodes[:reserved] + truncated = true + } + declared := max(page.Total, len(nodes)) + if truncated { + declared = max(declared, len(nodes)+1) + } + return localizationFileDeclarations{ + Nodes: nodes, + Declared: declared, + DeclaredKnown: !truncated, + Truncated: truncated, + } +} diff --git a/internal/mcp/localization_file_declarations_test.go b/internal/mcp/localization_file_declarations_test.go new file mode 100644 index 00000000..aa79fdc4 --- /dev/null +++ b/internal/mcp/localization_file_declarations_test.go @@ -0,0 +1,269 @@ +package mcp + +import ( + "context" + "errors" + "fmt" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/zzet/gortex/internal/graph" +) + +type localizationDeclarationCall struct { + file string + scope graph.LocalizationNodeScope + limit int +} + +type localizationDeclarationSpyReader struct { + graph.Reader + files map[string][]*graph.Node + totals map[string]int + fileCalls map[string]int + calls []localizationDeclarationCall + err error +} + +func (reader *localizationDeclarationSpyReader) GetFileNodes(string) []*graph.Node { + panic("localization declarations must never use the legacy full-row reader") +} + +func (reader *localizationDeclarationSpyReader) FindFileNodesBounded( + _ context.Context, + file string, + scope graph.LocalizationNodeScope, + limit int, +) (graph.BoundedNodeProjection, error) { + if reader.fileCalls == nil { + reader.fileCalls = make(map[string]int) + } + reader.fileCalls[file]++ + reader.calls = append(reader.calls, localizationDeclarationCall{file: file, scope: scope, limit: limit}) + if reader.err != nil { + return graph.BoundedNodeProjection{}, reader.err + } + + admitted := make([]*graph.Node, 0, min(len(reader.files[file]), limit)) + for _, node := range reader.files[file] { + if scope.Allows(node) { + admitted = append(admitted, node) + } + } + total := len(admitted) + if configured, ok := reader.totals[file]; ok { + total = configured + } + truncated := total > limit + observed := total + if observed > limit+1 { + observed = limit + 1 + } + if len(admitted) > limit { + admitted = admitted[:limit] + } + return graph.BoundedNodeProjection{Nodes: admitted, Total: observed, Truncated: truncated}, nil +} + +func newLocalizationDeclarationTestCache(reader graph.Reader) *localizationFileDeclarationCache { + return newLocalizationFileDeclarationCache( + context.Background(), reader, graph.LocalizationNodeScope{}, + ) +} + +func TestLocalizationFileDeclarationCacheReadsEachFileOnce(t *testing.T) { + function := &graph.Node{ID: "src/a.go::run", Name: "run", Kind: graph.KindFunction, FilePath: "src/a.go"} + reader := &localizationDeclarationSpyReader{files: map[string][]*graph.Node{"src/a.go": {function}}} + cache := newLocalizationDeclarationTestCache(reader) + + require.Equal(t, []*graph.Node{function}, cache.definitions("src/a.go").Nodes) + require.Equal(t, []*graph.Node{function}, cache.definitions("src/a.go").Nodes) + require.Equal(t, 1, reader.fileCalls["src/a.go"]) + require.Equal(t, localizationFileNodeLimit, reader.calls[0].limit) +} + +func TestLocalizationFileDeclarationCacheCachesMissesAndDistinctFiles(t *testing.T) { + reader := &localizationDeclarationSpyReader{files: map[string][]*graph.Node{}} + cache := newLocalizationDeclarationTestCache(reader) + + require.Empty(t, cache.definitions("src/missing.go").Nodes) + require.Empty(t, cache.definitions("src/missing.go").Nodes) + require.Empty(t, cache.definitions("src/other.go").Nodes) + require.Equal(t, map[string]int{"src/missing.go": 1, "src/other.go": 1}, reader.fileCalls) +} + +func TestLocalizationFileDeclarationCachePushesScopeAndDefinitionExclusions(t *testing.T) { + first := &graph.Node{ID: "src/a.go::first", Name: "first", Kind: graph.KindFunction, FilePath: "src/a.go", RepoPrefix: "repo"} + second := &graph.Node{ID: "src/a.go::Thing", Name: "Thing", Kind: graph.KindType, FilePath: "src/a.go", RepoPrefix: "repo"} + reader := &localizationDeclarationSpyReader{files: map[string][]*graph.Node{"src/a.go": { + {ID: "src/a.go", Name: "a.go", Kind: graph.KindFile, FilePath: "src/a.go", RepoPrefix: "repo"}, + first, + {ID: "src/a.go::arg", Name: "arg", Kind: graph.KindParam, FilePath: "src/a.go", RepoPrefix: "repo"}, + second, + {ID: "foreign/a.go::hidden", Name: "hidden", Kind: graph.KindFunction, FilePath: "src/a.go", RepoPrefix: "foreign"}, + }}} + cache := newLocalizationFileDeclarationCache( + context.Background(), reader, + graph.LocalizationNodeScope{RepoAllow: map[string]bool{"repo": true}}, + ) + + require.Equal(t, []*graph.Node{first, second}, cache.definitions("src/a.go").Nodes) + require.Len(t, reader.calls, 1) + for _, kind := range localizationDeclarationExcludedKinds { + require.Truef(t, reader.calls[0].scope.ExcludeKinds[kind], "kind %q was not pushed down", kind) + } + require.True(t, reader.calls[0].scope.RepoAllow["repo"]) +} + +func TestLocalizationFileDeclarationCacheBoundsRetainedDefinitions(t *testing.T) { + first := &graph.Node{ID: "src/a.go::first", Name: "first", Kind: graph.KindFunction, FilePath: "src/a.go"} + second := &graph.Node{ID: "src/a.go::second", Name: "second", Kind: graph.KindFunction, FilePath: "src/a.go"} + third := &graph.Node{ID: "src/a.go::third", Name: "third", Kind: graph.KindFunction, FilePath: "src/a.go"} + reader := &localizationDeclarationSpyReader{files: map[string][]*graph.Node{"src/a.go": {first, second, third}}} + cache := newBoundedLocalizationFileDeclarationCache( + context.Background(), reader, graph.LocalizationNodeScope{}, 2, + ) + + firstRead := cache.definitions("src/a.go") + require.Equal(t, []*graph.Node{first, second}, firstRead.Nodes) + require.Equal(t, 3, firstRead.Declared) + require.False(t, firstRead.DeclaredKnown) + require.True(t, firstRead.Truncated) + require.Equal(t, []*graph.Node{first, second}, cache.definitions("src/a.go").Nodes) + require.Equal(t, 1, reader.fileCalls["src/a.go"]) + require.Equal(t, 2, reader.calls[0].limit) +} + +func TestBoundedLocalizationDeclarationsRenderHonestLowerBounds(t *testing.T) { + const ( + file = "src/generated.go" + declared = 140 + retentionLimit = 128 + ) + nodes := make([]*graph.Node, 0, declared) + for index := 0; index < declared; index++ { + nodes = append(nodes, &graph.Node{ + ID: file + fmt.Sprintf("::declaration%03d", index), Name: fmt.Sprintf("declaration%03d", index), + Kind: graph.KindFunction, FilePath: file, StartLine: index + 1, + }) + } + reader := &localizationDeclarationSpyReader{files: map[string][]*graph.Node{file: nodes}} + cache := newBoundedLocalizationFileDeclarationCache( + context.Background(), reader, graph.LocalizationNodeScope{}, retentionLimit, + ) + + page := localizationPageOutlineProvider( + nil, []exploreTarget{{node: nodes[0]}}, nil, cache.definitions, + )() + + require.NotNil(t, page) + require.NotNil(t, page.Leading) + require.Equal(t, retentionLimit+1, page.Leading.Declared) + require.Equal(t, 1, page.Leading.Elided) + require.True(t, page.Leading.Truncated) + require.Len(t, page.Leading.Rows, retentionLimit) + retained := cache.byFile[file] + require.False(t, retained.DeclaredKnown) + require.True(t, retained.Truncated) + require.Len(t, retained.Nodes, retentionLimit) + require.Equal(t, 1, reader.fileCalls[file]) +} + +func TestLocalizationFileDeclarationCacheFailsClosed(t *testing.T) { + legacy := &struct{ graph.Reader }{Reader: graph.New()} + unsupported := newLocalizationDeclarationTestCache(legacy).definitions("src/a.go") + require.Empty(t, unsupported.Nodes) + require.True(t, unsupported.Truncated) + require.False(t, unsupported.DeclaredKnown) + + failedReader := &localizationDeclarationSpyReader{err: errors.New("read failed")} + failed := newLocalizationDeclarationTestCache(failedReader).definitions("src/a.go") + require.Empty(t, failed.Nodes) + require.True(t, failed.Truncated) + + cancelledCtx, cancel := context.WithCancel(context.Background()) + cancel() + cancelled := newLocalizationFileDeclarationCache( + cancelledCtx, &localizationDeclarationSpyReader{}, graph.LocalizationNodeScope{}, + ).definitions("src/a.go") + require.Empty(t, cancelled.Nodes) + require.True(t, cancelled.Truncated) +} + +func TestLocalizationFileDeclarationCacheSharesRequestBudgetInPriorityOrder(t *testing.T) { + reader := &localizationDeclarationSpyReader{files: make(map[string][]*graph.Node)} + ctx := withLocalizationFileRequestBudget(context.Background()) + cache := newLocalizationFileDeclarationCache(ctx, reader, graph.LocalizationNodeScope{}) + for index := 0; index < 5; index++ { + file := fmt.Sprintf("src/file-%d.go", index) + for declaration := 0; declaration < localizationFileNodeLimit; declaration++ { + reader.files[file] = append(reader.files[file], &graph.Node{ + ID: fmt.Sprintf("%s::owner-%04d", file, declaration), Name: "owner", + Kind: graph.KindFunction, FilePath: file, + }) + } + page := cache.definitions(file) + if index < localizationFileRequestLimit/localizationFileNodeLimit { + require.False(t, page.Truncated) + require.Len(t, page.Nodes, localizationFileNodeLimit) + } else { + require.True(t, page.Truncated) + require.Empty(t, page.Nodes) + } + } + require.Len(t, reader.calls, localizationFileRequestLimit/localizationFileNodeLimit) + for index, call := range reader.calls { + require.Equal(t, fmt.Sprintf("src/file-%d.go", index), call.file) + require.Equal(t, localizationFileNodeLimit, call.limit) + } +} + +func TestLocalizationFileDeclarationCacheBudgetsRankedOutlinesWithoutUpgrades(t *testing.T) { + reader := &localizationDeclarationSpyReader{files: make(map[string][]*graph.Node)} + ctx := withLocalizationFileRequestBudget(context.Background()) + cache := newLocalizationFileDeclarationCache(ctx, reader, graph.LocalizationNodeScope{}) + + for rank := 0; rank < localizationOutlineFileCap; rank++ { + file := fmt.Sprintf("src/outline-%d.go", rank) + limit := localizationOutlineFileFetchLimit(rank) + declared := limit + 1 + if rank == 0 { + declared = limit + } + for index := 0; index < declared; index++ { + reader.files[file] = append(reader.files[file], &graph.Node{ + ID: fmt.Sprintf("%s::owner-%04d", file, index), + Name: fmt.Sprintf("owner%04d", index), Kind: graph.KindFunction, + FilePath: file, StartLine: index + 1, + }) + } + + page := cache.outlineDefinitions(file, rank) + require.Len(t, page.Nodes, limit) + require.Equal(t, limit, reader.calls[rank].limit) + if rank == 0 { + require.False(t, page.Truncated) + require.True(t, page.DeclaredKnown) + require.Equal(t, localizationFileNodeLimit, page.Declared) + } else { + require.True(t, page.Truncated) + require.False(t, page.DeclaredKnown) + require.Equal(t, limit+1, page.Declared) + } + } + + require.Len(t, reader.calls, localizationOutlineFileCap) + shallow := cache.outlineDefinitions("src/outline-2.go", 0) + require.Len(t, shallow.Nodes, localizationOutlineTrailingFileFetchLimit) + require.Len(t, reader.calls, localizationOutlineFileCap, "a cached shallow page must not be upgraded") + + wantConsumed := localizationOutlineProtectedFileCount*localizationFileNodeLimit + + (localizationOutlineFileCap-localizationOutlineProtectedFileCount)*localizationOutlineTrailingFileFetchLimit + require.Equal(t, 3072, wantConsumed) + cache.budget.mu.Lock() + remaining := cache.budget.remaining + cache.budget.mu.Unlock() + require.Equal(t, localizationFileRequestLimit-wantConsumed, remaining) + require.LessOrEqual(t, wantConsumed, 3072) +} diff --git a/internal/mcp/localization_file_index_test.go b/internal/mcp/localization_file_index_test.go new file mode 100644 index 00000000..ce0f8e62 --- /dev/null +++ b/internal/mcp/localization_file_index_test.go @@ -0,0 +1,144 @@ +package mcp + +import ( + "context" + "fmt" + "testing" + + "github.com/zzet/gortex/internal/graph" + "github.com/zzet/gortex/internal/query" +) + +type boundedFileCall struct { + path string + scope graph.LocalizationNodeScope + limit int +} + +type boundedFileStoreProbe struct { + graph.Store + calls []boundedFileCall + read func(context.Context, string, graph.LocalizationNodeScope, int) (graph.BoundedNodeProjection, error) +} + +func (store *boundedFileStoreProbe) FindFileNodesBounded( + ctx context.Context, + path string, + scope graph.LocalizationNodeScope, + limit int, +) (graph.BoundedNodeProjection, error) { + store.calls = append(store.calls, boundedFileCall{path: path, scope: scope, limit: limit}) + return store.read(ctx, path, scope, limit) +} + +func TestBoundedFileSymbolIndexFailsClosedOnSaturationAndPushesScope(t *testing.T) { + const path = "repo/dense_test.go" + probe := &boundedFileStoreProbe{Store: graph.New()} + probe.read = func(_ context.Context, gotPath string, scope graph.LocalizationNodeScope, limit int) (graph.BoundedNodeProjection, error) { + if gotPath != path || limit != localizationFileNodeLimit { + t.Fatalf("bounded call = (%q, %d), want (%q, %d)", gotPath, limit, path, localizationFileNodeLimit) + } + if scope.WorkspaceID != "workspace" || scope.ProjectID != "project" || !scope.RepoAllow["repo"] { + t.Fatalf("scope was not pushed before cap: %#v", scope) + } + if scope.ExcludeTests { + t.Fatal("file localization excluded test declarations") + } + for _, kind := range localizationFileIndexKinds { + if !scope.Kinds[kind] { + t.Fatalf("kind %q missing from storage projection: %#v", kind, scope.Kinds) + } + } + owner := &graph.Node{ID: path + "::owner", Name: "owner", Kind: graph.KindFunction, FilePath: path, StartLine: 1, EndLine: 100} + return graph.BoundedNodeProjection{Nodes: []*graph.Node{owner}, Total: limit + 1, Truncated: true}, nil + } + server := &Server{graph: probe} + indexes := server.buildFileSymbolIndexForOrderedPathsScopedContext( + context.Background(), []string{path}, + query.QueryOptions{WorkspaceID: "workspace", ProjectID: "project", RepoAllow: map[string]bool{"repo": true}}, + ) + if index := indexes[path]; index == nil || !index.saturated || index.smallestEnclosing(10) != nil { + t.Fatalf("saturated file admitted an owner: %#v", index) + } +} + +func TestBoundedFileSymbolIndexStopsAtRequestBudget(t *testing.T) { + probe := &boundedFileStoreProbe{Store: graph.New()} + probe.read = func(_ context.Context, path string, _ graph.LocalizationNodeScope, limit int) (graph.BoundedNodeProjection, error) { + nodes := make([]*graph.Node, limit) + for index := range nodes { + nodes[index] = &graph.Node{ + ID: fmt.Sprintf("%s::owner-%04d", path, index), Name: "owner", + Kind: graph.KindFunction, FilePath: path, StartLine: index + 1, EndLine: index + 1, + } + } + return graph.BoundedNodeProjection{Nodes: nodes, Total: len(nodes)}, nil + } + server := &Server{graph: probe} + paths := []string{"repo/a.go", "repo/b.go", "repo/c.go", "repo/d.go", "repo/e.go"} + indexes := server.buildFileSymbolIndexForOrderedPathsScopedContext(context.Background(), paths, query.QueryOptions{}) + if len(probe.calls) != localizationFileRequestLimit/localizationFileNodeLimit { + t.Fatalf("bounded calls = %d, want request budget to stop after %d", len(probe.calls), localizationFileRequestLimit/localizationFileNodeLimit) + } + if index := indexes[paths[4]]; index == nil || !index.saturated { + t.Fatalf("budget-exhausted path was not failed closed: %#v", index) + } +} + +type unboundedFileStore struct{ graph.Store } + +func TestBoundedFileSymbolIndexFailsClosedWithoutCapabilityAndOnCancellation(t *testing.T) { + path := "repo/handler.go" + missing := (&Server{graph: &unboundedFileStore{Store: graph.New()}}).buildFileSymbolIndexForPaths( + map[string]struct{}{path: {}}, + ) + if index := missing[path]; index == nil || !index.saturated { + t.Fatalf("missing capability did not fail closed: %#v", index) + } + + probe := &boundedFileStoreProbe{Store: graph.New()} + probe.read = func(context.Context, string, graph.LocalizationNodeScope, int) (graph.BoundedNodeProjection, error) { + t.Fatal("cancelled request reached storage") + return graph.BoundedNodeProjection{}, nil + } + ctx, cancel := context.WithCancel(context.Background()) + cancel() + cancelled := (&Server{graph: probe}).buildFileSymbolIndexForOrderedPathsScopedContext(ctx, []string{path}, query.QueryOptions{}) + if index := cancelled[path]; index == nil || !index.saturated || len(probe.calls) != 0 { + t.Fatalf("cancelled lookup = %#v, calls=%d", index, len(probe.calls)) + } +} + +func TestBoundedFileSymbolIndexHonorsOverlayReplacementAndTombstone(t *testing.T) { + const path = "repo/handler.go" + base := graph.New() + base.AddNode(&graph.Node{ID: path + "::old", Name: "old", Kind: graph.KindFunction, FilePath: path, StartLine: 1, EndLine: 20}) + server := &Server{graph: base} + + replacement := graph.NewOverlayLayer() + replacement.MarkFile(path, false) + replacement.AddNode(path, &graph.Node{ID: path + "::new", Name: "new", Kind: graph.KindFunction, FilePath: path, StartLine: 1, EndLine: 20}) + replacementCtx := WithOverlayView(context.Background(), graph.NewOverlaidView(base, replacement)) + indexes := server.buildFileSymbolIndexForPathsScopedContext(replacementCtx, map[string]struct{}{path: {}}, query.QueryOptions{}) + if owner := indexes[path].smallestEnclosing(10); owner == nil || owner.Name != "new" { + t.Fatalf("overlay replacement owner = %#v, want new", owner) + } + + tombstone := graph.NewOverlayLayer() + tombstone.MarkFile(path, true) + tombstoneCtx := WithOverlayView(context.Background(), graph.NewOverlaidView(base, tombstone)) + indexes = server.buildFileSymbolIndexForPathsScopedContext(tombstoneCtx, map[string]struct{}{path: {}}, query.QueryOptions{}) + if index := indexes[path]; index != nil { + t.Fatalf("overlay tombstone leaked a file index: %#v", index) + } +} + +func TestSaturatedFileSymbolIndexCannotPromoteRange(t *testing.T) { + index := &fileSymbolIndex{ + saturated: true, + syms: []*graph.Node{{ID: "repo/a.go::wrong", Kind: graph.KindFunction, FilePath: "repo/a.go", StartLine: 1, EndLine: 20}}, + } + if got := exploreSourceRangeDefinitions(index, 10, 10); len(got) != 0 { + t.Fatalf("saturated range promoted omitted candidate: %#v", got) + } +} diff --git a/internal/mcp/localization_file_outline.go b/internal/mcp/localization_file_outline.go index 94ae1132..46515a1c 100644 --- a/internal/mcp/localization_file_outline.go +++ b/internal/mcp/localization_file_outline.go @@ -23,8 +23,13 @@ import ( // carries outlines; a terminal page's caller is answering, not choosing. const ( - // localizationOutlineRowCap bounds the rows the leading file's outline may - // serialize. + // localizationOutlineCompleteRows is the internal no-cap sentinel used only + // while re-projecting an already bounded outline during relief tests and + // compatibility paths. Page construction never selects it. + localizationOutlineCompleteRows = -1 + // localizationOutlineRowCap bounds the leading file's declaration index. + // Forty task-prioritized head/tail rows preserve useful navigation without + // letting one large file consume an otherwise healthy localization page. localizationOutlineRowCap = 40 // localizationOutlineHeadRows splits an elided outline between the file's // opening declarations and its closing ones. @@ -32,27 +37,37 @@ const ( // localizationOutlineFloorRows is the smallest index still worth its bytes. // Budget pressure shrinks an outline to this floor before it may drop one. localizationOutlineFloorRows = 8 - // localizationOutlineSecondFileRowCap is the depth the file ranked directly - // after the leading one gets. Every file after that starts at the floor. - localizationOutlineSecondFileRowCap = 12 - // localizationOutlineFileCap bounds how many of the page's distinct files - // are indexed at all, leading file included. - localizationOutlineFileCap = 6 + // Later page files need enough declarations to expose siblings without + // spending the request-wide projection budget reserved for the leading pair. + localizationOutlineTrailingFileFetchLimit = 128 + // localizationOutlineFileCap covers the first ten distinct ranked files: + // two complete bounded indexes plus eight shallow sibling indexes. + localizationOutlineFileCap = 10 + // The page's first two files keep their complete retained indexes until every + // lower-ranked file has yielded its expendable depth and breadth. + localizationOutlineProtectedFileCount = 2 ) -// localizationOutlineFileRowCap is the depth ladder over a page's files: the -// leading file keeps the whole index, the next one a third of it, and the rest -// start where shrinking would stop anyway. +// localizationOutlineFileRowCap is the initial wire depth over a page's files. +// The first two files expose every declaration retained by their bounded read; +// later files start at the relief floor. func localizationOutlineFileRowCap(rank int) int { - switch rank { - case 0: - return localizationOutlineRowCap - case 1: - return localizationOutlineSecondFileRowCap + if rank >= 0 && rank < localizationOutlineProtectedFileCount { + return localizationOutlineCompleteRows } return localizationOutlineFloorRows } +// localizationOutlineFileFetchLimit bounds the lightweight declaration read +// before wire projection. The first two distinct files receive the full typed +// backend page and the remaining eight receive a shallow page. +func localizationOutlineFileFetchLimit(rank int) int { + if rank >= 0 && rank < localizationOutlineProtectedFileCount { + return localizationFileNodeLimit + } + return localizationOutlineTrailingFileFetchLimit +} + // localizationPageOutline is the page's declaration index: the leading file's // outline and the outlines of the further page files, deepest first. type localizationPageOutline struct { @@ -66,23 +81,26 @@ type localizationOutlineRow struct { // Kind is a one-rune hint. The row is a locator, so a coarse hint costs a // byte and the name still decides. Kind string `json:"kind,omitempty"` + key string } // localizationFileOutline is the bounded declaration index of a page file. -// Declared counts what the file declares; Elided counts the rows the cap -// dropped from the middle. +// Declared is exact unless Truncated is true, in which case it is a lower bound. +// Elided follows the same rule and counts declarations not represented in Rows. // -// The unexported fields are the whole file in line order plus the task-term -// priority over it, retained so the same outline can be re-elided at a smaller -// cap without re-reading the graph. They are never serialized. +// The unexported fields are the retained declarations in line order plus the +// task-term priority over them, kept so the same outline can be re-elided at a +// smaller cap without re-reading the graph. They are never serialized. type localizationFileOutline struct { - File string `json:"file"` - Declared int `json:"declared"` - Elided int `json:"elided,omitempty"` - Rows []localizationOutlineRow `json:"rows"` + File string `json:"file"` + Declared int `json:"declared"` + Elided int `json:"elided,omitempty"` + Truncated bool `json:"truncated,omitempty"` + Rows []localizationOutlineRow `json:"rows"` all []localizationOutlineRow priority []int + rank int } // localizationPageAcceptsOutline admits the outline on the states whose caller @@ -103,9 +121,36 @@ func localizationPageOutlineProvider( pool []*rerank.Candidate, targets []exploreTarget, terms map[string]struct{}, - enumerate func(string) []*graph.Node, + enumerate any, ) func() *localizationPageOutline { - if enumerate == nil { + return localizationPageOutlineProviderWithCaps( + pool, targets, terms, enumerate, localizationOutlineFileRowCap, + ) +} + +// boundedLocalizationPageOutlineProvider applies the shared initial row policy +// to structured localize responses. Structured and task pages both start with +// the complete retained leading pair and relieve their own copies under budget. +func boundedLocalizationPageOutlineProvider( + pool []*rerank.Candidate, + targets []exploreTarget, + terms map[string]struct{}, + enumerate any, +) func() *localizationPageOutline { + return localizationPageOutlineProviderWithCaps( + pool, targets, terms, enumerate, localizationOutlineFileRowCap, + ) +} + +func localizationPageOutlineProviderWithCaps( + pool []*rerank.Candidate, + targets []exploreTarget, + terms map[string]struct{}, + enumerate any, + rowCap func(int) int, +) func() *localizationPageOutline { + enumerateDeclarations := localizationOutlineDeclarationEnumerator(enumerate) + if enumerateDeclarations == nil { return nil } var ( @@ -118,13 +163,19 @@ func localizationPageOutlineProvider( } built = true index := func(file string, rank int) *localizationFileOutline { - nodes := enumerate(file) - if len(nodes) == 0 { + declarations := enumerateDeclarations(file, rank) + if len(declarations.Nodes) == 0 && !declarations.Truncated { // Nothing to enumerate leaves the nodes this page already // fetched, which is the whole of what it knows about that file. - nodes = localizationOutlineFetchedNodes(pool, targets) + declarations = localizationFileDeclarations{ + Nodes: localizationOutlineFetchedNodes(pool, targets), + } } - return newLocalizationFileOutlineForTerms(file, nodes, terms, localizationOutlineFileRowCap(rank)) + outline := newLocalizationFileOutlineForDeclarations(file, declarations, terms, rowCap(rank)) + if outline != nil { + outline.rank = rank + } + return outline } // A page whose ranking never settled on one file has no leading slice to // give — but its rows still name files, and every one of them declares @@ -142,13 +193,18 @@ func localizationPageOutlineProvider( leading = "" } } - ranked := localizationOutlinePageRowCounts(targets) - for _, file := range localizationOutlineFollowingFiles(targets, leading, following) { - other := index(file, len(page.Others)+1) - if other == nil || other.Declared <= ranked[file] { - // A file whose every declaration is already a row on this page - // is a file the caller can already see. Indexing it costs bytes - // the deeper files would rather have. + followingFiles := localizationOutlineFollowingFiles(targets, leading, following) + for indexInPage, file := range followingFiles { + rank := indexInPage + if leading != "" { + rank++ + } + other := index(file, rank) + if other == nil || rank >= localizationOutlineProtectedFileCount && + !localizationOutlineAddsUnrankedDeclaration(other, targets) { + // Lower-ranked files whose declarations are all already visible + // add no navigation value. The top two remain explicit even when + // counts happen to match: identity, not cardinality, is the proof. continue } page.Others = append(page.Others, other) @@ -160,6 +216,26 @@ func localizationPageOutlineProvider( } } +// localizationOutlineDeclarationEnumerator keeps direct outline helpers that +// enumerate raw nodes source-compatible while allowing bounded caches to carry +// an exact declaration count alongside their retained page. +func localizationOutlineDeclarationEnumerator(enumerate any) func(string, int) localizationFileDeclarations { + switch enumerate := enumerate.(type) { + case func(string, int) localizationFileDeclarations: + return enumerate + case func(string) localizationFileDeclarations: + return func(file string, _ int) localizationFileDeclarations { + return enumerate(file) + } + case func(string) []*graph.Node: + return func(file string, _ int) localizationFileDeclarations { + return localizationFileDeclarations{Nodes: enumerate(file)} + } + default: + return nil + } +} + // localizationOutlineFollowingFiles names the page's other distinct files in // page order. The rows are already ranked, so their file order is the page's own // judgement of which neighbourhood the caller should look at next. @@ -207,19 +283,36 @@ func localizationOutlineLeadingFile(pool []*rerank.Candidate, targets []exploreT return exploreLeadingRankedFile(ranked) } -// localizationOutlinePageRowCounts counts how many of a file's declarations the -// page already names as ranked rows. -func localizationOutlinePageRowCounts(targets []exploreTarget) map[string]int { - counts := make(map[string]int, len(targets)) +// localizationOutlineAddsUnrankedDeclaration checks actual declaration +// identities. Counts cannot prove coverage when ranked rows repeat one identity +// or omit a different sibling. +func localizationOutlineAddsUnrankedDeclaration(outline *localizationFileOutline, targets []exploreTarget) bool { + if outline == nil { + return false + } + if outline.Truncated { + return true + } + visible := make(map[string]struct{}, len(targets)) for _, target := range targets { - if target.node == nil { + if target.node == nil || nodeDisplayPath(target.node) != outline.File { continue } - if file := nodeDisplayPath(target.node); file != "" { - counts[file]++ + key := target.node.ID + if key == "" { + key = strings.TrimSpace(target.node.QualName) + if key == "" { + key = strings.TrimSpace(target.node.Name) + } + } + visible[key] = struct{}{} + } + for _, row := range outline.all { + if _, covered := visible[row.key]; !covered { + return true } } - return counts + return false } func localizationOutlineFetchedNodes(pool []*rerank.Candidate, targets []exploreTarget) []*graph.Node { @@ -252,14 +345,23 @@ func newLocalizationFileOutlineForTerms( nodes []*graph.Node, terms map[string]struct{}, rowCap int, +) *localizationFileOutline { + return newLocalizationFileOutlineForDeclarations(file, localizationFileDeclarations{Nodes: nodes}, terms, rowCap) +} + +func newLocalizationFileOutlineForDeclarations( + file string, + declarations localizationFileDeclarations, + terms map[string]struct{}, + rowCap int, ) *localizationFileOutline { file = strings.TrimSpace(file) if file == "" { return nil } - rows := make([]localizationOutlineRow, 0, len(nodes)) - seen := make(map[string]struct{}, len(nodes)) - for _, node := range nodes { + rows := make([]localizationOutlineRow, 0, len(declarations.Nodes)) + seen := make(map[string]struct{}, len(declarations.Nodes)) + for _, node := range declarations.Nodes { if node == nil || nodeDisplayPath(node) != file || isNonDefinitionNode(node.Kind) { continue } @@ -282,6 +384,7 @@ func newLocalizationFileOutlineForTerms( Name: compactLocalizationField(name, localizationMaxNameRunes), Line: node.StartLine, Kind: localizationOutlineKindLetter(node.Kind), + key: key, }) } if len(rows) == 0 { @@ -293,11 +396,16 @@ func newLocalizationFileOutlineForTerms( } return rows[first].Name < rows[second].Name }) + declared := max(len(rows), declarations.Declared) + if declarations.Truncated { + declared = max(declared, len(rows)+1) + } outline := &localizationFileOutline{ - File: file, - Declared: len(rows), - all: rows, - priority: localizationOutlinePriority(rows, terms), + File: file, + Declared: declared, + Truncated: declarations.Truncated, + all: rows, + priority: localizationOutlinePriority(rows, terms), } outline.elide(rowCap) return outline @@ -311,9 +419,10 @@ func localizationOutlinePriority(rows []localizationOutlineRow, terms map[string return nil } type match struct{ index, matched, longest int } + wanted := localizationOutlineTerms(terms) matches := make([]match, 0, len(rows)) for index, row := range rows { - matched, longest := localizationOutlineRowTermMatch(row.Name, terms) + matched, longest := localizationOutlineRowTermMatchExpanded(row.Name, wanted) if matched == 0 { continue } @@ -339,30 +448,70 @@ func localizationOutlinePriority(rows []localizationOutlineRow, terms map[string // name carries, over the same camel/snake tokenization and the same root form // the task's own terms were built with. func localizationOutlineRowTermMatch(name string, terms map[string]struct{}) (matched, longest int) { - if len(terms) == 0 { + return localizationOutlineRowTermMatchExpanded(name, localizationOutlineTerms(terms)) +} + +func localizationOutlineRowTermMatchExpanded(name string, wanted map[string]struct{}) (matched, longest int) { + if len(wanted) == 0 { return 0, 0 } - seen := make(map[string]struct{}, len(terms)) + seen := make(map[string]struct{}, len(wanted)) for _, raw := range rerank.Tokenize(name) { - token := exploreTerminalTermRoot(strings.ToLower(strings.TrimSpace(raw))) - if token == "" { - continue + forms := localizationOutlineTermForms(strings.ToLower(strings.TrimSpace(raw))) + matchedForm := "" + for _, form := range forms { + if _, ok := wanted[form]; ok { + matchedForm = form + break + } } - if _, ok := terms[token]; !ok { + if matchedForm == "" { continue } - if _, duplicate := seen[token]; duplicate { + if _, duplicate := seen[matchedForm]; duplicate { continue } - seen[token] = struct{}{} + seen[matchedForm] = struct{}{} matched++ - if len(token) > longest { - longest = len(token) + if len(matchedForm) > longest { + longest = len(matchedForm) } } return matched, longest } +func localizationOutlineTerms(terms map[string]struct{}) map[string]struct{} { + forms := make(map[string]struct{}, len(terms)*2) + for term := range terms { + for _, form := range localizationOutlineTermForms(term) { + forms[form] = struct{}{} + } + } + return forms +} + +// localizationOutlineTermForms is deliberately outline-local: conservative +// inflection alternatives improve declaration indexing without widening answer +// readiness or the global semantic ranker. The exact spelling always remains. +func localizationOutlineTermForms(term string) []string { + term = strings.ToLower(strings.TrimSpace(term)) + if term == "" { + return nil + } + forms := []string{term} + for _, suffix := range []string{"ing", "ed", "er", "s"} { + if !strings.HasSuffix(term, suffix) { + continue + } + stem := strings.TrimSuffix(term, suffix) + if len([]rune(stem)) >= 4 && stem != term { + forms = append(forms, stem) + } + break + } + return forms +} + // elide re-projects the retained rows at rowCap. The declarations the task names // are kept first; the file's head and tail then fill whatever the cap has left, // so an index the caller cannot navigate by eye still shows both of its ends. @@ -370,12 +519,17 @@ func (o *localizationFileOutline) elide(rowCap int) { if o == nil { return } + if rowCap == localizationOutlineCompleteRows { + o.Rows = o.all + o.Elided = max(o.Declared-len(o.Rows), 0) + return + } if rowCap < 0 { rowCap = 0 } if len(o.all) <= rowCap { o.Rows = o.all - o.Elided = 0 + o.Elided = max(o.Declared-len(o.Rows), 0) return } // The task's own matches lead, but never take the whole index: a task whose @@ -410,7 +564,7 @@ func (o *localizationFileOutline) elide(rowCap int) { } } o.Rows = rows - o.Elided = len(o.all) - len(rows) + o.Elided = max(o.Declared-len(rows), 0) } // clone detaches an outline from the provider's cache. One request can pack @@ -438,47 +592,85 @@ func (p *localizationPageOutline) clone() *localizationPageOutline { return copied } -// relieve gives back one increment of outline payload. An index that no longer -// fits shrinks toward its floor before it is dropped: a shorter index still -// names the file and still carries the declarations the task asked about, where -// no index at all sends the caller back to search for what it is already -// looking at. Depth goes before breadth — the deepest index gives back first, -// and only once every file is at its floor does a file go, shallowest first. +// dropUnprotectedFloorFile gives back one rank-two-or-later outline after all +// of its depth has already yielded. Envelope packing uses this before stripping +// independently useful evidence detail; the protected leading pair is untouched. +func (p *localizationPageOutline) dropUnprotectedFloorFile() bool { + if p == nil { + return false + } + for index := len(p.Others) - 1; index >= 0; index-- { + outline := p.Others[index] + if outline == nil || p.otherRank(index) < localizationOutlineProtectedFileCount || + len(outline.Rows) > localizationOutlineFloorRows { + continue + } + p.Others = append(p.Others[:index], p.Others[index+1:]...) + return true + } + return false +} + +// relieve gives back one increment of outline payload. Every rank-two-or-later +// file shrinks and then drops before either of the top two files changes. Under +// further pressure rank one yields before rank zero. This makes completeness a +// real preference while leaving the envelope's hard byte bound authoritative. func (p *localizationPageOutline) relieve() { if p == nil { return } - floor := p.leadingFloor() - deepest, rows := (*localizationFileOutline)(nil), floor - if p.Leading != nil && len(p.Leading.Rows) > rows { - deepest, rows = p.Leading, len(p.Leading.Rows) - } - for _, other := range p.Others { - if other != nil && len(other.Rows) >= rows && len(other.Rows) > localizationOutlineFloorRows { - deepest, rows, floor = other, len(other.Rows), localizationOutlineFloorRows + // Lower-ranked depth yields from the tail inward. + for index := len(p.Others) - 1; index >= 0; index-- { + outline := p.Others[index] + if outline == nil || p.otherRank(index) < localizationOutlineProtectedFileCount || + len(outline.Rows) <= localizationOutlineFloorRows { + continue } + outline.elide(localizationOutlineNextRowCap(len(outline.Rows))) + return } - if deepest != nil { - deepest.elide(max(localizationOutlineNextRowCap(rows), floor)) + // Once all lower-ranked files reached their floor, give back their breadth + // before touching the protected pair. + for index := len(p.Others) - 1; index >= 0; index-- { + outline := p.Others[index] + if outline == nil || p.otherRank(index) < localizationOutlineProtectedFileCount { + continue + } + p.Others = append(p.Others[:index], p.Others[index+1:]...) return } - if last := len(p.Others); last > 0 { - p.Others = p.Others[:last-1] + // The second-ranked file is the next expendable block. + for index := len(p.Others) - 1; index >= 0; index-- { + outline := p.Others[index] + if outline == nil || p.otherRank(index) != 1 { + continue + } + if len(outline.Rows) > localizationOutlineFloorRows { + outline.elide(localizationOutlineNextRowCap(len(outline.Rows))) + } else { + p.Others = append(p.Others[:index], p.Others[index+1:]...) + } return } - p.Leading = nil + if p.Leading != nil { + if len(p.Leading.Rows) > localizationOutlineFloorRows { + p.Leading.elide(localizationOutlineNextRowCap(len(p.Leading.Rows))) + } else { + p.Leading = nil + } + } } -// leadingFloor is where the leading file's index stops shrinking. While the page -// still indexes further files, they are what a tight page gives back — measured, -// spending the leading file's depth on a neighbour's first twelve rows loses the -// answer more often than it finds one. Once no further file is left, the leading -// index may shrink to the same floor as any other. -func (p *localizationPageOutline) leadingFloor() int { - if len(p.Others) > 0 { - return localizationOutlineSecondFileRowCap +// otherRank preserves compatibility with directly-constructed outline blocks: +// an unset rank on an Others entry means its positional page rank. +func (p *localizationPageOutline) otherRank(index int) int { + if p == nil || index < 0 || index >= len(p.Others) || p.Others[index] == nil { + return index + 1 } - return localizationOutlineFloorRows + if rank := p.Others[index].rank; rank > 0 { + return rank + } + return index + 1 } // empty reports a block with nothing left to give. @@ -492,7 +684,7 @@ func (p *localizationPageOutline) atFloor() bool { if p == nil { return true } - if p.Leading != nil && len(p.Leading.Rows) > p.leadingFloor() { + if p.Leading != nil && len(p.Leading.Rows) > localizationOutlineFloorRows { return false } for _, other := range p.Others { diff --git a/internal/mcp/localization_file_outline_test.go b/internal/mcp/localization_file_outline_test.go index fa5a47ef..527725cd 100644 --- a/internal/mcp/localization_file_outline_test.go +++ b/internal/mcp/localization_file_outline_test.go @@ -300,8 +300,8 @@ func TestLocalizationBudgetLeavesRoomForOutlineBesideEvidence(t *testing.T) { t.Fatalf("envelope = %d bytes, budget = %d", bytes, localizationDefaultBudgetTokens*localizationEnvelopeBytesPerToken) } - if envelope.Outline == nil || len(envelope.Outline.Rows) != localizationOutlineRowCap { - t.Fatalf("outline = %#v, want %d rows beside the evidence", envelope.Outline, localizationOutlineRowCap) + if envelope.Outline == nil || len(envelope.Outline.Rows) != len(declared) || envelope.Outline.Elided != 0 { + t.Fatalf("outline = %#v, want all %d rows beside the evidence", envelope.Outline, len(declared)) } if len(envelope.Evidence) != len(targets) { t.Fatalf("evidence rows = %d, want all %d ranked rows beside the outline", @@ -774,17 +774,18 @@ func TestFurtherFilesGiveWayBeforeTheLeadingFilesDepth(t *testing.T) { }) } leading.Declared = len(leading.all) - leading.elide(localizationOutlineRowCap) + leading.elide(localizationOutlineCompleteRows) page := &localizationPageOutline{Leading: leading} - for _, file := range []string{"repo/second.go", "repo/third.go"} { - other := &localizationFileOutline{File: file} - for index := 0; index < 20; index++ { + for index, file := range []string{"repo/second.go", "repo/third.go"} { + rank := index + 1 + other := &localizationFileOutline{File: file, rank: rank} + for row := 0; row < 20; row++ { other.all = append(other.all, localizationOutlineRow{ - Name: fmt.Sprintf("Other%02d", index), Line: index + 1, Kind: "f", + Name: fmt.Sprintf("Other%02d", row), Line: row + 1, Kind: "f", }) } other.Declared = len(other.all) - other.elide(localizationOutlineSecondFileRowCap) + other.elide(localizationOutlineFileRowCap(rank)) page.Others = append(page.Others, other) } @@ -809,9 +810,9 @@ func TestFurtherFilesGiveWayBeforeTheLeadingFilesDepth(t *testing.T) { t.Fatalf("further files dropped %d times, want 2", len(depthWhenAFileWent)) } for _, depth := range depthWhenAFileWent { - if depth < localizationOutlineSecondFileRowCap { - t.Fatalf("a further file survived the leading file's depth falling to %d, floor %d", - depth, localizationOutlineSecondFileRowCap) + if depth != len(leading.all) { + t.Fatalf("a further file outlived the leading file's complete %d-row depth: %d", + len(leading.all), depth) } } } @@ -959,6 +960,256 @@ func TestOutlineElisionRanksStrongerTaskTermMatchesFirst(t *testing.T) { } } +func TestOutlineTopTwoFilesKeepAllRetainedDeclarations(t *testing.T) { + leading := outlineDeclaredFile(localizationOutlineRowCap + 7) + secondFile := "repo/second.go" + secondCount := localizationOutlineFloorRows + 13 + second := make([]*graph.Node, 0, secondCount) + for index := 0; index < secondCount; index++ { + second = append(second, outlineFileDeclaration(secondFile, fmt.Sprintf("Second%02d", index), index+1)) + } + targets := []exploreTarget{{node: leading[0]}, {node: second[0]}} + page := boundedLocalizationPageOutlineProvider( + outlinePool(leading[0]), targets, nil, + func(file string) []*graph.Node { + if file == outlineLeadingFile { + return leading + } + return second + }, + )() + + if page == nil || page.Leading == nil || len(page.Others) != 1 { + t.Fatalf("page = %#v, want complete retained top-two outlines", page) + } + if got := len(page.Leading.Rows); got != len(leading) { + t.Fatalf("leading rows = %d, want all %d retained declarations", got, len(leading)) + } + if page.Leading.Declared != len(leading) || page.Leading.Elided != 0 || page.Leading.Truncated { + t.Fatalf("leading counts = %#v, want exact complete retained outline", page.Leading) + } + if got := len(page.Others[0].Rows); got != len(second) { + t.Fatalf("second rows = %d, want all %d retained declarations", got, len(second)) + } + if page.Others[0].Declared != len(second) || page.Others[0].Elided != 0 || page.Others[0].Truncated { + t.Fatalf("second counts = %#v, want exact complete retained outline", page.Others[0]) + } +} + +func TestRankAwareOutlinePagesReportTruncationLowerBounds(t *testing.T) { + files := []string{outlineLeadingFile, "repo/second.go", "repo/third.go"} + byFile := make(map[string]localizationFileDeclarations, len(files)) + for rank, file := range files { + limit := localizationOutlineFileFetchLimit(rank) + nodes := make([]*graph.Node, 0, limit) + for index := 0; index < limit; index++ { + nodes = append(nodes, outlineFileDeclaration(file, fmt.Sprintf("Declaration%04d", index), index+1)) + } + byFile[file] = localizationFileDeclarations{ + Nodes: nodes, Declared: limit + 1, DeclaredKnown: false, Truncated: true, + } + } + targets := []exploreTarget{ + {node: byFile[files[0]].Nodes[0]}, + {node: byFile[files[1]].Nodes[0]}, + {node: byFile[files[2]].Nodes[0]}, + } + page := boundedLocalizationPageOutlineProvider( + outlinePool(targets[0].node), targets, nil, + func(file string, _ int) localizationFileDeclarations { return byFile[file] }, + )() + if page == nil || page.Leading == nil || len(page.Others) != 2 { + t.Fatalf("page = %#v, want three ranked outlines", page) + } + for _, outline := range []*localizationFileOutline{page.Leading, page.Others[0]} { + if !outline.Truncated || outline.Declared != localizationFileNodeLimit+1 || + outline.Elided != 1 || len(outline.Rows) != localizationFileNodeLimit { + t.Fatalf("protected truncated outline = %#v, want 1024 rows and lower bound 1025", outline) + } + } + trailing := page.Others[1] + wantTrailingElided := localizationOutlineTrailingFileFetchLimit + 1 - localizationOutlineFloorRows + if !trailing.Truncated || trailing.Declared != localizationOutlineTrailingFileFetchLimit+1 || + trailing.Elided != wantTrailingElided || len(trailing.Rows) != localizationOutlineFloorRows { + t.Fatalf("trailing truncated outline = %#v, want 8 rows and lower bound 129", trailing) + } +} + +func TestOutlineProviderEnumeratesAtMostTenDistinctFiles(t *testing.T) { + if localizationOutlineFileCap != 10 { + t.Fatalf("outline file cap = %d, want 10", localizationOutlineFileCap) + } + targets := outlineBreadthTargets(localizationOutlineFileCap + 1) + enumerated := 0 + page := localizationPageOutlineProvider(nil, targets, nil, func(file string) []*graph.Node { + enumerated++ + return []*graph.Node{ + outlineFileDeclaration(file, "Ranked", 1), + outlineFileDeclaration(file, "Unranked", 2), + } + })() + if page == nil { + t.Fatal("page outline is absent") + } + if enumerated != localizationOutlineFileCap { + t.Fatalf("enumerated %d files, want cap %d", enumerated, localizationOutlineFileCap) + } + if got := len(page.Others); got > localizationOutlineFileCap { + t.Fatalf("serialized %d outlines, cap %d", got, localizationOutlineFileCap) + } +} + +func TestOutlineSkipUsesDeclarationIdentity(t *testing.T) { + leading := outlineDeclaration("Lead", 1) + secondFile := "repo/second.go" + thirdFile := "repo/third.go" + fourthFile := "repo/fourth.go" + second := outlineFileDeclaration(secondFile, "Second", 1) + thirdRanked := outlineFileDeclaration(thirdFile, "ThirdRanked", 1) + thirdMissing := outlineFileDeclaration(thirdFile, "ThirdMissing", 2) + fourthA := outlineFileDeclaration(fourthFile, "FourthA", 1) + fourthB := outlineFileDeclaration(fourthFile, "FourthB", 2) + targets := []exploreTarget{ + {node: leading}, {node: second}, + {node: thirdRanked}, {node: thirdRanked}, + {node: fourthA}, {node: fourthB}, + } + byFile := map[string][]*graph.Node{ + outlineLeadingFile: {leading}, + secondFile: {second}, + thirdFile: {thirdRanked, thirdMissing}, + fourthFile: {fourthA, fourthB}, + } + page := localizationPageOutlineProvider( + outlinePool(leading), targets, nil, func(file string) []*graph.Node { return byFile[file] }, + )() + if page == nil || len(page.Others) != 2 { + t.Fatalf("others = %#v, want protected second plus identity-missing third", page) + } + if page.Others[0].File != secondFile || page.Others[1].File != thirdFile { + t.Fatalf("others = %#v, fully covered fourth file should be skipped", page.Others) + } +} + +func TestOutlineReliefProtectsTopTwoFiles(t *testing.T) { + outline := func(file string, rank int) *localizationFileOutline { + rows := make([]localizationOutlineRow, 30) + for index := range rows { + rows[index] = localizationOutlineRow{Name: fmt.Sprintf("Row%02d", index), Line: index + 1} + } + result := &localizationFileOutline{File: file, Declared: len(rows), all: rows, rank: rank} + result.elide(localizationOutlineCompleteRows) + return result + } + page := &localizationPageOutline{ + Leading: outline("repo/lead.go", 0), + Others: []*localizationFileOutline{ + outline("repo/second.go", 1), outline("repo/third.go", 2), + }, + } + leadingRows, secondRows := len(page.Leading.Rows), len(page.Others[0].Rows) + page.relieve() + if len(page.Leading.Rows) != leadingRows || len(page.Others[0].Rows) != secondRows || len(page.Others[1].Rows) >= 30 { + t.Fatalf("first relief changed protected outlines: %#v", page) + } + for len(page.Others) > 1 { + page.relieve() + } + if len(page.Leading.Rows) != leadingRows || len(page.Others[0].Rows) != secondRows { + t.Fatalf("lower-ranked file outlived protected depth: %#v", page) + } + page.relieve() + if len(page.Leading.Rows) != leadingRows || len(page.Others[0].Rows) >= secondRows { + t.Fatalf("rank one did not yield before rank zero: %#v", page) + } +} + +func TestOutlineDropsUnprotectedFloorBeforeEvidenceRelief(t *testing.T) { + outline := func(file string, rank int) *localizationFileOutline { + rows := make([]localizationOutlineRow, localizationOutlineFloorRows) + for index := range rows { + rows[index] = localizationOutlineRow{Name: fmt.Sprintf("Row%02d", index), Line: index + 1} + } + return &localizationFileOutline{File: file, Declared: len(rows), Rows: rows, all: rows, rank: rank} + } + page := &localizationPageOutline{ + Leading: outline("repo/lead.go", 0), + Others: []*localizationFileOutline{ + outline("repo/second.go", 1), outline("repo/third.go", 2), + }, + } + if !page.dropUnprotectedFloorFile() { + t.Fatal("rank-two floor outline was not dropped") + } + if len(page.Others) != 1 || page.Others[0].File != "repo/second.go" || page.Leading == nil { + t.Fatalf("protected outline changed: %#v", page) + } + if page.dropUnprotectedFloorFile() { + t.Fatal("protected leading pair was dropped") + } +} + +func TestOutlineTermFormsMatchConservativeInflections(t *testing.T) { + for _, test := range []struct { + name string + query string + row string + }{ + {name: "ing", query: "buffering", row: "buffer"}, + {name: "ed", query: "matched", row: "match"}, + {name: "er", query: "reader", row: "read"}, + {name: "exact retained", query: "render", row: "render"}, + } { + t.Run(test.name, func(t *testing.T) { + matched, _ := localizationOutlineRowTermMatch(test.row, map[string]struct{}{test.query: {}}) + if matched != 1 { + t.Fatalf("%q did not match %q", test.query, test.row) + } + }) + } + if matched, _ := localizationOutlineRowTermMatch("go", map[string]struct{}{"going": {}}); matched != 0 { + t.Fatal("a stem shorter than four characters matched") + } +} + +func TestTruncatedOutlineCarriesAnExplicitLowerBound(t *testing.T) { + declarations := outlineDeclaredFile(2) + outline := newLocalizationFileOutlineForDeclarations( + outlineLeadingFile, + localizationFileDeclarations{ + Nodes: declarations, Declared: len(declarations) + 1, + DeclaredKnown: false, Truncated: true, + }, + nil, localizationOutlineCompleteRows, + ) + if outline == nil || !outline.Truncated || outline.Declared != 3 || outline.Elided != 1 { + t.Fatalf("truncated outline = %#v, want lower bound 3 with at least one elided", outline) + } + if !localizationOutlineAddsUnrankedDeclaration( + outline, []exploreTarget{{node: declarations[0]}, {node: declarations[1]}}, + ) { + t.Fatal("a saturated outline was treated as completely covered by ranked rows") + } + body, err := json.Marshal(outline) + if err != nil { + t.Fatal(err) + } + if !strings.Contains(string(body), `"truncated":true`) { + t.Fatalf("truncation lower bound is absent from JSON: %s", body) + } +} + +func TestOutlinePrivateIdentityDoesNotSerialize(t *testing.T) { + outline := newLocalizationFileOutline(outlineLeadingFile, outlineDeclaredFile(2)) + body, err := json.Marshal(outline) + if err != nil { + t.Fatal(err) + } + if strings.Contains(string(body), `"key"`) || strings.Contains(string(body), `"rank"`) { + t.Fatalf("private outline fields leaked into JSON: %s", body) + } +} + func TestScatteredRankingIndexesThePageFilesWithoutALeadingOne(t *testing.T) { var pool []*rerank.Candidate var targets []exploreTarget @@ -1007,9 +1258,12 @@ func TestScatteredRankingIndexesThePageFilesWithoutALeadingOne(t *testing.T) { t.Fatalf("scattered ranking paid %d graph enumerations, cap %d", enumerated, localizationOutlineFileCap) } - for _, other := range page.Others { - if len(other.Rows) > localizationOutlineSecondFileRowCap { - t.Fatalf("a file on a page with no leading one took the deepest slice: %#v", other) + for rank, other := range page.Others { + if got, want := len(other.Rows), len(declared[other.File]); got != want { + t.Fatalf("rank %d retained %d rows, want all %d available declarations: %#v", rank, got, want, other) + } + if rank >= localizationOutlineProtectedFileCount && len(other.Rows) > localizationOutlineFloorRows { + t.Fatalf("unprotected scattered rank %d exceeded the relief floor: %#v", rank, other) } } } diff --git a/internal/mcp/localization_projection.go b/internal/mcp/localization_projection.go new file mode 100644 index 00000000..8c33b43c --- /dev/null +++ b/internal/mcp/localization_projection.go @@ -0,0 +1,121 @@ +package mcp + +import ( + "context" + + "github.com/zzet/gortex/internal/graph" + "github.com/zzet/gortex/internal/query" +) + +func (s *Server) localizationNodeScope( + ctx context.Context, + opts query.QueryOptions, + kinds ...graph.NodeKind, +) graph.LocalizationNodeScope { + return s.localizationNodeScopeWithTests(ctx, opts, true, kinds...) +} + +// localizationNodeScopeWithTests keeps request/session intersection identical +// across localization projections while letting file ownership include test +// declarations without transferring metadata from SQLite. +func (s *Server) localizationNodeScopeWithTests( + ctx context.Context, + opts query.QueryOptions, + excludeTests bool, + kinds ...graph.NodeKind, +) graph.LocalizationNodeScope { + if session, bound := s.sessionScopeOptions(ctx); bound { + switch { + case opts.WorkspaceID == "": + opts.WorkspaceID = session.WorkspaceID + case session.WorkspaceID != "" && opts.WorkspaceID != session.WorkspaceID: + opts.WorkspaceID = unresolvedWorkspacePrefix + "disjoint-localization-scope" + opts.RepoAllow = nil + } + switch { + case len(opts.RepoAllow) == 0: + opts.RepoAllow = session.RepoAllow + case len(session.RepoAllow) == 0: + // The session workspace boundary alone remains authoritative. + default: + if merged := intersectRepoSets(opts.RepoAllow, session.RepoAllow); len(merged) > 0 { + opts.RepoAllow = merged + } else { + opts.WorkspaceID = unresolvedWorkspacePrefix + "disjoint-localization-repos" + opts.RepoAllow = nil + } + } + } + scope := graph.LocalizationNodeScope{ + WorkspaceID: opts.WorkspaceID, + ProjectID: opts.ProjectID, + RepoAllow: opts.RepoAllow, + ExcludeTests: excludeTests, + } + if len(kinds) > 0 { + scope.Kinds = make(map[graph.NodeKind]bool, len(kinds)) + for _, kind := range kinds { + scope.Kinds[kind] = true + } + } + return scope +} + +// boundedLocalizationExactName deliberately has no Reader.FindNodesByName +// fallback. A backend without the typed capability returns no localization +// evidence instead of reintroducing an unbounded full-row read. +func boundedLocalizationExactName( + ctx context.Context, + reader graph.Reader, + name string, + scope graph.LocalizationNodeScope, + limit int, +) (graph.BoundedNodeProjection, bool) { + bounded, ok := reader.(graph.BoundedExactNameReader) + if !ok || ctx.Err() != nil { + return graph.BoundedNodeProjection{}, false + } + page, err := bounded.FindNodesByNameBounded(ctx, name, scope, limit) + if err != nil { + return graph.BoundedNodeProjection{}, false + } + return page, true +} + +// boundedLocalizationFileNodes performs one metadata-free file projection +// under the request-wide localization budget. Incomplete pages are not usable +// evidence: ranked-file recovery makes uniqueness and ownership decisions, so +// accepting a prefix could fabricate an unambiguous declaration. +func boundedLocalizationFileNodes( + ctx context.Context, + reader graph.Reader, + budget *localizationFileRequestBudget, + path string, + scope graph.LocalizationNodeScope, + limit int, +) (graph.BoundedNodeProjection, bool) { + bounded, ok := reader.(graph.BoundedFileNodeReader) + if !ok || path == "" || ctx.Err() != nil { + return graph.BoundedNodeProjection{}, false + } + if limit > localizationFileNodeLimit { + limit = localizationFileNodeLimit + } + reserved := budget.reserve(limit) + if reserved <= 0 { + return graph.BoundedNodeProjection{}, false + } + page, err := bounded.FindFileNodesBounded(ctx, path, scope, reserved) + if err != nil || ctx.Err() != nil { + return graph.BoundedNodeProjection{}, false + } + consumed := page.Total + if len(page.Nodes) > consumed { + consumed = len(page.Nodes) + } + budget.finish(reserved, consumed) + if page.Truncated { + return graph.BoundedNodeProjection{}, false + } + return page, true +} diff --git a/internal/mcp/localization_secondary_citation_test.go b/internal/mcp/localization_secondary_citation_test.go new file mode 100644 index 00000000..a49891a5 --- /dev/null +++ b/internal/mcp/localization_secondary_citation_test.go @@ -0,0 +1,119 @@ +package mcp + +import ( + "testing" + + "github.com/zzet/gortex/internal/graph" +) + +func localizationSecondaryCitationTestTarget(id, name string) exploreTarget { + return exploreTarget{node: &graph.Node{ + ID: id, Name: name, Kind: graph.KindFunction, + FilePath: "src/citations.go", StartLine: 1, + }} +} + +func TestLocalizationSecondaryTaskCitedTargetConsumesRepresentedMention(t *testing.T) { + format := localizationSecondaryCitationTestTarget("src/citations.go::format", "format") + format.node.QualName = "fmt::format" + formatOverload := localizationSecondaryCitationTestTarget("src/citations.go::format#2", "format") + formatOverload.node.QualName = "other::format" + macro := localizationSecondaryCitationTestTarget("src/citations.go::FMT_FORMAT_AS", "FMT_FORMAT_AS") + later := localizationSecondaryCitationTestTarget("src/citations.go::LATER_MACRO", "LATER_MACRO") + targets := []exploreTarget{formatOverload, later, macro, format} + task := "The call fmt::format(value) should use FMT_FORMAT_AS(Type) before LATER_MACRO(Type)." + + got, ok := localizationSecondaryTaskCitedTarget(task, targets, []exploreTarget{format}) + if !ok { + t.Fatal("expected an independent cited target") + } + if got.node == nil || got.node.ID != macro.node.ID { + t.Fatalf("secondary target = %#v, want %q", got.node, macro.node.ID) + } + if targets[0].node.ID != formatOverload.node.ID || targets[1].node.ID != later.node.ID { + t.Fatalf("helper mutated caller-owned order: %v", exploreLocalizationTargetSymbols(targets)) + } +} + +func TestLocalizationSecondaryTaskCitedTargetRejectsLooseProse(t *testing.T) { + target := localizationSecondaryCitationTestTarget("src/citations.go::target", "target") + if got, ok := localizationSecondaryTaskCitedTarget( + "Investigate the target behavior in src/citations.go.", + []exploreTarget{target}, nil, + ); ok { + t.Fatalf("loose prose selected %#v", got.node) + } + perform := localizationSecondaryCitationTestTarget("src/citations.go::performHydrate", "performHydrate") + if got, ok := localizationSecondaryTaskCitedTarget( + "Investigate performHydrateLater during startup.", + []exploreTarget{perform}, nil, + ); ok { + t.Fatalf("identifier suffix selected %#v", got.node) + } +} + +func TestLocalizationEvidenceTargetsFromDraftReservesOneIndependentCitation(t *testing.T) { + format := localizationSecondaryCitationTestTarget("src/citations.go::format", "format") + format.sourceRange = true + formatOverload := localizationSecondaryCitationTestTarget("src/citations.go::format#2", "format") + ordinaryA := localizationSecondaryCitationTestTarget("src/citations.go::ordinaryA", "ordinaryA") + ordinaryB := localizationSecondaryCitationTestTarget("src/citations.go::ordinaryB", "ordinaryB") + ordinaryC := localizationSecondaryCitationTestTarget("src/citations.go::ordinaryC", "ordinaryC") + ordinaryD := localizationSecondaryCitationTestTarget("src/citations.go::ordinaryD", "ordinaryD") + later := localizationSecondaryCitationTestTarget("src/citations.go::LATER_MACRO", "LATER_MACRO") + macro := localizationSecondaryCitationTestTarget("src/citations.go::FMT_FORMAT_AS", "FMT_FORMAT_AS") + targets := []exploreTarget{ + format, formatOverload, ordinaryA, ordinaryB, ordinaryC, ordinaryD, later, macro, + } + task := "The call fmt::format(value) should use FMT_FORMAT_AS(Type) before LATER_MACRO(Type)." + + ordered := localizationEvidenceTargetsFromDraft(task, "", targets, []exploreDraftEntry{}) + if len(ordered) != len(targets) { + t.Fatalf("ordered target count = %d, want %d", len(ordered), len(targets)) + } + got := exploreLocalizationTargetSymbols(ordered) + wantPrefix := []string{format.node.ID, macro.node.ID, formatOverload.node.ID} + for index, want := range wantPrefix { + if got[index] != want { + t.Fatalf("ordered IDs = %v, want prefix %v", got, wantPrefix) + } + } + laterIndex := -1 + for index, id := range got { + if id == later.node.ID { + laterIndex = index + break + } + } + if laterIndex <= 1 { + t.Fatalf("second independent citation was also reserved: %v", got) + } + + envelope := localizationExploreEnvelope{Evidence: make([]localizationEvidence, 0, len(ordered))} + for index, target := range ordered { + node := target.node + envelope.Evidence = append(envelope.Evidence, localizationEvidence{ + Rank: index + 1, ID: node.ID, Name: node.Name, QualName: node.QualName, + Kind: string(node.Kind), File: nodeDisplayPath(node), Line: node.StartLine, + }) + } + digest := newLocalizationEvidenceDigestForTask(task, envelope) + freezeLocalizationPrimaryCohort(task, &envelope, digest) + // LATER_MACRO is task-cited, so it rides one task-named extra seat past + // the ranked block; the ranked width itself stays at the primary limit. + if len(digest.primaryIDs) != localizationFinalResponsePrimaryLimit+1 { + t.Fatalf("PRIMARY count = %d, want %d: %v", len(digest.primaryIDs), localizationFinalResponsePrimaryLimit+1, digest.primaryIDs) + } + macroPrimary := false + for _, id := range digest.primaryIDs { + macroPrimary = macroPrimary || id == macro.node.ID + } + if !macroPrimary { + t.Fatalf("secondary cited target not frozen as PRIMARY: %v", digest.primaryIDs) + } + for _, row := range envelope.Evidence { + if row.ID == macro.node.ID && (row.Provenance != "" || row.taskCitedPrimaryEligible) { + t.Fatalf("secondary citation gained authority metadata: %#v", row) + } + } +} diff --git a/internal/mcp/localization_source_window.go b/internal/mcp/localization_source_window.go new file mode 100644 index 00000000..c9821bd4 --- /dev/null +++ b/internal/mcp/localization_source_window.go @@ -0,0 +1,289 @@ +package mcp + +import ( + "context" + "strings" + "unicode/utf8" + + "github.com/zzet/gortex/internal/review" +) + +const ( + localizationSourceWindowRadius = 100 + localizationSourceWindowMaxLines = 2*localizationSourceWindowRadius + 1 + localizationSourceWindowMaxBytes = 12 << 10 + localizationSourceWindowCandidateLimit = 64 +) + +// localizationSourceWindow is optional context around one exact source-literal +// match already discovered by localization. It is neither evidence nor proof: +// the ordinary evidence projection and digest remain authoritative. +type localizationSourceWindow struct { + Path string `json:"path"` + StartLine int `json:"start_line"` + EndLine int `json:"end_line"` + MatchLine int `json:"match_line"` + Content string `json:"content"` + Bytes int `json:"bytes"` + Lines int `json:"lines"` + SecretsRedacted bool `json:"secrets_redacted,omitempty"` + Truncated bool `json:"truncated,omitempty"` + AnchorSymbol string `json:"anchor_symbol"` +} + +// localizationSourceWindowHitCollector is request-local. Callers pass nil on +// ordinary explore requests, preserving their existing behavior and memory. +type localizationSourceWindowHitCollector struct { + hits []exploreSourceLiteralHit +} + +func (collector *localizationSourceWindowHitCollector) add(hits ...exploreSourceLiteralHit) { + if collector == nil { + return + } + for _, hit := range hits { + if len(collector.hits) >= localizationSourceWindowCandidateLimit { + return + } + if strings.TrimSpace(hit.nodeID) == "" || strings.TrimSpace(hit.matchPath) == "" || + hit.matchLine < 1 || strings.TrimSpace(hit.literal) == "" { + continue + } + duplicate := false + for _, existing := range collector.hits { + if existing.nodeID == hit.nodeID && existing.matchPath == hit.matchPath && + existing.matchLine == hit.matchLine && strings.EqualFold(existing.literal, hit.literal) { + duplicate = true + break + } + } + if !duplicate { + collector.hits = append(collector.hits, hit) + } + } +} + +// elect returns one exact hit whose source-literal candidate survived the final +// page target selection. The order is deterministic and favors settled, +// construction-aligned callsites before retrieval rank. +func (collector *localizationSourceWindowHitCollector) elect(targets []exploreTarget) (int, exploreSourceLiteralHit, bool) { + if collector == nil || len(collector.hits) == 0 || len(targets) == 0 { + return 0, exploreSourceLiteralHit{}, false + } + targetByID := make(map[string]int, len(targets)) + for index, target := range targets { + if target.node == nil || target.node.ID == "" { + continue + } + if _, exists := targetByID[target.node.ID]; !exists { + targetByID[target.node.ID] = index + } + } + bestIndex := -1 + var best exploreSourceLiteralHit + for _, hit := range collector.hits { + index, survives := targetByID[hit.nodeID] + if !survives { + continue + } + if !targets[index].sourceLiteral && !hit.callee { + continue + } + if bestIndex < 0 || localizationSourceWindowHitBetter(index, targets[index], hit, bestIndex, targets[bestIndex], best) { + bestIndex, best = index, hit + } + } + return bestIndex, best, bestIndex >= 0 +} + +func localizationSourceWindowHitBetter( + leftIndex int, + leftTarget exploreTarget, + left exploreSourceLiteralHit, + rightIndex int, + rightTarget exploreTarget, + right exploreSourceLiteralHit, +) bool { + if left.ambiguous != right.ambiguous { + return !left.ambiguous + } + if leftTarget.sourceLiteralAligned != rightTarget.sourceLiteralAligned { + return leftTarget.sourceLiteralAligned + } + if left.callee != right.callee { + return left.callee + } + if leftIndex != rightIndex { + return leftIndex < rightIndex + } + if left.anchor != right.anchor { + return left.anchor < right.anchor + } + if left.rank != right.rank { + return left.rank < right.rank + } + if left.matchPath != right.matchPath { + return left.matchPath < right.matchPath + } + if left.matchLine != right.matchLine { + return left.matchLine < right.matchLine + } + return strings.ToLower(left.literal) < strings.ToLower(right.literal) +} + +// localizationSourceWindowForHit validates the original match coordinate +// against the live overlay-aware file view, then redacts and bounds it. It never +// searches for a moved match: stale coordinates simply produce no window. +func (s *Server) localizationSourceWindowForHit( + ctx context.Context, + hit exploreSourceLiteralHit, + anchorSymbol string, +) *localizationSourceWindow { + if s == nil || ctx == nil || ctx.Err() != nil || hit.matchLine < 1 || + strings.TrimSpace(hit.matchPath) == "" || strings.TrimSpace(hit.literal) == "" || + strings.TrimSpace(anchorSymbol) == "" { + return nil + } + absPath, displayPath, err := s.resolveFilePath(hit.matchPath) + if err != nil { + return nil + } + content, startLine, _, err := s.readLinesForCtx( + ctx, absPath, hit.matchLine, hit.matchLine, localizationSourceWindowRadius, + ) + if err != nil || ctx.Err() != nil || content == "" || looksBinary([]byte(content)) || !utf8.ValidString(content) { + return nil + } + lines := localizationSourceWindowLines(content) + matchIndex := hit.matchLine - startLine + if matchIndex < 0 || matchIndex >= len(lines) || !exploreTextHasExactLiteral(lines[matchIndex], hit.literal) { + return nil + } + + redacted, redactions := review.RedactSecrets(strings.Join(lines, "\n")) + if !utf8.ValidString(redacted) { + return nil + } + redactedLines := localizationSourceWindowLines(redacted) + if len(redactedLines) != len(lines) || matchIndex >= len(redactedLines) || + !exploreTextHasExactLiteral(redactedLines[matchIndex], hit.literal) { + return nil + } + if displayPath == "" { + displayPath = hit.matchPath + } + window := &localizationSourceWindow{ + Path: displayPath, + StartLine: startLine, + EndLine: startLine + len(redactedLines) - 1, + MatchLine: hit.matchLine, + Content: strings.Join(redactedLines, "\n"), + SecretsRedacted: redactions > 0, + AnchorSymbol: anchorSymbol, + } + window.refreshSize() + for window.Lines > localizationSourceWindowMaxLines || window.Bytes > localizationSourceWindowMaxBytes { + if !window.shrinkOuterLine() { + return nil + } + } + return window +} + +func localizationSourceWindowLines(content string) []string { + lines := strings.Split(content, "\n") + if len(lines) > 1 && lines[len(lines)-1] == "" { + lines = lines[:len(lines)-1] + } + return lines +} + +func (window *localizationSourceWindow) clone() *localizationSourceWindow { + if window == nil { + return nil + } + copy := *window + return © +} + +func (window *localizationSourceWindow) refreshSize() { + if window == nil { + return + } + window.Bytes = len(window.Content) + if window.Content == "" { + window.Lines = 0 + return + } + window.Lines = strings.Count(window.Content, "\n") + 1 +} + +// shrinkOuterLine removes the farther outer line, preferring the trailing side +// on ties, while retaining the complete matching line. +func (window *localizationSourceWindow) shrinkOuterLine() bool { + if window == nil || window.Content == "" { + return false + } + lines := strings.Split(window.Content, "\n") + matchIndex := window.MatchLine - window.StartLine + if len(lines) <= 1 || matchIndex < 0 || matchIndex >= len(lines) { + return false + } + leftDistance := matchIndex + rightDistance := len(lines) - matchIndex - 1 + switch { + case rightDistance >= leftDistance && rightDistance > 0: + lines = lines[:len(lines)-1] + window.EndLine-- + case leftDistance > 0: + lines = lines[1:] + window.StartLine++ + default: + return false + } + window.Content = strings.Join(lines, "\n") + window.Truncated = true + window.refreshSize() + return true +} + +// localizationEnvelopePackingSourceWindow spends only bytes left in the normal +// localization envelope. It never uses the retired-read allowance and never +// removes existing evidence, outlines, bodies, or completion fields. +func localizationEnvelopePackingSourceWindow( + envelope localizationExploreEnvelope, + window *localizationSourceWindow, + maxBytes int, +) localizationExploreEnvelope { + if window == nil || maxBytes < 1 { + return envelope + } + anchorPresent := false + for _, symbol := range envelope.Symbols { + if symbol == window.AnchorSymbol { + anchorPresent = true + break + } + } + if !anchorPresent { + return envelope + } + candidate := envelope + candidate.SourceWindow = window.clone() + for { + if localizationEnvelopeFits(candidate, maxBytes) { + return candidate + } + if !candidate.SourceWindow.shrinkOuterLine() { + return envelope + } + } +} + +func localizationShedSourceWindow(envelope *localizationExploreEnvelope) bool { + if envelope == nil || envelope.SourceWindow == nil { + return false + } + envelope.SourceWindow = nil + return true +} diff --git a/internal/mcp/localization_source_window_test.go b/internal/mcp/localization_source_window_test.go new file mode 100644 index 00000000..a640ba89 --- /dev/null +++ b/internal/mcp/localization_source_window_test.go @@ -0,0 +1,255 @@ +package mcp + +import ( + "context" + "encoding/json" + "fmt" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/zzet/gortex/internal/config" + "github.com/zzet/gortex/internal/daemon" + "github.com/zzet/gortex/internal/graph" + "github.com/zzet/gortex/internal/graph/store_sqlite" + "github.com/zzet/gortex/internal/indexer" + "github.com/zzet/gortex/internal/parser" +) + +func newLocalizationSourceWindowServer(t testing.TB, rel, content string) (*Server, string) { + t.Helper() + root := t.TempDir() + abs := filepath.Join(root, filepath.FromSlash(rel)) + require.NoError(t, os.MkdirAll(filepath.Dir(abs), 0o755)) + require.NoError(t, os.WriteFile(abs, []byte(content), 0o644)) + store, err := store_sqlite.Open(filepath.Join(t.TempDir(), "graph.sqlite")) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, store.Close()) }) + idx := indexer.New(store, parser.NewRegistry(), config.IndexConfig{}, nil) + idx.SetRootPath(root) + return &Server{graph: store, indexer: idx}, abs +} + +func localizationSourceWindowFixture(lines, matchLine int, literal string) string { + out := make([]string, lines) + for index := range out { + out[index] = fmt.Sprintf("line-%03d", index+1) + } + out[matchLine-1] = "register(" + literal + ")" + return strings.Join(out, "\n") + "\n" +} + +func TestLocalizationSourceWindowCentersAndCapsAt201Lines(t *testing.T) { + const rel = "src/registry.go" + server, _ := newLocalizationSourceWindowServer(t, rel, localizationSourceWindowFixture(400, 200, "NEEDLE")) + window := server.localizationSourceWindowForHit(context.Background(), exploreSourceLiteralHit{ + nodeID: "src/registry.go::register", matchPath: rel, matchLine: 200, literal: "NEEDLE", + }, "src/registry.go::register") + + require.NotNil(t, window) + require.Equal(t, 100, window.StartLine) + require.Equal(t, 300, window.EndLine) + require.Equal(t, localizationSourceWindowMaxLines, window.Lines) + require.LessOrEqual(t, window.Bytes, localizationSourceWindowMaxBytes) + require.Contains(t, window.Content, "register(NEEDLE)") +} + +func TestLocalizationSourceWindowClampsAtFileStart(t *testing.T) { + const rel = "src/registry.go" + server, _ := newLocalizationSourceWindowServer(t, rel, localizationSourceWindowFixture(40, 2, "NEEDLE")) + window := server.localizationSourceWindowForHit(context.Background(), exploreSourceLiteralHit{ + nodeID: "src/registry.go::register", matchPath: rel, matchLine: 2, literal: "NEEDLE", + }, "src/registry.go::register") + + require.NotNil(t, window) + require.Equal(t, 1, window.StartLine) + require.Equal(t, 40, window.EndLine) + require.Equal(t, 40, window.Lines) +} + +func TestLocalizationSourceWindowShrinksTo12KiBAndRetainsMatch(t *testing.T) { + const rel = "src/registry.go" + lines := make([]string, 260) + for index := range lines { + lines[index] = fmt.Sprintf("%03d %s", index+1, strings.Repeat("x", 180)) + } + lines[129] = "register(NEEDLE)" + server, _ := newLocalizationSourceWindowServer(t, rel, strings.Join(lines, "\n")) + window := server.localizationSourceWindowForHit(context.Background(), exploreSourceLiteralHit{ + nodeID: "src/registry.go::register", matchPath: rel, matchLine: 130, literal: "NEEDLE", + }, "src/registry.go::register") + + require.NotNil(t, window) + require.LessOrEqual(t, window.Bytes, localizationSourceWindowMaxBytes) + require.True(t, window.Truncated) + require.Less(t, window.Lines, localizationSourceWindowMaxLines) + require.LessOrEqual(t, window.StartLine, window.MatchLine) + require.GreaterOrEqual(t, window.EndLine, window.MatchLine) + require.Contains(t, window.Content, "register(NEEDLE)") +} + +func TestLocalizationSourceWindowOmitsOversizedMatchLine(t *testing.T) { + const rel = "src/registry.go" + content := strings.Repeat("x", localizationSourceWindowMaxBytes) + " NEEDLE" + server, _ := newLocalizationSourceWindowServer(t, rel, content) + window := server.localizationSourceWindowForHit(context.Background(), exploreSourceLiteralHit{ + nodeID: "src/registry.go::register", matchPath: rel, matchLine: 1, literal: "NEEDLE", + }, "src/registry.go::register") + require.Nil(t, window) +} + +func TestLocalizationSourceWindowRejectsStaleMatchLine(t *testing.T) { + const rel = "src/registry.go" + server, _ := newLocalizationSourceWindowServer(t, rel, localizationSourceWindowFixture(20, 10, "OTHER")) + window := server.localizationSourceWindowForHit(context.Background(), exploreSourceLiteralHit{ + nodeID: "src/registry.go::register", matchPath: rel, matchLine: 10, literal: "NEEDLE", + }, "src/registry.go::register") + require.Nil(t, window) +} + +func TestLocalizationSourceWindowValidatesOverlayAwareContent(t *testing.T) { + const rel = "src/registry.go" + server, _ := newLocalizationSourceWindowServer(t, rel, localizationSourceWindowFixture(20, 10, "DISK")) + manager := daemon.NewOverlayManager(0) + require.NoError(t, manager.RegisterWithID("session", "")) + require.NoError(t, manager.Push("session", daemon.OverlayFile{ + Path: rel, Content: localizationSourceWindowFixture(20, 10, "OVERLAY"), + }, nil)) + server.overlays = manager + ctx, _, err := server.prepareOverlayRequest(WithSessionID(context.Background(), "session")) + require.NoError(t, err) + + overlayWindow := server.localizationSourceWindowForHit(ctx, exploreSourceLiteralHit{ + nodeID: "src/registry.go::register", matchPath: rel, matchLine: 10, literal: "OVERLAY", + }, "src/registry.go::register") + require.NotNil(t, overlayWindow) + require.Contains(t, overlayWindow.Content, "OVERLAY") + + staleDisk := server.localizationSourceWindowForHit(ctx, exploreSourceLiteralHit{ + nodeID: "src/registry.go::register", matchPath: rel, matchLine: 10, literal: "DISK", + }, "src/registry.go::register") + require.Nil(t, staleDisk) +} + +func TestLocalizationSourceWindowRedactsSecrets(t *testing.T) { + const rel = "src/registry.go" + content := "func register() {\n" + + " token := \"ghp_123456789012345678901234567890123456\"\n" + + " register(NEEDLE)\n" + + "}\n" + server, _ := newLocalizationSourceWindowServer(t, rel, content) + window := server.localizationSourceWindowForHit(context.Background(), exploreSourceLiteralHit{ + nodeID: "src/registry.go::register", matchPath: rel, matchLine: 3, literal: "NEEDLE", + }, "src/registry.go::register") + + require.NotNil(t, window) + require.True(t, window.SecretsRedacted) + require.NotContains(t, window.Content, "ghp_123456789012345678901234567890123456") + require.Contains(t, window.Content, "redacted") +} + +func TestLocalizationSourceWindowOmitsRedactedMatchLiteral(t *testing.T) { + const rel = "src/registry.go" + const secret = "ghp_123456789012345678901234567890123456" + server, _ := newLocalizationSourceWindowServer(t, rel, "register("+secret+")\n") + window := server.localizationSourceWindowForHit(context.Background(), exploreSourceLiteralHit{ + nodeID: "src/registry.go::register", matchPath: rel, matchLine: 1, literal: secret, + }, "src/registry.go::register") + require.Nil(t, window, "automatic redaction must not leave a window whose visible match vanished") +} + +func TestLocalizationSourceWindowCollectorElectsOnlySurvivingTargetDeterministically(t *testing.T) { + collector := &localizationSourceWindowHitCollector{} + collector.add( + exploreSourceLiteralHit{nodeID: "dropped", matchPath: "d.go", matchLine: 2, literal: "x"}, + exploreSourceLiteralHit{nodeID: "kept", rank: 4, ambiguous: true, matchPath: "z.go", matchLine: 8, literal: "x"}, + exploreSourceLiteralHit{nodeID: "kept", rank: 2, matchPath: "a.go", matchLine: 4, literal: "x"}, + ) + targets := []exploreTarget{{node: &graph.Node{ID: "kept"}, sourceLiteral: true}} + index, hit, ok := collector.elect(targets) + require.True(t, ok) + require.Zero(t, index) + require.Equal(t, "a.go", hit.matchPath) +} + +func TestLocalizationSourceWindowCollectorRejectsUnacceptedOwner(t *testing.T) { + collector := &localizationSourceWindowHitCollector{} + collector.add(exploreSourceLiteralHit{ + nodeID: "ordinary", matchPath: "ordinary.go", matchLine: 2, literal: "x", + }) + _, _, ok := collector.elect([]exploreTarget{{node: &graph.Node{ID: "ordinary"}}}) + require.False(t, ok) +} + +func TestLocalizationSourceWindowPackingUsesOnlyLeftoverAndPreservesEnvelope(t *testing.T) { + base := localizationExploreEnvelope{ + Completion: newLocalizationCompletion(true, ""), + Files: []string{"src/registry.go"}, + Symbols: []string{"src/registry.go::register"}, + Evidence: []localizationEvidence{{ + Rank: 1, ID: "src/registry.go::register", Name: "register", + File: "src/registry.go", Line: 1, + }}, + } + window := &localizationSourceWindow{ + Path: "src/registry.go", StartLine: 1, EndLine: 3, MatchLine: 2, + Content: "before\nregister(NEEDLE)\nafter", AnchorSymbol: "src/registry.go::register", + } + window.refreshSize() + without, err := json.Marshal(base) + require.NoError(t, err) + packed := localizationEnvelopePackingSourceWindow(base, window, len(without)) + require.Nil(t, packed.SourceWindow, "a window must not displace ordinary envelope bytes") + require.Equal(t, base.Evidence, packed.Evidence) + require.Equal(t, base.Completion, packed.Completion) + + packed = localizationEnvelopePackingSourceWindow(base, window, len(without)+1024) + require.NotNil(t, packed.SourceWindow) + require.Equal(t, base.Evidence, packed.Evidence) + require.Equal(t, base.Completion, packed.Completion) +} + +func TestLocalizationSourceWindowShedsBeforeExistingPayload(t *testing.T) { + envelope := localizationExploreEnvelope{ + Evidence: []localizationEvidence{{ID: "kept", Source: "important"}}, + Outline: &localizationFileOutline{File: "src/registry.go"}, + SourceWindow: &localizationSourceWindow{ + Path: "src/registry.go", Content: "optional", AnchorSymbol: "kept", + }, + } + require.True(t, localizationShedSourceWindow(&envelope)) + require.Nil(t, envelope.SourceWindow) + require.Equal(t, "important", envelope.Evidence[0].Source) + require.NotNil(t, envelope.Outline) +} + +func TestLocalizationBuilderOmitsWindowBeforeChangingEvidence(t *testing.T) { + node := &graph.Node{ + ID: "src/registry.go::register", Name: "register", Kind: graph.KindFunction, + FilePath: "src/registry.go", StartLine: 1, EndLine: 3, + } + window := &localizationSourceWindow{ + Path: "src/registry.go", StartLine: 1, EndLine: 1, MatchLine: 1, + Content: strings.Repeat("w", localizationSourceWindowMaxBytes), + AnchorSymbol: node.ID, + } + window.refreshSize() + completion := newLocalizationCompletion(true, "") + baseline, _, baselineDigest, baselineCompletion := buildLocalizationExploreResultForTaskFinalized( + completion, "find register", []exploreTarget{{node: node}}, exploreMinBudgetTokens, + ) + withWindow, _, withDigest, withCompletion := buildLocalizationExploreResultForTaskFinalized( + completion, "find register", []exploreTarget{{node: node, sourceWindow: window}}, exploreMinBudgetTokens, + ) + baselineText, ok := singleTextContent(baseline) + require.True(t, ok) + withText, ok := singleTextContent(withWindow) + require.True(t, ok) + require.Equal(t, baselineText, withText) + require.Equal(t, baselineDigest, withDigest) + require.Equal(t, baselineCompletion, withCompletion) + require.NotContains(t, withText, "source_window") +} diff --git a/internal/mcp/localization_task_aligned_seating_test.go b/internal/mcp/localization_task_aligned_seating_test.go new file mode 100644 index 00000000..83ddd393 --- /dev/null +++ b/internal/mcp/localization_task_aligned_seating_test.go @@ -0,0 +1,150 @@ +package mcp + +import ( + "strings" + "testing" +) + +func TestTaskNamedLeadingFileRowSeatsPrimaryOverRankFill(t *testing.T) { + rows := []localizationDigestRow{ + {ID: "repo/errors.go::ErrorSink.register", Name: "register", QualName: "ErrorSink.register", Kind: "method", File: "repo/errors.go", Line: 12}, + {ID: "repo/socket.go::PacketSocket.open", Name: "open", QualName: "PacketSocket.open", Kind: "method", File: "repo/socket.go", Line: 28}, + {ID: "repo/publisher.go::DatagramPublisher.start", Name: "start", QualName: "DatagramPublisher.start", Kind: "method", File: "repo/publisher.go", Line: 62, supportingOnly: true, leadingFileDepth: true}, + {ID: "repo/socket.go::PacketSocket.close", Name: "close", QualName: "PacketSocket.close", Kind: "method", File: "repo/socket.go", Line: 52}, + {ID: "repo/socket.go::PacketSocket.send", Name: "send", QualName: "PacketSocket.send", Kind: "method", File: "repo/socket.go", Line: 60}, + {ID: "repo/cube.go::CubeSink.connect", Name: "connect", QualName: "CubeSink.connect", Kind: "method", File: "repo/cube.go", Line: 73}, + {ID: "repo/header.go::HeaderCodec.encode", Name: "encode", QualName: "HeaderCodec.encode", Kind: "method", File: "repo/header.go", Line: 91}, + } + response := renderLocalizationFinalResponseForTask( + "undefined constant while starting DatagramPublisher", nil, rows) + want := "- PRIMARY — repo/publisher.go:62 — repo/publisher.go::DatagramPublisher.start" + if !strings.Contains(response, want) { + t.Fatalf("task-named leading-file row was not seated:\n%s", response) + } +} + +func TestTaskAlignedSeatsExtendWithoutDisplacingRankedRows(t *testing.T) { + rows := []localizationDigestRow{ + {ID: "repo/pipeline.go::Pipeline.ingest", Name: "ingest", QualName: "Pipeline.ingest", Kind: "method", File: "repo/pipeline.go", Line: 10}, + {ID: "repo/journal.go::JournalWriter.append", Name: "append", QualName: "JournalWriter.append", Kind: "method", File: "repo/journal.go", Line: 20}, + {ID: "repo/journal.go::JournalWriter.seal", Name: "seal", QualName: "JournalWriter.seal", Kind: "method", File: "repo/journal.go", Line: 34}, + {ID: "repo/journal.go::JournalWriter.compact", Name: "compact", QualName: "JournalWriter.compact", Kind: "method", File: "repo/journal.go", Line: 48}, + {ID: "repo/journal.go::JournalWriter.replayTail", Name: "replayTail", QualName: "JournalWriter.replayTail", Kind: "method", File: "repo/journal.go", Line: 61}, + {ID: "repo/journal.go::JournalWriter.truncate", Name: "truncate", QualName: "JournalWriter.truncate", Kind: "method", File: "repo/journal.go", Line: 75}, + } + response := renderLocalizationFinalResponseForTask( + "journal overflow: append, seal, compact, replayTail and truncate misbehave under JournalWriter load", nil, rows) + // The ranked lead keeps its seat even though it is the only row the task + // does not name; the sixth row rides an extra seat instead. + for _, want := range []string{ + "- PRIMARY — repo/pipeline.go:10 — repo/pipeline.go::Pipeline.ingest", + "- PRIMARY — repo/journal.go:75 — repo/journal.go::JournalWriter.truncate", + } { + if !strings.Contains(response, want) { + t.Fatalf("expected %q as primary:\n%s", want, response) + } + } + if got := strings.Count(response, "- PRIMARY —"); got != localizationFinalResponsePrimaryLimit+1 { + t.Fatalf("primary rows = %d, want %d:\n%s", got, localizationFinalResponsePrimaryLimit+1, response) + } +} + +func TestTaskMentionSeatsWhileScanLanesStaySupporting(t *testing.T) { + rows := []localizationDigestRow{ + {ID: "repo/loop.go::EventLoop.tick", Name: "tick", QualName: "EventLoop.tick", Kind: "method", File: "repo/loop.go", Line: 9}, + {ID: "repo/queue.go::WorkQueue.drain", Name: "drain", QualName: "WorkQueue.drain", Kind: "method", File: "repo/queue.go", Line: 21}, + {ID: "repo/state.go::StateCache.load", Name: "load", QualName: "StateCache.load", Kind: "method", File: "repo/state.go", Line: 33}, + {ID: "repo/codec.go::RingBufferCodec.flushSegment", Name: "flushSegment", QualName: "RingBufferCodec.flushSegment", Kind: "method", File: "repo/codec.go", Line: 47, + Provenance: localizationProvenanceTaskMention}, + {ID: "repo/codec.go::RingBufferCodec.reset", Name: "reset", QualName: "RingBufferCodec.reset", Kind: "method", File: "repo/codec.go", Line: 68, + Provenance: localizationProvenanceDirectAdjacency}, + {ID: "repo/codec.go::RingBufferCodec.scanTail", Name: "scanTail", QualName: "RingBufferCodec.scanTail", Kind: "method", File: "repo/codec.go", Line: 84, + Provenance: localizationProvenanceBodyMention, supportingOnly: true}, + } + response := renderLocalizationFinalResponseForTask( + "overflow inside RingBufferCodec during flush", nil, rows) + seated := "- PRIMARY — repo/codec.go:47 — repo/codec.go::RingBufferCodec.flushSegment" + if !strings.Contains(response, seated) { + t.Fatalf("task-mention row was not seated:\n%s", response) + } + for _, blocked := range []string{ + "- PRIMARY — repo/codec.go:68 — repo/codec.go::RingBufferCodec.reset", + "- PRIMARY — repo/codec.go:84 — repo/codec.go::RingBufferCodec.scanTail", + } { + if strings.Contains(response, blocked) { + t.Fatalf("scan-lane row claimed a primary seat:\n%s", response) + } + } +} + +func TestLocalizationTaskNamesIdentifierWholeWord(t *testing.T) { + cases := []struct { + task, identifier string + minRunes int + want bool + }{ + {"w.Written() is true after headers", "Written", 5, true}, + {"constructing SyslogUdpHandler fails", "SyslogUdpHandler", 6, true}, + {"a task about starting the daemon", "start", 5, false}, + {"the writer is written to disk", "Written", 5, false}, + {"getStorage returns the persist storage", "Storage", 5, false}, + {"Flush()", "Flush", 5, true}, + } + for _, tc := range cases { + if got := localizationTaskNamesIdentifier(tc.task, tc.identifier, tc.minRunes); got != tc.want { + t.Fatalf("localizationTaskNamesIdentifier(%q, %q) = %v, want %v", + tc.task, tc.identifier, got, tc.want) + } + } +} + +func TestTaskNamesRowIgnoresNamespaceAndRepoSegments(t *testing.T) { + task := "Localize bug: in the pre-fix Humanizer checkout, Taxes is singularized wrongly" + namespaceRow := localizationDigestRow{ + ID: "humanizer-1499/src/A.cs::src.Humanizer.TimeSpanHumanizeExtensions.GetNormalCaseTimeAsInteger", + Name: "GetNormalCaseTimeAsInteger", + QualName: "src.Humanizer.TimeSpanHumanizeExtensions.GetNormalCaseTimeAsInteger", + } + if localizationTaskNamesRow(task, namespaceRow) { + t.Fatalf("a namespace segment named the row") + } + repoOwnedType := localizationDigestRow{ + ID: "humanizer-1499/src/B.cs::src.Humanizer.NumberToWordsExtension", + Name: "NumberToWordsExtension", + QualName: "src.Humanizer.NumberToWordsExtension", + } + if localizationTaskNamesRow(task, repoOwnedType) { + t.Fatalf("the repository namespace named a top-level type") + } + declaredOwner := localizationDigestRow{ + ID: "monolog-1569/src/H.php::SyslogUdpHandler.__construct", + Name: "__construct", + QualName: "SyslogUdpHandler.__construct", + } + if !localizationTaskNamesRow("constructing SyslogUdpHandler fails", declaredOwner) { + t.Fatalf("the declaring owner did not name the row") + } +} + +func TestMergeLocalizationDigestRowEvidenceKeepsStrongestProvenance(t *testing.T) { + base := localizationDigestRow{ID: "repo/a.go::Parser.scan", File: "repo/a.go"} + adjacency := base + adjacency.Provenance = localizationProvenanceDirectAdjacency + literal := base + literal.Provenance = localizationProvenanceContentLiteral + structural := base + structural.Provenance = localizationProvenanceImplementationTarget + + if got := mergeLocalizationDigestRowEvidence(adjacency, literal).Provenance; got != localizationProvenanceContentLiteral { + t.Fatalf("literal re-observation lost to adjacency first-write: %q", got) + } + if got := mergeLocalizationDigestRowEvidence(literal, adjacency).Provenance; got != localizationProvenanceContentLiteral { + t.Fatalf("adjacency re-observation erased a literal mark: %q", got) + } + if got := mergeLocalizationDigestRowEvidence(literal, structural).Provenance; got != localizationProvenanceContentLiteral { + t.Fatalf("structural re-observation outranked a literal mark: %q", got) + } + if got := mergeLocalizationDigestRowEvidence(adjacency, structural).Provenance; got != localizationProvenanceImplementationTarget { + t.Fatalf("structural re-observation lost to adjacency first-write: %q", got) + } +} diff --git a/internal/mcp/localization_terminal_evidence_v8_test.go b/internal/mcp/localization_terminal_evidence_v8_test.go index 3cdb7376..ddd34506 100644 --- a/internal/mcp/localization_terminal_evidence_v8_test.go +++ b/internal/mcp/localization_terminal_evidence_v8_test.go @@ -409,7 +409,8 @@ func TestSearchTextCaptureResolvesInScopeOwnerThenFileEvidence(t *testing.T) { ownerGraph.AddNode(owner) ownerServer := &Server{graph: ownerGraph} ctx := withLocalizationPermittedEvidenceCapture(context.Background(), 71) - ownerServer.captureLocalizationSearchText(ctx, []enrichedTextMatch{{Path: "repo/config.go", Line: 24, Text: "register marker"}}) + ownerIndexes := ownerServer.buildFileSymbolIndexForPaths(map[string]struct{}{"repo/config.go": {}}) + ownerServer.captureLocalizationSearchText(ctx, []enrichedTextMatch{{Path: "repo/config.go", Line: 24, Text: "register marker"}}, ownerIndexes) rows, recorded := localizationEvidenceForPermittedCall(ctx, "search", "text", 71) if !recorded || len(rows) != 1 || rows[0].ID != owner.ID || rows[0].Line != 24 || rows[0].Provenance != "permitted_search_text_owner" { t.Fatalf("file-level text owner capture = %#v, recorded=%v", rows, recorded) @@ -420,14 +421,16 @@ func TestSearchTextCaptureResolvesInScopeOwnerThenFileEvidence(t *testing.T) { fileGraph.AddNode(fileOnly) fileServer := &Server{graph: fileGraph} fileCtx := withLocalizationPermittedEvidenceCapture(context.Background(), 72) - fileServer.captureLocalizationSearchText(fileCtx, []enrichedTextMatch{{Path: "repo/root.go", Line: 3, Text: "package marker"}}) + fileIndexes := fileServer.buildFileSymbolIndexForPaths(map[string]struct{}{"repo/root.go": {}}) + fileServer.captureLocalizationSearchText(fileCtx, []enrichedTextMatch{{Path: "repo/root.go", Line: 3, Text: "package marker"}}, fileIndexes) rows, recorded = localizationEvidenceForPermittedCall(fileCtx, "search", "text", 72) if !recorded || len(rows) != 1 || rows[0].ID != fileOnly.ID || rows[0].File != fileOnly.FilePath || rows[0].Line != 3 || rows[0].Provenance != "permitted_search_text_file" { t.Fatalf("file text evidence capture = %#v, recorded=%v", rows, recorded) } missingCtx := withLocalizationPermittedEvidenceCapture(context.Background(), 73) - fileServer.captureLocalizationSearchText(missingCtx, []enrichedTextMatch{{Path: "repo/missing.go", Line: 8, Text: "unattributed marker"}}) + missingIndexes := fileServer.buildFileSymbolIndexForPaths(map[string]struct{}{"repo/missing.go": {}}) + fileServer.captureLocalizationSearchText(missingCtx, []enrichedTextMatch{{Path: "repo/missing.go", Line: 8, Text: "unattributed marker"}}, missingIndexes) if rows, recorded = localizationEvidenceForPermittedCall(missingCtx, "search", "text", 73); !recorded || len(rows) != 0 { t.Fatalf("unvalidated text path became evidence: %#v, recorded=%v", rows, recorded) } diff --git a/internal/mcp/localization_text_index_test.go b/internal/mcp/localization_text_index_test.go new file mode 100644 index 00000000..0790e906 --- /dev/null +++ b/internal/mcp/localization_text_index_test.go @@ -0,0 +1,139 @@ +package mcp + +import ( + "context" + "path/filepath" + "testing" + + "github.com/zzet/gortex/internal/graph" + "github.com/zzet/gortex/internal/query" + "github.com/zzet/gortex/internal/search/trigram" +) + +func TestSearchTextEnrichmentAndCaptureShareOneBoundedFetch(t *testing.T) { + const path = "repo/handler_test.go" + backing := graph.New() + owner := &graph.Node{ + ID: path + "::testOwner", Name: "testOwner", Kind: graph.KindFunction, + FilePath: path, StartLine: 5, EndLine: 30, Meta: map[string]any{"is_test": true}, + } + backing.AddNode(owner) + probe := &boundedFileStoreProbe{Store: backing} + probe.read = func(_ context.Context, gotPath string, scope graph.LocalizationNodeScope, limit int) (graph.BoundedNodeProjection, error) { + if gotPath != path || limit != localizationFileNodeLimit || scope.ExcludeTests { + t.Fatalf("bounded search-text fetch = (%q, %#v, %d)", gotPath, scope, limit) + } + return backing.FindFileNodesBounded(context.Background(), gotPath, scope, limit) + } + server := &Server{graph: probe} + ctx := withLocalizationPermittedEvidenceCapture(context.Background(), 91) + enriched, indexes := server.enrichTextMatchesContext(ctx, []trigram.Match{{Path: path, Line: 10, Text: "needle"}}, query.QueryOptions{}) + if len(enriched) != 1 || enriched[0].SymbolID != owner.ID { + t.Fatalf("test-file owner enrichment = %#v", enriched) + } + server.captureLocalizationSearchText(ctx, enriched, indexes) + if len(probe.calls) != 1 { + t.Fatalf("enrichment + capture performed %d file fetches, want one", len(probe.calls)) + } + rows, recorded := localizationEvidenceForPermittedCall(ctx, "search", "text", 91) + if !recorded || len(rows) != 1 || rows[0].ID != owner.ID { + t.Fatalf("shared-index capture = %#v, recorded=%v", rows, recorded) + } +} + +func TestSearchTextEnrichmentPreservesFirstSeenFilePriority(t *testing.T) { + const ( + firstPath = "repo/z.go" + secondPath = "repo/a.go" + ) + backing := graph.New() + firstOwner := &graph.Node{ + ID: firstPath + "::first", Name: "first", Kind: graph.KindFunction, + FilePath: firstPath, StartLine: 1, EndLine: 20, + } + secondOwner := &graph.Node{ + ID: secondPath + "::second", Name: "second", Kind: graph.KindFunction, + FilePath: secondPath, StartLine: 1, EndLine: 20, + } + backing.AddNode(firstOwner) + backing.AddNode(secondOwner) + probe := &boundedFileStoreProbe{Store: backing} + probe.read = func(ctx context.Context, path string, scope graph.LocalizationNodeScope, limit int) (graph.BoundedNodeProjection, error) { + return backing.FindFileNodesBounded(ctx, path, scope, limit) + } + server := &Server{graph: probe} + + enriched, _ := server.enrichTextMatchesContext(context.Background(), []trigram.Match{ + {Path: firstPath, Line: 3, Text: "first"}, + {Path: secondPath, Line: 4, Text: "second"}, + {Path: firstPath, Line: 5, Text: "first again"}, + }, query.QueryOptions{}) + + if len(probe.calls) != 2 || probe.calls[0].path != firstPath || probe.calls[1].path != secondPath { + t.Fatalf("bounded fetch order = %#v, want first-seen paths [%q %q]", probe.calls, firstPath, secondPath) + } + if len(enriched) != 3 || enriched[0].SymbolID != firstOwner.ID || enriched[1].SymbolID != secondOwner.ID || enriched[2].SymbolID != firstOwner.ID { + t.Fatalf("first-seen enrichment = %#v", enriched) + } +} + +func TestGenericFileSymbolIndexMapOrderingRemainsDeterministic(t *testing.T) { + probe := &boundedFileStoreProbe{Store: graph.New()} + probe.read = func(_ context.Context, _ string, _ graph.LocalizationNodeScope, _ int) (graph.BoundedNodeProjection, error) { + return graph.BoundedNodeProjection{}, nil + } + server := &Server{graph: probe} + server.buildFileSymbolIndexForPathsScopedContext(context.Background(), map[string]struct{}{ + "repo/z.go": {}, + "repo/a.go": {}, + "repo/m.go": {}, + }, query.QueryOptions{}) + + if len(probe.calls) != 3 || probe.calls[0].path != "repo/a.go" || probe.calls[1].path != "repo/m.go" || probe.calls[2].path != "repo/z.go" { + t.Fatalf("generic map fetch order = %#v, want alphabetical determinism", probe.calls) + } +} + +func TestLocalizationTextMatchDirectSymbolIDSkipsFileIndex(t *testing.T) { + store := graph.New() + node := &graph.Node{ID: "repo/direct.go::direct", Name: "direct", Kind: graph.KindFunction, FilePath: "repo/direct.go"} + store.AddNode(node) + server := &Server{graph: store} + got, provenance := server.localizationTextMatchNode( + context.Background(), + enrichedTextMatch{Path: "repo/ignored.go", SymbolID: node.ID}, + map[string]*fileSymbolIndex{"repo/ignored.go": {saturated: true}}, + ) + if got != node || provenance != "permitted_search_text" { + t.Fatalf("direct SymbolID = (%#v, %q), want unchanged typed lookup", got, provenance) + } +} + +func TestLocalizationTextMatchUsesWindowsPathAlias(t *testing.T) { + if filepath.Separator == '/' { + t.Skip("Windows graph path spelling uses a distinct separator only on Windows") + } + server := &Server{graph: graph.New()} + matchPath := "repo/dir/handler.go" + aliasPath := graphMatchPathKey(matchPath, true) + owner := &graph.Node{ID: aliasPath + "::owner", Name: "owner", Kind: graph.KindFunction, FilePath: aliasPath, StartLine: 1, EndLine: 20} + indexes := map[string]*fileSymbolIndex{aliasPath: {syms: []*graph.Node{owner}}} + got, provenance := server.localizationTextMatchNode(context.Background(), enrichedTextMatch{Path: matchPath, Line: 10}, indexes) + if got != owner || provenance != "permitted_search_text_owner" { + t.Fatalf("Windows path alias = (%#v, %q), want owner", got, provenance) + } +} + +func TestLocalizationTextMatchFailsClosedForSaturatedExactPath(t *testing.T) { + server := &Server{graph: graph.New()} + path := "repo/dense.go" + wrong := &graph.Node{ID: path + "::wide", Name: "wide", Kind: graph.KindFunction, FilePath: path, StartLine: 1, EndLine: 100} + got, provenance := server.localizationTextMatchNode( + context.Background(), + enrichedTextMatch{Path: path, Line: 50}, + map[string]*fileSymbolIndex{path: {saturated: true, syms: []*graph.Node{wrong}, fileNode: &graph.Node{ID: path, Kind: graph.KindFile, FilePath: path}}}, + ) + if got != nil || provenance != "" { + t.Fatalf("saturated path misattributed owner/file: (%#v, %q)", got, provenance) + } +} diff --git a/internal/mcp/overlay.go b/internal/mcp/overlay.go index acc627dc..c01fb401 100644 --- a/internal/mcp/overlay.go +++ b/internal/mcp/overlay.go @@ -131,18 +131,23 @@ func (s *Server) wrapToolHandlerMode(h mcpserver.ToolHandlerFunc, injectOverlay // tool call (GORTEX_AUTOINDEX=1). Cheap getenv + sync.Once on the // request path; all real work runs on a background goroutine. s.maybeAutoIndexCWD() + // Every bounded file-summary lookup descended from this tools/call + // shares one request-local allowance. Overlay preparation and facade + // forwarding derive child contexts, so the pointer survives both paths; + // idempotence prevents nested preparation from resetting the budget. + ctx = withLocalizationFileRequestBudget(ctx) if injectOverlay { - view, err := s.buildOverlayViewForCtx(ctx) + var err error + ctx, _, err = s.prepareOverlayRequest(ctx) if err != nil { - // Drift surfaces as a structured tool error result so the - // client knows to re-read and resubmit. Return (result, - // nil) so the JSON-RPC framing carries the message rather - // than a transport error. + if ctxErr := requestContextError(ctx, err); ctxErr != nil { + return nil, ctxErr + } + // Drift and ownership failures surface as structured tool + // errors so the client can refresh and resubmit without a + // transport-level failure. return mcp.NewToolResultError(err.Error()), nil } - if view != nil { - ctx = WithOverlayView(ctx, view) - } } // Warmup fast path: when the daemon is still warming up and // this is a graph-querying tool, the handler still runs (so diff --git a/internal/mcp/overlay_bounded_extractor_test.go b/internal/mcp/overlay_bounded_extractor_test.go new file mode 100644 index 00000000..5b9a90be --- /dev/null +++ b/internal/mcp/overlay_bounded_extractor_test.go @@ -0,0 +1,459 @@ +package mcp + +import ( + "context" + "errors" + "fmt" + "os" + "path/filepath" + "sync" + "testing" + + "github.com/stretchr/testify/require" + "go.uber.org/zap" + + "github.com/zzet/gortex/internal/config" + "github.com/zzet/gortex/internal/daemon" + "github.com/zzet/gortex/internal/graph" + "github.com/zzet/gortex/internal/indexer" + "github.com/zzet/gortex/internal/parser" + "github.com/zzet/gortex/internal/query" +) + +type overlayBoundedExtractorCall struct { + path string + limits parser.ExtractionLimits +} + +type overlayBoundedTestExtractor struct { + mu sync.Mutex + calls []overlayBoundedExtractorCall + legacyCalls int + bounded func(int, context.Context, string, []byte, parser.ExtractionLimits) (*parser.ExtractionResult, parser.ExtractionUsage, error) +} + +func (e *overlayBoundedTestExtractor) Language() string { return "go" } +func (e *overlayBoundedTestExtractor) Extensions() []string { return []string{".go"} } + +func (e *overlayBoundedTestExtractor) Extract(string, []byte) (*parser.ExtractionResult, error) { + e.mu.Lock() + e.legacyCalls++ + e.mu.Unlock() + return nil, errors.New("legacy Extract must not be called for a bounded extractor") +} + +func (e *overlayBoundedTestExtractor) ExtractBounded( + ctx context.Context, + filePath string, + content []byte, + limits parser.ExtractionLimits, +) (*parser.ExtractionResult, parser.ExtractionUsage, error) { + e.mu.Lock() + call := len(e.calls) + e.calls = append(e.calls, overlayBoundedExtractorCall{path: filePath, limits: limits}) + fn := e.bounded + e.mu.Unlock() + if fn == nil { + return &parser.ExtractionResult{}, parser.ExtractionUsage{RawNodes: 1}, nil + } + return fn(call, ctx, filePath, content, limits) +} + +func (e *overlayBoundedTestExtractor) snapshotCalls() ([]overlayBoundedExtractorCall, int) { + e.mu.Lock() + defer e.mu.Unlock() + return append([]overlayBoundedExtractorCall(nil), e.calls...), e.legacyCalls +} + +func setupOverlayExtractorServer( + t *testing.T, + ext parser.Extractor, +) (*Server, string, string, *overlayLayerTestStore) { + t.Helper() + dir := t.TempDir() + extension := ".go" + if extensions := ext.Extensions(); len(extensions) > 0 { + extension = extensions[0] + } + targetFile := filepath.Join(dir, "target"+extension) + callerFile := filepath.Join(dir, "caller"+extension) + require.NoError(t, os.WriteFile(targetFile, []byte("overlay"), 0o644)) + require.NoError(t, os.WriteFile(callerFile, []byte("overlay"), 0o644)) + + base := graph.New() + registry := parser.NewRegistry() + registry.Register(ext) + cfg := config.Default() + idx := indexer.New(base, registry, cfg.Index, zap.NewNop()) + idx.SetRootPath(dir) + idx.SetRepoPrefix("repo") + + store := &overlayLayerTestStore{Store: base} + store.fileFn = func(context.Context, string, graph.LocalizationNodeScope, int) (graph.BoundedNodeProjection, error) { + return graph.BoundedNodeProjection{}, nil + } + store.nameFn = func(context.Context, string, graph.LocalizationNodeScope, int) (graph.BoundedNodeProjection, error) { + return graph.BoundedNodeProjection{}, nil + } + server := &Server{ + graph: store, + indexer: idx, + engine: query.NewEngine(store), + logger: zap.NewNop(), + } + return server, targetFile, callerFile, store +} + +func TestConstructOverlayLayerBoundedExtractorSharesRequestBudgets(t *testing.T) { + extractor := &overlayBoundedTestExtractor{} + server, targetFile, callerFile, _ := setupOverlayExtractorServer(t, extractor) + extractor.bounded = func( + call int, + _ context.Context, + filePath string, + _ []byte, + limits parser.ExtractionLimits, + ) (*parser.ExtractionResult, parser.ExtractionUsage, error) { + switch call { + case 0: + require.Equal(t, overlayLayerParsedNodesMax, limits.MaxNodes) + require.Equal(t, overlayLayerParsedEdgesMax, limits.MaxEdges) + require.Equal(t, overlayLayerParsedResultBytesMax, limits.MaxStdoutBytes) + require.Equal(t, overlayLayerParsedResultBytesMax, limits.MaxResultBytes) + return overlayRawExtractionResult(filePath, 1, 1), parser.ExtractionUsage{ + RawNodes: overlayLayerParsedNodesMax - 1, + RawEdges: overlayLayerParsedEdgesMax, + StdoutBytes: 10, + ResultBytes: overlayLayerParsedResultBytesMax, + }, nil + case 1: + require.Equal(t, 1, limits.MaxNodes) + require.Zero(t, limits.MaxEdges, "zero remaining edges must be passed strictly") + require.Zero(t, limits.MaxStdoutBytes, "zero remaining stdout must be passed strictly") + require.Zero(t, limits.MaxResultBytes, "zero remaining result bytes must be passed strictly") + return overlayRawExtractionResult(filePath, 1, 0), parser.ExtractionUsage{RawNodes: 1}, nil + default: + t.Fatalf("unexpected bounded extraction call %d", call) + return nil, parser.ExtractionUsage{}, nil + } + } + + layer, paths, err := server.constructOverlayLayer(context.Background(), []daemon.OverlayFile{ + {Path: targetFile, Content: "overlay"}, + {Path: callerFile, Content: "overlay"}, + }) + require.NoError(t, err) + require.NotNil(t, layer) + require.Len(t, paths, 2) + calls, legacyCalls := extractor.snapshotCalls() + require.Len(t, calls, 2) + require.Zero(t, legacyCalls) +} + +func TestConstructOverlayLayerRejectsInvalidBoundedUsageWithoutPartialState(t *testing.T) { + tests := []struct { + name string + result func(string) *parser.ExtractionResult + usage parser.ExtractionUsage + want string + }{ + { + name: "missing mandatory file charge", + result: func(string) *parser.ExtractionResult { + return overlayRawExtractionResult("target.go", 0, 0) + }, + usage: parser.ExtractionUsage{}, + want: "negative or missing usage", + }, + { + name: "negative edge charge", + result: func(path string) *parser.ExtractionResult { + return overlayRawExtractionResult(path, 1, 0) + }, + usage: parser.ExtractionUsage{RawNodes: 1, RawEdges: -1}, + want: "negative or missing usage", + }, + { + name: "stdout larger than result", + result: func(path string) *parser.ExtractionResult { + return overlayRawExtractionResult(path, 1, 0) + }, + usage: parser.ExtractionUsage{RawNodes: 1, StdoutBytes: 2, ResultBytes: 1}, + want: "result usage is smaller", + }, + { + name: "returned nodes exceed raw charge", + result: func(path string) *parser.ExtractionResult { + return overlayRawExtractionResult(path, 2, 0) + }, + usage: parser.ExtractionUsage{RawNodes: 1}, + want: "result exceeds reported raw usage", + }, + { + name: "raw charge exceeds passed remainder", + result: func(path string) *parser.ExtractionResult { + return overlayRawExtractionResult(path, 1, 0) + }, + usage: parser.ExtractionUsage{RawNodes: overlayLayerParsedNodesMax + 1}, + want: "parsed nodes exceeds limit", + }, + { + name: "negative node charge", + result: func(path string) *parser.ExtractionResult { + return overlayRawExtractionResult(path, 1, 0) + }, + usage: parser.ExtractionUsage{RawNodes: -1}, + want: "negative or missing usage", + }, + { + name: "negative stdout charge", + result: func(path string) *parser.ExtractionResult { + return overlayRawExtractionResult(path, 1, 0) + }, + usage: parser.ExtractionUsage{RawNodes: 1, StdoutBytes: -1}, + want: "negative or missing usage", + }, + { + name: "negative result charge", + result: func(path string) *parser.ExtractionResult { + return overlayRawExtractionResult(path, 1, 0) + }, + usage: parser.ExtractionUsage{RawNodes: 1, ResultBytes: -1}, + want: "negative or missing usage", + }, + { + name: "returned edges exceed raw charge", + result: func(path string) *parser.ExtractionResult { + return overlayRawExtractionResult(path, 1, 1) + }, + usage: parser.ExtractionUsage{RawNodes: 1}, + want: "result exceeds reported raw usage", + }, + { + name: "raw edges exceed passed remainder", + result: func(path string) *parser.ExtractionResult { + return overlayRawExtractionResult(path, 1, 0) + }, + usage: parser.ExtractionUsage{RawNodes: 1, RawEdges: overlayLayerParsedEdgesMax + 1}, + want: "parsed edges exceeds limit", + }, + { + name: "stdout exceeds passed remainder", + result: func(path string) *parser.ExtractionResult { + return overlayRawExtractionResult(path, 1, 0) + }, + usage: parser.ExtractionUsage{ + RawNodes: 1, StdoutBytes: overlayLayerParsedResultBytesMax + 1, + ResultBytes: overlayLayerParsedResultBytesMax + 1, + }, + want: "parsed stdout bytes exceeds limit", + }, + { + name: "result exceeds passed remainder", + result: func(path string) *parser.ExtractionResult { + return overlayRawExtractionResult(path, 1, 0) + }, + usage: parser.ExtractionUsage{RawNodes: 1, ResultBytes: overlayLayerParsedResultBytesMax + 1}, + want: "parsed result bytes exceeds limit", + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + extractor := &overlayBoundedTestExtractor{} + server, targetFile, _, store := setupOverlayExtractorServer(t, extractor) + var returned *parser.ExtractionResult + extractor.bounded = func( + _ int, _ context.Context, filePath string, _ []byte, _ parser.ExtractionLimits, + ) (*parser.ExtractionResult, parser.ExtractionUsage, error) { + returned = test.result(filePath) + return returned, test.usage, nil + } + + layer, paths, err := server.constructOverlayLayer(context.Background(), []daemon.OverlayFile{{Path: targetFile, Content: "overlay"}}) + require.ErrorContains(t, err, test.want) + require.Nil(t, layer) + require.Nil(t, paths) + require.NotNil(t, returned) + require.Nil(t, returned.Tree) + require.Nil(t, returned.ConstValues) + fileCalls, _ := store.snapshotFileCalls() + require.Empty(t, fileCalls) + }) + } +} + +func TestConstructOverlayLayerBoundedErrorAndCancellationCleanResult(t *testing.T) { + t.Run("result and error", func(t *testing.T) { + extractor := &overlayBoundedTestExtractor{} + server, targetFile, _, store := setupOverlayExtractorServer(t, extractor) + sentinel := errors.New("bounded extractor failed") + result := overlayRawExtractionResult("target.go", 1, 1) + extractor.bounded = func(int, context.Context, string, []byte, parser.ExtractionLimits) (*parser.ExtractionResult, parser.ExtractionUsage, error) { + return result, parser.ExtractionUsage{}, sentinel + } + layer, paths, err := server.constructOverlayLayer(context.Background(), []daemon.OverlayFile{{Path: targetFile, Content: "overlay"}}) + require.ErrorIs(t, err, sentinel) + require.Nil(t, layer) + require.Nil(t, paths) + require.Nil(t, result.Tree) + require.Nil(t, result.ConstValues) + fileCalls, _ := store.snapshotFileCalls() + require.Empty(t, fileCalls) + }) + + t.Run("parent cancellation wins", func(t *testing.T) { + extractor := &overlayBoundedTestExtractor{} + server, targetFile, _, store := setupOverlayExtractorServer(t, extractor) + ctx, cancel := context.WithCancel(context.Background()) + result := overlayRawExtractionResult("target.go", 1, 1) + extractor.bounded = func(int, context.Context, string, []byte, parser.ExtractionLimits) (*parser.ExtractionResult, parser.ExtractionUsage, error) { + cancel() + return result, parser.ExtractionUsage{}, errors.New("component failure") + } + layer, paths, err := server.constructOverlayLayer(ctx, []daemon.OverlayFile{{Path: targetFile, Content: "overlay"}}) + require.ErrorIs(t, err, context.Canceled) + require.Nil(t, layer) + require.Nil(t, paths) + require.Nil(t, result.Tree) + require.Nil(t, result.ConstValues) + fileCalls, _ := store.snapshotFileCalls() + require.Empty(t, fileCalls) + }) + + t.Run("component deadline remains an ordinary live-parent error", func(t *testing.T) { + extractor := &overlayBoundedTestExtractor{} + server, targetFile, _, store := setupOverlayExtractorServer(t, extractor) + ctx := context.Background() + result := overlayRawExtractionResult("target.go", 1, 1) + extractor.bounded = func(int, context.Context, string, []byte, parser.ExtractionLimits) (*parser.ExtractionResult, parser.ExtractionUsage, error) { + return result, parser.ExtractionUsage{}, context.DeadlineExceeded + } + layer, paths, err := server.constructOverlayLayer(ctx, []daemon.OverlayFile{{Path: targetFile, Content: "overlay"}}) + require.ErrorIs(t, err, context.DeadlineExceeded) + require.NoError(t, ctx.Err(), "the parent request remains live") + require.Nil(t, layer) + require.Nil(t, paths) + require.Nil(t, result.Tree) + require.Nil(t, result.ConstValues) + fileCalls, _ := store.snapshotFileCalls() + require.Empty(t, fileCalls) + }) +} + +func TestConstructOverlayLayerBoundedAggregatePlusOneFailsAtomically(t *testing.T) { + tests := []struct { + name string + first parser.ExtractionUsage + resource string + secondLimits func(*testing.T, parser.ExtractionLimits) + }{ + { + name: "raw nodes", + first: parser.ExtractionUsage{RawNodes: overlayLayerParsedNodesMax}, + resource: "nodes", + secondLimits: func(t *testing.T, limits parser.ExtractionLimits) { + require.Zero(t, limits.MaxNodes) + }, + }, + { + name: "raw edges", + first: parser.ExtractionUsage{RawNodes: 1, RawEdges: overlayLayerParsedEdgesMax}, + resource: "edges", + secondLimits: func(t *testing.T, limits parser.ExtractionLimits) { + require.Zero(t, limits.MaxEdges) + }, + }, + { + name: "result bytes", + first: parser.ExtractionUsage{RawNodes: 1, StdoutBytes: 1, ResultBytes: overlayLayerParsedResultBytesMax}, + resource: "result_bytes", + secondLimits: func(t *testing.T, limits parser.ExtractionLimits) { + require.Zero(t, limits.MaxStdoutBytes) + require.Zero(t, limits.MaxResultBytes) + }, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + extractor := &overlayBoundedTestExtractor{} + server, targetFile, callerFile, store := setupOverlayExtractorServer(t, extractor) + var first, overflow *parser.ExtractionResult + extractor.bounded = func( + call int, _ context.Context, filePath string, _ []byte, limits parser.ExtractionLimits, + ) (*parser.ExtractionResult, parser.ExtractionUsage, error) { + if call == 0 { + first = overlayRawExtractionResult(filePath, 1, 0) + return first, test.first, nil + } + test.secondLimits(t, limits) + overflow = overlayRawExtractionResult(filePath, 1, 0) + limit := int64(0) + switch test.resource { + case "nodes": + limit = int64(limits.MaxNodes) + case "edges": + limit = int64(limits.MaxEdges) + case "result_bytes": + limit = int64(limits.MaxResultBytes) + } + return overflow, parser.ExtractionUsage{}, &parser.ExtractionLimitError{ + Resource: test.resource, Limit: limit, Observed: limit + 1, + } + } + + layer, paths, err := server.constructOverlayLayer(context.Background(), []daemon.OverlayFile{ + {Path: targetFile, Content: "overlay"}, + {Path: callerFile, Content: "overlay"}, + }) + require.ErrorIs(t, err, parser.ErrExtractionLimit) + var limitErr *parser.ExtractionLimitError + require.ErrorAs(t, err, &limitErr, "wrapping must preserve the typed parser limit") + require.Equal(t, test.resource, limitErr.Resource) + require.Nil(t, layer) + require.Nil(t, paths) + require.NotNil(t, first) + require.NotNil(t, overflow) + require.Nil(t, first.Tree) + require.Nil(t, first.ConstValues) + require.Nil(t, overflow.Tree) + require.Nil(t, overflow.ConstValues) + fileCalls, _ := store.snapshotFileCalls() + require.Len(t, fileCalls, 1, "failed second extraction must not read or stage its base file") + }) + } +} + +func TestConstructOverlayLayerBoundedAcceptsFull257FileCohort(t *testing.T) { + extractor := &overlayBoundedTestExtractor{} + server, targetFile, _, store := setupOverlayExtractorServer(t, extractor) + extractor.bounded = func( + call int, _ context.Context, filePath string, _ []byte, limits parser.ExtractionLimits, + ) (*parser.ExtractionResult, parser.ExtractionUsage, error) { + require.Equal(t, overlayLayerParsedNodesMax-call, limits.MaxNodes) + require.Equal(t, overlayLayerParsedEdgesMax, limits.MaxEdges) + require.Equal(t, overlayLayerParsedResultBytesMax-call, limits.MaxStdoutBytes) + require.Equal(t, overlayLayerParsedResultBytesMax-call, limits.MaxResultBytes) + return overlayRawExtractionResult(filePath, 1, 0), parser.ExtractionUsage{ + RawNodes: 1, ResultBytes: 1, + }, nil + } + + files := make([]daemon.OverlayFile, overlayRequestSnapshotMaxFiles) + for index := range files { + files[index] = daemon.OverlayFile{ + Path: filepath.Join(filepath.Dir(targetFile), "bounded-"+fmt.Sprintf("%03d.go", index)), + Content: "overlay", + } + } + layer, paths, err := server.constructOverlayLayer(context.Background(), files) + require.NoError(t, err) + require.NotNil(t, layer) + require.Len(t, paths, overlayRequestSnapshotMaxFiles) + calls, legacyCalls := extractor.snapshotCalls() + require.Len(t, calls, overlayRequestSnapshotMaxFiles) + require.Zero(t, legacyCalls) + fileCalls, _ := store.snapshotFileCalls() + require.Len(t, fileCalls, overlayRequestSnapshotMaxFiles) +} diff --git a/internal/mcp/overlay_e2e_test.go b/internal/mcp/overlay_e2e_test.go index d8629268..ebad75d4 100644 --- a/internal/mcp/overlay_e2e_test.go +++ b/internal/mcp/overlay_e2e_test.go @@ -535,6 +535,180 @@ func TestOverlay_ListExposesExpiryMetadata(t *testing.T) { require.Contains(t, body, `"idle_ttl_seconds"`) } +func TestOverlay_RequestSnapshotPinsWorkspaceAndFilesAcrossConcurrentPush(t *testing.T) { + srv, _, targetFile, _ := setupOverlayServer(t) + srv.indexer.SetWorkspaceID("workspace-a") + const sessionID = "request-snapshot" + require.NoError(t, srv.OverlayManager().RegisterWithID(sessionID, "workspace-a")) + const firstContent = "package main\n\nfunc FirstBuffer() {}\n" + require.NoError(t, srv.OverlayManager().Push(sessionID, daemon.OverlayFile{ + Path: targetFile, Content: firstContent, + }, nil)) + + ctx, firstView, err := srv.prepareOverlayRequest(WithSessionID(context.Background(), sessionID)) + require.NoError(t, err) + require.NotNil(t, firstView) + snapshot, ok := overlayRequestSnapshotFromContext(ctx) + require.True(t, ok) + require.Equal(t, "workspace-a", snapshot.workspace) + require.Len(t, snapshot.files, 1) + require.Equal(t, firstContent, snapshot.files[0].Content) + + pushDone := make(chan error, 1) + go func() { + pushDone <- srv.OverlayManager().Push(sessionID, daemon.OverlayFile{ + Path: targetFile, Content: "package main\n\nfunc SecondBuffer() {}\n", + }, nil) + }() + require.NoError(t, <-pushDone) + + content, found := srv.overlayContentFor(ctx, targetFile) + require.True(t, found) + require.Equal(t, firstContent, content, "the prepared request must retain its immutable buffer cohort") + reusedCtx, reusedView, err := srv.prepareOverlayRequest(ctx) + require.NoError(t, err) + require.Same(t, firstView, reusedView) + reusedSnapshot, ok := overlayRequestSnapshotFromContext(reusedCtx) + require.True(t, ok) + require.Same(t, snapshot, reusedSnapshot) + + mismatchedCtx := WithSessionID(ctx, "different-session") + _, _, err = srv.prepareOverlayRequest(mismatchedCtx) + require.ErrorContains(t, err, "belongs to session") +} + +func TestOverlay_RequestSnapshotPinsEmptyAttemptAcrossLaterPush(t *testing.T) { + srv, _, targetFile, _ := setupOverlayServer(t) + srv.indexer.SetWorkspaceID("workspace-a") + const sessionID = "empty-request-snapshot" + require.NoError(t, srv.OverlayManager().RegisterWithID(sessionID, "workspace-a")) + + ctx, view, err := srv.prepareOverlayRequest(WithSessionID(context.Background(), sessionID)) + require.NoError(t, err) + require.Nil(t, view) + snapshot, ok := overlayRequestSnapshotFromContext(ctx) + require.True(t, ok, "even an empty attempt must be pinned") + require.Equal(t, "workspace-a", snapshot.workspace) + require.Empty(t, snapshot.files) + + require.NoError(t, srv.OverlayManager().Push(sessionID, daemon.OverlayFile{ + Path: targetFile, Content: "package main\n\nfunc LaterBuffer() {}\n", + }, nil)) + reusedCtx, reusedView, err := srv.prepareOverlayRequest(ctx) + require.NoError(t, err) + require.Nil(t, reusedView, "a nested call must not resnapshot buffers pushed mid-request") + reusedSnapshot, ok := overlayRequestSnapshotFromContext(reusedCtx) + require.True(t, ok) + require.Same(t, snapshot, reusedSnapshot) + _, found := srv.overlayContentFor(reusedCtx, targetFile) + require.False(t, found) +} + +func TestOverlay_RequestSnapshotRejectsForeignWorkspaceAndPath(t *testing.T) { + t.Run("workspace mismatch", func(t *testing.T) { + srv, _, targetFile, _ := setupOverlayServer(t) + srv.indexer.SetWorkspaceID("workspace-a") + const sessionID = "foreign-workspace" + require.NoError(t, srv.OverlayManager().RegisterWithID(sessionID, "workspace-b")) + require.NoError(t, srv.OverlayManager().Push(sessionID, daemon.OverlayFile{ + Path: targetFile, Deleted: true, + }, nil)) + + _, _, err := srv.prepareOverlayRequest(WithSessionID(context.Background(), sessionID)) + require.ErrorContains(t, err, "not registered workspace") + }) + + t.Run("untracked path", func(t *testing.T) { + srv, _, _, _ := setupOverlayServer(t) + srv.indexer.SetWorkspaceID("workspace-a") + const sessionID = "foreign-path" + require.NoError(t, srv.OverlayManager().RegisterWithID(sessionID, "workspace-a")) + require.NoError(t, srv.OverlayManager().Push(sessionID, daemon.OverlayFile{ + Path: filepath.Join(t.TempDir(), "foreign.go"), Deleted: true, + }, nil)) + + _, _, err := srv.prepareOverlayRequest(WithSessionID(context.Background(), sessionID)) + require.ErrorContains(t, err, "outside the registered workspace") + }) +} + +func TestOverlay_RequestSnapshotCanonicalizesEquivalentAliases(t *testing.T) { + srv, _, targetFile, _ := setupOverlayServer(t) + const sessionID = "canonical-aliases" + require.NoError(t, srv.OverlayManager().RegisterWithID(sessionID, "")) + const content = "package main\n\nfunc CanonicalBuffer() {}\n" + for _, alias := range []string{targetFile, filepath.Base(targetFile), "./" + filepath.Base(targetFile)} { + require.NoError(t, srv.OverlayManager().Push(sessionID, daemon.OverlayFile{ + Path: alias, Content: content, BaseSHA: "", + }, nil)) + } + + ctx, view, err := srv.prepareOverlayRequest(WithSessionID(context.Background(), sessionID)) + require.NoError(t, err) + require.NotNil(t, view) + snapshot, ok := overlayRequestSnapshotFromContext(ctx) + require.True(t, ok) + require.True(t, snapshot.canonical) + require.Len(t, snapshot.files, 1) + require.Equal(t, filepath.Base(targetFile), snapshot.files[0].Path) + require.Equal(t, content, snapshot.files[0].Content) + got, found := srv.overlayContentFor(ctx, targetFile) + require.True(t, found) + require.Equal(t, content, got) +} + +func TestOverlay_RequestSnapshotRejectsConflictingAliases(t *testing.T) { + for _, test := range []struct { + name string + first daemon.OverlayFile + second daemon.OverlayFile + }{ + { + name: "content", + first: daemon.OverlayFile{Content: "package main\nfunc First() {}\n"}, + second: daemon.OverlayFile{Content: "package main\nfunc Second() {}\n"}, + }, + { + name: "tombstone", + first: daemon.OverlayFile{Content: "package main\nfunc Present() {}\n"}, + second: daemon.OverlayFile{Deleted: true}, + }, + { + name: "base sha", + first: daemon.OverlayFile{Content: "package main\n", BaseSHA: "first"}, + second: daemon.OverlayFile{Content: "package main\n", BaseSHA: "second"}, + }, + } { + t.Run(test.name, func(t *testing.T) { + srv, _, targetFile, _ := setupOverlayServer(t) + const sessionID = "conflicting-aliases" + require.NoError(t, srv.OverlayManager().RegisterWithID(sessionID, "")) + first, second := test.first, test.second + first.Path = targetFile + second.Path = filepath.Base(targetFile) + require.NoError(t, srv.OverlayManager().Push(sessionID, first, nil)) + require.NoError(t, srv.OverlayManager().Push(sessionID, second, nil)) + + _, _, err := srv.prepareOverlayRequest(WithSessionID(context.Background(), sessionID)) + require.ErrorContains(t, err, "conflicting overlay aliases") + }) + } +} + +func TestOverlay_RequestSnapshotRejectsNonCanonicalSnapshotWithView(t *testing.T) { + srv, _, targetFile, _ := setupOverlayServer(t) + const sessionID = "noncanonical-view" + ctx := WithSessionID(context.Background(), sessionID) + ctx = withOverlayRequestSnapshot(ctx, &overlayRequestSnapshot{ + sessionID: sessionID, + files: []daemon.OverlayFile{{Path: targetFile, Deleted: true}}, + }) + ctx = WithOverlayView(ctx, graph.NewOverlaidView(srv.graph, graph.NewOverlayLayer())) + + _, _, err := srv.prepareOverlayRequest(ctx) + require.ErrorContains(t, err, "non-canonical request snapshot") +} + // baseNodeIDs returns a sorted slice of every node ID in the base // graph. Used to verify the shadow-graph design's load-bearing // invariant: base is never mutated during overlay processing. diff --git a/internal/mcp/overlay_layer_base_bounds_test.go b/internal/mcp/overlay_layer_base_bounds_test.go new file mode 100644 index 00000000..6e7f31cc --- /dev/null +++ b/internal/mcp/overlay_layer_base_bounds_test.go @@ -0,0 +1,222 @@ +package mcp + +import ( + "context" + "errors" + "fmt" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/zzet/gortex/internal/graph" +) + +type overlayLayerFileReaderFunc func( + context.Context, + string, + graph.LocalizationNodeScope, + int, +) (graph.BoundedNodeProjection, error) + +func (f overlayLayerFileReaderFunc) FindFileNodesBounded( + ctx context.Context, + filePath string, + scope graph.LocalizationNodeScope, + limit int, +) (graph.BoundedNodeProjection, error) { + return f(ctx, filePath, scope, limit) +} + +func overlayLayerProjection(count, limit int, filePath string) graph.BoundedNodeProjection { + kept := count + if kept > limit { + kept = limit + } + nodes := make([]*graph.Node, kept) + for index := range nodes { + nodes[index] = &graph.Node{ + ID: fmt.Sprintf("%s::node-%d", filePath, index), + Name: fmt.Sprintf("node-%d", index), + FilePath: filePath, + } + } + total := count + if total > limit+1 { + total = limit + 1 + } + return graph.BoundedNodeProjection{ + Nodes: nodes, + Total: total, + Truncated: count > limit, + } +} + +func TestReadOverlayBaseNodesPerFileBounds(t *testing.T) { + t.Run("exact", func(t *testing.T) { + total := 0 + reader := overlayLayerFileReaderFunc(func( + _ context.Context, + filePath string, + scope graph.LocalizationNodeScope, + limit int, + ) (graph.BoundedNodeProjection, error) { + require.Equal(t, graph.LocalizationNodeScope{}, scope) + require.Equal(t, overlayLayerBaseNodesPerFileMax, limit) + return overlayLayerProjection(overlayLayerBaseNodesPerFileMax, limit, filePath), nil + }) + + nodes, err := readOverlayBaseNodes(context.Background(), reader, "exact.go", &total) + require.NoError(t, err) + require.Len(t, nodes, overlayLayerBaseNodesPerFileMax) + require.Equal(t, overlayLayerBaseNodesPerFileMax, total) + }) + + t.Run("plus_one", func(t *testing.T) { + total := 0 + reader := overlayLayerFileReaderFunc(func( + _ context.Context, + filePath string, + _ graph.LocalizationNodeScope, + limit int, + ) (graph.BoundedNodeProjection, error) { + return overlayLayerProjection(overlayLayerBaseNodesPerFileMax+1, limit, filePath), nil + }) + + _, err := readOverlayBaseNodes(context.Background(), reader, "overflow.go", &total) + var limitErr *overlayLayerLimitError + require.ErrorAs(t, err, &limitErr) + require.Equal(t, "base nodes per file", limitErr.Resource) + require.Zero(t, total) + }) +} + +func TestReadOverlayBaseNodesAggregateSentinel(t *testing.T) { + counts := map[string]int{ + "one.go": 1024, "two.go": 1024, "three.go": 1024, "four.go": 1024, + "empty.go": 0, "one-more.go": 1, + } + var limits []int + reader := overlayLayerFileReaderFunc(func( + _ context.Context, + filePath string, + _ graph.LocalizationNodeScope, + limit int, + ) (graph.BoundedNodeProjection, error) { + limits = append(limits, limit) + return overlayLayerProjection(counts[filePath], limit, filePath), nil + }) + + total := 0 + for _, filePath := range []string{"one.go", "two.go", "three.go", "four.go"} { + _, err := readOverlayBaseNodes(context.Background(), reader, filePath, &total) + require.NoError(t, err) + } + require.Equal(t, overlayLayerBaseNodesMax, total) + + nodes, err := readOverlayBaseNodes(context.Background(), reader, "empty.go", &total) + require.NoError(t, err) + require.Empty(t, nodes) + require.Equal(t, 1, limits[len(limits)-1], "full budget must use a non-zero emptiness sentinel") + require.Equal(t, overlayLayerBaseNodesMax, total) + + _, err = readOverlayBaseNodes(context.Background(), reader, "one-more.go", &total) + var limitErr *overlayLayerLimitError + require.ErrorAs(t, err, &limitErr) + require.Equal(t, "base nodes", limitErr.Resource) + require.Equal(t, overlayLayerBaseNodesMax, total, "rejection must not retain the sentinel row") + require.Equal(t, 1, limits[len(limits)-1]) +} + +func TestReadOverlayBaseNodesRemainingAggregateBudget(t *testing.T) { + t.Run("exact_remaining", func(t *testing.T) { + total := 4000 + reader := overlayLayerFileReaderFunc(func( + _ context.Context, + filePath string, + _ graph.LocalizationNodeScope, + limit int, + ) (graph.BoundedNodeProjection, error) { + require.Equal(t, 96, limit) + return overlayLayerProjection(96, limit, filePath), nil + }) + _, err := readOverlayBaseNodes(context.Background(), reader, "remaining.go", &total) + require.NoError(t, err) + require.Equal(t, overlayLayerBaseNodesMax, total) + }) + + t.Run("plus_one_remaining", func(t *testing.T) { + total := 4000 + reader := overlayLayerFileReaderFunc(func( + _ context.Context, + filePath string, + _ graph.LocalizationNodeScope, + limit int, + ) (graph.BoundedNodeProjection, error) { + require.Equal(t, 96, limit) + return overlayLayerProjection(97, limit, filePath), nil + }) + _, err := readOverlayBaseNodes(context.Background(), reader, "overflow.go", &total) + var limitErr *overlayLayerLimitError + require.ErrorAs(t, err, &limitErr) + require.Equal(t, "base nodes", limitErr.Resource) + require.Equal(t, 4000, total) + }) +} + +func TestReadOverlayBaseNodesFailsClosed(t *testing.T) { + t.Run("unsupported", func(t *testing.T) { + total := 0 + _, err := readOverlayBaseNodes(context.Background(), nil, "file.go", &total) + require.ErrorIs(t, err, graph.ErrBoundedLocalizationUnavailable) + }) + + t.Run("reader_error", func(t *testing.T) { + sentinel := errors.New("read failed") + total := 0 + reader := overlayLayerFileReaderFunc(func( + context.Context, string, graph.LocalizationNodeScope, int, + ) (graph.BoundedNodeProjection, error) { + return graph.BoundedNodeProjection{}, sentinel + }) + _, err := readOverlayBaseNodes(context.Background(), reader, "file.go", &total) + require.ErrorIs(t, err, sentinel) + }) + + t.Run("invalid_projection", func(t *testing.T) { + total := 0 + reader := overlayLayerFileReaderFunc(func( + context.Context, string, graph.LocalizationNodeScope, int, + ) (graph.BoundedNodeProjection, error) { + return graph.BoundedNodeProjection{Total: 1}, nil + }) + _, err := readOverlayBaseNodes(context.Background(), reader, "file.go", &total) + require.ErrorContains(t, err, "invalid bounded projection") + require.Zero(t, total) + }) + + t.Run("nil_node", func(t *testing.T) { + total := 0 + reader := overlayLayerFileReaderFunc(func( + context.Context, string, graph.LocalizationNodeScope, int, + ) (graph.BoundedNodeProjection, error) { + return graph.BoundedNodeProjection{Nodes: []*graph.Node{nil}, Total: 1}, nil + }) + _, err := readOverlayBaseNodes(context.Background(), reader, "file.go", &total) + require.ErrorContains(t, err, "invalid nil node") + require.Zero(t, total) + }) + + t.Run("cancel_after_reader", func(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + total := 0 + reader := overlayLayerFileReaderFunc(func( + context.Context, string, graph.LocalizationNodeScope, int, + ) (graph.BoundedNodeProjection, error) { + cancel() + return graph.BoundedNodeProjection{}, nil + }) + _, err := readOverlayBaseNodes(ctx, reader, "file.go", &total) + require.ErrorIs(t, err, context.Canceled) + require.Zero(t, total) + }) +} diff --git a/internal/mcp/overlay_layer_construction_bounds_test.go b/internal/mcp/overlay_layer_construction_bounds_test.go new file mode 100644 index 00000000..d8432dd8 --- /dev/null +++ b/internal/mcp/overlay_layer_construction_bounds_test.go @@ -0,0 +1,331 @@ +package mcp + +import ( + "context" + "errors" + "fmt" + "os" + "path/filepath" + "testing" + + "github.com/stretchr/testify/require" + "go.uber.org/zap" + + "github.com/zzet/gortex/internal/config" + "github.com/zzet/gortex/internal/daemon" + "github.com/zzet/gortex/internal/graph" + "github.com/zzet/gortex/internal/indexer" + "github.com/zzet/gortex/internal/parser" + "github.com/zzet/gortex/internal/query" +) + +func setupOverlayLayerBoundServer(t *testing.T) (*Server, string, string, *overlayLayerTestStore, *overlayLayerTestExtractor) { + t.Helper() + dir := t.TempDir() + targetFile := filepath.Join(dir, "target.go") + callerFile := filepath.Join(dir, "caller.go") + require.NoError(t, os.WriteFile(targetFile, []byte("package main\n"), 0o644)) + require.NoError(t, os.WriteFile(callerFile, []byte("package main\n"), 0o644)) + + base := graph.New() + extractor := &overlayLayerTestExtractor{} + registry := parser.NewRegistry() + registry.Register(extractor) + cfg := config.Default() + idx := indexer.New(base, registry, cfg.Index, zap.NewNop()) + idx.SetRootPath(dir) + idx.SetRepoPrefix("repo") + + store := &overlayLayerTestStore{Store: base} + store.fileFn = func( + context.Context, string, graph.LocalizationNodeScope, int, + ) (graph.BoundedNodeProjection, error) { + return graph.BoundedNodeProjection{}, nil + } + store.nameFn = func( + context.Context, string, graph.LocalizationNodeScope, int, + ) (graph.BoundedNodeProjection, error) { + return graph.BoundedNodeProjection{}, nil + } + srv := &Server{ + graph: store, + indexer: idx, + engine: query.NewEngine(store), + logger: zap.NewNop(), + } + return srv, targetFile, callerFile, store, extractor +} + +func overlayRawExtractionResult(filePath string, nodeCount, edgeCount int) *parser.ExtractionResult { + nodes := make([]*graph.Node, nodeCount) + if nodeCount > 0 { + nodes[0] = &graph.Node{ + ID: filePath + "::node", + Name: "node-" + filePath, + FilePath: filePath, + Meta: map[string]any{"file": filePath}, + } + } + edges := make([]*graph.Edge, edgeCount) + if edgeCount > 0 { + edges[0] = &graph.Edge{ + From: filePath + "::node", + To: "external::target", + Meta: map[string]any{"file": filePath}, + } + } + return &parser.ExtractionResult{ + Nodes: nodes, + Edges: edges, + Tree: parser.NewParseTree(nil, []byte(filePath), "go"), + ConstValues: []parser.ConstValue{{NodeID: filePath + "::const", FilePath: filePath, Value: "large-sidecar"}}, + } +} + +func TestConstructOverlayLayerParsedEntryBounds(t *testing.T) { + t.Run("exact_aggregate_with_nil_entries", func(t *testing.T) { + srv, targetFile, callerFile, store, extractor := setupOverlayLayerBoundServer(t) + var results []*parser.ExtractionResult + extractor.extract = func(_ int, filePath string, _ []byte) (*parser.ExtractionResult, error) { + result := overlayRawExtractionResult(filePath, overlayLayerParsedNodesMax/2, overlayLayerParsedEdgesMax/2) + results = append(results, result) + return result, nil + } + layer, paths, err := srv.constructOverlayLayer(context.Background(), []daemon.OverlayFile{ + {Path: targetFile, Content: "package main"}, + {Path: callerFile, Content: "package main"}, + }) + require.NoError(t, err) + require.NotNil(t, layer) + require.Len(t, paths, 2) + require.Len(t, results, 2) + for _, result := range results { + require.Nil(t, result.Tree) + require.Nil(t, result.ConstValues) + } + fileCalls, legacyFileReads := store.snapshotFileCalls() + require.Len(t, fileCalls, 2) + require.Zero(t, legacyFileReads, "construction must not call GetFileNodes") + view := graph.NewOverlaidView(store, layer) + require.Len(t, view.GetFileNodes("target.go"), 1) + require.Len(t, view.GetFileNodes("caller.go"), 1) + }) + + t.Run("nodes_plus_one_is_atomic", func(t *testing.T) { + srv, targetFile, callerFile, _, extractor := setupOverlayLayerBoundServer(t) + var first, overflow *parser.ExtractionResult + extractor.extract = func(call int, filePath string, _ []byte) (*parser.ExtractionResult, error) { + count := overlayLayerParsedNodesMax / 2 + if call == 1 { + count++ + } + result := overlayRawExtractionResult(filePath, count, 1) + if call == 0 { + first = result + } else { + overflow = result + } + return result, nil + } + layer, paths, err := srv.constructOverlayLayer(context.Background(), []daemon.OverlayFile{ + {Path: targetFile, Content: "package main"}, + {Path: callerFile, Content: "package main"}, + }) + var limitErr *overlayLayerLimitError + require.ErrorAs(t, err, &limitErr) + require.Equal(t, "parsed nodes", limitErr.Resource) + require.Nil(t, layer) + require.Nil(t, paths) + require.Equal(t, "target.go::node", first.Nodes[0].ID) + require.Equal(t, "target.go::node", first.Edges[0].From) + require.Equal(t, "target.go", first.Nodes[0].Meta["file"]) + require.Nil(t, first.Tree) + require.Nil(t, first.ConstValues) + require.NotNil(t, overflow) + require.Nil(t, overflow.Tree) + require.Nil(t, overflow.ConstValues) + }) + + t.Run("edges_plus_one_is_atomic", func(t *testing.T) { + srv, targetFile, callerFile, _, extractor := setupOverlayLayerBoundServer(t) + var first, overflow *parser.ExtractionResult + extractor.extract = func(call int, filePath string, _ []byte) (*parser.ExtractionResult, error) { + count := overlayLayerParsedEdgesMax / 2 + if call == 1 { + count++ + } + result := overlayRawExtractionResult(filePath, 1, count) + if call == 0 { + first = result + } else { + overflow = result + } + return result, nil + } + layer, paths, err := srv.constructOverlayLayer(context.Background(), []daemon.OverlayFile{ + {Path: targetFile, Content: "package main"}, + {Path: callerFile, Content: "package main"}, + }) + var limitErr *overlayLayerLimitError + require.ErrorAs(t, err, &limitErr) + require.Equal(t, "parsed edges", limitErr.Resource) + require.Nil(t, layer) + require.Nil(t, paths) + require.Equal(t, "target.go::node", first.Nodes[0].ID) + require.Equal(t, "target.go::node", first.Edges[0].From) + require.NotNil(t, overflow) + require.Nil(t, overflow.Tree) + require.Nil(t, overflow.ConstValues) + }) +} + +func TestConstructOverlayLayerReleasesTreeOnEveryExtractorExit(t *testing.T) { + t.Run("result_and_error", func(t *testing.T) { + srv, targetFile, _, store, extractor := setupOverlayLayerBoundServer(t) + sentinel := errors.New("extract failed") + result := overlayRawExtractionResult("target.go", 1, 1) + extractor.extract = func(int, string, []byte) (*parser.ExtractionResult, error) { + return result, sentinel + } + layer, paths, err := srv.constructOverlayLayer(context.Background(), []daemon.OverlayFile{{Path: targetFile, Content: "package main"}}) + require.ErrorIs(t, err, sentinel) + require.Nil(t, layer) + require.Nil(t, paths) + require.Nil(t, result.Tree) + require.Nil(t, result.ConstValues) + fileCalls, _ := store.snapshotFileCalls() + require.Empty(t, fileCalls) + }) + + t.Run("cancel_after_extract", func(t *testing.T) { + srv, targetFile, _, store, extractor := setupOverlayLayerBoundServer(t) + ctx, cancel := context.WithCancel(context.Background()) + result := overlayRawExtractionResult("target.go", 1, 1) + extractor.extract = func(int, string, []byte) (*parser.ExtractionResult, error) { + cancel() + return result, nil + } + layer, paths, err := srv.constructOverlayLayer(ctx, []daemon.OverlayFile{{Path: targetFile, Content: "package main"}}) + require.ErrorIs(t, err, context.Canceled) + require.Nil(t, layer) + require.Nil(t, paths) + require.Nil(t, result.Tree) + require.Nil(t, result.ConstValues) + fileCalls, _ := store.snapshotFileCalls() + require.Empty(t, fileCalls) + }) + + t.Run("pre_canceled_skips_extract", func(t *testing.T) { + srv, targetFile, _, _, extractor := setupOverlayLayerBoundServer(t) + ctx, cancel := context.WithCancel(context.Background()) + cancel() + layer, paths, err := srv.constructOverlayLayer(ctx, []daemon.OverlayFile{{Path: targetFile, Content: "package main"}}) + require.ErrorIs(t, err, context.Canceled) + require.Nil(t, layer) + require.Nil(t, paths) + require.Zero(t, extractor.callCount()) + }) + + t.Run("nil_result_fails_closed", func(t *testing.T) { + srv, targetFile, _, _, extractor := setupOverlayLayerBoundServer(t) + extractor.extract = func(int, string, []byte) (*parser.ExtractionResult, error) { + return nil, nil + } + layer, paths, err := srv.constructOverlayLayer(context.Background(), []daemon.OverlayFile{{Path: targetFile, Content: "package main"}}) + require.ErrorContains(t, err, "extractor returned nil result") + require.Nil(t, layer) + require.Nil(t, paths) + }) +} + +func TestConstructOverlayLayerSnapshotsReusedExtractorPointers(t *testing.T) { + srv, targetFile, callerFile, store, extractor := setupOverlayLayerBoundServer(t) + sharedNode := &graph.Node{Meta: map[string]any{}} + sharedEdge := &graph.Edge{Meta: map[string]any{}} + shared := &parser.ExtractionResult{Nodes: []*graph.Node{sharedNode}, Edges: []*graph.Edge{sharedEdge}} + extractor.extract = func(call int, filePath string, _ []byte) (*parser.ExtractionResult, error) { + sharedNode.ID = filePath + "::node" + sharedNode.Name = fmt.Sprintf("name-%d", call) + sharedNode.FilePath = filePath + sharedNode.Meta["call"] = call + sharedEdge.From = sharedNode.ID + sharedEdge.To = "external::target" + sharedEdge.Meta["call"] = call + return shared, nil + } + + layer, _, err := srv.constructOverlayLayer(context.Background(), []daemon.OverlayFile{ + {Path: targetFile, Content: "package main"}, + {Path: callerFile, Content: "package main"}, + }) + require.NoError(t, err) + view := graph.NewOverlaidView(store, layer) + targetNodes := view.GetFileNodes("target.go") + callerNodes := view.GetFileNodes("caller.go") + require.Len(t, targetNodes, 1) + require.Len(t, callerNodes, 1) + require.Equal(t, "name-0", targetNodes[0].Name) + require.Equal(t, 0, targetNodes[0].Meta["call"]) + require.Equal(t, "name-1", callerNodes[0].Name) + require.Equal(t, 1, callerNodes[0].Meta["call"]) + targetEdges := layer.OutEdgesByFromAll()["repo/target.go::node"] + require.Len(t, targetEdges, 1) + require.Equal(t, 0, targetEdges[0].Meta["call"]) +} + +func TestConstructOverlayLayerRequiresBoundedBaseReader(t *testing.T) { + for _, deleted := range []bool{false, true} { + name := "replacement" + if deleted { + name = "tombstone" + } + t.Run(name, func(t *testing.T) { + srv, targetFile, _, store, extractor := setupOverlayLayerBoundServer(t) + var parsed *parser.ExtractionResult + extractor.extract = func(_ int, filePath string, _ []byte) (*parser.ExtractionResult, error) { + parsed = overlayRawExtractionResult(filePath, 1, 0) + return parsed, nil + } + srv.graph = &overlayLayerLegacyStore{Store: store} + layer, paths, err := srv.constructOverlayLayer(context.Background(), []daemon.OverlayFile{{ + Path: targetFile, Content: "package main", Deleted: deleted, + }}) + require.ErrorIs(t, err, graph.ErrBoundedLocalizationUnavailable) + require.Nil(t, layer) + require.Nil(t, paths) + _, legacyFileReads := store.snapshotFileCalls() + require.Zero(t, legacyFileReads) + if !deleted { + require.NotNil(t, parsed) + require.Nil(t, parsed.Tree) + require.Nil(t, parsed.ConstValues) + } + }) + } +} + +func TestConstructOverlayLayerPreservesTombstoneAndReplacementIdentity(t *testing.T) { + srv, targetFile, _, _, extractor := setupOverlayLayerBoundServer(t) + base := graph.New() + base.AddNode(&graph.Node{ID: "target.go::Old", Name: "Old", FilePath: "target.go"}) + store := &overlayLayerTestStore{Store: base} + srv.graph = store + extractor.extract = func(_ int, filePath string, _ []byte) (*parser.ExtractionResult, error) { + return &parser.ExtractionResult{Nodes: []*graph.Node{{ID: filePath + "::New", Name: "New", FilePath: filePath}}}, nil + } + + layer, _, err := srv.constructOverlayLayer(context.Background(), []daemon.OverlayFile{{Path: targetFile, Content: "package main"}}) + require.NoError(t, err) + _, legacyFileReads := store.snapshotFileCalls() + require.Zero(t, legacyFileReads) + view := graph.NewOverlaidView(store, layer) + require.Empty(t, view.FindNodesByName("Old")) + require.Len(t, view.FindNodesByName("New"), 1) + + layer, _, err = srv.constructOverlayLayer(context.Background(), []daemon.OverlayFile{{Path: targetFile, Deleted: true}}) + require.NoError(t, err) + view = graph.NewOverlaidView(store, layer) + require.True(t, layer.IsTombstone("target.go")) + require.Empty(t, view.GetFileNodes("target.go")) + require.Empty(t, view.FindNodesByName("Old")) +} diff --git a/internal/mcp/overlay_layer_integration_bounds_test.go b/internal/mcp/overlay_layer_integration_bounds_test.go new file mode 100644 index 00000000..9d8ebcf2 --- /dev/null +++ b/internal/mcp/overlay_layer_integration_bounds_test.go @@ -0,0 +1,135 @@ +package mcp + +import ( + "context" + "fmt" + "path/filepath" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/zzet/gortex/internal/daemon" + "github.com/zzet/gortex/internal/graph" + "github.com/zzet/gortex/internal/parser" + "github.com/zzet/gortex/internal/query" + "github.com/zzet/gortex/internal/semantic/lsp" +) + +func TestOverlayLayerFailureIsNotCached(t *testing.T) { + srv, targetFile, _, _, extractor := setupOverlayLayerBoundServer(t) + extractor.extract = func(call int, filePath string, _ []byte) (*parser.ExtractionResult, error) { + if call == 0 { + return overlayRawExtractionResult(filePath, overlayLayerParsedNodesMax+1, 0), nil + } + return overlayRawExtractionResult(filePath, 1, 0), nil + } + ctx := WithSessionID(context.Background(), "cache-session") + ctx = withOverlayRequestSnapshot(ctx, &overlayRequestSnapshot{ + sessionID: "cache-session", + files: []daemon.OverlayFile{{Path: targetFile, Content: "package main"}}, + canonical: true, + }) + + view, err := srv.buildOverlayViewForCtx(ctx) + require.Error(t, err) + require.Nil(t, view) + _, cached := srv.overlayLayerCache.Load("cache-session") + require.False(t, cached) + + view, err = srv.buildOverlayViewForCtx(ctx) + require.NoError(t, err) + require.NotNil(t, view) + _, cached = srv.overlayLayerCache.Load("cache-session") + require.True(t, cached) + require.Equal(t, 2, extractor.callCount()) +} + +func TestOverlayLayerFailureReturnsNoBranchOrSimulationPartial(t *testing.T) { + t.Run("branch", func(t *testing.T) { + srv, targetFile, _, _, extractor := setupOverlayLayerBoundServer(t) + var result *parser.ExtractionResult + extractor.extract = func(_ int, filePath string, _ []byte) (*parser.ExtractionResult, error) { + result = overlayRawExtractionResult(filePath, overlayLayerParsedNodesMax+1, 0) + return result, nil + } + ids, paths, err := srv.runBranchQuery( + context.Background(), + []daemon.OverlayFile{{Path: targetFile, Content: "package main"}}, + "get_callers", + "target.go::Target", + query.QueryOptions{}, + ) + require.Error(t, err) + require.Nil(t, ids) + require.Nil(t, paths) + require.Nil(t, result.Tree) + require.Nil(t, result.ConstValues) + }) + + t.Run("simulation", func(t *testing.T) { + srv, targetFile, _, _, extractor := setupOverlayLayerBoundServer(t) + var result *parser.ExtractionResult + extractor.extract = func(_ int, filePath string, _ []byte) (*parser.ExtractionResult, error) { + result = overlayRawExtractionResult(filePath, 1, overlayLayerParsedEdgesMax+1) + return result, nil + } + sim, err := srv.buildSimulation( + context.Background(), + []lsp.WorkspaceEdit{fullFileSimulationEdit(targetFile, "package main\n")}, + false, + ) + require.Error(t, err) + require.Nil(t, sim) + require.Nil(t, result.Tree) + require.Nil(t, result.ConstValues) + }) +} + +func TestConstructOverlayLayerResolverCancellationReturnsNoLayer(t *testing.T) { + srv, targetFile, _, store, extractor := setupOverlayLayerBoundServer(t) + ctx, cancel := context.WithCancel(context.Background()) + var result *parser.ExtractionResult + extractor.extract = func(_ int, filePath string, _ []byte) (*parser.ExtractionResult, error) { + result = &parser.ExtractionResult{ + Nodes: []*graph.Node{{ID: filePath + "::Caller", Name: "Caller", FilePath: filePath}}, + Edges: []*graph.Edge{{From: filePath + "::Caller", To: unresolvedPrefix + "call::Missing"}}, + Tree: parser.NewParseTree(nil, []byte(filePath), "go"), + } + return result, nil + } + store.nameFn = func( + context.Context, string, graph.LocalizationNodeScope, int, + ) (graph.BoundedNodeProjection, error) { + cancel() + return graph.BoundedNodeProjection{}, nil + } + + layer, paths, err := srv.constructOverlayLayer(ctx, []daemon.OverlayFile{{Path: targetFile, Content: "package main"}}) + require.ErrorIs(t, err, context.Canceled) + require.Nil(t, layer) + require.Nil(t, paths) + require.Equal(t, "target.go::Caller", result.Nodes[0].ID) + require.Equal(t, unresolvedPrefix+"call::Missing", result.Edges[0].To) + require.Nil(t, result.Tree) +} + +func TestConstructOverlayLayerAcceptsFullLiteralFileEnvelope(t *testing.T) { + srv, targetFile, _, store, extractor := setupOverlayLayerBoundServer(t) + extractor.extract = func(_ int, _ string, _ []byte) (*parser.ExtractionResult, error) { + return &parser.ExtractionResult{}, nil + } + files := make([]daemon.OverlayFile, overlayRequestSnapshotMaxFiles) + for index := range files { + files[index] = daemon.OverlayFile{ + Path: filepath.Join(filepath.Dir(targetFile), fmt.Sprintf("literal-%03d.go", index)), + Content: "package main", + } + } + layer, paths, err := srv.constructOverlayLayer(context.Background(), files) + require.NoError(t, err) + require.NotNil(t, layer) + require.Len(t, paths, overlayRequestSnapshotMaxFiles) + fileCalls, legacyFileReads := store.snapshotFileCalls() + require.Len(t, fileCalls, overlayRequestSnapshotMaxFiles) + require.Zero(t, legacyFileReads) +} diff --git a/internal/mcp/overlay_layer_resolver_bounds_test.go b/internal/mcp/overlay_layer_resolver_bounds_test.go new file mode 100644 index 00000000..cf2030be --- /dev/null +++ b/internal/mcp/overlay_layer_resolver_bounds_test.go @@ -0,0 +1,238 @@ +package mcp + +import ( + "context" + "errors" + "fmt" + "sort" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/zzet/gortex/internal/graph" +) + +func overlayLayerWithUnresolved(names ...string) (*graph.OverlayLayer, []string) { + layer := graph.NewOverlayLayer() + layer.MarkFile("overlay.go", false) + fromIDs := make([]string, 0, len(names)) + for index, name := range names { + from := fmt.Sprintf("overlay.go::caller-%d", index) + fromIDs = append(fromIDs, from) + layer.AddEdge(&graph.Edge{From: from, To: unresolvedPrefix + "call::" + name}) + } + return layer, fromIDs +} + +func TestResolveOverlayEdgesDedupesSortsAndRebuildsInbound(t *testing.T) { + layer, fromIDs := overlayLayerWithUnresolved("Beta", "Alpha", "Alpha") + store := &overlayLayerTestStore{Store: graph.New()} + store.nameFn = func( + _ context.Context, + name string, + _ graph.LocalizationNodeScope, + limit int, + ) (graph.BoundedNodeProjection, error) { + require.Equal(t, 1, limit) + return graph.BoundedNodeProjection{ + Nodes: []*graph.Node{{ID: "base.go::" + name, Name: name, FilePath: "base.go"}}, + Total: 1, + }, nil + } + + require.NoError(t, resolveOverlayEdges(context.Background(), store, layer)) + calls := store.snapshotNameCalls() + _, legacyNameReads := store.snapshotLegacyReads() + require.Zero(t, legacyNameReads, "resolver must never fall back to FindNodesByName") + require.Equal(t, []string{"Alpha", "Beta"}, []string{calls[0].name, calls[1].name}) + for _, call := range calls { + require.Equal(t, 1, call.limit) + } + + view := graph.NewOverlaidView(store, layer) + require.Equal(t, "base.go::Beta", view.GetOutEdges(fromIDs[0])[0].To) + require.Equal(t, "base.go::Alpha", view.GetOutEdges(fromIDs[1])[0].To) + require.Equal(t, "base.go::Alpha", view.GetOutEdges(fromIDs[2])[0].To) + require.Len(t, view.GetInEdges("base.go::Alpha"), 2) +} + +func TestResolveOverlayEdgesDistinctNameBounds(t *testing.T) { + names := make([]string, overlayLayerUnresolvedNamesMax) + for index := range names { + names[index] = fmt.Sprintf("Name%04d", index) + } + layer, _ := overlayLayerWithUnresolved(names...) + store := &overlayLayerTestStore{Store: graph.New()} + store.nameFn = func( + context.Context, string, graph.LocalizationNodeScope, int, + ) (graph.BoundedNodeProjection, error) { + return graph.BoundedNodeProjection{}, nil + } + require.NoError(t, resolveOverlayEdges(context.Background(), store, layer)) + calls := store.snapshotNameCalls() + require.Len(t, calls, overlayLayerUnresolvedNamesMax) + ordered := make([]string, 0, len(calls)) + for _, call := range calls { + ordered = append(ordered, call.name) + } + require.True(t, sort.StringsAreSorted(ordered)) + + names = append(names, "Name-overflow") + layer, _ = overlayLayerWithUnresolved(names...) + store = &overlayLayerTestStore{Store: graph.New()} + _, called := store.nameFn, false + store.nameFn = func( + context.Context, string, graph.LocalizationNodeScope, int, + ) (graph.BoundedNodeProjection, error) { + called = true + return graph.BoundedNodeProjection{}, nil + } + err := resolveOverlayEdges(context.Background(), store, layer) + var limitErr *overlayLayerLimitError + require.ErrorAs(t, err, &limitErr) + require.Equal(t, "unresolved names", limitErr.Resource) + require.False(t, called, "the complete distinct-name set must be capped before querying base") +} + +func TestResolveOverlayEdgesOverlayPrecedence(t *testing.T) { + t.Run("unique", func(t *testing.T) { + layer, fromIDs := overlayLayerWithUnresolved("Target") + layer.AddNode("overlay.go", &graph.Node{ID: "overlay.go::Target", Name: "Target", FilePath: "overlay.go"}) + store := &overlayLayerTestStore{Store: graph.New()} + store.nameFn = func( + context.Context, string, graph.LocalizationNodeScope, int, + ) (graph.BoundedNodeProjection, error) { + return graph.BoundedNodeProjection{}, errors.New("base must not be queried") + } + require.NoError(t, resolveOverlayEdges(context.Background(), store, layer)) + require.Empty(t, store.snapshotNameCalls()) + require.Equal(t, "overlay.go::Target", graph.NewOverlaidView(store, layer).GetOutEdges(fromIDs[0])[0].To) + }) + + t.Run("ambiguous", func(t *testing.T) { + layer, fromIDs := overlayLayerWithUnresolved("Target") + layer.AddNode("overlay.go", &graph.Node{ID: "overlay.go::Target-one", Name: "Target", FilePath: "overlay.go"}) + layer.AddNode("overlay.go", &graph.Node{ID: "overlay.go::Target-two", Name: "Target", FilePath: "overlay.go"}) + store := &overlayLayerTestStore{Store: graph.New()} + store.nameFn = func( + context.Context, string, graph.LocalizationNodeScope, int, + ) (graph.BoundedNodeProjection, error) { + return graph.BoundedNodeProjection{}, errors.New("base must not be queried") + } + require.NoError(t, resolveOverlayEdges(context.Background(), store, layer)) + require.Empty(t, store.snapshotNameCalls()) + require.Equal(t, unresolvedPrefix+"call::Target", graph.NewOverlaidView(store, layer).GetOutEdges(fromIDs[0])[0].To) + }) +} + +func TestResolveOverlayEdgesBaseCardinality(t *testing.T) { + for _, test := range []struct { + name string + projection graph.BoundedNodeProjection + wantTarget string + }{ + {name: "zero", projection: graph.BoundedNodeProjection{}, wantTarget: unresolvedPrefix + "call::Target"}, + { + name: "one", + projection: graph.BoundedNodeProjection{ + Nodes: []*graph.Node{{ID: "base.go::Target", Name: "Target", FilePath: "base.go"}}, + Total: 1, + }, + wantTarget: "base.go::Target", + }, + { + name: "ambiguous", + projection: graph.BoundedNodeProjection{ + Nodes: []*graph.Node{{ID: "base.go::Target", Name: "Target", FilePath: "base.go"}}, + Total: 2, + Truncated: true, + }, + wantTarget: unresolvedPrefix + "call::Target", + }, + } { + t.Run(test.name, func(t *testing.T) { + layer, fromIDs := overlayLayerWithUnresolved("Target") + store := &overlayLayerTestStore{Store: graph.New()} + store.nameFn = func( + context.Context, string, graph.LocalizationNodeScope, int, + ) (graph.BoundedNodeProjection, error) { + return test.projection, nil + } + require.NoError(t, resolveOverlayEdges(context.Background(), store, layer)) + require.Equal(t, test.wantTarget, graph.NewOverlaidView(store, layer).GetOutEdges(fromIDs[0])[0].To) + }) + } +} + +func TestResolveOverlayEdgesFiltersCoveredBaseFiles(t *testing.T) { + for _, tombstone := range []bool{false, true} { + name := "replacement" + if tombstone { + name = "tombstone" + } + t.Run(name, func(t *testing.T) { + layer := graph.NewOverlayLayer() + layer.MarkFile("base.go", tombstone) + layer.MarkFile("overlay.go", false) + from := "overlay.go::caller" + layer.AddEdge(&graph.Edge{From: from, To: unresolvedPrefix + "call::Stale"}) + store := &overlayLayerTestStore{Store: graph.New()} + store.nameFn = func( + _ context.Context, + _ string, + scope graph.LocalizationNodeScope, + _ int, + ) (graph.BoundedNodeProjection, error) { + require.True(t, scope.ExcludesFile("base.go")) + return graph.BoundedNodeProjection{}, nil + } + require.NoError(t, resolveOverlayEdges(context.Background(), store, layer)) + require.Equal(t, unresolvedPrefix+"call::Stale", graph.NewOverlaidView(store, layer).GetOutEdges(from)[0].To) + }) + } +} + +func TestResolveOverlayEdgesFailsClosedWithoutPartialRewrite(t *testing.T) { + sentinel := errors.New("query failed") + layer, fromIDs := overlayLayerWithUnresolved("Alpha", "Beta") + store := &overlayLayerTestStore{Store: graph.New()} + store.nameFn = func( + _ context.Context, + name string, + _ graph.LocalizationNodeScope, + _ int, + ) (graph.BoundedNodeProjection, error) { + if name == "Beta" { + return graph.BoundedNodeProjection{}, sentinel + } + return graph.BoundedNodeProjection{ + Nodes: []*graph.Node{{ID: "base.go::Alpha", Name: "Alpha", FilePath: "base.go"}}, + Total: 1, + }, nil + } + err := resolveOverlayEdges(context.Background(), store, layer) + require.ErrorIs(t, err, sentinel) + view := graph.NewOverlaidView(store, layer) + require.Equal(t, unresolvedPrefix+"call::Alpha", view.GetOutEdges(fromIDs[0])[0].To) + require.Equal(t, unresolvedPrefix+"call::Beta", view.GetOutEdges(fromIDs[1])[0].To) + + legacy := &overlayLayerLegacyStore{Store: graph.New()} + unsupportedLayer, _ := overlayLayerWithUnresolved("Target") + err = resolveOverlayEdges(context.Background(), legacy, unsupportedLayer) + require.ErrorIs(t, err, graph.ErrBoundedLocalizationUnavailable) +} + +func TestResolveOverlayEdgesPropagatesCancellation(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + layer, fromIDs := overlayLayerWithUnresolved("Target") + store := &overlayLayerTestStore{Store: graph.New()} + store.nameFn = func( + context.Context, string, graph.LocalizationNodeScope, int, + ) (graph.BoundedNodeProjection, error) { + cancel() + return graph.BoundedNodeProjection{}, nil + } + err := resolveOverlayEdges(ctx, store, layer) + require.ErrorIs(t, err, context.Canceled) + require.Equal(t, unresolvedPrefix+"call::Target", graph.NewOverlaidView(store, layer).GetOutEdges(fromIDs[0])[0].To) +} diff --git a/internal/mcp/overlay_layer_test_helpers_test.go b/internal/mcp/overlay_layer_test_helpers_test.go new file mode 100644 index 00000000..a97f7d40 --- /dev/null +++ b/internal/mcp/overlay_layer_test_helpers_test.go @@ -0,0 +1,136 @@ +package mcp + +import ( + "context" + "sync" + + "github.com/zzet/gortex/internal/graph" + "github.com/zzet/gortex/internal/parser" +) + +type overlayLayerFileCall struct { + path string + scope graph.LocalizationNodeScope + limit int +} + +type overlayLayerNameCall struct { + name string + scope graph.LocalizationNodeScope + limit int +} + +type overlayLayerTestStore struct { + graph.Store + + mu sync.Mutex + fileCalls []overlayLayerFileCall + nameCalls []overlayLayerNameCall + legacyFileReads int + legacyNameReads int + fileFn func(context.Context, string, graph.LocalizationNodeScope, int) (graph.BoundedNodeProjection, error) + nameFn func(context.Context, string, graph.LocalizationNodeScope, int) (graph.BoundedNodeProjection, error) +} + +func (s *overlayLayerTestStore) FindFileNodesBounded( + ctx context.Context, + filePath string, + scope graph.LocalizationNodeScope, + limit int, +) (graph.BoundedNodeProjection, error) { + s.mu.Lock() + s.fileCalls = append(s.fileCalls, overlayLayerFileCall{path: filePath, scope: scope, limit: limit}) + fn := s.fileFn + s.mu.Unlock() + if fn != nil { + return fn(ctx, filePath, scope, limit) + } + reader, ok := s.Store.(graph.BoundedFileNodeReader) + if !ok { + return graph.BoundedNodeProjection{}, graph.ErrBoundedLocalizationUnavailable + } + return reader.FindFileNodesBounded(ctx, filePath, scope, limit) +} + +func (s *overlayLayerTestStore) FindNodesByNameBounded( + ctx context.Context, + name string, + scope graph.LocalizationNodeScope, + limit int, +) (graph.BoundedNodeProjection, error) { + s.mu.Lock() + s.nameCalls = append(s.nameCalls, overlayLayerNameCall{name: name, scope: scope, limit: limit}) + fn := s.nameFn + s.mu.Unlock() + if fn != nil { + return fn(ctx, name, scope, limit) + } + reader, ok := s.Store.(graph.BoundedExactNameReader) + if !ok { + return graph.BoundedNodeProjection{}, graph.ErrBoundedLocalizationUnavailable + } + return reader.FindNodesByNameBounded(ctx, name, scope, limit) +} + +func (s *overlayLayerTestStore) GetFileNodes(filePath string) []*graph.Node { + s.mu.Lock() + s.legacyFileReads++ + s.mu.Unlock() + return s.Store.GetFileNodes(filePath) +} + +func (s *overlayLayerTestStore) FindNodesByName(name string) []*graph.Node { + s.mu.Lock() + s.legacyNameReads++ + s.mu.Unlock() + return s.Store.FindNodesByName(name) +} + +func (s *overlayLayerTestStore) snapshotFileCalls() ([]overlayLayerFileCall, int) { + s.mu.Lock() + defer s.mu.Unlock() + return append([]overlayLayerFileCall(nil), s.fileCalls...), s.legacyFileReads +} + +func (s *overlayLayerTestStore) snapshotNameCalls() []overlayLayerNameCall { + s.mu.Lock() + defer s.mu.Unlock() + return append([]overlayLayerNameCall(nil), s.nameCalls...) +} + +func (s *overlayLayerTestStore) snapshotLegacyReads() (fileReads, nameReads int) { + s.mu.Lock() + defer s.mu.Unlock() + return s.legacyFileReads, s.legacyNameReads +} + +type overlayLayerLegacyStore struct { + graph.Store +} + +type overlayLayerTestExtractor struct { + mu sync.Mutex + calls int + extract func(int, string, []byte) (*parser.ExtractionResult, error) +} + +func (e *overlayLayerTestExtractor) Language() string { return "go" } +func (e *overlayLayerTestExtractor) Extensions() []string { return []string{".go"} } + +func (e *overlayLayerTestExtractor) Extract(filePath string, content []byte) (*parser.ExtractionResult, error) { + e.mu.Lock() + call := e.calls + e.calls++ + fn := e.extract + e.mu.Unlock() + if fn == nil { + return &parser.ExtractionResult{}, nil + } + return fn(call, filePath, content) +} + +func (e *overlayLayerTestExtractor) callCount() int { + e.mu.Lock() + defer e.mu.Unlock() + return e.calls +} diff --git a/internal/mcp/overlay_localization_test.go b/internal/mcp/overlay_localization_test.go new file mode 100644 index 00000000..bc39e030 --- /dev/null +++ b/internal/mcp/overlay_localization_test.go @@ -0,0 +1,90 @@ +package mcp + +import ( + "context" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/zzet/gortex/internal/graph" + "github.com/zzet/gortex/internal/query" + "github.com/zzet/gortex/internal/search/trigram" +) + +func TestOverlayLocalizationFiltersShadowedDurableContent(t *testing.T) { + base := graph.New() + for _, deleted := range []bool{false, true} { + layer := graph.NewOverlayLayer() + layer.MarkFile("shadow.go", deleted) + ctx := WithOverlayView(context.Background(), graph.NewOverlaidView(base, layer)) + srv := &Server{graph: base} + hits := []graph.ContentHit{ + {NodeID: "shadow.go::stale", FilePath: "shadow.go", Snippet: "stale needle"}, + {NodeID: "live.go::fresh", FilePath: "live.go", Snippet: "fresh needle"}, + } + filtered := srv.filterOverlayContentHits(ctx, hits) + require.Equal(t, []graph.ContentHit{hits[1]}, filtered) + } +} + +func TestOverlayLocalizationMapsReplacementOwnerAndAdjacency(t *testing.T) { + base := graph.New() + baseOwner := &graph.Node{ID: "sample.go::Caller", Kind: graph.KindFunction, Name: "Caller", FilePath: "sample.go", StartLine: 1, EndLine: 5} + staleCallee := &graph.Node{ID: "sample.go::Stale", Kind: graph.KindFunction, Name: "Stale", FilePath: "sample.go", StartLine: 7, EndLine: 9} + base.AddBatch([]*graph.Node{baseOwner, staleCallee}, []*graph.Edge{{From: baseOwner.ID, To: staleCallee.ID, Kind: graph.EdgeCalls, Line: 3}}) + + layer := graph.NewOverlayLayer() + layer.MarkFile("sample.go", false) + overlayOwner := &graph.Node{ID: baseOwner.ID, Kind: graph.KindFunction, Name: "Caller", FilePath: "sample.go", StartLine: 1, EndLine: 5} + freshCallee := &graph.Node{ID: "sample.go::Fresh", Kind: graph.KindFunction, Name: "Fresh", FilePath: "sample.go", StartLine: 7, EndLine: 9} + layer.AddNode("sample.go", overlayOwner) + layer.AddNode("sample.go", freshCallee) + layer.AddEdge(&graph.Edge{From: overlayOwner.ID, To: freshCallee.ID, Kind: graph.EdgeCalls, Line: 3}) + view := graph.NewOverlaidView(base, layer) + ctx := WithOverlayView(context.Background(), view) + + srv := &Server{graph: base} + recall := srv.mapExploreSourceLiteralMatchesContext(ctx, "needle", []trigram.Match{{ + Path: "sample.go", Line: 3, Text: `Fresh("needle")`, + }}, query.QueryOptions{}) + require.Len(t, recall.hits, 1) + require.Equal(t, freshCallee.ID, recall.hits[0].nodeID) + require.True(t, recall.hits[0].callee) +} + +func TestOverlayLocalizationTombstoneSuppressesSourceOwner(t *testing.T) { + base := graph.New() + owner := &graph.Node{ID: "sample.go::Caller", Kind: graph.KindFunction, Name: "Caller", FilePath: "sample.go", StartLine: 1, EndLine: 5} + base.AddNode(owner) + layer := graph.NewOverlayLayer() + layer.MarkFile("sample.go", true) + ctx := WithOverlayView(context.Background(), graph.NewOverlaidView(base, layer)) + + srv := &Server{graph: base} + recall := srv.mapExploreSourceLiteralMatchesContext(ctx, "needle", []trigram.Match{{ + Path: "sample.go", Line: 3, Text: `Call("needle")`, + }}, query.QueryOptions{}) + require.Empty(t, recall.hits) +} + +func TestOverlayLocalizationQualifiedOwnerUsesReplacement(t *testing.T) { + base := graph.New() + staleOwner := &graph.Node{ID: "sample.go::Owner", Kind: graph.KindType, Name: "Owner", FilePath: "sample.go", StartLine: 1, EndLine: 8} + base.AddNode(staleOwner) + + layer := graph.NewOverlayLayer() + layer.MarkFile("sample.go", false) + layer.MarkRemoved(staleOwner.Name, staleOwner.ID) + freshOwner := &graph.Node{ID: "sample.go::OwnerV2", Kind: graph.KindType, Name: "Owner", FilePath: "sample.go", StartLine: 1, EndLine: 10} + member := &graph.Node{ID: "sample.go::Owner.run", Kind: graph.KindMethod, Name: "run", FilePath: "sample.go", StartLine: 3, EndLine: 5} + layer.AddNode("sample.go", freshOwner) + layer.AddNode("sample.go", member) + ctx := WithOverlayView(context.Background(), graph.NewOverlaidView(base, layer)) + + srv := &Server{graph: base} + candidate := srv.exploreQualifiedAnchorOwnerCandidate(ctx, exploreSyntacticAnchor{ + qualifiedName: "Owner.run", + }, member, query.QueryOptions{}, map[string]struct{}{}) + require.NotNil(t, candidate) + require.Equal(t, freshOwner.ID, candidate.Node.ID) +} diff --git a/internal/mcp/overlay_snapshot_bounds_test.go b/internal/mcp/overlay_snapshot_bounds_test.go new file mode 100644 index 00000000..5c530566 --- /dev/null +++ b/internal/mcp/overlay_snapshot_bounds_test.go @@ -0,0 +1,483 @@ +package mcp + +import ( + "context" + "errors" + "strings" + "sync/atomic" + "testing" + "time" + + mcplib "github.com/mark3labs/mcp-go/mcp" + "github.com/stretchr/testify/require" + + "github.com/zzet/gortex/internal/daemon" + "github.com/zzet/gortex/internal/semantic/lsp" +) + +func TestOverlayRequestSnapshotLimitsMatchLiteralEnvelope(t *testing.T) { + require.Equal(t, 257, overlayRequestSnapshotMaxFiles) + require.Equal(t, 4_259_840, overlayRequestSnapshotMaxBytes) +} + +func TestSnapshotOverlayRequestForCtxBoundsRawAliasesBeforeCanonicalization(t *testing.T) { + for _, tc := range []struct { + name string + fileCount int + wantError bool + }{ + {name: "exact", fileCount: overlayRequestSnapshotMaxFiles}, + {name: "plus_one", fileCount: overlayRequestSnapshotMaxFiles + 1, wantError: true}, + } { + t.Run(tc.name, func(t *testing.T) { + manager := daemon.NewOverlayManager(time.Hour) + const sessionID = "raw-alias-boundary" + require.NoError(t, manager.RegisterWithID(sessionID, "workspace")) + pushRawOverlayAliases(t, manager, sessionID, tc.fileCount) + + server := &Server{overlays: manager} + snapshot, err := server.snapshotOverlayRequestForCtx(WithSessionID(context.Background(), sessionID)) + if tc.wantError { + require.Nil(t, snapshot) + requireOverlaySnapshotBoundError(t, err, "files", overlayRequestSnapshotMaxFiles) + return + } + require.NoError(t, err) + require.NotNil(t, snapshot) + require.Len(t, snapshot.files, overlayRequestSnapshotMaxFiles) + require.False(t, snapshot.canonical) + }) + } +} + +func TestSnapshotOverlayRequestForCtxPropagatesByteBoundary(t *testing.T) { + manager := daemon.NewOverlayManager(time.Hour) + const sessionID = "raw-byte-boundary" + require.NoError(t, manager.RegisterWithID(sessionID, "workspace")) + server := &Server{overlays: manager} + ctx := WithSessionID(context.Background(), sessionID) + + exact := daemon.OverlayFile{ + Path: "p", + Content: strings.Repeat("x", overlayRequestSnapshotMaxBytes-2), + BaseSHA: "s", + } + require.NoError(t, manager.Push(sessionID, exact, nil)) + snapshot, err := server.snapshotOverlayRequestForCtx(ctx) + require.NoError(t, err) + require.Len(t, snapshot.files, 1) + + exact.Content += "x" + require.NoError(t, manager.Push(sessionID, exact, nil)) + snapshot, err = server.snapshotOverlayRequestForCtx(ctx) + require.Nil(t, snapshot) + requireOverlaySnapshotBoundError(t, err, "bytes", overlayRequestSnapshotMaxBytes) +} + +func TestSnapshotOverlayRequestForCtxPinsOnlyMissingSession(t *testing.T) { + server := &Server{overlays: daemon.NewOverlayManager(time.Hour)} + const sessionID = "missing-overlay-session" + + snapshot, err := server.snapshotOverlayRequestForCtx(WithSessionID(context.Background(), sessionID)) + require.NoError(t, err) + require.NotNil(t, snapshot) + require.Equal(t, sessionID, snapshot.sessionID) + require.Empty(t, snapshot.workspace) + require.Nil(t, snapshot.files) +} + +func TestSnapshotOverlayRequestForCtxChecksCancellationAroundAcquisition(t *testing.T) { + manager := daemon.NewOverlayManager(time.Hour) + const sessionID = "snapshot-cancellation" + require.NoError(t, manager.RegisterWithID(sessionID, "workspace")) + server := &Server{overlays: manager} + + preCanceled, cancel := context.WithCancel(WithSessionID(context.Background(), sessionID)) + cancel() + snapshot, err := server.snapshotOverlayRequestForCtx(preCanceled) + require.Nil(t, snapshot) + require.ErrorIs(t, err, context.Canceled) + + postCanceled := &cancelAfterNErrContext{ + Context: WithSessionID(context.Background(), sessionID), + allowedCalls: 1, + } + snapshot, err = server.snapshotOverlayRequestForCtx(postCanceled) + require.Nil(t, snapshot) + require.ErrorIs(t, err, context.Canceled) + require.Equal(t, int32(2), postCanceled.errCalls.Load()) +} + +func TestValidateOverlayBuildEnvelopeExactAndPlusOne(t *testing.T) { + exact := daemon.OverlayFile{ + Path: "p", + Content: strings.Repeat("x", overlayRequestSnapshotMaxBytes-2), + BaseSHA: "s", + } + require.NoError(t, validateOverlayBuildEnvelope([]daemon.OverlayFile{exact})) + require.NoError(t, validateOverlayBuildMapEnvelope(map[string]daemon.OverlayFile{"p": exact})) + + server := &Server{} + layer, paths, err := server.constructOverlayLayer(context.Background(), []daemon.OverlayFile{exact}) + require.NoError(t, err) + require.Nil(t, layer) + require.Nil(t, paths) + + tooLarge := exact + tooLarge.BaseSHA += "s" + requireOverlaySnapshotBoundError( + t, + validateOverlayBuildEnvelope([]daemon.OverlayFile{tooLarge}), + "bytes", + overlayRequestSnapshotMaxBytes, + ) + requireOverlaySnapshotBoundError( + t, + validateOverlayBuildMapEnvelope(map[string]daemon.OverlayFile{"p": tooLarge}), + "bytes", + overlayRequestSnapshotMaxBytes, + ) + layer, paths, err = server.constructOverlayLayer(context.Background(), []daemon.OverlayFile{tooLarge}) + require.Nil(t, layer) + require.Nil(t, paths) + requireOverlaySnapshotBoundError(t, err, "bytes", overlayRequestSnapshotMaxBytes) + + tooMany := make([]daemon.OverlayFile, overlayRequestSnapshotMaxFiles+1) + requireOverlaySnapshotBoundError(t, validateOverlayBuildEnvelope(tooMany), "files", overlayRequestSnapshotMaxFiles) + tooManyMap := make(map[string]daemon.OverlayFile, overlayRequestSnapshotMaxFiles+1) + for i := range tooMany { + path := strings.Repeat("x", i+1) + tooManyMap[path] = daemon.OverlayFile{Path: path} + } + requireOverlaySnapshotBoundError(t, validateOverlayBuildMapEnvelope(tooManyMap), "files", overlayRequestSnapshotMaxFiles) +} + +func TestConstructOverlayLayerCancellationPrecedesRawValidation(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + files := make([]daemon.OverlayFile, overlayRequestSnapshotMaxFiles+1) + + layer, paths, err := (&Server{}).constructOverlayLayer(ctx, files) + require.Nil(t, layer) + require.Nil(t, paths) + require.ErrorIs(t, err, context.Canceled) +} + +func TestBuildSimulationBoundsInheritedSnapshotAndPreservesPinnedEmpty(t *testing.T) { + t.Run("oversized inheritance", func(t *testing.T) { + manager := daemon.NewOverlayManager(time.Hour) + const sessionID = "oversized-inheritance" + require.NoError(t, manager.RegisterWithID(sessionID, "workspace")) + pushRawOverlayAliases(t, manager, sessionID, overlayRequestSnapshotMaxFiles+1) + + sim, err := (&Server{overlays: manager}).buildSimulation( + WithSessionID(context.Background(), sessionID), + nil, + true, + ) + require.Nil(t, sim) + requireOverlaySnapshotBoundError(t, err, "files", overlayRequestSnapshotMaxFiles) + }) + + t.Run("empty inheritance", func(t *testing.T) { + manager := daemon.NewOverlayManager(time.Hour) + const sessionID = "empty-inheritance" + require.NoError(t, manager.RegisterWithID(sessionID, "workspace")) + + sim, err := (&Server{overlays: manager}).buildSimulation( + WithSessionID(context.Background(), sessionID), + nil, + true, + ) + require.NoError(t, err) + require.NotNil(t, sim) + require.Empty(t, sim.initial) + require.Empty(t, sim.snapshots) + }) +} + +func TestBuildSimulationPreflightsStepBeforeRetainingSnapshot(t *testing.T) { + path := t.TempDir() + "/new.go" + for _, tc := range []struct { + name string + extraByte int + wantError bool + }{ + {name: "exact"}, + {name: "plus_one", extraByte: 1, wantError: true}, + } { + t.Run(tc.name, func(t *testing.T) { + contentBytes := overlayRequestSnapshotMaxBytes - len(path) + tc.extraByte + edit := fullFileSimulationEdit(path, strings.Repeat("x", contentBytes)) + + sim, err := (&Server{}).buildSimulation(context.Background(), []lsp.WorkspaceEdit{edit}, false) + if tc.wantError { + require.Nil(t, sim) + requireOverlaySnapshotBoundError(t, err, "bytes", overlayRequestSnapshotMaxBytes) + return + } + require.NoError(t, err) + require.NotNil(t, sim) + require.Len(t, sim.snapshots, 1) + require.Len(t, sim.snapshots[0], 1) + }) + } +} + +func TestSimulationStepLimitExactAndPlusOne(t *testing.T) { + exact := make([]lsp.WorkspaceEdit, overlaySimulationMaxSteps) + sim, err := (&Server{}).buildSimulation(context.Background(), exact, false) + require.NoError(t, err) + require.NotNil(t, sim) + require.Len(t, sim.steps, overlaySimulationMaxSteps) + require.Len(t, sim.snapshots, overlaySimulationMaxSteps) + + sim, err = (&Server{}).buildSimulation( + context.Background(), + make([]lsp.WorkspaceEdit, overlaySimulationMaxSteps+1), + false, + ) + require.Nil(t, sim) + require.ErrorContains(t, err, "simulation steps exceed limit 16") +} + +func TestSimulationRawInputExactAndPlusOne(t *testing.T) { + exactEdit := "{}" + strings.Repeat(" ", overlaySimulationInputMaxBytes-2) + _, err := parseWorkspaceEdit(exactEdit) + require.NoError(t, err) + _, err = parseWorkspaceEdit(exactEdit + " ") + require.ErrorContains(t, err, "workspace_edit input exceeds limit") + + exactArray := "[{}]" + strings.Repeat(" ", overlaySimulationInputMaxBytes-4) + _, err = parsePRReviewEdits(exactArray) + require.Error(t, err) + require.NotContains(t, err.Error(), "input exceeds limit") + _, err = parsePRReviewEdits(exactArray + " ") + require.ErrorContains(t, err, "edits input exceeds limit") + + server, _, _, _ := setupOverlayServer(t) + req := mcplib.CallToolRequest{} + req.Params.Arguments = map[string]any{"steps": exactArray} + result, callErr := server.handleSimulateChain(context.Background(), req) + require.NoError(t, callErr) + require.NotNil(t, result) + require.True(t, result.IsError) + require.NotContains(t, toolText(result), "input exceeds limit") + + req.Params.Arguments = map[string]any{"steps": exactArray + " "} + result, callErr = server.handleSimulateChain(context.Background(), req) + require.NoError(t, callErr) + require.NotNil(t, result) + require.True(t, result.IsError) + require.Contains(t, toolText(result), "steps input exceeds limit") +} + +func TestSimulationArrayStepLimitCheckedBeforeElementParsing(t *testing.T) { + exact := rawEmptyWorkspaceEditArray(overlaySimulationMaxSteps) + _, err := parsePRReviewEdits(exact) + require.Error(t, err) + require.NotContains(t, err.Error(), "exceed limit") + + _, err = parsePRReviewEdits(rawEmptyWorkspaceEditArray(overlaySimulationMaxSteps + 1)) + require.ErrorContains(t, err, "edits exceed limit 16") + + server, _, _, _ := setupOverlayServer(t) + req := mcplib.CallToolRequest{} + req.Params.Arguments = map[string]any{"steps": exact} + result, callErr := server.handleSimulateChain(context.Background(), req) + require.NoError(t, callErr) + require.NotContains(t, toolText(result), "steps exceed limit") + + req.Params.Arguments = map[string]any{"steps": rawEmptyWorkspaceEditArray(overlaySimulationMaxSteps + 1)} + result, callErr = server.handleSimulateChain(context.Background(), req) + require.NoError(t, callErr) + require.Contains(t, toolText(result), "steps exceed limit 16") +} + +func TestCompareBranchesPropagatesCancellationDuringConstruction(t *testing.T) { + server, sessionID, _, targetFile, _, baseCtx := branchBootstrap(t) + for _, branch := range []string{"a", "b"} { + _, err := server.OverlayManager().Fork(sessionID, daemon.ForkOptions{Name: branch}) + require.NoError(t, err) + require.NoError(t, server.OverlayManager().PushToBranch( + sessionID, + branch, + daemon.OverlayFile{Path: targetFile, Content: "package main\nfunc Target() {}\n"}, + nil, + )) + } + + ctx := &cancelAfterNErrContext{Context: baseCtx, allowedCalls: 4} + req := mcplib.CallToolRequest{} + req.Params.Arguments = map[string]any{ + "a": "a", "b": "b", "kind": "get_callers", "id": "target.go::Target", + } + result, err := server.handleCompareBranches(ctx, req) + require.Nil(t, result) + require.ErrorIs(t, err, context.Canceled) + require.GreaterOrEqual(t, ctx.errCalls.Load(), int32(5)) +} + +func TestRequestContextErrorUsesOnlyParentRequestState(t *testing.T) { + componentErr := errors.Join(errors.New("extractor timeout"), context.DeadlineExceeded) + require.NoError(t, requestContextError(context.Background(), componentErr)) + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + require.ErrorIs(t, requestContextError(ctx, errors.New("ordinary failure")), context.Canceled) +} + +func TestSimulationPublicHandlersPropagateCancellation(t *testing.T) { + server, _, targetFile, _ := setupOverlayServer(t) + rawEdit := buildSingleFileEdit(targetFile, "package main\nfunc Target() {}\n") + + for _, tc := range []struct { + name string + call func(context.Context) (*mcplib.CallToolResult, error) + }{ + { + name: "preview edit", + call: func(ctx context.Context) (*mcplib.CallToolResult, error) { + req := mcplib.CallToolRequest{} + req.Params.Arguments = map[string]any{"workspace_edit": rawEdit, "diagnostics": false} + return server.handlePreviewEdit(ctx, req) + }, + }, + { + name: "simulate chain", + call: func(ctx context.Context) (*mcplib.CallToolResult, error) { + req := mcplib.CallToolRequest{} + req.Params.Arguments = map[string]any{"steps": "[" + rawEdit + "]", "diagnostics": false} + return server.handleSimulateChain(ctx, req) + }, + }, + { + name: "change contract", + call: func(ctx context.Context) (*mcplib.CallToolResult, error) { + req := mcplib.CallToolRequest{} + req.Params.Arguments = map[string]any{"source": "edit", "workspace_edit": rawEdit} + return server.handleChangeContract(ctx, req) + }, + }, + } { + t.Run(tc.name, func(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + result, err := tc.call(ctx) + require.Nil(t, result) + require.ErrorIs(t, err, context.Canceled) + }) + } +} + +func TestCompareWithOverlayPropagatesPreparationCancellation(t *testing.T) { + server, _, targetFile, _ := setupOverlayServer(t) + const sessionID = "canceled-overlay-diff" + require.NoError(t, server.OverlayManager().RegisterWithID(sessionID, "")) + require.NoError(t, server.OverlayManager().Push(sessionID, daemon.OverlayFile{ + Path: targetFile, Content: "package main\nfunc Target() {}\n", + }, nil)) + + ctx, cancel := context.WithCancel(WithSessionID(context.Background(), sessionID)) + cancel() + req := mcplib.CallToolRequest{} + req.Params.Arguments = map[string]any{"kind": "get_callers", "id": "target.go::Target"} + result, err := server.handleCompareWithOverlay(ctx, req) + require.Nil(t, result) + require.ErrorIs(t, err, context.Canceled) +} + +func TestPRReviewSimulationPropagatesCancellation(t *testing.T) { + server, _, targetFile, _ := setupOverlayServer(t) + const sessionID = "canceled-pr-simulation" + require.NoError(t, server.OverlayManager().RegisterWithID(sessionID, "")) + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + req := mcplib.CallToolRequest{} + req.Params.Arguments = map[string]any{ + "session_id": sessionID, + "edits": "[" + buildSingleFileEdit(targetFile, "package main\nfunc Target() {}\n") + "]", + } + result, gate, err := server.buildPRReviewSimulation(ctx, req) + require.Nil(t, result) + require.Equal(t, reviewGate{}, gate) + require.ErrorIs(t, err, context.Canceled) +} + +func TestCompareBranchesBoundsNamedBranchSnapshot(t *testing.T) { + server, sessionID, _, _, _, ctx := branchBootstrap(t) + _, err := server.OverlayManager().Fork(sessionID, daemon.ForkOptions{Name: "oversized"}) + require.NoError(t, err) + for i := 0; i < overlayRequestSnapshotMaxFiles+1; i++ { + require.NoError(t, server.OverlayManager().PushToBranch( + sessionID, + "oversized", + daemon.OverlayFile{ + Path: strings.Repeat("./", i) + "sentinel.go", + Content: "package sentinel\n", + }, + nil, + )) + } + + result, body := invokeTool(t, server, ctx, "compare_branches", map[string]any{ + "a": "oversized", + "b": daemon.MainBranchName, + "kind": "get_callers", + "id": "sentinel.go::Sentinel", + }) + require.True(t, result.IsError) + require.Contains(t, body, "overlay snapshot too large: files") +} + +func rawEmptyWorkspaceEditArray(count int) string { + if count <= 0 { + return "[]" + } + return "[" + strings.TrimSuffix(strings.Repeat("{},", count), ",") + "]" +} + +func pushRawOverlayAliases(t *testing.T, manager *daemon.OverlayManager, sessionID string, count int) { + t.Helper() + for i := 0; i < count; i++ { + require.NoError(t, manager.Push(sessionID, daemon.OverlayFile{ + Path: strings.Repeat("./", i) + "sentinel.go", + Content: "package sentinel\n", + }, nil)) + } +} + +func requireOverlaySnapshotBoundError(t *testing.T, err error, resource string, limit int) { + t.Helper() + var tooLarge *daemon.ErrOverlaySnapshotTooLarge + require.Error(t, err) + require.True(t, errors.As(err, &tooLarge), "expected ErrOverlaySnapshotTooLarge, got %T: %v", err, err) + require.Equal(t, resource, tooLarge.Resource) + require.Equal(t, limit, tooLarge.Limit) +} + +type cancelAfterNErrContext struct { + context.Context + allowedCalls int32 + errCalls atomic.Int32 +} + +func (c *cancelAfterNErrContext) Err() error { + if c.errCalls.Add(1) > c.allowedCalls { + return context.Canceled + } + return nil +} + +func fullFileSimulationEdit(path, content string) lsp.WorkspaceEdit { + return lsp.WorkspaceEdit{Changes: map[string][]lsp.TextEdit{ + path: {{ + Range: lsp.Range{ + Start: lsp.Position{}, + End: lsp.Position{Line: 100_000}, + }, + NewText: content, + }}, + }} +} diff --git a/internal/mcp/overlay_subprocess_extractor_test.go b/internal/mcp/overlay_subprocess_extractor_test.go new file mode 100644 index 00000000..aac7c8bb --- /dev/null +++ b/internal/mcp/overlay_subprocess_extractor_test.go @@ -0,0 +1,126 @@ +package mcp + +import ( + "context" + "os" + "path/filepath" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/zzet/gortex/internal/config" + "github.com/zzet/gortex/internal/daemon" + "github.com/zzet/gortex/internal/graph" + "github.com/zzet/gortex/internal/parser" + parserlanguages "github.com/zzet/gortex/internal/parser/languages" +) + +func newOverlaySubprocessExtractor(command string, args ...string) *parserlanguages.SubprocessExtractor { + return parserlanguages.NewSubprocessExtractor(config.ExtractorPluginSpec{ + Language: "overlay-plugin", + Extensions: []string{".plug"}, + Command: command, + Args: args, + }, nil) +} + +func writeSparsePluginOutput(t *testing.T, size int64, prefix string) string { + t.Helper() + path := filepath.Join(t.TempDir(), "plugin-output.json") + file, err := os.Create(path) + require.NoError(t, err) + if prefix != "" { + _, err = file.WriteString(prefix) + require.NoError(t, err) + } + require.NoError(t, file.Truncate(size)) + require.NoError(t, file.Close()) + return path +} + +func TestOverlaySubprocessExtractorFailuresDoNotBuildOrCacheReplacement(t *testing.T) { + resultPrefix := `{"nodes":[{"kind":"function","name":"Synthesized"}]}` + resultLimited := writeSparsePluginOutput(t, int64(overlayLayerParsedResultBytesMax), resultPrefix) + stdoutOverflow := writeSparsePluginOutput(t, int64(overlayLayerParsedResultBytesMax)+1, "") + + tests := []struct { + name string + extractor *parserlanguages.SubprocessExtractor + limitResource string + }{ + { + name: "stdout limit", + extractor: newOverlaySubprocessExtractor("cat", stdoutOverflow), + limitResource: "stdout_bytes", + }, + { + name: "combined result limit", + extractor: newOverlaySubprocessExtractor("cat", resultLimited), + limitResource: "result_bytes", + }, + { + name: "malformed output", + extractor: newOverlaySubprocessExtractor("sh", "-c", "cat >/dev/null; printf 'not-json'"), + }, + { + name: "nonzero command", + extractor: newOverlaySubprocessExtractor("sh", "-c", "cat >/dev/null; printf 'null'; exit 7"), + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + server, targetFile, _, store := setupOverlayExtractorServer(t, test.extractor) + graphPath := filepath.Base(targetFile) + old := &graph.Node{ID: graphPath + "::Old", Kind: graph.KindFunction, Name: "Old", FilePath: graphPath} + store.AddNode(old) + + ctx := WithSessionID(context.Background(), "bounded-plugin-session") + ctx = withOverlayRequestSnapshot(ctx, &overlayRequestSnapshot{ + sessionID: "bounded-plugin-session", + files: []daemon.OverlayFile{{Path: targetFile, Content: "overlay"}}, + canonical: true, + }) + view, err := server.buildOverlayViewForCtx(ctx) + require.Error(t, err) + require.Nil(t, view) + if test.limitResource != "" { + require.ErrorIs(t, err, parser.ErrExtractionLimit) + var limitErr *parser.ExtractionLimitError + require.ErrorAs(t, err, &limitErr) + require.Equal(t, test.limitResource, limitErr.Resource) + } else { + require.NotErrorIs(t, err, parser.ErrExtractionLimit) + } + _, cached := server.overlayLayerCache.Load("bounded-plugin-session") + require.False(t, cached, "failed parsing must not cache a file-only replacement") + baseNodes := store.Store.GetFileNodes(graphPath) + require.Len(t, baseNodes, 1) + require.Equal(t, old.ID, baseNodes[0].ID, "durable declaration must remain visible in the base graph") + fileCalls, _ := store.snapshotFileCalls() + require.Empty(t, fileCalls, "strict subprocess failures must stop before replacement identity reads") + }) + } +} + +func TestOverlaySubprocessLegacyExtractStillDegradesOperationalFailure(t *testing.T) { + extractor := newOverlaySubprocessExtractor("sh", "-c", "cat >/dev/null; printf 'not-json'") + result, err := extractor.Extract("target.plug", []byte("overlay")) + require.NoError(t, err) + require.NotNil(t, result) + require.Len(t, result.Nodes, 1) + require.Equal(t, graph.KindFile, result.Nodes[0].Kind) + require.Empty(t, result.Edges) +} + +func TestConstructOverlayLayerTrustedLegacyExtractorPathIsUnchanged(t *testing.T) { + server, targetFile, _, _, extractor := setupOverlayLayerBoundServer(t) + extractor.extract = func(_ int, filePath string, _ []byte) (*parser.ExtractionResult, error) { + return overlayRawExtractionResult(filePath, 1, 0), nil + } + layer, paths, err := server.constructOverlayLayer(context.Background(), []daemon.OverlayFile{{Path: targetFile, Content: "overlay"}}) + require.NoError(t, err) + require.NotNil(t, layer) + require.Len(t, paths, 1) + require.Equal(t, 1, extractor.callCount()) +} diff --git a/internal/mcp/overlay_temporal_options_test.go b/internal/mcp/overlay_temporal_options_test.go index 1366188d..d53ef253 100644 --- a/internal/mcp/overlay_temporal_options_test.go +++ b/internal/mcp/overlay_temporal_options_test.go @@ -1,6 +1,7 @@ package mcp import ( + "context" "os" "path/filepath" "testing" @@ -41,7 +42,7 @@ func Run(ctx workflow.Context) { workflow.ExecuteActivity(ctx, name) } ` - layer, paths, err := srv.constructOverlayLayer([]daemon.OverlayFile{{Path: "overlay.go", Content: source}}) + layer, paths, err := srv.constructOverlayLayer(context.Background(), []daemon.OverlayFile{{Path: "overlay.go", Content: source}}) if err != nil { t.Fatal(err) } diff --git a/internal/mcp/overlay_view.go b/internal/mcp/overlay_view.go index 831f616d..54fe4a27 100644 --- a/internal/mcp/overlay_view.go +++ b/internal/mcp/overlay_view.go @@ -4,7 +4,9 @@ import ( "context" "crypto/sha256" "encoding/hex" + "errors" "fmt" + "path" "path/filepath" "sort" "strings" @@ -23,6 +25,95 @@ import ( // `s.readerFor(ctx)`. Unexported so external code can't smuggle a // view onto an unrelated context. type overlayViewCtxKey struct{} +type overlayRequestSnapshotCtxKey struct{} + +const ( + overlayRequestSnapshotMaxFiles = exploreSourceLiteralOverlayMaxCoveredFiles + 1 + overlayRequestSnapshotMaxBytes = exploreSourceLiteralOverlayMaxBytes + exploreSourceLiteralOverlayMaxLineBytes + overlaySimulationInputMaxBytes = exploreSourceLiteralOverlayMaxBytes + overlaySimulationMaxSteps = 16 + + overlayLayerBaseNodesPerFileMax = 1024 + // BaseNodesMax bounds summaries retained across the layer build. Once the + // budget is full, a later file may transiently return one emptiness sentinel. + overlayLayerBaseNodesMax = 4096 + overlayLayerParsedNodesMax = 4096 + overlayLayerParsedEdgesMax = 16384 + overlayLayerParsedResultBytesMax = 16 << 20 + overlayLayerUnresolvedNamesMax = 1024 +) + +type overlayLayerLimitError struct { + Resource string + Limit int +} + +func (e *overlayLayerLimitError) Error() string { + return fmt.Sprintf("overlay layer %s exceeds limit %d", e.Resource, e.Limit) +} + +func validateBoundedOverlayExtractionUsage( + result *parser.ExtractionResult, + usage parser.ExtractionUsage, + limits parser.ExtractionLimits, +) error { + if usage.RawNodes < 1 || usage.RawEdges < 0 || usage.StdoutBytes < 0 || usage.ResultBytes < 0 { + return fmt.Errorf("bounded extractor returned negative or missing usage") + } + if usage.StdoutBytes > usage.ResultBytes { + return fmt.Errorf("bounded extractor result usage is smaller than stdout usage") + } + if len(result.Nodes) > usage.RawNodes || len(result.Edges) > usage.RawEdges { + return fmt.Errorf("bounded extractor result exceeds reported raw usage") + } + if usage.RawNodes > limits.MaxNodes { + return &overlayLayerLimitError{Resource: "parsed nodes", Limit: overlayLayerParsedNodesMax} + } + if usage.RawEdges > limits.MaxEdges { + return &overlayLayerLimitError{Resource: "parsed edges", Limit: overlayLayerParsedEdgesMax} + } + if usage.StdoutBytes > int64(limits.MaxStdoutBytes) { + return &overlayLayerLimitError{Resource: "parsed stdout bytes", Limit: overlayLayerParsedResultBytesMax} + } + if usage.ResultBytes > int64(limits.MaxResultBytes) { + return &overlayLayerLimitError{Resource: "parsed result bytes", Limit: overlayLayerParsedResultBytesMax} + } + return nil +} + +type stagedOverlayFile struct { + graphPath string + deleted bool + repoPrefix string + baseNodes []*graph.Node + result *parser.ExtractionResult +} + +// overlayRequestSnapshot is the immutable raw-buffer cohort used to build one +// request's OverlaidView. OverlayFile strings are immutable; retaining this +// slice through handler return adds only the copied slice headers and prevents +// a second manager snapshot from observing a different editor state. +type overlayRequestSnapshot struct { + sessionID string + workspace string + files []daemon.OverlayFile + canonical bool +} + +func withOverlayRequestSnapshot(ctx context.Context, snapshot *overlayRequestSnapshot) context.Context { + if ctx == nil || snapshot == nil { + return ctx + } + return context.WithValue(ctx, overlayRequestSnapshotCtxKey{}, snapshot) +} + +func overlayRequestSnapshotFromContext(ctx context.Context) (*overlayRequestSnapshot, bool) { + if ctx == nil { + return nil, false + } + snapshot, ok := ctx.Value(overlayRequestSnapshotCtxKey{}).(*overlayRequestSnapshot) + return snapshot, ok && snapshot != nil +} // WithOverlayView returns a child context carrying the // shadow-graph view for the current `tools/call`. Tool handlers @@ -101,29 +192,167 @@ type overlayLayerCacheEntry struct { files []string } -// buildOverlayViewForCtx is the per-request entry called by -// wrapToolHandler. Returns (nil, nil) for non-overlay sessions, the -// overlay-view for overlay-active sessions, or (nil, err) when drift -// detection trips so the client knows to refresh and resubmit. -func (s *Server) buildOverlayViewForCtx(ctx context.Context) (*graph.OverlaidView, error) { +func (s *Server) snapshotOverlayRequestForCtx(ctx context.Context) (*overlayRequestSnapshot, error) { if s == nil || s.overlays == nil { return nil, nil } - sessID := SessionIDFromContext(ctx) - if sessID == "" { + sessionID := SessionIDFromContext(ctx) + if sessionID == "" { return nil, nil } - if s.overlays.FileCount(sessID) == 0 { - return nil, nil + if err := ctx.Err(); err != nil { + return nil, err + } + workspace, files, err := s.overlays.SnapshotForBounded( + sessionID, + overlayRequestSnapshotMaxFiles, + overlayRequestSnapshotMaxBytes, + ) + if ctxErr := ctx.Err(); ctxErr != nil { + return nil, ctxErr } - _, files, err := s.overlays.SnapshotFor(sessID) if err != nil { - // Session evaporated. Fast-path the non-overlay route. - return nil, nil + if !errors.Is(err, daemon.ErrSessionNotFound) { + return nil, err + } + // Pin the attempted empty cohort. A nested facade must not retry and + // accidentally observe buffers pushed after the outer request began. + return &overlayRequestSnapshot{sessionID: sessionID}, nil + } + return &overlayRequestSnapshot{sessionID: sessionID, workspace: workspace, files: files}, nil +} + +// prepareOverlayRequest pins one manager snapshot and the graph view derived +// from it onto the request context. Reusing an already-prepared context is +// idempotent, which keeps nested facade calls on the same editor-buffer cohort. +func (s *Server) prepareOverlayRequest(ctx context.Context) (context.Context, *graph.OverlaidView, error) { + if ctx == nil { + return ctx, nil, nil + } + snapshot, ok := overlayRequestSnapshotFromContext(ctx) + if ok && snapshot.sessionID != SessionIDFromContext(ctx) { + return ctx, nil, fmt.Errorf("overlay request snapshot belongs to session %q, not %q", snapshot.sessionID, SessionIDFromContext(ctx)) + } + if !ok { + if OverlayViewFromContext(ctx) != nil { + return ctx, nil, fmt.Errorf("overlay view has no pinned request snapshot") + } + var err error + snapshot, err = s.snapshotOverlayRequestForCtx(ctx) + if err != nil { + return ctx, nil, err + } + } + if snapshot == nil { + return ctx, OverlayViewFromContext(ctx), nil + } + if view := OverlayViewFromContext(ctx); view != nil && !snapshot.canonical { + return ctx, nil, fmt.Errorf("overlay view has a non-canonical request snapshot") + } + if err := s.canonicalizeOverlayRequestSnapshot(snapshot); err != nil { + return ctx, nil, err + } + ctx = withOverlayRequestSnapshot(ctx, snapshot) + if view := OverlayViewFromContext(ctx); view != nil { + return ctx, view, nil + } + view, err := s.buildOverlayViewForCtx(ctx) + if err != nil { + return ctx, nil, err + } + if view != nil { + ctx = WithOverlayView(ctx, view) + } + return ctx, view, nil +} + +func canonicalOverlayGraphPath(candidate string) string { + candidate = strings.TrimSpace(strings.ReplaceAll(candidate, "\\", "/")) + if candidate == "" { + return "" + } + candidate = path.Clean(candidate) + if candidate == "." { + return "" + } + return candidate +} + +// canonicalizeOverlayRequestSnapshot validates and normalizes the pinned raw +// buffer cohort exactly once. Aliases resolving to one graph path collapse +// only when their replacement state is identical. SnapshotFor does not retain +// push chronology, so conflicting aliases fail closed instead of guessing +// which editor state is newer. +func (s *Server) canonicalizeOverlayRequestSnapshot(snapshot *overlayRequestSnapshot) error { + if snapshot == nil || snapshot.canonical { + return nil + } + if len(snapshot.files) == 0 { + snapshot.canonical = true + return nil } - if len(files) == 0 { + + sort.SliceStable(snapshot.files, func(i, j int) bool { + return snapshot.files[i].Path < snapshot.files[j].Path + }) + type canonicalRecord struct { + file daemon.OverlayFile + rawPath string + } + byPath := make(map[string]canonicalRecord, len(snapshot.files)) + for _, file := range snapshot.files { + rawPath := file.Path + absPath, err := s.resolveOverlayAbsPath(rawPath) + if err != nil { + return err + } + owner := s.pickIndexerForPath(absPath) + if absPath == "" || owner == nil { + return fmt.Errorf("overlay path %q is outside the registered workspace", rawPath) + } + if snapshot.workspace != "" && owner.WorkspaceID() != snapshot.workspace { + return fmt.Errorf("overlay path %q belongs to workspace %q, not registered workspace %q", rawPath, owner.WorkspaceID(), snapshot.workspace) + } + graphPath := canonicalOverlayGraphPath(s.resolveOverlayGraphPath(rawPath, absPath)) + if graphPath == "" { + return fmt.Errorf("overlay path %q has no canonical graph path", rawPath) + } + if existing, ok := byPath[graphPath]; ok { + if existing.file.Content != file.Content || existing.file.Deleted != file.Deleted || existing.file.BaseSHA != file.BaseSHA { + return fmt.Errorf("conflicting overlay aliases %q and %q resolve to %q", existing.rawPath, rawPath, graphPath) + } + continue + } + file.Path = graphPath + byPath[graphPath] = canonicalRecord{file: file, rawPath: rawPath} + } + + paths := make([]string, 0, len(byPath)) + for graphPath := range byPath { + paths = append(paths, graphPath) + } + sort.Strings(paths) + files := make([]daemon.OverlayFile, 0, len(paths)) + for _, graphPath := range paths { + files = append(files, byPath[graphPath].file) + } + snapshot.files = files + snapshot.canonical = true + return nil +} + +// buildOverlayViewForCtx consumes only the raw snapshot installed by request +// preparation, ensuring parsing and later raw-buffer reads observe one cohort. +func (s *Server) buildOverlayViewForCtx(ctx context.Context) (*graph.OverlaidView, error) { + snapshot, ok := overlayRequestSnapshotFromContext(ctx) + if !ok || snapshot == nil || len(snapshot.files) == 0 { return nil, nil } + if !snapshot.canonical { + return nil, fmt.Errorf("overlay request snapshot is not canonical") + } + files := snapshot.files + sessID := SessionIDFromContext(ctx) // Drift check up front for every overlay that carries a BaseSHA. // We do it here, before parsing, so a stale overlay never costs @@ -167,7 +396,7 @@ func (s *Server) buildOverlayViewForCtx(ctx context.Context) (*graph.OverlaidVie } } - layer, paths, err := s.constructOverlayLayer(files) + layer, paths, err := s.constructOverlayLayer(ctx, files) if err != nil { return nil, err } @@ -276,14 +505,35 @@ func (s *Server) pickIndexerForPath(absPath string) *indexer.Indexer { // simple name resolution — but covers the common cases: direct // function calls, method calls in the same file, intra-package // references. -func (s *Server) constructOverlayLayer(files []daemon.OverlayFile) (*graph.OverlayLayer, []string, error) { +func (s *Server) constructOverlayLayer(ctx context.Context, files []daemon.OverlayFile) (*graph.OverlayLayer, []string, error) { + if ctx == nil { + ctx = context.Background() + } + if err := ctx.Err(); err != nil { + return nil, nil, err + } + if err := validateOverlayBuildEnvelope(files); err != nil { + return nil, nil, err + } if s.graph == nil { return nil, nil, nil } - layer := graph.NewOverlayLayer() - var coveredPaths []string + + // Stage every bounded read and extraction before constructing the layer. + // Extractors may reuse result pointers, so no prefix or graph mutation is + // allowed until every file has passed its caps and cancellation checks. + baseReader, _ := s.graph.(graph.BoundedFileNodeReader) + staged := make([]stagedOverlayFile, 0, len(files)) + coveredPaths := make([]string, 0, len(files)) + baseNodeCount := 0 + parsedNodeCount := 0 + parsedEdgeCount := 0 + parsedResultBytes := int64(0) for _, ov := range files { + if err := ctx.Err(); err != nil { + return nil, nil, err + } absPath, err := s.resolveOverlayAbsPath(ov.Path) if err != nil { return nil, nil, err @@ -295,11 +545,15 @@ func (s *Server) constructOverlayLayer(files []daemon.OverlayFile) (*graph.Overl coveredPaths = append(coveredPaths, graphPath) if ov.Deleted { - // Tombstone: hide every base node for this file. - for _, n := range s.graph.GetFileNodes(graphPath) { - layer.MarkRemoved(n.Name, n.ID) + baseNodes, err := readOverlayBaseNodes(ctx, baseReader, graphPath, &baseNodeCount) + if err != nil { + return nil, nil, err } - layer.MarkFile(graphPath, true) + staged = append(staged, stagedOverlayFile{ + graphPath: graphPath, + deleted: true, + baseNodes: baseNodes, + }) continue } @@ -311,7 +565,8 @@ func (s *Server) constructOverlayLayer(files []daemon.OverlayFile) (*graph.Overl if reg == nil { continue } - lang, ok := reg.DetectLanguageContent(absPath, []byte(ov.Content)) + content := []byte(ov.Content) + lang, ok := reg.DetectLanguageContent(absPath, content) if !ok { continue } @@ -324,41 +579,145 @@ func (s *Server) constructOverlayLayer(files []daemon.OverlayFile) (*graph.Overl relPath = filepath.ToSlash(r) } } - result, err := idx.ExtractBuffer(lang, relPath, []byte(ov.Content)) + if err := ctx.Err(); err != nil { + return nil, nil, err + } + + var ( + result *parser.ExtractionResult + extractErr error + boundedUsage parser.ExtractionUsage + boundedLimits parser.ExtractionLimits + bounded bool + ) + ext, extractorOK := reg.GetByLanguage(lang) + if boundedExtractor, isBounded := ext.(parser.BoundedExtractor); extractorOK && isBounded { + bounded = true + boundedLimits = parser.DefaultExtractionLimits() + boundedLimits.MaxNodes = overlayLayerParsedNodesMax - parsedNodeCount + boundedLimits.MaxEdges = overlayLayerParsedEdgesMax - parsedEdgeCount + remainingResultBytes := int64(overlayLayerParsedResultBytesMax) - parsedResultBytes + boundedLimits.MaxStdoutBytes = int(remainingResultBytes) + boundedLimits.MaxResultBytes = int(remainingResultBytes) + result, boundedUsage, extractErr = boundedExtractor.ExtractBounded( + ctx, relPath, content, boundedLimits, + ) + } else { + // Legacy extractors go through the indexer's admission lifecycle + // so overlay parses share its crash isolation. + result, extractErr = idx.ExtractBuffer(lang, relPath, content) + } + if result != nil { + // Overlay construction never retains parse trees or constant-value + // sidecars. Release/clear before every error and cancellation branch. + result.ReleaseTree() + result.ConstValues = nil + } + if err := ctx.Err(); err != nil { + return nil, nil, err + } + if extractErr != nil { + return nil, nil, fmt.Errorf("overlay parse %s: %w", ov.Path, extractErr) + } + if result == nil { + return nil, nil, fmt.Errorf("overlay parse %s: extractor returned nil result", ov.Path) + } + if bounded { + if err := validateBoundedOverlayExtractionUsage(result, boundedUsage, boundedLimits); err != nil { + return nil, nil, fmt.Errorf("overlay parse %s: %w", ov.Path, err) + } + } + + // Keep the raw returned-slice checks as defense in depth for both optional + // bounded implementations and trusted legacy extractors. + if len(result.Nodes) > overlayLayerParsedNodesMax-parsedNodeCount { + return nil, nil, &overlayLayerLimitError{ + Resource: "parsed nodes", + Limit: overlayLayerParsedNodesMax, + } + } + if len(result.Edges) > overlayLayerParsedEdgesMax-parsedEdgeCount { + return nil, nil, &overlayLayerLimitError{ + Resource: "parsed edges", + Limit: overlayLayerParsedEdgesMax, + } + } + if bounded { + parsedNodeCount += boundedUsage.RawNodes + parsedEdgeCount += boundedUsage.RawEdges + parsedResultBytes += boundedUsage.ResultBytes + } else { + parsedNodeCount += len(result.Nodes) + parsedEdgeCount += len(result.Edges) + } + baseNodes, err := readOverlayBaseNodes(ctx, baseReader, graphPath, &baseNodeCount) if err != nil { - return nil, nil, fmt.Errorf("overlay parse %s: %w", ov.Path, err) - } - // Track which base IDs disappear under the overlay so - // FindNodesByName / GetInEdges filter them. Build the set - // first from base, then mark every base ID that the overlay - // did NOT re-emit (by ID equality). - baseIDsByName := map[string]map[string]bool{} - for _, n := range s.graph.GetFileNodes(graphPath) { - set, ok := baseIDsByName[n.Name] - if !ok { - set = make(map[string]bool) - baseIDsByName[n.Name] = set + return nil, nil, err + } + staged = append(staged, stagedOverlayFile{ + graphPath: graphPath, + repoPrefix: idx.RepoPrefix(), + baseNodes: baseNodes, + result: cloneOverlayExtractionResult(result), + }) + } + + if len(coveredPaths) == 0 { + return nil, nil, nil + } + if err := ctx.Err(); err != nil { + return nil, nil, err + } + sort.Strings(coveredPaths) + + layer := graph.NewOverlayLayer() + for index, file := range staged { + if index&15 == 0 { + if err := ctx.Err(); err != nil { + return nil, nil, err + } + } + if file.deleted { + for _, node := range file.baseNodes { + layer.MarkRemoved(node.Name, node.ID) } - set[n.ID] = true - } - overlayIDsByName := map[string]map[string]bool{} - applyRepoPrefixToResult(result, idx.RepoPrefix()) - layer.MarkFile(graphPath, false) - for _, n := range result.Nodes { - layer.AddNode(graphPath, n) - set, ok := overlayIDsByName[n.Name] - if !ok { + layer.MarkFile(file.graphPath, true) + continue + } + + // The staged result is an owned, unprefixed clone captured before the + // next Extract call. Prefix it only after the full cohort preflight. + result := file.result + applyRepoPrefixToResult(result, file.repoPrefix) + + baseIDsByName := make(map[string]map[string]bool) + for _, node := range file.baseNodes { + set := baseIDsByName[node.Name] + if set == nil { set = make(map[string]bool) - overlayIDsByName[n.Name] = set + baseIDsByName[node.Name] = set } - set[n.ID] = true + set[node.ID] = true } - for _, e := range result.Edges { - layer.AddEdge(e) + overlayIDsByName := make(map[string]map[string]bool) + layer.MarkFile(file.graphPath, false) + if result != nil { + for _, node := range result.Nodes { + if node == nil { + continue + } + layer.AddNode(file.graphPath, node) + set := overlayIDsByName[node.Name] + if set == nil { + set = make(map[string]bool) + overlayIDsByName[node.Name] = set + } + set[node.ID] = true + } + for _, edge := range result.Edges { + layer.AddEdge(edge) + } } - // Names that existed in base but were not re-emitted by the - // overlay (same name, different ID, or absent entirely) get - // marked removed so FindNodesByName filters the base hits. for name, baseIDs := range baseIDsByName { overlayIDs := overlayIDsByName[name] for id := range baseIDs { @@ -369,17 +728,168 @@ func (s *Server) constructOverlayLayer(files []daemon.OverlayFile) (*graph.Overl } } - if len(coveredPaths) == 0 { - return nil, nil, nil + if err := resolveOverlayEdges(ctx, s.graph, layer); err != nil { + return nil, nil, err } - sort.Strings(coveredPaths) + return layer, coveredPaths, nil +} - // Local resolver pass: rewrite unresolved overlay edges to point - // at concrete IDs whenever a single best match exists in - // (overlay ∪ base). - s.resolveOverlayEdges(layer) +func readOverlayBaseNodes( + ctx context.Context, + reader graph.BoundedFileNodeReader, + graphPath string, + total *int, +) ([]*graph.Node, error) { + if reader == nil { + return nil, fmt.Errorf("overlay base identities for %s: %w", graphPath, graph.ErrBoundedLocalizationUnavailable) + } + if total == nil || *total < 0 || *total > overlayLayerBaseNodesMax { + return nil, fmt.Errorf("overlay base identities for %s: invalid aggregate count", graphPath) + } + if err := ctx.Err(); err != nil { + return nil, err + } - return layer, coveredPaths, nil + remaining := overlayLayerBaseNodesMax - *total + queryLimit := overlayLayerBaseNodesPerFileMax + aggregateLimited := remaining < queryLimit + if aggregateLimited { + queryLimit = remaining + if queryLimit == 0 { + // A one-row emptiness sentinel lets an actually empty later file pass + // without ever issuing an invalid limit-zero projection. + queryLimit = 1 + } + } + projection, err := reader.FindFileNodesBounded( + ctx, + graphPath, + graph.LocalizationNodeScope{}, + queryLimit, + ) + if ctxErr := ctx.Err(); ctxErr != nil { + return nil, ctxErr + } + if err != nil { + return nil, fmt.Errorf("overlay base identities for %s: %w", graphPath, err) + } + if projection.Truncated || projection.Total > queryLimit || len(projection.Nodes) > queryLimit { + resource := "base nodes per file" + limit := overlayLayerBaseNodesPerFileMax + if aggregateLimited { + resource = "base nodes" + limit = overlayLayerBaseNodesMax + } + return nil, &overlayLayerLimitError{Resource: resource, Limit: limit} + } + if projection.Total < 0 || projection.Total != len(projection.Nodes) { + return nil, fmt.Errorf("overlay base identities for %s: invalid bounded projection", graphPath) + } + if projection.Total > remaining { + return nil, &overlayLayerLimitError{ + Resource: "base nodes", + Limit: overlayLayerBaseNodesMax, + } + } + for _, node := range projection.Nodes { + if node == nil { + return nil, fmt.Errorf("overlay base identities for %s: invalid nil node", graphPath) + } + } + *total += projection.Total + return projection.Nodes, nil +} + +func cloneOverlayExtractionResult(result *parser.ExtractionResult) *parser.ExtractionResult { + if result == nil { + return nil + } + cloned := &parser.ExtractionResult{ + Nodes: make([]*graph.Node, len(result.Nodes)), + Edges: make([]*graph.Edge, len(result.Edges)), + } + for index, node := range result.Nodes { + if node == nil { + continue + } + copy := *node + copy.Meta = cloneOverlayMetadata(node.Meta) + cloned.Nodes[index] = © + } + for index, edge := range result.Edges { + if edge == nil { + continue + } + copy := *edge + copy.Meta = cloneOverlayMetadata(edge.Meta) + cloned.Edges[index] = © + } + return cloned +} + +func cloneOverlayMetadata(meta map[string]any) map[string]any { + if meta == nil { + return nil + } + cloned := make(map[string]any, len(meta)) + for key, value := range meta { + cloned[key] = value + } + return cloned +} + +func validateOverlayBuildEnvelope(files []daemon.OverlayFile) error { + if len(files) > overlayRequestSnapshotMaxFiles { + return overlayBuildEnvelopeError("files", overlayRequestSnapshotMaxFiles) + } + totalBytes := 0 + for _, file := range files { + if err := addOverlayBuildFileBytes(&totalBytes, file); err != nil { + return err + } + } + return nil +} + +// validateOverlayBuildMapEnvelope checks simulation state before it is copied +// into a retained, sorted snapshot. Keeping this map-shaped seam avoids +// materializing an already-oversized state merely to discover its size. +func validateOverlayBuildMapEnvelope(files map[string]daemon.OverlayFile) error { + if len(files) > overlayRequestSnapshotMaxFiles { + return overlayBuildEnvelopeError("files", overlayRequestSnapshotMaxFiles) + } + totalBytes := 0 + for _, file := range files { + if err := addOverlayBuildFileBytes(&totalBytes, file); err != nil { + return err + } + } + return nil +} + +func addOverlayBuildFileBytes(totalBytes *int, file daemon.OverlayFile) error { + for _, size := range [...]int{len(file.Path), len(file.Content), len(file.BaseSHA)} { + remaining := overlayRequestSnapshotMaxBytes - *totalBytes + if size > remaining { + return overlayBuildEnvelopeError("bytes", overlayRequestSnapshotMaxBytes) + } + *totalBytes += size + } + return nil +} + +func overlayBuildEnvelopeError(resource string, limit int) error { + return &daemon.ErrOverlaySnapshotTooLarge{Resource: resource, Limit: limit} +} + +// requestContextError preserves request transport cancellation instead of +// converting it into an MCP tool-error payload. Only the parent request context +// is authoritative: component-local cancellation remains an ordinary tool error. +func requestContextError(ctx context.Context, _ error) error { + if ctx == nil { + return nil + } + return ctx.Err() } // applyRepoPrefixToResult prepends repoPrefix to every node/edge in @@ -434,76 +944,140 @@ const unresolvedPrefix = "unresolved::" // overlay buffers are transient and the common case the editor // cares about is "I added a call to Foo; does find_usages of Foo // now include this site?" Direct name resolution covers that. -func (s *Server) resolveOverlayEdges(layer *graph.OverlayLayer) { +func resolveOverlayEdges(ctx context.Context, base graph.Reader, layer *graph.OverlayLayer) error { if layer == nil { - return + return nil } - // Collect every From → []Edge that the layer holds. We iterate - // over a copy of the map so we can rewrite layer edges - // in-place via AddEdge / removal pattern (layer is meant - // to be append-only post-construction; the resolver pass runs - // before the layer is handed to the View, so we still own it). - for _, edges := range layer.OutEdgesByFromAll() { - for _, e := range edges { - if !strings.HasPrefix(e.To, unresolvedPrefix) { - continue - } - target := strings.TrimPrefix(e.To, unresolvedPrefix) - // Strip kind segment if present (e.g. "call::FooBar"). - if i := strings.Index(target, "::"); i > 0 { - target = target[i+2:] + if ctx == nil { + ctx = context.Background() + } + if err := ctx.Err(); err != nil { + return err + } + + edgesByFrom := layer.OutEdgesByFromAll() + names := make(map[string]struct{}) + inspected := 0 + for _, edges := range edgesByFrom { + for _, edge := range edges { + inspected++ + if inspected&127 == 0 { + if err := ctx.Err(); err != nil { + return err + } } - // Strip trailing argument-count / disambiguator hints. - if i := strings.Index(target, "@"); i > 0 { - target = target[:i] + if edge == nil { + continue } - if target == "" { + name := overlayUnresolvedTargetName(edge.To) + if name == "" { continue } - resolved := s.lookupOverlayTarget(layer, target) - if resolved == "" { + if _, exists := names[name]; exists { continue } - e.To = resolved + if len(names) >= overlayLayerUnresolvedNamesMax { + return &overlayLayerLimitError{ + Resource: "unresolved names", + Limit: overlayLayerUnresolvedNamesMax, + } + } + names[name] = struct{}{} } } - // Rebuild the layer's inEdges index now that targets may have - // changed. The layer exposes a Rebuild helper so we don't have - // to know the internal map shape. - layer.RebuildInEdges() -} -// lookupOverlayTarget tries to find a unique node with the given -// short name in (layer ∪ base). Returns the node ID on a unique -// match, empty string otherwise. Tied matches return empty so the -// edge stays as a placeholder rather than picking the wrong target. -func (s *Server) lookupOverlayTarget(layer *graph.OverlayLayer, name string) string { - overlay := layer.NodesByName(name) - if len(overlay) == 1 { - return overlay[0].ID - } - if len(overlay) > 1 { - return "" + orderedNames := make([]string, 0, len(names)) + for name := range names { + orderedNames = append(orderedNames, name) } - if s.graph == nil { - return "" + sort.Strings(orderedNames) + resolvedByName := make(map[string]string, len(orderedNames)) + view := graph.NewOverlaidView(base, layer) + for _, name := range orderedNames { + if err := ctx.Err(); err != nil { + return err + } + overlay := layer.NodesByName(name) + if len(overlay) == 1 { + if overlay[0] == nil { + return fmt.Errorf("overlay target %q: invalid nil overlay node", name) + } + resolvedByName[name] = overlay[0].ID + continue + } + if len(overlay) > 1 { + continue + } + + projection, err := view.FindNodesByNameBounded( + ctx, + name, + graph.LocalizationNodeScope{}, + 1, + ) + if ctxErr := ctx.Err(); ctxErr != nil { + return ctxErr + } + if err != nil { + return fmt.Errorf("overlay target %q: %w", name, err) + } + if projection.Total < 0 || projection.Total < len(projection.Nodes) || len(projection.Nodes) > 1 { + return fmt.Errorf("overlay target %q: invalid bounded projection", name) + } + if projection.Truncated || projection.Total > 1 { + continue // ambiguous names remain unresolved + } + switch projection.Total { + case 0: + if len(projection.Nodes) != 0 { + return fmt.Errorf("overlay target %q: invalid empty projection", name) + } + case 1: + if len(projection.Nodes) != 1 || projection.Nodes[0] == nil { + return fmt.Errorf("overlay target %q: invalid unique projection", name) + } + resolvedByName[name] = projection.Nodes[0].ID + default: + return fmt.Errorf("overlay target %q: invalid bounded projection", name) + } } - hits := s.graph.FindNodesByName(name) - // Drop hits whose file is overlaid AND whose ID wasn't kept by - // the overlay — those are now-deleted symbols. - keep := hits[:0:0] - for _, n := range hits { - if layer.HasFile(graph.IDFile(n.ID)) { - if !layer.HasNode(n.ID) { + + inspected = 0 + for _, edges := range edgesByFrom { + for _, edge := range edges { + inspected++ + if inspected&127 == 0 { + if err := ctx.Err(); err != nil { + return err + } + } + if edge == nil { continue } + if resolved := resolvedByName[overlayUnresolvedTargetName(edge.To)]; resolved != "" { + edge.To = resolved + } } - keep = append(keep, n) } - if len(keep) == 1 { - return keep[0].ID + if err := ctx.Err(); err != nil { + return err } - return "" + layer.RebuildInEdges() + return ctx.Err() +} + +func overlayUnresolvedTargetName(target string) string { + if !strings.HasPrefix(target, unresolvedPrefix) { + return "" + } + target = strings.TrimPrefix(target, unresolvedPrefix) + if index := strings.Index(target, "::"); index > 0 { + target = target[index+2:] + } + if index := strings.Index(target, "@"); index > 0 { + target = target[:index] + } + return target } // hashOverlayFiles produces a stable content-hash of an overlay @@ -530,22 +1104,15 @@ func hashOverlayFiles(files []daemon.OverlayFile) string { // case, since a tombstone has no content to return and callers // should treat the file as absent. func (s *Server) overlayContentFor(ctx context.Context, absPath string) (string, bool) { - if s == nil || s.overlays == nil || ctx == nil { + if s == nil || ctx == nil { return "", false } - sessID := SessionIDFromContext(ctx) - if sessID == "" { - return "", false - } - if s.overlays.FileCount(sessID) == 0 { - return "", false - } - _, files, err := s.overlays.SnapshotFor(sessID) - if err != nil || len(files) == 0 { + snapshot, ok := overlayRequestSnapshotFromContext(ctx) + if !ok || snapshot == nil || !snapshot.canonical { return "", false } cleanedAbs := filepath.Clean(absPath) - for _, ov := range files { + for _, ov := range snapshot.files { if ov.Deleted { continue } @@ -560,6 +1127,38 @@ func (s *Server) overlayContentFor(ctx context.Context, absPath string) (string, return "", false } +// overlayShadowsGraphPath reports whether the active request replaces or +// tombstones the path represented by a base search hit. Search indexes are +// built from durable files, so their snippets must never authenticate stale +// content from a file whose editor overlay owns the request-time view. +func (s *Server) overlayShadowsGraphPath(ctx context.Context, graphPath, nodeID string) bool { + view := OverlayViewFromContext(ctx) + if view == nil || view.Layer() == nil { + return false + } + layer := view.Layer() + for _, candidate := range []string{graphPath, graph.IDFile(nodeID)} { + candidate = filepath.ToSlash(filepath.Clean(strings.TrimSpace(candidate))) + if candidate != "" && candidate != "." && layer.HasFile(candidate) { + return true + } + } + return false +} + +func (s *Server) filterOverlayContentHits(ctx context.Context, hits []graph.ContentHit) []graph.ContentHit { + if OverlayViewFromContext(ctx) == nil || len(hits) == 0 { + return hits + } + filtered := hits[:0:0] + for _, hit := range hits { + if !s.overlayShadowsGraphPath(ctx, hit.FilePath, hit.NodeID) { + filtered = append(filtered, hit) + } + } + return filtered +} + // overlayCacheInvalidate drops the cached layer for a session. Called // by overlay_push / overlay_delete / overlay_drop so the next tool // call re-parses with the fresh buffer state. diff --git a/internal/mcp/search_consumption_cursor_test.go b/internal/mcp/search_consumption_cursor_test.go new file mode 100644 index 00000000..a9ce5c08 --- /dev/null +++ b/internal/mcp/search_consumption_cursor_test.go @@ -0,0 +1,96 @@ +package mcp + +import ( + "context" + "encoding/json" + "testing" + + mcplib "github.com/mark3labs/mcp-go/mcp" + "go.uber.org/zap" + + "github.com/zzet/gortex/internal/graph" + "github.com/zzet/gortex/internal/query" +) + +func TestSearchSymbolsCursorRecordsReturnedPageAndClearsOutOfRange(t *testing.T) { + store := graph.New() + store.AddBatch([]*graph.Node{ + {ID: "repo/a.go::CursorNeedleAlpha", Kind: graph.KindFunction, Name: "CursorNeedleAlpha", FilePath: "repo/a.go"}, + {ID: "repo/b.go::CursorNeedleBeta", Kind: graph.KindFunction, Name: "CursorNeedleBeta", FilePath: "repo/b.go"}, + {ID: "repo/c.go::CursorNeedleGamma", Kind: graph.KindFunction, Name: "CursorNeedleGamma", FilePath: "repo/c.go"}, + }, nil) + server := NewServer(query.NewEngine(store), store, nil, nil, zap.NewNop(), nil) + ctx := context.Background() + + callPage := func(cursor string) map[string]any { + t.Helper() + arguments := map[string]any{"query": "CursorNeedle", "limit": 1} + if cursor != "" { + arguments["cursor"] = cursor + } + request := mcplib.CallToolRequest{} + request.Params.Name = "search_symbols" + request.Params.Arguments = arguments + result, err := findAndCallHandler(server, "search_symbols", ctx, request) + if err != nil { + t.Fatalf("search_symbols: %v", err) + } + if result.IsError { + t.Fatalf("search_symbols returned an error: %#v", result.Content) + } + text, ok := result.Content[0].(mcplib.TextContent) + if !ok { + t.Fatalf("search_symbols content = %T, want text", result.Content[0]) + } + var response map[string]any + if err := json.Unmarshal([]byte(text.Text), &response); err != nil { + t.Fatalf("decode search response: %v", err) + } + return response + } + + assertRecordedPage := func(response map[string]any) string { + t.Helper() + rows, ok := response["results"].([]any) + if !ok || len(rows) != 1 { + t.Fatalf("results = %#v, want one row", response["results"]) + } + row, ok := rows[0].(map[string]any) + if !ok { + t.Fatalf("result row = %T, want object", rows[0]) + } + id, _ := row["id"].(string) + if id == "" { + t.Fatalf("result row has no id: %#v", row) + } + + session := server.sessionFor(ctx) + session.mu.Lock() + recorded := append([]string(nil), session.lastSearch.returned...) + session.mu.Unlock() + if len(recorded) != 1 || recorded[0] != id { + t.Fatalf("recorded page = %#v, response id = %q", recorded, id) + } + return id + } + + firstID := assertRecordedPage(callPage("")) + secondID := assertRecordedPage(callPage(encodeCursor(1))) + if secondID == firstID { + t.Fatalf("cursor page repeated first result %q", firstID) + } + + response := callPage(encodeCursor(10_000)) + if raw := response["results"]; raw != nil { + if rows, ok := raw.([]any); !ok || len(rows) != 0 { + t.Fatalf("out-of-range results = %#v, want empty", raw) + } + } + session := server.sessionFor(ctx) + session.mu.Lock() + state := session.lastSearch + session.mu.Unlock() + if state.query != "" || len(state.returned) != 0 || len(state.returnedIDs) != 0 || len(state.fileIDs) != 0 || len(state.consumed) != 0 || !state.at.IsZero() { + t.Fatalf("out-of-range page retained attribution: %#v", state) + } +} diff --git a/internal/mcp/search_consumption_page_test.go b/internal/mcp/search_consumption_page_test.go new file mode 100644 index 00000000..9e20afa7 --- /dev/null +++ b/internal/mcp/search_consumption_page_test.go @@ -0,0 +1,176 @@ +package mcp + +import ( + "fmt" + "reflect" + "sync" + "testing" + "time" + + "github.com/zzet/gortex/internal/graph" +) + +func TestRecordLastSearchPageBoundsPositionsAndCanonicalFiles(t *testing.T) { + nodes := make([]*graph.Node, lastSearchPageCap+1) + for index := range nodes { + path := fmt.Sprintf("repo/file-%03d.go", index) + nodes[index] = &graph.Node{ID: path + "::Symbol", FilePath: path} + } + + var session sessionState + session.recordLastSearchPage("needle", nodes) + + session.mu.Lock() + state := session.lastSearch + gotReturned := append([]string(nil), state.returned...) + gotLastRank, gotLast := state.returnedIDs[nodes[lastSearchPageCap-1].ID] + _, gotOverflow := state.returnedIDs[nodes[lastSearchPageCap].ID] + _, gotOverflowFile := state.fileIDs[nodes[lastSearchPageCap].FilePath] + gotFileBuckets := len(state.fileIDs) + gotMemberships := 0 + for _, bucket := range state.fileIDs { + gotMemberships += len(bucket) + } + session.mu.Unlock() + + if len(gotReturned) != lastSearchPageCap { + t.Fatalf("returned len = %d, want %d", len(gotReturned), lastSearchPageCap) + } + if !gotLast || gotLastRank != lastSearchPageCap-1 { + t.Fatalf("last admitted rank = (%d, %v), want (%d, true)", gotLastRank, gotLast, lastSearchPageCap-1) + } + if gotOverflow || gotOverflowFile { + t.Fatalf("position %d escaped cap: id=%v file=%v", lastSearchPageCap, gotOverflow, gotOverflowFile) + } + if gotFileBuckets > lastSearchPageCap || gotMemberships > lastSearchPageCap { + t.Fatalf("file attribution escaped cap: buckets=%d memberships=%d cap=%d", gotFileBuckets, gotMemberships, lastSearchPageCap) + } +} + +func TestRecordLastSearchPageUsesRawPositionsAndCanonicalFallback(t *testing.T) { + nodes := make([]*graph.Node, lastSearchPageCap+1) + nodes[0] = nil + nodes[1] = &graph.Node{} + nodes[2] = &graph.Node{ID: "repo/a.go::A", FilePath: " ./repo/../repo/a.go "} + nodes[3] = &graph.Node{ID: "repo/a.go::A", FilePath: "wrong/duplicate.go"} + nodes[4] = &graph.Node{ID: "repo/fallback.go::B", FilePath: " . "} + for index := 5; index < lastSearchPageCap; index++ { + nodes[index] = nil + } + nodes[lastSearchPageCap] = &graph.Node{ID: "repo/overflow.go::C", FilePath: "repo/overflow.go"} + + var session sessionState + session.recordLastSearchPage("needle", nodes) + + session.mu.Lock() + gotReturned := append([]string(nil), session.lastSearch.returned...) + gotA := append([]string(nil), session.lastSearch.fileIDs["repo/a.go"]...) + gotFallback := append([]string(nil), session.lastSearch.fileIDs["repo/fallback.go"]...) + _, gotDuplicatePath := session.lastSearch.fileIDs["wrong/duplicate.go"] + _, gotOverflowPath := session.lastSearch.fileIDs["repo/overflow.go"] + session.mu.Unlock() + + if want := []string{"repo/a.go::A", "repo/fallback.go::B"}; !reflect.DeepEqual(gotReturned, want) { + t.Fatalf("returned = %#v, want %#v", gotReturned, want) + } + if want := []string{"repo/a.go::A"}; !reflect.DeepEqual(gotA, want) { + t.Fatalf("canonical file IDs = %#v, want %#v", gotA, want) + } + if want := []string{"repo/fallback.go::B"}; !reflect.DeepEqual(gotFallback, want) { + t.Fatalf("fallback file IDs = %#v, want %#v", gotFallback, want) + } + if gotDuplicatePath || gotOverflowPath { + t.Fatalf("inadmissible file membership recorded: duplicate=%v overflow=%v", gotDuplicatePath, gotOverflowPath) + } +} + +func TestAttributedFileConsumptionIsFreshExactAndOnce(t *testing.T) { + var session sessionState + session.recordLastSearchPage("needle", []*graph.Node{ + {ID: "repo/a.go::A", FilePath: "repo/a.go"}, + {ID: "repo/a.go::B", FilePath: "repo/a.go"}, + {ID: "repo/b.go::C", FilePath: "repo/b.go"}, + }) + + query, ids := session.attributedFileConsumption(" ./repo/a.go ") + if query != "needle" || !reflect.DeepEqual(ids, []string{"repo/a.go::A", "repo/a.go::B"}) { + t.Fatalf("first attribution = (%q, %#v)", query, ids) + } + if query, ids := session.attributedFileConsumption("repo/a.go"); query != "" || len(ids) != 0 { + t.Fatalf("repeat attribution = (%q, %#v), want empty", query, ids) + } + if query, ids := session.attributedFileConsumption("repo/missing.go"); query != "" || len(ids) != 0 { + t.Fatalf("other-file attribution = (%q, %#v), want empty", query, ids) + } + if query, ids := session.attributedFileConsumption("repo/b.go"); query != "needle" || !reflect.DeepEqual(ids, []string{"repo/b.go::C"}) { + t.Fatalf("second file attribution = (%q, %#v)", query, ids) + } + + session.recordLastSearchPage("stale", []*graph.Node{{ID: "repo/stale.go::S", FilePath: "repo/stale.go"}}) + session.mu.Lock() + session.lastSearch.at = time.Now().Add(-comboWindow - time.Second) + session.mu.Unlock() + if query, ids := session.attributedFileConsumption("repo/stale.go"); query != "" || len(ids) != 0 { + t.Fatalf("stale attribution = (%q, %#v), want empty", query, ids) + } +} + +func TestRecordLastSearchPageEmptyClearsPriorState(t *testing.T) { + var session sessionState + session.recordLastSearchPage("first", []*graph.Node{{ID: "repo/a.go::A", FilePath: "repo/a.go"}}) + session.recordLastSearchPage("second", nil) + + session.mu.Lock() + state := session.lastSearch + session.mu.Unlock() + if state.query != "" || len(state.returned) != 0 || len(state.returnedIDs) != 0 || len(state.fileIDs) != 0 || len(state.consumed) != 0 || !state.at.IsZero() { + t.Fatalf("empty page retained state: %#v", state) + } +} + +func TestRecordLastSearchPageOwnsInputsAndAttributionResults(t *testing.T) { + node := &graph.Node{ID: "repo/a.go::A", FilePath: "repo/a.go"} + page := []*graph.Node{node} + var session sessionState + session.recordLastSearchPage("needle", page) + + page[0] = nil + node.ID = "mutated" + node.FilePath = "mutated.go" + query, ids := session.attributedFileConsumption("repo/a.go") + if query != "needle" || !reflect.DeepEqual(ids, []string{"repo/a.go::A"}) { + t.Fatalf("owned attribution = (%q, %#v)", query, ids) + } + ids[0] = "mutated-result" + + session.mu.Lock() + gotReturned := append([]string(nil), session.lastSearch.returned...) + session.mu.Unlock() + if !reflect.DeepEqual(gotReturned, []string{"repo/a.go::A"}) { + t.Fatalf("stored IDs aliased caller/result: %#v", gotReturned) + } +} + +func TestRecordLastSearchPageConcurrentAccessIsBounded(t *testing.T) { + var session sessionState + var workers sync.WaitGroup + for worker := 0; worker < 8; worker++ { + worker := worker + workers.Add(1) + go func() { + defer workers.Done() + for iteration := 0; iteration < 32; iteration++ { + path := fmt.Sprintf("repo/%d.go", worker) + session.recordLastSearchPage("needle", []*graph.Node{{ID: fmt.Sprintf("%s::S%d", path, iteration), FilePath: path}}) + session.attributedFileConsumption(path) + } + }() + } + workers.Wait() + + session.mu.Lock() + defer session.mu.Unlock() + if len(session.lastSearch.returned) > lastSearchPageCap || len(session.lastSearch.returnedIDs) > lastSearchPageCap || len(session.lastSearch.fileIDs) > lastSearchPageCap { + t.Fatalf("concurrent state escaped cap: %#v", session.lastSearch) + } +} diff --git a/internal/mcp/search_symbols_content_fallback.go b/internal/mcp/search_symbols_content_fallback.go index debc8be6..63aa8147 100644 --- a/internal/mcp/search_symbols_content_fallback.go +++ b/internal/mcp/search_symbols_content_fallback.go @@ -88,6 +88,9 @@ func (s *Server) searchSymbolsContentFallback( search := s.searchExploreSourceLiteral( boundedCtx, term, "", scope, searchSymbolsContentFallbackMaxMatches, ) + if ctx.Err() != nil { + return nil + } if len(search.matches) == 0 { return nil } @@ -109,7 +112,7 @@ func (s *Server) searchSymbolsContentFallback( return nil } - enriched := s.enrichTextMatches(exact) + enriched, _ := s.enrichTextMatchesContext(boundedCtx, exact, scope) if len(enriched) == 0 { return nil } diff --git a/internal/mcp/server.go b/internal/mcp/server.go index cc84844a..59028282 100644 --- a/internal/mcp/server.go +++ b/internal/mcp/server.go @@ -6,6 +6,7 @@ import ( "fmt" "math" "os" + "path/filepath" "runtime" "sort" "strings" @@ -774,12 +775,14 @@ func (ss *sessionState) markCueOnce(key string) bool { type lastSearchState struct { query string // returned is the result IDs in rank order (0 = top); returnedIDs - // maps an ID to its rank for O(1) membership + rank lookup. consumed - // tracks which returned IDs the agent went on to use, so a later - // search can record an implicit "skip-above" negative for the + // maps an ID to its rank for O(1) membership + rank lookup. fileIDs + // maps each canonical returned Node.FilePath to the IDs on that page. + // consumed tracks which returned IDs the agent went on to use, so a + // later search can record an implicit "skip-above" negative for the // higher-ranked results that were passed over. returned []string returnedIDs map[string]int + fileIDs map[string][]string consumed map[string]struct{} at time.Time } @@ -1132,28 +1135,80 @@ func (s *Server) resolveSessionFormat(ctx context.Context) string { // with a T-second window; 5 minutes is long enough for agents that // interleave many tool calls but short enough that an unrelated later // consume doesn't get mis-attributed. -const comboWindow = 5 * time.Minute +const ( + comboWindow = 5 * time.Minute + lastSearchPageCap = 256 +) -// recordLastSearch captures the query + the IDs it returned so a later -// consume call can be credited to this query. Truncating to the top N -// results keeps the map small — only symbols the agent can plausibly -// have seen are eligible. -func (ss *sessionState) recordLastSearch(query string, ids []string) { +func normalizeSearchAttributionPath(path string) string { + path = strings.TrimSpace(path) + if path == "" { + return "" + } + path = filepath.ToSlash(filepath.Clean(path)) + if path == "." { + return "" + } + return path +} + +// recordLastSearch captures one actually returned page so a later consume call +// can be credited to this query. At most lastSearchPageCap page positions are +// inspected; ranked IDs and per-file memberships are both unique and bounded. +func (ss *sessionState) recordLastSearchPage(query string, nodes []*graph.Node) { ss.mu.Lock() defer ss.mu.Unlock() - set := make(map[string]int, len(ids)) - for i, id := range ids { - set[id] = i + + positions := min(len(nodes), lastSearchPageCap) + returned := make([]string, 0, positions) + returnedIDs := make(map[string]int, positions) + fileIDs := make(map[string][]string) + for index := 0; index < positions; index++ { + node := nodes[index] + if node == nil || node.ID == "" { + continue + } + if _, duplicate := returnedIDs[node.ID]; duplicate { + continue + } + returnedIDs[node.ID] = len(returned) + returned = append(returned, node.ID) + path := normalizeSearchAttributionPath(node.FilePath) + if path == "" { + path = normalizeSearchAttributionPath(graph.IDFile(node.ID)) + } + if path != "" { + fileIDs[path] = append(fileIDs[path], node.ID) + } + } + if strings.TrimSpace(query) == "" || len(returned) == 0 { + ss.lastSearch = lastSearchState{} + return } ss.lastSearch = lastSearchState{ query: query, - returned: append([]string(nil), ids...), - returnedIDs: set, + returned: returned, + returnedIDs: returnedIDs, + fileIDs: fileIDs, consumed: make(map[string]struct{}), at: time.Now(), } } +// recordLastSearch retains the legacy ID-only seam used by symbol consumers +// that do not have node metadata. File memberships are recovered from IDs when +// possible; search_symbols uses recordLastSearchPage with authoritative nodes. +func (ss *sessionState) recordLastSearch(query string, ids []string) { + nodes := make([]*graph.Node, 0, min(len(ids), lastSearchPageCap)) + for index, id := range ids { + if index >= lastSearchPageCap { + break + } + nodes = append(nodes, &graph.Node{ID: id}) + } + ss.recordLastSearchPage(query, nodes) +} + // attributedQuery returns the query string that should receive credit for // consuming symbolID, or "" if no recent search eligibly returned it. // Cleared from the caller's perspective but not from state — a single @@ -1177,42 +1232,57 @@ func (ss *sessionState) attributedQuery(symbolID string) string { return ss.lastSearch.query } -// attributedConsumptionBatch credits a set of symbol IDs to the recent -// search in one pass: it returns the search's query and the subset of -// ids that the search returned within the attribution window (marking -// each consumed). Used by the tool-call observer when the agent opens a -// file — every symbol in it that the search surfaced is credited at -// once. Returns ("", nil) when no fresh search is attributable. -func (ss *sessionState) attributedConsumptionBatch(ids []string) (string, []string) { +// attributedFileConsumption returns the recent search query and the exact +// returned-page IDs recorded for path, marking each as consumed. The lookup is +// graph-free: file membership was captured with the page under this same lock. +func (ss *sessionState) attributedFileConsumption(path string) (string, []string) { + path = normalizeSearchAttributionPath(path) ss.mu.Lock() defer ss.mu.Unlock() - if ss.lastSearch.query == "" || time.Since(ss.lastSearch.at) > comboWindow { + if path == "" || ss.lastSearch.query == "" || time.Since(ss.lastSearch.at) > comboWindow { + return "", nil + } + ids := ss.lastSearch.fileIDs[path] + if len(ids) == 0 { return "", nil } if ss.lastSearch.consumed == nil { ss.lastSearch.consumed = make(map[string]struct{}) } - var matched []string + matched := make([]string, 0, len(ids)) for _, id := range ids { - if id == "" { - continue - } - if _, ok := ss.lastSearch.returnedIDs[id]; !ok { + if _, consumed := ss.lastSearch.consumed[id]; consumed { continue } ss.lastSearch.consumed[id] = struct{}{} matched = append(matched, id) } + if len(matched) == 0 { + return "", nil + } return ss.lastSearch.query, matched } -// hasFreshSearch reports whether a search is recent enough to attribute -// a consume to. A cheap gate so file-open handlers skip the work of -// enumerating a file's symbols when nothing could be credited anyway. -func (ss *sessionState) hasFreshSearch() bool { +// attributedConsumptionBatch preserves the legacy ID-batch seam while file +// consumers migrate to the graph-free attributedFileConsumption lookup. +func (ss *sessionState) attributedConsumptionBatch(ids []string) (string, []string) { ss.mu.Lock() defer ss.mu.Unlock() - return ss.lastSearch.query != "" && time.Since(ss.lastSearch.at) <= comboWindow + if ss.lastSearch.query == "" || time.Since(ss.lastSearch.at) > comboWindow { + return "", nil + } + if ss.lastSearch.consumed == nil { + ss.lastSearch.consumed = make(map[string]struct{}) + } + matched := make([]string, 0, len(ids)) + for _, id := range ids { + if _, ok := ss.lastSearch.returnedIDs[id]; !ok || id == "" { + continue + } + ss.lastSearch.consumed[id] = struct{}{} + matched = append(matched, id) + } + return ss.lastSearch.query, matched } // drainSkippedNegatives computes the implicit "skip-above" negatives for diff --git a/internal/mcp/tools_analyze_named.go b/internal/mcp/tools_analyze_named.go index 603444c5..a5908056 100644 --- a/internal/mcp/tools_analyze_named.go +++ b/internal/mcp/tools_analyze_named.go @@ -122,25 +122,15 @@ func (s *Server) handleAnalyzeNamed(ctx context.Context, req mcp.CallToolRequest if err != nil { return mcp.NewToolResultError(err.Error()), nil } - fileSymbols := s.buildFileSymbolIndex(targets) - lookup := func(graphPath string, line int) (string, string) { - idx := fileSymbols[graphPath] - if idx == nil { - return "", "" - } - return idx.find(line) - } - rows := make([]sastRow, 0, 64) summary := make(map[string]*sastSummary, len(detectors)) var errsAcc []string for _, d := range detectors { opts := astquery.Options{ - Detector: d.Name, - Targets: targets, - SymbolLookup: lookup, - Resolver: astquery.DefaultLanguageResolver, - Limit: 5000, + Detector: d.Name, + Targets: targets, + Resolver: astquery.DefaultLanguageResolver, + Limit: 5000, } if excludeTestsSet { opts.ExcludeTests = excludeTests @@ -206,6 +196,13 @@ func (s *Server) handleAnalyzeNamed(ctx context.Context, req mcp.CallToolRequest rows = rows[:limit] truncated = true } + s.enrichASTSymbolIDsContext( + ctx, + len(rows), + func(index int) string { return rows[index].File }, + func(index int) int { return rows[index].Line }, + func(index int, id string) { rows[index].Symbol = id }, + ) summaries := make([]sastSummary, 0, len(summary)) for _, e := range summary { diff --git a/internal/mcp/tools_analyze_sast.go b/internal/mcp/tools_analyze_sast.go index 08873154..ca0c05f8 100644 --- a/internal/mcp/tools_analyze_sast.go +++ b/internal/mcp/tools_analyze_sast.go @@ -59,6 +59,39 @@ type sastCWEBucket struct { Count int `json:"count"` } +func sortSASTMatchesByPriority(matches []astquery.Match) { + sort.SliceStable(matches, func(i, j int) bool { + ri, rj := severityRank(matches[i].Severity), severityRank(matches[j].Severity) + if ri != rj { + return ri > rj + } + if matches[i].Detector != matches[j].Detector { + return matches[i].Detector < matches[j].Detector + } + if matches[i].File != matches[j].File { + return matches[i].File < matches[j].File + } + return matches[i].Line < matches[j].Line + }) +} + +func (s *Server) prepareReviewSASTMatchesContext( + ctx context.Context, + matches []astquery.Match, + limit int, + enrich bool, +) { + sortSASTMatchesByPriority(matches) + if !enrich { + return + } + enrichCount := len(matches) + if limit > 0 && enrichCount > limit { + enrichCount = limit + } + s.enrichASTMatchesContext(ctx, matches[:enrichCount]) +} + // handleAnalyzeSAST runs the category bundle (sast or hygiene). func (s *Server) handleAnalyzeSAST(ctx context.Context, req mcp.CallToolRequest, kind string) (*mcp.CallToolResult, error) { args := req.GetArguments() @@ -85,15 +118,6 @@ func (s *Server) handleAnalyzeSAST(ctx context.Context, req mcp.CallToolRequest, return mcp.NewToolResultError(err.Error()), nil } - fileSymbols := s.buildFileSymbolIndex(targets) - lookup := func(graphPath string, line int) (string, string) { - idx := fileSymbols[graphPath] - if idx == nil { - return "", "" - } - return idx.find(line) - } - bundle := astquery.DetectorsByCategory(kind) if len(bundle) == 0 { return mcp.NewToolResultError(fmt.Sprintf("analyze %s: no detectors registered for category %q", kind, kind)), nil @@ -158,11 +182,10 @@ func (s *Server) handleAnalyzeSAST(ctx context.Context, req mcp.CallToolRequest, detMeta[d.Name] = d opts := astquery.Options{ - Detector: d.Name, - Targets: targets, - SymbolLookup: lookup, - Resolver: astquery.DefaultLanguageResolver, - Limit: 5000, + Detector: d.Name, + Targets: targets, + Resolver: astquery.DefaultLanguageResolver, + Limit: 5000, } if excludeTestsSet { opts.ExcludeTests = excludeTests @@ -206,6 +229,11 @@ func (s *Server) handleAnalyzeSAST(ctx context.Context, req mcp.CallToolRequest, // the resolved call / loop metadata refutes. Only the review // bundle is grounded; sast / hygiene / domain pass through. if kind == "review" { + // Establish the same global priority the final rows use before the + // bounded enrichment pass. Grounding preserves this order. Only the + // requested result prefix is enriched: grounding may drop a prefix + // row, but admitting later rows would exceed the caller's work bound. + s.prepareReviewSASTMatchesContext(ctx, collected, limit, !kindsOnly) collected = review.GroundReviewMatches(s.graph, collected) } @@ -263,6 +291,15 @@ func (s *Server) handleAnalyzeSAST(ctx context.Context, req mcp.CallToolRequest, rows = rows[:limit] truncated = true } + if !kindsOnly && kind != "review" { + s.enrichASTSymbolIDsContext( + ctx, + len(rows), + func(index int) string { return rows[index].File }, + func(index int) int { return rows[index].Line }, + func(index int, id string) { rows[index].Symbol = id }, + ) + } summaries := make([]sastSummary, 0, len(summary)) for _, entry := range summary { diff --git a/internal/mcp/tools_analyze_unsafe_patterns.go b/internal/mcp/tools_analyze_unsafe_patterns.go index 6cc59c8f..4a63a13b 100644 --- a/internal/mcp/tools_analyze_unsafe_patterns.go +++ b/internal/mcp/tools_analyze_unsafe_patterns.go @@ -50,9 +50,9 @@ type unsafePatternSummary struct { // Filters (all optional): // - language — comma-separated subset (rust,python,javascript,typescript,go). // - detector — comma-separated subset (must be a member of -// astquery.UnsafePatternDetectors). Lets the agent -// narrow the bundle without falling back to -// individual search_ast calls. +// astquery.UnsafePatternDetectors). Lets the agent +// narrow the bundle without falling back to +// individual search_ast calls. // - severity — comma-separated subset (error,warning,info). // - path_prefix — keep matches whose file path starts with this. // - limit — cap rows (default 200; matches error_surface's UX). @@ -78,17 +78,6 @@ func (s *Server) handleAnalyzeUnsafePatterns(ctx context.Context, req mcp.CallTo return mcp.NewToolResultError(err.Error()), nil } - // Per-file enclosing-symbol index — shared across detectors so - // every row gets symbol enrichment without re-building. - fileSymbols := s.buildFileSymbolIndex(targets) - lookup := func(graphPath string, line int) (string, string) { - idx := fileSymbols[graphPath] - if idx == nil { - return "", "" - } - return idx.find(line) - } - // Reject unknown detector names early so the agent gets a // pointed error instead of an empty result. if len(detectorFilter) > 0 { @@ -123,10 +112,9 @@ func (s *Server) handleAnalyzeUnsafePatterns(ctx context.Context, req mcp.CallTo } } opts := astquery.Options{ - Detector: name, - Targets: targets, - SymbolLookup: lookup, - Resolver: astquery.DefaultLanguageResolver, + Detector: name, + Targets: targets, + Resolver: astquery.DefaultLanguageResolver, // Generous per-detector cap; the outer `limit` is the // agent-facing budget. Picking 5000 protects against a // pathological repo where one detector returns tens of @@ -212,6 +200,13 @@ func (s *Server) handleAnalyzeUnsafePatterns(ctx context.Context, req mcp.CallTo rows = rows[:limit] truncated = true } + s.enrichASTSymbolIDsContext( + ctx, + len(rows), + func(index int) string { return rows[index].File }, + func(index int) int { return rows[index].Line }, + func(index int, id string) { rows[index].Symbol = id }, + ) summaries := make([]unsafePatternSummary, 0, len(summary)) for _, name := range astquery.UnsafePatternDetectors { diff --git a/internal/mcp/tools_ast.go b/internal/mcp/tools_ast.go index 363fa911..d36b586a 100644 --- a/internal/mcp/tools_ast.go +++ b/internal/mcp/tools_ast.go @@ -97,28 +97,13 @@ func (s *Server) handleSearchAST(ctx context.Context, req mcp.CallToolRequest) ( return mcp.NewToolResultError(err.Error()), nil } - // Build a per-file enclosing-symbol index lazily. Each file - // is a small list of function/method/closure nodes; the - // lookup is amortised by caching the per-file index on first - // hit. The graph walk is single-pass so even big indexes - // pay it once per `search_ast` call. - fileSymbols := s.buildFileSymbolIndex(targets) - lookup := func(graphPath string, line int) (string, string) { - idx := fileSymbols[graphPath] - if idx == nil { - return "", "" - } - return idx.find(line) - } - opts := astquery.Options{ - Pattern: pattern, - Detector: detector, - Language: language, - Targets: targets, - SymbolLookup: lookup, - Resolver: astquery.DefaultLanguageResolver, - Limit: limit, + Pattern: pattern, + Detector: detector, + Language: language, + Targets: targets, + Resolver: astquery.DefaultLanguageResolver, + Limit: limit, } // Honor explicit override; otherwise let the engine apply // its per-mode default (true for detectors, false for raw @@ -133,6 +118,7 @@ func (s *Server) handleSearchAST(ctx context.Context, req mcp.CallToolRequest) ( if runErr != nil { return mcp.NewToolResultError(runErr.Error()), nil } + s.enrichASTMatchesContext(ctx, res.Matches) if minFanIn > 0 { res.Matches = filterByMinFanIn(s.graph, res.Matches, minFanIn) diff --git a/internal/mcp/tools_ast_test.go b/internal/mcp/tools_ast_test.go index 4c28493b..5f94e9af 100644 --- a/internal/mcp/tools_ast_test.go +++ b/internal/mcp/tools_ast_test.go @@ -85,6 +85,55 @@ func F() { } } +func TestSearchASTUsesBoundedEnclosingProjection(t *testing.T) { + srv, _ := setupTestServer(t) + abs := writeTempGoFile(t, "bounded.go", `package x + +func Bounded() { + panic("boom") +} +`) + quietAbs := writeTempGoFile(t, "quiet.go", `package x + +func Quiet() {} +`) + backing := graph.New() + backing.AddNode(&graph.Node{ + ID: abs, Kind: graph.KindFile, Name: abs, + FilePath: abs, Language: "go", StartLine: 1, EndLine: 5, + }) + backing.AddNode(&graph.Node{ + ID: quietAbs, Kind: graph.KindFile, Name: quietAbs, + FilePath: quietAbs, Language: "go", StartLine: 1, EndLine: 3, + }) + owner := &graph.Node{ + ID: abs + "::Bounded", Kind: graph.KindFunction, Name: "Bounded", + FilePath: abs, StartLine: 3, EndLine: 5, + } + backing.AddNode(owner) + probe := &astTargetBoundedStore{Store: backing} + probe.read = func(ctx context.Context, path string, scope graph.LocalizationNodeScope, limit int) (graph.BoundedNodeProjection, error) { + return backing.FindFileNodesBounded(ctx, path, scope, limit) + } + srv.graph = probe + + out := callSearchAST(t, srv, map[string]any{ + "pattern": `((call_expression function: (identifier) @fn) @match (#eq? @fn "panic"))`, + "language": "go", + }) + + if got, _ := out["total"].(float64); got != 1 { + t.Fatalf("expected 1 bounded match, got %v\n%v", got, out) + } + match := out["matches"].([]any)[0].(map[string]any) + if match["symbol_id"] != owner.ID || match["symbol_name"] != owner.Name { + t.Fatalf("bounded handler enrichment = %#v, want %s/%s", match, owner.ID, owner.Name) + } + if len(probe.calls) != 1 || probe.calls[0].path != abs { + t.Fatalf("bounded handler calls = %#v, want one read of %q", probe.calls, abs) + } +} + func TestSearchAST_BundledDetector_HardcodedSecret(t *testing.T) { srv, _ := setupTestServer(t) abs := writeTempGoFile(t, "creds.go", `package x diff --git a/internal/mcp/tools_core.go b/internal/mcp/tools_core.go index 6d9adef8..083346b9 100644 --- a/internal/mcp/tools_core.go +++ b/internal/mcp/tools_core.go @@ -837,21 +837,6 @@ func stripNonDefinitionNodes(sg *query.SubGraph) *query.SubGraph { } } -// fileDefinitionNodes enumerates one file's declared symbols through the same -// graph query and definition filter get_file_summary answers from, without the -// tool layer's freshness, encoding, and accounting passes. The path must -// already be in the graph's stored form. -func fileDefinitionNodes(eng *query.Engine, filePath string) []*graph.Node { - if eng == nil || strings.TrimSpace(filePath) == "" { - return nil - } - sg := stripNonDefinitionNodes(eng.GetFileSymbols(filePath)) - if sg == nil { - return nil - } - return sg.Nodes -} - // compactSubGraph formats a SubGraph as compact text. func compactSubGraph(sg *query.SubGraph) string { var b strings.Builder @@ -1917,10 +1902,6 @@ func (s *Server) handleSearchSymbols(ctx context.Context, req mcp.CallToolReques } } - // Remember the returned IDs for attribution on later consume calls. - // Cap at top limit so unseen "overflow" results don't get credited. - recordLastSearchFromNodes(sess, q, nodes, limit) - total := len(nodes) // Slice the (offset, limit) window. nextCursor is empty when the // last row in `nodes` is included. @@ -1932,6 +1913,9 @@ func (s *Server) handleSearchSymbols(ctx context.Context, req mcp.CallToolReques offset = total } page := nodes[offset:end] + // Record only the actual post-cursor page. Empty and out-of-range pages + // deliberately clear any prior attribution state. + recordLastSearchFromNodes(sess, q, page) // Decorate the page with absolute file paths so every output format // below surfaces an openable path alongside the repo-relative one. page = s.withAbsPaths(page) diff --git a/internal/mcp/tools_explore.go b/internal/mcp/tools_explore.go index b00ac0fa..111f70a9 100644 --- a/internal/mcp/tools_explore.go +++ b/internal/mcp/tools_explore.go @@ -50,12 +50,12 @@ const ( // unrestricted grep agent showed the answer present-but-unnamed far more // often than absent; serving bodies is what closes that, and the session // token spend stays well under the plain-search agent's. - localizationDefaultBudgetTokens = 12000 - exploreMinBudgetTokens = 1000 - exploreMaxBudgetTokens = 24000 - exploreDefaultMaxSymbols = 16 - exploreMaxMaxSymbols = 30 - exploreRingCap = 5 // callers / callees shown per target + localizationDefaultBudgetTokens = 12000 + exploreMinBudgetTokens = 1000 + exploreMaxBudgetTokens = 24000 + exploreDefaultMaxSymbols = 16 + exploreMaxMaxSymbols = 30 + exploreRingCap = 5 // callers / callees shown per target // Kept at eight while the page itself widens: this cap also sets the // envelope's shed floor and the tight-budget wire guarantee, so widening // it is a contract change, not a serving change. @@ -93,30 +93,77 @@ func (s *Server) registerExploreTool() { // exploreTarget is one ranked candidate plus its bounded neighborhood, // gathered before rendering so the renderer can honour the token budget. type exploreTarget struct { - node *graph.Node - score float64 - callers []*graph.Node - callees []*graph.Node - directCalleesComplete bool // false when the direct projection was truncated, bounded, or otherwise lower-bound - causalCallees []exploreCausalNeighbor - causalChangeBridge bool // one graph-proven caller/callee retained as provenance for a promoted continuation - causalChangeLeaf bool // graph-proven wrapper implementation or task-aligned cross-file change callable - causalChangeOwner bool // same-file type that encloses or is uniquely returned by the causal change callable - source string // full body (may be empty for non-source kinds) - divergentDefaultOwner bool // unique child constructor whose concrete default causes the queried behavior - divergentDefaultType bool // owning type paired with divergentDefaultOwner for coherent file/symbol output - conceptImplementation bool // primary identifier-backed callable; may establish answer readiness - conceptComplement bool // marginal concept callable protected as evidence, never as terminal proof - syntacticAnchor bool // task-spelled flag/identifier owner protected by bounded lexical competition - exactContent bool // verified full quoted-literal hit from content_fts - exactContentAmbiguous bool // exact evidence has visible or possibly truncated peers - sourceLiteral bool // exact source-body hit that must survive final envelope packing - sourceLiteralCallee bool // exact source callsite uniquely resolved to this invoked callable - sourceLiteralAligned bool // source-literal callee that instantiates the task's value; strongest literal owner - typedAnchorProjection bool // bounded field-owner-call proof promoted from a task-aligned typed field - foldedOwner bool // synthetic owner inserted by concept member folding - leadingFileDepth bool // sibling of the ranked leading file, admitted into the expendable breadth tail - localizationRelation string // direct_caller/direct_callee row promoted only into the bounded terminal projection + node *graph.Node + score float64 + callers []*graph.Node + callees []*graph.Node + directCalleesComplete bool // false when the direct projection was truncated, bounded, or otherwise lower-bound + causalCallees []exploreCausalNeighbor + causalChangeBridge bool // one graph-proven caller/callee retained as provenance for a promoted continuation + causalChangeLeaf bool // graph-proven wrapper implementation or task-aligned cross-file change callable + causalChangeOwner bool // same-file type that encloses or is uniquely returned by the causal change callable + source string // full body (may be empty for non-source kinds) + sourceWindow *localizationSourceWindow + divergentDefaultOwner bool // unique child constructor whose concrete default causes the queried behavior + divergentDefaultType bool // owning type paired with divergentDefaultOwner for coherent file/symbol output + conceptImplementation bool // primary identifier-backed callable; may establish answer readiness + conceptComplement bool // marginal concept callable protected as evidence, never as terminal proof + syntacticAnchor bool // task-spelled flag/identifier owner protected by bounded lexical competition + sourceRange bool // exact task path+line owner reserved for presentation, never completion authority + exactContent bool // verified full quoted-literal hit from content_fts + exactContentAmbiguous bool // exact evidence has visible or possibly truncated peers + sourceLiteral bool // exact source-body hit that must survive final envelope packing + sourceLiteralCallee bool // exact source callsite uniquely resolved to this invoked callable + sourceLiteralAligned bool // source-literal callee that instantiates the task's value; strongest literal owner + literalPrimaryEligible bool // authenticated unanchored literal may use the one bounded PRIMARY reserve + literalMatchCount int // distinct task literals matched; stable tie-break among eligible rows + typedAnchorProjection bool // bounded field-owner-call proof promoted from a task-aligned typed field + foldedOwner bool // synthetic owner inserted by concept member folding + leadingFileDepth bool // sibling of the ranked leading file, admitted into the expendable breadth tail + localizationRelation string // direct_caller/direct_callee row promoted only into the bounded terminal projection +} + +func exploreTargetFromCandidate( + candidate *rerank.Candidate, + protectedImplementationID string, + literalPrimaryEligible bool, + quotedLiteralTask bool, +) exploreTarget { + if candidate == nil || candidate.Node == nil { + return exploreTarget{} + } + target := exploreTarget{ + node: candidate.Node, score: candidate.Score, + conceptImplementation: candidate.Node.ID == protectedImplementationID, + } + if candidate.Signals == nil { + return target + } + target.conceptComplement = candidate.Signals[exploreConceptComplementSignal] > 0 + target.syntacticAnchor = candidate.Signals[exploreSyntacticAnchorSignal] > 0 + target.sourceRange = candidate.Signals[exploreSourceRangeSignal] > 0 + target.exactContent = candidate.Signals[exploreContentRecallExactSignal] > 0 + target.exactContentAmbiguous = candidate.Signals[exploreContentRecallAmbiguousSignal] > 0 + target.sourceLiteral = candidate.Signals[exploreSourceLiteralSignal] > 0 + target.sourceLiteralCallee = candidate.Signals[exploreSourceLiteralCalleeSignal] > 0 + target.sourceLiteralAligned = candidate.Signals[exploreSourceLiteralTaskAlignSignal] > 0 + target.typedAnchorProjection = candidate.Signals[exploreTypedAnchorProjectionSignal] > 0 + // Eligibility is per-row for evidence grounded in a literal the caller + // quoted: that the task also names a file or symbol does not make the + // quoted value's registration site less of an answer. Only the inferred + // bare-literal lane — whose literals are Gortex's own guess — remains + // gated on the page having no explicit anchor. + eligible := literalPrimaryEligible || + quotedLiteralTask && !target.exactContentAmbiguous + if eligible && (target.sourceLiteral || target.exactContent) { + target.literalPrimaryEligible = true + matches := candidate.Signals[exploreContentRecallTermSignal] + if coverage := candidate.Signals[exploreSourceLiteralCoverageSignal]; coverage > matches { + matches = coverage + } + target.literalMatchCount = int(matches) + } + return target } type exploreCausalNeighbor struct { @@ -2663,6 +2710,10 @@ func (s *Server) handleExplore(ctx context.Context, req mcp.CallToolRequest) (*m // paths, keys, flags, and environment names needed by config/CI searches. artifactIntent := classifyExploreArtifactIntent(task) localize := req.GetBool("localize", false) + var sourceWindowHits *localizationSourceWindowHitCollector + if localize { + sourceWindowHits = &localizationSourceWindowHitCollector{} + } // The same ranked spans, derived once, applied to source candidate paths. // Ordinary exploration keeps the zero value and so corroborates nothing. var pathProbes explorePathProbes @@ -2739,7 +2790,7 @@ func (s *Server) handleExplore(ctx context.Context, req mcp.CallToolRequest) (*m // graph-aware reranker: its centrality and edge hydration costs scale // with every candidate, not just the final response size. ranked = limitExploreCandidates(ranked, fetch*2) - if content := s.gatherExploreQuotedContentCandidates(ctx, task, ranked, fetch, opts); len(content) > 0 { + if content := s.gatherExploreQuotedContentCandidatesCollecting(ctx, task, ranked, fetch, opts, sourceWindowHits); len(content) > 0 { ranked = mergeExploreCandidates(ranked, content, 0) ranked = limitExploreCandidatesPreservingSourceLiteral(ranked, fetch*2) } @@ -2756,7 +2807,7 @@ func (s *Server) handleExplore(ctx context.Context, req mcp.CallToolRequest) (*m // supplies at most one candidate per matching file, and the final // source-evidence reservation keeps its strongest result without // reranking the already-ranked primary channel a second time. - if content := s.gatherExploreQuotedContentCandidates(ctx, task, ranked, fetch, opts); len(content) > 0 { + if content := s.gatherExploreQuotedContentCandidatesCollecting(ctx, task, ranked, fetch, opts, sourceWindowHits); len(content) > 0 { ranked = mergeExploreCandidates(ranked, content, 0) ranked = limitExploreCandidatesPreservingSourceLiteral(ranked, fetch*2) } @@ -2768,18 +2819,32 @@ func (s *Server) handleExplore(ctx context.Context, req mcp.CallToolRequest) (*m // index. var protectedSyntacticAnchors map[int]string if queryClass == rerank.QueryClassConcept { - anchorCandidates, protected := s.gatherExploreSyntacticAnchorCandidates( - ctx, task, ranked, eng, opts, rctx, + anchorCandidates, protected := s.gatherExploreSyntacticAnchorCandidatesCollecting( + ctx, task, ranked, eng, opts, rctx, sourceWindowHits, ) protectedSyntacticAnchors = protected if len(anchorCandidates) > 0 { ranked = mergeExploreCandidates(ranked, anchorCandidates, fetch) } } + // Only an unanchored concept task may infer bare source literals. Quoted + // literals, artifact evidence, explicit paths/names, and protected syntax + // already have stronger bounded lanes and make this path a strict no-op. + if exploreBareLiteralLaneEligible( + task, searchQuery, queryClass == rerank.QueryClassConcept, artifactLane.ready, + ranked, protectedSyntacticAnchors, + ) { + terms := exploreBareLiteralRecallTerms(task) + if content := s.gatherExploreContentCandidatesForTermsCollecting(ctx, task, terms, ranked, fetch, opts, sourceWindowHits); len(content) > 0 { + ranked = mergeExploreCandidates(ranked, content, 0) + ranked = limitExploreCandidatesPreservingSourceLiteral(ranked, fetch*2) + ranked = rerankExploreConceptCoverage(searchQuery, ranked) + } + } // Exact source citations in issue bodies are stronger than semantic ranking: // map each bounded file/line range to its smallest enclosing declaration and // place those task-spelled candidates at the head before final selection. - ranked = s.promoteExploreSourceRangeCandidates(ctx, task, ranked, opts) + ranked = s.promoteExploreSourceRangeCandidates(ctx, task, ranked, eng.Reader(), opts) // Resilience ladder: a warm-restarted daemon can transiently return an // empty scoped ranked result (workspace stamps not yet backfilled, or // search bundles served before their node payloads re-materialise) @@ -2900,19 +2965,11 @@ func (s *Server) handleExplore(ctx context.Context, req mcp.CallToolRequest) (*m Depth: 1, Limit: exploreRingCap * 3, Detail: "brief", WorkspaceID: resolved.WorkspaceID, ProjectID: resolved.ProjectID, RepoAllow: resolved.RepoAllow, } - explicitTarget := false - if _, hasPath := exploreQueryPathAnchors(searchQuery); hasPath { - explicitTarget = true - } - if !explicitTarget { - for _, c := range cands { - if c != nil && c.Node != nil && exploreLocalizationExplicitAnchor(searchQuery, c.Node) { - explicitTarget = true - break - } - } - } + explicitTarget := exploreHasExplicitCandidateTarget(searchQuery, cands) artifactTargets := artifactLane.targets + literalPrimaryEligible := exploreLiteralEvidenceEligible(searchQuery, cands, protectedSyntacticAnchors) + quotedLiteralTask := len(exploreQuotedRecallTerms(task)) > 0 || + len(exploreDistinctiveBareTokens(task, "")) > 0 targets := make([]exploreTarget, 0, len(artifactTargets)+len(cands)) // Artifact evidence leads only when the strong classifier activated its // lane. Ordinary source localization retains the exact prior target order. @@ -2922,20 +2979,7 @@ func (s *Server) handleExplore(ctx context.Context, req mcp.CallToolRequest) (*m continue } n := c.Node - t := exploreTarget{ - node: n, score: c.Score, - conceptImplementation: n.ID == protectedImplementationID, - } - if c.Signals != nil { - t.conceptComplement = c.Signals[exploreConceptComplementSignal] > 0 - t.syntacticAnchor = c.Signals[exploreSyntacticAnchorSignal] > 0 - t.exactContent = c.Signals[exploreContentRecallExactSignal] > 0 - t.exactContentAmbiguous = c.Signals[exploreContentRecallAmbiguousSignal] > 0 - t.sourceLiteral = c.Signals[exploreSourceLiteralSignal] > 0 - t.sourceLiteralCallee = c.Signals[exploreSourceLiteralCalleeSignal] > 0 - t.sourceLiteralAligned = c.Signals[exploreSourceLiteralTaskAlignSignal] > 0 - t.typedAnchorProjection = c.Signals[exploreTypedAnchorProjectionSignal] > 0 - } + t := exploreTargetFromCandidate(c, protectedImplementationID, literalPrimaryEligible, quotedLiteralTask) t.source = s.manifestSymbolSource(ctx, n) if callers := eng.GetCallers(n.ID, ringOpts); callers != nil { t.callers = ringNeighbors(callers.Nodes, n.ID, exploreRingCap) @@ -2989,9 +3033,12 @@ func (s *Server) handleExplore(ctx context.Context, req mcp.CallToolRequest) (*m } } if exploreQueryIsConceptTask(task) && len(targets) > len(artifactTargets) { - symbolTargets := promoteExploreDivergentDefaultOwner(task, targets[len(artifactTargets):], s.graph, maxSymbols, func(node *graph.Node) string { - return s.manifestSymbolSource(ctx, node) - }, s.divergentDefaultFallbackSLOOverride) + symbolTargets := promoteExploreDivergentDefaultOwner( + ctx, task, targets[len(artifactTargets):], eng.Reader(), s.localizationNodeScope(ctx, opts), maxSymbols, + func(node *graph.Node) string { + return s.manifestSymbolSource(ctx, node) + }, s.divergentDefaultFallbackSLOOverride, + ) targets = append(targets[:len(artifactTargets):len(artifactTargets)], symbolTargets...) } // Causal-change promotion is the second concept lane over the same ranked @@ -3001,17 +3048,20 @@ func (s *Server) handleExplore(ctx context.Context, req mcp.CallToolRequest) (*m if len(targets) > len(artifactTargets) { symbolTargets := targets[len(artifactTargets):] if exploreDivergentDefaultOwnerSymbol(symbolTargets) == "" { - symbolTargets = promoteExploreCausalChangeTargets(task, symbolTargets, s.graph, maxSymbols, func(node *graph.Node) string { - return s.manifestSymbolSource(ctx, node) - }, func(node *graph.Node) ([]*graph.Node, bool) { - callees := eng.GetCallChain(node.ID, ringOpts) - if callees == nil { - return nil, false - } - direct, projectionComplete := ringNeighborsProjection(callees.Nodes, node.ID, exploreRingCap) - complete := !callees.Truncated && !callees.BudgetHit && !callees.LowerBound && projectionComplete - return direct, complete - }) + symbolTargets = promoteExploreCausalChangeTargets( + ctx, task, symbolTargets, eng.Reader(), s.localizationNodeScope(ctx, opts), maxSymbols, + func(node *graph.Node) string { + return s.manifestSymbolSource(ctx, node) + }, func(node *graph.Node) ([]*graph.Node, bool) { + callees := eng.GetCallChain(node.ID, ringOpts) + if callees == nil { + return nil, false + } + direct, projectionComplete := ringNeighborsProjection(callees.Nodes, node.ID, exploreRingCap) + complete := !callees.Truncated && !callees.BudgetHit && !callees.LowerBound && projectionComplete + return direct, complete + }, + ) targets = append(targets[:len(artifactTargets):len(artifactTargets)], symbolTargets...) } } @@ -3020,21 +3070,26 @@ func (s *Server) handleExplore(ctx context.Context, req mcp.CallToolRequest) (*m // structured responses can reserve a source body without a broad read. targets = s.materializeExploreStructuralSource(ctx, task, targets, opts) + declarationScope := s.localizationNodeScopeWithTests(ctx, opts, false) if !localize { - return mcp.NewToolResultText(s.renderExplore(task, targets, budget)), nil + outlineProvider := newExploreTaskPageOutlineProvider( + ctx, eng.Reader(), task, declarationScope, + ) + return mcp.NewToolResultText(s.renderExploreTask(task, targets, budget, outlineProvider)), nil } + declarations := newLocalizationFileDeclarationCache(ctx, eng.Reader(), declarationScope) symbolTargets = targets[len(artifactTargets):] // An implementation-intent query expands abstract seeds into their // concrete implementors before terminality is judged, so the envelope // carries the code that changes and answer_ready can see it. if exploreImplementationIntent(task) { - symbolTargets = s.expandImplementationTargets(ctx, symbolTargets) + symbolTargets = s.expandImplementationTargets(ctx, symbolTargets, declarationScope) targets = append(targets[:len(artifactTargets):len(artifactTargets)], symbolTargets...) } else if queryClass == rerank.QueryClassConcept { // Concept answers prefer the owning type when several of its members // rank together; implementation-intent queries are exempt because // they ask for exactly those members. - symbolTargets = preserveExploreDivergentDefaultOrder(s.foldMemberOwners(ctx, symbolTargets)) + symbolTargets = preserveExploreDivergentDefaultOrder(s.foldMemberOwners(ctx, symbolTargets, declarationScope)) // Owner folding is weaker than a unique source-literal callsite whose // callee was resolved and hydrated. Re-promote that proof after folding // so terminality is judged against the same strongest evidence that the @@ -3050,7 +3105,7 @@ func (s *Server) handleExplore(ctx context.Context, req mcp.CallToolRequest) (*m answerReady := exploreAnswerReady(task, symbolTargets) || artifactLane.ready if answerReady && !artifactLane.ready { eng := s.engineFor(ctx) - if eng != nil && exploreImplementationAnswerBlocked(task, symbolTargets, eng.GetOutEdges, eng.GetSymbol) { + if eng != nil && exploreImplementationAnswerBlocked(ctx, task, symbolTargets, eng.Reader(), declarationScope) { // Only abstract declarations in evidence for an implementation // question: stay nonterminal so the permitted refinement read // can reach the concrete side. @@ -3067,14 +3122,19 @@ func (s *Server) handleExplore(ctx context.Context, req mcp.CallToolRequest) (*m return s.hydrateExploreLeadingFileDepthTarget(ctx, eng, node, ringOpts) }, ) + if index, hit, ok := sourceWindowHits.elect(pageTargets); ok { + if pageTargets[index].node != nil { + pageTargets[index].sourceWindow = s.localizationSourceWindowForHit(ctx, hit, pageTargets[index].node.ID) + } + } targets = append(targets[:len(artifactTargets):len(artifactTargets)], pageTargets...) // The same leading file, indexed rather than ranked. The enumeration is // deferred: only a page that stays non-terminal pays for it, and it pays // once however many envelopes this request packs. The task's own terms ride // along so a bounded index keeps the declarations the task named. - pageOutline := localizationPageOutlineProvider( + pageOutline := boundedLocalizationPageOutlineProvider( localizationRankedPool, pageTargets, exploreTerminalTerms(task), - func(file string) []*graph.Node { return fileDefinitionNodes(eng, file) }, + declarations.outlineDefinitions, ) // File evidence can make localization answer-ready, but it never becomes a // synthetic exact-symbol read. Exact reads remain declaration-only. @@ -3108,8 +3168,8 @@ func (s *Server) handleExplore(ctx context.Context, req mcp.CallToolRequest) (*m preferredSymbol := explorePreferredRoutedRefinementSymbol( preferred, symbolTargets, routes, ) - result, refinement, boundedRoutes, digest := buildLocalizationRefinementResultForTaskWithOutline( - preferredSymbol, task, targets, budget, routes, pageOutline, + result, refinement, boundedRoutes, digest := buildLocalizationRefinementResultForTaskWithOutlineAndDeclarations( + preferredSymbol, task, targets, budget, routes, pageOutline, declarations, ) if refinement.State != localizationStateNeedsRefinement { refinement.digest = digest @@ -3131,8 +3191,8 @@ func (s *Server) handleExplore(ctx context.Context, req mcp.CallToolRequest) (*m // and is retained for post-terminal replay — for the exact-read contract // too, whose success promotes to answer_ready with the evidence already // stashed. - result, _, digest, completion := buildLocalizationExploreResultForTaskFinalizedWithOutline( - completion, task, targets, budget, pageOutline, + result, _, digest, completion := buildLocalizationExploreResultForTaskFinalizedWithOutlineAndDeclarations( + completion, task, targets, budget, pageOutline, declarations, ) // Literal-driven terminality must show its evidence: when the verdict // rests on a quoted-literal match but the budgeted envelope shed the @@ -3145,8 +3205,8 @@ func (s *Server) handleExplore(ctx context.Context, req mcp.CallToolRequest) (*m preferredSymbol := explorePreferredRoutedRefinementSymbol( explorePreferredRefinementSymbol(task, symbolTargets), symbolTargets, routes, ) - refined, refinement, boundedRoutes, refinedDigest := buildLocalizationRefinementResultForTaskWithOutline( - preferredSymbol, task, targets, budget, routes, pageOutline, + refined, refinement, boundedRoutes, refinedDigest := buildLocalizationRefinementResultForTaskWithOutlineAndDeclarations( + preferredSymbol, task, targets, budget, routes, pageOutline, declarations, ) if refinement.State != localizationStateNeedsRefinement { refinement.digest = refinedDigest @@ -3182,8 +3242,9 @@ type localizationExploreEnvelope struct { // Outline indexes the page's leading file; Outlines indexes the page's // further files, deepest first. The split keeps the leading file where // consumers already read it. - Outline *localizationFileOutline `json:"outline,omitempty"` - Outlines []*localizationFileOutline `json:"outlines,omitempty"` + Outline *localizationFileOutline `json:"outline,omitempty"` + Outlines []*localizationFileOutline `json:"outlines,omitempty"` + SourceWindow *localizationSourceWindow `json:"source_window,omitempty"` } type localizationEvidence struct { @@ -3200,6 +3261,18 @@ type localizationEvidence struct { Callees []string `json:"callees,omitempty"` Provenance string `json:"provenance,omitempty"` Source string `json:"source,omitempty"` + + // Request-local presentation authority. These fields are deliberately not + // serialized: provenance remains truthful while supplemental rows and the + // frozen initial PRIMARY cohort keep their distinct presentation roles. + literalPrimaryEligible bool + taskCitedPrimaryEligible bool + primaryCohortOrder int + supportingOnly bool + // leadingFileDepth separates the two supportingOnly writers: a row the + // initial page admitted for leading-file completeness may still earn a + // task-named seat, while a post-terminal supplemental row never does. + leadingFileDepth bool } func (s *Server) completeEmptyLocalization(ctx context.Context, task string, budget int) *mcp.CallToolResult { @@ -3542,6 +3615,64 @@ func prioritizeLocalizationEvidenceTarget(requiredID string, targets []exploreTa return ordered } +// localizationSecondaryTaskCitedTarget reserves one independently cited +// identity already present in the bounded target window. Citation offsets used +// by the explicit target or exact source-range owners are consumed: one authored +// mention cannot seat several overloads that share the same bare name. The +// caller only reorders presentation evidence; no graph read, provenance, or +// completion authority is added. +func localizationSecondaryTaskCitedTarget( + task string, + targets, selected []exploreTarget, +) (exploreTarget, bool) { + if strings.TrimSpace(task) == "" || len(targets) == 0 { + return exploreTarget{}, false + } + selectedIDs := make(map[string]struct{}, len(selected)) + selectedNames := make(map[string]struct{}, len(selected)) + consumedOffsets := make(map[int]struct{}, len(selected)) + for _, target := range selected { + if target.node == nil { + continue + } + selectedIDs[exploreDraftNodeKey(target.node)] = struct{}{} + if name := strings.TrimSpace(target.node.Name); name != "" { + selectedNames[name] = struct{}{} + } + if offset := localizationDirectAdjacencyNodeTaskCitationOffset(task, target.node); offset >= 0 { + consumedOffsets[offset] = struct{}{} + } + } + bestIndex, bestOffset := -1, -1 + for index, target := range targets { + if target.node == nil { + continue + } + if _, selected := selectedIDs[exploreDraftNodeKey(target.node)]; selected { + continue + } + offset := localizationDirectAdjacencyNodeTaskCitationOffset(task, target.node) + if offset < 0 { + continue + } + name := strings.TrimSpace(target.node.Name) + if _, represented := selectedNames[name]; represented && + localizationTaskExactIdentifierOffset(task, name) == offset { + continue + } + if _, consumed := consumedOffsets[offset]; consumed { + continue + } + if bestIndex < 0 || offset < bestOffset { + bestIndex, bestOffset = index, offset + } + } + if bestIndex < 0 { + return exploreTarget{}, false + } + return targets[bestIndex], true +} + func localizationEvidenceTargetsFromDraft(task, exactID string, targets []exploreTarget, draft []exploreDraftEntry) []exploreTarget { if len(targets) == 0 { return targets @@ -3580,6 +3711,14 @@ func localizationEvidenceTargetsFromDraft(task, exactID string, targets []explor } } } + // An exact task path+line citation authenticates the enclosing declaration + // as presentation evidence. Reserve every already-bounded range owner before + // generic draft rows, while leaving completion authority unchanged. + for _, target := range targets { + if target.sourceRange { + appendTarget(target) + } + } // A graph-proven causal constructor and its owning type outrank the // downstream retrieval seed. Their explicit admission metadata survives // draft ranking, owner folding, and byte-budget packing, so this ordering is @@ -3608,6 +3747,33 @@ func localizationEvidenceTargetsFromDraft(task, exactID string, targets []explor } } } + // One further concrete identifier that the task independently cites may + // precede generic draft ranking. Graph-proven causal pairs, the retrieval + // head, and any prescribed exact symbol have already retained their order. + // Offsets represented above are consumed so one authored mention cannot + // reserve several same-name overloads. + if target, ok := localizationSecondaryTaskCitedTarget(task, targets, selected); ok { + appendTarget(target) + } + // Unanchored content evidence may occupy at most two PRIMARY seats. Candidate + // selection already bounds this lane; order its survivors by the number of + // distinct task literals they corroborate before the ordinary draft fills + // the projection. + literalPrimaries := make([]exploreTarget, 0, exploreSourceLiteralReservationMax) + for _, target := range targets { + if target.literalPrimaryEligible { + literalPrimaries = append(literalPrimaries, target) + } + } + sort.SliceStable(literalPrimaries, func(i, j int) bool { + return literalPrimaries[i].literalMatchCount > literalPrimaries[j].literalMatchCount + }) + for index, target := range literalPrimaries { + if index == exploreSourceLiteralReservationMax { + break + } + appendTarget(target) + } // A source-body literal is the only direct evidence that can identify an // implementation absent from symbol metadata. Reserve the strongest one // before draft promotion and byte-budget packing can consume every slot. @@ -3781,7 +3947,7 @@ func interleaveLocalizationDirectRelationsWithRoutes( if index < localizationDirectEvidenceReserve || target.node.ID == requiredID || target.causalChangeBridge || target.causalChangeLeaf || target.causalChangeOwner || target.divergentDefaultOwner || target.divergentDefaultType || target.conceptImplementation || target.conceptComplement || - target.exactContent || target.sourceLiteral || target.typedAnchorProjection || + target.sourceRange || target.exactContent || target.sourceLiteral || target.typedAnchorProjection || // Depth rows already paid for the tail they occupy; a relationship // row must not reclaim the same slot a second time. target.leadingFileDepth { @@ -3789,7 +3955,7 @@ func interleaveLocalizationDirectRelationsWithRoutes( } if target.node.ID == requiredID || target.causalChangeBridge || target.causalChangeLeaf || target.causalChangeOwner || target.divergentDefaultOwner || target.divergentDefaultType || - target.conceptImplementation || target.conceptComplement || + target.conceptImplementation || target.conceptComplement || target.sourceRange || target.exactContent || target.sourceLiteral || target.typedAnchorProjection { orderedPrefix = index + 1 } @@ -4052,6 +4218,19 @@ func buildLocalizationRefinementResultForTaskWithOutline( budget int, routes map[string]localizationRefinementRoute, outline func() *localizationPageOutline, +) (*mcp.CallToolResult, localizationCompletion, map[string]localizationRefinementRoute, *localizationEvidenceDigest) { + return buildLocalizationRefinementResultForTaskWithOutlineAndDeclarations( + preferredSymbol, task, targets, budget, routes, outline, nil, + ) +} + +func buildLocalizationRefinementResultForTaskWithOutlineAndDeclarations( + preferredSymbol, task string, + targets []exploreTarget, + budget int, + routes map[string]localizationRefinementRoute, + outline func() *localizationPageOutline, + declarations *localizationFileDeclarationCache, ) (*mcp.CallToolResult, localizationCompletion, map[string]localizationRefinementRoute, *localizationEvidenceDigest) { choosePreferred := func(symbols []string, requested string) (string, []string, map[string]localizationRefinementRoute) { authorized, bounded := boundedLocalizationRefinementRoutes(symbols, routes, requested) @@ -4075,8 +4254,8 @@ func buildLocalizationRefinementResultForTaskWithOutline( preferredSymbol, preauthorized, prebounded := choosePreferred(candidateSymbols, preferredSymbol) if preferredSymbol == "" { advisory := newLocalizationCompletion(true, "") - result, _, digest, packedCompletion := buildLocalizationExploreResultForTaskFinalizedWithOutline( - advisory, task, targets, budget, outline, + result, _, digest, packedCompletion := buildLocalizationExploreResultForTaskFinalizedWithOutlineAndDeclarations( + advisory, task, targets, budget, outline, declarations, ) return result, packedCompletion, nil, digest } @@ -4087,8 +4266,8 @@ func buildLocalizationRefinementResultForTaskWithOutline( budgetCompletion := newLocalizationRefinementCompletionForSymbols(preferredSymbol, preauthorized) budgetCompletion.refinementRoutes = prebounded var finalRoutes map[string]localizationRefinementRoute - result, _, digest, packedCompletion := buildLocalizationExploreResultForTaskFinalizedWithOutline( - budgetCompletion, task, targets, budget, outline, + result, _, digest, packedCompletion := buildLocalizationExploreResultForTaskFinalizedWithOutlineAndDeclarations( + budgetCompletion, task, targets, budget, outline, declarations, func(packed localizationExploreEnvelope) localizationCompletion { packedPreferred, allowedSymbols, bounded := choosePreferred(packed.Symbols, preferredSymbol) if packedPreferred == "" { @@ -4125,6 +4304,20 @@ func buildLocalizationExploreResultForTaskFinalizedWithOutline( budget int, outline func() *localizationPageOutline, finalize ...localizationCompletionFinalizer, +) (*mcp.CallToolResult, []string, *localizationEvidenceDigest, localizationCompletion) { + return buildLocalizationExploreResultForTaskFinalizedWithOutlineAndDeclarations( + completion, task, targets, budget, outline, nil, finalize..., + ) +} + +func buildLocalizationExploreResultForTaskFinalizedWithOutlineAndDeclarations( + completion localizationCompletion, + task string, + targets []exploreTarget, + budget int, + outline func() *localizationPageOutline, + declarations *localizationFileDeclarationCache, + finalize ...localizationCompletionFinalizer, ) (*mcp.CallToolResult, []string, *localizationEvidenceDigest, localizationCompletion) { var draft []exploreDraftEntry if strings.TrimSpace(task) != "" { @@ -4212,6 +4405,19 @@ func buildLocalizationExploreResultForTaskFinalizedWithOutline( } } } + literalMandatory := 0 + for index, target := range targets { + if !target.literalPrimaryEligible && !localizationStrongSourceLiteralCallee(target) { + continue + } + if index+1 > mandatoryCount { + mandatoryCount = index + 1 + } + literalMandatory++ + if literalMandatory == localizationFinalResponseLiteralReserve { + break + } + } for index, target := range targets { if target.sourceLiteral && index+1 > mandatoryCount { mandatoryCount = index + 1 @@ -4250,6 +4456,10 @@ func buildLocalizationExploreResultForTaskFinalizedWithOutline( Callers: boundedLocalizationNeighborIDs(target.callers, localizationMaxNeighborIDs), Callees: boundedLocalizationNeighborIDs(target.callees, localizationMaxNeighborIDs), Provenance: provenance, + + literalPrimaryEligible: target.literalPrimaryEligible || localizationStrongSourceLiteralCallee(target), + supportingOnly: target.leadingFileDepth, + leadingFileDepth: target.leadingFileDepth, } candidate := envelope @@ -4418,6 +4628,9 @@ func buildLocalizationExploreResultForTaskFinalizedWithOutline( // index. A trade that ends with no index bought nothing, so it is undone. var untraded []localizationEvidence for !localizationEnvelopeFits(envelope, shedBudget) { + if localizationShedSourceWindow(&envelope) { + continue + } block := &localizationPageOutline{Leading: envelope.Outline, Others: envelope.Outlines} if !block.empty() { // The outlines are a convenience over rows the caller can still @@ -4425,13 +4638,17 @@ func buildLocalizationExploreResultForTaskFinalizedWithOutline( // here — but they give way by degrees. A shorter index is worth far // more than none, and only pressure past the floor drops one. if block.atFloor() { - // The index has given back everything it can and the page still - // does not fit. A ranked page fills its budget with rows, so - // without this the floor is unreachable exactly on the pages - // that need an index most. The expendable breadth tail pays, - // in the detail a caller can re-derive from the identity that - // stays — never in a row, and never inside the reserve the - // refinement contract may name. + // A rank-two-or-later file at its floor is less valuable than + // independently useful evidence detail. Drop that expendable + // breadth before stripping a row; the leading pair stays protected. + if block.dropUnprotectedFloorFile() { + envelope.Outline, envelope.Outlines = block.Leading, block.Others + continue + } + // The protected index has given back everything it can and the + // page still does not fit. The expendable evidence tail now pays, + // in detail a caller can re-derive from the identity that stays — + // never in a row or refinement reserve. if untraded == nil { untraded = append([]localizationEvidence(nil), envelope.Evidence...) } @@ -4493,6 +4710,20 @@ func buildLocalizationExploreResultForTaskFinalizedWithOutline( } envelope.Evidence[shed].Source = "" } + for _, target := range acceptedTargets { + if target.sourceWindow == nil { + continue + } + envelope = localizationEnvelopePackingSourceWindow(envelope, target.sourceWindow, maxBytes) + break + } + // Freeze the ranked answer before adjacency/body/depth supplementation can + // rebuild the digest and accidentally promote a newly materialized row. + freezeLocalizationPrimaryCohort(task, &envelope, digest) + if declarations != nil { + envelope, digest = promoteLocalizationDirectAdjacency(task, envelope, declarations.reader, len(envelope.Evidence), shedBudget, digest) + } + envelope, digest = promoteLocalizationBodyMentions(task, envelope, declarations, shedBudget, digest) body, err := json.Marshal(envelope) if err != nil { return mcp.NewToolResultError("encode localization result: " + err.Error()), nil, nil, envelope.Completion @@ -5116,16 +5347,20 @@ func exploreLiteralWordRune(r rune) bool { // snippets are unhighlighted, so a bounded case-folded substring scan with // Unicode-aware word boundaries is sufficient and does not read source again. func exploreTextHasExactLiteral(text, literal string) bool { - literal = strings.TrimSpace(literal) - if text == "" || literal == "" { + return exploreLowerTextHasExactLiteral(strings.ToLower(text), literal) +} + +// exploreLowerTextHasExactLiteral is the allocation-aware sibling for bounded +// scans that test many declaration names against one already-lowercased body. +func exploreLowerTextHasExactLiteral(lowerText, literal string) bool { + literal = strings.ToLower(strings.TrimSpace(literal)) + if lowerText == "" || literal == "" { return false } - text = strings.ToLower(text) - literal = strings.ToLower(literal) first, _ := utf8.DecodeRuneInString(literal) last, _ := utf8.DecodeLastRuneInString(literal) - for offset := 0; offset <= len(text)-len(literal); { - relative := strings.Index(text[offset:], literal) + for offset := 0; offset <= len(lowerText)-len(literal); { + relative := strings.Index(lowerText[offset:], literal) if relative < 0 { return false } @@ -5133,12 +5368,12 @@ func exploreTextHasExactLiteral(text, literal string) bool { end := start + len(literal) beforeOK := true if start > 0 && exploreLiteralWordRune(first) { - previous, _ := utf8.DecodeLastRuneInString(text[:start]) + previous, _ := utf8.DecodeLastRuneInString(lowerText[:start]) beforeOK = !exploreLiteralWordRune(previous) } afterOK := true - if end < len(text) && exploreLiteralWordRune(last) { - next, _ := utf8.DecodeRuneInString(text[end:]) + if end < len(lowerText) && exploreLiteralWordRune(last) { + next, _ := utf8.DecodeRuneInString(lowerText[end:]) afterOK = !exploreLiteralWordRune(next) } if beforeOK && afterOK { @@ -5217,12 +5452,11 @@ const ( exploreSourceLiteralCoverageSignal = "explore_source_literal_coverage" exploreSourceLiteralTaskAlignSignal = "explore_source_literal_task_alignment" exploreSourceLiteralReservationMax = 2 - exploreQuotedRecallMaxTerms = 3 + exploreQuotedRecallMaxTerms = 6 // Literals mined from code blocks may only take the slots prose left - // unclaimed, up to this total, so a task with no code block searches - // exactly as many terms as before. - exploreQuotedRecallMaxMinedTerms = 5 - exploreQuotedRecallMaxPerTerm = 12 + // unclaimed. Explicit and mined anchors share one fixed request budget. + exploreQuotedRecallMaxMinedTerms = 6 + exploreQuotedRecallMaxPerTerm = 24 exploreQuotedRecallRetryMaxRows = 24 ) @@ -5404,16 +5638,12 @@ func exploreQuotedRecallHasExactSourceNode( if exploreDraftIsTestNode(node) && !exploreQueryHasTestIntent(task) { return false } - // The explicit-anchor short-circuit answers "is this node named by the - // request", not "is this term already covered". A compact value — a locale, - // protocol, or configuration code — is never a declaration identity, so a - // candidate that merely matches some other anchor of the same request must - // not mark that value as covered: doing so suppresses the one bounded lane - // able to find where the value is registered. Declaration text below still - // settles compact coverage exactly as before. - if !exploreQuotedRecallCompactTerms(terms) && exploreLocalizationExplicitAnchor(task, node) { - return true - } + // Coverage means the term is visible in this node's own declaration text — + // name, signature, or qualified name. That a candidate happens to be some + // other anchor of the same request says nothing about where a quoted value + // lives; treating it as coverage suppressed the one bounded lane able to + // find the value's registration site on exactly the tasks that both quote + // a literal and name a file. retrieval := node.RetrievalMetadata() for _, term := range terms { // Compact values such as locale and protocol codes are usually source @@ -5486,8 +5716,8 @@ func exploreSourceLiteralConstructionIntent(task string) bool { return false } -// gatherExploreQuotedContentCandidates performs at most four bounded content -// searches (three literals plus one adaptive retry). It scans source bodies only +// gatherExploreQuotedContentCandidates performs at most seven bounded content +// searches (six literals plus one adaptive retry). It scans source bodies only // when neither the ordinary nor content channels already contain an exact, // localizable code symbol. The fallback is request-local and never persists a // source-body index. @@ -5498,15 +5728,55 @@ func (s *Server) gatherExploreQuotedContentCandidates( limit int, scope query.QueryOptions, ) []*rerank.Candidate { - if s == nil || s.graph == nil || ctx.Err() != nil { + return s.gatherExploreQuotedContentCandidatesCollecting(ctx, task, ordinary, limit, scope, nil) +} + +func (s *Server) gatherExploreQuotedContentCandidatesCollecting( + ctx context.Context, + task string, + ordinary []*rerank.Candidate, + limit int, + scope query.QueryOptions, + collector *localizationSourceWindowHitCollector, +) []*rerank.Candidate { + return s.gatherExploreContentCandidatesForTermsCollecting( + ctx, task, exploreQuotedRecallTerms(task), ordinary, limit, scope, collector, + ) +} + +// gatherExploreContentCandidatesForTerms is the shared bounded content path for +// authored quoted literals and the separately-gated inferred bare lane. +func (s *Server) gatherExploreContentCandidatesForTerms( + ctx context.Context, + task string, + terms []string, + ordinary []*rerank.Candidate, + limit int, + scope query.QueryOptions, +) []*rerank.Candidate { + return s.gatherExploreContentCandidatesForTermsCollecting(ctx, task, terms, ordinary, limit, scope, nil) +} + +func (s *Server) gatherExploreContentCandidatesForTermsCollecting( + ctx context.Context, + task string, + terms []string, + ordinary []*rerank.Candidate, + limit int, + scope query.QueryOptions, + collector *localizationSourceWindowHitCollector, +) []*rerank.Candidate { + if s == nil || ctx.Err() != nil { return nil } - content, hasContent := s.graph.(graph.ContentSearcher) - terms := exploreQuotedRecallTerms(task) - if len(terms) == 0 { + reader := s.readerFor(ctx) + if reader == nil { return nil } - perTerm := clampInt(limit/4, 4, exploreQuotedRecallMaxPerTerm) + // Durable content FTS remains the discovery backend. Overlay-owned files + // are filtered below before any snippet can authenticate a candidate. + content, hasContent := s.graph.(graph.ContentSearcher) + perTerm := clampInt(limit/3, 4, exploreQuotedRecallMaxPerTerm) repoPrefix := "" if len(scope.RepoAllow) == 1 { for prefix, allowed := range scope.RepoAllow { @@ -5518,6 +5788,13 @@ func (s *Server) gatherExploreQuotedContentCandidates( if repoPrefix == "" { repoPrefix, _ = s.sessionLocality(ctx) } + // A task that quotes nothing still names its subject: identifier-shaped + // prose tokens feed the same bounded source-literal recall below, so an + // unquoted issue is not blind to source bodies. + bareTokens := exploreDistinctiveBareTokens(task, repoPrefix) + if len(terms) == 0 && len(bareTokens) == 0 { + return nil + } type recallPage struct { term string @@ -5536,7 +5813,9 @@ func (s *Server) gatherExploreQuotedContentCandidates( if err != nil { continue } - pages = append(pages, recallPage{term: term, hits: hits, saturated: len(hits) >= perTerm}) + saturated := len(hits) >= perTerm + hits = s.filterOverlayContentHits(ctx, hits) + pages = append(pages, recallPage{term: term, hits: hits, saturated: saturated}) } // A short or collision-heavy literal can fill the first page before its @@ -5570,8 +5849,8 @@ func (s *Server) gatherExploreQuotedContentCandidates( } if retry >= 0 && ctx.Err() == nil { if hits, err := content.SearchContent(pages[retry].term, repoPrefix, exploreQuotedRecallRetryMaxRows); err == nil { - pages[retry].hits = hits pages[retry].saturated = len(hits) >= exploreQuotedRecallRetryMaxRows + pages[retry].hits = s.filterOverlayContentHits(ctx, hits) } } @@ -5587,6 +5866,7 @@ func (s *Server) gatherExploreQuotedContentCandidates( sourceLiteralSettled := make(map[string]bool) sourceLiteralCallee := make(map[string]bool) sourceLiteralTaskAligned := make(map[string]bool) + sourceLiteralCoordinates := make(map[string][]exploreSourceLiteralHit) for _, page := range pages { seenForTerm := make(map[string]struct{}, len(page.hits)) exactIDs := make(map[string]struct{}) @@ -5621,14 +5901,20 @@ func (s *Server) gatherExploreQuotedContentCandidates( } nodes := make(map[string]*graph.Node, len(order)) if len(order) > 0 { - nodes = s.graph.GetNodesByIDs(order) + nodes = reader.GetNodesByIDs(order) } // Decide source coverage per quoted term. An exact metadata hit for one // symbol-like term must not suppress a different compact value whose only // useful evidence is inside a registration body. The selected missing term // still feeds one bounded source scan, so this preserves the fixed I/O cap. - sourceRecallTerms := make([]string, 0, len(terms)) - for _, term := range terms { + sourceRecallTerms := make([]string, 0, len(terms)+len(bareTokens)) + seenRecallTerms := make(map[string]struct{}, len(terms)+len(bareTokens)) + appendRecallTerm := func(term string) { + key := strings.ToLower(term) + if _, duplicate := seenRecallTerms[key]; duplicate { + return + } + seenRecallTerms[key] = struct{}{} oneTerm := []string{term} exactSourceFound := exploreQuotedRecallHasExactSourceCandidate(task, oneTerm, ordinary, scope) if !exactSourceFound { @@ -5643,6 +5929,18 @@ func (s *Server) gatherExploreQuotedContentCandidates( sourceRecallTerms = append(sourceRecallTerms, term) } } + for _, term := range terms { + appendRecallTerm(term) + } + // Quoted values keep priority in the bounded term budget; mined tokens + // fill behind them and are subject to the same declaration-coverage test, + // so a token already visible in ranked metadata never triggers a scan. + for _, term := range bareTokens { + if len(sourceRecallTerms) >= exploreSourceLiteralRecallMaxTerms+2 { + break + } + appendRecallTerm(term) + } // content_fts stores content-class nodes rather than ordinary source bodies. // An exact document hit therefore does not prove that the source declaration @@ -5652,6 +5950,7 @@ func (s *Server) gatherExploreQuotedContentCandidates( sourceRecall := s.gatherExploreSourceLiteralRecall(ctx, sourceRecallTerms, repoPrefix, scope) missingNodes := make([]string, 0, len(sourceRecall.hits)) for _, hit := range sourceRecall.hits { + sourceLiteralCoordinates[hit.nodeID] = append(sourceLiteralCoordinates[hit.nodeID], hit) if previous, exists := bestRank[hit.nodeID]; !exists { order = append(order, hit.nodeID) bestRank[hit.nodeID] = hit.rank @@ -5693,7 +5992,7 @@ func (s *Server) gatherExploreQuotedContentCandidates( } } if len(missingNodes) > 0 { - for id, node := range s.graph.GetNodesByIDs(missingNodes) { + for id, node := range reader.GetNodesByIDs(missingNodes) { nodes[id] = node } } @@ -5709,16 +6008,9 @@ func (s *Server) gatherExploreQuotedContentCandidates( calleeIDs = append(calleeIDs, id) } } - if len(calleeIDs) > 0 { - for id, edges := range s.graph.GetOutEdgesByNodeIDs(calleeIDs) { - for _, edge := range edges { - if edge != nil && edge.Kind == graph.EdgeInstantiates { - sourceLiteralTaskAligned[id] = true - break - } - } - } - } + sourceLiteralTaskAligned = projectExploreSourceLiteralConstructionAlignment( + ctx, reader, calleeIDs, + ) } } if len(order) == 0 { @@ -5742,6 +6034,7 @@ func (s *Server) gatherExploreQuotedContentCandidates( } } if sourceRank := sourceLiteralHit[id]; sourceRank > 0 { + collector.add(sourceLiteralCoordinates[id]...) signals[exploreSourceLiteralSignal] = sourceRank signals[exploreSourceLiteralCoverageSignal] = float64(len(sourceLiteralAnchors[id])) if sourceLiteralCallee[id] { @@ -6162,8 +6455,16 @@ func reserveExploreSourceLiteralCandidate(candidates, bounded []*rerank.Candidat sources := make([]*rerank.Candidate, 0, exploreSourceLiteralReservationMax) seenSourceIDs := make(map[string]struct{}, exploreSourceLiteralReservationMax) for _, candidate := range candidates { - if candidate == nil || candidate.Node == nil || candidate.Node.ID == "" || candidate.Signals == nil || - candidate.Signals[exploreSourceLiteralSignal] <= 0 { + if candidate == nil || candidate.Node == nil || candidate.Node.ID == "" || candidate.Signals == nil { + continue + } + // A code definition holding an exact content match is literal evidence + // with the same standing as a grep-lane owner: without a slot here it + // falls off the ranked cut on every non-concept query class. + literalEvidence := candidate.Signals[exploreSourceLiteralSignal] > 0 || + candidate.Signals[exploreContentRecallExactSignal] > 0 && + exploreCodeDefinitionKind(candidate.Node.Kind) + if !literalEvidence { continue } if _, duplicate := seenSourceIDs[candidate.Node.ID]; duplicate { @@ -6189,20 +6490,20 @@ func reserveExploreSourceLiteralCandidate(candidates, bounded []*rerank.Candidat if leftCoverage != rightCoverage { return leftCoverage > rightCoverage } - leftSettled := sources[i].Signals[exploreContentRecallAmbiguousSignal] <= 0 - rightSettled := sources[j].Signals[exploreContentRecallAmbiguousSignal] <= 0 - if leftSettled != rightSettled { - return leftSettled + leftAligned := sources[i].Signals[exploreSourceLiteralTaskAlignSignal] > 0 + rightAligned := sources[j].Signals[exploreSourceLiteralTaskAlignSignal] > 0 + if leftAligned != rightAligned { + return leftAligned } leftDirect := directProductionCallee(sources[i]) rightDirect := directProductionCallee(sources[j]) if leftDirect != rightDirect { return leftDirect } - leftAligned := sources[i].Signals[exploreSourceLiteralTaskAlignSignal] > 0 - rightAligned := sources[j].Signals[exploreSourceLiteralTaskAlignSignal] > 0 - if leftAligned != rightAligned { - return leftAligned + leftSettled := sources[i].Signals[exploreContentRecallAmbiguousSignal] <= 0 + rightSettled := sources[j].Signals[exploreContentRecallAmbiguousSignal] <= 0 + if leftSettled != rightSettled { + return leftSettled } leftRank := sources[i].Signals[exploreSourceLiteralSignal] rightRank := sources[j].Signals[exploreSourceLiteralSignal] @@ -6212,13 +6513,15 @@ func reserveExploreSourceLiteralCandidate(candidates, bounded []*rerank.Candidat return callableSpecificity(sources[i]) > callableSpecificity(sources[j]) }) selectedSources := sources[:0] + // A crowded source field is collision evidence: when more owners carry + // literal signals than the reservation could ever seat, a second + // single-anchor slot would spend the answer on one literal's noise. A + // lone extra settled owner is the opposite — a distinct proven site. + crowded := len(sources) > exploreSourceLiteralReservationMax for _, source := range sources { coverage := source.Signals[exploreSourceLiteralCoverageSignal] settled := source.Signals[exploreContentRecallAmbiguousSignal] <= 0 - // One ambiguous single-anchor owner may preserve recall, but allowing a - // second would spend most of a three-slot answer on collision evidence. - // Multi-anchor corroboration is strong enough to keep despite ambiguity. - if len(selectedSources) > 0 && coverage < 2 && !settled { + if len(selectedSources) > 0 && coverage < 2 && (!settled || crowded) { continue } selectedSources = append(selectedSources, source) diff --git a/internal/mcp/tools_explore_quoted_recall_test.go b/internal/mcp/tools_explore_quoted_recall_test.go index be452427..d68be61f 100644 --- a/internal/mcp/tools_explore_quoted_recall_test.go +++ b/internal/mcp/tools_explore_quoted_recall_test.go @@ -4,6 +4,7 @@ import ( "context" "fmt" "path/filepath" + "strings" "testing" "github.com/stretchr/testify/require" @@ -87,6 +88,93 @@ func candidateByID(candidates []*rerank.Candidate, id string) *rerank.Candidate return nil } +func TestExploreQuotedRecallTermsAdmitsSixExplicitLiterals(t *testing.T) { + terms := exploreQuotedRecallTerms(`find "alpha", "bravo", "charlie", "delta", "echo", "foxtrot", and "golf"`) + require.Equal(t, []string{"alpha", "bravo", "charlie", "delta", "echo", "foxtrot"}, terms) +} + +func TestGatherExploreQuotedContentCandidatesUsesBoundedWidePages(t *testing.T) { + terms := []string{"alpha", "bravo", "charlie", "delta", "echo", "foxtrot"} + hits := make(map[string][]graph.ContentHit, len(terms)) + for _, term := range terms { + hits[term] = quotedRecallHits(term, 1, 0) + } + server, store := newQuotedRecallCountingServer(t, hits) + candidates := server.gatherExploreQuotedContentCandidates( + context.Background(), `find "alpha", "bravo", "charlie", "delta", "echo", and "foxtrot"`, nil, 72, + query.QueryOptions{RepoAllow: map[string]bool{"demo": true}}, + ) + + require.Len(t, candidates, len(terms)) + require.Equal(t, []int{24, 24, 24, 24, 24, 24}, store.searchLimits) + require.Equal(t, 1, store.graphLookups, "all term pages must share one graph lookup") +} + +func TestGatherExploreContentCandidatesForBareTermsUsesSharedBounds(t *testing.T) { + hits := map[string][]graph.ContentHit{ + "DISKFULL": {{ + NodeID: "demo/storage.go::persist", FilePath: "demo/storage.go", + Snippet: `return errors.New("DISKFULL")`, + }}, + } + server, store := newQuotedRecallCountingServer(t, hits) + candidates := server.gatherExploreContentCandidatesForTerms( + context.Background(), "the daemon emits DISKFULL while persisting", []string{"DISKFULL"}, nil, 72, + query.QueryOptions{RepoAllow: map[string]bool{"demo": true}}, + ) + + require.Equal(t, []int{exploreQuotedRecallMaxPerTerm}, store.searchLimits) + require.Equal(t, 1, store.graphLookups) + exact := candidateByID(candidates, "demo/storage.go::persist") + require.NotNil(t, exact) + require.Equal(t, float64(1), exact.Signals[exploreContentRecallExactSignal]) +} + +func TestSourceWindowCollectionAddsNoContentSearchOrGraphLookup(t *testing.T) { + hits := map[string][]graph.ContentHit{ + "needle": {{ + NodeID: "demo/needle.go::candidate", FilePath: "demo/needle.go", + Snippet: `register("needle")`, + }}, + } + ordinary := []*rerank.Candidate{{Node: &graph.Node{ + ID: "demo/owner.go::needle", Name: "needle", Kind: graph.KindFunction, + FilePath: "demo/owner.go", RepoPrefix: "demo", + }}} + baselineServer, baselineStore := newQuotedRecallCountingServer(t, hits) + baseline := baselineServer.gatherExploreQuotedContentCandidates( + context.Background(), `find "needle"`, ordinary, 12, query.QueryOptions{}, + ) + collectingServer, collectingStore := newQuotedRecallCountingServer(t, hits) + collector := &localizationSourceWindowHitCollector{} + collected := collectingServer.gatherExploreQuotedContentCandidatesCollecting( + context.Background(), `find "needle"`, ordinary, 12, query.QueryOptions{}, collector, + ) + + require.Equal(t, baselineStore.searchLimits, collectingStore.searchLimits) + require.Equal(t, baselineStore.searchRows, collectingStore.searchRows) + require.Equal(t, baselineStore.graphLookups, collectingStore.graphLookups) + require.Equal(t, len(baseline), len(collected)) + require.Empty(t, collector.hits, "content metadata without an existing trigram coordinate cannot fabricate a window") +} + +func TestGatherExploreContentCandidatesForBareTermsCapsSearchesAndHydratesOnce(t *testing.T) { + terms := []string{"ALPHA", "BRAVO", "CHARLIE", "DELTA", "ECHO"} + hits := make(map[string][]graph.ContentHit, len(terms)) + for _, term := range terms { + hits[term] = quotedRecallHits(term, exploreQuotedRecallMaxPerTerm, -1) + } + server, store := newQuotedRecallCountingServer(t, hits) + candidates := server.gatherExploreContentCandidatesForTerms( + context.Background(), strings.Join(terms, " "), terms, nil, 72, + query.QueryOptions{RepoAllow: map[string]bool{"demo": true}}, + ) + + require.NotEmpty(t, candidates) + require.LessOrEqual(t, len(store.searchLimits), exploreBareLiteralMaxTerms+1) + require.Equal(t, 1, store.graphLookups, "all bare term pages must share one graph lookup") +} + func TestGatherExploreQuotedContentCandidatesRetriesOneSaturatedTermWithinBounds(t *testing.T) { hits := map[string][]graph.ContentHit{ "ku": quotedRecallHits("ku", exploreQuotedRecallRetryMaxRows, 17), @@ -100,8 +188,8 @@ func TestGatherExploreQuotedContentCandidatesRetriesOneSaturatedTermWithinBounds ) require.LessOrEqual(t, len(store.searchLimits), exploreQuotedRecallMaxTerms+1) - require.Equal(t, []int{5, 5, 5, exploreQuotedRecallRetryMaxRows}, store.searchLimits) - require.Equal(t, []int{5, 5, 5, exploreQuotedRecallRetryMaxRows}, store.searchRows) + require.Equal(t, []int{6, 6, 6, exploreQuotedRecallRetryMaxRows}, store.searchLimits) + require.Equal(t, []int{6, 6, 6, exploreQuotedRecallRetryMaxRows}, store.searchRows) require.LessOrEqual(t, store.searchRows[len(store.searchRows)-1], exploreQuotedRecallRetryMaxRows) require.Equal(t, 1, store.graphLookups, "all final pages must share one graph lookup") @@ -126,7 +214,7 @@ func TestGatherExploreQuotedContentCandidatesKeepsUniqueExactFastPath(t *testing query.QueryOptions{RepoAllow: map[string]bool{"demo": true}}, ) - require.Equal(t, []int{4}, store.searchLimits) + require.Equal(t, []int{5}, store.searchLimits) require.Equal(t, 1, store.graphLookups) exact := candidateByID(candidates, "demo/exact.go::candidate") require.NotNil(t, exact) diff --git a/internal/mcp/tools_explore_source_preservation_test.go b/internal/mcp/tools_explore_source_preservation_test.go index 8e6f9596..beb14e39 100644 --- a/internal/mcp/tools_explore_source_preservation_test.go +++ b/internal/mcp/tools_explore_source_preservation_test.go @@ -117,6 +117,23 @@ func TestSelectFinalExploreCandidatesPrefersTaskAlignedAmbiguousCallee(t *testin require.Nil(t, candidateByID(selected, third.Node.ID), "three-way collision noise must stay outside the bounded answer") } +func TestSelectFinalExploreCandidatesPrefersGraphResolvedCalleeWithinOneLiteralSeat(t *testing.T) { + owner := sourcePreservationCandidate("settled-owner", 4, 1) + owner.Signals[exploreSourceLiteralCoverageSignal] = 1 + callee := sourcePreservationCandidate("ambiguous-callee", 9, 0.5) + callee.Signals[exploreSourceLiteralCoverageSignal] = 1 + callee.Signals[exploreContentRecallAmbiguousSignal] = 1 + callee.Signals[exploreSourceLiteralCalleeSignal] = 1 + selected := selectFinalExploreCandidates([]*rerank.Candidate{ + sourcePreservationCandidate("semantic-head", 0, 0), owner, callee, + }, nil, 2) + + require.Len(t, selected, 2) + require.Equal(t, "semantic-head", selected[0].Node.ID) + require.Equal(t, callee.Node.ID, selected[1].Node.ID) + require.Nil(t, candidateByID(selected, owner.Node.ID)) +} + func TestSelectFinalExploreCandidatesDeduplicatesSourceOwnersAndHonorsSmallLimits(t *testing.T) { head := sourcePreservationCandidate("head", 0, 0) source := sourcePreservationCandidate("aligned-source", 8, 0.5) @@ -299,6 +316,24 @@ func TestExploreAnswerReadyKeepsQuotedNonExactConceptNonTerminal(t *testing.T) { require.False(t, exploreAnswerReady(`locate locale registry pipeline for "ku"`, []exploreTarget{head})) } +func TestLocalizationEvidenceOrdersEligibleLiteralRowsByDistinctMatches(t *testing.T) { + semantic := exploreTarget{node: &graph.Node{ID: "repo/semantic.go::Semantic", Name: "Semantic", FilePath: "repo/semantic.go"}} + low := exploreTarget{ + node: &graph.Node{ID: "repo/low.go::Low", Name: "Low", FilePath: "repo/low.go"}, + sourceLiteral: true, literalPrimaryEligible: true, literalMatchCount: 1, + } + high := exploreTarget{ + node: &graph.Node{ID: "repo/high.go::High", Name: "High", FilePath: "repo/high.go"}, + sourceLiteral: true, literalPrimaryEligible: true, literalMatchCount: 3, + } + + ordered := localizationEvidenceTargetsFromDraft("find the handler", "", []exploreTarget{semantic, low, high}, nil) + require.Len(t, ordered, 3) + require.Equal(t, semantic.node.ID, ordered[0].node.ID) + require.Equal(t, high.node.ID, ordered[1].node.ID) + require.Equal(t, low.node.ID, ordered[2].node.ID) +} + func BenchmarkLimitExploreCandidatesPreservingSourceLiteral80(b *testing.B) { candidates := make([]*rerank.Candidate, 0, 80) for i := 0; i < 79; i++ { diff --git a/internal/mcp/tools_overlay_branch.go b/internal/mcp/tools_overlay_branch.go index c17487c0..dc4587fa 100644 --- a/internal/mcp/tools_overlay_branch.go +++ b/internal/mcp/tools_overlay_branch.go @@ -413,7 +413,18 @@ func (s *Server) handleCompareBranches(ctx context.Context, req mcp.CallToolRequ limit := int(req.GetFloat("limit", 50)) opts := query.QueryOptions{Depth: depth, Limit: limit, Detail: "brief"} - aFiles, errA := s.overlays.FilesForBranch(id, a) + if err := ctx.Err(); err != nil { + return nil, err + } + aFiles, errA := s.overlays.FilesForBranchBounded( + id, + a, + overlayRequestSnapshotMaxFiles, + overlayRequestSnapshotMaxBytes, + ) + if err := ctx.Err(); err != nil { + return nil, err + } if errA != nil { switch { case errors.Is(errA, daemon.ErrSessionNotFound): @@ -424,7 +435,18 @@ func (s *Server) handleCompareBranches(ctx context.Context, req mcp.CallToolRequ return mcp.NewToolResultError(errA.Error()), nil } } - bFiles, errB := s.overlays.FilesForBranch(id, b) + if err := ctx.Err(); err != nil { + return nil, err + } + bFiles, errB := s.overlays.FilesForBranchBounded( + id, + b, + overlayRequestSnapshotMaxFiles, + overlayRequestSnapshotMaxBytes, + ) + if err := ctx.Err(); err != nil { + return nil, err + } if errB != nil { switch { case errors.Is(errB, daemon.ErrSessionNotFound): @@ -436,12 +458,18 @@ func (s *Server) handleCompareBranches(ctx context.Context, req mcp.CallToolRequ } } - aIDs, aPaths, err := s.runBranchQuery(aFiles, kind, symID, opts) + aIDs, aPaths, err := s.runBranchQuery(ctx, aFiles, kind, symID, opts) if err != nil { + if ctxErr := requestContextError(ctx, err); ctxErr != nil { + return nil, ctxErr + } return mcp.NewToolResultError(fmt.Sprintf("branch %q: %v", a, err)), nil } - bIDs, bPaths, err := s.runBranchQuery(bFiles, kind, symID, opts) + bIDs, bPaths, err := s.runBranchQuery(ctx, bFiles, kind, symID, opts) if err != nil { + if ctxErr := requestContextError(ctx, err); ctxErr != nil { + return nil, ctxErr + } return mcp.NewToolResultError(fmt.Sprintf("branch %q: %v", b, err)), nil } @@ -479,11 +507,11 @@ func (s *Server) handleCompareBranches(ctx context.Context, req mcp.CallToolRequ // branch file set and runs the query kind against it. Returns the // stable-sorted result IDs and the overlay path coverage. Empty // branches fall through to the base engine. -func (s *Server) runBranchQuery(files []daemon.OverlayFile, kind, symID string, opts query.QueryOptions) ([]string, []string, error) { +func (s *Server) runBranchQuery(ctx context.Context, files []daemon.OverlayFile, kind, symID string, opts query.QueryOptions) ([]string, []string, error) { eng := s.engine paths := []string{} if len(files) > 0 { - layer, p, err := s.constructOverlayLayer(files) + layer, p, err := s.constructOverlayLayer(ctx, files) if err != nil { return nil, nil, err } diff --git a/internal/mcp/tools_overlay_diff.go b/internal/mcp/tools_overlay_diff.go index 082d1859..d1964f70 100644 --- a/internal/mcp/tools_overlay_diff.go +++ b/internal/mcp/tools_overlay_diff.go @@ -50,8 +50,11 @@ func (s *Server) handleCompareWithOverlay(ctx context.Context, req mcp.CallToolR if SessionIDFromContext(ctx) == "" { return mcp.NewToolResultError("compare_with_overlay requires an MCP session; connect via the daemon or set Mcp-Session-Id"), nil } - view, viewErr := s.buildOverlayViewForCtx(ctx) + ctx, view, viewErr := s.prepareOverlayRequest(ctx) if viewErr != nil { + if ctxErr := requestContextError(ctx, viewErr); ctxErr != nil { + return nil, ctxErr + } return mcp.NewToolResultError(viewErr.Error()), nil } if view == nil { diff --git a/internal/mcp/tools_pr_review_context.go b/internal/mcp/tools_pr_review_context.go index a994a35e..36984dae 100644 --- a/internal/mcp/tools_pr_review_context.go +++ b/internal/mcp/tools_pr_review_context.go @@ -277,7 +277,10 @@ func (s *Server) handlePRReviewContext(ctx context.Context, req mcp.CallToolRequ // --- section: simulate_chain (gated on an explicit overlay session) --- if wantSimulate { - sim, gate := s.buildPRReviewSimulation(ctx, req) + sim, gate, simErr := s.buildPRReviewSimulation(ctx, req) + if simErr != nil { + return nil, simErr + } out.Simulation = sim out.Gates = append(out.Gates, gate) } @@ -431,7 +434,7 @@ func (s *Server) buildConfigAuditSection(repoRoot string) *audit.Report { // (via the `session_id` param or the request context) the section is omitted // with a note rather than run against a phantom session. The base graph is // never mutated regardless. -func (s *Server) buildPRReviewSimulation(ctx context.Context, req mcp.CallToolRequest) (*prReviewSimulation, reviewGate) { +func (s *Server) buildPRReviewSimulation(ctx context.Context, req mcp.CallToolRequest) (*prReviewSimulation, reviewGate, error) { editsArg := strings.TrimSpace(req.GetString("edits", "")) if editsArg == "" { return &prReviewSimulation{ @@ -440,7 +443,7 @@ func (s *Server) buildPRReviewSimulation(ctx context.Context, req mcp.CallToolRe }, reviewGate{ Name: "simulate_chain", Status: prReviewPass, Detail: "skipped: no edits supplied", - } + }, nil } sessionID := strings.TrimSpace(req.GetString("session_id", "")) @@ -454,7 +457,7 @@ func (s *Server) buildPRReviewSimulation(ctx context.Context, req mcp.CallToolRe }, reviewGate{ Name: "simulate_chain", Status: prReviewPass, Detail: "skipped: no overlay session id", - } + }, nil } edits, err := parsePRReviewEdits(editsArg) @@ -465,7 +468,7 @@ func (s *Server) buildPRReviewSimulation(ctx context.Context, req mcp.CallToolRe }, reviewGate{ Name: "simulate_chain", Status: prReviewWarn, Detail: "invalid edits: " + err.Error(), - } + }, nil } // Run the chain on top of the named session's overlay. buildSimulation @@ -474,13 +477,16 @@ func (s *Server) buildPRReviewSimulation(ctx context.Context, req mcp.CallToolRe simCtx := WithSessionID(ctx, sessionID) sim, simErr := s.buildSimulation(simCtx, edits, true) if simErr != nil { + if ctxErr := requestContextError(simCtx, simErr); ctxErr != nil { + return nil, reviewGate{}, ctxErr + } return &prReviewSimulation{ Ran: false, GraphUntouched: true, SessionID: sessionID, Note: "simulation failed: " + simErr.Error(), }, reviewGate{ Name: "simulate_chain", Status: prReviewWarn, Detail: "simulation failed: " + simErr.Error(), - } + }, nil } steps := make([]map[string]any, 0, len(sim.steps)) @@ -512,12 +518,15 @@ func (s *Server) buildPRReviewSimulation(ctx context.Context, req mcp.CallToolRe status = prReviewBlock detail = detailf("%d step(s) introduce broken callers / implementors", len(sim.steps)) } - return out, reviewGate{Name: "simulate_chain", Status: status, Detail: detail} + return out, reviewGate{Name: "simulate_chain", Status: status, Detail: detail}, nil } // parsePRReviewEdits parses the `edits` JSON array into WorkspaceEdits, // reusing the same per-edit parser the simulate_chain handler uses. func parsePRReviewEdits(raw string) ([]lsp.WorkspaceEdit, error) { + if len(raw) > overlaySimulationInputMaxBytes { + return nil, fmt.Errorf("edits input exceeds limit %d bytes", overlaySimulationInputMaxBytes) + } var rawEdits []json.RawMessage if err := json.Unmarshal([]byte(raw), &rawEdits); err != nil { return nil, errors.New("edits must be a JSON array of WorkspaceEdit objects: " + err.Error()) @@ -525,6 +534,9 @@ func parsePRReviewEdits(raw string) ([]lsp.WorkspaceEdit, error) { if len(rawEdits) == 0 { return nil, errors.New("edits array is empty") } + if len(rawEdits) > overlaySimulationMaxSteps { + return nil, fmt.Errorf("edits exceed limit %d", overlaySimulationMaxSteps) + } edits := make([]lsp.WorkspaceEdit, 0, len(rawEdits)) for i, re := range rawEdits { edit, err := parseWorkspaceEdit(string(re)) diff --git a/internal/mcp/tools_review.go b/internal/mcp/tools_review.go index e1c60990..5b5196f5 100644 --- a/internal/mcp/tools_review.go +++ b/internal/mcp/tools_review.go @@ -641,21 +641,11 @@ func (s *Server) reviewRulepackMatches(ctx context.Context, changedFiles []strin return nil } - fileSymbols := s.buildFileSymbolIndex(targets) - lookup := func(graphPath string, line int) (string, string) { - idx := fileSymbols[graphPath] - if idx == nil { - return "", "" - } - return idx.find(line) - } - var collected []astquery.Match for _, d := range bundle { res, runErr := astquery.Run(ctx, astquery.Options{ Detector: d.Name, Targets: targets, - SymbolLookup: lookup, Resolver: astquery.DefaultLanguageResolver, Limit: 5000, ExcludeTests: true, @@ -668,8 +658,10 @@ func (s *Server) reviewRulepackMatches(ctx context.Context, changedFiles []strin // Graph-grounding post-pass: drop the N+1 / check-then-act rows the resolved // call / loop metadata refutes. This is the same FP-reduction the analyze - // review path applies. Grounding keys off the match's symbol id, so the - // path is only rewritten once the rows that survive are known. + // review path applies. Grounding keys off the match's symbol id, so bounded + // post-match enrichment must precede it; path rewriting still happens only + // after the surviving rows are known. + s.enrichASTMatchesContext(ctx, collected) kept := review.GroundReviewMatches(s.graph, collected) for i := range kept { kept[i].File = reviewRepoRelPath(kept[i].File, repoPrefix) @@ -1174,7 +1166,10 @@ func (s *Server) handleReviewPack(ctx context.Context, req mcp.CallToolRequest) // Cost bound: run a speculative preview_edit for the high-risk (d=1-heavy) // changed symbols only — never the whole changeset. - previews := s.highRiskPreviews(ctx, diff, impact) + previews, previewErr := s.highRiskPreviews(ctx, diff, impact) + if previewErr != nil { + return nil, previewErr + } // Privacy-safe risk receipt over the whole changeset. scrub := requestBoolDefault(req, "scrub", false) @@ -1298,12 +1293,15 @@ func (s *Server) reviewTestTargets(ctx context.Context, ids []string) []string { // so the broken-callers / impact rollup is computed without touching disk. Only // the high-risk subset is simulated, so the pass never scales with the whole // changeset. -func (s *Server) highRiskPreviews(ctx context.Context, diff *analysis.DiffResult, impact map[string]*analysis.ImpactResult) []reviewPreview { +func (s *Server) highRiskPreviews(ctx context.Context, diff *analysis.DiffResult, impact map[string]*analysis.ImpactResult) ([]reviewPreview, error) { if diff == nil || impact == nil { - return nil + return nil, nil } var previews []reviewPreview for _, cs := range diff.ChangedSymbols { + if err := ctx.Err(); err != nil { + return nil, err + } ir := impact[cs.ID] if ir == nil || len(ir.ByDepth[1]) < highRiskD1Threshold { continue @@ -1313,7 +1311,13 @@ func (s *Server) highRiskPreviews(ctx context.Context, diff *analysis.DiffResult continue } sim, err := s.buildSimulation(ctx, []lsp.WorkspaceEdit{edit}, false) - if err != nil || len(sim.steps) == 0 { + if err != nil { + if ctxErr := requestContextError(ctx, err); ctxErr != nil { + return nil, ctxErr + } + continue + } + if len(sim.steps) == 0 { continue } step := sim.steps[0] @@ -1325,7 +1329,7 @@ func (s *Server) highRiskPreviews(ctx context.Context, diff *analysis.DiffResult }) } sort.SliceStable(previews, func(i, j int) bool { return previews[i].SymbolID < previews[j].SymbolID }) - return previews + return previews, nil } // identityEditForSymbol builds a no-op WorkspaceEdit that rewrites a symbol's own diff --git a/internal/mcp/tools_review_rulepack_test.go b/internal/mcp/tools_review_rulepack_test.go index 34f78b1b..121b57f6 100644 --- a/internal/mcp/tools_review_rulepack_test.go +++ b/internal/mcp/tools_review_rulepack_test.go @@ -99,6 +99,8 @@ func TestReviewRulepackMatches_JoinsRepoRelativeChangedFiles(t *testing.T) { // Findings travel onward to rule resolution, the risk ranking, and the // forge comment API — all of which speak repo-relative paths. for _, m := range matches { + require.NotEmpty(t, m.SymbolID, + "review grounding must receive post-match enclosing-symbol enrichment") require.Equal(t, "pkg/widget.go", m.File, "match paths must be repo-relative, not graph-prefixed") } diff --git a/internal/mcp/tools_search_text.go b/internal/mcp/tools_search_text.go index f1e39c33..c7b96672 100644 --- a/internal/mcp/tools_search_text.go +++ b/internal/mcp/tools_search_text.go @@ -109,8 +109,8 @@ func (s *Server) handleSearchText(ctx context.Context, req mcp.CallToolRequest) matches = limitTextMatches(matches, limit) } - enriched := s.enrichTextMatches(matches) - s.captureLocalizationSearchText(ctx, enriched) + enriched, fileIndexes := s.enrichTextMatchesContext(ctx, matches, queryOptionsForResolvedScope(resolved)) + s.captureLocalizationSearchText(ctx, enriched, fileIndexes) resp := map[string]any{ "query": query, "matches": enriched, @@ -228,35 +228,58 @@ func graphMatchPathKey(path string, repoPrefixed bool) string { return filepath.FromSlash(path) } -// enrichTextMatches decorates every trigram match with its enclosing -// graph symbol. It builds one per-file symbol index for the set of -// matched files, then resolves each match's line through it. -// -// The index is queried under both path spellings: file nodes are fetched -// and keyed by the graph's own FilePath, which on Windows is not the -// forward-slash path the match carries. -func (s *Server) enrichTextMatches(matches []trigram.Match) []enrichedTextMatch { +// enrichTextMatchesContext decorates every trigram match with its enclosing +// graph symbol through the bounded file projection, and returns the same file +// indexes used for enrichment so localization evidence capture never repeats +// the file scan. Storage receives the effective request/session scope; exact +// and Windows path spellings share the request-wide 4096-node budget. +func (s *Server) enrichTextMatchesContext( + ctx context.Context, + matches []trigram.Match, + opts query.QueryOptions, +) ([]enrichedTextMatch, map[string]*fileSymbolIndex) { out := make([]enrichedTextMatch, 0, len(matches)) - paths := make(map[string]struct{}, len(matches)) - for _, m := range matches { - paths[m.Path] = struct{}{} - if key := graphMatchPathKey(m.Path, true); key != m.Path { - paths[key] = struct{}{} + exactPaths := make([]string, 0, len(matches)) + aliasPaths := make([]string, 0, len(matches)) + exactSeen := make(map[string]struct{}, len(matches)) + aliasSeen := make(map[string]struct{}, len(matches)) + for _, match := range matches { + if _, duplicate := exactSeen[match.Path]; !duplicate { + exactSeen[match.Path] = struct{}{} + exactPaths = append(exactPaths, match.Path) } - } - idx := s.buildFileSymbolIndexForPaths(paths) - for _, m := range matches { - em := enrichedTextMatch{Path: m.Path, Line: m.Line, Text: m.Text} - fi := idx[m.Path] - if fi == nil { - if key := graphMatchPathKey(m.Path, true); key != m.Path { - fi = idx[key] + if alias := graphMatchPathKey(match.Path, true); alias != match.Path { + if _, duplicate := aliasSeen[alias]; !duplicate { + aliasSeen[alias] = struct{}{} + aliasPaths = append(aliasPaths, alias) } } - if fi != nil { - em.SymbolID, em.SymbolName = fi.find(m.Line) + } + orderedPaths := make([]string, 0, len(exactPaths)+len(aliasPaths)) + orderedPaths = append(orderedPaths, exactPaths...) + for _, alias := range aliasPaths { + if _, isExact := exactSeen[alias]; !isExact { + orderedPaths = append(orderedPaths, alias) } - out = append(out, em) } - return out + indexes := s.buildFileSymbolIndexForOrderedPathsScopedContext(ctx, orderedPaths, opts) + for _, match := range matches { + enriched := enrichedTextMatch{Path: match.Path, Line: match.Line, Text: match.Text} + index := fileSymbolIndexForPath(indexes, match.Path) + if index != nil { + enriched.SymbolID, enriched.SymbolName = index.find(match.Line) + } + out = append(out, enriched) + } + return out, indexes +} + +func fileSymbolIndexForPath(indexes map[string]*fileSymbolIndex, path string) *fileSymbolIndex { + if index := indexes[path]; index != nil { + return index + } + if key := graphMatchPathKey(path, true); key != path { + return indexes[key] + } + return nil } diff --git a/internal/mcp/tools_simulate.go b/internal/mcp/tools_simulate.go index 98ab9970..945bdb9e 100644 --- a/internal/mcp/tools_simulate.go +++ b/internal/mcp/tools_simulate.go @@ -108,6 +108,9 @@ func (s *Server) handlePreviewEdit(ctx context.Context, req mcp.CallToolRequest) sim, simErr := s.buildSimulation(ctx, []lsp.WorkspaceEdit{edit}, inherit) if simErr != nil { + if ctxErr := requestContextError(ctx, simErr); ctxErr != nil { + return nil, ctxErr + } return mcp.NewToolResultError(simErr.Error()), nil } step := sim.steps[0] @@ -149,6 +152,9 @@ func (s *Server) handleSimulateChain(ctx context.Context, req mcp.CallToolReques if err != nil { return mcp.NewToolResultError(err.Error()), nil } + if len(rawSteps) > overlaySimulationInputMaxBytes { + return mcp.NewToolResultError(fmt.Sprintf("steps input exceeds limit %d bytes", overlaySimulationInputMaxBytes)), nil + } var rawEdits []json.RawMessage if err := json.Unmarshal([]byte(rawSteps), &rawEdits); err != nil { return mcp.NewToolResultError("steps must be a JSON array of WorkspaceEdit objects: " + err.Error()), nil @@ -156,6 +162,9 @@ func (s *Server) handleSimulateChain(ctx context.Context, req mcp.CallToolReques if len(rawEdits) == 0 { return mcp.NewToolResultError("steps array is empty — pass at least one WorkspaceEdit"), nil } + if len(rawEdits) > overlaySimulationMaxSteps { + return mcp.NewToolResultError(fmt.Sprintf("steps exceed limit %d", overlaySimulationMaxSteps)), nil + } edits := make([]lsp.WorkspaceEdit, 0, len(rawEdits)) for i, raw := range rawEdits { edit, parseErr := parseWorkspaceEdit(string(raw)) @@ -176,6 +185,9 @@ func (s *Server) handleSimulateChain(ctx context.Context, req mcp.CallToolReques sim, simErr := s.buildSimulation(ctx, edits, inherit) if simErr != nil { + if ctxErr := requestContextError(ctx, simErr); ctxErr != nil { + return nil, ctxErr + } return mcp.NewToolResultError(simErr.Error()), nil } @@ -301,13 +313,36 @@ type simulationStep struct { // contents from prior steps' snapshots and never persists them // unless the caller asks for `keep`. func (s *Server) buildSimulation(ctx context.Context, edits []lsp.WorkspaceEdit, inherit bool) (*simulation, error) { + if ctx == nil { + ctx = context.Background() + } + if err := ctx.Err(); err != nil { + return nil, err + } + if len(edits) > overlaySimulationMaxSteps { + return nil, fmt.Errorf("simulation steps exceed limit %d", overlaySimulationMaxSteps) + } sim := &simulation{} current := map[string]daemon.OverlayFile{} if inherit { if sessID := SessionIDFromContext(ctx); sessID != "" && s.overlays != nil && s.overlays.Has(sessID) { - _, files, err := s.overlays.SnapshotFor(sessID) - if err == nil { + if err := ctx.Err(); err != nil { + return nil, err + } + _, files, err := s.overlays.SnapshotForBounded( + sessID, + overlayRequestSnapshotMaxFiles, + overlayRequestSnapshotMaxBytes, + ) + if ctxErr := ctx.Err(); ctxErr != nil { + return nil, ctxErr + } + if err != nil { + if !errors.Is(err, daemon.ErrSessionNotFound) { + return nil, fmt.Errorf("inherit overlay snapshot: %w", err) + } + } else { for _, f := range files { current[filepath.Clean(f.Path)] = f } @@ -392,7 +427,11 @@ func (s *Server) buildSimulation(ctx context.Context, edits []lsp.WorkspaceEdit, sort.Strings(step.deletedFiles) sort.Strings(step.missingFiles) - // 2. Snapshot the overlay state at this step. + // 2. Snapshot the overlay state at this step. Reject an oversized map + // before allocating or retaining its sorted slice representation. + if err := validateOverlayBuildMapEnvelope(current); err != nil { + return nil, fmt.Errorf("step %d: overlay snapshot: %w", stepIdx, err) + } snap := make([]daemon.OverlayFile, 0, len(current)) for _, f := range current { snap = append(snap, f) @@ -403,7 +442,7 @@ func (s *Server) buildSimulation(ctx context.Context, edits []lsp.WorkspaceEdit, // 3. Compute graph impact for this step: build the layer, // diff vs base, surface broken callers / implementors, // rank test targets. - layer, _, layerErr := s.constructOverlayLayer(snap) + layer, _, layerErr := s.constructOverlayLayer(ctx, snap) if layerErr != nil { return nil, fmt.Errorf("step %d: overlay parse: %w", stepIdx, layerErr) } @@ -1117,6 +1156,9 @@ func (s *Server) persistSimulationOverlay(ctx context.Context, sim *simulation) // `{uri: [TextEdit,...]}`) or `documentChanges` (modern form) — both // are valid LSP shapes. func parseWorkspaceEdit(raw string) (lsp.WorkspaceEdit, error) { + if len(raw) > overlaySimulationInputMaxBytes { + return lsp.WorkspaceEdit{}, fmt.Errorf("workspace_edit input exceeds limit %d bytes", overlaySimulationInputMaxBytes) + } raw = strings.TrimSpace(raw) if raw == "" { return lsp.WorkspaceEdit{}, errors.New("workspace_edit is empty") diff --git a/internal/mcp/tools_symbols_for_ranges.go b/internal/mcp/tools_symbols_for_ranges.go index a7e0d5be..359feac4 100644 --- a/internal/mcp/tools_symbols_for_ranges.go +++ b/internal/mcp/tools_symbols_for_ranges.go @@ -42,14 +42,26 @@ type rangeSpecJSON struct { EndLine int `json:"end_line"` } -// lowerRanges resolves each (file, range) spec to the symbols that enclose any -// line in the range, deduplicated across all specs by symbol ID. It returns -// the hits plus the display paths of any files that could not be resolved or -// carry no indexed symbols, so the caller can report partial coverage rather -// than silently dropping them. -func (s *Server) lowerRanges(specs []rangeSpec) ([]rangeSymbolHit, []string) { +// loweredRanges holds the symbols each (file, range) spec resolved to, +// deduplicated across all specs by symbol ID, plus the display paths of any +// files that could not be resolved or carry no indexed symbols, so the caller +// can report partial coverage rather than silently dropping them. +type loweredRanges struct { + hits []rangeSymbolHit + unresolved []string + saturated []string +} + +func (s *Server) lowerRangesContext(ctx context.Context, specs []rangeSpec) ([]rangeSymbolHit, []string) { + lowered := s.lowerRangesDetailedContext(ctx, specs) + unresolved := append(append([]string(nil), lowered.unresolved...), lowered.saturated...) + sort.Strings(unresolved) + return lowered.hits, dedupeStrings(unresolved) +} + +func (s *Server) lowerRangesDetailedContext(ctx context.Context, specs []rangeSpec) loweredRanges { if s == nil || s.graph == nil || len(specs) == 0 { - return nil, nil + return loweredRanges{} } byGraphPath := make(map[string][]rangeSpec) displayOf := make(map[string]string) @@ -65,19 +77,25 @@ func (s *Server) lowerRanges(specs []rangeSpec) ([]rangeSymbolHit, []string) { displayOf[gp] = relPath } if len(byGraphPath) == 0 { - return nil, unresolved + sort.Strings(unresolved) + return loweredRanges{unresolved: dedupeStrings(unresolved)} } want := make(map[string]struct{}, len(byGraphPath)) for gp := range byGraphPath { want[gp] = struct{}{} } - indexes := s.buildFileSymbolIndexForPaths(want) + indexes := s.buildFileSymbolIndexForPathsContext(ctx, want) seen := make(map[string]struct{}) var hits []rangeSymbolHit + var saturated []string for gp, specsForFile := range byGraphPath { idx := indexes[gp] + if idx != nil && idx.saturated { + saturated = append(saturated, displayOf[gp]) + continue + } if idx == nil { unresolved = append(unresolved, displayOf[gp]) continue @@ -106,7 +124,10 @@ func (s *Server) lowerRanges(specs []rangeSpec) ([]rangeSymbolHit, []string) { return hits[a].StartLine < hits[b].StartLine }) sort.Strings(unresolved) - return hits, dedupeStrings(unresolved) + sort.Strings(saturated) + return loweredRanges{ + hits: hits, unresolved: dedupeStrings(unresolved), saturated: dedupeStrings(saturated), + } } // parseRangeSpecs reads range specs from a request. Two forms are accepted: @@ -164,7 +185,7 @@ func (s *Server) handleSymbolsForRanges(ctx context.Context, req mcp.CallToolReq if err != nil { return mcp.NewToolResultError(err.Error()), nil } - hits, unresolved := s.lowerRanges(specs) + hits, unresolved := s.lowerRangesContext(ctx, specs) if hits == nil { hits = []rangeSymbolHit{} } diff --git a/internal/parser/extraction_limits.go b/internal/parser/extraction_limits.go new file mode 100644 index 00000000..5c549e53 --- /dev/null +++ b/internal/parser/extraction_limits.go @@ -0,0 +1,106 @@ +package parser + +import ( + "context" + "errors" + "fmt" +) + +const ( + defaultExtractionMaxNodes = 4096 + defaultExtractionMaxEdges = 16384 + defaultExtractionMaxStdoutBytes = 16 << 20 + defaultExtractionMaxResultBytes = 16 << 20 + defaultExtractionMaxStderrBytes = 64 << 10 +) + +// ExtractionLimits bounds one extractor invocation. MaxNodes includes the +// always-emitted file node. A zero field is a strict zero budget; callers that +// want the finite defaults start from DefaultExtractionLimits. +type ExtractionLimits struct { + MaxNodes int + MaxEdges int + // MaxStdoutBytes bounds plugin wire output. + MaxStdoutBytes int + // MaxResultBytes bounds the combined retained wire output and cumulative + // synthesized node-ID bytes. + MaxResultBytes int + // MaxStderrBytes bounds retained diagnostics only. Zero retains none; + // excess is discarded and does not fail an otherwise successful extraction. + MaxStderrBytes int +} + +// DefaultExtractionLimits returns the finite hard envelope used by legacy +// extraction. Bounded callers may narrow any field, including to zero. +func DefaultExtractionLimits() ExtractionLimits { + return ExtractionLimits{ + MaxNodes: defaultExtractionMaxNodes, + MaxEdges: defaultExtractionMaxEdges, + MaxStdoutBytes: defaultExtractionMaxStdoutBytes, + MaxResultBytes: defaultExtractionMaxResultBytes, + MaxStderrBytes: defaultExtractionMaxStderrBytes, + } +} + +// Clamped restricts every caller-provided field to its finite hard maximum. +// Negative values fail closed to zero; positive values can never widen the +// process envelope beyond DefaultExtractionLimits. +func (limits ExtractionLimits) Clamped() ExtractionLimits { + defaults := DefaultExtractionLimits() + limits.MaxNodes = clampExtractionLimit(limits.MaxNodes, defaults.MaxNodes) + limits.MaxEdges = clampExtractionLimit(limits.MaxEdges, defaults.MaxEdges) + limits.MaxStdoutBytes = clampExtractionLimit(limits.MaxStdoutBytes, defaults.MaxStdoutBytes) + limits.MaxResultBytes = clampExtractionLimit(limits.MaxResultBytes, defaults.MaxResultBytes) + limits.MaxStderrBytes = clampExtractionLimit(limits.MaxStderrBytes, defaults.MaxStderrBytes) + return limits +} + +func clampExtractionLimit(value, hardMax int) int { + if value < 0 { + return 0 + } + if value > hardMax { + return hardMax + } + return value +} + +// ErrExtractionLimit classifies every hard extraction-envelope failure. +var ErrExtractionLimit = errors.New("extraction limit exceeded") + +// ExtractionLimitError reports that an extractor crossed a caller-visible +// hard envelope. Callers must discard any partial result on this error. +type ExtractionLimitError struct { + Resource string + Limit int64 + Observed int64 +} + +func (e *ExtractionLimitError) Error() string { + if e == nil { + return ErrExtractionLimit.Error() + } + return fmt.Sprintf("%s: %s observed %d exceeds limit %d", ErrExtractionLimit, e.Resource, e.Observed, e.Limit) +} + +func (e *ExtractionLimitError) Unwrap() error { return ErrExtractionLimit } + +// ExtractionUsage reports conservative charges for one successful bounded +// extraction. RawNodes includes the mandatory file node and every inspected +// plugin node row; RawEdges includes every inspected edge row. Counts and +// ResultBytes are cumulative across duplicate fields (no refund). ResultBytes +// includes exact StdoutBytes plus synthesized node IDs. Failures report zero. +type ExtractionUsage struct { + RawNodes int + RawEdges int + StdoutBytes int64 + ResultBytes int64 +} + +// BoundedExtractor is an optional Extractor capability for request-scoped +// cancellation and hard output/result envelopes. Legacy callers continue to +// use Extract unchanged. +type BoundedExtractor interface { + Extractor + ExtractBounded(context.Context, string, []byte, ExtractionLimits) (*ExtractionResult, ExtractionUsage, error) +} diff --git a/internal/parser/languages/extractor_plugin.go b/internal/parser/languages/extractor_plugin.go index f4cadc1a..175e1a52 100644 --- a/internal/parser/languages/extractor_plugin.go +++ b/internal/parser/languages/extractor_plugin.go @@ -4,6 +4,9 @@ import ( "bytes" "context" "encoding/json" + "errors" + "fmt" + "io" "os" "os/exec" "strings" @@ -22,7 +25,10 @@ import ( // spec doesn't set TimeoutMs. Generous enough for a real extractor pass // over a single file, small enough that a hung plugin never stalls the // whole index. -const defaultExtractorPluginTimeout = 5000 * time.Millisecond +const ( + defaultExtractorPluginTimeout = 5000 * time.Millisecond + extractorPluginWaitDelay = 250 * time.Millisecond +) // pluginNode mirrors one entry of the plugin's JSON `nodes` array. The // field names are the documented wire shape; meta is free-form. @@ -49,8 +55,314 @@ type pluginEdge struct { // pluginDocument is the full JSON document a plugin writes to stdout. type pluginDocument struct { - Nodes []pluginNode `json:"nodes"` - Edges []pluginEdge `json:"edges"` + Nodes []pluginNode `json:"nodes"` + Edges []pluginEdge `json:"edges"` + rawNodes int + rawEdges int +} + +const pluginCancellationPollMask = 127 + +type cappedPluginBuffer struct { + buffer bytes.Buffer + limit int + truncated bool + onOverflow func() + overflowOnce sync.Once +} + +func (b *cappedPluginBuffer) Write(p []byte) (int, error) { + original := len(p) + remaining := b.limit - b.buffer.Len() + if remaining > 0 { + if remaining > original { + remaining = original + } + _, _ = b.buffer.Write(p[:remaining]) + } + if remaining < original { + b.truncated = true + b.overflowOnce.Do(func() { + if b.onOverflow != nil { + b.onOverflow() + } + }) + } + return original, nil +} + +func (b *cappedPluginBuffer) String() string { return b.buffer.String() } + +type pluginResultBudget struct { + limit int64 + observed int64 +} + +func newPluginResultBudget(limit, stdoutBytes int64) (*pluginResultBudget, error) { + budget := &pluginResultBudget{limit: limit} + if stdoutBytes < 0 || stdoutBytes > limit { + return nil, &parser.ExtractionLimitError{ + Resource: "result_bytes", + Limit: limit, + Observed: limit + 1, + } + } + budget.observed = stdoutBytes + return budget, nil +} + +func (b *pluginResultBudget) reserveSynthesizedID(filePath, name string) error { + if b == nil { + return nil + } + if err := b.reserveSynthesizedBytes(int64(len(filePath))); err != nil { + return err + } + if err := b.reserveSynthesizedBytes(2); err != nil { + return err + } + return b.reserveSynthesizedBytes(int64(len(name))) +} + +func (b *pluginResultBudget) reserveSynthesizedBytes(size int64) error { + if size < 0 || b.observed > b.limit || size > b.limit-b.observed { + return &parser.ExtractionLimitError{ + Resource: "result_bytes", + Limit: b.limit, + Observed: b.limit + 1, + } + } + b.observed += size + return nil +} + +func pollPluginContext(ctx context.Context, position int) error { + if position&pluginCancellationPollMask != 0 { + return nil + } + return ctx.Err() +} + +func decodePluginDocument(ctx context.Context, reader io.Reader, filePath string, limits parser.ExtractionLimits, budget *pluginResultBudget) (pluginDocument, error) { + if err := ctx.Err(); err != nil { + return pluginDocument{}, err + } + readLimit := int64(limits.MaxStdoutBytes) + 1 + limited := &io.LimitedReader{R: reader, N: readLimit} + decoder := json.NewDecoder(limited) + document, decodeErr := decodePluginObject(ctx, decoder, filePath, limits, budget) + if decodeErr != nil { + if err := ctx.Err(); err != nil { + return pluginDocument{}, err + } + observed := readLimit - limited.N + if observed > int64(limits.MaxStdoutBytes) { + return pluginDocument{}, stdoutLimitError(limits.MaxStdoutBytes) + } + return pluginDocument{}, decodeErr + } + + token, trailingErr := decoder.Token() + observed := readLimit - limited.N + if observed > int64(limits.MaxStdoutBytes) { + return pluginDocument{}, stdoutLimitError(limits.MaxStdoutBytes) + } + if err := ctx.Err(); err != nil { + return pluginDocument{}, err + } + if trailingErr != io.EOF { + if trailingErr != nil { + return pluginDocument{}, trailingErr + } + return pluginDocument{}, fmt.Errorf("unexpected trailing JSON token %v", token) + } + return document, nil +} + +func stdoutLimitError(limit int) error { + return &parser.ExtractionLimitError{ + Resource: "stdout_bytes", + Limit: int64(limit), + Observed: int64(limit) + 1, + } +} + +func decodePluginObject(ctx context.Context, decoder *json.Decoder, filePath string, limits parser.ExtractionLimits, budget *pluginResultBudget) (pluginDocument, error) { + opening, err := decoder.Token() + if err != nil { + return pluginDocument{}, err + } + if opening == nil { + return pluginDocument{}, nil + } + if delimiter, ok := opening.(json.Delim); !ok || delimiter != '{' { + return pluginDocument{}, fmt.Errorf("plugin output must be a JSON object or null") + } + + var document pluginDocument + rawNodes, rawEdges, fields := 0, 0, 0 + for decoder.More() { + if err := pollPluginContext(ctx, fields); err != nil { + return pluginDocument{}, err + } + fields++ + keyToken, err := decoder.Token() + if err != nil { + return pluginDocument{}, err + } + key, ok := keyToken.(string) + if !ok { + return pluginDocument{}, fmt.Errorf("plugin output contains a non-string field name") + } + switch { + case strings.EqualFold(key, "nodes"): + document.Nodes, err = decodePluginNodes(ctx, decoder, filePath, limits.MaxNodes, &rawNodes, budget) + case strings.EqualFold(key, "edges"): + document.Edges, err = decodePluginEdges(ctx, decoder, limits.MaxEdges, &rawEdges) + default: + err = skipPluginJSONValue(ctx, decoder) + } + if err != nil { + return pluginDocument{}, err + } + } + closing, err := decoder.Token() + if err != nil { + return pluginDocument{}, err + } + if delimiter, ok := closing.(json.Delim); !ok || delimiter != '}' { + return pluginDocument{}, fmt.Errorf("plugin output object is not closed") + } + document.rawNodes = rawNodes + document.rawEdges = rawEdges + return document, ctx.Err() +} + +func decodePluginNodes(ctx context.Context, decoder *json.Decoder, filePath string, maxTotal int, rawCount *int, budget *pluginResultBudget) ([]pluginNode, error) { + opening, err := decoder.Token() + if err != nil { + return nil, err + } + if opening == nil { + return nil, nil + } + if delimiter, ok := opening.(json.Delim); !ok || delimiter != '[' { + return nil, fmt.Errorf("plugin nodes must be an array or null") + } + + nodes := make([]pluginNode, 0) + for decoder.More() { + if err := pollPluginContext(ctx, *rawCount); err != nil { + return nil, err + } + if *rawCount >= maxTotal-1 { + return nil, &parser.ExtractionLimitError{ + Resource: "nodes", + Limit: int64(maxTotal), + Observed: int64(*rawCount) + 2, + } + } + (*rawCount)++ + var node pluginNode + if err := decoder.Decode(&node); err != nil { + return nil, err + } + if graph.ValidNodeKind(graph.NodeKind(strings.TrimSpace(node.Kind))) && strings.TrimSpace(node.ID) == "" { + fp := node.FilePath + if fp == "" { + fp = filePath + } + if err := budget.reserveSynthesizedID(fp, node.Name); err != nil { + return nil, err + } + } + nodes = append(nodes, node) + } + closing, err := decoder.Token() + if err != nil { + return nil, err + } + if delimiter, ok := closing.(json.Delim); !ok || delimiter != ']' { + return nil, fmt.Errorf("plugin nodes array is not closed") + } + return nodes, ctx.Err() +} + +func decodePluginEdges(ctx context.Context, decoder *json.Decoder, maxTotal int, rawCount *int) ([]pluginEdge, error) { + opening, err := decoder.Token() + if err != nil { + return nil, err + } + if opening == nil { + return nil, nil + } + if delimiter, ok := opening.(json.Delim); !ok || delimiter != '[' { + return nil, fmt.Errorf("plugin edges must be an array or null") + } + + edges := make([]pluginEdge, 0) + for decoder.More() { + if err := pollPluginContext(ctx, *rawCount); err != nil { + return nil, err + } + if *rawCount >= maxTotal { + return nil, &parser.ExtractionLimitError{ + Resource: "edges", + Limit: int64(maxTotal), + Observed: int64(*rawCount) + 1, + } + } + (*rawCount)++ + var edge pluginEdge + if err := decoder.Decode(&edge); err != nil { + return nil, err + } + edges = append(edges, edge) + } + closing, err := decoder.Token() + if err != nil { + return nil, err + } + if delimiter, ok := closing.(json.Delim); !ok || delimiter != ']' { + return nil, fmt.Errorf("plugin edges array is not closed") + } + return edges, ctx.Err() +} + +func skipPluginJSONValue(ctx context.Context, decoder *json.Decoder) error { + if err := ctx.Err(); err != nil { + return err + } + token, err := decoder.Token() + if err != nil { + return err + } + delimiter, ok := token.(json.Delim) + if !ok || delimiter != '{' && delimiter != '[' { + return nil + } + depth, tokens := 1, 0 + for depth > 0 { + if err := pollPluginContext(ctx, tokens); err != nil { + return err + } + tokens++ + token, err = decoder.Token() + if err != nil { + return err + } + delimiter, ok = token.(json.Delim) + if !ok { + continue + } + switch delimiter { + case '{', '[': + depth++ + case '}', ']': + depth-- + } + } + return ctx.Err() } // SubprocessExtractor is a parser.Extractor backed by an external @@ -109,57 +421,132 @@ func (e *SubprocessExtractor) fileNode(filePath string, src []byte) *graph.Node } } -// Extract runs the plugin command over src and returns the file node -// plus every valid node/edge the plugin emitted. Any error path returns -// just the file node and a nil error — the index is never failed by a -// plugin. +// Extract preserves the legacy extractor contract: plugin failures, including +// bounded-output failures, degrade to a file-only result and never fail indexing. func (e *SubprocessExtractor) Extract(filePath string, src []byte) (*parser.ExtractionResult, error) { - result := &parser.ExtractionResult{} - fileNode := e.fileNode(filePath, src) - result.Nodes = append(result.Nodes, fileNode) + result, _, err := e.ExtractBounded(context.Background(), filePath, src, parser.DefaultExtractionLimits()) + if err == nil { + return result, nil + } + e.log.Warn("extractor plugin: bounded extraction failed — degraded to file node only", + zap.String("language", e.language), zap.String("file", filePath), zap.Error(err)) + return e.fileOnlyResult(filePath, src), nil +} + +// ExtractBounded runs the configured plugin with hard wire and result limits. +// It is an all-or-nothing capability: caller cancellation, limit violations, +// command failures, timeouts and malformed output return no partial result. +// The legacy Extract wrapper alone converts those errors to a file-only result. +func (e *SubprocessExtractor) ExtractBounded(ctx context.Context, filePath string, src []byte, requested parser.ExtractionLimits) (*parser.ExtractionResult, parser.ExtractionUsage, error) { + var noUsage parser.ExtractionUsage + if ctx == nil { + ctx = context.Background() + } + limits := requested.Clamped() + if err := ctx.Err(); err != nil { + return nil, noUsage, err + } + if limits.MaxNodes < 1 { + return nil, noUsage, &parser.ExtractionLimitError{Resource: "nodes", Limit: 0, Observed: 1} + } + result := e.fileOnlyResult(filePath, src) if e.command == "" { - return result, nil + return result, parser.ExtractionUsage{RawNodes: 1}, nil } - ctx, cancel := context.WithTimeout(context.Background(), e.timeout) - defer cancel() + runCtx, runCancel := context.WithTimeout(ctx, e.timeout) + defer runCancel() + stdout := &cappedPluginBuffer{limit: limits.MaxStdoutBytes, onOverflow: runCancel} + stderr := &cappedPluginBuffer{limit: limits.MaxStderrBytes} // command/args are operator-declared config, not user-derived input. - cmd := exec.CommandContext(ctx, e.command, e.args...) //nolint:gosec + cmd := exec.CommandContext(runCtx, e.command, e.args...) //nolint:gosec platform.ConfigureBackgroundCommand(cmd) + cmd.WaitDelay = extractorPluginWaitDelay cmd.Stdin = bytes.NewReader(src) + cmd.Stdout = stdout + cmd.Stderr = stderr cmd.Env = append(os.Environ(), "GORTEX_FILE_PATH="+filePath) - var out, errBuf bytes.Buffer - cmd.Stdout = &out - cmd.Stderr = &errBuf - if err := cmd.Run(); err != nil { - e.log.Warn("extractor plugin: command failed — degraded to file node only", + + runErr := cmd.Run() + if err := ctx.Err(); err != nil { + return nil, noUsage, err + } + if stdout.truncated { + return nil, noUsage, stdoutLimitError(limits.MaxStdoutBytes) + } + if errors.Is(runCtx.Err(), context.DeadlineExceeded) { + return nil, noUsage, fmt.Errorf("extractor plugin timed out after %s", e.timeout) + } + if runErr != nil { + return nil, noUsage, fmt.Errorf("extractor plugin failed (stderr=%q, truncated=%t): %w", + strings.TrimSpace(stderr.String()), stderr.truncated, runErr) + } + if stderr.truncated { + e.log.Warn("extractor plugin: stderr truncated", zap.String("language", e.language), zap.String("file", filePath), - zap.String("stderr", strings.TrimSpace(errBuf.String())), zap.Error(err)) - return result, nil + zap.Int("retained_bytes", limits.MaxStderrBytes)) } - var doc pluginDocument - if err := json.Unmarshal(out.Bytes(), &doc); err != nil { - e.log.Warn("extractor plugin: invalid JSON output — degraded to file node only", - zap.String("language", e.language), zap.String("file", filePath), zap.Error(err)) - return result, nil + stdoutBytes := int64(stdout.buffer.Len()) + budget, budgetErr := newPluginResultBudget(int64(limits.MaxResultBytes), stdoutBytes) + if budgetErr != nil { + return nil, noUsage, budgetErr } + document, decodeErr := decodePluginDocument(ctx, bytes.NewReader(stdout.buffer.Bytes()), filePath, limits, budget) + if err := ctx.Err(); err != nil { + return nil, noUsage, err + } + if decodeErr != nil { + if errors.Is(decodeErr, parser.ErrExtractionLimit) { + return nil, noUsage, decodeErr + } + return nil, noUsage, fmt.Errorf("decode extractor plugin output: %w", decodeErr) + } + converted, err := e.convertPluginDocument(ctx, filePath, result, document) + if err != nil { + return nil, noUsage, err + } + usage := parser.ExtractionUsage{ + RawNodes: document.rawNodes + 1, + RawEdges: document.rawEdges, + StdoutBytes: stdoutBytes, + ResultBytes: budget.observed, + } + return converted, usage, nil +} - for _, pn := range doc.Nodes { - node := e.toNode(filePath, pn) - if node == nil { - continue +func (e *SubprocessExtractor) fileOnlyResult(filePath string, src []byte) *parser.ExtractionResult { + return &parser.ExtractionResult{Nodes: []*graph.Node{e.fileNode(filePath, src)}} +} + +func (e *SubprocessExtractor) convertPluginDocument(ctx context.Context, filePath string, fileOnly *parser.ExtractionResult, document pluginDocument) (*parser.ExtractionResult, error) { + result := &parser.ExtractionResult{ + Nodes: make([]*graph.Node, 0, len(fileOnly.Nodes)+len(document.Nodes)), + Edges: make([]*graph.Edge, 0, len(document.Edges)), + } + result.Nodes = append(result.Nodes, fileOnly.Nodes...) + for index, pluginNode := range document.Nodes { + if err := pollPluginContext(ctx, index); err != nil { + return nil, err + } + node := e.toNode(filePath, pluginNode) + if node != nil { + result.Nodes = append(result.Nodes, node) } - result.Nodes = append(result.Nodes, node) } - for _, pe := range doc.Edges { - edge := e.toEdge(filePath, pe) - if edge == nil { - continue + for index, pluginEdge := range document.Edges { + if err := pollPluginContext(ctx, index); err != nil { + return nil, err } - result.Edges = append(result.Edges, edge) + edge := e.toEdge(filePath, pluginEdge) + if edge != nil { + result.Edges = append(result.Edges, edge) + } + } + if err := ctx.Err(); err != nil { + return nil, err } return result, nil } @@ -224,7 +611,10 @@ func (e *SubprocessExtractor) toEdge(filePath string, pe pluginEdge) *graph.Edge } } -var _ parser.Extractor = (*SubprocessExtractor)(nil) +var ( + _ parser.Extractor = (*SubprocessExtractor)(nil) + _ parser.BoundedExtractor = (*SubprocessExtractor)(nil) +) // RegisterExtractorPlugins registers every configured subprocess // extractor plugin, mirroring RegisterCustomGrammars: a spec whose diff --git a/internal/parser/languages/extractor_plugin_test.go b/internal/parser/languages/extractor_plugin_test.go index 7a2a3d6c..334cc219 100644 --- a/internal/parser/languages/extractor_plugin_test.go +++ b/internal/parser/languages/extractor_plugin_test.go @@ -1,7 +1,12 @@ package languages import ( + "context" + "errors" + "math" + "strings" "testing" + "time" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -22,6 +27,37 @@ func shFor(language, ext, script string) config.ExtractorPluginSpec { } } +func TestExtractionLimitsClampAndClassify(t *testing.T) { + defaults := parser.DefaultExtractionLimits() + require.Equal(t, defaults, defaults.Clamped()) + + plusOne := parser.ExtractionLimits{ + MaxNodes: defaults.MaxNodes + 1, MaxEdges: defaults.MaxEdges + 1, + MaxStdoutBytes: defaults.MaxStdoutBytes + 1, MaxResultBytes: defaults.MaxResultBytes + 1, + MaxStderrBytes: defaults.MaxStderrBytes + 1, + } + require.Equal(t, defaults, plusOne.Clamped()) + + maxed := parser.ExtractionLimits{ + MaxNodes: math.MaxInt, MaxEdges: math.MaxInt, MaxStdoutBytes: math.MaxInt, + MaxResultBytes: math.MaxInt, MaxStderrBytes: math.MaxInt, + } + require.Equal(t, defaults, maxed.Clamped()) + require.Equal(t, parser.ExtractionLimits{}, parser.ExtractionLimits{}.Clamped(), "bounded zero must remain strict") + require.Equal(t, parser.ExtractionLimits{}, parser.ExtractionLimits{ + MaxNodes: -1, MaxEdges: -1, MaxStdoutBytes: -1, MaxResultBytes: -1, MaxStderrBytes: -1, + }.Clamped()) + + limitErr := &parser.ExtractionLimitError{Resource: "nodes", Limit: 7, Observed: 8} + require.ErrorIs(t, limitErr, parser.ErrExtractionLimit) + var typed *parser.ExtractionLimitError + require.ErrorAs(t, limitErr, &typed) + require.Equal(t, int64(8), typed.Observed) + var nilLimit *parser.ExtractionLimitError + require.Equal(t, parser.ErrExtractionLimit.Error(), nilLimit.Error()) + assert.True(t, errors.Is(limitErr, parser.ErrExtractionLimit)) +} + func TestSubprocessExtractor_EmitsNodesAndEdges(t *testing.T) { // One function node + one EdgeDefines from the file to it. The // script drains stdin so the parent's stdin pipe always closes. @@ -100,6 +136,331 @@ func TestSubprocessExtractor_BadJSONDegradesToFileNode(t *testing.T) { assert.Equal(t, graph.KindFile, res.Nodes[0].Kind) } +func decodePluginForTest(document string, limits parser.ExtractionLimits) (pluginDocument, error) { + limits = limits.Clamped() + budget, err := newPluginResultBudget(int64(limits.MaxResultBytes), int64(len(document))) + if err != nil { + return pluginDocument{}, err + } + return decodePluginDocument(context.Background(), strings.NewReader(document), "x", limits, budget) +} + +func requirePluginLimit(t *testing.T, err error, resource string, limit, observed int64) { + t.Helper() + require.ErrorIs(t, err, parser.ErrExtractionLimit) + var limitErr *parser.ExtractionLimitError + require.ErrorAs(t, err, &limitErr) + assert.Equal(t, resource, limitErr.Resource) + assert.Equal(t, limit, limitErr.Limit) + assert.Equal(t, observed, limitErr.Observed) +} + +func TestDecodePluginDocument_CompatibilityAndCumulativeCaps(t *testing.T) { + t.Run("top-level null", func(t *testing.T) { + document, err := decodePluginForTest("null", parser.DefaultExtractionLimits()) + require.NoError(t, err) + assert.Empty(t, document.Nodes) + assert.Empty(t, document.Edges) + }) + + t.Run("case-insensitive duplicate fields are last-wins", func(t *testing.T) { + limits := parser.DefaultExtractionLimits() + limits.MaxNodes = 3 + document, err := decodePluginForTest(`{"Nodes":[{"id":"a","kind":"function"}],"nOdEs":[{"id":"b","kind":"function"}]}`, limits) + require.NoError(t, err) + require.Len(t, document.Nodes, 1) + assert.Equal(t, "b", document.Nodes[0].ID) + }) + + t.Run("duplicate fields count cumulatively", func(t *testing.T) { + limits := parser.DefaultExtractionLimits() + limits.MaxNodes = 3 + _, err := decodePluginForTest(`{"nodes":[{"id":"a","kind":"function"}],"NODES":[{"id":"b","kind":"function"},{"id":"c","kind":"function"}]}`, limits) + requirePluginLimit(t, err, "nodes", 3, 4) + }) + + t.Run("invalid raw nodes still consume the cap", func(t *testing.T) { + limits := parser.DefaultExtractionLimits() + limits.MaxNodes = 2 + _, err := decodePluginForTest(`{"nodes":[{"kind":"invalid"},{"id":"b","kind":"function"}]}`, limits) + requirePluginLimit(t, err, "nodes", 2, 3) + }) + + t.Run("raw edges enforce exact plus one", func(t *testing.T) { + limits := parser.DefaultExtractionLimits() + limits.MaxEdges = 1 + _, err := decodePluginForTest(`{"edges":[{"from":"a","to":"b","kind":"calls"},{"from":"b","to":"c","kind":"calls"}]}`, limits) + requirePluginLimit(t, err, "edges", 1, 2) + }) +} + +func TestDecodePluginDocument_WireAndSynthesizedIDBounds(t *testing.T) { + t.Run("stdout exact and plus one", func(t *testing.T) { + limits := parser.DefaultExtractionLimits() + limits.MaxStdoutBytes = len("null") + _, err := decodePluginForTest("null", limits) + require.NoError(t, err) + _, err = decodePluginForTest("null ", limits) + requirePluginLimit(t, err, "stdout_bytes", int64(len("null")), int64(len("null")+1)) + }) + + t.Run("combined result starts with stdout", func(t *testing.T) { + limits := parser.DefaultExtractionLimits() + limits.MaxResultBytes = len("null") + _, err := decodePluginForTest("null", limits) + require.NoError(t, err) + limits.MaxResultBytes-- + _, err = decodePluginForTest("null", limits) + requirePluginLimit(t, err, "result_bytes", int64(len("null")-1), int64(len("null"))) + }) + + t.Run("synthesized id exact and plus one", func(t *testing.T) { + const wire = `{"nodes":[{"kind":"function","name":"abc"}]}` + const synthesized = "x::abc" + limits := parser.DefaultExtractionLimits() + limits.MaxNodes = 2 + limits.MaxResultBytes = len(wire) + len(synthesized) + document, err := decodePluginForTest(wire, limits) + require.NoError(t, err) + require.Len(t, document.Nodes, 1) + + limits.MaxResultBytes-- + _, err = decodePluginForTest(wire, limits) + requirePluginLimit(t, err, "result_bytes", int64(len(wire)+len(synthesized)-1), int64(len(wire)+len(synthesized))) + }) + + t.Run("overwritten synthesized ids are not refunded", func(t *testing.T) { + const wire = `{"nodes":[{"kind":"function","name":"abc"}],"nodes":[{"kind":"function","name":"abc"}]}` + const synthesized = "x::abc" + limits := parser.DefaultExtractionLimits() + limits.MaxNodes = 3 + limits.MaxResultBytes = len(wire) + len(synthesized) + _, err := decodePluginForTest(wire, limits) + requirePluginLimit(t, err, "result_bytes", int64(len(wire)+len(synthesized)), int64(len(wire)+len(synthesized)+1)) + }) +} + +func TestSubprocessExtractor_BoundedExactAndPlusOne(t *testing.T) { + run := func(wire string, limits parser.ExtractionLimits) (*parser.ExtractionResult, parser.ExtractionUsage, error) { + ext := NewSubprocessExtractor(shFor("ml", ".ml", "cat >/dev/null; printf '%s' '"+wire+"'"), nil) + return ext.ExtractBounded(context.Background(), "x.ml", nil, limits) + } + + t.Run("nodes", func(t *testing.T) { + const wire = `{"nodes":[{"id":"a","kind":"function"}]}` + limits := parser.DefaultExtractionLimits() + limits.MaxNodes = 2 + result, usage, err := run(wire, limits) + require.NoError(t, err) + require.Len(t, result.Nodes, 2) + assert.Equal(t, parser.ExtractionUsage{ + RawNodes: 2, StdoutBytes: int64(len(wire)), ResultBytes: int64(len(wire)), + }, usage) + + limits.MaxNodes = 1 + result, usage, err = run(wire, limits) + assert.Nil(t, result) + assert.Empty(t, usage) + requirePluginLimit(t, err, "nodes", 1, 2) + }) + + t.Run("edges including strict zero", func(t *testing.T) { + const emptyWire = `{"edges":[]}` + limits := parser.DefaultExtractionLimits() + limits.MaxEdges = 0 + result, usage, err := run(emptyWire, limits) + require.NoError(t, err) + require.Len(t, result.Nodes, 1) + assert.Equal(t, parser.ExtractionUsage{ + RawNodes: 1, StdoutBytes: int64(len(emptyWire)), ResultBytes: int64(len(emptyWire)), + }, usage) + + const wire = `{"edges":[{"from":"a","to":"b","kind":"calls"}]}` + limits.MaxEdges = 1 + result, usage, err = run(wire, limits) + require.NoError(t, err) + require.Len(t, result.Edges, 1) + assert.Equal(t, parser.ExtractionUsage{ + RawNodes: 1, RawEdges: 1, StdoutBytes: int64(len(wire)), ResultBytes: int64(len(wire)), + }, usage) + + limits.MaxEdges = 0 + result, usage, err = run(wire, limits) + assert.Nil(t, result) + assert.Empty(t, usage) + requirePluginLimit(t, err, "edges", 0, 1) + }) + + t.Run("stdout bytes", func(t *testing.T) { + const wire = "null" + limits := parser.DefaultExtractionLimits() + limits.MaxStdoutBytes = len(wire) + result, usage, err := run(wire, limits) + require.NoError(t, err) + require.Len(t, result.Nodes, 1) + assert.Equal(t, parser.ExtractionUsage{RawNodes: 1, StdoutBytes: 4, ResultBytes: 4}, usage) + + limits.MaxStdoutBytes-- + result, usage, err = run(wire, limits) + assert.Nil(t, result) + assert.Empty(t, usage) + requirePluginLimit(t, err, "stdout_bytes", int64(len(wire)-1), int64(len(wire))) + }) + + t.Run("combined result bytes", func(t *testing.T) { + const wire = `{"nodes":[{"kind":"function","name":"n"}]}` + const synthesized = "x.ml::n" + limits := parser.DefaultExtractionLimits() + limits.MaxNodes = 2 + limits.MaxResultBytes = len(wire) + len(synthesized) + result, usage, err := run(wire, limits) + require.NoError(t, err) + require.Len(t, result.Nodes, 2) + assert.Equal(t, parser.ExtractionUsage{ + RawNodes: 2, StdoutBytes: int64(len(wire)), ResultBytes: int64(len(wire) + len(synthesized)), + }, usage) + + limits.MaxResultBytes-- + result, usage, err = run(wire, limits) + assert.Nil(t, result) + assert.Empty(t, usage) + requirePluginLimit(t, err, "result_bytes", int64(len(wire)+len(synthesized)-1), int64(len(wire)+len(synthesized))) + }) +} + +func TestSubprocessExtractor_ReportsConservativeUsage(t *testing.T) { + limits := parser.DefaultExtractionLimits() + + t.Run("null charges only exact stdout and file node", func(t *testing.T) { + ext := NewSubprocessExtractor(shFor("ml", ".ml", "cat >/dev/null; printf 'null'"), nil) + result, usage, err := ext.ExtractBounded(context.Background(), "x.ml", nil, limits) + require.NoError(t, err) + require.Len(t, result.Nodes, 1) + assert.Equal(t, parser.ExtractionUsage{RawNodes: 1, StdoutBytes: 4, ResultBytes: 4}, usage) + }) + + t.Run("duplicate and invalid rows remain charged", func(t *testing.T) { + const wire = `{"nodes":[{"id":"ignored","kind":"invalid"},{"id":"old","kind":"function"}],"NODES":[{"kind":"function","name":"last"}],"edges":[{"from":"","to":"x","kind":"calls"}],"EDGES":[{"from":"a","to":"b","kind":"calls"}]}` + ext := NewSubprocessExtractor(shFor("ml", ".ml", "cat >/dev/null; printf '%s' '"+wire+"'"), nil) + result, usage, err := ext.ExtractBounded(context.Background(), "x.ml", nil, limits) + require.NoError(t, err) + require.Len(t, result.Nodes, 2) + require.Len(t, result.Edges, 1) + assert.Equal(t, parser.ExtractionUsage{ + RawNodes: 4, RawEdges: 2, StdoutBytes: int64(len(wire)), + ResultBytes: int64(len(wire) + len("x.ml::last")), + }, usage) + }) +} + +func TestSubprocessExtractor_BoundedProcessAndFailureSemantics(t *testing.T) { + limits := parser.DefaultExtractionLimits() + + t.Run("zero node budget rejects before command", func(t *testing.T) { + ext := NewSubprocessExtractor(shFor("ml", ".ml", "cat >/dev/null; printf 'null'"), nil) + zero := limits + zero.MaxNodes = 0 + result, usage, err := ext.ExtractBounded(context.Background(), "x.ml", nil, zero) + assert.Nil(t, result) + assert.Empty(t, usage) + requirePluginLimit(t, err, "nodes", 0, 1) + }) + + t.Run("bounded operational errors have no partial result", func(t *testing.T) { + for name, script := range map[string]string{ + "nonzero": "cat >/dev/null; printf 'null'; exit 3", + "malformed": "cat >/dev/null; printf 'not-json'", + } { + t.Run(name, func(t *testing.T) { + ext := NewSubprocessExtractor(shFor("ml", ".ml", script), nil) + result, usage, err := ext.ExtractBounded(context.Background(), "x.ml", nil, limits) + assert.Nil(t, result) + assert.Empty(t, usage) + require.Error(t, err) + }) + } + }) + + t.Run("stderr cap is retention only", func(t *testing.T) { + ext := NewSubprocessExtractor(shFor("ml", ".ml", "cat >/dev/null; printf 'abcdefghijklmnopqrstuvwxyz' >&2; printf 'null'"), nil) + for _, retained := range []int{0, 8} { + bounded := limits + bounded.MaxStderrBytes = retained + result, usage, err := ext.ExtractBounded(context.Background(), "x.ml", nil, bounded) + require.NoError(t, err) + require.Len(t, result.Nodes, 1) + assert.Equal(t, parser.ExtractionUsage{RawNodes: 1, StdoutBytes: 4, ResultBytes: 4}, usage) + } + }) + + t.Run("caller cancellation is prompt and exact", func(t *testing.T) { + ext := NewSubprocessExtractor(shFor("ml", ".ml", "cat >/dev/null; exec sleep 10"), nil) + ctx, cancel := context.WithCancel(context.Background()) + timer := time.AfterFunc(25*time.Millisecond, cancel) + defer timer.Stop() + started := time.Now() + result, usage, err := ext.ExtractBounded(ctx, "x.ml", nil, limits) + assert.Nil(t, result) + assert.Empty(t, usage) + require.ErrorIs(t, err, context.Canceled) + assert.Less(t, time.Since(started), 2*time.Second) + }) + + t.Run("caller deadline is prompt and exact", func(t *testing.T) { + ext := NewSubprocessExtractor(shFor("ml", ".ml", "cat >/dev/null; exec sleep 10"), nil) + ctx, cancel := context.WithTimeout(context.Background(), 25*time.Millisecond) + defer cancel() + started := time.Now() + result, usage, err := ext.ExtractBounded(ctx, "x.ml", nil, limits) + assert.Nil(t, result) + assert.Empty(t, usage) + require.ErrorIs(t, err, context.DeadlineExceeded) + assert.Less(t, time.Since(started), 2*time.Second) + }) + + t.Run("descendant-held stdout is bounded by WaitDelay", func(t *testing.T) { + ext := NewSubprocessExtractor(shFor("ml", ".ml", "cat >/dev/null; (sleep 1) & printf 'null'"), nil) + started := time.Now() + result, usage, err := ext.ExtractBounded(context.Background(), "x.ml", nil, limits) + assert.Nil(t, result) + assert.Empty(t, usage) + require.Error(t, err) + assert.Less(t, time.Since(started), 2*time.Second) + }) + + t.Run("stdout overflow cancels a sleeping plugin", func(t *testing.T) { + ext := NewSubprocessExtractor(shFor("ml", ".ml", "cat >/dev/null; printf '0123456789'; exec sleep 10"), nil) + bounded := limits + bounded.MaxStdoutBytes = 4 + started := time.Now() + result, usage, err := ext.ExtractBounded(context.Background(), "x.ml", nil, bounded) + assert.Nil(t, result) + assert.Empty(t, usage) + requirePluginLimit(t, err, "stdout_bytes", 4, 5) + assert.Less(t, time.Since(started), 2*time.Second) + }) + + t.Run("internal timeout is ordinary", func(t *testing.T) { + spec := shFor("ml", ".ml", "cat >/dev/null; exec sleep 10") + spec.TimeoutMs = 25 + ext := NewSubprocessExtractor(spec, nil) + result, usage, err := ext.ExtractBounded(context.Background(), "x.ml", nil, limits) + assert.Nil(t, result) + assert.Empty(t, usage) + require.Error(t, err) + assert.NotErrorIs(t, err, context.DeadlineExceeded) + assert.NotErrorIs(t, err, context.Canceled) + }) + + t.Run("legacy stdout overflow degrades to file only", func(t *testing.T) { + ext := NewSubprocessExtractor(shFor("ml", ".ml", "exec dd if=/dev/zero bs=1048576 count=17 2>/dev/null"), nil) + result, err := ext.Extract("x.ml", nil) + require.NoError(t, err) + require.Len(t, result.Nodes, 1) + assert.Empty(t, result.Edges) + }) +} + func TestRegisterExtractorPlugins_SkipsInvalidAndCollisions(t *testing.T) { reg := parser.NewRegistry() reg.Register(NewGoExtractor()) // claims "go" and ".go"