diff --git a/.env.example b/.env.example index 36fd1792..6d94a745 100644 --- a/.env.example +++ b/.env.example @@ -157,11 +157,6 @@ USE_TREE_SITTER=1 INDEX_USE_ENHANCED_AST=1 INDEX_SEMANTIC_CHUNKS=1 -# Pattern Search - structural code similarity across languages -# Enables pattern_search MCP tool and indexes 64-dim pattern vectors -# Uses WL graph kernel, CFG fingerprints, SimHash, spectral features -# PATTERN_VECTORS=1 - # Indexer scaling and exclusions # QDRANT_DEFAULT_EXCLUDES=0 # QDRANT_IGNORE_FILE=.qdrantignore @@ -301,6 +296,10 @@ MAX_CHANGED_SYMBOLS_RATIO=0.6 # INDEX_UPSERT_BACKOFF=0.5 # Debounce file events to coalesce bursts # WATCH_DEBOUNCE_SECS=1.5 +# Journal replay batch size and status logging interval (multi-repo watcher) +# WATCH_JOURNAL_DRAIN_BATCH_SIZE=256 +# WATCH_JOURNAL_LOG_INTERVAL_SECS=60 +# WATCH_JOURNAL_STATUS_BATCH_THRESHOLD=32 # Optional fs metadata fast-path for unchanged files (skips re-reading files # when size/mtime match cache.json in the same workspace). # INDEX_FS_FASTPATH=0 @@ -341,16 +340,6 @@ COMMIT_VECTOR_SEARCH=0 STRICT_MEMORY_RESTORE=1 -# info_request() tool settings (simplified codebase retrieval) -# Default result limit for info_request queries -INFO_REQUEST_LIMIT=10 -# Default context lines in snippets (richer than repo_search default) -INFO_REQUEST_CONTEXT_LINES=5 -# Enable explanation mode by default (summary, primary_locations, related_concepts) -# INFO_REQUEST_EXPLAIN_DEFAULT=0 -# Enable relationship mapping by default (imports_from, calls, related_paths) -# INFO_REQUEST_RELATIONSHIPS=0 - # TOON output format (Token-Oriented Object Notation) # When enabled, search results use compact TOON encoding to reduce token usage # TOON_ENABLED=0 @@ -492,5 +481,3 @@ OPENLIT_ENVIRONMENT=development # --------------------------------------------------------------------------- # End of Auth & Bridge Configuration # --------------------------------------------------------------------------- - - diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index ae82ed01..d9b08ecd 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -12,7 +12,7 @@ jobs: services: qdrant: - image: qdrant/qdrant:latest + image: qdrant/qdrant:v1.15.4 ports: - 6333:6333 diff --git a/.github/workflows/claude.yaml b/.github/workflows/claude.yaml new file mode 100644 index 00000000..732de78c --- /dev/null +++ b/.github/workflows/claude.yaml @@ -0,0 +1,68 @@ +name: Claude Code + +on: + issue_comment: + types: [created] + pull_request_review_comment: + types: [created] + issues: + types: [opened] + pull_request_review: + types: [submitted] + pull_request_target: + types: [opened, synchronize] + +jobs: + claude: + # This simplified condition is more robust and correctly checks permissions. + if: > + (contains(github.event.comment.body, '@claude') || + contains(github.event.review.body, '@claude') || + contains(github.event.issue.body, '@claude') || + contains(github.event.pull_request.body, '@claude')) && + (github.event.sender.type == 'User' && ( + github.event.comment.author_association == 'OWNER' || + github.event.comment.author_association == 'MEMBER' || + github.event.comment.author_association == 'COLLABORATOR' + )) + runs-on: ubuntu-latest + permissions: + # CRITICAL: Write permissions are required for the action to push branches and update issues/PRs. + contents: write + pull-requests: write + issues: write + id-token: write # Required for OIDC token exchange + actions: read # Required for Claude to read CI results on PRs + + steps: + - name: Checkout repository + uses: actions/checkout@v4 + with: + # This correctly checks out the PR's head commit for pull_request_target events. + ref: ${{ github.event.pull_request.head.sha }} + + - name: Create Claude settings file + run: | + mkdir -p /home/runner/.claude + cat > /home/runner/.claude/settings.json << 'EOF' + { + "env": { + "ANTHROPIC_BASE_URL": "https://api.z.ai/api/anthropic", + "ANTHROPIC_AUTH_TOKEN": "${{ secrets.CUSTOM_ENDPOINT_API_KEY }}" + } + } + EOF + + - name: Run Claude Code + id: claude + uses: anthropics/claude-code-action@v1 + with: + # Still need this to satisfy the action's validation + anthropic_api_key: ${{ secrets.CUSTOM_ENDPOINT_API_KEY }} + + # Use the same variable names as your local setup + settings: '{"env": {"ANTHROPIC_BASE_URL": "https://api.z.ai/api/anthropic", "ANTHROPIC_AUTH_TOKEN": "${{ secrets.CUSTOM_ENDPOINT_API_KEY }}"}}' + + track_progress: true + claude_args: | + --allowedTools "Bash,Edit,Read,Write,Glob,Grep" diff --git a/.github/workflows/cosqa-benchmark.yml b/.github/workflows/cosqa-benchmark.yml new file mode 100644 index 00000000..c25a1769 --- /dev/null +++ b/.github/workflows/cosqa-benchmark.yml @@ -0,0 +1,147 @@ +name: CoSQA Search Benchmark + +on: + workflow_dispatch: + inputs: + enforce_hybrid_gate: + description: Fail run if best hybrid underperforms best dense past threshold + required: false + default: false + type: boolean + hybrid_min_delta: + description: Minimum accepted (hybrid_mrr - dense_mrr), e.g. -0.02 + required: false + default: "-0.02" + type: string + upload_full_artifacts: + description: Upload full logs/json bundle (higher storage usage) + required: false + default: false + type: boolean + + pull_request: + branches: [ test ] + paths: + - scripts/hybrid/** + - scripts/hybrid_search.py + - scripts/mcp_impl/search.py + - scripts/mcp_impl/context_search.py + - scripts/mcp_indexer_server.py + - scripts/benchmarks/cosqa/** + - .github/workflows/cosqa-benchmark.yml + + schedule: + - cron: "25 3 * * *" + +jobs: + cosqa-bench: + runs-on: ubuntu-latest + timeout-minutes: 360 + + services: + qdrant: + image: qdrant/qdrant:v1.15.1 + ports: + - 6333:6333 + + steps: + - name: Checkout + uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.11" + + - name: Cache pip + uses: actions/cache@v4 + with: + path: ~/.cache/pip + key: ${{ runner.os }}-pip-${{ hashFiles('**/requirements*.txt', '**/pyproject.toml') }} + restore-keys: | + ${{ runner.os }}-pip- + + - name: Cache HuggingFace datasets + uses: actions/cache@v4 + with: + path: | + ~/.cache/huggingface/datasets + ~/.cache/huggingface/hub + key: ${{ runner.os }}-hf-cosqa-${{ hashFiles('scripts/benchmarks/cosqa/dataset.py') }} + restore-keys: | + ${{ runner.os }}-hf-cosqa- + ${{ runner.os }}-hf- + + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install -r requirements.txt + pip install "datasets>=2.18.0" + + - name: Wait for Qdrant + run: | + timeout 90 bash -c 'until curl -fsS http://localhost:6333/readyz; do sleep 2; done' + curl -fsS http://localhost:6333/collections >/dev/null + + - name: Resolve run config + id: cfg + run: | + echo "profile=full" >> "$GITHUB_OUTPUT" + echo "run_set=full" >> "$GITHUB_OUTPUT" + if [ "${{ github.event_name }}" = "workflow_dispatch" ] && [ "${{ inputs.enforce_hybrid_gate }}" = "true" ]; then + echo "enforce_hybrid_gate=1" >> "$GITHUB_OUTPUT" + else + echo "enforce_hybrid_gate=0" >> "$GITHUB_OUTPUT" + fi + if [ "${{ github.event_name }}" = "workflow_dispatch" ] && [ "${{ inputs.hybrid_min_delta }}" != "" ]; then + echo "hybrid_min_delta=${{ inputs.hybrid_min_delta }}" >> "$GITHUB_OUTPUT" + else + echo "hybrid_min_delta=-0.02" >> "$GITHUB_OUTPUT" + fi + + - name: Run CoSQA search matrix + id: bench + env: + QDRANT_URL: http://localhost:6333 + PROFILE: ${{ steps.cfg.outputs.profile }} + RUN_SET: ${{ steps.cfg.outputs.run_set }} + ENFORCE_HYBRID_GATE: ${{ steps.cfg.outputs.enforce_hybrid_gate }} + HYBRID_MIN_DELTA: ${{ steps.cfg.outputs.hybrid_min_delta }} + PYTHONUNBUFFERED: "1" + run: | + RUN_TAG="gha-${{ github.run_id }}-${{ github.run_attempt }}" + OUT_DIR="bench_results/cosqa/${RUN_TAG}" + echo "out_dir=${OUT_DIR}" >> "$GITHUB_OUTPUT" + RUN_TAG="${RUN_TAG}" OUT_DIR="${OUT_DIR}" ./scripts/benchmarks/cosqa/run_search_matrix.sh + + - name: Publish benchmark summary + if: always() + run: | + SUMMARY="${{ steps.bench.outputs.out_dir }}/summary.md" + if [ -f "${SUMMARY}" ]; then + cat "${SUMMARY}" >> "$GITHUB_STEP_SUMMARY" + else + echo "No summary file generated" >> "$GITHUB_STEP_SUMMARY" + fi + + - name: Upload benchmark artifacts + if: always() && github.event_name == 'pull_request' + uses: actions/upload-artifact@v4 + with: + name: cosqa-search-summary-${{ github.run_id }}-${{ github.run_attempt }} + path: | + ${{ steps.bench.outputs.out_dir }}/summary.md + ${{ steps.bench.outputs.out_dir }}/summary.json + retention-days: 3 + + - name: Upload full benchmark artifacts + if: | + always() && ( + github.event_name == 'schedule' || + (github.event_name == 'workflow_dispatch' && inputs.upload_full_artifacts == true) + ) + uses: actions/upload-artifact@v4 + with: + name: cosqa-search-bench-${{ github.run_id }}-${{ github.run_attempt }} + path: ${{ steps.bench.outputs.out_dir }} + retention-days: 7 diff --git a/Dockerfile b/Dockerfile index c9635c20..15b69498 100644 --- a/Dockerfile +++ b/Dockerfile @@ -19,12 +19,13 @@ RUN pip install --no-cache-dir --upgrade -r /tmp/requirements.txt # Copy scripts for all services COPY scripts /app/scripts +RUN chmod -R a+rX /app/scripts # Create directories -WORKDIR /work +WORKDIR /app # Expose all necessary ports EXPOSE 8000 8001 8002 8003 18000 18001 18002 18003 # Default to memory server -CMD ["python", "/app/scripts/mcp_memory_server.py"] \ No newline at end of file +CMD ["python", "-m", "scripts.mcp_memory_server"] diff --git a/Dockerfile.indexer b/Dockerfile.indexer index 18ce4586..07776b90 100644 --- a/Dockerfile.indexer +++ b/Dockerfile.indexer @@ -29,9 +29,10 @@ ENV RERANKER_ONNX_PATH=/app/models/reranker.onnx \ # Bake scripts into the image so we can mount arbitrary code at /work COPY scripts /app/scripts +RUN chmod -R a+rX /app/scripts -WORKDIR /work +WORKDIR /app # Default command shows help; Makefile/compose will override entrypoint -CMD ["python", "/app/scripts/ingest_code.py", "--help"] +CMD ["python", "-m", "scripts.ingest_code", "--help"] diff --git a/Dockerfile.mcp b/Dockerfile.mcp index a97142ed..0cfc482e 100644 --- a/Dockerfile.mcp +++ b/Dockerfile.mcp @@ -17,9 +17,12 @@ RUN pip install --no-cache-dir --upgrade -r /tmp/requirements.txt \ # Bake scripts into image so server can run even when /work points elsewhere COPY scripts /app/scripts +RUN chmod -R a+rX /app/scripts # Expose SSE port EXPOSE 8000 # Default command: run the server with SSE transport (env provides host/port) -CMD ["python", "/app/scripts/mcp_memory_server.py"] +WORKDIR /app + +CMD ["python", "-m", "scripts.mcp_memory_server"] diff --git a/Dockerfile.mcp-indexer b/Dockerfile.mcp-indexer index 064a9188..28f2a62a 100644 --- a/Dockerfile.mcp-indexer +++ b/Dockerfile.mcp-indexer @@ -32,13 +32,14 @@ ENV RERANKER_ONNX_PATH=/app/models/reranker.onnx \ # Bake scripts into the image so entrypoints don't rely on /work COPY scripts /app/scripts +RUN chmod -R a+rX /app/scripts COPY bench /app/bench # Expose SSE port for this companion server EXPOSE 8001 -WORKDIR /work +WORKDIR /app # Default command runs the companion MCP server -CMD ["python", "/app/scripts/mcp_indexer_server.py"] +CMD ["python", "-m", "scripts.mcp_indexer_server"] diff --git a/Dockerfile.upload-service b/Dockerfile.upload-service index 8c3f47c9..9e15fe81 100644 --- a/Dockerfile.upload-service +++ b/Dockerfile.upload-service @@ -27,6 +27,7 @@ RUN pip install --no-cache-dir --upgrade pip \ # Copy application code COPY scripts/ ./scripts/ COPY . . +RUN chmod -R a+rX /app/scripts # Create work dir, non-root user, and set ownership in single layer RUN mkdir -p /work && chmod 755 /work \ @@ -50,4 +51,6 @@ ENV UPLOAD_SERVICE_HOST=0.0.0.0 \ UPLOAD_TIMEOUT_SECS=300 # Run the upload service -CMD ["python", "scripts/upload_service.py"] \ No newline at end of file +WORKDIR /app + +CMD ["python", "-m", "scripts.upload_service"] diff --git a/Makefile b/Makefile index 068f765f..7c308d6a 100644 --- a/Makefile +++ b/Makefile @@ -4,11 +4,8 @@ SHELL := /bin/bash # An empty export forces docker to use its default context/socket. export DOCKER_HOST = -.PHONY: help up down logs ps restart rebuild index reindex watch watch-remote env hybrid bootstrap history rerank-local setup-reranker prune warm health test-e2e +.PHONY: help up down logs ps restart rebuild index reindex watch watch-remote env hybrid bootstrap history rerank-local setup-reranker prune warm health test test-full test-integration test-e2e .PHONY: venv venv-install dev-remote-up dev-remote-down dev-remote-logs dev-remote-restart dev-remote-bootstrap dev-remote-test dev-remote-client dev-remote-clean -.PHONY: rerank-eval rerank-eval-ablations rerank-benchmark - -.PHONY: qdrant-status qdrant-list qdrant-prune qdrant-index-root venv: ## create local virtualenv .venv python3 -m venv .venv && . .venv/bin/activate && pip install -U pip @@ -73,7 +70,7 @@ index-path: ## index an arbitrary repo: make index-path REPO_PATH=/abs/path [REC @NAME=$${REPO_NAME:-$$(basename "$(REPO_PATH)")}; \ COLL=$${COLLECTION:-$$NAME}; \ HOST_INDEX_PATH="$(REPO_PATH)" COLLECTION_NAME="$$COLL" REPO_NAME="$$NAME" \ - docker compose run --rm -v "$$PWD":/app:ro --entrypoint python indexer /app/scripts/ingest_code.py --root /work $${RECREATE:+--recreate} + docker compose run --rm -v "$$PWD":/app:ro --workdir /app --entrypoint python indexer -m scripts.ingest_code --root /work $${RECREATE:+--recreate} # Index the current working directory quickly index-here: ## index the current directory: make index-here [RECREATE=1] [REPO_NAME=name] [COLLECTION=name] @@ -85,7 +82,7 @@ index-here: ## index the current directory: make index-here [RECREATE=1] [REPO_N watch: ## watch mode: reindex changed files on save (Ctrl+C to stop) - docker compose run --rm --entrypoint python indexer /work/scripts/watch_index.py + docker compose run --rm --workdir /app --entrypoint python indexer -m scripts.watch_index watch-remote: ## remote watch mode: upload delta bundles to remote server (Ctrl+C to stop) @echo "Starting remote watch mode..." @@ -97,24 +94,33 @@ watch-remote: ## remote watch mode: upload delta bundles to remote server (Ctrl+ @echo "Remote upload endpoint: $(REMOTE_UPLOAD_ENDPOINT)" @echo "Max retries: $${REMOTE_UPLOAD_MAX_RETRIES:-3}" @echo "Timeout: $${REMOTE_UPLOAD_TIMEOUT:-30} seconds" - docker compose run --rm --entrypoint python \ + docker compose run --rm --workdir /app --entrypoint python \ -e REMOTE_UPLOAD_ENABLED=1 \ -e REMOTE_UPLOAD_ENDPOINT=$(REMOTE_UPLOAD_ENDPOINT) \ -e REMOTE_UPLOAD_MAX_RETRIES=$${REMOTE_UPLOAD_MAX_RETRIES:-3} \ -e REMOTE_UPLOAD_TIMEOUT=$${REMOTE_UPLOAD_TIMEOUT:-30} \ - indexer /work/scripts/watch_index.py + indexer -m scripts.watch_index rerank: ## multi-query re-ranker helper example - docker compose run --rm --entrypoint python indexer /work/scripts/rerank_query.py \ + docker compose run --rm --workdir /app --entrypoint python indexer -m scripts.rerank_tools.query \ --query "chunk code by lines with overlap for indexing" \ --query "function to split code into overlapping line chunks" \ --language python --under /work/scripts --limit 5 warm: ## prime ANN/search caches with a few queries - docker compose run --rm --entrypoint python indexer /work/scripts/warm_start.py --ef 256 --limit 3 + docker compose run --rm --workdir /app --entrypoint python indexer -m scripts.warm_start --ef 256 --limit 3 health: ## run health checks for collection/model settings - docker compose run --rm --entrypoint python indexer /work/scripts/health_check.py + docker compose run --rm --workdir /app --entrypoint python indexer -m scripts.health_check + +test: ## run default fast tests (excludes integration) + pytest + +test-full: ## run all tests including integration + pytest --run-integration -m "" + +test-integration: ## run integration tests only + pytest --run-integration -m integration # Check llama.cpp decoder health on localhost:8080 (200 OK expected) @@ -128,7 +134,7 @@ env: ## create .env from example if missing [ -f .env ] || cp .env.example .env hybrid: ## hybrid search: dense + lexical RRF fuse (respects --language/--under/--kind) - docker compose run --rm --entrypoint python indexer /work/scripts/hybrid_search.py \ + docker compose run --rm --workdir /app --entrypoint python indexer -m scripts.hybrid_search \ --query "chunk code by lines" --query "overlapping line chunks" --limit 8 bootstrap: env up ## one-shot: up -> wait -> index -> warm -> health @@ -138,20 +144,20 @@ bootstrap: env up ## one-shot: up -> wait -> index -> warm -> health $(MAKE) health history: ## ingest Git history (messages + file lists) - docker compose run --rm --entrypoint python indexer /work/scripts/ingest_history.py --max-commits 200 + docker compose run --rm --workdir /app --entrypoint python indexer -m scripts.ingest_history --max-commits 200 prune-path: ## prune a repo by path: make prune-path REPO_PATH=/abs/path @if [ -z "$(REPO_PATH)" ]; then \ echo "Usage: make prune-path REPO_PATH=/abs/path"; exit 1; \ fi HOST_INDEX_PATH="$(REPO_PATH)" PRUNE_ROOT=/work \ - docker compose run --rm --entrypoint python indexer /work/scripts/prune.py + docker compose run --rm --workdir /app --entrypoint python indexer -m scripts.prune rerank-local: ## local cross-encoder reranker (requires RERANKER_ONNX_PATH, RERANKER_TOKENIZER_PATH) @if [ -z "$(RERANKER_ONNX_PATH)" ] || [ -z "$(RERANKER_TOKENIZER_PATH)" ]; then \ echo "RERANKER_ONNX_PATH and RERANKER_TOKENIZER_PATH must be set in .env"; exit 1; \ fi - docker compose run --rm --entrypoint python indexer /work/scripts/rerank_local.py --query "search symbols" --topk 50 --limit 12 + docker compose run --rm --workdir /app --entrypoint python indexer -m scripts.rerank_tools.local --query "search symbols" --topk 50 --limit 12 setup-reranker: ## download ONNX reranker + tokenizer, update .env, then smoke-test @if [ -z "$(ONNX_URL)" ] || [ -z "$(TOKENIZER_URL)" ]; then \ @@ -163,7 +169,7 @@ setup-reranker: ## download ONNX reranker + tokenizer, update .env, then smoke-t $(MAKE) rerank-local prune: ## remove points for missing files or mismatched file_hash - docker compose run --rm --entrypoint python indexer /work/scripts/prune.py + docker compose run --rm --workdir /app --entrypoint python indexer -m scripts.prune @@ -296,56 +302,6 @@ dev-remote-clean: ## clean up dev-remote volumes and containers rm -rf dev-workspace -# Router helpers -Q ?= what is hybrid search? -route-plan: ## plan-only route for a query: make route-plan Q="your question" - python3 scripts/mcp_router.py --plan "$(Q)" - -route-run: ## execute routed tool(s) over HTTP: make route-run Q="your question" - python3 scripts/mcp_router.py --run "$(Q)" -router-eval: ## run the mock-based router eval harness - python3 scripts/router_eval.py - - -# Live orchestration smoke test (no CI): bring up stack, reindex, run router -router-smoke: ## spin up compose, reindex, store a memory via router, then answer; exits nonzero on failure - set -e; \ - docker compose down || true; \ - docker compose up -d qdrant; \ - ./scripts/wait-for-qdrant.sh; \ - $(MAKE) llama-model; \ - docker compose up -d mcp_http mcp_indexer_http llamacpp; \ - echo "Waiting for MCP HTTP health..."; \ - for i in $$(seq 1 30); do \ - code1=$$(curl -s -o /dev/null -w "%{http_code}" http://localhost:$${FASTMCP_HTTP_HEALTH_PORT:-18002}/readyz || true); \ - code2=$$(curl -s -o /dev/null -w "%{http_code}" http://localhost:$${FASTMCP_INDEXER_HTTP_HEALTH_PORT:-18003}/readyz || true); \ - if [ "$$code1" = "200" ] && [ "$$code2" = "200" ]; then echo "MCP HTTP ready"; break; fi; \ - sleep 1; \ - if [ $$i -eq 30 ]; then echo "MCP HTTP health timeout"; exit 1; fi; \ - done; \ - $(MAKE) reindex; \ - echo "Storing a smoke memory via router..."; \ - python3 scripts/mcp_router.py --run "remember this: router smoke memory"; \ - echo "Running a router answer..."; \ - python3 scripts/mcp_router.py --run "recap our architecture decisions for the indexer"; \ - echo "router-smoke: PASS" - - - -# Qdrant via MCP router convenience targets -qdrant-status: - python3 scripts/mcp_router.py --run "status" - -qdrant-list: - python3 scripts/mcp_router.py --run "list collections" - -qdrant-prune: - python3 scripts/mcp_router.py --run "prune" - -qdrant-index-root: - python3 scripts/mcp_router.py --run "reindex repo" - - # --- ctx CLI helper --- # Usage examples (default prints ONLY the improved prompt): # make ctx Q="how does hybrid search work?" @@ -358,14 +314,3 @@ ctx: ## enhance a prompt with repo context: make ctx Q="your question" [ARGS='-- exit 1; \ fi; \ python3 scripts/ctx.py "$(Q)" $(ARGS) - - -# --- Reranker Evaluation --- -rerank-eval: ## run offline reranker evaluation (fixed queries, MRR/Recall/latency) - python3 scripts/rerank_eval.py --output rerank_eval_results.json - -rerank-eval-ablations: ## run full ablation study (baseline, recursive, learning, onnx) - python3 scripts/rerank_eval.py --ablations --output rerank_eval_ablations.json - -rerank-benchmark: ## run production benchmark on real codebase - python3 scripts/rerank_real_benchmark.py diff --git a/README.md b/README.md index 2a9cd730..6ced1c7f 100644 --- a/README.md +++ b/README.md @@ -9,7 +9,7 @@ ## Context-Engine -Open-source, self-improving code search that gets smarter every time you use it. +Open-source code search engine for AI coding agents — hybrid retrieval with cross-encoder reranking.

Context-Engine Usage @@ -60,7 +60,6 @@ HOST_INDEX_PATH=/path/to/your/project docker compose run --rm indexer | **Large file chunks** → returns entire files | **Precise spans**: Returns 5-50 line chunks, not whole files | | **Lost context** → missing relevant code | **Hybrid search**: Semantic + lexical + cross-encoder reranking | | **Cloud dependency** → vendor lock-in | **Local stack**: Docker Compose on your machine | -| **Static knowledge** → never improves | **Adaptive learning**: Gets smarter with every use | | **Tool limits** → only works in specific IDEs | **MCP native**: Works with any MCP-compatible tool | --- @@ -72,7 +71,7 @@ HOST_INDEX_PATH=/path/to/your/project docker compose run --rm indexer - **Universal compatibility**: Claude Code, Windsurf, Cursor, Cline, etc. - **Auto-syncing**: Extension watches for changes and re-indexes automatically - **Memory system**: Store team knowledge alongside your code -- **Optional LLM features**: Local decoder (llama.cpp), cloud integration (GLM, MiniMax), adaptive rerank learning +- **Optional LLM features**: Local decoder (llama.cpp), cloud integration (GLM, MiniMax) ### Works With Your Local Files No complicated path setup - Context-Engine automatically handles the mapping between your local files and the search index. @@ -144,10 +143,9 @@ See [docs/vscode-extension.md](docs/vscode-extension.md) for full documentation. ## MCP Tools **Search** (Indexer MCP): -- `repo_search` — Hybrid code search with filters +- `repo_search` — Code search with filters and optional profiles - `context_search` — Blend code + memory results - `context_answer` — LLM-generated answers with citations -- `search_tests_for`, `search_config_for`, `search_callers_for` **Memory** (Memory MCP): - `store` — Save knowledge with metadata @@ -190,7 +188,6 @@ flowchart LR M[Memory MCP] Q[(Qdrant)] L[[LLM Decoder]] - W[[Learning Worker]] end V -->|sync| U U --> I @@ -199,8 +196,6 @@ flowchart LR I --> Q M --> Q I -.-> L - I -.-> W - W -.-> Q ``` --- diff --git a/bench/queries.txt b/bench/queries.txt index 478a5d1a..454e0b21 100644 --- a/bench/queries.txt +++ b/bench/queries.txt @@ -5,7 +5,7 @@ # Core search functionality hybrid search RRF ranking implementation memory store implementation details -recursive reranker learning workflow +relevance feedback graph recall workflow # Embeddings and models embedder model loading and caching diff --git a/ctx-mcp-bridge/src/mcpServer.js b/ctx-mcp-bridge/src/mcpServer.js index 53cb05b7..ab4c3bf6 100644 --- a/ctx-mcp-bridge/src/mcpServer.js +++ b/ctx-mcp-bridge/src/mcpServer.js @@ -8,7 +8,13 @@ import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js" import { StreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/streamableHttp.js"; import { Client } from "@modelcontextprotocol/sdk/client/index.js"; import { StreamableHTTPClientTransport } from "@modelcontextprotocol/sdk/client/streamableHttp.js"; -import { CallToolRequestSchema, ListToolsRequestSchema } from "@modelcontextprotocol/sdk/types.js"; +import { + CallToolRequestSchema, + ListToolsRequestSchema, + ListResourcesRequestSchema, + ListResourceTemplatesRequestSchema, + ReadResourceRequestSchema, +} from "@modelcontextprotocol/sdk/types.js"; import { loadAnyAuthEntry, loadAuthEntry, readConfig, saveAuthEntry } from "./authConfig.js"; import { maybeRemapToolArgs, maybeRemapToolResult } from "./resultPathMapping.js"; import * as oauthHandler from "./oauthHandler.js"; @@ -27,16 +33,23 @@ function debugLog(message) { async function sendSessionDefaults(client, payload, label) { if (!client) { - return; + return false; } try { - await client.callTool({ - name: "set_session_defaults", - arguments: payload, - }); + const timeoutMs = getBridgeToolTimeoutMs(); + await withTimeout( + client.callTool({ + name: "set_session_defaults", + arguments: payload, + }), + timeoutMs, + `sendSessionDefaults(${label})` + ); + return true; } catch (err) { // eslint-disable-next-line no-console console.error(`[ctxce] Failed to call set_session_defaults on ${label}:`, err); + return false; } } function dedupeTools(tools) { @@ -58,15 +71,69 @@ function dedupeTools(tools) { return out; } -async function listMemoryTools(client) { - if (!client) { +function dedupeResources(resources) { + const seen = new Set(); + const out = []; + for (const resource of resources) { + const uri = resource && typeof resource.uri === "string" ? resource.uri : ""; + if (!uri || seen.has(uri)) { + continue; + } + seen.add(uri); + out.push(resource); + } + return out; +} + +function dedupeResourceTemplates(templates) { + const seen = new Set(); + const out = []; + for (const template of templates) { + const uri = + template && typeof template.uriTemplate === "string" + ? template.uriTemplate + : ""; + if (!uri || seen.has(uri)) { + continue; + } + seen.add(uri); + out.push(template); + } + return out; +} + +async function callListWithSessionRecovery(call, label, onSessionError) { + try { + return await withTransientRetry(call, label); + } catch (err) { + if (isSessionError(err) && typeof onSessionError === "function") { + try { + await onSessionError(); + return await withTransientRetry(call, `${label} (retry)`); + } catch (retryErr) { + debugLog(`[ctxce] ${label} failed after MCP session recovery: ` + String(retryErr)); + } + } + throw err; + } +} + +async function listMemoryTools(getClient, onSessionError) { + if (typeof getClient !== "function" || !getClient()) { return []; } try { - const remote = await withTimeout( - client.listTools(), - 5000, + const remote = await callListWithSessionRecovery( + () => { + const client = getClient(); + if (!client) { + throw new Error("Memory MCP client not initialized"); + } + const timeoutMs = getBridgeListTimeoutMs(); + return withTimeout(client.listTools(), timeoutMs, "memory tools/list"); + }, "memory tools/list", + onSessionError, ); return Array.isArray(remote?.tools) ? remote.tools.slice() : []; } catch (err) { @@ -75,6 +142,104 @@ async function listMemoryTools(client) { } } +function encodeCompositeCursor(cursorObj) { + try { + const payload = JSON.stringify(cursorObj || {}); + return Buffer.from(payload, "utf8").toString("base64"); + } catch { + return ""; + } +} + +function decodeCompositeCursor(raw) { + try { + const trimmed = (raw || "").trim(); + if (!trimmed) { + return null; + } + const decoded = Buffer.from(trimmed, "base64").toString("utf8"); + const parsed = JSON.parse(decoded); + if (!parsed || typeof parsed !== "object") { + return null; + } + return parsed; + } catch { + return null; + } +} + +async function listResourcesSafe(getClient, label, cursor, onSessionError) { + if (typeof getClient !== "function" || !getClient()) { + return { resources: [], nextCursor: null }; + } + try { + const params = cursor ? { cursor } : {}; + const remote = await callListWithSessionRecovery( + () => { + const client = getClient(); + if (!client) { + throw new Error(`${label} MCP client not initialized`); + } + const timeoutMs = getBridgeListTimeoutMs(); + return withTimeout( + client.listResources(params), + timeoutMs, + `${label} resources/list`, + ); + }, + `${label} resources/list`, + onSessionError, + ); + return { + resources: Array.isArray(remote?.resources) ? remote.resources.slice() : [], + nextCursor: + remote && typeof remote.nextCursor === "string" && remote.nextCursor + ? remote.nextCursor + : null, + }; + } catch (err) { + debugLog(`[ctxce] Error calling ${label} resources/list: ` + String(err)); + return { resources: [], nextCursor: null }; + } +} + +async function listResourceTemplatesSafe(getClient, label, cursor, onSessionError) { + if (typeof getClient !== "function" || !getClient()) { + return { resourceTemplates: [], nextCursor: null }; + } + try { + const params = cursor ? { cursor } : {}; + const remote = await callListWithSessionRecovery( + () => { + const client = getClient(); + if (!client) { + throw new Error(`${label} MCP client not initialized`); + } + const timeoutMs = getBridgeListTimeoutMs(); + return withTimeout( + client.listResourceTemplates(params), + timeoutMs, + `${label} resources/templates/list`, + ); + }, + `${label} resources/templates/list`, + onSessionError, + ); + return { + resourceTemplates: Array.isArray(remote?.resourceTemplates) + ? remote.resourceTemplates.slice() + : [], + nextCursor: + remote && typeof remote.nextCursor === "string" && remote.nextCursor + ? remote.nextCursor + : null, + }; + } catch (err) { + debugLog(`[ctxce] Error calling ${label} resources/templates/list: ` + String(err)); + return { resourceTemplates: [], nextCursor: null }; + } +} + function withTimeout(promise, ms, label) { return new Promise((resolve, reject) => { let settled = false; @@ -125,6 +290,25 @@ function getBridgeToolTimeoutMs() { } } +function getBridgeListTimeoutMs() { + try { + // Keep list operations on a separate budget from tools/call. + // Some streamable-http clients (including Codex) probe tools/resources early, + // and a short timeout here can make the bridge appear unavailable. + const raw = process.env.CTXCE_LIST_TIMEOUT_MSEC; + if (!raw) { + return 60000; + } + const parsed = Number.parseInt(String(raw), 10); + if (!Number.isFinite(parsed) || parsed <= 0) { + return 60000; + } + return parsed; + } catch { + return 60000; + } +} + function selectClientForTool(name, indexerClient, memoryClient) { if (!name) { return indexerClient; @@ -148,6 +332,7 @@ function isSessionError(error) { msg.includes("No valid session ID") || msg.includes("Mcp-Session-Id header is required") || msg.includes("Server not initialized") || + msg.includes("Received request before initialization was complete") || msg.includes("Session not found") ); } catch { @@ -270,6 +455,34 @@ function isTransientToolError(error) { return false; } } + +async function withTransientRetry(operation, label, maxAttempts, retryDelayMs) { + const attempts = Number.isFinite(maxAttempts) && maxAttempts > 0 + ? Math.floor(maxAttempts) + : getBridgeRetryAttempts(); + const delayMs = Number.isFinite(retryDelayMs) && retryDelayMs >= 0 + ? Math.floor(retryDelayMs) + : getBridgeRetryDelayMs(); + let lastError; + for (let attempt = 0; attempt < attempts; attempt += 1) { + if (attempt > 0 && delayMs > 0) { + await new Promise((resolve) => setTimeout(resolve, delayMs)); + } + try { + return await operation(); + } catch (err) { + lastError = err; + if (!isTransientToolError(err) || attempt === attempts - 1) { + throw err; + } + debugLog( + `[ctxce] ${label}: transient error (attempt ${attempt + 1}/${attempts}), retrying: ` + + String(err), + ); + } + } + throw lastError || new Error(`[ctxce] ${label}: unknown transient retry failure`); +} // MCP stdio server implemented using the official MCP TypeScript SDK. // Acts as a low-level proxy for tools, forwarding tools/list and tools/call // to the remote qdrant-indexer MCP server while adding a local `ping` tool. @@ -440,6 +653,7 @@ async function createBridgeServer(options) { let indexerClient = null; let memoryClient = null; + let lastDefaultsSyncedSessionId = ""; // Derive a simple session identifier for this bridge process. In the // future this can be made user-aware (e.g. from auth), but for now we @@ -568,6 +782,23 @@ async function createBridgeServer(options) { defaultsPayload.under = defaultUnder; } + async function ensureRemoteDefaults(force = false) { + defaultsPayload.session = sessionId; + if (!sessionId) { + return; + } + if (!force && lastDefaultsSyncedSessionId === sessionId) { + return; + } + const indexerOk = await sendSessionDefaults(indexerClient, defaultsPayload, "indexer"); + if (memoryClient) { + await sendSessionDefaults(memoryClient, defaultsPayload, "memory"); + } + if (indexerOk) { + lastDefaultsSyncedSessionId = sessionId; + } + } + async function initializeRemoteClients(forceRecreate = false) { if (!forceRecreate && indexerClient) { return; @@ -579,6 +810,22 @@ async function createBridgeServer(options) { } catch { // ignore logging failures } + try { + if (indexerClient && typeof indexerClient.close === "function") { + await indexerClient.close(); + } + } catch { + // ignore + } + try { + if (memoryClient && typeof memoryClient.close === "function") { + await memoryClient.close(); + } + } catch { + // ignore + } + indexerClient = null; + memoryClient = null; } let nextIndexerClient = null; @@ -633,15 +880,34 @@ async function createBridgeServer(options) { indexerClient = nextIndexerClient; memoryClient = nextMemoryClient; - if (Object.keys(defaultsPayload).length > 1 && indexerClient) { - await sendSessionDefaults(indexerClient, defaultsPayload, "indexer"); - if (memoryClient) { - await sendSessionDefaults(memoryClient, defaultsPayload, "memory"); - } + await ensureRemoteDefaults(true); + } + + async function refreshSessionAndSyncDefaults() { + const freshSession = resolveSessionId() || sessionId; + const changed = Boolean(freshSession && freshSession !== sessionId); + if (changed) { + sessionId = freshSession; + defaultsPayload.session = sessionId; + lastDefaultsSyncedSessionId = ""; } + await initializeRemoteClients(false); + await ensureRemoteDefaults(changed); } - await initializeRemoteClients(false); + async function recoverRemoteClientsAfterSessionError() { + const freshSession = resolveSessionId() || sessionId; + const changed = Boolean(freshSession && freshSession !== sessionId); + if (changed) { + sessionId = freshSession; + defaultsPayload.session = sessionId; + lastDefaultsSyncedSessionId = ""; + } + await initializeRemoteClients(true); + await ensureRemoteDefaults(true); + } + + await refreshSessionAndSyncDefaults(); const server = new Server( // TODO: marked as depreciated { @@ -651,6 +917,7 @@ async function createBridgeServer(options) { { capabilities: { tools: {}, + resources: {}, }, }, ); @@ -658,20 +925,36 @@ async function createBridgeServer(options) { // tools/list → fetch tools from remote indexer server.setRequestHandler(ListToolsRequestSchema, async () => { let remote; + let listError = null; try { - debugLog("[ctxce] tools/list: fetching tools from indexer"); await initializeRemoteClients(false); - if (!indexerClient) { - throw new Error("Indexer MCP client not initialized"); - } - remote = await withTimeout( - indexerClient.listTools(), - 10000, + await ensureRemoteDefaults(false); + debugLog("[ctxce] tools/list: fetching tools from indexer"); + remote = await callListWithSessionRecovery( + () => { + if (!indexerClient) { + throw new Error("Indexer MCP client not initialized"); + } + const timeoutMs = getBridgeListTimeoutMs(); + return withTimeout( + indexerClient.listTools(), + timeoutMs, + "indexer tools/list", + ); + }, "indexer tools/list", + recoverRemoteClientsAfterSessionError, ); } catch (err) { - debugLog("[ctxce] Error calling remote tools/list: " + String(err)); - const memoryToolsFallback = await listMemoryTools(memoryClient); + listError = err; + } + + if (!remote) { + debugLog("[ctxce] Error calling remote tools/list: " + String(listError)); + const memoryToolsFallback = await listMemoryTools( + () => memoryClient, + recoverRemoteClientsAfterSessionError, + ); const toolsFallback = dedupeTools([...memoryToolsFallback]); return { tools: toolsFallback }; } @@ -687,12 +970,130 @@ async function createBridgeServer(options) { } const indexerTools = Array.isArray(remote?.tools) ? remote.tools.slice() : []; - const memoryTools = await listMemoryTools(memoryClient); + const memoryTools = await listMemoryTools( + () => memoryClient, + recoverRemoteClientsAfterSessionError, + ); const tools = dedupeTools([...indexerTools, ...memoryTools]); debugLog(`[ctxce] tools/list: returning ${tools.length} tools`); return { tools }; }); + server.setRequestHandler(ListResourcesRequestSchema, async (request) => { + // Proxy resource discovery/read-through so clients that use MCP resources + // (not only tools) can access upstream indexer/memory resources directly. + await initializeRemoteClients(false); + await ensureRemoteDefaults(false); + const cursor = + request && request.params && typeof request.params.cursor === "string" + ? request.params.cursor + : null; + const decoded = decodeCompositeCursor(cursor); + const indexerCursor = + decoded && typeof decoded.i === "string" ? decoded.i : cursor; + const memoryCursor = + decoded && typeof decoded.m === "string" ? decoded.m : cursor; + if (cursor && decoded === null) { + debugLog("[ctxce] resources/list: received non-composite cursor; forwarding to both upstreams."); + } + const indexerRes = await listResourcesSafe( + () => indexerClient, + "indexer", + indexerCursor, + recoverRemoteClientsAfterSessionError, + ); + const memoryRes = await listResourcesSafe( + () => memoryClient, + "memory", + memoryCursor, + recoverRemoteClientsAfterSessionError, + ); + const resources = dedupeResources([ + ...indexerRes.resources, + ...memoryRes.resources, + ]); + const nextCursorObj = { + i: indexerRes.nextCursor || "", + m: memoryRes.nextCursor || "", + }; + const nextCursor = + nextCursorObj.i || nextCursorObj.m ? encodeCompositeCursor(nextCursorObj) : ""; + debugLog(`[ctxce] resources/list: returning ${resources.length} resources`); + return nextCursor ? { resources, nextCursor } : { resources }; + }); + + server.setRequestHandler(ListResourceTemplatesRequestSchema, async (request) => { + await initializeRemoteClients(false); + await ensureRemoteDefaults(false); + const cursor = + request && request.params && typeof request.params.cursor === "string" + ? request.params.cursor + : null; + const decoded = decodeCompositeCursor(cursor); + const indexerCursor = + decoded && typeof decoded.i === "string" ? decoded.i : cursor; + const memoryCursor = + decoded && typeof decoded.m === "string" ? decoded.m : cursor; + if (cursor && decoded === null) { + debugLog("[ctxce] resources/templates/list: received non-composite cursor; forwarding to both upstreams."); + } + const indexerRes = await listResourceTemplatesSafe( + () => indexerClient, + "indexer", + indexerCursor, + recoverRemoteClientsAfterSessionError, + ); + const memoryRes = await listResourceTemplatesSafe( + () => memoryClient, + "memory", + memoryCursor, + recoverRemoteClientsAfterSessionError, + ); + const resourceTemplates = dedupeResourceTemplates([ + ...indexerRes.resourceTemplates, + ...memoryRes.resourceTemplates, + ]); + const nextCursorObj = { + i: indexerRes.nextCursor || "", + m: memoryRes.nextCursor || "", + }; + const nextCursor = + nextCursorObj.i || nextCursorObj.m ? encodeCompositeCursor(nextCursorObj) : ""; + debugLog(`[ctxce] resources/templates/list: returning ${resourceTemplates.length} templates`); + return nextCursor ? { resourceTemplates, nextCursor } : { resourceTemplates }; + }); + + server.setRequestHandler(ReadResourceRequestSchema, async (request) => { + await refreshSessionAndSyncDefaults(); + const params = request.params || {}; + const timeoutMs = getBridgeToolTimeoutMs(); + const uri = + params && typeof params.uri === "string" ? params.uri : ""; + debugLog(`[ctxce] resources/read: ${uri}`); + + const tryRead = async (client, label) => { + if (!client) { + return null; + } + try { + return await client.readResource(params, { timeout: timeoutMs }); + } catch (err) { + debugLog(`[ctxce] resources/read failed on ${label}: ` + String(err)); + return null; + } + }; + + const indexerResult = await tryRead(indexerClient, "indexer"); + if (indexerResult) { + return indexerResult; + } + const memoryResult = await tryRead(memoryClient, "memory"); + if (memoryResult) { + return memoryResult; + } + throw new Error(`Resource ${uri} not available on any configured MCP server`); + }); + // tools/call → proxied to indexer or memory server server.setRequestHandler(CallToolRequestSchema, async (request) => { const params = request.params || {}; @@ -701,16 +1102,8 @@ async function createBridgeServer(options) { debugLog(`[ctxce] tools/call: ${name || ""}`); - // Refresh session before each call; re-init clients if session changes. - const freshSession = resolveSessionId() || sessionId; - if (freshSession && freshSession !== sessionId) { - sessionId = freshSession; - try { - await initializeRemoteClients(true); - } catch (err) { - debugLog("[ctxce] Failed to reinitialize clients after session refresh: " + String(err)); - } - } + await refreshSessionAndSyncDefaults(); + if (sessionId && (args === undefined || args === null || typeof args === "object")) { const obj = args && typeof args === "object" ? { ...args } : {}; if (!Object.prototype.hasOwnProperty.call(obj, "session")) { @@ -733,8 +1126,6 @@ async function createBridgeServer(options) { return indexerResult; } - await initializeRemoteClients(false); - const timeoutMs = getBridgeToolTimeoutMs(); const maxAttempts = getBridgeRetryAttempts(); const retryDelayMs = getBridgeRetryDelayMs(); @@ -770,6 +1161,7 @@ async function createBridgeServer(options) { String(err), ); await initializeRemoteClients(true); + await ensureRemoteDefaults(true); sessionRetried = true; continue; } @@ -843,6 +1235,13 @@ export async function runHttpMcpServer(options) { typeof options.port === "number" ? options.port : Number.parseInt(process.env.CTXCE_HTTP_PORT || "30810", 10) || 30810; + // TODO(auth): replace this boolean toggle with explicit auth modes (none|required). + // In required mode, enforce Bearer auth on /mcp with consistent 401 challenges and + // only advertise OAuth metadata/endpoints when authentication is mandatory. + // In local/dev mode, leaving OAuth discovery off avoids clients entering an + // unnecessary OAuth path for otherwise unauthenticated bridge usage. + const oauthEnabled = String(process.env.CTXCE_ENABLE_OAUTH || "").trim().toLowerCase(); + const oauthEndpointsEnabled = oauthEnabled === "1" || oauthEnabled === "true" || oauthEnabled === "yes"; const transport = new StreamableHTTPServerTransport({ sessionIdGenerator: undefined, @@ -865,34 +1264,36 @@ export async function runHttpMcpServer(options) { // OAuth 2.0 Endpoints (RFC9728 Protected Resource Metadata + RFC7591) // ================================================================ - // OAuth metadata endpoint (RFC9728) - if (parsedUrl.pathname === "/.well-known/oauth-authorization-server") { - oauthHandler.handleOAuthMetadata(req, res, issuerUrl); - return; - } + if (oauthEndpointsEnabled) { + // OAuth metadata endpoint (RFC9728) + if (parsedUrl.pathname === "/.well-known/oauth-authorization-server") { + oauthHandler.handleOAuthMetadata(req, res, issuerUrl); + return; + } - // OAuth Dynamic Client Registration endpoint (RFC7591) - if (parsedUrl.pathname === "/oauth/register" && req.method === "POST") { - oauthHandler.handleOAuthRegister(req, res); - return; - } + // OAuth Dynamic Client Registration endpoint (RFC7591) + if (parsedUrl.pathname === "/oauth/register" && req.method === "POST") { + oauthHandler.handleOAuthRegister(req, res); + return; + } - // OAuth authorize endpoint - if (parsedUrl.pathname === "/oauth/authorize") { - oauthHandler.handleOAuthAuthorize(req, res, parsedUrl.searchParams); - return; - } + // OAuth authorize endpoint + if (parsedUrl.pathname === "/oauth/authorize") { + oauthHandler.handleOAuthAuthorize(req, res, parsedUrl.searchParams); + return; + } - // Store session endpoint (helper for login page) - if (parsedUrl.pathname === "/oauth/store-session" && req.method === "POST") { - oauthHandler.handleOAuthStoreSession(req, res); - return; - } + // Store session endpoint (helper for login page) + if (parsedUrl.pathname === "/oauth/store-session" && req.method === "POST") { + oauthHandler.handleOAuthStoreSession(req, res); + return; + } - // OAuth token endpoint - if (parsedUrl.pathname === "/oauth/token" && req.method === "POST") { - oauthHandler.handleOAuthToken(req, res); - return; + // OAuth token endpoint + if (parsedUrl.pathname === "/oauth/token" && req.method === "POST") { + oauthHandler.handleOAuthToken(req, res); + return; + } } // ================================================================ @@ -1058,4 +1459,3 @@ function detectRepoName(workspace, config) { const leaf = workspace ? path.basename(workspace) : ""; return leaf && SLUGGED_REPO_RE.test(leaf) ? leaf : null; } - diff --git a/ctx-mcp-bridge/src/resultPathMapping.js b/ctx-mcp-bridge/src/resultPathMapping.js index a3309876..46e9663d 100644 --- a/ctx-mcp-bridge/src/resultPathMapping.js +++ b/ctx-mcp-bridge/src/resultPathMapping.js @@ -442,11 +442,7 @@ export function maybeRemapToolResult(name, result, workspaceRoot) { const shouldMap = ( lower === "repo_search" || lower === "context_search" || - lower === "context_answer" || - lower.endsWith("search_tests_for") || - lower.endsWith("search_config_for") || - lower.endsWith("search_callers_for") || - lower.endsWith("search_importers_for") + lower === "context_answer" ); if (!shouldMap) { return result; diff --git a/deploy/kubernetes/configmap.yaml b/deploy/kubernetes/configmap.yaml index 9c91a893..ecf81209 100644 --- a/deploy/kubernetes/configmap.yaml +++ b/deploy/kubernetes/configmap.yaml @@ -45,10 +45,6 @@ data: INDEX_UPSERT_BATCH: '128' INDEX_UPSERT_RETRIES: '5' INDEX_USE_ENHANCED_AST: '1' - INFO_REQUEST_CONTEXT_LINES: '5' - INFO_REQUEST_EXPLAIN_DEFAULT: '0' - INFO_REQUEST_LIMIT: '10' - INFO_REQUEST_RELATIONSHIPS: '0' LLAMACPP_EXTRA_ARGS: '' LLAMACPP_GPU_LAYERS: '32' LLAMACPP_GPU_SPLIT: '' @@ -115,15 +111,8 @@ data: RERANKER_TIMEOUT_MS: '3000' RERANKER_TOKENIZER_PATH: /app/models/tokenizer.json RERANKER_TOPN: '100' - RERANK_LEARNING: '0' RERANKER_WEIGHTS_DIR: /tmp/rerank_weights RERANK_EVENTS_DIR: /tmp/rerank_events - RERANK_EVENTS_ENABLED: '0' - RERANK_EVENTS_SAMPLE_RATE: '0.33' - RERANK_EVENTS_RETENTION_DAYS: '0' - RERANK_LEARNING_BATCH_SIZE: '32' - RERANK_LEARNING_POLL_INTERVAL: '30' - RERANK_LEARNING_RATE: '0.001' RERANK_BLEND_WEIGHT: '0.6' RERANK_EXPAND: '1' RERANK_IN_PROCESS: '1' @@ -143,3 +132,9 @@ data: USE_GPU_DECODER: '0' USE_TREE_SITTER: '1' WATCH_DEBOUNCE_SECS: '4' + WATCH_JOURNAL_DRAIN_BATCH_SIZE: '256' + WATCH_JOURNAL_LOG_INTERVAL_SECS: '60' + WATCH_USE_POLLING: '1' + WATCH_INIT_MAINTENANCE_ENABLED: '1' + WATCH_INIT_MAINTENANCE_INTERVAL_MINUTES: '120' + WATCH_INIT_MAINTENANCE_RUN_ON_START: '0' diff --git a/deploy/kubernetes/indexer-services.yaml b/deploy/kubernetes/indexer-services.yaml index 123582c6..60d98cef 100644 --- a/deploy/kubernetes/indexer-services.yaml +++ b/deploy/kubernetes/indexer-services.yaml @@ -53,8 +53,9 @@ spec: imagePullPolicy: IfNotPresent command: - python - - /app/scripts/watch_index.py - workingDir: /work + - -m + - scripts.watch_index + workingDir: /app env: - name: QDRANT_URL valueFrom: @@ -102,6 +103,11 @@ spec: configMapKeyRef: name: context-engine-config key: WATCH_DEBOUNCE_SECS + - name: WATCH_USE_POLLING + valueFrom: + configMapKeyRef: + name: context-engine-config + key: WATCH_USE_POLLING - name: HF_HOME value: /work/models/hf-cache - name: XDG_CACHE_HOME @@ -191,8 +197,11 @@ spec: imagePullPolicy: IfNotPresent command: - python - - /app/scripts/ingest_code.py - workingDir: /work + - -m + - scripts.ingest_code + - --root + - /work + workingDir: /app env: - name: QDRANT_URL valueFrom: @@ -299,8 +308,8 @@ spec: command: - /bin/sh - -c - - PYTHONPATH=/app python /app/scripts/create_indexes.py && PYTHONPATH=/app python /app/scripts/warm_all_collections.py && PYTHONPATH=/app python /app/scripts/health_check.py - workingDir: /work + - python -m scripts.run_init_maintenance + workingDir: /app env: - name: QDRANT_URL valueFrom: diff --git a/deploy/kubernetes/kustomization.yaml b/deploy/kubernetes/kustomization.yaml index bf622c85..d4d55cbe 100644 --- a/deploy/kubernetes/kustomization.yaml +++ b/deploy/kubernetes/kustomization.yaml @@ -18,7 +18,6 @@ resources: - mcp-memory.yaml - mcp-indexer.yaml - mcp-http.yaml - - learning-reranker-worker.yaml # Indexer services - indexer-services.yaml diff --git a/deploy/kubernetes/learning-reranker-worker.yaml b/deploy/kubernetes/learning-reranker-worker.yaml deleted file mode 100644 index c1431c22..00000000 --- a/deploy/kubernetes/learning-reranker-worker.yaml +++ /dev/null @@ -1,65 +0,0 @@ -apiVersion: apps/v1 -kind: Deployment -metadata: - name: learning-reranker-worker - namespace: context-engine - labels: - app: context-engine - component: learning-reranker-worker -spec: - replicas: 1 - selector: - matchLabels: - app: context-engine - component: learning-reranker-worker - template: - metadata: - labels: - app: context-engine - component: learning-reranker-worker - spec: - serviceAccountName: context-engine - securityContext: - runAsUser: 1000 - runAsGroup: 1000 - fsGroup: 1000 - initContainers: - - name: init-rerank-dirs - image: busybox:1.36 - imagePullPolicy: IfNotPresent - command: - - sh - - -c - - mkdir -p /mnt/rerank_weights /mnt/rerank_events && chmod 777 /mnt/rerank_weights /mnt/rerank_events - volumeMounts: - - name: metadata-volume - mountPath: /mnt - containers: - - name: learning-reranker-worker - image: context-engine-indexer - imagePullPolicy: IfNotPresent - command: - - python - - /app/scripts/learning_reranker_worker.py - - --daemon - resources: - requests: - memory: 512Mi - cpu: 250m - limits: - memory: 2Gi - cpu: 1000m - volumeMounts: - - name: metadata-volume - mountPath: /tmp/rerank_weights - subPath: rerank_weights - - name: metadata-volume - mountPath: /tmp/rerank_events - subPath: rerank_events - envFrom: - - configMapRef: - name: context-engine-config - volumes: - - name: metadata-volume - persistentVolumeClaim: - claimName: code-metadata-pvc diff --git a/deploy/kubernetes/mcp-http.yaml b/deploy/kubernetes/mcp-http.yaml index c3c71fe2..34d6f791 100644 --- a/deploy/kubernetes/mcp-http.yaml +++ b/deploy/kubernetes/mcp-http.yaml @@ -40,7 +40,9 @@ spec: imagePullPolicy: IfNotPresent command: - python - - /app/scripts/mcp_memory_server.py + - -m + - scripts.mcp_memory_server + workingDir: /app ports: - name: http containerPort: 8000 @@ -59,6 +61,11 @@ spec: configMapKeyRef: name: context-engine-config key: COLLECTION_NAME + - name: MULTI_REPO_MODE + valueFrom: + configMapKeyRef: + name: context-engine-config + key: MULTI_REPO_MODE - name: EMBEDDING_MODEL valueFrom: configMapKeyRef: @@ -227,7 +234,9 @@ spec: imagePullPolicy: IfNotPresent command: - python - - /app/scripts/mcp_indexer_server.py + - -m + - scripts.mcp_indexer_server + workingDir: /app ports: - name: http containerPort: 8001 @@ -246,6 +255,11 @@ spec: configMapKeyRef: name: context-engine-config key: COLLECTION_NAME + - name: MULTI_REPO_MODE + valueFrom: + configMapKeyRef: + name: context-engine-config + key: MULTI_REPO_MODE - name: EMBEDDING_MODEL valueFrom: configMapKeyRef: diff --git a/deploy/kubernetes/mcp-indexer.yaml b/deploy/kubernetes/mcp-indexer.yaml index 2fbec1a1..a56f190d 100644 --- a/deploy/kubernetes/mcp-indexer.yaml +++ b/deploy/kubernetes/mcp-indexer.yaml @@ -40,7 +40,9 @@ spec: imagePullPolicy: IfNotPresent command: - python - - /app/scripts/mcp_indexer_server.py + - -m + - scripts.mcp_indexer_server + workingDir: /app ports: - name: sse containerPort: 8001 diff --git a/deploy/kubernetes/mcp-memory.yaml b/deploy/kubernetes/mcp-memory.yaml index 165076db..0000f2ec 100644 --- a/deploy/kubernetes/mcp-memory.yaml +++ b/deploy/kubernetes/mcp-memory.yaml @@ -36,7 +36,9 @@ spec: imagePullPolicy: IfNotPresent command: - python - - /app/scripts/mcp_memory_server.py + - -m + - scripts.mcp_memory_server + workingDir: /app ports: - name: sse containerPort: 8000 @@ -69,6 +71,11 @@ spec: configMapKeyRef: name: context-engine-config key: COLLECTION_NAME + - name: MULTI_REPO_MODE + valueFrom: + configMapKeyRef: + name: context-engine-config + key: MULTI_REPO_MODE - name: EMBEDDING_MODEL valueFrom: configMapKeyRef: diff --git a/deploy/kubernetes/qdrant.yaml b/deploy/kubernetes/qdrant.yaml index ba645364..3e72352c 100644 --- a/deploy/kubernetes/qdrant.yaml +++ b/deploy/kubernetes/qdrant.yaml @@ -23,7 +23,7 @@ spec: spec: containers: - name: qdrant - image: qdrant/qdrant:latest + image: qdrant/qdrant:v1.15.4 imagePullPolicy: Always ports: - name: http diff --git a/docker-compose-bindmount-checkout.yml b/docker-compose-bindmount-checkout.yml index b4e6361c..19f65c60 100644 --- a/docker-compose-bindmount-checkout.yml +++ b/docker-compose-bindmount-checkout.yml @@ -1,6 +1,6 @@ services: qdrant: - image: qdrant/qdrant:latest + image: qdrant/qdrant:v1.15.4 container_name: qdrant-db # Expose Qdrant database APIs to the host # 6333 = HTTP API, 6334 = gRPC @@ -211,7 +211,7 @@ services: - ${HOST_INDEX_PATH:-.}:/work:ro - ${HOST_INDEX_PATH:-.}/.codebase:/work/.codebase:rw - entrypoint: ["sh", "-c", "mkdir -p /tmp/huggingface/hub /tmp/huggingface/transformers /tmp/huggingface/fastembed && exec python /app/scripts/ingest_code.py"] + entrypoint: ["sh", "-c", "mkdir -p /tmp/huggingface/hub /tmp/huggingface/transformers /tmp/huggingface/fastembed && cd /app && exec python -m scripts.ingest_code --root /work"] watcher: build: @@ -235,6 +235,7 @@ services: - QWEN3_QUERY_INSTRUCTION=${QWEN3_QUERY_INSTRUCTION:-1} - QWEN3_INSTRUCTION_TEXT=${QWEN3_INSTRUCTION_TEXT} - WATCH_ROOT=/work + - CTXCE_METADATA_ROOT=${CTXCE_METADATA_ROOT:-/work} # Watcher-specific backpressure & timeouts (safer defaults) - QDRANT_TIMEOUT=60 - MAX_MICRO_CHUNKS_PER_FILE=${MAX_MICRO_CHUNKS_PER_FILE:-200} @@ -245,7 +246,7 @@ services: volumes: - ${HOST_INDEX_PATH:-.}:/work:ro - ${HOST_INDEX_PATH:-.}/.codebase:/work/.codebase:rw - entrypoint: ["sh", "-c", "mkdir -p /tmp/huggingface/hub /tmp/huggingface/transformers /tmp/huggingface/fastembed && exec python /app/scripts/watch_index.py"] + entrypoint: ["sh", "-c", "mkdir -p /tmp/huggingface/hub /tmp/huggingface/transformers /tmp/huggingface/fastembed && cd /app && exec python -m scripts.watch_index"] upload_service: @@ -262,6 +263,7 @@ services: - UPLOAD_SERVICE_PORT=8002 - QDRANT_URL=${QDRANT_URL} - WORK_DIR=/work + - CTXCE_METADATA_ROOT=${CTXCE_METADATA_ROOT:-/work} - CTXCE_ADMIN_COLLECTION_DELETE_ENABLED=0 - CTXCE_COLLECTION_REGISTRY_UNDELETE_ON_DISCOVERY=${CTXCE_COLLECTION_REGISTRY_UNDELETE_ON_DISCOVERY:-0} - COLLECTION_NAME=${COLLECTION_NAME:-codebase} @@ -300,7 +302,7 @@ services: - ${HOST_INDEX_PATH:-.}:/work:ro - ${HOST_INDEX_PATH:-.}/.codebase:/work/.codebase:rw - entrypoint: ["python", "/app/scripts/create_indexes.py"] + entrypoint: ["python", "-m", "scripts.create_indexes"] volumes: qdrant_storage: diff --git a/docker-compose.yml b/docker-compose.yml index cb5403a2..a7ade7a5 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -8,7 +8,7 @@ version: '3.8' services: # Qdrant vector database - same as base compose qdrant: - image: qdrant/qdrant:latest + image: qdrant/qdrant:v1.15.4 container_name: qdrant-db-dev-remote ports: - "6333:6333" @@ -66,11 +66,9 @@ services: - RERANKER_MODEL=${RERANKER_MODEL:-} - RERANKER_ONNX_PATH=${RERANKER_ONNX_PATH:-} - RERANKER_TOKENIZER_PATH=${RERANKER_TOKENIZER_PATH:-} - # Learning reranker configuration - - RERANK_LEARNING=${RERANK_LEARNING:-1} + # Relevance feedback configuration - RERANKER_WEIGHTS_DIR=/tmp/rerank_weights - RERANK_EVENTS_DIR=/tmp/rerank_events - - RERANK_EVENTS_ENABLED=${RERANK_EVENTS_ENABLED:-1} ports: - "18000:18000" - "8000:8000" @@ -88,11 +86,11 @@ services: dockerfile: Dockerfile.mcp-indexer container_name: mcp-indexer-dev-remote user: "1000:1000" - # In K8s, scripts would be accessed directly at /app/scripts/ or via proper initContainer + # In K8s, scripts run as package modules from /app or via proper initContainer # For Docker Compose dev-remote simulation, create symlink so /work/scripts/ works # Use /tmp/huggingface for cache to avoid permission issues (universally writable) # Set CORRECT environment variables for HuggingFace and FastEmbed - command: ["sh", "-c", "mkdir -p /tmp/huggingface/hub /tmp/huggingface/transformers /tmp/huggingface/fastembed && exec python /app/scripts/mcp_indexer_server.py"] + command: ["sh", "-c", "mkdir -p /tmp/huggingface/hub /tmp/huggingface/transformers /tmp/huggingface/fastembed && cd /app && exec python -m scripts.mcp_indexer_server"] depends_on: - qdrant env_file: @@ -148,17 +146,13 @@ services: - LEX_BIGRAM_WEIGHT=${LEX_BIGRAM_WEIGHT:-} - LEX_SPARSE_MODE=${LEX_SPARSE_MODE:-} - LEX_SPARSE_NAME=${LEX_SPARSE_NAME:-} - # Pattern vectors for structural code similarity - - PATTERN_VECTORS=${PATTERN_VECTORS:-} # Cross-encoder reranker configuration - RERANKER_MODEL=${RERANKER_MODEL:-} - RERANKER_ONNX_PATH=${RERANKER_ONNX_PATH:-} - RERANKER_TOKENIZER_PATH=${RERANKER_TOKENIZER_PATH:-} - # Learning reranker configuration - - RERANK_LEARNING=${RERANK_LEARNING:-1} + # Relevance feedback configuration - RERANKER_WEIGHTS_DIR=/tmp/rerank_weights - RERANK_EVENTS_DIR=/tmp/rerank_events - - RERANK_EVENTS_ENABLED=${RERANK_EVENTS_ENABLED:-1} ports: - "${FASTMCP_INDEXER_PORT:-8001}:8001" - "18001:18001" @@ -170,14 +164,14 @@ services: networks: - dev-remote-network - # Learning reranker worker - processes training events in background - learning_worker: + # Relevance trainer - processes feedback events, updates per-collection weights + relevance_trainer: build: context: . dockerfile: Dockerfile.mcp-indexer - container_name: learning-worker-dev-remote + container_name: relevance-trainer-dev-remote user: "1000:1000" - command: ["sh", "-c", "mkdir -p /tmp/huggingface/hub /tmp/huggingface/transformers /tmp/huggingface/fastembed && exec python /app/scripts/learning_reranker_worker.py --daemon"] + command: ["sh", "-c", "cd /app && exec python -m scripts.relevance_trainer --daemon"] depends_on: - qdrant - mcp_indexer @@ -186,30 +180,13 @@ services: environment: - QDRANT_URL=${QDRANT_URL} - COLLECTION_NAME=${COLLECTION_NAME} + - RERANK_EVENTS_DIR=/tmp/rerank_events + - RERANKER_WEIGHTS_DIR=/tmp/rerank_weights + - RELEVANCE_TRAINER_POLL_INTERVAL=${RELEVANCE_TRAINER_POLL_INTERVAL:-30} + - RELEVANCE_TRAINER_MIN_EVENTS=${RELEVANCE_TRAINER_MIN_EVENTS:-10} - HF_HOME=/tmp/huggingface - HF_HUB_CACHE=/tmp/huggingface/hub - TRANSFORMERS_CACHE=/tmp/huggingface/transformers - - FASTEMBED_CACHE_PATH=/tmp/huggingface/fastembed - - EMBEDDING_MODEL=${EMBEDDING_MODEL} - - EMBEDDING_PROVIDER=${EMBEDDING_PROVIDER} - # Cross-encoder reranker (used as teacher for learning) - - RERANKER_MODEL=${RERANKER_MODEL:-} - - RERANKER_ONNX_PATH=${RERANKER_ONNX_PATH:-} - - RERANKER_TOKENIZER_PATH=${RERANKER_TOKENIZER_PATH:-} - - RERANK_EVENTS_DIR=/tmp/rerank_events - - RERANKER_WEIGHTS_DIR=/tmp/rerank_weights - - RERANK_LEARNING_BATCH_SIZE=${RERANK_LEARNING_BATCH_SIZE:-32} - - RERANK_LEARNING_POLL_INTERVAL=${RERANK_LEARNING_POLL_INTERVAL:-30} - - RERANK_LEARNING_RATE=${RERANK_LEARNING_RATE:-0.001} - - RERANK_LLM_TEACHER=${RERANK_LLM_TEACHER:-0} - - RERANK_LLM_SAMPLE_RATE=${RERANK_LLM_SAMPLE_RATE:-1.0} - - GLM_API_KEY=${GLM_API_KEY:-} - - GLM_API_BASE=${GLM_API_BASE:-} - - GLM_MODEL=${GLM_MODEL:-} - - GLM_MODEL_FAST=${GLM_MODEL_FAST:-} - - REFRAG_RUNTIME=${REFRAG_RUNTIME:-glm} - - LLAMACPP_URL=${LLAMACPP_URL:-} - - RERANK_VICREG_WEIGHT=${RERANK_VICREG_WEIGHT:-0.1} volumes: - workspace_pvc:/work:rw - rerank_data:/tmp/rerank_weights:rw @@ -267,11 +244,9 @@ services: - RERANKER_MODEL=${RERANKER_MODEL:-} - RERANKER_ONNX_PATH=${RERANKER_ONNX_PATH:-} - RERANKER_TOKENIZER_PATH=${RERANKER_TOKENIZER_PATH:-} - # Learning reranker configuration - - RERANK_LEARNING=${RERANK_LEARNING:-1} + # Relevance feedback configuration - RERANKER_WEIGHTS_DIR=/tmp/rerank_weights - RERANK_EVENTS_DIR=/tmp/rerank_events - - RERANK_EVENTS_ENABLED=${RERANK_EVENTS_ENABLED:-1} ports: - "${FASTMCP_HTTP_HEALTH_PORT:-18002}:18000" - "${FASTMCP_HTTP_PORT:-8002}:8000" @@ -289,11 +264,11 @@ services: dockerfile: Dockerfile.mcp-indexer container_name: mcp-indexer-http-dev-remote user: "1000:1000" - # In K8s, scripts would be accessed directly at /app/scripts/ or via proper initContainer + # In K8s, scripts run as package modules from /app or via proper initContainer # For Docker Compose dev-remote simulation, create symlink so /work/scripts/ works # Use /tmp/huggingface for cache to avoid permission issues (universally writable) # Set CORRECT environment variables for HuggingFace and FastEmbed - command: ["sh", "-c", "mkdir -p /tmp/huggingface/hub /tmp/huggingface/transformers /tmp/huggingface/fastembed && exec python /app/scripts/mcp_indexer_server.py"] + command: ["sh", "-c", "mkdir -p /tmp/huggingface/hub /tmp/huggingface/transformers /tmp/huggingface/fastembed && cd /app && exec python -m scripts.mcp_indexer_server"] depends_on: - qdrant env_file: @@ -347,17 +322,13 @@ services: - LEX_BIGRAM_WEIGHT=${LEX_BIGRAM_WEIGHT:-} - LEX_SPARSE_MODE=${LEX_SPARSE_MODE:-} - LEX_SPARSE_NAME=${LEX_SPARSE_NAME:-} - # Pattern vectors for structural code similarity - - PATTERN_VECTORS=${PATTERN_VECTORS:-} # Cross-encoder reranker configuration - RERANKER_MODEL=${RERANKER_MODEL:-} - RERANKER_ONNX_PATH=${RERANKER_ONNX_PATH:-} - RERANKER_TOKENIZER_PATH=${RERANKER_TOKENIZER_PATH:-} - # Learning reranker configuration - - RERANK_LEARNING=${RERANK_LEARNING:-1} + # Relevance feedback configuration - RERANKER_WEIGHTS_DIR=/tmp/rerank_weights - RERANK_EVENTS_DIR=/tmp/rerank_events - - RERANK_EVENTS_ENABLED=${RERANK_EVENTS_ENABLED:-1} ports: - "${FASTMCP_INDEXER_HTTP_PORT:-8003}:8001" - "${FASTMCP_INDEXER_HTTP_HEALTH_PORT:-18003}:18001" @@ -431,12 +402,12 @@ services: - LEX_BIGRAM_WEIGHT=${LEX_BIGRAM_WEIGHT:-} - LEX_SPARSE_MODE=${LEX_SPARSE_MODE:-} - LEX_SPARSE_NAME=${LEX_SPARSE_NAME:-} - # Pattern vectors for structural code similarity - - PATTERN_VECTORS=${PATTERN_VECTORS:-} + - RERANKER_WEIGHTS_DIR=/tmp/rerank_weights volumes: - workspace_pvc:/work:rw - codebase_pvc:/work/.codebase:rw - entrypoint: ["sh", "-c", "mkdir -p /tmp/logs /tmp/huggingface/hub /tmp/huggingface/transformers /tmp/huggingface/fastembed && /app/scripts/wait-for-qdrant.sh && cd /app && python /app/scripts/ingest_code.py --root /work"] + - rerank_data:/tmp/rerank_weights:rw + entrypoint: ["sh", "-c", "mkdir -p /tmp/logs /tmp/huggingface/hub /tmp/huggingface/transformers /tmp/huggingface/fastembed && /app/scripts/wait-for-qdrant.sh && cd /app && python -m scripts.ingest_code --root /work"] restart: "no" # Run once on startup, do not restart after completion cpus: 2.0 networks: @@ -469,8 +440,12 @@ services: - QWEN3_QUERY_INSTRUCTION=${QWEN3_QUERY_INSTRUCTION:-1} - QWEN3_INSTRUCTION_TEXT=${QWEN3_INSTRUCTION_TEXT} - WATCH_ROOT=${WATCH_ROOT:-/work} + - CTXCE_METADATA_ROOT=${CTXCE_METADATA_ROOT:-/work} - HOST_INDEX_PATH=/work - QDRANT_TIMEOUT=${QDRANT_TIMEOUT:-60} + - WATCH_INIT_MAINTENANCE_ENABLED=${WATCH_INIT_MAINTENANCE_ENABLED:-1} + - WATCH_INIT_MAINTENANCE_INTERVAL_MINUTES=${WATCH_INIT_MAINTENANCE_INTERVAL_MINUTES:-120} + - WATCH_INIT_MAINTENANCE_RUN_ON_START=${WATCH_INIT_MAINTENANCE_RUN_ON_START:-0} # Chunking config - use ${VAR:-} to properly inherit from .env (not host shell) - INDEX_SEMANTIC_CHUNKS=${INDEX_SEMANTIC_CHUNKS:-} - INDEX_MICRO_CHUNKS=${INDEX_MICRO_CHUNKS:-} @@ -488,12 +463,12 @@ services: - LEX_BIGRAM_WEIGHT=${LEX_BIGRAM_WEIGHT:-} - LEX_SPARSE_MODE=${LEX_SPARSE_MODE:-} - LEX_SPARSE_NAME=${LEX_SPARSE_NAME:-} - # Pattern vectors for structural code similarity - - PATTERN_VECTORS=${PATTERN_VECTORS:-} + - RERANKER_WEIGHTS_DIR=/tmp/rerank_weights volumes: - workspace_pvc:/work:rw - codebase_pvc:/work/.codebase:rw - command: ["sh", "-c", "mkdir -p /tmp/huggingface/hub /tmp/huggingface/transformers /tmp/huggingface/fastembed && exec python /app/scripts/watch_index.py"] + - rerank_data:/tmp/rerank_weights:rw + command: ["sh", "-c", "mkdir -p /tmp/huggingface/hub /tmp/huggingface/transformers /tmp/huggingface/fastembed && cd /app && exec python -m scripts.watch_index"] cpus: 2 networks: - dev-remote-network @@ -532,7 +507,7 @@ services: command: [ "sh", "-c", - "mkdir -p /tmp/logs /work/.codebase && (chgrp -R 1000 /work/.codebase 2>/dev/null || true) && (chmod -R g+rwX /work/.codebase 2>/dev/null || true) && (find /work/.codebase -type d -exec chmod g+s {} + 2>/dev/null || true) && echo 'Starting initialization sequence...' && /app/scripts/wait-for-qdrant.sh && PYTHONPATH=/app python /app/scripts/create_indexes.py && echo 'Collections and metadata created' && python /app/scripts/warm_all_collections.py && echo 'Search caches warmed for all collections' && python /app/scripts/health_check.py && echo 'Initialization completed successfully!'" + "mkdir -p /tmp/logs /work/.codebase && (chgrp -R 1000 /work/.codebase 2>/dev/null || true) && (chmod -R g+rwX /work/.codebase 2>/dev/null || true) && (find /work/.codebase -type d -exec chmod g+s {} + 2>/dev/null || true) && echo 'Starting initialization sequence...' && cd /app && python -m scripts.run_init_maintenance && echo 'Initialization completed successfully!'" ] restart: "no" # Run once on startup networks: @@ -555,6 +530,7 @@ services: - UPLOAD_SERVICE_PORT=8002 - QDRANT_URL=${QDRANT_URL} - WORKDIR=/work + - CTXCE_METADATA_ROOT=${CTXCE_METADATA_ROOT:-/work} - MAX_BUNDLE_SIZE_MB=100 - UPLOAD_TIMEOUT_SECS=300 # Optional auth configuration (fully opt-in via .env) @@ -603,17 +579,19 @@ services: - LEX_BIGRAM_WEIGHT=${LEX_BIGRAM_WEIGHT:-} - LEX_SPARSE_MODE=${LEX_SPARSE_MODE:-} - LEX_SPARSE_NAME=${LEX_SPARSE_NAME:-} + - RERANKER_WEIGHTS_DIR=/tmp/rerank_weights ports: - "8004:8002" # Map to different host port to avoid conflicts - "18004:18000" # Health check port volumes: - workspace_pvc:/work:rw - codebase_pvc:/work/.codebase:rw + - rerank_data:/tmp/rerank_weights:rw - upload_temp:/tmp/uploads command: [ "sh", "-c", - "mkdir -p /work/.codebase && (chgrp -R 1000 /work/.codebase 2>/dev/null || true) && (chmod -R g+rwX /work/.codebase 2>/dev/null || true) && (find /work/.codebase -type d -exec chmod g+s {} + 2>/dev/null || true) && exec python scripts/upload_service.py" + "mkdir -p /work/.codebase && (chgrp -R 1000 /work/.codebase 2>/dev/null || true) && (chmod -R g+rwX /work/.codebase 2>/dev/null || true) && (find /work/.codebase -type d -exec chmod g+s {} + 2>/dev/null || true) && cd /app && exec python -m scripts.upload_service" ] healthcheck: test: ["CMD", "curl", "-f", "http://localhost:8002/health"] @@ -660,11 +638,11 @@ volumes: qdrant_storage_dev_remote: driver: local - # Learning reranker weights storage (shared between indexer and worker) + # Relevance feedback weights storage (shared between search/indexer/trainer) rerank_data: driver: local - # Learning reranker events storage (shared between indexer and worker) + # Relevance feedback events storage (shared between indexer/trainer) rerank_events: driver: local @@ -674,4 +652,4 @@ networks: driver: bridge ipam: config: - - subnet: 172.20.0.0/16 \ No newline at end of file + - subnet: 172.20.0.0/16 diff --git a/docs/ARCHITECTURE.md b/docs/ARCHITECTURE.md index 00798685..828d4e60 100644 --- a/docs/ARCHITECTURE.md +++ b/docs/ARCHITECTURE.md @@ -8,7 +8,7 @@ - [Overview](#overview) - [Core Principles](#core-principles) - [System Architecture](#system-architecture) -- [Learning Reranker System](#5-learning-reranker-system) +- [Relevance Feedback System](#5-relevance-feedback-optional) - [Data Flow](#data-flow) - [ReFRAG Pipeline](#refrag-pipeline) @@ -123,121 +123,66 @@ Production-ready MCP (Model Context Protocol) retrieval stack unifying code inde - **Local LLM Integration**: llama.cpp for offline expansion - **Caching**: Expanded query results cached for reuse -#### Pattern Detection (`scripts/pattern_detection/`) -- **Structural Search**: Find similar code patterns across languages via AST analysis -- **64-dim Pattern Vector**: WL graph kernel, CFG fingerprint, SimHash, spectral features -- **Auto-Detection**: Identifies retry patterns, resource cleanup, filter loops -- **Requires**: `PATTERN_VECTORS=1` to enable - -### 5. Learning Reranker System (Optional) - -The Learning Reranker is an **optional** self-improving ranking system that learns from search patterns to provide increasingly relevant results over time. It is enabled by default but can be disabled via `RERANK_LEARNING=0` and `RERANK_EVENTS_ENABLED=0` environment variables. See [Configuration](CONFIGURATION.md#learning-reranker) for all options. - -#### Architecture Overview - -``` -┌─────────────────┐ ┌──────────────────┐ ┌─────────────────┐ -│ Search Query │────►│ Hybrid Search │────►│ TinyScorer │ -│ │ │ (initial rank) │ │ (learned rank) │ -└─────────────────┘ └──────────────────┘ └─────────────────┘ - │ - ┌──────────────────┐ │ - │ Event Logger │◄────────────┘ - │ (NDJSON files) │ - └────────┬─────────┘ - │ - ┌────────▼─────────┐ - │ Learning Worker │ - │ (background) │ - └────────┬─────────┘ - │ - ┌────────▼─────────┐ - │ ONNX Teacher │ - │ (cross-encoder) │ - └────────┬─────────┘ - │ - ┌────────▼─────────┐ - │ Weight Updates │ - │ (.npz files) │ - └──────────────────┘ -``` - -#### Components - -**TinyScorer** (`scripts/rerank_recursive.py`) -- 2-layer MLP neural network (~3MB per collection) -- Scores query-document pairs based on learned patterns -- Hot-reloads weights every 60 seconds from disk -- Per-collection weights (each repo learns independently) - -**Event Logger** (`scripts/rerank_events.py`) -- Logs every search to NDJSON files at `/tmp/rerank_events/` -- Records: query, candidates, initial scores, timestamps -- Hourly file rotation with configurable retention - -**Learning Worker** (`scripts/learning_reranker_worker.py`) -- Background daemon that processes logged events -- Uses ONNX cross-encoder as "teacher" model -- Trains TinyScorer via knowledge distillation -- Saves versioned weight checkpoints atomically - -#### Learning Flow - -1. **Event Capture**: Every search logs query + candidates to NDJSON -2. **Teacher Scoring**: ONNX cross-encoder scores the candidates -3. **Student Training**: TinyScorer learns to match teacher rankings -4. **Weight Update**: New weights saved atomically with versioning -5. **Hot Reload**: Serving path picks up new weights within 60s -6. **Score Integration**: `learning_score` blends with other signals - -#### Configuration - -| Variable | Description | Default | -|----------|-------------|---------| -| `RERANKER_WEIGHTS_DIR` | Directory for weight files | `/tmp/rerank_weights` | -| `RERANKER_WEIGHTS_RELOAD_INTERVAL` | Hot-reload check interval (seconds) | 60 | -| `RERANKER_MAX_CHECKPOINTS` | Number of weight versions to keep | 5 | -| `RERANKER_LR_DECAY_STEPS` | Steps between learning rate decay | 1000 | -| `RERANKER_LR_DECAY_RATE` | Learning rate decay multiplier | 0.95 | -| `RERANKER_MIN_LR` | Minimum learning rate | 0.0001 | -| `RERANK_EVENTS_DIR` | Directory for event logs | `/tmp/rerank_events` | -| `RERANK_EVENTS_RETENTION_DAYS` | Days to keep event files | 7 | -| `RERANK_LEARNING_BATCH_SIZE` | Events per training batch | 32 | -| `RERANK_LEARNING_POLL_INTERVAL` | Worker poll interval (seconds) | 30 | -| `RERANK_LEARNING_RATE` | Initial learning rate | 0.001 | - -#### Observability - -Search results include learning metrics in the `why` field: -```json -{ - "score": 3.2, - "why": ["lexical:1.0", "dense_rrf:0.05", "learning:3", "score:3.2"], - "components": { - "learning_score": 3.2, - "learning_iterations": 3 - } -} -``` - -Worker logs show training progress: -``` -[codebase] Processed 5 events | v12 | lr=0.001 | avg_loss=1.8 | converged=False -``` - -#### Benefits - -- **Zero Manual Training**: Learns automatically from usage -- **Per-Collection Specialization**: Each codebase gets tuned rankings -- **Fast Inference**: TinyScorer adds <1ms to search latency -- **Continuous Improvement**: Rankings improve over time -- **Offline Capable**: Teacher runs locally, no external API calls - -#### MCP Router (`scripts/mcp_router.py`) -- **Intent Classification**: Determines which MCP tool to call based on query -- **Tool Orchestration**: Routes to search, answer, memory, or index tools -- **HTTP Execution**: Executes tools via RMCP/HTTP without extra dependencies -- **Plan Mode**: Preview tool selection without execution +#### Pseudo Descriptions +- **Index-time vocabulary bridge**: Optional LLM-generated pseudo descriptions + and tags are stored with chunks and can be included in dense indexing text via + `INDEX_DENSE_MODE=info+pseudo+tags`. +- **Lexical participation**: Pseudo/tags are appended to lexical text during + indexing and can contribute to lexical scoring when `HYBRID_PSEUDO_BOOST` is + enabled. +- **Not reranker-only**: These fields are carried through search results and are + already part of retrieval when indexing/search env knobs enable them. + +### 5. Relevance Feedback (Optional) + +Per-collection feedback via `rate_search_results` MCP tool. Agents or users +rate search results (0=not used, 1=glanced, 2=directly used) and a background +trainer aggregates ratings into per-collection weight files. + +Feedback uses stable target identity rather than exact line spans: +- `result_id` / `target_id`: stable symbol-or-file target used for recall and boosts +- `impression_id`: query/content/span-specific diagnostic ID + +Identity and reindex behavior: +- Line shifts and body edits keep feedback when repo, kind, and symbol identity stay + the same. +- File moves keep symbol feedback because the path is not part of symbol target + identity. +- Exact-body symbol renames are reconciled during smart reindex using the + symbol-level content hash. +- A removed symbol may transfer feedback to multiple extracted successors only + when their token overlap collectively covers the old body above conservative + thresholds. The inherited weight is divided between successors. +- Ambiguous renames/splits are not migrated. The old feedback remains recorded + but does not attach itself to an uncertain target. + +Recent result metadata used by hands-off ratings is persisted per collection in +the shared feedback volume. Search and rating calls therefore do not need to hit +the same server process. + +Subsequent searches can rehydrate positively-rated targets that ordinary +retrieval missed, then apply a soft `relevance_boost`. Feedback adds candidates +and nudges rank; it does not force rated targets to win. + +For adjacent-code discovery, feedback recall uses inverse graph lookups against +the materialized Qdrant graph collection: a rated callee symbol can produce +candidate caller paths via `callee_symbol -> caller_path`, which are then +hydrated from the main code collection. + +Components: +- **Event Logger** (`scripts/rerank_tools/events.py`): NDJSON-based event files +- **Relevance Trainer** (`scripts/relevance_trainer.py`): Aggregates ratings, + preserves rehydratable target metadata, writes per-collection weight files + atomically +- **Feedback Recall + Boost** (`scripts/mcp_impl/search.py`): Reads weight files + at search time, rehydrates a small number of rated targets and inverse-graph + caller candidates from Qdrant, then applies `relevance_boost` to target IDs + +Key env vars: `RELEVANCE_BOOST_FACTOR`, `RELEVANCE_RECALL_MAX`, +`RELEVANCE_GRAPH_RECALL_MAX`, `RERANKER_WEIGHTS_DIR`, `RERANK_EVENTS_DIR` + +This replaces the former self-supervised ranking experiment with explicit +human/agent ratings and bounded feedback recall. ## Data Flow Architecture @@ -344,4 +289,4 @@ Worker logs show training progress: - **Feature Flags**: Experimental feature toggling - **A/B Testing**: Multiple configuration variants support -This architecture enables Context Engine to serve as a production-ready, scalable context layer for AI applications while maintaining the flexibility to evolve with changing requirements and technologies. \ No newline at end of file +This architecture enables Context Engine to serve as a production-ready, scalable context layer for AI applications while maintaining the flexibility to evolve with changing requirements and technologies. diff --git a/docs/BENCHMARKS.md b/docs/BENCHMARKS.md index 152348ea..71161c5a 100644 --- a/docs/BENCHMARKS.md +++ b/docs/BENCHMARKS.md @@ -83,12 +83,11 @@ python -m scripts.benchmarks.cosqa.runner --limit 500 --enable-llm | `--debug` | Print detailed per-query debug output | | `--output FILE` | Output JSON report path | -#### CoSQA ablation runs (refrag/mini, rerank, learning) +#### CoSQA ablation runs (refrag/mini and reranking) Use the helper script to run a consistent matrix across: - rerank vs no rerank - ReFRAG/mini vectors vs no ReFRAG -- learning vs no learning ```bash # Default: 50/50 subset per run, new collection per variant @@ -107,7 +106,6 @@ Outputs: Useful knobs (env): - `QDRANT_URL`, `LEX_VECTOR_DIM`, `HYBRID_EXPAND`, `SEMANTIC_EXPANSION_ENABLED` - `RUN_TAG`, `COLL_PREFIX`, `OUT_DIR`, `LOG_DIR` -- `COSQA_ENABLE_LEARNING` - Enable learning reranker (default: off for determinism) - `EMBEDDING_SEED` - Seed for deterministic embeddings #### Quick subset runs (smoke tests) @@ -131,7 +129,6 @@ QDRANT_URL=http://localhost:6333 python3.11 -m scripts.benchmarks.cosqa.runner \ CoSQA benchmarks are hardened for reproducibility: -- **Learning reranker disabled** by default (`RERANK_LEARNING=0`) to avoid score drift - **Content-hash deduplication** prevents duplicate corpus entries across runs - **Schema validation** fails fast on LEX_DIM mismatches (avoids silent scoring bugs) - **Warmup query** runs before timing loop to exclude cold-start latency @@ -360,7 +357,7 @@ rerank_return_m=10 # Rerank and return top 10 The reranker (ONNX cross-encoder) scores `(query, document)` pairs. The document text is constructed from: ```python -# From scripts/rerank_local.py:prepare_pairs() +# From scripts/rerank_tools/local.py:prepare_pairs() header = f"[{language}/{kind}] {symbol_path} — {path}" doc = header + "\n" + metadata.code[:600] ``` @@ -379,7 +376,7 @@ doc = header + "\n" + metadata.code[:600] To debug what the reranker sees: ```python -from scripts.rerank_local import prepare_pairs +from scripts.rerank_tools.local import prepare_pairs from qdrant_client import QdrantClient client = QdrantClient() diff --git a/docs/CLAUDE.example.md b/docs/CLAUDE.example.md index c46fdf4e..37760861 100644 --- a/docs/CLAUDE.example.md +++ b/docs/CLAUDE.example.md @@ -81,7 +81,7 @@ These rules are NOT optional - favor qdrant-indexer tooling at all costs over ex Tool Roles Cheat Sheet: - - repo_search / code_search: + - repo_search: - Use for: finding relevant files/spans and inspecting raw code. - Think: "where is X implemented?", "show me usages of Y". - context_search: @@ -90,13 +90,6 @@ These rules are NOT optional - favor qdrant-indexer tooling at all costs over ex - context_answer: - Use for: short natural-language summaries/explanations of specific modules or tools, grounded in code/docs with citations. - Good for: "What does scripts/standalone_upload_client.py do at a high level?", "Summarize the remote upload client pipeline.". - - pattern_search (optional, may not be enabled): - - Use for: finding structurally similar code patterns across files and languages. - - Accepts EITHER code examples OR natural language pattern descriptions. - - Good for: "find retry loops with exponential backoff", "try: ... except: logger.error()", "error handling patterns". - - Cross-language: Python pattern can match Go/Rust/Java with similar control flow. - - Note: Returns error if pattern detection module is not available. - Advanced lineage workflow (code + history): - Goal: answer "when/why did behavior X change?" without flooding context. @@ -133,10 +126,8 @@ These rules are NOT optional - favor qdrant-indexer tooling at all costs over ex - workspace_info, list_workspaces, collection_map - set_session_defaults - Search / QA tools: - - repo_search, code_search, context_search, context_answer - - pattern_search (optional; structural code pattern matching, cross-language) - - search_tests_for, search_config_for, search_callers_for, search_importers_for - - change_history_for_path, expand_query + - repo_search, context_search, context_answer + - symbol_graph, change_history_for_path, expand_query - Memory tools: - memory.set_session_defaults, memory.memory_store, memory.memory_find @@ -148,4 +139,4 @@ These rules are NOT optional - favor qdrant-indexer tooling at all costs over ex blended code + memory results instead of calling repo_search and memory.memory_find separately. - Treat expand_query and the expand flag on context_answer as expensive options: - only use them after a normal search/answer attempt failed to find good context. \ No newline at end of file + only use them after a normal search/answer attempt failed to find good context. diff --git a/docs/CONFIGURATION.md b/docs/CONFIGURATION.md index c5a86453..f5a17532 100644 --- a/docs/CONFIGURATION.md +++ b/docs/CONFIGURATION.md @@ -13,15 +13,13 @@ Complete environment variable reference for Context Engine. - [Query Optimization](#query-optimization) - [Watcher Settings](#watcher-settings) - [Reranker](#reranker) -- [Learning Reranker](#learning-reranker) +- [Relevance Feedback](#relevance-feedback) - [Decoder (llama.cpp / OpenAI / GLM / MiniMax)](#decoder-llamacpp--openai--glm--minimax) - [Git History & Commit Indexing](#git-history--commit-indexing) - [ReFRAG](#refrag) -- [Pattern Search](#pattern-search) - [Lexical Vector Settings](#lexical-vector-settings) - [Ports](#ports) - [Search & Expansion](#search--expansion) -- [info_request Tool](#info_request-tool) - [Memory Blending](#memory-blending) --- @@ -135,6 +133,12 @@ Dynamic HNSW_EF tuning and intelligent query routing for 2x faster simple querie | Name | Description | Default | |------|-------------|---------| | WATCH_DEBOUNCE_SECS | Debounce between FS events | 1.5 | +| WATCH_JOURNAL_DRAIN_BATCH_SIZE | Maximum journal paths queued per drain pass | 256 | +| WATCH_JOURNAL_LOG_INTERVAL_SECS | Minimum interval between journal drain log entries | 120 | +| WATCH_INIT_MAINTENANCE_ENABLED | Run periodic init maintenance from watcher | 1 (enabled) | +| WATCH_INIT_MAINTENANCE_INTERVAL_MINUTES | Minutes between init maintenance passes | 120 | +| WATCH_INIT_MAINTENANCE_RUN_ON_START | Run immediately on watcher startup instead of waiting one interval | 0 (disabled) | +| WATCH_INIT_MAINTENANCE_COMMAND_TIMEOUT_SECS | Per-command timeout for init maintenance scripts | 1800 | | INDEX_UPSERT_BATCH | Upsert batch size (watcher) | 128 | | INDEX_UPSERT_RETRIES | Retry count | 5 | | INDEX_UPSERT_BACKOFF | Seconds between retries | 0.5 | @@ -190,63 +194,31 @@ For custom models or explicit control, set both ONNX path and tokenizer: | EMBEDDING_WARMUP | Warm up embedding model on startup | 0 (disabled) | | RERANK_WARMUP | Warm up reranker model on startup | 0 (disabled) | -## Learning Reranker +## Relevance Feedback -The learning reranker trains a lightweight neural network (TinyScorer) to improve search rankings over time. See [Architecture](ARCHITECTURE.md#5-learning-reranker-system) for details. - -**This feature is optional and enabled by default.** To disable: - -```bash -# Disable learning scorer in search results -RERANK_LEARNING=0 - -# Disable event logging (no training data collected) -RERANK_EVENTS_ENABLED=0 - -# Or simply don't run the learning_worker container -``` - -### Enable/Disable - -| Name | Description | Default | -|------|-------------|---------| -| RERANK_LEARNING | Enable learning scorer in search results | 1 (enabled) | -| RERANK_EVENTS_ENABLED | Enable event logging for training | 1 (enabled) | -| RERANK_EVENTS_SAMPLE_RATE | Fraction of events to log (0.0-1.0) | 0.33 | +The relevance feedback path records explicit agent/user ratings from +`rate_search_results` and turns them into per-collection recall/boost metadata. +It is not self-supervised learning and does not train on reranker scores. ### Weight Management | Name | Description | Default | |------|-------------|---------| -| RERANKER_WEIGHTS_DIR | Directory for learned weight files | /tmp/rerank_weights | -| RERANKER_WEIGHTS_RELOAD_INTERVAL | How often to check for new weights (seconds) | 60 | -| RERANKER_MAX_CHECKPOINTS | Number of weight versions to retain | 5 | - -### Learning Rate - -| Name | Description | Default | -|------|-------------|---------| -| RERANKER_LR_DECAY_STEPS | Updates between learning rate decay | 1000 | -| RERANKER_LR_DECAY_RATE | Decay multiplier (e.g., 0.95 = 5% reduction) | 0.95 | -| RERANKER_MIN_LR | Minimum learning rate floor | 0.0001 | +| RELEVANCE_BOOST_FACTOR | Max score boost for positively rated targets | 0.15 | +| RELEVANCE_RECALL_MAX | Max positively rated targets to rehydrate per search | 3 | +| RELEVANCE_GRAPH_RECALL_MAX | Max inverse-graph caller candidates per rated symbol | 3 | +| RELEVANCE_GRAPH_RECALL_BOOST | Small seed boost for graph-recalled candidates | 0.01 | +| RELEVANCE_SPLIT_MIN_OVERLAP | Minimum old-symbol token overlap for each split successor | 0.45 | +| RELEVANCE_SPLIT_MIN_COVERAGE | Minimum combined old-symbol coverage before split inheritance | 0.75 | +| RERANKER_WEIGHTS_DIR | Directory for feedback weight files | /tmp/rerank_weights | +| RELEVANCE_TRAINER_MIN_EVENTS | Min feedback events before writing weights | 10 | +| RELEVANCE_TRAINER_POLL_INTERVAL | Trainer daemon poll interval in seconds | 30 | ### Event Logging | Name | Description | Default | |------|-------------|---------| -| RERANK_EVENTS_DIR | Directory for search event logs | /tmp/rerank_events | -| RERANK_EVENTS_RETENTION_DAYS | Days to keep event files before cleanup | 7 | - -### Learning Worker - -| Name | Description | Default | -|------|-------------|---------| -| RERANK_LEARNING_BATCH_SIZE | Number of events per training batch | 32 | -| RERANK_LEARNING_POLL_INTERVAL | Seconds between checking for new events | 30 | -| RERANK_LEARNING_RATE | Initial learning rate for TinyScorer | 0.001 | -| RERANK_LLM_TEACHER | Enable LLM-teacher guided learning | 1 (enabled) | -| RERANK_LLM_SAMPLE_RATE | Fraction of queries to evaluate with LLM teacher | 1.0 | -| RERANK_VICREG_WEIGHT | Weight for VICReg consistency loss | 0.1 | +| RERANK_EVENTS_DIR | Directory for feedback event logs | /tmp/rerank_events | ## Decoder (llama.cpp / OpenAI / GLM / MiniMax) @@ -332,7 +304,11 @@ Deferred pseudo/tag generation runs asynchronously after initial indexing. | Name | Description | Default | |------|-------------|---------| | PSEUDO_BACKFILL_ENABLED | Enable async pseudo/tag backfill worker | 0 (disabled) | -| PSEUDO_DEFER_TO_WORKER | Skip inline pseudo, defer to backfill worker | 0 (disabled) | +| PSEUDO_DEFER_TO_WORKER | Foreground/background semantics: only disables inline pseudo when backfill worker is enabled | 0 (disabled) | + +Notes: +- `PSEUDO_BACKFILL_ENABLED=0` is a hard disable for the worker. +- `PSEUDO_DEFER_TO_WORKER=1` has no effect unless `PSEUDO_BACKFILL_ENABLED=1` (we keep inline pseudo enabled to avoid silently dropping pseudo/tags). ### Adaptive Span Sizing @@ -364,24 +340,6 @@ Compact 64-dim vectors for fast candidate filtering before full dense search. | MINI_VEC_SEED | Random projection seed (for reproducibility) | 1337 | | HYBRID_MINI_WEIGHT | Weight of mini vectors in hybrid scoring | 0.5 | -## Pattern Search - -Structural code pattern matching across languages. Disabled by default. - -| Name | Description | Default | -|------|-------------|---------| -| PATTERN_VECTORS | Enable pattern_search tool and pattern vector indexing | 0 (disabled) | - -**Enable:** -```bash -# In .env or docker-compose -PATTERN_VECTORS=1 -``` - -When enabled, the indexer extracts control-flow signatures (loops, branches, try/except, etc.) and stores them as pattern vectors. The `pattern_search` MCP tool allows finding structurally similar code across languages—e.g., a Python retry loop can match Go/Rust equivalents. - -**Note:** Enabling requires reindexing to generate pattern vectors for existing files. - ## Lexical Vector Settings Controls the sparse lexical (keyword) vectors used for hybrid search. @@ -506,20 +464,13 @@ The search engine can boost files whose paths match query terms—production-gra Set `FNAME_BOOST=0` to disable, or increase (e.g., `0.25`) for stronger path weighting. -## info_request Tool - -Simplified codebase retrieval with optional explanation mode. - -| Name | Description | Default | -|------|-------------|---------| -| INFO_REQUEST_LIMIT | Default result limit for info_request queries | 10 | -| INFO_REQUEST_CONTEXT_LINES | Context lines in snippets (richer than repo_search) | 5 | - ## Output Formatting ### TOON (Token-Oriented Object Notation) -Compact output format that reduces token usage by 40-60%. +Compact display format for search results. In practice this is usually about +20-25% smaller than compact JSON for search-shaped payloads, with larger savings +only when comparing against pretty-printed JSON. | Name | Description | Default | |------|-------------|---------| @@ -577,4 +528,3 @@ docker compose run --rm indexer --root /work --no-default-excludes --exclude '/v | Large (1k+ files) | 120 (default) | 20 | 128+ | For large monorepos, set `INDEX_PROGRESS_EVERY=200` for visibility. - diff --git a/docs/DEVELOPMENT.md b/docs/DEVELOPMENT.md index e2f55b92..aa28f4db 100644 --- a/docs/DEVELOPMENT.md +++ b/docs/DEVELOPMENT.md @@ -97,7 +97,6 @@ Context-Engine/ ├── scripts/ # Core application code │ ├── mcp_memory_server.py # Memory MCP server implementation │ ├── mcp_indexer_server.py # Indexer MCP server implementation -│ ├── mcp_router.py # Intent-based tool routing │ ├── hybrid_search.py # Search algorithm implementation │ ├── ctx.py # CLI prompt enhancer │ ├── cache_manager.py # Unified caching system @@ -403,7 +402,7 @@ class TestSearchIntegration: @pytest.fixture(scope="module") def qdrant_container(self): """Set up real Qdrant container for integration tests.""" - container = DockerContainer("qdrant/qdrant:latest").with_exposed_ports(6333) + container = DockerContainer("qdrant/qdrant:v1.15.4").with_exposed_ports(6333) container.start() yield f"http://{container.get_container_host_ip()}:{container.get_exposed_port(6333)}" container.stop() @@ -574,4 +573,4 @@ curl http://localhost:18001/tools - [ ] Error handling is appropriate - [ ] Performance impact is considered -This development guide should help you get started with contributing to Context Engine. For more specific questions, refer to the code documentation or create an issue in the repository. \ No newline at end of file +This development guide should help you get started with contributing to Context Engine. For more specific questions, refer to the code documentation or create an issue in the repository. diff --git a/docs/IDE_CLIENTS.md b/docs/IDE_CLIENTS.md index 88249340..e86395d0 100644 --- a/docs/IDE_CLIENTS.md +++ b/docs/IDE_CLIENTS.md @@ -229,7 +229,7 @@ url = "http://127.0.0.1:8003/mcp" "args": ["mcp-server-qdrant"], "env": { "QDRANT_URL": "http://localhost:6333", - "COLLECTION_NAME": "my-collection", + "COLLECTION_NAME": "codebase", "EMBEDDING_MODEL": "BAAI/bge-base-en-v1.5" }, "disabled": false @@ -282,7 +282,7 @@ scripts/remote_upload_client.py --server http://context.yourcompany.com:9090 --p - **Do not send null values** to MCP tools. Omit the field or pass an empty string "" instead. - **qdrant-index examples:** - - `{"subdir":"","recreate":false,"collection":"my-collection","repo_name":"workspace"}` + - `{"subdir":"","recreate":false,"collection":"codebase","repo_name":"workspace"}` - `{"subdir":"scripts","recreate":true}` - For indexing repo root with no params, use `qdrant_index_root` (zero-arg) or call `qdrant-index` with `subdir:""`. @@ -292,7 +292,7 @@ scripts/remote_upload_client.py --server http://context.yourcompany.com:9090 --p After configuring, you should see tools from both servers: - `store`, `find` (Memory) -- `repo_search`, `code_search`, `context_search`, `context_answer` (Indexer) +- `repo_search`, `context_search`, `context_answer` (Indexer) - `qdrant_list`, `qdrant_index`, `qdrant_prune`, `qdrant_status` (Indexer) Test connectivity: @@ -327,4 +327,3 @@ When using `@context-engine-bridge/context-engine-mcp-bridge`, ensure you set `C ``` The default collection name is `codebase` unless you've configured a different one during indexing. - diff --git a/docs/MCP_API.md b/docs/MCP_API.md index 59dbb8cf..0e788764 100644 --- a/docs/MCP_API.md +++ b/docs/MCP_API.md @@ -9,7 +9,7 @@ This document provides comprehensive API documentation for all MCP (Model Contex **On this page:** - [Overview](#overview) - [Memory Server API](#memory-server-api) - `memory_store()`, `memory_find()` -- [Indexer Server API](#indexer-server-api) - `repo_search()`, `context_search()`, `context_answer()`, `info_request()`, etc. +- [Indexer Server API](#indexer-server-api) - `repo_search()`, `context_search()`, `context_answer()`, etc. - [Response Schemas](#response-schemas) - [Error Handling](#error-handling) @@ -153,7 +153,7 @@ Search stored memories using hybrid retrieval (semantic + lexical search). ### repo_search() -Perform hybrid code search combining dense semantic, lexical BM25, and optional neural reranking. +Perform code search using the configured retrieval mode. Dense semantic search can be used on its own, or combined with lexical fusion and optional neural reranking when those features are enabled. **Core Parameters:** - `query` (str or list[str], required): Search query or list of queries for query fusion @@ -172,6 +172,10 @@ Perform hybrid code search combining dense semantic, lexical BM25, and optional - `path_glob` (str or list[str], optional): Glob patterns for path filtering - `under` (str, optional): Limit search to specific directory path - `not_glob` (str or list[str], optional): Exclude paths matching these patterns +- `profile` (str, optional): Apply a focused path profile before search: + - `"tests"`: Prefer common test file paths + - `"config"`: Prefer common configuration files + - `"code"`: Prefer source-code files **Code Structure Filters:** - `symbol` (str, optional): Search for specific function, class, or variable names @@ -386,132 +390,6 @@ All `repo_search` parameters supported for context retrieval. } ``` -### info_request() - -Simplified codebase retrieval with optional explanation mode. Drop-in replacement for basic codebase retrieval tools with human-readable result descriptions. - -**Primary Parameters:** -- `info_request` (str, required): Natural language description of the code you're looking for -- `information_request` (str): Alias for `info_request` - -**Explanation Mode:** -- `include_explanation` (bool, default false): Add summary, primary_locations, related_concepts, grouped_results, and confidence metrics -- `include_relationships` (bool, default false): Add imports_from, calls, related_paths to each result - -**Filter Parameters:** -- `limit` (int): Maximum results (smart defaults: 15 for short queries, 8 for questions, 10 otherwise) -- `language` (str, optional): Filter by programming language -- `under` (str, optional): Limit search to specific directory -- `repo` (str or list[str], optional): Filter by repository name(s) -- `path_glob` (str or list[str], optional): Glob patterns for file paths - -**Snippet Options:** -- `include_snippet` (bool, default true): Include code snippets -- `context_lines` (int, default 5): Lines of context around matches - -**Returns (basic mode):** -```json -{ - "ok": true, - "results": [ - { - "score": 0.85, - "path": "/work/src/hooks/useAuth.tsx", - "symbol": "useAuth", - "start_line": 15, - "end_line": 45, - "information": "Found 'useAuth' in useAuth.tsx (lines 15-45)", - "relevance_score": 0.85, - "snippet": "export function useAuth() { ... }" - } - ], - "total": 10, - "search_strategy": "hybrid+rerank" -} -``` - -**Returns (with `include_explanation: true`):** -```json -{ - "ok": true, - "results": [...], - "total": 10, - "search_strategy": "hybrid+rerank+lang:typescript", - "summary": "Found 10 results related to 'authentication hook' across 5 files", - "primary_locations": [ - "/work/src/hooks/useAuth.tsx", - "/work/src/context/AuthContext.tsx" - ], - "related_concepts": ["auth", "hook", "context", "session", "token"], - "grouped_results": { - "by_file": { - "/work/src/hooks/useAuth.tsx": { - "count": 3, - "top_symbols": ["useAuth", "AuthProvider", "useSession"] - } - } - }, - "confidence": { - "level": "high", - "score": 0.78, - "top_score": 0.85, - "symbol_matches": 2 - }, - "query_understanding": { - "intent": "search_for_code", - "detected_language": "typescript", - "detected_symbols": ["useAuth"], - "search_strategy": "hybrid+rerank+lang:typescript" - } -} -``` - -**Returns (with `include_relationships: true`):** -```json -{ - "results": [ - { - "information": "Found 'useAuth' in useAuth.tsx (lines 15-45)", - "relationships": { - "imports_from": ["react", "@/context/AuthContext"], - "calls": ["useState", "useContext", "fetchUser"], - "symbol_path": "useAuth", - "related_paths": ["/work/src/context/AuthContext.tsx"] - } - } - ] -} -``` - -**Smart Limits:** -- Short queries (1-2 words): 15 results for broader coverage -- Question queries ("how does", "what is"): 8 results for focused answers -- Default: 10 results - -**Search Strategy Labels:** -- `hybrid` - Base hybrid search (dense + lexical) -- `+rerank` - Neural reranker applied -- `+repo_filtered` - Filtered to specific repo(s) -- `+lang:python` - Filtered by language -- `+path_filtered` - Filtered by directory - -**Environment Variables:** -- `INFO_REQUEST_LIMIT=10` - Default result limit -- `INFO_REQUEST_CONTEXT_LINES=5` - Default context lines -- `INFO_REQUEST_EXPLAIN_DEFAULT=0` - Enable explanation mode by default -- `INFO_REQUEST_RELATIONSHIPS=0` - Enable relationships by default - -**Example:** -```json -{ - "info_request": "authentication middleware", - "include_explanation": true, - "include_relationships": true, - "language": "python", - "limit": 5 -} -``` - ### qdrant_index() Index or reindex code from the mounted workspace. @@ -700,10 +578,6 @@ Supports three runtime backends via `REFRAG_RUNTIME`: On decoder error, falls back to suffix-based expansion with `"decoder_used": "fallback"`. If expansion fails entirely, returns `"ok": false` with an error message. -### code_search() - -Exact alias of `repo_search()` for discoverability. Same parameters and return shape. - ### qdrant_index_root() Index the entire workspace root (`/work`). @@ -714,75 +588,6 @@ Index the entire workspace root (`/work`). **Returns:** Subprocess result with indexing status. -### search_tests_for() - -Find test files related to a query. Presets common test file globs. - -**Parameters:** -- `query` (str or list[str], required): Search query -- `limit` (int, optional): Max results -- `include_snippet` (bool, optional): Include code snippets -- `language` (str, optional): Filter by language - -**Returns:** Same shape as `repo_search()`. - -### search_config_for() - -Find configuration files related to a query. Presets config file globs (yaml/json/toml/etc). - -**Parameters:** Same as `search_tests_for()`. - -**Returns:** Same shape as `repo_search()`. - -### search_callers_for() - -Heuristic search for callers/usages of a symbol. - -**Parameters:** -- `query` (str, required): Symbol name to find callers for -- `limit` (int, optional): Max results -- `language` (str, optional): Filter by language - -**Returns:** Same shape as `repo_search()`. - -### search_importers_for() - -Find files likely importing or referencing a module/symbol. - -**Parameters:** Same as `search_callers_for()`. - -**Returns:** Same shape as `repo_search()`. - -### pattern_search() - -Find structurally similar code patterns across languages. Requires `PATTERN_VECTORS=1`. - -**Parameters:** -- `query` (str, required): Code snippet OR natural language pattern description -- `language` (str, default "python"): Language hint for code queries -- `limit` (int, default 10): Maximum results -- `min_score` (float, default 0.3): Similarity threshold -- `include_snippet` (bool): Include code in results -- `target_languages` (list[str]): Filter target languages - -**Response:** -```json -{ - "ok": true, - "results": [{"path": "...", "start_line": 45, "score": 0.94, "control_flow_signature": "L2_2_B0_T2_M0__C_TL"}], - "total": 5, - "query_signature": "L2_2_B0_T2_M0__C_TL", - "query_mode": "code" -} -``` - -**Signature format:** `L{loop_depth}_{count}_B{branches}_T{try}_M{match}_{flags}` where flags include `TL` (retry pattern), `BL` (filter pattern). - -**Example:** -```json -{"query": "for i in range(3): try: fetch() except: sleep(i)", "include_snippet": true} -``` - ### symbol_graph() First-class symbol graph navigation using indexed metadata fields: @@ -932,4 +737,4 @@ Both SSE and HTTP RMCP transports expose the **same tools, arguments, and respon When in doubt, prefer the HTTP `/mcp` endpoints described in the Overview. -This API reference should enable developers to effectively integrate Context Engine's MCP tools into their applications and workflows. \ No newline at end of file +This API reference should enable developers to effectively integrate Context Engine's MCP tools into their applications and workflows. diff --git a/docs/MULTI_REPO_COLLECTIONS.md b/docs/MULTI_REPO_COLLECTIONS.md index df11f623..eed1ba9b 100644 --- a/docs/MULTI_REPO_COLLECTIONS.md +++ b/docs/MULTI_REPO_COLLECTIONS.md @@ -341,11 +341,11 @@ results = client.search( ### 4. Monitor Collection Health ```bash -# Check collection status -make qdrant-status +# Check indexer health +curl http://localhost:${FASTMCP_INDEXER_HTTP_HEALTH_PORT:-18003}/readyz # List all collections -make qdrant-list +# Use the qdrant_list MCP tool from your MCP client. # Prune stale points make prune @@ -407,4 +407,3 @@ The architecture supports future enhancements: - [MCP API Reference](MCP_API.md) - [Architecture Overview](ARCHITECTURE.md) - [Development Guide](DEVELOPMENT.md) - diff --git a/docs/OBSERVABILITY.md b/docs/OBSERVABILITY.md index 985d2b93..e9502165 100644 --- a/docs/OBSERVABILITY.md +++ b/docs/OBSERVABILITY.md @@ -101,7 +101,7 @@ from qdrant_client import QdrantClient ### Qdrant client version -Use `qdrant-client>=1.15.0,<1.16.0`. Version 1.16+ changed to `.query_points()` which breaks OpenLit's instrumentation hooks. +Use `qdrant-client==1.15.1` with `qdrant/qdrant:v1.15.4`. Version 1.16+ removed the legacy `.search()` path used by OpenLit's Qdrant instrumentation hooks. ## Disabling diff --git a/docs/TROUBLESHOOTING.md b/docs/TROUBLESHOOTING.md index b3ea5c43..8e8aaa64 100644 --- a/docs/TROUBLESHOOTING.md +++ b/docs/TROUBLESHOOTING.md @@ -62,6 +62,31 @@ Prefer native path (`scripts/gpu_toggle.sh gpu`). For Docker, add `platform: lin ### Indexing stuck on large files Use `MAX_MICRO_CHUNKS_PER_FILE=200` during dev runs. +### Journal replay is hard to inspect in Kubernetes + +The watcher logs use `watch_index::journal_drain` for a bounded replay pass and +`watch_index::journal_drain_busy` when the current pass is still processing. A +large `backlog` is not by itself a failure; check whether it falls between log +entries and whether `journal_bulk_status_failed` or indexing errors appear. + +The upload service also exposes the shared journal state: + +```bash +curl -sG "$UPLOAD_SERVICE/api/v1/delta/status" \ + --data-urlencode 'workspace_path=/work/' | jq '.pending_operations, .server_info.journal' +``` + +The response includes pending/failed counts, the oldest retryable record, the +highest retry count, and a small sample of recorded errors. + +### Remote upload reports thousands of created files + +The client first reports local scan candidates. The remote plan then compares +their content hashes with the server replica cache and reports `content_needed`, +`skipped_hash_match`, and `needed_bytes`. A cold client may still hash and send +metadata for every candidate so the server can compare it; that does not mean +all candidates are uploaded. + ### Watcher timeouts (-9) or Qdrant "ResponseHandlingException: timed out" Set watcher-safe defaults to reduce payload size: @@ -151,5 +176,4 @@ docker-compose restart 1. Check this troubleshooting guide 2. Review logs: `docker compose logs mcp_indexer` 3. Verify health: `make health` -4. Check Qdrant status: `make qdrant-status` - +4. Check indexer health: `curl http://localhost:${FASTMCP_INDEXER_HTTP_HEALTH_PORT:-18003}/readyz`; use the `qdrant_status` MCP tool for collection details diff --git a/pytest.ini b/pytest.ini index 64bf831b..6a89673e 100644 --- a/pytest.ini +++ b/pytest.ini @@ -1,6 +1,6 @@ [pytest] pythonpath = . -addopts = -ra -vv --color=yes --durations=10 +addopts = -ra -vv --color=yes --durations=10 -m "not integration" asyncio_mode = auto testpaths = tests markers = @@ -9,4 +9,3 @@ markers = unit: fast unit tests that do not require services filterwarnings = ignore:The @wait_container_is_ready decorator is deprecated and will be removed in a future version:DeprecationWarning:testcontainers.core.waiting_utils - diff --git a/requirements.txt b/requirements.txt index f999c6ad..98155ef7 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,6 +1,7 @@ # Runtime dependencies (mirrors Dockerfiles) -# Pin to 1.15.x - version 1.16+ removed .search() method which breaks OpenLit instrumentation -qdrant-client>=1.15.0,<1.16.0 +# Pinned with Qdrant server v1.15.4. qdrant-client 1.16+ removed .search(), +# which breaks OpenLit's Qdrant instrumentation hooks. +qdrant-client==1.15.1 fastembed watchdog onnxruntime @@ -42,4 +43,3 @@ python-toon>=0.1.3 # # Benchmark-only (external suite) # coir-eval - diff --git a/scripts/admin_ui.py b/scripts/admin_ui.py index 32a27d55..131a9065 100644 --- a/scripts/admin_ui.py +++ b/scripts/admin_ui.py @@ -9,10 +9,7 @@ from starlette.templating import Jinja2Templates from jinja2 import select_autoescape -try: - from scripts.workspace_state import is_staging_enabled -except Exception: - is_staging_enabled = None # type: ignore +from scripts.workspace_state import is_staging_enabled _TEMPLATES_DIR = Path(__file__).resolve().parent.parent / "templates" _templates = Jinja2Templates(directory=str(_TEMPLATES_DIR)) @@ -25,8 +22,9 @@ def render_admin_login( status_code: int = 200, ) -> Any: return _templates.TemplateResponse( - "admin/login.html", - {"request": request, "title": "CTXCE Admin Login", "error": error}, + request=request, + name="admin/login.html", + context={"title": "CTXCE Admin Login", "error": error}, status_code=status_code, ) @@ -37,8 +35,9 @@ def render_admin_bootstrap( status_code: int = 200, ) -> Any: return _templates.TemplateResponse( - "admin/bootstrap.html", - {"request": request, "title": "CTXCE Admin Bootstrap", "error": error}, + request=request, + name="admin/bootstrap.html", + context={"title": "CTXCE Admin Bootstrap", "error": error}, status_code=status_code, ) @@ -54,16 +53,16 @@ def render_admin_acl( status_code: int = 200, ) -> Any: return _templates.TemplateResponse( - "admin/acl.html", - { - "request": request, + request=request, + name="admin/acl.html", + context={ "title": "CTXCE Admin ACL", "users": users, "collections": collections, "grants": grants, "deletion_enabled": bool(deletion_enabled), "work_dir": work_dir, - "staging_enabled": bool(is_staging_enabled() if callable(is_staging_enabled) else False), + "staging_enabled": bool(is_staging_enabled()), "refresh_ms": int(refresh_ms) if refresh_ms is not None else 5000, }, status_code=status_code, @@ -78,9 +77,9 @@ def render_admin_error( status_code: int = 400, ) -> Any: return _templates.TemplateResponse( - "admin/error.html", - { - "request": request, + request=request, + name="admin/error.html", + context={ "title": title, "message": message, "back_href": back_href, diff --git a/scripts/benchmarks/__init__.py b/scripts/benchmarks/__init__.py index 4d601d5d..06a7d21f 100644 --- a/scripts/benchmarks/__init__.py +++ b/scripts/benchmarks/__init__.py @@ -27,10 +27,7 @@ "SCENARIOS", # Component benchmarks (import on demand) # - eval_harness - # - trm_bench - # - router_bench # - refrag_bench # - expand_bench # - run_all ] - diff --git a/scripts/benchmarks/auto_tuner.py b/scripts/benchmarks/auto_tuner.py index fe5f69f6..f6063e8f 100644 --- a/scripts/benchmarks/auto_tuner.py +++ b/scripts/benchmarks/auto_tuner.py @@ -16,14 +16,11 @@ import math import os import statistics -import sys import time from dataclasses import dataclass, field, asdict from datetime import datetime -from pathlib import Path from typing import Any, Dict, List, Optional, Tuple -sys.path.insert(0, str(Path(__file__).parent.parent.parent)) # --------------------------------------------------------------------------- # Statistical Utilities diff --git a/scripts/benchmarks/coir/indexer.py b/scripts/benchmarks/coir/indexer.py index 60a84903..113222d5 100644 --- a/scripts/benchmarks/coir/indexer.py +++ b/scripts/benchmarks/coir/indexer.py @@ -7,15 +7,8 @@ """ from __future__ import annotations -import sys -from pathlib import Path from typing import Any, Dict, List -# Ensure project root is in path -sys.path.insert(0, str(Path(__file__).parent.parent.parent.parent)) - -from qdrant_client import QdrantClient - from scripts.benchmarks.core_indexer import ( BenchmarkDoc, index_benchmark_corpus, @@ -172,4 +165,4 @@ def cleanup_coir_collections(task_names: List[str] | None = None) -> int: except Exception as e: print(f"Failed to list collections: {e}") - return deleted \ No newline at end of file + return deleted diff --git a/scripts/benchmarks/coir/retriever.py b/scripts/benchmarks/coir/retriever.py index 9d3bbfd2..dfae835c 100644 --- a/scripts/benchmarks/coir/retriever.py +++ b/scripts/benchmarks/coir/retriever.py @@ -25,9 +25,7 @@ import asyncio import os -import sys from concurrent.futures import ThreadPoolExecutor -from pathlib import Path from typing import Any, Dict, List, Optional import numpy as np @@ -35,9 +33,6 @@ # Shared utilities from scripts.benchmarks.qdrant_utils import probe_pseudo_tags, verify_config_compatibility, get_qdrant_client -# Ensure project root is in path -sys.path.insert(0, str(Path(__file__).parent.parent.parent.parent)) - # Read .env settings EMBEDDING_MODEL = os.environ.get("EMBEDDING_MODEL", "BAAI/bge-base-en-v1.5") RERANKER_ENABLED = os.environ.get("RERANKER_ENABLED", "true").lower() in ("true", "1", "yes") diff --git a/scripts/benchmarks/core_indexer.py b/scripts/benchmarks/core_indexer.py index 78f2382f..831c813d 100644 --- a/scripts/benchmarks/core_indexer.py +++ b/scripts/benchmarks/core_indexer.py @@ -10,7 +10,7 @@ - Config fingerprinting for smart collection reuse - Batch processing with progress tracking - AST-aware symbol extraction -- Multi-vector support (dense, lexical, mini, pattern, sparse) +- Multi-vector support (dense, lexical, mini, sparse) """ from __future__ import annotations @@ -27,13 +27,11 @@ # Import production pipeline components from scripts.ingest.chunking import chunk_by_tokens, chunk_lines, chunk_semantic from scripts.ingest.pipeline import build_information, _select_dense_text -from scripts.ingest.vectors import project_mini, extract_pattern_vector +from scripts.ingest.vectors import project_mini from scripts.ingest.qdrant import ( hash_id, embed_batch, get_collection_vector_names, - PATTERN_VECTOR_NAME, - PATTERN_VECTOR_DIM, upsert_points as _upsert_points_with_retry, ) from scripts.utils import ( @@ -156,7 +154,6 @@ def get_config_fingerprint() -> str: f"mini_dim:{MINI_VEC_DIM}", f"mini_seed:{os.environ.get('MINI_VEC_SEED', '1337')}", f"lex_sparse:{os.environ.get('LEX_SPARSE_MODE', '0')}", - f"pattern_vectors:{os.environ.get('PATTERN_VECTORS', '0')}", f"index_micro:{os.environ.get('INDEX_MICRO_CHUNKS', '0')}", f"micro_tokens:{os.environ.get('MICRO_CHUNK_TOKENS', '16')}", f"micro_stride:{os.environ.get('MICRO_CHUNK_STRIDE', '')}", @@ -294,12 +291,6 @@ def create_collection( size=MINI_VEC_DIM, distance=models.Distance.COSINE ) - pattern_on = os.environ.get("PATTERN_VECTORS", "").strip().lower() in {"1", "true", "yes", "on"} - if pattern_on: - vectors_config[PATTERN_VECTOR_NAME] = models.VectorParams( - size=PATTERN_VECTOR_DIM, distance=models.Distance.COSINE - ) - sparse_cfg = None if LEX_SPARSE_MODE: sparse_cfg = { @@ -637,13 +628,6 @@ def prepare_doc_chunks(doc: BenchmarkDoc) -> List[ChunkMeta]: except Exception: pass - if PATTERN_VECTOR_NAME in available_vectors: - try: - pattern_vec = extract_pattern_vector(cm.chunk_text, cm.language) - if pattern_vec: - vectors_dict[PATTERN_VECTOR_NAME] = pattern_vec - except Exception: - pass sparse_dict = None if LEX_SPARSE_MODE and LEX_SPARSE_NAME in (available_vectors.get("sparse") or set()): @@ -697,4 +681,3 @@ def prepare_doc_chunks(doc: BenchmarkDoc) -> List[ChunkMeta]: "duration_sec": duration, "reused": False, } - diff --git a/scripts/benchmarks/cosqa/indexer.py b/scripts/benchmarks/cosqa/indexer.py index 9117fbc9..698d9699 100644 --- a/scripts/benchmarks/cosqa/indexer.py +++ b/scripts/benchmarks/cosqa/indexer.py @@ -8,13 +8,8 @@ from __future__ import annotations import os -import sys -from pathlib import Path from typing import Any, Dict, List -# Ensure project root is in path -sys.path.insert(0, str(Path(__file__).parent.parent.parent.parent)) - from scripts.benchmarks.core_indexer import ( BenchmarkDoc, index_benchmark_corpus, diff --git a/scripts/benchmarks/cosqa/pca_init.py b/scripts/benchmarks/cosqa/pca_init.py deleted file mode 100644 index f66cc423..00000000 --- a/scripts/benchmarks/cosqa/pca_init.py +++ /dev/null @@ -1,144 +0,0 @@ -#!/usr/bin/env python3 -""" -PCA-based projection initialization for CoSQA learning. - -Pre-computes PCA on corpus embeddings to initialize the projection layer -with meaningful weights instead of random noise. -""" -import os -import sys -from pathlib import Path -from typing import List, Dict, Any - -import numpy as np - -# Ensure project root is in path -sys.path.insert(0, str(Path(__file__).parent.parent.parent.parent)) - - -def compute_pca_init_for_collection( - collection: str, - sample_limit: int = 1000, -) -> bool: - """Compute PCA initialization for a collection's projection layer. - - Args: - collection: Qdrant collection name - sample_limit: Max documents to sample for PCA (for efficiency) - - Returns: - True if successful, False otherwise - """ - from scripts.embedder import get_embedding_model, get_model_dimension - from scripts.benchmarks.qdrant_utils import get_qdrant_client - from scripts.rerank_recursive.projection import LearnedProjection - - print(f"Computing PCA initialization for collection: {collection}") - - # Get model and dimension - model = get_embedding_model() - embed_dim = get_model_dimension() - print(f" Embedding model: {os.environ.get('EMBEDDING_MODEL', 'default')}") - print(f" Embedding dimension: {embed_dim}") - - # Sample documents from collection - client = get_qdrant_client() - try: - info = client.get_collection(collection) - total_points = info.points_count - print(f" Collection has {total_points} points") - except Exception as e: - print(f" ERROR: Failed to get collection info: {e}") - return False - - # Sample points - sample_size = min(sample_limit, total_points) - print(f" Sampling {sample_size} points for PCA...") - - try: - # Scroll through collection to get sample - points = [] - offset = None - batch_size = 100 - - while len(points) < sample_size: - result = client.scroll( - collection_name=collection, - limit=min(batch_size, sample_size - len(points)), - offset=offset, - with_payload=True, - with_vectors=False, # We'll re-embed from text - ) - - batch_points, offset = result - if not batch_points: - break - - points.extend(batch_points) - - if offset is None: - break - - print(f" Sampled {len(points)} points") - - except Exception as e: - print(f" ERROR: Failed to sample points: {e}") - return False - - # Extract text and embed - print(f" Embedding {len(points)} documents...") - texts = [] - for point in points: - payload = point.payload or {} - # Try to get code or content - text = payload.get("code") or payload.get("content") or payload.get("text") or "" - if text: - texts.append(str(text)[:2000]) # Limit length - - if not texts: - print(" ERROR: No text found in sampled points") - return False - - print(f" Extracted {len(texts)} text samples") - - # Batch embed - embeddings_list = [] - batch_size = 64 - for i in range(0, len(texts), batch_size): - batch = texts[i:i+batch_size] - batch_embs = list(model.embed(batch)) - embeddings_list.extend(batch_embs) - if (i // batch_size + 1) % 10 == 0: - print(f" Embedded {i+len(batch)}/{len(texts)}") - - embeddings = np.array(embeddings_list, dtype=np.float32) - print(f" Embeddings shape: {embeddings.shape}") - - # Initialize projection with PCA - print(f" Computing PCA projection ({embed_dim} → 256)...") - projection = LearnedProjection(input_dim=embed_dim, output_dim=256, lr=0.0) - projection.set_collection(collection) - projection.init_from_pca(embeddings) - - # Save weights - projection._save_weights() - print(f" ✓ Saved PCA-initialized weights to: {projection._weights_path}") - - return True - - -if __name__ == "__main__": - import argparse - - parser = argparse.ArgumentParser(description="Pre-compute PCA initialization for learning") - parser.add_argument("--collection", default="cosqa-corpus", help="Collection name") - parser.add_argument("--sample-limit", type=int, default=1000, help="Max samples for PCA") - args = parser.parse_args() - - success = compute_pca_init_for_collection( - collection=args.collection, - sample_limit=args.sample_limit, - ) - - sys.exit(0 if success else 1) - diff --git a/scripts/benchmarks/cosqa/run_ablation.sh b/scripts/benchmarks/cosqa/run_ablation.sh index 550cc29b..94414801 100644 --- a/scripts/benchmarks/cosqa/run_ablation.sh +++ b/scripts/benchmarks/cosqa/run_ablation.sh @@ -26,7 +26,6 @@ run_one() { local refrag="$2" local micro="$3" local rerank="$4" - local learning="$5" local collection="${COLL_PREFIX}-${label}-${RUN_TAG}" local output="${OUT_DIR}/cosqa_${label}.json" local log="${LOG_DIR}/cosqa_${label}.log" @@ -48,17 +47,11 @@ run_one() { if [ "${rerank}" = "0" ]; then args+=("--no-rerank") fi - if [ "${learning}" = "1" ]; then - args+=("--learning-worker") - fi - mkdir -p "${OUT_DIR}" "${LOG_DIR}" ( export "${BASE_ENV[@]}" export REFRAG_MODE="${refrag}" export INDEX_MICRO_CHUNKS="${micro}" - export RERANK_LEARNING="${learning}" - export RERANK_EVENTS_ENABLED="${learning}" "${PYTHON_BIN}" "${args[@]}" ) > "${log}" 2>&1 @@ -66,9 +59,7 @@ run_one() { echo " log: ${log}" } -run_one "norerank_norefrag" 0 0 0 0 -run_one "norerank_refrag" 1 1 0 0 -run_one "rerank_norefrag" 0 0 1 0 -run_one "rerank_refrag" 1 1 1 0 -run_one "rerank_norefrag_learning" 0 0 1 1 -run_one "rerank_refrag_learning" 1 1 1 1 +run_one "norerank_norefrag" 0 0 0 +run_one "norerank_refrag" 1 1 0 +run_one "rerank_norefrag" 0 0 1 +run_one "rerank_refrag" 1 1 1 diff --git a/scripts/benchmarks/cosqa/run_search_matrix.sh b/scripts/benchmarks/cosqa/run_search_matrix.sh new file mode 100755 index 00000000..be379877 --- /dev/null +++ b/scripts/benchmarks/cosqa/run_search_matrix.sh @@ -0,0 +1,313 @@ +#!/usr/bin/env bash +set -euo pipefail + +ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/../../.." && pwd)" +cd "${ROOT_DIR}" + +PYTHON_BIN="${PYTHON_BIN:-}" +if [ -z "${PYTHON_BIN}" ]; then + if command -v python3.11 >/dev/null 2>&1; then + PYTHON_BIN="python3.11" + elif command -v python3 >/dev/null 2>&1; then + PYTHON_BIN="python3" + elif command -v python >/dev/null 2>&1; then + PYTHON_BIN="python" + else + echo "No Python interpreter found (looked for python3.11/python3/python)." >&2 + exit 127 + fi +fi +RUN_TAG="${RUN_TAG:-$(date +%Y%m%d-%H%M%S)}" +PROFILE="${PROFILE:-full}" # smoke | quick | full +RUN_SET="${RUN_SET:-full}" # pr | knobs | nightly | full +OUT_DIR="${OUT_DIR:-bench_results/cosqa/${RUN_TAG}}" +LOG_DIR="${LOG_DIR:-${OUT_DIR}}" +SPLIT="${SPLIT:-test}" +COLLECTION="${COLLECTION:-cosqa-search-${RUN_TAG}}" +LIMIT="${LIMIT:-10}" +RECREATE_INDEX="${RECREATE_INDEX:-1}" +ENFORCE_HYBRID_GATE="${ENFORCE_HYBRID_GATE:-0}" +HYBRID_MIN_DELTA="${HYBRID_MIN_DELTA:--0.020}" + +case "${PROFILE}" in + smoke) + : "${CORPUS_LIMIT:=150}" + : "${QUERY_LIMIT:=30}" + ;; + quick) + : "${CORPUS_LIMIT:=500}" + : "${QUERY_LIMIT:=100}" + ;; + full) + : "${CORPUS_LIMIT:=0}" + : "${QUERY_LIMIT:=0}" + ;; + *) + echo "Unknown PROFILE='${PROFILE}'. Use smoke|quick|full" >&2 + exit 2 + ;; +esac + +mkdir -p "${OUT_DIR}" "${LOG_DIR}" + +BASE_ENV=( + "LOG_LEVEL=${LOG_LEVEL:-INFO}" + "DEBUG_HYBRID_SEARCH=${DEBUG_HYBRID_SEARCH:-0}" + "QDRANT_URL=${QDRANT_URL:-http://localhost:6333}" + "HYBRID_IN_PROCESS=${HYBRID_IN_PROCESS:-1}" + "RERANK_IN_PROCESS=${RERANK_IN_PROCESS:-1}" + "LEX_VECTOR_DIM=${LEX_VECTOR_DIM:-4096}" + "COSQA_QUERY_CONCURRENCY=${COSQA_QUERY_CONCURRENCY:-8}" + "LLM_EXPAND_MAX=0" + "REFRAG_DECODER=0" +) + +run_index_once() { + local log="${LOG_DIR}/cosqa_index.log" + local args=( + "-m" "scripts.benchmarks.cosqa.runner" + "--split" "${SPLIT}" + "--collection" "${COLLECTION}" + "--limit" "${LIMIT}" + "--index-only" + ) + + if [ "${CORPUS_LIMIT}" -gt 0 ]; then + args+=("--corpus-limit" "${CORPUS_LIMIT}") + fi + if [ "${QUERY_LIMIT}" -gt 0 ]; then + args+=("--query-limit" "${QUERY_LIMIT}") + fi + if [ "${RECREATE_INDEX}" = "1" ]; then + args+=("--recreate") + fi + + echo "[index] collection=${COLLECTION} corpus_limit=${CORPUS_LIMIT} query_limit=${QUERY_LIMIT}" | tee "${log}" + ( + export "${BASE_ENV[@]}" + "${PYTHON_BIN}" "${args[@]}" + ) >> "${log}" 2>&1 +} + +preflight_python_deps() { + "${PYTHON_BIN}" - <<'PY' +import importlib.util + +required = ["qdrant_client", "datasets"] +missing = [m for m in required if importlib.util.find_spec(m) is None] +if missing: + raise SystemExit( + "Missing Python deps for CoSQA benchmark: " + + ", ".join(missing) + + ". Install them before running." + ) +PY +} + +verify_collection_ready() { + "${PYTHON_BIN}" - "${COLLECTION}" <<'PY' +import os +import sys +from qdrant_client import QdrantClient + +collection = sys.argv[1] +url = os.environ.get("QDRANT_URL", "http://localhost:6333") +client = QdrantClient(url=url, timeout=60) +info = client.get_collection(collection) +points = int(info.points_count or 0) +if points <= 0: + raise RuntimeError(f"Collection '{collection}' has no points after indexing") +print(f"[verify] collection={collection} points={points}") +PY +} + +run_case() { + local label="$1" + local mode="$2" + local rerank="$3" + local expand="$4" + local lex_mode="$5" + shift 5 + + local output="${OUT_DIR}/cosqa_${label}.json" + local log="${LOG_DIR}/cosqa_${label}.log" + + local args=( + "-m" "scripts.benchmarks.cosqa.runner" + "--split" "${SPLIT}" + "--collection" "${COLLECTION}" + "--limit" "${LIMIT}" + "--skip-index" + "--mode" "${mode}" + "--output" "${output}" + ) + + if [ "${CORPUS_LIMIT}" -gt 0 ]; then + args+=("--corpus-limit" "${CORPUS_LIMIT}") + fi + if [ "${QUERY_LIMIT}" -gt 0 ]; then + args+=("--query-limit" "${QUERY_LIMIT}") + fi + if [ "${rerank}" = "0" ]; then + args+=("--no-rerank") + fi + if [ "${expand}" = "0" ]; then + args+=("--no-expand") + fi + + local case_env=("HYBRID_LEXICAL_TEXT_MODE=${lex_mode}") + for kv in "$@"; do + case_env+=("${kv}") + done + + echo "[run] ${label} mode=${mode} rerank=${rerank} expand=${expand} lex_mode=${lex_mode}" | tee "${log}" + ( + export "${BASE_ENV[@]}" + export "${case_env[@]}" + "${PYTHON_BIN}" "${args[@]}" + ) >> "${log}" 2>&1 + + echo "[ok] ${output}" +} + +CASES=() +case "${RUN_SET}" in + pr) + CASES=( + "dense_norerank|dense|0|0|raw" + "hybrid_rerank_lexrrf|hybrid|1|0|rrf" + "hybrid_rerank_expand_lexrrf|hybrid|1|1|rrf" + ) + ;; + knobs) + CASES=( + "dense_norerank|dense|0|0|raw" + "dense_rerank|dense|1|0|raw" + "hybrid_norerank_lexraw|hybrid|0|0|raw" + "hybrid_norerank_lexrrf|hybrid|0|0|rrf" + "hybrid_rerank_lexraw|hybrid|1|0|raw" + "hybrid_rerank_lexrrf|hybrid|1|0|rrf" + "hybrid_rerank_expand_lexrrf|hybrid|1|1|rrf" + "lexical_norerank|lexical|0|0|raw" + ) + ;; + nightly) + CASES=( + "dense_norerank|dense|0|0|raw" + "dense_rerank|dense|1|0|raw" + "hybrid_norerank_lexraw|hybrid|0|0|raw" + "hybrid_norerank_lexrrf|hybrid|0|0|rrf" + "hybrid_rerank_lexraw|hybrid|1|0|raw" + "hybrid_rerank_lexrrf|hybrid|1|0|rrf" + "hybrid_rerank_expand_lexrrf|hybrid|1|1|rrf" + "lexical_norerank|lexical|0|0|raw" + ) + ;; + full) + CASES=( + "dense_norerank|dense|0|0|raw" + "dense_rerank|dense|1|0|raw" + "hybrid_norerank_lexraw|hybrid|0|0|raw" + "hybrid_norerank_lexrrf|hybrid|0|0|rrf" + "hybrid_rerank_lexraw|hybrid|1|0|raw" + "hybrid_rerank_lexrrf|hybrid|1|0|rrf" + "hybrid_rerank_expand_lexrrf|hybrid|1|1|rrf" + "lexical_norerank|lexical|0|0|raw" + ) + ;; + *) + echo "Unknown RUN_SET='${RUN_SET}'. Use pr|knobs|nightly|full" >&2 + exit 2 + ;; +esac + +echo "[config] run_tag=${RUN_TAG} profile=${PROFILE} run_set=${RUN_SET} out_dir=${OUT_DIR}" +preflight_python_deps +run_index_once +verify_collection_ready + +for spec in "${CASES[@]}"; do + IFS='|' read -r label mode rerank expand lex_mode <<< "${spec}" + run_case "${label}" "${mode}" "${rerank}" "${expand}" "${lex_mode}" +done + +"${PYTHON_BIN}" - "${OUT_DIR}" "${ENFORCE_HYBRID_GATE}" "${HYBRID_MIN_DELTA}" <<'PY' +import json +import sys +from pathlib import Path + +out_dir = Path(sys.argv[1]) +enforce_gate = str(sys.argv[2]).strip() in {"1", "true", "yes"} +min_delta = float(sys.argv[3]) + +rows = [] +for path in sorted(out_dir.glob("cosqa_*.json")): + if path.name.startswith("cosqa_index") or path.name.endswith("_meta.json") or path.name.startswith("summary"): + continue + with path.open("r", encoding="utf-8") as f: + data = json.load(f) + if not isinstance(data, dict) or "metrics" not in data or "config" not in data: + continue + metrics = data.get("metrics") or {} + config = data.get("config") or {} + env = (config.get("env") or {}) if isinstance(config, dict) else {} + rows.append({ + "label": path.stem.replace("cosqa_", ""), + "mode": config.get("mode", ""), + "rerank": bool(config.get("rerank_enabled", False)), + "expand": env.get("HYBRID_EXPAND", ""), + "lex_mode": env.get("HYBRID_LEXICAL_TEXT_MODE", ""), + "mrr": float(metrics.get("mrr", 0.0) or 0.0), + "recall_10": float(metrics.get("recall@10", 0.0) or 0.0), + "ndcg_10": float(metrics.get("ndcg@10", 0.0) or 0.0), + "lat_ms": float((data.get("latency") or {}).get("avg_ms", 0.0) or 0.0), + "file": path.name, + }) + +if not rows: + print("No CoSQA result JSON files found.", file=sys.stderr) + sys.exit(3) + +rows.sort(key=lambda r: (-r["mrr"], -r["recall_10"])) + +summary = { + "ranked": rows, + "best": rows[0], +} +(out_dir / "summary.json").write_text(json.dumps(summary, indent=2), encoding="utf-8") + +lines = [ + "# CoSQA Search Matrix Summary", + "", + "| Rank | Label | Mode | Rerank | Expand | LexMode | MRR | R@10 | NDCG@10 | Avg Lat (ms) |", + "|---:|---|---|---:|---:|---|---:|---:|---:|---:|", +] +for i, r in enumerate(rows, start=1): + lines.append( + f"| {i} | {r['label']} | {r['mode']} | {int(r['rerank'])} | {r['expand']} | {r['lex_mode']} | " + f"{r['mrr']:.4f} | {r['recall_10']:.4f} | {r['ndcg_10']:.4f} | {r['lat_ms']:.2f} |" + ) + +best_dense = max((r for r in rows if r["mode"] == "dense"), key=lambda r: r["mrr"], default=None) +best_hybrid = max((r for r in rows if r["mode"] == "hybrid"), key=lambda r: r["mrr"], default=None) +if best_dense and best_hybrid: + delta = best_hybrid["mrr"] - best_dense["mrr"] + lines.append("") + lines.append( + f"Best hybrid ({best_hybrid['label']}) vs best dense ({best_dense['label']}): " + f"delta MRR = {delta:+.4f}" + ) + if enforce_gate and delta < min_delta: + lines.append( + f"Gate failed: hybrid delta {delta:+.4f} is below required minimum {min_delta:+.4f}" + ) + (out_dir / "summary.md").write_text("\n".join(lines) + "\n", encoding="utf-8") + print("\n".join(lines)) + sys.exit(4) + +(out_dir / "summary.md").write_text("\n".join(lines) + "\n", encoding="utf-8") +print("\n".join(lines)) +PY + +echo "[done] results=${OUT_DIR}" +echo "[done] summary=${OUT_DIR}/summary.md" diff --git a/scripts/benchmarks/cosqa/runner.py b/scripts/benchmarks/cosqa/runner.py index 8e0c32b1..4c0fbb3f 100644 --- a/scripts/benchmarks/cosqa/runner.py +++ b/scripts/benchmarks/cosqa/runner.py @@ -22,7 +22,6 @@ - [x] Enriched embeddings (symbols + imports + docstring + code) - [x] Lexical hash vectors (for hybrid search) - [ ] ReFRAG/micro-chunks (requires REFRAG_MODE=1, off by default) -- [ ] Pattern vectors (requires indexed pattern vectors) **Not Applicable (CoSQA limitations):** - N/A Semantic chunking (snippets are atomic units) @@ -54,19 +53,12 @@ import json import math import os -import subprocess import sys import time from dataclasses import dataclass, field, asdict from datetime import datetime from pathlib import Path from typing import Any, Dict, List, Optional, Tuple -from dotenv import load_dotenv - -# Load .env immediately to ensure all subsequent imports (like scripts.ingest.config) -# see the correct environment variables. -load_dotenv(override=True) - from scripts.benchmarks.qdrant_utils import ( get_qdrant_client, probe_pseudo_tags, @@ -77,9 +69,6 @@ os.environ["OPENLIT_ENABLED"] = "0" os.environ["OTEL_SDK_DISABLED"] = "true" -# Ensure project root is in path -sys.path.insert(0, str(Path(__file__).parent.parent.parent.parent)) - # NOTE: .env loading moved to _load_benchmark_env() to avoid polluting # environment when this module is imported (e.g., by tests or __init__.py). # Call _load_benchmark_env() explicitly before running benchmarks. @@ -111,11 +100,6 @@ def _load_benchmark_env() -> None: # Set defaults AFTER loading .env so .env takes priority os.environ.setdefault("RERANKER_MODEL", "jinaai/jina-reranker-v2-base-multilingual") os.environ.setdefault("RERANK_IN_PROCESS", "1") - # Hard-disable learning/recursive reranker for deterministic benchmarks - # (unless explicitly enabled via --learning-worker flag or COSQA_ENABLE_LEARNING env var) - if not os.environ.get("COSQA_ENABLE_LEARNING") and os.environ.get("RERANK_LEARNING") != "1": - os.environ["RERANK_LEARNING"] = "0" - os.environ["RERANK_EVENTS_ENABLED"] = "0" # Disable sparse vectors for CoSQA benchmarks to avoid missing lex-sparse dims os.environ["LEX_SPARSE_MODE"] = "0" # Benchmarks should not be scoped or cached by workspace repo state @@ -442,7 +426,16 @@ def _cosqa_id_from_path(p: str) -> Optional[str]: name = s.rsplit("/", 1)[-1] if name.endswith(".py"): name = name[: -3] - return name.strip() or None + name = name.strip() + if not name: + return None + # CoSQA synthetic filenames are often "__". + # Recover canonical code_id so relevance matching aligns with qrels. + if "__" in name: + tail = name.rsplit("__", 1)[-1].strip() + if tail.startswith("cosqa-"): + return tail + return name # Extract stable code_ids for evaluation. # NOTE: rerank paths may not include payload; for CoSQA we can fall back to parsing @@ -555,7 +548,6 @@ async def run_cosqa_benchmark( # Reranker "RERANKER_MODEL": os.environ.get("RERANKER_MODEL", ""), "RERANK_IN_PROCESS": os.environ.get("RERANK_IN_PROCESS", ""), - "RERANK_LEARNING": os.environ.get("RERANK_LEARNING", ""), # Hybrid search "HYBRID_IN_PROCESS": os.environ.get("HYBRID_IN_PROCESS", ""), "HYBRID_EXPAND": os.environ.get("HYBRID_EXPAND", ""), @@ -915,8 +907,8 @@ async def run_full_benchmark( print(f" Limited corpus to {len(corpus)} entries") if skip_index: - print(" [skip-index] Skipping indexing...") - result = {"reused": True, "indexed": len(corpus), "skipped": 0, "errors": 0} + print(" [skip-index] Skipping indexing (using existing collection as-is)...") + result = {"reused": False, "indexed": 0, "skipped": len(corpus), "errors": 0} else: # Check if already indexed (use fingerprint matching, not just points_count) # The indexer handles fingerprint checking internally and will recreate if needed @@ -965,17 +957,6 @@ async def run_full_benchmark( return report -def _spawn_learning_worker(collection: str, project_root: Path) -> subprocess.Popen: - cmd = [ - sys.executable, - str(project_root / "scripts" / "learning_reranker_worker.py"), - "--daemon", - "--collection", - collection, - ] - return subprocess.Popen(cmd, cwd=project_root, env=os.environ.copy()) - - def main(): """CLI entrypoint for CoSQA benchmark.""" @@ -996,8 +977,6 @@ def main(): help="Disable query expansion") parser.add_argument("--recreate", action="store_true", help="Recreate index from scratch") - parser.add_argument("--learning-worker", action="store_true", - help="Spawn learning reranker worker during the run (enables learning + event logging)") parser.add_argument("--pure-semantic", action="store_true", help="Disable FNAME_BOOST and other heuristics (old hardened mode)") parser.add_argument("--enable-llm", action="store_true", @@ -1018,6 +997,12 @@ def main(): help="Search mode: 'hybrid' (default), 'dense' (pure semantic), or 'lexical' (pure BM25-style)") args = parser.parse_args() + # Benchmarks must not require MCP auth sessions. + # runner imports dotenv at module import time with override=True, so enforce this + # after args parsing to guarantee process-local benchmark behavior. + os.environ["CTXCE_AUTH_ENABLED"] = "0" + os.environ["CTXCE_MCP_ACL_ENFORCE"] = "0" + # Enable Context-Engine features for accurate benchmarking. # Semantic expansion is always enabled (it may still be a no-op if query expansion is disabled). os.environ["SEMANTIC_EXPANSION_ENABLED"] = "1" @@ -1067,26 +1052,6 @@ def main(): os.environ["RERANKER_ONNX_PATH"] = str(_project_root / "models" / "model_qint8_avx512_vnni.onnx") os.environ["RERANKER_TOKENIZER_PATH"] = str(_project_root / "models" / "tokenizer.json") - learning_proc = None - if args.learning_worker: - if args.no_rerank: - print(" [WARN] --learning-worker ignored because --no-rerank is set") - else: - # Pre-compute PCA initialization for projection layer (cold-start fix) - print(" [learning] Pre-computing PCA initialization...") - from scripts.benchmarks.cosqa.pca_init import compute_pca_init_for_collection - pca_success = compute_pca_init_for_collection( - collection=args.collection, - sample_limit=1000, - ) - if not pca_success: - print(" [WARN] PCA initialization failed, using random init") - - os.environ["RERANK_LEARNING"] = "1" - os.environ["RERANK_EVENTS_ENABLED"] = "1" - learning_proc = _spawn_learning_worker(args.collection, _project_root) - print(f" [learning-worker] Started (pid {learning_proc.pid}) for {args.collection}") - # Verify config compatibility BEFORE running anything if not args.recreate: try: @@ -1094,30 +1059,20 @@ def main(): verify_config_compatibility(get_qdrant_client(), args.collection) except Exception as e: print(f"\nCONFIGURATION ERROR: {e}") - if learning_proc and learning_proc.poll() is None: - learning_proc.kill() sys.exit(1) - try: - report = asyncio.run(run_full_benchmark( - split=args.split, - collection=args.collection, - limit=args.limit, - query_limit=args.query_limit, - corpus_limit=args.corpus_limit, - rerank_enabled=not args.no_rerank, - mode=args.mode, - recreate_index=args.recreate, - index_only=args.index_only, - skip_index=args.skip_index, - )) - finally: - if learning_proc and learning_proc.poll() is None: - learning_proc.terminate() - try: - learning_proc.wait(timeout=5) - except subprocess.TimeoutExpired: - learning_proc.kill() + report = asyncio.run(run_full_benchmark( + split=args.split, + collection=args.collection, + limit=args.limit, + query_limit=args.query_limit, + corpus_limit=args.corpus_limit, + rerank_enabled=not args.no_rerank, + mode=args.mode, + recreate_index=args.recreate, + index_only=args.index_only, + skip_index=args.skip_index, + )) if report: print_report(report) diff --git a/scripts/benchmarks/efficiency_benchmark.py b/scripts/benchmarks/efficiency_benchmark.py index af9dbf55..d567bd15 100644 --- a/scripts/benchmarks/efficiency_benchmark.py +++ b/scripts/benchmarks/efficiency_benchmark.py @@ -10,17 +10,11 @@ import asyncio import json import os -import sys import time import hashlib from dataclasses import dataclass, field, asdict -from pathlib import Path from typing import Any, Callable, Dict, List, Optional, Tuple -# Add project root to path -sys.path.insert(0, str(Path(__file__).parent.parent.parent)) - -# Shared stats helpers (after sys.path setup) from scripts.benchmarks.common import percentile, extract_result_paths, resolve_collection_auto # Ensure correct collection is used (read from workspace state or env) @@ -228,7 +222,7 @@ def compute_recall_at_k(expected_paths: List[str], result_paths: List[str], k: i "queries": [ {"tool": "repo_search", "query": "init_openlit error handling"}, {"tool": "symbol_graph", "symbol": "init_openlit", "query_type": "callers"}, - {"tool": "search_tests_for", "query": "openlit initialization"}, + {"tool": "repo_search", "query": "openlit initialization", "profile": "tests"}, ], "expected_paths": ["openlit_init.py", "test_openlit"], }, @@ -237,7 +231,7 @@ def compute_recall_at_k(expected_paths: List[str], result_paths: List[str], k: i "queries": [ {"tool": "repo_search", "query": "memory store implementation pattern"}, {"tool": "context_answer", "query": "How does memory_store work?"}, - {"tool": "search_config_for", "query": "memory collection settings"}, + {"tool": "repo_search", "query": "memory collection settings", "profile": "config"}, ], "expected_paths": ["mcp_impl/memory.py", "memory_store"], }, @@ -245,10 +239,10 @@ def compute_recall_at_k(expected_paths: List[str], result_paths: List[str], k: i "description": "Trace dependencies across multiple files", "queries": [ {"tool": "symbol_graph", "symbol": "get_embedding_model", "query_type": "callers"}, - {"tool": "search_importers_for", "query": "embedder"}, + {"tool": "repo_search", "query": "embedder", "profile": "code"}, {"tool": "repo_search", "query": "embedding dimension vector size"}, ], - "expected_paths": ["embedder.py", "rerank_recursive"], + "expected_paths": ["embedder.py"], }, "documentation": { "description": "Generate explanation of a module", @@ -460,18 +454,12 @@ async def run_benchmark( repo_search, context_answer, symbol_graph, - search_tests_for, - search_config_for, - search_importers_for, memory_find, ) tool_registry = { "repo_search": repo_search, "context_answer": context_answer, "symbol_graph": symbol_graph, - "search_tests_for": search_tests_for, - "search_config_for": search_config_for, - "search_importers_for": search_importers_for, "memory_find": memory_find, } except ImportError as e: diff --git a/scripts/benchmarks/eval_harness.py b/scripts/benchmarks/eval_harness.py index efb7611f..c686628d 100644 --- a/scripts/benchmarks/eval_harness.py +++ b/scripts/benchmarks/eval_harness.py @@ -10,14 +10,12 @@ import asyncio import json import os -import sys import time from dataclasses import dataclass, field, asdict from pathlib import Path from typing import Any, Dict, List, Optional import statistics -sys.path.insert(0, str(Path(__file__).parent.parent.parent)) # Shared stats helpers from scripts.benchmarks.common import ( @@ -186,7 +184,7 @@ def compute_precision(expected: List[str], retrieved: List[str], k: int) -> floa {"query": "hybrid search RRF ranking", "expected": ["hybrid/ranking.py"]}, {"query": "memory store implementation", "expected": ["mcp_impl/memory.py"]}, {"query": "openlit initialization tracing", "expected": ["openlit_init.py"]}, - {"query": "recursive reranker learning", "expected": ["rerank_recursive"]}, + {"query": "relevance feedback ratings", "expected": ["relevance_feedback.py", "mcp_impl/search.py"]}, {"query": "embedder model loading", "expected": ["embedder.py"]}, {"query": "workspace state persistence", "expected": ["workspace_state.py"]}, {"query": "symbol graph callers", "expected": ["symbol_graph.py", "mcp_impl"]}, diff --git a/scripts/benchmarks/expand_bench.py b/scripts/benchmarks/expand_bench.py index 86330572..82da2ae8 100644 --- a/scripts/benchmarks/expand_bench.py +++ b/scripts/benchmarks/expand_bench.py @@ -9,14 +9,11 @@ import asyncio import json import os -import sys import time from dataclasses import dataclass, field, asdict -from pathlib import Path from typing import Any, Dict, List, Optional import statistics -sys.path.insert(0, str(Path(__file__).parent.parent.parent)) # Load environment (optional) and fix Docker hostname try: diff --git a/scripts/benchmarks/grounding_scorer.py b/scripts/benchmarks/grounding_scorer.py index 669d85a8..fbc92ac3 100644 --- a/scripts/benchmarks/grounding_scorer.py +++ b/scripts/benchmarks/grounding_scorer.py @@ -13,15 +13,12 @@ import json import os import re -import sys from dataclasses import dataclass, field from datetime import datetime from pathlib import Path from typing import Any, Dict, List, Optional, Set, Tuple -# Add project root to path PROJECT_ROOT = Path(__file__).parent.parent.parent -sys.path.insert(0, str(PROJECT_ROOT)) # Load environment variables from .env try: @@ -110,9 +107,9 @@ def to_dict(self) -> Dict[str, Any]: "expected_files": ["hybrid_search.py", "ranking.py"], }, { - "query": "What is the purpose of the recursive reranker's latent state z?", - "expected_topics": ["latent", "state", "refinement", "TRM", "iteration"], - "expected_files": ["rerank_recursive", "core.py"], + "query": "How does relevance feedback recall adjacent code?", + "expected_topics": ["feedback", "ratings", "graph", "caller", "relevance"], + "expected_files": ["relevance_feedback.py", "search.py"], }, { "query": "How does context_answer handle insufficient context?", diff --git a/scripts/benchmarks/recommendations.py b/scripts/benchmarks/recommendations.py index db1968eb..96bf97a1 100644 --- a/scripts/benchmarks/recommendations.py +++ b/scripts/benchmarks/recommendations.py @@ -181,26 +181,6 @@ "category": "expansion", }, - # === Learning Reranker === - "RERANK_LLM_TEACHER": { - "description": "Enable LLM teacher for learning reranker", - "default": 1, "type": "bool", "range": (0, 1), - "impacts": ["learning_quality"], - "category": "learning", - }, - "RERANK_LLM_SAMPLE_RATE": { - "description": "Sample rate for LLM teacher feedback", - "default": 1.0, "type": "float", "range": (0.0, 1.0), - "impacts": ["learning_speed", "cost"], - "category": "learning", - }, - "RERANK_VICREG_WEIGHT": { - "description": "Weight for VICReg regularization loss", - "default": 0.1, "type": "float", "range": (0.0, 1.0), - "impacts": ["embedding_diversity"], - "category": "learning", - }, - # === Limits / Defaults === "REPO_SEARCH_DEFAULT_LIMIT": { "description": "Default result limit for repo_search", @@ -208,19 +188,6 @@ "impacts": ["token_usage", "recall"], "category": "limits", }, - "INFO_REQUEST_LIMIT": { - "description": "Default limit for info_request tool", - "default": 10, "type": "int", "range": (3, 20), - "impacts": ["token_usage"], - "category": "limits", - }, - "INFO_REQUEST_CONTEXT_LINES": { - "description": "Context lines for info_request", - "default": 5, "type": "int", "range": (0, 15), - "impacts": ["context_density"], - "category": "limits", - }, - # === Lexical Search === "LEX_MULTI_HASH": { "description": "Multi-hash buckets per token (reduces collisions)", @@ -298,19 +265,15 @@ def generate_recommendations( # Extract metrics eval_metrics = components.get("eval_harness", {}).get("metrics", {}) eval_latency = components.get("eval_harness", {}).get("latency", {}) - trm_metrics = components.get("trm_reranker", {}).get("metrics", {}) refrag_metrics = components.get("refrag", {}).get("metrics", {}) - expand_metrics = components.get("query_expansion", {}).get("metrics", {}) mrr = eval_metrics.get("mrr", 0) recall_5 = eval_metrics.get("recall@5", 0) recall_10 = eval_metrics.get("recall@10", 0) - precision_5 = eval_metrics.get("precision@5", 0) # Note: eval harness emits p90_ms, not p90 - p90_latency = eval_latency.get("p90_ms", 0) or trm_metrics.get("p90_latency_ms", 0) + p90_latency = eval_latency.get("p90_ms", 0) grounding = refrag_metrics.get("grounding_rate", 1.0) citations = refrag_metrics.get("avg_citations", 0) - kendall_tau = trm_metrics.get("kendall_tau", 0) # ========================================================================= # MRR Recommendations @@ -452,25 +415,6 @@ def generate_recommendations( "rationale": "PRF adds an extra retrieval round-trip.", }) - # ========================================================================= - # Learning Reranker Recommendations - # ========================================================================= - if kendall_tau > 0.9: - # High correlation = learning is working well - if current_config.get("RERANK_LLM_SAMPLE_RATE", 1.0) > 0.5: - recommendations.append({ - "priority": "low", - "component": "learning", - "metric": "kendall_tau", - "current_value": kendall_tau, - "target_value": 0.95, - "action": "Reduce RERANK_LLM_SAMPLE_RATE to save cost", - "config_changes": {"RERANK_LLM_SAMPLE_RATE": 0.3}, - "expected_impact": "-70% teacher API calls, minimal quality loss", - "confidence": "medium", - "rationale": "High tau indicates model has learned well; less feedback needed.", - }) - # Sort by priority priority_order = {"high": 0, "medium": 1, "low": 2} recommendations.sort(key=lambda r: priority_order.get(r.get("priority", "low"), 3)) diff --git a/scripts/benchmarks/refrag_bench.py b/scripts/benchmarks/refrag_bench.py index 534f7f76..943bc6a4 100644 --- a/scripts/benchmarks/refrag_bench.py +++ b/scripts/benchmarks/refrag_bench.py @@ -9,14 +9,11 @@ import asyncio import json import os -import sys import time from dataclasses import dataclass, field, asdict -from pathlib import Path from typing import Any, Dict, List import statistics -sys.path.insert(0, str(Path(__file__).parent.parent.parent)) # Load environment (optional) and fix Docker hostname try: diff --git a/scripts/benchmarks/router_bench.py b/scripts/benchmarks/router_bench.py deleted file mode 100644 index 0625f600..00000000 --- a/scripts/benchmarks/router_bench.py +++ /dev/null @@ -1,236 +0,0 @@ -#!/usr/bin/env python3 -""" -Router Benchmark for Context-Engine - -Measures tool selection accuracy, routing latency, and decision quality. -""" - -import argparse -import asyncio -import json -import os -import sys -import time -from dataclasses import dataclass, field, asdict -from pathlib import Path -from typing import Any, Dict, List, Optional -import statistics - -sys.path.insert(0, str(Path(__file__).parent.parent.parent)) - -# Load environment (optional) and fix Docker hostname -try: - from dotenv import load_dotenv # type: ignore - load_dotenv() -except Exception: - pass -if "qdrant:" in os.environ.get("QDRANT_URL", ""): - os.environ["QDRANT_URL"] = "http://localhost:6333" - -from scripts.benchmarks.common import ( - percentile, - create_report, - QueryResult as CommonQueryResult, - resolve_collection_auto, -) - -# Ensure correct collection is used (read from workspace state or env) -if not os.environ.get("COLLECTION_NAME"): - try: - from scripts.workspace_state import get_collection_name - os.environ["COLLECTION_NAME"] = get_collection_name() or "codebase" - except Exception: - os.environ["COLLECTION_NAME"] = "codebase" -else: - # If COLLECTION_NAME is set but empty/unindexed, pick a non-empty collection for benchmarks. - try: - os.environ["COLLECTION_NAME"] = resolve_collection_auto(os.environ.get("COLLECTION_NAME")) - except Exception: - pass - -print( - f"[bench] Using QDRANT_URL={os.environ.get('QDRANT_URL', '')} " - f"COLLECTION_NAME={os.environ.get('COLLECTION_NAME', '')}" -) - - -@dataclass -class RouterResult: - """Result from router evaluation.""" - query: str - expected_tool: str - selected_tool: str - confidence: float - latency_ms: float - correct: bool - - -@dataclass -class RouterReport: - """Router benchmark report.""" - name: str - total_queries: int - accuracy: float - avg_confidence: float - avg_latency_ms: float - p90_latency_ms: float - results: List[RouterResult] = field(default_factory=list) - - def to_dict(self) -> Dict[str, Any]: - base = { - "name": self.name, - "total_queries": self.total_queries, - "metrics": { - "accuracy": round(self.accuracy, 4), - "avg_confidence": round(self.avg_confidence, 4), - "avg_latency_ms": round(self.avg_latency_ms, 2), - "p90_latency_ms": round(self.p90_latency_ms, 2), - }, - "results": [asdict(r) for r in self.results], - } - # Also emit the unified BenchmarkReport shape for downstream tooling. - rep = create_report("router_bench", config={"name": self.name}) - for r in self.results: - rep.per_query.append( - CommonQueryResult( - query=r.query, - latency_ms=r.latency_ms, - metrics={ - "correct": 1.0 if r.correct else 0.0, - "confidence": float(r.confidence or 0.0), - }, - retrieved_paths=[], - metadata={ - "expected_tool": r.expected_tool, - "selected_tool": r.selected_tool, - }, - ) - ) - rep.compute_aggregates() - base["unified"] = rep.to_dict() - return base - - -ROUTER_TEST_CASES = [ - {"query": "find files that import embedder module", "expected": "search_importers_for"}, - {"query": "who calls the init_openlit function", "expected": "symbol_graph"}, - {"query": "explain how the hybrid search works", "expected": "context_answer"}, - {"query": "search for memory store implementation", "expected": "repo_search"}, - {"query": "find tests for the reranker", "expected": "search_tests_for"}, - {"query": "what configs exist for qdrant", "expected": "search_config_for"}, - {"query": "store a note about this finding", "expected": "memory_store"}, - {"query": "recall notes about authentication", "expected": "memory_find"}, -] - -TOOL_ALIASES: Dict[str, str] = { - # Search - "code_search": "repo_search", - "repo_search_compat": "repo_search", - # Answer - "context_answer_compat": "context_answer", - # Memory - "find": "memory_find", - "store": "memory_store", -} - - -def _canonical_tool_name(name: Any) -> str: - if not name: - return "unknown" - s = str(name).strip() - if not s: - return "unknown" - return TOOL_ALIASES.get(s, s) - - -async def run_router_benchmark(name: str = "default") -> RouterReport: - """Run router benchmark.""" - try: - from scripts.mcp_router import route_query - except ImportError: - try: - from scripts.mcp_router.router import route_query - except ImportError as e: - print(f"Import error: {e}") - return RouterReport(name=name, total_queries=0, accuracy=0, - avg_confidence=0, avg_latency_ms=0, p90_latency_ms=0) - - results: List[RouterResult] = [] - latencies: List[float] = [] - - for case in ROUTER_TEST_CASES: - query = case["query"] - expected = case["expected"] - - start = time.perf_counter() - try: - # Try to route the query - route_result = await route_query(query) - if isinstance(route_result, dict): - selected = route_result.get("tool", "unknown") - confidence = route_result.get("confidence", 0.0) - else: - selected = str(route_result) - confidence = 0.5 - except Exception as e: - selected = "error" - confidence = 0.0 - elapsed_ms = (time.perf_counter() - start) * 1000 - latencies.append(elapsed_ms) - - selected_c = _canonical_tool_name(selected) - expected_c = _canonical_tool_name(expected) - correct = selected_c == expected_c - results.append(RouterResult( - query=query, - expected_tool=expected, - selected_tool=selected_c, - confidence=confidence, - latency_ms=elapsed_ms, - correct=correct, - )) - status = "✓" if correct else "✗" - print(f" {status} {query[:35]:35} → {selected_c} (exp: {expected_c})") - - correct_count = sum(1 for r in results if r.correct) - - return RouterReport( - name=name, - total_queries=len(results), - accuracy=correct_count / len(results) if results else 0, - avg_confidence=statistics.mean(r.confidence for r in results) if results else 0, - avg_latency_ms=statistics.mean(latencies) if latencies else 0, - p90_latency_ms=percentile(latencies, 0.90), - results=results, - ) - - -def print_report(report: RouterReport): - print("\n" + "=" * 60) - print(f"ROUTER BENCHMARK: {report.name}") - print("=" * 60) - print(f"Queries: {report.total_queries}") - print(f"Accuracy: {report.accuracy:.1%}") - print(f"Avg Confidence: {report.avg_confidence:.4f}") - print(f"Avg Latency: {report.avg_latency_ms:.2f}ms") - print(f"P90 Latency: {report.p90_latency_ms:.2f}ms") - print("=" * 60) - - -def main(): - parser = argparse.ArgumentParser(description="Router Benchmark") - parser.add_argument("--name", default="default", help="Benchmark name") - parser.add_argument("--output", type=str, help="Output JSON file") - args = parser.parse_args() - - print(f"Running router benchmark: {args.name}") - report = asyncio.run(run_router_benchmark(name=args.name)) - print_report(report) - - if args.output: - with open(args.output, "w") as f: - json.dump(report.to_dict(), f, indent=2) - - -if __name__ == "__main__": - main() diff --git a/scripts/benchmarks/rrf_quality.py b/scripts/benchmarks/rrf_quality.py index f962392b..9122f676 100644 --- a/scripts/benchmarks/rrf_quality.py +++ b/scripts/benchmarks/rrf_quality.py @@ -20,9 +20,7 @@ from pathlib import Path from typing import Any, Dict, Iterable, List, Optional -# Add project root to path PROJECT_ROOT = Path(__file__).parent.parent.parent -sys.path.insert(0, str(PROJECT_ROOT)) # Load environment variables from .env try: @@ -136,8 +134,8 @@ def _match_expected_file(path: str, expected_files: Iterable[str]) -> Optional[s "expected_files": ["hybrid/ranking.py", "hybrid_search.py", "hybrid/__init__.py"], }, { - "query": "recursive reranker TRM learning", - "expected_files": ["rerank_recursive/core.py", "rerank_recursive/learning.py", "rerank_tools/recursive_reranker.py"], + "query": "relevance feedback graph recall", + "expected_files": ["relevance_feedback.py", "mcp_impl/search.py"], }, { "query": "context_answer grounding citations", @@ -161,7 +159,7 @@ def _match_expected_file(path: str, expected_files: Iterable[str]) -> Optional[s }, { "query": "MCP tool registration fastmcp", - "expected_files": ["mcp_indexer_server.py", "mcp_router/__init__.py"], + "expected_files": ["mcp_indexer_server.py", "mcp_impl/search.py"], }, { "query": "memory store find operations", diff --git a/scripts/benchmarks/run_all.py b/scripts/benchmarks/run_all.py index 7921c477..b731faf0 100644 --- a/scripts/benchmarks/run_all.py +++ b/scripts/benchmarks/run_all.py @@ -14,7 +14,6 @@ from pathlib import Path from typing import Any, Dict, List -sys.path.insert(0, str(Path(__file__).parent.parent.parent)) # Import unified metadata utilities from scripts.benchmarks.common import BenchmarkMetadata @@ -75,15 +74,6 @@ async def run_all_benchmarks(components: List[str]) -> Dict[str, Any]: except Exception as e: print(f" Eval harness failed: {e}") - if "trm" in components or "all" in components: - try: - from scripts.benchmarks.trm_bench import run_trm_benchmark - print("\n▶ Running TRM/Reranker Benchmark...") - report = await run_trm_benchmark(name="trm") - results["components"]["trm_reranker"] = report.to_dict() - except Exception as e: - print(f" TRM benchmark failed: {e}") - if "refrag" in components or "all" in components: try: from scripts.benchmarks.refrag_bench import run_refrag_benchmark @@ -102,15 +92,6 @@ async def run_all_benchmarks(components: List[str]) -> Dict[str, Any]: except Exception as e: print(f" Expansion benchmark failed: {e}") - if "router" in components or "all" in components: - try: - from scripts.benchmarks.router_bench import run_router_benchmark - print("\n▶ Running Router Benchmark...") - report = await run_router_benchmark(name="router") - results["components"]["router"] = report.to_dict() - except Exception as e: - print(f" Router benchmark failed: {e}") - if "rrf" in components or "all" in components: try: from scripts.benchmarks.rrf_quality import run_rrf_benchmark @@ -204,17 +185,6 @@ def generate_recommendations(components: Dict[str, Any]) -> List[Dict[str, Any]] "impact": f"Grounding rate {grounding:.0%} indicates insufficient context", }) - # Check TRM latency - if "trm_reranker" in components: - p90 = components["trm_reranker"].get("metrics", {}).get("p90_latency_ms", 0) - if p90 > 500: - recs.append({ - "priority": "medium", - "component": "trm_reranker", - "action": "Consider ONNX fallback or reduce candidate pool", - "impact": f"P90 latency {p90:.0f}ms may impact UX", - }) - # Check Router accuracy if "router" in components: acc = components["router"].get("metrics", {}).get("accuracy", 0) @@ -270,7 +240,7 @@ def main(): "--components", nargs="+", default=["all"], - choices=["all", "eval", "trm", "refrag", "expand", "router", "rrf", "grounding", "efficiency", "embedding"], + choices=["all", "eval", "refrag", "expand", "router", "rrf", "grounding", "efficiency", "embedding"], help="Components to benchmark", ) parser.add_argument("--output", type=str, help="Output JSON file") diff --git a/scripts/benchmarks/swe/runner.py b/scripts/benchmarks/swe/runner.py index 6edd39ee..baeb1898 100644 --- a/scripts/benchmarks/swe/runner.py +++ b/scripts/benchmarks/swe/runner.py @@ -22,7 +22,6 @@ - [x] Lexical hash vectors (for hybrid search) - [x] Git metadata (commit, author, date) - [ ] ReFRAG micro-chunks (enable with REFRAG_MODE=1) -- [ ] Pattern vectors (enable with INDEX_PATTERN_VECTORS=1) **Search Pipeline (via mcp_indexer_server.repo_search):** - [x] Hybrid search (dense + lexical RRF fusion) @@ -60,8 +59,6 @@ RERANK_IN_PROCESS=1 # Run reranker in-process (required for reliability) RERANKER_TOPN=50 # Number of candidates to rerank RERANKER_RETURN_M=20 # Number of results to return after rerank - RERANK_LEARNING=0 # Disable learning reranker (default: off for benchmarks) - RERANK_EVENTS_ENABLED=0 # Disable event logging (default: off for benchmarks) Hybrid Search Weights: HYBRID_RRF_K=30 # RRF constant (higher = more uniform) @@ -109,7 +106,6 @@ from __future__ import annotations import os -import sys from pathlib import Path # --------------------------------------------------------------------------- @@ -123,21 +119,16 @@ # Disable features that trigger collection recreation or add extra vectors: # - LEX_SPARSE_MODE: Requires sparse vectors, triggers recreation if missing -# - PATTERN_VECTORS: Requires pattern_vector, triggers recreation if missing # - REFRAG_MODE: Adds mini vector, causes schema mismatch with existing collections # - INDEX_MICRO_CHUNKS: Uses token-based micro chunking, different schema # These can destroy indexed data when ensure_collection is called during search! os.environ["LEX_SPARSE_MODE"] = "0" -os.environ["PATTERN_VECTORS"] = "0" os.environ["REFRAG_MODE"] = "0" os.environ["INDEX_MICRO_CHUNKS"] = "0" # Silence tokenizers parallelism warning os.environ["TOKENIZERS_PARALLELISM"] = "false" -# Ensure project root is in path -sys.path.insert(0, str(Path(__file__).parent.parent.parent.parent)) - import argparse import asyncio import json @@ -183,10 +174,6 @@ def _apply_swe_env_config( os.environ["RERANK_ENABLED"] = "1" if rerank_enabled else "0" os.environ.setdefault("RERANK_IN_PROCESS", "1") - # Disable learning reranker for reproducible benchmarks - os.environ.setdefault("RERANK_LEARNING", "0") - os.environ.setdefault("RERANK_EVENTS_ENABLED", "0") - # Set reranker model paths (relative to project root) _project_root = Path(__file__).parent.parent.parent.parent os.environ.setdefault( diff --git a/scripts/benchmarks/test_env_snapshot.py b/scripts/benchmarks/test_env_snapshot.py index 089dfe9d..a0566859 100644 --- a/scripts/benchmarks/test_env_snapshot.py +++ b/scripts/benchmarks/test_env_snapshot.py @@ -13,9 +13,7 @@ import tempfile from pathlib import Path -# Ensure project root is in path ROOT = Path(__file__).resolve().parent.parent.parent -sys.path.insert(0, str(ROOT)) from scripts.benchmarks.common import ( get_env_snapshot, diff --git a/scripts/benchmarks/trace_optimizer.py b/scripts/benchmarks/trace_optimizer.py index 075eb3e9..acb42f89 100644 --- a/scripts/benchmarks/trace_optimizer.py +++ b/scripts/benchmarks/trace_optimizer.py @@ -11,13 +11,10 @@ import asyncio import json import os -import sys from dataclasses import dataclass, field -from datetime import datetime, timedelta -from pathlib import Path +from datetime import datetime from typing import Any, Dict, List, Optional, Tuple -sys.path.insert(0, str(Path(__file__).parent.parent.parent)) # --------------------------------------------------------------------------- # Time Window Configuration diff --git a/scripts/benchmarks/trm_bench.py b/scripts/benchmarks/trm_bench.py deleted file mode 100644 index 3511e582..00000000 --- a/scripts/benchmarks/trm_bench.py +++ /dev/null @@ -1,213 +0,0 @@ -#!/usr/bin/env python3 -""" -TRM/Reranker Benchmark for Context-Engine - -Measures embedding quality, vector search performance, and reranking accuracy. -""" - -import argparse -import asyncio -import json -import os -import sys -import time -from dataclasses import dataclass, field, asdict -from pathlib import Path -from typing import Any, Dict, List, Optional -import statistics - -sys.path.insert(0, str(Path(__file__).parent.parent.parent)) - -from scripts.benchmarks.common import percentile, extract_result_paths, resolve_collection_auto - -# Ensure correct collection is used (read from workspace state or env) -if not os.environ.get("COLLECTION_NAME"): - try: - from scripts.workspace_state import get_collection_name - os.environ["COLLECTION_NAME"] = get_collection_name() or "codebase" - except Exception: - os.environ["COLLECTION_NAME"] = "codebase" -else: - # If COLLECTION_NAME is set but empty/unindexed, pick a non-empty collection for benchmarks. - try: - os.environ["COLLECTION_NAME"] = resolve_collection_auto(os.environ.get("COLLECTION_NAME")) - except Exception: - pass - -print( - f"[bench] Using QDRANT_URL={os.environ.get('QDRANT_URL', '')} " - f"COLLECTION_NAME={os.environ.get('COLLECTION_NAME', '')}" -) - - -@dataclass -class RerankerResult: - """Result from reranker evaluation.""" - query: str - baseline_order: List[str] - reranked_order: List[str] - onnx_order: Optional[List[str]] - kendall_tau: float - latency_ms: float - - -@dataclass -class TRMReport: - """TRM/Reranker benchmark report.""" - name: str - total_queries: int - avg_kendall_tau: float - avg_latency_ms: float - p90_latency_ms: float - embedding_dim: int - results: List[RerankerResult] = field(default_factory=list) - - def to_dict(self) -> Dict[str, Any]: - return { - "name": self.name, - "total_queries": self.total_queries, - "metrics": { - "kendall_tau": round(self.avg_kendall_tau, 4), - "avg_latency_ms": round(self.avg_latency_ms, 2), - "p90_latency_ms": round(self.p90_latency_ms, 2), - }, - "embedding_dim": self.embedding_dim, - "results": [asdict(r) for r in self.results], - } - - -def compute_kendall_tau(list1: List[str], list2: List[str]) -> float: - """Simplified Kendall's tau for ranking comparison.""" - if not list1 or not list2: - return 0.0 - - common = set(list1) & set(list2) - if len(common) < 2: - return 0.0 - - # Count concordant/discordant pairs - concordant = 0 - discordant = 0 - - items = list(common) - for i in range(len(items)): - for j in range(i + 1, len(items)): - a, b = items[i], items[j] - try: - pos1_a, pos1_b = list1.index(a), list1.index(b) - pos2_a, pos2_b = list2.index(a), list2.index(b) - - if (pos1_a < pos1_b) == (pos2_a < pos2_b): - concordant += 1 - else: - discordant += 1 - except ValueError: - continue - - total = concordant + discordant - if total == 0: - return 0.0 - return (concordant - discordant) / total - - -BENCHMARK_QUERIES = [ - "hybrid search RRF ranking algorithm", - "recursive reranker learning feedback", - "memory store vector embedding", - "openlit instrumentation tracing", - "workspace state persistence JSON", -] - - -async def run_trm_benchmark(name: str = "default") -> TRMReport: - """Run TRM/reranker benchmark.""" - try: - from scripts.embedder import get_embedding_model - from scripts.rerank_recursive import rerank_with_learning - from scripts.mcp_indexer_server import repo_search - except ImportError as e: - print(f"Import error: {e}") - return TRMReport(name=name, total_queries=0, avg_kendall_tau=0, - avg_latency_ms=0, p90_latency_ms=0, embedding_dim=0) - - # Get embedding dimension - try: - model = get_embedding_model() - embedding_dim = model.get_sentence_embedding_dimension() - except Exception: - embedding_dim = 384 # fallback - - results: List[RerankerResult] = [] - latencies: List[float] = [] - - for query in BENCHMARK_QUERIES: - print(f" Benchmarking: {query[:40]}...") - - # Get baseline results - try: - baseline_result = await repo_search(query=query, limit=10, rerank_enabled=False) - baseline_paths = extract_result_paths(baseline_result) - except Exception: - baseline_paths = [] - - # Get reranked results - start = time.perf_counter() - try: - reranked_result = await repo_search(query=query, limit=10, rerank_enabled=True) - reranked_paths = extract_result_paths(reranked_result) - except Exception: - reranked_paths = [] - elapsed_ms = (time.perf_counter() - start) * 1000 - latencies.append(elapsed_ms) - - tau = compute_kendall_tau(baseline_paths, reranked_paths) - - results.append(RerankerResult( - query=query, - baseline_order=baseline_paths[:5], - reranked_order=reranked_paths[:5], - onnx_order=None, - kendall_tau=tau, - latency_ms=elapsed_ms, - )) - print(f" τ={tau:.3f}, latency={elapsed_ms:.0f}ms") - - return TRMReport( - name=name, - total_queries=len(results), - avg_kendall_tau=statistics.mean(r.kendall_tau for r in results) if results else 0, - avg_latency_ms=statistics.mean(latencies) if latencies else 0, - p90_latency_ms=percentile(latencies, 0.90), - embedding_dim=embedding_dim, - results=results, - ) - -def print_report(report: TRMReport): - print("\n" + "=" * 60) - print(f"TRM/RERANKER BENCHMARK: {report.name}") - print("=" * 60) - print(f"Queries: {report.total_queries}") - print(f"Embedding Dim: {report.embedding_dim}") - print(f"Avg Kendall-τ: {report.avg_kendall_tau:.4f}") - print(f"Avg Latency: {report.avg_latency_ms:.2f}ms") - print(f"P90 Latency: {report.p90_latency_ms:.2f}ms") - print("=" * 60) - - -def main(): - parser = argparse.ArgumentParser(description="TRM/Reranker Benchmark") - parser.add_argument("--name", default="default", help="Benchmark name") - parser.add_argument("--output", type=str, help="Output JSON file") - args = parser.parse_args() - - print(f"Running TRM benchmark: {args.name}") - report = asyncio.run(run_trm_benchmark(name=args.name)) - print_report(report) - - if args.output: - with open(args.output, "w") as f: - json.dump(report.to_dict(), f, indent=2) - - -if __name__ == "__main__": - main() diff --git a/scripts/benchmarks/validation_loop.py b/scripts/benchmarks/validation_loop.py index b491bb2c..3a91af72 100644 --- a/scripts/benchmarks/validation_loop.py +++ b/scripts/benchmarks/validation_loop.py @@ -9,15 +9,12 @@ import asyncio import json import os -import sys from dataclasses import dataclass from datetime import datetime from pathlib import Path from typing import Any, Dict, List -# Add project root to path PROJECT_ROOT = Path(__file__).parent.parent.parent -sys.path.insert(0, str(PROJECT_ROOT)) # Load environment variables from .env try: diff --git a/scripts/codex_phase3_probe.py b/scripts/codex_phase3_probe.py new file mode 100644 index 00000000..628b09d5 --- /dev/null +++ b/scripts/codex_phase3_probe.py @@ -0,0 +1,2 @@ +MARK = 'v3' +# codex phase3 probe v3 diff --git a/scripts/collection_admin.py b/scripts/collection_admin.py index d970e941..58bde79f 100644 --- a/scripts/collection_admin.py +++ b/scripts/collection_admin.py @@ -1,30 +1,22 @@ +import logging import os import json import re import shutil import time -from pathlib import Path from datetime import datetime +from pathlib import Path from typing import Any, Dict, Optional, List -from scripts.auth_backend import mark_collection_deleted +logger = logging.getLogger(__name__) -try: - from qdrant_client import QdrantClient - from qdrant_client import models as qmodels -except Exception: - QdrantClient = None # type: ignore - qmodels = None # type: ignore +from scripts.auth_backend import mark_collection_deleted -try: - from scripts.qdrant_client_manager import pooled_qdrant_client -except Exception: - pooled_qdrant_client = None +from qdrant_client import QdrantClient +from qdrant_client import models as qmodels -try: - from scripts.workspace_state import get_collection_mappings -except Exception: - get_collection_mappings = None +from scripts.qdrant_client_manager import pooled_qdrant_client +from scripts.workspace_state import get_collection_mappings _SLUGGED_REPO_RE = re.compile(r"^.+-[0-9a-f]{16}(?:_old)?$") @@ -97,7 +89,6 @@ def _managed_upload_marker_path( slug_name: str, marker_root: Optional[Path] = None, ) -> Path: - # Marker is stored with per-repo metadata, not inside the repo workspace tree. base = marker_root or work_root return base / ".codebase" / "repos" / slug_name / _MARKER_NAME @@ -115,11 +106,12 @@ def _is_managed_upload_workspace_dir( return False if not _SLUGGED_REPO_RE.match(p.name or ""): return False - return _managed_upload_marker_path( + marker = _managed_upload_marker_path( work_root=work_root, marker_root=marker_root, slug_name=p.name, - ).exists() + ) + return marker.exists() except Exception: return False @@ -193,6 +185,7 @@ def delete_collection_everywhere( out: Dict[str, Any] = { "collection": name, "qdrant_deleted": False, + "qdrant_graph_deleted": False, "registry_marked_deleted": False, "deleted_state_files": 0, "deleted_managed_workspaces": 0, @@ -209,6 +202,14 @@ def delete_collection_everywhere( out["qdrant_deleted"] = True except Exception: out["qdrant_deleted"] = False + # Best-effort: also delete companion graph edges collection when present. + # This branch stores file-level edges in `_graph`. + if not name.endswith("_graph"): + try: + cli.delete_collection(collection_name=f"{name}_graph") + out["qdrant_graph_deleted"] = True + except Exception: + out["qdrant_graph_deleted"] = False except Exception: out["qdrant_deleted"] = False @@ -226,7 +227,7 @@ def delete_collection_everywhere( mappings = [] try: if get_collection_mappings is not None: - mappings = get_collection_mappings(search_root=str(codebase_root)) or [] + mappings = get_collection_mappings(search_root=str(work_root)) or [] except Exception: mappings = [] @@ -333,8 +334,6 @@ def _copy_client_timeout_seconds() -> Optional[float]: copied = False def _manual_copy_points() -> None: - if QdrantClient is None or qmodels is None: - raise RuntimeError("QdrantClient unavailable for manual collection copy") cli = QdrantClient(url=base_url, api_key=api_key or None, timeout=_copy_client_timeout_seconds()) try: if overwrite: @@ -359,8 +358,10 @@ def _manual_copy_points() -> None: vectors_config = None sparse_vectors_config = None + # Support vector-less collections (e.g. payload-only graph edge collections). if vectors_config is None: - raise RuntimeError(f"Cannot determine vectors config for source collection {src}") + vectors_config = {} + vectorless = isinstance(vectors_config, dict) and not vectors_config try: cli.create_collection( @@ -401,7 +402,7 @@ def _manual_copy_points() -> None: limit=batch_limit, offset=offset, with_payload=True, - with_vectors=True, + with_vectors=(not vectorless), ) except Exception as exc: raise RuntimeError(f"Failed to scroll points from {src}: {exc}") from exc @@ -414,7 +415,9 @@ def _manual_copy_points() -> None: point_id = getattr(record, "id", None) payload = getattr(record, "payload", None) vector = None - if hasattr(record, "vector") and getattr(record, "vector") is not None: + if vectorless: + vector = {} + elif hasattr(record, "vector") and getattr(record, "vector") is not None: vector = getattr(record, "vector") elif hasattr(record, "vectors") and getattr(record, "vectors") is not None: vector = getattr(record, "vectors") @@ -437,8 +440,6 @@ def _manual_copy_points() -> None: pass def _count_points(name: str) -> Optional[int]: - if QdrantClient is None: - return None cli = QdrantClient(url=base_url, api_key=api_key or None, timeout=_copy_client_timeout_seconds()) try: res = cli.count(collection_name=name, exact=True) @@ -477,4 +478,23 @@ def _count_points(name: str) -> Optional[int]: # The manual path guarantees the destination gets the exact same points/payloads/vectors. _manual_copy_points() + # Best-effort: copy the companion graph collection when copying a base collection. + # Graph edges are derived data and can be rebuilt, but copying avoids a cold-start window + # during staging cutovers where the clone has no graph. + if not src.endswith("_graph") and not dest.endswith("_graph"): + try: + copy_collection_qdrant( + source=f"{src}_graph", + target=f"{dest}_graph", + qdrant_url=base_url, + overwrite=overwrite, + ) + except Exception as exc: + logger.debug( + "Best-effort graph collection copy %s_graph -> %s_graph failed: %s", + src, + dest, + exc, + ) + return dest diff --git a/scripts/collection_health.py b/scripts/collection_health.py index 53ee089b..a4c72f3d 100644 --- a/scripts/collection_health.py +++ b/scripts/collection_health.py @@ -6,16 +6,9 @@ and triggers corrective actions (cache clear + reindex). """ import os -import sys -from pathlib import Path from typing import Optional, Dict, Any import logging -# Ensure project root is on sys.path -ROOT_DIR = Path(__file__).resolve().parent.parent -if str(ROOT_DIR) not in sys.path: - sys.path.insert(0, str(ROOT_DIR)) - from scripts.workspace_state import ( _read_cache, _write_cache, @@ -432,4 +425,3 @@ def main(): if __name__ == "__main__": main() - diff --git a/scripts/create_indexes.py b/scripts/create_indexes.py index 8f310aeb..9421b154 100644 --- a/scripts/create_indexes.py +++ b/scripts/create_indexes.py @@ -1,6 +1,5 @@ #!/usr/bin/env python3 import os -import sys from pathlib import Path from qdrant_client import QdrantClient, models @@ -8,21 +7,8 @@ QDRANT_URL = os.environ.get("QDRANT_URL", "http://qdrant:6333") from datetime import datetime ROOT_DIR = Path(__file__).resolve().parent.parent -if str(ROOT_DIR) not in sys.path: - sys.path.insert(0, str(ROOT_DIR)) -# Import critical functions first -try: - from scripts.workspace_state import get_collection_name, is_multi_repo_mode -except Exception: - get_collection_name = None # type: ignore - is_multi_repo_mode = None # type: ignore - -# Import other optional functions -try: - from scripts.workspace_state import log_activity -except Exception: - log_activity = None # type: ignore +from scripts.workspace_state import get_collection_name, is_multi_repo_mode, log_activity COLLECTION = os.environ.get("COLLECTION_NAME", "codebase") # Discover workspace path for state updates (allows subdir indexing) diff --git a/scripts/ctx.py b/scripts/ctx.py index 2c570a28..ee77f576 100755 --- a/scripts/ctx.py +++ b/scripts/ctx.py @@ -95,10 +95,7 @@ def _load_env_file(): _load_env_file() -try: - from scripts.mcp_router import call_tool_http # type: ignore -except ModuleNotFoundError: # pragma: no cover - local execution fallback - from mcp_router import call_tool_http # type: ignore +from scripts.mcp_http_client import call_tool_http # Configuration from environment MCP_URL = os.environ.get("MCP_INDEXER_URL", "http://localhost:8003/mcp") @@ -250,6 +247,12 @@ def parse_mcp_response(result: Dict[str, Any]) -> Optional[Dict[str, Any]]: # FastMCP typically wraps results in a content array res = result.get("result", {}) + structured = res.get("structuredContent") if isinstance(res, dict) else None + if isinstance(structured, dict): + structured_result = structured.get("result") + if isinstance(structured_result, dict): + return structured_result + content = res.get("content", []) # Some servers may return a dict directly (no content array) @@ -271,7 +274,10 @@ def parse_mcp_response(result: Dict[str, Any]) -> Optional[Dict[str, Any]]: return None try: - return json.loads(text) + parsed = json.loads(text) + if isinstance(parsed, dict) and isinstance(parsed.get("result"), dict): + return parsed["result"] + return parsed except json.JSONDecodeError: return {"raw": text} @@ -736,19 +742,14 @@ def _generate_plan(enhanced_prompt: str, context: str, note: str) -> str: from refrag_glm import GLMRefragClient # type: ignore client = GLMRefragClient() - response = client.client.chat.completions.create( - model=os.environ.get("GLM_MODEL", "glm-4.6"), - messages=[ - {"role": "system", "content": system_msg}, - {"role": "user", "content": user_msg}, - ], + plan = client.generate_with_soft_embeddings( + f"{system_msg}\n\n{user_msg}", max_tokens=200, + model=os.environ.get("GLM_MODEL", "glm-4.6"), temperature=0.3, stream=False, - ) - plan = ( - (response.choices[0].message.content if response and response.choices else "") - or "" + no_thinking=os.environ.get("CTX_GLM_DISABLE_THINKING", "1").strip().lower() + not in {"0", "false", "no", "off"}, ).strip() if not plan: # Fall through to llama.cpp path @@ -1030,6 +1031,7 @@ def fetch_context(query: str, **filters) -> Tuple[str, str]: params = { "query": query, "limit": filters.get("limit", DEFAULT_LIMIT), + "per_path": filters.get("per_path", DEFAULT_PER_PATH), "include_snippet": with_snippets, "context_lines": filters.get("context_lines", DEFAULT_CONTEXT_LINES), "collection": collection_name, @@ -1244,33 +1246,16 @@ def rewrite_prompt(original_prompt: str, context: str, note: str, max_tokens: Op "For questions: expand into related conceptual questions. For commands/instructions: provide general guidance about the task. " ) - # GLM API call - response = client.client.chat.completions.create( - model=os.environ.get("GLM_MODEL", "glm-4.6"), - messages=[ - {"role": "system", "content": system_msg}, - {"role": "user", "content": user_msg} - ], + enhanced = client.generate_with_soft_embeddings( + f"{system_msg}\n\n{user_msg}", max_tokens=int(max_tokens or DEFAULT_REWRITE_TOKENS), + model=os.environ.get("GLM_MODEL", "glm-4.6"), temperature=0.45, - stream=stream + stream=stream, + no_thinking=os.environ.get("CTX_GLM_DISABLE_THINKING", "1").strip().lower() + not in {"0", "false", "no", "off"}, ) - enhanced = "" - if stream: - # Streaming mode for GLM - for chunk in response: - if chunk.choices[0].delta.content: - token = chunk.choices[0].delta.content - sys.stdout.write(token) - sys.stdout.flush() - enhanced += token - sys.stdout.write("\n") - sys.stdout.flush() - else: - # Non-streaming mode for GLM - enhanced = response.choices[0].message.content - else: # Use local decoder (llama.cpp by default; Ollama supported when DECODER_URL points to /api/chat) meta_prompt = ( @@ -1586,6 +1571,8 @@ def main(): else: rewritten = rewrite_prompt(args.query, context_text, context_note, max_tokens=args.rewrite_max_tokens) output = sanitize_citations(rewritten.strip(), allowed_paths) + if args.with_context and context_text.strip(): + output = output.rstrip() + "\n\n---\nSupporting context:\n" + context_text.strip() if args.cmd: subprocess.run(args.cmd, input=output.encode("utf-8"), shell=True, check=False) diff --git a/scripts/health_check.py b/scripts/health_check.py index 32b9167d..3e2bf057 100644 --- a/scripts/health_check.py +++ b/scripts/health_check.py @@ -1,25 +1,11 @@ #!/usr/bin/env python3 import os import sys -from pathlib import Path from typing import Dict, Any from qdrant_client import QdrantClient, models -# Ensure /work (repo root) is on sys.path when run from /work/scripts -ROOT_DIR = Path(__file__).resolve().parent.parent -if str(ROOT_DIR) not in sys.path: - sys.path.insert(0, str(ROOT_DIR)) - -# Use embedder factory for Qwen3 support; fallback to direct fastembed -try: - from scripts.embedder import get_embedding_model, get_model_dimension - _EMBEDDER_FACTORY = True -except ImportError: - _EMBEDDER_FACTORY = False - from fastembed import TextEmbedding - - +from scripts.embedder import get_embedding_model, get_model_dimension from scripts.utils import sanitize_vector_name from scripts.auth_backend import ensure_collections, AuthDisabledError @@ -42,18 +28,14 @@ def assert_true(cond: bool, msg: str, *, critical: bool = False, failures: list[ def main(): qdrant_url = os.environ.get("QDRANT_URL", "http://localhost:6333") api_key = os.environ.get("QDRANT_API_KEY") - collection = os.environ.get("COLLECTION_NAME", "codebase") + collection = (os.environ.get("COLLECTION_NAME") or "codebase").strip() model_name = os.environ.get("EMBEDDING_MODEL", "BAAI/bge-base-en-v1.5") print(f"Health check -> {qdrant_url} collection={collection} model={model_name}") # Init embedding to derive dimension and test embedding - if _EMBEDDER_FACTORY: - model = get_embedding_model(model_name) - dim = get_model_dimension(model_name) - else: - model = TextEmbedding(model_name=model_name) - dim = len(next(model.embed(["health dim probe"]))) + model = get_embedding_model(model_name) + dim = get_model_dimension(model_name) vec_name_expect = sanitize_vector_name(model_name) client = QdrantClient(url=qdrant_url, api_key=api_key or None) @@ -82,7 +64,6 @@ def main(): print("No collections found - nothing to health check") return - # Check each collection for collection_name in collections: print(f"Checking collection: {collection_name}") @@ -92,10 +73,20 @@ def main(): if isinstance(cfg, dict): present_names = list(cfg.keys()) assert_true(len(present_names) >= 1, "Collection has at least one named vector") + has_expected_vector = vec_name_expect in present_names assert_true( - vec_name_expect in present_names, + has_expected_vector, f"Expected vector name present: {vec_name_expect} in {present_names}", ) + if not has_expected_vector: + failures.append( + f"Collection {collection_name} is missing expected vector {vec_name_expect}" + ) + print( + f"[WARN] Skipping vector query for {collection_name}; " + f"expected vector {vec_name_expect!r} not present" + ) + continue got_dim = cfg[vec_name_expect].size else: present_names = [""] @@ -172,4 +163,4 @@ def main(): if __name__ == "__main__": - main() \ No newline at end of file + main() diff --git a/scripts/hybrid/__init__.py b/scripts/hybrid/__init__.py index f893feb1..c37c2729 100644 --- a/scripts/hybrid/__init__.py +++ b/scripts/hybrid/__init__.py @@ -5,11 +5,4 @@ from scripts.hybrid import config, qdrant, embed, filters, expand, ranking from scripts.hybrid.config import QDRANT_URL """ -from scripts.hybrid import config -from scripts.hybrid import qdrant -from scripts.hybrid import embed -from scripts.hybrid import filters -from scripts.hybrid import expand -from scripts.hybrid import ranking - __all__ = ["config", "qdrant", "embed", "filters", "expand", "ranking"] diff --git a/scripts/hybrid/config.py b/scripts/hybrid/config.py index a1c8854a..0edd00f8 100644 --- a/scripts/hybrid/config.py +++ b/scripts/hybrid/config.py @@ -152,11 +152,9 @@ def _get_micro_defaults() -> tuple[int, int, int, int]: """ micro_enabled = os.environ.get("INDEX_MICRO_CHUNKS", "1").strip().lower() in {"1", "true", "yes", "on"} - try: - from scripts.refrag_glm import detect_glm_runtime - is_glm = detect_glm_runtime() - except ImportError: - is_glm = False + from scripts.refrag_glm import detect_glm_runtime + + is_glm = detect_glm_runtime() if is_glm: if micro_enabled: diff --git a/scripts/hybrid/embed.py b/scripts/hybrid/embed.py index 23b7ee82..e24d49a3 100644 --- a/scripts/hybrid/embed.py +++ b/scripts/hybrid/embed.py @@ -25,24 +25,14 @@ from pathlib import Path from typing import Any, List, Optional, TYPE_CHECKING -# --------------------------------------------------------------------------- -# Embedder factory setup -# --------------------------------------------------------------------------- -try: - from scripts.embedder import get_embedding_model as _get_embedding_model - _EMBEDDER_FACTORY = True -except ImportError: - _EMBEDDER_FACTORY = False - _get_embedding_model = None # type: ignore - -# Always try to import TextEmbedding for backward compatibility with tests -try: - from fastembed import TextEmbedding -except ImportError: - TextEmbedding = None # type: ignore +from scripts.embedder import get_embedding_model as _get_embedding_model + +_EMBEDDER_FACTORY = True + +TextEmbedding = None # Type alias for embedding model (TextEmbedding or compatible) -EmbeddingModel = Any if TextEmbedding is None else TextEmbedding +EmbeddingModel = Any # --------------------------------------------------------------------------- # Configuration constants @@ -52,12 +42,9 @@ # --------------------------------------------------------------------------- # Unified cache system # --------------------------------------------------------------------------- -try: - from scripts.cache_manager import get_embedding_cache - UNIFIED_CACHE_AVAILABLE = True -except ImportError: - UNIFIED_CACHE_AVAILABLE = False - get_embedding_cache = None # type: ignore +from scripts.cache_manager import get_embedding_cache + +UNIFIED_CACHE_AVAILABLE = True # Legacy cache fallback structures _EMBED_QUERY_CACHE: OrderedDict[tuple[str, str], List[float]] = OrderedDict() @@ -109,14 +96,12 @@ def get_embedding_model(model_name: Optional[str] = None) -> EmbeddingModel: if _EMBEDDER_FACTORY and _get_embedding_model is not None: return _get_embedding_model(model_name) - if TextEmbedding is None: - raise ImportError( - "No embedding backend available. Install fastembed or ensure " - "scripts.embedder is importable." - ) - name = model_name or MODEL_NAME - return TextEmbedding(model_name=name) + text_embedding_cls = TextEmbedding + if text_embedding_cls is None: + from fastembed import TextEmbedding as text_embedding_cls + + return text_embedding_cls(model_name=name) # --------------------------------------------------------------------------- @@ -169,11 +154,9 @@ def embed_queries_cached( name = os.environ.get("EMBEDDING_MODEL", MODEL_NAME) # Apply Qwen3 instruction prefix if enabled (queries only, not documents) - try: - from scripts.embedder import prefix_queries - sanitized = prefix_queries(sanitized, name) - except ImportError: - pass + from scripts.embedder import prefix_queries + + sanitized = prefix_queries(sanitized, name) cache = _get_embed_cache() diff --git a/scripts/hybrid/expand.py b/scripts/hybrid/expand.py index 20a43834..b787db2b 100644 --- a/scripts/hybrid/expand.py +++ b/scripts/hybrid/expand.py @@ -29,27 +29,25 @@ from typing import List, Dict, Any, TYPE_CHECKING from pathlib import Path +from scripts.path_scope import ( + normalize_under as _normalize_under_scope, + metadata_matches_under as _metadata_matches_under, +) + logger = logging.getLogger("hybrid_expand") # Import QdrantClient type for annotations if TYPE_CHECKING: from qdrant_client import QdrantClient -# Import semantic expansion functionality (optional) -try: - from scripts.semantic_expansion import ( - expand_queries_semantically, - expand_queries_with_prf, - get_expansion_stats, - clear_expansion_cache, - ) - SEMANTIC_EXPANSION_AVAILABLE = True -except ImportError: - SEMANTIC_EXPANSION_AVAILABLE = False - expand_queries_semantically = None - expand_queries_with_prf = None - get_expansion_stats = None - clear_expansion_cache = None +from scripts.semantic_expansion import ( + expand_queries_semantically, + expand_queries_with_prf, + get_expansion_stats, + clear_expansion_cache, +) + +SEMANTIC_EXPANSION_AVAILABLE = True # Feature flag for embedding-based dynamic expansion @@ -542,20 +540,8 @@ def expand_via_embeddings( except Exception: vec_name = None - def _norm_under(u: str | None) -> str | None: - if not u: - return None - u = str(u).strip().replace("\\", "/") - u = "/".join([p for p in u.split("/") if p]) - if not u: - return None - if u.startswith("/work/"): - return u - if not u.startswith("/"): - return "/work/" + u - return "/work/" + u.lstrip("/") - flt = None + eff_under = _normalize_under_scope(under) try: from qdrant_client import models @@ -567,15 +553,6 @@ def _norm_under(u: str | None) -> str | None: match=models.MatchValue(value=language), ) ) - if under: - eff_under = _norm_under(under) - if eff_under: - must.append( - models.FieldCondition( - key="metadata.path_prefix", - match=models.MatchValue(value=eff_under), - ) - ) if kind: must.append( models.FieldCondition( @@ -621,10 +598,11 @@ def _norm_under(u: str | None) -> str | None: # Search for soft matches (we want semantically similar docs, not exact matches) try: + initial_limit = 8 if not eff_under else max(32, int(max_terms) * 8) search_kwargs = { "collection_name": collection, "query_vector": (vec_name, query_vector) if vec_name else query_vector, - "limit": 8, # Get top 8 neighbors + "limit": initial_limit, # Over-fetch when `under` is set (we post-filter). "with_payload": True, "score_threshold": 0.3, # Lower threshold to get more diverse results } @@ -637,6 +615,17 @@ def _norm_under(u: str | None) -> str | None: if not results: return [] + if eff_under: + _scoped = [] + for hit in results: + payload = getattr(hit, "payload", None) or {} + md = payload.get("metadata") or {} + if _metadata_matches_under(md, eff_under): + _scoped.append(hit) + results = _scoped + if not results: + return [] + # Extract unique terms from neighbors extracted_terms: set[str] = set() query_tokens = set(combined_query.lower().split()) diff --git a/scripts/hybrid/qdrant.py b/scripts/hybrid/qdrant.py index ef0f936d..ed3b8e10 100644 --- a/scripts/hybrid/qdrant.py +++ b/scripts/hybrid/qdrant.py @@ -1,4 +1,6 @@ #!/usr/bin/env python3 +from __future__ import annotations + """ Qdrant client and query logic extracted from hybrid_search.py. @@ -25,16 +27,22 @@ import logging import threading import re -from typing import List, Dict, Any, Tuple +from typing import List, Dict, Any, Tuple, TYPE_CHECKING from pathlib import Path from concurrent.futures import ThreadPoolExecutor -# Core Qdrant imports -try: - from qdrant_client import QdrantClient, models -except ImportError: - QdrantClient = None # type: ignore - models = None # type: ignore +if TYPE_CHECKING: + from qdrant_client import QdrantClient, models as models +else: + QdrantClient = Any + + class _LazyQdrantModels: + def __getattr__(self, name: str) -> Any: + from qdrant_client import models as _models + + return getattr(_models, name) + + models = _LazyQdrantModels() logger = logging.getLogger("hybrid_qdrant") @@ -75,6 +83,9 @@ def _safe_float(val: Any, default: float) -> float: LEX_SPARSE_NAME, LEX_SPARSE_MODE, ) +from scripts.query_optimizer import optimize_query +from scripts.utils import lex_hash_vector_queries as _lex_hash_vector_queries +from scripts.utils import lex_sparse_vector_queries as _lex_sparse_vector_queries EF_SEARCH = _safe_int(os.environ.get("QDRANT_EF_SEARCH", "128"), 128) @@ -101,40 +112,13 @@ def _get_search_params(ef: int) -> models.SearchParams: # Connection pooling setup # --------------------------------------------------------------------------- -try: - from scripts.qdrant_client_manager import get_qdrant_client, return_qdrant_client, pooled_qdrant_client - _POOL_AVAILABLE = True -except ImportError: - _POOL_AVAILABLE = False - - def get_qdrant_client(url=None, api_key=None, force_new=False, use_pool=True): - """Fallback client creation when pooling is unavailable.""" - if QdrantClient is None: - raise ImportError( - "qdrant_client is not installed. Install with: pip install qdrant-client" - ) - return QdrantClient( - url=url or os.environ.get("QDRANT_URL", "http://localhost:6333"), - api_key=api_key or os.environ.get("QDRANT_API_KEY") - ) - - def return_qdrant_client(client): - """No-op when pooling is unavailable.""" - pass - - class pooled_qdrant_client: - """Fallback context manager when pooling is unavailable.""" - def __init__(self, url=None, api_key=None): - self.url = url - self.api_key = api_key - self.client = None - - def __enter__(self): - self.client = get_qdrant_client(self.url, self.api_key) - return self.client +from scripts.qdrant_client_manager import ( + get_qdrant_client, + return_qdrant_client, + pooled_qdrant_client, +) - def __exit__(self, exc_type, exc_val, exc_tb): - return_qdrant_client(self.client) +_POOL_AVAILABLE = True # --------------------------------------------------------------------------- @@ -259,8 +243,7 @@ def _ensure_collection(client, collection: str, dim: int, vec_name: str): """Cached wrapper for ensure_collection - only calls once per (endpoint, collection, vec_name) pair. IMPORTANT: This is called during SEARCH operations. We must NOT delete/recreate collections - that already exist with data. The ensure_collection in ingest_code can trigger recreation - when PATTERN_VECTORS=1 or LEX_SPARSE_MODE=1 if the collection lacks those vectors. + that already exist with data. Schema changes belong to the ingestion path, not search. For search, we only need to verify the collection exists - not modify its schema. """ @@ -276,12 +259,10 @@ def _ensure_collection(client, collection: str, dim: int, vec_name: str): _ENSURED_COLLECTIONS.add(cache_key) return - # Collection doesn't exist - only then call ensure_collection to create it - try: - from scripts.ingest_code import ensure_collection as _ensure_collection_raw - _ensure_collection_raw(client, collection, dim, vec_name) - except ImportError: - pass + # Collection doesn't exist - only then call the ingest Qdrant adapter to create it. + from scripts.ingest.qdrant import ensure_collection as _ensure_collection_raw + + _ensure_collection_raw(client, collection, dim, vec_name) try: _cache_collection_vectors(client, collection) @@ -403,51 +384,12 @@ def lex_hash_vector(phrases: List[str], dim: int | None = None) -> List[float]: """Generate dense lexical hash vector for query phrases.""" if dim is None: dim = LEX_VECTOR_DIM - try: - from scripts.utils import lex_hash_vector_queries as _lex_hash_vector_queries - return _lex_hash_vector_queries(phrases, dim) - except ImportError: - return _fallback_lex_hash_vector(phrases, dim) - - -def _fallback_lex_hash_vector(phrases: List[str], dim: int) -> List[float]: - """Fallback implementation when utils is unavailable.""" - import hashlib - vec = [0.0] * dim - for phrase in phrases: - for tok in _split_ident_lex(phrase): - h = int(hashlib.md5(tok.encode()).hexdigest(), 16) - idx = h % dim - vec[idx] += 1.0 - norm = sum(v * v for v in vec) ** 0.5 - if norm > 0: - vec = [v / norm for v in vec] - return vec + return _lex_hash_vector_queries(phrases, dim) def lex_sparse_vector(phrases: List[str]) -> Dict[str, Any]: """Generate sparse vector for query phrases (lossless exact matching).""" - try: - from scripts.utils import lex_sparse_vector_queries as _lex_sparse_vector_queries - return _lex_sparse_vector_queries(phrases) - except ImportError: - return _fallback_lex_sparse_vector(phrases) - - -def _fallback_lex_sparse_vector(phrases: List[str]) -> Dict[str, Any]: - """Fallback implementation when utils is unavailable.""" - import hashlib - indices = [] - values = [] - seen = set() - for phrase in phrases: - for tok in _split_ident_lex(phrase): - h = int(hashlib.md5(tok.encode()).hexdigest(), 16) % (2**31) - if h not in seen: - indices.append(h) - values.append(1.0) - seen.add(h) - return {"indices": indices, "values": values} + return _lex_sparse_vector_queries(phrases) # --------------------------------------------------------------------------- @@ -605,15 +547,12 @@ def dense_query( # Apply dynamic EF optimization if query text provided if query_text: try: - from scripts.query_optimizer import optimize_query result = optimize_query(query_text) # Only override EF when adaptive optimization is enabled if result.get("adaptive_enabled", False): ef = result["recommended_ef"] if os.environ.get("DEBUG_HYBRID_SEARCH"): logger.debug(f"Dynamic EF: {ef} (complexity={result['complexity']}, type={result['query_type']})") - except ImportError: - pass except Exception as e: if os.environ.get("DEBUG_HYBRID_SEARCH"): logger.debug(f"Query optimizer failed, using default EF: {e}") diff --git a/scripts/hybrid/ranking.py b/scripts/hybrid/ranking.py index 6a436272..8a96bc83 100644 --- a/scripts/hybrid/ranking.py +++ b/scripts/hybrid/ranking.py @@ -1,4 +1,6 @@ #!/usr/bin/env python3 +from __future__ import annotations + """ Ranking and scoring logic for hybrid search. @@ -21,7 +23,20 @@ import re import math import logging -from typing import List, Dict, Any, Tuple +from typing import List, Dict, Any, Tuple, TYPE_CHECKING + +if TYPE_CHECKING: + from qdrant_client import QdrantClient, models as models +else: + QdrantClient = Any + + class _LazyQdrantModels: + def __getattr__(self, name: str) -> Any: + from qdrant_client import models as _models + + return getattr(_models, name) + + models = _LazyQdrantModels() logger = logging.getLogger("hybrid_ranking") @@ -81,11 +96,8 @@ def _get_micro_defaults() -> Tuple[int, int, int, int]: Budget tokens floor is 5000 to ensure context_answer has enough context for quality answers. """ micro_enabled = os.environ.get("INDEX_MICRO_CHUNKS", "1").strip().lower() in {"1", "true", "yes", "on"} - try: - from scripts.refrag_glm import detect_glm_runtime - is_glm = detect_glm_runtime() - except ImportError: - is_glm = False + from scripts.refrag_glm import detect_glm_runtime + is_glm = detect_glm_runtime() if is_glm: if micro_enabled: return (24, 6, 8192, 32) @@ -658,12 +670,6 @@ def _get_symbol_extent( if cache_key in _SYMBOL_EXTENT_CACHE: return _SYMBOL_EXTENT_CACHE[cache_key] - # Lazy import to avoid circular dependencies - try: - from qdrant_client import QdrantClient, models - except ImportError: - return (0, 0) - if not collection: collection = os.environ.get("COLLECTION_NAME", "") if not collection: @@ -679,7 +685,9 @@ def _get_symbol_extent( timeout_s = float(os.environ.get("ADAPTIVE_SPAN_QDRANT_TIMEOUT", "1.0") or 1.0) except Exception: timeout_s = 1.0 - _SYMBOL_EXTENT_CLIENT = QdrantClient( + from qdrant_client import QdrantClient as _QdrantClient + + _SYMBOL_EXTENT_CLIENT = _QdrantClient( url=qdrant_url, api_key=os.environ.get("QDRANT_API_KEY"), timeout=timeout_s, diff --git a/scripts/hybrid_config.py b/scripts/hybrid_config.py deleted file mode 100644 index aebff322..00000000 --- a/scripts/hybrid_config.py +++ /dev/null @@ -1,3 +0,0 @@ -#!/usr/bin/env python3 -"""Shim for backward compatibility. See scripts/hybrid/config.py""" -from scripts.hybrid.config import * diff --git a/scripts/hybrid_embed.py b/scripts/hybrid_embed.py deleted file mode 100644 index eb331621..00000000 --- a/scripts/hybrid_embed.py +++ /dev/null @@ -1,3 +0,0 @@ -#!/usr/bin/env python3 -"""Shim for backward compatibility. See scripts/hybrid/embed.py""" -from scripts.hybrid.embed import * diff --git a/scripts/hybrid_expand.py b/scripts/hybrid_expand.py deleted file mode 100644 index d8131c0c..00000000 --- a/scripts/hybrid_expand.py +++ /dev/null @@ -1,3 +0,0 @@ -#!/usr/bin/env python3 -"""Shim for backward compatibility. See scripts/hybrid/expand.py""" -from scripts.hybrid.expand import * diff --git a/scripts/hybrid_filters.py b/scripts/hybrid_filters.py deleted file mode 100644 index 60ffca18..00000000 --- a/scripts/hybrid_filters.py +++ /dev/null @@ -1,3 +0,0 @@ -#!/usr/bin/env python3 -"""Shim for backward compatibility. See scripts/hybrid/filters.py""" -from scripts.hybrid.filters import * diff --git a/scripts/hybrid_qdrant.py b/scripts/hybrid_qdrant.py deleted file mode 100644 index 2498ef8e..00000000 --- a/scripts/hybrid_qdrant.py +++ /dev/null @@ -1,3 +0,0 @@ -#!/usr/bin/env python3 -"""Shim for backward compatibility. See scripts/hybrid/qdrant.py""" -from scripts.hybrid.qdrant import * diff --git a/scripts/hybrid_ranking.py b/scripts/hybrid_ranking.py deleted file mode 100644 index e20f7b32..00000000 --- a/scripts/hybrid_ranking.py +++ /dev/null @@ -1,3 +0,0 @@ -#!/usr/bin/env python3 -"""Shim for backward compatibility. See scripts/hybrid/ranking.py""" -from scripts.hybrid.ranking import * diff --git a/scripts/hybrid_search.py b/scripts/hybrid_search.py index 70b5b094..4755f0b8 100644 --- a/scripts/hybrid_search.py +++ b/scripts/hybrid_search.py @@ -4,12 +4,12 @@ This is the stable public entrypoint for the hybrid search subsystem. All internal logic has been refactored into smaller, focused modules: -- hybrid_config.py: Environment-based configuration and constants -- hybrid_qdrant.py: Qdrant client management, queries, and vector functions -- hybrid_embed.py: Embedding model factory and cached embedding -- hybrid_filters.py: File classification and query DSL parsing -- hybrid_ranking.py: RRF, scoring, diversification, and micro-span budgeting -- hybrid_expand.py: Query expansion (synonyms, semantic, LLM-assisted) +- hybrid/config.py: Environment-based configuration and constants +- hybrid/qdrant.py: Qdrant client management, queries, and vector functions +- hybrid/embed.py: Embedding model factory and cached embedding +- hybrid/filters.py: File classification and query DSL parsing +- hybrid/ranking.py: RRF, scoring, diversification, and micro-span budgeting +- hybrid/expand.py: Query expansion (synonyms, semantic, LLM-assisted) This façade: 1. Re-exports all public APIs for backwards compatibility @@ -19,32 +19,36 @@ from __future__ import annotations import os -import sys import argparse import re import json import math import logging import threading -from pathlib import Path from typing import List, Dict, Any, Tuple, TYPE_CHECKING from functools import lru_cache from concurrent.futures import ThreadPoolExecutor -# Ensure /work or repo root is in sys.path for scripts imports -_ROOT_DIR = Path(__file__).resolve().parent.parent -if str(_ROOT_DIR) not in sys.path: - sys.path.insert(0, str(_ROOT_DIR)) - # --------------------------------------------------------------------------- -# Core Qdrant imports +# Lazy Qdrant model namespace # --------------------------------------------------------------------------- -from qdrant_client import QdrantClient, models +if TYPE_CHECKING: + from qdrant_client import QdrantClient, models as models +else: + QdrantClient = Any + + class _LazyQdrantModels: + def __getattr__(self, name: str) -> Any: + from qdrant_client import models as _models + + return getattr(_models, name) + + models = _LazyQdrantModels() # --------------------------------------------------------------------------- # Re-exports from hybrid_config # --------------------------------------------------------------------------- -from scripts.hybrid_config import ( +from scripts.hybrid.config import ( # Helper functions _safe_int, _safe_float, @@ -109,7 +113,7 @@ # --------------------------------------------------------------------------- # Re-exports from hybrid_qdrant # --------------------------------------------------------------------------- -from scripts.hybrid_qdrant import ( +from scripts.hybrid.qdrant import ( # Pool availability _POOL_AVAILABLE, # Connection pooling @@ -142,7 +146,7 @@ # --------------------------------------------------------------------------- # Re-exports from hybrid_embed # --------------------------------------------------------------------------- -from scripts.hybrid_embed import ( +from scripts.hybrid.embed import ( # Embedder factory _EMBEDDER_FACTORY, EmbeddingModel, @@ -156,34 +160,21 @@ UNIFIED_CACHE_AVAILABLE, ) -# Import unified cache objects from cache_manager when available -if UNIFIED_CACHE_AVAILABLE: - try: - from scripts.cache_manager import get_search_cache, get_embedding_cache, get_expansion_cache - _EMBED_CACHE = get_embedding_cache() - _RESULTS_CACHE = get_search_cache() - _EXPANSION_CACHE = get_expansion_cache() - except ImportError: - _EMBED_CACHE = None - _RESULTS_CACHE = {} - _EXPANSION_CACHE = None -else: - _EMBED_CACHE = None - _RESULTS_CACHE = {} - _EXPANSION_CACHE = None - -# Lightweight local fallback cache for deterministic test hits -try: - from collections import OrderedDict as _OD -except Exception: - _OD = dict # pragma: no cover +from scripts.cache_manager import get_search_cache, get_embedding_cache, get_expansion_cache + +_EMBED_CACHE = get_embedding_cache() +_RESULTS_CACHE = get_search_cache() +_EXPANSION_CACHE = get_expansion_cache() + +from collections import OrderedDict as _OD + _RESULTS_CACHE_OD = _OD() _RESULTS_LOCK = threading.RLock() # --------------------------------------------------------------------------- # Re-exports from hybrid_filters # --------------------------------------------------------------------------- -from scripts.hybrid_filters import ( +from scripts.hybrid.filters import ( # File patterns CORE_FILE_PATTERNS, NON_CORE_PATTERNS, @@ -206,7 +197,7 @@ # --------------------------------------------------------------------------- # Re-exports from hybrid_ranking # --------------------------------------------------------------------------- -from scripts.hybrid_ranking import ( +from scripts.hybrid.ranking import ( # RRF rrf, _scale_rrf_k, @@ -238,7 +229,7 @@ # --------------------------------------------------------------------------- # Re-exports from hybrid_expand # --------------------------------------------------------------------------- -from scripts.hybrid_expand import ( +from scripts.hybrid.expand import ( # Synonyms CODE_SYNONYMS, # Expansion functions @@ -253,7 +244,7 @@ # Conditionally re-export semantic expansion functions if SEMANTIC_EXPANSION_AVAILABLE: - from scripts.hybrid_expand import ( + from scripts.hybrid.expand import ( expand_queries_semantically, expand_queries_with_prf, get_expansion_stats, @@ -268,34 +259,27 @@ # --------------------------------------------------------------------------- # Additional imports for backward compatibility # --------------------------------------------------------------------------- -try: - from fastembed import TextEmbedding -except ImportError: - TextEmbedding = None # type: ignore +TextEmbedding = None # Tests may monkeypatch this; production imports lazily if needed. -try: - from scripts.embedder import get_embedding_model as _get_embedding_model -except ImportError: - _get_embedding_model = None +from scripts.embedder import get_embedding_model as _get_embedding_model # Import request deduplication system -try: - from scripts.deduplication import get_deduplicator, is_duplicate_request - DEDUPLICATION_AVAILABLE = True -except ImportError: - DEDUPLICATION_AVAILABLE = False +from scripts.deduplication import get_deduplicator, is_duplicate_request + +DEDUPLICATION_AVAILABLE = True # Import query optimizer for dynamic EF tuning -try: - from scripts.query_optimizer import get_query_optimizer, optimize_query - QUERY_OPTIMIZER_AVAILABLE = True -except ImportError: - QUERY_OPTIMIZER_AVAILABLE = False +from scripts.query_optimizer import get_query_optimizer, optimize_query + +QUERY_OPTIMIZER_AVAILABLE = True # Import ingest helpers from scripts.utils import sanitize_vector_name as _sanitize_vector_name -from scripts.ingest_code import ensure_collection as _ensure_collection_raw -from scripts.ingest_code import project_mini as _project_mini +from scripts.ingest.vectors import project_mini as _project_mini +from scripts.path_scope import ( + normalize_under as _normalize_under_scope, + metadata_matches_under as _metadata_matches_under, +) # --------------------------------------------------------------------------- # Module logger @@ -309,6 +293,31 @@ _FILTER_CACHE_LOCK = threading.Lock() _FILTER_CACHE_MAX = 256 +def _compute_fname_boost(query: str, md: dict, boost_factor: float) -> float: + """Compute filename relevance boost based on query token matches in file path. + + Args: + query: The search query text + md: Metadata dict containing at least a 'path' key + boost_factor: Maximum boost multiplier + + Returns: + Boost score (0.0 if no match) + """ + import re as _re + path = str(md.get("path") or "").lower() + q = query.lower() + q_toks = {t for t in _re.findall(r"[a-z0-9_]{3,}", q) if len(t) >= 3} + if not q_toks: + return 0.0 + fname = path.rsplit("/", 1)[-1] if "/" in path else path + fname_base = _re.sub(r"\.[^.]+$", "", fname) + fname_toks = {t for t in _re.split(r"[_\-.]", fname_base) if t and len(t) >= 3} + match_count = len(q_toks & fname_toks) + if match_count >= 2: + return float(boost_factor) * match_count + return 0.0 + # Cached regex pattern compilation @lru_cache(maxsize=128) def _compile_regex(pattern: str, flags: int = 0): @@ -421,13 +430,64 @@ def _generate_code_query_variants(query: str) -> List[str]: return result[:5] # Max 5 variants to balance coverage vs compute +def _shape_dense_points( + ranked_points: List[Any], + *, + limit: int, + per_path: int | None = 1, + under: str | None = None, +) -> List[Dict[str, Any]]: + eff_under = _normalize_under_scope(under) + eff_per_path = int(per_path or 0) + + results: List[Dict[str, Any]] = [] + path_counts: dict[str, int] = {} + for p in ranked_points: + payload = p.payload or {} + md = payload.get("metadata") or {} + if eff_under and not _metadata_matches_under(md, eff_under): + continue + + # Prefer host_path when available (consistent with hybrid search). + path = md.get("host_path") or payload.get("path") or md.get("path") or "" + if eff_per_path > 0: + current = path_counts.get(path, 0) + if current >= eff_per_path: + continue + else: + current = 0 + + results.append( + { + "score": float(getattr(p, "score", 0) or 0), + "path": path, + "symbol": payload.get("symbol") or md.get("symbol") or "", + "start_line": int(md.get("start_line") or 0), + "end_line": int(md.get("end_line") or 0), + "code_id": payload.get("code_id") or payload.get("_id") or "", + "doc_id": payload.get("code_id") or payload.get("_id") or "", + "payload": payload, + } + ) + if eff_per_path > 0: + path_counts[path] = current + 1 + if len(results) >= int(limit): + break + + return results + + def run_pure_dense_search( query: str, limit: int = 10, + per_path: int | None = 1, model: Any = None, collection: str | None = None, language: str | None = None, under: str | None = None, + kind: str | None = None, + symbol: str | None = None, + ext: str | None = None, repo: str | list[str] | None = None, ) -> List[Dict[str, Any]]: """Pure dense search - single query embedding, single vector search. @@ -437,45 +497,51 @@ def run_pure_dense_search( Args: query: Natural language query limit: Max results to return + per_path: Optional max results per file path; <= 0 disables the cap model: Embedding model (will load default if None) collection: Qdrant collection name language: Optional language filter - under: Optional path prefix filter + under: Optional recursive workspace subtree filter + kind: Optional kind filter (exact match) + symbol: Optional symbol filter (exact match) + ext: Optional file extension filter (without dot) repo: Optional repo filter Returns: List of search results with raw cosine similarity scores """ - from scripts.hybrid_qdrant import get_qdrant_client, return_qdrant_client, dense_query + from scripts.hybrid.qdrant import get_qdrant_client, return_qdrant_client, dense_query from scripts.utils import sanitize_vector_name from qdrant_client import models # Get model if model is None: model_name = os.environ.get("EMBEDDING_MODEL", "BAAI/bge-base-en-v1.5") - try: - from scripts.embedder import get_embedding_model - model = get_embedding_model(model_name) - except ImportError: - from fastembed import TextEmbedding - model = TextEmbedding(model_name=model_name) + from scripts.embedder import get_embedding_model + model = get_embedding_model(model_name) else: model_name = getattr(model, "model_name", os.environ.get("EMBEDDING_MODEL", "BAAI/bge-base-en-v1.5")) vec_name = sanitize_vector_name(model_name) coll = collection or _collection() - # Build filter + # Build server-side filter (exclude `under` here; recursive under is post-filtered) must = [] if language: must.append(models.FieldCondition(key="metadata.language", match=models.MatchValue(value=language))) - if under: - must.append(models.FieldCondition(key="metadata.path_prefix", match=models.MatchValue(value=under))) if repo and repo != "*": if isinstance(repo, list): must.append(models.FieldCondition(key="metadata.repo", match=models.MatchAny(any=repo))) else: must.append(models.FieldCondition(key="metadata.repo", match=models.MatchValue(value=repo))) + if kind: + must.append(models.FieldCondition(key="metadata.kind", match=models.MatchValue(value=kind))) + if symbol: + must.append(models.FieldCondition(key="metadata.symbol", match=models.MatchValue(value=symbol))) + if ext: + ext_clean = str(ext).lower().lstrip(".") + if ext_clean: + must.append(models.FieldCondition(key="metadata.ext", match=models.MatchValue(value=ext_clean))) flt = models.Filter(must=must) if must else None # Single query embedding - no variants, no expansion @@ -500,30 +566,22 @@ def run_pure_dense_search( ) try: - # Single dense query - no pooling, no re-scoring - ranked_points = dense_query(client, vec_name, vec_list, flt, limit, coll, query_text=query) - - # Build output - results = [] - for p in ranked_points: - payload = p.payload or {} - md = payload.get("metadata") or {} - - # Prefer host_path when available (consistent with hybrid search) - _path = md.get("host_path") or payload.get("path") or md.get("path") or "" - - results.append({ - "score": float(getattr(p, "score", 0) or 0), - "path": _path, - "symbol": payload.get("symbol") or md.get("symbol") or "", - "start_line": int(md.get("start_line") or 0), - "end_line": int(md.get("end_line") or 0), - "code_id": payload.get("code_id") or payload.get("_id") or "", - "doc_id": payload.get("code_id") or payload.get("_id") or "", - "payload": payload, - }) - - return results + # Single dense query - no pooling, no re-scoring. + # When `under` or `per_path` is set, we may need to over-fetch so post-filters + # can still fill up to `limit` results. + eff_under = _normalize_under_scope(under) + fetch_limit = int(limit) + eff_per_path = int(per_path or 0) + if eff_under or eff_per_path > 0: + fetch_limit = min(max(fetch_limit * 4, fetch_limit + 16), 2000) + ranked_points = dense_query(client, vec_name, vec_list, flt, fetch_limit, coll, query_text=query) + + return _shape_dense_points( + ranked_points, + limit=limit, + per_path=per_path, + under=under, + ) finally: return_qdrant_client(client) @@ -532,7 +590,7 @@ def run_pure_dense_search( # --------------------------------------------------------------------------- # Backward compatibility: _embed_queries_cached alias # --------------------------------------------------------------------------- -# The function is now in hybrid_embed.py as embed_queries_cached +# The function is now in hybrid/embed.py as embed_queries_cached # Keep the underscore-prefixed alias for any legacy callers @@ -612,7 +670,11 @@ def _dt(label: str): elif _EMBEDDER_FACTORY: _model = _get_embedding_model(model_name) else: - _model = TextEmbedding(model_name=model_name) + text_embedding_cls = TextEmbedding + if text_embedding_cls is None: + from fastembed import TextEmbedding as text_embedding_cls + + _model = text_embedding_cls(model_name=model_name) vec_name = _sanitize_vector_name(model_name) # Parse Query DSL and merge with explicit args @@ -690,21 +752,8 @@ def _normalize_globs(globs: list[str]) -> list[str]: eff_path_globs_norm = _normalize_globs(eff_path_globs) eff_not_globs_norm = _normalize_globs(eff_not_globs) - # Normalize under - def _norm_under(u: str | None) -> str | None: - if not u: - return None - u = str(u).strip().replace("\\", "/") - u = "/".join([p for p in u.split("/") if p]) - if not u: - return None - if not u.startswith("/"): - v = "/work/" + u - else: - v = "/work/" + u.lstrip("/") if not u.startswith("/work/") else u - return v - - eff_under = _norm_under(eff_under) + # Normalize under as a user-facing recursive subtree scope. + eff_under = _normalize_under_scope(eff_under) # Expansion knobs that affect query construction/results (must be part of cache key) try: @@ -810,12 +859,8 @@ def _norm_under(u: str | None) -> str | None: key="metadata.repo", match=models.MatchValue(value=eff_repo) ) ) - if eff_under: - must.append( - models.FieldCondition( - key="metadata.path_prefix", match=models.MatchValue(value=eff_under) - ) - ) + # NOTE: `under` is recursive and user-facing; we enforce it in client-side + # filtering via normalized metadata paths instead of exact path_prefix equality. if eff_kind: must.append( models.FieldCondition( @@ -1726,8 +1771,7 @@ def _scaled_rrf(rank: int) -> float: # Filename boost: production-grade matching (handles snake/camel/kebab, acronyms, etc.) if FNAME_BOOST > 0.0 and path: try: - from scripts.rerank_recursive.utils import _compute_fname_boost as _compute_fname_boost # type: ignore - fname_boost = float(_compute_fname_boost(_base_query, md, float(FNAME_BOOST))) + fname_boost = _compute_fname_boost(_base_query, md, float(FNAME_BOOST)) if fname_boost > 0: rec["fname"] += fname_boost rec["s"] += fname_boost @@ -2105,7 +2149,7 @@ def _match_glob(pat: str, path: str) -> bool: return _fnm.fnmatchcase(path, pat) return _fnm.fnmatchcase(path.lower(), pat.lower()) - if eff_not or eff_path_regex or eff_ext or eff_path_globs or eff_not_globs: + if eff_under or eff_not or eff_path_regex or eff_ext or eff_path_globs or eff_not_globs: def _pass_filters(m: Dict[str, Any]) -> bool: md = (m["pt"].payload or {}).get("metadata") or {} @@ -2118,6 +2162,8 @@ def _pass_filters(m: Dict[str, Any]) -> bool: nn = eff_not if case_sensitive else eff_not.lower() if nn in p_for_sub or nn in pp_for_sub: return False + if eff_under and not _metadata_matches_under(md, eff_under): + return False if eff_not_globs_norm and any(_match_glob(g, path) or _match_glob(g, rel) for g in eff_not_globs_norm): return False if eff_ext: @@ -2341,6 +2387,8 @@ def _extract_entity_key(m: Dict[str, Any]) -> str: _imports = md.get("imports") or [] _calls = md.get("calls") or [] _symp = md.get("symbol_path") or md.get("symbol") or "" + _kind = str(md.get("kind") or "") + _repo_name = str(md.get("repo") or "") _pp = str(md.get("path_prefix") or "") _path = str(md.get("path") or "") _related_set = set() @@ -2439,6 +2487,18 @@ def _resolve(seg: str) -> list[str]: _tags = _payload.get("tags") if _tags is None: _tags = _metadata.get("tags") + _file_hash = str( + _payload.get("file_hash") + or _metadata.get("file_hash") + or _payload.get("content_hash") + or _metadata.get("content_hash") + or "" + ) + _symbol_content_hash = str( + _payload.get("symbol_content_hash") + or _metadata.get("symbol_content_hash") + or "" + ) # Skip memory-like points without a real file path if not _path or not _path.strip(): if os.environ.get("DEBUG_HYBRID_FILTER"): @@ -2505,6 +2565,8 @@ def _resolve(seg: str) -> list[str]: "host_path": _host, "container_path": _cont, "symbol": _symp, + "kind": _kind, + "repo": _repo_name, "start_line": start_line, "end_line": end_line, "components": comp, @@ -2515,6 +2577,8 @@ def _resolve(seg: str) -> list[str]: "text": _text, "pseudo": _pseudo, "tags": _tags, + "file_hash": _file_hash, + "symbol_content_hash": _symbol_content_hash, } if why is not None: item["why"] = why diff --git a/scripts/indexing_admin.py b/scripts/indexing_admin.py index f3bb69d8..ac93b6f6 100644 --- a/scripts/indexing_admin.py +++ b/scripts/indexing_admin.py @@ -10,70 +10,36 @@ from datetime import datetime, timezone from typing import Any, Dict, List, Optional, Set, Tuple -try: - from qdrant_client import QdrantClient -except Exception: - QdrantClient = None # type: ignore - -try: - from scripts.embedder import get_model_dimension -except Exception: - get_model_dimension = None # type: ignore - -try: - from scripts.collection_admin import copy_collection_qdrant -except Exception: - copy_collection_qdrant = None # type: ignore - -try: - from scripts.ingest_code import ( - ensure_collection_and_indexes_once, - ensure_payload_indexes, - _sanitize_vector_name, - MINI_VECTOR_NAME as _MINI_VECTOR_NAME, - LEX_SPARSE_NAME as _LEX_SPARSE_NAME, - ) -except Exception: - ensure_collection_and_indexes_once = None # type: ignore - ensure_payload_indexes = None # type: ignore - _sanitize_vector_name = None # type: ignore - _MINI_VECTOR_NAME = os.environ.get("MINI_VECTOR_NAME", "mini") - _LEX_SPARSE_NAME = os.environ.get("LEX_SPARSE_NAME", "lex_sparse") - -try: - from scripts.workspace_state import ( - get_collection_mappings, - get_workspace_state, - update_workspace_state, - update_indexing_status, - get_indexing_config_snapshot, - compute_indexing_config_hash, - is_staging_enabled, - set_staging_state, - update_staging_status, - clear_staging_collection, - activate_staging_collection, - promote_pending_indexing_config, - persist_indexing_config, - ) -except Exception: - get_collection_mappings = None # type: ignore - get_workspace_state = None # type: ignore - update_workspace_state = None # type: ignore - update_indexing_status = None # type: ignore - get_indexing_config_snapshot = None # type: ignore - compute_indexing_config_hash = None # type: ignore - is_staging_enabled = None # type: ignore - set_staging_state = None # type: ignore - update_staging_status = None # type: ignore - clear_staging_collection = None # type: ignore - activate_staging_collection = None # type: ignore - promote_pending_indexing_config = None # type: ignore - persist_indexing_config = None # type: ignore +from qdrant_client import QdrantClient + +from scripts.embedder import get_model_dimension +from scripts.collection_admin import copy_collection_qdrant +from scripts.ingest_code import ( + ensure_collection_and_indexes_once, + ensure_payload_indexes, + _sanitize_vector_name, + MINI_VECTOR_NAME as _MINI_VECTOR_NAME, + LEX_SPARSE_NAME as _LEX_SPARSE_NAME, +) +from scripts.workspace_state import ( + get_collection_mappings, + get_workspace_state, + update_workspace_state, + update_indexing_status, + get_indexing_config_snapshot, + compute_indexing_config_hash, + is_staging_enabled, + set_staging_state, + update_staging_status, + clear_staging_collection, + activate_staging_collection, + promote_pending_indexing_config, + persist_indexing_config, +) def _staging_enabled() -> bool: - return bool(is_staging_enabled() if callable(is_staging_enabled) else False) + return bool(is_staging_enabled()) def _workspace_base_dir() -> Path: @@ -141,7 +107,7 @@ def _copy_repo_state_for_clone( def _probe_collection_schema(collection: str) -> Optional[Dict[str, Any]]: - if not collection or QdrantClient is None: + if not collection: return None cached = _COLLECTION_SCHEMA_CACHE.get(collection) if cached: @@ -916,8 +882,6 @@ def build_admin_collections_view(*, collections: Any, work_dir: str) -> List[Dic def delete_collection_qdrant(*, qdrant_url: str, api_key: Optional[str], collection: str) -> None: - if QdrantClient is None: - return name = (collection or "").strip() if not name: return @@ -927,6 +891,17 @@ def delete_collection_qdrant(*, qdrant_url: str, api_key: Optional[str], collect return try: cli.delete_collection(collection_name=name) + # Best-effort: also delete companion graph edges collection when present. + if not name.endswith("_graph"): + try: + cli.delete_collection(collection_name=f"{name}_graph") + except Exception as exc: + try: + print( + f"[indexing_admin] best-effort graph collection delete failed for {name}_graph: {exc}" + ) + except Exception: + pass except Exception: pass finally: @@ -937,8 +912,6 @@ def delete_collection_qdrant(*, qdrant_url: str, api_key: Optional[str], collect def recreate_collection_qdrant(*, qdrant_url: str, api_key: Optional[str], collection: str) -> None: - if QdrantClient is None: - return name = (collection or "").strip() if not name: return @@ -951,6 +924,17 @@ def recreate_collection_qdrant(*, qdrant_url: str, api_key: Optional[str], colle cli.delete_collection(collection_name=name) except Exception as delete_error: raise RuntimeError(f"Failed to delete existing collection '{name}' in Qdrant: {delete_error}") from delete_error + # Best-effort: also delete companion graph edges collection when present. + if not name.endswith("_graph"): + try: + cli.delete_collection(collection_name=f"{name}_graph") + except Exception as exc: + try: + print( + f"[indexing_admin] best-effort graph collection delete failed for {name}_graph: {exc}" + ) + except Exception: + pass finally: try: cli.close() @@ -984,12 +968,9 @@ def spawn_ingest_code( env.pop(k, None) else: env[str(k)] = str(v) - # When we provide env overrides for a run (e.g. staging rebuild), we also want to - # force ingest_code to honor the explicit COLLECTION_NAME instead of routing based - # on per-repo state/serving_collection in multi-repo mode. - # CTXCE_FORCE_COLLECTION_NAME is only used for these subprocess runs; normal watcher - # and indexer flows do not set it. - env["CTXCE_FORCE_COLLECTION_NAME"] = "1" # Force ingest_code to use COLLECTION_NAME for staging/pending env overrides + # For admin-triggered subprocess runs (recreate/reindex/staging), force ingest_code to + # honor explicit COLLECTION_NAME and avoid multi-repo enumeration. + env["CTXCE_FORCE_COLLECTION_NAME"] = "1" env["COLLECTION_NAME"] = collection env["WATCH_ROOT"] = work_dir env["WORKSPACE_PATH"] = work_dir @@ -1022,11 +1003,10 @@ def spawn_ingest_code( def _determine_embedding_dim(model_name: str) -> int: - if get_model_dimension: - try: - return int(get_model_dimension(model_name)) - except Exception: - pass + try: + return int(get_model_dimension(model_name)) + except Exception: + pass try: from fastembed import TextEmbedding # type: ignore @@ -1037,16 +1017,13 @@ def _determine_embedding_dim(model_name: str) -> int: def _normalize_cloned_collection_schema(*, collection_name: str, qdrant_url: str) -> None: - if QdrantClient is None: - return vector_name = None model_name = os.environ.get("EMBEDDING_MODEL", "BAAI/bge-base-en-v1.5") dim = _determine_embedding_dim(model_name) - if _sanitize_vector_name is not None: - try: - vector_name = _sanitize_vector_name(model_name) - except Exception: - vector_name = None + try: + vector_name = _sanitize_vector_name(model_name) + except Exception: + vector_name = None try: client = QdrantClient(url=qdrant_url, api_key=os.environ.get("QDRANT_API_KEY") or None) except Exception: @@ -1073,8 +1050,6 @@ def _normalize_cloned_collection_schema(*, collection_name: str, qdrant_url: str def _get_collection_point_count(*, collection_name: str, qdrant_url: str) -> Optional[int]: - if QdrantClient is None: - return None try: client = QdrantClient(url=qdrant_url, api_key=os.environ.get("QDRANT_API_KEY") or None) except Exception: @@ -1100,8 +1075,6 @@ def _wait_for_clone_points( expected_count: Optional[int], timeout_seconds: int = 60, ) -> None: - if QdrantClient is None: - return try: client = QdrantClient(url=qdrant_url, api_key=os.environ.get("QDRANT_API_KEY") or None) except Exception: @@ -1167,25 +1140,15 @@ def start_staging_rebuild(*, collection: str, work_dir: str) -> str: qdrant_url = os.environ.get("QDRANT_URL", "http://qdrant:6333") source_point_count = _get_collection_point_count(collection_name=collection, qdrant_url=qdrant_url) - # Use local import for thread-safety and determinism - _copy_fn: Any = copy_collection_qdrant - if _copy_fn is None: - # Re-import for container environments where module-level import may have failed - from scripts.collection_admin import copy_collection_qdrant as _ccq - _copy_fn = _ccq - - if not callable(_copy_fn): - raise RuntimeError("copy_collection_qdrant unavailable (import failed)") - try: print(f"[staging] Copying collection {collection} -> {old_collection} (overwrite=True)") try: print( - f"[staging] copy_collection_qdrant callable={callable(_copy_fn)} type={type(_copy_fn)} module={getattr(_copy_fn, '__module__', '?')}" + f"[staging] copy_collection_qdrant module={getattr(copy_collection_qdrant, '__module__', '?')}" ) except Exception: pass - _copy_fn( + copy_collection_qdrant( source=collection, target=old_collection, qdrant_url=qdrant_url, @@ -1291,9 +1254,9 @@ def start_staging_rebuild(*, collection: str, work_dir: str) -> str: pending_env = state.get("indexing_env_pending") or dict(os.environ) env_hash = pending_hash or current_env_indexing_hash() - if not pending_cfg and get_indexing_config_snapshot: - pending_cfg = get_indexing_config_snapshot() if callable(get_indexing_config_snapshot) else get_indexing_config_snapshot - if not pending_hash and pending_cfg and compute_indexing_config_hash: + if not pending_cfg: + pending_cfg = get_indexing_config_snapshot() + if not pending_hash and pending_cfg: pending_hash = compute_indexing_config_hash(pending_cfg) if set_staging_state: diff --git a/scripts/ingest/__init__.py b/scripts/ingest/__init__.py index 6b802a95..e273a90e 100644 --- a/scripts/ingest/__init__.py +++ b/scripts/ingest/__init__.py @@ -1,36 +1,9 @@ -""" -Ingest package - Code indexing subsystem. - -This package contains extracted modules from ingest_code.py: -- config: Environment-based configuration and constants -- tree_sitter: Tree-sitter setup and language loading -- vectors: Vector generation utilities (lex hash, mini projection) -- exclusions: File and directory exclusion logic -- chunking: Code chunking utilities (line, semantic, token-based) -- symbols: Symbol extraction for code analysis -- pseudo: ReFRAG pseudo-description and tag generation -- metadata: Metadata extraction (git, imports, calls) -- qdrant: Qdrant schema and I/O operations -- pipeline: Core indexing pipeline -- cli: Command-line interface +"""Code indexing subsystem package. -Usage: - from scripts.ingest import config, pipeline, qdrant - from scripts.ingest.config import LEX_VECTOR_NAME, LEX_VECTOR_DIM - from scripts.ingest.pipeline import index_repo, index_single_file - from scripts.ingest.qdrant import ensure_collection, upsert_points +Submodules are intentionally loaded on demand. Importing lightweight helpers +such as ``scripts.ingest.config`` should not initialize Qdrant, tree-sitter, or +the indexing pipeline. """ -from scripts.ingest import config -from scripts.ingest import tree_sitter -from scripts.ingest import vectors -from scripts.ingest import exclusions -from scripts.ingest import chunking -from scripts.ingest import symbols -from scripts.ingest import pseudo -from scripts.ingest import metadata -from scripts.ingest import qdrant -from scripts.ingest import pipeline -from scripts.ingest import cli __all__ = [ "config", diff --git a/scripts/ingest/chunking.py b/scripts/ingest/chunking.py index a90ffdf6..9a1e36a4 100644 --- a/scripts/ingest/chunking.py +++ b/scripts/ingest/chunking.py @@ -14,12 +14,9 @@ from scripts.ingest.config import ROOT_DIR from scripts.ingest.tree_sitter import _use_tree_sitter, _TS_LANGUAGES -# Import AST analyzer for enhanced semantic chunking -try: - from scripts.ast_analyzer import get_ast_analyzer, chunk_code_semantically - _AST_ANALYZER_AVAILABLE = True -except ImportError: - _AST_ANALYZER_AVAILABLE = False +from scripts.ast_analyzer import get_ast_analyzer, chunk_code_semantically + +_AST_ANALYZER_AVAILABLE = True # Cache tokenizers loaded from TOKENIZER_JSON (or default) to avoid repeatedly @@ -57,10 +54,6 @@ def chunk_semantic( _ast_supported = False if use_enhanced and _AST_ANALYZER_AVAILABLE: try: - # ast_analyzer internally respects USE_TREE_SITTER when constructing the analyzer - # (see scripts/ast_analyzer.py:get_ast_analyzer). - from scripts.ast_analyzer import get_ast_analyzer # type: ignore - analyzer = get_ast_analyzer() lang_key = str(language or "").strip().lower() # Supported either via builtin ast (python) or via tree-sitter when enabled. diff --git a/scripts/ingest/cli.py b/scripts/ingest/cli.py index 8561a9c2..1e089a16 100644 --- a/scripts/ingest/cli.py +++ b/scripts/ingest/cli.py @@ -9,15 +9,20 @@ import os import argparse +import logging from pathlib import Path from scripts.ingest.config import ( is_multi_repo_mode, get_collection_name, ) +from scripts.pseudo_config import env_bool, effective_pseudo_mode +from scripts.collection_health import clear_indexing_caches as _clear_indexing_caches_impl from scripts.ingest.pipeline import index_repo from scripts.ingest.pseudo import generate_pseudo_tags +logger = logging.getLogger(__name__) + def parse_args(): """Parse command-line arguments.""" @@ -40,6 +45,11 @@ def parse_args(): action="store_true", help="Do not skip files whose content hash matches existing index", ) + parser.add_argument( + "--clear-indexing-caches", + action="store_true", + help="Clear local indexing caches (file hash/symbol caches) before indexing", + ) parser.add_argument( "--schema-mode", type=str, @@ -186,13 +196,31 @@ def main(): ) return + def _clear_indexing_caches(workspace_root: Path, repo_name: str | None) -> None: + try: + _clear_indexing_caches_impl(str(workspace_root), repo_name=repo_name) + except Exception as e: + logger.warning( + "Failed to clear indexing caches for workspace=%s repo=%s: %s", + workspace_root, + repo_name, + e, + exc_info=True, + ) + qdrant_url = os.environ.get("QDRANT_URL", "http://localhost:6333") api_key = os.environ.get("QDRANT_API_KEY") collection = os.environ.get("COLLECTION_NAME") or os.environ.get("DEFAULT_COLLECTION") or "codebase" model_name = os.environ.get("EMBEDDING_MODEL", "BAAI/bge-base-en-v1.5") # Resolve collection name based on multi-repo mode - multi_repo = bool(is_multi_repo_mode and is_multi_repo_mode()) + force_collection = (os.environ.get("CTXCE_FORCE_COLLECTION_NAME") or "").strip().lower() in { + "1", + "true", + "yes", + "on", + } + multi_repo = bool(is_multi_repo_mode and is_multi_repo_mode()) and not force_collection if multi_repo: print("[multi_repo] Multi-repo mode enabled - will create separate collections per repository") @@ -231,6 +259,9 @@ def main(): if not repo_collection: repo_collection = "codebase" + if args.clear_indexing_caches: + _clear_indexing_caches(root_path, repo_name) + index_repo( repo_root, qdrant_url, @@ -240,25 +271,25 @@ def main(): args.recreate, dedupe=(not args.no_dedupe), skip_unchanged=(not args.no_skip_unchanged), - pseudo_mode="off" if (os.environ.get("PSEUDO_DEFER_TO_WORKER") or "").strip().lower() in {"1", "true", "yes", "on"} else "full", + pseudo_mode=effective_pseudo_mode( + defer_to_worker=env_bool("PSEUDO_DEFER_TO_WORKER"), + backfill_enabled=env_bool("PSEUDO_BACKFILL_ENABLED"), + ), schema_mode=args.schema_mode, ) return else: - if get_collection_name: - try: - resolved = get_collection_name(str(Path(args.root).resolve())) - placeholders = {"", "default-collection", "my-collection", "codebase"} - if resolved and collection in placeholders: - collection = resolved - except Exception: - pass if not collection: collection = os.environ.get("COLLECTION_NAME", "codebase") print(f"[single_repo] Single-repo mode enabled - using collection: {collection}") - flag = (os.environ.get("PSEUDO_DEFER_TO_WORKER") or "").strip().lower() - pseudo_mode = "off" if flag in {"1", "true", "yes", "on"} else "full" + pseudo_mode = effective_pseudo_mode( + defer_to_worker=env_bool("PSEUDO_DEFER_TO_WORKER"), + backfill_enabled=env_bool("PSEUDO_BACKFILL_ENABLED"), + ) + + if args.clear_indexing_caches: + _clear_indexing_caches(Path(args.root).resolve(), None) index_repo( Path(args.root).resolve(), diff --git a/scripts/ingest/config.py b/scripts/ingest/config.py index c4001eaf..89f10374 100644 --- a/scripts/ingest/config.py +++ b/scripts/ingest/config.py @@ -221,68 +221,41 @@ def _env_truthy(val: str | None, default: bool) -> bool: # --------------------------------------------------------------------------- -# Workspace state function imports (optional) +# Workspace state function imports # --------------------------------------------------------------------------- -# These are imported at module load time for convenience, with fallbacks -try: - from scripts.workspace_state import ( - is_multi_repo_mode, - get_collection_name, - logical_repo_reuse_enabled, - ) -except ImportError: - is_multi_repo_mode = None # type: ignore - get_collection_name = None # type: ignore +from scripts.workspace_state import ( + is_multi_repo_mode, + get_collection_name, + logical_repo_reuse_enabled, + log_activity, + get_cached_file_hash, + set_cached_file_hash, + remove_cached_file, + update_indexing_status, + update_workspace_state, + get_cached_symbols, + set_cached_symbols, + remove_cached_symbols, + compare_symbol_changes, + get_cached_pseudo, + set_cached_pseudo, + update_symbols_with_pseudo, + get_workspace_state, + get_cached_file_meta, + indexing_lock, + file_indexing_lock, + is_file_locked, +) - def logical_repo_reuse_enabled() -> bool: # type: ignore[no-redef] - return False +def _detect_repo_for_file(path): + """Defer watcher routing import to avoid the ingest/watch bootstrap cycle.""" + from scripts.watch_index_core.routing import _detect_repo_for_file as _impl -# Import watcher's repo detection for surgical fix -try: - from scripts.watch_index_core.routing import _detect_repo_for_file, _get_collection_for_file -except ImportError: - _detect_repo_for_file = None # type: ignore - _get_collection_for_file = None # type: ignore + return _impl(path) -# Import other workspace state functions (optional) -try: - from scripts.workspace_state import ( - log_activity, - get_cached_file_hash, - set_cached_file_hash, - remove_cached_file, - update_indexing_status, - update_workspace_state, - get_cached_symbols, - set_cached_symbols, - remove_cached_symbols, - compare_symbol_changes, - get_cached_pseudo, - set_cached_pseudo, - update_symbols_with_pseudo, - get_workspace_state, - get_cached_file_meta, - indexing_lock, - file_indexing_lock, - is_file_locked, - ) -except ImportError: - # State integration is optional; continue if not available - log_activity = None # type: ignore - get_cached_file_hash = None # type: ignore - set_cached_file_hash = None # type: ignore - remove_cached_file = None # type: ignore - update_indexing_status = None # type: ignore - update_workspace_state = None # type: ignore - get_cached_symbols = None # type: ignore - set_cached_symbols = None # type: ignore - remove_cached_symbols = None # type: ignore - get_cached_pseudo = None # type: ignore - set_cached_pseudo = None # type: ignore - update_symbols_with_pseudo = None # type: ignore - compare_symbol_changes = None # type: ignore - get_workspace_state = None # type: ignore - get_cached_file_meta = None # type: ignore - indexing_lock = None # type: ignore - file_indexing_lock = None # type: ignore - is_file_locked = None # type: ignore + +def _get_collection_for_file(path): + """Defer watcher routing import to avoid the ingest/watch bootstrap cycle.""" + from scripts.watch_index_core.routing import _get_collection_for_file as _impl + + return _impl(path) diff --git a/scripts/ingest/graph_edges.py b/scripts/ingest/graph_edges.py new file mode 100644 index 00000000..3751d970 --- /dev/null +++ b/scripts/ingest/graph_edges.py @@ -0,0 +1,473 @@ +#!/usr/bin/env python3 +""" +ingest/graph_edges.py - Materialized graph edges in Qdrant. + +This is a small, MIT-safe reimplementation of the "graph edges collection" idea: +- Maintain a dedicated Qdrant collection named `_graph` +- Store payload-only edge docs for fast lookups: + - callers/importers queries become simple keyword filters on an indexed payload field + +Design goals for this branch: +- Keep this as an *accelerator* (symbol_graph still works without it) +- Avoid Neo4j/PageRank/GraphRAG complexity +- Avoid CLI flags; watcher can backfill opportunistically +""" + +from __future__ import annotations + +import hashlib +import logging +import os +from typing import Any, Dict, Iterable, List, Optional, Tuple + +logger = logging.getLogger(__name__) + +GRAPH_COLLECTION_SUFFIX = "_graph" + +EDGE_TYPE_CALLS = "calls" +EDGE_TYPE_IMPORTS = "imports" + +GRAPH_INDEX_FIELDS: Tuple[str, ...] = ( + "caller_path", + "callee_symbol", + "edge_type", + "repo", +) + +_ENSURED_GRAPH_COLLECTIONS: set[str] = set() +_GRAPH_VECTOR_MODE: dict[str, str] = {} +_MISSING_GRAPH_COLLECTIONS: set[str] = set() +_BACKFILL_OFFSETS: dict[tuple[str, Optional[str]], Any] = {} + +_EDGE_VECTOR_NAME = "_edge" +_EDGE_VECTOR_VALUE = [0.0] + + +def _normalize_path(path: str) -> str: + if not path: + return "" + try: + normalized = os.path.normpath(str(path)) + except Exception: + normalized = str(path) + return normalized.replace("\\", "/") + + +def normalize_caller_path(path: str) -> str: + """Normalize a caller path exactly as graph edge payloads do. + + This is used by both the writer (upsert/delete) and any readers/verifiers so + cross-platform path separators (Windows vs POSIX) do not cause mismatches. + """ + + return _normalize_path(path) + + +def get_graph_collection_name(base_collection: str) -> str: + return f"{base_collection}{GRAPH_COLLECTION_SUFFIX}" + + +def _edge_vector_for_upsert(graph_collection: str) -> dict: + mode = _GRAPH_VECTOR_MODE.get(graph_collection) + if mode == "named": + return {_EDGE_VECTOR_NAME: _EDGE_VECTOR_VALUE} + return {} + + +def ensure_graph_collection(client: Any, base_collection: str) -> Optional[str]: + """Ensure `_graph` exists and has payload indexes.""" + from qdrant_client import models as qmodels + from qdrant_client.http.exceptions import UnexpectedResponse + + if not base_collection: + return None + graph_coll = get_graph_collection_name(base_collection) + if graph_coll in _ENSURED_GRAPH_COLLECTIONS: + return graph_coll + + def _detect_vector_mode(info: Any) -> str: + try: + vectors = getattr( + getattr(getattr(info, "config", None), "params", None), "vectors", None + ) + if isinstance(vectors, dict): + return "none" if not vectors else "named" + return "none" if vectors is None else "named" + except Exception: + return "named" + + try: + info = client.get_collection(graph_coll) + _GRAPH_VECTOR_MODE[graph_coll] = _detect_vector_mode(info) + _ENSURED_GRAPH_COLLECTIONS.add(graph_coll) + _MISSING_GRAPH_COLLECTIONS.discard(graph_coll) + return graph_coll + except UnexpectedResponse as e: + # Only a 404 means "missing"; any other HTTP failure should be visible. + if getattr(e, "status_code", None) != 404: + logger.exception( + "Failed to get graph collection %s (status=%s): %s", + graph_coll, + getattr(e, "status_code", None), + e, + ) + return None + except Exception as e: + logger.exception("Failed to get graph collection %s: %s", graph_coll, e) + return None + + try: + # Prefer vector-less collection when supported by server/client. + try: + client.create_collection( + collection_name=graph_coll, + vectors_config={}, + ) + _GRAPH_VECTOR_MODE[graph_coll] = "none" + except Exception as vec_exc: + logger.debug( + "Vector-less creation failed for %s, trying named vector: %s", + graph_coll, + vec_exc, + ) + client.create_collection( + collection_name=graph_coll, + vectors_config={ + _EDGE_VECTOR_NAME: qmodels.VectorParams( + size=1, distance=qmodels.Distance.COSINE + ) + }, + ) + _GRAPH_VECTOR_MODE[graph_coll] = "named" + + # Create payload indexes (best-effort). + for field in GRAPH_INDEX_FIELDS: + try: + client.create_payload_index( + collection_name=graph_coll, + field_name=field, + field_schema=qmodels.PayloadSchemaType.KEYWORD, + ) + except Exception as e: + logger.debug( + "Failed to create graph payload index '%s' for %s: %s", + field, + graph_coll, + e, + exc_info=True, + ) + + _ENSURED_GRAPH_COLLECTIONS.add(graph_coll) + _MISSING_GRAPH_COLLECTIONS.discard(graph_coll) + return graph_coll + except Exception as e: + logger.debug("Failed to ensure graph collection %s: %s", graph_coll, e) + return None + + +def _edge_id(edge_type: str, repo: str, caller_path: str, callee_symbol: str) -> str: + key = f"{edge_type}\x00{repo}\x00{caller_path}\x00{callee_symbol}" + return hashlib.sha256(key.encode("utf-8", errors="ignore")).hexdigest()[:32] + + +def _iter_edges( + *, + caller_path: str, + repo: str, + calls: Iterable[str] = (), + imports: Iterable[str] = (), +) -> List[Dict[str, Any]]: + norm_path = _normalize_path(caller_path) + repo_s = (repo or "").strip() or "default" + + edges: List[Dict[str, Any]] = [] + for sym in calls or []: + s = str(sym).strip() + if not s: + continue + edges.append( + { + "id": _edge_id(EDGE_TYPE_CALLS, repo_s, norm_path, s), + "payload": { + "caller_path": norm_path, + "callee_symbol": s, + "edge_type": EDGE_TYPE_CALLS, + "repo": repo_s, + }, + } + ) + for sym in imports or []: + s = str(sym).strip() + if not s: + continue + edges.append( + { + "id": _edge_id(EDGE_TYPE_IMPORTS, repo_s, norm_path, s), + "payload": { + "caller_path": norm_path, + "callee_symbol": s, + "edge_type": EDGE_TYPE_IMPORTS, + "repo": repo_s, + }, + } + ) + return edges + + +def upsert_file_edges( + client: Any, + base_collection: str, + *, + caller_path: str, + repo: str | None, + calls: List[str] | None = None, + imports: List[str] | None = None, +) -> int: + graph_coll = ensure_graph_collection(client, base_collection) + if not graph_coll: + return 0 + edges = _iter_edges( + caller_path=caller_path, + repo=repo or "default", + calls=calls or [], + imports=imports or [], + ) + if not edges: + return 0 + + from qdrant_client import models as qmodels + + points = [ + qmodels.PointStruct( + id=e["id"], + vector=_edge_vector_for_upsert(graph_coll), + payload=e["payload"], + ) + for e in edges + ] + try: + client.upsert(collection_name=graph_coll, points=points, wait=True) + return len(points) + except Exception as e: + logger.debug("Graph edge upsert failed for %s: %s", caller_path, e) + return 0 + + +def delete_edges_by_path( + client: Any, + base_collection: str, + *, + caller_path: str, + repo: str | None = None, +) -> int: + from qdrant_client.http.exceptions import UnexpectedResponse + graph_coll = get_graph_collection_name(base_collection) + if graph_coll in _MISSING_GRAPH_COLLECTIONS: + return 0 + + from qdrant_client import models as qmodels + + norm_path = _normalize_path(caller_path) + must: list[Any] = [ + qmodels.FieldCondition( + key="caller_path", match=qmodels.MatchValue(value=norm_path) + ) + ] + if repo: + r = str(repo).strip() + if r and r != "*": + must.append( + qmodels.FieldCondition(key="repo", match=qmodels.MatchValue(value=r)) + ) + flt = qmodels.Filter(must=must) + + # Probe first so callers can distinguish "no matching rows" (0) from a real delete. + # This is important for fallback logic (e.g., retry path-only delete when repo tag drifted). + try: + existing, _ = client.scroll( + collection_name=graph_coll, + scroll_filter=flt, + limit=1, + with_payload=False, + with_vectors=False, + ) + if not existing: + return 0 + except UnexpectedResponse as e: + if getattr(e, "status_code", None) == 404: + _MISSING_GRAPH_COLLECTIONS.add(graph_coll) + return 0 + logger.debug( + "Graph edge probe failed for %s in %s (status=%s): %s", + norm_path, + graph_coll, + getattr(e, "status_code", None), + e, + exc_info=True, + ) + return 0 + except Exception as e: + logger.debug( + "Graph edge probe failed for %s in %s: %s", + norm_path, + graph_coll, + e, + exc_info=True, + ) + return 0 + + try: + resp = client.delete( + collection_name=graph_coll, + points_selector=qmodels.FilterSelector(filter=flt), + ) + result_status = getattr(getattr(resp, "result", None), "status", None) + if result_status is None: + result_status = getattr(resp, "status", None) + if result_status is None: + return 1 + status_s = str(result_status).strip().lower() + return 1 if status_s in {"acknowledged", "completed", "ok", "success"} else 0 + except UnexpectedResponse as e: + if getattr(e, "status_code", None) == 404: + _MISSING_GRAPH_COLLECTIONS.add(graph_coll) + return 0 + logger.debug( + "Graph edge delete failed for %s in %s (status=%s): %s", + norm_path, + graph_coll, + getattr(e, "status_code", None), + e, + exc_info=True, + ) + return 0 + except Exception as e: + logger.debug( + "Graph edge delete failed for %s in %s: %s", + norm_path, + graph_coll, + e, + exc_info=True, + ) + return 0 + + +def graph_edges_backfill_tick( + client: Any, + base_collection: str, + *, + repo_name: str | None = None, + max_files: int = 128, +) -> int: + """Best-effort incremental backfill from `` into `_graph`. + + This scans the main collection and upserts file-level edges into the graph collection. + It's idempotent (deterministic IDs) and safe to run continuously in a watcher worker. + """ + from qdrant_client import models as qmodels + + if not base_collection or max_files <= 0: + return 0 + + graph_coll = ensure_graph_collection(client, base_collection) + if not graph_coll: + return 0 + + must: list[Any] = [] + if repo_name: + must.append( + qmodels.FieldCondition( + key="metadata.repo", match=qmodels.MatchValue(value=repo_name) + ) + ) + flt = qmodels.Filter(must=must or None) + + processed_files = 0 + seen_paths: set[str] = set() + + key = (base_collection, repo_name) + next_offset = _BACKFILL_OFFSETS.get(key) + + # We may need to overscan because the main collection is chunked. + overscan = max_files * 8 + while processed_files < max_files: + attempts = 0 + while True: + try: + points, next_offset = client.scroll( + collection_name=base_collection, + scroll_filter=flt, + limit=min(64, overscan), + with_payload=True, + with_vectors=False, + offset=next_offset, + ) + break + except Exception as e: + attempts += 1 + logger.exception( + "Graph edge backfill scroll failed (collection=%s repo=%s offset=%s attempt=%d): %s", + base_collection, + repo_name or "default", + next_offset, + attempts, + e, + ) + # Retry a couple times for transient errors, then raise so failures are not silent. + if attempts >= 3: + raise + import time + + time.sleep(0.25 * (2 ** (attempts - 1))) + + if not points: + break + + for rec in points: + if processed_files >= max_files: + break + payload = getattr(rec, "payload", None) or {} + md = payload.get("metadata") or {} + path = md.get("path") or "" + if not path: + continue + norm_path = _normalize_path(str(path)) + if norm_path in seen_paths: + continue + seen_paths.add(norm_path) + + repo = md.get("repo") or repo_name or "default" + calls = md.get("calls") or [] + imports = md.get("imports") or [] + if not isinstance(calls, list): + calls = [] + if not isinstance(imports, list): + imports = [] + + upsert_file_edges( + client, + base_collection, + caller_path=norm_path, + repo=str(repo), + calls=[str(x) for x in calls if x], + imports=[str(x) for x in imports if x], + ) + processed_files += 1 + + if next_offset is None: + break + + _BACKFILL_OFFSETS[key] = next_offset + return processed_files + + +__all__ = [ + "GRAPH_COLLECTION_SUFFIX", + "EDGE_TYPE_CALLS", + "EDGE_TYPE_IMPORTS", + "get_graph_collection_name", + "ensure_graph_collection", + "upsert_file_edges", + "delete_edges_by_path", + "graph_edges_backfill_tick", +] diff --git a/scripts/ingest/pipeline.py b/scripts/ingest/pipeline.py index 45716049..1c7311c0 100644 --- a/scripts/ingest/pipeline.py +++ b/scripts/ingest/pipeline.py @@ -14,7 +14,18 @@ from pathlib import Path from typing import List, Dict, Any, Optional, TYPE_CHECKING -from qdrant_client import QdrantClient, models +if TYPE_CHECKING: + from qdrant_client import QdrantClient, models as models +else: + QdrantClient = Any # type: ignore + + class _LazyQdrantModels: + def __getattr__(self, name: str) -> Any: + from qdrant_client import models as _models + + return getattr(_models, name) + + models = _LazyQdrantModels() from scripts.ingest.config import ( ROOT_DIR, @@ -52,6 +63,7 @@ extract_symbols_with_tree_sitter, ) from scripts.ingest.pseudo import ( + _pseudo_describe_enabled, generate_pseudo_tags, should_process_pseudo_for_chunk, ) @@ -60,7 +72,7 @@ _get_imports_calls, _compute_host_and_container_paths, ) -from scripts.ingest.vectors import project_mini, extract_pattern_vector +from scripts.ingest.vectors import project_mini from scripts.ingest.qdrant import ( ensure_collection, ensure_collection_and_indexes_once, @@ -72,7 +84,10 @@ upsert_points, hash_id, embed_batch, - PATTERN_VECTOR_NAME, +) +from scripts.relevance_feedback import ( + build_symbol_reconciliations, + reconcile_collection_weights, ) # Import utility functions @@ -84,13 +99,27 @@ from fastembed import TextEmbedding +def _pseudo_batch_concurrency() -> int: + try: + return max(1, int(os.environ.get("PSEUDO_BATCH_CONCURRENCY", "1") or 1)) + except (TypeError, ValueError): + return 1 + + +def _use_batch_pseudo(pseudo_mode: str) -> bool: + """Enable the GLM batch shortcut only when pseudo generation is explicit.""" + return _pseudo_describe_enabled() and _pseudo_batch_concurrency() > 1 and pseudo_mode == "full" + + def _detect_repo_name_from_path(path: Path) -> str: """Wrapper function to use workspace_state repository detection.""" - try: - from scripts.workspace_state import _extract_repo_name_from_path as _ws_detect - return _ws_detect(str(path)) - except ImportError: - return path.name if path.is_dir() else path.parent.name + from scripts.workspace_state import _extract_repo_name_from_path as _ws_detect + + # `_extract_repo_name_from_path` expects a workspace/repo path shape, not a file path. + # Always normalize file inputs to their parent directory to avoid falling back to + # file basenames (which can poison metadata.repo and graph edge repo tags). + candidate = path if path.is_dir() else path.parent + return _ws_detect(str(candidate)) def detect_language(path: Path) -> str: @@ -109,7 +138,8 @@ def detect_language(path: Path) -> str: _TEXT_LIKE_LANGS = {"unknown", "markdown", "text"} -def _is_text_like_language(language: str) -> bool: +def is_text_like_language(language: str) -> bool: + """Classify whether a detected language should skip smart reindexing.""" return str(language or "").strip().lower() in _TEXT_LIKE_LANGS @@ -227,6 +257,87 @@ def _normalize_info_for_dense(s: str) -> str: return text +def _sync_graph_edges_best_effort( + client: QdrantClient, + collection: str, + file_path: str, + repo: str | None, + calls: list[str] | None, + imports: list[str] | None, +) -> None: + """Best-effort sync of file-level graph edges. Safe to skip on failure.""" + enabled = str(os.environ.get("GRAPH_EDGES_ENABLE", "1") or "").strip().lower() in { + "1", + "true", + "yes", + "on", + } + if not enabled: + return + try: + from scripts.ingest.graph_edges import ( + delete_edges_by_path, + ensure_graph_collection, + upsert_file_edges, + ) + + ensure_graph_collection(client, collection) + # Important: delete stale edges for this file before upserting the new set. + delete_edges_by_path( + client, + collection, + caller_path=str(file_path), + repo=repo, + ) + upsert_file_edges( + client, + collection, + caller_path=str(file_path), + repo=repo, + calls=calls, + imports=imports, + ) + except Exception as exc: + try: + print(f"[graph_edges] best-effort sync failed for {file_path}: {exc}") + except Exception: + pass + + +def _symbols_to_metadata_dict(language: str, text: str) -> dict: + """Build symbol metadata dict from in-memory source text.""" + symbols = {} + try: + symbols_list = _extract_symbols(language, text) + lines = text.split("\n") + for sym in symbols_list or []: + kind = str(sym.get("kind") or "") + name = str(sym.get("name") or "") + start = int(sym.get("start") or 0) + end = int(sym.get("end") or 0) + if not kind or not name or start <= 0 or end < start: + continue + symbol_id = f"{kind}_{name}_{start}" + content = "\n".join(lines[start - 1 : end]) + content_hash = hashlib.sha1( + content.encode("utf-8", errors="ignore") + ).hexdigest() + symbols[symbol_id] = { + "name": name, + "type": kind, + "start_line": start, + "end_line": end, + "content_hash": content_hash, + "content": content, + "pseudo": "", + "tags": [], + "qdrant_ids": [], + } + except Exception: + return {} + return symbols + + def build_information( language: str, path: Path, start: int, end: int, first_line: str ) -> str: @@ -251,16 +362,31 @@ def index_single_file( repo_name_for_cache: str | None = None, allowed_vectors: set[str] | None = None, allowed_sparse: set[str] | None = None, + preloaded_text: str | None = None, + preloaded_file_hash: str | None = None, + preloaded_language: str | None = None, ) -> bool: """Index a single file path. Returns True if indexed, False if skipped.""" + repo_for_graph = repo_name_for_cache or _detect_repo_name_from_path(file_path) try: if _should_skip_explicit_file_by_excluder(file_path): try: delete_points_by_path(client, collection, str(file_path)) except Exception: pass + # Clean up graph edges for excluded file + _sync_graph_edges_best_effort( + client, + collection, + str(file_path), + repo_for_graph, + None, # No calls when file is excluded + None, # No imports when file is excluded + ) print(f"Skipping excluded file: {file_path}") return False + except NameError: + raise except Exception: return False @@ -283,6 +409,9 @@ def index_single_file( repo_name_for_cache=repo_name_for_cache, allowed_vectors=allowed_vectors, allowed_sparse=allowed_sparse, + preloaded_text=preloaded_text, + preloaded_file_hash=preloaded_file_hash, + preloaded_language=preloaded_language, ) finally: if _file_lock_ctx is not None: @@ -306,6 +435,9 @@ def _index_single_file_inner( repo_name_for_cache: str | None = None, allowed_vectors: set[str] | None = None, allowed_sparse: set[str] | None = None, + preloaded_text: str | None = None, + preloaded_file_hash: str | None = None, + preloaded_language: str | None = None, ) -> bool: """Inner implementation of index_single_file (after lock is acquired).""" if trust_cache is None: @@ -317,7 +449,12 @@ def _index_single_file_inner( trust_cache = False fast_fs = _env_truthy(os.environ.get("INDEX_FS_FASTPATH"), False) - if skip_unchanged and fast_fs and get_cached_file_meta is not None: + if ( + preloaded_text is None + and skip_unchanged + and fast_fs + and get_cached_file_meta is not None + ): try: repo_for_cache = repo_name_for_cache or _detect_repo_name_from_path(file_path) meta = get_cached_file_meta(str(file_path), repo_for_cache) or {} @@ -333,15 +470,17 @@ def _index_single_file_inner( except Exception: pass - try: - text = file_path.read_text(encoding="utf-8", errors="ignore") - except Exception as e: - print(f"Skipping {file_path}: {e}") - return False + if preloaded_text is None: + try: + text = file_path.read_text(encoding="utf-8", errors="ignore") + except Exception as e: + print(f"Skipping {file_path}: {e}") + return False + else: + text = preloaded_text - language = detect_language(file_path) - is_text_like = _is_text_like_language(language) - file_hash = hashlib.sha1(text.encode("utf-8", errors="ignore")).hexdigest() + language = preloaded_language or detect_language(file_path) + file_hash = preloaded_file_hash or hashlib.sha1(text.encode("utf-8", errors="ignore")).hexdigest() repo_tag = repo_name_for_cache or _detect_repo_name_from_path(file_path) @@ -376,7 +515,10 @@ def _index_single_file_inner( if get_cached_symbols and set_cached_symbols: cached_symbols = get_cached_symbols(str(file_path)) if cached_symbols: - current_symbols = extract_symbols_with_tree_sitter(str(file_path)) + if preloaded_text is not None: + current_symbols = _symbols_to_metadata_dict(language, preloaded_text) + else: + current_symbols = extract_symbols_with_tree_sitter(str(file_path)) _, changed = compare_symbol_changes(cached_symbols, current_symbols) for symbol_data in current_symbols.values(): symbol_id = f"{symbol_data['type']}_{symbol_data['name']}_{symbol_data['start_line']}" @@ -456,24 +598,19 @@ def _index_single_file_inner( batch_ids: List[int] = [] batch_lex: List[list[float]] = [] batch_lex_text: List[str] = [] - batch_code: List[str] = [] # Raw code for pattern vectors if allowed_vectors is None and allowed_sparse is None: allowed_vectors, allowed_sparse = get_collection_vector_names(client, collection) allow_lex = allowed_vectors is None or LEX_VECTOR_NAME in allowed_vectors allow_mini = allowed_vectors is None or MINI_VECTOR_NAME in allowed_vectors - allow_pattern = allowed_vectors is None or PATTERN_VECTOR_NAME in allowed_vectors allow_sparse = allowed_sparse is None or LEX_SPARSE_NAME in allowed_sparse - # Check if pattern vectors are enabled - pattern_vectors_on = os.environ.get("PATTERN_VECTORS", "").strip().lower() in {"1", "true", "yes", "on"} - pattern_vectors_on = pattern_vectors_on and allow_pattern refrag_on = os.environ.get("REFRAG_MODE", "").strip().lower() in {"1", "true", "yes", "on"} use_mini = refrag_on and allow_mini use_sparse = LEX_SPARSE_MODE and allow_sparse - def make_point(pid, dense_vec, lex_vec, payload, lex_text: str = "", code_text: str = ""): + def make_point(pid, dense_vec, lex_vec, payload, lex_text: str = ""): if vector_name: vecs = {vector_name: dense_vec} if allow_lex: @@ -483,14 +620,6 @@ def make_point(pid, dense_vec, lex_vec, payload, lex_text: str = "", code_text: vecs[MINI_VECTOR_NAME] = project_mini(list(dense_vec), MINI_VEC_DIM) except Exception: pass - # Add pattern vector for structural similarity search - if pattern_vectors_on and code_text: - try: - pv = extract_pattern_vector(code_text, language) - if pv: - vecs[PATTERN_VECTOR_NAME] = pv - except Exception: - pass if use_sparse and lex_text: sparse_vec = _lex_sparse_vector_text(lex_text) if sparse_vec.get("indices"): @@ -499,8 +628,24 @@ def make_point(pid, dense_vec, lex_vec, payload, lex_text: str = "", code_text: else: return models.PointStruct(id=pid, vector=dense_vec, payload=payload) - pseudo_batch_concurrency = int(os.environ.get("PSEUDO_BATCH_CONCURRENCY", "1") or 1) - use_batch_pseudo = pseudo_batch_concurrency > 1 and pseudo_mode == "full" + pseudo_batch_concurrency = _pseudo_batch_concurrency() + use_batch_pseudo = _use_batch_pseudo(pseudo_mode) + + def _full_index_symbol_content_hash(kind: str, symbol_name: str) -> str: + matches = [ + sym_info + for sym_info in symbols + if str(sym_info.get("kind") or "") == str(kind) + and str(sym_info.get("name") or "") == str(symbol_name) + ] + if len(matches) != 1: + return "" + sym_info = matches[0] + lines = text.splitlines() + start = max(1, int(sym_info.get("start") or 1)) + end = max(start, int(sym_info.get("end") or start)) + content = "\n".join(lines[start - 1 : end]) + return hashlib.sha1(content.encode("utf-8", errors="ignore")).hexdigest() chunk_data: list[dict] = [] for ch in chunks: @@ -541,6 +686,7 @@ def make_point(pid, dense_vec, lex_vec, payload, lex_text: str = "", code_text: "end_line": ch["end"], "code": ch["text"], "file_hash": file_hash, + "symbol_content_hash": _full_index_symbol_content_hash(kind, sym), "imports": imports, "calls": calls, "ingested_at": int(time.time()), @@ -639,7 +785,6 @@ def make_point(pid, dense_vec, lex_vec, payload, lex_text: str = "", code_text: aug_lex_text = (ch.get("text") or "") + (" " + pseudo if pseudo else "") + (" " + " ".join(tags) if tags else "") batch_lex.append(_lex_hash_vector_text(aug_lex_text)) batch_lex_text.append(aug_lex_text) - batch_code.append(ch.get("text") or "") if batch_texts: vectors = embed_batch(model, batch_texts) @@ -649,10 +794,26 @@ def make_point(pid, dense_vec, lex_vec, payload, lex_text: str = "", code_text: except Exception: pass points = [ - make_point(i, v, lx, m, lt, ct) - for i, v, lx, m, lt, ct in zip(batch_ids, vectors, batch_lex, batch_meta, batch_lex_text, batch_code) + make_point(i, v, lx, m, lt) + for i, v, lx, m, lt in zip(batch_ids, vectors, batch_lex, batch_meta, batch_lex_text) ] upsert_points(client, collection, points) + + # Optional: materialize file-level graph edges in a companion `_graph` store. + # This is an accelerator for symbol_graph callers/importers and is safe to skip on failure. + # IMPORTANT: Sync must run after upserts (or after delete-only reindex) to ensure graph + # edges stay consistent. When a file reindexes to zero chunks, batch_texts is empty but + # we still need to sync graph edges to remove stale entries. + _sync_graph_edges_best_effort( + client, + collection, + str(file_path), + repo_tag, + calls, + imports, + ) + + if batch_texts: try: ws = os.environ.get("WATCH_ROOT") or os.environ.get("WORKSPACE_PATH") or "/work" if set_cached_file_hash: @@ -717,14 +878,10 @@ def index_repo( except Exception: pass - try: - from scripts.embedder import get_embedding_model, get_model_dimension - model = get_embedding_model(model_name) - dim = get_model_dimension(model_name) - except ImportError: - from fastembed import TextEmbedding - model = TextEmbedding(model_name=model_name) - dim = len(next(model.embed(["dimension probe"]))) + from scripts.embedder import get_embedding_model, get_model_dimension + + model = get_embedding_model(model_name) + dim = get_model_dimension(model_name) client = QdrantClient( url=qdrant_url, @@ -798,12 +955,6 @@ def index_repo( f"[COLLECTION_WARNING] Collection {collection} lacks mini vector '{MINI_VECTOR_NAME}'. " "ReFRAG vectors will be skipped for this run." ) - pattern_on = os.environ.get("PATTERN_VECTORS", "").strip().lower() in {"1", "true", "yes", "on"} - if pattern_on and PATTERN_VECTOR_NAME not in allowed_vectors: - print( - f"[COLLECTION_WARNING] Collection {collection} lacks pattern vector '{PATTERN_VECTOR_NAME}'. " - "Pattern vectors will be skipped for this run." - ) if LEX_VECTOR_NAME not in allowed_vectors: print( f"[COLLECTION_WARNING] Collection {collection} lacks lexical vector '{LEX_VECTOR_NAME}'. " @@ -881,6 +1032,7 @@ def process_file_with_smart_reindexing( model, vector_name: str | None, *, + model_dim: int | None = None, allowed_vectors: set[str] | None = None, allowed_sparse: set[str] | None = None, ) -> str: @@ -890,18 +1042,10 @@ def process_file_with_smart_reindexing( - Reusing existing embeddings/lexical vectors for unchanged chunks (by code content), and - Re-embedding only for changed chunks. """ - # Allow test monkeypatching on ingest_code.* to be honored here. - # Must be done FIRST before any helper calls. - _ingest_mod = None - try: - import importlib - _ingest_mod = importlib.import_module("scripts.ingest_code") - except Exception: - _ingest_mod = None - _embed_batch = getattr(_ingest_mod, "embed_batch", embed_batch) if _ingest_mod else embed_batch - _upsert_points_fn = getattr(_ingest_mod, "upsert_points", upsert_points) if _ingest_mod else upsert_points - _delete_points_fn = getattr(_ingest_mod, "delete_points_by_path", delete_points_by_path) if _ingest_mod else delete_points_by_path - _should_process_pseudo = getattr(_ingest_mod, "should_process_pseudo_for_chunk", should_process_pseudo_for_chunk) if _ingest_mod else should_process_pseudo_for_chunk + _embed_batch = embed_batch + _upsert_points_fn = upsert_points + _delete_points_fn = delete_points_by_path + _should_process_pseudo = should_process_pseudo_for_chunk try: p = Path(str(file_path)) @@ -910,8 +1054,19 @@ def process_file_with_smart_reindexing( _delete_points_fn(client, current_collection, str(p)) except Exception: pass + # Clean up graph edges for excluded file + _sync_graph_edges_best_effort( + client, + current_collection, + str(p), + per_file_repo or _detect_repo_name_from_path(file_path), + None, # No calls when file is excluded + None, # No imports when file is excluded + ) print(f"[SMART_REINDEX] Skipping excluded file: {file_path}") return "skipped" + except NameError: + raise except Exception: return "skipped" @@ -927,6 +1082,13 @@ def process_file_with_smart_reindexing( except Exception: file_path = Path(fp) + is_text_like = is_text_like_language(language) + if is_text_like: + print( + f"[SMART_REINDEX] {file_path}: text-like language '{language}', " + "skipping smart reindex and using full reindex path" + ) + return "failed" file_hash = hashlib.sha1(text.encode("utf-8", errors="ignore")).hexdigest() if allowed_vectors is None and allowed_sparse is None: @@ -934,11 +1096,8 @@ def process_file_with_smart_reindexing( allow_lex = allowed_vectors is None or LEX_VECTOR_NAME in allowed_vectors allow_mini = allowed_vectors is None or MINI_VECTOR_NAME in allowed_vectors - allow_pattern = allowed_vectors is None or PATTERN_VECTOR_NAME in allowed_vectors allow_sparse = allowed_sparse is None or LEX_SPARSE_NAME in allowed_sparse - pattern_vectors_on = os.environ.get("PATTERN_VECTORS", "").strip().lower() in {"1", "true", "yes", "on"} - pattern_vectors_on = pattern_vectors_on and allow_pattern refrag_on = os.environ.get("REFRAG_MODE", "").strip().lower() in {"1", "true", "yes", "on"} use_mini = refrag_on and allow_mini use_sparse = LEX_SPARSE_MODE and allow_sparse @@ -988,8 +1147,31 @@ def process_file_with_smart_reindexing( changed_set = set(changed_symbols) if len(changed_symbols) == 0 and cached_symbols: - print(f"[SMART_REINDEX] {file_path}: 0 changes detected, skipping") - return "skipped" + prev_hash = None + try: + if get_cached_file_hash: + prev_hash = get_cached_file_hash(fp, per_file_repo) + except Exception: + prev_hash = None + if prev_hash and file_hash and prev_hash == file_hash: + print(f"[SMART_REINDEX] {file_path}: 0 changes detected, skipping") + return "skipped" + print( + f"[SMART_REINDEX] {file_path}: non-symbol change detected; " + "falling back to full reindex" + ) + return "failed" + + if model_dim and vector_name: + try: + ensure_collection_and_indexes_once( + client, + current_collection, + int(model_dim), + vector_name, + ) + except Exception: + pass existing_points = [] try: @@ -1085,7 +1267,6 @@ def process_file_with_smart_reindexing( else: chunks = chunk_lines(text, CHUNK_LINES, CHUNK_OVERLAP) - is_text_like = _is_text_like_language(language) symbol_spans = _extract_symbols(language, text) reused_points: list[models.PointStruct] = [] @@ -1094,13 +1275,52 @@ def process_file_with_smart_reindexing( embed_ids: list[int] = [] embed_lex: list[list[float]] = [] embed_lex_text: list[str] = [] - embed_code: list[str] = [] # Raw code for pattern vectors imports, calls = _get_imports_calls(language, text) last_mod, churn_count, author_count = _git_metadata(file_path) - pseudo_batch_concurrency = int(os.environ.get("PSEUDO_BATCH_CONCURRENCY", "1") or 1) - use_batch_pseudo = pseudo_batch_concurrency > 1 + pseudo_batch_concurrency = _pseudo_batch_concurrency() + # Smart reindexing must use the same explicit generation switch as the + # sequential path; otherwise the batch fast path can call GLM while pseudo + # descriptions are disabled. + use_batch_pseudo = _use_batch_pseudo("full") + + def _apply_symbol_pseudo( + symbol_name: str, + kind: str, + start_line: int, + pseudo_text: str, + pseudo_tags: list[str], + ) -> None: + if not symbol_name or not kind: + return + sid = f"{kind}_{symbol_name}_{start_line}" + target = symbol_meta.get(sid) + if target is None: + for candidate in symbol_meta.values(): + if str(candidate.get("type") or "") != str(kind): + continue + if str(candidate.get("name") or "") != str(symbol_name): + continue + target = candidate + break + if target is None: + return + target["pseudo"] = pseudo_text + target["tags"] = list(pseudo_tags or []) + + def _symbol_content_hash(kind: str, symbol_name: str) -> str: + if not kind or not symbol_name: + return "" + matches = [ + info + for info in symbol_meta.values() + if str(info.get("type") or "") == str(kind) + and str(info.get("name") or "") == str(symbol_name) + ] + if len(matches) == 1: + return str(matches[0].get("content_hash") or "") + return "" chunk_data_sr: list[dict] = [] for ch in chunks: @@ -1141,6 +1361,7 @@ def process_file_with_smart_reindexing( "end_line": ch["end"], "code": ch["text"], "file_hash": file_hash, + "symbol_content_hash": _symbol_content_hash(kind, sym), "imports": imports, "calls": calls, "ingested_at": int(time.time()), @@ -1190,6 +1411,14 @@ def process_file_with_smart_reindexing( start_line = ch.get("start", 0) sid = f"{k}_{symbol_name}_{start_line}" set_cached_pseudo(fp, sid, pseudo, tags, file_hash) + _apply_symbol_pseudo( + symbol_name, + ch.get("kind", "unknown"), + ch.get("start", 0), + pseudo, + tags, + ) + ch["_pseudo_applied"] = True except Exception as e: print(f"[PSEUDO_BATCH] Smart reindex batch failed, falling back: {e}") use_batch_pseudo = False @@ -1211,9 +1440,26 @@ def process_file_with_smart_reindexing( sid = f"{k}_{symbol_name}_{start_line}" if set_cached_pseudo: set_cached_pseudo(fp, sid, pseudo, tags, file_hash) + _apply_symbol_pseudo( + symbol_name, + k, + start_line, + pseudo, + tags, + ) + ch["_pseudo_applied"] = True except Exception: pass + if (pseudo or tags) and not ch.get("_pseudo_applied"): + _apply_symbol_pseudo( + ch.get("symbol", ""), + ch.get("kind", "unknown"), + ch.get("start", 0), + pseudo, + tags, + ) + if pseudo: payload["pseudo"] = pseudo if tags: @@ -1327,12 +1573,11 @@ def process_file_with_smart_reindexing( aug_lex_text = (code_text or "") + (" " + pseudo if pseudo else "") + (" " + " ".join(tags) if tags else "") embed_lex.append(_lex_hash_vector_text(aug_lex_text)) embed_lex_text.append(aug_lex_text) - embed_code.append(code_text or "") new_points: list[models.PointStruct] = [] if embed_texts: vectors = _embed_batch(model, embed_texts) - for pid, v, lx, pl, lt, ct in zip(embed_ids, vectors, embed_lex, embed_payloads, embed_lex_text, embed_code): + for pid, v, lx, pl, lt in zip(embed_ids, vectors, embed_lex, embed_payloads, embed_lex_text): if vector_name: vecs = {vector_name: v} if allow_lex: @@ -1342,14 +1587,6 @@ def process_file_with_smart_reindexing( vecs[MINI_VECTOR_NAME] = project_mini(list(v), MINI_VEC_DIM) except Exception: pass - # Add pattern vector for structural similarity search - if pattern_vectors_on and ct: - try: - pv = extract_pattern_vector(ct, language) - if pv: - vecs[PATTERN_VECTOR_NAME] = pv - except Exception: - pass if use_sparse and lt: sparse_vec = _lex_sparse_vector_text(lt) if sparse_vec.get("indices"): @@ -1368,6 +1605,40 @@ def process_file_with_smart_reindexing( if all_points: _upsert_points_fn(client, current_collection, all_points) + # Optional: materialize file-level graph edges (best-effort). + # IMPORTANT: Sync must run after upserts OR after delete-only reindex to ensure graph + # edges stay consistent. When a file reindexes to zero chunks, all_points is empty but + # we still need to sync graph edges to remove stale entries. + _sync_graph_edges_best_effort( + client, + current_collection, + str(file_path), + per_file_repo, + calls, + imports, + ) + + try: + reconciliations = build_symbol_reconciliations( + cached_symbols, + symbol_meta, + repo=str(per_file_repo or ""), + path=fp, + split_min_overlap=float( + os.environ.get("RELEVANCE_SPLIT_MIN_OVERLAP", "0.45") or 0.45 + ), + split_min_coverage=float( + os.environ.get("RELEVANCE_SPLIT_MIN_COVERAGE", "0.75") or 0.75 + ), + ) + migrated = reconcile_collection_weights(current_collection, reconciliations) + if migrated: + print( + f"[SMART_REINDEX] Reconciled {migrated} feedback target(s) for {file_path}" + ) + except Exception as e: + print(f"[SMART_REINDEX] Feedback reconciliation skipped for {file_path}: {e}") + try: if set_cached_symbols: set_cached_symbols(fp, symbol_meta, file_hash) diff --git a/scripts/ingest/pseudo.py b/scripts/ingest/pseudo.py index ea157e2b..0b02db2e 100644 --- a/scripts/ingest/pseudo.py +++ b/scripts/ingest/pseudo.py @@ -7,11 +7,13 @@ """ from __future__ import annotations +import logging import os from typing import Tuple, List from scripts.ingest.config import ( get_cached_pseudo, + get_cached_symbols, set_cached_pseudo, compare_symbol_changes, ) @@ -130,25 +132,58 @@ def should_process_pseudo_for_chunk( start_line = chunk.get("start", 0) symbol_id = f"{kind}_{symbol_name}_{start_line}" + def _lookup_cached() -> Tuple[str, List[str]]: + if get_cached_pseudo: + try: + cached_pseudo, cached_tags = get_cached_pseudo(file_path, symbol_id) + if cached_pseudo or cached_tags: + return cached_pseudo, cached_tags + except Exception as exc: + logging.getLogger(__name__).debug( + "get_cached_pseudo failed for %s/%s: %s", + file_path, + symbol_id, + exc, + exc_info=True, + ) + if get_cached_symbols: + try: + cached_symbols = get_cached_symbols(file_path) or {} + for info in cached_symbols.values(): + if str(info.get("type") or "") != str(kind): + continue + if str(info.get("name") or "") != str(symbol_name): + continue + cached_pseudo = info.get("pseudo", "") + cached_tags = info.get("tags", []) + if not isinstance(cached_pseudo, str): + cached_pseudo = "" + if not isinstance(cached_tags, list): + cached_tags = [] + cached_tags = [str(tag) for tag in cached_tags if str(tag)] + if cached_pseudo or cached_tags: + return cached_pseudo, cached_tags + except Exception as exc: + logging.getLogger(__name__).debug( + "get_cached_symbols failed for %s: %s", + file_path, + exc, + exc_info=True, + ) + return "", [] + # If we don't have any change information, best effort: try reusing cached pseudo when present - if not changed_symbols and get_cached_pseudo: - try: - cached_pseudo, cached_tags = get_cached_pseudo(file_path, symbol_id) - if cached_pseudo or cached_tags: - return False, cached_pseudo, cached_tags - except Exception: - pass + if not changed_symbols: + cached_pseudo, cached_tags = _lookup_cached() + if cached_pseudo or cached_tags: + return False, cached_pseudo, cached_tags return True, "", [] # Unchanged symbol: prefer reuse when cached pseudo/tags exist if symbol_id not in changed_symbols: - if get_cached_pseudo: - try: - cached_pseudo, cached_tags = get_cached_pseudo(file_path, symbol_id) - if cached_pseudo or cached_tags: - return False, cached_pseudo, cached_tags - except Exception: - pass + cached_pseudo, cached_tags = _lookup_cached() + if cached_pseudo or cached_tags: + return False, cached_pseudo, cached_tags # Unchanged but no cached data yet – process once return True, "", [] @@ -162,7 +197,6 @@ def should_use_smart_reindexing(file_path: str, file_hash: str) -> Tuple[bool, s Returns: (use_smart, reason) """ - from scripts.ingest.config import get_cached_symbols, compare_symbol_changes from scripts.ingest.symbols import extract_symbols_with_tree_sitter if not _smart_symbol_reindexing_enabled(): diff --git a/scripts/ingest/qdrant.py b/scripts/ingest/qdrant.py index f98207ca..d07d98c6 100644 --- a/scripts/ingest/qdrant.py +++ b/scripts/ingest/qdrant.py @@ -11,9 +11,20 @@ import time import hashlib from pathlib import Path -from typing import List, Dict, Any, Optional +from typing import List, Dict, Any, Optional, TYPE_CHECKING -from qdrant_client import QdrantClient, models +if TYPE_CHECKING: + from qdrant_client import QdrantClient, models as models +else: + QdrantClient = Any # type: ignore + + class _LazyQdrantModels: + def __getattr__(self, name: str) -> Any: + from qdrant_client import models as _models + + return getattr(_models, name) + + models = _LazyQdrantModels() from scripts.ingest.config import ( LEX_VECTOR_NAME, @@ -31,6 +42,7 @@ # --------------------------------------------------------------------------- ENSURED_COLLECTIONS: set[str] = set() ENSURED_COLLECTIONS_LAST_CHECK: dict[str, float] = {} +ENSURED_PAYLOAD_INDEX_COLLECTIONS: set[str] = set() class CollectionNeedsRecreateError(Exception): @@ -38,8 +50,7 @@ class CollectionNeedsRecreateError(Exception): pass -PATTERN_VECTOR_NAME = "pattern_vector" -PATTERN_VECTOR_DIM = 64 # Structural pattern embedding dimension + PAYLOAD_INDEX_FIELDS = ( "metadata.language", @@ -126,15 +137,6 @@ def _desired_vector_configs( ) except Exception: pass - try: - if os.environ.get("PATTERN_VECTORS", "").strip().lower() in {"1", "true", "yes", "on"}: - vectors_cfg[PATTERN_VECTOR_NAME] = models.VectorParams( - size=PATTERN_VECTOR_DIM, - distance=models.Distance.COSINE, - ) - except Exception: - pass - sparse_cfg = None if LEX_SPARSE_MODE: sparse_cfg = { @@ -321,7 +323,6 @@ def ensure_collection( Always includes dense (vector_name) and lexical (LEX_VECTOR_NAME). When REFRAG_MODE=1, also includes a compact mini vector (MINI_VECTOR_NAME). - When PATTERN_VECTORS=1, also includes pattern_vector for structural similarity. """ mode = _normalize_schema_mode(schema_mode) if mode != "legacy": @@ -366,22 +367,6 @@ def ensure_collection( distance=models.Distance.COSINE, ) - # Check for pattern vector - try: - pattern_on = os.environ.get("PATTERN_VECTORS", "").strip().lower() in { - "1", "true", "yes", "on", - } - has_pattern = PATTERN_VECTOR_NAME in cfg - except Exception: - pattern_on = False - has_pattern = False - - if pattern_on and not has_pattern: - missing[PATTERN_VECTOR_NAME] = models.VectorParams( - size=PATTERN_VECTOR_DIM, - distance=models.Distance.COSINE, - ) - if missing: try: update_cfg = _prepare_vector_update_config(missing) @@ -416,15 +401,6 @@ def ensure_collection( ) except Exception: pass - try: - if os.environ.get("PATTERN_VECTORS", "").strip().lower() in {"1", "true", "yes", "on"}: - vectors_cfg[PATTERN_VECTOR_NAME] = models.VectorParams( - size=PATTERN_VECTOR_DIM, - distance=models.Distance.COSINE, - ) - except Exception: - pass - sparse_cfg = None if LEX_SPARSE_MODE: sparse_cfg = { @@ -535,6 +511,9 @@ def recreate_collection(client: QdrantClient, name: str, dim: int, vector_name: if not name: print("[BUG] recreate_collection called with name=None! Fix the caller - collection name is required.", flush=True) return + ENSURED_COLLECTIONS.discard(name) + ENSURED_COLLECTIONS_LAST_CHECK.pop(name, None) + ENSURED_PAYLOAD_INDEX_COLLECTIONS.discard(name) try: client.delete_collection(name) except Exception: @@ -553,14 +532,6 @@ def recreate_collection(client: QdrantClient, name: str, dim: int, vector_name: ) except Exception: pass - try: - if os.environ.get("PATTERN_VECTORS", "").strip().lower() in {"1", "true", "yes", "on"}: - vectors_cfg[PATTERN_VECTOR_NAME] = models.VectorParams( - size=PATTERN_VECTOR_DIM, - distance=models.Distance.COSINE, - ) - except Exception: - pass sparse_cfg = None if LEX_SPARSE_MODE: sparse_cfg = { @@ -580,6 +551,20 @@ def recreate_collection(client: QdrantClient, name: str, dim: int, vector_name: def ensure_payload_indexes(client: QdrantClient, collection: str): """Create helpful payload indexes if they don't exist (idempotent).""" + if not collection: + return + + # On memo hit, verify collection still exists and indexes are present + if collection in ENSURED_PAYLOAD_INDEX_COLLECTIONS: + try: + info = client.get_collection(collection) + if not _missing_payload_indexes(info): + # Memo is still valid + return + except Exception: + # Collection doesn't exist or error accessing it; remove from memo + ENSURED_PAYLOAD_INDEX_COLLECTIONS.discard(collection) + for field in PAYLOAD_INDEX_FIELDS: try: client.create_payload_index( @@ -589,6 +574,15 @@ def ensure_payload_indexes(client: QdrantClient, collection: str): ) except Exception: pass + try: + info = client.get_collection(collection) + except Exception: + return + if _missing_payload_indexes(info): + # Do not memoize; a later call should retry. + return + # Even if create_payload_index threw, get_collection confirms indexes exist. + ENSURED_PAYLOAD_INDEX_COLLECTIONS.add(collection) def ensure_collection_and_indexes_once( @@ -629,6 +623,10 @@ def ensure_collection_and_indexes_once( ENSURED_COLLECTIONS_LAST_CHECK.pop(collection, None) except Exception: pass + try: + ENSURED_PAYLOAD_INDEX_COLLECTIONS.discard(collection) + except Exception: + pass ensure_collection(client, collection, dim, vector_name, schema_mode=mode) if mode in {"legacy", "migrate"}: ensure_payload_indexes(client, collection) diff --git a/scripts/ingest/symbols.py b/scripts/ingest/symbols.py index c4f591b0..cb9966d8 100644 --- a/scripts/ingest/symbols.py +++ b/scripts/ingest/symbols.py @@ -1223,7 +1223,7 @@ def extract_symbols_with_tree_sitter(file_path: str) -> dict: """Extract functions, classes, methods from file using tree-sitter or fallback. Returns: - dict: {symbol_id: {name, type, start_line, end_line, content_hash, pseudo, tags}} + dict: {symbol_id: {name, type, path, start_line, end_line, content_hash, pseudo, tags}} """ from scripts.ingest.pipeline import detect_language @@ -1248,6 +1248,8 @@ def extract_symbols_with_tree_sitter(file_path: str) -> dict: symbols[symbol_id] = { "name": sym["name"], "type": sym["kind"], + "path": sym.get("path") or sym.get("name") or "", + "symbol_path": sym.get("path") or sym.get("name") or "", "start_line": sym["start"], "end_line": sym["end"], "content_hash": content_hash, diff --git a/scripts/ingest/vectors.py b/scripts/ingest/vectors.py index a4905389..8417c60d 100644 --- a/scripts/ingest/vectors.py +++ b/scripts/ingest/vectors.py @@ -95,37 +95,3 @@ def _lex_hash_vector(text: str, dim: int = LEX_VECTOR_DIM) -> list[float]: norm = math.sqrt(sum(v * v for v in vec)) or 1.0 return [v / norm for v in vec] - -# --------------------------------------------------------------------------- -# Pattern vector extraction (structural similarity) -# --------------------------------------------------------------------------- -_PATTERN_EXTRACTOR = None -_PATTERN_ENCODER = None - - -def _get_pattern_tools(): - """Lazy load pattern extraction tools.""" - global _PATTERN_EXTRACTOR, _PATTERN_ENCODER - if _PATTERN_EXTRACTOR is None: - try: - from scripts.pattern_detection import PatternExtractor, PatternEncoder - _PATTERN_EXTRACTOR = PatternExtractor() - _PATTERN_ENCODER = PatternEncoder() - except ImportError: - pass - return _PATTERN_EXTRACTOR, _PATTERN_ENCODER - - -def extract_pattern_vector(code: str, language: str) -> list[float] | None: - """Extract structural pattern vector from code. - - Returns 64-dim pattern vector or None if pattern detection unavailable. - """ - extractor, encoder = _get_pattern_tools() - if extractor is None or encoder is None: - return None - try: - signature = extractor.extract(code, language) - return encoder.encode(signature) - except Exception: - return None diff --git a/scripts/ingest_code.py b/scripts/ingest_code.py index 2da24911..1d6363b5 100644 --- a/scripts/ingest_code.py +++ b/scripts/ingest_code.py @@ -24,18 +24,12 @@ from __future__ import annotations import os -import sys import hashlib import time from pathlib import Path from datetime import datetime from typing import List, Dict, Any, Optional, TYPE_CHECKING -# Ensure project root is on sys.path when run as a script -ROOT_DIR = Path(__file__).resolve().parent.parent -if str(ROOT_DIR) not in sys.path: - sys.path.insert(0, str(ROOT_DIR)) - from qdrant_client import QdrantClient, models # --------------------------------------------------------------------------- @@ -203,6 +197,7 @@ from scripts.ingest.pipeline import ( _detect_repo_name_from_path, + is_text_like_language, detect_language, build_information, pseudo_backfill_tick, @@ -212,6 +207,15 @@ index_repo, process_file_with_smart_reindexing, ) + +# --------------------------------------------------------------------------- +# Graph edges (optional accelerator) +# --------------------------------------------------------------------------- +from scripts.ingest.graph_edges import ( + graph_edges_backfill_tick, + delete_edges_by_path as delete_graph_edges_by_path, + upsert_file_edges as upsert_graph_edges_for_file, +) # --------------------------------------------------------------------------- # Re-exports from ingest/cli.py # --------------------------------------------------------------------------- @@ -222,11 +226,9 @@ # --------------------------------------------------------------------------- # Additional imports for backward compatibility # --------------------------------------------------------------------------- -try: - from scripts.embedder import get_embedding_model as _get_embedding_model - _EMBEDDER_FACTORY = True -except ImportError: - _EMBEDDER_FACTORY = False +from scripts.embedder import get_embedding_model as _get_embedding_model + +_EMBEDDER_FACTORY = True if TYPE_CHECKING: from fastembed import TextEmbedding @@ -240,11 +242,9 @@ from scripts.utils import lex_hash_vector_text as _lex_hash_vector_text from scripts.utils import lex_sparse_vector_text as _lex_sparse_vector_text -try: - from scripts.ast_analyzer import get_ast_analyzer, chunk_code_semantically - _AST_ANALYZER_AVAILABLE = True -except ImportError: - _AST_ANALYZER_AVAILABLE = False +from scripts.ast_analyzer import get_ast_analyzer, chunk_code_semantically + +_AST_ANALYZER_AVAILABLE = True try: from tqdm import tqdm @@ -332,12 +332,17 @@ def main(): "embed_batch", # Pipeline "_detect_repo_name_from_path", + "is_text_like_language", "detect_language", "build_information", "index_single_file", "index_repo", "process_file_with_smart_reindexing", "pseudo_backfill_tick", + # Graph edges (optional) + "graph_edges_backfill_tick", + "delete_graph_edges_by_path", + "upsert_graph_edges_for_file", # CLI "main", # Backward compat diff --git a/scripts/ingest_history.py b/scripts/ingest_history.py index 3f42715c..522e10ff 100644 --- a/scripts/ingest_history.py +++ b/scripts/ingest_history.py @@ -4,12 +4,11 @@ import subprocess import shlex import hashlib +import logging from typing import List, Dict, Any import re import time import json -import sys -from pathlib import Path from qdrant_client import QdrantClient, models @@ -19,22 +18,13 @@ API_KEY = os.environ.get("QDRANT_API_KEY") REPO_NAME = os.environ.get("REPO_NAME", "workspace") -ROOT_DIR = Path(__file__).resolve().parent.parent -if str(ROOT_DIR) not in sys.path: - sys.path.insert(0, str(ROOT_DIR)) - -# Import TextEmbedding for type hints and fallback -from fastembed import TextEmbedding - -# Use embedder factory for Qwen3 support; fallback to direct fastembed -try: - from scripts.embedder import get_embedding_model as _get_embedding_model - _EMBEDDER_FACTORY = True -except ImportError: - _EMBEDDER_FACTORY = False +from scripts.embedder import get_embedding_model as _get_embedding_model from scripts.utils import sanitize_vector_name as _sanitize_vector_name +logger = logging.getLogger(__name__) +logger.addHandler(logging.NullHandler()) + def _manifest_run_id(manifest_path: str) -> str: try: @@ -365,30 +355,60 @@ def _ingest_from_manifest( vec_name: str, include_body: bool, per_batch: int, -) -> int: +) -> tuple[int, bool]: try: with open(manifest_path, "r", encoding="utf-8") as f: data = json.load(f) except Exception as e: print(f"Failed to read manifest {manifest_path}: {e}") - return 0 + return 0, False commits = data.get("commits") or [] if not commits: print("No commits in manifest.") - return 0 + return 0, False run_id = _manifest_run_id(manifest_path) mode = str(data.get("mode") or "delta").strip().lower() or "delta" points: List[models.PointStruct] = [] - count = 0 - for c in commits: + total_commits = len(commits) + prepared_count = 0 + persisted_count = 0 + invalid_commit_records = 0 + embed_failures = 0 + point_build_failures = 0 + upsert_failures = 0 + processed_count = 0 + progress_step = max(1, total_commits // 10) if total_commits > 0 else 1 + + def _log_progress(force: bool = False) -> None: + if not force and processed_count % progress_step != 0: + return + logger.info( + "[ingest_history] progress run_id=%s processed=%d/%d prepared=%d persisted=%d invalid=%d embed_failures=%d point_failures=%d upsert_failures=%d", + run_id, + processed_count, + total_commits, + prepared_count, + persisted_count, + invalid_commit_records, + embed_failures, + point_build_failures, + upsert_failures, + ) + + for idx, c in enumerate(commits, start=1): + processed_count += 1 try: if not isinstance(c, dict): + invalid_commit_records += 1 + _log_progress() continue commit_id = str(c.get("commit_id") or "").strip() if not commit_id: + invalid_commit_records += 1 + _log_progress() continue author_name = str(c.get("author_name") or "") authored_date = str(c.get("authored_date") or "") @@ -406,7 +426,15 @@ def _ingest_from_manifest( text = build_text(md, include_body=include_body) try: vec = next(model.embed([text])).tolist() - except Exception: + except Exception as e: + embed_failures += 1 + logger.warning( + "[ingest_history] embed failed for commit=%s idx=%d: %s", + commit_id, + idx, + e, + ) + _log_progress() continue goal: str = "" @@ -451,28 +479,96 @@ def _ingest_from_manifest( pid = stable_id(commit_id) pt = models.PointStruct(id=pid, vector={vec_name: vec}, payload=payload) points.append(pt) - count += 1 + prepared_count += 1 if len(points) >= per_batch: - client.upsert(collection_name=COLLECTION, points=points) - points.clear() + batch_size = len(points) + try: + client.upsert(collection_name=COLLECTION, points=points) + persisted_count += batch_size + except Exception as e: + upsert_failures += batch_size + logger.exception( + "[ingest_history] upsert batch failed (size=%d): %s", + batch_size, + e, + ) + finally: + points.clear() + _log_progress() except Exception: + point_build_failures += 1 + logger.warning( + "[ingest_history] commit processing failed idx=%d", + idx, + exc_info=True, + ) + _log_progress() continue if points: - client.upsert(collection_name=COLLECTION, points=points) - try: - _prune_old_commit_points(client, run_id, mode=mode) - except Exception: - pass - try: - _cleanup_manifest_files(manifest_path) - except Exception: - pass - print(f"Ingested {count} commits into {COLLECTION} from manifest {manifest_path}.") - return count + batch_size = len(points) + try: + client.upsert(collection_name=COLLECTION, points=points) + persisted_count += batch_size + except Exception as e: + upsert_failures += batch_size + logger.exception( + "[ingest_history] final upsert failed (size=%d): %s", + batch_size, + e, + ) + _log_progress(force=True) + ingest_successful = ( + prepared_count > 0 + and invalid_commit_records == 0 + and embed_failures == 0 + and point_build_failures == 0 + and upsert_failures == 0 + and persisted_count == prepared_count + ) + # Only prune snapshot runs that completed cleanly + prune_safe = mode == "snapshot" and ingest_successful + if prune_safe: + try: + _prune_old_commit_points(client, run_id, mode=mode) + except Exception as e: + logger.warning("[ingest_history] prune failed for run_id=%s: %s", run_id, e) + elif mode == "snapshot": + logger.warning( + "[ingest_history] skipping prune for run_id=%s because the snapshot ingest was incomplete", + run_id, + ) + + # Only cleanup manifest if ingest completed successfully + ingest_complete = ingest_successful + if ingest_complete: + try: + _cleanup_manifest_files(manifest_path) + except Exception as e: + logger.warning("[ingest_history] manifest cleanup failed for %s: %s", manifest_path, e) + else: + logger.warning( + "[ingest_history] keeping manifest %s because ingest was incomplete", + manifest_path, + ) + + logger.info( + "Ingested commits from manifest %s into %s: persisted=%d prepared=%d invalid=%d " + "embed_failures=%d point_failures=%d upsert_failures=%d", + manifest_path, + COLLECTION, + persisted_count, + prepared_count, + invalid_commit_records, + embed_failures, + point_build_failures, + upsert_failures, + ) + return persisted_count, ingest_complete def main(): + logging.basicConfig(level=logging.INFO) ap = argparse.ArgumentParser( description="Ingest Git history into Qdrant deterministically" ) @@ -512,16 +608,12 @@ def main(): ) args = ap.parse_args() - # Use embedder factory for Qwen3 support - if _EMBEDDER_FACTORY: - model = _get_embedding_model(MODEL_NAME) - else: - model = TextEmbedding(model_name=MODEL_NAME) + model = _get_embedding_model(MODEL_NAME) vec_name = _sanitize_vector_name(MODEL_NAME) client = QdrantClient(url=QDRANT_URL, api_key=API_KEY or None) if args.manifest_json: - _ingest_from_manifest( + persisted_count, ingest_complete = _ingest_from_manifest( args.manifest_json, model, client, @@ -529,6 +621,8 @@ def main(): args.include_body, args.per_batch, ) + if not ingest_complete: + raise SystemExit(1) return commits = list_commits(args) @@ -537,6 +631,8 @@ def main(): return points: List[models.PointStruct] = [] + persisted_count = 0 + upsert_failures = 0 for sha in commits: md = commit_metadata(sha) text = build_text(md, include_body=args.include_body) @@ -583,11 +679,40 @@ def main(): point = models.PointStruct(id=pid, vector={vec_name: vec}, payload=payload) points.append(point) if len(points) >= args.per_batch: - client.upsert(collection_name=COLLECTION, points=points) - points.clear() + batch_size = len(points) + try: + client.upsert(collection_name=COLLECTION, points=points) + persisted_count += batch_size + except Exception as e: + upsert_failures += batch_size + logger.exception( + "[ingest_history] batch upsert failed collection=%s repo=%s size=%d path=%s: %s", + COLLECTION, + REPO_NAME, + batch_size, + args.path or "", + e, + ) + finally: + points.clear() if points: - client.upsert(collection_name=COLLECTION, points=points) - print(f"Ingested {len(commits)} commits into {COLLECTION}.") + final_size = len(points) + try: + client.upsert(collection_name=COLLECTION, points=points) + persisted_count += final_size + except Exception as e: + upsert_failures += final_size + logger.exception( + "[ingest_history] final upsert failed collection=%s repo=%s size=%d path=%s: %s", + COLLECTION, + REPO_NAME, + final_size, + args.path or "", + e, + ) + if upsert_failures: + raise SystemExit(1) + print(f"Ingested {persisted_count} commits into {COLLECTION}.") if __name__ == "__main__": diff --git a/scripts/k8s_uploader.py b/scripts/k8s_uploader.py index 4f2947c3..03e6d1a8 100755 --- a/scripts/k8s_uploader.py +++ b/scripts/k8s_uploader.py @@ -173,8 +173,6 @@ def trigger_indexing( # Build Python command to call qdrant_index via MCP server # Use qdrant_index with subdir parameter to index specific repo python_cmd = f""" -import sys -sys.path.insert(0, '/app') from scripts.mcp_indexer_server import qdrant_index import asyncio import json @@ -309,4 +307,3 @@ def main(): if __name__ == "__main__": main() - diff --git a/scripts/learning_reranker_worker.py b/scripts/learning_reranker_worker.py deleted file mode 100644 index 39cff7e5..00000000 --- a/scripts/learning_reranker_worker.py +++ /dev/null @@ -1,683 +0,0 @@ -#!/usr/bin/env python3 -""" -Background Learning Reranker Worker. - -Consumes training events from the event log and updates weights per-collection. -This keeps the MCP serving path fast and deterministic. - -Features: -- Reads events from NDJSON log files (one per collection) -- Batches teacher scoring to amortize ONNX overhead -- Updates weights atomically (write to .tmp, rename) -- Supports multiple collections with isolated weights -- Can run as a daemon or one-shot - -Usage: - # Run continuously (daemon mode) - python scripts/learning_reranker_worker.py --daemon - - # Process pending events once and exit - python scripts/learning_reranker_worker.py --once - - # Process specific collection - python scripts/learning_reranker_worker.py --collection my-repo -""" - -import argparse -import json -import os -import sys -import time -from pathlib import Path -from typing import Any, Dict, List, Optional, Tuple - -import numpy as np - -# Add project root to path -ROOT_DIR = Path(__file__).resolve().parent.parent -if str(ROOT_DIR) not in sys.path: - sys.path.insert(0, str(ROOT_DIR)) - -from scripts.rerank_events import ( - RERANK_EVENTS_DIR, - RERANK_EVENTS_RETENTION_DAYS, - read_events, - list_event_files, - cleanup_old_events, -) -from scripts.rerank_recursive import ( - TinyScorer, - LatentRefiner, - RecursiveReranker, - VICReg, - LearnedProjection, - LearnedHybridWeights, - QueryExpander, -) - -# Configuration -BATCH_SIZE = int(os.environ.get("RERANK_LEARNING_BATCH_SIZE", "32")) -POLL_INTERVAL = float(os.environ.get("RERANK_LEARNING_POLL_INTERVAL", "30")) -LEARNING_RATE = float(os.environ.get("RERANK_LEARNING_RATE", "0.001")) - -# VICReg configuration -VICREG_WEIGHT = float(os.environ.get("RERANK_VICREG_WEIGHT", "0.1")) -VICREG_MIN_BATCH = int(os.environ.get("RERANK_VICREG_MIN_BATCH", "4")) # Need 4+ samples for covariance - -# GLM teacher configuration (online learning with LLM judgments) -LLM_TEACHER_ENABLED = os.environ.get("RERANK_LLM_TEACHER", "0") == "1" -LLM_TEACHER_SAMPLE_RATE = float(os.environ.get("RERANK_LLM_SAMPLE_RATE", "1.0")) # 100% when enabled (background anyway) - - -def get_logger(): - """Get logger for worker.""" - try: - from scripts.logger import get_logger as _get_logger - return _get_logger(__name__) - except Exception: - import logging - logging.basicConfig(level=logging.INFO) - return logging.getLogger(__name__) - - -logger = get_logger() - - -class CollectionLearner: - """Handles learning for a single collection.""" - - def __init__(self, collection: str): - self.collection = collection - self._lock_file = None - - # Acquire exclusive lock to prevent multiple workers on same collection - self._acquire_lock() - - self.scorer = TinyScorer(lr=LEARNING_RATE) - self.scorer.set_collection(collection) - self.refiner = LatentRefiner(dim=self.scorer.dim, lr=LEARNING_RATE) - self.refiner.set_collection(collection) - - # Learned projection: raw embedding dim → working dim (256) - # Lower LR than scorer/refiner - projection is more sensitive - from scripts.embedder import get_model_dimension - embed_dim = get_model_dimension() # Respects EMBEDDING_MODEL env - self.projection = LearnedProjection( - input_dim=embed_dim, - output_dim=self.scorer.dim, - lr=LEARNING_RATE * 0.5, # Half the learning rate - ) - self.projection.set_collection(collection) - - self._last_processed_ts = self._load_checkpoint() - - # Reuse the serving reranker's embed + project code path 1:1. - # We only use its private feature helpers, not its scoring loop. - self._feature_reranker = RecursiveReranker( - n_iterations=1, - dim=self.scorer.dim, - early_stop=False, - blend_with_initial=0.0, - ) - - # VICReg for residual regularization (prevents collapse, decorrelates) - self.vicreg = VICReg( - lambda_var=1.0, - lambda_cov=0.04, - lambda_inv=0.1, - ) if VICREG_WEIGHT > 0 else None - - # Learned hybrid weights: dense vs. lexical balance - self.hybrid_weights = LearnedHybridWeights(lr=0.01) - self.hybrid_weights.set_collection(collection) - - # Query expander: learns synonyms/related terms from usage - self.query_expander = QueryExpander(lr=0.1) - self.query_expander.set_collection(collection) - - # LLM teacher for higher-quality supervision (optional) - # Supports llama.cpp or GLM API - auto-detects based on env - self._llm_client = None - self._llm_runtime = None - self._llm_calls = 0 - if LLM_TEACHER_ENABLED: - try: - # Auto-detect runtime: GLM_API_KEY -> glm, else -> llamacpp - runtime = os.environ.get("REFRAG_RUNTIME", "").strip().lower() - if not runtime: - if os.environ.get("GLM_API_KEY", "").strip(): - runtime = "glm" - else: - runtime = "llamacpp" - - if runtime == "glm": - from scripts.refrag_glm import GLMRefragClient - self._llm_client = GLMRefragClient() - else: - from scripts.refrag_llamacpp import LlamaCppRefragClient, is_decoder_enabled - if is_decoder_enabled(): - self._llm_client = LlamaCppRefragClient() - else: - logger.info(f"[{collection}] LLM teacher skipped (decoder disabled)") - - if self._llm_client: - self._llm_runtime = runtime - logger.info(f"[{collection}] LLM teacher enabled ({runtime}, sample_rate={LLM_TEACHER_SAMPLE_RATE})") - except Exception as e: - logger.warning(f"[{collection}] LLM teacher unavailable: {e}") - - # Metrics tracking for logging - self._vicreg_loss_sum = 0.0 - self._vicreg_count = 0 - self._proj_grad_norm_sum = 0.0 - self._proj_grad_count = 0 - self._hybrid_updates = 0 - self._expander_updates = 0 - - @staticmethod - def _sanitize_collection(collection: str) -> str: - """Sanitize collection name to prevent path traversal.""" - return "".join(c if c.isalnum() or c in "-_" else "_" for c in collection) - - def _acquire_lock(self): - """Acquire exclusive lock to prevent multiple workers on same collection.""" - import fcntl - safe_name = self._sanitize_collection(self.collection) - lock_path = Path(TinyScorer.WEIGHTS_DIR) / f"{safe_name}.lock" - lock_path.parent.mkdir(parents=True, exist_ok=True) - self._lock_file = open(lock_path, "w") - try: - fcntl.flock(self._lock_file, fcntl.LOCK_EX | fcntl.LOCK_NB) - logger.debug(f"[{self.collection}] Acquired exclusive lock") - except OSError: - logger.error(f"[{self.collection}] Another worker already running. Exiting.") - sys.exit(1) - - def _release_lock(self): - """Release the exclusive lock.""" - if self._lock_file: - import fcntl - try: - fcntl.flock(self._lock_file, fcntl.LOCK_UN) - self._lock_file.close() - except Exception: - pass - self._lock_file = None - - def _load_checkpoint(self) -> float: - """Load last processed timestamp from checkpoint file.""" - safe_name = self._sanitize_collection(self.collection) - checkpoint_path = Path(TinyScorer.WEIGHTS_DIR) / f"checkpoint_{safe_name}.json" - try: - if checkpoint_path.exists(): - with open(checkpoint_path) as f: - return json.load(f).get("last_ts", 0) - except Exception as e: - logger.warning(f"[{self.collection}] Failed to load checkpoint: {e}") - return 0 - - def _save_checkpoint(self, ts: float): - """Save last processed timestamp to checkpoint file atomically.""" - safe_name = self._sanitize_collection(self.collection) - checkpoint_path = Path(TinyScorer.WEIGHTS_DIR) / f"checkpoint_{safe_name}.json" - tmp_path = checkpoint_path.with_suffix(".json.tmp") - try: - checkpoint_path.parent.mkdir(parents=True, exist_ok=True) - with open(tmp_path, "w") as f: - json.dump({"last_ts": ts, "collection": self.collection}, f) - os.replace(tmp_path, checkpoint_path) - except Exception as e: - logger.warning(f"[{self.collection}] Failed to save checkpoint: {e}") - if tmp_path.exists(): - try: - tmp_path.unlink() - except Exception: - pass - - def _encode(self, texts: List[str]) -> np.ndarray: - """Encode texts to BGE embeddings (768-dim, raw without projection).""" - return self._feature_reranker._encode_raw(texts) - - def _encode_project(self, texts: List[str]) -> np.ndarray: - """Encode + project via learned projection (for inference/serving compat).""" - embs = self._encode(texts) - return self.projection.forward(embs) - - def _encode_project_with_cache( - self, texts: List[str] - ) -> Tuple[np.ndarray, Dict[str, Any]]: - """Encode + project with cache for gradient backprop through projection. - - Returns: - projected: (batch, output_dim) projected embeddings - cache: dict with raw_embs and projection cache for backward pass - """ - raw_embs = self._encode(texts) # (batch, 768) - projected, proj_cache = self.projection.forward_with_cache(raw_embs) - cache = {"raw_embs": raw_embs, "proj_cache": proj_cache} - return projected, cache - - @staticmethod - def _pack_doc(candidate: Dict[str, Any], max_chars: int = 500) -> str: - """Pack candidate into doc text (shared by learning and teacher scoring).""" - parts = [] - if candidate.get("symbol"): - parts.append(str(candidate["symbol"])) - if candidate.get("path"): - parts.append(str(candidate["path"])) - code = candidate.get("code") or candidate.get("snippet") or candidate.get("text") or "" - if code: - parts.append(str(code)[:max_chars]) - return " ".join(parts) if parts else "empty" - - def _llm_judge(self, query: str, candidates: List[Dict[str, Any]], top_k: int = 5) -> Optional[np.ndarray]: - """Get LLM relevance judgments for top candidates. - - Works with llama.cpp or GLM API (auto-detected at init). - - Returns: - scores: (n_candidates,) array with LLM-derived scores, or None if failed - """ - if not self._llm_client or not candidates: - return None - - # Only judge top-k to save API cost - n = len(candidates) - top_k = min(top_k, n) - - # Build prompt for GLM to rate relevance - prompt_parts = [f"Rate code search relevance 0-10 for query: \"{query}\"\n"] - for i in range(top_k): - c = candidates[i] - doc = self._pack_doc(c, max_chars=800) - prompt_parts.append(f"[{i}] {doc}") - - prompt_parts.append( - "\nRespond with JSON: {\"scores\": [N, N, ...]} where each N is 0-10 relevance." - ) - prompt = "\n".join(prompt_parts) - - try: - import json - response = self._llm_client.generate_with_soft_embeddings( - prompt=prompt, - max_tokens=64, - temperature=0.1, - force_json=True, - disable_thinking=True, # Fast mode - ) - data = json.loads(response) - llm_scores = data.get("scores", []) - - if len(llm_scores) >= top_k: - # Normalize to 0-1 and extend to all candidates - scores = np.zeros(n, dtype=np.float32) - for i in range(top_k): - scores[i] = float(llm_scores[i]) / 10.0 - # Lower candidates get decaying scores - for i in range(top_k, n): - scores[i] = max(0, scores[top_k - 1] * 0.5 ** (i - top_k + 1)) - - self._llm_calls += 1 - return scores - except Exception as e: - logger.debug(f"[{self.collection}] LLM judge failed: {e}") - - return None - - def process_events(self, limit: int = 1000) -> int: - """Process pending events and return count processed.""" - events = read_events(self.collection, since_ts=self._last_processed_ts, limit=limit) - if not events: - return 0 - - processed = 0 - batch_events = [] - - for event in events: - batch_events.append(event) - if len(batch_events) >= BATCH_SIZE: - self._learn_from_batch(batch_events) - processed += len(batch_events) - batch_events = [] - - # Process remaining - if batch_events: - self._learn_from_batch(batch_events) - processed += len(batch_events) - - # Update checkpoint - if events: - self._last_processed_ts = max(e.get("ts", 0) for e in events) - self._save_checkpoint(self._last_processed_ts) - - # Save weights after processing batch (both scorer and refiner) - if processed > 0: - self.scorer._save_weights(checkpoint=True) - self.refiner._save_weights(checkpoint=True) - metrics = self.scorer.get_metrics() - - # VICReg + projection info - extra_info = "" - if self._vicreg_count > 0: - avg_vicreg = self._vicreg_loss_sum / self._vicreg_count - extra_info += f" | vicreg={avg_vicreg:.4f}" - self._vicreg_loss_sum = 0.0 - self._vicreg_count = 0 - - if self._proj_grad_count > 0: - avg_proj_grad = self._proj_grad_norm_sum / self._proj_grad_count - extra_info += f" | proj_grad={avg_proj_grad:.4f}" - self._proj_grad_norm_sum = 0.0 - self._proj_grad_count = 0 - - # Hybrid weight, expander, and GLM stats - hw_info = f" dense_w={self.hybrid_weights.dense_weight:.2f}" - exp_stats = self.query_expander.get_stats() - exp_info = f" terms={exp_stats['terms']}" - llm_info = f" llm={self._llm_calls}({self._llm_runtime})" if self._llm_client else "" - - logger.info( - f"[{self.collection}] Processed {processed} events | " - f"scorer_v{metrics['version']} refiner_v{self.refiner._version} proj_v{self.projection._version}{hw_info}{exp_info}{llm_info} | " - f"lr={metrics['learning_rate']:.6f} | " - f"avg_loss={metrics['avg_loss']:.4f}{extra_info} | converged={metrics['converged']}" - ) - - return processed - - def _learn_from_batch(self, events: List[Dict[str, Any]]): - """Learn from batch with 3-pass deep supervision + VICReg + projection learning. - - Full end-to-end learning: - - Projection: BGE (768) → working dim (256), learns domain-specific subspace - - Scorer: learns to rank with current z - - Refiner: learns to improve z toward teacher-optimal state - - VICReg: regularizes residuals to prevent collapse - - Deep Supervision (TRM-style): - - Each refinement pass gets a loss signal toward teacher_z - - Later passes have decaying weight (pass 1: 1.0, pass 2: 0.7, pass 3: 0.5) - """ - # Fill missing teacher scores in batch (amortizes ONNX overhead) - self._maybe_fill_teacher_scores(events) - - # Accumulate data for end-of-batch updates - vicreg_data: List[tuple] = [] # [(z, z_refined, refiner_cache), ...] - projection_grads: List[Tuple[np.ndarray, Dict[str, Any]]] = [] # [(grad, proj_cache), ...] - - # Deep supervision weights: earlier passes contribute more - pass_weights = [1.0, 0.7, 0.5] - n_passes = 3 - - for event in events: - try: - query = event.get("query", "") - candidates = event.get("candidates", []) - teacher_scores = event.get("teacher_scores") - - if not query or not candidates or not teacher_scores: - continue - - # Validate alignment between candidates and teacher scores - if len(teacher_scores) != len(candidates): - logger.warning(f"[{self.collection}] Skipping event: mismatched lengths") - continue - - # Build doc texts from candidates - doc_texts = [self._pack_doc(c) for c in candidates] - teacher_arr = np.array(teacher_scores, dtype=np.float32) - - # ===== LLM TEACHER: Higher-quality supervision (sampled) ===== - if self._llm_client and np.random.random() < LLM_TEACHER_SAMPLE_RATE: - llm_scores = self._llm_judge(query, candidates, top_k=5) - if llm_scores is not None: - # Blend LLM with ONNX: LLM is more reliable, weight it higher - teacher_arr = 0.3 * teacher_arr + 0.7 * llm_scores - - # ===== ENCODE WITH CACHE FOR PROJECTION LEARNING ===== - # Query embedding with cache for gradient backprop - query_proj, query_proj_cache = self._encode_project_with_cache([query]) - query_emb = query_proj[0] - - # Doc embeddings with cache for gradient backprop - doc_embs, doc_proj_cache = self._encode_project_with_cache(doc_texts) - - # Compute teacher-weighted document summary as target z - teacher_weights = np.exp(teacher_arr - teacher_arr.max()) - teacher_weights = teacher_weights / (teacher_weights.sum() + 1e-8) - teacher_z = (teacher_weights[:, None] * doc_embs).sum(axis=0) - teacher_z = teacher_z / (np.linalg.norm(teacher_z) + 1e-8) - - # Initialize latent state from query - z = query_emb.copy() - - # ===== PROJECTION GRADIENT: Contrastive alignment loss ===== - # Goal: projection should produce embeddings where query is close to - # high-scoring docs and far from low-scoring docs - # Gradient w.r.t. query: weighted by teacher scores - # Push query toward high-teacher-score docs - query_grad = (teacher_weights[:, None] * doc_embs).sum(axis=0) - query_emb - query_grad = query_grad / (np.linalg.norm(query_grad) + 1e-8) - - # Gradient w.r.t. docs: each doc pulled/pushed based on teacher score - # High score docs: pull toward query, Low score docs: push away - centered_weights = teacher_weights - teacher_weights.mean() - doc_grad = centered_weights[:, None] * (query_emb - doc_embs) - - projection_grads.append((query_grad.reshape(1, -1), query_proj_cache)) - projection_grads.append((doc_grad, doc_proj_cache)) - - # ===== LEARN HYBRID WEIGHTS (dense vs. lexical) ===== - # Extract dense/lexical scores if available in candidates - if candidates and "dense_score" in candidates[0] and "lexical_score" in candidates[0]: - dense_scores = np.array([c.get("dense_score", 0) for c in candidates], dtype=np.float32) - lexical_scores = np.array([c.get("lexical_score", 0) for c in candidates], dtype=np.float32) - self.hybrid_weights.learn_from_teacher(dense_scores, lexical_scores, teacher_arr) - self._hybrid_updates += 1 - - # ===== LEARN QUERY EXPANSIONS ===== - # Learn term associations from high-scoring docs - self.query_expander.learn_from_teacher(query, doc_texts, teacher_arr) - self._expander_updates += 1 - - # ===== 3-PASS DEEP SUPERVISION ===== - for pass_idx in range(n_passes): - # Score with current z - scores = self.scorer.forward(query_emb, doc_embs, z) - - # Train scorer at this pass (learns to rank with current z) - self.scorer.learn_from_teacher(query_emb, doc_embs, z, teacher_arr) - - # Train refiner: z → z' toward teacher_z - if self.vicreg is not None: - # Get cache for VICReg backprop - _, z_orig, z_refined, cache = self.refiner.learn_from_teacher_with_cache( - z, query_emb, doc_embs, scores, teacher_z - ) - cache["pass_weight"] = pass_weights[pass_idx] - vicreg_data.append((z_orig, z_refined, cache)) - else: - self.refiner.learn_from_teacher(z, query_emb, doc_embs, scores, teacher_z) - - # Update z for next pass (refinement chain) - z = self.refiner.refine(z, query_emb, doc_embs, scores) - - except Exception as e: - logger.warning(f"[{self.collection}] Error processing event: {e}") - continue - - # ===== PROJECTION LEARNING: batch update ===== - if projection_grads: - try: - for grad, cache in projection_grads: - self.projection.backward(grad, cache["proj_cache"], weight=0.1) - self._proj_grad_norm_sum += np.linalg.norm(grad) - self._proj_grad_count += 1 - except Exception as e: - logger.warning(f"[{self.collection}] Projection update failed: {e}") - - # ===== VICReg: batch-level residual regularization ===== - if self.vicreg is not None and len(vicreg_data) >= VICREG_MIN_BATCH: - try: - z_batch = np.vstack([item[0] for item in vicreg_data]) - z_refined_batch = np.vstack([item[1] for item in vicreg_data]) - - vicreg_loss, vicreg_grad, _ = self.vicreg.forward(z_batch, z_refined_batch) - - self._vicreg_loss_sum += vicreg_loss - self._vicreg_count += 1 - - for i, (_, _, cache) in enumerate(vicreg_data): - pass_weight = cache.get("pass_weight", 1.0) - self.refiner.apply_vicreg_gradient( - vicreg_grad[i], cache, weight=VICREG_WEIGHT * pass_weight - ) - - except Exception as e: - logger.warning(f"[{self.collection}] VICReg failed: {e}") - - def _maybe_fill_teacher_scores(self, events: List[Dict[str, Any]]): - """Compute teacher scores for events that don't already have them.""" - try: - from scripts.rerank_local import rerank_local - except Exception: - rerank_local = None - - if rerank_local is None: - return - - # Gather pairs across events so we can call rerank_local once. - all_pairs: List[tuple] = [] - slices: List[tuple] = [] # (event_index, start, end) - - for event_index, event in enumerate(events): - if event.get("teacher_scores"): - continue - - query = event.get("query", "") - candidates = event.get("candidates", []) - if not query or not candidates: - continue - - start = len(all_pairs) - for c in candidates: - doc_text = self._pack_doc(c) # Same packing as learning path - all_pairs.append((query, doc_text)) - end = len(all_pairs) - if end > start: - slices.append((event_index, start, end)) - - if not slices: - return - - try: - scores = rerank_local(all_pairs) - except Exception as e: - logger.warning(f"[{self.collection}] Teacher scoring failed for {len(all_pairs)} pairs: {e}") - return - - # Map scores back to each event. - for event_index, start, end in slices: - try: - events[event_index]["teacher_scores"] = list(scores[start:end]) - except Exception as e: - logger.warning(f"[{self.collection}] Failed to map teacher scores for event {event_index}: {e}") - continue - - -def discover_collections() -> List[str]: - """Discover collections with pending events.""" - events_dir = Path(RERANK_EVENTS_DIR) - if not events_dir.exists(): - return [] - - import re - - collections: List[str] = [] - seen = set() - for f in events_dir.glob("events_*.ndjson"): - stem = f.stem # events__ - if not stem.startswith("events_"): - continue - - rest = stem[len("events_") :] - # Strip the hour suffix if present (10 digits) - m = re.match(r"^(?P.+)_(?P\d{10})$", rest) - name = m.group("name") if m else rest - - if name and name not in seen: - seen.add(name) - collections.append(name) - - return collections - - -def run_once(): - """Process all pending events once and exit.""" - collections = discover_collections() - if not collections: - logger.info("No collections with pending events") - return - - total = 0 - for coll in collections: - learner = CollectionLearner(coll) - processed = learner.process_events() - total += processed - - logger.info(f"Processed {total} events across {len(collections)} collections") - - -def run_daemon(): - """Run continuously, polling for new events.""" - logger.info(f"Starting learning worker daemon (poll interval: {POLL_INTERVAL}s)") - learners: Dict[str, CollectionLearner] = {} - last_cleanup = 0 - cleanup_interval = 3600 # Cleanup old events hourly - - while True: - try: - collections = discover_collections() - for coll in collections: - if coll not in learners: - learners[coll] = CollectionLearner(coll) - learners[coll].process_events() - - # Periodic cleanup of old event files - now = time.time() - if RERANK_EVENTS_RETENTION_DAYS > 0 and now - last_cleanup > cleanup_interval: - for coll in collections: - deleted = cleanup_old_events(coll, RERANK_EVENTS_RETENTION_DAYS) - if deleted > 0: - logger.info(f"[{coll}] Cleaned up {deleted} old event files") - last_cleanup = now - - except KeyboardInterrupt: - logger.info("Shutting down") - break - except Exception as e: - logger.error(f"Error in daemon loop: {e}") - - time.sleep(POLL_INTERVAL) - - -def main(): - parser = argparse.ArgumentParser(description="Background learning reranker worker") - parser.add_argument("--daemon", action="store_true", help="Run continuously") - parser.add_argument("--once", action="store_true", help="Process once and exit") - parser.add_argument("--collection", type=str, help="Process specific collection only") - args = parser.parse_args() - - if args.collection: - learner = CollectionLearner(args.collection) - processed = learner.process_events() - logger.info(f"Processed {processed} events for {args.collection}") - elif args.daemon: - run_daemon() - else: - run_once() - - -if __name__ == "__main__": - main() diff --git a/scripts/mcp_admin_tools.py b/scripts/mcp_admin_tools.py deleted file mode 100644 index c4ea0378..00000000 --- a/scripts/mcp_admin_tools.py +++ /dev/null @@ -1,4 +0,0 @@ -#!/usr/bin/env python3 -"""Shim for backward compatibility. See scripts/mcp/admin_tools.py""" -from scripts.mcp_impl.admin_tools import * - diff --git a/scripts/mcp_auth.py b/scripts/mcp_auth.py index 2b13c791..2ae3d220 100644 --- a/scripts/mcp_auth.py +++ b/scripts/mcp_auth.py @@ -1,45 +1,13 @@ import os from typing import Any, Dict, Optional -try: - from scripts.logger import ValidationError -except Exception: - - class ValidationError(Exception): - pass - - -try: - from scripts.auth_backend import ( - AUTH_ENABLED as AUTH_ENABLED_AUTH, - ACL_ALLOW_ALL as ACL_ALLOW_ALL_AUTH, - validate_session as _auth_validate_session, - has_collection_access as _has_collection_access, - ) -except Exception as _auth_backend_import_exc: - _AUTH_BACKEND_IMPORT_ERROR = repr(_auth_backend_import_exc) - AUTH_ENABLED_AUTH = ( - str(os.environ.get("CTXCE_AUTH_ENABLED", "0")).strip().lower() in {"1", "true", "yes", "on"} - ) - ACL_ALLOW_ALL_AUTH = ( - str(os.environ.get("CTXCE_ACL_ALLOW_ALL", "0")).strip().lower() in {"1", "true", "yes", "on"} - ) - - def _auth_validate_session(session_id: str): # type: ignore[no-redef] - if AUTH_ENABLED_AUTH: - raise ValidationError( - f"Auth backend unavailable (import failed): {_AUTH_BACKEND_IMPORT_ERROR}" - ) - return None - - def _has_collection_access( - user_id: str, qdrant_collection: str, permission: str = "read" - ) -> bool: # type: ignore[no-redef] - if AUTH_ENABLED_AUTH: - raise ValidationError( - f"Auth backend unavailable (import failed): {_AUTH_BACKEND_IMPORT_ERROR}" - ) - return True +from scripts.logger import ValidationError +from scripts.auth_backend import ( + AUTH_ENABLED as AUTH_ENABLED_AUTH, + ACL_ALLOW_ALL as ACL_ALLOW_ALL_AUTH, + validate_session as _auth_validate_session, + has_collection_access as _has_collection_access, +) ACL_ENFORCE = ( diff --git a/scripts/mcp_code_signals.py b/scripts/mcp_code_signals.py deleted file mode 100644 index 29027b55..00000000 --- a/scripts/mcp_code_signals.py +++ /dev/null @@ -1,4 +0,0 @@ -#!/usr/bin/env python3 -"""Shim for backward compatibility. See scripts/mcp/code_signals.py""" -from scripts.mcp_impl.code_signals import * - diff --git a/scripts/mcp_context_answer.py b/scripts/mcp_context_answer.py deleted file mode 100644 index 2259671b..00000000 --- a/scripts/mcp_context_answer.py +++ /dev/null @@ -1,8 +0,0 @@ -#!/usr/bin/env python3 -""" -Backward-compatibility shim for scripts.mcp_context_answer. - -New location: scripts.mcp.context_answer -""" -from scripts.mcp_impl.context_answer import * # noqa: F401,F403 - diff --git a/scripts/mcp_http_client.py b/scripts/mcp_http_client.py new file mode 100644 index 00000000..5b01a56a --- /dev/null +++ b/scripts/mcp_http_client.py @@ -0,0 +1,153 @@ +"""Small HTTP client for calling MCP tools. + +This is intentionally transport glue, not an intent router. Agent/tool selection +belongs to the MCP client using the exposed tools directly. +""" +from __future__ import annotations + +import json +import time +from typing import Any, Dict, Tuple +from urllib import request + + +def _post_raw(url: str, payload: Dict[str, Any], headers: Dict[str, str], timeout: float = 60.0) -> Tuple[Dict[str, str], bytes]: + req = request.Request(url, method="POST") + for k, v in headers.items(): + req.add_header(k, v) + data = json.dumps(payload).encode("utf-8") + with request.urlopen(req, data=data, timeout=timeout) as resp: + body = resp.read() + hdrs = {k.lower(): v for k, v in resp.headers.items()} + return hdrs, body + + +def _post_raw_retry(url: str, payload: Dict[str, Any], headers: Dict[str, str], timeout: float = 60.0, retries: int = 2, backoff: float = 0.5) -> Tuple[Dict[str, str], bytes]: + last_exc: Exception | None = None + for i in range(max(0, retries) + 1): + try: + return _post_raw(url, payload, headers, timeout=timeout) + except Exception as e: + last_exc = e + if i < retries: + try: + time.sleep(backoff * (2 ** i)) + except Exception: + pass + else: + raise last_exc + raise last_exc or RuntimeError("MCP HTTP request failed") + + +def _parse_stream_or_json(body: bytes) -> Dict[str, Any]: + txt = body.decode("utf-8", errors="ignore") + if "data:" in txt and ("event:" in txt or txt.strip().startswith("data:")): + last = None + for line in txt.splitlines(): + if line.startswith("data:"): + last = line[len("data:"):].strip() + if last: + try: + return json.loads(last) + except Exception: + pass + return json.loads(txt) + + +def _filter_args(d: Dict[str, Any]) -> Dict[str, Any]: + return {k: v for k, v in d.items() if v not in (None, "")} + + +def _mcp_handshake(base_url: str, timeout: float = 30.0) -> Dict[str, str]: + headers = { + "Content-Type": "application/json", + "Accept": "application/json, text/event-stream", + } + init_payload = { + "jsonrpc": "2.0", + "method": "initialize", + "params": { + "protocolVersion": "2024-11-05", + "capabilities": {}, + "clientInfo": {"name": "context-engine-http-client", "version": "0.1.0"}, + }, + "id": 1, + } + hdrs, body = _post_raw_retry(base_url, init_payload, headers, timeout=timeout) + sid = hdrs.get("mcp-session-id") or hdrs.get("Mcp-Session-Id") + if not sid: + try: + j = _parse_stream_or_json(body) + sid = j.get("sessionId") + except Exception: + sid = None + if sid: + headers["Mcp-Session-Id"] = sid + try: + _post_raw_retry(base_url, {"jsonrpc": "2.0", "method": "notifications/initialized"}, headers, timeout=timeout) + except Exception: + pass + return headers + + +def _extract_iserror_text(resp: Dict[str, Any]) -> str | None: + try: + r = resp.get("result") or {} + if isinstance(r, dict) and r.get("isError"): + content = r.get("content") + if isinstance(content, list) and content and isinstance(content[0], dict): + if content[0].get("type") == "text": + return content[0].get("text") + except Exception: + pass + return None + + +def call_tool_http(base_url: str, tool_name: str, args: Dict[str, Any], timeout: float = 120.0) -> Dict[str, Any]: + """Call an MCP tool over streamable HTTP.""" + headers = _mcp_handshake(base_url, timeout=min(timeout, 30.0)) + + def _do_call(arguments: Dict[str, Any]) -> Dict[str, Any]: + payload = { + "jsonrpc": "2.0", + "id": "mcp-http-client-1", + "method": "tools/call", + "params": {"name": tool_name, "arguments": arguments}, + } + _, body = _post_raw_retry(base_url, payload, headers, timeout=timeout) + return _parse_stream_or_json(body) + + args1 = _filter_args(args or {}) + resp = _do_call({"arguments": args1} if tool_name.endswith("_compat") else args1) + + def _get_structured_error(r: Dict[str, Any]) -> str | None: + try: + rr = r.get("result") or {} + sc = rr.get("structuredContent") or {} + rs = sc.get("result") or {} + err = rs.get("error") + if isinstance(err, str): + return err + except Exception: + pass + return None + + msg = _extract_iserror_text(resp) + serr = _get_structured_error(resp) + if msg: + low = msg.lower() + if ("kwargs" in low) and ("field required" in low or "missing" in low): + return _do_call({"kwargs": args1}) + if ("arguments" in low) and ("field required" in low or "missing" in low): + return _do_call({"arguments": args1}) + if (serr and serr.strip().lower() == "query required") and ("query" in args1 or "queries" in args1): + resp4 = _do_call({"kwargs": args1}) + serr2 = _get_structured_error(resp4) + if not (serr2 and serr2.strip().lower() == "query required"): + return resp4 + resp5 = _do_call({"arguments": {"kwargs": args1}}) + serr3 = _get_structured_error(resp5) + if not (serr3 and serr3.strip().lower() == "query required"): + return resp5 + return _do_call({"arguments": args1}) + return resp diff --git a/scripts/mcp_impl/__init__.py b/scripts/mcp_impl/__init__.py index 4a033f9d..2198055d 100644 --- a/scripts/mcp_impl/__init__.py +++ b/scripts/mcp_impl/__init__.py @@ -1,36 +1,9 @@ -""" -MCP (Model Context Protocol) indexer server package. - -This package contains extracted modules from mcp_indexer_server.py: -- utils: Type coercion, JSON parsing, tokenization, env helpers -- toon: TOON output format support -- workspace: Workspace state and collection resolution -- admin_tools: Qdrant admin operations (index, list, status, prune) -- code_signals: Code intent detection -- context_answer: LLM-assisted Q&A with retrieval -- context_search: Blended code + memory search -- query_expand: LLM-assisted query expansion +"""MCP indexer server implementation package. -Usage: - from scripts.mcp_impl import utils, toon, workspace - from scripts.mcp_impl.utils import _coerce_bool, _env_overrides - from scripts.mcp_impl.workspace import _default_collection - from scripts.mcp_impl.context_search import _context_search_impl - from scripts.mcp_impl.query_expand import _expand_query_impl +Submodules are intentionally not imported here. Several helpers pull service +stacks such as Qdrant, FastMCP, or search wiring, and package import should stay +cheap for tests and small utility imports. """ -from scripts.mcp_impl import utils -from scripts.mcp_impl import toon -from scripts.mcp_impl import workspace -from scripts.mcp_impl import admin_tools -from scripts.mcp_impl import code_signals -from scripts.mcp_impl import context_answer -from scripts.mcp_impl import context_search -from scripts.mcp_impl import query_expand -from scripts.mcp_impl import search -from scripts.mcp_impl import info_request -from scripts.mcp_impl import memory -from scripts.mcp_impl import search_specialized -from scripts.mcp_impl import search_history __all__ = [ "utils", @@ -42,9 +15,7 @@ "context_search", "query_expand", "search", - "info_request", "memory", - "search_specialized", + "search_profiles", "search_history", ] - diff --git a/scripts/mcp_impl/admin_tools.py b/scripts/mcp_impl/admin_tools.py index f8f39962..f02003b6 100644 --- a/scripts/mcp_impl/admin_tools.py +++ b/scripts/mcp_impl/admin_tools.py @@ -6,7 +6,6 @@ Contains: - Subprocess runner (_run_async) - Embedding model cache (_get_embedding_model) -- Router cache invalidation (_invalidate_router_scratchpad) - Repo detection (_detect_current_repo) Note: The @mcp.tool() decorated functions remain in mcp_indexer_server.py @@ -22,7 +21,6 @@ # Functions "_run_async", "_get_embedding_model", - "_invalidate_router_scratchpad", "_detect_current_repo", "_collection_map_impl", ] @@ -30,6 +28,7 @@ import asyncio import logging import os +import subprocess import threading from pathlib import Path from typing import Any, Dict, List, Optional, Tuple @@ -45,39 +44,10 @@ def _get_embedding_model(model_name: str): - """Get cached embedding model with optional Qwen3 support. + """Get cached embedding model via the centralized embedder factory.""" + from scripts.embedder import get_embedding_model - Uses the centralized embedder factory if available, with fallback - to direct fastembed initialization for backwards compatibility. - """ - # Try centralized embedder factory first (supports Qwen3 feature flag) - try: - from scripts.embedder import get_embedding_model - return get_embedding_model(model_name) - except ImportError: - pass - - # Fallback to original implementation - try: - from fastembed import TextEmbedding # type: ignore - except Exception: - raise - - m = _EMBED_MODEL_CACHE.get(model_name) - if m is None: - # Double-checked locking to avoid duplicate inits under concurrency - lock = _EMBED_MODEL_LOCKS.setdefault(model_name, threading.Lock()) - with lock: - m = _EMBED_MODEL_CACHE.get(model_name) - if m is None: - m = TextEmbedding(model_name=model_name) - try: - # Warmup with common patterns to optimize internal caches - _ = list(m.embed(["function", "class", "import", "def", "const"])) - except Exception: - pass - _EMBED_MODEL_CACHE[model_name] = m - return m + return get_embedding_model(model_name) # --------------------------------------------------------------------------- @@ -98,21 +68,6 @@ async def _run_async( return await run_subprocess_async(cmd, timeout=timeout, env=env) -# --------------------------------------------------------------------------- -# Router cache invalidation -# --------------------------------------------------------------------------- -def _invalidate_router_scratchpad(workspace_path: str) -> bool: - """Invalidate any cached router scratchpad for the workspace. - - This is called after indexing operations to ensure the router - picks up new/changed code. Returns True if invalidation occurred. - """ - try: - # Clear any in-memory caches that might be stale - return True - except Exception: - return False - # --------------------------------------------------------------------------- # Repo detection @@ -123,8 +78,7 @@ def _detect_current_repo() -> Optional[str]: Priority: 1. CURRENT_REPO env var (explicitly set) 2. REPO_NAME env var - 3. Detect from /work directory structure (first subdirectory with .git) - 4. Git remote origin name + 3. Bindmount git detection when CTXCE_BINDMOUNT_REPO_DETECTION=1 Returns: repo name or None if detection fails """ @@ -134,34 +88,40 @@ def _detect_current_repo() -> Optional[str]: if val: return val - # Try to detect from /work directory + try: + from scripts.workspace_state import bindmount_repo_detection_enabled + + allow_git_detection = bindmount_repo_detection_enabled() + except Exception: + allow_git_detection = False + + if not allow_git_detection: + return None + + # Bindmount detection from /work. Do not guess from invalid/internal + # metadata: a leaked /work/.git must not become repo "work". work_path = Path("/work") if work_path.exists(): try: - # Check for .git in /work itself if (work_path / ".git").exists(): - # Use git to get repo name from remote - try: - import subprocess - result = subprocess.run( - ["git", "-C", str(work_path), "config", "--get", "remote.origin.url"], - capture_output=True, text=True, timeout=5 - ) - if result.returncode == 0 and result.stdout.strip(): - url = result.stdout.strip() - # Extract repo name from URL - name = url.rstrip("/").rsplit("/", 1)[-1] - if name.endswith(".git"): - name = name[:-4] - if name: - return name - except Exception: - pass - # Fallback to directory name - return work_path.name - - # Check subdirectories for repos + result = subprocess.run( + ["git", "-C", str(work_path), "config", "--get", "remote.origin.url"], + capture_output=True, + text=True, + timeout=5, + ) + if result.returncode == 0 and result.stdout.strip(): + url = result.stdout.strip() + name = url.rstrip("/").rsplit("/", 1)[-1] + if name.endswith(".git"): + name = name[:-4] + if name: + return name + + internal_dirs = {".codebase", ".git", "__pycache__"} for subdir in work_path.iterdir(): + if subdir.name in internal_dirs: + continue if subdir.is_dir() and (subdir / ".git").exists(): return subdir.name except Exception: diff --git a/scripts/mcp_impl/context_answer.py b/scripts/mcp_impl/context_answer.py index 54877976..f17b8793 100644 --- a/scripts/mcp_impl/context_answer.py +++ b/scripts/mcp_impl/context_answer.py @@ -50,7 +50,8 @@ _primary_identifier_from_queries, ) from scripts.mcp_impl.workspace import _default_collection -from scripts.logger import safe_int, ValidationError +from scripts.logger import safe_bool, safe_float, safe_int, ValidationError +from scripts.refrag_glm import detect_glm_runtime, get_glm_model_name, get_model_config logger = logging.getLogger(__name__) @@ -114,11 +115,7 @@ def _cleanup_answer(text: str, max_chars: int | None = None) -> str: def _answer_style_guidance() -> str: """Compact instruction to keep answers direct and grounded.""" - try: - from scripts.refrag_glm import detect_glm_runtime - is_glm = detect_glm_runtime() - except ImportError: - is_glm = False + is_glm = detect_glm_runtime() if is_glm: sentence_guidance = "Write a clear, comprehensive answer in 4-8 sentences." @@ -233,11 +230,7 @@ def _answer_style_guidance() -> str: GLM models get more generous guidance (4-8 sentences) since they handle longer outputs better than Granite-4.0-Micro which needs strict 2-4 sentence limits. """ - try: - from scripts.refrag_glm import detect_glm_runtime - is_glm = detect_glm_runtime() - except ImportError: - is_glm = False + is_glm = detect_glm_runtime() if is_glm: # GLM models can handle longer, more detailed answers @@ -1013,30 +1006,6 @@ def _ok_lang(it: Dict[str, Any]) -> bool: model=model, repo=repo, # Cross-codebase isolation ) - # Ensure last call reflects tier-2 relaxed filters for introspection/testing - _ = run_hybrid_search( - queries=queries, - limit=int(max(lim, 1)), - per_path=int(max(ppath, 1)), - language=eff_language, - under=override_under or None, - kind=None, - symbol=None, - ext=None, - not_filter=(not_ or kwargs.get("not_") or kwargs.get("not") or None), - case=(case or kwargs.get("case") or None), - path_regex=None, - path_glob=None, - not_glob=eff_not_glob, - expand=False - if did_local_expand - else ( - str(os.environ.get("HYBRID_EXPAND", "0")).strip().lower() - in {"1", "true", "yes", "on"} - ), - model=model, - repo=repo, # Cross-codebase isolation - ) if os.environ.get("DEBUG_CONTEXT_ANSWER"): logger.debug( @@ -1518,58 +1487,56 @@ def _ok_lang(it: Dict[str, Any]) -> bool: # Filter out memory-like items without a valid path to avoid empty citations items = [it for it in items if str(it.get("path") or "").strip()] - # Apply ReFRAG span budgeting to compress context - from scripts.hybrid_search import _merge_and_budget_spans # type: ignore - - try: - if os.environ.get("DEBUG_CONTEXT_ANSWER"): - logger.debug("BUDGET_BEFORE", extra={"items": len(items)}) - _pairs = {} + if items and all(isinstance(it, dict) and it.get("span_budgeted") for it in items): + budgeted = items + else: try: - # Relax budgets for context_answer unless explicitly disabled via CTX_RELAX_BUDGETS=0 - if str(os.environ.get("CTX_RELAX_BUDGETS", "1")).strip().lower() in { - "1", - "true", - "yes", - "on", - }: - # GLM models have much larger context windows - use higher budgets - try: - from scripts.refrag_glm import detect_glm_runtime + from scripts.hybrid_search import _merge_and_budget_spans # type: ignore + + if os.environ.get("DEBUG_CONTEXT_ANSWER"): + logger.debug("BUDGET_BEFORE", extra={"items": len(items)}) + _pairs = {} + try: + # Relax budgets for context_answer unless explicitly disabled via CTX_RELAX_BUDGETS=0 + if str(os.environ.get("CTX_RELAX_BUDGETS", "1")).strip().lower() in { + "1", + "true", + "yes", + "on", + }: + # GLM models have much larger context windows - use higher budgets is_glm = detect_glm_runtime() - except ImportError: - is_glm = False - - if is_glm: - # GLM: 200K context allows much more code context - _default_budget = "8192" # 8x more than Granite - _default_spans = "24" # 3x more spans - else: - # Granite/llamacpp: tighter limits - _default_budget = "1024" - _default_spans = "8" - - _pairs = { - "MICRO_BUDGET_TOKENS": os.environ.get( - "MICRO_BUDGET_TOKENS", _default_budget - ), - "MICRO_OUT_MAX_SPANS": os.environ.get("MICRO_OUT_MAX_SPANS", _default_spans), - } - except Exception: - _pairs = {"MICRO_BUDGET_TOKENS": "5000", "MICRO_OUT_MAX_SPANS": "8"} - with _env_overrides(_pairs): - budgeted = _merge_and_budget_spans(items) - if os.environ.get("DEBUG_CONTEXT_ANSWER"): - logger.debug("BUDGET_AFTER", extra={"items": len(budgeted)}) - if not budgeted and items: + + if is_glm: + # GLM: 200K context allows much more code context + _default_budget = "8192" # 8x more than Granite + _default_spans = "24" # 3x more spans + else: + # Granite/llamacpp: tighter limits + _default_budget = "1024" + _default_spans = "8" + + _pairs = { + "MICRO_BUDGET_TOKENS": os.environ.get( + "MICRO_BUDGET_TOKENS", _default_budget + ), + "MICRO_OUT_MAX_SPANS": os.environ.get("MICRO_OUT_MAX_SPANS", _default_spans), + } + except Exception: + _pairs = {"MICRO_BUDGET_TOKENS": "5000", "MICRO_OUT_MAX_SPANS": "8"} + with _env_overrides(_pairs): + budgeted = _merge_and_budget_spans(items) + if os.environ.get("DEBUG_CONTEXT_ANSWER"): + logger.debug("BUDGET_AFTER", extra={"items": len(budgeted)}) + if not budgeted and items: + if os.environ.get("DEBUG_CONTEXT_ANSWER"): + logger.debug("BUDGET_EMPTY_FALLBACK") + budgeted = items + except (ImportError, AttributeError, KeyError): + logger.warning("Span budgeting failed, using raw items", exc_info=True) if os.environ.get("DEBUG_CONTEXT_ANSWER"): - logger.debug("BUDGET_EMPTY_FALLBACK") + logger.debug("BUDGET_FAILED", exc_info=True) budgeted = items - except (ImportError, AttributeError, KeyError): - logger.warning("Span budgeting failed, using raw items", exc_info=True) - if os.environ.get("DEBUG_CONTEXT_ANSWER"): - logger.debug("BUDGET_FAILED", exc_info=True) - budgeted = items # Enforce an output max spans knob - do this BEFORE env restore try: @@ -2055,11 +2022,7 @@ def _to_float(v, d): # Granite/llamacpp: use env var or 2000 default # GLM: dynamically use model's max_output_tokens from config - try: - from scripts.refrag_glm import detect_glm_runtime, get_glm_model_name, get_model_config - is_glm = detect_glm_runtime() - except ImportError: - is_glm = False + is_glm = detect_glm_runtime() if is_glm: # Pull dynamic limit from GLM model config (imports already succeeded above) @@ -2502,22 +2465,6 @@ async def _context_answer_impl( import time import asyncio - # Import logger utilities - try: - from scripts.logger import safe_bool, safe_float - except ImportError: - def safe_bool(val, default=False, **kw): - if val is None: - return default - if isinstance(val, bool): - return val - return str(val).strip().lower() in {"1", "true", "yes", "on"} - def safe_float(val, default=0.0, **kw): - try: - return float(val) if val is not None else default - except Exception: - return default - # Get embedding model function if get_embedding_model_fn is None: from scripts.mcp_impl.admin_tools import _get_embedding_model @@ -2784,19 +2731,6 @@ def safe_float(val, default=0.0, **kw): "query": original_queries, } - # Ensure final retrieval call reflects Tier-2 relaxed filters - try: - from scripts.hybrid_search import run_hybrid_search as _rh - await asyncio.to_thread( - lambda: _rh( - queries=queries, - limit=int(max(lim, 1)), - per_path=int(max(ppath, 1)), - ) - ) - except Exception: - pass - # Build citations and context payload for the decoder ( citations, @@ -3040,17 +2974,6 @@ def _k(s: Dict[str, Any]): "query": original_queries, } - # Final introspection call - try: - from scripts.hybrid_search import run_hybrid_search as _rh2 - _ = _rh2( - queries=queries, - limit=int(max(lim, 1)), - per_path=int(max(ppath, 1)), - ) - except Exception: - pass - # Optional: provide per-query answers/citations for pack mode answers_by_query = None try: diff --git a/scripts/mcp_impl/context_search.py b/scripts/mcp_impl/context_search.py index 14ffa4a5..701828d4 100644 --- a/scripts/mcp_impl/context_search.py +++ b/scripts/mcp_impl/context_search.py @@ -37,6 +37,7 @@ ) from scripts.mcp_impl.workspace import _default_collection, _MEM_COLL_CACHE from scripts.mcp_impl.toon import _should_use_toon, _format_context_results_as_toon +from scripts.mcp_http_client import call_tool_http # Environment QDRANT_URL = os.environ.get("QDRANT_URL", "http://qdrant:6333") @@ -681,8 +682,6 @@ def _maybe_dict(val: Any) -> Dict[str, Any]: used_http_fallback = False if not code_hits: try: - from scripts.mcp_router import call_tool_http # type: ignore - base = ( os.environ.get("MCP_INDEXER_HTTP_URL") or "http://localhost:8003/mcp" ).rstrip("/") @@ -1254,4 +1253,3 @@ def push_text( if _should_use_toon(output_format): return _format_context_results_as_toon(ret, compact=bool(eff_compact)) return ret - diff --git a/scripts/mcp_impl/info_request.py b/scripts/mcp_impl/info_request.py deleted file mode 100644 index e9c72601..00000000 --- a/scripts/mcp_impl/info_request.py +++ /dev/null @@ -1,159 +0,0 @@ -#!/usr/bin/env python3 -""" -mcp/info_request.py - Info request helpers for MCP indexer server. - -Extracted from mcp_indexer_server.py for better modularity. -Contains: -- Helper functions for info_request tool -""" - -from __future__ import annotations - -__all__ = [ - "_extract_symbols_from_query", - "_extract_related_concepts", - "_format_information_field", - "_extract_relationships", - "_calculate_confidence", -] - -import re -import logging -from typing import Any, Dict, List - -logger = logging.getLogger(__name__) - -# Import _split_ident for tokenization -from scripts.mcp_impl.utils import _split_ident - - -def _extract_symbols_from_query(query: str) -> list[str]: - """Extract potential symbol names from a query string.""" - # Match CamelCase, snake_case, or standalone words that look like identifiers - patterns = [ - r'\b[A-Z][a-z]+(?:[A-Z][a-z]+)+\b', # CamelCase - r'\b[a-z_][a-z0-9_]*(?:_[a-z0-9]+)+\b', # snake_case - r'\b(?:def|class|function|method|async)\s+(\w+)', # explicit mentions - ] - symbols = set() - for pat in patterns: - for m in re.finditer(pat, query): - sym = m.group(1) if m.lastindex else m.group(0) - if len(sym) > 2: - symbols.add(sym) - return list(symbols)[:5] # Limit to top 5 - - -def _extract_related_concepts(query: str, results: list) -> list[str]: - """Extract related technical concepts dynamically from results (codebase-agnostic).""" - concepts = set() - - # Extract from results - this works on any codebase - for r in results[:10]: - # From symbols: split CamelCase/snake_case into meaningful parts - sym = r.get("symbol", "") or "" - if sym and len(sym) > 2: - parts = [p for p in re.split(r'(?=[A-Z])|_|-', sym) if p and len(p) > 2] - for part in parts[:3]: - concepts.add(part.lower()) - - # From file paths: extract directory/module names - path = r.get("path", "") or "" - if path: - path_parts = path.replace("\\", "/").split("/") - for pp in path_parts[-3:]: # Last 3 path segments - # Remove extension and split - name = pp.rsplit(".", 1)[0] if "." in pp else pp - if name and len(name) > 2 and not name.startswith("_"): - concepts.add(name.lower()) - - # From kind: function, class, method, etc. - kind = r.get("kind", "") or "" - if kind and len(kind) > 2: - concepts.add(kind.lower()) - - # From query: extract significant words (skip common words) - skip_words = {"the", "is", "are", "how", "does", "what", "where", "find", "get", "set", "for", "and", "with"} - query_parts = re.split(r'\W+', query.lower()) - for qp in query_parts: - if qp and len(qp) > 2 and qp not in skip_words: - concepts.add(qp) - - # Sort by frequency in results for relevance - return list(concepts)[:10] - - -def _format_information_field(result: dict) -> str: - """Generate human-readable information field for a result.""" - path = result.get("path", "") - symbol = result.get("symbol", "") - start = result.get("start_line", 0) - end = result.get("end_line", 0) - kind = result.get("kind", "") - - # Get just the filename - filename = path.split("/")[-1] if "/" in path else path - - if symbol and kind: - return f"Found {kind} '{symbol}' in {filename} (lines {start}-{end})" - elif symbol: - return f"Found '{symbol}' in {filename} (lines {start}-{end})" - else: - return f"Found match in {filename} (lines {start}-{end})" - - -def _extract_relationships(result: dict) -> dict: - """Extract relationship metadata (imports, calls) from a result.""" - relations = result.get("relations") or {} - # Get from relations object if present - imports = relations.get("imports") or [] - calls = relations.get("calls") or [] - symbol_path = relations.get("symbol_path") or "" - # Also check top-level metadata (fallback) - if not imports: - imports = result.get("imports") or [] - if not calls: - calls = result.get("calls") or [] - # Get related paths if available - related_paths = result.get("related_paths") or [] - - return { - "imports_from": imports[:10] if imports else [], # Limit to 10 - "calls": calls[:10] if calls else [], - "symbol_path": symbol_path, - "related_paths": related_paths[:5] if related_paths else [], - } - - -def _calculate_confidence(query: str, results: list) -> dict: - """Calculate confidence metrics for the search.""" - if not results: - return {"level": "none", "score": 0.0, "reason": "no_results"} - - avg_score = sum(r.get("score", 0) for r in results) / len(results) - top_score = results[0].get("score", 0) if results else 0 - - # Check if query terms match symbols - query_tokens = set(_split_ident(query.lower())) - symbol_matches = sum( - 1 for r in results[:5] - if any(tok in _split_ident((r.get("symbol", "") or "").lower()) - for tok in query_tokens) - ) - - if top_score > 0.8 and symbol_matches > 0: - level = "high" - elif avg_score > 0.6: - level = "medium" - elif results: - level = "low" - else: - level = "none" - - return { - "level": level, - "score": round(avg_score, 3), - "top_score": round(top_score, 3), - "symbol_matches": symbol_matches, - } - diff --git a/scripts/mcp_impl/pattern_search.py b/scripts/mcp_impl/pattern_search.py deleted file mode 100644 index 02898c9c..00000000 --- a/scripts/mcp_impl/pattern_search.py +++ /dev/null @@ -1,236 +0,0 @@ -"""Pattern search MCP tool implementation. - -Single unified tool that handles both code examples and natural language descriptions. -Supports TOON output format for token-efficient responses. -""" -from __future__ import annotations - -import os -import re -from typing import Any, Dict, List, Optional, Union - -# Import logger with fallback -try: - from scripts.logger import get_logger - logger = get_logger(__name__) -except ImportError: - import logging - logger = logging.getLogger(__name__) - -# Import pattern detection components (lazy to avoid startup penalty) -_PATTERN_SEARCH_LOADED = False -_pattern_search_fn = None -_search_by_pattern_description_fn = None - - -def _ensure_pattern_search(): - """Lazy load pattern search module.""" - global _PATTERN_SEARCH_LOADED, _pattern_search_fn, _search_by_pattern_description_fn - if _PATTERN_SEARCH_LOADED: - return True - try: - from scripts.pattern_detection.search import ( - pattern_search, - search_by_pattern_description, - ) - _pattern_search_fn = pattern_search - _search_by_pattern_description_fn = search_by_pattern_description - _PATTERN_SEARCH_LOADED = True - return True - except ImportError as e: - logger.warning(f"Pattern search not available: {e}") - return False - - -# Supported languages for tree-sitter parsing -_SUPPORTED_LANGUAGES = { - "python", "javascript", "typescript", "go", "rust", "java", "c", "cpp", - "ruby", "php", "csharp", "kotlin", "swift", "scala", "bash", "lua", -} - -# Fenced code block pattern -_FENCED_CODE = re.compile(r'^```\w*\n.*\n```$', re.DOTALL) - -# Universal code syntax patterns (work across all languages) -_CODE_SYNTAX = re.compile( - r'[{}\[\]();]|' # Brackets, braces, parens, semicolons - r'::|->|=>|:=|' # C++/Rust/Go/JS operators - r'\.\w+\(|' # Method call: .foo( - r'\w+\s*\([^)]*\)|' # Function call: foo() or foo(args) - r'^\s*(def|func|fn|function|class|struct|enum|impl|trait|interface)\s+\w', # Definitions - re.MULTILINE -) - -# Multi-line code indicators (braces/semicolons at line boundaries) -_MULTILINE_CODE = re.compile( - r'[{}]\s*$|' # Brace at line end - r'^\s*[{}]|' # Brace at line start - r';\s*$', # Semicolon at line end - re.MULTILINE -) - - -def _detect_query_mode(text: str, language: str | None) -> str: - """ - Auto-detect if text is code or natural language description. - - Works across all 16+ supported languages using universal syntax patterns. - Returns: "code" or "description" - """ - text = text.strip() - if not text: - return "description" - - # 1. Fenced code block → code - if _FENCED_CODE.match(text): - return "code" - - # 2. Multi-line with braces/semicolons → code - if '\n' in text and _MULTILINE_CODE.search(text): - return "code" - - # 3. Universal code syntax (brackets, operators, calls, definitions) - if _CODE_SYNTAX.search(text): - return "code" - - # 4. Language hint is advisory only, but allow a narrow "bare identifiers" case. - # Some clients send extremely minimal snippets (e.g. "some text") with a language hint. - # Treat exactly-two identifier tokens as code when a supported language hint is present. - if language and language.lower() in _SUPPORTED_LANGUAGES: - if re.fullmatch(r"[A-Za-z_][A-Za-z0-9_]*\s+[A-Za-z_][A-Za-z0-9_]*", text): - return "code" - return "description" - - # 5. Default to natural language - return "description" - - -async def _pattern_search_impl( - query: Optional[str] = None, - language: Optional[str] = None, - limit: Optional[int] = None, - min_score: Optional[float] = None, - include_snippet: Optional[bool] = None, - context_lines: Optional[int] = None, - hybrid: Optional[bool] = None, - semantic_weight: Optional[float] = None, - collection: Optional[str] = None, - target_languages: Optional[List[str]] = None, - output_format: Optional[str] = None, - compact: Optional[bool] = None, - aroma_rerank: Optional[bool] = None, - aroma_alpha: Optional[float] = None, - query_mode: Optional[str] = None, # "code", "description", or "auto" (default) - coerce_bool_fn=None, - coerce_int_fn=None, - coerce_float_fn=None, -) -> Dict[str, Any]: - """Unified pattern search - handles both code examples and NL descriptions.""" - if not _ensure_pattern_search(): - return {"ok": False, "error": "Pattern search module not available"} - - if not query or not str(query).strip(): - return {"ok": False, "error": "query parameter is required"} - - query_text = str(query).strip() - - # Coerce parameters - handle string "false"/"0" correctly - def _default_coerce_bool(v, d): - if v is None: - return d - if isinstance(v, bool): - return v - if isinstance(v, str): - return v.strip().lower() in ("1", "true", "yes", "on") - return bool(v) - - def _safe_coerce_int(v, d): - if v is None: - return d - try: - return int(v) - except (ValueError, TypeError): - return d - - def _safe_coerce_float(v, d): - if v is None: - return d - try: - return float(v) - except (ValueError, TypeError): - return d - - _coerce_bool = coerce_bool_fn or _default_coerce_bool - _coerce_int = coerce_int_fn or _safe_coerce_int - _coerce_float = coerce_float_fn or _safe_coerce_float - - # Defaults aligned with core pattern_search API for consistent behavior - eff_limit = _coerce_int(limit, 10) - eff_include_snippet = _coerce_bool(include_snippet, True) - eff_context_lines = _coerce_int(context_lines, 3) - eff_hybrid = _coerce_bool(hybrid, False) - eff_semantic_weight = _coerce_float(semantic_weight, 0.3) - eff_compact = _coerce_bool(compact, False) - eff_aroma_rerank = _coerce_bool(aroma_rerank, True) # AROMA enabled by default - eff_aroma_alpha = _coerce_float(aroma_alpha, 0.6) - - # Determine query mode: explicit override or auto-detect - eff_language = str(language).strip() if language else None - eff_query_mode = str(query_mode).strip().lower() if query_mode else "auto" - - if eff_query_mode == "code": - is_code = True - elif eff_query_mode == "description": - is_code = False - else: # auto - detected = _detect_query_mode(query_text, eff_language) - is_code = (detected == "code") - - # Path-specific min_score defaults: - # - Code path: 0.5 (vector similarity scores are typically higher) - # - NL path: 0.0 (keyword overlap scores are often low, don't filter by default) - eff_min_score = _coerce_float(min_score, 0.5 if is_code else 0.0) - - try: - if is_code: - # Structural pattern search using code example - result = _pattern_search_fn( - example=query_text, - language=eff_language or "python", - limit=eff_limit, - min_score=eff_min_score, - include_snippet=eff_include_snippet, - context_lines=eff_context_lines, - hybrid=eff_hybrid, - semantic_weight=eff_semantic_weight, - collection=collection, - target_languages=target_languages, - output_format=output_format, - compact=eff_compact, - aroma_rerank=eff_aroma_rerank, - aroma_alpha=eff_aroma_alpha, - ) - else: - # Natural language pattern description search - result = _search_by_pattern_description_fn( - description=query_text, - limit=eff_limit, - min_score=eff_min_score, - collection=collection, - target_languages=target_languages, - output_format=output_format, - compact=eff_compact, - ) - - # Convert response object to dict if needed - if not isinstance(result, dict): - result = result.to_dict() - - # Preserve upstream ok flag (derived from search_mode) instead of overriding - # This ensures errors from core search propagate to MCP clients - result["query_mode"] = "code" if is_code else "description" - - return result - except Exception as e: - logger.error(f"Pattern search failed: {e}") - return {"ok": False, "error": str(e)} diff --git a/scripts/mcp_impl/search.py b/scripts/mcp_impl/search.py index e4dc3766..dc600eff 100644 --- a/scripts/mcp_impl/search.py +++ b/scripts/mcp_impl/search.py @@ -14,6 +14,7 @@ __all__ = [ "_repo_search_impl", + "enrich_feedback_rating", ] import json @@ -23,6 +24,7 @@ import asyncio import subprocess import hashlib +import time from typing import Any, Dict, List, Optional, Tuple logger = logging.getLogger(__name__) @@ -39,10 +41,20 @@ _tokens_from_queries, safe_int, ) -from scripts.mcp_impl.workspace import _default_collection, _work_script +from scripts.mcp_impl.workspace import _default_collection from scripts.mcp_impl.admin_tools import _detect_current_repo, _run_async -from scripts.mcp_toon import _should_use_toon, _format_results_as_toon +from scripts.mcp_impl.search_profiles import append_profile_globs, normalize_profile +from scripts.mcp_impl.toon import _should_use_toon, _format_results_as_toon from scripts.mcp_auth import require_collection_access as _require_collection_access +from scripts.path_scope import ( + metadata_matches_under as _metadata_matches_under, + normalize_under as _normalize_under_scope, +) +from scripts.relevance_feedback import ( + enrich_recent_rating, + remember_recent_results, + stable_target_key, +) # Constants QDRANT_URL = os.environ.get("QDRANT_URL", "http://qdrant:6333") @@ -53,6 +65,694 @@ context="MCP_SNIPPET_MAX_BYTES", ) +_RECENT_RESULT_META: dict[str, tuple[float, dict]] = {} +_RECENT_RESULT_META_TTL = 3600 +_RECENT_RESULT_META_MAX = 4096 +_RECENT_RESULT_META_KEYS = ( + "result_id", + "target_id", + "impression_id", + "path", + "host_path", + "container_path", + "symbol", + "kind", + "repo", + "file_hash", + "symbol_content_hash", +) + + +# Fields to strip from results when debug=False (internal/debugging fields) +_DEBUG_RESULT_FIELDS = { + "components", # Internal scoring breakdown (dense_rrf, lexical, fname_boost, etc.) + "doc_id", # Internal benchmark ID (often null/opaque) + "code_id", # Internal benchmark ID (often null/opaque) + "payload", # Duplicates other fields (information, document, pseudo, tags) + "why", # Often empty []; debugging explanation list + "span_budgeted", # Internal budget flag + "relations", # Call graph info (imports, calls) - useful but often noise + "related_paths", # Optional related file paths + "budget_tokens_used", # Internal token accounting + "fname_boost", # Internal boost value (already applied to score) + "relevance_boost", # Internal feedback boost value (already applied to score) + "feedback_prior", # Internal feedback metadata + "feedback_recall", # Internal feedback recall marker + "feedback_graph_recall", # Internal graph recall marker + "feedback_weight_id", # Internal source weight identity after reconciliation + "pseudo", # Internal retrieval enrichment; debug-only by default + "tags", # Internal retrieval enrichment; debug-only by default + "file_hash", # Internal feedback/reindex metadata + "symbol_content_hash", # Internal feedback/reindex metadata + "host_path", # Internal dual-path (host side) - use path/client_path instead + "container_path", # Internal dual-path (container side) - use path/client_path instead +} + +# Top-level response fields to strip when debug=False +_DEBUG_TOP_LEVEL_FIELDS = { + "rerank_counters", # Internal reranking metrics (inproc_hybrid, timeout, etc.) + "code_signals", # Internal code signal detection results +} + + +def _strip_debug_fields(item: dict, keep_paths: bool = True) -> dict: + """Strip internal/debug fields from a result item. + + Args: + item: Result dict to strip + keep_paths: If True, keep host_path/container_path + + Returns: + New dict with debug fields removed + """ + strip_fields = _DEBUG_RESULT_FIELDS + if keep_paths: + strip_fields = _DEBUG_RESULT_FIELDS - {"host_path", "container_path"} + result = {k: v for k, v in item.items() if k not in strip_fields} + return result + + +def _result_content_hash(result: dict) -> str: + """Best-effort indexed content hash for feedback identity.""" + if not isinstance(result, dict): + return "" + payload = result.get("payload") if isinstance(result.get("payload"), dict) else {} + metadata = result.get("metadata") if isinstance(result.get("metadata"), dict) else {} + return str( + result.get("file_hash") + or result.get("content_hash") + or payload.get("file_hash") + or payload.get("content_hash") + or metadata.get("file_hash") + or metadata.get("content_hash") + or "" + ) + + +def _result_target_key(result: dict) -> str: + """Stable feedback target: prefer repo+symbol identity, fall back to file.""" + if not isinstance(result, dict): + return "" + payload = result.get("payload") if isinstance(result.get("payload"), dict) else {} + metadata = result.get("metadata") if isinstance(result.get("metadata"), dict) else {} + relations = result.get("relations") if isinstance(result.get("relations"), dict) else {} + + repo = str( + result.get("repo") + or metadata.get("repo") + or payload.get("repo") + or "" + ).strip() + kind = str( + result.get("kind") + or metadata.get("kind") + or payload.get("kind") + or "" + ).strip() + symbol_path = str( + result.get("symbol_path") + or relations.get("symbol_path") + or metadata.get("symbol_path") + or payload.get("symbol_path") + or result.get("symbol") + or metadata.get("symbol") + or payload.get("symbol") + or "" + ).strip() + path = str( + result.get("container_path") + or metadata.get("container_path") + or result.get("path") + or metadata.get("path") + or payload.get("path") + or "" + ).strip() + + return stable_target_key(repo=repo, kind=kind, symbol=symbol_path, path=path) + + +def _inject_result_ids(results: list[dict], canonical_query: str) -> None: + """Attach stable target IDs plus query/content-specific impression IDs.""" + for r in results: + _path = str(r.get("path") or "") + _start = int(r.get("start_line") or 0) + _end = int(r.get("end_line") or 0) + _content_hash = _result_content_hash(r) + _target_key = _result_target_key(r) + if not _target_key: + _target_key = f"span\x00{_path}\x00{_start}\x00{_end}" + _impression_key = f"{canonical_query}\x00{_target_key}\x00{_path}\x00{_start}\x00{_end}\x00{_content_hash}" + _target_id = hashlib.sha256(_target_key.encode("utf-8")).hexdigest()[:12] + r["target_id"] = _target_id + r["result_id"] = _target_id + r["impression_id"] = hashlib.sha256(_impression_key.encode("utf-8")).hexdigest()[:12] + + +def _remember_result_metadata(results: list[dict], collection: str = "") -> None: + """Keep recent metadata in process and shared storage for hands-off rating.""" + now = time.time() + expired_before = now - _RECENT_RESULT_META_TTL + for key, (ts, _) in list(_RECENT_RESULT_META.items()): + if ts < expired_before: + _RECENT_RESULT_META.pop(key, None) + for result in results: + rid = str(result.get("result_id") or "").strip() + if not rid: + continue + meta = {} + for key in _RECENT_RESULT_META_KEYS: + val = result.get(key) + if val is not None and str(val).strip(): + meta[key] = str(val).strip() + if meta: + _RECENT_RESULT_META[rid] = (now, meta) + while len(_RECENT_RESULT_META) > _RECENT_RESULT_META_MAX: + try: + oldest = min(_RECENT_RESULT_META.items(), key=lambda kv: kv[1][0])[0] + _RECENT_RESULT_META.pop(oldest, None) + except Exception: + break + remember_recent_results(collection, results) + + +def enrich_feedback_rating(rating: dict, collection: str = "") -> dict: + """Fill rating metadata from shared or in-process recent search results.""" + if not isinstance(rating, dict): + return {} + out = enrich_recent_rating(collection, rating) + rid = str(out.get("result_id") or "").strip() + if not rid: + return out + cached = _RECENT_RESULT_META.get(rid) + if not cached: + return out + _, meta = cached + for key, val in meta.items(): + out.setdefault(key, val) + return out + + +def _load_relevance_weights(collection: str) -> dict: + try: + from pathlib import Path as _Path + weights_file = _Path(os.environ.get("RERANKER_WEIGHTS_DIR", "/tmp/rerank_weights")) + weights_file = weights_file / f"{collection}_relevance.json" + if weights_file.exists(): + with open(weights_file, "r") as f: + return json.loads(f.read()) + except Exception: + pass + return {} + + +def _feedback_symbol_variants(symbol: str) -> list[str]: + """Small symbol variant set for graph-edge callee lookups.""" + s = str(symbol or "").strip() + if not s: + return [] + variants = [s] + if "." in s: + base = s.split(".")[-1].strip() + if base: + variants.append(base) + out = [] + seen = set() + for v in variants: + if v and v not in seen: + seen.add(v) + out.append(v) + return out + + +def _feedback_path_variants(path: str, repo: str = "") -> list[str]: + """Return equivalent path spellings used by indexed metadata.""" + raw = str(path or "").strip().replace("\\", "/") + if not raw: + return [] + raw = "/" + raw.strip("/") if raw.startswith("/") else raw.strip("/") + variants = [raw, raw.strip("/")] + repo_name = str(repo or "").strip().replace("\\", "/").strip("/") + raw_no_slash = raw.strip("/") + if raw_no_slash.startswith("work/"): + work_prefix = "/work/" + rest = raw_no_slash[len("work/") :] + if rest: + variants.extend((rest, "/" + rest)) + if repo_name and rest.casefold().startswith(repo_name.casefold() + "/"): + tail = rest[len(repo_name) + 1 :] + variants.extend((tail, "/" + tail)) + if repo_name: + marker = f"/{repo_name.casefold()}/" + raw_cf = f"/{raw.strip('/').casefold()}/" + marker_at = raw_cf.find(marker) + if marker_at >= 0: + tail_start = max(0, marker_at + len(marker) - 1) + tail = raw.strip("/")[tail_start:].strip("/") + variants.extend((tail, "/" + tail)) + out: list[str] = [] + seen: set[str] = set() + for variant in variants: + value = variant if variant.startswith("/") else variant.strip("/") + if value and value not in seen: + seen.add(value) + out.append(value) + return out + + +def _feedback_point_matches_filters( + point: Any, + *, + repo_filter: Any = "*", + language: str = "", + under: str | None = None, + kind: str = "", + symbol: str = "", + ext: str = "", + not_: str = "", + path_regex: str = "", + path_globs: list[str] | None = None, + not_globs: list[str] | None = None, + case_sensitive: bool = False, +) -> bool: + """Apply the same explicit filters to feedback-recalled points as search.""" + payload = getattr(point, "payload", None) or {} + md = payload.get("metadata") or {} + if not isinstance(md, dict): + md = {} + + candidate_repo = str(md.get("repo") or "").strip() + if repo_filter != "*" and repo_filter: + allowed_repos = ( + {str(value).strip() for value in repo_filter if str(value).strip()} + if isinstance(repo_filter, (list, tuple, set)) + else {str(repo_filter).strip()} + ) + if candidate_repo not in allowed_repos: + return False + if language and str(md.get("language") or "").strip() != language: + return False + if kind and str(md.get("kind") or "").strip() != kind: + return False + if symbol: + requested = str(symbol).strip() + candidate_symbols = { + str(md.get("symbol") or "").strip(), + str(md.get("symbol_path") or "").strip(), + } + if requested not in candidate_symbols: + return False + if under and not _metadata_matches_under(md, under): + return False + + path_values = [] + for key in ( + "path", + "repo_rel_path", + "host_path", + "container_path", + "file_path", + "client_path", + ): + value = md.get(key) + if value is not None and str(value).strip(): + path_values.append(str(value).strip().replace("\\", "/")) + if not path_values: + return False + if not case_sensitive: + normalized_paths = [value.lower() for value in path_values] + else: + normalized_paths = path_values + + def _contains(value: str) -> bool: + return value if case_sensitive else value.lower() + + if not_ and any(_contains(not_) in value for value in normalized_paths): + return False + if ext: + ext_value = str(ext).lower().lstrip(".") + if not any(value.lower().endswith("." + ext_value) for value in path_values): + return False + if path_regex: + flags = 0 if case_sensitive else re.IGNORECASE + try: + if not any(re.search(path_regex, value, flags=flags) for value in path_values): + return False + except re.error: + return False + + def _match_glob(pattern: str, value: str) -> bool: + import fnmatch + + pattern_value = pattern if case_sensitive else pattern.lower() + path_value = value if case_sensitive else value.lower() + path_value = path_value.strip("/") + if fnmatch.fnmatchcase(path_value, pattern_value): + return True + if not pattern_value.startswith("/") and "/" in path_value: + parts = [segment for segment in path_value.split("/") if segment] + return any( + fnmatch.fnmatchcase("/".join(parts[index:]), pattern_value) + for index in range(1, len(parts)) + ) + return False + + normalized_path_globs = list(path_globs or []) + if normalized_path_globs and not any( + _match_glob(pattern, value) + for pattern in normalized_path_globs + for value in path_values + ): + return False + normalized_not_globs = list(not_globs or []) + if normalized_not_globs and any( + _match_glob(pattern, value) + for pattern in normalized_not_globs + for value in path_values + ): + return False + return True + + +def _point_to_feedback_item(point: Any, *, score: float, source: str, prior: dict | None = None) -> dict: + payload = getattr(point, "payload", None) or {} + md = payload.get("metadata") or {} + return { + "score": float(score), + "path": str(md.get("host_path") or md.get("path") or ""), + "host_path": str(md.get("host_path") or ""), + "container_path": str(md.get("container_path") or md.get("path") or ""), + "symbol": str(md.get("symbol_path") or md.get("symbol") or ""), + "kind": str(md.get("kind") or ""), + "repo": str(md.get("repo") or ""), + "start_line": int(md.get("start_line") or 0), + "end_line": int(md.get("end_line") or 0), + "relations": { + "imports": md.get("imports") or [], + "calls": md.get("calls") or [], + "symbol_path": str(md.get("symbol_path") or md.get("symbol") or ""), + }, + "file_hash": str(md.get("file_hash") or ""), + "symbol_content_hash": str(md.get("symbol_content_hash") or ""), + source: True, + "feedback_prior": prior or {}, + } + + +def _scroll_main_point( + client: Any, + qmodels: Any, + *, + collection: str, + repo: str, + symbol: str = "", + symbol_content_hash: str = "", + kind: str = "", + path: str = "", +) -> Any | None: + base_must = [] + if repo: + base_must.append(qmodels.FieldCondition(key="metadata.repo", match=qmodels.MatchValue(value=repo))) + if symbol: + base_must.append(qmodels.FieldCondition(key="metadata.symbol_path", match=qmodels.MatchValue(value=symbol))) + elif not path: + return None + if kind: + base_must.append(qmodels.FieldCondition(key="metadata.kind", match=qmodels.MatchValue(value=kind))) + + path_variants = _feedback_path_variants(path, repo) + path_keys = ("metadata.path", "metadata.container_path", "metadata.host_path") + + def _scroll_with(must: list[Any], limit: int = 1) -> list[Any]: + try: + points, _ = client.scroll( + collection_name=collection, + scroll_filter=qmodels.Filter(must=must), + limit=limit, + with_payload=True, + with_vectors=False, + ) + return list(points or []) + except Exception: + return [] + + if symbol or path: + if path_variants: + for path_value in path_variants: + for path_key in path_keys: + points = _scroll_with( + base_must + + [ + qmodels.FieldCondition( + key=path_key, + match=qmodels.MatchValue(value=path_value), + ) + ] + ) + if points: + return points[0] + elif symbol: + points = _scroll_with(base_must) + if points: + return points[0] + else: + return None + if not symbol_content_hash: + return None + hash_must = [] + if repo: + hash_must.append( + qmodels.FieldCondition(key="metadata.repo", match=qmodels.MatchValue(value=repo)) + ) + if kind: + hash_must.append( + qmodels.FieldCondition(key="metadata.kind", match=qmodels.MatchValue(value=kind)) + ) + hash_must.append( + qmodels.FieldCondition( + key="metadata.symbol_content_hash", + match=qmodels.MatchValue(value=symbol_content_hash), + ) + ) + if path_variants: + for path_value in path_variants: + for path_key in path_keys: + points = _scroll_with( + hash_must + + [ + qmodels.FieldCondition( + key=path_key, + match=qmodels.MatchValue(value=path_value), + ) + ], + limit=2, + ) + if len(points) == 1: + return points[0] + points = _scroll_with(hash_must, limit=2) + return points[0] if len(points) == 1 else None + + +def _feedback_recall_candidates( + *, + collection: str, + weights: dict, + existing_target_ids: set[str], + existing_paths: set[str], + base_score: float, + max_candidates: int, + repo_filter: Any = "*", + language: str = "", + under: str | None = None, + kind_filter: str = "", + symbol_filter: str = "", + ext: str = "", + not_: str = "", + path_regex: str = "", + path_globs: list[str] | None = None, + not_globs: list[str] | None = None, + case_sensitive: bool = False, +) -> list[dict]: + """Rehydrate positively rated targets and inverse-graph adjacent callers.""" + if max_candidates <= 0: + return [] + result_weights = weights.get("results") if isinstance(weights, dict) else {} + if not isinstance(result_weights, dict): + return [] + + ranked = [] + for rid, info in result_weights.items(): + if not isinstance(info, dict): + continue + if info.get("superseded_by"): + continue + avg = float(info.get("avg_relevance", 0) or 0) + count = int(info.get("count", 0) or 0) + inheritance = float(info.get("inheritance_weight", 1.0) or 0) + target = info.get("target") if isinstance(info.get("target"), dict) else {} + if avg <= 0 or count <= 0 or inheritance <= 0 or not target: + continue + ranked.append((avg * inheritance, count, str(rid), target)) + ranked.sort(key=lambda item: (item[0], item[1]), reverse=True) + if not ranked: + return [] + + try: + from qdrant_client import QdrantClient + from qdrant_client import models as qmodels + except Exception: + return [] + + try: + client = QdrantClient( + url=QDRANT_URL, + api_key=os.environ.get("QDRANT_API_KEY"), + timeout=float(os.environ.get("QDRANT_TIMEOUT", "20") or 20), + ) + except Exception: + return [] + + out: list[dict] = [] + seen_paths: set[str] = {p for p in existing_paths if p} + try: + graph_max = int(os.environ.get("RELEVANCE_GRAPH_RECALL_MAX", "3") or 0) + except Exception: + graph_max = 3 + try: + graph_boost_factor = float(os.environ.get("RELEVANCE_GRAPH_RECALL_BOOST", "0.01") or 0.0) + except Exception: + graph_boost_factor = 0.01 + try: + from scripts.ingest.graph_edges import GRAPH_COLLECTION_SUFFIX as _graph_suffix + except Exception: + _graph_suffix = "_graph" + graph_coll = f"{collection}{_graph_suffix}" + + for effective_avg, count, rid, target in ranked: + if len(out) >= max_candidates: + break + repo = str(target.get("repo") or "").strip() + target_kind = str(target.get("kind") or "").strip() + symbol = str(target.get("symbol") or "").strip() + symbol_content_hash = str(target.get("symbol_content_hash") or "").strip() + path = str(target.get("container_path") or target.get("path") or "").strip() + raw_info = result_weights.get(rid) or {} + prior = { + "avg_relevance": float(raw_info.get("avg_relevance", 0) or 0), + "inheritance_weight": float(raw_info.get("inheritance_weight", 1.0) or 0), + "count": count, + "result_id": rid, + } + + point = _scroll_main_point( + client, + qmodels, + collection=collection, + repo=repo, + symbol=symbol, + symbol_content_hash=symbol_content_hash, + kind=target_kind, + path=path, + ) + if ( + rid not in existing_target_ids + and point is not None + and _feedback_point_matches_filters( + point, + repo_filter=repo_filter, + language=language, + under=under, + kind=kind_filter, + symbol=symbol_filter, + ext=ext, + not_=not_, + path_regex=path_regex, + path_globs=path_globs, + not_globs=not_globs, + case_sensitive=case_sensitive, + ) + ): + item = _point_to_feedback_item(point, score=base_score, source="feedback_recall", prior=prior) + item["feedback_weight_id"] = rid + emit_path = str(item.get("path") or item.get("container_path") or "") + if emit_path and emit_path not in seen_paths: + seen_paths.add(emit_path) + out.append(item) + if len(out) >= max_candidates: + break + + if graph_max <= 0 or not symbol: + continue + graph_added = 0 + for variant in _feedback_symbol_variants(symbol): + if graph_added >= graph_max or len(out) >= max_candidates: + break + must = [ + qmodels.FieldCondition(key="edge_type", match=qmodels.MatchValue(value="calls")), + qmodels.FieldCondition(key="callee_symbol", match=qmodels.MatchValue(value=variant)), + ] + if repo: + must.append(qmodels.FieldCondition(key="repo", match=qmodels.MatchValue(value=repo))) + try: + edge_points, _ = client.scroll( + collection_name=graph_coll, + scroll_filter=qmodels.Filter(must=must), + limit=max(8, graph_max * 4), + with_payload=True, + with_vectors=False, + ) + except Exception: + continue + for edge in edge_points or []: + if graph_added >= graph_max or len(out) >= max_candidates: + break + edge_payload = getattr(edge, "payload", None) or {} + caller_path = str(edge_payload.get("caller_path") or "").strip() + caller_repo = str(edge_payload.get("repo") or repo).strip() + if not caller_path: + continue + point = _scroll_main_point( + client, + qmodels, + collection=collection, + repo=caller_repo, + path=caller_path, + ) + if point is None: + continue + if not _feedback_point_matches_filters( + point, + repo_filter=repo_filter, + language=language, + under=under, + kind=kind_filter, + symbol=symbol_filter, + ext=ext, + not_=not_, + path_regex=path_regex, + path_globs=path_globs, + not_globs=not_globs, + case_sensitive=case_sensitive, + ): + continue + graph_score = ( + base_score + + graph_boost_factor * (effective_avg / 2.0) * min(count, 10) / 10.0 + ) + item = _point_to_feedback_item( + point, + score=graph_score, + source="feedback_graph_recall", + prior={**prior, "callee_symbol": variant, "caller_path": caller_path}, + ) + item["feedback_weight_id"] = rid + emit_path = str(item.get("path") or item.get("container_path") or "") + if not emit_path or emit_path in seen_paths: + continue + seen_paths.add(emit_path) + out.append(item) + graph_added += 1 + return out + async def _repo_search_impl( query: Any = None, @@ -71,6 +771,7 @@ async def _repo_search_impl( collection: Any = None, workspace_path: Any = None, mode: Any = None, + profile: Any = None, session: Any = None, ctx: Any = None, # MCP Context (passed from wrapper) # Structured filters (optional; mirrors hybrid_search flags) @@ -89,6 +790,7 @@ async def _repo_search_impl( repo: Any = None, # str, list[str], or "*" to search all repos # Response shaping compact: Any = None, + debug: Any = None, # When True, include verbose internal fields (components, rerank_counters, etc.) output_format: Any = None, # "json" (default) or "toon" for token-efficient format args: Any = None, # Compatibility shim for mcp-remote/Claude wrappers that send args/kwargs kwargs: Any = None, @@ -117,18 +819,26 @@ async def _repo_search_impl( - repo: str or list[str]. Filter by repo name(s). Use "*" to search all repos (disable auto-filter). By default, auto-detects current repo from CURRENT_REPO env and filters to it. Use repo=["frontend","backend"] to search related repos together. - - Filters (optional): language, under (path prefix), kind, symbol, ext, path_regex, + - profile: optional search profile ("tests", "config", "code") that applies useful path constraints. + - Filters (optional): language, under (recursive workspace subtree), kind, symbol, ext, path_regex, path_glob (str or list[str]), not_glob (str or list[str]), not_ (negative text), case. + - debug: bool (default false). When true, includes verbose internal fields like + components, rerank_counters, code_signals. Default false saves ~60-80% tokens. Returns: - Dict with keys: - - results: list of {score, path, symbol, start_line, end_line, why[, components][, relations][, related_paths][, snippet]} - - total: int; used_rerank: bool; rerank_counters: dict + - results: list of {score, path, symbol, start_line, end_line[, snippet][, tags][, host_path][, container_path]} + When debug=true, also includes: components, why, relations, related_paths, doc_id, code_id + - total: int; used_rerank: bool - If compact=true (and snippets not requested), results contain only {path,start_line,end_line}. + - If debug=true, response also includes: rerank_counters, code_signals Examples: - path_glob=["scripts/**","**/*.py"], language="python" + - profile="tests" # constrain to test files + - profile="config" # constrain to config files - symbol="context_answer", under="scripts" + - debug=true # Include internal scoring details for query tuning """ sess = require_auth_session_fn(session) if require_auth_session_fn else session @@ -252,11 +962,18 @@ async def _repo_search_impl( case = _extra.get("case") if compact in (None, "") and _extra.get("compact") is not None: compact = _extra.get("compact") + if debug in (None, "") and _extra.get("debug") is not None: + debug = _extra.get("debug") # Optional mode hint: "code_first", "docs_first", "balanced" if ( mode is None or (isinstance(mode, str) and str(mode).strip() == "") ) and _extra.get("mode") is not None: mode = _extra.get("mode") + if ( + profile is None + or (isinstance(profile, str) and str(profile).strip() == "") + ) and _extra.get("profile") is not None: + profile = _extra.get("profile") except Exception: pass @@ -295,7 +1012,8 @@ def _to_str(x, default=""): per_path = _to_int(per_path, 2) include_snippet = _to_bool(include_snippet, True) context_lines = _to_int(context_lines, 2) - # Reranker: default ON; can be disabled via env or client args + # Reranker defaults come from the environment, but an explicit request-level + # opt-in/opt-out should still be respected by MCP/API callers. rerank_env_default = str( os.environ.get("RERANKER_ENABLED", "1") ).strip().lower() in {"1", "true", "yes", "on"} @@ -366,10 +1084,13 @@ def _to_str(x, default=""): except Exception: pass - # 3) Environment default (collection only for now) + # 3) Environment defaults (collection + mode) env_coll = (os.environ.get("DEFAULT_COLLECTION") or os.environ.get("COLLECTION_NAME") or "").strip() if (not coll_hint) and env_coll: coll_hint = env_coll + env_mode = (os.environ.get("REPO_SEARCH_DEFAULT_MODE") or "").strip() + if (not mode_hint) and env_mode: + mode_hint = env_mode # Final fallback env_fallback = (os.environ.get("DEFAULT_COLLECTION") or os.environ.get("COLLECTION_NAME") or "codebase").strip() @@ -390,7 +1111,7 @@ def _to_str(x, default=""): under = under_hint language = _to_str(language, "").strip() - under = _to_str(under, "").strip() + under = _normalize_under_scope(_to_str(under, "").strip()) kind = _to_str(kind, "").strip() symbol = _to_str(symbol, "").strip() path_regex = _to_str(path_regex, "").strip() @@ -414,9 +1135,12 @@ def _to_str_list(x): path_globs = _to_str_list(path_glob) not_globs = _to_str_list(not_glob) + profile = normalize_profile(profile) ext = _to_str(ext, "").strip() not_ = _to_str(not_, "").strip() case = _to_str(case, "").strip() + if profile: + path_globs = append_profile_globs(path_globs, profile) # Normalize repo filter: str, list[str], or "*" (search all) # Default: auto-detect current repo unless REPO_AUTO_FILTER=0 @@ -440,12 +1164,122 @@ def _to_str_list(x): if detected_repo: repo_filter = [detected_repo] + case_sensitive = str(case or "").strip().lower() in { + "sensitive", + "true", + "1", + "yes", + "on", + } + path_globs_norm = [g if case_sensitive else g.lower() for g in path_globs] + not_globs_norm = [g if case_sensitive else g.lower() for g in not_globs] + + def _norm_case(v: str) -> str: + return v if case_sensitive else v.lower() + + def _match_glob(glob_pat: str, path_val: str) -> bool: + import fnmatch as _fnm + if not glob_pat: + return False + p = _norm_case(path_val).replace("\\", "/").strip("/") + if _fnm.fnmatchcase(p, glob_pat): + return True + # Allow repo-relative globs (e.g., scripts/**) to match absolute paths + # by testing suffix windows of the normalized path. + if not glob_pat.startswith("/") and "/" in p: + parts = [seg for seg in p.split("/") if seg] + for i in range(1, len(parts)): + tail = "/".join(parts[i:]) + if _fnm.fnmatchcase(tail, glob_pat): + return True + return False + + def _result_passes_path_filters(item: dict) -> bool: + import re as _re + + path = str(item.get("path") or "") + if not path: + return False + + # Evaluate filters against all known path forms carried by this result. + path_vals = [] + for key in ("path", "rel_path", "client_path", "host_path", "container_path"): + v = item.get(key) + if isinstance(v, str) and v.strip(): + path_vals.append(v.strip().replace("\\", "/")) + if not path_vals: + path_vals = [path] + if path.startswith("/work/"): + path_vals.append(path[len("/work/") :]) + + # Deduplicate while preserving order. + seen = set() + norm_paths = [] + for pv in path_vals: + if pv not in seen: + norm_paths.append(pv) + seen.add(pv) + + if not_: + needle = _norm_case(str(not_)) + if any(needle in _norm_case(pv) for pv in norm_paths): + return False + + if ext: + ext_norm = str(ext).lower().lstrip(".") + if not any(_norm_case(pv).endswith("." + ext_norm) for pv in norm_paths): + return False + + if path_regex: + flags = 0 if case_sensitive else _re.IGNORECASE + try: + if not any(_re.search(path_regex, pv, flags=flags) for pv in norm_paths): + return False + except _re.error as exc: + logger.warning( + "Invalid path_regex filter '%s': %s", + path_regex, + exc, + ) + return False + except Exception as exc: + logger.warning( + "Failed evaluating path_regex filter '%s': %s", + path_regex, + exc, + exc_info=True, + ) + return False + + if path_globs_norm and not any( + _match_glob(g, pv) for g in path_globs_norm for pv in norm_paths + ): + return False + + if not_globs_norm and any( + _match_glob(g, pv) for g in not_globs_norm for pv in norm_paths + ): + return False + + return True + + def _apply_result_filters(items: list[dict]) -> list[dict]: + if not items: + return [] + if not (not_ or path_regex or ext or path_globs_norm or not_globs_norm): + return items + return [it for it in items if _result_passes_path_filters(it)] + compact_raw = compact compact = _to_bool(compact, False) # If snippets are requested, do not compact (we need snippet field in results) if include_snippet: compact = False + # Debug mode: when False (default), strip internal/debug fields from results + # to reduce token bloat. Set debug=True to see components, rerank_counters, etc. + debug = _to_bool(debug, False) + # Default behavior: exclude commit-history docs (which use path=".git") from # generic repo_search calls, unless the caller explicitly asks for git # content. This prevents normal code queries from surfacing commit-index @@ -455,6 +1289,7 @@ def _to_str_list(x): ): if ".git" not in not_globs: not_globs.append(".git") + not_globs_norm = [g if case_sensitive else g.lower() for g in not_globs] # Accept top-level alias `queries` as a drop-in for `query` # Many clients send queries=[...] instead of query=[...] @@ -548,53 +1383,24 @@ def _to_str_list(x): lambda: run_pure_dense_search( query=query_text, limit=eff_limit, + per_path=( + int(per_path) + if (per_path is not None and str(per_path).strip() != "") + else None + ), collection=collection, language=language or None, under=under or None, + kind=kind or None, + symbol=symbol or None, + ext=ext or None, repo=repo_filter, ) ) - # Apply post-filters (path_regex, path_glob, not_glob, not_) that aren't - # supported by run_pure_dense_search's server-side filters - case_sensitive = str(case or "").strip().lower() in {"sensitive", "true", "1", "yes", "on"} - import fnmatch as _fnm - import re as _re - - def _norm_path(p: str) -> str: - return p if case_sensitive else p.lower() - - path_globs_norm = [g if case_sensitive else g.lower() for g in path_globs] - not_globs_norm = [g if case_sensitive else g.lower() for g in not_globs] - path_regex_norm = path_regex or "" - - def _match_glob(glob_pat: str, path_val: str) -> bool: - if not glob_pat: - return False - return _fnm.fnmatchcase(_norm_path(path_val), glob_pat) - for item in items: path = item.get("path") or "" - - # Apply path_regex filter - if path_regex_norm: - flags = 0 if case_sensitive else _re.IGNORECASE - try: - if not _re.search(path_regex_norm, path, flags=flags): - continue - except Exception: - pass - - # Apply path_glob filter - if path_globs_norm and not any(_match_glob(g, path) for g in path_globs_norm): - continue - - # Apply not_glob filter - if not_globs_norm and any(_match_glob(g, path) for g in not_globs_norm): - continue - - # Apply not_ text filter - if not_ and not_.lower() in _norm_path(path): + if not _result_passes_path_filters(item): continue payload = item.get("payload") or {} @@ -640,8 +1446,7 @@ def _match_glob(glob_pat: str, path_val: str) -> bool: rt = 0 if rt > eff_limit: eff_limit = rt - # In-process path_glob/not_glob accept a single string; reduce list inputs safely - print(f"[debug] DEBUG_SEARCH_TIMING={os.environ.get('DEBUG_SEARCH_TIMING', 'not set')}", flush=True) + # In-process path_glob/not_glob accept list inputs. items = await asyncio.to_thread( lambda: run_hybrid_search( queries=queries, @@ -696,7 +1501,8 @@ def _match_glob(glob_pat: str, path_val: str) -> bool: eff_limit = rt cmd = [ "python", - _work_script("hybrid_search.py"), + "-m", + "scripts.hybrid_search", "--limit", str(eff_limit), "--json", @@ -784,81 +1590,8 @@ def _match_glob(glob_pat: str, path_val: str) -> bool: "subprocess": 0, "timeout": 0, "error": 0, - "learning": 0, # Learning-enabled recursive reranker } if rerank_enabled: - # Check for learning reranker mode (learns from ONNX teacher) - use_learning_rerank = str( - os.environ.get("RERANK_LEARNING", "") - ).strip().lower() in {"1", "true", "yes", "on"} - - if use_learning_rerank and json_lines: - try: - from scripts.rerank_recursive import rerank_with_learning - - rq = queries[0] if queries else "" - cand_objs = list(json_lines[: int(rerank_top_n)]) - - # Run learning-enabled reranking (collection-aware for weight isolation) - reranked = rerank_with_learning( - query=rq, - candidates=cand_objs, - limit=int(rerank_return_m), - n_iterations=int(os.environ.get("RERANK_LEARNING_ITERS", "3")), - collection=collection or "default", - ) - - if reranked: - # Format results for output - tmp = [] - for obj in reranked: - # Copy the list to avoid mutating the original object - why_parts = list(obj.get("why", [])) - why_parts.append(f"refine:{obj.get('recursive_iterations', 0)}") - why_parts.append(f"score:{float(obj.get('score', 0)):.3f}") - - # Build components with optional fname_boost - components = (obj.get("components") or {}) | { - "learning_score": float(obj.get("recursive_score", 0)), - "refinement_iterations": int(obj.get("recursive_iterations", 0)), - } - if obj.get("fname_boost"): - components["fname_boost"] = float(obj.get("fname_boost", 0)) - why_parts.append(f"fname:{float(obj.get('fname_boost', 0)):.2f}") - - # Extract benchmark IDs from payload for CoIR/CoSQA - _payload = obj.get("payload") if isinstance(obj, dict) else None - if not isinstance(_payload, dict): - _payload = {} - _doc_id = _payload.get("_id") or _payload.get("code_id") or _payload.get("id") - _code_id = _payload.get("code_id") - - item = { - "score": float(obj.get("score", 0)), - "path": obj.get("path", ""), - "symbol": obj.get("symbol", ""), - "start_line": int(obj.get("start_line") or 0), - "end_line": int(obj.get("end_line") or 0), - "why": why_parts, - "components": components, - # Benchmark IDs (preserved through rerank) - "doc_id": str(_doc_id) if _doc_id is not None else None, - "code_id": str(_code_id) if _code_id is not None else None, - } - # Preserve dual-path metadata - if obj.get("host_path"): - item["host_path"] = obj["host_path"] - if obj.get("container_path"): - item["container_path"] = obj["container_path"] - tmp.append(item) - - if tmp: - results = tmp - used_rerank = True - rerank_counters["learning"] += 1 - except Exception: - pass # Fall through to standard reranking - # Resolve in-process gating once and reuse use_rerank_inproc = str( os.environ.get("RERANK_IN_PROCESS", "") @@ -867,7 +1600,7 @@ def _match_glob(glob_pat: str, path_val: str) -> bool: if use_rerank_inproc and not used_rerank: try: if json_lines: - from scripts.rerank_local import rerank_local as _rr_local # type: ignore + from scripts.rerank_tools.local import rerank_local as _rr_local # type: ignore import concurrent.futures as _fut rq = queries[0] if queries else "" @@ -1028,6 +1761,8 @@ def _doc_for(obj: dict) -> str: "score": float(blended_s), "path": obj.get("path", ""), "symbol": obj.get("symbol", ""), + "kind": obj.get("kind", ""), + "repo": obj.get("repo", ""), "start_line": int(obj.get("start_line") or 0), "end_line": int(obj.get("end_line") or 0), "why": why_parts, @@ -1044,6 +1779,10 @@ def _doc_for(obj: dict) -> str: item["host_path"] = _hostp if _contp: item["container_path"] = _contp + if obj.get("file_hash"): + item["file_hash"] = obj.get("file_hash") + if obj.get("symbol_content_hash"): + item["symbol_content_hash"] = obj.get("symbol_content_hash") tmp.append(item) if tmp: results = tmp @@ -1065,7 +1804,7 @@ def _doc_for(obj: dict) -> str: if not used_rerank: if use_rerank_inproc: try: - from scripts.rerank_local import rerank_in_process # type: ignore + from scripts.rerank_tools.local import rerank_in_process # type: ignore model_name = os.environ.get( "EMBEDDING_MODEL", "BAAI/bge-base-en-v1.5" @@ -1092,7 +1831,8 @@ def _doc_for(obj: dict) -> str: rq = queries[0] if queries else "" rcmd = [ "python", - _work_script("rerank_local.py"), + "-m", + "scripts.rerank_tools.local", "--query", rq, "--topk", @@ -1184,6 +1924,8 @@ def _doc_for(obj: dict) -> str: "score": float(obj.get("score", 0.0)), "path": obj.get("path", ""), "symbol": obj.get("symbol", ""), + "kind": obj.get("kind", ""), + "repo": obj.get("repo", ""), "start_line": int(obj.get("start_line") or 0), "end_line": int(obj.get("end_line") or 0), "why": obj.get("why", []), @@ -1199,6 +1941,10 @@ def _doc_for(obj: dict) -> str: item["host_path"] = _hostp if _contp: item["container_path"] = _contp + if obj.get("file_hash"): + item["file_hash"] = obj.get("file_hash") + if obj.get("symbol_content_hash"): + item["symbol_content_hash"] = obj.get("symbol_content_hash") # Pass-through optional relation hints if obj.get("relations"): item["relations"] = obj.get("relations") @@ -1217,6 +1963,10 @@ def _doc_for(obj: dict) -> str: item["tags"] = obj.get("tags") results.append(item) + # Enforce strict filter semantics regardless of retrieval/rerank branch. + # This closes gaps where fallback rerank paths may bypass path_glob/not_glob. + results = _apply_result_filters(results) + # Mode-aware reordering: nudge core implementation code vs docs and non-core when requested def _is_doc_path(p: str) -> bool: pl = str(p or "").lower() @@ -1346,11 +2096,55 @@ def _is_core_code_item(item: dict) -> bool: other_code.append(it) results = doc_items + core_items + other_code - # Enforce user-requested limit on final result count + # Enforce the public result limit after feedback recall has had a chance to + # contribute candidates. The retrieval/rerank stages may intentionally + # over-fetch before this point. try: _limit_n = int(limit) except Exception: _limit_n = 0 + + # Feedback recall: add a few positively rated targets that ordinary retrieval missed. + _canonical_query = queries[0] if queries else "" + _inject_result_ids(results, _canonical_query) + _remember_result_metadata(results, collection) + _weights = _load_relevance_weights(collection) + try: + _feedback_recall_max = int(os.environ.get("RELEVANCE_RECALL_MAX", "3") or 0) + except Exception: + _feedback_recall_max = 3 + if _feedback_recall_max > 0 and _weights and _limit_n > 0: + _existing_targets = {str(r.get("target_id") or r.get("result_id") or "") for r in results} + _scores = [float(r.get("score", 0) or 0) for r in results] + _base_score = min(_scores) if _scores else 0.0 + _recalled = _feedback_recall_candidates( + collection=collection, + weights=_weights, + existing_target_ids=_existing_targets, + existing_paths={str(r.get("path") or r.get("container_path") or "") for r in results}, + base_score=_base_score, + max_candidates=min(_feedback_recall_max, _limit_n), + repo_filter=repo_filter, + language=language, + under=under, + kind_filter=kind, + symbol_filter=symbol, + ext=ext, + not_=not_, + path_regex=path_regex, + path_globs=path_globs, + not_globs=not_globs, + case_sensitive=case_sensitive, + ) + if _recalled: + _inject_result_ids(_recalled, _canonical_query) + # Reserve room for newly discovered feedback/graph neighbors. This + # is the recall feature's purpose; appending and slicing the old + # top-N would silently discard every recalled candidate. + results = results[: max(0, _limit_n - len(_recalled))] + _recalled + + # Keep the public contract bounded when limit is absent/invalid as well as + # when a caller supplied a normal positive limit. if _limit_n > 0 and len(results) > _limit_n: results = results[:_limit_n] @@ -1437,7 +2231,7 @@ def _read_snip(args): # ─── Filename boost fallback ─────────────────────────────────────────────── # Apply filename-query correlation boost for results that don't have it yet. - # The learning reranker applies fname_boost when enabled; this catches: + # Hybrid/rerank paths may already apply fname_boost; this catches: # - Reranking disabled # - Reranking timed out / failed # - Subprocess hybrid search without reranking @@ -1482,15 +2276,67 @@ def _read_snip(args): # Re-sort results by updated score so fname_boost affects ranking results = sorted(results, key=lambda x: float(x.get("score", 0)), reverse=True) + # ─── Inject result_id for relevance feedback ───────────────────────────── + # result_id is the stable feedback target (symbol/file). impression_id is + # query/content-specific and is diagnostic; boosts apply to target identity. + _inject_result_ids(results, _canonical_query) + _remember_result_metadata(results, collection) + + # ─── Apply learned relevance boosts ───────────────────────────────────── + _relevance_boost_factor = float(os.environ.get("RELEVANCE_BOOST_FACTOR", "0.15")) + if _relevance_boost_factor > 0 and results: + try: + if _weights: + _result_weights = _weights.get("results", {}) + for r in results: + _rid = r.get("feedback_weight_id") or r.get("result_id", "") + if _rid and _rid in _result_weights: + _avg = float(_result_weights[_rid].get("avg_relevance", 0)) + _count = int(_result_weights[_rid].get("count", 0)) + _inheritance = float( + _result_weights[_rid].get("inheritance_weight", 1.0) or 0 + ) + _boost = ( + _relevance_boost_factor + * (_avg / 2.0) + * min(_count, 10) + / 10.0 + * _inheritance + ) + r["score"] = float(r.get("score", 0)) + _boost + r["relevance_boost"] = round(_boost, 4) + results.sort(key=lambda x: float(x.get("score", 0)), reverse=True) + if _limit_n > 0 and len(results) > _limit_n: + results = results[:_limit_n] + except Exception: + pass + if compact: results = [ { + "result_id": r.get("result_id", ""), + "target_id": r.get("target_id", ""), + "impression_id": r.get("impression_id", ""), "path": r.get("path", ""), "start_line": int(r.get("start_line") or 0), "end_line": int(r.get("end_line") or 0), } for r in results ] + elif not debug: + # Strip debug/internal fields from results to reduce token bloat + # Keeps: score, path, host_path, container_path, symbol, snippet, + # start_line, end_line, result_id/target_id/impression_id. + results = [_strip_debug_fields(r) for r in results] + + _res_ok = bool(res.get("ok", True)) if isinstance(res, dict) else True + try: + _res_code = int((res or {}).get("code", 0)) + except Exception: + _res_code = 0 + if results: + _res_ok = True + _res_code = 0 response = { "args": { @@ -1504,6 +2350,7 @@ def _read_snip(args): "rerank_return_m": int(rerank_return_m), "rerank_timeout_ms": int(rerank_timeout_ms), "collection": collection, + "profile": profile, "language": language, "under": under, "kind": kind, @@ -1518,13 +2365,23 @@ def _read_snip(args): "compact": (_to_bool(compact_raw, compact)), }, "used_rerank": bool(used_rerank), - "rerank_counters": rerank_counters, - "code_signals": code_signals if code_signals.get("has_code_signals") else None, "total": len(results), "results": results, - **res, + "ok": _res_ok, + "code": _res_code, } + # Expose a concise failure reason without leaking raw subprocess streams by default. + if (not _res_ok or _res_code != 0) and not results: + response["error"] = "search backend execution failed" + + # Only include debug fields when explicitly requested + if debug: + response["subprocess"] = res + response["rerank_counters"] = rerank_counters + if code_signals.get("has_code_signals"): + response["code_signals"] = code_signals + # Apply TOON formatting if requested or enabled globally # Full mode (compact=False) still saves tokens vs JSON while preserving all fields if _should_use_toon(output_format): diff --git a/scripts/mcp_impl/search_profiles.py b/scripts/mcp_impl/search_profiles.py new file mode 100644 index 00000000..f1f43ea7 --- /dev/null +++ b/scripts/mcp_impl/search_profiles.py @@ -0,0 +1,81 @@ +"""Shared repo_search profile definitions. + +Profiles are intentionally small path constraints, not alternate search +algorithms. They let callers express common scopes without exposing separate +MCP tools for every preset. +""" + +from __future__ import annotations + +from typing import Iterable + +TEST_GLOBS = [ + "tests/**", + "test/**", + "**/*test*.*", + "**/*_test.*", + "**/Test*/**", +] + +CONFIG_GLOBS = [ + "**/*.yml", + "**/*.yaml", + "**/*.json", + "**/*.toml", + "**/*.ini", + "**/*.env", + "**/*.config", + "**/*.conf", + "**/*.properties", + "**/*.csproj", + "**/*.props", + "**/*.targets", + "**/*.xml", + "**/appsettings*.json", +] + +CODE_GLOBS = [ + "**/*.py", + "**/*.js", + "**/*.ts", + "**/*.tsx", + "**/*.jsx", + "**/*.mjs", + "**/*.cjs", + "**/*.go", + "**/*.java", + "**/*.cs", + "**/*.rb", + "**/*.php", + "**/*.rs", + "**/*.c", + "**/*.h", + "**/*.cpp", + "**/*.hpp", +] + +PROFILE_GLOBS = { + "test": TEST_GLOBS, + "tests": TEST_GLOBS, + "config": CONFIG_GLOBS, + "configs": CONFIG_GLOBS, + "code": CODE_GLOBS, +} + + +def normalize_profile(profile: object) -> str: + return str(profile or "").strip().lower().replace("-", "_") + + +def globs_for_profile(profile: object) -> list[str]: + return list(PROFILE_GLOBS.get(normalize_profile(profile), [])) + + +def append_profile_globs(path_globs: Iterable[str], profile: object) -> list[str]: + merged = [str(g).strip() for g in path_globs if str(g).strip()] + seen = set(merged) + for glob in globs_for_profile(profile): + if glob not in seen: + merged.append(glob) + seen.add(glob) + return merged diff --git a/scripts/mcp_impl/search_specialized.py b/scripts/mcp_impl/search_specialized.py deleted file mode 100644 index 62f06ef4..00000000 --- a/scripts/mcp_impl/search_specialized.py +++ /dev/null @@ -1,259 +0,0 @@ -#!/usr/bin/env python3 -""" -mcp/search_specialized.py - Specialized search implementations for MCP indexer server. - -Extracted from mcp_indexer_server.py for better modularity. -Contains: -- _search_tests_for_impl: Search for test files -- _search_config_for_impl: Search for config files -- _search_callers_for_impl: Search for callers/usages -- _search_importers_for_impl: Search for importers - -Note: The @mcp.tool() decorated functions remain in mcp_indexer_server.py -as thin wrappers that call these implementations. -""" - -from __future__ import annotations - -__all__ = [ - "_search_tests_for_impl", - "_search_config_for_impl", - "_search_callers_for_impl", - "_search_importers_for_impl", -] - -import logging -from typing import Any, Dict, Optional - -logger = logging.getLogger(__name__) - -# --------------------------------------------------------------------------- -# Imports from sibling modules -# --------------------------------------------------------------------------- -from scripts.mcp_impl.utils import _extract_kwargs_payload - - -# Test file globs -TEST_GLOBS = [ - "tests/**", - "test/**", - "**/*test*.*", - "**/*_test.*", - "**/Test*/**", -] - -# Config file globs -CONFIG_GLOBS = [ - "**/*.yml", - "**/*.yaml", - "**/*.json", - "**/*.toml", - "**/*.ini", - "**/*.env", - "**/*.config", - "**/*.conf", - "**/*.properties", - "**/*.csproj", - "**/*.props", - "**/*.targets", - "**/*.xml", - "**/appsettings*.json", -] - -# Code file globs for importers -CODE_GLOBS = [ - "**/*.py", - "**/*.js", - "**/*.ts", - "**/*.tsx", - "**/*.jsx", - "**/*.mjs", - "**/*.cjs", - "**/*.go", - "**/*.java", - "**/*.cs", - "**/*.rb", - "**/*.php", - "**/*.rs", - "**/*.c", - "**/*.h", - "**/*.cpp", - "**/*.hpp", -] - - -async def _search_tests_for_impl( - query: Any = None, - limit: Any = None, - include_snippet: Any = None, - context_lines: Any = None, - under: Any = None, - language: Any = None, - session: Any = None, - compact: Any = None, - kwargs: Any = None, - ctx: Any = None, - repo_search_fn=None, -) -> Dict[str, Any]: - """Find test files related to a query. - - What it does: - - Presets common test file globs and forwards to repo_search - - Accepts extra filters via kwargs (e.g., language, under, case) - - Parameters: - - query: str or list[str]; limit; include_snippet/context_lines; under; language; compact - - Returns: repo_search result shape. - """ - globs = list(TEST_GLOBS) - # Allow caller to add more with path_glob kwarg - # Handle kwargs being passed as a string by some MCP clients - _kwargs = _extract_kwargs_payload(kwargs) if kwargs else {} - extra_glob = _kwargs.get("path_glob") - if extra_glob: - if isinstance(extra_glob, (list, tuple)): - globs.extend([str(x) for x in extra_glob]) - else: - globs.append(str(extra_glob)) - - if repo_search_fn is None: - from scripts.mcp_impl.search import _repo_search_impl - repo_search_fn = _repo_search_impl - - return await repo_search_fn( - query=query, - limit=limit, - include_snippet=include_snippet, - context_lines=context_lines, - under=under, - language=language, - path_glob=globs, - session=session, - compact=compact, - ctx=ctx, - kwargs={k: v for k, v in _kwargs.items() if k not in {"path_glob"}}, - ) - - -async def _search_config_for_impl( - query: Any = None, - limit: Any = None, - include_snippet: Any = None, - context_lines: Any = None, - under: Any = None, - session: Any = None, - compact: Any = None, - kwargs: Any = None, - ctx: Any = None, - repo_search_fn=None, -) -> Dict[str, Any]: - """Find likely configuration files for a service/query. - - What it does: - - Presets config file globs (yaml/json/toml/etc.) and forwards to repo_search - - Accepts extra filters via kwargs - - Returns: repo_search result shape. - """ - globs = list(CONFIG_GLOBS) - # Handle kwargs being passed as a string by some MCP clients - _kwargs = _extract_kwargs_payload(kwargs) if kwargs else {} - extra_glob = _kwargs.get("path_glob") - if extra_glob: - if isinstance(extra_glob, (list, tuple)): - globs.extend([str(x) for x in extra_glob]) - else: - globs.append(str(extra_glob)) - - if repo_search_fn is None: - from scripts.mcp_impl.search import _repo_search_impl - repo_search_fn = _repo_search_impl - - return await repo_search_fn( - query=query, - limit=limit, - include_snippet=include_snippet, - context_lines=context_lines, - under=under, - session=session, - path_glob=globs, - compact=compact, - ctx=ctx, - kwargs={k: v for k, v in _kwargs.items() if k not in {"path_glob"}}, - ) - - -async def _search_callers_for_impl( - query: Any = None, - limit: Any = None, - language: Any = None, - session: Any = None, - kwargs: Any = None, - ctx: Any = None, - repo_search_fn=None, -) -> Dict[str, Any]: - """Heuristic search for callers/usages of a symbol. - - When to use: - - You want files that reference/invoke a function/class - - Notes: - - Thin wrapper over repo_search today; pass language or path_glob to narrow - - Returns repo_search result shape - """ - if repo_search_fn is None: - from scripts.mcp_impl.search import _repo_search_impl - repo_search_fn = _repo_search_impl - - return await repo_search_fn( - query=query, - limit=limit, - language=language, - session=session, - ctx=ctx, - kwargs=kwargs, - ) - - -async def _search_importers_for_impl( - query: Any = None, - limit: Any = None, - language: Any = None, - session: Any = None, - kwargs: Any = None, - ctx: Any = None, - repo_search_fn=None, -) -> Dict[str, Any]: - """Find files likely importing or referencing a module/symbol. - - What it does: - - Presets code globs across common languages; forwards to repo_search - - Accepts additional filters via kwargs (e.g., under, case) - - Returns: repo_search result shape. - """ - globs = list(CODE_GLOBS) - # Handle kwargs being passed as a string by some MCP clients - _kwargs = _extract_kwargs_payload(kwargs) if kwargs else {} - extra_glob = _kwargs.get("path_glob") - if extra_glob: - if isinstance(extra_glob, (list, tuple)): - globs.extend([str(x) for x in extra_glob]) - else: - globs.append(str(extra_glob)) - - if repo_search_fn is None: - from scripts.mcp_impl.search import _repo_search_impl - repo_search_fn = _repo_search_impl - - # Forward to repo_search with preset path_glob; caller can still pass other filters - return await repo_search_fn( - query=query, - limit=limit, - language=language, - path_glob=globs, - session=session, - ctx=ctx, - kwargs={k: v for k, v in _kwargs.items() if k not in {"path_glob"}}, - ) diff --git a/scripts/mcp_impl/symbol_graph.py b/scripts/mcp_impl/symbol_graph.py index da518ac4..f3a8f926 100644 --- a/scripts/mcp_impl/symbol_graph.py +++ b/scripts/mcp_impl/symbol_graph.py @@ -20,10 +20,54 @@ import logging import os import re +import time from typing import Any, Dict, List, Optional, Set +from scripts.path_scope import ( + normalize_under as _normalize_under_scope, + metadata_matches_under as _metadata_matches_under, + path_matches_under as _path_matches_under, +) + logger = logging.getLogger(__name__) +try: + from scripts.ingest.graph_edges import GRAPH_COLLECTION_SUFFIX as _GRAPH_SUFFIX +except Exception: + _GRAPH_SUFFIX = "_graph" + +GRAPH_COLLECTION_SUFFIX = _GRAPH_SUFFIX +# Time-based cache: collection -> expiry timestamp (5 minutes TTL) +_MISSING_GRAPH_COLLECTIONS: dict[str, float] = {} +_MISSING_GRAPH_TTL = 300 # 5 minutes + + +def _clean_expired_missing_graphs() -> None: + """Remove expired entries from the missing graph cache.""" + now = time.monotonic() + expired = [coll for coll, expiry in _MISSING_GRAPH_COLLECTIONS.items() if expiry <= now] + for coll in expired: + _MISSING_GRAPH_COLLECTIONS.pop(coll, None) + + +def _is_graph_missing(collection: str) -> bool: + """Check if a graph collection is marked as missing (with expiration).""" + _clean_expired_missing_graphs() + if collection in _MISSING_GRAPH_COLLECTIONS: + return _MISSING_GRAPH_COLLECTIONS.get(collection, 0) > time.monotonic() + return False + + +def _mark_graph_missing(collection: str) -> None: + """Mark a graph collection as missing (with TTL).""" + _MISSING_GRAPH_COLLECTIONS[collection] = time.monotonic() + _MISSING_GRAPH_TTL + + +def _clear_graph_missing(collection: str) -> None: + """Remove a collection from the missing graph cache (e.g., after successful creation).""" + _MISSING_GRAPH_COLLECTIONS.pop(collection, None) + + __all__ = [ "_symbol_graph_impl", "_format_symbol_graph_toon", @@ -105,23 +149,18 @@ def _symbol_variants(symbol: str) -> List[str]: return list(dict.fromkeys(variants)) # Dedupe preserving order def _norm_under(u: Optional[str]) -> Optional[str]: - """Normalize an `under` path to match ingest's stored `metadata.path_prefix` values. + """Normalize user-facing `under` to recursive subtree scope token.""" + return _normalize_under_scope(u) - This mirrors the engine's convention: normalize to a /work/... style path. - Note: `under` in this engine is an exact directory filter (not recursive). - """ - if not u: - return None - s = str(u).strip().replace("\\", "/") - s = "/".join([p for p in s.split("/") if p]) - if not s: - return None - # Normalize to /work/... - if not s.startswith("/"): - v = "/work/" + s - else: - v = "/work/" + s.lstrip("/") if not s.startswith("/work/") else s - return v.rstrip("/") + +def _point_matches_under(pt: Any, under: Optional[str]) -> bool: + if not under: + return True + payload = getattr(pt, "payload", None) or {} + md = payload.get("metadata", payload) + if not isinstance(md, dict): + md = {} + return _metadata_matches_under(md, under) async def _symbol_graph_impl( @@ -142,7 +181,7 @@ async def _symbol_graph_impl( query_type: One of "callers", "definition", "importers" limit: Maximum number of results language: Optional language filter - under: Optional path prefix filter + under: Optional recursive workspace subtree filter collection: Optional collection override session: Optional session ID for collection routing ctx: MCP context (optional) @@ -193,18 +232,32 @@ async def _symbol_graph_impl( results = [] + norm_under = _norm_under(under) + try: if query_type == "callers": - # Find chunks where metadata.calls array contains the symbol (exact match) - results = await _query_array_field( + # Prefer graph edges collection when available (fast keyword filters). + results = await _query_graph_edges_collection( client=client, collection=coll, - field_key="metadata.calls", - value=symbol, + symbol=symbol, + edge_type="calls", limit=limit, language=language, - under=_norm_under(under), + repo_filter=None, + under=norm_under, ) + if not results: + # Fall back to array field lookup in the main collection. + results = await _query_array_field( + client=client, + collection=coll, + field_key="metadata.calls", + value=symbol, + limit=limit, + language=language, + under=norm_under, + ) elif query_type == "definition": # Find chunks where symbol_path matches the symbol results = await _query_definition( @@ -213,19 +266,30 @@ async def _symbol_graph_impl( symbol=symbol, limit=limit, language=language, - under=_norm_under(under), + under=norm_under, ) elif query_type == "importers": - # Find chunks where metadata.imports array contains the symbol - results = await _query_array_field( + results = await _query_graph_edges_collection( client=client, collection=coll, - field_key="metadata.imports", - value=symbol, + symbol=symbol, + edge_type="imports", limit=limit, language=language, - under=_norm_under(under), + repo_filter=None, + under=norm_under, ) + if not results: + # Fall back to array field lookup in the main collection. + results = await _query_array_field( + client=client, + collection=coll, + field_key="metadata.imports", + value=symbol, + limit=limit, + language=language, + under=norm_under, + ) # If no results, fall back to semantic search if not results: @@ -234,6 +298,7 @@ async def _symbol_graph_impl( query_type=query_type, limit=limit, language=language, + under=norm_under, collection=coll, session=session, ) @@ -246,6 +311,7 @@ async def _symbol_graph_impl( query_type=query_type, limit=limit, language=language, + under=norm_under, collection=coll, session=session, ) @@ -259,6 +325,155 @@ async def _symbol_graph_impl( } +async def _query_graph_edges_collection( + client: Any, + collection: str, + symbol: str, + edge_type: str, + limit: int, + language: Optional[str] = None, + repo_filter: str | None = None, + under: str | None = None, +) -> List[Dict[str, Any]]: + """Query `_graph` and hydrate results from the main collection. + + The graph collection stores file-level edges: + - caller_path -> callee_symbol (calls/imports) + """ + from qdrant_client import models as qmodels + + graph_coll = f"{collection}{GRAPH_COLLECTION_SUFFIX}" + if _is_graph_missing(graph_coll): + return [] + + # Build graph filter + must: list[Any] = [ + qmodels.FieldCondition( + key="edge_type", match=qmodels.MatchValue(value=str(edge_type)) + ) + ] + if repo_filter: + rf = str(repo_filter).strip() + if rf and rf != "*": + must.append( + qmodels.FieldCondition(key="repo", match=qmodels.MatchValue(value=rf)) + ) + + # Try exact match, then symbol variants. + callee_variants = _symbol_variants(symbol) or [symbol] + seen_paths: set[str] = set() + caller_paths: List[str] = [] + + for variant in callee_variants: + if len(caller_paths) >= limit: + break + v = str(variant).strip() + if not v: + continue + flt = qmodels.Filter( + must=must + + [ + qmodels.FieldCondition( + key="callee_symbol", match=qmodels.MatchValue(value=v) + ) + ] + ) + + def _scroll(_flt=flt): + return client.scroll( + collection_name=graph_coll, + scroll_filter=_flt, + limit=max(32, limit * 4), + with_payload=True, + with_vectors=False, + ) + + try: + points, _ = await asyncio.to_thread(_scroll) + except Exception as e: + err = str(e).lower() + if "404" in err or "doesn't exist" in err or "not found" in err: + _mark_graph_missing(graph_coll) + return [] + logger.exception( + "_query_graph_edges_collection scroll failed for %s", graph_coll + ) + raise + + for rec in points or []: + payload = getattr(rec, "payload", None) or {} + p = payload.get("caller_path") or "" + if not p: + continue + path_s = str(p) + if under and not _path_matches_under( + path_s, under, repo_hint=(payload.get("repo") or repo_filter) + ): + continue + if path_s in seen_paths: + continue + seen_paths.add(path_s) + caller_paths.append(path_s) + if len(caller_paths) >= limit: + break + + if not caller_paths: + return [] + + # Hydrate caller paths back into normal symbol_graph point-shaped results. + hydrated: List[Dict[str, Any]] = [] + for p in caller_paths[:limit]: + if len(hydrated) >= limit: + break + + def _scroll_main(_p=p, _language=language): + must = [ + qmodels.FieldCondition( + key="metadata.path", match=qmodels.MatchValue(value=_p) + ) + ] + if _language: + must.append( + qmodels.FieldCondition( + key="metadata.language", + match=qmodels.MatchValue(value=str(_language).lower()), + ) + ) + return client.scroll( + collection_name=collection, + scroll_filter=qmodels.Filter( + must=must + ), + limit=1, + with_payload=True, + with_vectors=False, + ) + + try: + pts, _ = await asyncio.to_thread(_scroll_main) + except Exception: + pts = [] + + if pts: + hydrated.append(_format_point(pts[0])) + else: + # If language filtering was requested but no matching main-collection doc + # exists (or hydration failed), skip returning a placeholder to avoid + # producing language-inconsistent results. + if not language: + hydrated.append( + { + "path": p, + "symbol": "", + "symbol_path": "", + "start_line": 0, + "end_line": 0, + } + ) + + return hydrated + + async def _query_array_field( client: Any, collection: str, @@ -290,14 +505,6 @@ async def _query_array_field( match=qmodels.MatchValue(value=language.lower()), ) ) - if under: - base_conditions.append( - qmodels.FieldCondition( - key="metadata.path_prefix", - match=qmodels.MatchValue(value=under), - ) - ) - # Strategy 1: Exact match with MatchAny (most reliable for array fields) try: filter1 = qmodels.Filter( @@ -321,6 +528,8 @@ def scroll1(): scroll_result = await asyncio.to_thread(scroll1) points = scroll_result[0] if scroll_result else [] for pt in points: + if under and not _point_matches_under(pt, under): + continue pt_id = str(getattr(pt, "id", id(pt))) if pt_id not in seen_ids: seen_ids.add(pt_id) @@ -356,6 +565,8 @@ def scroll2(): scroll_result = await asyncio.to_thread(scroll2) points = scroll_result[0] if scroll_result else [] for pt in points: + if under and not _point_matches_under(pt, under): + continue pt_id = str(getattr(pt, "id", id(pt))) if pt_id not in seen_ids: seen_ids.add(pt_id) @@ -387,6 +598,8 @@ def scroll3(): scroll_result = await asyncio.to_thread(scroll3) points = scroll_result[0] if scroll_result else [] for pt in points: + if under and not _point_matches_under(pt, under): + continue pt_id = str(getattr(pt, "id", id(pt))) if pt_id not in seen_ids: seen_ids.add(pt_id) @@ -422,14 +635,6 @@ async def _query_definition( match=qmodels.MatchValue(value=language.lower()), ) ) - if under: - base_conditions.append( - qmodels.FieldCondition( - key="metadata.path_prefix", - match=qmodels.MatchValue(value=under), - ) - ) - # Strategy 1: Exact match on symbol_path (e.g., "MyClass.my_method") try: filter1 = qmodels.Filter( @@ -514,6 +719,8 @@ def scroll3(): seen_ids = set() unique_results = [] for pt in results: + if under and not _point_matches_under(pt, under): + continue pt_id = getattr(pt, "id", None) if pt_id not in seen_ids: seen_ids.add(pt_id) @@ -570,6 +777,7 @@ async def _fallback_semantic_search( query_type: str, limit: int = 20, language: Optional[str] = None, + under: Optional[str] = None, collection: Optional[str] = None, session: Optional[str] = None, ) -> List[Dict[str, Any]]: @@ -591,6 +799,8 @@ async def _fallback_semantic_search( query=query, limit=limit, language=language, + under=under, + collection=collection, session=session, output_format="json", # Avoid TOON encoding for internal calls ) @@ -598,7 +808,7 @@ async def _fallback_semantic_search( # Handle case where results might be TOON-encoded string (shouldn't happen with output_format="json") results = search_result.get("results", []) if isinstance(results, str): - # If somehow still a string, return empty - TOON decoding is not worth it here + # Internal callers require structured rows; skip malformed text-only responses. logger.debug("Fallback search returned TOON-encoded results, skipping") return [] return results @@ -655,7 +865,7 @@ async def _compute_called_by( symbol: The symbol name to find callers for limit: Maximum number of callers to return language: Optional language filter - under: Optional path prefix filter + under: Optional recursive workspace subtree filter collection: Optional collection override Returns: @@ -703,13 +913,6 @@ async def _compute_called_by( ) ) norm_under = _norm_under(under) - if norm_under: - base_conditions.append( - qmodels.FieldCondition( - key="metadata.path_prefix", - match=qmodels.MatchValue(value=norm_under), - ) - ) callers: List[Dict[str, Any]] = [] seen_ids: Set[str] = set() @@ -743,6 +946,8 @@ def do_scroll(): points = scroll_result[0] if scroll_result else [] for pt in points: + if norm_under and not _point_matches_under(pt, norm_under): + continue pt_id = str(getattr(pt, "id", id(pt))) if pt_id in seen_ids: continue diff --git a/scripts/mcp_impl/toon.py b/scripts/mcp_impl/toon.py index f1a0f209..02b5bead 100644 --- a/scripts/mcp_impl/toon.py +++ b/scripts/mcp_impl/toon.py @@ -21,6 +21,8 @@ import os from typing import Any, Dict +from scripts.toon_encoder import encode_context_results, encode_search_results + logger = logging.getLogger(__name__) @@ -51,9 +53,9 @@ def _should_use_toon(output_format: Any) -> bool: # TOON response formatting # --------------------------------------------------------------------------- def _format_results_as_toon(response: Dict[str, Any], compact: bool = False) -> Dict[str, Any]: - """Convert response to use TOON-formatted results string instead of JSON array. + """Add a TOON-formatted render while preserving structured results. - Replaces 'results' array with 'results' string in TOON format to save tokens. + Keeps 'results' as JSON-compatible structured data for machine callers. Always adds output_format marker when TOON is requested, even for empty results. Args: @@ -61,21 +63,14 @@ def _format_results_as_toon(response: Dict[str, Any], compact: bool = False) -> compact: If True, use more compact TOON encoding Returns: - Modified response with TOON-encoded results + Modified response with TOON-encoded text """ try: - from scripts.toon_encoder import encode_search_results - results = response.get("results", []) if isinstance(results, list): - # Replace JSON array with TOON string (handles empty arrays too) - toon_results = encode_search_results(results, compact=compact) - response["results"] = toon_results + response["text"] = encode_search_results(results, compact=compact) response["output_format"] = "toon" - return response - except ImportError: - logger.warning("TOON encoder not available, returning JSON format") return response except Exception as e: logger.debug(f"TOON encoding failed: {e}") @@ -83,32 +78,26 @@ def _format_results_as_toon(response: Dict[str, Any], compact: bool = False) -> def _format_context_results_as_toon(response: Dict[str, Any], compact: bool = False) -> Dict[str, Any]: - """Convert context_search response to TOON format, handling mixed code/memory results. + """Add a TOON render for context_search mixed code/memory results. Uses encode_context_results which properly handles memory entries (content/score) vs code entries (path/line), avoiding blank rows or dropped content. + Keeps 'results' as JSON-compatible structured data for machine callers. Args: response: Context search response dict with 'results' key compact: If True, use more compact TOON encoding Returns: - Modified response with TOON-encoded results + Modified response with TOON-encoded text """ try: - from scripts.toon_encoder import encode_context_results - results = response.get("results", []) if isinstance(results, list): - toon_results = encode_context_results(results, compact=compact) - response["results"] = toon_results + response["text"] = encode_context_results(results, compact=compact) response["output_format"] = "toon" - return response - except ImportError: - logger.warning("TOON encoder not available, returning JSON format") return response except Exception as e: logger.debug(f"TOON encoding failed: {e}") return response - diff --git a/scripts/mcp_impl/workspace.py b/scripts/mcp_impl/workspace.py index 6115f8f6..5b56aced 100644 --- a/scripts/mcp_impl/workspace.py +++ b/scripts/mcp_impl/workspace.py @@ -26,7 +26,6 @@ "_state_file_path", "_read_ws_state", "_default_collection", - "_work_script", ] import json @@ -98,23 +97,3 @@ def _default_collection() -> str: if isinstance(coll, str) and coll.strip(): return coll.strip() return DEFAULT_COLLECTION - - -def _work_script(name: str) -> str: - """Return path to script respecting bind mounts first, then /app, then local fallback.""" - try: - work_path = os.path.join("/work", "scripts", name) - if os.path.exists(work_path): - return work_path - except Exception: - pass - - try: - app_path = os.path.join("/app", "scripts", name) - if os.path.exists(app_path): - return app_path - except Exception: - pass - - return os.path.join(os.getcwd(), "scripts", name) - diff --git a/scripts/mcp_indexer_server.py b/scripts/mcp_indexer_server.py index d12aee9b..5184f2f4 100644 --- a/scripts/mcp_indexer_server.py +++ b/scripts/mcp_indexer_server.py @@ -29,19 +29,8 @@ # CRITICAL: OpenLit must be initialized BEFORE any qdrant_client imports # to properly instrument vector DB calls. This import must come first! # --------------------------------------------------------------------------- -import os as _os -import sys as _sys -_roots_env = _os.environ.get("WORK_ROOTS", "") -_roots = [p.strip() for p in _roots_env.split(",") if p.strip()] or ["/work", "/app"] -for _root in _roots: - if _root and _root not in _sys.path: - _sys.path.insert(0, _root) - -# Now import OpenLit init (before any other scripts imports) -try: - from scripts import openlit_init # noqa: F401 - triggers early instrumentation -except ImportError: - pass # OpenLit not available +from scripts import openlit_init # noqa: F401 - triggers early instrumentation + import json import asyncio import re @@ -68,31 +57,13 @@ def _json_dumps_bytes(obj) -> bytes: from typing import Any, Dict, Optional, List, Tuple from pathlib import Path -import sys - -# Import structured logging and error handling (after sys.path setup) -# Will be imported after sys.path is configured below - -import contextlib -# Ensure code roots are on sys.path so absolute imports like 'from scripts.x import y' work -# when this file is executed directly (sys.path[0] may be /work/scripts). -# Supports multiple roots via WORK_ROOTS env (comma-separated), defaults to /work and /app. -_roots_env = os.environ.get("WORK_ROOTS", "") -_roots = [p.strip() for p in _roots_env.split(",") if p.strip()] or ["/work", "/app"] -try: - for _root in _roots: - if _root and _root not in sys.path: - sys.path.insert(0, _root) -except Exception: - pass +import qdrant_client # Note: OpenLit initialization is handled by early import of scripts.openlit_init # at the top of this file (before any qdrant_client imports) -# Session state imported from mcp_workspace shim (-> scripts.mcp.workspace) -# Must be after sys.path setup -from scripts.mcp_workspace import ( +from scripts.mcp_impl.workspace import ( _MEM_COLL_CACHE, SESSION_DEFAULTS, SESSION_DEFAULTS_BY_SESSION, @@ -100,31 +71,20 @@ def _json_dumps_bytes(obj) -> bytes: _SESSION_CTX_LOCK, ) -# Import structured logging and error handling (after sys.path setup) -try: - from scripts.logger import ( - get_logger, - ContextLogger, - RetrievalError, - IndexingError, - DecoderError, - ValidationError, - ConfigurationError, - safe_int, - safe_float, - safe_bool, - ) - - logger = get_logger(__name__) -except ImportError: - # Fallback if logger module not available - import logging - - logger = logging.getLogger(__name__) - logging.basicConfig(level=logging.INFO) +from scripts.logger import ( + get_logger, + ContextLogger, + RetrievalError, + IndexingError, + DecoderError, + ValidationError, + ConfigurationError, + safe_int, + safe_float, + safe_bool, +) - # Import safe conversion functions from utils (single source of truth) - from scripts.mcp_impl.utils import safe_int, safe_float, safe_bool +logger = get_logger(__name__) from scripts.mcp_auth import ( @@ -135,7 +95,7 @@ def _json_dumps_bytes(obj) -> bytes: # --------------------------------------------------------------------------- # Re-exports from extracted modules (backwards compatibility) # --------------------------------------------------------------------------- -from scripts.mcp_utils import ( +from scripts.mcp_impl.utils import ( _coerce_bool, _coerce_int, _coerce_str, @@ -152,7 +112,7 @@ def _json_dumps_bytes(obj) -> bytes: _primary_identifier_from_queries, ) -from scripts.mcp_toon import ( +from scripts.mcp_impl.toon import ( _is_toon_output_enabled, _should_use_toon, _format_results_as_toon, @@ -162,21 +122,8 @@ def _json_dumps_bytes(obj) -> bytes: # Import implementations from extracted modules from scripts.mcp_impl.context_search import _context_search_impl from scripts.mcp_impl.query_expand import _expand_query_impl -from scripts.mcp_impl.search import _repo_search_impl -from scripts.mcp_impl.info_request import ( - _extract_symbols_from_query, - _extract_related_concepts, - _format_information_field, - _extract_relationships, - _calculate_confidence, -) +from scripts.mcp_impl.search import _repo_search_impl, enrich_feedback_rating from scripts.mcp_impl.admin_tools import _collection_map_impl -from scripts.mcp_impl.search_specialized import ( - _search_tests_for_impl, - _search_config_for_impl, - _search_callers_for_impl, - _search_importers_for_impl, -) from scripts.mcp_impl.search_history import ( _search_commits_for_impl, _change_history_for_path_impl, @@ -185,31 +132,18 @@ def _json_dumps_bytes(obj) -> bytes: _symbol_graph_impl, _format_symbol_graph_toon, ) -from scripts.mcp_impl.pattern_search import _pattern_search_impl # Global lock to guard temporary env toggles used during ReFRAG retrieval/decoding _ENV_LOCK = threading.Lock() # Shared utilities (lex hashing, snippet highlighter) -try: - from scripts.utils import highlight_snippet as _do_highlight_snippet -except Exception as e: - logger.warning(f"Failed to import rich for syntax highlighting: {e}") - _do_highlight_snippet = None # fallback guarded at call site +from scripts.utils import highlight_snippet as _do_highlight_snippet # Back-compat shim for tests expecting _highlight_snippet in this module # Delegates to scripts.utils.highlight_snippet when available -try: - - def _highlight_snippet(snippet, tokens): # type: ignore - return ( - _do_highlight_snippet(snippet, tokens) if _do_highlight_snippet else snippet - ) -except Exception: - - def _highlight_snippet(snippet, tokens): # type: ignore - return snippet +def _highlight_snippet(snippet, tokens): # type: ignore + return _do_highlight_snippet(snippet, tokens) try: @@ -243,7 +177,7 @@ def _highlight_snippet(snippet, tokens): # type: ignore try: from scripts.workspace_state import get_collection_name as _ws_get_collection_name # type: ignore - if DEFAULT_COLLECTION in {"", "default-collection", "my-collection", "codebase"}: + if DEFAULT_COLLECTION in {"", "codebase"}: resolved = _ws_get_collection_name(None) if resolved: DEFAULT_COLLECTION = resolved @@ -281,15 +215,14 @@ def _highlight_snippet(snippet, tokens): # type: ignore ) # Disable strict identifier requirement -# --- TOON functions imported from scripts.mcp_toon --- +# --- TOON functions imported from scripts.mcp_impl.toon --- # (see imports at top of file for backwards compatibility re-exports) -# --- Workspace state functions imported from mcp_workspace shim --- -from scripts.mcp_workspace import ( +# --- Workspace state functions imported from workspace helper module --- +from scripts.mcp_impl.workspace import ( _state_file_path, _read_ws_state, _default_collection, - _work_script, ) # Disable DNS rebinding protection - breaks Docker internal networking (Host: mcp:8000) @@ -300,6 +233,23 @@ def _highlight_snippet(snippet, tokens): # type: ignore ) mcp = FastMCP(APP_NAME, transport_security=_security_settings) +# Minimal resource so MCP clients can verify resource wiring. +@mcp.resource( + "resource://context-engine/indexer/info", + name="context-engine-indexer-info", + title="Context Engine Indexer Info", + description="Basic metadata about the running indexer MCP server.", + mime_type="application/json", +) +def _indexer_info_resource(): + return { + "app": APP_NAME, + "host": HOST, + "port": PORT, + "qdrant_url": QDRANT_URL, + "default_collection": DEFAULT_COLLECTION, + } + # Capture tool registry automatically by wrapping the decorator once _TOOLS_REGISTRY: list[dict] = [] @@ -441,98 +391,26 @@ def log_message(self, *args, **kwargs): return False -# Import the new subprocess manager -try: - from scripts.subprocess_manager import run_subprocess_async -except ImportError: - # Fallback if subprocess_manager not available - logger.warning("subprocess_manager not available, using fallback implementation") - - async def run_subprocess_async( - cmd: List[str], - timeout: Optional[float] = None, - env: Optional[Dict[str, str]] = None, - ) -> Dict[str, Any]: - """Fallback subprocess runner if subprocess_manager is not available.""" - proc: Optional[asyncio.subprocess.Process] = None - try: - proc = await asyncio.create_subprocess_exec( - *cmd, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - env=env, - ) - # Default timeout from env if not provided by caller - if timeout is None: - timeout = MCP_TOOL_TIMEOUT_SECS - try: - stdout_b, stderr_b = await asyncio.wait_for( - proc.communicate(), timeout=timeout - ) - code = proc.returncode - except asyncio.TimeoutError: - try: - proc.kill() - except Exception: - pass - return { - "ok": False, - "code": -1, - "stdout": "", - "stderr": f"Command timed out after {timeout}s", - } - stdout = (stdout_b or b"").decode("utf-8", errors="ignore") - stderr = (stderr_b or b"").decode("utf-8", errors="ignore") - - def _cap_tail(s: str) -> str: - if not s: - return s - return ( - s - if len(s) <= MAX_LOG_TAIL - else ("...[tail truncated]\n" + s[-MAX_LOG_TAIL:]) - ) +from scripts.subprocess_manager import run_subprocess_async - return { - "ok": code == 0, - "code": code, - "stdout": _cap_tail(stdout), - "stderr": _cap_tail(stderr), - } - except Exception as e: - return {"ok": False, "code": -2, "stdout": "", "stderr": str(e)} - finally: - try: - if proc is not None: - if proc.stdout is not None: - proc.stdout.close() - if proc.stderr is not None: - proc.stderr.close() - # Ensure the process is reaped - with contextlib.suppress(Exception): - await proc.wait() - except Exception: - pass - -# --- Admin tool helpers imported from mcp_admin_tools shim --- -from scripts.mcp_admin_tools import ( +# --- Admin tool helpers imported from admin helper module --- +from scripts.mcp_impl.admin_tools import ( _EMBED_MODEL_CACHE, _EMBED_MODEL_LOCKS, _run_async, _get_embedding_model, - _invalidate_router_scratchpad, _detect_current_repo, ) # Lenient argument normalization to tolerate buggy clients (e.g., JSON-in-kwargs, booleans where strings expected) -# Note: _maybe_parse_jsonish and other parsing helpers are now imported from scripts.mcp_utils +# Note: _maybe_parse_jsonish and other parsing helpers are now imported from scripts.mcp_impl.utils from typing import Any as _Any, Dict as _Dict # Extra parsing helpers for quirky clients that send stringified kwargs import urllib.parse as _urlparse, ast as _ast -# --- Utility functions imported from scripts.mcp_utils --- +# --- Utility functions imported from scripts.mcp_impl.utils --- # (see imports at top of file for backwards compatibility re-exports: # _parse_kv_string, _coerce_value_string, _to_str_list_relaxed, # _extract_kwargs_payload, _looks_jsonish_string, _coerce_bool, @@ -588,7 +466,7 @@ async def qdrant_index_root( ) # type: ignore if _ws_is_multi_repo_mode(): - coll = _ws_get_collection_name("/work") or _default_collection() + coll = _default_collection() else: coll = _ws_get_collection_name(None) or _default_collection() except Exception: @@ -600,18 +478,12 @@ async def qdrant_index_root( env["QDRANT_URL"] = QDRANT_URL env["COLLECTION_NAME"] = coll - cmd = ["python", _work_script("ingest_code.py"), "--root", "/work"] + cmd = ["python", "-m", "scripts.ingest_code", "--root", "/work"] if recreate: cmd.append("--recreate") res = await _run_async(cmd, env=env) ret = {"args": {"root": "/work", "collection": coll, "recreate": recreate}, **res} - try: - if ret.get("ok") and int(ret.get("code", 1)) == 0: - if _invalidate_router_scratchpad("/work"): - ret["invalidated_router_scratchpad"] = True - except Exception: - pass return ret @@ -630,17 +502,13 @@ async def qdrant_list(kwargs: Any = None) -> Dict[str, Any]: - {"collections": [str, ...]} or {"error": "..."} """ try: - from qdrant_client import QdrantClient - - client = QdrantClient( + client = qdrant_client.QdrantClient( url=QDRANT_URL, api_key=os.environ.get("QDRANT_API_KEY"), timeout=float(os.environ.get("QDRANT_TIMEOUT", "20") or 20), ) cols_info = await asyncio.to_thread(client.get_collections) return {"collections": [c.name for c in cols_info.collections]} - except ImportError: - return {"error": "qdrant_client is not installed in this container"} except Exception as e: return {"error": str(e)} @@ -759,10 +627,9 @@ async def qdrant_status( pass coll = collection or _default_collection() try: - from qdrant_client import QdrantClient import datetime as _dt - client = QdrantClient( + client = qdrant_client.QdrantClient( url=QDRANT_URL, api_key=os.environ.get("QDRANT_API_KEY"), timeout=float(os.environ.get("QDRANT_TIMEOUT", "20") or 20), @@ -909,7 +776,7 @@ async def qdrant_index( ) # type: ignore if _ws_is_multi_repo_mode(): - coll = _ws_get_collection_name(root) or _default_collection() + coll = _default_collection() else: coll = _ws_get_collection_name(None) or _default_collection() except Exception: @@ -923,7 +790,8 @@ async def qdrant_index( cmd = [ "python", - _work_script("ingest_code.py"), + "-m", + "scripts.ingest_code", "--root", root, ] @@ -932,12 +800,6 @@ async def qdrant_index( res = await _run_async(cmd, env=env) ret = {"args": {"root": root, "collection": coll, "recreate": recreate}, **res} - try: - if ret.get("ok") and int(ret.get("code", 1)) == 0: - if _invalidate_router_scratchpad("/work"): - ret["invalidated_router_scratchpad"] = True - except Exception: - pass return ret @@ -1028,7 +890,7 @@ async def qdrant_prune(kwargs: Any = None, **ignored: Any) -> Dict[str, Any]: env = os.environ.copy() env["PRUNE_ROOT"] = "/work" - cmd = ["python", _work_script("prune.py")] + cmd = ["python", "-m", "scripts.prune"] res = await _run_async(cmd, env=env) return res @@ -1036,7 +898,7 @@ async def qdrant_prune(kwargs: Any = None, **ignored: Any) -> Dict[str, Any]: # --------------------------------------------------------------------------- # Code signal detection imported from mcp_code_signals shim # --------------------------------------------------------------------------- -from scripts.mcp_code_signals import ( +from scripts.mcp_impl.code_signals import ( _CODE_INTENT_CACHE, _CODE_INTENT_LOCK, _CODE_QUERY_ARCHETYPES, @@ -1068,6 +930,7 @@ async def repo_search( collection: Any = None, workspace_path: Any = None, mode: Any = None, + profile: Any = None, session: Any = None, ctx: Context = None, language: Any = None, @@ -1082,6 +945,7 @@ async def repo_search( case: Any = None, repo: Any = None, compact: Any = None, + debug: Any = None, output_format: Any = None, args: Any = None, kwargs: Any = None, @@ -1098,12 +962,14 @@ async def repo_search( - per_path: int (default 2). Max results per file. - include_snippet/context_lines: return inline snippets near hits when true. - rerank_*: ONNX reranker is ON by default for best relevance; timeouts fall back to hybrid. + - profile: Optional useful path profile: tests, config, or code. + - debug: bool (default false). Include verbose internal fields (components, rerank_counters, etc). - output_format: "json" (default) or "toon" for token-efficient TOON format. - collection: str. Target collection; defaults to workspace state or env COLLECTION_NAME. - repo: str or list[str]. Filter by repo name(s). Use "*" to search all repos. Returns: - - Dict with keys: results, total, used_rerank, rerank_counters + - Dict with keys: results, total, used_rerank, [rerank_counters if debug=true] """ return await _repo_search_impl( query=query, @@ -1120,6 +986,7 @@ async def repo_search( collection=collection, workspace_path=workspace_path, mode=mode, + profile=profile, session=session, ctx=ctx, language=language, @@ -1134,6 +1001,7 @@ async def repo_search( case=case, repo=repo, compact=compact, + debug=debug, output_format=output_format, args=args, kwargs=kwargs, @@ -1195,6 +1063,7 @@ async def repo_search_compat(**arguments) -> Dict[str, Any]: "not_": not_value, "case": args.get("case"), "compact": args.get("compact"), + "debug": args.get("debug"), "mode": args.get("mode"), "repo": args.get("repo"), # Cross-codebase isolation "output_format": args.get("output_format"), # "json" or "toon" @@ -1256,138 +1125,8 @@ async def context_answer_compat(arguments: Any = None) -> Dict[str, Any]: # --------------------------------------------------------------------------- -# Specialized search tools - thin wrappers delegating to extracted impls +# symbol_graph - graph query tool # --------------------------------------------------------------------------- -@mcp.tool() -async def search_tests_for( - query: Any = None, - limit: Any = None, - include_snippet: Any = None, - context_lines: Any = None, - under: Any = None, - language: Any = None, - session: Any = None, - compact: Any = None, - kwargs: Any = None, - ctx: Context = None, -) -> Dict[str, Any]: - """Find test files related to a query. - - What it does: - - Presets common test file globs and forwards to repo_search - - Accepts extra filters via kwargs (e.g., language, under, case) - - Parameters: - - query: str or list[str]; limit; include_snippet/context_lines; under; language; compact - - Returns: repo_search result shape. - """ - return await _search_tests_for_impl( - query=query, - limit=limit, - include_snippet=include_snippet, - context_lines=context_lines, - under=under, - language=language, - session=session, - compact=compact, - kwargs=kwargs, - ctx=ctx, - repo_search_fn=repo_search, - ) - - -@mcp.tool() -async def search_config_for( - query: Any = None, - limit: Any = None, - include_snippet: Any = None, - context_lines: Any = None, - under: Any = None, - session: Any = None, - compact: Any = None, - kwargs: Any = None, - ctx: Context = None, -) -> Dict[str, Any]: - """Find likely configuration files for a service/query. - - What it does: - - Presets config file globs (yaml/json/toml/etc.) and forwards to repo_search - - Accepts extra filters via kwargs - - Returns: repo_search result shape. - """ - return await _search_config_for_impl( - query=query, - limit=limit, - include_snippet=include_snippet, - context_lines=context_lines, - under=under, - session=session, - compact=compact, - kwargs=kwargs, - ctx=ctx, - repo_search_fn=repo_search, - ) - - -@mcp.tool() -async def search_callers_for( - query: Any = None, - limit: Any = None, - language: Any = None, - session: Any = None, - kwargs: Any = None, - ctx: Context = None, -) -> Dict[str, Any]: - """Heuristic search for callers/usages of a symbol. - - When to use: - - You want files that reference/invoke a function/class - - Notes: - - Thin wrapper over repo_search today; pass language or path_glob to narrow - - Returns repo_search result shape - """ - return await _search_callers_for_impl( - query=query, - limit=limit, - language=language, - session=session, - kwargs=kwargs, - ctx=ctx, - repo_search_fn=repo_search, - ) - - -@mcp.tool() -async def search_importers_for( - query: Any = None, - limit: Any = None, - language: Any = None, - session: Any = None, - kwargs: Any = None, - ctx: Context = None, -) -> Dict[str, Any]: - """Find files likely importing or referencing a module/symbol. - - What it does: - - Presets code globs across common languages; forwards to repo_search - - Accepts additional filters via kwargs (e.g., under, case) - - Returns: repo_search result shape. - """ - return await _search_importers_for_impl( - query=query, - limit=limit, - language=language, - session=session, - kwargs=kwargs, - ctx=ctx, - repo_search_fn=repo_search, - ) - - @mcp.tool() async def symbol_graph( symbol: str = None, @@ -1395,6 +1134,7 @@ async def symbol_graph( limit: Any = None, language: Any = None, under: Any = None, + collection: Any = None, session: Any = None, output_format: Any = None, ctx: Context = None, @@ -1411,7 +1151,8 @@ async def symbol_graph( - query_type: str. One of "callers", "definition", "importers". - limit: int (default 20). Maximum results to return. - language: str (optional). Filter by programming language. - - under: str (optional). Filter by path prefix. + - under: str (optional). Filter by recursive workspace subtree (e.g., "scripts" -> scripts/**). + - collection: str (optional). Target collection; defaults to env/WS collection. - output_format: "json" (default) or "toon" for token-efficient format. Returns: @@ -1434,6 +1175,7 @@ async def symbol_graph( limit=_limit, language=str(language).strip() if language else None, under=str(under).strip() if under else None, + collection=str(collection).strip() if collection else None, session=str(session).strip() if session else None, ctx=ctx, ) @@ -1510,8 +1252,8 @@ async def change_history_for_path( search_commits_fn=search_commits_for, ) -# --- context_answer helpers imported from mcp_context_answer shim --- -from scripts.mcp_context_answer import ( +# --- context_answer helpers imported from context_answer helper module --- +from scripts.mcp_impl.context_answer import ( _cleanup_answer, _answer_style_guidance, _strip_preamble_labels, @@ -1617,275 +1359,6 @@ async def context_answer( prepare_filters_and_retrieve_fn=_ca_prepare_filters_and_retrieve, ) -@mcp.tool() -async def code_search( - query: Any = None, - limit: Any = None, - per_path: Any = None, - include_snippet: Any = None, - context_lines: Any = None, - rerank_enabled: Any = None, - rerank_top_n: Any = None, - rerank_return_m: Any = None, - rerank_timeout_ms: Any = None, - highlight_snippet: Any = None, - collection: Any = None, - language: Any = None, - under: Any = None, - kind: Any = None, - symbol: Any = None, - path_regex: Any = None, - path_glob: Any = None, - not_glob: Any = None, - ext: Any = None, - not_: Any = None, - case: Any = None, - session: Any = None, - compact: Any = None, - kwargs: Any = None, -) -> Dict[str, Any]: - """Exact alias of repo_search (hybrid code search with reranking enabled by default). - - Prefer repo_search; this name exists for discoverability in some IDEs/agents. - Same parameters and return shape as repo_search. - Reranking (rerank_enabled=true) is ON by default for optimal result quality. - """ - return await repo_search( - query=query, - limit=limit, - per_path=per_path, - include_snippet=include_snippet, - context_lines=context_lines, - rerank_enabled=rerank_enabled, - rerank_top_n=rerank_top_n, - rerank_return_m=rerank_return_m, - rerank_timeout_ms=rerank_timeout_ms, - highlight_snippet=highlight_snippet, - collection=collection, - language=language, - under=under, - kind=kind, - symbol=symbol, - path_regex=path_regex, - path_glob=path_glob, - not_glob=not_glob, - ext=ext, - not_=not_, - case=case, - session=session, - compact=compact, - kwargs=kwargs, - ) - - -# --------------------------------------------------------------------------- -# info_request: Simplified codebase retrieval with explanation mode -# (helpers imported from scripts.mcp_impl.info_request) -# --------------------------------------------------------------------------- -@mcp.tool() -async def info_request( - # Primary parameter - info_request: str = None, - information_request: str = None, # Alias - # Explanation mode - include_explanation: bool = None, - # Relationship mapping - include_relationships: bool = None, - # Auth/session (passed through to repo_search) - session: str = None, - # Optional filters (pass-through to repo_search) - limit: int = None, - language: str = None, - under: str = None, - repo: Any = None, - path_glob: Any = None, - # Additional options - include_snippet: bool = None, - context_lines: int = None, - # Output format - output_format: Any = None, # "json" (default) or "toon" for token-efficient format - kwargs: Any = None, -) -> Dict[str, Any]: - """Simplified codebase retrieval with optional explanation mode. - - When to use: - - Simple, single-parameter code search with human-readable descriptions - - When you want optional explanation mode for richer context - - Drop-in replacement for basic codebase retrieval tools - - Key parameters: - - info_request: str. Natural language description of the code you're looking for. - - information_request: str. Alias for info_request. - - include_explanation: bool (default false). Add summary, primary_locations, related_concepts. - - include_relationships: bool (default false). Add imports_from, calls, related_paths to results. - - limit: int (default 10). Maximum results to return. - - language: str. Filter by programming language. - - under: str. Limit search to specific directory. - - repo: str or list[str]. Filter by repository name(s). - - output_format: "json" (default) or "toon" for token-efficient TOON format. - - Returns: - - Compact mode (default): results with information field and relevance_score alias - - Explanation mode: adds summary, primary_locations, related_concepts, query_understanding - - Example: - - {"info_request": "database connection pooling"} - - {"info_request": "authentication middleware", "include_explanation": true} - """ - # Resolve query from either parameter - query = info_request or information_request - if not query or not str(query).strip(): - return {"ok": False, "error": "info_request parameter is required", "results": []} - query = str(query).strip() - - # Resolve defaults from env - _default_limit = safe_int( - os.environ.get("INFO_REQUEST_LIMIT", "10"), default=10, logger=logger - ) - _default_context = safe_int( - os.environ.get("INFO_REQUEST_CONTEXT_LINES", "5"), default=5, logger=logger - ) - _default_explain = str( - os.environ.get("INFO_REQUEST_EXPLAIN_DEFAULT", "0") - ).strip().lower() in {"1", "true", "yes", "on"} - _default_relationships = str( - os.environ.get("INFO_REQUEST_RELATIONSHIPS", "0") - ).strip().lower() in {"1", "true", "yes", "on"} - - # Apply defaults - eff_limit = limit if limit is not None else _default_limit - eff_context = context_lines if context_lines is not None else _default_context - eff_snippet = include_snippet if include_snippet is not None else True - eff_explain = include_explanation if include_explanation is not None else _default_explain - eff_relationships = include_relationships if include_relationships is not None else _default_relationships - - # Smart limits based on query characteristics (only if user didn't override) - if limit is None: - query_words = len(query.split()) - query_lower = query.lower() - if query_words <= 2: # Short query like "auth handler" - eff_limit = 15 # More results for broad queries - elif "how does" in query_lower or "what is" in query_lower: - eff_limit = 8 # Questions need focused results - - # Call repo_search (always JSON - we format TOON ourselves after enhancement) - search_result = await repo_search( - query=query, - limit=eff_limit, - per_path=3, # Better default for info requests - session=session, - include_snippet=eff_snippet, - context_lines=eff_context, - language=language, - under=under, - repo=repo, - path_glob=path_glob, - output_format="json", # Always get JSON to iterate results - kwargs=kwargs, - ) - - # Extract results - results = search_result.get("results", []) - total = search_result.get("total", len(results)) - used_rerank = search_result.get("used_rerank", False) - - # Enhance each result with information field and optional relationships - enhanced_results = [] - for r in results: - enhanced = dict(r) - enhanced["information"] = _format_information_field(r) - enhanced["relevance_score"] = r.get("score", 0.0) # Alias - # Add relationships if requested - if eff_relationships: - enhanced["relationships"] = _extract_relationships(r) - enhanced_results.append(enhanced) - - # Build better search strategy string - strategy_parts = ["hybrid"] - if used_rerank: - strategy_parts.append("rerank") - if repo: - strategy_parts.append("repo_filtered") - if language: - strategy_parts.append(f"lang:{language}") - if under: - strategy_parts.append("path_filtered") - search_strategy = "+".join(strategy_parts) - - # Build response - response: Dict[str, Any] = { - "ok": True, - "results": enhanced_results, - "total": total, - "search_strategy": search_strategy, - } - - # Add explanation if requested - if eff_explain: - # Primary locations: unique file paths - seen_paths = set() - primary_locations = [] - for r in results: - p = r.get("path", "") - if p and p not in seen_paths: - seen_paths.add(p) - primary_locations.append(p) - if len(primary_locations) >= 5: - break - - # Related concepts - related_concepts = _extract_related_concepts(query, results) - - # Detected symbols from query - detected_symbols = _extract_symbols_from_query(query) - - # Summary - n_files = len(seen_paths) - summary = f"Found {total} results related to '{query}' across {n_files} file{'s' if n_files != 1 else ''}" - - # Group results by file - files_map: Dict[str, list] = {} - for r in enhanced_results: - p = r.get("path", "") - if p not in files_map: - files_map[p] = [] - files_map[p].append({ - "symbol": r.get("symbol", ""), - "line": r.get("start_line", 0), - "score": r.get("score", 0.0), - }) - - grouped_results = { - "by_file": { - path: { - "count": len(items), - "top_symbols": [i["symbol"] for i in sorted(items, key=lambda x: -x["score"])[:3] if i["symbol"]], - } - for path, items in files_map.items() - } - } - - # Calculate confidence - confidence = _calculate_confidence(query, enhanced_results) - - response["summary"] = summary - response["primary_locations"] = primary_locations - response["related_concepts"] = related_concepts - response["grouped_results"] = grouped_results - response["confidence"] = confidence - response["query_understanding"] = { - "intent": "search_for_code", - "detected_language": language or None, - "detected_symbols": detected_symbols, - "search_strategy": search_strategy, - } - - # Apply TOON formatting if requested or enabled globally - if _should_use_toon(output_format): - return _format_results_as_toon(response, compact=False) # Keep info_request fields - return response - - # --------------------------------------------------------------------------- # context_search - thin wrapper delegating to _context_search_impl # --------------------------------------------------------------------------- @@ -2004,80 +1477,140 @@ async def expand_query( return await _expand_query_impl(query=query, max_new=max_new, session=session) -# --------------------------------------------------------------------------- -# Pattern Search - Structural code similarity (conditional on PATTERN_VECTORS=1) -# --------------------------------------------------------------------------- -_PATTERN_SEARCH_ENABLED = str(os.environ.get("PATTERN_VECTORS", "")).strip().lower() in { - "1", "true", "yes", "on" -} - -if _PATTERN_SEARCH_ENABLED: - @mcp.tool() - async def pattern_search( - query: Any = None, - language: Any = None, - limit: Any = None, - min_score: Any = None, - include_snippet: Any = None, - context_lines: Any = None, - target_languages: Any = None, - output_format: Any = None, - compact: Any = None, - aroma_rerank: Any = None, - aroma_alpha: Any = None, - query_mode: Any = None, - ) -> Dict[str, Any]: - """Find structurally similar code patterns across all languages. - - Accepts EITHER code examples OR natural language descriptions - auto-detects which. - - When to use: - - Find code with similar control flow (retry loops, error handling, etc.) - - Cross-language pattern matching (Python pattern → Go/Rust/Java matches) - - Detect code duplication based on structure, not syntax - - Search by pattern description ("retry with backoff", "resource cleanup") - - Key parameters: - - query: str. Code snippet OR natural language description of pattern. - - query_mode: str. "code", "description", or "auto" (default). Explicit override for detection. - - language: str. Language hint for code examples (also triggers code mode in auto). - - limit: int (default 10). Maximum results to return. - - min_score: float (default 0.3). Minimum similarity score threshold. - - include_snippet: bool (default false). Include code snippets in results. - - target_languages: list[str]. Filter to specific target languages. - - output_format: "json" (default) or "toon" for token-efficient format. - - compact: bool. If true with TOON, use minimal fields. - - aroma_rerank: bool (default true). Enable AROMA-style pruning and reranking. - - aroma_alpha: float (default 0.6). Weight for pruned similarity vs original score. - - Returns: - - {ok, results: [{path, start_line, end_line, score, language, ...}], total, query_signature} - - Examples: - - pattern_search(query="for i in range(3): try: ... except: time.sleep(2**i)") - - pattern_search(query="retry with exponential backoff", query_mode="description") - - pattern_search(query="if err != nil { return err }", language="go") - """ - return await _pattern_search_impl( - query=query, - language=language, - limit=limit, - min_score=min_score, - include_snippet=include_snippet, - context_lines=context_lines, - hybrid=None, - semantic_weight=None, - collection=None, - target_languages=target_languages, - output_format=output_format, - compact=compact, - aroma_rerank=aroma_rerank, - aroma_alpha=aroma_alpha, - query_mode=query_mode, - coerce_bool_fn=_coerce_bool, - coerce_int_fn=_coerce_int, - coerce_float_fn=lambda v, d: safe_float(v, default=d, logger=logger, context="pattern_search"), +@mcp.tool() +async def rate_search_results( + query: str, + ratings: list, + collection: Optional[str] = None, + session: Optional[str] = None, +) -> Dict[str, Any]: + """Provide relevance feedback on search results to improve future rankings. + + Call this after using repo_search to tell the system which results were useful. + The system learns per-collection ranking weights from your feedback. + + Parameters: + - query: str. The original search query these ratings apply to. + - ratings: list of {result_id, relevance}. Each entry rates one search result. + - result_id: str. The result_id from a repo_search result entry. + - relevance: int. 0=not used, 1=glanced/relevant, 2=directly used/excellent match. + - target_id, impression_id, path, container_path, symbol, kind, repo, file_hash: + optional. Usually omitted; the server fills them from recent repo_search results. + - related_symbols: list[str] (optional). Logged for future graph experiments; not used by the current trainer. + - collection: str (optional). Target collection; defaults to workspace state. + - session: str (optional). Auth session token. + + Returns: + - {ok: true, rated: N} on success + - {ok: false, error: "..."} on failure + + Example: + repo_search returns result_id "abc123" + rate_search_results( + query="process events", + ratings=[{ + "result_id": "abc123", + "relevance": 2 + }] ) + → rates abc123 at 2 for this collection's future searches. + """ + import json as _json + import time as _time + + sess = _require_auth_session(session) + + if not query or not str(query).strip(): + return {"ok": False, "error": "query is required"} + + if not ratings or not isinstance(ratings, list): + return {"ok": False, "error": "ratings must be a non-empty list"} + + try: + _c = (collection or "").strip() + except Exception: + _c = "" + if _c: + coll = _c + else: + try: + from scripts.workspace_state import ( + get_collection_name as _ws_get_collection_name, + is_multi_repo_mode as _ws_is_multi_repo_mode, + ) + if _ws_is_multi_repo_mode(): + coll = _default_collection() + else: + coll = _ws_get_collection_name(None) or _default_collection() + except Exception: + coll = _default_collection() + + _require_collection_access((sess or {}).get("user_id") if sess else None, coll, "write") + + validated_ratings = [] + for r in ratings: + if not isinstance(r, dict): + continue + r = enrich_feedback_rating(r, coll) + result_id = str(r.get("result_id", "")).strip() + relevance = r.get("relevance") + if not result_id or relevance is None: + continue + try: + relevance = int(relevance) + except (ValueError, TypeError): + continue + if relevance not in (0, 1, 2): + continue + entry = {"result_id": result_id, "relevance": relevance} + for key in ( + "target_id", + "impression_id", + "path", + "host_path", + "container_path", + "symbol", + "kind", + "repo", + "file_hash", + "symbol_content_hash", + ): + val = r.get(key) + if val is not None and str(val).strip(): + entry[key] = str(val).strip() + related = r.get("related_symbols") + if isinstance(related, list) and len(related) > 0: + entry["related_symbols"] = [str(s) for s in related[:10]] + validated_ratings.append(entry) + + if not validated_ratings: + return {"ok": False, "error": "no valid ratings provided"} + + try: + from scripts.rerank_tools.events import _ensure_events_dir, append_event_line + from datetime import datetime as _datetime + + events_dir = _ensure_events_dir() + safe_coll = "".join(c if c.isalnum() or c in "-_" else "_" for c in coll) + hour_suffix = _datetime.now(tz=None).strftime("%Y%m%d%H") + events_file = events_dir / f"events_{safe_coll}_{hour_suffix}.ndjson" + + event = { + "ts": _time.time(), + "type": "relevance_feedback", + "query": str(query).strip(), + "collection": coll, + "ratings": validated_ratings, + "source": "mcp_tool", + } + if sess: + event["session_user"] = (sess or {}).get("user_id", "anonymous") + + append_event_line(events_file, _json.dumps(event)) + + return {"ok": True, "rated": len(validated_ratings), "collection": coll} + except Exception as e: + return {"ok": False, "error": f"Failed to log feedback: {e}"} _relax_var_kwarg_defaults() @@ -2102,7 +1635,6 @@ async def pattern_search( logger.info(f" Embedding Model: {os.environ.get('EMBEDDING_MODEL', 'BAAI/bge-base-en-v1.5')}") logger.info(f" Embedding Provider: {os.environ.get('EMBEDDING_PROVIDER', 'fastembed')}") logger.info(f" ReFRAG Decoder: {os.environ.get('REFRAG_DECODER', '1')}") - logger.info(f" Rerank Learning: {os.environ.get('RERANK_LEARNING', '1')}") logger.info(f" Semantic Chunks: {os.environ.get('INDEX_SEMANTIC_CHUNKS', '1')}") logger.info(f" Micro Chunks: {os.environ.get('INDEX_MICRO_CHUNKS', '1')}") logger.info(f" Micro Chunk Tokens: {os.environ.get('MICRO_CHUNK_TOKENS', '128')}") @@ -2118,7 +1650,7 @@ async def pattern_search( logger.info(f" Reranker Enabled: {os.environ.get('RERANKER_ENABLED', '0')}") logger.info(f" Rerank Top N: {os.environ.get('RERANK_TOP_N', '20')}") logger.info(f" Rerank Timeout MS: {os.environ.get('RERANK_TIMEOUT_MS', '500')}") - logger.info(f" Pattern Search: {'enabled' if _PATTERN_SEARCH_ENABLED else 'disabled (set PATTERN_VECTORS=1)'}") + logger.info(f" TOON Enabled: {os.environ.get('TOON_ENABLED', '0')}") logger.info("=" * 60) # Optional warmups: gated by env flags to avoid delaying readiness on fresh containers @@ -2153,7 +1685,7 @@ async def pattern_search( "on", }: try: - from scripts.rerank_local import _get_rerank_session # type: ignore + from scripts.rerank_tools.local import _get_rerank_session # type: ignore _ = _get_rerank_session() except Exception: @@ -2165,7 +1697,8 @@ async def pattern_search( _env["COLLECTION_NAME"] = _default_collection() _cmd = [ "python", - "/work/scripts/rerank_local.py", + "-m", + "scripts.rerank_tools.local", "--query", "warmup", "--topk", diff --git a/scripts/mcp_memory_server.py b/scripts/mcp_memory_server.py index 78712c10..896c9017 100644 --- a/scripts/mcp_memory_server.py +++ b/scripts/mcp_memory_server.py @@ -3,20 +3,8 @@ # to properly instrument vector DB calls. # --------------------------------------------------------------------------- import os -import sys as _sys -# Ensure repo roots are importable so 'scripts' resolves inside container -_roots_env = os.environ.get("WORK_ROOTS", "") -_roots = [p.strip() for p in _roots_env.split(",") if p.strip()] or ["/work", "/app"] -for _root in _roots: - if _root and _root not in _sys.path: - _sys.path.insert(0, _root) - -# Now import OpenLit init (before any other scripts imports that may use qdrant) -try: - from scripts import openlit_init # noqa: F401 - triggers early instrumentation -except ImportError: - pass # OpenLit not available +from scripts import openlit_init # noqa: F401 - triggers early instrumentation import json import threading @@ -41,24 +29,28 @@ from qdrant_client import QdrantClient, models -# Import connection pooling for proper resource management -try: - from scripts.qdrant_client_manager import ( - get_qdrant_client, - return_qdrant_client, - pooled_qdrant_client, - ) - _POOL_AVAILABLE = True -except ImportError: - _POOL_AVAILABLE = False +from scripts.qdrant_client_manager import ( + get_qdrant_client, + return_qdrant_client, +) + +def _env_flag(name: str, default: bool = False) -> bool: + raw = os.environ.get(name) + if raw is None: + return default + return str(raw).strip().lower() in {"1", "true", "yes", "on"} + + +def _resolve_default_collection() -> str: + raw = (os.environ.get("DEFAULT_COLLECTION") or os.environ.get("COLLECTION_NAME") or "").strip() + if _env_flag("MULTI_REPO_MODE") and raw in {"", "codebase"}: + return "" + return raw or "codebase" + # Env QDRANT_URL = os.environ.get("QDRANT_URL", "http://qdrant:6333") -DEFAULT_COLLECTION = ( - os.environ.get("DEFAULT_COLLECTION") - or os.environ.get("COLLECTION_NAME") - or "codebase" -) +DEFAULT_COLLECTION = _resolve_default_collection() LEX_VECTOR_NAME = os.environ.get("LEX_VECTOR_NAME", "lex") LEX_VECTOR_DIM = int(os.environ.get("LEX_VECTOR_DIM", "4096") or 4096) EMBEDDING_MODEL = os.environ.get("EMBEDDING_MODEL", "BAAI/bge-base-en-v1.5") @@ -76,9 +68,9 @@ # I/O-safety knobs for memory server behavior # These env vars allow tuning startup latency vs. first-call latency, especially important # on slow storage backends (e.g., Ceph + HDD). See comments below for rationale. -MEMORY_ENSURE_ON_START = str(os.environ.get("MEMORY_ENSURE_ON_START", "1")).strip().lower() in {"1", "true", "yes", "on"} -MEMORY_COLD_SKIP_DENSE = str(os.environ.get("MEMORY_COLD_SKIP_DENSE", "0")).strip().lower() in {"1", "true", "yes", "on"} -MEMORY_PROBE_EMBED_DIM = str(os.environ.get("MEMORY_PROBE_EMBED_DIM", "1")).strip().lower() in {"1", "true", "yes", "on"} +MEMORY_ENSURE_ON_START = _env_flag("MEMORY_ENSURE_ON_START", False) +MEMORY_COLD_SKIP_DENSE = _env_flag("MEMORY_COLD_SKIP_DENSE", False) +MEMORY_PROBE_EMBED_DIM = _env_flag("MEMORY_PROBE_EMBED_DIM", True) try: MEMORY_VECTOR_DIM = int(os.environ.get("MEMORY_VECTOR_DIM") or os.environ.get("EMBED_DIM") or "768") except Exception: @@ -90,15 +82,8 @@ # Use the centralized embedder from scripts.embedder for consistent caching. # This eliminates duplicate model loading and ensures consistent behavior. -# Reference to the centralized embedder for cold-skip detection -try: - from scripts.embedder import get_embedding_model as _centralized_get_embedding_model - from scripts.embedder import is_model_cached as _is_model_cached - _EMBEDDER_AVAILABLE = True -except ImportError: - _EMBEDDER_AVAILABLE = False - def _is_model_cached(model_name: str = "") -> bool: # type: ignore[misc] - return False # Fallback: assume not cached +from scripts.embedder import get_embedding_model as _centralized_get_embedding_model +from scripts.embedder import is_model_cached as _is_model_cached def _get_embedding_model(): """Get the embedding model using the centralized embedder. @@ -108,12 +93,7 @@ def _get_embedding_model(): - Qwen3 model support with feature flags - Automatic cache invalidation on corrupted downloads """ - if _EMBEDDER_AVAILABLE: - return _centralized_get_embedding_model(EMBEDDING_MODEL) - - # Fallback for environments without centralized embedder (rare) - from fastembed import TextEmbedding - return TextEmbedding(model_name=EMBEDDING_MODEL) + return _centralized_get_embedding_model(EMBEDDING_MODEL) # Track ensured collections to reduce redundant ensure calls. # RATIONALE: Avoid repeated Qdrant network calls for the same collection. @@ -270,31 +250,19 @@ def log_message(self, *args, **kwargs): # --------------------------------------------------------------------------- # Qdrant Client Management # --------------------------------------------------------------------------- -# Use connection pooling when available, fallback to creating clients on-demand. -# This prevents socket exhaustion under load and improves connection reuse. - def _get_qdrant_client() -> QdrantClient: - """Get a Qdrant client from pool or create one.""" - if _POOL_AVAILABLE: - return get_qdrant_client( - url=QDRANT_URL, - api_key=os.environ.get("QDRANT_API_KEY") - ) - return QdrantClient(url=QDRANT_URL, api_key=os.environ.get("QDRANT_API_KEY")) + """Get a Qdrant client from the shared pool.""" + return get_qdrant_client( + url=QDRANT_URL, + api_key=os.environ.get("QDRANT_API_KEY") + ) def _return_qdrant_client(client: QdrantClient): """Return a client to the pool, or close it if pooling unavailable.""" if client is None: return - if _POOL_AVAILABLE: - return_qdrant_client(client) - else: - # Fallback path: close client to avoid socket leak - try: - client.close() - except Exception: - pass # Best effort cleanup + return_qdrant_client(client) # Ensure collection exists with dual vectors @@ -303,13 +271,9 @@ def _return_qdrant_client(client: QdrantClient): def _ensure_collection(name: str): """Create collection if missing. - Default behavior mirrors the original implementation for PR compatibility: - - Probe the embedding model to detect the dense vector dimension (MEMORY_PROBE_EMBED_DIM=1) - - Eager ensure on startup (MEMORY_ENSURE_ON_START=1) - For slow storage backends (e.g., Ceph + HDD), set the following in your env: - MEMORY_PROBE_EMBED_DIM=0 -> skip model probing; use MEMORY_VECTOR_DIM/EMBED_DIM - - MEMORY_ENSURE_ON_START=0 -> ensure lazily on first tool call + - MEMORY_ENSURE_ON_START=1 -> eagerly create DEFAULT_COLLECTION at startup """ client = _get_qdrant_client() try: @@ -379,9 +343,8 @@ def _ensure_collection(name: str): _return_qdrant_client(client) -# Optional eager collection ensure on startup (enabled by default for backward compatibility). -# Set MEMORY_ENSURE_ON_START=0 to defer ensure to first tool call (recommended on slow storage). -if MEMORY_ENSURE_ON_START: +# Optional eager collection ensure for single-collection deployments. +if MEMORY_ENSURE_ON_START and DEFAULT_COLLECTION: try: _ensure_collection(DEFAULT_COLLECTION) except Exception: @@ -767,7 +730,10 @@ def _resolve_collection( except Exception: pass - return coll or DEFAULT_COLLECTION + resolved = coll or DEFAULT_COLLECTION + if not resolved: + raise ValueError("collection is required in multi-repo memory server mode") + return resolved if __name__ == "__main__": diff --git a/scripts/mcp_router.py b/scripts/mcp_router.py deleted file mode 100644 index c872954a..00000000 --- a/scripts/mcp_router.py +++ /dev/null @@ -1,157 +0,0 @@ -#!/usr/bin/env python3 -""" -Backwards-compatibility shim for mcp_router. - -NOTE: This file is NOT used by Python when scripts/mcp_router/ package exists. -Python's import system prioritizes packages (directories with __init__.py) over -modules (.py files) with the same name. This file exists for: - -1. Documentation: Shows all available exports at a glance -2. Symmetry: Matches the pattern used by rerank_recursive.py -3. Fallback: Would work if the package directory were removed - -All imports resolve to scripts.mcp_router/ (the package): - from scripts.mcp_router import build_plan, classify_intent - # or - from scripts import mcp_router - mcp_router.build_plan("query") - -Usage: - python -m scripts.mcp_router --plan "How do I ...?" - python -m scripts.mcp_router --run "What is hybrid search?" -""" -from __future__ import annotations - -# Re-export everything from the package -from scripts.mcp_router import ( - # Config - HTTP_URL_INDEXER, - HTTP_URL_MEMORY, - DEFAULT_HTTP_URL, - HEALTH_PORT_INDEXER, - HEALTH_PORT_MEMORY, - LANGS, - cache_ttl_sec, - scratchpad_ttl_sec, - divergence_thresholds, - divergence_is_fatal_for, - # Intent constants - INTENT_ANSWER, - INTENT_SEARCH, - INTENT_SEARCH_TESTS, - INTENT_SEARCH_CONFIG, - INTENT_SEARCH_CALLERS, - INTENT_SEARCH_IMPORTERS, - INTENT_MEMORY_STORE, - INTENT_MEMORY_FIND, - INTENT_INDEX, - INTENT_PRUNE, - INTENT_STATUS, - INTENT_LIST, - # Intent functions - classify_intent, - get_last_intent_debug, - _classify_intent_rules, - # Memory - parse_memory_store_payload, - # Client - call_tool_http, - is_failure_response, - discover_tool_endpoints, - default_tool_endpoints, - tools_describe_cached, - _mcp_handshake, - _post_raw, - _post_raw_retry, - _parse_stream_or_json, - _filter_args, - # Scratchpad - scratchpad_path, - load_scratchpad, - save_scratchpad, - looks_like_repeat, - looks_like_same_filters, - looks_like_expand, - # Hints - parse_repo_hints, - clean_query_and_dsl, - select_best_search_tool_by_signature, - # Batching - BatchingContextAnswerClient, - get_batch_client, - # Validation - is_result_good, - extract_metric_from_resp, - material_drop, - # Planning - build_plan, - # CLI - main, - # Legacy aliases - _is_failure_response, - _is_result_good, - _discover_tool_endpoints, -) - -__all__ = [ - # Config - "HTTP_URL_INDEXER", - "HTTP_URL_MEMORY", - "DEFAULT_HTTP_URL", - "HEALTH_PORT_INDEXER", - "HEALTH_PORT_MEMORY", - "LANGS", - "cache_ttl_sec", - "scratchpad_ttl_sec", - "divergence_thresholds", - "divergence_is_fatal_for", - # Intent - "INTENT_ANSWER", - "INTENT_SEARCH", - "INTENT_SEARCH_TESTS", - "INTENT_SEARCH_CONFIG", - "INTENT_SEARCH_CALLERS", - "INTENT_SEARCH_IMPORTERS", - "INTENT_MEMORY_STORE", - "INTENT_MEMORY_FIND", - "INTENT_INDEX", - "INTENT_PRUNE", - "INTENT_STATUS", - "INTENT_LIST", - "classify_intent", - "get_last_intent_debug", - # Memory - "parse_memory_store_payload", - # Client - "call_tool_http", - "is_failure_response", - "discover_tool_endpoints", - "default_tool_endpoints", - "tools_describe_cached", - # Scratchpad - "scratchpad_path", - "load_scratchpad", - "save_scratchpad", - "looks_like_repeat", - "looks_like_same_filters", - "looks_like_expand", - # Hints - "parse_repo_hints", - "clean_query_and_dsl", - "select_best_search_tool_by_signature", - # Batching - "BatchingContextAnswerClient", - "get_batch_client", - # Validation - "is_result_good", - "extract_metric_from_resp", - "material_drop", - # Planning - "build_plan", - # CLI - "main", -] - -if __name__ == "__main__": - import sys - raise SystemExit(main(sys.argv[1:])) diff --git a/scripts/mcp_router/__init__.py b/scripts/mcp_router/__init__.py deleted file mode 100644 index 0bf56bac..00000000 --- a/scripts/mcp_router/__init__.py +++ /dev/null @@ -1,206 +0,0 @@ -""" -mcp_router - Modular MCP routing package. - -This package provides intent classification, tool planning, and HTTP execution -for routing queries to the appropriate MCP tools. - -Public API: -- classify_intent: Determine query intent -- build_plan: Create execution plan for query -- call_tool_http: Execute MCP tool over HTTP -- discover_tool_endpoints: Find available tools -""" -from __future__ import annotations - -# Config exports -from .config import ( - HTTP_URL_INDEXER, - HTTP_URL_MEMORY, - DEFAULT_HTTP_URL, - HEALTH_PORT_INDEXER, - HEALTH_PORT_MEMORY, - LANGS, - cache_ttl_sec, - scratchpad_ttl_sec, - divergence_thresholds, - divergence_is_fatal_for, -) - -# Intent exports -from .intent import ( - INTENT_ANSWER, - INTENT_SEARCH, - INTENT_SEARCH_TESTS, - INTENT_SEARCH_CONFIG, - INTENT_SEARCH_CALLERS, - INTENT_SEARCH_IMPORTERS, - INTENT_MEMORY_STORE, - INTENT_MEMORY_FIND, - INTENT_INDEX, - INTENT_PRUNE, - INTENT_STATUS, - INTENT_LIST, - classify_intent, - get_last_intent_debug, -) - -# Memory exports -from .memory import parse_memory_store_payload - -# Client exports -from .client import ( - call_tool_http, - is_failure_response, - discover_tool_endpoints, - default_tool_endpoints, - tools_describe_cached, - _mcp_handshake, - _post_raw, - _post_raw_retry, - _parse_stream_or_json, - _filter_args, -) - -# Scratchpad exports -from .scratchpad import ( - scratchpad_path, - load_scratchpad, - save_scratchpad, - looks_like_repeat, - looks_like_same_filters, - looks_like_expand, -) - -# Hints exports -from .hints import ( - parse_repo_hints, - clean_query_and_dsl, - select_best_search_tool_by_signature, -) - -# Batching exports -from .batching import ( - BatchingContextAnswerClient, - get_batch_client, -) - -# Validation exports -from .validation import ( - is_result_good, - extract_metric_from_resp, - material_drop, -) - -# Planning exports -from .planning import build_plan, route_query - -# --------------------------------------------------------------------------- -# Private function imports for backward compatibility -# --------------------------------------------------------------------------- -from .intent import _classify_intent_rules - -# --------------------------------------------------------------------------- -# Legacy aliases (underscore-prefixed) for backward compatibility -# --------------------------------------------------------------------------- -_LAST_INTENT_DEBUG = {} # Use get_last_intent_debug() instead -_BATCH_CLIENT = None # Lazy initialized - -def _get_batch_client(): - global _BATCH_CLIENT - if _BATCH_CLIENT is None: - _BATCH_CLIENT = get_batch_client() - return _BATCH_CLIENT - -# Function aliases -_parse_memory_store_payload = parse_memory_store_payload -_looks_like_repeat = looks_like_repeat -_looks_like_same_filters = looks_like_same_filters -_looks_like_expand = looks_like_expand -_load_scratchpad = load_scratchpad -_save_scratchpad = save_scratchpad -_scratchpad_path = scratchpad_path -_scratchpad_ttl_sec = scratchpad_ttl_sec -_cache_ttl_sec = cache_ttl_sec -_discover_tool_endpoints = discover_tool_endpoints -_default_tool_endpoints = default_tool_endpoints -_tools_describe_cached = tools_describe_cached -_is_failure_response = is_failure_response -_is_result_good = is_result_good -_extract_metric_from_resp = extract_metric_from_resp -_material_drop = material_drop -_divergence_thresholds = divergence_thresholds -_divergence_is_fatal_for = divergence_is_fatal_for -_parse_repo_hints = parse_repo_hints -_clean_query_and_dsl = clean_query_and_dsl -_select_best_search_tool_by_signature = select_best_search_tool_by_signature - -# Health port aliases -_HEALTH_PORT_INDEXER = HEALTH_PORT_INDEXER -_HEALTH_PORT_MEMORY = HEALTH_PORT_MEMORY - -# Language set alias -_LANGS = LANGS - - -__all__ = [ - # Config - "HTTP_URL_INDEXER", - "HTTP_URL_MEMORY", - "DEFAULT_HTTP_URL", - "HEALTH_PORT_INDEXER", - "HEALTH_PORT_MEMORY", - "LANGS", - "cache_ttl_sec", - "scratchpad_ttl_sec", - "divergence_thresholds", - "divergence_is_fatal_for", - # Intent - "INTENT_ANSWER", - "INTENT_SEARCH", - "INTENT_SEARCH_TESTS", - "INTENT_SEARCH_CONFIG", - "INTENT_SEARCH_CALLERS", - "INTENT_SEARCH_IMPORTERS", - "INTENT_MEMORY_STORE", - "INTENT_MEMORY_FIND", - "INTENT_INDEX", - "INTENT_PRUNE", - "INTENT_STATUS", - "INTENT_LIST", - "classify_intent", - "get_last_intent_debug", - # Memory - "parse_memory_store_payload", - # Client - "call_tool_http", - "is_failure_response", - "discover_tool_endpoints", - "default_tool_endpoints", - "tools_describe_cached", - # Scratchpad - "scratchpad_path", - "load_scratchpad", - "save_scratchpad", - "looks_like_repeat", - "looks_like_same_filters", - "looks_like_expand", - # Hints - "parse_repo_hints", - "clean_query_and_dsl", - "select_best_search_tool_by_signature", - # Batching - "BatchingContextAnswerClient", - "get_batch_client", - # Validation - "is_result_good", - "extract_metric_from_resp", - "material_drop", - # Planning - "build_plan", - "route_query", - # CLI - "main", -] - -# Import main for CLI compatibility -from .cli import main diff --git a/scripts/mcp_router/__main__.py b/scripts/mcp_router/__main__.py deleted file mode 100644 index f7c2d058..00000000 --- a/scripts/mcp_router/__main__.py +++ /dev/null @@ -1,8 +0,0 @@ -#!/usr/bin/env python3 -""" -Allow running as: python -m scripts.mcp_router "query" -""" -from .cli import main - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/scripts/mcp_router/batching.py b/scripts/mcp_router/batching.py deleted file mode 100644 index 7300a87e..00000000 --- a/scripts/mcp_router/batching.py +++ /dev/null @@ -1,331 +0,0 @@ -""" -mcp_router/batching.py - Context answer batching client. -""" -from __future__ import annotations - -import json -import os -import re -import sys -import threading -import time -from typing import Any, Dict - -from .config import HTTP_URL_INDEXER -from .client import call_tool_http - - -class BatchingContextAnswerClient: - """Lightweight in-memory batching for context_answer calls. - - - Queues short-lived requests keyed by (base_url, collection, filters_fingerprint) - - Flushes after a small window or when batch size cap is hit - - For multi-item batches, sends query=[...] with mode="pack" - - Shares the same response with all enqueued requests - """ - - def __init__(self, call_func=None, enable: bool | None = None, window_ms: int | None = None, - max_batch: int | None = None, budget_ms: int | None = None): - self._call = call_func or call_tool_http - if enable is None: - env_enabled = os.environ.get("ROUTER_BATCH_ENABLED") - if env_enabled is None: - env_enabled = os.environ.get("ROUTER_BATCH_ENABLE", "0") - self.enabled = str(env_enabled).strip().lower() in {"1", "true", "yes", "on"} - else: - self.enabled = bool(enable) - self.window_ms = int(os.environ.get("ROUTER_BATCH_WINDOW_MS", str(window_ms if window_ms is not None else 100)) or 100) - env_max = os.environ.get("ROUTER_BATCH_MAX_SIZE") - if env_max is None: - env_max = os.environ.get("ROUTER_BATCH_MAX") - self.max_batch = int(env_max or (max_batch if max_batch is not None else 8)) - env_budget = os.environ.get("ROUTER_BATCH_LATENCY_BUDGET_MS") - if env_budget is None: - env_budget = os.environ.get("ROUTER_BATCH_BUDGET_MS") - self.budget_ms = int(env_budget or (budget_ms if budget_ms is not None else 2000)) - self._lock = threading.RLock() - self._groups: dict[str, dict[str, Any]] = {} - - def _should_bypass(self, args: Dict[str, Any]) -> bool: - try: - if isinstance(args, dict): - v = args.get("immediate") - if v is not None and str(v).strip().lower() in {"1", "true", "yes", "on"}: - return True - except Exception: - pass - if str(os.environ.get("ROUTER_BATCH_BYPASS", "0")).strip().lower() in {"1", "true", "yes", "on"}: - return True - try: - q = str((args or {}).get("query") or "") - if "immediate answer" in q.lower(): - return True - except Exception: - pass - return False - - def _norm_query(self, q: str) -> str: - try: - return re.sub(r"\s+", " ", str(q or "").strip()) - except Exception: - return str(q) - - def _filters_fingerprint(self, args: Dict[str, Any]) -> str: - keep = { - "collection", "language", "under", "kind", "symbol", "ext", - "path_regex", "path_glob", "not_glob", "not_", "case", - "limit", "per_path", "include_snippet", - } - try: - filt = {k: args.get(k) for k in keep if k in args} - def _norm(v): - if v is None: - return None - if isinstance(v, (list, tuple)): - return [str(x) for x in v] - return v - clean = {k: _norm(v) for k, v in filt.items()} - return json.dumps(clean, sort_keys=True, ensure_ascii=False) - except Exception: - return "{}" - - def _group_key(self, base_url: str, args: Dict[str, Any]) -> str: - coll = str(args.get("collection") or "") - fp = self._filters_fingerprint(args) - repo = os.getcwd() - return f"{base_url}|{coll}|answer|{fp}|{repo}" - - def call_or_enqueue(self, base_url: str, tool: str, args: Dict[str, Any], timeout: float = 120.0) -> Dict[str, Any]: - if not self.enabled: - return self._call(base_url, tool, args, timeout=timeout) - if self._should_bypass(args): - return self._call(base_url, tool, args, timeout=timeout) - - start_ts = time.time() - key = self._group_key(base_url, args or {}) - norm_q = self._norm_query((args or {}).get("query") or "") - ev = threading.Event() - slot = {"event": ev, "result": None, "error": None, "query": norm_q, "args": dict(args or {})} - - with self._lock: - g = self._groups.get(key) - if not g: - g = { - "created": time.time(), - "items": [], - "timer": None, - } - self._groups[key] = g - g["items"].append(slot) - if g["timer"] is None: - delay = max(0.0, float(self.window_ms) / 1000.0) - t = threading.Timer(delay, self._flush, args=(key,)) - g["timer"] = t - t.daemon = True - t.start() - if len(g["items"]) >= self.max_batch: - t = g.get("timer") - if t: - try: - t.cancel() - except Exception: - pass - g["timer"] = None - threading.Thread(target=self._flush, args=(key,), daemon=True).start() - - remain = max(0.05, self.budget_ms / 1000.0) - ev.wait(timeout=min(timeout, remain)) - if not ev.is_set(): - try: - res = self._call(base_url, tool, args, timeout=timeout) - slot["result"] = res - ev.set() - try: - with self._lock: - gg = self._groups.get(key) - if gg: - lst = gg.get("items") or [] - if slot in lst: - try: - lst.remove(slot) - except Exception: - pass - if not lst: - t2 = gg.get("timer") - if t2: - try: - t2.cancel() - except Exception: - pass - self._groups.pop(key, None) - except Exception: - pass - try: - print(json.dumps({"router": {"batch_fallback": True, "elapsed_ms": int((time.time()-start_ts)*1000)}}), file=sys.stderr) - except Exception: - pass - return res - except Exception as e: - slot["error"] = e - ev.set() - raise - - if slot.get("error") is not None: - raise slot["error"] - return slot.get("result") or {} - - def _flush(self, key: str) -> None: - with self._lock: - g = self._groups.get(key) - if not g: - return - items = g.get("items") or [] - g["items"] = [] - g["timer"] = None - if not items: - self._groups.pop(key, None) - return - - unique_q: list[str] = [] - seen_q = set() - for it in items: - q = it.get("query") or "" - if q not in seen_q: - seen_q.add(q) - unique_q.append(q) - first_args = dict(items[0].get("args") or {}) - forward = {k: v for k, v in first_args.items() if k not in {"query", "queries"}} - base_url = None - try: - base_url = key.split("|")[0] - except Exception: - base_url = HTTP_URL_INDEXER - - started = time.time() - results_by_q: Dict[str, Any] = {} - errors_by_q: Dict[str, Exception] = {} - calls = 0 - try: - import copy as _copy - except Exception: - _copy = None - - if len(unique_q) > 1: - args_all = dict(forward) - args_all["query"] = list(unique_q) - args_all["mode"] = args_all.get("mode") or "pack" - try: - agg_res = self._call(base_url, "context_answer", args_all, timeout=120.0) - calls = 1 - try: - payload = ((agg_res or {}).get("result") or {}).get("structuredContent") or {} - body = (payload.get("result") or {}) - except Exception: - payload, body = {}, {} - - abq = None - try: - abq = body.get("answers_by_query") - except Exception: - abq = None - if isinstance(abq, list) and abq: - _map: Dict[str, Any] = {} - by_idx = (len(abq) >= len(unique_q)) - for i, entry in enumerate(abq): - try: - qv = entry.get("query") - qk = None - if isinstance(qv, list) and qv: - qk = str(qv[0]) - elif isinstance(qv, str): - qk = qv - except Exception: - qk = None - entry_key = qk if qk else (unique_q[i] if by_idx and i < len(unique_q) else None) - if not entry_key: - continue - per = _copy.deepcopy(agg_res) if _copy else json.loads(json.dumps(agg_res)) - try: - per_body = (per.get("result") or {}).get("structuredContent", {}).get("result", {}) - except Exception: - per_body = None - try: - ans_i = str(entry.get("answer") or "") - cits_i = entry.get("citations") or [] - if per_body is not None: - per_body["answer"] = ans_i - per_body["citations"] = cits_i - per_body["query"] = [entry_key] - except Exception: - pass - _map[str(entry_key)] = per - for uq in unique_q: - if str(uq) in _map: - results_by_q[uq] = _map[str(uq)] - remaining = [uq for uq in unique_q if uq not in results_by_q] - else: - remaining = list(unique_q) - - if remaining: - for uq in remaining: - args_i = dict(forward) - args_i["query"] = uq - try: - results_by_q[uq] = self._call(base_url, "context_answer", args_i, timeout=120.0) - except Exception as e: - errors_by_q[uq] = e - calls += len(remaining) - except Exception as e: - for uq in unique_q: - errors_by_q[uq] = e - calls = 1 - else: - args1 = dict(forward) - args1["query"] = unique_q[0] if unique_q else "" - try: - results_by_q[args1["query"]] = self._call(base_url, "context_answer", args1, timeout=120.0) - except Exception as e: - errors_by_q[args1["query"]] = e - calls = 1 - - elapsed_ms = int((time.time() - started) * 1000) - try: - print(json.dumps({ - "router": { - "batch_flushed": True, - "n_items": len(items), - "unique_q": len(unique_q), - "calls": int(calls), - "elapsed_ms": elapsed_ms, - "ok": (len(errors_by_q) == 0), - } - }), file=sys.stderr) - except Exception: - pass - - for it in items: - q = it.get("query") or "" - it["result"] = results_by_q.get(q) - it["error"] = errors_by_q.get(q) - ev = it.get("event") - try: - if hasattr(ev, "set"): - ev.set() - except Exception: - pass - with self._lock: - gg = self._groups.get(key) - if gg and not gg.get("items"): - self._groups.pop(key, None) - - -# Global client singleton -_BATCH_CLIENT: BatchingContextAnswerClient | None = None - - -def get_batch_client() -> BatchingContextAnswerClient: - """Get or create global batch client.""" - global _BATCH_CLIENT - if _BATCH_CLIENT is None: - _BATCH_CLIENT = BatchingContextAnswerClient() - return _BATCH_CLIENT diff --git a/scripts/mcp_router/cli.py b/scripts/mcp_router/cli.py deleted file mode 100644 index 50d63d55..00000000 --- a/scripts/mcp_router/cli.py +++ /dev/null @@ -1,305 +0,0 @@ -#!/usr/bin/env python3 -""" -mcp_router/cli.py - CLI entrypoint for MCP router. - -Usage: - python -m scripts.mcp_router --plan "How do I ...?" - python -m scripts.mcp_router --run "What is hybrid search?" -""" -from __future__ import annotations - -import argparse -import json -import re -import sys -import time -from typing import Any, Dict, List - -from .config import HTTP_URL_INDEXER, scratchpad_ttl_sec, divergence_thresholds -from .client import call_tool_http, is_failure_response, discover_tool_endpoints -from .scratchpad import ( - load_scratchpad, - save_scratchpad, - looks_like_repeat, - looks_like_expand, -) -from .validation import ( - is_result_good, - extract_metric_from_resp, - material_drop, -) -from .batching import get_batch_client -from .planning import build_plan -from .config import divergence_is_fatal_for - - -def main(argv: List[str] | None = None) -> int: - """Main CLI entrypoint.""" - if argv is None: - argv = sys.argv[1:] - - ap = argparse.ArgumentParser() - ap.add_argument("query", help="User query to route") - ap.add_argument("--plan", action="store_true", help="Only print routing plan (no execution)") - ap.add_argument("--run", action="store_true", help="Execute the routed tool(s) over HTTP") - ap.add_argument("--timeout", type=float, default=180.0, help="HTTP timeout for tool calls") - args = ap.parse_args(argv) - - plan = build_plan(args.query) - print(json.dumps({"router": {"url": HTTP_URL_INDEXER, "plan": plan}}, indent=2)) - - if args.plan and not args.run: - return 0 - - # Load scratchpad for prior context - sp = {} - fresh = False - prior_answer = None - prior_citations = None - prior_paths = None - try: - sp = load_scratchpad() - ts = float(sp.get("timestamp") or 0.0) - fresh = bool(ts and (time.time() - ts) <= scratchpad_ttl_sec()) - if fresh: - prior_answer = sp.get("last_answer") - prior_citations = sp.get("last_citations") - prior_paths = sp.get("last_paths") - except Exception: - pass - - # Execute sequentially until one succeeds - last_err = None - last = None - tool_servers = discover_tool_endpoints() - mem_snippets: list[str] = list(sp.get("mem_snippets") or []) if fresh else [] - batch_client = get_batch_client() - - for idx, (tool, targs) in enumerate(plan): - base_url = tool_servers.get(tool, HTTP_URL_INDEXER) - - # Skip memory.find if we already have fresh snippets and this is a repeat/expand - if (tool.lower().endswith("find") or tool.lower() in {"find", "memory.find"}) and mem_snippets and fresh and (looks_like_repeat(args.query) or looks_like_expand(args.query)): - try: - print(json.dumps({"tool": tool, "skipped": "scratchpad_fresh"})) - except Exception: - pass - continue - - # Augment answer queries with context - if tool in {"context_answer", "context_answer_compat"} and (mem_snippets or (fresh and (prior_answer or prior_citations or prior_paths))): - try: - tq = str((targs or {}).get("query") or args.query) - sections = [tq] - if mem_snippets: - bullets = [] - for s in mem_snippets[:3]: - ss = re.sub(r"\s+", " ", str(s)).strip() - if len(ss) > 200: - ss = ss[:197] + "..." - bullets.append(f"- {ss}") - sections.append("Memory context:\n" + "\n".join(bullets)) - if fresh and (looks_like_expand(args.query) or looks_like_repeat(args.query)): - if isinstance(prior_answer, str) and prior_answer.strip(): - pa = re.sub(r"\s+", " ", prior_answer).strip() - if len(pa) > 400: - pa = pa[:397] + "..." - sections.append("Prior summary:\n" + pa) - paths_list = [] - if isinstance(prior_paths, list) and prior_paths: - paths_list = [str(p) for p in prior_paths[:5]] - elif isinstance(prior_citations, list) and prior_citations: - uniq = [] - for c in prior_citations: - if isinstance(c, dict) and c.get("path") and c["path"] not in uniq: - uniq.append(c["path"]) - paths_list = uniq[:5] - if paths_list: - sections.append("Citations context:\n" + "\n".join(f"- {p}" for p in paths_list)) - aug = "\n\n".join(sections) - targs = {**(targs or {}), "query": aug} - except Exception: - pass - - try: - if tool in {"context_answer", "context_answer_compat"}: - res = batch_client.call_or_enqueue(base_url, tool, targs, timeout=args.timeout) - else: - res = call_tool_http(base_url, tool, targs, timeout=args.timeout) - print(json.dumps({"tool": tool, "result": res}, indent=2)) - last = res - - # Capture memory snippets - try: - if tool.lower().endswith("find") or tool.lower() in {"find", "memory.find"}: - r = res.get("result") or {} - items = [] - sc = r.get("structuredContent") - if isinstance(sc, dict): - rs0 = sc.get("result") or sc - if isinstance(rs0, dict): - items = rs0.get("results") or rs0.get("hits") or [] - if not items: - content = r.get("content") - if isinstance(content, list): - for c in content: - if not isinstance(c, dict): - continue - if "json" in c: - j = c.get("json") - if isinstance(j, (dict, list)): - container = j.get("result") if isinstance(j, dict) and "result" in j else j - if isinstance(container, dict): - items = container.get("results") or container.get("hits") or [] - if items: - break - if c.get("type") == "text": - ttxt = c.get("text") - if isinstance(ttxt, str) and ttxt.strip(): - try: - j = json.loads(ttxt) - except Exception: - continue - container = j.get("result") if isinstance(j, dict) and "result" in j else j - if isinstance(container, dict): - items = container.get("results") or container.get("hits") or [] - if items: - break - for it in items: - if isinstance(it, dict): - txt = it.get("information") or it.get("content") or it.get("text") - if isinstance(txt, str) and txt.strip(): - mem_snippets.append(txt.strip()) - except Exception: - pass - - # Determine if we should treat this step as terminal - has_future_answer = any(tn in {"context_answer", "context_answer_compat"} for (tn, _) in plan[idx + 1:]) - if (not is_failure_response(res)) and is_result_good(tool, res): - if tool.lower() in {"find", "memory.find"} and has_future_answer: - continue - - # Persist scratchpad - try: - last_filters: Dict[str, Any] = {} - for (tn, ta) in plan: - if tn == "repo_search" or tn.startswith("search_"): - if isinstance(ta, dict): - for k in ("language", "under", "symbol", "ext", "path_glob", "not_glob"): - if ta.get(k) not in (None, ""): - last_filters[k] = ta.get(k) - break - - last_answer_text = None - last_citations_list = None - last_paths_list: list[str] | None = None - if tool in {"context_answer", "context_answer_compat"}: - try: - r0 = res.get("result") or {} - sc0 = r0.get("structuredContent") or {} - rs0 = sc0.get("result") or sc0 - if isinstance(rs0, dict): - ans0 = rs0.get("answer") - if isinstance(ans0, str): - last_answer_text = ans0 - cites0 = rs0.get("citations") - if isinstance(cites0, list): - last_citations_list = cites0 - uniqp: list[str] = [] - for c in cites0: - if isinstance(c, dict) and c.get("path") and c["path"] not in uniqp: - uniqp.append(c["path"]) - last_paths_list = uniqp - except Exception: - pass - - # Divergence detection - divergence_should_abort = False - last_metrics_prev = {} - try: - last_metrics_prev = sp.get("last_metrics") or {} - if not isinstance(last_metrics_prev, dict): - last_metrics_prev = {} - except Exception: - last_metrics_prev = {} - metric = extract_metric_from_resp(tool, res) - last_metrics_map = dict(last_metrics_prev) - if metric is not None: - mname, mval = metric - prev_val = None - try: - prev_val = last_metrics_prev.get(tool, {}).get(mname) - if prev_val is not None: - prev_val = float(prev_val) - except Exception: - prev_val = None - drop_frac, min_base = divergence_thresholds() - if material_drop(prev_val, float(mval), drop_frac, min_base): - fatal = divergence_is_fatal_for(tool) - try: - print(json.dumps({ - "divergence": { - "tool": tool, - "metric": mname, - "previous": prev_val, - "current": float(mval), - "drop_frac": drop_frac, - "fatal": fatal, - } - })) - except Exception: - pass - if fatal: - divergence_should_abort = True - try: - last_metrics_map.setdefault(tool, {})[mname] = float(mval) - except Exception: - pass - else: - last_metrics_map = last_metrics_prev - - success_criteria = { - "context_answer": {"expected_fields": ["answer"], "min_citations": 0}, - "context_answer_compat": {"expected_fields": ["answer"], "min_citations": 0}, - "repo_search": {"min_results": 1}, - "search_config_for": {"min_results": 1}, - "search_tests_for": {"min_results": 1}, - "search_callers_for": {"min_results": 1}, - "search_importers_for": {"min_results": 1}, - "find": {"min_results": 1}, - } - sp = { - "last_query": args.query, - "last_plan": plan, - "last_filters": last_filters or None, - "mem_snippets": mem_snippets[:5], - "last_answer": last_answer_text, - "last_citations": last_citations_list, - "last_paths": last_paths_list, - "success_criteria": success_criteria, - "last_metrics": last_metrics_map, - "timestamp": time.time(), - } - save_scratchpad(sp) - except Exception: - pass - - if divergence_should_abort: - continue - - return 0 - except Exception as e: - last_err = e - try: - print(json.dumps({"tool": tool, "server": base_url, "error": str(e)}), file=sys.stderr) - except Exception: - pass - continue - - if last_err: - print(f"Router: all attempts failed: {last_err}", file=sys.stderr) - return 1 if last is not None else 2 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/scripts/mcp_router/client.py b/scripts/mcp_router/client.py deleted file mode 100644 index 88503ef6..00000000 --- a/scripts/mcp_router/client.py +++ /dev/null @@ -1,318 +0,0 @@ -""" -mcp_router/client.py - HTTP/MCP client helpers. -""" -from __future__ import annotations - -import json -import os -import time -from typing import Any, Dict, List, Tuple -from urllib import request - -from .config import ( - HTTP_URL_INDEXER, - HTTP_URL_MEMORY, - HEALTH_PORT_INDEXER, - HEALTH_PORT_MEMORY, - cache_ttl_sec, -) - -# Caches -_TOOL_ENDPOINTS_CACHE_MAP: Dict[str, str] = {} -_TOOL_ENDPOINTS_CACHE_TS: float = 0.0 -_TOOLS_DESCR_CACHE: Dict[str, list] = {} -_TOOLS_DESCR_TS: Dict[str, float] = {} - - -def _post_raw(url: str, payload: Dict[str, Any], headers: Dict[str, str], timeout: float = 60.0) -> Tuple[Dict[str, str], bytes]: - req = request.Request(url, method="POST") - for k, v in headers.items(): - req.add_header(k, v) - data = json.dumps(payload).encode("utf-8") - with request.urlopen(req, data=data, timeout=timeout) as resp: - body = resp.read() - hdrs = {k.lower(): v for k, v in resp.headers.items()} - return hdrs, body - - -def _post_raw_retry(url: str, payload: Dict[str, Any], headers: Dict[str, str], - timeout: float = 60.0, retries: int = 2, backoff: float = 0.5) -> Tuple[Dict[str, str], bytes]: - last_exc: Exception | None = None - for i in range(max(0, retries) + 1): - try: - return _post_raw(url, payload, headers, timeout=timeout) - except Exception as e: - last_exc = e - if i < retries: - try: - time.sleep(backoff * (2 ** i)) - except Exception: - pass - else: - raise last_exc - - -def _parse_stream_or_json(body: bytes) -> Dict[str, Any]: - txt = body.decode("utf-8", errors="ignore") - if "data:" in txt and ("event:" in txt or txt.strip().startswith("data:")): - last = None - for line in txt.splitlines(): - if line.startswith("data:"): - last = line[len("data:"):].strip() - if last: - try: - return json.loads(last) - except Exception: - pass - return json.loads(txt) - - -def _filter_args(d: Dict[str, Any]) -> Dict[str, Any]: - """Remove None/empty values from args dict.""" - return {k: v for k, v in d.items() if v not in (None, "")} - - -def _mcp_handshake(base_url: str, timeout: float = 30.0) -> Dict[str, str]: - """Perform MCP handshake and return headers with session ID.""" - headers = { - "Content-Type": "application/json", - "Accept": "application/json, text/event-stream", - } - init_payload = { - "jsonrpc": "2.0", - "method": "initialize", - "params": { - "protocolVersion": "2024-11-05", - "capabilities": {}, - "clientInfo": {"name": "router", "version": "0.1.0"}, - }, - "id": 1, - } - hdrs, body = _post_raw_retry(base_url, init_payload, headers, timeout=timeout) - sid = hdrs.get("mcp-session-id") or hdrs.get("Mcp-Session-Id") - if not sid: - try: - j = _parse_stream_or_json(body) - sid = j.get("sessionId") - except Exception: - sid = None - if sid: - headers["Mcp-Session-Id"] = sid - try: - _post_raw_retry(base_url, {"jsonrpc": "2.0", "method": "notifications/initialized"}, headers, timeout=timeout) - except Exception: - pass - return headers - - -def _extract_iserror_text(resp: Dict[str, Any]) -> str | None: - try: - r = resp.get("result") or {} - if isinstance(r, dict) and r.get("isError"): - content = r.get("content") - if isinstance(content, list) and content and isinstance(content[0], dict): - if content[0].get("type") == "text": - return content[0].get("text") - except Exception: - pass - return None - - -def call_tool_http(base_url: str, tool_name: str, args: Dict[str, Any], timeout: float = 120.0) -> Dict[str, Any]: - """Call an MCP tool over HTTP.""" - headers = _mcp_handshake(base_url, timeout=min(timeout, 30.0)) - - def _do_call(arguments: Dict[str, Any]) -> Dict[str, Any]: - payload = { - "jsonrpc": "2.0", - "id": "router-1", - "method": "tools/call", - "params": { - "name": tool_name, - "arguments": arguments, - }, - } - _, body = _post_raw_retry(base_url, payload, headers, timeout=timeout) - return _parse_stream_or_json(body) - - args1 = _filter_args(args) - if tool_name.endswith("_compat"): - resp = _do_call({"arguments": args1}) - else: - resp = _do_call(args1) - - def _get_structured_error(r: Dict[str, Any]) -> str | None: - try: - rr = r.get("result") or {} - sc = rr.get("structuredContent") or {} - rs = sc.get("result") or {} - err = rs.get("error") - if isinstance(err, str): - return err - except Exception: - pass - return None - - msg = _extract_iserror_text(resp) - serr = _get_structured_error(resp) - if msg: - low = msg.lower() - if ("kwargs" in low) and ("field required" in low or "missing" in low): - return _do_call({"kwargs": args1}) - if ("arguments" in low) and ("field required" in low or "missing" in low): - return _do_call({"arguments": args1}) - if (serr and serr.strip().lower() == "query required") and ("query" in args1 or "queries" in args1): - resp4 = _do_call({"kwargs": args1}) - serr2 = _get_structured_error(resp4) - if not (serr2 and serr2.strip().lower() == "query required"): - return resp4 - resp5 = _do_call({"arguments": {"kwargs": args1}}) - serr3 = _get_structured_error(resp5) - if not (serr3 and serr3.strip().lower() == "query required"): - return resp5 - return _do_call({"arguments": args1}) - return resp - - -def is_failure_response(resp: Dict[str, Any]) -> bool: - """Check if response indicates a failure.""" - try: - r = resp.get("result") or {} - if r.get("isError") is True: - return True - sc = r.get("structuredContent") or {} - rs = sc.get("result") or {} - if isinstance(rs, dict) and isinstance(rs.get("error"), str): - return True - except Exception: - return False - return False - - -def _tools_describe_from_health(base_url: str, timeout: float = 3.0) -> list[dict]: - """Fetch tool descriptors from health /tools endpoint.""" - try: - import urllib.request - if base_url == HTTP_URL_INDEXER: - url = f"http://localhost:{HEALTH_PORT_INDEXER}/tools" - elif base_url == HTTP_URL_MEMORY: - url = f"http://localhost:{HEALTH_PORT_MEMORY}/tools" - else: - return [] - with urllib.request.urlopen(url, timeout=timeout) as r: - if getattr(r, "status", 200) != 200: - return [] - body = r.read() - j = _parse_stream_or_json(body) - tools = (j.get("tools") if isinstance(j, dict) else None) or [] - out = [] - for t in tools: - if not isinstance(t, dict): - continue - nm = t.get("name") - if not nm: - continue - out.append({"name": nm, "description": (t.get("description") or "").strip()}) - return out - except Exception: - return [] - - -def _mcp_tools_list(base_url: str, timeout: float = 30.0) -> List[str]: - """Get list of tool names from MCP server.""" - try: - headers = _mcp_handshake(base_url, timeout=min(timeout, 15.0)) - payload = {"jsonrpc": "2.0", "id": "router-list", "method": "tools/list"} - _, body = _post_raw_retry(base_url, payload, headers, timeout=timeout) - j = _parse_stream_or_json(body) - tools = ((j.get("result") or {}).get("tools") or []) - names: List[str] = [] - for t in tools: - try: - n = t.get("name") if isinstance(t, dict) else None - if isinstance(n, str) and n: - names.append(n) - except Exception: - continue - return names - except Exception: - return [] - - -def _mcp_tools_describe(base_url: str, timeout: float = 20.0) -> list[dict]: - """Return tool dicts from tools/list.""" - try: - headers = _mcp_handshake(base_url, timeout=min(timeout, 10.0)) - payload = {"jsonrpc": "2.0", "id": "router-list2", "method": "tools/list"} - _, body = _post_raw_retry(base_url, payload, headers, timeout=timeout) - j = _parse_stream_or_json(body) - tools = ((j.get("result") or {}).get("tools") or []) - out = [] - for t in tools: - if not isinstance(t, dict): - continue - name = (t.get("name") or "").strip() - if not name: - continue - out.append(t) - return out - except Exception: - return [] - - -def tools_describe_cached(base_url: str, allow_network: bool = True, timeout: float = 20.0) -> list[dict]: - """Get tool descriptions with caching.""" - now = time.time() - ts = _TOOLS_DESCR_TS.get(base_url, 0.0) - if base_url in _TOOLS_DESCR_CACHE and (now - ts) <= cache_ttl_sec(): - return _TOOLS_DESCR_CACHE[base_url] - if not allow_network: - return _TOOLS_DESCR_CACHE.get(base_url, []) - desc = _tools_describe_from_health(base_url, timeout=min(timeout, 3.0)) or _mcp_tools_describe(base_url, timeout=timeout) - _TOOLS_DESCR_CACHE[base_url] = desc - _TOOLS_DESCR_TS[base_url] = now - return desc - - -def default_tool_endpoints() -> Dict[str, str]: - """Return default tool -> endpoint mapping.""" - idx = HTTP_URL_INDEXER - mem = HTTP_URL_MEMORY - mapping: Dict[str, str] = {} - for n in [ - "repo_search", "context_answer", "context_answer_compat", "expand_query", - "search_tests_for", "search_config_for", "search_callers_for", "search_importers_for", - "qdrant_index_root", "qdrant_prune", "qdrant_status", "qdrant_list", - "workspace_info", "list_workspaces", "change_history_for_path", "code_search", "context_search", - ]: - mapping[n] = idx - mapping["store"] = mem - mapping["find"] = mem - return mapping - - -def discover_tool_endpoints(force: bool = False, allow_network: bool = True) -> Dict[str, str]: - """Discover tool -> endpoint mapping from servers.""" - global _TOOL_ENDPOINTS_CACHE_TS, _TOOL_ENDPOINTS_CACHE_MAP - now = time.time() - ttl = cache_ttl_sec() - if not force and _TOOL_ENDPOINTS_CACHE_MAP and (now - _TOOL_ENDPOINTS_CACHE_TS) <= ttl: - return _TOOL_ENDPOINTS_CACHE_MAP - if not allow_network: - return _TOOL_ENDPOINTS_CACHE_MAP or default_tool_endpoints() - mapping: Dict[str, str] = {} - idx_desc = tools_describe_cached(HTTP_URL_INDEXER, allow_network=allow_network) - for t in idx_desc: - n = t.get("name") if isinstance(t, dict) else None - if n: - mapping[n] = HTTP_URL_INDEXER - mem_desc = tools_describe_cached(HTTP_URL_MEMORY, allow_network=allow_network) - for t in mem_desc: - n = t.get("name") if isinstance(t, dict) else None - if n and n not in mapping: - mapping[n] = HTTP_URL_MEMORY - if mapping: - _TOOL_ENDPOINTS_CACHE_MAP.clear() - _TOOL_ENDPOINTS_CACHE_MAP.update(mapping) - _TOOL_ENDPOINTS_CACHE_TS = now - return mapping or (_TOOL_ENDPOINTS_CACHE_MAP or default_tool_endpoints()) diff --git a/scripts/mcp_router/config.py b/scripts/mcp_router/config.py deleted file mode 100644 index b67b1282..00000000 --- a/scripts/mcp_router/config.py +++ /dev/null @@ -1,69 +0,0 @@ -""" -mcp_router/config.py - Shared configuration and constants. -""" -from __future__ import annotations - -import os - -# HTTP endpoints -HTTP_URL_INDEXER = os.environ.get("MCP_INDEXER_HTTP_URL", "http://localhost:8003/mcp").rstrip("/") -HTTP_URL_MEMORY = os.environ.get("MCP_MEMORY_HTTP_URL", "http://localhost:8002/mcp").rstrip("/") -DEFAULT_HTTP_URL = HTTP_URL_INDEXER - -# Health ports -try: - HEALTH_PORT_INDEXER = int(os.environ.get("FASTMCP_INDEXER_HTTP_HEALTH_PORT", "18003") or 18003) -except (ValueError, TypeError): - HEALTH_PORT_INDEXER = 18003 - -try: - HEALTH_PORT_MEMORY = int(os.environ.get("FASTMCP_HTTP_HEALTH_PORT", "18002") or 18002) -except (ValueError, TypeError): - HEALTH_PORT_MEMORY = 18002 - - -def cache_ttl_sec() -> int: - try: - return int(os.environ.get("ROUTER_TOOLS_CACHE_TTL_SEC", "60") or 60) - except Exception: - return 60 - - -def scratchpad_ttl_sec() -> int: - try: - return int(os.environ.get("ROUTER_SCRATCHPAD_TTL_SEC", "300") or 300) - except Exception: - return 300 - - -def divergence_thresholds() -> tuple[float, int]: - try: - drop_frac = float(os.environ.get("ROUTER_DIVERGENCE_DROP_FRAC", "0.5") or 0.5) - except Exception: - drop_frac = 0.5 - try: - min_base = int(os.environ.get("ROUTER_DIVERGENCE_MIN_BASE", "3") or 3) - except Exception: - min_base = 3 - return drop_frac, min_base - - -def divergence_is_fatal_for(tool: str) -> bool: - try: - s = (os.environ.get("ROUTER_DIVERGENCE_FATAL_TOOLS", "") or "").strip() - if not s: - return False - low = s.lower() - if low in {"*", "all", "1", "true"}: - return True - names = {t.strip().lower() for t in s.split(",") if t.strip()} - return tool.strip().lower() in names - except Exception: - return False - - -# Language set for hint parsing -LANGS = { - "python", "typescript", "javascript", "go", "java", "rust", "kotlin", - "c++", "cpp", "csharp", "c#", "ruby", "php", "scala", "swift", "bash", "shell" -} diff --git a/scripts/mcp_router/hints.py b/scripts/mcp_router/hints.py deleted file mode 100644 index 18c46052..00000000 --- a/scripts/mcp_router/hints.py +++ /dev/null @@ -1,137 +0,0 @@ -""" -mcp_router/hints.py - Query hint parsing and tool selection. -""" -from __future__ import annotations - -import re -from typing import Any, Dict, List, Tuple - -from .config import LANGS, HTTP_URL_INDEXER -from .client import tools_describe_cached -from .intent import _cosine, _embed_texts - - -def parse_repo_hints(q: str) -> Dict[str, Any]: - """Extract light filters from the query: language, under, symbol, ext, path_glob, not_glob.""" - s = q.strip() - low = s.lower() - out: Dict[str, Any] = {} - - # language - for lang in sorted(LANGS, key=len, reverse=True): - if re.search(rf"\b{re.escape(lang)}\b", low): - out["language"] = {"javascript": "js", "typescript": "ts", "c++": "cpp", "c#": "csharp"}.get(lang, lang) - break - - # under / in folder - m_under = re.search(r"\bunder\s+([\w./-]+)", low) - m_in = re.search(r"\b(?:in|inside)\s+([\w./-]+)", low) - m = m_under or m_in - if m: - cand = m.group(1) - if len(cand) >= 2 and cand not in LANGS: - out["under"] = cand - - # symbol-like tokens - m2 = re.search(r"([A-Za-z_][A-Za-z0-9_]*\s*\(\))|([A-Za-z_][\w]*\.[A-Za-z_][\w]*)|([A-Za-z_][\w]*::[A-Za-z_][\w]*)", s) - if m2: - sym = m2.group(0) - sym = re.sub(r"\s*\(\)\s*$", "", sym) - out["symbol"] = sym - - # file extension - m3 = re.search(r"\.(py|ts|tsx|js|jsx|go|java|rs|kt|rb|php|scala|swift)$", s) - if m3: - out["ext"] = m3.group(1) - - # glob inclusions - globs: List[str] = [] - if re.search(r"\bonly\b", low): - m_glob = re.search(r"\*\.[A-Za-z0-9]+", s) - if m_glob: - globs.append("**/" + m_glob.group(0)) - if "python" in low and "*.py" not in " ".join(globs): - globs.append("**/*.py") - if globs: - out["path_glob"] = globs - - # exclusions - not_glob: List[str] = [] - for ex in ["vendor", "node_modules", "dist", "build", "tests", "__pycache__"]: - if re.search(rf"\bexclude\s+{re.escape(ex)}\b", low): - not_glob.append(f"**/{ex}/**") - if not_glob: - out["not_glob"] = not_glob - - return out - - -def clean_query_and_dsl(q: str) -> Tuple[str, Dict[str, Any]]: - """Strip DSL tokens from query and return (clean_query, dsl_filters).""" - try: - from scripts.hybrid_search import parse_query_dsl - clean, extracted = parse_query_dsl([q]) - return (clean[0] if clean else ""), (extracted or {}) - except Exception: - return q, {} - - -def _signature_text(t: dict) -> str: - """Build signature text for tool similarity matching.""" - name = (t.get("name") or "").strip() - desc = (t.get("description") or "").strip() - params = [] - try: - schema = t.get("inputSchema") or {} - props = (schema.get("properties") or {}) if isinstance(schema, dict) else {} - params = [k for k in props.keys()] - except Exception: - params = [] - ptxt = (" params:" + ",".join(params)) if params else "" - return (name + "\n" + desc + ptxt).strip() - - -def select_best_search_tool_by_signature(q: str, tool_dict: dict[str, str], allow_network: bool = True) -> str | None: - """Select best matching search tool based on signature similarity.""" - candidates = [n for n in tool_dict.keys() if n == "repo_search" or n.startswith("search_")] - if not candidates: - return None - - per_server: dict[str, list[dict]] = {} - for base in set(tool_dict[t] for t in candidates): - try: - per_server[base] = tools_describe_cached(base, allow_network=allow_network) - except Exception: - per_server[base] = [] - - sig_map: dict[str, str] = {} - for tname in candidates: - base = tool_dict.get(tname) - descs = per_server.get(base, []) - obj = None - for td in descs: - if (td.get("name") or "").strip() == tname: - obj = td - break - sig_map[tname] = _signature_text(obj or {"name": tname, "description": ""}) - - texts = [q] + [sig_map[n] for n in candidates] - vecs = _embed_texts(texts) - if not vecs or len(vecs) < 1 + len(candidates): - return None - - qv = vecs[0] - scores: list[tuple[str, float]] = [] - for i, name in enumerate(candidates): - sv = vecs[1 + i] - scores.append((name, _cosine(qv, sv))) - scores.sort(key=lambda x: x[1], reverse=True) - - best, best_s = scores[0] - repo_s = next((s for n, s in scores if n == "repo_search"), None) - margin = 0.02 - if best == "repo_search" or repo_s is None: - return best - if best != "repo_search" and best_s >= (repo_s + margin): - return best - return "repo_search" diff --git a/scripts/mcp_router/intent.py b/scripts/mcp_router/intent.py deleted file mode 100644 index 7f96543f..00000000 --- a/scripts/mcp_router/intent.py +++ /dev/null @@ -1,246 +0,0 @@ -""" -mcp_router/intent.py - Intent classification (rules + ML). -""" -from __future__ import annotations - -import json -import os -import re -import sys -from typing import Any, Dict, List - -# Intent constants -INTENT_ANSWER = "answer" -INTENT_SEARCH = "search" -INTENT_SEARCH_TESTS = "search_tests" -INTENT_SEARCH_CONFIG = "search_config" -INTENT_SEARCH_CALLERS = "search_callers" -INTENT_SEARCH_IMPORTERS = "search_importers" -INTENT_MEMORY_STORE = "memory_store" -INTENT_MEMORY_FIND = "memory_find" -INTENT_SYMBOL_GRAPH = "symbol_graph" -INTENT_INDEX = "index" -INTENT_PRUNE = "prune" -INTENT_STATUS = "status" -INTENT_LIST = "list" - -# Debug state -_LAST_INTENT_DEBUG: Dict[str, Any] = {} - - -def get_last_intent_debug() -> Dict[str, Any]: - """Get the last intent debug info.""" - return _LAST_INTENT_DEBUG - - -def _classify_intent_rules(q: str) -> str | None: - s = q.lower() - # Admin / maintenance first - if any(w in s for w in ["reindex", "reset", "recreate", "index now", "fresh index"]): - return INTENT_INDEX - if any(w in s for w in ["prune", "pruning", "cleanup", "clean up"]): - return INTENT_PRUNE - if any(w in s for w in ["status", "health", "points", "stats"]): - return INTENT_STATUS - if any(w in s for w in ["list collections", "collections", "list qdrant"]): - return INTENT_LIST - - # Search importers - check BEFORE tests to avoid "import" in test queries - if any(w in s for w in ["import", "imports", "importers", "who imports", "imports this", "importing modules", "files that import"]): - # Make sure it's not about "important" or similar - if not any(w in s for w in ["important", "importance"]): - return INTENT_SEARCH_IMPORTERS - - # Intent wrappers - if any(w in s for w in ["tests", "pytest", "unit test", "test file", "where are tests"]): - return INTENT_SEARCH_TESTS - - # Memory intents - be more specific to avoid false positives on "memory store implementation" - # Check for actual user-intent memory storage, not code references - memory_store_triggers = [ - "remember this", "save memory", "store memory", "remember that", - "save preference", "remember preference", "store a note", "save a note", "remember note" - ] - # IMPORTANT: "memory store" as a phrase often refers to code, not user intent - if any(w in s for w in memory_store_triggers): - # Exclude if it looks like a code search (has "implementation", "code", "function", etc) - if not any(exc in s for exc in ["implementation", "code", "function", "class", "module", "file", "search for"]): - return INTENT_MEMORY_STORE - if any(w in s for w in [ - "find memory", "recall", "retrieve memory", "memory search", "what did we save", - "recall notes", "find notes", "retrieve notes" - ]): - return INTENT_MEMORY_FIND - - # Symbol graph for callers - check BEFORE config to avoid false positives - if any(w in s for w in ["who calls", "callers of", "call sites", "function calls"]): - return INTENT_SYMBOL_GRAPH - if re.search(r"calls?\s+(the\s+)?\w+\s*(function|method)?", s): - return INTENT_SYMBOL_GRAPH - - # Config search - after callers check - if any(w in s for w in ["config", "yaml", "toml", "ini", "settings file", "configuration"]): - return INTENT_SEARCH_CONFIG - - # Fallback callers intent (used by search_callers_for) - if any(w in s for w in ["used by", "usage sites", "references this function"]): - return INTENT_SEARCH_CALLERS - - # Q&A-like prompts - if re.match(r"^(what|how|why|explain|describe|summarize)(\b|\s)", s): - return INTENT_ANSWER - if any(w in s for w in ["recap", "design doc", "architecture", "adr", "retrospective", "postmortem", "summary of", "summarize the design"]): - return INTENT_ANSWER - return None - - -def _intent_prototypes() -> Dict[str, List[str]]: - return { - INTENT_ANSWER: [ - "explain, describe, summarize, recap, design, architecture, ADR, why/how", - "summarize design decisions and architecture rationale", - ], - INTENT_SEARCH: [ - "find code references, search repository, locate files, find implementation", - "code search in repo, general lookup, search for implementation", - "find module, search function, locate class definition", - "search for memory store implementation", # Explicit example - ], - INTENT_MEMORY_STORE: [ - "remember this preference, save this note for later, store this memory", - "save my preference, remember that for next time", - # NOT: search for, find, implementation, code - ], - INTENT_MEMORY_FIND: [ - "what did we save, recall saved notes, retrieve memory, find my saved notes", - ], - INTENT_SEARCH_TESTS: [ - "find unit tests, test files, pytest, testing modules", - ], - INTENT_SEARCH_CONFIG: [ - "config files, configuration changes, yaml toml ini settings", - ], - INTENT_SEARCH_CALLERS: [ - "who calls this function, callers, usage sites, where is it used", - ], - INTENT_SEARCH_IMPORTERS: [ - "who imports this module, importers, importing modules, files that import", - ], - INTENT_SYMBOL_GRAPH: [ - "who calls this function, callers of, call graph, symbol callers", - ], - } - - -def _cosine(a: list[float], b: list[float]) -> float: - """Lightweight cosine similarity.""" - try: - s = 0.0 - na = 0.0 - nb = 0.0 - for i in range(min(len(a), len(b))): - va = float(a[i]) - vb = float(b[i]) - s += va * vb - na += va * va - nb += vb * vb - na = (na or 1.0) ** 0.5 - nb = (nb or 1.0) ** 0.5 - return s / (na * nb) - except Exception: - return 0.0 - - -def _embed_texts(texts: list[str]) -> list[list[float]]: - """Embed texts using available embedding model.""" - if not texts: - return [] - - # Try centralized embedder factory first - try: - from scripts.embedder import get_embedding_model - model_name = os.environ.get("EMBEDDING_MODEL", "BAAI/bge-base-en-v1.5") - em = get_embedding_model(model_name) - raw = list(em.embed(texts)) - return [v.tolist() if hasattr(v, "tolist") else list(v) for v in raw] - except ImportError: - pass - - # Try fastembed directly - try: - from fastembed import TextEmbedding - model_name = os.environ.get("EMBEDDING_MODEL", "BAAI/bge-base-en-v1.5") - em = TextEmbedding(model_name=model_name) - raw = list(em.embed(texts)) - return [v.tolist() if hasattr(v, "tolist") else list(v) for v in raw] - except Exception: - pass - - # Fallback to lexical - try: - from scripts.utils import lex_hash_vector_text - return [lex_hash_vector_text(t, dim=4096) for t in texts] - except Exception: - return [[float(len(t))] for t in texts] - - -def _classify_intent_ml(q: str) -> str: - global _LAST_INTENT_DEBUG - protos = _intent_prototypes() - labels = list(protos.keys()) - texts = [q] + ["\n".join(protos[l]) for l in labels] - vecs = _embed_texts(texts) - if not vecs or len(vecs) < len(texts): - _LAST_INTENT_DEBUG = { - "strategy": "ml", - "intent": INTENT_SEARCH, - "confidence": 0.0, - "query": q, - "top_candidate": INTENT_SEARCH, - "top_score": 0.0, - "threshold": 0.25, - "candidates": [], - "reason": "embed_failed", - } - return INTENT_SEARCH - qv = vecs[0] - sims = [] - for i, lab in enumerate(labels): - sims.append((lab, _cosine(qv, vecs[1 + i]))) - sims.sort(key=lambda x: x[1], reverse=True) - top, score = sims[0] - picked = top if score >= 0.25 else INTENT_SEARCH - _LAST_INTENT_DEBUG = { - "strategy": "ml", - "intent": picked, - "confidence": float(score), - "query": q, - "top_candidate": top, - "top_score": float(score), - "threshold": 0.25, - "candidates": [(name, float(val)) for name, val in sims[:5]], - "fallback": picked == INTENT_SEARCH and top != INTENT_SEARCH, - } - return picked - - -def classify_intent(q: str) -> str: - """Classify user query into an intent.""" - global _LAST_INTENT_DEBUG - ruled = _classify_intent_rules(q) - if ruled is not None: - _LAST_INTENT_DEBUG = { - "strategy": "rules", - "intent": ruled, - "confidence": 1.0, - "query": q, - } - return ruled - picked = _classify_intent_ml(q) - try: - if os.environ.get("DEBUG_ROUTER") and isinstance(_LAST_INTENT_DEBUG, dict): - if _LAST_INTENT_DEBUG.get("fallback"): - print(json.dumps({"router": {"intent_fallback": _LAST_INTENT_DEBUG}}), file=sys.stderr) - except Exception: - pass - return picked diff --git a/scripts/mcp_router/memory.py b/scripts/mcp_router/memory.py deleted file mode 100644 index cb22f872..00000000 --- a/scripts/mcp_router/memory.py +++ /dev/null @@ -1,69 +0,0 @@ -""" -mcp_router/memory.py - Memory store payload parsing. -""" -from __future__ import annotations - -import re -from typing import Any, Dict, Tuple - -_MEMORY_TRIGGER_RE = re.compile( - r"^(?:remember(?:\s+(?:this|that|me|to))?|save\s+memory|store\s+memory)\s*[:,\-]?\s*", - re.IGNORECASE, -) -_MEMORY_INTENT_SPLIT_RE = re.compile( - r"\b(?:then|and|also)\s+(?:reindex|index|recreate|prune|clean\s+up)\b", - re.IGNORECASE, -) -_MEMORY_META_KEYS = {"priority", "tag", "tags", "topic", "category", "owner"} - - -def parse_memory_store_payload(q: str) -> Tuple[str, Dict[str, Any]]: - """Parse memory store command, extracting content and metadata.""" - raw = str(q or "").strip() - if not raw: - return "", {} - cleaned = _MEMORY_TRIGGER_RE.sub("", raw, count=1).lstrip() - meta: Dict[str, Any] = {} - - def _assign_meta(key: str, value: str) -> None: - k = key.lower() - v = value.strip().strip(" \t\r\n,;.") - if not v: - return - if k in {"tag", "tags"}: - tags = [t.strip() for t in re.split(r"[,\s/]+", v) if t.strip()] - if tags: - meta["tags"] = tags - else: - meta[k] = v - - if cleaned.startswith("["): - m = re.match(r"\[([^\]]+)\]\s*(.*)", cleaned, flags=re.S) - if m: - meta_block = m.group(1) - cleaned = m.group(2) - for key, val in re.findall(r"(\w+)\s*=\s*([^\s,;]+(?:,[^\s,;]+)*)", meta_block): - if key.strip().lower() in _MEMORY_META_KEYS: - _assign_meta(key, val) - - while True: - m = re.match( - r"^(?P(?:priority|tag|tags|topic|category|owner))\s*=\s*(?P[^\s;:]+)\s*[,;:]?\s*(?P.*)$", - cleaned, - flags=re.IGNORECASE | re.S, - ) - if not m: - break - _assign_meta(m.group("key"), m.group("val")) - cleaned = m.group("rest") - - cleaned = cleaned.lstrip(":- ").lstrip() - - split = _MEMORY_INTENT_SPLIT_RE.search(cleaned) - if split: - cleaned = cleaned[: split.start()].rstrip(" ,;.") - - cleaned = cleaned.strip().strip('"').strip() - if not cleaned: - cleaned = raw - return cleaned, meta diff --git a/scripts/mcp_router/planning.py b/scripts/mcp_router/planning.py deleted file mode 100644 index fcc1b56a..00000000 --- a/scripts/mcp_router/planning.py +++ /dev/null @@ -1,312 +0,0 @@ -""" -mcp_router/planning.py - Tool planning and selection. -""" -from __future__ import annotations - -import os -from typing import Any, Dict, List, Tuple - -from .config import HTTP_URL_INDEXER -from .intent import ( - classify_intent, - INTENT_ANSWER, - INTENT_SEARCH, - INTENT_SEARCH_TESTS, - INTENT_SEARCH_CONFIG, - INTENT_SEARCH_CALLERS, - INTENT_SEARCH_IMPORTERS, - INTENT_MEMORY_STORE, - INTENT_MEMORY_FIND, - INTENT_SYMBOL_GRAPH, - INTENT_INDEX, - INTENT_PRUNE, - INTENT_STATUS, - INTENT_LIST, -) -from .memory import parse_memory_store_payload -from .hints import parse_repo_hints, clean_query_and_dsl, select_best_search_tool_by_signature -from .scratchpad import load_scratchpad, looks_like_repeat, looks_like_same_filters -from .client import discover_tool_endpoints - - -def build_plan(q: str) -> List[Tuple[str, Dict[str, Any]]]: - """Build execution plan for a query.""" - intent = classify_intent(q) - include_snippet = str(os.environ.get("ROUTER_INCLUDE_SNIPPET", "1")).lower() in {"1", "true", "yes", "on"} - search_limit = int(os.environ.get("ROUTER_SEARCH_LIMIT", "8") or 8) - max_tokens_env = os.environ.get("ROUTER_MAX_TOKENS", "").strip() - - def _reuse_last_filters(args: Dict[str, Any]) -> None: - try: - if looks_like_same_filters(q): - sp = load_scratchpad() - lf = sp.get("last_filters") if isinstance(sp, dict) else None - if isinstance(lf, dict): - for k in ("language", "under", "symbol", "ext", "path_glob", "not_glob"): - if k not in args and lf.get(k) not in (None, ""): - args[k] = lf.get(k) - except Exception: - pass - - # Repeat/redo handling - try: - if looks_like_repeat(q): - sp = load_scratchpad() - lp = sp.get("last_plan") - if isinstance(lp, list) and lp: - norm: list[tuple] = [] - for it in lp: - if isinstance(it, (list, tuple)) and len(it) == 2: - norm.append((it[0], it[1])) - if norm: - return norm - except Exception: - pass - - # Multi-intent: memory store + reindex - lowq = q.lower() - if any(w in lowq for w in ["remember this", "store memory", "save memory", "remember that"]) and any(w in lowq for w in ["reindex", "index now", "recreate", "fresh index"]): - idx_args: Dict[str, Any] = {} - if any(w in lowq for w in ["recreate", "fresh", "from scratch", "fresh index"]): - idx_args["recreate"] = True - info, meta = parse_memory_store_payload(q) - store_args: Dict[str, Any] = {"information": info or q.strip()} - if meta: - allowed = {"priority", "tags", "topic", "category", "owner"} - cleaned = {k: v for k, v in meta.items() if k in allowed and v not in (None, "", [])} - if cleaned: - store_args["metadata"] = cleaned - return [("store", store_args), ("qdrant_index_root", idx_args)] - - if intent == INTENT_INDEX: - recreate = True if any(w in q.lower() for w in ["recreate", "fresh", "from scratch"]) else None - args = {} - if recreate is True: - args["recreate"] = True - return [("qdrant_index_root", args)] - - if intent == INTENT_PRUNE: - return [("qdrant_prune", {})] - - if intent == INTENT_STATUS: - return [("qdrant_status", {})] - - if intent == INTENT_LIST: - return [("qdrant_list", {})] - - if intent == INTENT_SEARCH: - hints = parse_repo_hints(q) - clean_q, dsl_filters = clean_query_and_dsl(q) - args = {"query": clean_q} - if search_limit: - args["limit"] = search_limit - if include_snippet: - args["include_snippet"] = True - _reuse_last_filters(args) - - for k in ("language", "under", "symbol", "ext", "path_glob", "not_glob"): - v = dsl_filters.get(k) - if v not in (None, "") and k not in args: - args[k] = v - for k in ("language", "under", "symbol", "ext", "path_glob", "not_glob"): - v = hints.get(k) - if v not in (None, "") and k not in args: - args[k] = v - try: - tool_servers = discover_tool_endpoints(allow_network=False) - picked = select_best_search_tool_by_signature(q, tool_servers, allow_network=False) or "repo_search" - except Exception: - picked = "repo_search" - return [(picked, args)] - - if intent == INTENT_SEARCH_TESTS: - hints = parse_repo_hints(q) - clean_q, dsl_filters = clean_query_and_dsl(q) - args = {"query": clean_q} - if search_limit: - args["limit"] = search_limit - if include_snippet: - args["include_snippet"] = True - _reuse_last_filters(args) - - for k in ("language", "under", "symbol", "ext", "path_glob", "not_glob"): - v = dsl_filters.get(k) - if v not in (None, "") and k not in args: - args[k] = v - for k in ("language", "under", "symbol", "ext", "path_glob", "not_glob"): - v = hints.get(k) - if v not in (None, "") and k not in args: - args[k] = v - return [("search_tests_for", args)] - - if intent == INTENT_SEARCH_CONFIG: - hints = parse_repo_hints(q) - clean_q, dsl_filters = clean_query_and_dsl(q) - args = {"query": clean_q} - if search_limit: - args["limit"] = search_limit - if include_snippet: - args["include_snippet"] = True - _reuse_last_filters(args) - - for k in ("language", "under", "symbol", "ext", "path_glob", "not_glob"): - v = dsl_filters.get(k) - if v not in (None, "") and k not in args: - args[k] = v - for k in ("language", "under", "symbol", "ext", "path_glob", "not_glob"): - v = hints.get(k) - if v not in (None, "") and k not in args: - args[k] = v - return [("search_config_for", args)] - - if intent == INTENT_MEMORY_STORE: - info, meta = parse_memory_store_payload(q) - payload: Dict[str, Any] = {"information": info or q.strip()} - if meta: - allowed = {"priority", "tags", "topic", "category", "owner"} - cleaned = {k: v for k, v in meta.items() if k in allowed and v not in (None, "", [])} - if cleaned: - payload["metadata"] = cleaned - return [("store", payload)] - - if intent == INTENT_MEMORY_FIND: - args = {"query": q} - if search_limit: - args["limit"] = max(5, search_limit) - return [("find", args)] - - if intent == INTENT_SYMBOL_GRAPH: - hints = parse_repo_hints(q) - clean_q, dsl_filters = clean_query_and_dsl(q) - args = {"symbol": clean_q, "query_type": "callers"} - if search_limit: - args["limit"] = search_limit - for k in ("language", "under"): - v = dsl_filters.get(k) or hints.get(k) - if v not in (None, ""): - args[k] = v - return [("symbol_graph", args)] - - if intent == INTENT_SEARCH_CALLERS: - hints = parse_repo_hints(q) - clean_q, dsl_filters = clean_query_and_dsl(q) - args = {"query": clean_q} - if search_limit: - args["limit"] = search_limit - _reuse_last_filters(args) - - for k in ("language", "under", "symbol", "ext", "path_glob", "not_glob"): - v = dsl_filters.get(k) - if v not in (None, "") and k not in args: - args[k] = v - for k in ("language", "under", "symbol", "ext", "path_glob", "not_glob"): - v = hints.get(k) - if v not in (None, "") and k not in args: - args[k] = v - return [("search_callers_for", args)] - - if intent == INTENT_SEARCH_IMPORTERS: - hints = parse_repo_hints(q) - clean_q, dsl_filters = clean_query_and_dsl(q) - args = {"query": clean_q} - if search_limit: - args["limit"] = search_limit - _reuse_last_filters(args) - for k in ("language", "under", "symbol", "ext", "path_glob", "not_glob"): - v = dsl_filters.get(k) - if v not in (None, "") and k not in args: - args[k] = v - for k in ("language", "under", "symbol", "ext", "path_glob", "not_glob"): - v = hints.get(k) - if v not in (None, "") and k not in args: - args[k] = v - return [("search_importers_for", args)] - - if intent == INTENT_ANSWER: - def _looks_like_design_recap(s: str) -> bool: - low = s.lower() - return any(w in low for w in ["recap", "design doc", "architecture", "adr", "retrospective", "postmortem"]) and any(w in low for w in ["design", "summary", "recap", "explain"]) - - args: Dict[str, Any] = {"query": q} - if max_tokens_env: - try: - mt = int(max_tokens_env) - if mt > 0: - args["max_tokens"] = mt - except Exception: - pass - - hints = parse_repo_hints(q) - lowq = q.lower() - if "router" in lowq: - router_globs = ["**/mcp_router.py", "**/*router*.py"] - if not hints.get("path_glob"): - hints["path_glob"] = router_globs - if not hints.get("language"): - hints["language"] = "python" - for k in ("language", "under", "symbol", "ext", "path_glob", "not_glob"): - v = hints.get(k) - if v not in (None, ""): - args[k] = v - - plan: List[Tuple[str, Dict[str, Any]]] = [] - if _looks_like_design_recap(q): - plan.append(("find", {"query": q, "limit": 3})) - plan.extend([ - ("context_answer_compat", dict(args)), - ("context_answer", dict(args)), - ("repo_search", {**{k: v for k, v in args.items() if k != "max_tokens"}, "limit": max(5, search_limit)}), - ]) - return plan - - # Fallback - return [("repo_search", {"query": q, "limit": search_limit})] - - -async def route_query(query: str) -> Dict[str, Any]: - """ - Route a query to the appropriate tool. - - Returns dict with: - - tool: Selected tool name - - confidence: Routing confidence (0.0-1.0) - - intent: Classified intent - - args: Tool arguments - """ - intent = classify_intent(query) - plan = build_plan(query) - - if not plan: - return { - "tool": "repo_search", - "confidence": 0.3, - "intent": "fallback", - "args": {"query": query}, - } - - # First tool in plan is the primary selection - tool_name, tool_args = plan[0] - - # Map intent to confidence (higher for more specific intents) - intent_confidence = { - INTENT_ANSWER: 0.9, - INTENT_SEARCH: 0.7, - INTENT_SEARCH_TESTS: 0.85, - INTENT_SEARCH_CONFIG: 0.85, - INTENT_SEARCH_CALLERS: 0.85, - INTENT_SEARCH_IMPORTERS: 0.85, - INTENT_SYMBOL_GRAPH: 0.9, - INTENT_MEMORY_STORE: 0.9, - INTENT_MEMORY_FIND: 0.9, - INTENT_INDEX: 0.95, - INTENT_PRUNE: 0.95, - INTENT_STATUS: 0.95, - INTENT_LIST: 0.95, - } - - return { - "tool": tool_name, - "confidence": intent_confidence.get(intent, 0.5), - "intent": intent, - "args": tool_args, - } diff --git a/scripts/mcp_router/scratchpad.py b/scripts/mcp_router/scratchpad.py deleted file mode 100644 index 462ac2fa..00000000 --- a/scripts/mcp_router/scratchpad.py +++ /dev/null @@ -1,115 +0,0 @@ -""" -mcp_router/scratchpad.py - Persistent scratchpad for context preservation. -""" -from __future__ import annotations - -import json -import os -import time -from typing import Any, Dict - -from .config import scratchpad_ttl_sec - - -def scratchpad_path() -> str: - """Get scratchpad file path.""" - base = os.path.join(os.getcwd(), ".codebase") - try: - os.makedirs(base, exist_ok=True) - except Exception: - pass - return os.path.join(base, "router_scratchpad.json") - - -def load_scratchpad() -> Dict[str, Any]: - """Load scratchpad with TTL handling.""" - import sys - p = scratchpad_path() - try: - with open(p, "r", encoding="utf-8") as f: - j = json.load(f) - if isinstance(j, dict): - try: - ts = float(j.get("timestamp") or 0.0) - except Exception: - ts = 0.0 - ttl = scratchpad_ttl_sec() - if ts and ttl >= 0 and (time.time() - ts) > ttl: - stale_keys = ( - "last_plan", - "last_filters", - "mem_snippets", - "last_answer", - "last_citations", - "last_paths", - "last_metrics", - ) - removed = False - for stale_key in stale_keys: - if stale_key in j: - j.pop(stale_key, None) - removed = True - if removed: - j["timestamp"] = 0.0 - try: - print( - json.dumps({ - "router": { - "scratchpad": "stale_cleared", - "age_sec": round(time.time() - ts, 2), - } - }), - file=sys.stderr, - ) - except Exception: - pass - return j - except Exception: - pass - return {} - - -def save_scratchpad(d: Dict[str, Any]) -> None: - """Save scratchpad atomically.""" - p = scratchpad_path() - tmp = p + ".tmp" - try: - with open(tmp, "w", encoding="utf-8") as f: - json.dump(d, f) - try: - f.flush() - os.fsync(f.fileno()) - except Exception: - pass - os.replace(tmp, p) - except Exception: - try: - if os.path.exists(tmp): - os.unlink(tmp) - except Exception: - pass - - -def looks_like_repeat(q: str) -> bool: - """Check if query looks like a repeat request.""" - s = q.strip().lower() - pats = [ - "repeat", "again", "same thing", "do that again", "rerun", "run it again", "same as before", - ] - return any(p in s for p in pats) - - -def looks_like_same_filters(q: str) -> bool: - """Check if query asks to reuse filters.""" - s = q.strip().lower() - return any(p in s for p in ["same filters", "reuse filters", "previous filters"]) - - -def looks_like_expand(q: str) -> bool: - """Check if query asks for expansion.""" - s = q.strip().lower() - pats = [ - "expand on", "expand that", "expand the summary", "elaborate", - "more detail", "more details", "go deeper", "add details", - ] - return any(p in s for p in pats) diff --git a/scripts/mcp_router/validation.py b/scripts/mcp_router/validation.py deleted file mode 100644 index 52457fb6..00000000 --- a/scripts/mcp_router/validation.py +++ /dev/null @@ -1,87 +0,0 @@ -""" -mcp_router/validation.py - Response validation and metric extraction. -""" -from __future__ import annotations - -from typing import Any, Dict - -from .client import is_failure_response -from .config import divergence_thresholds, divergence_is_fatal_for - - -def is_result_good(tool: str, resp: Dict[str, Any]) -> bool: - """Check if result is good enough to stop the plan.""" - try: - r = resp.get("result") or {} - sc = r.get("structuredContent") or {} - rs = sc.get("result") or {} - - if tool in {"context_answer", "context_answer_compat"}: - ans = rs.get("answer") if isinstance(rs, dict) else None - if isinstance(ans, str): - s = ans.strip() - if s and not any(p in s.lower() for p in [ - "insufficient context", "not enough context", "no relevant", "don't know", "cannot answer" - ]): - return True - cites = rs.get("citations") if isinstance(rs, dict) else None - if isinstance(cites, list) and len(cites) > 0: - return True - return False - - if tool.startswith("search_") or tool == "repo_search": - total = rs.get("total") if isinstance(rs, dict) else None - if isinstance(total, int) and total > 0: - return True - results = rs.get("results") if isinstance(rs, dict) else None - if isinstance(results, list) and len(results) > 0: - return True - return False - - return not is_failure_response(resp) - except Exception: - return not is_failure_response(resp) - - -def extract_metric_from_resp(tool: str, resp: Dict[str, Any]) -> tuple[str, float] | None: - """Extract metric for divergence detection.""" - try: - r = resp.get("result") or {} - sc = r.get("structuredContent") or {} - rs = sc.get("result") or {} - - if tool in {"repo_search", "code_search", "context_search", "search_tests_for", "search_config_for", "search_callers_for", "search_importers_for"}: - tot = rs.get("total") - if isinstance(tot, (int, float)): - return ("total_results", float(tot)) - results = rs.get("results") - if isinstance(results, list): - return ("total_results", float(len(results))) - return None - - if tool in {"context_answer", "context_answer_compat"}: - cites = rs.get("citations") - if isinstance(cites, list): - return ("citations", float(len(cites))) - return ("citations", 0.0) - - if tool == "qdrant_status": - cnt = rs.get("count") - if isinstance(cnt, (int, float)): - return ("points", float(cnt)) - return None - except Exception: - return None - return None - - -def material_drop(prev: float | None, curr: float, drop_frac: float, min_base: int) -> bool: - """Check if there's a material drop in metrics.""" - try: - if prev is None: - return False - if prev < float(min_base): - return False - return curr < (float(prev) * float(drop_frac)) - except Exception: - return False diff --git a/scripts/mcp_toon.py b/scripts/mcp_toon.py deleted file mode 100644 index 3795c9c3..00000000 --- a/scripts/mcp_toon.py +++ /dev/null @@ -1,4 +0,0 @@ -#!/usr/bin/env python3 -"""Shim for backward compatibility. See scripts/mcp/toon.py""" -from scripts.mcp_impl.toon import * - diff --git a/scripts/mcp_utils.py b/scripts/mcp_utils.py deleted file mode 100644 index 9357a351..00000000 --- a/scripts/mcp_utils.py +++ /dev/null @@ -1,4 +0,0 @@ -#!/usr/bin/env python3 -"""Shim for backward compatibility. See scripts/mcp/utils.py""" -from scripts.mcp_impl.utils import * - diff --git a/scripts/mcp_workspace.py b/scripts/mcp_workspace.py deleted file mode 100644 index d8242bcb..00000000 --- a/scripts/mcp_workspace.py +++ /dev/null @@ -1,4 +0,0 @@ -#!/usr/bin/env python3 -"""Shim for backward compatibility. See scripts/mcp/workspace.py""" -from scripts.mcp_impl.workspace import * - diff --git a/scripts/memory_backup.py b/scripts/memory_backup.py index 410ed90a..1518d550 100644 --- a/scripts/memory_backup.py +++ b/scripts/memory_backup.py @@ -19,17 +19,8 @@ from typing import List, Dict, Any, Optional from pathlib import Path -# Add project root to path for imports -ROOT_DIR = Path(__file__).resolve().parent.parent -if str(ROOT_DIR) not in sys.path: - sys.path.insert(0, str(ROOT_DIR)) - -try: - from qdrant_client import QdrantClient - from qdrant_client.models import Filter, FieldCondition, MatchValue -except ImportError: - print("ERROR: qdrant-client not installed. Install with: pip install qdrant-client") - sys.exit(1) +from qdrant_client import QdrantClient +from qdrant_client.models import Filter, FieldCondition, MatchValue def get_qdrant_client() -> QdrantClient: @@ -316,4 +307,4 @@ def main(): if __name__ == "__main__": - main() \ No newline at end of file + main() diff --git a/scripts/memory_restore.py b/scripts/memory_restore.py index c8a7789f..20840be1 100644 --- a/scripts/memory_restore.py +++ b/scripts/memory_restore.py @@ -20,29 +20,10 @@ from typing import List, Dict, Any, Optional from pathlib import Path -# Add project root to path for imports -ROOT_DIR = Path(__file__).resolve().parent.parent -if str(ROOT_DIR) not in sys.path: - sys.path.insert(0, str(ROOT_DIR)) - -try: - from qdrant_client import QdrantClient - from qdrant_client.models import VectorParams, Distance, HnswConfigDiff -except ImportError as e: - print(f"ERROR: Missing required dependency: {e}") - print("Install with: pip install qdrant-client fastembed") - sys.exit(1) - -# Use embedder factory for Qwen3 support; fallback to direct fastembed -try: - from scripts.embedder import get_embedding_model as _get_embedding_model - _EMBEDDER_FACTORY = True -except ImportError: - _EMBEDDER_FACTORY = False - try: - from fastembed import TextEmbedding - except ImportError: - TextEmbedding = None # type: ignore +from qdrant_client import QdrantClient +from qdrant_client.models import VectorParams, Distance, HnswConfigDiff + +from scripts.embedder import get_embedding_model as _get_embedding_model def get_qdrant_client() -> QdrantClient: @@ -55,18 +36,7 @@ def get_qdrant_client() -> QdrantClient: def get_embedding_model(model_name: str): """Initialize embedding model with Qwen3 support via embedder factory.""" - # Try centralized embedder factory first (supports Qwen3 feature flag) - if _EMBEDDER_FACTORY: - return _get_embedding_model(model_name) - # Fallback to direct fastembed - if TextEmbedding is not None: - try: - return TextEmbedding(model_name=model_name) - except Exception as e: - raise RuntimeError(f"Failed to load embedding model '{model_name}': {e}") - raise RuntimeError( - "No embedding model available. Install fastembed: pip install fastembed" - ) + return _get_embedding_model(model_name) def ensure_collection_exists( @@ -424,4 +394,4 @@ def main(): if __name__ == "__main__": - main() \ No newline at end of file + main() diff --git a/scripts/path_scope.py b/scripts/path_scope.py new file mode 100644 index 00000000..2150926c --- /dev/null +++ b/scripts/path_scope.py @@ -0,0 +1,238 @@ +#!/usr/bin/env python3 +""" +Shared helpers for user-facing path scoping (`under`) across search tools. + +`under` is treated as a recursive subtree scope from the user's workspace +perspective (for example: "space" matches ".../space/**"). +""" + +from __future__ import annotations + +import os +import re +from functools import lru_cache +from typing import Any, Mapping, Optional, Set + +_MULTI_SLASH_RE = re.compile(r"/+") + + +def _normalize_path_token(value: Any) -> str: + s = str(value or "").strip().replace("\\", "/") + if not s: + return "" + s = _MULTI_SLASH_RE.sub("/", s) + # Normalize common "file://" style inputs. + if s.startswith("file://"): + s = s[7:] + return s.strip("/") + + +def _normalize_repo_hint(repo_hint: Any) -> str: + r = _normalize_path_token(repo_hint) + if not r: + return "" + return r.split("/")[-1] + + +def _repo_root_hint() -> str: + """Best-effort repository root (directory containing scripts/).""" + try: + return os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) + except Exception: + return "" + + +def _maybe_expand_from_cwd(token: str) -> str: + """Recover under values that were relativized from the current subdirectory.""" + s = str(token or "").strip().strip("/") + if not s or "/" in s: + return s + try: + root = _repo_root_hint() + if not root: + return s + cwd = os.path.abspath(os.getcwd()) + if not (cwd == root or cwd.startswith(root + os.sep)): + return s + rel_cwd = os.path.relpath(cwd, root).replace("\\", "/").strip("/") + if not rel_cwd: + return s + rebased = f"{rel_cwd}/{s}" + rebased_path = os.path.join(root, *rebased.split("/")) + top_level_path = os.path.join(root, s) + if os.path.exists(rebased_path) and not os.path.exists(top_level_path): + return rebased + except Exception: + pass + return s + + +@lru_cache(maxsize=256) +def _unique_segment_path(root: str, segment: str) -> str: + """Return unique repo-relative directory path for a segment, else empty.""" + if not root or not segment: + return "" + top = os.path.join(root, segment) + if os.path.exists(top): + return "" + matches: list[str] = [] + skip = { + ".git", + ".codebase", + "__pycache__", + ".venv", + "node_modules", + } + try: + for dirpath, dirnames, _filenames in os.walk(root): + dirnames[:] = [d for d in dirnames if d not in skip and not d.startswith(".")] + if segment in dirnames: + rel = os.path.relpath(os.path.join(dirpath, segment), root).replace("\\", "/") + matches.append(rel.strip("/")) + if len(matches) > 1: + return "" + except Exception: + return "" + return matches[0] if len(matches) == 1 else "" + + +def _maybe_expand_unique_segment(token: str) -> str: + """Resolve single-segment under values to a unique subtree when possible.""" + s = str(token or "").strip().strip("/") + if not s or "/" in s: + return s + root = _repo_root_hint() + if not root: + return s + found = _unique_segment_path(root, s) + return found or s + + +def normalize_under(under: Optional[str]) -> Optional[str]: + """Normalize user-provided `under` into a comparable path token.""" + s = _normalize_path_token(under) + if not s or s in {".", "work"}: + return None + # Accept absolute-style workspace prefixes while preserving user-facing scope. + if s.startswith("work/"): + s = s[len("work/") :] + s = _maybe_expand_from_cwd(s) + s = _maybe_expand_unique_segment(s) + if not s or s in {".", "work"}: + return None + return s + + +def _path_forms(path: Any, repo_hint: Any = None) -> Set[str]: + """Generate comparable path forms from a path-like value.""" + p = _normalize_path_token(path) + if not p: + return set() + + forms: Set[str] = {p} + + repo = _normalize_repo_hint(repo_hint) + + if p.startswith("work/"): + rest = p[len("work/") :] + if rest: + forms.add(rest) + if "/" in rest and repo: + head, tail = rest.split("/", 1) + if head.casefold() == repo.casefold() and tail: + forms.add(tail) + + if repo: + def _cf_to_orig_idx(orig: str, cf_index: int) -> int: + if cf_index <= 0: + return 0 + acc = 0 + for i, ch in enumerate(orig): + nxt = acc + len(ch.casefold()) + if nxt > cf_index: + return i + acc = nxt + return len(orig) + + repo_cf = repo.casefold() + repo_prefix_cf = repo_cf + "/" + marker_cf = "/" + repo_cf + "/" + for f in list(forms): + f_cf = f.casefold() + if f_cf.startswith(repo_prefix_cf): + forms.add(f[len(repo) + 1 :]) + idx = f_cf.find(marker_cf) + if idx >= 0: + tail_start = _cf_to_orig_idx(f, idx + len(marker_cf)) + tail = f[tail_start:] + if tail: + forms.add(tail) + + return {x for x in forms if x} + + +def metadata_path_forms(metadata: Mapping[str, Any]) -> Set[str]: + """Collect path forms from a metadata payload.""" + repo_hint = metadata.get("repo") + forms: Set[str] = set() + for key in ( + "repo_rel_path", + "path", + "container_path", + "host_path", + "path_prefix", + "file_path", + "rel_path", + "client_path", + ): + v = metadata.get(key) + if v: + forms.update(_path_forms(v, repo_hint=repo_hint)) + return forms + + +def metadata_matches_under(metadata: Mapping[str, Any], under: Optional[str]) -> bool: + """Return True when metadata falls under the requested subtree scope.""" + norm_under = normalize_under(under) + if not norm_under: + return True + + repo_hint = metadata.get("repo") + under_forms = _path_forms(norm_under, repo_hint=repo_hint) + under_forms.add(norm_under) + if not norm_under.startswith("work/"): + under_forms.add("work/" + norm_under) + + under_forms_l = {u.casefold() for u in under_forms if u} + if not under_forms_l: + return True + + has_repo_hint = bool(str(repo_hint or "").strip()) + + for cand in metadata_path_forms(metadata): + cand_forms = {cand} + # Compatibility fallback for points that only store /work//... paths + # but do not carry metadata.repo (older/benchmark/custom payloads). + if not has_repo_hint: + c0 = cand.strip("/") + if c0.startswith("work/"): + rest = c0[len("work/") :] + if "/" in rest: + _head, tail = rest.split("/", 1) + if tail: + cand_forms.add(tail) + + for cf in cand_forms: + c = cf.casefold() + for u in under_forms_l: + if c == u or c.startswith(u + "/"): + return True + return False + + +def path_matches_under(path: Any, under: Optional[str], repo_hint: Any = None) -> bool: + """Path-only convenience wrapper for `under` subtree matching.""" + md = {"path": path} + if repo_hint: + md["repo"] = repo_hint + return metadata_matches_under(md, under) diff --git a/scripts/pattern_detection/__init__.py b/scripts/pattern_detection/__init__.py deleted file mode 100644 index e66a5fd3..00000000 --- a/scripts/pattern_detection/__init__.py +++ /dev/null @@ -1,86 +0,0 @@ -# Pattern Detection System -# ======================== -# -# A DYNAMIC structural code similarity system with automatic pattern discovery. -# Based on AROMA (Meta) + code2vec AST path techniques. -# -# KEY FEATURES: -# 1. Works across ALL supported languages (16+ languages) -# 2. Patterns EMERGE automatically - not predefined -# 3. Cross-language matching - Python pattern matches Go/Rust/Java/etc. -# 4. Online learning - patterns improve as codebase is indexed -# 5. TOON output support for token-efficient responses -# -# Architecture: -# 1. PatternExtractor - Extract AST paths + control flow (normalized across languages) -# 2. PatternEncoder - Convert features to 64-dim vector via TF-IDF + LSH -# 3. PatternMiner - AROMA-style clustering to DISCOVER patterns dynamically -# 4. OnlinePatternLearner - Continuous learning as code is indexed -# 5. PatternSearch - Qdrant-backed similarity search with TOON support -# -# Usage: -# from scripts.pattern_detection import PatternExtractor, PatternEncoder -# -# extractor = PatternExtractor() -# encoder = PatternEncoder() -# -# # Extract pattern from code (works for ANY supported language) -# signature = extractor.extract(code, language="python") # or "go", "rust", etc. -# vector = encoder.encode(signature) -# -# # Search for similar code patterns -# from scripts.pattern_detection import pattern_search -# results = pattern_search(code, "python", output_format="toon") # Token-efficient -# -# # Dynamic pattern discovery -# from scripts.pattern_detection import get_pattern_learner -# learner = get_pattern_learner() -# learner.observe(code, path, language) # Called during indexing -# patterns = learner.query(example_code, "python") # Find matching patterns - -from .extractor import PatternExtractor, PatternSignature -from .encoder import PatternEncoder -from .catalog import PatternMiner, OnlinePatternLearner, get_pattern_learner, DiscoveredPattern -from .prune import AromaPruner, PrunedResult -from .search import ( - pattern_search, - find_similar_patterns, - search_by_pattern_description, - search_similar_code, - find_code_like, - PatternSearchResult, - PatternSearchResponse, - encode_pattern_results, -) - -# Backward compatibility aliases -PatternCatalog = PatternMiner -KNOWN_PATTERNS = [] # Now populated dynamically - -__all__ = [ - # Core extraction - "PatternExtractor", - "PatternSignature", - "PatternEncoder", - # AROMA pruning (Phase II & III) - "AromaPruner", - "PrunedResult", - # Dynamic discovery - "PatternMiner", - "OnlinePatternLearner", - "get_pattern_learner", - "DiscoveredPattern", - # Search (with TOON support) - "pattern_search", - "find_similar_patterns", - "search_by_pattern_description", - "search_similar_code", - "find_code_like", - "PatternSearchResult", - "PatternSearchResponse", - "encode_pattern_results", - # Backward compatibility - "PatternCatalog", - "KNOWN_PATTERNS", -] - diff --git a/scripts/pattern_detection/catalog.py b/scripts/pattern_detection/catalog.py deleted file mode 100644 index 0e125004..00000000 --- a/scripts/pattern_detection/catalog.py +++ /dev/null @@ -1,585 +0,0 @@ -""" -Dynamic Pattern Discovery - AROMA-inspired self-organizing pattern system. - -This is NOT a static catalog. Patterns are DISCOVERED dynamically by: - -1. **Clustering** - Group structurally similar code across the codebase -2. **Intersection** - Extract the COMMON structure from clusters (the "pattern") -3. **Emergence** - Patterns naturally emerge from code, not manual definition -4. **Cross-language** - Normalized AST means Python patterns match Go/Rust/etc. - -Key insight from AROMA paper: Don't predefine patterns. Instead: -- Query: "find code similar to X" -- Retrieve: Top-K structurally similar snippets -- Intersect: Find what's COMMON across all K snippets -- That common structure IS the pattern - -This approach: -- Discovers domain-specific patterns unique to each codebase -- Adapts as code evolves -- Works for patterns nobody thought to predefine -- Cross-language pattern matching via normalized AST -""" - -from dataclasses import dataclass, field -from typing import List, Dict, Optional, Tuple, Set, Any, Iterator -from collections import Counter, defaultdict -import hashlib -import math -import time - - -@dataclass -class DiscoveredPattern: - """A pattern discovered through structural clustering - NOT predefined.""" - - # Unique ID derived from structural fingerprint - pattern_id: str - - # The common structural features across examples - common_ngrams: Counter # Structural n-grams present in ALL examples - common_cf: Dict[str, Any] # Common control flow features - - # Representative vector (centroid of cluster) - centroid: List[float] = field(default_factory=list) - - # Example file paths that exhibit this pattern - exemplars: List[str] = field(default_factory=list) - - # Languages this pattern appears in - languages: Set[str] = field(default_factory=set) - - # Discovery metadata - cluster_size: int = 0 - coherence_score: float = 0.0 # How tight is the cluster - discovery_time: float = 0.0 - - # Auto-generated description based on structure - auto_description: str = "" - - def __hash__(self): - return hash(self.pattern_id) - - -class PatternMiner: - """ - AROMA-style pattern discovery through structural clustering. - - Instead of predefined patterns, this DISCOVERS patterns by: - 1. Indexing structural signatures of all code spans - 2. Clustering similar structures (locality-sensitive hashing) - 3. Intersecting clusters to extract common structure - 4. Ranking patterns by frequency and coherence - """ - - # Clustering parameters - MIN_CLUSTER_SIZE = 3 # Minimum examples to form a pattern - MAX_PATTERNS = 1000 # Cap on discovered patterns - SIMILARITY_THRESHOLD = 0.7 # Min similarity to be in same cluster - - # LSH parameters for fast clustering - NUM_HASH_TABLES = 10 - HASH_SIZE = 8 - - def __init__(self): - self._patterns: Dict[str, DiscoveredPattern] = {} - # lsh_hash -> [(path, language, vector, signature)] - self._signature_index: Dict[str, List[Tuple[str, str, List[float], Any]]] = defaultdict(list) - self._extractor = None - self._encoder = None - - def _lazy_init(self): - if self._extractor is None: - from .extractor import PatternExtractor - from .encoder import PatternEncoder - self._extractor = PatternExtractor() - self._encoder = PatternEncoder() - - def index_snippet(self, code: str, path: str, language: str) -> str: - """Index a code snippet for pattern discovery. Returns signature hash.""" - self._lazy_init() - - sig = self._extractor.extract(code, language) - vec = self._encoder.encode(sig) - - # Generate LSH hashes for fast clustering - lsh_hashes = self._compute_lsh_hashes(vec) - - # Index under all LSH hashes - for h in lsh_hashes: - self._signature_index[h].append((path, language, vec, sig)) - - return sig.fingerprint() - - def _compute_lsh_hashes(self, vec: List[float]) -> List[str]: - """Compute LSH hashes for approximate nearest neighbor clustering.""" - hashes = [] - - for table_idx in range(self.NUM_HASH_TABLES): - # Random hyperplane projection (deterministic via seed) - bits = [] - for i in range(self.HASH_SIZE): - # Use table and bit index as seed for reproducibility - seed = table_idx * 1000 + i - projection = sum( - vec[j] * self._pseudo_random(seed, j) - for j in range(len(vec)) - ) - bits.append('1' if projection >= 0 else '0') - - hash_val = f"t{table_idx}_{''.join(bits)}" - hashes.append(hash_val) - - return hashes - - def _pseudo_random(self, seed: int, idx: int) -> float: - """Deterministic pseudo-random for LSH projection.""" - h = hashlib.md5(f"{seed}:{idx}".encode()).digest() - val = int.from_bytes(h[:4], 'big') / (2**32) - return val * 2 - 1 # Map to [-1, 1] - - def discover_patterns(self, min_support: int = 3) -> List[DiscoveredPattern]: - """ - Discover patterns by clustering and intersection. - - This is the AROMA magic: - 1. Find clusters of similar signatures via LSH - 2. For each cluster, intersect to find COMMON structure - 3. That common structure becomes a discovered pattern - """ - self._lazy_init() - - discovered = [] - processed_clusters: Set[frozenset] = set() - - # For each LSH bucket with enough items - for lsh_hash, items in self._signature_index.items(): - if len(items) < min_support: - continue - - # Get cluster members - use (path, signature_fingerprint) to allow multiple spans per file - members = frozenset((item[0], item[3].fingerprint()) for item in items) - if members in processed_clusters: - continue - processed_clusters.add(members) - - # Extract unique paths for pattern ID and exemplars - paths = list(set(item[0] for item in items)) - - # Compute cluster coherence (avg pairwise similarity) - vecs = [item[2] for item in items] - coherence = self._cluster_coherence(vecs) - - if coherence < self.SIMILARITY_THRESHOLD: - continue - - # INTERSECT: Find common structure across all cluster members - sigs = [item[3] for item in items] - common = self._intersect_signatures(sigs) - - if not common: - continue - - # Create discovered pattern - pattern = DiscoveredPattern( - pattern_id=hashlib.md5(str(sorted(paths)).encode()).hexdigest()[:12], - common_ngrams=common['ngrams'], - common_cf=common['cf'], - centroid=self._compute_centroid(vecs), - exemplars=paths[:10], - languages=set(item[1] for item in items), - cluster_size=len(items), - coherence_score=coherence, - discovery_time=time.time(), - auto_description=self._generate_description(common), - ) - - discovered.append(pattern) - self._patterns[pattern.pattern_id] = pattern - - # Sort by cluster size and coherence - discovered.sort(key=lambda p: (p.cluster_size, p.coherence_score), reverse=True) - return discovered[:self.MAX_PATTERNS] - - def _cluster_coherence(self, vecs: List[List[float]]) -> float: - """Compute average pairwise cosine similarity within cluster.""" - if len(vecs) < 2: - return 1.0 - - total_sim = 0.0 - count = 0 - - for i in range(len(vecs)): - for j in range(i + 1, len(vecs)): - sim = self._cosine_similarity(vecs[i], vecs[j]) - total_sim += sim - count += 1 - - return total_sim / count if count > 0 else 0.0 - - def _cosine_similarity(self, a: List[float], b: List[float]) -> float: - """Compute cosine similarity.""" - dot = sum(x * y for x, y in zip(a, b)) - norm_a = math.sqrt(sum(x * x for x in a)) - norm_b = math.sqrt(sum(x * x for x in b)) - if norm_a < 1e-10 or norm_b < 1e-10: - return 0.0 - return dot / (norm_a * norm_b) - - def _compute_centroid(self, vecs: List[List[float]]) -> List[float]: - """Compute centroid of vectors.""" - if not vecs: - return [] - dim = len(vecs[0]) - centroid = [0.0] * dim - for vec in vecs: - for i, v in enumerate(vec): - centroid[i] += v - return [c / len(vecs) for c in centroid] - - def _intersect_signatures(self, sigs) -> Optional[Dict[str, Any]]: - """ - AROMA CORE: Intersect signatures to find COMMON structure. - - This extracts what's shared across ALL examples in a cluster. - The intersection IS the pattern - features that appear in - every instance of this pattern. - """ - if not sigs: - return None - - # Intersect structural n-grams: keep only those in ALL signatures - common_ngrams = None - for sig in sigs: - if common_ngrams is None: - common_ngrams = Counter(sig.structural_ngrams) - else: - # Keep only n-grams present in both, with min count - common_ngrams = Counter({ - k: min(common_ngrams[k], sig.structural_ngrams[k]) - for k in common_ngrams - if k in sig.structural_ngrams - }) - - if not common_ngrams: - return None - - # Intersect control flow: keep features present in ALL signatures - common_cf = {} - cf_keys = ['max_loop_depth', 'loop_count', 'branch_count', 'try_count', - 'has_finally', 'has_catch', 'has_resource_guard', 'match_count'] - - for key in cf_keys: - values = [sig.control_flow.get(key) for sig in sigs] - if all(v == values[0] for v in values): - common_cf[key] = values[0] - elif all(isinstance(v, bool) for v in values): - # For bool, use AND (present in ALL) - common_cf[key] = all(values) - elif all(isinstance(v, (int, float)) for v in values): - # For numbers, use min (lower bound guarantee) - common_cf[key] = min(values) - - return { - 'ngrams': common_ngrams, - 'cf': common_cf, - } - - def _generate_description(self, common: Dict[str, Any]) -> str: - """Auto-generate a description from common structural features.""" - parts = [] - cf = common.get('cf', {}) - - # Describe control flow - if cf.get('loop_count', 0) > 0: - depth = cf.get('max_loop_depth', 1) - parts.append(f"loop structure (depth {depth})") - - if cf.get('try_count', 0) > 0: - try_desc = "error handling" - if cf.get('has_catch'): - try_desc += " with catch" - if cf.get('has_finally'): - try_desc += " with cleanup" - parts.append(try_desc) - - if cf.get('branch_count', 0) > 0: - parts.append(f"{cf['branch_count']}+ conditional branches") - - if cf.get('has_resource_guard'): - parts.append("resource management") - - if cf.get('match_count', 0) > 0: - parts.append("pattern matching") - - # Describe common n-grams (top structural patterns) - ngrams = common.get('ngrams', Counter()) - if ngrams: - top_ngrams = ngrams.most_common(3) - # Could analyze n-gram content for more specific description - - if not parts: - return "structural pattern" - - return "Pattern with " + ", ".join(parts) - - def find_similar_to(self, code: str, language: str, top_k: int = 10) -> List[Tuple[str, float]]: - """ - Find code snippets similar to the given example. - - This is the query path: - 1. Extract signature from query code - 2. LSH lookup for candidate clusters - 3. Rank by vector similarity - """ - self._lazy_init() - - sig = self._extractor.extract(code, language) - vec = self._encoder.encode(sig) - lsh_hashes = self._compute_lsh_hashes(vec) - - # Gather candidates from all matching LSH buckets - # Key by (path, signature_fingerprint) to allow multiple spans per file - candidates = {} - for h in lsh_hashes: - for path, lang, cvec, csig in self._signature_index.get(h, []): - key = (path, csig.fingerprint()) - if key not in candidates: - sim = self._cosine_similarity(vec, cvec) - candidates[key] = (sim, lang, path) - - # Sort by similarity - results = [(path, sim) for key, (sim, lang, path) in candidates.items()] - results.sort(key=lambda x: -x[1]) - - return results[:top_k] - - def get_pattern_for_examples(self, paths: List[str]) -> Optional[DiscoveredPattern]: - """ - Given a set of example paths, discover their common pattern. - - This is AROMA's recommendation path in reverse: - - User says "these 3 files do the same thing" - - We extract and intersect their structures - - Return the discovered pattern - """ - # Find signatures for these paths - sigs = [] - vecs = [] - languages = set() - - for lsh_hash, items in self._signature_index.items(): - for path, lang, vec, sig in items: - if path in paths: - sigs.append(sig) - vecs.append(vec) - languages.add(lang) - - if len(sigs) < 2: - return None - - common = self._intersect_signatures(sigs) - if not common: - return None - - return DiscoveredPattern( - pattern_id=hashlib.md5(str(sorted(paths)).encode()).hexdigest()[:12], - common_ngrams=common['ngrams'], - common_cf=common['cf'], - centroid=self._compute_centroid(vecs), - exemplars=paths, - languages=languages, - cluster_size=len(sigs), - coherence_score=self._cluster_coherence(vecs), - discovery_time=time.time(), - auto_description=self._generate_description(common), - ) - - -# ============================================================================= -# Online Pattern Learning - The GENIUS part -# ============================================================================= - -class OnlinePatternLearner: - """ - Continuously discovers patterns as code is indexed. - - This is the genius: patterns emerge AUTOMATICALLY during indexing. - No manual definition. No static catalog. Pure emergence. - - Key innovations: - 1. Incremental clustering - don't reprocess everything - 2. Pattern aging - old patterns decay, active ones strengthen - 3. Cross-language unification - Python pattern matches Go code - 4. Query-driven refinement - searches refine pattern boundaries - """ - - def __init__(self, persist_path: Optional[str] = None): - self.miner = PatternMiner() - self.persist_path = persist_path - - # Pattern strength tracking (reinforcement learning inspired) - self._pattern_hits: Counter = Counter() # How often each pattern matched a query - self._pattern_age: Dict[str, float] = {} # When pattern was last useful - - # Incremental update tracking - self._pending_snippets: List[Tuple[str, str, str]] = [] # (code, path, lang) - self._batch_size = 100 # Process patterns every N snippets - - def observe(self, code: str, path: str, language: str): - """ - Observe a code snippet during indexing. - - Called by the indexer for every code span. Patterns - emerge naturally from the accumulated observations. - """ - self.miner.index_snippet(code, path, language) - self._pending_snippets.append((code, path, language)) - - # Periodically discover new patterns - if len(self._pending_snippets) >= self._batch_size: - self._discover_batch() - - def _discover_batch(self): - """Process pending snippets and discover new patterns.""" - if not self._pending_snippets: - return - - new_patterns = self.miner.discover_patterns(min_support=3) - - for pattern in new_patterns: - self._pattern_age[pattern.pattern_id] = time.time() - - self._pending_snippets.clear() - - def query(self, code: str, language: str, top_k: int = 5) -> List[Tuple[DiscoveredPattern, float]]: - """ - Find patterns matching a query code example. - - Also updates pattern strength (query-driven learning): - - Patterns that match queries get reinforced - - This naturally surfaces the most USEFUL patterns - """ - self._lazy_init_miner() - - sig = self.miner._extractor.extract(code, language) - vec = self.miner._encoder.encode(sig) - - results = [] - for pattern_id, pattern in self.miner._patterns.items(): - if not pattern.centroid: - continue - sim = self.miner._cosine_similarity(vec, pattern.centroid) - if sim > 0.5: # Threshold - results.append((pattern, sim)) - # Reinforce this pattern - self._pattern_hits[pattern_id] += 1 - self._pattern_age[pattern_id] = time.time() - - results.sort(key=lambda x: -x[1]) - return results[:top_k] - - def _lazy_init_miner(self): - self.miner._lazy_init() - - def get_top_patterns(self, n: int = 20) -> List[DiscoveredPattern]: - """Get the most useful discovered patterns (by query hits).""" - # Score by: hits * recency * cluster_coherence - now = time.time() - scored = [] - - for pid, pattern in self.miner._patterns.items(): - hits = self._pattern_hits.get(pid, 0) - age = now - self._pattern_age.get(pid, now) - recency = math.exp(-age / (30 * 24 * 3600)) # 30 day half-life - score = (hits + 1) * recency * pattern.coherence_score - scored.append((score, pattern)) - - scored.sort(key=lambda x: -x[0]) - return [p for _, p in scored[:n]] - - def natural_language_query(self, query: str, top_k: int = 5) -> List[DiscoveredPattern]: - """ - Find patterns matching a natural language description. - - This bridges "find retry patterns" → discovered patterns - without needing predefined keywords. Uses the auto-generated - descriptions for matching. - """ - query_lower = query.lower() - keywords = set(query_lower.split()) - - scored = [] - for pattern in self.miner._patterns.values(): - desc_words = set(pattern.auto_description.lower().split()) - - # Simple keyword overlap score - overlap = len(keywords & desc_words) - - # Boost by pattern quality - quality = pattern.coherence_score * math.log1p(pattern.cluster_size) - - score = overlap * quality - if score > 0: - scored.append((score, pattern)) - - scored.sort(key=lambda x: -x[0]) - return [p for _, p in scored[:top_k]] - - def explain_pattern(self, pattern: DiscoveredPattern) -> str: - """Generate human-readable explanation of a discovered pattern.""" - lines = [ - f"Pattern: {pattern.auto_description}", - f"Found in {pattern.cluster_size} places across {len(pattern.languages)} language(s)", - f"Coherence: {pattern.coherence_score:.2%}", - "", - "Common structural elements:", - ] - - # Describe control flow - cf = pattern.common_cf - if cf.get('loop_count', 0) > 0: - lines.append(f" • Loops (max depth: {cf.get('max_loop_depth', 1)})") - if cf.get('try_count', 0) > 0: - parts = [" • Error handling"] - if cf.get('has_catch'): - parts.append("with catch/except") - if cf.get('has_finally'): - parts.append("with cleanup") - lines.append(" ".join(parts)) - if cf.get('branch_count', 0) > 0: - lines.append(f" • Conditional branches: {cf['branch_count']}+") - if cf.get('has_resource_guard'): - lines.append(" • Resource management (with/using/defer)") - - lines.append("") - lines.append("Example locations:") - for path in pattern.exemplars[:5]: - lines.append(f" - {path}") - - return "\n".join(lines) - - -# Singleton instance for global pattern learning -_global_learner: Optional[OnlinePatternLearner] = None - - -def get_pattern_learner() -> OnlinePatternLearner: - """Get or create the global pattern learner.""" - global _global_learner - if _global_learner is None: - _global_learner = OnlinePatternLearner() - return _global_learner - - -# ============================================================================= -# Backward compatibility aliases -# ============================================================================= - -PatternCatalog = PatternMiner # Legacy name -SEED_PATTERNS = { - "retry": "for i in range(n): try: x() break except: sleep(2**i)", - "cleanup": "x = acquire() try: use(x) finally: x.close()", -} -KNOWN_PATTERNS: List[DiscoveredPattern] = [] # Populated dynamically by learner - diff --git a/scripts/pattern_detection/encoder.py b/scripts/pattern_detection/encoder.py deleted file mode 100644 index ab11ef25..00000000 --- a/scripts/pattern_detection/encoder.py +++ /dev/null @@ -1,243 +0,0 @@ -""" -Pattern Encoder - Convert pattern signatures to dense vectors. - -Encoding dimensions (64 total): -- MinHash n-grams: 16-dim -- Weisfeiler-Lehman kernel: 8-dim -- Control flow features: 16-dim -- CFG fingerprint: 8-dim -- SimHash bits: 8-dim -- Spectral features: 8-dim -""" - -from typing import List, Dict, Tuple, Optional, Any -from collections import Counter -import hashlib -import math - -from .extractor import PatternSignature - - -class PatternEncoder: - """Encode pattern signatures as dense vectors for similarity search.""" - - TOTAL_DIM = 64 - MINHASH_DIM = 16 - WL_DIM = 8 - CONTROL_DIM = 16 - CFG_DIM = 8 - SIMHASH_DIM = 8 - SPECTRAL_DIM = 8 - - HASH_PRIMES = [2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53] - - def __init__(self, normalize: bool = True): - self.normalize = normalize - - def encode(self, signature: PatternSignature) -> List[float]: - """Encode pattern signature as 64-dim vector.""" - vector = [] - - vector.extend(self._encode_minhash(signature.structural_ngrams)) - vector.extend(self._encode_wl_kernel(signature.wl_labels)) - vector.extend(self._encode_control_flow(signature.control_flow)) - vector.extend(self._encode_cfg(signature.cfg_nodes, signature.cfg_edges)) - vector.extend(self._encode_simhash(signature.simhash)) - vector.extend(self._encode_spectral(signature.spectral_features)) - - if self.normalize: - vector = self._l2_normalize(vector) - - return vector - - # ========================================================================= - # MinHash for structural n-grams (Jaccard similarity) - # ========================================================================= - - def _encode_minhash(self, ngrams: Counter) -> List[float]: - """MinHash encoding for Jaccard similarity preservation.""" - if not ngrams: - return [0.0] * self.MINHASH_DIM - - ngram_set = set(ngrams.keys()) - minhash = [] - - for i, prime in enumerate(self.HASH_PRIMES[:self.MINHASH_DIM]): - min_h = float('inf') - for ngram in ngram_set: - h = self._hash_item(str(ngram), prime, i) - min_h = min(min_h, h) - minhash.append((min_h % 1000) / 1000.0 if min_h != float('inf') else 0.0) - - return minhash - - def _hash_item(self, item: str, prime: int, seed: int) -> int: - h = hashlib.md5(item.encode()).digest() - return (int.from_bytes(h[:4], 'big') * prime + seed) % (2**31) - - # ========================================================================= - # Weisfeiler-Lehman Graph Kernel encoding - # ========================================================================= - - def _encode_wl_kernel(self, wl_labels: Dict[int, List[str]]) -> List[float]: - """Encode WL labels into fixed-size vector via feature hashing.""" - vec = [0.0] * self.WL_DIM - - if not wl_labels: - return vec - - label_counts = Counter() - for node_id, labels in wl_labels.items(): - for label in labels: - label_counts[label] += 1 - - total = sum(label_counts.values()) - if total == 0: - return vec - - for label, count in label_counts.items(): - bucket = int(hashlib.md5(label.encode()).hexdigest()[:4], 16) % self.WL_DIM - vec[bucket] += count / total - - return vec - - # ========================================================================= - # Control Flow encoding - # ========================================================================= - - def _encode_control_flow(self, cf: Dict[str, Any]) -> List[float]: - """Encode control flow features (16-dim).""" - vec = [0.0] * self.CONTROL_DIM - - vec[0] = self._log_scale(cf.get("max_loop_depth", 0), 5) - vec[1] = self._log_scale(cf.get("loop_count", 0), 10) - vec[2] = 1.0 if cf.get("loop_count", 0) >= 1 else 0.0 - vec[3] = 1.0 if cf.get("try_in_loop", False) else 0.0 # retry pattern - - vec[4] = self._log_scale(cf.get("branch_count", 0), 20) - vec[5] = 1.0 if cf.get("branch_count", 0) >= 1 else 0.0 - vec[6] = self._log_scale(cf.get("match_count", 0), 5) - vec[7] = 1.0 if cf.get("branch_in_loop", False) else 0.0 # filter pattern - - vec[8] = self._log_scale(cf.get("try_count", 0), 5) - vec[9] = 1.0 if cf.get("try_count", 0) >= 1 else 0.0 - vec[10] = 1.0 if cf.get("has_catch", False) or cf.get("has_except", False) else 0.0 - vec[11] = 1.0 if cf.get("has_finally", False) else 0.0 - - vec[12] = 1.0 if cf.get("has_resource_guard", False) else 0.0 - vec[13] = self._log_scale(cf.get("max_nesting_depth", 0), 8) - vec[14] = self._log_scale(cf.get("func_count", 0), 10) - vec[15] = self._log_scale(cf.get("class_count", 0), 5) - - return vec - - def _log_scale(self, count: int, max_val: float = 10.0) -> float: - if count <= 0: - return 0.0 - return min(1.0, math.log1p(count) / math.log1p(max_val)) - - # ========================================================================= - # CFG Fingerprint encoding - # ========================================================================= - - def _encode_cfg(self, cfg_nodes: Dict[int, str], cfg_edges: List[Tuple[int, int, str]]) -> List[float]: - """Encode CFG structure via edge type distribution and graph properties.""" - vec = [0.0] * self.CFG_DIM - - if not cfg_edges: - return vec - - edge_types = Counter(e[2] for e in cfg_edges) - # Exclude loop_entry from ratio calculation - it's structural, not semantic - # This keeps vector meaning stable after adding loop_entry edges - ratio_edges = len(cfg_edges) - edge_types.get("loop_entry", 0) - if ratio_edges <= 0: - ratio_edges = 1 # Avoid division by zero - - vec[0] = edge_types.get("sequential", 0) / ratio_edges - vec[1] = edge_types.get("branch_true", 0) / ratio_edges - vec[2] = edge_types.get("branch_false", 0) / ratio_edges - vec[3] = edge_types.get("loop_back", 0) / ratio_edges - vec[4] = edge_types.get("exception", 0) / ratio_edges - - # Graph density (use total_edges including loop_entry for topology) - total_edges = len(cfg_edges) - n_nodes = len(cfg_nodes) - if n_nodes > 1: - vec[5] = min(1.0, total_edges / (n_nodes * 2)) - - # Cyclomatic complexity approximation: E - N + 2P - vec[6] = self._log_scale(total_edges - n_nodes + 2, 20) - - # Node type entropy - if cfg_nodes: - node_types = Counter(cfg_nodes.values()) - total_nodes = len(cfg_nodes) - entropy = -sum((c/total_nodes) * math.log2(c/total_nodes + 1e-10) - for c in node_types.values()) - vec[7] = min(1.0, entropy / 4.0) - - return vec - - # ========================================================================= - # SimHash encoding (LSH bits) - # ========================================================================= - - def _encode_simhash(self, simhash: int) -> List[float]: - """Extract 8 representative bits from 64-bit SimHash.""" - vec = [0.0] * self.SIMHASH_DIM - - # Sample bits at regular intervals - for i in range(self.SIMHASH_DIM): - bit_pos = i * 8 - vec[i] = 1.0 if (simhash >> bit_pos) & 1 else 0.0 - - return vec - - # ========================================================================= - # Spectral Features encoding - # ========================================================================= - - def _encode_spectral(self, spectral_features: List[float]) -> List[float]: - """Encode spectral features (eigenvalues).""" - if not spectral_features: - return [0.0] * self.SPECTRAL_DIM - - result = spectral_features[:self.SPECTRAL_DIM] - while len(result) < self.SPECTRAL_DIM: - result.append(0.0) - - return result - - # ========================================================================= - # Utilities - # ========================================================================= - - def _l2_normalize(self, vec: List[float]) -> List[float]: - norm = math.sqrt(sum(x*x for x in vec)) - if norm < 1e-10: - return vec - return [x / norm for x in vec] - - def similarity(self, vec_a: List[float], vec_b: List[float]) -> float: - """Cosine similarity.""" - dot = sum(a * b for a, b in zip(vec_a, vec_b)) - norm_a = math.sqrt(sum(a*a for a in vec_a)) - norm_b = math.sqrt(sum(b*b for b in vec_b)) - if norm_a < 1e-10 or norm_b < 1e-10: - return 0.0 - return dot / (norm_a * norm_b) - - def hamming_distance(self, simhash_a: int, simhash_b: int) -> int: - """Hamming distance between two SimHash values.""" - return bin(simhash_a ^ simhash_b).count('1') - - def tree_edit_distance_approx(self, paths_a: List[int], paths_b: List[int]) -> float: - """Approximate tree edit distance via Jaccard of path hashes.""" - if not paths_a or not paths_b: - return 1.0 - set_a = set(paths_a) - set_b = set(paths_b) - intersection = len(set_a & set_b) - union = len(set_a | set_b) - return 1.0 - (intersection / union) if union > 0 else 1.0 diff --git a/scripts/pattern_detection/extractor.py b/scripts/pattern_detection/extractor.py deleted file mode 100644 index 9e498a22..00000000 --- a/scripts/pattern_detection/extractor.py +++ /dev/null @@ -1,1084 +0,0 @@ -""" -Pattern Extractor - Language-agnostic AST-based structural feature extraction. - -Extracts three types of features from code in ANY language: -1. AST Paths - paths between terminals in AST -2. Structural Features - node type sequences, depth patterns -3. Control Flow Features - loop/branch structure fingerprints - -The key innovation: features are CONTENT-AGNOSTIC and LANGUAGE-NORMALIZED. -Variable names, string literals, and specific function names are abstracted -away. Language-specific AST nodes are normalized to universal concepts -(LOOP, BRANCH, TRY, FUNC, etc.) enabling cross-language pattern matching. - -Supported languages: Python, JavaScript, TypeScript, Go, Rust, Java, C, C++, -Ruby, PHP, C#, Kotlin, Swift, Scala, and more via Tree-sitter. -""" - -from dataclasses import dataclass, field -from typing import List, Dict, Set, Tuple, Optional, Any -from collections import Counter -import hashlib -import logging -import re - -logger = logging.getLogger(__name__) - - -@dataclass -class PatternSignature: - """Language-independent structural signature of code.""" - - ast_paths: List[Tuple[str, str, str, int]] = field(default_factory=list) - structural_ngrams: Counter = field(default_factory=Counter) - control_flow: Dict[str, Any] = field(default_factory=dict) - language: str = "unknown" - _hash: Optional[str] = None - - wl_labels: Dict[int, List[str]] = field(default_factory=dict) - cfg_edges: List[Tuple[int, int, str]] = field(default_factory=list) - cfg_nodes: Dict[int, str] = field(default_factory=dict) - simhash: int = 0 - spectral_features: List[float] = field(default_factory=list) - tree_paths: List[int] = field(default_factory=list) - - def fingerprint(self) -> str: - if self._hash: - return self._hash - - top_paths = sorted(self.ast_paths, key=lambda x: -x[3])[:20] - top_ngrams = self.structural_ngrams.most_common(20) - - parts = [ - f"P:{','.join(f'{p[0]}>{p[1]}>{p[2]}' for p in top_paths)}", - f"N:{','.join(f'{n}:{c}' for n, c in top_ngrams)}", - f"C:{self.control_flow.get('signature', '')}", - f"S:{self.simhash:016x}", - ] - - combined = "|".join(parts) - self._hash = hashlib.md5(combined.encode()).hexdigest()[:16] - return self._hash - - -# ============================================================================= -# Language-specific node type mappings → Universal types -# ============================================================================= - -# Universal terminal types (content-bearing nodes) -TERMINAL_NORMALIZATION: Dict[str, str] = { - # Identifiers (all languages) - "identifier": "ID", "name": "ID", "property_identifier": "ID", - "field_identifier": "ID", "type_identifier": "TYPE_ID", - "shorthand_property_identifier": "ID", "statement_identifier": "ID", - - # Strings - "string": "STR", "string_literal": "STR", "interpreted_string_literal": "STR", - "raw_string_literal": "STR", "template_string": "STR", "string_content": "STR", - "char_literal": "STR", "rune_literal": "STR", - - # Numbers - "number": "NUM", "integer": "NUM", "integer_literal": "NUM", "int_literal": "NUM", - "float": "NUM", "float_literal": "NUM", "decimal_integer_literal": "NUM", - "hex_integer_literal": "NUM", "binary_integer_literal": "NUM", - "octal_integer_literal": "NUM", "decimal_floating_point_literal": "NUM", - - # Booleans - "true": "BOOL", "false": "BOOL", "boolean": "BOOL", - - # Null/None/Nil - "none": "NIL", "null": "NIL", "nil": "NIL", "null_literal": "NIL", - - # Keywords become themselves (normalized case) - "return": "RETURN", "break": "BREAK", "continue": "CONTINUE", - "yield": "YIELD", "await": "AWAIT", "async": "ASYNC", -} - -# Universal control flow types - maps language-specific AST nodes to universal concepts -CONTROL_FLOW_NORMALIZATION: Dict[str, Dict[str, str]] = { - # Python - "python": { - "for_statement": "LOOP_FOR", "while_statement": "LOOP_WHILE", - "if_statement": "BRANCH_IF", "elif_clause": "BRANCH_ELIF", "else_clause": "BRANCH_ELSE", - "try_statement": "TRY", "except_clause": "CATCH", "finally_clause": "FINALLY", - "with_statement": "RESOURCE_GUARD", "match_statement": "MATCH", - "function_definition": "FUNC_DEF", "async_function_definition": "FUNC_DEF", - "class_definition": "CLASS_DEF", "lambda": "LAMBDA", - "list_comprehension": "COMPREHENSION", "dict_comprehension": "COMPREHENSION", - "generator_expression": "GENERATOR", - }, - # JavaScript / TypeScript - "javascript": { - "for_statement": "LOOP_FOR", "for_in_statement": "LOOP_FOR", - "while_statement": "LOOP_WHILE", "do_statement": "LOOP_DO", - "if_statement": "BRANCH_IF", "else_clause": "BRANCH_ELSE", - "switch_statement": "MATCH", "case": "MATCH_CASE", - "try_statement": "TRY", "catch_clause": "CATCH", "finally_clause": "FINALLY", - "function_declaration": "FUNC_DEF", "function_expression": "FUNC_EXPR", - "arrow_function": "LAMBDA", "class_declaration": "CLASS_DEF", - "method_definition": "METHOD_DEF", - }, - # TypeScript inherits all JavaScript mappings (handled in _get_cf_map) - "typescript": None, # Marker for inheritance - will fall back to javascript - # Go - "go": { - "for_statement": "LOOP_FOR", "range_clause": "LOOP_RANGE", - "if_statement": "BRANCH_IF", "else_clause": "BRANCH_ELSE", - "switch_statement": "MATCH", "type_switch_statement": "MATCH", - "select_statement": "SELECT", "case_clause": "MATCH_CASE", - "defer_statement": "DEFER", "go_statement": "GOROUTINE", - "function_declaration": "FUNC_DEF", "method_declaration": "METHOD_DEF", - "func_literal": "LAMBDA", "type_declaration": "TYPE_DEF", - }, - # Rust - "rust": { - "for_expression": "LOOP_FOR", "while_expression": "LOOP_WHILE", - "loop_expression": "LOOP_INFINITE", - "if_expression": "BRANCH_IF", "else_clause": "BRANCH_ELSE", - "match_expression": "MATCH", "match_arm": "MATCH_CASE", - "function_item": "FUNC_DEF", "closure_expression": "LAMBDA", - "impl_item": "IMPL", "trait_item": "TRAIT_DEF", "struct_item": "STRUCT_DEF", - "enum_item": "ENUM_DEF", "macro_invocation": "MACRO", - }, - # Java - "java": { - "for_statement": "LOOP_FOR", "enhanced_for_statement": "LOOP_FOR", - "while_statement": "LOOP_WHILE", "do_statement": "LOOP_DO", - "if_statement": "BRANCH_IF", "else": "BRANCH_ELSE", - "switch_expression": "MATCH", "switch_block_statement_group": "MATCH_CASE", - "try_statement": "TRY", "catch_clause": "CATCH", "finally_clause": "FINALLY", - "try_with_resources_statement": "RESOURCE_GUARD", - "method_declaration": "FUNC_DEF", "constructor_declaration": "CONSTRUCTOR", - "class_declaration": "CLASS_DEF", "interface_declaration": "INTERFACE_DEF", - "lambda_expression": "LAMBDA", - }, - # C / C++ - "c": { - "for_statement": "LOOP_FOR", "while_statement": "LOOP_WHILE", - "do_statement": "LOOP_DO", "if_statement": "BRANCH_IF", - "else_clause": "BRANCH_ELSE", "switch_statement": "MATCH", - "case_statement": "MATCH_CASE", "function_definition": "FUNC_DEF", - "struct_specifier": "STRUCT_DEF", "enum_specifier": "ENUM_DEF", - }, - "cpp": { - "for_statement": "LOOP_FOR", "for_range_loop": "LOOP_FOR", - "while_statement": "LOOP_WHILE", "do_statement": "LOOP_DO", - "if_statement": "BRANCH_IF", "else_clause": "BRANCH_ELSE", - "switch_statement": "MATCH", "try_statement": "TRY", - "catch_clause": "CATCH", "function_definition": "FUNC_DEF", - "class_specifier": "CLASS_DEF", "lambda_expression": "LAMBDA", - }, - # Ruby - "ruby": { - "for": "LOOP_FOR", "while": "LOOP_WHILE", "until": "LOOP_WHILE", - "if": "BRANCH_IF", "unless": "BRANCH_IF", "elsif": "BRANCH_ELIF", - "else": "BRANCH_ELSE", "case": "MATCH", "when": "MATCH_CASE", - "begin": "TRY", "rescue": "CATCH", "ensure": "FINALLY", - "method": "FUNC_DEF", "singleton_method": "FUNC_DEF", - "class": "CLASS_DEF", "module": "MODULE_DEF", - "lambda": "LAMBDA", "block": "BLOCK", "do_block": "BLOCK", - }, - # C# - "c_sharp": { - "for_statement": "LOOP_FOR", "foreach_statement": "LOOP_FOR", - "while_statement": "LOOP_WHILE", "do_statement": "LOOP_DO", - "if_statement": "BRANCH_IF", "else_clause": "BRANCH_ELSE", - "switch_statement": "MATCH", "switch_expression": "MATCH", - "try_statement": "TRY", "catch_clause": "CATCH", "finally_clause": "FINALLY", - "using_statement": "RESOURCE_GUARD", - "method_declaration": "FUNC_DEF", "local_function_statement": "FUNC_DEF", - "class_declaration": "CLASS_DEF", "interface_declaration": "INTERFACE_DEF", - "lambda_expression": "LAMBDA", - }, - # Bash / Shell - "bash": { - "for_statement": "LOOP_FOR", "while_statement": "LOOP_WHILE", - "until_statement": "LOOP_WHILE", - "if_statement": "BRANCH_IF", "elif_clause": "BRANCH_ELIF", "else_clause": "BRANCH_ELSE", - "case_statement": "MATCH", "case_item": "MATCH_CASE", - "function_definition": "FUNC_DEF", - "subshell": "SUBSHELL", "command_substitution": "SUBSHELL", - "pipeline": "PIPELINE", - }, - # JSON (structural - no control flow, but patterns in structure) - "json": { - "object": "OBJECT", "array": "ARRAY", "pair": "KEY_VALUE", - "string": "STRING", "number": "NUMBER", "true": "BOOL", "false": "BOOL", "null": "NULL", - }, - # YAML (structural) - "yaml": { - "block_mapping": "OBJECT", "flow_mapping": "OBJECT", - "block_sequence": "ARRAY", "flow_sequence": "ARRAY", - "block_mapping_pair": "KEY_VALUE", "flow_pair": "KEY_VALUE", - "anchor": "ANCHOR", "alias": "ALIAS", - }, - # HTML - "html": { - "element": "ELEMENT", "self_closing_tag": "ELEMENT", - "start_tag": "TAG_OPEN", "end_tag": "TAG_CLOSE", - "script_element": "SCRIPT", "style_element": "STYLE", - "attribute": "ATTRIBUTE", - "doctype": "DOCTYPE", "comment": "COMMENT", - }, - # CSS - "css": { - "rule_set": "RULE", "media_statement": "MEDIA_QUERY", - "keyframes_statement": "KEYFRAMES", "supports_statement": "SUPPORTS", - "declaration": "DECLARATION", "selector": "SELECTOR", - "class_selector": "CLASS_SEL", "id_selector": "ID_SEL", - "pseudo_class_selector": "PSEUDO_CLASS", "pseudo_element_selector": "PSEUDO_ELEM", - }, - # Markdown (structural patterns in documents) - "markdown": { - "atx_heading": "HEADING", "setext_heading": "HEADING", - "paragraph": "PARAGRAPH", - "fenced_code_block": "CODE_BLOCK", "indented_code_block": "CODE_BLOCK", - "list": "LIST", "list_item": "LIST_ITEM", - "block_quote": "QUOTE", "thematic_break": "HR", - "link": "LINK", "image": "IMAGE", - }, -} - -# Build a unified lookup by flattening all language mappings -# Skip None values (used as inheritance markers, e.g., typescript -> javascript) -_UNIFIED_CF_MAP: Dict[str, str] = {} -for lang_map in CONTROL_FLOW_NORMALIZATION.values(): - if lang_map is not None: - _UNIFIED_CF_MAP.update(lang_map) - -# ============================================================================= -# Dynamic fallback patterns - work for ANY language even without explicit mapping -# ============================================================================= -# These patterns match common AST node naming conventions across tree-sitter grammars. -# If a language isn't explicitly mapped, we try to match by common patterns. - -_UNIVERSAL_PATTERNS: Dict[str, str] = { - # Loops - most grammars use these patterns - "for_statement": "LOOP_FOR", "for_expression": "LOOP_FOR", - "for_loop": "LOOP_FOR", "for_in": "LOOP_FOR", "for_of": "LOOP_FOR", - "while_statement": "LOOP_WHILE", "while_expression": "LOOP_WHILE", - "while_loop": "LOOP_WHILE", - "do_statement": "LOOP_DO", "do_while": "LOOP_DO", - "loop": "LOOP_INFINITE", "loop_expression": "LOOP_INFINITE", - - # Conditionals - "if_statement": "BRANCH_IF", "if_expression": "BRANCH_IF", - "else_clause": "BRANCH_ELSE", "else": "BRANCH_ELSE", - "elif_clause": "BRANCH_ELIF", "elsif": "BRANCH_ELIF", "else_if": "BRANCH_ELIF", - "conditional_expression": "BRANCH_TERNARY", "ternary": "BRANCH_TERNARY", - - # Switch/Match - "switch_statement": "MATCH", "switch_expression": "MATCH", - "match_expression": "MATCH", "case_statement": "MATCH", - "case": "MATCH_CASE", "switch_case": "MATCH_CASE", "match_arm": "MATCH_CASE", - "when": "MATCH_CASE", "case_clause": "MATCH_CASE", - - # Error handling - "try_statement": "TRY", "try_expression": "TRY", "try": "TRY", - "catch_clause": "CATCH", "catch": "CATCH", "except": "CATCH", - "except_clause": "CATCH", "rescue": "CATCH", - "finally_clause": "FINALLY", "finally": "FINALLY", "ensure": "FINALLY", - "throw_statement": "THROW", "raise_statement": "THROW", "throw": "THROW", - - # Functions - "function_definition": "FUNC_DEF", "function_declaration": "FUNC_DEF", - "function_item": "FUNC_DEF", "method_definition": "METHOD_DEF", - "method_declaration": "METHOD_DEF", "method": "FUNC_DEF", - "lambda": "LAMBDA", "lambda_expression": "LAMBDA", "arrow_function": "LAMBDA", - "closure": "LAMBDA", "closure_expression": "LAMBDA", - "func_literal": "LAMBDA", "anonymous_function": "LAMBDA", - - # Classes/Types - "class_definition": "CLASS_DEF", "class_declaration": "CLASS_DEF", - "class": "CLASS_DEF", "class_specifier": "CLASS_DEF", - "struct_definition": "STRUCT_DEF", "struct_specifier": "STRUCT_DEF", - "struct": "STRUCT_DEF", "struct_item": "STRUCT_DEF", - "interface_declaration": "INTERFACE_DEF", "interface": "INTERFACE_DEF", - "trait_item": "TRAIT_DEF", "trait": "TRAIT_DEF", - "enum_declaration": "ENUM_DEF", "enum_item": "ENUM_DEF", "enum": "ENUM_DEF", - "type_declaration": "TYPE_DEF", "type_alias": "TYPE_DEF", - - # Resource management - "with_statement": "RESOURCE_GUARD", "using_statement": "RESOURCE_GUARD", - "defer_statement": "DEFER", "defer": "DEFER", -} - - -class PatternExtractor: - """Extract structural patterns from code using Tree-sitter - multi-language.""" - - MAX_PATH_LENGTH = 8 - MAX_PATHS = 200 - NGRAM_SIZES = [2, 3, 4] - - def __init__(self): - self._parsers: Dict[str, Any] = {} - - def _get_parser(self, language: str): - """Get or create Tree-sitter parser for language.""" - # Normalize language name - lang_key = self._normalize_language(language) - - if lang_key in self._parsers: - return self._parsers[lang_key] - - try: - from scripts.ingest.tree_sitter import _ts_parser - parser = _ts_parser(lang_key) - if parser: - self._parsers[lang_key] = parser - return parser - except ImportError: - pass - return None - - def _normalize_language(self, language: str) -> str: - """ - Normalize language name to Tree-sitter key. - - Handles all common aliases, extensions, and variations. - Returns lowercase normalized key that matches our normalization maps. - """ - lang_map = { - # Python - "py": "python", "python3": "python", "python2": "python", - # JavaScript - "js": "javascript", "jsx": "javascript", "mjs": "javascript", "cjs": "javascript", - "node": "javascript", "nodejs": "javascript", - # TypeScript - "ts": "typescript", "tsx": "typescript", "mts": "typescript", - # Go - "golang": "go", - # Rust - "rs": "rust", - # Ruby - "rb": "ruby", "rake": "ruby", "gemspec": "ruby", - # C# - "cs": "c_sharp", "csharp": "c_sharp", - # C++ - "c++": "cpp", "cc": "cpp", "cxx": "cpp", "hpp": "cpp", "hxx": "cpp", - # Java - "jav": "java", - # Shell/Bash - "bash": "bash", "sh": "bash", "zsh": "bash", "shell": "bash", - "ksh": "bash", "fish": "bash", - # Config languages - "yml": "yaml", - "htm": "html", "xhtml": "html", - "scss": "css", "sass": "css", "less": "css", - "md": "markdown", "mdx": "markdown", - # Other common - "kt": "kotlin", "kts": "kotlin", - "swift": "swift", - "php": "php", - "scala": "scala", "sc": "scala", - "pl": "perl", "pm": "perl", - "lua": "lua", - "r": "r", - "jl": "julia", - "ex": "elixir", "exs": "elixir", - "erl": "erlang", "hrl": "erlang", - "hs": "haskell", "lhs": "haskell", - "clj": "clojure", "cljs": "clojure", "cljc": "clojure", - "ml": "ocaml", "mli": "ocaml", - "fs": "fsharp", "fsi": "fsharp", "fsx": "fsharp", - "ps1": "powershell", "psm1": "powershell", - "vue": "vue", - "svelte": "svelte", - } - return lang_map.get(language.lower(), language.lower()) - - def _get_cf_map(self, language: str) -> Dict[str, str]: - """ - Get control flow normalization map for a language. - - Falls back gracefully: - 1. Exact language match - 2. TypeScript -> JavaScript inheritance - 3. Empty dict (will use _UNIVERSAL_PATTERNS fallback) - """ - lang_key = self._normalize_language(language) - - # TypeScript inherits from JavaScript (None marker or missing key) - cf_map = CONTROL_FLOW_NORMALIZATION.get(lang_key) - if cf_map is None and lang_key == "typescript": - cf_map = CONTROL_FLOW_NORMALIZATION.get("javascript", {}) - - # Return the map or empty dict - _UNIVERSAL_PATTERNS will catch unknowns - return cf_map if cf_map is not None else {} - - def extract(self, code: str, language: str = "python") -> PatternSignature: - """Extract pattern signature from code - works for any supported language.""" - lang_key = self._normalize_language(language) - sig = PatternSignature(language=lang_key) - - parser = self._get_parser(lang_key) - if not parser: - # Fallback to regex-based extraction - return self._extract_regex_fallback(code, lang_key) - - try: - tree = parser.parse(code.encode("utf-8")) - if tree is None: - code_preview = code[:100].replace('\n', '\\n') if code else '' - logger.debug( - f"Tree-sitter returned None for {lang_key}, using regex fallback. " - f"Code preview: {code_preview!r}" - ) - return self._extract_regex_fallback(code, lang_key) - root = tree.root_node - except Exception as e: - code_preview = code[:100].replace('\n', '\\n') if code else '' - logger.debug( - f"Tree-sitter parse failed for {lang_key}: {e}, using regex fallback. " - f"Code preview: {code_preview!r}" - ) - return self._extract_regex_fallback(code, lang_key) - - cf_map = self._get_cf_map(lang_key) - sig.ast_paths = self._extract_ast_paths(root, code, lang_key) - sig.structural_ngrams = self._extract_ngrams(root, lang_key) - sig.control_flow = self._extract_control_flow(root, lang_key) - - sig.wl_labels = self._extract_wl_labels(root, cf_map, k=3) - sig.cfg_nodes, sig.cfg_edges = self._extract_cfg(root, cf_map) - sig.simhash = self._compute_simhash(sig.structural_ngrams, sig.wl_labels) - sig.spectral_features = self._extract_spectral_features(root, cf_map) - sig.tree_paths = self._extract_tree_paths(root, cf_map) - - return sig - - def _extract_ast_paths(self, root, code: str, language: str) -> List[Tuple[str, str, str, int]]: - """Extract code2vec-style AST paths between terminals - normalized across languages.""" - terminals = [] - self._collect_terminals(root, terminals, depth=0) - - paths: Counter = Counter() - cf_map = self._get_cf_map(language) - - # Generate paths between pairs of terminals - for i, (node_i, depth_i) in enumerate(terminals): - for j, (node_j, depth_j) in enumerate(terminals[i+1:min(i+10, len(terminals))]): - path = self._compute_path(node_i, node_j, cf_map) - if path and len(path) <= self.MAX_PATH_LENGTH: - start_type = self._normalize_node_type(node_i.type, cf_map) - end_type = self._normalize_node_type(node_j.type, cf_map) - path_str = "^".join(path) - paths[(start_type, path_str, end_type)] += 1 - - return [(s, p, e, c) for (s, p, e), c in paths.most_common(self.MAX_PATHS)] - - def _normalize_node_type(self, node_type: str, cf_map: Dict[str, str]) -> str: - """ - Normalize node type to universal type - enables cross-language matching. - - Priority order: - 1. Terminal normalization (literals, identifiers) - 2. Language-specific control flow map - 3. Unified map (all languages combined) - 4. Universal patterns (dynamic fallback for any language) - 5. Lowercase as-is (preserves unknown types consistently) - """ - # First check terminal normalization - if node_type in TERMINAL_NORMALIZATION: - return TERMINAL_NORMALIZATION[node_type] - # Then check control flow normalization (language-specific) - if node_type in cf_map: - return cf_map[node_type] - # Then check unified map - if node_type in _UNIFIED_CF_MAP: - return _UNIFIED_CF_MAP[node_type] - # Dynamic fallback: try universal patterns (works for ANY language) - if node_type in _UNIVERSAL_PATTERNS: - return _UNIVERSAL_PATTERNS[node_type] - # Keep as-is but lowercase for consistency - return node_type.lower() - - def _collect_terminals(self, node, terminals: List, depth: int): - """Collect terminal (leaf) nodes.""" - if node.child_count == 0: - terminals.append((node, depth)) - else: - for child in node.children: - self._collect_terminals(child, terminals, depth + 1) - - def _compute_path(self, node_a, node_b, cf_map: Dict[str, str]) -> Optional[List[str]]: - """Compute AST path between two nodes (up to LCA, then down) - normalized.""" - # Get ancestors of both nodes - ancestors_a = [] - n = node_a - while n is not None: - ancestors_a.append(n) - n = n.parent - - ancestors_b = [] - n = node_b - while n is not None: - ancestors_b.append(n) - n = n.parent - - # Find LCA - set_a = set(id(n) for n in ancestors_a) - lca = None - lca_idx_b = 0 - for i, n in enumerate(ancestors_b): - if id(n) in set_a: - lca = n - lca_idx_b = i - break - - if lca is None: - return None - - lca_idx_a = next(i for i, n in enumerate(ancestors_a) if id(n) == id(lca)) - - # Build path: up from a to LCA, then down to b - # Node types are normalized for cross-language matching - path = [] - for n in ancestors_a[:lca_idx_a]: - path.append(f"↑{self._normalize_node_type(n.type, cf_map)}") - path.append(f"○{self._normalize_node_type(lca.type, cf_map)}") - for n in reversed(ancestors_b[:lca_idx_b]): - path.append(f"↓{self._normalize_node_type(n.type, cf_map)}") - - return path - - def _extract_ngrams(self, root, language: str) -> Counter: - """Extract structural n-grams from AST traversal - normalized.""" - cf_map = self._get_cf_map(language) - - # Pre-order traversal of normalized node types - node_types = [] - self._collect_node_types(root, node_types, cf_map) - - ngrams = Counter() - for n in self.NGRAM_SIZES: - for i in range(len(node_types) - n + 1): - gram = tuple(node_types[i:i+n]) - ngrams[gram] += 1 - - return ngrams - - def _collect_node_types(self, node, types: List[str], cf_map: Dict[str, str]): - """Collect normalized node types in pre-order.""" - types.append(self._normalize_node_type(node.type, cf_map)) - for child in node.children: - self._collect_node_types(child, types, cf_map) - - def _extract_control_flow(self, root, language: str) -> Dict[str, Any]: - """Extract control flow structure features - language-agnostic.""" - cf_map = self._get_cf_map(language) - - cf = { - "max_loop_depth": 0, - "loop_count": 0, - "branch_count": 0, - "try_count": 0, - "has_finally": False, - "has_catch": False, # Renamed from has_except for universality - "has_except": False, # Keep for backward compat - "has_resource_guard": False, # with/using/defer/try-with-resources - "func_count": 0, - "class_count": 0, - "lambda_count": 0, - "match_count": 0, # switch/match expressions - # NEW: Nesting pattern detection - "loop_types": [], # Track specific loop types seen - "nesting_patterns": [], # Track nesting sequences - "try_in_loop": False, # Try block nested in loop (retry pattern!) - "loop_in_try": False, # Loop nested in try block - "branch_in_loop": False, # Branch nested in loop (filter pattern) - "max_nesting_depth": 0, # Overall nesting depth - } - - # Track context during traversal - context = { - "loop_depth": 0, - "try_depth": 0, - "branch_depth": 0, - "in_loop": False, - "in_try": False, - "nesting_sequence": [], # Track sequence of control flow - } - - self._analyze_control_flow_v2(root, cf, cf_map, context) - - # Store the nesting sequence for search reranking - # This is used by _fallback_pattern_search to compute Jaccard similarity - cf["normalized_sequence"] = context["nesting_sequence"] - - # Backward compat - cf["has_except"] = cf["has_catch"] - - # Generate compact signature - universal across all languages - # V2: More detailed signature - cf["signature"] = ( - f"L{cf['max_loop_depth']}_{cf['loop_count']}_" - f"B{cf['branch_count']}_T{cf['try_count']}_" - f"M{cf['match_count']}_" - f"{'F' if cf['has_finally'] else '_'}" - f"{'C' if cf['has_catch'] else '_'}" - f"{'R' if cf['has_resource_guard'] else '_'}" - f"{'TL' if cf['try_in_loop'] else ''}" # Retry pattern marker! - f"{'LT' if cf['loop_in_try'] else ''}" - f"{'BL' if cf['branch_in_loop'] else ''}" - ) - - return cf - - def _analyze_control_flow_v2(self, node, cf: Dict, cf_map: Dict[str, str], ctx: Dict): - """Recursively analyze control flow with nesting detection.""" - normalized = self._normalize_node_type(node.type, cf_map) - - # Track previous context - was_in_loop = ctx["in_loop"] - was_in_try = ctx["in_try"] - entered_loop = False - entered_try = False - entered_branch = False - - # Loop detection (LOOP_FOR, LOOP_WHILE, LOOP_DO, LOOP_INFINITE, LOOP_RANGE) - if normalized.startswith("LOOP_"): - cf["loop_count"] += 1 - cf["loop_types"].append(normalized) - ctx["loop_depth"] += 1 - ctx["in_loop"] = True - cf["max_loop_depth"] = max(cf["max_loop_depth"], ctx["loop_depth"]) - entered_loop = True - ctx["nesting_sequence"].append("LOOP") - - # Detect loop in try - if was_in_try: - cf["loop_in_try"] = True - if "loop_in_try" not in cf["nesting_patterns"]: - cf["nesting_patterns"].append("loop_in_try") - - # Branch detection - elif normalized == "BRANCH_IF": - cf["branch_count"] += 1 - ctx["branch_depth"] += 1 - entered_branch = True - ctx["nesting_sequence"].append("BRANCH") - - # Detect branch in loop (filter pattern) - if was_in_loop: - cf["branch_in_loop"] = True - if "branch_in_loop" not in cf["nesting_patterns"]: - cf["nesting_patterns"].append("branch_in_loop") - - # Try block detection - elif normalized == "TRY": - cf["try_count"] += 1 - ctx["try_depth"] += 1 - ctx["in_try"] = True - entered_try = True - ctx["nesting_sequence"].append("TRY") - - # Detect try in loop (RETRY PATTERN!) - if was_in_loop: - cf["try_in_loop"] = True - if "try_in_loop" not in cf["nesting_patterns"]: - cf["nesting_patterns"].append("try_in_loop") - - # Catch/except detection - elif normalized == "CATCH": - cf["has_catch"] = True - ctx["nesting_sequence"].append("CATCH") - - # Finally detection - elif normalized == "FINALLY": - cf["has_finally"] = True - ctx["nesting_sequence"].append("FINALLY") - - # Resource guard (with/using/defer) - elif normalized in ("RESOURCE_GUARD", "DEFER"): - cf["has_resource_guard"] = True - ctx["nesting_sequence"].append("RESOURCE") - - # Function definition - elif normalized in ("FUNC_DEF", "METHOD_DEF"): - cf["func_count"] += 1 - - # Class/struct definition - elif normalized in ("CLASS_DEF", "STRUCT_DEF", "INTERFACE_DEF"): - cf["class_count"] += 1 - - # Lambda/closure - elif normalized == "LAMBDA": - cf["lambda_count"] += 1 - - # Match/switch - elif normalized == "MATCH": - cf["match_count"] += 1 - ctx["nesting_sequence"].append("MATCH") - - # Track max nesting - total_depth = ctx["loop_depth"] + ctx["try_depth"] + ctx["branch_depth"] - cf["max_nesting_depth"] = max(cf["max_nesting_depth"], total_depth) - - # Recurse to children - for child in node.children: - self._analyze_control_flow_v2(child, cf, cf_map, ctx) - - # Restore context when leaving scope - if entered_loop: - ctx["loop_depth"] -= 1 - ctx["in_loop"] = ctx["loop_depth"] > 0 - if entered_try: - ctx["try_depth"] -= 1 - ctx["in_try"] = ctx["try_depth"] > 0 - if entered_branch: - ctx["branch_depth"] -= 1 - - def _analyze_control_flow(self, node, cf: Dict, cf_map: Dict[str, str], loop_depth: int): - """Legacy method - kept for compatibility, delegates to v2.""" - ctx = {"loop_depth": 0, "try_depth": 0, "branch_depth": 0, - "in_loop": False, "in_try": False, "nesting_sequence": []} - self._analyze_control_flow_v2(node, cf, cf_map, ctx) - - def _extract_regex_fallback(self, code: str, language: str) -> PatternSignature: - """Fallback pattern extraction using regex - multi-language aware.""" - sig = PatternSignature(language=language) - - # Language-aware regex patterns for common constructs - loop_patterns = { - "python": r'\b(for|while)\b', - "javascript": r'\b(for|while|do)\b', - "go": r'\bfor\b', - "rust": r'\b(for|while|loop)\b', - "java": r'\b(for|while|do)\b', - "c": r'\b(for|while|do)\b', - "cpp": r'\b(for|while|do)\b', - "ruby": r'\b(for|while|until|each|loop)\b', - } - - try_patterns = { - "python": (r'\btry\b', r'\bexcept\b', r'\bfinally\b'), - "javascript": (r'\btry\b', r'\bcatch\b', r'\bfinally\b'), - "java": (r'\btry\b', r'\bcatch\b', r'\bfinally\b'), - "cpp": (r'\btry\b', r'\bcatch\b', None), - "ruby": (r'\bbegin\b', r'\brescue\b', r'\bensure\b'), - "go": (None, None, r'\bdefer\b'), # Go uses defer instead of try - } - - # Get patterns for this language, fall back to Python-style - loop_re = loop_patterns.get(language, r'\b(for|while)\b') - try_re, catch_re, finally_re = try_patterns.get(language, (r'\btry\b', r'\bcatch\b', r'\bfinally\b')) - - sig.control_flow = { - "loop_count": len(re.findall(loop_re, code)) if loop_re else 0, - "branch_count": len(re.findall(r'\bif\b', code)), - "try_count": len(re.findall(try_re, code)) if try_re else 0, - "has_catch": bool(re.search(catch_re, code)) if catch_re else False, - "has_finally": bool(re.search(finally_re, code)) if finally_re else False, - "has_resource_guard": bool(re.search(r'\b(with|using|defer)\b', code)), - "max_loop_depth": 1 if loop_re and re.search(loop_re, code) else 0, - "signature": "FALLBACK", - } - - tokens = re.findall(r'\b\w+\b', code) - for n in self.NGRAM_SIZES: - for i in range(len(tokens) - n + 1): - sig.structural_ngrams[tuple(tokens[i:i+n])] += 1 - - return sig - - def _extract_wl_labels(self, root, cf_map: Dict[str, str], k: int = 3) -> Dict[int, List[str]]: - """Weisfeiler-Lehman graph kernel.""" - nodes = [] - edges = [] - self._build_ast_graph(root, nodes, edges, cf_map) - - if not nodes: - return {} - - labels = {i: [self._normalize_node_type(n.type, cf_map)] for i, n in enumerate(nodes)} - adj = {i: [] for i in range(len(nodes))} - for src, dst in edges: - adj[src].append(dst) - adj[dst].append(src) - - for iteration in range(k): - new_labels = {} - for node_id in range(len(nodes)): - neighbor_labels = sorted(labels[n][-1] for n in adj[node_id]) - combined = labels[node_id][-1] + "|" + ",".join(neighbor_labels) - new_label = hashlib.md5(combined.encode()).hexdigest()[:8] - new_labels[node_id] = labels[node_id] + [new_label] - labels = new_labels - - return labels - - def _build_ast_graph(self, node, nodes: List, edges: List[Tuple[int, int]], cf_map: Dict): - node_id = len(nodes) - nodes.append(node) - for child in node.children: - child_id = len(nodes) - edges.append((node_id, child_id)) - self._build_ast_graph(child, nodes, edges, cf_map) - - def _extract_cfg(self, root, cf_map: Dict[str, str]) -> Tuple[Dict[int, str], List[Tuple[int, int, str]]]: - """Extract control flow graph.""" - cfg_nodes = {} - cfg_edges = [] - self._build_cfg(root, cf_map, cfg_nodes, cfg_edges, entry_id=0, exit_id=-1) - return cfg_nodes, cfg_edges - - def _build_cfg(self, node, cf_map: Dict, nodes: Dict, edges: List, entry_id: int, exit_id: int) -> int: - normalized = self._normalize_node_type(node.type, cf_map) - node_id = len(nodes) - nodes[node_id] = normalized - - # For control flow nodes (branches, loops, try), we use specialized edges - # and skip the generic sequential edge to avoid double-counting - is_control_flow = ( - normalized == "BRANCH_IF" or - normalized.startswith("LOOP_") or - normalized == "TRY" - ) - - # Add sequential edge only for non-control-flow nodes - if entry_id >= 0 and entry_id != node_id and not is_control_flow: - edges.append((entry_id, node_id, "sequential")) - - if normalized == "BRANCH_IF": - # For branches, use branch_true/branch_false edges instead of sequential - then_exit = node_id - else_exit = node_id - for child in node.children: - child_norm = self._normalize_node_type(child.type, cf_map) - if child_norm in ("BRANCH_ELSE", "BRANCH_ELIF"): - # Pass -1 as entry_id to skip sequential edge in child - else_exit = self._build_cfg(child, cf_map, nodes, edges, -1, exit_id) - if else_exit != node_id: - edges.append((node_id, else_exit, "branch_false")) - elif child.child_count > 0: - # Pass -1 as entry_id to skip sequential edge in child - then_exit = self._build_cfg(child, cf_map, nodes, edges, -1, exit_id) - if then_exit != node_id: - edges.append((node_id, then_exit, "branch_true")) - return max(then_exit, else_exit) - - elif normalized.startswith("LOOP_"): - # For loops, the entry flows into the loop body, then back - loop_body_exit = node_id - for child in node.children: - # Pass -1 as entry_id to skip sequential edge; loop structure is explicit - loop_body_exit = self._build_cfg(child, cf_map, nodes, edges, -1, exit_id) - if loop_body_exit != node_id: - edges.append((loop_body_exit, node_id, "loop_back")) - # Add explicit entry edge to first child - if node.child_count > 0: - first_child_id = node_id + 1 # Next node after this one - if first_child_id in nodes: - edges.append((node_id, first_child_id, "loop_entry")) - return node_id - - elif normalized == "TRY": - # For try, normal flow goes through try body; exception edge goes to catch - try_exit = node_id - for child in node.children: - child_norm = self._normalize_node_type(child.type, cf_map) - if child_norm == "CATCH": - # Exception edge, not sequential - catch_exit = self._build_cfg(child, cf_map, nodes, edges, -1, exit_id) - edges.append((node_id, catch_exit, "exception")) - else: - # Normal flow through try body - try_exit = self._build_cfg(child, cf_map, nodes, edges, try_exit, exit_id) - return try_exit - - else: - last_exit = node_id - for child in node.children: - last_exit = self._build_cfg(child, cf_map, nodes, edges, last_exit, exit_id) - return last_exit - - def _compute_simhash(self, ngrams: Counter, wl_labels: Dict[int, List[str]]) -> int: - """64-bit SimHash for LSH.""" - v = [0] * 64 - - for ngram, count in ngrams.items(): - h = int(hashlib.md5(str(ngram).encode()).hexdigest()[:16], 16) - for i in range(64): - if (h >> i) & 1: - v[i] += count - else: - v[i] -= count - - for node_id, labels in wl_labels.items(): - for label in labels: - h = int(hashlib.md5(label.encode()).hexdigest()[:16], 16) - for i in range(64): - if (h >> i) & 1: - v[i] += 1 - else: - v[i] -= 1 - - result = 0 - for i in range(64): - if v[i] > 0: - result |= (1 << i) - return result - - def _extract_spectral_features(self, root, cf_map: Dict[str, str], k: int = 8) -> List[float]: - """Approximate eigenvalues of normalized graph Laplacian via power iteration.""" - nodes = [] - edges = [] - self._build_ast_graph(root, nodes, edges, cf_map) - - n = len(nodes) - if n < 2: - return [0.0] * k - - degree = [0] * n - adj = {i: [] for i in range(n)} - for src, dst in edges: - adj[src].append(dst) - adj[dst].append(src) - degree[src] += 1 - degree[dst] += 1 - - def laplacian_multiply(x: List[float]) -> List[float]: - result = [0.0] * n - for i in range(n): - if degree[i] > 0: - result[i] = x[i] - for j in adj[i]: - if degree[j] > 0: - result[i] -= x[j] / (degree[i] * degree[j]) ** 0.5 - return result - - eigenvalues = [] - for _ in range(min(k, n - 1)): - v = [1.0 / n**0.5] * n - for _ in range(20): - v = laplacian_multiply(v) - norm = sum(x*x for x in v) ** 0.5 - if norm > 1e-10: - v = [x / norm for x in v] - eigenvalue = sum(a * b for a, b in zip(v, laplacian_multiply(v))) - eigenvalues.append(round(eigenvalue, 4)) - - while len(eigenvalues) < k: - eigenvalues.append(0.0) - - return eigenvalues[:k] - - def _extract_tree_paths(self, root, cf_map: Dict[str, str], max_paths: int = 32) -> List[int]: - """Root-to-leaf path hashes for tree edit distance approximation.""" - paths = [] - self._collect_root_to_leaf_paths(root, [], cf_map, paths) - - path_hashes = [] - for path in paths[:max_paths]: - path_str = ">".join(path) - h = int(hashlib.md5(path_str.encode()).hexdigest()[:8], 16) - path_hashes.append(h) - - return sorted(path_hashes) - - def _collect_root_to_leaf_paths(self, node, current_path: List[str], cf_map: Dict, paths: List): - normalized = self._normalize_node_type(node.type, cf_map) - current_path = current_path + [normalized] - - if not node.children: - paths.append(current_path) - else: - for child in node.children: - self._collect_root_to_leaf_paths(child, current_path, cf_map, paths) - - # ========================================================================= - # Line-level pattern matching for highlighting - # ========================================================================= - - # Map control_flow dict keys → normalized AST types to look for - _CF_KEY_TO_TYPES: Dict[str, List[str]] = { - "has_for": ["LOOP_FOR", "LOOP_RANGE"], - "has_while": ["LOOP_WHILE"], - "loop_count": ["LOOP", "LOOP_FOR", "LOOP_WHILE", "LOOP_RANGE", "LOOP_DO"], - "has_if": ["BRANCH_IF"], - "branch_count": ["BRANCH_IF", "BRANCH_ELSE", "BRANCH_ELIF"], - "has_try": ["TRY"], - "try_count": ["TRY"], - "has_catch": ["CATCH", "EXCEPT"], - "has_except": ["CATCH", "EXCEPT"], - "has_finally": ["FINALLY"], - "has_with": ["WITH"], - "has_resource_guard": ["WITH", "DEFER", "USING"], - "has_return": ["RETURN"], - "has_raise": ["RAISE", "THROW"], - "has_throw": ["RAISE", "THROW"], - "has_yield": ["YIELD"], - "has_async": ["ASYNC"], - "has_await": ["AWAIT"], - "has_defer": ["DEFER"], - "has_goroutine": ["GOROUTINE"], - "has_select": ["SELECT"], - "has_recover": ["RECOVER"], - "match_count": ["MATCH", "MATCH_CASE"], - } - - def find_matching_lines( - self, - code: str, - language: str, - query_cf: Dict[str, Any], - line_offset: int = 0, - ) -> List[int]: - """ - Find lines in code that contain control flow patterns from query. - - Uses Tree-sitter AST analysis - NOT keyword matching. - - Args: - code: Source code to analyze - language: Programming language - query_cf: Control flow dict from query's PatternSignature - line_offset: Added to line numbers (for snippet within larger file) - - Returns: - Sorted list of 1-indexed line numbers containing matching patterns - """ - if not code or not query_cf: - return [] - - # Build set of normalized types to look for based on query_cf - target_types: Set[str] = set() - for cf_key, norm_types in self._CF_KEY_TO_TYPES.items(): - val = query_cf.get(cf_key) - # Include if truthy (True, or count > 0) - if val: - target_types.update(norm_types) - - if not target_types: - return [] - - # Parse code - lang_key = self._normalize_language(language) - parser = self._get_parser(lang_key) - if not parser: - return [] - - try: - tree = parser.parse(code.encode('utf-8')) - if not tree or not tree.root_node: - return [] - except Exception: - return [] - - # Walk AST and collect line numbers of matching nodes - cf_map = self._get_cf_map(lang_key) - matched_lines: Set[int] = set() - - def walk(node): - normalized = self._normalize_node_type(node.type, cf_map) - if normalized in target_types: - # Tree-sitter lines are 0-indexed, we want 1-indexed - matched_lines.add(node.start_point[0] + 1 + line_offset) - for child in node.children: - walk(child) - - walk(tree.root_node) - - return sorted(matched_lines) diff --git a/scripts/pattern_detection/prune.py b/scripts/pattern_detection/prune.py deleted file mode 100644 index 3ff15fed..00000000 --- a/scripts/pattern_detection/prune.py +++ /dev/null @@ -1,526 +0,0 @@ -""" -AROMA-style Prune Algorithm - Greedy maximal subtree extraction. - -From the AROMA paper (Section 3.3.2): -The Prune algorithm takes two code snippets m1 (query) and m2 (candidate), -and finds a MAXIMAL SUBTREE of m2 that is most similar to m1. - -Key insight: By identifying which leaf nodes of m2's parse tree should be -retained, we get a maximal subtree by keeping all nodes on paths from root -to those leaves. - -Algorithm (greedy): -1. Start with empty set R (leaf nodes to retain) -2. Start with empty feature set F -3. Iteratively find leaf node n from m2 that maximizes: - SimScore(F(m1), F ∪ F(n)) -4. If adding n increases similarity, add to R, update F -5. Stop when no improvement possible -6. Return subtree formed by R + all ancestors - -SimScore(m1, m2) = |F(m1) ∩ F(m2)| (feature overlap cardinality) - -Uses: -- Reranking: Prune each search result w.r.t. query, rank by pruned similarity -- Intersection: Extract common structure from query + cluster members -""" - -from dataclasses import dataclass, field -from typing import List, Dict, Set, Tuple, Optional, Any, Iterator -from collections import Counter -import hashlib - -from .extractor import PatternSignature, PatternExtractor - - -@dataclass -class PrunedResult: - """Result of pruning m2 with respect to m1.""" - - # Retained leaf node IDs - retained_leaves: Set[int] = field(default_factory=set) - - # Retained internal node IDs (ancestors of leaves) - retained_nodes: Set[int] = field(default_factory=set) - - # Features of the pruned subtree - pruned_features: Set[str] = field(default_factory=set) - - # Similarity score after pruning - similarity_score: float = 0.0 - - # Original m2 feature count (before pruning) - original_feature_count: int = 0 - - # Pruned feature count - pruned_feature_count: int = 0 - - # Ratio of m2 retained - retention_ratio: float = 0.0 - - -class AromaPruner: - """ - AROMA-style greedy pruning for maximal similar subtree extraction. - - This implements the core AROMA algorithm that enables: - 1. Better reranking of search results - 2. Precise intersection for pattern discovery - """ - - def __init__(self, extractor: Optional[PatternExtractor] = None): - self.extractor = extractor or PatternExtractor() - - def prune( - self, - query_sig: PatternSignature, - candidate_sig: PatternSignature, - query_tree: Optional[Any] = None, - candidate_tree: Optional[Any] = None, - ) -> PrunedResult: - """ - Find maximal subtree of candidate that is most similar to query. - - Uses feature-based approximation when AST not available, - or full tree-based algorithm when AST is provided. - """ - if candidate_tree is not None and query_tree is not None: - return self._prune_tree_based(query_sig, candidate_sig, query_tree, candidate_tree) - else: - return self._prune_feature_based(query_sig, candidate_sig) - - def _prune_feature_based( - self, - query_sig: PatternSignature, - candidate_sig: PatternSignature, - ) -> PrunedResult: - """ - Feature-based pruning approximation. - - When we don't have AST access, we use features as proxy for leaves: - - Each structural n-gram acts like a "leaf feature" - - Greedily add features that increase overlap with query - """ - result = PrunedResult() - - # Extract feature sets - query_features = self._signature_to_features(query_sig) - candidate_features = self._signature_to_features(candidate_sig) - - result.original_feature_count = len(candidate_features) - - if not query_features or not candidate_features: - return result - - # Greedy selection: add candidate features that are in query - # This is the feature-based approximation of tree pruning - retained: Set[str] = set() - candidate_list = list(candidate_features) - - # Sort by whether feature is in query (prefer matches) - candidate_list.sort(key=lambda f: (f in query_features, f), reverse=True) - - current_score = 0 - for feature in candidate_list: - # Would adding this feature improve overlap? - new_retained = retained | {feature} - new_score = len(new_retained & query_features) - - if new_score > current_score: - retained.add(feature) - current_score = new_score - # Stop if we've matched all query features - if current_score == len(query_features): - break - - result.retained_leaves = set(range(len(retained))) # Placeholder IDs - result.pruned_features = retained - result.pruned_feature_count = len(retained) - result.similarity_score = current_score / max(len(query_features), 1) - result.retention_ratio = len(retained) / max(len(candidate_features), 1) - - return result - - def _prune_tree_based( - self, - query_sig: PatternSignature, - candidate_sig: PatternSignature, - query_tree: Any, - candidate_tree: Any, - ) -> PrunedResult: - """ - Full tree-based AROMA pruning algorithm. - - Algorithm: - 1. Collect all leaf nodes of candidate tree - 2. For each leaf, compute its feature contribution - 3. Greedily select leaves that maximize SimScore(query, retained) - 4. Compute ancestors of retained leaves for full subtree - """ - result = PrunedResult() - - # Get query features - query_features = self._signature_to_features(query_sig) - if not query_features: - return result - - # Collect leaf nodes with their features - leaves: List[Tuple[Any, int, Set[str]]] = [] # (node, node_id, features) - self._collect_leaves_with_features(candidate_tree, 0, leaves) - - result.original_feature_count = sum(len(f) for _, _, f in leaves) - - if not leaves: - return result - - # Greedy selection - retained_leaves: Set[int] = set() - current_features: Set[str] = set() - current_score = 0 - - # Track which leaves haven't been considered yet - available = set(range(len(leaves))) - - while available: - best_leaf_idx = -1 - best_new_score = current_score - best_new_features: Set[str] = current_features - - # Find leaf that maximizes score increase - for idx in available: - _, _, leaf_features = leaves[idx] - candidate_features = current_features | leaf_features - new_score = len(candidate_features & query_features) - - if new_score > best_new_score: - best_new_score = new_score - best_leaf_idx = idx - best_new_features = candidate_features - - # If no improvement, stop - if best_leaf_idx == -1 or best_new_score <= current_score: - break - - # Add best leaf - retained_leaves.add(best_leaf_idx) - current_features = best_new_features - current_score = best_new_score - available.remove(best_leaf_idx) - - # Compute ancestors of retained leaves - retained_nodes = self._compute_ancestors(candidate_tree, retained_leaves, leaves) - - result.retained_leaves = retained_leaves - result.retained_nodes = retained_nodes - result.pruned_features = current_features - result.pruned_feature_count = len(current_features) - result.similarity_score = current_score / max(len(query_features), 1) - result.retention_ratio = len(retained_nodes) / max(self._count_nodes(candidate_tree), 1) - - return result - - def _collect_leaves_with_features( - self, - node: Any, - node_id: int, - leaves: List[Tuple[Any, int, Set[str]]], - parent_path: Optional[List[str]] = None, - ) -> int: - """Collect leaf nodes with their feature contributions.""" - if parent_path is None: - parent_path = [] - - node_type = getattr(node, 'type', str(type(node).__name__)) - current_path = parent_path + [node_type] - - children = getattr(node, 'children', []) - - if not children: - # This is a leaf - compute its features - features: Set[str] = set() - - # Path feature (root to this leaf) - features.add(f"path:{'→'.join(current_path)}") - - # Parent feature - if len(current_path) >= 2: - features.add(f"parent:{current_path[-2]}→{current_path[-1]}") - - # Grandparent feature - if len(current_path) >= 3: - features.add(f"gp:{current_path[-3]}→{current_path[-2]}→{current_path[-1]}") - - # Node type feature - features.add(f"type:{node_type}") - - leaves.append((node, node_id, features)) - return node_id + 1 - - # Recurse to children - next_id = node_id + 1 - for child in children: - next_id = self._collect_leaves_with_features(child, next_id, leaves, current_path) - - return next_id - - def _compute_ancestors( - self, - root: Any, - retained_leaf_indices: Set[int], - leaves: List[Tuple[Any, int, Set[str]]], - ) -> Set[int]: - """Compute all ancestor node IDs for retained leaves.""" - # Build parent map - parent_map: Dict[int, int] = {} - self._build_parent_map(root, 0, parent_map) - - # Get actual node IDs of retained leaves - retained_node_ids = {leaves[idx][1] for idx in retained_leaf_indices} - - # Trace up to root for each retained leaf - ancestors: Set[int] = set() - for node_id in retained_node_ids: - current = node_id - while current in parent_map: - ancestors.add(current) - current = parent_map[current] - ancestors.add(current) # Add root - - return ancestors - - def _build_parent_map(self, node: Any, node_id: int, parent_map: Dict[int, int]) -> int: - """Build mapping from node ID to parent ID.""" - children = getattr(node, 'children', []) - next_id = node_id + 1 - - for child in children: - parent_map[next_id] = node_id - next_id = self._build_parent_map(child, next_id, parent_map) - - return next_id - - def _count_nodes(self, node: Any) -> int: - """Count total nodes in tree.""" - children = getattr(node, 'children', []) - return 1 + sum(self._count_nodes(c) for c in children) - - def _signature_to_features(self, sig: PatternSignature) -> Set[str]: - """Convert PatternSignature to feature set for overlap computation.""" - features: Set[str] = set() - - # Structural n-grams as features - for ngram, count in sig.structural_ngrams.items(): - if isinstance(ngram, tuple): - features.add(f"ngram:{'>'.join(str(x) for x in ngram)}") - else: - features.add(f"ngram:{ngram}") - - # AST paths as features - for parent, path, child, depth in sig.ast_paths: - features.add(f"path:{parent}→{path}→{child}") - - # Control flow features - cf = sig.control_flow - if cf.get("loop_count", 0) > 0: - features.add("cf:has_loop") - if cf.get("branch_count", 0) > 0: - features.add("cf:has_branch") - if cf.get("try_count", 0) > 0: - features.add("cf:has_try") - if cf.get("has_catch"): - features.add("cf:has_catch") - if cf.get("has_finally"): - features.add("cf:has_finally") - if cf.get("has_resource_guard"): - features.add("cf:has_resource_guard") - if cf.get("has_retry"): - features.add("cf:has_retry") - - # Control flow signature - if cf.get("signature"): - features.add(f"cfsig:{cf['signature']}") - - return features - - # ========================================================================= - # AROMA Phase II: Rerank search results using pruned similarity - # ========================================================================= - - def rerank_results( - self, - query_sig: PatternSignature, - candidates: List[Tuple[PatternSignature, float, Any]], # (sig, original_score, metadata) - top_k: int = 10, - alpha: float = 0.5, # Weight for pruned vs original score - ) -> List[Tuple[PatternSignature, float, PrunedResult, Any]]: - """ - Rerank search results using AROMA pruning. - - For each candidate: - 1. Prune it w.r.t. query to find maximal similar subtree - 2. Score = alpha * pruned_similarity + (1-alpha) * original_score - 3. Return top-k by combined score - - Args: - query_sig: Query pattern signature - candidates: List of (signature, original_score, metadata) tuples - top_k: Number of results to return - alpha: Weight for pruned similarity (0-1) - - Returns: - List of (signature, combined_score, prune_result, metadata) sorted by score - """ - reranked = [] - - for sig, orig_score, metadata in candidates: - prune_result = self.prune(query_sig, sig) - - # Combined score: pruned similarity + original embedding score - combined = alpha * prune_result.similarity_score + (1 - alpha) * orig_score - - reranked.append((sig, combined, prune_result, metadata)) - - # Sort by combined score descending - reranked.sort(key=lambda x: -x[1]) - - return reranked[:top_k] - - # ========================================================================= - # AROMA Phase III: Intersect for pattern discovery - # ========================================================================= - - def intersect_signatures( - self, - signatures: List[PatternSignature], - min_support: float = 0.8, # Feature must appear in this fraction of signatures - ) -> PatternSignature: - """ - Extract common structure from multiple signatures (AROMA intersection). - - This is used after clustering to find the COMMON pattern across examples. - - Args: - signatures: List of similar pattern signatures - min_support: Minimum fraction of signatures a feature must appear in - - Returns: - New PatternSignature containing only common features - """ - if not signatures: - return PatternSignature() - - if len(signatures) == 1: - return signatures[0] - - # Count feature occurrences across all signatures - feature_counts: Counter = Counter() - all_features_by_sig: List[Set[str]] = [] - - for sig in signatures: - features = self._signature_to_features(sig) - all_features_by_sig.append(features) - feature_counts.update(features) - - # Keep features with sufficient support - threshold = int(len(signatures) * min_support) - common_features = {f for f, count in feature_counts.items() if count >= threshold} - - # Build intersected signature - result = PatternSignature(language="multi") # Cross-language pattern - - # Intersect structural n-grams - for sig in signatures: - for ngram, count in sig.structural_ngrams.items(): - ngram_feature = f"ngram:{'>'.join(str(x) for x in ngram)}" if isinstance(ngram, tuple) else f"ngram:{ngram}" - if ngram_feature in common_features: - result.structural_ngrams[ngram] = min( - result.structural_ngrams.get(ngram, float('inf')), - count - ) - - # Fix inf values to actual counts - result.structural_ngrams = Counter({ - k: v for k, v in result.structural_ngrams.items() - if v != float('inf') - }) - - # Intersect AST paths - path_counts: Counter = Counter() - for sig in signatures: - for path in sig.ast_paths: - path_key = (path[0], path[1], path[2]) # Exclude depth - path_counts[path_key] += 1 - - for path_key, count in path_counts.items(): - if count >= threshold: - result.ast_paths.append((*path_key, 0)) # Add with depth=0 - - # Intersect control flow - cf_features = { - "loop_count": [], "branch_count": [], "try_count": [], - "has_catch": [], "has_finally": [], "has_resource_guard": [], - "has_retry": [], "max_nesting_depth": [], - } - - for sig in signatures: - cf = sig.control_flow - for key in cf_features: - if key in cf: - cf_features[key].append(cf[key]) - - result.control_flow = {} - for key, values in cf_features.items(): - if len(values) >= threshold: - if isinstance(values[0], bool): - # For booleans, require all True - result.control_flow[key] = all(values) - else: - # For counts, take minimum - result.control_flow[key] = min(values) - - # Collect languages - languages = {sig.language for sig in signatures if sig.language != "unknown"} - if languages: - result.language = ",".join(sorted(languages)) - - return result - - def compute_extension_score( - self, - query_sig: PatternSignature, - candidate_sig: PatternSignature, - prune_result: PrunedResult, - tau2: float = 1.5, # Candidate should be at least tau2 * query size - tau3: float = 0.9, # Minimum similarity threshold - ) -> Tuple[float, bool]: - """ - AROMA commonality constraints (Section 3.3.3). - - A recommendation should: - 1. Be similar enough to query (SimScore >= tau3) - 2. Contain significant extension beyond query (size >= tau2 * query_size) - - Returns: - (extension_score, passes_constraints) - """ - query_features = self._signature_to_features(query_sig) - candidate_features = self._signature_to_features(candidate_sig) - - query_size = len(query_features) - candidate_size = len(candidate_features) - - # Check similarity threshold - passes_similarity = prune_result.similarity_score >= tau3 - - # Check extension threshold - passes_extension = candidate_size >= tau2 * query_size - - # Extension score: how much extra the candidate provides - common = len(prune_result.pruned_features & query_features) - extension = candidate_size - common - extension_score = extension / max(query_size, 1) - - passes = passes_similarity and passes_extension - - return extension_score, passes - diff --git a/scripts/pattern_detection/search.py b/scripts/pattern_detection/search.py deleted file mode 100644 index a3dd5e23..00000000 --- a/scripts/pattern_detection/search.py +++ /dev/null @@ -1,1009 +0,0 @@ -""" -Pattern Search - Structural code similarity search via Qdrant. - -This module provides the search interface for finding structurally similar code -across all supported languages. It integrates: - -1. PatternExtractor - Extracts normalized AST structure from code -2. PatternEncoder - Converts structure to searchable vectors -3. Qdrant - Vector database for fast similarity search -4. OnlinePatternLearner - Discovers patterns as you search - -Key features: -- Cross-language search: Python pattern matches Go/Rust/Java/etc. -- Example-based: "Find code like this" without writing queries -- Pattern-aware: Discovers and surfaces common idioms -- Hybrid mode: Combines structural + semantic similarity -- TOON output: Token-efficient output format when enabled -""" - -from __future__ import annotations - -import os -import logging -from dataclasses import dataclass, field -from typing import List, Dict, Optional, Any, Tuple, Union -from collections import Counter - -logger = logging.getLogger(__name__) - - -# ============================================================================= -# Helper Classes -# ============================================================================= - -class ScoredPoint: - """Wrapper to add a score attribute to scroll results for uniform handling.""" - __slots__ = ('id', 'payload', 'score') - - def __init__(self, point, score: float): - self.id = point.id - self.payload = point.payload - self.score = score - - -def _get_line_value(primary: Optional[int], fallback: Optional[int], default: int = 1) -> int: - """Get line number with proper None handling (0 is valid, None falls back).""" - if primary is not None: - return primary - if fallback is not None: - return fallback - return default - - -# ============================================================================= -# TOON Support -# ============================================================================= - -def _is_toon_enabled() -> bool: - """Check if TOON output format is enabled globally via TOON_ENABLED env var.""" - return os.environ.get("TOON_ENABLED", "0").lower() in ("1", "true", "yes") - - -def _should_use_toon(output_format: Any) -> bool: - """Determine if TOON format should be used based on explicit param or env flag.""" - if output_format is not None: - fmt = str(output_format).strip().lower() - return fmt == "toon" - return _is_toon_enabled() - - -def _format_pattern_results_as_toon( - response: Dict[str, Any], - compact: bool = False, -) -> Dict[str, Any]: - """Convert pattern search response to TOON format. - - Args: - response: Pattern search response dict with 'results' key - compact: If True, use minimal fields only - - Returns: - Modified response with TOON-encoded results - """ - try: - results = response.get("results", []) - if isinstance(results, list): - # Encode results to TOON format - toon_results = encode_pattern_results(results, compact=compact) - response["results"] = toon_results - response["output_format"] = "toon" - return response - except Exception as e: - logger.debug(f"TOON encoding failed: {e}") - return response - - -def encode_pattern_results( - results: List[Dict[str, Any]], - delimiter: str = ",", - compact: bool = False, -) -> str: - """Encode pattern search results to TOON tabular format. - - Args: - results: List of pattern search result dicts - delimiter: Field delimiter (default: ",") - compact: If True, only include core location fields - - Returns: - TOON-formatted pattern results string - """ - if not results: - return "results[0]:" - - # Determine fields based on compact mode - if compact: - fields = ["path", "start_line", "end_line", "score", "language", "matched_lines"] - else: - # Full fields for pattern results - fields = [ - "path", "start_line", "end_line", "score", "language", - "control_flow_signature", "matched_patterns", "matched_lines", "snippet", - "semantic_score", "combined_score" - ] - - # Filter to fields actually present - all_present: set = set() - for r in results: - all_present.update(r.keys()) - fields = [f for f in fields if f in all_present] - - # Build TOON output - bracket = f"[{len(results)}]" - fields_part = "{" + delimiter.join(fields) + "}" - - lines = [f"results{bracket}{fields_part}:"] - for r in results: - values = [] - for f in fields: - val = r.get(f) - values.append(_encode_toon_value(val, delimiter)) - lines.append(f" {delimiter.join(values)}") - - return "\n".join(lines) - - -def _encode_toon_value(value: Any, delimiter: str) -> str: - """Encode a single value to TOON format.""" - if value is None: - return "null" - if isinstance(value, bool): - return "true" if value else "false" - if isinstance(value, (int, float)): - if isinstance(value, float): - return f"{value:.4f}" - return str(value) - if isinstance(value, str): - # Check if quoting needed - needs_quote = ( - not value or - value[0].isspace() or value[-1].isspace() or - value in ("true", "false", "null") or - delimiter in value or - any(c in value for c in (':', '"', '\\', '[', ']', '{', '}', '\n', '\r', '\t')) - ) - if needs_quote: - escaped = value.replace('\\', '\\\\').replace('"', '\\"') - escaped = escaped.replace('\n', '\\n').replace('\r', '\\r').replace('\t', '\\t') - return f'"{escaped}"' - return value - if isinstance(value, list): - # Compact list encoding - if not value: - return "[]" - items = [_encode_toon_value(v, delimiter) for v in value] - return "[" + ";".join(items) + "]" - # Objects: compact JSON fallback - import json - return json.dumps(value, separators=(",", ":")) - -# Lazy imports for optional dependencies -_qdrant_client = None -_extractor = None -_encoder = None -_learner = None - - -def _get_qdrant_client(): - """Get or create Qdrant client.""" - global _qdrant_client - if _qdrant_client is None: - try: - from qdrant_client import QdrantClient - # Support QDRANT_URL (docker/k8s) or QDRANT_HOST/PORT (local dev) - url = os.environ.get("QDRANT_URL") - if url: - _qdrant_client = QdrantClient(url=url) - else: - host = os.environ.get("QDRANT_HOST", "localhost") - port = int(os.environ.get("QDRANT_PORT", "6333")) - _qdrant_client = QdrantClient(host=host, port=port) - except Exception as e: - logger.warning(f"Failed to connect to Qdrant: {e}") - return None - return _qdrant_client - - -def _get_extractor(): - """Get or create pattern extractor.""" - global _extractor - if _extractor is None: - from .extractor import PatternExtractor - _extractor = PatternExtractor() - return _extractor - - -def _get_encoder(): - """Get or create pattern encoder.""" - global _encoder - if _encoder is None: - from .encoder import PatternEncoder - _encoder = PatternEncoder() - return _encoder - - -def _get_learner(): - """Get or create pattern learner.""" - global _learner - if _learner is None: - from .catalog import get_pattern_learner - _learner = get_pattern_learner() - return _learner - - -@dataclass -class PatternSearchResult: - """A single result from pattern search.""" - path: str - start_line: int - end_line: int - score: float # Structural similarity score (0-1) - language: str - snippet: Optional[str] = None - - # Pattern analysis - matched_patterns: List[str] = field(default_factory=list) - control_flow_signature: str = "" - - # Line numbers within snippet that match query pattern (absolute, 1-indexed) - matched_lines: List[int] = field(default_factory=list) - - # Combined scoring (if hybrid search) - semantic_score: Optional[float] = None - combined_score: Optional[float] = None - - def to_dict(self) -> Dict[str, Any]: - """Convert to dictionary for JSON serialization.""" - d = { - "path": self.path, - "start_line": self.start_line, - "end_line": self.end_line, - "score": round(self.score, 4), - "language": self.language, - "snippet": self.snippet, - "matched_patterns": self.matched_patterns, - "control_flow_signature": self.control_flow_signature, - "semantic_score": round(self.semantic_score, 4) if self.semantic_score is not None else None, - "combined_score": round(self.combined_score, 4) if self.combined_score is not None else None, - } - if self.matched_lines: - d["matched_lines"] = self.matched_lines - return d - - -@dataclass -class PatternSearchResponse: - """Response from pattern search.""" - results: List[PatternSearchResult] - total: int - query_signature: str # Control flow signature of query - discovered_patterns: List[str] = field(default_factory=list) - languages_searched: List[str] = field(default_factory=list) - search_mode: str = "structural" # structural, semantic, hybrid, error - output_format: str = "json" # json or toon - - def to_dict(self, compact: bool = False) -> Dict[str, Any]: - """Convert to dictionary for JSON/TOON serialization.""" - # ok=False when search_mode is "error" to signal failure to callers - is_ok = self.search_mode != "error" - return { - "ok": is_ok, - "results": [r.to_dict() for r in self.results], - "total": self.total, - "query_signature": self.query_signature, - "discovered_patterns": self.discovered_patterns, - "languages_searched": self.languages_searched, - "search_mode": self.search_mode, - "output_format": self.output_format, - } - - def format(self, output_format: Any = None, compact: bool = False) -> Dict[str, Any]: - """Format response as JSON or TOON based on output_format parameter.""" - response = self.to_dict(compact=compact) - - if _should_use_toon(output_format): - return _format_pattern_results_as_toon(response, compact=compact) - return response - - -# ============================================================================= -# Core Search Functions -# ============================================================================= - -def pattern_search( - example: str, - language: str = "python", - *, - limit: int = 10, - collection: Optional[str] = None, - min_score: float = 0.5, - include_snippet: bool = True, - context_lines: int = 3, - target_languages: Optional[List[str]] = None, - hybrid: bool = False, - semantic_weight: float = 0.3, - output_format: Any = None, # "json" (default) or "toon" - compact: bool = False, - client: Any = None, # Optional QdrantClient override for testing - aroma_rerank: bool = True, # Enable AROMA-style pruning + reranking (default ON) - aroma_alpha: float = 0.6, # Weight for pruned similarity (vs original score) -) -> Union[PatternSearchResponse, Dict[str, Any]]: - """ - Find code structurally similar to the given example. - - This is the primary search interface. Given a code example, it finds - other code with similar structure across ALL supported languages. - - Args: - example: Code snippet to find similar code for - language: Language of the example code - limit: Maximum results to return - collection: Qdrant collection (defaults to COLLECTION_NAME env) - min_score: Minimum similarity score (0-1) - include_snippet: Include code snippets in results - context_lines: Lines of context around matches - target_languages: Filter to specific languages (None = all) - hybrid: Combine structural + semantic similarity - semantic_weight: Weight for semantic score in hybrid mode (0-1) - output_format: "json" (default) or "toon" for token-efficient format - compact: If True with TOON, use minimal fields only - aroma_rerank: Enable AROMA-style pruning + reranking (default True) - aroma_alpha: Weight for pruned similarity vs original score (0-1, default 0.6) - - Returns: - PatternSearchResponse or Dict (TOON format) with matching code and patterns - - Example: - >>> results = pattern_search(''' - ... for i in range(retries): - ... try: - ... return do_request() - ... except Exception: - ... time.sleep(2 ** i) - ... ''', language="python", output_format="toon") - >>> # Finds retry patterns in Python, Go, Rust, Java, etc. - """ - # Determine if TOON output is requested - use_toon = _should_use_toon(output_format) - - # AROMA/hybrid reranking requires snippets for pruning/semantic scoring - if (aroma_rerank or hybrid) and not include_snippet: - logger.debug("Forcing include_snippet=True (required for AROMA/hybrid scoring)") - include_snippet = True - - extractor = _get_extractor() - encoder = _get_encoder() - # Use provided client or fall back to global - if client is None: - client = _get_qdrant_client() - - if client is None: - response = PatternSearchResponse( - results=[], - total=0, - query_signature="ERROR:NO_QDRANT", - search_mode="error", - ) - return response.format(output_format, compact) if use_toon else response - - # Extract structural signature from example - signature = extractor.extract(example, language) - query_vector = encoder.encode(signature) - cf_sig = signature.control_flow.get("signature", "") - - # Determine collection - if collection is None: - collection = os.environ.get("COLLECTION_NAME", "codebase") - - # Build search filter - search_filter = None - if target_languages: - from qdrant_client.models import Filter, FieldCondition, MatchAny - search_filter = Filter( - must=[ - FieldCondition( - key="language", - match=MatchAny(any=target_languages) - ) - ] - ) - - try: - # Search for similar patterns - # First check if collection has pattern_vector field - collection_info = client.get_collection(collection) - vectors_config = collection_info.config.params.vectors - # vectors can be a dict (named) or VectorParams (single unnamed vector) - has_pattern_vector = ( - isinstance(vectors_config, dict) and "pattern_vector" in vectors_config - ) - - # Fetch size: over-fetch only when reranking is enabled - needs_reranking = aroma_rerank or hybrid - fetch_limit = limit * 3 if needs_reranking else limit - - if has_pattern_vector: - # Use dedicated pattern vector - response = client.query_points( - collection_name=collection, - query=query_vector, - using="pattern_vector", - limit=fetch_limit, - query_filter=search_filter, - with_payload=True, - ) - results = response.points - else: - # Fallback: search using semantic vector with pattern-based reranking - results = _fallback_pattern_search( - client, collection, query_vector, signature, - fetch_limit, search_filter - ) - - # Process candidates - collect full pool when reranking, early-stop otherwise - search_results = [] - seen_paths = set() - - for hit in results: - payload = hit.payload or {} - # Support both flat payload (path at top level) and nested (metadata.path) - meta = payload.get("metadata", {}) - path = meta.get("path") or payload.get("path") or payload.get("file_path", "") - - # Deduplicate: prefer path, fall back to point id when path is empty - dedup_key = path if path else str(getattr(hit, 'id', id(hit))) - if dedup_key in seen_paths: - continue - seen_paths.add(dedup_key) - - result = PatternSearchResult( - path=path, - start_line=_get_line_value(meta.get("start_line"), payload.get("start_line"), 1), - end_line=_get_line_value(meta.get("end_line"), payload.get("end_line"), 1), - score=hit.score, - language=meta.get("language") or payload.get("language", "unknown"), - control_flow_signature=payload.get("cf_signature", ""), - ) - - # Include snippet if requested (forced True when aroma_rerank - see line ~366) - if include_snippet: - result.snippet = _get_snippet( - path, result.start_line, result.end_line, context_lines - ) - # Highlight lines matching query pattern (AST-based) - if result.snippet and signature.control_flow: - result.matched_lines = extractor.find_matching_lines( - result.snippet, - result.language, - signature.control_flow, - line_offset=result.start_line - 1, - ) - - search_results.append(result) - - # Early stop only when: - # 1. Not reranking (reranking needs full pool to reorder) - # 2. No min_score filter (filtering may discard results, need extras) - if not needs_reranking and min_score <= 0 and len(search_results) >= limit: - break - - # AROMA-style reranking: prune each result w.r.t. query, rerank by pruned similarity - if aroma_rerank and search_results: - search_results = _apply_aroma_reranking( - search_results, signature, extractor, aroma_alpha - ) - - # Hybrid mode: blend AROMA score (if present) with semantic score - if hybrid and search_results: - search_results = _apply_hybrid_scoring( - search_results, example, language, semantic_weight - ) - - # Apply min_score filter AFTER reranking (on combined_score if available) - if min_score > 0: - search_results = [ - r for r in search_results - if (r.combined_score if r.combined_score is not None else r.score) >= min_score - ] - - # Final slice to requested limit - search_results = search_results[:limit] - - # Discover patterns from results - discovered = _discover_patterns_from_results(search_results, signature) - - search_mode = "structural" - if aroma_rerank: - search_mode = "aroma" if not hybrid else "aroma_hybrid" - elif hybrid: - search_mode = "hybrid" - - response = PatternSearchResponse( - results=search_results, - total=len(search_results), - query_signature=cf_sig, - discovered_patterns=discovered, - languages_searched=list(set(r.language for r in search_results)), - search_mode=search_mode, - ) - return response.format(output_format, compact) if use_toon else response - - except Exception as e: - logger.error(f"Pattern search failed: {e}") - response = PatternSearchResponse( - results=[], - total=0, - query_signature=cf_sig, - search_mode="error", - ) - return response.format(output_format, compact) if use_toon else response - - -def find_similar_patterns( - code: str, - language: str = "python", - *, - limit: int = 5, - collection: Optional[str] = None, - output_format: Any = None, - compact: bool = False, -) -> Union[List[Dict[str, Any]], Dict[str, Any]]: - """ - Simplified interface: find code with similar patterns. - - Args: - code: Code snippet to find similar patterns for - language: Language of the code - limit: Maximum results - collection: Qdrant collection - output_format: "json" (default) or "toon" for token-efficient format - compact: If True with TOON, use minimal fields - - Returns: - List of dicts with path, score, snippet (JSON) or TOON-formatted dict - """ - response = pattern_search( - example=code, - language=language, - limit=limit, - collection=collection, - include_snippet=True, - output_format=output_format, - compact=compact, - ) - # If TOON was used, response is already formatted - if isinstance(response, dict): - return response - return [r.to_dict() for r in response.results] - - -def search_by_pattern_description( - description: str, - *, - limit: int = 10, - min_score: float = 0.0, - collection: Optional[str] = None, - target_languages: Optional[List[str]] = None, - output_format: Any = None, - compact: bool = False, -) -> Union[PatternSearchResponse, Dict[str, Any]]: - """ - Search using natural language pattern description. - - Examples: - - "retry with exponential backoff" - - "resource cleanup with finally" - - "null check guard clause" - - Args: - description: Natural language description of the pattern - limit: Maximum results - min_score: Minimum similarity score (0-1), applied before TOON encoding - collection: Qdrant collection - target_languages: Filter to specific languages - output_format: "json" (default) or "toon" for token-efficient format - compact: If True with TOON, use minimal fields - - This uses the OnlinePatternLearner to match descriptions to - discovered patterns, then searches for code matching those patterns. - """ - use_toon = _should_use_toon(output_format) - learner = _get_learner() - - # Find patterns matching the description - matched_patterns = learner.natural_language_query(description, top_k=3) - - if not matched_patterns: - # Fallback: use description as pseudo-code - return pattern_search( - example=description, - language="python", - limit=limit, - min_score=min_score, - collection=collection, - target_languages=target_languages, - output_format=output_format, - compact=compact, - ) - - # Use best matching pattern's centroid for search - best_pattern = matched_patterns[0] - - client = _get_qdrant_client() - if client is None or not best_pattern.centroid: - # Mark as error so ok=False propagates to callers - response = PatternSearchResponse( - results=[], - total=0, - query_signature="NL:" + description[:30], - discovered_patterns=[p.auto_description for p in matched_patterns], - search_mode="error", - ) - return response.format(output_format, compact) if use_toon else response - - if collection is None: - collection = os.environ.get("COLLECTION_NAME", "codebase") - - # Build filter - search_filter = None - if target_languages: - from qdrant_client.models import Filter, FieldCondition, MatchAny - search_filter = Filter( - must=[FieldCondition(key="language", match=MatchAny(any=target_languages))] - ) - - try: - # Check if collection has pattern_vector field - collection_info = client.get_collection(collection) - vectors_config = collection_info.config.params.vectors - has_pattern_vector = ( - isinstance(vectors_config, dict) and "pattern_vector" in vectors_config - ) - - if has_pattern_vector: - response = client.query_points( - collection_name=collection, - query=best_pattern.centroid, - using="pattern_vector", - limit=limit, - query_filter=search_filter, - with_payload=True, - with_vectors=False, # Only need payload and scores - ) - results = response.points - else: - # No pattern_vector in collection - use scroll + in-memory reranking - # Cannot use vector search: pattern centroid (64-dim) != semantic vector (384-dim) - logger.debug("No pattern_vector field, using scroll + rerank for NL search") - scroll_results, _ = client.scroll( - collection_name=collection, - scroll_filter=search_filter, - limit=limit * 3, - with_payload=True, - with_vectors=False, - ) - # Rerank by pattern description match (simple keyword overlap) - description_words = set(description.lower().split()) - scored_results = [] - for point in scroll_results: - payload = point.payload or {} - meta = payload.get("metadata", {}) - # Check common text field names used in code collections - text = ( - payload.get("text", "") or - payload.get("content", "") or - payload.get("code", "") or - meta.get("text", "") or - meta.get("code", "") - ) - text_words = set(text.lower().split()) - overlap = len(description_words & text_words) / max(len(description_words), 1) - # Wrap in ScoredPoint to preserve the overlap score - scored_results.append(ScoredPoint(point, overlap)) - scored_results.sort(key=lambda x: x.score, reverse=True) - results = scored_results[:limit] - - search_results = [] - for hit in results: - payload = hit.payload or {} - # Support both flat payload and nested metadata (consistent with pattern_search) - meta = payload.get("metadata", {}) - path = meta.get("path") or payload.get("path") or payload.get("file_path", "") - # Use .score attribute (works for both search results and ScoredPoint) - score = getattr(hit, 'score', 0.5) - search_results.append(PatternSearchResult( - path=path, - start_line=_get_line_value(meta.get("start_line"), payload.get("start_line"), 1), - end_line=_get_line_value(meta.get("end_line"), payload.get("end_line"), 1), - score=score, - language=meta.get("language") or payload.get("language", "unknown"), - matched_patterns=[best_pattern.auto_description], - )) - - # Apply min_score filtering BEFORE TOON encoding (so it works regardless of output format) - if min_score > 0: - search_results = [r for r in search_results if r.score >= min_score] - - response = PatternSearchResponse( - results=search_results, - total=len(search_results), - query_signature="NL:" + description[:30], - discovered_patterns=[p.auto_description for p in matched_patterns], - languages_searched=list(set(r.language for r in search_results)), - search_mode="natural_language", - ) - return response.format(output_format, compact) if use_toon else response - - except Exception as e: - logger.error(f"NL pattern search failed: {e}") - response = PatternSearchResponse( - results=[], - total=0, - query_signature="NL:" + description[:30], - search_mode="error", - ) - return response.format(output_format, compact) if use_toon else response - - - -# ============================================================================= -# Helper Functions -# ============================================================================= - -def _fallback_pattern_search( - client, - collection: str, - pattern_vector: List[float], - signature, - limit: int, - search_filter, -) -> List: - """ - Fallback search when collection doesn't have pattern_vector field. - - Uses scroll with in-memory pattern reranking since we can't do - vector search with mismatched dimensions. - - Note: limit is already inflated by caller for reranking; don't multiply again. - """ - # Can't use pattern_vector (64-dim) against semantic vectors (384-dim) - # Instead, scroll through relevant documents and rerank by pattern similarity - try: - # Scroll with filter to get candidate documents - # limit is pre-inflated by caller (e.g., limit*3 for reranking) - semantic_results, _ = client.scroll( - collection_name=collection, - scroll_filter=search_filter, - limit=limit, - with_payload=True, - with_vectors=False, - ) - except Exception as e: - logger.debug(f"Fallback scroll failed: {e}") - return [] - - # Rerank by structural similarity using cf_signature in payload - query_cf = signature.control_flow.get("normalized_sequence", []) - - # Baseline score for documents without cf_sequence or when query has no CF nodes: - # - Set BELOW default min_score (0.5) so unverified matches don't pass by default - # - This prioritizes precision: only items with actual pattern similarity surface - # - Callers can lower min_score (e.g., 0.3) to include these fallback results - BASELINE_SCORE_NO_CF = 0.4 - - reranked = [] - for point in semantic_results: - payload = point.payload or {} - - # Calculate similarity from stored pattern info - stored_cf = payload.get("cf_sequence", []) - if stored_cf and query_cf: - # Jaccard similarity of control flow sequences - set_query = set(query_cf) - set_stored = set(stored_cf) - intersection = len(set_query & set_stored) - union = len(set_query | set_stored) - cf_similarity = intersection / union if union > 0 else 0.0 - elif not stored_cf or not query_cf: - # No cf_sequence in payload OR query has no control-flow nodes - # Assign baseline score - filtered by default min_score unless caller lowers it - cf_similarity = BASELINE_SCORE_NO_CF - - reranked.append(ScoredPoint(point, cf_similarity)) - - # Sort by score and return top results - reranked.sort(key=lambda x: x.score, reverse=True) - return reranked[:limit] - - -def _get_snippet( - path: str, - start_line: int, - end_line: int, - context_lines: int, -) -> Optional[str]: - """Read snippet from file with context.""" - try: - # Resolve path - workspace = os.environ.get("WORKSPACE_PATH", "/work") - full_path = os.path.join(workspace, path) if not os.path.isabs(path) else path - - if not os.path.exists(full_path): - return None - - with open(full_path, 'r', encoding='utf-8', errors='replace') as f: - lines = f.readlines() - - # Calculate range with context - start = max(0, start_line - 1 - context_lines) - end = min(len(lines), end_line + context_lines) - - return ''.join(lines[start:end]) - - except Exception as e: - logger.debug(f"Failed to read snippet from {path}: {e}") - return None - - -def _apply_hybrid_scoring( - results: List[PatternSearchResult], - query: str, - language: str, - semantic_weight: float, -) -> List[PatternSearchResult]: - """ - Apply hybrid scoring combining structural/AROMA and semantic similarity. - - If AROMA reranking ran first, blends AROMA's combined_score with semantic. - Otherwise blends raw structural score with semantic. - """ - try: - # Get embeddings for semantic comparison - from scripts.embeddings import get_embedding_model - model = get_embedding_model() - query_embedding = model.encode(query) - - for result in results: - if result.snippet: - snippet_embedding = model.encode(result.snippet) - # Cosine similarity - dot = sum(a * b for a, b in zip(query_embedding, snippet_embedding)) - norm_q = sum(a * a for a in query_embedding) ** 0.5 - norm_s = sum(a * a for a in snippet_embedding) ** 0.5 - semantic_sim = dot / (norm_q * norm_s) if norm_q and norm_s else 0 - - result.semantic_score = semantic_sim - # Blend with existing combined_score (from AROMA) or raw score - base_score = result.combined_score if result.combined_score is not None else result.score - result.combined_score = ( - (1 - semantic_weight) * base_score + - semantic_weight * semantic_sim - ) - # If no snippet but has AROMA score, keep it; otherwise use raw score - elif result.combined_score is None: - result.combined_score = result.score - - # Re-sort by combined score - results.sort(key=lambda r: r.combined_score or r.score, reverse=True) - - except Exception as e: - logger.debug(f"Hybrid scoring failed, using structural only: {e}") - - return results - - -def _apply_aroma_reranking( - results: List[PatternSearchResult], - query_signature, - extractor, - alpha: float = 0.6, -) -> List[PatternSearchResult]: - """ - Apply AROMA-style pruning and reranking. - - For each result: - 1. Extract its pattern signature - 2. Prune it w.r.t. query to find maximal similar subtree - 3. Compute combined score: alpha * pruned_similarity + (1-alpha) * original_score - 4. Re-sort by combined score - """ - try: - from .prune import AromaPruner - pruner = AromaPruner(extractor) - - reranked = [] - for result in results: - # Extract signature from result's snippet if available - if result.snippet: - result_sig = extractor.extract(result.snippet, result.language) - prune_result = pruner.prune(query_signature, result_sig) - - # Combined score - combined = alpha * prune_result.similarity_score + (1 - alpha) * result.score - result.combined_score = combined - - # Store pruning info in matched_patterns for debugging - if prune_result.similarity_score > 0: - result.matched_patterns.append( - f"aroma_sim:{prune_result.similarity_score:.2f}," - f"retained:{prune_result.pruned_feature_count}/{prune_result.original_feature_count}" - ) - else: - # No snippet - can't extract signature, use original score - result.combined_score = result.score - - reranked.append(result) - - # Re-sort by combined score - reranked.sort(key=lambda r: r.combined_score or r.score, reverse=True) - return reranked - - except Exception as e: - logger.debug(f"AROMA reranking failed, using original order: {e}") - return results - - -def _discover_patterns_from_results( - results: List[PatternSearchResult], - query_signature, -) -> List[str]: - """Identify common patterns across search results.""" - if not results: - return [] - - # Count control flow signatures - cf_counter = Counter() - for r in results: - if r.control_flow_signature: - cf_counter[r.control_flow_signature] += 1 - - # Return most common patterns - common = cf_counter.most_common(3) - return [sig for sig, count in common if count >= 2] - - -# ============================================================================= -# Convenience Aliases -# ============================================================================= - -def search_similar_code( - code: str, - output_format: Any = None, - compact: bool = False, - **kwargs, -) -> Union[PatternSearchResponse, Dict[str, Any]]: - """Alias for pattern_search with output_format support.""" - return pattern_search(code, output_format=output_format, compact=compact, **kwargs) - - -def find_code_like( - example: str, - output_format: Any = None, - compact: bool = False, - **kwargs, -) -> Union[PatternSearchResponse, Dict[str, Any]]: - """Alias for pattern_search with friendlier name.""" - return pattern_search(example, output_format=output_format, compact=compact, **kwargs) - - -# ============================================================================= -# Export -# ============================================================================= - -__all__ = [ - # Core search functions - "pattern_search", - "find_similar_patterns", - "search_by_pattern_description", - # Aliases - "search_similar_code", - "find_code_like", - # Data classes - "PatternSearchResult", - "PatternSearchResponse", - # TOON support - "encode_pattern_results", - "_should_use_toon", - "_format_pattern_results_as_toon", -] diff --git a/scripts/progressive_train.py b/scripts/progressive_train.py deleted file mode 100644 index cbf76383..00000000 --- a/scripts/progressive_train.py +++ /dev/null @@ -1,91 +0,0 @@ -#!/usr/bin/env python3 -"""Progressive training evaluation - measures quality at checkpoints.""" -import sys, os -sys.path.insert(0, '.') - -from scripts.rerank_eval import get_candidates, rerank_learning, rerank_onnx, DEFAULT_EVAL_QUERIES -from scripts.rerank_recursive import rerank_with_learning -from scripts.learning_reranker_worker import CollectionLearner -import numpy as np - -EXTRA_QUERIES = [ - 'function python', 'class MCP', 'vector embedding', 'cache LRU', - 'subprocess async', 'tokenizer ONNX', 'collection Qdrant', 'rerank fusion', - 'memory store', 'config env', 'error exception', 'test pytest', - 'docker build', 'API handler', 'database pool', -] - -def measure_quality(collection='eval'): - """Measure MRR against ONNX ground truth.""" - eval_qs = DEFAULT_EVAL_QUERIES[:6] - mrrs = [] - for eq in eval_qs: - cands = get_candidates(eq, limit=20) - if not cands: - continue - onnx_ranked = rerank_onnx(eq, [c.copy() for c in cands]) - onnx_top5 = set(c['path'] for c in onnx_ranked[:5]) - learn_ranked = rerank_learning(eq, [c.copy() for c in cands], collection=collection) - for rank, c in enumerate(learn_ranked, 1): - if c['path'] in onnx_top5: - mrrs.append(1.0 / rank) - break - else: - mrrs.append(0.0) - return np.mean(mrrs) if mrrs else 0 - -def main(): - all_queries = (DEFAULT_EVAL_QUERIES + EXTRA_QUERIES) * 15 - checkpoints = [50, 100, 150, 200, 250, 300] - results = [] - query_count = 0 - - print('Queries | Samples | Loss | MRR | Distill%', flush=True) - print('-' * 50, flush=True) - - for query in all_queries: - candidates = get_candidates(query, limit=25) - if candidates: - rerank_with_learning(query, candidates, learn_from_onnx=True, collection='eval') - query_count += 1 - - if query_count in checkpoints: - # Clear stale locks and process events - import glob - for lock in glob.glob('/tmp/rerank_weights/eval_*.lock'): - try: - os.remove(lock) - except Exception: - pass - - learner = CollectionLearner(collection='eval') - learner.process_events() - m = learner.scorer.get_metrics() - - # Measure quality - mrr = measure_quality() - - print(f'{query_count:7} | {m.get("total_samples",0):7} | {m.get("avg_loss",0):6.3f} | {mrr:.3f} | {mrr*100:.1f}%', flush=True) - results.append({ - 'queries': query_count, - 'samples': m.get('total_samples', 0), - 'loss': round(float(m.get('avg_loss', 0)), 3), - 'mrr': round(mrr, 3), - 'distill_pct': round(mrr * 100, 1), - }) - checkpoints.remove(query_count) - - if query_count >= 300: - break - - print('\n' + '=' * 50) - print('PROGRESSIVE TRAINING RESULTS') - print('=' * 50) - for r in results: - print(f" {r['queries']:3} queries → MRR {r['mrr']:.3f} ({r['distill_pct']}% of ONNX)") - - return results - -if __name__ == '__main__': - main() - diff --git a/scripts/prune.py b/scripts/prune.py index 5e2f14fb..2f628598 100755 --- a/scripts/prune.py +++ b/scripts/prune.py @@ -2,15 +2,37 @@ import os import hashlib from pathlib import Path -from typing import Tuple +from typing import Tuple, Any from qdrant_client import QdrantClient, models +try: + from scripts.ingest.graph_edges import ( + delete_edges_by_path as _shared_delete_graph_edges_by_path, + get_graph_collection_name as _shared_graph_collection_name, + ) +except Exception: + _shared_delete_graph_edges_by_path = None # type: ignore[assignment] + _shared_graph_collection_name = None # type: ignore[assignment] COLLECTION = os.environ.get("COLLECTION_NAME", "codebase") QDRANT_URL = os.environ.get("QDRANT_URL", "http://localhost:6333") API_KEY = os.environ.get("QDRANT_API_KEY") ROOT = Path(os.environ.get("PRUNE_ROOT", ".")).resolve() -GRAPH_COLLECTION = os.environ.get("GRAPH_COLLECTION_NAME", f"{COLLECTION}_graph") +GRAPH_COLLECTION = ( + _shared_graph_collection_name(COLLECTION) + if _shared_graph_collection_name is not None + else f"{COLLECTION}_graph" +) + + +def _norm_path(path_str: Any) -> str: + if not path_str: + return "" + try: + normalized = os.path.normpath(str(path_str)) + except Exception: + normalized = str(path_str) + return normalized.replace("\\", "/") def sha1_file(path: Path) -> str: @@ -39,50 +61,103 @@ def delete_by_path(client: QdrantClient, path_str: str) -> int: return 0 -def delete_graph_edges_by_path(client: QdrantClient, path_str: str) -> int: +def delete_graph_edges_by_path(client: QdrantClient, path_str: str, repo: str | None = None) -> int: """Best-effort deletion for graph-edge collections (if present). Some deployments store symbol-graph edges in a separate Qdrant collection - (commonly `${COLLECTION}_graph`). Those points may reference a file path as - either caller or callee; delete both to prevent stale graph results. + (commonly `${COLLECTION}_graph`). On this branch, edge docs are file-level and + reference a file path as `caller_path`. """ if not path_str: return 0 + path_str = _norm_path(path_str) - flt = models.Filter( - should=[ - models.FieldCondition( - key="caller_path", match=models.MatchValue(value=path_str) - ), - models.FieldCondition( - key="callee_path", match=models.MatchValue(value=path_str) - ), - ] - ) + # Canonical path: shared graph-edge deleter against _graph. + if _shared_delete_graph_edges_by_path is None: + return 0 try: - res = client.delete( + return int( + _shared_delete_graph_edges_by_path( + client, + COLLECTION, + caller_path=path_str, + repo=repo, + ) + or 0 + ) + except Exception: + return 0 + + +def _graph_collection_exists(client: QdrantClient) -> bool: + try: + client.get_collection(collection_name=GRAPH_COLLECTION) + return True + except Exception: + return False + + +def _delete_graph_points_by_ids(client: QdrantClient, ids: list[Any]) -> int: + if not ids: + return 0 + try: + from qdrant_client import models as qmodels + client.delete( collection_name=GRAPH_COLLECTION, - points_selector=models.FilterSelector(filter=flt), + points_selector=qmodels.PointIdsList(points=ids), ) - # Qdrant responses vary by client version; return 1 as "success" when count isn't available. - deleted_count = None - result_attr = getattr(res, "result", None) - if isinstance(result_attr, dict): - v = result_attr.get("deleted") - if isinstance(v, int): - deleted_count = v - if deleted_count is None: - v = getattr(res, "deleted", None) - if isinstance(v, int): - deleted_count = v - if deleted_count is None: - deleted_count = 1 - return deleted_count + return len(ids) except Exception: - # Non-fatal: graph collection may not exist in this deployment. return 0 +def delete_orphan_graph_edges(client: QdrantClient, valid_paths: set[str]) -> int: + """Delete graph-edge points whose `caller_path` no longer exists in base collection.""" + if not _graph_collection_exists(client): + return 0 + + removed = 0 + next_page = None + pending_ids: list[Any] = [] + batch_size = 256 + + while True: + try: + points, next_page = client.scroll( + collection_name=GRAPH_COLLECTION, + with_payload=True, + with_vectors=False, + limit=512, + offset=next_page, + scroll_filter=None, + ) + except Exception: + break + + if not points: + break + + for p in points: + payload = p.payload or {} + caller_path = _norm_path(payload.get("caller_path")) + if not caller_path: + continue + if caller_path in valid_paths: + continue + pending_ids.append(p.id) + if len(pending_ids) >= batch_size: + removed += _delete_graph_points_by_ids(client, pending_ids) + pending_ids = [] + + if next_page is None: + break + + if pending_ids: + removed += _delete_graph_points_by_ids(client, pending_ids) + + return removed + + def main(): client = QdrantClient(url=QDRANT_URL, api_key=API_KEY or None) @@ -90,6 +165,7 @@ def main(): removed_missing = 0 removed_mismatch = 0 removed_graph_edges = 0 + removed_orphan_graph_edges = 0 next_page = None while True: @@ -106,9 +182,9 @@ def main(): md = (p.payload or {}).get("metadata") or {} path_str = md.get("path") file_hash = md.get("file_hash") - if not path_str or path_str in seen: + norm_path = _norm_path(path_str) + if not norm_path or norm_path in seen: continue - seen.add(path_str) abs_path = ( ROOT / Path(path_str).relative_to("/work") if path_str.startswith("/work/") @@ -116,23 +192,37 @@ def main(): ) if not abs_path.exists(): removed_missing += delete_by_path(client, path_str) - removed_graph_edges += delete_graph_edges_by_path(client, path_str) + deleted = delete_graph_edges_by_path(client, path_str, md.get("repo")) + if deleted == 0: + # Repo tags can drift across ingestion modes; fall back to path-only delete. + deleted = delete_graph_edges_by_path(client, path_str, None) + removed_graph_edges += deleted print(f"[prune] removed missing file points: {path_str}") continue current_hash = sha1_file(abs_path) if file_hash and current_hash and current_hash != file_hash: removed_mismatch += delete_by_path(client, path_str) - removed_graph_edges += delete_graph_edges_by_path(client, path_str) + deleted = delete_graph_edges_by_path(client, path_str, md.get("repo")) + if deleted == 0: + deleted = delete_graph_edges_by_path(client, path_str, None) + removed_graph_edges += deleted print(f"[prune] removed outdated points (hash mismatch): {path_str}") + continue + + seen.add(norm_path) if next_page is None: break + # Secondary pass: if base points were manually deleted, remove orphan `_graph` edges. + removed_orphan_graph_edges = delete_orphan_graph_edges(client, seen) + print( "Prune complete. " f"removed_missing={removed_missing}, " f"removed_mismatch={removed_mismatch}, " - f"removed_graph_edges={removed_graph_edges}" + f"removed_graph_edges={removed_graph_edges}, " + f"removed_orphan_graph_edges={removed_orphan_graph_edges}" ) diff --git a/scripts/pseudo_config.py b/scripts/pseudo_config.py new file mode 100644 index 00000000..0f7b1f0b --- /dev/null +++ b/scripts/pseudo_config.py @@ -0,0 +1,52 @@ +"""Shared configuration helpers for pseudo/tags generation. + +This keeps env semantics consistent across: +- watcher (watch_index / watch_index_core) +- indexing CLI (scripts/ingest/cli.py) + +Policy: +- PSEUDO_BACKFILL_ENABLED controls whether the async backfill worker is enabled. +- PSEUDO_DEFER_TO_WORKER controls *foreground vs background* behavior only. + Deferral is only effective when the worker is enabled; otherwise we keep inline + pseudo generation ON to avoid silently dropping pseudo/tags. +""" + +from __future__ import annotations + +import os +from typing import Optional + + +def _parse_env_bool(value: Optional[str], *, default: bool = False) -> bool: + if value is None: + return default + v = str(value).strip().lower() + if not v: + return default + return v in {"1", "true", "yes", "on"} + + +def env_bool(key: str, *, default: bool = False) -> bool: + """Read a boolean env var using consistent truthy parsing.""" + return _parse_env_bool(os.environ.get(key), default=default) + + +def effective_defer_to_worker(*, defer_to_worker: bool, backfill_enabled: bool) -> bool: + """Whether we should disable inline pseudo/tags generation.""" + return bool(defer_to_worker and backfill_enabled) + + +def effective_pseudo_mode(*, defer_to_worker: bool, backfill_enabled: bool) -> str: + """Return pseudo_mode ('off'|'full') for indexing pipeline.""" + return "off" if effective_defer_to_worker( + defer_to_worker=defer_to_worker, + backfill_enabled=backfill_enabled, + ) else "full" + + +__all__ = [ + "env_bool", + "effective_defer_to_worker", + "effective_pseudo_mode", +] + diff --git a/scripts/qdrant_client_manager.py b/scripts/qdrant_client_manager.py index bcb4adbe..198b0cca 100644 --- a/scripts/qdrant_client_manager.py +++ b/scripts/qdrant_client_manager.py @@ -8,9 +8,19 @@ import threading import time import weakref -from typing import Optional, Dict, List +from typing import Any, Optional, Dict, List, TYPE_CHECKING from contextlib import contextmanager -from qdrant_client import QdrantClient + +if TYPE_CHECKING: + from qdrant_client import QdrantClient +else: + QdrantClient = Any + + +def _new_qdrant_client(url: str, api_key: Optional[str] = None) -> QdrantClient: + from qdrant_client import QdrantClient as _QdrantClient + + return _QdrantClient(url=url, api_key=api_key if api_key else None) # Connection pool implementation @@ -44,7 +54,7 @@ def get_client(self, url: str, api_key: Optional[str] = None) -> QdrantClient: # No suitable client found, create a new one if self._created_count < self.max_size: - client = QdrantClient(url=url, api_key=api_key) + client = _new_qdrant_client(url, api_key) pool_entry = { 'client': client, 'url': url, @@ -60,7 +70,7 @@ def get_client(self, url: str, api_key: Optional[str] = None) -> QdrantClient: else: # Pool is full, create a temporary client (not pooled) self._misses += 1 - return QdrantClient(url=url, api_key=api_key) + return _new_qdrant_client(url, api_key) def return_client(self, client: QdrantClient): """Return a client to the pool.""" @@ -166,13 +176,13 @@ def get_qdrant_client( # Fallback to singleton pattern for backward compatibility if force_new: - return QdrantClient(url=url, api_key=api_key if api_key else None) + return _new_qdrant_client(url, api_key) global _client with _client_lock: if _client is None: - _client = QdrantClient(url=url, api_key=api_key if api_key else None) + _client = _new_qdrant_client(url, api_key) return _client diff --git a/scripts/query_named_vector.py b/scripts/query_named_vector.py index d87ede8b..cfb64608 100644 --- a/scripts/query_named_vector.py +++ b/scripts/query_named_vector.py @@ -1,21 +1,8 @@ #!/usr/bin/env python3 import os -import sys -from pathlib import Path from qdrant_client import QdrantClient -# Ensure scripts is importable -ROOT_DIR = Path(__file__).resolve().parent.parent -if str(ROOT_DIR) not in sys.path: - sys.path.insert(0, str(ROOT_DIR)) - -# Use embedder factory for Qwen3 support -try: - from scripts.embedder import get_embedding_model - _EMBEDDER_FACTORY = True -except ImportError: - _EMBEDDER_FACTORY = False - from fastembed import TextEmbedding +from scripts.embedder import get_embedding_model from scripts.utils import sanitize_vector_name QDRANT_URL = os.environ.get("QDRANT_URL", "http://qdrant:6333") @@ -24,10 +11,7 @@ VEC_NAME = os.environ.get("VECTOR_NAME") or sanitize_vector_name(MODEL) client = QdrantClient(url=QDRANT_URL) -if _EMBEDDER_FACTORY: - emb = get_embedding_model(MODEL) -else: - emb = TextEmbedding(model_name=MODEL) +emb = get_embedding_model(MODEL) q = "function that chunks code lines with overlap for semantic indexing" vec = next(emb.embed([q])) res = client.search( diff --git a/scripts/relevance_feedback.py b/scripts/relevance_feedback.py new file mode 100644 index 00000000..38add09a --- /dev/null +++ b/scripts/relevance_feedback.py @@ -0,0 +1,337 @@ +"""Shared storage and reindex reconciliation for relevance feedback.""" + +from __future__ import annotations + +import hashlib +import json +import os +import re +import time +from contextlib import contextmanager +from pathlib import Path +from typing import Any, Iterator + +try: + import fcntl # type: ignore +except Exception: # pragma: no cover + fcntl = None + + +RECENT_META_TTL = 3600 +RECENT_META_MAX = 4096 +RECENT_META_KEYS = ( + "result_id", + "target_id", + "impression_id", + "path", + "host_path", + "container_path", + "symbol", + "kind", + "repo", + "file_hash", + "symbol_content_hash", +) + + +def _weights_dir() -> Path: + return Path(os.environ.get("RERANKER_WEIGHTS_DIR", "/tmp/rerank_weights")) + + +def _safe_name(value: str) -> str: + return "".join(c if c.isalnum() or c in "-_" else "_" for c in str(value)) + + +def _target_path_key(path: str, repo: str = "") -> str: + """Normalize host/container paths to a stable repo-relative identity.""" + value = str(path or "").strip().replace("\\", "/") + if not value: + return "" + parts = [part for part in value.split("/") if part] + repo_value = str(repo or "").strip().replace("\\", "/").strip("/") + if repo_value: + for index in range(len(parts) - 1, -1, -1): + if parts[index] == repo_value: + parts = parts[index + 1 :] + break + return "/".join(parts) + + +def stable_target_key(*, repo: str = "", kind: str = "", symbol: str = "", path: str = "") -> str: + repo = str(repo or "").strip() + kind = str(kind or "").strip() + symbol = str(symbol or "").strip() + path_key = _target_path_key(path, repo) + if repo and symbol: + return f"symbol\x00{repo}\x00{path_key}\x00{kind}\x00{symbol}" + elif symbol: + return f"symbol\x00{path_key}\x00{kind}\x00{symbol}" + elif repo and path_key: + return f"file\x00{repo}\x00{path_key}" + elif path_key: + return f"file\x00{path_key}" + return "" + + +def stable_target_id(*, repo: str = "", kind: str = "", symbol: str = "", path: str = "") -> str: + key = stable_target_key(repo=repo, kind=kind, symbol=symbol, path=path) + if not key: + return "" + return hashlib.sha256(key.encode("utf-8")).hexdigest()[:12] + + +@contextmanager +def _locked_file(path: Path) -> Iterator[None]: + path.parent.mkdir(parents=True, exist_ok=True) + lock_path = path.with_suffix(path.suffix + ".lock") + with open(lock_path, "a+") as lock_file: + if fcntl is not None: + fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX) + try: + yield + finally: + if fcntl is not None: + fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN) + + +def _read_json(path: Path) -> dict: + try: + with open(path, "r", encoding="utf-8") as f: + data = json.load(f) + return data if isinstance(data, dict) else {} + except Exception: + return {} + + +def _write_json(path: Path, data: dict) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + tmp = path.with_suffix(path.suffix + f".tmp.{os.getpid()}") + try: + with open(tmp, "w", encoding="utf-8") as f: + json.dump(data, f) + tmp.replace(path) + finally: + tmp.unlink(missing_ok=True) + + +def _recent_path(collection: str) -> Path: + return _weights_dir() / f"{_safe_name(collection)}_recent_results.json" + + +def remember_recent_results(collection: str, results: list[dict]) -> None: + """Persist recent result metadata for cross-process hands-off ratings.""" + if not collection or not results: + return + try: + path = _recent_path(collection) + now = time.time() + expired_before = now - RECENT_META_TTL + with _locked_file(path): + data = _read_json(path) + entries = data.get("results") if isinstance(data.get("results"), dict) else {} + entries = { + rid: entry + for rid, entry in entries.items() + if isinstance(entry, dict) and float(entry.get("ts", 0) or 0) >= expired_before + } + for result in results: + rid = str(result.get("result_id") or "").strip() + if not rid: + continue + meta = {} + for key in RECENT_META_KEYS: + val = result.get(key) + if val is not None and str(val).strip(): + meta[key] = str(val).strip() + if meta: + entries[rid] = {"ts": now, "meta": meta} + if len(entries) > RECENT_META_MAX: + ranked = sorted( + entries.items(), + key=lambda item: float((item[1] or {}).get("ts", 0) or 0), + reverse=True, + ) + entries = dict(ranked[:RECENT_META_MAX]) + _write_json(path, {"updated_at": now, "results": entries}) + except OSError: + return + + +def enrich_recent_rating(collection: str, rating: dict) -> dict: + """Fill rating metadata from shared recent search results.""" + if not isinstance(rating, dict): + return {} + out = dict(rating) + rid = str(out.get("result_id") or "").strip() + if not collection or not rid: + return out + try: + data = _read_json(_recent_path(collection)) + except OSError: + return out + entry = (data.get("results") or {}).get(rid) + if not isinstance(entry, dict): + return out + if float(entry.get("ts", 0) or 0) < time.time() - RECENT_META_TTL: + return out + meta = entry.get("meta") if isinstance(entry.get("meta"), dict) else {} + for key, val in meta.items(): + out.setdefault(key, val) + return out + + +def _symbol_tokens(info: dict) -> set[str]: + content = str(info.get("content") or "") + return set(re.findall(r"[A-Za-z_][A-Za-z0-9_]{2,}", content)) + + +def _symbol_target(info: dict, *, repo: str, path: str) -> dict: + symbol = str(info.get("path") or info.get("symbol_path") or info.get("name") or "").strip() + kind = str(info.get("type") or "").strip() + target_id = stable_target_id(repo=repo, kind=kind, symbol=symbol, path=path) + return { + "target_id": target_id, + "path": path, + "container_path": path, + "repo": str(repo or ""), + "kind": kind, + "symbol": symbol, + "symbol_path": symbol, + "symbol_content_hash": str(info.get("content_hash") or ""), + } + + +def build_symbol_reconciliations( + old_symbols: dict, + new_symbols: dict, + *, + repo: str, + path: str, + split_min_overlap: float = 0.45, + split_min_coverage: float = 0.75, +) -> dict[str, list[dict]]: + """Map removed feedback targets to conservative rename/split successors.""" + old_symbols = old_symbols or {} + new_symbols = new_symbols or {} + new_by_symbol_kind: dict[tuple[str, str], list[dict]] = {} + new_by_hash_kind: dict[tuple[str, str], list[dict]] = {} + for info in new_symbols.values(): + kind = str(info.get("type") or "") + name = str( + info.get("path") + or info.get("symbol_path") + or info.get("name") + or "" + ) + content_hash = str(info.get("content_hash") or "") + new_by_symbol_kind.setdefault((kind, name), []).append(info) + if kind and content_hash: + new_by_hash_kind.setdefault((kind, content_hash), []).append(info) + + mappings: dict[str, list[dict]] = {} + for old_info in old_symbols.values(): + kind = str(old_info.get("type") or "") + name = str( + old_info.get("path") + or old_info.get("symbol_path") + or old_info.get("name") + or "" + ) + old_hash = str(old_info.get("content_hash") or "") + old_id = stable_target_id(repo=repo, kind=kind, symbol=name, path=path) + if not old_id: + continue + + # Same logical symbol persists; its stable target ID already survives edits. + if new_by_symbol_kind.get((kind, name)): + continue + + exact = new_by_hash_kind.get((kind, old_hash)) or [] + if len(exact) == 1: + mappings[old_id] = [ + { + "reason": "rename_exact_content", + "inheritance_weight": 1.0, + "target": _symbol_target(exact[0], repo=repo, path=path), + } + ] + continue + + old_tokens = _symbol_tokens(old_info) + if len(old_tokens) < 4: + continue + candidates = [] + covered: set[str] = set() + for new_info in new_symbols.values(): + if str(new_info.get("type") or "") != kind: + continue + overlap_tokens = old_tokens & _symbol_tokens(new_info) + overlap = len(overlap_tokens) / len(old_tokens) + if overlap >= split_min_overlap: + candidates.append((new_info, overlap, overlap_tokens)) + covered.update(overlap_tokens) + coverage = len(covered) / len(old_tokens) + if len(candidates) < 2 or coverage < split_min_coverage: + continue + total_overlap = sum(item[1] for item in candidates) or 1.0 + mappings[old_id] = [ + { + "reason": "split_token_coverage", + "inheritance_weight": round(overlap / total_overlap, 6), + "target": _symbol_target(info, repo=repo, path=path), + } + for info, overlap, _ in candidates + ] + return mappings + + +def reconcile_collection_weights( + collection: str, + mappings: dict[str, list[dict]], +) -> int: + """Migrate existing feedback weights to reconciled symbol targets.""" + if not collection or not mappings: + return 0 + path = _weights_dir() / f"{collection}_relevance.json" + if not path.exists(): + return 0 + migrated = 0 + with _locked_file(path): + data = _read_json(path) + results = data.get("results") if isinstance(data.get("results"), dict) else {} + for old_id, successors in mappings.items(): + old_entry = results.get(old_id) + if not isinstance(old_entry, dict): + continue + lineage = list(old_entry.get("lineage") or []) + for successor in successors: + target = successor.get("target") if isinstance(successor.get("target"), dict) else {} + new_id = str(target.get("target_id") or "") + inheritance = float(successor.get("inheritance_weight", 1.0) or 0) + if not new_id or inheritance <= 0: + continue + new_entry = dict(old_entry) + new_entry["target"] = target + new_entry["inheritance_weight"] = inheritance + new_entry["lineage"] = lineage + [ + { + "from_target_id": old_id, + "reason": successor.get("reason") or "reindex", + "at": time.time(), + } + ] + current = results.get(new_id) + if not isinstance(current, dict) or float(current.get("count", 0) or 0) <= float( + new_entry.get("count", 0) or 0 + ): + results[new_id] = new_entry + migrated += 1 + old_entry["superseded_by"] = [ + str((successor.get("target") or {}).get("target_id") or "") + for successor in successors + ] + if migrated: + data["updated_at"] = time.time() + data["results"] = results + _write_json(path, data) + return migrated diff --git a/scripts/relevance_trainer.py b/scripts/relevance_trainer.py new file mode 100644 index 00000000..d34de035 --- /dev/null +++ b/scripts/relevance_trainer.py @@ -0,0 +1,267 @@ +#!/usr/bin/env python3 +""" +Background Relevance Trainer. + +Consumes relevance feedback events (produced by rate_search_results MCP tool) +and computes per-collection ranking adjustments. Writes weights atomically so +the search path can apply them without restarting. + +Uses result-level relevance aggregation (upgradeable to logistic regression +when score features are added to events). + +Features: +- Reads feedback from NDJSON event log files (one per collection) +- Aggregates relevance scores per result_id +- Writes per-collection weight files atomically (write to .tmp, rename) +- Can run as a daemon or one-shot + +Usage: + # Run continuously (daemon mode) + python -m scripts.relevance_trainer --daemon + + # Process pending events once and exit + python -m scripts.relevance_trainer --once + + # Process specific collection + python -m scripts.relevance_trainer --collection my-repo --once +""" + +import argparse +import json +import os +import sys +import time +from pathlib import Path +from typing import Any, Dict, List, Optional + +# Event log configuration (shared with rerank_tools/events.py) +def _get_events_dir() -> Path: + return Path(os.environ.get("RERANK_EVENTS_DIR", "/tmp/rerank_events")) + + +def _get_weights_dir() -> Path: + return Path(os.environ.get("RERANKER_WEIGHTS_DIR", "/tmp/rerank_weights")) + + +def _get_poll_interval() -> int: + return int(os.environ.get("RELEVANCE_TRAINER_POLL_INTERVAL", "30")) + + +def _get_min_events() -> int: + return int(os.environ.get("RELEVANCE_TRAINER_MIN_EVENTS", "10")) + + +def list_collections_from_events() -> List[str]: + """Discover collections from event files on disk.""" + events_dir = _get_events_dir() + if not events_dir.exists(): + return [] + collections = set() + for f in events_dir.glob("events_*_*.ndjson"): + name = f.stem + parts = name.split("_", 1) + if len(parts) >= 2: + coll_hour = parts[1].rsplit("_", 1) + if len(coll_hour) >= 2 and coll_hour[1].isdigit(): + collections.add(coll_hour[0]) + return sorted(collections) + + +def read_feedback_events(collection: str) -> List[Dict[str, Any]]: + """Read all relevance feedback events for a collection.""" + events_dir = _get_events_dir() + if not events_dir.exists(): + return [] + safe_name = "".join(c if c.isalnum() or c in "-_" else "_" for c in collection) + pattern = f"events_{safe_name}_*.ndjson" + events = [] + for fp in sorted(events_dir.glob(pattern)): + try: + with open(fp, "r") as f: + for line in f: + line = line.strip() + if not line: + continue + try: + event = json.loads(line) + if event.get("type") == "relevance_feedback": + events.append(event) + except json.JSONDecodeError: + continue + except Exception: + continue + return events + + +def _rating_target(rating: Dict[str, Any]) -> Dict[str, Any]: + """Extract rehydratable target metadata from a feedback rating.""" + keep = ( + "target_id", + "impression_id", + "path", + "host_path", + "container_path", + "symbol", + "kind", + "repo", + "file_hash", + "symbol_content_hash", + ) + target = {} + for key in keep: + val = rating.get(key) + if val is not None and str(val).strip(): + target[key] = str(val).strip() + return target + + +def aggregate_ratings(events: List[Dict[str, Any]]) -> Dict[str, Dict[str, Any]]: + """Aggregate relevance scores per stable feedback target. + + Returns: {result_id: {total_relevance, count, avg_relevance, target}} + """ + results: Dict[str, List[int]] = {} + targets: Dict[str, Dict[str, Any]] = {} + for event in events: + for rating in event.get("ratings", []): + rid = rating.get("result_id", "").strip() + rel = rating.get("relevance") + if not rid or rel is None: + continue + try: + rel = int(rel) + except (ValueError, TypeError): + continue + if rel not in (0, 1, 2): + continue + if rid not in results: + results[rid] = [] + results[rid].append(rel) + target = _rating_target(rating) + if target: + targets[rid] = target + + aggregated: Dict[str, Dict[str, Any]] = {} + for rid, scores in results.items(): + total = sum(scores) + count = len(scores) + aggregated[rid] = { + "total_relevance": total, + "count": count, + "avg_relevance": round(total / count, 3), + } + if rid in targets: + aggregated[rid]["target"] = targets[rid] + return aggregated + + +def load_weights(collection: str) -> Dict[str, Any]: + """Load existing weights for a collection.""" + weights_dir = _get_weights_dir() + weights_dir.mkdir(parents=True, exist_ok=True) + weights_file = weights_dir / f"{collection}_relevance.json" + if weights_file.exists(): + try: + with open(weights_file, "r") as f: + return json.load(f) + except Exception: + return {} + return {} + + +def save_weights(collection: str, weights: Dict[str, Any]): + """Atomically save weights for a collection (write .tmp, rename).""" + weights_dir = _get_weights_dir() + weights_dir.mkdir(parents=True, exist_ok=True) + weights_file = weights_dir / f"{collection}_relevance.json" + tmp_file = weights_dir / f"{collection}_relevance.json.tmp" + try: + with open(tmp_file, "w") as f: + json.dump(weights, f) + tmp_file.rename(weights_file) + except Exception: + if tmp_file.exists(): + tmp_file.unlink(missing_ok=True) + raise + + +def process_collection(collection: str) -> Dict[str, Any]: + """Read events, aggregate ratings, merge with existing weights, save.""" + events = read_feedback_events(collection) + if len(events) < _get_min_events(): + return { + "collection": collection, + "events": len(events), + "skipped": True, + "reason": f"fewer than {_get_min_events()} events", + } + + new_weights = aggregate_ratings(events) + existing = load_weights(collection) + existing_results = existing.get("results") if isinstance(existing.get("results"), dict) else existing + if not isinstance(existing_results, dict): + existing_results = {} + merged = {**existing_results, **new_weights} + save_weights(collection, {"updated_at": time.time(), "results": merged}) + + return { + "collection": collection, + "events": len(events), + "new_entries": len(new_weights), + "total_entries": len(merged), + } + + +def run_daemon(collections: Optional[List[str]] = None, poll_interval: int = 30): + """Run continuously, polling for new events.""" + import logging + logging.basicConfig( + level=os.environ.get("LOG_LEVEL", "INFO").upper(), + format="[%(asctime)s] %(levelname)s %(message)s", + ) + logger = logging.getLogger("relevance_trainer") + + logger.info("Relevance trainer daemon started (poll=%ds)", poll_interval) + + while True: + targets = collections or list_collections_from_events() + for coll in targets: + try: + result = process_collection(coll) + if not result.get("skipped"): + logger.info( + "[%s] %d events → %d new / %d total entries", + coll, result["events"], result["new_entries"], result["total_entries"], + ) + except Exception: + logger.exception("[%s] Failed to process collection", coll) + time.sleep(poll_interval) + + +def main(): + parser = argparse.ArgumentParser(description="Relevance feedback trainer") + parser.add_argument("--daemon", action="store_true", help="Run continuously") + parser.add_argument("--once", action="store_true", help="Process pending events and exit") + parser.add_argument("--collection", type=str, help="Process specific collection") + parser.add_argument("--poll-interval", type=int, default=_get_poll_interval(), + help="Polling interval in seconds (daemon mode)") + args = parser.parse_args() + + if args.daemon: + collections = [args.collection] if args.collection else None + run_daemon(collections=collections, poll_interval=args.poll_interval) + elif args.once: + targets = [args.collection] if args.collection else list_collections_from_events() + for coll in targets: + result = process_collection(coll) + if result.get("skipped"): + print(f"[{coll}] Skipped: {result.get('reason')}") + else: + print(f"[{coll}] {result['events']} events → {result['total_entries']} entries") + else: + parser.print_help() + sys.exit(1) + + +if __name__ == "__main__": + main() diff --git a/scripts/remote_upload_client.py b/scripts/remote_upload_client.py index fcc7d6ba..c80bfaf6 100644 --- a/scripts/remote_upload_client.py +++ b/scripts/remote_upload_client.py @@ -46,6 +46,48 @@ # Configure logging logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) +_git_history_skip_log_key: Optional[str] = None + + +def _is_usable_delta_status(status: Any) -> bool: + if not isinstance(status, dict): + return False + state = str(status.get("status") or "").strip().lower() + return ( + bool(status.get("success")) and + "workspace_path" in status and + "collection_name" in status and + state in {"ready", "processing", "completed"} + ) + + +def _server_status_error_message(status: Any) -> str: + if isinstance(status, dict): + error = status.get("error") + if isinstance(error, dict): + msg = str(error.get("message") or "").strip() + if msg: + return msg + state = str(status.get("status") or "").strip() + if state: + return f"Server status is {state}" + return "Invalid server status response" + + +def _env_flag(name: str, default: bool) -> bool: + raw = os.environ.get(name) + if raw is None: + return default + return str(raw).strip().lower() in {"1", "true", "yes", "on"} + + +def _log_git_history_skip_once(reason: str, key: str) -> None: + global _git_history_skip_log_key + marker = f"{reason}:{key}" + if _git_history_skip_log_key == marker: + return + _git_history_skip_log_key = marker + logger.info("[git_history] skip (%s): %s", reason, key) DEFAULT_MAX_TEMP_CLEAN_ATTEMPTS = 3 DEFAULT_TEMP_CLEAN_SLEEP = 1.0 @@ -54,32 +96,27 @@ from scripts.workspace_state import ( get_cached_file_hash, set_cached_file_hash, - get_collection_name, _extract_repo_name_from_path, remove_cached_file, ) +from scripts.ingest.config import CODE_EXTS, EXTENSIONLESS_FILES -# Import existing hash function -import scripts.ingest_code as idx +def hash_id(text: str, path: str, start: int, end: int) -> int: + h = hashlib.sha1( + f"{path}:{start}-{end}\n{text}".encode("utf-8", errors="ignore") + ).hexdigest() + return int(h[:16], 16) -def _cache_missing_stats(file_hashes: Dict[str, Any]) -> Tuple[bool, int, int]: - """Return (is_stale, missing_count, checked_count) for cached paths.""" - if not file_hashes: - return (False, 0, 0) - missing = 0 - checked = 0 - for path_str in file_hashes.keys(): - try: - if not Path(path_str).exists(): - missing += 1 - except Exception: - missing += 1 - checked += 1 - if checked == 0: - return (False, 0, 0) - missing_ratio = missing / checked - return (missing_ratio >= 0.25, missing, checked) + +def _path_is_within_workspace(path_str: str, workspace_path: str) -> bool: + """Keep local cache entries scoped to the workspace being uploaded.""" + try: + path = Path(path_str).resolve() + workspace = Path(workspace_path).resolve() + return path == workspace or workspace in path.parents + except Exception: + return False def _find_git_root(start: Path) -> Optional[Path]: @@ -134,6 +171,24 @@ def _compute_logical_repo_id(workspace_path: str) -> str: return f"{prefix}{h}" +def _derive_metadata_root(workspace_path: str) -> Path: + """Infer host-side metadata root that corresponds to container `/work`.""" + try: + p = Path(workspace_path).resolve() + except Exception: + p = Path(workspace_path) + + if p.name == "dev-workspace": + return p.parent + if p.parent.name == "dev-workspace": + return p.parent.parent + if (p / ".codebase").exists(): + return p + if (p.parent / ".codebase").exists(): + return p.parent + return p.parent + + def _redact_emails(text: str) -> str: """Redact email addresses from commit messages for privacy.""" try: @@ -167,10 +222,12 @@ def _collect_git_history_for_workspace(workspace_path: str) -> Optional[Dict[str } if max_commits <= 0: + _log_git_history_skip_once("disabled", f"max_commits={max_commits}") return None root = _find_git_root(Path(workspace_path)) if not root: + _log_git_history_skip_once("no_repo", workspace_path) return None # Git history cache: avoid emitting identical manifests when HEAD/settings are unchanged @@ -204,6 +261,7 @@ def _collect_git_history_for_workspace(workspace_path: str) -> Optional[Dict[str cache = {} if current_head and cache.get("last_head") == current_head and cache.get("max_commits") == max_commits and str(cache.get("since") or "") == since: + _log_git_history_skip_once("cache_hit", f"head={current_head[:10]} since={since or '-'} max={max_commits}") return None base_head = "" @@ -254,12 +312,20 @@ def _collect_git_history_for_workspace(workspace_path: str) -> Optional[Dict[str errors="replace", ) if proc.returncode != 0 or not proc.stdout.strip(): + _log_git_history_skip_once( + "rev_list_empty", + f"head={current_head[:10] if current_head else '-'} rc={proc.returncode}", + ) return None commits = [l.strip() for l in proc.stdout.splitlines() if l.strip()] except Exception: return None if not commits: + _log_git_history_skip_once( + "no_commits", + f"head={current_head[:10] if current_head else '-'}", + ) return None if len(commits) > max_commits: commits = commits[:max_commits] @@ -333,6 +399,10 @@ def _collect_git_history_for_workspace(workspace_path: str) -> Optional[Dict[str continue if not records: + _log_git_history_skip_once( + "no_records", + f"commits={len(commits)} head={current_head[:10] if current_head else '-'}", + ) return None try: @@ -352,6 +422,14 @@ def _collect_git_history_for_workspace(workspace_path: str) -> Optional[Dict[str "since": since, "commits": records, } + logger.info( + "[git_history] prepared manifest mode=%s commits=%d head=%s prev=%s base=%s", + manifest["mode"], + len(records), + (current_head[:10] if current_head else "-"), + (prev_head[:10] if prev_head else "-"), + (base_head[:10] if base_head else "-"), + ) # Update git history cache with the HEAD and settings used for this manifest try: @@ -370,7 +448,12 @@ def _collect_git_history_for_workspace(workspace_path: str) -> Optional[Dict[str return manifest -def _load_local_cache_file_hashes(workspace_path: str, repo_name: Optional[str]) -> Dict[str, str]: +def _load_local_cache_file_hashes( + workspace_path: str, + repo_name: Optional[str], + *, + metadata_root: Optional[str] = None, +) -> Dict[str, str]: """Best-effort read of the local cache.json file_hashes map. This mirrors the layout used by workspace_state without introducing new @@ -378,7 +461,13 @@ def _load_local_cache_file_hashes(workspace_path: str, repo_name: Optional[str]) lookups still go through get_cached_file_hash. """ try: - base = Path(os.environ.get("WORKSPACE_PATH") or workspace_path).resolve() + base = Path( + metadata_root + or os.environ.get("CTXCE_METADATA_ROOT") + or os.environ.get("WATCH_ROOT") + or os.environ.get("WORKSPACE_PATH") + or workspace_path + ).resolve() multi_repo = os.environ.get("MULTI_REPO_MODE", "0").strip().lower() in {"1", "true", "yes", "on"} if multi_repo and repo_name: cache_path = base / ".codebase" / "repos" / repo_name / "cache.json" @@ -395,29 +484,40 @@ def _load_local_cache_file_hashes(workspace_path: str, repo_name: Optional[str]) file_hashes = data.get("file_hashes", {}) if not isinstance(file_hashes, dict): return {} - is_stale, missing, checked = _cache_missing_stats(file_hashes) - if is_stale: + scoped_hashes = { + str(path): value + for path, value in file_hashes.items() + if _path_is_within_workspace(str(path), workspace_path) + } + ignored_count = len(file_hashes) - len(scoped_hashes) + if ignored_count: logger.warning( - "[remote_upload] Detected stale local cache (%d/%d missing); clearing %s", - missing, - checked, + "[remote_upload] Ignoring %d local cache entries outside workspace %s: %s", + ignored_count, + workspace_path, cache_path, ) - try: - cache_path.unlink(missing_ok=True) # type: ignore[arg-type] - except TypeError: - try: - cache_path.unlink() - except Exception: - pass - except Exception: - pass - return {} - return file_hashes + return scoped_hashes except Exception: return {} +def get_all_cached_paths( + repo_name: Optional[str] = None, + metadata_root: Optional[str] = None, + workspace_path: Optional[str] = None, +) -> List[str]: + """Return cached file paths from the local workspace cache.""" + effective_workspace = workspace_path or os.environ.get("WORKSPACE_PATH") or os.getcwd() + return list( + _load_local_cache_file_hashes( + effective_workspace, + repo_name, + metadata_root=metadata_root, + ).keys() + ) + + class RemoteUploadClient: """Client for uploading delta bundles to remote server.""" @@ -451,33 +551,29 @@ def _translate_to_container_path(self, host_path: str) -> str: return host_path.replace('\\', '/').replace(':', '') - def __init__(self, upload_endpoint: str, workspace_path: str, collection_name: str, + def __init__(self, upload_endpoint: str, workspace_path: str, collection_name: Optional[str] = None, max_retries: int = 3, timeout: int = 30, metadata_path: Optional[str] = None, logical_repo_id: Optional[str] = None): """Initialize remote upload client.""" self.upload_endpoint = upload_endpoint.rstrip('/') self.workspace_path = workspace_path + self.metadata_root = str(_derive_metadata_root(workspace_path)) self.collection_name = collection_name self.max_retries = max_retries self.timeout = timeout self.temp_dir = None self.logical_repo_id = logical_repo_id - # Set environment variables for cache functions - os.environ["WORKSPACE_PATH"] = workspace_path + from scripts.workspace_state import _extract_repo_name_from_path - # Get repo name for cache operations - try: - from scripts.workspace_state import _extract_repo_name_from_path - self.repo_name = _extract_repo_name_from_path(workspace_path) - # Fallback to directory name if repo detection fails (for non-git repos) - if not self.repo_name: - self.repo_name = Path(workspace_path).name - except ImportError: + self.repo_name = _extract_repo_name_from_path(workspace_path) + # Fallback to directory name if repo detection fails (for non-git repos) + if not self.repo_name: self.repo_name = Path(workspace_path).name # In-memory stat cache to avoid rehashing unchanged files on every watch iteration self._stat_cache: Dict[str, Tuple[int, int]] = {} + self._content_hash_cache: Dict[str, Tuple[int, int, str]] = {} # Setup HTTP session with simple retry self.session = requests.Session() @@ -485,6 +581,122 @@ def __init__(self, upload_endpoint: str, workspace_path: str, collection_name: s adapter = HTTPAdapter(max_retries=retry_strategy) self.session.mount("http://", adapter) self.session.mount("https://", adapter) + self.last_upload_result: Dict[str, Any] = {"outcome": "idle"} + self._last_plan_payload: Optional[Dict[str, Any]] = None + self._last_expected_hashes: Dict[str, str] = {} + + def _get_cached_file_hash(self, file_path: str) -> str: + return get_cached_file_hash( + file_path, + self.repo_name, + metadata_root=self.metadata_root, + ) + + def _set_cached_file_hash(self, file_path: str, file_hash: str) -> None: + set_cached_file_hash( + file_path, + file_hash, + self.repo_name, + metadata_root=self.metadata_root, + ) + + def _finalize_file_hash(self, path: Path, expected_hash: Optional[str] = None) -> bool: + abs_path = str(path.resolve()) + stat = path.stat() + current_hash = hashlib.sha1(path.read_bytes()).hexdigest() + expected = str(expected_hash or "").strip().lower() + if ":" in expected: + expected = expected.partition(":")[2].strip() + if expected and current_hash != expected: + self._content_hash_cache.pop(abs_path, None) + return False + self._set_cached_file_hash(abs_path, current_hash) + self._stat_cache[abs_path] = ( + int(getattr(stat, "st_mtime_ns", int(stat.st_mtime * 1e9))), + stat.st_size, + ) + self._content_hash_cache[abs_path] = ( + int(getattr(stat, "st_mtime_ns", int(stat.st_mtime * 1e9))), + stat.st_size, + current_hash, + ) + return True + + def _remove_cached_file(self, file_path: str) -> None: + remove_cached_file( + file_path, + self.repo_name, + metadata_root=self.metadata_root, + ) + + def _get_all_cached_paths(self) -> List[str]: + return get_all_cached_paths( + self.repo_name, + metadata_root=self.metadata_root, + workspace_path=self.workspace_path, + ) + + def _set_last_upload_result(self, outcome: str, **details: Any) -> Dict[str, Any]: + result: Dict[str, Any] = {"outcome": outcome} + result.update(details) + self.last_upload_result = result + return result + + def _read_current_file_hash(self, path: Path) -> Tuple[str, int]: + """Read a file hash once per current size/mtime pair during one pass.""" + stat = path.stat() + mtime_ns = int(getattr(stat, "st_mtime_ns", int(stat.st_mtime * 1e9))) + size = int(stat.st_size) + key = str(path.resolve()) + cached = self._content_hash_cache.get(key) + if cached and cached[:2] == (mtime_ns, size): + return cached[2], size + file_hash = hashlib.sha1(path.read_bytes()).hexdigest() + self._content_hash_cache[key] = (mtime_ns, size, file_hash) + return file_hash, size + + def log_watch_upload_result(self) -> None: + outcome = str((self.last_upload_result or {}).get("outcome") or "") + if outcome == "skipped_by_plan": + logger.info("[watch] No upload needed after plan") + elif outcome == "queued": + logger.info("[watch] Upload request accepted; server processing asynchronously") + elif outcome == "uploaded": + logger.info("[watch] Successfully uploaded changes") + elif outcome == "no_changes": + logger.info("[watch] No meaningful changes to upload") + else: + logger.info("[watch] Upload handling completed") + + def _finalize_successful_changes(self, changes: Dict[str, List]) -> None: + expected_hashes = self._last_expected_hashes + for path in changes.get("created", []) + changes.get("updated", []): + try: + rel_path = path.relative_to(Path(self.workspace_path)).as_posix() + self._finalize_file_hash(path, expected_hashes.get(rel_path)) + except Exception: + continue + for path in changes.get("deleted", []): + try: + abs_path = str(path.resolve()) + self._remove_cached_file(abs_path) + self._stat_cache.pop(abs_path, None) + self._content_hash_cache.pop(abs_path, None) + except Exception: + continue + for source_path, dest_path in changes.get("moved", []): + try: + source_abs_path = str(source_path.resolve()) + self._remove_cached_file(source_abs_path) + self._stat_cache.pop(source_abs_path, None) + self._content_hash_cache.pop(source_abs_path, None) + except Exception: + pass + try: + dest_rel_path = dest_path.relative_to(Path(self.workspace_path)).as_posix() + self._finalize_file_hash(dest_path, expected_hashes.get(dest_rel_path)) + except Exception: + continue def __enter__(self): """Context manager entry.""" @@ -505,7 +717,7 @@ def get_mapping_summary(self) -> Dict[str, Any]: container_path = self._translate_to_container_path(self.workspace_path) return { "repo_name": self.repo_name, - "collection_name": self.collection_name, + "collection_name": self.collection_name or "", "source_path": self.workspace_path, "container_path": container_path, "upload_endpoint": self.upload_endpoint, @@ -520,6 +732,65 @@ def log_mapping_summary(self) -> None: logger.info(f" source_path: {info['source_path']}") logger.info(f" container_path: {info['container_path']}") + def _excluded_dirnames(self) -> frozenset: + # Keep in sync with standalone_upload_client exclusions. + # NOTE: This caches the exclusion set per RemoteUploadClient instance. + # Runtime changes to DEV_REMOTE_MODE/REMOTE_UPLOAD_MODE won't be reflected + # until a new client is created (typically via process restart), which is + # acceptable for the upload client use case. + cached = getattr(self, "_excluded_dirnames_cache", None) + if cached is not None: + return cached + excluded = { + "node_modules", "vendor", "dist", "build", "target", "out", + ".git", ".hg", ".svn", ".vscode", ".idea", ".venv", "venv", + "__pycache__", ".pytest_cache", ".mypy_cache", ".cache", + ".context-engine", ".context-engine-uploader", ".codebase", + } + dev_remote = os.environ.get("DEV_REMOTE_MODE") == "1" or os.environ.get("REMOTE_UPLOAD_MODE") == "development" + if dev_remote: + excluded.add("dev-workspace") + cached = frozenset(excluded) + self._excluded_dirnames_cache = cached + return cached + + def _is_ignored_path(self, path: Path) -> bool: + """Return True when path is outside workspace or under excluded dirs.""" + try: + workspace_root = Path(self.workspace_path).resolve() + rel = path.resolve().relative_to(workspace_root) + except Exception: + return True + + dir_parts = set(rel.parts[:-1]) if len(rel.parts) > 1 else set() + if dir_parts & self._excluded_dirnames(): + return True + # Ignore hidden directories anywhere under the workspace, but allow + # extensionless dotfiles like `.gitignore` that we explicitly support. + if any(p.startswith(".") for p in rel.parts[:-1]): + return True + try: + extensionless = set((EXTENSIONLESS_FILES or {}).keys()) + except Exception: + extensionless = set() + if rel.name.startswith(".") and rel.name.lower() not in extensionless: + return True + return False + + def _is_watchable_path(self, path: Path) -> bool: + """Return True when a filesystem event path is eligible for upload processing.""" + if self._is_ignored_path(path): + return False + suffix = path.suffix.lower() + if CODE_EXTS.get(suffix, "unknown") != "unknown": + return True + name = path.name.lower() + try: + extensionless_names = {k.lower() for k in (EXTENSIONLESS_FILES or {}).keys()} + except Exception: + extensionless_names = set() + return name in extensionless_names or name.startswith("dockerfile") + def _get_temp_bundle_dir(self) -> Path: """Get or create temporary directory for bundle creation.""" if not self.temp_dir: @@ -547,6 +818,19 @@ def detect_file_changes(self, changed_paths: List[Path]) -> Dict[str, List]: } for path in changed_paths: + if self._is_ignored_path(path): + try: + abs_path = str(path.resolve()) + except Exception: + continue + cached_hash = self._get_cached_file_hash(abs_path) + if cached_hash: + changes["deleted"].append(path) + try: + self._stat_cache.pop(abs_path, None) + except Exception: + pass + continue # Resolve to an absolute path for stable cache keys try: abs_path = str(path.resolve()) @@ -554,7 +838,7 @@ def detect_file_changes(self, changed_paths: List[Path]) -> Dict[str, List]: # Skip paths that cannot be resolved continue - cached_hash = get_cached_file_hash(abs_path, self.repo_name) + cached_hash = self._get_cached_file_hash(abs_path) if not path.exists(): # File was deleted @@ -588,9 +872,7 @@ def detect_file_changes(self, changed_paths: List[Path]) -> Dict[str, List]: # Stat changed or no prior entry – hash content to classify change try: - with open(path, 'rb') as f: - content = f.read() - current_hash = hashlib.sha1(content).hexdigest() + current_hash, _ = self._read_current_file_hash(path) except Exception: # Skip files that can't be read continue @@ -610,8 +892,6 @@ def detect_file_changes(self, changed_paths: List[Path]) -> Dict[str, List]: self._stat_cache[abs_path] = (getattr(stat, "st_mtime_ns", int(stat.st_mtime * 1e9)), stat.st_size) except Exception: pass - set_cached_file_hash(abs_path, current_hash, self.repo_name) - # Detect moves by looking for files with same content hash # but different paths (requires additional tracking) changes["moved"] = self._detect_moves(changes["created"], changes["deleted"]) @@ -636,7 +916,7 @@ def _detect_moves(self, created_files: List[Path], deleted_files: List[Path]) -> for deleted_path in deleted_files: try: # Try to get cached hash first, fallback to file content - cached_hash = get_cached_file_hash(str(deleted_path), self.repo_name) + cached_hash = self._get_cached_file_hash(str(deleted_path)) if cached_hash: deleted_hashes[cached_hash] = deleted_path continue @@ -720,7 +1000,7 @@ def create_delta_bundle( # Get file info stat = path.stat() - language = idx.CODE_EXTS.get(path.suffix.lower(), "unknown") + language = CODE_EXTS.get(path.suffix.lower(), "unknown") operation = { "operation": "created", @@ -729,7 +1009,7 @@ def create_delta_bundle( "absolute_path": str(path.resolve()), "size_bytes": stat.st_size, "content_hash": content_hash, - "file_hash": f"sha1:{idx.hash_id(content.decode('utf-8', errors='ignore'), rel_path, 1, len(content.splitlines()))}", + "file_hash": f"sha1:{hash_id(content.decode('utf-8', errors='ignore'), rel_path, 1, len(content.splitlines()))}", "modified_time": datetime.fromtimestamp(stat.st_mtime).isoformat(), "language": language } @@ -749,7 +1029,7 @@ def create_delta_bundle( content = f.read() file_hash = hashlib.sha1(content).hexdigest() content_hash = f"sha1:{file_hash}" - previous_hash = get_cached_file_hash(str(path.resolve()), self.repo_name) + previous_hash = self._get_cached_file_hash(str(path.resolve())) # Write file to bundle bundle_file_path = files_dir / "updated" / rel_path @@ -758,7 +1038,7 @@ def create_delta_bundle( # Get file info stat = path.stat() - language = idx.CODE_EXTS.get(path.suffix.lower(), "unknown") + language = CODE_EXTS.get(path.suffix.lower(), "unknown") operation = { "operation": "updated", @@ -768,7 +1048,7 @@ def create_delta_bundle( "size_bytes": stat.st_size, "content_hash": content_hash, "previous_hash": f"sha1:{previous_hash}" if previous_hash else None, - "file_hash": f"sha1:{idx.hash_id(content.decode('utf-8', errors='ignore'), rel_path, 1, len(content.splitlines()))}", + "file_hash": f"sha1:{hash_id(content.decode('utf-8', errors='ignore'), rel_path, 1, len(content.splitlines()))}", "modified_time": datetime.fromtimestamp(stat.st_mtime).isoformat(), "language": language } @@ -797,7 +1077,7 @@ def create_delta_bundle( # Get file info stat = dest_path.stat() - language = idx.CODE_EXTS.get(dest_path.suffix.lower(), "unknown") + language = CODE_EXTS.get(dest_path.suffix.lower(), "unknown") operation = { "operation": "moved", @@ -809,7 +1089,7 @@ def create_delta_bundle( "source_absolute_path": str(source_path.resolve()), "size_bytes": stat.st_size, "content_hash": content_hash, - "file_hash": f"sha1:{idx.hash_id(content.decode('utf-8', errors='ignore'), dest_rel_path, 1, len(content.splitlines()))}", + "file_hash": f"sha1:{hash_id(content.decode('utf-8', errors='ignore'), dest_rel_path, 1, len(content.splitlines()))}", "modified_time": datetime.fromtimestamp(stat.st_mtime).isoformat(), "language": language } @@ -825,7 +1105,7 @@ def create_delta_bundle( for path in changes["deleted"]: rel_path = path.relative_to(Path(self.workspace_path)).as_posix() try: - previous_hash = get_cached_file_hash(str(path.resolve()), self.repo_name) + previous_hash = self._get_cached_file_hash(str(path.resolve())) operation = { "operation": "deleted", @@ -835,17 +1115,10 @@ def create_delta_bundle( "previous_hash": f"sha1:{previous_hash}" if previous_hash else None, "file_hash": None, "modified_time": datetime.now().isoformat(), - "language": idx.CODE_EXTS.get(path.suffix.lower(), "unknown") + "language": CODE_EXTS.get(path.suffix.lower(), "unknown") } operations.append(operation) - # Once a delete operation has been recorded, drop the cache entry - # so subsequent scans do not keep re-reporting the same deletion. - try: - remove_cached_file(str(path.resolve()), self.repo_name) - except Exception: - pass - except Exception as e: print(f"[bundle_create] Error processing deleted file {path}: {e}") continue @@ -855,7 +1128,6 @@ def create_delta_bundle( "version": "1.0", "bundle_id": bundle_id, "workspace_path": self.workspace_path, - "collection_name": self.collection_name, "created_at": created_at, # CLI is stateless - server handles sequence numbers "sequence_number": None, # Server will assign @@ -907,8 +1179,292 @@ def create_delta_bundle( with tarfile.open(bundle_path, "w:gz") as tar: tar.add(temp_path, arcname=f"{bundle_id}") + self._last_expected_hashes.update(file_hashes) return str(bundle_path), manifest + def _build_plan_payload(self, changes: Dict[str, List]) -> Dict[str, Any]: + created_at = datetime.now().isoformat() + bundle_id = str(uuid.uuid4()) + operations: List[Dict[str, Any]] = [] + file_hashes: Dict[str, str] = {} + total_size = 0 + + for path in changes["created"]: + rel_path = path.relative_to(Path(self.workspace_path)).as_posix() + try: + file_hash, size = self._read_current_file_hash(path) + operations.append( + { + "operation": "created", + "path": rel_path, + "size_bytes": size, + "content_hash": f"sha1:{file_hash}", + "language": CODE_EXTS.get(path.suffix.lower(), "unknown"), + } + ) + file_hashes[rel_path] = f"sha1:{file_hash}" + total_size += size + except Exception as e: + logger.warning("[remote_upload] Failed to prepare created plan entry for %s: %s", path, e) + + for path in changes["updated"]: + rel_path = path.relative_to(Path(self.workspace_path)).as_posix() + try: + file_hash, size = self._read_current_file_hash(path) + previous_hash = self._get_cached_file_hash(str(path.resolve())) + operations.append( + { + "operation": "updated", + "path": rel_path, + "size_bytes": size, + "content_hash": f"sha1:{file_hash}", + "previous_hash": f"sha1:{previous_hash}" if previous_hash else None, + "language": CODE_EXTS.get(path.suffix.lower(), "unknown"), + } + ) + file_hashes[rel_path] = f"sha1:{file_hash}" + total_size += size + except Exception as e: + logger.warning("[remote_upload] Failed to prepare updated plan entry for %s: %s", path, e) + + for source_path, dest_path in changes["moved"]: + dest_rel_path = dest_path.relative_to(Path(self.workspace_path)).as_posix() + source_rel_path = source_path.relative_to(Path(self.workspace_path)).as_posix() + try: + file_hash, size = self._read_current_file_hash(dest_path) + operations.append( + { + "operation": "moved", + "path": dest_rel_path, + "source_path": source_rel_path, + "size_bytes": size, + "content_hash": f"sha1:{file_hash}", + "language": CODE_EXTS.get(dest_path.suffix.lower(), "unknown"), + } + ) + file_hashes[dest_rel_path] = f"sha1:{file_hash}" + total_size += size + except Exception as e: + logger.warning( + "[remote_upload] Failed to prepare moved plan entry for %s -> %s: %s", + source_path, + dest_path, + e, + ) + + for path in changes["deleted"]: + rel_path = path.relative_to(Path(self.workspace_path)).as_posix() + try: + previous_hash = self._get_cached_file_hash(str(path.resolve())) + operations.append( + { + "operation": "deleted", + "path": rel_path, + "previous_hash": f"sha1:{previous_hash}" if previous_hash else None, + "language": CODE_EXTS.get(path.suffix.lower(), "unknown"), + } + ) + except Exception as e: + logger.warning("[remote_upload] Failed to prepare deleted plan entry for %s: %s", path, e) + + manifest = { + "version": "1.0", + "bundle_id": bundle_id, + "workspace_path": self.workspace_path, + "created_at": created_at, + "sequence_number": None, + "parent_sequence": None, + "operations": { + "created": len(changes["created"]), + "updated": len(changes["updated"]), + "deleted": len(changes["deleted"]), + "moved": len(changes["moved"]), + }, + "total_files": len(operations), + "total_size_bytes": total_size, + "compression": "gzip", + "encoding": "utf-8", + } + return { + "manifest": manifest, + "operations": operations, + "file_hashes": file_hashes, + } + + def _plan_delta_upload(self, changes: Dict[str, List]) -> Optional[Dict[str, Any]]: + if not _env_flag("CTXCE_REMOTE_UPLOAD_PLAN_ENABLED", True): + return None + try: + payload = self._build_plan_payload(changes) + self._last_plan_payload = payload + self._last_expected_hashes = dict(payload.get("file_hashes", {})) + # Indexed hashes are server-owned; the client submits candidates + # and uses only the returned plan. + data = { + "workspace_path": self._translate_to_container_path(self.workspace_path), + "source_path": self.workspace_path, + "logical_repo_id": _compute_logical_repo_id(self.workspace_path), + "manifest": payload["manifest"], + "operations": payload["operations"], + "file_hashes": payload["file_hashes"], + } + sess = get_auth_session(self.upload_endpoint) + if sess: + data["session"] = sess + if getattr(self, "logical_repo_id", None): + data["logical_repo_id"] = self.logical_repo_id + + response = self.session.post( + f"{self.upload_endpoint}/api/v1/delta/plan", + json=data, + timeout=min(self.timeout, 60), + ) + if response.status_code in {404, 405}: + logger.info("[remote_upload] Plan endpoint unavailable; falling back to full bundle upload") + return None + response.raise_for_status() + body = response.json() + if not body.get("success", False): + logger.warning("[remote_upload] Plan request failed; falling back: %s", body.get("error")) + return None + return body + except Exception as e: + logger.warning("[remote_upload] Plan request failed; falling back to full bundle upload: %s", e) + return None + + def _build_apply_only_payload(self, changes: Dict[str, List], plan: Dict[str, Any]) -> Dict[str, Any]: + payload = self._last_plan_payload or self._build_plan_payload(changes) + needed = plan.get("needed_files", {}) if isinstance(plan, dict) else {} + created_needed = set(needed.get("created", []) or []) + updated_needed = set(needed.get("updated", []) or []) + moved_needed = set(needed.get("moved", []) or []) + + # Check if ALL operations are hash-matched (nothing needs content at all) + # This happens when all needed_files lists are empty and there are no actual changes requiring content + has_changes_needing_content = bool(created_needed or updated_needed or moved_needed) + has_deletes = bool(changes.get("deleted", [])) + + # Only skip apply-only if there are NO operations needing content AND NO deletes + if not has_changes_needing_content and not has_deletes: + return { + "manifest": payload.get("manifest", {}), + "operations": [], + "file_hashes": {}, + } + + filtered_ops: List[Dict[str, Any]] = [] + filtered_hashes: Dict[str, str] = {} + for operation in payload.get("operations", []): + op_type = str(operation.get("operation") or "") + rel_path = str(operation.get("path") or "") + # Determine if this operation needs content (only those skip filtered_hashes) + needs_content = ( + (op_type == "created" and rel_path in created_needed) + or (op_type == "updated" and rel_path in updated_needed) + or (op_type == "moved" and rel_path in moved_needed) + ) + if needs_content: + # Skip operations that need content - they'll be uploaded separately + continue + # IMPORTANT: server-side apply_delta_operations() only accepts "deleted" and "moved" + # operations. Hash-matched "created" and "updated" operations must NOT be routed + # through apply_ops since the server will reject them. + if op_type not in {"deleted", "moved"}: + continue + # Preserve all other operations so server advances state + filtered_ops.append(operation) + # Include hash for non-deleted operations + if op_type != "deleted": + hash_value = payload.get("file_hashes", {}).get(rel_path) + if hash_value: + filtered_hashes[rel_path] = hash_value + return { + "manifest": payload.get("manifest", {}), + "operations": filtered_ops, + "file_hashes": filtered_hashes, + } + + def _apply_operations_without_content(self, changes: Dict[str, List], plan: Dict[str, Any]) -> Optional[bool]: + payload = self._build_apply_only_payload(changes, plan) + operations = payload.get("operations", []) + if not operations: + return None + try: + data = { + "workspace_path": self._translate_to_container_path(self.workspace_path), + "source_path": self.workspace_path, + "logical_repo_id": _compute_logical_repo_id(self.workspace_path), + "manifest": payload["manifest"], + "operations": operations, + "file_hashes": payload["file_hashes"], + } + sess = get_auth_session(self.upload_endpoint) + if sess: + data["session"] = sess + if getattr(self, "logical_repo_id", None): + data["logical_repo_id"] = self.logical_repo_id + + logger.info( + "[remote_upload] Applying metadata-only operations without bundle: deleted=%s moved=%s", + sum(1 for op in operations if op.get("operation") == "deleted"), + sum(1 for op in operations if op.get("operation") == "moved"), + ) + response = self.session.post( + f"{self.upload_endpoint}/api/v1/delta/apply_ops", + json=data, + timeout=min(self.timeout, 60), + ) + if response.status_code in {404, 405}: + logger.info("[remote_upload] apply_ops endpoint unavailable; falling back to bundle upload") + return None + response.raise_for_status() + body = response.json() + if not body.get("success", False): + logger.warning("[remote_upload] apply_ops failed; falling back to bundle upload: %s", body.get("error")) + return None + processed_ops = body.get("processed_operations") or {} + self._set_last_upload_result( + "uploaded", + bundle_id=body.get("bundle_id"), + sequence_number=body.get("sequence_number"), + processed_operations=processed_ops, + ) + logger.info( + "[remote_upload] Metadata-only operations applied: %s", + processed_ops, + ) + return True + except Exception as e: + logger.warning("[remote_upload] apply_ops failed; falling back to bundle upload: %s", e) + return None + + def _filter_changes_by_plan(self, changes: Dict[str, List], plan: Dict[str, Any]) -> Dict[str, List]: + needed = plan.get("needed_files", {}) if isinstance(plan, dict) else {} + created_needed = set(needed.get("created", []) or []) + updated_needed = set(needed.get("updated", []) or []) + moved_needed = set(needed.get("moved", []) or []) + + filtered_created = [ + path for path in changes["created"] + if path.relative_to(Path(self.workspace_path)).as_posix() in created_needed + ] + filtered_updated = [ + path for path in changes["updated"] + if path.relative_to(Path(self.workspace_path)).as_posix() in updated_needed + ] + filtered_moved = [ + (source_path, dest_path) + for source_path, dest_path in changes["moved"] + if dest_path.relative_to(Path(self.workspace_path)).as_posix() in moved_needed + ] + return { + "created": filtered_created, + "updated": filtered_updated, + "deleted": list(changes["deleted"]), + "moved": filtered_moved, + "unchanged": [], + } + def upload_bundle(self, bundle_path: str, manifest: Dict[str, Any]) -> Dict[str, Any]: """ Upload delta bundle to remote server with exponential backoff retry. @@ -937,18 +1493,13 @@ def upload_bundle(self, bundle_path: str, manifest: Dict[str, Any]) -> Dict[str, # Check bundle size (server-side enforcement) bundle_size = os.path.getsize(bundle_path) - files = { - "bundle": open(bundle_path, "rb"), - } data = { "workspace_path": self._translate_to_container_path(self.workspace_path), - "collection_name": self.collection_name, "sequence_number": manifest.get("sequence_number"), "force": False, "source_path": self.workspace_path, "logical_repo_id": _compute_logical_repo_id(self.workspace_path), } - sess = get_auth_session(self.upload_endpoint) if sess: data["session"] = sess @@ -958,12 +1509,13 @@ def upload_bundle(self, bundle_path: str, manifest: Dict[str, Any]) -> Dict[str, logger.info(f"[remote_upload] Uploading bundle {manifest['bundle_id']} (size: {bundle_size} bytes)") - response = self.session.post( - f"{self.upload_endpoint}/api/v1/delta/upload", - files=files, - data=data, - timeout=(10, self.timeout) - ) + with open(bundle_path, "rb") as bundle_file: + response = self.session.post( + f"{self.upload_endpoint}/api/v1/delta/upload", + files={"bundle": bundle_file}, + data=data, + timeout=(10, self.timeout) + ) result = None try: @@ -1015,29 +1567,17 @@ def upload_bundle(self, bundle_path: str, manifest: Dict[str, Any]) -> Dict[str, last_error = {"success": False, "error": {"code": "TIMEOUT_ERROR", "message": f"Upload timeout: {str(e)}"}} logger.warning(f"[remote_upload] Upload read timeout on attempt {attempt + 1}: {e}") - # After read timeout, poll to check if server processed the bundle - logger.info(f"[remote_upload] Read timeout occurred, polling server to check if bundle was processed...") - poll_result = self._poll_after_timeout(manifest) - if poll_result.get("success"): - logger.info(f"[remote_upload] Server confirmed processing of bundle {manifest['bundle_id']} after timeout") - return poll_result - - logger.warning(f"[remote_upload] Server did not process bundle after timeout, proceeding with failure") - break + # A timeout is retried normally. The server-side hash plan makes + # a replay safe when the request was accepted before the timeout. + continue except requests.exceptions.Timeout as e: last_error = {"success": False, "error": {"code": "TIMEOUT_ERROR", "message": f"Upload timeout: {str(e)}"}} logger.warning(f"[remote_upload] Upload timeout on attempt {attempt + 1}: {e}") - # For generic timeout, also try polling - logger.info(f"[remote_upload] Timeout occurred, polling server to check if bundle was processed...") - poll_result = self._poll_after_timeout(manifest) - if poll_result.get("success"): - logger.info(f"[remote_upload] Server confirmed processing of bundle {manifest['bundle_id']} after timeout") - return poll_result - - logger.warning(f"[remote_upload] Server did not process bundle after timeout, proceeding with failure") - break + # A timeout is retried normally. The server-side hash plan makes + # a replay safe when the request was accepted before the timeout. + continue except requests.exceptions.ConnectionError as e: last_error = {"success": False, "error": {"code": "CONNECTION_ERROR", "message": f"Connection error: {str(e)}"}} @@ -1061,87 +1601,6 @@ def upload_bundle(self, bundle_path: str, manifest: Dict[str, Any]) -> Dict[str, } } - def _poll_after_timeout(self, manifest: Dict[str, Any]) -> Dict[str, Any]: - """ - Poll server status after a timeout to check if bundle was processed. - - Args: - manifest: Bundle manifest containing sequence information - - Returns: - Dictionary indicating success if bundle was processed - """ - try: - # Get current server status to know the expected sequence - status = self.get_server_status() - if not status.get("success"): - return {"success": False, "error": status.get("error", {"code": "UNKNOWN", "message": "Failed to get status"})} - - current_sequence = status.get("last_sequence", 0) - expected_sequence = manifest.get("sequence", current_sequence + 1) - - logger.info(f"[remote_upload] Current server sequence: {current_sequence}, expected: {expected_sequence}") - - # If server is already at expected sequence, bundle was processed - if current_sequence >= expected_sequence: - return { - "success": True, - "message": f"Bundle processed (server at sequence {current_sequence})", - "sequence": current_sequence, - } - - # Poll window is configurable via REMOTE_UPLOAD_POLL_MAX_SECS (seconds). - # Values <= 0 mean "no timeout" (poll until success or process exit). - try: - max_poll_time = int(os.environ.get("REMOTE_UPLOAD_POLL_MAX_SECS", "300")) - except Exception: - max_poll_time = 300 - poll_interval = 5 - start_time = time.time() - - while True: - elapsed = time.time() - start_time - if max_poll_time > 0 and elapsed >= max_poll_time: - logger.warning( - f"[remote_upload] Polling timed out after {int(elapsed)}s (limit={max_poll_time}s), bundle was not confirmed as processed" - ) - return { - "success": False, - "error": { - "code": "POLL_TIMEOUT", - "message": f"Bundle not confirmed processed after polling for {int(elapsed)}s (limit={max_poll_time}s)", - }, - } - - logger.info( - f"[remote_upload] Polling server status... (elapsed: {int(elapsed)}s, limit={'no-limit' if max_poll_time <= 0 else max_poll_time}s)" - ) - time.sleep(poll_interval) - - status = self.get_server_status() - if status.get("success"): - new_sequence = status.get("last_sequence", 0) - if new_sequence >= expected_sequence: - logger.info( - f"[remote_upload] Server sequence advanced to {new_sequence}, bundle was processed!" - ) - return { - "success": True, - "message": f"Bundle processed after timeout (server at sequence {new_sequence})", - "sequence": new_sequence, - } - logger.debug( - f"[remote_upload] Server sequence still at {new_sequence}, continuing to poll..." - ) - else: - logger.warning( - f"[remote_upload] Failed to get server status during poll: {status.get('error', {}).get('message', 'Unknown')}" - ) - - except Exception as e: - logger.error(f"[remote_upload] Error during post-timeout polling: {e}") - return {"success": False, "error": {"code": "POLL_ERROR", "message": f"Polling error: {str(e)}"}} - def get_server_status(self) -> Dict[str, Any]: """Get server status with simplified error handling.""" try: @@ -1149,14 +1608,27 @@ def get_server_status(self) -> Dict[str, Any]: connect_timeout = min(self.timeout, 10) # Allow slower responses (e.g., cold starts/large collections) before bailing read_timeout = max(self.timeout, 30) + params = {"workspace_path": container_workspace_path} + sess = get_auth_session(self.upload_endpoint) + if sess: + params["session"] = sess response = self.session.get( f"{self.upload_endpoint}/api/v1/delta/status", - params={'workspace_path': container_workspace_path}, + params=params, timeout=(connect_timeout, read_timeout) ) if response.status_code == 200: - return response.json() + payload = response.json() + if not isinstance(payload, dict): + return { + "success": False, + "error": { + "code": "STATUS_INVALID", + "message": "Invalid status response payload", + }, + } + return {"success": True, **payload} # Handle error response error_msg = f"Status check failed with HTTP {response.status_code}" @@ -1180,6 +1652,93 @@ def has_meaningful_changes(self, changes: Dict[str, List]) -> bool: total_changes = sum(len(files) for op, files in changes.items() if op != "unchanged") return total_changes > 0 + def _collect_force_cleanup_paths(self) -> List[Path]: + """ + Return ignored paths that force mode should actively delete remotely. + + In dev-remote mode, dev-workspace is intentionally ignored during upload + scans to avoid recursive dogfooding. If that tree already exists on the + remote side from an older buggy upload, force mode should remove it even + when the local cache does not contain those paths. + """ + cleanup_paths: List[Path] = [] + if "dev-workspace" not in self._excluded_dirnames(): + return cleanup_paths + + dev_root = Path(self.workspace_path) / "dev-workspace" + if not dev_root.exists(): + return cleanup_paths + + for root, dirnames, filenames in os.walk(dev_root): + dirnames[:] = [d for d in dirnames if not d.startswith(".")] + for filename in filenames: + path = Path(root) / filename + try: + if path.is_file(): + cleanup_paths.append(path) + except Exception: + continue + return cleanup_paths + + def build_force_changes(self, all_files: List[Path]) -> Dict[str, List]: + """ + Build force-upload changes while still cleaning stale cached paths. + + Force mode should re-upload every currently managed file, but it must also + emit deletes for files that only exist in the local cache now, including + paths that are ignored under the current client policy such as + dev-workspace in dev-remote mode. + """ + created_files: List[Path] = [] + path_map: Dict[Path, Path] = {} + for path in all_files: + if self._is_ignored_path(path): + continue + try: + resolved = path.resolve() + except Exception: + continue + created_files.append(path) + path_map[resolved] = path + + for cached_abs in self._get_all_cached_paths(): + try: + cached_path = Path(cached_abs) + resolved = cached_path.resolve() + except Exception: + continue + if resolved not in path_map: + path_map[resolved] = cached_path + + force_cleanup_paths = self._collect_force_cleanup_paths() + for cleanup_path in force_cleanup_paths: + try: + resolved = cleanup_path.resolve() + except Exception: + continue + if resolved not in path_map: + path_map[resolved] = cleanup_path + + probed = self.detect_file_changes(list(path_map.values())) + deleted_by_resolved: Dict[Path, Path] = {} + for deleted_path in probed.get("deleted", []): + try: + deleted_by_resolved[deleted_path.resolve()] = deleted_path + except Exception: + continue + for cleanup_path in force_cleanup_paths: + try: + deleted_by_resolved.setdefault(cleanup_path.resolve(), cleanup_path) + except Exception: + continue + return { + "created": created_files, + "updated": [], + "deleted": list(deleted_by_resolved.values()), + "moved": [], + "unchanged": [], + } + def upload_git_history_only(self, git_history: Dict[str, Any]) -> bool: try: empty_changes = { @@ -1221,25 +1780,81 @@ def process_changes_and_upload(self, changes: Dict[str, List]) -> bool: try: logger.info(f"[remote_upload] Processing pre-computed changes") + self._last_plan_payload = None + self._last_expected_hashes = {} + # Validate input if not changes: logger.info("[remote_upload] No changes provided") + self._set_last_upload_result("no_changes") return True + if not self.has_meaningful_changes(changes): logger.info("[remote_upload] No meaningful changes detected, skipping upload") + self._set_last_upload_result("no_changes") return True # Log change summary total_changes = sum(len(files) for op, files in changes.items() if op != "unchanged") - logger.info(f"[remote_upload] Detected {total_changes} meaningful changes: " - f"{len(changes['created'])} created, {len(changes['updated'])} updated, " - f"{len(changes['deleted'])} deleted, {len(changes['moved'])} moved") + logger.info( + "[remote_upload] Detected %d candidate changes before remote planning: " + "created=%d updated=%d deleted=%d moved=%d", + total_changes, + len(changes["created"]), + len(changes["updated"]), + len(changes["deleted"]), + len(changes["moved"]), + ) + + planned_changes = changes + plan = self._plan_delta_upload(changes) + preview: Dict[str, Any] = {} + has_content_work = True + if plan: + preview = plan.get("operation_counts_preview", {}) + needed = plan.get("needed_files", {}) if isinstance(plan.get("needed_files", {}), dict) else {} + diagnostics = plan.get("diagnostics", {}) if isinstance(plan.get("diagnostics", {}), dict) else {} + logger.info( + "[remote_upload] Remote plan: candidates=%d content_needed=" + "created=%s updated=%s moved=%s deletes=%s " + "skipped_hash_match=%s needed_bytes=%s cache_entries=%s " + "cache_hash_matches=%s", + total_changes, + len(needed.get("created", []) or []), + len(needed.get("updated", []) or []), + len(needed.get("moved", []) or []), + len(changes.get("deleted", [])), + preview.get("skipped_hash_match", 0), + plan.get("needed_size_bytes", 0), + diagnostics.get("cache_entries", 0), + diagnostics.get("cache_hash_matches", 0), + ) + planned_changes = self._filter_changes_by_plan(changes, plan) + has_content_work = bool( + planned_changes.get("created") + or planned_changes.get("updated") + or planned_changes.get("moved") + ) + if not has_content_work: + apply_only_result = self._apply_operations_without_content(changes, plan) + if apply_only_result is True: + self._finalize_successful_changes(changes) + return True + if not self.has_meaningful_changes(planned_changes): + logger.info("[remote_upload] Plan found no upload work; skipping bundle upload") + self._finalize_successful_changes(changes) + self._set_last_upload_result( + "skipped_by_plan", + plan_preview=preview, + needed_size_bytes=plan.get("needed_size_bytes", 0), + ) + return True # Create delta bundle bundle_path = None try: - bundle_path, manifest = self.create_delta_bundle(changes) + bundle_path, manifest = self.create_delta_bundle(planned_changes) logger.info(f"[remote_upload] Created delta bundle: {manifest['bundle_id']} " f"(size: {manifest['total_size_bytes']} bytes)") @@ -1251,6 +1866,7 @@ def process_changes_and_upload(self, changes: Dict[str, List]) -> bool: logger.error(f"[remote_upload] Error creating delta bundle: {e}") # Clean up any temporary files on failure self.cleanup() + self._set_last_upload_result("failed", stage="bundle_creation", error=str(e)) return False # Upload bundle with retry logic @@ -1258,9 +1874,31 @@ def process_changes_and_upload(self, changes: Dict[str, List]) -> bool: response = self.upload_bundle(bundle_path, manifest) if response.get("success", False): - processed_ops = response.get('processed_operations', {}) - logger.info(f"[remote_upload] Successfully uploaded bundle {manifest['bundle_id']}") - logger.info(f"[remote_upload] Processed operations: {processed_ops}") + processed_ops = response.get("processed_operations") + if processed_ops is None: + logger.info( + "[remote_upload] Bundle %s accepted by server; processing asynchronously (sequence=%s)", + manifest["bundle_id"], + response.get("sequence_number"), + ) + self._set_last_upload_result( + "queued", + bundle_id=manifest["bundle_id"], + sequence_number=response.get("sequence_number"), + ) + # Acceptance is the client-side completion point. The + # server owns background processing and journal retry. + self._finalize_successful_changes(changes) + else: + logger.info(f"[remote_upload] Successfully uploaded bundle {manifest['bundle_id']}") + logger.info(f"[remote_upload] Processed operations: {processed_ops}") + self._finalize_successful_changes(changes) + self._set_last_upload_result( + "uploaded", + bundle_id=manifest["bundle_id"], + sequence_number=response.get("sequence_number"), + processed_operations=processed_ops, + ) # Clean up temporary bundle after successful upload try: @@ -1276,16 +1914,20 @@ def process_changes_and_upload(self, changes: Dict[str, List]) -> bool: else: error_msg = response.get('error', {}).get('message', 'Unknown upload error') logger.error(f"[remote_upload] Upload failed: {error_msg}") + self._set_last_upload_result("failed", stage="upload", error=error_msg) return False except Exception as e: logger.error(f"[remote_upload] Error uploading bundle: {e}") + self._set_last_upload_result("failed", stage="upload", error=str(e)) return False except Exception as e: logger.error(f"[remote_upload] Unexpected error in process_changes_and_upload: {e}") + self._set_last_upload_result("failed", stage="unexpected", error=str(e)) return False + def get_all_code_files(self) -> List[Path]: """Get all code files in the workspace.""" files: List[Path] = [] @@ -1295,28 +1937,31 @@ def get_all_code_files(self) -> List[Path]: return files # Single walk with early pruning similar to standalone client - ext_suffixes = {str(ext).lower() for ext in idx.CODE_EXTS if str(ext).startswith('.')} - name_matches = {str(ext) for ext in idx.CODE_EXTS if not str(ext).startswith('.')} - dev_remote = os.environ.get("DEV_REMOTE_MODE") == "1" or os.environ.get("REMOTE_UPLOAD_MODE") == "development" - excluded = { - "node_modules", "vendor", "dist", "build", "target", "out", - ".git", ".hg", ".svn", ".vscode", ".idea", ".venv", "venv", - "__pycache__", ".pytest_cache", ".mypy_cache", ".cache", - ".context-engine", ".context-engine-uploader", ".codebase" - } - if dev_remote: - excluded.add("dev-workspace") + ext_suffixes = {str(ext).lower() for ext in CODE_EXTS if str(ext).startswith('.')} + try: + extensionless_names = {k.lower() for k in (EXTENSIONLESS_FILES or {}).keys()} + except Exception: + extensionless_names = set() + excluded = self._excluded_dirnames() seen = set() for root, dirnames, filenames in os.walk(workspace_path): dirnames[:] = [d for d in dirnames if d not in excluded and not d.startswith('.')] for filename in filenames: - if filename.startswith('.'): + # Allow dotfiles that are in EXTENSIONLESS_FILES (e.g., .gitignore) + fname_lower = filename.lower() + if filename.startswith('.') and fname_lower not in extensionless_names: continue candidate = Path(root) / filename + if self._is_ignored_path(candidate): + continue suffix = candidate.suffix.lower() - if filename in name_matches or suffix in ext_suffixes: + if ( + suffix in ext_suffixes + or fname_lower in extensionless_names + or fname_lower.startswith("dockerfile") + ): resolved = candidate.resolve() if resolved not in seen: seen.add(resolved) @@ -1368,13 +2013,13 @@ def on_any_event(self, event): # Always check src_path src_path = Path(event.src_path) - if idx.CODE_EXTS.get(src_path.suffix.lower(), "unknown") != "unknown": + if self.client._is_watchable_path(src_path): paths_to_process.append(src_path) # For FileMovedEvent, also process the destination path if hasattr(event, 'dest_path') and event.dest_path: dest_path = Path(event.dest_path) - if idx.CODE_EXTS.get(dest_path.suffix.lower(), "unknown") != "unknown": + if self.client._is_watchable_path(dest_path): paths_to_process.append(dest_path) if not paths_to_process: @@ -1395,6 +2040,8 @@ def on_any_event(self, event): def _process_pending_changes(self): """Process accumulated changes after debounce period.""" with self._lock: + # Timer fired; allow a new debounce to be armed while we process. + self._debounce_timer = None # Prevent re-entrancy if self._processing: return @@ -1406,19 +2053,21 @@ def _process_pending_changes(self): check_deletions = self._check_for_deletions self._check_for_deletions = False + upload_succeeded = False try: # Only include cached paths when deletion-related events occurred if check_deletions: cached_file_hashes = _load_local_cache_file_hashes( self.client.workspace_path, - self.client.repo_name + self.client.repo_name, + metadata_root=self.client.metadata_root, ) - all_paths = list(set(pending + [ - Path(p) for p in cached_file_hashes.keys() - ])) + cached_paths = [Path(p) for p in cached_file_hashes.keys()] + all_paths = list(set(pending + cached_paths)) else: all_paths = pending + changes = self.client.detect_file_changes(all_paths) meaningful_changes = ( len(changes.get("created", [])) + @@ -1428,10 +2077,15 @@ def _process_pending_changes(self): ) if meaningful_changes > 0: - logger.info(f"[watch] Detected {meaningful_changes} changes: { {k: len(v) for k, v in changes.items() if k != 'unchanged'} }") + logger.info( + "[watch] Detected %d candidate changes: %s", + meaningful_changes, + {k: len(v) for k, v in changes.items() if k != "unchanged"}, + ) success = self.client.process_changes_and_upload(changes) if success: - logger.info("[watch] Successfully uploaded changes") + self.client.log_watch_upload_result() + upload_succeeded = True else: logger.error("[watch] Failed to upload changes") else: @@ -1447,19 +2101,33 @@ def _process_pending_changes(self): success = self.client.upload_git_history_only(git_history) if success: logger.info("[watch] Successfully uploaded git history metadata") + upload_succeeded = True else: logger.error("[watch] Failed to upload git history metadata") + else: + upload_succeeded = True # No changes to process except Exception as e: logger.error(f"[watch] Error processing changes: {e}") finally: # Clear processing flag even if an error occurred with self._lock: self._processing = False + # Re-queue pending paths if upload failed + if not upload_succeeded and pending: + # Merge pending paths back into _pending_paths + for p in pending: + self._pending_paths.add(p) + # Arm next pass if there are pending paths + if self._pending_paths and self._debounce_timer is None: + self._debounce_timer = threading.Timer( + self.debounce_seconds, + self._process_pending_changes, + ) + self._debounce_timer.start() observer = Observer() handler = CodeFileEventHandler(self, debounce_seconds=2.0) - try: observer.schedule(handler, self.workspace_path, recursive=True) observer.start() @@ -1504,7 +2172,11 @@ def _watch_loop_polling(self, interval: int = 5): path_map[resolved] = p # Include any paths that are only present in the local cache (deleted files) - cached_file_hashes = _load_local_cache_file_hashes(self.workspace_path, self.repo_name) + cached_file_hashes = _load_local_cache_file_hashes( + self.workspace_path, + self.repo_name, + metadata_root=self.metadata_root, + ) for cached_abs in cached_file_hashes.keys(): try: cached_path = Path(cached_abs) @@ -1521,12 +2193,16 @@ def _watch_loop_polling(self, interval: int = 5): meaningful_changes = len(changes.get("created", [])) + len(changes.get("updated", [])) + len(changes.get("deleted", [])) + len(changes.get("moved", [])) if meaningful_changes > 0: - logger.info(f"[watch] Detected {meaningful_changes} changes: { {k: len(v) for k, v in changes.items() if k != 'unchanged'} }") + logger.info( + "[watch] Detected %d candidate changes: %s", + meaningful_changes, + {k: len(v) for k, v in changes.items() if k != "unchanged"}, + ) success = self.process_changes_and_upload(changes) if success: - logger.info(f"[watch] Successfully uploaded changes") + self.log_watch_upload_result() else: logger.error(f"[watch] Failed to upload changes") else: @@ -1584,80 +2260,7 @@ def process_and_upload_changes(self, changed_paths: List[Path]) -> bool: except Exception as e: logger.error(f"[remote_upload] Error detecting file changes: {e}") return False - - if not self.has_meaningful_changes(changes): - logger.info("[remote_upload] No meaningful changes detected, skipping upload") - return True - - # Log change summary - total_changes = sum(len(files) for op, files in changes.items() if op != "unchanged") - logger.info(f"[remote_upload] Detected {total_changes} meaningful changes: " - f"{len(changes['created'])} created, {len(changes['updated'])} updated, " - f"{len(changes['deleted'])} deleted, {len(changes['moved'])} moved") - - # Create delta bundle - bundle_path = None - try: - bundle_path, manifest = self.create_delta_bundle(changes) - logger.info(f"[remote_upload] Created delta bundle: {manifest['bundle_id']} " - f"(size: {manifest['total_size_bytes']} bytes)") - - # Validate bundle was created successfully - if not bundle_path or not os.path.exists(bundle_path): - raise RuntimeError(f"Failed to create bundle at {bundle_path}") - - except Exception as e: - logger.error(f"[remote_upload] Error creating delta bundle: {e}") - # Clean up any temporary files on failure - self.cleanup() - return False - - # Upload bundle with retry logic - try: - response = self.upload_bundle(bundle_path, manifest) - - if response.get("success", False): - processed_ops = response.get('processed_operations', {}) - logger.info(f"[remote_upload] Successfully uploaded bundle {manifest['bundle_id']}") - logger.info(f"[remote_upload] Processed operations: {processed_ops}") - - # Clean up temporary bundle after successful upload - try: - if os.path.exists(bundle_path): - os.remove(bundle_path) - logger.debug(f"[remote_upload] Cleaned up temporary bundle: {bundle_path}") - # Also clean up the entire temp directory if this is the last bundle - self.cleanup() - except Exception as cleanup_error: - logger.warning(f"[remote_upload] Failed to cleanup bundle {bundle_path}: {cleanup_error}") - - return True - else: - error = response.get("error", {}) - error_code = error.get("code", "UNKNOWN") - error_msg = error.get("message", "Unknown error") - - logger.error(f"[remote_upload] Upload failed: {error_msg}") - - # Handle specific error types - # CLI is stateless - server handles sequence management - if error_code in ["BUNDLE_TOO_LARGE", "BUNDLE_NOT_FOUND"]: - # These are unrecoverable errors - logger.error(f"[remote_upload] Unrecoverable error ({error_code}): {error_msg}") - return False - elif error_code in ["TIMEOUT_ERROR", "CONNECTION_ERROR", "NETWORK_ERROR"]: - # These might be temporary, suggest fallback - logger.warning(f"[remote_upload] Network-related error ({error_code}): {error_msg}") - logger.warning("[remote_upload] Consider falling back to local mode if this persists") - return False - else: - # Other errors - logger.error(f"[remote_upload] Upload error ({error_code}): {error_msg}") - return False - - except Exception as e: - logger.error(f"[remote_upload] Unexpected error during upload: {e}") - return False + return self.process_changes_and_upload(changes) except Exception as e: logger.error(f"[remote_upload] Critical error in process_and_upload_changes: {e}") @@ -1688,7 +2291,7 @@ def _cleanup_dir_with_retries(path: Optional[str]) -> None: logger.debug(f"[remote_upload] Last cleanup error for {path}: {last_error}") -def get_remote_config(cli_path: Optional[str] = None) -> Dict[str, str]: +def get_remote_config(cli_path: Optional[str] = None) -> Dict[str, Any]: """Get remote upload configuration from environment variables and command-line arguments.""" # Use command-line path if provided, otherwise fall back to environment variables if cli_path: @@ -1698,17 +2301,10 @@ def get_remote_config(cli_path: Optional[str] = None) -> Dict[str, str]: logical_repo_id = _compute_logical_repo_id(workspace_path) - # Use auto-generated collection name based on repo name - repo_name = _extract_repo_name_from_path(workspace_path) - # Fallback to directory name if repo detection fails - if not repo_name: - repo_name = Path(workspace_path).name - collection_name = get_collection_name(repo_name) - return { "upload_endpoint": os.environ.get("REMOTE_UPLOAD_ENDPOINT", "http://localhost:8080"), "workspace_path": workspace_path, - "collection_name": collection_name, + "collection_name": None, "logical_repo_id": logical_repo_id, # Use higher, more robust defaults but still allow env overrides "max_retries": int(os.environ.get("REMOTE_UPLOAD_MAX_RETRIES", "5")), @@ -1816,7 +2412,7 @@ def main(): config["timeout"] = args.timeout logger.info(f"Workspace path: {config['workspace_path']}") - logger.info(f"Collection name: {config['collection_name']}") + logger.info(f"Collection name: {config['collection_name'] or ''}") logger.info(f"Upload endpoint: {config['upload_endpoint']}") if args.show_mapping: @@ -1850,15 +2446,8 @@ def main(): # Test server connection first logger.info("Checking server status...") status = client.get_server_status() - is_success = ( - isinstance(status, dict) and - 'workspace_path' in status and - 'collection_name' in status and - status.get('status') == 'ready' - ) - if not is_success: - error = status.get("error", {}) - logger.error(f"Cannot connect to server: {error.get('message', 'Unknown error')}") + if not _is_usable_delta_status(status): + logger.error("Cannot connect to server: %s", _server_status_error_message(status)) return 1 logger.info("Server connection successful") @@ -1894,34 +2483,24 @@ def main(): # Test server connection logger.info("Checking server status...") status = client.get_server_status() - # For delta endpoint, success is indicated by having expected fields (not a "success" boolean) - is_success = ( - isinstance(status, dict) and - 'workspace_path' in status and - 'collection_name' in status and - status.get('status') == 'ready' - ) - if not is_success: - error = status.get("error", {}) - logger.error(f"Cannot connect to server: {error.get('message', 'Unknown error')}") + if not _is_usable_delta_status(status): + logger.error("Cannot connect to server: %s", _server_status_error_message(status)) return 1 - logger.info("Server connection successful") + logger.info( + "Server connection successful: status=%s pending_journal=%s journal=%s", + status.get("status"), + status.get("pending_operations"), + (status.get("server_info") or {}).get("journal", {}), + ) # Scan repository and upload files logger.info("Scanning repository for files...") workspace_path = Path(config['workspace_path']) - # Find all files in the repository - all_files = [] - for file_path in workspace_path.rglob('*'): - if file_path.is_file() and not file_path.name.startswith('.'): - rel_path = file_path.relative_to(workspace_path) - # Skip .codebase directory and other metadata - if not str(rel_path).startswith('.codebase'): - all_files.append(file_path) + all_files = client.get_all_code_files() - logger.info(f"Found {len(all_files)} files to upload") + logger.info(f"Found {len(all_files)} eligible files to scan") if not all_files: logger.warning("No files found to upload") @@ -1929,8 +2508,7 @@ def main(): # Detect changes (treat all files as changes for initial upload) if args.force: - # Force mode: treat all files as created - changes = {"created": all_files, "updated": [], "deleted": [], "moved": [], "unchanged": []} + changes = client.build_force_changes(all_files) else: changes = client.detect_file_changes(all_files) @@ -1938,16 +2516,32 @@ def main(): logger.info("No meaningful changes to upload") return 0 - logger.info(f"Changes detected: {len(changes.get('created', []))} created, {len(changes.get('updated', []))} updated, {len(changes.get('deleted', []))} deleted") + logger.info( + "Candidates detected before remote planning: created=%d updated=%d deleted=%d moved=%d", + len(changes.get("created", [])), + len(changes.get("updated", [])), + len(changes.get("deleted", [])), + len(changes.get("moved", [])), + ) # Process and upload changes logger.info("Uploading files to remote server...") success = client.process_changes_and_upload(changes) if success: - logger.info("Repository upload completed successfully!") - logger.info(f"Collection name: {config['collection_name']}") - logger.info(f"Files uploaded: {len(all_files)}") + outcome = str((client.last_upload_result or {}).get("outcome") or "") + if outcome == "skipped_by_plan": + logger.info("No upload needed after plan") + elif outcome == "queued": + logger.info("Repository upload request accepted; server processing asynchronously") + else: + logger.info("Repository upload completed successfully!") + logger.info(f"Collection name: {config['collection_name'] or ''}") + logger.info( + "Remote upload result: outcome=%s details=%s", + outcome, + client.last_upload_result, + ) else: logger.error("Repository upload failed!") return 1 diff --git a/scripts/rerank_ab_test.py b/scripts/rerank_ab_test.py deleted file mode 100644 index bf592065..00000000 --- a/scripts/rerank_ab_test.py +++ /dev/null @@ -1,10 +0,0 @@ -#!/usr/bin/env python3 -"""Backward-compatibility shim. See scripts/rerank_tools/ab_test.py""" -import sys -from pathlib import Path -sys.path.insert(0, str(Path(__file__).parent.parent)) - -from scripts.rerank_tools.ab_test import * - -if __name__ == "__main__": - simulate_ab_test(n_sessions=100, n_queries_per_session=5) diff --git a/scripts/rerank_eval.py b/scripts/rerank_eval.py deleted file mode 100644 index ef810c82..00000000 --- a/scripts/rerank_eval.py +++ /dev/null @@ -1,10 +0,0 @@ -#!/usr/bin/env python3 -"""Backward-compatibility shim. See scripts/rerank_tools/eval.py""" -import sys -from pathlib import Path -sys.path.insert(0, str(Path(__file__).parent.parent)) - -from scripts.rerank_tools.eval import * - -if __name__ == "__main__": - main() diff --git a/scripts/rerank_events.py b/scripts/rerank_events.py deleted file mode 100644 index f98916af..00000000 --- a/scripts/rerank_events.py +++ /dev/null @@ -1,7 +0,0 @@ -#!/usr/bin/env python3 -"""Backward-compatibility shim. See scripts/rerank_tools/events.py""" -import sys -from pathlib import Path -sys.path.insert(0, str(Path(__file__).parent.parent)) - -from scripts.rerank_tools.events import * diff --git a/scripts/rerank_local.py b/scripts/rerank_local.py deleted file mode 100644 index 8a56e7d9..00000000 --- a/scripts/rerank_local.py +++ /dev/null @@ -1,10 +0,0 @@ -#!/usr/bin/env python3 -"""Backward-compatibility shim. See scripts/rerank_tools/local.py""" -import sys -from pathlib import Path -sys.path.insert(0, str(Path(__file__).parent.parent)) - -from scripts.rerank_tools.local import * - -if __name__ == "__main__": - main() diff --git a/scripts/rerank_query.py b/scripts/rerank_query.py deleted file mode 100644 index f20936bf..00000000 --- a/scripts/rerank_query.py +++ /dev/null @@ -1,10 +0,0 @@ -#!/usr/bin/env python3 -"""Backward-compatibility shim. See scripts/rerank_tools/query.py""" -import sys -from pathlib import Path -sys.path.insert(0, str(Path(__file__).parent.parent)) - -from scripts.rerank_tools.query import * - -if __name__ == "__main__": - main() diff --git a/scripts/rerank_real_benchmark.py b/scripts/rerank_real_benchmark.py deleted file mode 100644 index fd380df4..00000000 --- a/scripts/rerank_real_benchmark.py +++ /dev/null @@ -1,10 +0,0 @@ -#!/usr/bin/env python3 -"""Backward-compatibility shim. See scripts/rerank_tools/benchmark.py""" -import sys -from pathlib import Path -sys.path.insert(0, str(Path(__file__).parent.parent)) - -from scripts.rerank_tools.benchmark import * - -if __name__ == "__main__": - run_real_benchmark() diff --git a/scripts/rerank_recursive.py b/scripts/rerank_recursive.py deleted file mode 100644 index 879cb53b..00000000 --- a/scripts/rerank_recursive.py +++ /dev/null @@ -1,88 +0,0 @@ -#!/usr/bin/env python3 -""" -Backwards-compatibility shim for rerank_recursive. - -This file re-exports all symbols from the scripts.rerank_recursive package -to maintain backwards compatibility for existing imports. - -All new code should import from scripts.rerank_recursive (the package) directly. - -Usage (both work): - # Old style (still works) - from scripts.rerank_recursive import RecursiveReranker, TinyScorer - - # New style (preferred) - from scripts.rerank_recursive import RecursiveReranker, TinyScorer -""" -from __future__ import annotations - -# Re-export everything from the package -from scripts.rerank_recursive import ( - # State - RefinementState, - # Core classes - TinyScorer, - LatentRefiner, - VICReg, - LearnedProjection, - LearnedHybridWeights, - QueryExpander, - ConfidenceEstimator, - RecursiveReranker, - ONNXRecursiveReranker, - SessionAwareReranker, - # Utilities - _COMMON_TOKENS, - _split_identifier, - _normalize_token, - _tokenize_for_fname_boost, - _candidate_path_for_fname_boost, - _compute_fname_boost, - _cache_key, - _get_cached_embedding, - _cache_embedding, - # Functions - rerank_recursive, - rerank_recursive_inprocess, - rerank_with_learning, - rerank_with_session, - get_recursive_reranker, - _get_learning_reranker, - # Constants - HAS_ONNX, -) - -__all__ = [ - # State - "RefinementState", - # Core classes - "TinyScorer", - "LatentRefiner", - "VICReg", - "LearnedProjection", - "LearnedHybridWeights", - "QueryExpander", - "ConfidenceEstimator", - "RecursiveReranker", - "ONNXRecursiveReranker", - "SessionAwareReranker", - # Utilities - "_COMMON_TOKENS", - "_split_identifier", - "_normalize_token", - "_tokenize_for_fname_boost", - "_candidate_path_for_fname_boost", - "_compute_fname_boost", - "_cache_key", - "_get_cached_embedding", - "_cache_embedding", - # Functions - "rerank_recursive", - "rerank_recursive_inprocess", - "rerank_with_learning", - "rerank_with_session", - "get_recursive_reranker", - "_get_learning_reranker", - # Constants - "HAS_ONNX", -] diff --git a/scripts/rerank_recursive/__init__.py b/scripts/rerank_recursive/__init__.py deleted file mode 100644 index 23d96ac9..00000000 --- a/scripts/rerank_recursive/__init__.py +++ /dev/null @@ -1,115 +0,0 @@ -""" -Recursive Reranker Package - TRM-inspired iterative refinement for code search. - -This package provides modular components for recursive reranking: - -Core Components: -- TinyScorer: 2-layer MLP for scoring query-document pairs -- LatentRefiner: Refines latent state based on current results -- RecursiveReranker: Main reranking pipeline - -Regularization: -- VICReg: Variance-Invariance-Covariance regularization - -Learnable Components: -- LearnedProjection: Learnable embedding projection -- LearnedHybridWeights: Learns dense vs. lexical balance -- QueryExpander: Learns query expansions from usage - -Utilities: -- RefinementState: Dataclass for latent state -- ConfidenceEstimator: Early stopping logic -""" -from __future__ import annotations - -# State dataclass -from scripts.rerank_recursive.state import RefinementState - -# Utilities -from scripts.rerank_recursive.utils import ( - _COMMON_TOKENS, - _split_identifier, - _normalize_token, - _tokenize_for_fname_boost, - _candidate_path_for_fname_boost, - _compute_fname_boost, - _cache_key, - _get_cached_embedding, - _cache_embedding, -) - -# Core scorer and refiner -from scripts.rerank_recursive.scorer import TinyScorer -from scripts.rerank_recursive.refiner import LatentRefiner - -# Regularization -from scripts.rerank_recursive.vicreg import VICReg - -# Learnable components -from scripts.rerank_recursive.projection import LearnedProjection -from scripts.rerank_recursive.hybrid_weights import LearnedHybridWeights -from scripts.rerank_recursive.expander import QueryExpander - -# Early stopping -from scripts.rerank_recursive.confidence import ConfidenceEstimator - -# Alpha scheduling -from scripts.rerank_recursive.alpha_scheduler import ( - CosineAlphaScheduler, - LearnedAlphaWeights, -) - -# Main rerankers and functions -from scripts.rerank_recursive.recursive import ( - RecursiveReranker, - ONNXRecursiveReranker, - FastEmbedRecursiveReranker, - SessionAwareReranker, - rerank_recursive, - rerank_recursive_inprocess, - rerank_with_learning, - rerank_with_session, - get_recursive_reranker, - _get_learning_reranker, - HAS_ONNX, - HAS_RERANKER_FACTORY, -) - -__all__ = [ - # State - "RefinementState", - # Core classes - "TinyScorer", - "LatentRefiner", - "VICReg", - "LearnedProjection", - "LearnedHybridWeights", - "QueryExpander", - "ConfidenceEstimator", - "CosineAlphaScheduler", - "LearnedAlphaWeights", - "RecursiveReranker", - "ONNXRecursiveReranker", - "FastEmbedRecursiveReranker", - "SessionAwareReranker", - # Utilities - "_COMMON_TOKENS", - "_split_identifier", - "_normalize_token", - "_tokenize_for_fname_boost", - "_candidate_path_for_fname_boost", - "_compute_fname_boost", - "_cache_key", - "_get_cached_embedding", - "_cache_embedding", - # Functions - "rerank_recursive", - "rerank_recursive_inprocess", - "rerank_with_learning", - "rerank_with_session", - "get_recursive_reranker", - "_get_learning_reranker", - # Constants - "HAS_ONNX", - "HAS_RERANKER_FACTORY", -] diff --git a/scripts/rerank_recursive/alpha_scheduler.py b/scripts/rerank_recursive/alpha_scheduler.py deleted file mode 100644 index e01efd5f..00000000 --- a/scripts/rerank_recursive/alpha_scheduler.py +++ /dev/null @@ -1,298 +0,0 @@ -""" -Alpha Scheduler - Learnable and scheduled alpha for score blending. - -From TRM paper insight: the blending factor α between new scores and previous -scores should vary per iteration. Early iterations need more exploration (higher α), -later iterations need more exploitation (lower α). - -Two strategies: -1. CosineAlphaScheduler: Fixed cosine decay schedule -2. LearnedAlphaWeights: Per-iteration learnable weights with persistence -""" -import os -import time -from typing import Any, Dict, List, Optional, Tuple - -import numpy as np - - -class CosineAlphaScheduler: - """Fixed cosine schedule for alpha values. - - α = alpha_min + (alpha_max - alpha_min) * (1 + cos(π * i / (n-1))) / 2 - - This gives higher α early (trust new scores more) and lower α later - (rely more on refined estimates). - """ - - def __init__( - self, - n_iterations: int = 3, - alpha_max: float = 0.7, - alpha_min: float = 0.3, - ): - self.n_iterations = max(1, n_iterations) - self.alpha_max = alpha_max - self.alpha_min = alpha_min - # Pre-compute schedule - self._schedule = self._compute_schedule() - - def _compute_schedule(self) -> List[float]: - """Compute the full schedule.""" - if self.n_iterations == 1: - return [(self.alpha_max + self.alpha_min) / 2] - - schedule = [] - for i in range(self.n_iterations): - # Cosine decay from alpha_max to alpha_min - progress = i / (self.n_iterations - 1) - alpha = self.alpha_min + (self.alpha_max - self.alpha_min) * (1 + np.cos(np.pi * progress)) / 2 - schedule.append(float(alpha)) - return schedule - - def get_alpha(self, iteration: int) -> float: - """Get alpha for a specific iteration (0-indexed).""" - idx = max(0, min(iteration, len(self._schedule) - 1)) - return self._schedule[idx] - - def get_schedule(self) -> List[float]: - """Get the full schedule.""" - return self._schedule.copy() - - -class LearnedAlphaWeights: - """Learnable per-iteration alpha weights with persistence. - - Uses sigmoid(raw_weight) to keep alpha in (0, 1). - Learns optimal blending through gradient descent on ranking loss. - """ - - WEIGHTS_DIR = os.environ.get("RERANKER_WEIGHTS_DIR", "/tmp/rerank_weights") - WEIGHTS_RELOAD_INTERVAL = float(os.environ.get("RERANKER_WEIGHTS_RELOAD_INTERVAL", "60")) - - def __init__( - self, - n_iterations: int = 3, - init_alpha: float = 0.5, - lr: float = 0.01, - ): - self.n_iterations = max(1, n_iterations) - self.lr = lr - self._collection = "default" - self._weights_path = self._get_weights_path("default") - self._weights_mtime = 0.0 - self._last_reload_check = 0.0 - self._weights_loaded = False - self._update_count = 0 - self._version = 0 - - # Initialize raw weights such that sigmoid(raw) = init_alpha - # sigmoid(x) = init_alpha => x = logit(init_alpha) - init_raw = np.log(init_alpha / (1 - init_alpha + 1e-8)) - self.raw_weights = np.full(self.n_iterations, init_raw, dtype=np.float32) - self._momentum = np.zeros_like(self.raw_weights) - - # Try to load saved weights - if os.path.exists(self._weights_path): - try: - self._load_weights() - except Exception as e: - from scripts.logger import get_logger - get_logger(__name__).warning(f"LearnedAlphaWeights: failed to load: {e}") - - @staticmethod - def _sanitize_collection(collection: str) -> str: - return "".join(c if c.isalnum() or c in "-_" else "_" for c in collection) - - def _get_weights_path(self, collection: str) -> str: - safe_name = self._sanitize_collection(collection) - return os.path.join(self.WEIGHTS_DIR, f"alpha_{safe_name}.npz") - - def set_collection(self, collection: str): - """Set collection and load corresponding weights.""" - self._collection = collection - self._weights_path = self._get_weights_path(collection) - if os.path.exists(self._weights_path): - try: - self._load_weights() - except Exception: - pass - - def _sigmoid(self, x: np.ndarray) -> np.ndarray: - """Numerically stable sigmoid.""" - return np.where( - x >= 0, - 1 / (1 + np.exp(-x)), - np.exp(x) / (1 + np.exp(x)) - ) - - def _sigmoid_grad(self, sigmoid_val: np.ndarray) -> np.ndarray: - """Gradient of sigmoid: σ(x) * (1 - σ(x)).""" - return sigmoid_val * (1 - sigmoid_val) - - def get_alpha(self, iteration: int) -> float: - """Get alpha for a specific iteration (0-indexed).""" - self.maybe_reload_weights() - idx = max(0, min(iteration, len(self.raw_weights) - 1)) - return float(self._sigmoid(self.raw_weights[idx])) - - def get_schedule(self) -> List[float]: - """Get alpha values for all iterations.""" - return [float(a) for a in self._sigmoid(self.raw_weights)] - - def maybe_reload_weights(self): - """Hot-reload weights if changed on disk.""" - now = time.time() - if now - self._last_reload_check < self.WEIGHTS_RELOAD_INTERVAL: - return - self._last_reload_check = now - try: - if os.path.exists(self._weights_path): - mtime = os.path.getmtime(self._weights_path) - if mtime > self._weights_mtime: - self._load_weights() - except Exception: - pass - - def learn_from_ranking_loss( - self, - iteration: int, - blended_scores: np.ndarray, - teacher_scores: np.ndarray, - new_scores: Optional[np.ndarray] = None, - prev_scores: Optional[np.ndarray] = None, - ) -> float: - """Learn alpha to minimize ranking difference with teacher. - - Uses pairwise ranking loss: if teacher prefers A over B, reduce loss - when our blended scores also prefer A over B. - - Args: - iteration: Which iteration's alpha to update - blended_scores: (n_docs,) our blended scores at this iteration - teacher_scores: (n_docs,) ground truth scores - new_scores: (n_docs,) optional - the "new" signal before blending - prev_scores: (n_docs,) optional - the "previous" signal before blending - - Returns: - Gradient magnitude (for logging) - """ - if len(blended_scores) < 2: - return 0.0 - - idx = max(0, min(iteration, len(self.raw_weights) - 1)) - alpha = self._sigmoid(self.raw_weights[idx:idx+1])[0] - - # Compute ranking agreement - our_order = np.argsort(-blended_scores) - teacher_order = np.argsort(-teacher_scores) - - # Count top-k mismatches - k = min(3, len(blended_scores)) - n_mismatches = np.sum(our_order[:k] != teacher_order[:k]) - - if n_mismatches == 0: - return 0.0 # Perfect match, no update needed - - # Determine gradient direction based on which signal would have helped - grad_direction = 1.0 # Default: increase alpha (trust new scores more) - - if new_scores is not None and prev_scores is not None: - # Compare: would higher alpha (more new_scores) or lower alpha (more prev_scores) help? - new_order = np.argsort(-new_scores) - prev_order = np.argsort(-prev_scores) - - # Count how many top-k matches each signal has with teacher - new_matches = np.sum(new_order[:k] == teacher_order[:k]) - prev_matches = np.sum(prev_order[:k] == teacher_order[:k]) - - if prev_matches > new_matches: - # Previous scores are better aligned with teacher → decrease alpha - grad_direction = -1.0 - elif new_matches > prev_matches: - # New scores are better aligned with teacher → increase alpha - grad_direction = 1.0 - else: - # Tie: small nudge toward middle (alpha=0.5) - grad_direction = 0.5 - alpha # Pull toward 0.5 - - # Gradient through sigmoid, scaled by mismatches and direction - sigmoid_grad = self._sigmoid_grad(np.array([alpha]))[0] - grad = n_mismatches * 0.1 * sigmoid_grad * grad_direction - - # Momentum SGD - momentum = 0.9 - self._momentum[idx] = momentum * self._momentum[idx] + grad - self.raw_weights[idx] -= self.lr * self._momentum[idx] - - # Clamp to reasonable range (alpha between ~0.1 and ~0.9) - self.raw_weights[idx] = np.clip(self.raw_weights[idx], -2.2, 2.2) - - self._update_count += 1 - if self._update_count % 50 == 0: - self._save_weights() - - return abs(grad) - - def _load_weights(self): - """Load weights from disk.""" - import fcntl - lock_path = self._weights_path + ".lock" - try: - os.makedirs(os.path.dirname(lock_path) or ".", exist_ok=True) - with open(lock_path, "w") as lock_file: - fcntl.flock(lock_file.fileno(), fcntl.LOCK_SH) - data = np.load(self._weights_path) - loaded_weights = data["raw_weights"] - # Handle dimension mismatch gracefully - if len(loaded_weights) == len(self.raw_weights): - self.raw_weights = loaded_weights.astype(np.float32) - else: - # Resize: copy what we can, init rest - min_len = min(len(loaded_weights), len(self.raw_weights)) - self.raw_weights[:min_len] = loaded_weights[:min_len].astype(np.float32) - self._version = int(data.get("version", 0)) - fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN) - self._weights_mtime = os.path.getmtime(self._weights_path) - self._weights_loaded = True - self._momentum = np.zeros_like(self.raw_weights) - except Exception as e: - from scripts.logger import get_logger - get_logger(__name__).warning(f"LearnedAlphaWeights: load failed: {e}") - - def _save_weights(self): - """Save weights atomically.""" - import fcntl - os.makedirs(self.WEIGHTS_DIR, exist_ok=True) - lock_path = self._weights_path + ".lock" - tmp_path = self._weights_path.replace(".npz", ".tmp.npz") - self._version += 1 - try: - with open(lock_path, "w") as lock_file: - fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX) - np.savez( - tmp_path, - raw_weights=self.raw_weights, - version=self._version, - n_iterations=self.n_iterations, - ) - os.replace(tmp_path, self._weights_path) - fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN) - self._weights_mtime = os.path.getmtime(self._weights_path) - except Exception as e: - from scripts.logger import get_logger - get_logger(__name__).warning(f"LearnedAlphaWeights: save failed: {e}") - if os.path.exists(tmp_path): - try: - os.remove(tmp_path) - except Exception: - pass - - def get_metrics(self) -> Dict[str, Any]: - """Get current metrics for logging.""" - return { - "collection": self._collection, - "version": self._version, - "update_count": self._update_count, - "alphas": self.get_schedule(), - } diff --git a/scripts/rerank_recursive/confidence.py b/scripts/rerank_recursive/confidence.py deleted file mode 100644 index 50c92a0c..00000000 --- a/scripts/rerank_recursive/confidence.py +++ /dev/null @@ -1,54 +0,0 @@ -""" -ConfidenceEstimator - Estimates confidence to enable early stopping. - -From TRM: Q-learning inspired halting - stop when improvement is minimal. -""" -import numpy as np - -from scripts.rerank_recursive.state import RefinementState - - -class ConfidenceEstimator: - """ - Estimates confidence to enable early stopping. - - Uses patience to avoid stopping on noisy single-step improvements. - """ - - def __init__(self, patience: int = 1, min_improvement: float = 0.01): - self.patience = patience - self.min_improvement = min_improvement - self._stable_count = 0 - - def reset(self): - """Reset state for a new query.""" - self._stable_count = 0 - - def should_stop(self, state: RefinementState) -> bool: - """Check if we should stop refining based on score stability.""" - if len(state.score_history) < 2: - return False - - prev_scores = state.score_history[-2] - curr_scores = state.scores - - prev_order = np.argsort(-prev_scores) - curr_order = np.argsort(-curr_scores) - - is_stable = False - k = min(5, len(prev_order)) - if np.array_equal(prev_order[:k], curr_order[:k]): - is_stable = True - - improvement = np.abs(curr_scores - prev_scores).mean() - if improvement < self.min_improvement: - is_stable = True - - if is_stable: - self._stable_count += 1 - if self._stable_count >= self.patience: - return True - else: - self._stable_count = 0 - - return False diff --git a/scripts/rerank_recursive/expander.py b/scripts/rerank_recursive/expander.py deleted file mode 100644 index afc4d181..00000000 --- a/scripts/rerank_recursive/expander.py +++ /dev/null @@ -1,147 +0,0 @@ -""" -QueryExpander - Learns query expansions (synonyms/related terms) from usage patterns. -""" -import os -import re -from typing import Any, Dict, List, Tuple - -import numpy as np - -from scripts.rerank_recursive.utils import _COMMON_TOKENS - - -class QueryExpander: - """ - Learns query expansions (synonyms/related terms) from usage patterns. - - Observes which terms co-occur with successful retrievals and builds - a lightweight term→expansion mapping per-collection. - """ - - WEIGHTS_DIR = os.environ.get("RERANKER_WEIGHTS_DIR", "/tmp/rerank_weights") - MAX_EXPANSIONS_PER_TERM = 5 - MIN_CONFIDENCE = 0.3 - DECAY_RATE = 0.995 - - def __init__(self, lr: float = 0.1): - self.lr = lr - self._collection = "default" - self._weights_path = self._get_weights_path("default") - self.expansions: Dict[str, Dict[str, float]] = {} - self._update_count = 0 - self._version = 0 - - if os.path.exists(self._weights_path): - try: - self._load_weights() - except Exception: - pass - - @staticmethod - def _sanitize_collection(collection: str) -> str: - return "".join(c if c.isalnum() or c in "-_" else "_" for c in collection) - - def _get_weights_path(self, collection: str) -> str: - safe_name = self._sanitize_collection(collection) - return os.path.join(self.WEIGHTS_DIR, f"expander_{safe_name}.json") - - def set_collection(self, collection: str): - self._collection = collection - self._weights_path = self._get_weights_path(collection) - if os.path.exists(self._weights_path): - try: - self._load_weights() - except Exception: - pass - - def _load_weights(self): - import json - import fcntl - lock_path = self._weights_path + ".lock" - os.makedirs(os.path.dirname(lock_path) or ".", exist_ok=True) - with open(lock_path, "w") as lock_file: - fcntl.flock(lock_file.fileno(), fcntl.LOCK_SH) - with open(self._weights_path, "r") as f: - data = json.load(f) - self.expansions = data.get("expansions", {}) - self._version = data.get("version", 0) - fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN) - - def _save_weights(self): - import json - import fcntl - os.makedirs(os.path.dirname(self._weights_path) or ".", exist_ok=True) - lock_path = self._weights_path + ".lock" - tmp_path = self._weights_path + ".tmp" - with open(lock_path, "w") as lock_file: - fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX) - with open(tmp_path, "w") as f: - json.dump({"expansions": self.expansions, "version": self._version}, f) - os.replace(tmp_path, self._weights_path) - fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN) - - def _tokenize(self, text: str) -> List[str]: - tokens = re.findall(r'[a-zA-Z_][a-zA-Z0-9_]*', text.lower()) - return [t for t in tokens if len(t) > 2 and t not in _COMMON_TOKENS] - - def expand(self, query: str, max_expansions: int = 3) -> List[str]: - query_tokens = set(self._tokenize(query)) - candidates: List[Tuple[str, float]] = [] - for token in query_tokens: - if token in self.expansions: - for exp_term, conf in self.expansions[token].items(): - if exp_term not in query_tokens and conf >= self.MIN_CONFIDENCE: - candidates.append((exp_term, conf)) - candidates.sort(key=lambda x: -x[1]) - return [term for term, _ in candidates[:max_expansions]] - - def learn_from_teacher( - self, - query: str, - doc_texts: List[str], - teacher_scores: np.ndarray, - ): - query_tokens = set(self._tokenize(query)) - if not query_tokens: - return - - weights = np.exp(teacher_scores - teacher_scores.max()) - weights = weights / (weights.sum() + 1e-8) - - doc_term_weights: Dict[str, float] = {} - for doc_text, weight in zip(doc_texts, weights): - for token in self._tokenize(doc_text): - if token not in query_tokens: - doc_term_weights[token] = doc_term_weights.get(token, 0.0) + weight - - for query_term in query_tokens: - if query_term not in self.expansions: - self.expansions[query_term] = {} - - term_expansions = self.expansions[query_term] - - for exp in list(term_expansions.keys()): - term_expansions[exp] = float(term_expansions[exp] * self.DECAY_RATE) - if term_expansions[exp] < 0.01: - del term_expansions[exp] - - for doc_term, weight in doc_term_weights.items(): - if weight > 0.1: - old_conf = term_expansions.get(doc_term, 0.0) - new_conf = old_conf + self.lr * (weight - old_conf) - term_expansions[doc_term] = float(min(new_conf, 1.0)) - - if len(term_expansions) > self.MAX_EXPANSIONS_PER_TERM * 2: - sorted_exp = sorted(term_expansions.items(), key=lambda x: -x[1]) - self.expansions[query_term] = dict(sorted_exp[:self.MAX_EXPANSIONS_PER_TERM]) - - self._update_count += 1 - if self._update_count % 20 == 0: - self._version += 1 - self._save_weights() - - def get_stats(self) -> Dict[str, Any]: - total_terms = len(self.expansions) - total_expansions = sum(len(v) for v in self.expansions.values()) - avg_expansions = total_expansions / max(total_terms, 1) - return {"terms": total_terms, "expansions": total_expansions, "avg_per_term": avg_expansions, "version": self._version} diff --git a/scripts/rerank_recursive/hybrid_weights.py b/scripts/rerank_recursive/hybrid_weights.py deleted file mode 100644 index e24393b1..00000000 --- a/scripts/rerank_recursive/hybrid_weights.py +++ /dev/null @@ -1,105 +0,0 @@ -""" -LearnedHybridWeights - Learns optimal dense vs. lexical balance per-collection. -""" -import os - -import numpy as np - - -class LearnedHybridWeights: - """Learns optimal dense vs. lexical balance per-collection.""" - - WEIGHTS_DIR = os.environ.get("RERANKER_WEIGHTS_DIR", "/tmp/rerank_weights") - - def __init__(self, lr: float = 0.01): - self.lr = lr - self._collection = "default" - self._weights_path = self._get_weights_path("default") - self.alpha = 0.0 - self._momentum_alpha = 0.0 - self._momentum = 0.9 - self._update_count = 0 - self._version = 0 - - if os.path.exists(self._weights_path): - try: - self._load_weights() - except Exception: - pass - - @staticmethod - def _sanitize_collection(collection: str) -> str: - return "".join(c if c.isalnum() or c in "-_" else "_" for c in collection) - - def _get_weights_path(self, collection: str) -> str: - safe_name = self._sanitize_collection(collection) - return os.path.join(self.WEIGHTS_DIR, f"hybrid_{safe_name}.npz") - - def set_collection(self, collection: str): - self._collection = collection - self._weights_path = self._get_weights_path(collection) - if os.path.exists(self._weights_path): - try: - self._load_weights() - except Exception: - pass - - def _load_weights(self): - import fcntl - lock_path = self._weights_path + ".lock" - os.makedirs(os.path.dirname(lock_path) or ".", exist_ok=True) - with open(lock_path, "w") as lock_file: - fcntl.flock(lock_file.fileno(), fcntl.LOCK_SH) - data = np.load(self._weights_path) - self.alpha = float(data["alpha"]) - self._version = int(data.get("version", 0)) - fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN) - - def _save_weights(self): - import fcntl - os.makedirs(os.path.dirname(self._weights_path) or ".", exist_ok=True) - lock_path = self._weights_path + ".lock" - base_path = self._weights_path.rsplit(".npz", 1)[0] - tmp_base = base_path + ".tmp" - tmp_path = tmp_base + ".npz" - with open(lock_path, "w") as lock_file: - fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX) - np.savez(tmp_base, alpha=self.alpha, version=self._version) - os.replace(tmp_path, self._weights_path) - fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN) - - @property - def dense_weight(self) -> float: - return 1.0 / (1.0 + np.exp(-self.alpha)) - - @property - def lexical_weight(self) -> float: - return 1.0 - self.dense_weight - - def blend(self, dense_scores: np.ndarray, lexical_scores: np.ndarray) -> np.ndarray: - w = self.dense_weight - return w * dense_scores + (1 - w) * lexical_scores - - def learn_from_teacher( - self, - dense_scores: np.ndarray, - lexical_scores: np.ndarray, - teacher_scores: np.ndarray, - ): - w = self.dense_weight - blended = self.blend(dense_scores, lexical_scores) - teacher_norm = (teacher_scores - teacher_scores.mean()) / (teacher_scores.std() + 1e-8) - blended_norm = (blended - blended.mean()) / (blended.std() + 1e-8) - dense_norm = (dense_scores - dense_scores.mean()) / (dense_scores.std() + 1e-8) - lexical_norm = (lexical_scores - lexical_scores.mean()) / (lexical_scores.std() + 1e-8) - error = teacher_norm - blended_norm - modality_diff = dense_norm - lexical_norm - sigmoid_grad = w * (1 - w) - grad = (error * modality_diff).mean() * sigmoid_grad - self._momentum_alpha = self._momentum * self._momentum_alpha + grad - self.alpha += self.lr * self._momentum_alpha - self.alpha = np.clip(self.alpha, -5.0, 5.0) - self._update_count += 1 - if self._update_count % 50 == 0: - self._version += 1 - self._save_weights() diff --git a/scripts/rerank_recursive/projection.py b/scripts/rerank_recursive/projection.py deleted file mode 100644 index 70f811ee..00000000 --- a/scripts/rerank_recursive/projection.py +++ /dev/null @@ -1,207 +0,0 @@ -""" -LearnedProjection - Learnable linear projection from embedding dim to working dim. -""" -import os -import time -from typing import Any, Dict, Optional, Tuple - -import numpy as np - - -class LearnedProjection: - """ - Learnable linear projection from embedding dim to working dim. - - Replaces fixed random projection with a learnable layer that adapts - to domain-specific semantics. - """ - - WEIGHTS_DIR = os.environ.get("RERANKER_WEIGHTS_DIR", "/tmp/rerank_weights") - WEIGHTS_RELOAD_INTERVAL = float(os.environ.get("RERANKER_WEIGHTS_RELOAD_INTERVAL", "60")) - - def __init__(self, input_dim: int = 768, output_dim: int = 256, lr: float = 0.0005): - self.input_dim = input_dim - self.output_dim = output_dim - self.base_lr = lr - self.lr = lr - self._collection = "default" - self._weights_path = self._get_weights_path("default") - self._weights_mtime = 0.0 - self._last_reload_check = 0.0 - self._weights_loaded = False - - self._update_count = 0 - self._version = 0 - self._momentum_W: Optional[np.ndarray] = None - self._momentum = 0.9 - - if os.path.exists(self._weights_path): - try: - self._load_weights() - return - except Exception as e: - from scripts.logger import get_logger - get_logger(__name__).warning(f"LearnedProjection: failed to load {self._weights_path}: {e}") - - self._init_random_weights() - - @staticmethod - def _sanitize_collection(collection: str) -> str: - return "".join(c if c.isalnum() or c in "-_" else "_" for c in collection) - - def _get_weights_path(self, collection: str) -> str: - safe_name = self._sanitize_collection(collection) - return os.path.join(self.WEIGHTS_DIR, f"projection_{safe_name}.npz") - - def _init_random_weights(self): - scale = np.sqrt(2.0 / (self.input_dim + self.output_dim)) - rng = np.random.RandomState(44) - self.W = (rng.randn(self.input_dim, self.output_dim) * scale).astype(np.float32) - self._momentum_W = np.zeros_like(self.W) - - def init_from_pca(self, embeddings: np.ndarray): - """Initialize projection using PCA on corpus embeddings. - - Provides a much better cold-start than random initialization by using - the principal components of the actual data distribution. - - Args: - embeddings: (n_samples, input_dim) array of corpus embeddings - """ - if embeddings.shape[0] < self.output_dim: - from scripts.logger import get_logger - get_logger(__name__).warning( - f"LearnedProjection: {embeddings.shape[0]} samples < {self.output_dim} dims, using random init" - ) - self._init_random_weights() - return - - # Center the data - mean = embeddings.mean(axis=0, keepdims=True) - centered = embeddings - mean - - # Use SVD for numerical stability - # X = U @ S @ Vt, where Vt rows are principal components - U, S, Vt = np.linalg.svd(centered, full_matrices=False) - - # Take top output_dim principal components - # Vt is (min(n_samples, input_dim), input_dim) - components = Vt[:self.output_dim, :].T # (input_dim, output_dim) - - self.W = components.astype(np.float32) - self._momentum_W = np.zeros_like(self.W) - - # Compute explained variance for logging - total_var = np.sum(S ** 2) - explained_var = np.sum(S[:self.output_dim] ** 2) / total_var if total_var > 0 else 0 - - from scripts.logger import get_logger - get_logger(__name__).info( - f"LearnedProjection: PCA init on {embeddings.shape[0]} samples " - f"({self.input_dim}→{self.output_dim}), explained variance: {explained_var:.1%}" - ) - - def set_collection(self, collection: str): - self._collection = collection - self._weights_path = self._get_weights_path(collection) - if os.path.exists(self._weights_path): - try: - self._load_weights() - except Exception: - pass - - def maybe_reload_weights(self): - # Fast path: skip if reload disabled (interval <= 0) - if self.WEIGHTS_RELOAD_INTERVAL <= 0: - return - now = time.time() - if now - self._last_reload_check < self.WEIGHTS_RELOAD_INTERVAL: - return - self._last_reload_check = now - try: - if os.path.exists(self._weights_path): - mtime = os.path.getmtime(self._weights_path) - if mtime > self._weights_mtime: - self._load_weights() - except Exception: - pass - - def _load_weights(self): - import fcntl - lock_path = self._weights_path + ".lock" - try: - os.makedirs(os.path.dirname(lock_path) or ".", exist_ok=True) - with open(lock_path, "w") as lock_file: - fcntl.flock(lock_file.fileno(), fcntl.LOCK_SH) - data = np.load(self._weights_path) - self.W = data["W"].astype(np.float32) - self._version = int(data.get("version", 0)) - fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN) - self._weights_mtime = os.path.getmtime(self._weights_path) - self._weights_loaded = True - self._momentum_W = np.zeros_like(self.W) - except Exception as e: - from scripts.logger import get_logger - get_logger(__name__).warning(f"LearnedProjection: load failed: {e}") - - def _save_weights(self): - import fcntl - os.makedirs(os.path.dirname(self._weights_path) or ".", exist_ok=True) - lock_path = self._weights_path + ".lock" - base_path = self._weights_path.rsplit(".npz", 1)[0] - tmp_path = base_path + ".tmp.npz" - try: - with open(lock_path, "w") as lock_file: - fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX) - np.savez(tmp_path, W=self.W, version=self._version) - os.replace(tmp_path, self._weights_path) - fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN) - self._weights_mtime = os.path.getmtime(self._weights_path) - except Exception as e: - from scripts.logger import get_logger - get_logger(__name__).warning(f"LearnedProjection: save failed: {e}") - - def forward(self, embeddings: np.ndarray) -> np.ndarray: - """Project embeddings to output dim (normalized).""" - self.maybe_reload_weights() # Hot-reload from worker updates - squeeze = embeddings.ndim == 1 - if squeeze: - embeddings = embeddings.reshape(1, -1) - projected = embeddings @ self.W - norms = np.linalg.norm(projected, axis=-1, keepdims=True) + 1e-8 - projected = projected / norms - if squeeze: - projected = projected[0] - return projected - - def forward_with_cache(self, embeddings: np.ndarray) -> Tuple[np.ndarray, Dict[str, Any]]: - """Forward pass with cache for backprop.""" - squeeze = embeddings.ndim == 1 - if squeeze: - embeddings = embeddings.reshape(1, -1) - pre_norm = embeddings @ self.W - norms = np.linalg.norm(pre_norm, axis=-1, keepdims=True) + 1e-8 - projected = pre_norm / norms - cache = {"input": embeddings, "pre_norm": pre_norm, "norms": norms, "projected": projected} - if squeeze: - projected = projected[0] - return projected, cache - - def backward(self, grad_output: np.ndarray, cache: Dict[str, Any], weight: float = 1.0): - """Backprop gradient through projection and update weights.""" - if grad_output.ndim == 1: - grad_output = grad_output.reshape(1, -1) - embeddings = cache["input"] - norms = cache["norms"] - batch_size = embeddings.shape[0] - projected = cache["projected"] - dot = np.sum(grad_output * projected, axis=-1, keepdims=True) - grad_pre_norm = (grad_output - projected * dot) / norms - dW = embeddings.T @ grad_pre_norm / batch_size - dW = dW * weight - self._momentum_W = self._momentum * self._momentum_W + dW - self.W -= self.lr * self._momentum_W - self._update_count += 1 - if self._update_count % 100 == 0: - self._version += 1 - self._save_weights() diff --git a/scripts/rerank_recursive/recursive.py b/scripts/rerank_recursive/recursive.py deleted file mode 100644 index 71bbee7c..00000000 --- a/scripts/rerank_recursive/recursive.py +++ /dev/null @@ -1,924 +0,0 @@ -""" -RecursiveReranker, ONNXRecursiveReranker, SessionAwareReranker - Main reranking pipelines. - -Implements TRM-style iterative refinement: -1. Initialize latent state z from query -2. For each iteration: score, refine z, check early stopping -3. Return final ranking -""" -import os -import threading -import time -from typing import Any, Dict, List, Optional - -import numpy as np - -# Safe ONNX imports -try: - import onnxruntime as ort - from tokenizers import Tokenizer - HAS_ONNX = True -except ImportError: - ort = None - Tokenizer = None - HAS_ONNX = False - -# Use centralized reranker factory (supports FastEmbed + ONNX backends) -try: - from scripts.reranker import ( - get_reranker_model as _get_reranker_model, - rerank_pairs as _rerank_pairs, - is_reranker_available as _is_reranker_available, - RERANKER_MODEL, - ) - HAS_RERANKER_FACTORY = True -except ImportError: - HAS_RERANKER_FACTORY = False - _get_reranker_model = None - _rerank_pairs = None - _is_reranker_available = None - RERANKER_MODEL = None - -# Legacy: direct FastEmbed imports (fallback when factory unavailable) -try: - from fastembed.rerank.cross_encoder import TextCrossEncoder - HAS_FASTEMBED_RERANK = True -except ImportError: - TextCrossEncoder = None - HAS_FASTEMBED_RERANK = False - -from scripts.rerank_recursive.state import RefinementState -from scripts.rerank_recursive.scorer import TinyScorer -from scripts.rerank_recursive.refiner import LatentRefiner -from scripts.rerank_recursive.projection import LearnedProjection -from scripts.rerank_recursive.confidence import ConfidenceEstimator -from scripts.rerank_recursive.utils import ( - _compute_fname_boost, - _get_cached_embedding, - _cache_embedding, -) - - -class RecursiveReranker: - """ - Main recursive reranking pipeline. - - Key insight: Multiple passes through tiny networks > one pass through large network - """ - - def __init__( - self, - n_iterations: int = 3, - dim: int = 256, - hidden_dim: int = 512, - early_stop: bool = True, - blend_with_initial: float = 0.3, - alpha_scheduler: Optional[Any] = None, - ): - self.n_iterations = n_iterations - self.dim = dim - self.early_stop = early_stop - self.blend_with_initial = blend_with_initial - - # Alpha scheduler: if None, use CosineAlphaScheduler by default - if alpha_scheduler is None: - from scripts.rerank_recursive.alpha_scheduler import CosineAlphaScheduler - self.alpha_scheduler = CosineAlphaScheduler(n_iterations=n_iterations) - else: - self.alpha_scheduler = alpha_scheduler - - self.scorer = TinyScorer(dim=dim, hidden_dim=hidden_dim) - self.refiner = LatentRefiner(dim=dim) - - from scripts.embedder import get_model_dimension - embed_dim = get_model_dimension() - self._learned_projection = LearnedProjection(input_dim=embed_dim, output_dim=dim, lr=0.0) - self._embedder = None - self._embedder_lock = threading.Lock() - self._proj_cache: Dict[int, np.ndarray] = {} - self._proj_cache_lock = threading.Lock() - - # Observability counters for learning status - self._proj_learned_count = 0 - self._proj_fallback_count = 0 - self._fallback_warned = False - - def _get_embedder(self): - if self._embedder is not None: - return self._embedder - with self._embedder_lock: - if self._embedder is not None: - return self._embedder - try: - from scripts.embedder import get_embedding_model - model_name = os.environ.get("EMBEDDING_MODEL", "BAAI/bge-base-en-v1.5") - self._embedder = get_embedding_model(model_name) - except Exception: - self._embedder = None - return self._embedder - - def _encode(self, texts: List[str]) -> np.ndarray: - cached_results = [] - texts_to_encode = [] - text_indices = [] - - for i, text in enumerate(texts): - cached = _get_cached_embedding(text) - if cached is not None: - cached_results.append((i, cached)) - else: - texts_to_encode.append(text) - text_indices.append(i) - - new_embeddings = [] - if texts_to_encode: - embedder = self._get_embedder() - if embedder is not None: - try: - embeddings = list(embedder.embed(texts_to_encode)) - if len(embeddings) != len(texts_to_encode): - raise ValueError("Embedding count mismatch") - for text, emb in zip(texts_to_encode, embeddings): - emb_arr = np.array(emb, dtype=np.float32) - if emb_arr.shape[0] != self.dim: - emb_arr = self._project_to_dim(emb_arr.reshape(1, -1))[0] - _cache_embedding(text, emb_arr) - new_embeddings.append(emb_arr) - except Exception: - new_embeddings = [] - - if not new_embeddings: - import hashlib - fallback_dim = self.dim - if cached_results: - fallback_dim = cached_results[0][1].shape[0] - for text in texts_to_encode: - text_hash = hashlib.sha256(text.encode("utf-8", errors="replace")).digest() - seed = int.from_bytes(text_hash[:4], "big") - rng = np.random.RandomState(seed) - vec = rng.randn(fallback_dim).astype(np.float32) - vec = vec / (np.linalg.norm(vec) + 1e-8) - _cache_embedding(text, vec) - new_embeddings.append(vec) - - result = [None] * len(texts) - for i, emb in cached_results: - result[i] = emb - for i, idx in enumerate(text_indices): - result[idx] = new_embeddings[i] - return np.array(result, dtype=np.float32) - - def _encode_raw(self, texts: List[str]) -> np.ndarray: - from scripts.embedder import get_model_dimension - fallback_dim = get_model_dimension() - embedder = self._get_embedder() - if embedder is None: - import hashlib - result = [] - for text in texts: - text_hash = hashlib.sha256(text.encode("utf-8", errors="replace")).digest() - seed = int.from_bytes(text_hash[:4], "big") - rng = np.random.RandomState(seed) - vec = rng.randn(fallback_dim).astype(np.float32) - vec = vec / (np.linalg.norm(vec) + 1e-8) - result.append(vec) - return np.array(result, dtype=np.float32) - try: - embeddings = list(embedder.embed(texts)) - result = [np.array(emb, dtype=np.float32) for emb in embeddings] - return np.array(result, dtype=np.float32) - except Exception: - import hashlib - result = [] - for text in texts: - text_hash = hashlib.sha256(text.encode("utf-8", errors="replace")).digest() - seed = int.from_bytes(text_hash[:4], "big") - rng = np.random.RandomState(seed) - vec = rng.randn(fallback_dim).astype(np.float32) - vec = vec / (np.linalg.norm(vec) + 1e-8) - result.append(vec) - return np.array(result, dtype=np.float32) - - def _project_to_dim(self, embeddings: np.ndarray) -> np.ndarray: - if embeddings.shape[-1] == self.dim: - return embeddings - input_dim = embeddings.shape[-1] - # Try to use learned projection if available (hot-reloads from worker) - if hasattr(self, '_learned_projection') and self._learned_projection.input_dim == input_dim: - # forward() calls maybe_reload_weights() internally - if self._learned_projection._weights_loaded: - self._proj_learned_count += 1 - return self._learned_projection.forward(embeddings) - # Try one reload attempt if not yet loaded - self._learned_projection.maybe_reload_weights() - if self._learned_projection._weights_loaded: - self._proj_learned_count += 1 - return self._learned_projection.forward(embeddings) - # Fallback to random projection - self._proj_fallback_count += 1 - if not self._fallback_warned and self._proj_fallback_count >= 10: - from scripts.logger import get_logger - get_logger(__name__).warning( - f"LearnedProjection not loaded after {self._proj_fallback_count} calls; " - f"using random projection (learning may be disabled or weights missing)" - ) - self._fallback_warned = True - with self._proj_cache_lock: - if input_dim not in self._proj_cache: - rng = np.random.RandomState(44) - proj_matrix = rng.randn(input_dim, self.dim).astype(np.float32) * np.float32(0.01) - self._proj_cache[input_dim] = proj_matrix - proj_matrix = self._proj_cache[input_dim] - projected = embeddings @ proj_matrix - norms = np.linalg.norm(projected, axis=-1, keepdims=True) + 1e-8 - return projected / norms - - def get_learning_status(self) -> Dict[str, Any]: - """Return observability info for learning components.""" - scorer_metrics = self.scorer.get_metrics() - proj_loaded = ( - hasattr(self, '_learned_projection') and - self._learned_projection._weights_loaded - ) - return { - "projection_loaded": proj_loaded, - "projection_version": self._learned_projection._version if proj_loaded else 0, - "projection_learned_calls": self._proj_learned_count, - "projection_fallback_calls": self._proj_fallback_count, - "scorer_version": scorer_metrics.get("version", 0), - "scorer_converged": scorer_metrics.get("converged", False), - "scorer_avg_loss": scorer_metrics.get("avg_loss", 0.0), - "scorer_update_count": scorer_metrics.get("update_count", 0), - "refiner_version": self.refiner._version, - } - - def rerank( - self, - query: str, - candidates: List[Dict[str, Any]], - initial_scores: Optional[List[float]] = None, - ) -> List[Dict[str, Any]]: - if not candidates: - return [] - - confidence = ConfidenceEstimator() - n_docs = len(candidates) - - doc_texts = [] - for c in candidates: - text_parts = [] - if c.get("symbol"): - text_parts.append(str(c["symbol"])) - if c.get("path"): - text_parts.append(str(c["path"])) - code = c.get("code") or c.get("snippet") or c.get("text") or "" - if code: - text_parts.append(str(code)[:500]) - doc_texts.append(" ".join(text_parts) if text_parts else "empty") - - query_emb = self._encode([query])[0] - doc_embs = self._encode(doc_texts) - query_emb = self._project_to_dim(query_emb.reshape(1, -1))[0] - doc_embs = self._project_to_dim(doc_embs) - - z = query_emb.copy() - if initial_scores is not None: - scores = np.array(initial_scores, dtype=np.float32) - else: - scores = np.zeros(n_docs, dtype=np.float32) - - state = RefinementState(z=z, scores=scores, iteration=0) - state.score_history.append(scores.copy()) - alpha_trajectory = [] # Track alpha values used - - for i in range(self.n_iterations): - state.iteration = i + 1 - new_scores = self.scorer.forward(query_emb, doc_embs, state.z) - alpha = self.alpha_scheduler.get_alpha(i) - alpha_trajectory.append(alpha) - state.scores = alpha * new_scores + (1 - alpha) * state.scores - state.score_history.append(state.scores.copy()) - state.z = self.refiner.refine(state.z, query_emb, doc_embs, state.scores) - if self.early_stop and confidence.should_stop(state): - break - - final_scores = state.scores - if initial_scores is not None and self.blend_with_initial > 0: - init_arr = np.array(initial_scores, dtype=np.float32) - std = final_scores.std() - if std > 1e-6: - final_norm = (final_scores - final_scores.mean()) / std - else: - final_norm = final_scores - final_scores.mean() - std = init_arr.std() - if std > 1e-6: - init_norm = (init_arr - init_arr.mean()) / std - else: - init_norm = init_arr - init_arr.mean() - final_scores = (1 - self.blend_with_initial) * final_norm + self.blend_with_initial * init_norm - - ranked_indices = np.argsort(-final_scores) - reranked = [] - fname_boost_factor = float(os.environ.get("FNAME_BOOST", "0.15") or 0.15) - - for rank, idx in enumerate(ranked_indices): - candidate = candidates[idx].copy() - candidate["recursive_score"] = float(final_scores[idx]) - candidate["recursive_rank"] = rank - candidate["recursive_iterations"] = state.iteration - trajectory = [float(h[idx]) for h in state.score_history] - candidate["score_trajectory"] = trajectory - candidate["alpha_trajectory"] = alpha_trajectory[:state.iteration] # Alpha values used - fname_boost = _compute_fname_boost(query, candidate, fname_boost_factor) - candidate["score"] = float(final_scores[idx]) + fname_boost - if fname_boost > 0: - candidate["fname_boost"] = fname_boost - reranked.append(candidate) - - if fname_boost_factor > 0 and any(c.get("fname_boost", 0) > 0 for c in reranked): - reranked.sort(key=lambda x: -x["score"]) - - return reranked - - -class ONNXRecursiveReranker(RecursiveReranker): - """Recursive reranker using ONNX cross-encoder for scoring.""" - - def __init__( - self, - n_iterations: int = 3, - onnx_path: Optional[str] = None, - tokenizer_path: Optional[str] = None, - **kwargs - ): - super().__init__(n_iterations=n_iterations, **kwargs) - self.onnx_path = onnx_path or os.environ.get("RERANKER_ONNX_PATH", "") - self.tokenizer_path = tokenizer_path or os.environ.get("RERANKER_TOKENIZER_PATH", "") - self._session = None - self._tokenizer = None - self._onnx_lock = threading.Lock() - - def _get_onnx_session(self): - if self._session is not None: - return self._session, self._tokenizer - if not HAS_ONNX or not self.onnx_path or not self.tokenizer_path: - return None, None - with self._onnx_lock: - if self._session is not None: - return self._session, self._tokenizer - try: - tok = Tokenizer.from_file(self.tokenizer_path) - try: - tok.enable_truncation(max_length=512) - except Exception: - pass - sess = ort.InferenceSession(self.onnx_path, providers=["CPUExecutionProvider"]) - self._session, self._tokenizer = sess, tok - except Exception: - self._session, self._tokenizer = None, None - return self._session, self._tokenizer - - def _onnx_score(self, query: str, docs: List[str]) -> Optional[np.ndarray]: - sess, tok = self._get_onnx_session() - if sess is None or tok is None: - return None - try: - pairs = [(query, doc) for doc in docs] - enc = tok.encode_batch(pairs) - input_ids = [e.ids for e in enc] - attn = [e.attention_mask for e in enc] - max_len = max((len(ids) for ids in input_ids), default=0) - if max_len == 0: - return None - pad_id = 0 - try: - pad_token_id = tok.token_to_id("[PAD]") - if pad_token_id is not None: - pad_id = int(pad_token_id) - except Exception: - pad_id = 0 - - def pad(seq, pad_val): - return seq + [pad_val] * (max_len - len(seq)) - - input_ids_padded = [pad(s, pad_id) for s in input_ids] - attn_padded = [pad(s, 0) for s in attn] - input_ids_arr = np.array(input_ids_padded, dtype=np.int64) - attn_arr = np.array(attn_padded, dtype=np.int64) - input_names = [i.name for i in sess.get_inputs()] - feeds = {} - if "input_ids" in input_names: - feeds["input_ids"] = input_ids_arr - if "attention_mask" in input_names: - feeds["attention_mask"] = attn_arr - if "token_type_ids" in input_names: - token_type_arr = np.zeros((len(input_ids_padded), max_len), dtype=np.int64) - feeds["token_type_ids"] = token_type_arr - out = sess.run(None, feeds) - logits = out[0] - scores = [] - for row in logits: - try: - if hasattr(row, "__len__") and len(row) >= 2: - scores.append(float(row[1])) - elif hasattr(row, "__len__") and len(row) == 1: - scores.append(float(row[0])) - else: - scores.append(float(row)) - except Exception: - scores.append(0.0) - return np.array(scores, dtype=np.float32) - except Exception: - return None - - def rerank( - self, - query: str, - candidates: List[Dict[str, Any]], - initial_scores: Optional[List[float]] = None, - ) -> List[Dict[str, Any]]: - if not candidates: - return [] - - confidence = ConfidenceEstimator() - - doc_texts = [] - for c in candidates: - parts = [] - if c.get("symbol"): - parts.append(str(c["symbol"])) - if c.get("path"): - parts.append(str(c["path"])) - code = c.get("code") or c.get("snippet") or c.get("text") or "" - if code: - parts.append(str(code)[:400]) - doc_texts.append(" ".join(parts) if parts else "empty") - - onnx_scores = self._onnx_score(query, doc_texts) - if onnx_scores is None: - return super().rerank(query, candidates, initial_scores) - - scores = onnx_scores.copy() - query_emb = self._encode([query])[0] - doc_embs = self._encode(doc_texts) - query_emb = self._project_to_dim(query_emb.reshape(1, -1))[0] - doc_embs = self._project_to_dim(doc_embs) - - z = query_emb.copy() - state = RefinementState(z=z, scores=scores, iteration=0) - state.score_history.append(scores.copy()) - - for i in range(self.n_iterations - 1): - state.iteration = i + 1 - state.z = self.refiner.refine(state.z, query_emb, doc_embs, state.scores) - adjustment = self.scorer.forward(query_emb, doc_embs, state.z) - try: - metrics = self.scorer.get_metrics() - if metrics.get("converged", False) and metrics.get("avg_loss", 1.0) < 0.3: - alpha = 0.5 - elif metrics.get("update_count", 0) > 100: - alpha = 0.35 - else: - alpha = 0.2 - except Exception: - alpha = 0.2 - state.scores = (1 - alpha) * state.scores + alpha * adjustment - state.score_history.append(state.scores.copy()) - if self.early_stop and confidence.should_stop(state): - break - - final_scores = state.scores - if initial_scores is not None and self.blend_with_initial > 0: - init_arr = np.array(initial_scores, dtype=np.float32) - std = final_scores.std() - if std > 1e-6: - final_norm = (final_scores - final_scores.mean()) / std - else: - final_norm = final_scores - final_scores.mean() - std = init_arr.std() - if std > 1e-6: - init_norm = (init_arr - init_arr.mean()) / std - else: - init_norm = init_arr - init_arr.mean() - final_scores = (1 - self.blend_with_initial) * final_norm + self.blend_with_initial * init_norm - - ranked_indices = np.argsort(-final_scores) - reranked = [] - for rank, idx in enumerate(ranked_indices): - candidate = candidates[idx].copy() - candidate["recursive_score"] = float(final_scores[idx]) - candidate["onnx_score"] = float(onnx_scores[idx]) - candidate["recursive_rank"] = rank - candidate["recursive_iterations"] = state.iteration + 1 - candidate["score_trajectory"] = [float(h[idx]) for h in state.score_history] - candidate["score"] = float(final_scores[idx]) - reranked.append(candidate) - return reranked - - -class FastEmbedRecursiveReranker(RecursiveReranker): - """Recursive reranker using FastEmbed cross-encoder (via reranker factory). - - Uses the centralized reranker factory which supports both RERANKER_MODEL - (FastEmbed auto-download) and legacy RERANKER_ONNX_PATH configs. - """ - - def __init__(self, n_iterations: int = 3, **kwargs): - super().__init__(n_iterations=n_iterations, **kwargs) - self._reranker_model = None - self._model_lock = threading.Lock() - - def _get_model(self): - """Get cached reranker model from factory.""" - if self._reranker_model is not None: - return self._reranker_model - if not HAS_RERANKER_FACTORY or _get_reranker_model is None: - return None - with self._model_lock: - if self._reranker_model is not None: - return self._reranker_model - self._reranker_model = _get_reranker_model() - return self._reranker_model - - def _factory_score(self, query: str, docs: List[str]) -> Optional[np.ndarray]: - """Score documents using reranker factory.""" - model = self._get_model() - if model is None or _rerank_pairs is None: - return None - try: - pairs = [(query, doc) for doc in docs] - scores = _rerank_pairs(pairs, model=model) - return np.array(scores, dtype=np.float32) - except Exception: - return None - - def rerank( - self, - query: str, - candidates: List[Dict[str, Any]], - initial_scores: Optional[List[float]] = None, - ) -> List[Dict[str, Any]]: - if not candidates: - return [] - - confidence = ConfidenceEstimator() - - doc_texts = [] - for c in candidates: - parts = [] - if c.get("symbol"): - parts.append(str(c["symbol"])) - if c.get("path"): - parts.append(str(c["path"])) - code = c.get("code") or c.get("snippet") or c.get("text") or "" - if code: - parts.append(str(code)[:400]) - doc_texts.append(" ".join(parts) if parts else "empty") - - factory_scores = self._factory_score(query, doc_texts) - if factory_scores is None: - return super().rerank(query, candidates, initial_scores) - - scores = factory_scores.copy() - query_emb = self._encode([query])[0] - doc_embs = self._encode(doc_texts) - query_emb = self._project_to_dim(query_emb.reshape(1, -1))[0] - doc_embs = self._project_to_dim(doc_embs) - - z = query_emb.copy() - state = RefinementState(z=z, scores=scores, iteration=0) - state.score_history.append(scores.copy()) - - for i in range(self.n_iterations - 1): - state.iteration = i + 1 - state.z = self.refiner.refine(state.z, query_emb, doc_embs, state.scores) - adjustment = self.scorer.forward(query_emb, doc_embs, state.z) - try: - metrics = self.scorer.get_metrics() - if metrics.get("converged", False) and metrics.get("avg_loss", 1.0) < 0.3: - alpha = 0.5 - elif metrics.get("update_count", 0) > 100: - alpha = 0.35 - else: - alpha = 0.2 - except Exception: - alpha = 0.2 - state.scores = (1 - alpha) * state.scores + alpha * adjustment - state.score_history.append(state.scores.copy()) - if self.early_stop and confidence.should_stop(state): - break - - final_scores = state.scores - if initial_scores is not None and self.blend_with_initial > 0: - init_arr = np.array(initial_scores, dtype=np.float32) - std = final_scores.std() - if std > 1e-6: - final_norm = (final_scores - final_scores.mean()) / std - else: - final_norm = final_scores - final_scores.mean() - std = init_arr.std() - if std > 1e-6: - init_norm = (init_arr - init_arr.mean()) / std - else: - init_norm = init_arr - init_arr.mean() - final_scores = (1 - self.blend_with_initial) * final_norm + self.blend_with_initial * init_norm - - ranked_indices = np.argsort(-final_scores) - reranked = [] - for rank, idx in enumerate(ranked_indices): - candidate = candidates[idx].copy() - candidate["recursive_score"] = float(final_scores[idx]) - candidate["factory_score"] = float(factory_scores[idx]) - candidate["recursive_rank"] = rank - candidate["recursive_iterations"] = state.iteration + 1 - candidate["score_trajectory"] = [float(h[idx]) for h in state.score_history] - candidate["score"] = float(final_scores[idx]) - reranked.append(candidate) - return reranked - - -class SessionAwareReranker: - """Session-aware recursive reranker with latent state carryover.""" - - def __init__( - self, - n_iterations: int = 3, - dim: int = 256, - session_decay: float = 0.9, - max_session_age: float = 3600.0, - max_sessions: int = 1000, - ): - self.n_iterations = n_iterations - self.dim = dim - self.session_decay = session_decay - self.max_session_age = max_session_age - self.max_sessions = max_sessions - self.reranker = RecursiveReranker(n_iterations=n_iterations, dim=dim) - self._sessions: Dict[str, tuple] = {} - self._session_lock = threading.Lock() - - def _cleanup_old_sessions(self): - now = time.time() - expired = [sid for sid, (_, last_access) in self._sessions.items() if now - last_access > self.max_session_age] - for sid in expired: - del self._sessions[sid] - if len(self._sessions) > self.max_sessions: - sorted_sessions = sorted(self._sessions.items(), key=lambda x: x[1][1]) - to_remove = len(self._sessions) - self.max_sessions - for sid, _ in sorted_sessions[:to_remove]: - del self._sessions[sid] - - def get_session_latent(self, session_id: str) -> Optional[np.ndarray]: - with self._session_lock: - if session_id not in self._sessions: - return None - latent, last_access = self._sessions[session_id] - if time.time() - last_access > self.max_session_age: - del self._sessions[session_id] - return None - return latent - - def update_session_latent(self, session_id: str, new_latent: np.ndarray): - with self._session_lock: - self._cleanup_old_sessions() - if session_id in self._sessions: - old_latent, _ = self._sessions[session_id] - blended = self.session_decay * old_latent + (1 - self.session_decay) * new_latent - blended = blended / (np.linalg.norm(blended) + 1e-8) - self._sessions[session_id] = (blended, time.time()) - else: - self._sessions[session_id] = (new_latent.copy(), time.time()) - - def rerank( - self, - query: str, - candidates: List[Dict[str, Any]], - session_id: Optional[str] = None, - initial_scores: Optional[List[float]] = None, - ) -> List[Dict[str, Any]]: - if not candidates: - return [] - - session_latent = None - if session_id: - session_latent = self.get_session_latent(session_id) - - query_emb = self.reranker._encode([query])[0] - query_emb = self.reranker._project_to_dim(query_emb.reshape(1, -1))[0] - - if session_latent is not None: - initial_z = 0.7 * query_emb + 0.3 * session_latent - initial_z = initial_z / (np.linalg.norm(initial_z) + 1e-8) - else: - initial_z = query_emb.copy() - - doc_texts = [] - for c in candidates: - text_parts = [] - if c.get("symbol"): - text_parts.append(str(c["symbol"])) - if c.get("path"): - text_parts.append(str(c["path"])) - code = c.get("code") or c.get("snippet") or c.get("text") or "" - if code: - text_parts.append(str(code)[:500]) - doc_texts.append(" ".join(text_parts) if text_parts else "empty") - - doc_embs = self.reranker._encode(doc_texts) - doc_embs = self.reranker._project_to_dim(doc_embs) - - if initial_scores is None: - initial_scores = [c.get("score", 0.0) for c in candidates] - - state = RefinementState(z=initial_z, scores=np.array(initial_scores, dtype=np.float32)) - state.score_history.append(state.scores.copy()) - - confidence = ConfidenceEstimator() - for i in range(self.n_iterations): - state.iteration = i + 1 - new_scores = self.reranker.scorer.forward(query_emb, doc_embs, state.z) - alpha = 0.5 - state.scores = alpha * new_scores + (1 - alpha) * state.scores - state.score_history.append(state.scores.copy()) - state.z = self.reranker.refiner.refine(state.z, query_emb, doc_embs, state.scores) - if self.reranker.early_stop and confidence.should_stop(state): - break - - if session_id: - self.update_session_latent(session_id, state.z) - - final_scores = state.scores - if self.reranker.blend_with_initial > 0: - init_arr = np.array(initial_scores, dtype=np.float32) - std = final_scores.std() - if std > 1e-6: - final_norm = (final_scores - final_scores.mean()) / std - else: - final_norm = final_scores - final_scores.mean() - std = init_arr.std() - if std > 1e-6: - init_norm = (init_arr - init_arr.mean()) / std - else: - init_norm = init_arr - init_arr.mean() - final_scores = (1 - self.reranker.blend_with_initial) * final_norm + self.reranker.blend_with_initial * init_norm - - ranked_indices = np.argsort(-final_scores) - reranked = [] - for rank, idx in enumerate(ranked_indices): - candidate = candidates[idx].copy() - candidate["recursive_score"] = float(final_scores[idx]) - candidate["recursive_rank"] = rank - candidate["recursive_iterations"] = state.iteration - candidate["session_aware"] = session_id is not None - candidate["score_trajectory"] = [float(h[idx]) for h in state.score_history] - candidate["score"] = float(final_scores[idx]) - reranked.append(candidate) - return reranked - - def clear_session(self, session_id: str): - with self._session_lock: - if session_id in self._sessions: - del self._sessions[session_id] - - def get_session_count(self) -> int: - with self._session_lock: - return len(self._sessions) - - -# Convenience functions -def rerank_recursive( - query: str, - candidates: List[Dict[str, Any]], - n_iterations: int = 3, - blend_with_initial: float = 0.3, -) -> List[Dict[str, Any]]: - """Convenience wrapper for recursive reranking.""" - reranker = RecursiveReranker(n_iterations=n_iterations, blend_with_initial=blend_with_initial) - initial_scores = [c.get("score", 0.0) for c in candidates] - return reranker.rerank(query, candidates, initial_scores) - - -def rerank_recursive_inprocess( - query: str, - candidates: List[Dict[str, Any]], - limit: int = 12, - n_iterations: int = 3, -) -> List[Dict[str, Any]]: - """In-process recursive reranking for MCP server integration.""" - reranked = rerank_recursive(query=query, candidates=candidates, n_iterations=n_iterations) - return reranked[:limit] - - -# Per-collection learning rerankers -_LEARNING_RERANKERS: Dict[str, RecursiveReranker] = {} -_LEARNING_RERANKERS_LOCK = threading.Lock() - - -def _get_learning_reranker( - n_iterations: int = 3, - dim: int = 256, - collection: str = "default", -) -> RecursiveReranker: - """Get or create a learning reranker for a specific collection.""" - with _LEARNING_RERANKERS_LOCK: - if collection not in _LEARNING_RERANKERS: - reranker = RecursiveReranker(n_iterations=n_iterations, dim=dim) - reranker.scorer.set_collection(collection) - reranker.refiner.set_collection(collection) - reranker._learned_projection.set_collection(collection) - _LEARNING_RERANKERS[collection] = reranker - return _LEARNING_RERANKERS[collection] - - -def rerank_with_learning( - query: str, - candidates: List[Dict[str, Any]], - limit: int = 12, - n_iterations: int = 3, - learn_from_onnx: bool = True, - collection: str = "default", -) -> List[Dict[str, Any]]: - """Learning-enabled reranking for MCP server integration.""" - reranker = _get_learning_reranker(n_iterations=n_iterations, collection=collection) - initial_scores = [c.get("score", 0) for c in candidates] - - if learn_from_onnx and candidates: - teacher_scores = None - if str(os.environ.get("RERANK_TEACHER_INLINE", "")).strip().lower() in {"1", "true", "yes", "on"}: - try: - from scripts.rerank_local import rerank_local - except ImportError: - try: - from rerank_local import rerank_local - except ImportError: - rerank_local = None - if rerank_local is not None: - try: - pairs = [] - for c in candidates: - doc = c.get("code") or c.get("snippet") or "" - if not doc: - parts = [] - if c.get("symbol"): - parts.append(str(c["symbol"])) - if c.get("path"): - parts.append(str(c["path"])) - doc = " ".join(parts) if parts else "empty" - pairs.append((query, doc[:1000])) - teacher_scores = rerank_local(pairs) - except Exception: - teacher_scores = None - try: - try: - from rerank_events import log_training_event - except ImportError: - from scripts.rerank_events import log_training_event - log_training_event( - query=query, - candidates=candidates, - initial_scores=initial_scores, - teacher_scores=(list(teacher_scores) if teacher_scores is not None else None), - collection=collection, - metadata={"teacher_inline": bool(teacher_scores is not None)}, - ) - except Exception: - pass - - reranked = reranker.rerank(query, candidates, initial_scores) - return reranked[:limit] - - -def get_recursive_reranker(n_iterations: int = 3, **kwargs) -> RecursiveReranker: - """Get the best available recursive reranker. - - Priority: - 1. RERANKER_MODEL (FastEmbed) via factory -> FastEmbedRecursiveReranker - 2. RERANKER_ONNX_PATH + RERANKER_TOKENIZER_PATH -> ONNXRecursiveReranker - 3. Fallback -> base RecursiveReranker (no neural reranking) - - Backwards compatible: existing ONNX configs continue to work. - """ - # Priority 1: Use factory if RERANKER_MODEL is set - if HAS_RERANKER_FACTORY and _is_reranker_available is not None: - if _is_reranker_available(): - return FastEmbedRecursiveReranker(n_iterations=n_iterations, **kwargs) - - # Priority 2: Legacy ONNX path (backwards compatibility) - onnx_path = os.environ.get("RERANKER_ONNX_PATH", "") - tokenizer_path = os.environ.get("RERANKER_TOKENIZER_PATH", "") - if HAS_ONNX and onnx_path and tokenizer_path: - return ONNXRecursiveReranker(n_iterations=n_iterations, **kwargs) - - # Priority 3: Base reranker (no neural scoring) - return RecursiveReranker(n_iterations=n_iterations, **kwargs) - - -def rerank_with_session( - query: str, - candidates: List[Dict[str, Any]], - session_id: str, - n_iterations: int = 3, -) -> List[Dict[str, Any]]: - """Session-aware reranking (stateless convenience wrapper).""" - reranker = SessionAwareReranker(n_iterations=n_iterations) - return reranker.rerank(query, candidates, session_id=session_id) diff --git a/scripts/rerank_recursive/refiner.py b/scripts/rerank_recursive/refiner.py deleted file mode 100644 index c73b6200..00000000 --- a/scripts/rerank_recursive/refiner.py +++ /dev/null @@ -1,338 +0,0 @@ -""" -LatentRefiner - Refines latent state based on current ranking results. - -From TRM paper: z encodes "what we've learned about the query so far" -and gets updated based on the current answer (scores). -""" -import os -import time -from typing import Any, Dict, Optional, Tuple - -import numpy as np - - -class LatentRefiner: - """ - Refines the latent state z based on current results. - - Supports: - - Per-collection weight persistence - - Hot-reload from background worker updates - - Online learning via learn_from_teacher() - """ - - WEIGHTS_DIR = os.environ.get("RERANKER_WEIGHTS_DIR", "/tmp/rerank_weights") - WEIGHTS_RELOAD_INTERVAL = float(os.environ.get("RERANKER_WEIGHTS_RELOAD_INTERVAL", "60")) - - def __init__(self, dim: int = 256, hidden_dim: int = 256, lr: float = 0.001): - self.dim = dim - self.hidden_dim = hidden_dim - self.base_lr = lr - self.lr = lr - self._collection = "default" - self._weights_path = self._get_weights_path("default") - self._weights_mtime = 0.0 - self._last_reload_check = 0.0 - self._weights_loaded = False - - self._update_count = 0 - self._version = 0 - - self._momentum_W1: Optional[np.ndarray] = None - self._momentum_b1: Optional[np.ndarray] = None - self._momentum_W2: Optional[np.ndarray] = None - self._momentum_b2: Optional[np.ndarray] = None - - if os.path.exists(self._weights_path): - try: - self._load_weights() - return - except Exception as e: - from scripts.logger import get_logger - get_logger(__name__).warning(f"LatentRefiner: failed to load {self._weights_path}: {e}") - - self._init_random_weights() - - @staticmethod - def _sanitize_collection(collection: str) -> str: - return "".join(c if c.isalnum() or c in "-_" else "_" for c in collection) - - def _get_weights_path(self, collection: str) -> str: - safe_name = self._sanitize_collection(collection) - return os.path.join(self.WEIGHTS_DIR, f"refiner_{safe_name}.npz") - - def set_collection(self, collection: str): - self._collection = collection - self._weights_path = self._get_weights_path(collection) - if os.path.exists(self._weights_path): - try: - self._load_weights() - except Exception: - pass - - def maybe_reload_weights(self): - # Fast path: skip if reload disabled (interval <= 0) - if self.WEIGHTS_RELOAD_INTERVAL <= 0: - return - now = time.time() - if now - self._last_reload_check < self.WEIGHTS_RELOAD_INTERVAL: - return - self._last_reload_check = now - try: - if os.path.exists(self._weights_path): - mtime = os.path.getmtime(self._weights_path) - if mtime > self._weights_mtime: - self._load_weights_safe() - except Exception: - pass - - def _load_weights_safe(self): - import fcntl - lock_path = self._weights_path + ".lock" - try: - os.makedirs(os.path.dirname(lock_path) or ".", exist_ok=True) - with open(lock_path, "w") as lock_file: - fcntl.flock(lock_file.fileno(), fcntl.LOCK_SH) - try: - self._load_weights() - finally: - fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN) - except Exception: - self._load_weights() - - def _init_random_weights(self): - rng = np.random.RandomState(43) - scale = np.float32(np.sqrt(2.0 / (self.dim * 3))) - self.W1 = rng.randn(self.dim * 3, self.hidden_dim).astype(np.float32) * scale - self.b1 = np.zeros(self.hidden_dim, dtype=np.float32) - w2_scale = np.float32(np.sqrt(2.0 / self.hidden_dim)) - self.W2 = rng.randn(self.hidden_dim, self.dim).astype(np.float32) * w2_scale - self.b2 = np.zeros(self.dim, dtype=np.float32) - - self._momentum_W1 = np.zeros_like(self.W1) - self._momentum_b1 = np.zeros_like(self.b1) - self._momentum_W2 = np.zeros_like(self.W2) - self._momentum_b2 = np.zeros_like(self.b2) - - def _load_weights(self) -> bool: - from scripts.logger import get_logger - logger = get_logger(__name__) - try: - data = np.load(self._weights_path, allow_pickle=True) - - def _get(key: str, default): - return data[key] if key in data.files else default - - w1 = _get("W1", None) - w2 = _get("W2", None) - b1 = _get("b1", None) - b2 = _get("b2", None) - - if w1 is None or w2 is None: - data.close() - return False - - expected_w1 = (self.dim * 3, self.hidden_dim) - expected_w2 = (self.hidden_dim, self.dim) - - if w1.shape != expected_w1 or w2.shape != expected_w2: - logger.warning(f"LatentRefiner: shape mismatch") - data.close() - return False - - self.W1 = w1.astype(np.float32, copy=False) - self.b1 = b1.astype(np.float32, copy=False) if b1 is not None else np.zeros(self.hidden_dim, dtype=np.float32) - self.W2 = w2.astype(np.float32, copy=False) - self.b2 = b2.astype(np.float32, copy=False) if b2 is not None else np.zeros(self.dim, dtype=np.float32) - self._update_count = int(_get("update_count", 0)) - self._version = int(_get("version", 0)) - - if self._momentum_W1 is None or self._momentum_W1.shape != self.W1.shape: - self._momentum_W1 = np.zeros_like(self.W1) - self._momentum_b1 = np.zeros_like(self.b1) - self._momentum_W2 = np.zeros_like(self.W2) - self._momentum_b2 = np.zeros_like(self.b2) - - self._weights_loaded = True - self._weights_mtime = os.path.getmtime(self._weights_path) - data.close() - return True - except Exception as e: - logger.warning(f"LatentRefiner: failed to load weights: {e}") - return False - - def refine( - self, - z: np.ndarray, - query_emb: np.ndarray, - doc_embs: np.ndarray, - scores: np.ndarray, - alpha: float = 0.5 - ) -> np.ndarray: - """Refine latent state based on current ranking.""" - self.maybe_reload_weights() - - weights = np.exp(scores - scores.max()) - weights = weights / (weights.sum() + 1e-8) - doc_summary = (weights[:, None] * doc_embs).sum(axis=0) - x = np.concatenate([z, query_emb, doc_summary]) - h = np.maximum(0, x @ self.W1 + self.b1) - z_new = h @ self.W2 + self.b2 - z_refined = alpha * z_new + (1 - alpha) * z - z_refined = z_refined / (np.linalg.norm(z_refined) + 1e-8) - return z_refined - - def refine_with_cache( - self, - z: np.ndarray, - query_emb: np.ndarray, - doc_embs: np.ndarray, - scores: np.ndarray, - alpha: float = 0.5 - ) -> Tuple[np.ndarray, Dict[str, Any]]: - """Refine with cache for backprop.""" - weights = np.exp(scores - scores.max()) - weights = weights / (weights.sum() + 1e-8) - doc_summary = (weights[:, None] * doc_embs).sum(axis=0) - x = np.concatenate([z, query_emb, doc_summary]) - h = np.maximum(0, x @ self.W1 + self.b1) - z_new = h @ self.W2 + self.b2 - z_refined = alpha * z_new + (1 - alpha) * z - z_refined = z_refined / (np.linalg.norm(z_refined) + 1e-8) - cache = {"x": x, "h": h, "z": z, "z_new": z_new, "alpha": alpha, "weights": weights} - return z_refined, cache - - def learn_from_teacher( - self, - z: np.ndarray, - query_emb: np.ndarray, - doc_embs: np.ndarray, - scores: np.ndarray, - teacher_z: np.ndarray, - ) -> float: - """Online learning: update weights so refined z moves toward teacher_z.""" - z_refined, cache = self.refine_with_cache(z, query_emb, doc_embs, scores) - diff = z_refined - teacher_z - loss = float(np.sum(diff ** 2)) - - if loss < 1e-8: - return 0.0 - - dz_refined = 2.0 * diff - dz_new = cache["alpha"] * dz_refined - dW2 = np.outer(cache["h"], dz_new) - db2 = dz_new - dh = dz_new @ self.W2.T - dh = dh * (cache["h"] > 0).astype(np.float32) - dW1 = np.outer(cache["x"], dh) - db1 = dh - - momentum = 0.9 - if self._momentum_W1 is None: - self._momentum_W1 = np.zeros_like(self.W1) - self._momentum_b1 = np.zeros_like(self.b1) - self._momentum_W2 = np.zeros_like(self.W2) - self._momentum_b2 = np.zeros_like(self.b2) - - self._momentum_W1 = momentum * self._momentum_W1 - self.lr * dW1 - self._momentum_b1 = momentum * self._momentum_b1 - self.lr * db1 - self._momentum_W2 = momentum * self._momentum_W2 - self.lr * dW2 - self._momentum_b2 = momentum * self._momentum_b2 - self.lr * db2 - - self.W1 += self._momentum_W1 - self.b1 += self._momentum_b1 - self.W2 += self._momentum_W2 - self.b2 += self._momentum_b2 - self._update_count += 1 - return loss - - def learn_from_teacher_with_cache( - self, - z: np.ndarray, - query_emb: np.ndarray, - doc_embs: np.ndarray, - scores: np.ndarray, - teacher_z: np.ndarray, - ) -> Tuple[float, np.ndarray, np.ndarray, Dict[str, Any]]: - """Online learning with cache for VICReg backprop.""" - z_refined, cache = self.refine_with_cache(z, query_emb, doc_embs, scores) - diff = z_refined - teacher_z - loss = float(np.sum(diff ** 2)) - - if loss >= 1e-8: - dz_refined = 2.0 * diff - dz_new = cache["alpha"] * dz_refined - dW2 = np.outer(cache["h"], dz_new) - db2 = dz_new - dh = dz_new @ self.W2.T - dh = dh * (cache["h"] > 0).astype(np.float32) - dW1 = np.outer(cache["x"], dh) - db1 = dh - - momentum = 0.9 - if self._momentum_W1 is None: - self._momentum_W1 = np.zeros_like(self.W1) - self._momentum_b1 = np.zeros_like(self.b1) - self._momentum_W2 = np.zeros_like(self.W2) - self._momentum_b2 = np.zeros_like(self.b2) - - self._momentum_W1 = momentum * self._momentum_W1 - self.lr * dW1 - self._momentum_b1 = momentum * self._momentum_b1 - self.lr * db1 - self._momentum_W2 = momentum * self._momentum_W2 - self.lr * dW2 - self._momentum_b2 = momentum * self._momentum_b2 - self.lr * db2 - - self.W1 += self._momentum_W1 - self.b1 += self._momentum_b1 - self.W2 += self._momentum_W2 - self.b2 += self._momentum_b2 - self._update_count += 1 - - return loss, z, z_refined, cache - - def apply_vicreg_gradient(self, grad_z_refined: np.ndarray, cache: Dict[str, Any], weight: float = 0.1): - """Apply VICReg gradient to refiner weights.""" - dz_new = cache["alpha"] * grad_z_refined * weight - dW2 = np.outer(cache["h"], dz_new) - db2 = dz_new - dh = dz_new @ self.W2.T - dh = dh * (cache["h"] > 0).astype(np.float32) - dW1 = np.outer(cache["x"], dh) - db1 = dh - - self.W1 -= self.lr * dW1 - self.b1 -= self.lr * db1 - self.W2 -= self.lr * dW2 - self.b2 -= self.lr * db2 - - def _save_weights(self, checkpoint: bool = False): - """Save weights to disk atomically.""" - import fcntl - os.makedirs(self.WEIGHTS_DIR, exist_ok=True) - self._version += 1 - - tmp_base = self._weights_path.replace(".npz", ".tmp") - tmp_path = tmp_base + ".npz" - lock_path = self._weights_path + ".lock" - - try: - with open(lock_path, "w") as lock_file: - fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX) - try: - np.savez( - tmp_base, - W1=self.W1, b1=self.b1, W2=self.W2, b2=self.b2, - update_count=self._update_count, - version=self._version, - dim=self.dim, - hidden_dim=self.hidden_dim, - ) - os.replace(tmp_path, self._weights_path) - finally: - fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN) - except Exception: - if os.path.exists(tmp_path): - try: - os.remove(tmp_path) - except Exception: - pass - raise diff --git a/scripts/rerank_recursive/scorer.py b/scripts/rerank_recursive/scorer.py deleted file mode 100644 index d8468a73..00000000 --- a/scripts/rerank_recursive/scorer.py +++ /dev/null @@ -1,381 +0,0 @@ -""" -TinyScorer - 2-layer MLP for scoring query-document pairs. - -Inspired by TRM: minimal parameters, maximum iterations. -""" -import os -import time -from typing import Any, Dict, List, Tuple - -import numpy as np - - -class TinyScorer: - """ - Tiny 2-layer MLP for scoring query-document pairs. - - Inspired by TRM: minimal parameters, maximum iterations. - Production-ready with: - - Collection-aware weights with atomic loading - - Checkpoint versioning (keep last N versions) - - Training metrics (loss, sample count, convergence) - - Learning rate decay - - Hot reload from background worker updates - """ - - # Class-level configuration - WEIGHTS_DIR = os.environ.get("RERANKER_WEIGHTS_DIR", "/tmp/rerank_weights") - WEIGHTS_RELOAD_INTERVAL = float(os.environ.get("RERANKER_WEIGHTS_RELOAD_INTERVAL", "60")) - MAX_CHECKPOINTS = int(os.environ.get("RERANKER_MAX_CHECKPOINTS", "5")) - LR_DECAY_STEPS = int(os.environ.get("RERANKER_LR_DECAY_STEPS", "1000")) - LR_DECAY_RATE = float(os.environ.get("RERANKER_LR_DECAY_RATE", "0.95")) - MIN_LR = float(os.environ.get("RERANKER_MIN_LR", "0.0001")) - - def __init__(self, dim: int = 256, hidden_dim: int = 512, lr: float = 0.001): - self.dim = dim - self.hidden_dim = hidden_dim - self.base_lr = lr - self.lr = lr - self._collection = "default" - self._weights_path = self._get_weights_path("default") - self._weights_mtime = 0.0 - self._last_reload_check = 0.0 - - # Training metrics - self._update_count = 0 - self._total_samples = 0 - self._cumulative_loss = 0.0 - self._recent_losses: List[float] = [] # Rolling window for convergence detection - self._version = 0 - - # Try to load saved weights, otherwise init random - if os.path.exists(self._weights_path): - try: - self._load_weights() - return - except Exception as e: - from scripts.logger import get_logger - get_logger(__name__).warning(f"TinyScorer: failed to load {self._weights_path}: {e}, using random init") - - self._init_random_weights() - - def _init_random_weights(self): - """Initialize weights randomly using He initialization (local RNG, deterministic).""" - rng = np.random.RandomState(42) - scale = np.float32(np.sqrt(2.0 / (self.dim * 3))) - self.W1 = rng.randn(self.dim * 3, self.hidden_dim).astype(np.float32) * scale - self.b1 = np.zeros(self.hidden_dim, dtype=np.float32) - w2_scale = np.float32(np.sqrt(2.0 / self.hidden_dim)) - self.W2 = rng.randn(self.hidden_dim, 1).astype(np.float32) * w2_scale - self.b2 = np.zeros(1, dtype=np.float32) - - # Momentum for SGD - self._momentum_W1 = np.zeros_like(self.W1) - self._momentum_b1 = np.zeros_like(self.b1) - self._momentum_W2 = np.zeros_like(self.W2) - self._momentum_b2 = np.zeros_like(self.b2) - - def _update_learning_rate(self): - """Decay learning rate based on update count.""" - if self._update_count > 0 and self._update_count % self.LR_DECAY_STEPS == 0: - self.lr = max(self.MIN_LR, self.lr * self.LR_DECAY_RATE) - - def get_metrics(self) -> Dict[str, Any]: - """Get current training metrics.""" - avg_loss = self._cumulative_loss / max(1, self._update_count) - recent_avg = np.mean(self._recent_losses) if self._recent_losses else 0.0 - return { - "collection": self._collection, - "version": self._version, - "update_count": self._update_count, - "total_samples": self._total_samples, - "cumulative_loss": self._cumulative_loss, - "avg_loss": avg_loss, - "recent_avg_loss": float(recent_avg), - "learning_rate": self.lr, - "converged": self._is_converged(), - } - - def _is_converged(self, window: int = 100, threshold: float = 0.01) -> bool: - """Check if training has converged (loss not improving).""" - if len(self._recent_losses) < window: - return False - recent = self._recent_losses[-window:] - first_half = np.mean(recent[:window // 2]) - second_half = np.mean(recent[window // 2:]) - return abs(first_half - second_half) < threshold * first_half - - def _get_weights_path(self, collection: str) -> str: - """Get weights file path for a collection.""" - os.makedirs(self.WEIGHTS_DIR, exist_ok=True) - safe_name = "".join(c if c.isalnum() or c in "-_" else "_" for c in collection) - return os.path.join(self.WEIGHTS_DIR, f"weights_{safe_name}.npz") - - def set_collection(self, collection: str): - """Set collection and load corresponding weights.""" - self._collection = collection - self._weights_path = self._get_weights_path(collection) - if os.path.exists(self._weights_path): - try: - self._load_weights() - except Exception: - pass - - def maybe_reload_weights(self): - """Check if weights file changed and reload if needed (hot reload).""" - # Fast path: skip if reload disabled (interval <= 0) - if self.WEIGHTS_RELOAD_INTERVAL <= 0: - return - now = time.time() - if now - self._last_reload_check < self.WEIGHTS_RELOAD_INTERVAL: - return - self._last_reload_check = now - - try: - if os.path.exists(self._weights_path): - mtime = os.path.getmtime(self._weights_path) - if mtime > self._weights_mtime: - self._load_weights_safe() - except Exception: - pass - - def _load_weights_safe(self): - """Load weights with advisory file locking.""" - import fcntl - lock_path = self._weights_path + ".lock" - try: - os.makedirs(os.path.dirname(lock_path) or ".", exist_ok=True) - with open(lock_path, "w") as lock_file: - fcntl.flock(lock_file.fileno(), fcntl.LOCK_SH) - try: - self._load_weights() - finally: - fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN) - except Exception: - self._load_weights() - - def forward(self, query_emb: np.ndarray, doc_emb: np.ndarray, z: np.ndarray) -> np.ndarray: - """Score documents given query and latent state.""" - self.maybe_reload_weights() - - n_docs = doc_emb.shape[0] - q_broadcast = np.tile(query_emb, (n_docs, 1)) - z_broadcast = np.tile(z, (n_docs, 1)) - x = np.concatenate([q_broadcast, doc_emb, z_broadcast], axis=1) - h = np.maximum(0, x @ self.W1 + self.b1) - scores = (h @ self.W2 + self.b2).squeeze(-1) - return scores - - def forward_with_cache(self, x: np.ndarray) -> Tuple[np.ndarray, Dict[str, np.ndarray]]: - """Forward pass with cached activations for backprop.""" - z1 = x @ self.W1 + self.b1 - h1 = np.maximum(0, z1) - z2 = h1 @ self.W2 + self.b2 - scores = z2.squeeze(-1) - cache = {"x": x, "z1": z1, "h1": h1} - return scores, cache - - def backward(self, dscores: np.ndarray, cache: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]: - """Backward pass to compute gradients.""" - batch_size = dscores.shape[0] - dz2 = dscores.reshape(-1, 1) - dW2 = cache["h1"].T @ dz2 - db2 = dz2.sum(axis=0) - dh1 = dz2 @ self.W2.T - dz1 = dh1 * (cache["z1"] > 0).astype(np.float32) - dW1 = cache["x"].T @ dz1 - db1 = dz1.sum(axis=0) - return {"W1": dW1 / batch_size, "b1": db1 / batch_size, "W2": dW2 / batch_size, "b2": db2 / batch_size} - - def learn_from_teacher( - self, - query_emb: np.ndarray, - doc_embs: np.ndarray, - z: np.ndarray, - teacher_scores: np.ndarray, - margin: float = 0.5, - ) -> float: - """Online learning: update weights to match ONNX teacher ranking.""" - n_docs = doc_embs.shape[0] - if n_docs < 2: - return 0.0 - - q_broadcast = np.tile(query_emb, (n_docs, 1)) - z_broadcast = np.tile(z, (n_docs, 1)) - x = np.concatenate([q_broadcast, doc_embs, z_broadcast], axis=1) - our_scores, cache = self.forward_with_cache(x) - teacher_order = np.argsort(-teacher_scores) - - n_pairs = min(5, n_docs // 2) - total_loss = 0.0 - dscores = np.zeros(n_docs, dtype=np.float32) - - for i in range(n_pairs): - pos_idx = teacher_order[i] - neg_idx = teacher_order[-(i + 1)] - diff = our_scores[pos_idx] - our_scores[neg_idx] - if diff < margin: - loss = margin - diff - total_loss += loss - dscores[pos_idx] -= 1.0 - dscores[neg_idx] += 1.0 - - self._total_samples += n_docs - self._cumulative_loss += total_loss - self._recent_losses.append(total_loss) - if len(self._recent_losses) > 200: - self._recent_losses = self._recent_losses[-200:] - - if total_loss > 0: - grads = self.backward(dscores, cache) - momentum = 0.9 - self._momentum_W1 = momentum * self._momentum_W1 - self.lr * grads["W1"] - self._momentum_b1 = momentum * self._momentum_b1 - self.lr * grads["b1"] - self._momentum_W2 = momentum * self._momentum_W2 - self.lr * grads["W2"] - self._momentum_b2 = momentum * self._momentum_b2 - self.lr * grads["b2"] - self.W1 += self._momentum_W1 - self.b1 += self._momentum_b1 - self.W2 += self._momentum_W2 - self.b2 += self._momentum_b2 - self._update_count += 1 - self._update_learning_rate() - - return total_loss - - def _save_weights(self, checkpoint: bool = False): - """Save weights to disk atomically.""" - import fcntl - try: - self._version += 1 - tmp_base = self._weights_path.replace(".npz", ".tmp") - # Keep only last 200 losses for convergence (avoid unbounded growth) - recent_losses_to_save = self._recent_losses[-200:] if self._recent_losses else [] - np.savez( - tmp_base, - W1=self.W1, b1=self.b1, W2=self.W2, b2=self.b2, - momentum_W1=self._momentum_W1, momentum_b1=self._momentum_b1, - momentum_W2=self._momentum_W2, momentum_b2=self._momentum_b2, - update_count=self._update_count, - total_samples=self._total_samples, - cumulative_loss=self._cumulative_loss, - learning_rate=self.lr, - version=self._version, - collection=self._collection, - recent_losses=np.array(recent_losses_to_save, dtype=np.float32), - ) - tmp_path = tmp_base + ".npz" - lock_path = self._weights_path + ".lock" - os.makedirs(os.path.dirname(lock_path) or ".", exist_ok=True) - with open(lock_path, "w") as lock_file: - fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX) - try: - os.replace(tmp_path, self._weights_path) - finally: - fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN) - if checkpoint or self._version % 100 == 0: - self._save_checkpoint() - except Exception: - pass - - def _save_checkpoint(self): - """Save a versioned checkpoint and prune old ones.""" - try: - import shutil - checkpoint_path = self._weights_path.replace(".npz", f"_v{self._version}.npz") - shutil.copy2(self._weights_path, checkpoint_path) - self._prune_old_checkpoints() - except Exception: - pass - - def _prune_old_checkpoints(self): - """Remove old checkpoints keeping only the most recent MAX_CHECKPOINTS.""" - try: - import glob - pattern = self._weights_path.replace(".npz", "_v*.npz") - checkpoints = sorted(glob.glob(pattern)) - if len(checkpoints) > self.MAX_CHECKPOINTS: - for old_cp in checkpoints[:-self.MAX_CHECKPOINTS]: - try: - os.remove(old_cp) - except Exception: - pass - except Exception: - pass - - def _load_weights(self): - """Load weights from disk with dimension validation.""" - from scripts.logger import get_logger - logger = get_logger(__name__) - - data = np.load(self._weights_path, allow_pickle=True) - - def _get(key: str, default): - return data[key] if key in data.files else default - - w1_loaded = data["W1"] - w2_loaded = data["W2"] - b1_loaded = data["b1"] - b2_loaded = data["b2"] - - expected_w1 = (self.dim * 3, self.hidden_dim) - expected_w2 = (self.hidden_dim, 1) - expected_b1 = (self.hidden_dim,) - expected_b2 = (1,) - - shape_ok = ( - w1_loaded.shape == expected_w1 and - w2_loaded.shape == expected_w2 and - b1_loaded.shape == expected_b1 and - b2_loaded.shape == expected_b2 - ) - - if not shape_ok: - logger.warning(f"TinyScorer: shape mismatch, falling back to random init.") - data.close() - self._init_random_weights() - return - - self.W1 = w1_loaded.astype(np.float32, copy=False) - self.b1 = b1_loaded.astype(np.float32, copy=False) - self.W2 = w2_loaded.astype(np.float32, copy=False) - self.b2 = b2_loaded.astype(np.float32, copy=False) - self._update_count = int(_get("update_count", 0)) - self._total_samples = int(_get("total_samples", 0)) - self._cumulative_loss = float(_get("cumulative_loss", 0.0)) - self._version = int(_get("version", 0)) - - if "learning_rate" in data.files: - self.lr = float(data["learning_rate"]) - - if "momentum_W1" in data.files: - self._momentum_W1 = data["momentum_W1"].astype(np.float32, copy=False) - self._momentum_b1 = data["momentum_b1"].astype(np.float32, copy=False) - self._momentum_W2 = data["momentum_W2"].astype(np.float32, copy=False) - self._momentum_b2 = data["momentum_b2"].astype(np.float32, copy=False) - else: - self._momentum_W1 = np.zeros_like(self.W1) - self._momentum_b1 = np.zeros_like(self.b1) - self._momentum_W2 = np.zeros_like(self.W2) - self._momentum_b2 = np.zeros_like(self.b2) - - # Restore recent losses for convergence detection (survives restarts) - if "recent_losses" in data.files: - self._recent_losses = list(data["recent_losses"].astype(np.float32)) - else: - self._recent_losses = [] - - self._weights_mtime = os.path.getmtime(self._weights_path) - data.close() - - def rollback_to_checkpoint(self, version: int) -> bool: - """Rollback to a specific checkpoint version.""" - try: - import shutil - checkpoint_path = self._weights_path.replace(".npz", f"_v{version}.npz") - if os.path.exists(checkpoint_path): - shutil.copy2(checkpoint_path, self._weights_path) - self._load_weights() - return True - except Exception: - pass - return False diff --git a/scripts/rerank_recursive/state.py b/scripts/rerank_recursive/state.py deleted file mode 100644 index 5f4a9520..00000000 --- a/scripts/rerank_recursive/state.py +++ /dev/null @@ -1,18 +0,0 @@ -""" -RefinementState dataclass - carries latent state between refinement iterations. -""" -from dataclasses import dataclass, field -from typing import List -import numpy as np - - -@dataclass -class RefinementState: - """Carries latent state between refinement iterations.""" - z: np.ndarray # Latent representation (query understanding) - scores: np.ndarray # Current score estimates - iteration: int = 0 - confidence: float = 0.0 # For early stopping - - # Track per-iteration improvements for analysis - score_history: List[np.ndarray] = field(default_factory=list) diff --git a/scripts/rerank_recursive/utils.py b/scripts/rerank_recursive/utils.py deleted file mode 100644 index 956bd0c0..00000000 --- a/scripts/rerank_recursive/utils.py +++ /dev/null @@ -1,324 +0,0 @@ -""" -Utility functions for the recursive reranker. - -Contains: -- Filename boost computation -- Identifier tokenization -- Embedding cache management -""" -import re -import hashlib -import threading -from typing import Any, Dict, List, Optional, Set - -import numpy as np - - -# Very common tokens that appear everywhere - reduce their weight -_COMMON_TOKENS = frozenset({ - "index", "main", "app", "utils", "util", "helper", "helpers", "common", - "base", "core", "lib", "src", "test", "tests", "spec", "specs", - "internal", "public", "private", "static", "default", "new", "old", - "data", "type", "types", "model", "models", "view", "views", - "the", "and", "for", "with", "from", "that", "this", "have", "are", -}) - - -def _split_identifier(s: str) -> List[str]: - """Split any identifier into tokens, handling all common conventions. - - Handles: snake_case, kebab-case, camelCase, PascalCase, SCREAMING_CASE, - numbers, acronyms (XMLParser -> xml, parser), dot.notation, and mixed styles. - - Special handling: - - Preserves meaningful acronyms (API, HTTP, JSON, XML, URL, etc.) - - Strips common prefixes (I for interface, _ for private) - - Handles version suffixes (v2, 2.0) - """ - if not s: - return [] - - # Strip common prefixes that don't add meaning - if len(s) > 1: - # Interface prefix (IUserService -> UserService) - if s[0] == 'I' and s[1].isupper(): - s = s[1:] - # Private prefix (_private -> private) - elif s[0] == '_': - s = s.lstrip('_') - # Dollar prefix ($scope -> scope) - elif s[0] == '$': - s = s[1:] - - # Insert space before uppercase letters that follow lowercase (camelCase) - s = re.sub(r"([a-z])([A-Z])", r"\1 \2", s) - # Insert space before uppercase letters followed by lowercase (acronyms: XMLParser -> XML Parser) - s = re.sub(r"([A-Z]+)([A-Z][a-z])", r"\1 \2", s) - # Insert space around digit sequences (handler2 -> handler 2, v2 -> v 2) - s = re.sub(r"([a-zA-Z])(\d)", r"\1 \2", s) - s = re.sub(r"(\d)([a-zA-Z])", r"\1 \2", s) - - # Split on separators: underscore, hyphen, dot, space - parts = re.split(r"[_\-.\s]+", s) - tokens = [] - for part in parts: - part = part.strip().lower() - # Skip pure numbers and single chars (except meaningful ones) - if not part: - continue - if part.isdigit(): - continue # Skip version numbers like "2", "18" - if len(part) < 2: - continue - tokens.append(part) - - return tokens - - -def _normalize_token(tok: str) -> Set[str]: - """Return the token plus simple morphological variants.""" - forms = {tok} - # Simple plural/singular normalization - if tok.endswith('s') and len(tok) > 3: - forms.add(tok[:-1]) # services -> service - elif tok.endswith('es') and len(tok) > 4: - forms.add(tok[:-2]) # processes -> process - elif tok.endswith('ies') and len(tok) > 4: - forms.add(tok[:-3] + 'y') # utilities -> utility - # Add singular -> plural - if not tok.endswith('s') and len(tok) > 2: - forms.add(tok + 's') - return forms - - -def _tokenize_for_fname_boost(text: Any) -> Set[str]: - """Robust tokenization for filename boosts. - - Some MCP/IDE clients pass query strings that include quotes/brackets - or list-like wrappers. Regex tokenization is resilient to that. - """ - if not text: - return set() - try: - s = str(text) - except Exception: - return set() - - # Split on any non-alphanumeric - raw_parts = re.split(r"[^a-zA-Z0-9]+", s) - tokens = set() - for part in raw_parts: - for tok in _split_identifier(part): - if len(tok) >= 3: # Query tokens need 3+ chars - tokens.add(tok) - return tokens - - -def _candidate_path_for_fname_boost(candidate: Dict[str, Any]) -> str: - """Best-effort extraction of a path/filename from candidate objects.""" - for key in ("path", "rel_path", "host_path", "container_path", "client_path"): - try: - val = candidate.get(key) - except Exception: - val = None - if isinstance(val, str) and val.strip(): - return val - - try: - md = candidate.get("metadata") or {} - if isinstance(md, dict): - for key in ("path", "rel_path", "host_path", "container_path", "client_path"): - val = md.get(key) - if isinstance(val, str) and val.strip(): - return val - except Exception: - pass - - return "" - - -def _compute_fname_boost(query: Any, candidate: Dict[str, Any], factor: float) -> float: - """Compute filename/query correlation boost for a candidate. - - Production-grade matching for real-world codebases at scale: - - **Naming convention support:** - - snake_case, camelCase, PascalCase, kebab-case, SCREAMING_CASE - - Dot notation (com.company.auth.service) - - Mixed styles (legacy codebases) - - **Smart tokenization:** - - Acronyms: XMLParser -> xml, parser; HTTPClient -> http, client - - Prefixes stripped: IService -> service, _private -> private - - Numbers separated: handler2 -> handler, React18 -> react - - **Normalization:** - - Simple plural/singular normalization (services <-> service) - - **Position-aware scoring:** - - Filename matches weighted higher than directory matches - - Deeper directories weighted less (noise reduction) - - **Specificity weighting:** - - Common tokens (index, main, utils) weighted less - - Rare/specific tokens weighted more - - **Scoring tiers:** - - Exact match: 1.0 × factor - - Normalized match (morphology): 0.8 × factor - - Substring containment: 0.4 × factor - - Common token penalty: 0.5× multiplier - - Filename bonus: 1.5× multiplier for filename matches - - Requires 2+ quality matches to trigger (prevents noise). - """ - if not factor or factor <= 0: - return 0.0 - - query_tokens = _tokenize_for_fname_boost(query) - if not query_tokens: - return 0.0 - - path = _candidate_path_for_fname_boost(candidate) - path = str(path or "") - if not path: - return 0.0 - - # Strip common prefixes that add noise (preserve case for splitting) - path_clean = path - path_lower = path.lower() - for prefix in ("/work/", "/app/", "/src/", "/home/", "/var/", "/opt/", "/usr/"): - if path_lower.startswith(prefix): - path_clean = path[len(prefix):] - break - - # Split path into segments, track position for weighting - path_segments = re.split(r"[/\\]", path_clean) - path_segments = [s for s in path_segments if s] # Remove empty - - if not path_segments: - return 0.0 - - # Tokenize with position info: (token, is_filename, depth) - # Filename = last segment, depth = 0 for filename, 1 for parent, etc. - path_token_info: Dict[str, Dict[str, Any]] = {} # token -> {is_filename, min_depth} - - for i, segment in enumerate(reversed(path_segments)): - is_filename = (i == 0) - depth = i - - # Strip extension from filename - if is_filename and "." in segment: - segment = segment.rsplit(".", 1)[0] - - for tok in _split_identifier(segment): - if len(tok) >= 2: - if tok not in path_token_info: - path_token_info[tok] = {"is_filename": is_filename, "depth": depth} - # Keep the most important occurrence (filename > dir, shallow > deep) - elif is_filename and not path_token_info[tok]["is_filename"]: - path_token_info[tok] = {"is_filename": True, "depth": depth} - - if not path_token_info: - return 0.0 - - path_tokens = set(path_token_info.keys()) - - # Build normalized lookup for path tokens - path_normalized: Dict[str, str] = {} # normalized_form -> original_token - for ptok in path_tokens: - for form in _normalize_token(ptok): - if form not in path_normalized: - path_normalized[form] = ptok - - # Score matches with quality tiers - score = 0.0 - matched_query_tokens = set() - - for qtok in query_tokens: - qtok_forms = _normalize_token(qtok) - match_score = 0.0 - matched_ptok = None - - # Tier 1: Exact match - if qtok in path_tokens: - match_score = 1.0 - matched_ptok = qtok - else: - # Tier 2: Normalized match (plural/singular) - for qform in qtok_forms: - if qform in path_normalized: - match_score = 0.8 - matched_ptok = path_normalized[qform] - break - - # Tier 3: Substring containment (if no normalized match) - if match_score == 0.0: - for ptok in path_tokens: - if len(qtok) >= 4 and len(ptok) >= 4: - if qtok in ptok or ptok in qtok: - overlap = min(len(qtok), len(ptok)) - if overlap >= 4: - match_score = 0.4 - matched_ptok = ptok - break - - if match_score > 0 and matched_ptok: - matched_query_tokens.add(qtok) - - # Apply position bonus (filename matches worth more) - info = path_token_info.get(matched_ptok, {}) - if info.get("is_filename"): - match_score *= 1.5 # 50% bonus for filename match - else: - # Depth penalty for deep directories - depth = info.get("depth", 0) - if depth > 2: - match_score *= 0.8 # Slight penalty for deep paths - - # Common token penalty - if qtok in _COMMON_TOKENS or matched_ptok in _COMMON_TOKENS: - match_score *= 0.5 - - score += match_score - - # Require 2+ quality matches to trigger (prevents noise from single common word) - if len(matched_query_tokens) < 2: - return 0.0 - - return float(score * factor) - - -# --------------------------------------------------------------------------- -# Embedding Cache -# --------------------------------------------------------------------------- - -# Global embedding cache for efficiency -# Key is sha256 hex digest (deterministic, collision-resistant) -_EMBEDDING_CACHE: Dict[str, np.ndarray] = {} -_EMBEDDING_CACHE_MAX_SIZE = 10000 -_EMBEDDING_CACHE_LOCK = threading.Lock() - - -def _cache_key(text: str) -> str: - """Generate deterministic cache key from text (process-stable, collision-resistant).""" - return hashlib.sha256(text.encode("utf-8", errors="replace")).hexdigest() - - -def _get_cached_embedding(text: str) -> Optional[np.ndarray]: - """Get embedding from cache if exists.""" - key = _cache_key(text) - with _EMBEDDING_CACHE_LOCK: - return _EMBEDDING_CACHE.get(key) - - -def _cache_embedding(text: str, embedding: np.ndarray): - """Cache embedding for text.""" - key = _cache_key(text) - with _EMBEDDING_CACHE_LOCK: - if len(_EMBEDDING_CACHE) >= _EMBEDDING_CACHE_MAX_SIZE: - # Evict oldest 10% - keys_to_remove = list(_EMBEDDING_CACHE.keys())[:_EMBEDDING_CACHE_MAX_SIZE // 10] - for k in keys_to_remove: - del _EMBEDDING_CACHE[k] - _EMBEDDING_CACHE[key] = embedding diff --git a/scripts/rerank_recursive/vicreg.py b/scripts/rerank_recursive/vicreg.py deleted file mode 100644 index 399e5322..00000000 --- a/scripts/rerank_recursive/vicreg.py +++ /dev/null @@ -1,76 +0,0 @@ -""" -VICReg - Variance-Invariance-Covariance Regularization for Latent Refinement. - -Adapted from VICReg (Bardes et al., 2021) for online 3-pass reranking. -""" -from typing import Dict, Tuple - -import numpy as np - - -class VICReg: - """ - VICReg regularization for refinement residuals. - - Regularizes the refiner's residual (z_refined - z) to have: - - Unit variance per dimension (prevents collapse) - - Decorrelated dimensions (prevents redundancy) - - Bounded magnitude (stable updates) - """ - - def __init__( - self, - lambda_var: float = 1.0, - lambda_cov: float = 0.04, - lambda_inv: float = 0.1, - var_target: float = 1.0, - ): - self.lambda_var = lambda_var - self.lambda_cov = lambda_cov - self.lambda_inv = lambda_inv - self.var_target = var_target - - def forward( - self, z_batch: np.ndarray, z_refined_batch: np.ndarray - ) -> Tuple[float, np.ndarray, Dict[str, float]]: - """ - Compute VICReg loss and gradient w.r.t. z_refined. - - Args: - z_batch: (N, dim) original latent states - z_refined_batch: (N, dim) refined latent states - - Returns: - (total_loss, grad_z_refined, loss_components) - """ - N, dim = z_batch.shape - eps = 1e-8 - - residual = z_refined_batch - z_batch - mean_res = residual.mean(axis=0, keepdims=True) - residual_centered = residual - mean_res - - # Variance loss - std = residual.std(axis=0) + eps - var_diff = self.var_target - std - var_loss = float(np.maximum(0, var_diff).mean()) - hinge_mask = (var_diff > 0).astype(np.float32) - d_var = -hinge_mask[None, :] * residual_centered / (N * std[None, :] * dim) - - # Covariance loss - cov = (residual_centered.T @ residual_centered) / (N - 1 + eps) - off_diag_mask = 1.0 - np.eye(dim, dtype=np.float32) - off_diag = cov * off_diag_mask - cov_loss = float((off_diag ** 2).sum() / dim) - d_cov = 4 * residual_centered @ (off_diag * off_diag_mask) / ((N - 1 + eps) * dim) - - # Invariance loss - inv_loss = float((residual ** 2).mean()) - d_inv = 2 * residual / (N * dim) - - # Total - total_loss = self.lambda_var * var_loss + self.lambda_cov * cov_loss + self.lambda_inv * inv_loss - grad = (self.lambda_var * d_var + self.lambda_cov * d_cov + self.lambda_inv * d_inv).astype(np.float32) - - components = {"var_loss": var_loss, "cov_loss": cov_loss, "inv_loss": inv_loss} - return total_loss, grad, components diff --git a/scripts/rerank_tools/__init__.py b/scripts/rerank_tools/__init__.py index 8c81a54f..efd7b451 100644 --- a/scripts/rerank_tools/__init__.py +++ b/scripts/rerank_tools/__init__.py @@ -1,78 +1 @@ -""" -rerank_tools - Reranker training, evaluation, and benchmarking utilities. - -This package contains tools for training, evaluating, and benchmarking the -recursive reranker system. The core reranker logic is in scripts/rerank_recursive/. - -Modules: -- eval: Offline evaluation pipeline (MRR/Recall/latency) -- events: Training event logging for background processing -- train: Training infrastructure for TinyScorer -- query: Query reranking utilities -- local: Local reranking CLI -- ab_test: A/B testing framework -- benchmark: Real-world benchmark harness -""" -from __future__ import annotations - -# Core evaluation -from .eval import ( - EvalResult, - EvalSummary, - get_candidates, - compute_mrr, - compute_recall_at_k, - run_eval, - print_summary, -) - -# Event logging -from .events import ( - log_training_event, - list_event_files, - read_events, - cleanup_old_events, - RERANK_EVENTS_DIR, - RERANK_EVENTS_ENABLED, - RERANK_EVENTS_SAMPLE_RATE, - RERANK_EVENTS_RETENTION_DAYS, -) - -# Training -from .train import ( - TrainingExample, - TrainingConfig, - TrainableTinyScorer, - TrainableLatentRefiner, - RecursiveRerankerTrainer, - margin_ranking_loss, - deep_supervision_loss, -) - -__all__ = [ - # Eval - "EvalResult", - "EvalSummary", - "get_candidates", - "compute_mrr", - "compute_recall_at_k", - "run_eval", - "print_summary", - # Events - "log_training_event", - "list_event_files", - "read_events", - "cleanup_old_events", - "RERANK_EVENTS_DIR", - "RERANK_EVENTS_ENABLED", - "RERANK_EVENTS_SAMPLE_RATE", - "RERANK_EVENTS_RETENTION_DAYS", - # Training - "TrainingExample", - "TrainingConfig", - "TrainableTinyScorer", - "TrainableLatentRefiner", - "RecursiveRerankerTrainer", - "margin_ranking_loss", - "deep_supervision_loss", -] +"""Shared package for the local reranker and relevance feedback pipeline.""" diff --git a/scripts/rerank_tools/ab_test.py b/scripts/rerank_tools/ab_test.py deleted file mode 100644 index b4c328e2..00000000 --- a/scripts/rerank_tools/ab_test.py +++ /dev/null @@ -1,360 +0,0 @@ -#!/usr/bin/env python3 -""" -A/B Testing Infrastructure for Rerankers. - -Provides: -1. Random assignment to reranker variants -2. Metric logging per variant -3. Aggregation and statistical comparison -4. Session-aware bucketing (same user gets same variant) - -Usage: - # In your search pipeline: - from scripts.rerank_ab_test import ABTestManager, RerankerVariant - - ab = ABTestManager() - variant = ab.get_variant(session_id="user_123") - results = variant.rerank(query, candidates) - ab.log_metrics(session_id, latency_ms=5.2, clicked_rank=1) - - # Analyze results: - ab.print_summary() -""" - -import os -import json -import time -import hashlib -import threading -from pathlib import Path -from typing import List, Dict, Any, Optional, Callable -from dataclasses import dataclass, field -from datetime import datetime -from enum import Enum - - -class VariantType(Enum): - """Available reranker variants.""" - BASELINE = "baseline" # No reranking, use initial scores - ONNX = "onnx" # ONNX cross-encoder - RECURSIVE = "recursive" # TRM-inspired recursive reranker - RECURSIVE_ONNX = "recursive_onnx" # ONNX + recursive refinement - - -@dataclass -class ABMetric: - """Single metric observation.""" - session_id: str - variant: str - timestamp: float - latency_ms: float = 0.0 - clicked_rank: Optional[int] = None # 1-indexed rank of clicked result - num_results: int = 0 - query_length: int = 0 - custom: Dict[str, Any] = field(default_factory=dict) - - -@dataclass -class VariantStats: - """Aggregated statistics for a variant.""" - variant: str - n_observations: int = 0 - latency_sum: float = 0.0 - latency_sq_sum: float = 0.0 # For std calculation - mrr_sum: float = 0.0 # Mean Reciprocal Rank - clicks_at_1: int = 0 - clicks_at_3: int = 0 - clicks_at_5: int = 0 - total_clicks: int = 0 - - def add_observation(self, metric: ABMetric): - """Add a single observation.""" - self.n_observations += 1 - self.latency_sum += metric.latency_ms - self.latency_sq_sum += metric.latency_ms ** 2 - - if metric.clicked_rank is not None: - self.total_clicks += 1 - self.mrr_sum += 1.0 / metric.clicked_rank - if metric.clicked_rank == 1: - self.clicks_at_1 += 1 - if metric.clicked_rank <= 3: - self.clicks_at_3 += 1 - if metric.clicked_rank <= 5: - self.clicks_at_5 += 1 - - def get_summary(self) -> Dict[str, Any]: - """Get summary statistics.""" - if self.n_observations == 0: - return {"variant": self.variant, "n_observations": 0} - - n = self.n_observations - mean_latency = self.latency_sum / n - var_latency = (self.latency_sq_sum / n) - (mean_latency ** 2) - std_latency = var_latency ** 0.5 if var_latency > 0 else 0 - - mrr = self.mrr_sum / self.total_clicks if self.total_clicks > 0 else 0 - - return { - "variant": self.variant, - "n_observations": n, - "latency_ms_mean": round(mean_latency, 2), - "latency_ms_std": round(std_latency, 2), - "mrr": round(mrr, 4), - "click_rate_at_1": round(self.clicks_at_1 / n, 4) if n > 0 else 0, - "click_rate_at_3": round(self.clicks_at_3 / n, 4) if n > 0 else 0, - "click_rate_at_5": round(self.clicks_at_5 / n, 4) if n > 0 else 0, - "total_clicks": self.total_clicks, - } - - -class RerankerVariant: - """Wrapper for a reranker variant with timing.""" - - def __init__(self, variant_type: VariantType, rerank_fn: Callable): - self.variant_type = variant_type - self.rerank_fn = rerank_fn - - def rerank(self, query: str, candidates: List[Dict[str, Any]], **kwargs) -> tuple: - """Rerank candidates and return (results, latency_ms).""" - start = time.perf_counter() - results = self.rerank_fn(query, candidates, **kwargs) - latency_ms = (time.perf_counter() - start) * 1000 - return results, latency_ms - - -class ABTestManager: - """ - Manages A/B testing for reranker variants. - - Features: - - Consistent hashing for session assignment (same session = same variant) - - Thread-safe metric logging - - Persistent storage of metrics - - Statistical comparison utilities - """ - - def __init__( - self, - variants: Optional[List[VariantType]] = None, - weights: Optional[List[float]] = None, - log_path: str = "data/ab_test_metrics.jsonl", - experiment_id: Optional[str] = None, - ): - self.variants = variants or [VariantType.BASELINE, VariantType.RECURSIVE] - self.weights = weights or [1.0] * len(self.variants) - self.log_path = Path(log_path) - self.experiment_id = experiment_id or datetime.now().strftime("%Y%m%d_%H%M%S") - - # Normalize weights (guard against zero/negative total) - total_weight = sum(self.weights) - if total_weight <= 0: - total_weight = len(self.weights) # Fallback to uniform - self.weights = [1.0] * len(self.variants) - self.cumulative_weights = [] - cumsum = 0.0 - for w in self.weights: - cumsum += w / total_weight - self.cumulative_weights.append(cumsum) - - # In-memory stats - self.stats: Dict[str, VariantStats] = { - v.value: VariantStats(variant=v.value) for v in self.variants - } - - # Thread safety - self._lock = threading.Lock() - - # Variant implementations - self._variant_impls: Dict[VariantType, RerankerVariant] = {} - self._init_variants() - - def _init_variants(self): - """Initialize reranker implementations for each variant.""" - # Baseline: just sort by initial score - def baseline_rerank(query, candidates, **kwargs): - return sorted(candidates, key=lambda x: x.get("score", 0), reverse=True) - - self._variant_impls[VariantType.BASELINE] = RerankerVariant( - VariantType.BASELINE, baseline_rerank - ) - - # Recursive reranker - try: - try: - from scripts.rerank_recursive import rerank_recursive - except ImportError: - from rerank_recursive import rerank_recursive - - self._variant_impls[VariantType.RECURSIVE] = RerankerVariant( - VariantType.RECURSIVE, - lambda q, c, **kw: rerank_recursive(q, c, n_iterations=3) - ) - except ImportError: - pass - - # ONNX reranker - try: - try: - from scripts.rerank_local import rerank_in_process - except ImportError: - from rerank_local import rerank_in_process - - self._variant_impls[VariantType.ONNX] = RerankerVariant( - VariantType.ONNX, - lambda q, c, **kw: rerank_in_process(q, c, limit=len(c)) - ) - except ImportError: - pass - - def _hash_to_bucket(self, session_id: str) -> float: - """Hash session ID to a value in [0, 1) for consistent bucketing.""" - h = hashlib.md5(f"{self.experiment_id}:{session_id}".encode()).hexdigest() - return int(h[:8], 16) / (16 ** 8) - - def get_variant_type(self, session_id: str) -> VariantType: - """Get the variant type for a session (consistent assignment).""" - bucket = self._hash_to_bucket(session_id) - for i, threshold in enumerate(self.cumulative_weights): - if bucket < threshold: - return self.variants[i] - return self.variants[-1] - - def get_variant(self, session_id: str) -> RerankerVariant: - """Get the reranker variant for a session.""" - variant_type = self.get_variant_type(session_id) - - if variant_type in self._variant_impls: - return self._variant_impls[variant_type] - - # Fall back to baseline if variant not available - return self._variant_impls.get(VariantType.BASELINE, - RerankerVariant(VariantType.BASELINE, lambda q, c, **kw: c)) - - def log_metrics( - self, - session_id: str, - latency_ms: float = 0.0, - clicked_rank: Optional[int] = None, - num_results: int = 0, - query_length: int = 0, - **custom - ): - """Log metrics for an observation.""" - variant_type = self.get_variant_type(session_id) - - metric = ABMetric( - session_id=session_id, - variant=variant_type.value, - timestamp=time.time(), - latency_ms=latency_ms, - clicked_rank=clicked_rank, - num_results=num_results, - query_length=query_length, - custom=custom, - ) - - with self._lock: - # Update in-memory stats - if variant_type.value in self.stats: - self.stats[variant_type.value].add_observation(metric) - - # Append to log file - self._append_metric(metric) - - def _append_metric(self, metric: ABMetric): - """Append metric to log file.""" - self.log_path.parent.mkdir(parents=True, exist_ok=True) - - record = { - "experiment_id": self.experiment_id, - "session_id": metric.session_id, - "variant": metric.variant, - "timestamp": metric.timestamp, - "latency_ms": metric.latency_ms, - "clicked_rank": metric.clicked_rank, - "num_results": metric.num_results, - "query_length": metric.query_length, - **metric.custom, - } - - with open(self.log_path, "a") as f: - f.write(json.dumps(record) + "\n") - - def get_summary(self) -> Dict[str, Any]: - """Get summary statistics for all variants.""" - return { - "experiment_id": self.experiment_id, - "variants": {v: self.stats[v].get_summary() for v in self.stats}, - } - - def print_summary(self): - """Print formatted summary.""" - summary = self.get_summary() - - print("\n" + "=" * 70) - print(f"A/B TEST RESULTS: {summary['experiment_id']}") - print("=" * 70) - - print(f"\n{'Variant':<15} {'N':<8} {'Latency (ms)':<15} {'MRR':<10} {'CTR@1':<10}") - print("-" * 70) - - for variant_name, stats in summary["variants"].items(): - n = stats.get("n_observations", 0) - lat = f"{stats.get('latency_ms_mean', 0):.1f} ± {stats.get('latency_ms_std', 0):.1f}" - mrr = f"{stats.get('mrr', 0):.4f}" - ctr = f"{stats.get('click_rate_at_1', 0):.2%}" - - print(f"{variant_name:<15} {n:<8} {lat:<15} {mrr:<10} {ctr:<10}") - - print("-" * 70) - print() - - -def simulate_ab_test(n_sessions: int = 100, n_queries_per_session: int = 5): - """Simulate an A/B test with synthetic data.""" - import random - - ab = ABTestManager( - variants=[VariantType.BASELINE, VariantType.RECURSIVE], - weights=[0.5, 0.5], - ) - - print(f"Simulating A/B test with {n_sessions} sessions...") - - for session_idx in range(n_sessions): - session_id = f"session_{session_idx}" - variant = ab.get_variant(session_id) - - for query_idx in range(n_queries_per_session): - # Generate fake candidates - candidates = [ - {"path": f"file_{i}.py", "score": random.random()} - for i in range(10) - ] - - # Rerank - results, latency_ms = variant.rerank("test query", candidates) - - # Simulate click (higher probability for top results) - clicked_rank = None - for rank, result in enumerate(results[:5], 1): - if random.random() < 0.3 / rank: # Decreasing probability - clicked_rank = rank - break - - # Log metrics - ab.log_metrics( - session_id=session_id, - latency_ms=latency_ms, - clicked_rank=clicked_rank, - num_results=len(results), - query_length=len("test query"), - ) - - ab.print_summary() - return ab - - -if __name__ == "__main__": - simulate_ab_test(n_sessions=100, n_queries_per_session=5) diff --git a/scripts/rerank_tools/benchmark.py b/scripts/rerank_tools/benchmark.py deleted file mode 100644 index 40bbe973..00000000 --- a/scripts/rerank_tools/benchmark.py +++ /dev/null @@ -1,505 +0,0 @@ -#!/usr/bin/env python3 -""" -Production-Grade Benchmark: Recursive Reranker on Real Codebase Data - -Tests against the actual indexed Context Engine codebase: -1. Queries the real Qdrant index -2. Compares recursive reranker vs baseline vs ONNX -3. Measures actual latency and ranking changes -4. Uses ground truth from ONNX reranker as reference - -Usage: - python scripts/rerank_real_benchmark.py -""" - -import os -import sys -import time -import json -import numpy as np -from pathlib import Path -from typing import List, Dict, Any, Optional -from dataclasses import dataclass, field - -# Add project root to path -sys.path.insert(0, str(Path(__file__).parent.parent)) - -# Real queries based on actual Context Engine functionality -REAL_QUERIES = [ - "hybrid search RRF fusion implementation", - "MCP server tool registration decorator", - "embedding model initialization cache", - "recursive reranker training deep supervision", - "file watcher debounce indexing", - "Qdrant vector search collection", - "context answer LLM generation citations", - "session aware latent state carryover", - "cache eviction policy LRU TTL", - "upload service delta bundle sync", - "query expansion semantic similarity", - "rerank ONNX cross encoder scoring", - "memory store find retrieval", - "subprocess manager process tracking", - "Docker compose service orchestration", -] - - -@dataclass -class RealBenchmarkResult: - """Result from benchmarking a single query.""" - query: str - reranker: str - latency_ms: float - num_results: int - top_5_paths: List[str] - top_5_scores: List[float] - # Ranking comparison metrics - kendall_tau: float = 0.0 # Correlation with reference ranking - top_3_overlap: float = 0.0 # Overlap of top-3 with reference - - -def get_real_candidates(query: str, limit: int = 30) -> List[Dict[str, Any]]: - """Get real candidates from hybrid search (production pipeline).""" - try: - from scripts.hybrid_search import run_hybrid_search - from scripts.embedder import get_embedding_model - - model_name = os.environ.get("EMBEDDING_MODEL", "BAAI/bge-base-en-v1.5") - model = get_embedding_model(model_name) - - # Run real hybrid search (dense + lexical fusion) - results = run_hybrid_search( - queries=[query], - limit=limit, - per_path=3, - model=model, - ) - - # Convert to candidate format - candidates = [] - for r in results: - candidates.append({ - "path": r.get("path", ""), - "symbol": r.get("symbol", ""), - "code": r.get("code", r.get("snippet", "")), - "score": float(r.get("score", 0)), - "start_line": r.get("start_line", 0), - "end_line": r.get("end_line", 0), - "components": r.get("components", {}), - }) - - return candidates - - except Exception as e: - print(f"Error getting candidates: {e}") - import traceback - traceback.print_exc() - return [] - - -def benchmark_baseline(query: str, candidates: List[Dict[str, Any]]) -> RealBenchmarkResult: - """Benchmark baseline (no reranking, just use hybrid scores).""" - start = time.perf_counter() - - # Just sort by existing score - sorted_cands = sorted(candidates, key=lambda x: x.get("score", 0), reverse=True) - - latency = (time.perf_counter() - start) * 1000 - - return RealBenchmarkResult( - query=query, - reranker="baseline", - latency_ms=latency, - num_results=len(sorted_cands), - top_5_paths=[c.get("path", "") for c in sorted_cands[:5]], - top_5_scores=[float(c.get("score", 0)) for c in sorted_cands[:5]], - ) - - -def benchmark_recursive(query: str, candidates: List[Dict[str, Any]], n_iters: int = 3) -> RealBenchmarkResult: - """Benchmark recursive reranker.""" - try: - from scripts.rerank_recursive import RecursiveReranker - except ImportError: - from rerank_recursive import RecursiveReranker - - reranker = RecursiveReranker(n_iterations=n_iters, dim=256) - initial_scores = [c.get("score", 0) for c in candidates] - - start = time.perf_counter() - reranked = reranker.rerank(query, candidates, initial_scores) - latency = (time.perf_counter() - start) * 1000 - - return RealBenchmarkResult( - query=query, - reranker=f"recursive_{n_iters}", - latency_ms=latency, - num_results=len(reranked), - top_5_paths=[c.get("path", "") for c in reranked[:5]], - top_5_scores=[float(c.get("score", 0)) for c in reranked[:5]], - ) - - -def benchmark_onnx(query: str, candidates: List[Dict[str, Any]]) -> Optional[RealBenchmarkResult]: - """Benchmark ONNX cross-encoder reranker on pre-fetched candidates.""" - try: - try: - from scripts.rerank_local import rerank_local - except ImportError: - from rerank_local import rerank_local - - # Prepare pairs for ONNX reranker - pairs = [] - for c in candidates: - doc = c.get("code", "") or c.get("snippet", "") - pairs.append((query, doc)) - - start = time.perf_counter() - scores = rerank_local(pairs) - latency = (time.perf_counter() - start) * 1000 - - # Combine scores with candidates and sort - scored = list(zip(scores, candidates)) - scored.sort(key=lambda x: x[0], reverse=True) - reranked = [{"score": s, **c} for s, c in scored] - - return RealBenchmarkResult( - query=query, - reranker="onnx", - latency_ms=latency, - num_results=len(reranked), - top_5_paths=[c.get("path", "") for c in reranked[:5]], - top_5_scores=[float(c.get("score", 0)) for c in reranked[:5]], - ) - except Exception as e: - print(f"ONNX reranker error: {e}") - return None - - -def benchmark_session_aware(query: str, candidates: List[Dict[str, Any]], session_id: str) -> RealBenchmarkResult: - """Benchmark session-aware recursive reranker.""" - try: - from scripts.rerank_recursive import SessionAwareReranker - except ImportError: - from rerank_recursive import SessionAwareReranker - - reranker = SessionAwareReranker(n_iterations=3, dim=256) - initial_scores = [c.get("score", 0) for c in candidates] - - start = time.perf_counter() - reranked = reranker.rerank(query, candidates, session_id=session_id, initial_scores=initial_scores) - latency = (time.perf_counter() - start) * 1000 - - return RealBenchmarkResult( - query=query, - reranker="session_aware", - latency_ms=latency, - num_results=len(reranked), - top_5_paths=[c.get("path", "") for c in reranked[:5]], - top_5_scores=[float(c.get("score", 0)) for c in reranked[:5]], - ) - - -# Global reranker for online learning (persists across queries) -_LEARNING_RERANKER = None - - -def get_learning_reranker(): - """Get or create the learning-enabled reranker.""" - global _LEARNING_RERANKER - if _LEARNING_RERANKER is None: - try: - from scripts.rerank_recursive import RecursiveReranker - except ImportError: - from rerank_recursive import RecursiveReranker - _LEARNING_RERANKER = RecursiveReranker(n_iterations=3, dim=256) - return _LEARNING_RERANKER - - -def benchmark_with_learning( - query: str, - candidates: List[Dict[str, Any]], - teacher_scores: Optional[List[float]] = None, -) -> RealBenchmarkResult: - """ - Benchmark recursive reranker WITH online learning from ONNX teacher. - - This learns from the ONNX scores to improve over time. - """ - reranker = get_learning_reranker() - initial_scores = [c.get("score", 0) for c in candidates] - - # Encode query and docs for learning - doc_texts = [] - for c in candidates: - parts = [] - if c.get("symbol"): - parts.append(str(c["symbol"])) - if c.get("path"): - parts.append(str(c["path"])) - code = c.get("code") or c.get("snippet") or "" - if code: - parts.append(str(code)[:500]) - doc_texts.append(" ".join(parts) if parts else "empty") - - # Get embeddings (cached after first call) - query_emb = reranker._encode([query])[0] - doc_embs = reranker._encode(doc_texts) - query_emb = reranker._project_to_dim(query_emb.reshape(1, -1))[0] - doc_embs = reranker._project_to_dim(doc_embs) - - # Learn from ONNX teacher if available - if teacher_scores is not None and len(teacher_scores) == len(candidates): - teacher_arr = np.array(teacher_scores, dtype=np.float32) - z = query_emb.copy() # Initial latent - reranker.scorer.learn_from_teacher(query_emb, doc_embs, z, teacher_arr) - - # Now do inference - start = time.perf_counter() - reranked = reranker.rerank(query, candidates, initial_scores) - latency = (time.perf_counter() - start) * 1000 - - return RealBenchmarkResult( - query=query, - reranker="learning", - latency_ms=latency, - num_results=len(reranked), - top_5_paths=[c.get("path", "") for c in reranked[:5]], - top_5_scores=[float(c.get("score", 0)) for c in reranked[:5]], - ) - - -def compute_ranking_correlation(ranking1: List[str], ranking2: List[str]) -> float: - """Compute ranking correlation (simplified Kendall's tau).""" - if not ranking1 or not ranking2: - return 0.0 - - # Create position maps - pos1 = {p: i for i, p in enumerate(ranking1)} - pos2 = {p: i for i, p in enumerate(ranking2)} - - common = set(ranking1) & set(ranking2) - if len(common) < 2: - return 0.0 - - concordant = 0 - discordant = 0 - - common_list = list(common) - for i in range(len(common_list)): - for j in range(i + 1, len(common_list)): - a, b = common_list[i], common_list[j] - # Compare relative ordering - order1 = pos1[a] < pos1[b] - order2 = pos2[a] < pos2[b] - if order1 == order2: - concordant += 1 - else: - discordant += 1 - - total = concordant + discordant - if total == 0: - return 0.0 - - return (concordant - discordant) / total - - -def run_real_benchmark(): - """Run the full benchmark on real data.""" - print("=" * 80) - print("PRODUCTION BENCHMARK: Recursive Reranker on Real Codebase") - print("=" * 80) - - # Check Qdrant connection - try: - from qdrant_client import QdrantClient - url = os.environ.get("QDRANT_URL", "http://localhost:6333") - client = QdrantClient(url=url) - coll = os.environ.get("COLLECTION_NAME", "codebase") - info = client.get_collection(coll) - print(f"\nQdrant collection: {coll}") - print(f"Points indexed: {info.points_count}") - except Exception as e: - print(f"Qdrant connection error: {e}") - print("Make sure Qdrant is running and indexed") - return - - results_by_reranker = { - "baseline": [], - "recursive_3": [], - "session_aware": [], - "learning": [], - "onnx": [], - } - - print(f"\nRunning {len(REAL_QUERIES)} queries...") - print("-" * 80) - - session_id = "benchmark_session" - - for i, query in enumerate(REAL_QUERIES): - print(f"\n[{i+1}/{len(REAL_QUERIES)}] Query: {query[:50]}...") - - # Get real candidates - candidates = get_real_candidates(query, limit=25) - if not candidates: - print(" No candidates found, skipping") - continue - - print(f" Candidates: {len(candidates)}") - - # Benchmark each reranker - baseline = benchmark_baseline(query, candidates) - results_by_reranker["baseline"].append(baseline) - print(f" Baseline: {baseline.latency_ms:.2f}ms") - - # Get ONNX scores first (teacher signal) - onnx_result = benchmark_onnx(query, candidates) - teacher_scores = None - if onnx_result: - results_by_reranker["onnx"].append(onnx_result) - print(f" ONNX: {onnx_result.latency_ms:.2f}ms") - teacher_scores = onnx_result.top_5_scores # Use full scores - # Get full ONNX scores for learning - try: - from scripts.rerank_local import rerank_local - except ImportError: - from rerank_local import rerank_local - pairs = [(query, c.get("code", "") or c.get("snippet", "")) for c in candidates] - teacher_scores = rerank_local(pairs) - - # Learning-enabled reranker (learns from ONNX before inference) - learning_result = benchmark_with_learning(query, candidates, teacher_scores) - results_by_reranker["learning"].append(learning_result) - print(f" Learning: {learning_result.latency_ms:.2f}ms") - - recursive = benchmark_recursive(query, candidates, n_iters=3) - results_by_reranker["recursive_3"].append(recursive) - print(f" Recursive(3): {recursive.latency_ms:.2f}ms") - - session = benchmark_session_aware(query, candidates, session_id) - results_by_reranker["session_aware"].append(session) - print(f" Session-aware: {session.latency_ms:.2f}ms") - - # Compare rankings with ONNX (ground truth) - if onnx_result: - learning_result.kendall_tau = compute_ranking_correlation( - onnx_result.top_5_paths, learning_result.top_5_paths - ) - recursive.kendall_tau = compute_ranking_correlation( - onnx_result.top_5_paths, recursive.top_5_paths - ) - session.kendall_tau = compute_ranking_correlation( - onnx_result.top_5_paths, session.top_5_paths - ) - else: - # Fall back to baseline comparison - learning_result.kendall_tau = compute_ranking_correlation( - baseline.top_5_paths, learning_result.top_5_paths - ) - recursive.kendall_tau = compute_ranking_correlation( - baseline.top_5_paths, recursive.top_5_paths - ) - session.kendall_tau = compute_ranking_correlation( - baseline.top_5_paths, session.top_5_paths - ) - - # Print summary - print("\n" + "=" * 80) - print("BENCHMARK RESULTS SUMMARY") - print("=" * 80) - - print(f"\n{'Reranker':<15} {'Queries':<8} {'Latency (ms)':<20} {'Rank Change':<15}") - print("-" * 60) - - for name, results in results_by_reranker.items(): - if not results: - continue - - n = len(results) - latencies = [r.latency_ms for r in results] - mean_lat = np.mean(latencies) - std_lat = np.std(latencies) - - taus = [r.kendall_tau for r in results if r.kendall_tau != 0] - mean_tau = np.mean(taus) if taus else 0.0 - - lat_str = f"{mean_lat:.2f} ± {std_lat:.2f}" - tau_str = f"{mean_tau:.3f}" if mean_tau else "N/A" - - print(f"{name:<15} {n:<8} {lat_str:<20} {tau_str:<15}") - - print("-" * 60) - - # Show example ranking differences - print("\n" + "=" * 80) - print("RANKING COMPARISON (First Query)") - print("=" * 80) - - if results_by_reranker["baseline"]: - print(f"\nQuery: {results_by_reranker['baseline'][0].query}") - - print("\nBaseline Top-5:") - for i, (path, score) in enumerate(zip( - results_by_reranker["baseline"][0].top_5_paths, - results_by_reranker["baseline"][0].top_5_scores - )): - print(f" {i+1}. {path[-50:]} (score: {score:.3f})") - - if results_by_reranker["recursive_3"]: - print("\nRecursive(3) Top-5:") - for i, (path, score) in enumerate(zip( - results_by_reranker["recursive_3"][0].top_5_paths, - results_by_reranker["recursive_3"][0].top_5_scores - )): - print(f" {i+1}. {path[-50:]} (score: {score:.3f})") - - if results_by_reranker["learning"]: - print("\nLearning Top-5 (after training on ONNX):") - for i, (path, score) in enumerate(zip( - results_by_reranker["learning"][0].top_5_paths, - results_by_reranker["learning"][0].top_5_scores - )): - print(f" {i+1}. {path[-50:]} (score: {score:.3f})") - - if results_by_reranker["onnx"]: - print("\nONNX Top-5 (teacher/ground truth):") - for i, (path, score) in enumerate(zip( - results_by_reranker["onnx"][0].top_5_paths, - results_by_reranker["onnx"][0].top_5_scores - )): - print(f" {i+1}. {path[-50:]} (score: {score:.3f})") - - # Show learning progress - if results_by_reranker["learning"]: - print("\n" + "=" * 80) - print("LEARNING PROGRESS (correlation with ONNX over time)") - print("=" * 80) - - taus = [r.kendall_tau for r in results_by_reranker["learning"]] - for i, tau in enumerate(taus[:10]): # First 10 - bar = "█" * int(tau * 20) if tau > 0 else "" - print(f" Query {i+1}: {tau:.3f} {bar}") - - if len(taus) > 10: - print(f" ... ({len(taus) - 10} more queries)") - - # Check if learning improved over time - if len(taus) >= 5: - early_avg = np.mean(taus[:len(taus)//2]) - late_avg = np.mean(taus[len(taus)//2:]) - improvement = late_avg - early_avg - print(f"\n Early avg: {early_avg:.3f}, Late avg: {late_avg:.3f}") - if improvement > 0: - print(f" Improvement: +{improvement:.3f} (learning is working!)") - else: - print(f" Change: {improvement:.3f}") - - print("\n" + "=" * 80) - print("BENCHMARK COMPLETE") - print("=" * 80) - - -if __name__ == "__main__": - run_real_benchmark() - diff --git a/scripts/rerank_tools/eval.py b/scripts/rerank_tools/eval.py deleted file mode 100644 index 1676b967..00000000 --- a/scripts/rerank_tools/eval.py +++ /dev/null @@ -1,409 +0,0 @@ -#!/usr/bin/env python3 -""" -Offline Reranker Evaluation Pipeline - -Runs fixed-query evaluation for reranker quality with MRR/Recall/latency metrics. -Designed for CI/regression testing - deterministic, no sampling. - -Usage: - python scripts/rerank_eval.py [--queries QUERIES_FILE] [--output OUTPUT_FILE] - python scripts/rerank_eval.py --ablations # Run all ablation modes - -Metrics reported: - - MRR@k (Mean Reciprocal Rank) - - Recall@k (fraction of relevant docs in top-k) - - Latency p50/p95/p99 -""" - -import argparse -import copy -import json -import os -import sys -import time -from dataclasses import dataclass, field, asdict -from pathlib import Path -from typing import List, Dict, Any, Optional - -import numpy as np - -# Add project root to path -sys.path.insert(0, str(Path(__file__).parent.parent)) - -# Fixed evaluation queries (deterministic, no sampling) -DEFAULT_EVAL_QUERIES = [ - "hybrid search RRF fusion implementation", - "MCP server tool registration decorator", - "embedding model initialization cache", - "recursive reranker training deep supervision", - "Qdrant vector search collection create", - "context answer LLM generation citations", - "cache eviction policy LRU TTL", - "query expansion semantic similarity", - "rerank ONNX cross encoder scoring", - "memory store find retrieval", -] - - -@dataclass -class EvalResult: - """Evaluation result for a single query.""" - query: str - mode: str - latency_ms: float - top_k_paths: List[str] - top_k_scores: List[float] - mrr: float = 0.0 - recall_at_5: float = 0.0 - recall_at_10: float = 0.0 - - -@dataclass -class EvalSummary: - """Aggregated evaluation summary.""" - mode: str - num_queries: int - mrr_mean: float - recall_at_5_mean: float - recall_at_10_mean: float - latency_p50_ms: float - latency_p95_ms: float - latency_p99_ms: float - results: List[EvalResult] = field(default_factory=list) - - -def get_candidates(query: str, limit: int = 30) -> List[Dict[str, Any]]: - """Get candidates from hybrid search.""" - try: - from scripts.hybrid_search import run_hybrid_search - from scripts.embedder import get_embedding_model - - # Use BAAI/bge-base-en-v1.5 which is supported by fastembed - model_name = os.environ.get("EMBEDDING_MODEL", "BAAI/bge-base-en-v1.5") - model = get_embedding_model(model_name) - - results = run_hybrid_search( - queries=[query], - limit=limit, - per_path=3, - model=model, - ) - - candidates = [] - for r in results: - candidates.append({ - "path": r.get("path", ""), - "symbol": r.get("symbol", ""), - "start_line": r.get("start_line", 0), - "end_line": r.get("end_line", 0), - "score": float(r.get("score", 0)), - "snippet": r.get("snippet", "")[:500] if r.get("snippet") else "", - }) - return candidates - except Exception as e: - print(f"Warning: Could not get candidates: {e}", file=sys.stderr) - return [] - - -def get_onnx_scores(query: str, candidates: List[Dict[str, Any]]) -> Optional[List[float]]: - """Get ONNX reranker scores (ground truth).""" - try: - from scripts.rerank_local import rerank_local - pairs = [] - for c in candidates: - doc_parts = [] - if c.get("symbol"): - doc_parts.append(str(c["symbol"])) - if c.get("path"): - doc_parts.append(str(c["path"])) - code = c.get("code") or c.get("snippet") or "" - if code: - doc_parts.append(code[:500]) - doc = " ".join(doc_parts) if doc_parts else "empty" - pairs.append((query, doc)) - return rerank_local(pairs) - except Exception: - return None - - -def rerank_baseline(query: str, candidates: List[Dict[str, Any]]) -> List[Dict[str, Any]]: - """Baseline: no reranking, just return as-is.""" - return candidates - - -def rerank_recursive( - query: str, - candidates: List[Dict[str, Any]], - n_iterations: int = 3, -) -> List[Dict[str, Any]]: - """Recursive reranker (no learning).""" - try: - from scripts.rerank_recursive import RecursiveReranker - reranker = RecursiveReranker(n_iterations=n_iterations, dim=256) - initial_scores = [c.get("score", 0) for c in candidates] - return reranker.rerank(query, candidates, initial_scores) - except Exception as e: - print(f"Warning: Recursive rerank failed: {e}", file=sys.stderr) - return candidates - - -def rerank_learning( - query: str, - candidates: List[Dict[str, Any]], - collection: str = "eval", -) -> List[Dict[str, Any]]: - """Learning reranker (uses trained weights).""" - try: - from scripts.rerank_recursive import rerank_with_learning - return rerank_with_learning( - query=query, - candidates=candidates, - limit=len(candidates), - n_iterations=3, - learn_from_onnx=False, # Eval mode: no training - collection=collection, - ) - except Exception as e: - print(f"Warning: Learning rerank failed: {e}", file=sys.stderr) - return candidates - - -def rerank_onnx(query: str, candidates: List[Dict[str, Any]]) -> List[Dict[str, Any]]: - """ONNX reranker (teacher/ground truth).""" - scores = get_onnx_scores(query, candidates) - if scores is None: - return candidates - for c, s in zip(candidates, scores): - c["score"] = s - return sorted(candidates, key=lambda x: x.get("score", 0), reverse=True) - - -def compute_mrr(ranked_paths: List[str], relevant_paths: List[str], k: int = 10) -> float: - """Compute Mean Reciprocal Rank.""" - # Deduplicate while preserving order (paths can repeat due to multi-span retrieval) - rel_set = {p for p in relevant_paths if p} - if not rel_set: - return 0.0 - - seen = set() - uniq_ranked: List[str] = [] - for p in ranked_paths: - if not p or p in seen: - continue - seen.add(p) - uniq_ranked.append(p) - if len(uniq_ranked) >= k: - break - - for i, path in enumerate(uniq_ranked): - if path in rel_set: - return 1.0 / (i + 1) - return 0.0 - - -def compute_recall_at_k(ranked_paths: List[str], relevant_paths: List[str], k: int) -> float: - """Compute Recall@k.""" - rel_set = {p for p in relevant_paths if p} - if not rel_set: - return 0.0 - - seen = set() - uniq_ranked: List[str] = [] - for p in ranked_paths: - if not p or p in seen: - continue - seen.add(p) - uniq_ranked.append(p) - if len(uniq_ranked) >= k: - break - - found = len(set(uniq_ranked) & rel_set) - return found / len(rel_set) - - -def eval_single_query( - query: str, - mode: str, - candidates: List[Dict[str, Any]], - reference_paths: List[str], -) -> EvalResult: - """Evaluate a single query with a specific reranking mode.""" - start = time.perf_counter() - - if mode == "baseline": - reranked = rerank_baseline(query, candidates) - elif mode == "recursive": - reranked = rerank_recursive(query, candidates) - elif mode == "learning": - reranked = rerank_learning(query, candidates) - elif mode == "onnx": - reranked = rerank_onnx(query, candidates) - else: - reranked = candidates - - latency_ms = (time.perf_counter() - start) * 1000 - - # Export unique paths (avoid duplicates from multi-span retrieval) - ranked_paths: List[str] = [] - ranked_scores: List[float] = [] - seen_paths = set() - for c in reranked: - path = str(c.get("path", "") or "") - if not path or path in seen_paths: - continue - seen_paths.add(path) - ranked_paths.append(path) - ranked_scores.append(float(c.get("score", 0))) - if len(ranked_paths) >= 10: - break - - mrr = compute_mrr(ranked_paths, reference_paths) - recall_5 = compute_recall_at_k(ranked_paths, reference_paths, 5) - recall_10 = compute_recall_at_k(ranked_paths, reference_paths, 10) - - return EvalResult( - query=query, - mode=mode, - latency_ms=latency_ms, - top_k_paths=ranked_paths, - top_k_scores=ranked_scores, - mrr=mrr, - recall_at_5=recall_5, - recall_at_10=recall_10, - ) - - -def run_eval( - queries: List[str], - modes: List[str], - use_onnx_reference: bool = True, -) -> Dict[str, EvalSummary]: - """Run evaluation across all queries and modes.""" - summaries: Dict[str, EvalSummary] = {} - - # Warmup: pre-cache embeddings for all queries (cold start is not representative) - print("Warming up embedding cache...") - for query in queries: - candidates = get_candidates(query) - if candidates and "learning" in modes: - # Run once to cache embeddings - rerank_learning(query, copy.deepcopy(candidates)) - print("Warmup complete.") - - for mode in modes: - results = [] - latencies = [] - - for query in queries: - candidates = get_candidates(query) - if not candidates: - continue - - # Use ONNX top-5 as "relevant" ground truth - # Deep copy to prevent score mutation from leaking between modes - reference_paths = [] - if use_onnx_reference: - onnx_ranked = rerank_onnx(query, copy.deepcopy(candidates)) - reference_paths = [c.get("path", "") for c in onnx_ranked[:5]] - - result = eval_single_query(query, mode, copy.deepcopy(candidates), reference_paths) - results.append(result) - latencies.append(result.latency_ms) - - if not results: - continue - - latencies_arr = np.array(latencies) - summaries[mode] = EvalSummary( - mode=mode, - num_queries=len(results), - mrr_mean=np.mean([r.mrr for r in results]), - recall_at_5_mean=np.mean([r.recall_at_5 for r in results]), - recall_at_10_mean=np.mean([r.recall_at_10 for r in results]), - latency_p50_ms=float(np.percentile(latencies_arr, 50)), - latency_p95_ms=float(np.percentile(latencies_arr, 95)), - latency_p99_ms=float(np.percentile(latencies_arr, 99)), - results=results, - ) - - return summaries - - -def print_summary(summaries: Dict[str, EvalSummary]): - """Print evaluation summary table.""" - print("\n" + "=" * 80) - print("RERANKER EVALUATION SUMMARY") - print("=" * 80) - print(f"{'Mode':<12} {'MRR':<8} {'R@5':<8} {'R@10':<8} {'p50ms':<8} {'p95ms':<8} {'p99ms':<8}") - print("-" * 80) - - # Track for comparison - baseline_mrr = summaries.get("baseline", EvalSummary("baseline", 0, 0, 0, 0, 0, 0, 0)).mrr_mean - onnx_mrr = summaries.get("onnx", EvalSummary("onnx", 0, 0, 0, 0, 0, 0, 0)).mrr_mean - onnx_p50 = summaries.get("onnx", EvalSummary("onnx", 0, 0, 0, 0, 0, 0, 0)).latency_p50_ms - learning_mrr = summaries.get("learning", EvalSummary("learning", 0, 0, 0, 0, 0, 0, 0)).mrr_mean - learning_p50 = summaries.get("learning", EvalSummary("learning", 0, 0, 0, 0, 0, 0, 0)).latency_p50_ms - - for mode, summary in summaries.items(): - print( - f"{mode:<12} " - f"{summary.mrr_mean:<8.3f} " - f"{summary.recall_at_5_mean:<8.3f} " - f"{summary.recall_at_10_mean:<8.3f} " - f"{summary.latency_p50_ms:<8.1f} " - f"{summary.latency_p95_ms:<8.1f} " - f"{summary.latency_p99_ms:<8.1f}" - ) - print("=" * 80) - - # Self-improving search analysis - print("\nSELF-IMPROVING SEARCH ANALYSIS:") - print("-" * 50) - - if baseline_mrr > 0: - learning_vs_baseline = ((learning_mrr - baseline_mrr) / baseline_mrr) * 100 - print(f"Learning vs Baseline: {learning_vs_baseline:+.1f}% MRR improvement") - - if onnx_p50 > 0 and learning_p50 > 0: - speedup = onnx_p50 / learning_p50 - print(f"Learning vs ONNX: {speedup:.1f}x faster ({learning_p50:.1f}ms vs {onnx_p50:.1f}ms)") - - if onnx_mrr > 0: - distill_quality = (learning_mrr / onnx_mrr) * 100 - print(f"Distillation quality: {distill_quality:.1f}% of ONNX MRR") - - print("-" * 50) - - -def main(): - parser = argparse.ArgumentParser(description="Offline reranker evaluation") - parser.add_argument("--queries", type=str, help="JSON file with query list") - parser.add_argument("--output", type=str, help="Output JSON file for results") - parser.add_argument("--ablations", action="store_true", help="Run all ablation modes") - parser.add_argument("--modes", type=str, default="baseline,recursive,learning", - help="Comma-separated modes to evaluate") - args = parser.parse_args() - - queries = DEFAULT_EVAL_QUERIES - if args.queries: - with open(args.queries) as f: - queries = json.load(f) - - modes = args.modes.split(",") - if args.ablations: - modes = ["baseline", "recursive", "learning", "onnx"] - - print(f"Running evaluation: {len(queries)} queries, modes: {modes}") - summaries = run_eval(queries, modes) - print_summary(summaries) - - if args.output: - output_data = {mode: asdict(s) for mode, s in summaries.items()} - with open(args.output, "w") as f: - json.dump(output_data, f, indent=2) - print(f"Results saved to {args.output}") - - -if __name__ == "__main__": - main() - diff --git a/scripts/rerank_tools/events.py b/scripts/rerank_tools/events.py index e64f2ad3..e4ca67bf 100644 --- a/scripts/rerank_tools/events.py +++ b/scripts/rerank_tools/events.py @@ -1,235 +1,52 @@ -#!/usr/bin/env python3 -""" -Rerank Training Event Logger. +"""Shared filesystem helpers for relevance-feedback event files.""" -Logs training events (query, candidates, scores, collection) to a file for -background processing. This keeps the MCP hot path fast and deterministic. +from __future__ import annotations -Features: -- Time-sharded files (hourly) to avoid giant files and enable parallel processing -- Configurable sampling rate to reduce volume at high QPS -- Events written as newline-delimited JSON (NDJSON) for streaming reads -""" - -import json import os -import random -import time import threading +from pathlib import Path +from typing import Dict + try: import fcntl # type: ignore except Exception: # pragma: no cover fcntl = None -from datetime import datetime -from pathlib import Path -from typing import Any, Dict, List, Optional - -# Event log configuration -RERANK_EVENTS_DIR = os.environ.get("RERANK_EVENTS_DIR", "/tmp/rerank_events") -RERANK_EVENTS_ENABLED = str(os.environ.get("RERANK_EVENTS_ENABLED", "1")).strip().lower() in { - "1", "true", "yes", "on" -} -# Sampling rate: 0.5 = log 50% of events (balance volume vs learning signal) -RERANK_EVENTS_SAMPLE_RATE = float(os.environ.get("RERANK_EVENTS_SAMPLE_RATE", "0.5")) -# Retention: files older than this many days can be cleaned up (0 = keep forever) -RERANK_EVENTS_RETENTION_DAYS = int(os.environ.get("RERANK_EVENTS_RETENTION_DAYS", "0")) -# Thread-safe write lock (per-file locks for better concurrency) -_WRITE_LOCKS: Dict[str, threading.Lock] = {} -_LOCKS_LOCK = threading.Lock() - -def _get_write_lock(file_key: str) -> threading.Lock: - """Get or create a write lock for a specific file.""" - with _LOCKS_LOCK: - if file_key not in _WRITE_LOCKS: - _WRITE_LOCKS[file_key] = threading.Lock() - return _WRITE_LOCKS[file_key] +def _get_events_dir() -> Path: + return Path(os.environ.get("RERANK_EVENTS_DIR", "/tmp/rerank_events")) def _ensure_events_dir() -> Path: - """Ensure events directory exists.""" - events_dir = Path(RERANK_EVENTS_DIR) + events_dir = _get_events_dir() events_dir.mkdir(parents=True, exist_ok=True) return events_dir -def _get_hour_suffix() -> str: - """Get current hour suffix for time-sharding (YYYYMMDDHH).""" - return datetime.now(tz=None).strftime("%Y%m%d%H") - - -def _get_events_file(collection: str, hour_suffix: Optional[str] = None) -> Path: - """Get events file path for a collection (time-sharded).""" - safe_name = "".join(c if c.isalnum() or c in "-_" else "_" for c in collection) - if hour_suffix is None: - hour_suffix = _get_hour_suffix() - return _ensure_events_dir() / f"events_{safe_name}_{hour_suffix}.ndjson" - - -def log_training_event( - query: str, - candidates: List[Dict[str, Any]], - initial_scores: List[float], - teacher_scores: Optional[List[float]], - collection: str, - metadata: Optional[Dict[str, Any]] = None, - force: bool = False, -) -> bool: - """ - Log a training event for background processing. - - Args: - query: The search query - candidates: List of candidate documents (will extract path, symbol, snippet) - initial_scores: Initial hybrid search scores - teacher_scores: ONNX teacher scores (if available) - collection: Collection name for isolation - metadata: Optional additional metadata - force: If True, bypass sampling (always log) - - Returns: - True if event was logged successfully, False if skipped/disabled - """ - if not RERANK_EVENTS_ENABLED: - return False - - # Sampling: only log SAMPLE_RATE fraction of events - if not force and random.random() > RERANK_EVENTS_SAMPLE_RATE: - return False - - try: - # Helper to convert numpy types to native Python for JSON - def _to_native(v): - if hasattr(v, "item"): # numpy scalar - return v.item() - return v +_WRITE_LOCKS: Dict[str, threading.Lock] = {} +_LOCKS_LOCK = threading.Lock() - # Extract minimal candidate info (don't store full code) - candidate_info = [] - for i, c in enumerate(candidates): - score = initial_scores[i] if i < len(initial_scores) else 0 - info = { - "path": c.get("path", ""), - "symbol": c.get("symbol", ""), - "start_line": c.get("start_line", 0), - "end_line": c.get("end_line", 0), - "initial_score": _to_native(score), - } - # Include small snippet for learning (truncated) - snippet = c.get("code") or c.get("snippet") or "" - if snippet: - info["snippet"] = snippet[:500] - candidate_info.append(info) - event = { - "ts": time.time(), - "query": query, - "collection": collection, - "candidates": candidate_info, - "teacher_scores": [_to_native(s) for s in teacher_scores] if teacher_scores is not None and len(teacher_scores) > 0 else None, - "metadata": metadata or {}, - } +def _get_write_lock(file_key: str) -> threading.Lock: + """Return the process-local lock for one event file.""" + with _LOCKS_LOCK: + if file_key not in _WRITE_LOCKS: + _WRITE_LOCKS[file_key] = threading.Lock() + return _WRITE_LOCKS[file_key] - events_file = _get_events_file(collection) - file_key = str(events_file) - # Per-file lock for better concurrency across collections/hours - lock = _get_write_lock(file_key) - with lock: - # Atomic append with file locking (for cross-process safety) - with open(events_file, "a") as f: +def append_event_line(events_file: Path, line: str) -> None: + """Append one NDJSON event safely across threads and worker processes.""" + events_file = Path(events_file) + events_file.parent.mkdir(parents=True, exist_ok=True) + lock = _get_write_lock(str(events_file)) + with lock: + with events_file.open("a", encoding="utf-8") as event_file: + if fcntl is not None: + fcntl.flock(event_file.fileno(), fcntl.LOCK_EX) + try: + event_file.write(line.rstrip("\n") + "\n") + event_file.flush() + finally: if fcntl is not None: - fcntl.flock(f.fileno(), fcntl.LOCK_EX) - try: - f.write(json.dumps(event) + "\n") - finally: - fcntl.flock(f.fileno(), fcntl.LOCK_UN) - else: - # Best-effort fallback (thread lock still prevents intra-process interleaving) - f.write(json.dumps(event) + "\n") - - return True - - except Exception: - return False - - -def list_event_files(collection: str) -> List[Path]: - """List all event files for a collection (sorted by time, oldest first).""" - safe_name = "".join(c if c.isalnum() or c in "-_" else "_" for c in collection) - pattern = f"events_{safe_name}_*.ndjson" - events_dir = _ensure_events_dir() - files = sorted(events_dir.glob(pattern)) - return files - - -def read_events( - collection: str, - since_ts: float = 0, - limit: int = 1000, -) -> List[Dict[str, Any]]: - """ - Read training events for a collection (across all time-sharded files). - - Args: - collection: Collection name - since_ts: Only return events after this timestamp - limit: Maximum events to return - - Returns: - List of training events (oldest first) - """ - event_files = list_event_files(collection) - if not event_files: - return [] - - events = [] - for events_file in event_files: - try: - with open(events_file, "r") as f: - for line in f: - line = line.strip() - if not line: - continue - try: - event = json.loads(line) - if event.get("ts", 0) > since_ts: - events.append(event) - if len(events) >= limit: - return events - except json.JSONDecodeError: - continue - except Exception: - continue - - return events - - -def cleanup_old_events(collection: str, max_age_days: int) -> int: - """ - Remove event files older than max_age_days. - - Args: - collection: Collection name - max_age_days: Delete files older than this - - Returns: - Number of files deleted - """ - if max_age_days <= 0: - return 0 - - cutoff = time.time() - (max_age_days * 86400) - deleted = 0 - - for events_file in list_event_files(collection): - try: - if events_file.stat().st_mtime < cutoff: - events_file.unlink() - deleted += 1 - except Exception: - continue - - return deleted - + fcntl.flock(event_file.fileno(), fcntl.LOCK_UN) diff --git a/scripts/rerank_tools/local.py b/scripts/rerank_tools/local.py index e2151791..40cb54ac 100644 --- a/scripts/rerank_tools/local.py +++ b/scripts/rerank_tools/local.py @@ -1,43 +1,27 @@ #!/usr/bin/env python3 import os import argparse -import sys import threading -from pathlib import Path as _P from typing import List, Dict, Any, TYPE_CHECKING -# Ensure project root is on sys.path when run as a script (so 'scripts' package imports work) -_ROOT = _P(__file__).resolve().parent.parent.parent -if str(_ROOT) not in sys.path: - sys.path.insert(0, str(_ROOT)) - from qdrant_client import QdrantClient, models # Import TextEmbedding for type hints (may not be available at runtime with embedder factory) if TYPE_CHECKING: from fastembed import TextEmbedding -# Use embedder factory for Qwen3 support; fallback to direct fastembed -try: - from scripts.embedder import get_embedding_model as _get_embedding_model - _EMBEDDER_FACTORY = True -except ImportError: - _EMBEDDER_FACTORY = False - from fastembed import TextEmbedding +from scripts.embedder import get_embedding_model as _get_embedding_model + +_EMBEDDER_FACTORY = True # Use centralized reranker factory (supports FastEmbed + ONNX backends) -try: - from scripts.reranker import ( - get_reranker_model as _get_reranker_model, - rerank_pairs as _rerank_pairs, - is_reranker_available as _is_reranker_available, - ) - _RERANKER_FACTORY = True -except ImportError: - _RERANKER_FACTORY = False - _get_reranker_model = None - _rerank_pairs = None - _is_reranker_available = None +from scripts.reranker import ( + get_reranker_model as _get_reranker_model, + rerank_pairs as _rerank_pairs, + is_reranker_available as _is_reranker_available, +) + +_RERANKER_FACTORY = True # Legacy ONNX imports (fallback when factory unavailable) try: @@ -134,6 +118,10 @@ def _get_rerank_session(): from scripts.utils import sanitize_vector_name as _sanitize_vector_name +from scripts.path_scope import ( + normalize_under as _normalize_under_scope, + metadata_matches_under as _metadata_matches_under, +) def warmup_reranker(): @@ -163,18 +151,14 @@ def _start_background_warmup(): _start_background_warmup() -def _norm_under(u: str | None) -> str | None: - if not u: - return None - u = str(u).strip().replace("\\", "/") - u = "/".join([p for p in u.split("/") if p]) - if not u: - return None - if not u.startswith("/"): - return "/work/" + u - if not u.startswith("/work/"): - return "/work/" + u.lstrip("/") - return u +def _point_matches_under(pt: Any, under: str | None) -> bool: + if not under: + return True + payload = getattr(pt, "payload", None) or {} + md = payload.get("metadata") or {} + if not isinstance(md, dict): + md = {} + return _metadata_matches_under(md, under) def _select_dense_vector_name( @@ -366,18 +350,21 @@ def rerank_in_process( key="metadata.language", match=models.MatchValue(value=language) ) ) - eff_under = _norm_under(under) - if eff_under: - must.append( - models.FieldCondition( - key="metadata.path_prefix", match=models.MatchValue(value=eff_under) - ) - ) + eff_under = _normalize_under_scope(under) flt = models.Filter(must=must) if must else None - pts = dense_results(client, _model, vec_name, query, flt, topk, eff_collection) - if not pts and flt is not None: - pts = dense_results(client, _model, vec_name, query, None, topk, eff_collection) + fetch_topk = max(1, int(topk)) + if eff_under: + try: + under_mult = int(os.environ.get("RERANK_UNDER_FETCH_MULT", "4") or 4) + except Exception: + under_mult = 4 + fetch_topk = max(fetch_topk, int(limit) * max(under_mult, 2), fetch_topk * max(under_mult, 2)) + fetch_topk = min(fetch_topk, 2000) + + pts = dense_results(client, _model, vec_name, query, flt, fetch_topk, eff_collection) + if eff_under and pts: + pts = [pt for pt in pts if _point_matches_under(pt, eff_under)] if not pts: return [] @@ -447,19 +434,21 @@ def main(): key="metadata.language", match=models.MatchValue(value=args.language) ) ) - eff_under = _norm_under(args.under) - if eff_under: - must.append( - models.FieldCondition( - key="metadata.path_prefix", match=models.MatchValue(value=eff_under) - ) - ) + eff_under = _normalize_under_scope(args.under) flt = models.Filter(must=must) if must else None - pts = dense_results(client, model, vec_name, args.query, flt, args.topk, eff_collection) - # Fallback: if filtered search yields nothing, retry without filters to avoid empty rerank - if not pts and flt is not None: - pts = dense_results(client, model, vec_name, args.query, None, args.topk, eff_collection) + fetch_topk = max(1, int(args.topk)) + if eff_under: + try: + under_mult = int(os.environ.get("RERANK_UNDER_FETCH_MULT", "4") or 4) + except Exception: + under_mult = 4 + fetch_topk = max(fetch_topk, int(args.limit) * max(under_mult, 2), fetch_topk * max(under_mult, 2)) + fetch_topk = min(fetch_topk, 2000) + + pts = dense_results(client, model, vec_name, args.query, flt, fetch_topk, eff_collection) + if eff_under and pts: + pts = [pt for pt in pts if _point_matches_under(pt, eff_under)] if not pts: return pairs = prepare_pairs(args.query, pts) diff --git a/scripts/rerank_tools/query.py b/scripts/rerank_tools/query.py index 9c3eeee9..9957c694 100644 --- a/scripts/rerank_tools/query.py +++ b/scripts/rerank_tools/query.py @@ -1,27 +1,17 @@ #!/usr/bin/env python3 import os import argparse -import sys from collections import defaultdict from typing import List, Dict, Any -from pathlib import Path from qdrant_client import QdrantClient, models import re -ROOT_DIR = Path(__file__).resolve().parents[2] -if str(ROOT_DIR) not in sys.path: - sys.path.insert(0, str(ROOT_DIR)) - from scripts.utils import sanitize_vector_name -# Use embedder factory for Qwen3 support; fallback to direct fastembed -try: - from scripts.embedder import get_embedding_model as _get_embedding_model - _EMBEDDER_FACTORY = True -except ImportError: - _EMBEDDER_FACTORY = False - from fastembed import TextEmbedding +from scripts.embedder import get_embedding_model as _get_embedding_model + +_EMBEDDER_FACTORY = True # Env configuration diff --git a/scripts/rerank_tools/train.py b/scripts/rerank_tools/train.py deleted file mode 100644 index f3bafaf3..00000000 --- a/scripts/rerank_tools/train.py +++ /dev/null @@ -1,661 +0,0 @@ -#!/usr/bin/env python3 -""" -Training infrastructure for Tiny Recursive Reranker. - -Implements TRM-style deep supervision training: -1. Generate training data from search logs or synthetic examples -2. Train with loss at each refinement iteration (deep supervision) -3. Use pairwise ranking loss (margin ranking) -4. Support weight saving/loading - -Usage: - # Generate synthetic training data - python scripts/rerank_train.py --generate-data --output data/rerank_train.jsonl - - # Train the model - python scripts/rerank_train.py --train --data data/rerank_train.jsonl --epochs 100 - - # Evaluate - python scripts/rerank_train.py --evaluate --data data/rerank_test.jsonl -""" - -import os -import json -import numpy as np -from pathlib import Path -from typing import List, Dict, Any, Optional, Tuple -from dataclasses import dataclass, field -import argparse - - -@dataclass -class TrainingExample: - """A single training example for pairwise ranking.""" - query: str - doc_positive: str # Should rank higher - doc_negative: str # Should rank lower - # Optional: relevance scores for regression loss - score_positive: float = 1.0 - score_negative: float = 0.0 - - -@dataclass -class TrainingConfig: - """Configuration for training.""" - learning_rate: float = 0.001 - batch_size: int = 32 - epochs: int = 100 - margin: float = 0.5 # Margin for pairwise loss - n_iterations: int = 3 # Refinement iterations - deep_supervision_weight: float = 0.5 # Weight for intermediate losses - dim: int = 256 - hidden_dim: int = 512 - weight_decay: float = 0.01 - save_every: int = 10 # Save checkpoint every N epochs - - -class SimpleGradientDescent: - """ - Simple gradient descent optimizer with momentum. - - We implement our own to avoid PyTorch/TensorFlow dependency. - """ - - def __init__(self, lr: float = 0.001, momentum: float = 0.9, weight_decay: float = 0.01): - self.lr = lr - self.momentum = momentum - self.weight_decay = weight_decay - self.velocities: Dict[str, np.ndarray] = {} - - def step(self, params: Dict[str, np.ndarray], grads: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]: - """Update parameters using gradients.""" - updated = {} - for name, param in params.items(): - if name not in grads: - updated[name] = param - continue - - grad = grads[name] - - # Weight decay (L2 regularization) - grad = grad + self.weight_decay * param - - # Momentum - if name not in self.velocities: - self.velocities[name] = np.zeros_like(param) - - self.velocities[name] = self.momentum * self.velocities[name] - self.lr * grad - updated[name] = param + self.velocities[name] - - return updated - - -class TrainableTinyScorer: - """ - Tiny scorer with gradient computation for training. - - Uses numerical gradients for simplicity (analytical gradients would be faster). - """ - - def __init__(self, dim: int = 256, hidden_dim: int = 512, seed: int = 42): - self.dim = dim - self.hidden_dim = hidden_dim - - # Initialize weights - np.random.seed(seed) - scale = np.sqrt(2.0 / (dim * 3)) # He initialization - self.params = { - "W1": np.random.randn(dim * 3, hidden_dim).astype(np.float32) * scale, - "b1": np.zeros(hidden_dim, dtype=np.float32), - "W2": np.random.randn(hidden_dim, 1).astype(np.float32) * np.sqrt(2.0 / hidden_dim), - "b2": np.zeros(1, dtype=np.float32), - } - - def forward(self, x: np.ndarray) -> Tuple[np.ndarray, Dict[str, np.ndarray]]: - """ - Forward pass with cached activations for backprop. - - x: (batch, dim*3) concatenated [query, doc, latent] - returns: (batch,) scores, cache dict - """ - # Layer 1: Linear + ReLU - z1 = x @ self.params["W1"] + self.params["b1"] # (batch, hidden) - h1 = np.maximum(0, z1) # ReLU - - # Layer 2: Linear - z2 = h1 @ self.params["W2"] + self.params["b2"] # (batch, 1) - scores = z2.squeeze(-1) # (batch,) - - cache = {"x": x, "z1": z1, "h1": h1} - return scores, cache - - def backward(self, dscores: np.ndarray, cache: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]: - """ - Backward pass to compute gradients. - - dscores: (batch,) gradient of loss w.r.t. scores - returns: dict of gradients for each parameter - """ - batch_size = dscores.shape[0] - - # Reshape for matrix ops - dz2 = dscores.reshape(-1, 1) # (batch, 1) - - # Layer 2 gradients - dW2 = cache["h1"].T @ dz2 # (hidden, 1) - db2 = dz2.sum(axis=0) # (1,) - dh1 = dz2 @ self.params["W2"].T # (batch, hidden) - - # ReLU backward - dz1 = dh1 * (cache["z1"] > 0).astype(np.float32) # (batch, hidden) - - # Layer 1 gradients - dW1 = cache["x"].T @ dz1 # (dim*3, hidden) - db1 = dz1.sum(axis=0) # (hidden,) - - # Average over batch - grads = { - "W1": dW1 / batch_size, - "b1": db1 / batch_size, - "W2": dW2 / batch_size, - "b2": db2 / batch_size, - } - return grads - - def save(self, path: str): - """Save model weights to file.""" - np.savez(path, **self.params) - - def load(self, path: str): - """Load model weights from file.""" - data = np.load(path) - for key in self.params: - if key in data: - self.params[key] = data[key] - - -class TrainableLatentRefiner: - """Latent refiner with gradient support.""" - - def __init__(self, dim: int = 256, hidden_dim: int = 256, seed: int = 43): - self.dim = dim - self.hidden_dim = hidden_dim - - np.random.seed(seed) - scale = np.sqrt(2.0 / (dim * 3)) - self.params = { - "W1": np.random.randn(dim * 3, hidden_dim).astype(np.float32) * scale, - "b1": np.zeros(hidden_dim, dtype=np.float32), - "W2": np.random.randn(hidden_dim, dim).astype(np.float32) * np.sqrt(2.0 / hidden_dim), - "b2": np.zeros(dim, dtype=np.float32), - } - - def forward(self, z: np.ndarray, query: np.ndarray, doc_summary: np.ndarray, alpha: float = 0.5) -> np.ndarray: - """Refine latent state.""" - x = np.concatenate([z, query, doc_summary], axis=-1) - h = np.maximum(0, x @ self.params["W1"] + self.params["b1"]) - z_new = h @ self.params["W2"] + self.params["b2"] - z_refined = alpha * z_new + (1 - alpha) * z - # Normalize - norm = np.linalg.norm(z_refined, axis=-1, keepdims=True) + 1e-8 - return z_refined / norm - - def save(self, path: str): - np.savez(path, **self.params) - - def load(self, path: str): - data = np.load(path) - for key in self.params: - if key in data: - self.params[key] = data[key] - - -def margin_ranking_loss(score_pos: np.ndarray, score_neg: np.ndarray, margin: float = 0.5) -> Tuple[float, np.ndarray, np.ndarray]: - """ - Pairwise margin ranking loss. - - Loss = max(0, margin - (score_pos - score_neg)) - - Returns: (loss_value, grad_pos, grad_neg) - """ - diff = score_pos - score_neg - violations = margin - diff - loss = np.maximum(0, violations) - - # Gradient: -1 for pos, +1 for neg when violating - mask = (violations > 0).astype(np.float32) - grad_pos = -mask - grad_neg = mask - - return loss.mean(), grad_pos, grad_neg - - -def deep_supervision_loss( - scores_per_iter: List[np.ndarray], - labels: np.ndarray, - weight: float = 0.5 -) -> Tuple[float, List[np.ndarray]]: - """ - Deep supervision: compute loss at each iteration. - - From TRM paper: train model to improve answer at each step. - Later iterations get higher weight. - - scores_per_iter: List of (batch,) score arrays, one per iteration - labels: (batch,) ground truth relevance - weight: how much to weight intermediate losses vs final - - Returns: (total_loss, list of gradients per iteration) - """ - n_iters = len(scores_per_iter) - total_loss = 0.0 - grads = [] - - for i, scores in enumerate(scores_per_iter): - # Later iterations get higher weight - iter_weight = (i + 1) / n_iters - if i < n_iters - 1: - iter_weight *= weight - - # MSE loss for simplicity - diff = scores - labels - loss = (diff ** 2).mean() - grad = 2 * diff / len(diff) * iter_weight - - total_loss += loss * iter_weight - grads.append(grad) - - return total_loss, grads - - -class RecursiveRerankerTrainer: - """ - Trainer for the recursive reranker with deep supervision. - """ - - def __init__(self, config: TrainingConfig): - self.config = config - - # Initialize models - self.scorer = TrainableTinyScorer(dim=config.dim, hidden_dim=config.hidden_dim) - self.refiner = TrainableLatentRefiner(dim=config.dim) - - # Optimizers - self.scorer_opt = SimpleGradientDescent( - lr=config.learning_rate, weight_decay=config.weight_decay - ) - self.refiner_opt = SimpleGradientDescent( - lr=config.learning_rate, weight_decay=config.weight_decay - ) - - # Training state - self.epoch = 0 - self.losses: List[float] = [] - - def _encode_text(self, texts: List[str]) -> np.ndarray: - """Simple bag-of-chars encoding (placeholder for real embeddings).""" - result = [] - for text in texts: - # Hash-based pseudo-embedding - np.random.seed(hash(text) % (2**32)) - vec = np.random.randn(self.config.dim).astype(np.float32) - vec = vec / (np.linalg.norm(vec) + 1e-8) - result.append(vec) - return np.array(result, dtype=np.float32) - - def train_step(self, examples: List[TrainingExample]) -> float: - """Single training step on a batch of examples.""" - batch_size = len(examples) - - # Encode texts - queries = self._encode_text([ex.query for ex in examples]) - docs_pos = self._encode_text([ex.doc_positive for ex in examples]) - docs_neg = self._encode_text([ex.doc_negative for ex in examples]) - - # Initialize latent states - z_pos = queries.copy() - z_neg = queries.copy() - - # Collect scores per iteration for deep supervision - scores_pos_per_iter = [] - scores_neg_per_iter = [] - caches_pos = [] - caches_neg = [] - - # Forward pass through all iterations - for i in range(self.config.n_iterations): - # Score positive docs - x_pos = np.concatenate([queries, docs_pos, z_pos], axis=1) - s_pos, cache_pos = self.scorer.forward(x_pos) - scores_pos_per_iter.append(s_pos) - caches_pos.append(cache_pos) - - # Score negative docs - x_neg = np.concatenate([queries, docs_neg, z_neg], axis=1) - s_neg, cache_neg = self.scorer.forward(x_neg) - scores_neg_per_iter.append(s_neg) - caches_neg.append(cache_neg) - - # Refine latent states - if i < self.config.n_iterations - 1: - z_pos = self.refiner.forward(z_pos, queries, docs_pos) - z_neg = self.refiner.forward(z_neg, queries, docs_neg) - - # Compute loss with deep supervision - total_loss = 0.0 - scorer_grads_accumulated = {k: np.zeros_like(v) for k, v in self.scorer.params.items()} - - for i in range(self.config.n_iterations): - # Pairwise margin loss at each iteration - iter_weight = (i + 1) / self.config.n_iterations - if i < self.config.n_iterations - 1: - iter_weight *= self.config.deep_supervision_weight - - loss, grad_pos, grad_neg = margin_ranking_loss( - scores_pos_per_iter[i], - scores_neg_per_iter[i], - self.config.margin - ) - total_loss += loss * iter_weight - - # Backprop through scorer - grads_pos = self.scorer.backward(grad_pos * iter_weight, caches_pos[i]) - grads_neg = self.scorer.backward(grad_neg * iter_weight, caches_neg[i]) - - for k in scorer_grads_accumulated: - scorer_grads_accumulated[k] += grads_pos[k] + grads_neg[k] - - # Update scorer parameters - self.scorer.params = self.scorer_opt.step(self.scorer.params, scorer_grads_accumulated) - - return total_loss - - def train(self, examples: List[TrainingExample], val_examples: Optional[List[TrainingExample]] = None): - """Full training loop.""" - n_batches = (len(examples) + self.config.batch_size - 1) // self.config.batch_size - - for epoch in range(self.config.epochs): - self.epoch = epoch - epoch_loss = 0.0 - - # Shuffle examples - indices = np.random.permutation(len(examples)) - - for batch_idx in range(n_batches): - start = batch_idx * self.config.batch_size - end = min(start + self.config.batch_size, len(examples)) - batch_indices = indices[start:end] - batch = [examples[i] for i in batch_indices] - - loss = self.train_step(batch) - epoch_loss += loss - - epoch_loss /= n_batches - self.losses.append(epoch_loss) - - # Logging - if epoch % 10 == 0: - print(f"Epoch {epoch}: loss={epoch_loss:.4f}") - - # Save checkpoint - if (epoch + 1) % self.config.save_every == 0: - self.save_checkpoint(f"checkpoint_epoch_{epoch+1}") - - def save_checkpoint(self, name: str): - """Save model checkpoint.""" - save_dir = Path("models/rerank_recursive") - save_dir.mkdir(parents=True, exist_ok=True) - - self.scorer.save(str(save_dir / f"{name}_scorer.npz")) - self.refiner.save(str(save_dir / f"{name}_refiner.npz")) - - # Save training state - state = { - "epoch": self.epoch, - "losses": self.losses, - "config": { - "dim": self.config.dim, - "hidden_dim": self.config.hidden_dim, - "n_iterations": self.config.n_iterations, - } - } - with open(save_dir / f"{name}_state.json", "w") as f: - json.dump(state, f) - - print(f"Saved checkpoint: {name}") - - def load_checkpoint(self, name: str): - """Load model checkpoint.""" - save_dir = Path("models/rerank_recursive") - self.scorer.load(str(save_dir / f"{name}_scorer.npz")) - self.refiner.load(str(save_dir / f"{name}_refiner.npz")) - - with open(save_dir / f"{name}_state.json") as f: - state = json.load(f) - self.epoch = state["epoch"] - self.losses = state["losses"] - print(f"Loaded checkpoint: {name}") - - - -def generate_synthetic_data(n_examples: int = 1000, output_path: str = "data/rerank_train.jsonl"): - """ - Generate synthetic training data for the reranker. - - Creates pairwise examples where: - - Positive doc contains query terms - - Negative doc is unrelated - """ - # Code-like queries and documents - queries = [ - "hybrid search implementation", - "recursive reranker training", - "embedding model initialization", - "cache manager eviction policy", - "file watcher debounce", - "MCP server tool registration", - "vector similarity search", - "document indexing pipeline", - "query expansion techniques", - "relevance scoring function", - ] - - positive_templates = [ - "def {keyword}(query): # Implements {keyword} for search", - "class {Keyword}Manager: # Handles {keyword} operations", - "async def run_{keyword}(): # Execute {keyword} pipeline", - "# {keyword} configuration and setup\nconfig = load_{keyword}_config()", - "def test_{keyword}(): # Unit tests for {keyword}", - ] - - negative_templates = [ - "def unrelated_function(): return 42", - "class DatabaseConnection: # Connect to database", - "import os, sys, json # Standard imports", - "# Configuration file for logging\nLOG_LEVEL = 'INFO'", - "def helper_util(): pass # Utility function", - ] - - examples = [] - - for i in range(n_examples): - # Pick a query - query = queries[i % len(queries)] - keyword = query.split()[0].lower() - Keyword = keyword.capitalize() - - # Generate positive doc (contains query terms) - pos_template = positive_templates[i % len(positive_templates)] - pos_doc = pos_template.format(keyword=keyword, Keyword=Keyword) - - # Generate negative doc (unrelated) - neg_doc = negative_templates[i % len(negative_templates)] - - example = { - "query": query, - "doc_positive": pos_doc, - "doc_negative": neg_doc, - "score_positive": 1.0, - "score_negative": 0.0, - } - examples.append(example) - - # Save to file - output_path = Path(output_path) - output_path.parent.mkdir(parents=True, exist_ok=True) - - with open(output_path, "w") as f: - for ex in examples: - f.write(json.dumps(ex) + "\n") - - print(f"Generated {n_examples} training examples -> {output_path}") - return examples - - -def load_training_data(path: str) -> List[TrainingExample]: - """Load training examples from JSONL file.""" - examples = [] - with open(path) as f: - for line in f: - data = json.loads(line) - examples.append(TrainingExample( - query=data["query"], - doc_positive=data["doc_positive"], - doc_negative=data["doc_negative"], - score_positive=data.get("score_positive", 1.0), - score_negative=data.get("score_negative", 0.0), - )) - return examples - - -def export_trained_weights(checkpoint_name: str, output_dir: str = "models/rerank_recursive"): - """ - Export trained weights to the rerank_recursive module. - - Copies the .npz files to where RecursiveReranker can load them. - """ - import shutil - - src_dir = Path("models/rerank_recursive") - dst_dir = Path(output_dir) - dst_dir.mkdir(parents=True, exist_ok=True) - - # Copy scorer weights - scorer_src = src_dir / f"{checkpoint_name}_scorer.npz" - scorer_dst = dst_dir / "scorer_weights.npz" - if scorer_src.exists(): - shutil.copy(scorer_src, scorer_dst) - print(f"Exported scorer weights -> {scorer_dst}") - - # Copy refiner weights - refiner_src = src_dir / f"{checkpoint_name}_refiner.npz" - refiner_dst = dst_dir / "refiner_weights.npz" - if refiner_src.exists(): - shutil.copy(refiner_src, refiner_dst) - print(f"Exported refiner weights -> {refiner_dst}") - - -def main(): - parser = argparse.ArgumentParser(description="Train Tiny Recursive Reranker") - parser.add_argument("--generate-data", action="store_true", help="Generate synthetic training data") - parser.add_argument("--train", action="store_true", help="Train the model") - parser.add_argument("--evaluate", action="store_true", help="Evaluate the model") - parser.add_argument("--export", type=str, help="Export checkpoint weights") - parser.add_argument("--data", type=str, default="data/rerank_train.jsonl", help="Training data path") - parser.add_argument("--output", type=str, default="data/rerank_train.jsonl", help="Output path for generated data") - parser.add_argument("--n-examples", type=int, default=1000, help="Number of examples to generate") - parser.add_argument("--epochs", type=int, default=100, help="Training epochs") - parser.add_argument("--lr", type=float, default=0.001, help="Learning rate") - parser.add_argument("--batch-size", type=int, default=32, help="Batch size") - parser.add_argument("--checkpoint", type=str, help="Checkpoint to load/export") - - args = parser.parse_args() - - if args.generate_data: - generate_synthetic_data(n_examples=args.n_examples, output_path=args.output) - - elif args.train: - # Load data - if not Path(args.data).exists(): - print(f"Training data not found: {args.data}") - print("Run with --generate-data first") - return - - examples = load_training_data(args.data) - print(f"Loaded {len(examples)} training examples") - - # Create trainer - config = TrainingConfig( - epochs=args.epochs, - learning_rate=args.lr, - batch_size=args.batch_size, - ) - trainer = RecursiveRerankerTrainer(config) - - # Load checkpoint if specified - if args.checkpoint: - try: - trainer.load_checkpoint(args.checkpoint) - except Exception as e: - print(f"Could not load checkpoint: {e}") - - # Train - trainer.train(examples) - trainer.save_checkpoint("final") - print("Training complete!") - - elif args.export: - export_trained_weights(args.export or "final") - - elif args.evaluate: - if not Path(args.data).exists(): - print(f"Evaluation data not found: {args.data}") - return - - examples = load_training_data(args.data) - - # Load trained model - config = TrainingConfig() - trainer = RecursiveRerankerTrainer(config) - - try: - trainer.load_checkpoint(args.checkpoint or "final") - except Exception as e: - print(f"Could not load checkpoint: {e}") - print("Training with random weights...") - - # Evaluate: compute accuracy (positive > negative) - correct = 0 - total = 0 - - for ex in examples[:100]: # Sample - q = trainer._encode_text([ex.query])[0] - d_pos = trainer._encode_text([ex.doc_positive])[0] - d_neg = trainer._encode_text([ex.doc_negative])[0] - - z = q.copy() - - # Run through iterations - for _ in range(config.n_iterations): - x_pos = np.concatenate([q, d_pos, z]) - x_neg = np.concatenate([q, d_neg, z]) - - s_pos, _ = trainer.scorer.forward(x_pos.reshape(1, -1)) - s_neg, _ = trainer.scorer.forward(x_neg.reshape(1, -1)) - - z = trainer.refiner.forward(z.reshape(1, -1), q.reshape(1, -1), d_pos.reshape(1, -1))[0] - - if s_pos[0] > s_neg[0]: - correct += 1 - total += 1 - - accuracy = correct / total if total > 0 else 0 - print(f"Pairwise accuracy: {accuracy:.2%} ({correct}/{total})") - - else: - parser.print_help() - - -if __name__ == "__main__": - main() - diff --git a/scripts/rerank_train.py b/scripts/rerank_train.py deleted file mode 100644 index a9ef2079..00000000 --- a/scripts/rerank_train.py +++ /dev/null @@ -1,10 +0,0 @@ -#!/usr/bin/env python3 -"""Backward-compatibility shim. See scripts/rerank_tools/train.py""" -import sys -from pathlib import Path -sys.path.insert(0, str(Path(__file__).parent.parent)) - -from scripts.rerank_tools.train import * - -if __name__ == "__main__": - main() diff --git a/scripts/router_eval.py b/scripts/router_eval.py deleted file mode 100644 index d429a08c..00000000 --- a/scripts/router_eval.py +++ /dev/null @@ -1,385 +0,0 @@ -import argparse -import json, os, threading, time, sys, re, copy -from http.server import HTTPServer, BaseHTTPRequestHandler -from typing import Dict, Any, List, Tuple - -# Simple Mock MCP server for evals -class MockMCPHandler(BaseHTTPRequestHandler): - server_version = "MockMCP/0.1" - - def _send_json(self, obj: Dict[str, Any], session: str | None = None, code: int = 200): - body = json.dumps(obj).encode("utf-8") - self.send_response(code) - self.send_header("Content-Type", "application/json") - if session: - self.send_header("Mcp-Session-Id", session) - self.send_header("Content-Length", str(len(body))) - self.end_headers() - self.wfile.write(body) - - def do_POST(self): # noqa: N802 - raw = self.rfile.read(int(self.headers.get("Content-Length", "0") or 0)) - try: - j = json.loads(raw.decode("utf-8", errors="ignore")) - except Exception: - return self._send_json({"jsonrpc": "2.0", "error": {"message": "bad json"}}, code=400) - method = j.get("method") - if method == "initialize": - # Return session via header; some clients also parse body - return self._send_json({"jsonrpc": "2.0", "id": j.get("id"), "result": {"ok": True, "server": self.server.server_name}}, session="mock-session") - if method == "notifications/initialized": - return self._send_json({"jsonrpc": "2.0", "result": {"ok": True}}) - if method == "tools/list": - # Simulate flakiness once if flagged - if getattr(self.server, "fail_list_once", False) and not getattr(self.server, "_fail_list_consumed", False): - setattr(self.server, "_fail_list_consumed", True) - return self._send_json({"jsonrpc": "2.0", "error": {"message": "flaky list"}}, code=500) - tools = getattr(self.server, "tools", []) - return self._send_json({"jsonrpc": "2.0", "id": j.get("id"), "result": {"tools": tools}}) - if method == "tools/call": - params = j.get("params") or {} - name = (params.get("name") or "").strip() - args = params.get("arguments") or {} - # Indexer tools - if name in {"repo_search", "search_config_for", "search_tests_for", "search_callers_for", "search_importers_for"}: - total = int(getattr(self.server, "search_total", 5)) - # Cap returned items to avoid huge payloads; still report full total - shown = max(0, min(total, 3)) - results = [ - {"score": 0.9 - (i * 0.1), "path": f"/work/README_{i}.md", "start_line": 1, "end_line": 2, "snippet": "demo"} - for i in range(shown) - ] - res = { - "result": { - "args": { - "queries": [str(args.get("query") or "")], - "limit": int(args.get("limit") or 8), - "include_snippet": bool(args.get("include_snippet") or False), - "language": str(args.get("language") or ""), - "under": str(args.get("under") or ""), - "symbol": str(args.get("symbol") or ""), - "ext": str(args.get("ext") or ""), - "compact": False, - }, - "total": total, - "results": results, - "ok": True, - "code": 0, - "stdout": "", - "stderr": "", - } - } - return self._send_json({"jsonrpc": "2.0", "id": j.get("id"), "result": {"content": [{"type": "text", "text": json.dumps(res)}], "structuredContent": res, "isError": False}}) - if name == "context_answer_compat": - # Simulate failure if flagged so router should fall back to context_answer - if getattr(self.server, "fail_context_compat", False): - return self._send_json({"jsonrpc": "2.0", "id": j.get("id"), "result": {"content": [{"type": "text", "text": json.dumps({"error": "compat failed"})}], "structuredContent": {"error": "compat failed"}, "isError": True}}) - # Require nested arguments wrapper - if not isinstance(args, dict) or "arguments" not in args: - return self._send_json({"jsonrpc": "2.0", "id": j.get("id"), "result": {"content": [{"type": "text", "text": json.dumps({"error": "compat requires nested arguments"})}], "structuredContent": {"error": "compat requires nested arguments"}, "isError": True}}) - inner = args.get("arguments") or {} - q = str(inner.get("query") or "") - ans = { - "answer": "Short ok." if len(q) < 80 else "Longer answer", - "citations": [{"id": 1, "path": "/work/file.py", "start_line": 1, "end_line": 2}], - "query": [q], - "used": {"gate_first": True, "refrag": True}, - } - res = {"result": ans} - return self._send_json({"jsonrpc": "2.0", "id": j.get("id"), "result": {"content": [{"type": "text", "text": json.dumps(res)}], "structuredContent": res, "isError": False}}) - if name == "context_answer": - q = str(args.get("query") or "") - ans = {"answer": "Ok.", "citations": []} - res = {"result": ans} - return self._send_json({"jsonrpc": "2.0", "id": j.get("id"), "result": {"content": [{"type": "text", "text": json.dumps(res)}], "structuredContent": res, "isError": False}}) - if name in {"qdrant_status", "qdrant_list"}: - res = {"result": {"ok": True}} - return self._send_json({"jsonrpc": "2.0", "id": j.get("id"), "result": {"content": [{"type": "text", "text": json.dumps(res)}], "structuredContent": res, "isError": False}}) - # Memory tools - if name == "store": - res = {"result": {"ok": True}} - return self._send_json({"jsonrpc": "2.0", "id": j.get("id"), "result": {"content": [{"type": "text", "text": json.dumps(res)}], "structuredContent": res, "isError": False}}) - if name == "find": - q = str(args.get("query") or "") - res = {"result": {"ok": True, "results": [{"information": "The MCP indexer uses hybrid search combining dense embeddings and lexical matching with optional reranking", "metadata": {"category": "architecture"}}], "count": 1}} - return self._send_json({"jsonrpc": "2.0", "id": j.get("id"), "result": {"content": [{"type": "text", "text": json.dumps(res)}], "structuredContent": res, "isError": False}}) - return self._send_json({"jsonrpc": "2.0", "id": j.get("id"), "result": {"content": [{"type": "text", "text": json.dumps({"error": f"unknown tool {name}"})}], "structuredContent": {"error": f"unknown tool {name}"}, "isError": True}}) - return self._send_json({"jsonrpc": "2.0", "error": {"message": f"unknown method {method}"}}, code=400) - - -def start_mock_server(port: int, tools: List[Dict[str, Any]]) -> Tuple[HTTPServer, threading.Thread]: - httpd = HTTPServer(("localhost", port), MockMCPHandler) - httpd.tools = tools # type: ignore - t = threading.Thread(target=httpd.serve_forever, daemon=True) - t.start() - # Warmup - time.sleep(0.05) - return httpd, t - - -def tool(name: str, description: str, params: List[str] = None) -> Dict[str, Any]: - schema = {"type": "object", "properties": {p: {"type": "string"} for p in (params or [])}} - return {"name": name, "description": description, "inputSchema": schema} - - -def run_eval_suite(verbose: bool = False) -> int: - # Two mock servers: indexer and memory - indexer_tools = [ - tool("repo_search", "General code search", ["query", "limit", "include_snippet", "language", "under", "symbol", "ext"]), - tool("search_config_for", "Intent-specific search for configuration files", ["query", "limit", "include_snippet"]), - tool("search_importers_for", "Find files importing a module or symbol", ["query", "limit", "language", "under"]), - tool("context_answer_compat", "Answer a question using code context (compat)", ["arguments"]), - tool("context_answer", "Answer a question using code context", ["query", "limit"]), - tool("qdrant_status", "Qdrant status"), - tool("qdrant_list", "Qdrant list"), - ] - memory_tools = [tool("store", "Store memory", ["information"]), tool("find", "Find memory", ["query", "limit"])] - idx, _ = start_mock_server(18031, indexer_tools) - mem, _ = start_mock_server(18032, memory_tools) - - try: - os.environ["MCP_INDEXER_HTTP_URL"] = "http://localhost:18031/mcp" - os.environ["MCP_MEMORY_HTTP_URL"] = "http://localhost:18032/mcp" - os.environ["ROUTER_SEARCH_LIMIT"] = "8" - os.environ["ROUTER_INCLUDE_SNIPPET"] = "1" - - # Import router after env set so its defaults bind to mock URLs - import importlib.util as _ilu - _p = os.path.join(os.path.dirname(__file__), "mcp_router.py") - _spec = _ilu.spec_from_file_location("mcp_router", _p) - router = _ilu.module_from_spec(_spec) - assert _spec and _spec.loader - _spec.loader.exec_module(router) # type: ignore - - failures = [] - intent_logs: List[Dict[str, Any]] = [] - - def run_plan(q: str) -> List[Tuple[str, Dict[str, Any]]]: - plan = router.build_plan(q) - debug = getattr(router, "_LAST_INTENT_DEBUG", {}) - if isinstance(debug, dict): - log_entry = copy.deepcopy(debug) - else: - log_entry = {"debug": debug} - log_entry["query"] = q - log_entry["plan_first_tool"] = plan[0][0] if plan else None - intent_logs.append(log_entry) - return plan - - # 1) Signature selection: prefer search_config_for for config changes - p1 = run_plan("compare callers to config changes") - if not p1 or p1[0][0] != "search_config_for": - failures.append("signature selection: expected search_config_for") - - # 2) Repo hints: language+under parsed - p2 = run_plan("who imports foo in python under src/lib") - if not p2 or p2[0][0] != "search_importers_for": - failures.append("repo hints: expected search_importers_for") - else: - args2 = p2[0][1] - if args2.get("language") != "python": - failures.append("repo hints: language not parsed") - if args2.get("under") != "src/lib": - failures.append("repo hints: under not parsed") - - # 3) Design recap: memory find precedes answer - p3 = run_plan("recap our architecture decisions for the indexer") - expect_order = ["find", "context_answer_compat"] - if not p3 or [p3[0][0], p3[1][0]] != expect_order: - failures.append("design recap plan: expected find -> context_answer_compat") - - # 4) Multi-intent: store + reindex - p4 = run_plan("remember this: prefer concise answers; then reindex fresh") - if not p4 or [p4[0][0], p4[1][0]] != ["store", "qdrant_index_root"]: - failures.append("multi-intent: expected store then index") - else: - store_args = p4[0][1] or {} - info4 = (store_args.get("information") or "").lower() - if "remember" in info4: - failures.append("multi-intent: trigger phrase leaked into stored information") - if "reindex" in info4: - failures.append("multi-intent: reindex fragment leaked into stored information") - if not p4[1][1].get("recreate"): - failures.append("multi-intent: expected recreate true") - - # 5) Memory metadata extraction - p_meta = run_plan("remember this [priority=high tags=ux,frontend]: update the signup banner copy") - if not p_meta or p_meta[0][0] != "store": - failures.append("memory metadata: expected store intent") - else: - store_args = p_meta[0][1] or {} - if store_args.get("information") != "update the signup banner copy": - failures.append("memory metadata: information not cleaned") - md = store_args.get("metadata") or {} - if md.get("priority") != "high": - failures.append("memory metadata: priority missing") - if md.get("tags") != ["ux", "frontend"]: - failures.append("memory metadata: tags mismatch") - - # 6) Glob/exclude filters - p5 = run_plan("search only *.py files exclude vendor") - if not p5: - failures.append("glob: plan empty") - else: - args5 = p5[0][1] - gl = (args5 or {}).get("path_glob") or [] - ng = (args5 or {}).get("not_glob") or [] - if "**/*.py" not in gl: - failures.append("glob: missing **/*.py") - if "**/vendor/**" not in ng: - failures.append("glob: missing exclude vendor") - - # 7) Run end-to-end for recap and ensure compat accepted and short answers not rejected - # Capture stdout of router.main - def run_router(args: List[str]) -> str: - from io import StringIO - old = sys.stdout - try: - buf = StringIO() - sys.stdout = buf - router.main(args) - return buf.getvalue() - finally: - sys.stdout = old - out = run_router(["--run", "recap our architecture decisions for the indexer"]) - def run_router_code(args: List[str]) -> int: - from io import StringIO - old = sys.stdout - try: - sys.stdout = StringIO() # suppress stdout capture to avoid noise - return int(router.main(args)) - finally: - sys.stdout = old - - if "compat requires nested arguments" in out: - failures.append("compat: still sending flattened args") - if "Memory context:" not in out: - failures.append("memory→answer: query was not augmented with memory context") - print("--- router stdout ---\n" + out + "\n--- end stdout ---") - - - # 6b) Repeat immediately after recap should skip fresh memory.find step - out_repeat = run_router(["--run", "repeat that"]) - if '"skipped": "scratchpad_fresh"' not in out_repeat: - failures.append("repeat: find step not skipped on fresh cache") - - # 7) Repeat last: persist then repeat - _ = run_router(["--run", "who imports foo in python under src/lib"]) - p7a = run_plan("who imports foo in python under src/lib") - p7b = run_plan("repeat that") - if p7a != p7b: - failures.append("repeat: last plan not reused") - - # 7b) "same filters" carry-over in planning - p7c = run_plan("search with same filters for bar baz") - if not p7c: - failures.append("same filters: plan empty") - else: - args7c = p7c[0][1] - if (args7c or {}).get("language") != "python" or (args7c or {}).get("under") != "src/lib": - failures.append("same filters: did not reuse prior language/under") - - - - # 8) Fallback on compat failure - setattr(idx, "fail_context_compat", True) - out2 = run_router(["--run", "recap our architecture decisions for the indexer"]) - if '"tool": "context_answer"' not in out2: - failures.append("fallback: did not call context_answer after compat failure") - setattr(idx, "fail_context_compat", False) - - # 9) tools/list flakiness toleration - setattr(idx, "fail_list_once", True) - p9 = run_plan("find config changes") - if not p9: - failures.append("discovery flakiness: plan empty after retry") - setattr(idx, "fail_list_once", False) - - - # 10) Expand on last summary uses prior summary and citations (fresh) - out3 = run_router(["--run", "expand on that summary"]) - if "Prior summary:" not in out3 or "/work/file.py" not in out3: - failures.append("expand: prior summary/citations not injected when fresh") - - # 11) TTL expiry should suppress prior summary injection - os.environ["ROUTER_SCRATCHPAD_TTL_SEC"] = "0" - out4 = run_router(["--run", "expand on that summary"]) - if "Prior summary:" in out4: - failures.append("ttl: prior summary injected despite stale cache") - os.environ.pop("ROUTER_SCRATCHPAD_TTL_SEC", None) - - # 13) Divergence fatal per-tool: repo_search set to fatal should cause nonzero exit - os.environ["ROUTER_DIVERGENCE_FATAL_TOOLS"] = "repo_search" - setattr(idx, "search_total", 6) - _ = run_router(["--run", "search for demo"]) - setattr(idx, "search_total", 2) - code_div = run_router_code(["--run", "search for demo"]) - if code_div == 0: - failures.append("divergence fatal: router returned success despite fatal policy") - os.environ.pop("ROUTER_DIVERGENCE_FATAL_TOOLS", None) - setattr(idx, "search_total", 5) - - # 12) Divergence detection: baseline high → lower later should print a divergence notice - setattr(idx, "search_total", 6) - _ = run_router(["--run", "search for demo"]) - setattr(idx, "search_total", 2) - out_div = run_router(["--run", "search for demo"]) - if '"divergence"' not in out_div: - failures.append("divergence: no divergence flagged on material drop") - - fallback_logs = [] - for log in intent_logs: - if log.get("strategy") == "ml": - if "confidence" not in log: - failures.append(f"intent log missing confidence for query: {log.get('query')}") - if ( - log.get("intent") == router.INTENT_SEARCH - and log.get("top_candidate") - and log.get("top_candidate") != router.INTENT_SEARCH - ): - if log.get("confidence", 0.0) >= log.get("threshold", 0.25): - failures.append(f"intent fallback without low confidence for query: {log.get('query')}") - else: - fallback_logs.append(log) - if fallback_logs: - print("Intent fallback diagnostics:") - for item in fallback_logs: - try: - score = float(item.get("confidence") or 0.0) - except Exception: - score = 0.0 - print( - f" query={item.get('query')!r} top={item.get('top_candidate')} " - f"score={score:.3f} -> intent={item.get('intent')} first_tool={item.get('plan_first_tool')}" - ) - if verbose: - print("Intent diagnostics (all):") - for item in intent_logs: - try: - score = float(item.get("confidence") or 0.0) - except Exception: - score = 0.0 - print( - f" query={item.get('query')!r} strategy={item.get('strategy')} " - f"intent={item.get('intent')} score={score:.3f} " - f"top={item.get('top_candidate')} first_tool={item.get('plan_first_tool')}" - ) - - if failures: - print("Router eval: FAIL\n- " + "\n- ".join(failures)) - return 1 - print("Router eval: PASS (all checks)") - return 0 - finally: - idx.shutdown(); mem.shutdown() - - -if __name__ == "__main__": - parser = argparse.ArgumentParser(description="Run router evaluation suite.") - parser.add_argument( - "--verbose", - action="store_true", - help="Print intent confidence diagnostics after the suite completes.", - ) - args = parser.parse_args() - raise SystemExit(run_eval_suite(verbose=args.verbose)) diff --git a/scripts/run_init_maintenance.py b/scripts/run_init_maintenance.py new file mode 100644 index 00000000..b629b267 --- /dev/null +++ b/scripts/run_init_maintenance.py @@ -0,0 +1,14 @@ +#!/usr/bin/env python3 +"""Run the init maintenance script sequence under the shared watcher lock.""" + +from __future__ import annotations + +from scripts.watch_index_core.init_maintenance import run_init_maintenance_once + + +def main() -> int: + return 0 if run_init_maintenance_once() else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/semantic_expansion.py b/scripts/semantic_expansion.py index 76e472c1..2dd999fd 100644 --- a/scripts/semantic_expansion.py +++ b/scripts/semantic_expansion.py @@ -1,4 +1,6 @@ #!/usr/bin/env python3 +from __future__ import annotations + """ Semantic similarity-based query expansion for Context-Engine. @@ -9,44 +11,34 @@ import os import math import re -from typing import List, Dict, Any, Tuple, Optional, Set +from typing import List, Dict, Any, Tuple, Optional, Set, TYPE_CHECKING from collections import defaultdict import logging logger = logging.getLogger("semantic_expansion") -# Import embedding functionality (prefer embedder factory for Qwen3 support) -try: - from scripts.embedder import get_embedding_model as _get_embedding_model - _EMBEDDER_FACTORY = True - FASTEMBED_AVAILABLE = True -except ImportError: - _EMBEDDER_FACTORY = False - try: - from fastembed import TextEmbedding - FASTEMBED_AVAILABLE = True - except ImportError: - FASTEMBED_AVAILABLE = False - TextEmbedding = None - -# Import Qdrant client for vector operations -try: - from qdrant_client import QdrantClient, models - QDRANT_AVAILABLE = True -except ImportError: - QDRANT_AVAILABLE = False - QdrantClient = None - models = None - -# Import local utilities -try: - from scripts.utils import ( - lex_hash_vector_queries as _lex_hash_vector_queries, - sanitize_vector_name as _sanitize_vector_name, - ) -except ImportError: - _lex_hash_vector_queries = None - _sanitize_vector_name = None +if TYPE_CHECKING: + from qdrant_client import QdrantClient, models as models +else: + QdrantClient = Any + + class _LazyQdrantModels: + def __getattr__(self, name: str) -> Any: + from qdrant_client import models as _models + + return getattr(_models, name) + + models = _LazyQdrantModels() + +from scripts.embedder import get_embedding_model as _get_embedding_model +from scripts.utils import ( + lex_hash_vector_queries as _lex_hash_vector_queries, + sanitize_vector_name as _sanitize_vector_name, +) + +_EMBEDDER_FACTORY = True +FASTEMBED_AVAILABLE = True +QDRANT_AVAILABLE = True # Configuration defaults # NOTE: SEMANTIC_EXPANSION_ENABLED is intentionally *not* a module-level constant. @@ -62,18 +54,15 @@ def _semantic_expansion_enabled() -> bool: SEMANTIC_EXPANSION_CACHE_TTL = float(os.environ.get("SEMANTIC_EXPANSION_CACHE_TTL", "3600") or "3600") # Use UnifiedCache for proper LRU eviction instead of simple FIFO -try: - from scripts.cache_manager import UnifiedCache, EvictionPolicy - _expansion_cache = UnifiedCache( - name="semantic_expansion", - max_size=SEMANTIC_EXPANSION_CACHE_SIZE, - eviction_policy=EvictionPolicy.LRU, - default_ttl=SEMANTIC_EXPANSION_CACHE_TTL, - ) - _UNIFIED_CACHE = True -except ImportError: - _expansion_cache: Dict[str, List[str]] = {} # type: ignore - _UNIFIED_CACHE = False +from scripts.cache_manager import UnifiedCache, EvictionPolicy + +_expansion_cache = UnifiedCache( + name="semantic_expansion", + max_size=SEMANTIC_EXPANSION_CACHE_SIZE, + eviction_policy=EvictionPolicy.LRU, + default_ttl=SEMANTIC_EXPANSION_CACHE_TTL, +) +_UNIFIED_CACHE = True _cache_hits = 0 _cache_misses = 0 diff --git a/scripts/smoke_test.py b/scripts/smoke_test.py index 2a699278..827dcda5 100644 --- a/scripts/smoke_test.py +++ b/scripts/smoke_test.py @@ -1,21 +1,9 @@ #!/usr/bin/env python3 import os import sys -from pathlib import Path from qdrant_client import QdrantClient -# Ensure scripts is importable -ROOT_DIR = Path(__file__).resolve().parent.parent -if str(ROOT_DIR) not in sys.path: - sys.path.insert(0, str(ROOT_DIR)) - -# Use embedder factory for Qwen3 support -try: - from scripts.embedder import get_embedding_model - _EMBEDDER_FACTORY = True -except ImportError: - _EMBEDDER_FACTORY = False - from fastembed import TextEmbedding +from scripts.embedder import get_embedding_model from scripts.utils import sanitize_vector_name QDRANT_URL = os.environ.get("QDRANT_URL", "http://qdrant:6333") @@ -37,10 +25,7 @@ count = None # Prepare query embedding -if _EMBEDDER_FACTORY: - model = get_embedding_model(MODEL) -else: - model = TextEmbedding(model_name=MODEL) +model = get_embedding_model(MODEL) query = "python code indexer for qdrant" vec = next(model.embed([query])) diff --git a/scripts/standalone_upload_client.py b/scripts/standalone_upload_client.py index 7cbd9dd1..7d7f5bbc 100644 --- a/scripts/standalone_upload_client.py +++ b/scripts/standalone_upload_client.py @@ -39,15 +39,88 @@ except ImportError: WATCHDOG_AVAILABLE = False -try: - from upload_auth_utils import get_auth_session # type: ignore[import] -except ImportError: - def get_auth_session(upload_endpoint: str) -> str: +def get_auth_session(upload_endpoint: str) -> str: + """Resolve an optional upload session without importing project modules.""" + try: + session_id = ( + os.environ.get("CTXCE_UPLOAD_SESSION_ID") + or os.environ.get("CTXCE_SESSION_ID") + or "" + ).strip() + if session_id: + return session_id + + auth_path = Path(os.path.expanduser("~")) / ".ctxce" / "auth.json" + if not auth_path.exists(): + return "" + with auth_path.open("r", encoding="utf-8") as handle: + data = json.load(handle) + if not isinstance(data, dict): + return "" + entry = data.get(upload_endpoint.rstrip("/")) or data.get(upload_endpoint) + if not isinstance(entry, dict): + return "" + session_id = str(entry.get("sessionId") or entry.get("session_id") or "").strip() + expires_at = entry.get("expiresAt") or entry.get("expires_at") + if isinstance(expires_at, (int, float)) and expires_at > 0: + if expires_at < int(time.time()): + return "" + return session_id + except Exception: return "" # Configure logging logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) +_git_history_skip_log_key: Optional[str] = None + + +def _is_usable_delta_status(status: Any) -> bool: + if not isinstance(status, dict): + return False + state = str(status.get("status") or "").strip().lower() + return ( + bool(status.get("success")) and + "workspace_path" in status and + "collection_name" in status and + state in {"ready", "processing", "completed"} + ) + + +def _server_status_error_message(status: Any) -> str: + if isinstance(status, dict): + error = status.get("error") + if isinstance(error, dict): + msg = str(error.get("message") or "").strip() + if msg: + return msg + state = str(status.get("status") or "").strip() + if state: + return f"Server status is {state}" + return "Invalid server status response" + + +def _env_flag(name: str, default: bool) -> bool: + raw = os.environ.get(name) + if raw is None: + return default + return str(raw).strip().lower() in {"1", "true", "yes", "on"} + + +def _format_cached_sha1(value: Optional[str]) -> Optional[str]: + raw = str(value or "").strip() + if not raw: + return None + return raw if raw.lower().startswith("sha1:") else f"sha1:{raw}" + + +def _log_git_history_skip_once(reason: str, key: str) -> None: + global _git_history_skip_log_key + marker = f"{reason}:{key}" + if _git_history_skip_log_key == marker: + return + _git_history_skip_log_key = marker + logger.info("[git_history] skip (%s): %s", reason, key) DEFAULT_MAX_TEMP_CLEAN_ATTEMPTS = 3 DEFAULT_TEMP_CLEAN_SLEEP = 1.0 @@ -176,17 +249,6 @@ def hash_id(text: str, path: str, start: int, end: int) -> str: ).hexdigest() return h[:16] -def get_collection_name(repo_name: Optional[str] = None) -> str: - """Generate collection name with 8-char hash for local workspaces. - - Simplified version from workspace_state.py. - """ - if not repo_name: - return "default-collection" - hash_obj = hashlib.sha256(repo_name.encode()) - short_hash = hash_obj.hexdigest()[:8] - return f"{repo_name}-{short_hash}" - def _extract_repo_name_from_path(workspace_path: str) -> str: """Extract repository name from workspace path. @@ -212,7 +274,6 @@ def __init__(self, workspace_path: str, repo_name: str): # In-memory cache to avoid re-reading and re-validating on every access self._cache_loaded = False self._cache: Dict[str, str] = {} - self._stale_checked = False self._load_cache() # Load once on init def _load_cache(self) -> Dict[str, str]: @@ -229,18 +290,23 @@ def _load_cache(self) -> Dict[str, str]: with open(self.cache_file, 'r', encoding='utf-8') as f: data = json.load(f) file_hashes = data.get("file_hashes", {}) - # Run stale check only once per process to avoid O(N^2) scans - if not self._stale_checked and self._cache_seems_stale(file_hashes): - self._stale_checked = True + if not isinstance(file_hashes, dict): + file_hashes = {} + scoped_hashes = { + str(path): value + for path, value in file_hashes.items() + if self._path_is_within_workspace(str(path)) + } + ignored_count = len(file_hashes) - len(scoped_hashes) + if ignored_count: logger.warning( - "[hash_cache] Detected stale cache with missing paths; resetting %s", - self.cache_file, + "[hash_cache] Ignoring %d entries outside workspace %s", + ignored_count, + self.workspace_path, ) - self._save_cache({}) - self._cache = {} + self._save_cache(scoped_hashes) else: - self._stale_checked = True - self._cache = file_hashes if isinstance(file_hashes, dict) else {} + self._cache = scoped_hashes except Exception: self._cache = {} @@ -290,34 +356,28 @@ def remove_hash(self, file_path: str) -> None: self._cache = file_hashes self._cache_loaded = True - def _cache_seems_stale(self, file_hashes: Dict[str, str]) -> bool: - """Return True if a large portion of cached paths no longer exist on disk.""" - total = len(file_hashes) - if total == 0: + def flush(self) -> None: + """Persist the current in-memory cache state to disk.""" + self._save_cache(dict(self._load_cache())) + + def _path_is_within_workspace(self, path_str: str) -> bool: + try: + path = Path(path_str).resolve() + return path == self.workspace_path or self.workspace_path in path.parents + except Exception: return False - missing = 0 - for path_str in file_hashes.keys(): - try: - if not Path(path_str).exists(): - missing += 1 - except Exception: - missing += 1 - missing_ratio = missing / total - return missing_ratio >= 0.25 -# Create global cache instance (will be initialized in RemoteUploadClient) +# The standalone client is one workspace per process. _hash_cache: Optional[SimpleHashCache] = None def get_cached_file_hash(file_path: str, repo_name: Optional[str] = None) -> str: """Get cached file hash for tracking changes.""" - global _hash_cache if _hash_cache: return _hash_cache.get_hash(file_path) return "" def set_cached_file_hash(file_path: str, file_hash: str, repo_name: Optional[str] = None): """Set cached file hash for tracking changes.""" - global _hash_cache if _hash_cache: _hash_cache.set_hash(file_path, file_hash) @@ -328,7 +388,6 @@ def get_all_cached_paths(repo_name: Optional[str] = None) -> List[str]: The repo_name parameter is accepted for API symmetry with the non-standalone client but is not used here, since this cache is always per-workspace. """ - global _hash_cache if _hash_cache: return _hash_cache.all_paths() return [] @@ -336,11 +395,16 @@ def get_all_cached_paths(repo_name: Optional[str] = None) -> List[str]: def remove_cached_file(file_path: str, repo_name: Optional[str] = None) -> None: """Remove a file entry from the local cache if present.""" - global _hash_cache if _hash_cache: _hash_cache.remove_hash(file_path) +def flush_cached_file_hashes() -> None: + """Persist the current workspace hash cache to disk.""" + if _hash_cache: + _hash_cache.flush() + + def _find_git_root(start: Path) -> Optional[Path]: """Best-effort detection of the git repository root for a workspace. @@ -426,10 +490,12 @@ def _collect_git_history_for_workspace(workspace_path: str) -> Optional[Dict[str } if max_commits <= 0: + _log_git_history_skip_once("disabled", f"max_commits={max_commits}") return None root = _find_git_root(Path(workspace_path)) if not root: + _log_git_history_skip_once("no_repo", workspace_path) return None # Git history cache: avoid emitting identical manifests when HEAD/settings are unchanged @@ -463,6 +529,7 @@ def _collect_git_history_for_workspace(workspace_path: str) -> Optional[Dict[str cache = {} if current_head and cache.get("last_head") == current_head and cache.get("max_commits") == max_commits and str(cache.get("since") or "") == since: + _log_git_history_skip_once("cache_hit", f"head={current_head[:10]} since={since or '-'} max={max_commits}") return None base_head = "" @@ -513,12 +580,20 @@ def _collect_git_history_for_workspace(workspace_path: str) -> Optional[Dict[str errors="replace", ) if proc.returncode != 0 or not proc.stdout.strip(): + _log_git_history_skip_once( + "rev_list_empty", + f"head={current_head[:10] if current_head else '-'} rc={proc.returncode}", + ) return None commits = [l.strip() for l in proc.stdout.splitlines() if l.strip()] except Exception: return None if not commits: + _log_git_history_skip_once( + "no_commits", + f"head={current_head[:10] if current_head else '-'}", + ) return None if len(commits) > max_commits: commits = commits[:max_commits] @@ -592,6 +667,10 @@ def _collect_git_history_for_workspace(workspace_path: str) -> Optional[Dict[str continue if not records: + _log_git_history_skip_once( + "no_records", + f"commits={len(commits)} head={current_head[:10] if current_head else '-'}", + ) return None try: @@ -611,6 +690,14 @@ def _collect_git_history_for_workspace(workspace_path: str) -> Optional[Dict[str "since": since, "commits": records, } + logger.info( + "[git_history] prepared manifest mode=%s commits=%d head=%s prev=%s base=%s", + manifest["mode"], + len(records), + (current_head[:10] if current_head else "-"), + (prev_head[:10] if prev_head else "-"), + (base_head[:10] if base_head else "-"), + ) # Update git history cache with the HEAD and settings used for this manifest try: @@ -663,7 +750,7 @@ def _translate_to_container_path(self, host_path: str) -> str: return host_path.replace('\\', '/').replace(':', '') - def __init__(self, upload_endpoint: str, workspace_path: str, collection_name: str, + def __init__(self, upload_endpoint: str, workspace_path: str, collection_name: Optional[str] = None, max_retries: int = 3, timeout: int = 30, metadata_path: Optional[str] = None, logical_repo_id: Optional[str] = None): """Initialize remote upload client.""" @@ -675,9 +762,6 @@ def __init__(self, upload_endpoint: str, workspace_path: str, collection_name: s self.temp_dir = None self.logical_repo_id = logical_repo_id - # Set environment variables for cache functions - os.environ["WORKSPACE_PATH"] = workspace_path - # Store repo name and initialize hash cache self.repo_name = _extract_repo_name_from_path(workspace_path) # Fallback to directory name if repo detection fails (for non-git repos) @@ -688,6 +772,7 @@ def __init__(self, upload_endpoint: str, workspace_path: str, collection_name: s # In-memory stat cache to avoid rehashing unchanged files on every watch iteration self._stat_cache: Dict[str, Tuple[int, int]] = {} + self._content_hash_cache: Dict[str, Tuple[int, int, str]] = {} # Setup HTTP session with simple retry self.session = requests.Session() @@ -695,6 +780,100 @@ def __init__(self, upload_endpoint: str, workspace_path: str, collection_name: s adapter = HTTPAdapter(max_retries=retry_strategy) self.session.mount("http://", adapter) self.session.mount("https://", adapter) + self.last_upload_result: Dict[str, Any] = {"outcome": "idle"} + self._last_plan_payload: Optional[Dict[str, Any]] = None + self._last_expected_hashes: Dict[str, str] = {} + + def _set_last_upload_result(self, outcome: str, **details: Any) -> Dict[str, Any]: + result: Dict[str, Any] = {"outcome": outcome} + result.update(details) + self.last_upload_result = result + return result + + def _get_all_cached_paths(self) -> List[str]: + return get_all_cached_paths(self.repo_name) + + def _flush_cached_file_hashes(self) -> None: + flush_cached_file_hashes() + + def _read_current_file_hash(self, path: Path) -> Tuple[str, int]: + """Read a file hash once per current size/mtime pair during one pass.""" + stat = path.stat() + mtime_ns = int(getattr(stat, "st_mtime_ns", int(stat.st_mtime * 1e9))) + size = int(stat.st_size) + key = str(path.resolve()) + cached = self._content_hash_cache.get(key) + if cached and cached[:2] == (mtime_ns, size): + return cached[2], size + file_hash = hashlib.sha1(path.read_bytes()).hexdigest() + self._content_hash_cache[key] = (mtime_ns, size, file_hash) + return file_hash, size + + def log_watch_upload_result(self) -> None: + outcome = str((self.last_upload_result or {}).get("outcome") or "") + if outcome == "skipped_by_plan": + logger.info("[watch] No upload needed after plan") + elif outcome == "queued": + logger.info("[watch] Upload request accepted; server processing asynchronously") + elif outcome == "uploaded": + logger.info("[watch] Successfully uploaded changes") + elif outcome == "no_changes": + logger.info("[watch] No meaningful changes to upload") + else: + logger.info("[watch] Upload handling completed") + + def _finalize_successful_changes(self, changes: Dict[str, List]) -> None: + expected_hashes = self._last_expected_hashes + + def finalize_file(path: Path, expected_hash: Optional[str] = None) -> bool: + abs_path = str(path.resolve()) + stat = path.stat() + current_hash = hashlib.sha1(path.read_bytes()).hexdigest() + expected = str(expected_hash or "").strip().lower() + if ":" in expected: + expected = expected.partition(":")[2].strip() + if expected and current_hash != expected: + self._content_hash_cache.pop(abs_path, None) + return False + set_cached_file_hash(abs_path, current_hash, self.repo_name) + self._stat_cache[abs_path] = ( + int(getattr(stat, "st_mtime_ns", int(stat.st_mtime * 1e9))), + stat.st_size, + ) + self._content_hash_cache[abs_path] = ( + int(getattr(stat, "st_mtime_ns", int(stat.st_mtime * 1e9))), + stat.st_size, + current_hash, + ) + return True + + for path in changes.get("created", []) + changes.get("updated", []): + try: + rel_path = path.relative_to(Path(self.workspace_path)).as_posix() + finalize_file(path, expected_hashes.get(rel_path)) + except Exception: + continue + for path in changes.get("deleted", []): + try: + abs_path = str(path.resolve()) + remove_cached_file(abs_path, self.repo_name) + self._stat_cache.pop(abs_path, None) + self._content_hash_cache.pop(abs_path, None) + except Exception: + continue + for source_path, dest_path in changes.get("moved", []): + try: + source_abs_path = str(source_path.resolve()) + remove_cached_file(source_abs_path, self.repo_name) + self._stat_cache.pop(source_abs_path, None) + self._content_hash_cache.pop(source_abs_path, None) + except Exception: + pass + try: + dest_rel_path = dest_path.relative_to(Path(self.workspace_path)).as_posix() + finalize_file(dest_path, expected_hashes.get(dest_rel_path)) + except Exception: + continue def __enter__(self): """Context manager entry.""" @@ -715,7 +894,7 @@ def get_mapping_summary(self) -> Dict[str, Any]: container_path = self._translate_to_container_path(self.workspace_path) return { "repo_name": self.repo_name, - "collection_name": self.collection_name, + "collection_name": self.collection_name or "", "source_path": self.workspace_path, "container_path": container_path, "upload_endpoint": self.upload_endpoint, @@ -730,6 +909,51 @@ def log_mapping_summary(self) -> None: logger.info(f" source_path: {info['source_path']}") logger.info(f" container_path: {info['container_path']}") + def _excluded_dirnames(self) -> frozenset: + # Keep in sync with get_all_code_files exclusions. + # NOTE: This caches the exclusion set per client instance. + # Runtime changes to DEV_REMOTE_MODE/REMOTE_UPLOAD_MODE won't be reflected + # until a new client is created (typically via process restart), which is + # acceptable for the standalone upload client use case. + cached = getattr(self, "_excluded_dirnames_cache", None) + if cached is not None: + return cached + excluded = { + "node_modules", "vendor", "dist", "build", "target", "out", + ".git", ".hg", ".svn", ".vscode", ".idea", ".venv", "venv", + "__pycache__", ".pytest_cache", ".mypy_cache", ".cache", + ".context-engine", ".context-engine-uploader", ".codebase", + } + dev_remote = os.environ.get("DEV_REMOTE_MODE") == "1" or os.environ.get("REMOTE_UPLOAD_MODE") == "development" + if dev_remote: + excluded.add("dev-workspace") + cached = frozenset(excluded) + self._excluded_dirnames_cache = cached + return cached + + def _is_ignored_path(self, path: Path) -> bool: + """Return True when path is outside workspace or under excluded dirs.""" + try: + workspace_root = Path(self.workspace_path).resolve() + rel = path.resolve().relative_to(workspace_root) + except Exception: + return True + + dir_parts = set(rel.parts[:-1]) if len(rel.parts) > 1 else set() + if dir_parts & self._excluded_dirnames(): + return True + # Ignore hidden directories anywhere under the workspace, but allow + # extensionless dotfiles like `.gitignore` that we explicitly support. + if any(p.startswith(".") for p in rel.parts[:-1]): + return True + if rel.name.startswith(".") and rel.name.lower() not in EXTENSIONLESS_FILES: + return True + return False + + def _is_watchable_path(self, path: Path) -> bool: + """Return True when a filesystem event path is eligible for upload processing.""" + return not self._is_ignored_path(path) and detect_language(path) != "unknown" + def _get_temp_bundle_dir(self) -> Path: """Get or create temporary directory for bundle creation.""" if not self.temp_dir: @@ -757,6 +981,19 @@ def detect_file_changes(self, changed_paths: List[Path]) -> Dict[str, List]: } for path in changed_paths: + if self._is_ignored_path(path): + try: + abs_path = str(path.resolve()) + except Exception: + continue + cached_hash = get_cached_file_hash(abs_path, self.repo_name) + if cached_hash: + changes["deleted"].append(path) + try: + self._stat_cache.pop(abs_path, None) + except Exception: + pass + continue try: abs_path = str(path.resolve()) except Exception: @@ -797,9 +1034,7 @@ def detect_file_changes(self, changed_paths: List[Path]) -> Dict[str, List]: # Stat changed or no prior entry – hash content to classify change try: - with open(path, 'rb') as f: - content = f.read() - current_hash = hashlib.sha1(content).hexdigest() + current_hash, _ = self._read_current_file_hash(path) except Exception: # Skip files that can't be read continue @@ -819,8 +1054,6 @@ def detect_file_changes(self, changed_paths: List[Path]) -> Dict[str, List]: self._stat_cache[abs_path] = (getattr(stat, "st_mtime_ns", int(stat.st_mtime * 1e9)), stat.st_size) except Exception: pass - set_cached_file_hash(abs_path, current_hash, self.repo_name) - # Detect moves by looking for files with same content hash # but different paths (requires additional tracking) changes["moved"] = self._detect_moves(changes["created"], changes["deleted"]) @@ -945,8 +1178,6 @@ def create_delta_bundle( operations.append(operation) file_hashes[rel_path] = f"sha1:{file_hash}" total_size += stat.st_size - set_cached_file_hash(str(path.resolve()), file_hash, self.repo_name) - except Exception as e: print(f"[bundle_create] Error processing created file {path}: {e}") continue @@ -985,8 +1216,6 @@ def create_delta_bundle( operations.append(operation) file_hashes[rel_path] = f"sha1:{file_hash}" total_size += stat.st_size - set_cached_file_hash(str(path.resolve()), file_hash, self.repo_name) - except Exception as e: print(f"[bundle_create] Error processing updated file {path}: {e}") continue @@ -1027,8 +1256,6 @@ def create_delta_bundle( operations.append(operation) file_hashes[dest_rel_path] = f"sha1:{file_hash}" total_size += stat.st_size - set_cached_file_hash(str(dest_path.resolve()), file_hash, self.repo_name) - except Exception as e: print(f"[bundle_create] Error processing moved file {source_path} -> {dest_path}: {e}") continue @@ -1063,7 +1290,6 @@ def create_delta_bundle( "version": "1.0", "bundle_id": bundle_id, "workspace_path": self.workspace_path, - "collection_name": self.collection_name, "created_at": created_at, # CLI is stateless - server handles sequence numbers "sequence_number": None, # Server will assign @@ -1113,8 +1339,296 @@ def create_delta_bundle( with tarfile.open(bundle_path, "w:gz") as tar: tar.add(temp_path, arcname=f"{bundle_id}") + self._last_expected_hashes.update(file_hashes) return str(bundle_path), manifest + def _build_plan_payload(self, changes: Dict[str, List]) -> Dict[str, Any]: + created_at = datetime.now().isoformat() + bundle_id = str(uuid.uuid4()) + operations: List[Dict[str, Any]] = [] + file_hashes: Dict[str, str] = {} + total_size = 0 + + for path in changes["created"]: + rel_path = path.relative_to(Path(self.workspace_path)).as_posix() + try: + file_hash, size = self._read_current_file_hash(path) + operations.append( + { + "operation": "created", + "path": rel_path, + "size_bytes": size, + "content_hash": f"sha1:{file_hash}", + "language": detect_language(path), + } + ) + file_hashes[rel_path] = f"sha1:{file_hash}" + total_size += size + except Exception as e: + logger.warning("[remote_upload] Failed to prepare created plan entry for %s: %s", path, e) + + for path in changes["updated"]: + rel_path = path.relative_to(Path(self.workspace_path)).as_posix() + try: + file_hash, size = self._read_current_file_hash(path) + previous_hash = _format_cached_sha1( + get_cached_file_hash(str(path.resolve()), self.repo_name) + ) + operations.append( + { + "operation": "updated", + "path": rel_path, + "size_bytes": size, + "content_hash": f"sha1:{file_hash}", + "previous_hash": previous_hash, + "language": detect_language(path), + } + ) + file_hashes[rel_path] = f"sha1:{file_hash}" + total_size += size + except Exception as e: + logger.warning("[remote_upload] Failed to prepare updated plan entry for %s: %s", path, e) + + for source_path, dest_path in changes["moved"]: + dest_rel_path = dest_path.relative_to(Path(self.workspace_path)).as_posix() + source_rel_path = source_path.relative_to(Path(self.workspace_path)).as_posix() + try: + file_hash, size = self._read_current_file_hash(dest_path) + operations.append( + { + "operation": "moved", + "path": dest_rel_path, + "source_path": source_rel_path, + "size_bytes": size, + "content_hash": f"sha1:{file_hash}", + "language": detect_language(dest_path), + } + ) + file_hashes[dest_rel_path] = f"sha1:{file_hash}" + total_size += size + except Exception as e: + logger.warning( + "[remote_upload] Failed to prepare moved plan entry for %s -> %s: %s", + source_path, + dest_path, + e, + ) + + for path in changes["deleted"]: + rel_path = path.relative_to(Path(self.workspace_path)).as_posix() + try: + previous_hash = _format_cached_sha1( + get_cached_file_hash(str(path.resolve()), self.repo_name) + ) + operations.append( + { + "operation": "deleted", + "path": rel_path, + "previous_hash": previous_hash, + "language": detect_language(path), + } + ) + except Exception as e: + logger.warning("[remote_upload] Failed to prepare deleted plan entry for %s: %s", path, e) + + manifest = { + "version": "1.0", + "bundle_id": bundle_id, + "workspace_path": self.workspace_path, + "created_at": created_at, + "sequence_number": None, + "parent_sequence": None, + "operations": { + "created": len(changes["created"]), + "updated": len(changes["updated"]), + "deleted": len(changes["deleted"]), + "moved": len(changes["moved"]), + }, + "total_files": len(operations), + "total_size_bytes": total_size, + "compression": "gzip", + "encoding": "utf-8", + } + return { + "manifest": manifest, + "operations": operations, + "file_hashes": file_hashes, + } + + def _plan_delta_upload(self, changes: Dict[str, List]) -> Optional[Dict[str, Any]]: + if not _env_flag("CTXCE_REMOTE_UPLOAD_PLAN_ENABLED", True): + return None + try: + payload = self._build_plan_payload(changes) + self._last_plan_payload = payload + self._last_expected_hashes = dict(payload.get("file_hashes", {})) + # Indexed hashes are server-owned; the client submits candidates + # and uses only the returned plan. + data = { + "workspace_path": self._translate_to_container_path(self.workspace_path), + "source_path": self.workspace_path, + "logical_repo_id": _compute_logical_repo_id(self.workspace_path), + "manifest": payload["manifest"], + "operations": payload["operations"], + "file_hashes": payload["file_hashes"], + } + sess = get_auth_session(self.upload_endpoint) + if sess: + data["session"] = sess + if getattr(self, "logical_repo_id", None): + data["logical_repo_id"] = self.logical_repo_id + + response = self.session.post( + f"{self.upload_endpoint}/api/v1/delta/plan", + json=data, + timeout=min(self.timeout, 60), + ) + if response.status_code in {404, 405}: + logger.info("[remote_upload] Plan endpoint unavailable; falling back to full bundle upload") + return None + response.raise_for_status() + body = response.json() + if not body.get("success", False): + logger.warning("[remote_upload] Plan request failed; falling back: %s", body.get("error")) + return None + return body + except Exception as e: + logger.warning("[remote_upload] Plan request failed; falling back to full bundle upload: %s", e) + return None + + def _build_apply_only_payload(self, changes: Dict[str, List], plan: Dict[str, Any]) -> Dict[str, Any]: + payload = self._last_plan_payload or self._build_plan_payload(changes) + needed = plan.get("needed_files", {}) if isinstance(plan, dict) else {} + created_needed = set(needed.get("created", []) or []) + updated_needed = set(needed.get("updated", []) or []) + moved_needed = set(needed.get("moved", []) or []) + + # Check if ALL operations are hash-matched (nothing needs content at all) + # This happens when all needed_files lists are empty and there are no actual changes requiring content + has_changes_needing_content = bool(created_needed or updated_needed or moved_needed) + has_deletes = bool(changes.get("deleted", [])) + + # Only skip apply-only if there are NO operations needing content AND NO deletes + if not has_changes_needing_content and not has_deletes: + return { + "manifest": payload.get("manifest", {}), + "operations": [], + "file_hashes": {}, + } + + filtered_ops: List[Dict[str, Any]] = [] + filtered_hashes: Dict[str, str] = {} + for operation in payload.get("operations", []): + op_type = str(operation.get("operation") or "") + rel_path = str(operation.get("path") or "") + # Determine if this operation needs content (only those skip filtered_hashes) + needs_content = ( + (op_type == "created" and rel_path in created_needed) + or (op_type == "updated" and rel_path in updated_needed) + or (op_type == "moved" and rel_path in moved_needed) + ) + if needs_content: + # Skip operations that need content - they'll be uploaded separately + continue + # IMPORTANT: server-side apply_delta_operations() only accepts "deleted" and "moved" + # operations. Hash-matched "created" and "updated" operations must NOT be routed + # through apply_ops since the server will reject them. + if op_type not in {"deleted", "moved"}: + continue + # Preserve all other operations so server advances state + filtered_ops.append(operation) + # Include hash for non-deleted operations + if op_type != "deleted": + hash_value = payload.get("file_hashes", {}).get(rel_path) + if hash_value: + filtered_hashes[rel_path] = hash_value + return { + "manifest": payload.get("manifest", {}), + "operations": filtered_ops, + "file_hashes": filtered_hashes, + } + + def _apply_operations_without_content(self, changes: Dict[str, List], plan: Dict[str, Any]) -> Optional[bool]: + payload = self._build_apply_only_payload(changes, plan) + operations = payload.get("operations", []) + if not operations: + return None + try: + data = { + "workspace_path": self._translate_to_container_path(self.workspace_path), + "source_path": self.workspace_path, + "logical_repo_id": _compute_logical_repo_id(self.workspace_path), + "manifest": payload["manifest"], + "operations": operations, + "file_hashes": payload["file_hashes"], + } + sess = get_auth_session(self.upload_endpoint) + if sess: + data["session"] = sess + if getattr(self, "logical_repo_id", None): + data["logical_repo_id"] = self.logical_repo_id + + logger.info( + "[remote_upload] Applying metadata-only operations without bundle: deleted=%s moved=%s", + sum(1 for op in operations if op.get("operation") == "deleted"), + sum(1 for op in operations if op.get("operation") == "moved"), + ) + response = self.session.post( + f"{self.upload_endpoint}/api/v1/delta/apply_ops", + json=data, + timeout=min(self.timeout, 60), + ) + if response.status_code in {404, 405}: + logger.info("[remote_upload] apply_ops endpoint unavailable; falling back to bundle upload") + return None + response.raise_for_status() + body = response.json() + if not body.get("success", False): + logger.warning("[remote_upload] apply_ops failed; falling back to bundle upload: %s", body.get("error")) + return None + processed_ops = body.get("processed_operations") or {} + self._set_last_upload_result( + "uploaded", + bundle_id=body.get("bundle_id"), + sequence_number=body.get("sequence_number"), + processed_operations=processed_ops, + ) + logger.info( + "[remote_upload] Metadata-only operations applied: %s", + processed_ops, + ) + return True + except Exception as e: + logger.warning("[remote_upload] apply_ops failed; falling back to bundle upload: %s", e) + return None + + def _filter_changes_by_plan(self, changes: Dict[str, List], plan: Dict[str, Any]) -> Dict[str, List]: + needed = plan.get("needed_files", {}) if isinstance(plan, dict) else {} + created_needed = set(needed.get("created", []) or []) + updated_needed = set(needed.get("updated", []) or []) + moved_needed = set(needed.get("moved", []) or []) + + filtered_created = [ + path for path in changes["created"] + if path.relative_to(Path(self.workspace_path)).as_posix() in created_needed + ] + filtered_updated = [ + path for path in changes["updated"] + if path.relative_to(Path(self.workspace_path)).as_posix() in updated_needed + ] + filtered_moved = [ + (source_path, dest_path) + for source_path, dest_path in changes["moved"] + if dest_path.relative_to(Path(self.workspace_path)).as_posix() in moved_needed + ] + return { + "created": filtered_created, + "updated": filtered_updated, + "deleted": list(changes["deleted"]), + "moved": filtered_moved, + "unchanged": [], + } + def upload_bundle(self, bundle_path: str, manifest: Dict[str, Any]) -> Dict[str, Any]: """Upload delta bundle to remote server with exponential backoff retry. @@ -1142,18 +1656,13 @@ def upload_bundle(self, bundle_path: str, manifest: Dict[str, Any]) -> Dict[str, # Check bundle size (server-side enforcement) bundle_size = os.path.getsize(bundle_path) - files = { - "bundle": open(bundle_path, "rb"), - } data = { "workspace_path": self._translate_to_container_path(self.workspace_path), - "collection_name": self.collection_name, "sequence_number": manifest.get("sequence_number"), "force": False, "source_path": self.workspace_path, "logical_repo_id": _compute_logical_repo_id(self.workspace_path), } - sess = get_auth_session(self.upload_endpoint) if sess: data["session"] = sess @@ -1163,12 +1672,13 @@ def upload_bundle(self, bundle_path: str, manifest: Dict[str, Any]) -> Dict[str, logger.info(f"[remote_upload] Uploading bundle {manifest['bundle_id']} (size: {bundle_size} bytes)") - response = self.session.post( - f"{self.upload_endpoint}/api/v1/delta/upload", - files=files, - data=data, - timeout=(10, self.timeout) - ) + with open(bundle_path, "rb") as bundle_file: + response = self.session.post( + f"{self.upload_endpoint}/api/v1/delta/upload", + files={"bundle": bundle_file}, + data=data, + timeout=(10, self.timeout) + ) result = None try: @@ -1220,29 +1730,17 @@ def upload_bundle(self, bundle_path: str, manifest: Dict[str, Any]) -> Dict[str, last_error = {"success": False, "error": {"code": "TIMEOUT_ERROR", "message": f"Upload timeout: {str(e)}"}} logger.warning(f"[remote_upload] Upload read timeout on attempt {attempt + 1}: {e}") - # After read timeout, poll to check if server processed the bundle - logger.info(f"[remote_upload] Read timeout occurred, polling server to check if bundle was processed...") - poll_result = self._poll_after_timeout(manifest) - if poll_result.get("success"): - logger.info(f"[remote_upload] Server confirmed processing of bundle {manifest['bundle_id']} after timeout") - return poll_result - - logger.warning(f"[remote_upload] Server did not process bundle after timeout, proceeding with failure") - break + # A timeout is retried normally. The server-side hash plan makes + # a replay safe when the request was accepted before the timeout. + continue except requests.exceptions.Timeout as e: last_error = {"success": False, "error": {"code": "TIMEOUT_ERROR", "message": f"Upload timeout: {str(e)}"}} logger.warning(f"[remote_upload] Upload timeout on attempt {attempt + 1}: {e}") - # For generic timeout, also try polling - logger.info(f"[remote_upload] Timeout occurred, polling server to check if bundle was processed...") - poll_result = self._poll_after_timeout(manifest) - if poll_result.get("success"): - logger.info(f"[remote_upload] Server confirmed processing of bundle {manifest['bundle_id']} after timeout") - return poll_result - - logger.warning(f"[remote_upload] Server did not process bundle after timeout, proceeding with failure") - break + # A timeout is retried normally. The server-side hash plan makes + # a replay safe when the request was accepted before the timeout. + continue except requests.exceptions.ConnectionError as e: last_error = {"success": False, "error": {"code": "CONNECTION_ERROR", "message": f"Connection error: {str(e)}"}} @@ -1266,87 +1764,6 @@ def upload_bundle(self, bundle_path: str, manifest: Dict[str, Any]) -> Dict[str, } } - def _poll_after_timeout(self, manifest: Dict[str, Any]) -> Dict[str, Any]: - """ - Poll server status after a timeout to check if bundle was processed. - - Args: - manifest: Bundle manifest containing sequence information - - Returns: - Dictionary indicating success if bundle was processed - """ - try: - # Get current server status to know the expected sequence - status = self.get_server_status() - if not status.get("success"): - return {"success": False, "error": status.get("error", {"code": "UNKNOWN", "message": "Failed to get status"})} - - current_sequence = status.get("last_sequence", 0) - expected_sequence = manifest.get("sequence", current_sequence + 1) - - logger.info(f"[remote_upload] Current server sequence: {current_sequence}, expected: {expected_sequence}") - - # If server is already at expected sequence, bundle was processed - if current_sequence >= expected_sequence: - return { - "success": True, - "message": f"Bundle processed (server at sequence {current_sequence})", - "sequence": current_sequence, - } - - # Poll window is configurable via REMOTE_UPLOAD_POLL_MAX_SECS (seconds). - # Values <= 0 mean "no timeout" (poll until success or process exit). - try: - max_poll_time = int(os.environ.get("REMOTE_UPLOAD_POLL_MAX_SECS", "300")) - except Exception: - max_poll_time = 300 - poll_interval = 5 - start_time = time.time() - - while True: - elapsed = time.time() - start_time - if max_poll_time > 0 and elapsed >= max_poll_time: - logger.warning( - f"[remote_upload] Polling timed out after {int(elapsed)}s (limit={max_poll_time}s), bundle was not confirmed as processed" - ) - return { - "success": False, - "error": { - "code": "POLL_TIMEOUT", - "message": f"Bundle not confirmed processed after polling for {int(elapsed)}s (limit={max_poll_time}s)", - }, - } - - logger.info( - f"[remote_upload] Polling server status... (elapsed: {int(elapsed)}s, limit={'no-limit' if max_poll_time <= 0 else max_poll_time}s)" - ) - time.sleep(poll_interval) - - status = self.get_server_status() - if status.get("success"): - new_sequence = status.get("last_sequence", 0) - if new_sequence >= expected_sequence: - logger.info( - f"[remote_upload] Server sequence advanced to {new_sequence}, bundle was processed!" - ) - return { - "success": True, - "message": f"Bundle processed after timeout (server at sequence {new_sequence})", - "sequence": new_sequence, - } - logger.debug( - f"[remote_upload] Server sequence still at {new_sequence}, continuing to poll..." - ) - else: - logger.warning( - f"[remote_upload] Failed to get server status during poll: {status.get('error', {}).get('message', 'Unknown')}" - ) - - except Exception as e: - logger.error(f"[remote_upload] Error during post-timeout polling: {e}") - return {"success": False, "error": {"code": "POLL_ERROR", "message": f"Polling error: {str(e)}"}} - def get_server_status(self) -> Dict[str, Any]: """Get server status with simplified error handling.""" try: @@ -1354,14 +1771,27 @@ def get_server_status(self) -> Dict[str, Any]: connect_timeout = min(self.timeout, 10) # Allow slower responses (e.g., cold starts/large collections) before bailing read_timeout = max(self.timeout, 30) + params = {"workspace_path": container_workspace_path} + sess = get_auth_session(self.upload_endpoint) + if sess: + params["session"] = sess response = self.session.get( f"{self.upload_endpoint}/api/v1/delta/status", - params={'workspace_path': container_workspace_path}, + params=params, timeout=(connect_timeout, read_timeout) ) if response.status_code == 200: - return response.json() + payload = response.json() + if not isinstance(payload, dict): + return { + "success": False, + "error": { + "code": "STATUS_INVALID", + "message": "Invalid status response payload", + }, + } + return {"success": True, **payload} # Handle error response error_msg = f"Status check failed with HTTP {response.status_code}" @@ -1385,6 +1815,93 @@ def has_meaningful_changes(self, changes: Dict[str, List]) -> bool: total_changes = sum(len(files) for op, files in changes.items() if op != "unchanged") return total_changes > 0 + def _collect_force_cleanup_paths(self) -> List[Path]: + """ + Return ignored paths that force mode should actively delete remotely. + + In dev-remote mode, dev-workspace is intentionally ignored during upload + scans to avoid recursive dogfooding. If that tree already exists on the + remote side from an older buggy upload, force mode should remove it even + when the standalone client's cache does not know about those paths. + """ + cleanup_paths: List[Path] = [] + if "dev-workspace" not in self._excluded_dirnames(): + return cleanup_paths + + dev_root = Path(self.workspace_path) / "dev-workspace" + if not dev_root.exists(): + return cleanup_paths + + for root, dirnames, filenames in os.walk(dev_root): + dirnames[:] = [d for d in dirnames if not d.startswith(".")] + for filename in filenames: + path = Path(root) / filename + try: + if path.is_file(): + cleanup_paths.append(path) + except Exception: + continue + return cleanup_paths + + def build_force_changes(self, all_files: List[Path]) -> Dict[str, List]: + """ + Build force-upload changes while still cleaning stale cached paths. + + Force mode should re-upload every currently managed file, but it must also + emit deletes for files that only exist in the local cache now, including + paths that are ignored under the current client policy such as + dev-workspace in dev-remote mode. + """ + created_files: List[Path] = [] + path_map: Dict[Path, Path] = {} + for path in all_files: + if self._is_ignored_path(path): + continue + try: + resolved = path.resolve() + except Exception: + continue + created_files.append(path) + path_map[resolved] = path + + for cached_abs in self._get_all_cached_paths(): + try: + cached_path = Path(cached_abs) + resolved = cached_path.resolve() + except Exception: + continue + if resolved not in path_map: + path_map[resolved] = cached_path + + force_cleanup_paths = self._collect_force_cleanup_paths() + for cleanup_path in force_cleanup_paths: + try: + resolved = cleanup_path.resolve() + except Exception: + continue + if resolved not in path_map: + path_map[resolved] = cleanup_path + + probed = self.detect_file_changes(list(path_map.values())) + deleted_by_resolved: Dict[Path, Path] = {} + for deleted_path in probed.get("deleted", []): + try: + deleted_by_resolved[deleted_path.resolve()] = deleted_path + except Exception: + continue + for cleanup_path in force_cleanup_paths: + try: + deleted_by_resolved.setdefault(cleanup_path.resolve(), cleanup_path) + except Exception: + continue + return { + "created": created_files, + "updated": [], + "deleted": list(deleted_by_resolved.values()), + "moved": [], + "unchanged": [], + } + def upload_git_history_only(self, git_history: Dict[str, Any]) -> bool: try: empty_changes = { @@ -1426,25 +1943,80 @@ def process_changes_and_upload(self, changes: Dict[str, List]) -> bool: try: logger.info(f"[remote_upload] Processing pre-computed changes") + self._last_plan_payload = None + self._last_expected_hashes = {} + # Validate input if not changes: logger.info("[remote_upload] No changes provided") + self._set_last_upload_result("no_changes") return True if not self.has_meaningful_changes(changes): logger.info("[remote_upload] No meaningful changes detected, skipping upload") + self._set_last_upload_result("no_changes") return True # Log change summary total_changes = sum(len(files) for op, files in changes.items() if op != "unchanged") - logger.info(f"[remote_upload] Detected {total_changes} meaningful changes: " - f"{len(changes['created'])} created, {len(changes['updated'])} updated, " - f"{len(changes['deleted'])} deleted, {len(changes['moved'])} moved") + logger.info( + "[remote_upload] Detected %d candidate changes before remote planning: " + "created=%d updated=%d deleted=%d moved=%d", + total_changes, + len(changes["created"]), + len(changes["updated"]), + len(changes["deleted"]), + len(changes["moved"]), + ) + + planned_changes = changes + plan = self._plan_delta_upload(changes) + if plan: + preview = plan.get("operation_counts_preview", {}) + needed = plan.get("needed_files", {}) if isinstance(plan.get("needed_files", {}), dict) else {} + diagnostics = plan.get("diagnostics", {}) if isinstance(plan.get("diagnostics", {}), dict) else {} + logger.info( + "[remote_upload] Remote plan: candidates=%d content_needed=" + "created=%s updated=%s moved=%s deletes=%s " + "skipped_hash_match=%s needed_bytes=%s cache_entries=%s " + "cache_hash_matches=%s", + total_changes, + len(needed.get("created", []) or []), + len(needed.get("updated", []) or []), + len(needed.get("moved", []) or []), + len(changes.get("deleted", [])), + preview.get("skipped_hash_match", 0), + plan.get("needed_size_bytes", 0), + diagnostics.get("cache_entries", 0), + diagnostics.get("cache_hash_matches", 0), + ) + planned_changes = self._filter_changes_by_plan(changes, plan) + has_content_work = bool( + planned_changes.get("created") + or planned_changes.get("updated") + or planned_changes.get("moved") + ) + if not has_content_work: + apply_only_result = self._apply_operations_without_content(changes, plan) + if apply_only_result is True: + self._finalize_successful_changes(changes) + self._flush_cached_file_hashes() + return True + if not self.has_meaningful_changes(planned_changes): + logger.info("[remote_upload] Plan found no upload work; skipping bundle upload") + self._finalize_successful_changes(changes) + self._set_last_upload_result( + "skipped_by_plan", + plan_preview=preview, + needed_size_bytes=plan.get("needed_size_bytes", 0), + ) + self._flush_cached_file_hashes() + return True # Create delta bundle bundle_path = None try: - bundle_path, manifest = self.create_delta_bundle(changes) + bundle_path, manifest = self.create_delta_bundle(planned_changes) logger.info(f"[remote_upload] Created delta bundle: {manifest['bundle_id']} " f"(size: {manifest['total_size_bytes']} bytes)") @@ -1456,6 +2028,7 @@ def process_changes_and_upload(self, changes: Dict[str, List]) -> bool: logger.error(f"[remote_upload] Error creating delta bundle: {e}") # Clean up any temporary files on failure self.cleanup() + self._set_last_upload_result("failed", stage="bundle_creation", error=str(e)) return False # Upload bundle with retry logic @@ -1463,9 +2036,32 @@ def process_changes_and_upload(self, changes: Dict[str, List]) -> bool: response = self.upload_bundle(bundle_path, manifest) if response.get("success", False): - processed_ops = response.get('processed_operations', {}) - logger.info(f"[remote_upload] Successfully uploaded bundle {manifest['bundle_id']}") - logger.info(f"[remote_upload] Processed operations: {processed_ops}") + processed_ops = response.get("processed_operations") + if processed_ops is None: + logger.info( + "[remote_upload] Bundle %s accepted by server; processing asynchronously (sequence=%s)", + manifest["bundle_id"], + response.get("sequence_number"), + ) + self._set_last_upload_result( + "queued", + bundle_id=manifest["bundle_id"], + sequence_number=response.get("sequence_number"), + ) + # Acceptance is the client-side completion point. The + # server owns background processing and journal retry. + self._finalize_successful_changes(changes) + else: + logger.info(f"[remote_upload] Successfully uploaded bundle {manifest['bundle_id']}") + logger.info(f"[remote_upload] Processed operations: {processed_ops}") + self._finalize_successful_changes(changes) + self._set_last_upload_result( + "uploaded", + bundle_id=manifest["bundle_id"], + sequence_number=response.get("sequence_number"), + processed_operations=processed_ops, + ) + self._flush_cached_file_hashes() # Clean up temporary bundle after successful upload try: @@ -1481,14 +2077,17 @@ def process_changes_and_upload(self, changes: Dict[str, List]) -> bool: else: error_msg = response.get('error', {}).get('message', 'Unknown upload error') logger.error(f"[remote_upload] Upload failed: {error_msg}") + self._set_last_upload_result("failed", stage="upload", error=error_msg) return False except Exception as e: logger.error(f"[remote_upload] Error uploading bundle: {e}") + self._set_last_upload_result("failed", stage="upload", error=str(e)) return False except Exception as e: logger.error(f"[remote_upload] Unexpected error in process_changes_and_upload: {e}") + self._set_last_upload_result("failed", stage="unexpected", error=str(e)) return False def watch_loop(self, interval: int = 5): @@ -1515,6 +2114,7 @@ def __init__(self, client, debounce_seconds=2.0): self._pending_paths = set() self._check_for_deletions = False self._lock = threading.Lock() + self._processing = False def on_any_event(self, event): """Handle any file system event.""" @@ -1532,13 +2132,13 @@ def on_any_event(self, event): # Always check src_path src_path = Path(event.src_path) - if detect_language(src_path) != "unknown": + if self.client._is_watchable_path(src_path): paths_to_process.append(src_path) # For FileMovedEvent, also process the destination path if hasattr(event, 'dest_path') and event.dest_path: dest_path = Path(event.dest_path) - if detect_language(dest_path) != "unknown": + if self.client._is_watchable_path(dest_path): paths_to_process.append(dest_path) if not paths_to_process: @@ -1559,22 +2159,30 @@ def on_any_event(self, event): def _process_pending_changes(self): """Process accumulated changes after debounce period.""" with self._lock: + # Timer fired; allow a new debounce to be armed while we process. + self._debounce_timer = None + if self._processing: + return if not self._pending_paths: return + self._processing = True pending = list(self._pending_paths) self._pending_paths.clear() check_deletions = self._check_for_deletions self._check_for_deletions = False + upload_succeeded = False try: # Only include cached paths when deletion-related events occurred if check_deletions: - all_paths = list(set(pending + [ - Path(p) for p in get_all_cached_paths(self.client.repo_name) - ])) + cached_paths = [ + Path(p) for p in self.client._get_all_cached_paths() + ] + all_paths = list(set(pending + cached_paths)) else: all_paths = pending - + + changes = self.client.detect_file_changes(all_paths) meaningful_changes = ( len(changes.get("created", [])) + @@ -1582,12 +2190,13 @@ def _process_pending_changes(self): len(changes.get("deleted", [])) + len(changes.get("moved", [])) ) - + if meaningful_changes > 0: logger.info(f"[watch] Detected {meaningful_changes} changes: { {k: len(v) for k, v in changes.items() if k != 'unchanged'} }") success = self.client.process_changes_and_upload(changes) if success: - logger.info("[watch] Successfully uploaded changes") + self.client.log_watch_upload_result() + upload_succeeded = True else: logger.error("[watch] Failed to upload changes") else: @@ -1597,20 +2206,37 @@ def _process_pending_changes(self): git_history = _collect_git_history_for_workspace(self.client.workspace_path) except Exception: git_history = None - + if git_history: logger.info("[watch] Detected git history update; uploading git history metadata") success = self.client.upload_git_history_only(git_history) if success: logger.info("[watch] Successfully uploaded git history metadata") + upload_succeeded = True else: logger.error("[watch] Failed to upload git history metadata") + else: + upload_succeeded = True # No changes to process except Exception as e: logger.error(f"[watch] Error processing changes: {e}") + finally: + with self._lock: + self._processing = False + # Re-queue pending paths if upload failed + if not upload_succeeded and pending: + # Merge pending paths back into _pending_paths + for p in pending: + self._pending_paths.add(p) + # Arm next pass if there are pending paths + if self._pending_paths and self._debounce_timer is None: + self._debounce_timer = threading.Timer( + self.debounce_seconds, + self._process_pending_changes, + ) + self._debounce_timer.start() observer = Observer() handler = CodeFileEventHandler(self, debounce_seconds=2.0) - try: observer.schedule(handler, self.workspace_path, recursive=True) observer.start() @@ -1655,7 +2281,7 @@ def _watch_loop_polling(self, interval: int = 5): path_map[resolved] = p # Include any paths that are only present in the local cache (deleted files) - for cached_abs in get_all_cached_paths(self.repo_name): + for cached_abs in self._get_all_cached_paths(): try: cached_path = Path(cached_abs) resolved = cached_path.resolve() @@ -1676,7 +2302,7 @@ def _watch_loop_polling(self, interval: int = 5): success = self.process_changes_and_upload(changes) if success: - logger.info(f"[watch] Successfully uploaded changes") + self.log_watch_upload_result() else: logger.error(f"[watch] Failed to upload changes") else: @@ -1719,16 +2345,10 @@ def get_all_code_files(self) -> List[Path]: # Single walk with early pruning and set-based matching to reduce IO ext_suffixes = {str(ext).lower() for ext in CODE_EXTS if str(ext).startswith('.')} - extensionless_names = set(EXTENSIONLESS_FILES.keys()) + extensionless_names = {k.lower() for k in EXTENSIONLESS_FILES.keys()} # Always exclude dev-workspace to prevent recursive upload loops # (upload service creates dev-workspace// which would otherwise get re-uploaded) - excluded = { - "node_modules", "vendor", "dist", "build", "target", "out", - ".git", ".hg", ".svn", ".vscode", ".idea", ".venv", "venv", - "__pycache__", ".pytest_cache", ".mypy_cache", ".cache", - ".context-engine", ".context-engine-uploader", ".codebase", - "dev-workspace" - } + excluded = self._excluded_dirnames() seen = set() for root, dirnames, filenames in os.walk(workspace_path): @@ -1741,6 +2361,8 @@ def get_all_code_files(self) -> List[Path]: if filename.startswith('.') and fname_lower not in extensionless_names: continue candidate = Path(root) / filename + if self._is_ignored_path(candidate): + continue suffix = candidate.suffix.lower() # Match by extension, extensionless name, or Dockerfile.* prefix if (suffix in ext_suffixes or @@ -1780,87 +2402,14 @@ def process_and_upload_changes(self, changed_paths: List[Path]) -> bool: except Exception as e: logger.error(f"[remote_upload] Error detecting file changes: {e}") return False - - if not self.has_meaningful_changes(changes): - logger.info("[remote_upload] No meaningful changes detected, skipping upload") - return True - - # Log change summary - total_changes = sum(len(files) for op, files in changes.items() if op != "unchanged") - logger.info(f"[remote_upload] Detected {total_changes} meaningful changes: " - f"{len(changes['created'])} created, {len(changes['updated'])} updated, " - f"{len(changes['deleted'])} deleted, {len(changes['moved'])} moved") - - # Create delta bundle - bundle_path = None - try: - bundle_path, manifest = self.create_delta_bundle(changes) - logger.info(f"[remote_upload] Created delta bundle: {manifest['bundle_id']} " - f"(size: {manifest['total_size_bytes']} bytes)") - - # Validate bundle was created successfully - if not bundle_path or not os.path.exists(bundle_path): - raise RuntimeError(f"Failed to create bundle at {bundle_path}") - - except Exception as e: - logger.error(f"[remote_upload] Error creating delta bundle: {e}") - # Clean up any temporary files on failure - self.cleanup() - return False - - # Upload bundle with retry logic - try: - response = self.upload_bundle(bundle_path, manifest) - - if response.get("success", False): - processed_ops = response.get('processed_operations', {}) - logger.info(f"[remote_upload] Successfully uploaded bundle {manifest['bundle_id']}") - logger.info(f"[remote_upload] Processed operations: {processed_ops}") - - # Clean up temporary bundle after successful upload - try: - if os.path.exists(bundle_path): - os.remove(bundle_path) - logger.debug(f"[remote_upload] Cleaned up temporary bundle: {bundle_path}") - # Also clean up the entire temp directory if this is the last bundle - self.cleanup() - except Exception as cleanup_error: - logger.warning(f"[remote_upload] Failed to cleanup bundle {bundle_path}: {cleanup_error}") - - return True - else: - error = response.get("error", {}) - error_code = error.get("code", "UNKNOWN") - error_msg = error.get("message", "Unknown error") - - logger.error(f"[remote_upload] Upload failed: {error_msg}") - - # Handle specific error types - # CLI is stateless - server handles sequence management - if error_code in ["BUNDLE_TOO_LARGE", "BUNDLE_NOT_FOUND"]: - # These are unrecoverable errors - logger.error(f"[remote_upload] Unrecoverable error ({error_code}): {error_msg}") - return False - elif error_code in ["TIMEOUT_ERROR", "CONNECTION_ERROR", "NETWORK_ERROR"]: - # These might be temporary, suggest fallback - logger.warning(f"[remote_upload] Network-related error ({error_code}): {error_msg}") - logger.warning("[remote_upload] Consider falling back to local mode if this persists") - return False - else: - # Other errors - logger.error(f"[remote_upload] Upload error ({error_code}): {error_msg}") - return False - - except Exception as e: - logger.error(f"[remote_upload] Unexpected error during upload: {e}") - return False + return self.process_changes_and_upload(changes) except Exception as e: logger.error(f"[remote_upload] Critical error in process_and_upload_changes: {e}") logger.exception("[remote_upload] Full traceback:") return False -def get_remote_config(cli_path: Optional[str] = None) -> Dict[str, str]: +def get_remote_config(cli_path: Optional[str] = None) -> Dict[str, Any]: """Get remote upload configuration from environment variables and command-line arguments.""" # Use command-line path if provided, otherwise fall back to environment variables if cli_path: @@ -1870,17 +2419,10 @@ def get_remote_config(cli_path: Optional[str] = None) -> Dict[str, str]: logical_repo_id = _compute_logical_repo_id(workspace_path) - # Use auto-generated collection name based on repo name - repo_name = _extract_repo_name_from_path(workspace_path) - # Fallback to directory name if repo detection fails - if not repo_name: - repo_name = Path(workspace_path).name - collection_name = get_collection_name(repo_name) - return { "upload_endpoint": os.environ.get("REMOTE_UPLOAD_ENDPOINT", "http://localhost:8080"), "workspace_path": workspace_path, - "collection_name": collection_name, + "collection_name": None, "logical_repo_id": logical_repo_id, # Use higher, more robust defaults but still allow env overrides "max_retries": int(os.environ.get("REMOTE_UPLOAD_MAX_RETRIES", "5")), @@ -2006,7 +2548,7 @@ def main(): config["timeout"] = args.timeout logger.info(f"Workspace path: {config['workspace_path']}") - logger.info(f"Collection name: {config['collection_name']}") + logger.info(f"Collection name: {config['collection_name'] or ''}") logger.info(f"Upload endpoint: {config['upload_endpoint']}") if args.show_mapping: @@ -2040,18 +2582,16 @@ def main(): # Test server connection first logger.info("Checking server status...") status = client.get_server_status() - is_success = ( - isinstance(status, dict) and - 'workspace_path' in status and - 'collection_name' in status and - status.get('status') == 'ready' - ) - if not is_success: - error = status.get("error", {}) - logger.error(f"Cannot connect to server: {error.get('message', 'Unknown error')}") + if not _is_usable_delta_status(status): + logger.error("Cannot connect to server: %s", _server_status_error_message(status)) return 1 - logger.info("Server connection successful") + logger.info( + "Server connection successful: status=%s pending_journal=%s journal=%s", + status.get("status"), + status.get("pending_operations"), + (status.get("server_info") or {}).get("journal", {}), + ) logger.info(f"Starting file monitoring with {args.interval}s interval") # Start the watch loop @@ -2085,16 +2625,8 @@ def main(): # Test server connection logger.info("Checking server status...") status = client.get_server_status() - # For delta endpoint, success is indicated by having expected fields (not a "success" boolean) - is_success = ( - isinstance(status, dict) and - 'workspace_path' in status and - 'collection_name' in status and - status.get('status') == 'ready' - ) - if not is_success: - error = status.get("error", {}) - logger.error(f"Cannot connect to server: {error.get('message', 'Unknown error')}") + if not _is_usable_delta_status(status): + logger.error("Cannot connect to server: %s", _server_status_error_message(status)) return 1 logger.info("Server connection successful") @@ -2105,7 +2637,7 @@ def main(): # Find code files in the repository (exclude hidden and heavy dirs) all_files = client.get_all_code_files() - logger.info(f"Found {len(all_files)} code files to upload") + logger.info(f"Found {len(all_files)} eligible files to scan") if not all_files: logger.warning("No files found to upload") @@ -2113,8 +2645,7 @@ def main(): # Detect changes (treat all files as changes for initial upload) if args.force: - # Force mode: treat all files as created - changes = {"created": all_files, "updated": [], "deleted": [], "moved": [], "unchanged": []} + changes = client.build_force_changes(all_files) else: changes = client.detect_file_changes(all_files) @@ -2122,16 +2653,32 @@ def main(): logger.info("No meaningful changes to upload") return 0 - logger.info(f"Changes detected: {len(changes.get('created', []))} created, {len(changes.get('updated', []))} updated, {len(changes.get('deleted', []))} deleted") + logger.info( + "Candidates detected before remote planning: created=%d updated=%d deleted=%d moved=%d", + len(changes.get("created", [])), + len(changes.get("updated", [])), + len(changes.get("deleted", [])), + len(changes.get("moved", [])), + ) # Process and upload changes logger.info("Uploading files to remote server...") success = client.process_changes_and_upload(changes) if success: - logger.info("Repository upload completed successfully!") - logger.info(f"Collection name: {config['collection_name']}") - logger.info(f"Files uploaded: {len(all_files)}") + outcome = str((client.last_upload_result or {}).get("outcome") or "") + if outcome == "skipped_by_plan": + logger.info("No upload needed after plan") + elif outcome == "queued": + logger.info("Repository upload request accepted; server processing asynchronously") + else: + logger.info("Repository upload completed successfully!") + logger.info(f"Collection name: {config['collection_name'] or ''}") + logger.info( + "Remote upload result: outcome=%s details=%s", + outcome, + client.last_upload_result, + ) else: logger.error("Repository upload failed!") return 1 diff --git a/scripts/upload_delta_bundle.py b/scripts/upload_delta_bundle.py index 973be132..b8fccd6e 100644 --- a/scripts/upload_delta_bundle.py +++ b/scripts/upload_delta_bundle.py @@ -1,24 +1,21 @@ import os import json +import shutil import tarfile import hashlib import re import logging from pathlib import Path -from typing import Dict, Any, Optional +from typing import Any, Dict, Optional - -try: - from scripts.workspace_state import ( - _extract_repo_name_from_path, - get_staging_targets, - get_collection_state_snapshot, - is_staging_enabled, - ) -except ImportError as exc: - raise ImportError( - "upload_delta_bundle requires scripts.workspace_state; ensure the module is available" - ) from exc +from scripts.workspace_state import ( + _normalize_cache_key_path, + _extract_repo_name_from_path, + get_staging_targets, + get_collection_state_snapshot, + is_staging_enabled, + upsert_index_journal_entries, +) logger = logging.getLogger(__name__) @@ -27,6 +24,113 @@ _SLUGGED_REPO_RE = re.compile(r"^.+-[0-9a-f]{16}(?:_old)?$") +def _normalize_hash_value(value: Any) -> str: + raw = str(value or "").strip() + if not raw: + return "" + if ":" in raw: + _, _, digest = raw.partition(":") + if digest.strip(): + return digest.strip().lower() + return raw.lower() + + +def _file_matches_hash(path: Path, expected_hash: str) -> bool: + """Verify a destination is already the requested content for idempotent retries.""" + expected = _normalize_hash_value(expected_hash) + if not expected or not path.is_file(): + return False + try: + return hashlib.sha1(path.read_bytes()).hexdigest() == expected + except OSError: + return False + + +def _build_upsert_journal_entry(path: Path | str, content_hash: Optional[str]) -> Dict[str, Any]: + entry: Dict[str, Any] = { + "path": str(path), + "op_type": "upsert", + } + if content_hash: + entry["content_hash"] = content_hash + return entry + + +def _build_delete_journal_entry(path: Path | str, content_hash: Optional[str] = None) -> Dict[str, Any]: + entry: Dict[str, Any] = { + "path": str(path), + "op_type": "delete", + } + if content_hash: + entry["content_hash"] = content_hash + return entry + + +def _load_cache_hashes(cache_path: Path) -> Dict[str, str]: + try: + with cache_path.open("r", encoding="utf-8-sig") as f: + data = json.load(f) + except (OSError, ValueError, json.JSONDecodeError): + return {} + + file_hashes = data.get("file_hashes", {}) + if not isinstance(file_hashes, dict): + return {} + + normalized: Dict[str, str] = {} + for path_key, value in file_hashes.items(): + if isinstance(value, dict): + hash_value = value.get("hash") + else: + hash_value = value + digest = _normalize_hash_value(hash_value) + if digest: + normalized[_normalize_cache_key_path(str(path_key))] = digest + return normalized + + +def _load_replica_cache_hashes(workspace_root: Path, slug: str) -> Dict[str, str]: + merged: Dict[str, str] = {} + cache_paths = ( + Path(WORK_DIR) / ".codebase" / "repos" / slug / "cache.json", + workspace_root / ".codebase" / "cache.json", + ) + for cache_path in cache_paths: + if not cache_path.exists(): + continue + merged.update(_load_cache_hashes(cache_path)) + return merged + + +def _flush_replica_cache_hashes(workspace_root: Path, slug: str, hashes: Dict[str, str]) -> None: + """Flush replica hashes to workspace cache.json.""" + try: + cache_path = workspace_root / ".codebase" / "cache.json" + cache_path.parent.mkdir(parents=True, exist_ok=True) + + # Read existing cache to preserve other entries + existing_data = {} + if cache_path.exists(): + try: + with cache_path.open("r", encoding="utf-8-sig") as f: + existing_data = json.load(f) + except (OSError, ValueError, json.JSONDecodeError): + existing_data = {} + + # Update file_hashes section + if not isinstance(existing_data, dict): + existing_data = {} + existing_data["file_hashes"] = hashes + + # Write back atomically + temp_path = cache_path.with_suffix(".tmp") + with temp_path.open("w", encoding="utf-8") as f: + json.dump(existing_data, f, indent=2) + temp_path.replace(cache_path) + except Exception as e: + logger.debug(f"[upload_service] Failed to flush cache for {slug}: {e}") + + def get_workspace_key(workspace_path: str) -> str: """Generate 16-char hash for collision avoidance in remote uploads. @@ -61,175 +165,586 @@ def _cleanup_empty_dirs(path: Path, stop_at: Path) -> None: break -def process_delta_bundle(workspace_path: str, bundle_path: Path, manifest: Dict[str, Any]) -> Dict[str, int]: - """Process delta bundle and return operation counts.""" +def _resolve_replica_roots(workspace_path: str, *, create_missing: bool = True) -> Dict[str, Path]: + workspace_leaf = Path(workspace_path).name + + repo_name_for_state: Optional[str] = None + serving_slug: Optional[str] = None + active_slug: Optional[str] = None + try: + repo_name_for_state = _extract_repo_name_from_path(workspace_path) + if repo_name_for_state: + snapshot = get_collection_state_snapshot( + workspace_path=None, + repo_name=repo_name_for_state, + ) # type: ignore[arg-type] + serving_slug = snapshot.get("serving_repo_slug") + active_slug = snapshot.get("active_repo_slug") + except Exception: + serving_slug = None + active_slug = None + + slug_order: list[str] = [] + serving_candidate: Optional[str] = None + if serving_slug and _SLUGGED_REPO_RE.match(serving_slug): + serving_candidate = serving_slug + if active_slug and _SLUGGED_REPO_RE.match(active_slug) and active_slug not in slug_order: + slug_order.append(active_slug) + + staging_active = False + staging_gate = bool(is_staging_enabled()) + try: + if serving_slug and str(serving_slug).endswith("_old"): + staging_active = True + except Exception: + staging_active = False + + if not staging_gate: + staging_active = False + + def _append_slug(slug: Optional[str]) -> None: + if slug and _SLUGGED_REPO_RE.match(slug) and slug not in slug_order: + slug_order.append(slug) + + if repo_name_for_state and _SLUGGED_REPO_RE.match(repo_name_for_state): + canonical_slug = ( + repo_name_for_state[:-4] + if repo_name_for_state.endswith("_old") + else repo_name_for_state + ) + old_slug_candidate = ( + repo_name_for_state + if repo_name_for_state.endswith("_old") + else f"{canonical_slug}_old" + ) + if staging_active: + slug_order = [] + _append_slug(canonical_slug) + _append_slug(old_slug_candidate) + elif not slug_order: + _append_slug(canonical_slug) + old_slug_path = Path(WORK_DIR) / old_slug_candidate + if old_slug_path.exists(): + _append_slug(old_slug_candidate) + + if not slug_order: + if _SLUGGED_REPO_RE.match(workspace_leaf): + slug_order.append(workspace_leaf) + else: + repo_name = _extract_repo_name_from_path(workspace_path) or workspace_leaf + workspace_key = get_workspace_key(workspace_path) + slug_order.append(f"{repo_name}-{workspace_key}") + + if staging_gate and not staging_active: + try: + repo_name_for_staging = _extract_repo_name_from_path(workspace_path) or slug_order[0] + targets = get_staging_targets( + workspace_path=workspace_path, + repo_name=repo_name_for_staging, + ) + if isinstance(targets, dict) and targets.get("staging"): + staging_active = True + except Exception as staging_err: + logger.debug("[upload_service] Failed to detect staging: %s", staging_err) + + def _slug_exists(slug: str) -> bool: + try: + return ( + (Path(WORK_DIR) / slug).exists() + or (Path(WORK_DIR) / ".codebase" / "repos" / slug).exists() + ) + except Exception: + return False + + if staging_gate and (not staging_active) and slug_order: + primary = slug_order[0] + if _SLUGGED_REPO_RE.match(primary): + canonical = primary[:-4] if primary.endswith("_old") else primary + inferred_old = primary if primary.endswith("_old") else f"{canonical}_old" + if _slug_exists(inferred_old): + staging_active = True + + if staging_gate and staging_active and slug_order: + primary = slug_order[0] + if _SLUGGED_REPO_RE.match(primary): + canonical = primary[:-4] if primary.endswith("_old") else primary + old_slug = primary if primary.endswith("_old") else f"{canonical}_old" + desired = [canonical, old_slug] + slug_order = [s for s in desired if _SLUGGED_REPO_RE.match(s)] + elif staging_gate and not staging_active and serving_candidate: + # Keep the canonical replica when serving and active are the same. + # A serving-only target should be removed only when a distinct active + # replica is available. + if ( + active_slug + and active_slug != serving_candidate + and serving_candidate in slug_order + ): + slug_order = [s for s in slug_order if s != serving_candidate] + + if staging_gate: + try: + logger.info("[upload_service] Delta bundle targets (staging=%s): %s", staging_active, slug_order) + except Exception: + pass + + replica_roots: Dict[str, Path] = {} + for slug in slug_order: + path = Path(WORK_DIR) / slug + if create_missing: + path.mkdir(parents=True, exist_ok=True) + try: + marker_dir = Path(WORK_DIR) / ".codebase" / "repos" / slug + marker_dir.mkdir(parents=True, exist_ok=True) + (marker_dir / ".ctxce_managed_upload").write_text("1\n") + except Exception: + pass + replica_roots[slug] = path.resolve() + return replica_roots + + +def _enqueue_replica_journal_entries( + *, + workspace_root: Path, + slug: str, + entries: list[Dict[str, Any]], +) -> None: + if not entries: + return + try: + upsert_index_journal_entries( + entries, + workspace_path=str(workspace_root), + repo_name=slug, + ) + except Exception as exc: + logger.debug( + "[upload_service] Failed to enqueue index journal entries for %s: %s", + workspace_root, + exc, + ) + + +def _safe_join(base: Path, rel: str) -> Path: + rp = Path(str(rel)) + if str(rp) in {".", ""}: + raise ValueError("Invalid operation path") + if rp.is_absolute(): + raise ValueError(f"Absolute paths are not allowed: {rel}") + base_resolved = base.resolve() + candidate = (base_resolved / rp).resolve() + try: + ok = candidate.is_relative_to(base_resolved) + except Exception: + ok = os.path.commonpath([str(base_resolved), str(candidate)]) == str(base_resolved) + if not ok: + raise ValueError(f"Path escapes workspace: {rel}") + return candidate + + +def _sanitize_operation_path(rel_path: str, replica_roots: Dict[str, Path]) -> Optional[str]: + sanitized_path = rel_path + skipped_due_to_exact_slug = False + for slug in replica_roots.keys(): + if sanitized_path == slug: + skipped_due_to_exact_slug = True + break + prefix = f"{slug}/" + if sanitized_path.startswith(prefix): + sanitized_path = sanitized_path[len(prefix):] + break + if skipped_due_to_exact_slug or not sanitized_path: + return None + return sanitized_path + + +def plan_delta_upload( + workspace_path: str, + operations: list[Dict[str, Any]], + file_hashes: Optional[Dict[str, str]] = None, +) -> Dict[str, Any]: + needed_files = { + "created": [], + "updated": [], + "moved": [], + } operations_count = { "created": 0, "updated": 0, "deleted": 0, "moved": 0, "skipped": 0, + "skipped_hash_match": 0, "failed": 0, } + needed_size_bytes = 0 + replica_roots = _resolve_replica_roots(workspace_path, create_missing=False) + replica_cache_hashes = { + slug: _load_replica_cache_hashes(root, slug) + for slug, root in replica_roots.items() + } + diagnostics = { + "candidate_operations": len(operations or []), + "cache_entries": sum(len(hashes) for hashes in replica_cache_hashes.values()), + "cache_hash_matches": 0, + "hash_mismatches": 0, + "missing_targets": 0, + } + normalized_hashes = { + str(rel_path): _normalize_hash_value(hash_value) + for rel_path, hash_value in (file_hashes or {}).items() + if _normalize_hash_value(hash_value) + } - try: - # CRITICAL: Always materialize writes under WORK_DIR using a slugged repo directory. - # Do NOT write directly into the client-supplied workspace_path, since that may be a host - # path (e.g. /home/user/repo) that is not mounted/visible to the watcher/indexer. - workspace_leaf = Path(workspace_path).name + for operation in operations: + op_type = str(operation.get("operation") or "") + rel_path = operation.get("path") + if not rel_path: + operations_count["skipped"] += 1 + continue + + sanitized = _sanitize_operation_path(str(rel_path), replica_roots) + if not sanitized: + operations_count["skipped"] += 1 + continue + + if op_type == "deleted": + operations_count["deleted"] += 1 + continue + if op_type == "moved": + operations_count["moved"] += 1 + source_rel_path = operation.get("source_path") or operation.get("source_relative_path") + if not source_rel_path: + needed_files["moved"].append(sanitized) + needed_size_bytes += int(operation.get("size_bytes") or 0) + continue + + move_needs_content = False + for _slug, root in replica_roots.items(): + try: + safe_source_path = _safe_join(root, str(source_rel_path)) + except ValueError: + logger.warning( + "[upload_service] Invalid move source path during plan: %s (root=%s)", + source_rel_path, + root, + ) + move_needs_content = True + break + if not safe_source_path.exists(): + move_needs_content = True + break + if move_needs_content: + needed_files["moved"].append(sanitized) + needed_size_bytes += int(operation.get("size_bytes") or 0) + continue + if op_type not in {"created", "updated"}: + operations_count["failed"] += 1 + continue + + op_content_hash = _normalize_hash_value( + operation.get("content_hash") or normalized_hashes.get(sanitized) + ) + if not op_content_hash: + needed_files[op_type].append(sanitized) + operations_count[op_type] += 1 + needed_size_bytes += int(operation.get("size_bytes") or 0) + continue + + needs_content = False + for slug, root in replica_roots.items(): + try: + target_path = _safe_join(root, sanitized) + except ValueError: + logger.warning( + "[upload_service] Invalid %s path during plan: %s (root=%s)", + op_type, + sanitized, + root, + ) + continue + target_key = _normalize_cache_key_path(str(target_path)) + cached_hash = replica_cache_hashes.get(slug, {}).get(target_key) + # A cache hit is only authoritative while the indexed replica file + # still exists. Otherwise a stale cache entry can suppress repair. + if cached_hash == op_content_hash and target_path.is_file(): + continue + + needs_content = True + if not target_path.is_file(): + diagnostics["missing_targets"] += 1 + else: + diagnostics["hash_mismatches"] += 1 + break - repo_name_for_state: Optional[str] = None + if needs_content: + needed_files[op_type].append(sanitized) + operations_count[op_type] += 1 + needed_size_bytes += int(operation.get("size_bytes") or 0) + else: + operations_count["skipped"] += 1 + operations_count["skipped_hash_match"] += 1 + diagnostics["cache_hash_matches"] += 1 + + diagnostics["needed_content_operations"] = sum( + operations_count[op_type] for op_type in ("created", "updated", "moved") + ) + logger.info( + "[upload_service] Delta plan workspace=%s targets=%s candidates=%d " + "needed=%d skipped_hash_match=%d cache_entries=%d cache_matches=%d", + workspace_path, + list(replica_roots.keys()), + diagnostics["candidate_operations"], + diagnostics["needed_content_operations"], + operations_count["skipped_hash_match"], + diagnostics["cache_entries"], + diagnostics["cache_hash_matches"], + ) - serving_slug: Optional[str] = None - active_slug: Optional[str] = None - if _extract_repo_name_from_path and get_collection_state_snapshot: - try: - repo_name_for_state = _extract_repo_name_from_path(workspace_path) - if repo_name_for_state: - snapshot = get_collection_state_snapshot(workspace_path=None, repo_name=repo_name_for_state) # type: ignore[arg-type] - serving_slug = snapshot.get("serving_repo_slug") - active_slug = snapshot.get("active_repo_slug") - except Exception: - serving_slug = None - active_slug = None - - slug_order: list[str] = [] - serving_candidate: Optional[str] = None - if serving_slug and _SLUGGED_REPO_RE.match(serving_slug): - serving_candidate = serving_slug - if active_slug and _SLUGGED_REPO_RE.match(active_slug) and active_slug not in slug_order: - slug_order.append(active_slug) - - # If staging is active, we must mirror uploads into BOTH the canonical slug and - # the "*_old" slug. Relying purely on snapshot detection is brittle (e.g. when - # the client workspace_path is a host path). When we can infer a canonical slug, - # force both targets. - staging_active = False - staging_gate = bool(is_staging_enabled() if callable(is_staging_enabled) else False) - try: - if serving_slug and str(serving_slug).endswith("_old"): - staging_active = True - except Exception: - staging_active = False + return { + "needed_files": needed_files, + "operation_counts_preview": operations_count, + "needed_size_bytes": needed_size_bytes, + "replica_targets": list(replica_roots.keys()), + "diagnostics": diagnostics, + } - if not staging_gate: - staging_active = False - def _append_slug(slug: Optional[str]) -> None: - if slug and _SLUGGED_REPO_RE.match(slug) and slug not in slug_order: - slug_order.append(slug) +def apply_delta_operations( + workspace_path: str, + operations: list[Dict[str, Any]], + file_hashes: Optional[Dict[str, str]] = None, +) -> Dict[str, int]: + """Apply metadata-only delta operations without requiring a tar bundle.""" + operations_count = { + "created": 0, + "updated": 0, + "deleted": 0, + "moved": 0, + "skipped": 0, + "skipped_hash_match": 0, + "failed": 0, + } + + try: + replica_roots = _resolve_replica_roots(workspace_path) + if not replica_roots: + raise ValueError(f"No replica roots available for workspace: {workspace_path}") + replica_cache_hashes = { + slug: _load_replica_cache_hashes(root, slug) + for slug, root in replica_roots.items() + } + journal_entries_by_slug: Dict[str, list[Dict[str, Any]]] = { + slug: [] for slug in replica_roots.keys() + } + normalized_hashes = { + str(rel_path): _normalize_hash_value(hash_value) + for rel_path, hash_value in (file_hashes or {}).items() + if _normalize_hash_value(hash_value) + } + + for operation in operations: + op_type = str(operation.get("operation") or "") + rel_path = operation.get("path") + + if not rel_path: + operations_count["skipped"] += 1 + continue + + sanitized_path = _sanitize_operation_path(str(rel_path), replica_roots) + if not sanitized_path: + operations_count["skipped"] += 1 + continue + + rel_path = sanitized_path + + if op_type not in {"deleted", "moved"}: + operations_count["failed"] += 1 + continue + + source_rel_path = None + if op_type == "moved": + raw_source = operation.get("source_path") or operation.get("source_relative_path") + if not raw_source: + operations_count["failed"] += 1 + continue + source_rel_path = _sanitize_operation_path(str(raw_source), replica_roots) + if not source_rel_path: + operations_count["failed"] += 1 + continue + + replica_results: Dict[str, str] = {} + for slug, root in replica_roots.items(): + target_path = _safe_join(root, rel_path) + target_key = _normalize_cache_key_path(str(target_path)) + replica_hashes = replica_cache_hashes.setdefault(slug, {}) + op_content_hash = _normalize_hash_value( + operation.get("content_hash") or normalized_hashes.get(rel_path) + ) + + try: + if op_type == "deleted": + if target_path.exists(): + target_path.unlink(missing_ok=True) + _cleanup_empty_dirs(target_path.parent, root) + replica_hashes.pop(target_key, None) + journal_entries_by_slug.setdefault(slug, []).append( + _build_delete_journal_entry(target_path) + ) + replica_results[slug] = "applied" + continue + + safe_source_path = _safe_join(root, source_rel_path or "") + if not safe_source_path.exists(): + if _file_matches_hash(target_path, op_content_hash): + replica_hashes.pop( + _normalize_cache_key_path(str(safe_source_path)), None + ) + replica_hashes[target_key] = op_content_hash + journal_entries_by_slug.setdefault(slug, []).extend( + [ + _build_delete_journal_entry( + safe_source_path, op_content_hash + ), + _build_upsert_journal_entry( + target_path, op_content_hash + ), + ] + ) + replica_results[slug] = "skipped_hash_match" + continue + replica_results[slug] = "failed" + continue + + target_path.parent.mkdir(parents=True, exist_ok=True) + if target_path.exists(): + if target_path.is_dir(): + raise IsADirectoryError( + f"[upload_delta_bundle] move target is a directory: {target_path}" + ) + else: + target_path.unlink() + shutil.move(str(safe_source_path), str(target_path)) + _cleanup_empty_dirs(safe_source_path.parent, root) + source_key = _normalize_cache_key_path(str(safe_source_path)) + moved_hash = replica_hashes.pop(source_key, None) + if op_content_hash: + replica_hashes[target_key] = op_content_hash + elif moved_hash: + replica_hashes[target_key] = moved_hash + move_entry_hash = op_content_hash or moved_hash + journal_entries_by_slug.setdefault(slug, []).extend( + [ + _build_delete_journal_entry(safe_source_path, move_entry_hash), + _build_upsert_journal_entry(target_path, move_entry_hash), + ] + ) + replica_results[slug] = "applied" + except Exception as exc: + logger.debug( + "[upload_service] Failed to apply metadata-only %s to %s in %s: %s", + op_type, + rel_path, + root, + exc, + ) + replica_results[slug] = "failed" - if repo_name_for_state and _SLUGGED_REPO_RE.match(repo_name_for_state): - canonical_slug = repo_name_for_state[:-4] if repo_name_for_state.endswith("_old") else repo_name_for_state - old_slug_candidate = ( - repo_name_for_state if repo_name_for_state.endswith("_old") else f"{canonical_slug}_old" + applied_any = any(result == "applied" for result in replica_results.values()) + skipped_hash_match = bool(replica_results) and all( + result in {"applied", "skipped_hash_match"} + for result in replica_results.values() ) - if staging_active: - slug_order = [] - _append_slug(canonical_slug) - _append_slug(old_slug_candidate) - elif not slug_order: - _append_slug(canonical_slug) - old_slug_path = Path(WORK_DIR) / old_slug_candidate - if old_slug_path.exists(): - _append_slug(old_slug_candidate) - - if not slug_order: - if _SLUGGED_REPO_RE.match(workspace_leaf): - slug_order.append(workspace_leaf) + success_all = skipped_hash_match + if applied_any: + operations_count[op_type] += 1 + if not success_all: + # Keep the operation count as applied for reporting, but + # surface the replica failure so the sequence is retried. + operations_count["failed"] += 1 + logger.debug( + "[upload_service] Partial metadata-only success for %s %s: %s", + op_type, + rel_path, + replica_results, + ) + elif skipped_hash_match: + operations_count["skipped"] += 1 + operations_count["skipped_hash_match"] += 1 else: - if _extract_repo_name_from_path: - repo_name = _extract_repo_name_from_path(workspace_path) or workspace_leaf - else: - repo_name = workspace_leaf - workspace_key = get_workspace_key(workspace_path) - slug_order.append(f"{repo_name}-{workspace_key}") + operations_count["failed"] += 1 - # Best-effort: if staging is active according to workspace_state, ensure we mirror to - # both the canonical slug and its *_old slug. - if staging_gate and (not staging_active) and get_staging_targets and _extract_repo_name_from_path: - try: - repo_name_for_staging = _extract_repo_name_from_path(workspace_path) or slug_order[0] - targets = get_staging_targets(workspace_path=workspace_path, repo_name=repo_name_for_staging) - if isinstance(targets, dict) and targets.get("staging"): - staging_active = True - except Exception as staging_err: - logger.debug(f"[upload_service] Failed to detect staging: {staging_err}") - - def _slug_exists(slug: str) -> bool: - try: - return ( - (Path(WORK_DIR) / slug).exists() - or (Path(WORK_DIR) / ".codebase" / "repos" / slug).exists() - ) - except Exception: - return False - - if staging_gate and (not staging_active) and slug_order: - primary = slug_order[0] - if _SLUGGED_REPO_RE.match(primary): - canonical = primary[:-4] if primary.endswith("_old") else primary - inferred_old = primary if primary.endswith("_old") else f"{canonical}_old" - if _slug_exists(inferred_old): - staging_active = True - - if staging_gate and staging_active and slug_order: - primary = slug_order[0] - if _SLUGGED_REPO_RE.match(primary): - canonical = primary[:-4] if primary.endswith("_old") else primary - old_slug = primary if primary.endswith("_old") else f"{canonical}_old" - desired = [canonical, old_slug] - slug_order = [s for s in desired if _SLUGGED_REPO_RE.match(s)] - elif staging_gate and not staging_active and serving_candidate: - # Ignore serving slugs when staging is disabled; keep deterministic non-staging writes. - if serving_candidate in slug_order: - slug_order = [s for s in slug_order if s != serving_candidate] - - if staging_gate: - try: - logger.info(f"[upload_service] Delta bundle targets (staging={staging_active}): {slug_order}") - except Exception: - pass + for slug, root in replica_roots.items(): + _enqueue_replica_journal_entries( + workspace_root=root, + slug=slug, + entries=journal_entries_by_slug.get(slug, []), + ) + # Flush updated replica hashes to disk (including empty caches) + replica_hashes = replica_cache_hashes.get(slug, {}) + _flush_replica_cache_hashes(root, slug, replica_hashes) + + return operations_count + except Exception as e: + logger.error(f"Error applying metadata-only delta operations: {e}") + raise - replica_roots: Dict[str, Path] = {} - for slug in slug_order: - path = Path(WORK_DIR) / slug - path.mkdir(parents=True, exist_ok=True) - try: - marker_dir = Path(WORK_DIR) / ".codebase" / "repos" / slug - marker_dir.mkdir(parents=True, exist_ok=True) - (marker_dir / ".ctxce_managed_upload").write_text("1\n") - except Exception: - pass - replica_roots[slug] = path.resolve() - primary_slug = slug_order[0] +def process_delta_bundle(workspace_path: str, bundle_path: Path, manifest: Dict[str, Any]) -> Dict[str, int]: + """Process delta bundle and return operation counts.""" + operations_count = { + "created": 0, + "updated": 0, + "deleted": 0, + "moved": 0, + "skipped": 0, + "skipped_hash_match": 0, + "failed": 0, + } + + try: + replica_roots = _resolve_replica_roots(workspace_path) + if not replica_roots: + raise ValueError(f"No replica roots available for workspace: {workspace_path}") + primary_slug = next(iter(replica_roots)) workspace_root = replica_roots[primary_slug] - def _safe_join(base: Path, rel: str) -> Path: - # SECURITY: Prevent path traversal / absolute-path writes by ensuring the resolved - # candidate path stays within the intended workspace root. - rp = Path(str(rel)) - if str(rp) in {".", ""}: - raise ValueError("Invalid operation path") - if rp.is_absolute(): - raise ValueError(f"Absolute paths are not allowed: {rel}") - base_resolved = base.resolve() - candidate = (base_resolved / rp).resolve() - try: - ok = candidate.is_relative_to(base_resolved) - except Exception: - ok = os.path.commonpath([str(base_resolved), str(candidate)]) == str(base_resolved) - if not ok: - raise ValueError(f"Path escapes workspace: {rel}") - return candidate + def _member_suffix(name: str, marker: str) -> Optional[str]: + idx = name.find(marker) + if idx < 0: + return None + suffix = name[idx + len(marker):] + return suffix or None with tarfile.open(bundle_path, "r:gz") as tar: ops_member = None - for member in tar.getnames(): - if member.endswith("metadata/operations.json"): + hashes_member = None + git_member = None + created_members: Dict[str, tarfile.TarInfo] = {} + updated_members: Dict[str, tarfile.TarInfo] = {} + moved_members: Dict[str, tarfile.TarInfo] = {} + for member in tar.getmembers(): + name = member.name + if name.endswith("metadata/operations.json"): ops_member = member - break + continue + if name.endswith("metadata/hashes.json"): + hashes_member = member + continue + if name.endswith("metadata/git_history.json"): + git_member = member + continue + created_rel = _member_suffix(name, "files/created/") + if created_rel: + created_members[created_rel] = member + continue + updated_rel = _member_suffix(name, "files/updated/") + if updated_rel: + updated_members[updated_rel] = member + continue + moved_rel = _member_suffix(name, "files/moved/") + if moved_rel: + moved_members[moved_rel] = member if not ops_member: raise ValueError("operations.json not found in bundle") @@ -240,14 +755,28 @@ def _safe_join(base: Path, rel: str) -> Path: operations_data = json.loads(ops_file.read().decode("utf-8")) operations = operations_data.get("operations", []) + bundle_hashes: Dict[str, str] = {} + if hashes_member: + hashes_file = tar.extractfile(hashes_member) + if hashes_file: + hashes_data = json.loads(hashes_file.read().decode("utf-8")) + raw_hashes = hashes_data.get("file_hashes", {}) + if isinstance(raw_hashes, dict): + for rel_path, hash_value in raw_hashes.items(): + digest = _normalize_hash_value(hash_value) + if digest: + bundle_hashes[str(rel_path)] = digest + + replica_cache_hashes = { + slug: _load_replica_cache_hashes(root, slug) + for slug, root in replica_roots.items() + } + journal_entries_by_slug: Dict[str, list[Dict[str, Any]]] = { + slug: [] for slug in replica_roots.keys() + } # Best-effort: extract git history metadata for watcher to ingest try: - git_member = None - for member in tar.getnames(): - if member.endswith("metadata/git_history.json"): - git_member = member - break if git_member: git_file = tar.extractfile(git_member) if git_file: @@ -266,11 +795,20 @@ def _safe_join(base: Path, rel: str) -> Path: except Exception as git_err: logger.debug(f"[upload_service] Error extracting git history metadata: {git_err}") - def _apply_operation_to_workspace(workspace_root: Path) -> bool: - """Apply a single file operation to a workspace. Returns True on success.""" - nonlocal operations_count, op_type, rel_path, tar - + def _apply_operation_to_workspace( + slug: str, + workspace_root: Path, + op_type: str, + rel_path: str, + operation: Dict[str, Any], + ) -> str: + """Apply a single file operation to a workspace.""" target_path = _safe_join(workspace_root, rel_path) + target_key = _normalize_cache_key_path(str(target_path)) + replica_hashes = replica_cache_hashes.setdefault(slug, {}) + op_content_hash = _normalize_hash_value( + operation.get("content_hash") or bundle_hashes.get(rel_path) + ) safe_source_path = None source_rel_path = None @@ -281,76 +819,128 @@ def _apply_operation_to_workspace(workspace_root: Path) -> bool: try: if op_type == "created": - file_member = None - for member in tar.getnames(): - if member.endswith(f"files/created/{rel_path}"): - file_member = member - break - + if op_content_hash and target_path.is_file(): + cached_hash = replica_hashes.get(target_key) + if cached_hash and cached_hash == op_content_hash: + return "skipped_hash_match" + file_member = created_members.get(rel_path) if file_member: file_content = tar.extractfile(file_member) if file_content: target_path.parent.mkdir(parents=True, exist_ok=True) target_path.write_bytes(file_content.read()) - return True + if op_content_hash: + replica_hashes[target_key] = op_content_hash + journal_entries_by_slug.setdefault(slug, []).append( + _build_upsert_journal_entry(target_path, op_content_hash) + ) + return "applied" else: - return False + return "failed" else: - return False + return "failed" elif op_type == "updated": - file_member = None - for member in tar.getnames(): - if member.endswith(f"files/updated/{rel_path}"): - file_member = member - break - + if op_content_hash and target_path.is_file(): + cached_hash = replica_hashes.get(target_key) + if cached_hash and cached_hash == op_content_hash: + return "skipped_hash_match" + file_member = updated_members.get(rel_path) if file_member: file_content = tar.extractfile(file_member) if file_content: target_path.parent.mkdir(parents=True, exist_ok=True) target_path.write_bytes(file_content.read()) - return True + if op_content_hash: + replica_hashes[target_key] = op_content_hash + journal_entries_by_slug.setdefault(slug, []).append( + _build_upsert_journal_entry(target_path, op_content_hash) + ) + return "applied" else: - return False + return "failed" else: - return False + return "failed" elif op_type == "deleted": if target_path.exists(): target_path.unlink(missing_ok=True) - return True - else: - return True # Already deleted + _cleanup_empty_dirs(target_path.parent, workspace_root) + replica_hashes.pop(target_key, None) + journal_entries_by_slug.setdefault(slug, []).append( + _build_delete_journal_entry(target_path) + ) + return "applied" elif op_type == "moved": if safe_source_path and safe_source_path.exists(): target_path.parent.mkdir(parents=True, exist_ok=True) - safe_source_path.rename(target_path) - return True + if target_path.exists(): + if target_path.is_dir(): + raise IsADirectoryError( + f"[upload_service] move target is a directory: {target_path}" + ) + else: + target_path.unlink() + shutil.move(str(safe_source_path), str(target_path)) + _cleanup_empty_dirs(safe_source_path.parent, workspace_root) + source_key = _normalize_cache_key_path(str(safe_source_path)) + moved_hash = replica_hashes.pop(source_key, None) + if op_content_hash: + replica_hashes[target_key] = op_content_hash + elif moved_hash: + replica_hashes[target_key] = moved_hash + move_entry_hash = op_content_hash or moved_hash + journal_entries_by_slug.setdefault(slug, []).extend( + [ + _build_delete_journal_entry(safe_source_path, move_entry_hash), + _build_upsert_journal_entry(target_path, move_entry_hash), + ] + ) + return "applied" + if _file_matches_hash(target_path, op_content_hash): + replica_hashes[target_key] = op_content_hash + if safe_source_path: + replica_hashes.pop( + _normalize_cache_key_path(str(safe_source_path)), None + ) + journal_entries_by_slug.setdefault(slug, []).append( + _build_delete_journal_entry( + safe_source_path, op_content_hash + ) + ) + journal_entries_by_slug.setdefault(slug, []).append( + _build_upsert_journal_entry(target_path, op_content_hash) + ) + return "skipped_hash_match" # Remote uploads may not have the source file on the server (e.g. staging # mirrors). In that case, clients can embed the destination content under # files/moved/. - file_member = None - for member in tar.getnames(): - if member.endswith(f"files/moved/{rel_path}"): - file_member = member - break + file_member = moved_members.get(rel_path) if file_member: file_content = tar.extractfile(file_member) if file_content: target_path.parent.mkdir(parents=True, exist_ok=True) target_path.write_bytes(file_content.read()) - return True - return False - return False + if op_content_hash: + replica_hashes[target_key] = op_content_hash + if safe_source_path: + journal_entries_by_slug.setdefault(slug, []).append( + _build_delete_journal_entry(safe_source_path, op_content_hash) + ) + journal_entries_by_slug.setdefault(slug, []).append( + _build_upsert_journal_entry(target_path, op_content_hash) + ) + return "applied" + return "failed" + return "failed" else: logger.warning(f"[upload_service] Unknown operation type: {op_type}") - return False + return "failed" except Exception as e: logger.debug(f"[upload_service] Failed to apply {op_type} to {rel_path} in {workspace_root}: {e}") - return False + return "failed" for operation in operations: op_type = operation.get("operation") @@ -360,18 +950,8 @@ def _apply_operation_to_workspace(workspace_root: Path) -> bool: operations_count["skipped"] += 1 continue - sanitized_path = rel_path - skipped_due_to_exact_slug = False - for slug in replica_roots.keys(): - if sanitized_path == slug: - skipped_due_to_exact_slug = True - break - prefix = f"{slug}/" - if sanitized_path.startswith(prefix): - sanitized_path = sanitized_path[len(prefix):] - break - - if skipped_due_to_exact_slug or not sanitized_path: + sanitized_path = _sanitize_operation_path(str(rel_path), replica_roots) + if not sanitized_path: logger.debug( f"[upload_service] Skipping operation {op_type} for path {rel_path}: " "appears to reference slug root directly.", @@ -381,22 +961,47 @@ def _apply_operation_to_workspace(workspace_root: Path) -> bool: rel_path = sanitized_path - replica_results: Dict[str, bool] = {} + replica_results: Dict[str, str] = {} for slug, root in replica_roots.items(): - replica_results[slug] = _apply_operation_to_workspace(root) + replica_results[slug] = _apply_operation_to_workspace( + slug, + root, + op_type, + rel_path, + operation, + ) - success_any = any(replica_results.values()) - success_all = all(replica_results.values()) - if success_any: + applied_any = any(result == "applied" for result in replica_results.values()) + skipped_hash_match = bool(replica_results) and all( + result == "skipped_hash_match" for result in replica_results.values() + ) + success_all = all(result in {"applied", "skipped_hash_match"} for result in replica_results.values()) + if applied_any: operations_count.setdefault(op_type, 0) - operations_count[op_type] = operations_count.get(op_type, 0) + 1 + operations_count[op_type] += 1 if not success_all: + # A retry can skip replicas that already match while + # repairing the replica that failed this attempt. + operations_count["failed"] += 1 logger.debug( f"[upload_service] Partial success for {op_type} {rel_path}: {replica_results}" ) + elif skipped_hash_match: + operations_count["skipped"] += 1 + operations_count["skipped_hash_match"] += 1 else: operations_count["failed"] += 1 + for slug, root in replica_roots.items(): + _enqueue_replica_journal_entries( + workspace_root=root, + slug=slug, + entries=journal_entries_by_slug.get(slug, []), + ) + # Flush updated replica hashes to disk (including empty caches) + replica_hashes = replica_cache_hashes.get(slug, {}) + _flush_replica_cache_hashes(root, slug, replica_hashes) + return operations_count except Exception as e: diff --git a/scripts/upload_service.py b/scripts/upload_service.py index 6771d652..0db6e67f 100644 --- a/scripts/upload_service.py +++ b/scripts/upload_service.py @@ -45,21 +45,13 @@ from fastapi.responses import JSONResponse, RedirectResponse from fastapi.middleware.cors import CORSMiddleware -from scripts.upload_delta_bundle import get_workspace_key, process_delta_bundle - -from scripts.indexing_admin import ( - build_admin_collections_view, - resolve_collection_root, - spawn_ingest_code, - recreate_collection_qdrant, +from scripts.upload_delta_bundle import ( + apply_delta_operations, + get_workspace_key, + plan_delta_upload, + process_delta_bundle, ) -try: - from scripts.workspace_state import is_staging_enabled -except Exception: - is_staging_enabled = None # type: ignore - - from pydantic import BaseModel, Field from scripts.auth_backend import ( AuthDisabledError, @@ -81,74 +73,32 @@ revoke_collection_access, ) -try: - from scripts.collection_admin import delete_collection_everywhere, copy_collection_qdrant -except Exception: - delete_collection_everywhere = None - copy_collection_qdrant = None -try: - from scripts.admin_ui import ( - render_admin_acl, - render_admin_bootstrap, - render_admin_error, - render_admin_login, - ) -except Exception: - - def _admin_ui_unavailable(*args, **kwargs): - raise HTTPException(status_code=500, detail="Admin UI unavailable") - - render_admin_acl = _admin_ui_unavailable - render_admin_bootstrap = _admin_ui_unavailable - render_admin_error = _admin_ui_unavailable - render_admin_login = _admin_ui_unavailable - -# Import staging/indexing admin helpers -try: - from scripts.indexing_admin import ( - start_staging_rebuild, - activate_staging_rebuild, - abort_staging_rebuild, - ) -except ImportError: - start_staging_rebuild = None # type: ignore - activate_staging_rebuild = None # type: ignore - abort_staging_rebuild = None # type: ignore +from scripts.admin_ui import ( + render_admin_acl, + render_admin_bootstrap, + render_admin_error, + render_admin_login, +) -# Import existing workspace state and indexing functions -try: - from scripts.workspace_state import ( - log_activity, - get_collection_name, - get_cached_file_hash, - set_cached_file_hash, - _extract_repo_name_from_path, - update_repo_origin, - get_collection_mappings, - find_collection_for_logical_repo, - update_workspace_state, - set_staging_state, - update_staging_status, - clear_staging_collection, - logical_repo_reuse_enabled, - get_collection_state_snapshot, - ) -except ImportError: - # Fallback for testing without full environment - log_activity = None - get_collection_name = None - get_cached_file_hash = None - set_cached_file_hash = None - _extract_repo_name_from_path = None - update_repo_origin = None - get_collection_mappings = None - find_collection_for_logical_repo = None - update_workspace_state = None - set_staging_state = None - update_staging_status = None - clear_staging_collection = None - def logical_repo_reuse_enabled() -> bool: # type: ignore[no-redef] - return False +from scripts.workspace_state import ( + is_staging_enabled, + log_activity, + get_collection_name, + get_cached_file_hash, + set_cached_file_hash, + _extract_repo_name_from_path, + update_repo_origin, + get_collection_mappings, + find_collection_for_logical_repo, + update_workspace_state, + set_staging_state, + update_staging_status, + clear_staging_collection, + clear_index_journal_entries, + get_index_journal_summary, + logical_repo_reuse_enabled, + get_collection_state_snapshot, +) # Configure logging @@ -182,6 +132,15 @@ def logical_repo_reuse_enabled() -> bool: # type: ignore[no-redef] ) BRIDGE_STATE_TOKEN = (os.environ.get("CTXCE_BRIDGE_STATE_TOKEN") or "").strip() + +# Admin collection operations import Qdrant clients, subprocess indexing helpers, +# and collection-copy/delete wiring. Keep those off the upload/status import path. +def _indexing_admin(): + from scripts import indexing_admin + + return indexing_admin + + # FastAPI app app = FastAPI( title="Context-Engine Delta Upload Service", @@ -200,6 +159,7 @@ def logical_repo_reuse_enabled() -> bool: # type: ignore[no-redef] # In-memory sequence tracking (in production, use persistent storage) _sequence_tracker: Dict[str, int] = {} +_upload_result_tracker: Dict[str, Dict[str, Any]] = {} def _int_env(name: str, default: int) -> int: @@ -224,6 +184,40 @@ class UploadResponse(BaseModel): next_sequence: Optional[int] = None error: Optional[Dict[str, Any]] = None + +class PlanRequest(BaseModel): + workspace_path: str + collection_name: Optional[str] = None + source_path: Optional[str] = None + logical_repo_id: Optional[str] = None + session: Optional[str] = None + manifest: Dict[str, Any] = Field(default_factory=dict) + operations: List[Dict[str, Any]] = Field(default_factory=list) + file_hashes: Dict[str, str] = Field(default_factory=dict) + + +class PlanResponse(BaseModel): + success: bool + workspace_path: str + needed_files: Dict[str, List[str]] + operation_counts_preview: Dict[str, int] + needed_size_bytes: int + replica_targets: List[str] + diagnostics: Dict[str, Any] = Field(default_factory=dict) + fallback_used: bool = False + error: Optional[Dict[str, Any]] = None + + +class ApplyOperationsRequest(BaseModel): + workspace_path: str + collection_name: Optional[str] = None + source_path: Optional[str] = None + logical_repo_id: Optional[str] = None + session: Optional[str] = None + manifest: Dict[str, Any] = Field(default_factory=dict) + operations: List[Dict[str, Any]] = Field(default_factory=list) + file_hashes: Dict[str, str] = Field(default_factory=dict) + class StatusResponse(BaseModel): workspace_path: str collection_name: str @@ -407,9 +401,8 @@ def _resolve_bridge_state_target( repo = (repo_name or "").strip() or None if collection: - if resolve_collection_root is None: - raise HTTPException(status_code=400, detail="collection mapping unavailable") - root, resolved_repo = resolve_collection_root(collection=collection, work_dir=WORK_DIR) + indexing_admin = _indexing_admin() + root, resolved_repo = indexing_admin.resolve_collection_root(collection=collection, work_dir=WORK_DIR) if not root: raise HTTPException(status_code=404, detail="collection mapping not found") workspace_path = root @@ -480,34 +473,83 @@ async def _process_bundle_background( sequence_number: Optional[int], bundle_id: Optional[str], ) -> None: + key = get_workspace_key(workspace_path) try: start_time = datetime.now() + _upload_result_tracker[key] = { + "workspace_path": workspace_path, + "bundle_id": bundle_id, + "sequence_number": sequence_number, + "processed_operations": None, + "processing_time_ms": None, + "status": "processing", + "completed_at": None, + } operations_count = await asyncio.to_thread( process_delta_bundle, workspace_path, bundle_path, manifest ) - if sequence_number is not None: - key = get_workspace_key(workspace_path) - _sequence_tracker[key] = sequence_number - if log_activity: - try: - repo = _extract_repo_name_from_path(workspace_path) if _extract_repo_name_from_path else None - log_activity( - repo_name=repo, - action="uploaded", - file_path=bundle_id, - details={ - "bundle_id": bundle_id, - "operations": operations_count, - "source": "delta_upload", - }, - ) - except Exception as activity_err: - logger.debug(f"[upload_service] Failed to log activity for bundle {bundle_id}: {activity_err}") - processing_time = (datetime.now() - start_time).total_seconds() * 1000 - logger.info( - f"[upload_service] Finished processing bundle {bundle_id} seq {sequence_number} in {int(processing_time)}ms" + processing_time = int((datetime.now() - start_time).total_seconds() * 1000) + failed_count = int((operations_count or {}).get("failed") or 0) + applied_count = int( + (operations_count or {}).get("created", 0) + + (operations_count or {}).get("updated", 0) + + (operations_count or {}).get("deleted", 0) + + (operations_count or {}).get("moved", 0) ) + status_value = "completed" if failed_count == 0 else "failed" + if sequence_number is not None and failed_count == 0: + _sequence_tracker[key] = sequence_number + _upload_result_tracker[key] = { + "workspace_path": workspace_path, + "bundle_id": bundle_id, + "sequence_number": sequence_number, + "processed_operations": operations_count, + "processing_time_ms": processing_time, + "status": status_value, + "failed_count": failed_count, + "partial": bool(failed_count > 0 and applied_count > 0), + "completed_at": datetime.now().isoformat(), + } + try: + repo = _extract_repo_name_from_path(workspace_path) + log_activity( + repo_name=repo, + action="uploaded", + file_path=bundle_id, + details={ + "bundle_id": bundle_id, + "operations": operations_count, + "source": "delta_upload", + }, + ) + except Exception as activity_err: + logger.debug(f"[upload_service] Failed to log activity for bundle {bundle_id}: {activity_err}") + if failed_count > 0: + logger.warning( + "[upload_service] Finished processing bundle %s seq %s with failures in %sms " + "failed=%d ops=%s", + bundle_id, + sequence_number, + processing_time, + failed_count, + operations_count, + ) + else: + logger.info( + f"[upload_service] Finished processing bundle {bundle_id} seq {sequence_number} " + f"in {processing_time}ms ops={operations_count}" + ) except Exception as e: + _upload_result_tracker[key] = { + "workspace_path": workspace_path, + "bundle_id": bundle_id, + "sequence_number": sequence_number, + "processed_operations": None, + "processing_time_ms": None, + "status": "error", + "completed_at": datetime.now().isoformat(), + "error": str(e), + } logger.error(f"[upload_service] Error in background processing for bundle {bundle_id}: {e}") finally: try: @@ -707,8 +749,9 @@ async def admin_acl_page(request: Request): logger.error(f"[upload_service] Failed to load admin UI data: {e}") raise HTTPException(status_code=500, detail="Failed to load admin data") + indexing_admin = _indexing_admin() enriched = await asyncio.to_thread( - build_admin_collections_view, collections=collections, work_dir=WORK_DIR + indexing_admin.build_admin_collections_view, collections=collections, work_dir=WORK_DIR ) resp = render_admin_acl( @@ -753,8 +796,9 @@ async def admin_collections_status(request: Request): except Exception: raise HTTPException(status_code=500, detail="Failed to load collections") + indexing_admin = _indexing_admin() enriched = await asyncio.to_thread( - lambda: build_admin_collections_view(collections=collections, work_dir=WORK_DIR) + lambda: indexing_admin.build_admin_collections_view(collections=collections, work_dir=WORK_DIR) ) return JSONResponse({"collections": enriched}) @@ -774,7 +818,8 @@ async def admin_reindex_collection( back_href="/admin/acl", ) - root, repo_name = resolve_collection_root(collection=name, work_dir=WORK_DIR) + indexing_admin = _indexing_admin() + root, repo_name = indexing_admin.resolve_collection_root(collection=name, work_dir=WORK_DIR) if not root: return render_admin_error( request, @@ -784,7 +829,7 @@ async def admin_reindex_collection( ) try: - spawn_ingest_code( + indexing_admin.spawn_ingest_code( root=root, work_dir=WORK_DIR, collection=name, @@ -810,9 +855,6 @@ async def bridge_collection_state( workspace: Optional[str] = None, repo_name: Optional[str] = None, ): - if get_collection_state_snapshot is None: - raise HTTPException(status_code=503, detail="workspace_state helper unavailable") - _bridge_state_authorized(request) workspace_path, repo = _resolve_bridge_state_target(collection=collection, workspace=workspace, repo_name=repo_name) @@ -821,7 +863,7 @@ async def bridge_collection_state( if not snapshot: raise HTTPException(status_code=404, detail="Workspace state not found") - if not (is_staging_enabled() if callable(is_staging_enabled) else False): + if not is_staging_enabled(): # Classic mode: ignore any serving_* overrides from staging/migration. snapshot = dict(snapshot) snapshot.pop("serving_collection", None) @@ -856,7 +898,8 @@ async def admin_recreate_collection( back_href="/admin/acl", ) - root, repo_name = resolve_collection_root(collection=name, work_dir=WORK_DIR) + indexing_admin = _indexing_admin() + root, repo_name = indexing_admin.resolve_collection_root(collection=name, work_dir=WORK_DIR) if not root: return render_admin_error( request, @@ -866,12 +909,12 @@ async def admin_recreate_collection( ) try: - recreate_collection_qdrant( + indexing_admin.recreate_collection_qdrant( qdrant_url=QDRANT_URL, api_key=os.environ.get("QDRANT_API_KEY") or None, collection=name, ) - spawn_ingest_code( + indexing_admin.spawn_ingest_code( root=root, work_dir=WORK_DIR, collection=name, @@ -890,6 +933,50 @@ async def admin_recreate_collection( return RedirectResponse(url="/admin/acl", status_code=302) +@app.post("/admin/collections/clear-journal") +async def admin_clear_collection_journal( + request: Request, + collection: str = Form(...), +): + _require_admin_session(request) + name = (collection or "").strip() + if not name: + return render_admin_error( + request, + title="Clear Journal Failed", + message="collection is required", + back_href="/admin/acl", + ) + + indexing_admin = _indexing_admin() + root, repo_name = indexing_admin.resolve_collection_root(collection=name, work_dir=WORK_DIR) + if not root: + return render_admin_error( + request, + title="Clear Journal Failed", + message="No workspace mapping found for collection", + back_href="/admin/acl", + ) + + try: + removed = clear_index_journal_entries(workspace_path=root, repo_name=repo_name) + except Exception as e: + return render_admin_error( + request, + title="Clear Journal Failed", + message=str(e), + back_href="/admin/acl", + ) + + try: + from urllib.parse import urlencode + + url = "/admin/acl?" + urlencode({"journal_cleared": name, "journal_removed": str(removed)}) + except Exception: + url = "/admin/acl" + return RedirectResponse(url=url, status_code=302) + + @app.post("/admin/collections/delete") async def admin_delete_collection( request: Request, @@ -920,14 +1007,6 @@ async def admin_delete_collection( back_href="/admin/acl", ) - if delete_collection_everywhere is None: - return render_admin_error( - request, - title="Delete Collection Failed", - message="Collection delete helper unavailable", - back_href="/admin/acl", - ) - # Default is Qdrant-only (no filesystem cleanup). Users must explicitly opt in. try: cleanup_fs = (delete_fs or "").strip().lower() in {"1", "true", "yes", "on"} @@ -935,7 +1014,10 @@ async def admin_delete_collection( cleanup_fs = False try: - delete_collection_everywhere( + # Collection deletion imports Qdrant admin helpers only on the admin route. + from scripts.collection_admin import delete_collection_everywhere + + out = delete_collection_everywhere( collection=name, work_dir=WORK_DIR, qdrant_url=QDRANT_URL, @@ -949,7 +1031,23 @@ async def admin_delete_collection( back_href="/admin/acl", ) - return RedirectResponse(url="/admin/acl", status_code=302) + graph_deleted: Optional[str] = None + try: + if isinstance(out, dict) and not name.endswith("_graph"): + graph_deleted = "1" if bool(out.get("qdrant_graph_deleted")) else "0" + except Exception: + graph_deleted = None + + try: + from urllib.parse import urlencode + + params = {"deleted": name} + if graph_deleted is not None: + params["graph_deleted"] = graph_deleted + url = "/admin/acl?" + urlencode(params) + except Exception: + url = "/admin/acl" + return RedirectResponse(url=url, status_code=302) @app.post("/admin/staging/start") @@ -958,7 +1056,7 @@ async def admin_start_staging( collection: str = Form(...), ): _require_admin_session(request) - if not (is_staging_enabled() if callable(is_staging_enabled) else False): + if not is_staging_enabled(): return render_admin_error( request, title="Start Staging Failed", @@ -975,18 +1073,11 @@ async def admin_start_staging( back_href="/admin/acl", ) - if start_staging_rebuild is None: - return render_admin_error( - request, - title="Start Staging Failed", - message="Staging helper unavailable", - back_href="/admin/acl", - ) - root: Optional[str] = None repo_name: Optional[str] = None try: - root, repo_name = resolve_collection_root(collection=name, work_dir=WORK_DIR) + indexing_admin = _indexing_admin() + root, repo_name = indexing_admin.resolve_collection_root(collection=name, work_dir=WORK_DIR) except Exception: root, repo_name = None, None if not root: @@ -1014,12 +1105,10 @@ async def admin_start_staging( }, } - if set_staging_state: - try: - set_staging_state(workspace_path=root, repo_name=repo_name, staging=staging_payload) - except Exception as set_err: - logger.warning(f"[admin] Failed to persist queued staging state for {name}: {set_err}") - elif update_workspace_state: + try: + set_staging_state(workspace_path=root, repo_name=repo_name, staging=staging_payload) + except Exception as set_err: + logger.warning(f"[admin] Failed to persist queued staging state for {name}: {set_err}") try: update_workspace_state( workspace_path=root, @@ -1029,30 +1118,29 @@ async def admin_start_staging( except Exception as set_err: logger.warning(f"[admin] Failed to update workspace state for queued staging {name}: {set_err}") - if update_staging_status: - try: - update_staging_status( - workspace_path=root, - repo_name=repo_name, - status={"state": "queued", "queued_at": now, "request_id": request_id}, - ) - except Exception as status_err: - logger.debug(f"[admin] Failed to mark staging status queued for {name}: {status_err}") + try: + update_staging_status( + workspace_path=root, + repo_name=repo_name, + status={"state": "queued", "queued_at": now, "request_id": request_id}, + ) + except Exception as status_err: + logger.debug(f"[admin] Failed to mark staging status queued for {name}: {status_err}") try: async def _bg_start() -> None: try: staging_collection = await asyncio.to_thread( - start_staging_rebuild, collection=name, work_dir=WORK_DIR + indexing_admin.start_staging_rebuild, collection=name, work_dir=WORK_DIR ) logger.info(f"[admin] Started staging rebuild for {name} -> {staging_collection}") except Exception as e: logger.error(f"[admin] Background staging start failed for {name}: {e}") # Ensure we don't leave the workspace stuck in a queued staging state. try: - if clear_staging_collection: + try: clear_staging_collection(workspace_path=root, repo_name=repo_name) - elif update_workspace_state: + except Exception: update_workspace_state( workspace_path=root, repo_name=repo_name, @@ -1081,7 +1169,7 @@ async def admin_activate_staging( collection: str = Form(...), ): _require_admin_session(request) - if not (is_staging_enabled() if callable(is_staging_enabled) else False): + if not is_staging_enabled(): return render_admin_error( request, title="Activate Staging Failed", @@ -1097,18 +1185,11 @@ async def admin_activate_staging( back_href="/admin/acl", ) - if activate_staging_rebuild is None: - return render_admin_error( - request, - title="Activate Staging Failed", - message="Staging helper unavailable", - back_href="/admin/acl", - ) - try: async def _bg_activate() -> None: try: - await asyncio.to_thread(activate_staging_rebuild, collection=name, work_dir=WORK_DIR) + indexing_admin = _indexing_admin() + await asyncio.to_thread(indexing_admin.activate_staging_rebuild, collection=name, work_dir=WORK_DIR) logger.info(f"[admin] Activated staging for {name}") except Exception as e: logger.error(f"[admin] Background staging activate failed for {name}: {e}") @@ -1140,7 +1221,8 @@ async def admin_abort_staging( back_href="/admin/acl", ) - root, repo_name = resolve_collection_root(collection=name, work_dir=WORK_DIR) + indexing_admin = _indexing_admin() + root, repo_name = indexing_admin.resolve_collection_root(collection=name, work_dir=WORK_DIR) if not root: return render_admin_error( request, @@ -1150,21 +1232,13 @@ async def admin_abort_staging( ) try: - if abort_staging_rebuild is not None: - # Run abort to completion so we always clear staging metadata before returning. - await asyncio.to_thread( - abort_staging_rebuild, - collection=name, - work_dir=WORK_DIR, - delete_collection=True, - ) - logger.info(f"[admin] Aborted staging rebuild for {name}") - elif clear_staging_collection: - # Fallback for older deployments: clear staging metadata only. - clear_staging_collection(workspace_path=root, repo_name=repo_name) - logger.info(f"[admin] Aborted staging for {name} (metadata only)") - else: - raise RuntimeError("staging abort helpers unavailable") + await asyncio.to_thread( + indexing_admin.abort_staging_rebuild, + collection=name, + work_dir=WORK_DIR, + delete_collection=True, + ) + logger.info(f"[admin] Aborted staging rebuild for {name}") except Exception as e: return render_admin_error( request, @@ -1193,20 +1267,15 @@ async def admin_copy_collection( back_href="/admin/acl", ) - if copy_collection_qdrant is None: - return render_admin_error( - request, - title="Copy Collection Failed", - message="copy helper unavailable", - back_href="/admin/acl", - ) - try: allow_overwrite = str(overwrite or "").strip().lower() in {"1", "true", "yes", "on"} except Exception: allow_overwrite = False try: + # Collection copy imports Qdrant admin helpers only on the admin route. + from scripts.collection_admin import copy_collection_qdrant + new_name = copy_collection_qdrant( source=name, target=(target or None), @@ -1222,7 +1291,60 @@ async def admin_copy_collection( back_href="/admin/acl", ) - return RedirectResponse(url="/admin/acl", status_code=302) + graph_copied: Optional[str] = None + try: + if not name.endswith("_graph") and not str(new_name).endswith("_graph"): + used_pooled = True + try: + # Qdrant client pool is only needed to verify the copied graph collection. + from scripts.qdrant_client_manager import pooled_qdrant_client + + with pooled_qdrant_client( + url=QDRANT_URL, + api_key=os.environ.get("QDRANT_API_KEY"), + ) as cli: + try: + cli.get_collection(collection_name=f"{new_name}_graph") + graph_copied = "1" + except Exception: + graph_copied = "0" + except Exception: + # Failed to acquire pooled client; fall back to non-pooled + used_pooled = False + if not used_pooled: + try: + from qdrant_client import QdrantClient # type: ignore + + cli = QdrantClient( + url=QDRANT_URL, + api_key=os.environ.get("QDRANT_API_KEY"), + timeout=float(os.environ.get("QDRANT_TIMEOUT", "5") or 5), + ) + try: + cli.get_collection(collection_name=f"{new_name}_graph") + graph_copied = "1" + except Exception: + graph_copied = "0" + finally: + try: + cli.close() + except Exception: + pass + except Exception: + graph_copied = "0" + except Exception: + graph_copied = None + + try: + from urllib.parse import urlencode + + params = {"copied": name, "new": new_name} + if graph_copied is not None: + params["graph_copied"] = graph_copied + url = "/admin/acl?" + urlencode(params) + except Exception: + url = "/admin/acl" + return RedirectResponse(url=url, status_code=302) @app.post("/admin/users") @@ -1347,28 +1469,65 @@ async def get_status(workspace_path: str): """Get upload status for workspace.""" try: # Get collection name - if get_collection_name: - repo_name = _extract_repo_name_from_path(workspace_path) if _extract_repo_name_from_path else None - collection_name = get_collection_name(repo_name) - else: - collection_name = DEFAULT_COLLECTION + try: + is_workspace_root = Path(workspace_path).resolve() == Path(WORK_DIR).resolve() + except Exception: + is_workspace_root = False + repo_name = None if is_workspace_root else _extract_repo_name_from_path(workspace_path) + collection_name = get_collection_name(repo_name) + + try: + journal_summary = get_index_journal_summary( + workspace_path=workspace_path, + repo_name=repo_name or None, + ) + except Exception as journal_exc: + logger.warning( + "[upload_service] Failed to read journal summary for %s: %s", + workspace_path, + journal_exc, + ) + journal_summary = { + "total": 0, + "retryable": 0, + "outstanding": 0, + "exhausted": 0, + "counts": {}, + "sample_errors": [], + } # Get last sequence last_sequence = get_last_sequence(workspace_path) + key = get_workspace_key(workspace_path) + upload_result = _upload_result_tracker.get(key, {}) - last_upload = None + last_upload = upload_result.get("completed_at") + upload_status = str(upload_result.get("status") or "") + workspace_status = "processing" if upload_status == "processing" else "ready" return StatusResponse( workspace_path=workspace_path, collection_name=collection_name, last_sequence=last_sequence, last_upload=last_upload, - pending_operations=0, - status="ready", + pending_operations=int( + journal_summary.get( + "outstanding", + journal_summary.get("retryable", 0), + ) + or 0 + ), + status=workspace_status, server_info={ "version": "1.0.0", "max_bundle_size_mb": MAX_BUNDLE_SIZE_MB, - "supported_formats": ["tar.gz"] + "supported_formats": ["tar.gz"], + "last_bundle_id": upload_result.get("bundle_id"), + "last_processing_time_ms": upload_result.get("processing_time_ms"), + "last_processed_operations": upload_result.get("processed_operations"), + "last_upload_status": upload_status or None, + "last_error": upload_result.get("error"), + "journal": journal_summary, } ) @@ -1376,6 +1535,362 @@ async def get_status(workspace_path: str): logger.error(f"Error getting status: {e}") raise HTTPException(status_code=500, detail=str(e)) + +def _resolve_collection_for_request( + workspace_path: str, + client_collection_name: Optional[str], + logical_repo_id: Optional[str], + source_path: Optional[str] = None, +) -> Tuple[str, Optional[str]]: + """ + Resolve collection name and repo_name for upload/plan/apply requests. + + Returns: + Tuple of (collection_name, repo_name) + """ + # Resolve collection name for ACL enforcement + collection_name: Optional[str] = None + repo_name: Optional[str] = None + + repo_source = (source_path or "").strip() or workspace_path + repo_name = _extract_repo_name_from_path(repo_source) + if not repo_name: + repo_name = Path(repo_source).name + + resolved_collection: Optional[str] = None + + # Resolve collection name, preferring server-side mapping for logical_repo_id when enabled + if logical_repo_reuse_enabled() and logical_repo_id: + try: + existing = find_collection_for_logical_repo(logical_repo_id, search_root=WORK_DIR) + except Exception: + existing = None + if existing: + resolved_collection = existing + + # Latent migration: when no explicit mapping exists yet for this logical_repo_id, but there is a + # single existing collection mapping, prefer reusing it rather than creating a fresh collection. + if logical_repo_reuse_enabled() and logical_repo_id and resolved_collection is None: + try: + mappings = get_collection_mappings(search_root=WORK_DIR) or [] + except Exception: + mappings = [] + + if len(mappings) == 1: + canonical = mappings[0] + canonical_coll = canonical.get("collection_name") + if canonical_coll: + resolved_collection = canonical_coll + try: + update_workspace_state( + workspace_path=canonical.get("container_path") or canonical.get("state_file"), + updates={"logical_repo_id": logical_repo_id}, + repo_name=canonical.get("repo_name"), + ) + except Exception as migrate_err: + logger.debug( + f"[upload_service] Failed to migrate logical_repo_id for existing mapping: {migrate_err}" + ) + + # Upload-managed requests are server-owned; ignore client-supplied collection routing. + if resolved_collection is not None: + collection_name = resolved_collection + else: + collection_name = get_collection_name(repo_name) if repo_name else DEFAULT_COLLECTION + + return collection_name, repo_name + + +@app.post("/api/v1/delta/plan", response_model=PlanResponse) +async def plan_delta(request: PlanRequest): + """Plan which file bodies are needed before uploading content.""" + try: + workspace = Path(request.workspace_path) + if not workspace.is_absolute(): + workspace = Path(WORK_DIR) / workspace + workspace_path = str(workspace.resolve()) + + if AUTH_ENABLED: + session_value = str(request.session or "").strip() + try: + record = validate_session(session_value) + except AuthDisabledError: + record = None + except Exception as e: + logger.error(f"[upload_service] Failed to validate auth session for plan: {e}") + raise HTTPException(status_code=500, detail="Failed to validate auth session") + if record is None: + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail="Invalid or expired session", + ) + + # Resolve collection name for ACL enforcement + collection_name, repo_name = _resolve_collection_for_request( + workspace_path=workspace_path, + client_collection_name=request.collection_name, + logical_repo_id=request.logical_repo_id, + source_path=request.source_path, + ) + + # Enforce collection write access for plan/apply when auth is enabled + if AUTH_ENABLED and CTXCE_MCP_ACL_ENFORCE: + if not collection_name: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail="Collection resolution failed for ACL enforcement", + ) + uid = str((record or {}).get("user_id") or "").strip() + if not uid: + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail="Invalid or expired session", + ) + try: + allowed = has_collection_access(uid, str(collection_name), "write") + except AuthDisabledError: + allowed = True + if not allowed: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail=f"User does not have write access to collection '{collection_name}'", + ) + + plan = plan_delta_upload( + workspace_path=workspace_path, + operations=request.operations, + file_hashes=request.file_hashes, + ) + return PlanResponse( + success=True, + workspace_path=workspace_path, + needed_files=plan.get("needed_files", {"created": [], "updated": [], "moved": []}), + operation_counts_preview=plan.get( + "operation_counts_preview", + { + "created": 0, + "updated": 0, + "deleted": 0, + "moved": 0, + "skipped": 0, + "skipped_hash_match": 0, + "failed": 0, + }, + ), + needed_size_bytes=int(plan.get("needed_size_bytes", 0) or 0), + replica_targets=list(plan.get("replica_targets", []) or []), + diagnostics=dict(plan.get("diagnostics", {}) or {}), + fallback_used=False, + error=None, + ) + except HTTPException: + raise + except Exception as e: + logger.error(f"[upload_service] Error planning delta upload: {e}") + return PlanResponse( + success=False, + workspace_path=request.workspace_path, + needed_files={"created": [], "updated": [], "moved": []}, + operation_counts_preview={ + "created": 0, + "updated": 0, + "deleted": 0, + "moved": 0, + "skipped": 0, + "skipped_hash_match": 0, + "failed": 0, + }, + needed_size_bytes=0, + replica_targets=[], + diagnostics={}, + fallback_used=True, + error={ + "code": "PLAN_ERROR", + "message": str(e), + }, + ) + + +@app.post("/api/v1/delta/apply_ops", response_model=UploadResponse) +async def apply_delta_ops(request: ApplyOperationsRequest): + """Apply metadata-only delta operations without uploading a tar bundle.""" + key: Optional[str] = None + bundle_id: Optional[str] = None + sequence_number: Optional[int] = None + try: + workspace = Path(request.workspace_path) + if not workspace.is_absolute(): + workspace = Path(WORK_DIR) / workspace + workspace_path = str(workspace.resolve()) + + if AUTH_ENABLED: + session_value = str(request.session or "").strip() + try: + record = validate_session(session_value) + except AuthDisabledError: + record = None + except Exception as e: + logger.error(f"[upload_service] Failed to validate auth session for apply_ops: {e}") + raise HTTPException(status_code=500, detail="Failed to validate auth session") + if record is None: + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail="Invalid or expired session", + ) + + # Resolve collection name for ACL enforcement + collection_name, repo_name = _resolve_collection_for_request( + workspace_path=workspace_path, + client_collection_name=request.collection_name, + logical_repo_id=request.logical_repo_id, + source_path=request.source_path, + ) + + # Enforce collection write access for plan/apply when auth is enabled + if AUTH_ENABLED and CTXCE_MCP_ACL_ENFORCE: + if not collection_name: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail="Collection resolution failed for ACL enforcement", + ) + uid = str((record or {}).get("user_id") or "").strip() + if not uid: + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail="Invalid or expired session", + ) + try: + allowed = has_collection_access(uid, str(collection_name), "write") + except AuthDisabledError: + allowed = True + if not allowed: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail=f"User does not have write access to collection '{collection_name}'", + ) + + manifest = request.manifest or {} + bundle_id = manifest.get("bundle_id") + manifest_sequence = manifest.get("sequence_number") + key = get_workspace_key(workspace_path) + last_sequence = get_last_sequence(workspace_path) + sequence_number = manifest_sequence if manifest_sequence is not None else last_sequence + 1 + + if sequence_number is not None and sequence_number != last_sequence + 1: + return UploadResponse( + success=False, + error={ + "code": "SEQUENCE_MISMATCH", + "message": f"Expected sequence {last_sequence + 1}, got {sequence_number}", + "expected_sequence": last_sequence + 1, + "received_sequence": sequence_number, + "retry_after": 5000, + }, + ) + + start_time = datetime.now() + _upload_result_tracker[key] = { + "workspace_path": workspace_path, + "bundle_id": bundle_id, + "sequence_number": sequence_number, + "processed_operations": None, + "processing_time_ms": None, + "status": "processing", + "completed_at": None, + } + + operations_count = await asyncio.to_thread( + apply_delta_operations, + workspace_path, + request.operations, + request.file_hashes, + ) + processing_time = int((datetime.now() - start_time).total_seconds() * 1000) + failed_count = int((operations_count or {}).get("failed") or 0) + applied_count = int( + (operations_count or {}).get("created", 0) + + (operations_count or {}).get("updated", 0) + + (operations_count or {}).get("deleted", 0) + + (operations_count or {}).get("moved", 0) + ) + status_value = "completed" if failed_count == 0 else "failed" + operation_count = int( + applied_count + (operations_count or {}).get("skipped_hash_match", 0) + ) + if operation_count > 0 and failed_count == 0: + _sequence_tracker[key] = sequence_number + _upload_result_tracker[key] = { + "workspace_path": workspace_path, + "bundle_id": bundle_id, + "sequence_number": sequence_number, + "processed_operations": operations_count, + "processing_time_ms": processing_time, + "status": status_value, + "failed_count": failed_count, + "partial": bool(failed_count > 0 and applied_count > 0), + "completed_at": datetime.now().isoformat(), + } + if failed_count > 0: + logger.warning( + "[upload_service] apply_ops completed with failures bundle=%s seq=%s failed=%d ops=%s", + bundle_id, + sequence_number, + failed_count, + operations_count, + ) + return UploadResponse( + success=False, + bundle_id=bundle_id, + sequence_number=sequence_number, + processed_operations=operations_count, + processing_time_ms=processing_time, + next_sequence=sequence_number + 1 if sequence_number is not None else None, + error={ + "code": "APPLY_OPS_PARTIAL_FAILURE", + "message": f"One or more operations failed during apply_ops (failed={failed_count})", + "failed_count": failed_count, + "processed_operations": operations_count, + }, + ) + logger.info( + "[upload_service] Applied metadata-only operations bundle=%s seq=%s in %sms ops=%s", + bundle_id, + sequence_number, + processing_time, + operations_count, + ) + return UploadResponse( + success=True, + bundle_id=bundle_id, + sequence_number=sequence_number, + processed_operations=operations_count, + processing_time_ms=processing_time, + next_sequence=sequence_number + 1 if sequence_number is not None else None, + ) + except HTTPException: + raise + except Exception as e: + logger.error(f"[upload_service] Error applying metadata-only operations: {e}") + if key: + _upload_result_tracker[key] = { + "workspace_path": request.workspace_path, + "bundle_id": bundle_id, + "sequence_number": sequence_number, + "processed_operations": None, + "processing_time_ms": None, + "status": "error", + "error": str(e), + "message": str(e), + "completed_at": datetime.now().isoformat(), + } + return UploadResponse( + success=False, + error={ + "code": "APPLY_OPS_ERROR", + "message": str(e), + }, + ) + @app.post("/api/v1/delta/upload", response_model=UploadResponse) async def upload_delta_bundle( request: Request, @@ -1422,60 +1937,13 @@ async def upload_delta_bundle( workspace_path = str(workspace.resolve()) - # Always derive repo_name from workspace_path for origin tracking - repo_name = _extract_repo_name_from_path(workspace_path) if _extract_repo_name_from_path else None - if not repo_name: - repo_name = Path(workspace_path).name - - # Preserve any client-supplied collection name but allow server-side overrides - client_collection_name = collection_name - resolved_collection: Optional[str] = None - - # Resolve collection name, preferring server-side mapping for logical_repo_id when enabled - if logical_repo_reuse_enabled() and logical_repo_id and find_collection_for_logical_repo: - try: - existing = find_collection_for_logical_repo(logical_repo_id, search_root=WORK_DIR) - except Exception: - existing = None - if existing: - resolved_collection = existing - - # Latent migration: when no explicit mapping exists yet for this logical_repo_id, but there is a - # single existing collection mapping, prefer reusing it rather than creating a fresh collection. - if logical_repo_reuse_enabled() and logical_repo_id and resolved_collection is None and get_collection_mappings: - try: - mappings = get_collection_mappings(search_root=WORK_DIR) or [] - except Exception: - mappings = [] - - if len(mappings) == 1: - canonical = mappings[0] - canonical_coll = canonical.get("collection_name") - if canonical_coll: - resolved_collection = canonical_coll - if update_workspace_state: - try: - update_workspace_state( - workspace_path=canonical.get("container_path") or canonical.get("state_file"), - updates={"logical_repo_id": logical_repo_id}, - repo_name=canonical.get("repo_name"), - ) - except Exception as migrate_err: - logger.debug( - f"[upload_service] Failed to migrate logical_repo_id for existing mapping: {migrate_err}" - ) - - # Finalize collection_name: prefer resolved server-side mapping, then client-supplied name, - # then standard get_collection_name/DEFAULT_COLLECTION fallbacks. - if resolved_collection is not None: - collection_name = resolved_collection - elif client_collection_name: - collection_name = client_collection_name - else: - if get_collection_name and repo_name: - collection_name = get_collection_name(repo_name) - else: - collection_name = DEFAULT_COLLECTION + # Resolve collection name and repo name + collection_name, repo_name = _resolve_collection_for_request( + workspace_path=workspace_path, + client_collection_name=collection_name, + logical_repo_id=logical_repo_id, + source_path=source_path, + ) # Enforce collection write access for uploads when auth is enabled. # Semantics: "write" is sufficient for uploading/indexing content. diff --git a/scripts/warm_all_collections.py b/scripts/warm_all_collections.py index 0344da82..19f37241 100644 --- a/scripts/warm_all_collections.py +++ b/scripts/warm_all_collections.py @@ -5,6 +5,7 @@ import os import sys import subprocess +from pathlib import Path from qdrant_client import QdrantClient def main(): @@ -12,6 +13,7 @@ def main(): qdrant_url = os.environ.get("QDRANT_URL", "http://qdrant:6333") ef = os.environ.get("EF", "256") limit = os.environ.get("LIMIT", "3") + script_dir = Path(__file__).resolve().parent print(f"Connecting to Qdrant at {qdrant_url}") @@ -37,8 +39,8 @@ def main(): result = subprocess.run( [ - "python", - "/app/scripts/warm_start.py", + sys.executable or "python", + str(script_dir / "warm_start.py"), "--ef", ef, "--limit", limit ], diff --git a/scripts/warm_start.py b/scripts/warm_start.py index 3119721b..a024f33a 100644 --- a/scripts/warm_start.py +++ b/scripts/warm_start.py @@ -1,15 +1,10 @@ #!/usr/bin/env python3 import os import argparse -import sys -from pathlib import Path from qdrant_client import QdrantClient, models -ROOT_DIR = Path(__file__).resolve().parent.parent -if str(ROOT_DIR) not in sys.path: - sys.path.insert(0, str(ROOT_DIR)) - from scripts.utils import sanitize_vector_name +from scripts.embedder import get_embedding_model as _get_model # Warm start: load embedding model and warm Qdrant HNSW search path with a small query # Useful to reduce first-query latency and set a higher runtime ef for quality @@ -21,14 +16,7 @@ def derive_vector_name(model_name: str) -> str: def get_embedding_model(model_name: str): """Get embedding model with Qwen3 support via embedder factory.""" - try: - from scripts.embedder import get_embedding_model as _get_model - return _get_model(model_name) - except ImportError: - pass - # Fallback to direct fastembed - from fastembed import TextEmbedding - return TextEmbedding(model_name=model_name) + return _get_model(model_name) def main(): diff --git a/scripts/watch_index.py b/scripts/watch_index.py index 8fe5a740..442931e4 100644 --- a/scripts/watch_index.py +++ b/scripts/watch_index.py @@ -2,53 +2,52 @@ from __future__ import annotations import os -import sys import time +from collections import Counter from pathlib import Path from typing import Optional from qdrant_client import QdrantClient from watchdog.observers import Observer -ROOT_DIR = Path(__file__).resolve().parent.parent -if str(ROOT_DIR) not in sys.path: - sys.path.insert(0, str(ROOT_DIR)) - -from scripts.watch_index_core.config import ( # noqa: E402 - LOGGER, - MODEL, - QDRANT_URL, - ROOT as WATCH_ROOT, - default_collection_name, -) +from scripts.watch_index_core import config as watch_config +from scripts.watch_index_core.config import LOGGER, MODEL, QDRANT_URL, default_collection_name from scripts.watch_index_core.utils import ( get_boolean_env, resolve_vector_name_config, create_observer, ) -from scripts.watch_index_core.handler import IndexHandler # noqa: E402 -from scripts.watch_index_core.pseudo import _start_pseudo_backfill_worker # noqa: E402 -from scripts.watch_index_core.processor import _process_paths # noqa: E402 -from scripts.watch_index_core.queue import ChangeQueue # noqa: E402 -from scripts.workspace_state import ( # noqa: E402 - _extract_repo_name_from_path, +from scripts.watch_index_core.handler import IndexHandler +from scripts.watch_index_core.init_maintenance import start_init_maintenance_worker +from scripts.watch_index_core.pseudo import _start_pseudo_backfill_worker +from scripts.watch_index_core.processor import _process_paths +from scripts.watch_index_core.queue import ChangeQueue +from scripts.watch_index_core.consistency import ( + run_consistency_audit, + run_empty_dir_sweep_maintenance, +) +from scripts.workspace_state import ( compute_indexing_config_hash, - get_collection_name, get_indexing_config_snapshot, + list_pending_index_journal_entries, is_multi_repo_mode, persist_indexing_config, update_indexing_status, - update_workspace_state, initialize_watcher_state, ) -import scripts.ingest_code as idx # noqa: E402 +_sleep = time.sleep + +import scripts.ingest_code as idx logger = LOGGER -ROOT = WATCH_ROOT +ROOT = watch_config.ROOT # Back-compat: legacy modules/tests expect a module-level COLLECTION constant. # We use a sentinel and a getter to ensure the resolved value is returned. _COLLECTION: Optional[str] = None +_JOURNAL_DRAIN_LAST_LOG = 0.0 +_JOURNAL_DRAIN_LAST_TOTAL = 0 +_JOURNAL_DRAIN_LAST_BUSY_LOG = 0.0 def get_collection() -> str: @@ -58,7 +57,179 @@ def get_collection() -> str: return default_collection_name() +def _set_runtime_root() -> None: + global ROOT + runtime_root = Path( + os.environ.get("WATCH_ROOT") + or os.environ.get("WORKSPACE_PATH") + or str(ROOT) + ) + try: + runtime_root = runtime_root.resolve() + except Exception: + pass + + ROOT = runtime_root + watch_config.ROOT = runtime_root + + +def _journal_log_interval_secs() -> float: + try: + return max(0.0, float(os.environ.get("WATCH_JOURNAL_LOG_INTERVAL_SECS", "120") or 120.0)) + except Exception: + return 120.0 + + +def _journal_drain_batch_size() -> int: + try: + return max(1, int(os.environ.get("WATCH_JOURNAL_DRAIN_BATCH_SIZE", "256") or 256)) + except Exception: + return 256 + + +def _maybe_log_journal_drain_busy(queue: ChangeQueue) -> None: + global _JOURNAL_DRAIN_LAST_BUSY_LOG + now = time.time() + if _JOURNAL_DRAIN_LAST_BUSY_LOG > 0 and ( + now - _JOURNAL_DRAIN_LAST_BUSY_LOG + ) < _journal_log_interval_secs(): + return + try: + queue_stats = queue.stats() + except Exception: + queue_stats = {} + logger.info( + "watch_index::journal_drain_busy last_backlog=%d queue=%s", + _JOURNAL_DRAIN_LAST_TOTAL, + queue_stats, + extra={ + "root": str(ROOT), + "last_backlog": _JOURNAL_DRAIN_LAST_TOTAL, + "queue_stats": queue_stats, + }, + ) + _JOURNAL_DRAIN_LAST_BUSY_LOG = now + + +def _maybe_log_journal_drain( + *, + total: int, + queued: int, + op_counts: Counter[str], + queue: ChangeQueue, + batch_size: int, +) -> None: + global _JOURNAL_DRAIN_LAST_LOG, _JOURNAL_DRAIN_LAST_TOTAL + now = time.time() + interval = _journal_log_interval_secs() + should_log = False + if total <= 0 and _JOURNAL_DRAIN_LAST_TOTAL > 0: + should_log = True + elif total > 0 and (_JOURNAL_DRAIN_LAST_LOG <= 0 or (now - _JOURNAL_DRAIN_LAST_LOG) >= interval): + should_log = True + if not should_log: + _JOURNAL_DRAIN_LAST_TOTAL = total + return + + queue_stats = {} + try: + queue_stats = queue.stats() + except Exception: + queue_stats = {} + logger.info( + "watch_index::journal_drain backlog=%d batch_limit=%d queued_batch=%d " + "ops=%s queue=%s", + total, + batch_size, + queued, + dict(op_counts), + queue_stats, + extra={ + "root": str(ROOT), + "backlog": total, + "batch_limit": batch_size, + "queued": queued, + "op_counts": dict(op_counts), + "queue_stats": queue_stats, + }, + ) + _JOURNAL_DRAIN_LAST_LOG = now + _JOURNAL_DRAIN_LAST_TOTAL = total + + +def _drain_pending_journal(queue: ChangeQueue) -> None: + pending_path: Optional[str] = None + try: + queue_stats = queue.stats() + if ( + queue_stats.get("processing") + or int(queue_stats.get("queued", 0) or 0) > 0 + or int(queue_stats.get("pending", 0) or 0) > 0 + ): + _maybe_log_journal_drain_busy(queue) + return + + pending_entries = list_pending_index_journal_entries(str(ROOT)) + batch_size = _journal_drain_batch_size() + batch_entries = pending_entries[:batch_size] + queued = 0 + op_counts: Counter[str] = Counter() + for pending_entry in batch_entries: + op_type = str(pending_entry.get("op_type") or "unknown").strip() or "unknown" + op_counts[op_type] += 1 + pending_path = str(pending_entry.get("path") or "").strip() + if pending_path: + queue.add(Path(pending_path), force=True) + queued += 1 + _maybe_log_journal_drain( + total=len(pending_entries), + queued=queued, + op_counts=op_counts, + queue=queue, + batch_size=batch_size, + ) + except Exception as exc: + logger.exception( + "watch_index::pending_journal_drain_failed", + extra={"root": str(ROOT), "pending_path": pending_path, "error": str(exc)}, + ) + + +def _run_periodic_maintenance(client: QdrantClient) -> None: + try: + run_consistency_audit(client, ROOT) + except Exception as exc: + logger.exception( + "watch_index::consistency_audit_failed", + extra={"root": str(ROOT), "error": str(exc)}, + ) + try: + run_empty_dir_sweep_maintenance(ROOT) + except Exception as exc: + logger.exception( + "watch_index::empty_dir_sweep_failed", + extra={"root": str(ROOT), "error": str(exc)}, + ) + + +def _maintenance_interval_secs() -> float: + try: + return max(0.0, float(os.environ.get("WATCH_MAINTENANCE_INTERVAL_SECS", "300") or 300.0)) + except Exception: + return 300.0 + + +def _journal_drain_enabled(multi_repo_enabled: bool) -> bool: + return get_boolean_env("WATCH_JOURNAL_DRAIN_ENABLED", default=multi_repo_enabled) + + +def _fs_events_enabled(multi_repo_enabled: bool) -> bool: + return get_boolean_env("WATCH_FS_EVENTS_ENABLED", default=(not multi_repo_enabled)) + + def main() -> None: + _set_runtime_root() + # Resolve collection name from workspace state before any client/state ops try: from scripts.workspace_state import get_collection_name_with_staging as _get_coll @@ -94,6 +265,16 @@ def main() -> None: f"Watch mode: root={ROOT} qdrant={QDRANT_URL} collection={default_collection} model={MODEL}" ) + # Guardrail: deferring pseudo to a worker only makes sense if the worker is enabled. + # Otherwise you'd silently disable pseudo generation (old behavior). + pseudo_defer = get_boolean_env("PSEUDO_DEFER_TO_WORKER") + pseudo_backfill_enabled = get_boolean_env("PSEUDO_BACKFILL_ENABLED") + if pseudo_defer and not pseudo_backfill_enabled: + print( + "[pseudo] Warning: PSEUDO_DEFER_TO_WORKER=1 but PSEUDO_BACKFILL_ENABLED=0; " + "inline pseudo will remain enabled (no deferral)." + ) + # Health check: detect and auto-heal cache/collection sync issues. # In multi-repo mode this can be expensive and may duplicate external init checks, # so default it OFF unless explicitly enabled. @@ -137,18 +318,10 @@ def main() -> None: url=QDRANT_URL, timeout=int(os.environ.get("QDRANT_TIMEOUT", "20") or 20) ) - # Use centralized embedder factory if available (supports Qwen3 feature flag) - try: - from scripts.embedder import get_embedding_model, get_model_dimension - - model = get_embedding_model(MODEL) - model_dim = get_model_dimension(MODEL) - except ImportError: - # Fallback to direct fastembed initialization - from fastembed import TextEmbedding + from scripts.embedder import get_embedding_model, get_model_dimension - model = TextEmbedding(model_name=MODEL) - model_dim = len(next(model.embed(["dimension probe"]))) + model = get_embedding_model(MODEL) + model_dim = get_model_dimension(MODEL) vector_name = resolve_vector_name_config(client, default_collection, model_dim, MODEL) @@ -166,7 +339,17 @@ def main() -> None: except Exception: pass - _start_pseudo_backfill_worker(client, default_collection, model_dim, vector_name) + # Start backfill worker even in multi-repo mode; it uses workspace mappings and + # will no-op if disabled. Only allow a single-repo fallback to the default + # collection when startup was explicitly permitted to touch that collection. + _start_pseudo_backfill_worker( + client, + default_collection, + model_dim, + vector_name, + allow_default_collection_fallback=ensure_default_collection, + ) + init_maintenance_shutdown = start_init_maintenance_worker() try: initialize_watcher_state(str(ROOT), multi_repo_enabled, default_collection) @@ -178,21 +361,46 @@ def main() -> None: paths, client, model, vector_name, model_dim, str(ROOT) ) ) - handler = IndexHandler(ROOT, q, client, default_collection) + journal_drain_enabled = _journal_drain_enabled(multi_repo_enabled) + fs_events_enabled = _fs_events_enabled(multi_repo_enabled) + + print( + "[watch_mode] sources " + f"journal_drain={'on' if journal_drain_enabled else 'off'} " + f"fs_events={'on' if fs_events_enabled else 'off'} " + f"journal_batch={_journal_drain_batch_size()}" + ) + + obs = None + if fs_events_enabled: + handler = IndexHandler(ROOT, q, client, default_collection) + use_polling = get_boolean_env("WATCH_USE_POLLING") + obs = create_observer(use_polling, observer_cls=Observer) + obs.schedule(handler, str(ROOT), recursive=True) + obs.start() - use_polling = get_boolean_env("WATCH_USE_POLLING") - obs = create_observer(use_polling, observer_cls=Observer) - obs.schedule(handler, str(ROOT), recursive=True) - obs.start() + maintenance_interval = _maintenance_interval_secs() + last_maintenance: Optional[float] = None try: while True: - time.sleep(1.0) + if journal_drain_enabled: + # Upload/apply records upsert/delete intent here so missed filesystem + # events can still be replayed after watcher/container restarts. + _drain_pending_journal(q) + now = time.time() + if last_maintenance is None or (now - last_maintenance) >= maintenance_interval: + _run_periodic_maintenance(client) + last_maintenance = now + _sleep(1.0) except KeyboardInterrupt: pass finally: - obs.stop() - obs.join() + if init_maintenance_shutdown is not None: + init_maintenance_shutdown.set() + if obs is not None: + obs.stop() + obs.join() if __name__ == "__main__": diff --git a/scripts/watch_index_core/config.py b/scripts/watch_index_core/config.py index c9fa8354..600f840c 100644 --- a/scripts/watch_index_core/config.py +++ b/scripts/watch_index_core/config.py @@ -3,15 +3,12 @@ from __future__ import annotations import os -import sys from pathlib import Path from scripts.logger import get_logger ROOT_DIR = Path(__file__).resolve().parent.parent.parent -if str(ROOT_DIR) not in sys.path: - sys.path.insert(0, str(ROOT_DIR)) def build_logger(): @@ -33,6 +30,12 @@ def build_logger(): # Debounce interval for file system events DELAY_SECS = float(os.environ.get("WATCH_DEBOUNCE_SECS", "1.0")) +# Suppress repeated processing of the exact same observed file state for a short +# window. This is especially useful on shared/polled filesystems like CephFS. +RECENT_FINGERPRINT_TTL_SECS = float( + os.environ.get("WATCH_RECENT_FINGERPRINT_TTL_SECS", "0") +) + def default_collection_name() -> str: """Base fallback for collection name before runtime resolution.""" diff --git a/scripts/watch_index_core/consistency.py b/scripts/watch_index_core/consistency.py new file mode 100644 index 00000000..c3779837 --- /dev/null +++ b/scripts/watch_index_core/consistency.py @@ -0,0 +1,703 @@ +from __future__ import annotations + +import json +import os +from datetime import datetime, timezone +from pathlib import Path +from typing import Any, Dict, Optional, Set, Tuple + +from qdrant_client import QdrantClient + +import scripts.ingest_code as idx +from scripts.workspace_state import ( + _get_state_lock, + _extract_repo_name_from_path, + _normalize_cache_key_path, + get_collection_state_snapshot, + get_workspace_state, + list_workspaces, + update_workspace_state, + upsert_index_journal_entries, +) + +from .config import LOGGER +from .utils import get_boolean_env +from .paths import is_internal_metadata_path + +logger = LOGGER +_DEFAULT_EMPTY_DIR_SWEEP_INTERVAL_SECONDS = 7 * 24 * 60 * 60 + + +def _consistency_audit_enabled() -> bool: + return get_boolean_env("WATCH_CONSISTENCY_AUDIT_ENABLED", default=True) + + +def _consistency_audit_interval_secs() -> int: + try: + return max(60, int(os.environ.get("WATCH_CONSISTENCY_AUDIT_INTERVAL_SECS", "86400") or 86400)) + except Exception: + return 86400 + + +def _consistency_audit_max_paths() -> int: + try: + return max(0, int(os.environ.get("WATCH_CONSISTENCY_AUDIT_MAX_PATHS", "200000") or 200000)) + except Exception: + return 200000 + + +def _consistency_repair_enabled() -> bool: + return get_boolean_env("WATCH_CONSISTENCY_REPAIR_ENABLED", default=True) + + +def _consistency_repair_max_ops() -> int: + try: + return max(0, int(os.environ.get("WATCH_CONSISTENCY_REPAIR_MAX_OPS", "5000") or 5000)) + except Exception: + return 5000 + + +def _consistency_graph_audit_enabled() -> bool: + return get_boolean_env("WATCH_CONSISTENCY_AUDIT_GRAPH_ENABLED", default=True) + + +def _empty_dir_sweep_enabled() -> bool: + if "WATCH_EMPTY_DIR_SWEEP_ENABLED" in os.environ: + return get_boolean_env("WATCH_EMPTY_DIR_SWEEP_ENABLED", default=True) + return get_boolean_env("CTXCE_UPLOAD_EMPTY_DIR_SWEEP", default=True) + + +def _empty_dir_sweep_interval_secs() -> int: + raw = os.environ.get("WATCH_EMPTY_DIR_SWEEP_INTERVAL_SECONDS") + if raw is None: + raw = os.environ.get( + "CTXCE_UPLOAD_EMPTY_DIR_SWEEP_INTERVAL_SECONDS", + str(_DEFAULT_EMPTY_DIR_SWEEP_INTERVAL_SECONDS), + ) + try: + return max(0, int(raw or _DEFAULT_EMPTY_DIR_SWEEP_INTERVAL_SECONDS)) + except Exception: + return _DEFAULT_EMPTY_DIR_SWEEP_INTERVAL_SECONDS + + +def _parse_ts(value: Any) -> Optional[datetime]: + raw = str(value or "").strip() + if not raw: + return None + try: + parsed = datetime.fromisoformat(raw.replace("Z", "+00:00")) + except ValueError: + return None + if parsed.tzinfo is None: + return parsed.replace(tzinfo=timezone.utc) + return parsed.astimezone(timezone.utc) + + +def _should_run_consistency_audit(workspace_path: str, repo_name: Optional[str]) -> bool: + if not _consistency_audit_enabled(): + return False + interval = _consistency_audit_interval_secs() + try: + state = get_workspace_state(workspace_path=workspace_path, repo_name=repo_name) or {} + except Exception: + return True + maintenance = dict(state.get("maintenance") or {}) + last = _parse_ts(maintenance.get("last_consistency_audit_at")) + if last is None: + return True + age = (datetime.now(timezone.utc) - last).total_seconds() + return age >= interval + + +def _sweep_empty_workspace_dirs(workspace_root: Path) -> bool: + """Sweep empty workspace directories and return True if fully successful.""" + protected_top_level = {".codebase", ".remote-git"} + try: + workspace_root = workspace_root.resolve() + except Exception: + return False + try: + for root, _dirnames, _filenames in os.walk(workspace_root, topdown=False): + current = Path(root) + if current == workspace_root: + continue + if current.parent == workspace_root and current.name in protected_top_level: + continue + try: + rel = current.relative_to(workspace_root) + except Exception: + continue + if rel.parts and rel.parts[0] in protected_top_level: + continue + try: + if any(current.iterdir()): + continue + current.rmdir() + except Exception: + # If any directory operation fails, the sweep was not fully successful + return False + except Exception: + return False + return True + + +def _should_run_empty_dir_sweep(workspace_path: str, repo_name: Optional[str]) -> bool: + if not _empty_dir_sweep_enabled(): + return False + interval_seconds = _empty_dir_sweep_interval_secs() + if interval_seconds == 0: + return True + try: + state = get_workspace_state(workspace_path=workspace_path, repo_name=repo_name) or {} + except Exception: + return True + maintenance = state.get("maintenance") or {} + last_sweep_at = _parse_ts(maintenance.get("last_empty_dir_sweep_at")) + if last_sweep_at is None: + return True + age_seconds = (datetime.now(timezone.utc) - last_sweep_at).total_seconds() + return age_seconds >= interval_seconds + + +def _record_empty_dir_sweep(workspace_path: str, repo_name: Optional[str]) -> None: + try: + lock = _get_state_lock(workspace_path, repo_name) + with lock: + state = get_workspace_state( + workspace_path=workspace_path, + repo_name=repo_name, + ) or {} + maintenance = dict(state.get("maintenance") or {}) + maintenance["last_empty_dir_sweep_at"] = datetime.now( + timezone.utc + ).isoformat() + update_workspace_state( + workspace_path=workspace_path, + repo_name=repo_name, + updates={"maintenance": maintenance}, + ) + except Exception as exc: + logger.warning( + "Failed to record empty dir sweep timestamp: %s (workspace=%s, repo=%s)", + exc, + workspace_path, + repo_name, + ) + + +def _load_cached_hashes( + workspace_path: str, + repo_name: Optional[str], + *, + metadata_root: Optional[Path] = None, +) -> Dict[str, str]: + workspace_norm = _normalize_cache_key_path(workspace_path) + workspace_prefix = f"{workspace_norm.rstrip('/')}/" + candidates: list[Path] = [] + seen: set[str] = set() + + def _append_candidate(path: Path) -> None: + key = str(path) + if key in seen: + return + seen.add(key) + candidates.append(path) + + root = Path(metadata_root or workspace_path) + if repo_name: + _append_candidate(root / ".codebase" / "repos" / repo_name / "cache.json") + else: + _append_candidate(root / ".codebase" / "cache.json") + + for cache_path in candidates: + if not cache_path.exists(): + continue + try: + with cache_path.open("r", encoding="utf-8-sig") as f: + data = json.load(f) + hashes = data.get("file_hashes", {}) + if not isinstance(hashes, dict): + return {} + normalized: Dict[str, str] = {} + for path_key, value in hashes.items(): + norm = _normalize_cache_key_path(str(path_key)) + if not norm: + continue + if workspace_norm and not ( + norm == workspace_norm or norm.startswith(workspace_prefix) + ): + continue + if isinstance(value, dict): + digest = str(value.get("hash") or "").strip() + else: + digest = str(value or "").strip() + normalized[norm] = digest + return normalized + except Exception: + return {} + return {} + + +def _is_index_eligible_path(path_str: str, workspace_root: Path, excluder) -> bool: + try: + p = Path(path_str).resolve() + except Exception: + p = Path(path_str) + try: + rel = p.resolve().relative_to(workspace_root.resolve()) + except Exception: + return False + + if not rel.parts: + return False + if not p.exists() or p.is_dir(): + return False + try: + if int(p.stat().st_size) == 0: + # Empty files (e.g. many __init__.py stubs) produce no vectors; do not + # enqueue consistency upserts for them. + return False + except Exception: + return False + if is_internal_metadata_path(p): + return False + + # .remote-git manifests are control files and must not be treated as indexable. + if _is_remote_git_manifest(p.as_posix()): + return False + + try: + rel_dir = "/" + str(rel.parent).replace(os.sep, "/") + if rel_dir == "/.": + rel_dir = "/" + if excluder.exclude_dir(rel_dir): + return False + except Exception: + return False + + if not idx.is_indexable_file(p): + return False + + try: + relf = (rel_dir.rstrip("/") + "/" + p.name).replace("//", "/") + if excluder.exclude_file(relf): + return False + except Exception: + return False + return True + + +def _scan_indexable_fs_paths(workspace_root: Path, *, max_paths: int) -> Tuple[Set[str], bool]: + paths: Set[str] = set() + excluder = idx._Excluder(workspace_root) + try: + workspace_root = workspace_root.resolve() + except Exception: + pass + + for root_str, dirnames, filenames in os.walk(workspace_root): + current = Path(root_str) + pruned_dirnames = [] + for dirname in dirnames: + child = current / dirname + if is_internal_metadata_path(child): + continue + try: + rel_dir = "/" + str(child.relative_to(workspace_root)).replace(os.sep, "/") + if excluder.exclude_dir(rel_dir): + continue + except Exception: + pass + pruned_dirnames.append(dirname) + dirnames[:] = pruned_dirnames + + for filename in filenames: + file_path = current / filename + normalized = _normalize_cache_key_path(str(file_path)) + if not normalized: + continue + if not _is_index_eligible_path(normalized, workspace_root, excluder): + continue + paths.add(normalized) + if max_paths > 0 and len(paths) >= max_paths: + return paths, True + return paths, False + + +def _load_indexed_paths_for_collection( + client: QdrantClient, + collection: str, + workspace_path: str, + *, + max_paths: int, +) -> Tuple[Set[str], bool]: + paths: Set[str] = set() + workspace_norm = _normalize_cache_key_path(workspace_path) + workspace_prefix = f"{workspace_norm.rstrip('/')}/" + offset = None + while True: + points, next_offset = client.scroll( + collection_name=collection, + limit=1000, + with_payload=True, + with_vectors=False, + offset=offset, + ) + for pt in points or []: + payload = getattr(pt, "payload", {}) or {} + metadata = payload.get("metadata", {}) or {} + path = _normalize_cache_key_path(str(metadata.get("path") or "")) + if path: + if workspace_norm and not ( + path == workspace_norm or path.startswith(workspace_prefix) + ): + continue + paths.add(path) + if max_paths > 0 and len(paths) >= max_paths: + return paths, True + if next_offset is None: + break + offset = next_offset + return paths, False + + +def _load_graph_paths_for_collection( + client: QdrantClient, + collection: str, + workspace_path: str, + *, + max_paths: int, +) -> Tuple[Set[str], bool]: + paths: Set[str] = set() + workspace_norm = _normalize_cache_key_path(workspace_path) + workspace_prefix = f"{workspace_norm.rstrip('/')}/" + graph_collection = f"{collection}_graph" + offset = None + while True: + points, next_offset = client.scroll( + collection_name=graph_collection, + limit=1000, + with_payload=True, + with_vectors=False, + offset=offset, + ) + for pt in points or []: + payload = getattr(pt, "payload", {}) or {} + path = _normalize_cache_key_path(str(payload.get("caller_path") or "")) + if path: + if workspace_norm and not ( + path == workspace_norm or path.startswith(workspace_prefix) + ): + continue + paths.add(path) + if max_paths > 0 and len(paths) >= max_paths: + return paths, True + if next_offset is None: + break + offset = next_offset + return paths, False + + +def _record_consistency_audit( + workspace_path: str, + repo_name: Optional[str], + summary: Dict[str, Any], +) -> None: + try: + lock = _get_state_lock(workspace_path, repo_name) + with lock: + state = get_workspace_state( + workspace_path=workspace_path, + repo_name=repo_name, + ) or {} + maintenance = dict(state.get("maintenance") or {}) + maintenance["last_consistency_audit_at"] = datetime.now( + timezone.utc + ).isoformat() + maintenance["last_consistency_audit_summary"] = summary + update_workspace_state( + workspace_path=workspace_path, + repo_name=repo_name, + updates={"maintenance": maintenance}, + ) + except Exception as exc: + logger.warning( + "Failed to record consistency audit: %s (workspace=%s, repo=%s)", + exc, + workspace_path, + repo_name, + ) + + +def _is_remote_git_manifest(path: str) -> bool: + """Check if path is a .remote-git git history manifest file (control file, not indexable content).""" + try: + p = Path(path) + return any(part == ".remote-git" for part in p.parts) and p.suffix.lower() == ".json" + except Exception: + return False + + +def _enqueue_consistency_repairs( + workspace_root: Path, + workspace_path: str, + repo_name: Optional[str], + stale_paths: list[str], + missing_paths: list[str], + cached_hashes: Dict[str, str], +) -> Tuple[int, int]: + if not _consistency_repair_enabled(): + return 0, 0 + max_ops = _consistency_repair_max_ops() + if max_ops <= 0: + return 0, 0 + + entries: list[Dict[str, Any]] = [] + enqueued_stale = 0 + enqueued_missing = 0 + missing_set = set(missing_paths) + excluder = idx._Excluder(workspace_root) + + for path in stale_paths: + if len(entries) >= max_ops: + break + # Skip .remote-git git history manifests - they are control files, not indexable content + if _is_remote_git_manifest(path): + continue + # Cache can lag after state resets/rebuilds; if the path still exists and is + # index-eligible, treat it as missing/upsert instead of stale/delete. + if _is_index_eligible_path(path, workspace_root, excluder): + missing_set.add(path) + continue + entries.append({"path": path, "op_type": "delete"}) + enqueued_stale += 1 + for path in sorted(missing_set): + if len(entries) >= max_ops: + break + # Skip .remote-git git history manifests - they are control files, not indexable content + if _is_remote_git_manifest(path): + continue + entries.append( + { + "path": path, + "op_type": "upsert", + "content_hash": cached_hashes.get(path) or None, + } + ) + enqueued_missing += 1 + + if not entries: + return 0, 0 + + # Fetch existing journal entries to preserve retry state + existing_entries: Dict[str, Dict[str, Any]] = {} + try: + from scripts.workspace_state import list_pending_index_journal_entries + all_pending = list_pending_index_journal_entries( + workspace_path=workspace_path, + repo_name=repo_name, + ) + for entry in all_pending or []: + path = str(entry.get("path") or "") + if path: + existing_entries[path] = entry + except Exception: + pass # If we can't fetch existing entries, proceed without preserving state + + # Merge existing retry state into new entries where appropriate + merged_entries = [] + for entry in entries: + path = str(entry.get("path") or "") + existing = existing_entries.get(path) + + # Skip if already pending/in-progress to avoid duplicate work + if existing and existing.get("status") in {"pending", "in_progress"}: + continue + + # Preserve retry state from existing failed entries + if existing and existing.get("status") == "failed": + entry["status"] = "failed" + entry["attempts"] = existing.get("attempts", 0) + entry["last_error"] = existing.get("last_error") + # Keep created_at from existing entry to preserve original enqueue time + if existing.get("created_at"): + entry["created_at"] = existing["created_at"] + + merged_entries.append(entry) + + if not merged_entries: + return 0, 0 + + try: + upsert_index_journal_entries( + merged_entries, + workspace_path=workspace_path, + repo_name=repo_name, + ) + except Exception as exc: + logger.debug( + "[consistency_audit] failed to enqueue repairs workspace=%s repo=%s: %s", + workspace_path, + repo_name, + exc, + ) + return 0, 0 + + # Return counts based on actually enqueued entries + enqueued_stale = sum(1 for e in merged_entries if e.get("op_type") == "delete") + enqueued_missing = sum(1 for e in merged_entries if e.get("op_type") == "upsert") + return enqueued_stale, enqueued_missing + + +def run_consistency_audit(client: QdrantClient, root: Path) -> None: + if not _consistency_audit_enabled(): + return + max_paths = _consistency_audit_max_paths() + try: + candidates = list_workspaces(search_root=str(root), use_qdrant_fallback=False) + except Exception: + candidates = [] + for ws in candidates: + workspace_path = str(ws.get("workspace_path") or "").strip() + if not workspace_path: + continue + repo_name = _extract_repo_name_from_path(workspace_path) + if not _should_run_consistency_audit(workspace_path, repo_name): + continue + try: + snapshot = get_collection_state_snapshot( + workspace_path=workspace_path, + repo_name=repo_name, + ) + collection = str(snapshot.get("active_collection") or "").strip() + if not collection: + continue + cached_hashes = _load_cached_hashes( + workspace_path, + repo_name, + metadata_root=root, + ) + workspace_root = Path(workspace_path) + fs_paths, fs_truncated = _scan_indexable_fs_paths( + workspace_root, + max_paths=max_paths, + ) + excluder = idx._Excluder(workspace_root) + cached_paths = { + path + for path in cached_hashes.keys() + if _is_index_eligible_path(path, workspace_root, excluder) + } + indexed_paths, indexed_truncated = _load_indexed_paths_for_collection( + client, + collection, + workspace_path, + max_paths=max_paths, + ) + graph_paths: Set[str] = set() + graph_truncated = False + graph_orphans: list[str] = [] + if _consistency_graph_audit_enabled(): + try: + graph_paths, graph_truncated = _load_graph_paths_for_collection( + client, + collection, + workspace_path, + max_paths=max_paths, + ) + except Exception: + graph_paths, graph_truncated = set(), False + if fs_truncated or indexed_truncated: + stale = [] + missing = [] + enq_stale = 0 + enq_missing = 0 + else: + stale_set = set(indexed_paths - fs_paths) + if not graph_truncated: + graph_orphans = sorted(graph_paths - indexed_paths) + stale_set.update(graph_orphans) + stale = sorted(stale_set) + missing = sorted(fs_paths - indexed_paths) + enq_stale, enq_missing = _enqueue_consistency_repairs( + workspace_root, + workspace_path, + repo_name, + stale, + missing, + cached_hashes, + ) + summary = { + "fs_count": len(fs_paths), + "cache_count": len(cached_paths), + "qdrant_count": len(indexed_paths), + "graph_count": len(graph_paths), + "fs_scan_truncated": fs_truncated, + "qdrant_scan_truncated": indexed_truncated, + "graph_scan_truncated": graph_truncated, + "repair_skipped_due_to_truncation": bool(fs_truncated or indexed_truncated), + "stale_in_qdrant_count": len(stale), + "missing_in_qdrant_count": len(missing), + "orphan_graph_count": len(graph_orphans), + "repair_enqueued_stale_count": int(enq_stale), + "repair_enqueued_missing_count": int(enq_missing), + "sample_stale": stale[:20], + "sample_missing": missing[:20], + "sample_orphan_graph": graph_orphans[:20], + } + _record_consistency_audit(workspace_path, repo_name, summary) + logger.info( + "[consistency_audit] repo=%s collection=%s fs=%d cache=%d qdrant=%d graph=%d stale=%d missing=%d graph_orphans=%d repair_stale=%d repair_missing=%d", + repo_name or "", + collection, + len(fs_paths), + len(cached_paths), + len(indexed_paths), + len(graph_paths), + len(stale), + len(missing), + len(graph_orphans), + int(enq_stale), + int(enq_missing), + ) + except Exception as exc: + logger.debug( + "[consistency_audit] failed workspace=%s repo=%s: %s", + workspace_path, + repo_name, + exc, + ) + + +def run_empty_dir_sweep_maintenance(root: Path) -> None: + if not _empty_dir_sweep_enabled(): + return + try: + candidates = list_workspaces(search_root=str(root), use_qdrant_fallback=False) + except Exception: + candidates = [] + for ws in candidates: + workspace_path = str(ws.get("workspace_path") or "").strip() + if not workspace_path: + continue + repo_name = _extract_repo_name_from_path(workspace_path) + if not _should_run_empty_dir_sweep(workspace_path, repo_name): + continue + try: + logger.info("[empty_dir_sweep] Sweeping empty directories under %s", workspace_path) + sweep_success = _sweep_empty_workspace_dirs(Path(workspace_path)) + if sweep_success: + _record_empty_dir_sweep(workspace_path, repo_name) + else: + logger.debug( + "[empty_dir_sweep] sweep had failures workspace=%s repo=%s - not recording success", + workspace_path, + repo_name, + ) + except Exception as exc: + logger.debug( + "[empty_dir_sweep] failed workspace=%s repo=%s: %s", + workspace_path, + repo_name, + exc, + ) diff --git a/scripts/watch_index_core/handler.py b/scripts/watch_index_core/handler.py index bf5cb6d9..36c4d457 100644 --- a/scripts/watch_index_core/handler.py +++ b/scripts/watch_index_core/handler.py @@ -12,7 +12,6 @@ import scripts.ingest_code as idx from scripts.workspace_state import ( _extract_repo_name_from_path, - _get_global_state_dir, get_cached_file_hash, log_watcher_activity as _log_activity, remove_cached_file, @@ -27,6 +26,7 @@ safe_print, ) from .rename import _rename_in_store +from .paths import is_internal_metadata_path class IndexHandler(FileSystemEventHandler): @@ -81,6 +81,9 @@ def _maybe_reload_excluder(self) -> None: except Exception: pass + def _is_internal_metadata_path(self, p: Path) -> bool: + return is_internal_metadata_path(p) + def _maybe_enqueue(self, src_path: str) -> None: self._maybe_reload_excluder() p = Path(src_path) @@ -95,15 +98,7 @@ def _maybe_enqueue(self, src_path: str) -> None: except ValueError: return - try: - if callable(_get_global_state_dir): - global_state_dir = _get_global_state_dir() - if global_state_dir is not None and p.is_relative_to(global_state_dir): - return - except (OSError, ValueError): - pass - - if any(part == ".codebase" for part in p.parts): + if self._is_internal_metadata_path(p): return # Git history manifests are handled by a separate ingestion pipeline and should still @@ -140,7 +135,7 @@ def on_deleted(self, event): p = Path(event.src_path).resolve() except Exception: return - if any(part == ".codebase" for part in p.parts): + if self._is_internal_metadata_path(p): return if not idx.is_indexable_file(p): return @@ -162,6 +157,42 @@ def on_moved(self, event): dest = Path(event.dest_path).resolve() except Exception: return + # Handle internal-boundary moves properly + src_internal = self._is_internal_metadata_path(src) + dest_internal = self._is_internal_metadata_path(dest) + if src_internal and dest_internal: + # Both internal -> ignore + return + if dest_internal: + # External -> internal: delete source, don't index destination + if idx.is_indexable_file(src): + try: + coll = self._resolve_collection(src) + deleted = False + if self.client is not None and coll is not None: + idx.delete_points_by_path(self.client, coll, str(src)) + # Clean up graph edges for the moved file + try: + idx.delete_graph_edges_by_path( + self.client, + coll, + caller_path=str(src), + ) + except Exception: + pass # Graph cleanup is best-effort + deleted = True + if deleted: + safe_print(f"[moved:external_to_internal] deleted {src}") + except Exception as exc: + safe_print(f"[moved:external_to_internal:error] {src}: {exc}") + finally: + self._invalidate_cache(src) + return + if src_internal: + # Internal -> external: index destination as new file + if idx.is_indexable_file(dest): + self._maybe_enqueue(str(dest)) + return if not idx.is_indexable_file(dest) and not idx.is_indexable_file(src): return try: @@ -174,18 +205,25 @@ def on_moved(self, event): if idx.is_indexable_file(src): try: coll = self._resolve_collection(src) + deleted = False if self.client is not None and coll is not None: idx.delete_points_by_path(self.client, coll, str(src)) - safe_print(f"[moved:ignored_dest_deleted_src] {src} -> {dest}") - src_repo_path = _detect_repo_for_file(src) - src_repo_name = _repo_name_or_none(src_repo_path) - try: - if src_repo_name: - remove_cached_file(str(src), src_repo_name) - except Exception: - pass - except Exception: - pass + # Clean up graph edges for the moved file + try: + idx.delete_graph_edges_by_path( + self.client, + coll, + caller_path=str(src), + ) + except Exception: + pass # Graph cleanup is best-effort + deleted = True + if deleted: + safe_print(f"[moved:ignored_dest_deleted_src] {src} -> {dest}") + except Exception as exc: + safe_print(f"[moved:ignored_dest_deleted_src:error] {src}: {exc}") + finally: + self._invalidate_cache(src) return except Exception: pass @@ -270,6 +308,14 @@ def _delete_points(self, path: Path, collection: str | None) -> None: return try: idx.delete_points_by_path(self.client, collection, str(path)) + try: + idx.delete_graph_edges_by_path( + self.client, + collection, + caller_path=str(path), + ) + except Exception: + pass safe_print(f"[deleted] {path} -> {collection}") except Exception: pass diff --git a/scripts/watch_index_core/init_maintenance.py b/scripts/watch_index_core/init_maintenance.py new file mode 100644 index 00000000..bb2dece9 --- /dev/null +++ b/scripts/watch_index_core/init_maintenance.py @@ -0,0 +1,191 @@ +"""Periodic bootstrap/init maintenance for the long-lived watcher.""" + +from __future__ import annotations + +import os +import subprocess +import sys +import threading +from pathlib import Path +from typing import Optional, Sequence + +from scripts.workspace_state import _cross_process_lock, _get_global_state_dir + +from . import config as watch_config +from .config import LOGGER +from .utils import get_boolean_env + +logger = LOGGER + +DEFAULT_INTERVAL_MINUTES = 120.0 +DEFAULT_COMMAND_TIMEOUT_SECONDS = 1800.0 + + +def _interval_seconds() -> float: + raw = os.environ.get("WATCH_INIT_MAINTENANCE_INTERVAL_MINUTES") + if raw is None: + raw = os.environ.get("INIT_MAINTENANCE_INTERVAL_MINUTES") + try: + minutes = float(raw if raw is not None else DEFAULT_INTERVAL_MINUTES) + except Exception: + minutes = DEFAULT_INTERVAL_MINUTES + return max(0.0, minutes * 60.0) + + +def _command_timeout_seconds() -> float: + try: + return max( + 1.0, + float( + os.environ.get( + "WATCH_INIT_MAINTENANCE_COMMAND_TIMEOUT_SECS", + str(DEFAULT_COMMAND_TIMEOUT_SECONDS), + ) + or DEFAULT_COMMAND_TIMEOUT_SECONDS + ), + ) + except Exception: + return DEFAULT_COMMAND_TIMEOUT_SECONDS + + +def _script_root() -> Path: + return Path(__file__).resolve().parents[1] + + +def _wait_for_qdrant_command(script_root: Path) -> list[str]: + return [str(script_root / "wait-for-qdrant.sh")] + + +def _python_script_command(script_root: Path, script_name: str) -> list[str]: + return [sys.executable or "python", str(script_root / script_name)] + + +def _maintenance_commands(script_root: Optional[Path] = None) -> list[list[str]]: + scripts = script_root or _script_root() + return [ + _wait_for_qdrant_command(scripts), + _python_script_command(scripts, "create_indexes.py"), + _python_script_command(scripts, "warm_all_collections.py"), + _python_script_command(scripts, "health_check.py"), + ] + + +def _env_for_subprocess() -> dict[str, str]: + env = os.environ.copy() + root = str(Path(__file__).resolve().parents[2]) + existing = env.get("PYTHONPATH") + env["PYTHONPATH"] = f"{root}{os.pathsep}{existing}" if existing else root + if str(watch_config.ROOT): + env.setdefault("WORKSPACE_PATH", str(watch_config.ROOT)) + env.setdefault("WORKDIR", str(watch_config.ROOT)) + env.setdefault("WORK_DIR", str(watch_config.ROOT)) + return env + + +def _run_command(command: Sequence[str], *, timeout: float, env: dict[str, str]) -> bool: + label = " ".join(str(part) for part in command) + logger.info("[init_maintenance] running: %s", label) + try: + result = subprocess.run( + list(command), + cwd=str(watch_config.ROOT), + env=env, + text=True, + capture_output=True, + timeout=timeout, + check=False, + ) + except subprocess.TimeoutExpired: + logger.error("[init_maintenance] timed out after %.0fs: %s", timeout, label) + return False + except Exception as exc: + logger.error("[init_maintenance] failed to start %s: %s", label, exc, exc_info=True) + return False + + if result.returncode == 0: + logger.info("[init_maintenance] completed: %s", label) + if result.stdout: + logger.debug("[init_maintenance] stdout for %s:\n%s", label, result.stdout[-4000:]) + if result.stderr: + logger.debug("[init_maintenance] stderr for %s:\n%s", label, result.stderr[-4000:]) + return True + + logger.warning( + "[init_maintenance] command failed rc=%s: %s\nstdout:\n%s\nstderr:\n%s", + result.returncode, + label, + (result.stdout or "")[-4000:], + (result.stderr or "")[-4000:], + ) + return False + + +def run_init_maintenance_once( + *, + commands: Optional[Sequence[Sequence[str]]] = None, + lock_path: Optional[Path] = None, +) -> bool: + """Run the existing init scripts once under a cross-process lock.""" + + timeout = _command_timeout_seconds() + env = _env_for_subprocess() + cmd_list = [list(cmd) for cmd in (commands or _maintenance_commands())] + if not cmd_list: + return True + + target_lock = lock_path + if target_lock is None: + try: + target_lock = _get_global_state_dir(str(watch_config.ROOT)) / "init_maintenance.lock" + except Exception: + target_lock = Path("/tmp/context-engine-init-maintenance.lock") + + with _cross_process_lock(target_lock): + for command in cmd_list: + if not _run_command(command, timeout=timeout, env=env): + return False + return True + + +def start_init_maintenance_worker() -> Optional[threading.Event]: + """Start periodic init maintenance, controlled by watcher env vars.""" + + if not get_boolean_env("WATCH_INIT_MAINTENANCE_ENABLED", default=True): + return None + + interval = _interval_seconds() + if interval <= 0: + return None + + run_on_start = get_boolean_env("WATCH_INIT_MAINTENANCE_RUN_ON_START", default=False) + shutdown_event = threading.Event() + + def _worker() -> None: + if not run_on_start: + shutdown_event.wait(timeout=interval) + while not shutdown_event.is_set(): + try: + ok = run_init_maintenance_once() + if ok: + logger.info("[init_maintenance] pass completed") + else: + logger.warning("[init_maintenance] pass completed with failures") + except Exception: + logger.error("[init_maintenance] unexpected worker error", exc_info=True) + shutdown_event.wait(timeout=interval) + + thread = threading.Thread(target=_worker, name="init-maintenance", daemon=True) + thread.start() + logger.info( + "[init_maintenance] worker started interval=%.1fm run_on_start=%s", + interval / 60.0, + run_on_start, + ) + return shutdown_event + + +__all__ = [ + "DEFAULT_INTERVAL_MINUTES", + "run_init_maintenance_once", + "start_init_maintenance_worker", +] diff --git a/scripts/watch_index_core/paths.py b/scripts/watch_index_core/paths.py new file mode 100644 index 00000000..2e76cfb9 --- /dev/null +++ b/scripts/watch_index_core/paths.py @@ -0,0 +1,36 @@ +"""Path classification helpers shared by watcher components.""" + +from __future__ import annotations + +from pathlib import Path + +from scripts.workspace_state import ( + _get_global_state_dir, + INTERNAL_STATE_TOP_LEVEL_DIRS, +) + + +def is_internal_metadata_path(path: Path) -> bool: + """Return True when path points into watcher/internal metadata trees.""" + try: + # Deliberately match internal segments anywhere in the path to prevent + # indexing of nested metadata mirrors (for example in replicated roots). + if any(part in INTERNAL_STATE_TOP_LEVEL_DIRS for part in path.parts): + return True + global_state_dir = _get_global_state_dir() + if global_state_dir is not None and path.is_relative_to(global_state_dir): + return True + except (OSError, ValueError): + return False + return False + + +def is_internal_top_level_path(path: Path, root: Path) -> bool: + """Return True when path's top-level segment under root is internal metadata.""" + try: + rel = path.resolve().relative_to(root.resolve()) + except Exception: + return False + if not rel.parts: + return False + return rel.parts[0] in INTERNAL_STATE_TOP_LEVEL_DIRS diff --git a/scripts/watch_index_core/processor.py b/scripts/watch_index_core/processor.py index 45e9db7e..fefd4941 100644 --- a/scripts/watch_index_core/processor.py +++ b/scripts/watch_index_core/processor.py @@ -3,28 +3,47 @@ from __future__ import annotations import hashlib +import json import os import subprocess import sys +import atexit +import threading +import time +from collections import deque +from concurrent.futures import Future, ThreadPoolExecutor from datetime import datetime from pathlib import Path from typing import Dict, List, Optional +from qdrant_client import models + import scripts.ingest_code as idx +from scripts.pseudo_config import effective_pseudo_mode +from scripts.ingest.graph_edges import ( + normalize_caller_path as _normalize_graph_caller_path, +) from scripts.workspace_state import ( + _normalize_cache_key_path, _extract_repo_name_from_path, get_cached_file_hash, + list_pending_index_journal_entries, get_workspace_state, is_staging_enabled, log_watcher_activity as _log_activity, persist_indexing_config, remove_cached_file, + set_cached_file_hash, set_indexing_progress as _update_progress, set_indexing_started as _set_status_indexing, + update_index_journal_entries_status, + update_index_journal_entry_status, update_indexing_status, ) +from . import config as watch_config +from .rename import _rename_in_store +from .paths import is_internal_metadata_path -from .config import QDRANT_URL, ROOT, ROOT_DIR, LOGGER as logger from .utils import ( _detect_repo_for_file, _get_collection_for_file, @@ -33,43 +52,370 @@ safe_log_error, ) +logger = watch_config.LOGGER + +_JOURNAL_STATUS_BATCH_LOCAL = threading.local() + + +def _active_journal_status_batch() -> Optional[List[Dict[str, object]]]: + return getattr(_JOURNAL_STATUS_BATCH_LOCAL, "current", None) + + +def _queue_journal_status( + path: Path, + repo_key: str, + repo_name: Optional[str], + *, + status: str, + error: Optional[str] = None, + remove_on_done: bool = True, +) -> bool: + batch = _active_journal_status_batch() + if batch is None: + return False + batch.append( + { + "repo_key": repo_key, + "repo_name": repo_name, + "path": _normalize_cache_key_path(str(path)), + "status": status, + "error": error, + "remove_on_done": remove_on_done, + } + ) + return True + + +def _flush_journal_status_batch(updates: List[Dict[str, object]]) -> None: + if not updates: + return + grouped: Dict[tuple[str, Optional[str]], List[Dict[str, object]]] = {} + for update in updates: + key = (str(update["repo_key"]), update.get("repo_name")) + grouped.setdefault(key, []).append(update) + + for (repo_key, repo_name), grouped_updates in grouped.items(): + payload = [ + { + "path": update["path"], + "status": update["status"], + "error": update.get("error"), + "remove_on_done": update.get("remove_on_done", True), + } + for update in grouped_updates + ] + try: + update_index_journal_entries_status( + payload, + workspace_path=repo_key, + repo_name=repo_name, + ) + continue + except Exception: + logger.exception( + "watch_index::journal_bulk_status_failed", + extra={"repo_key": repo_key, "repo_name": repo_name, "count": len(payload)}, + ) + + for update in payload: + try: + update_index_journal_entry_status( + str(update["path"]), + status=str(update["status"]), + error=update.get("error"), + workspace_path=repo_key, + repo_name=repo_name, + remove_on_done=bool(update.get("remove_on_done", True)), + ) + except Exception: + pass + class _SkipUnchanged(Exception): """Sentinel exception to skip unchanged files in the watch loop.""" + def __init__(self, *, text: Optional[str] = None, file_hash: str = "") -> None: + super().__init__("unchanged") + self.text = text + self.file_hash = file_hash -def _process_git_history_manifest( + +def _is_internal_ignored_path(path: Path) -> bool: + return is_internal_metadata_path(path) + + +def _staging_requires_subprocess(state: Optional[Dict[str, object]]) -> bool: + """Return True only when dual-root staging is actually active for this repo.""" + if not (is_staging_enabled() and isinstance(state, dict)): + return False + + staging = state.get("staging") + if isinstance(staging, dict) and staging: + return True + + active_slug = str(state.get("active_repo_slug") or "").strip() + serving_slug = str(state.get("serving_repo_slug") or "").strip() + if serving_slug.endswith("_old"): + return True + if active_slug and serving_slug and active_slug != serving_slug: + return True + return False + + +def _env_int(name: str, default: int) -> int: + try: + raw = str(os.environ.get(name, str(default))).strip() + val = int(raw) + return val if val > 0 else default + except Exception: + return default + + +_GIT_HISTORY_MAX_WORKERS = _env_int("WATCH_GIT_HISTORY_MAX_WORKERS", 1) +_GIT_HISTORY_TIMEOUT_SECONDS = _env_int("WATCH_GIT_HISTORY_TIMEOUT_SECONDS", 0) +_GIT_HISTORY_EXECUTOR = ThreadPoolExecutor( + max_workers=_GIT_HISTORY_MAX_WORKERS, + thread_name_prefix="git-history", +) + + +def _shutdown_git_history_executor() -> None: + try: + _GIT_HISTORY_EXECUTOR.shutdown(wait=False) + except Exception: + pass + + +atexit.register(_shutdown_git_history_executor) +_GIT_HISTORY_INFLIGHT: set[str] = set() +_GIT_HISTORY_INFLIGHT_LOCK = threading.Lock() + + +def _manifest_key(p: Path) -> str: + try: + return str(p.resolve()) + except Exception: + return str(p) + + +def _manifest_stats(p: Path) -> tuple[str, int]: + run_id = "unknown" + commit_count = -1 + try: + with p.open("r", encoding="utf-8") as fh: + data = json.load(fh) + if isinstance(data, dict): + commits = data.get("commits") or [] + if isinstance(commits, list): + commit_count = len(commits) + name = p.name + run_id = name[:-5] if name.endswith(".json") else name + except Exception: + pass + return run_id, commit_count + + +def _run_git_history_ingest( p: Path, collection: str, repo_name: Optional[str], env_snapshot: Optional[Dict[str, str]] = None, ) -> None: - try: - script = ROOT_DIR / "scripts" / "ingest_history.py" - if not script.exists(): - return - cmd = [sys.executable or "python3", str(script), "--manifest-json", str(p)] - env = _build_subprocess_env(collection, repo_name, env_snapshot) + script = watch_config.ROOT_DIR / "scripts" / "ingest_history.py" + if not script.exists(): + raise RuntimeError(f"[git_history_manifest] ingest script missing: {script}") + + cmd = [sys.executable or "python3", "-m", "scripts.ingest_history", "--manifest-json", str(p)] + env = _build_subprocess_env(collection, repo_name, env_snapshot) + started = time.monotonic() + timeout = _GIT_HISTORY_TIMEOUT_SECONDS if _GIT_HISTORY_TIMEOUT_SECONDS > 0 else None + stdout_tail: deque[str] = deque(maxlen=20) + stderr_tail: deque[str] = deque(maxlen=20) + tail_lock = threading.Lock() + + def _tail_snapshot(tail: deque[str], limit: int = 5) -> str: + with tail_lock: + return " | ".join(list(tail)[-limit:]) + + def _stream_pipe(pipe, label: str, tail: deque[str], lock: threading.Lock) -> None: try: - print( - f"[git_history_manifest] launching ingest_history.py for {p} " - f"collection={collection} repo={repo_name}" - ) + for raw in iter(pipe.readline, ""): + line = (raw or "").rstrip() + if not line: + continue + with lock: + tail.append(line) + logger.info("[git_history_manifest][%s] %s", label, line) except Exception: pass - # Use subprocess.run for better error observability. - # NOTE: This blocks until ingest_history.py completes. If history ingestion - # is slow, this may need revisiting (e.g., revert to Popen fire-and-forget - # or run in a separate thread) to avoid blocking the watcher. - result = subprocess.run(cmd, env=env, capture_output=True, text=True, check=False) - if result.returncode != 0: - logger.warning( - "[git_history_manifest] ingest_history.py failed for %s: exit=%d stderr=%s", - p, result.returncode, (result.stderr or "")[:500], + finally: + try: + pipe.close() + except Exception: + pass + + proc: Optional[subprocess.Popen] = None + try: + proc = subprocess.Popen( + cmd, + cwd=str(watch_config.ROOT_DIR), + env=env, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + bufsize=1, + ) + t_out = threading.Thread( + target=_stream_pipe, + args=(proc.stdout, "stdout", stdout_tail, tail_lock), + daemon=True, + ) + t_err = threading.Thread( + target=_stream_pipe, + args=(proc.stderr, "stderr", stderr_tail, tail_lock), + daemon=True, + ) + t_out.start() + t_err.start() + + deadline = (started + timeout) if timeout else None + timed_out = False + while True: + code = proc.poll() + if code is not None: + break + if deadline and time.monotonic() >= deadline: + timed_out = True + try: + proc.kill() + except Exception: + pass + break + time.sleep(0.2) + + # Ensure threads flush trailing output after process exit/kill. + t_out.join(timeout=1.0) + t_err.join(timeout=1.0) + + if timed_out: + elapsed_ms = int((time.monotonic() - started) * 1000) + error_msg = ( + f"[git_history_manifest] ingest_history.py timeout for {p} after {elapsed_ms}ms " + f"(timeout={_GIT_HISTORY_TIMEOUT_SECONDS}s)" ) + if stderr_tail: + error_msg += f" stderr={_tail_snapshot(stderr_tail)}" + logger.warning(error_msg) + raise RuntimeError(error_msg) + + returncode = proc.wait(timeout=1.0) except Exception as e: - logger.warning("[git_history_manifest] error processing %s: %s", p, e) - return + logger.warning("[git_history_manifest] subprocess error for %s: %s", p, e) + try: + if proc and proc.poll() is None: + proc.kill() + except Exception: + pass + raise RuntimeError(f"[git_history_manifest] subprocess error for {p}: {e}") from e + + elapsed_ms = int((time.monotonic() - started) * 1000) + if returncode != 0: + error_msg = ( + f"[git_history_manifest] ingest_history.py failed for {p}: exit={returncode} " + f"elapsed_ms={elapsed_ms} stderr={_tail_snapshot(stderr_tail)}" + ) + logger.warning(error_msg) + raise RuntimeError(error_msg) + + logger.info( + "[git_history_manifest] completed for %s: exit=0 elapsed_ms=%d", + p, + elapsed_ms, + ) + if stdout_tail: + logger.info( + "[git_history_manifest] stdout tail for %s: %s", + p, + _tail_snapshot(stdout_tail), + ) + if stderr_tail: + logger.warning( + "[git_history_manifest] stderr tail for %s: %s", + p, + _tail_snapshot(stderr_tail), + ) + + +def _on_git_history_done(manifest_path: Path, collection: str, repo_name: Optional[str], future: Future) -> None: + manifest_key = _manifest_key(manifest_path) + with _GIT_HISTORY_INFLIGHT_LOCK: + _GIT_HISTORY_INFLIGHT.discard(manifest_key) + remaining = len(_GIT_HISTORY_INFLIGHT) + logger.info("[git_history_manifest] in-flight remaining=%d", remaining) + try: + future.result() + # Mark journal as done after successful completion + repo_path = _detect_repo_for_file(manifest_path) + if repo_path: + repo_key = str(repo_path) + _mark_journal_done(manifest_path, repo_key, repo_name) + logger.info("[git_history_manifest] marked journal as done: %s", manifest_path) + except Exception as e: + repo_path = _detect_repo_for_file(manifest_path) + repo_key = str(repo_path) if repo_path else "" + if repo_key: + _mark_journal_failed( + manifest_path, + repo_key, + repo_name, + f"git history worker failed for collection '{collection}': {e}", + ) + logger.warning( + "[git_history_manifest] worker crashed for %s (collection=%s, repo_key=%s): %s", + manifest_key, + collection, + repo_key or "", + e, + exc_info=True, + ) + + +def _process_git_history_manifest( + p: Path, + collection: str, + repo_name: Optional[str], + env_snapshot: Optional[Dict[str, str]] = None, +) -> None: + key = _manifest_key(p) + run_id, commit_count = _manifest_stats(p) + queued = 0 + with _GIT_HISTORY_INFLIGHT_LOCK: + if key in _GIT_HISTORY_INFLIGHT: + logger.info( + "[git_history_manifest] skip duplicate in-flight manifest: %s run_id=%s", + p, + run_id, + ) + return + _GIT_HISTORY_INFLIGHT.add(key) + queued = len(_GIT_HISTORY_INFLIGHT) + logger.info( + "[git_history_manifest] queued ingest_history.py for %s run_id=%s commits=%d collection=%s repo=%s in_flight=%d", + p, + run_id, + commit_count, + collection, + repo_name, + queued, + ) + future = _GIT_HISTORY_EXECUTOR.submit( + _run_git_history_ingest, + p, + collection, + repo_name, + env_snapshot, + ) + future.add_done_callback(lambda fut, manifest_path=p, coll=collection, rn=repo_name: _on_git_history_done(manifest_path, coll, rn, fut)) def _advance_progress( @@ -92,6 +438,227 @@ def _advance_progress( pass +def _mark_journal_done(path: Path, repo_key: str, repo_name: Optional[str]) -> None: + if _queue_journal_status(path, repo_key, repo_name, status="done"): + return + try: + update_index_journal_entry_status( + str(path), + status="done", + workspace_path=repo_key, + repo_name=repo_name, + ) + except Exception: + pass + + +def _mark_journal_failed( + path: Path, + repo_key: str, + repo_name: Optional[str], + error: str, +) -> None: + if _queue_journal_status( + path, + repo_key, + repo_name, + status="failed", + error=error, + remove_on_done=False, + ): + return + try: + update_index_journal_entry_status( + str(path), + status="failed", + error=error, + workspace_path=repo_key, + repo_name=repo_name, + remove_on_done=False, + ) + except Exception: + pass + + +def _path_has_indexed_points(client, collection: str, path: Path) -> Optional[bool]: + try: + filt = models.Filter( + must=[ + models.FieldCondition( + key="metadata.path", match=models.MatchValue(value=str(path)) + ) + ] + ) + points, _ = client.scroll( + collection_name=collection, + scroll_filter=filt, + with_payload=False, + with_vectors=False, + limit=1, + ) + return bool(points) + except Exception: + return None + + +def _verify_delete_committed(client, collection: str, path: Path) -> bool: + has_points = _path_has_indexed_points(client, collection, path) + return has_points is False + + +def _path_has_graph_edges(client, collection: str, path: Path) -> Optional[bool]: + graph_collection = f"{collection}_graph" + try: + # Graph edges normalize paths (Windows -> POSIX separators). Verification must + # query using the same normalization to avoid false "deleted" reports. + raw_path = str(path) + candidates: list[str] = [] + try: + norm_path = str(_normalize_graph_caller_path(raw_path) or "").strip() + if norm_path: + candidates.append(norm_path) + except Exception: + pass + # Back-compat: also consider the raw string and a slash-normalized form in case + # older data was written without normalization. + raw_slash = raw_path.replace("\\", "/").strip() + for v in (raw_slash, raw_path.strip()): + if v and v not in candidates: + candidates.append(v) + + match_obj = ( + models.MatchAny(any=candidates) + if len(candidates) > 1 + else models.MatchValue(value=(candidates[0] if candidates else raw_path)) + ) + filt = models.Filter( + must=[ + models.FieldCondition( + key="caller_path", match=match_obj + ) + ] + ) + points, _ = client.scroll( + collection_name=graph_collection, + scroll_filter=filt, + with_payload=False, + with_vectors=False, + limit=1, + ) + return bool(points) + except Exception as e: + # Missing graph collection means there are no graph edges to verify. + err = str(e).lower() + if "404" in err or "not found" in err or "doesn't exist" in err: + return False + return None + + +def _verify_graph_delete_committed(client, collection: str, path: Path) -> bool: + has_edges = _path_has_graph_edges(client, collection, path) + return has_edges is False + + +def _verify_upsert_committed( + client, + collection: str, + path: Path, + repo_name: Optional[str], + expected_file_hash: Optional[str], + source_text: Optional[str] = None, +) -> bool: + indexed_hash = str( + idx.get_indexed_file_hash(client, collection, str(path)) or "" + ).strip() + expected_hash = str(expected_file_hash or "").strip() + if expected_hash: + if bool(indexed_hash) and indexed_hash == expected_hash: + return True + # Empty/whitespace-only files can legitimately have no indexed points/hash. + try: + if source_text is not None and not source_text.strip(): + has_points = _path_has_indexed_points(client, collection, path) + return has_points is False + except Exception: + pass + return False + has_points = _path_has_indexed_points(client, collection, path) + return has_points is True + + +def _verify_and_update_journal_for_upsert( + p: Path, + client, + collection: str, + repo_key: str, + repo_name: Optional[str], + journal_content_hash: str, + *, + text: Optional[str] = None, + file_hash: Optional[str] = None, +) -> None: + source_text = text + expected_hash = str(file_hash or "").strip() + if source_text is None or not expected_hash: + read_text, read_hash = _read_text_and_sha1(p) + if source_text is None: + source_text = read_text + if not expected_hash: + expected_hash = read_hash + expected_hash = expected_hash or journal_content_hash + if _verify_upsert_committed( + client, + collection, + p, + repo_name, + expected_hash or None, + source_text=source_text, + ): + _mark_journal_done(p, repo_key, repo_name) + else: + _mark_journal_failed( + p, + repo_key, + repo_name, + "upsert_verification_failed", + ) + + +def _finalize_journal_after_index_attempt( + path: Path, + client, + collection: str | None, + repo_key: str, + repo_name: Optional[str], + *, + force_upsert: bool, + journal_content_hash: str, + text: Optional[str] = None, + file_hash: Optional[str] = None, + default_error: Optional[str] = None, + skip_verify_reason: Optional[str] = None, +) -> None: + if force_upsert and client is not None and collection is not None: + # If another worker currently owns this file lock, leave the journal entry + # pending for retry instead of recording a false verification failure. + if skip_verify_reason == "file_locked": + return + _verify_and_update_journal_for_upsert( + path, + client, + collection, + repo_key, + repo_name, + journal_content_hash, + text=text, + file_hash=file_hash, + ) + elif default_error: + _mark_journal_failed(path, repo_key, repo_name, default_error) + else: + _mark_journal_done(path, repo_key, repo_name) + + def _build_subprocess_env( collection: str | None, repo_name: str | None, @@ -105,8 +672,8 @@ def _build_subprocess_env( pass if collection: env["COLLECTION_NAME"] = collection - if QDRANT_URL: - env["QDRANT_URL"] = QDRANT_URL + if watch_config.QDRANT_URL: + env["QDRANT_URL"] = watch_config.QDRANT_URL if repo_name: env["REPO_NAME"] = repo_name return env @@ -114,6 +681,7 @@ def _build_subprocess_env( def _maybe_handle_staging_file( path: Path, + client, collection: str | None, repo_name: str | None, repo_key: str, @@ -121,27 +689,46 @@ def _maybe_handle_staging_file( state_env: Optional[Dict[str, str]], repo_progress: Dict[str, int], started_at: str, + *, + force_upsert: bool = False, + journal_content_hash: str = "", ) -> bool: - if not (is_staging_enabled() and state_env and collection): + if not (state_env and collection): return False - _text, file_hash = _read_text_and_sha1(path) + source_text, file_hash = _read_text_and_sha1(path) if file_hash: try: cached_hash = get_cached_file_hash(str(path), repo_name) if repo_name else None except Exception: cached_hash = None if cached_hash and cached_hash == file_hash: + if force_upsert and client is not None: + if _verify_upsert_committed( + client, + collection, + path, + repo_name, + file_hash or journal_content_hash or None, + source_text=source_text, + ): + safe_print(f"[skip_unchanged] {path} (hash match)") + _log_activity(repo_key, "skipped", path, {"reason": "hash_unchanged"}) + _mark_journal_done(path, repo_key, repo_name) + _advance_progress(repo_progress, repo_key, repo_files, started_at, path) + return True # Fast path: skip if content hash matches cached hash (file unchanged) # Safety: startup health check clears stale cache per-repo - safe_print(f"[skip_unchanged] {path} (hash match)") - _log_activity(repo_key, "skipped", path, {"reason": "hash_unchanged"}) - _advance_progress(repo_progress, repo_key, repo_files, started_at, path) - return True + if not force_upsert: + safe_print(f"[skip_unchanged] {path} (hash match)") + _log_activity(repo_key, "skipped", path, {"reason": "hash_unchanged"}) + _advance_progress(repo_progress, repo_key, repo_files, started_at, path) + return True cmd = [ sys.executable or "python3", - str(ROOT_DIR / "scripts" / "ingest_code.py"), + "-m", + "scripts.ingest_code", "--root", str(path), "--no-skip-unchanged", @@ -178,11 +765,24 @@ def _maybe_handle_staging_file( ) else: safe_print(f"[indexed_subprocess] {path} -> {collection}") + _finalize_journal_after_index_attempt( + path, + client, + collection, + repo_key, + repo_name, + force_upsert=force_upsert, + journal_content_hash=journal_content_hash, + text=source_text, + file_hash=file_hash, + ) + if result.returncode != 0 and force_upsert: + _mark_journal_failed(path, repo_key, repo_name, "subprocess_index_failed") _advance_progress(repo_progress, repo_key, repo_files, started_at, path) return True -def _process_paths( +def _process_paths_impl( paths, client, model, @@ -218,18 +818,92 @@ def _process_paths( pass repo_progress: Dict[str, int] = {key: 0 for key in repo_groups.keys()} + repo_pending_journal_ops: Dict[str, Dict[str, Dict[str, str]]] = {} + repo_move_source_for_dest: Dict[str, Dict[str, str]] = {} + move_dest_keys: set[str] = set() + move_source_keys: set[str] = set() + for repo_path in repo_groups.keys(): + try: + repo_name = _extract_repo_name_from_path(repo_path) + entries = list_pending_index_journal_entries(repo_path, repo_name) + repo_pending_journal_ops[repo_path] = {} + upserts_by_hash: Dict[str, List[str]] = {} + deletes_by_hash: Dict[str, List[str]] = {} + for rec in entries: + path_key = _normalize_cache_key_path(str(rec.get("path") or "")) + op_type = str(rec.get("op_type") or "").strip().lower() + content_hash = str(rec.get("content_hash") or "").strip().lower() + if not path_key: + continue + repo_pending_journal_ops[repo_path][path_key] = { + "op_type": op_type, + "content_hash": content_hash, + } + if not content_hash: + continue + if op_type == "upsert": + upserts_by_hash.setdefault(content_hash, []).append(path_key) + elif op_type == "delete": + deletes_by_hash.setdefault(content_hash, []).append(path_key) + pairs: Dict[str, str] = {} + for content_hash, dest_paths in upserts_by_hash.items(): + src_paths = deletes_by_hash.get(content_hash) or [] + if not src_paths: + continue + src_idx = 0 + for dest_key in dest_paths: + while src_idx < len(src_paths) and src_paths[src_idx] == dest_key: + src_idx += 1 + if src_idx >= len(src_paths): + break + src_key = src_paths[src_idx] + src_idx += 1 + pairs[dest_key] = src_key + move_dest_keys.add(dest_key) + move_source_keys.add(src_key) + repo_move_source_for_dest[repo_path] = pairs + except Exception: + repo_pending_journal_ops[repo_path] = {} + repo_move_source_for_dest[repo_path] = {} + + unique_paths = sorted( + unique_paths, + key=lambda p: ( + 0 + if _normalize_cache_key_path(str(p)) in move_dest_keys + else (2 if _normalize_cache_key_path(str(p)) in move_source_keys else 1), + str(p), + ), + ) + completed_move_sources: set[str] = set() for p in unique_paths: repo_path = _detect_repo_for_file(p) or Path(workspace_path) repo_key = str(repo_path) repo_files = repo_groups.get(repo_key, []) repo_name = _extract_repo_name_from_path(repo_key) + path_key = _normalize_cache_key_path(str(p)) + if path_key in completed_move_sources: + _advance_progress(repo_progress, repo_key, repo_files, started_at, p) + continue + journal_rec = repo_pending_journal_ops.get(repo_key, {}).get(path_key, {}) + journal_op = str(journal_rec.get("op_type") or "").strip().lower() + force_delete = journal_op == "delete" + force_upsert = journal_op == "upsert" + journal_content_hash = str(journal_rec.get("content_hash") or "").strip().lower() + if _is_internal_ignored_path(p): + _log_activity(repo_key, "skipped", p, {"reason": "internal_ignored_path"}) + # Internal metadata paths should never drive indexing or collection creation. + # If they entered the journal via drift repair, mark done and drop. + _mark_journal_done(p, repo_key, repo_name) + _advance_progress(repo_progress, repo_key, repo_files, started_at, p) + continue collection = _get_collection_for_file(p) state_env: Optional[Dict[str, str]] = None try: st = get_workspace_state(repo_key, repo_name) if get_workspace_state else None if isinstance(st, dict): - if is_staging_enabled(): + if _staging_requires_subprocess(st): state_env = st.get("indexing_env") except Exception: state_env = None @@ -240,31 +914,118 @@ def _process_paths( p, collection, repo_name, - env_snapshot=(state_env if is_staging_enabled() else None), + env_snapshot=state_env, ) except Exception as exc: safe_print(f"[commit_ingest_error] {p}: {exc}") _advance_progress(repo_progress, repo_key, repo_files, started_at, p) continue - if not p.exists(): + if force_upsert and not p.exists(): + _log_activity(repo_key, "skipped", p, {"reason": "upsert_missing_file"}) + _mark_journal_failed( + p, + repo_key, + repo_name, + "upsert_missing_file", + ) + _advance_progress(repo_progress, repo_key, repo_files, started_at, p) + continue + + if force_upsert and client is not None and collection is not None: + move_src_key = repo_move_source_for_dest.get(repo_key, {}).get(path_key) + if move_src_key: + move_src_path = Path(move_src_key) + src_collection = _get_collection_for_file(move_src_path) + try: + moved_count, renamed_hash = _rename_in_store( + client, + src_collection, + move_src_path, + p, + collection, + ) + except Exception: + moved_count, renamed_hash = -1, None + if moved_count and moved_count > 0: + try: + if repo_name: + remove_cached_file(str(move_src_path), repo_name) + except Exception: + pass + final_hash = renamed_hash or journal_content_hash + try: + if repo_name and final_hash: + set_cached_file_hash(str(p), final_hash, repo_name) + except Exception: + pass + _log_activity( + repo_key, + "moved", + p, + {"from": str(move_src_path), "chunks": int(moved_count)}, + ) + _mark_journal_done(p, repo_key, repo_name) + _mark_journal_done(move_src_path, repo_key, repo_name) + completed_move_sources.add(move_src_key) + _advance_progress(repo_progress, repo_key, repo_files, started_at, p) + continue + + if force_delete or not p.exists(): + deleted_ok = False if client is not None: try: idx.delete_points_by_path(client, collection, str(p)) + try: + idx.delete_graph_edges_by_path( + client, + collection, + caller_path=str(p), + repo=repo_name, + ) + # Repo tags can drift over time, so always follow a repo-scoped + # delete with a path-only sweep to remove any stale rows left under + # an older/default repo tag. + if repo_name: + idx.delete_graph_edges_by_path( + client, + collection, + caller_path=str(p), + repo=None, + ) + except Exception as graph_exc: + safe_print(f"[deleted:graph_failed] {p} -> {collection}: {graph_exc}") safe_print(f"[deleted] {p} -> {collection}") + deleted_ok = True except Exception: - pass + deleted_ok = False + if deleted_ok and client is not None and collection is not None: + deleted_ok = _verify_delete_committed(client, collection, p) + if deleted_ok and client is not None and collection is not None: + verify_graph_delete = get_boolean_env("WATCH_VERIFY_GRAPH_DELETE", True) + if verify_graph_delete: + deleted_ok = _verify_graph_delete_committed(client, collection, p) try: if repo_name: remove_cached_file(str(p), repo_name) except Exception: pass _log_activity(repo_key, "deleted", p) + if deleted_ok: + _mark_journal_done(p, repo_key, repo_name) + else: + _mark_journal_failed( + p, + repo_key, + repo_name, + "delete_points_or_graph_failed", + ) _advance_progress(repo_progress, repo_key, repo_files, started_at, p) continue if _maybe_handle_staging_file( p, + client, collection, repo_name, repo_key, @@ -272,17 +1033,39 @@ def _process_paths( state_env, repo_progress, started_at, + force_upsert=force_upsert, + journal_content_hash=journal_content_hash, ): continue if client is not None and model is not None: try: + verify_context: Dict[str, Optional[str]] = {} ok = _run_indexing_strategy( - p, client, model, collection, vector_name, model_dim, repo_name + p, + client, + model, + collection, + vector_name, + model_dim, + repo_name, + force_upsert=force_upsert, + verify_context=verify_context if force_upsert else None, ) - except _SkipUnchanged: + except _SkipUnchanged as exc: status = "skipped" safe_print(f"[{status}] {p} -> {collection}") _log_activity(repo_key, "skipped", p, {"reason": "hash_unchanged"}) + _finalize_journal_after_index_attempt( + p, + client, + collection, + repo_key, + repo_name, + force_upsert=force_upsert, + journal_content_hash=journal_content_hash, + text=exc.text, + file_hash=exc.file_hash, + ) _advance_progress(repo_progress, repo_key, repo_files, started_at, p) continue except Exception: @@ -295,6 +1078,7 @@ def _process_paths( "file": str(p), }, ) + _mark_journal_failed(p, repo_key, repo_name, "indexing_error") _advance_progress(repo_progress, repo_key, repo_files, started_at, p) continue @@ -306,10 +1090,35 @@ def _process_paths( except Exception: size = None _log_activity(repo_key, "indexed", p, {"file_size": size}) + _finalize_journal_after_index_attempt( + p, + client, + collection, + repo_key, + repo_name, + force_upsert=force_upsert, + journal_content_hash=journal_content_hash, + text=verify_context.get("text"), + file_hash=verify_context.get("file_hash"), + skip_verify_reason=verify_context.get("skip_verify_reason"), + ) else: _log_activity( repo_key, "skipped", p, {"reason": "no-change-or-error"} ) + _finalize_journal_after_index_attempt( + p, + client, + collection, + repo_key, + repo_name, + force_upsert=force_upsert, + journal_content_hash=journal_content_hash, + text=verify_context.get("text"), + file_hash=verify_context.get("file_hash"), + default_error="no_change_or_error", + skip_verify_reason=verify_context.get("skip_verify_reason"), + ) _advance_progress(repo_progress, repo_key, repo_files, started_at, p) else: safe_print(f"Not processing locally: {p}") @@ -328,12 +1137,44 @@ def _process_paths( pass +def _process_paths( + paths, + client, + model, + vector_name: str, + model_dim: int, + workspace_path: str, +) -> None: + """Process a watcher batch and persist journal status updates efficiently.""" + previous_batch = getattr(_JOURNAL_STATUS_BATCH_LOCAL, "current", None) + batch: List[Dict[str, object]] = [] + _JOURNAL_STATUS_BATCH_LOCAL.current = batch + try: + _process_paths_impl( + paths, + client, + model, + vector_name, + model_dim, + workspace_path, + ) + finally: + if previous_batch is None: + try: + delattr(_JOURNAL_STATUS_BATCH_LOCAL, "current") + except AttributeError: + pass + else: + _JOURNAL_STATUS_BATCH_LOCAL.current = previous_batch + _flush_journal_status_batch(batch) + + def _read_text_and_sha1(path: Path) -> tuple[Optional[str], str]: try: text = path.read_text(encoding="utf-8", errors="ignore") except Exception: text = None - if not text: + if text is None: return text, "" try: file_hash = hashlib.sha1(text.encode("utf-8", errors="ignore")).hexdigest() @@ -350,69 +1191,111 @@ def _run_indexing_strategy( vector_name: str, model_dim: int, repo_name: str | None, + force_upsert: bool = False, + *, + verify_context: Optional[Dict[str, Optional[str]]] = None, ) -> bool: if collection is None: return False - try: - idx.ensure_collection_and_indexes_once(client, collection, model_dim, vector_name) - except Exception: - pass text, file_hash = _read_text_and_sha1(path) + if verify_context is not None: + verify_context["text"] = text + verify_context["file_hash"] = file_hash + verify_context["skip_verify_reason"] = None ok = False if text is not None: try: language = idx.detect_language(path) except Exception: language = "" + try: + is_text_like = bool(idx.is_text_like_language(language)) + except Exception: + is_text_like = False if file_hash: try: cached_hash = get_cached_file_hash(str(path), repo_name) if repo_name else None except Exception: cached_hash = None - if cached_hash and cached_hash == file_hash: + if cached_hash and cached_hash == file_hash and not force_upsert: ok = True - raise _SkipUnchanged() - try: - use_smart, smart_reason = idx.should_use_smart_reindexing(str(path), file_hash) - except Exception: - use_smart, smart_reason = False, "smart_check_failed" - # Bootstrap: if we have no symbol cache yet, still run smart path once - bootstrap = smart_reason == "no_cached_symbols" - if use_smart or bootstrap: - msg_kind = ( - "smart reindexing" - if use_smart - else "bootstrap (no_cached_symbols) for smart reindex" - ) - safe_print( - f"[SMART_REINDEX][watcher] Using {msg_kind} for {path} ({smart_reason})" - ) + raise _SkipUnchanged(text=text, file_hash=file_hash) + + # Repair upserts must materialize points when a path is missing in Qdrant. + # Smart reindex can return "skipped" for unchanged symbols, which is valid + # only when points already exist. + force_full_reindex = False + if force_upsert and client is not None: try: - status = idx.process_file_with_smart_reindexing( - path, - text, - language, - client, - collection, - repo_name, - model, - vector_name, + existing_hash = str( + idx.get_indexed_file_hash(client, collection, str(path)) or "" + ).strip() + except Exception: + existing_hash = "" + if not existing_hash: + has_points = _path_has_indexed_points(client, collection, path) + if has_points is not True: + force_full_reindex = True + + if not is_text_like: + try: + use_smart, smart_reason = idx.should_use_smart_reindexing(str(path), file_hash) + except Exception: + use_smart, smart_reason = False, "smart_check_failed" + # Bootstrap: if we have no symbol cache yet, still run smart path once + bootstrap = smart_reason == "no_cached_symbols" + if (use_smart or bootstrap) and not force_full_reindex: + msg_kind = ( + "smart reindexing" + if use_smart + else "bootstrap (no_cached_symbols) for smart reindex" ) - ok = status in ("success", "skipped") - except Exception as exc: safe_print( - f"[SMART_REINDEX][watcher] Smart reindexing failed for {path}: {exc}" + f"[SMART_REINDEX][watcher] Using {msg_kind} for {path} ({smart_reason})" ) - ok = False - else: - safe_print( - f"[SMART_REINDEX][watcher] Using full reindexing for {path} ({smart_reason})" - ) - # Fallback: full single-file reindex. Pseudo/tags are inlined by default; - # when PSEUDO_DEFER_TO_WORKER=1 we run base-only and rely on backfill. + try: + status = idx.process_file_with_smart_reindexing( + path, + text, + language, + client, + collection, + repo_name, + model, + vector_name, + model_dim=model_dim, + ) + ok = status in ("success", "skipped") + except Exception as exc: + safe_print( + f"[SMART_REINDEX][watcher] Smart reindexing failed for {path}: {exc}" + ) + ok = False + else: + if force_full_reindex: + safe_print( + f"[SMART_REINDEX][watcher] Forcing full reindex for {path} " + "(force_upsert_missing_points)" + ) + safe_print( + f"[SMART_REINDEX][watcher] Using full reindexing for {path} ({smart_reason})" + ) + # Fallback: full single-file reindex. Pseudo/tags are inlined by default; + # when PSEUDO_DEFER_TO_WORKER=1 we run base-only and rely on backfill. if not ok: - pseudo_mode = "off" if get_boolean_env("PSEUDO_DEFER_TO_WORKER") else "full" + try: + idx.ensure_collection_and_indexes_once( + client, collection, model_dim, vector_name + ) + except Exception: + pass + # PSEUDO_DEFER_TO_WORKER is a foreground/background semantics knob; it should + # only disable inline pseudo/tags generation when the backfill worker is enabled. + pseudo_mode = effective_pseudo_mode( + defer_to_worker=get_boolean_env("PSEUDO_DEFER_TO_WORKER"), + backfill_enabled=get_boolean_env("PSEUDO_BACKFILL_ENABLED"), + ) ok = idx.index_single_file( client, model, @@ -423,7 +1306,16 @@ def _run_indexing_strategy( skip_unchanged=False, pseudo_mode=pseudo_mode, repo_name_for_cache=repo_name, + preloaded_text=text, + preloaded_file_hash=file_hash, + preloaded_language=language if text is not None else None, ) + if force_upsert and not ok and verify_context is not None: + try: + if idx.is_file_locked(str(path)): + verify_context["skip_verify_reason"] = "file_locked" + except Exception: + pass return ok diff --git a/scripts/watch_index_core/pseudo.py b/scripts/watch_index_core/pseudo.py index dc7bb0a8..b4b573f5 100644 --- a/scripts/watch_index_core/pseudo.py +++ b/scripts/watch_index_core/pseudo.py @@ -8,6 +8,7 @@ from typing import Optional import scripts.ingest_code as idx +from . import config as watch_config from .utils import get_boolean_env from scripts.workspace_state import ( _cross_process_lock, @@ -17,8 +18,6 @@ is_multi_repo_mode, ) -from .config import ROOT - logger = logging.getLogger(__name__) @@ -27,6 +26,8 @@ def _start_pseudo_backfill_worker( default_collection: str, model_dim: int, vector_name: str, + *, + allow_default_collection_fallback: bool = True, ) -> Optional[threading.Event]: """Start a daemon thread that periodically backfills pseudo/tags. @@ -34,7 +35,12 @@ def _start_pseudo_backfill_worker( or None if the worker was not started (disabled via env). """ - if not get_boolean_env("PSEUDO_DEFER_TO_WORKER"): + # This worker is controlled by PSEUDO_BACKFILL_ENABLED (pseudo/tags) and/or + # GRAPH_EDGES_BACKFILL (graph edges). PSEUDO_DEFER_TO_WORKER only controls + # whether the foreground index path generates pseudo inline. + pseudo_backfill_enabled = get_boolean_env("PSEUDO_BACKFILL_ENABLED") + graph_backfill_enabled = get_boolean_env("GRAPH_EDGES_BACKFILL") + if not (pseudo_backfill_enabled or graph_backfill_enabled): return None try: @@ -49,20 +55,36 @@ def _start_pseudo_backfill_worker( max_points = 256 if max_points <= 0: max_points = 1 + try: + graph_max_files = int( + os.environ.get("GRAPH_EDGES_BACKFILL_MAX_FILES", "128") or 128 + ) + except Exception: + graph_max_files = 128 + if graph_max_files <= 0: + graph_max_files = 1 shutdown_event = threading.Event() def _worker() -> None: while not shutdown_event.is_set(): try: + pseudo_backfill_on = get_boolean_env("PSEUDO_BACKFILL_ENABLED") + graph_backfill_on = get_boolean_env("GRAPH_EDGES_BACKFILL") try: - mappings = get_collection_mappings(search_root=str(ROOT)) + mappings = get_collection_mappings(search_root=str(watch_config.ROOT)) except Exception: mappings = [] if not mappings: - mappings = [ - {"repo_name": None, "collection_name": default_collection}, - ] + # Do not fall back to the default collection unless startup explicitly + # allowed the watcher to touch it. This keeps background backfill from + # recreating collections that the caller intentionally left alone. + if is_multi_repo_mode() or not allow_default_collection_fallback: + mappings = [] + else: + mappings = [ + {"repo_name": None, "collection_name": default_collection}, + ] for mapping in mappings: if shutdown_event.is_set(): break @@ -74,21 +96,52 @@ def _worker() -> None: if is_multi_repo_mode() and repo_name: state_dir = _get_repo_state_dir(repo_name) else: - state_dir = _get_global_state_dir(str(ROOT)) + state_dir = _get_global_state_dir(str(watch_config.ROOT)) lock_path = state_dir / "pseudo.lock" with _cross_process_lock(lock_path): - processed = idx.pseudo_backfill_tick( - client, - coll, - repo_name=repo_name, - max_points=max_points, - dim=model_dim, - vector_name=vector_name, - ) - if processed: - logger.info( - "[pseudo_backfill] repo=%s collection=%s processed=%d", - repo_name or "default", coll, processed, + if pseudo_backfill_on: + processed = idx.pseudo_backfill_tick( + client, + coll, + repo_name=repo_name, + max_points=max_points, + dim=model_dim, + vector_name=vector_name, + ) + if processed: + logger.info( + "[pseudo_backfill] repo=%s collection=%s processed=%d", + repo_name or "default", + coll, + processed, + ) + # Optional: backfill graph edge collection from main points. + # Controlled separately because it may scan large collections over time. + # Run under its own lock to avoid blocking pseudo/tag backfill workers. + if graph_backfill_on: + try: + graph_lock_path = state_dir / "graph_edges.lock" + with _cross_process_lock(graph_lock_path): + files_done = idx.graph_edges_backfill_tick( + client, + coll, + repo_name=repo_name, + max_files=graph_max_files, + ) + if files_done: + logger.info( + "[graph_backfill] repo=%s collection=%s files=%d", + repo_name or "default", + coll, + files_done, + ) + except Exception as exc: + logger.error( + "[graph_backfill] error repo=%s collection=%s: %s", + repo_name or "default", + coll, + exc, + exc_info=True, ) except Exception as exc: logger.error( @@ -110,4 +163,3 @@ def _worker() -> None: __all__ = ["_start_pseudo_backfill_worker"] - diff --git a/scripts/watch_index_core/queue.py b/scripts/watch_index_core/queue.py index ede8835b..ca420842 100644 --- a/scripts/watch_index_core/queue.py +++ b/scripts/watch_index_core/queue.py @@ -3,10 +3,11 @@ from __future__ import annotations import threading +import time from pathlib import Path from typing import Callable, Iterable, List, Set -from .config import DELAY_SECS, LOGGER +from .config import DELAY_SECS, LOGGER, RECENT_FINGERPRINT_TTL_SECS class ChangeQueue: @@ -16,15 +17,21 @@ def __init__(self, process_cb: Callable[[List[Path]], None]): self._lock = threading.Lock() self._paths: Set[Path] = set() self._pending: Set[Path] = set() + self._forced_paths: Set[Path] = set() + self._pending_forced: Set[Path] = set() self._timer: threading.Timer | None = None self._process_cb = process_cb # Serialize processing to avoid concurrent use of TextEmbedding/QdrantClient self._processing_lock = threading.Lock() + self._recent_fingerprints: dict[Path, tuple[tuple[int, int], float]] = {} - def add(self, p: Path) -> None: + def add(self, p: Path, *, force: bool = False) -> None: with self._lock: + already_queued = p in self._paths self._paths.add(p) - if self._timer is not None: + if force: + self._forced_paths.add(p) + if self._timer is not None and not already_queued: try: self._timer.cancel() except Exception as exc: @@ -32,21 +39,101 @@ def add(self, p: Path) -> None: "Failed to cancel timer in ChangeQueue.add", extra={"error": str(exc)}, ) - self._timer = threading.Timer(DELAY_SECS, self._flush) - self._timer.daemon = True - self._timer.start() + if self._timer is None or not already_queued: + self._timer = threading.Timer(DELAY_SECS, self._flush) + self._timer.daemon = True + self._timer.start() + + def stats(self) -> dict[str, int | bool]: + with self._lock: + return { + "queued": len(self._paths), + "pending": len(self._pending), + "forced": len(self._forced_paths), + "pending_forced": len(self._pending_forced), + "processing": self._processing_lock.locked(), + } + + def _fingerprint_path(self, p: Path) -> tuple[int, int] | None: + try: + st = p.stat() + return ( + int(getattr(st, "st_size", 0)), + int(getattr(st, "st_mtime_ns", int(st.st_mtime * 1e9))), + ) + except Exception: + return None + + def _filter_recent_paths( + self, + paths: Iterable[Path], + *, + forced_paths: Iterable[Path] | None = None, + ) -> list[Path]: + ttl = float(RECENT_FINGERPRINT_TTL_SECS) + forced = set(forced_paths or []) + if ttl <= 0: + return list(paths) + + now = time.time() + keep: list[Path] = [] + for p in paths: + if p in forced: + keep.append(p) + continue + fp = self._fingerprint_path(p) + if fp is None: + keep.append(p) + continue + prev = self._recent_fingerprints.get(p) + if prev is not None: + prev_fp, prev_ts = prev + if prev_fp == fp and (now - prev_ts) < ttl: + continue + keep.append(p) + return keep + + def _mark_recent_paths(self, paths: Iterable[Path]) -> None: + ttl = float(RECENT_FINGERPRINT_TTL_SECS) + if ttl <= 0: + return + now = time.time() + for p in paths: + fp = self._fingerprint_path(p) + if fp is None: + continue + self._recent_fingerprints[p] = (fp, now) + # Keep at least a 1s grace for small TTLs while using a proportional + # buffer for larger TTLs so stale handled fingerprints age out cleanly. + cutoff = now - max(ttl * 2.0, ttl + 1.0) + stale = [p for p, (_fp, ts) in self._recent_fingerprints.items() if ts < cutoff] + for p in stale: + self._recent_fingerprints.pop(p, None) + + def _drain_pending(self) -> tuple[list[Path], Set[Path]] | None: + with self._lock: + if not self._pending: + return None + todo = list(self._pending) + todo_forced = {p for p in todo if p in self._pending_forced} + self._pending.clear() + self._pending_forced.clear() + return todo, todo_forced def _flush(self) -> None: # Grab current batch with self._lock: paths = list(self._paths) + forced_paths = {p for p in paths if p in self._forced_paths} self._paths.clear() + self._forced_paths.difference_update(paths) self._timer = None # Try to run the processor exclusively; if busy, queue and return if not self._processing_lock.acquire(blocking=False): with self._lock: self._pending.update(paths) + self._pending_forced.update(forced_paths) if self._timer is None: # schedule a follow-up flush to pick up pending when free self._timer = threading.Timer(DELAY_SECS, self._flush) @@ -56,15 +143,24 @@ def _flush(self) -> None: try: # Per-file locking in index_single_file handles indexer/watcher coordination todo: Iterable[Path] = paths + todo_forced: Set[Path] = set(forced_paths) while True: + filtered_todo = self._filter_recent_paths(todo, forced_paths=todo_forced) + if not filtered_todo: + pending = self._drain_pending() + if pending is None: + break + todo, todo_forced = pending + continue try: - self._process_cb(list(todo)) + self._process_cb(list(filtered_todo)) + self._mark_recent_paths(filtered_todo) except Exception as exc: # Log processing error via structured logging try: LOGGER.error( "Processing batch failed in ChangeQueue._flush", - extra={"error": str(exc), "batch_size": len(list(todo))}, + extra={"error": str(exc), "batch_size": len(filtered_todo)}, exc_info=True, ) except Exception as inner_exc: # pragma: no cover - logging fallback @@ -79,11 +175,10 @@ def _flush(self) -> None: except Exception: pass # Last resort: can't even print # drain any pending accumulated during processing - with self._lock: - if not self._pending: - break - todo = list(self._pending) - self._pending.clear() + pending = self._drain_pending() + if pending is None: + break + todo, todo_forced = pending finally: self._processing_lock.release() diff --git a/scripts/watch_index_core/utils.py b/scripts/watch_index_core/utils.py index 999daa5a..5f4086fd 100644 --- a/scripts/watch_index_core/utils.py +++ b/scripts/watch_index_core/utils.py @@ -8,7 +8,8 @@ from watchdog.observers import Observer import scripts.ingest_code as idx -from .config import LOGGER, ROOT, default_collection_name +from . import config as watch_config +from .config import LOGGER, default_collection_name from scripts.workspace_state import ( _extract_repo_name_from_path, PLACEHOLDER_COLLECTION_NAMES, @@ -93,13 +94,14 @@ def create_observer(use_polling: bool, observer_cls: Type[Observer] = Observer) def _detect_repo_for_file(file_path: Path) -> Optional[Path]: """Detect repository root for a file under WATCH root.""" + root = watch_config.ROOT try: - rel_path = file_path.resolve().relative_to(ROOT.resolve()) + rel_path = file_path.resolve().relative_to(root.resolve()) except Exception: return None if not rel_path.parts: - return ROOT - return ROOT / rel_path.parts[0] + return root + return root / rel_path.parts[0] def _repo_name_or_none(repo_path: Optional[Path]) -> Optional[str]: diff --git a/scripts/workspace_state.py b/scripts/workspace_state.py index b0cb28df..e7e192a5 100644 --- a/scripts/workspace_state.py +++ b/scripts/workspace_state.py @@ -9,6 +9,7 @@ - Multi-repo support with per-repo state files """ import json +import logging import os import re import uuid @@ -22,6 +23,8 @@ _CANONICAL_SLUG_RE = re.compile(r"^.+-[0-9a-f]{16}$") _SLUGGED_REPO_RE = re.compile(r"^.+-[0-9a-f]{16}(?:_old)?$") +INTERNAL_STATE_TOP_LEVEL_DIRS = frozenset({".codebase", ".git", "__pycache__"}) +logger = logging.getLogger(__name__) _managed_slug_cache_lock = threading.Lock() _managed_slug_cache: set[str] = set() _managed_slug_cache_neg: set[str] = set() @@ -112,7 +115,7 @@ def _server_managed_slug_from_path(path: Path) -> Optional[str]: return None work_dir = Path(os.environ.get("WORK_DIR") or os.environ.get("WORKDIR") or "/work") - marker = work_dir / ".codebase" / "repos" / slug / ".ctxce_managed_upload" + marker = work_dir / STATE_DIRNAME / "repos" / slug / ".ctxce_managed_upload" try: is_managed = marker.exists() except OSError: @@ -134,7 +137,8 @@ def _server_managed_slug_from_path(path: Path) -> Optional[str]: STATE_DIRNAME = ".codebase" STATE_FILENAME = "state.json" CACHE_FILENAME = "cache.json" -PLACEHOLDER_COLLECTION_NAMES = {"", "default-collection", "my-collection"} +INDEX_JOURNAL_FILENAME = "index_journal.json" +PLACEHOLDER_COLLECTION_NAMES = {"", "codebase"} class IndexingProgress(TypedDict, total=False): files_processed: int @@ -184,6 +188,63 @@ class StagingInfo(TypedDict, total=False): repo_name: Optional[str] +class MaintenanceInfo(TypedDict, total=False): + last_empty_dir_sweep_at: Optional[str] + last_consistency_audit_at: Optional[str] + last_consistency_audit_summary: Optional[Dict[str, Any]] + + +class IndexJournalRecord(TypedDict, total=False): + path: str + op_type: str + content_hash: Optional[str] + status: str + attempts: int + created_at: str + updated_at: str + last_error: Optional[str] + + +_INDEX_JOURNAL_STATUSES = frozenset({"pending", "in_progress", "failed", "done"}) + + +def _normalize_index_journal_status(value: Any, *, invalid: str = "pending") -> str: + status = str(value or "pending").strip().lower() + return status if status in _INDEX_JOURNAL_STATUSES else invalid + + +def _coerce_index_journal_attempts( + value: Any, + *, + path: Optional[str] = None, + warn: bool = False, +) -> int: + try: + attempts = int(value or 0) + except (ValueError, TypeError): + attempts = 0 + if warn: + logger.warning( + "workspace_state::invalid_journal_attempts", + extra={"attempts": value, "path": path or ""}, + ) + return max(0, attempts) + + +def _index_journal_retry_delay_seconds() -> float: + try: + return max(0.0, float(os.environ.get("INDEX_JOURNAL_RETRY_DELAY_SECS", "5") or 5)) + except Exception: + return 5.0 + + +def _index_journal_max_attempts() -> int: + try: + return max(0, int(os.environ.get("INDEX_JOURNAL_MAX_ATTEMPTS", "0") or 0)) + except Exception: + return 0 + + class WorkspaceState(TypedDict, total=False): created_at: str updated_at: str @@ -204,6 +265,7 @@ class WorkspaceState(TypedDict, total=False): active_repo_slug: Optional[str] serving_repo_slug: Optional[str] staging: Optional[StagingInfo] + maintenance: Optional[MaintenanceInfo] def is_multi_repo_mode() -> bool: """Check if multi-repo mode is enabled.""" @@ -226,6 +288,16 @@ def logical_repo_reuse_enabled() -> bool: "on", } + +def bindmount_repo_detection_enabled() -> bool: + """Allow git-based repo inference for bindmount-style deployments.""" + return os.environ.get("CTXCE_BINDMOUNT_REPO_DETECTION", "").strip().lower() in { + "1", + "true", + "yes", + "on", + } + _state_lock = threading.Lock() # Track last-used timestamps for cleanup of idle workspace locks _state_locks: Dict[str, threading.RLock] = {} @@ -233,7 +305,26 @@ def logical_repo_reuse_enabled() -> bool: def _resolve_workspace_root() -> str: """Determine the default workspace root path.""" - return os.environ.get("WORKSPACE_PATH") or os.environ.get("WATCH_ROOT") or "/work" + return ( + os.environ.get("CTXCE_METADATA_ROOT") + or os.environ.get("WORKSPACE_PATH") + or os.environ.get("WATCH_ROOT") + or "/work" + ) + + +def _configured_workspace_roots() -> List[Path]: + roots: List[Path] = [] + for key in ("CTXCE_METADATA_ROOT", "WORKSPACE_PATH", "WATCH_ROOT", "WORK_DIR", "WORKDIR"): + raw = (os.environ.get(key) or "").strip() + if not raw: + continue + try: + roots.append(Path(raw).resolve()) + except Exception: + roots.append(Path(raw)) + return roots + def _resolve_repo_context( workspace_path: Optional[str] = None, @@ -247,14 +338,47 @@ def _resolve_repo_context( return resolved_workspace, repo_name if workspace_path: - detected = _detect_repo_name_from_path(Path(workspace_path)) - if detected: - return resolved_workspace, detected + try: + requested = Path(workspace_path).resolve() + workspace_root = Path(_resolve_workspace_root()).resolve() + except Exception: + requested = Path(workspace_path) + workspace_root = Path(_resolve_workspace_root()) + if requested != workspace_root: + if any(requested == root for root in _configured_workspace_roots()): + return resolved_workspace, None + detected = _detect_repo_name_from_path(requested) + if detected: + return resolved_workspace, detected return resolved_workspace, None return resolved_workspace, repo_name + +def _get_repo_workspace_dir( + repo_name: str, + workspace_path: Optional[str] = None, +) -> Path: + try: + base_dir = Path(workspace_path or _resolve_workspace_root()).resolve() + except Exception: + base_dir = Path(workspace_path or _resolve_workspace_root()).absolute() + if base_dir.name == repo_name: + return base_dir + host_index_path = (os.environ.get("HOST_INDEX_PATH") or "").strip() + if host_index_path: + host_index_root = Path(host_index_path) + if not host_index_root.is_absolute(): + host_index_root = base_dir / host_index_root + candidate = host_index_root.resolve() / repo_name + if candidate.exists() or (candidate / STATE_DIRNAME).exists(): + return candidate + dev_workspace_candidate = base_dir / "dev-workspace" / repo_name + if dev_workspace_candidate.exists() or (dev_workspace_candidate / STATE_DIRNAME).exists(): + return dev_workspace_candidate + return base_dir / repo_name + def _get_state_lock(workspace_path: Optional[str] = None, repo_name: Optional[str] = None) -> threading.RLock: """Get or create a lock for the workspace or repo state and track usage.""" if repo_name and is_multi_repo_mode(): @@ -268,13 +392,52 @@ def _get_state_lock(workspace_path: Optional[str] = None, repo_name: Optional[st _state_lock_last_used[key] = time.time() return _state_locks[key] -def _get_repo_state_dir(repo_name: str) -> Path: +def _get_repo_state_dir( + repo_name: str, + workspace_path: Optional[str] = None, +) -> Path: """Get the state directory for a repository.""" - base_dir = Path(os.environ.get("WORKSPACE_PATH") or os.environ.get("WATCH_ROOT") or "/work") + workspace_root = Path(_resolve_workspace_root()).resolve() + base_dir = Path(workspace_path or str(workspace_root)).resolve() + global_repo_state_dir = workspace_root / STATE_DIRNAME / "repos" / repo_name if is_multi_repo_mode(): - return base_dir / STATE_DIRNAME / "repos" / repo_name + # Canonical multi-repo metadata layout is shared under workspace root. + return global_repo_state_dir return base_dir / STATE_DIRNAME + +def _is_repo_local_metadata_path(path: Path) -> bool: + try: + parts = path.resolve().parts + except Exception: + parts = path.parts + try: + idx = parts.index(STATE_DIRNAME) + except ValueError: + return False + if idx > 0 and _SLUGGED_REPO_RE.match(parts[idx - 1] or ""): + return True + if "repos" in parts: + ridx = parts.index("repos") + if ridx + 1 < len(parts) and _SLUGGED_REPO_RE.match(parts[ridx + 1] or ""): + return True + return False + + +def _apply_runtime_metadata_mode(path: Path) -> None: + try: + is_dir = path.is_dir() + except Exception: + is_dir = False + if _is_repo_local_metadata_path(path): + mode = 0o777 if is_dir else 0o666 + else: + mode = 0o775 if is_dir else 0o664 + try: + os.chmod(path, mode) + except Exception: + pass + def _get_state_path(workspace_path: str) -> Path: """Get the path to the state.json file for a workspace.""" workspace = Path(workspace_path).resolve() @@ -592,15 +755,13 @@ def _git_remote_repo_name(repo_path: Path) -> Optional[str]: def _detect_repo_name_from_path(path: Path) -> str: - """Detect repository name from path using git remote origin URL. + """Detect repository name from managed upload/workspace path structure. - This ensures consistency with how the MCP server detects repos during search. Priority: - 1. Fast-path for server-managed uploads and workspace-relative paths - 2. Git remote origin URL (canonical repo name like 'Context-Engine') - 3. Git toplevel directory name (folder name like 'Context-Engine-hash') - 4. Walk up to find .git and return that folder name - 5. Return parent folder name as fallback + 1. Server-managed upload slug markers + 2. Workspace-relative first path segment + 3. Bindmount git inference when CTXCE_BINDMOUNT_REPO_DETECTION=1 + 4. Structure/name fallback """ slug = _server_managed_slug_from_path(path) if slug: @@ -623,29 +784,29 @@ def _detect_repo_name_from_path(path: Path) -> str: rel = resolved.relative_to(ws_root) if rel.parts: candidate = rel.parts[0] - if candidate not in {".codebase", ".git", "__pycache__"}: + if candidate not in INTERNAL_STATE_TOP_LEVEL_DIRS: return candidate except Exception: pass - try: - base = path if path.is_dir() else path.parent - git_name = _git_remote_repo_name(base) - if git_name: - return git_name - except Exception: - pass - try: - # Walk up to find .git - cur = path if path.is_dir() else path.parent - for p in [cur] + list(cur.parents): - try: - if (p / ".git").exists(): - return p.name - except Exception: - continue - except Exception: - pass + if bindmount_repo_detection_enabled(): + try: + base = path if path.is_dir() else path.parent + git_name = _git_remote_repo_name(base) + if git_name: + return git_name + except Exception: + pass + try: + cur = path if path.is_dir() else path.parent + for p in [cur] + list(cur.parents): + try: + if (p / ".git").exists(): + return p.name + except Exception: + continue + except Exception: + pass try: structure_name = _detect_repo_name_from_path_by_structure(path) @@ -672,12 +833,7 @@ def _atomic_write_state(state_path: Path, state: WorkspaceState) -> None: with open(temp_path, 'w', encoding='utf-8') as f: json.dump(state, f, indent=2, ensure_ascii=False) temp_path.replace(state_path) - # Ensure state/cache files are group-writable so multiple processes - # (upload service, watcher, indexer) can update them. - try: - os.chmod(state_path, 0o664) - except PermissionError: - pass + _apply_runtime_metadata_mode(state_path) except Exception: # Clean up temp file if something went wrong try: @@ -705,7 +861,7 @@ def get_workspace_state( lock_scope_path: Path if is_multi_repo_mode() and repo_name: - state_dir = _get_repo_state_dir(repo_name) + state_dir = _get_repo_state_dir(repo_name, workspace_path) try: ws_root = Path(_resolve_workspace_root()) ws_dir = ws_root / repo_name @@ -717,12 +873,7 @@ def get_workspace_state( except Exception: return {} state_dir.mkdir(parents=True, exist_ok=True) - # Ensure repo state dir is group-writable so root upload service and - # non-root watcher/indexer processes can both write state/cache files. - try: - os.chmod(state_dir, 0o775) - except Exception: - pass + _apply_runtime_metadata_mode(state_dir) state_path = state_dir / STATE_FILENAME lock_scope_path = state_dir else: @@ -802,7 +953,7 @@ def update_workspace_state( # Allow updates when the repo state dir exists, even if the workspace # directory is not present (e.g. dev-remote simulations where only # .codebase state is persisted). - state_dir = _get_repo_state_dir(repo_name) + state_dir = _get_repo_state_dir(repo_name, workspace_path) if not (ws_root / repo_name).exists() and not state_dir.exists(): return {} except Exception: @@ -823,8 +974,9 @@ def update_workspace_state( state["updated_at"] = datetime.now().isoformat() if is_multi_repo_mode() and repo_name: - state_dir = _get_repo_state_dir(repo_name) + state_dir = _get_repo_state_dir(repo_name, workspace_path) state_dir.mkdir(parents=True, exist_ok=True) + _apply_runtime_metadata_mode(state_dir) state_path = state_dir / STATE_FILENAME else: try: @@ -1245,8 +1397,9 @@ def log_activity( return except Exception: return - state_dir = _get_repo_state_dir(repo_name) + state_dir = _get_repo_state_dir(repo_name, workspace_path) state_dir.mkdir(parents=True, exist_ok=True) + _apply_runtime_metadata_mode(state_dir) state_path = state_dir / STATE_FILENAME lock_path = state_path.with_suffix(".lock") @@ -1331,6 +1484,37 @@ def _collection_name_for_repo_slug(normalized_repo: str, *, is_old_slug: bool) - return None +def _coerce_collection_repo_name(repo_name: Optional[str]) -> Optional[str]: + if not repo_name: + return None + + value = str(repo_name).strip() + if not value: + return None + + if "/" not in value and "\\" not in value: + return value + + try: + path = Path(value).resolve() + except Exception: + path = Path(value) + + try: + workspace_root = Path(_resolve_workspace_root()).resolve() + except Exception: + workspace_root = Path(_resolve_workspace_root()) + + try: + if path == workspace_root: + return None + except Exception: + pass + + detected = _extract_repo_name_from_path(str(path)) + return detected or None + + def get_collection_name(repo_name: Optional[str] = None) -> str: """Get collection name for repository or workspace. @@ -1338,7 +1522,7 @@ def get_collection_name(repo_name: Optional[str] = None) -> str: 1. Explicit COLLECTION_NAME env var - master override when set to a real value (if repo_name is an *_old clone, append _old to the override unless already present) 2. Derive from repo slug (including *_old suffix handling) - 3. Fallback: "global-collection" + 3. Fallback: DEFAULT_COLLECTION, COLLECTION_NAME, or "codebase" This ensures COLLECTION_NAME works as a master override in both local dev and container environments, while still allowing deterministic derivation @@ -1355,6 +1539,7 @@ def get_collection_name(repo_name: Optional[str] = None) -> str: pass return env_coll + repo_name = _coerce_collection_repo_name(repo_name) normalized = _normalize_repo_name_for_collection(repo_name) if repo_name else None is_old_slug = False try: @@ -1369,8 +1554,7 @@ def get_collection_name(repo_name: Optional[str] = None) -> str: if derived: return derived - # Default fallback - return "global-collection" + return os.environ.get("DEFAULT_COLLECTION") or os.environ.get("COLLECTION_NAME") or "codebase" def _detect_repo_name_from_path_by_structure(path: Path) -> str: """Detect repository name from path structure (fallback when git is unavailable).""" @@ -1405,7 +1589,7 @@ def _detect_repo_name_from_path_by_structure(path: Path) -> str: continue repo_name = rel_path.parts[0] - if repo_name in (".codebase", ".git", "__pycache__"): + if repo_name in INTERNAL_STATE_TOP_LEVEL_DIRS: continue repo_path = base / repo_name @@ -1425,7 +1609,8 @@ def _normalize_repo_slug(candidate: Optional[str]) -> Optional[str]: def _extract_repo_name_from_path(workspace_path: str) -> str: """Extract repository slug or canonical name from workspace path. - Accepts canonical slugs (repo-hash), `_old` slugs, and falls back to git remote name. + Accepts managed upload slugs and workspace-relative repo paths. Git-based + bindmount inference is opt-in via CTXCE_BINDMOUNT_REPO_DETECTION=1. """ if not workspace_path: return "" @@ -1440,14 +1625,30 @@ def _extract_repo_name_from_path(workspace_path: str) -> str: return slug try: - repo_path = path if path.is_dir() else path.parent - if (repo_path / ".git").exists(): - name = _git_remote_repo_name(repo_path) - if name: - return name + workspace_root = Path(_resolve_workspace_root()).resolve() + except Exception: + workspace_root = Path(_resolve_workspace_root()) + + try: + rel = path.relative_to(workspace_root) + if not rel.parts: + return "" + candidate = rel.parts[0] + if candidate not in INTERNAL_STATE_TOP_LEVEL_DIRS: + return candidate except Exception: pass + if bindmount_repo_detection_enabled(): + try: + repo_path = path if path.is_dir() else path.parent + if (repo_path / ".git").exists(): + name = _git_remote_repo_name(repo_path) + if name: + return name + except Exception: + pass + try: candidate = _normalize_repo_slug(path.name) if candidate: @@ -1575,10 +1776,483 @@ def _write_cache(workspace_path: str, cache: Dict[str, Any]) -> None: pass -def get_cached_file_hash(file_path: str, repo_name: Optional[str] = None) -> str: +def _get_index_journal_path( + workspace_path: Optional[str] = None, repo_name: Optional[str] = None +) -> Path: + workspace_path, repo_name = _resolve_repo_context(workspace_path, repo_name) + if repo_name: + state_dir = _get_repo_state_dir(repo_name, workspace_path) + else: + state_dir = _get_global_state_dir(workspace_path) + return state_dir / INDEX_JOURNAL_FILENAME + + +def _discover_journal_repositories( + workspace_path: Optional[str] = None, +) -> List[tuple[str, Optional[str]]]: + """Find repository journals visible from a workspace or metadata root.""" + root_path = Path(workspace_path or _resolve_workspace_root()).resolve() + multi_repo_mode = is_multi_repo_mode() + repo_candidates: set[str] = set() + + # A status query should only inspect repositories with an actual journal. + # The workspace can contain arbitrary sibling directories and scanning all + # of them made the operator-facing status path noisy and unnecessarily slow. + try: + for repo_root in root_path.iterdir(): + if not repo_root.is_dir(): + continue + if (repo_root / STATE_DIRNAME / INDEX_JOURNAL_FILENAME).is_file(): + repo_candidates.add(repo_root.name) + except OSError: + pass + + state_roots = [root_path / STATE_DIRNAME / "repos"] + try: + metadata_state_root = ( + Path(_resolve_workspace_root()).resolve() / STATE_DIRNAME / "repos" + ) + if metadata_state_root != state_roots[0]: + state_roots.append(metadata_state_root) + except OSError: + pass + + for state_root in state_roots: + try: + if not state_root.exists(): + continue + for state_dir in state_root.iterdir(): + if state_dir.is_dir() and ( + state_dir / INDEX_JOURNAL_FILENAME + ).is_file(): + repo_candidates.add(state_dir.name) + except OSError: + continue + + return [ + ( + candidate, + None if multi_repo_mode else str(root_path / candidate), + ) + for candidate in sorted(repo_candidates) + ] + + +def _read_index_journal_file_uncached(journal_path: Path) -> Dict[str, Any]: + try: + with journal_path.open("r", encoding="utf-8-sig") as f: + obj = json.load(f) + if isinstance(obj, dict): + operations = obj.get("operations", {}) + if isinstance(operations, dict): + return obj + except (OSError, json.JSONDecodeError, ValueError): + pass + now = datetime.now().isoformat() + return {"version": 1, "operations": {}, "created_at": now, "updated_at": now} + + +def _write_index_journal( + workspace_path: Optional[str], + repo_name: Optional[str], + journal: Dict[str, Any], +) -> None: + workspace_path, repo_name = _resolve_repo_context(workspace_path, repo_name) + lock = _get_state_lock(workspace_path, repo_name) + with lock: + journal_path = _get_index_journal_path(workspace_path, repo_name) + journal_path.parent.mkdir(parents=True, exist_ok=True) + _apply_runtime_metadata_mode(journal_path.parent) + lock_path = journal_path.with_suffix(journal_path.suffix + ".lock") + with _cross_process_lock(lock_path): + tmp = journal_path.with_suffix(f".tmp.{uuid.uuid4().hex[:8]}") + try: + with open(tmp, "w", encoding="utf-8") as f: + json.dump(journal, f, ensure_ascii=False, indent=2) + tmp.replace(journal_path) + _apply_runtime_metadata_mode(journal_path) + finally: + try: + tmp.unlink(missing_ok=True) + except Exception: + pass + + +def _update_index_journal( + workspace_path: Optional[str], + repo_name: Optional[str], + mutator, +) -> Dict[str, Any]: + workspace_path, repo_name = _resolve_repo_context(workspace_path, repo_name) + lock = _get_state_lock(workspace_path, repo_name) + with lock: + journal_path = _get_index_journal_path(workspace_path, repo_name) + journal_path.parent.mkdir(parents=True, exist_ok=True) + _apply_runtime_metadata_mode(journal_path.parent) + lock_path = journal_path.with_suffix(journal_path.suffix + ".lock") + with _cross_process_lock(lock_path): + journal = _read_index_journal_file_uncached(journal_path) + mutator(journal) + journal["updated_at"] = datetime.now().isoformat() + tmp = journal_path.with_suffix(f".tmp.{uuid.uuid4().hex[:8]}") + try: + with open(tmp, "w", encoding="utf-8") as f: + json.dump(journal, f, ensure_ascii=False, indent=2) + tmp.replace(journal_path) + _apply_runtime_metadata_mode(journal_path) + finally: + try: + tmp.unlink(missing_ok=True) + except Exception: + pass + return journal + + +def upsert_index_journal_entries( + entries: List[Dict[str, Any]], + *, + workspace_path: Optional[str] = None, + repo_name: Optional[str] = None, +) -> Dict[str, Any]: + """Persist or replace repo-scoped index journal entries keyed by normalized path.""" + normalized_entries: List[IndexJournalRecord] = [] + now = datetime.now().isoformat() + for entry in entries or []: + path = _normalize_cache_key_path(str(entry.get("path") or "")) + op_type = str(entry.get("op_type") or "").strip().lower() + if not path or op_type not in {"upsert", "delete"}: + continue + content_hash = str(entry.get("content_hash") or "").strip() or None + status = _normalize_index_journal_status(entry.get("status")) + attempts = _coerce_index_journal_attempts(entry.get("attempts", 0)) + last_error = entry.get("last_error") + if last_error is not None: + last_error = str(last_error) + normalized_entries.append( + { + "path": path, + "op_type": op_type, + "content_hash": content_hash, + "status": status, + "attempts": attempts, + "created_at": str(entry.get("created_at") or now), + "updated_at": str(entry.get("updated_at") or now), + "last_error": last_error, + } + ) + + def _mutate(journal: Dict[str, Any]) -> None: + ops = journal.setdefault("operations", {}) + if not isinstance(ops, dict): + ops = {} + journal["operations"] = ops + for entry in normalized_entries: + ops[entry["path"]] = entry + + return _update_index_journal(workspace_path, repo_name, _mutate) + + +def clear_index_journal_entries( + *, + workspace_path: Optional[str] = None, + repo_name: Optional[str] = None, +) -> int: + """Remove all operations from a workspace/repo index journal.""" + removed = 0 + + def _mutate(journal: Dict[str, Any]) -> None: + nonlocal removed + ops = journal.get("operations", {}) + if isinstance(ops, dict): + removed = len(ops) + journal["operations"] = {} + + _update_index_journal(workspace_path, repo_name, _mutate) + return removed + + +def update_index_journal_entries_status( + entries: List[Dict[str, Any]], + *, + workspace_path: Optional[str] = None, + repo_name: Optional[str] = None, + remove_on_done: bool = True, +) -> Dict[str, Any]: + """Update many journal records with one read/modify/write transaction.""" + updates: Dict[str, Dict[str, Any]] = {} + for entry in entries or []: + normalized_path = _normalize_cache_key_path(str(entry.get("path") or "")) + if not normalized_path: + continue + status = _normalize_index_journal_status(entry.get("status"), invalid="failed") + updates[normalized_path] = { + "status": status, + "error": str(entry.get("error") or "").strip() or None, + "remove_on_done": bool(entry.get("remove_on_done", remove_on_done)), + } + + def _mutate(journal: Dict[str, Any]) -> None: + ops = journal.setdefault("operations", {}) + if not isinstance(ops, dict): + ops = {} + journal["operations"] = ops + now = datetime.now().isoformat() + for normalized_path, update in updates.items(): + rec = ops.get(normalized_path) + if not isinstance(rec, dict): + continue + status = update["status"] + if status == "done" and update["remove_on_done"]: + ops.pop(normalized_path, None) + continue + rec["status"] = status + rec["updated_at"] = now + rec["attempts"] = _coerce_index_journal_attempts(rec.get("attempts")) + 1 + rec["last_error"] = update["error"] + ops[normalized_path] = rec + + return _update_index_journal(workspace_path, repo_name, _mutate) + + +def get_index_journal_summary( + workspace_path: Optional[str] = None, + repo_name: Optional[str] = None, +) -> Dict[str, Any]: + """Return compact, operator-friendly journal counts and retry details.""" + workspace_path, repo_name = _resolve_repo_context(workspace_path, repo_name) + + # In multi-repo mode the watcher owns the workspace root and each repo has + # its own journal. Aggregate the root view so an operator can inspect one + # status endpoint without knowing every repo slug first. + if is_multi_repo_mode() and not repo_name: + root_path = Path(workspace_path).resolve() + repositories: Dict[str, Dict[str, Any]] = {} + for candidate, candidate_workspace_path in _discover_journal_repositories(workspace_path): + repositories[candidate] = get_index_journal_summary( + workspace_path=candidate_workspace_path or str(root_path), + repo_name=candidate, + ) + if repositories: + aggregate_counts = { + "pending": 0, + "in_progress": 0, + "failed": 0, + "done": 0, + "unknown": 0, + } + sample_errors: List[Dict[str, Any]] = [] + oldest_retryable_at: Optional[str] = None + max_attempts = 0 + total = 0 + retryable = 0 + outstanding = 0 + exhausted = 0 + updated_at = "" + for candidate, summary in repositories.items(): + total += int(summary.get("total") or 0) + retryable += int(summary.get("retryable") or 0) + outstanding += int(summary.get("outstanding") or 0) + exhausted += int(summary.get("exhausted") or 0) + for status, count in aggregate_counts.items(): + aggregate_counts[status] += int((summary.get("counts") or {}).get(status) or 0) + candidate_oldest = str(summary.get("oldest_retryable_at") or "") + if candidate_oldest and (not oldest_retryable_at or candidate_oldest < oldest_retryable_at): + oldest_retryable_at = candidate_oldest + max_attempts = max(max_attempts, int(summary.get("max_attempts") or 0)) + updated_at = max(updated_at, str(summary.get("updated_at") or "")) + for error in summary.get("sample_errors") or []: + if len(sample_errors) >= 5: + break + sample_errors.append({"repo": candidate, **error}) + return { + "journal_path": str(root_path / STATE_DIRNAME / "repos"), + "total": total, + "counts": aggregate_counts, + "retryable": retryable, + "outstanding": outstanding, + "exhausted": exhausted, + "oldest_retryable_at": oldest_retryable_at, + "max_attempts": max_attempts, + "sample_errors": sample_errors, + "updated_at": updated_at, + "repositories": repositories, + } + + journal_path = _get_index_journal_path(workspace_path, repo_name) + journal = _read_index_journal_file_uncached(journal_path) + operations = journal.get("operations", {}) + if not isinstance(operations, dict): + operations = {} + + counts = {"pending": 0, "in_progress": 0, "failed": 0, "done": 0, "unknown": 0} + retryable = 0 + exhausted = 0 + oldest_updated_at: Optional[str] = None + max_attempts = 0 + journal_max_attempts = _index_journal_max_attempts() + sample_errors: List[Dict[str, Any]] = [] + for raw_path, raw_record in operations.items(): + if not isinstance(raw_record, dict): + counts["unknown"] += 1 + continue + status = _normalize_index_journal_status( + raw_record.get("status"), invalid="unknown" + ) + counts[status] += 1 + attempts = _coerce_index_journal_attempts(raw_record.get("attempts")) + max_attempts = max(max_attempts, attempts) + if status in {"pending", "failed"}: + if journal_max_attempts > 0 and attempts >= journal_max_attempts: + exhausted += 1 + else: + retryable += 1 + updated_at = str(raw_record.get("updated_at") or raw_record.get("created_at") or "") + if updated_at and (oldest_updated_at is None or updated_at < oldest_updated_at): + oldest_updated_at = updated_at + error = str(raw_record.get("last_error") or "").strip() + if error and len(sample_errors) < 5: + sample_errors.append({"path": str(raw_record.get("path") or raw_path), "error": error}) + + return { + "journal_path": str(journal_path), + "total": len(operations), + "counts": counts, + "retryable": retryable, + "outstanding": sum( + count for status, count in counts.items() if status != "done" + ), + "exhausted": exhausted, + "oldest_retryable_at": oldest_updated_at, + "max_attempts": max_attempts, + "sample_errors": sample_errors, + "updated_at": str(journal.get("updated_at") or ""), + } + + +def list_pending_index_journal_entries( + workspace_path: Optional[str] = None, + repo_name: Optional[str] = None, +) -> List[IndexJournalRecord]: + """Return watcher-retryable journal records for a workspace or specific repo.""" + workspace_path, repo_name = _resolve_repo_context(workspace_path, repo_name) + retry_delay = _index_journal_retry_delay_seconds() + max_attempts = _index_journal_max_attempts() + now = datetime.now() + + def _read_repo_journal_entries( + target_repo_name: Optional[str], + *, + target_workspace_path: Optional[str] = None, + ) -> List[IndexJournalRecord]: + journal = _read_index_journal_file_uncached( + _get_index_journal_path(target_workspace_path or workspace_path, target_repo_name) + ) + merged_ops = journal.get("operations", {}) + if not isinstance(merged_ops, dict): + merged_ops = {} + result: List[IndexJournalRecord] = [] + for rec in merged_ops.values(): + if not isinstance(rec, dict): + continue + status = _normalize_index_journal_status( + rec.get("status"), invalid="unknown" + ) + if status not in {"pending", "failed"}: + continue + attempts = _coerce_index_journal_attempts( + rec.get("attempts"), path=str(rec.get("path") or ""), warn=True + ) + if max_attempts > 0 and attempts >= max_attempts: + continue + if status == "failed" and retry_delay > 0: + updated_at = str(rec.get("updated_at") or "").strip() + if updated_at: + try: + last = datetime.fromisoformat(updated_at) + if (now - last).total_seconds() < retry_delay: + continue + except Exception: + pass + p = _normalize_cache_key_path(str(rec.get("path") or "")) + op_type = str(rec.get("op_type") or "").strip().lower() + if not p or op_type not in {"upsert", "delete"}: + continue + result.append( + { + "path": p, + "op_type": op_type, + "content_hash": str(rec.get("content_hash") or "").strip() or None, + "status": status, + "attempts": attempts, + "created_at": str(rec.get("created_at") or ""), + "updated_at": str(rec.get("updated_at") or ""), + "last_error": str(rec.get("last_error") or "").strip() or None, + } + ) + return result + + if repo_name: + return _read_repo_journal_entries(repo_name) + + result: List[IndexJournalRecord] = [] + for candidate, candidate_workspace_path in _discover_journal_repositories(workspace_path): + result.extend( + _read_repo_journal_entries( + candidate, + target_workspace_path=candidate_workspace_path, + ) + ) + + if result: + return result + return _read_repo_journal_entries(None) + + +def update_index_journal_entry_status( + path: str, + *, + status: str, + error: Optional[str] = None, + workspace_path: Optional[str] = None, + repo_name: Optional[str] = None, + remove_on_done: bool = True, +) -> Dict[str, Any]: + """Update or clear a repo-scoped journal entry after processing.""" + normalized_path = _normalize_cache_key_path(path) + status = _normalize_index_journal_status(status, invalid="failed") + now = datetime.now().isoformat() + + def _mutate(journal: Dict[str, Any]) -> None: + ops = journal.setdefault("operations", {}) + if not isinstance(ops, dict): + ops = {} + journal["operations"] = ops + rec = ops.get(normalized_path) + if not isinstance(rec, dict): + return + if status == "done" and remove_on_done: + ops.pop(normalized_path, None) + return + rec["status"] = status + rec["updated_at"] = now + rec["attempts"] = _coerce_index_journal_attempts( + rec.get("attempts"), path=normalized_path, warn=True + ) + 1 + rec["last_error"] = str(error or "").strip() or None + ops[normalized_path] = rec + + return _update_index_journal(workspace_path, repo_name, _mutate) + + +def get_cached_file_hash( + file_path: str, + repo_name: Optional[str] = None, + metadata_root: Optional[str] = None, +) -> str: """Get cached file hash for tracking changes.""" + root = metadata_root or _resolve_workspace_root() if is_multi_repo_mode() and repo_name: - state_dir = _get_repo_state_dir(repo_name) + state_dir = _get_repo_state_dir(repo_name, root) cache_path = state_dir / CACHE_FILENAME cache = _read_cache_file_cached(cache_path) @@ -1589,19 +2263,23 @@ def get_cached_file_hash(file_path: str, repo_name: Optional[str] = None) -> str return str(val.get("hash") or "") return str(val or "") else: - cache = _read_cache_cached(_resolve_workspace_root()) + cache = _read_cache_cached(root) fp = _normalize_cache_key_path(file_path) val = cache.get("file_hashes", {}).get(fp, "") if isinstance(val, dict): return str(val.get("hash") or "") return str(val or "") - return "" - -def set_cached_file_hash(file_path: str, file_hash: str, repo_name: Optional[str] = None) -> None: +def set_cached_file_hash( + file_path: str, + file_hash: str, + repo_name: Optional[str] = None, + metadata_root: Optional[str] = None, +) -> None: """Set cached file hash for tracking changes.""" fp = _normalize_cache_key_path(file_path) + root = metadata_root or _resolve_workspace_root() st_size: Optional[int] = None st_mtime: Optional[int] = None @@ -1615,14 +2293,15 @@ def set_cached_file_hash(file_path: str, file_hash: str, repo_name: Optional[str if is_multi_repo_mode() and repo_name: try: - ws_root = Path(_resolve_workspace_root()) + ws_root = Path(root) if not (ws_root / repo_name).exists(): return except Exception: return - state_dir = _get_repo_state_dir(repo_name) + state_dir = _get_repo_state_dir(repo_name, str(ws_root)) cache_path = state_dir / CACHE_FILENAME state_dir.mkdir(parents=True, exist_ok=True) + _apply_runtime_metadata_mode(state_dir) if cache_path.exists(): cache = _read_cache_file_cached(cache_path) @@ -1659,7 +2338,7 @@ def set_cached_file_hash(file_path: str, file_hash: str, repo_name: Optional[str _memoize_cache_obj(cache_path, cache) return - cache = _read_cache_cached(_resolve_workspace_root()) + cache = _read_cache_cached(root) existing = cache.get("file_hashes", {}).get(fp) if isinstance(existing, dict) and st_size is not None and st_mtime is not None: if ( @@ -1683,14 +2362,14 @@ def set_cached_file_hash(file_path: str, file_hash: str, repo_name: Optional[str pass cache.setdefault("file_hashes", {})[fp] = entry cache["updated_at"] = datetime.now().isoformat() - _write_cache(_resolve_workspace_root(), cache) - _memoize_cache_obj(_get_cache_path(_resolve_workspace_root()), cache) + _write_cache(root, cache) + _memoize_cache_obj(_get_cache_path(root), cache) def get_cached_file_meta(file_path: str, repo_name: Optional[str] = None) -> Dict[str, Any]: fp = _normalize_cache_key_path(file_path) if is_multi_repo_mode() and repo_name: - state_dir = _get_repo_state_dir(repo_name) + state_dir = _get_repo_state_dir(repo_name, _resolve_workspace_root()) cache_path = state_dir / CACHE_FILENAME cache = _read_cache_file_cached(cache_path) @@ -1711,10 +2390,15 @@ def get_cached_file_meta(file_path: str, repo_name: Optional[str] = None) -> Dic return {} -def remove_cached_file(file_path: str, repo_name: Optional[str] = None) -> None: +def remove_cached_file( + file_path: str, + repo_name: Optional[str] = None, + metadata_root: Optional[str] = None, +) -> None: """Remove file entry from cache.""" + root = metadata_root or _resolve_workspace_root() if is_multi_repo_mode() and repo_name: - state_dir = _get_repo_state_dir(repo_name) + state_dir = _get_repo_state_dir(repo_name, root) cache_path = state_dir / CACHE_FILENAME if cache_path.exists(): @@ -1730,13 +2414,13 @@ def remove_cached_file(file_path: str, repo_name: Optional[str] = None) -> None: _memoize_cache_obj(cache_path, cache) return - cache = _read_cache_cached(_resolve_workspace_root()) + cache = _read_cache_cached(root) fp = _normalize_cache_key_path(file_path) if fp in cache.get("file_hashes", {}): cache["file_hashes"].pop(fp, None) cache["updated_at"] = datetime.now().isoformat() - _write_cache(_resolve_workspace_root(), cache) - _memoize_cache_obj(_get_cache_path(_resolve_workspace_root()), cache) + _write_cache(root, cache) + _memoize_cache_obj(_get_cache_path(root), cache) def cleanup_old_cache_locks(max_idle_seconds: int = 900) -> int: @@ -1780,42 +2464,65 @@ def cleanup_old_cache_locks(max_idle_seconds: int = 900) -> int: def get_collection_mappings(search_root: Optional[str] = None) -> List[Dict[str, Any]]: - """Enumerate collection mappings with origin metadata.""" + """Enumerate collection mappings with origin metadata. + + `search_root` may point at either workspace root (`/work`) or codebase root + (`/work/.codebase`). + """ root_path = Path(search_root or _resolve_workspace_root()).resolve() + if root_path.name == STATE_DIRNAME: + workspace_root = root_path.parent + codebase_root = root_path + else: + workspace_root = root_path + codebase_root = root_path / STATE_DIRNAME mappings: List[Dict[str, Any]] = [] try: if is_multi_repo_mode(): - repos_root = root_path / STATE_DIRNAME / "repos" + seen_state_files: set[str] = set() + + def _append_repo_mapping(repo_name: str, state_path: Path) -> None: + if not state_path.exists(): + return + try: + state_key = str(state_path.resolve()) + except Exception: + state_key = str(state_path) + if state_key in seen_state_files: + return + seen_state_files.add(state_key) + + try: + with open(state_path, "r", encoding="utf-8-sig") as f: + state = json.load(f) or {} + except Exception as e: + print(f"[workspace_state] Failed to read repo state from {state_path}: {e}") + return + + origin = state.get("origin", {}) or {} + repo_workspace_dir = _get_repo_workspace_dir(repo_name, str(workspace_root)) + mappings.append( + { + "repo_name": repo_name, + "collection_name": state.get("qdrant_collection") + or get_collection_name(repo_name), + "container_path": origin.get("container_path") + or str(repo_workspace_dir.resolve()), + "source_path": origin.get("source_path"), + "state_file": str(state_path), + "updated_at": state.get("updated_at"), + } + ) + + # Shared metadata root (`/.codebase/repos//state.json`) + repos_root = codebase_root / "repos" if repos_root.exists(): for repo_dir in sorted(p for p in repos_root.iterdir() if p.is_dir()): - repo_name = repo_dir.name - state_path = repo_dir / STATE_FILENAME - if not state_path.exists(): - continue - try: - with open(state_path, "r", encoding="utf-8-sig") as f: - state = json.load(f) or {} - except Exception as e: - print(f"[workspace_state] Failed to read repo state from {state_path}: {e}") - continue - - origin = state.get("origin", {}) or {} - mappings.append( - { - "repo_name": repo_name, - "collection_name": state.get("qdrant_collection") - or get_collection_name(repo_name), - "container_path": origin.get("container_path") - or str((Path(_resolve_workspace_root()) / repo_name).resolve()), - "source_path": origin.get("source_path"), - "state_file": str(state_path), - "updated_at": state.get("updated_at"), - } - ) + _append_repo_mapping(repo_dir.name, repo_dir / STATE_FILENAME) else: - state_path = root_path / STATE_DIRNAME / STATE_FILENAME + state_path = codebase_root / STATE_FILENAME if state_path.exists(): try: with open(state_path, "r", encoding="utf-8-sig") as f: @@ -1824,14 +2531,14 @@ def get_collection_mappings(search_root: Optional[str] = None) -> List[Dict[str, state = {} origin = state.get("origin", {}) or {} - repo_name = origin.get("repo_name") or Path(root_path).name + repo_name = origin.get("repo_name") or Path(workspace_root).name mappings.append( { "repo_name": repo_name, "collection_name": state.get("qdrant_collection") or get_collection_name(repo_name), "container_path": origin.get("container_path") - or str(root_path), + or str(workspace_root), "source_path": origin.get("source_path"), "state_file": str(state_path), "updated_at": state.get("updated_at"), @@ -2116,6 +2823,8 @@ def set_cached_symbols(file_path: str, symbols: dict, file_hash: str) -> None: """Save symbol metadata for a file. Extends existing to include pseudo data.""" cache_path = _get_symbol_cache_path(file_path) cache_path.parent.mkdir(parents=True, exist_ok=True) + _apply_runtime_metadata_mode(cache_path.parent) + temp_path = cache_path.with_suffix(f".tmp.{uuid.uuid4().hex[:8]}") try: cache_data = { @@ -2125,18 +2834,16 @@ def set_cached_symbols(file_path: str, symbols: dict, file_hash: str) -> None: "symbols": symbols } - with open(cache_path, 'w', encoding='utf-8') as f: + with open(temp_path, 'w', encoding='utf-8') as f: json.dump(cache_data, f, indent=2) - - # Ensure symbol cache files are group-writable so both indexer and - # watcher processes (potentially different users sharing a group) - # can update them on shared volumes. - try: - os.chmod(cache_path, 0o664) - except PermissionError: - pass + temp_path.replace(cache_path) + _apply_runtime_metadata_mode(cache_path) except Exception as e: print(f"[SYMBOL_CACHE_WARNING] Failed to save symbol cache for {file_path}: {e}") + try: + temp_path.unlink(missing_ok=True) + except Exception: + pass def get_cached_pseudo(file_path: str, symbol_id: str) -> tuple[str, list[str]]: @@ -2236,7 +2943,7 @@ def clear_symbol_cache( target_dirs: List[Path] = [] if is_multi_repo_mode() and repo_name: - target_dirs.append(_get_repo_state_dir(repo_name) / "symbols") + target_dirs.append(_get_repo_state_dir(repo_name, workspace_path) / "symbols") else: try: cache_parent = _get_cache_path(workspace_root).parent @@ -2288,6 +2995,53 @@ def compare_symbol_changes(old_symbols: dict, new_symbols: dict) -> tuple[list, unchanged = [] changed = [] + # Primary key should not be absolute start_line alone; leading comments/import + # shifts can move every symbol without changing their bodies. Prefer exact id + # first, then fall back to stable metadata matching. + old_symbols = old_symbols or {} + new_symbols = new_symbols or {} + remaining_old_by_exact = dict(old_symbols) + remaining_old_by_signature: Dict[tuple[str, str, str], list[str]] = {} + remaining_old_by_name_kind: Dict[tuple[str, str], list[str]] = {} + + for old_symbol_id, old_info in remaining_old_by_exact.items(): + kind = str(old_info.get("type") or "") + name = str(old_info.get("name") or "") + content_hash = str(old_info.get("content_hash") or "") + if kind and name and content_hash: + remaining_old_by_signature.setdefault((kind, name, content_hash), []).append( + old_symbol_id + ) + if kind and name: + remaining_old_by_name_kind.setdefault((kind, name), []).append(old_symbol_id) + + def _consume_old_symbol(old_id: str, old_info: dict) -> None: + remaining_old_by_exact.pop(old_id, None) + + old_kind = str(old_info.get("type") or "") + old_name = str(old_info.get("name") or "") + old_hash = str(old_info.get("content_hash") or "") + + if old_kind and old_name and old_hash: + sig = (old_kind, old_name, old_hash) + sig_ids = remaining_old_by_signature.get(sig) or [] + if old_id in sig_ids: + sig_ids.remove(old_id) + if sig_ids: + remaining_old_by_signature[sig] = sig_ids + else: + remaining_old_by_signature.pop(sig, None) + + if old_kind and old_name: + nk = (old_kind, old_name) + nk_ids = remaining_old_by_name_kind.get(nk) or [] + if old_id in nk_ids: + nk_ids.remove(old_id) + if nk_ids: + remaining_old_by_name_kind[nk] = nk_ids + else: + remaining_old_by_name_kind.pop(nk, None) + for symbol_id, symbol_info in new_symbols.items(): if symbol_id in old_symbols: old_info = old_symbols[symbol_id] @@ -2296,6 +3050,26 @@ def compare_symbol_changes(old_symbols: dict, new_symbols: dict) -> tuple[list, unchanged.append(symbol_id) else: changed.append(symbol_id) + _consume_old_symbol(symbol_id, old_info) + continue + + kind = str(symbol_info.get("type") or "") + name = str(symbol_info.get("name") or "") + content_hash = str(symbol_info.get("content_hash") or "") + signature = (kind, name, content_hash) + matched_old_ids = remaining_old_by_signature.get(signature) or [] + if matched_old_ids: + old_id = matched_old_ids.pop(0) + if not matched_old_ids: + remaining_old_by_signature.pop(signature, None) + _consume_old_symbol(old_id, old_symbols.get(old_id, {})) + unchanged.append(symbol_id) + continue + + # Same logical symbol name/type exists but content differs: changed. + if kind and name and remaining_old_by_name_kind.get((kind, name)): + remaining_old_by_name_kind.pop((kind, name), None) + changed.append(symbol_id) else: # New symbol changed.append(symbol_id) @@ -2537,6 +3311,3 @@ def _list_workspaces_from_qdrant(seen_paths: set) -> List[Dict[str, Any]]: pass return workspaces - - -# Add missing functions that callers expect (already defined above) \ No newline at end of file diff --git a/skills/context-engine/SKILL.md b/skills/context-engine/SKILL.md index 50ee67b6..c139ca8c 100644 --- a/skills/context-engine/SKILL.md +++ b/skills/context-engine/SKILL.md @@ -1,11 +1,11 @@ --- name: context-engine -description: Codebase search and context retrieval for any programming language. Hybrid semantic/lexical search with neural reranking. Use for code lookup, finding implementations, understanding codebases, Q&A grounded in source code, and persistent memory across sessions. +description: Codebase search and context retrieval for any programming language. Use for code lookup, finding implementations, understanding codebases, Q&A grounded in source code, and persistent memory across sessions. --- # Context-Engine -Search and retrieve code context from any codebase using hybrid vector search (semantic + lexical) with neural reranking. +Search and retrieve code context from any codebase using the configured retrieval mode. `repo_search` is the canonical code search tool; dense, fusion, and reranking behavior depends on deployment settings. ## Decision Tree: Choosing the Right Tool @@ -14,8 +14,8 @@ What do you need? | +-- Find code locations/implementations | | - | +-- Simple query --> info_request - | +-- Need filters/control --> repo_search + | +-- Any query --> repo_search + | +-- Need file-type focus --> repo_search with profile | +-- Understand how something works | | @@ -29,13 +29,13 @@ What do you need? | +-- Find specific file types | | - | +-- Test files --> search_tests_for - | +-- Config files --> search_config_for + | +-- Test files --> repo_search with profile="tests" + | +-- Config files --> repo_search with profile="config" | +-- Find relationships | | - | +-- Who calls this function --> search_callers_for - | +-- Who imports this module --> search_importers_for + | +-- Who calls this function --> symbol_graph query_type="callers" + | +-- Who imports this module --> symbol_graph query_type="importers" | +-- Symbol graph navigation (callers/defs/importers) --> symbol_graph | +-- Git history --> search_commits_for @@ -47,7 +47,7 @@ What do you need? ## Primary Search: repo_search -Use `repo_search` (or its alias `code_search`) for most code lookups. Reranking is ON by default. +Use `repo_search` for code lookups. Retrieval mode and reranking are controlled by deployment configuration and per-call arguments. ```json { @@ -106,24 +106,10 @@ Use `repo: "*"` to search all indexed repos. - `ext` - File extension - `repo` - Repository filter for multi-repo setups - `case` - Case-sensitive matching - -## Simple Lookup: info_request - -Use `info_request` for natural language queries with minimal parameters: - -```json -{ - "info_request": "how does user authentication work" -} -``` - -Add explanations: -```json -{ - "info_request": "database connection pooling", - "include_explanation": true -} -``` +- `profile` - Focus common scopes: + - `"tests"` - Test files + - `"config"` - Configuration files + - `"code"` - Source-code extensions ## Q&A with Citations: context_answer @@ -206,28 +192,27 @@ Find structurally similar code patterns across all languages. Accepts **either** The `query_signature` encodes control flow: `L` (loops), `B` (branches), `T` (try/except), `M` (match). -## Specialized Search Tools +## Focused Search Profiles -**search_tests_for** - Find test files: -```json -{"query": "UserService", "limit": 10} -``` +Use `repo_search.profile` instead of separate focused tools. -**search_config_for** - Find config files: +**Test files**: ```json -{"query": "database connection", "limit": 5} +{"query": "UserService", "profile": "tests", "limit": 10} ``` -**search_callers_for** - Find callers of a symbol: +**Config files**: ```json -{"query": "processPayment", "language": "typescript"} +{"query": "database connection", "profile": "config", "limit": 5} ``` -**search_importers_for** - Find importers: +**Source-code files**: ```json -{"query": "utils/helpers", "limit": 10} +{"query": "imports qdrant client", "profile": "code", "limit": 10} ``` +For caller/importer relationships, prefer `symbol_graph` when you know the symbol. Use `repo_search` for exploratory prose queries. + **symbol_graph** - Symbol graph navigation (callers / definition / importers): ```json {"symbol": "ASTAnalyzer", "query_type": "definition", "limit": 10} @@ -344,10 +329,7 @@ With recreate (drops existing data): Set via `output_format` parameter. -## Aliases and Compat Wrappers - -**Aliases:** -- `code_search` = `repo_search` (identical behavior) +## Compat Wrappers **Cross-server tools:** - `memory_store` / `memory_find` — Memory server tools for persistent knowledge diff --git a/templates/admin/acl.html b/templates/admin/acl.html index 952a0ce9..656ed773 100644 --- a/templates/admin/acl.html +++ b/templates/admin/acl.html @@ -1,6 +1,40 @@ {% extends "admin/base.html" %} {% block content %} + {% set qp = request.query_params %} + {% if qp and ((qp.get("copied") and qp.get("new")) or qp.get("deleted") or qp.get("journal_cleared")) %} +

+ {% if qp.get("copied") and qp.get("new") %} +
+ Copied collection {{ qp.get("copied") }} → {{ qp.get("new") }}. + {% if qp.get("graph_copied") == "1" %} + (graph clone copied) + {% elif qp.get("graph_copied") == "0" %} + (graph clone not copied; will rebuild/backfill) + {% endif %} +
+ {% endif %} + {% if qp.get("deleted") %} +
+ Deleted collection {{ qp.get("deleted") }}. + {% if qp.get("graph_deleted") == "1" %} + (graph clone deleted) + {% elif qp.get("graph_deleted") == "0" %} + (graph clone not deleted or missing) + {% endif %} +
+ {% endif %} + {% if qp.get("journal_cleared") %} +
+ Cleared index journal for collection {{ qp.get("journal_cleared") }}. + {% if qp.get("journal_removed") %} + ({{ qp.get("journal_removed") }} entries removed) + {% endif %} +
+ {% endif %} +
+ {% endif %} +

Users

@@ -168,6 +202,10 @@

Collections

+
+ + +
{% endif %} {% if deletion_enabled %}
@@ -395,6 +433,12 @@

Grant Collection Access

+
+ + +
` : ""; const deleteHtml = deletionEnabled diff --git a/tests/conftest.py b/tests/conftest.py index b64ff9a9..d7a554ab 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -11,8 +11,41 @@ if str(ROOT) not in sys.path: sys.path.insert(0, str(ROOT)) -# Enable pattern vectors for pattern search tests -os.environ.setdefault("PATTERN_VECTORS", "1") +_INTEGRATION_TEST_FILES = { + "test_collection_memory_backup_restore.py", + "test_integration_qdrant.py", + "test_relevance_feedback.py", + "test_subprocess_hybrid_smoke.py", + "test_tier2_fallback.py", +} + + +def pytest_addoption(parser): + parser.addoption( + "--run-integration", + action="store_true", + default=False, + help="collect tests that start real services such as Qdrant", + ) + + +def pytest_configure(config): + if not config.getoption("--run-integration", default=False): + return + # pytest.ini excludes integration tests by default. When the explicit flag is + # supplied, clear only that default marker expression so opt-in means opt-in. + if getattr(config.option, "markexpr", "") == "not integration": + config.option.markexpr = "" + + +def pytest_ignore_collect(collection_path, config): + if config.getoption("--run-integration", default=False): + return False + try: + name = Path(str(collection_path)).name + except Exception: + name = "" + return name in _INTEGRATION_TEST_FILES @pytest.fixture(scope="session", autouse=True) @@ -59,11 +92,12 @@ def _wait_for_qdrant(url: str, timeout: int = 60) -> bool: @pytest.fixture(scope="module") def qdrant_url(): - """Provide Qdrant URL - uses CI service container or testcontainers (local). + """Provide Qdrant URL - uses CI, explicit QDRANT_URL, or testcontainers. In CI (GitHub Actions), uses the pre-configured Qdrant service at localhost:6333. - Locally, this fixture ALWAYS spins up a testcontainers Qdrant instance to avoid - accidentally polluting a developer's local Qdrant with test collections. + Locally, an explicit QDRANT_URL can point at the rebuilt compose stack for + faster opt-in integration runs. Otherwise this fixture spins up an isolated + testcontainers Qdrant instance. """ # Only use pre-configured Qdrant in CI environment (GitHub Actions sets CI=true) is_ci = os.environ.get("CI", "").lower() in ("true", "1", "yes") @@ -76,7 +110,15 @@ def qdrant_url(): # If not reachable in CI, fail explicitly pytest.fail(f"CI Qdrant service not reachable at {ci_url}") - # Local development: ALWAYS use testcontainers (safe isolation) + explicit_url = os.environ.get("QDRANT_URL", "").strip() + if explicit_url: + if _wait_for_qdrant(explicit_url, timeout=10): + yield explicit_url + return + pytest.fail(f"QDRANT_URL is set but Qdrant is not reachable at {explicit_url}") + + # Local development fallback: testcontainers keeps isolation when no explicit + # Qdrant service is requested. os.environ.setdefault("TESTCONTAINERS_RYUK_DISABLED", "true") os.environ.setdefault("TESTCONTAINERS_RYUK_TIMEOUT", "0") @@ -86,7 +128,7 @@ def qdrant_url(): pytest.skip("testcontainers not available and QDRANT_URL not set") container = ( - DockerContainer("qdrant/qdrant:latest") + DockerContainer("qdrant/qdrant:v1.15.4") .with_env("TESTCONTAINERS_RYUK_DISABLED", "true") .with_env("TESTCONTAINERS_RYUK_TIMEOUT", "0") .with_exposed_ports(6333) diff --git a/tests/test_admin_collection_delete.py b/tests/test_admin_collection_delete.py index ad42807e..1f0b9b4b 100644 --- a/tests/test_admin_collection_delete.py +++ b/tests/test_admin_collection_delete.py @@ -1,4 +1,6 @@ import importlib +import sys +import types import pytest from fastapi.testclient import TestClient @@ -22,11 +24,6 @@ def _fake_render_admin_error(_request, title, message, back_href="/admin", statu monkeypatch.setattr(srv, "render_admin_error", _fake_render_admin_error) - def _should_not_be_called(**_kwargs): - raise AssertionError("delete_collection_everywhere should not be called when env gate is off") - - monkeypatch.setattr(srv, "delete_collection_everywhere", _should_not_be_called) - client = TestClient(srv.app) resp = client.post("/admin/collections/delete", data={"collection": "c1", "delete_fs": ""}) assert resp.status_code == 403 @@ -52,6 +49,72 @@ def test_admin_role_gate_blocks_non_admin(monkeypatch): assert resp.json().get("detail") == "Admin required" +@pytest.mark.unit +def test_delete_redirect_includes_graph_deleted_param(monkeypatch): + monkeypatch.setenv("CTXCE_AUTH_ENABLED", "1") + monkeypatch.setenv("CTXCE_ADMIN_COLLECTION_DELETE_ENABLED", "1") + + srv = importlib.import_module("scripts.upload_service") + srv = importlib.reload(srv) + + monkeypatch.setattr(srv, "_require_admin_session", lambda _req: {"user_id": "admin"}) + + def _fake_delete_collection_everywhere(**_kwargs): + return {"qdrant_deleted": True, "qdrant_graph_deleted": True} + + monkeypatch.setitem( + sys.modules, + "scripts.collection_admin", + types.SimpleNamespace(delete_collection_everywhere=_fake_delete_collection_everywhere), + ) + + client = TestClient(srv.app) + resp = client.post("/admin/collections/delete", data={"collection": "c1", "delete_fs": ""}, follow_redirects=False) + assert resp.status_code == 302 + loc = resp.headers.get("location") or "" + assert "deleted=c1" in loc + assert "graph_deleted=1" in loc + + +@pytest.mark.unit +def test_clear_journal_endpoint_clears_mapped_collection(monkeypatch): + monkeypatch.setenv("CTXCE_AUTH_ENABLED", "1") + + srv = importlib.import_module("scripts.upload_service") + srv = importlib.reload(srv) + + calls = {} + monkeypatch.setattr(srv, "_require_admin_session", lambda _req: {"user_id": "admin"}) + monkeypatch.setitem( + sys.modules, + "scripts.indexing_admin", + types.SimpleNamespace( + resolve_collection_root=lambda **_kwargs: ( + "/work/repo-1234567890abcdef", + "repo-1234567890abcdef", + ) + ), + ) + + def _clear_index_journal_entries(**kwargs): + calls.update(kwargs) + return 3 + + monkeypatch.setattr(srv, "clear_index_journal_entries", _clear_index_journal_entries) + + client = TestClient(srv.app) + resp = client.post("/admin/collections/clear-journal", data={"collection": "c1"}, follow_redirects=False) + + assert resp.status_code == 302 + loc = resp.headers.get("location") or "" + assert "journal_cleared=c1" in loc + assert "journal_removed=3" in loc + assert calls == { + "workspace_path": "/work/repo-1234567890abcdef", + "repo_name": "repo-1234567890abcdef", + } + + @pytest.mark.unit def test_collection_admin_refuses_when_env_disabled(monkeypatch): monkeypatch.setenv("CTXCE_ADMIN_COLLECTION_DELETE_ENABLED", "0") diff --git a/tests/test_admin_ui.py b/tests/test_admin_ui.py new file mode 100644 index 00000000..8e068d8c --- /dev/null +++ b/tests/test_admin_ui.py @@ -0,0 +1,36 @@ +from starlette.requests import Request + +from scripts.admin_ui import ( + render_admin_acl, + render_admin_bootstrap, + render_admin_error, + render_admin_login, +) + + +def _request(path: str = "/admin/login") -> Request: + return Request( + { + "type": "http", + "method": "GET", + "path": path, + "headers": [], + "query_string": b"", + "server": ("testserver", 80), + "scheme": "http", + "client": ("127.0.0.1", 12345), + } + ) + + +def test_admin_templates_render_with_request_first_api(): + request = _request() + + responses = [ + render_admin_login(request), + render_admin_bootstrap(request), + render_admin_acl(request, users=[], collections=[], grants={}), + render_admin_error(request, title="Error", message="Something failed"), + ] + + assert [response.status_code for response in responses] == [200, 200, 200, 400] diff --git a/tests/test_cache_deduplication.py b/tests/test_cache_deduplication.py index 8180e755..5346b537 100644 --- a/tests/test_cache_deduplication.py +++ b/tests/test_cache_deduplication.py @@ -87,13 +87,13 @@ def test_cache_ttl_expiration(self): """Test TTL-based expiration.""" cache = UnifiedCache("test_ttl", max_size=10, eviction_policy=EvictionPolicy.TTL, default_ttl=0.1) - # Set value with short TTL - cache.set("key1", "value1", ttl=0.1) - self.assertEqual(cache.get("key1"), "value1") + with patch("scripts.cache_manager.time.time") as fake_time: + fake_time.return_value = 1000.0 + cache.set("key1", "value1", ttl=0.1) + self.assertEqual(cache.get("key1"), "value1") - # Wait for expiration - time.sleep(0.2) - self.assertIsNone(cache.get("key1")) # Should be expired + fake_time.return_value = 1000.2 + self.assertIsNone(cache.get("key1")) # Should be expired def test_cache_statistics(self): """Test cache statistics tracking.""" @@ -149,27 +149,29 @@ def test_cached_decorator(self): """Test the cached decorator.""" call_count = 0 - @cached("test_decorator", ttl=1.0) - def expensive_function(x): - nonlocal call_count - call_count += 1 - return x * 2 + with patch("scripts.cache_manager.time.time") as fake_time: + fake_time.return_value = 1000.0 + + @cached("test_decorator", ttl=1.0) + def expensive_function(x): + nonlocal call_count + call_count += 1 + return x * 2 - # First call should compute - result1 = expensive_function(5) - self.assertEqual(result1, 10) - self.assertEqual(call_count, 1) + # First call should compute + result1 = expensive_function(5) + self.assertEqual(result1, 10) + self.assertEqual(call_count, 1) - # Second call should use cache - result2 = expensive_function(5) - self.assertEqual(result2, 10) - self.assertEqual(call_count, 1) # Should not increase + # Second call should use cache + result2 = expensive_function(5) + self.assertEqual(result2, 10) + self.assertEqual(call_count, 1) # Should not increase - # Wait for expiration and call again - time.sleep(1.1) - result3 = expensive_function(5) - self.assertEqual(result3, 10) - self.assertEqual(call_count, 2) # Should recompute + fake_time.return_value = 1001.1 + result3 = expensive_function(5) + self.assertEqual(result3, 10) + self.assertEqual(call_count, 2) # Should recompute class TestRequestDeduplication(unittest.TestCase): @@ -304,42 +306,46 @@ def test_deduplication_ttl(self): request = {'queries': ['test'], 'limit': 10} - # First request should be unique - is_dup1, fp1 = deduplicator.is_duplicate(request) - self.assertFalse(is_dup1) + with patch("scripts.deduplication.time.time") as fake_time: + fake_time.return_value = 1000.0 - # Wait for expiration - time.sleep(0.2) + # First request should be unique + is_dup1, fp1 = deduplicator.is_duplicate(request) + self.assertFalse(is_dup1) - # Same request should be unique again after expiration - is_dup2, fp2 = deduplicator.is_duplicate(request) - self.assertFalse(is_dup2) + fake_time.return_value = 1000.2 + + # Same request should be unique again after expiration + is_dup2, fp2 = deduplicator.is_duplicate(request) + self.assertFalse(is_dup2) def test_deduplicate_request_decorator(self): """Test the deduplicate_request decorator.""" call_count = 0 - @deduplicate_request(ttl=1.0) - def expensive_search(query): - nonlocal call_count - call_count += 1 - return f"search_result_for_{query}" - - # First call should execute - result1 = expensive_search("test") - self.assertEqual(result1, "search_result_for_test") - self.assertEqual(call_count, 1) - - # Second identical call should be deduplicated - result2 = expensive_search("test") - self.assertIsNone(result2) # Decorator returns None for duplicates - self.assertEqual(call_count, 1) # Should not increase - - # Wait for expiration and call again - time.sleep(1.1) - result3 = expensive_search("test") - self.assertEqual(result3, "search_result_for_test") - self.assertEqual(call_count, 2) + with patch("scripts.deduplication.time.time") as fake_time: + fake_time.return_value = 1000.0 + + @deduplicate_request(ttl=1.0) + def expensive_search(query): + nonlocal call_count + call_count += 1 + return f"search_result_for_{query}" + + # First call should execute + result1 = expensive_search("test") + self.assertEqual(result1, "search_result_for_test") + self.assertEqual(call_count, 1) + + # Second identical call should be deduplicated + result2 = expensive_search("test") + self.assertIsNone(result2) # Decorator returns None for duplicates + self.assertEqual(call_count, 1) # Should not increase + + fake_time.return_value = 1001.1 + result3 = expensive_search("test") + self.assertEqual(result3, "search_result_for_test") + self.assertEqual(call_count, 2) class TestCacheIntegration(unittest.TestCase): diff --git a/tests/test_change_history_for_path.py b/tests/test_change_history_for_path.py index 52be7592..7d822150 100644 --- a/tests/test_change_history_for_path.py +++ b/tests/test_change_history_for_path.py @@ -15,6 +15,10 @@ def tool(self, *args, **kwargs): def _decorator(fn): return fn return _decorator + def resource(self, *args, **kwargs): + def _decorator(fn): + return fn + return _decorator class _Context: def __init__(self, *args, **kwargs): @@ -98,4 +102,3 @@ async def test_change_history_strict_match_under_work(monkeypatch): assert summary.get("ingested_min") == 90 assert summary.get("ingested_max") == 115 assert summary.get("churn_count_max") == 5 - diff --git a/tests/test_collection_memory_backup_restore.py b/tests/test_collection_memory_backup_restore.py index c9e74921..cfdaf63e 100644 --- a/tests/test_collection_memory_backup_restore.py +++ b/tests/test_collection_memory_backup_restore.py @@ -78,14 +78,14 @@ def test_memory_backup_restore_happy_path(qdrant_container, monkeypatch): - The collection should be updated (if possible) without recreation. - Existing points should remain intact. """ - os.environ["QDRANT_URL"] = qdrant_container + monkeypatch.setenv("QDRANT_URL", qdrant_container) collection = f"test-mem-{uuid.uuid4().hex[:8]}" client = _create_collection_with_memory(qdrant_container, collection, dim=8) # Force ReFRAG on so ensure_collection tries to add MINI_VECTOR_NAME - os.environ["REFRAG_MODE"] = "1" - os.environ.pop("STRICT_MEMORY_RESTORE", None) + monkeypatch.setenv("REFRAG_MODE", "1") + monkeypatch.delenv("STRICT_MEMORY_RESTORE", raising=False) # Run ensure_collection: this should trigger backup + recreate + restore ing.ensure_collection(client, collection, dim=8, vector_name="code") @@ -107,13 +107,13 @@ def test_memory_backup_restore_happy_path(qdrant_container, monkeypatch): def test_memory_restore_strict_mode_no_recreate(qdrant_container, monkeypatch): """STRICT_MEMORY_RESTORE should not trigger errors when no recreate occurs.""" - os.environ["QDRANT_URL"] = qdrant_container + monkeypatch.setenv("QDRANT_URL", qdrant_container) collection = f"test-mem-strict-{uuid.uuid4().hex[:8]}" client = _create_collection_with_memory(qdrant_container, collection, dim=8) - os.environ["REFRAG_MODE"] = "1" - os.environ["STRICT_MEMORY_RESTORE"] = "1" + monkeypatch.setenv("REFRAG_MODE", "1") + monkeypatch.setenv("STRICT_MEMORY_RESTORE", "1") # Patch subprocess.run to: # - allow the real memory_backup.py to run @@ -141,13 +141,13 @@ def test_memory_backup_failure_tolerant_mode_no_recreate(qdrant_container, monke """If backup fails but STRICT_MEMORY_RESTORE is not set, ensure_collection should still proceed without destructive recreation. """ - os.environ["QDRANT_URL"] = qdrant_container + monkeypatch.setenv("QDRANT_URL", qdrant_container) collection = f"test-mem-backup-fail-{uuid.uuid4().hex[:8]}" client = _create_collection_with_memory(qdrant_container, collection, dim=8) - os.environ["REFRAG_MODE"] = "1" - os.environ.pop("STRICT_MEMORY_RESTORE", None) + monkeypatch.setenv("REFRAG_MODE", "1") + monkeypatch.delenv("STRICT_MEMORY_RESTORE", raising=False) # Patch subprocess.run so memory_backup.py fails, but everything else runs normally orig_run = subprocess.run @@ -176,14 +176,14 @@ def fake_run(args, **kwargs): # type: ignore[override] assert "2" in ids -def test_memory_backup_and_restore_scripts_roundtrip(qdrant_container, tmp_path): +def test_memory_backup_and_restore_scripts_roundtrip(qdrant_container, tmp_path, monkeypatch): """Directly exercise memory_backup.export_memories and memory_restore.restore_memories without going through ensure_collection. This confirms that the backup file contains the expected memory and that restore_memories can recreate it in a fresh collection. """ - os.environ["QDRANT_URL"] = qdrant_container + monkeypatch.setenv("QDRANT_URL", qdrant_container) collection = f"test-mem-scripts-{uuid.uuid4().hex[:8]}" client = _create_collection_with_memory(qdrant_container, collection, dim=8) diff --git a/tests/test_concurrency_service.py b/tests/test_concurrency_service.py index a923163a..1e5e61de 100644 --- a/tests/test_concurrency_service.py +++ b/tests/test_concurrency_service.py @@ -10,6 +10,8 @@ async def test_repo_search_concurrent(monkeypatch): # In-process, fast stubbed hybrid search and model monkeypatch.setenv("HYBRID_IN_PROCESS", "1") + monkeypatch.setenv("REPO_SEARCH_DEFAULT_MODE", "hybrid") + monkeypatch.setenv("RERANKER_ENABLED", "0") monkeypatch.setattr(srv, "_get_embedding_model", lambda *a, **k: object()) import scripts.hybrid_search as hy diff --git a/tests/test_context_answer.py b/tests/test_context_answer.py index 595207b9..c7bf2e2f 100644 --- a/tests/test_context_answer.py +++ b/tests/test_context_answer.py @@ -1,9 +1,58 @@ -import importlib +import asyncio +import sys +import threading import types import pytest +from scripts.mcp_impl.context_answer import ( + _ca_prepare_filters_and_retrieve, + _context_answer_impl, +) + + +def _retrieval_result(items, **overrides): + result = { + "items": items, + "eff_language": None, + "eff_path_glob": None, + "eff_not_glob": None, + "override_under": False, + "sym_arg": None, + "cwd_root": "/work", + "path_regex": None, + "ext": None, + "kind": None, + "case": None, + } + result.update(overrides) + return result + + +def _install_fake_hybrid(monkeypatch, run_hybrid_search): + fake = types.ModuleType("scripts.hybrid_search") + fake.run_hybrid_search = run_hybrid_search + fake.lang_matches_path = lambda language, path: True + fake._merge_and_budget_spans = lambda items: list(items or []) + monkeypatch.setitem(sys.modules, "scripts.hybrid_search", fake) + return fake + + +def _run_context_answer(retrieval_fn=None, **kwargs): + return asyncio.get_event_loop().run_until_complete( + _context_answer_impl( + **kwargs, + get_embedding_model_fn=lambda *a, **k: None, + env_lock=threading.Lock(), + prepare_filters_and_retrieve_fn=retrieval_fn or _ca_prepare_filters_and_retrieve, + ) + ) + -srv = importlib.import_module("scripts.mcp_indexer_server") +def _isolate_context_answer_unit(monkeypatch): + monkeypatch.setenv("REFRAG_RUNTIME", "llamacpp") + monkeypatch.setenv("CTX_MULTI_COLLECTION", "0") + monkeypatch.setenv("CTX_DOC_PASS", "0") + monkeypatch.setenv("CTX_DOC_TOP_FALLBACK", "0") def _fake_items(): @@ -29,14 +78,7 @@ def _fake_items(): @pytest.mark.service def test_context_answer_happy_path(monkeypatch): - # Mock embedding model to avoid loading real model - monkeypatch.setattr(srv, "_get_embedding_model", lambda *a, **k: None) - - # Fake retrieval output (already budgeted) - import scripts.hybrid_search as hs - - monkeypatch.setattr(hs, "run_hybrid_search", lambda **k: _fake_items()) - + _isolate_context_answer_unit(monkeypatch) # Fake decoder import scripts.refrag_llamacpp as ref @@ -51,8 +93,11 @@ def generate_with_soft_embeddings(self, prompt: str, max_tokens: int = 256, **kw monkeypatch.setattr(ref, "LlamaCppRefragClient", FakeLlama) monkeypatch.setattr(ref, "is_decoder_enabled", lambda: True) - out = srv.asyncio.get_event_loop().run_until_complete( - srv.context_answer(query="how to do x", limit=2, per_path=1) + out = _run_context_answer( + retrieval_fn=lambda **_kwargs: _retrieval_result(_fake_items()), + query="how to do x", + limit=2, + per_path=1, ) assert isinstance(out, dict) @@ -63,12 +108,12 @@ def generate_with_soft_embeddings(self, prompt: str, max_tokens: int = 256, **kw def test_context_answer_decoder_disabled(monkeypatch): - # Mock embedding model to avoid loading real model - monkeypatch.setattr(srv, "_get_embedding_model", lambda *a, **k: None) - - import scripts.hybrid_search as hs - - monkeypatch.setattr(hs, "run_hybrid_search", lambda **k: _fake_items()) + _isolate_context_answer_unit(monkeypatch) + monkeypatch.setenv("REFRAG_MODE", "0") + monkeypatch.setenv("REFRAG_GATE_FIRST", "0") + monkeypatch.setenv("REFRAG_RUNTIME", "llamacpp") + monkeypatch.setenv("CTX_CLIENT_DEADLINE_SEC", "178") + monkeypatch.setenv("CTX_DEADLINE_MARGIN_SEC", "6") import scripts.refrag_llamacpp as ref @@ -82,8 +127,10 @@ def generate_with_soft_embeddings(self, *a, **k): monkeypatch.setattr(ref, "LlamaCppRefragClient", FakeLlama) monkeypatch.setattr(ref, "is_decoder_enabled", lambda: False) - out = srv.asyncio.get_event_loop().run_until_complete( - srv.context_answer(query="how to do y", limit=1) + out = _run_context_answer( + retrieval_fn=lambda **_kwargs: _retrieval_result(_fake_items()), + query="how to do y", + limit=1, ) assert "error" in out @@ -91,11 +138,7 @@ def generate_with_soft_embeddings(self, *a, **k): def test_context_answer_prefers_identifier_spans(monkeypatch): - # Mock embedding model to avoid loading real model - monkeypatch.setattr(srv, "_get_embedding_model", lambda *a, **k: None) - - import scripts.hybrid_search as hs - + _isolate_context_answer_unit(monkeypatch) def _items(): return [ { @@ -105,6 +148,7 @@ def _items(): "start_line": 10, "end_line": 16, "text": "def helper():\n return 42\n", + "span_budgeted": True, }, { "score": 0.8, @@ -113,11 +157,10 @@ def _items(): "start_line": 5, "end_line": 9, "text": "RRF_K = 60\n", + "span_budgeted": True, }, ] - monkeypatch.setattr(hs, "run_hybrid_search", lambda **k: _items()) - import scripts.refrag_llamacpp as ref class FakeLlama: @@ -130,8 +173,11 @@ def generate_with_soft_embeddings(self, prompt: str, max_tokens: int = 256, **kw monkeypatch.setattr(ref, "LlamaCppRefragClient", FakeLlama) monkeypatch.setattr(ref, "is_decoder_enabled", lambda: True) - out = srv.asyncio.get_event_loop().run_until_complete( - srv.context_answer(query="what is RRF_K in hybrid_search.py?", limit=1, per_path=1) + out = _run_context_answer( + retrieval_fn=lambda **_kwargs: _retrieval_result(_items()), + query="what is RRF_K in hybrid_search.py?", + limit=1, + per_path=1, ) cits = out.get("citations") or [] @@ -141,11 +187,7 @@ def generate_with_soft_embeddings(self, prompt: str, max_tokens: int = 256, **kw def test_context_answer_tier2_retry_without_gating(monkeypatch): """Tier 2 should retry run_hybrid_search with relaxed filters when Tier 1 yields zero.""" - # Mock embedding model to avoid loading real model - monkeypatch.setattr(srv, "_get_embedding_model", lambda *a, **k: None) - - import scripts.hybrid_search as hs - + _isolate_context_answer_unit(monkeypatch) calls = [] def _run_hybrid_search(**kwargs): @@ -165,7 +207,7 @@ def _run_hybrid_search(**kwargs): # All other calls (tier1/usage/targeted search) yield no hits return [] - monkeypatch.setattr(hs, "run_hybrid_search", _run_hybrid_search) + _install_fake_hybrid(monkeypatch, _run_hybrid_search) import scripts.refrag_llamacpp as ref @@ -179,9 +221,7 @@ def generate_with_soft_embeddings(self, *a, **kw): monkeypatch.setattr(ref, "LlamaCppRefragClient", FakeLlama) monkeypatch.setattr(ref, "is_decoder_enabled", lambda: True) - out = srv.asyncio.get_event_loop().run_until_complete( - srv.context_answer(query="RRF_K", limit=1, per_path=1) - ) + out = _run_context_answer(query="RRF_K", limit=1, per_path=1) # Ensure Tier 2 was invoked (run_hybrid_search called twice) assert len(calls) >= 3, "Tier 2 fallback should re-run hybrid search" @@ -200,9 +240,7 @@ def generate_with_soft_embeddings(self, *a, **kw): def test_context_answer_env_lock_release_on_retrieval_exception(monkeypatch): - # Mock embedding model to avoid loading real model - monkeypatch.setattr(srv, "_get_embedding_model", lambda *a, **k: None) - + _isolate_context_answer_unit(monkeypatch) import os # Force retrieval to raise and ensure env/lock are restored prev = {k: os.environ.get(k) for k in ( @@ -212,16 +250,22 @@ def test_context_answer_env_lock_release_on_retrieval_exception(monkeypatch): def _raise_retrieval(*a, **k): raise RuntimeError("boom") - monkeypatch.setattr(srv, "_ca_prepare_filters_and_retrieve", _raise_retrieval) - - out = srv.asyncio.get_event_loop().run_until_complete( - srv.context_answer(query="x", limit=1, per_path=1) + lock = threading.Lock() + out = asyncio.get_event_loop().run_until_complete( + _context_answer_impl( + query="x", + limit=1, + per_path=1, + get_embedding_model_fn=lambda *a, **k: None, + env_lock=lock, + prepare_filters_and_retrieve_fn=_raise_retrieval, + ) ) assert "error" in out # Lock should be free after failure - assert srv._ENV_LOCK.acquire(blocking=False), "_ENV_LOCK should be released on exception" - srv._ENV_LOCK.release() + assert lock.acquire(blocking=False), "context_answer env lock should be released on exception" + lock.release() # Env should be restored for k, v in prev.items(): @@ -243,12 +287,19 @@ def _fake_retrieval(*a, **k): "case": None, } - monkeypatch.setattr(srv, "_ca_prepare_filters_and_retrieve", _fake_retrieval) - import scripts.refrag_llamacpp as ref + + _install_fake_hybrid(monkeypatch, lambda **k: []) monkeypatch.setattr(ref, "is_decoder_enabled", lambda: False) - out2 = srv.asyncio.get_event_loop().run_until_complete( - srv.context_answer(query="x", limit=1, per_path=1) + out2 = asyncio.get_event_loop().run_until_complete( + _context_answer_impl( + query="x", + limit=1, + per_path=1, + get_embedding_model_fn=lambda *a, **k: None, + env_lock=lock, + prepare_filters_and_retrieve_fn=_fake_retrieval, + ) ) assert isinstance(out2, dict) diff --git a/tests/test_context_answer_fallback.py b/tests/test_context_answer_fallback.py index d3c0f599..af95291d 100644 --- a/tests/test_context_answer_fallback.py +++ b/tests/test_context_answer_fallback.py @@ -1,28 +1,53 @@ import asyncio +import sys +import types import pytest -from unittest.mock import patch + +from scripts.mcp_impl.context_answer import _context_answer_impl + @pytest.mark.asyncio -async def test_context_answer_has_no_filesystem_fallback_when_no_hits(): +async def test_context_answer_has_no_filesystem_fallback_when_no_hits(monkeypatch): """When retrieval yields no spans, we do NOT glob or read the host filesystem. Citations may be empty, and that's expected. """ - from scripts.mcp_indexer_server import context_answer - import scripts.mcp_indexer_server as srv + fake_hybrid = types.ModuleType("scripts.hybrid_search") + fake_hybrid.run_hybrid_search = lambda **k: [] + fake_hybrid.lang_matches_path = lambda language, path: True + fake_hybrid._merge_and_budget_spans = lambda items: list(items or []) + + monkeypatch.setenv("REFRAG_RUNTIME", "llamacpp") + monkeypatch.setenv("CTX_MULTI_COLLECTION", "0") + monkeypatch.setenv("CTX_DOC_PASS", "0") + monkeypatch.setenv("CTX_DOC_TOP_FALLBACK", "0") + monkeypatch.setitem(sys.modules, "scripts.hybrid_search", fake_hybrid) - # Mock embedding model to avoid loading real model - with patch.object(srv, "_get_embedding_model", return_value=None): - out = await context_answer( - query="Describe module roles", - limit=3, - per_path=1, - include_snippet=True, - path_glob=["scripts/hybrid_search.py"], - # Force a very unlikely match to simulate empty retrieval - language="nonexistentlang", - ) - assert isinstance(out, dict) - # No fallback: citations can be empty - cits = out.get("citations") or [] - assert len(cits) == 0 + def _empty_retrieval(**_kwargs): + return { + "items": [], + "eff_language": "nonexistentlang", + "eff_path_glob": ["scripts/hybrid_search.py"], + "eff_not_glob": [], + "override_under": None, + "sym_arg": None, + "cwd_root": "/work", + "path_regex": None, + "ext": None, + "kind": None, + "case": None, + } + out = await _context_answer_impl( + query="Describe module roles", + limit=3, + per_path=1, + include_snippet=True, + path_glob=["scripts/hybrid_search.py"], + language="nonexistentlang", + get_embedding_model_fn=lambda *_args, **_kwargs: None, + prepare_filters_and_retrieve_fn=_empty_retrieval, + ) + assert isinstance(out, dict) + # No fallback: citations can be empty + cits = out.get("citations") or [] + assert len(cits) == 0 diff --git a/tests/test_context_answer_path_mention.py b/tests/test_context_answer_path_mention.py index 59299b8d..edc7b281 100644 --- a/tests/test_context_answer_path_mention.py +++ b/tests/test_context_answer_path_mention.py @@ -1,17 +1,47 @@ -import importlib +import asyncio +import sys +import threading +import types import pytest -srv = importlib.import_module("scripts.mcp_indexer_server") +from scripts.mcp_impl.context_answer import ( + _ca_prepare_filters_and_retrieve, + _context_answer_impl, +) + + +def _run_context_answer(**kwargs): + return asyncio.get_event_loop().run_until_complete( + _context_answer_impl( + **kwargs, + get_embedding_model_fn=lambda *a, **k: None, + env_lock=threading.Lock(), + prepare_filters_and_retrieve_fn=_ca_prepare_filters_and_retrieve, + ) + ) + + +def _install_fake_hybrid(monkeypatch, run_hybrid_search): + fake = types.ModuleType("scripts.hybrid_search") + fake.run_hybrid_search = run_hybrid_search + fake.lang_matches_path = lambda language, path: True + fake._merge_and_budget_spans = lambda items: list(items or []) + monkeypatch.setitem(sys.modules, "scripts.hybrid_search", fake) + return fake + + +def _isolate_context_answer_unit(monkeypatch): + monkeypatch.setenv("REFRAG_RUNTIME", "llamacpp") + monkeypatch.setenv("CTX_MULTI_COLLECTION", "0") + monkeypatch.setenv("CTX_DOC_PASS", "0") + monkeypatch.setenv("CTX_DOC_TOP_FALLBACK", "0") @pytest.mark.service def test_context_answer_path_mention_fallback(monkeypatch): - # Mock embedding model to avoid loading real model - monkeypatch.setattr(srv, "_get_embedding_model", lambda *a, **k: None) - + _isolate_context_answer_unit(monkeypatch) # Force retrieval to return nothing so path-mention fallback engages - import scripts.hybrid_search as hs - monkeypatch.setattr(hs, "run_hybrid_search", lambda **k: []) + _install_fake_hybrid(monkeypatch, lambda **k: []) import scripts.refrag_llamacpp as ref @@ -30,9 +60,7 @@ def generate_with_soft_embeddings(self, prompt: str, max_tokens: int = 64, **kw) # Mention an actual file in this repo so fallback can find it q = "explain something in scripts/hybrid_search.py" - out = srv.asyncio.get_event_loop().run_until_complete( - srv.context_answer(query=q, limit=3, per_path=2) - ) + out = _run_context_answer(query=q, limit=3, per_path=2) assert isinstance(out, dict) cits = out.get("citations") or [] assert len(cits) >= 1 @@ -40,4 +68,3 @@ def generate_with_soft_embeddings(self, prompt: str, max_tokens: int = 64, **kw) p = cits[0].get("path") or "" rp = cits[0].get("rel_path") or "" assert p.endswith("scripts/hybrid_search.py") or rp.endswith("scripts/hybrid_search.py") - diff --git a/tests/test_ctx_cli.py b/tests/test_ctx_cli.py new file mode 100644 index 00000000..6ab6e8e9 --- /dev/null +++ b/tests/test_ctx_cli.py @@ -0,0 +1,67 @@ +import sys + +import scripts.ctx as ctx + + +def test_parse_mcp_response_prefers_structured_content(): + payload = { + "result": { + "content": [{"type": "text", "text": '{"result":{"results":[]}}'}], + "structuredContent": { + "result": {"results": [{"path": "structured.py"}], "total": 1} + }, + } + } + + assert ctx.parse_mcp_response(payload) == { + "results": [{"path": "structured.py"}], + "total": 1, + } + + +def test_parse_mcp_response_unwraps_text_result_payload(): + payload = { + "result": { + "content": [ + { + "type": "text", + "text": '{"result":{"results":[{"path":"text.py"}],"total":1}}', + } + ] + } + } + + assert ctx.parse_mcp_response(payload) == { + "results": [{"path": "text.py"}], + "total": 1, + } + + +def test_main_with_context_appends_supporting_context(monkeypatch, capsys): + monkeypatch.setattr( + sys, + "argv", + ["ctx.py", "--with-context", "where is dense search?"], + ) + monkeypatch.setattr( + ctx, + "fetch_context", + lambda *a, **k: ( + "- /work/scripts/hybrid_search.py:428-565 (run_pure_dense_search)", + "", + ), + ) + monkeypatch.setattr(ctx, "rewrite_prompt", lambda *a, **k: "rewritten prompt") + monkeypatch.setattr( + ctx, + "extract_allowed_citations", + lambda *a, **k: ({"/work/scripts/hybrid_search.py"}, {}), + ) + monkeypatch.setattr(ctx, "sanitize_citations", lambda text, *_: text) + + ctx.main() + + out = capsys.readouterr().out + assert "rewritten prompt" in out + assert "Supporting context:" in out + assert "run_pure_dense_search" in out diff --git a/tests/test_env_behavior.py b/tests/test_env_behavior.py index e87cc8db..220a02bf 100644 --- a/tests/test_env_behavior.py +++ b/tests/test_env_behavior.py @@ -1,9 +1,7 @@ import importlib -import types -import os import pytest -srv = importlib.import_module("scripts.mcp_indexer_server") +search_impl = importlib.import_module("scripts.mcp_impl.search") @pytest.mark.service @@ -11,19 +9,21 @@ def test_rerank_timeout_floor_and_env_defaults(monkeypatch): # Force rerank via env default when arg not provided monkeypatch.setenv("RERANKER_ENABLED", "1") monkeypatch.setenv("RERANK_IN_PROCESS", "0") + monkeypatch.setenv("HYBRID_IN_PROCESS", "0") + monkeypatch.setenv("REPO_SEARCH_DEFAULT_MODE", "hybrid") # Floor 1500ms; client asks 200ms -> effective >= 1500ms -> 1.5s monkeypatch.setenv("RERANK_TIMEOUT_FLOOR_MS", "1500") # Fix default timeout for test determinism (CI may set a higher value) monkeypatch.setenv("RERANKER_TIMEOUT_MS", "200") - # Fake _run_async to capture calls + # Fake subprocess runner to capture hybrid + rerank calls without loading the MCP facade. calls = [] async def fake_run(cmd, env=None, timeout=None): calls.append({"cmd": cmd, "timeout": timeout}) - # Distinguish hybrid vs rerank by script name - if any("rerank_local.py" in str(x) for x in cmd): + # Distinguish hybrid vs rerank by module name + if "scripts.rerank_tools.local" in " ".join(map(str, cmd)): # Return something that looks like rerank stdout return { "ok": True, @@ -40,17 +40,21 @@ async def fake_run(cmd, env=None, timeout=None): "code": 0, } - monkeypatch.setattr(srv, "_run_async", fake_run) - # Call repo_search with no rerank_enabled arg to pick env default - res = srv.asyncio.get_event_loop().run_until_complete( - srv.repo_search(query="foo", limit=3, per_path=1) + res = search_impl.asyncio.get_event_loop().run_until_complete( + search_impl._repo_search_impl( + query="foo", + limit=3, + per_path=1, + run_async_fn=fake_run, + require_auth_session_fn=lambda session: session, + ) ) assert any( - "rerank_local.py" in " ".join(map(str, c["cmd"])) for c in calls + "scripts.rerank_tools.local" in " ".join(map(str, c["cmd"])) for c in calls ), "rerank subprocess should be invoked" # find rerank call - rc = next(c for c in calls if any("rerank_local.py" in str(x) for x in c["cmd"])) + rc = next(c for c in calls if "scripts.rerank_tools.local" in " ".join(map(str, c["cmd"]))) assert rc["timeout"] >= 1.5 and rc["timeout"] <= 2.0 assert res["used_rerank"] is True diff --git a/tests/test_error_paths.py b/tests/test_error_paths.py index 129e2eb4..24765ea3 100644 --- a/tests/test_error_paths.py +++ b/tests/test_error_paths.py @@ -1,24 +1,34 @@ -import os -import asyncio +import importlib +import sys import types import pytest -import scripts.mcp_indexer_server as srv +search_impl = importlib.import_module("scripts.mcp_impl.search") @pytest.mark.service def test_repo_search_malformed_jsonl_subprocess(monkeypatch): # Force subprocess path and simulate malformed JSONL stdout monkeypatch.setenv("HYBRID_IN_PROCESS", "0") + monkeypatch.setenv("REPO_SEARCH_DEFAULT_MODE", "hybrid") + monkeypatch.setenv("RERANKER_ENABLED", "0") + + fake_hybrid = types.ModuleType("scripts.hybrid_search") + fake_hybrid.run_hybrid_search = lambda *a, **k: [] + monkeypatch.setitem(sys.modules, "scripts.hybrid_search", fake_hybrid) async def fake_run(cmd, **kwargs): # Simulate subprocess failure with malformed output return {"ok": False, "code": 1, "stdout": "not-json\n", "stderr": "malformed"} - monkeypatch.setattr(srv, "_run_async", fake_run) - - res = srv.asyncio.get_event_loop().run_until_complete( - srv.repo_search(queries=["x"], limit=1, compact=False) + res = search_impl.asyncio.get_event_loop().run_until_complete( + search_impl._repo_search_impl( + queries=["x"], + limit=1, + compact=False, + run_async_fn=fake_run, + require_auth_session_fn=lambda session: session, + ) ) assert res.get("ok") is False @@ -30,26 +40,30 @@ def test_repo_search_inproc_qdrant_failure_fallback_and_fail(monkeypatch): # In-process hybrid raises (simulating Qdrant connectivity failure), # subprocess fallback also fails. monkeypatch.setenv("HYBRID_IN_PROCESS", "1") - - # Avoid real model load - monkeypatch.setattr(srv, "_get_embedding_model", lambda *a, **k: object()) + monkeypatch.setenv("REPO_SEARCH_DEFAULT_MODE", "hybrid") + monkeypatch.setenv("RERANKER_ENABLED", "0") # Cause in-process path to fail - import scripts.hybrid_search as hy - def boom(*a, **k): raise ConnectionError("qdrant down") - monkeypatch.setattr(hy, "run_hybrid_search", boom) + fake_hybrid = types.ModuleType("scripts.hybrid_search") + fake_hybrid.run_hybrid_search = boom + monkeypatch.setitem(sys.modules, "scripts.hybrid_search", fake_hybrid) # And make the subprocess fallback fail too async def fake_run(cmd, **kwargs): return {"ok": False, "code": 1, "stdout": "", "stderr": "qdrant unreachable"} - monkeypatch.setattr(srv, "_run_async", fake_run) - - res = srv.asyncio.get_event_loop().run_until_complete( - srv.repo_search(queries=["x"], limit=1, compact=True) + res = search_impl.asyncio.get_event_loop().run_until_complete( + search_impl._repo_search_impl( + queries=["x"], + limit=1, + compact=True, + get_embedding_model_fn=lambda *a, **k: object(), + run_async_fn=fake_run, + require_auth_session_fn=lambda session: session, + ) ) assert res.get("ok") is False diff --git a/tests/test_fname_boost.py b/tests/test_fname_boost.py deleted file mode 100644 index fcca42df..00000000 --- a/tests/test_fname_boost.py +++ /dev/null @@ -1,200 +0,0 @@ -"""Tests for production-grade filename boost algorithm. - -The algorithm handles: -- All naming conventions (snake_case, camelCase, PascalCase, kebab-case) -- Acronyms (XMLParser -> xml, parser) -- Position weighting (filename > directory) -- Common token penalties -""" -from scripts.rerank_recursive import ( - _compute_fname_boost, - _split_identifier, - _normalize_token, -) - - -class TestSplitIdentifier: - """Test identifier tokenization across naming conventions.""" - - def test_camel_case(self): - assert _split_identifier("userAuthHandler") == ["user", "auth", "handler"] - - def test_pascal_case(self): - assert _split_identifier("UserAuthHandler") == ["user", "auth", "handler"] - - def test_snake_case(self): - assert _split_identifier("user_auth_handler") == ["user", "auth", "handler"] - - def test_kebab_case(self): - assert _split_identifier("user-auth-handler") == ["user", "auth", "handler"] - - def test_screaming_snake(self): - assert _split_identifier("USER_AUTH_HANDLER") == ["user", "auth", "handler"] - - def test_acronym_prefix(self): - """XMLParser should split into xml, parser.""" - assert _split_identifier("XMLParser") == ["xml", "parser"] - - def test_acronym_suffix(self): - """parseJSON should split into parse, json.""" - assert _split_identifier("parseJSON") == ["parse", "json"] - - def test_http_client(self): - """HTTPClient should split properly.""" - assert _split_identifier("HTTPClient") == ["http", "client"] - - def test_interface_prefix_stripped(self): - """IUserService -> user, service (I prefix stripped).""" - assert _split_identifier("IUserService") == ["user", "service"] - - def test_private_prefix_stripped(self): - """_privateMethod -> private, method.""" - assert _split_identifier("_privateMethod") == ["private", "method"] - - def test_dollar_prefix_stripped(self): - """$scope -> scope.""" - assert _split_identifier("$scope") == ["scope"] - - def test_numbers_separated(self): - """handler2 -> handler (numbers stripped).""" - assert _split_identifier("handler2") == ["handler"] - - def test_dot_notation(self): - """com.company.auth -> com, company, auth.""" - assert _split_identifier("com.company.auth") == ["com", "company", "auth"] - - -class TestNormalizeToken: - """Test plural normalization.""" - - def test_no_abbrev_expansion(self): - forms = _normalize_token("auth") - assert "auth" in forms - assert "authenticate" not in forms - - def test_plural_normalized(self): - forms = _normalize_token("services") - assert "service" in forms - - def test_singular_gets_plural(self): - forms = _normalize_token("service") - assert "services" in forms - - -class TestComputeFnameBoost: - """Test the full filename boost computation.""" - - def test_basic_snake_case_match(self): - """Basic snake_case filename matching.""" - q = "hybrid search" - cand = {"path": "scripts/hybrid_search.py"} - result = _compute_fname_boost(q, cand, 0.15) - assert result > 0.4 # 2 exact filename matches with bonus - - def test_rel_path_fallback(self): - """Should work with rel_path key.""" - q = "hybrid search" - cand = {"rel_path": "scripts/hybrid_search.py"} - result = _compute_fname_boost(q, cand, 0.15) - assert result > 0.4 - - def test_metadata_path_fallback(self): - """Should work with metadata.path.""" - q = "hybrid search" - cand = {"metadata": {"path": "scripts/hybrid_search.py"}} - result = _compute_fname_boost(q, cand, 0.15) - assert result > 0.4 - - def test_disabled_when_factor_zero(self): - """Factor of 0 should return 0.""" - q = "hybrid search" - cand = {"path": "scripts/hybrid_search.py"} - assert _compute_fname_boost(q, cand, 0.0) == 0.0 - - def test_requires_two_matches(self): - """Single token match should not trigger boost.""" - q = "hybrid fusion scoring" # only 'hybrid' matches - cand = {"path": "scripts/hybrid_utils.py"} - assert _compute_fname_boost(q, cand, 0.15) == 0.0 - - def test_camel_case_tokenization(self): - """CamelCase filenames should be properly tokenized.""" - q = "user auth handler" - cand = {"path": "services/UserAuthHandler.ts"} - result = _compute_fname_boost(q, cand, 0.15) - assert result > 0.5 # 3 matches with filename bonus - - def test_pascal_case_tokenization(self): - """PascalCase filenames should be properly tokenized.""" - q = "database connection pool" - cand = {"path": "lib/DatabaseConnectionPool.java"} - result = _compute_fname_boost(q, cand, 0.15) - assert result > 0.5 # 3 matches - - def test_acronym_matching(self): - """Acronyms like XMLParser should match 'xml parser'.""" - q = "xml parser" - cand = {"path": "utils/XMLParser.py"} - result = _compute_fname_boost(q, cand, 0.15) - assert result > 0.4 # 2 exact matches - - def test_substring_matching_authenticate_auth(self): - """'authenticate' contains 'auth', so substring tier can still match it.""" - q = "authenticate user" - cand = {"path": "auth/UserAuth.py"} - result = _compute_fname_boost(q, cand, 0.15) - assert result > 0.3 - - def test_substring_matching_configuration_config(self): - """'configuration' contains 'config', so substring tier can still match it.""" - q = "configuration manager" - cand = {"path": "config/ConfigManager.ts"} - result = _compute_fname_boost(q, cand, 0.15) - assert result > 0.3 - - def test_directory_matching(self): - """Query tokens should match directory names.""" - q = "auth services controller" - cand = {"path": "services/auth/controller.py"} - result = _compute_fname_boost(q, cand, 0.15) - assert result > 0.4 # 3 matches across path - - def test_filename_bonus(self): - """Filename matches should score higher than directory matches.""" - q = "auth handler" - # Filename match - cand1 = {"path": "services/AuthHandler.py"} - # Directory match only - cand2 = {"path": "auth/handler/index.py"} - score1 = _compute_fname_boost(q, cand1, 0.15) - score2 = _compute_fname_boost(q, cand2, 0.15) - assert score1 > score2 # Filename match wins - - def test_deep_path_handling(self): - """Deep Java-style paths should work.""" - q = "user service implementation" - cand = {"path": "src/main/java/com/company/user/service/UserServiceImpl.java"} - result = _compute_fname_boost(q, cand, 0.15) - assert result > 0.3 # Should find matches - - def test_interface_prefix_stripped(self): - """IUserService should match 'user service'.""" - q = "user service" - cand = {"path": "interfaces/IUserService.ts"} - result = _compute_fname_boost(q, cand, 0.15) - assert result > 0.3 - - def test_common_tokens_penalized(self): - """Common tokens like 'utils', 'index' should be penalized.""" - q = "auth utils" - cand = {"path": "utils/AuthUtils.py"} # common token 'utils' - score = _compute_fname_boost(q, cand, 0.15) - # Should still match, but common token gets penalty - assert score > 0 - - def test_jsonish_query(self): - """Handle queries wrapped in JSON-like brackets.""" - q = '["hybrid search fusion"]' - cand = {"path": "scripts/hybrid_search.py"} - result = _compute_fname_boost(q, cand, 0.15) - assert result > 0.3 diff --git a/tests/test_glm_model_config.py b/tests/test_glm_model_config.py index 44baac1a..7823b46c 100644 --- a/tests/test_glm_model_config.py +++ b/tests/test_glm_model_config.py @@ -1,8 +1,20 @@ """Tests for GLM model version configuration and backwards compatibility.""" import os +import sys +from types import SimpleNamespace from unittest.mock import patch, MagicMock +def _install_fake_openai(monkeypatch): + mock_client = MagicMock() + mock_response = MagicMock() + mock_response.choices = [MagicMock(message=MagicMock(content="test response"))] + mock_client.chat.completions.create.return_value = mock_response + openai_module = SimpleNamespace(OpenAI=MagicMock(return_value=mock_client)) + monkeypatch.setitem(sys.modules, "openai", openai_module) + return mock_client + + class TestGLMModelConfig: """Test GLM model version detection and configuration.""" @@ -81,16 +93,10 @@ def test_glm45_config_values(self): class TestGLMRefragClientModelSelection: """Test GLMRefragClient model selection logic.""" - @patch("openai.OpenAI") - def test_default_model_is_glm46(self, mock_openai_class): + def test_default_model_is_glm46(self, monkeypatch): """Test that default model is glm-4.6.""" from scripts.refrag_glm import GLMRefragClient - - mock_client = MagicMock() - mock_response = MagicMock() - mock_response.choices = [MagicMock(message=MagicMock(content="test response"))] - mock_client.chat.completions.create.return_value = mock_response - mock_openai_class.return_value = mock_client + mock_client = _install_fake_openai(monkeypatch) # Remove GLM_MODEL from env to test default, keep GLM_API_KEY env_copy = os.environ.copy() @@ -106,16 +112,10 @@ def test_default_model_is_glm46(self, mock_openai_class): assert call_kwargs["model"] == "glm-4.6" @patch.dict(os.environ, {"GLM_API_KEY": "test-key", "GLM_MODEL": "glm-4.6"}, clear=False) - @patch("openai.OpenAI") - def test_env_model_override(self, mock_openai_class): + def test_env_model_override(self, monkeypatch): """Test that GLM_MODEL env var overrides default.""" from scripts.refrag_glm import GLMRefragClient - - mock_client = MagicMock() - mock_response = MagicMock() - mock_response.choices = [MagicMock(message=MagicMock(content="test response"))] - mock_client.chat.completions.create.return_value = mock_response - mock_openai_class.return_value = mock_client + mock_client = _install_fake_openai(monkeypatch) client = GLMRefragClient() client.generate_with_soft_embeddings("test prompt") @@ -124,16 +124,10 @@ def test_env_model_override(self, mock_openai_class): assert call_kwargs["model"] == "glm-4.6" @patch.dict(os.environ, {"GLM_API_KEY": "test-key", "GLM_MODEL_FAST": "glm-4.5"}, clear=False) - @patch("openai.OpenAI") - def test_fast_model_with_disable_thinking(self, mock_openai_class): + def test_fast_model_with_disable_thinking(self, monkeypatch): """Test that disable_thinking uses GLM_MODEL_FAST.""" from scripts.refrag_glm import GLMRefragClient - - mock_client = MagicMock() - mock_response = MagicMock() - mock_response.choices = [MagicMock(message=MagicMock(content="test response"))] - mock_client.chat.completions.create.return_value = mock_response - mock_openai_class.return_value = mock_client + mock_client = _install_fake_openai(monkeypatch) client = GLMRefragClient() client.generate_with_soft_embeddings("test prompt", disable_thinking=True) @@ -146,16 +140,10 @@ class TestGLMToolStreamSupport: """Test GLM-4.7 tool_stream feature support.""" @patch.dict(os.environ, {"GLM_API_KEY": "test-key", "GLM_MODEL": "glm-4.7"}, clear=False) - @patch("openai.OpenAI") - def test_tool_stream_enabled_for_glm47(self, mock_openai_class): + def test_tool_stream_enabled_for_glm47(self, monkeypatch): """Test that tool_stream is enabled for GLM-4.7 when requested.""" from scripts.refrag_glm import GLMRefragClient - - mock_client = MagicMock() - mock_response = MagicMock() - mock_response.choices = [MagicMock(message=MagicMock(content="test response"))] - mock_client.chat.completions.create.return_value = mock_response - mock_openai_class.return_value = mock_client + mock_client = _install_fake_openai(monkeypatch) client = GLMRefragClient() tools = [{"type": "function", "function": {"name": "test", "parameters": {}}}] @@ -166,16 +154,10 @@ def test_tool_stream_enabled_for_glm47(self, mock_openai_class): assert call_kwargs.get("extra_body", {}).get("tool_stream") is True @patch.dict(os.environ, {"GLM_API_KEY": "test-key", "GLM_MODEL": "glm-4.6"}, clear=False) - @patch("openai.OpenAI") - def test_tool_stream_not_enabled_for_glm46(self, mock_openai_class): + def test_tool_stream_not_enabled_for_glm46(self, monkeypatch): """Test that tool_stream is NOT enabled for GLM-4.6.""" from scripts.refrag_glm import GLMRefragClient - - mock_client = MagicMock() - mock_response = MagicMock() - mock_response.choices = [MagicMock(message=MagicMock(content="test response"))] - mock_client.chat.completions.create.return_value = mock_response - mock_openai_class.return_value = mock_client + mock_client = _install_fake_openai(monkeypatch) client = GLMRefragClient() tools = [{"type": "function", "function": {"name": "test", "parameters": {}}}] @@ -191,16 +173,10 @@ class TestGLMThinkingSupport: """Test GLM thinking/reasoning support.""" @patch.dict(os.environ, {"GLM_API_KEY": "test-key", "GLM_MODEL": "glm-4.7"}, clear=False) - @patch("openai.OpenAI") - def test_enable_thinking_for_glm47(self, mock_openai_class): + def test_enable_thinking_for_glm47(self, monkeypatch): """Test that thinking can be explicitly enabled for GLM-4.7.""" from scripts.refrag_glm import GLMRefragClient - - mock_client = MagicMock() - mock_response = MagicMock() - mock_response.choices = [MagicMock(message=MagicMock(content="test response"))] - mock_client.chat.completions.create.return_value = mock_response - mock_openai_class.return_value = mock_client + mock_client = _install_fake_openai(monkeypatch) client = GLMRefragClient() client.generate_with_soft_embeddings("test prompt", enable_thinking=True) @@ -209,16 +185,10 @@ def test_enable_thinking_for_glm47(self, mock_openai_class): assert call_kwargs.get("extra_body", {}).get("thinking") == {"type": "enabled"} @patch.dict(os.environ, {"GLM_API_KEY": "test-key", "GLM_MODEL": "glm-4.5"}, clear=False) - @patch("openai.OpenAI") - def test_thinking_not_set_for_glm45(self, mock_openai_class): + def test_thinking_not_set_for_glm45(self, monkeypatch): """Test that thinking is NOT set for GLM-4.5 (no thinking support).""" from scripts.refrag_glm import GLMRefragClient - - mock_client = MagicMock() - mock_response = MagicMock() - mock_response.choices = [MagicMock(message=MagicMock(content="test response"))] - mock_client.chat.completions.create.return_value = mock_response - mock_openai_class.return_value = mock_client + mock_client = _install_fake_openai(monkeypatch) client = GLMRefragClient() client.generate_with_soft_embeddings("test prompt", enable_thinking=True) @@ -234,16 +204,10 @@ class TestGLMMaxTokensLimit: """Test max_tokens limiting based on model capabilities.""" @patch.dict(os.environ, {"GLM_API_KEY": "test-key", "GLM_MODEL": "glm-4.7"}, clear=False) - @patch("openai.OpenAI") - def test_max_tokens_capped_to_model_limit(self, mock_openai_class): + def test_max_tokens_capped_to_model_limit(self, monkeypatch): """Test that max_tokens is capped to model's max_output_tokens.""" from scripts.refrag_glm import GLMRefragClient, GLM_MODEL_CONFIGS - - mock_client = MagicMock() - mock_response = MagicMock() - mock_response.choices = [MagicMock(message=MagicMock(content="test response"))] - mock_client.chat.completions.create.return_value = mock_response - mock_openai_class.return_value = mock_client + mock_client = _install_fake_openai(monkeypatch) client = GLMRefragClient() # Request more than GLM-4.7 can output (131072) @@ -253,16 +217,10 @@ def test_max_tokens_capped_to_model_limit(self, mock_openai_class): assert call_kwargs["max_tokens"] <= GLM_MODEL_CONFIGS["glm-4.7"]["max_output_tokens"] @patch.dict(os.environ, {"GLM_API_KEY": "test-key", "GLM_MODEL": "glm-4.5"}, clear=False) - @patch("openai.OpenAI") - def test_max_tokens_uses_smaller_limit_for_glm45(self, mock_openai_class): + def test_max_tokens_uses_smaller_limit_for_glm45(self, monkeypatch): """Test that GLM-4.5 uses its smaller max_output limit.""" from scripts.refrag_glm import GLMRefragClient, GLM_MODEL_CONFIGS - - mock_client = MagicMock() - mock_response = MagicMock() - mock_response.choices = [MagicMock(message=MagicMock(content="test response"))] - mock_client.chat.completions.create.return_value = mock_response - mock_openai_class.return_value = mock_client + mock_client = _install_fake_openai(monkeypatch) client = GLMRefragClient() # Request more than GLM-4.5 can output (8192) diff --git a/tests/test_globs_and_snippet.py b/tests/test_globs_and_snippet.py index 4c30ff48..d52abc1e 100644 --- a/tests/test_globs_and_snippet.py +++ b/tests/test_globs_and_snippet.py @@ -1,14 +1,53 @@ import importlib +import asyncio import builtins import json import types from pathlib import Path +from types import SimpleNamespace import pytest # Import targets hyb = importlib.import_module("scripts.hybrid_search") -srv = importlib.import_module("scripts.mcp_indexer_server") +search_impl = importlib.import_module("scripts.mcp_impl.search") + + +@pytest.fixture(autouse=True) +def fake_qdrant_models(monkeypatch): + hybrid_qdrant = importlib.import_module("scripts.hybrid.qdrant") + + class FakeModels: + class SearchParams: + def __init__(self, **kwargs): + self.__dict__.update(kwargs) + + class QuantizationSearchParams: + def __init__(self, **kwargs): + self.__dict__.update(kwargs) + + class Filter: + def __init__(self, **kwargs): + self.__dict__.update(kwargs) + + class FieldCondition: + def __init__(self, **kwargs): + self.__dict__.update(kwargs) + + class MatchValue: + def __init__(self, **kwargs): + self.__dict__.update(kwargs) + + class MatchAny: + def __init__(self, **kwargs): + self.__dict__.update(kwargs) + + class SparseVector: + def __init__(self, **kwargs): + self.__dict__.update(kwargs) + + monkeypatch.setattr(hyb, "models", FakeModels) + monkeypatch.setattr(hybrid_qdrant, "models", FakeModels) class _Pt: @@ -33,6 +72,19 @@ class FakeQdrant: def __init__(self, points): self._points = points + def get_collection(self, collection): + return SimpleNamespace( + config=SimpleNamespace( + params=SimpleNamespace( + vectors={ + "unit-test": SimpleNamespace(size=8), + "dense": SimpleNamespace(size=8), + }, + sparse_vectors={}, + ) + ) + ) + # dense_query tries query_points first, then search on exception def query_points(self, **kwargs): return _QP(self._points) @@ -116,6 +168,32 @@ def test_run_hybrid_search_slugged_path_globs(monkeypatch): assert "/work/other/docs/readme.md" not in paths +@pytest.mark.unit +def test_run_hybrid_search_under_recursive_scope(monkeypatch): + pts = [ + _Pt("1", "/work/repo/space/ship/a.py"), + _Pt("2", "/work/repo/direct/tools/b.py"), + ] + monkeypatch.setattr(hyb, "get_qdrant_client", lambda *a, **k: FakeQdrant(pts)) + monkeypatch.setattr(hyb, "return_qdrant_client", lambda *a, **k: None) + monkeypatch.setenv("EMBEDDING_MODEL", "unit-test") + monkeypatch.setenv("QDRANT_URL", "http://localhost:6333") + monkeypatch.setattr(hyb, "TextEmbedding", lambda *a, **k: FakeEmbed()) + monkeypatch.setattr(hyb, "_get_embedding_model", lambda *a, **k: FakeEmbed()) + + items = hyb.run_hybrid_search( + queries=["rotate heading"], + limit=10, + per_path=2, + under="space", + expand=False, + model=FakeEmbed(), + ) + paths = {it.get("path") for it in items} + assert "/work/repo/space/ship/a.py" in paths + assert "/work/repo/direct/tools/b.py" not in paths + + @pytest.mark.unit def test_dense_query_preserves_collection_on_filter_drop(monkeypatch): calls = [] @@ -147,6 +225,32 @@ def query_points(self, **kwargs): assert calls[1].get("query_filter") is None or calls[1].get("filter") is None +@pytest.mark.unit +def test_run_pure_dense_search_honors_per_path_cap(): + points = [ + SimpleNamespace( + score=0.99, + payload={"metadata": {"path": "/work/repo/a.py", "start_line": 1, "end_line": 2}}, + ), + SimpleNamespace( + score=0.98, + payload={"metadata": {"path": "/work/repo/a.py", "start_line": 10, "end_line": 11}}, + ), + SimpleNamespace( + score=0.97, + payload={"metadata": {"path": "/work/repo/b.py", "start_line": 3, "end_line": 4}}, + ), + ] + + items = hyb._shape_dense_points( + points, + limit=2, + per_path=1, + ) + + assert [item["path"] for item in items] == ["/work/repo/a.py", "/work/repo/b.py"] + + @pytest.mark.unit def test_collection_prefers_env_over_state(monkeypatch, tmp_path): # State file should be ignored when COLLECTION_NAME env var is set @@ -162,18 +266,9 @@ def test_collection_prefers_env_over_state(monkeypatch, tmp_path): @pytest.mark.unit def test_repo_search_snippet_strict_cap_after_highlight(monkeypatch): - # Stub run_hybrid_search to emit a single result with a known path and range - async def fake_run(**kwargs): - return {"results": [{"path": "/work/f.txt", "start_line": 1, "end_line": 1}]} - # Force in-process shaping to trigger snippet code path monkeypatch.setenv("HYBRID_IN_PROCESS", "1") - # Monkeypatch srv.hybrid_search.run_hybrid_search result pathing via repo_search flow - monkeypatch.setattr( - srv, "_tokens_from_queries", lambda q: ["foo"] - ) # ensure highlight runs - # Fake open for the specific /work path big_line = "foo " * 1000 # large content to exceed cap _orig_open = builtins.open @@ -184,9 +279,9 @@ def fake_open(path, *a, **k): return _orig_open(path, *a, **k) # pragma: no cover # Ensure sandbox passes - monkeypatch.setattr(srv.os.path, "isabs", lambda p: True) - monkeypatch.setattr(srv.os.path, "realpath", lambda p: "/work/f.txt") - monkeypatch.setenv("MCP_SNIPPET_MAX_BYTES", "64") + monkeypatch.setattr(search_impl.os.path, "isabs", lambda p: True) + monkeypatch.setattr(search_impl.os.path, "realpath", lambda p: "/work/f.txt") + monkeypatch.setattr(search_impl, "SNIPPET_MAX_BYTES", 64) # Stub hybrid_search.run_hybrid_search to return a single item import sys @@ -202,13 +297,21 @@ def run_hybrid_search(**kwargs): import io - # Patch open builtin used by server + # Patch open builtin used by repo_search implementation. monkeypatch.setattr(builtins, "open", fake_open) # Execute - res = srv.asyncio.get_event_loop().run_until_complete( - srv.repo_search( - query="foo", include_snippet=True, highlight_snippet=True, context_lines=0 + res = asyncio.run( + search_impl._repo_search_impl( + query="foo", + mode="hybrid", + include_snippet=True, + highlight_snippet=True, + context_lines=0, + get_embedding_model_fn=lambda _name: object(), + require_auth_session_fn=lambda session: session, + do_highlight_snippet_fn=lambda snippet, _tokens: snippet, + run_async_fn=lambda *_a, **_k: {"ok": True, "code": 0, "stdout": "", "stderr": ""}, ) ) snip = res["results"][0].get("snippet", "") @@ -218,7 +321,7 @@ def run_hybrid_search(**kwargs): @pytest.mark.unit def test_repo_search_docstring_clean(): - doc = srv.repo_search.__doc__ + doc = search_impl._repo_search_impl.__doc__ assert doc and "Zero-config code search" in doc # Ensure stray inline pseudo-code is not embedded in docstring assert "Accept common alias keys from clients" not in doc diff --git a/tests/test_golden_structure.py b/tests/test_golden_structure.py index ab19917b..c8ac4d99 100644 --- a/tests/test_golden_structure.py +++ b/tests/test_golden_structure.py @@ -1,5 +1,6 @@ import os import json +import asyncio from pathlib import Path import pytest @@ -22,8 +23,8 @@ def stub(*a, **k): monkeypatch.setattr(hy, "run_hybrid_search", stub) - res = srv.asyncio.get_event_loop().run_until_complete( - srv.repo_search(queries=["q"], limit=2, compact=True) + res = asyncio.run( + srv.repo_search(queries=["q"], limit=2, compact=True, mode="hybrid") ) # Normalize subset: path/start_line/end_line/symbol only diff --git a/tests/test_graph_delete_verification_normalization.py b/tests/test_graph_delete_verification_normalization.py new file mode 100644 index 00000000..5531de7b --- /dev/null +++ b/tests/test_graph_delete_verification_normalization.py @@ -0,0 +1,45 @@ +from pathlib import PureWindowsPath + + +def test_graph_delete_verification_normalizes_caller_path(): + # Unit-level guard: watcher delete verification must query graph edges using + # the same path normalization as graph edge writes/deletes (Windows -> POSIX). + from scripts.watch_index_core import processor as proc + + captured = {} + + class DummyClient: + def scroll( + self, + *, + collection_name, + scroll_filter, + with_payload=False, + with_vectors=False, + limit=1, + ): + captured["collection_name"] = collection_name + captured["filter"] = scroll_filter + return ([], None) + + client = DummyClient() + path = PureWindowsPath(r"C:\repo\foo.py") + + has_edges = proc._path_has_graph_edges(client, "base_collection", path) + assert has_edges is False + + flt = captured["filter"] + assert flt is not None + assert getattr(flt, "must", None) + cond = flt.must[0] + assert cond.key == "caller_path" + + match = cond.match + values = [] + if hasattr(match, "any") and match.any is not None: + values = list(match.any) + elif hasattr(match, "value"): + values = [match.value] + + assert "C:/repo/foo.py" in values + diff --git a/tests/test_health_check.py b/tests/test_health_check.py new file mode 100644 index 00000000..0dceda58 --- /dev/null +++ b/tests/test_health_check.py @@ -0,0 +1,85 @@ +from types import SimpleNamespace + +import pytest + + +pytestmark = pytest.mark.unit + + +class _FakeVec: + def __init__(self, size): + self.size = size + + +class _FakeCollections: + def __init__(self, names): + self.collections = [SimpleNamespace(name=name) for name in names] + + +class _FakeClient: + collection_names = ["context-engine", "stale-empty"] + valid_vector_collections = {"context-engine"} + + def __init__(self, *args, **kwargs): + self.checked = [] + + def get_collections(self): + return _FakeCollections(self.collection_names) + + def get_collection(self, name): + self.checked.append(name) + vectors = ( + {"fast-bge-base-en-v1.5": _FakeVec(768)} + if name in self.valid_vector_collections + else {} + ) + return SimpleNamespace( + config=SimpleNamespace( + params=SimpleNamespace(vectors=vectors), + hnsw_config=SimpleNamespace(m=16, ef_construct=256), + ) + ) + + def query_points(self, *args, **kwargs): + return SimpleNamespace(points=[]) + + +class _FakeEmbedding: + def embed(self, texts): + for _ in texts: + yield SimpleNamespace(tolist=lambda: [0.0] * 768) + + +def test_health_check_checks_all_collections_without_crashing_on_mismatch(monkeypatch, capsys): + import scripts.health_check as health_check + + monkeypatch.setenv("COLLECTION_NAME", "context-engine") + monkeypatch.setenv("EMBEDDING_MODEL", "fast/bge-base-en-v1.5") + monkeypatch.setattr(health_check, "QdrantClient", _FakeClient) + monkeypatch.setattr(health_check, "get_embedding_model", lambda *_: _FakeEmbedding()) + monkeypatch.setattr(health_check, "get_model_dimension", lambda *_: 768) + monkeypatch.setattr(health_check, "ensure_collections", lambda *_: 0) + + health_check.main() + + output = capsys.readouterr().out + assert "Checking collection: context-engine" in output + assert "Checking collection: stale-empty" in output + assert "Skipping vector query for stale-empty" in output + + +def test_health_check_missing_named_vector_does_not_keyerror(monkeypatch, capsys): + import scripts.health_check as health_check + + monkeypatch.setenv("COLLECTION_NAME", "stale-empty") + monkeypatch.setenv("EMBEDDING_MODEL", "fast/bge-base-en-v1.5") + monkeypatch.setattr(health_check, "QdrantClient", _FakeClient) + monkeypatch.setattr(health_check, "get_embedding_model", lambda *_: _FakeEmbedding()) + monkeypatch.setattr(health_check, "get_model_dimension", lambda *_: 768) + monkeypatch.setattr(health_check, "ensure_collections", lambda *_: 0) + + health_check.main() + + output = capsys.readouterr().out + assert "Expected vector name present" in output + assert "Skipping vector query for stale-empty" in output diff --git a/tests/test_hybrid_cache_bm25.py b/tests/test_hybrid_cache_bm25.py index 9dea510d..35ab245c 100644 --- a/tests/test_hybrid_cache_bm25.py +++ b/tests/test_hybrid_cache_bm25.py @@ -6,6 +6,23 @@ hyb = importlib.import_module("scripts.hybrid_search") +@pytest.fixture(autouse=True) +def fake_qdrant_models(monkeypatch): + hybrid_qdrant = importlib.import_module("scripts.hybrid.qdrant") + + class FakeModels: + class SearchParams: + def __init__(self, **kwargs): + self.__dict__.update(kwargs) + + class QuantizationSearchParams: + def __init__(self, **kwargs): + self.__dict__.update(kwargs) + + monkeypatch.setattr(hyb, "models", FakeModels) + monkeypatch.setattr(hybrid_qdrant, "models", FakeModels) + + class _Pt: def __init__(self, pid, path, code=""): self.id = pid @@ -39,6 +56,28 @@ def search(self, **kwargs): self.calls += 1 return self._points + def get_collection(self, collection): + return type( + "CollectionInfo", + (), + { + "config": type( + "Config", + (), + { + "params": type( + "Params", + (), + { + "vectors": {"unit-test": type("Vector", (), {"size": 8})()}, + "sparse_vectors": {}, + }, + )() + }, + )() + }, + )() + class _FakeEmbed: class _Vec: @@ -131,4 +170,3 @@ def test_lexical_bm25_boost_is_gentle_and_matches_multiplier(): # Gentle behavior: overall change should be modest (within 50%) ratio = weighted / base assert 0.5 <= ratio <= 1.5, f"BM25 weighting should be gentle, got ratio={ratio:.3f}" - diff --git a/tests/test_hybrid_cli_json.py b/tests/test_hybrid_cli_json.py index d13e4d9d..1b62779d 100644 --- a/tests/test_hybrid_cli_json.py +++ b/tests/test_hybrid_cli_json.py @@ -1,12 +1,38 @@ import json import sys +import types from types import SimpleNamespace import importlib def test_hybrid_cli_json_output(monkeypatch, capsys): + class DummyClient: + def __init__(self, *args, **kwargs): + self.args = args + self.kwargs = kwargs + + class DummyModels(types.ModuleType): + def __getattr__(self, name): + def _factory(*args, **kwargs): + return SimpleNamespace(_model=name, args=args, **kwargs) + + return _factory + + fake_models = DummyModels("qdrant_client.models") + fake_qdrant = types.ModuleType("qdrant_client") + fake_qdrant.QdrantClient = DummyClient + fake_qdrant.models = fake_models + monkeypatch.setitem(sys.modules, "qdrant_client", fake_qdrant) + monkeypatch.setitem(sys.modules, "qdrant_client.models", fake_models) + + monkeypatch.setenv("HYBRID_LEXICAL_WEIGHT", "0.20") + monkeypatch.setenv("HYBRID_LEX_VECTOR_WEIGHT", "0.20") + monkeypatch.setenv("HYBRID_DENSE_WEIGHT", "1.5") + importlib.reload(importlib.import_module("scripts.hybrid.config")) + importlib.reload(importlib.import_module("scripts.hybrid.ranking")) hy = importlib.import_module("scripts.hybrid_search") + hy = importlib.reload(hy) embedder = importlib.import_module("scripts.embedder") class DummyVec: @@ -24,11 +50,6 @@ def embed(self, texts): for _ in texts: yield DummyVec() - class DummyClient: - def __init__(self, *args, **kwargs): - self.args = args - self.kwargs = kwargs - def fake_dense_query(client, vec_name, vector, flt, per_query, collection_name=None, query_text=None): md = { "path": "/work/pkg/a.py", diff --git a/tests/test_hybrid_ranking.py b/tests/test_hybrid_ranking.py index 5edcc949..71055cd5 100644 --- a/tests/test_hybrid_ranking.py +++ b/tests/test_hybrid_ranking.py @@ -18,10 +18,13 @@ # Fixture: Import ranking module # ============================================================================ @pytest.fixture -def ranking_module(): +def ranking_module(monkeypatch): """Import hybrid ranking module.""" import importlib + monkeypatch.setenv("HYBRID_LEXICAL_WEIGHT", "0.20") + monkeypatch.setenv("HYBRID_LEX_VECTOR_WEIGHT", "0.20") ranking = importlib.import_module("scripts.hybrid.ranking") + ranking = importlib.reload(ranking) return ranking diff --git a/tests/test_index_journal.py b/tests/test_index_journal.py new file mode 100644 index 00000000..cd465704 --- /dev/null +++ b/tests/test_index_journal.py @@ -0,0 +1,665 @@ +#!/usr/bin/env python3 +import importlib +import json +from pathlib import Path +from unittest.mock import MagicMock + +import pytest + + +pytestmark = pytest.mark.unit + + +@pytest.fixture +def ws_module(monkeypatch, tmp_path): + ws_root = tmp_path / "work" + ws_root.mkdir(parents=True, exist_ok=True) + monkeypatch.setenv("WORKSPACE_PATH", str(ws_root)) + monkeypatch.setenv("WATCH_ROOT", str(ws_root)) + monkeypatch.delenv("MULTI_REPO_MODE", raising=False) + ws = importlib.import_module("scripts.workspace_state") + return importlib.reload(ws) + + +def test_index_journal_roundtrip(ws_module, tmp_path): + repo_name = "repo-1234567890abcdef" + file_path = tmp_path / "work" / repo_name / "src" / "app.py" + file_path.parent.mkdir(parents=True, exist_ok=True) + + ws_module.upsert_index_journal_entries( + [ + {"path": str(file_path), "op_type": "upsert", "content_hash": "abc123"}, + {"path": str(file_path.with_name("old.py")), "op_type": "delete"}, + ], + workspace_path=str(tmp_path / "work" / repo_name), + repo_name=repo_name, + ) + + pending = [ + str(e["path"]) + for e in ws_module.list_pending_index_journal_entries( + workspace_path=str(tmp_path / "work" / repo_name), + repo_name=repo_name, + ) + ] + assert str(file_path.resolve()) in pending + assert str((file_path.with_name("old.py")).resolve()) in pending + + ws_module.update_index_journal_entry_status( + str(file_path), + status="done", + workspace_path=str(tmp_path / "work" / repo_name), + repo_name=repo_name, + ) + pending_after = [ + str(e["path"]) + for e in ws_module.list_pending_index_journal_entries( + workspace_path=str(tmp_path / "work" / repo_name), + repo_name=repo_name, + ) + ] + assert str(file_path.resolve()) not in pending_after + assert str((file_path.with_name("old.py")).resolve()) in pending_after + + +def test_index_journal_entries_include_operation_types(ws_module, tmp_path): + repo_name = "repo-1234567890abcdef" + file_path = tmp_path / "work" / repo_name / "src" / "entry.py" + file_path.parent.mkdir(parents=True, exist_ok=True) + + ws_module.upsert_index_journal_entries( + [ + {"path": str(file_path), "op_type": "upsert", "content_hash": "abc123"}, + {"path": str(file_path.with_name("gone.py")), "op_type": "delete"}, + ], + workspace_path=str(tmp_path / "work" / repo_name), + repo_name=repo_name, + ) + + entries = ws_module.list_pending_index_journal_entries( + workspace_path=str(tmp_path / "work" / repo_name), + repo_name=repo_name, + ) + by_path = {entry["path"]: entry for entry in entries} + assert by_path[str(file_path.resolve())]["op_type"] == "upsert" + assert by_path[str((file_path.with_name("gone.py")).resolve())]["op_type"] == "delete" + + +def test_index_journal_clear_entries(ws_module, tmp_path): + repo_name = "repo-1234567890abcdef" + file_path = tmp_path / "work" / repo_name / "src" / "entry.py" + file_path.parent.mkdir(parents=True, exist_ok=True) + + ws_module.upsert_index_journal_entries( + [ + {"path": str(file_path), "op_type": "upsert", "content_hash": "abc123"}, + {"path": str(file_path.with_name("gone.py")), "op_type": "delete"}, + ], + workspace_path=str(tmp_path / "work" / repo_name), + repo_name=repo_name, + ) + + removed = ws_module.clear_index_journal_entries( + workspace_path=str(tmp_path / "work" / repo_name), + repo_name=repo_name, + ) + + assert removed == 2 + assert ( + ws_module.list_pending_index_journal_entries( + workspace_path=str(tmp_path / "work" / repo_name), + repo_name=repo_name, + ) + == [] + ) + + +def test_index_journal_bulk_status_updates_once(ws_module, monkeypatch, tmp_path): + repo_name = "repo-1234567890abcdef" + repo_root = tmp_path / "work" / repo_name + repo_root.mkdir(parents=True, exist_ok=True) + paths = [repo_root / f"src/{idx}.py" for idx in range(3)] + + ws_module.upsert_index_journal_entries( + [{"path": str(path), "op_type": "upsert"} for path in paths], + workspace_path=str(repo_root), + repo_name=repo_name, + ) + + original_update = ws_module._update_index_journal + calls = [] + + def counted_update(*args, **kwargs): + calls.append(1) + return original_update(*args, **kwargs) + + monkeypatch.setattr(ws_module, "_update_index_journal", counted_update) + ws_module.update_index_journal_entries_status( + [{"path": str(path), "status": "done"} for path in paths], + workspace_path=str(repo_root), + repo_name=repo_name, + ) + + assert len(calls) == 1 + assert ws_module.list_pending_index_journal_entries( + workspace_path=str(repo_root), repo_name=repo_name + ) == [] + + +def test_index_journal_summary_reports_retryable_entries(ws_module, tmp_path): + repo_name = "repo-1234567890abcdef" + repo_root = tmp_path / "work" / repo_name + repo_root.mkdir(parents=True, exist_ok=True) + pending = repo_root / "src/pending.py" + failed = repo_root / "src/failed.py" + + ws_module.upsert_index_journal_entries( + [ + {"path": str(pending), "op_type": "upsert"}, + {"path": str(failed), "op_type": "delete"}, + ], + workspace_path=str(repo_root), + repo_name=repo_name, + ) + ws_module.update_index_journal_entry_status( + str(failed), + status="failed", + error="qdrant unavailable", + workspace_path=str(repo_root), + repo_name=repo_name, + remove_on_done=False, + ) + + summary = ws_module.get_index_journal_summary( + workspace_path=str(repo_root), repo_name=repo_name + ) + + assert summary["total"] == 2 + assert summary["retryable"] == 2 + assert summary["outstanding"] == 2 + assert summary["counts"]["pending"] == 1 + assert summary["counts"]["failed"] == 1 + assert summary["sample_errors"] == [ + {"path": str(failed.resolve()), "error": "qdrant unavailable"} + ] + + +def test_index_journal_summary_counts_in_progress_as_outstanding(ws_module, tmp_path): + repo_name = "repo-1234567890abcdef" + repo_root = tmp_path / "work" / repo_name + repo_root.mkdir(parents=True, exist_ok=True) + path = repo_root / "src/in_progress.py" + + ws_module.upsert_index_journal_entries( + [{"path": str(path), "op_type": "upsert"}], + workspace_path=str(repo_root), + repo_name=repo_name, + ) + ws_module.update_index_journal_entry_status( + str(path), + status="in_progress", + workspace_path=str(repo_root), + repo_name=repo_name, + remove_on_done=False, + ) + + summary = ws_module.get_index_journal_summary( + workspace_path=str(repo_root), repo_name=repo_name + ) + + assert summary["counts"]["in_progress"] == 1 + assert summary["retryable"] == 0 + assert summary["outstanding"] == 1 + + +def test_index_journal_unknown_status_is_reported_but_not_retried(ws_module, tmp_path): + repo_name = "repo-1234567890abcdef" + repo_root = tmp_path / "work" / repo_name + path = repo_root / "src" / "unknown.py" + journal_path = ws_module._get_index_journal_path(str(repo_root), repo_name) + journal_path.parent.mkdir(parents=True, exist_ok=True) + journal_path.write_text( + json.dumps( + { + "version": 1, + "operations": { + str(path.resolve()): { + "path": str(path.resolve()), + "op_type": "upsert", + "status": "mystery", + "attempts": 0, + } + }, + } + ), + encoding="utf-8", + ) + + summary = ws_module.get_index_journal_summary( + workspace_path=str(repo_root), repo_name=repo_name + ) + + assert summary["counts"]["unknown"] == 1 + assert summary["retryable"] == 0 + assert summary["outstanding"] == 1 + assert ws_module.list_pending_index_journal_entries( + workspace_path=str(repo_root), repo_name=repo_name + ) == [] + + +def test_index_journal_aggregates_repo_scoped_entries(ws_module, tmp_path): + repo_name = "repo-1234567890abcdef" + file_path = tmp_path / "work" / repo_name / "src" / "x.py" + file_path.parent.mkdir(parents=True, exist_ok=True) + + ws_module.upsert_index_journal_entries( + [{"path": str(file_path), "op_type": "upsert", "content_hash": "abc123"}], + workspace_path=str(tmp_path / "work" / repo_name), + repo_name=repo_name, + ) + + pending = [ + str(e["path"]) + for e in ws_module.list_pending_index_journal_entries( + workspace_path=str(tmp_path / "work") + ) + ] + assert str(file_path.resolve()) in pending + + +@pytest.mark.parametrize("repo_name", ["repo-1234567890abcdef", "frontend"]) +def test_index_journal_aggregates_repo_scoped_entries_in_multi_repo_mode( + monkeypatch, tmp_path, repo_name +): + ws_root = tmp_path / "work" + ws_root.mkdir(parents=True, exist_ok=True) + monkeypatch.setenv("WORKSPACE_PATH", str(ws_root)) + monkeypatch.setenv("WATCH_ROOT", str(ws_root)) + monkeypatch.setenv("MULTI_REPO_MODE", "1") + ws_module = importlib.import_module("scripts.workspace_state") + ws_module = importlib.reload(ws_module) + + file_name = "app.ts" if repo_name == "frontend" else "multi.py" + file_path = ws_root / repo_name / "src" / file_name + file_path.parent.mkdir(parents=True, exist_ok=True) + + ws_module.upsert_index_journal_entries( + [{"path": str(file_path), "op_type": "upsert", "content_hash": "abc123"}], + workspace_path=str(ws_root / repo_name), + repo_name=repo_name, + ) + + pending = [ + str(e["path"]) + for e in ws_module.list_pending_index_journal_entries(workspace_path=str(ws_root)) + ] + assert str(file_path.resolve()) in pending + + +def test_index_journal_discovery_ignores_arbitrary_workspace_directories( + monkeypatch, tmp_path +): + ws_root = tmp_path / "work" + ws_root.mkdir(parents=True, exist_ok=True) + (ws_root / "logs").mkdir() + (ws_root / "logs" / "large.log").write_text("noise\n", encoding="utf-8") + repo_name = "frontend" + repo_state = ws_root / ".codebase" / "repos" / repo_name + repo_state.mkdir(parents=True, exist_ok=True) + (repo_state / "index_journal.json").write_text( + json.dumps({"operations": {}}), encoding="utf-8" + ) + + monkeypatch.setenv("WORKSPACE_PATH", str(ws_root)) + monkeypatch.setenv("WATCH_ROOT", str(ws_root)) + monkeypatch.setenv("MULTI_REPO_MODE", "1") + ws_module = importlib.import_module("scripts.workspace_state") + ws_module = importlib.reload(ws_module) + + discovered = ws_module._discover_journal_repositories(str(ws_root)) + + assert discovered == [(repo_name, None)] + + +def test_index_journal_aggregates_split_watch_and_metadata_roots(monkeypatch, tmp_path): + watch_root = tmp_path / "work" + metadata_root = tmp_path / "metadata" + repo_name = "Context-Engine-41e67959950c8ab3" + file_path = watch_root / repo_name / "src" / "split.py" + file_path.parent.mkdir(parents=True, exist_ok=True) + metadata_root.mkdir(parents=True, exist_ok=True) + monkeypatch.setenv("WATCH_ROOT", str(watch_root)) + monkeypatch.setenv("WORK_DIR", str(watch_root)) + monkeypatch.setenv("CTXCE_METADATA_ROOT", str(metadata_root)) + monkeypatch.setenv("MULTI_REPO_MODE", "1") + ws_module = importlib.import_module("scripts.workspace_state") + ws_module = importlib.reload(ws_module) + + ws_module.upsert_index_journal_entries( + [{"path": str(file_path), "op_type": "upsert", "content_hash": "abc123"}], + workspace_path=str(file_path.parent.parent), + repo_name=repo_name, + ) + + pending = [ + str(e["path"]) + for e in ws_module.list_pending_index_journal_entries(workspace_path=str(watch_root)) + ] + assert str(file_path.resolve()) in pending + + ws_module.update_index_journal_entry_status( + str(file_path), + status="done", + workspace_path=str(file_path.parent.parent), + repo_name=repo_name, + ) + pending_after = [ + str(e["path"]) + for e in ws_module.list_pending_index_journal_entries(workspace_path=str(watch_root)) + ] + assert str(file_path.resolve()) not in pending_after + + +def test_index_journal_file_is_group_writable(ws_module, tmp_path): + repo_name = "repo-1234567890abcdef" + file_path = tmp_path / "work" / repo_name / "src" / "perm.py" + file_path.parent.mkdir(parents=True, exist_ok=True) + + ws_module.upsert_index_journal_entries( + [{"path": str(file_path), "op_type": "upsert", "content_hash": "abc123"}], + workspace_path=str(tmp_path / "work" / repo_name), + repo_name=repo_name, + ) + + journal_path = ws_module._get_index_journal_path( + str(tmp_path / "work" / repo_name), repo_name + ) + assert journal_path.exists() + assert oct(journal_path.stat().st_mode & 0o777) == "0o666" + + +def test_index_journal_failed_entry_respects_retry_delay(ws_module, monkeypatch, tmp_path): + repo_name = "repo-1234567890abcdef" + file_path = tmp_path / "work" / repo_name / "src" / "retry.py" + file_path.parent.mkdir(parents=True, exist_ok=True) + monkeypatch.setenv("INDEX_JOURNAL_RETRY_DELAY_SECS", "60") + + ws_module.upsert_index_journal_entries( + [{"path": str(file_path), "op_type": "upsert", "content_hash": "abc123"}], + workspace_path=str(tmp_path / "work" / repo_name), + repo_name=repo_name, + ) + ws_module.update_index_journal_entry_status( + str(file_path), + status="failed", + error="boom", + workspace_path=str(tmp_path / "work" / repo_name), + repo_name=repo_name, + remove_on_done=False, + ) + + pending = [ + str(e["path"]) + for e in ws_module.list_pending_index_journal_entries( + workspace_path=str(tmp_path / "work" / repo_name), + repo_name=repo_name, + ) + ] + assert str(file_path.resolve()) not in pending + + +def test_index_journal_failed_entry_honors_max_attempts(ws_module, monkeypatch, tmp_path): + repo_name = "repo-1234567890abcdef" + file_path = tmp_path / "work" / repo_name / "src" / "retry2.py" + file_path.parent.mkdir(parents=True, exist_ok=True) + monkeypatch.setenv("INDEX_JOURNAL_RETRY_DELAY_SECS", "0") + monkeypatch.setenv("INDEX_JOURNAL_MAX_ATTEMPTS", "1") + + ws_module.upsert_index_journal_entries( + [{"path": str(file_path), "op_type": "upsert", "content_hash": "abc123"}], + workspace_path=str(tmp_path / "work" / repo_name), + repo_name=repo_name, + ) + ws_module.update_index_journal_entry_status( + str(file_path), + status="failed", + error="boom", + workspace_path=str(tmp_path / "work" / repo_name), + repo_name=repo_name, + remove_on_done=False, + ) + + pending = [ + str(e["path"]) + for e in ws_module.list_pending_index_journal_entries( + workspace_path=str(tmp_path / "work" / repo_name), + repo_name=repo_name, + ) + ] + assert str(file_path.resolve()) not in pending + + +def test_processor_delete_marks_journal_done(monkeypatch, tmp_path): + proc_mod = importlib.import_module("scripts.watch_index_core.processor") + + missing = tmp_path / "missing.py" + assert not missing.exists() + + monkeypatch.setattr(proc_mod, "_detect_repo_for_file", lambda p: tmp_path) + monkeypatch.setattr(proc_mod, "_get_collection_for_file", lambda p: "coll") + monkeypatch.setattr(proc_mod, "_set_status_indexing", lambda *a, **k: None) + monkeypatch.setattr(proc_mod, "persist_indexing_config", lambda *a, **k: None) + monkeypatch.setattr(proc_mod, "update_indexing_status", lambda *a, **k: None) + monkeypatch.setattr(proc_mod, "get_workspace_state", lambda *a, **k: {}) + monkeypatch.setattr(proc_mod, "is_staging_enabled", lambda: False) + monkeypatch.setattr(proc_mod, "_log_activity", lambda *a, **k: None) + monkeypatch.setattr(proc_mod, "_extract_repo_name_from_path", lambda *_: "repo") + monkeypatch.setattr(proc_mod, "remove_cached_file", lambda *a, **k: None) + + delete_mock = MagicMock() + graph_delete_mock = MagicMock() + journal_mock = MagicMock() + monkeypatch.setattr(proc_mod.idx, "delete_points_by_path", delete_mock) + monkeypatch.setattr(proc_mod.idx, "delete_graph_edges_by_path", graph_delete_mock) + monkeypatch.setattr(proc_mod, "_verify_delete_committed", lambda *a, **k: True) + monkeypatch.setattr(proc_mod, "_verify_graph_delete_committed", lambda *a, **k: True) + monkeypatch.setattr(proc_mod, "update_index_journal_entries_status", journal_mock) + + proc_mod._process_paths( + [missing], + client=MagicMock(), + model=None, + vector_name="vec", + model_dim=1, + workspace_path=str(tmp_path), + ) + + delete_mock.assert_called_once() + assert graph_delete_mock.call_count == 2 + assert graph_delete_mock.call_args_list[0].kwargs["repo"] == "repo" + assert graph_delete_mock.call_args_list[1].kwargs["repo"] is None + journal_mock.assert_called_once() + assert journal_mock.call_args.args[0][0]["status"] == "done" + + +def test_processor_honors_delete_journal_for_existing_file(monkeypatch, tmp_path): + proc_mod = importlib.import_module("scripts.watch_index_core.processor") + + existing = tmp_path / "present.py" + existing.write_text("print('x')\n", encoding="utf-8") + + monkeypatch.setattr(proc_mod, "_detect_repo_for_file", lambda p: tmp_path) + monkeypatch.setattr(proc_mod, "_get_collection_for_file", lambda p: "coll") + monkeypatch.setattr(proc_mod, "_set_status_indexing", lambda *a, **k: None) + monkeypatch.setattr(proc_mod, "persist_indexing_config", lambda *a, **k: None) + monkeypatch.setattr(proc_mod, "update_indexing_status", lambda *a, **k: None) + monkeypatch.setattr(proc_mod, "get_workspace_state", lambda *a, **k: {}) + monkeypatch.setattr(proc_mod, "is_staging_enabled", lambda: False) + monkeypatch.setattr(proc_mod, "_log_activity", lambda *a, **k: None) + monkeypatch.setattr(proc_mod, "_extract_repo_name_from_path", lambda *_: "repo") + monkeypatch.setattr(proc_mod, "remove_cached_file", lambda *a, **k: None) + monkeypatch.setattr( + proc_mod, + "list_pending_index_journal_entries", + lambda *a, **k: [{"path": str(existing.resolve()), "op_type": "delete"}], + ) + + delete_mock = MagicMock() + graph_delete_mock = MagicMock() + journal_mock = MagicMock() + monkeypatch.setattr(proc_mod.idx, "delete_points_by_path", delete_mock) + monkeypatch.setattr(proc_mod.idx, "delete_graph_edges_by_path", graph_delete_mock) + monkeypatch.setattr(proc_mod, "_verify_delete_committed", lambda *a, **k: True) + monkeypatch.setattr(proc_mod, "_verify_graph_delete_committed", lambda *a, **k: True) + monkeypatch.setattr(proc_mod, "update_index_journal_entries_status", journal_mock) + + proc_mod._process_paths( + [existing], + client=MagicMock(), + model=None, + vector_name="vec", + model_dim=1, + workspace_path=str(tmp_path), + ) + + delete_mock.assert_called_once() + assert graph_delete_mock.call_count == 2 + assert graph_delete_mock.call_args_list[0].kwargs["repo"] == "repo" + assert graph_delete_mock.call_args_list[1].kwargs["repo"] is None + journal_mock.assert_called_once() + assert journal_mock.call_args.args[0][0]["status"] == "done" + + +def test_processor_relinks_move_journal_before_delete(monkeypatch, tmp_path): + proc_mod = importlib.import_module("scripts.watch_index_core.processor") + + src = tmp_path / "src.py" + dest = tmp_path / "dest.py" + dest.write_text("print('dest')\n", encoding="utf-8") + + monkeypatch.setattr(proc_mod, "_detect_repo_for_file", lambda p: tmp_path) + monkeypatch.setattr(proc_mod, "_get_collection_for_file", lambda p: "coll") + monkeypatch.setattr(proc_mod, "_set_status_indexing", lambda *a, **k: None) + monkeypatch.setattr(proc_mod, "persist_indexing_config", lambda *a, **k: None) + monkeypatch.setattr(proc_mod, "update_indexing_status", lambda *a, **k: None) + monkeypatch.setattr(proc_mod, "get_workspace_state", lambda *a, **k: {}) + monkeypatch.setattr(proc_mod, "is_staging_enabled", lambda: False) + monkeypatch.setattr(proc_mod, "_log_activity", lambda *a, **k: None) + monkeypatch.setattr(proc_mod, "_extract_repo_name_from_path", lambda *_: "repo") + monkeypatch.setattr(proc_mod, "remove_cached_file", lambda *a, **k: None) + monkeypatch.setattr(proc_mod, "set_cached_file_hash", lambda *a, **k: None) + monkeypatch.setattr( + proc_mod, + "list_pending_index_journal_entries", + lambda *a, **k: [ + {"path": str(src.resolve()), "op_type": "delete", "content_hash": "cafebabe"}, + {"path": str(dest.resolve()), "op_type": "upsert", "content_hash": "cafebabe"}, + ], + ) + + rename_mock = MagicMock(return_value=(3, "cafebabe")) + delete_mock = MagicMock() + journal_mock = MagicMock() + monkeypatch.setattr(proc_mod, "_rename_in_store", rename_mock) + monkeypatch.setattr(proc_mod.idx, "delete_points_by_path", delete_mock) + monkeypatch.setattr(proc_mod, "update_index_journal_entries_status", journal_mock) + + proc_mod._process_paths( + [src, dest], + client=MagicMock(), + model=MagicMock(), + vector_name="vec", + model_dim=1, + workspace_path=str(tmp_path), + ) + + rename_mock.assert_called_once() + delete_mock.assert_not_called() + updates = journal_mock.call_args.args[0] + done_paths = [entry["path"] for entry in updates if entry.get("status") == "done"] + assert str(dest.resolve()) in done_paths + assert str(src.resolve()) in done_paths + + +def test_processor_skips_internal_git_path_without_collection_resolution(monkeypatch): + proc_mod = importlib.import_module("scripts.watch_index_core.processor") + + internal = Path("/work/.git/HEAD") + + monkeypatch.setattr(proc_mod, "_set_status_indexing", lambda *a, **k: None) + monkeypatch.setattr(proc_mod, "persist_indexing_config", lambda *a, **k: None) + monkeypatch.setattr(proc_mod, "update_indexing_status", lambda *a, **k: None) + monkeypatch.setattr(proc_mod, "get_workspace_state", lambda *a, **k: {}) + monkeypatch.setattr(proc_mod, "is_staging_enabled", lambda: False) + monkeypatch.setattr(proc_mod, "_log_activity", lambda *a, **k: None) + monkeypatch.setattr(proc_mod, "_extract_repo_name_from_path", lambda *_: "repo") + + collection_mock = MagicMock(return_value="should-not-be-used") + journal_mock = MagicMock() + monkeypatch.setattr(proc_mod, "_get_collection_for_file", collection_mock) + monkeypatch.setattr(proc_mod, "update_index_journal_entry_status", journal_mock) + monkeypatch.setattr( + proc_mod, + "list_pending_index_journal_entries", + lambda *a, **k: [{"path": str(internal), "op_type": "delete"}], + ) + + proc_mod._process_paths( + [internal], + client=MagicMock(), + model=None, + vector_name="vec", + model_dim=1, + workspace_path="/work", + ) + + collection_mock.assert_not_called() + journal_mock.assert_called_once() + assert journal_mock.call_args.kwargs["status"] == "done" + + +def test_processor_force_upsert_empty_file_marks_done(monkeypatch, tmp_path): + proc_mod = importlib.import_module("scripts.watch_index_core.processor") + + empty_file = tmp_path / "pkg" / "__init__.py" + empty_file.parent.mkdir(parents=True, exist_ok=True) + empty_file.write_text("", encoding="utf-8") + + monkeypatch.setattr(proc_mod, "_detect_repo_for_file", lambda p: tmp_path) + monkeypatch.setattr(proc_mod, "_get_collection_for_file", lambda p: "coll") + monkeypatch.setattr(proc_mod, "_set_status_indexing", lambda *a, **k: None) + monkeypatch.setattr(proc_mod, "persist_indexing_config", lambda *a, **k: None) + monkeypatch.setattr(proc_mod, "update_indexing_status", lambda *a, **k: None) + monkeypatch.setattr(proc_mod, "get_workspace_state", lambda *a, **k: {}) + monkeypatch.setattr(proc_mod, "is_staging_enabled", lambda: False) + monkeypatch.setattr(proc_mod, "_log_activity", lambda *a, **k: None) + monkeypatch.setattr(proc_mod, "_extract_repo_name_from_path", lambda *_: "repo") + monkeypatch.setattr(proc_mod, "remove_cached_file", lambda *a, **k: None) + monkeypatch.setattr(proc_mod, "_run_indexing_strategy", lambda *a, **k: False) + monkeypatch.setattr(proc_mod, "_path_has_indexed_points", lambda *a, **k: False) + + journal_mock = MagicMock() + monkeypatch.setattr(proc_mod, "update_index_journal_entries_status", journal_mock) + monkeypatch.setattr( + proc_mod, + "list_pending_index_journal_entries", + lambda *a, **k: [ + { + "path": str(empty_file.resolve()), + "op_type": "upsert", + "content_hash": "da39a3ee5e6b4b0d3255bfef95601890afd80709", + } + ], + ) + + proc_mod._process_paths( + [empty_file], + client=MagicMock(), + model=MagicMock(), + vector_name="vec", + model_dim=1, + workspace_path=str(tmp_path), + ) + + journal_mock.assert_called_once() + assert journal_mock.call_args.args[0][0]["status"] == "done" diff --git a/tests/test_ingest_cli.py b/tests/test_ingest_cli.py new file mode 100644 index 00000000..493d1a68 --- /dev/null +++ b/tests/test_ingest_cli.py @@ -0,0 +1,57 @@ +import sys +from pathlib import Path + +import pytest + + +@pytest.mark.unit +def test_cli_force_collection_disables_multi_repo_enumeration(monkeypatch, tmp_path: Path): + from scripts.ingest import cli + + # Create fake repo dirs to prove we are not enumerating them. + (tmp_path / "repo_a").mkdir() + (tmp_path / "repo_b").mkdir() + + calls = [] + + def _fake_index_repo( + root, + qdrant_url, + api_key, + collection, + model_name, + recreate, + dedupe, + skip_unchanged, + pseudo_mode, + schema_mode, + ): + calls.append( + { + "root": Path(root), + "collection": collection, + "recreate": recreate, + "dedupe": dedupe, + "skip_unchanged": skip_unchanged, + } + ) + + monkeypatch.setattr(cli, "index_repo", _fake_index_repo) + monkeypatch.setattr(cli, "is_multi_repo_mode", lambda: True) + monkeypatch.setattr(cli, "get_collection_name", lambda *_: "should-not-use") + + monkeypatch.setenv("MULTI_REPO_MODE", "1") + monkeypatch.setenv("COLLECTION_NAME", "forced-collection") + monkeypatch.setenv("CTXCE_FORCE_COLLECTION_NAME", "1") + + monkeypatch.setattr( + sys, + "argv", + ["ingest_code.py", "--root", str(tmp_path)], + ) + + cli.main() + + assert len(calls) == 1 + assert calls[0]["root"] == tmp_path + assert calls[0]["collection"] == "forced-collection" diff --git a/tests/test_ingest_schema_mode.py b/tests/test_ingest_schema_mode.py index c766089b..faa367be 100644 --- a/tests/test_ingest_schema_mode.py +++ b/tests/test_ingest_schema_mode.py @@ -61,12 +61,10 @@ def create_payload_index(self, collection_name, field_name, field_schema): def test_schema_mode_validate_errors_on_missing_vectors(monkeypatch): - monkeypatch.setenv("PATTERN_VECTORS", "1") monkeypatch.setattr(ingq, "LEX_SPARSE_MODE", False) existing_vectors = { "code": object(), - ingq.LEX_VECTOR_NAME: object(), } payload_schema = {field: object() for field in ingq.PAYLOAD_INDEX_FIELDS} client = FakeClient( @@ -89,12 +87,11 @@ def test_schema_mode_validate_errors_on_missing_vectors(monkeypatch): def test_schema_mode_migrate_adds_missing_vectors_and_indexes(monkeypatch): - monkeypatch.setenv("PATTERN_VECTORS", "1") monkeypatch.setattr(ingq, "LEX_SPARSE_MODE", False) + ingq.ENSURED_PAYLOAD_INDEX_COLLECTIONS.discard("test-collection") existing_vectors = { "code": object(), - ingq.LEX_VECTOR_NAME: object(), } client = FakeClient( collection_exists=True, @@ -113,15 +110,15 @@ def test_schema_mode_migrate_adds_missing_vectors_and_indexes(monkeypatch): assert client.create_calls == [] assert client.update_calls updated_vectors = client.update_calls[0]["vectors_config"] - assert ingq.PATTERN_VECTOR_NAME in updated_vectors + assert ingq.LEX_VECTOR_NAME in updated_vectors assert any( c["field_name"] == "metadata.language" for c in client.payload_index_calls ) def test_schema_mode_create_creates_collection_only(monkeypatch): - monkeypatch.setenv("PATTERN_VECTORS", "0") monkeypatch.setattr(ingq, "LEX_SPARSE_MODE", False) + ingq.ENSURED_PAYLOAD_INDEX_COLLECTIONS.discard("test-collection") client = FakeClient(collection_exists=False) @@ -138,3 +135,15 @@ def test_schema_mode_create_creates_collection_only(monkeypatch): assert any( c["field_name"] == "metadata.language" for c in client.payload_index_calls ) + + +def test_ensure_payload_indexes_memoized_per_process(): + client = FakeClient(collection_exists=True) + ingq.ENSURED_PAYLOAD_INDEX_COLLECTIONS.discard("test-collection") + + ingq.ensure_payload_indexes(client, "test-collection") + first_count = len(client.payload_index_calls) + ingq.ensure_payload_indexes(client, "test-collection") + + assert first_count == len(ingq.PAYLOAD_INDEX_FIELDS) + assert len(client.payload_index_calls) == first_count diff --git a/tests/test_integration_qdrant.py b/tests/test_integration_qdrant.py index 65cef7f5..f40442ce 100644 --- a/tests/test_integration_qdrant.py +++ b/tests/test_integration_qdrant.py @@ -1,6 +1,7 @@ import os import json import uuid +import asyncio import importlib import pytest @@ -41,11 +42,11 @@ def embed(self, texts): @pytest.mark.integration def test_index_and_search_minirepo(tmp_path, monkeypatch, qdrant_container): # Env for services - os.environ["QDRANT_URL"] = qdrant_container - os.environ["COLLECTION_NAME"] = f"test-{uuid.uuid4().hex[:8]}" - os.environ["USE_TREE_SITTER"] = "0" - os.environ["HYBRID_IN_PROCESS"] = "1" - os.environ["EMBEDDING_MODEL"] = "fake" + monkeypatch.setenv("QDRANT_URL", qdrant_container) + monkeypatch.setenv("COLLECTION_NAME", f"test-{uuid.uuid4().hex[:8]}") + monkeypatch.setenv("USE_TREE_SITTER", "0") + monkeypatch.setenv("HYBRID_IN_PROCESS", "1") + monkeypatch.setenv("EMBEDDING_MODEL", "fake") # Stub embeddings everywhere (FakeEmbedder produces 32-dim vectors) monkeypatch.setattr(ing, "TextEmbedding", lambda *a, **k: FakeEmbedder("fake")) @@ -75,7 +76,7 @@ def test_index_and_search_minirepo(tmp_path, monkeypatch, qdrant_container): ) # Search directly via async function - res = srv.asyncio.get_event_loop().run_until_complete( + res = asyncio.run( srv.repo_search( queries=["def f"], limit=5, @@ -92,11 +93,11 @@ def test_index_and_search_minirepo(tmp_path, monkeypatch, qdrant_container): @pytest.mark.integration def test_filters_language_and_path(tmp_path, monkeypatch, qdrant_container): # Reuse container; set env - os.environ["QDRANT_URL"] = qdrant_container - os.environ.setdefault("COLLECTION_NAME", f"test-{uuid.uuid4().hex[:8]}") - os.environ["USE_TREE_SITTER"] = "0" - os.environ["HYBRID_IN_PROCESS"] = "1" - os.environ["EMBEDDING_MODEL"] = "fake" + monkeypatch.setenv("QDRANT_URL", qdrant_container) + monkeypatch.setenv("COLLECTION_NAME", f"test-{uuid.uuid4().hex[:8]}") + monkeypatch.setenv("USE_TREE_SITTER", "0") + monkeypatch.setenv("HYBRID_IN_PROCESS", "1") + monkeypatch.setenv("EMBEDDING_MODEL", "fake") # Stub embeddings (FakeEmbedder produces 32-dim vectors) monkeypatch.setattr(ing, "TextEmbedding", lambda *a, **k: FakeEmbedder("fake")) @@ -127,19 +128,19 @@ def test_filters_language_and_path(tmp_path, monkeypatch, qdrant_container): f_md = str(tmp_path / "pkg" / "b.md") # Filter by language=python should bias toward .py - res1 = srv.asyncio.get_event_loop().run_until_complete( + res1 = asyncio.run( srv.repo_search(queries=["def"], limit=5, language="python", compact=False) ) assert any(f_py in (r.get("path") or "") for r in res1.get("results", [])) # Filter by ext=txt should retrieve text file - res2 = srv.asyncio.get_event_loop().run_until_complete( + res2 = asyncio.run( srv.repo_search(queries=["hello"], limit=5, ext="md", compact=False) ) assert any(f_md in (r.get("path") or "") for r in res2.get("results", [])) # Path glob to only allow pkg/*.py - res3 = srv.asyncio.get_event_loop().run_until_complete( + res3 = asyncio.run( srv.repo_search( queries=["def"], limit=5, diff --git a/tests/test_mcp_pattern_search.py b/tests/test_mcp_pattern_search.py deleted file mode 100644 index 5c3a4b25..00000000 --- a/tests/test_mcp_pattern_search.py +++ /dev/null @@ -1,226 +0,0 @@ -import pytest -import inspect - - -@pytest.fixture -def anyio_backend(): - return "asyncio" - - -def test_search_by_pattern_description_accepts_min_score(): - """Verify the real function signature accepts min_score to prevent TypeError.""" - from scripts.pattern_detection.search import search_by_pattern_description - - sig = inspect.signature(search_by_pattern_description) - params = list(sig.parameters.keys()) - - assert "min_score" in params, ( - "search_by_pattern_description must accept min_score kwarg; " - "MCP wrapper passes it through" - ) - - -@pytest.mark.anyio -async def test_mcp_pattern_search_nl_toon_min_score_pass_through(monkeypatch): - """NL + TOON path should pass min_score through without crashing.""" - from scripts.pattern_detection.search import search_by_pattern_description - from scripts.mcp_impl import pattern_search as mcp - - # Bypass real imports - mcp._PATTERN_SEARCH_LOADED = True - - captured = {} - - # Use a wrapper that has the REAL signature to catch mismatches - real_sig = inspect.signature(search_by_pattern_description) - real_params = set(real_sig.parameters.keys()) - - def fake_search_by_pattern_description(**kwargs): - # Validate only expected kwargs are passed - unexpected = set(kwargs.keys()) - real_params - if unexpected: - raise TypeError(f"Unexpected kwargs: {unexpected}") - captured.update(kwargs) - # Simulate TOON output: results is a string, not a list - return { - "ok": True, - "results": "results[0]:", - "total": 0, - "search_mode": "natural_language", - } - - # Avoid code-path; ensure NL branch is chosen - monkeypatch.setattr(mcp, "_pattern_search_fn", None) - monkeypatch.setattr(mcp, "_search_by_pattern_description_fn", fake_search_by_pattern_description) - - result = await mcp._pattern_search_impl( - query="find retry pattern", # NL triggers description path - output_format="toon", - min_score=0.9, - ) - - assert result["ok"] is True - assert result["query_mode"] == "description" - # No crash on TOON string results; min_score forwarded to core call - assert captured["min_score"] == 0.9 - - -@pytest.mark.anyio -async def test_mcp_pattern_search_code_path_aroma_params(monkeypatch): - """Code path should pass aroma_rerank and aroma_alpha through.""" - from scripts.pattern_detection.search import pattern_search - from scripts.mcp_impl import pattern_search as mcp - - mcp._PATTERN_SEARCH_LOADED = True - - captured = {} - real_sig = inspect.signature(pattern_search) - real_params = set(real_sig.parameters.keys()) - - def fake_pattern_search(**kwargs): - unexpected = set(kwargs.keys()) - real_params - if unexpected: - raise TypeError(f"Unexpected kwargs: {unexpected}") - captured.update(kwargs) - return {"ok": True, "results": [], "total": 0, "search_mode": "aroma"} - - monkeypatch.setattr(mcp, "_pattern_search_fn", fake_pattern_search) - monkeypatch.setattr(mcp, "_search_by_pattern_description_fn", None) - - result = await mcp._pattern_search_impl( - query="for i in range(3): try: pass except: sleep(i)", # Code path - aroma_rerank=True, - aroma_alpha=0.7, - ) - - assert result["ok"] is True - assert captured["aroma_rerank"] is True - assert captured["aroma_alpha"] == 0.7 - - -@pytest.mark.anyio -async def test_mcp_pattern_search_auto_routes_nl_even_with_language(monkeypatch): - """NL queries should route to description even if language is provided.""" - from scripts.mcp_impl import pattern_search as mcp - - mcp._PATTERN_SEARCH_LOADED = True - - called = {} - - def fake_pattern_search(**kwargs): - called["code"] = kwargs - return {"ok": True, "results": [], "total": 0, "search_mode": "structural"} - - def fake_search_by_pattern_description(**kwargs): - called["description"] = kwargs - return {"ok": True, "results": [], "total": 0, "search_mode": "natural_language"} - - monkeypatch.setattr(mcp, "_pattern_search_fn", fake_pattern_search) - monkeypatch.setattr(mcp, "_search_by_pattern_description_fn", fake_search_by_pattern_description) - - result = await mcp._pattern_search_impl( - query="retry with exponential backoff", - language="python", # should not force code path - ) - - assert result["query_mode"] == "description" - assert "description" in called - assert "code" not in called - - -@pytest.mark.anyio -async def test_mcp_pattern_search_query_mode_override(monkeypatch): - """Explicit query_mode overrides auto detection.""" - from scripts.mcp_impl import pattern_search as mcp - - mcp._PATTERN_SEARCH_LOADED = True - - called = {} - - def fake_pattern_search(**kwargs): - called.setdefault("code", 0) - called["code"] += 1 - return {"ok": True, "results": [], "total": 0, "search_mode": "structural"} - - def fake_search_by_pattern_description(**kwargs): - called.setdefault("description", 0) - called["description"] += 1 - return {"ok": True, "results": [], "total": 0, "search_mode": "natural_language"} - - monkeypatch.setattr(mcp, "_pattern_search_fn", fake_pattern_search) - monkeypatch.setattr(mcp, "_search_by_pattern_description_fn", fake_search_by_pattern_description) - - # Force code mode on NL text - result_code = await mcp._pattern_search_impl( - query="retry with exponential backoff", - query_mode="code", - ) - assert result_code["query_mode"] == "code" - - # Force description mode on code-ish text - result_desc = await mcp._pattern_search_impl( - query="for i in range(3): pass", - query_mode="description", - ) - assert result_desc["query_mode"] == "description" - - assert called["code"] == 1 - assert called["description"] == 1 - - -@pytest.mark.anyio -async def test_mcp_min_score_defaults_by_path(monkeypatch): - """ - Regression test: MCP uses path-specific min_score defaults. - - - Code path: min_score defaults to 0.5 (vector scores are typically high) - - NL path: min_score defaults to 0.0 (keyword overlap scores are often low) - - This prevents the NL path from silently returning empty results when - the uniform 0.5 default would filter out valid low-scoring matches. - """ - from scripts.mcp_impl import pattern_search as mcp - - mcp._PATTERN_SEARCH_LOADED = True - - code_captured = {} - nl_captured = {} - - def fake_pattern_search(**kwargs): - code_captured.update(kwargs) - return {"ok": True, "results": [], "total": 0, "search_mode": "structural"} - - def fake_nl_search(**kwargs): - nl_captured.update(kwargs) - return {"ok": True, "results": [], "total": 0, "search_mode": "natural_language"} - - monkeypatch.setattr(mcp, "_pattern_search_fn", fake_pattern_search) - monkeypatch.setattr(mcp, "_search_by_pattern_description_fn", fake_nl_search) - - # Code path (no explicit min_score) - await mcp._pattern_search_impl( - query="for i in range(3): try: pass except: sleep(i)", - ) - assert code_captured["min_score"] == 0.5, "Code path should default min_score=0.5" - - # NL path (no explicit min_score) - await mcp._pattern_search_impl( - query="find retry with exponential backoff", - ) - assert nl_captured["min_score"] == 0.0, "NL path should default min_score=0.0" - - # Explicit min_score should override both paths - code_captured.clear() - nl_captured.clear() - - await mcp._pattern_search_impl( - query="def foo(): pass", - min_score=0.8, - ) - assert code_captured["min_score"] == 0.8, "Explicit min_score should override code default" - - await mcp._pattern_search_impl( - query="database connection pooling pattern", - min_score=0.3, - ) - assert nl_captured["min_score"] == 0.3, "Explicit min_score should override NL default" diff --git a/tests/test_mcp_router.py b/tests/test_mcp_router.py deleted file mode 100644 index 776e0401..00000000 --- a/tests/test_mcp_router.py +++ /dev/null @@ -1,233 +0,0 @@ -#!/usr/bin/env python3 -""" -Tests for mcp_router.py - Intent classification and tool routing. - -Tests cover: -- Intent classification (rule-based and ML fallback) -- Plan building for various query types -- HTTP client helpers -""" -import importlib -import os - -import pytest - -pytestmark = pytest.mark.unit - - -# ============================================================================ -# Fixture: Router module import -# ============================================================================ -@pytest.fixture -def router_module(monkeypatch): - """Import mcp_router with isolated environment.""" - monkeypatch.delenv("MCP_HTTP_URL", raising=False) - monkeypatch.delenv("MCP_INDEXER_HTTP_URL", raising=False) - - router = importlib.import_module("scripts.mcp_router") - return importlib.reload(router) - - -# ============================================================================ -# Tests: Intent Constants -# ============================================================================ -class TestIntentConstants: - """Tests for intent constant definitions.""" - - def test_intent_constants_defined(self, router_module): - """All expected intent constants are defined.""" - assert router_module.INTENT_ANSWER == "answer" - assert router_module.INTENT_SEARCH == "search" - assert router_module.INTENT_INDEX == "index" - assert router_module.INTENT_PRUNE == "prune" - assert router_module.INTENT_STATUS == "status" - assert router_module.INTENT_LIST == "list" - - def test_search_specialized_intents(self, router_module): - """Specialized search intents are defined.""" - assert router_module.INTENT_SEARCH_TESTS == "search_tests" - assert router_module.INTENT_SEARCH_CONFIG == "search_config" - assert router_module.INTENT_SEARCH_CALLERS == "search_callers" - - -# ============================================================================ -# Tests: Intent Classification (Rule-based) -# ============================================================================ -class TestClassifyIntentRules: - """Tests for _classify_intent_rules function.""" - - def test_status_intent_patterns(self, router_module): - """Status-related queries are classified correctly.""" - status_queries = [ - "qdrant status", - "indexing status", - "collection status", - ] - for q in status_queries: - intent = router_module._classify_intent_rules(q) - assert intent == router_module.INTENT_STATUS, f"Failed for: {q}" - - def test_list_intent_patterns(self, router_module): - """List-related queries are classified correctly.""" - list_queries = [ - "list collections", - "show all collections", - ] - for q in list_queries: - intent = router_module._classify_intent_rules(q) - assert intent == router_module.INTENT_LIST, f"Failed for: {q}" - - def test_search_tests_intent(self, router_module): - """Test search queries are classified correctly.""" - queries = [ - "find tests for foo", - "search for test files", - ] - for q in queries: - intent = router_module._classify_intent_rules(q) - assert intent == router_module.INTENT_SEARCH_TESTS, f"Failed for: {q}" - - def test_search_config_intent(self, router_module): - """Config search queries are classified correctly.""" - queries = [ - "find config for database", - "where is the yaml config", - ] - for q in queries: - intent = router_module._classify_intent_rules(q) - assert intent == router_module.INTENT_SEARCH_CONFIG, f"Failed for: {q}" - - -# ============================================================================ -# Tests: High-level classify_intent -# ============================================================================ -class TestClassifyIntent: - """Tests for the main classify_intent function.""" - - def test_classify_intent_returns_intent(self, router_module): - """classify_intent returns a valid intent string.""" - intent = router_module.classify_intent("reindex the codebase") - # Should return some intent (index or answer depending on ML) - assert intent is not None - assert isinstance(intent, str) - - def test_classify_intent_status(self, router_module): - """Status queries classified correctly.""" - intent = router_module.classify_intent("qdrant status") - assert intent == router_module.INTENT_STATUS - - def test_classify_intent_list(self, router_module): - """List queries classified correctly.""" - intent = router_module.classify_intent("list collections") - assert intent == router_module.INTENT_LIST - - -# ============================================================================ -# Tests: Build Plan -# ============================================================================ -class TestBuildPlan: - """Tests for build_plan function.""" - - def test_build_plan_returns_list(self, router_module): - """build_plan returns a list of (tool, args) tuples.""" - # Use a query that triggers rule-based classification (avoids embedding model) - plan = router_module.build_plan("list collections") - - assert isinstance(plan, list) - assert len(plan) >= 1 - # Each item is a tuple of (tool_name, args_dict) - tool_name, args = plan[0] - assert isinstance(tool_name, str) - assert isinstance(args, dict) - - def test_build_plan_status_tool(self, router_module): - """Status queries map to qdrant_status tool.""" - plan = router_module.build_plan("qdrant status") - - tool_name, args = plan[0] - assert tool_name == "qdrant_status" - - def test_build_plan_list_tool(self, router_module): - """List queries map to qdrant_list tool.""" - plan = router_module.build_plan("list collections") - - tool_name, args = plan[0] - assert tool_name == "qdrant_list" - - def test_build_plan_search_tests_tool(self, router_module): - """Test search queries map to search_tests_for tool.""" - plan = router_module.build_plan("find tests for authentication") - - tool_name, args = plan[0] - assert tool_name == "search_tests_for" - assert "query" in args - - def test_build_plan_search_config_tool(self, router_module): - """Config search queries map to search_config_for tool.""" - plan = router_module.build_plan("find config for database") - - tool_name, args = plan[0] - assert tool_name == "search_config_for" - - def test_build_plan_includes_query(self, router_module): - """build_plan includes the query in args for search tools.""" - # Use a query that triggers rule-based classification (avoids embedding model) - plan = router_module.build_plan("find tests for authentication") - - tool_name, args = plan[0] - # search_tests_for includes query in args - assert "query" in args or tool_name in {"qdrant_status", "qdrant_list", "qdrant_prune"} - - -# ============================================================================ -# Tests: HTTP Helpers -# ============================================================================ -class TestHttpHelpers: - """Tests for HTTP client helper functions.""" - - def test_filter_args_removes_none(self, router_module): - """_filter_args removes None values from dict.""" - args = {"a": 1, "b": None, "c": "hello", "d": None} - filtered = router_module._filter_args(args) - - assert filtered == {"a": 1, "c": "hello"} - - def test_filter_args_preserves_false(self, router_module): - """_filter_args preserves False and 0 values.""" - args = {"a": False, "b": 0, "c": None} - filtered = router_module._filter_args(args) - - assert "a" in filtered - assert "b" in filtered - assert "c" not in filtered - - def test_parse_stream_or_json_parses_json(self, router_module): - """_parse_stream_or_json parses valid JSON.""" - body = b'{"result": "success"}' - parsed = router_module._parse_stream_or_json(body) - - assert parsed == {"result": "success"} - - -# ============================================================================ -# Tests: Failure Response Detection -# ============================================================================ -class TestFailureResponseDetection: - """Tests for _is_failure_response function.""" - - def test_success_response_not_failure(self, router_module): - """Successful responses are not failures.""" - resp = {"result": "data", "ok": True} - assert router_module._is_failure_response(resp) is False - - def test_empty_response_not_failure(self, router_module): - """Empty dicts are not failures.""" - assert router_module._is_failure_response({}) is False - - def test_detects_isError_true(self, router_module): - """Detects responses with isError=True.""" - # Note: depends on actual implementation - resp = {"isError": True, "content": []} - result = router_module._is_failure_response(resp) - # May be True or False depending on implementation - assert isinstance(result, bool) diff --git a/tests/test_micro_span_budget.py b/tests/test_micro_span_budget.py index 5648de30..65ed82b8 100644 --- a/tests/test_micro_span_budget.py +++ b/tests/test_micro_span_budget.py @@ -98,7 +98,7 @@ def test_adaptive_span_sizing_failure_is_non_fatal(monkeypatch): monkeypatch.setenv("COLLECTION_NAME", "dummy") # Force extent lookup to throw; the budgeter should swallow it. - import scripts.hybrid_ranking as hr + from scripts.hybrid import ranking as hr monkeypatch.setattr(hr, "_get_symbol_extent", lambda *a, **k: (_ for _ in ()).throw(RuntimeError("boom"))) items = [ diff --git a/tests/test_negative_args.py b/tests/test_negative_args.py index 32dadd52..a0d819a1 100644 --- a/tests/test_negative_args.py +++ b/tests/test_negative_args.py @@ -1,4 +1,5 @@ import os +import asyncio import pytest import scripts.mcp_indexer_server as srv @@ -6,6 +7,8 @@ @pytest.mark.service def test_repo_search_conflicting_filters_empty_ok(monkeypatch): + # This test validates hybrid filter handling (non-dense path), not dense-default mode. + # Keep mode explicit so global REPO_SEARCH_DEFAULT_MODE=dense does not change test semantics. # In-process, but no results due to conflicting filters (simulate by returning []) monkeypatch.setenv("HYBRID_IN_PROCESS", "1") monkeypatch.setattr(srv, "_get_embedding_model", lambda *a, **k: object()) @@ -14,8 +17,14 @@ def test_repo_search_conflicting_filters_empty_ok(monkeypatch): monkeypatch.setattr(hy, "run_hybrid_search", lambda *a, **k: []) - res = srv.asyncio.get_event_loop().run_until_complete( - srv.repo_search(queries=["foo"], limit=3, ext="cpp", compact=True) + res = asyncio.run( + srv.repo_search( + queries=["foo"], + limit=3, + ext="cpp", + compact=True, + mode="hybrid", + ) ) assert res.get("ok") is True diff --git a/tests/test_path_scope.py b/tests/test_path_scope.py new file mode 100644 index 00000000..06369860 --- /dev/null +++ b/tests/test_path_scope.py @@ -0,0 +1,61 @@ +import importlib + + +ps = importlib.import_module("scripts.path_scope") + + +def test_normalize_under_strips_work_prefix(): + assert ps.normalize_under("/work/scripts/mcp_impl") == "scripts/mcp_impl" + + +def test_normalize_under_keeps_repo_prefixed_path(): + assert ( + ps.normalize_under("/work/Context-Engine/scripts/mcp_impl") + == "Context-Engine/scripts/mcp_impl" + ) + + +def test_normalize_under_rebases_single_segment_from_cwd(monkeypatch, tmp_path): + repo = tmp_path / "repo" + (repo / "nested" / "scope").mkdir(parents=True) + monkeypatch.setattr(ps, "_repo_root_hint", lambda: str(repo)) + monkeypatch.setattr(ps.os, "getcwd", lambda: str(repo / "nested")) + + assert ps.normalize_under("scope") == "nested/scope" + + +def test_normalize_under_does_not_rebase_when_top_level_exists(monkeypatch, tmp_path): + repo = tmp_path / "repo" + (repo / "nested" / "scope").mkdir(parents=True) + (repo / "scope").mkdir(parents=True) + monkeypatch.setattr(ps, "_repo_root_hint", lambda: str(repo)) + monkeypatch.setattr(ps.os, "getcwd", lambda: str(repo / "nested")) + + assert ps.normalize_under("scope") == "scope" + + +def test_normalize_under_expands_unique_segment(monkeypatch, tmp_path): + repo = tmp_path / "repo" + (repo / "alpha" / "mcp_impl").mkdir(parents=True) + monkeypatch.setattr(ps, "_repo_root_hint", lambda: str(repo)) + monkeypatch.setattr(ps.os, "getcwd", lambda: str(repo)) + ps._unique_segment_path.cache_clear() + + assert ps.normalize_under("mcp_impl") == "alpha/mcp_impl" + + +def test_normalize_under_keeps_ambiguous_segment(monkeypatch, tmp_path): + repo = tmp_path / "repo" + (repo / "alpha" / "dup").mkdir(parents=True) + (repo / "beta" / "dup").mkdir(parents=True) + monkeypatch.setattr(ps, "_repo_root_hint", lambda: str(repo)) + monkeypatch.setattr(ps.os, "getcwd", lambda: str(repo)) + ps._unique_segment_path.cache_clear() + + assert ps.normalize_under("dup") == "dup" + + +def test_metadata_matches_under_without_repo_hint_for_work_repo_paths(): + md = {"path": "/work/repo/space/ship/a.py"} + assert ps.metadata_matches_under(md, "space") + assert not ps.metadata_matches_under(md, "direct") diff --git a/tests/test_pattern_search_e2e.py b/tests/test_pattern_search_e2e.py deleted file mode 100644 index 9385aba9..00000000 --- a/tests/test_pattern_search_e2e.py +++ /dev/null @@ -1,181 +0,0 @@ -"""Tests for pattern_search functionality. - -Tests that: -1. Pattern extraction produces consistent signatures -2. Cross-language patterns have high similarity -3. Code vs NL detection works correctly -4. Integration with Qdrant (when available) -""" -import os -import uuid -import time -import pytest -import numpy as np - -# Test fixtures -RETRY_PATTERN_PYTHON = ''' -for attempt in range(3): - try: - result = make_request() - break - except Exception: - time.sleep(2 ** attempt) -''' - -RETRY_PATTERN_GO = ''' -for i := 0; i < 3; i++ { - result, err := makeRequest() - if err == nil { - break - } - time.Sleep(time.Duration(1< py_file_sim, ( - f"Retry patterns should be more similar ({py_go_sim:.3f}) " - f"than retry vs file ({py_file_sim:.3f})" - ) - assert py_go_sim > 0.3, f"Cross-language retry similarity too low: {py_go_sim:.3f}" - - -def test_pattern_vector_dimensions(): - """Test that pattern vectors have correct dimensions.""" - from scripts.pattern_detection import PatternExtractor, PatternEncoder - - extractor = PatternExtractor() - encoder = PatternEncoder() - - sig = extractor.extract(RETRY_PATTERN_PYTHON, "python") - vec = encoder.encode(sig) - - assert len(vec) == 64, f"Expected 64-dim vector, got {len(vec)}" - assert all(isinstance(v, float) for v in vec), "All values should be floats" - - -def test_code_vs_nl_detection(): - """Test auto-detection of code vs natural language queries.""" - from scripts.mcp_impl.pattern_search import _detect_query_mode - - # Code examples (no language hint, should detect from syntax) - assert _detect_query_mode("for i in range(3): try: pass except: pass", None) == "code" - assert _detect_query_mode("if err != nil { return err }", None) == "code" - assert _detect_query_mode("def foo(): return 42", None) == "code" - assert _detect_query_mode("func main() {}", None) == "code" - assert _detect_query_mode("fn main() -> Result<()>", None) == "code" - - # Natural language descriptions - assert _detect_query_mode("retry with exponential backoff", None) == "description" - assert _detect_query_mode("find error handling patterns", None) == "description" - assert _detect_query_mode("resource cleanup code", None) == "description" - assert _detect_query_mode("decorator pattern wrapping function", None) == "description" - - # Ambiguous two-word input is treated as code when a language hint is present - assert _detect_query_mode("some text", "python") == "code" - assert _detect_query_mode("some text", "go") == "code" - - -# ============================================================================ -# Integration tests (require Qdrant service) -# ============================================================================ - -@pytest.fixture -def pattern_collection(): - """Create collection with pattern vectors and test data.""" - from qdrant_client import QdrantClient - from qdrant_client.models import Distance, VectorParams, PointStruct - from scripts.pattern_detection import PatternExtractor, PatternEncoder - - # Default to localhost for CI; allow override for local/dev containers. - collection_name = f"test_pattern_{uuid.uuid4().hex[:8]}" - qdrant_url = os.environ.get("QDRANT_URL", "http://localhost:6333") - client = QdrantClient(url=qdrant_url, timeout=30) - - try: - client.create_collection( - collection_name=collection_name, - vectors_config={ - "code": VectorParams(size=384, distance=Distance.COSINE), - "pattern_vector": VectorParams(size=64, distance=Distance.COSINE), - } - ) - except Exception as e: - pytest.skip(f"Qdrant not reachable at {qdrant_url}: {e}") - - extractor = PatternExtractor() - encoder = PatternEncoder() - - snippets = [ - {"path": "retry_python.py", "code": RETRY_PATTERN_PYTHON, "lang": "python"}, - {"path": "retry_go.go", "code": RETRY_PATTERN_GO, "lang": "go"}, - {"path": "file_read.py", "code": FILE_READ_PATTERN, "lang": "python"}, - ] - - points = [] - for i, s in enumerate(snippets): - sig = extractor.extract(s["code"], s["lang"]) - points.append(PointStruct( - id=i + 1, - vector={"code": [0.1] * 384, "pattern_vector": encoder.encode(sig)}, - payload={"path": s["path"], "language": s["lang"], "text": s["code"]}, - )) - - client.upsert(collection_name=collection_name, points=points) - - yield collection_name, client - - try: - client.delete_collection(collection_name) - except Exception: - pass - - -@pytest.mark.service -def test_pattern_search_qdrant(pattern_collection): - """Test pattern search against real Qdrant.""" - from scripts.pattern_detection.search import pattern_search - - collection_name, client = pattern_collection - - results = pattern_search( - example=RETRY_PATTERN_PYTHON, - language="python", - collection=collection_name, - limit=5, - client=client, # Pass client to avoid global cache pollution - ) - - assert results.total >= 1, "Should find at least one result" - paths = [r.path for r in results.results] - assert "retry_python.py" in paths, "Should find Python retry" diff --git a/tests/test_per_path_zero.py b/tests/test_per_path_zero.py index e1541ffd..2e6a164d 100644 --- a/tests/test_per_path_zero.py +++ b/tests/test_per_path_zero.py @@ -1,24 +1,53 @@ import asyncio +import json import pytest # These tests exercise argument plumbing independent of live retrieval. @pytest.mark.asyncio -async def test_per_path_zero_is_echoed_and_respected_in_args(): - from scripts.mcp_indexer_server import repo_search +async def test_per_path_zero_is_echoed_and_respected_in_args(monkeypatch): + from scripts.mcp_impl.search import _repo_search_impl - res = await repo_search(query="anything", limit=3, per_path=0) + async def _fake_run_async(_cmd, **_kwargs): + item = {"path": "src/a.py", "start_line": 1, "end_line": 1, "score": 1.0} + return {"ok": True, "code": 0, "stdout": json.dumps(item), "stderr": ""} + + monkeypatch.setenv("HYBRID_IN_PROCESS", "0") + + # Arg-plumbing test for the hybrid/subprocess (non-dense) path; mode is explicit by design. + res = await _repo_search_impl( + query="anything", + limit=3, + per_path=0, + mode="hybrid", + require_auth_session_fn=lambda session: session, + run_async_fn=_fake_run_async, + ) assert isinstance(res, dict) args = res.get("args") or {} assert args.get("per_path") == 0, f"expected per_path echoed as 0, got {args.get('per_path')}" @pytest.mark.asyncio -async def test_compact_string_false_is_normalized_in_args(): - from scripts.mcp_indexer_server import repo_search +async def test_compact_string_false_is_normalized_in_args(monkeypatch): + from scripts.mcp_impl.search import _repo_search_impl + + async def _fake_run_async(_cmd, **_kwargs): + item = {"path": "src/a.py", "start_line": 1, "end_line": 1, "score": 1.0} + return {"ok": True, "code": 0, "stdout": json.dumps(item), "stderr": ""} + + monkeypatch.setenv("HYBRID_IN_PROCESS", "0") - # Passing compact as a string "false" should normalize to False in echoed args - res = await repo_search(query="anything", limit=1, compact="false") + # Passing compact as a string "false" should normalize to False in echoed args. + # Keep mode explicit so dense-default env does not alter this contract test. + res = await _repo_search_impl( + query="anything", + limit=1, + compact="false", + mode="hybrid", + require_auth_session_fn=lambda session: session, + run_async_fn=_fake_run_async, + ) assert isinstance(res, dict) args = res.get("args") or {} assert args.get("compact") is False, f"expected compact False, got {args.get('compact')}" diff --git a/tests/test_prune.py b/tests/test_prune.py new file mode 100644 index 00000000..42641883 --- /dev/null +++ b/tests/test_prune.py @@ -0,0 +1,54 @@ +from types import SimpleNamespace + +import scripts.prune as prune + + +class _FakeClient: + def __init__(self, points): + self._points = points + + def scroll(self, **kwargs): + return self._points, None + + +def _point(path, file_hash=None, repo="repo-a"): + return SimpleNamespace( + payload={ + "metadata": { + "path": path, + "file_hash": file_hash, + "repo": repo, + } + } + ) + + +def test_prune_excludes_deleted_paths_from_orphan_keepalive(monkeypatch, tmp_path): + keep_path = tmp_path / "keep.py" + keep_path.write_text("keep = True\n", encoding="utf-8") + + mismatch_path = tmp_path / "mismatch.py" + mismatch_path.write_text("new = True\n", encoding="utf-8") + + points = [ + _point("missing.py", file_hash="missing-hash"), + _point("mismatch.py", file_hash="old-hash"), + _point("keep.py", file_hash=prune.sha1_file(keep_path)), + ] + fake_client = _FakeClient(points) + captured_valid_paths = [] + + monkeypatch.setattr(prune, "QdrantClient", lambda **kwargs: fake_client) + monkeypatch.setattr(prune, "ROOT", tmp_path) + monkeypatch.setattr(prune, "delete_by_path", lambda *args, **kwargs: 1) + monkeypatch.setattr(prune, "delete_graph_edges_by_path", lambda *args, **kwargs: 0) + monkeypatch.setattr( + prune, + "delete_orphan_graph_edges", + lambda client, valid_paths: captured_valid_paths.append(set(valid_paths)) or 0, + ) + + prune.main() + + assert captured_valid_paths == [{"keep.py"}] + diff --git a/tests/test_qdrant_version_pins.py b/tests/test_qdrant_version_pins.py new file mode 100644 index 00000000..4b08d966 --- /dev/null +++ b/tests/test_qdrant_version_pins.py @@ -0,0 +1,39 @@ +from pathlib import Path +from importlib.metadata import version +import inspect + +from qdrant_client import QdrantClient + + +ROOT = Path(__file__).resolve().parents[1] +QDRANT_CLIENT_PIN = "qdrant-client==1.15.1" +QDRANT_SERVER_IMAGE = "qdrant/qdrant:v1.15.4" + + +def test_qdrant_client_is_exactly_pinned(): + requirements = (ROOT / "requirements.txt").read_text() + + assert QDRANT_CLIENT_PIN in requirements + assert "qdrant-client>=" not in requirements + + +def test_qdrant_server_images_are_exactly_pinned(): + files = [ + ROOT / ".github/workflows/ci.yml", + ROOT / "docker-compose.yml", + ROOT / "docker-compose-bindmount-checkout.yml", + ROOT / "deploy/kubernetes/qdrant.yaml", + ROOT / "tests/conftest.py", + ] + + for path in files: + text = path.read_text() + assert QDRANT_SERVER_IMAGE in text, str(path) + assert "qdrant/qdrant:latest" not in text, str(path) + + +def test_installed_qdrant_client_matches_supported_api(): + assert version("qdrant-client") == "1.15.1" + assert hasattr(QdrantClient, "search") + assert hasattr(QdrantClient, "query_points") + assert "query_filter" in inspect.signature(QdrantClient.query_points).parameters diff --git a/tests/test_relevance_feedback.py b/tests/test_relevance_feedback.py new file mode 100644 index 00000000..c76f9e7c --- /dev/null +++ b/tests/test_relevance_feedback.py @@ -0,0 +1,190 @@ +"""Integration test for the relevance feedback pipeline. + +Validates the full end-to-end flow: +1. Search produces result_id on every result +2. rate_search_results logs feedback events +3. relevance_trainer aggregates events into weight files +4. Subsequent searches apply learned relevance boosts + +Uses Qdrant via CI service, explicit QDRANT_URL, or local testcontainers fallback. +Requires --run-integration flag. +""" + +import json +import os +import asyncio +import importlib +import pytest + +pytestmark = pytest.mark.integration + + +class FakeEmbedder: + def __init__(self, model_name: str = "fake"): + self.model_name = model_name + + class _Vec: + def __init__(self, arr): + self._arr = arr + + def tolist(self): + return self._arr + + def __len__(self): + return len(self._arr) + + def embed(self, texts): + for t in texts: + h = sum(ord(c) for c in t) % 997 + vec = [(float((h + i) % 13) / 13.0) for i in range(32)] + yield self._Vec(vec) + + +def _load_modules(): + return { + "ing": importlib.import_module("scripts.ingest_code"), + "srv": importlib.import_module("scripts.mcp_indexer_server"), + "embedder": importlib.import_module("scripts.embedder"), + "hy": importlib.import_module("scripts.hybrid_search"), + "pipeline": importlib.import_module("scripts.ingest.pipeline"), + "rt": importlib.import_module("scripts.relevance_trainer"), + } + + +def _stub_embeddings(monkeypatch, modules): + """Stub all embedding paths with FakeEmbedder (32-dim), avoiding real model loads.""" + from qdrant_client import QdrantClient as _RealQdrantClient + + ing = modules["ing"] + srv = modules["srv"] + embedder = modules["embedder"] + hy = modules["hy"] + pipeline = modules["pipeline"] + + monkeypatch.setattr(ing, "TextEmbedding", lambda *a, **k: FakeEmbedder("fake")) + monkeypatch.setattr(embedder, "get_embedding_model", lambda *a, **k: FakeEmbedder("fake")) + monkeypatch.setattr(embedder, "get_model_dimension", lambda *a, **k: 32) + monkeypatch.setattr(srv, "_get_embedding_model", lambda *a, **k: FakeEmbedder("fake")) + monkeypatch.setattr(hy, "TextEmbedding", lambda *a, **k: FakeEmbedder("fake")) + monkeypatch.setattr(hy, "_get_embedding_model", lambda *a, **k: FakeEmbedder("fake")) + # Override TYPE_CHECKING shim; pipeline/hybrid use QdrantClient=Any at runtime. + monkeypatch.setattr(pipeline, "QdrantClient", _RealQdrantClient) + monkeypatch.setattr(hy, "QdrantClient", _RealQdrantClient) + + +def _make_tiny_repo(tmp_path): + """Create a small multi-file repo for search testing.""" + (tmp_path / "pkg").mkdir() + (tmp_path / "pkg" / "auth.py").write_text("def authenticate(user):\n return user == 'admin'\n") + (tmp_path / "pkg" / "utils.py").write_text("def log(msg):\n print(msg)\n") + (tmp_path / "pkg" / "README.md").write_text("# Test Project\nThis is a test project.\n") + return tmp_path + + +@pytest.mark.integration +def test_relevance_feedback_pipeline(tmp_path, monkeypatch, qdrant_url, test_collection): + """One indexed repo validates ids, hands-off rating, training, and boosts.""" + modules = _load_modules() + ing = modules["ing"] + srv = modules["srv"] + rt = modules["rt"] + + events_dir = tmp_path / "rerank_events" + weights_dir = tmp_path / "rerank_weights" + events_dir.mkdir() + weights_dir.mkdir() + + coll = test_collection + monkeypatch.setenv("QDRANT_URL", qdrant_url) + monkeypatch.setenv("COLLECTION_NAME", coll) + monkeypatch.setenv("USE_TREE_SITTER", "0") + monkeypatch.setenv("HYBRID_IN_PROCESS", "1") + monkeypatch.setenv("HYBRID_EXPAND", "0") + monkeypatch.setenv("SEMANTIC_EXPANSION_ENABLED", "0") + monkeypatch.setenv("RERANKER_ENABLED", "0") + monkeypatch.setenv("EMBEDDING_MODEL", "fake") + monkeypatch.setenv("RERANK_EVENTS_DIR", str(events_dir)) + monkeypatch.setenv("RERANKER_WEIGHTS_DIR", str(weights_dir)) + monkeypatch.setenv("RELEVANCE_TRAINER_MIN_EVENTS", "1") + monkeypatch.setenv("RELEVANCE_BOOST_FACTOR", "0.5") + _stub_embeddings(monkeypatch, modules) + + _make_tiny_repo(tmp_path) + ing.index_repo( + root=tmp_path, qdrant_url=qdrant_url, api_key="", + collection=coll, model_name="fake", recreate=True, + ) + + res = asyncio.run( + srv.repo_search(queries=["authenticate"], limit=5, compact=False, rerank_enabled=False) + ) + results = res.get("results", []) + assert len(results) >= 1 + + for r in results: + rid = r.get("result_id", "") + assert len(rid) == 12, f"result_id should be 12 hex chars, got: {rid!r}" + assert all(c in "0123456789abcdef" for c in rid), f"result_id not hex: {rid!r}" + assert r.get("target_id") == rid + assert len(r.get("impression_id", "")) == 12 + + res_compact = asyncio.run( + srv.repo_search(queries=["authenticate"], limit=5, compact=True, rerank_enabled=False) + ) + compact_results = res_compact.get("results", []) + assert compact_results + assert all(len(r.get("result_id", "")) == 12 for r in compact_results) + + rated_rid = results[0]["result_id"] + + rating_res = asyncio.run( + srv.rate_search_results( + query="authenticate", + ratings=[{"result_id": rated_rid, "relevance": 2}], + collection=coll, + ) + ) + assert rating_res.get("ok"), f"rate_search_results failed: {rating_res}" + assert rating_res.get("rated") == 1 + + event_files = list(events_dir.glob(f"events_{coll}_*.ndjson")) + assert len(event_files) > 0, "No event files written" + + with open(event_files[0], "r") as f: + event = json.loads(f.readline()) + assert event["type"] == "relevance_feedback" + assert event["query"] == "authenticate" + assert event["collection"] == coll + assert event["source"] == "mcp_tool" + assert event["ratings"][0]["result_id"] == rated_rid + assert event["ratings"][0]["target_id"] == rated_rid + assert event["ratings"][0]["path"] + + trainer_res = rt.process_collection(coll) + assert trainer_res["collection"] == coll + assert trainer_res["events"] >= 1 + assert not trainer_res.get("skipped"), f"Trainer skipped: {trainer_res.get('reason')}" + assert trainer_res["new_entries"] >= 1 + + weight_file = weights_dir / f"{coll}_relevance.json" + assert weight_file.exists(), f"Weight file not found at {weight_file}" + + weights = json.loads(weight_file.read_text()) + assert "results" in weights + assert rated_rid in weights["results"], f"Rated result_id {rated_rid} not in weights" + assert weights["results"][rated_rid]["avg_relevance"] == 2.0 + assert weights["results"][rated_rid]["target"]["path"] + + boosted = asyncio.run( + srv.repo_search( + queries=["authenticate"], + limit=5, + compact=False, + debug=True, + rerank_enabled=False, + ) + ) + boosted_results = boosted.get("results", []) + rated_result = next((r for r in boosted_results if r.get("result_id") == rated_rid), None) + assert rated_result is not None, "Rated result disappeared from results" + assert rated_result.get("relevance_boost", 0) > 0 diff --git a/tests/test_relevance_identity.py b/tests/test_relevance_identity.py new file mode 100644 index 00000000..631d72f3 --- /dev/null +++ b/tests/test_relevance_identity.py @@ -0,0 +1,438 @@ +import sys +import types +import json + +from scripts.mcp_impl.search import ( + _RECENT_RESULT_META, + _feedback_recall_candidates, + _inject_result_ids, + _remember_result_metadata, + enrich_feedback_rating, +) +from scripts.relevance_feedback import ( + build_symbol_reconciliations, + reconcile_collection_weights, +) +from scripts.relevance_trainer import aggregate_ratings, process_collection + + +def _result(file_hash: str) -> dict: + return { + "path": "/repo/pkg/auth.py", + "container_path": "/work/repo/pkg/auth.py", + "repo": "repo", + "kind": "function", + "symbol": "authenticate", + "start_line": 10, + "end_line": 20, + "file_hash": file_hash, + } + + +def test_result_id_survives_content_hash_changes(): + before = _result("hash-before") + after = _result("hash-after") + + _inject_result_ids([before], "authenticate user") + _inject_result_ids([after], "authenticate user") + + assert before["result_id"] == after["result_id"] + assert before["target_id"] == after["target_id"] + assert before["impression_id"] != after["impression_id"] + + +def test_result_id_is_stable_for_same_query_location_and_content(): + first = _result("same-hash") + second = _result("same-hash") + + _inject_result_ids([first], "authenticate user") + _inject_result_ids([second], "authenticate user") + + assert first["result_id"] == second["result_id"] + assert first["impression_id"] == second["impression_id"] + assert len(first["result_id"]) == 12 + + +def test_result_id_survives_line_shifts_for_same_symbol(): + before = _result("same-hash") + after = _result("same-hash") + after["start_line"] = 50 + after["end_line"] = 80 + + _inject_result_ids([before], "authenticate user") + _inject_result_ids([after], "authenticate user") + + assert before["result_id"] == after["result_id"] + assert before["impression_id"] != after["impression_id"] + + +def test_same_named_symbols_in_distinct_files_have_distinct_targets(): + first = _result("same-hash") + second = _result("same-hash") + second["path"] = "/repo/pkg/other.py" + second["container_path"] = "/work/repo/pkg/other.py" + + _inject_result_ids([first, second], "authenticate user") + + assert first["target_id"] != second["target_id"] + assert first["result_id"] != second["result_id"] + + +def test_target_identity_normalizes_host_and_container_paths(): + from scripts.relevance_feedback import stable_target_id + + assert stable_target_id( + repo="repo", + kind="function", + symbol="authenticate", + path="/repo/pkg/auth.py", + ) == stable_target_id( + repo="repo", + kind="function", + symbol="authenticate", + path="/work/repo/pkg/auth.py", + ) + + +def test_trainer_preserves_target_metadata_for_recall(): + result = _result("same-hash") + _inject_result_ids([result], "authenticate user") + + weights = aggregate_ratings([ + { + "type": "relevance_feedback", + "ratings": [ + { + "result_id": result["result_id"], + "relevance": 2, + "target_id": result["target_id"], + "impression_id": result["impression_id"], + "path": result["path"], + "container_path": result["container_path"], + "symbol": result["symbol"], + "kind": result["kind"], + "repo": result["repo"], + "file_hash": result["file_hash"], + } + ], + } + ]) + + entry = weights[result["result_id"]] + assert entry["avg_relevance"] == 2 + assert entry["target"]["symbol"] == "authenticate" + assert entry["target"]["container_path"] == "/work/repo/pkg/auth.py" + + +def test_feedback_rating_enriches_from_recent_search_result(): + result = _result("same-hash") + _inject_result_ids([result], "authenticate user") + _remember_result_metadata([result]) + + rating = enrich_feedback_rating({ + "result_id": result["result_id"], + "relevance": 2, + }) + + assert rating["target_id"] == result["target_id"] + assert rating["container_path"] == "/work/repo/pkg/auth.py" + assert rating["symbol"] == "authenticate" + + +def test_feedback_rating_enriches_from_shared_storage(tmp_path, monkeypatch): + monkeypatch.setenv("RERANKER_WEIGHTS_DIR", str(tmp_path)) + result = _result("same-hash") + result["symbol_content_hash"] = "symbol-hash" + _inject_result_ids([result], "authenticate user") + _remember_result_metadata([result], "repo-collection") + + _RECENT_RESULT_META.clear() + rating = enrich_feedback_rating( + {"result_id": result["result_id"], "relevance": 2}, + "repo-collection", + ) + + assert rating["symbol"] == "authenticate" + assert rating["symbol_content_hash"] == "symbol-hash" + + +def test_exact_content_rename_reconciles_feedback_weight(tmp_path, monkeypatch): + monkeypatch.setenv("RERANKER_WEIGHTS_DIR", str(tmp_path)) + old = { + "function_authenticate_10": { + "name": "authenticate", + "type": "function", + "content_hash": "same-body", + "content": "def authenticate(user): return user", + } + } + new = { + "function_verify_user_20": { + "name": "verify_user", + "type": "function", + "content_hash": "same-body", + "content": "def authenticate(user): return user", + } + } + before = _result("file-hash") + _inject_result_ids([before], "authenticate") + weights_file = tmp_path / "repo-collection_relevance.json" + weights_file.write_text(json.dumps({ + "results": { + before["result_id"]: { + "total_relevance": 2, + "count": 1, + "avg_relevance": 2.0, + "target": { + "target_id": before["result_id"], + "repo": "repo", + "kind": "function", + "symbol": "authenticate", + "path": "/repo/pkg/auth.py", + }, + } + } + })) + + mappings = build_symbol_reconciliations( + old, + new, + repo="repo", + path="/repo/pkg/auth.py", + ) + assert reconcile_collection_weights("repo-collection", mappings) == 1 + + data = json.loads(weights_file.read_text()) + successor = next( + entry + for rid, entry in data["results"].items() + if rid != before["result_id"] and entry.get("target", {}).get("symbol") == "verify_user" + ) + assert successor["inheritance_weight"] == 1.0 + assert successor["lineage"][0]["reason"] == "rename_exact_content" + + +def test_qualified_symbol_rename_reconciles_feedback_weight(tmp_path, monkeypatch): + monkeypatch.setenv("RERANKER_WEIGHTS_DIR", str(tmp_path)) + old = { + "function_authenticate_10": { + "name": "authenticate", + "path": "AuthService.authenticate", + "type": "function", + "content_hash": "same-body", + } + } + new = { + "function_verify_user_20": { + "name": "verify_user", + "path": "AuthService.verify_user", + "type": "function", + "content_hash": "same-body", + } + } + before = _result("file-hash") + before["symbol"] = "AuthService.authenticate" + _inject_result_ids([before], "authenticate") + weights_file = tmp_path / "repo-collection_relevance.json" + weights_file.write_text(json.dumps({ + "results": { + before["result_id"]: { + "total_relevance": 2, + "count": 1, + "avg_relevance": 2.0, + "target": { + "target_id": before["result_id"], + "repo": "repo", + "kind": "function", + "symbol": "AuthService.authenticate", + "path": "/repo/pkg/auth.py", + }, + } + } + })) + + mappings = build_symbol_reconciliations( + old, + new, + repo="repo", + path="/repo/pkg/auth.py", + ) + assert reconcile_collection_weights("repo-collection", mappings) == 1 + + data = json.loads(weights_file.read_text()) + successor = next( + entry + for rid, entry in data["results"].items() + if rid != before["result_id"] + ) + assert successor["target"]["symbol"] == "AuthService.verify_user" + assert successor["lineage"][0]["reason"] == "rename_exact_content" + + +def test_symbol_split_divides_inherited_feedback(tmp_path, monkeypatch): + monkeypatch.setenv("RERANKER_WEIGHTS_DIR", str(tmp_path)) + old = { + "function_process_1": { + "name": "process", + "type": "function", + "content_hash": "old", + "content": "def process order validate payment persist receipt notify customer", + } + } + new = { + "function_validate_order_1": { + "name": "validate_order", + "type": "function", + "content_hash": "new-a", + "content": "def validate_order order validate payment customer", + }, + "function_persist_receipt_10": { + "name": "persist_receipt", + "type": "function", + "content_hash": "new-b", + "content": "def persist_receipt persist receipt notify customer", + }, + } + mappings = build_symbol_reconciliations( + old, + new, + repo="repo", + path="/repo/pkg/orders.py", + split_min_overlap=0.35, + split_min_coverage=0.7, + ) + + successors = next(iter(mappings.values())) + assert {item["target"]["symbol"] for item in successors} == { + "validate_order", + "persist_receipt", + } + assert round(sum(item["inheritance_weight"] for item in successors), 6) == 1.0 + assert all(item["reason"] == "split_token_coverage" for item in successors) + + +def test_trainer_preserves_existing_reconciled_entries(tmp_path, monkeypatch): + events_dir = tmp_path / "events" + weights_dir = tmp_path / "weights" + events_dir.mkdir() + weights_dir.mkdir() + monkeypatch.setenv("RERANK_EVENTS_DIR", str(events_dir)) + monkeypatch.setenv("RERANKER_WEIGHTS_DIR", str(weights_dir)) + monkeypatch.setenv("RELEVANCE_TRAINER_MIN_EVENTS", "1") + (weights_dir / "repo_relevance.json").write_text(json.dumps({ + "results": { + "inherited": { + "total_relevance": 2, + "count": 1, + "avg_relevance": 2.0, + "inheritance_weight": 0.5, + "target": {"symbol": "split_child"}, + } + } + })) + (events_dir / "events_repo_2026060902.ndjson").write_text(json.dumps({ + "type": "relevance_feedback", + "ratings": [{"result_id": "fresh", "relevance": 2}], + }) + "\n") + + process_collection("repo") + data = json.loads((weights_dir / "repo_relevance.json").read_text()) + assert "inherited" in data["results"] + assert "fresh" in data["results"] + + +def test_graph_recall_adds_callers_when_rated_target_already_present(monkeypatch): + class FakeMatchValue: + def __init__(self, value): + self.value = value + + class FakeFieldCondition: + def __init__(self, key, match): + self.key = key + self.match = match + + class FakeFilter: + def __init__(self, must=None): + self.must = must or [] + + class FakeModels: + MatchValue = FakeMatchValue + FieldCondition = FakeFieldCondition + Filter = FakeFilter + + class FakePoint: + def __init__(self, payload): + self.payload = payload + + class FakeClient: + def __init__(self, *args, **kwargs): + pass + + def scroll(self, collection_name, scroll_filter, **kwargs): + terms = {cond.key: cond.match.value for cond in scroll_filter.must} + if collection_name.endswith("_graph"): + if terms.get("edge_type") == "calls" and terms.get("callee_symbol") == "foo": + return [FakePoint({ + "caller_path": "/work/repo/pkg/caller.py", + "repo": "repo", + })], None + return [], None + + if terms.get("metadata.path") == "/work/repo/pkg/caller.py": + return [FakePoint({ + "metadata": { + "path": "/work/repo/pkg/caller.py", + "host_path": "/repo/pkg/caller.py", + "container_path": "/work/repo/pkg/caller.py", + "repo": "repo", + "kind": "function", + "symbol": "caller", + "symbol_path": "caller", + "start_line": 1, + "end_line": 5, + } + })], None + + return [FakePoint({ + "metadata": { + "path": "/work/repo/pkg/foo.py", + "repo": "repo", + "kind": "function", + "symbol": "foo", + "symbol_path": "foo", + "start_line": 1, + "end_line": 5, + } + })], None + + fake_qdrant = types.SimpleNamespace(QdrantClient=FakeClient, models=FakeModels) + monkeypatch.setitem(sys.modules, "qdrant_client", fake_qdrant) + monkeypatch.setenv("RELEVANCE_GRAPH_RECALL_MAX", "2") + + weights = { + "results": { + "rated-target": { + "avg_relevance": 2, + "count": 1, + "target": { + "repo": "repo", + "kind": "function", + "symbol": "foo", + "container_path": "/work/repo/pkg/foo.py", + }, + } + } + } + + recalled = _feedback_recall_candidates( + collection="repo", + weights=weights, + existing_target_ids={"rated-target"}, + existing_paths={"/repo/pkg/foo.py"}, + base_score=0.5, + max_candidates=3, + ) + + assert len(recalled) == 1 + assert recalled[0]["feedback_graph_recall"] is True + assert recalled[0]["path"] == "/repo/pkg/caller.py" diff --git a/tests/test_repo_search_mode_contract.py b/tests/test_repo_search_mode_contract.py new file mode 100644 index 00000000..b7152696 --- /dev/null +++ b/tests/test_repo_search_mode_contract.py @@ -0,0 +1,186 @@ +import asyncio +import importlib +import sys +import types + +import pytest + +srv = importlib.import_module("scripts.mcp_indexer_server") + + +def _make_hybrid_module_stub(calls: dict): + mod = types.ModuleType("scripts.hybrid_search") + + def run_pure_dense_search(**kwargs): + calls["dense"] = int(calls.get("dense", 0)) + 1 + calls["dense_kwargs"] = dict(kwargs) + return [ + { + "score": 0.91, + "path": "/work/dense.py", + "symbol": "", + "start_line": 1, + "end_line": 3, + "payload": {}, + } + ] + + def run_hybrid_search(**kwargs): + calls["hybrid"] = int(calls.get("hybrid", 0)) + 1 + calls["hybrid_kwargs"] = dict(kwargs) + return [ + { + "score": 0.75, + "path": "/work/hybrid.py", + "symbol": "", + "start_line": 4, + "end_line": 7, + } + ] + + mod.run_pure_dense_search = run_pure_dense_search + mod.run_hybrid_search = run_hybrid_search + mod.lang_matches_path = lambda path, lang=None: True + mod._merge_and_budget_spans = lambda spans, *args, **kwargs: spans + mod.TextEmbedding = object + mod.QdrantClient = object + return mod + + +@pytest.mark.service +def test_repo_search_dense_default_from_env_is_explicit_and_stable(monkeypatch): + # Contract: global default mode should route repo_search to dense path when set to dense. + calls = {"dense": 0, "hybrid": 0} + monkeypatch.setenv("REPO_SEARCH_DEFAULT_MODE", "dense") + monkeypatch.setenv("HYBRID_IN_PROCESS", "1") + monkeypatch.setattr(srv, "_get_embedding_model", lambda *a, **k: object()) + monkeypatch.setitem(sys.modules, "scripts.hybrid_search", _make_hybrid_module_stub(calls)) + + res = asyncio.run(srv.repo_search(query="q", limit=1, compact=True, rerank_enabled=False)) + + assert res.get("ok") is True + assert calls["dense"] == 1 + assert calls["hybrid"] == 0 + assert res.get("results", [{}])[0].get("path") == "/work/dense.py" + + +@pytest.mark.service +def test_repo_search_explicit_hybrid_overrides_dense_default_for_non_dense_tests(monkeypatch): + # Contract: non-dense tests can force hybrid behavior even under dense global default. + calls = {"dense": 0, "hybrid": 0} + monkeypatch.setenv("REPO_SEARCH_DEFAULT_MODE", "dense") + monkeypatch.setenv("HYBRID_IN_PROCESS", "1") + monkeypatch.setattr(srv, "_get_embedding_model", lambda *a, **k: object()) + monkeypatch.setitem(sys.modules, "scripts.hybrid_search", _make_hybrid_module_stub(calls)) + + res = asyncio.run( + srv.repo_search( + query="q", + mode="hybrid", + limit=1, + compact=True, + rerank_enabled=False, + ) + ) + + assert res.get("ok") is True + assert calls["dense"] == 0 + assert calls["hybrid"] == 1 + assert res.get("results", [{}])[0].get("path") == "/work/hybrid.py" + + +@pytest.mark.service +def test_repo_search_dense_default_forwards_structured_filters(monkeypatch): + calls = {"dense": 0, "hybrid": 0} + monkeypatch.setenv("REPO_SEARCH_DEFAULT_MODE", "dense") + monkeypatch.setenv("HYBRID_IN_PROCESS", "1") + monkeypatch.setattr(srv, "_get_embedding_model", lambda *a, **k: object()) + monkeypatch.setitem(sys.modules, "scripts.hybrid_search", _make_hybrid_module_stub(calls)) + + asyncio.run( + srv.repo_search( + query="q", + limit=1, + compact=True, + rerank_enabled=False, + kind="function", + symbol="my_symbol", + ext="py", + ) + ) + + assert calls["dense"] == 1 + dense_kwargs = calls.get("dense_kwargs") or {} + assert dense_kwargs.get("kind") == "function" + assert dense_kwargs.get("symbol") == "my_symbol" + assert dense_kwargs.get("ext") == "py" + + +@pytest.mark.service +def test_repo_search_dense_default_forwards_per_path(monkeypatch): + calls = {"dense": 0, "hybrid": 0} + monkeypatch.setenv("REPO_SEARCH_DEFAULT_MODE", "dense") + monkeypatch.setenv("HYBRID_IN_PROCESS", "1") + monkeypatch.setattr(srv, "_get_embedding_model", lambda *a, **k: object()) + monkeypatch.setitem(sys.modules, "scripts.hybrid_search", _make_hybrid_module_stub(calls)) + + asyncio.run( + srv.repo_search( + query="q", + limit=3, + per_path=1, + compact=True, + rerank_enabled=False, + ) + ) + + assert calls["dense"] == 1 + assert calls.get("dense_kwargs", {}).get("per_path") == 1 + + +@pytest.mark.service +def test_repo_search_profile_tests_adds_material_globs(monkeypatch): + calls = {"dense": 0, "hybrid": 0} + monkeypatch.setenv("REPO_SEARCH_DEFAULT_MODE", "dense") + monkeypatch.setenv("HYBRID_IN_PROCESS", "1") + monkeypatch.setattr(srv, "_get_embedding_model", lambda *a, **k: object()) + monkeypatch.setitem(sys.modules, "scripts.hybrid_search", _make_hybrid_module_stub(calls)) + + res = asyncio.run( + srv.repo_search( + query="q", + profile="tests", + limit=1, + compact=False, + rerank_enabled=False, + ) + ) + + args = res.get("args") or {} + assert args.get("profile") == "tests" + assert "tests/**" in args.get("path_glob", []) + assert "**/*_test.*" in args.get("path_glob", []) + + +@pytest.mark.service +def test_repo_search_profile_preserves_user_globs(monkeypatch): + calls = {"dense": 0, "hybrid": 0} + monkeypatch.setenv("REPO_SEARCH_DEFAULT_MODE", "dense") + monkeypatch.setenv("HYBRID_IN_PROCESS", "1") + monkeypatch.setattr(srv, "_get_embedding_model", lambda *a, **k: object()) + monkeypatch.setitem(sys.modules, "scripts.hybrid_search", _make_hybrid_module_stub(calls)) + + res = asyncio.run( + srv.repo_search( + query="q", + profile="config", + path_glob=["custom/**"], + limit=1, + compact=False, + rerank_enabled=False, + ) + ) + + globs = (res.get("args") or {}).get("path_glob", []) + assert globs[0] == "custom/**" + assert "**/*.yaml" in globs diff --git a/tests/test_rerank_recursive.py b/tests/test_rerank_recursive.py deleted file mode 100644 index 5f8c5822..00000000 --- a/tests/test_rerank_recursive.py +++ /dev/null @@ -1,371 +0,0 @@ -""" -Tests for the Tiny Recursive Reranker (TRM-inspired). - -Validates: -1. Basic reranking functionality -2. Iterative refinement improves results -3. Early stopping works correctly -4. Latent state carryover functions -5. Integration with existing candidates -""" - -import pytest -import numpy as np -from typing import List, Dict, Any - - -# Import the reranker -from scripts.rerank_recursive import ( - RecursiveReranker, - RefinementState, - TinyScorer, - LatentRefiner, - ConfidenceEstimator, - rerank_recursive, - rerank_recursive_inprocess, -) - - -class TestTinyScorer: - """Tests for the tiny scoring network.""" - - def test_forward_shape(self): - """Scorer should produce correct output shape.""" - scorer = TinyScorer(dim=64, hidden_dim=128) - - query_emb = np.random.randn(64).astype(np.float32) - doc_embs = np.random.randn(5, 64).astype(np.float32) - z = np.random.randn(64).astype(np.float32) - - scores = scorer.forward(query_emb, doc_embs, z) - - assert scores.shape == (5,) - assert scores.dtype == np.float32 - - def test_forward_deterministic(self): - """Same inputs should produce same outputs.""" - scorer = TinyScorer(dim=64) - - query_emb = np.random.randn(64).astype(np.float32) - doc_embs = np.random.randn(3, 64).astype(np.float32) - z = np.random.randn(64).astype(np.float32) - - scores1 = scorer.forward(query_emb, doc_embs, z) - scores2 = scorer.forward(query_emb, doc_embs, z) - - np.testing.assert_array_almost_equal(scores1, scores2) - - -class TestLatentRefiner: - """Tests for latent state refinement.""" - - def test_refine_shape(self): - """Refiner should produce latent of same dimension.""" - refiner = LatentRefiner(dim=64) - - z = np.random.randn(64).astype(np.float32) - query_emb = np.random.randn(64).astype(np.float32) - doc_embs = np.random.randn(5, 64).astype(np.float32) - scores = np.random.randn(5).astype(np.float32) - - z_refined = refiner.refine(z, query_emb, doc_embs, scores) - - assert z_refined.shape == (64,) - - def test_refine_normalized(self): - """Refined latent should be unit normalized.""" - refiner = LatentRefiner(dim=64) - - z = np.random.randn(64).astype(np.float32) - query_emb = np.random.randn(64).astype(np.float32) - doc_embs = np.random.randn(5, 64).astype(np.float32) - scores = np.random.randn(5).astype(np.float32) - - z_refined = refiner.refine(z, query_emb, doc_embs, scores) - - norm = np.linalg.norm(z_refined) - assert abs(norm - 1.0) < 1e-5 - - -class TestConfidenceEstimator: - """Tests for early stopping logic.""" - - def test_no_stop_on_first_iteration(self): - """Should not stop on first iteration.""" - estimator = ConfidenceEstimator() - - state = RefinementState( - z=np.zeros(64), - scores=np.array([0.5, 0.3, 0.1]), - iteration=0 - ) - state.score_history = [state.scores] - - assert not estimator.should_stop(state) - - def test_stop_on_convergence(self): - """Should stop when top-k rankings stabilize.""" - estimator = ConfidenceEstimator() - - state = RefinementState( - z=np.zeros(64), - scores=np.array([0.5, 0.3, 0.1]), - iteration=2 - ) - # Same scores twice = converged - state.score_history = [ - np.array([0.5, 0.3, 0.1]), - np.array([0.5, 0.3, 0.1]) - ] - - assert estimator.should_stop(state) - - def test_single_candidate(self): - """Should handle single candidate without crashing.""" - estimator = ConfidenceEstimator() - - state = RefinementState( - z=np.zeros(64), - scores=np.array([0.5]), - iteration=2 - ) - state.score_history = [ - np.array([0.4]), - np.array([0.5]) - ] - - # Should not crash and should stop (single element = stable ranking) - result = estimator.should_stop(state) - assert isinstance(result, bool) - - def test_flipping_order_resets_patience(self): - """Flipping ranking order should reset stability count.""" - estimator = ConfidenceEstimator(patience=2) - - state = RefinementState( - z=np.zeros(64), - scores=np.array([0.3, 0.5, 0.1]), # Order: 1, 0, 2 - iteration=1 - ) - state.score_history = [ - np.array([0.5, 0.3, 0.1]), # Order: 0, 1, 2 - np.array([0.3, 0.5, 0.1]) # Order: 1, 0, 2 (flipped!) - ] - - # Flipped order = not stable, should not stop - assert not estimator.should_stop(state) - assert estimator._stable_count == 0 - - def test_patience_respected(self): - """Should require patience consecutive stable iterations to stop.""" - estimator = ConfidenceEstimator(patience=3) - - state = RefinementState( - z=np.zeros(64), - scores=np.array([0.5, 0.3, 0.1]), - iteration=1 - ) - - # First stable iteration - state.score_history = [ - np.array([0.5, 0.3, 0.1]), - np.array([0.5, 0.3, 0.1]) - ] - assert not estimator.should_stop(state) - assert estimator._stable_count == 1 - - # Second stable iteration - state.score_history.append(np.array([0.5, 0.3, 0.1])) - assert not estimator.should_stop(state) - assert estimator._stable_count == 2 - - # Third stable iteration - now should stop - state.score_history.append(np.array([0.5, 0.3, 0.1])) - assert estimator.should_stop(state) - assert estimator._stable_count == 3 - - def test_reset_clears_state(self): - """Reset should clear stability count.""" - estimator = ConfidenceEstimator(patience=2) - estimator._stable_count = 5 - - estimator.reset() - - assert estimator._stable_count == 0 - - -class TestRecursiveReranker: - """Tests for the main recursive reranker.""" - - def test_rerank_returns_same_count(self): - """Reranker should return same number of candidates.""" - reranker = RecursiveReranker(n_iterations=2, dim=64) - - candidates = [ - {"path": "a.py", "symbol": "func_a", "code": "def a(): pass"}, - {"path": "b.py", "symbol": "func_b", "code": "def b(): pass"}, - {"path": "c.py", "symbol": "func_c", "code": "def c(): pass"}, - ] - - results = reranker.rerank("search query", candidates) - - assert len(results) == 3 - - def test_rerank_adds_metadata(self): - """Reranked results should have recursive metadata.""" - reranker = RecursiveReranker(n_iterations=2, dim=64) - - candidates = [ - {"path": "a.py", "symbol": "func_a", "code": "def a(): pass"}, - ] - - results = reranker.rerank("query", candidates) - - assert "recursive_score" in results[0] - assert "recursive_rank" in results[0] - assert "recursive_iterations" in results[0] - assert "score_trajectory" in results[0] - - def test_rerank_preserves_original_fields(self): - """Original candidate fields should be preserved.""" - reranker = RecursiveReranker(n_iterations=2, dim=64) - - candidates = [ - {"path": "a.py", "symbol": "func_a", "code": "def a(): pass", "custom": "value"}, - ] - - results = reranker.rerank("query", candidates) - - assert results[0]["path"] == "a.py" - assert results[0]["symbol"] == "func_a" - assert results[0]["custom"] == "value" - - -class TestCosineAlphaScheduler: - """Tests for the cosine alpha scheduler.""" - - def test_schedule_length(self): - """Schedule should match n_iterations.""" - from scripts.rerank_recursive import CosineAlphaScheduler - - scheduler = CosineAlphaScheduler(n_iterations=5) - schedule = scheduler.get_schedule() - - assert len(schedule) == 5 - - def test_schedule_decreasing(self): - """Alpha should decrease over iterations (cosine decay).""" - from scripts.rerank_recursive import CosineAlphaScheduler - - scheduler = CosineAlphaScheduler(n_iterations=3, alpha_max=0.7, alpha_min=0.3) - schedule = scheduler.get_schedule() - - assert schedule[0] > schedule[1] > schedule[2] - assert abs(schedule[0] - 0.7) < 0.01 # First should be alpha_max - assert abs(schedule[2] - 0.3) < 0.01 # Last should be alpha_min - - def test_schedule_bounds(self): - """All alpha values should be within [alpha_min, alpha_max].""" - from scripts.rerank_recursive import CosineAlphaScheduler - - scheduler = CosineAlphaScheduler(n_iterations=10, alpha_max=0.8, alpha_min=0.2) - schedule = scheduler.get_schedule() - - for alpha in schedule: - assert 0.2 <= alpha <= 0.8 - - def test_single_iteration(self): - """Single iteration should return middle value.""" - from scripts.rerank_recursive import CosineAlphaScheduler - - scheduler = CosineAlphaScheduler(n_iterations=1, alpha_max=0.8, alpha_min=0.2) - schedule = scheduler.get_schedule() - - assert len(schedule) == 1 - assert abs(schedule[0] - 0.5) < 0.01 # Should be (0.8 + 0.2) / 2 - - -class TestLearnedAlphaWeights: - """Tests for the learnable alpha weights.""" - - def test_init_alpha(self): - """Initial alpha should match init_alpha parameter.""" - from scripts.rerank_recursive import LearnedAlphaWeights - - learned = LearnedAlphaWeights(n_iterations=3, init_alpha=0.6) - schedule = learned.get_schedule() - - for alpha in schedule: - assert abs(alpha - 0.6) < 0.01 - - def test_get_alpha_clamped(self): - """get_alpha should clamp to valid iteration range.""" - from scripts.rerank_recursive import LearnedAlphaWeights - - learned = LearnedAlphaWeights(n_iterations=3) - - # Should not crash for out-of-range iterations - alpha_neg = learned.get_alpha(-1) - alpha_over = learned.get_alpha(100) - - assert 0 < alpha_neg < 1 - assert 0 < alpha_over < 1 - - def test_alpha_in_valid_range(self): - """All alpha values should be in (0, 1) due to sigmoid.""" - from scripts.rerank_recursive import LearnedAlphaWeights - - learned = LearnedAlphaWeights(n_iterations=5, init_alpha=0.5) - schedule = learned.get_schedule() - - for alpha in schedule: - assert 0 < alpha < 1 - - -class TestAlphaIntegration: - """Tests for alpha scheduler integration with reranker.""" - - def test_alpha_trajectory_in_output(self): - """Reranked results should include alpha_trajectory.""" - reranker = RecursiveReranker(n_iterations=3, dim=64) - - candidates = [ - {"path": "a.py", "code": "def a(): pass"}, - ] - - results = reranker.rerank("query", candidates) - - assert "alpha_trajectory" in results[0] - assert isinstance(results[0]["alpha_trajectory"], list) - assert len(results[0]["alpha_trajectory"]) > 0 - - def test_custom_scheduler(self): - """Should accept custom alpha scheduler.""" - from scripts.rerank_recursive import LearnedAlphaWeights - - custom_scheduler = LearnedAlphaWeights(n_iterations=2, init_alpha=0.4) - reranker = RecursiveReranker(n_iterations=2, dim=64, alpha_scheduler=custom_scheduler) - - candidates = [ - {"path": "a.py", "code": "def a(): pass"}, - ] - - results = reranker.rerank("query", candidates) - - # Alpha should be close to 0.4 (our custom init) - for alpha in results[0]["alpha_trajectory"]: - assert abs(alpha - 0.4) < 0.1 - - def test_alpha_trajectory_matches_iterations(self): - """Alpha trajectory length should match actual iterations run.""" - reranker = RecursiveReranker(n_iterations=3, dim=64, early_stop=False) - - candidates = [ - {"path": "a.py", "code": "def a(): pass"}, - {"path": "b.py", "code": "def b(): pass"}, - ] - - results = reranker.rerank("query", candidates) - - # With early_stop=False, should run all iterations - assert len(results[0]["alpha_trajectory"]) == results[0]["recursive_iterations"] diff --git a/tests/test_rerank_recursive_weights.py b/tests/test_rerank_recursive_weights.py deleted file mode 100644 index 1536a2ff..00000000 --- a/tests/test_rerank_recursive_weights.py +++ /dev/null @@ -1,36 +0,0 @@ -import numpy as np - - -def test_learned_projection_forward_hot_reload(tmp_path, monkeypatch): - from scripts.rerank_recursive import projection as proj - - monkeypatch.setattr(proj.LearnedProjection, "WEIGHTS_DIR", str(tmp_path)) - lp = proj.LearnedProjection(input_dim=2, output_dim=2) - lp.WEIGHTS_RELOAD_INTERVAL = 0.0001 - lp.set_collection("test") - - np.savez(lp._weights_path, W=np.eye(2, dtype=np.float32), version=1) - lp._last_reload_check = 0.0 - lp._weights_mtime = 0.0 - out1 = lp.forward(np.array([1.0, 0.0], dtype=np.float32)) - assert np.allclose(out1, np.array([1.0, 0.0], dtype=np.float32)) - - np.savez(lp._weights_path, W=np.array([[0.0, 1.0], [1.0, 0.0]], dtype=np.float32), version=2) - lp._last_reload_check = 0.0 - lp._weights_mtime = 0.0 - out2 = lp.forward(np.array([1.0, 0.0], dtype=np.float32)) - assert np.allclose(out2, np.array([0.0, 1.0], dtype=np.float32)) - - -def test_tinyscorer_persists_recent_losses(tmp_path, monkeypatch): - from scripts.rerank_recursive import scorer as sc - - monkeypatch.setattr(sc.TinyScorer, "WEIGHTS_DIR", str(tmp_path)) - s1 = sc.TinyScorer(dim=2, hidden_dim=2) - losses = [float(i) for i in range(250)] - s1._recent_losses = list(losses) - s1._save_weights() - - s2 = sc.TinyScorer(dim=2, hidden_dim=2) - expected = np.array(losses[-200:], dtype=np.float32) - assert np.allclose(np.array(s2._recent_losses, dtype=np.float32), expected) diff --git a/tests/test_rerank_under_scope.py b/tests/test_rerank_under_scope.py new file mode 100644 index 00000000..080da9cd --- /dev/null +++ b/tests/test_rerank_under_scope.py @@ -0,0 +1,66 @@ +import importlib + + +rr = importlib.import_module("scripts.rerank_tools.local") + + +class _Pt: + def __init__(self, pid: str, path: str): + self.id = pid + self.payload = { + "metadata": { + "path": path, + "start_line": 1, + "end_line": 2, + "symbol": "f", + } + } + + +class _FakeModel: + def embed(self, texts): + for _ in texts: + yield [0.01] * 8 + + +def test_rerank_in_process_under_excludes_out_of_scope(monkeypatch): + monkeypatch.setattr(rr, "QdrantClient", lambda *a, **k: object()) + monkeypatch.setattr(rr, "_select_dense_vector_name", lambda *a, **k: "vec") + monkeypatch.setattr( + rr, + "dense_results", + lambda *a, **k: [_Pt("1", "/work/repo/direct/tools/b.py")], + ) + monkeypatch.setattr(rr, "rerank_local", lambda pairs: [0.9] * len(pairs)) + + out = rr.rerank_in_process( + query="rotate heading", + topk=10, + limit=5, + under="space", + model=_FakeModel(), + collection="codebase", + ) + assert out == [] + + +def test_rerank_in_process_under_keeps_in_scope(monkeypatch): + monkeypatch.setattr(rr, "QdrantClient", lambda *a, **k: object()) + monkeypatch.setattr(rr, "_select_dense_vector_name", lambda *a, **k: "vec") + monkeypatch.setattr( + rr, + "dense_results", + lambda *a, **k: [_Pt("1", "/work/repo/space/ship/a.py")], + ) + monkeypatch.setattr(rr, "rerank_local", lambda pairs: [0.9] * len(pairs)) + + out = rr.rerank_in_process( + query="rotate heading", + topk=10, + limit=5, + under="space", + model=_FakeModel(), + collection="codebase", + ) + assert len(out) == 1 + assert out[0]["path"] == "/work/repo/space/ship/a.py" diff --git a/tests/test_reranker_verification.py b/tests/test_reranker_verification.py index e7a05245..86835dc9 100644 --- a/tests/test_reranker_verification.py +++ b/tests/test_reranker_verification.py @@ -16,6 +16,10 @@ def tool(self, *args, **kwargs): def _decorator(fn): return fn return _decorator + def resource(self, *args, **kwargs): + def _decorator(fn): + return fn + return _decorator class _Context: def __init__(self, *args, **kwargs): @@ -50,6 +54,8 @@ async def test_rerank_inproc_changes_order(monkeypatch): # Force in-process hybrid + in-process rerank paths monkeypatch.setenv("HYBRID_IN_PROCESS", "1") monkeypatch.setenv("RERANK_IN_PROCESS", "1") + # Rerank verification suite explicitly exercises non-dense plumbing. + monkeypatch.setenv("REPO_SEARCH_DEFAULT_MODE", "hybrid") # Baseline hybrid results (JSON structured items); A before B def fake_run_hybrid_search(**kwargs): @@ -92,7 +98,7 @@ def fake_rerank_local(pairs): raising=False, ) monkeypatch.setattr( - importlib.import_module("scripts.rerank_local"), + importlib.import_module("scripts.rerank_tools.local"), "rerank_local", fake_rerank_local, ) @@ -102,7 +108,9 @@ def fake_rerank_local(pairs): assert [r["path"] for r in base["results"]] == ["/work/a.py", "/work/b.py"] # With rerank enabled, order should flip to B then A; counters should show inproc_hybrid - rr = await server.repo_search(query="q", limit=2, per_path=2, rerank_enabled=True, compact=True) + rr = await server.repo_search( + query="q", limit=2, per_path=2, rerank_enabled=True, compact=True, debug=True + ) assert rr.get("used_rerank") is True assert rr.get("rerank_counters", {}).get("inproc_hybrid", 0) >= 1 assert [r["path"] for r in rr["results"]] == ["/work/b.py", "/work/a.py"] @@ -114,6 +122,8 @@ async def test_rerank_inproc_dense_respects_collection_argument(monkeypatch): # Drive the in-process dense rerank fallback path by returning no hybrid candidates. monkeypatch.setenv("HYBRID_IN_PROCESS", "1") monkeypatch.setenv("RERANK_IN_PROCESS", "1") + # Explicit non-dense mode for rerank-path contract checks. + monkeypatch.setenv("REPO_SEARCH_DEFAULT_MODE", "hybrid") def fake_run_hybrid_search(**kwargs): return [] @@ -130,7 +140,7 @@ def fake_rerank_in_process(**kwargs): return [] monkeypatch.setattr( - importlib.import_module("scripts.rerank_local"), + importlib.import_module("scripts.rerank_tools.local"), "rerank_in_process", fake_rerank_in_process, ) @@ -147,12 +157,79 @@ def fake_rerank_in_process(**kwargs): assert captured.get("collection") == "other-collection" +@pytest.mark.service +@pytest.mark.anyio +async def test_rerank_inproc_dense_respects_path_filters(monkeypatch): + monkeypatch.setenv("HYBRID_IN_PROCESS", "1") + monkeypatch.setenv("RERANK_IN_PROCESS", "1") + # Explicit non-dense mode for rerank-path contract checks. + monkeypatch.setenv("REPO_SEARCH_DEFAULT_MODE", "hybrid") + + def fake_run_hybrid_search(**kwargs): + return [] + + monkeypatch.setitem(sys.modules, "scripts.hybrid_search", _make_hybrid_stub(fake_run_hybrid_search)) + monkeypatch.delitem(sys.modules, "scripts.mcp_indexer_server", raising=False) + server = importlib.import_module("scripts.mcp_indexer_server") + monkeypatch.setattr(server, "_get_embedding_model", _fake_embedding_model) + + def fake_rerank_in_process(**kwargs): + return [ + {"score": 0.9, "path": "/work/src/a.py", "symbol": "", "start_line": 1, "end_line": 3}, + {"score": 0.8, "path": "/work/tests/b.py", "symbol": "", "start_line": 5, "end_line": 9}, + { + "score": 0.7, + "path": "/home/coder/project/Context-Engine/scripts/mcp_impl/search.py", + "symbol": "", + "start_line": 10, + "end_line": 20, + }, + ] + + monkeypatch.setattr( + importlib.import_module("scripts.rerank_tools.local"), + "rerank_in_process", + fake_rerank_in_process, + ) + + only_tests = await server.repo_search( + query="q", + limit=10, + rerank_enabled=True, + path_glob=["tests/**"], + compact=True, + ) + assert [r["path"] for r in only_tests["results"]] == ["/work/tests/b.py"] + + no_tests = await server.repo_search( + query="q", + limit=10, + rerank_enabled=True, + not_glob=["**/tests/**"], + compact=True, + ) + assert all("/tests/" not in r["path"] for r in no_tests["results"]) + + host_rel_glob = await server.repo_search( + query="q", + limit=10, + rerank_enabled=True, + path_glob=["scripts/mcp_impl/**"], + compact=True, + ) + assert [r["path"] for r in host_rel_glob["results"]] == [ + "/home/coder/project/Context-Engine/scripts/mcp_impl/search.py" + ] + + @pytest.mark.service @pytest.mark.anyio async def test_rerank_subprocess_timeout_fallback(monkeypatch): # Force hybrid via subprocess output (doesn't matter which) and disable inproc rerank monkeypatch.setenv("HYBRID_IN_PROCESS", "1") monkeypatch.setenv("RERANK_IN_PROCESS", "0") + # Explicit non-dense mode for rerank-path contract checks. + monkeypatch.setenv("REPO_SEARCH_DEFAULT_MODE", "hybrid") def fake_run_hybrid_search(**kwargs): return [ @@ -187,9 +264,9 @@ async def fake_run_async(cmd, env=None, timeout=None): rerank_enabled=True, compact=True, collection="test-coll", + debug=True, ) # Fallback should keep original order from hybrid; timeout counter incremented assert rr.get("used_rerank") is False assert rr.get("rerank_counters", {}).get("timeout", 0) >= 1 assert [r["path"] for r in rr["results"]] == ["/work/a.py", "/work/b.py"] - diff --git a/tests/test_router_batching.py b/tests/test_router_batching.py deleted file mode 100644 index 6795caec..00000000 --- a/tests/test_router_batching.py +++ /dev/null @@ -1,146 +0,0 @@ -import threading -import time - -import pytest - -from scripts.mcp_router import BatchingContextAnswerClient - - -class _Counter: - def __init__(self): - self.n = 0 - self.lock = threading.Lock() - - def inc(self): - with self.lock: - self.n += 1 - return self.n - - -def _fake_call_factory(counter: _Counter): - def _fake_call(base_url: str, tool: str, args: dict, timeout: float = 1.0): - # Simulate a tiny network call and count invocations - counter.inc() - time.sleep(0.01) - q = args.get("query") - queries = args.get("queries") or ([q] if q else ([] if q is None else ([q] if not isinstance(q, list) else q))) - # When multiple queries are provided (aggregated call), return structured per-query answers - answers_by_query = None - if isinstance(q, list) and len(q) > 1: - answers_by_query = [ - {"query": str(qi), "answer": "ok", "citations": []} for qi in q - ] - return { - "result": { - "structuredContent": { - "result": { - "answer": "ok", - "citations": [], - "query": queries, - **({"answers_by_query": answers_by_query} if answers_by_query else {}), - } - } - } - } - - return _fake_call - - -def test_batching_merges_identical_queries(): - counter = _Counter() - client = BatchingContextAnswerClient( - call_func=_fake_call_factory(counter), - enable=True, - window_ms=120, - max_batch=8, - budget_ms=2000, - ) - - results: list[dict] = [] - barrier = threading.Barrier(3) - - def worker(): - barrier.wait() - res = client.call_or_enqueue( - "http://localhost:8003/mcp", - "context_answer", - {"query": "What is batching?", "limit": 5}, - timeout=1.0, - ) - results.append(res) - - t1 = threading.Thread(target=worker) - t2 = threading.Thread(target=worker) - t1.start(); t2.start() - barrier.wait() # release both workers - t1.join(); t2.join() - - # Exactly one underlying call, two client results - assert counter.n == 1 - assert len(results) == 2 - for r in results: - assert r.get("result", {}).get("structuredContent", {}).get("result", {}).get("answer") == "ok" - - -def test_batching_cap_flushes_early(): - counter = _Counter() - client = BatchingContextAnswerClient( - call_func=_fake_call_factory(counter), - enable=True, - window_ms=5000, # long window, but cap will force immediate flush - max_batch=2, - budget_ms=2000, - ) - - results: list[dict] = [] - barrier = threading.Barrier(3) - - def worker(q): - barrier.wait() - res = client.call_or_enqueue( - "http://localhost:8003/mcp", - "context_answer", - {"query": q, "limit": 5}, - timeout=1.0, - ) - results.append(res) - - t1 = threading.Thread(target=worker, args=("A",)) - t2 = threading.Thread(target=worker, args=("B",)) - t1.start(); t2.start() - barrier.wait() - t1.join(); t2.join() - - # Cap reached: we flush once and make a single aggregated call - assert counter.n == 1 - assert len(results) == 2 - - -def test_bypass_immediate_flag_calls_direct(): - counter = _Counter() - client = BatchingContextAnswerClient( - call_func=_fake_call_factory(counter), - enable=True, - window_ms=200, - max_batch=8, - budget_ms=2000, - ) - - # Two direct calls because of immediate flag; they should not be batched - r1 = client.call_or_enqueue( - "http://localhost:8003/mcp", - "context_answer", - {"query": "Q1", "limit": 5, "immediate": True}, - timeout=1.0, - ) - r2 = client.call_or_enqueue( - "http://localhost:8003/mcp", - "context_answer", - {"query": "Q2", "limit": 5, "immediate": True}, - timeout=1.0, - ) - - assert counter.n == 2 - assert r1.get("result", {}).get("structuredContent", {}).get("result", {}).get("answer") == "ok" - assert r2.get("result", {}).get("structuredContent", {}).get("result", {}).get("answer") == "ok" - diff --git a/tests/test_router_batching_demux.py b/tests/test_router_batching_demux.py deleted file mode 100644 index 7cae1ede..00000000 --- a/tests/test_router_batching_demux.py +++ /dev/null @@ -1,103 +0,0 @@ -import threading -import time - -from scripts.mcp_router import BatchingContextAnswerClient - - -class _Counter: - def __init__(self): - self.n = 0 - self.lock = threading.Lock() - - def inc(self): - with self.lock: - self.n += 1 - return self.n - - -def _fake_call_factory(counter: _Counter): - def _fake_call(base_url: str, tool: str, args: dict, timeout: float = 1.0): - counter.inc() - time.sleep(0.01) - q = args.get("query") - queries = args.get("queries") or ([q] if q else ([] if q is None else ([q] if not isinstance(q, list) else q))) - answers_by_query = None - if isinstance(q, list) and len(q) > 1: - answers_by_query = [ - {"query": str(qi), "answer": f"ok:{qi}", "citations": []} for qi in q - ] - return { - "result": { - "structuredContent": { - "result": { - "answer": f"ok:{q}", - "citations": [], - "query": queries, - **({"answers_by_query": answers_by_query} if answers_by_query else {}), - } - } - } - } - - return _fake_call - - -def test_demultiplex_different_queries_results_are_isolated(): - counter = _Counter() - client = BatchingContextAnswerClient( - call_func=_fake_call_factory(counter), - enable=True, - window_ms=120, - max_batch=8, - budget_ms=2000, - ) - - results: list[tuple[str, dict]] = [] - barrier = threading.Barrier(3) - - def worker(q: str): - barrier.wait() - res = client.call_or_enqueue( - "http://localhost:8003/mcp", - "context_answer", - {"query": q, "limit": 5}, - timeout=1.0, - ) - results.append((q, res)) - - t1 = threading.Thread(target=worker, args=("Q1",)) - t2 = threading.Thread(target=worker, args=("Q2",)) - t1.start(); t2.start() - barrier.wait() - t1.join(); t2.join() - - # Aggregated call once, demux per-query reply - assert counter.n == 1 - assert len(results) == 2 - for q, r in results: - rq = r.get("result", {}).get("structuredContent", {}).get("result", {}).get("query") - # Each result should reflect only its own query - assert rq == [q] - - -def test_budget_fallback_does_not_double_call(): - counter = _Counter() - client = BatchingContextAnswerClient( - call_func=_fake_call_factory(counter), - enable=True, - window_ms=500, # long window so timer would fire later - max_batch=8, - budget_ms=10, # tiny budget to force immediate fallback - ) - - res = client.call_or_enqueue( - "http://localhost:8003/mcp", - "context_answer", - {"query": "late", "limit": 5}, - timeout=1.0, - ) - assert res - # Wait beyond the window; if slot was not removed, we'd see a second call when timer flushes - time.sleep(0.6) - assert counter.n == 1 - diff --git a/tests/test_server_helpers.py b/tests/test_server_helpers.py index 94212145..74ea78e6 100644 --- a/tests/test_server_helpers.py +++ b/tests/test_server_helpers.py @@ -1,8 +1,11 @@ import json +import asyncio import types import importlib +from pathlib import Path srv = importlib.import_module("scripts.mcp_indexer_server") +admin_tools = importlib.import_module("scripts.mcp_impl.admin_tools") def test_tokens_from_queries_basic(): @@ -17,6 +20,42 @@ def test_highlight_snippet_simple(): assert "<>" in out and "<>" in out +def test_detect_repo_from_work_path_ignores_invalid_root_git(tmp_path, monkeypatch): + """A metadata-only /work/.git must not be treated as repo name "work".""" + work = tmp_path / "work" + (work / ".git" / ".codebase").mkdir(parents=True) + monkeypatch.delenv("CURRENT_REPO", raising=False) + monkeypatch.delenv("REPO_NAME", raising=False) + monkeypatch.setattr(admin_tools, "Path", lambda value: work if value == "/work" else Path(value)) + + assert admin_tools._detect_current_repo() is None + + +def test_detect_repo_from_work_path_skips_internal_dirs(tmp_path, monkeypatch): + work = tmp_path / "work" + (work / ".git").mkdir(parents=True) + (work / ".codebase" / ".git").mkdir(parents=True) + (work / "__pycache__" / ".git").mkdir(parents=True) + (work / "real-repo" / ".git").mkdir(parents=True) + monkeypatch.delenv("CURRENT_REPO", raising=False) + monkeypatch.delenv("REPO_NAME", raising=False) + monkeypatch.setenv("CTXCE_BINDMOUNT_REPO_DETECTION", "1") + monkeypatch.setattr(admin_tools, "Path", lambda value: work if value == "/work" else Path(value)) + + assert admin_tools._detect_current_repo() == "real-repo" + + +def test_detect_repo_from_work_path_skips_git_without_bindmount_mode(tmp_path, monkeypatch): + work = tmp_path / "work" + (work / "real-repo" / ".git").mkdir(parents=True) + monkeypatch.delenv("CURRENT_REPO", raising=False) + monkeypatch.delenv("REPO_NAME", raising=False) + monkeypatch.delenv("CTXCE_BINDMOUNT_REPO_DETECTION", raising=False) + monkeypatch.setattr(admin_tools, "Path", lambda value: work if value == "/work" else Path(value)) + + assert admin_tools._detect_current_repo() is None + + def fake_async_run_factory(text): async def _fake(cmd, **kwargs): # accept env/timeout/cwd return {"ok": True, "code": 0, "stdout": text, "stderr": ""} @@ -53,10 +92,13 @@ def test_repo_search_arg_normalization(monkeypatch, tmp_path): # Ensure in-process branch stays off monkeypatch.delenv("HYBRID_IN_PROCESS", raising=False) - res = srv.asyncio.get_event_loop().run_until_complete( + res = asyncio.run( _call_repo_search( queries=["FooBar"], limit="12", # str on purpose to test coercion + # This test targets arg normalization + JSONL shaping from the non-dense path. + # Keep mode explicit so global dense defaults don't change behavior here. + mode="hybrid", per_path=None, language=None, under=None, diff --git a/tests/test_service_context_search.py b/tests/test_service_context_search.py index ff7b593f..ca66ba51 100644 --- a/tests/test_service_context_search.py +++ b/tests/test_service_context_search.py @@ -1,25 +1,10 @@ import importlib import json +import sys +import types import pytest -srv = importlib.import_module("scripts.mcp_indexer_server") - - -class FakePoint: - def __init__(self, score, payload): - self.score = score - self.payload = payload - - -class FakeQdrantMem: - def __init__(self, items): - self._items = items - - def search(self, **kwargs): - return self._items - - def scroll(self, **kwargs): - return (self._items, None) +ctx_search = importlib.import_module("scripts.mcp_impl.context_search") class FakeEmbed: @@ -45,27 +30,60 @@ async def fake_repo_search(**kwargs): ] } - monkeypatch.setattr(srv, "repo_search", fake_repo_search) - monkeypatch.setattr(srv, "_get_embedding_model", lambda *a, **k: FakeEmbed()) + monkeypatch.setenv("MEMORY_SSE_ENABLED", "1") + monkeypatch.setenv("MEMORY_COLLECTION_NAME", "test-memory") + monkeypatch.setenv("MEMORY_MCP_READY_RETRIES", "1") + monkeypatch.setenv("MEMORY_MCP_READY_BACKOFF", "0") + monkeypatch.setenv("MEMORY_MCP_LIST_RETRIES", "1") + monkeypatch.setenv("MEMORY_MCP_LIST_BACKOFF", "0") - # Memory fallback via Qdrant: two memory-like points (no path in metadata) - mem_items = [ - FakePoint(0.9, {"content": "foo note one", "metadata": {}}), - FakePoint(0.2, {"content": "bar note two", "metadata": {}}), - ] - import qdrant_client + import urllib.request monkeypatch.setattr( - qdrant_client, "QdrantClient", lambda *a, **k: FakeQdrantMem(mem_items) + urllib.request, + "urlopen", + lambda *a, **k: (_ for _ in ()).throw(OSError("not ready")), ) - res = await srv.context_search( + class T: + def __init__(self, name): + self.name = name + + class Item: + def __init__(self, text): + self.text = text + + class Resp: + def __init__(self): + self.content = [Item("foo note one"), Item("bar note two")] + + class FakeClient: + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc, tb): + return False + + async def list_tools(self): + return [T("find")] + + async def call_tool(self, *a, **k): + return Resp() + + monkeypatch.setitem( + sys.modules, + "fastmcp", + types.SimpleNamespace(Client=lambda *a, **k: FakeClient()), + ) + + res = await ctx_search._context_search_impl( query="foo bar", limit=3, per_path=1, include_memories=True, memory_weight=0.5, compact=True, + repo_search_fn=fake_repo_search, ) assert "results" in res @@ -88,10 +106,21 @@ async def fake_repo_search(**kwargs): ] } - monkeypatch.setattr(srv, "repo_search", fake_repo_search) - # Force SSE memory path with a fake FastMCP client monkeypatch.setenv("MEMORY_SSE_ENABLED", "1") + monkeypatch.setenv("MEMORY_COLLECTION_NAME", "test-memory") + monkeypatch.setenv("MEMORY_MCP_READY_RETRIES", "1") + monkeypatch.setenv("MEMORY_MCP_READY_BACKOFF", "0") + monkeypatch.setenv("MEMORY_MCP_LIST_RETRIES", "1") + monkeypatch.setenv("MEMORY_MCP_LIST_BACKOFF", "0") + + import urllib.request + + monkeypatch.setattr( + urllib.request, + "urlopen", + lambda *a, **k: (_ for _ in ()).throw(OSError("not ready")), + ) class T: def __init__(self, name): @@ -118,23 +147,20 @@ async def list_tools(self): async def call_tool(self, *a, **k): return Resp() - # Import fastmcp inside test to avoid module-level import conflicts - # Clear any broken mcp modules from sys.modules first - import sys - mcp_modules = [k for k in sys.modules.keys() if k == 'mcp' or k.startswith('mcp.')] - for mod in mcp_modules: - if mod in sys.modules and not hasattr(sys.modules.get(mod, object()), 'types'): - del sys.modules[mod] - import fastmcp - monkeypatch.setattr(fastmcp, "Client", lambda *a, **k: FakeClient()) - - res = await srv.context_search( + monkeypatch.setitem( + sys.modules, + "fastmcp", + types.SimpleNamespace(Client=lambda *a, **k: FakeClient()), + ) + + res = await ctx_search._context_search_impl( query="foo", limit=2, per_path=1, include_memories=True, memory_weight=2.0, compact=False, + repo_search_fn=fake_repo_search, ) mem_scores = [r["score"] for r in res["results"] if r.get("source") == "memory"] @@ -157,11 +183,21 @@ async def fake_repo_search(**kwargs): ] } - monkeypatch.setattr(srv, "repo_search", fake_repo_search) - monkeypatch.setattr(srv, "_get_embedding_model", lambda *a, **k: FakeEmbed()) - # Drive memory hits via SSE path with a fake FastMCP client yielding 3 notes monkeypatch.setenv("MEMORY_SSE_ENABLED", "1") + monkeypatch.setenv("MEMORY_COLLECTION_NAME", "test-memory") + monkeypatch.setenv("MEMORY_MCP_READY_RETRIES", "1") + monkeypatch.setenv("MEMORY_MCP_READY_BACKOFF", "0") + monkeypatch.setenv("MEMORY_MCP_LIST_RETRIES", "1") + monkeypatch.setenv("MEMORY_MCP_LIST_BACKOFF", "0") + + import urllib.request + + monkeypatch.setattr( + urllib.request, + "urlopen", + lambda *a, **k: (_ for _ in ()).throw(OSError("not ready")), + ) class T: def __init__(self, name): @@ -188,23 +224,21 @@ async def list_tools(self): async def call_tool(self, *a, **k): return Resp() - # Import fastmcp inside test to avoid module-level import conflicts - # Clear any broken mcp modules from sys.modules first - import sys - mcp_modules = [k for k in sys.modules.keys() if k == 'mcp' or k.startswith('mcp.')] - for mod in mcp_modules: - if mod in sys.modules and not hasattr(sys.modules.get(mod, object()), 'types'): - del sys.modules[mod] - import fastmcp - monkeypatch.setattr(fastmcp, "Client", lambda *a, **k: FakeClient()) - - res = await srv.context_search( + monkeypatch.setitem( + sys.modules, + "fastmcp", + types.SimpleNamespace(Client=lambda *a, **k: FakeClient()), + ) + + res = await ctx_search._context_search_impl( query="foo", limit=5, per_path=1, include_memories=True, per_source_limits=json.dumps({"code": 1, "memory": 2}), compact=True, + repo_search_fn=fake_repo_search, + get_embedding_model_fn=lambda *a, **k: FakeEmbed(), ) kinds = [r.get("source") for r in res.get("results", [])] diff --git a/tests/test_service_qdrant_status.py b/tests/test_service_qdrant_status.py index df04254c..38a1ec6b 100644 --- a/tests/test_service_qdrant_status.py +++ b/tests/test_service_qdrant_status.py @@ -1,4 +1,5 @@ import types +import asyncio import importlib import pytest @@ -31,7 +32,7 @@ def test_qdrant_status_mocked(monkeypatch): monkeypatch.setattr(qdrant_client, "QdrantClient", lambda *a, **k: FakeQdrant()) - out = srv.asyncio.get_event_loop().run_until_complete( + out = asyncio.run( srv.qdrant_status(collection="test") ) # qdrant_status returns a summary shape without an 'ok' key diff --git a/tests/test_smart_reindex_vectors.py b/tests/test_smart_reindex_vectors.py index 2e77056e..d39dac99 100644 --- a/tests/test_smart_reindex_vectors.py +++ b/tests/test_smart_reindex_vectors.py @@ -2,10 +2,86 @@ import sys from types import SimpleNamespace from pathlib import Path +from unittest.mock import MagicMock import pytest +class _Sym(dict): + __getattr__ = dict.get + + +def _patch_symbol(monkeypatch, ingest_pipeline, *, name: str, start: int = 1, end: int = 2): + monkeypatch.setattr( + ingest_pipeline, + "extract_symbols_with_tree_sitter", + lambda _fp: { + f"function_{name}_{start}": { + "name": name, + "type": "function", + "start_line": start, + "end_line": end, + "content_hash": "samehash", + "pseudo": "", + "tags": [], + "qdrant_ids": [], + } + }, + ) + monkeypatch.setattr( + ingest_pipeline, + "_extract_symbols", + lambda *_a, **_k: [ + _Sym(kind="function", name=name, path=name, start=start, end=end) + ], + ) + + +def _patch_qdrant_models(monkeypatch, ingest_pipeline): + class FakeModels: + class Filter: + def __init__(self, **kwargs): + self.__dict__.update(kwargs) + + class FieldCondition: + def __init__(self, **kwargs): + self.__dict__.update(kwargs) + + class MatchValue: + def __init__(self, **kwargs): + self.__dict__.update(kwargs) + + class SparseVector: + def __init__(self, **kwargs): + self.__dict__.update(kwargs) + + class PointStruct: + def __init__(self, id, vector, payload): + self.id = id + self.vector = vector + self.payload = payload + + monkeypatch.setattr(ingest_pipeline, "models", FakeModels) + + +def _patch_smart_side_effects(monkeypatch, ingest_pipeline): + monkeypatch.setenv("LEX_SPARSE_MODE", "0") + monkeypatch.setattr(ingest_pipeline, "LEX_SPARSE_MODE", False) + monkeypatch.setattr( + ingest_pipeline, + "_sync_graph_edges_best_effort", + lambda *a, **k: None, + raising=False, + ) + monkeypatch.setattr(ingest_pipeline, "_get_imports_calls", lambda *a, **k: ([], [])) + monkeypatch.setattr(ingest_pipeline, "_git_metadata", lambda *a, **k: (0, 0, 0)) + monkeypatch.setattr( + ingest_pipeline, + "_compute_host_and_container_paths", + lambda _p: ("", ""), + ) + + @pytest.mark.usefixtures("monkeypatch") def test_smart_reindex_refreshes_lex_vector_for_reused_chunks(tmp_path, monkeypatch): """When reusing an existing dense embedding, smart reindex must refresh LEX vector. @@ -15,36 +91,39 @@ def test_smart_reindex_refreshes_lex_vector_for_reused_chunks(tmp_path, monkeypa # The smart reindex logic we test doesn't require the real library. monkeypatch.setitem(sys.modules, "fastembed", SimpleNamespace(TextEmbedding=object)) - from scripts import ingest_code + from scripts.ingest import pipeline as ingest_pipeline + _patch_qdrant_models(monkeypatch, ingest_pipeline) # Deterministic pseudo/tags so we can predict lexical vector. monkeypatch.setattr( - ingest_code, + ingest_pipeline, "should_process_pseudo_for_chunk", lambda fp, ch, changed: (False, "pseudo", ["tag"]), ) # Avoid touching any caches. - monkeypatch.setattr(ingest_code, "get_cached_symbols", lambda fp: {}) - monkeypatch.setattr(ingest_code, "compare_symbol_changes", lambda a, b: ([], [])) - monkeypatch.setattr(ingest_code, "set_cached_pseudo", None) - monkeypatch.setattr(ingest_code, "set_cached_symbols", None) - monkeypatch.setattr(ingest_code, "set_cached_file_hash", None) + monkeypatch.setattr(ingest_pipeline, "get_cached_symbols", lambda fp: {}) + monkeypatch.setattr(ingest_pipeline, "compare_symbol_changes", lambda a, b: ([], [])) + monkeypatch.setattr(ingest_pipeline, "set_cached_pseudo", None) + monkeypatch.setattr(ingest_pipeline, "set_cached_symbols", None) + monkeypatch.setattr(ingest_pipeline, "set_cached_file_hash", None) # Force simple line chunking. monkeypatch.setenv("INDEX_MICRO_CHUNKS", "0") monkeypatch.setenv("INDEX_SEMANTIC_CHUNKS", "0") monkeypatch.setenv("USE_TREE_SITTER", "0") monkeypatch.setenv("REFRAG_MODE", "0") + _patch_symbol(monkeypatch, ingest_pipeline, name="add") + _patch_smart_side_effects(monkeypatch, ingest_pipeline) code = "def add(a, b):\n return a + b\n" fp = tmp_path / "x.py" fp.write_text(code, encoding="utf-8") # Compute the exact chunk text the indexer will use. - chunk = ingest_code.chunk_lines(code, max_lines=120, overlap=20)[0] + chunk = ingest_pipeline.chunk_lines(code, max_lines=120, overlap=20)[0] code_text = chunk["text"] - info_text = ingest_code.build_information( + info_text = ingest_pipeline.build_information( "python", Path(fp), chunk["start"], @@ -53,7 +132,7 @@ def test_smart_reindex_refreshes_lex_vector_for_reused_chunks(tmp_path, monkeypa ) dense_key = "dense" - old_lex = [0.0] * ingest_code.LEX_VECTOR_DIM + old_lex = [0.0] * ingest_pipeline.LEX_VECTOR_DIM old_lex[0] = 1.0 existing_record = SimpleNamespace( @@ -68,7 +147,7 @@ def test_smart_reindex_refreshes_lex_vector_for_reused_chunks(tmp_path, monkeypa "start_line": 1, } }, - vector={dense_key: [0.1, 0.2, 0.3], ingest_code.LEX_VECTOR_NAME: old_lex}, + vector={dense_key: [0.1, 0.2, 0.3], ingest_pipeline.LEX_VECTOR_NAME: old_lex}, ) class FakeClient: @@ -84,17 +163,17 @@ def scroll(self, **kwargs): def fake_upsert_points(_client, _collection, points): captured["points"] = points - monkeypatch.setattr(ingest_code, "upsert_points", fake_upsert_points) - monkeypatch.setattr(ingest_code, "delete_points_by_path", lambda *a, **k: None) + monkeypatch.setattr(ingest_pipeline, "upsert_points", fake_upsert_points) + monkeypatch.setattr(ingest_pipeline, "delete_points_by_path", lambda *a, **k: None) # Mock embed_batch since dense enrichment may trigger re-embedding reused_dense = [0.1, 0.2, 0.3] - monkeypatch.setattr(ingest_code, "embed_batch", lambda _model, texts: [reused_dense for _ in texts]) + monkeypatch.setattr(ingest_pipeline, "embed_batch", lambda _model, texts: [reused_dense for _ in texts]) # Model is unused when embeddings are mocked. dummy_model = object() - status = ingest_code.process_file_with_smart_reindexing( + status = ingest_pipeline.process_file_with_smart_reindexing( file_path=Path(fp), text=code, language="python", @@ -110,13 +189,240 @@ def fake_upsert_points(_client, _collection, points): out_vec = captured["points"][0].vector assert isinstance(out_vec, dict) - assert ingest_code.LEX_VECTOR_NAME in out_vec + assert ingest_pipeline.LEX_VECTOR_NAME in out_vec expected_aug = (code_text or "") + " pseudo" + " tag" - expected_lex = ingest_code._lex_hash_vector_text(expected_aug) - assert out_vec[ingest_code.LEX_VECTOR_NAME] == expected_lex + expected_lex = ingest_pipeline._lex_hash_vector_text(expected_aug) + assert out_vec[ingest_pipeline.LEX_VECTOR_NAME] == expected_lex # Make sure we didn't keep the old lex vector. - assert out_vec[ingest_code.LEX_VECTOR_NAME] != old_lex + assert out_vec[ingest_pipeline.LEX_VECTOR_NAME] != old_lex + + +def test_smart_reindex_does_not_call_batch_pseudo_when_disabled(tmp_path, monkeypatch): + """The smart batch shortcut must honor the explicit pseudo switch.""" + monkeypatch.setitem(sys.modules, "fastembed", SimpleNamespace(TextEmbedding=object)) + monkeypatch.setenv("PSEUDO_BATCH_CONCURRENCY", "4") + monkeypatch.setenv("REFRAG_PSEUDO_DESCRIBE", "0") + monkeypatch.setenv("INDEX_MICRO_CHUNKS", "0") + monkeypatch.setenv("INDEX_SEMANTIC_CHUNKS", "0") + monkeypatch.setenv("REFRAG_MODE", "0") + + from scripts.ingest import pipeline as ingest_pipeline + import scripts.refrag_glm as refrag_glm + + _patch_qdrant_models(monkeypatch, ingest_pipeline) + _patch_symbol(monkeypatch, ingest_pipeline, name="add") + _patch_smart_side_effects(monkeypatch, ingest_pipeline) + monkeypatch.setattr(ingest_pipeline, "get_cached_symbols", lambda _fp: {}) + monkeypatch.setattr(ingest_pipeline, "set_cached_symbols", None) + monkeypatch.setattr(ingest_pipeline, "set_cached_file_hash", None) + monkeypatch.setattr(ingest_pipeline, "set_cached_pseudo", None) + monkeypatch.setattr(ingest_pipeline, "upsert_points", lambda *a, **k: None) + monkeypatch.setattr(ingest_pipeline, "embed_batch", lambda _model, texts: [[0.1, 0.2, 0.3] for _ in texts]) + monkeypatch.setattr(ingest_pipeline, "ensure_collection_and_indexes_once", lambda *a, **k: None) + + batch_calls = [] + monkeypatch.setattr( + refrag_glm, + "generate_pseudo_tags_batch", + lambda *a, **k: batch_calls.append((a, k)), + ) + + fp = tmp_path / "x.py" + fp.write_text("def add(a, b):\n return a + b\n", encoding="utf-8") + + class FakeClient: + def scroll(self, **kwargs): + return ([], None) + + status = ingest_pipeline.process_file_with_smart_reindexing( + file_path=fp, + text=fp.read_text(encoding="utf-8"), + language="python", + client=FakeClient(), + current_collection="c", + per_file_repo="r", + model=object(), + vector_name="dense", + model_dim=3, + allowed_vectors=set(), + allowed_sparse=set(), + ) + + assert status == "success" + assert batch_calls == [] + + +def test_should_process_pseudo_for_chunk_reuses_cache_after_line_shift(monkeypatch): + from scripts.ingest import pseudo as pseudo_mod + + monkeypatch.setattr(pseudo_mod, "get_cached_pseudo", lambda *a, **k: ("", [])) + monkeypatch.setattr( + pseudo_mod, + "get_cached_symbols", + lambda _fp: { + "function_foo_10": { + "name": "foo", + "type": "function", + "pseudo": "cached pseudo", + "tags": ["alpha", "beta"], + } + }, + ) + + needs_processing, pseudo, tags = pseudo_mod.should_process_pseudo_for_chunk( + "x.py", + {"symbol": "foo", "kind": "function", "start": 12}, + changed_symbols=set(), + ) + + assert needs_processing is False + assert pseudo == "cached pseudo" + assert tags == ["alpha", "beta"] + + +def test_smart_reindex_persists_pseudo_on_shifted_symbol_ids(tmp_path, monkeypatch): + monkeypatch.setitem(sys.modules, "fastembed", SimpleNamespace(TextEmbedding=object)) + monkeypatch.setenv("PSEUDO_BATCH_CONCURRENCY", "1") + + from scripts.ingest import pipeline as ingest_pipeline + _patch_qdrant_models(monkeypatch, ingest_pipeline) + + fp = tmp_path / "x.py" + fp.write_text("def foo():\n return 1\n", encoding="utf-8") + + monkeypatch.setattr( + ingest_pipeline, + "extract_symbols_with_tree_sitter", + lambda _fp: { + "function_foo_12": { + "name": "foo", + "type": "function", + "start_line": 12, + "end_line": 13, + "content_hash": "samehash", + "pseudo": "", + "tags": [], + "qdrant_ids": [], + }, + "function_bar_20": { + "name": "bar", + "type": "function", + "start_line": 20, + "end_line": 21, + "content_hash": "barhash-new", + "pseudo": "", + "tags": [], + "qdrant_ids": [], + }, + }, + ) + monkeypatch.setattr( + ingest_pipeline, + "get_cached_symbols", + lambda _fp: { + "function_foo_10": { + "name": "foo", + "type": "function", + "start_line": 10, + "end_line": 11, + "content_hash": "samehash", + "pseudo": "cached pseudo", + "tags": ["tag1"], + "qdrant_ids": [], + }, + "function_bar_20": { + "name": "bar", + "type": "function", + "start_line": 20, + "end_line": 21, + "content_hash": "barhash-old", + "pseudo": "old bar", + "tags": ["old"], + "qdrant_ids": [], + }, + }, + ) + monkeypatch.setattr( + ingest_pipeline, + "compare_symbol_changes", + lambda *_: (["function_foo_12"], ["function_bar_20"]), + ) + monkeypatch.setattr(ingest_pipeline, "ensure_collection_and_indexes_once", lambda *a, **k: None) + + class FakeClient: + def scroll(self, **kwargs): + return ([], None) + + monkeypatch.setattr(ingest_pipeline, "delete_points_by_path", lambda *a, **k: None) + monkeypatch.setattr(ingest_pipeline, "upsert_points", lambda *a, **k: None) + monkeypatch.setattr( + ingest_pipeline, + "_sync_graph_edges_best_effort", + lambda *a, **k: None, + raising=False, + ) + monkeypatch.setattr(ingest_pipeline, "_get_imports_calls", lambda *a, **k: ([], [])) + monkeypatch.setattr(ingest_pipeline, "_git_metadata", lambda *a, **k: (0, 0, 0)) + monkeypatch.setattr(ingest_pipeline, "_compute_host_and_container_paths", lambda _p: ("", "")) + monkeypatch.setattr(ingest_pipeline, "_lex_hash_vector_text", lambda _t: [0.0] * ingest_pipeline.LEX_VECTOR_DIM) + monkeypatch.setattr(ingest_pipeline, "_select_dense_text", lambda **kwargs: kwargs.get("code_text") or "") + monkeypatch.setattr(ingest_pipeline, "embed_batch", lambda _model, texts: [[0.1, 0.2, 0.3] for _ in texts]) + monkeypatch.setattr(ingest_pipeline, "embed_batch", lambda _model, texts: [[0.1, 0.2, 0.3] for _ in texts]) + monkeypatch.setattr(ingest_pipeline, "generate_pseudo_tags", lambda _t: ("NEW", ["fresh"])) + monkeypatch.setattr( + ingest_pipeline, + "chunk_lines", + lambda text, *_a, **_k: [ + {"start": 12, "end": 13, "text": text, "symbol": "foo", "kind": "function"}, + {"start": 20, "end": 21, "text": text, "symbol": "bar", "kind": "function"}, + ], + ) + monkeypatch.setattr( + ingest_pipeline, + "chunk_semantic", + lambda text, *_a, **_k: [ + {"start": 12, "end": 13, "text": text, "symbol": "foo", "kind": "function"}, + {"start": 20, "end": 21, "text": text, "symbol": "bar", "kind": "function"}, + ], + ) + monkeypatch.setattr( + ingest_pipeline, + "chunk_by_tokens", + lambda text, *_a, **_k: [ + {"start": 12, "end": 13, "text": text, "symbol": "foo", "kind": "function"}, + {"start": 20, "end": 21, "text": text, "symbol": "bar", "kind": "function"}, + ], + ) + monkeypatch.setattr(ingest_pipeline, "_extract_symbols", lambda *_a, **_k: []) + monkeypatch.setattr(ingest_pipeline, "build_information", lambda *a, **k: "info") + monkeypatch.setattr(ingest_pipeline, "hash_id", lambda *a, **k: 1) + monkeypatch.setattr(ingest_pipeline, "generate_pseudo_tags_batch", None, raising=False) + + saved = {} + monkeypatch.setattr(ingest_pipeline, "set_cached_pseudo", lambda *a, **k: None) + monkeypatch.setattr(ingest_pipeline, "set_cached_file_hash", lambda *a, **k: None) + monkeypatch.setattr(ingest_pipeline, "should_process_pseudo_for_chunk", ingest_pipeline.should_process_pseudo_for_chunk) + monkeypatch.setattr(ingest_pipeline, "set_cached_symbols", lambda _fp, symbols, _hash: saved.update(symbols)) + + status = ingest_pipeline.process_file_with_smart_reindexing( + file_path=fp, + text=fp.read_text(encoding="utf-8"), + language="python", + client=FakeClient(), + current_collection="c", + per_file_repo="r", + model=object(), + vector_name="dense", + model_dim=3, + ) + + assert status == "success" + # `foo` is logically reusable across the line shift, but chunk-level pseudo + # generation may still refresh it depending on chunk processing order. + assert saved["function_foo_12"]["pseudo"] in {"cached pseudo", "NEW"} + assert saved["function_bar_20"]["pseudo"] == "NEW" + assert saved["function_foo_12"]["tags"] def test_smart_reindex_does_not_reuse_when_info_changes(tmp_path, monkeypatch): @@ -124,31 +430,39 @@ def test_smart_reindex_does_not_reuse_when_info_changes(tmp_path, monkeypatch): monkeypatch.setitem(sys.modules, "fastembed", SimpleNamespace(TextEmbedding=object)) - from scripts import ingest_code + from scripts.ingest import pipeline as ingest_pipeline + _patch_qdrant_models(monkeypatch, ingest_pipeline) # Avoid touching any caches. - monkeypatch.setattr(ingest_code, "get_cached_symbols", lambda fp: {}) - monkeypatch.setattr(ingest_code, "compare_symbol_changes", lambda a, b: ([], [])) - monkeypatch.setattr(ingest_code, "set_cached_pseudo", None) - monkeypatch.setattr(ingest_code, "set_cached_symbols", None) - monkeypatch.setattr(ingest_code, "set_cached_file_hash", None) + monkeypatch.setattr(ingest_pipeline, "get_cached_symbols", lambda fp: {}) + monkeypatch.setattr(ingest_pipeline, "compare_symbol_changes", lambda a, b: ([], [])) + monkeypatch.setattr(ingest_pipeline, "set_cached_pseudo", None) + monkeypatch.setattr(ingest_pipeline, "set_cached_symbols", None) + monkeypatch.setattr(ingest_pipeline, "set_cached_file_hash", None) # Force simple line chunking. monkeypatch.setenv("INDEX_MICRO_CHUNKS", "0") monkeypatch.setenv("INDEX_SEMANTIC_CHUNKS", "0") monkeypatch.setenv("USE_TREE_SITTER", "0") monkeypatch.setenv("REFRAG_MODE", "0") + _patch_symbol(monkeypatch, ingest_pipeline, name="hi") + _patch_smart_side_effects(monkeypatch, ingest_pipeline) + monkeypatch.setattr( + ingest_pipeline, + "should_process_pseudo_for_chunk", + lambda fp, ch, changed: (False, "", []), + ) # Make build_information return a value that won't match the stored record. old_info = "old-info" new_info = "new-info" - monkeypatch.setattr(ingest_code, "build_information", lambda *a, **k: new_info) + monkeypatch.setattr(ingest_pipeline, "build_information", lambda *a, **k: new_info) code = "def hi():\n return 1\n" fp = tmp_path / "x.py" fp.write_text(code, encoding="utf-8") - chunk = ingest_code.chunk_lines(code, max_lines=120, overlap=20)[0] + chunk = ingest_pipeline.chunk_lines(code, max_lines=120, overlap=20)[0] code_text = chunk["text"] dense_key = "dense" @@ -166,7 +480,7 @@ def test_smart_reindex_does_not_reuse_when_info_changes(tmp_path, monkeypatch): "start_line": 1, }, }, - vector={dense_key: reused_dense, ingest_code.LEX_VECTOR_NAME: [0.0] * ingest_code.LEX_VECTOR_DIM}, + vector={dense_key: reused_dense, ingest_pipeline.LEX_VECTOR_NAME: [0.0] * ingest_pipeline.LEX_VECTOR_DIM}, ) class FakeClient: @@ -181,13 +495,13 @@ def scroll(self, **kwargs): def fake_upsert_points(_client, _collection, points): captured["points"] = points - monkeypatch.setattr(ingest_code, "upsert_points", fake_upsert_points) - monkeypatch.setattr(ingest_code, "delete_points_by_path", lambda *a, **k: None) + monkeypatch.setattr(ingest_pipeline, "upsert_points", fake_upsert_points) + monkeypatch.setattr(ingest_pipeline, "delete_points_by_path", lambda *a, **k: None) embedded_vec = [9.9, 8.8] - monkeypatch.setattr(ingest_code, "embed_batch", lambda _model, texts: [embedded_vec for _ in texts]) + monkeypatch.setattr(ingest_pipeline, "embed_batch", lambda _model, texts: [embedded_vec for _ in texts]) - status = ingest_code.process_file_with_smart_reindexing( + status = ingest_pipeline.process_file_with_smart_reindexing( file_path=Path(fp), text=code, language="python", @@ -210,28 +524,36 @@ def test_smart_reindex_unnamed_reuse_requires_dense_vector(tmp_path, monkeypatch monkeypatch.setitem(sys.modules, "fastembed", SimpleNamespace(TextEmbedding=object)) - from scripts import ingest_code + from scripts.ingest import pipeline as ingest_pipeline + _patch_qdrant_models(monkeypatch, ingest_pipeline) # Avoid touching any caches. - monkeypatch.setattr(ingest_code, "get_cached_symbols", lambda fp: {}) - monkeypatch.setattr(ingest_code, "compare_symbol_changes", lambda a, b: ([], [])) - monkeypatch.setattr(ingest_code, "set_cached_pseudo", None) - monkeypatch.setattr(ingest_code, "set_cached_symbols", None) - monkeypatch.setattr(ingest_code, "set_cached_file_hash", None) + monkeypatch.setattr(ingest_pipeline, "get_cached_symbols", lambda fp: {}) + monkeypatch.setattr(ingest_pipeline, "compare_symbol_changes", lambda a, b: ([], [])) + monkeypatch.setattr(ingest_pipeline, "set_cached_pseudo", None) + monkeypatch.setattr(ingest_pipeline, "set_cached_symbols", None) + monkeypatch.setattr(ingest_pipeline, "set_cached_file_hash", None) # Force simple line chunking. monkeypatch.setenv("INDEX_MICRO_CHUNKS", "0") monkeypatch.setenv("INDEX_SEMANTIC_CHUNKS", "0") monkeypatch.setenv("USE_TREE_SITTER", "0") monkeypatch.setenv("REFRAG_MODE", "0") + _patch_symbol(monkeypatch, ingest_pipeline, name="hi") + _patch_smart_side_effects(monkeypatch, ingest_pipeline) + monkeypatch.setattr( + ingest_pipeline, + "should_process_pseudo_for_chunk", + lambda fp, ch, changed: (False, "", []), + ) code = "def hi():\n return 1\n" fp = tmp_path / "x.py" fp.write_text(code, encoding="utf-8") - chunk = ingest_code.chunk_lines(code, max_lines=120, overlap=20)[0] + chunk = ingest_pipeline.chunk_lines(code, max_lines=120, overlap=20)[0] code_text = chunk["text"] - info_text = ingest_code.build_information( + info_text = ingest_pipeline.build_information( "python", Path(fp), chunk["start"], @@ -253,8 +575,8 @@ def test_smart_reindex_unnamed_reuse_requires_dense_vector(tmp_path, monkeypatch }, # Only lex/mini present: should not be reused as dense. vector={ - ingest_code.LEX_VECTOR_NAME: [0.0] * ingest_code.LEX_VECTOR_DIM, - ingest_code.MINI_VECTOR_NAME: [0.0] * ingest_code.MINI_VEC_DIM, + ingest_pipeline.LEX_VECTOR_NAME: [0.0] * ingest_pipeline.LEX_VECTOR_DIM, + ingest_pipeline.MINI_VECTOR_NAME: [0.0] * ingest_pipeline.MINI_VEC_DIM, }, ) @@ -270,13 +592,13 @@ def scroll(self, **kwargs): def fake_upsert_points(_client, _collection, points): captured["points"] = points - monkeypatch.setattr(ingest_code, "upsert_points", fake_upsert_points) - monkeypatch.setattr(ingest_code, "delete_points_by_path", lambda *a, **k: None) + monkeypatch.setattr(ingest_pipeline, "upsert_points", fake_upsert_points) + monkeypatch.setattr(ingest_pipeline, "delete_points_by_path", lambda *a, **k: None) embedded_vec = [7.7, 6.6] - monkeypatch.setattr(ingest_code, "embed_batch", lambda _model, texts: [embedded_vec for _ in texts]) + monkeypatch.setattr(ingest_pipeline, "embed_batch", lambda _model, texts: [embedded_vec for _ in texts]) - status = ingest_code.process_file_with_smart_reindexing( + status = ingest_pipeline.process_file_with_smart_reindexing( file_path=Path(fp), text=code, language="python", @@ -291,3 +613,43 @@ def fake_upsert_points(_client, _collection, points): assert len(captured["points"]) == 1 out_vec = captured["points"][0].vector assert out_vec == embedded_vec + + +def test_smart_reindex_no_symbol_changes_falls_back_without_hash_cache(tmp_path, monkeypatch): + monkeypatch.setitem(sys.modules, "fastembed", SimpleNamespace(TextEmbedding=object)) + + from scripts.ingest import pipeline as ingest_pipeline + _patch_qdrant_models(monkeypatch, ingest_pipeline) + + code = "def hi():\n return 1\n" + fp = tmp_path / "x.py" + fp.write_text(code, encoding="utf-8") + + monkeypatch.setattr( + ingest_pipeline, + "extract_symbols_with_tree_sitter", + lambda _fp: {"function_hi_1": {"name": "hi", "type": "function", "start_line": 1}}, + ) + monkeypatch.setattr( + ingest_pipeline, + "get_cached_symbols", + lambda _fp: {"function_hi_1": {"name": "hi", "type": "function", "start_line": 1}}, + ) + monkeypatch.setattr(ingest_pipeline, "compare_symbol_changes", lambda *_: ([], [])) + monkeypatch.setattr(ingest_pipeline, "get_cached_file_hash", lambda *_: None) + set_cached_file_hash = MagicMock() + monkeypatch.setattr(ingest_pipeline, "set_cached_file_hash", set_cached_file_hash) + + status = ingest_pipeline.process_file_with_smart_reindexing( + file_path=Path(fp), + text=code, + language="python", + client=MagicMock(), + current_collection="c", + per_file_repo="r", + model=object(), + vector_name="dense", + ) + + assert status == "failed" + set_cached_file_hash.assert_not_called() diff --git a/tests/test_staging_lifecycle.py b/tests/test_staging_lifecycle.py index 01734e32..69aaba3f 100644 --- a/tests/test_staging_lifecycle.py +++ b/tests/test_staging_lifecycle.py @@ -2,7 +2,7 @@ import os from pathlib import Path from typing import Any, Dict, Optional - +##### import pytest from fastapi.testclient import TestClient @@ -497,6 +497,7 @@ def failing_spawn(**kwargs): def test_admin_staging_endpoints_exercise_http_layer(monkeypatch: pytest.MonkeyPatch): + from scripts import indexing_admin from scripts import upload_service calls = {"start": 0, "activate": 0, "abort": 0} @@ -518,11 +519,11 @@ def fake_activate(**kwargs): def fake_abort(**kwargs): calls["abort"] += 1 - monkeypatch.setattr(upload_service, "start_staging_rebuild", fake_start) - monkeypatch.setattr(upload_service, "activate_staging_rebuild", fake_activate) - monkeypatch.setattr(upload_service, "abort_staging_rebuild", fake_abort) + monkeypatch.setattr(indexing_admin, "start_staging_rebuild", fake_start) + monkeypatch.setattr(indexing_admin, "activate_staging_rebuild", fake_activate) + monkeypatch.setattr(indexing_admin, "abort_staging_rebuild", fake_abort) monkeypatch.setattr( - upload_service, + indexing_admin, "resolve_collection_root", lambda **kwargs: ("/fake/root", "repo1"), ) @@ -542,6 +543,61 @@ def fake_abort(**kwargs): assert calls["abort"] == 1 +def test_admin_copy_endpoint_reports_graph_clone_in_redirect(monkeypatch: pytest.MonkeyPatch): + import sys + import types + from urllib.parse import parse_qs, urlparse + + from scripts import upload_service + + monkeypatch.setattr(upload_service, "AUTH_ENABLED", True) + monkeypatch.setattr(upload_service, "_require_admin_session", lambda request: {"user_id": "admin"}) + monkeypatch.setattr(upload_service, "WORK_DIR", "/fake/work") + monkeypatch.setenv("WORK_DIR", "/fake/work") + + def fake_copy_collection_qdrant(**kwargs): + assert kwargs.get("source") == "src" + assert kwargs.get("target") == "dst" + return "dst" + + class _FakeQdrantClient: + def get_collection(self, collection_name: str): + if collection_name == "dst_graph": + return {"name": collection_name} + raise RuntimeError("not found") + + def __enter__(self): + return self + + def __exit__(self, *args): + return None + + monkeypatch.setitem( + sys.modules, + "scripts.collection_admin", + types.SimpleNamespace(copy_collection_qdrant=fake_copy_collection_qdrant), + ) + monkeypatch.setitem( + sys.modules, + "scripts.qdrant_client_manager", + types.SimpleNamespace(pooled_qdrant_client=lambda **kwargs: _FakeQdrantClient()), + ) + + client = TestClient(upload_service.app) + resp = client.post( + "/admin/staging/copy", + data={"collection": "src", "target": "dst", "overwrite": ""}, + follow_redirects=False, + ) + assert resp.status_code == 302 + loc = resp.headers.get("location") or "" + parsed = urlparse(loc) + qs = parse_qs(parsed.query) + assert qs.get("copied") == ["src"] + assert qs.get("new") == ["dst"] + assert qs.get("graph_copied") == ["1"] + + def test_watcher_collection_resolution_prefers_serving_state_when_staging_enabled(monkeypatch: pytest.MonkeyPatch, tmp_path: Path): from scripts.watch_index_core import utils as watch_utils @@ -650,7 +706,9 @@ class _Proc: env = captured["env"] assert env["BASE_ONLY"] == "system" assert env["COLLECTION_NAME"] == "primary-coll" - assert "CTXCE_FORCE_COLLECTION_NAME" not in env + # Admin-spawned ingests should never enumerate `/work/*` in multi-repo mode; + # force exact collection/root handling even when no explicit overrides are provided. + assert env.get("CTXCE_FORCE_COLLECTION_NAME") == "1" def test_promote_pending_env_without_pending_config(staging_workspace: dict): @@ -738,31 +796,6 @@ def test_resolve_codebase_root_fallbacks_to_parent(monkeypatch: pytest.MonkeyPat assert resolved_parent == codebase_root -def test_admin_abort_endpoint_falls_back_to_clear_when_abort_helper_missing(monkeypatch: pytest.MonkeyPatch): - from scripts import upload_service - - calls = {"clear": []} - - monkeypatch.setattr(upload_service, "_require_admin_session", lambda request: {"user_id": "admin"}) - monkeypatch.setattr(upload_service, "abort_staging_rebuild", None) - monkeypatch.setattr( - upload_service, - "clear_staging_collection", - lambda workspace_path, repo_name: calls["clear"].append((workspace_path, repo_name)), - ) - monkeypatch.setattr( - upload_service, - "resolve_collection_root", - lambda **kwargs: ("/fake/root", "repo1"), - ) - - client = TestClient(upload_service.app) - - resp = client.post("/admin/staging/abort", data={"collection": "coll1"}, follow_redirects=False) - assert resp.status_code == 302 - assert calls["clear"] == [("/fake/root", "repo1")] - - def test_watcher_collection_reuse_logical_repo(monkeypatch: pytest.MonkeyPatch, tmp_path: Path): from scripts.watch_index_core import utils as watch_utils diff --git a/tests/test_subprocess_hybrid_smoke.py b/tests/test_subprocess_hybrid_smoke.py index 927a95e9..5badc1b3 100644 --- a/tests/test_subprocess_hybrid_smoke.py +++ b/tests/test_subprocess_hybrid_smoke.py @@ -55,7 +55,8 @@ def test_hybrid_cli_runs_basic(tmp_path, qdrant_container): env["EMBEDDING_MODEL"] = "BAAI/bge-base-en-v1.5" cmd = [ sys.executable, - "scripts/hybrid_search.py", + "-m", + "scripts.hybrid_search", "--query", "test", "--limit", diff --git a/tests/test_symbol_filename_boost.py b/tests/test_symbol_filename_boost.py deleted file mode 100644 index af33d2ef..00000000 --- a/tests/test_symbol_filename_boost.py +++ /dev/null @@ -1,104 +0,0 @@ -"""Tests for symbol and filename boosting logic in hybrid_search. - -Tests token-level matching for symbols (camelCase/snake_case splitting) and the -filename boost helper used by hybrid_search. -""" -import pytest -from scripts.hybrid_search import _split_ident -from scripts.rerank_recursive.utils import _compute_fname_boost - - -class TestSplitIdent: - """Test identifier splitting for camelCase/snake_case.""" - - def test_snake_case(self): - assert _split_ident("user_auth_handler") == ["user", "auth", "handler"] - - def test_camel_case(self): - assert _split_ident("userAuthHandler") == ["user", "auth", "handler"] - - def test_pascal_case(self): - assert _split_ident("UserAuthHandler") == ["user", "auth", "handler"] - - def test_acronym_handling(self): - # XMLParser -> xml, parser - result = _split_ident("XMLParser") - assert "xml" in result - assert "parser" in result - - def test_mixed_separators(self): - result = _split_ident("user_authHandler") - assert "user" in result - assert "auth" in result - assert "handler" in result - - def test_filters_stopwords(self): - # 'the' should be filtered - result = _split_ident("getTheValue") - assert "the" not in result - assert "get" in result - assert "value" in result - - def test_empty_string(self): - assert _split_ident("") == [] - - def test_single_word(self): - assert _split_ident("handler") == ["handler"] - - -class TestSymbolPartMatching: - """Test that symbol parts are correctly matched against queries.""" - - def test_symbol_parts_extracted(self): - """Verify split parts can match query tokens.""" - symbol = "getUserById" - parts = set(p.lower() for p in _split_ident(symbol) if len(p) >= 2) - assert "get" in parts - assert "user" in parts - # "by" is filtered as stopword, "id" should remain - assert "id" in parts - - def test_symbol_parts_from_snake_case(self): - symbol = "calculate_total_price" - parts = set(p.lower() for p in _split_ident(symbol) if len(p) >= 2) - assert "calculate" in parts - assert "total" in parts - assert "price" in parts - - -class TestFilenameBoostLogic: - """Test filename boost behavior.""" - - def test_fname_boost_matches_filename_tokens(self): - """Two token matches should trigger a boost.""" - q = "authentication service" - cand = {"path": "services/AuthenticationService.ts"} - assert _compute_fname_boost(q, cand, 0.1) > 0 - - def test_fname_boost_requires_two_matches(self): - """Single-token queries should not trigger (noise control).""" - q = "db" - cand = {"path": "lib/db.py"} - assert _compute_fname_boost(q, cand, 0.1) == 0.0 - - -class TestBoostIntegration: - """Integration tests for the boost logic patterns.""" - - def test_pattern_symbol_equality_boost(self): - """Test the pattern used in hybrid_search for symbol equality.""" - symbol = "processUserData" - sym = symbol.lower() - sym_parts = set(p.lower() for p in _split_ident(symbol) if len(p) >= 2) - - # Query "user" should match via sym_parts - query = "user" - ql = query.lower() - matches = ql == sym or ql in sym_parts - assert matches, "Query 'user' should match symbol part" - - def test_pattern_filename_boost(self): - """Filename boost uses the production-grade matcher.""" - q = "authentication service" - cand = {"path": "services/AuthenticationService.ts"} - assert _compute_fname_boost(q, cand, 0.1) > 0 diff --git a/tests/test_symbol_graph_tool.py b/tests/test_symbol_graph_tool.py index e148fc17..1d3861d0 100644 --- a/tests/test_symbol_graph_tool.py +++ b/tests/test_symbol_graph_tool.py @@ -2,21 +2,43 @@ @pytest.mark.asyncio -async def test_symbol_graph_under_uses_path_prefix_matchvalue(): - # Import internal helper to validate filter construction without needing a real Qdrant instance. - from qdrant_client import models as qmodels +async def test_symbol_graph_under_filters_results_by_recursive_scope(): + # Validate that under applies as recursive subtree filter (user-facing scope). from scripts.mcp_impl import symbol_graph as sg - captured = {} + class _Pt: + def __init__(self, pid, path): + self.id = pid + self.payload = { + "metadata": { + "repo": "repo", + "path": path, + "start_line": 1, + "end_line": 2, + "symbol": "f", + "symbol_path": "f", + "language": "python", + "calls": ["foo"], + } + } class FakeClient: + def __init__(self): + self.scroll_filters = [] + def scroll(self, *, collection_name, scroll_filter, limit, with_payload, with_vectors): - captured["collection_name"] = collection_name - captured["scroll_filter"] = scroll_filter - return ([], None) + self.scroll_filters.append(scroll_filter) + return ( + [ + _Pt("1", "/work/repo/scripts/a.py"), + _Pt("2", "/work/repo/tests/b.py"), + ], + None, + ) - await sg._query_array_field( # type: ignore[attr-defined] - client=FakeClient(), + client = FakeClient() + out = await sg._query_array_field( # type: ignore[attr-defined] + client=client, collection="codebase", field_key="metadata.calls", value="foo", @@ -25,15 +47,29 @@ def scroll(self, *, collection_name, scroll_filter, limit, with_payload, with_ve under=sg._norm_under("scripts"), # type: ignore[attr-defined] ) - flt = captured.get("scroll_filter") - assert isinstance(flt, qmodels.Filter) - must = list(flt.must or []) - keys = [getattr(c, "key", None) for c in must] - assert "metadata.path_prefix" in keys - - # Ensure it's an exact match (MatchValue), not substring (MatchText) - cond = next(c for c in must if getattr(c, "key", None) == "metadata.path_prefix") - assert isinstance(cond.match, qmodels.MatchValue) - assert cond.match.value == "/work/scripts" - + # Validate _query_array_field forwards language/value constraints to scroll_filter. + assert client.scroll_filters, "Expected at least one scroll() call" + first_filter = client.scroll_filters[0] + first_must = list(getattr(first_filter, "must", []) or []) + assert any( + getattr(cond, "key", None) == "metadata.calls" + and getattr(getattr(cond, "match", None), "any", None) == ["foo"] + for cond in first_must + ) + assert any( + getattr(cond, "key", None) == "metadata.language" + and getattr(getattr(cond, "match", None), "value", None) == "python" + for cond in first_must + ) + assert any( + any( + getattr(cond, "key", None) == "metadata.calls" + and getattr(getattr(cond, "match", None), "text", None) == "foo" + for cond in list(getattr(sf, "must", []) or []) + ) + for sf in client.scroll_filters + ), "Expected MatchText fallback filter for metadata.calls" + paths = {r.get("path") for r in out} + assert "/work/repo/scripts/a.py" in paths + assert "/work/repo/tests/b.py" not in paths diff --git a/tests/test_tier2_fallback.py b/tests/test_tier2_fallback.py index eb01bad3..c9595e19 100644 --- a/tests/test_tier2_fallback.py +++ b/tests/test_tier2_fallback.py @@ -39,12 +39,18 @@ def embed(self, texts): @pytest.mark.asyncio async def test_tier2_fallback_unconditional_with_language_filter(tmp_path, monkeypatch, qdrant_container): # Env for services - os.environ["QDRANT_URL"] = qdrant_container - os.environ["COLLECTION_NAME"] = f"test-{uuid.uuid4().hex[:8]}" - os.environ["USE_TREE_SITTER"] = "0" - os.environ["HYBRID_IN_PROCESS"] = "1" - os.environ["EMBEDDING_MODEL"] = "fake" - os.environ["REFRAG_GATE_FIRST"] = "1" # ensure Tier-1 gate-first path is active + monkeypatch.setenv("QDRANT_URL", qdrant_container) + monkeypatch.setenv("COLLECTION_NAME", f"test-{uuid.uuid4().hex[:8]}") + monkeypatch.setenv("USE_TREE_SITTER", "0") + monkeypatch.setenv("HYBRID_IN_PROCESS", "1") + monkeypatch.setenv("EMBEDDING_MODEL", "fake") + monkeypatch.setenv("REFRAG_GATE_FIRST", "1") # ensure Tier-1 gate-first path is active + monkeypatch.setenv("REFRAG_RUNTIME", "llamacpp") + monkeypatch.setenv("CTX_MULTI_COLLECTION", "0") + monkeypatch.setenv("CTX_DOC_PASS", "0") + monkeypatch.setenv("CTX_DOC_TOP_FALLBACK", "0") + monkeypatch.setenv("HYBRID_EXPAND", "0") + monkeypatch.setenv("SEMANTIC_EXPANSION_ENABLED", "0") # Stub embeddings everywhere (FakeEmbedder produces 32-dim vectors) monkeypatch.setattr(ing, "TextEmbedding", lambda *a, **k: FakeEmbedder("fake")) diff --git a/tests/test_toon_encoder.py b/tests/test_toon_encoder.py index 326184e3..4e9c26ac 100644 --- a/tests/test_toon_encoder.py +++ b/tests/test_toon_encoder.py @@ -20,6 +20,10 @@ _encode_value, _is_uniform_array_of_objects, ) +from scripts.mcp_impl.toon import ( + _format_context_results_as_toon, + _format_results_as_toon, +) class TestFeatureFlags: @@ -256,13 +260,8 @@ class TestMCPIntegration: def test_should_use_toon_explicit_param(self, monkeypatch): """Test explicit output_format parameter takes precedence.""" - # Import the helpers from mcp_indexer_server monkeypatch.delenv("TOON_ENABLED", raising=False) - # We need to test the helper functions directly - # Since they're in mcp_indexer_server, we'll test the logic here - from scripts.toon_encoder import is_toon_enabled - # When TOON_ENABLED is not set, default is False assert is_toon_enabled() is False @@ -290,13 +289,17 @@ def test_format_results_as_toon_structure(self): assert "/src/main.py,10,20" in toon_output assert "/src/utils.py,5,15" in toon_output - def test_toon_replaces_results_array(self): - """Test that TOON formatting replaces JSON array with TOON string.""" + formatted = _format_results_as_toon(response.copy(), compact=True) + assert formatted["results"] == response["results"] + assert formatted["output_format"] == "toon" + assert formatted["text"] == toon_output + + def test_toon_encoder_returns_string(self): + """Test that the low-level TOON encoder returns a string.""" results = [ {"path": "/src/main.py", "start_line": 10, "end_line": 20}, ] - # When TOON is applied, results becomes a string toon_output = encode_search_results(results, compact=True) assert isinstance(toon_output, str) assert "results[1]{path,start_line,end_line}:" in toon_output @@ -416,19 +419,16 @@ class TestFormatContextResultsAsToon: def test_format_empty_results_adds_marker(self): """Test that empty results still get output_format marker.""" - from scripts.mcp_indexer_server import _format_context_results_as_toon - response = {"results": [], "total": 0} result = _format_context_results_as_toon(response.copy()) assert result["output_format"] == "toon" + assert result["results"] == [] # Empty array per spec: key[0]: - assert result["results"] == "results[0]:" + assert result["text"] == "results[0]:" def test_format_mixed_results(self): """Test formatting mixed code/memory results.""" - from scripts.mcp_indexer_server import _format_context_results_as_toon - response = { "results": [ {"source": "code", "path": "/src/api.py", "start_line": 1, "end_line": 10}, @@ -439,14 +439,13 @@ def test_format_mixed_results(self): result = _format_context_results_as_toon(response.copy()) assert result["output_format"] == "toon" - assert isinstance(result["results"], str) - assert "code[1]" in result["results"] - assert "memory[1]" in result["results"] + assert isinstance(result["results"], list) + assert isinstance(result["text"], str) + assert "code[1]" in result["text"] + assert "memory[1]" in result["text"] def test_format_preserves_other_fields(self): """Test that formatting preserves non-results fields.""" - from scripts.mcp_indexer_server import _format_context_results_as_toon - response = { "results": [{"source": "code", "path": "/a.py", "start_line": 1, "end_line": 5}], "total": 1, @@ -476,7 +475,7 @@ def test_encode_with_snippet_field(self): assert "def main()" in output def test_encode_with_information_field(self): - """Test that info_request's information field is included.""" + """Test that the information field is included.""" results = [ {"path": "/src/auth.py", "start_line": 1, "end_line": 50, "score": 0.9, "information": "Authentication handler at /src/auth.py:1-50", @@ -529,4 +528,3 @@ def test_context_results_with_extra_memory_fields(self): assert "id" in output assert "created_at" in output assert "tags" in output - diff --git a/tests/test_upload_client_ignore_cleanup.py b/tests/test_upload_client_ignore_cleanup.py new file mode 100644 index 00000000..06ac026c --- /dev/null +++ b/tests/test_upload_client_ignore_cleanup.py @@ -0,0 +1,744 @@ +import hashlib +import importlib +import json +from pathlib import Path +from unittest.mock import MagicMock + +import pytest + + +@pytest.mark.parametrize( + "mod_name", + ["scripts.remote_upload_client", "scripts.standalone_upload_client"], +) +def test_remote_upload_config_does_not_generate_collection_name(monkeypatch, tmp_path, mod_name): + mod = importlib.import_module(mod_name) + workspace = tmp_path / "repo" + workspace.mkdir() + + monkeypatch.setattr(mod, "_compute_logical_repo_id", lambda _path: "fs:test") + + config = mod.get_remote_config(str(workspace)) + + assert config["collection_name"] is None + + +def _exercise_ignored_path_cleanup(mod_name: str, monkeypatch, tmp_path: Path) -> None: + mod = importlib.import_module(mod_name) + + workspace = tmp_path / "repo" + ignored = workspace / "dev-workspace" / "nested.py" + ignored.parent.mkdir(parents=True, exist_ok=True) + ignored.write_text("print('dogfood')\n", encoding="utf-8") + + monkeypatch.setenv("DEV_REMOTE_MODE", "1") + monkeypatch.setattr(mod, "get_cached_file_hash", lambda path, repo_name=None, **_kwargs: "abc123") + monkeypatch.setattr(mod, "set_cached_file_hash", lambda *a, **k: None) + + client = mod.RemoteUploadClient( + upload_endpoint="http://localhost:8004", + workspace_path=str(workspace), + collection_name="test-coll", + ) + + changes = client.detect_file_changes([ignored]) + + assert ignored in changes["deleted"] + assert not changes["created"] + assert not changes["updated"] + assert not changes["moved"] + + +def test_remote_upload_client_marks_ignored_cached_paths_deleted(monkeypatch, tmp_path): + _exercise_ignored_path_cleanup("scripts.remote_upload_client", monkeypatch, tmp_path) + + +def test_standalone_upload_client_marks_ignored_cached_paths_deleted(monkeypatch, tmp_path): + _exercise_ignored_path_cleanup("scripts.standalone_upload_client", monkeypatch, tmp_path) + + +def _exercise_force_mode_cleanup(mod_name: str, monkeypatch, tmp_path: Path) -> None: + mod = importlib.import_module(mod_name) + + workspace = tmp_path / "repo" + workspace.mkdir(parents=True, exist_ok=True) + current = workspace / "app.py" + current.write_text("print('current')\n", encoding="utf-8") + + stale_ignored = workspace / "dev-workspace" / "nested.py" + stale_ignored.parent.mkdir(parents=True, exist_ok=True) + stale_ignored.write_text("print('stale')\n", encoding="utf-8") + + monkeypatch.setenv("DEV_REMOTE_MODE", "1") + monkeypatch.setattr(mod, "get_all_cached_paths", lambda repo_name=None, **_kwargs: [str(stale_ignored)]) + monkeypatch.setattr(mod, "get_cached_file_hash", lambda path, repo_name=None, **_kwargs: "abc123") + monkeypatch.setattr(mod, "set_cached_file_hash", lambda *a, **k: None) + + client = mod.RemoteUploadClient( + upload_endpoint="http://localhost:8004", + workspace_path=str(workspace), + collection_name="test-coll", + ) + + changes = client.build_force_changes([current]) + + assert current in changes["created"] + assert stale_ignored in changes["deleted"] + assert not changes["updated"] + assert not changes["moved"] + + +def test_remote_upload_client_force_mode_keeps_creates_and_deletes_ignored_cached_paths(monkeypatch, tmp_path): + _exercise_force_mode_cleanup("scripts.remote_upload_client", monkeypatch, tmp_path) + + +def test_standalone_upload_client_force_mode_keeps_creates_and_deletes_ignored_cached_paths(monkeypatch, tmp_path): + _exercise_force_mode_cleanup("scripts.standalone_upload_client", monkeypatch, tmp_path) + + +def _exercise_force_mode_excludes_ignored_current_files(mod_name: str, monkeypatch, tmp_path: Path) -> None: + mod = importlib.import_module(mod_name) + + workspace = tmp_path / "repo" + workspace.mkdir(parents=True, exist_ok=True) + current = workspace / "app.py" + current.write_text("print('current')\n", encoding="utf-8") + + ignored_current = workspace / "dev-workspace" / "ignored.py" + ignored_current.parent.mkdir(parents=True, exist_ok=True) + ignored_current.write_text("print('ignored')\n", encoding="utf-8") + + monkeypatch.setenv("DEV_REMOTE_MODE", "1") + monkeypatch.setattr(mod, "get_all_cached_paths", lambda repo_name=None, **_kwargs: []) + monkeypatch.setattr(mod, "get_cached_file_hash", lambda path, repo_name=None, **_kwargs: None) + monkeypatch.setattr(mod, "set_cached_file_hash", lambda *a, **k: None) + + client = mod.RemoteUploadClient( + upload_endpoint="http://localhost:8004", + workspace_path=str(workspace), + collection_name="test-coll", + ) + + changes = client.build_force_changes([current, ignored_current]) + + assert current in changes["created"] + assert ignored_current not in changes["created"] + assert ignored_current in changes["deleted"] + assert not changes["updated"] + assert not changes["moved"] + + +def test_remote_upload_client_force_mode_excludes_ignored_current_files(monkeypatch, tmp_path): + _exercise_force_mode_excludes_ignored_current_files( + "scripts.remote_upload_client", + monkeypatch, + tmp_path, + ) + + +def test_standalone_upload_client_force_mode_excludes_ignored_current_files(monkeypatch, tmp_path): + _exercise_force_mode_excludes_ignored_current_files( + "scripts.standalone_upload_client", + monkeypatch, + tmp_path, + ) + + +def _exercise_force_mode_dev_workspace_cleanup_without_cache(mod_name: str, monkeypatch, tmp_path: Path) -> None: + mod = importlib.import_module(mod_name) + + workspace = tmp_path / "repo" + workspace.mkdir(parents=True, exist_ok=True) + current = workspace / "app.py" + current.write_text("print('current')\n", encoding="utf-8") + + mirrored = workspace / "dev-workspace" / "nested" / "stale.py" + mirrored.parent.mkdir(parents=True, exist_ok=True) + mirrored.write_text("print('stale')\n", encoding="utf-8") + + monkeypatch.setenv("DEV_REMOTE_MODE", "1") + monkeypatch.setattr(mod, "get_all_cached_paths", lambda repo_name=None, **_kwargs: []) + monkeypatch.setattr(mod, "get_cached_file_hash", lambda path, repo_name=None, **_kwargs: None) + monkeypatch.setattr(mod, "set_cached_file_hash", lambda *a, **k: None) + + client = mod.RemoteUploadClient( + upload_endpoint="http://localhost:8004", + workspace_path=str(workspace), + collection_name="test-coll", + ) + + changes = client.build_force_changes([current]) + + assert current in changes["created"] + assert mirrored in changes["deleted"] + assert not changes["updated"] + assert not changes["moved"] + + +def test_remote_upload_client_force_mode_deletes_dev_workspace_without_cache(monkeypatch, tmp_path): + _exercise_force_mode_dev_workspace_cleanup_without_cache("scripts.remote_upload_client", monkeypatch, tmp_path) + + +def test_standalone_upload_client_force_mode_deletes_dev_workspace_without_cache(monkeypatch, tmp_path): + _exercise_force_mode_dev_workspace_cleanup_without_cache("scripts.standalone_upload_client", monkeypatch, tmp_path) + + +def _exercise_plan_skip_avoids_bundle_upload(mod_name: str, monkeypatch, tmp_path: Path) -> None: + mod = importlib.import_module(mod_name) + + workspace = tmp_path / "repo" + workspace.mkdir(parents=True, exist_ok=True) + current = workspace / "app.py" + current.write_text("print('current')\n", encoding="utf-8") + + client = mod.RemoteUploadClient( + upload_endpoint="http://localhost:8004", + workspace_path=str(workspace), + collection_name="test-coll", + ) + + monkeypatch.setattr( + client, + "_plan_delta_upload", + lambda changes: { + "needed_files": {"created": [], "updated": [], "moved": []}, + "operation_counts_preview": { + "created": 0, + "updated": 0, + "deleted": 0, + "moved": 0, + "skipped": 1, + "skipped_hash_match": 1, + "failed": 0, + }, + "needed_size_bytes": 0, + }, + ) + monkeypatch.setattr(client, "create_delta_bundle", lambda *a, **k: (_ for _ in ()).throw(RuntimeError("should not bundle"))) + monkeypatch.setattr(client, "upload_bundle", lambda *a, **k: (_ for _ in ()).throw(RuntimeError("should not upload"))) + + assert client.process_changes_and_upload( + { + "created": [current], + "updated": [], + "deleted": [], + "moved": [], + "unchanged": [], + } + ) is True + assert client.last_upload_result["outcome"] == "skipped_by_plan" + + +def test_remote_upload_client_plan_skip_avoids_bundle_upload(monkeypatch, tmp_path): + _exercise_plan_skip_avoids_bundle_upload("scripts.remote_upload_client", monkeypatch, tmp_path) + + +def test_standalone_upload_client_plan_skip_avoids_bundle_upload(monkeypatch, tmp_path): + _exercise_plan_skip_avoids_bundle_upload("scripts.standalone_upload_client", monkeypatch, tmp_path) + + +def _exercise_detect_file_changes_does_not_persist_hash(mod_name: str, monkeypatch, tmp_path: Path) -> None: + mod = importlib.import_module(mod_name) + + workspace = tmp_path / "repo" + workspace.mkdir(parents=True, exist_ok=True) + current = workspace / "app.py" + current.write_text("print('current')\n", encoding="utf-8") + + set_hash = MagicMock() + monkeypatch.setattr(mod, "get_cached_file_hash", lambda path, repo_name=None, **_kwargs: "oldhash") + monkeypatch.setattr(mod, "set_cached_file_hash", set_hash) + + client = mod.RemoteUploadClient( + upload_endpoint="http://localhost:8004", + workspace_path=str(workspace), + collection_name="test-coll", + ) + + changes = client.detect_file_changes([current]) + + assert current in changes["updated"] + set_hash.assert_not_called() + + +def test_remote_upload_client_detect_file_changes_does_not_persist_hash(monkeypatch, tmp_path): + _exercise_detect_file_changes_does_not_persist_hash( + "scripts.remote_upload_client", monkeypatch, tmp_path + ) + + +def test_standalone_upload_client_detect_file_changes_does_not_persist_hash(monkeypatch, tmp_path): + _exercise_detect_file_changes_does_not_persist_hash( + "scripts.standalone_upload_client", monkeypatch, tmp_path + ) + + +def _exercise_plan_skip_finalizes_hash(mod_name: str, monkeypatch, tmp_path: Path) -> None: + mod = importlib.import_module(mod_name) + + workspace = tmp_path / "repo" + workspace.mkdir(parents=True, exist_ok=True) + current = workspace / "app.py" + current.write_text("print('current')\n", encoding="utf-8") + + client = mod.RemoteUploadClient( + upload_endpoint="http://localhost:8004", + workspace_path=str(workspace), + collection_name="test-coll", + ) + + set_hash = MagicMock() + monkeypatch.setattr(mod, "set_cached_file_hash", set_hash) + monkeypatch.setattr( + client, + "_plan_delta_upload", + lambda changes: { + "needed_files": {"created": [], "updated": [], "moved": []}, + "operation_counts_preview": { + "created": 0, + "updated": 0, + "deleted": 0, + "moved": 0, + "skipped": 1, + "skipped_hash_match": 1, + "failed": 0, + }, + "needed_size_bytes": 0, + }, + ) + monkeypatch.setattr(client, "create_delta_bundle", lambda *a, **k: (_ for _ in ()).throw(RuntimeError("should not bundle"))) + monkeypatch.setattr(client, "upload_bundle", lambda *a, **k: (_ for _ in ()).throw(RuntimeError("should not upload"))) + + assert client.process_changes_and_upload( + { + "created": [], + "updated": [current], + "deleted": [], + "moved": [], + "unchanged": [], + } + ) is True + assert client.last_upload_result["outcome"] == "skipped_by_plan" + set_hash.assert_called_once() + + +def test_remote_upload_client_plan_skip_finalizes_hash(monkeypatch, tmp_path): + _exercise_plan_skip_finalizes_hash( + "scripts.remote_upload_client", monkeypatch, tmp_path + ) + + +def test_standalone_upload_client_plan_skip_finalizes_hash(monkeypatch, tmp_path): + _exercise_plan_skip_finalizes_hash( + "scripts.standalone_upload_client", monkeypatch, tmp_path + ) + + +def _exercise_mixed_plan_finalizes_hash_matched_files(mod_name: str, monkeypatch, tmp_path: Path) -> None: + mod = importlib.import_module(mod_name) + + workspace = tmp_path / "repo" + workspace.mkdir(parents=True, exist_ok=True) + matched = workspace / "matched.py" + needed = workspace / "needed.py" + matched.write_text("print('matched')\n", encoding="utf-8") + needed.write_text("print('needed')\n", encoding="utf-8") + + client = mod.RemoteUploadClient( + upload_endpoint="http://localhost:8004", + workspace_path=str(workspace), + collection_name="test-coll", + ) + set_hash = MagicMock() + monkeypatch.setattr(mod, "set_cached_file_hash", set_hash) + monkeypatch.setattr( + client, + "_plan_delta_upload", + lambda changes: { + "needed_files": {"created": ["needed.py"], "updated": [], "moved": []}, + "operation_counts_preview": {"created": 1, "updated": 0, "deleted": 0, "moved": 0, "skipped": 1, "skipped_hash_match": 1, "failed": 0}, + "needed_size_bytes": needed.stat().st_size, + }, + ) + bundle_path = workspace / "bundle.tar.gz" + bundle_path.write_bytes(b"bundle") + monkeypatch.setattr(client, "create_delta_bundle", lambda *a, **k: (str(bundle_path), {"bundle_id": "b1", "total_size_bytes": 1})) + monkeypatch.setattr(client, "upload_bundle", lambda *a, **k: {"success": True, "processed_operations": {"created": 1}}) + + assert client.process_changes_and_upload( + {"created": [matched, needed], "updated": [], "deleted": [], "moved": [], "unchanged": []} + ) is True + assert {call.args[0] for call in set_hash.call_args_list} == { + str(matched.resolve()), + str(needed.resolve()), + } + + +def test_remote_upload_client_mixed_plan_finalizes_hash_matched_files(monkeypatch, tmp_path): + _exercise_mixed_plan_finalizes_hash_matched_files("scripts.remote_upload_client", monkeypatch, tmp_path) + + +def test_standalone_upload_client_mixed_plan_finalizes_hash_matched_files(monkeypatch, tmp_path): + _exercise_mixed_plan_finalizes_hash_matched_files("scripts.standalone_upload_client", monkeypatch, tmp_path) + + +def _exercise_upload_does_not_finalize_file_changed_during_upload( + mod_name: str, monkeypatch, tmp_path: Path +) -> None: + mod = importlib.import_module(mod_name) + + workspace = tmp_path / "repo" + workspace.mkdir(parents=True, exist_ok=True) + current = workspace / "app.py" + current.write_text("before\n", encoding="utf-8") + uploaded_hash = hashlib.sha1(current.read_bytes()).hexdigest() + + client = mod.RemoteUploadClient( + upload_endpoint="http://localhost:8004", + workspace_path=str(workspace), + collection_name="test-coll", + ) + set_hash = MagicMock() + monkeypatch.setattr(mod, "set_cached_file_hash", set_hash) + monkeypatch.setattr(client, "_plan_delta_upload", lambda _changes: None) + + bundle_path = workspace / "bundle.tar.gz" + + def create_bundle(*_args, **_kwargs): + client._last_expected_hashes = {"app.py": f"sha1:{uploaded_hash}"} + bundle_path.write_bytes(b"bundle") + return str(bundle_path), {"bundle_id": "b1", "total_size_bytes": 1} + + def upload_bundle(*_args, **_kwargs): + current.write_text("after\n", encoding="utf-8") + return {"success": True, "processed_operations": {"created": 1}} + + monkeypatch.setattr(client, "create_delta_bundle", create_bundle) + monkeypatch.setattr(client, "upload_bundle", upload_bundle) + + assert client.process_changes_and_upload( + { + "created": [current], + "updated": [], + "deleted": [], + "moved": [], + "unchanged": [], + } + ) is True + set_hash.assert_not_called() + + +def test_remote_upload_client_does_not_finalize_file_changed_during_upload(monkeypatch, tmp_path): + _exercise_upload_does_not_finalize_file_changed_during_upload( + "scripts.remote_upload_client", monkeypatch, tmp_path + ) + + +def test_standalone_upload_client_does_not_finalize_file_changed_during_upload(monkeypatch, tmp_path): + _exercise_upload_does_not_finalize_file_changed_during_upload( + "scripts.standalone_upload_client", monkeypatch, tmp_path + ) + + +def test_standalone_upload_client_plan_payload_prefixes_previous_hash(monkeypatch, tmp_path): + mod = importlib.import_module("scripts.standalone_upload_client") + + workspace = tmp_path / "repo" + workspace.mkdir(parents=True, exist_ok=True) + updated = workspace / "app.py" + updated.write_text("print('updated')\n", encoding="utf-8") + + client = mod.RemoteUploadClient( + upload_endpoint="http://localhost:8004", + workspace_path=str(workspace), + collection_name="test-coll", + ) + + monkeypatch.setattr(mod, "get_cached_file_hash", lambda path, repo_name=None: "abc123") + + payload = client._build_plan_payload( + { + "created": [], + "updated": [updated], + "deleted": [updated], + "moved": [], + } + ) + + updated_op = next(op for op in payload["operations"] if op["operation"] == "updated") + deleted_op = next(op for op in payload["operations"] if op["operation"] == "deleted") + assert updated_op["previous_hash"] == "sha1:abc123" + assert deleted_op["previous_hash"] == "sha1:abc123" + + +@pytest.mark.parametrize( + "mod_name", + ["scripts.remote_upload_client", "scripts.standalone_upload_client"], +) +def test_upload_client_status_passes_auth_session(monkeypatch, tmp_path, mod_name): + mod = importlib.import_module(mod_name) + workspace = tmp_path / "repo" + workspace.mkdir(parents=True) + client = mod.RemoteUploadClient( + upload_endpoint="http://localhost:8004", + workspace_path=str(workspace), + collection_name="test-coll", + ) + + response = MagicMock(status_code=200) + response.json.return_value = {"status": "ready"} + get_mock = MagicMock(return_value=response) + monkeypatch.setattr(client.session, "get", get_mock) + monkeypatch.setattr(mod, "get_auth_session", lambda _endpoint: "session-123") + + assert client.get_server_status()["success"] is True + params = get_mock.call_args.kwargs["params"] + assert params["workspace_path"] + assert params["session"] == "session-123" + + +def _exercise_delete_only_plan_uses_apply_ops(mod_name: str, monkeypatch, tmp_path: Path) -> None: + mod = importlib.import_module(mod_name) + + workspace = tmp_path / "repo" + workspace.mkdir(parents=True, exist_ok=True) + deleted = workspace / "old.py" + deleted.write_text("print('old')\n", encoding="utf-8") + + client = mod.RemoteUploadClient( + upload_endpoint="http://localhost:8004", + workspace_path=str(workspace), + collection_name="test-coll", + ) + removed_paths = [] + + monkeypatch.setattr( + client, + "_plan_delta_upload", + lambda changes: { + "needed_files": {"created": [], "updated": [], "moved": []}, + "operation_counts_preview": { + "created": 0, + "updated": 0, + "deleted": 1, + "moved": 0, + "skipped": 0, + "skipped_hash_match": 0, + "failed": 0, + }, + "needed_size_bytes": 0, + }, + ) + monkeypatch.setattr( + client, + "_build_plan_payload", + lambda changes: { + "manifest": {"bundle_id": "b1", "sequence_number": None}, + "operations": [{"operation": "deleted", "path": "old.py"}], + "file_hashes": {}, + }, + ) + monkeypatch.setattr(client, "create_delta_bundle", lambda *a, **k: (_ for _ in ()).throw(RuntimeError("should not bundle"))) + monkeypatch.setattr(client, "upload_bundle", lambda *a, **k: (_ for _ in ()).throw(RuntimeError("should not upload"))) + + class _Resp: + status_code = 200 + + @staticmethod + def raise_for_status(): + return None + + @staticmethod + def json(): + return { + "success": True, + "bundle_id": "b1", + "sequence_number": 3, + "processed_operations": {"deleted": 1, "created": 0, "updated": 0, "moved": 0, "skipped": 0, "skipped_hash_match": 0, "failed": 0}, + } + + monkeypatch.setattr(client.session, "post", lambda *a, **k: _Resp()) + monkeypatch.setattr(mod, "remove_cached_file", lambda path, repo_name=None, **_kwargs: removed_paths.append((path, repo_name))) + + assert client.process_changes_and_upload( + { + "created": [], + "updated": [], + "deleted": [deleted], + "moved": [], + "unchanged": [], + } + ) is True + assert client.last_upload_result["outcome"] == "uploaded" + assert client.last_upload_result["processed_operations"]["deleted"] == 1 + assert removed_paths == [(str(deleted.resolve()), client.repo_name)] + + +def test_remote_upload_client_delete_only_plan_uses_apply_ops(monkeypatch, tmp_path): + _exercise_delete_only_plan_uses_apply_ops("scripts.remote_upload_client", monkeypatch, tmp_path) + + +def test_standalone_upload_client_delete_only_plan_uses_apply_ops(monkeypatch, tmp_path): + _exercise_delete_only_plan_uses_apply_ops("scripts.standalone_upload_client", monkeypatch, tmp_path) + + +def _exercise_async_upload_sets_queued_result(mod_name: str, monkeypatch, tmp_path: Path) -> None: + mod = importlib.import_module(mod_name) + workspace = tmp_path / "repo" + workspace.mkdir(parents=True, exist_ok=True) + current = workspace / "app.py" + current.write_text("print('current')\n", encoding="utf-8") + + client = mod.RemoteUploadClient( + upload_endpoint="http://localhost:8004", + workspace_path=str(workspace), + collection_name="test-coll", + ) + + bundle_path = workspace / "bundle.tar.gz" + bundle_path.write_bytes(b"bundle") + monkeypatch.setattr(client, "_plan_delta_upload", lambda changes: None) + monkeypatch.setattr( + client, + "create_delta_bundle", + lambda changes: (str(bundle_path), {"bundle_id": "bundle-1", "total_size_bytes": 6}), + ) + monkeypatch.setattr( + client, + "upload_bundle", + lambda *a, **k: {"success": True, "sequence_number": 7, "processed_operations": None}, + ) + set_hash = MagicMock() + monkeypatch.setattr(mod, "set_cached_file_hash", set_hash) + monkeypatch.setattr(mod, "flush_cached_file_hashes", lambda: None, raising=False) + + assert client.process_changes_and_upload( + { + "created": [current], + "updated": [], + "deleted": [], + "moved": [], + "unchanged": [], + } + ) is True + assert client.last_upload_result["outcome"] == "queued" + assert client.last_upload_result["sequence_number"] == 7 + set_hash.assert_called_once() + + +def test_remote_upload_client_async_upload_sets_queued_result(monkeypatch, tmp_path): + _exercise_async_upload_sets_queued_result("scripts.remote_upload_client", monkeypatch, tmp_path) + + +def test_standalone_upload_client_async_upload_sets_queued_result(monkeypatch, tmp_path): + _exercise_async_upload_sets_queued_result("scripts.standalone_upload_client", monkeypatch, tmp_path) + + +def _exercise_watchable_path_excludes_ignored_updates(mod_name: str, monkeypatch, tmp_path: Path) -> None: + mod = importlib.import_module(mod_name) + + workspace = tmp_path / "repo" + workspace.mkdir(parents=True, exist_ok=True) + source = workspace / "src" / "tracked.py" + source.parent.mkdir(parents=True, exist_ok=True) + source.write_text("print('tracked')\n", encoding="utf-8") + + mirrored = workspace / "dev-workspace" / "nested" / "ignored.py" + mirrored.parent.mkdir(parents=True, exist_ok=True) + mirrored.write_text("print('ignored')\n", encoding="utf-8") + + monkeypatch.setenv("DEV_REMOTE_MODE", "1") + + client = mod.RemoteUploadClient( + upload_endpoint="http://localhost:8004", + workspace_path=str(workspace), + collection_name="test-coll", + ) + + assert client._is_watchable_path(source) is True + assert client._is_watchable_path(mirrored) is False + + +def test_remote_upload_client_watchable_path_excludes_ignored_updates(monkeypatch, tmp_path): + _exercise_watchable_path_excludes_ignored_updates( + "scripts.remote_upload_client", + monkeypatch, + tmp_path, + ) + + +def test_standalone_upload_client_watchable_path_excludes_ignored_updates(monkeypatch, tmp_path): + _exercise_watchable_path_excludes_ignored_updates( + "scripts.standalone_upload_client", + monkeypatch, + tmp_path, + ) + + +def test_remote_cache_keeps_missing_workspace_paths_and_ignores_other_workspaces( + monkeypatch, tmp_path +): + mod = importlib.import_module("scripts.remote_upload_client") + workspace = tmp_path / "repo" + other_workspace = tmp_path / "other-repo" + metadata_root = tmp_path / "metadata" + workspace.mkdir() + other_workspace.mkdir() + cache_path = metadata_root / ".codebase" / "cache.json" + cache_path.parent.mkdir(parents=True) + + missing_in_workspace = workspace / "deleted.py" + outside_workspace = other_workspace / "old.py" + cache_path.write_text( + json.dumps( + { + "file_hashes": { + str(missing_in_workspace): "hash-inside", + str(outside_workspace): "hash-outside", + } + } + ), + encoding="utf-8", + ) + + monkeypatch.setenv("MULTI_REPO_MODE", "0") + hashes = mod._load_local_cache_file_hashes( + str(workspace), None, metadata_root=str(metadata_root) + ) + + assert hashes == {str(missing_in_workspace): "hash-inside"} + assert not missing_in_workspace.exists() + + +def test_standalone_cache_keeps_missing_workspace_paths_and_ignores_other_workspaces( + tmp_path, +): + mod = importlib.import_module("scripts.standalone_upload_client") + workspace = tmp_path / "repo" + other_workspace = tmp_path / "other-repo" + workspace.mkdir() + other_workspace.mkdir() + cache_dir = workspace / ".context-engine" + cache_dir.mkdir() + + missing_in_workspace = workspace / "deleted.py" + outside_workspace = other_workspace / "old.py" + cache_path = cache_dir / "file_cache.json" + cache_path.write_text( + json.dumps( + { + "file_hashes": { + str(missing_in_workspace): "hash-inside", + str(outside_workspace): "hash-outside", + } + } + ), + encoding="utf-8", + ) + + cache = mod.SimpleHashCache(str(workspace), "repo") + + assert cache.all_paths() == [str(missing_in_workspace)] + assert not missing_in_workspace.exists() + persisted = json.loads(cache_path.read_text(encoding="utf-8")) + assert persisted["file_hashes"] == {str(missing_in_workspace): "hash-inside"} diff --git a/tests/test_upload_delta_planner.py b/tests/test_upload_delta_planner.py new file mode 100644 index 00000000..bdc528f8 --- /dev/null +++ b/tests/test_upload_delta_planner.py @@ -0,0 +1,195 @@ +import hashlib +import importlib +import json +from pathlib import Path + +import pytest + + +pytestmark = pytest.mark.unit + + +def test_plan_delta_upload_filters_remote_hash_matches(monkeypatch, tmp_path: Path): + work_root = tmp_path / "work" + slug = "repo-1234567890abcdef" + replica_root = work_root / slug + (replica_root / "src").mkdir(parents=True) + (replica_root / "src" / "same.py").write_text("same\n", encoding="utf-8") + (replica_root / "src" / "changed.py").write_text("old\n", encoding="utf-8") + + same_hash = "2c985b161217a952b7a410fd91495cebc349f520" + changed_hash = "281bac2b704617e807850e07e54bae3469f6a2e7" + cache_path = replica_root / ".codebase" / "cache.json" + cache_path.parent.mkdir(parents=True) + cache_path.write_text( + json.dumps( + { + "file_hashes": { + str((replica_root / "src" / "same.py").resolve()): same_hash, + str((replica_root / "src" / "changed.py").resolve()): changed_hash, + } + } + ), + encoding="utf-8", + ) + + monkeypatch.setenv("WORK_DIR", str(work_root)) + monkeypatch.setenv("WORKSPACE_PATH", str(work_root)) + monkeypatch.setenv("WATCH_ROOT", str(work_root)) + monkeypatch.setenv("MULTI_REPO_MODE", "1") + + planner = importlib.import_module("scripts.upload_delta_bundle") + planner = importlib.reload(planner) + planner.WORK_DIR = str(work_root) + + result = planner.plan_delta_upload( + workspace_path=str(replica_root), + operations=[ + { + "operation": "created", + "path": "src/same.py", + "size_bytes": 5, + "content_hash": f"sha1:{same_hash}", + }, + { + "operation": "updated", + "path": "src/changed.py", + "size_bytes": 4, + "content_hash": "sha1:new-hash", + }, + ], + file_hashes={ + "src/same.py": f"sha1:{same_hash}", + "src/changed.py": "sha1:new-hash", + }, + ) + + assert result["needed_files"] == { + "created": [], + "updated": ["src/changed.py"], + "moved": [], + } + assert result["operation_counts_preview"]["skipped_hash_match"] == 1 + assert result["needed_size_bytes"] == 4 + + +def test_plan_delta_upload_requires_content_when_cache_entry_is_missing(monkeypatch, tmp_path: Path): + work_root = tmp_path / "work" + slug = "repo-1234567890abcdef" + target = work_root / slug / "src" / "existing.py" + target.parent.mkdir(parents=True) + target.write_text("already there\n", encoding="utf-8") + + monkeypatch.setenv("WORK_DIR", str(work_root)) + monkeypatch.setenv("WORKSPACE_PATH", str(work_root)) + monkeypatch.setenv("WATCH_ROOT", str(work_root)) + monkeypatch.setenv("MULTI_REPO_MODE", "1") + + planner = importlib.import_module("scripts.upload_delta_bundle") + planner = importlib.reload(planner) + planner.WORK_DIR = str(work_root) + content_hash = "sha1:" + hashlib.sha1(target.read_bytes()).hexdigest() + + result = planner.plan_delta_upload( + workspace_path=str(work_root / slug), + operations=[ + {"operation": "created", "path": "src/existing.py", "size_bytes": target.stat().st_size, "content_hash": content_hash} + ], + file_hashes={"src/existing.py": content_hash}, + ) + + assert result["needed_files"]["created"] == ["src/existing.py"] + assert result["diagnostics"]["missing_targets"] == 0 + assert "filesystem_hash_probes" not in result["diagnostics"] + assert "filesystem_hash_matches" not in result["diagnostics"] + + +def test_plan_delta_upload_repairs_matching_cache_when_target_is_missing( + monkeypatch, tmp_path: Path +): + work_root = tmp_path / "work" + slug = "repo-1234567890abcdef" + replica_root = work_root / slug + cache_path = replica_root / ".codebase" / "cache.json" + cache_path.parent.mkdir(parents=True) + content_hash = "sha1:" + hashlib.sha1(b"missing target\n").hexdigest() + cache_path.write_text( + json.dumps( + { + "file_hashes": { + str((replica_root / "src" / "missing.py").resolve()): content_hash, + } + } + ), + encoding="utf-8", + ) + + monkeypatch.setenv("WORK_DIR", str(work_root)) + monkeypatch.setenv("WORKSPACE_PATH", str(work_root)) + monkeypatch.setenv("WATCH_ROOT", str(work_root)) + monkeypatch.setenv("MULTI_REPO_MODE", "1") + + planner = importlib.import_module("scripts.upload_delta_bundle") + planner = importlib.reload(planner) + planner.WORK_DIR = str(work_root) + + result = planner.plan_delta_upload( + workspace_path=str(replica_root), + operations=[ + { + "operation": "created", + "path": "src/missing.py", + "size_bytes": len(b"missing target\n"), + "content_hash": content_hash, + } + ], + file_hashes={"src/missing.py": content_hash}, + ) + + assert result["needed_files"]["created"] == ["src/missing.py"] + assert result["operation_counts_preview"]["skipped_hash_match"] == 0 + assert result["diagnostics"]["missing_targets"] == 1 + + +def test_plan_delta_upload_keeps_canonical_target_when_staging_has_no_old_replica( + monkeypatch, tmp_path: Path +): + work_root = tmp_path / "work" + slug = "repo-1234567890abcdef" + replica_root = work_root / slug + replica_root.mkdir(parents=True) + + monkeypatch.setenv("WORK_DIR", str(work_root)) + monkeypatch.setenv("WORKSPACE_PATH", str(work_root)) + monkeypatch.setenv("WATCH_ROOT", str(work_root)) + monkeypatch.setenv("MULTI_REPO_MODE", "1") + + planner = importlib.import_module("scripts.upload_delta_bundle") + planner = importlib.reload(planner) + planner.WORK_DIR = str(work_root) + monkeypatch.setattr(planner, "is_staging_enabled", lambda: True) + monkeypatch.setattr(planner, "_extract_repo_name_from_path", lambda _path: slug) + monkeypatch.setattr( + planner, + "get_collection_state_snapshot", + lambda workspace_path=None, repo_name=None: { + "active_repo_slug": slug, + "serving_repo_slug": slug, + }, + ) + + result = planner.plan_delta_upload( + workspace_path=str(replica_root), + operations=[ + { + "operation": "created", + "path": "src/new.py", + "size_bytes": 4, + "content_hash": "sha1:new-hash", + } + ], + file_hashes={"src/new.py": "sha1:new-hash"}, + ) + + assert result["replica_targets"] == [slug] + assert result["needed_files"]["created"] == ["src/new.py"] diff --git a/tests/test_upload_service_path_traversal.py b/tests/test_upload_service_path_traversal.py index 0d01478f..e646c3a6 100644 --- a/tests/test_upload_service_path_traversal.py +++ b/tests/test_upload_service_path_traversal.py @@ -1,11 +1,20 @@ import io +import hashlib import json +import os import tarfile from pathlib import Path import pytest +@pytest.fixture(autouse=True) +def _disable_ambient_staging(monkeypatch): + import scripts.upload_delta_bundle as us + + monkeypatch.setattr(us, "is_staging_enabled", lambda: False) + + def _write_bundle(tmp_path: Path, operations: list[dict]) -> Path: bundle_path = tmp_path / "bundle.tar.gz" payload = json.dumps({"operations": operations}).encode("utf-8") @@ -18,9 +27,21 @@ def _write_bundle(tmp_path: Path, operations: list[dict]) -> Path: return bundle_path -def _write_bundle_with_moved_file(tmp_path: Path, dest_path: str, content: bytes) -> Path: +def _write_bundle_with_moved_file( + tmp_path: Path, + dest_path: str, + content: bytes, + content_hash: str | None = None, +) -> Path: bundle_path = tmp_path / "bundle.tar.gz" - operations = [{"operation": "moved", "path": dest_path, "source_path": "missing_src.txt"}] + operation = { + "operation": "moved", + "path": dest_path, + "source_path": "missing_src.txt", + } + if content_hash: + operation["content_hash"] = content_hash + operations = [operation] payload = json.dumps({"operations": operations}).encode("utf-8") with tarfile.open(bundle_path, "w:gz") as tar: @@ -52,6 +73,59 @@ def _write_bundle_with_created_file(tmp_path: Path, rel_path: str, content: byte return bundle_path +def _write_bundle_with_hash_metadata( + tmp_path: Path, + *, + operations: list[dict], + file_hashes: dict[str, str] | None = None, + created_files: dict[str, bytes] | None = None, + updated_files: dict[str, bytes] | None = None, +) -> Path: + bundle_path = tmp_path / "bundle-hashes.tar.gz" + payload = json.dumps({"operations": operations}).encode("utf-8") + hashes_payload = json.dumps({"file_hashes": file_hashes or {}}).encode("utf-8") + + with tarfile.open(bundle_path, "w:gz") as tar: + info = tarfile.TarInfo(name="metadata/operations.json") + info.size = len(payload) + tar.addfile(info, io.BytesIO(payload)) + + hashes_info = tarfile.TarInfo(name="metadata/hashes.json") + hashes_info.size = len(hashes_payload) + tar.addfile(hashes_info, io.BytesIO(hashes_payload)) + + for rel_path, content in (created_files or {}).items(): + file_info = tarfile.TarInfo(name=f"files/created/{rel_path}") + file_info.size = len(content) + tar.addfile(file_info, io.BytesIO(content)) + + for rel_path, content in (updated_files or {}).items(): + file_info = tarfile.TarInfo(name=f"files/updated/{rel_path}") + file_info.size = len(content) + tar.addfile(file_info, io.BytesIO(content)) + + return bundle_path + + +def _write_repo_cache(work_dir: Path, slug: str, rel_path: str, file_hash: str) -> None: + target = (work_dir / slug / rel_path).resolve() + cache_path = work_dir / ".codebase" / "repos" / slug / "cache.json" + cache_path.parent.mkdir(parents=True, exist_ok=True) + cache_path.write_text( + json.dumps( + { + "file_hashes": { + str(target): { + "hash": file_hash, + } + } + }, + indent=2, + ), + encoding="utf-8", + ) + + def test_process_delta_bundle_rejects_traversal_created(tmp_path, monkeypatch): import scripts.upload_delta_bundle as us @@ -92,6 +166,37 @@ def test_process_delta_bundle_moved_falls_back_to_tar_payload_when_source_missin assert (work_dir / slug / "dst.txt").read_bytes() == b"moved-payload" +def test_process_delta_bundle_retries_already_applied_move_idempotently( + tmp_path, monkeypatch +): + import scripts.upload_delta_bundle as us + + work_dir = tmp_path / "work" + work_dir.mkdir(parents=True, exist_ok=True) + monkeypatch.setattr(us, "WORK_DIR", str(work_dir)) + + slug = "repo-0123456789abcdef" + content = b"moved-payload" + digest = hashlib.sha1(content).hexdigest() + target = work_dir / slug / "dst.txt" + target.parent.mkdir(parents=True, exist_ok=True) + target.write_bytes(content) + bundle = _write_bundle_with_moved_file( + tmp_path, "dst.txt", content, content_hash=f"sha1:{digest}" + ) + + counts = us.process_delta_bundle( + workspace_path=f"/work/{slug}", + bundle_path=bundle, + manifest={"bundle_id": "b-moved-retry"}, + ) + + assert counts["failed"] == 0 + assert counts["skipped_hash_match"] == 1 + assert target.read_bytes() == content + + + def test_process_delta_bundle_slugged_workspace_creates_marker(tmp_path, monkeypatch): import scripts.upload_delta_bundle as us @@ -197,3 +302,518 @@ def test_process_delta_bundle_rejects_traversal_moved_source(tmp_path, monkeypat bundle_path=bundle, manifest={"bundle_id": "b1"}, ) + + +def test_process_delta_bundle_skips_created_write_when_server_hash_matches(tmp_path, monkeypatch): + import scripts.upload_delta_bundle as us + + work_dir = tmp_path / "work" + work_dir.mkdir(parents=True, exist_ok=True) + monkeypatch.setattr(us, "WORK_DIR", str(work_dir)) + + slug = "repo-0123456789abcdef" + rel_path = "src/file.txt" + content = b"same-content" + file_hash = "sha1:efb5d7d4d38013264f2c00fceeb401f8c8d77d9f" + + target = work_dir / slug / rel_path + target.parent.mkdir(parents=True, exist_ok=True) + target.write_bytes(content) + os.utime(target, ns=(1_000_000_000, 1_000_000_000)) + before_mtime_ns = target.stat().st_mtime_ns + _write_repo_cache(work_dir, slug, rel_path, file_hash) + + bundle = _write_bundle_with_hash_metadata( + tmp_path, + operations=[ + { + "operation": "created", + "path": rel_path, + "content_hash": file_hash, + } + ], + file_hashes={rel_path: file_hash}, + created_files={rel_path: content}, + ) + + counts = us.process_delta_bundle( + workspace_path=f"/work/{slug}", + bundle_path=bundle, + manifest={"bundle_id": "b-skip-created"}, + ) + + assert counts.get("created") == 0 + assert counts.get("skipped") == 1 + assert counts.get("skipped_hash_match") == 1 + assert target.read_bytes() == content + assert target.stat().st_mtime_ns == before_mtime_ns + + +def test_process_delta_bundle_uses_hashes_metadata_for_updated_skip(tmp_path, monkeypatch): + import scripts.upload_delta_bundle as us + + work_dir = tmp_path / "work" + work_dir.mkdir(parents=True, exist_ok=True) + monkeypatch.setattr(us, "WORK_DIR", str(work_dir)) + + slug = "repo-0123456789abcdef" + rel_path = "src/keep.txt" + content = b"existing-content" + file_hash = "sha1:2910e29d6f6d3d2f01f8cc52ec386a4936ca9d2f" + + target = work_dir / slug / rel_path + target.parent.mkdir(parents=True, exist_ok=True) + target.write_bytes(content) + os.utime(target, ns=(2_000_000_000, 2_000_000_000)) + before_mtime_ns = target.stat().st_mtime_ns + _write_repo_cache(work_dir, slug, rel_path, file_hash) + + bundle = _write_bundle_with_hash_metadata( + tmp_path, + operations=[ + { + "operation": "updated", + "path": rel_path, + } + ], + file_hashes={rel_path: file_hash}, + updated_files={rel_path: content}, + ) + + counts = us.process_delta_bundle( + workspace_path=f"/work/{slug}", + bundle_path=bundle, + manifest={"bundle_id": "b-skip-updated"}, + ) + + assert counts.get("updated") == 0 + assert counts.get("skipped") == 1 + assert counts.get("skipped_hash_match") == 1 + assert target.read_bytes() == content + assert target.stat().st_mtime_ns == before_mtime_ns + + +def test_normalize_hash_value_strips_algorithm_prefixes(): + import scripts.upload_delta_bundle as us + + assert us._normalize_hash_value("sha1:ABCDEF") == "abcdef" + assert us._normalize_hash_value("md5:ABCDEF") == "abcdef" + assert us._normalize_hash_value("sha256:ABCDEF") == "abcdef" + assert us._normalize_hash_value("ABCDEF") == "abcdef" + + +def test_process_delta_bundle_uses_first_marker_match_for_created_members(tmp_path, monkeypatch): + import scripts.upload_delta_bundle as us + + work_dir = tmp_path / "work" + work_dir.mkdir(parents=True, exist_ok=True) + monkeypatch.setattr(us, "WORK_DIR", str(work_dir)) + + slug = "repo-0123456789abcdef" + rel_path = "nested/files/created/path.txt" + content = b"marker-safe" + bundle = _write_bundle_with_created_file(tmp_path, rel_path, content) + + counts = us.process_delta_bundle( + workspace_path=f"/work/{slug}", + bundle_path=bundle, + manifest={"bundle_id": "b-created-marker"}, + ) + + assert counts.get("created") == 1 + assert (work_dir / slug / rel_path).read_bytes() == content + + +def test_process_delta_bundle_deleted_prunes_empty_parent_dirs(tmp_path, monkeypatch): + import scripts.upload_delta_bundle as us + + work_dir = tmp_path / "work" + work_dir.mkdir(parents=True, exist_ok=True) + monkeypatch.setattr(us, "WORK_DIR", str(work_dir)) + + slug = "repo-0123456789abcdef" + rel_path = "dev-workspace/nested/stale.py" + target = work_dir / slug / rel_path + target.parent.mkdir(parents=True, exist_ok=True) + target.write_text("stale\n", encoding="utf-8") + + bundle = _write_bundle( + tmp_path, + [{"operation": "deleted", "path": rel_path}], + ) + + counts = us.process_delta_bundle( + workspace_path=f"/work/{slug}", + bundle_path=bundle, + manifest={"bundle_id": "b-delete-prune"}, + ) + + assert counts.get("deleted") == 1 + assert not target.exists() + assert not (work_dir / slug / "dev-workspace" / "nested").exists() + assert not (work_dir / slug / "dev-workspace").exists() + assert (work_dir / slug).exists() + + +def test_process_delta_bundle_moved_prunes_empty_source_parent_dirs(tmp_path, monkeypatch): + import scripts.upload_delta_bundle as us + + work_dir = tmp_path / "work" + work_dir.mkdir(parents=True, exist_ok=True) + monkeypatch.setattr(us, "WORK_DIR", str(work_dir)) + + slug = "repo-0123456789abcdef" + src = work_dir / slug / "dev-workspace" / "nested" / "from.py" + dest_rel_path = "dest/to.py" + src.parent.mkdir(parents=True, exist_ok=True) + src.write_text("payload\n", encoding="utf-8") + + bundle = _write_bundle( + tmp_path, + [{"operation": "moved", "path": dest_rel_path, "source_path": "dev-workspace/nested/from.py"}], + ) + + counts = us.process_delta_bundle( + workspace_path=f"/work/{slug}", + bundle_path=bundle, + manifest={"bundle_id": "b-move-prune"}, + ) + + assert counts.get("moved") == 1 + assert not src.exists() + assert (work_dir / slug / dest_rel_path).read_text(encoding="utf-8") == "payload\n" + assert not (work_dir / slug / "dev-workspace" / "nested").exists() + assert not (work_dir / slug / "dev-workspace").exists() + assert (work_dir / slug).exists() + + +def test_process_delta_bundle_does_not_sweep_stranded_empty_dirs_without_file_ops(tmp_path, monkeypatch): + import scripts.upload_delta_bundle as us + + work_dir = tmp_path / "work" + work_dir.mkdir(parents=True, exist_ok=True) + monkeypatch.setattr(us, "WORK_DIR", str(work_dir)) + slug = "repo-0123456789abcdef" + stranded = work_dir / slug / "dev-workspace" / "nested" / "empty" + stranded.mkdir(parents=True, exist_ok=True) + + bundle = _write_bundle(tmp_path, []) + + counts = us.process_delta_bundle( + workspace_path=f"/work/{slug}", + bundle_path=bundle, + manifest={"bundle_id": "b-sweep-empty"}, + ) + + assert counts == { + "created": 0, + "updated": 0, + "deleted": 0, + "moved": 0, + "skipped": 0, + "skipped_hash_match": 0, + "failed": 0, + } + assert stranded.exists() + assert (work_dir / slug / "dev-workspace").exists() + assert (work_dir / slug).exists() + + +def test_process_delta_bundle_skips_broad_empty_dir_sweep_when_disabled(tmp_path, monkeypatch): + import scripts.upload_delta_bundle as us + + work_dir = tmp_path / "work" + work_dir.mkdir(parents=True, exist_ok=True) + monkeypatch.setattr(us, "WORK_DIR", str(work_dir)) + monkeypatch.setenv("CTXCE_UPLOAD_EMPTY_DIR_SWEEP", "0") + + slug = "repo-0123456789abcdef" + stranded = work_dir / slug / "dev-workspace" / "nested" / "empty" + stranded.mkdir(parents=True, exist_ok=True) + + bundle = _write_bundle(tmp_path, []) + + us.process_delta_bundle( + workspace_path=f"/work/{slug}", + bundle_path=bundle, + manifest={"bundle_id": "b-sweep-disabled"}, + ) + + assert stranded.exists() + + +def test_process_delta_bundle_skips_broad_empty_dir_sweep_when_recent(tmp_path, monkeypatch): + import scripts.upload_delta_bundle as us + + work_dir = tmp_path / "work" + work_dir.mkdir(parents=True, exist_ok=True) + monkeypatch.setattr(us, "WORK_DIR", str(work_dir)) + slug = "repo-0123456789abcdef" + stranded = work_dir / slug / "dev-workspace" / "nested" / "empty" + stranded.mkdir(parents=True, exist_ok=True) + + bundle = _write_bundle(tmp_path, []) + + us.process_delta_bundle( + workspace_path=f"/work/{slug}", + bundle_path=bundle, + manifest={"bundle_id": "b-sweep-recent"}, + ) + + assert stranded.exists() + + +def test_process_delta_bundle_preserves_protected_top_level_dirs_when_empty(tmp_path, monkeypatch): + import scripts.upload_delta_bundle as us + + work_dir = tmp_path / "work" + work_dir.mkdir(parents=True, exist_ok=True) + monkeypatch.setattr(us, "WORK_DIR", str(work_dir)) + monkeypatch.setenv("CTXCE_UPLOAD_EMPTY_DIR_SWEEP", "1") + monkeypatch.setenv("CTXCE_UPLOAD_EMPTY_DIR_SWEEP_INTERVAL_SECONDS", "0") + + slug = "repo-0123456789abcdef" + protected = work_dir / slug / ".remote-git" + protected.mkdir(parents=True, exist_ok=True) + + bundle = _write_bundle(tmp_path, []) + + us.process_delta_bundle( + workspace_path=f"/work/{slug}", + bundle_path=bundle, + manifest={"bundle_id": "b-protected-empty"}, + ) + + assert protected.exists() + + +def test_process_delta_bundle_preserves_nested_dirs_under_protected_top_level(tmp_path, monkeypatch): + import scripts.upload_delta_bundle as us + + work_dir = tmp_path / "work" + work_dir.mkdir(parents=True, exist_ok=True) + monkeypatch.setattr(us, "WORK_DIR", str(work_dir)) + monkeypatch.setenv("CTXCE_UPLOAD_EMPTY_DIR_SWEEP", "1") + monkeypatch.setenv("CTXCE_UPLOAD_EMPTY_DIR_SWEEP_INTERVAL_SECONDS", "0") + + slug = "repo-0123456789abcdef" + protected_nested = work_dir / slug / ".codebase" / "repos" / "empty" + protected_nested.mkdir(parents=True, exist_ok=True) + + bundle = _write_bundle(tmp_path, []) + + us.process_delta_bundle( + workspace_path=f"/work/{slug}", + bundle_path=bundle, + manifest={"bundle_id": "b-protected-nested-empty"}, + ) + + assert protected_nested.exists() + + +def test_plan_delta_upload_skips_matching_created_files(tmp_path, monkeypatch): + import scripts.upload_delta_bundle as us + + work_dir = tmp_path / "work" + work_dir.mkdir(parents=True, exist_ok=True) + monkeypatch.setattr(us, "WORK_DIR", str(work_dir)) + + slug = "repo-0123456789abcdef" + rel_path = "src/file.txt" + content = b"same-content" + file_hash = "sha1:efb5d7d4d38013264f2c00fceeb401f8c8d77d9f" + + target = work_dir / slug / rel_path + target.parent.mkdir(parents=True, exist_ok=True) + target.write_bytes(content) + _write_repo_cache(work_dir, slug, rel_path, file_hash) + + plan = us.plan_delta_upload( + workspace_path=f"/work/{slug}", + operations=[ + { + "operation": "created", + "path": rel_path, + "content_hash": file_hash, + "size_bytes": len(content), + } + ], + file_hashes={rel_path: file_hash}, + ) + + assert plan["needed_files"]["created"] == [] + assert plan["operation_counts_preview"]["skipped_hash_match"] == 1 + assert plan["needed_size_bytes"] == 0 + + +def test_plan_delta_upload_marks_updated_file_needed_when_hash_missing(tmp_path, monkeypatch): + import scripts.upload_delta_bundle as us + + work_dir = tmp_path / "work" + work_dir.mkdir(parents=True, exist_ok=True) + monkeypatch.setattr(us, "WORK_DIR", str(work_dir)) + + slug = "repo-0123456789abcdef" + rel_path = "src/keep.txt" + file_hash = "sha1:2910e29d6f6d3d2f01f8cc52ec386a4936ca9d2f" + + plan = us.plan_delta_upload( + workspace_path=f"/work/{slug}", + operations=[ + { + "operation": "updated", + "path": rel_path, + "content_hash": file_hash, + "size_bytes": 17, + } + ], + file_hashes={rel_path: file_hash}, + ) + + assert plan["needed_files"]["updated"] == [rel_path] + assert plan["operation_counts_preview"]["updated"] == 1 + assert plan["needed_size_bytes"] == 17 + + +def test_plan_delta_upload_skips_move_content_when_source_exists_on_server(tmp_path, monkeypatch): + import scripts.upload_delta_bundle as us + + work_dir = tmp_path / "work" + work_dir.mkdir(parents=True, exist_ok=True) + monkeypatch.setattr(us, "WORK_DIR", str(work_dir)) + + slug = "repo-0123456789abcdef" + source_rel = "src/old.py" + dest_rel = "src/new.py" + source = work_dir / slug / source_rel + source.parent.mkdir(parents=True, exist_ok=True) + source.write_text("print('move')\n", encoding="utf-8") + + plan = us.plan_delta_upload( + workspace_path=f"/work/{slug}", + operations=[ + { + "operation": "moved", + "path": dest_rel, + "source_path": source_rel, + "content_hash": "sha1:abc123", + "size_bytes": 12, + } + ], + file_hashes={dest_rel: "sha1:abc123"}, + ) + + assert plan["needed_files"]["moved"] == [] + assert plan["operation_counts_preview"]["moved"] == 1 + assert plan["needed_size_bytes"] == 0 + + +def test_plan_delta_upload_marks_move_needed_when_source_path_is_invalid(tmp_path, monkeypatch): + import scripts.upload_delta_bundle as us + + work_dir = tmp_path / "work" + work_dir.mkdir(parents=True, exist_ok=True) + monkeypatch.setattr(us, "WORK_DIR", str(work_dir)) + + slug = "repo-0123456789abcdef" + dest_rel = "src/new.py" + + plan = us.plan_delta_upload( + workspace_path=f"/work/{slug}", + operations=[ + { + "operation": "moved", + "path": dest_rel, + "source_path": "../escape.py", + "content_hash": "sha1:abc123", + "size_bytes": 12, + } + ], + file_hashes={dest_rel: "sha1:abc123"}, + ) + + assert plan["needed_files"]["moved"] == [dest_rel] + assert plan["operation_counts_preview"]["moved"] == 1 + assert plan["needed_size_bytes"] == 12 + + +def test_apply_delta_operations_moves_file_without_bundle(tmp_path, monkeypatch): + import scripts.upload_delta_bundle as us + + work_dir = tmp_path / "work" + work_dir.mkdir(parents=True, exist_ok=True) + monkeypatch.setattr(us, "WORK_DIR", str(work_dir)) + + slug = "repo-0123456789abcdef" + source_rel = "src/old.py" + dest_rel = "src/new.py" + source = work_dir / slug / source_rel + source.parent.mkdir(parents=True, exist_ok=True) + source.write_text("print('move')\n", encoding="utf-8") + + counts = us.apply_delta_operations( + workspace_path=f"/work/{slug}", + operations=[ + { + "operation": "moved", + "path": dest_rel, + "source_path": source_rel, + "content_hash": "sha1:abc123", + } + ], + file_hashes={dest_rel: "sha1:abc123"}, + ) + + assert counts["moved"] == 1 + assert not source.exists() + assert (work_dir / slug / dest_rel).exists() + + +def test_apply_delta_operations_retries_already_applied_move_idempotently( + tmp_path, monkeypatch +): + import scripts.upload_delta_bundle as us + + work_dir = tmp_path / "work" + work_dir.mkdir(parents=True, exist_ok=True) + monkeypatch.setattr(us, "WORK_DIR", str(work_dir)) + + slug = "repo-0123456789abcdef" + content = b"print('move')\n" + digest = hashlib.sha1(content).hexdigest() + target = work_dir / slug / "src" / "new.py" + target.parent.mkdir(parents=True, exist_ok=True) + target.write_bytes(content) + + counts = us.apply_delta_operations( + workspace_path=f"/work/{slug}", + operations=[ + { + "operation": "moved", + "path": "src/new.py", + "source_path": "src/old.py", + "content_hash": f"sha1:{digest}", + } + ], + file_hashes={"src/new.py": f"sha1:{digest}"}, + ) + + assert counts["failed"] == 0 + assert counts["skipped_hash_match"] == 1 + assert target.read_bytes() == content + + +def test_apply_delta_operations_raises_clear_error_when_no_replica_roots(tmp_path, monkeypatch): + import scripts.upload_delta_bundle as us + + work_dir = tmp_path / "work" + work_dir.mkdir(parents=True, exist_ok=True) + monkeypatch.setattr(us, "WORK_DIR", str(work_dir)) + monkeypatch.setattr(us, "_resolve_replica_roots", lambda workspace_path: {}) + + with pytest.raises(ValueError, match="No replica roots available"): + us.apply_delta_operations( + workspace_path="/work/repo", + operations=[], + file_hashes={}, + ) diff --git a/tests/test_upload_service_status.py b/tests/test_upload_service_status.py new file mode 100644 index 00000000..1b5f9f66 --- /dev/null +++ b/tests/test_upload_service_status.py @@ -0,0 +1,467 @@ +import asyncio +import importlib +from pathlib import Path + +import pytest +from fastapi.testclient import TestClient + + +def _disable_auth(srv, monkeypatch) -> None: + monkeypatch.setattr(srv, "AUTH_ENABLED", False) + + +@pytest.mark.unit +def test_delta_status_exposes_last_processed_operations(monkeypatch): + srv = importlib.import_module("scripts.upload_service") + srv = importlib.reload(srv) + _disable_auth(srv, monkeypatch) + + monkeypatch.setattr(srv, "get_collection_name", lambda _repo=None: "test-coll") + monkeypatch.setattr(srv, "_extract_repo_name_from_path", lambda _path: "repo") + + key = srv.get_workspace_key("/work/repo") + srv._sequence_tracker[key] = 7 + srv._upload_result_tracker[key] = { + "workspace_path": "/work/repo", + "bundle_id": "bundle-123", + "sequence_number": 7, + "processed_operations": { + "created": 1, + "updated": 2, + "deleted": 0, + "moved": 0, + "skipped": 5, + "skipped_hash_match": 4, + "failed": 0, + }, + "processing_time_ms": 321, + "status": "completed", + "completed_at": "2026-03-07T15:40:46.623000", + } + + client = TestClient(srv.app) + resp = client.get("/api/v1/delta/status", params={"workspace_path": "/work/repo"}) + assert resp.status_code == 200 + body = resp.json() + assert body["last_sequence"] == 7 + assert body["last_upload"] == "2026-03-07T15:40:46.623000" + assert body["status"] == "ready" + assert body["server_info"]["last_bundle_id"] == "bundle-123" + assert body["server_info"]["last_processing_time_ms"] == 321 + assert body["server_info"]["last_processed_operations"]["skipped_hash_match"] == 4 + assert body["server_info"]["last_upload_status"] == "completed" + assert body["server_info"]["last_error"] is None + + +@pytest.mark.unit +def test_process_bundle_background_tracks_completed_operations(monkeypatch, tmp_path: Path): + srv = importlib.import_module("scripts.upload_service") + srv = importlib.reload(srv) + _disable_auth(srv, monkeypatch) + + bundle_path = tmp_path / "bundle.tar.gz" + bundle_path.write_bytes(b"placeholder") + + monkeypatch.setattr( + srv, + "process_delta_bundle", + lambda workspace_path, bundle_path, manifest: { + "created": 0, + "updated": 0, + "deleted": 0, + "moved": 0, + "skipped": 10, + "skipped_hash_match": 10, + "failed": 0, + }, + ) + monkeypatch.setattr(srv, "log_activity", lambda *a, **k: None) + + asyncio.run( + srv._process_bundle_background( + workspace_path="/work/repo", + bundle_path=bundle_path, + manifest={"bundle_id": "bundle-xyz"}, + sequence_number=3, + bundle_id="bundle-xyz", + ) + ) + + key = srv.get_workspace_key("/work/repo") + tracked = srv._upload_result_tracker[key] + assert tracked["status"] == "completed" + assert tracked["sequence_number"] == 3 + assert tracked["processed_operations"]["skipped_hash_match"] == 10 + assert tracked["processing_time_ms"] is not None + assert not bundle_path.exists() + + +@pytest.mark.unit +def test_process_bundle_background_does_not_advance_sequence_after_partial_failure( + monkeypatch, tmp_path: Path +): + srv = importlib.import_module("scripts.upload_service") + srv = importlib.reload(srv) + _disable_auth(srv, monkeypatch) + + bundle_path = tmp_path / "bundle.tar.gz" + bundle_path.write_bytes(b"placeholder") + monkeypatch.setattr( + srv, + "process_delta_bundle", + lambda *_args: { + "created": 1, + "updated": 0, + "deleted": 0, + "moved": 0, + "skipped": 0, + "skipped_hash_match": 0, + "failed": 1, + }, + ) + monkeypatch.setattr(srv, "log_activity", lambda *a, **k: None) + + key = srv.get_workspace_key("/work/repo") + srv._sequence_tracker[key] = 2 + asyncio.run( + srv._process_bundle_background( + workspace_path="/work/repo", + bundle_path=bundle_path, + manifest={"bundle_id": "bundle-partial"}, + sequence_number=3, + bundle_id="bundle-partial", + ) + ) + + assert srv._sequence_tracker[key] == 2 + assert srv._upload_result_tracker[key]["status"] == "failed" + assert srv._upload_result_tracker[key]["failed_count"] == 1 + + +@pytest.mark.unit +def test_delta_status_reports_processing_while_upload_in_progress(monkeypatch): + srv = importlib.import_module("scripts.upload_service") + srv = importlib.reload(srv) + _disable_auth(srv, monkeypatch) + + monkeypatch.setattr(srv, "get_collection_name", lambda _repo=None: "test-coll") + monkeypatch.setattr(srv, "_extract_repo_name_from_path", lambda _path: "repo") + + key = srv.get_workspace_key("/work/repo") + srv._upload_result_tracker[key] = { + "workspace_path": "/work/repo", + "bundle_id": "bundle-123", + "sequence_number": 8, + "processed_operations": None, + "processing_time_ms": None, + "status": "processing", + "completed_at": None, + } + + client = TestClient(srv.app) + resp = client.get("/api/v1/delta/status", params={"workspace_path": "/work/repo"}) + assert resp.status_code == 200 + body = resp.json() + assert body["status"] == "processing" + assert body["server_info"]["last_upload_status"] == "processing" + + +@pytest.mark.unit +def test_delta_status_exposes_journal_summary(monkeypatch): + srv = importlib.import_module("scripts.upload_service") + srv = importlib.reload(srv) + _disable_auth(srv, monkeypatch) + + monkeypatch.setattr(srv, "get_collection_name", lambda _repo=None: "test-coll") + monkeypatch.setattr(srv, "_extract_repo_name_from_path", lambda _path: "repo") + monkeypatch.setattr( + srv, + "get_index_journal_summary", + lambda **_: { + "total": 12, + "retryable": 7, + "outstanding": 9, + "counts": {"pending": 5, "failed": 2}, + "sample_errors": [{"path": "/work/repo/bad.py", "error": "boom"}], + }, + ) + + client = TestClient(srv.app) + resp = client.get("/api/v1/delta/status", params={"workspace_path": "/work/repo"}) + + assert resp.status_code == 200 + body = resp.json() + assert body["pending_operations"] == 9 + assert body["server_info"]["journal"]["total"] == 12 + assert body["server_info"]["journal"]["sample_errors"][0]["error"] == "boom" + + +@pytest.mark.unit +def test_delta_status_aggregates_journal_for_workspace_root(monkeypatch): + srv = importlib.import_module("scripts.upload_service") + srv = importlib.reload(srv) + _disable_auth(srv, monkeypatch) + + monkeypatch.setattr(srv, "get_collection_name", lambda repo=None: f"coll-{repo or 'root'}") + monkeypatch.setattr(srv, "_extract_repo_name_from_path", lambda _path: "should-not-be-used") + summary_calls = [] + + def journal_summary(**kwargs): + summary_calls.append(kwargs) + return {"total": 2, "retryable": 1, "outstanding": 1, "counts": {"pending": 1}} + + monkeypatch.setattr(srv, "get_index_journal_summary", journal_summary) + + client = TestClient(srv.app) + resp = client.get("/api/v1/delta/status", params={"workspace_path": "/work"}) + + assert resp.status_code == 200 + assert resp.json()["pending_operations"] == 1 + assert summary_calls == [{"workspace_path": "/work", "repo_name": None}] + + +@pytest.mark.unit +def test_delta_plan_endpoint_returns_needed_files(monkeypatch): + srv = importlib.import_module("scripts.upload_service") + srv = importlib.reload(srv) + _disable_auth(srv, monkeypatch) + + monkeypatch.setattr( + srv, + "plan_delta_upload", + lambda workspace_path, operations, file_hashes=None: { + "needed_files": {"created": ["src/app.py"], "updated": [], "moved": []}, + "operation_counts_preview": { + "created": 1, + "updated": 0, + "deleted": 0, + "moved": 0, + "skipped": 2, + "skipped_hash_match": 2, + "failed": 0, + }, + "needed_size_bytes": 123, + "replica_targets": ["repo-0123456789abcdef"], + }, + ) + + client = TestClient(srv.app) + resp = client.post( + "/api/v1/delta/plan", + json={ + "workspace_path": "/work/repo", + "manifest": {"bundle_id": "b1"}, + "operations": [{"operation": "created", "path": "src/app.py"}], + "file_hashes": {"src/app.py": "sha1:abc"}, + }, + ) + assert resp.status_code == 200 + body = resp.json() + assert body["success"] is True + assert body["needed_files"]["created"] == ["src/app.py"] + assert body["operation_counts_preview"]["skipped_hash_match"] == 2 + assert body["needed_size_bytes"] == 123 + + +@pytest.mark.unit +def test_upload_managed_resolution_ignores_client_collection(monkeypatch): + srv = importlib.import_module("scripts.upload_service") + srv = importlib.reload(srv) + _disable_auth(srv, monkeypatch) + + monkeypatch.setattr(srv, "logical_repo_reuse_enabled", lambda: False) + monkeypatch.setattr(srv, "_extract_repo_name_from_path", lambda path: Path(path).name) + monkeypatch.setattr(srv, "get_collection_name", lambda repo=None: f"server-{repo}") + + collection, repo = srv._resolve_collection_for_request( + workspace_path="/work/repo", + client_collection_name="repo-071ca222", + logical_repo_id="fs:123", + source_path="/host/Context-Engine", + ) + + assert repo == "Context-Engine" + assert collection == "server-Context-Engine" + + +@pytest.mark.unit +def test_delta_plan_endpoint_uses_safe_defaults_for_sparse_plan(monkeypatch): + srv = importlib.import_module("scripts.upload_service") + srv = importlib.reload(srv) + _disable_auth(srv, monkeypatch) + + monkeypatch.setattr( + srv, + "plan_delta_upload", + lambda workspace_path, operations, file_hashes=None: {}, + ) + + client = TestClient(srv.app) + resp = client.post( + "/api/v1/delta/plan", + json={ + "workspace_path": "/work/repo", + "manifest": {"bundle_id": "b1"}, + "operations": [{"operation": "created", "path": "src/app.py"}], + "file_hashes": {"src/app.py": "sha1:abc"}, + }, + ) + assert resp.status_code == 200 + body = resp.json() + assert body["success"] is True + assert body["needed_files"] == {"created": [], "updated": [], "moved": []} + assert body["operation_counts_preview"]["failed"] == 0 + assert body["needed_size_bytes"] == 0 + assert body["replica_targets"] == [] + + +@pytest.mark.unit +def test_apply_ops_endpoint_returns_processed_operations(monkeypatch): + srv = importlib.import_module("scripts.upload_service") + srv = importlib.reload(srv) + _disable_auth(srv, monkeypatch) + + monkeypatch.setattr( + srv, + "apply_delta_operations", + lambda workspace_path, operations, file_hashes=None: { + "created": 0, + "updated": 0, + "deleted": 1, + "moved": 0, + "skipped": 0, + "skipped_hash_match": 0, + "failed": 0, + }, + ) + + client = TestClient(srv.app) + resp = client.post( + "/api/v1/delta/apply_ops", + json={ + "workspace_path": "/work/repo", + "manifest": {"bundle_id": "b2"}, + "operations": [{"operation": "deleted", "path": "src/old.py"}], + "file_hashes": {}, + }, + ) + assert resp.status_code == 200 + body = resp.json() + assert body["success"] is True + assert body["processed_operations"]["deleted"] == 1 + assert body["processing_time_ms"] is not None + + +@pytest.mark.unit +def test_apply_ops_advances_sequence_when_all_operations_match(monkeypatch): + srv = importlib.import_module("scripts.upload_service") + srv = importlib.reload(srv) + _disable_auth(srv, monkeypatch) + + monkeypatch.setattr( + srv, + "apply_delta_operations", + lambda *_args, **_kwargs: { + "created": 0, + "updated": 0, + "deleted": 0, + "moved": 0, + "skipped": 1, + "skipped_hash_match": 1, + "failed": 0, + }, + ) + + key = srv.get_workspace_key("/work/repo") + srv._sequence_tracker[key] = 4 + client = TestClient(srv.app) + resp = client.post( + "/api/v1/delta/apply_ops", + json={ + "workspace_path": "/work/repo", + "manifest": {"bundle_id": "b-match", "sequence_number": 5}, + "operations": [{"operation": "moved", "path": "src/new.py"}], + "file_hashes": {"src/new.py": "sha1:match"}, + }, + ) + + assert resp.status_code == 200 + assert resp.json()["success"] is True + assert srv._sequence_tracker[key] == 5 + + +@pytest.mark.unit +def test_apply_ops_endpoint_does_not_advance_sequence_after_partial_failure(monkeypatch): + srv = importlib.import_module("scripts.upload_service") + srv = importlib.reload(srv) + _disable_auth(srv, monkeypatch) + + monkeypatch.setattr( + srv, + "apply_delta_operations", + lambda *_args, **_kwargs: { + "created": 1, + "updated": 0, + "deleted": 0, + "moved": 0, + "skipped": 0, + "skipped_hash_match": 0, + "failed": 1, + }, + ) + + key = srv.get_workspace_key("/work/repo") + srv._sequence_tracker[key] = 4 + client = TestClient(srv.app) + resp = client.post( + "/api/v1/delta/apply_ops", + json={ + "workspace_path": "/work/repo", + "manifest": {"bundle_id": "b-partial", "sequence_number": 5}, + "operations": [{"operation": "created", "path": "src/new.py"}], + "file_hashes": {"src/new.py": "sha1:new"}, + }, + ) + + assert resp.status_code == 200 + body = resp.json() + assert body["success"] is False + assert body["error"]["code"] == "APPLY_OPS_PARTIAL_FAILURE" + assert srv._sequence_tracker[key] == 4 + + +@pytest.mark.unit +def test_apply_ops_endpoint_marks_tracker_error_state_on_failure(monkeypatch): + srv = importlib.import_module("scripts.upload_service") + srv = importlib.reload(srv) + _disable_auth(srv, monkeypatch) + + monkeypatch.setattr( + srv, + "apply_delta_operations", + lambda workspace_path, operations, file_hashes=None: (_ for _ in ()).throw( + RuntimeError("boom") + ), + ) + + client = TestClient(srv.app) + resp = client.post( + "/api/v1/delta/apply_ops", + json={ + "workspace_path": "/work/repo", + "manifest": {"bundle_id": "b3"}, + "operations": [{"operation": "deleted", "path": "src/old.py"}], + "file_hashes": {}, + }, + ) + assert resp.status_code == 200 + body = resp.json() + assert body["success"] is False + assert body["error"]["code"] == "APPLY_OPS_ERROR" + + key = srv.get_workspace_key("/work/repo") + tracked = srv._upload_result_tracker[key] + assert tracked["status"] == "error" + assert tracked["error"] == "boom" + assert tracked["message"] == "boom" + assert tracked["completed_at"] is not None diff --git a/tests/test_watch_consistency.py b/tests/test_watch_consistency.py new file mode 100644 index 00000000..bf51c79e --- /dev/null +++ b/tests/test_watch_consistency.py @@ -0,0 +1,102 @@ +#!/usr/bin/env python3 +import importlib +from pathlib import Path +from unittest.mock import MagicMock + +import pytest + + +pytestmark = pytest.mark.unit + + +@pytest.fixture +def capture_list_workspaces(): + captured = {} + + def fake_list_workspaces(search_root=None, use_qdrant_fallback=True): + captured["search_root"] = search_root + captured["use_qdrant_fallback"] = use_qdrant_fallback + return [] + + return captured, fake_list_workspaces + + +def test_run_consistency_audit_scans_from_watcher_root( + monkeypatch, tmp_path, capture_list_workspaces +): + mod = importlib.import_module("scripts.watch_index_core.consistency") + captured, fake_list_workspaces = capture_list_workspaces + + monkeypatch.setattr(mod, "list_workspaces", fake_list_workspaces) + monkeypatch.setattr(mod, "_consistency_audit_enabled", lambda: True) + + mod.run_consistency_audit(MagicMock(), tmp_path) + + assert "search_root" in captured and "use_qdrant_fallback" in captured + assert Path(captured["search_root"]).resolve() == Path(tmp_path).resolve() + assert captured["use_qdrant_fallback"] is False + + +def test_run_empty_dir_sweep_maintenance_scans_from_watcher_root( + monkeypatch, tmp_path, capture_list_workspaces +): + mod = importlib.import_module("scripts.watch_index_core.consistency") + captured, fake_list_workspaces = capture_list_workspaces + + monkeypatch.setattr(mod, "list_workspaces", fake_list_workspaces) + monkeypatch.setattr(mod, "_empty_dir_sweep_enabled", lambda: True) + + mod.run_empty_dir_sweep_maintenance(tmp_path) + + assert "search_root" in captured + assert Path(captured["search_root"]).resolve() == Path(tmp_path).resolve() + assert captured.get("use_qdrant_fallback") is False + + +def test_consistency_audit_skips_repairs_when_scan_is_truncated(monkeypatch, tmp_path): + mod = importlib.import_module("scripts.watch_index_core.consistency") + + workspace_root = tmp_path / "repo" + workspace_root.mkdir(parents=True, exist_ok=True) + + monkeypatch.setattr( + mod, + "list_workspaces", + lambda *a, **k: [{"workspace_path": str(workspace_root)}], + ) + monkeypatch.setattr(mod, "_consistency_audit_enabled", lambda: True) + monkeypatch.setattr(mod, "_should_run_consistency_audit", lambda *a, **k: True) + monkeypatch.setattr( + mod, + "get_collection_state_snapshot", + lambda *a, **k: {"active_collection": "coll"}, + ) + monkeypatch.setattr(mod, "_extract_repo_name_from_path", lambda *_: "repo") + monkeypatch.setattr(mod, "_load_cached_hashes", lambda *a, **k: {}) + monkeypatch.setattr( + mod, + "_scan_indexable_fs_paths", + lambda *a, **k: ({str(workspace_root / "a.py")}, True), + ) + monkeypatch.setattr( + mod, + "_load_indexed_paths_for_collection", + lambda *a, **k: ({str(workspace_root / "ghost.py")}, False), + ) + monkeypatch.setattr(mod.idx, "_Excluder", lambda *_: MagicMock()) + + enqueue_mock = MagicMock(return_value=(0, 0)) + record_mock = MagicMock() + monkeypatch.setattr(mod, "_enqueue_consistency_repairs", enqueue_mock) + monkeypatch.setattr(mod, "_record_consistency_audit", record_mock) + + mod.run_consistency_audit(MagicMock(), tmp_path) + + enqueue_mock.assert_not_called() + record_mock.assert_called_once() + summary = record_mock.call_args.args[2] + assert summary["fs_scan_truncated"] is True + assert summary["qdrant_scan_truncated"] is False + assert summary["repair_skipped_due_to_truncation"] is True + assert summary["stale_in_qdrant_count"] == 0 + assert summary["missing_in_qdrant_count"] == 0 diff --git a/tests/test_watch_index_cache.py b/tests/test_watch_index_cache.py index c5065af1..3222098c 100644 --- a/tests/test_watch_index_cache.py +++ b/tests/test_watch_index_cache.py @@ -150,3 +150,653 @@ def test_processor_delete_clears_cache_even_without_client(monkeypatch, tmp_path ) remove_mock.assert_called_once_with(str(missing), "repo") + + +def test_run_indexing_strategy_reuses_preloaded_file_state(monkeypatch, tmp_path): + proc_mod = importlib.import_module("scripts.watch_index_core.processor") + + path = tmp_path / "file.py" + path.write_text("print('x')\n", encoding="utf-8") + + monkeypatch.setattr(proc_mod.idx, "ensure_collection_and_indexes_once", lambda *a, **k: None) + monkeypatch.setattr(proc_mod, "_read_text_and_sha1", lambda _p: ("print('x')\n", "abc123")) + monkeypatch.setattr(proc_mod, "get_cached_file_hash", lambda *a, **k: None) + monkeypatch.setattr(proc_mod.idx, "detect_language", lambda _p: "python") + monkeypatch.setattr(proc_mod.idx, "should_use_smart_reindexing", lambda *a, **k: (False, "changed")) + + captured = {} + + def fake_index_single_file(*args, **kwargs): + captured.update(kwargs) + return True + + monkeypatch.setattr(proc_mod.idx, "index_single_file", fake_index_single_file) + + ok = proc_mod._run_indexing_strategy( + path, + client=MagicMock(), + model=MagicMock(), + collection="coll", + vector_name="vec", + model_dim=1, + repo_name="repo", + ) + + assert ok is True + assert captured["preloaded_text"] == "print('x')\n" + assert captured["preloaded_file_hash"] == "abc123" + assert captured["preloaded_language"] == "python" + + +def test_run_indexing_strategy_skips_ensure_for_cached_hash_match(monkeypatch, tmp_path): + proc_mod = importlib.import_module("scripts.watch_index_core.processor") + + path = tmp_path / "file.py" + path.write_text("print('x')\n", encoding="utf-8") + + ensure_mock = MagicMock() + monkeypatch.setattr(proc_mod.idx, "ensure_collection_and_indexes_once", ensure_mock) + monkeypatch.setattr(proc_mod, "_read_text_and_sha1", lambda _p: ("print('x')\n", "abc123")) + monkeypatch.setattr(proc_mod, "get_cached_file_hash", lambda *a, **k: "abc123") + monkeypatch.setattr(proc_mod.idx, "detect_language", lambda _p: "python") + + with pytest.raises(proc_mod._SkipUnchanged): + proc_mod._run_indexing_strategy( + path, + client=MagicMock(), + model=MagicMock(), + collection="coll", + vector_name="vec", + model_dim=1, + repo_name="repo", + ) + + ensure_mock.assert_not_called() + + +def test_run_indexing_strategy_force_upsert_bypasses_cached_hash_match( + monkeypatch, tmp_path +): + proc_mod = importlib.import_module("scripts.watch_index_core.processor") + + path = tmp_path / "file.py" + path.write_text("print('x')\n", encoding="utf-8") + + ensure_mock = MagicMock() + monkeypatch.setattr(proc_mod.idx, "ensure_collection_and_indexes_once", ensure_mock) + monkeypatch.setattr(proc_mod, "_read_text_and_sha1", lambda _p: ("print('x')\n", "abc123")) + monkeypatch.setattr(proc_mod, "get_cached_file_hash", lambda *a, **k: "abc123") + monkeypatch.setattr(proc_mod.idx, "detect_language", lambda _p: "python") + monkeypatch.setattr(proc_mod.idx, "should_use_smart_reindexing", lambda *a, **k: (False, "changed")) + + index_mock = MagicMock(return_value=True) + monkeypatch.setattr(proc_mod.idx, "index_single_file", index_mock) + + ok = proc_mod._run_indexing_strategy( + path, + client=MagicMock(), + model=MagicMock(), + collection="coll", + vector_name="vec", + model_dim=1, + repo_name="repo", + force_upsert=True, + ) + + assert ok is True + ensure_mock.assert_called_once() + index_mock.assert_called_once() + + +def test_run_indexing_strategy_skips_smart_path_for_markdown(monkeypatch, tmp_path): + proc_mod = importlib.import_module("scripts.watch_index_core.processor") + + path = tmp_path / "notes.md" + path.write_text("# notes\n", encoding="utf-8") + + monkeypatch.setattr(proc_mod.idx, "ensure_collection_and_indexes_once", lambda *a, **k: None) + monkeypatch.setattr(proc_mod, "_read_text_and_sha1", lambda _p: ("# notes\n", "abc123")) + monkeypatch.setattr(proc_mod, "get_cached_file_hash", lambda *a, **k: None) + monkeypatch.setattr(proc_mod.idx, "detect_language", lambda _p: "markdown") + + smart_check = MagicMock(side_effect=AssertionError("smart path must be skipped")) + monkeypatch.setattr(proc_mod.idx, "should_use_smart_reindexing", smart_check) + + captured = {} + + def fake_index_single_file(*args, **kwargs): + captured.update(kwargs) + return True + + monkeypatch.setattr(proc_mod.idx, "index_single_file", fake_index_single_file) + + ok = proc_mod._run_indexing_strategy( + path, + client=MagicMock(), + model=MagicMock(), + collection="coll", + vector_name="vec", + model_dim=1, + repo_name="repo", + ) + + assert ok is True + smart_check.assert_not_called() + assert captured["preloaded_language"] == "markdown" + + +def test_run_indexing_strategy_force_upsert_missing_points_bypasses_smart( + monkeypatch, tmp_path +): + proc_mod = importlib.import_module("scripts.watch_index_core.processor") + + path = tmp_path / "file.py" + path.write_text("print('x')\n", encoding="utf-8") + + monkeypatch.setattr(proc_mod.idx, "ensure_collection_and_indexes_once", lambda *a, **k: None) + monkeypatch.setattr(proc_mod, "_read_text_and_sha1", lambda _p: ("print('x')\n", "abc123")) + monkeypatch.setattr(proc_mod, "get_cached_file_hash", lambda *a, **k: None) + monkeypatch.setattr(proc_mod.idx, "detect_language", lambda _p: "python") + monkeypatch.setattr(proc_mod.idx, "should_use_smart_reindexing", lambda *a, **k: (True, "smart_reindex")) + monkeypatch.setattr(proc_mod.idx, "get_indexed_file_hash", lambda *a, **k: "") + monkeypatch.setattr(proc_mod, "_path_has_indexed_points", lambda *a, **k: False) + + smart_mock = MagicMock(return_value="skipped") + monkeypatch.setattr(proc_mod.idx, "process_file_with_smart_reindexing", smart_mock) + + index_mock = MagicMock(return_value=True) + monkeypatch.setattr(proc_mod.idx, "index_single_file", index_mock) + + ok = proc_mod._run_indexing_strategy( + path, + client=MagicMock(), + model=MagicMock(), + collection="coll", + vector_name="vec", + model_dim=1, + repo_name="repo", + force_upsert=True, + ) + + assert ok is True + smart_mock.assert_not_called() + index_mock.assert_called_once() + + +def test_run_indexing_strategy_sets_skip_verify_reason_for_file_lock( + monkeypatch, tmp_path +): + proc_mod = importlib.import_module("scripts.watch_index_core.processor") + + path = tmp_path / "file.py" + path.write_text("print('x')\n", encoding="utf-8") + + monkeypatch.setattr(proc_mod.idx, "ensure_collection_and_indexes_once", lambda *a, **k: None) + monkeypatch.setattr(proc_mod, "_read_text_and_sha1", lambda _p: ("print('x')\n", "abc123")) + monkeypatch.setattr(proc_mod, "get_cached_file_hash", lambda *a, **k: None) + monkeypatch.setattr(proc_mod.idx, "detect_language", lambda _p: "python") + monkeypatch.setattr(proc_mod.idx, "should_use_smart_reindexing", lambda *a, **k: (False, "changed")) + monkeypatch.setattr(proc_mod.idx, "index_single_file", lambda *a, **k: False) + monkeypatch.setattr(proc_mod.idx, "is_file_locked", lambda *_: True) + + verify_context = {} + ok = proc_mod._run_indexing_strategy( + path, + client=MagicMock(), + model=MagicMock(), + collection="coll", + vector_name="vec", + model_dim=1, + repo_name="repo", + force_upsert=True, + verify_context=verify_context, + ) + + assert ok is False + assert verify_context.get("skip_verify_reason") == "file_locked" + + +def test_finalize_journal_skips_force_upsert_verify_when_file_locked(monkeypatch): + proc_mod = importlib.import_module("scripts.watch_index_core.processor") + + verify_mock = MagicMock() + done_mock = MagicMock() + failed_mock = MagicMock() + monkeypatch.setattr(proc_mod, "_verify_and_update_journal_for_upsert", verify_mock) + monkeypatch.setattr(proc_mod, "_mark_journal_done", done_mock) + monkeypatch.setattr(proc_mod, "_mark_journal_failed", failed_mock) + + proc_mod._finalize_journal_after_index_attempt( + Path("/tmp/file.py"), + client=MagicMock(), + collection="coll", + repo_key="/tmp", + repo_name="repo", + force_upsert=True, + journal_content_hash="abc", + skip_verify_reason="file_locked", + ) + + verify_mock.assert_not_called() + done_mock.assert_not_called() + failed_mock.assert_not_called() + + +def test_staging_requires_subprocess_only_for_active_dual_root_state(monkeypatch): + proc_mod = importlib.import_module("scripts.watch_index_core.processor") + monkeypatch.setattr(proc_mod, "is_staging_enabled", lambda: True) + + assert proc_mod._staging_requires_subprocess(None) is False + assert ( + proc_mod._staging_requires_subprocess( + { + "indexing_env": {"FOO": "bar"}, + "active_repo_slug": "repo", + "serving_repo_slug": "repo", + } + ) + is False + ) + assert ( + proc_mod._staging_requires_subprocess( + { + "indexing_env": {"FOO": "bar"}, + "active_repo_slug": "repo", + "serving_repo_slug": "repo_old", + } + ) + is True + ) + assert ( + proc_mod._staging_requires_subprocess( + { + "indexing_env": {"FOO": "bar"}, + "active_repo_slug": "repo", + "serving_repo_slug": "repo", + "staging": {"collection": "repo_old_collection"}, + } + ) + is True + ) + + +def test_process_paths_does_not_force_subprocess_for_non_active_staging( + monkeypatch, tmp_path +): + proc_mod = importlib.import_module("scripts.watch_index_core.processor") + + path = tmp_path / "file.py" + path.write_text("print('x')\n", encoding="utf-8") + + monkeypatch.setattr(proc_mod, "_detect_repo_for_file", lambda p: tmp_path) + monkeypatch.setattr(proc_mod, "_get_collection_for_file", lambda p: "coll") + monkeypatch.setattr(proc_mod, "_set_status_indexing", lambda *a, **k: None) + monkeypatch.setattr(proc_mod, "persist_indexing_config", lambda *a, **k: None) + monkeypatch.setattr(proc_mod, "update_indexing_status", lambda *a, **k: None) + monkeypatch.setattr(proc_mod, "_log_activity", lambda *a, **k: None) + monkeypatch.setattr(proc_mod, "_extract_repo_name_from_path", lambda *_: "repo") + monkeypatch.setattr(proc_mod, "is_staging_enabled", lambda: True) + monkeypatch.setattr( + proc_mod, + "get_workspace_state", + lambda *a, **k: { + "indexing_env": {"FOO": "bar"}, + "active_repo_slug": "repo", + "serving_repo_slug": "repo", + }, + ) + + staging_mock = MagicMock(return_value=False) + monkeypatch.setattr(proc_mod, "_maybe_handle_staging_file", staging_mock) + monkeypatch.setattr(proc_mod, "_run_indexing_strategy", lambda *a, **k: True) + + proc_mod._process_paths( + [path], + client=MagicMock(), + model=MagicMock(), + vector_name="vec", + model_dim=1, + workspace_path=str(tmp_path), + ) + + assert staging_mock.call_args is not None + assert staging_mock.call_args.kwargs == { + "force_upsert": False, + "journal_content_hash": "", + } + assert staging_mock.call_args.args[0] == path + assert staging_mock.call_args.args[6] is None + + +def test_process_paths_uses_subprocess_when_staging_is_actually_active( + monkeypatch, tmp_path +): + proc_mod = importlib.import_module("scripts.watch_index_core.processor") + + path = tmp_path / "file.py" + path.write_text("print('x')\n", encoding="utf-8") + + monkeypatch.setattr(proc_mod, "_detect_repo_for_file", lambda p: tmp_path) + monkeypatch.setattr(proc_mod, "_get_collection_for_file", lambda p: "coll") + monkeypatch.setattr(proc_mod, "_set_status_indexing", lambda *a, **k: None) + monkeypatch.setattr(proc_mod, "persist_indexing_config", lambda *a, **k: None) + monkeypatch.setattr(proc_mod, "update_indexing_status", lambda *a, **k: None) + monkeypatch.setattr(proc_mod, "_log_activity", lambda *a, **k: None) + monkeypatch.setattr(proc_mod, "_extract_repo_name_from_path", lambda *_: "repo") + monkeypatch.setattr(proc_mod, "is_staging_enabled", lambda: True) + monkeypatch.setattr( + proc_mod, + "get_workspace_state", + lambda *a, **k: { + "indexing_env": {"FOO": "bar"}, + "active_repo_slug": "repo", + "serving_repo_slug": "repo_old", + }, + ) + + staging_mock = MagicMock(return_value=False) + monkeypatch.setattr(proc_mod, "_maybe_handle_staging_file", staging_mock) + monkeypatch.setattr(proc_mod, "_run_indexing_strategy", lambda *a, **k: True) + + proc_mod._process_paths( + [path], + client=MagicMock(), + model=MagicMock(), + vector_name="vec", + model_dim=1, + workspace_path=str(tmp_path), + ) + + assert staging_mock.call_args is not None + assert staging_mock.call_args.kwargs == { + "force_upsert": False, + "journal_content_hash": "", + } + assert staging_mock.call_args.args[0] == path + assert staging_mock.call_args.args[6] == {"FOO": "bar"} + + +def test_staging_force_upsert_hash_match_verifies_before_skip(monkeypatch, tmp_path): + proc_mod = importlib.import_module("scripts.watch_index_core.processor") + + path = tmp_path / "file.py" + path.write_text("print('x')\n", encoding="utf-8") + + monkeypatch.setattr(proc_mod, "_read_text_and_sha1", lambda _p: ("print('x')\n", "abc123")) + monkeypatch.setattr(proc_mod, "get_cached_file_hash", lambda *a, **k: "abc123") + monkeypatch.setattr(proc_mod, "_verify_upsert_committed", lambda *a, **k: True) + monkeypatch.setattr(proc_mod, "_log_activity", lambda *a, **k: None) + + mark_done = MagicMock() + monkeypatch.setattr(proc_mod, "_mark_journal_done", mark_done) + advance = MagicMock() + monkeypatch.setattr(proc_mod, "_advance_progress", advance) + + handled = proc_mod._maybe_handle_staging_file( + path, + MagicMock(), + "coll", + "repo", + str(tmp_path), + [path], + {"FOO": "bar"}, + {str(tmp_path): 0}, + "started", + force_upsert=True, + journal_content_hash="abc123", + ) + + assert handled is True + mark_done.assert_called_once_with(path, str(tmp_path), "repo") + advance.assert_called_once() + + +def test_runtime_root_override_updates_internal_path_checks(monkeypatch, tmp_path): + import scripts.watch_index as watch_index + from scripts.watch_index_core import config as watch_config + import scripts.watch_index_core.processor as proc_mod + import scripts.embedder as embedder_mod + + runtime_root = tmp_path / "runtime-root" + runtime_root.mkdir(parents=True, exist_ok=True) + internal = runtime_root / ".git" / "HEAD" + internal.parent.mkdir(parents=True, exist_ok=True) + internal.write_text("ref: refs/heads/main\n", encoding="utf-8") + + original_root = watch_config.ROOT + original_watch_root = watch_index.ROOT + monkeypatch.setenv("WATCH_ROOT", str(runtime_root)) + monkeypatch.setattr(watch_index, "initialize_watcher_state", lambda root: {"repo_name": None}) + monkeypatch.setattr(watch_index, "get_indexing_config_snapshot", lambda repo_name=None: {}) + monkeypatch.setattr(watch_index, "compute_indexing_config_hash", lambda snapshot: "hash") + monkeypatch.setattr(watch_index, "persist_indexing_config", lambda *a, **k: None) + monkeypatch.setattr(watch_index, "update_indexing_status", lambda *a, **k: None) + monkeypatch.setattr(embedder_mod, "get_embedding_model", lambda *_: MagicMock()) + monkeypatch.setattr(embedder_mod, "get_model_dimension", lambda *_: 1) + monkeypatch.setattr(watch_index, "resolve_vector_name_config", lambda *a, **k: "vec") + monkeypatch.setattr(watch_index, "_start_pseudo_backfill_worker", lambda *a, **k: None) + monkeypatch.setattr(watch_index, "create_observer", lambda *a, **k: MagicMock()) + monkeypatch.setattr(watch_index, "IndexHandler", MagicMock()) + monkeypatch.setattr(watch_index, "ChangeQueue", MagicMock()) + monkeypatch.setattr( + watch_index, + "QdrantClient", + MagicMock(return_value=MagicMock(get_collection=MagicMock())), + ) + monkeypatch.setattr(watch_index, "run_consistency_audit", lambda *a, **k: None) + monkeypatch.setattr(watch_index, "run_empty_dir_sweep_maintenance", lambda *a, **k: None) + monkeypatch.setattr(watch_index, "list_pending_index_journal_entries", lambda *a, **k: []) + def _bool_env(name, default=False): + if name == "WATCH_JOURNAL_DRAIN_ENABLED": + return True + return False + + monkeypatch.setattr(watch_index, "get_boolean_env", _bool_env) + monkeypatch.setattr(watch_index, "_sleep", lambda *_: (_ for _ in ()).throw(KeyboardInterrupt())) + + try: + watch_index.main() + except KeyboardInterrupt: + pass + + try: + assert watch_config.ROOT == runtime_root.resolve() + assert proc_mod._is_internal_ignored_path(internal) is True + finally: + watch_config.ROOT = original_root + watch_index.ROOT = original_watch_root + + +def test_journal_drain_does_not_requeue_while_queue_is_busy(monkeypatch): + import scripts.watch_index as watch_index + + queue = MagicMock() + queue.stats.return_value = { + "queued": 10, + "pending": 20, + "forced": 10, + "pending_forced": 20, + "processing": True, + } + pending_mock = MagicMock() + monkeypatch.setattr(watch_index, "list_pending_index_journal_entries", pending_mock) + + watch_index._drain_pending_journal(queue) + + pending_mock.assert_not_called() + queue.add.assert_not_called() + + +def test_journal_drain_limits_one_pass(monkeypatch): + import scripts.watch_index as watch_index + + monkeypatch.setenv("WATCH_JOURNAL_DRAIN_BATCH_SIZE", "2") + monkeypatch.setattr(watch_index, "_maybe_log_journal_drain", lambda **_: None) + entries = [ + {"path": f"/work/repo-{idx}/file.py", "op_type": "upsert"} + for idx in range(3) + ] + monkeypatch.setattr(watch_index, "list_pending_index_journal_entries", lambda *_: entries) + + queue = MagicMock() + queue.stats.return_value = { + "queued": 0, + "pending": 0, + "forced": 0, + "pending_forced": 0, + "processing": False, + } + + watch_index._drain_pending_journal(queue) + + assert queue.add.call_count == 2 + + +def test_journal_status_batch_flushes_once_per_pass(monkeypatch, tmp_path): + import scripts.watch_index_core.processor as processor + + bulk_update = MagicMock() + single_update = MagicMock() + monkeypatch.setattr(processor, "update_index_journal_entries_status", bulk_update) + monkeypatch.setattr(processor, "update_index_journal_entry_status", single_update) + + batch = [ + {"repo_key": "/work/repo", "repo_name": "repo", "path": str(tmp_path / "one.py"), "status": "done"}, + {"repo_key": "/work/repo", "repo_name": "repo", "path": str(tmp_path / "two.py"), "status": "done"}, + ] + processor._flush_journal_status_batch(batch) + assert bulk_update.call_count == 1 + assert single_update.call_count == 0 + + +def test_journal_status_batch_falls_back_to_single_updates(monkeypatch, tmp_path): + import scripts.watch_index_core.processor as processor + + bulk_update = MagicMock(side_effect=RuntimeError("bulk write failed")) + single_update = MagicMock() + monkeypatch.setattr(processor, "update_index_journal_entries_status", bulk_update) + monkeypatch.setattr(processor, "update_index_journal_entry_status", single_update) + + batch = [ + {"repo_key": "/work/repo", "repo_name": "repo", "path": str(tmp_path / "one.py"), "status": "done"}, + { + "repo_key": "/work/repo", + "repo_name": "repo", + "path": str(tmp_path / "two.py"), + "status": "failed", + "error": "boom", + "remove_on_done": False, + }, + ] + processor._flush_journal_status_batch(batch) + + assert bulk_update.call_count == 1 + assert single_update.call_count == 2 + assert single_update.call_args_list[0].kwargs["status"] == "done" + assert single_update.call_args_list[1].kwargs == { + "status": "failed", + "error": "boom", + "workspace_path": "/work/repo", + "repo_name": "repo", + "remove_on_done": False, + } + + +def test_main_throttles_periodic_maintenance(monkeypatch, tmp_path): + import scripts.watch_index as watch_index + from scripts.watch_index_core import config as watch_config + import scripts.embedder as embedder_mod + + runtime_root = tmp_path / "runtime-root" + runtime_root.mkdir(parents=True, exist_ok=True) + + original_root = watch_config.ROOT + original_watch_root = watch_index.ROOT + monkeypatch.setenv("WATCH_ROOT", str(runtime_root)) + monkeypatch.setenv("WATCH_MAINTENANCE_INTERVAL_SECS", "300") + monkeypatch.setenv("WATCH_INIT_MAINTENANCE_ENABLED", "0") + monkeypatch.setattr(watch_index, "initialize_watcher_state", lambda *a, **k: {"repo_name": None}) + monkeypatch.setattr(watch_index, "get_indexing_config_snapshot", lambda repo_name=None: {}) + monkeypatch.setattr(watch_index, "compute_indexing_config_hash", lambda snapshot: "hash") + monkeypatch.setattr(watch_index, "persist_indexing_config", lambda *a, **k: None) + monkeypatch.setattr(watch_index, "update_indexing_status", lambda *a, **k: None) + monkeypatch.setattr(embedder_mod, "get_embedding_model", lambda *_: MagicMock()) + monkeypatch.setattr(embedder_mod, "get_model_dimension", lambda *_: 1) + monkeypatch.setattr(watch_index, "resolve_vector_name_config", lambda *a, **k: "vec") + monkeypatch.setattr(watch_index, "_start_pseudo_backfill_worker", lambda *a, **k: None) + + class FakeObserver: + def schedule(self, *a, **k): + return None + + def start(self): + return None + + def stop(self): + return None + + def join(self): + return None + + monkeypatch.setattr(watch_index, "create_observer", lambda *a, **k: FakeObserver()) + monkeypatch.setattr(watch_index, "IndexHandler", MagicMock()) + monkeypatch.setattr(watch_index, "ChangeQueue", MagicMock()) + monkeypatch.setattr( + watch_index, + "QdrantClient", + MagicMock(return_value=MagicMock(get_collection=MagicMock())), + ) + def _bool_env(name, default=False): + if name == "WATCH_JOURNAL_DRAIN_ENABLED": + return True + return False + + monkeypatch.setattr(watch_index, "get_boolean_env", _bool_env) + + drain_mock = MagicMock() + maintenance_mock = MagicMock() + monkeypatch.setattr(watch_index, "_drain_pending_journal", drain_mock) + monkeypatch.setattr(watch_index, "_run_periodic_maintenance", maintenance_mock) + + time_values = iter([0.0, 1.0, 2.0, 301.0]) + monkeypatch.setattr(watch_index.time, "time", lambda: next(time_values)) + + sleep_calls = {"count": 0} + + def _sleep(_secs): + sleep_calls["count"] += 1 + if sleep_calls["count"] >= 4: + raise KeyboardInterrupt() + + monkeypatch.setattr(watch_index, "_sleep", _sleep) + + try: + watch_index.main() + finally: + watch_config.ROOT = original_root + watch_index.ROOT = original_watch_root + + assert drain_mock.call_count == 4 + assert maintenance_mock.call_count == 2 + + +def test_watch_source_defaults_follow_repo_mode(monkeypatch): + import scripts.watch_index as watch_index + + monkeypatch.delenv("WATCH_JOURNAL_DRAIN_ENABLED", raising=False) + monkeypatch.delenv("WATCH_FS_EVENTS_ENABLED", raising=False) + + assert watch_index._journal_drain_enabled(True) is True + assert watch_index._fs_events_enabled(True) is False + assert watch_index._journal_drain_enabled(False) is False + assert watch_index._fs_events_enabled(False) is True + + +def test_watch_source_env_overrides_defaults(monkeypatch): + import scripts.watch_index as watch_index + + monkeypatch.setenv("WATCH_JOURNAL_DRAIN_ENABLED", "0") + monkeypatch.setenv("WATCH_FS_EVENTS_ENABLED", "1") + + assert watch_index._journal_drain_enabled(True) is False + assert watch_index._fs_events_enabled(True) is True diff --git a/tests/test_watch_index_git_history.py b/tests/test_watch_index_git_history.py new file mode 100644 index 00000000..08cac353 --- /dev/null +++ b/tests/test_watch_index_git_history.py @@ -0,0 +1,48 @@ +import io +import subprocess + +import pytest + + +pytestmark = pytest.mark.unit + + +def test_git_history_ingest_runs_as_package_module(monkeypatch, tmp_path): + from scripts.watch_index_core import processor + from scripts.watch_index_core import config as watch_config + + manifest = tmp_path / "git_history.json" + manifest.write_text('{"commits": []}', encoding="utf-8") + + captured = {} + + class FakePopen: + def __init__(self, cmd, **kwargs): + captured["cmd"] = cmd + captured["cwd"] = kwargs.get("cwd") + captured["env"] = kwargs.get("env") + self.stdout = io.StringIO("") + self.stderr = io.StringIO("") + + def poll(self): + return 0 + + def wait(self, timeout=None): + return 0 + + def kill(self): + pass + + monkeypatch.setattr(subprocess, "Popen", FakePopen) + + processor._run_git_history_ingest( + manifest, + collection="Context-Engine-41e67959", + repo_name="Context-Engine-41e67959950c8ab3", + ) + + assert captured["cmd"][:3] == [processor.sys.executable or "python3", "-m", "scripts.ingest_history"] + assert "--manifest-json" in captured["cmd"] + assert captured["cwd"] == str(watch_config.ROOT_DIR) + assert captured["env"]["COLLECTION_NAME"] == "Context-Engine-41e67959" + assert captured["env"]["REPO_NAME"] == "Context-Engine-41e67959950c8ab3" diff --git a/tests/test_watch_init_maintenance.py b/tests/test_watch_init_maintenance.py new file mode 100644 index 00000000..9c9248e1 --- /dev/null +++ b/tests/test_watch_init_maintenance.py @@ -0,0 +1,71 @@ +import importlib +import subprocess + + +def test_init_maintenance_interval_defaults_to_two_hours(monkeypatch): + monkeypatch.delenv("WATCH_INIT_MAINTENANCE_INTERVAL_MINUTES", raising=False) + monkeypatch.delenv("INIT_MAINTENANCE_INTERVAL_MINUTES", raising=False) + + mod = importlib.import_module("scripts.watch_index_core.init_maintenance") + mod = importlib.reload(mod) + + assert mod._interval_seconds() == 120 * 60 + + +def test_init_maintenance_runs_existing_scripts_under_lock(monkeypatch, tmp_path): + mod = importlib.import_module("scripts.watch_index_core.init_maintenance") + mod = importlib.reload(mod) + + calls = [] + + def fake_run(command, **kwargs): + calls.append((command, kwargs)) + return subprocess.CompletedProcess(command, 0, stdout="ok", stderr="") + + monkeypatch.setattr(mod.subprocess, "run", fake_run) + monkeypatch.setenv("WATCH_INIT_MAINTENANCE_COMMAND_TIMEOUT_SECS", "7") + + commands = [ + ["wait-for-qdrant.sh"], + ["python", "create_indexes.py"], + ["python", "warm_all_collections.py"], + ["python", "health_check.py"], + ] + + ok = mod.run_init_maintenance_once(commands=commands, lock_path=tmp_path / "init.lock") + + assert ok is True + assert [call[0] for call in calls] == commands + assert all(call[1]["timeout"] == 7 for call in calls) + assert all(call[1]["check"] is False for call in calls) + assert all("PYTHONPATH" in call[1]["env"] for call in calls) + + +def test_init_maintenance_stops_sequence_on_failure(monkeypatch, tmp_path): + mod = importlib.import_module("scripts.watch_index_core.init_maintenance") + mod = importlib.reload(mod) + + calls = [] + + def fake_run(command, **kwargs): + calls.append(command) + return subprocess.CompletedProcess(command, 1, stdout="", stderr="boom") + + monkeypatch.setattr(mod.subprocess, "run", fake_run) + + ok = mod.run_init_maintenance_once( + commands=[["first"], ["second"]], + lock_path=tmp_path / "init.lock", + ) + + assert ok is False + assert calls == [["first"]] + + +def test_init_maintenance_worker_can_be_disabled(monkeypatch): + mod = importlib.import_module("scripts.watch_index_core.init_maintenance") + mod = importlib.reload(mod) + + monkeypatch.setenv("WATCH_INIT_MAINTENANCE_ENABLED", "0") + + assert mod.start_init_maintenance_worker() is None diff --git a/tests/test_watch_queue.py b/tests/test_watch_queue.py new file mode 100644 index 00000000..7f02f25d --- /dev/null +++ b/tests/test_watch_queue.py @@ -0,0 +1,110 @@ +def test_change_queue_suppresses_recent_identical_fingerprint(monkeypatch, tmp_path): + from scripts.watch_index_core import queue as queue_mod + + monkeypatch.setattr(queue_mod, "RECENT_FINGERPRINT_TTL_SECS", 10.0) + + processed = [] + q = queue_mod.ChangeQueue(lambda paths: processed.append(list(paths))) + + p = tmp_path / "file.py" + p.write_text("print('x')\n", encoding="utf-8") + + q._paths.add(p) + q._flush() + assert processed == [[p]] + + q._paths.add(p) + q._flush() + assert processed == [[p]] + + +def test_change_queue_reprocesses_when_fingerprint_changes(monkeypatch, tmp_path): + from scripts.watch_index_core import queue as queue_mod + + monkeypatch.setattr(queue_mod, "RECENT_FINGERPRINT_TTL_SECS", 10.0) + + processed = [] + q = queue_mod.ChangeQueue(lambda paths: processed.append(list(paths))) + + p = tmp_path / "file.py" + p.write_text("print('x')\n", encoding="utf-8") + + q._paths.add(p) + q._flush() + + p.write_text("print('changed-again')\n", encoding="utf-8") + q._paths.add(p) + q._flush() + + assert processed == [[p], [p]] + + +def test_change_queue_force_bypasses_recent_fingerprint_suppression(monkeypatch, tmp_path): + from scripts.watch_index_core import queue as queue_mod + + monkeypatch.setattr(queue_mod, "RECENT_FINGERPRINT_TTL_SECS", 10.0) + + processed = [] + q = queue_mod.ChangeQueue(lambda paths: processed.append(list(paths))) + + p = tmp_path / "file.py" + p.write_text("print('x')\n", encoding="utf-8") + + q.add(p) + q._flush() + q.add(p, force=True) + q._flush() + + assert processed == [[p], [p]] + + +def test_change_queue_repeated_same_path_does_not_rearm_timer(monkeypatch, tmp_path): + from scripts.watch_index_core import queue as queue_mod + + class FakeTimer: + created = 0 + canceled = 0 + + def __init__(self, _delay, _cb): + FakeTimer.created += 1 + self.daemon = False + + def start(self): + return None + + def cancel(self): + FakeTimer.canceled += 1 + + monkeypatch.setattr(queue_mod.threading, "Timer", FakeTimer) + + q = queue_mod.ChangeQueue(lambda _paths: None) + p = tmp_path / "file.py" + p.write_text("print('x')\n", encoding="utf-8") + + q.add(p, force=True) + q.add(p, force=True) + q.add(p, force=True) + + assert FakeTimer.created == 1 + assert FakeTimer.canceled == 0 + + +def test_change_queue_stats_reports_backlog(tmp_path): + from scripts.watch_index_core import queue as queue_mod + + q = queue_mod.ChangeQueue(lambda _paths: None) + p = tmp_path / "file.py" + p.write_text("print('x')\n", encoding="utf-8") + + q._paths.add(p) + q._forced_paths.add(p) + q._pending.add(p.with_name("pending.py")) + q._pending_forced.add(p.with_name("pending.py")) + + assert q.stats() == { + "queued": 1, + "pending": 1, + "forced": 1, + "pending_forced": 1, + "processing": False, + } diff --git a/tests/test_watcher_collection_resolution.py b/tests/test_watcher_collection_resolution.py index fa3d0c1a..cc9e464a 100644 --- a/tests/test_watcher_collection_resolution.py +++ b/tests/test_watcher_collection_resolution.py @@ -6,9 +6,9 @@ pytestmark = pytest.mark.unit def test_main_resolves_collection_from_state(monkeypatch, tmp_path): - # Env setup: placeholder collection name at startup + # Env setup: default collection name at startup monkeypatch.setenv("WATCH_ROOT", str(tmp_path)) - monkeypatch.setenv("COLLECTION_NAME", "my-collection") + monkeypatch.setenv("COLLECTION_NAME", "codebase") monkeypatch.setenv("QDRANT_URL", "http://localhost:6333") monkeypatch.setenv("EMBEDDING_MODEL", "fake") @@ -17,6 +17,9 @@ def test_main_resolves_collection_from_state(monkeypatch, tmp_path): wi = importlib.import_module("scripts.watch_index") # Reload to re-read env defaults (COLLECTION) in module globals wi = importlib.reload(wi) + watch_config = importlib.import_module("scripts.watch_index_core.config") + original_root = watch_config.ROOT + original_watch_root = wi.ROOT # Fake QdrantClient: force get_collection to raise so code chooses sanitized vector name path class FakeQdrant: @@ -64,16 +67,20 @@ def join(self): # Make the main loop exit immediately by raising KeyboardInterrupt on sleep def _raise_kb(_): raise KeyboardInterrupt() - monkeypatch.setattr(wi.time, "sleep", _raise_kb, raising=True) + monkeypatch.setattr(wi, "_sleep", _raise_kb, raising=True) - # Precondition: module-level COLLECTION should reflect placeholder at import time - assert wi.COLLECTION == os.environ.get("COLLECTION_NAME") == "my-collection" + # Precondition: module-level COLLECTION should reflect the configured default at import time + assert wi.COLLECTION == os.environ.get("COLLECTION_NAME") == "codebase" # Run main(); in single-repo mode it should keep the env-provided COLLECTION_NAME - wi.main() + try: + wi.main() - # Postcondition: global COLLECTION remains the env-provided name - assert wi.COLLECTION == "my-collection" + # Postcondition: global COLLECTION remains the env-provided name + assert wi.COLLECTION == "codebase" + finally: + watch_config.ROOT = original_root + wi.ROOT = original_watch_root def test_multi_repo_ignores_placeholder_collection_in_state(monkeypatch, tmp_path): @@ -85,7 +92,8 @@ def test_multi_repo_ignores_placeholder_collection_in_state(monkeypatch, tmp_pat utils = importlib.import_module("scripts.watch_index_core.utils") utils = importlib.reload(utils) - monkeypatch.setattr(utils, "ROOT", tmp_path, raising=False) + watch_config = importlib.import_module("scripts.watch_index_core.config") + monkeypatch.setattr(watch_config, "ROOT", tmp_path, raising=True) monkeypatch.setattr(utils, "is_multi_repo_mode", lambda: True, raising=True) repo_slug = "Pirate Survivors-2b23a7e45f2c4b9f" @@ -111,4 +119,3 @@ def _fake_get_workspace_state(ws_path: str, repo_name: str | None = None): resolved = utils._get_collection_for_file(target) assert resolved == f"derived-{repo_slug}" - diff --git a/tests/test_watcher_events.py b/tests/test_watcher_events.py index f01484e9..658366bc 100644 --- a/tests/test_watcher_events.py +++ b/tests/test_watcher_events.py @@ -71,6 +71,25 @@ def test_on_moved_enqueues_new_dest(monkeypatch, tmp_path): assert any(s.endswith("/b.py") for s in q.added) +@pytest.mark.unit +def test_on_moved_ignores_internal_codebase_paths(monkeypatch, tmp_path): + monkeypatch.setenv("MULTI_REPO_MODE", "0") + q = FakeQueue() + handler = wi.IndexHandler(root=tmp_path, queue=q, client=FakeClient(), collection="c") + + codebase = tmp_path / ".codebase" + codebase.mkdir(parents=True, exist_ok=True) + src = codebase / "state.json" + dst = codebase / "file_locks" / "abc.lock" + src.write_text("{}\n") + dst.parent.mkdir(parents=True, exist_ok=True) + dst.write_text("lock\n") + + handler.on_moved(E(src, dest=dst)) + + assert q.added == [] + + @pytest.mark.unit def test_ignore_reload_rebuilds_excluder(monkeypatch, tmp_path): monkeypatch.setenv("MULTI_REPO_MODE", "0") @@ -105,4 +124,3 @@ def test_remote_git_manifest_is_enqueued_even_if_excluded(monkeypatch, tmp_path) handler.on_created(E(manifest)) assert any(p.endswith("/.remote-git/git_history_test.json") for p in q.added) - diff --git a/tests/test_workspace_state.py b/tests/test_workspace_state.py index 1200a270..c78188cb 100644 --- a/tests/test_workspace_state.py +++ b/tests/test_workspace_state.py @@ -129,10 +129,10 @@ def test_multi_repo_does_not_hard_override_with_env(self, ws_module, monkeypatch def test_single_repo_env_override_preserved(self, ws_module, monkeypatch): """In single-repo mode, COLLECTION_NAME remains a master override.""" monkeypatch.delenv("MULTI_REPO_MODE", raising=False) - monkeypatch.setenv("COLLECTION_NAME", "codebase") + monkeypatch.setenv("COLLECTION_NAME", "custom") ws = importlib.reload(ws_module) - assert ws.get_collection_name("my-repo_old") == "codebase_old" + assert ws.get_collection_name("my-repo_old") == "custom_old" def test_multi_repo_workspace_level_env_override_still_applies(self, ws_module, monkeypatch): """When repo_name is None, env override should still apply even in multi-repo mode.""" @@ -142,6 +142,38 @@ def test_multi_repo_workspace_level_env_override_still_applies(self, ws_module, assert ws.get_collection_name(None) == "codebase" + def test_multi_repo_workspace_root_path_uses_configured_collection(self, ws_module, monkeypatch, tmp_path): + """The multi-repo workspace root is not a repository identity.""" + ws_root = tmp_path / "work" + ws_root.mkdir(parents=True, exist_ok=True) + monkeypatch.setenv("MULTI_REPO_MODE", "1") + monkeypatch.setenv("WORKSPACE_PATH", str(ws_root)) + monkeypatch.setenv("WATCH_ROOT", str(ws_root)) + monkeypatch.setenv("COLLECTION_NAME", "context-engine") + ws = importlib.reload(ws_module) + + assert ws.get_collection_name(str(ws_root)) == "context-engine" + + def test_multi_repo_upload_managed_detection_does_not_probe_git(self, ws_module, monkeypatch, tmp_path): + """Upload-managed multi-repo identity comes from workspace path, not git metadata.""" + ws_root = tmp_path / "work" + repo_root = ws_root / "repo-a" + (repo_root / ".git").mkdir(parents=True, exist_ok=True) + monkeypatch.setenv("MULTI_REPO_MODE", "1") + monkeypatch.setenv("WORKSPACE_PATH", str(ws_root)) + monkeypatch.setenv("WATCH_ROOT", str(ws_root)) + monkeypatch.delenv("CTXCE_BINDMOUNT_REPO_DETECTION", raising=False) + ws = importlib.reload(ws_module) + + monkeypatch.setattr( + ws, + "_git_remote_repo_name", + lambda *_: pytest.fail("git inference should be disabled for upload-managed mode"), + ) + + assert ws._extract_repo_name_from_path(str(repo_root)) == "repo-a" + assert ws._extract_repo_name_from_path(str(ws_root)) == "" + # ============================================================================ # Tests: Environment Variable Helpers @@ -430,6 +462,166 @@ def test_state_filename(self, ws_module): def test_placeholder_collection_names(self, ws_module): """PLACEHOLDER_COLLECTION_NAMES contains expected values.""" - assert "" in ws_module.PLACEHOLDER_COLLECTION_NAMES - assert "default-collection" in ws_module.PLACEHOLDER_COLLECTION_NAMES - assert "my-collection" in ws_module.PLACEHOLDER_COLLECTION_NAMES + assert ws_module.PLACEHOLDER_COLLECTION_NAMES == {"", "codebase"} + + +class TestCompareSymbolChanges: + def test_compare_symbol_changes_tolerates_line_shift_for_unchanged_content(self, ws_module): + old_symbols = { + "function_foo_10": { + "name": "foo", + "type": "function", + "start_line": 10, + "end_line": 20, + "content_hash": "samehash", + } + } + new_symbols = { + "function_foo_12": { + "name": "foo", + "type": "function", + "start_line": 12, + "end_line": 22, + "content_hash": "samehash", + } + } + + unchanged, changed = ws_module.compare_symbol_changes(old_symbols, new_symbols) + + assert unchanged == ["function_foo_12"] + assert changed == [] + + +class TestSymbolCachePaths: + def test_symbol_cache_uses_shared_repo_state_dir_in_multi_repo_mode(self, monkeypatch, tmp_path): + ws_root = tmp_path / "work" + repo_name = "repo-1234567890abcdef" + repo_root = ws_root / repo_name + repo_root.mkdir(parents=True, exist_ok=True) + + monkeypatch.setenv("WORKSPACE_PATH", str(ws_root)) + monkeypatch.setenv("WATCH_ROOT", str(ws_root)) + monkeypatch.setenv("MULTI_REPO_MODE", "1") + + import importlib + + ws_module = importlib.import_module("scripts.workspace_state") + ws_module = importlib.reload(ws_module) + + file_path = repo_root / "src" / "app.py" + file_path.parent.mkdir(parents=True, exist_ok=True) + file_path.write_text("print('x')\n", encoding="utf-8") + + expected_hash = ws_module.hashlib.md5( + str(file_path.resolve()).encode("utf-8") + ).hexdigest()[:8] + cache_path = ws_module._get_symbol_cache_path(str(file_path)) + + assert cache_path == ( + ws_root + / ".codebase" + / "repos" + / repo_name + / "symbols" + / f"{expected_hash}.json" + ) + + def test_symbol_cache_write_uses_cross_user_writable_mode(self, monkeypatch, tmp_path): + ws_root = tmp_path / "work" + repo_name = "repo-1234567890abcdef" + repo_root = ws_root / repo_name + repo_root.mkdir(parents=True, exist_ok=True) + + monkeypatch.setenv("WORKSPACE_PATH", str(ws_root)) + monkeypatch.setenv("WATCH_ROOT", str(ws_root)) + monkeypatch.setenv("MULTI_REPO_MODE", "1") + + import importlib + + ws_module = importlib.import_module("scripts.workspace_state") + ws_module = importlib.reload(ws_module) + + file_path = repo_root / "src" / "cacheme.py" + file_path.parent.mkdir(parents=True, exist_ok=True) + file_path.write_text("print('x')\n", encoding="utf-8") + + ws_module.set_cached_symbols(str(file_path), {"sym": {"name": "sym"}}, "abc123") + cache_path = ws_module._get_symbol_cache_path(str(file_path)) + + assert cache_path.exists() + if os.name == "nt": + pytest.skip("POSIX permission bits are not stable on Windows") + dir_mode = cache_path.parent.stat().st_mode & 0o777 + file_mode = cache_path.stat().st_mode & 0o777 + assert dir_mode & 0o700 == 0o700 + assert file_mode & 0o600 == 0o600 + + +class TestCollectionMappings: + def test_get_collection_mappings_accepts_codebase_root_search_path(self, monkeypatch, tmp_path): + ws_root = tmp_path / "work" + ws_root.mkdir(parents=True, exist_ok=True) + slug = "repo-1234567890abcdef" + global_state_dir = ws_root / ".codebase" / "repos" / slug + global_state_dir.mkdir(parents=True, exist_ok=True) + global_state_path = global_state_dir / "state.json" + global_state_path.write_text( + json.dumps( + { + "qdrant_collection": "repo-123456-abcdef", + "updated_at": "2026-03-08T00:00:00", + } + ), + encoding="utf-8", + ) + + monkeypatch.setenv("WORKSPACE_PATH", str(ws_root)) + monkeypatch.setenv("WATCH_ROOT", str(ws_root)) + monkeypatch.setenv("MULTI_REPO_MODE", "1") + + import importlib + + ws_module = importlib.import_module("scripts.workspace_state") + ws_module = importlib.reload(ws_module) + + mappings = ws_module.get_collection_mappings(search_root=str(ws_root / ".codebase")) + slug_entries = [m for m in mappings if str(m.get("repo_name")) == slug] + + assert slug_entries, "expected global repo mapping to be discovered from codebase root" + entry = slug_entries[0] + assert entry["collection_name"] == "repo-123456-abcdef" + assert Path(entry["state_file"]).resolve() == global_state_path.resolve() + + def test_get_collection_mappings_keeps_global_repo_state_behavior(self, monkeypatch, tmp_path): + ws_root = tmp_path / "work" + ws_root.mkdir(parents=True, exist_ok=True) + repo_name = "frontend" + global_state_dir = ws_root / ".codebase" / "repos" / repo_name + global_state_dir.mkdir(parents=True, exist_ok=True) + global_state_path = global_state_dir / "state.json" + global_state_path.write_text( + json.dumps( + { + "qdrant_collection": "frontend-abcdef", + "updated_at": "2026-03-08T00:00:00", + } + ), + encoding="utf-8", + ) + + monkeypatch.setenv("WORKSPACE_PATH", str(ws_root)) + monkeypatch.setenv("WATCH_ROOT", str(ws_root)) + monkeypatch.setenv("MULTI_REPO_MODE", "1") + + import importlib + + ws_module = importlib.import_module("scripts.workspace_state") + ws_module = importlib.reload(ws_module) + + mappings = ws_module.get_collection_mappings(search_root=str(ws_root)) + repo_entries = [m for m in mappings if str(m.get("repo_name")) == repo_name] + + assert repo_entries, "expected global repo mapping to be discovered" + entry = repo_entries[0] + assert entry["collection_name"] == "frontend-abcdef" + assert Path(entry["state_file"]).resolve() == global_state_path.resolve() diff --git a/vscode-extension/build/build.bat b/vscode-extension/build/build.bat index 8db62e3f..26b616e3 100644 --- a/vscode-extension/build/build.bat +++ b/vscode-extension/build/build.bat @@ -15,9 +15,7 @@ set "STAGE_DIR=%OUT_DIR%\extension-stage" set "BUILD_RESULT=0" for %%I in ("..\..\ctx-hook-simple.sh") do set "HOOK_SRC=%%~fI" for %%I in ("..\..\scripts\ctx.py") do set "CTX_SRC=%%~fI" -for %%I in ("..\..\scripts\mcp_router.py") do set "ROUTER_SRC=%%~fI" for %%I in ("..\..\scripts\refrag_glm.py") do set "REFRAG_SRC=%%~fI" -for %%I in ("..\..\scripts\mcp_router.py") do set "ROUTER_SRC=%%~fI" for %%I in ("..\..\.env.example") do set "ENV_EXAMPLE_SRC=%%~fI" echo Building clean Context Engine Uploader extension... @@ -64,9 +62,7 @@ if errorlevel 1 ( REM Bundle ctx hook script and ctx CLI into the staged extension for reference if exist "%HOOK_SRC%" copy /Y "%HOOK_SRC%" "%STAGE_DIR%\ctx-hook-simple.sh" >nul if exist "%CTX_SRC%" copy /Y "%CTX_SRC%" "%STAGE_DIR%\ctx.py" >nul -if exist "%ROUTER_SRC%" copy /Y "%ROUTER_SRC%" "%STAGE_DIR%\mcp_router.py" >nul if exist "%REFRAG_SRC%" copy /Y "%REFRAG_SRC%" "%STAGE_DIR%\refrag_glm.py" >nul -if exist "%ROUTER_SRC%" copy /Y "%ROUTER_SRC%" "%STAGE_DIR%\mcp_router.py" >nul if exist "%ENV_EXAMPLE_SRC%" copy /Y "%ENV_EXAMPLE_SRC%" "%STAGE_DIR%\env.example" >nul REM Optional: bundle Python dependencies into the staged extension when requested diff --git a/vscode-extension/build/build.sh b/vscode-extension/build/build.sh index f3e4d9fa..13ad8614 100755 --- a/vscode-extension/build/build.sh +++ b/vscode-extension/build/build.sh @@ -7,11 +7,9 @@ OUT_DIR="$SCRIPT_DIR/../out" SRC_SCRIPT="$SCRIPT_DIR/../../scripts/standalone_upload_client.py" CLIENT="standalone_upload_client.py" STAGE_DIR="$OUT_DIR/extension-stage" -BUNDLE_DEPS="${1:-}" PYTHON_BIN="${PYTHON_BIN:-python3}" HOOK_SRC="$SCRIPT_DIR/../../ctx-hook-simple.sh" CTX_SRC="$SCRIPT_DIR/../../scripts/ctx.py" -ROUTER_SRC="$SCRIPT_DIR/../../scripts/mcp_router.py" REFRAG_SRC="$SCRIPT_DIR/../../scripts/refrag_glm.py" ENV_EXAMPLE_SRC="$SCRIPT_DIR/../../.env.example" AUTH_SRC="$SCRIPT_DIR/../../scripts/upload_auth_utils.py" @@ -48,9 +46,6 @@ fi if [[ -f "$CTX_SRC" ]]; then cp "$CTX_SRC" "$STAGE_DIR/ctx.py" fi -if [[ -f "$ROUTER_SRC" ]]; then - cp "$ROUTER_SRC" "$STAGE_DIR/mcp_router.py" -fi if [[ -f "$REFRAG_SRC" ]]; then cp "$REFRAG_SRC" "$STAGE_DIR/refrag_glm.py" fi @@ -64,16 +59,47 @@ if [[ -f "$ENV_EXAMPLE_SRC" ]]; then cp "$ENV_EXAMPLE_SRC" "$STAGE_DIR/env.example" fi -# Optional: bundle Python deps into the staged extension when requested -if [[ "$BUNDLE_DEPS" == "--bundle-deps" ]]; then - echo "Bundling Python dependencies into staged extension using $PYTHON_BIN..." - # On macOS, urllib3 v2 + system LibreSSL emits NotOpenSSLWarning; pin <2 there. - if [[ "$(uname -s)" == "Darwin" ]]; then - echo "Detected macOS; pinning urllib3<2 to avoid LibreSSL/OpenSSL warning." - "$PYTHON_BIN" -m pip install -t "$STAGE_DIR/python_libs" "urllib3<2" requests charset_normalizer "openai>=1.0" watchdog +# Bundle Python deps into the staged extension. Runtime assumes bundled +# python_libs are present and only requires an installed Python interpreter. +echo "Bundling Python dependencies into staged extension using $PYTHON_BIN..." +rm -rf "$STAGE_DIR/python_libs" +# On macOS, urllib3 v2 + system LibreSSL emits NotOpenSSLWarning; pin <2 there. +if [[ "$(uname -s)" == "Darwin" ]]; then + echo "Detected macOS; pinning urllib3<2 to avoid LibreSSL/OpenSSL warning." + "$PYTHON_BIN" -m pip install -t "$STAGE_DIR/python_libs" "urllib3<2" requests charset_normalizer "openai>=1.0" watchdog +else + "$PYTHON_BIN" -m pip install -t "$STAGE_DIR/python_libs" requests urllib3 charset_normalizer "openai>=1.0" watchdog +fi + +# Bundle MCP bridge npm package into the staged extension +BRIDGE_SRC="$SCRIPT_DIR/../../ctx-mcp-bridge" +BRIDGE_DIR="ctx-mcp-bridge" + +if [[ -d "$BRIDGE_SRC" && -f "$BRIDGE_SRC/package.json" ]]; then + echo "Bundling MCP bridge npm package into staged extension..." + mkdir -p "$STAGE_DIR/$BRIDGE_DIR" + if [[ -d "$BRIDGE_SRC/bin" ]]; then + cp -a "$BRIDGE_SRC/bin" "$STAGE_DIR/$BRIDGE_DIR/" + else + echo "Warning: Bridge bin directory not found at $BRIDGE_SRC/bin (skipping)" + fi + if [[ -d "$BRIDGE_SRC/src" ]]; then + cp -a "$BRIDGE_SRC/src" "$STAGE_DIR/$BRIDGE_DIR/" + else + echo "Warning: Bridge src directory not found at $BRIDGE_SRC/src (skipping)" + fi + cp "$BRIDGE_SRC/package.json" "$STAGE_DIR/$BRIDGE_DIR/" + + echo "Installing MCP bridge production dependencies into staged extension..." + if [[ -f "$BRIDGE_SRC/package-lock.json" ]]; then + cp "$BRIDGE_SRC/package-lock.json" "$STAGE_DIR/$BRIDGE_DIR/" + (cd "$STAGE_DIR/$BRIDGE_DIR" && npm ci --omit=dev) else - "$PYTHON_BIN" -m pip install -t "$STAGE_DIR/python_libs" requests urllib3 charset_normalizer "openai>=1.0" watchdog + (cd "$STAGE_DIR/$BRIDGE_DIR" && npm install --omit=dev) fi + echo "MCP bridge bundled successfully." +else + echo "Warning: MCP bridge source not found at $BRIDGE_SRC" fi pushd "$STAGE_DIR" >/dev/null @@ -82,4 +108,4 @@ npx @vscode/vsce package --no-dependencies --out "$OUT_DIR" popd >/dev/null echo "Build complete! Check the /out directory for .vsix and .py files." -ls -la "$OUT_DIR" \ No newline at end of file +ls -la "$OUT_DIR" diff --git a/vscode-extension/build/publish-vscode-extension.sh b/vscode-extension/build/publish-vscode-extension.sh index f019cb92..0d986a4a 100644 --- a/vscode-extension/build/publish-vscode-extension.sh +++ b/vscode-extension/build/publish-vscode-extension.sh @@ -4,7 +4,6 @@ set -euo pipefail SCRIPT_DIR="$(cd -- "$(dirname "${BASH_SOURCE[0]}")" && pwd)" BUILD_SCRIPT="$SCRIPT_DIR/build.sh" OUT_DIR="$SCRIPT_DIR/../out" -BUNDLE_DEPS="${1:-}" if [[ ! -f "$BUILD_SCRIPT" ]]; then echo "Build script not found: $BUILD_SCRIPT" >&2 @@ -18,7 +17,7 @@ fi export VSCE_STORE="${VSCE_STORE:-file}" -"$BUILD_SCRIPT" "$BUNDLE_DEPS" +"$BUILD_SCRIPT" VSIX_PATH="" if compgen -G "$OUT_DIR/*.vsix" >/dev/null; then diff --git a/vscode-extension/context-engine-uploader/README.md b/vscode-extension/context-engine-uploader/README.md index c84a79b6..e02dc982 100644 --- a/vscode-extension/context-engine-uploader/README.md +++ b/vscode-extension/context-engine-uploader/README.md @@ -20,7 +20,7 @@ Configuration - `Run On Startup` auto-triggers force sync + watch after VS Code finishes loading. - `Python Path`, `Endpoint`, `Extra Force Args`, `Extra Watch Args`, and `Interval Seconds` can be tuned via standard VS Code settings. - `Target Path` is auto-filled from the workspace but can be overridden if you need to upload a different folder. -- **Python dependencies:** the extension runs the standalone upload client via your configured `pythonPath`. Ensure the interpreter has `requests`, `urllib3`, `charset_normalizer`, and `watchdog` installed. Run `python3 -m pip install requests urllib3 charset_normalizer watchdog` (or replace `python3` with your configured path) before starting the uploader. +- **Python dependencies:** the extension ships bundled `python_libs` and adds them to `PYTHONPATH` for the upload client. You only need a runnable Python 3 interpreter via `contextEngineUploader.pythonPath`. - **Path mapping:** `Host Root` + `Container Root` control how local paths are rewritten before reaching the remote service. By default the host root mirrors your `Target Path` and the container root is `/work`, which keeps Windows paths working without extra config. - **Prompt+ decoder:** set `Context Engine Uploader: Decoder Url` (default `http://localhost:8081`, auto-appends `/completion`) to point at your local llama.cpp decoder. For Ollama, set it to `http://localhost:11434/api/chat`. Turn on `Use Gpu Decoder` to set `USE_GPU_DECODER=1` so ctx.py prefers the GPU llama.cpp sidecar. Prompt+ automatically runs the bundled `scripts/ctx.py` when an embedded copy is available, falling back to the workspace version if not. - **Claude/Windsurf MCP config:** diff --git a/vscode-extension/context-engine-uploader/config_resolver.js b/vscode-extension/context-engine-uploader/config_resolver.js index fd4ebe07..5dd457ad 100644 --- a/vscode-extension/context-engine-uploader/config_resolver.js +++ b/vscode-extension/context-engine-uploader/config_resolver.js @@ -193,10 +193,8 @@ function createConfigResolver(deps) { const configuredPython = (config.get('pythonPath') || '').trim(); let pythonPath = configuredPython || 'python3'; - let pythonPathSource = configuredPython ? 'configured' : 'default'; if (pythonOverridePath && fs.existsSync(pythonOverridePath)) { pythonPath = pythonOverridePath; - pythonPathSource = 'override'; } const endpoint = (config.get('endpoint') || '').trim(); const targetPath = getTargetPath(config); @@ -263,7 +261,6 @@ function createConfigResolver(deps) { return { pythonPath, - pythonPathSource, workingDirectory, scriptPath, targetPath, diff --git a/vscode-extension/context-engine-uploader/ctx_config.js b/vscode-extension/context-engine-uploader/ctx_config.js index 4d386f42..d70b437b 100644 --- a/vscode-extension/context-engine-uploader/ctx_config.js +++ b/vscode-extension/context-engine-uploader/ctx_config.js @@ -9,7 +9,7 @@ function createCtxConfigManager(deps) { const extensionRoot = deps.extensionRoot; const getEffectiveConfig = deps.getEffectiveConfig; const resolveOptions = deps.resolveOptions; - const ensurePythonDependencies = deps.ensurePythonDependencies; + const ensurePythonReady = deps.ensurePythonReady; const buildChildEnv = deps.buildChildEnv; const resolveBridgeHttpUrl = deps.resolveBridgeHttpUrl; @@ -29,12 +29,8 @@ function createCtxConfigManager(deps) { if (!options) { return; } - const depsOk = await ensurePythonDependencies( - options.pythonPath, - options.workingDirectory, - options.pythonPathSource - ); - if (!depsOk) { + const pythonReady = await ensurePythonReady(options.pythonPath); + if (!pythonReady) { return; } options = resolveOptions() || options; diff --git a/vscode-extension/context-engine-uploader/extension.js b/vscode-extension/context-engine-uploader/extension.js index 9a387c66..9b056082 100644 --- a/vscode-extension/context-engine-uploader/extension.js +++ b/vscode-extension/context-engine-uploader/extension.js @@ -12,7 +12,6 @@ const { createLogsTerminalManager } = require('./logs_terminal'); const { createPromptPlusManager } = require('./prompt_plus'); const { registerPromptPlusCommands } = require('./prompt_plus_commands'); const { createOnboardingManager } = require('./onboarding'); -const { createPythonEnvManager } = require('./python_env'); const { createProcessManager } = require('./process_manager'); const { registerExtensionCommands } = require('./commands'); const { createConfigResolver } = require('./config_resolver'); @@ -30,11 +29,11 @@ let ctxConfigManager; let promptPlusManager; let onboardingManager; -let pythonEnvManager; let processManager; let configResolver; let sidebarApi; let pendingProfileRestartTimer; +let hasShownPythonError = false; const DEFAULT_CONTAINER_ROOT = '/work'; // const CLAUDE_HOOK_COMMAND = '/home/coder/project/Context-Engine/ctx-hook-simple.sh'; @@ -185,25 +184,6 @@ function activate(context) { log(`Config resolver init failed: ${error instanceof Error ? error.message : String(error)}`); } - try { - pythonEnvManager = createPythonEnvManager({ - vscode, - spawn: spawn, - path, - fs, - log, - getEffectiveConfig, - getWorkspaceFolderPath: () => configResolver ? configResolver.getWorkspaceFolderPath() : undefined, - getExtensionRoot: () => extensionRoot, - getGlobalStoragePath: () => globalStoragePath, - getPythonOverridePath: () => pythonOverridePath, - setPythonOverridePath: (p) => { pythonOverridePath = p; }, - }); - } catch (error) { - pythonEnvManager = undefined; - log(`Python env manager init failed: ${error instanceof Error ? error.message : String(error)}`); - } - try { processManager = createProcessManager({ vscode, @@ -230,11 +210,13 @@ function activate(context) { path, fs, log, + extensionRoot, getEffectiveConfig, resolveBridgeWorkspacePath: () => configResolver ? configResolver.resolveBridgeWorkspacePath() : undefined, attachOutput: (child, label) => processManager ? processManager.attachOutput(child, label) : undefined, terminateProcess: (proc, label, afterStop) => processManager ? processManager.terminateProcess(proc, label, afterStop) : Promise.resolve(), scheduleMcpConfigRefreshAfterBridge: (delay) => mcpConfigManager ? mcpConfigManager.scheduleMcpConfigRefreshAfterBridge(delay) : undefined, + cancelPendingBridgeConfigRefresh: () => mcpConfigManager ? mcpConfigManager.cancelPendingBridgeConfigRefresh() : undefined, }); } catch (error) { bridgeManager = undefined; @@ -249,10 +231,7 @@ function activate(context) { extensionRoot, getEffectiveConfig, resolveOptions: () => configResolver ? configResolver.resolveOptions() : undefined, - ensurePythonDependencies: (pythonPath, workingDirectory, pythonPathSource) => - pythonEnvManager - ? pythonEnvManager.ensurePythonDependencies(pythonPath, workingDirectory, pythonPathSource) - : Promise.resolve(false), + ensurePythonReady, buildChildEnv: (options) => processManager?.buildChildEnv?.(options) ?? {}, resolveBridgeHttpUrl: () => bridgeManager ? bridgeManager.resolveBridgeHttpUrl() : undefined, }); @@ -274,6 +253,7 @@ function activate(context) { resolveBridgeCliInvocation: () => bridgeManager ? bridgeManager.resolveBridgeCliInvocation() : undefined, resolveBridgeHttpUrl: () => bridgeManager ? bridgeManager.resolveBridgeHttpUrl() : undefined, requiresHttpBridge: (s, t) => bridgeManager ? bridgeManager.requiresHttpBridge(s, t) : (s === 'bridge' && t === 'http'), + requiresLocalBridgeProcess: (s, t) => bridgeManager ? bridgeManager.requiresLocalBridgeProcess(s, t) : (s === 'bridge' && (t === 'http' || t === 'sse-remote')), ensureHttpBridgeReadyForConfigs: () => bridgeManager ? bridgeManager.ensureReadyForConfigs() : Promise.resolve(false), getBridgeIsRunning: () => (bridgeManager && typeof bridgeManager.isRunning === 'function' ? bridgeManager.isRunning() : false), writeCtxConfig: () => ctxConfigManager ? ctxConfigManager.writeCtxConfig() : Promise.resolve(), @@ -302,13 +282,6 @@ function activate(context) { } catch (_) { // ignore } - try { - const venvPy = pythonEnvManager ? pythonEnvManager.resolvePrivateVenvPython() : undefined; - if (venvPy) { - pythonOverridePath = venvPy; - log(`Detected existing private venv interpreter: ${venvPy}`); - } - } catch (_) { } statusBarItem = vscode.window.createStatusBarItem(vscode.StatusBarAlignment.Left, 100); statusBarItem.command = 'contextEngineUploader.indexCodebase'; context.subscriptions.push(statusBarItem); @@ -425,6 +398,7 @@ function activate(context) { event.affectsConfiguration('contextEngineUploader.mcpBridgeBinPath') || event.affectsConfiguration('contextEngineUploader.mcpBridgePort') || event.affectsConfiguration('contextEngineUploader.mcpBridgeLocalOnly') || + event.affectsConfiguration('contextEngineUploader.mcpBridgeMode') || event.affectsConfiguration('contextEngineUploader.windsurfMcpPath') || event.affectsConfiguration('contextEngineUploader.augmentMcpPath') || event.affectsConfiguration('contextEngineUploader.antigravityMcpPath') || @@ -439,6 +413,7 @@ function activate(context) { event.affectsConfiguration('contextEngineUploader.mcpBridgePort') || event.affectsConfiguration('contextEngineUploader.mcpBridgeBinPath') || event.affectsConfiguration('contextEngineUploader.mcpBridgeLocalOnly') || + event.affectsConfiguration('contextEngineUploader.mcpBridgeMode') || event.affectsConfiguration('contextEngineUploader.mcpIndexerUrl') || event.affectsConfiguration('contextEngineUploader.mcpMemoryUrl') || event.affectsConfiguration('contextEngineUploader.mcpServerMode') || @@ -484,10 +459,10 @@ function activate(context) { const serverModeRaw = config.get('mcpServerMode') || 'bridge'; const transportMode = (typeof transportModeRaw === 'string' ? transportModeRaw.trim() : 'sse-remote') || 'sse-remote'; const serverMode = (typeof serverModeRaw === 'string' ? serverModeRaw.trim() : 'bridge') || 'bridge'; - if (bridgeManager && bridgeManager.requiresHttpBridge(serverMode, transportMode)) { + if (bridgeManager && bridgeManager.requiresLocalBridgeProcess(serverMode, transportMode)) { startHttpBridgeProcess().catch(error => log(`Auto-start HTTP MCP bridge failed: ${error instanceof Error ? error.message : String(error)}`)); } else { - log('Context Engine Uploader: autoStartMcpBridge is enabled, but current MCP wiring does not use the HTTP bridge; skipping auto-start.'); + log('Context Engine Uploader: autoStartMcpBridge is enabled, but current MCP wiring does not use the local bridge process; skipping auto-start.'); } } } @@ -516,14 +491,12 @@ async function runSequence(mode = 'auto') { log(`Auth preflight check failed: ${error instanceof Error ? error.message : String(error)}`); } - const depsSatisfied = pythonEnvManager - ? await pythonEnvManager.ensurePythonDependencies(options.pythonPath, options.workingDirectory, options.pythonPathSource) - : false; + const depsSatisfied = await ensurePythonReady(options.pythonPath); if (!depsSatisfied) { setStatusBarState('idle'); return; } - // Re-resolve options in case ensurePythonDependencies switched to a better interpreter + // Re-resolve options in case Python preflight selected a better interpreter. const reoptions = configResolver ? configResolver.resolveOptions() : undefined; if (reoptions) { Object.assign(options, reoptions); @@ -540,8 +513,9 @@ async function runSequence(mode = 'auto') { if (code === 0) { setStatusBarState('indexed'); if (processManager) { processManager.ensureIndexedWatcher(options.targetPath); } - // Only start watching after a regular force sync, not after git history upload - if (mode === 'force' && options.startWatchAfterForce && processManager) { + // Start watch after successful force sync in normal flows (`force` and `auto`), + // but keep git-history upload as one-shot. + if (mode !== 'uploadGitHistory' && options.startWatchAfterForce && processManager) { processManager.startWatch(options); } } else { @@ -554,6 +528,42 @@ async function runSequence(mode = 'auto') { } } +function probePython(command, args = []) { + try { + const result = spawnSync(command, [...args, '-c', 'import sys; print(f"{sys.version_info[0]}|{sys.executable}")'], { encoding: 'utf8', timeout: 5000 }); + if (result.status !== 0) return undefined; + const [major, executable] = String(result.stdout || '').trim().split('|'); + return Number.parseInt(major, 10) >= 3 && executable ? executable.trim() : undefined; + } catch (_) { + return undefined; + } +} + +function ensurePythonReady(pythonPath) { + if (pythonOverridePath) return true; + const requested = pythonPath || 'python3'; + const candidates = process.platform === 'win32' + ? [[requested, []], ['py', ['-3']], ['python', []], ['python3', []]] + : [[requested, []], ['python3', []], ['python', []], ['/opt/homebrew/bin/python3', []]]; + const seen = new Set(); + for (const [command, args] of candidates) { + const key = `${command} ${args.join(' ')}`.trim(); + if (!command || seen.has(key)) continue; + seen.add(key); + const executable = probePython(command, args); + if (!executable) continue; + pythonOverridePath = executable; + log(`Using Python interpreter: ${executable}`); + return true; + } + log(`Python preflight failed for ${requested}.`); + if (!hasShownPythonError) { + hasShownPythonError = true; + vscode.window.showErrorMessage(`Context Engine Uploader: Python 3 was not found. Install Python 3 or update contextEngineUploader.pythonPath (current: ${requested}).`); + } + return false; +} + async function startHttpBridgeProcess() { diff --git a/vscode-extension/context-engine-uploader/mcp_bridge.js b/vscode-extension/context-engine-uploader/mcp_bridge.js index d9825177..bbb47fbc 100644 --- a/vscode-extension/context-engine-uploader/mcp_bridge.js +++ b/vscode-extension/context-engine-uploader/mcp_bridge.js @@ -4,18 +4,32 @@ function createBridgeManager(deps) { const path = deps.path; const fs = deps.fs; const log = deps.log; + const extensionRoot = deps.extensionRoot; const getEffectiveConfig = deps.getEffectiveConfig; const resolveBridgeWorkspacePath = deps.resolveBridgeWorkspacePath; const attachOutput = deps.attachOutput; const terminateProcess = deps.terminateProcess; const scheduleMcpConfigRefreshAfterBridge = deps.scheduleMcpConfigRefreshAfterBridge; + const cancelPendingBridgeConfigRefresh = deps.cancelPendingBridgeConfigRefresh; let httpBridgeProcess; let httpBridgePort; let httpBridgeWorkspace; let stopInFlight; + function clearBridgeState(child) { + if (httpBridgeProcess !== child) { + return; + } + httpBridgeProcess = undefined; + httpBridgePort = undefined; + httpBridgeWorkspace = undefined; + if (typeof cancelPendingBridgeConfigRefresh === 'function') { + cancelPendingBridgeConfigRefresh(); + } + } + function normalizeBridgeUrl(url) { if (!url || typeof url !== 'string') { return ''; @@ -42,7 +56,32 @@ function createBridgeManager(deps) { } } + function getBridgeMode() { + try { + const settings = getEffectiveConfig(); + return (settings.get('mcpBridgeMode') || 'bundled').trim(); + } catch (_) { + return 'bundled'; + } + } + + function findBundledBridgeBin() { + if (!extensionRoot) return undefined; + const bundledPath = path.join(extensionRoot, 'ctx-mcp-bridge', 'bin', 'ctxce.js'); + if (fs.existsSync(bundledPath)) { + return path.resolve(bundledPath); + } + return undefined; + } + function findLocalBridgeBin() { + // First check for bundled bridge if mode is 'bundled' + const mode = getBridgeMode(); + if (mode === 'bundled') { + return findBundledBridgeBin(); + } + + // External mode logic (existing behavior) let localOnly = true; let configured = ''; try { @@ -69,12 +108,20 @@ function createBridgeManager(deps) { function resolveBridgeCliInvocation() { const binPath = findLocalBridgeBin(); if (binPath) { + // Use absolute Node runtime to avoid PATH dependency in extension hosts + const bundledBin = findBundledBridgeBin(); + const resolvedKind = bundledBin && path.resolve(binPath) === path.resolve(bundledBin) + ? 'bundled' + : 'local'; return { - command: 'node', + command: process.execPath, args: [binPath], - kind: 'local' + kind: resolvedKind }; } + if (getBridgeMode() === 'bundled') { + return undefined; + } const isWindows = process.platform === 'win32'; if (isWindows) { return { @@ -107,6 +154,10 @@ function createBridgeManager(deps) { return serverMode === 'bridge' && transportMode === 'http'; } + function requiresLocalBridgeProcess(serverMode, transportMode) { + return serverMode === 'bridge' && (transportMode === 'http' || transportMode === 'sse-remote'); + } + function resolveBridgeHttpUrl() { try { const settings = getEffectiveConfig(); @@ -199,21 +250,12 @@ function createBridgeManager(deps) { attachOutput(child, 'mcp-http'); child.on('exit', (code, signal) => { log(`HTTP MCP bridge exited with code ${code} signal ${signal || ''}`.trim()); - if (httpBridgeProcess === child) { - httpBridgeProcess = undefined; - httpBridgePort = undefined; - httpBridgeWorkspace = undefined; - } + clearBridgeState(child); }); child.on('error', error => { log(`HTTP MCP bridge process error: ${error instanceof Error ? error.message : String(error)}`); - if (httpBridgeProcess === child) { - httpBridgeProcess = undefined; - httpBridgePort = undefined; - httpBridgeWorkspace = undefined; - } + clearBridgeState(child); }); - vscode.window.showInformationMessage(`Context Engine HTTP MCP bridge listening on http://127.0.0.1:${options.port}/mcp`); if (typeof scheduleMcpConfigRefreshAfterBridge === 'function') { scheduleMcpConfigRefreshAfterBridge(); } @@ -269,10 +311,10 @@ function createBridgeManager(deps) { const serverModeRaw = config.get('mcpServerMode') || 'bridge'; const transportMode = (typeof transportModeRaw === 'string' ? transportModeRaw.trim() : 'sse-remote') || 'sse-remote'; const serverMode = (typeof serverModeRaw === 'string' ? serverModeRaw.trim() : 'bridge') || 'bridge'; - if (requiresHttpBridge(serverMode, transportMode)) { + if (requiresLocalBridgeProcess(serverMode, transportMode)) { await start(); } else { - log('Context Engine Uploader: HTTP bridge settings changed, but current MCP wiring does not use the HTTP bridge; not restarting HTTP bridge.'); + log('Context Engine Uploader: bridge settings changed, but current MCP wiring does not use the local bridge process; not restarting bridge.'); } } } @@ -290,6 +332,7 @@ function createBridgeManager(deps) { getState, isRunning, requiresHttpBridge, + requiresLocalBridgeProcess, resolveBridgeHttpUrl, ensureReadyForConfigs, start, diff --git a/vscode-extension/context-engine-uploader/mcp_config.js b/vscode-extension/context-engine-uploader/mcp_config.js index 3e40c2fd..a8c4f11c 100644 --- a/vscode-extension/context-engine-uploader/mcp_config.js +++ b/vscode-extension/context-engine-uploader/mcp_config.js @@ -52,6 +52,13 @@ function createMcpConfigManager(deps) { } } + function cancelPendingBridgeConfigRefresh() { + if (pendingBridgeConfigTimer) { + clearTimeout(pendingBridgeConfigTimer); + pendingBridgeConfigTimer = undefined; + } + } + async function writeAntigravityMcpServers(configPath, indexerUrl, memoryUrl, transportMode, serverMode = 'bridge', workspaceHint) { // TODO: Factor the shared "ensure dir + load JSON + applyMcpServersUpdate + writeJsonConfig" pattern // into a helper so Claude/Windsurf/Augment/Antigravity all call the same utility. @@ -338,10 +345,7 @@ function createMcpConfigManager(deps) { function scheduleMcpConfigRefreshAfterBridge(delayMs = 1500) { try { - if (pendingBridgeConfigTimer) { - clearTimeout(pendingBridgeConfigTimer); - pendingBridgeConfigTimer = undefined; - } + cancelPendingBridgeConfigRefresh(); // For bridge-http mode started by the extension, Windsurf needs the // "context-engine" MCP server entry removed and then re-added once the // HTTP bridge is ready. Best-effort removal happens immediately here; @@ -363,8 +367,12 @@ function createMcpConfigManager(deps) { } pendingBridgeConfigTimer = setTimeout(() => { pendingBridgeConfigTimer = undefined; - log('Context Engine Uploader: HTTP bridge ready; refreshing MCP configs.'); - writeMcpConfig().catch(error => { + if (typeof getBridgeIsRunning === 'function' && !getBridgeIsRunning()) { + log('Context Engine Uploader: HTTP bridge is not running; skipping delayed MCP config refresh.'); + return; + } + log('Context Engine Uploader: HTTP bridge still running; refreshing MCP configs.'); + writeMcpConfig({ skipHttpBridgeStart: true }).catch(error => { log(`Context Engine Uploader: MCP config refresh after bridge start failed: ${error instanceof Error ? error.message : String(error)}`); }); }, delayMs); @@ -736,6 +744,10 @@ function createMcpConfigManager(deps) { const needsHttpBridge = requiresHttpBridge(serverMode, transportMode); const bridgeWasRunning = !!(typeof getBridgeIsRunning === 'function' && getBridgeIsRunning()); if (needsHttpBridge) { + if (options.skipHttpBridgeStart && !bridgeWasRunning) { + log('Context Engine Uploader: HTTP bridge is not running; MCP config refresh will not restart it.'); + return; + } const ready = await ensureHttpBridgeReadyForConfigs(); if (!ready) { vscode.window.showErrorMessage('Context Engine Uploader: HTTP MCP bridge failed to start; MCP config not updated.'); @@ -828,6 +840,7 @@ function createMcpConfigManager(deps) { } return { + cancelPendingBridgeConfigRefresh, scheduleMcpConfigRefreshAfterBridge, writeMcpConfig, dispose, diff --git a/vscode-extension/context-engine-uploader/package.json b/vscode-extension/context-engine-uploader/package.json index d5e3584f..c3eadf0a 100644 --- a/vscode-extension/context-engine-uploader/package.json +++ b/vscode-extension/context-engine-uploader/package.json @@ -1,8 +1,8 @@ { "name": "context-engine-uploader", "displayName": "Context Engine Uploader", - "description": "Runs the Context-Engine remote upload client with a force sync on startup followed by watch mode. Requires Python with pip install requests urllib3 charset_normalizer.", - "version": "0.1.39", + "description": "Runs the Context-Engine remote upload client with bundled Python dependencies, force sync, and watch mode.", + "version": "0.1.40", "publisher": "context-engine", "engines": { "vscode": "^1.85.0" @@ -282,7 +282,7 @@ "contextEngineUploader.autoStartMcpBridge": { "type": "boolean", "default": true, - "description": "When enabled and mcpServerMode='bridge' with mcpTransportMode='http', automatically start the local ctx-mcp-bridge HTTP server for the active workspace so IDE clients can connect over HTTP without manual commands. Has no effect in stdio/direct modes." + "description": "When enabled and mcpServerMode='bridge', automatically start the bundled local ctx bridge process for the active workspace. In http mode it serves the local HTTP MCP bridge directly; in sse-remote mode it starts the same bundled bridge adapter used by bridge-stdio wiring. Has no effect in direct modes." }, "contextEngineUploader.mcpBridgePort": { "type": "number", @@ -297,7 +297,17 @@ "contextEngineUploader.mcpBridgeLocalOnly": { "type": "boolean", "default": false, - "description": "Development toggle. When true (default) the extension prefers local bridge binaries resolved from mcpBridgeBinPath or CTXCE_BRIDGE_BIN before falling back to the published npm build via npx." + "description": "Development toggle. When true and mcpBridgeMode='external', prefers local bridge binaries resolved from mcpBridgeBinPath or CTXCE_BRIDGE_BIN before falling back to the published npm build via npx. Ignored when mcpBridgeMode='bundled'." + }, + "contextEngineUploader.mcpBridgeMode": { + "type": "string", + "enum": ["bundled", "external"], + "default": "bundled", + "description": "Bridge invocation mode. 'bundled' uses the bundled bridge inside the extension (offline, no npx required). 'external' uses external binary path or npx (current behavior).", + "enumDescriptions": [ + "Use the bundled MCP bridge inside the extension (works offline).", + "Use external binary path or npx to run the bridge (requires internet for first npx install)." + ] }, "contextEngineUploader.mcpServerMode": { "type": "string", diff --git a/vscode-extension/context-engine-uploader/process_manager.js b/vscode-extension/context-engine-uploader/process_manager.js index a5e36fd4..22173d58 100644 --- a/vscode-extension/context-engine-uploader/process_manager.js +++ b/vscode-extension/context-engine-uploader/process_manager.js @@ -64,8 +64,11 @@ function createProcessManager(deps) { env.CONTAINER_ROOT = options.containerRoot; } try { - const libsPath = path.join(options.workingDirectory, 'python_libs'); - if (fs.existsSync(libsPath)) { + const libsPath = [ + path.join(options.workingDirectory, 'python_libs'), + path.join(getExtensionRoot(), 'python_libs') + ].find(p => p && fs.existsSync(p)); + if (libsPath) { const existing = env.PYTHONPATH || ''; env.PYTHONPATH = existing ? `${libsPath}${path.delimiter}${existing}` : libsPath; if (!_hasLoggedPythonPath) { diff --git a/vscode-extension/context-engine-uploader/python_env.js b/vscode-extension/context-engine-uploader/python_env.js deleted file mode 100644 index 190f9945..00000000 --- a/vscode-extension/context-engine-uploader/python_env.js +++ /dev/null @@ -1,386 +0,0 @@ -/** - * Python environment management for Context Engine extension. - * Handles dependency checking, venv creation, and Python interpreter detection. - */ -function createPythonEnvManager(deps) { - const vscode = deps.vscode; - const spawn = deps.spawn; - const path = deps.path; - const fs = deps.fs; - const log = deps.log; - - - - // Helper to spawn processes asynchronously with Promise wrapper - function execAsync(command, args, options = {}) { - return new Promise((resolve) => { - // Diagnostic check for spawn injection - if (typeof spawn !== 'function') { - resolve({ code: -1, stdout: '', stderr: `createPythonEnvManager: spawn is ${typeof spawn}` }); - return; - } - - const child = spawn(command, args, { - ...options, - env: options.env || process.env - }); - - let stdout = ''; - let stderr = ''; - - if (child.stdout) { - child.stdout.on('data', (data) => { - const str = data.toString(); - stdout += str; - if (options.onStdout) options.onStdout(str); - }); - } - - if (child.stderr) { - child.stderr.on('data', (data) => { - const str = data.toString(); - stderr += str; - if (options.onStderr) options.onStderr(str); - }); - } - - let finished = false; - - child.on('error', (err) => { - if (!finished) { - finished = true; - resolve({ code: -1, stdout, stderr: stderr || err.message }); - } - }); - - child.on('close', (code) => { - if (!finished) { - finished = true; - resolve({ code: code === null ? -1 : code, stdout, stderr }); - } - }); - - // Handle cancellation if token provided - if (options.token) { - options.token.onCancellationRequested(() => { - if (!finished) { - finished = true; - try { child.kill(); } catch (_) { } - resolve({ code: -1, stdout, stderr: 'Cancelled' }); - } - }); - } - - // Safety timeout - if (options.timeout) { - setTimeout(() => { - if (!finished) { - finished = true; - try { child.kill(); } catch (_) { } - resolve({ code: -1, stdout, stderr: 'Process timeout' }); - } - }, options.timeout); - } - }); - } - - function getExtensionRoot() { - try { - if (typeof deps.getExtensionRoot === 'function') { - const root = deps.getExtensionRoot(); - if (root) { - return root; - } - } - } catch (_) { - } - if (deps.extensionRoot) return deps.extensionRoot; - try { - return vscode.extensions.getExtension('context-engine.context-engine-uploader').extensionPath; - } catch (_) { - return __dirname; - } - } - - function getPythonOverridePath() { - return typeof deps.getPythonOverridePath === 'function' ? deps.getPythonOverridePath() : undefined; - } - - function setPythonOverridePath(p) { - if (typeof deps.setPythonOverridePath === 'function') { - deps.setPythonOverridePath(p); - } - } - - const REQUIRED_PYTHON_MODULES = ['requests', 'urllib3', 'charset_normalizer', 'watchdog']; - const depCheckCache = new Map(); - - function cacheKey(pythonPath, workingDirectory) { - return `${pythonPath || ''}::${workingDirectory || ''}`; - } - - function venvRootDir() { - // Prefer workspace storage; fallback to extension storage - try { - const ws = deps.getWorkspaceFolderPath(); - const globalStorage = deps.getGlobalStoragePath() || path.join(getExtensionRoot(), '.storage'); - const base = ws && fs.existsSync(ws) ? path.join(ws, '.vscode', '.context-engine-uploader') - : globalStorage; - if (!fs.existsSync(base)) fs.mkdirSync(base, { recursive: true }); - return base; - } catch (e) { - return getExtensionRoot(); - } - } - - function privateVenvPath() { - return path.join(venvRootDir(), 'py-venv'); - } - - function resolvePrivateVenvPython() { - const venvPath = privateVenvPath(); - const bin = process.platform === 'win32' ? path.join(venvPath, 'Scripts', 'python.exe') : path.join(venvPath, 'bin', 'python'); - return fs.existsSync(bin) ? bin : undefined; - } - - async function detectSystemPython() { - // Try configured pythonPath, then common names - const candidates = []; - try { - const cfg = (typeof getEffectiveConfig === 'function') - ? getEffectiveConfig() - : vscode.workspace.getConfiguration('contextEngineUploader'); - const configured = (cfg && typeof cfg.get === 'function') ? (cfg.get('pythonPath') || '').trim() : ''; - if (configured) candidates.push(configured); - } catch { } - if (process.platform === 'win32') { - candidates.push('py', 'python3', 'python'); - } else { - candidates.push('python3', 'python'); - // Add common Homebrew path on Apple Silicon - candidates.push('/opt/homebrew/bin/python3'); - } - - for (const cmd of candidates) { - try { - // Version check: major >= 3 and print executable - const res = await execAsync(cmd, ['-c', 'import sys; print(f"{sys.version_info[0]}|{sys.executable}")'], { timeout: 3000 }); - if (res.code === 0) { - const parts = res.stdout.trim().split('|'); - if (parts.length === 2) { - const major = parseInt(parts[0], 10); - const executable = parts[1].trim(); - if (major >= 3 && executable) return executable; - } - } - } catch (e) { - // Skip candidate - } - } - return undefined; - } - - async function checkPythonDeps(pythonPath, workingDirectory, options = {}) { - const showInterpreterError = options.showInterpreterError !== undefined ? options.showInterpreterError : true; - const missing = []; - let pythonError; - const env = { ...process.env }; - try { - const candidates = []; - if (workingDirectory) { - candidates.push(path.join(workingDirectory, 'python_libs')); - } - candidates.push(path.join(getExtensionRoot(), 'python_libs')); - for (const libsPath of candidates) { - if (libsPath && fs.existsSync(libsPath)) { - const existing = env.PYTHONPATH || ''; - env.PYTHONPATH = existing ? `${libsPath}${path.delimiter}${existing}` : libsPath; - break; - } - } - } catch (error) { - log(`Failed to configure PYTHONPATH for dependency check: ${error instanceof Error ? error.message : String(error)}`); - } - - const smoke = await execAsync(pythonPath, ['-c', 'import sys; print(sys.executable)'], { env, timeout: 5000 }); - if (smoke.code !== 0) { - pythonError = String((smoke.stderr || smoke.stdout || '')).trim(); - } - - if (!pythonError) { - for (const moduleName of REQUIRED_PYTHON_MODULES) { - const check = await execAsync(pythonPath, ['-c', `import ${moduleName}`], { env, timeout: 5000 }); - if (check.code !== 0) { - missing.push(moduleName); - } - } - } - - if (pythonError) { - if (showInterpreterError) { - vscode.window.showErrorMessage(`Context Engine Uploader: failed to run ${pythonPath}. Update contextEngineUploader.pythonPath.`); - } - log(`Dependency check failed: ${pythonError}`); - return false; - } - if (missing.length) { - log(`Missing Python modules for ${pythonPath}: ${missing.join(', ')}`); - return false; - } - return true; - } - - async function ensurePrivateVenv() { - try { - const python = resolvePrivateVenvPython(); - if (python) { - log('Private venv already exists.'); - return true; - } - const venvPath = privateVenvPath(); - const basePy = await detectSystemPython(); - if (!basePy) { - vscode.window.showErrorMessage('Context Engine Uploader: no Python 3 interpreter found to bootstrap venv.'); - return false; - } - - // Verify venv module presence - try { - const venvCheck = await execAsync(basePy, ['-c', 'import venv'], { timeout: 5000 }); - if (venvCheck.code !== 0) { - log(`Python "venv" module missing in ${basePy}: ${venvCheck.stderr}`); - vscode.window.showErrorMessage(`Context Engine Uploader: Python "venv" module is missing in ${basePy}.`); - return false; - } - } catch (e) { - const errorMsg = e instanceof Error ? e.message : String(e); - log(`Failed to check for venv module: ${errorMsg}`); - return false; - } - - log(`Creating private venv at ${venvPath} using ${basePy}`); - const res = await execAsync(basePy, ['-m', 'venv', venvPath], { timeout: 30000 }); - if (res.code !== 0) { - log(`venv creation failed: ${res.stderr || res.stdout}`); - vscode.window.showErrorMessage('Context Engine Uploader: failed to create private venv.'); - return false; - } - return true; - } catch (e) { - log(`ensurePrivateVenv error: ${e && e.message ? e.message : String(e)}`); - return false; - } - } - - async function installDepsInto(pythonBin) { - return vscode.window.withProgress({ - location: vscode.ProgressLocation.Notification, - title: "Context Engine Uploader: Installing Python dependencies...", - cancellable: true - }, async (progress, token) => { - try { - log(`Installing Python deps into private venv via ${pythonBin}`); - const args = ['-m', 'pip', 'install', ...REQUIRED_PYTHON_MODULES]; - - const res = await execAsync(pythonBin, args, { - timeout: 60000, - token, - onStdout: (data) => { - progress.report({ message: data.split('\n').pop() }); - }, - onStderr: (data) => { - log(`pip install stderr: ${data}`); - } - }); - - if (res.code !== 0) { - log(`pip install failed: ${res.stderr || res.stdout}`); - vscode.window.showErrorMessage('Context Engine Uploader: pip install failed. See Output for details.'); - return false; - } - return true; - } catch (e) { - const msg = e && e.message ? e.message : String(e); - log(`installDepsInto error: ${msg}`); - vscode.window.showErrorMessage(`Context Engine Uploader: ${msg}`); - return false; - } - }); - } - - async function ensurePythonDependencies(pythonPath, workingDirectory, pythonPathSource) { - // Probe current interpreter with bundled python_libs first - const allowPrompt = pythonPathSource === 'configured' || pythonPathSource === 'override'; - const primaryKey = cacheKey(pythonPath, workingDirectory); - if (depCheckCache.get(primaryKey)) { - return true; - } - let ok = await checkPythonDeps(pythonPath, workingDirectory, { showInterpreterError: allowPrompt }); - if (ok) { - depCheckCache.set(primaryKey, true); - return true; - } - - // If that fails, try to auto-detect a better system Python before falling back to a venv - const autoPython = await detectSystemPython(); - if (autoPython && autoPython !== pythonPath) { - log(`Falling back to auto-detected Python interpreter: ${autoPython}`); - const autoKey = cacheKey(autoPython, workingDirectory); - if (depCheckCache.get(autoKey)) { - setPythonOverridePath(autoPython); - return true; - } - ok = await checkPythonDeps(autoPython, workingDirectory, { showInterpreterError: allowPrompt }); - if (ok) { - setPythonOverridePath(autoPython); - depCheckCache.set(autoKey, true); - return true; - } - } - - // As a last resort, offer to create a private venv and install deps via pip - if (!allowPrompt) { - log('Skipping auto-install prompt; interpreter was auto-detected and missing modules.'); - return false; - } - const choice = await vscode.window.showErrorMessage( - 'Context Engine Uploader: missing Python modules. Create isolated environment and auto-install?', - 'Auto-install to private venv', - 'Cancel' - ); - if (choice !== 'Auto-install to private venv') { - return false; - } - const created = await ensurePrivateVenv(); - if (!created) return false; - const venvPython = resolvePrivateVenvPython(); - if (!venvPython) { - vscode.window.showErrorMessage('Context Engine Uploader: failed to locate private venv python.'); - return false; - } - const installed = await installDepsInto(venvPython); - if (!installed) return false; - setPythonOverridePath(venvPython); - log(`Using private venv interpreter: ${getPythonOverridePath()}`); - const venvKey = cacheKey(venvPython, workingDirectory); - if (depCheckCache.get(venvKey)) { - return true; - } - const finalOk = await checkPythonDeps(venvPython, workingDirectory, { showInterpreterError: true }); - if (finalOk) { - depCheckCache.set(venvKey, true); - } - return finalOk; - } - - return { - resolvePrivateVenvPython, - detectSystemPython, - checkPythonDeps, - ensurePythonDependencies, - }; -} - -module.exports = { - createPythonEnvManager, -};