litellm/qa_sticky_session.sh
Krrish Dholakia 6c21029cb7
feat(sandbox): reuse e2b container across requests when metadata.session_id is set (#31688)
* feat(sandbox): reuse e2b container across requests when metadata.session_id is set

When a client passes `metadata.session_id` in a /chat/completions request
alongside a code_interpreter tool, the proxy now routes all requests sharing
that session_id to the same sandbox container. State (variables, imports,
installed packages) persists across requests within the session.

Without a session_id the existing ephemeral behavior is unchanged: one
container per agentic loop, deleted immediately after.

The sandbox key is derived from session_id rather than a per-request UUID.
The cleanup and post-loop hooks skip deletion for session-scoped containers.
TTL-based pruning (15 min idle) still applies and refreshes on every use,
so an active session never expires mid-use. The session_id-scoped key is
registered in all_litellm_params and the proxy strip-list so it never
leaks to the upstream LLM provider.

* fix(sandbox): scope session sandbox key to API key identity; add per-identity LRU cap

Two security issues addressed:

1. Cross-user sandbox isolation: the session_id supplied by the client is now
   combined with the server-minted user_api_key_hash to form the cache key
   (format: "{hash}:{session_id}" when authenticated, bare session_id for
   non-proxy use). Two tenants sharing the same session_id no longer share a
   sandbox.

2. Bounded session allocation: each API key identity is capped at
   _SESSION_SCOPED_PER_IDENTITY_CAP (10) live session-scoped containers. When
   a new session is opened beyond the cap, the least-recently-used entry for
   that identity is evicted and its sandbox deleted, preventing unbounded
   accumulation via rotating session IDs.

The container cache tuple gains a fourth element (identity: str | None) so
eviction can filter by identity without parsing key formats. Tests added for
both properties.
2026-06-30 18:58:09 -07:00

59 lines
2 KiB
Bash
Executable file

#!/usr/bin/env bash
# QA: code interpreter sandbox stickiness via metadata.session_id
# bash qa_sticky_session.sh
# LITELLM_BASE_URL=http://localhost:4000 LITELLM_KEY=sk-1234 bash qa_sticky_session.sh
set -euo pipefail
BASE="${LITELLM_BASE_URL:-http://localhost:4000}"
KEY="${LITELLM_KEY:-sk-1234}"
MODEL="${LITELLM_MODEL:-gpt-4o-mini}"
# proxy running at http://localhost:4000 (master key: sk-1234)
SESSION_A="qa-session-$(date +%s)-A"
SESSION_B="qa-session-$(date +%s)-B"
content() {
echo "$1" | python3 -c "import sys,json; d=json.load(sys.stdin); print(d.get('choices',[{}])[0].get('message',{}).get('content','<error>'))"
}
call() {
local session="${1:-}" code="$2" meta=""
[[ -n "$session" ]] && meta=", \"metadata\": {\"session_id\": \"$session\"}"
curl -s -X POST "$BASE/chat/completions" \
-H "Content-Type: application/json" \
-H "Authorization: Bearer $KEY" \
-d "{\"model\":\"$MODEL\"$meta,\"tools\":[{\"type\":\"code_interpreter\"}],\"messages\":[{\"role\":\"user\",\"content\":\"Run this Python code and tell me the result: $code\"}]}"
}
assert_match() {
local label="$1" body="$2" pattern="$3"
if echo "$body" | grep -qiE "$pattern"; then
echo "PASS $label"
else
echo "FAIL $label (expected /$pattern/)"
echo " $(content "$body")"
exit 1
fi
}
echo "=== Sticky Session Sandbox QA ==="
echo "base: $BASE session A: $SESSION_A session B: $SESSION_B"
echo
R=$(call "$SESSION_A" "x = 42; print(x)")
assert_match "same session_id reuses sandbox (set x=42)" "$R" "42"
R=$(call "$SESSION_A" "print(x)")
assert_match "same session_id keeps state (x still 42)" "$R" "42"
R=$(call "$SESSION_B" "print(x)")
assert_match "different session_id is isolated" "$R" "not defined|NameError|undefined|error"
R=$(call "" "y = 99; print(y)")
assert_match "no session_id runs code" "$R" "99"
R=$(call "" "print(y)")
assert_match "no session_id gets fresh sandbox each request" "$R" "not defined|NameError|undefined|error"
echo
echo "All checks passed."