From e4250df3516407bfe623610b501afbe3635f908b Mon Sep 17 00:00:00 2001 From: Emilio Garcia Date: Tue, 18 Nov 2025 16:37:11 -0500 Subject: [PATCH] fix: context token explosion bug fixed --- src/llama_stack/core/utils/context.py | 29 +++++++++++++++++++++------ 1 file changed, 23 insertions(+), 6 deletions(-) diff --git a/src/llama_stack/core/utils/context.py b/src/llama_stack/core/utils/context.py index 0c3e41f00..9221979b9 100644 --- a/src/llama_stack/core/utils/context.py +++ b/src/llama_stack/core/utils/context.py @@ -30,10 +30,22 @@ def preserve_contexts_async_generator[T]( # This is needed to propagate context across async generator boundaries for context_var in context_vars: try: - previous_values[context_var] = context_var.get() + current_value = context_var.get() + previous_values[context_var] = current_value + # Only set if value actually changed to avoid creating unnecessary tokens + target_value = initial_context_values[context_var.name] + if current_value != target_value: + tokens[context_var] = context_var.set(target_value) + else: + tokens[context_var] = None except LookupError: previous_values[context_var] = _MISSING - tokens[context_var] = context_var.set(initial_context_values[context_var.name]) + # Var was unset, set it to initial value + target_value = initial_context_values[context_var.name] + if target_value is not None: + tokens[context_var] = context_var.set(target_value) + else: + tokens[context_var] = None def _restore_context_var(context_var: ContextVar, *, _tokens=tokens, _prev=previous_values) -> None: token = _tokens.get(context_var) @@ -45,10 +57,15 @@ def preserve_contexts_async_generator[T]( except (RuntimeError, ValueError): pass - if previous_value is _MISSING: - context_var.set(None) - else: - context_var.set(previous_value) + # Manual restoration when token reset fails + # Avoid creating unnecessary tokens to prevent accumulation that can cause deadlocks + if previous_value is not _MISSING: + try: + current_value = context_var.get() + if current_value != previous_value: + context_var.set(previous_value) + except LookupError: + context_var.set(previous_value) try: item = await gen.__anext__()