fix(python): preserve max_output_tokens in ModelLimits

python/copilot/client.py

@@ -1010,18 +1010,21 @@ class ModelLimits:
     """Model limits"""
 
     max_prompt_tokens: int | None = None
+    max_output_tokens: int | None = field(default=None, kw_only=True)
     max_context_window_tokens: int | None = None
     vision: ModelVisionLimits | None = None
 
     @staticmethod
     def from_dict(obj: Any) -> ModelLimits:
         assert isinstance(obj, dict)
         max_prompt_tokens = obj.get("max_prompt_tokens")
+        max_output_tokens = obj.get("max_output_tokens")
         max_context_window_tokens = obj.get("max_context_window_tokens")
         vision_dict = obj.get("vision")
         vision = ModelVisionLimits.from_dict(vision_dict) if vision_dict else None
         return ModelLimits(
             max_prompt_tokens=max_prompt_tokens,
+            max_output_tokens=max_output_tokens,
             max_context_window_tokens=max_context_window_tokens,
             vision=vision,
         )
@@ -1030,6 +1033,8 @@ def to_dict(self) -> dict:
         result: dict = {}
         if self.max_prompt_tokens is not None:
             result["max_prompt_tokens"] = self.max_prompt_tokens
+        if self.max_output_tokens is not None:
+            result["max_output_tokens"] = self.max_output_tokens
         if self.max_context_window_tokens is not None:
             result["max_context_window_tokens"] = self.max_context_window_tokens
         if self.vision is not None:

python/test_client.py

@@ -39,6 +39,7 @@
     ModelInfo,
     ModelLimits,
     ModelSupports,
+    ModelVisionLimits,
 )
 from copilot.generated.rpc import AutoTier as AutoTierEnum
 from copilot.generated.rpc import (
@@ -2068,6 +2069,30 @@ def test_long_context_empty_object_round_trip(self):
         assert billing.to_dict() == {"tokenPrices": {"longContext": {}}}
 
 
+class TestModelLimits:
+    def test_max_output_tokens_round_trip(self):
+        wire = {
+            "max_prompt_tokens": 128000,
+            "max_output_tokens": 16384,
+            "max_context_window_tokens": 144384,
+        }
+
+        limits = ModelLimits.from_dict(wire)
+
+        assert limits.max_output_tokens == 16384
+        assert limits.to_dict() == wire
+
+    def test_max_output_tokens_preserves_positional_constructor(self):
+        vision = ModelVisionLimits()
+
+        limits = ModelLimits(128000, 144384, vision)
+
+        assert limits.max_prompt_tokens == 128000
+        assert limits.max_context_window_tokens == 144384
+        assert limits.vision is vision
+        assert limits.max_output_tokens is None
+
+
 class TestOnListModels:
     @pytest.mark.asyncio
     async def test_list_models_with_custom_handler(self):