Marcel Hofer
b87d17e34d
Add native OpenAI Decisions and OpenRouter Decisions API routing to the image and text Decision Router nodes while keeping local in-process llama.cpp inference as the default. Discover OpenRouter decision models by input modality and decision output capability, then enforce server profile image policy and optional model allowlists both in the UI model catalog and before inference. Keep credentials and endpoints in server-side profiles; apply profile input, image, and timeout limits; never fall back to chat-generated JSON. Preserve saved workflow widget order, release any retained local model when switching to API mode, and normalize typed route choices, probabilities, and confidence into decision_json. Add API setup documentation and adapter/router/security regression coverage. Validation: 240 tests passed across six test modules run individually in bundled ComfyUI Python 3.12; Python and JavaScript syntax checks, API catalog validation, version metadata checks, and git diff --check passed. Provider inference is covered with mocked responses; no authenticated live inference was run.
380 lines
18 KiB
Python
380 lines
18 KiB
Python
from __future__ import annotations
|
|
|
|
import base64
|
|
import importlib.util
|
|
import json
|
|
import math
|
|
import os
|
|
import sys
|
|
import tempfile
|
|
import types
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest import mock
|
|
|
|
HERE = Path(__file__).resolve().parent
|
|
PKG = HERE.parent
|
|
MODULE_PATH = PKG / 'nodes' / 'openai_compatible_api.py'
|
|
if str(PKG) not in sys.path:
|
|
sys.path.insert(0, str(PKG))
|
|
|
|
output_cleaner = types.ModuleType('AILab_OutputCleaner')
|
|
class OutputCleanConfig:
|
|
def __init__(self, **kwargs):
|
|
self.__dict__.update(kwargs)
|
|
output_cleaner.OutputCleanConfig = OutputCleanConfig
|
|
output_cleaner.clean_model_output = lambda text, config=None: text.replace('<think>x</think>', '').strip()
|
|
sys.modules['AILab_OutputCleaner'] = output_cleaner
|
|
|
|
stream_display = types.ModuleType('AILab_StreamDisplay')
|
|
stream_display.compose_streamed_model_output = lambda reasoning, content: (
|
|
f'<think>\n{reasoning}\n</think>\n\n{content}' if reasoning and content
|
|
else f'<think>\n{reasoning}\n</think>' if reasoning else content
|
|
)
|
|
stream_display.extract_stream_token = lambda chunk: {
|
|
'reasoning': ((chunk.get('choices') or [{}])[0].get('delta') or {}).get('reasoning_content', ''),
|
|
'content': ((chunk.get('choices') or [{}])[0].get('delta') or {}).get('content', ''),
|
|
}
|
|
stream_display.get_thinking_stream_display = lambda: None
|
|
sys.modules['AILab_StreamDisplay'] = stream_display
|
|
|
|
SPEC = importlib.util.spec_from_file_location('thinkingllm_openai_compatible_api_test', MODULE_PATH)
|
|
assert SPEC and SPEC.loader
|
|
api_node = importlib.util.module_from_spec(SPEC)
|
|
SPEC.loader.exec_module(api_node)
|
|
|
|
class SecurityTests(unittest.TestCase):
|
|
def profile(self, **overrides):
|
|
raw = {
|
|
'provider': 'Test',
|
|
'base_url': 'https://example.test/v1',
|
|
'auth': 'bearer_env',
|
|
'api_key_env': 'TEST_API_KEY',
|
|
}
|
|
raw.update(overrides)
|
|
return api_node._normalize_profile('Production', raw)
|
|
|
|
def test_no_workflow_secret_or_endpoint_inputs(self):
|
|
inputs = api_node.ThinkingLLMOpenAICompatibleAPI.INPUT_TYPES()
|
|
names = set(inputs.get('required', {})) | set(inputs.get('optional', {}))
|
|
self.assertIn('api_profile', names)
|
|
self.assertIn('image', names)
|
|
self.assertIn('image_url', names)
|
|
self.assertNotIn('video', names)
|
|
self.assertNotIn('video_url', names)
|
|
for forbidden in ('api_key', 'api_key_env', 'base_url', 'headers'):
|
|
self.assertNotIn(forbidden, names)
|
|
|
|
def test_model_name_is_a_curated_combo(self):
|
|
inputs = api_node.ThinkingLLMOpenAICompatibleAPI.INPUT_TYPES()
|
|
model_def = inputs['required']['model_name']
|
|
options = model_def[0]
|
|
self.assertIsInstance(options, list)
|
|
self.assertGreaterEqual(len(options), 1)
|
|
# The dropdown should carry real curated models, not a placeholder.
|
|
self.assertNotIn('provider/model-id', options)
|
|
# The dropdown should contain at least one model from the current provider catalog.
|
|
catalog_models = api_node._catalog_models_for_profile('OrcaRouter')
|
|
self.assertTrue(catalog_models)
|
|
self.assertIn(catalog_models[0], options)
|
|
|
|
def test_custom_model_field_overrides_dropdown(self):
|
|
inputs = api_node.ThinkingLLMOpenAICompatibleAPI.INPUT_TYPES()
|
|
self.assertIn('custom_model_name', inputs.get('optional', {}))
|
|
p = self.profile()
|
|
self.assertEqual(api_node._validate_model(p, 'custom/provider-model'), 'custom/provider-model')
|
|
self.assertEqual(api_node._validate_model(p, 'any/custom-id'), 'any/custom-id')
|
|
|
|
def test_authenticated_profiles_require_https(self):
|
|
with self.assertRaisesRegex(ValueError, 'must use HTTPS'):
|
|
api_node._normalize_profile('Unsafe', {
|
|
'base_url': 'http://example.test/v1',
|
|
'auth': 'bearer_env',
|
|
'api_key_env': 'TEST_KEY',
|
|
})
|
|
|
|
def test_native_decision_profile_requires_a_supported_protocol_and_endpoint_pair(self):
|
|
profile = self.profile(
|
|
decision_protocol='openrouter_decisions',
|
|
decision_endpoint='https://openrouter.ai/api/alpha/decisions',
|
|
)
|
|
self.assertEqual(profile['decision_protocol'], 'openrouter_decisions')
|
|
self.assertEqual(profile['decision_endpoint'], 'https://openrouter.ai/api/alpha/decisions')
|
|
with self.assertRaisesRegex(ValueError, 'both decision_protocol and decision_endpoint'):
|
|
self.profile(decision_protocol='openai_decisions')
|
|
with self.assertRaisesRegex(ValueError, 'unsupported decision_protocol'):
|
|
self.profile(decision_protocol='chat_json', decision_endpoint='https://example.test/decisions')
|
|
|
|
def test_plain_http_non_loopback_requires_explicit_opt_in(self):
|
|
with self.assertRaisesRegex(ValueError, 'Plain HTTP'):
|
|
api_node._normalize_profile('LAN', {
|
|
'base_url': 'http://192.168.1.10:8000/v1',
|
|
'auth': 'none',
|
|
})
|
|
p = api_node._normalize_profile('LAN', {
|
|
'base_url': 'http://192.168.1.10:8000/v1',
|
|
'auth': 'none',
|
|
'allow_insecure_http': True,
|
|
})
|
|
self.assertTrue(p['allow_insecure_http'])
|
|
|
|
def test_loopback_http_allowed(self):
|
|
p = api_node._normalize_profile('Local', {
|
|
'base_url': 'http://127.0.0.1:8000/v1',
|
|
'auth': 'none',
|
|
})
|
|
self.assertEqual(p['base_url'], 'http://127.0.0.1:8000/v1')
|
|
|
|
def test_unknown_profile_fields_fail_closed(self):
|
|
with self.assertRaisesRegex(ValueError, 'unknown fields'):
|
|
self.profile(max_token_limit=10)
|
|
|
|
def test_deprecated_boolean_extra_body_switch_rejected(self):
|
|
with self.assertRaisesRegex(ValueError, 'deprecated allow_extra_body'):
|
|
self.profile(allow_extra_body=True)
|
|
|
|
def test_exact_extra_body_allowlist(self):
|
|
p = self.profile(allowed_extra_body_fields=['reasoning'])
|
|
parsed = api_node._parse_extra_json('{"reasoning":{"effort":"high"}}', allowed_fields=p['allowed_extra_body_fields'])
|
|
self.assertIn('reasoning', parsed)
|
|
with self.assertRaisesRegex(ValueError, 'not allowed'):
|
|
api_node._parse_extra_json('{"n":4}', allowed_fields=p['allowed_extra_body_fields'])
|
|
with self.assertRaisesRegex(ValueError, 'protected request fields'):
|
|
api_node._parse_extra_json('{"model":"evil"}', allowed_fields=['model'])
|
|
|
|
def test_extra_json_size_limit(self):
|
|
p = self.profile(allowed_extra_body_fields=['x'])
|
|
payload = '{"x":"' + ('a' * api_node.MAX_EXTRA_BODY_JSON_CHARS) + '"}'
|
|
with self.assertRaisesRegex(ValueError, 'safety limit'):
|
|
api_node._parse_extra_json(payload, allowed_fields=p['allowed_extra_body_fields'])
|
|
|
|
def test_nonfinite_extra_json_rejected(self):
|
|
with self.assertRaisesRegex(ValueError, 'Non-finite'):
|
|
api_node._parse_extra_json('{"reasoning":NaN}', allowed_fields=['reasoning'])
|
|
|
|
def test_model_and_limits(self):
|
|
p = self.profile(allowed_models=['approved/model'], max_tokens_limit=2048, max_input_chars=20, max_timeout_seconds=30)
|
|
self.assertEqual(api_node._validate_model(p, 'approved/model'), 'approved/model')
|
|
with self.assertRaisesRegex(ValueError, 'not allowed'):
|
|
api_node._validate_model(p, 'other/model')
|
|
with self.assertRaisesRegex(ValueError, 'control characters'):
|
|
api_node._validate_model_identifier('bad\x1bmodel')
|
|
with self.assertRaisesRegex(ValueError, 'exceeds server-side limit'):
|
|
api_node._validate_max_tokens(p, 2049)
|
|
with self.assertRaisesRegex(ValueError, 'max_input_chars'):
|
|
api_node._validate_input_text(p, '1234567890', '12345678901')
|
|
with self.assertRaisesRegex(ValueError, 'exceeds server-side limit'):
|
|
api_node._validate_timeout(p, 31)
|
|
|
|
def test_sampling_rejects_nan_and_wrong_types(self):
|
|
with self.assertRaisesRegex(ValueError, 'finite'):
|
|
api_node._validate_sampling(float('nan'), 0.9, 1, 0)
|
|
with self.assertRaisesRegex(ValueError, 'seed must be an integer'):
|
|
api_node._validate_sampling(0.5, 0.9, 1.2, 0)
|
|
with self.assertRaisesRegex(ValueError, 'temperature must be a numeric'):
|
|
api_node._validate_sampling('0.5', 0.9, 1, 0)
|
|
|
|
def test_thinking_budget_has_server_side_cap(self):
|
|
p = self.profile(max_thinking_budget=4096)
|
|
self.assertEqual(api_node._validate_thinking_budget(p, 4096), 4096)
|
|
with self.assertRaisesRegex(ValueError, 'exceeds server-side limit'):
|
|
api_node._validate_thinking_budget(p, 4097)
|
|
|
|
def test_allowed_models_requires_strings(self):
|
|
with self.assertRaisesRegex(ValueError, 'list of strings'):
|
|
self.profile(allowed_models=[123])
|
|
|
|
def test_profile_booleans_are_strict(self):
|
|
with self.assertRaisesRegex(ValueError, 'JSON boolean'):
|
|
self.profile(send_seed='false')
|
|
with self.assertRaisesRegex(ValueError, 'JSON boolean'):
|
|
self.profile(allow_images='true')
|
|
|
|
def test_environment_variable_name_is_validated(self):
|
|
with self.assertRaisesRegex(ValueError, 'environment-variable name'):
|
|
self.profile(api_key_env='BAD-NAME')
|
|
|
|
def test_missing_key_does_not_expose_env_name_in_client_error(self):
|
|
p = self.profile(api_key_env='THINKINGLLM_TEST_SECRET_KEY')
|
|
with mock.patch.dict(os.environ, {}, clear=False):
|
|
os.environ.pop('THINKINGLLM_TEST_SECRET_KEY', None)
|
|
with self.assertRaises(RuntimeError) as ctx:
|
|
api_node._authorization_headers(p)
|
|
self.assertNotIn('THINKINGLLM_TEST_SECRET_KEY', str(ctx.exception))
|
|
|
|
def test_redirects_are_rejected(self):
|
|
handler = api_node._NoRedirectHandler()
|
|
req = __import__('urllib.request').request.Request('https://example.test/v1/chat/completions')
|
|
self.assertIsNone(handler.redirect_request(req, None, 302, 'Found', {}, 'https://evil.test/steal'))
|
|
|
|
def test_terminal_copy_strips_controls(self):
|
|
value = api_node._terminal_display_safe('hello\x1b[31mred\x07\rworld')
|
|
self.assertNotIn('\x1b', value)
|
|
self.assertNotIn('\x07', value)
|
|
self.assertNotIn('\r', value)
|
|
self.assertIn('[31mredworld', value)
|
|
|
|
def test_openrouter_reasoning_details_are_captured(self):
|
|
chunk = {
|
|
'choices': [{
|
|
'delta': {
|
|
'reasoning_details': [
|
|
{'type': 'reasoning.text', 'text': 'step one '},
|
|
{'type': 'reasoning.summary', 'summary': 'step two'},
|
|
],
|
|
'content': 'answer',
|
|
}
|
|
}]
|
|
}
|
|
token = api_node._extract_stream_token_extended(chunk)
|
|
self.assertEqual(token['reasoning'], 'step one step two')
|
|
self.assertEqual(token['content'], 'answer')
|
|
|
|
def test_stream_error_payload_detected(self):
|
|
self.assertIn('429', api_node._provider_error_from_payload({'error': {'code': 429, 'message': 'rate limit'}}))
|
|
|
|
def test_remote_node_always_reruns(self):
|
|
self.assertTrue(math.isnan(api_node.ThinkingLLMOpenAICompatibleAPI.IS_CHANGED()))
|
|
|
|
def test_image_support_is_fail_closed_for_custom_profiles(self):
|
|
p = self.profile()
|
|
self.assertFalse(p['allow_images'])
|
|
self.assertFalse(p['allow_image_url'])
|
|
with self.assertRaisesRegex(ValueError, 'Image input is disabled'):
|
|
api_node._validate_image_url(p, 'https://example.com/image.png')
|
|
|
|
def test_image_url_requires_image_capability_and_https(self):
|
|
with self.assertRaisesRegex(ValueError, 'cannot enable allow_image_url'):
|
|
self.profile(allow_image_url=True)
|
|
p = self.profile(allow_images=True, allow_image_url=True)
|
|
signed = 'https://example.com/image.png?signature=abc123'
|
|
self.assertEqual(api_node._validate_image_url(p, signed), signed)
|
|
with self.assertRaisesRegex(ValueError, 'public HTTPS'):
|
|
api_node._validate_image_url(p, 'http://example.com/image.png')
|
|
with self.assertRaisesRegex(ValueError, 'public HTTPS'):
|
|
api_node._validate_image_url(p, 'data:image/png;base64,AAAA')
|
|
|
|
def test_comfyui_image_is_encoded_as_png_base64(self):
|
|
import numpy as np
|
|
p = self.profile(allow_images=True, max_image_pixels=4, max_image_bytes=10000)
|
|
image = np.zeros((1, 2, 2, 3), dtype=np.float32)
|
|
image[..., 0] = 1.0
|
|
data_url = api_node._image_to_data_url(p, image)
|
|
self.assertTrue(data_url.startswith('data:image/png;base64,'))
|
|
png = base64.b64decode(data_url.split(',', 1)[1])
|
|
self.assertTrue(png.startswith(b'\x89PNG\r\n\x1a\n'))
|
|
|
|
def test_image_size_and_batch_limits_are_enforced(self):
|
|
import numpy as np
|
|
p = self.profile(allow_images=True, max_image_pixels=3, max_image_bytes=10000)
|
|
with self.assertRaisesRegex(ValueError, 'max_image_pixels'):
|
|
api_node._image_to_data_url(p, np.zeros((1, 2, 2, 3), dtype=np.float32))
|
|
p2 = self.profile(allow_images=True)
|
|
with self.assertRaisesRegex(ValueError, 'exactly one IMAGE'):
|
|
api_node._image_to_data_url(p2, np.zeros((2, 1, 1, 3), dtype=np.float32))
|
|
|
|
def test_image_and_image_url_are_alternatives(self):
|
|
p = self.profile(allow_images=True, allow_image_url=True)
|
|
with self.assertRaisesRegex(ValueError, 'either the IMAGE input or image_url'):
|
|
api_node._build_user_content(p, 'describe', image=object(), image_url='https://example.com/x.png')
|
|
|
|
def test_multimodal_content_uses_text_first_then_image_url(self):
|
|
p = self.profile(allow_images=True, allow_image_url=True)
|
|
content, supplied = api_node._build_user_content(p, 'describe', image_url='https://example.com/x.png')
|
|
self.assertTrue(supplied)
|
|
self.assertEqual(content[0], {'type': 'text', 'text': 'describe'})
|
|
self.assertEqual(content[1]['type'], 'image_url')
|
|
self.assertEqual(content[1]['image_url']['url'], 'https://example.com/x.png')
|
|
|
|
def test_image_rejection_message_is_actionable(self):
|
|
p = self.profile(allow_images=True)
|
|
message = api_node._image_rejection_message(p, 'text-only/model', 400)
|
|
self.assertIn('Image input was rejected', message)
|
|
self.assertIn('vision-capable', message)
|
|
|
|
def test_builtins_enable_standard_image_transport(self):
|
|
profiles = api_node._load_api_profiles()
|
|
self.assertTrue(profiles['OpenRouter']['allow_images'])
|
|
self.assertTrue(profiles['OpenRouter']['allow_image_url'])
|
|
|
|
def test_builtins_can_be_disabled(self):
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
with mock.patch.dict(os.environ, {
|
|
api_node.DISABLE_BUILTINS_ENV: '1',
|
|
api_node.PROFILE_FILE_ENV: str(Path(tmp) / 'missing.json'),
|
|
}, clear=False):
|
|
self.assertEqual(api_node._load_api_profiles(), {})
|
|
|
|
def test_invalid_custom_override_does_not_fall_back_to_builtin(self):
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
path = Path(tmp) / 'profiles.json'
|
|
path.write_text(json.dumps({'profiles': {'OpenRouter': {
|
|
'provider': 'OpenRouter',
|
|
'base_url': 'https://openrouter.ai/api/v1',
|
|
'auth': 'bearer_env',
|
|
'api_key_env': 'OPENROUTER_API_KEY',
|
|
'max_token_limit': 1
|
|
}}}), encoding='utf-8')
|
|
with mock.patch.dict(os.environ, {api_node.PROFILE_FILE_ENV: str(path)}, clear=False):
|
|
profiles = api_node._load_api_profiles()
|
|
self.assertNotIn('OpenRouter', profiles)
|
|
|
|
def test_custom_profile_file_has_configuration_not_secret(self):
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
path = Path(tmp) / 'profiles.json'
|
|
path.write_text(json.dumps({'profiles': {'Production': {
|
|
'provider': 'Test',
|
|
'base_url': 'https://example.test/v1',
|
|
'auth': 'bearer_env',
|
|
'api_key_env': 'TEST_API_KEY',
|
|
'allowed_models': ['approved/model'],
|
|
'max_tokens_limit': 1024,
|
|
'max_input_chars': 10000,
|
|
'max_timeout_seconds': 60,
|
|
'allow_images': True,
|
|
'allow_image_url': False,
|
|
'max_image_pixels': 16777216,
|
|
'max_image_bytes': 20000000,
|
|
}}}), encoding='utf-8')
|
|
with mock.patch.dict(os.environ, {
|
|
api_node.DISABLE_BUILTINS_ENV: '1',
|
|
api_node.PROFILE_FILE_ENV: str(path),
|
|
}, clear=False):
|
|
profiles = api_node._load_api_profiles()
|
|
self.assertEqual(list(profiles), ['Production'])
|
|
self.assertEqual(profiles['Production']['api_key_env'], 'TEST_API_KEY')
|
|
self.assertTrue(profiles['Production']['allow_images'])
|
|
self.assertNotIn('api_key', profiles['Production'])
|
|
|
|
def test_catalog_is_ux_only_and_never_becomes_an_allowlist(self):
|
|
# The curated catalog must not be interpreted as a security boundary.
|
|
catalog = api_node._load_model_catalog()
|
|
self.assertIsInstance(catalog, dict)
|
|
self.assertIn('OrcaRouter', catalog)
|
|
orca = catalog['OrcaRouter']
|
|
self.assertIsInstance(orca.get('models'), list)
|
|
catalog_models = api_node._catalog_models_for_profile('OrcaRouter')
|
|
self.assertTrue(catalog_models)
|
|
self.assertTrue(any(
|
|
isinstance(model, dict) and model.get('id') == catalog_models[0]
|
|
for model in orca['models']
|
|
))
|
|
|
|
def test_catalog_missing_or_broken_is_failure_tolerant(self):
|
|
with mock.patch.object(api_node, 'CATALOG_PATH', Path('does/not/exist.json')):
|
|
api_node._CATALOG_CACHE = None
|
|
self.assertEqual(api_node._load_model_catalog(), {})
|
|
self.assertEqual(api_node._catalog_models_for_profile('OrcaRouter'), [])
|
|
api_node._CATALOG_CACHE = None
|
|
|
|
def test_catalog_does_not_replace_allowed_models(self):
|
|
# Even if a profile has no allowlist, a request for a non-listed model must
|
|
# be passed through unchanged — the catalog is UI only.
|
|
p = self.profile()
|
|
self.assertIsNone(p['allowed_models'])
|
|
self.assertEqual(api_node._validate_model(p, 'custom/provider-model'), 'custom/provider-model')
|
|
self.assertEqual(api_node._validate_model(p, 'anything/else'), 'anything/else')
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|