Skip to content

Fix safety finding leaks in LangChain cache wrappers and LlamaIndex streaming depth counters #18

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
manas-fortifyroot merged 1 commit into from
Apr 3, 2026
Merged

Fix safety finding leaks in LangChain cache wrappers and LlamaIndex streaming depth counters #18

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions ...opentelemetry-instrumentation-langchain/opentelemetry/instrumentation/langchain/safety.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -70,6 +70,8 @@ def _safety_on_worker():


def base_chat_model_generate_with_cache_wrapper(wrapped, instance, args, kwargs):
from opentelemetry.instrumentation.fortifyroot import discard_deferred_findings
discard_deferred_findings() # FR: prevent stale findings from prior request on same thread
response = wrapped(*args, **kwargs)
_apply_chat_result_completion_safety(instance, response)
return response
Expand All @@ -91,6 +93,8 @@ def _completion_safety_on_worker():


def base_llm_generate_helper_wrapper(wrapped, instance, args, kwargs):
from opentelemetry.instrumentation.fortifyroot import discard_deferred_findings
discard_deferred_findings() # FR: prevent stale findings from prior request on same thread
response = wrapped(*args, **kwargs)
_apply_llm_result_completion_safety(instance, response)
return response
Expand Down
51 changes: 22 additions & 29 deletions ...entelemetry-instrumentation-llamaindex/opentelemetry/instrumentation/llamaindex/safety.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -153,16 +153,12 @@ def llm_stream_chat_wrapper(wrapped, instance, args, kwargs): # FR: streaming s
_enter_safety()
try:
updated_args, updated_kwargs = _apply_chat_prompt_safety(instance, args, kwargs)
except Exception:
span = trace.get_current_span()
span_name = f"{instance.__class__.__name__}.stream_chat"
safety = LlamaIndexStreamingSafety(span, span_name, LLMRequestTypeValues.CHAT.value)
return wrap_stream(wrapped(*updated_args, **updated_kwargs), safety)
finally:
_exit_safety()
raise
# NOTE: _exit_safety NOT called here — the generator from wrap_stream may
# outlive this call frame. The depth counter stays +1 during iteration,
# which is harmless (prevents spurious clears during nested calls).
span = trace.get_current_span()
span_name = f"{instance.__class__.__name__}.stream_chat"
safety = LlamaIndexStreamingSafety(span, span_name, LLMRequestTypeValues.CHAT.value)
return wrap_stream(wrapped(*updated_args, **updated_kwargs), safety)


async def llm_astream_chat_wrapper(wrapped, instance, args, kwargs): # FR: streaming safety
Expand All @@ -183,14 +179,13 @@ def _safety_on_worker():

(updated_args, updated_kwargs), worker_findings = await asyncio.to_thread(_safety_on_worker)
inject_deferred_findings(worker_findings)
except Exception:
result = wrapped(*updated_args, **updated_kwargs)
if _inspect.iscoroutine(result): # coroutine returning async gen (standard LI pattern)
result = await result
safety = LlamaIndexStreamingSafety(span, span_name, LLMRequestTypeValues.CHAT.value)
return make_async_stream(result, safety)
finally:
_exit_safety()
raise
result = wrapped(*updated_args, **updated_kwargs)
if _inspect.iscoroutine(result): # coroutine returning async gen (standard LI pattern)
result = await result
safety = LlamaIndexStreamingSafety(span, span_name, LLMRequestTypeValues.CHAT.value)
return make_async_stream(result, safety)


def llm_stream_complete_wrapper(wrapped, instance, args, kwargs): # FR: streaming safety
Expand All @@ -201,13 +196,12 @@ def llm_stream_complete_wrapper(wrapped, instance, args, kwargs): # FR: streami
_enter_safety()
try:
updated_args, updated_kwargs = _apply_completion_prompt_safety(instance, args, kwargs)
except Exception:
span = trace.get_current_span()
span_name = f"{instance.__class__.__name__}.stream_complete"
safety = LlamaIndexStreamingSafety(span, span_name, LLMRequestTypeValues.COMPLETION.value)
return wrap_stream(wrapped(*updated_args, **updated_kwargs), safety)
finally:
_exit_safety()
raise
span = trace.get_current_span()
span_name = f"{instance.__class__.__name__}.stream_complete"
safety = LlamaIndexStreamingSafety(span, span_name, LLMRequestTypeValues.COMPLETION.value)
return wrap_stream(wrapped(*updated_args, **updated_kwargs), safety)


async def llm_astream_complete_wrapper(wrapped, instance, args, kwargs): # FR: streaming safety
Expand All @@ -228,14 +222,13 @@ def _safety_on_worker():

(updated_args, updated_kwargs), worker_findings = await asyncio.to_thread(_safety_on_worker)
inject_deferred_findings(worker_findings)
except Exception:
result = wrapped(*updated_args, **updated_kwargs)
if _inspect.iscoroutine(result):
result = await result
safety = LlamaIndexStreamingSafety(span, span_name, LLMRequestTypeValues.COMPLETION.value)
return make_async_stream(result, safety)
finally:
_exit_safety()
raise
result = wrapped(*updated_args, **updated_kwargs)
if _inspect.iscoroutine(result):
result = await result
safety = LlamaIndexStreamingSafety(span, span_name, LLMRequestTypeValues.COMPLETION.value)
return make_async_stream(result, safety)


_METHOD_WRAPPERS = {
Expand Down
Loading