wip: improve coverage

Author: dsfaccini

pydantic_ai_slim/pydantic_ai/_parts_manager.py

@@ -73,6 +73,25 @@ def get_parts(self) -> list[ModelResponsePart]:
         """
         return [p for p in self._parts if not isinstance(p, ToolCallPartDelta)]
 
+    def has_incomplete_parts(self) -> bool:
+        """Check if there are any incomplete ToolCallPartDeltas being managed.
+
+        Returns:
+            True if there are any ToolCallPartDelta objects in the internal parts list.
+        """
+        return any(isinstance(p, ToolCallPartDelta) for p in self._parts)
+
+    def is_vendor_id_mapped(self, vendor_id: VendorId) -> bool:
+        """Check if a vendor ID is currently mapped to a part index.
+
+        Args:
+            vendor_id: The vendor ID to check.
+
+        Returns:
+            True if the vendor ID is mapped to a part index, False otherwise.
+        """
+        return vendor_id in self._vendor_id_to_part_index
+
     def finalize(self) -> Generator[ModelResponseStreamEvent, None, None]:
         """Flush any buffered content, appending to ThinkingParts or creating TextParts.
 
@@ -106,7 +125,7 @@ def finalize(self) -> Generator[ModelResponseStreamEvent, None, None]:
 
         # flush any remaining buffered content
         for vendor_part_id, buffered_content in list(self._thinking_tag_buffer.items()):
-            if buffered_content:
+            if buffered_content:  # pragma: no branch - buffer should never contain empty string
                 part_index = self._vendor_id_to_part_index.get(vendor_part_id)
 
                 # If buffered content belongs to a ThinkingPart, append it to the ThinkingPart
@@ -208,33 +227,7 @@ def _handle_text_delta_simple(  # noqa: C901
             if part_index is not None:
                 existing_part = self._parts[part_index]
 
-                if thinking_tags and isinstance(existing_part, ThinkingPart):
-                    end_tag = thinking_tags[1]
-                    if end_tag in content:
-                        before_end, after_end = content.split(end_tag, 1)
-
-                        if before_end:
-                            yield from self.handle_thinking_delta(vendor_part_id=vendor_part_id, content=before_end)
-
-                        self._vendor_id_to_part_index.pop(vendor_part_id)
-
-                        if after_end:
-                            yield from self._handle_text_delta_simple(
-                                vendor_part_id=vendor_part_id,
-                                content=after_end,
-                                id=id,
-                                thinking_tags=thinking_tags,
-                                ignore_leading_whitespace=ignore_leading_whitespace,
-                            )
-                        return
-
-                    if content == end_tag:
-                        self._vendor_id_to_part_index.pop(vendor_part_id)
-                        return
-
-                    yield from self.handle_thinking_delta(vendor_part_id=vendor_part_id, content=content)
-                    return
-                elif isinstance(existing_part, TextPart):
+                if isinstance(existing_part, TextPart):
                     existing_text_part_and_index = existing_part, part_index
                 else:
                     raise UnexpectedModelBehavior(f'Cannot apply a text delta to {existing_part=}')
@@ -267,19 +260,17 @@ def _handle_text_delta_simple(  # noqa: C901
             # Create ThinkingPart but defer PartStartEvent until there is content
             new_part_index = len(self._parts)
             part = ThinkingPart(content='')
-            if vendor_part_id is not None:
-                self._vendor_id_to_part_index[vendor_part_id] = new_part_index
             self._parts.append(part)
 
-            if after_start:
+            if after_start:  # pragma: no branch
                 yield from self.handle_thinking_delta(vendor_part_id=vendor_part_id, content=after_start)
             return
 
         if existing_text_part_and_index is None:
             # This is a workaround for models that emit `<think>\n</think>\n\n` or an empty text part ahead of tool calls (e.g. Ollama + Qwen3),
             # which we don't want to end up treating as a final result when using `run_stream` with `str` as a valid `output_type`.
             if ignore_leading_whitespace and (len(content) == 0 or content.isspace()):
-                return None
+                return
 
             new_part_index = len(self._parts)
             part = TextPart(content=content, id=id)
@@ -294,7 +285,9 @@ def _handle_text_delta_simple(  # noqa: C901
 
             updated_text_part = part_delta.apply(existing_text_part)
             self._parts[part_index] = updated_text_part
-            if part_index not in self._started_part_indices:
+            if (
+                part_index not in self._started_part_indices
+            ):  # pragma: no cover - defensive: TextPart should always be started
                 self._started_part_indices.add(part_index)
                 yield PartStartEvent(index=part_index, part=updated_text_part)
             else:
@@ -458,6 +451,10 @@ def handle_thinking_delta(
                 latest_part = self._parts[part_index]
                 if isinstance(latest_part, ThinkingPart):
                     existing_thinking_part_and_index = latest_part, part_index
+                elif isinstance(latest_part, TextPart):
+                    raise UnexpectedModelBehavior(
+                        'Cannot create ThinkingPart after TextPart: thinking must come before text in response'
+                    )
         else:
             # Otherwise, attempt to look up an existing ThinkingPart by vendor_part_id
             part_index = self._vendor_id_to_part_index.get(vendor_part_id)

tests/test_parts_manager.py

@@ -581,6 +581,9 @@ def test_handle_thinking_delta_new_part_with_vendor_id():
     parts = manager.get_parts()
     assert parts == snapshot([ThinkingPart(content='new thought')])
 
+    # Verify vendor_part_id was mapped to the part index
+    assert manager.is_vendor_id_mapped('thinking')
+
 
 def test_handle_thinking_delta_no_content():
     manager = ModelResponsePartsManager()
@@ -603,6 +606,98 @@ def test_handle_thinking_delta_no_content_or_signature():
             pass
 
 
+def test_handle_text_delta_append_to_thinking_part_without_vendor_id():
+    """Test appending to ThinkingPart when vendor_part_id is None (lines 202-203)."""
+    manager = ModelResponsePartsManager()
+    thinking_tags = ('<think>', '</think>')
+
+    # Create a ThinkingPart using handle_text_delta with thinking tags and vendor_part_id=None
+    events = list(manager.handle_text_delta(vendor_part_id=None, content='<think>initial', thinking_tags=thinking_tags))
+    assert len(events) == 1
+    assert isinstance(events[0], PartStartEvent)
+    assert isinstance(events[0].part, ThinkingPart)
+    assert events[0].part.content == 'initial'
+
+    # Now append more content with vendor_part_id=None - should append to existing ThinkingPart
+    events = list(manager.handle_text_delta(vendor_part_id=None, content=' reasoning', thinking_tags=thinking_tags))
+    assert len(events) == 1
+    assert isinstance(events[0], PartDeltaEvent)
+    assert events[0].index == 0
+
+    parts = manager.get_parts()
+    assert len(parts) == 1
+    assert isinstance(parts[0], ThinkingPart)
+    assert parts[0].content == 'initial reasoning'
+
+
+def test_simple_path_whitespace_handling():
+    """Test whitespace-only prefix with ignore_leading_whitespace in simple path (S10 → S11).
+
+    This tests the branch where whitespace before a start tag is ignored when
+    vendor_part_id=None (which routes to simple path).
+    """
+    manager = ModelResponsePartsManager()
+    thinking_tags = ('<think>', '</think>')
+
+    events = list(
+        manager.handle_text_delta(
+            vendor_part_id=None,
+            content='  \n<think>reasoning',
+            thinking_tags=thinking_tags,
+            ignore_leading_whitespace=True,
+        )
+    )
+
+    assert len(events) == 1
+    assert isinstance(events[0], PartStartEvent)
+    assert isinstance(events[0].part, ThinkingPart)
+    assert events[0].part.content == 'reasoning'
+
+    parts = manager.get_parts()
+    assert len(parts) == 1
+    assert isinstance(parts[0], ThinkingPart)
+    assert parts[0].content == 'reasoning'
+
+
+def test_simple_path_text_prefix_rejection():
+    """Test that text before start tag disables thinking tag detection in simple path (S12).
+
+    When there's non-whitespace text before the start tag, the entire content should be
+    treated as a TextPart with the tag included as literal text.
+    """
+    manager = ModelResponsePartsManager()
+    thinking_tags = ('<think>', '</think>')
+
+    events = list(
+        manager.handle_text_delta(vendor_part_id=None, content='foo<think>reasoning', thinking_tags=thinking_tags)
+    )
+
+    assert len(events) == 1
+    assert isinstance(events[0], PartStartEvent)
+    assert isinstance(events[0].part, TextPart)
+    assert events[0].part.content == 'foo<think>reasoning'
+
+    parts = manager.get_parts()
+    assert len(parts) == 1
+    assert isinstance(parts[0], TextPart)
+    assert parts[0].content == 'foo<think>reasoning'
+
+
+def test_empty_whitespace_content_with_ignore_leading_whitespace():
+    """Test that empty/whitespace content is ignored when ignore_leading_whitespace=True (line 282)."""
+    manager = ModelResponsePartsManager()
+
+    # Empty content with ignore_leading_whitespace should yield no events
+    events = list(manager.handle_text_delta(vendor_part_id='id1', content='', ignore_leading_whitespace=True))
+    assert len(events) == 0
+    assert manager.get_parts() == []
+
+    # Whitespace-only content with ignore_leading_whitespace should yield no events
+    events = list(manager.handle_text_delta(vendor_part_id='id2', content='  \n\t', ignore_leading_whitespace=True))
+    assert len(events) == 0
+    assert manager.get_parts() == []
+
+
 def test_handle_part():
     manager = ModelResponsePartsManager()
 
@@ -632,3 +727,60 @@ def test_handle_part():
     event = manager.handle_part(vendor_part_id=None, part=part3)
     assert event == snapshot(PartStartEvent(index=1, part=part3))
     assert manager.get_parts() == snapshot([part2, part3])
+
+
+def test_handle_tool_call_delta_no_vendor_id_with_non_tool_latest_part():
+    """Test handle_tool_call_delta with vendor_part_id=None when latest part is NOT a tool call (line 515->526)."""
+    manager = ModelResponsePartsManager()
+
+    # Create a TextPart first
+    for _ in manager.handle_text_delta(vendor_part_id=None, content='some text'):
+        pass
+
+    # Try to send a tool call delta with vendor_part_id=None and tool_name=None
+    # Since latest part is NOT a tool call, this should create a new incomplete tool call delta
+    event = manager.handle_tool_call_delta(vendor_part_id=None, tool_name=None, args='{"arg":')
+
+    # Since tool_name is None for a new part, we get a ToolCallPartDelta with no event
+    assert event is None
+
+    # The ToolCallPartDelta is created internally but not returned by get_parts() since it's incomplete
+    assert manager.has_incomplete_parts()
+    assert len(manager.get_parts()) == 1
+    assert isinstance(manager.get_parts()[0], TextPart)
+
+
+def test_handle_thinking_delta_raises_error_when_thinking_after_text():
+    """Test that handle_thinking_delta raises error when trying to create ThinkingPart after TextPart."""
+    manager = ModelResponsePartsManager()
+
+    # Create a TextPart first
+    for _ in manager.handle_text_delta(vendor_part_id=None, content='some text'):
+        pass
+
+    # Now try to create a ThinkingPart with vendor_part_id=None
+    # This should raise an error because thinking must come before text
+    with pytest.raises(
+        UnexpectedModelBehavior, match='Cannot create ThinkingPart after TextPart: thinking must come before text'
+    ):
+        for _ in manager.handle_thinking_delta(vendor_part_id=None, content='thinking'):
+            pass
+
+
+def test_handle_thinking_delta_create_new_part_with_no_vendor_id():
+    """Test creating new ThinkingPart when vendor_part_id is None and no parts exist yet."""
+    manager = ModelResponsePartsManager()
+
+    # Create ThinkingPart with vendor_part_id=None (no parts exist yet, so no constraint violation)
+    events = list(manager.handle_thinking_delta(vendor_part_id=None, content='thinking'))
+
+    assert len(events) == 1
+    assert isinstance(events[0], PartStartEvent)
+    assert events[0].index == 0
+
+    parts = manager.get_parts()
+    assert len(parts) == 1
+    assert parts[0] == snapshot(ThinkingPart(content='thinking'))
+
+    # Verify vendor_part_id was NOT mapped (it's None)
+    assert not manager.is_vendor_id_mapped('thinking')

tests/test_parts_manager_split_tags.py

@@ -305,3 +305,68 @@ def test_thinking_interrupted_by_incomplete_end_tag_and_vendor_switch():
     assert parts[0].content == 'thinking foo</th'
     assert isinstance(parts[1], TextPart)
     assert parts[1].content == 'new content'
+
+
+def test_split_end_tag_with_content_before():
+    """Test content before split end tag in buffered chunks (line 337)."""
+    events, parts = stream_text_deltas(['<think>', 'reasoning content</th', 'ink>'])
+
+    assert len(parts) == 1
+    assert isinstance(parts[0], ThinkingPart)
+    assert parts[0].content == 'reasoning content'
+
+    # Verify events
+    assert any(isinstance(e, PartStartEvent) and isinstance(e.part, ThinkingPart) for e in events)
+
+
+def test_split_end_tag_with_content_after():
+    """Test content after split end tag in buffered chunks (line 343)."""
+    events, parts = stream_text_deltas(['<think>', 'reasoning', '</thi', 'nk>after text'])
+
+    assert len(parts) == 2
+    assert isinstance(parts[0], ThinkingPart)
+    assert parts[0].content == 'reasoning'
+    assert isinstance(parts[1], TextPart)
+    assert parts[1].content == 'after text'
+
+    # Verify events
+    assert any(isinstance(e, PartStartEvent) and isinstance(e.part, ThinkingPart) for e in events)
+    assert any(isinstance(e, PartStartEvent) and isinstance(e.part, TextPart) for e in events)
+
+
+def test_split_end_tag_with_content_before_and_after():
+    """Test content both before and after split end tag."""
+    _, parts = stream_text_deltas(['<think>', 'reason</th', 'ink>after'])
+
+    assert len(parts) == 2
+    assert isinstance(parts[0], ThinkingPart)
+    assert parts[0].content == 'reason'
+    assert isinstance(parts[1], TextPart)
+    assert parts[1].content == 'after'
+
+
+def test_cross_path_end_tag_handling():
+    """Test end tag handling when buffering fallback delegates to simple path (C2 → S5).
+
+    This tests the scenario where buffering creates a ThinkingPart, then non-matching
+    content triggers the C2 fallback to simple path, which then handles the end tag.
+    """
+    _, parts = stream_text_deltas(['<think>initial', 'x', 'more</think>after'])
+
+    assert len(parts) == 2
+    assert isinstance(parts[0], ThinkingPart)
+    assert parts[0].content == 'initialxmore'
+    assert isinstance(parts[1], TextPart)
+    assert parts[1].content == 'after'
+
+
+def test_cross_path_bare_end_tag():
+    """Test bare end tag when buffering fallback delegates to simple path (C2 → S5).
+
+    This tests the specific branch where content equals exactly the end tag.
+    """
+    _, parts = stream_text_deltas(['<think>done', 'x', '</think>'])
+
+    assert len(parts) == 1
+    assert isinstance(parts[0], ThinkingPart)
+    assert parts[0].content == 'donex'

tests/test_streaming.py

@@ -1892,3 +1892,39 @@ async def ret_a(x: str) -> str:
             AgentRunResultEvent(result=AgentRunResult(output='{"ret_a":"a-apple"}')),
         ]
     )
+
+
+async def test_streaming_finalize_with_incomplete_thinking_tag():
+    """Test that incomplete thinking tags are flushed via finalize during streaming (lines 585-591 in models/__init__.py)."""
+
+    async def stream_with_incomplete_thinking(
+        _messages: list[ModelMessage], _agent_info: AgentInfo
+    ) -> AsyncIterator[str]:
+        # Stream incomplete thinking tag that will be buffered
+        yield '<thi'
+
+    agent = Agent(FunctionModel(stream_function=stream_with_incomplete_thinking))
+
+    events: list[AgentStreamEvent] = []
+
+    async def event_stream_handler(_ctx: RunContext[None], stream: AsyncIterable[AgentStreamEvent]):
+        async for event in stream:
+            events.append(event)
+
+    # This will trigger the finalize logic in models/__init__.py when the stream completes
+    result = await agent.run('Hello', event_stream_handler=event_stream_handler)
+
+    # The incomplete tag should be flushed as TextPart
+    assert result.output == '<thi'
+
+    # Verify that PartStartEvent and PartEndEvent were emitted from finalize
+    part_start_events = [e for e in events if isinstance(e, PartStartEvent)]
+    part_end_events = [e for e in events if isinstance(e, PartEndEvent)]
+
+    assert len(part_start_events) == 1
+    assert isinstance(part_start_events[0].part, TextPart)
+    assert part_start_events[0].part.content == '<thi'
+
+    assert len(part_end_events) == 1
+    assert isinstance(part_end_events[0].part, TextPart)
+    assert part_end_events[0].part.content == '<thi'