Loosen sandbox compaction model parsing

Author: sdcoffey

src/agents/sandbox/capabilities/compaction.py

@@ -10,16 +10,21 @@
 from .capability import Capability
 
 _DEFAULT_COMPACT_THRESHOLD = 240_000
+_MODEL_NAME_SEPARATOR_TRANSLATION = str.maketrans("", "", ".-")
 
 
-class CompactionModelInfo(BaseModel):
-    context_window: int
+def _model_lookup_key(model: str) -> str:
+    normalized_model = model.strip().lower().removeprefix("openai/")
+    return normalized_model.translate(_MODEL_NAME_SEPARATOR_TRANSLATION)
 
-    @classmethod
-    def for_model(cls, model: str) -> CompactionModelInfo:
-        normalized_model = model.removeprefix("openai/")
 
-        if normalized_model in (
+def _model_context_windows(models: tuple[str, ...], context_window: int) -> dict[str, int]:
+    return {_model_lookup_key(model): context_window for model in models}
+
+
+_MODEL_CONTEXT_WINDOWS: dict[str, int] = {
+    **_model_context_windows(
+        (
             "gpt-5.4",
             "gpt-5.4-2026-03-05",
             "gpt-5.4-pro",
@@ -30,9 +35,11 @@ def for_model(cls, model: str) -> CompactionModelInfo:
             "gpt-4.1-mini-2025-04-14",
             "gpt-4.1-nano",
             "gpt-4.1-nano-2025-04-14",
-        ):
-            return cls(context_window=1_047_576)
-        if normalized_model in (
+        ),
+        1_047_576,
+    ),
+    **_model_context_windows(
+        (
             "gpt-5",
             "gpt-5-2025-08-07",
             "gpt-5-codex",
@@ -57,9 +64,11 @@ def for_model(cls, model: str) -> CompactionModelInfo:
             "gpt-5.4-mini-2026-03-17",
             "gpt-5.4-nano",
             "gpt-5.4-nano-2026-03-17",
-        ):
-            return cls(context_window=400_000)
-        if normalized_model in (
+        ),
+        400_000,
+    ),
+    **_model_context_windows(
+        (
             "codex-mini-latest",
             "o1",
             "o1-2024-12-17",
@@ -77,9 +86,11 @@ def for_model(cls, model: str) -> CompactionModelInfo:
             "o4-mini-2025-04-16",
             "o4-mini-deep-research",
             "o4-mini-deep-research-2025-06-26",
-        ):
-            return cls(context_window=200_000)
-        if normalized_model in (
+        ),
+        200_000,
+    ),
+    **_model_context_windows(
+        (
             "gpt-4o",
             "gpt-4o-2024-05-13",
             "gpt-4o-2024-08-06",
@@ -90,9 +101,27 @@ def for_model(cls, model: str) -> CompactionModelInfo:
             "gpt-5.1-chat-latest",
             "gpt-5.2-chat-latest",
             "gpt-5.3-chat-latest",
-        ):
-            return cls(context_window=128_000)
+        ),
+        128_000,
+    ),
+}
 
+
+class CompactionModelInfo(BaseModel):
+    context_window: int
+
+    @classmethod
+    def maybe_for_model(cls, model: str) -> CompactionModelInfo | None:
+        context_window = _MODEL_CONTEXT_WINDOWS.get(_model_lookup_key(model))
+        if context_window is None:
+            return None
+        return cls(context_window=context_window)
+
+    @classmethod
+    def for_model(cls, model: str) -> CompactionModelInfo:
+        model_info = cls.maybe_for_model(model)
+        if model_info is not None:
+            return model_info
         raise ValueError(f"Unknown context window for model: {model!r}")
 
 
@@ -153,7 +182,11 @@ def sampling_params(self, sampling_params: dict[str, Any]) -> dict[str, Any]:
         if policy is None:
             model = sampling_params.get("model")
             if isinstance(model, str) and model:
-                policy = DynamicCompactionPolicy(model_info=CompactionModelInfo.for_model(model))
+                model_info = CompactionModelInfo.maybe_for_model(model)
+                if model_info is None:
+                    policy = StaticCompactionPolicy()
+                else:
+                    policy = DynamicCompactionPolicy(model_info=model_info)
             else:
                 policy = StaticCompactionPolicy()
 

tests/sandbox/capabilities/test_compaction_capability.py

@@ -26,6 +26,34 @@ def test_sampling_params_uses_static_threshold(self) -> None:
         }
         assert isinstance(capability.policy, StaticCompactionPolicy)
 
+    def test_sampling_params_infers_hyphenated_model_threshold(self) -> None:
+        capability = Compaction()
+
+        sampling_params = capability.sampling_params({"model": "gpt-5-2"})
+
+        assert sampling_params == {
+            "context_management": [
+                {
+                    "type": "compaction",
+                    "compact_threshold": 360_000,
+                }
+            ]
+        }
+
+    def test_sampling_params_falls_back_for_unknown_model(self) -> None:
+        capability = Compaction()
+
+        sampling_params = capability.sampling_params({"model": "azure-prod-deployment"})
+
+        assert sampling_params == {
+            "context_management": [
+                {
+                    "type": "compaction",
+                    "compact_threshold": 240_000,
+                }
+            ]
+        }
+
     def test_process_context_keeps_items_from_last_compaction(self) -> None:
         """Tests compaction truncates history to the last compaction item, inclusive."""
 

tests/sandbox/test_compaction.py

@@ -14,6 +14,10 @@
         ("o3", 200_000),
         ("gpt-4o", 128_000),
         ("openai/gpt-5.4", 1_047_576),
+        ("gpt-5-2", 400_000),
+        ("gpt-5-4", 1_047_576),
+        ("openai/gpt-5-4-mini", 400_000),
+        ("gpt-4-1-mini", 1_047_576),
     ],
 )
 def test_compaction_model_info_for_model_returns_context_window(
@@ -26,3 +30,7 @@ def test_compaction_model_info_for_model_returns_context_window(
 def test_compaction_model_info_for_model_rejects_unknown_model() -> None:
     with pytest.raises(ValueError, match="Unknown context window for model"):
         CompactionModelInfo.for_model("not-a-model")
+
+
+def test_compaction_model_info_maybe_for_model_returns_none_for_unknown_model() -> None:
+    assert CompactionModelInfo.maybe_for_model("not-a-model") is None