From e9eea8682c49043f02253794dde56f2e711cb095 Mon Sep 17 00:00:00 2001 From: Tom Aisthorpe Date: Thu, 30 Apr 2026 13:31:46 +0100 Subject: [PATCH 1/2] Add protect mode validation --- aikido_zen/__init__.py | 7 +++++++ aikido_zen/init_test.py | 10 ++++++++++ 2 files changed, 17 insertions(+) diff --git a/aikido_zen/__init__.py b/aikido_zen/__init__.py index abf507e00..81c07d60d 100644 --- a/aikido_zen/__init__.py +++ b/aikido_zen/__init__.py @@ -24,6 +24,9 @@ from aikido_zen.helpers.aikido_disabled_flag_active import aikido_disabled_flag_active +VALID_MODES = ("daemon", "daemon_only", "daemon_disabled") + + def protect(mode="daemon", token=""): """ Mode can be set to : @@ -32,6 +35,10 @@ def protect(mode="daemon", token=""): - daemon_disabled : This will import sinks/sources but won't start a background process Protect user's application """ + if mode not in VALID_MODES: + raise ValueError( + f"Invalid mode {mode!r}, expected one of {VALID_MODES}. To pass a token, use protect(token=...)" + ) if aikido_disabled_flag_active(): # Do not run any aikido code when the disabled flag is on return diff --git a/aikido_zen/init_test.py b/aikido_zen/init_test.py index 5a99f7e98..f9ec3ae13 100644 --- a/aikido_zen/init_test.py +++ b/aikido_zen/init_test.py @@ -17,3 +17,13 @@ def test_protect_with_django(monkeypatch, caplog): def test_protect_sets_token(): aikido_zen.protect(token="MY_TOKEN_1") assert get_token_from_env().token == "MY_TOKEN_1" + + +def test_protect_rejects_invalid_mode(): + with pytest.raises(ValueError, match=r"Invalid mode .*protect\(token=\.\.\.\)"): + aikido_zen.protect("AIK_RUNTIME_some-token-string") + + +@pytest.mark.parametrize("mode", ["daemon", "daemon_only", "daemon_disabled"]) +def test_protect_accepts_valid_modes(mode): + aikido_zen.protect(mode=mode) From c02e3ba8a6e6d7a27f1d89ac752e5fe5e8b70424 Mon Sep 17 00:00:00 2001 From: Tom Aisthorpe Date: Thu, 30 Apr 2026 14:03:43 +0100 Subject: [PATCH 2/2] Fix langchain test using wrong mode --- aikido_zen/sinks/tests/langchain_test.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/aikido_zen/sinks/tests/langchain_test.py b/aikido_zen/sinks/tests/langchain_test.py index e1b7a6940..55d1f4155 100644 --- a/aikido_zen/sinks/tests/langchain_test.py +++ b/aikido_zen/sinks/tests/langchain_test.py @@ -5,7 +5,7 @@ import aikido_zen from aikido_zen.thread.thread_cache import get_cache -aikido_zen.protect(mode="daemon-disabled") +aikido_zen.protect(mode="daemon_disabled") skip_no_openai_key = pytest.mark.skipif( "OPENAI_API_KEY" not in os.environ,