autobyteus 1.2.1__py3-none-any.whl → 1.2.3__py3-none-any.whl

autobyteus/agent/streaming/parser/states/xml_patch_file_tool_parsing_state.py

@@ -0,0 +1,328 @@
+"""
+XmlPatchFileToolParsingState: Streams <tool name="patch_file"> blocks.
+
+This state specializes the generic XmlToolParsingState to stream patch/diff content
+and capture the path for display. Argument parsing is handled later by the
+ToolInvocationAdapter.
+"""
+from typing import TYPE_CHECKING, Optional
+
+from .xml_tool_parsing_state import XmlToolParsingState
+from ..events import SegmentType
+
+if TYPE_CHECKING:
+    from ..parser_context import ParserContext
+
+
+class XmlPatchFileToolParsingState(XmlToolParsingState):
+    """
+    Streams <tool name="patch_file"> tool calls.
+    
+    This state operates similarly to XmlWriteFileToolParsingState but provides
+    a distinct type (PATCH_FILE) for unified diff content streaming.
+    
+    Key differences from write_file:
+    - Content arg name: 'patch' instead of 'content'
+    - Sentinel markers: __START_PATCH__ / __END_PATCH__
+    - Segment type: PATCH_FILE
+    """
+    
+    SEGMENT_TYPE = SegmentType.PATCH_FILE
+    START_CONTENT_MARKER = "__START_PATCH__"
+    END_CONTENT_MARKER = "__END_PATCH__"
+    CONTENT_ARG_CLOSE_TAG = "</arg>"
+    
+    def __init__(self, context: "ParserContext", opening_tag: str):
+        super().__init__(context, opening_tag)
+        if self._tool_name != "patch_file":
+            pass
+            
+        # Internal state for streaming
+        self._found_content_start = False
+        self._content_buffering = "" 
+        self._captured_path: Optional[str] = None
+        self._defer_start = True  # Defer emission until path is found
+        self._swallowing_remaining = False  # Swallow closing tags
+        self._content_mode = "seek_marker"
+        self._content_seek_buffer = ""
+        self._marker_tail = ""
+        
+    def run(self) -> None:
+        """
+        Custom run loop to stream ONLY the patch argument content.
+        """
+        from .text_state import TextState
+        
+        if self._swallowing_remaining:
+            self._handle_swallowing()
+            return
+
+        # Note: We do NOT emit start immediately anymore.
+        
+        if not self.context.has_more_chars():
+            return
+
+        chunk = self.context.consume_remaining()
+        
+        if not self._found_content_start:
+            self._content_buffering += chunk
+            
+            import re
+            
+            # 1. Try to find path if missing
+            if not self._captured_path:
+                path_match = re.search(r'<arg\s+name=["\']path["\']>([^<]+)</arg>', self._content_buffering, re.IGNORECASE)
+                if path_match:
+                    self._captured_path = path_match.group(1).strip()
+                    # Now we have path, we can emit start if we were waiting for it
+                    if self._defer_start and not self._segment_started:
+                        # Construct metadata with path
+                        meta = self._get_start_metadata()
+                        meta["path"] = self._captured_path
+                        self.context.emit_segment_start(self.SEGMENT_TYPE, **meta)
+                        self._segment_started = True
+                        self._defer_start = False
+
+            # 2. Look for patch content start (note: arg name is 'patch', not 'content')
+            match = re.search(r'<arg\s+name=["\']patch["\']>', self._content_buffering, re.IGNORECASE)
+            
+            if match:
+                self._found_content_start = True
+                end_of_tag = match.end()
+                
+                # If we still haven't emitted start (e.g. no path found but content started), emit now without path
+                if not self._segment_started:
+                    self.context.emit_segment_start(self.SEGMENT_TYPE, **self._get_start_metadata())
+                    self._segment_started = True
+                
+                # Update path in metadata if we found it late (redundant but safe)
+                if self._captured_path:
+                    self.context.update_current_segment_metadata(path=self._captured_path)
+                
+                real_content = self._content_buffering[end_of_tag:]
+                self._content_buffering = "" 
+                self._content_mode = "seek_marker"
+                self._content_seek_buffer = ""
+                self._marker_tail = ""
+                self._tail = ""
+                self._process_content_chunk(real_content)
+            else:
+                # If closing tool and still no content
+                if "</tool>" in self._content_buffering:
+                    # If start never happened, force it
+                    if not self._segment_started:
+                        self.context.emit_segment_start(self.SEGMENT_TYPE, **self._get_start_metadata())
+                        self._segment_started = True
+                        
+                    self._on_segment_complete() 
+                    self.context.emit_segment_end()
+                    self.context.transition_to(TextState(self.context))
+        else:
+            self._process_content_chunk(chunk)
+
+    def _process_content_chunk(self, chunk: str) -> None:
+        """Process content chunk, supporting optional content markers."""
+        if not chunk:
+            return
+
+        if self._content_mode == "marker":
+            self._process_marker_content(chunk)
+            return
+
+        if self._content_mode == "default":
+            self._process_default_content(chunk)
+            return
+
+        self._process_seek_marker_content(chunk)
+
+    def _process_seek_marker_content(self, chunk: str) -> None:
+        """Seek __START_PATCH__ before committing to default parsing."""
+        self._content_seek_buffer += chunk
+
+        start_idx = self._content_seek_buffer.find(self.START_CONTENT_MARKER)
+        if start_idx != -1:
+            after_start = self._content_seek_buffer[start_idx + len(self.START_CONTENT_MARKER):]
+            # Strip leading newline after marker to avoid empty first line
+            if after_start.startswith("\n"):
+                after_start = after_start[1:]
+            self._content_seek_buffer = ""
+            self._content_mode = "marker"
+            self._marker_tail = ""
+            self._tail = ""
+            if after_start:
+                self._process_marker_content(after_start)
+            return
+
+        closing_idx = self._content_seek_buffer.find(self.CONTENT_ARG_CLOSE_TAG)
+        if closing_idx != -1:
+            buffered = self._content_seek_buffer
+            self._content_seek_buffer = ""
+            self._content_mode = "default"
+            self._tail = ""
+            self._process_default_content(buffered)
+            return
+
+        stripped = self._content_seek_buffer.lstrip()
+        if stripped and not self.START_CONTENT_MARKER.startswith(stripped):
+            buffered = self._content_seek_buffer
+            self._content_seek_buffer = ""
+            self._content_mode = "default"
+            self._tail = ""
+            self._process_default_content(buffered)
+
+    def _process_default_content(self, chunk: str) -> None:
+        """Process content chunk, stripping closing tags."""
+        closing_tag = self.CONTENT_ARG_CLOSE_TAG
+        combined = self._tail + chunk
+
+        idx = combined.find(closing_tag)
+
+        if idx != -1:
+            actual_content = combined[:idx]
+            if actual_content:
+                self.context.emit_segment_content(actual_content)
+
+            # We found the end of the content argument.
+            # Switch to swallowing mode to eat </arguments></tool>
+            self._tail = ""
+            remainder = combined[idx + len(closing_tag):]
+            self._content_buffering = remainder
+            self._swallowing_remaining = True
+
+            # Immediately try to finish if we have the closing tags
+            self._handle_swallowing()
+            return
+
+        holdback_len = len(closing_tag) - 1
+        if len(combined) > holdback_len:
+            safe = combined[:-holdback_len]
+            if safe:
+                self.context.emit_segment_content(safe)
+            self._tail = combined[-holdback_len:]
+        else:
+            self._tail = combined
+
+    def _process_marker_content(self, chunk: str) -> None:
+        """Process content chunk when inside __START_PATCH__/__END_PATCH__ markers.
+        
+        The __END_PATCH__ sentinel is only valid if followed by optional whitespace
+        and then </arg>. This prevents false positives when patch content contains
+        the literal __END_PATCH__ string.
+        """
+        import re
+        
+        combined = self._marker_tail + chunk
+        end_marker = self.END_CONTENT_MARKER
+        closing_tag = self.CONTENT_ARG_CLOSE_TAG
+
+        # Priority 1: Check for the explicit end marker WITH lookahead validation
+        search_start = 0
+        while True:
+            idx = combined.find(end_marker, search_start)
+            if idx == -1:
+                break
+                
+            remainder_after_marker = combined[idx + len(end_marker):]
+            
+            # Validate: must be followed by whitespace* + </arg>
+            if re.match(r'^\s*</arg>', remainder_after_marker):
+                # Valid sentinel - emit content and transition
+                actual_content = combined[:idx]
+                if actual_content:
+                    self.context.emit_segment_content(actual_content)
+
+                self._marker_tail = ""
+                remainder = combined[idx + len(end_marker):]
+                self._content_buffering = remainder
+                self._swallowing_remaining = True
+                self._handle_swallowing()
+                return
+            elif remainder_after_marker.strip() == "":
+                # Indeterminate - need more data
+                # Hold back from idx onwards
+                if idx > 0:
+                    safe_content = combined[:idx]
+                    if safe_content:
+                        self.context.emit_segment_content(safe_content)
+                    self._marker_tail = combined[idx:]
+                else:
+                    self._marker_tail = combined
+                return
+            else:
+                # False positive
+                search_start = idx + len(end_marker)
+
+        # Priority 2: Check for closing arg tag as fallback (missing sentinel case)
+        idx_close = combined.find(closing_tag)
+        if idx_close != -1:
+            remainder_after_close = combined[idx_close + len(closing_tag):]
+            
+            # Check if followed by standard XML closure (ignoring whitespace)
+            is_valid_closure = False
+            if re.match(r'^\s*(?:</arguments>|</tool>)', remainder_after_close):
+                is_valid_closure = True
+            
+            # If we have indeterminate whitespace, we must hold back to be sure
+            elif remainder_after_close.strip() == "":
+                 self._marker_tail = combined
+                 return
+
+            if is_valid_closure:
+                actual_content = combined[:idx_close]
+                
+                # Remove trailing newline for clean file content
+                if re.search(r'\n\s*$', actual_content):
+                    actual_content = re.sub(r'\n\s*$', '', actual_content)
+                
+                if actual_content:
+                    self.context.emit_segment_content(actual_content)
+                
+                self._marker_tail = ""
+                remainder = combined[idx_close + len(closing_tag):]
+                self._content_buffering = remainder
+                self._swallowing_remaining = True
+                self._handle_swallowing()
+                return
+
+        # Holdback logic - hold back enough to detect EITHER marker OR closing_tag + context
+        # __END_PATCH__ is shorter than __END_CONTENT__ but still needs buffer
+        max_holdback = 35
+        
+        if len(combined) > max_holdback:
+            safe = combined[:-max_holdback]
+            if safe:
+                self.context.emit_segment_content(safe)
+            self._marker_tail = combined[-max_holdback:]
+        else:
+            self._marker_tail = combined
+
+    def _handle_swallowing(self) -> None:
+        """Consume stream until </tool> is found."""
+        from .text_state import TextState
+        
+        # Add any new data to buffer
+        self._content_buffering += self.context.consume_remaining()
+        
+        closing_tag = "</tool>"
+        idx = self._content_buffering.find(closing_tag)
+        
+        if idx != -1:
+            # We found the end!
+            # Anything after </tool> belongs to the next state (TextState)
+            remainder = self._content_buffering[idx + len(closing_tag):]
+            
+            self._on_segment_complete()
+            self.context.emit_segment_end()
+            if remainder:
+                # Rewind so the next state can parse the remainder (e.g., another tool tag).
+                self.context.rewind_by(len(remainder))
+            self.context.transition_to(TextState(self.context))
+        else:
+            # Nothing yet, keep swallowing
+            holdback_len = len(closing_tag) - 1
+            if len(self._content_buffering) > holdback_len:
+                # Discard safe prefix
+                self._content_buffering = self._content_buffering[-holdback_len:]
+
+    def _on_segment_complete(self) -> None:
+        return None

autobyteus/agent/streaming/parser/states/xml_run_bash_tool_parsing_state.py

@@ -0,0 +1,129 @@
+"""
+XmlRunBashToolParsingState: Streams <tool name="run_bash"> blocks.
+
+This state specializes the generic XmlToolParsingState to stream the command
+content without parsing arguments. Argument parsing is handled later by the
+ToolInvocationAdapter.
+"""
+from typing import TYPE_CHECKING
+
+from .xml_tool_parsing_state import XmlToolParsingState
+from ..events import SegmentType
+
+if TYPE_CHECKING:
+    from ..parser_context import ParserContext
+
+
+class XmlRunBashToolParsingState(XmlToolParsingState):
+    """
+    Parses <tool name="run_bash"> tool calls.
+    
+    This state operates identically to XmlToolParsingState but provides
+    a distinct type (RUN_BASH) and specialized metadata handling if needed.
+    """
+    
+    SEGMENT_TYPE = SegmentType.RUN_BASH
+    
+    def __init__(self, context: "ParserContext", opening_tag: str):
+        super().__init__(context, opening_tag)
+        if self._tool_name != "run_bash":
+            pass
+            
+        self._found_content_start = False
+        self._content_buffering = "" 
+        self._swallowing_remaining = False
+        
+    def run(self) -> None:
+        """
+        Custom run loop to stream ONLY the command argument.
+        """
+        from .text_state import TextState
+        
+        if self._swallowing_remaining:
+            self._handle_swallowing()
+            return
+
+        if not self._segment_started:
+            self.context.emit_segment_start(self.SEGMENT_TYPE, **self._get_start_metadata())
+            self._segment_started = True
+
+        if not self.context.has_more_chars():
+            return
+
+        chunk = self.context.consume_remaining()
+        
+        if not self._found_content_start:
+            self._content_buffering += chunk
+            
+            import re
+            # Look for command start
+            match = re.search(r'<arg\s+name=["\']command["\']>', self._content_buffering, re.IGNORECASE)
+            
+            if match:
+                self._found_content_start = True
+                end_of_tag = match.end()
+                
+                real_content = self._content_buffering[end_of_tag:]
+                self._content_buffering = "" 
+                self._process_content_chunk(real_content)
+            else:
+                if "</tool>" in self._content_buffering:
+                    self._on_segment_complete() 
+                    self.context.emit_segment_end()
+                    self.context.transition_to(TextState(self.context))
+        else:
+            self._process_content_chunk(chunk)
+
+    def _process_content_chunk(self, chunk: str) -> None:
+        """Process content chunk, stripping closing tags."""
+        
+        closing_tag = "</arg>"
+        combined = self._tail + chunk
+        
+        idx = combined.find(closing_tag)
+        if idx != -1:
+            actual_content = combined[:idx]
+            if actual_content:
+                self.context.emit_segment_content(actual_content)
+            
+            self._tail = ""
+            remainder = combined[idx + len(closing_tag):]
+            self._content_buffering = remainder 
+            self._swallowing_remaining = True
+            
+            self._handle_swallowing()
+            return
+            
+        holdback_len = len(closing_tag) - 1
+        if len(combined) > holdback_len:
+            safe = combined[:-holdback_len]
+            self.context.emit_segment_content(safe)
+            self._tail = combined[-holdback_len:]
+        else:
+            self._tail = combined
+
+    def _handle_swallowing(self) -> None:
+        """Consume stream until </tool> is found."""
+        from .text_state import TextState
+        
+        self._content_buffering += self.context.consume_remaining()
+        
+        closing_tag = "</tool>"
+        idx = self._content_buffering.find(closing_tag)
+        
+        if idx != -1:
+            remainder = self._content_buffering[idx + len(closing_tag):]
+            
+            self._on_segment_complete()
+            self.context.emit_segment_end()
+            if remainder:
+                # Rewind so the next state can parse the remainder (e.g., another tool tag).
+                self.context.rewind_by(len(remainder))
+            self.context.transition_to(TextState(self.context))
+        else:
+            holdback_len = len(closing_tag) - 1
+            if len(self._content_buffering) > holdback_len:
+                self._content_buffering = self._content_buffering[-holdback_len:]
+
+    def _on_segment_complete(self) -> None:
+        return None

autobyteus/agent/streaming/parser/states/xml_tag_initialization_state.py

@@ -0,0 +1,151 @@
+"""
+XmlTagInitializationState: Analyzes potential XML tags after a '<' is detected.
+
+This state buffers characters to identify special tags like <tool, <write_file,
+<run_bash and transitions to the appropriate specialized state.
+
+UNIFORM HANDOFF: All content states receive the complete opening_tag and handle
+their own initialization consistently.
+"""
+from typing import TYPE_CHECKING
+
+from .base_state import BaseState
+from ..events import SegmentType
+
+if TYPE_CHECKING:
+    from ..parser_context import ParserContext
+
+
+class XmlTagInitializationState(BaseState):
+    """
+    Analyzes a potential XML tag to determine the correct specialized state.
+    
+    This state is entered when a '<' is detected. It buffers characters
+    to identify tags like <tool, <write_file, <run_bash.
+    
+    If no known tag is detected, the buffered content is emitted as text.
+    
+    UNIFORM HANDOFF PATTERN:
+    All content-parsing states receive (context, opening_tag) and handle
+    their own buffer initialization consistently.
+    """
+    
+    # Known tag prefixes (lowercase for case-insensitive matching)
+    POSSIBLE_WRITE_FILE = "<write_file"
+    POSSIBLE_RUN_BASH = "<run_bash"
+    POSSIBLE_TOOL = "<tool"
+    
+    def __init__(self, context: "ParserContext"):
+        super().__init__(context)
+        # Consume the '<' that triggered this state
+        self.context.advance()
+        self._tag_buffer = "<"
+    
+    def run(self) -> None:
+        """
+        Buffer characters and identify the tag type.
+        
+        Transitions to specialized states or reverts to text if unknown.
+        """
+        from .text_state import TextState
+        from .custom_xml_tag_write_file_parsing_state import CustomXmlTagWriteFileParsingState
+        from .custom_xml_tag_run_bash_parsing_state import CustomXmlTagRunBashParsingState
+        from .xml_tool_parsing_state import XmlToolParsingState
+        from .xml_write_file_tool_parsing_state import XmlWriteFileToolParsingState
+        from .xml_run_bash_tool_parsing_state import XmlRunBashToolParsingState
+
+        if not self.context.has_more_chars():
+            return
+
+        start_pos = self.context.get_position()
+        end_idx = self.context.find(">", start_pos)
+
+        if end_idx == -1:
+            self._tag_buffer += self.context.consume_remaining()
+
+            lower_buffer = self._tag_buffer.lower()
+            could_be_write_file = (
+                self.POSSIBLE_WRITE_FILE.startswith(lower_buffer) or 
+                lower_buffer.startswith(self.POSSIBLE_WRITE_FILE)
+            )
+            could_be_run_bash = (
+                self.POSSIBLE_RUN_BASH.startswith(lower_buffer) or 
+                lower_buffer.startswith(self.POSSIBLE_RUN_BASH)
+            )
+            could_be_tool = (
+                self.POSSIBLE_TOOL.startswith(lower_buffer) or 
+                lower_buffer.startswith(self.POSSIBLE_TOOL)
+            )
+
+            if not (could_be_write_file or could_be_run_bash or could_be_tool):
+                self.context.append_text_segment(self._tag_buffer)
+                self.context.transition_to(TextState(self.context))
+            return
+
+        self._tag_buffer += self.context.consume(end_idx - start_pos + 1)
+        lower_buffer = self._tag_buffer.lower()
+
+        # Handle legacy <write_file> tag
+        if lower_buffer.startswith(self.POSSIBLE_WRITE_FILE):
+            if self.context.get_current_segment_type() == SegmentType.TEXT:
+                self.context.emit_segment_end()
+            self.context.transition_to(
+                CustomXmlTagWriteFileParsingState(self.context, self._tag_buffer)
+            )
+            return
+
+        if lower_buffer.startswith(self.POSSIBLE_RUN_BASH):
+            if self.context.get_current_segment_type() == SegmentType.TEXT:
+                self.context.emit_segment_end()
+            self.context.transition_to(
+                CustomXmlTagRunBashParsingState(self.context, self._tag_buffer)
+            )
+            return
+
+        if lower_buffer.startswith(self.POSSIBLE_TOOL):
+            if self.context.parse_tool_calls:
+                if self.context.get_current_segment_type() == SegmentType.TEXT:
+                    self.context.emit_segment_end()
+
+                # Extract tool name
+                import re
+                name_match = re.search(r'name\s*=\s*["\']([^"\']+)["\']', self._tag_buffer, re.IGNORECASE)
+                
+                if name_match:
+                    tool_name = name_match.group(1).lower()
+                    
+                    # --- Registry Lookup ---
+                    from ..xml_tool_parsing_state_registry import XmlToolParsingStateRegistry
+                    registry = XmlToolParsingStateRegistry()
+                    
+                    # Dispatch
+                    state_class = registry.get_state_for_tool(tool_name)
+                    if state_class:
+                        self.context.transition_to(state_class(self.context, self._tag_buffer))
+                    else:
+                        # Fallback to generic tool state
+                        self.context.transition_to(XmlToolParsingState(self.context, self._tag_buffer))
+                else:
+                    # No name found, generic fallback
+                    self.context.transition_to(XmlToolParsingState(self.context, self._tag_buffer))
+            else:
+                self.context.append_text_segment(self._tag_buffer)
+                self.context.transition_to(TextState(self.context))
+            return
+
+        self.context.append_text_segment(self._tag_buffer)
+        self.context.transition_to(TextState(self.context))
+
+    def finalize(self) -> None:
+        """
+        Called when stream ends while in this state.
+        
+        Emit any buffered content as text.
+        """
+        from .text_state import TextState
+        
+        if self._tag_buffer:
+            self.context.append_text_segment(self._tag_buffer)
+            self._tag_buffer = ""
+        
+        self.context.transition_to(TextState(self.context))

autobyteus/agent/streaming/parser/states/xml_tool_parsing_state.py

@@ -0,0 +1,63 @@
+"""
+XmlToolParsingState: Streams <tool name="...">...</tool> blocks.
+
+This state handles tool call boundaries and content streaming. Tool arguments
+are parsed later by the ToolInvocationAdapter.
+"""
+import re
+from typing import TYPE_CHECKING, Optional
+
+from .delimited_content_state import DelimitedContentState
+from ..events import SegmentType
+
+if TYPE_CHECKING:
+    from ..parser_context import ParserContext
+
+
+class XmlToolParsingState(DelimitedContentState):
+    """
+    Streams tool call blocks.
+    
+    Expected format: <tool name="..."><arg>value</arg></tool>
+    
+    Supports two argument formats:
+    1. Wrapped: <arguments><arg1>value1</arg1></arguments>
+    2. Direct: <arg1>value1</arg1><arg2>value2</arg2>
+    
+    The state:
+    1. Extracts tool name from the opening tag
+    2. Emits SEGMENT_START with tool metadata
+    3. Streams raw content for real-time display
+    4. Emits SEGMENT_END when </tool> is found
+    """
+    
+    # Pattern to extract tool name
+    NAME_PATTERN = re.compile(r'name\s*=\s*["\']([^"\']+)["\']', re.IGNORECASE)
+    CLOSING_TAG = "</tool>"
+    SEGMENT_TYPE = SegmentType.TOOL_CALL
+    def __init__(self, context: "ParserContext", opening_tag: str):
+        """
+        Initialize the tool parsing state.
+        
+        Args:
+            context: The parser context.
+            opening_tag: The complete opening tag (e.g., '<tool name="read_file">').
+        """
+        super().__init__(context, opening_tag)
+        self._tool_name: Optional[str] = None
+        
+        # Extract tool name from opening tag
+        match = self.NAME_PATTERN.search(opening_tag)
+        if match:
+            self._tool_name = match.group(1)
+
+    def _can_start_segment(self) -> bool:
+        return self._tool_name is not None
+
+    def _get_start_metadata(self) -> dict:
+        return {"tool_name": self._tool_name} if self._tool_name else {}
+
+    def _on_segment_complete(self) -> None:
+        return None
+
+    # finalize inherited from DelimitedContentState