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

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

@@ -0,0 +1,343 @@
+"""
+XmlWriteFileToolParsingState: Streams <tool name="write_file"> blocks.
+
+This state specializes the generic XmlToolParsingState to stream file 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 XmlWriteFileToolParsingState(XmlToolParsingState):
+    """
+    Streams <tool name="write_file"> tool calls.
+    
+    This state operates identically to XmlToolParsingState but provides
+    a distinct type (WRITE_FILE) and specialized metadata handling if needed.
+    """
+    
+    SEGMENT_TYPE = SegmentType.WRITE_FILE
+    START_CONTENT_MARKER = "__START_CONTENT__"
+    END_CONTENT_MARKER = "__END_CONTENT__"
+    CONTENT_ARG_CLOSE_TAG = "</arg>"
+    
+    def __init__(self, context: "ParserContext", opening_tag: str):
+        super().__init__(context, opening_tag)
+        if self._tool_name != "write_file":
+            pass
+            
+        # Internal state for streaming
+        self._found_content_start = False
+        self._content_buffering = "" 
+        self._captured_path: Optional[str] = None
+        self._defer_start = True # New flag to defer emission
+        self._swallowing_remaining = False # New flag to 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 content argument.
+        """
+        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 content start
+            match = re.search(r'<arg\s+name=["\']content["\']>', 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_CONTENT__ 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.
+            # Instead of stopping, we 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_CONTENT__/__END_CONTENT__ markers.
+        
+        The __END_CONTENT__ sentinel is only valid if followed by optional whitespace
+        and then </arg>. This prevents false positives when file content contains
+        the literal __END_CONTENT__ 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
+        # We need to find __END_CONTENT__ that is followed by whitespace* + </arg>
+        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 to decide if this is valid
+                # Hold back from idx onwards
+                if idx > 0:
+                    # Emit safe content before the potential marker
+                    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 - this __END_CONTENT__ is part of the file content
+                # Continue searching for the next occurrence
+                search_start = idx + len(end_marker)
+
+        # Priority 2: Check for closing arg tag as fallback (missing sentinel case)
+        # Only treat </arg> as terminator if it looks like the actual XML structure end
+        # (i.e. followed by </arguments> or </tool>)
+        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)
+            # We match if we see the start of the next tag, OR if we have only whitespace (ambiguous - wait)
+            
+            is_valid_closure = False
+            # If we see the next tag immediately start
+            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() == "":
+                 # We can't decide yet. Hold back everything from existing tag start.
+                 # But we can allow partial emit of previous content if we separate it.
+                 # Let's just hold back the whole combined tail to be safe/simple.
+                 self._marker_tail = combined
+                 return
+
+            if is_valid_closure:
+                actual_content = combined[:idx_close]
+                
+                # User request: remove the last \n for the file content
+                # This specifically handles the indented </arg> case logic
+                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
+        # We need to hold back enough to detect EITHER marker OR closing_tag + context
+        # </arg> (6) + \s + </tool> (7) ~= 15-20 chars holdback
+        # But to be safe against splitting </arguments>, let's hold back ~20 chars.
+        # Also need to account for __END_CONTENT__ (16 chars) + whitespace + </arg> (6) = ~25 chars
+        
+        max_holdback = 35 # Safe buffer for regex lookahead
+        
+        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!
+            # We are done with this tool.
+            
+            # 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 (clearing buffer to avoid memory issues if valid)
+            # But we need to keep a holdback in case </tool> is split?
+            # </tool> is 7 chars.
+            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/strategies/__init__.py

@@ -0,0 +1,17 @@
+"""
+Detection strategies for streaming parser.
+"""
+
+from .base import DetectionStrategy
+from .json_tool_strategy import JsonToolStrategy
+from .sentinel_strategy import SentinelStrategy
+from .xml_tag_strategy import XmlTagStrategy
+from .registry import create_detection_strategies
+
+__all__ = [
+    "DetectionStrategy",
+    "JsonToolStrategy",
+    "SentinelStrategy",
+    "XmlTagStrategy",
+    "create_detection_strategies",
+]

autobyteus/agent/streaming/parser/strategies/base.py

@@ -0,0 +1,24 @@
+"""
+Detection strategy interfaces for streaming parser.
+"""
+from __future__ import annotations
+
+from typing import Protocol, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from ..parser_context import ParserContext
+    from ..states.base_state import BaseState
+
+
+class DetectionStrategy(Protocol):
+    """Protocol for detection strategies."""
+
+    name: str
+
+    def next_marker(self, context: "ParserContext", start_pos: int) -> int:
+        """Return next marker index or -1 if not applicable."""
+        ...
+
+    def create_state(self, context: "ParserContext") -> "BaseState":
+        """Create the state to handle parsing after this marker."""
+        ...

autobyteus/agent/streaming/parser/strategies/json_tool_strategy.py

@@ -0,0 +1,26 @@
+"""
+Detection strategy for JSON tool calls.
+"""
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from ..parser_context import ParserContext
+    from ..states.base_state import BaseState
+
+
+class JsonToolStrategy:
+    name = "json_tool"
+
+    def next_marker(self, context: "ParserContext", start_pos: int) -> int:
+        if not context.parse_tool_calls:
+            return -1
+        next_curly = context.find("{", start_pos)
+        next_bracket = context.find("[", start_pos)
+        candidates = [idx for idx in (next_curly, next_bracket) if idx != -1]
+        return min(candidates) if candidates else -1
+
+    def create_state(self, context: "ParserContext") -> "BaseState":
+        from ..states.json_initialization_state import JsonInitializationState
+        return JsonInitializationState(context)

autobyteus/agent/streaming/parser/strategies/registry.py

@@ -0,0 +1,28 @@
+"""
+Registry for detection strategies.
+"""
+from __future__ import annotations
+
+from typing import Dict, List
+
+from .base import DetectionStrategy
+from .json_tool_strategy import JsonToolStrategy
+from .sentinel_strategy import SentinelStrategy
+from .xml_tag_strategy import XmlTagStrategy
+
+
+_STRATEGY_REGISTRY: Dict[str, DetectionStrategy] = {
+    "sentinel": SentinelStrategy(),
+    "xml_tag": XmlTagStrategy(),
+    "json_tool": JsonToolStrategy(),
+}
+
+
+def create_detection_strategies(strategy_order: List[str]) -> List[DetectionStrategy]:
+    """Create ordered detection strategies."""
+    strategies: List[DetectionStrategy] = []
+    for name in strategy_order:
+        strategy = _STRATEGY_REGISTRY.get(name)
+        if strategy:
+            strategies.append(strategy)
+    return strategies

autobyteus/agent/streaming/parser/strategies/sentinel_strategy.py

@@ -0,0 +1,23 @@
+"""
+Detection strategy for sentinel-formatted segments.
+"""
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from ..sentinel_format import START_MARKER
+
+if TYPE_CHECKING:
+    from ..parser_context import ParserContext
+    from ..states.base_state import BaseState
+
+
+class SentinelStrategy:
+    name = "sentinel"
+
+    def next_marker(self, context: "ParserContext", start_pos: int) -> int:
+        return context.find(START_MARKER, start_pos)
+
+    def create_state(self, context: "ParserContext") -> "BaseState":
+        from ..states.sentinel_initialization_state import SentinelInitializationState
+        return SentinelInitializationState(context)

autobyteus/agent/streaming/parser/strategies/xml_tag_strategy.py

@@ -0,0 +1,21 @@
+"""
+Detection strategy for XML-like tags such as <file>, <bash>, <tool>, <!doctype>.
+"""
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from ..parser_context import ParserContext
+    from ..states.base_state import BaseState
+
+
+class XmlTagStrategy:
+    name = "xml_tag"
+
+    def next_marker(self, context: "ParserContext", start_pos: int) -> int:
+        return context.find("<", start_pos)
+
+    def create_state(self, context: "ParserContext") -> "BaseState":
+        from ..states.xml_tag_initialization_state import XmlTagInitializationState
+        return XmlTagInitializationState(context)

autobyteus/agent/streaming/parser/stream_scanner.py

@@ -0,0 +1,167 @@
+"""
+StreamScanner: A class to manage a string buffer and cursor position for sequential reading.
+
+This encapsulates the logic of navigating a stream of text, preventing direct
+manipulation of the cursor and buffer from multiple state classes.
+"""
+from typing import Optional
+
+
+class StreamScanner:
+    """
+    Manages a string buffer and a cursor position for sequential reading.
+    
+    This class encapsulates navigation logic for streaming text, providing
+    a clean interface for state machine states to read characters without
+    directly manipulating the buffer.
+    """
+    
+    def __init__(self, initial_buffer: str = ""):
+        """Initialize the scanner with an optional initial buffer."""
+        self._buffer: str = initial_buffer
+        self._pos: int = 0
+
+    def append(self, text: str) -> None:
+        """Append more text to the end of the buffer."""
+        self._buffer += text
+
+    def peek(self) -> Optional[str]:
+        """
+        Look at the character at the current cursor position without advancing.
+        
+        Returns:
+            The character at the cursor, or None if at the end.
+        """
+        if self._pos < len(self._buffer):
+            return self._buffer[self._pos]
+        return None
+
+    def advance(self) -> None:
+        """Move the cursor forward by one position."""
+        if self.has_more_chars():
+            self._pos += 1
+
+    def advance_by(self, count: int) -> None:
+        """
+        Move the cursor forward by a specified number of positions.
+        
+        Args:
+            count: The number of characters to advance.
+        """
+        self._pos = min(len(self._buffer), self._pos + count)
+
+    def has_more_chars(self) -> bool:
+        """
+        Check if there are more characters to read from the buffer.
+        
+        Returns:
+            True if the cursor is not at the end of the buffer.
+        """
+        return self._pos < len(self._buffer)
+
+    def substring(self, start: int, end: Optional[int] = None) -> str:
+        """
+        Extract a substring from the buffer.
+        
+        Args:
+            start: The starting index.
+            end: The ending index (exclusive). If None, reads to end of buffer.
+            
+        Returns:
+            The extracted substring.
+        """
+        if end is None:
+            return self._buffer[start:]
+        return self._buffer[start:end]
+
+    def find(self, sub: str, start: Optional[int] = None) -> int:
+        """
+        Find a substring in the buffer starting from a position.
+
+        Args:
+            sub: The substring to search for.
+            start: Optional start index. Defaults to current position.
+
+        Returns:
+            The index of the substring, or -1 if not found.
+        """
+        if start is None:
+            start = self._pos
+        return self._buffer.find(sub, start)
+
+    def consume(self, count: int) -> str:
+        """
+        Consume a number of characters from the current position.
+
+        Args:
+            count: Number of characters to consume.
+
+        Returns:
+            The consumed substring.
+        """
+        if count <= 0:
+            return ""
+        end = min(len(self._buffer), self._pos + count)
+        result = self._buffer[self._pos:end]
+        self._pos = end
+        return result
+
+    def consume_remaining(self) -> str:
+        """
+        Consume all remaining characters in the buffer.
+
+        Returns:
+            The remaining substring from current position to end.
+        """
+        if self._pos >= len(self._buffer):
+            return ""
+        result = self._buffer[self._pos:]
+        self._pos = len(self._buffer)
+        return result
+
+    def get_position(self) -> int:
+        """
+        Return the current zero-based position of the cursor.
+        
+        Returns:
+            The current cursor position.
+        """
+        return self._pos
+
+    def get_buffer_length(self) -> int:
+        """
+        Return the total length of the internal buffer.
+        
+        Returns:
+            The length of the buffer.
+        """
+        return len(self._buffer)
+
+    def set_position(self, position: int) -> None:
+        """
+        Set the cursor to a specific position.
+        
+        Args:
+            position: The new cursor position (clamped to valid range).
+        """
+        self._pos = max(0, min(len(self._buffer), position))
+
+    def compact(self, min_prefix: int = 65536) -> None:
+        """
+        Compact the buffer by dropping consumed prefix data.
+
+        If all data is consumed, clears the buffer entirely.
+        Otherwise, drops the consumed prefix when it exceeds min_prefix.
+
+        Args:
+            min_prefix: Minimum consumed prefix length to trigger compaction.
+        """
+        if self._pos == 0:
+            return
+        if self._pos >= len(self._buffer):
+            self._buffer = ""
+            self._pos = 0
+            return
+        if self._pos >= min_prefix:
+            self._buffer = self._buffer[self._pos:]
+            self._pos = 0