Files
goodguy1963-ComfyUI-Thinkin…/tests/test_openai_api_security.py
T
Marcel Hofer 7d0600f141 chore: prepare v3.0.0 release
Add separate image and text Decision Router nodes with configurable system prompts, up to 52 user-defined routes, per-item decisions, normalized decision_json output, and native IMAGE/STRING route sockets.

Run Clef-Flash 9B, Clef 27B, and OpenJev 27B in-process through llama.cpp. Add decision-head, logits, MTMD vision, and dynamic CUDA backend capability checks before model downloads. Keep model weights and vision projectors out of the repository; do not launch a separate llama-server or upgrade PyTorch/Transformers.

Block inactive route outputs so unused downstream branches are skipped. Keep dynamic route names synchronized with connected outputs and preserve stable socket identities. Register the commercial UI variant while keeping non-commercial OpenJev unavailable in commercial builds.

Update all 3.0.0 package metadata, README, runtime and workflow guides, the illustrated router example, and the detailed update log. Refresh regression expectations for commercial router registration and the current provider model catalog.

Validation: all 220 tests pass in the bundled ComfyUI Python 3.12 runtime. Python compilation, JavaScript syntax, version metadata, README image, and staged diff checks pass. The Image Router was manually exercised with Clef-Flash 9B on fox, dog, cat, and none routes.

Coverage limits: full-weight inference for Clef 27B and OpenJev 27B, text-mode smoke tests, and clean installs across other OS/Python/CUDA combinations remain unverified. The documented Windows x64, CPython 3.12, CUDA 13.0 JamePeng runtime is the only verified build matrix.
2026-10-28 22:29:51 +01:00

368 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_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()