test: fix sys.modules stub leakage and apply repo formatting

test_claude_agents_* now remove their placeholder packages after loading, which previously blocked real imports of the temporal plugin tree in later-collected tests. Formats touched files with ruff and narrows span.output types in new tests for pyright.

Author: levilentz

src/agentex/lib/adk/_modules/tracing.py

@@ -247,9 +247,7 @@ async def turn_span(
 
         Example::
 
-            async with adk.tracing.turn_span(
-                trace_id=task.id, name="turn", input={...}, task_id=task.id
-            ) as turn:
+            async with adk.tracing.turn_span(trace_id=task.id, name="turn", input={...}, task_id=task.id) as turn:
                 result = await run_llm_calls()
                 turn.output = {"response": result.text}
                 turn.record_usage(usage=result.usage, cost_usd=result.cost_usd)

src/agentex/lib/core/temporal/plugins/openai_agents/models/temporal_tracing_model.py

@@ -7,6 +7,7 @@
 The key innovation is that these are thin wrappers around the standard OpenAI models,
 avoiding code duplication while adding tracing capabilities.
 """
+
 from __future__ import annotations
 
 import logging
@@ -51,24 +52,28 @@ def _serialize_item(item: Any) -> dict[str, Any]:
     Uses model_dump() for Pydantic models, otherwise extracts attributes manually.
     Filters out internal Pydantic fields that can't be serialized.
     """
-    if hasattr(item, 'model_dump'):
+    if hasattr(item, "model_dump"):
         # Pydantic model - use model_dump for proper serialization
         try:
-            return item.model_dump(mode='json', exclude_unset=True)
+            return item.model_dump(mode="json", exclude_unset=True)
         except Exception:
             # Fallback to dict conversion
-            return dict(item) if hasattr(item, '__iter__') else {}
+            return dict(item) if hasattr(item, "__iter__") else {}
     else:
         # Not a Pydantic model - extract attributes manually
         item_dict = {}
         for attr_name in dir(item):
-            if not attr_name.startswith('_') and attr_name not in ('model_fields', 'model_config', 'model_computed_fields'):
+            if not attr_name.startswith("_") and attr_name not in (
+                "model_fields",
+                "model_config",
+                "model_computed_fields",
+            ):
                 try:
                     attr_value = getattr(item, attr_name, None)
                     # Skip methods and None values
                     if attr_value is not None and not callable(attr_value):
                         # Convert to JSON-serializable format
-                        if hasattr(attr_value, 'model_dump'):
+                        if hasattr(attr_value, "model_dump"):
                             item_dict[attr_name] = attr_value.model_dump()
                         elif isinstance(attr_value, (str, int, float, bool, list, dict)):
                             item_dict[attr_name] = attr_value
@@ -106,7 +111,9 @@ def __init__(self, openai_client: Optional[AsyncOpenAI] = None, **kwargs):
         # Initialize tracer for all models
         agentex_client = create_async_agentex_client()
         self._tracer = AsyncTracer(agentex_client)
-        logger.info(f"[TemporalTracingModelProvider] Initialized with AgentEx tracer, custom_client={openai_client is not None}")
+        logger.info(
+            f"[TemporalTracingModelProvider] Initialized with AgentEx tracer, custom_client={openai_client is not None}"
+        )
 
     @override
     def get_model(self, model_name: Optional[str]) -> Model:
@@ -126,10 +133,14 @@ def get_model(self, model_name: Optional[str]) -> Model:
             logger.info(f"[TemporalTracingModelProvider] Wrapping OpenAIResponsesModel '{model_name}' with tracing")
             return TemporalTracingResponsesModel(base_model, self._tracer)  # type: ignore[abstract]
         elif isinstance(base_model, OpenAIChatCompletionsModel):
-            logger.info(f"[TemporalTracingModelProvider] Wrapping OpenAIChatCompletionsModel '{model_name}' with tracing")
+            logger.info(
+                f"[TemporalTracingModelProvider] Wrapping OpenAIChatCompletionsModel '{model_name}' with tracing"
+            )
             return TemporalTracingChatCompletionsModel(base_model, self._tracer)  # type: ignore[abstract]
         else:
-            logger.warning(f"[TemporalTracingModelProvider] Unknown model type, returning without tracing: {type(base_model)}")
+            logger.warning(
+                f"[TemporalTracingModelProvider] Unknown model type, returning without tracing: {type(base_model)}"
+            )
             return base_model
 
 
@@ -181,7 +192,9 @@ async def get_response(
 
         # If we have tracing context, wrap with span
         if trace_id and parent_span_id:
-            logger.debug(f"[TemporalTracingResponsesModel] Adding tracing span for task_id={task_id}, trace_id={trace_id}")
+            logger.debug(
+                f"[TemporalTracingResponsesModel] Adding tracing span for task_id={task_id}, trace_id={trace_id}"
+            )
 
             trace = self._tracer.trace(trace_id)
 
@@ -199,7 +212,9 @@ async def get_response(
                         "temperature": model_settings.temperature,
                         "max_tokens": model_settings.max_tokens,
                         "reasoning": model_settings.reasoning,
-                    } if model_settings else None,
+                    }
+                    if model_settings
+                    else None,
                 },
             ) as span:
                 try:
@@ -222,7 +237,7 @@ async def get_response(
                     new_items = []
                     final_output = None
 
-                    if hasattr(response, 'output') and response.output:
+                    if hasattr(response, "output") and response.output:
                         response_output = response.output if isinstance(response.output, list) else [response.output]
 
                         for item in response_output:
@@ -232,12 +247,12 @@ async def get_response(
                                     new_items.append(item_dict)
 
                                     # Extract final_output from message type if available
-                                    if item_dict.get('type') == 'message' and not final_output:
-                                        content = item_dict.get('content', [])
+                                    if item_dict.get("type") == "message" and not final_output:
+                                        content = item_dict.get("content", [])
                                         if content and isinstance(content, list):
                                             for content_part in content:
-                                                if isinstance(content_part, dict) and 'text' in content_part:
-                                                    final_output = content_part['text']
+                                                if isinstance(content_part, dict) and "text" in content_part:
+                                                    final_output = content_part["text"]
                                                     break
                             except Exception as e:
                                 logger.warning(f"Failed to serialize item in temporal tracing model: {e}")
@@ -329,7 +344,9 @@ async def get_response(
 
         # If we have tracing context, wrap with span
         if trace_id and parent_span_id:
-            logger.debug(f"[TemporalTracingChatCompletionsModel] Adding tracing span for task_id={task_id}, trace_id={trace_id}")
+            logger.debug(
+                f"[TemporalTracingChatCompletionsModel] Adding tracing span for task_id={task_id}, trace_id={trace_id}"
+            )
 
             trace = self._tracer.trace(trace_id)
 
@@ -346,7 +363,9 @@ async def get_response(
                     "model_settings": {
                         "temperature": model_settings.temperature,
                         "max_tokens": model_settings.max_tokens,
-                    } if model_settings else None,
+                    }
+                    if model_settings
+                    else None,
                 },
             ) as span:
                 try:
@@ -366,7 +385,7 @@ async def get_response(
                     new_items = []
                     final_output = None
 
-                    if hasattr(response, 'output') and response.output:
+                    if hasattr(response, "output") and response.output:
                         response_output = response.output if isinstance(response.output, list) else [response.output]
 
                         for item in response_output:
@@ -376,12 +395,12 @@ async def get_response(
                                     new_items.append(item_dict)
 
                                     # Extract final_output from message type if available
-                                    if item_dict.get('type') == 'message' and not final_output:
-                                        content = item_dict.get('content', [])
+                                    if item_dict.get("type") == "message" and not final_output:
+                                        content = item_dict.get("content", [])
                                         if content and isinstance(content, list):
                                             for content_part in content:
-                                                if isinstance(content_part, dict) and 'text' in content_part:
-                                                    final_output = content_part['text']
+                                                if isinstance(content_part, dict) and "text" in content_part:
+                                                    final_output = content_part["text"]
                                                     break
                             except Exception as e:
                                 logger.warning(f"Failed to serialize item in temporal tracing model: {e}")
@@ -406,7 +425,9 @@ async def get_response(
                     raise
         else:
             # No tracing context, just pass through
-            logger.debug("[TemporalTracingChatCompletionsModel] No tracing context available, calling base model directly")
+            logger.debug(
+                "[TemporalTracingChatCompletionsModel] No tracing context available, calling base model directly"
+            )
             return await self._base_model.get_response(
                 system_instructions=system_instructions,
                 input=input,
@@ -422,4 +443,4 @@ async def get_response(
     def stream_response(self, *args, **kwargs):
         """Streaming is handled via get_response in Temporal activities.
         Required so the class is concrete and instantiable at runtime."""
-        raise NotImplementedError("stream_response is not used in Temporal activities - use get_response instead")
+        raise NotImplementedError("stream_response is not used in Temporal activities - use get_response instead")

src/agentex/lib/utils/completions.py

@@ -26,8 +26,7 @@ def _(a: Completion, b: Completion) -> Completion:
     # final chunk carries usage but no choices. Keep the unpaired side instead
     # of truncating to the shorter list.
     a.choices = [
-        x if y is None else y if x is None else _concat_chunks(x, y)
-        for x, y in zip_longest(a.choices, b.choices)
+        x if y is None else y if x is None else _concat_chunks(x, y) for x, y in zip_longest(a.choices, b.choices)
     ]
     a.usage = _concat_chunks(a.usage, b.usage)
 
@@ -45,6 +44,7 @@ def _(a: Choice, b: Choice) -> Choice:
     a.finish_reason = a.finish_reason or b.finish_reason
     return a
 
+
 @_concat_chunks.register
 def _(a: Usage | None, b: Usage | None) -> Usage | None:
     if a is not None and b is not None:
@@ -68,9 +68,7 @@ def _(a: Delta, b: Delta) -> Delta:
             if tool_call.index not in grouped_tool_calls:
                 grouped_tool_calls[tool_call.index] = tool_call
             else:
-                grouped_tool_calls[tool_call.index] = _concat_chunks(
-                    grouped_tool_calls[tool_call.index], tool_call
-                )
+                grouped_tool_calls[tool_call.index] = _concat_chunks(grouped_tool_calls[tool_call.index], tool_call)
 
         a.tool_calls = list(grouped_tool_calls.values())
     elif hasattr(b, "tool_calls") and b.tool_calls:
@@ -88,11 +86,7 @@ def _(a: ToolCallRequest, b: ToolCallRequest) -> ToolCallRequest:
     index_val = a.index if hasattr(a, "index") and a.index is not None else b.index
 
     # Concatenate the function part
-    function_val = (
-        _concat_chunks(a.function, b.function)
-        if a.function and b.function
-        else a.function or b.function
-    )
+    function_val = _concat_chunks(a.function, b.function) if a.function and b.function else a.function or b.function
 
     # Set all properties
     a.id = id_val

tests/lib/adk/providers/test_litellm_usage.py

@@ -7,13 +7,13 @@
 
 from __future__ import annotations
 
+from typing import Any
 from datetime import UTC, datetime
 from contextlib import asynccontextmanager
 from unittest.mock import AsyncMock, MagicMock
 
 from agentex.types.span import Span
 from agentex.types.task_message import TaskMessage
-from agentex.types.task_message_content import TextContent
 from agentex.lib.types.llm_messages import (
     Delta,
     Usage,
@@ -22,6 +22,7 @@
     Completion,
     AssistantMessage,
 )
+from agentex.types.task_message_content import TextContent
 from agentex.lib.core.services.adk.providers.litellm import (
     LiteLLMService,
     _stream_kwargs_with_usage,
@@ -54,6 +55,11 @@ def trace(self, trace_id):
         return self.trace_obj
 
 
+def _output_dict(span: Span) -> dict[str, Any]:
+    assert isinstance(span.output, dict)
+    return span.output
+
+
 def _make_streaming_service():
     streaming_context = MagicMock()
     streaming_context.task_message = TaskMessage(
@@ -142,7 +148,7 @@ async def test_span_output_carries_usage(self):
         )
 
         span = tracer.trace_obj.spans[0]
-        assert span.output["usage"] == {"prompt_tokens": 7, "completion_tokens": 3, "total_tokens": 10}
+        assert _output_dict(span)["usage"] == {"prompt_tokens": 7, "completion_tokens": 3, "total_tokens": 10}
 
     async def test_span_output_omits_usage_when_absent(self):
         completion = Completion(
@@ -158,7 +164,7 @@ async def test_span_output_omits_usage_when_absent(self):
             trace_id="trace-1",
         )
 
-        assert "usage" not in tracer.trace_obj.spans[0].output
+        assert "usage" not in _output_dict(tracer.trace_obj.spans[0])
 
 
 class TestChatCompletionStream:
@@ -177,8 +183,8 @@ async def test_stream_requests_usage_and_span_output_carries_it(self):
         assert len(results) == 3
         assert captured_kwargs["stream_options"] == {"include_usage": True}
         span = tracer.trace_obj.spans[0]
-        assert span.output["usage"] == {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}
-        assert span.output["choices"][0]["message"]["content"] == "Hello!"
+        assert _output_dict(span)["usage"] == {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}
+        assert _output_dict(span)["choices"][0]["message"]["content"] == "Hello!"
 
 
 class TestChatCompletionStreamAutoSend:
@@ -195,9 +201,9 @@ async def test_usage_only_final_chunk_reaches_span_output(self):
 
         assert captured_kwargs["stream_options"] == {"include_usage": True}
         span = tracer.trace_obj.spans[0]
-        assert span.output["usage"] == {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}
+        assert _output_dict(span)["usage"] == {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}
         # TaskMessage dump is still the base of the span output
-        assert span.output["id"] == "msg-1"
+        assert _output_dict(span)["id"] == "msg-1"
 
     async def test_stream_without_usage_chunk_omits_usage(self):
         captured_kwargs: dict = {}
@@ -210,4 +216,4 @@ async def test_stream_without_usage_chunk_omits_usage(self):
             trace_id="trace-1",
         )
 
-        assert "usage" not in tracer.trace_obj.spans[0].output
+        assert "usage" not in _output_dict(tracer.trace_obj.spans[0])

tests/lib/adk/test_tracing_module.py

@@ -314,9 +314,7 @@ async def test_turn_span_preserves_existing_data(self):
         mock_service.end_span.return_value = started
 
         with patch.object(_tracing_mod, "in_temporal_workflow", return_value=False):
-            async with module.turn_span(
-                trace_id="trace-123", name="turn", data={"custom": "value"}
-            ) as turn:
+            async with module.turn_span(trace_id="trace-123", name="turn", data={"custom": "value"}) as turn:
                 turn.record_usage(usage={"prompt_tokens": 3, "completion_tokens": 4})
 
         ended_span = mock_service.end_span.call_args.kwargs["span"]