diff --git a/.env.example b/.env.example index 7ddba7a..70ef6eb 100644 --- a/.env.example +++ b/.env.example @@ -33,3 +33,7 @@ PINECONE_READ_ONLY_MCP_LOG_FORMAT=text # Optional: Request timeout for Pinecone calls, in milliseconds (default: 15000) # PINECONE_REQUEST_TIMEOUT_MS=15000 + +# Optional: Multi-source mode (replaces PINECONE_API_KEY + PINECONE_INDEX_NAME when set) +# PINECONE_SOURCES=api_key_1:${PINECONE_API_KEY_1}:index_name_1;api_key_2:${PINECONE_API_KEY_2}:index_name_2 +# Or: PINECONE_CONFIG_FILE=./pinecone-sources.json diff --git a/CHANGELOG.md b/CHANGELOG.md index 5a31cb5..a363ea5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,10 @@ Tagged releases are published to npm from GitHub Actions when a **GitHub Release ## [Unreleased] +### Added + +- **Multi-source mode:** configure multiple Pinecone API keys / indexes in one MCP server via `PINECONE_SOURCES`, `--sources`, or a JSON config file (`PINECONE_CONFIG_FILE` / `--config-file`). New `list_sources` tool (when more than one source is configured). Optional `source` parameter on discovery and query tools; `list_namespaces` aggregates across sources and tags each namespace with `source`. See [CONFIGURATION.md](docs/CONFIGURATION.md#multi-source-mode), [TOOLS.md](docs/TOOLS.md#multi-source-mode), and deployment profiles in [CONFIGURATION.md](docs/CONFIGURATION.md#deployment-profiles). + ### Changed - **Library:** `resolveConfig()` returns `CoreServerConfig`; `resolveAllianceConfig()` returns `AllianceServerConfig`. `setupCoreServer` / `setupAllianceServer` accept only their respective branded config and context types (`CoreServerContext` / `AllianceServerContext`). `ServerConfig` remains an alias for `ServerConfigBase` on read paths (`ctx.getConfig()`). See [MIGRATION.md](docs/MIGRATION.md#unreleased-branded-serverconfig-types). diff --git a/README.md b/README.md index 9cea0c3..35c72ed 100644 --- a/README.md +++ b/README.md @@ -57,6 +57,7 @@ For successful `query`, `query_documents`, and `guided_query` payloads, **rerank - **Hybrid Search**: Combines dense and sparse embeddings for superior recall - **Semantic Reranking**: Uses BGE reranker model for improved precision - **Dynamic Namespace Discovery**: Automatically discovers available namespaces in your Pinecone index +- **Multi-Source (optional)**: Connect multiple Pinecone projects in one MCP server via `PINECONE_SOURCES` or a JSON config file; use `list_sources` and per-result `source` tags for routing - **Metadata Filtering**: Supports optional metadata filters for refined searches - **Fast presets**: Lazy initialization, connection pooling, and efficient result merging; use the `query` tool `preset=fast | detailed | full` to trade latency vs quality (no published benchmarks yet — treat descriptions as qualitative). - **Production-oriented defaults**: Input validation, error handling, and configurable logging; upgrading from **0.1.x** — see [MIGRATION.md](docs/MIGRATION.md). @@ -122,6 +123,8 @@ The codebase is split into two layers: You need a **Pinecone API key**. **Index** (`PINECONE_INDEX_NAME` or `--index-name`) is required for core/library use; the **published CLI** defaults to `rag-hybrid` when unset (Alliance deployment). Sparse index defaults to `{index}-sparse`. **Rerank:** set `PINECONE_RERANK_MODEL` to enable; the CLI defaults to `bge-reranker-v2-m3` when unset. See [docs/CONFIGURATION.md](docs/CONFIGURATION.md) (core vs Alliance table). +**Multiple projects:** set `PINECONE_SOURCES` or `PINECONE_CONFIG_FILE` instead of a single `PINECONE_API_KEY` — see [Multi-source mode](docs/CONFIGURATION.md#multi-source-mode). + Quick reference (published CLI / Alliance — core embedders require index, no index/rerank defaults): | Variable | Required | Default (Alliance CLI) | diff --git a/benchmarks/latency.ts b/benchmarks/latency.ts index 56a57c8..c012581 100644 --- a/benchmarks/latency.ts +++ b/benchmarks/latency.ts @@ -12,6 +12,7 @@ import { performance } from 'node:perf_hooks'; import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'; import { PineconeClient } from '../src/pinecone-client.js'; import { setLogLevel } from '../src/logger.js'; +import { exitOnDemoFailure } from '../examples/lib/exit-on-failure.js'; import { setPineconeClient } from '../src/server/client-context.js'; import { invalidateNamespacesCache, getNamespacesWithCache } from '../src/server/namespaces-cache.js'; import { registerGuidedQueryTool } from '../src/server/tools/guided-query-tool.js'; @@ -323,7 +324,4 @@ async function main(): Promise { console.log(`Wrote ${baselinePath}`); } -main().catch((err) => { - console.error(err); - process.exit(1); -}); +main().catch(exitOnDemoFailure('latency benchmark')); diff --git a/docs/CONFIGURATION.md b/docs/CONFIGURATION.md index d030944..23ac441 100644 --- a/docs/CONFIGURATION.md +++ b/docs/CONFIGURATION.md @@ -25,8 +25,10 @@ Configuration is built from **CLI flags** (when using the binary), **environment | `requestTimeoutMs` | `requestTimeoutMs` / `PINECONE_REQUEST_TIMEOUT_MS` | `15000` | | `disableSuggestFlow` | `disableSuggestFlow` / `PINECONE_DISABLE_SUGGEST_FLOW` | **Core `resolveConfig`:** `true` (gate off). **Alliance `resolveAllianceConfig` / CLI:** `false` (gate on). Bool parsing: true/1/yes/on | | `checkIndexes` | `checkIndexes` / `PINECONE_CHECK_INDEXES` | `false` | +| `sources` | `sources` / `PINECONE_SOURCES` or JSON config file | Omitted in single-key mode; see [Multi-source mode](#multi-source-mode) | +| `defaultSource` | JSON config `defaultSource` only | First source when using inline `PINECONE_SOURCES` | -**Throws** if `apiKey` or `indexName` is missing after trim. +**Throws** if `apiKey` or `indexName` is missing after trim (single-key mode). In multi-source mode, `PINECONE_API_KEY` is ignored when `PINECONE_SOURCES` or a config file is set; credentials come from each source entry. For the full Alliance tool surface (including `suggest_query_params`, `guided_query`, and built-in URL generators), import from `@will-cppa/pinecone-read-only-mcp/alliance` and use the three-step instance-first recipe at [Library embedding](#library-embedding) below. @@ -43,6 +45,113 @@ C++ Alliance deployers can copy [examples/alliance/.env.example](../examples/all --- +## Multi-source mode + +Use **one MCP server entry** with multiple Pinecone API keys / projects when `PINECONE_SOURCES`, `--sources`, or a JSON config file (`PINECONE_CONFIG_FILE` / `--config-file`) is set. + +### Inline format (`PINECONE_SOURCES` / `--sources`) + +Semicolon-separated entries: `name:apiKey:indexName` + +```bash +PINECONE_SOURCES=api_key_1:${PINECONE_API_KEY_1}:index_name_1;api_key_2:${PINECONE_API_KEY_2}:index_name_2 +``` + +API keys may contain colons; the parser treats the last `:` segment as `indexName` and everything between `name:` and `:indexName` as the key. + +### JSON config file + +Set `PINECONE_CONFIG_FILE` (or `--config-file`) to a path such as [examples/multi-source/pinecone-sources.json.example](../examples/multi-source/pinecone-sources.json.example): + +```json +{ + "defaultSource": "api_key_1", + "sources": { + "api_key_1": { "apiKey": "${PINECONE_API_KEY_1}", "indexName": "index_name_1" }, + "api_key_2": { "apiKey": "${PINECONE_API_KEY_2}", "indexName": "index_name_2" } + } +} +``` + +Values support `${ENV_VAR}` indirection (resolved at startup). Per-source `sparseIndexName` and `rerankModel` are optional; Alliance defaults apply when omitted. + +### MCP tools and routing + +| Tool | `source` parameter | +| ---- | ------------------ | +| `list_sources` | Registered only when more than one source is configured | +| `list_namespaces`, `namespace_router` | Omit to aggregate all sources; results include `source` when tagged | +| `query`, `count`, `query_documents`, `keyword_search`, `generate_urls`, `suggest_query_params`, `guided_query` | Omit when the namespace uniquely identifies one source; required when the same namespace exists on multiple sources | + +Discovery responses may include `source_errors` when one project fails but others succeed. Suggest-flow state uses compound keys `source:namespace` in multi-source mode. + +Single-key deployments (`PINECONE_API_KEY` + `PINECONE_INDEX_NAME` only) are unchanged — no `source` field on responses and no `list_sources` tool. + +### Deployment profiles + +Multi-source mode supports two operational profiles. **Never** ship a merged internal config through the same channel used for external partners. + +| Profile | Who | Config | Risk if mis-shared | +| ------- | --- | ------ | ------------------ | +| **External (public-only)** | External companies, public MCP distribution | `PINECONE_API_KEY` + `PINECONE_INDEX_NAME`, or `PINECONE_SOURCES` with **one** entry | Low — single public key only | +| **Internal (merged)** | Staff machines with access to private data | `PINECONE_SOURCES` or JSON config with **two+** entries | **High** — private API key and private namespace names exposed | + +**External MCP config (unchanged):** + +```json +{ + "mcpServers": { + "pinecone-search": { + "command": "npx", + "args": ["-y", "@will-cppa/pinecone-read-only-mcp"], + "env": { + "PINECONE_API_KEY": "your-public-key", + "PINECONE_INDEX_NAME": "rag-hybrid" + } + } + } +} +``` + +**Internal MCP config (merged sources):** + +```json +{ + "mcpServers": { + "pinecone-search": { + "command": "npx", + "args": ["-y", "@will-cppa/pinecone-read-only-mcp"], + "env": { + "PINECONE_SOURCES": "api_key_1:${PINECONE_API_KEY_1}:index_name_1;api_key_2:${PINECONE_API_KEY_2}:index_name_2", + "PINECONE_API_KEY_1": "pcsk_...", + "PINECONE_API_KEY_2": "pcsk_..." + } + } + } +} +``` + +Prefer `PINECONE_CONFIG_FILE` with `${ENV_VAR}` indirection over inline API keys in `PINECONE_SOURCES`. See [SECURITY.md](./SECURITY.md). + +### Architecture decision + +**Chosen:** Option A — multi-source server with optional `source` parameter on tools (`SourceRegistry` + `PINECONE_SOURCES` / JSON config). + +**Rejected:** + +- **Option B (Cursor routing rules only):** Does not fix the UX problem when users forget which MCP entry to use; no code changes. +- **Option C (thin proxy MCP):** Extra package and latency; duplicates routing logic already in `ServerContext`. + +**Rationale:** Pinecone SDK v8 supports multiple `Pinecone({ apiKey })` instances per process; MCP has no barrier to aggregating backends. Security is enforced by **deployment profiles** (public-only vs merged config), not per-query MCP authorization. All multi-source results include `source` for LLM provenance. + +**Execution semantics:** Discovery tools aggregate all sources when `source` is omitted. Execution tools (`query`, `count`, etc.) call `resolveSource`: infer source when the namespace exists on exactly one project; return `VALIDATION` when ambiguous. They do **not** fan out one query to all sources. `guided_query` without `namespace` routes via the aggregated namespace list and sets `selected_source` in `decision_trace`. + +**Audit logging:** When a tool resolves a specific source, stderr logs `toolname [source=name]` at INFO (execution tools and discovery tools that pass an explicit `source` filter). Aggregated discovery without `source` does not log per-source lines. + +See also [TOOLS.md § Multi-source mode](./TOOLS.md#multi-source-mode). + +--- + ## CLI flags (`parseCli` / `src/cli.ts`) | Flag | Maps to | @@ -58,6 +167,8 @@ C++ Alliance deployers can copy [examples/alliance/.env.example](../examples/all | `--request-timeout-ms` | `requestTimeoutMs` | | `--disable-suggest-flow` | `disableSuggestFlow: true` | | `--check-indexes` | `checkIndexes: true` | +| `--sources` | `sources` (inline multi-source string) | +| `--config-file` | `configFile` / `PINECONE_CONFIG_FILE` | | `--help` / `-h` | Print help and exit | | `--version` / `-v` | Print version and exit | diff --git a/docs/MIGRATION.md b/docs/MIGRATION.md index 1fb013a..7a34884 100644 --- a/docs/MIGRATION.md +++ b/docs/MIGRATION.md @@ -165,6 +165,26 @@ await setupCoreServer({ context: ctx }); See [deprecation-policy.md § Active deprecations](./deprecation-policy.md#active-deprecations-legacy-module-facades) for the full inventory. +### Multi-source mode and legacy facades + +When `PINECONE_SOURCES` or a config file is active, legacy module facades (`getNamespacesWithCache`, `registerUrlGenerator`, `markSuggested`, `requireSuggested`, etc.) operate on the **default source only** (first inline source, or `defaultSource` from JSON). They do not aggregate across projects. Prefer an explicit `ServerContext` built with `sourceRegistry` from `buildSourceRegistry()` for multi-tenant or multi-project embedding. + +## Unreleased: Multi-source Pinecone projects + +**Who is affected:** Operators running separate MCP entries per Pinecone project, or library embedders needing multiple indexes. + +**After (single MCP server):** + +```bash +PINECONE_SOURCES=api_key_1:${PINECONE_API_KEY_1}:index_name_1;api_key_2:${PINECONE_API_KEY_2}:index_name_2 +``` + +Or point `PINECONE_CONFIG_FILE` at a JSON file (see [examples/multi-source/pinecone-sources.json.example](../examples/multi-source/pinecone-sources.json.example)). + +**Client flow:** `list_sources` → `list_namespaces` (note `source` on each row) → pass `source` on `query` / `count` / etc. when a namespace name is ambiguous. + +Single-key configs are unchanged; no migration required when using one API key. + ## Unreleased: trimmed library exports **Who is affected:** Library embedders that imported `buildQueryExperimental` or `buildGuidedQueryExperimental` from `@will-cppa/pinecone-read-only-mcp` or `/alliance`. diff --git a/docs/README.md b/docs/README.md index 0b183a0..f587083 100644 --- a/docs/README.md +++ b/docs/README.md @@ -4,8 +4,8 @@ Guides for operators, MCP client authors, and library embedders. | Guide | Description | | ----- | ----------- | -| [TOOLS.md](./TOOLS.md) | All MCP tools, parameters, success/error shapes, suggest-flow | -| [CONFIGURATION.md](./CONFIGURATION.md) | Environment variables, CLI flags, `resolveConfig`, precedence | +| [TOOLS.md](./TOOLS.md) | All MCP tools, parameters, success/error shapes, suggest-flow, multi-source | +| [CONFIGURATION.md](./CONFIGURATION.md) | Environment variables, CLI flags, `resolveConfig`, multi-source mode, deployment profiles | | [SECURITY.md](./SECURITY.md) | API keys, log redaction, Docker hardening, reporting issues | | [CONTRIBUTING.md](../CONTRIBUTING.md) | Local dev, tests, lint/format, PR expectations | | [CI_CD.md](./CI_CD.md) | GitHub Actions, CodeQL, SBOM, Codecov, releases | diff --git a/docs/SECURITY.md b/docs/SECURITY.md index ff02c97..bbf1b51 100644 --- a/docs/SECURITY.md +++ b/docs/SECURITY.md @@ -2,8 +2,9 @@ ## API keys -- **Never** commit real Pinecone API keys. Use environment variables (`PINECONE_API_KEY`) or secret managers in CI. -- The CLI and `resolveConfig` read keys only from argv/env/overrides — logs must not echo raw keys. +- **Never** commit real Pinecone API keys. Use environment variables (`PINECONE_API_KEY`, or per-source keys referenced from `PINECONE_SOURCES` / a JSON config file) or secret managers in CI. +- The CLI and `resolveConfig` read keys only from argv/env/overrides — logs must not echo raw keys. In multi-source mode, each source may use a different API key; all are redacted in logs and MCP responses. +- Use **separate deployment profiles** for external (public-only) vs internal (merged) MCP configs — see [CONFIGURATION.md § Deployment profiles](./CONFIGURATION.md#deployment-profiles). ## Log redaction diff --git a/docs/TOOLS.md b/docs/TOOLS.md index cd28e82..9797f90 100644 --- a/docs/TOOLS.md +++ b/docs/TOOLS.md @@ -8,15 +8,18 @@ Success payloads separate **stable** fields (safe across minor bumps after `1.0. | Tool | Stable | Experimental | | ---- | ------ | ------------ | -| `list_namespaces` | `status`, `cache_hit`, `cache_ttl_seconds`, `expires_at_iso`, `count`, `namespaces` | _(none)_ | -| `namespace_router` | `status`, `cache_hit`, `user_query`, `suggestions`, `recommended_namespace` | _(none)_ | -| `suggest_query_params` | `status`, `cache_hit`, `suggested_fields`, `recommended_tool`, `use_count_tool`, `explanation`, `namespace_found` | _(none)_ | -| `count` | `status`, `count`, `truncated`, `namespace`, `metadata_filter` | _(none)_ | -| `query` | `status`, `mode`, `query`, `namespace`, `metadata_filter`, `result_count`, `results`, `fields` | `experimental.degraded`, `experimental.degradation_reason`, `experimental.hybrid_leg_failed`, `experimental.rerank_skipped_reason` | -| `keyword_search` | `status`, `query`, `namespace`, `index`, `metadata_filter`, `result_count`, `results`, `fields` | _(none)_ | -| `query_documents` | `status`, `query`, `namespace`, `metadata_filter`, `result_count`, `documents` | Same experimental degradation fields as `query` | -| `guided_query` | `status`, `result` (count or query-shaped stable fields) | `experimental.decision_trace`; query-path `result.experimental` degradation fields | -| `generate_urls` | `status`, `namespace`, `count`, `results` | _(none)_ | +| `list_sources` | `status`, `sources`, `default` | _(none)_ | +| `list_namespaces` | `status`, `cache_hit`, `cache_ttl_seconds`, `expires_at_iso`, `count`, `namespaces`, optional `source_errors` | _(none)_ | +| `namespace_router` | `status`, `cache_hit`, `user_query`, `suggestions`, `recommended_namespace`, optional `recommended_source` | _(none)_ | +| `suggest_query_params` | `status`, `cache_hit`, `suggested_fields`, `recommended_tool`, `use_count_tool`, `explanation`, `namespace_found`, optional `source` | _(none)_ | +| `count` | `status`, `count`, `truncated`, `namespace`, `metadata_filter`, optional `source` | _(none)_ | +| `query` | `status`, `mode`, `query`, `namespace`, `metadata_filter`, `result_count`, `results`, `fields`, optional `source` on payload and rows | `experimental.degraded`, `experimental.degradation_reason`, `experimental.hybrid_leg_failed`, `experimental.rerank_skipped_reason` | +| `keyword_search` | `status`, `query`, `namespace`, `index`, `metadata_filter`, `result_count`, `results`, `fields`, optional `source` | _(none)_ | +| `query_documents` | `status`, `query`, `namespace`, `metadata_filter`, `result_count`, `documents`, optional `source` | Same experimental degradation fields as `query` | +| `guided_query` | `status`, `result` (count or query-shaped stable fields) | `experimental.decision_trace` (includes optional `selected_source`); query-path `result.experimental` degradation fields | +| `generate_urls` | `status`, `namespace`, `count`, `results`, optional `source` | _(none)_ | + +Multi-source-only stable fields (`source` on namespace rows, `source_errors`, `recommended_source`) are **omitted** in single-key deployments. `selected_source` appears only under `guided_query` → `experimental.decision_trace` (experimental, not stable). Promotion process: [deprecation-policy.md § Stable vs experimental](./deprecation-policy.md#stable-vs-experimental-mcp-response-fields). @@ -24,8 +27,8 @@ Promotion process: [deprecation-policy.md § Stable vs experimental](./deprecati | Setup | Tools | MCP instructions | | ----- | ----- | ------------------ | -| `setupCoreServer` (package root) | **8:** `list_namespaces`, `namespace_router`, `count`, `query`, `keyword_search`, `query_documents`, `generate_urls`, `guided_query` | `CORE_SERVER_INSTRUCTIONS` — includes `guided_query`; no `suggest_query_params` | -| `setupAllianceServer` / published CLI | **9:** core tools plus `suggest_query_params` | `ALLIANCE_SERVER_INSTRUCTIONS` — includes suggest-flow quickstart | +| `setupCoreServer` (package root) | **8** (single-key) or **9** (multi-source adds `list_sources`): `list_namespaces`, `namespace_router`, `count`, `query`, `keyword_search`, `query_documents`, `generate_urls`, `guided_query` | `CORE_SERVER_INSTRUCTIONS` — includes `guided_query`; no `suggest_query_params` | +| `setupAllianceServer` / published CLI | **9** (single-key) or **10** (multi-source): core tools plus `suggest_query_params` | `ALLIANCE_SERVER_INSTRUCTIONS` — includes suggest-flow quickstart | ## Suggest-flow gate @@ -37,6 +40,48 @@ When **`disableSuggestFlow`** is **`true`** (core default via `resolveConfig`), **Core:** gate off by default; set `PINECONE_DISABLE_SUGGEST_FLOW=false` or `disableSuggestFlow: false` to enable the gate. **Alliance:** gate on by default; set `PINECONE_DISABLE_SUGGEST_FLOW=true` or CLI `--disable-suggest-flow` to bypass (not recommended for production). +## Multi-source mode + +Activated when `PINECONE_SOURCES`, `--sources`, or a JSON config file (`PINECONE_CONFIG_FILE` / `--config-file`) is set. See [CONFIGURATION.md § Multi-source mode](./CONFIGURATION.md#multi-source-mode) and [Deployment profiles](./CONFIGURATION.md#deployment-profiles). + +**Shared `source` parameter** (optional on most tools): + +> Pinecone source name (from `list_sources`). Omit on discovery tools to search all sources. On query tools, omit only when the namespace uniquely identifies one source. + +| Category | Tools | When `source` is omitted | +| -------- | ----- | ------------------------ | +| **Discovery** | `list_sources`, `list_namespaces`, `namespace_router` | Aggregate or list across **all** configured sources; rows include `source` | +| **Orchestrator** | `guided_query` (no `namespace`) | Route using aggregated namespace lists; `experimental.decision_trace.selected_source` | +| **Execution** | `suggest_query_params`, `query`, `count`, `query_documents`, `keyword_search`, `generate_urls`, `guided_query` (with `namespace`) | `resolveSource`: infer source when namespace is unique; `VALIDATION` (`field: source`) when ambiguous | + +**Typical multi-source flow:** + +```text +list_sources → list_namespaces → (optional) namespace_router → suggest_query_params → query | count | query_documents +``` + +Or single-shot: `guided_query` (routes across sources when `namespace` is omitted). + +**Suggest-flow in multi-source mode:** gate state uses compound keys `source:namespace`. Pass the same `source` + `namespace` pair for `suggest_query_params` and gated tools when the namespace exists on multiple sources. + +--- + +## `list_sources` (multi-source only) + +Registered only when more than one Pinecone source is configured. + +| | | +| --- | --- | +| **Input** | _(empty object)_ | +| **Success** | `{ status: 'success', sources: string[], default: string }` | +| **Errors** | `LIFECYCLE` when not in multi-source mode | + +**Example:** + +```json +{} +``` + --- ## 1. `list_namespaces` @@ -45,16 +90,22 @@ When **`disableSuggestFlow`** is **`true`** (core default via `resolveConfig`), | | | | --- | --- | -| **Input** | _(empty object)_ | -| **Success** | `{ status: 'success', cache_hit, cache_ttl_seconds, expires_at_iso, count, namespaces: [{ name, record_count, metadata_fields }] }` | +| **Input** | Optional `source` — filter to one configured project | +| **Success** | `{ status: 'success', cache_hit, cache_ttl_seconds, expires_at_iso, count, namespaces: [{ name, record_count, metadata_fields, source? }], source_errors? }` | | **Errors** | `PINECONE_ERROR`, `TIMEOUT`, etc. | -**Example (conceptual MCP params):** +**Example (multi-source, all projects):** ```json {} ``` +**Example (multi-source, one project):** + +```json +{ "source": "api_key_1" } +``` + --- ## 2. `namespace_router` @@ -65,8 +116,9 @@ When **`disableSuggestFlow`** is **`true`** (core default via `resolveConfig`), | ----- | ---- | -------- | ----------- | | `user_query` | string | yes | User question / intent | | `top_n` | int | no (default 3) | Max suggestions, 1–5 | +| `source` | string | no | Restrict ranking to one configured source (multi-source) | -**Success:** `{ status: 'success', cache_hit, user_query, suggestions, recommended_namespace }`. +**Success:** `{ status: 'success', cache_hit, user_query, suggestions: [{ namespace, score, record_count, reasons, source? }], recommended_namespace, recommended_source? }`. **Example:** @@ -84,8 +136,9 @@ When **`disableSuggestFlow`** is **`true`** (core default via `resolveConfig`), | ----- | ---- | -------- | ----------- | | `namespace` | string | yes | Target namespace (must exist in cached `list_namespaces`) | | `user_query` | string | yes | Natural-language task | +| `source` | string | no | Pinecone source (multi-source; required when namespace is ambiguous) | -**Success:** `{ status: 'success', cache_hit, ...suggestQueryParams fields including suggested_fields, recommended_tool, use_count_tool, explanation, namespace_found }`. +**Success:** `{ status: 'success', cache_hit, ...suggestQueryParams fields including suggested_fields, recommended_tool, use_count_tool, explanation, namespace_found, source? }`. **Example:** @@ -107,8 +160,9 @@ When **`disableSuggestFlow`** is **`true`** (core default via `resolveConfig`), | `namespace` | string | yes | Namespace | | `query_text` | string | yes | Query text (use broad text like `"document"` for metadata-only counts) | | `metadata_filter` | object | no | Pinecone metadata filter | +| `source` | string | no | Pinecone source (multi-source) | -**Success:** `{ status: 'success', count, truncated, namespace, metadata_filter? }`. +**Success:** `{ status: 'success', count, truncated, namespace, metadata_filter?, source? }`. --- @@ -125,8 +179,9 @@ When **`disableSuggestFlow`** is **`true`** (core default via `resolveConfig`), | `use_reranking` | boolean | no | When preset allows reranking | | `metadata_filter` | object | no | Metadata filter | | `fields` | string[] | no | Pinecone fields to return | +| `source` | string | no | Pinecone source (multi-source) | -**Success (`QueryResponse`):** `{ status: 'success', mode?, query, namespace, metadata_filter?, result_count, results[], fields?, experimental?: { degraded?, degradation_reason?, hybrid_leg_failed?, rerank_skipped_reason? } }`. +**Success (`QueryResponse`):** `{ status: 'success', mode?, query, namespace, metadata_filter?, result_count, results[], fields?, source?, experimental?: { ... } }`. Each row: `document_id`, `paper_number` (deprecated alias), `title`, `author`, `url`, `content`, `score`, `reranked`, optional `metadata`. @@ -168,8 +223,9 @@ Treat **`experimental.degraded: true`** as lower confidence even when `status` i | `top_k` | int | no | 1–100 | | `metadata_filter` | object | no | Filter | | `fields` | string[] | no | Returned fields | +| `source` | string | no | Pinecone source (multi-source) | -**Success:** Similar row shape to `query` (`KeywordSearchResponse`). +**Success:** Similar row shape to `query` (`KeywordSearchResponse`); optional `source` on payload and rows. --- @@ -184,8 +240,9 @@ Treat **`experimental.degraded: true`** as lower confidence even when `status` i | `top_k` | int | no | Documents to return (see constants, default 5, max 20) | | `metadata_filter` | object | no | Filter | | `max_chunks_per_document` | int | no | Cap merged chunks per doc (default 200, max 500) | +| `source` | string | no | Pinecone source (multi-source) | -**Success:** `{ status: 'success', query, namespace, metadata_filter?, result_count, documents[], experimental?: { degraded?, degradation_reason?, hybrid_leg_failed?, rerank_skipped_reason? } }`. +**Success:** `{ status: 'success', query, namespace, metadata_filter?, result_count, documents[], source?, experimental?: { ... } }`. --- @@ -201,10 +258,11 @@ Treat **`experimental.degraded: true`** as lower confidence even when `status` i | `top_k` | int | no | For query paths | | `preferred_tool` | `auto` \| `count` \| `fast` \| `detailed` \| `full` | no | Override automated tool choice | | `enrich_urls` | boolean | no (default true) | Run URL generator when metadata lacks `url` | +| `source` | string | no | Pinecone source (multi-source; pin routing when namespace is ambiguous) | **Success:** `{ status: 'success', experimental: { decision_trace }, result }` where `result` is either a count payload or a `QueryResponse`-shaped query payload. -**`experimental.decision_trace` fields (non-exhaustive):** `cache_hit`, `input_namespace`, `routed_namespace`, `selected_namespace`, `ranked_namespaces`, `suggested_fields`, `suggested_tool`, `selected_tool`, `explanation`, `enrich_urls`, `rerank_status` (`success` \| `skipped` \| `skipped_no_model` \| `failed`). +**`experimental.decision_trace` fields (non-exhaustive):** `cache_hit`, `input_namespace`, `routed_namespace`, `selected_namespace`, `selected_source?`, `ranked_namespaces`, `suggested_fields`, `suggested_tool`, `selected_tool`, `explanation`, `enrich_urls`, `rerank_status` (`success` \| `skipped` \| `skipped_no_model` \| `failed`). When the inner query path runs, `result.experimental` includes the same degradation fields as `query` (see [Rerank and hybrid degradation](#rerank-and-hybrid-degradation)). @@ -227,8 +285,9 @@ When the inner query path runs, `result.experimental` includes the same degradat | ----- | ---- | -------- | ----------- | | `namespace` | string | yes | Namespace | | `records` | object[] | yes | Up to 500 records (metadata object or `{ metadata: {...} }`) | +| `source` | string | no | Pinecone source (multi-source) | -**Success:** `{ status: 'success', namespace, count, results: [{ index, url, method, reason, metadata }] }`. +**Success:** `{ status: 'success', namespace, count, results: [...], source? }`. --- @@ -238,6 +297,9 @@ When the inner query path runs, `result.experimental` includes the same degradat Typical manual flow: list_namespaces → (optional) namespace_router → suggest_query_params → query | count | query_documents +Multi-source manual flow: + list_sources → list_namespaces → (optional) namespace_router → suggest_query_params → query | count | query_documents + Keyword-only: list_namespaces → keyword_search # no suggest gate diff --git a/examples/alliance/.env.example b/examples/alliance/.env.example index cb15c32..3bb804b 100644 --- a/examples/alliance/.env.example +++ b/examples/alliance/.env.example @@ -12,3 +12,7 @@ PINECONE_INDEX_NAME=rag-hybrid # Alliance rerank model (CLI default when unset: bge-reranker-v2-m3) # PINECONE_RERANK_MODEL=bge-reranker-v2-m3 + +# Optional: multiple Pinecone projects in one server (see docs/CONFIGURATION.md#multi-source-mode) +# PINECONE_SOURCES=api_key_1:${PINECONE_API_KEY_1}:index_name_1;api_key_2:${PINECONE_API_KEY_2}:index_name_2 +# PINECONE_CONFIG_FILE=./pinecone-sources.json diff --git a/examples/alliance/custom-url-generator.ts b/examples/alliance/custom-url-generator.ts index e105b71..a82676e 100644 --- a/examples/alliance/custom-url-generator.ts +++ b/examples/alliance/custom-url-generator.ts @@ -20,6 +20,7 @@ import { type UrlGenerationResult, } from '@will-cppa/pinecone-read-only-mcp'; import { resolveAllianceConfig, setupAllianceServer } from '@will-cppa/pinecone-read-only-mcp/alliance'; +import { exitOnDemoFailure } from '../lib/exit-on-failure.js'; async function main(): Promise { const apiKey = process.env['PINECONE_API_KEY']?.trim(); @@ -64,7 +65,4 @@ async function main(): Promise { console.log('Custom URL generator registered for namespace "product-docs".'); } -main().catch((err) => { - console.error(err); - process.exit(1); -}); +main().catch(exitOnDemoFailure('custom-url-generator')); diff --git a/examples/alliance/guided-query-demo.ts b/examples/alliance/guided-query-demo.ts index 3f34f7e..6d58763 100644 --- a/examples/alliance/guided-query-demo.ts +++ b/examples/alliance/guided-query-demo.ts @@ -20,6 +20,7 @@ import { PineconeClient, } from '@will-cppa/pinecone-read-only-mcp'; import { resolveAllianceConfig, setupAllianceServer } from '@will-cppa/pinecone-read-only-mcp/alliance'; +import { exitOnDemoFailure } from '../lib/exit-on-failure.js'; async function main(): Promise { const apiKey = process.env['PINECONE_API_KEY']?.trim(); @@ -50,7 +51,4 @@ async function main(): Promise { console.log('Server ready — call guided_query({ user_query, preferred_tool?: "auto" }).'); } -main().catch((err) => { - console.error(err); - process.exit(1); -}); +main().catch(exitOnDemoFailure('guided-query-demo')); diff --git a/examples/alliance/library-embedding-demo.ts b/examples/alliance/library-embedding-demo.ts index fe8e224..5ae7480 100644 --- a/examples/alliance/library-embedding-demo.ts +++ b/examples/alliance/library-embedding-demo.ts @@ -21,6 +21,7 @@ import { createServer, PineconeClient } from '@will-cppa/pinecone-read-only-mcp'; import { resolveAllianceConfig, setupAllianceServer } from '@will-cppa/pinecone-read-only-mcp/alliance'; +import { exitOnDemoFailure } from '../lib/exit-on-failure.js'; async function main(): Promise { const apiKey = process.env['PINECONE_API_KEY']?.trim(); @@ -51,7 +52,4 @@ async function main(): Promise { console.log('Embedded server constructed — connect your MCP transport (stdio, HTTP, etc.).'); } -main().catch((err) => { - console.error(err); - process.exit(1); -}); +main().catch(exitOnDemoFailure('library-embedding-demo')); diff --git a/examples/alliance/suggest-flow-demo.ts b/examples/alliance/suggest-flow-demo.ts index 9433d57..b060d92 100644 --- a/examples/alliance/suggest-flow-demo.ts +++ b/examples/alliance/suggest-flow-demo.ts @@ -21,6 +21,7 @@ import { PineconeClient, } from '@will-cppa/pinecone-read-only-mcp'; import { resolveAllianceConfig, setupAllianceServer } from '@will-cppa/pinecone-read-only-mcp/alliance'; +import { exitOnDemoFailure } from '../lib/exit-on-failure.js'; async function main(): Promise { const apiKey = process.env['PINECONE_API_KEY']?.trim(); @@ -54,7 +55,4 @@ async function main(): Promise { console.log('Server ready — connect a transport and issue suggest_query_params then query.'); } -main().catch((err) => { - console.error(err); - process.exit(1); -}); +main().catch(exitOnDemoFailure('suggest-flow-demo')); diff --git a/examples/lib/exit-on-failure.ts b/examples/lib/exit-on-failure.ts new file mode 100644 index 0000000..af108e9 --- /dev/null +++ b/examples/lib/exit-on-failure.ts @@ -0,0 +1,13 @@ +import { redactApiKey } from '../../src/logger.js'; + +/** + * Exit a demo script with a redacted error detail (CodeQL: js/clear-text-logging). + */ +export function exitOnDemoFailure(label: string): (err: unknown) => void { + return (err: unknown) => { + const detail = redactApiKey(err instanceof Error ? err.message : String(err)); + console.error(`${label} failed: ${detail}`); + console.error('Check credentials and index configuration.'); + process.exitCode = 1; + }; +} diff --git a/examples/multi-source/pinecone-sources.json.example b/examples/multi-source/pinecone-sources.json.example new file mode 100644 index 0000000..095b9cc --- /dev/null +++ b/examples/multi-source/pinecone-sources.json.example @@ -0,0 +1,13 @@ +{ + "defaultSource": "api_key_1", + "sources": { + "api_key_1": { + "apiKey": "${PINECONE_API_KEY_1}", + "indexName": "index_name_1" + }, + "api_key_2": { + "apiKey": "${PINECONE_API_KEY_2}", + "indexName": "index_name_2" + } + } +} diff --git a/examples/quickstart/mcp-demo.ts b/examples/quickstart/mcp-demo.ts index 80229c5..3c39ae1 100644 --- a/examples/quickstart/mcp-demo.ts +++ b/examples/quickstart/mcp-demo.ts @@ -13,6 +13,7 @@ import { resolveConfig, setupCoreServer, } from '@will-cppa/pinecone-read-only-mcp'; +import { exitOnDemoFailure } from '../lib/exit-on-failure.js'; import { dirname, join } from 'node:path'; import { fileURLToPath } from 'node:url'; import { createLinkedTransports } from '../mcp-linked-transport.js'; @@ -130,7 +131,4 @@ main() .then(() => { process.exit(0); }) - .catch((err) => { - console.error(err); - process.exit(1); - }); + .catch(exitOnDemoFailure('mcp-demo')); diff --git a/examples/quickstart/seed-data.ts b/examples/quickstart/seed-data.ts index 6737a8b..59e61f6 100644 --- a/examples/quickstart/seed-data.ts +++ b/examples/quickstart/seed-data.ts @@ -9,6 +9,7 @@ import { config as loadEnv } from 'dotenv'; import { Pinecone } from '@pinecone-database/pinecone'; import { dirname, join, resolve } from 'node:path'; import { fileURLToPath } from 'node:url'; +import { exitOnDemoFailure } from '../lib/exit-on-failure.js'; const __dirname = dirname(fileURLToPath(import.meta.url)); loadEnv({ path: join(__dirname, '.env') }); @@ -169,8 +170,5 @@ const isDirectRun = process.argv[1] !== undefined && resolve(process.argv[1]) === fileURLToPath(import.meta.url); if (isDirectRun) { - main().catch((err) => { - console.error(err); - process.exit(1); - }); + main().catch(exitOnDemoFailure('seed-data')); } diff --git a/scripts/test-search.ts b/scripts/test-search.ts index 0294b2b..d63d7a3 100644 --- a/scripts/test-search.ts +++ b/scripts/test-search.ts @@ -12,6 +12,7 @@ */ import { PineconeClient } from '../src/pinecone-client.js'; +import { redactApiKey } from '../src/logger.js'; async function test() { const apiKey = process.env.PINECONE_API_KEY || process.argv[2]; @@ -214,16 +215,17 @@ async function test() { } console.log(` Reranking overhead: ${duration2 - duration1}ms`); } catch (error) { - console.error('\n❌ Error during testing:', error); - if (error instanceof Error) { - console.error(' Message:', error.message); - } + console.error('\n❌ Error during testing.'); + const message = error instanceof Error ? error.message : String(error); + console.error(' Message:', redactApiKey(message)); process.exit(1); } } // Run the test test().catch((error) => { - console.error('Fatal error:', error); + console.error('Fatal error during test-search.'); + const message = error instanceof Error ? error.message : String(error); + console.error(' Message:', redactApiKey(message)); process.exit(1); }); diff --git a/src/alliance/config.test.ts b/src/alliance/config.test.ts index df26c61..e6ba694 100644 --- a/src/alliance/config.test.ts +++ b/src/alliance/config.test.ts @@ -1,3 +1,6 @@ +import { mkdtempSync, writeFileSync } from 'node:fs'; +import { join } from 'node:path'; +import { tmpdir } from 'node:os'; import { describe, expect, it } from 'vitest'; import { ALLIANCE_DEFAULT_INDEX_NAME, @@ -55,4 +58,25 @@ describe('resolveAllianceConfig', () => { ); expect(cfg.indexName).toBe('custom-index'); }); + + it('applies resolved Alliance index default to multi-source entries missing indexName', () => { + const dir = mkdtempSync(join(tmpdir(), 'pinecone-alliance-sources-')); + const filePath = join(dir, 'sources.json'); + writeFileSync( + filePath, + JSON.stringify({ + defaultSource: 'api_key_1', + sources: { + api_key_1: { apiKey: 'k1', indexName: 'explicit-index' }, + api_key_2: { apiKey: 'k2' }, + }, + }) + ); + const cfg = resolveAllianceConfig( + { configFile: filePath, indexName: 'cli-default-index' }, + { PINECONE_API_KEY: 'ignored' } + ); + const second = cfg.sources?.find((s) => s.name === 'api_key_2'); + expect(second?.indexName).toBe('cli-default-index'); + }); }); diff --git a/src/alliance/config.ts b/src/alliance/config.ts index 0e8224e..2126f9b 100644 --- a/src/alliance/config.ts +++ b/src/alliance/config.ts @@ -10,6 +10,7 @@ import { type AllianceServerConfig, type ConfigOverrides, } from '../core/config.js'; +import type { ParseSourcesOptions } from '../core/server/source-config.js'; /** C++ Alliance default dense index when env/CLI omit `PINECONE_INDEX_NAME`. */ export const ALLIANCE_DEFAULT_INDEX_NAME = 'rag-hybrid'; @@ -47,7 +48,13 @@ export function resolveAllianceConfig( trimOptional(overrides.rerankModel) ?? trimOptional(env['PINECONE_RERANK_MODEL']) ?? ALLIANCE_DEFAULT_RERANK_MODEL; - const cfg = resolveConfig({ ...overrides, indexName, rerankModel }, env); + const allianceParseOptions: ParseSourcesOptions = { + allianceDefaults: { + indexName, + rerankModel, + }, + }; + const cfg = resolveConfig({ ...overrides, indexName, rerankModel }, env, allianceParseOptions); const disableSuggestFlow = overrides.disableSuggestFlow ?? asBool(env['PINECONE_DISABLE_SUGGEST_FLOW'], false); return brandAllianceConfig({ ...cfg, disableSuggestFlow }); diff --git a/src/alliance/tools/suggest-query-params-tool.test.ts b/src/alliance/tools/suggest-query-params-tool.test.ts index 9f7035c..36e580c 100644 --- a/src/alliance/tools/suggest-query-params-tool.test.ts +++ b/src/alliance/tools/suggest-query-params-tool.test.ts @@ -69,6 +69,21 @@ describe('suggest_query_params tool handler', () => { expect(mockedMarkSuggested).not.toHaveBeenCalled(); }); + it('returns VALIDATION when source is provided without ServerContext', async () => { + const server = createMockServer(); + registerSuggestQueryParamsTool(server as never); + const err = assertToolErrorCode( + await server.getHandler('suggest_query_params')!({ + namespace: 'wg21', + user_query: 'hello', + source: 'api_key_1', + }), + 'VALIDATION' + ); + expect(err.field).toBe('source'); + expect(mockedGetNamespaces).not.toHaveBeenCalled(); + }); + it('marks suggestion flow and returns success when namespace exists in cache', async () => { mockedGetNamespaces.mockResolvedValue({ data: [makeNamespaceCacheEntry('wg21')], diff --git a/src/alliance/tools/suggest-query-params-tool.ts b/src/alliance/tools/suggest-query-params-tool.ts index 09f846f..8112245 100644 --- a/src/alliance/tools/suggest-query-params-tool.ts +++ b/src/alliance/tools/suggest-query-params-tool.ts @@ -5,6 +5,13 @@ import { getNamespacesWithCache } from '../../core/server/namespaces-cache.js'; import { suggestQueryParams } from '../../core/server/query-suggestion.js'; import type { ServerContext } from '../../core/server/server-context.js'; import { markSuggested } from '../../core/server/suggestion-flow.js'; +import { + optionalSourceField, + rejectSourceWithoutContext, + resolveSourceForTool, + sourceParamSchema, + sourceValidationError, +} from '../../core/server/source-tool-utils.js'; import { classifyToolCatchError, lifecycleToolError, @@ -38,6 +45,7 @@ export function registerSuggestQueryParamsTool(server: McpServer, ctx?: ServerCo .describe( 'The user\'s natural language question or intent (e.g. "list documents by author X with titles and links", "how many records match Y?", "what do the docs say about Z?").' ), + source: sourceParamSchema, }, }, async (params) => { @@ -45,7 +53,7 @@ export function registerSuggestQueryParamsTool(server: McpServer, ctx?: ServerCo if (ctx?.disposed) { return jsonErrorResponse(lifecycleToolError('ServerContext has been disposed')); } - const { namespace, user_query } = params; + const { namespace, user_query, source } = params; if (!user_query?.trim()) { return jsonErrorResponse(validationToolError('user_query cannot be empty', 'user_query')); } @@ -57,8 +65,25 @@ export function registerSuggestQueryParamsTool(server: McpServer, ctx?: ServerCo }) ); } - const { data: namespacesInfo, cache_hit } = ctx - ? await ctx.getNamespacesWithCache() + + const sourceError = rejectSourceWithoutContext(source, ctx); + if (sourceError) { + return jsonErrorResponse(sourceError); + } + + let activeCtx = ctx; + let activeSource: string | undefined; + if (ctx?.isMultiSource()) { + const resolved = await resolveSourceForTool(ctx, source, nsNorm); + if (!resolved.ok) { + return jsonErrorResponse(sourceValidationError(resolved.code, resolved.message)); + } + activeCtx = resolved.ctx; + activeSource = resolved.source; + } + + const { data: namespacesInfo, cache_hit } = activeCtx + ? await activeCtx.getNamespacesWithCache(activeSource) : await getNamespacesWithCache(); const ns = namespacesInfo.find( (n) => n.namespace === nsNorm || normalizeNamespace(n.namespace) === nsNorm @@ -66,12 +91,16 @@ export function registerSuggestQueryParamsTool(server: McpServer, ctx?: ServerCo const metadataFields = ns?.metadata ?? null; const result = suggestQueryParams(metadataFields, user_query.trim()); if (result.namespace_found) { - if (ctx) { - ctx.markSuggested(nsNorm, { - recommended_tool: result.recommended_tool, - suggested_fields: result.suggested_fields, - user_query: user_query.trim(), - }); + if (activeCtx) { + activeCtx.markSuggested( + nsNorm, + { + recommended_tool: result.recommended_tool, + suggested_fields: result.suggested_fields, + user_query: user_query.trim(), + }, + activeCtx.isMultiSource() ? activeSource : undefined + ); } else { markSuggested(nsNorm, { recommended_tool: result.recommended_tool, @@ -84,6 +113,7 @@ export function registerSuggestQueryParamsTool(server: McpServer, ctx?: ServerCo ...result, status: 'success', cache_hit, + ...optionalSourceField(activeCtx, activeSource), }; return validatedJsonResponse(suggestQueryParamsResponseSchema, response); } catch (error) { diff --git a/src/cli.ts b/src/cli.ts index 89bbb95..8e7b05b 100644 --- a/src/cli.ts +++ b/src/cli.ts @@ -119,6 +119,22 @@ export function parseCli(argv: string[] = process.argv.slice(2)): ParseCliResult case '--check-indexes': overrides.checkIndexes = true; break; + case '--sources': { + const v = readOptionValue(argv, i); + if (v !== undefined) { + overrides.sources = v; + i++; + } + break; + } + case '--config-file': { + const v = readOptionValue(argv, i); + if (v !== undefined) { + overrides.configFile = v; + i++; + } + break; + } default: break; } @@ -146,6 +162,8 @@ Options: --request-timeout-ms N Per Pinecone call timeout [env: PINECONE_REQUEST_TIMEOUT_MS] --disable-suggest-flow Bypass suggest_query_params gate (PINECONE_DISABLE_SUGGEST_FLOW) --check-indexes Verify dense + sparse indexes then exit 0/1 (PINECONE_CHECK_INDEXES) + --sources TEXT Multi-source inline config (PINECONE_SOURCES) + --config-file PATH Multi-source JSON config file (PINECONE_CONFIG_FILE) --help, -h Show this message --version, -v Print package version diff --git a/src/constants.ts b/src/constants.ts index 5fad13a..a317888 100644 --- a/src/constants.ts +++ b/src/constants.ts @@ -44,6 +44,7 @@ Features: - URL Generation: Use generate_urls to synthesize URLs for namespaces that have a registered generator when metadata lacks url. - Document reassembly: Use query_documents to get whole documents (chunks grouped and merged by document_number/doc_id/url) for content analysis or summarization. query_documents reranks when a rerank model is configured. - Keyword search: Use keyword_search to query the sparse index for lexical/keyword-only retrieval without reranking. +- Multi-Source: When PINECONE_SOURCES or a config file is set, multiple Pinecone projects are available in one server. Use list_sources and pass source on tools; list_namespaces tags each namespace with its source. Notes: - Result rows include both \`document_id\` (canonical) and \`paper_number\` (deprecated alias kept for one minor cycle; will be removed in the next major release). Prefer \`document_id\` in new code. @@ -58,7 +59,9 @@ Usage: 1. Prefer guided_query for single-call retrieval (no prerequisite tools). 2. Use list_namespaces (cached) to discover available namespaces in the index. The response includes \`expires_at_iso\` so you know when to refresh. 3. Optionally use namespace_router to choose candidate namespace(s) from user intent. -4. Use count for count questions, \`query\` with the appropriate preset for chunk-level retrieval, query_documents for full-document content, keyword_search for lexical retrieval, or generate_urls when records need synthesized URLs.`; +4. Use count for count questions, \`query\` with the appropriate preset for chunk-level retrieval, query_documents for full-document content, keyword_search for lexical retrieval, or generate_urls when records need synthesized URLs. + +Multi-source (when configured): call list_sources, then list_namespaces (all sources unless source is set). Pass source when a namespace may exist on multiple projects. Treat source on results as provenance.`; /** Alliance-only supplement appended to core instructions for {@link setupAllianceServer}. */ export const ALLIANCE_INSTRUCTIONS_APPENDIX = ` diff --git a/src/core/config.ts b/src/core/config.ts index 000d8d6..0f59a3d 100644 --- a/src/core/config.ts +++ b/src/core/config.ts @@ -8,6 +8,8 @@ */ import { DEFAULT_TOP_K, FLOW_CACHE_TTL_MS } from '../constants.js'; +import type { ParseSourcesOptions, SourceDefinition } from './server/source-config.js'; +import { resolveSourceDefinitions } from './server/source-config.js'; /** Allowed log levels, in ascending severity. */ export type LogLevel = 'DEBUG' | 'INFO' | 'WARN' | 'ERROR'; @@ -55,6 +57,10 @@ export interface ServerConfigBase { disableSuggestFlow: boolean; /** When true, on-startup probe verifies dense + sparse indexes exist. */ checkIndexes: boolean; + /** Named Pinecone sources when multi-project mode is enabled. */ + sources?: SourceDefinition[]; + /** Default source name when multi-source mode is enabled. */ + defaultSource?: string; } /** Backward-compatible alias for {@link ServerConfigBase} (read paths, docs). */ @@ -144,6 +150,8 @@ export interface ConfigOverrides { requestTimeoutMs?: number; disableSuggestFlow?: boolean; checkIndexes?: boolean; + sources?: string; + configFile?: string; } /** @@ -167,8 +175,53 @@ export interface ConfigOverrides { */ export function resolveConfig( overrides: ConfigOverrides, - env: NodeJS.ProcessEnv = process.env + env: NodeJS.ProcessEnv = process.env, + parseSourcesOptions?: ParseSourcesOptions ): CoreServerConfig { + const defaultTopK = overrides.defaultTopK ?? asPositiveInt(env['PINECONE_TOP_K'], DEFAULT_TOP_K); + const logLevel = asLogLevel( + overrides.logLevel ?? env['PINECONE_READ_ONLY_MCP_LOG_LEVEL'], + 'INFO' + ); + const logFormat = asLogFormat( + overrides.logFormat ?? env['PINECONE_READ_ONLY_MCP_LOG_FORMAT'], + 'text' + ); + const cacheTtlSeconds = + overrides.cacheTtlSeconds ?? + asPositiveInt(env['PINECONE_CACHE_TTL_SECONDS'], FLOW_CACHE_TTL_MS / 1000); + const requestTimeoutMs = + overrides.requestTimeoutMs ?? + asPositiveInt(env['PINECONE_REQUEST_TIMEOUT_MS'], DEFAULT_REQUEST_TIMEOUT_MS); + const disableSuggestFlow = + overrides.disableSuggestFlow ?? asBool(env['PINECONE_DISABLE_SUGGEST_FLOW'], true); + const checkIndexes = overrides.checkIndexes ?? asBool(env['PINECONE_CHECK_INDEXES'], false); + + const multi = resolveSourceDefinitions(overrides, env, parseSourcesOptions); + if (multi) { + const primary = multi.sources.find((s) => s.name === multi.defaultSource) ?? multi.sources[0]!; + if (trimOptional(env['PINECONE_API_KEY']) || trimOptional(overrides.apiKey)) { + process.stderr.write( + 'Warning: PINECONE_SOURCES / config file is active; PINECONE_API_KEY is ignored.\n' + ); + } + return brandCoreConfig({ + apiKey: primary.apiKey, + indexName: primary.indexName, + sparseIndexName: primary.sparseIndexName ?? `${primary.indexName}-sparse`, + ...(primary.rerankModel !== undefined ? { rerankModel: primary.rerankModel } : {}), + defaultTopK, + logLevel, + logFormat, + cacheTtlMs: cacheTtlSeconds * 1000, + requestTimeoutMs, + disableSuggestFlow, + checkIndexes, + sources: multi.sources, + defaultSource: multi.defaultSource, + }); + } + const apiKey = (overrides.apiKey ?? env['PINECONE_API_KEY'] ?? '').trim(); if (!apiKey) { throw new Error( @@ -189,25 +242,6 @@ export function resolveConfig( const rerankModel = trimOptional(overrides.rerankModel ?? env['PINECONE_RERANK_MODEL']); - const defaultTopK = overrides.defaultTopK ?? asPositiveInt(env['PINECONE_TOP_K'], DEFAULT_TOP_K); - const logLevel = asLogLevel( - overrides.logLevel ?? env['PINECONE_READ_ONLY_MCP_LOG_LEVEL'], - 'INFO' - ); - const logFormat = asLogFormat( - overrides.logFormat ?? env['PINECONE_READ_ONLY_MCP_LOG_FORMAT'], - 'text' - ); - const cacheTtlSeconds = - overrides.cacheTtlSeconds ?? - asPositiveInt(env['PINECONE_CACHE_TTL_SECONDS'], FLOW_CACHE_TTL_MS / 1000); - const requestTimeoutMs = - overrides.requestTimeoutMs ?? - asPositiveInt(env['PINECONE_REQUEST_TIMEOUT_MS'], DEFAULT_REQUEST_TIMEOUT_MS); - const disableSuggestFlow = - overrides.disableSuggestFlow ?? asBool(env['PINECONE_DISABLE_SUGGEST_FLOW'], true); - const checkIndexes = overrides.checkIndexes ?? asBool(env['PINECONE_CHECK_INDEXES'], false); - return brandCoreConfig({ apiKey, indexName, @@ -222,3 +256,5 @@ export function resolveConfig( checkIndexes, }); } + +export type { SourceDefinition } from './server/source-config.js'; diff --git a/src/core/index.ts b/src/core/index.ts index 67a9339..f603331 100644 --- a/src/core/index.ts +++ b/src/core/index.ts @@ -66,6 +66,9 @@ export { } from './server/url-registry.js'; export type { UrlGenerationResult, UrlGenerator, UrlGeneratorFn } from './server/url-registry.js'; export { resolveConfig, trimOptional } from './config.js'; +export type { SourceDefinition } from './server/source-config.js'; +export { SourceRegistry, buildSourceRegistry } from './server/source-registry.js'; +export type { AggregatedCacheResult, PerSourceCacheResult } from './server/source-registry.js'; export type { ServerConfig, ServerConfigBase, diff --git a/src/core/pinecone-client.ts b/src/core/pinecone-client.ts index 85ce4b7..7ba81d3 100644 --- a/src/core/pinecone-client.ts +++ b/src/core/pinecone-client.ts @@ -36,7 +36,11 @@ export class PineconeClient { * built via {@link resolveConfig} / CLI); this class does not read `process.env`. */ constructor(config: PineconeClientConfig) { - this.indexSession = new PineconeIndexSession(config.apiKey, config.indexName); + this.indexSession = new PineconeIndexSession( + config.apiKey, + config.indexName, + config.sparseIndexName + ); const normalizedRerankModel = config.rerankModel?.trim(); this.rerankModel = normalizedRerankModel ? normalizedRerankModel : undefined; this.defaultTopK = config.defaultTopK ?? DEFAULT_TOP_K; diff --git a/src/core/pinecone/indexes.ts b/src/core/pinecone/indexes.ts index 60dbcc0..63099ed 100644 --- a/src/core/pinecone/indexes.ts +++ b/src/core/pinecone/indexes.ts @@ -37,12 +37,13 @@ export class PineconeIndexSession { constructor( private readonly apiKey: string, - private readonly indexName: string + private readonly indexName: string, + private readonly sparseIndexName?: string ) {} - /** Same as hybrid sparse index name: `{indexName}-sparse`. */ + /** Sparse index name; defaults to `{indexName}-sparse`. */ getSparseIndexName(): string { - return `${this.indexName}-sparse`; + return this.sparseIndexName ?? `${this.indexName}-sparse`; } /** Ensure Pinecone client is initialized */ diff --git a/src/core/server/format-query-result.context.test.ts b/src/core/server/format-query-result.context.test.ts index af8f24e..f58c301 100644 --- a/src/core/server/format-query-result.context.test.ts +++ b/src/core/server/format-query-result.context.test.ts @@ -33,4 +33,23 @@ describe('formatSearchResultAsRow (ServerContext instance path)', () => { }); expect(row.url).toBe('https://ctx.example/papers/doc-1'); }); + + it('omits row source in single-source mode even when source option is set', () => { + isolateFromDefaultContext(); + const ctx = new ServerContext(resolveTestConfig()); + const doc: SearchResult = { + id: 'v1', + content: 'body', + score: 0.9, + metadata: { document_number: 'DOC-1', title: 'T', author: 'A' }, + reranked: false, + }; + + const row = formatSearchResultAsRow(doc, { + namespace: 'papers', + source: 'api_key_1', + ctx, + }); + expect(row.source).toBeUndefined(); + }); }); diff --git a/src/core/server/format-query-result.ts b/src/core/server/format-query-result.ts index 98f105e..957dc51 100644 --- a/src/core/server/format-query-result.ts +++ b/src/core/server/format-query-result.ts @@ -11,6 +11,7 @@ import { generateUrlForNamespace } from './url-registry.js'; export type FormatQueryResultOptions = { namespace?: string; + source?: string; enrichUrls?: boolean; contentMaxLength?: number; ctx?: ServerContext; @@ -46,7 +47,7 @@ export function formatSearchResultAsRow( if (options?.enrichUrls && options?.namespace) { const generated = options.ctx - ? options.ctx.generateUrlForNamespace(options.namespace, metadata) + ? options.ctx.generateUrlForNamespace(options.namespace, metadata, options.source) : generateUrlForNamespace(options.namespace, metadata); const existingUrl = metadata['url']; const urlIsBlank = typeof existingUrl !== 'string' || existingUrl.trim() === ''; @@ -81,6 +82,7 @@ export function formatSearchResultAsRow( score: Math.round(doc.score * 10000) / 10000, reranked: doc.reranked, metadata, + ...(options?.source && options?.ctx?.isMultiSource() ? { source: options.source } : {}), }; } diff --git a/src/core/server/multi-source.context.test.ts b/src/core/server/multi-source.context.test.ts new file mode 100644 index 0000000..dba98b4 --- /dev/null +++ b/src/core/server/multi-source.context.test.ts @@ -0,0 +1,87 @@ +import { describe, expect, it, vi } from 'vitest'; +import { makeMockPineconeClient, createMultiSourceTestContext } from './tools/test-helpers.js'; + +describe('multi-source ServerContext', () => { + it('resolveSource returns AMBIGUOUS_NAMESPACE when namespace exists on multiple sources', async () => { + const { ctx } = createMultiSourceTestContext(); + const result = await ctx.resolveSource(undefined, 'shared'); + expect(result).toEqual({ + ok: false, + code: 'AMBIGUOUS_NAMESPACE', + message: 'Namespace "shared" exists on multiple sources. Pass source explicitly.', + }); + }); + + it('resolveSource fails when namespace inference uses partial aggregation', async () => { + const client1 = makeMockPineconeClient(['wg21']); + const client2 = { + listNamespacesWithMetadata: vi.fn().mockRejectedValue(new Error('api_key_2 unreachable')), + query: vi.fn(), + count: vi.fn(), + keywordSearch: vi.fn(), + checkIndexes: vi.fn().mockResolvedValue({ ok: true, errors: [] }), + getSparseIndexName: () => 'sparse', + }; + const clients = new Map([ + ['api_key_1', client1], + ['api_key_2', client2], + ]); + const { ctx } = createMultiSourceTestContext({ clients }); + const result = await ctx.resolveSource(undefined, 'wg21'); + expect(result).toEqual({ + ok: false, + code: 'PARTIAL_SOURCE_AGGREGATION', + message: + 'Namespace discovery is incomplete because one or more sources failed. Pass source explicitly or retry after resolving source_errors.', + }); + }); + + it('resolveSource infers source when namespace is unique', async () => { + const { ctx } = createMultiSourceTestContext(); + const result = await ctx.resolveSource(undefined, 'wg21'); + expect(result).toEqual({ ok: true, source: 'api_key_1' }); + }); + + it('isolates compound suggest-flow keys per source', () => { + const { ctx } = createMultiSourceTestContext(); + ctx.markSuggested( + 'shared', + { recommended_tool: 'fast', suggested_fields: ['title'], user_query: 'q1' }, + 'api_key_1' + ); + ctx.markSuggested( + 'shared', + { recommended_tool: 'count', suggested_fields: [], user_query: 'q2' }, + 'api_key_2' + ); + expect(ctx.requireSuggested('shared', 'api_key_1').ok).toBe(true); + expect(ctx.requireSuggested('shared', 'api_key_2').ok).toBe(true); + expect(ctx.requireSuggested('shared', 'api_key_1').flow?.user_query).toBe('q1'); + expect(ctx.requireSuggested('shared', 'api_key_2').flow?.user_query).toBe('q2'); + }); + + it('registers URL generators per source without collision', () => { + const { ctx } = createMultiSourceTestContext(); + ctx.registerUrlGenerator( + 'shared', + () => ({ url: 'https://api-key-1.example', method: 'generated.custom' }), + 'api_key_1' + ); + ctx.registerUrlGenerator( + 'shared', + () => ({ url: 'https://api-key-2.example', method: 'generated.custom' }), + 'api_key_2' + ); + expect(ctx.generateUrlForNamespace('shared', {}, 'api_key_1').url).toBe( + 'https://api-key-1.example' + ); + expect(ctx.generateUrlForNamespace('shared', {}, 'api_key_2').url).toBe( + 'https://api-key-2.example' + ); + expect(ctx.hasUrlGenerator('shared')).toBe(true); + expect(ctx.hasUrlGenerator('shared', 'api_key_1')).toBe(true); + expect(ctx.unregisterUrlGenerator('shared', 'api_key_1')).toBe(true); + expect(ctx.hasUrlGenerator('shared', 'api_key_1')).toBe(false); + expect(ctx.hasUrlGenerator('shared', 'api_key_2')).toBe(true); + }); +}); diff --git a/src/core/server/namespace-router.ts b/src/core/server/namespace-router.ts index 5db484a..a7a093c 100644 --- a/src/core/server/namespace-router.ts +++ b/src/core/server/namespace-router.ts @@ -5,6 +5,7 @@ export type RankedNamespace = { score: number; record_count: number; reasons: string[]; + source?: string; }; /** @@ -66,6 +67,7 @@ export function rankNamespacesByQuery( score, record_count: ns.recordCount, reasons, + ...(ns.source !== undefined ? { source: ns.source } : {}), }; }) .sort((a, b) => { diff --git a/src/core/server/response-schemas.ts b/src/core/server/response-schemas.ts index 889b8ec..d428a9b 100644 --- a/src/core/server/response-schemas.ts +++ b/src/core/server/response-schemas.ts @@ -32,6 +32,7 @@ export const queryResultRowSchema = z.object({ score: z.number(), reranked: z.boolean(), metadata: z.record(z.string(), pineconeMetadataValueSchema).optional(), + source: z.string().optional(), }); export type QueryResultRowShape = z.infer; @@ -51,6 +52,7 @@ const rankedNamespaceSchema = z.object({ score: z.number(), record_count: z.number(), reasons: z.array(z.string()), + source: z.string().optional(), }); export const guidedQueryDecisionTraceSchema = z.object({ @@ -58,6 +60,7 @@ export const guidedQueryDecisionTraceSchema = z.object({ input_namespace: z.string().nullable(), routed_namespace: z.string().nullable(), selected_namespace: z.string(), + selected_source: z.string().optional(), ranked_namespaces: z.array(rankedNamespaceSchema), suggested_fields: z.array(z.string()), suggested_tool: z.enum(['count', 'fast', 'detailed', 'full']), @@ -80,6 +83,7 @@ export const querySuccessResponseSchema = z.object({ mode: z.enum(['query', 'query_fast', 'query_detailed']), query: z.string(), namespace: z.string(), + source: z.string().optional(), metadata_filter: z.record(z.string(), z.unknown()).optional(), result_count: z.number(), fields: z.array(z.string()).optional(), @@ -106,9 +110,11 @@ export const listNamespacesResponseSchema = z.object({ cache_ttl_seconds: z.number(), expires_at_iso: z.string(), count: z.number(), + source_errors: z.record(z.string(), z.string()).optional(), namespaces: z.array( z.object({ name: z.string(), + source: z.string().optional(), record_count: z.number(), metadata_fields: z.record(z.string(), z.string()), }) @@ -123,6 +129,7 @@ export const namespaceRouterResponseSchema = z.object({ user_query: z.string(), suggestions: z.array(rankedNamespaceSchema), recommended_namespace: z.string().nullable(), + recommended_source: z.string().optional(), }); export type NamespaceRouterResponse = z.infer; @@ -130,6 +137,7 @@ export type NamespaceRouterResponse = z.infer; +export const listSourcesResponseSchema = z.object({ + status: z.literal('success'), + sources: z.array(z.string()), + default: z.string(), +}); + +export type ListSourcesResponse = z.infer; + /** * Assemble optional `experimental` block for query-shaped tool responses. * Omits the key entirely when no experimental fields are present. diff --git a/src/core/server/server-context.ts b/src/core/server/server-context.ts index c4ed4f7..7ef2dcb 100644 --- a/src/core/server/server-context.ts +++ b/src/core/server/server-context.ts @@ -10,13 +10,27 @@ import { warnLegacyFacade } from './legacy-facade-warn.js'; import { normalizeNamespace } from './namespace-utils.js'; import type { RecommendedTool } from './query-suggestion.js'; import type { UrlGenerationResult, UrlGeneratorFn } from './url-registry.js'; +import { buildSourceRegistry, type SourceRegistry } from './source-registry.js'; export type NamespaceInfo = { namespace: string; recordCount: number; metadata: Record; + source?: string; }; +export type ResolveSourceResult = + | { ok: true; source: string } + | { + ok: false; + code: + | 'UNKNOWN_SOURCE' + | 'AMBIGUOUS_NAMESPACE' + | 'NAMESPACE_NOT_FOUND' + | 'PARTIAL_SOURCE_AGGREGATION'; + message: string; + }; + /** Public seed shape for namespace cache injection (not the internal {@link NamespaceInfo} type). {@link ServerContext} copies `data` at construction so callers may reuse or mutate seed buffers afterward. */ export type NamespaceCacheSeed = { data: Array<{ namespace: string; recordCount: number; metadata: Record }>; @@ -43,6 +57,7 @@ export type ServerContextInitOptions = { /** Pre-built dependencies accepted by {@link ServerContext} and factory helpers. */ export interface ServerContextComposition { client?: PineconeClient; + sourceRegistry?: SourceRegistry; urlGenerators?: Iterable; namespaceCacheSeed?: NamespaceCacheSeed; suggestionFlowSeed?: SuggestionFlowSeedEntry[]; @@ -76,6 +91,24 @@ function buildPineconeClient(config: ServerConfigBase): PineconeClient { }); } +function flowKey(source: string | undefined, namespace: string, multiSource: boolean): string { + if (multiSource && source) { + return `${source}:${namespace}`; + } + return namespace; +} + +function urlRegistryKey( + namespace: string, + source: string | undefined, + multiSource: boolean +): string { + if (multiSource && source) { + return `${source}:${namespace.trim()}`; + } + return namespace.trim(); +} + /** * Encapsulates per-server state: Pinecone client, config, URL registry, * suggest-flow gate, and namespaces cache. @@ -87,6 +120,7 @@ export class ServerContext< private toolsRegistered = false; private client: PineconeClient | null = null; private clientExplicitlySet = false; + private sourceRegistry: SourceRegistry | null = null; private configValue: T | null = null; private readonly unconfiguredAlliance: boolean; private readonly urlGenerators = new Map(); @@ -95,9 +129,15 @@ export class ServerContext< constructor(config?: T, composition?: ServerContextComposition, init?: ServerContextInitOptions) { this.unconfiguredAlliance = init?.unconfiguredAlliance ?? false; + if (composition?.client && composition?.sourceRegistry) { + throw new Error('Cannot pass both client and sourceRegistry in ServerContextComposition.'); + } if (config) { this.configValue = config; } + if (composition?.sourceRegistry) { + this.sourceRegistry = composition.sourceRegistry; + } if (composition?.client) { this.client = composition.client; this.clientExplicitlySet = true; @@ -174,10 +214,153 @@ export class ServerContext< private invalidateConfigDerivedState(): void { this.client = null; this.clientExplicitlySet = false; + this.sourceRegistry = null; this.namespacesCache = null; this.suggestionFlow.clear(); } + private ensureSourceRegistry(): SourceRegistry { + if (this.sourceRegistry) { + return this.sourceRegistry; + } + const cfg = this.getConfig(); + if (!cfg.sources || cfg.sources.length === 0) { + throw new Error('Multi-source registry requested but config.sources is not set.'); + } + this.sourceRegistry = buildSourceRegistry({ + sources: cfg.sources, + defaultSource: cfg.defaultSource ?? cfg.sources[0]!.name, + cacheTtlMs: cfg.cacheTtlMs, + defaultTopK: cfg.defaultTopK, + requestTimeoutMs: cfg.requestTimeoutMs, + }); + return this.sourceRegistry; + } + + isMultiSource(): boolean { + const cfg = this.configValue; + if (cfg?.sources && cfg.sources.length > 1) { + return true; + } + if (this.sourceRegistry) { + return this.sourceRegistry.isMultiSource(); + } + return false; + } + + listSources(): string[] { + if (this.sourceRegistry) { + return this.sourceRegistry.listSources(); + } + const cfg = this.getConfigIfSet() ?? this.getConfig(); + if (cfg.sources && cfg.sources.length > 0) { + return cfg.sources.map((s) => s.name); + } + return []; + } + + getDefaultSourceName(): string { + if (this.sourceRegistry) { + return this.sourceRegistry.getDefaultName(); + } + const cfg = this.getConfig(); + return cfg.defaultSource ?? cfg.sources?.[0]?.name ?? 'default'; + } + + getClientForSource(source: string): PineconeClient { + if (this.sourceRegistry) { + return this.sourceRegistry.get(source); + } + if (!this.isMultiSource()) { + return this.getClient(); + } + return this.ensureSourceRegistry().get(source); + } + + /** Return the Pinecone client, lazily constructing from config when unset. */ + getClient(): PineconeClient { + if (this.sourceRegistry) { + return this.sourceRegistry.getDefault(); + } + if (this.getConfig().sources && this.getConfig().sources!.length > 0) { + return this.ensureSourceRegistry().getDefault(); + } + if (!this.client) { + this.client = buildPineconeClient(this.getConfig()); + } + return this.client; + } + + async resolveSource(source?: string, namespace?: string): Promise { + if ( + !this.isMultiSource() && + !(this.getConfig().sources && this.getConfig().sources!.length > 0) + ) { + return { ok: true, source: this.getDefaultSourceName() }; + } + const registry = this.sourceRegistry ?? this.ensureSourceRegistry(); + const trimmedSource = source?.trim(); + if (trimmedSource) { + if (!registry.listSources().includes(trimmedSource)) { + return { + ok: false, + code: 'UNKNOWN_SOURCE', + message: `Unknown source "${trimmedSource}". Call list_sources for configured names.`, + }; + } + if (namespace !== undefined) { + const nsNorm = normalizeNamespace(namespace); + if (!nsNorm) { + return { ok: false, code: 'NAMESPACE_NOT_FOUND', message: 'namespace cannot be empty.' }; + } + const { data } = await registry.getNamespacesWithCache(trimmedSource); + const found = data.some((n) => normalizeNamespace(n.namespace) === nsNorm); + if (!found) { + return { + ok: false, + code: 'NAMESPACE_NOT_FOUND', + message: `Namespace "${namespace}" not found on source "${trimmedSource}".`, + }; + } + } + return { ok: true, source: trimmedSource }; + } + if (namespace !== undefined) { + const nsNorm = normalizeNamespace(namespace); + if (!nsNorm) { + return { ok: false, code: 'NAMESPACE_NOT_FOUND', message: 'namespace cannot be empty.' }; + } + const cacheResult = await this.getNamespacesWithCache(); + const sourceErrors = cacheResult.source_errors; + if (sourceErrors && Object.keys(sourceErrors).length > 0) { + return { + ok: false, + code: 'PARTIAL_SOURCE_AGGREGATION', + message: + 'Namespace discovery is incomplete because one or more sources failed. Pass source explicitly or retry after resolving source_errors.', + }; + } + const { data } = cacheResult; + const matches = data.filter((n) => normalizeNamespace(n.namespace) === nsNorm); + if (matches.length === 1) { + return { ok: true, source: matches[0]!.source ?? registry.getDefaultName() }; + } + if (matches.length > 1) { + return { + ok: false, + code: 'AMBIGUOUS_NAMESPACE', + message: `Namespace "${namespace}" exists on multiple sources. Pass source explicitly.`, + }; + } + return { + ok: false, + code: 'NAMESPACE_NOT_FOUND', + message: `Namespace "${namespace}" not found. Call list_namespaces first.`, + }; + } + return { ok: true, source: registry.getDefaultName() }; + } + setClient(client: PineconeClient): void { this.client = client; this.clientExplicitlySet = true; @@ -201,19 +384,21 @@ export class ServerContext< return this.client; } - /** Return the Pinecone client, lazily constructing from config when unset. */ - getClient(): PineconeClient { - if (!this.client) { - this.client = buildPineconeClient(this.getConfig()); + async checkAllIndexes(): Promise<{ ok: boolean; errors: string[] }> { + if (this.sourceRegistry) { + return this.sourceRegistry.checkAllIndexes(); } - return this.client; + if (this.getConfig().sources && this.getConfig().sources!.length > 0) { + return this.ensureSourceRegistry().checkAllIndexes(); + } + return this.getClient().checkIndexes(); } resetUrlGenerators(): void { this.urlGenerators.clear(); } - registerUrlGenerator(namespace: string, generator: UrlGeneratorFn): void { + registerUrlGenerator(namespace: string, generator: UrlGeneratorFn, source?: string): void { const normalizedNamespace = namespace.trim(); if (normalizedNamespace.length === 0) { throw new TypeError('namespace must be a non-empty string'); @@ -221,27 +406,83 @@ export class ServerContext< if (typeof generator !== 'function') { throw new TypeError('generator must be a function'); } + const multi = this.isMultiSource(); + if (multi) { + const sources = + this.sourceRegistry?.listSources() ?? + this.configValue?.sources?.map((entry) => entry.name) ?? + []; + if (source) { + this.urlGenerators.set(urlRegistryKey(normalizedNamespace, source, true), generator); + } else { + for (const src of sources) { + this.urlGenerators.set(urlRegistryKey(normalizedNamespace, src, true), generator); + } + } + return; + } this.urlGenerators.set(normalizedNamespace, generator); } - unregisterUrlGenerator(namespace: string): boolean { - return this.urlGenerators.delete(namespace.trim()); + unregisterUrlGenerator(namespace: string, source?: string): boolean { + const trimmed = namespace.trim(); + if (!this.isMultiSource()) { + return this.urlGenerators.delete(trimmed); + } + const sources = + this.sourceRegistry?.listSources() ?? + this.configValue?.sources?.map((entry) => entry.name) ?? + []; + if (source) { + return this.urlGenerators.delete(urlRegistryKey(trimmed, source, true)); + } + let removed = false; + for (const src of sources) { + if (this.urlGenerators.delete(urlRegistryKey(trimmed, src, true))) { + removed = true; + } + } + return removed || this.urlGenerators.delete(trimmed); } - hasUrlGenerator(namespace: string): boolean { - return this.urlGenerators.has(namespace.trim()); + hasUrlGenerator(namespace: string, source?: string): boolean { + const trimmed = namespace.trim(); + if (!this.isMultiSource()) { + return this.urlGenerators.has(trimmed); + } + if (source) { + return ( + this.urlGenerators.has(urlRegistryKey(trimmed, source, true)) || + this.urlGenerators.has(trimmed) + ); + } + const sources = + this.sourceRegistry?.listSources() ?? + this.configValue?.sources?.map((entry) => entry.name) ?? + []; + return ( + sources.some((src) => this.urlGenerators.has(urlRegistryKey(trimmed, src, true))) || + this.urlGenerators.has(trimmed) + ); } generateUrlForNamespace( namespace: string, - metadata: Record + metadata: Record, + source?: string ): UrlGenerationResult { const existingUrl = asString(metadata['url']); if (existingUrl) { return { url: existingUrl, method: 'metadata.url' }; } - const generator = this.urlGenerators.get(namespace.trim()); + const multi = this.isMultiSource(); + const trimmed = namespace.trim(); + const key = urlRegistryKey(trimmed, source, multi); + let generator = this.urlGenerators.get(key); + if (!generator && multi && source) { + generator = this.urlGenerators.get(trimmed); + } if (generator) { return generator(metadata); } @@ -263,19 +504,23 @@ export class ServerContext< } } - markSuggested(namespace: string, state: Omit): void { + markSuggested(namespace: string, state: Omit, source?: string): void { const key = normalizeNamespace(namespace); if (!key) { throw new Error('markSuggested: namespace must not be empty after trim'); } this.sweepExpiredSuggestionFlow(); - this.suggestionFlow.set(key, { + const flowKeyValue = flowKey(source, key, this.isMultiSource()); + this.suggestionFlow.set(flowKeyValue, { ...state, updatedAt: Date.now(), }); } - requireSuggested(namespace: string): + requireSuggested( + namespace: string, + source?: string + ): | { ok: true; flow: FlowState; @@ -304,7 +549,8 @@ export class ServerContext< }; } - const state = this.suggestionFlow.get(key); + const flowKeyValue = flowKey(source, key, this.isMultiSource()); + const state = this.suggestionFlow.get(flowKeyValue); if (!state) { return { ok: false, @@ -316,7 +562,7 @@ export class ServerContext< const cfg = this.getConfig(); const now = Date.now(); if (now - state.updatedAt > cfg.cacheTtlMs) { - this.suggestionFlow.delete(key); + this.suggestionFlow.delete(flowKeyValue); return { ok: false, message: @@ -331,11 +577,25 @@ export class ServerContext< this.suggestionFlow.clear(); } - async getNamespacesWithCache(): Promise<{ + async getNamespacesWithCache(source?: string): Promise<{ data: NamespaceInfo[]; cache_hit: boolean; expires_at: number; + source_errors?: Record; }> { + if ( + this.isMultiSource() || + (this.getConfig().sources && this.getConfig().sources!.length > 0) + ) { + const registry = this.sourceRegistry ?? this.ensureSourceRegistry(); + if (source) { + const result = await registry.getNamespacesWithCache(source); + return result; + } + const aggregated = await registry.getAllNamespacesWithCache(); + return aggregated; + } + const now = Date.now(); if (this.namespacesCache && now < this.namespacesCache.expiresAt) { return { @@ -353,7 +613,27 @@ export class ServerContext< return { data, cache_hit: false, expires_at: expiresAt }; } - invalidateNamespacesCache(): void { + async getNamespacesWithCacheForSource(source: string): Promise<{ + data: NamespaceInfo[]; + cache_hit: boolean; + expires_at: number; + }> { + return this.getNamespacesWithCache(source) as Promise<{ + data: NamespaceInfo[]; + cache_hit: boolean; + expires_at: number; + }>; + } + + invalidateNamespacesCache(source?: string): void { + if (this.sourceRegistry) { + this.sourceRegistry.invalidateNamespacesCache(source); + return; + } + if (this.isMultiSource()) { + this.ensureSourceRegistry().invalidateNamespacesCache(source); + return; + } this.namespacesCache = null; } @@ -385,6 +665,7 @@ export class ServerContext< this.toolsRegistered = false; this.client = null; this.clientExplicitlySet = false; + this.sourceRegistry = null; this.configValue = null; this.urlGenerators.clear(); this.suggestionFlow.clear(); diff --git a/src/core/server/source-config.test.ts b/src/core/server/source-config.test.ts new file mode 100644 index 0000000..fe79cbc --- /dev/null +++ b/src/core/server/source-config.test.ts @@ -0,0 +1,196 @@ +import { mkdtempSync, writeFileSync } from 'node:fs'; +import { join } from 'node:path'; +import { tmpdir } from 'node:os'; +import { describe, expect, it, vi } from 'vitest'; +import { + parseInlineSources, + parseSourcesConfigFile, + resolveEnvIndirection, + resolveSourceDefinitions, +} from './source-config.js'; +import { resolveConfig } from '../config.js'; + +describe('source-config', () => { + it('resolves env indirection', () => { + const env = { PINECONE_API_KEY_1: 'key-one' }; + expect(resolveEnvIndirection('${PINECONE_API_KEY_1}', env)).toBe('key-one'); + }); + + it('parses inline sources', () => { + const env = { + K1: 'api-1', + K2: 'api-2', + }; + const sources = parseInlineSources( + 'api_key_1:${K1}:index_name_1;api_key_2:${K2}:index_name_2', + env + ); + expect(sources).toHaveLength(2); + expect(sources[0]).toMatchObject({ + name: 'api_key_1', + apiKey: 'api-1', + indexName: 'index_name_1', + sparseIndexName: 'index_name_1-sparse', + }); + expect(sources[1]?.name).toBe('api_key_2'); + }); + + it('throws on duplicate source name in inline string', () => { + expect(() => parseInlineSources('api_key_1:sk-a:idx-a;api_key_1:sk-b:idx-b', {})).toThrow( + /Duplicate source name "api_key_1"/ + ); + }); + + it('parseSourcesConfigFile resolves env indirection and defaultSource', () => { + const dir = mkdtempSync(join(tmpdir(), 'pinecone-sources-')); + const filePath = join(dir, 'sources.json'); + writeFileSync( + filePath, + JSON.stringify({ + defaultSource: 'api_key_2', + sources: { + api_key_1: { apiKey: '${K1}', indexName: 'index_name_1' }, + api_key_2: { apiKey: '${K2}', indexName: 'index_name_2' }, + }, + }) + ); + const env = { K1: 'api-1', K2: 'api-2' }; + const parsed = parseSourcesConfigFile(filePath, env); + expect(parsed.defaultSource).toBe('api_key_2'); + expect(parsed.sources).toHaveLength(2); + const first = parsed.sources.find((s) => s.name === 'api_key_1'); + expect(first).toMatchObject({ + apiKey: 'api-1', + indexName: 'index_name_1', + sparseIndexName: 'index_name_1-sparse', + }); + }); + + it('parseSourcesConfigFile treats null sparseIndexName and rerankModel as omitted', () => { + const dir = mkdtempSync(join(tmpdir(), 'pinecone-sources-')); + const filePath = join(dir, 'null-sparse.json'); + writeFileSync( + filePath, + JSON.stringify({ + defaultSource: 'api_key_1', + sources: { + api_key_1: { + apiKey: 'k', + indexName: 'index_name_1', + sparseIndexName: null, + rerankModel: null, + }, + }, + }) + ); + const parsed = parseSourcesConfigFile(filePath, {}); + const first = parsed.sources.find((s) => s.name === 'api_key_1'); + expect(first?.sparseIndexName).toBe('index_name_1-sparse'); + expect(first?.rerankModel).toBeUndefined(); + }); + + it('throws when defaultSource is not a configured source name', () => { + const dir = mkdtempSync(join(tmpdir(), 'pinecone-sources-')); + const filePath = join(dir, 'bad-default.json'); + writeFileSync( + filePath, + JSON.stringify({ + defaultSource: 'missing', + sources: { + api_key_1: { apiKey: 'k', indexName: 'idx' }, + }, + }) + ); + expect(() => parseSourcesConfigFile(filePath, {})).toThrow(/defaultSource/); + }); + + it('resolveConfig uses PINECONE_SOURCES when set', () => { + vi.stubEnv('PINECONE_SOURCES', 'api_key_1:sk-test:my-index'); + vi.stubEnv('PINECONE_API_KEY', 'ignored'); + try { + const cfg = resolveConfig({}); + expect(cfg.sources).toHaveLength(1); + expect(cfg.sources?.[0]?.name).toBe('api_key_1'); + expect(cfg.apiKey).toBe('sk-test'); + } finally { + vi.unstubAllEnvs(); + } + }); + + it('resolveConfig prefers PINECONE_SOURCES over PINECONE_API_KEY', () => { + vi.stubEnv('PINECONE_SOURCES', 'api_key_1:from-sources:my-index'); + vi.stubEnv('PINECONE_API_KEY', 'from-single-key'); + vi.stubEnv('PINECONE_INDEX_NAME', 'single-index'); + try { + const cfg = resolveConfig({}); + expect(cfg.sources).toHaveLength(1); + expect(cfg.apiKey).toBe('from-sources'); + expect(cfg.indexName).toBe('my-index'); + } finally { + vi.unstubAllEnvs(); + } + }); + + it('resolveConfig uses defaultSource for top-level index fields', () => { + const dir = mkdtempSync(join(tmpdir(), 'pinecone-sources-')); + const filePath = join(dir, 'sources.json'); + writeFileSync( + filePath, + JSON.stringify({ + defaultSource: 'api_key_2', + sources: { + api_key_1: { apiKey: 'k1', indexName: 'index_name_1' }, + api_key_2: { apiKey: 'k2', indexName: 'index_name_2' }, + }, + }) + ); + vi.stubEnv('PINECONE_CONFIG_FILE', filePath); + try { + const cfg = resolveConfig({}); + expect(cfg.defaultSource).toBe('api_key_2'); + expect(cfg.indexName).toBe('index_name_2'); + expect(cfg.apiKey).toBe('k2'); + } finally { + vi.unstubAllEnvs(); + } + }); + + it('resolveSourceDefinitions prefers config file over inline env', () => { + const dir = mkdtempSync(join(tmpdir(), 'pinecone-sources-')); + const filePath = join(dir, 'sources.json'); + writeFileSync( + filePath, + JSON.stringify({ + defaultSource: 'from_file', + sources: { + from_file: { apiKey: 'file-key', indexName: 'file-index' }, + }, + }) + ); + const env = { + PINECONE_SOURCES: 'from_inline:sk:inline-index', + PINECONE_CONFIG_FILE: filePath, + }; + const parsed = resolveSourceDefinitions({}, env); + expect(parsed?.defaultSource).toBe('from_file'); + expect(parsed?.sources[0]?.indexName).toBe('file-index'); + }); + + it('resolveSourceDefinitions prefers overrides.configFile over env PINECONE_SOURCES', () => { + const dir = mkdtempSync(join(tmpdir(), 'pinecone-sources-')); + const filePath = join(dir, 'sources.json'); + writeFileSync( + filePath, + JSON.stringify({ + defaultSource: 'from_override', + sources: { + from_override: { apiKey: 'override-key', indexName: 'override-index' }, + }, + }) + ); + const env = { PINECONE_SOURCES: 'from_env:sk:env-index' }; + const parsed = resolveSourceDefinitions({ configFile: filePath }, env); + expect(parsed?.defaultSource).toBe('from_override'); + expect(parsed?.sources[0]?.indexName).toBe('override-index'); + }); +}); diff --git a/src/core/server/source-config.ts b/src/core/server/source-config.ts new file mode 100644 index 0000000..5a64c1a --- /dev/null +++ b/src/core/server/source-config.ts @@ -0,0 +1,213 @@ +/** + * Multi-source Pinecone configuration parsing (inline PINECONE_SOURCES and JSON config file). + */ + +import { readFileSync } from 'node:fs'; +import { resolve } from 'node:path'; +import { trimOptional } from '../config.js'; + +/** Named Pinecone project connection (one API key + index pair). */ +export interface SourceDefinition { + name: string; + apiKey: string; + indexName: string; + sparseIndexName?: string; + rerankModel?: string; +} + +export type ParseSourcesOptions = { + /** Apply Alliance defaults for indexName / rerankModel when omitted per entry. */ + allianceDefaults?: { + indexName?: string; + rerankModel?: string; + }; +}; + +const SOURCE_NAME_PATTERN = /^[a-zA-Z0-9_-]+$/; + +/** Resolve `${ENV_VAR}` references in a string value. */ +export function resolveEnvIndirection(value: string, env: NodeJS.ProcessEnv): string { + const trimmed = value.trim(); + const match = /^\$\{([A-Za-z_][A-Za-z0-9_]*)\}$/.exec(trimmed); + if (!match) { + return trimmed; + } + const envKey = match[1]!; + const resolved = env[envKey]?.trim(); + if (!resolved) { + throw new Error(`Environment variable ${envKey} is not set (referenced as ${trimmed}).`); + } + return resolved; +} + +function validateSourceName(name: string): void { + if (!name || !SOURCE_NAME_PATTERN.test(name)) { + throw new Error( + `Invalid source name "${name}": use alphanumeric characters, hyphens, or underscores only.` + ); + } +} + +function normalizeSourceEntry( + name: string, + raw: { + apiKey: string; + indexName: string; + sparseIndexName?: string; + rerankModel?: string; + }, + env: NodeJS.ProcessEnv, + allianceDefaults?: ParseSourcesOptions['allianceDefaults'] +): SourceDefinition { + validateSourceName(name); + const apiKey = resolveEnvIndirection(raw.apiKey, env); + if (!apiKey) { + throw new Error(`Source "${name}": apiKey is required.`); + } + const indexName = + trimOptional(resolveEnvIndirection(raw.indexName, env)) ?? + trimOptional(allianceDefaults?.indexName); + if (!indexName) { + throw new Error(`Source "${name}": indexName is required.`); + } + const sparseRaw = raw.sparseIndexName + ? resolveEnvIndirection(raw.sparseIndexName, env) + : undefined; + const sparseIndexName = trimOptional(sparseRaw) ?? `${indexName}-sparse`; + const rerankRaw = raw.rerankModel ? resolveEnvIndirection(raw.rerankModel, env) : undefined; + const rerankModel = trimOptional(rerankRaw) ?? trimOptional(allianceDefaults?.rerankModel); + return { + name, + apiKey, + indexName, + sparseIndexName, + ...(rerankModel !== undefined ? { rerankModel } : {}), + }; +} + +/** Parse inline `name:apiKey:indexName[;name2:...]` format. */ +export function parseInlineSources( + inline: string, + env: NodeJS.ProcessEnv = process.env, + options?: ParseSourcesOptions +): SourceDefinition[] { + const segments = inline + .split(';') + .map((s) => s.trim()) + .filter(Boolean); + if (segments.length === 0) { + throw new Error('PINECONE_SOURCES is empty.'); + } + const sources: SourceDefinition[] = []; + const seen = new Set(); + for (const segment of segments) { + const parts = segment.split(':'); + if (parts.length < 3) { + throw new Error( + `Invalid PINECONE_SOURCES segment "${segment}": expected name:apiKey:indexName (optional fields after index not supported in inline format).` + ); + } + const name = parts[0]?.trim() ?? ''; + const apiKey = parts.slice(1, -1).join(':').trim(); + const indexName = parts[parts.length - 1]?.trim() ?? ''; + if (!name || !apiKey || !indexName) { + throw new Error( + `Invalid PINECONE_SOURCES segment "${segment}": name, apiKey, and indexName are required.` + ); + } + if (seen.has(name)) { + throw new Error(`Duplicate source name "${name}" in PINECONE_SOURCES.`); + } + seen.add(name); + sources.push(normalizeSourceEntry(name, { apiKey, indexName }, env, options?.allianceDefaults)); + } + return sources; +} + +type JsonSourceFile = { + defaultSource?: string; + sources: Record< + string, + { + apiKey: string; + indexName: string; + sparseIndexName?: string; + rerankModel?: string; + } + >; +}; + +/** Parse JSON config file for multi-source setup. */ +export function parseSourcesConfigFile( + filePath: string, + env: NodeJS.ProcessEnv = process.env, + options?: ParseSourcesOptions +): { sources: SourceDefinition[]; defaultSource: string } { + const absolute = resolve(filePath); + let parsed: JsonSourceFile; + try { + const raw = readFileSync(absolute, 'utf8'); + parsed = JSON.parse(raw) as JsonSourceFile; + } catch (err) { + const message = err instanceof Error ? err.message : String(err); + throw new Error(`Failed to read PINECONE config file "${filePath}": ${message}`); + } + if (!parsed.sources || typeof parsed.sources !== 'object') { + throw new Error(`PINECONE config file "${filePath}" must contain a "sources" object.`); + } + const entries = Object.entries(parsed.sources); + if (entries.length === 0) { + throw new Error(`PINECONE config file "${filePath}" has no sources.`); + } + const sources: SourceDefinition[] = []; + const seen = new Set(); + for (const [name, cfg] of entries) { + if (seen.has(name)) { + throw new Error(`Duplicate source name "${name}" in config file.`); + } + seen.add(name); + if (!cfg || typeof cfg !== 'object') { + throw new Error(`Source "${name}" in config file must be an object.`); + } + sources.push( + normalizeSourceEntry( + name, + { + apiKey: String(cfg.apiKey ?? ''), + indexName: String(cfg.indexName ?? ''), + ...(cfg.sparseIndexName != null ? { sparseIndexName: String(cfg.sparseIndexName) } : {}), + ...(cfg.rerankModel != null ? { rerankModel: String(cfg.rerankModel) } : {}), + }, + env, + options?.allianceDefaults + ) + ); + } + const defaultSource = trimOptional(parsed.defaultSource) ?? sources[0]?.name; + if (!defaultSource || !seen.has(defaultSource)) { + throw new Error( + `defaultSource "${parsed.defaultSource ?? ''}" is not a configured source name.` + ); + } + return { sources, defaultSource }; +} + +/** Resolve sources from overrides/env/file; config file wins over inline when both are set. */ +export function resolveSourceDefinitions( + overrides: { sources?: string; configFile?: string }, + env: NodeJS.ProcessEnv = process.env, + options?: ParseSourcesOptions +): { sources: SourceDefinition[]; defaultSource: string } | null { + const inline = trimOptional(overrides.sources) ?? trimOptional(env['PINECONE_SOURCES']); + const configFile = + trimOptional(overrides.configFile) ?? trimOptional(env['PINECONE_CONFIG_FILE']); + + if (configFile) { + return parseSourcesConfigFile(configFile, env, options); + } + if (inline) { + const sources = parseInlineSources(inline, env, options); + return { sources, defaultSource: sources[0]!.name }; + } + return null; +} diff --git a/src/core/server/source-registry.test.ts b/src/core/server/source-registry.test.ts new file mode 100644 index 0000000..c9ab679 --- /dev/null +++ b/src/core/server/source-registry.test.ts @@ -0,0 +1,85 @@ +import { describe, expect, it, vi } from 'vitest'; +import { buildSourceRegistry } from './source-registry.js'; +import type { SourceDefinition } from './source-config.js'; + +const sources: SourceDefinition[] = [ + { name: 'api_key_1', apiKey: 'k1', indexName: 'idx-a', sparseIndexName: 'idx-a-sparse' }, + { name: 'api_key_2', apiKey: 'k2', indexName: 'idx-b', sparseIndexName: 'idx-b-sparse' }, +]; + +function mockClient(name: string) { + return { + listNamespacesWithMetadata: vi + .fn() + .mockResolvedValue([{ namespace: 'wg21', recordCount: 10, metadata: { title: 'string' } }]), + checkIndexes: vi.fn().mockResolvedValue({ ok: true, errors: [] }), + getSparseIndexName: () => `${name}-sparse`, + }; +} + +describe('SourceRegistry', () => { + it('aggregates namespaces from all sources', async () => { + const clients = new Map([ + ['api_key_1', mockClient('api_key_1') as never], + ['api_key_2', mockClient('api_key_2') as never], + ]); + const registry = buildSourceRegistry({ + sources, + defaultSource: 'api_key_1', + cacheTtlMs: 60_000, + defaultTopK: 10, + requestTimeoutMs: 15_000, + clients, + }); + const result = await registry.getAllNamespacesWithCache(); + expect(result.data).toHaveLength(2); + expect(result.data.map((n) => n.source).sort()).toEqual(['api_key_1', 'api_key_2']); + }); + + it('isolates per-source namespace caches', async () => { + const client1 = mockClient('api_key_1'); + const client2 = mockClient('api_key_2'); + const clients = new Map([ + ['api_key_1', client1 as never], + ['api_key_2', client2 as never], + ]); + const registry = buildSourceRegistry({ + sources, + defaultSource: 'api_key_1', + cacheTtlMs: 60_000, + defaultTopK: 10, + requestTimeoutMs: 15_000, + clients, + }); + await registry.getNamespacesWithCache('api_key_1'); + await registry.getNamespacesWithCache('api_key_2'); + expect(client1.listNamespacesWithMetadata).toHaveBeenCalledTimes(1); + expect(client2.listNamespacesWithMetadata).toHaveBeenCalledTimes(1); + }); + + it('returns partial namespaces and source_errors when one source fails', async () => { + const client1 = mockClient('api_key_1'); + const client2 = { + listNamespacesWithMetadata: vi.fn().mockRejectedValue(new Error('api_key_2 unreachable')), + checkIndexes: vi.fn().mockResolvedValue({ ok: true, errors: [] }), + getSparseIndexName: () => 'api_key_2-sparse', + }; + const clients = new Map([ + ['api_key_1', client1 as never], + ['api_key_2', client2 as never], + ]); + const registry = buildSourceRegistry({ + sources, + defaultSource: 'api_key_1', + cacheTtlMs: 60_000, + defaultTopK: 10, + requestTimeoutMs: 15_000, + clients, + }); + const result = await registry.getAllNamespacesWithCache(); + expect(result.data).toHaveLength(1); + expect(result.data[0]?.source).toBe('api_key_1'); + expect(result.source_errors).toEqual({ api_key_2: 'api_key_2 unreachable' }); + expect(result.cache_hit).toBe(false); + }); +}); diff --git a/src/core/server/source-registry.ts b/src/core/server/source-registry.ts new file mode 100644 index 0000000..a45a487 --- /dev/null +++ b/src/core/server/source-registry.ts @@ -0,0 +1,189 @@ +/** + * Registry of named Pinecone sources (one client + namespace cache per source). + */ + +import { PineconeClient } from '../pinecone-client.js'; +import type { NamespaceInfo } from './server-context.js'; +import type { SourceDefinition } from './source-config.js'; + +type CacheEntry = { + data: NamespaceInfo[]; + expiresAt: number; +}; + +export type PerSourceCacheResult = { + data: NamespaceInfo[]; + cache_hit: boolean; + expires_at: number; +}; + +export type AggregatedCacheResult = { + data: NamespaceInfo[]; + cache_hit: boolean; + expires_at: number; + source_errors?: Record; +}; + +export type BuildSourceRegistryOptions = { + sources: SourceDefinition[]; + defaultSource: string; + cacheTtlMs: number; + defaultTopK: number; + requestTimeoutMs: number; + clients?: Map; +}; + +export class SourceRegistry { + private readonly entries: Map< + string, + { client: PineconeClient; cache: CacheEntry | null; definition: SourceDefinition } + >; + private readonly defaultSourceName: string; + private readonly cacheTtlMs: number; + + constructor(options: BuildSourceRegistryOptions) { + this.cacheTtlMs = options.cacheTtlMs; + this.defaultSourceName = options.defaultSource; + this.entries = new Map(); + for (const def of options.sources) { + const client = + options.clients?.get(def.name) ?? + new PineconeClient({ + apiKey: def.apiKey, + indexName: def.indexName, + sparseIndexName: def.sparseIndexName, + rerankModel: def.rerankModel, + defaultTopK: options.defaultTopK, + requestTimeoutMs: options.requestTimeoutMs, + }); + this.entries.set(def.name, { client, cache: null, definition: def }); + } + if (!this.entries.has(this.defaultSourceName)) { + throw new Error(`Default source "${this.defaultSourceName}" is not configured.`); + } + } + + isMultiSource(): boolean { + return this.entries.size > 1; + } + + listSources(): string[] { + return [...this.entries.keys()]; + } + + getDefaultName(): string { + return this.defaultSourceName; + } + + getDefinition(name: string): SourceDefinition { + const entry = this.entries.get(name); + if (!entry) { + throw new Error(`Unknown Pinecone source "${name}".`); + } + return entry.definition; + } + + get(name: string): PineconeClient { + const entry = this.entries.get(name); + if (!entry) { + throw new Error(`Unknown Pinecone source "${name}".`); + } + return entry.client; + } + + getDefault(): PineconeClient { + return this.get(this.defaultSourceName); + } + + async getNamespacesWithCache(source: string): Promise { + const entry = this.entries.get(source); + if (!entry) { + throw new Error(`Unknown Pinecone source "${source}".`); + } + const now = Date.now(); + if (entry.cache && now < entry.cache.expiresAt) { + return { + data: entry.cache.data, + cache_hit: true, + expires_at: entry.cache.expiresAt, + }; + } + const raw = await entry.client.listNamespacesWithMetadata(); + const data: NamespaceInfo[] = raw.map((ns) => ({ + namespace: ns.namespace, + recordCount: ns.recordCount, + metadata: ns.metadata, + source, + })); + const expiresAt = now + this.cacheTtlMs; + entry.cache = { data, expiresAt }; + return { data, cache_hit: false, expires_at: expiresAt }; + } + + async getAllNamespacesWithCache(): Promise { + const names = this.listSources(); + const settled = await Promise.allSettled( + names.map(async (name) => { + const result = await this.getNamespacesWithCache(name); + return { name, result }; + }) + ); + const data: NamespaceInfo[] = []; + const source_errors: Record = {}; + let cache_hit = true; + let maxExpires = 0; + for (let i = 0; i < settled.length; i++) { + const outcome = settled[i]!; + const name = names[i]!; + if (outcome.status === 'fulfilled') { + data.push(...outcome.value.result.data); + if (!outcome.value.result.cache_hit) { + cache_hit = false; + } + maxExpires = Math.max(maxExpires, outcome.value.result.expires_at); + } else { + cache_hit = false; + const message = + outcome.reason instanceof Error ? outcome.reason.message : String(outcome.reason); + source_errors[name] = message; + } + } + const expires_at = maxExpires > 0 ? maxExpires : Date.now() + this.cacheTtlMs; + return { + data, + cache_hit, + expires_at, + ...(Object.keys(source_errors).length > 0 ? { source_errors } : {}), + }; + } + + invalidateNamespacesCache(source?: string): void { + if (source !== undefined) { + const entry = this.entries.get(source); + if (entry) { + entry.cache = null; + } + return; + } + for (const entry of this.entries.values()) { + entry.cache = null; + } + } + + async checkAllIndexes(): Promise<{ ok: boolean; errors: string[] }> { + const errors: string[] = []; + for (const name of this.listSources()) { + const result = await this.get(name).checkIndexes(); + if (!result.ok) { + for (const err of result.errors) { + errors.push(`[${name}] ${err}`); + } + } + } + return { ok: errors.length === 0, errors }; + } +} + +export function buildSourceRegistry(options: BuildSourceRegistryOptions): SourceRegistry { + return new SourceRegistry(options); +} diff --git a/src/core/server/source-tool-utils.ts b/src/core/server/source-tool-utils.ts new file mode 100644 index 0000000..505eb7e --- /dev/null +++ b/src/core/server/source-tool-utils.ts @@ -0,0 +1,112 @@ +/** + * Shared helpers for multi-source tool handlers. + */ + +import { z } from 'zod'; +import { getPineconeClient } from './client-context.js'; +import type { PineconeClient } from '../pinecone-client.js'; +import type { ServerContext } from './server-context.js'; +import { logToolInvocation, validationToolError, type ToolError } from './tool-error.js'; + +export const sourceParamSchema = z + .string() + .optional() + .describe( + 'Pinecone source name (from list_sources). Omit on discovery tools to search all sources. ' + + 'On query tools, omit only when namespace uniquely identifies one source.' + ); + +export type ResolveSourceFailureCode = + | 'UNKNOWN_SOURCE' + | 'AMBIGUOUS_NAMESPACE' + | 'NAMESPACE_NOT_FOUND' + | 'PARTIAL_SOURCE_AGGREGATION'; + +function validationFieldForSourceCode(code: ResolveSourceFailureCode): 'source' | 'namespace' { + switch (code) { + case 'NAMESPACE_NOT_FOUND': + case 'AMBIGUOUS_NAMESPACE': + return 'namespace'; + case 'UNKNOWN_SOURCE': + case 'PARTIAL_SOURCE_AGGREGATION': + return 'source'; + } +} + +export async function resolveSourceForTool( + ctx: ServerContext | undefined, + source: string | undefined, + namespace: string | undefined +): Promise< + | { ok: true; source: string; ctx: ServerContext } + | { ok: false; code: ResolveSourceFailureCode; message: string } +> { + if (!ctx) { + return { + ok: false, + code: 'UNKNOWN_SOURCE', + message: 'ServerContext is required for multi-source resolution.', + }; + } + const resolved = await ctx.resolveSource(source, namespace); + if (!resolved.ok) { + return resolved; + } + return { ok: true, source: resolved.source, ctx }; +} + +export function sourceValidationError(code: ResolveSourceFailureCode, message: string) { + const field = validationFieldForSourceCode(code); + const suggestion = + code === 'AMBIGUOUS_NAMESPACE' + ? 'Call list_namespaces and pass source explicitly when the namespace exists on multiple projects.' + : code === 'UNKNOWN_SOURCE' + ? 'Call list_sources to see configured source names.' + : code === 'PARTIAL_SOURCE_AGGREGATION' + ? 'Pass source explicitly, or fix failing sources listed in list_namespaces source_errors and retry.' + : 'Use list_namespaces to discover valid namespace and source pairs.'; + return validationToolError(message, field, { suggestion }); +} + +/** Reject `source` when legacy tool registration has no ServerContext. */ +export function rejectSourceWithoutContext( + source: string | undefined, + ctx: ServerContext | undefined +): ToolError | null { + if (source && !ctx) { + return validationToolError( + 'source parameter requires ServerContext (multi-source mode).', + 'source' + ); + } + return null; +} + +/** Pinecone client for a resolved source, or the context/default client in single-source mode. */ +export function getClientForResolvedSource( + ctx: ServerContext | undefined, + source: string | undefined, + toolName?: string +): PineconeClient { + if (toolName && source && ctx?.isMultiSource()) { + logToolInvocation(toolName, source); + } + if (ctx) { + if (ctx.isMultiSource() && source) { + return ctx.getClientForSource(source); + } + return ctx.getClient(); + } + return getPineconeClient(); +} + +/** Include `source` on responses only in multi-source mode. */ +export function optionalSourceField( + ctx: ServerContext | undefined, + source: string | undefined +): { source?: string } { + if (source && ctx?.isMultiSource()) { + return { source }; + } + return {}; +} diff --git a/src/core/server/tool-error.ts b/src/core/server/tool-error.ts index b241071..0d3f034 100644 --- a/src/core/server/tool-error.ts +++ b/src/core/server/tool-error.ts @@ -4,7 +4,7 @@ */ import { z } from 'zod'; -import { getLogLevel, error as logError } from '../../logger.js'; +import { getLogLevel, error as logError, info } from '../../logger.js'; /** User-facing error message: detailed in DEBUG, generic otherwise. */ export function getToolErrorMessage(error: unknown, fallbackMessage: string): string { @@ -13,8 +13,15 @@ export function getToolErrorMessage(error: unknown, fallbackMessage: string): st } /** Log tool failure to stderr via the level-based logger. */ -export function logToolError(toolName: string, error: unknown): void { - logError(`Error in ${toolName} tool`, error); +export function logToolError(toolName: string, error: unknown, source?: string): void { + const suffix = source ? ` [source=${source}]` : ''; + logError(`Error in ${toolName} tool${suffix}`, error); +} + +/** Log resolved source at INFO when multi-source routing is active. */ +export function logToolInvocation(toolName: string, source: string | undefined): void { + if (!source) return; + info(`${toolName} [source=${source}]`); } export const toolErrorCodeSchema = z.enum([ diff --git a/src/core/server/tools/count-tool.test.ts b/src/core/server/tools/count-tool.test.ts index 34c6d23..67df681 100644 --- a/src/core/server/tools/count-tool.test.ts +++ b/src/core/server/tools/count-tool.test.ts @@ -46,6 +46,21 @@ describe('count tool handler', () => { expect(err.field).toBe('namespace'); }); + it('returns VALIDATION when source is provided without ServerContext', async () => { + const server = createMockServer(); + registerCountTool(server as never); + const err = assertToolErrorCode( + await server.getHandler('count')!({ + namespace: 'wg21', + query_text: 'doc', + source: 'api_key_1', + }), + 'VALIDATION' + ); + expect(err.field).toBe('source'); + expect(mockedGetClient().count).not.toHaveBeenCalled(); + }); + it('returns VALIDATION when query_text is empty', async () => { const server = createMockServer(); registerCountTool(server as never); diff --git a/src/core/server/tools/count-tool.ts b/src/core/server/tools/count-tool.ts index 97dac0a..f3dd33f 100644 --- a/src/core/server/tools/count-tool.ts +++ b/src/core/server/tools/count-tool.ts @@ -1,10 +1,17 @@ import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'; import { z } from 'zod'; -import { getPineconeClient } from '../client-context.js'; import { metadataFilterSchema, validateMetadataFilterDetailed } from '../metadata-filter.js'; import { normalizeNamespace } from '../namespace-utils.js'; import type { ServerContext } from '../server-context.js'; import { requireSuggested } from '../suggestion-flow.js'; +import { + getClientForResolvedSource, + optionalSourceField, + rejectSourceWithoutContext, + resolveSourceForTool, + sourceParamSchema, + sourceValidationError, +} from '../source-tool-utils.js'; import { classifyToolCatchError, flowGateToolError, @@ -19,6 +26,7 @@ type CountExecParams = { namespace: string; query_text: string; metadata_filter?: Record; + source?: string; }; async function executeCount(params: CountExecParams, ctx?: ServerContext) { @@ -26,7 +34,7 @@ async function executeCount(params: CountExecParams, ctx?: ServerContext) { if (ctx?.disposed) { return jsonErrorResponse(lifecycleToolError('ServerContext has been disposed')); } - const { namespace, query_text, metadata_filter } = params; + const { namespace, query_text, metadata_filter, source } = params; const nsNorm = normalizeNamespace(namespace); if (!nsNorm) { return jsonErrorResponse( @@ -44,11 +52,28 @@ async function executeCount(params: CountExecParams, ctx?: ServerContext) { return jsonErrorResponse(validationToolError(err.message, err.field)); } } - const flowCheck = ctx ? ctx.requireSuggested(nsNorm) : requireSuggested(nsNorm); + const sourceError = rejectSourceWithoutContext(source, ctx); + if (sourceError) { + return jsonErrorResponse(sourceError); + } + let activeCtx = ctx; + let activeSource: string | undefined; + if (ctx) { + const resolved = await resolveSourceForTool(ctx, source, nsNorm); + if (!resolved.ok) { + return jsonErrorResponse(sourceValidationError(resolved.code, resolved.message)); + } + activeCtx = resolved.ctx; + activeSource = resolved.source; + } + + const flowCheck = activeCtx + ? activeCtx.requireSuggested(nsNorm, activeSource) + : requireSuggested(nsNorm); if (!flowCheck.ok) { return jsonErrorResponse(flowGateToolError(nsNorm, flowCheck.message)); } - const client = ctx ? ctx.getClient() : getPineconeClient(); + const client = getClientForResolvedSource(activeCtx, activeSource, 'count'); const { count, truncated } = await client.count({ query: query_text.trim(), namespace: nsNorm, @@ -59,6 +84,7 @@ async function executeCount(params: CountExecParams, ctx?: ServerContext) { count, truncated, namespace: nsNorm, + ...optionalSourceField(activeCtx, activeSource ?? ''), metadata_filter, }; return validatedJsonResponse(countResponseSchema, response); @@ -96,6 +122,7 @@ export function registerCountTool(server: McpServer, ctx?: ServerContext): void .describe( 'Optional metadata filter. Use exact field names from list_namespaces. E.g. {"author": {"$in": ["Alex Doe", "A. Doe"]}} to count by author.' ), + source: sourceParamSchema, }, }, async (params) => executeCount(params, ctx) diff --git a/src/core/server/tools/generate-urls-tool.test.ts b/src/core/server/tools/generate-urls-tool.test.ts index 74e1b75..379b529 100644 --- a/src/core/server/tools/generate-urls-tool.test.ts +++ b/src/core/server/tools/generate-urls-tool.test.ts @@ -19,6 +19,20 @@ describe('generate_urls tool handler', () => { expect(err.field).toBe('namespace'); }); + it('returns VALIDATION when source is provided without ServerContext', async () => { + const server = createMockServer(); + registerGenerateUrlsTool(server as never); + const err = assertToolErrorCode( + await server.getHandler('generate_urls')!({ + namespace: 'mailing', + records: [{ document_number: 'P1234' }], + source: 'api_key_1', + }), + 'VALIDATION' + ); + expect(err.field).toBe('source'); + }); + it('returns PINECONE_ERROR when generateUrlForNamespace throws', async () => { vi.spyOn(urlRegistry, 'generateUrlForNamespace').mockImplementation(() => { throw new Error('generator boom'); diff --git a/src/core/server/tools/generate-urls-tool.ts b/src/core/server/tools/generate-urls-tool.ts index 25cd928..c78673e 100644 --- a/src/core/server/tools/generate-urls-tool.ts +++ b/src/core/server/tools/generate-urls-tool.ts @@ -3,6 +3,12 @@ import { z } from 'zod'; import { normalizeNamespace } from '../namespace-utils.js'; import type { ServerContext } from '../server-context.js'; import { generateUrlForNamespace } from '../url-registry.js'; +import { + rejectSourceWithoutContext, + resolveSourceForTool, + sourceParamSchema, + sourceValidationError, +} from '../source-tool-utils.js'; import { classifyToolCatchError, lifecycleToolError, @@ -41,6 +47,7 @@ export function registerGenerateUrlsTool(server: McpServer, ctx?: ServerContext) .describe( 'Array of records from retrieval results. Each item may be either metadata itself or an object containing a metadata field.' ), + source: sourceParamSchema, }, }, async (params) => { @@ -48,7 +55,7 @@ export function registerGenerateUrlsTool(server: McpServer, ctx?: ServerContext) if (ctx?.disposed) { return jsonErrorResponse(lifecycleToolError('ServerContext has been disposed')); } - const { namespace, records } = params; + const { namespace, records, source } = params; const nsNorm = normalizeNamespace(namespace); if (!nsNorm) { return jsonErrorResponse( @@ -57,10 +64,25 @@ export function registerGenerateUrlsTool(server: McpServer, ctx?: ServerContext) }) ); } + const sourceError = rejectSourceWithoutContext(source, ctx); + if (sourceError) { + return jsonErrorResponse(sourceError); + } + let activeCtx = ctx; + let activeSource: string | undefined; + if (ctx) { + const resolved = await resolveSourceForTool(ctx, source, nsNorm); + if (!resolved.ok) { + return jsonErrorResponse(sourceValidationError(resolved.code, resolved.message)); + } + activeCtx = resolved.ctx; + activeSource = resolved.source; + } + const results = records.map((record, index) => { const metadata = extractMetadata(record); - const generated = ctx - ? ctx.generateUrlForNamespace(nsNorm, metadata) + const generated = activeCtx + ? activeCtx.generateUrlForNamespace(nsNorm, metadata, activeSource) : generateUrlForNamespace(nsNorm, metadata); return { index, diff --git a/src/core/server/tools/guided-query-tool.multi-source.test.ts b/src/core/server/tools/guided-query-tool.multi-source.test.ts new file mode 100644 index 0000000..4c5e331 --- /dev/null +++ b/src/core/server/tools/guided-query-tool.multi-source.test.ts @@ -0,0 +1,79 @@ +import { describe, expect, it, vi } from 'vitest'; +import { guidedQueryResponseSchema } from '../response-schemas.js'; +import { registerGuidedQueryTool } from './guided-query-tool.js'; +import { + assertToolErrorCode, + createMockServer, + createMultiSourceTestContext, + expectMatchesResponseSchema, + makeHybridQueryResult, + makeMockPineconeClient, + parseToolJson, +} from './test-helpers.js'; + +describe('guided_query tool (multi-source)', () => { + it('sets selected_source in decision_trace when namespace is auto-routed', async () => { + const query = vi.fn().mockResolvedValue(makeHybridQueryResult()); + const client1 = makeMockPineconeClient(['papers'], { query }); + const client2 = makeMockPineconeClient(['internal']); + const clients = new Map([ + ['api_key_1', client1], + ['api_key_2', client2], + ]); + const { ctx } = createMultiSourceTestContext({ + namespacesBySource: { api_key_1: ['papers'], api_key_2: ['internal'] }, + clients, + }); + + const server = createMockServer(); + registerGuidedQueryTool(server as never, ctx); + const body = parseToolJson( + await server.getHandler('guided_query')!({ + user_query: 'What do papers say about contracts?', + preferred_tool: 'fast', + enrich_urls: false, + }) + ); + expectMatchesResponseSchema(guidedQueryResponseSchema, body); + const trace = (body['experimental'] as Record)['decision_trace'] as Record< + string, + unknown + >; + expect(trace['selected_source']).toBe('api_key_1'); + expect(trace['selected_namespace']).toBe('papers'); + expect(query).toHaveBeenCalledOnce(); + }); + + it('returns VALIDATION on namespace when namespace is missing on all sources', async () => { + const { ctx } = createMultiSourceTestContext(); + const server = createMockServer(); + registerGuidedQueryTool(server as never, ctx); + const err = assertToolErrorCode( + await server.getHandler('guided_query')!({ + user_query: 'contracts', + namespace: 'missing-ns', + preferred_tool: 'fast', + enrich_urls: false, + }), + 'VALIDATION' + ); + expect(err.field).toBe('namespace'); + }); + + it('returns VALIDATION on namespace when namespace exists on multiple sources', async () => { + const { ctx } = createMultiSourceTestContext(); + const server = createMockServer(); + registerGuidedQueryTool(server as never, ctx); + const err = assertToolErrorCode( + await server.getHandler('guided_query')!({ + user_query: 'contracts', + namespace: 'shared', + preferred_tool: 'fast', + enrich_urls: false, + }), + 'VALIDATION' + ); + expect(err.field).toBe('namespace'); + expect(err.message).toMatch(/multiple sources/i); + }); +}); diff --git a/src/core/server/tools/guided-query-tool.test.ts b/src/core/server/tools/guided-query-tool.test.ts index b43c4ab..ac5a540 100644 --- a/src/core/server/tools/guided-query-tool.test.ts +++ b/src/core/server/tools/guided-query-tool.test.ts @@ -49,6 +49,20 @@ describe('guided_query tool handler', () => { } as never); }); + it('returns VALIDATION when source is provided without ServerContext', async () => { + const server = createMockServer(); + registerGuidedQueryTool(server as never); + const err = assertToolErrorCode( + await server.getHandler('guided_query')!({ + user_query: 'What does the paper say?', + source: 'api_key_1', + }), + 'VALIDATION' + ); + expect(err.field).toBe('source'); + expect(mockedGetNamespaces).not.toHaveBeenCalled(); + }); + it('guided_query: surfaces rerank failure in decision_trace', async () => { mockedGetClient.mockReturnValue({ query: vi.fn().mockResolvedValue( diff --git a/src/core/server/tools/guided-query-tool.ts b/src/core/server/tools/guided-query-tool.ts index 0528e17..f59ef9c 100644 --- a/src/core/server/tools/guided-query-tool.ts +++ b/src/core/server/tools/guided-query-tool.ts @@ -2,7 +2,6 @@ import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'; import { z } from 'zod'; import { FAST_QUERY_FIELDS, MAX_TOP_K, MIN_TOP_K } from '../../../constants.js'; import { guidedRerankStatus } from '../../rerank-trace.js'; -import { getPineconeClient } from '../client-context.js'; import { formatQueryResultRows } from '../format-query-result.js'; import { metadataFilterSchema, validateMetadataFilterDetailed } from '../metadata-filter.js'; import { rankNamespacesByQuery } from '../namespace-router.js'; @@ -10,6 +9,14 @@ import { getNamespacesWithCache } from '../namespaces-cache.js'; import { normalizeNamespace } from '../namespace-utils.js'; import { suggestQueryParams } from '../query-suggestion.js'; import type { ServerContext } from '../server-context.js'; +import { + getClientForResolvedSource, + optionalSourceField, + rejectSourceWithoutContext, + resolveSourceForTool, + sourceParamSchema, + sourceValidationError, +} from '../source-tool-utils.js'; import { buildGuidedQueryExperimental, buildQueryExperimental, @@ -82,6 +89,7 @@ export function registerGuidedQueryTool(server: McpServer, ctx?: ServerContext): .describe( 'If true, enrich result URLs using the namespace URL generator when metadata.url is missing (if supported for that namespace).' ), + source: sourceParamSchema, }, }, async (params) => { @@ -96,6 +104,7 @@ export function registerGuidedQueryTool(server: McpServer, ctx?: ServerContext): top_k, preferred_tool, enrich_urls, + source: inputSource, } = params; if (!user_query?.trim()) { @@ -109,21 +118,41 @@ export function registerGuidedQueryTool(server: McpServer, ctx?: ServerContext): } } + const sourceError = rejectSourceWithoutContext(inputSource, ctx); + if (sourceError) { + return jsonErrorResponse(sourceError); + } + const queryText = user_query.trim(); const { data: namespaces, cache_hit } = ctx - ? await ctx.getNamespacesWithCache() + ? await ctx.getNamespacesWithCache(inputSource) : await getNamespacesWithCache(); const ranked = rankNamespacesByQuery(queryText, namespaces, 3); let namespace: string | null = null; + let selectedSource: string | undefined; if (inputNamespace !== undefined) { namespace = normalizeNamespace(inputNamespace); if (!namespace) { return jsonErrorResponse(validationToolError('namespace cannot be empty', 'namespace')); } + if (ctx?.isMultiSource() || inputSource) { + const resolved = await resolveSourceForTool(ctx, inputSource, namespace); + if (!resolved.ok) { + return jsonErrorResponse(sourceValidationError(resolved.code, resolved.message)); + } + selectedSource = resolved.source; + } } else { - const top = ranked[0]?.namespace; - namespace = top ? normalizeNamespace(top) : null; + const top = ranked[0]; + namespace = top ? normalizeNamespace(top.namespace) : null; + selectedSource = top?.source; + if (namespace && !selectedSource && ctx) { + const resolved = await ctx.resolveSource(inputSource, namespace); + if (resolved.ok) { + selectedSource = resolved.source; + } + } } /* * ToolError mapping: empty index / no routable namespace is backend/data state @@ -143,7 +172,9 @@ export function registerGuidedQueryTool(server: McpServer, ctx?: ServerContext): } const ns = namespaces.find( - (n) => n.namespace === namespace || normalizeNamespace(n.namespace) === namespace + (n) => + (n.namespace === namespace || normalizeNamespace(n.namespace) === namespace) && + (selectedSource === undefined || n.source === selectedSource) ); const suggestion = suggestQueryParams(ns?.metadata ?? null, queryText); if (!suggestion.namespace_found) { @@ -160,7 +191,17 @@ export function registerGuidedQueryTool(server: McpServer, ctx?: ServerContext): } const selectedTool: GuidedToolName = resolveGuidedToolName(preferred_tool, suggestion); - if (ctx) { + if (ctx && selectedSource) { + ctx.markSuggested( + namespace, + { + recommended_tool: selectedTool, + suggested_fields: suggestion.suggested_fields, + user_query: queryText, + }, + selectedSource + ); + } else if (ctx) { ctx.markSuggested(namespace, { recommended_tool: selectedTool, suggested_fields: suggestion.suggested_fields, @@ -174,13 +215,14 @@ export function registerGuidedQueryTool(server: McpServer, ctx?: ServerContext): }); } - const client = ctx ? ctx.getClient() : getPineconeClient(); + const client = getClientForResolvedSource(ctx, selectedSource, 'guided_query'); const enrichUrls = enrich_urls ?? true; const baseTrace = { cache_hit, input_namespace: inputNamespace ?? null, routed_namespace: ranked[0]?.namespace ?? null, selected_namespace: namespace, + ...(selectedSource !== undefined ? { selected_source: selectedSource } : {}), ranked_namespaces: ranked, suggested_fields: suggestion.suggested_fields, suggested_tool: suggestion.recommended_tool, @@ -204,6 +246,7 @@ export function registerGuidedQueryTool(server: McpServer, ctx?: ServerContext): result: { tool: 'count', namespace, + ...optionalSourceField(ctx, selectedSource ?? ''), query: queryText, ...(metadata_filter !== undefined ? { metadata_filter } : {}), count, @@ -239,12 +282,14 @@ export function registerGuidedQueryTool(server: McpServer, ctx?: ServerContext): namespace, enrichUrls: enrichUrls, ctx, + ...(selectedSource !== undefined ? { source: selectedSource } : {}), }); const result: QueryResponse = { status: 'success', mode, query: queryText, namespace, + ...optionalSourceField(ctx, selectedSource ?? ''), metadata_filter: metadata_filter, result_count: formattedResults.length, ...(fields?.length ? { fields } : {}), diff --git a/src/core/server/tools/keyword-search-tool.ts b/src/core/server/tools/keyword-search-tool.ts index f5e7f0a..44d51c6 100644 --- a/src/core/server/tools/keyword-search-tool.ts +++ b/src/core/server/tools/keyword-search-tool.ts @@ -1,16 +1,23 @@ import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'; import { z } from 'zod'; import { MAX_TOP_K, MIN_TOP_K } from '../../../constants.js'; -import { getPineconeClient } from '../client-context.js'; import { formatQueryResultRows } from '../format-query-result.js'; -import type { ServerContext } from '../server-context.js'; import { metadataFilterSchema, validateMetadataFilterDetailed } from '../metadata-filter.js'; -import type { ToolError } from '../tool-error.js'; +import type { ServerContext } from '../server-context.js'; +import { + getClientForResolvedSource, + optionalSourceField, + rejectSourceWithoutContext, + resolveSourceForTool, + sourceParamSchema, + sourceValidationError, +} from '../source-tool-utils.js'; import { classifyToolCatchError, lifecycleToolError, logToolError, validationToolError, + type ToolError, } from '../tool-error.js'; import { keywordSearchSuccessResponseSchema, @@ -33,13 +40,14 @@ async function executeKeywordSearch( top_k: number; metadata_filter?: Record; fields?: string[]; + source?: string; }, ctx?: ServerContext ): Promise { if (ctx?.disposed) { return { ok: false, error: lifecycleToolError('ServerContext has been disposed') }; } - const { query_text, namespace, top_k, metadata_filter, fields } = params; + const { query_text, namespace, top_k, metadata_filter, fields, source } = params; const normalizedQuery = query_text.trim(); const normalizedNamespace = namespace?.trim() ?? ''; @@ -68,7 +76,26 @@ async function executeKeywordSearch( } } - const client = ctx ? ctx.getClient() : getPineconeClient(); + const sourceError = rejectSourceWithoutContext(source, ctx); + if (sourceError) { + return { ok: false, error: sourceError }; + } + + let activeCtx = ctx; + let activeSource: string | undefined; + if (ctx) { + const resolved = await resolveSourceForTool(ctx, source, normalizedNamespace); + if (!resolved.ok) { + return { + ok: false, + error: sourceValidationError(resolved.code, resolved.message), + }; + } + activeCtx = resolved.ctx; + activeSource = resolved.source; + } + + const client = getClientForResolvedSource(activeCtx, activeSource, 'keyword_search'); const results = await client.keywordSearch({ query: normalizedQuery, namespace: normalizedNamespace, @@ -78,14 +105,16 @@ async function executeKeywordSearch( }); const formattedResults = formatQueryResultRows(results, { + ctx: activeCtx, namespace: normalizedNamespace, - ctx, + source: activeSource, }); const response: KeywordSearchSuccessResponse = { status: 'success', query: normalizedQuery, namespace: normalizedNamespace, + ...(activeCtx && activeSource ? optionalSourceField(activeCtx, activeSource) : {}), index: client.getSparseIndexName(), metadata_filter: metadata_filter, result_count: formattedResults.length, @@ -130,6 +159,7 @@ export function registerKeywordSearchTool(server: McpServer, ctx?: ServerContext .describe( 'Optional field names to return. Omit for all fields; use suggest_query_params for suggestions.' ), + source: sourceParamSchema, }, }, async (params) => { @@ -141,6 +171,7 @@ export function registerKeywordSearchTool(server: McpServer, ctx?: ServerContext top_k: params.top_k, metadata_filter: params.metadata_filter, fields: params.fields, + source: params.source, }, ctx ); diff --git a/src/core/server/tools/list-namespaces-tool.context.test.ts b/src/core/server/tools/list-namespaces-tool.context.test.ts index 3d196c2..c02c0f2 100644 --- a/src/core/server/tools/list-namespaces-tool.context.test.ts +++ b/src/core/server/tools/list-namespaces-tool.context.test.ts @@ -3,8 +3,10 @@ import { registerListNamespacesTool } from './list-namespaces-tool.js'; import { listNamespacesResponseSchema } from '../response-schemas.js'; import { createMockServer, + createMultiSourceTestContext, createTestServerContext, expectMatchesResponseSchema, + makeMockPineconeClient, parseToolJson, } from './test-helpers.js'; @@ -62,3 +64,32 @@ describe('list_namespaces tool handler (ServerContext instance path)', () => { expect(listNamespacesWithMetadata).toHaveBeenCalledOnce(); }); }); + +describe('list_namespaces tool handler (multi-source)', () => { + it('tags namespaces with source and propagates source_errors on partial failure', async () => { + const client1 = makeMockPineconeClient(['wg21']); + const client2 = { + listNamespacesWithMetadata: vi.fn().mockRejectedValue(new Error('api_key_2 unreachable')), + query: vi.fn(), + count: vi.fn(), + keywordSearch: vi.fn(), + checkIndexes: vi.fn().mockResolvedValue({ ok: true, errors: [] }), + getSparseIndexName: () => 'sparse', + }; + const clients = new Map([ + ['api_key_1', client1], + ['api_key_2', client2], + ]); + const { ctx } = createMultiSourceTestContext({ clients }); + + const server = createMockServer(); + registerListNamespacesTool(server as never, ctx); + const body = parseToolJson(await server.getHandler('list_namespaces')!({})); + expectMatchesResponseSchema(listNamespacesResponseSchema, body); + const namespaces = body['namespaces'] as Array<{ name: string; source?: string }>; + expect(namespaces).toHaveLength(1); + expect(namespaces[0]).toMatchObject({ name: 'wg21', source: 'api_key_1' }); + expect(body['source_errors']).toEqual({ api_key_2: 'api_key_2 unreachable' }); + expect(body['cache_hit']).toBe(false); + }); +}); diff --git a/src/core/server/tools/list-namespaces-tool.test.ts b/src/core/server/tools/list-namespaces-tool.test.ts index e77d3a2..c925eb3 100644 --- a/src/core/server/tools/list-namespaces-tool.test.ts +++ b/src/core/server/tools/list-namespaces-tool.test.ts @@ -56,6 +56,16 @@ describe('list_namespaces tool handler', () => { expect(body.count).toBe(1); }); + it('returns VALIDATION when source is provided without ServerContext', async () => { + const server = createMockServer(); + registerListNamespacesTool(server as never); + const err = assertToolErrorCode( + await server.getHandler('list_namespaces')!({ source: 'api_key_1' }), + 'VALIDATION' + ); + expect(err.field).toBe('source'); + }); + it('returns error payload when getNamespacesWithCache throws', async () => { mockedGetNamespaces.mockRejectedValue(new Error('network down')); diff --git a/src/core/server/tools/list-namespaces-tool.ts b/src/core/server/tools/list-namespaces-tool.ts index c93ea7d..b9855ef 100644 --- a/src/core/server/tools/list-namespaces-tool.ts +++ b/src/core/server/tools/list-namespaces-tool.ts @@ -1,23 +1,42 @@ import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'; import { getNamespacesWithCache } from '../namespaces-cache.js'; import type { ServerContext } from '../server-context.js'; -import { classifyToolCatchError, lifecycleToolError, logToolError } from '../tool-error.js'; +import { rejectSourceWithoutContext, sourceParamSchema } from '../source-tool-utils.js'; +import { + classifyToolCatchError, + lifecycleToolError, + logToolError, + logToolInvocation, +} from '../tool-error.js'; import { listNamespacesResponseSchema, type ListNamespacesSuccessResponse, } from '../response-schemas.js'; import { jsonErrorResponse, validatedJsonResponse } from '../tool-response.js'; -async function executeListNamespaces(ctx?: ServerContext) { +async function executeListNamespaces(source: string | undefined, ctx?: ServerContext) { try { if (ctx?.disposed) { return jsonErrorResponse(lifecycleToolError('ServerContext has been disposed')); } - const { - data: namespacesInfo, - cache_hit, - expires_at, - } = ctx ? await ctx.getNamespacesWithCache() : await getNamespacesWithCache(); + const sourceError = rejectSourceWithoutContext(source, ctx); + if (sourceError) { + return jsonErrorResponse(sourceError); + } + if (source) { + logToolInvocation('list_namespaces', source); + } + const cacheResult = ctx + ? await ctx.getNamespacesWithCache(source) + : await getNamespacesWithCache(); + const { data: namespacesInfo, cache_hit, expires_at } = cacheResult; + const rawSourceErrors = 'source_errors' in cacheResult ? cacheResult.source_errors : undefined; + const source_errors = + rawSourceErrors !== undefined && + rawSourceErrors !== null && + typeof rawSourceErrors === 'object' + ? (rawSourceErrors as Record) + : undefined; const now = Date.now(); const ttlSeconds = Math.max(0, Math.floor((expires_at - now) / 1000)); @@ -27,10 +46,12 @@ async function executeListNamespaces(ctx?: ServerContext) { cache_ttl_seconds: ttlSeconds, expires_at_iso: new Date(expires_at).toISOString(), count: namespacesInfo.length, + ...(source_errors !== undefined && source_errors !== null ? { source_errors } : {}), namespaces: namespacesInfo.map((ns) => ({ name: ns.namespace, record_count: ns.recordCount, metadata_fields: ns.metadata, + ...(ns.source !== undefined ? { source: ns.source } : {}), })), }; @@ -50,9 +71,12 @@ export function registerListNamespacesTool(server: McpServer, ctx?: ServerContex 'List all available namespaces in the Pinecone index with their metadata fields and record counts. ' + 'Returns detailed information about each namespace including available metadata fields that can be used for filtering in queries. ' + 'Use this tool first to discover which namespaces exist and what metadata fields are available for filtering. ' + + 'In multi-source mode, omit source to list namespaces from all configured projects (each tagged with source). ' + 'Results are cached in-memory for 30 minutes for better performance.', - inputSchema: {}, + inputSchema: { + source: sourceParamSchema, + }, }, - async () => executeListNamespaces(ctx) + async (params) => executeListNamespaces(params.source, ctx) ); } diff --git a/src/core/server/tools/list-sources-tool.ts b/src/core/server/tools/list-sources-tool.ts new file mode 100644 index 0000000..99a5641 --- /dev/null +++ b/src/core/server/tools/list-sources-tool.ts @@ -0,0 +1,39 @@ +import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'; +import type { ServerContext } from '../server-context.js'; +import { classifyToolCatchError, lifecycleToolError, logToolError } from '../tool-error.js'; +import { listSourcesResponseSchema, type ListSourcesResponse } from '../response-schemas.js'; +import { jsonErrorResponse, validatedJsonResponse } from '../tool-response.js'; + +/** Register list_sources when multiple Pinecone projects are configured. */ +export function registerListSourcesTool(server: McpServer, ctx?: ServerContext): void { + server.registerTool( + 'list_sources', + { + description: + 'List configured Pinecone source names when multiple API keys/projects are active. ' + + 'Returns source ids and the default source.', + inputSchema: {}, + }, + async () => { + try { + if (ctx?.disposed) { + return jsonErrorResponse(lifecycleToolError('ServerContext has been disposed')); + } + if (!ctx?.isMultiSource()) { + return jsonErrorResponse( + lifecycleToolError('list_sources is only available in multi-source mode.') + ); + } + const response: ListSourcesResponse = { + status: 'success', + sources: ctx.listSources(), + default: ctx.getDefaultSourceName(), + }; + return validatedJsonResponse(listSourcesResponseSchema, response); + } catch (error) { + logToolError('list_sources', error); + return jsonErrorResponse(classifyToolCatchError(error, 'Failed to list sources')); + } + } + ); +} diff --git a/src/core/server/tools/namespace-router-tool.ts b/src/core/server/tools/namespace-router-tool.ts index 2ffb271..99a6c69 100644 --- a/src/core/server/tools/namespace-router-tool.ts +++ b/src/core/server/tools/namespace-router-tool.ts @@ -3,10 +3,12 @@ import { z } from 'zod'; import { getNamespacesWithCache } from '../namespaces-cache.js'; import { rankNamespacesByQuery } from '../namespace-router.js'; import type { ServerContext } from '../server-context.js'; +import { rejectSourceWithoutContext, sourceParamSchema } from '../source-tool-utils.js'; import { classifyToolCatchError, lifecycleToolError, logToolError, + logToolInvocation, validationToolError, } from '../tool-error.js'; import { @@ -22,7 +24,8 @@ export function registerNamespaceRouterTool(server: McpServer, ctx?: ServerConte { description: 'Suggest likely namespace(s) for a user query using namespace names, metadata fields, and keyword heuristics. ' + - 'Use before suggest_query_params when namespace is unclear.', + 'Use before suggest_query_params when namespace is unclear. ' + + 'In multi-source mode, ranks across all sources unless source is set.', inputSchema: { user_query: z .string() @@ -34,6 +37,7 @@ export function registerNamespaceRouterTool(server: McpServer, ctx?: ServerConte .max(5) .default(3) .describe('Maximum number of suggested namespaces (1-5).'), + source: sourceParamSchema, }, }, async (params) => { @@ -41,21 +45,30 @@ export function registerNamespaceRouterTool(server: McpServer, ctx?: ServerConte if (ctx?.disposed) { return jsonErrorResponse(lifecycleToolError('ServerContext has been disposed')); } - const { user_query, top_n } = params; + const { user_query, top_n, source } = params; + const sourceError = rejectSourceWithoutContext(source, ctx); + if (sourceError) { + return jsonErrorResponse(sourceError); + } if (!user_query?.trim()) { return jsonErrorResponse(validationToolError('user_query cannot be empty', 'user_query')); } + if (source) { + logToolInvocation('namespace_router', source); + } const { data, cache_hit } = ctx - ? await ctx.getNamespacesWithCache() + ? await ctx.getNamespacesWithCache(source) : await getNamespacesWithCache(); const ranked = rankNamespacesByQuery(user_query.trim(), data, top_n); + const top = ranked[0]; const response: NamespaceRouterResponse = { status: 'success', cache_hit, user_query: user_query.trim(), suggestions: ranked, - recommended_namespace: ranked[0]?.namespace ?? null, + recommended_namespace: top?.namespace ?? null, + ...(top?.source !== undefined ? { recommended_source: top.source } : {}), }; return validatedJsonResponse(namespaceRouterResponseSchema, response); } catch (error) { diff --git a/src/core/server/tools/query-documents-tool.test.ts b/src/core/server/tools/query-documents-tool.test.ts index 2affd46..ff84b45 100644 --- a/src/core/server/tools/query-documents-tool.test.ts +++ b/src/core/server/tools/query-documents-tool.test.ts @@ -54,6 +54,21 @@ describe('query_documents tool handler', () => { vi.restoreAllMocks(); }); + it('returns VALIDATION when source is provided without ServerContext', async () => { + const server = createMockServer(); + registerQueryDocumentsTool(server as never); + const err = assertToolErrorCode( + await server.getHandler('query_documents')!({ + query_text: 'semantic question', + namespace: 'wg21', + source: 'api_key_1', + }), + 'VALIDATION' + ); + expect(err.field).toBe('source'); + expect(mockedGetClient().query).not.toHaveBeenCalled(); + }); + it('happy path: queries chunks, reassembles, and returns documents', async () => { const server = createMockServer(); registerQueryDocumentsTool(server as never); diff --git a/src/core/server/tools/query-documents-tool.ts b/src/core/server/tools/query-documents-tool.ts index 8d201cc..0d78b12 100644 --- a/src/core/server/tools/query-documents-tool.ts +++ b/src/core/server/tools/query-documents-tool.ts @@ -5,12 +5,19 @@ import { MAX_QUERY_DOCUMENTS_TOP_K, QUERY_DOCUMENTS_MAX_CHUNKS, } from '../../../constants.js'; -import { getPineconeClient } from '../client-context.js'; import { metadataFilterSchema, validateMetadataFilterDetailed } from '../metadata-filter.js'; import { normalizeNamespace } from '../namespace-utils.js'; import { reassembleByDocument } from '../reassemble-documents.js'; import type { ServerContext } from '../server-context.js'; import { requireSuggested } from '../suggestion-flow.js'; +import { + getClientForResolvedSource, + optionalSourceField, + rejectSourceWithoutContext, + resolveSourceForTool, + sourceParamSchema, + sourceValidationError, +} from '../source-tool-utils.js'; import { classifyToolCatchError, flowGateToolError, @@ -76,6 +83,7 @@ export function registerQueryDocumentsTool(server: McpServer, ctx?: ServerContex .describe( 'Max chunks to merge per document (default 200). Lower for shorter merged_content.' ), + source: sourceParamSchema, }, }, async (params) => { @@ -89,6 +97,7 @@ export function registerQueryDocumentsTool(server: McpServer, ctx?: ServerContex top_k = DEFAULT_QUERY_DOCUMENTS_TOP_K, metadata_filter, max_chunks_per_document, + source, } = params; if (!query_text?.trim()) { @@ -111,13 +120,31 @@ export function registerQueryDocumentsTool(server: McpServer, ctx?: ServerContex ); } - const flowCheck = ctx ? ctx.requireSuggested(nsNorm) : requireSuggested(nsNorm); + const sourceError = rejectSourceWithoutContext(source, ctx); + if (sourceError) { + return jsonErrorResponse(sourceError); + } + + let activeCtx = ctx; + let activeSource: string | undefined; + if (ctx) { + const resolved = await resolveSourceForTool(ctx, source, nsNorm); + if (!resolved.ok) { + return jsonErrorResponse(sourceValidationError(resolved.code, resolved.message)); + } + activeCtx = resolved.ctx; + activeSource = resolved.source; + } + + const flowCheck = activeCtx + ? activeCtx.requireSuggested(nsNorm, activeSource) + : requireSuggested(nsNorm); if (!flowCheck.ok) { return jsonErrorResponse(flowGateToolError(nsNorm, flowCheck.message)); } const chunkLimit = Math.min(QUERY_DOCUMENTS_MAX_CHUNKS, top_k * CHUNKS_PER_DOCUMENT); - const client = ctx ? ctx.getClient() : getPineconeClient(); + const client = getClientForResolvedSource(activeCtx, activeSource, 'query_documents'); const queryOutcome = await client.query({ query: query_text.trim(), topK: chunkLimit, @@ -139,6 +166,7 @@ export function registerQueryDocumentsTool(server: McpServer, ctx?: ServerContex status: 'success', query: query_text.trim(), namespace: nsNorm, + ...optionalSourceField(activeCtx, activeSource ?? ''), metadata_filter, result_count: topDocuments.length, ...buildQueryExperimental(queryOutcome), diff --git a/src/core/server/tools/query-tool.test.ts b/src/core/server/tools/query-tool.test.ts index 6bcbcb4..1db2625 100644 --- a/src/core/server/tools/query-tool.test.ts +++ b/src/core/server/tools/query-tool.test.ts @@ -78,6 +78,22 @@ describe('query tool handler (preset-driven)', () => { expect(query).not.toHaveBeenCalled(); }); + it('returns VALIDATION when source is provided without ServerContext', async () => { + const server = createMockServer(); + registerQueryTool(server as never); + const err = assertToolErrorCode( + await server.getHandler('query')!({ + query_text: 'hello', + namespace: 'wg21', + top_k: 5, + preset: 'full', + source: 'api_key_1', + }), + 'VALIDATION' + ); + expect(err.field).toBe('source'); + }); + it('query (preset=full): happy path calls client.query and returns formatted rows', async () => { const server = createMockServer(); registerQueryTool(server as never); diff --git a/src/core/server/tools/query-tool.ts b/src/core/server/tools/query-tool.ts index 4be6030..5c3bb87 100644 --- a/src/core/server/tools/query-tool.ts +++ b/src/core/server/tools/query-tool.ts @@ -1,12 +1,19 @@ import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'; import { z } from 'zod'; import { FAST_QUERY_FIELDS, MAX_TOP_K, MIN_TOP_K } from '../../../constants.js'; -import { getPineconeClient } from '../client-context.js'; import { formatQueryResultRows } from '../format-query-result.js'; import { metadataFilterSchema, validateMetadataFilterDetailed } from '../metadata-filter.js'; import { normalizeNamespace } from '../namespace-utils.js'; import type { ServerContext } from '../server-context.js'; import { requireSuggested } from '../suggestion-flow.js'; +import { + getClientForResolvedSource, + optionalSourceField, + rejectSourceWithoutContext, + resolveSourceForTool, + sourceParamSchema, + sourceValidationError, +} from '../source-tool-utils.js'; import { classifyToolCatchError, flowGateToolError, @@ -31,11 +38,13 @@ type QueryExecParams = { metadata_filter?: Record; fields?: string[]; mode: QueryMode; + source?: string; }; /** Run the query tool: validate flow, call Pinecone, format and return results. */ async function executeQuery(params: QueryExecParams, ctx?: ServerContext) { - const { query_text, namespace, top_k, use_reranking, metadata_filter, fields, mode } = params; + const { query_text, namespace, top_k, use_reranking, metadata_filter, fields, mode, source } = + params; try { if (ctx?.disposed) { return jsonErrorResponse(lifecycleToolError('ServerContext has been disposed')); @@ -62,12 +71,30 @@ async function executeQuery(params: QueryExecParams, ctx?: ServerContext) { ); } - const flowCheck = ctx ? ctx.requireSuggested(nsNorm) : requireSuggested(nsNorm); + const sourceError = rejectSourceWithoutContext(source, ctx); + if (sourceError) { + return jsonErrorResponse(sourceError); + } + + let activeCtx = ctx; + let activeSource: string | undefined; + if (ctx) { + const resolved = await resolveSourceForTool(ctx, source, nsNorm); + if (!resolved.ok) { + return jsonErrorResponse(sourceValidationError(resolved.code, resolved.message)); + } + activeCtx = resolved.ctx; + activeSource = resolved.source; + } + + const flowCheck = activeCtx + ? activeCtx.requireSuggested(nsNorm, activeSource) + : requireSuggested(nsNorm); if (!flowCheck.ok) { return jsonErrorResponse(flowGateToolError(nsNorm, flowCheck.message)); } - const client = ctx ? ctx.getClient() : getPineconeClient(); + const client = getClientForResolvedSource(activeCtx, activeSource, 'query'); const queryOutcome = await client.query({ query: query_text.trim(), topK: top_k, @@ -77,13 +104,18 @@ async function executeQuery(params: QueryExecParams, ctx?: ServerContext) { fields: fields?.length ? fields : undefined, }); - const formattedResults = formatQueryResultRows(queryOutcome.results, { ctx }); + const formattedResults = formatQueryResultRows(queryOutcome.results, { + ctx: activeCtx, + namespace: nsNorm, + source: activeSource, + }); const response: QuerySuccessResponse = { status: 'success', mode, query: query_text, namespace: nsNorm, + ...optionalSourceField(activeCtx, activeSource ?? ''), metadata_filter: metadata_filter, result_count: formattedResults.length, results: formattedResults, @@ -122,6 +154,7 @@ const baseSchema = { .describe( 'Optional field names to return from Pinecone. Use suggest_query_params suggested_fields for better performance.' ), + source: sourceParamSchema, }; /** @@ -180,6 +213,7 @@ export function registerQueryTool(server: McpServer, ctx?: ServerContext): void metadata_filter: params.metadata_filter, fields, mode, + source: params.source, }, ctx ); diff --git a/src/core/server/tools/test-helpers.ts b/src/core/server/tools/test-helpers.ts index f7c2698..cdcd1ab 100644 --- a/src/core/server/tools/test-helpers.ts +++ b/src/core/server/tools/test-helpers.ts @@ -1,8 +1,11 @@ import type { z } from 'zod'; +import { vi } from 'vitest'; import type { HybridQueryResult, SearchResult } from '../../../types.js'; import { resolveConfig } from '../../config.js'; import type { PineconeClient } from '../../pinecone-client.js'; import type { ConfigOverrides, CoreServerConfig } from '../../config.js'; +import type { SourceDefinition } from '../source-config.js'; +import { buildSourceRegistry } from '../source-registry.js'; import { ServerContext, teardownDefaultServerContext, @@ -164,3 +167,73 @@ export function createTestServerContext(options?: { export function expectMatchesResponseSchema(schema: z.ZodType, body: unknown): T { return schema.parse(body); } + +/** Default two-source definitions for multi-source tests (overlapping `shared` namespace). */ +export const DEFAULT_MULTI_SOURCE_DEFINITIONS: SourceDefinition[] = [ + { name: 'api_key_1', apiKey: 'k1', indexName: 'idx-a', sparseIndexName: 'idx-a-sparse' }, + { name: 'api_key_2', apiKey: 'k2', indexName: 'idx-b', sparseIndexName: 'idx-b-sparse' }, +]; + +/** Minimal mock {@link PineconeClient} with configurable namespace list. */ +export function makeMockPineconeClient( + namespaces: string[], + options?: { query?: ReturnType } +) { + return { + listNamespacesWithMetadata: vi.fn().mockResolvedValue( + namespaces.map((namespace) => ({ + namespace, + recordCount: 1, + metadata: { title: 'string' }, + })) + ), + query: options?.query ?? vi.fn(), + count: vi.fn().mockResolvedValue({ count: 0, truncated: false }), + keywordSearch: vi.fn().mockResolvedValue([]), + checkIndexes: vi.fn().mockResolvedValue({ ok: true, errors: [] }), + getSparseIndexName: () => 'sparse', + }; +} + +export type MultiSourceTestContext = { + ctx: CoreServerContext; + clients: Map>; + sources: SourceDefinition[]; + registry: ReturnType; +}; + +/** Build an isolated multi-source {@link ServerContext} with injectable per-source mock clients. */ +export function createMultiSourceTestContext(options?: { + sources?: SourceDefinition[]; + namespacesBySource?: Record; + clients?: Map>; + defaultSource?: string; + config?: ConfigOverrides; +}): MultiSourceTestContext { + const sources = options?.sources ?? DEFAULT_MULTI_SOURCE_DEFINITIONS; + const defaultSource = options?.defaultSource ?? sources[0]!.name; + const namespacesBySource = + options?.namespacesBySource ?? + Object.fromEntries( + sources.map((s, i) => [s.name, i === 0 ? ['wg21', 'shared'] : ['shared', 'internal']]) + ); + const clients = + options?.clients ?? + new Map(sources.map((s) => [s.name, makeMockPineconeClient(namespacesBySource[s.name] ?? [])])); + const registry = buildSourceRegistry({ + sources, + defaultSource, + cacheTtlMs: 60_000, + defaultTopK: 10, + requestTimeoutMs: 15_000, + clients: clients as unknown as Map, + }); + const inlineSources = sources.map((s) => `${s.name}:${s.apiKey}:${s.indexName}`).join(';'); + const config = resolveConfig({ + sources: inlineSources, + disableSuggestFlow: false, + ...options?.config, + }); + const ctx = new ServerContext(config, { sourceRegistry: registry }); + return { ctx, clients, sources, registry }; +} diff --git a/src/core/server/url-registry.ts b/src/core/server/url-registry.ts index b852a94..84408c6 100644 --- a/src/core/server/url-registry.ts +++ b/src/core/server/url-registry.ts @@ -73,9 +73,9 @@ export function registerUrlGenerator(namespace: string, generator: UrlGeneratorF * {@link createServer} instead. See docs/MIGRATION.md#030-legacy-module-facade-deprecations. * @see ServerContext.unregisterUrlGenerator */ -export function unregisterUrlGenerator(namespace: string): boolean { +export function unregisterUrlGenerator(namespace: string, source?: string): boolean { warnLegacyFacade('unregisterUrlGenerator'); - return resolveDefaultServerContext().unregisterUrlGenerator(namespace); + return resolveDefaultServerContext().unregisterUrlGenerator(namespace, source); } /** @@ -86,9 +86,9 @@ export function unregisterUrlGenerator(namespace: string): boolean { * instead. See docs/MIGRATION.md#030-legacy-module-facade-deprecations. * @see ServerContext.hasUrlGenerator */ -export function hasUrlGenerator(namespace: string): boolean { +export function hasUrlGenerator(namespace: string, source?: string): boolean { warnLegacyFacade('hasUrlGenerator'); - return resolveDefaultServerContext().hasUrlGenerator(namespace); + return resolveDefaultServerContext().hasUrlGenerator(namespace, source); } /** diff --git a/src/core/setup-multi-instance.test.ts b/src/core/setup-multi-instance.test.ts index d95113c..c153d62 100644 --- a/src/core/setup-multi-instance.test.ts +++ b/src/core/setup-multi-instance.test.ts @@ -1,8 +1,9 @@ import { afterEach, describe, expect, it, vi } from 'vitest'; import { resolveAllianceConfig } from '../alliance/config.js'; import { setupAllianceServer } from '../alliance/setup.js'; -import { createIsolatedContext } from './server/server-context.js'; +import { ServerContext, createIsolatedContext } from './server/server-context.js'; import { setupCoreServer, teardownServer } from './setup.js'; +import * as listSourcesTool from './server/tools/list-sources-tool.js'; import { createTestServerContext, isolateFromDefaultContext, @@ -131,4 +132,19 @@ describe('setup multi-instance (phase 4)', () => { await expect(setupCoreServer({ context: ctxB })).resolves.toBeDefined(); expect(ctxB.disposed).toBe(false); }); + + it('registers list_sources when context resolves multi-source config from env', async () => { + isolateFromDefaultContext(); + const registerSpy = vi.spyOn(listSourcesTool, 'registerListSourcesTool'); + vi.stubEnv('PINECONE_SOURCES', 'api_key_1:sk-a:index_name_1;api_key_2:sk-b:index_name_2'); + try { + const ctx = new ServerContext(); + await setupCoreServer({ context: ctx }); + expect(ctx.isMultiSource()).toBe(true); + expect(registerSpy).toHaveBeenCalledOnce(); + } finally { + registerSpy.mockRestore(); + vi.unstubAllEnvs(); + } + }); }); diff --git a/src/core/setup.ts b/src/core/setup.ts index 48cea3d..fdf23b2 100644 --- a/src/core/setup.ts +++ b/src/core/setup.ts @@ -14,6 +14,7 @@ import { import { registerCountTool } from './server/tools/count-tool.js'; import { registerGenerateUrlsTool } from './server/tools/generate-urls-tool.js'; import { registerGuidedQueryTool } from './server/tools/guided-query-tool.js'; +import { registerListSourcesTool } from './server/tools/list-sources-tool.js'; import { registerKeywordSearchTool } from './server/tools/keyword-search-tool.js'; import { registerListNamespacesTool } from './server/tools/list-namespaces-tool.js'; import { registerNamespaceRouterTool } from './server/tools/namespace-router-tool.js'; @@ -164,6 +165,11 @@ async function registerCoreToolSurface( registerGenerateUrlsTool(server, ctx); registerGuidedQueryTool(server, ctx); + ctx.getConfig(); + if (ctx.isMultiSource()) { + registerListSourcesTool(server, ctx); + } + ctx.markToolsRegistered(); const handle = server as ServerHandle; diff --git a/src/index.ts b/src/index.ts index b330f1c..c231539 100644 --- a/src/index.ts +++ b/src/index.ts @@ -3,7 +3,7 @@ /** * Pinecone Read-Only MCP CLI entry point. * - * Thin composition root: parseCli() -> resolveAllianceConfig() -> createServer(config, { client }) + * Thin composition root: parseCli() -> resolveAllianceConfig() -> createServer(config, composition) * -> setupAllianceServer({ context: ctx }) -> connect to stdio transport. */ @@ -14,8 +14,15 @@ import type { AllianceServerConfig } from './alliance/config.js'; import { resolveAllianceConfig } from './alliance/config.js'; import { PineconeClient } from './core/pinecone-client.js'; import { createServer } from './core/server/server-context.js'; +import { buildSourceRegistry } from './core/server/source-registry.js'; import { setupAllianceServer } from './alliance/setup.js'; -import { setLogFormat, setLogLevel, warn as logWarn } from './logger.js'; +import { + error as logError, + redactApiKey, + setLogFormat, + setLogLevel, + warn as logWarn, +} from './logger.js'; dotenv.config(); @@ -37,7 +44,7 @@ function buildConfigOrExit(): AllianceServerConfig { try { return resolveAllianceConfig(parsed.overrides); } catch (err) { - const message = err instanceof Error ? err.message : String(err); + const message = redactApiKey(err instanceof Error ? err.message : String(err)); process.stderr.write(`Error: ${message}\n`); process.exit(1); } @@ -56,33 +63,59 @@ async function main(): Promise { ); } - const client = new PineconeClient({ - apiKey: config.apiKey, - indexName: config.indexName, - sparseIndexName: config.sparseIndexName, - rerankModel: config.rerankModel, - defaultTopK: config.defaultTopK, - requestTimeoutMs: config.requestTimeoutMs, - }); - const ctx = createServer(config, { client }); + let ctx; + if (config.sources && config.sources.length > 0) { + const sourceRegistry = buildSourceRegistry({ + sources: config.sources, + defaultSource: config.defaultSource ?? config.sources[0]!.name, + cacheTtlMs: config.cacheTtlMs, + defaultTopK: config.defaultTopK, + requestTimeoutMs: config.requestTimeoutMs, + }); + ctx = createServer(config, { sourceRegistry }); + } else { + const client = new PineconeClient({ + apiKey: config.apiKey, + indexName: config.indexName, + sparseIndexName: config.sparseIndexName, + rerankModel: config.rerankModel, + defaultTopK: config.defaultTopK, + requestTimeoutMs: config.requestTimeoutMs, + }); + ctx = createServer(config, { client }); + } if (config.checkIndexes) { - const result = await client.checkIndexes(); + const result = await ctx.checkAllIndexes(); if (!result.ok) { for (const err of result.errors) { - process.stderr.write(`--check-indexes: ${err}\n`); + process.stderr.write(`--check-indexes: ${redactApiKey(err)}\n`); } process.exit(1); } - process.stderr.write( - `--check-indexes: dense index "${config.indexName}" and sparse index "${config.sparseIndexName}" reachable.\n` - ); + if (config.sources && config.sources.length > 0) { + process.stderr.write( + `--check-indexes: all ${config.sources.length} source(s) reachable.\n` + ); + } else { + process.stderr.write( + `--check-indexes: dense index "${config.indexName}" and sparse index "${config.sparseIndexName}" reachable.\n` + ); + } + process.exit(0); } process.stderr.write(`Starting Pinecone Read-Only MCP server with stdio transport\n`); - process.stderr.write( - `Using Pinecone index: ${config.indexName} (sparse: ${config.sparseIndexName})\n` - ); + if (config.sources && config.sources.length > 0) { + const names = config.sources.map((s) => s.name).join(', '); + process.stderr.write( + `Multi-source mode: [${names}] (default: ${config.defaultSource ?? config.sources[0]!.name})\n` + ); + } else { + process.stderr.write( + `Using Pinecone index: ${config.indexName} (sparse: ${config.sparseIndexName})\n` + ); + } if (config.rerankModel) { process.stderr.write(`Rerank model: ${config.rerankModel}\n`); } @@ -104,7 +137,8 @@ async function main(): Promise { process.exit(0); }); } catch (error) { - process.stderr.write(`Fatal error in main(): ${(error as Error)?.stack ?? String(error)}\n`); + const summary = redactApiKey(error instanceof Error ? error.message : String(error)); + logError(`Fatal error in main(): ${summary}`); process.exit(1); } }