mirror of
https://github.com/google-gemini/gemini-cli.git
synced 2026-05-14 13:53:02 -07:00
feat(security): Introduce Conseca framework (#13193)
This commit is contained in:
committed by
GitHub
parent
05bc0399f3
commit
dde844dbe1
@@ -0,0 +1,279 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { describe, it, expect, beforeEach, vi } from 'vitest';
|
||||
import { ConsecaSafetyChecker } from './conseca.js';
|
||||
import { SafetyCheckDecision } from '../protocol.js';
|
||||
import type { SafetyCheckInput } from '../protocol.js';
|
||||
import {
|
||||
logConsecaPolicyGeneration,
|
||||
logConsecaVerdict,
|
||||
} from '../../telemetry/index.js';
|
||||
import type { Config } from '../../config/config.js';
|
||||
import * as policyGenerator from './policy-generator.js';
|
||||
import * as policyEnforcer from './policy-enforcer.js';
|
||||
|
||||
vi.mock('../../telemetry/index.js', () => ({
|
||||
logConsecaPolicyGeneration: vi.fn(),
|
||||
ConsecaPolicyGenerationEvent: vi.fn(),
|
||||
logConsecaVerdict: vi.fn(),
|
||||
ConsecaVerdictEvent: vi.fn(),
|
||||
}));
|
||||
|
||||
vi.mock('./policy-generator.js');
|
||||
vi.mock('./policy-enforcer.js');
|
||||
|
||||
describe('ConsecaSafetyChecker', () => {
|
||||
let checker: ConsecaSafetyChecker;
|
||||
let mockConfig: Config;
|
||||
|
||||
beforeEach(() => {
|
||||
// Reset singleton instance to ensure clean state
|
||||
ConsecaSafetyChecker.resetInstance();
|
||||
// Get the fresh singleton instance
|
||||
checker = ConsecaSafetyChecker.getInstance();
|
||||
|
||||
mockConfig = {
|
||||
enableConseca: true,
|
||||
getToolRegistry: vi.fn().mockReturnValue({
|
||||
getFunctionDeclarations: vi.fn().mockReturnValue([]),
|
||||
}),
|
||||
} as unknown as Config;
|
||||
checker.setConfig(mockConfig);
|
||||
vi.clearAllMocks();
|
||||
|
||||
// Default mock implementations
|
||||
vi.mocked(policyGenerator.generatePolicy).mockResolvedValue({ policy: {} });
|
||||
vi.mocked(policyEnforcer.enforcePolicy).mockResolvedValue({
|
||||
decision: SafetyCheckDecision.ALLOW,
|
||||
});
|
||||
});
|
||||
|
||||
it('should be a singleton', () => {
|
||||
const instance1 = ConsecaSafetyChecker.getInstance();
|
||||
const instance2 = ConsecaSafetyChecker.getInstance();
|
||||
expect(instance1).toBe(instance2);
|
||||
});
|
||||
|
||||
it('should return ALLOW when no user prompt is present in context', async () => {
|
||||
const input: SafetyCheckInput = {
|
||||
protocolVersion: '1.0.0',
|
||||
toolCall: { name: 'testTool' },
|
||||
context: {
|
||||
environment: { cwd: '/tmp', workspaces: [] },
|
||||
},
|
||||
};
|
||||
|
||||
const result = await checker.check(input);
|
||||
expect(result.decision).toBe(SafetyCheckDecision.ALLOW);
|
||||
});
|
||||
|
||||
it('should return ALLOW if enableConseca is false', async () => {
|
||||
const disabledConfig = {
|
||||
enableConseca: false,
|
||||
} as unknown as Config;
|
||||
checker.setConfig(disabledConfig);
|
||||
|
||||
const input: SafetyCheckInput = {
|
||||
protocolVersion: '1.0.0',
|
||||
toolCall: { name: 'testTool' },
|
||||
context: {
|
||||
environment: { cwd: '/tmp', workspaces: [] },
|
||||
},
|
||||
};
|
||||
|
||||
const result = await checker.check(input);
|
||||
expect(result.decision).toBe(SafetyCheckDecision.ALLOW);
|
||||
expect(result.reason).toBe('Conseca is disabled');
|
||||
expect(policyGenerator.generatePolicy).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('getPolicy should return cached policy if user prompt matches', async () => {
|
||||
const mockPolicy = {
|
||||
tool: {
|
||||
permissions: SafetyCheckDecision.ALLOW,
|
||||
constraints: 'None',
|
||||
rationale: 'Test',
|
||||
},
|
||||
};
|
||||
vi.mocked(policyGenerator.generatePolicy).mockResolvedValue({
|
||||
policy: mockPolicy,
|
||||
});
|
||||
|
||||
const policy1 = await checker.getPolicy('prompt', 'trusted', mockConfig);
|
||||
const policy2 = await checker.getPolicy('prompt', 'trusted', mockConfig);
|
||||
|
||||
expect(policy1).toBe(mockPolicy);
|
||||
expect(policy2).toBe(mockPolicy);
|
||||
expect(policyGenerator.generatePolicy).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
it('getPolicy should generate new policy if user prompt changes', async () => {
|
||||
const mockPolicy1 = {
|
||||
tool1: {
|
||||
permissions: SafetyCheckDecision.ALLOW,
|
||||
constraints: 'None',
|
||||
rationale: 'Test',
|
||||
},
|
||||
};
|
||||
const mockPolicy2 = {
|
||||
tool2: {
|
||||
permissions: SafetyCheckDecision.ALLOW,
|
||||
constraints: 'None',
|
||||
rationale: 'Test',
|
||||
},
|
||||
};
|
||||
vi.mocked(policyGenerator.generatePolicy)
|
||||
.mockResolvedValueOnce({ policy: mockPolicy1 })
|
||||
.mockResolvedValueOnce({ policy: mockPolicy2 });
|
||||
|
||||
const policy1 = await checker.getPolicy('prompt1', 'trusted', mockConfig);
|
||||
const policy2 = await checker.getPolicy('prompt2', 'trusted', mockConfig);
|
||||
|
||||
expect(policy1).toBe(mockPolicy1);
|
||||
expect(policy2).toBe(mockPolicy2);
|
||||
expect(policyGenerator.generatePolicy).toHaveBeenCalledTimes(2);
|
||||
});
|
||||
|
||||
it('check should call getPolicy and enforcePolicy', async () => {
|
||||
const mockPolicy = {
|
||||
tool: {
|
||||
permissions: SafetyCheckDecision.ALLOW,
|
||||
constraints: 'None',
|
||||
rationale: 'Test',
|
||||
},
|
||||
};
|
||||
vi.mocked(policyGenerator.generatePolicy).mockResolvedValue({
|
||||
policy: mockPolicy,
|
||||
});
|
||||
vi.mocked(policyEnforcer.enforcePolicy).mockResolvedValue({
|
||||
decision: SafetyCheckDecision.ALLOW,
|
||||
});
|
||||
|
||||
const input: SafetyCheckInput = {
|
||||
protocolVersion: '1.0.0',
|
||||
toolCall: { name: 'tool', args: {} },
|
||||
context: {
|
||||
environment: { cwd: '.', workspaces: [] },
|
||||
history: {
|
||||
turns: [
|
||||
{
|
||||
user: { text: 'user prompt' },
|
||||
model: {},
|
||||
},
|
||||
],
|
||||
},
|
||||
},
|
||||
};
|
||||
|
||||
const result = await checker.check(input);
|
||||
|
||||
expect(policyGenerator.generatePolicy).toHaveBeenCalledWith(
|
||||
'user prompt',
|
||||
expect.any(String),
|
||||
mockConfig,
|
||||
);
|
||||
expect(policyEnforcer.enforcePolicy).toHaveBeenCalledWith(
|
||||
mockPolicy,
|
||||
input.toolCall,
|
||||
mockConfig,
|
||||
);
|
||||
expect(result.decision).toBe(SafetyCheckDecision.ALLOW);
|
||||
});
|
||||
|
||||
it('check should return ALLOW if no user prompt found (fallback)', async () => {
|
||||
const input: SafetyCheckInput = {
|
||||
protocolVersion: '1.0.0',
|
||||
toolCall: { name: 'tool', args: {} },
|
||||
context: {
|
||||
environment: { cwd: '.', workspaces: [] },
|
||||
},
|
||||
};
|
||||
|
||||
const result = await checker.check(input);
|
||||
|
||||
expect(policyGenerator.generatePolicy).not.toHaveBeenCalled();
|
||||
expect(result.decision).toBe(SafetyCheckDecision.ALLOW);
|
||||
});
|
||||
|
||||
// Test state helpers
|
||||
it('should expose current state via helpers', async () => {
|
||||
const mockPolicy = {
|
||||
tool: {
|
||||
permissions: SafetyCheckDecision.ALLOW,
|
||||
constraints: 'None',
|
||||
rationale: 'Test',
|
||||
},
|
||||
};
|
||||
vi.mocked(policyGenerator.generatePolicy).mockResolvedValue({
|
||||
policy: mockPolicy,
|
||||
});
|
||||
|
||||
await checker.getPolicy('prompt', 'trusted', mockConfig);
|
||||
|
||||
expect(checker.getCurrentPolicy()).toBe(mockPolicy);
|
||||
expect(checker.getActiveUserPrompt()).toBe('prompt');
|
||||
});
|
||||
it('should log policy generation event when config is set', async () => {
|
||||
const mockPolicy = {
|
||||
tool: {
|
||||
permissions: SafetyCheckDecision.ALLOW,
|
||||
constraints: 'None',
|
||||
rationale: 'Test',
|
||||
},
|
||||
};
|
||||
vi.mocked(policyGenerator.generatePolicy).mockResolvedValue({
|
||||
policy: mockPolicy,
|
||||
});
|
||||
|
||||
await checker.getPolicy('telemetry_prompt', 'trusted', mockConfig);
|
||||
|
||||
expect(logConsecaPolicyGeneration).toHaveBeenCalledWith(
|
||||
mockConfig,
|
||||
expect.anything(),
|
||||
);
|
||||
});
|
||||
|
||||
it('should log verdict event on check', async () => {
|
||||
const mockPolicy = {
|
||||
tool: {
|
||||
permissions: SafetyCheckDecision.ALLOW,
|
||||
constraints: 'None',
|
||||
rationale: 'Test',
|
||||
},
|
||||
};
|
||||
vi.mocked(policyGenerator.generatePolicy).mockResolvedValue({
|
||||
policy: mockPolicy,
|
||||
});
|
||||
vi.mocked(policyEnforcer.enforcePolicy).mockResolvedValue({
|
||||
decision: SafetyCheckDecision.ALLOW,
|
||||
reason: 'Allowed by policy',
|
||||
});
|
||||
|
||||
const input: SafetyCheckInput = {
|
||||
protocolVersion: '1.0.0',
|
||||
toolCall: { name: 'tool', args: {} },
|
||||
context: {
|
||||
environment: { cwd: '.', workspaces: [] },
|
||||
history: {
|
||||
turns: [
|
||||
{
|
||||
user: { text: 'user prompt' },
|
||||
model: {},
|
||||
},
|
||||
],
|
||||
},
|
||||
},
|
||||
};
|
||||
|
||||
await checker.check(input);
|
||||
|
||||
expect(logConsecaVerdict).toHaveBeenCalledWith(
|
||||
mockConfig,
|
||||
expect.anything(),
|
||||
);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,170 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import type { InProcessChecker } from '../built-in.js';
|
||||
import type { SafetyCheckInput, SafetyCheckResult } from '../protocol.js';
|
||||
import { SafetyCheckDecision } from '../protocol.js';
|
||||
|
||||
import {
|
||||
logConsecaPolicyGeneration,
|
||||
ConsecaPolicyGenerationEvent,
|
||||
logConsecaVerdict,
|
||||
ConsecaVerdictEvent,
|
||||
} from '../../telemetry/index.js';
|
||||
import { debugLogger } from '../../utils/debugLogger.js';
|
||||
import type { Config } from '../../config/config.js';
|
||||
|
||||
import { generatePolicy } from './policy-generator.js';
|
||||
import { enforcePolicy } from './policy-enforcer.js';
|
||||
import type { SecurityPolicy } from './types.js';
|
||||
|
||||
export class ConsecaSafetyChecker implements InProcessChecker {
|
||||
private static instance: ConsecaSafetyChecker | undefined;
|
||||
private currentPolicy: SecurityPolicy | null = null;
|
||||
private activeUserPrompt: string | null = null;
|
||||
private config: Config | null = null;
|
||||
|
||||
/**
|
||||
* Private constructor to enforce singleton pattern.
|
||||
* Use `getInstance()` to access the instance.
|
||||
*/
|
||||
private constructor() {}
|
||||
|
||||
static getInstance(): ConsecaSafetyChecker {
|
||||
if (!ConsecaSafetyChecker.instance) {
|
||||
ConsecaSafetyChecker.instance = new ConsecaSafetyChecker();
|
||||
}
|
||||
return ConsecaSafetyChecker.instance;
|
||||
}
|
||||
|
||||
/**
|
||||
* Resets the singleton instance. Use only in tests.
|
||||
*/
|
||||
static resetInstance(): void {
|
||||
ConsecaSafetyChecker.instance = undefined;
|
||||
}
|
||||
|
||||
setConfig(config: Config): void {
|
||||
this.config = config;
|
||||
}
|
||||
|
||||
async check(input: SafetyCheckInput): Promise<SafetyCheckResult> {
|
||||
debugLogger.debug(
|
||||
`[Conseca] check called. History is: ${JSON.stringify(input.context.history)}`,
|
||||
);
|
||||
|
||||
if (!this.config) {
|
||||
debugLogger.debug('[Conseca] check failed: Config not initialized');
|
||||
return {
|
||||
decision: SafetyCheckDecision.ALLOW,
|
||||
reason: 'Config not initialized',
|
||||
};
|
||||
}
|
||||
|
||||
if (!this.config.enableConseca) {
|
||||
debugLogger.debug('[Conseca] check skipped: Conseca is not enabled.');
|
||||
return {
|
||||
decision: SafetyCheckDecision.ALLOW,
|
||||
reason: 'Conseca is disabled',
|
||||
};
|
||||
}
|
||||
|
||||
const userPrompt = this.extractUserPrompt(input);
|
||||
let trustedContent = '';
|
||||
|
||||
const toolRegistry = this.config.getToolRegistry();
|
||||
if (toolRegistry) {
|
||||
const tools = toolRegistry.getFunctionDeclarations();
|
||||
trustedContent = JSON.stringify(tools, null, 2);
|
||||
}
|
||||
|
||||
if (userPrompt) {
|
||||
await this.getPolicy(userPrompt, trustedContent, this.config);
|
||||
} else {
|
||||
debugLogger.debug(
|
||||
`[Conseca] Skipping policy generation because userPrompt is null`,
|
||||
);
|
||||
}
|
||||
|
||||
let result: SafetyCheckResult;
|
||||
|
||||
if (!this.currentPolicy) {
|
||||
result = {
|
||||
decision: SafetyCheckDecision.ALLOW, // Fallback if no policy generated yet
|
||||
reason: 'No security policy generated.',
|
||||
error: 'No security policy generated.',
|
||||
};
|
||||
} else {
|
||||
result = await enforcePolicy(
|
||||
this.currentPolicy,
|
||||
input.toolCall,
|
||||
this.config,
|
||||
);
|
||||
}
|
||||
|
||||
logConsecaVerdict(
|
||||
this.config,
|
||||
new ConsecaVerdictEvent(
|
||||
userPrompt || '',
|
||||
JSON.stringify(this.currentPolicy || {}),
|
||||
JSON.stringify(input.toolCall),
|
||||
result.decision,
|
||||
result.reason || '',
|
||||
'error' in result ? result.error : undefined,
|
||||
),
|
||||
);
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
async getPolicy(
|
||||
userPrompt: string,
|
||||
trustedContent: string,
|
||||
config: Config,
|
||||
): Promise<SecurityPolicy> {
|
||||
if (this.activeUserPrompt === userPrompt && this.currentPolicy) {
|
||||
return this.currentPolicy;
|
||||
}
|
||||
|
||||
const { policy, error } = await generatePolicy(
|
||||
userPrompt,
|
||||
trustedContent,
|
||||
config,
|
||||
);
|
||||
this.currentPolicy = policy;
|
||||
this.activeUserPrompt = userPrompt;
|
||||
|
||||
logConsecaPolicyGeneration(
|
||||
config,
|
||||
new ConsecaPolicyGenerationEvent(
|
||||
userPrompt,
|
||||
trustedContent,
|
||||
JSON.stringify(policy),
|
||||
error,
|
||||
),
|
||||
);
|
||||
|
||||
return policy;
|
||||
}
|
||||
|
||||
private extractUserPrompt(input: SafetyCheckInput): string | null {
|
||||
const prompt = input.context.history?.turns.at(-1)?.user.text;
|
||||
if (prompt) {
|
||||
return prompt;
|
||||
}
|
||||
debugLogger.debug(`[Conseca] extractUserPrompt failed.`);
|
||||
return null;
|
||||
}
|
||||
|
||||
// Helper methods for testing state
|
||||
getCurrentPolicy(): SecurityPolicy | null {
|
||||
return this.currentPolicy;
|
||||
}
|
||||
|
||||
getActiveUserPrompt(): string | null {
|
||||
return this.activeUserPrompt;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,21 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { describe, it, expect } from 'vitest';
|
||||
import { ConsecaSafetyChecker } from './conseca.js';
|
||||
import { InProcessCheckerType } from '../../policy/types.js';
|
||||
import { CheckerRegistry } from '../registry.js';
|
||||
|
||||
describe('Conseca Integration', () => {
|
||||
it('should be registered and resolvable via CheckerRegistry', () => {
|
||||
const registry = new CheckerRegistry('.');
|
||||
const checker = registry.resolveInProcess(InProcessCheckerType.CONSECA);
|
||||
|
||||
expect(checker).toBeDefined();
|
||||
expect(checker).toBeInstanceOf(ConsecaSafetyChecker);
|
||||
expect(checker).toBe(ConsecaSafetyChecker.getInstance());
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,167 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { describe, it, expect, vi, beforeEach } from 'vitest';
|
||||
import { enforcePolicy } from './policy-enforcer.js';
|
||||
import type { Config } from '../../config/config.js';
|
||||
import type { ContentGenerator } from '../../core/contentGenerator.js';
|
||||
import { SafetyCheckDecision } from '../protocol.js';
|
||||
import type { FunctionCall } from '@google/genai';
|
||||
import { LlmRole } from '../../telemetry/index.js';
|
||||
|
||||
describe('policy_enforcer', () => {
|
||||
let mockConfig: Config;
|
||||
let mockContentGenerator: ContentGenerator;
|
||||
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
mockContentGenerator = {
|
||||
generateContent: vi.fn(),
|
||||
} as unknown as ContentGenerator;
|
||||
|
||||
mockConfig = {
|
||||
getContentGenerator: vi.fn().mockReturnValue(mockContentGenerator),
|
||||
} as unknown as Config;
|
||||
});
|
||||
|
||||
it('should return ALLOW when content generator returns ALLOW', async () => {
|
||||
mockContentGenerator.generateContent = vi.fn().mockResolvedValue({
|
||||
candidates: [
|
||||
{
|
||||
content: {
|
||||
parts: [
|
||||
{ text: JSON.stringify({ decision: 'allow', reason: 'Safe' }) },
|
||||
],
|
||||
},
|
||||
},
|
||||
],
|
||||
});
|
||||
|
||||
const toolCall: FunctionCall = { name: 'testTool', args: {} };
|
||||
const policy = {
|
||||
testTool: {
|
||||
permissions: SafetyCheckDecision.ALLOW,
|
||||
constraints: 'None',
|
||||
rationale: 'Test',
|
||||
},
|
||||
};
|
||||
const result = await enforcePolicy(policy, toolCall, mockConfig);
|
||||
|
||||
expect(mockConfig.getContentGenerator).toHaveBeenCalled();
|
||||
expect(mockContentGenerator.generateContent).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
model: expect.any(String),
|
||||
config: expect.objectContaining({
|
||||
responseMimeType: 'application/json',
|
||||
responseSchema: expect.any(Object),
|
||||
}),
|
||||
contents: expect.arrayContaining([
|
||||
expect.objectContaining({
|
||||
role: 'user',
|
||||
parts: expect.arrayContaining([
|
||||
expect.objectContaining({
|
||||
text: expect.stringContaining('Security Policy:'),
|
||||
}),
|
||||
]),
|
||||
}),
|
||||
]),
|
||||
}),
|
||||
'conseca-policy-enforcement',
|
||||
LlmRole.SUBAGENT,
|
||||
);
|
||||
expect(result.decision).toBe(SafetyCheckDecision.ALLOW);
|
||||
});
|
||||
|
||||
it('should handle missing content generator gracefully (error case)', async () => {
|
||||
vi.mocked(mockConfig.getContentGenerator).mockReturnValue(
|
||||
undefined as unknown as ContentGenerator,
|
||||
);
|
||||
|
||||
const toolCall: FunctionCall = { name: 'testTool', args: {} };
|
||||
const policy = {
|
||||
testTool: {
|
||||
permissions: SafetyCheckDecision.ALLOW,
|
||||
constraints: 'None',
|
||||
rationale: 'Test',
|
||||
},
|
||||
};
|
||||
const result = await enforcePolicy(policy, toolCall, mockConfig);
|
||||
|
||||
expect(result.decision).toBe(SafetyCheckDecision.ALLOW);
|
||||
});
|
||||
|
||||
it('should ALLOW if tool name is missing with the reason and error as tool name is missing', async () => {
|
||||
const toolCall = { args: {} } as FunctionCall;
|
||||
const policy = {};
|
||||
const result = await enforcePolicy(policy, toolCall, mockConfig);
|
||||
|
||||
expect(result.decision).toBe(SafetyCheckDecision.ALLOW);
|
||||
expect(result.reason).toBe('Tool name is missing');
|
||||
if (result.decision === SafetyCheckDecision.ALLOW) {
|
||||
expect(result.error).toBe('Tool name is missing');
|
||||
}
|
||||
});
|
||||
|
||||
it('should handle empty policy by checking with LLM (fail-open/check behavior)', async () => {
|
||||
// Even if policy is empty for the tool, we currently send it to LLM.
|
||||
// The LLM might ALLOW or DENY based on its own judgment of "no policy".
|
||||
// We simulate the LLM allowing the action to match the current fail-open strategy.
|
||||
mockContentGenerator.generateContent = vi.fn().mockResolvedValue({
|
||||
candidates: [
|
||||
{
|
||||
content: {
|
||||
parts: [
|
||||
{
|
||||
text: JSON.stringify({
|
||||
decision: 'allow',
|
||||
reason: 'No restrictions',
|
||||
}),
|
||||
},
|
||||
],
|
||||
},
|
||||
},
|
||||
],
|
||||
});
|
||||
|
||||
const toolCall: FunctionCall = { name: 'unknownTool', args: {} };
|
||||
const policy = {}; // Empty policy
|
||||
const result = await enforcePolicy(policy, toolCall, mockConfig);
|
||||
|
||||
expect(result.decision).toBe(SafetyCheckDecision.ALLOW);
|
||||
expect(mockContentGenerator.generateContent).toHaveBeenCalled();
|
||||
if (result.decision === SafetyCheckDecision.ALLOW) {
|
||||
expect(result.error).toBeUndefined();
|
||||
}
|
||||
});
|
||||
|
||||
it('should handle malformed JSON response from LLM by failing open (ALLOW)', async () => {
|
||||
mockContentGenerator.generateContent = vi.fn().mockResolvedValue({
|
||||
candidates: [
|
||||
{
|
||||
content: {
|
||||
parts: [{ text: 'This is not JSON' }],
|
||||
},
|
||||
},
|
||||
],
|
||||
});
|
||||
|
||||
const toolCall: FunctionCall = { name: 'testTool', args: {} };
|
||||
const policy = {
|
||||
testTool: {
|
||||
permissions: SafetyCheckDecision.ALLOW,
|
||||
constraints: 'None',
|
||||
rationale: 'Test',
|
||||
},
|
||||
};
|
||||
const result = await enforcePolicy(policy, toolCall, mockConfig);
|
||||
|
||||
expect(result.decision).toBe(SafetyCheckDecision.ALLOW);
|
||||
expect(result.reason).toContain('JSON Parse Error');
|
||||
if (result.decision === SafetyCheckDecision.ALLOW) {
|
||||
expect(result.error).toContain('JSON Parse Error');
|
||||
}
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,164 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import type { Config } from '../../config/config.js';
|
||||
import type { FunctionCall } from '@google/genai';
|
||||
import { SafetyCheckDecision, type SafetyCheckResult } from '../protocol.js';
|
||||
import type { SecurityPolicy } from './types.js';
|
||||
import { getResponseText } from '../../utils/partUtils.js';
|
||||
import { safeTemplateReplace } from '../../utils/textUtils.js';
|
||||
|
||||
import { DEFAULT_GEMINI_FLASH_MODEL } from '../../config/models.js';
|
||||
import { debugLogger } from '../../utils/debugLogger.js';
|
||||
|
||||
import { LlmRole } from '../../telemetry/index.js';
|
||||
|
||||
const CONSECA_ENFORCEMENT_PROMPT = `
|
||||
You are a security enforcement engine. Your goal is to check if a specific tool call complies with a given security policy.
|
||||
|
||||
Input:
|
||||
1. **Security Policy:** A set of rules defining allowed and denied actions for this specific tool.
|
||||
2. **Tool Call:** The actual function call the system intends to execute.
|
||||
|
||||
Security Policy:
|
||||
{{policy}}
|
||||
|
||||
Tool Call:
|
||||
{{tool_call}}
|
||||
|
||||
Evaluate the tool call against the policy.
|
||||
1. Check if the tool is allowed.
|
||||
2. Check if the arguments match the constraints.
|
||||
3. Output a JSON object with:
|
||||
- "decision": "allow", "deny", or "ask_user".
|
||||
- "reason": A brief explanation.
|
||||
|
||||
Output strictly JSON.
|
||||
`;
|
||||
|
||||
import { z } from 'zod';
|
||||
import { zodToJsonSchema } from 'zod-to-json-schema';
|
||||
|
||||
const EnforcementResultSchema = z.object({
|
||||
decision: z.enum(['allow', 'deny', 'ask_user']),
|
||||
reason: z.string(),
|
||||
});
|
||||
|
||||
/**
|
||||
* Enforces the security policy for a given tool call.
|
||||
*/
|
||||
export async function enforcePolicy(
|
||||
policy: SecurityPolicy,
|
||||
toolCall: FunctionCall,
|
||||
config: Config,
|
||||
): Promise<SafetyCheckResult> {
|
||||
const model = DEFAULT_GEMINI_FLASH_MODEL;
|
||||
const contentGenerator = config.getContentGenerator();
|
||||
|
||||
if (!contentGenerator) {
|
||||
return {
|
||||
decision: SafetyCheckDecision.ALLOW,
|
||||
reason: 'Content generator not initialized',
|
||||
error: 'Content generator not initialized',
|
||||
};
|
||||
}
|
||||
|
||||
const toolName = toolCall.name;
|
||||
// If tool name is missing, we cannot enforce the policy. Allow by default.
|
||||
if (!toolName) {
|
||||
return {
|
||||
decision: SafetyCheckDecision.ALLOW,
|
||||
reason: 'Tool name is missing',
|
||||
error: 'Tool name is missing',
|
||||
};
|
||||
}
|
||||
|
||||
const toolPolicyStr = JSON.stringify(policy[toolName] || {}, null, 2);
|
||||
const toolCallStr = JSON.stringify(toolCall, null, 2);
|
||||
debugLogger.debug(
|
||||
`[Conseca] Enforcing policy for tool: ${toolName}`,
|
||||
toolCall,
|
||||
toolPolicyStr,
|
||||
toolCallStr,
|
||||
);
|
||||
|
||||
try {
|
||||
const result = await contentGenerator.generateContent(
|
||||
{
|
||||
model,
|
||||
config: {
|
||||
responseMimeType: 'application/json',
|
||||
responseSchema: zodToJsonSchema(EnforcementResultSchema, {
|
||||
target: 'openApi3',
|
||||
}),
|
||||
},
|
||||
contents: [
|
||||
{
|
||||
role: 'user',
|
||||
parts: [
|
||||
{
|
||||
text: safeTemplateReplace(CONSECA_ENFORCEMENT_PROMPT, {
|
||||
policy: toolPolicyStr,
|
||||
tool_call: toolCallStr,
|
||||
}),
|
||||
},
|
||||
],
|
||||
},
|
||||
],
|
||||
},
|
||||
'conseca-policy-enforcement',
|
||||
LlmRole.SUBAGENT,
|
||||
);
|
||||
|
||||
const responseText = getResponseText(result);
|
||||
debugLogger.debug(`[Conseca] Enforcement Raw Response: ${responseText}`);
|
||||
|
||||
if (!responseText) {
|
||||
return {
|
||||
decision: SafetyCheckDecision.ALLOW,
|
||||
reason: 'Empty response from policy enforcer',
|
||||
error: 'Empty response from policy enforcer',
|
||||
};
|
||||
}
|
||||
|
||||
try {
|
||||
const parsed = EnforcementResultSchema.parse(JSON.parse(responseText));
|
||||
debugLogger.debug(`[Conseca] Enforcement Parsed:`, parsed);
|
||||
|
||||
let decision: SafetyCheckDecision;
|
||||
switch (parsed.decision) {
|
||||
case 'allow':
|
||||
decision = SafetyCheckDecision.ALLOW;
|
||||
break;
|
||||
case 'ask_user':
|
||||
decision = SafetyCheckDecision.ASK_USER;
|
||||
break;
|
||||
case 'deny':
|
||||
default:
|
||||
decision = SafetyCheckDecision.DENY;
|
||||
break;
|
||||
}
|
||||
|
||||
return {
|
||||
decision,
|
||||
reason: parsed.reason,
|
||||
};
|
||||
} catch (parseError) {
|
||||
return {
|
||||
decision: SafetyCheckDecision.ALLOW,
|
||||
reason: 'JSON Parse Error in enforcement response',
|
||||
error: `JSON Parse Error: ${parseError instanceof Error ? parseError.message : String(parseError)}. Raw: ${responseText}`,
|
||||
};
|
||||
}
|
||||
} catch (error) {
|
||||
debugLogger.error('Policy enforcement failed:', error);
|
||||
return {
|
||||
decision: SafetyCheckDecision.ALLOW,
|
||||
reason: 'Policy enforcement failed',
|
||||
error: `Policy enforcement failed: ${error instanceof Error ? error.message : String(error)}`,
|
||||
};
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,116 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { describe, it, expect, vi, beforeEach } from 'vitest';
|
||||
import { generatePolicy } from './policy-generator.js';
|
||||
import { SafetyCheckDecision } from '../protocol.js';
|
||||
import type { Config } from '../../config/config.js';
|
||||
import type { ContentGenerator } from '../../core/contentGenerator.js';
|
||||
import { LlmRole } from '../../telemetry/index.js';
|
||||
|
||||
describe('policy_generator', () => {
|
||||
let mockConfig: Config;
|
||||
let mockContentGenerator: ContentGenerator;
|
||||
|
||||
beforeEach(() => {
|
||||
mockContentGenerator = {
|
||||
generateContent: vi.fn(),
|
||||
} as unknown as ContentGenerator;
|
||||
|
||||
mockConfig = {
|
||||
getContentGenerator: vi.fn().mockReturnValue(mockContentGenerator),
|
||||
} as unknown as Config;
|
||||
});
|
||||
|
||||
it('should return a policy object when content generator is available', async () => {
|
||||
const mockPolicy = {
|
||||
read_file: {
|
||||
permissions: SafetyCheckDecision.ALLOW,
|
||||
constraints: 'None',
|
||||
rationale: 'Test',
|
||||
},
|
||||
};
|
||||
mockContentGenerator.generateContent = vi.fn().mockResolvedValue({
|
||||
candidates: [
|
||||
{
|
||||
content: {
|
||||
parts: [
|
||||
{
|
||||
text: JSON.stringify({
|
||||
policies: [
|
||||
{
|
||||
tool_name: 'read_file',
|
||||
policy: mockPolicy.read_file,
|
||||
},
|
||||
],
|
||||
}),
|
||||
},
|
||||
],
|
||||
},
|
||||
},
|
||||
],
|
||||
});
|
||||
|
||||
const result = await generatePolicy(
|
||||
'test prompt',
|
||||
'trusted content',
|
||||
mockConfig,
|
||||
);
|
||||
|
||||
expect(mockConfig.getContentGenerator).toHaveBeenCalled();
|
||||
expect(mockContentGenerator.generateContent).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
model: expect.any(String),
|
||||
config: expect.objectContaining({
|
||||
responseMimeType: 'application/json',
|
||||
responseSchema: expect.any(Object),
|
||||
}),
|
||||
contents: expect.any(Array),
|
||||
}),
|
||||
'conseca-policy-generation',
|
||||
LlmRole.SUBAGENT,
|
||||
);
|
||||
expect(result.policy).toEqual(mockPolicy);
|
||||
expect(result.error).toBeUndefined();
|
||||
});
|
||||
|
||||
it('should handle missing content generator gracefully', async () => {
|
||||
vi.mocked(mockConfig.getContentGenerator).mockReturnValue(
|
||||
undefined as unknown as ContentGenerator,
|
||||
);
|
||||
|
||||
const result = await generatePolicy(
|
||||
'test prompt',
|
||||
'trusted content',
|
||||
mockConfig,
|
||||
);
|
||||
|
||||
expect(result.policy).toEqual({});
|
||||
expect(result.error).toBe('Content generator not initialized');
|
||||
});
|
||||
it('should prevent template injection (double interpolation)', async () => {
|
||||
mockContentGenerator.generateContent = vi.fn().mockResolvedValue({});
|
||||
|
||||
const userPrompt = '{{trusted_content}}';
|
||||
const trustedContent = 'SECRET_DATA';
|
||||
|
||||
await generatePolicy(userPrompt, trustedContent, mockConfig);
|
||||
|
||||
const generateContentCall = vi.mocked(mockContentGenerator.generateContent)
|
||||
.mock.calls[0];
|
||||
const request = generateContentCall[0] as {
|
||||
contents: Array<{ parts: Array<{ text: string }> }>;
|
||||
};
|
||||
const promptText = request.contents[0].parts[0].text;
|
||||
|
||||
// The user prompt should contain the literal placeholder, NOT the secret data
|
||||
expect(promptText).toContain('User Prompt: "{{trusted_content}}"');
|
||||
expect(promptText).not.toContain('User Prompt: "SECRET_DATA"');
|
||||
|
||||
// The trusted tools section SHOULD contain the secret data
|
||||
expect(promptText).toContain('Trusted Tools (Context):\nSECRET_DATA');
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,178 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import type { Config } from '../../config/config.js';
|
||||
import type { SecurityPolicy } from './types.js';
|
||||
import { getResponseText } from '../../utils/partUtils.js';
|
||||
import { safeTemplateReplace } from '../../utils/textUtils.js';
|
||||
import { DEFAULT_GEMINI_FLASH_MODEL } from '../../config/models.js';
|
||||
import { debugLogger } from '../../utils/debugLogger.js';
|
||||
import { SafetyCheckDecision } from '../protocol.js';
|
||||
|
||||
import { LlmRole } from '../../telemetry/index.js';
|
||||
|
||||
const CONSECA_POLICY_GENERATION_PROMPT = `
|
||||
You are a security expert responsible for generating fine-grained security policies for a large language model integrated into a command-line tool. Your role is to act as a "policy generator" that creates temporary, context-specific rules based on a user's prompt and the tools available to the main LLM.
|
||||
|
||||
Your primary goal is to enforce the principle of least privilege. The policies you create should be as restrictive as possible while still allowing the main LLM to complete the user's requested task.
|
||||
|
||||
For each tool that is relevant to the user's prompt, you must generate a policy object.
|
||||
|
||||
### Output Format
|
||||
You must return a JSON object with a "policies" key, which is an array of objects. Each object must have:
|
||||
- "tool_name": The name of the tool.
|
||||
- "policy": An object with:
|
||||
- "permissions": "allow" | "deny" | "ask_user"
|
||||
- "constraints": A detailed description of conditions (e.g. allowed files, arguments).
|
||||
- "rationale": Explanation for the policy.
|
||||
|
||||
Example JSON:
|
||||
\`\`\`json
|
||||
{
|
||||
"policies": [
|
||||
{
|
||||
"tool_name": "read_file",
|
||||
"policy": {
|
||||
"permissions": "allow",
|
||||
"constraints": "Only allow reading 'main.py'.",
|
||||
"rationale": "User asked to read main.py"
|
||||
}
|
||||
},
|
||||
{
|
||||
"tool_name": "run_shell_command",
|
||||
"policy": {
|
||||
"permissions": "deny",
|
||||
"constraints": "None",
|
||||
"rationale": "Shell commands are not needed for this task"
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
\`\`\`
|
||||
|
||||
### Guiding Principles:
|
||||
1. **Permissions:**
|
||||
* **allow:** Required tools for the task.
|
||||
* **deny:** Tools clearly outside the scope.
|
||||
* **ask_user:** Destructive actions or ambiguity.
|
||||
|
||||
2. **Constraints:**
|
||||
* Be specific! Restrict file paths, command arguments, etc.
|
||||
|
||||
3. **Rationale:**
|
||||
* Reference the user's prompt.
|
||||
|
||||
User Prompt: "{{user_prompt}}"
|
||||
|
||||
Trusted Tools (Context):
|
||||
{{trusted_content}}
|
||||
`;
|
||||
|
||||
import { z } from 'zod';
|
||||
import { zodToJsonSchema } from 'zod-to-json-schema';
|
||||
|
||||
const ToolPolicySchema = z.object({
|
||||
permissions: z.nativeEnum(SafetyCheckDecision),
|
||||
constraints: z.string(),
|
||||
rationale: z.string(),
|
||||
});
|
||||
|
||||
const SecurityPolicyResponseSchema = z.object({
|
||||
policies: z.array(
|
||||
z.object({
|
||||
tool_name: z.string(),
|
||||
policy: ToolPolicySchema,
|
||||
}),
|
||||
),
|
||||
});
|
||||
|
||||
export interface PolicyGenerationResult {
|
||||
policy: SecurityPolicy;
|
||||
error?: string;
|
||||
}
|
||||
|
||||
/**
|
||||
* Generates a security policy for the given user prompt and trusted content.
|
||||
*/
|
||||
export async function generatePolicy(
|
||||
userPrompt: string,
|
||||
trustedContent: string,
|
||||
config: Config,
|
||||
): Promise<PolicyGenerationResult> {
|
||||
const model = DEFAULT_GEMINI_FLASH_MODEL;
|
||||
const contentGenerator = config.getContentGenerator();
|
||||
|
||||
if (!contentGenerator) {
|
||||
return { policy: {}, error: 'Content generator not initialized' };
|
||||
}
|
||||
|
||||
try {
|
||||
const result = await contentGenerator.generateContent(
|
||||
{
|
||||
model,
|
||||
config: {
|
||||
responseMimeType: 'application/json',
|
||||
responseSchema: zodToJsonSchema(SecurityPolicyResponseSchema, {
|
||||
target: 'openApi3',
|
||||
}),
|
||||
},
|
||||
contents: [
|
||||
{
|
||||
role: 'user',
|
||||
parts: [
|
||||
{
|
||||
text: safeTemplateReplace(CONSECA_POLICY_GENERATION_PROMPT, {
|
||||
user_prompt: userPrompt,
|
||||
trusted_content: trustedContent,
|
||||
}),
|
||||
},
|
||||
],
|
||||
},
|
||||
],
|
||||
},
|
||||
'conseca-policy-generation',
|
||||
LlmRole.SUBAGENT,
|
||||
);
|
||||
|
||||
const responseText = getResponseText(result);
|
||||
debugLogger.debug(
|
||||
`[Conseca] Policy Generation Raw Response: ${responseText}`,
|
||||
);
|
||||
|
||||
if (!responseText) {
|
||||
return { policy: {}, error: 'Empty response from policy generator' };
|
||||
}
|
||||
|
||||
try {
|
||||
const parsed = SecurityPolicyResponseSchema.parse(
|
||||
JSON.parse(responseText),
|
||||
);
|
||||
const policiesList = parsed.policies;
|
||||
const policy: SecurityPolicy = {};
|
||||
for (const item of policiesList) {
|
||||
policy[item.tool_name] = item.policy;
|
||||
}
|
||||
|
||||
debugLogger.debug(`[Conseca] Policy Generation Parsed:`, policy);
|
||||
return { policy };
|
||||
} catch (parseError) {
|
||||
debugLogger.debug(
|
||||
`[Conseca] Policy Generation JSON Parse Error:`,
|
||||
parseError,
|
||||
);
|
||||
return {
|
||||
policy: {},
|
||||
error: `JSON Parse Error: ${parseError instanceof Error ? parseError.message : String(parseError)}. Raw: ${responseText}`,
|
||||
};
|
||||
}
|
||||
} catch (error) {
|
||||
debugLogger.error('Policy generation failed:', error);
|
||||
return {
|
||||
policy: {},
|
||||
error: `Policy generation failed: ${error instanceof Error ? error.message : String(error)}`,
|
||||
};
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,18 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import type { SafetyCheckDecision } from '../protocol.js';
|
||||
|
||||
export interface ToolPolicy {
|
||||
permissions: SafetyCheckDecision;
|
||||
constraints: string;
|
||||
rationale: string;
|
||||
}
|
||||
|
||||
/**
|
||||
* A map of tool names to their specific security policies.
|
||||
*/
|
||||
export type SecurityPolicy = Record<string, ToolPolicy>;
|
||||
@@ -7,50 +7,139 @@
|
||||
import { describe, it, expect, vi, beforeEach } from 'vitest';
|
||||
import { ContextBuilder } from './context-builder.js';
|
||||
import type { Config } from '../config/config.js';
|
||||
import type { ConversationTurn } from './protocol.js';
|
||||
import type { Content, FunctionCall } from '@google/genai';
|
||||
|
||||
describe('ContextBuilder', () => {
|
||||
let contextBuilder: ContextBuilder;
|
||||
let mockConfig: Config;
|
||||
const mockHistory: ConversationTurn[] = [
|
||||
{ user: { text: 'hello' }, model: { text: 'hi' } },
|
||||
];
|
||||
let mockConfig: Partial<Config>;
|
||||
let mockHistory: Content[];
|
||||
const mockCwd = '/home/user/project';
|
||||
const mockWorkspaces = ['/home/user/project'];
|
||||
|
||||
beforeEach(() => {
|
||||
vi.spyOn(process, 'cwd').mockReturnValue(mockCwd);
|
||||
mockHistory = [];
|
||||
|
||||
mockConfig = {
|
||||
getWorkspaceContext: vi.fn().mockReturnValue({
|
||||
getDirectories: vi.fn().mockReturnValue(mockWorkspaces),
|
||||
}),
|
||||
apiKey: 'secret-api-key',
|
||||
somePublicConfig: 'public-value',
|
||||
nested: {
|
||||
secretToken: 'hidden',
|
||||
public: 'visible',
|
||||
},
|
||||
} as unknown as Config;
|
||||
contextBuilder = new ContextBuilder(mockConfig, mockHistory);
|
||||
getQuestion: vi.fn().mockReturnValue('mock question'),
|
||||
getGeminiClient: vi.fn().mockReturnValue({
|
||||
getHistory: vi.fn().mockImplementation(() => mockHistory),
|
||||
}),
|
||||
};
|
||||
contextBuilder = new ContextBuilder(mockConfig as unknown as Config);
|
||||
});
|
||||
|
||||
it('should build full context with all fields', () => {
|
||||
it('should build full context with empty history', () => {
|
||||
mockHistory = [];
|
||||
// Should inject current question
|
||||
const context = contextBuilder.buildFullContext();
|
||||
expect(context.environment.cwd).toBe(mockCwd);
|
||||
expect(context.environment.workspaces).toEqual(mockWorkspaces);
|
||||
expect(context.history?.turns).toEqual(mockHistory);
|
||||
expect(context.history?.turns).toEqual([
|
||||
{
|
||||
user: { text: 'mock question' },
|
||||
model: {},
|
||||
},
|
||||
]);
|
||||
});
|
||||
|
||||
it('should build minimal context with only required keys', () => {
|
||||
it('should build full context with existing history (User -> Model)', () => {
|
||||
mockHistory = [
|
||||
{ role: 'user', parts: [{ text: 'Hello' }] },
|
||||
{ role: 'model', parts: [{ text: 'Hi there' }] },
|
||||
];
|
||||
// Should NOT inject current question if history exists
|
||||
const context = contextBuilder.buildFullContext();
|
||||
expect(context.history?.turns).toHaveLength(1);
|
||||
expect(context.history?.turns[0]).toEqual({
|
||||
user: { text: 'Hello' },
|
||||
model: { text: 'Hi there', toolCalls: [] },
|
||||
});
|
||||
});
|
||||
|
||||
it('should handle history with tool calls', () => {
|
||||
const mockToolCall: FunctionCall = {
|
||||
id: 'call_1',
|
||||
name: 'list_files',
|
||||
args: { path: '.' },
|
||||
};
|
||||
mockHistory = [
|
||||
{ role: 'user', parts: [{ text: 'List files' }] },
|
||||
{
|
||||
role: 'model',
|
||||
parts: [
|
||||
{ text: 'Sure, listing files.' },
|
||||
{ functionCall: mockToolCall },
|
||||
],
|
||||
},
|
||||
];
|
||||
|
||||
const context = contextBuilder.buildFullContext();
|
||||
expect(context.history?.turns).toHaveLength(1);
|
||||
expect(context.history?.turns[0].model.toolCalls).toEqual([mockToolCall]);
|
||||
expect(context.history?.turns[0].model.text).toBe('Sure, listing files.');
|
||||
});
|
||||
|
||||
it('should handle orphan model response (Model starts conversation)', () => {
|
||||
mockHistory = [
|
||||
{ role: 'model', parts: [{ text: 'Welcome!' }] },
|
||||
{ role: 'user', parts: [{ text: 'Thanks' }] },
|
||||
];
|
||||
|
||||
const context = contextBuilder.buildFullContext();
|
||||
// 1. Orphan model response -> Turn 1: User="" Model="Welcome!"
|
||||
// 2. User "Thanks" -> Turn 2: User="Thanks" Model={} (pending)
|
||||
expect(context.history?.turns).toHaveLength(2);
|
||||
expect(context.history?.turns[0]).toEqual({
|
||||
user: { text: '' },
|
||||
model: { text: 'Welcome!', toolCalls: [] },
|
||||
});
|
||||
expect(context.history?.turns[1]).toEqual({
|
||||
user: { text: 'Thanks' },
|
||||
model: {},
|
||||
});
|
||||
});
|
||||
|
||||
it('should handle multiple user turns in a row', () => {
|
||||
mockHistory = [
|
||||
{ role: 'user', parts: [{ text: 'Q1' }] },
|
||||
{ role: 'user', parts: [{ text: 'Q2' }] },
|
||||
{ role: 'model', parts: [{ text: 'A2' }] },
|
||||
];
|
||||
|
||||
const context = contextBuilder.buildFullContext();
|
||||
// 1. "Q1" -> Turn 1: User="Q1" Model={}
|
||||
// 2. "Q2" -> Turn 2: User="Q2" Model="A2"
|
||||
expect(context.history?.turns).toHaveLength(2);
|
||||
expect(context.history?.turns[0]).toEqual({
|
||||
user: { text: 'Q1' },
|
||||
model: {},
|
||||
});
|
||||
expect(context.history?.turns[1]).toEqual({
|
||||
user: { text: 'Q2' },
|
||||
model: { text: 'A2', toolCalls: [] },
|
||||
});
|
||||
});
|
||||
|
||||
it('should build minimal context', () => {
|
||||
mockHistory = [{ role: 'user', parts: [{ text: 'test' }] }];
|
||||
const context = contextBuilder.buildMinimalContext(['environment']);
|
||||
|
||||
expect(context).toHaveProperty('environment');
|
||||
expect(context).not.toHaveProperty('config');
|
||||
expect(context).not.toHaveProperty('history');
|
||||
});
|
||||
|
||||
it('should handle missing history', () => {
|
||||
contextBuilder = new ContextBuilder(mockConfig);
|
||||
it('should handle undefined parts gracefully', () => {
|
||||
mockHistory = [
|
||||
{ role: 'user', parts: undefined as unknown as [] },
|
||||
{ role: 'model', parts: undefined as unknown as [] },
|
||||
];
|
||||
const context = contextBuilder.buildFullContext();
|
||||
expect(context.history?.turns).toEqual([]);
|
||||
expect(context.history?.turns).toHaveLength(1);
|
||||
expect(context.history?.turns[0]).toEqual({
|
||||
user: { text: '' },
|
||||
model: { text: '', toolCalls: [] },
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
@@ -6,20 +6,39 @@
|
||||
|
||||
import type { SafetyCheckInput, ConversationTurn } from './protocol.js';
|
||||
import type { Config } from '../config/config.js';
|
||||
import { debugLogger } from '../utils/debugLogger.js';
|
||||
import type { Content, FunctionCall } from '@google/genai';
|
||||
|
||||
/**
|
||||
* Builds context objects for safety checkers, ensuring sensitive data is filtered.
|
||||
*/
|
||||
export class ContextBuilder {
|
||||
constructor(
|
||||
private readonly config: Config,
|
||||
private readonly conversationHistory: ConversationTurn[] = [],
|
||||
) {}
|
||||
constructor(private readonly config: Config) {}
|
||||
|
||||
/**
|
||||
* Builds the full context object with all available data.
|
||||
*/
|
||||
buildFullContext(): SafetyCheckInput['context'] {
|
||||
const clientHistory = this.config.getGeminiClient()?.getHistory() || [];
|
||||
const history = this.convertHistoryToTurns(clientHistory);
|
||||
|
||||
debugLogger.debug(
|
||||
`[ContextBuilder] buildFullContext called. Converted history length: ${history.length}`,
|
||||
);
|
||||
|
||||
// ContextBuilder's responsibility is to provide the *current* context.
|
||||
// If the conversation hasn't started (history is empty), we check if there's a pending question.
|
||||
// However, if the history is NOT empty, we trust it reflects the true state.
|
||||
const currentQuestion = this.config.getQuestion();
|
||||
if (currentQuestion && history.length === 0) {
|
||||
history.push({
|
||||
user: {
|
||||
text: currentQuestion,
|
||||
},
|
||||
model: {},
|
||||
});
|
||||
}
|
||||
|
||||
return {
|
||||
environment: {
|
||||
cwd: process.cwd(),
|
||||
@@ -29,7 +48,7 @@ export class ContextBuilder {
|
||||
.getDirectories() as string[],
|
||||
},
|
||||
history: {
|
||||
turns: this.conversationHistory,
|
||||
turns: history,
|
||||
},
|
||||
};
|
||||
}
|
||||
@@ -53,4 +72,51 @@ export class ContextBuilder {
|
||||
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
|
||||
return minimalContext as SafetyCheckInput['context'];
|
||||
}
|
||||
|
||||
// Helper to convert Google GenAI Content[] to Safety Protocol ConversationTurn[]
|
||||
private convertHistoryToTurns(history: Content[]): ConversationTurn[] {
|
||||
const turns: ConversationTurn[] = [];
|
||||
let currentUserRequest: { text: string } | undefined;
|
||||
|
||||
for (const content of history) {
|
||||
if (content.role === 'user') {
|
||||
if (currentUserRequest) {
|
||||
// Previous user turn didn't have a matching model response (or it was filtered out)
|
||||
// Push it as a turn with empty model response
|
||||
turns.push({ user: currentUserRequest, model: {} });
|
||||
}
|
||||
currentUserRequest = {
|
||||
text: content.parts?.map((p) => p.text).join('') || '',
|
||||
};
|
||||
} else if (content.role === 'model') {
|
||||
const modelResponse = {
|
||||
text:
|
||||
content.parts
|
||||
?.filter((p) => p.text)
|
||||
.map((p) => p.text)
|
||||
.join('') || '',
|
||||
toolCalls:
|
||||
content.parts
|
||||
?.filter((p) => 'functionCall' in p)
|
||||
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
|
||||
.map((p) => p.functionCall as FunctionCall) || [],
|
||||
};
|
||||
|
||||
if (currentUserRequest) {
|
||||
turns.push({ user: currentUserRequest, model: modelResponse });
|
||||
currentUserRequest = undefined;
|
||||
} else {
|
||||
// Model response without preceding user request.
|
||||
// This creates a turn with empty user text.
|
||||
turns.push({ user: { text: '' }, model: modelResponse });
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (currentUserRequest) {
|
||||
turns.push({ user: currentUserRequest, model: {} });
|
||||
}
|
||||
|
||||
return turns;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -89,6 +89,10 @@ export type SafetyCheckResult =
|
||||
* This will be shown to the user.
|
||||
*/
|
||||
reason?: string;
|
||||
/**
|
||||
* Optional error message if the decision was made due to a system failure (fail-open).
|
||||
*/
|
||||
error?: string;
|
||||
}
|
||||
| {
|
||||
decision: SafetyCheckDecision.DENY;
|
||||
|
||||
@@ -8,6 +8,7 @@ import { describe, it, expect, beforeEach } from 'vitest';
|
||||
import { CheckerRegistry } from './registry.js';
|
||||
import { InProcessCheckerType } from '../policy/types.js';
|
||||
import { AllowedPathChecker } from './built-in.js';
|
||||
import { ConsecaSafetyChecker } from './conseca/conseca.js';
|
||||
|
||||
describe('CheckerRegistry', () => {
|
||||
let registry: CheckerRegistry;
|
||||
@@ -18,10 +19,15 @@ describe('CheckerRegistry', () => {
|
||||
});
|
||||
|
||||
it('should resolve built-in in-process checkers', () => {
|
||||
const checker = registry.resolveInProcess(
|
||||
const allowedPathChecker = registry.resolveInProcess(
|
||||
InProcessCheckerType.ALLOWED_PATH,
|
||||
);
|
||||
expect(checker).toBeInstanceOf(AllowedPathChecker);
|
||||
expect(allowedPathChecker).toBeInstanceOf(AllowedPathChecker);
|
||||
|
||||
const consecaChecker = registry.resolveInProcess(
|
||||
InProcessCheckerType.CONSECA,
|
||||
);
|
||||
expect(consecaChecker).toBeInstanceOf(ConsecaSafetyChecker);
|
||||
});
|
||||
|
||||
it('should throw for unknown in-process checkers', () => {
|
||||
|
||||
@@ -9,6 +9,8 @@ import * as fs from 'node:fs';
|
||||
import { type InProcessChecker, AllowedPathChecker } from './built-in.js';
|
||||
import { InProcessCheckerType } from '../policy/types.js';
|
||||
|
||||
import { ConsecaSafetyChecker } from './conseca/conseca.js';
|
||||
|
||||
/**
|
||||
* Registry for managing safety checker resolution.
|
||||
*/
|
||||
@@ -17,10 +19,22 @@ export class CheckerRegistry {
|
||||
// No external built-ins for now
|
||||
]);
|
||||
|
||||
private static readonly BUILT_IN_IN_PROCESS_CHECKERS = new Map<
|
||||
string,
|
||||
InProcessChecker
|
||||
>([[InProcessCheckerType.ALLOWED_PATH, new AllowedPathChecker()]]);
|
||||
private static BUILT_IN_IN_PROCESS_CHECKERS:
|
||||
| Map<string, InProcessChecker>
|
||||
| undefined;
|
||||
|
||||
private static getBuiltInInProcessCheckers(): Map<string, InProcessChecker> {
|
||||
if (!CheckerRegistry.BUILT_IN_IN_PROCESS_CHECKERS) {
|
||||
CheckerRegistry.BUILT_IN_IN_PROCESS_CHECKERS = new Map<
|
||||
string,
|
||||
InProcessChecker
|
||||
>([
|
||||
[InProcessCheckerType.ALLOWED_PATH, new AllowedPathChecker()],
|
||||
[InProcessCheckerType.CONSECA, ConsecaSafetyChecker.getInstance()],
|
||||
]);
|
||||
}
|
||||
return CheckerRegistry.BUILT_IN_IN_PROCESS_CHECKERS;
|
||||
}
|
||||
|
||||
// Regex to validate checker names (alphanumeric and hyphens only)
|
||||
private static readonly VALID_NAME_PATTERN = /^[a-z0-9-]+$/;
|
||||
@@ -58,14 +72,14 @@ export class CheckerRegistry {
|
||||
throw new Error(`Invalid checker name "${name}".`);
|
||||
}
|
||||
|
||||
const checker = CheckerRegistry.BUILT_IN_IN_PROCESS_CHECKERS.get(name);
|
||||
const checker = CheckerRegistry.getBuiltInInProcessCheckers().get(name);
|
||||
if (checker) {
|
||||
return checker;
|
||||
}
|
||||
|
||||
throw new Error(
|
||||
`Unknown in-process checker "${name}". Available: ${Array.from(
|
||||
CheckerRegistry.BUILT_IN_IN_PROCESS_CHECKERS.keys(),
|
||||
CheckerRegistry.getBuiltInInProcessCheckers().keys(),
|
||||
).join(', ')}`,
|
||||
);
|
||||
}
|
||||
@@ -77,7 +91,7 @@ export class CheckerRegistry {
|
||||
static getBuiltInCheckers(): string[] {
|
||||
return [
|
||||
...Array.from(this.BUILT_IN_EXTERNAL_CHECKERS.keys()),
|
||||
...Array.from(this.BUILT_IN_IN_PROCESS_CHECKERS.keys()),
|
||||
...Array.from(this.getBuiltInInProcessCheckers().keys()),
|
||||
];
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user