From 40a057f3d4866ef47deaa9dd468b9ae3b235c364 Mon Sep 17 00:00:00 2001 From: Chandrasekharan M Date: Tue, 21 Jul 2026 16:50:12 +0530 Subject: [PATCH 1/2] UN-3636 [FIX] Drop the ENVIRONMENT gate on the LLM mock It did not defend the case it was added for. The threat was a worker env block copied out of the test overlay into a real deployment, but the gate was written into that same block, so a copy carries it. Base compose also sets ENVIRONMENT=development on both workers that run the injection, so any deployment derived from it satisfied the gate regardless. That left one real case -- the mock var set alone somewhere that sets no ENVIRONMENT at all -- which holds by accident rather than design, in exchange for depending on a variable nothing else in the codebase reads. What actually guards the hatch is unchanged: it is off unless someone sets UNSTRACT_LLM_MOCK_RESPONSE, and it warns once per process while active. Making mocked spend distinguishable downstream is the defence worth having, and it belongs on the usage record rather than in a config check. Co-Authored-By: Claude Opus 4.8 --- tests/README.md | 2 +- tests/compose/docker-compose.test.yaml | 7 +--- unstract/sdk1/src/unstract/sdk1/llm.py | 21 ----------- unstract/sdk1/tests/test_mock_response.py | 46 ----------------------- 4 files changed, 2 insertions(+), 74 deletions(-) diff --git a/tests/README.md b/tests/README.md index 982b12f553..5034954e7a 100644 --- a/tests/README.md +++ b/tests/README.md @@ -158,7 +158,7 @@ The `platform` pytest fixture in `tests/e2e/conftest.py` reads those env vars; e Execute-path e2e tests must not call a real provider, so the rig sets `UNSTRACT_LLM_MOCK_RESPONSE` (default `MOCK_LLM_OK`) before boot, for any runtime and treating an exported empty string as unset. The test overlay forwards it into the workers, and `unstract.sdk1.llm` passes it to litellm as `mock_response`: for the non-streaming completion path litellm returns the string verbatim with fixed usage (10 prompt / 20 completion / 30 total), so both the answer and the token counts are exact-assertable. Streaming (`stream_complete`) goes through a different litellm path whose usage differs, so don't assume 10/20/30 there. Sentinels like `litellm.RateLimitError` force error paths. Unset (the production default) the hook is a no-op. -Mocking needs two conditions, not one: `ENVIRONMENT` must also be `test` or `development`. Production sets neither variable, so a stray mock var alone cannot fake completions and their billing — and a refusal is logged rather than silent, since setting the var at all means someone expected mocking. `development` is allowed because that is what base compose sets on the workers that run the injection; the test overlay pins `ENVIRONMENT=test` on those same two workers explicitly, so the tier can't lose its mock to a base-compose edit. +The variable being set is the only condition, deliberately. A second gate on `ENVIRONMENT` was tried and removed: it did not defend against the case it was written for — a worker env block copied out of this overlay carries both variables together — and base compose sets `ENVIRONMENT=development` on the two workers that run the injection, so any deployment derived from it satisfied the gate anyway. What remains is that the hatch is off unless someone sets it, and that it warns once per process while active. Making fake spend safe to *detect* rather than trying to prevent the config wants the usage record itself to say it was mocked. A CI/dev override wins (the rig only fills an unset value). Running these tests under the rig **fails** if the var is missing; running **without** the rig just skips the execute-path tests — export it on both sides (your shell and the workers) if you boot the stack yourself. diff --git a/tests/compose/docker-compose.test.yaml b/tests/compose/docker-compose.test.yaml index feb46d6eda..f089791fe8 100644 --- a/tests/compose/docker-compose.test.yaml +++ b/tests/compose/docker-compose.test.yaml @@ -22,18 +22,13 @@ services: environment: - ENVIRONMENT=test - # Execute-path e2e must never reach a real provider. The delay gives each - # mocked completion a known cost so per-file durations stay comparable, and - # ENVIRONMENT is the mock's second condition — set here rather than inherited - # so the tier can't lose its mock to a base-compose edit. + # Execute-path e2e must never reach a real provider. worker-executor-v2: environment: - - ENVIRONMENT=test - UNSTRACT_LLM_MOCK_RESPONSE=${UNSTRACT_LLM_MOCK_RESPONSE:-} worker-file-processing-v2: environment: - - ENVIRONMENT=test - UNSTRACT_LLM_MOCK_RESPONSE=${UNSTRACT_LLM_MOCK_RESPONSE:-} # Only the worker's fallback if it can't reach the backend; kept in step so it diff --git a/unstract/sdk1/src/unstract/sdk1/llm.py b/unstract/sdk1/src/unstract/sdk1/llm.py index 1befc96eb3..b0a712b49f 100644 --- a/unstract/sdk1/src/unstract/sdk1/llm.py +++ b/unstract/sdk1/src/unstract/sdk1/llm.py @@ -35,10 +35,6 @@ # Lets tests force a deterministic completion without a provider or a secret. # Unset in production, where this is a no-op. _MOCK_RESPONSE_ENV = "UNSTRACT_LLM_MOCK_RESPONSE" -# Second condition on the hatch, so a stray mock var alone can't fake -# completions and their billing. Deployments that set neither fail closed. -_ENVIRONMENT_ENV = "ENVIRONMENT" -_MOCK_ALLOWED_ENVIRONMENTS = frozenset({"test", "development"}) @lru_cache(maxsize=1) @@ -52,27 +48,10 @@ def _warn_mock_active() -> None: ) -@lru_cache(maxsize=1) -def _warn_mock_refused(environment: str) -> None: - # Loud rather than silent: the var being set at all means someone expected - # mocking, and they need to know why the bill is real. - logger.warning( - "%s is set but %s=%r is not one of %s — calling the provider for real.", - _MOCK_RESPONSE_ENV, - _ENVIRONMENT_ENV, - environment, - sorted(_MOCK_ALLOWED_ENVIRONMENTS), - ) - - def _inject_mock_response(completion_kwargs: dict[str, object]) -> None: mock = os.getenv(_MOCK_RESPONSE_ENV) if not mock or "mock_response" in completion_kwargs: return - environment = os.getenv(_ENVIRONMENT_ENV, "").strip().lower() - if environment not in _MOCK_ALLOWED_ENVIRONMENTS: - _warn_mock_refused(environment) - return _warn_mock_active() completion_kwargs["mock_response"] = mock diff --git a/unstract/sdk1/tests/test_mock_response.py b/unstract/sdk1/tests/test_mock_response.py index 6a4e0b1a59..807d7d05f1 100644 --- a/unstract/sdk1/tests/test_mock_response.py +++ b/unstract/sdk1/tests/test_mock_response.py @@ -38,13 +38,6 @@ def _inject(kwargs: dict[str, object]) -> dict[str, object]: @pytest.fixture(autouse=True) def _reset_warn_cache() -> None: _load_llm_module()._warn_mock_active.cache_clear() - _load_llm_module()._warn_mock_refused.cache_clear() - - -@pytest.fixture(autouse=True) -def _allowed_environment(monkeypatch: pytest.MonkeyPatch) -> None: - """The hatch needs a permitted ENVIRONMENT; the guard itself is tested below.""" - monkeypatch.setenv("ENVIRONMENT", "test") def test_inject_is_noop_when_env_unset(monkeypatch: pytest.MonkeyPatch) -> None: @@ -67,45 +60,6 @@ def test_inject_sets_mock_response_when_env_set( assert _inject({"model": "gpt-4o"})["mock_response"] == "canned answer" -@pytest.mark.parametrize("environment", ["production", "staging", "", " "]) -def test_mock_refused_outside_allowed_environments( - monkeypatch: pytest.MonkeyPatch, environment: str -) -> None: - # The whole point of the guard: a stray mock var in a real deployment must - # not fake completions and their billing. - monkeypatch.setenv("UNSTRACT_LLM_MOCK_RESPONSE", "canned") - monkeypatch.setenv("ENVIRONMENT", environment) - assert _inject({"model": "gpt-4o"}) == {"model": "gpt-4o"} - - -def test_mock_refused_when_environment_unset(monkeypatch: pytest.MonkeyPatch) -> None: - # Production k8s sets no ENVIRONMENT at all, so unset must fail closed. - monkeypatch.setenv("UNSTRACT_LLM_MOCK_RESPONSE", "canned") - monkeypatch.delenv("ENVIRONMENT", raising=False) - assert _inject({"model": "gpt-4o"}) == {"model": "gpt-4o"} - - -def test_refusal_is_warned_so_a_real_bill_is_never_silent( - monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture -) -> None: - monkeypatch.setenv("UNSTRACT_LLM_MOCK_RESPONSE", "canned") - monkeypatch.setenv("ENVIRONMENT", "production") - with caplog.at_level("WARNING", logger="unstract.sdk1.llm"): - _inject({"model": "gpt-4o"}) - assert any("not one of" in r.message for r in caplog.records), caplog.text - - -@pytest.mark.parametrize("environment", ["test", "development", "TEST", " Development "]) -def test_mock_applies_in_allowed_environments( - monkeypatch: pytest.MonkeyPatch, environment: str -) -> None: - # `development` is what base compose sets on the workers that run the - # injection; a guard that only accepted `test` would kill the local stack. - monkeypatch.setenv("UNSTRACT_LLM_MOCK_RESPONSE", "canned") - monkeypatch.setenv("ENVIRONMENT", environment) - assert _inject({"model": "gpt-4o"})["mock_response"] == "canned" - - def test_inject_does_not_clobber_explicit_mock_response( monkeypatch: pytest.MonkeyPatch, ) -> None: From 9ae7a645e6776ff9fb946dc77038f875848f2e6d Mon Sep 17 00:00:00 2001 From: Chandrasekharan M Date: Tue, 21 Jul 2026 17:23:09 +0530 Subject: [PATCH 2/2] UN-3636 [MISC] Drop the unread ENVIRONMENT variable from compose MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Nothing reads it: no service, worker, frontend or plugin looks the variable up, and the one consumer it ever had — the LLM mock gate — was removed in the previous commit. Dropping it everywhere keeps a dead knob from looking load-bearing to the next reader. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01Ryg9chVDJQggCybpq3YoY3 --- docker/docker-compose.yaml | 15 --------------- tests/README.md | 2 +- tests/compose/docker-compose.test.yaml | 5 ----- 3 files changed, 1 insertion(+), 21 deletions(-) diff --git a/docker/docker-compose.yaml b/docker/docker-compose.yaml index aa36cbf961..e85b9bf22c 100644 --- a/docker/docker-compose.yaml +++ b/docker/docker-compose.yaml @@ -39,7 +39,6 @@ services: - ./workflow_data:/data - ${TOOL_REGISTRY_CONFIG_SRC_PATH}:/data/tool_registry_config environment: - - ENVIRONMENT=development - APPLICATION_NAME=unstract-backend labels: - traefik.enable=true @@ -60,7 +59,6 @@ services: - rabbitmq - db environment: - - ENVIRONMENT=development - APPLICATION_NAME=unstract-worker-metrics labels: @@ -83,7 +81,6 @@ services: - redis - rabbitmq environment: - - ENVIRONMENT=development - APPLICATION_NAME=unstract-worker-ide-callback - WORKER_TYPE=ide_callback - WORKER_NAME=ide-callback-worker @@ -107,7 +104,6 @@ services: ports: - "5555:5555" environment: - - ENVIRONMENT=development - APPLICATION_NAME=unstract-celery-flower volumes: - unstract_data:/data @@ -129,7 +125,6 @@ services: - db - rabbitmq environment: - - ENVIRONMENT=development - APPLICATION_NAME=unstract-celery-beat # Frontend React app @@ -143,7 +138,6 @@ services: - backend - reverse-proxy environment: - - ENVIRONMENT=development # Running platform version, surfaced on the profile page. Reuses the same # ${VERSION} that tags the image, so it always matches what's deployed # (and updates on RC->stable promotion, which only retags — doesn't rebuild). @@ -219,7 +213,6 @@ services: - redis - rabbitmq environment: - - ENVIRONMENT=development - APPLICATION_NAME=unstract-worker-api-deployment-v2 - WORKER_TYPE=api_deployment - CELERY_QUEUES_API_DEPLOYMENT=${CELERY_QUEUES_API_DEPLOYMENT:-celery_api_deployments} @@ -252,7 +245,6 @@ services: - redis - rabbitmq environment: - - ENVIRONMENT=development - APPLICATION_NAME=unstract-worker-callback-v2 - WORKER_TYPE=callback - WORKER_NAME=callback-worker-v2 @@ -294,7 +286,6 @@ services: - redis - rabbitmq environment: - - ENVIRONMENT=development - APPLICATION_NAME=unstract-worker-file-processing-v2 - WORKER_TYPE=file_processing - WORKER_MODE=oss @@ -332,7 +323,6 @@ services: - redis - rabbitmq environment: - - ENVIRONMENT=development - APPLICATION_NAME=unstract-worker-general-v2 - WORKER_TYPE=general - WORKER_NAME=general-worker-v2 @@ -360,7 +350,6 @@ services: - redis - rabbitmq environment: - - ENVIRONMENT=development - APPLICATION_NAME=unstract-worker-notification-v2 - WORKER_TYPE=notification - WORKER_NAME=notification-worker-v2 @@ -409,7 +398,6 @@ services: - redis - rabbitmq environment: - - ENVIRONMENT=development - APPLICATION_NAME=unstract-worker-log-consumer-v2 - WORKER_TYPE=log_consumer - WORKER_NAME=log-consumer-worker-v2 @@ -458,7 +446,6 @@ services: - redis - rabbitmq environment: - - ENVIRONMENT=development - APPLICATION_NAME=unstract-worker-log-history-scheduler-v2 # Scheduler interval in seconds - LOG_HISTORY_CONSUMER_INTERVAL=${LOG_HISTORY_CONSUMER_INTERVAL:-5} @@ -483,7 +470,6 @@ services: - redis - rabbitmq environment: - - ENVIRONMENT=development - APPLICATION_NAME=unstract-worker-scheduler-v2 - WORKER_TYPE=scheduler - WORKER_NAME=scheduler-worker-v2 @@ -529,7 +515,6 @@ services: - rabbitmq - platform-service environment: - - ENVIRONMENT=development - APPLICATION_NAME=unstract-worker-executor-v2 - WORKER_TYPE=executor - WORKER_NAME=executor-worker-v2 diff --git a/tests/README.md b/tests/README.md index 5034954e7a..3628cd82c1 100644 --- a/tests/README.md +++ b/tests/README.md @@ -158,7 +158,7 @@ The `platform` pytest fixture in `tests/e2e/conftest.py` reads those env vars; e Execute-path e2e tests must not call a real provider, so the rig sets `UNSTRACT_LLM_MOCK_RESPONSE` (default `MOCK_LLM_OK`) before boot, for any runtime and treating an exported empty string as unset. The test overlay forwards it into the workers, and `unstract.sdk1.llm` passes it to litellm as `mock_response`: for the non-streaming completion path litellm returns the string verbatim with fixed usage (10 prompt / 20 completion / 30 total), so both the answer and the token counts are exact-assertable. Streaming (`stream_complete`) goes through a different litellm path whose usage differs, so don't assume 10/20/30 there. Sentinels like `litellm.RateLimitError` force error paths. Unset (the production default) the hook is a no-op. -The variable being set is the only condition, deliberately. A second gate on `ENVIRONMENT` was tried and removed: it did not defend against the case it was written for — a worker env block copied out of this overlay carries both variables together — and base compose sets `ENVIRONMENT=development` on the two workers that run the injection, so any deployment derived from it satisfied the gate anyway. What remains is that the hatch is off unless someone sets it, and that it warns once per process while active. Making fake spend safe to *detect* rather than trying to prevent the config wants the usage record itself to say it was mocked. +The variable being set is the only condition, deliberately. A second gate on `ENVIRONMENT` was tried and removed: it did not defend against the case it was written for — a worker env block copied out of this overlay carries both variables together — and compose declared `ENVIRONMENT` on every service anyway, so any deployment derived from it satisfied the gate. Nothing read that variable, so it was dropped everywhere along with the gate. What remains is that the hatch is off unless someone sets it, and that it warns once per process while active. Making fake spend safe to *detect* rather than trying to prevent the config wants the usage record itself to say it was mocked. A CI/dev override wins (the rig only fills an unset value). Running these tests under the rig **fails** if the var is missing; running **without** the rig just skips the execute-path tests — export it on both sides (your shell and the workers) if you boot the stack yourself. diff --git a/tests/compose/docker-compose.test.yaml b/tests/compose/docker-compose.test.yaml index f089791fe8..8b18a3ca4b 100644 --- a/tests/compose/docker-compose.test.yaml +++ b/tests/compose/docker-compose.test.yaml @@ -7,20 +7,15 @@ services: # contributor run without setting it. image: unstract/backend:${UNSTRACT_TEST_VERSION:-latest} environment: - - ENVIRONMENT=test # Workers ask the backend for this, so this is the value that normally # takes effect for fan-out. - MAX_PARALLEL_FILE_BATCHES=${MAX_PARALLEL_FILE_BATCHES:-3} platform-service: image: unstract/platform-service:${UNSTRACT_TEST_VERSION:-latest} - environment: - - ENVIRONMENT=test runner: image: unstract/runner:${UNSTRACT_TEST_VERSION:-latest} - environment: - - ENVIRONMENT=test # Execute-path e2e must never reach a real provider. worker-executor-v2: