Skip to content

fix(router): propagate session_id/user_id to LiteLLM agent-completion traces #325

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
sheepdestroyer merged 3 commits into from
Jul 13, 2026
Merged

fix(router): propagate session_id/user_id to LiteLLM agent-completion traces #325

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
5ccdf8e
fix(router): propagate session_id/user_id to LiteLLM agent-completion…
Jul 13, 2026
1c63e77
fix(router): remove underscored headers, deep-copy metadata dict
Jul 13, 2026
2cc6791
chore: global consistency — remove duplicate globals, redundant impor…
Jul 13, 2026
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
25 changes: 12 additions & 13 deletions router/main.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -44,12 +44,6 @@
_REDIS_RETRY_INTERVAL_SECONDS = 5.0




_redis_client = None
_redis_last_init_attempt = 0.0
_REDIS_RETRY_INTERVAL_SECONDS = 5.0

def _valkey_port() -> int:
"""Resolve the Valkey cache port from env, preferring VALKEY_CACHE_PORT."""
port_str = os.getenv("VALKEY_CACHE_PORT") or os.getenv("VALKEY_PORT", "6379")
Expand Down Expand Up @@ -334,7 +328,7 @@ def _end_parent_obs(parent_obs, output=None, metadata=None) -> None:
Non-fatal — swallows all exceptions.
"""
if parent_obs is None:
return None
return
try:
update_kwargs = {}
if output is not None:
Expand All @@ -347,7 +341,7 @@ def _end_parent_obs(parent_obs, output=None, metadata=None) -> None:
except Exception:
logger.debug("_end_parent_obs failed (non-fatal)", exc_info=True)
pass
return None
return


def _end_child_span(span, output=None, metadata=None) -> None:
Expand Down Expand Up @@ -383,7 +377,7 @@ def _close_prop_ctx(prop_ctx):
except Exception:
logger.debug("_close_prop_ctx failed (non-fatal)", exc_info=True)
pass
return None
return


def _make_prop_ctx(session_id, user_id):
Expand Down Expand Up @@ -2231,8 +2225,7 @@ async def chat_completions(request: Request):
# Real native stream generator
async def native_agy_stream_generator(stream_gen, model_name):
"""Asynchronous generator yielding native OpenAI-compatible streaming chunks from the real agy daemon."""
import uuid

import time
created_time = int(time.time())
chunk_id = f"chatcmpl-{uuid.uuid4().hex[:12]}"
token_count = 0
Expand Down Expand Up @@ -2393,8 +2386,6 @@ async def native_agy_stream_generator(stream_gen, model_name):

async def agy_stream_generator():
"""Asynchronous generator yielding simulated OpenAI-compatible streaming chunks from a static agy response."""
import uuid

created_time = int(time.time())
chunk_id = f"chatcmpl-{uuid.uuid4().hex[:12]}"
chunk_size = 40
Expand Down Expand Up @@ -2612,7 +2603,15 @@ async def execute_proxy(model_name: str):
body_to_send["metadata"], dict
):
body_to_send["metadata"] = {}
else:
# Deep-copy to avoid mutating original body's metadata
# during fallback retries (shallow copy shares the dict)
body_to_send["metadata"] = dict(body_to_send["metadata"])
body_to_send["metadata"]["trace_name"] = "agent-completion"
Comment on lines 2604 to 2610

@gemini-code-assist gemini-code-assist Bot Jul 13, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Since body_to_send is created as a shallow copy of body (body_to_send = body.copy()), mutating its nested "metadata" dictionary directly will also mutate the original body dictionary. If the request fails and falls back to another model (which calls execute_proxy again with the same body object), the subsequent calls will receive the mutated metadata.

To prevent unintended side effects and ensure clean fallback behavior, copy the "metadata" dictionary if it already exists.

Suggested change
):
body_to_send["metadata"] = {}
body_to_send["metadata"]["trace_name"] = "agent-completion"
):
body_to_send['metadata'] = {}
else:
body_to_send['metadata'] = body_to_send['metadata'].copy()
body_to_send['metadata']['trace_name'] = 'agent-completion'

if _trace_session_id:
body_to_send["metadata"]["session_id"] = _trace_session_id
if _trace_user_id:
body_to_send["metadata"]["trace_user_id"] = _trace_user_id

if body.get("stream", False):
logger.info(f"Proxying streaming to LiteLLM as model={model_name}")
Expand Down