mirror of
https://github.com/n8n-io/n8n.git
synced 2026-10-17 05:29:44 +02:00
feat(Databricks Node): Add hidden Databricks Chat Model sub-node (#37299)
This commit is contained in:
Yen Su
GitHub
parent
ab48cfc2ad
commit
233f1ca0e7
|
|
@ -55,8 +55,36 @@ describe('fetchFollowingRedirects', () => {
|
|||
});
|
||||
|
||||
expect(onBeforeHop).toHaveBeenCalledTimes(2);
|
||||
expect(onBeforeHop).toHaveBeenNthCalledWith(1, 'https://example.com/start');
|
||||
expect(onBeforeHop).toHaveBeenNthCalledWith(2, 'https://example.com/relative');
|
||||
expect(onBeforeHop).toHaveBeenNthCalledWith(1, 'https://example.com/start', {
|
||||
crossedOrigin: false,
|
||||
});
|
||||
expect(onBeforeHop).toHaveBeenNthCalledWith(2, 'https://example.com/relative', {
|
||||
crossedOrigin: false,
|
||||
});
|
||||
});
|
||||
|
||||
it('treats crossedOrigin as permanent, even when the chain returns to the original origin', async () => {
|
||||
const fetcher = mockFetcher()
|
||||
.mockResolvedValueOnce(makeRedirect(307, 'https://other.example.com/away'))
|
||||
.mockResolvedValueOnce(makeRedirect(307, 'https://example.com/back'))
|
||||
.mockResolvedValueOnce(makeResponse(200));
|
||||
const onBeforeHop = vi.fn();
|
||||
|
||||
await fetchFollowingRedirects(
|
||||
fetcher,
|
||||
'https://example.com',
|
||||
{ headers: { authorization: 'Bearer secret' } },
|
||||
{ onBeforeHop },
|
||||
);
|
||||
|
||||
expect(onBeforeHop.mock.calls.map((call) => call[1])).toEqual([
|
||||
{ crossedOrigin: false },
|
||||
{ crossedOrigin: true },
|
||||
{ crossedOrigin: true },
|
||||
]);
|
||||
// Credential headers stripped on the cross-origin hop are not restored either
|
||||
const hop3Headers = new Headers((fetcher.mock.calls[2][1] as RequestInit).headers);
|
||||
expect(hop3Headers.get('authorization')).toBeNull();
|
||||
});
|
||||
|
||||
it('aborts the chain when onBeforeHop throws', async () => {
|
||||
|
|
@ -144,6 +172,40 @@ describe('fetchFollowingRedirects', () => {
|
|||
);
|
||||
});
|
||||
|
||||
it('strips credential headers when a redirect crosses origins', async () => {
|
||||
const fetcher = mockFetcher()
|
||||
.mockResolvedValueOnce(makeRedirect(307, 'https://other.example.com/target'))
|
||||
.mockResolvedValueOnce(makeResponse(200));
|
||||
|
||||
await fetchFollowingRedirects(fetcher, 'https://example.com', {
|
||||
headers: {
|
||||
authorization: 'Bearer secret',
|
||||
cookie: 'session=abc',
|
||||
'proxy-authorization': 'Basic xyz',
|
||||
'content-type': 'application/json',
|
||||
},
|
||||
});
|
||||
|
||||
const hop2Headers = new Headers((fetcher.mock.calls[1][1] as RequestInit).headers);
|
||||
expect(hop2Headers.get('authorization')).toBeNull();
|
||||
expect(hop2Headers.get('cookie')).toBeNull();
|
||||
expect(hop2Headers.get('proxy-authorization')).toBeNull();
|
||||
expect(hop2Headers.get('content-type')).toBe('application/json');
|
||||
});
|
||||
|
||||
it('keeps credential headers on a same-origin redirect', async () => {
|
||||
const fetcher = mockFetcher()
|
||||
.mockResolvedValueOnce(makeRedirect(307, '/moved'))
|
||||
.mockResolvedValueOnce(makeResponse(200));
|
||||
|
||||
await fetchFollowingRedirects(fetcher, 'https://example.com/start', {
|
||||
headers: { authorization: 'Bearer secret' },
|
||||
});
|
||||
|
||||
const hop2Headers = new Headers((fetcher.mock.calls[1][1] as RequestInit).headers);
|
||||
expect(hop2Headers.get('authorization')).toBe('Bearer secret');
|
||||
});
|
||||
|
||||
it('preserves a URL input object on the first hop', async () => {
|
||||
const fetcher = mockFetcher().mockResolvedValue(makeResponse(200));
|
||||
const url = new URL('https://example.com');
|
||||
|
|
|
|||
|
|
@ -33,10 +33,10 @@ describe('createRefreshingAuthFetch', () => {
|
|||
expect(init.redirect).toBeUndefined();
|
||||
});
|
||||
|
||||
it('forwards auth headers to a redirect after its URL is validated', async () => {
|
||||
it('forwards auth headers to a same-origin redirect after its URL is validated', async () => {
|
||||
const baseFetch = vi
|
||||
.fn()
|
||||
.mockResolvedValueOnce(makeRedirect('https://redirected.example/mcp'))
|
||||
.mockResolvedValueOnce(makeRedirect('https://example.com/v2/mcp'))
|
||||
.mockResolvedValueOnce(new Response('ok'));
|
||||
const assertAllowedUrl = vi.fn();
|
||||
const fetchWithAuth = createRefreshingAuthFetch({
|
||||
|
|
@ -57,13 +57,55 @@ describe('createRefreshingAuthFetch', () => {
|
|||
|
||||
const [, secondInit] = baseFetch.mock.calls[1] as [URL, RequestInit];
|
||||
const headers = new Headers(secondInit.headers);
|
||||
expect(assertAllowedUrl).toHaveBeenLastCalledWith('https://redirected.example/mcp');
|
||||
expect(assertAllowedUrl).toHaveBeenLastCalledWith('https://example.com/v2/mcp');
|
||||
expect(headers.get('authorization')).toBe('Bearer token');
|
||||
expect(headers.get('x-api-key')).toBe('secret');
|
||||
expect(headers.get('cookie')).toBe('session=secret');
|
||||
expect(headers.get('x-request')).toBe('value');
|
||||
});
|
||||
|
||||
it('withholds auth headers once a redirect crosses origins, even to a validated URL', async () => {
|
||||
const baseFetch = vi
|
||||
.fn()
|
||||
.mockResolvedValueOnce(makeRedirect('https://other.example/mcp'))
|
||||
.mockResolvedValueOnce(new Response('ok'));
|
||||
const fetchWithAuth = createRefreshingAuthFetch({
|
||||
baseFetch,
|
||||
initialHeaders: { Authorization: 'Bearer token', 'X-Api-Key': 'secret' },
|
||||
assertAllowedUrl: vi.fn(),
|
||||
});
|
||||
|
||||
await fetchWithAuth('https://example.com/mcp', {
|
||||
headers: { 'X-Request': 'value' },
|
||||
});
|
||||
|
||||
const [, secondInit] = baseFetch.mock.calls[1] as [URL, RequestInit];
|
||||
const headers = new Headers(secondInit.headers);
|
||||
expect(headers.get('authorization')).toBeNull();
|
||||
expect(headers.get('x-api-key')).toBeNull();
|
||||
expect(headers.get('x-request')).toBe('value');
|
||||
});
|
||||
|
||||
it('does not refresh on a 401 from a cross-origin redirect target', async () => {
|
||||
const baseFetch = vi
|
||||
.fn()
|
||||
.mockResolvedValueOnce(makeRedirect('https://other.example/mcp'))
|
||||
.mockResolvedValueOnce(makeUnauthorized());
|
||||
const refreshHeaders = vi.fn();
|
||||
const fetchWithAuth = createRefreshingAuthFetch({
|
||||
baseFetch,
|
||||
initialHeaders: { Authorization: 'Bearer token' },
|
||||
refreshHeaders,
|
||||
assertAllowedUrl: vi.fn(),
|
||||
});
|
||||
|
||||
const response = await fetchWithAuth('https://example.com/mcp');
|
||||
|
||||
expect(response.status).toBe(401);
|
||||
expect(refreshHeaders).not.toHaveBeenCalled();
|
||||
expect(baseFetch).toHaveBeenCalledTimes(2);
|
||||
});
|
||||
|
||||
describe('401 handling', () => {
|
||||
it('refreshes the token and retries the request once', async () => {
|
||||
const baseFetch = vi
|
||||
|
|
@ -194,7 +236,7 @@ describe('createRefreshingAuthFetch', () => {
|
|||
it('retries with the refreshed token on a redirect-validated request', async () => {
|
||||
const baseFetch = vi
|
||||
.fn()
|
||||
.mockResolvedValueOnce(makeRedirect('https://redirected.example/mcp'))
|
||||
.mockResolvedValueOnce(makeRedirect('https://example.com/v2/mcp'))
|
||||
.mockResolvedValueOnce(makeUnauthorized())
|
||||
.mockResolvedValueOnce(new Response('ok'));
|
||||
const refreshHeaders = vi.fn().mockResolvedValue({ Authorization: 'Bearer fresh' });
|
||||
|
|
@ -211,9 +253,9 @@ describe('createRefreshingAuthFetch', () => {
|
|||
expect(response.status).toBe(200);
|
||||
expect(authorizationOf(baseFetch.mock.calls[1])).toBe('Bearer stale');
|
||||
expect(authorizationOf(baseFetch.mock.calls[2])).toBe('Bearer fresh');
|
||||
expect(assertAllowedUrl.mock.calls.flat()).toEqual([
|
||||
expect(assertAllowedUrl.mock.calls.map((call) => call[0])).toEqual([
|
||||
'https://example.com/mcp',
|
||||
'https://redirected.example/mcp',
|
||||
'https://example.com/v2/mcp',
|
||||
]);
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -3,14 +3,22 @@ import { OperationalError } from 'n8n-workflow';
|
|||
const DEFAULT_MAX_REDIRECTS = 20;
|
||||
|
||||
export interface FollowRedirectsOptions {
|
||||
/** Called before each hop; throw to reject. */
|
||||
onBeforeHop?: (url: string) => void | Promise<void>;
|
||||
/**
|
||||
* Called before each hop; throw to reject. `crossedOrigin` turns true once
|
||||
* the chain has left its original origin and stays true — fetchers that
|
||||
* inject their own credential headers must withhold them from then on,
|
||||
* mirroring how this helper strips credentials it carries in `init`.
|
||||
*/
|
||||
onBeforeHop?: (url: string, info: { crossedOrigin: boolean }) => void | Promise<void>;
|
||||
maxRedirects?: number;
|
||||
}
|
||||
|
||||
const CREDENTIAL_HEADERS = ['authorization', 'cookie', 'proxy-authorization'];
|
||||
|
||||
/**
|
||||
* Manual redirect handling so each hop can be validated before the request is
|
||||
* sent. 301/302/303 demote unsafe methods to GET per fetch spec.
|
||||
* sent. Per fetch spec, 301/302/303 demote unsafe methods to GET, and
|
||||
* credential headers are stripped when a redirect crosses origins.
|
||||
*/
|
||||
export async function fetchFollowingRedirects(
|
||||
fetcher: (input: RequestInfo | URL, init?: RequestInit) => Promise<Response>,
|
||||
|
|
@ -22,11 +30,12 @@ export async function fetchFollowingRedirects(
|
|||
let currentInput: string | URL = url;
|
||||
let currentInit: RequestInit = { ...init };
|
||||
let hops = 0;
|
||||
let crossedOrigin = false;
|
||||
|
||||
while (true) {
|
||||
const currentUrlString = currentInput instanceof URL ? currentInput.href : currentInput;
|
||||
if (options?.onBeforeHop) {
|
||||
await options.onBeforeHop(currentUrlString);
|
||||
await options.onBeforeHop(currentUrlString, { crossedOrigin });
|
||||
}
|
||||
|
||||
const response = await fetcher(currentInput, {
|
||||
|
|
@ -52,6 +61,13 @@ export async function fetchFollowingRedirects(
|
|||
|
||||
currentInput = new URL(location, currentUrlString);
|
||||
|
||||
if (currentInput.origin !== new URL(currentUrlString).origin) {
|
||||
crossedOrigin = true;
|
||||
const headers = new Headers(currentInit.headers);
|
||||
for (const header of CREDENTIAL_HEADERS) headers.delete(header);
|
||||
currentInit = { ...currentInit, headers };
|
||||
}
|
||||
|
||||
const method = (currentInit.method ?? 'GET').toUpperCase();
|
||||
const isUnsafe = method !== 'GET' && method !== 'HEAD';
|
||||
if (
|
||||
|
|
|
|||
|
|
@ -44,6 +44,9 @@ export function createRefreshingAuthFetch({
|
|||
|
||||
return async (input: RequestInfo | URL, init?: RequestInit): Promise<Response> => {
|
||||
let retried = false;
|
||||
// Auth headers follow redirects within the starting origin, but per fetch
|
||||
// spec are withheld once a hop crosses origins, even to an allowed host
|
||||
let sendAuth = true;
|
||||
const authedFetch = async (
|
||||
requestInput: RequestInfo | URL,
|
||||
requestInit?: RequestInit,
|
||||
|
|
@ -52,12 +55,13 @@ export function createRefreshingAuthFetch({
|
|||
const execute = async () =>
|
||||
await baseFetch(requestInput instanceof Request ? requestInput.clone() : requestInput, {
|
||||
...requestInit,
|
||||
// Include auth headers for redirect requests too
|
||||
headers: mergeHeaders(getInputHeaders(requestInput, requestInit), authHeaders),
|
||||
headers: sendAuth
|
||||
? mergeHeaders(getInputHeaders(requestInput, requestInit), authHeaders)
|
||||
: getInputHeaders(requestInput, requestInit),
|
||||
});
|
||||
|
||||
const response = await execute();
|
||||
if (response.status !== 401 || !refreshHeaders || retried) return response;
|
||||
if (response.status !== 401 || !refreshHeaders || retried || !sendAuth) return response;
|
||||
|
||||
retried = true;
|
||||
const canRetry = authVersion !== requestAuthVersion || (await refresh());
|
||||
|
|
@ -73,7 +77,10 @@ export function createRefreshingAuthFetch({
|
|||
// unwrapped to their URL so the redirect loop can carry a stable input.
|
||||
const startUrl = input instanceof Request ? input.url : input;
|
||||
return await fetchFollowingRedirects(authedFetch, startUrl, init, {
|
||||
onBeforeHop: assertAllowedUrl,
|
||||
onBeforeHop: async (hopUrl, { crossedOrigin }) => {
|
||||
sendAuth = !crossedOrigin;
|
||||
await assertAllowedUrl(hopUrl);
|
||||
},
|
||||
});
|
||||
};
|
||||
}
|
||||
|
|
|
|||
|
|
@ -13,4 +13,6 @@ export const passthroughEgressFilter: NodeEgressFilter = {
|
|||
validateUrl: async () => await Promise.resolve(createResultOk(undefined)),
|
||||
// Same function reference every call, so callers can compare it by identity.
|
||||
createSecureLookup: (): LookupFunction => lookup,
|
||||
// No egress policy configured, so every redirect target is allowed.
|
||||
validateRedirectSync: () => {},
|
||||
};
|
||||
|
|
|
|||
|
|
@ -52,9 +52,11 @@ export interface ClientOAuth2Options {
|
|||
/**
|
||||
* When provided, token endpoint requests are validated against the host's
|
||||
* outbound network policy before dispatch, at DNS resolution time, and on
|
||||
* every redirect. Omit to leave the request unchecked.
|
||||
* every redirect. Omit to leave the request unchecked. Narrowed to the
|
||||
* methods this client uses so node-facing filters (`NodeEgressFilter`)
|
||||
* qualify as well as the full bridge.
|
||||
*/
|
||||
ssrfBridge?: SsrfBridge;
|
||||
ssrfBridge?: Pick<SsrfBridge, 'validateUrl' | 'validateRedirectSync' | 'createSecureLookup'>;
|
||||
}
|
||||
|
||||
export class ResponseError extends Error {
|
||||
|
|
|
|||
|
|
@ -71,6 +71,7 @@ describe('EmbeddingsOllama', () => {
|
|||
const egressFilter: NodeEgressFilter = {
|
||||
validateUrl: vi.fn().mockResolvedValue(createResultOk(undefined)),
|
||||
createSecureLookup: vi.fn().mockReturnValue(secureLookup),
|
||||
validateRedirectSync: vi.fn(),
|
||||
};
|
||||
const ctx = setupContext({ baseUrl: 'http://ollama.example.com:11434' }, egressFilter);
|
||||
|
||||
|
|
@ -93,6 +94,7 @@ describe('EmbeddingsOllama', () => {
|
|||
const egressFilter: NodeEgressFilter = {
|
||||
validateUrl: vi.fn().mockResolvedValue(createResultOk(undefined)),
|
||||
createSecureLookup: vi.fn().mockReturnValue(vi.fn()),
|
||||
validateRedirectSync: vi.fn(),
|
||||
};
|
||||
const ctx = setupContext(
|
||||
{
|
||||
|
|
@ -111,6 +113,7 @@ describe('EmbeddingsOllama', () => {
|
|||
const egressFilter: NodeEgressFilter = {
|
||||
validateUrl: vi.fn().mockResolvedValue(createResultOk(undefined)),
|
||||
createSecureLookup: vi.fn().mockReturnValue(vi.fn()),
|
||||
validateRedirectSync: vi.fn(),
|
||||
};
|
||||
const ctx = setupContext(
|
||||
{
|
||||
|
|
|
|||
|
|
@ -72,6 +72,7 @@ describe('LmChatOllama', () => {
|
|||
const egressFilter: NodeEgressFilter = {
|
||||
validateUrl: vi.fn().mockResolvedValue(createResultOk(undefined)),
|
||||
createSecureLookup: vi.fn().mockReturnValue(secureLookup),
|
||||
validateRedirectSync: vi.fn(),
|
||||
};
|
||||
const ctx = setupContext({ baseUrl: 'http://ollama.example.com:11434' }, egressFilter);
|
||||
|
||||
|
|
@ -92,6 +93,7 @@ describe('LmChatOllama', () => {
|
|||
const egressFilter: NodeEgressFilter = {
|
||||
validateUrl: vi.fn().mockResolvedValue(createResultOk(undefined)),
|
||||
createSecureLookup: vi.fn().mockReturnValue(vi.fn()),
|
||||
validateRedirectSync: vi.fn(),
|
||||
};
|
||||
const ctx = setupContext(
|
||||
{
|
||||
|
|
@ -110,6 +112,7 @@ describe('LmChatOllama', () => {
|
|||
const egressFilter: NodeEgressFilter = {
|
||||
validateUrl: vi.fn().mockResolvedValue(createResultOk(undefined)),
|
||||
createSecureLookup: vi.fn().mockReturnValue(vi.fn()),
|
||||
validateRedirectSync: vi.fn(),
|
||||
};
|
||||
const ctx = setupContext(
|
||||
{
|
||||
|
|
|
|||
|
|
@ -72,6 +72,7 @@ describe('LmOllama', () => {
|
|||
const egressFilter: NodeEgressFilter = {
|
||||
validateUrl: vi.fn().mockResolvedValue(createResultOk(undefined)),
|
||||
createSecureLookup: vi.fn().mockReturnValue(secureLookup),
|
||||
validateRedirectSync: vi.fn(),
|
||||
};
|
||||
const ctx = setupContext({ baseUrl: 'http://ollama.example.com:11434' }, egressFilter);
|
||||
|
||||
|
|
@ -92,6 +93,7 @@ describe('LmOllama', () => {
|
|||
const egressFilter: NodeEgressFilter = {
|
||||
validateUrl: vi.fn().mockResolvedValue(createResultOk(undefined)),
|
||||
createSecureLookup: vi.fn().mockReturnValue(vi.fn()),
|
||||
validateRedirectSync: vi.fn(),
|
||||
};
|
||||
const ctx = setupContext(
|
||||
{
|
||||
|
|
@ -110,6 +112,7 @@ describe('LmOllama', () => {
|
|||
const egressFilter: NodeEgressFilter = {
|
||||
validateUrl: vi.fn().mockResolvedValue(createResultOk(undefined)),
|
||||
createSecureLookup: vi.fn().mockReturnValue(vi.fn()),
|
||||
validateRedirectSync: vi.fn(),
|
||||
};
|
||||
const ctx = setupContext(
|
||||
{
|
||||
|
|
|
|||
|
|
@ -0,0 +1,346 @@
|
|||
import { ChatOpenAI, type ClientOptions } from '@langchain/openai';
|
||||
import {
|
||||
getProxyAgent,
|
||||
makeN8nLlmFailedAttemptHandler,
|
||||
N8nLlmTracing,
|
||||
getConnectionHintNoticeField,
|
||||
} from '@n8n/ai-utilities';
|
||||
import {
|
||||
NodeConnectionTypes,
|
||||
NodeOperationError,
|
||||
type ILoadOptionsFunctions,
|
||||
type INodeListSearchResult,
|
||||
type INodeType,
|
||||
type INodeTypeDescription,
|
||||
type ISupplyDataFunctions,
|
||||
type SupplyData,
|
||||
} from 'n8n-workflow';
|
||||
|
||||
import type { DatabricksOAuth2Credential } from './token-provider';
|
||||
import { createDatabricksFetch, getDatabricksTokenProvider } from './token-provider';
|
||||
import { openAiFailedAttemptHandler } from '../../vendors/OpenAi/helpers/error-handling';
|
||||
|
||||
// Every request carries a secret (bearer token, or the client secret on the
|
||||
// mint path), so an http host would ship it in cleartext
|
||||
function assertHttpsHost(ctx: ILoadOptionsFunctions | ISupplyDataFunctions, host: string) {
|
||||
if (!URL.canParse(host) || new URL(host).protocol !== 'https:') {
|
||||
throw new NodeOperationError(ctx.getNode(), 'Databricks host must use https');
|
||||
}
|
||||
}
|
||||
|
||||
interface ServingEndpointsResponse {
|
||||
endpoints?: Array<{
|
||||
name: string;
|
||||
task?: string;
|
||||
config?: {
|
||||
served_entities?: Array<{
|
||||
external_model?: { name: string };
|
||||
foundation_model?: { name: string };
|
||||
}>;
|
||||
};
|
||||
}>;
|
||||
}
|
||||
|
||||
async function searchModels(
|
||||
this: ILoadOptionsFunctions,
|
||||
filter?: string,
|
||||
): Promise<INodeListSearchResult> {
|
||||
const credentials = await this.getCredentials<DatabricksOAuth2Credential>('databricksOAuth2Api');
|
||||
assertHttpsHost(this, credentials.host);
|
||||
const host = credentials.host.replace(/\/$/, '');
|
||||
|
||||
const response: ServingEndpointsResponse = await this.helpers.httpRequestWithAuthentication.call(
|
||||
this,
|
||||
'databricksOAuth2Api',
|
||||
{
|
||||
method: 'GET',
|
||||
url: `${host}/api/2.0/serving-endpoints`,
|
||||
headers: { Accept: 'application/json' },
|
||||
json: true,
|
||||
},
|
||||
);
|
||||
|
||||
const endpoints = response.endpoints ?? [];
|
||||
|
||||
const allResults = endpoints
|
||||
// Covers llm/v1/chat (foundation/external models) and agent/*/chat; custom
|
||||
// endpoints without a task are reachable via the resourceLocator's ID mode
|
||||
.filter((endpoint) => endpoint.task?.includes('chat'))
|
||||
.map((endpoint) => {
|
||||
const modelNames = (endpoint.config?.served_entities ?? [])
|
||||
.map((entity) => entity.external_model?.name ?? entity.foundation_model?.name)
|
||||
.filter(Boolean)
|
||||
.join(', ');
|
||||
|
||||
return {
|
||||
name: endpoint.name,
|
||||
value: endpoint.name,
|
||||
url: `${host}/ml/endpoints/${endpoint.name}`,
|
||||
description: modelNames || 'Model serving endpoint',
|
||||
};
|
||||
});
|
||||
|
||||
if (filter) {
|
||||
const filterLower = filter.toLowerCase();
|
||||
return {
|
||||
results: allResults.filter(
|
||||
(r) =>
|
||||
r.name.toLowerCase().includes(filterLower) ||
|
||||
r.description.toLowerCase().includes(filterLower),
|
||||
),
|
||||
};
|
||||
}
|
||||
|
||||
return { results: allResults };
|
||||
}
|
||||
|
||||
export class LmChatDatabricks implements INodeType {
|
||||
methods = {
|
||||
listSearch: {
|
||||
searchModels,
|
||||
},
|
||||
};
|
||||
|
||||
description: INodeTypeDescription = {
|
||||
displayName: 'Databricks Chat Model',
|
||||
name: 'lmChatDatabricks',
|
||||
hidden: true,
|
||||
icon: { light: 'file:databricks.svg', dark: 'file:databricks.dark.svg' },
|
||||
group: ['transform'],
|
||||
version: [1],
|
||||
description: 'For advanced usage with an AI chain',
|
||||
defaults: {
|
||||
name: 'Databricks Chat Model',
|
||||
},
|
||||
codex: {
|
||||
categories: ['AI'],
|
||||
subcategories: {
|
||||
AI: ['Language Models', 'Root Nodes'],
|
||||
'Language Models': ['Chat Models (Recommended)'],
|
||||
},
|
||||
resources: {
|
||||
primaryDocumentation: [
|
||||
{
|
||||
url: 'https://docs.n8n.io/integrations/builtin/cluster-nodes/sub-nodes/n8n-nodes-langchain.lmchatdatabricks/',
|
||||
},
|
||||
],
|
||||
},
|
||||
},
|
||||
|
||||
inputs: [],
|
||||
|
||||
outputs: [NodeConnectionTypes.AiLanguageModel],
|
||||
outputNames: ['Model'],
|
||||
credentials: [
|
||||
{
|
||||
name: 'databricksOAuth2Api',
|
||||
required: true,
|
||||
},
|
||||
],
|
||||
properties: [
|
||||
getConnectionHintNoticeField([NodeConnectionTypes.AiChain, NodeConnectionTypes.AiAgent]),
|
||||
{
|
||||
displayName:
|
||||
'If using JSON response format, you must include word "json" in the prompt in your chain or agent. Also, make sure the selected endpoint supports JSON mode.',
|
||||
name: 'notice',
|
||||
type: 'notice',
|
||||
default: '',
|
||||
displayOptions: {
|
||||
show: {
|
||||
'/options.responseFormat': ['json_object'],
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
displayName: 'Model',
|
||||
name: 'model',
|
||||
type: 'resourceLocator',
|
||||
default: { mode: 'list', value: '' },
|
||||
required: true,
|
||||
modes: [
|
||||
{
|
||||
displayName: 'From List',
|
||||
name: 'list',
|
||||
type: 'list',
|
||||
placeholder: 'Select a model...',
|
||||
typeOptions: {
|
||||
searchListMethod: 'searchModels',
|
||||
searchable: true,
|
||||
},
|
||||
},
|
||||
{
|
||||
displayName: 'ID',
|
||||
name: 'id',
|
||||
type: 'string',
|
||||
placeholder: 'my-serving-endpoint',
|
||||
},
|
||||
],
|
||||
description: 'The serving endpoint. Choose from the list, or specify an ID.',
|
||||
},
|
||||
{
|
||||
displayName: 'Options',
|
||||
name: 'options',
|
||||
placeholder: 'Add Option',
|
||||
description: 'Additional options to add',
|
||||
type: 'collection',
|
||||
default: {},
|
||||
options: [
|
||||
{
|
||||
displayName: 'Frequency Penalty',
|
||||
name: 'frequencyPenalty',
|
||||
default: 0,
|
||||
typeOptions: { maxValue: 2, minValue: -2, numberPrecision: 1 },
|
||||
description:
|
||||
"Positive values penalize new tokens based on their existing frequency in the text so far, decreasing the model's likelihood to repeat the same line verbatim",
|
||||
type: 'number',
|
||||
},
|
||||
{
|
||||
displayName: 'Maximum Number of Tokens',
|
||||
name: 'maxTokens',
|
||||
default: -1,
|
||||
description:
|
||||
'The maximum number of tokens to generate in the completion. Most models have a context length of 2048 tokens (except for the newest models, which support 32,768).',
|
||||
type: 'number',
|
||||
typeOptions: {
|
||||
maxValue: 32768,
|
||||
},
|
||||
},
|
||||
{
|
||||
displayName: 'Response Format',
|
||||
name: 'responseFormat',
|
||||
default: 'text',
|
||||
type: 'options',
|
||||
options: [
|
||||
{
|
||||
name: 'Text',
|
||||
value: 'text',
|
||||
description: 'Regular text response',
|
||||
},
|
||||
{
|
||||
name: 'JSON',
|
||||
value: 'json_object',
|
||||
description:
|
||||
'Enables JSON mode, which should guarantee the message the model generates is valid JSON',
|
||||
},
|
||||
],
|
||||
},
|
||||
{
|
||||
displayName: 'Presence Penalty',
|
||||
name: 'presencePenalty',
|
||||
default: 0,
|
||||
typeOptions: { maxValue: 2, minValue: -2, numberPrecision: 1 },
|
||||
description:
|
||||
"Positive values penalize new tokens based on whether they appear in the text so far, increasing the model's likelihood to talk about new topics",
|
||||
type: 'number',
|
||||
},
|
||||
{
|
||||
displayName: 'Sampling Temperature',
|
||||
name: 'temperature',
|
||||
default: 0.7,
|
||||
typeOptions: { maxValue: 2, minValue: 0, numberPrecision: 1 },
|
||||
description:
|
||||
'Controls randomness: Lowering results in less random completions. As the temperature approaches zero, the model will become deterministic and repetitive.',
|
||||
type: 'number',
|
||||
},
|
||||
{
|
||||
displayName: 'Timeout',
|
||||
name: 'timeout',
|
||||
default: 360000,
|
||||
description: 'Maximum amount of time a request is allowed to take in milliseconds',
|
||||
type: 'number',
|
||||
},
|
||||
{
|
||||
displayName: 'Max Retries',
|
||||
name: 'maxRetries',
|
||||
default: 2,
|
||||
description: 'Maximum number of retries to attempt',
|
||||
type: 'number',
|
||||
},
|
||||
{
|
||||
displayName: 'Top P',
|
||||
name: 'topP',
|
||||
default: 1,
|
||||
typeOptions: { maxValue: 1, minValue: 0, numberPrecision: 1 },
|
||||
description:
|
||||
'Controls diversity via nucleus sampling: 0.5 means half of all likelihood-weighted options are considered. We generally recommend altering this or temperature but not both.',
|
||||
type: 'number',
|
||||
},
|
||||
],
|
||||
},
|
||||
],
|
||||
};
|
||||
|
||||
async supplyData(this: ISupplyDataFunctions, itemIndex: number): Promise<SupplyData> {
|
||||
const credential = await this.getCredentials<DatabricksOAuth2Credential>('databricksOAuth2Api');
|
||||
|
||||
// ENT-381 adds user-grant support; until then the in-node mint only works
|
||||
// with a client secret
|
||||
if (credential.grantType === 'authorizationCode') {
|
||||
throw new NodeOperationError(
|
||||
this.getNode(),
|
||||
'User (Authorization Code) login is not supported by this node yet - use a Client Credentials (Service Principal) credential',
|
||||
);
|
||||
}
|
||||
assertHttpsHost(this, credential.host);
|
||||
|
||||
const baseURL = `${credential.host.replace(/\/$/, '')}/serving-endpoints`;
|
||||
|
||||
const modelName = this.getNodeParameter('model', itemIndex, '', {
|
||||
extractValue: true,
|
||||
}) as string;
|
||||
|
||||
const options = this.getNodeParameter('options', itemIndex, {}) as {
|
||||
frequencyPenalty?: number;
|
||||
maxTokens?: number;
|
||||
maxRetries: number;
|
||||
timeout: number;
|
||||
presencePenalty?: number;
|
||||
temperature?: number;
|
||||
topP?: number;
|
||||
responseFormat?: 'text' | 'json_object';
|
||||
};
|
||||
|
||||
const egressFilter = this.helpers.getSecureEgressFilter();
|
||||
|
||||
const timeout = options.timeout;
|
||||
const configuration: ClientOptions = {
|
||||
baseURL,
|
||||
fetch: createDatabricksFetch(
|
||||
getDatabricksTokenProvider(this.getNode(), credential, egressFilter),
|
||||
egressFilter,
|
||||
),
|
||||
fetchOptions: {
|
||||
dispatcher: getProxyAgent(
|
||||
baseURL,
|
||||
{
|
||||
headersTimeout: timeout,
|
||||
bodyTimeout: timeout,
|
||||
},
|
||||
egressFilter?.createSecureLookup(),
|
||||
),
|
||||
},
|
||||
};
|
||||
|
||||
const modelKwargs: Record<string, unknown> = {};
|
||||
if (options.responseFormat) {
|
||||
modelKwargs.response_format = { type: options.responseFormat };
|
||||
}
|
||||
|
||||
const model = new ChatOpenAI({
|
||||
// Placeholder only - the fetch wrapper overwrites the Authorization header
|
||||
apiKey: 'databricks-oauth',
|
||||
model: modelName,
|
||||
...options,
|
||||
timeout,
|
||||
maxRetries: options.maxRetries ?? 2,
|
||||
configuration,
|
||||
callbacks: [new N8nLlmTracing(this)],
|
||||
modelKwargs: Object.keys(modelKwargs).length > 0 ? modelKwargs : undefined,
|
||||
onFailedAttempt: makeN8nLlmFailedAttemptHandler(this, openAiFailedAttemptHandler),
|
||||
});
|
||||
|
||||
return {
|
||||
response: model,
|
||||
};
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,19 @@
|
|||
<svg version="1.1" id="Layer_1" xmlns:x="ns_extend;" xmlns:i="ns_ai;" xmlns:graph="ns_graphs;" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" x="0px" y="0px" viewBox="0 0 40.1 42" style="enable-background:new 0 0 40.1 42;" xml:space="preserve">
|
||||
<style type="text/css">
|
||||
.st0{fill:#FF3621;}
|
||||
</style>
|
||||
<metadata>
|
||||
<sfw xmlns="ns_sfw;">
|
||||
<slices>
|
||||
</slices>
|
||||
<sliceSourceBounds bottomLeftOrigin="true" height="42" width="40.1" x="-69.1" y="-10.5">
|
||||
</sliceSourceBounds>
|
||||
</sfw>
|
||||
</metadata>
|
||||
<g>
|
||||
<path class="st0" d="M40.1,31.1v-7.4l-0.8-0.5L20.1,33.7l-18.2-10l0-4.3l18.2,9.9l20.1-10.9v-7.3l-0.8-0.5L20.1,21.2L2.6,11.6
|
||||
L20.1,2l14.1,7.7l1.1-0.6V8.3L20.1,0L0,10.9V12L20.1,23l18.2-10v4.4l-18.2,10L0.8,16.8L0,17.3v7.4l20.1,10.9l18.2-9.9v4.3l-18.2,10
|
||||
L0.8,29.5L0,30v1.1L20.1,42L40.1,31.1z">
|
||||
</path>
|
||||
</g>
|
||||
</svg>
|
||||
|
After Width: | Height: | Size: 855 B |
|
|
@ -0,0 +1,19 @@
|
|||
<svg version="1.1" id="Layer_1" xmlns:x="ns_extend;" xmlns:i="ns_ai;" xmlns:graph="ns_graphs;" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" x="0px" y="0px" viewBox="0 0 40.1 42" style="enable-background:new 0 0 40.1 42;" xml:space="preserve">
|
||||
<style type="text/css">
|
||||
.st0{fill:#FF3621;}
|
||||
</style>
|
||||
<metadata>
|
||||
<sfw xmlns="ns_sfw;">
|
||||
<slices>
|
||||
</slices>
|
||||
<sliceSourceBounds bottomLeftOrigin="true" height="42" width="40.1" x="-69.1" y="-10.5">
|
||||
</sliceSourceBounds>
|
||||
</sfw>
|
||||
</metadata>
|
||||
<g>
|
||||
<path class="st0" d="M40.1,31.1v-7.4l-0.8-0.5L20.1,33.7l-18.2-10l0-4.3l18.2,9.9l20.1-10.9v-7.3l-0.8-0.5L20.1,21.2L2.6,11.6
|
||||
L20.1,2l14.1,7.7l1.1-0.6V8.3L20.1,0L0,10.9V12L20.1,23l18.2-10v4.4l-18.2,10L0.8,16.8L0,17.3v7.4l20.1,10.9l18.2-9.9v4.3l-18.2,10
|
||||
L0.8,29.5L0,30v1.1L20.1,42L40.1,31.1z">
|
||||
</path>
|
||||
</g>
|
||||
</svg>
|
||||
|
After Width: | Height: | Size: 855 B |
|
|
@ -0,0 +1,334 @@
|
|||
/* eslint-disable n8n-nodes-base/node-filename-against-convention */
|
||||
/* eslint-disable n8n-nodes-base/node-param-description-lowercase-first-char */
|
||||
/* eslint-disable @typescript-eslint/no-unsafe-member-access */
|
||||
/* eslint-disable @typescript-eslint/no-unsafe-assignment */
|
||||
|
||||
/* eslint-disable @typescript-eslint/unbound-method */
|
||||
import { ChatOpenAI } from '@langchain/openai';
|
||||
import { makeN8nLlmFailedAttemptHandler, getProxyAgent } from '@n8n/ai-utilities';
|
||||
import { createMockExecuteFunction } from 'n8n-nodes-base/test/nodes/Helpers';
|
||||
import type { ILoadOptionsFunctions, INode, ISupplyDataFunctions } from 'n8n-workflow';
|
||||
import { NodeOperationError } from 'n8n-workflow';
|
||||
import type { Mocked } from 'vitest';
|
||||
|
||||
import { LmChatDatabricks } from '../LmChatDatabricks.node';
|
||||
import { createDatabricksFetch, getDatabricksTokenProvider } from '../token-provider';
|
||||
|
||||
vi.mock('@langchain/openai');
|
||||
vi.mock('@n8n/ai-utilities');
|
||||
vi.mock('../token-provider');
|
||||
|
||||
const MockedChatOpenAI = vi.mocked(ChatOpenAI);
|
||||
const mockedMakeN8nLlmFailedAttemptHandler = vi.mocked(makeN8nLlmFailedAttemptHandler);
|
||||
const mockedGetProxyAgent = vi.mocked(getProxyAgent);
|
||||
const mockedGetDatabricksTokenProvider = vi.mocked(getDatabricksTokenProvider);
|
||||
const mockedCreateDatabricksFetch = vi.mocked(createDatabricksFetch);
|
||||
|
||||
const mockTokenProvider = vi.fn(async () => 'test-token');
|
||||
const mockFetch = vi.fn() as unknown as typeof fetch;
|
||||
|
||||
const mockCredential = {
|
||||
host: 'https://my.databricks.com/',
|
||||
grantType: 'clientCredentials',
|
||||
clientId: 'test-client-id',
|
||||
clientSecret: 'test-client-secret',
|
||||
accessTokenUrl: 'https://my.databricks.com/oidc/v1/token',
|
||||
scope: 'all-apis',
|
||||
authentication: 'header',
|
||||
};
|
||||
|
||||
describe('LmChatDatabricks', () => {
|
||||
let node: LmChatDatabricks;
|
||||
|
||||
const mockNodeDef: INode = {
|
||||
id: '1',
|
||||
name: 'Databricks Chat Model',
|
||||
typeVersion: 1,
|
||||
type: '@n8n/n8n-nodes-langchain.lmChatDatabricks',
|
||||
position: [0, 0],
|
||||
parameters: {},
|
||||
};
|
||||
|
||||
const setupMockContext = (credentialOverrides: Partial<typeof mockCredential> = {}) => {
|
||||
const ctx = createMockExecuteFunction<ISupplyDataFunctions>(
|
||||
{},
|
||||
mockNodeDef,
|
||||
) as Mocked<ISupplyDataFunctions>;
|
||||
|
||||
ctx.getCredentials = vi.fn().mockResolvedValue({ ...mockCredential, ...credentialOverrides });
|
||||
ctx.getNode = vi.fn().mockReturnValue(mockNodeDef);
|
||||
ctx.helpers.getSecureEgressFilter = vi.fn().mockReturnValue(undefined);
|
||||
ctx.getNodeParameter = vi.fn().mockImplementation((paramName: string) => {
|
||||
if (paramName === 'model') return 'my-chat-endpoint';
|
||||
if (paramName === 'options') return {};
|
||||
return undefined;
|
||||
});
|
||||
|
||||
mockedMakeN8nLlmFailedAttemptHandler.mockReturnValue(vi.fn());
|
||||
mockedGetProxyAgent.mockReturnValue({} as any);
|
||||
mockedGetDatabricksTokenProvider.mockReturnValue(mockTokenProvider);
|
||||
mockedCreateDatabricksFetch.mockReturnValue(mockFetch);
|
||||
return ctx;
|
||||
};
|
||||
|
||||
beforeEach(() => {
|
||||
node = new LmChatDatabricks();
|
||||
vi.clearAllMocks();
|
||||
});
|
||||
|
||||
describe('node description', () => {
|
||||
it('should have correct node properties', () => {
|
||||
expect(node.description).toMatchObject({
|
||||
displayName: 'Databricks Chat Model',
|
||||
name: 'lmChatDatabricks',
|
||||
group: ['transform'],
|
||||
version: [1],
|
||||
hidden: true,
|
||||
credentials: [{ name: 'databricksOAuth2Api', required: true }],
|
||||
outputs: ['ai_languageModel'],
|
||||
outputNames: ['Model'],
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe('supplyData', () => {
|
||||
it('should create ChatOpenAI with basic configuration', async () => {
|
||||
const ctx = setupMockContext();
|
||||
|
||||
const result = await node.supplyData.call(ctx, 0);
|
||||
|
||||
expect(ctx.getCredentials).toHaveBeenCalledWith('databricksOAuth2Api');
|
||||
expect(MockedChatOpenAI).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
apiKey: 'databricks-oauth',
|
||||
model: 'my-chat-endpoint',
|
||||
maxRetries: 2,
|
||||
callbacks: expect.arrayContaining([expect.any(Object)]),
|
||||
onFailedAttempt: expect.any(Function),
|
||||
}),
|
||||
);
|
||||
expect(result).toEqual({ response: expect.any(Object) });
|
||||
});
|
||||
|
||||
it('should strip the trailing slash from the host in baseURL', async () => {
|
||||
const ctx = setupMockContext();
|
||||
|
||||
await node.supplyData.call(ctx, 0);
|
||||
|
||||
const callArgs = MockedChatOpenAI.mock.calls[0][0];
|
||||
expect(callArgs?.configuration?.baseURL).toBe('https://my.databricks.com/serving-endpoints');
|
||||
});
|
||||
|
||||
it('should wire the token-provider fetch wrapper into ChatOpenAI', async () => {
|
||||
const ctx = setupMockContext();
|
||||
|
||||
await node.supplyData.call(ctx, 0);
|
||||
|
||||
expect(mockedGetDatabricksTokenProvider).toHaveBeenCalledWith(
|
||||
mockNodeDef,
|
||||
mockCredential,
|
||||
undefined,
|
||||
);
|
||||
expect(mockedCreateDatabricksFetch).toHaveBeenCalledWith(mockTokenProvider, undefined);
|
||||
const callArgs = MockedChatOpenAI.mock.calls[0][0];
|
||||
expect(callArgs?.configuration?.fetch).toBe(mockFetch);
|
||||
});
|
||||
|
||||
it('should thread the egress filter into the token provider and proxy agent', async () => {
|
||||
const ctx = setupMockContext();
|
||||
const secureLookup = vi.fn();
|
||||
const egressFilter = {
|
||||
validateUrl: vi.fn(),
|
||||
validateRedirectSync: vi.fn(),
|
||||
createSecureLookup: vi.fn().mockReturnValue(secureLookup),
|
||||
};
|
||||
ctx.helpers.getSecureEgressFilter = vi.fn().mockReturnValue(egressFilter);
|
||||
|
||||
await node.supplyData.call(ctx, 0);
|
||||
|
||||
expect(mockedGetDatabricksTokenProvider).toHaveBeenCalledWith(
|
||||
mockNodeDef,
|
||||
mockCredential,
|
||||
egressFilter,
|
||||
);
|
||||
expect(mockedCreateDatabricksFetch).toHaveBeenCalledWith(mockTokenProvider, egressFilter);
|
||||
expect(mockedGetProxyAgent).toHaveBeenCalledWith(
|
||||
'https://my.databricks.com/serving-endpoints',
|
||||
expect.any(Object),
|
||||
secureLookup,
|
||||
);
|
||||
});
|
||||
|
||||
it('should read the model via resourceLocator value extraction', async () => {
|
||||
const ctx = setupMockContext();
|
||||
|
||||
await node.supplyData.call(ctx, 0);
|
||||
|
||||
expect(ctx.getNodeParameter).toHaveBeenCalledWith('model', 0, '', { extractValue: true });
|
||||
});
|
||||
|
||||
it('should pass options to ChatOpenAI', async () => {
|
||||
const ctx = setupMockContext();
|
||||
ctx.getNodeParameter = vi.fn().mockImplementation((paramName: string) => {
|
||||
if (paramName === 'model') return 'my-chat-endpoint';
|
||||
if (paramName === 'options')
|
||||
return {
|
||||
temperature: 0.5,
|
||||
maxTokens: 2000,
|
||||
topP: 0.9,
|
||||
frequencyPenalty: 0.3,
|
||||
presencePenalty: 0.2,
|
||||
timeout: 60000,
|
||||
maxRetries: 5,
|
||||
};
|
||||
return undefined;
|
||||
});
|
||||
|
||||
await node.supplyData.call(ctx, 0);
|
||||
|
||||
expect(MockedChatOpenAI).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
temperature: 0.5,
|
||||
maxTokens: 2000,
|
||||
topP: 0.9,
|
||||
frequencyPenalty: 0.3,
|
||||
presencePenalty: 0.2,
|
||||
timeout: 60000,
|
||||
maxRetries: 5,
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it('should set response_format in modelKwargs when responseFormat is provided', async () => {
|
||||
const ctx = setupMockContext();
|
||||
ctx.getNodeParameter = vi.fn().mockImplementation((paramName: string) => {
|
||||
if (paramName === 'model') return 'my-chat-endpoint';
|
||||
if (paramName === 'options') return { responseFormat: 'json_object' };
|
||||
return undefined;
|
||||
});
|
||||
|
||||
await node.supplyData.call(ctx, 0);
|
||||
|
||||
expect(MockedChatOpenAI).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
modelKwargs: { response_format: { type: 'json_object' } },
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it('should reject authorizationCode credentials', async () => {
|
||||
const ctx = setupMockContext({ grantType: 'authorizationCode' });
|
||||
|
||||
await expect(node.supplyData.call(ctx, 0)).rejects.toThrow(NodeOperationError);
|
||||
expect(MockedChatOpenAI).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should reject non-https hosts', async () => {
|
||||
const ctx = setupMockContext({ host: 'http://my.databricks.com' });
|
||||
|
||||
await expect(node.supplyData.call(ctx, 0)).rejects.toThrow(NodeOperationError);
|
||||
expect(MockedChatOpenAI).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should accept an uppercase HTTPS scheme', async () => {
|
||||
const ctx = setupMockContext({ host: 'HTTPS://my.databricks.com' });
|
||||
|
||||
await node.supplyData.call(ctx, 0);
|
||||
|
||||
expect(MockedChatOpenAI).toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
||||
describe('searchModels', () => {
|
||||
const endpointsResponse = {
|
||||
endpoints: [
|
||||
{
|
||||
name: 'chat-endpoint',
|
||||
task: 'llm/v1/chat',
|
||||
config: { served_entities: [{ foundation_model: { name: 'llama-3' } }] },
|
||||
},
|
||||
{ name: 'embeddings-endpoint', task: 'llm/v1/embeddings' },
|
||||
{
|
||||
name: 'agent-endpoint',
|
||||
task: 'agent/v1/chat',
|
||||
config: { served_entities: [{ external_model: { name: 'gpt-4o' } }] },
|
||||
},
|
||||
{ name: 'custom-endpoint' },
|
||||
],
|
||||
};
|
||||
|
||||
let mockContext: Mocked<ILoadOptionsFunctions>;
|
||||
let httpRequestWithAuthentication: ReturnType<typeof vi.fn>;
|
||||
|
||||
const setupSearchContext = (host: string, response: unknown) => {
|
||||
httpRequestWithAuthentication = vi.fn().mockResolvedValue(response);
|
||||
mockContext = {
|
||||
getCredentials: vi.fn().mockResolvedValue({ ...mockCredential, host }),
|
||||
getNode: vi.fn().mockReturnValue(mockNodeDef),
|
||||
helpers: { httpRequestWithAuthentication },
|
||||
} as unknown as Mocked<ILoadOptionsFunctions>;
|
||||
};
|
||||
|
||||
it('should reject non-https hosts before requesting the endpoint list', async () => {
|
||||
setupSearchContext('http://my.databricks.com', endpointsResponse);
|
||||
|
||||
await expect(node.methods.listSearch.searchModels.call(mockContext)).rejects.toThrow(
|
||||
'Databricks host must use https',
|
||||
);
|
||||
expect(httpRequestWithAuthentication).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should list only chat-capable endpoints', async () => {
|
||||
setupSearchContext('https://my.databricks.com', endpointsResponse);
|
||||
|
||||
const result = await node.methods.listSearch.searchModels.call(mockContext);
|
||||
|
||||
expect(httpRequestWithAuthentication).toHaveBeenCalledWith(
|
||||
'databricksOAuth2Api',
|
||||
expect.objectContaining({
|
||||
method: 'GET',
|
||||
url: 'https://my.databricks.com/api/2.0/serving-endpoints',
|
||||
}),
|
||||
);
|
||||
expect(result.results).toEqual([
|
||||
{
|
||||
name: 'chat-endpoint',
|
||||
value: 'chat-endpoint',
|
||||
url: 'https://my.databricks.com/ml/endpoints/chat-endpoint',
|
||||
description: 'llama-3',
|
||||
},
|
||||
{
|
||||
name: 'agent-endpoint',
|
||||
value: 'agent-endpoint',
|
||||
url: 'https://my.databricks.com/ml/endpoints/agent-endpoint',
|
||||
description: 'gpt-4o',
|
||||
},
|
||||
]);
|
||||
});
|
||||
|
||||
it('should apply the substring filter', async () => {
|
||||
setupSearchContext('https://my.databricks.com', endpointsResponse);
|
||||
|
||||
const result = await node.methods.listSearch.searchModels.call(mockContext, 'agent');
|
||||
|
||||
expect(result.results).toEqual([expect.objectContaining({ name: 'agent-endpoint' })]);
|
||||
});
|
||||
|
||||
it('should return no results when the response has no endpoints field', async () => {
|
||||
setupSearchContext('https://my.databricks.com', {});
|
||||
|
||||
const result = await node.methods.listSearch.searchModels.call(mockContext);
|
||||
|
||||
expect(result).toEqual({ results: [] });
|
||||
});
|
||||
|
||||
it('should strip the trailing slash from the host in the request URL', async () => {
|
||||
setupSearchContext('https://my.databricks.com/', endpointsResponse);
|
||||
|
||||
await node.methods.listSearch.searchModels.call(mockContext);
|
||||
|
||||
const [, requestOptions] = httpRequestWithAuthentication.mock.calls[0];
|
||||
expect(requestOptions.url).toBe('https://my.databricks.com/api/2.0/serving-endpoints');
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
@ -0,0 +1,409 @@
|
|||
import type { INode, NodeEgressFilter } from 'n8n-workflow';
|
||||
import { NodeOperationError } from 'n8n-workflow';
|
||||
|
||||
import type { DatabricksOAuth2Credential } from '../token-provider';
|
||||
import { createDatabricksFetch, getDatabricksTokenProvider } from '../token-provider';
|
||||
|
||||
const { MockClientOAuth2, mockGetToken } = vi.hoisted(() => {
|
||||
const mockGetToken = vi.fn();
|
||||
|
||||
class MockClientOAuth2 {
|
||||
credentials = { getToken: mockGetToken };
|
||||
|
||||
constructor(readonly options: unknown) {
|
||||
MockClientOAuth2.init(options);
|
||||
}
|
||||
|
||||
static init = vi.fn();
|
||||
}
|
||||
|
||||
return { MockClientOAuth2, mockGetToken };
|
||||
});
|
||||
|
||||
vi.mock('@n8n/client-oauth2', () => ({
|
||||
ClientOAuth2: MockClientOAuth2,
|
||||
}));
|
||||
|
||||
const mockNode: INode = {
|
||||
id: '1',
|
||||
name: 'Databricks Chat Model',
|
||||
typeVersion: 1,
|
||||
type: '@n8n/n8n-nodes-langchain.lmChatDatabricks',
|
||||
position: [0, 0],
|
||||
parameters: {},
|
||||
};
|
||||
|
||||
const SENTINEL_SECRET = 'sentinel-client-secret-xyz';
|
||||
|
||||
const mockCredential: DatabricksOAuth2Credential = {
|
||||
host: 'https://my.databricks.com',
|
||||
grantType: 'clientCredentials',
|
||||
clientId: 'test-client-id',
|
||||
clientSecret: SENTINEL_SECRET,
|
||||
scope: 'all-apis',
|
||||
authentication: 'header',
|
||||
};
|
||||
|
||||
// Realistic mint failure: an Error instance shaped like the token endpoint's
|
||||
// error response (status + body own props)
|
||||
class MockResponseError extends Error {
|
||||
status = 401;
|
||||
|
||||
body = { error: 'invalid_client', error_description: 'Client authentication failed' };
|
||||
|
||||
constructor() {
|
||||
super('401: Client authentication failed (invalid_client)');
|
||||
}
|
||||
}
|
||||
|
||||
function tokenResponse(accessToken: string, expiresIn?: number | string) {
|
||||
return {
|
||||
accessToken,
|
||||
data: { access_token: accessToken, expires_in: expiresIn },
|
||||
};
|
||||
}
|
||||
|
||||
describe('getDatabricksTokenProvider', () => {
|
||||
afterEach(() => {
|
||||
vi.clearAllMocks();
|
||||
vi.useRealTimers();
|
||||
});
|
||||
|
||||
it('should mint a token via ClientOAuth2 built from the decrypted credential', async () => {
|
||||
mockGetToken.mockResolvedValue(tokenResponse('token-a', 3600));
|
||||
const getToken = getDatabricksTokenProvider(mockNode, mockCredential);
|
||||
|
||||
await expect(getToken()).resolves.toBe('token-a');
|
||||
expect(MockClientOAuth2.init).toHaveBeenCalledWith({
|
||||
clientId: 'test-client-id',
|
||||
clientSecret: SENTINEL_SECRET,
|
||||
accessTokenUri: 'https://my.databricks.com/oidc/v1/token',
|
||||
scopes: ['all-apis'],
|
||||
authentication: 'header',
|
||||
});
|
||||
});
|
||||
|
||||
it('should derive the mint URL from the host, ignoring a stored accessTokenUrl', async () => {
|
||||
mockGetToken.mockResolvedValue(tokenResponse('token-a', 3600));
|
||||
const poisoned = {
|
||||
...mockCredential,
|
||||
host: 'https://my.databricks.com/',
|
||||
accessTokenUrl: 'http://attacker.example/token',
|
||||
} as DatabricksOAuth2Credential;
|
||||
const getToken = getDatabricksTokenProvider(mockNode, poisoned);
|
||||
|
||||
await expect(getToken()).resolves.toBe('token-a');
|
||||
expect(MockClientOAuth2.init).toHaveBeenCalledWith(
|
||||
expect.objectContaining({ accessTokenUri: 'https://my.databricks.com/oidc/v1/token' }),
|
||||
);
|
||||
});
|
||||
|
||||
it('should pass the egress filter to ClientOAuth2 as its ssrfBridge', async () => {
|
||||
mockGetToken.mockResolvedValue(tokenResponse('token-a', 3600));
|
||||
const egressFilter = {
|
||||
validateUrl: vi.fn(),
|
||||
validateRedirectSync: vi.fn(),
|
||||
createSecureLookup: vi.fn(),
|
||||
};
|
||||
const getToken = getDatabricksTokenProvider(mockNode, mockCredential, egressFilter);
|
||||
|
||||
await expect(getToken()).resolves.toBe('token-a');
|
||||
expect(MockClientOAuth2.init).toHaveBeenCalledWith(
|
||||
expect.objectContaining({ ssrfBridge: egressFilter }),
|
||||
);
|
||||
});
|
||||
|
||||
it('should reuse the cached token before the 60s early-expiry buffer', async () => {
|
||||
vi.useFakeTimers();
|
||||
mockGetToken.mockResolvedValue(tokenResponse('token-a', 3600));
|
||||
const getToken = getDatabricksTokenProvider(mockNode, mockCredential);
|
||||
|
||||
await expect(getToken()).resolves.toBe('token-a');
|
||||
vi.advanceTimersByTime((3600 - 120) * 1000);
|
||||
await expect(getToken()).resolves.toBe('token-a');
|
||||
expect(mockGetToken).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
it('should re-mint within the last 60s of the token window', async () => {
|
||||
vi.useFakeTimers();
|
||||
mockGetToken
|
||||
.mockResolvedValueOnce(tokenResponse('token-a', 3600))
|
||||
.mockResolvedValueOnce(tokenResponse('token-b', 3600));
|
||||
const getToken = getDatabricksTokenProvider(mockNode, mockCredential);
|
||||
|
||||
await expect(getToken()).resolves.toBe('token-a');
|
||||
vi.advanceTimersByTime((3600 - 30) * 1000);
|
||||
await expect(getToken()).resolves.toBe('token-b');
|
||||
expect(mockGetToken).toHaveBeenCalledTimes(2);
|
||||
});
|
||||
|
||||
it('should re-mint every call when expires_in is missing', async () => {
|
||||
mockGetToken
|
||||
.mockResolvedValueOnce(tokenResponse('token-a'))
|
||||
.mockResolvedValueOnce(tokenResponse('token-b'));
|
||||
const getToken = getDatabricksTokenProvider(mockNode, mockCredential);
|
||||
|
||||
await expect(getToken()).resolves.toBe('token-a');
|
||||
await expect(getToken()).resolves.toBe('token-b');
|
||||
expect(mockGetToken).toHaveBeenCalledTimes(2);
|
||||
});
|
||||
|
||||
it('should cache when expires_in is a numeric string', async () => {
|
||||
mockGetToken.mockResolvedValue(tokenResponse('token-a', '3600'));
|
||||
const getToken = getDatabricksTokenProvider(mockNode, mockCredential);
|
||||
|
||||
await expect(getToken()).resolves.toBe('token-a');
|
||||
await expect(getToken()).resolves.toBe('token-a');
|
||||
expect(mockGetToken).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
it('should share a single in-flight mint between concurrent callers', async () => {
|
||||
let resolveMint!: (value: ReturnType<typeof tokenResponse>) => void;
|
||||
mockGetToken.mockImplementation(
|
||||
async () => await new Promise((resolve) => (resolveMint = resolve)),
|
||||
);
|
||||
const getToken = getDatabricksTokenProvider(mockNode, mockCredential);
|
||||
|
||||
const first = getToken();
|
||||
const second = getToken();
|
||||
resolveMint(tokenResponse('token-a', 3600));
|
||||
|
||||
await expect(first).resolves.toBe('token-a');
|
||||
await expect(second).resolves.toBe('token-a');
|
||||
expect(mockGetToken).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
it('should clear the cache slot on a rejected mint so the next call re-mints', async () => {
|
||||
mockGetToken
|
||||
.mockRejectedValueOnce(new MockResponseError())
|
||||
.mockResolvedValueOnce(tokenResponse('token-a', 3600));
|
||||
const getToken = getDatabricksTokenProvider(mockNode, mockCredential);
|
||||
|
||||
await expect(getToken()).rejects.toThrow(NodeOperationError);
|
||||
await expect(getToken()).resolves.toBe('token-a');
|
||||
expect(mockGetToken).toHaveBeenCalledTimes(2);
|
||||
});
|
||||
|
||||
it('should wrap mint failures in a NodeOperationError that leaks no secret', async () => {
|
||||
mockGetToken.mockRejectedValue(new MockResponseError());
|
||||
const getToken = getDatabricksTokenProvider(mockNode, mockCredential);
|
||||
|
||||
const error = await getToken().then(
|
||||
() => {
|
||||
throw new Error('expected getToken to reject');
|
||||
},
|
||||
(e: NodeOperationError) => e,
|
||||
);
|
||||
|
||||
expect(error).toBeInstanceOf(NodeOperationError);
|
||||
expect(error.message).toBe('Failed to retrieve Databricks access token');
|
||||
expect(error.description).toContain('Client authentication failed');
|
||||
|
||||
const serialized = JSON.stringify(error, Object.getOwnPropertyNames(error));
|
||||
for (const text of [serialized, error.message, error.description ?? '']) {
|
||||
expect(text).not.toContain(SENTINEL_SECRET);
|
||||
expect(text).not.toContain('Bearer ');
|
||||
}
|
||||
});
|
||||
|
||||
it('should scrub the client secret from an echoed mint error message', async () => {
|
||||
mockGetToken.mockRejectedValue(
|
||||
new Error(`400: invalid request "client_secret=${SENTINEL_SECRET}"`),
|
||||
);
|
||||
const getToken = getDatabricksTokenProvider(mockNode, mockCredential);
|
||||
|
||||
const error = await getToken().then(
|
||||
() => {
|
||||
throw new Error('expected getToken to reject');
|
||||
},
|
||||
(e: NodeOperationError) => e,
|
||||
);
|
||||
|
||||
expect(error.description).not.toContain(SENTINEL_SECRET);
|
||||
expect(error.description).toContain('***');
|
||||
});
|
||||
});
|
||||
|
||||
describe('createDatabricksFetch', () => {
|
||||
const origFetch = globalThis.fetch;
|
||||
|
||||
afterEach(() => {
|
||||
globalThis.fetch = origFetch;
|
||||
vi.clearAllMocks();
|
||||
vi.useRealTimers();
|
||||
});
|
||||
|
||||
it('should set the Authorization header, overwriting an existing one', async () => {
|
||||
const mockFetch = vi.fn().mockResolvedValue(new Response('ok'));
|
||||
globalThis.fetch = mockFetch;
|
||||
const wrappedFetch = createDatabricksFetch(async () => 'fresh-token');
|
||||
|
||||
await wrappedFetch('https://my.databricks.com/serving-endpoints/chat/completions', {
|
||||
method: 'POST',
|
||||
headers: { authorization: 'Bearer stale-token' },
|
||||
});
|
||||
|
||||
const [, init] = mockFetch.mock.calls[0] as [string, RequestInit];
|
||||
expect(new Headers(init.headers).get('authorization')).toBe('Bearer fresh-token');
|
||||
expect(init.method).toBe('POST');
|
||||
});
|
||||
|
||||
it('should preserve the headers of a Request input when init sets none', async () => {
|
||||
const mockFetch = vi.fn().mockResolvedValue(new Response('ok'));
|
||||
globalThis.fetch = mockFetch;
|
||||
const wrappedFetch = createDatabricksFetch(async () => 'fresh-token');
|
||||
|
||||
const request = new Request('https://my.databricks.com/serving-endpoints/chat/completions', {
|
||||
method: 'POST',
|
||||
headers: { 'content-type': 'application/json' },
|
||||
});
|
||||
await wrappedFetch(request);
|
||||
|
||||
const [, init] = mockFetch.mock.calls[0] as [Request, RequestInit];
|
||||
const sentHeaders = new Headers(init.headers);
|
||||
expect(sentHeaders.get('content-type')).toBe('application/json');
|
||||
expect(sentHeaders.get('authorization')).toBe('Bearer fresh-token');
|
||||
});
|
||||
|
||||
it('should preserve the method and body of a Request input passed without init', async () => {
|
||||
const mockFetch = vi.fn().mockResolvedValue(new Response('ok'));
|
||||
globalThis.fetch = mockFetch;
|
||||
const wrappedFetch = createDatabricksFetch(async () => 'fresh-token');
|
||||
|
||||
const request = new Request('https://my.databricks.com/serving-endpoints/chat/completions', {
|
||||
method: 'POST',
|
||||
body: JSON.stringify({ messages: [] }),
|
||||
});
|
||||
await wrappedFetch(request);
|
||||
|
||||
const [, init] = mockFetch.mock.calls[0] as [string, RequestInit];
|
||||
expect(init.method).toBe('POST');
|
||||
expect(Buffer.from(init.body as ArrayBuffer).toString()).toBe('{"messages":[]}');
|
||||
});
|
||||
|
||||
it('should return the exact Response instance with an unread body', async () => {
|
||||
const response = new Response('data: chunk\n\n');
|
||||
globalThis.fetch = vi.fn().mockResolvedValue(response);
|
||||
const wrappedFetch = createDatabricksFetch(async () => 'fresh-token');
|
||||
|
||||
const result = await wrappedFetch('https://my.databricks.com/serving-endpoints');
|
||||
|
||||
expect(result).toBe(response);
|
||||
expect(result.bodyUsed).toBe(false);
|
||||
});
|
||||
|
||||
it('should return an error response unmodified without throwing', async () => {
|
||||
const response = new Response('{"error":"denied"}', { status: 401 });
|
||||
globalThis.fetch = vi.fn().mockResolvedValue(response);
|
||||
const wrappedFetch = createDatabricksFetch(async () => 'fresh-token');
|
||||
|
||||
const result = await wrappedFetch('https://my.databricks.com/serving-endpoints');
|
||||
|
||||
expect(result).toBe(response);
|
||||
});
|
||||
|
||||
it('should propagate network failures without leaking the token', async () => {
|
||||
globalThis.fetch = vi.fn().mockRejectedValue(new TypeError('fetch failed'));
|
||||
const wrappedFetch = createDatabricksFetch(async () => 'minted-token-abc');
|
||||
|
||||
const error = await wrappedFetch('https://my.databricks.com/serving-endpoints').then(
|
||||
() => {
|
||||
throw new Error('expected wrapped fetch to reject');
|
||||
},
|
||||
(e: Error) => e,
|
||||
);
|
||||
|
||||
const serialized = JSON.stringify(error, Object.getOwnPropertyNames(error));
|
||||
expect(serialized).not.toContain('minted-token-abc');
|
||||
expect(serialized).not.toContain('Bearer ');
|
||||
});
|
||||
|
||||
it('should validate each redirect hop against the egress filter before following it', async () => {
|
||||
const mockFetch = vi
|
||||
.fn()
|
||||
.mockResolvedValueOnce(
|
||||
new Response(null, { status: 302, headers: { location: 'https://internal.evil/steal' } }),
|
||||
)
|
||||
.mockResolvedValue(new Response('ok'));
|
||||
globalThis.fetch = mockFetch;
|
||||
const validateUrl = vi
|
||||
.fn()
|
||||
.mockResolvedValueOnce({ ok: true })
|
||||
.mockResolvedValueOnce({ ok: false, error: new Error('egress blocked') });
|
||||
const wrappedFetch = createDatabricksFetch(async () => 'fresh-token', {
|
||||
validateUrl,
|
||||
} as unknown as NodeEgressFilter);
|
||||
|
||||
await expect(wrappedFetch('https://my.databricks.com/serving-endpoints')).rejects.toThrow(
|
||||
'egress blocked',
|
||||
);
|
||||
|
||||
expect(validateUrl).toHaveBeenNthCalledWith(2, 'https://internal.evil/steal');
|
||||
// The redirect target must never be fetched
|
||||
expect(mockFetch).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
it('should follow a redirect the egress filter allows', async () => {
|
||||
const finalResponse = new Response('ok');
|
||||
const mockFetch = vi
|
||||
.fn()
|
||||
.mockResolvedValueOnce(
|
||||
new Response(null, {
|
||||
status: 302,
|
||||
headers: { location: 'https://my.databricks.com/moved' },
|
||||
}),
|
||||
)
|
||||
.mockResolvedValueOnce(finalResponse);
|
||||
globalThis.fetch = mockFetch;
|
||||
const wrappedFetch = createDatabricksFetch(async () => 'fresh-token', {
|
||||
validateUrl: vi.fn().mockResolvedValue({ ok: true }),
|
||||
} as unknown as NodeEgressFilter);
|
||||
|
||||
const result = await wrappedFetch('https://my.databricks.com/serving-endpoints');
|
||||
|
||||
expect(result).toBe(finalResponse);
|
||||
expect(mockFetch.mock.calls[1][0]).toEqual(new URL('https://my.databricks.com/moved'));
|
||||
});
|
||||
|
||||
it('should not send the bearer token to a cross-origin redirect target', async () => {
|
||||
const mockFetch = vi
|
||||
.fn()
|
||||
.mockResolvedValueOnce(
|
||||
new Response(null, {
|
||||
status: 307,
|
||||
headers: { location: 'https://other-allowed.example.com/endpoint' },
|
||||
}),
|
||||
)
|
||||
.mockResolvedValueOnce(new Response('ok'));
|
||||
globalThis.fetch = mockFetch;
|
||||
const wrappedFetch = createDatabricksFetch(async () => 'fresh-token', {
|
||||
validateUrl: vi.fn().mockResolvedValue({ ok: true }),
|
||||
} as unknown as NodeEgressFilter);
|
||||
|
||||
await wrappedFetch('https://my.databricks.com/serving-endpoints');
|
||||
|
||||
const hop2Headers = new Headers((mockFetch.mock.calls[1][1] as RequestInit).headers);
|
||||
expect(hop2Headers.get('authorization')).toBeNull();
|
||||
});
|
||||
|
||||
it('should send a rotated token after the previous one expires mid-execution', async () => {
|
||||
vi.useFakeTimers();
|
||||
const mockFetch = vi.fn().mockResolvedValue(new Response('ok'));
|
||||
globalThis.fetch = mockFetch;
|
||||
mockGetToken
|
||||
.mockResolvedValueOnce(tokenResponse('token-a', 3600))
|
||||
.mockResolvedValueOnce(tokenResponse('token-b', 3600));
|
||||
const wrappedFetch = createDatabricksFetch(
|
||||
getDatabricksTokenProvider(mockNode, mockCredential),
|
||||
);
|
||||
|
||||
await wrappedFetch('https://my.databricks.com/serving-endpoints');
|
||||
vi.advanceTimersByTime((3600 - 30) * 1000);
|
||||
await wrappedFetch('https://my.databricks.com/serving-endpoints');
|
||||
|
||||
const sentAuth = mockFetch.mock.calls.map((call) =>
|
||||
new Headers((call[1] as RequestInit).headers).get('authorization'),
|
||||
);
|
||||
expect(sentAuth).toEqual(['Bearer token-a', 'Bearer token-b']);
|
||||
});
|
||||
});
|
||||
|
|
@ -0,0 +1,123 @@
|
|||
import { fetchFollowingRedirects } from '@n8n/ai-utilities';
|
||||
import { ClientOAuth2 } from '@n8n/client-oauth2';
|
||||
import type { INode, NodeEgressFilter } from 'n8n-workflow';
|
||||
import { NodeOperationError } from 'n8n-workflow';
|
||||
|
||||
export interface DatabricksOAuth2Credential {
|
||||
host: string;
|
||||
grantType: 'clientCredentials' | 'authorizationCode';
|
||||
clientId: string;
|
||||
clientSecret: string;
|
||||
scope?: string;
|
||||
authentication?: 'header' | 'body';
|
||||
}
|
||||
|
||||
/**
|
||||
* Mints Databricks service-principal tokens on demand. Concurrent callers
|
||||
* share one in-flight mint, and tokens re-mint 60s before expiry so requests
|
||||
* near the end of the token window don't fail.
|
||||
*
|
||||
* Self-mints (instead of using core's `oauthTokenData`) because the model
|
||||
* client bypasses core's refresh machinery, so a core-issued token would
|
||||
* expire mid-run. The mint URL is derived from the https-validated `host`
|
||||
* (matching the credential's default) so a stored `accessTokenUrl` cannot
|
||||
* redirect the client secret elsewhere.
|
||||
*/
|
||||
export function getDatabricksTokenProvider(
|
||||
node: INode,
|
||||
credential: DatabricksOAuth2Credential,
|
||||
egressFilter?: NodeEgressFilter,
|
||||
): () => Promise<string> {
|
||||
const tokenUrl = `${credential.host.replace(/\/$/, '')}/oidc/v1/token`;
|
||||
|
||||
let cached: Promise<string> | undefined;
|
||||
let expiresAt = 0;
|
||||
|
||||
const mint = async (): Promise<string> => {
|
||||
try {
|
||||
// The bridge enforces the egress policy inside the client: pre-flight
|
||||
// URL validation, connect-time DNS pinning, and per-redirect checks
|
||||
const oAuthClient = new ClientOAuth2({
|
||||
clientId: credential.clientId,
|
||||
clientSecret: credential.clientSecret,
|
||||
accessTokenUri: tokenUrl,
|
||||
scopes: credential.scope?.split(' '),
|
||||
authentication: credential.authentication,
|
||||
ssrfBridge: egressFilter,
|
||||
});
|
||||
const token = await oAuthClient.credentials.getToken();
|
||||
const expiresIn = Number(token.data.expires_in);
|
||||
// ponytail: early-expiry buffer only; if server-side revocation mid-run
|
||||
// ever matters, add invalidate-and-retry-once on 401/403 in createDatabricksFetch
|
||||
expiresAt = Number.isNaN(expiresIn) ? 0 : Date.now() + (expiresIn - 60) * 1000;
|
||||
return token.accessToken;
|
||||
} catch (error) {
|
||||
cached = undefined;
|
||||
expiresAt = 0;
|
||||
// The description is persisted in execution error JSON; scrub the secret
|
||||
// in case the token endpoint or a proxy echoes the request back
|
||||
const message = error instanceof Error ? error.message : undefined;
|
||||
throw new NodeOperationError(node, 'Failed to retrieve Databricks access token', {
|
||||
description:
|
||||
message && credential.clientSecret
|
||||
? message.replaceAll(credential.clientSecret, '***')
|
||||
: message,
|
||||
});
|
||||
}
|
||||
};
|
||||
|
||||
return async () => {
|
||||
if (!cached || Date.now() >= expiresAt) {
|
||||
// Infinity until the mint resolves, so concurrent first callers join it
|
||||
expiresAt = Infinity;
|
||||
cached = mint();
|
||||
}
|
||||
return await cached;
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Wraps fetch to inject a fresh bearer token per request. Never reads or
|
||||
* clones the body, so streaming responses pass through untouched. Redirects
|
||||
* are followed manually so every hop is validated against the egress filter
|
||||
* before the token is sent to it, matching the MCP client's fetch wrapper;
|
||||
* the redirect helper also drops the bearer on cross-origin hops.
|
||||
*/
|
||||
export function createDatabricksFetch(
|
||||
getToken: () => Promise<string>,
|
||||
egressFilter?: NodeEgressFilter,
|
||||
): typeof globalThis.fetch {
|
||||
return async (input, init) => {
|
||||
// Passing headers in init replaces a Request input's own headers, so
|
||||
// carry those over when init doesn't set any
|
||||
const headers = new Headers(
|
||||
init?.headers ?? (input instanceof Request ? input.headers : undefined),
|
||||
);
|
||||
headers.set('authorization', `Bearer ${await getToken()}`);
|
||||
// The redirect loop takes a URL, so unwrap a Request input and carry its
|
||||
// method/body/signal over (init still wins, per fetch spec). The body is
|
||||
// buffered, which also lets 307/308 hops replay it.
|
||||
const requestInit: RequestInit = { ...init };
|
||||
if (input instanceof Request) {
|
||||
requestInit.method ??= input.method;
|
||||
requestInit.signal ??= input.signal;
|
||||
if (requestInit.body === undefined && input.body) {
|
||||
requestInit.body = await input.arrayBuffer();
|
||||
}
|
||||
}
|
||||
const startUrl = input instanceof Request ? input.url : input;
|
||||
return await fetchFollowingRedirects(
|
||||
fetch,
|
||||
startUrl,
|
||||
{ ...requestInit, headers },
|
||||
{
|
||||
onBeforeHop: async (hopUrl) => {
|
||||
if (egressFilter) {
|
||||
const result = await egressFilter.validateUrl(hopUrl);
|
||||
if (!result.ok) throw result.error;
|
||||
}
|
||||
},
|
||||
},
|
||||
);
|
||||
};
|
||||
}
|
||||
|
|
@ -38,6 +38,7 @@ const mockedProxyFetch = proxyFetch as MockedFunction<typeof proxyFetch>;
|
|||
const createTestEgressFilter = (): NodeEgressFilter => ({
|
||||
validateUrl: vi.fn().mockResolvedValue(createResultOk(undefined)),
|
||||
createSecureLookup: vi.fn().mockReturnValue(vi.fn()),
|
||||
validateRedirectSync: vi.fn(),
|
||||
});
|
||||
|
||||
const egressHelpers = <T extends { helpers: unknown }>(): Partial<T> =>
|
||||
|
|
|
|||
|
|
@ -38,6 +38,7 @@ const MockedClient = Client as MockedClass<typeof Client>;
|
|||
const createTestEgressFilter = (): NodeEgressFilter => ({
|
||||
validateUrl: vi.fn().mockResolvedValue(createResultOk(undefined)),
|
||||
createSecureLookup: vi.fn(),
|
||||
validateRedirectSync: vi.fn(),
|
||||
});
|
||||
|
||||
describe('utils', () => {
|
||||
|
|
@ -669,6 +670,35 @@ describe('utils', () => {
|
|||
expect(new Headers(options.init?.headers).get('authorization')).toBe('Bearer my-token');
|
||||
});
|
||||
|
||||
it('should not send auth headers to a cross-origin redirect target', async () => {
|
||||
mockClient.connect.mockResolvedValue(undefined);
|
||||
mockedProxyFetch
|
||||
.mockResolvedValueOnce(
|
||||
new Response(null, {
|
||||
status: 307,
|
||||
headers: { location: 'https://other.example.com/moved' },
|
||||
}),
|
||||
)
|
||||
.mockResolvedValueOnce(new Response('ok', { status: 200 }));
|
||||
|
||||
await connectMcpClient({
|
||||
serverTransport: transport,
|
||||
secureEgressFilter: createTestEgressFilter(),
|
||||
endpointUrl: 'https://example.com',
|
||||
headers: { Authorization: 'Bearer my-token' },
|
||||
name: 'test-client',
|
||||
version: 1,
|
||||
});
|
||||
|
||||
const [, opts] = (TransportClass as Mock).mock.calls[0];
|
||||
await opts.fetch('https://example.com/mcp', {});
|
||||
|
||||
expect(mockedProxyFetch).toHaveBeenCalledTimes(2);
|
||||
const [secondCallOptions] = mockedProxyFetch.mock.calls[1];
|
||||
expect(secondCallOptions.input.toString()).toBe('https://other.example.com/moved');
|
||||
expect(new Headers(secondCallOptions.init?.headers).get('authorization')).toBeNull();
|
||||
});
|
||||
|
||||
it('should preserve SDK headers passed as Headers instance', async () => {
|
||||
mockClient.connect.mockResolvedValue(undefined);
|
||||
|
||||
|
|
@ -851,6 +881,7 @@ describe('utils', () => {
|
|||
const secureLookup = vi.fn();
|
||||
const egressFilter: NodeEgressFilter = {
|
||||
validateUrl: vi.fn().mockResolvedValue(createResultError(new Error('Egress blocked'))),
|
||||
validateRedirectSync: vi.fn(),
|
||||
createSecureLookup: vi.fn().mockReturnValue(secureLookup),
|
||||
};
|
||||
|
||||
|
|
@ -884,6 +915,7 @@ describe('utils', () => {
|
|||
const secureLookup = vi.fn();
|
||||
const egressFilter: NodeEgressFilter = {
|
||||
validateUrl: vi.fn().mockResolvedValue(createResultOk(undefined)),
|
||||
validateRedirectSync: vi.fn(),
|
||||
createSecureLookup: vi.fn().mockReturnValue(secureLookup),
|
||||
};
|
||||
|
||||
|
|
|
|||
|
|
@ -49,6 +49,7 @@ const baseConnectionConfig: McpConnectionConfig = {
|
|||
const createTestEgressFilter = (): NodeEgressFilter => ({
|
||||
validateUrl: vi.fn().mockResolvedValue(createResultOk(undefined)),
|
||||
createSecureLookup: vi.fn(),
|
||||
validateRedirectSync: vi.fn(),
|
||||
});
|
||||
|
||||
const egressHelpers = <T extends { helpers: unknown }>(): Partial<T> =>
|
||||
|
|
|
|||
|
|
@ -297,7 +297,9 @@ function headersToRecord(headers: HeadersInit | undefined): Record<string, strin
|
|||
* - validates the initial URL and every redirect hop against `allowedDomains`
|
||||
* so credentials are never sent to a host the credential doesn't allow,
|
||||
* - validates the initial URL and every redirect hop against the instance
|
||||
* `secureEgressFilter`, and pins the connection to the validated address.
|
||||
* `secureEgressFilter`, and pins the connection to the validated address,
|
||||
* - withholds the auth headers once a redirect crosses origins, so
|
||||
* credentials never reach a host other than the one the request started on.
|
||||
*/
|
||||
function createAuthFetch(
|
||||
initialHeaders: Record<string, string> | undefined,
|
||||
|
|
|
|||
|
|
@ -124,6 +124,7 @@
|
|||
"dist/nodes/llms/LmChatAzureOpenAi/LmChatAzureOpenAi.node.js",
|
||||
"dist/nodes/llms/LmChatAwsBedrock/LmChatAwsBedrock.node.js",
|
||||
"dist/nodes/llms/LmChatCohere/LmChatCohere.node.js",
|
||||
"dist/nodes/llms/LmChatDatabricks/LmChatDatabricks.node.js",
|
||||
"dist/nodes/llms/LmChatDeepSeek/LmChatDeepSeek.node.js",
|
||||
"dist/nodes/llms/LmChatGoogleGemini/LmChatGoogleGemini.node.js",
|
||||
"dist/nodes/llms/LmChatGoogleVertex/LmChatGoogleVertex.node.js",
|
||||
|
|
|
|||
|
|
@ -194,6 +194,25 @@ describe('createAuthFetch — allowedDomains', () => {
|
|||
expect(baseFetchMock).toHaveBeenCalledTimes(2);
|
||||
});
|
||||
|
||||
it('withholds auth headers after a redirect hop crosses origins', async () => {
|
||||
baseFetchMock
|
||||
.mockResolvedValueOnce(makeRedirect('https://other.test/moved'))
|
||||
.mockResolvedValueOnce(makeOk());
|
||||
|
||||
const fetchFn = createAuthFetch({
|
||||
baseFetch,
|
||||
initialHeaders: { Authorization: 'Bearer A' },
|
||||
allowedDomains: { mode: 'domains', domains: 'example.test,other.test' },
|
||||
});
|
||||
await fetchFn('https://example.test/mcp');
|
||||
|
||||
expect(baseFetchMock).toHaveBeenCalledTimes(2);
|
||||
expect(new Headers(baseFetchMock.mock.calls[0][1].headers).get('authorization')).toBe(
|
||||
'Bearer A',
|
||||
);
|
||||
expect(new Headers(baseFetchMock.mock.calls[1][1].headers).get('authorization')).toBeNull();
|
||||
});
|
||||
|
||||
it('blocks all requests when mode is none', async () => {
|
||||
const fetchFn = createAuthFetch({
|
||||
baseFetch,
|
||||
|
|
|
|||
|
|
@ -60,7 +60,9 @@ function assertDomainPolicyAllowsUrl(url: string, policy: AuthFetchDomainPolicy)
|
|||
* HTTP_PROXY settings apply uniformly),
|
||||
* 2. injects the latest auth headers on every request,
|
||||
* 3. on a single 401, calls `onUnauthorized` to refresh the token and
|
||||
* retries the request once with the new headers.
|
||||
* retries the request once with the new headers,
|
||||
* 4. when a domain policy is set, follows redirects manually, validating
|
||||
* every hop and withholding the auth headers once a hop crosses origins.
|
||||
*
|
||||
* This mirrors the langchain MCP node's `createAuthFetch` so an agent's MCP
|
||||
* connection behaves identically to one configured via the workflow editor.
|
||||
|
|
|
|||
|
|
@ -1040,6 +1040,7 @@ describe('Kafka Utils', () => {
|
|||
const validateUrl = vi.fn().mockResolvedValue(result);
|
||||
const filter: NodeEgressFilter = {
|
||||
validateUrl,
|
||||
validateRedirectSync: vi.fn(),
|
||||
createSecureLookup: () => lookup,
|
||||
};
|
||||
return { filter, lookup, validateUrl };
|
||||
|
|
|
|||
|
|
@ -953,14 +953,16 @@ interface NodeHelperFunctions {
|
|||
|
||||
/**
|
||||
* Egress filter exposed to nodes whose embedded HTTP clients cannot go through
|
||||
* `httpRequest`. Mirrors the two layers n8n's own egress uses: a pre-flight URL
|
||||
* validation and a connect-time secure DNS lookup.
|
||||
* `httpRequest`. Mirrors the layers n8n's own egress uses: a pre-flight URL
|
||||
* validation, a connect-time secure DNS lookup, and per-redirect validation.
|
||||
*/
|
||||
export interface NodeEgressFilter {
|
||||
/** Validate a target URL before any connection. Resolves hostnames; direct IP literals are checked without DNS. */
|
||||
validateUrl(url: string | URL): Promise<Result<void, Error>>;
|
||||
/** DNS lookup drop-in that validates resolved addresses against the configured egress rules. */
|
||||
createSecureLookup(): LookupFunction;
|
||||
/** Validate a redirect hop synchronously; throws when the target is not allowed. */
|
||||
validateRedirectSync(url: string): void;
|
||||
}
|
||||
|
||||
export interface RequestHelperFunctions {
|
||||
|
|
|
|||
Loading…
Reference in New Issue
Block a user