From 478eb315df70d7efbc403c90ccac7f82b204222b Mon Sep 17 00:00:00 2001 From: Jason Robert Date: Mon, 4 May 2026 13:03:40 -0400 Subject: [PATCH 1/2] feat(engine): source context_window_max from provider SDKs at runtime MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The dashboard's "context window remaining" bar (and any future enforcement) reads context_window_max emitted in agent_started, agent_completed, and parallel_agent_completed events. Until now this value came from a hand-maintained context_window field on ModelPricing entries — but that field reported the *theoretical model max*, not the SDK's actual max_prompt_tokens. Concrete impact (from the GitHub community discussion #186340 raw API response on 2026-02-04): Model context_window max_prompt Bar wrong by gpt-5.x variants 400K 128K 3.1x claude-opus-4.5 200K 128K 1.6x claude-opus-4.6 1M 200K 5.0x This change replaces the static lookup with a per-provider runtime query so the bar always reflects the SDK's actual cap. - AgentProvider.get_max_prompt_tokens(model) -> int | None Concrete default returns None; Copilot and Claude providers override to query their SDKs (cached). - WorkflowEngine resolves the provider via the same path as execution (single-provider or registry) and prefers output.model over agent.model when an output is available. - Failures (SDK error, missing client, misconfigured mock) are swallowed and return None; metadata is best-effort and must never block workflow execution. Cleanup of dead code per scope: - Removed context_window field from ModelPricing and from all ~30 entries in DEFAULT_PRICING. - Updated fuzzy-match warning text to drop "context_window". - Deleted tests/test_providers/test_context_window.py (the entire file tested the static-table lookup that no longer exists). - New TestGetMaxPromptTokens classes in test_copilot.py and test_claude.py covering known model, unknown model, SDK failure, cache-after-first-call, and mock-mode fallback. - Updated test_context_window_events.py to inject expected values via a fake get_max_prompt_tokens (mock-handler mode no longer has a static table to fall back to). Added two new resolution-order tests covering default-model fallback and output.model preference. Pricing data ($/Mtok) is still hand-maintained — neither SDK exposes per-token dollar amounts, so DEFAULT_PRICING still needs entries for new models for cost math. Only the context-window field becomes self-correcting. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- src/conductor/engine/pricing.py | 64 +++------- src/conductor/engine/workflow.py | 67 ++++++++-- src/conductor/providers/base.py | 30 +++++ src/conductor/providers/claude.py | 37 ++++++ src/conductor/providers/copilot.py | 23 ++++ .../test_engine/test_context_window_events.py | 120 +++++++++++++++++- tests/test_providers/test_claude.py | 71 +++++++++++ tests/test_providers/test_context_window.py | 86 ------------- tests/test_providers/test_copilot.py | 65 ++++++++++ 9 files changed, 414 insertions(+), 149 deletions(-) delete mode 100644 tests/test_providers/test_context_window.py diff --git a/src/conductor/engine/pricing.py b/src/conductor/engine/pricing.py index 4e6787ba..ac1dd5fb 100644 --- a/src/conductor/engine/pricing.py +++ b/src/conductor/engine/pricing.py @@ -30,7 +30,7 @@ def _warn_fuzzy_match(requested: str, matched_key: str, strategy: str) -> None: _FUZZY_MATCH_WARNED.add(requested) logger.warning( "Pricing for model %r resolved via %s fallback to %r. " - "Cost and context_window metadata may be inaccurate. " + "Cost calculation may be inaccurate. " "Add %r to DEFAULT_PRICING or pass an override to silence this warning.", requested, strategy, @@ -41,66 +41,59 @@ def _warn_fuzzy_match(requested: str, matched_key: str, strategy: str) -> None: @dataclass(frozen=True) class ModelPricing: - """Pricing and metadata per model. + """Pricing per model. Attributes: input_per_mtok: Cost per million input tokens (USD). output_per_mtok: Cost per million output tokens (USD). cache_read_per_mtok: Cost per million cache read tokens (USD). cache_write_per_mtok: Cost per million cache write tokens (USD). - context_window: Context window size in tokens, or None if unknown. """ input_per_mtok: float output_per_mtok: float cache_read_per_mtok: float = 0.0 cache_write_per_mtok: float = 0.0 - context_window: int | None = None -# Default model table (pricing + context window metadata) -# Sources: OpenAI pricing page, Anthropic pricing page, provider docs +# Default model table (pricing only). +# Context-window metadata is sourced from each provider's SDK at runtime via +# ``AgentProvider.get_max_prompt_tokens()`` — see ``providers/base.py``. +# Sources: OpenAI pricing page, Anthropic pricing page, provider docs. DEFAULT_PRICING: dict[str, ModelPricing] = { # OpenAI / Copilot models - "gpt-4-turbo": ModelPricing( - input_per_mtok=10.00, output_per_mtok=30.00, context_window=128_000 - ), - "gpt-4o": ModelPricing(input_per_mtok=2.50, output_per_mtok=10.00, context_window=128_000), - "gpt-4o-mini": ModelPricing(input_per_mtok=0.15, output_per_mtok=0.60, context_window=128_000), - "gpt-4.1": ModelPricing(input_per_mtok=2.00, output_per_mtok=8.00, context_window=1_047_576), - "gpt-4.1-mini": ModelPricing( - input_per_mtok=0.15, output_per_mtok=0.60, context_window=1_047_576 - ), - "gpt-4": ModelPricing(input_per_mtok=30.00, output_per_mtok=60.00, context_window=8_192), - "gpt-3.5-turbo": ModelPricing(input_per_mtok=0.50, output_per_mtok=1.50, context_window=16_385), - "gpt-5.2": ModelPricing(input_per_mtok=2.00, output_per_mtok=8.00, context_window=400_000), - "gpt-5.1": ModelPricing(input_per_mtok=2.00, output_per_mtok=8.00, context_window=400_000), + "gpt-4-turbo": ModelPricing(input_per_mtok=10.00, output_per_mtok=30.00), + "gpt-4o": ModelPricing(input_per_mtok=2.50, output_per_mtok=10.00), + "gpt-4o-mini": ModelPricing(input_per_mtok=0.15, output_per_mtok=0.60), + "gpt-4.1": ModelPricing(input_per_mtok=2.00, output_per_mtok=8.00), + "gpt-4.1-mini": ModelPricing(input_per_mtok=0.15, output_per_mtok=0.60), + "gpt-4": ModelPricing(input_per_mtok=30.00, output_per_mtok=60.00), + "gpt-3.5-turbo": ModelPricing(input_per_mtok=0.50, output_per_mtok=1.50), + "gpt-5.2": ModelPricing(input_per_mtok=2.00, output_per_mtok=8.00), + "gpt-5.1": ModelPricing(input_per_mtok=2.00, output_per_mtok=8.00), # O-series - "o1": ModelPricing(input_per_mtok=15.00, output_per_mtok=60.00, context_window=200_000), - "o1-mini": ModelPricing(input_per_mtok=3.00, output_per_mtok=12.00, context_window=128_000), - "o1-preview": ModelPricing(input_per_mtok=15.00, output_per_mtok=60.00, context_window=128_000), - "o3-mini": ModelPricing(input_per_mtok=1.10, output_per_mtok=4.40, context_window=200_000), + "o1": ModelPricing(input_per_mtok=15.00, output_per_mtok=60.00), + "o1-mini": ModelPricing(input_per_mtok=3.00, output_per_mtok=12.00), + "o1-preview": ModelPricing(input_per_mtok=15.00, output_per_mtok=60.00), + "o3-mini": ModelPricing(input_per_mtok=1.10, output_per_mtok=4.40), # Claude 4.5 Series (newest) "claude-opus-4-5": ModelPricing( input_per_mtok=5.00, output_per_mtok=25.00, cache_read_per_mtok=0.50, cache_write_per_mtok=6.25, - context_window=200_000, ), "claude-sonnet-4-5": ModelPricing( input_per_mtok=3.00, output_per_mtok=15.00, cache_read_per_mtok=0.30, cache_write_per_mtok=3.75, - context_window=200_000, ), "claude-haiku-4-5": ModelPricing( input_per_mtok=1.00, output_per_mtok=5.00, cache_read_per_mtok=0.10, cache_write_per_mtok=1.25, - context_window=200_000, ), # Short aliases for Claude 4.5 Series (used in workflow files) "opus-4.5": ModelPricing( @@ -108,21 +101,18 @@ class ModelPricing: output_per_mtok=25.00, cache_read_per_mtok=0.50, cache_write_per_mtok=6.25, - context_window=200_000, ), "sonnet-4.5": ModelPricing( input_per_mtok=3.00, output_per_mtok=15.00, cache_read_per_mtok=0.30, cache_write_per_mtok=3.75, - context_window=200_000, ), "haiku-4.5": ModelPricing( input_per_mtok=1.00, output_per_mtok=5.00, cache_read_per_mtok=0.10, cache_write_per_mtok=1.25, - context_window=200_000, ), # Claude 4.6 Series "claude-opus-4.6": ModelPricing( @@ -130,21 +120,18 @@ class ModelPricing: output_per_mtok=25.00, cache_read_per_mtok=0.50, cache_write_per_mtok=6.25, - context_window=1_000_000, ), "claude-opus-4.6-1m": ModelPricing( input_per_mtok=5.00, output_per_mtok=25.00, cache_read_per_mtok=0.50, cache_write_per_mtok=6.25, - context_window=1_000_000, ), "claude-sonnet-4.6": ModelPricing( input_per_mtok=3.00, output_per_mtok=15.00, cache_read_per_mtok=0.30, cache_write_per_mtok=3.75, - context_window=1_000_000, ), # Claude 4 Series "claude-opus-4": ModelPricing( @@ -152,21 +139,18 @@ class ModelPricing: output_per_mtok=75.00, cache_read_per_mtok=1.50, cache_write_per_mtok=18.75, - context_window=200_000, ), "claude-sonnet-4": ModelPricing( input_per_mtok=3.00, output_per_mtok=15.00, cache_read_per_mtok=0.30, cache_write_per_mtok=3.75, - context_window=200_000, ), "claude-haiku-4": ModelPricing( input_per_mtok=0.25, output_per_mtok=1.25, cache_read_per_mtok=0.03, cache_write_per_mtok=0.30, - context_window=200_000, ), # Claude 3.x Series "claude-3-7-sonnet": ModelPricing( @@ -174,69 +158,59 @@ class ModelPricing: output_per_mtok=15.00, cache_read_per_mtok=0.30, cache_write_per_mtok=3.75, - context_window=200_000, ), "claude-3.7-sonnet": ModelPricing( input_per_mtok=3.00, output_per_mtok=15.00, cache_read_per_mtok=0.30, cache_write_per_mtok=3.75, - context_window=200_000, ), "claude-3-5-sonnet": ModelPricing( input_per_mtok=3.00, output_per_mtok=15.00, cache_read_per_mtok=0.30, cache_write_per_mtok=3.75, - context_window=200_000, ), "claude-3.5-sonnet": ModelPricing( input_per_mtok=3.00, output_per_mtok=15.00, cache_read_per_mtok=0.30, cache_write_per_mtok=3.75, - context_window=200_000, ), "claude-3-5-haiku": ModelPricing( input_per_mtok=0.80, output_per_mtok=4.00, cache_read_per_mtok=0.08, cache_write_per_mtok=1.00, - context_window=200_000, ), "claude-3.5-haiku": ModelPricing( input_per_mtok=0.80, output_per_mtok=4.00, cache_read_per_mtok=0.08, cache_write_per_mtok=1.00, - context_window=200_000, ), "claude-3-opus": ModelPricing( input_per_mtok=15.00, output_per_mtok=75.00, cache_read_per_mtok=1.50, cache_write_per_mtok=18.75, - context_window=200_000, ), "claude-3-sonnet": ModelPricing( input_per_mtok=3.00, output_per_mtok=15.00, cache_read_per_mtok=0.30, cache_write_per_mtok=3.75, - context_window=200_000, ), "claude-3-haiku": ModelPricing( input_per_mtok=0.25, output_per_mtok=1.25, cache_read_per_mtok=0.03, cache_write_per_mtok=0.30, - context_window=200_000, ), # Gemini "gemini-3.1-pro-preview": ModelPricing( input_per_mtok=1.25, output_per_mtok=5.00, - context_window=1_000_000, ), } diff --git a/src/conductor/engine/workflow.py b/src/conductor/engine/workflow.py index 1037b054..17d19f51 100644 --- a/src/conductor/engine/workflow.py +++ b/src/conductor/engine/workflow.py @@ -766,16 +766,55 @@ async def _execute_subworkflow_with_inputs( usage = child_engine.usage_tracker.get_summary() return output, usage - def _get_context_window_for_agent(self, agent: AgentDef) -> int | None: - """Return the context window size for an agent's model.""" - from conductor.engine.pricing import get_pricing - - model = agent.model + async def _get_provider_for_agent(self, agent: AgentDef) -> AgentProvider | None: + """Resolve the provider that will (or did) execute ``agent``. + + Mirrors the executor-resolution logic in ``_get_executor_for_agent`` + so context-window metadata lookups go through the same provider that + handles execution. Returns ``None`` only when no provider can be + determined (e.g. transient registry failures); callers must treat + ``None`` as "metadata unavailable". + """ + if self._registry is not None: + try: + return await self._registry.get_provider(agent) + except Exception as e: + logger.debug("Provider lookup via registry failed for %s: %s", agent.name, e) + return None + return self._single_provider + + async def _get_context_window_for_agent( + self, agent: AgentDef, output: AgentOutput | None = None + ) -> int | None: + """Return the SDK-reported max prompt tokens for an agent. + + Resolves the model in priority order: the model the SDK actually used + (``output.model``), then the agent's configured model, then the + workflow's runtime default. Returns ``None`` when no model can be + resolved, no provider can be reached, or the provider's metadata call + fails — context-window metadata is best-effort and must never break + workflow execution. + """ + model = ( + (output.model if output is not None else None) + or agent.model + or self.config.workflow.runtime.default_model + ) if not model: return None - - pricing = get_pricing(model) - return pricing.context_window if pricing else None + provider = await self._get_provider_for_agent(agent) + if provider is None: + return None + try: + return await provider.get_max_prompt_tokens(model) + except Exception as e: + logger.debug( + "get_max_prompt_tokens(%r) raised on provider for agent %s: %s", + model, + agent.name, + e, + ) + return None async def run(self, inputs: dict[str, Any]) -> dict[str, Any]: """Execute the workflow from entry_point to $end. @@ -1520,7 +1559,9 @@ async def _execute_loop(self, current_agent_name: str) -> dict[str, Any]: "agent_name": agent.name, "iteration": agent_execution_count, "agent_type": agent.type or "agent", - "context_window_max": self._get_context_window_for_agent(agent), + "context_window_max": await self._get_context_window_for_agent( + agent + ), }, ) @@ -1844,7 +1885,9 @@ async def _execute_loop(self, current_agent_name: str) -> dict[str, Any]: "output": output.content, "output_keys": output_keys, "context_window_used": output.input_tokens, - "context_window_max": self._get_context_window_for_agent(agent), + "context_window_max": await self._get_context_window_for_agent( + agent, output + ), }, ) @@ -2472,7 +2515,9 @@ async def execute_single_agent(agent: AgentDef) -> tuple[str, Any]: "tokens": output.tokens_used, "cost_usd": usage.cost_usd, "context_window_used": output.input_tokens, - "context_window_max": self._get_context_window_for_agent(agent), + "context_window_max": await self._get_context_window_for_agent( + agent, output + ), }, ) diff --git a/src/conductor/providers/base.py b/src/conductor/providers/base.py index 20d8c915..4aa4e963 100644 --- a/src/conductor/providers/base.py +++ b/src/conductor/providers/base.py @@ -144,3 +144,33 @@ async def close(self) -> None: such as HTTP clients or session state. """ ... + + async def get_max_prompt_tokens(self, model: str) -> int | None: + """Return the SDK-reported maximum input (prompt) tokens for ``model``. + + This is the authoritative cap on prompt size enforced by the underlying + SDK or backend (e.g. the Copilot ``max_prompt_tokens`` field, or the + Anthropic ``max_input_tokens`` field). It is typically lower than the + model's theoretical context window — for example, the Copilot SDK + currently caps most GPT-5 variants at 128K despite a 400K model max. + + Implementations should: + + * Query their SDK's model-listing endpoint (cached after the first call). + * Return ``None`` when the model is unknown to the provider, when the + SDK call fails, or when no metadata is available. + * Never raise — context-window metadata is best-effort and must not + interrupt workflow execution. + + The default implementation returns ``None``, which causes the + dashboard's context-window bar to be hidden and any future enforcement + to be skipped — both safe degradations. + + Args: + model: The model identifier as it would be sent to the SDK + (e.g. ``"gpt-5.2"``, ``"claude-sonnet-4-5-20250929"``). + + Returns: + The maximum prompt (input) tokens the SDK will accept, or ``None``. + """ + return None diff --git a/src/conductor/providers/claude.py b/src/conductor/providers/claude.py index 4f8af2bf..d0a0bc56 100644 --- a/src/conductor/providers/claude.py +++ b/src/conductor/providers/claude.py @@ -178,6 +178,13 @@ def __init__( self._mcp_servers_config = mcp_servers self._mcp_manager: MCPManager | None = None + # Cache of model_id -> max_input_tokens populated lazily on first + # get_max_prompt_tokens() call. Guarded by an asyncio.Lock to avoid + # racing concurrent first-callers and emitting duplicate models.list() + # requests. + self._max_input_cache: dict[str, int | None] | None = None + self._max_input_cache_lock = asyncio.Lock() + # Initialize the client (sync initialization) self._initialize_client() @@ -312,6 +319,36 @@ async def _log_available_models(self) -> None: except Exception as e: logger.warning(f"Could not list available models (discovery failed): {e}") + async def get_max_prompt_tokens(self, model: str) -> int | None: + """Return the Anthropic SDK's ``max_input_tokens`` for ``model``. + + Lazily populates a per-instance cache on first call by enumerating + ``client.models.list()``. Subsequent calls are dictionary lookups. + Returns ``None`` if the SDK is unavailable, the client failed to + initialize, the model isn't listed, or the listing call raises — + context-window metadata must never block workflow execution. + + Note: the value reflects the API's *default* input window. Claude + models with a 1M-context beta require an explicit beta header, which + Conductor does not set today; for those models the API still reports + the default window. + """ + if not ANTHROPIC_SDK_AVAILABLE or self._client is None: + return None + async with self._max_input_cache_lock: + if self._max_input_cache is None: + cache: dict[str, int | None] = {} + try: + page = await self._client.models.list() + except Exception as e: + logger.debug("Failed to list Anthropic models: %s", e) + self._max_input_cache = cache # cache the empty result + return None + for info in page.data: + cache[info.id] = getattr(info, "max_input_tokens", None) + self._max_input_cache = cache + return self._max_input_cache.get(model) + async def _ensure_mcp_connected(self) -> None: """Connect to MCP servers if configured. diff --git a/src/conductor/providers/copilot.py b/src/conductor/providers/copilot.py index 94878040..311eed0f 100644 --- a/src/conductor/providers/copilot.py +++ b/src/conductor/providers/copilot.py @@ -1775,6 +1775,29 @@ async def close(self) -> None: self._call_history.clear() self._retry_history.clear() + async def get_max_prompt_tokens(self, model: str) -> int | None: + """Return the Copilot SDK's ``max_prompt_tokens`` for ``model``. + + Queries ``client.list_models()`` (cached internally by the SDK) and + returns ``capabilities.limits.max_prompt_tokens`` for the matching + model. Returns ``None`` in mock-handler mode, when the SDK is + unavailable, when the model is unknown, or when any SDK call fails — + context-window metadata must never block workflow execution. + """ + if self._mock_handler is not None or not COPILOT_SDK_AVAILABLE: + return None + try: + await self._ensure_client_started() + models = await self._client.list_models() + except Exception as e: + logger.debug("Failed to list Copilot models for %r: %s", model, e) + return None + for info in models: + if info.id == model: + limits = getattr(info.capabilities, "limits", None) + return getattr(limits, "max_prompt_tokens", None) if limits else None + return None + def get_session_ids(self) -> dict[str, str]: """Get tracked session IDs for all executed agents. diff --git a/tests/test_engine/test_context_window_events.py b/tests/test_engine/test_context_window_events.py index c24e3562..1c581da4 100644 --- a/tests/test_engine/test_context_window_events.py +++ b/tests/test_engine/test_context_window_events.py @@ -1,4 +1,10 @@ -"""Tests that workflow events include context_window fields.""" +"""Tests that workflow events include context_window fields. + +Context-window metadata is now sourced from each provider's SDK at runtime +(``AgentProvider.get_max_prompt_tokens``). In mock-handler mode the Copilot +provider has no SDK to query and returns ``None`` by default, so these tests +monkeypatch the provider method to inject the values being asserted. +""" from __future__ import annotations @@ -45,6 +51,18 @@ def _make_emitter_and_collector() -> tuple[WorkflowEventEmitter, EventCollector] return emitter, collector +def _provider_with_max_prompt(values: dict[str, int | None]) -> CopilotProvider: + """Build a mock-handler Copilot provider whose ``get_max_prompt_tokens`` + returns values from ``values`` (or ``None`` for unknown models).""" + provider = CopilotProvider(mock_handler=lambda a, p, c: {"answer": "hi", "result": a.name}) + + async def fake_get_max_prompt_tokens(model: str) -> int | None: + return values.get(model) + + provider.get_max_prompt_tokens = fake_get_max_prompt_tokens # type: ignore[method-assign] + return provider + + class TestAgentStartedContextWindow: """agent_started event includes context_window_max.""" @@ -70,7 +88,7 @@ async def test_agent_started_has_context_window_max(self) -> None: ], output={"answer": "{{ a1.output.answer }}"}, ) - provider = CopilotProvider(mock_handler=lambda a, p, c: {"answer": "hi"}) + provider = _provider_with_max_prompt({"gpt-4o": 128000}) engine = WorkflowEngine(config, provider, event_emitter=emitter) await engine.run({}) @@ -104,7 +122,7 @@ async def test_agent_completed_has_context_window_fields(self) -> None: ], output={"answer": "{{ a1.output.answer }}"}, ) - provider = CopilotProvider(mock_handler=lambda a, p, c: {"answer": "hi"}) + provider = _provider_with_max_prompt({"gpt-4o": 128000}) engine = WorkflowEngine(config, provider, event_emitter=emitter) await engine.run({}) @@ -115,7 +133,7 @@ async def test_agent_completed_has_context_window_fields(self) -> None: class TestContextWindowNoneForUnknownModel: - """context_window_max is None when model is unknown.""" + """context_window_max is None when the provider has no metadata for the model.""" @pytest.mark.asyncio async def test_unknown_model_returns_none(self) -> None: @@ -139,7 +157,8 @@ async def test_unknown_model_returns_none(self) -> None: ], output={"answer": "{{ a1.output.answer }}"}, ) - provider = CopilotProvider(mock_handler=lambda a, p, c: {"answer": "hi"}) + # Empty metadata table — every lookup returns None. + provider = _provider_with_max_prompt({}) engine = WorkflowEngine(config, provider, event_emitter=emitter) await engine.run({}) @@ -184,7 +203,7 @@ async def test_parallel_agent_completed_has_context_window_fields(self) -> None: ], output={"result": "done"}, ) - provider = CopilotProvider(mock_handler=lambda a, p, c: {"result": a.name}) + provider = _provider_with_max_prompt({"gpt-4o": 128000}) engine = WorkflowEngine(config, provider, event_emitter=emitter) await engine.run({}) @@ -229,10 +248,97 @@ async def test_parallel_agent_unknown_model_context_window_none(self) -> None: ], output={"result": "done"}, ) - provider = CopilotProvider(mock_handler=lambda a, p, c: {"result": "ok"}) + provider = _provider_with_max_prompt({}) engine = WorkflowEngine(config, provider, event_emitter=emitter) await engine.run({}) events = collector.of_type("parallel_agent_completed") for event in events: assert event.data["context_window_max"] is None + + +class TestContextWindowResolutionOrder: + """The model is resolved from output.model first, then agent.model, then default.""" + + @pytest.mark.asyncio + async def test_default_model_used_when_agent_has_no_model(self) -> None: + emitter, collector = _make_emitter_and_collector() + config = WorkflowConfig( + workflow=WorkflowDef( + name="test-default", + entry_point="a1", + runtime=RuntimeConfig(provider="copilot", default_model="gpt-4o"), + context=ContextConfig(mode="accumulate"), + limits=LimitsConfig(max_iterations=10), + ), + agents=[ + AgentDef( + name="a1", + prompt="Hello", + output={"answer": OutputField(type="string")}, + routes=[RouteDef(to="$end")], + ), + ], + output={"answer": "{{ a1.output.answer }}"}, + ) + provider = _provider_with_max_prompt({"gpt-4o": 128000}) + engine = WorkflowEngine(config, provider, event_emitter=emitter) + await engine.run({}) + + # agent_started has no output yet, but should resolve via default_model. + assert collector.first("agent_started").data["context_window_max"] == 128000 + + @pytest.mark.asyncio + async def test_output_model_preferred_over_configured(self) -> None: + emitter, collector = _make_emitter_and_collector() + config = WorkflowConfig( + workflow=WorkflowDef( + name="test-output-model", + entry_point="a1", + runtime=RuntimeConfig(provider="copilot"), + context=ContextConfig(mode="accumulate"), + limits=LimitsConfig(max_iterations=10), + ), + agents=[ + AgentDef( + name="a1", + model="gpt-4o", + prompt="Hello", + output={"answer": OutputField(type="string")}, + routes=[RouteDef(to="$end")], + ), + ], + output={"answer": "{{ a1.output.answer }}"}, + ) + + # Mock handler reports a different model than the agent requested + # (e.g. the SDK aliased or substituted it). + def handler(agent, prompt, context): # type: ignore[no-untyped-def] + return {"answer": "hi"} + + provider = CopilotProvider(mock_handler=handler) + + async def fake_get_max_prompt_tokens(model: str) -> int | None: + return {"gpt-4o": 128000, "gpt-5.2": 400000}.get(model) + + provider.get_max_prompt_tokens = fake_get_max_prompt_tokens # type: ignore[method-assign] + + # Force the AgentOutput.model field via a wrapper. The simplest hook + # here is to set the model on the mock-execute return — done by + # patching the provider's execute to override output.model. + original_execute = provider.execute + + async def execute_with_model(*args, **kwargs): # type: ignore[no-untyped-def] + output = await original_execute(*args, **kwargs) + output.model = "gpt-5.2" + return output + + provider.execute = execute_with_model # type: ignore[method-assign] + + engine = WorkflowEngine(config, provider, event_emitter=emitter) + await engine.run({}) + + # agent_started runs before execution, no output yet — uses agent.model + assert collector.first("agent_started").data["context_window_max"] == 128000 + # agent_completed has output.model — uses that + assert collector.first("agent_completed").data["context_window_max"] == 400000 diff --git a/tests/test_providers/test_claude.py b/tests/test_providers/test_claude.py index bab3fabe..147f7979 100644 --- a/tests/test_providers/test_claude.py +++ b/tests/test_providers/test_claude.py @@ -2302,3 +2302,74 @@ def __init__(self) -> None: # Verify retry-after header was used (delay should be 5.0) assert len(provider._retry_history) == 1 assert provider._retry_history[0]["delay"] == 5.0 + + +@patch("conductor.providers.claude.ANTHROPIC_SDK_AVAILABLE", True) +@patch("conductor.providers.claude.AsyncAnthropic") +class TestClaudeGetMaxPromptTokens: + """Tests for ClaudeProvider.get_max_prompt_tokens.""" + + @pytest.mark.asyncio + async def test_returns_max_input_tokens_for_known_model( + self, mock_anthropic_class: Mock + ) -> None: + mock_client = Mock() + mock_client.models.list = AsyncMock( + return_value=Mock( + data=[ + Mock(id="claude-sonnet-4-5", max_input_tokens=200_000), + Mock(id="claude-opus-4-5", max_input_tokens=200_000), + ] + ) + ) + mock_anthropic_class.return_value = mock_client + + provider = ClaudeProvider() + assert await provider.get_max_prompt_tokens("claude-sonnet-4-5") == 200_000 + + @pytest.mark.asyncio + async def test_returns_none_for_unknown_model(self, mock_anthropic_class: Mock) -> None: + mock_client = Mock() + mock_client.models.list = AsyncMock(return_value=Mock(data=[])) + mock_anthropic_class.return_value = mock_client + + provider = ClaudeProvider() + assert await provider.get_max_prompt_tokens("unknown-x") is None + + @pytest.mark.asyncio + async def test_sdk_failure_returns_none(self, mock_anthropic_class: Mock) -> None: + """An exception from models.list() must not propagate.""" + mock_client = Mock() + mock_client.models.list = AsyncMock(side_effect=RuntimeError("network down")) + mock_anthropic_class.return_value = mock_client + + provider = ClaudeProvider() + assert await provider.get_max_prompt_tokens("claude-sonnet-4-5") is None + + @pytest.mark.asyncio + async def test_caches_after_first_call(self, mock_anthropic_class: Mock) -> None: + """Second call must hit the cache, not the SDK.""" + mock_client = Mock() + mock_client.models.list = AsyncMock( + return_value=Mock(data=[Mock(id="claude-sonnet-4-5", max_input_tokens=200_000)]) + ) + mock_anthropic_class.return_value = mock_client + + provider = ClaudeProvider() + await provider.get_max_prompt_tokens("claude-sonnet-4-5") + await provider.get_max_prompt_tokens("claude-sonnet-4-5") + await provider.get_max_prompt_tokens("anything-else") + + assert mock_client.models.list.await_count == 1 + + @pytest.mark.asyncio + @patch("conductor.providers.claude.ANTHROPIC_SDK_AVAILABLE", False) + async def test_returns_none_when_sdk_unavailable(self, mock_anthropic_class: Mock) -> None: + # Need a workaround: ANTHROPIC_SDK_AVAILABLE is False so __init__ + # raises. Build an instance bypassing the init guard by patching + # only at call time. + with patch("conductor.providers.claude.ANTHROPIC_SDK_AVAILABLE", True): + provider = ClaudeProvider() + + with patch("conductor.providers.claude.ANTHROPIC_SDK_AVAILABLE", False): + assert await provider.get_max_prompt_tokens("claude-sonnet-4-5") is None diff --git a/tests/test_providers/test_context_window.py b/tests/test_providers/test_context_window.py deleted file mode 100644 index cc2b2d6f..00000000 --- a/tests/test_providers/test_context_window.py +++ /dev/null @@ -1,86 +0,0 @@ -"""Tests for context window lookups via the unified pricing registry.""" - -from __future__ import annotations - -from conductor.engine.pricing import DEFAULT_PRICING, get_pricing - - -def _context_window(model: str) -> int | None: - """Helper: look up context window via the pricing registry.""" - pricing = get_pricing(model) - return pricing.context_window if pricing else None - - -class TestExactMatch: - """Exact model name lookups.""" - - def test_claude_sonnet_4(self) -> None: - assert _context_window("claude-sonnet-4") == 200_000 - - def test_claude_opus_4_6_1m(self) -> None: - assert _context_window("claude-opus-4.6-1m") == 1_000_000 - - def test_gpt_4o(self) -> None: - assert _context_window("gpt-4o") == 128_000 - - def test_gpt_4_legacy(self) -> None: - assert _context_window("gpt-4") == 8_192 - - def test_gpt_4_1(self) -> None: - assert _context_window("gpt-4.1") == 1_047_576 - - def test_short_alias(self) -> None: - assert _context_window("sonnet-4.5") == 200_000 - - def test_gemini(self) -> None: - assert _context_window("gemini-3.1-pro-preview") == 1_000_000 - - -class TestPrefixMatch: - """Prefix-based fuzzy matching.""" - - def test_dated_suffix(self) -> None: - assert _context_window("claude-sonnet-4-20250514") == 200_000 - - def test_latest_suffix(self) -> None: - assert _context_window("claude-3-5-sonnet-latest") == 200_000 - - def test_preview_suffix(self) -> None: - assert _context_window("claude-3-5-sonnet-preview") == 200_000 - - def test_o1_mini_prefers_longer_key(self) -> None: - """o1-mini must match 'o1-mini' (128K), not 'o1' (200K).""" - assert _context_window("o1-mini") == 128_000 - - def test_o1_mini_with_date(self) -> None: - assert _context_window("o1-mini-20240101") == 128_000 - - def test_dot_and_dash_notation(self) -> None: - assert _context_window("claude-3.5-sonnet") == 200_000 - assert _context_window("claude-3-5-sonnet") == 200_000 - - -class TestUnknownModel: - """Unknown models return None.""" - - def test_completely_unknown(self) -> None: - assert _context_window("totally-unknown-model") is None - - def test_empty_string(self) -> None: - assert _context_window("") is None - - def test_partial_gpt(self) -> None: - assert _context_window("gpt") is None - - -class TestTableConsistency: - """Sanity checks on the registry.""" - - def test_all_context_windows_positive(self) -> None: - for model, pricing in DEFAULT_PRICING.items(): - if pricing.context_window is not None: - assert pricing.context_window > 0, f"{model} has non-positive context window" - - def test_all_have_context_window(self) -> None: - for model, pricing in DEFAULT_PRICING.items(): - assert pricing.context_window is not None, f"{model} is missing context_window" diff --git a/tests/test_providers/test_copilot.py b/tests/test_providers/test_copilot.py index e1d0664e..801a0003 100644 --- a/tests/test_providers/test_copilot.py +++ b/tests/test_providers/test_copilot.py @@ -695,3 +695,68 @@ def test_runs_on_unix(self, monkeypatch: pytest.MonkeyPatch) -> None: # early — it proceeded past the guard and attempted the import. with contextlib.suppress(ModuleNotFoundError): provider._fix_pipe_blocking_mode() + + +class TestGetMaxPromptTokens: + """Tests for CopilotProvider.get_max_prompt_tokens.""" + + @pytest.mark.asyncio + async def test_mock_handler_mode_returns_none(self) -> None: + """Mock-handler mode has no SDK to query — must return None.""" + provider = CopilotProvider(mock_handler=stub_handler) + assert await provider.get_max_prompt_tokens("gpt-4o") is None + + @pytest.mark.asyncio + async def test_returns_max_prompt_tokens_for_known_model(self) -> None: + """Looks up the matching model and returns its max_prompt_tokens.""" + + class _Limits: + max_prompt_tokens = 128000 + + class _Caps: + limits = _Limits() + + class _Model: + id = "gpt-4o" + capabilities = _Caps() + + class _FakeClient: + async def list_models(self) -> list[Any]: + return [_Model()] + + provider = CopilotProvider(mock_handler=stub_handler) + provider._mock_handler = None # disable mock-handler short-circuit + provider._client = _FakeClient() + # Skip _ensure_client_started by marking as already-started. + provider._started = True + + result = await provider.get_max_prompt_tokens("gpt-4o") + assert result == 128000 + + @pytest.mark.asyncio + async def test_returns_none_for_unknown_model(self) -> None: + class _FakeClient: + async def list_models(self) -> list[Any]: + return [] + + provider = CopilotProvider(mock_handler=stub_handler) + provider._mock_handler = None + provider._client = _FakeClient() + provider._started = True + + assert await provider.get_max_prompt_tokens("anything") is None + + @pytest.mark.asyncio + async def test_sdk_failure_returns_none(self) -> None: + """An exception from the SDK must not propagate; metadata is best-effort.""" + + class _BoomClient: + async def list_models(self) -> list[Any]: + raise RuntimeError("network down") + + provider = CopilotProvider(mock_handler=stub_handler) + provider._mock_handler = None + provider._client = _BoomClient() + provider._started = True + + assert await provider.get_max_prompt_tokens("gpt-4o") is None From 4965eb50e1dbbf3ac9ab4178ea009d1ff91e6124 Mon Sep 17 00:00:00 2001 From: Jason Robert Date: Mon, 4 May 2026 13:54:50 -0400 Subject: [PATCH 2/2] =?UTF-8?q?fix(providers):=20address=20PR=20review=20?= =?UTF-8?q?=E2=80=94=20alias=20matching,=20retry=20chain,=20narrow=20excep?= =?UTF-8?q?tions?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Code review and rubber-duck reviews surfaced six issues on the SDK-driven context window refactor. This commit addresses all of them and applies the simplifier's cleanups. Issue 1: lost alias resolution (Important) Names like "claude-3-5-sonnet-latest" and short aliases used in our own examples (claude-haiku-4.5, claude-sonnet-4.5) silently lost context_window_max because the SDK's models.list() returns dated IDs (claude-3-5-sonnet-20241022), not the alias names users configure. Added match_model_id() in providers/base.py: a small alias-aware matcher (exact -> boundary-prefix in either direction -> suffix-strip retry). Both Copilot and Claude providers route SDK lookups through it. Covered by 11 unit tests in test_base.py. Issue 2: output.model precedence was a string-choice, not a retry chain If output.model was an SDK-unknown variant (e.g. a reasoning-effort tier the provider doesn't list), the resolver returned None instead of falling back to agent.model. WorkflowEngine._get_context_window_for_agent now tries each candidate in order (output.model -> agent.model -> default) and returns the first non-None lookup. Covered by a new test in test_context_window_events.py. Issue 3: lock held across the SDK round-trip (Minor) ClaudeProvider.get_max_prompt_tokens held _max_input_cache_lock while awaiting client.models.list(); parallel first-callers all blocked behind the same SDK call. Refactored: fetch outside the lock, lock only around the dict install. Also seed the cache from validate_connection() so callers who go through normal connection setup never pay for the round-trip. Issue 4: transient SDK failure cached forever (Minor) Previous implementation installed an empty cache on exception, so every subsequent call returned None forever. Now: don't cache failures. The cache stays None and the next call retries. Covered by test_sdk_failure_returns_none_and_does_not_cache. Issue 5: race in ClaudeProvider.close() (Critical) close() didn't acquire the cache lock and could close the client while a concurrent get_max_prompt_tokens() was awaiting models.list(). Now: drop the _client reference *before* awaiting close(), and invalidate _max_input_cache. In-flight requests will error and be swallowed by the metadata path's narrow except. Issue 6: error swallowing too broad (Minor) Three layers of `except Exception: return None` (provider x 2 + engine) hid genuine bugs. Provider methods now catch only SDK/transport errors narrowly (AnthropicError | OSError | TimeoutError for Claude; ProviderError | OSError | RuntimeError | TimeoutError for Copilot). The engine keeps its broad outer catch as the safety net so unexpected provider bugs still don't break workflows. Covered by test_unexpected_exception_propagates. Simplifier cleanups applied: - Dropped redundant `if limits else None` ternary in copilot.py (getattr with default never raises). - Dropped `or {}` and `if matched_id else None` redundancies in claude.py get_max_prompt_tokens. - Inlined the one-line _model_known() helper. - Consolidated 5 Copilot test methods through shared _make_model and _provider_with_list_models helpers (~60 lines of duplication removed). Validation: - 2029 passed, 9 skipped (1977 + 52 new/updated). - ruff check + format pass. - ty type check passes. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- src/conductor/engine/workflow.py | 55 +++++---- src/conductor/providers/base.py | 43 ++++++- src/conductor/providers/claude.py | 108 ++++++++++++------ src/conductor/providers/copilot.py | 29 +++-- .../test_engine/test_context_window_events.py | 49 ++++++++ tests/test_providers/test_base.py | 53 ++++++++- tests/test_providers/test_claude.py | 72 +++++++++++- tests/test_providers/test_copilot.py | 107 ++++++++++------- 8 files changed, 398 insertions(+), 118 deletions(-) diff --git a/src/conductor/engine/workflow.py b/src/conductor/engine/workflow.py index 17d19f51..c6076aad 100644 --- a/src/conductor/engine/workflow.py +++ b/src/conductor/engine/workflow.py @@ -788,33 +788,42 @@ async def _get_context_window_for_agent( ) -> int | None: """Return the SDK-reported max prompt tokens for an agent. - Resolves the model in priority order: the model the SDK actually used - (``output.model``), then the agent's configured model, then the - workflow's runtime default. Returns ``None`` when no model can be - resolved, no provider can be reached, or the provider's metadata call - fails — context-window metadata is best-effort and must never break - workflow execution. + Tries each candidate model in priority order — the model the SDK + actually used (``output.model``), the agent's configured model, the + workflow's runtime default — and returns the first non-``None`` + result. This is a real fallback chain: if ``output.model`` is an + SDK-specific variant the provider doesn't know about, the lookup + retries with ``agent.model`` before giving up. + + Returns ``None`` when no candidate resolves, no provider can be + reached, or the provider's metadata call fails — context-window + metadata is best-effort and must never break workflow execution. """ - model = ( - (output.model if output is not None else None) - or agent.model - or self.config.workflow.runtime.default_model - ) - if not model: - return None provider = await self._get_provider_for_agent(agent) if provider is None: return None - try: - return await provider.get_max_prompt_tokens(model) - except Exception as e: - logger.debug( - "get_max_prompt_tokens(%r) raised on provider for agent %s: %s", - model, - agent.name, - e, - ) - return None + candidates: list[str] = [] + if output is not None and output.model: + candidates.append(output.model) + if agent.model and agent.model not in candidates: + candidates.append(agent.model) + default = self.config.workflow.runtime.default_model + if default and default not in candidates: + candidates.append(default) + for model in candidates: + try: + value = await provider.get_max_prompt_tokens(model) + except Exception as e: + logger.debug( + "get_max_prompt_tokens(%r) raised on provider for agent %s: %s", + model, + agent.name, + e, + ) + continue + if value is not None: + return value + return None async def run(self, inputs: dict[str, Any]) -> dict[str, Any]: """Execute the workflow from entry_point to $end. diff --git a/src/conductor/providers/base.py b/src/conductor/providers/base.py index 4aa4e963..2b56c6b5 100644 --- a/src/conductor/providers/base.py +++ b/src/conductor/providers/base.py @@ -7,8 +7,9 @@ from __future__ import annotations import asyncio +import re from abc import ABC, abstractmethod -from collections.abc import Callable +from collections.abc import Callable, Iterable from dataclasses import dataclass from typing import TYPE_CHECKING, Any @@ -20,6 +21,46 @@ EventCallback = Callable[[str, dict[str, Any]], None] +# Suffixes that providers may strip when matching aliased model names against +# their SDK's canonical IDs (e.g. "claude-3-5-sonnet-latest" -> base name). +_VERSION_SUFFIX_RE = re.compile(r"-(\d{8}|latest|preview)$") + + +def match_model_id(requested: str, known_ids: Iterable[str]) -> str | None: + """Find the canonical SDK ID matching a possibly aliased model name. + + Match strategies, in order: + + 1. Exact match. + 2. Boundary prefix match (longest first), in either direction. Handles + both ``"claude-3-5-sonnet-20241022"`` for requested + ``"claude-3-5-sonnet"`` *and* the reverse, where the SDK lists a + dated ID and the user specified the base name. + 3. Suffix-strip (``-YYYYMMDD``, ``-latest``, ``-preview``) on the + requested name, then re-try strategies 1 and 2. + + Returns the matching SDK ID, or ``None`` if no strategy succeeds. + """ + ids = [str(i) for i in known_ids] + if not ids: + return None + if requested in ids: + return requested + sorted_ids = sorted(ids, key=lambda s: len(s), reverse=True) + for known in sorted_ids: + if requested.startswith(known + "-") or known.startswith(requested + "-"): + return known + simplified = _VERSION_SUFFIX_RE.sub("", requested) + if simplified == requested: + return None + if simplified in ids: + return simplified + for known in sorted_ids: + if simplified.startswith(known + "-") or known.startswith(simplified + "-"): + return known + return None + + @dataclass class AgentOutput: """Normalized output from any SDK provider. diff --git a/src/conductor/providers/claude.py b/src/conductor/providers/claude.py index d0a0bc56..b3feec81 100644 --- a/src/conductor/providers/claude.py +++ b/src/conductor/providers/claude.py @@ -30,7 +30,7 @@ from conductor.exceptions import ProviderError, ValidationError from conductor.executor.output import validate_output -from conductor.providers.base import AgentOutput, AgentProvider, EventCallback +from conductor.providers.base import AgentOutput, AgentProvider, EventCallback, match_model_id if TYPE_CHECKING: from conductor.config.schema import AgentDef, OutputField @@ -39,13 +39,14 @@ # Try to import the Anthropic SDK try: import anthropic - from anthropic import AsyncAnthropic + from anthropic import AnthropicError, AsyncAnthropic ANTHROPIC_SDK_AVAILABLE = True except ImportError: ANTHROPIC_SDK_AVAILABLE = False AsyncAnthropic = None # type: ignore[misc, assignment] anthropic = None # type: ignore[assignment] + AnthropicError = Exception # type: ignore[misc, assignment] logger = logging.getLogger(__name__) @@ -295,7 +296,8 @@ async def validate_connection(self) -> bool: async def _log_available_models(self) -> None: """List and log available models, warn if default model is unavailable. - Consolidated from the former _verify_available_models method. + Also seeds ``_max_input_cache`` so the first call to + :meth:`get_max_prompt_tokens` doesn't pay for an extra round-trip. """ if self._client is None: return @@ -304,50 +306,73 @@ async def _log_available_models(self) -> None: # Call client.models.list() to get available models (async) logger.debug("Discovering available Claude models via client.models.list()...") models_page = await self._client.models.list() - available_models = [model.id for model in models_page.data] + except (TimeoutError, AnthropicError, OSError) as e: + logger.warning(f"Could not list available models (discovery failed): {e}") + return - logger.info(f"Available Claude models: {', '.join(available_models)}") + available_models = [model.id for model in models_page.data] + logger.info(f"Available Claude models: {', '.join(available_models)}") - # Warn if default model not in list - if self._default_model not in available_models: - logger.warning( - f"Requested model '{self._default_model}' is not in the list of " - f"available models. API calls may fail. Available: {available_models}" - ) - else: - logger.debug(f"Default model '{self._default_model}' verified in available models") - except Exception as e: - logger.warning(f"Could not list available models (discovery failed): {e}") + # Warn if default model not in list (after stripping aliases like -latest). + if match_model_id(self._default_model, available_models) is None: + logger.warning( + f"Requested model '{self._default_model}' is not in the list of " + f"available models. API calls may fail. Available: {available_models}" + ) + else: + logger.debug(f"Default model '{self._default_model}' verified in available models") + + # Seed the metadata cache so get_max_prompt_tokens() is a pure lookup. + self._install_max_input_cache(models_page.data) + + def _install_max_input_cache(self, models_data: list[Any]) -> None: + """Replace ``_max_input_cache`` with a fresh mapping of id -> max_input.""" + self._max_input_cache = { + info.id: getattr(info, "max_input_tokens", None) for info in models_data + } async def get_max_prompt_tokens(self, model: str) -> int | None: """Return the Anthropic SDK's ``max_input_tokens`` for ``model``. - Lazily populates a per-instance cache on first call by enumerating - ``client.models.list()``. Subsequent calls are dictionary lookups. - Returns ``None`` if the SDK is unavailable, the client failed to - initialize, the model isn't listed, or the listing call raises — - context-window metadata must never block workflow execution. + On first call, populates a per-instance cache by enumerating + ``client.models.list()``; subsequent calls are dictionary lookups. + ``validate_connection()`` already populates the cache, so callers + that go through normal connection setup never pay for an extra + round-trip. + + Resolves aliases (``-latest``, dated suffixes, base/versioned name + mismatches) via :func:`match_model_id`. Returns ``None`` when the + SDK is unavailable, the model can't be resolved, or the listing + call fails — context-window metadata must never block workflow + execution. Note: the value reflects the API's *default* input window. Claude - models with a 1M-context beta require an explicit beta header, which - Conductor does not set today; for those models the API still reports - the default window. + models with a 1M-context beta require an explicit beta header, + which Conductor does not set today; for those models the API still + reports the default window. """ if not ANTHROPIC_SDK_AVAILABLE or self._client is None: return None - async with self._max_input_cache_lock: - if self._max_input_cache is None: - cache: dict[str, int | None] = {} - try: - page = await self._client.models.list() - except Exception as e: - logger.debug("Failed to list Anthropic models: %s", e) - self._max_input_cache = cache # cache the empty result - return None - for info in page.data: - cache[info.id] = getattr(info, "max_input_tokens", None) - self._max_input_cache = cache - return self._max_input_cache.get(model) + + if self._max_input_cache is None: + # Fetch outside the lock so concurrent callers don't all queue + # behind a slow round-trip; the lock only guards the install. + try: + page = await self._client.models.list() + except (TimeoutError, AnthropicError, OSError) as e: + # Don't cache the failure — let the next call retry. + logger.debug("Failed to list Anthropic models: %s", e) + return None + async with self._max_input_cache_lock: + if self._max_input_cache is None: + self._install_max_input_cache(page.data) + + # The block above either returned early on failure or installed the + # cache, so it's guaranteed non-None here. + cache = self._max_input_cache + assert cache is not None + matched_id = match_model_id(model, cache.keys()) + return cache.get(matched_id) if matched_id is not None else None async def _ensure_mcp_connected(self) -> None: """Connect to MCP servers if configured. @@ -435,11 +460,18 @@ async def close(self) -> None: logger.debug("MCP manager closed") if self._client is not None: - # AsyncAnthropic uses httpx AsyncClient internally which should be closed - await self._client.close() + # Drop the client reference *before* awaiting close() so any + # in-flight get_max_prompt_tokens() observes None on its next + # access and skips the SDK call. Already-issued requests will + # error and be swallowed by the metadata path's narrow except. + client = self._client self._client = None + await client.close() logger.debug("Claude provider closed") + # Drop cached metadata so a re-initialized provider re-fetches. + self._max_input_cache = None + async def execute( self, agent: AgentDef, diff --git a/src/conductor/providers/copilot.py b/src/conductor/providers/copilot.py index 311eed0f..fbb8fba6 100644 --- a/src/conductor/providers/copilot.py +++ b/src/conductor/providers/copilot.py @@ -18,7 +18,7 @@ from typing import TYPE_CHECKING, Any from conductor.exceptions import ProviderError, ValidationError -from conductor.providers.base import AgentOutput, AgentProvider, EventCallback +from conductor.providers.base import AgentOutput, AgentProvider, EventCallback, match_model_id if TYPE_CHECKING: from conductor.config.schema import AgentDef, OutputField @@ -1778,25 +1778,30 @@ async def close(self) -> None: async def get_max_prompt_tokens(self, model: str) -> int | None: """Return the Copilot SDK's ``max_prompt_tokens`` for ``model``. - Queries ``client.list_models()`` (cached internally by the SDK) and - returns ``capabilities.limits.max_prompt_tokens`` for the matching - model. Returns ``None`` in mock-handler mode, when the SDK is - unavailable, when the model is unknown, or when any SDK call fails — - context-window metadata must never block workflow execution. + Queries ``client.list_models()`` (cached internally by the SDK), + resolves any aliases (e.g. ``-latest``, dated suffixes, base-name + vs versioned-name) via :func:`match_model_id`, and returns + ``capabilities.limits.max_prompt_tokens`` for the matched entry. + + Returns ``None`` in mock-handler mode, when the SDK is unavailable, + when no match is found, or when the SDK call fails — context-window + metadata must never block workflow execution. """ if self._mock_handler is not None or not COPILOT_SDK_AVAILABLE: return None try: await self._ensure_client_started() models = await self._client.list_models() - except Exception as e: + except (TimeoutError, ProviderError, OSError, RuntimeError) as e: logger.debug("Failed to list Copilot models for %r: %s", model, e) return None - for info in models: - if info.id == model: - limits = getattr(info.capabilities, "limits", None) - return getattr(limits, "max_prompt_tokens", None) if limits else None - return None + by_id = {info.id: info for info in models} + matched_id = match_model_id(model, by_id.keys()) + if matched_id is None: + return None + info = by_id[matched_id] + limits = getattr(info.capabilities, "limits", None) + return getattr(limits, "max_prompt_tokens", None) def get_session_ids(self) -> dict[str, str]: """Get tracked session IDs for all executed agents. diff --git a/tests/test_engine/test_context_window_events.py b/tests/test_engine/test_context_window_events.py index 1c581da4..ec31bd7e 100644 --- a/tests/test_engine/test_context_window_events.py +++ b/tests/test_engine/test_context_window_events.py @@ -342,3 +342,52 @@ async def execute_with_model(*args, **kwargs): # type: ignore[no-untyped-def] assert collector.first("agent_started").data["context_window_max"] == 128000 # agent_completed has output.model — uses that assert collector.first("agent_completed").data["context_window_max"] == 400000 + + @pytest.mark.asyncio + async def test_falls_back_to_agent_model_when_output_model_unknown(self) -> None: + """If output.model is an SDK-unknown variant (e.g. a reasoning-effort + tier the provider doesn't list), the chain retries with agent.model + rather than returning None.""" + emitter, collector = _make_emitter_and_collector() + config = WorkflowConfig( + workflow=WorkflowDef( + name="test-fallback", + entry_point="a1", + runtime=RuntimeConfig(provider="copilot"), + context=ContextConfig(mode="accumulate"), + limits=LimitsConfig(max_iterations=10), + ), + agents=[ + AgentDef( + name="a1", + model="claude-opus-4.7", + prompt="Hello", + output={"answer": OutputField(type="string")}, + routes=[RouteDef(to="$end")], + ), + ], + output={"answer": "{{ a1.output.answer }}"}, + ) + + provider = CopilotProvider(mock_handler=lambda a, p, c: {"answer": "hi"}) + + # Provider only knows the base name, not the reasoning-effort variant. + async def fake_get_max_prompt_tokens(model: str) -> int | None: + return {"claude-opus-4.7": 200_000}.get(model) + + provider.get_max_prompt_tokens = fake_get_max_prompt_tokens # type: ignore[method-assign] + + original_execute = provider.execute + + async def execute_with_variant_model(*args, **kwargs): # type: ignore[no-untyped-def] + output = await original_execute(*args, **kwargs) + output.model = "claude-opus-4.7-xhigh" # SDK doesn't know this name + return output + + provider.execute = execute_with_variant_model # type: ignore[method-assign] + + engine = WorkflowEngine(config, provider, event_emitter=emitter) + await engine.run({}) + + # output.model returned None; chain fell back to agent.model. + assert collector.first("agent_completed").data["context_window_max"] == 200_000 diff --git a/tests/test_providers/test_base.py b/tests/test_providers/test_base.py index 86a8cbea..aaf4910f 100644 --- a/tests/test_providers/test_base.py +++ b/tests/test_providers/test_base.py @@ -1,6 +1,6 @@ """Unit tests for the AgentProvider ABC and AgentOutput dataclass.""" -from conductor.providers.base import AgentOutput +from conductor.providers.base import AgentOutput, match_model_id class TestAgentOutput: @@ -67,3 +67,54 @@ def test_agent_output_with_complex_content(self) -> None: assert output.content["analysis"]["score"] == 8.5 assert output.content["analysis"]["issues"] == ["minor", "cosmetic"] assert output.content["analysis"]["approved"] is True + + +class TestMatchModelId: + """Unit tests for the alias-aware model ID matcher.""" + + def test_exact_match(self) -> None: + assert match_model_id("gpt-4o", ["gpt-4o", "gpt-4.1"]) == "gpt-4o" + + def test_returns_none_when_no_known_ids(self) -> None: + assert match_model_id("gpt-4o", []) is None + + def test_returns_none_for_unrelated_name(self) -> None: + assert match_model_id("totally-different", ["gpt-4o"]) is None + + def test_versioned_suffix_matches_base(self) -> None: + # Requested name has a dated suffix; SDK lists the base name. + assert ( + match_model_id("claude-3-5-sonnet-20241022", ["claude-3-5-sonnet"]) + == "claude-3-5-sonnet" + ) + + def test_base_matches_versioned_sdk_id(self) -> None: + # Requested name is the base; SDK lists a dated/aliased variant. + assert ( + match_model_id("claude-3-5-sonnet", ["claude-3-5-sonnet-20241022"]) + == "claude-3-5-sonnet-20241022" + ) + + def test_latest_alias_strips_and_matches(self) -> None: + assert ( + match_model_id("claude-3-5-sonnet-latest", ["claude-3-5-sonnet-20241022"]) + == "claude-3-5-sonnet-20241022" + ) + + def test_preview_alias_strips_and_matches(self) -> None: + assert match_model_id("gemini-3.1-pro-preview", ["gemini-3.1-pro"]) == "gemini-3.1-pro" + + def test_longest_match_wins(self) -> None: + # "o1-mini-20240101" must match "o1-mini" (longer), not "o1". + assert match_model_id("o1-mini-20240101", ["o1", "o1-mini"]) == "o1-mini" + + def test_boundary_check_prevents_cross_family_match(self) -> None: + # "claude-opus-4.7" must NOT match "claude-opus-4" (different family). + # No valid match -> None. + assert match_model_id("claude-opus-4.7", ["claude-opus-4"]) is None + + def test_boundary_check_prevents_cross_family_with_suffix(self) -> None: + assert match_model_id("claude-opus-4.7-high", ["claude-opus-4"]) is None + + def test_unknown_after_suffix_strip_returns_none(self) -> None: + assert match_model_id("totally-different-latest", ["gpt-4o"]) is None diff --git a/tests/test_providers/test_claude.py b/tests/test_providers/test_claude.py index 147f7979..8365ac3e 100644 --- a/tests/test_providers/test_claude.py +++ b/tests/test_providers/test_claude.py @@ -2337,14 +2337,61 @@ async def test_returns_none_for_unknown_model(self, mock_anthropic_class: Mock) assert await provider.get_max_prompt_tokens("unknown-x") is None @pytest.mark.asyncio - async def test_sdk_failure_returns_none(self, mock_anthropic_class: Mock) -> None: - """An exception from models.list() must not propagate.""" + async def test_sdk_failure_returns_none_and_does_not_cache( + self, mock_anthropic_class: Mock + ) -> None: + """An SDK exception is swallowed and not cached, so a later call retries.""" + from anthropic import APIConnectionError + + # APIConnectionError requires a request kwarg; build a minimal one. + err = APIConnectionError(request=Mock()) + mock_client = Mock() - mock_client.models.list = AsyncMock(side_effect=RuntimeError("network down")) + # First call raises, second call succeeds — proves the failure isn't + # cached as "no metadata" forever. + mock_client.models.list = AsyncMock( + side_effect=[ + err, + Mock(data=[Mock(id="claude-sonnet-4-5", max_input_tokens=200_000)]), + ] + ) mock_anthropic_class.return_value = mock_client provider = ClaudeProvider() assert await provider.get_max_prompt_tokens("claude-sonnet-4-5") is None + assert await provider.get_max_prompt_tokens("claude-sonnet-4-5") == 200_000 + assert mock_client.models.list.await_count == 2 + + @pytest.mark.asyncio + async def test_unexpected_exception_propagates(self, mock_anthropic_class: Mock) -> None: + """Non-SDK exceptions (programming errors) are not swallowed by the + provider — they bubble up so the engine's outer safety net handles them.""" + mock_client = Mock() + mock_client.models.list = AsyncMock(side_effect=RuntimeError("bug")) + mock_anthropic_class.return_value = mock_client + + provider = ClaudeProvider() + with pytest.raises(RuntimeError): + await provider.get_max_prompt_tokens("claude-sonnet-4-5") + + @pytest.mark.asyncio + async def test_alias_resolves_via_match_model_id(self, mock_anthropic_class: Mock) -> None: + """``-latest`` and dated suffix aliases resolve to the SDK's listed ID.""" + mock_client = Mock() + mock_client.models.list = AsyncMock( + return_value=Mock( + data=[ + Mock(id="claude-3-5-sonnet-20241022", max_input_tokens=200_000), + ] + ) + ) + mock_anthropic_class.return_value = mock_client + + provider = ClaudeProvider() + # `-latest` strips to base, then prefix-matches the dated SDK ID. + assert await provider.get_max_prompt_tokens("claude-3-5-sonnet-latest") == 200_000 + # The base name (no dated suffix) also matches the dated SDK ID. + assert await provider.get_max_prompt_tokens("claude-3-5-sonnet") == 200_000 @pytest.mark.asyncio async def test_caches_after_first_call(self, mock_anthropic_class: Mock) -> None: @@ -2362,6 +2409,25 @@ async def test_caches_after_first_call(self, mock_anthropic_class: Mock) -> None assert mock_client.models.list.await_count == 1 + @pytest.mark.asyncio + async def test_validate_connection_seeds_cache(self, mock_anthropic_class: Mock) -> None: + """``validate_connection()`` populates the cache so the first + ``get_max_prompt_tokens()`` call is a pure dict lookup.""" + mock_client = Mock() + mock_client.models.list = AsyncMock( + return_value=Mock(data=[Mock(id="claude-sonnet-4-5", max_input_tokens=200_000)]) + ) + mock_anthropic_class.return_value = mock_client + + provider = ClaudeProvider() + assert await provider.validate_connection() is True + # validate_connection itself called list (once for the API check, then + # _log_available_models reuses the response). Reset the counter to + # prove get_max_prompt_tokens doesn't add another call. + before = mock_client.models.list.await_count + assert await provider.get_max_prompt_tokens("claude-sonnet-4-5") == 200_000 + assert mock_client.models.list.await_count == before + @pytest.mark.asyncio @patch("conductor.providers.claude.ANTHROPIC_SDK_AVAILABLE", False) async def test_returns_none_when_sdk_unavailable(self, mock_anthropic_class: Mock) -> None: diff --git a/tests/test_providers/test_copilot.py b/tests/test_providers/test_copilot.py index 801a0003..cc90515e 100644 --- a/tests/test_providers/test_copilot.py +++ b/tests/test_providers/test_copilot.py @@ -700,6 +700,33 @@ def test_runs_on_unix(self, monkeypatch: pytest.MonkeyPatch) -> None: class TestGetMaxPromptTokens: """Tests for CopilotProvider.get_max_prompt_tokens.""" + @staticmethod + def _make_model(model_id: str, max_prompt_tokens: int) -> Any: + from types import SimpleNamespace + + return SimpleNamespace( + id=model_id, + capabilities=SimpleNamespace( + limits=SimpleNamespace(max_prompt_tokens=max_prompt_tokens) + ), + ) + + @staticmethod + def _provider_with_list_models(list_models_impl: Any) -> CopilotProvider: + """Build a provider with the SDK short-circuit disabled and a fake client. + + Uses ``stub_handler`` for ``mock_handler`` to satisfy the constructor, + then nulls ``_mock_handler`` so ``get_max_prompt_tokens`` falls through + to the SDK path. ``_started=True`` skips ``_ensure_client_started``. + """ + from types import SimpleNamespace + + provider = CopilotProvider(mock_handler=stub_handler) + provider._mock_handler = None + provider._client = SimpleNamespace(list_models=list_models_impl) + provider._started = True + return provider + @pytest.mark.asyncio async def test_mock_handler_mode_returns_none(self) -> None: """Mock-handler mode has no SDK to query — must return None.""" @@ -708,55 +735,55 @@ async def test_mock_handler_mode_returns_none(self) -> None: @pytest.mark.asyncio async def test_returns_max_prompt_tokens_for_known_model(self) -> None: - """Looks up the matching model and returns its max_prompt_tokens.""" + async def list_models() -> list[Any]: + return [self._make_model("gpt-4o", 128000)] - class _Limits: - max_prompt_tokens = 128000 + provider = self._provider_with_list_models(list_models) + assert await provider.get_max_prompt_tokens("gpt-4o") == 128000 - class _Caps: - limits = _Limits() - - class _Model: - id = "gpt-4o" - capabilities = _Caps() - - class _FakeClient: - async def list_models(self) -> list[Any]: - return [_Model()] + @pytest.mark.asyncio + async def test_returns_none_for_unknown_model(self) -> None: + async def list_models() -> list[Any]: + return [] - provider = CopilotProvider(mock_handler=stub_handler) - provider._mock_handler = None # disable mock-handler short-circuit - provider._client = _FakeClient() - # Skip _ensure_client_started by marking as already-started. - provider._started = True + provider = self._provider_with_list_models(list_models) + assert await provider.get_max_prompt_tokens("anything") is None - result = await provider.get_max_prompt_tokens("gpt-4o") - assert result == 128000 + @pytest.mark.asyncio + async def test_oserror_returns_none_and_does_not_cache(self) -> None: + """A transport-level error is swallowed; the next call retries.""" + calls = 0 + + async def list_models() -> list[Any]: + nonlocal calls + calls += 1 + if calls == 1: + raise OSError("network down") + return [self._make_model("gpt-4o", 128000)] + + provider = self._provider_with_list_models(list_models) + assert await provider.get_max_prompt_tokens("gpt-4o") is None + assert await provider.get_max_prompt_tokens("gpt-4o") == 128000 @pytest.mark.asyncio - async def test_returns_none_for_unknown_model(self) -> None: - class _FakeClient: - async def list_models(self) -> list[Any]: - return [] + async def test_unexpected_exception_propagates(self) -> None: + """Non-SDK exceptions (programming errors) are not swallowed by the + provider — they bubble up so the engine's outer safety net handles them.""" - provider = CopilotProvider(mock_handler=stub_handler) - provider._mock_handler = None - provider._client = _FakeClient() - provider._started = True + async def list_models() -> list[Any]: + raise ValueError("bug") - assert await provider.get_max_prompt_tokens("anything") is None + provider = self._provider_with_list_models(list_models) + with pytest.raises(ValueError): + await provider.get_max_prompt_tokens("gpt-4o") @pytest.mark.asyncio - async def test_sdk_failure_returns_none(self) -> None: - """An exception from the SDK must not propagate; metadata is best-effort.""" - - class _BoomClient: - async def list_models(self) -> list[Any]: - raise RuntimeError("network down") + async def test_alias_resolves_via_match_model_id(self) -> None: + """Versioned-suffix aliases resolve to the SDK's listed ID.""" - provider = CopilotProvider(mock_handler=stub_handler) - provider._mock_handler = None - provider._client = _BoomClient() - provider._started = True + async def list_models() -> list[Any]: + return [self._make_model("claude-3-5-sonnet", 200_000)] - assert await provider.get_max_prompt_tokens("gpt-4o") is None + provider = self._provider_with_list_models(list_models) + assert await provider.get_max_prompt_tokens("claude-3-5-sonnet-latest") == 200_000 + assert await provider.get_max_prompt_tokens("claude-3-5-sonnet-20241022") == 200_000