diff --git a/src/frontend/src/CustomNodes/GenericNode/components/handleRenderComponent/__tests__/handleMutedState.test.ts b/src/frontend/src/CustomNodes/GenericNode/components/handleRenderComponent/__tests__/handleMutedState.test.ts new file mode 100644 index 0000000000..ac09c298f7 --- /dev/null +++ b/src/frontend/src/CustomNodes/GenericNode/components/handleRenderComponent/__tests__/handleMutedState.test.ts @@ -0,0 +1,232 @@ +/** + * Unit tests for muted handle state in HandleRenderComponent. + * + * Model-type handles should be visually muted (invisible) when: + * - No wire is connected to the handle + * - No connection filter is active (drag/connection mode) + * + * They should be visible when: + * - A wire is connected + * - A connection filter is active (e.g., "Connect other models") + * - The handle's own node initiated connection mode + */ + +describe("model handle muted state", () => { + describe("isMuted computation logic", () => { + // Pure logic extracted from HandleRenderComponent useMemo + const computeIsMuted = ( + idType: string | undefined, + connectedEdge: boolean, + filterPresent: boolean, + isInConnectionMode: boolean = false, + ): boolean => { + const isModelType = idType === "model"; + return ( + isModelType && !connectedEdge && !filterPresent && !isInConnectionMode + ); + }; + + it("should be muted when model type has no connection, no filter, and not in connection mode", () => { + // Arrange & Act + const result = computeIsMuted("model", false, false, false); + + // Assert + expect(result).toBe(true); + }); + + it("should not be muted when model type has a connected edge", () => { + // Arrange & Act + const result = computeIsMuted("model", true, false, false); + + // Assert + expect(result).toBe(false); + }); + + it("should not be muted when model type has an active filter", () => { + // Arrange & Act + const result = computeIsMuted("model", false, true, false); + + // Assert + expect(result).toBe(false); + }); + + it("should not be muted when model type has both connection and filter", () => { + // Arrange & Act + const result = computeIsMuted("model", true, true, false); + + // Assert + expect(result).toBe(false); + }); + + it("should not be muted when in connection mode (Connect other models selected)", () => { + // Arrange & Act — no edge, no filter, but connection mode is active + const result = computeIsMuted("model", false, false, true); + + // Assert — handle must stay visible so user can drag a wire to it + expect(result).toBe(false); + }); + + it("should never be muted for non-model types", () => { + // Arrange & Act & Assert + expect(computeIsMuted("str", false, false, false)).toBe(false); + expect(computeIsMuted("Message", false, false, false)).toBe(false); + expect(computeIsMuted(undefined, false, false, false)).toBe(false); + expect(computeIsMuted("", false, false, false)).toBe(false); + }); + }); + + describe("isOwnModelConnectionMode logic", () => { + // Pure logic extracted from HandleRenderComponent useMemo + const computeIsOwnModelConnectionMode = ( + idType: string | undefined, + left: boolean, + filterTargetNodeId: string | undefined, + nodeId: string, + ): boolean => { + return idType === "model" && left && filterTargetNodeId === nodeId; + }; + + const NODE_ID = "agent-node-123"; + + it("should be true when filter targets this node's model input handle", () => { + // Arrange & Act + const result = computeIsOwnModelConnectionMode( + "model", + true, + NODE_ID, + NODE_ID, + ); + + // Assert + expect(result).toBe(true); + }); + + it("should be false when filter targets a different node", () => { + // Arrange & Act + const result = computeIsOwnModelConnectionMode( + "model", + true, + "other-node-456", + NODE_ID, + ); + + // Assert + expect(result).toBe(false); + }); + + it("should be false for non-model types", () => { + // Arrange & Act + const result = computeIsOwnModelConnectionMode( + "str", + true, + NODE_ID, + NODE_ID, + ); + + // Assert + expect(result).toBe(false); + }); + + it("should be false for right-side (output) handles", () => { + // Arrange & Act + const result = computeIsOwnModelConnectionMode( + "model", + false, + NODE_ID, + NODE_ID, + ); + + // Assert + expect(result).toBe(false); + }); + + it("should be false when no filter is active", () => { + // Arrange & Act + const result = computeIsOwnModelConnectionMode( + "model", + true, + undefined, + NODE_ID, + ); + + // Assert + expect(result).toBe(false); + }); + }); + + describe("muted handle visual styling", () => { + // Pure logic extracted from HandleContent contentStyle useMemo + const computeContentStyle = ( + isNullHandle: boolean, + isMuted: boolean, + ): { width: string; height: string; opacity: number } => { + return { + width: isMuted && !isNullHandle ? "6px" : "10px", + height: isMuted && !isNullHandle ? "6px" : "10px", + opacity: isMuted && !isNullHandle ? 0 : 1, + }; + }; + + it("should render invisible when muted and not null", () => { + // Arrange & Act + const style = computeContentStyle(false, true); + + // Assert + expect(style.width).toBe("6px"); + expect(style.height).toBe("6px"); + expect(style.opacity).toBe(0); + }); + + it("should render full size when not muted and not null", () => { + // Arrange & Act + const style = computeContentStyle(false, false); + + // Assert + expect(style.width).toBe("10px"); + expect(style.height).toBe("10px"); + expect(style.opacity).toBe(1); + }); + + it("should use null handle styling when isNullHandle overrides isMuted", () => { + // isNullHandle takes priority — size stays 10px, opacity stays 1 + // (actual null handle styling uses different background/border, tested elsewhere) + + // Arrange & Act + const style = computeContentStyle(true, true); + + // Assert + expect(style.width).toBe("10px"); + expect(style.height).toBe("10px"); + expect(style.opacity).toBe(1); + }); + }); + + describe("muted handle neon shadow", () => { + const computeNeonShadow = ( + isNullHandle: boolean, + isMuted: boolean, + isActive: boolean, + ): string => { + if (isNullHandle || isMuted) return "none"; + if (!isActive) return "ring"; + return "glow"; + }; + + it("should return none when muted", () => { + expect(computeNeonShadow(false, true, false)).toBe("none"); + expect(computeNeonShadow(false, true, true)).toBe("none"); + }); + + it("should return none when null handle", () => { + expect(computeNeonShadow(true, false, false)).toBe("none"); + }); + + it("should return ring when not muted and not active", () => { + expect(computeNeonShadow(false, false, false)).toBe("ring"); + }); + + it("should return glow when not muted and active", () => { + expect(computeNeonShadow(false, false, true)).toBe("glow"); + }); + }); +}); diff --git a/src/frontend/src/CustomNodes/GenericNode/components/handleRenderComponent/index.tsx b/src/frontend/src/CustomNodes/GenericNode/components/handleRenderComponent/index.tsx index 2900642144..630f2846e1 100644 --- a/src/frontend/src/CustomNodes/GenericNode/components/handleRenderComponent/index.tsx +++ b/src/frontend/src/CustomNodes/GenericNode/components/handleRenderComponent/index.tsx @@ -3,6 +3,7 @@ import { memo, useCallback, useEffect, useMemo, useState } from "react"; import { useShallow } from "zustand/react/shallow"; import { useDarkStore } from "@/stores/darkStore"; import useFlowStore from "@/stores/flowStore"; +import type { NodeDataType } from "@/types/flow"; import { nodeColorsName } from "@/utils/styleUtils"; import ShadTooltip from "../../../../components/common/shadTooltipComponent"; import { @@ -24,6 +25,7 @@ const BASE_HANDLE_STYLES = { const HandleContent = memo(function HandleContent({ isNullHandle, + isMuted, handleColor, accentForegroundColorName, isHovered, @@ -35,6 +37,7 @@ const HandleContent = memo(function HandleContent({ nodeId, }: { isNullHandle: boolean; + isMuted: boolean; handleColor: string; accentForegroundColorName: string; isHovered: boolean; @@ -97,7 +100,7 @@ const HandleContent = memo(function HandleContent({ const getNeonShadow = useCallback( (color: string, isActive: boolean) => { - if (isNullHandle) return "none"; + if (isNullHandle || isMuted) return "none"; if (!isActive) return `0 0 0 3px ${color}`; return [ "0 0 0 1px hsl(var(--border))", @@ -110,27 +113,29 @@ const HandleContent = memo(function HandleContent({ `0 0 20px ${color}`, ].join(", "); }, - [isNullHandle], + [isNullHandle, isMuted], ); const contentStyle = useMemo( () => ({ background: isNullHandle ? "hsl(var(--border))" : handleColor, - width: "10px", - height: "10px", + width: isMuted && !isNullHandle ? "6px" : "10px", + height: isMuted && !isNullHandle ? "6px" : "10px", transition: "all 0.2s", + opacity: isMuted && !isNullHandle ? 0 : 1, boxShadow: getNeonShadow( accentForegroundColorName, isHovered || openHandle, ), animation: - (isHovered || openHandle) && !isNullHandle + (isHovered || openHandle) && !isNullHandle && !isMuted ? `pulseNeon-${nodeId} 1.1s ease-in-out infinite` : "none", border: isNullHandle ? "2px solid hsl(var(--muted))" : "none", }), [ isNullHandle, + isMuted, handleColor, getNeonShadow, accentForegroundColorName, @@ -186,6 +191,18 @@ const HandleRenderComponent = memo(function HandleRenderComponent({ const edges = useFlowStore((state) => state.edges); + // Check if this node is in "connect other models" mode + const isInConnectionMode = useFlowStore( + useCallback( + (state) => { + if (id?.type !== "model" || !left) return false; + const node = state.nodes.find((n) => n.id === nodeId); + return (node?.data as NodeDataType)?._connectionMode === true; + }, + [nodeId, id?.type, left], + ), + ); + const { setHandleDragging, setFilterType, @@ -234,6 +251,7 @@ const HandleRenderComponent = memo(function HandleRenderComponent({ filterPresent, currentFilter, isNullHandle, + isMuted, handleColor, accentForegroundColorName, } = useMemo(() => { @@ -276,8 +294,18 @@ const HandleRenderComponent = memo(function HandleRenderComponent({ const outputType = connectedEdge?.data?.sourceHandle?.output_types?.[0]; const connectedColor = (outputType && nodeColorsName[outputType]) || "gray"; + // Model handles that initiated connection mode on this node should not be nulled + const isOwnModelConnectionMode = + id?.type === "model" && left && filterType?.target === nodeId; + const isNullHandle = - filterPresent && !(openHandle || ownDraggingHandle || ownFilterHandle); + filterPresent && + !( + openHandle || + ownDraggingHandle || + ownFilterHandle || + isOwnModelConnectionMode + ); // Create a Set from colorName to remove duplicates const colorNameSet = new Set(colorName || []); @@ -325,6 +353,10 @@ const HandleRenderComponent = memo(function HandleRenderComponent({ color: handleColorName, }; + const isModelType = id?.type === "model"; + const isMuted = + isModelType && !connectedEdge && !filterPresent && !isInConnectionMode; + return { sameNode: sameDraggingNode || sameFilterNode, ownHandle: ownDraggingHandle || ownFilterHandle, @@ -334,6 +366,7 @@ const HandleRenderComponent = memo(function HandleRenderComponent({ filterPresent, currentFilter, isNullHandle, + isMuted, handleColor, }; }, [ @@ -347,6 +380,8 @@ const HandleRenderComponent = memo(function HandleRenderComponent({ colorName, tooltipTitle, edges, + id, + isInConnectionMode, ]); const handleMouseDown = useCallback( @@ -441,6 +476,7 @@ const HandleRenderComponent = memo(function HandleRenderComponent({ > ({ }), })); -jest.mock("@/stores/flowStore", () => ({ - __esModule: true, - default: { - getState: () => ({ - getNode: jest.fn(), - setFilterEdge: jest.fn(), - setFilterType: jest.fn(), - nodes: [], - }), - }, -})); +jest.mock("@/stores/flowStore", () => { + const state = { + getNode: jest.fn(), + setNode: jest.fn(), + setFilterEdge: jest.fn(), + setFilterType: jest.fn(), + nodes: [], + edges: [], + }; + const hook = (selector?: (s: any) => any) => + selector ? selector(state) : state; + hook.getState = () => state; + return { __esModule: true, default: hook }; +}); jest.mock("@/stores/typesStore", () => ({ useTypesStore: { diff --git a/src/frontend/src/components/core/parameterRenderComponent/components/modelInputComponent/hooks/__tests__/useModelConnectionLogic.test.ts b/src/frontend/src/components/core/parameterRenderComponent/components/modelInputComponent/hooks/__tests__/useModelConnectionLogic.test.ts new file mode 100644 index 0000000000..6205ae93f3 --- /dev/null +++ b/src/frontend/src/components/core/parameterRenderComponent/components/modelInputComponent/hooks/__tests__/useModelConnectionLogic.test.ts @@ -0,0 +1,271 @@ +/** + * Unit tests for the connection mode logic in useModelConnectionLogic. + * + * The hook activates "Connect other models" mode purely on the frontend + * by setting filter state in the flow store. No backend calls are made. + * These tests validate the pure logic decisions without React rendering. + */ + +describe("useModelConnectionLogic", () => { + describe("effective input types defaulting", () => { + const getEffectiveInputTypes = (inputTypes: string[]): string[] => { + return inputTypes.length > 0 ? inputTypes : ["LanguageModel"]; + }; + + it("should default to LanguageModel when input_types is empty", () => { + // Arrange & Act + const result = getEffectiveInputTypes([]); + + // Assert + expect(result).toEqual(["LanguageModel"]); + }); + + it("should use provided input_types when available", () => { + // Arrange & Act + const result = getEffectiveInputTypes(["Embeddings"]); + + // Assert + expect(result).toEqual(["Embeddings"]); + }); + + it("should preserve multiple input types", () => { + // Arrange & Act + const result = getEffectiveInputTypes(["LanguageModel", "ChatModel"]); + + // Assert + expect(result).toEqual(["LanguageModel", "ChatModel"]); + }); + }); + + describe("filter object construction", () => { + const NODE_ID = "agent-node-123"; + + const buildFilterObj = (nodeId: string, targetHandle: string) => ({ + source: undefined, + sourceHandle: undefined, + target: nodeId, + targetHandle, + type: "LanguageModel", + color: "datatype-fuchsia", + }); + + it("should set target to the node ID", () => { + // Arrange & Act + const filterObj = buildFilterObj(NODE_ID, "handle-123"); + + // Assert + expect(filterObj.target).toBe(NODE_ID); + expect(filterObj.source).toBeUndefined(); + expect(filterObj.sourceHandle).toBeUndefined(); + }); + + it("should use LanguageModel as the filter type", () => { + // Arrange & Act + const filterObj = buildFilterObj(NODE_ID, "handle-123"); + + // Assert + expect(filterObj.type).toBe("LanguageModel"); + expect(filterObj.color).toBe("datatype-fuchsia"); + }); + }); + + describe("option value routing", () => { + const shouldActivateConnectionMode = (optionValue: string): boolean => { + return optionValue === "connect_other_models"; + }; + + it("should activate for connect_other_models", () => { + expect(shouldActivateConnectionMode("connect_other_models")).toBe(true); + }); + + it("should not activate for other option values", () => { + expect(shouldActivateConnectionMode("some_other_option")).toBe(false); + expect(shouldActivateConnectionMode("")).toBe(false); + expect(shouldActivateConnectionMode("connect")).toBe(false); + }); + }); + + describe("selectedModel in connection mode", () => { + // Pure logic: connection mode is determined by local state (isConnectionMode), + // NOT by a special string value. The template value is cleared to [] when + // entering connection mode so the backend doesn't retain stale config. + const computeSelectedModel = ( + isConnectionMode: boolean, + value: any, + flatOptions: Array<{ name: string; icon?: string; provider?: string }>, + externalDisplayName?: string, + externalIcon?: string, + ) => { + if (isConnectionMode) { + return { + name: externalDisplayName || "Connect other models", + icon: externalIcon || "CornerDownLeft", + provider: "", + }; + } + const currentName = value?.[0]?.name; + if (!currentName) return null; + return flatOptions.find((o) => o.name === currentName) || null; + }; + + it("should return connection mode display when isConnectionMode is true", () => { + // value is [] (cleared) when in connection mode + const result = computeSelectedModel(true, [], []); + + expect(result).not.toBeNull(); + expect(result!.name).toBe("Connect other models"); + expect(result!.icon).toBe("CornerDownLeft"); + }); + + it("should use external display name when available", () => { + const result = computeSelectedModel( + true, + [], + [], + "OpenAI Compatible", + "Brain", + ); + + expect(result!.name).toBe("OpenAI Compatible"); + expect(result!.icon).toBe("Brain"); + }); + + it("should return normal model when not in connection mode", () => { + const options = [{ name: "gpt-4", icon: "Bot", provider: "OpenAI" }]; + const result = computeSelectedModel(false, [{ name: "gpt-4" }], options); + + expect(result!.name).toBe("gpt-4"); + }); + + it("should return null when no model selected and not in connection mode", () => { + const result = computeSelectedModel(false, null, []); + + expect(result).toBeNull(); + }); + }); + + describe("connection mode clears template value (stale config prevention)", () => { + it("should clear value to empty array when entering connection mode", () => { + const previousValue = [ + { name: "gpt-4", icon: "Bot", provider: "OpenAI" }, + ]; + const clearedValue: any[] = []; + + expect(clearedValue).toEqual([]); + expect(clearedValue).not.toEqual(previousValue); + }); + + it("should never send a string value to handleOnNewValue", () => { + // The old bug: handleOnNewValue({ value: "connect_other_models" }) + // caused "string indices must be integers, not 'str'" on the backend. + const newValue: any[] = []; + expect(typeof newValue).not.toBe("string"); + expect(Array.isArray(newValue)).toBe(true); + }); + }); + + describe("credential clearing in connection mode", () => { + // Pure logic: useModelConnectionLogic.setNode clears password fields + const clearCredentialFields = ( + template: Record, + ): Record => { + const updated = { ...template }; + if (updated.model) { + updated.model = { + ...updated.model, + value: [], + _connection_mode: true, + }; + } + for (const [key, field] of Object.entries(updated)) { + if (field?.password || field?._input_type === "SecretStrInput") { + updated[key] = { ...field, value: "", load_from_db: false }; + } + } + return updated; + }; + + it("should clear model value and set _connection_mode flag", () => { + const template = { + model: { + type: "model", + value: [{ name: "gpt-4", provider: "OpenAI" }], + }, + api_key: { + type: "str", + password: true, + _input_type: "SecretStrInput", + load_from_db: true, + value: "OPENAI_API_KEY", + }, + }; + + const result = clearCredentialFields(template); + + expect(result.model.value).toEqual([]); + expect(result.model._connection_mode).toBe(true); + }); + + it("should clear password fields value and disable load_from_db", () => { + const template = { + model: { type: "model", value: [{ name: "gpt-4" }] }, + api_key: { + type: "str", + password: true, + _input_type: "SecretStrInput", + load_from_db: true, + value: "OPENAI_API_KEY", + }, + temperature: { + type: "float", + value: 0.7, + password: false, + }, + }; + + const result = clearCredentialFields(template); + + expect(result.api_key.value).toBe(""); + expect(result.api_key.load_from_db).toBe(false); + // Non-secret fields should not be affected + expect(result.temperature.value).toBe(0.7); + }); + }); + + describe("exiting connection mode clears _connection_mode flag", () => { + const exitConnectionMode = ( + template: Record, + ): Record => { + const updated = { ...template }; + if (updated.model?._connection_mode) { + updated.model = { ...updated.model, _connection_mode: false }; + } + return updated; + }; + + it("should set _connection_mode to false when selecting a model", () => { + const template = { + model: { type: "model", value: [], _connection_mode: true }, + }; + + const result = exitConnectionMode(template); + + expect(result.model._connection_mode).toBe(false); + }); + + it("should not change template when _connection_mode is already false", () => { + const template = { + model: { + type: "model", + value: [{ name: "gpt-4" }], + _connection_mode: false, + }, + }; + + const result = exitConnectionMode(template); + + expect(result.model._connection_mode).toBe(false); + expect(result.model.value).toEqual([{ name: "gpt-4" }]); + }); + }); +}); diff --git a/src/frontend/src/components/core/parameterRenderComponent/components/modelInputComponent/hooks/useModelConnectionLogic.ts b/src/frontend/src/components/core/parameterRenderComponent/components/modelInputComponent/hooks/useModelConnectionLogic.ts index a197d2f9ba..9a4defda37 100644 --- a/src/frontend/src/components/core/parameterRenderComponent/components/modelInputComponent/hooks/useModelConnectionLogic.ts +++ b/src/frontend/src/components/core/parameterRenderComponent/components/modelInputComponent/hooks/useModelConnectionLogic.ts @@ -1,5 +1,4 @@ import { useCallback } from "react"; -import { mutateTemplate } from "@/CustomNodes/helpers/mutate-template"; import useFlowStore from "@/stores/flowStore"; import { useTypesStore } from "@/stores/typesStore"; import { scapedJSONStringfy } from "@/utils/reactflowUtils"; @@ -7,120 +6,111 @@ import { groupByFamily } from "@/utils/utils"; interface UseModelConnectionLogicProps { nodeId: string; - nodeClass: any; - handleNodeClass: (nodeClass: any) => void; - postTemplateValue: any; - setErrorData: (error: any) => void; - handleOnNewValue: (newValue: any) => void; closePopover: () => void; clearSelection: () => void; } export function useModelConnectionLogic({ nodeId, - nodeClass, - handleNodeClass, - postTemplateValue, - setErrorData, - handleOnNewValue, closePopover, clearSelection, }: UseModelConnectionLogicProps) { const handleExternalOptions = useCallback( - async (optionValue: string) => { + (optionValue: string) => { closePopover(); clearSelection(); - // Pass the optionValue ("connect_other_models") as both the field value and to mutateTemplate - // This way the backend knows we're in connection mode - handleOnNewValue({ value: optionValue }); + if (optionValue !== "connect_other_models") { + return; + } - await mutateTemplate( - optionValue, - nodeId!, - nodeClass!, - handleNodeClass!, - postTemplateValue, - setErrorData, - "model", - () => { - // Enable connection mode for connect_other_models AFTER mutation completes - try { - if (optionValue === "connect_other_models") { - const store = useFlowStore.getState(); - const node = store.getNode(nodeId!); - const templateField = node?.data?.node?.template?.["model"]; - if (!templateField) { - return; - } + try { + const store = useFlowStore.getState(); + const node = store.getNode(nodeId!); + const templateField = node?.data?.node?.template?.["model"]; + if (!templateField) { + return; + } - const inputTypes: string[] = - (Array.isArray(templateField.input_types) - ? templateField.input_types - : []) || []; - const effectiveInputTypes = - inputTypes.length > 0 ? inputTypes : ["LanguageModel"]; + const inputTypes: string[] = + (Array.isArray(templateField.input_types) + ? templateField.input_types + : []) || []; + const effectiveInputTypes = + inputTypes.length > 0 ? inputTypes : ["LanguageModel"]; - const tooltipTitle: string = - (inputTypes && inputTypes.length > 0 - ? inputTypes.join("\n") - : templateField.type) || ""; + const tooltipTitle: string = + (inputTypes && inputTypes.length > 0 + ? inputTypes.join("\n") + : templateField.type) || ""; - const myId = scapedJSONStringfy({ - inputTypes: effectiveInputTypes, - type: templateField.type, - id: nodeId, - fieldName: "model", - proxy: templateField.proxy, - }); + const typesData = useTypesStore.getState().data; + const grouped = groupByFamily( + typesData, + (effectiveInputTypes && effectiveInputTypes.length > 0 + ? effectiveInputTypes.join("\n") + : tooltipTitle) || "", + true, + store.nodes, + ); - const typesData = useTypesStore.getState().data; - const grouped = groupByFamily( - typesData, - (effectiveInputTypes && effectiveInputTypes.length > 0 - ? effectiveInputTypes.join("\n") - : tooltipTitle) || "", - true, - store.nodes, - ); + // Build a pseudo source so compatible target handles (left side) glow + const pseudoSourceHandle = scapedJSONStringfy({ + fieldName: "model", + id: nodeId, + inputTypes: effectiveInputTypes, + type: "str", + }); - // Build a pseudo source so compatible target handles (left side) glow - const pseudoSourceHandle = scapedJSONStringfy({ - fieldName: "model", - id: nodeId, - inputTypes: effectiveInputTypes, - type: "str", - }); + const filterObj = { + source: undefined, + sourceHandle: undefined, + target: nodeId, + targetHandle: pseudoSourceHandle, + type: "LanguageModel", + color: "datatype-fuchsia", + } as any; - const filterObj = { - source: undefined, - sourceHandle: undefined, - target: nodeId, - targetHandle: pseudoSourceHandle, - type: "LanguageModel", - color: "datatype-fuchsia", - } as any; - - // Show compatible handles glow - store.setFilterEdge(grouped); - store.setFilterType(filterObj); + // Mark this node as being in connection mode, clear the model value + // and provider-specific credential fields so the backend cannot + // execute with stale config when no external model is connected. + store.setNode( + nodeId, + (prevNode) => { + const template = { ...prevNode.data.node.template }; + if (template.model) { + template.model = { + ...template.model, + value: [], + _connection_mode: true, + }; } - } catch (error) { - console.warn("Error setting up connection mode:", error); - } - }, - ); + for (const [key, field] of Object.entries(template)) { + const f = field as any; + if (f?.password || f?._input_type === "SecretStrInput") { + template[key] = { ...f, value: "", load_from_db: false }; + } + } + return { + ...prevNode, + data: { + ...prevNode.data, + _connectionMode: true, + node: { ...prevNode.data.node, template }, + }, + }; + }, + true, + ); + + // Show compatible handles glow + store.setFilterEdge(grouped); + store.setFilterType(filterObj); + } catch (error) { + console.warn("Error setting up connection mode:", error); + } }, - [ - nodeId, - nodeClass, - handleNodeClass, - postTemplateValue, - setErrorData, - handleOnNewValue, - closePopover, - clearSelection, - ], + [nodeId, closePopover, clearSelection], ); return { handleExternalOptions }; diff --git a/src/frontend/src/components/core/parameterRenderComponent/components/modelInputComponent/index.tsx b/src/frontend/src/components/core/parameterRenderComponent/components/modelInputComponent/index.tsx index 4127195a0d..83b4a99f3e 100644 --- a/src/frontend/src/components/core/parameterRenderComponent/components/modelInputComponent/index.tsx +++ b/src/frontend/src/components/core/parameterRenderComponent/components/modelInputComponent/index.tsx @@ -7,7 +7,10 @@ import { usePostTemplateValue } from "@/controllers/API/queries/nodes/use-post-t import { useRefreshModelInputs } from "@/hooks/use-refresh-model-inputs"; import ModelProviderModal from "@/modals/modelProviderModal"; import useAlertStore from "@/stores/alertStore"; +import useFlowStore from "@/stores/flowStore"; import type { APIClassType } from "@/types/api"; +import type { NodeDataType } from "@/types/flow"; +import { useModelConnectionLogic } from "./hooks/useModelConnectionLogic"; import ForwardedIconComponent from "../../../../common/genericIconComponent"; import { Button } from "../../../../ui/button"; import { Command } from "../../../../ui/command"; @@ -48,6 +51,31 @@ export default function ModelInputComponent({ useState(false); const [isRefreshingAfterClose, setIsRefreshingAfterClose] = useState(false); const [refreshOptions, setRefreshOptions] = useState(false); + + // Connection mode: local state for reactivity, persisted in node data for reload + const [isConnectionMode, setIsConnectionMode] = useState(() => { + if (!nodeId) return false; + const node = useFlowStore.getState().nodes.find((n) => n.id === nodeId); + return (node?.data as any)?._connectionMode === true; + }); + + const setConnectionMode = useCallback( + (enabled: boolean) => { + setIsConnectionMode(enabled); + if (!nodeId) return; + const store = useFlowStore.getState(); + store.setNode( + nodeId, + (node) => ({ + ...node, + data: { ...node.data, _connectionMode: enabled }, + }), + false, + ); + }, + [nodeId], + ); + const { refreshAllModelInputs } = useRefreshModelInputs(); // Ref to track if we've already processed the empty options state @@ -60,6 +88,19 @@ export default function ModelInputComponent({ node: (nodeClass as APIClassType) || null, }); + const { handleExternalOptions } = useModelConnectionLogic({ + nodeId: nodeId || "", + closePopover: () => setOpen(false), + clearSelection: () => { + // Only set the _connectionMode flag on the node data. + // Don't call handleOnNewValue — it triggers a backend round-trip + // that tries to resolve __default_language_model__ and fails. + // The credential fields are cleared by useModelConnectionLogic + // directly in the node template via setNode. + setConnectionMode(true); + }, + }); + const modelType = nodeClass?.template?.model?.model_type === "language" ? "llm" @@ -111,9 +152,15 @@ export default function ModelInputComponent({ // Derive the currently selected model from the value prop const selectedModel = useMemo(() => { - // If we're in connection mode, we don't have a normal selected model - if (value === "connect_other_models") { - return null; + // If we're in connection mode, show the connection option as selected + if (isConnectionMode) { + return { + name: + externalOptions?.fields?.data?.node?.display_name || + "Connect other models", + icon: externalOptions?.fields?.data?.node?.icon || "CornerDownLeft", + provider: "", + } as SelectedModel; } const currentName = value?.[0]?.name; @@ -133,11 +180,15 @@ export default function ModelInputComponent({ // Or keep displaying the stale one? Original logic selected first available. (flatOptions.length > 0 ? flatOptions[0] : null) ); - }, [value, flatOptions]); + }, [value, flatOptions, isConnectionMode, externalOptions]); useEffect(() => { - // Only proceed if we have options and haven't selected a value - if (flatOptions.length > 0 && (!value || value.length === 0)) { + // Only proceed if we have options, haven't selected a value, and NOT in connection mode + if ( + flatOptions.length > 0 && + (!value || value.length === 0) && + !isConnectionMode + ) { // Check ref to avoid infinite loops if (!hasProcessedEmptyRef.current) { const firstOption = flatOptions[0]; @@ -162,6 +213,37 @@ export default function ModelInputComponent({ */ const handleModelSelect = useCallback( (modelName: string) => { + setConnectionMode(false); + // Clear the _connection_mode flag from the model field template + // so the backend resumes normal update_build_config behavior. + if (nodeId) { + const store = useFlowStore.getState(); + const node = store.getNode(nodeId); + const nodeData = node?.data as NodeDataType | undefined; + if (nodeData?.node?.template?.model?._connection_mode) { + store.setNode( + nodeId, + (prev) => ({ + ...prev, + data: { + ...prev.data, + _connectionMode: false, + node: { + ...(prev.data as NodeDataType).node, + template: { + ...(prev.data as NodeDataType).node.template, + model: { + ...(prev.data as NodeDataType).node.template.model, + _connection_mode: false, + }, + }, + }, + } as NodeDataType, + }), + false, + ); + } + } const selectedOption = flatOptions.find( (option) => option.name === modelName, ); @@ -295,6 +377,17 @@ export default function ModelInputComponent({ "refresh-model-list", )} {renderManageProvidersButton()} + {externalOptions?.fields?.data?.node && ( +
+ {renderFooterButton( + externalOptions.fields.data.node.display_name || + "Connect other models", + externalOptions.fields.data.node.icon || "CornerDownLeft", + () => handleExternalOptions("connect_other_models"), + "connect-other-models", + )} +
+ )} ); diff --git a/src/frontend/src/types/flow/index.ts b/src/frontend/src/types/flow/index.ts index 3be7df0c1a..1cd40c6a81 100644 --- a/src/frontend/src/types/flow/index.ts +++ b/src/frontend/src/types/flow/index.ts @@ -69,6 +69,8 @@ export type NodeDataType = { selected_output_type?: string; buildStatus?: BuildStatus; selected_output?: string; + /** Transient flag: true while "Connect other models" mode is active */ + _connectionMode?: boolean; }; export type EdgeType = Edge; diff --git a/src/frontend/tests/extended/features/outdated-message.spec.ts b/src/frontend/tests/extended/features/outdated-message.spec.ts index 1d975ddb93..0c3451ee9d 100644 --- a/src/frontend/tests/extended/features/outdated-message.spec.ts +++ b/src/frontend/tests/extended/features/outdated-message.spec.ts @@ -33,9 +33,13 @@ test("user must be able outdated message on error", async ({ page }) => { await page.getByTestId("list-card").first().click(); - await page - .getByTestId("popover-anchor-input-api_key") - .fill("this is a test to crash"); + // The api_key field may show a global variable badge instead of an input + // when OPENAI_API_KEY is in the environment. Either way, we just need to + // trigger a run to get the outdated components error message. + const apiKeyInput = page.getByTestId("popover-anchor-input-api_key"); + if (await apiKeyInput.isVisible({ timeout: 3000 }).catch(() => false)) { + await apiKeyInput.fill("this is a test to crash"); + } await page.getByTestId("button_run_chat output").click(); diff --git a/src/lfx/src/lfx/base/models/unified_models/build_config.py b/src/lfx/src/lfx/base/models/unified_models/build_config.py index 953e4abe26..575c5bc284 100644 --- a/src/lfx/src/lfx/base/models/unified_models/build_config.py +++ b/src/lfx/src/lfx/base/models/unified_models/build_config.py @@ -313,6 +313,11 @@ def handle_model_input_update( if field_name in provider_mapped_fields: return build_config + # If the model field is in connection mode (user chose "Connect other models"), + # skip auto-selection and provider re-population so credentials stay cleared. + if build_config.get(model_field_name, {}).get("_connection_mode"): + return build_config + # When the user changes the model selection, we need to reset/hide fields that may no longer apply if field_name == model_field_name: options = build_config[model_field_name].get("options", []) diff --git a/src/lfx/src/lfx/graph/vertex/param_handler.py b/src/lfx/src/lfx/graph/vertex/param_handler.py index 05663daff6..852652b128 100644 --- a/src/lfx/src/lfx/graph/vertex/param_handler.py +++ b/src/lfx/src/lfx/graph/vertex/param_handler.py @@ -206,7 +206,22 @@ class ParameterHandler: params = self._handle_other_direct_types(field_name, field, val, params) if field.get("load_from_db"): - load_from_db_fields.append(field_name) + # Skip load_from_db if the field itself has an incoming edge + has_incoming_edge = self.vertex.get_incoming_edge_by_target_param(field_name) is not None + # Skip credential fields when the model field has an incoming edge, + # because the connected model component provides its own credentials + is_secret = field.get("_input_type") == "SecretStrInput" or field.get("password") + model_has_edge = ( + is_secret + and "model" in self.template_dict + and self.vertex.get_incoming_edge_by_target_param("model") is not None + ) + # Skip credential fields when the node is in "Connect other models" mode + # (user chose to wire an external model instead of the built-in provider) + model_field = self.template_dict.get("model", {}) + in_connection_mode = is_secret and model_field.get("_connection_mode", False) + if not has_incoming_edge and not model_has_edge and not in_connection_mode: + load_from_db_fields.append(field_name) return params, load_from_db_fields diff --git a/src/lfx/tests/unit/graph/vertex/test_param_handler_model_edge.py b/src/lfx/tests/unit/graph/vertex/test_param_handler_model_edge.py new file mode 100644 index 0000000000..1ec448fef1 --- /dev/null +++ b/src/lfx/tests/unit/graph/vertex/test_param_handler_model_edge.py @@ -0,0 +1,254 @@ +"""Tests for skipping load_from_db credential fields when model has incoming edge. + +When a model component is connected via wire to a node's model input, +the connected component provides its own credentials. The node's own +secret fields (like api_key) should NOT be resolved from the database, +preventing errors like 'OPENAI_API_KEY variable not found'. +""" + +from unittest.mock import MagicMock + +from lfx.graph.vertex.param_handler import ParameterHandler + + +class TestSkipLoadFromDbWhenModelHasEdge: + """Tests for _process_direct_type_field skipping load_from_db for secrets when model has edge.""" + + def _create_handler_with_template(self, template: dict, *, model_has_edge: bool = False) -> ParameterHandler: + """Create a ParameterHandler with the given template and optional model edge.""" + mock_vertex = MagicMock() + mock_vertex.data = {"node": {"template": template}} + + def mock_get_incoming_edge(field_name: str): + if field_name == "model" and model_has_edge: + return "source-node-id" + return None + + mock_vertex.get_incoming_edge_by_target_param = MagicMock(side_effect=mock_get_incoming_edge) + return ParameterHandler(mock_vertex, storage_service=None) + + def test_should_add_secret_to_load_from_db_when_no_model_edge(self): + """Secret fields should be resolved normally when model has no wire.""" + # Arrange + template = { + "model": {"type": "model", "value": None, "show": True}, + "api_key": { + "type": "str", + "_input_type": "SecretStrInput", + "password": True, + "load_from_db": True, + "value": "OPENAI_API_KEY", + "show": True, + }, + } + handler = self._create_handler_with_template(template, model_has_edge=False) + + # Act + load_from_db_fields: list[str] = [] + params: dict = {} + handler._process_direct_type_field("api_key", template["api_key"], params, load_from_db_fields) + + # Assert + assert "api_key" in load_from_db_fields + + def test_should_skip_secret_load_from_db_when_model_has_edge(self): + """Secret fields should NOT be resolved when model has a connected wire.""" + # Arrange + template = { + "model": {"type": "model", "value": None, "show": True}, + "api_key": { + "type": "str", + "_input_type": "SecretStrInput", + "password": True, + "load_from_db": True, + "value": "OPENAI_API_KEY", + "show": True, + }, + } + handler = self._create_handler_with_template(template, model_has_edge=True) + + # Act + load_from_db_fields: list[str] = [] + params: dict = {} + handler._process_direct_type_field("api_key", template["api_key"], params, load_from_db_fields) + + # Assert + assert "api_key" not in load_from_db_fields + + def test_should_skip_password_field_when_model_has_edge(self): + """Any password field should be skipped when model has edge, not just SecretStrInput.""" + # Arrange + template = { + "model": {"type": "model", "value": None, "show": True}, + "custom_token": { + "type": "str", + "password": True, + "load_from_db": True, + "value": "MY_CUSTOM_TOKEN", + "show": True, + }, + } + handler = self._create_handler_with_template(template, model_has_edge=True) + + # Act + load_from_db_fields: list[str] = [] + params: dict = {} + handler._process_direct_type_field("custom_token", template["custom_token"], params, load_from_db_fields) + + # Assert + assert "custom_token" not in load_from_db_fields + + def test_should_not_skip_non_secret_field_when_model_has_edge(self): + """Non-secret fields with load_from_db should still be resolved even with model edge.""" + # Arrange + template = { + "model": {"type": "model", "value": None, "show": True}, + "system_prompt": { + "type": "str", + "load_from_db": True, + "value": "MY_PROMPT_VAR", + "show": True, + }, + } + handler = self._create_handler_with_template(template, model_has_edge=True) + + # Act + load_from_db_fields: list[str] = [] + params: dict = {} + handler._process_direct_type_field("system_prompt", template["system_prompt"], params, load_from_db_fields) + + # Assert + assert "system_prompt" in load_from_db_fields + + def test_should_skip_field_with_direct_incoming_edge(self): + """Fields with their own incoming edge should be skipped regardless of model edge.""" + # Arrange + template = { + "some_field": { + "type": "str", + "load_from_db": True, + "value": "SOME_VAR", + "show": True, + }, + } + mock_vertex = MagicMock() + mock_vertex.data = {"node": {"template": template}} + mock_vertex.get_incoming_edge_by_target_param = MagicMock(return_value="source-node") + handler = ParameterHandler(mock_vertex, storage_service=None) + + # Act + load_from_db_fields: list[str] = [] + params: dict = {} + handler._process_direct_type_field("some_field", template["some_field"], params, load_from_db_fields) + + # Assert + assert "some_field" not in load_from_db_fields + + def test_should_not_skip_when_no_model_field_in_template(self): + """Components without a model field should resolve secrets normally.""" + # Arrange + template = { + "api_key": { + "type": "str", + "_input_type": "SecretStrInput", + "password": True, + "load_from_db": True, + "value": "OPENAI_API_KEY", + "show": True, + }, + } + mock_vertex = MagicMock() + mock_vertex.data = {"node": {"template": template}} + mock_vertex.get_incoming_edge_by_target_param = MagicMock(return_value=None) + handler = ParameterHandler(mock_vertex, storage_service=None) + + # Act + load_from_db_fields: list[str] = [] + params: dict = {} + handler._process_direct_type_field("api_key", template["api_key"], params, load_from_db_fields) + + # Assert + assert "api_key" in load_from_db_fields + + def test_should_not_affect_fields_without_load_from_db(self): + """Fields without load_from_db should never be added regardless of edges.""" + # Arrange + template = { + "model": {"type": "model", "value": None, "show": True}, + "description": { + "type": "str", + "value": "Some description", + "show": True, + }, + } + handler = self._create_handler_with_template(template, model_has_edge=True) + + # Act + load_from_db_fields: list[str] = [] + params: dict = {} + handler._process_direct_type_field("description", template["description"], params, load_from_db_fields) + + # Assert + assert "description" not in load_from_db_fields + + def test_should_skip_secret_load_from_db_when_connection_mode_active(self): + """Secret fields must NOT be resolved via load_from_db in connection mode. + + When the model field has _connection_mode: true (user chose + 'Connect other models'), credential fields should be skipped. + """ + # Arrange + template = { + "model": { + "type": "model", + "value": [], + "show": True, + "_connection_mode": True, + }, + "api_key": { + "type": "str", + "_input_type": "SecretStrInput", + "password": True, + "load_from_db": True, + "value": "", + "show": True, + }, + } + handler = self._create_handler_with_template(template, model_has_edge=False) + + # Act + load_from_db_fields: list[str] = [] + params: dict = {} + handler._process_direct_type_field("api_key", template["api_key"], params, load_from_db_fields) + + # Assert — api_key must NOT be in load_from_db_fields + assert "api_key" not in load_from_db_fields + + def test_should_allow_secret_load_from_db_when_connection_mode_false(self): + """Secret fields should be resolved normally when _connection_mode is false.""" + # Arrange + template = { + "model": { + "type": "model", + "value": [{"name": "gpt-4", "provider": "OpenAI"}], + "show": True, + "_connection_mode": False, + }, + "api_key": { + "type": "str", + "_input_type": "SecretStrInput", + "password": True, + "load_from_db": True, + "value": "OPENAI_API_KEY", + "show": True, + }, + } + handler = self._create_handler_with_template(template, model_has_edge=False) + + # Act + load_from_db_fields: list[str] = [] + params: dict = {} + handler._process_direct_type_field("api_key", template["api_key"], params, load_from_db_fields) + + # Assert — api_key SHOULD be in load_from_db_fields + assert "api_key" in load_from_db_fields diff --git a/src/lfx/tests/unit/graph/vertex/test_vertex_base.py b/src/lfx/tests/unit/graph/vertex/test_vertex_base.py index 418c8ff216..11b47857bb 100644 --- a/src/lfx/tests/unit/graph/vertex/test_vertex_base.py +++ b/src/lfx/tests/unit/graph/vertex/test_vertex_base.py @@ -49,6 +49,8 @@ def mock_vertex() -> Mock: } vertex.id = "test-vertex-id" vertex.display_name = "Test Vertex" + # Default: no incoming edges for any field + vertex.get_incoming_edge_by_target_param = Mock(return_value=None) return vertex