phoenix-oss/llama-stack-mirror
0
0
Fork 1
mirror of https://github.com/meta-llama/llama-stack.git synced 2025-12-03 09:53:45 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Revert "fix: context token explosion bug fixed"

Browse source 
This reverts commit 914bf84a1d.
This commit is contained in:
Emilio Garcia 2025-11-18 19:29:01 -05:00
parent e4250df351
commit fc4e4a838b
1 changed files with 6 additions and 23 deletions
Download patch file Download diff file Expand all files Collapse all files

29
src/llama_stack/core/utils/context.py
View file

@ -30,22 +30,10 @@ def preserve_contexts_async_generator[T](
# This is needed to propagate context across async generator boundaries # This is needed to propagate context across async generator boundaries
for context_var in context_vars: for context_var in context_vars:
try: try:
current_value = context_var.get() previous_values[context_var] = 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: except LookupError:
previous_values[context_var] = _MISSING previous_values[context_var] = _MISSING
# Var was unset, set it to initial value tokens[context_var] = context_var.set(initial_context_values[context_var.name])
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: def _restore_context_var(context_var: ContextVar, *, _tokens=tokens, _prev=previous_values) -> None:
token = _tokens.get(context_var) token = _tokens.get(context_var)
@ -57,15 +45,10 @@ def preserve_contexts_async_generator[T](
except (RuntimeError, ValueError): except (RuntimeError, ValueError):
pass pass
# Manual restoration when token reset fails if previous_value is _MISSING:
# Avoid creating unnecessary tokens to prevent accumulation that can cause deadlocks context_var.set(None)
if previous_value is not _MISSING: else:
try: context_var.set(previous_value)
current_value = context_var.get()
if current_value != previous_value:
context_var.set(previous_value)
except LookupError:
context_var.set(previous_value)
try: try:
item = await gen.__anext__() item = await gen.__anext__()