feat(sdp): Add unit tests

[agduan]

(6/6) Adds unit tests for the Sensitive Data Protection (SDP) middleware.

View diff

Checklist (if applicable):

• PR title is following https://www.conventionalcommits.org/en/v1.0.0/
• Tested (manually, unit tested, etc.)
• Docs updated (updated docs or a docs bug required)

[gemini-code-assist]

Code Review

This pull request introduces a Sensitive Data Protection (SDP) middleware plugin for Google Cloud in Genkit, allowing input and output message redaction using Google Cloud DLP. Feedback on the implementation highlights a critical security issue where using the logical OR operator (||) could leak sensitive data if the redacted output is an empty string, recommending the nullish coalescing operator (??) instead. Additionally, the reviewer suggested optimizing client initialization to avoid redundant credential checks, removing a @ts-ignore comment by providing a safe fallback array, and expanding the test suite to cover output redaction.

[gemini-code-assist]

Using the logical OR operator (||) here can lead to a security vulnerability. If the Sensitive Data Protection (SDP) API successfully redacts the content by completely removing it (resulting in an empty string ""), the expression response.item?.value || text will evaluate to the original sensitive text because "" is falsy.

Using the nullish coalescing operator (??) ensures that we only fall back to the original text if the response value is strictly null or undefined.

Suggested change

	return response.item?.value || text;
	return response.item?.value ?? text;

[gemini-code-assist]

The createDlpClient helper function can be simplified significantly. Since the caller already resolves the credentials and project ID from the environment, we don't need to duplicate that logic here. Passing projectId and credentials directly as arguments makes the function cleaner and avoids potential duplicate calls to credentialsFromEnvironment().

async function createDlpClient(
  projectId?: string,
  credentials?: any
): Promise<v2.DlpServiceClient> {
  let dlpModule;
  try {
    dlpModule = await import('@google-cloud/dlp');
  } catch (e) {
    throw new Error(
      'Please install the @google-cloud/dlp package to use the SDP middleware.'
    );
  }

  return new dlpModule.v2.DlpServiceClient({
    credentials: credentials as any,
    projectId,
  });
}

[gemini-code-assist]

We can optimize the client initialization by only calling credentialsFromEnvironment() if either projectId or credentials is missing from the options. This avoids unnecessary environment/file checks when the user has explicitly configured the middleware with inline options. Additionally, we can pass these resolved values directly to the simplified createDlpClient function.

            const envAuth =
              !opts.projectId || !opts.credentials
                ? await credentialsFromEnvironment()
                : null;
            const projectId = opts.projectId || envAuth?.projectId;
            const credentials = opts.credentials || envAuth?.credentials;

            if (!projectId) {
              throw new Error(
                'Project ID is required for Sensitive Data Protection. Please set the projectId option or configure it in your environment.'
              );
            }

            const client = await createDlpClient(projectId, credentials);

[gemini-code-assist]

We can eliminate the @ts-ignore comment by using message.content || []. This provides a safe fallback array and allows TypeScript to correctly infer the types for flatMap.

Suggested change

	// @ts-ignore - flatMap creates a flat array of content parts
	.flatMap((message) => message.content!)
	// ignore multimedia content and content that has been cleaned already
	// extract all message content into a single array
	.flatMap((message) => message.content || [])

[gemini-code-assist]

The current test suite only covers input redaction (where the prompt contains sensitive data). There is no test verifying that output redaction works correctly when the model itself generates sensitive data. Adding a test case for this scenario ensures full coverage of the middleware's output interception path.

  });

  it('should redact sensitive data generated by the model (output redaction)', async () => {
    ai.defineModel({ name: 'sensitiveModel' }, async () => {
      return {
        message: {
          role: 'model',
          content: [{ text: 'Here is my credit card' }],
        },
      };
    });

    const response = await ai.generate({
      model: 'sensitiveModel',
      prompt: 'Hello',
      use: [sensitiveDataProtection({})],
    });

    expect(response.text).toBe('Here is [REDACTED_BY_MOCK]');
  });
});