Fix merged context vars and add test

Author: sanjayprabhu

projects/fal/src/fal/app.py

@@ -83,12 +83,10 @@ async def open_isolate_channel(address: str) -> async_grpc.Channel | None:
     return channel
 
 
-def merge_contextvars(logger_labels: dict[str, str]) -> dict[str, str]:
-    ctx = contextvars.copy_context()
-    for k in ctx:
-        if k.name.startswith(LOG_CONTEXT_PREFIX) and ctx[k] is not Ellipsis:
-            logger_labels.setdefault(k.name[len(LOG_CONTEXT_PREFIX) :], ctx[k])
-    return logger_labels
+def merge_contextvars(logger_labels: dict[str, str]) -> None:
+    for k, v in logger_labels.items():
+        contextvar = contextvars.ContextVar(f"{LOG_CONTEXT_PREFIX}{k}")
+        contextvar.set(v)
 
 
 def clear_contextvars() -> None:

projects/fal/tests/unit/test_app.py

@@ -1,11 +1,13 @@
 from __future__ import annotations
 
+import contextvars
 import os
 
 import pytest
 
 import fal
 from fal import App, endpoint
+from fal.app import merge_contextvars, clear_contextvars, LOG_CONTEXT_PREFIX
 from fal.container import ContainerImage
 
 
@@ -114,6 +116,39 @@ class LeakCheckApp(App):
     assert "app_auth" not in hk
 
 
+def test_merge_context_vars():
+    labels = {"fal_request_id": "123", "fal_endpoint": "/"}
+    request_id_var = f"{LOG_CONTEXT_PREFIX}fal_request_id"
+    endpoint_var = f"{LOG_CONTEXT_PREFIX}fal_endpoint"
+    unrelated_var = "unrelated_key"
+    contextvars.ContextVar(unrelated_var).set("value")
+
+    # We have to convert to dict and lookup by name because each ContextVar
+    # is a different object. Since merge_contextvars creates new ContextVars,
+    # we can't just do direct lookups.
+    vars = dict((k.name, v) for k, v in contextvars.copy_context().items())
+
+    assert vars.get(unrelated_var) == "value"
+
+    assert vars.get(request_id_var) is None
+    assert vars.get(endpoint_var) is None
+
+    merge_contextvars(labels)
+    vars = dict((k.name, v) for k, v in contextvars.copy_context().items())
+
+    assert vars.get(request_id_var) == "123"
+    assert vars.get(endpoint_var) == "/"
+
+    clear_contextvars()
+    vars = dict((k.name, v) for k, v in contextvars.copy_context().items())
+
+    # Cleared contextvars are set to Ellipsis
+    assert vars.get(request_id_var) is Ellipsis
+    assert vars.get(endpoint_var) is Ellipsis
+    # Does not clear unrelated contextvars
+    assert vars.get("unrelated_key") == "value"
+
+
 @pytest.mark.asyncio
 async def test_runner_state_lifecycle_complete():
     """Test that FAL_RUNNER_STATE transitions through all phases correctly"""