fix(prompt): isolate truncation warnings per context

Follow-up to salvaged PR #41619: replace the module-global
_truncation_warnings list with a contextvars.ContextVar so concurrent
gateway-session prompt builds can't drain or clear each other's pending
warnings (cross-session leak). Adds a context-isolation test.
This commit is contained in:
teknium
2026-06-16 11:28:35 -07:00
committed by Teknium
parent f6a42b1acf
commit 6ebc449915
2 changed files with 53 additions and 6 deletions
+25 -6
View File
@@ -8,6 +8,7 @@ import json
import logging
import os
import threading
import contextvars
from collections import OrderedDict
from pathlib import Path
@@ -976,14 +977,32 @@ def _get_context_file_max_chars() -> int:
return CONTEXT_FILE_MAX_CHARS
# Collect truncation warnings so the caller (run_agent) can surface them.
_truncation_warnings: list = []
# A ContextVar (not a module-global list) isolates accumulation per thread /
# per async task, so concurrent gateway-session prompt builds can't drain or
# clear each other's pending warnings (cross-session leak). Each build runs in
# its own context, collects its own warnings, and drains them synchronously.
_truncation_warnings: "contextvars.ContextVar[Optional[list]]" = contextvars.ContextVar(
"context_file_truncation_warnings", default=None
)
def _record_truncation_warning(msg: str) -> None:
"""Append a truncation warning to the current context's accumulator."""
warnings = _truncation_warnings.get()
if warnings is None:
warnings = []
_truncation_warnings.set(warnings)
warnings.append(msg)
def drain_truncation_warnings() -> list:
"""Return and clear any truncation warnings accumulated since last drain."""
warnings = _truncation_warnings.copy()
_truncation_warnings.clear()
return warnings
"""Return and clear any truncation warnings accumulated in this context."""
warnings = _truncation_warnings.get()
if not warnings:
return []
drained = list(warnings)
warnings.clear()
return drained
# =========================================================================
@@ -1503,7 +1522,7 @@ def _truncate_content(content: str, filename: str, max_chars: Optional[int] = No
f"increase context_file_max_chars or trim the file!"
)
logger.warning(msg)
_truncation_warnings.append(msg)
_record_truncation_warning(msg)
head_chars = int(max_chars * CONTEXT_TRUNCATE_HEAD_RATIO)
tail_chars = int(max_chars * CONTEXT_TRUNCATE_TAIL_RATIO)
head = content[:head_chars]