From 6a6d1d77553b2df08349a2fa8e0cf052a64488c6 Mon Sep 17 00:00:00 2001 From: Abhishek Kumar Date: Tue, 30 Jun 2026 12:59:43 +0530 Subject: [PATCH] Fix short masked key preservation --- .../configuration/ai_model_configuration.py | 13 +- api/services/configuration/masking.py | 6 + api/services/configuration/merge.py | 13 +- api/tests/test_ai_model_configuration_v2.py | 94 +++++++++++++ api/tests/test_resolve_effective_config.py | 18 +++ api/tests/test_resolve_masked_api_keys.py | 123 ++++++++++++++++++ 6 files changed, 257 insertions(+), 10 deletions(-) create mode 100644 api/tests/test_resolve_masked_api_keys.py diff --git a/api/services/configuration/ai_model_configuration.py b/api/services/configuration/ai_model_configuration.py index 3aabdae6..258df551 100644 --- a/api/services/configuration/ai_model_configuration.py +++ b/api/services/configuration/ai_model_configuration.py @@ -28,6 +28,7 @@ from api.schemas.ai_model_configuration import ( ) from api.services.configuration.masking import ( SERVICE_SECRET_FIELDS, + _secret_values_differ, contains_masked_key, mask_key, resolve_masked_api_keys, @@ -244,11 +245,13 @@ def merge_ai_model_configuration_v2_secrets( existing_dograh = existing_dict.get("dograh") or {} incoming_key = incoming_dograh.get("api_key") existing_key = existing_dograh.get("api_key") - if incoming_key and existing_key and contains_masked_key(incoming_key): - incoming_dograh["api_key"] = resolve_masked_api_keys( + if incoming_key and existing_key: + resolved_key = resolve_masked_api_keys( incoming_key, existing_key, ) + if _secret_values_differ(incoming_key, resolved_key): + incoming_dograh["api_key"] = resolved_key if incoming_dict.get("mode") == "byok" and existing_dict.get("mode") == "byok": _merge_byok_secret_fields(incoming_dict.get("byok"), existing_dict.get("byok")) @@ -381,11 +384,13 @@ def _merge_service_secret_fields(incoming: dict, existing: dict): existing_secret = existing[secret_field] if incoming_secret is None: incoming[secret_field] = existing_secret - elif contains_masked_key(incoming_secret): - incoming[secret_field] = resolve_masked_api_keys( + else: + resolved_secret = resolve_masked_api_keys( incoming_secret, existing_secret, ) + if _secret_values_differ(incoming_secret, resolved_secret): + incoming[secret_field] = resolved_secret def _raise_if_masked_secret(value): diff --git a/api/services/configuration/masking.py b/api/services/configuration/masking.py index a7e1af6a..7df5b8d3 100644 --- a/api/services/configuration/masking.py +++ b/api/services/configuration/masking.py @@ -80,6 +80,12 @@ def is_mask_of(masked: str, real_key: str) -> bool: return mask_key(real_key) == masked +def _secret_values_differ(left: str | list[str], right: str | list[str]) -> bool: + left_values = left if isinstance(left, list) else [left] + right_values = right if isinstance(right, list) else [right] + return left_values != right_values + + def resolve_masked_api_keys( incoming: str | list[str], existing: str | list[str] ) -> str | list[str]: diff --git a/api/services/configuration/merge.py b/api/services/configuration/merge.py index 3100fa45..77e608ca 100644 --- a/api/services/configuration/merge.py +++ b/api/services/configuration/merge.py @@ -11,7 +11,7 @@ from api.schemas.ai_model_configuration import EffectiveAIModelConfiguration from api.services.configuration.masking import ( MODEL_OVERRIDE_FIELDS, SERVICE_SECRET_FIELDS, - contains_masked_key, + _secret_values_differ, resolve_masked_api_keys, ) @@ -50,14 +50,15 @@ def _merge_service_secret_fields( incoming_secret = incoming_cfg.get(secret_field) existing_secret = existing_cfg[secret_field] if incoming_secret is not None: - if contains_masked_key(incoming_secret): + resolved_secret = resolve_masked_api_keys( + incoming_secret, + existing_secret, + ) + if _secret_values_differ(incoming_secret, resolved_secret): incoming_cfg[secret_field] = ( existing_secret if masked_value_preserves_full_secret - else resolve_masked_api_keys( - incoming_secret, - existing_secret, - ) + else resolved_secret ) elif preserve_missing: incoming_cfg[secret_field] = existing_secret diff --git a/api/tests/test_ai_model_configuration_v2.py b/api/tests/test_ai_model_configuration_v2.py index 296872fa..4f87c18e 100644 --- a/api/tests/test_ai_model_configuration_v2.py +++ b/api/tests/test_ai_model_configuration_v2.py @@ -191,6 +191,100 @@ def test_masked_dograh_key_is_preserved_when_saving_same_mode(): check_for_masked_keys_in_ai_model_configuration_v2(merged) +def test_short_masked_dograh_key_is_preserved_when_saving_same_mode(): + existing = OrganizationAIModelConfigurationV2( + mode="dograh", + dograh=DograhManagedAIModelConfiguration(api_key="SECRET"), + ) + incoming = OrganizationAIModelConfigurationV2( + mode="dograh", + dograh=DograhManagedAIModelConfiguration(api_key=mask_key("SECRET")), + ) + + assert incoming.dograh.api_key == "**CRET" + + merged = merge_ai_model_configuration_v2_secrets(incoming, existing) + + assert merged.dograh.api_key == "SECRET" + check_for_masked_keys_in_ai_model_configuration_v2(merged) + + +def test_unmasked_dograh_key_replaces_existing_when_saving_same_mode(): + existing = OrganizationAIModelConfigurationV2( + mode="dograh", + dograh=DograhManagedAIModelConfiguration(api_key="SECRET"), + ) + incoming = OrganizationAIModelConfigurationV2( + mode="dograh", + dograh=DograhManagedAIModelConfiguration(api_key="NEWKEY"), + ) + + merged = merge_ai_model_configuration_v2_secrets(incoming, existing) + + assert merged.dograh.api_key == "NEWKEY" + + +def test_short_masked_byok_service_key_is_preserved_when_saving_same_mode(): + existing = OrganizationAIModelConfigurationV2.model_validate( + { + "mode": "byok", + "byok": { + "mode": "pipeline", + "pipeline": { + "llm": { + "provider": "openai", + "api_key": "SECRET", + "model": "gpt-4.1", + }, + "tts": { + "provider": "elevenlabs", + "api_key": "el-real-secret", + "model": "eleven_flash_v2_5", + "voice": "Rachel", + }, + "stt": { + "provider": "deepgram", + "api_key": "dg-real-secret", + "model": "nova-3-general", + }, + }, + }, + } + ) + incoming = OrganizationAIModelConfigurationV2.model_validate( + { + "mode": "byok", + "byok": { + "mode": "pipeline", + "pipeline": { + "llm": { + "provider": "openai", + "api_key": mask_key("SECRET"), + "model": "gpt-4.1", + }, + "tts": { + "provider": "elevenlabs", + "api_key": "el-real-secret", + "model": "eleven_flash_v2_5", + "voice": "Rachel", + }, + "stt": { + "provider": "deepgram", + "api_key": "dg-real-secret", + "model": "nova-3-general", + }, + }, + }, + } + ) + + assert incoming.byok.pipeline.llm.api_key == "**CRET" + + merged = merge_ai_model_configuration_v2_secrets(incoming, existing) + + assert merged.byok.pipeline.llm.api_key == "SECRET" + + def test_masked_v2_configuration_masks_nested_service_keys(): config = OrganizationAIModelConfigurationV2( mode="byok", diff --git a/api/tests/test_resolve_effective_config.py b/api/tests/test_resolve_effective_config.py index 1b9ad8c6..59b7d4b2 100644 --- a/api/tests/test_resolve_effective_config.py +++ b/api/tests/test_resolve_effective_config.py @@ -563,6 +563,24 @@ class TestWorkflowConfigurationSecrets: assert merged["model_overrides"]["tts"]["voice"] == "Bella" assert incoming["model_overrides"]["tts"]["api_key"] != "el-real-tts-key" + def test_restores_short_masked_model_override_secret_from_existing_config(self): + existing = { + "model_overrides": { + "llm": { + "provider": "openai", + "api_key": "SECRET", + "model": "gpt-4.1-mini", + } + } + } + incoming = mask_workflow_configurations(existing) + + assert incoming["model_overrides"]["llm"]["api_key"] == "**CRET" + + merged = merge_workflow_configuration_secrets(incoming, existing) + + assert merged["model_overrides"]["llm"]["api_key"] == "SECRET" + def test_single_masked_key_preserves_existing_multi_key_override(self): existing = { "model_overrides": { diff --git a/api/tests/test_resolve_masked_api_keys.py b/api/tests/test_resolve_masked_api_keys.py new file mode 100644 index 00000000..44f8ce91 --- /dev/null +++ b/api/tests/test_resolve_masked_api_keys.py @@ -0,0 +1,123 @@ +"""Unit tests for ``resolve_masked_api_keys``. + +The resolver restores real API keys when the client echoes back the masked +placeholder, while leaving genuinely new keys untouched. It must handle adds, +removes, reorders, and partial replacements across both scalar and list values. +""" + +from api.services.configuration.masking import ( + _secret_values_differ, + mask_key, + resolve_masked_api_keys, +) + +REAL_A = "sk-real-key-aaaaaaaaaaaa1111" +REAL_B = "sk-real-key-bbbbbbbbbbbb2222" +REAL_C = "sk-real-key-cccccccccccc3333" +REAL_D_SMALL = "ABCDE" + +MASK_A = mask_key(REAL_A) +MASK_B = mask_key(REAL_B) +MASK_C = mask_key(REAL_C) +MASK_D_SMALL = mask_key(REAL_D_SMALL) + + +class TestScalar: + def test_masked_scalar_restores_existing(self): + assert resolve_masked_api_keys(MASK_A, REAL_A) == REAL_A + + def test_unmasked_scalar_is_kept(self): + new_key = "sk-brand-new-real-key-9999" + assert resolve_masked_api_keys(new_key, REAL_A) == new_key + + def test_mask_of_different_key_is_kept_verbatim(self): + # Incoming is a mask, but not the mask of the existing key -> kept as-is. + assert resolve_masked_api_keys(MASK_B, REAL_A) == MASK_B + + def test_small_mask_is_resolved(self): + assert resolve_masked_api_keys(MASK_D_SMALL, REAL_D_SMALL) == REAL_D_SMALL + + +class TestList: + def test_all_masked_restores_all_reals(self): + result = resolve_masked_api_keys([MASK_A, MASK_B], [REAL_A, REAL_B]) + assert result == [REAL_A, REAL_B] + + def test_reorder_is_preserved(self): + # Client sent the keys back in a different order than stored. + result = resolve_masked_api_keys([MASK_B, MASK_A], [REAL_A, REAL_B]) + assert result == [REAL_B, REAL_A] + + def test_new_key_added_alongside_masked(self): + new_key = "sk-brand-new-real-key-9999" + result = resolve_masked_api_keys([MASK_A, new_key], [REAL_A]) + assert result == [REAL_A, new_key] + + def test_key_removed(self): + # Only one of two stored keys is echoed back. + result = resolve_masked_api_keys([MASK_B], [REAL_A, REAL_B]) + assert result == [REAL_B] + + def test_partial_replacement(self): + new_key = "sk-brand-new-real-key-9999" + result = resolve_masked_api_keys( + [MASK_A, new_key, MASK_C], [REAL_A, REAL_B, REAL_C] + ) + assert result == [REAL_A, new_key, REAL_C] + + def test_unmasked_keys_kept_verbatim(self): + result = resolve_masked_api_keys([REAL_A, REAL_B], [REAL_C]) + assert result == [REAL_A, REAL_B] + + +class TestUsedDeduplication: + def test_each_existing_key_consumed_at_most_once(self): + # Two identical masked placeholders, but only one matching real key + # exists. The first consumes REAL_A; the second has no unused match and + # is therefore kept verbatim rather than duplicating the real key. + result = resolve_masked_api_keys([MASK_A, MASK_A], [REAL_A]) + assert result == [REAL_A, MASK_A] + + def test_two_identical_masks_two_identical_reals(self): + result = resolve_masked_api_keys([MASK_A, MASK_A], [REAL_A, REAL_A]) + assert result == [REAL_A, REAL_A] + + +class TestMixedTypes: + def test_scalar_incoming_list_existing(self): + # A masked scalar resolved against a list of stored keys. + result = resolve_masked_api_keys(MASK_B, [REAL_A, REAL_B]) + assert result == [REAL_B] + + def test_list_incoming_scalar_existing(self): + result = resolve_masked_api_keys([MASK_A], REAL_A) + assert result == [REAL_A] + + +class TestEmpty: + def test_empty_incoming_list(self): + assert resolve_masked_api_keys([], [REAL_A]) == [] + + def test_no_existing_keys_keeps_incoming(self): + result = resolve_masked_api_keys([MASK_A, MASK_B], []) + assert result == [MASK_A, MASK_B] + + +class TestSecretValuesDiffer: + def test_same_scalar_values_do_not_differ(self): + assert _secret_values_differ(REAL_A, REAL_A) is False + + def test_different_scalar_values_differ(self): + assert _secret_values_differ(REAL_A, REAL_B) is True + + def test_same_list_values_do_not_differ(self): + assert _secret_values_differ([REAL_A, REAL_B], [REAL_A, REAL_B]) is False + + def test_same_list_values_in_different_order_differ(self): + assert _secret_values_differ([REAL_A, REAL_B], [REAL_B, REAL_A]) is True + + def test_scalar_and_single_item_list_do_not_differ(self): + assert _secret_values_differ(REAL_A, [REAL_A]) is False + + def test_scalar_and_multi_item_list_differ(self): + assert _secret_values_differ(REAL_A, [REAL_A, REAL_B]) is True