Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 30 additions & 0 deletions src/alliance/tools/guided-query-tool.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -107,6 +107,36 @@ describe('guided_query tool handler', () => {
expect(trace.rerank_status).toBe('skipped_no_model');
});

it('guided_query: redacts API keys from returned MCP content', async () => {
const apiKey = 'pcsk_guided_secret_1234567890abcdef';
mockedGetClient.mockReturnValue({
query: vi.fn().mockResolvedValue(
makeHybridQueryResult({
degraded: true,
degradation_reason: `rerank_failed: ${apiKey}`,
results: [makeSearchResult({ reranked: false })],
})
),
count: vi.fn().mockResolvedValue({ count: 7, truncated: false }),
} as never);

const server = createMockServer();
registerGuidedQueryTool(server as never);

const raw = await server.getHandler('guided_query')!({
user_query: 'What does the paper say about contracts?',
namespace: 'papers',
top_k: 8,
preferred_tool: 'auto',
enrich_urls: false,
});
const body = parseToolJson(raw);
const result = body.result as Record<string, unknown>;

expect(JSON.stringify(raw)).not.toContain(apiKey);
expect(result.degradation_reason).toBe('rerank_failed: ***');
});

it('runs query_detailed path on auto when user asks for content', async () => {
const server = createMockServer();
registerGuidedQueryTool(server as never);
Expand Down
61 changes: 61 additions & 0 deletions src/core/server/tool-response-redaction.test.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,61 @@
import { afterEach, describe, expect, it } from 'vitest';
import { redactApiKey, setLogLevel } from '../../logger.js';
import { classifyToolCatchError } from './tool-error.js';
import { jsonErrorResponse, jsonResponse } from './tool-response.js';
import { assertToolError, parseToolJson } from './tools/test-helpers.js';

describe('MCP response redaction', () => {
const apiKey = 'pcsk_response_secret_1234567890abcdef';

afterEach(() => {
setLogLevel('INFO');
});

it('redacts pcsk_ Pinecone API key strings directly', () => {
expect(redactApiKey(`connection failed for ${apiKey}`)).toBe('connection failed for ***');
});

it('redacts tool error message and suggestion fields before returning MCP content', () => {
const raw = jsonErrorResponse({
code: 'PINECONE_ERROR',
message: `Pinecone SDK error for ${apiKey}`,
recoverable: false,
suggestion: `Rotate ${apiKey}`,
});
const err = assertToolError(raw);

expect(JSON.stringify(raw)).not.toContain(apiKey);
expect(err.message).toBe('Pinecone SDK error for ***');
expect(err.suggestion).toBe('Rotate ***');
});

it('redacts catch-all DEBUG error messages before returning MCP content', () => {
setLogLevel('DEBUG');
const raw = jsonErrorResponse(
classifyToolCatchError(new Error(`Pinecone connection string ${apiKey}`), 'fallback')
);
const err = assertToolError(raw);

expect(JSON.stringify(raw)).not.toContain(apiKey);
expect(err.message).toBe('Pinecone connection string ***');
});

it('redacts nested success payload metadata before returning MCP content arrays', () => {
const body = parseToolJson(
jsonResponse({
status: 'success',
degradation_reason: `rerank_failed: ${apiKey}`,
decision_trace: {
note: `Authorization: Bearer ${apiKey}`,
},
diagnostic_metadata: {
nested: [`api_key=${apiKey}`],
},
})
);

expect(JSON.stringify(body)).not.toContain(apiKey);
expect(body.degradation_reason).toBe('rerank_failed: ***');
expect(JSON.stringify(body)).toContain('***');
});
});
5 changes: 3 additions & 2 deletions src/core/server/tool-response.ts
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import type { ToolError } from './tool-error.js';
import { toolErrorSchema } from './tool-error.js';
import { redactValue } from '../../logger.js';

export type TextPayload = {
content: Array<{ type: 'text'; text: string }>;
Expand All @@ -12,15 +13,15 @@ export function jsonResponse(payload: unknown): TextPayload {
content: [
{
type: 'text',
text: JSON.stringify(payload, null, 2),
text: JSON.stringify(redactValue(payload), null, 2),
},
],
};
}

/** Build an MCP tool error payload with JSON-stringified {@link ToolError} and isError: true. */
export function jsonErrorResponse(err: ToolError): TextPayload {
const validated = toolErrorSchema.parse(err);
const validated = toolErrorSchema.parse(redactValue(err));
return {
isError: true,
content: [
Expand Down
27 changes: 27 additions & 0 deletions src/core/server/tools/query-tool.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -245,6 +245,33 @@ describe('query tool handler (preset-driven)', () => {
expect(body.degradation_reason).toMatch(/rerank_skipped_no_model/);
});

it('query: redacts API keys from degradation reasons before returning content', async () => {
const apiKey = 'pcsk_query_secret_1234567890abcdef';
mockedGetClient.mockReturnValue({
query: vi.fn().mockResolvedValue(
makeHybridQueryResult({
degraded: true,
degradation_reason: `rerank_failed: ${apiKey}`,
})
),
count: vi.fn(),
} as never);

const server = createMockServer();
registerQueryTool(server as never);

const raw = await server.getHandler('query')!({
query_text: 'hello',
namespace: 'wg21',
top_k: 3,
preset: 'detailed',
});
const body = parseToolJson(raw);

expect(JSON.stringify(raw)).not.toContain(apiKey);
expect(body.degradation_reason).toBe('rerank_failed: ***');
});

it('query: surfaces unreranked hits when client returns reranked:false (rerank fallback shape)', async () => {
mockedGetClient.mockReturnValue({
query: vi.fn().mockResolvedValue(
Expand Down
3 changes: 2 additions & 1 deletion src/logger.ts
Original file line number Diff line number Diff line change
Expand Up @@ -59,13 +59,14 @@ export function redactApiKey(s: string): string {
/[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}/g,
'***'
);
out = out.replace(/\bpcsk_[A-Za-z0-9_-]+\b/g, '***');
out = out.replace(/(api[_-]?key["':\s=]+)([^\s"',}]+)/gi, '$1***');
out = out.replace(/(Authorization:\s*Bearer\s+)([^\s"',}]+)/gi, '$1***');
return out;
}

/** Recursively redact API keys from a serializable value. Returns a deep copy. */
function redactValue(value: unknown, seen: WeakSet<object> = new WeakSet()): unknown {
export function redactValue(value: unknown, seen: WeakSet<object> = new WeakSet()): unknown {
if (typeof value === 'string') return redactApiKey(value);
if (value === null || value === undefined) return value;
if (typeof value === 'object') {
Expand Down