open-context-engine 0.1.0

package/scripts/mcp-opencontextengine.mjs

@@ -0,0 +1,37 @@
+import { parseArgs } from 'node:util';
+import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js';
+import { createMcpServer } from '../src/mcp.mjs';
+import { clientConfig } from '../src/client.mjs';
+import { createWorkspaceManager } from '../src/workspaces.mjs';
+
+const {values} = parseArgs({options:{root:{type:'string'},state:{type:'string'},connect:{type:'boolean'}}});
+if (values.connect && (values.root !== undefined || values.state !== undefined) || values.root === '' || values.state === '') {
+  throw new Error('Usage: node scripts/mcp-opencontextengine.mjs [--root /repository] [--state /outside/index] OR --connect');
+}
+let workspaces, server, stopping = false;
+async function shutdown(code = 0) {
+  if (stopping) return;
+  stopping = true;
+  try {
+    await server?.close();
+  } finally {
+    await workspaces?.close();
+    process.exit(code);
+  }
+}
+process.once('SIGINT', () => {void shutdown();});
+process.once('SIGTERM', () => {void shutdown();});
+process.stdin.once('end', () => {void shutdown();});
+try {
+  if (values.connect) {
+    server = createMcpServer(clientConfig());
+  } else {
+    workspaces = createWorkspaceManager(values);
+    if (values.root) await workspaces.get();
+    if (!stopping) server = createMcpServer(undefined, {resolveConfig:workspaces.get, automatic:!values.root});
+  }
+  if (!stopping) await server.connect(new StdioServerTransport());
+} catch (error) {
+  process.stderr.write(`OpenContextEngine MCP startup failed: ${error.message}\n`);
+  await shutdown(1);
+}

package/scripts/retrieval-server.py

@@ -0,0 +1,135 @@
+"""Authenticated, persistent retrieval worker; model inference stays remote."""
+import hashlib
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+import json
+from pathlib import Path
+import re
+import secrets
+import sys
+import threading
+import time
+from urllib.request import urlopen
+import numpy as np
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT / 'src' / 'retrieval'))
+from routed import RoutedEngine, VERSION
+from live import LiveIndex, IndexUnavailable
+
+
+def plan_query(query):
+    parts = [part.strip() for part in re.split(r'[:：;；，]|,\s+(?=how|which|why|where|what)|\s+and\s+(?=how|which|why|where|what)', query, flags=re.I) if len(part.strip()) > 7]
+    facets = parts if 1 < len(parts) <= 4 else [query]
+    return {'intent':query,'facets':[{'question':part,'terms':re.findall(r'[A-Za-z][A-Za-z0-9_]*',part)} for part in facets]}
+
+
+def serve(config):
+    initialized = time.monotonic()
+    live = LiveIndex(config).start() if config.get('root') else None
+    state = Path(config['state'])
+    if live:
+        index, retrieval = None, None
+    else:
+        units = json.loads((state / 'units.json').read_text())
+        index = json.loads((state / 'metadata.json').read_text())
+        retrieval = RoutedEngine(units, np.load(state/'vectors.npy'), config['embeddingUrl'], config['reranker'], config['embeddingKey'])
+    health = {'status':'ready','engine':VERSION,'index':index,'queryCache':False,
+        'initializationMs':round((time.monotonic()-initialized)*1000),
+        'sourceSha256':{name:hashlib.sha256((ROOT/name).read_bytes()).hexdigest() for name in
+            ['src/retrieval/engine.py','src/retrieval/batched.py','src/retrieval/routed.py','src/retrieval/reranker.py','scripts/retrieval-server.py',
+             'src/retrieval/languages/__init__.py','src/retrieval/languages/schema.py',
+             'src/retrieval/languages/text.py','src/retrieval/languages/files.py',
+             'src/retrieval/languages/python.py','src/retrieval/languages/go.py','src/retrieval/languages/go_ast.go','src/retrieval/languages/go_types.go',
+             'src/retrieval/languages/typescript.py',
+             'src/retrieval/languages/typescript.mjs','package.json','package-lock.json']
+            if name != 'package-lock.json' or (ROOT/name).is_file()}}
+    model_base = config['reranker']['baseUrl'].removesuffix('/v1')
+    if not live:
+        with urlopen(model_base+'/healthz',timeout=10) as response:
+            health['reranker'] = json.load(response)
+    lock = threading.Lock()
+
+    class Handler(BaseHTTPRequestHandler):
+        def log_message(self, *args):
+            pass
+
+        def reply(self, status, data):
+            body = json.dumps(data,ensure_ascii=False).encode()
+            self.send_response(status)
+            self.send_header('Content-Type','application/json; charset=utf-8')
+            self.send_header('Content-Length',str(len(body)))
+            self.end_headers()
+            self.wfile.write(body)
+
+        def do_GET(self):
+            if self.path == '/healthz':
+                return self.reply(200, {'status':'running','engine':VERSION,'mode':'live'} if live else health)
+            if self.path == '/status':
+                if not secrets.compare_digest(self.headers.get('Authorization',''), 'Bearer '+config['serviceKey']):
+                    return self.reply(401, {'error':'Unauthorized'})
+                return self.reply(200, live.status() if live else {'status':'ready','mode':'frozen','generation':index})
+            self.reply(404,{'error':'Not found'})
+
+        def do_POST(self):
+            if self.path!='/search':
+                return self.reply(404,{'error':'Not found'})
+            if not secrets.compare_digest(self.headers.get('Authorization',''), 'Bearer '+config['serviceKey']):
+                return self.reply(401,{'error':'Unauthorized'})
+            try:
+                length = int(self.headers.get('Content-Length','0'))
+                if not 1<=length<=32768:
+                    raise ValueError('Invalid body size')
+                body = json.loads(self.rfile.read(length))
+                if not isinstance(body,dict) or set(body)-{'query','budget','trace','freshnessWaitMs'}:
+                    raise ValueError('Unknown fields')
+                query,budget = body.get('query'),body.get('budget',4000)
+                wait_ms = body.get('freshnessWaitMs', 30000)
+                if type(wait_ms) is not int or not 0 <= wait_ms <= 120000:
+                    raise ValueError('Invalid freshness wait')
+                if not isinstance(query,str) or not query.strip() or len(query)>8192 or type(budget) is not int or not 256<=budget<=8000 or type(body.get('trace',False)) is not bool:
+                    raise ValueError('Invalid search input')
+            except (ValueError,TypeError):
+                return self.reply(422,{'error':'Expected query text and a token budget between 256 and 8000'})
+            start = time.monotonic()
+            if not lock.acquire(timeout=5):
+                return self.reply(429,{'error':'Retrieval worker busy'})
+            try:
+                queued = round((time.monotonic()-start)*1000)
+                generation = live.current(wait_ms/1000) if live else None
+                engine = generation.engine if live else retrieval
+                if engine is None:
+                    raw,debug = '', {'tokens':0,'elapsedMs':0}
+                else:
+                    raw,debug = engine.search(plan_query(query),budget=budget)
+                if live:
+                    live.verify(generation)
+                response = {'context':raw,'tokens':debug['tokens'],'engine':VERSION,
+                    'retrievalMs':debug['elapsedMs'],'queueMs':queued,
+                    'serverElapsedMs':round((time.monotonic()-start)*1000),'queryCache':False,
+                    'index': {'mode':'live','identity':generation.identity,'freshness':'verified-after-search',
+                              'completedAt':generation.info['completedAt']} if live else {'mode':'frozen'}}
+                if body.get('trace'):
+                    response['diagnostics'] = debug
+                self.reply(200,response)
+            except IndexUnavailable as error:
+                self.reply(503,{'error':str(error),'index':live.status()})
+            except Exception as error:
+                print(json.dumps({'event':'search-failed','type':type(error).__name__}),flush=True)
+                self.reply(502,{'error':'Retrieval or model request failed'})
+            finally:
+                lock.release()
+
+    server = ThreadingHTTPServer(('127.0.0.1',config.get('port',23505)),Handler)
+    server.daemon_threads = True
+    print(json.dumps({'listening':f'http://127.0.0.1:{server.server_port}',
+                      'health':{'status':'running','mode':'live'} if live else health}),flush=True)
+    try:
+        server.serve_forever()
+    finally:
+        server.server_close()
+        if live:
+            live.close()
+
+
+if __name__ == '__main__':
+    serve(json.loads(sys.stdin.readline()))

package/src/client.mjs

@@ -0,0 +1,37 @@
+import { loadEnvironment } from './config.mjs';
+
+export function clientConfig(environment=process.env) {
+  const env=loadEnvironment(environment);
+  const url=new URL(env.OCE_BASE_URL || 'http://127.0.0.1:45005');
+  if (url.username || url.password || url.search || url.hash ||
+      !(url.protocol==='https:' || url.protocol==='http:' && url.hostname==='127.0.0.1')) throw new Error('Use HTTPS or an explicit loopback SSH forward');
+  const apiKey=env.OCE_API_KEY || env.RERANK_API_KEY;
+  if (!apiKey) throw new Error('Set OCE_API_KEY or the existing project reranker key');
+  return {baseUrl:url.href.replace(/\/$/,''),apiKey};
+}
+
+export async function search(query,{budget=4000,trace=false,freshnessWaitMs=30000,config=clientConfig(),signal}={}) {
+  const started=performance.now();
+  const timeout = AbortSignal.timeout(freshnessWaitMs + 30000);
+  const response=await fetch(`${config.baseUrl}/search`,{method:'POST',redirect:'error',signal:signal ? AbortSignal.any([signal,timeout]) : timeout,
+    headers:{'content-type':'application/json',authorization:`Bearer ${config.apiKey}`},body:JSON.stringify({query,budget,trace,freshnessWaitMs})});
+  if (!response.ok) {
+    if (response.status === 503) {
+      const body = await response.json();
+      throw new Error(`Index unavailable: ${body.error || 'update pending'}`);
+    }
+    throw new Error(`Retrieval HTTP ${response.status}`);
+  }
+  const result=await response.json();
+  if (typeof result.context!=='string' || !Number.isFinite(result.retrievalMs)) throw new Error('Invalid retrieval response');
+  return {...result,clientElapsedMs:Math.round(performance.now()-started)};
+}
+
+export async function indexStatus({config=clientConfig(),signal}={}) {
+  const timeout = AbortSignal.timeout(10000);
+  const response = await fetch(`${config.baseUrl}/status`, {redirect:'error',
+    signal:signal ? AbortSignal.any([signal,timeout]) : timeout,
+    headers:{authorization:`Bearer ${config.apiKey}`}});
+  if (!response.ok) throw new Error(`Retrieval HTTP ${response.status}`);
+  return response.json();
+}

package/src/config.mjs

@@ -0,0 +1,49 @@
+import { existsSync, readFileSync, mkdirSync, writeFileSync, renameSync, rmSync } from 'node:fs';
+import { homedir } from 'node:os';
+import { resolve, join } from 'node:path';
+import { fileURLToPath } from 'node:url';
+import { parseEnv } from 'node:util';
+import { randomUUID } from 'node:crypto';
+import { normalizeEnvironment } from './environment.mjs';
+
+export const projectRoot = fileURLToPath(new URL('../', import.meta.url));
+export const configKeys = new Set([
+  'EMBEDDING_BASE_URL','EMBEDDING_API_KEY','EMBEDDING_MODEL','EMBEDDING_SSH_TUNNEL_URL','EMBEDDING_SSH_REMOTE',
+  'RERANK_BASE_URL','RERANK_API_KEY','RERANK_MODEL','RERANK_REMOTE_RUNTIME_URL','RERANK_REMOTE_HOST',
+  'RERANK_SSH_TUNNEL_URL','RERANK_SSH_REMOTE',
+  'OCE_ALLOW_HTTP','OCE_RERANK_API','OCE_RERANK_CONCURRENCY','OCE_RERANK_MAX_DOCUMENTS','OCE_EMBEDDING_DIMENSIONS',
+  'OCE_EMBEDDING_REVISION','OCE_PYTHON','OCE_GO_BINARY','OCE_LANGUAGE_OPTIONS','OCE_POLL_SECONDS',
+  'OCE_DEBOUNCE_SECONDS','OCE_API_KEY','OCE_BASE_URL','TIKTOKEN_CACHE_DIR',
+]);
+export function configDirectory(environment = process.env) {
+  return resolve(environment.OCE_CONFIG_HOME || join(homedir(), '.config', 'opencontextengine'));
+}
+export function readUserConfig(environment = process.env) {
+  const path = join(configDirectory(environment), 'config.json');
+  if (!existsSync(path)) return {};
+  let data;
+  try {data = JSON.parse(readFileSync(path, 'utf8'));} catch {throw new Error(`Invalid OpenContextEngine configuration: ${path}`);}
+  if (data?.schemaVersion !== 1 || !data.env || typeof data.env !== 'object' || Array.isArray(data.env)
+      || Object.entries(data.env).some(([key,value]) => !configKeys.has(key) || typeof value !== 'string')) {
+    throw new Error(`Invalid OpenContextEngine configuration: ${path}`);
+  }
+  return data.env;
+}
+export function loadEnvironment(environment = process.env, checkout = projectRoot) {
+  const file = join(checkout, '.env');
+  return {...normalizeEnvironment(existsSync(file) ? parseEnv(readFileSync(file,'utf8')) : {}),
+    ...readUserConfig(environment),
+    ...normalizeEnvironment(environment)};
+}
+export function saveUserConfig(values, environment = process.env) {
+  const directory = configDirectory(environment), path = join(directory,'config.json');
+  const env = Object.fromEntries(Object.entries(values).filter(([key]) => configKeys.has(key)));
+  if (Object.values(env).some(value => typeof value !== 'string')) throw new Error('Configuration values must be strings');
+  mkdirSync(directory, {recursive:true,mode:0o700});
+  const temporary = join(directory, `.config-${randomUUID()}.tmp`);
+  try {
+    writeFileSync(temporary, JSON.stringify({schemaVersion:1,env},null,2)+'\n', {flag:'wx',mode:0o600});
+    renameSync(temporary,path);
+  } finally {rmSync(temporary,{force:true});}
+  return path;
+}

package/src/environment.mjs

@@ -0,0 +1,10 @@
+// Keep existing deployments working while publishing only the OCE_* interface.
+export function normalizeEnvironment(environment) {
+  const result = {...environment};
+  for (const [key,value] of Object.entries(environment)) {
+    if (!key.startsWith('REPONERVE_')) continue;
+    const current = 'OCE_' + key.slice('REPONERVE_'.length);
+    if (!Object.hasOwn(environment,current)) result[current] = value;
+  }
+  return result;
+}

package/src/eval/remote-models.mjs

@@ -0,0 +1,70 @@
+import os from 'node:os';
+
+// User requirement: model inference must run on a separately authorized server.
+export function remoteModelUrl(value, {allowHttp = false} = {}) {
+  const url = new URL(value);
+  const host = url.hostname.toLowerCase().replace(/^\[|\]$/g, '');
+  const local = new Set(['localhost', 'localhost.localdomain', '0.0.0.0', '::', '::1', os.hostname().toLowerCase()]);
+  for (const addresses of Object.values(os.networkInterfaces())) {
+    for (const address of addresses ?? []) local.add(address.address.toLowerCase());
+  }
+  if (!(url.protocol === 'https:' || (allowHttp && url.protocol === 'http:')) || url.username || url.password || url.search || url.hash ||
+      local.has(host) || /^127\./.test(host) || host.endsWith('.localhost') || host.endsWith('.local')) {
+    throw new Error('Model endpoint must be remote HTTPS (or explicitly enable OCE_ALLOW_HTTP=1); local model inference is prohibited');
+  }
+  return url.href.replace(/\/$/, '');
+}
+
+export function remoteRerankerConfig(env) {
+  if (!env.RERANK_BASE_URL || !env.RERANK_MODEL || !env.RERANK_API_KEY) {
+    throw new Error('A user-approved remote reranker is required. No local model will be installed or started.');
+  }
+  const api = env.OCE_RERANK_API ?? 'rerank';
+  if (!['rerank', 'rerank-batch'].includes(api)) throw new Error('OCE_RERANK_API must be rerank or rerank-batch');
+  const integer = (name, fallback, maximum) => {
+    const value = Number(env[name] ?? fallback);
+    if (!Number.isInteger(value) || value < 1 || value > maximum) throw new Error(`${name} must be an integer from 1 to ${maximum}`);
+    return value;
+  };
+  return { baseUrl: remoteModelUrl(env.RERANK_BASE_URL, {allowHttp:env.OCE_ALLOW_HTTP === '1'}), model: env.RERANK_MODEL, apiKey: env.RERANK_API_KEY,
+    api, concurrency: integer('OCE_RERANK_CONCURRENCY', 2, 8),
+    maxDocuments: integer('OCE_RERANK_MAX_DOCUMENTS', 128, 1024) };
+}
+
+// An explicitly named authorized Linux worker can reach its own remote GPU
+// service without a round trip through the public HTTPS proxy. This does not
+// permit launching a model process, nor accepting a loopback model on the Mac.
+export function rerankerExecutionTransport(env, host = { platform: process.platform, hostname: os.hostname() }) {
+  const baseUrl = remoteModelUrl(env.RERANK_BASE_URL, {allowHttp:env.OCE_ALLOW_HTTP === '1'});
+  if (env.RERANK_SSH_TUNNEL_URL) {
+    if (env.RERANK_REMOTE_RUNTIME_URL) throw new Error('Choose either SSH forwarding or same-host reranker transport');
+    return sshForward(baseUrl, env.RERANK_SSH_TUNNEL_URL, env.RERANK_SSH_REMOTE, new URL(baseUrl).pathname);
+  }
+  if (!env.RERANK_REMOTE_RUNTIME_URL) return { baseUrl, requestBaseUrl: baseUrl, transport: new URL(baseUrl).protocol.slice(0,-1) };
+  const url = new URL(env.RERANK_REMOTE_RUNTIME_URL);
+  if (host.platform !== 'linux' || env.RERANK_REMOTE_HOST !== host.hostname ||
+      url.protocol !== 'http:' || url.hostname !== '127.0.0.1' || !url.port ||
+      url.pathname !== '/v1' || url.username || url.password || url.search || url.hash) {
+    throw new Error('Direct GPU transport requires the explicitly named remote Linux execution host');
+  }
+  return { baseUrl, requestBaseUrl: url.href, transport: 'authorized-remote-host-loopback', executionHost: host.hostname };
+}
+
+// Explicit SSH forwarding reaches a remote service; no local model is loaded.
+// Cache identity remains the logical provider URL when transport changes.
+export function embeddingTransportConfig(env) {
+  const baseUrl = remoteModelUrl(env.EMBEDDING_BASE_URL, {allowHttp:env.OCE_ALLOW_HTTP === '1'});
+  if (!env.EMBEDDING_SSH_TUNNEL_URL) return { baseUrl, requestBaseUrl: baseUrl, transport: new URL(baseUrl).protocol.slice(0,-1) };
+  return sshForward(baseUrl, env.EMBEDDING_SSH_TUNNEL_URL, env.EMBEDDING_SSH_REMOTE, '/v1');
+}
+
+function sshForward(baseUrl, value, remote, path) {
+  const tunnel = new URL(value);
+  if (tunnel.protocol !== 'http:' || tunnel.hostname !== '127.0.0.1' || !tunnel.port ||
+      tunnel.username || tunnel.password || tunnel.search || tunnel.hash ||
+      tunnel.pathname.replace(/\/$/, '') !== path.replace(/\/$/, '') ||
+      !remote || !/^[a-zA-Z0-9_.-]+@[a-zA-Z0-9.-]+:[0-9]+$/.test(remote)) {
+    throw new Error('SSH model transport requires an explicit loopback forward and remote SSH authority');
+  }
+  return { baseUrl, requestBaseUrl: tunnel.href.replace(/\/$/, ''), transport: 'ssh-tunnel', remote };
+}

package/src/mcp.mjs

@@ -0,0 +1,58 @@
+import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
+import { z } from 'zod';
+import { search, indexStatus } from './client.mjs';
+
+export function createMcpServer(config, {resolveConfig, automatic = false} = {}) {
+  const workspaceInstructions = automatic
+    ? 'Pass directory_path as the absolute directory of the project the user is working on with every tool call. '
+      + 'Use the workspace path supplied by your host or inspect the current project directory; do not guess it. '
+      + 'A new project is indexed on first access. The same server can search multiple projects independently. '
+    : 'This server searches one configured repository. ';
+  async function selectConfig(directoryPath) {
+    if (resolveConfig) return resolveConfig(directoryPath);
+    if (directoryPath !== undefined) throw new Error('Connected-service mode uses its configured repository; omit directory_path');
+    return config;
+  }
+  const pathSchema = z.string().min(1).describe('Absolute path to the project directory. Required in automatic workspace mode.');
+  const directoryPath = automatic ? pathSchema : pathSchema.optional();
+  const server = new McpServer({name:'open-context-engine',version:'0.1.0'}, {
+    instructions:workspaceInstructions + 'Search for source evidence. Results include source paths and line numbers. '
+      + 'Search waits for saved file changes to be indexed. If an update is pending or fails, inspect index_status and retry after it completes. '
+      + 'Read target files again before editing, because code may change after a search.',
+  });
+  const annotations = {readOnlyHint:true,destructiveHint:false,idempotentHint:true,openWorldHint:false};
+  server.registerTool('search_code', {
+    title:'Search repository source',
+    description:'Find code implementing a task or behavior in a repository. Returns source evidence under a token budget. '
+      + workspaceInstructions
+      + 'Uses saved working-tree content, including uncommitted changes; rejects stale results when synchronization fails.',
+    inputSchema:{directory_path:directoryPath,query:z.string().trim().min(1).max(8192),budget:z.number().int().min(256).max(8000).default(4000),
+      freshnessWaitMs:z.number().int().min(0).max(120000).default(30000)},
+    annotations,
+  }, async ({directory_path,query,budget,freshnessWaitMs}, extra) => {
+    try {
+      const selected = await selectConfig(directory_path);
+      const result = await search(query,{budget,freshnessWaitMs,config:selected,signal:extra.signal});
+      if (result.index?.mode !== 'live') throw new Error('This service uses a frozen index; connect to a service started with --root');
+      return {content:[{type:'text',text:result.context || 'No matching source context.'}],structuredContent:result};
+    } catch (error) {
+      return {isError:true,content:[{type:'text',text:error.message}]};
+    }
+  });
+  server.registerTool('index_status', {
+    title:'Repository index status',
+    description:'Inspect repository synchronization, current generation, embedding reuse and last update error. '
+      + workspaceInstructions
+      + 'Status is the latest background observation; search_code actively checks source freshness.',
+    inputSchema:{directory_path:directoryPath},annotations,
+  }, async ({directory_path},extra) => {
+    try {
+      const selected = await selectConfig(directory_path);
+      const result = await indexStatus({config:selected,signal:extra.signal});
+      return {content:[{type:'text',text:JSON.stringify(result)}],structuredContent:result};
+    } catch (error) {
+      return {isError:true,content:[{type:'text',text:error.message}]};
+    }
+  });
+  return server;
+}

package/src/retrieval/batched.py

@@ -0,0 +1,220 @@
+"""Compile the reference retrieval DAG into two batched neural scoring waves.
+
+Candidate retention reserves distinct structural neighbors before neural scoring.
+Scoring uses ordinary rerank calls or optional multi-query batches. Duplicate
+pairs are reused only within one request. No response or cross-query cache.
+"""
+from collections import defaultdict
+import time
+import numpy as np
+
+from engine import Engine, document, post
+from reranker import rerank_pairs
+
+VERSION = 'batched-dag-v6'
+
+
+class BatchedEngine(Engine):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.structural_neighbors = [dict() for _ in self.units]
+        for unit in self.units:
+            for relation in unit.get('relations', []):
+                target = relation['target']
+                kind = relation['kind']
+                # Follow execution and snippet completeness, not broad type/import hubs.
+                supported = kind in {'calls', 'same_symbol'} or (
+                    kind == 'member_of' and unit['kind'] in {'function', 'method'}
+                    and self.units[target]['kind'] in {'function', 'method'})
+                if supported:
+                    for source, neighbor in ((unit['id'], target), (target, unit['id'])):
+                        edges = self.structural_neighbors[source]
+                        edges[neighbor] = max(edges.get(neighbor, 0), relation['confidence'])
+        self.context_bundles = self._context_bundles()
+
+    def _context_bundles(self):
+        """Keep small callable bodies and their bounded caller context together."""
+        by_symbol = defaultdict(list)
+        for unit in self.units:
+            by_symbol[unit['symbol']].append(unit['id'])
+        owners = []
+        for unit in self.units:
+            symbol, parent, seen = unit['symbol'], unit.get('owner'), set()
+            while parent in by_symbol and parent not in seen:
+                seen.add(parent)
+                ancestor = self.units[by_symbol[parent][0]]
+                if ancestor['kind'] not in {'function', 'method'}:
+                    break
+                symbol, parent = parent, ancestor.get('owner')
+            owners.append(symbol)
+        families = defaultdict(list)
+        for uid, symbol in enumerate(owners):
+            families[symbol].append(uid)
+        small = {symbol: ids for symbol, ids in families.items()
+                 if self.units[by_symbol[symbol][0]]['kind'] in {'function', 'method'}
+                 and sum(self.costs[uid] for uid in ids) <= 256}
+        callers = defaultdict(set)
+        for unit in self.units:
+            for relation in unit.get('relations', []):
+                source, target = owners[unit['id']], owners[relation['target']]
+                if relation['kind'] == 'calls' and source != target:
+                    callers[target].add(source)
+        bundles = []
+        for unit in self.units:
+            root = owners[unit['id']]
+            if root not in small:
+                bundles.append([unit['id']]); continue
+            included, frontier = {root}, [root]
+            total = sum(self.costs[uid] for uid in small[root])
+            # Bounded source-only context, never a substitute for task labels.
+            for _ in range(2):
+                following = []
+                for symbol in frontier:
+                    if len(callers[symbol]) > 4:
+                        continue
+                    for caller in sorted(callers[symbol], key=lambda s: families[s][0]):
+                        if caller in included or caller not in small:
+                            continue
+                        cost = sum(self.costs[uid] for uid in small[caller])
+                        if total + cost <= 768:
+                            total += cost; included.add(caller); following.append(caller)
+                frontier = following
+            bundles.append(sorted(uid for symbol in included for uid in small[symbol]))
+        return bundles
+
+    def retain_candidates(self, ranked, expanded, facet_scores, fused, limit=80, graph_slots=16):
+        core = ranked[:limit - graph_slots]
+        core_set = set(core)
+        support = {}
+        # An unscored neighbor inherits evidence from its scored anchor. Requiring
+        # the neighbor's own first-wave score would discard the point of expansion.
+        for anchor in core:
+            score = max(scores.get(anchor, 0) for scores in facet_scores)
+            neighbors = {uid: confidence for uid, confidence in self.structural_neighbors[anchor].items()
+                         if uid not in core_set}
+            for uid, confidence in neighbors.items():
+                support[uid] = max(support.get(uid, 0), score * confidence / len(neighbors))
+        novel = (set(expanded) | set(support)) - core_set
+        graph = sorted(novel, key=lambda uid: (-support.get(uid, 0), -fused[uid], uid))[:graph_slots]
+        retained = list(dict.fromkeys(core + graph + ranked))[:limit]
+        return retained, {'core': core, 'graph': graph,
+                          'novelGraphCandidates': len(novel),
+                          'support': {uid: support.get(uid, 0) for uid in graph}}
+
+    def scoring_policy(self, cache, queries, second_wave):
+        return {}
+
+    def scoring_pairs(self, requested, queries, dense, policy):
+        return requested
+
+    def pair_score(self, cache, query, uid, queries, dense, policy):
+        return cache[(query, uid)]
+
+    def search(self, plan, budget=4000):
+        start = time.monotonic()
+        facets = plan['facets']
+        queries = [plan['intent']] + [f['question'] for f in facets]
+        unique = list(dict.fromkeys(queries))
+        embedded = post(self.embed_url+'/embeddings',{'model':self.embedding_model,
+            'input':['Instruct: Retrieve source code implementing the requested behavior.\nQuery: '+q for q in unique]},self.embedding_key)
+        vectors = [r['embedding'] for r in sorted(embedded['data'],key=lambda r:r['index'])]
+        embedding_at = time.monotonic()
+        dense = self.vectors @ np.asarray([vectors[unique.index(q)] for q in queries],dtype=np.float32).T
+        pools, fused = [], defaultdict(float)
+        for col, query in enumerate(queries):
+            lexical = self.lexical(query + (' ' + ' '.join(facets[col-1]['terms']) if col else ''))
+            local = defaultdict(float)
+            for scores in (dense[:, col], lexical):
+                for rank, uid in enumerate(np.argsort(-scores)[:40]):
+                    if scores[uid] > 0:
+                        local[int(uid)] += 1 / (30 + rank)
+            pools.append(sorted(local,key=local.get,reverse=True)[:28])
+            for uid,score in local.items():
+                fused[uid] += score
+        recalled_at = time.monotonic()
+        pair_cache, waves = {}, []
+
+        def rank_wave(jobs):
+            requested = [(query,uid) for query,ids in jobs for uid in ids]
+            policy = self.scoring_policy(pair_cache, queries, bool(waves))
+            scoring = self.scoring_pairs(requested, queries, dense, policy)
+            needed = list(dict.fromkeys(pair for pair in scoring if pair not in pair_cache))
+            if needed:
+                query_list = list(dict.fromkeys(q for q,uid in needed))
+                ids = list(dict.fromkeys(uid for q,uid in needed))
+                qmap = {q:i for i,q in enumerate(query_list)}
+                imap = {uid:i for i,uid in enumerate(ids)}
+                pairs = [(qmap[q],imap[uid]) for q,uid in needed]
+                before = time.monotonic()
+                data = rerank_pairs(self.reranker, query_list,
+                    [document(self.units[uid],5000) for uid in ids], pairs, post)
+                rows = data['results']
+                for row in rows:
+                    value = row['relevance_score']
+                    pair_cache[needed[row['index']]] = value
+                waves.append({'elapsedMs':round((time.monotonic()-before)*1000),
+                    'requests':data['meta']['request_count'],
+                    'modelMs':data.get('meta',{}).get('elapsed_ms'),'pairs':len(needed),
+                    'reusedPairs':len(requested)-len(needed),'inputTokens':data.get('usage',{}).get('input_tokens'),
+                    'maxBatchSize':data.get('meta',{}).get('max_batch_size'),
+                    'batchTokenBudget':data.get('meta',{}).get('batch_token_budget'), 'scoringPolicy': policy})
+            return [{uid:self.pair_score(pair_cache, query, uid, queries, dense, policy) for uid in ids} for query,ids in jobs]
+
+        facet_scores = rank_wave(list(zip(queries,pools)))
+        seeds = set()
+        for scores in facet_scores:
+            seeds.update(sorted(scores,key=scores.get,reverse=True)[:3])
+        candidates = set().union(*(set(s) for s in facet_scores))
+        expanded = set()
+        for uid in seeds:
+            neighbors = self.units[uid]['edges'] + self.incoming[uid]
+            expanded.update(sorted(set(neighbors),key=lambda x:fused[x],reverse=True)[:10])
+        candidates.update(expanded)
+        ranked = sorted(candidates,key=lambda uid:max(s.get(uid,0) for s in facet_scores)+min(.15,fused[uid]),reverse=True)
+        retained, retention = self.retain_candidates(ranked, expanded, facet_scores, fused)
+        expanded.update(retention['graph'])
+        candidates.update(retention['graph'])
+        jobs = [(plan['intent'],retained)] + [(facet['question'],retained) for facet in facets]
+        outputs = rank_wave(jobs)
+        overall = outputs[0]
+        for col,scores in enumerate(outputs[1:],1):
+            facet_scores[col].update(scores)
+        ranked_at = time.monotonic()
+        covered = np.zeros(len(facets))
+        selected, selected_set, spent, trace = [], set(), 0, []
+        available = set(retained)
+        while available:
+            choices = []
+            for uid in available:
+                bundle = [item for item in self.context_bundles[uid] if item not in selected_set]
+                cost = sum(self.costs[item] for item in bundle)
+                if spent+cost>budget:
+                    continue
+                values = np.asarray([s.get(uid,0) for s in facet_scores[1:]])
+                gain = float(np.sum(values/(1+covered)))/len(facets)
+                gain = .7*gain+.3*overall.get(uid,0)
+                gain /= (max(120,cost)/300)**.35
+                choices.append((gain,uid,values,bundle,cost))
+            if not choices:
+                break
+            gain,uid,values,bundle,cost = max(choices,key=lambda item:(item[0],-item[1]))
+            if gain<getattr(self, 'min_gain', .015):
+                break
+            available.difference_update(bundle);selected.extend(bundle);selected_set.update(bundle)
+            spent+=cost;covered+=values
+            for item in bundle:
+                trace.append({'id':item,'overall':overall.get(item,0),
+                    'facets':[s.get(item,0) for s in facet_scores[1:]],
+                    'tokens':self.costs[item],'gain':gain,'graphExpanded':item in expanded,
+                    'selectionAnchor':uid,'contextOnly':item not in retained})
+        raw = '\n'.join(self.render(self.units[uid]) for uid in selected)
+        end = time.monotonic()
+        return raw,{'version':getattr(self, 'version', VERSION),'elapsedMs':round((end-start)*1000),'tokens':len(self.encoding.encode(raw)),
+            'candidateCount':len(candidates),'rerankedCount':len(retained),'expandedCount':len(expanded),
+            'retention':retention,
+            'modelRequests':{'embedding':1,'rerank':sum(w['requests'] for w in waves)},
+            'rerankApi':self.reranker.get('api','rerank'),'queryCache':False,'pairCacheScope':'one-search-only',
+            'waves':waves,'timingMs':{'embedding':round((embedding_at-start)*1000),'recall':round((recalled_at-embedding_at)*1000),
+                'rerank':sum(w['elapsedMs'] for w in waves),'rerankModel':sum(w['modelMs'] or 0 for w in waves),
+                'graphAndBookkeeping':round((ranked_at-recalled_at)*1000)-sum(w['elapsedMs'] for w in waves),
+                'selection':round((end-ranked_at)*1000)},'selected':trace,'plan':plan}