fast-agent-mcp 0.3.16__py3-none-any.whl → 0.3.17__py3-none-any.whl

fast_agent/ui/enhanced_prompt.py

@@ -9,6 +9,7 @@ import shlex
 import subprocess
 import tempfile
 from importlib.metadata import version
+from pathlib import Path
 from typing import TYPE_CHECKING, Any, Dict, List, Optional
 
 from prompt_toolkit import PromptSession
@@ -23,7 +24,7 @@ from rich import print as rich_print
 from fast_agent.agents.agent_types import AgentType
 from fast_agent.constants import FAST_AGENT_ERROR_CHANNEL, FAST_AGENT_REMOVED_METADATA_CHANNEL
 from fast_agent.core.exceptions import PromptExitError
-from fast_agent.llm.model_info import get_model_info
+from fast_agent.llm.model_info import ModelInfo
 from fast_agent.mcp.types import McpAgentProtocol
 from fast_agent.ui.mcp_display import render_mcp_status
 
@@ -723,18 +724,9 @@ async def get_enhanced_input(
             # Build TDV capability segment based on model database
             info = None
             if llm:
-                try:
-                    info = get_model_info(llm)
-                except TypeError:
-                    info = None
+                info = ModelInfo.from_llm(llm)
             if not info and model_name:
-                try:
-                    info = get_model_info(model_name)
-                except TypeError:
-                    info = None
-                except Exception as exc:
-                    print(f"[toolbar debug] get_model_info failed for '{agent_name}': {exc}")
-                    info = None
+                info = ModelInfo.from_name(model_name)
 
             # Default to text-only if info resolution fails for any reason
             t, d, v = (True, False, False)
@@ -965,6 +957,28 @@ async def get_enhanced_input(
         if shell_enabled:
             modes_display = ", ".join(shell_access_modes or ("direct",))
             shell_display = f"{modes_display}, {shell_name}" if shell_name else modes_display
+
+            # Add working directory info
+            shell_runtime = getattr(shell_agent, "_shell_runtime", None)
+            if shell_runtime:
+                working_dir = shell_runtime.working_directory()
+                try:
+                    # Try to show relative to cwd for cleaner display
+                    working_dir_display = str(working_dir.relative_to(Path.cwd()))
+                    if working_dir_display == ".":
+                        # Show last 2 parts of the path (e.g., "source/fast-agent")
+                        parts = Path.cwd().parts
+                        if len(parts) >= 2:
+                            working_dir_display = "/".join(parts[-2:])
+                        elif len(parts) == 1:
+                            working_dir_display = parts[0]
+                        else:
+                            working_dir_display = str(Path.cwd())
+                except ValueError:
+                    # If not relative to cwd, show absolute path
+                    working_dir_display = str(working_dir)
+                shell_display = f"{shell_display} | cwd: {working_dir_display}"
+
             rich_print(f"[yellow]Shell Access ({shell_display})[/yellow]")
 
         rich_print()

fast_agent/ui/markdown_helpers.py

@@ -0,0 +1,104 @@
+from __future__ import annotations
+
+from typing import Any, Iterable, Iterator
+
+HTML_ESCAPE_CHARS: dict[str, str] = {
+    "&": "&",
+    "<": "&lt;",
+    ">": "&gt;",
+    '"': "&quot;",
+    "'": "&#39;",
+}
+
+
+def _flatten_tokens(tokens: Iterable[Any]) -> Iterator[Any]:
+    """Recursively flatten markdown-it token trees."""
+    for token in tokens:
+        yield token
+        if token.children:
+            yield from _flatten_tokens(token.children)
+
+
+def prepare_markdown_content(content: str, escape_xml: bool = True) -> str:
+    """Prepare content for markdown rendering, escaping HTML/XML outside code blocks."""
+    if not escape_xml or not isinstance(content, str):
+        return content
+
+    from markdown_it import MarkdownIt
+
+    parser = MarkdownIt()
+    try:
+        tokens = parser.parse(content)
+    except Exception:
+        result = content
+        for char, replacement in HTML_ESCAPE_CHARS.items():
+            result = result.replace(char, replacement)
+        return result
+
+    protected_ranges: list[tuple[int, int]] = []
+    lines = content.split("\n")
+
+    for token in _flatten_tokens(tokens):
+        if token.map is not None:
+            if token.type in ("fence", "code_block"):
+                start_line = token.map[0]
+                end_line = token.map[1]
+                start_pos = sum(len(line) + 1 for line in lines[:start_line])
+                end_pos = sum(len(line) + 1 for line in lines[:end_line])
+                protected_ranges.append((start_pos, end_pos))
+
+        if token.type == "code_inline":
+            code_content = token.content
+            if code_content:
+                pattern = f"`{code_content}`"
+                start = 0
+                while True:
+                    pos = content.find(pattern, start)
+                    if pos == -1:
+                        break
+                    in_protected = any(s <= pos < e for s, e in protected_ranges)
+                    if not in_protected:
+                        protected_ranges.append((pos, pos + len(pattern)))
+                    start = pos + len(pattern)
+
+    import re
+
+    fence_pattern = r"^```"
+    fences = list(re.finditer(fence_pattern, content, re.MULTILINE))
+
+    if len(fences) % 2 == 1:
+        last_fence_pos = fences[-1].start()
+        in_protected = any(s <= last_fence_pos < e for s, e in protected_ranges)
+        if not in_protected:
+            protected_ranges.append((last_fence_pos, len(content)))
+
+    protected_ranges.sort(key=lambda x: x[0])
+
+    merged_ranges: list[tuple[int, int]] = []
+    for start, end in protected_ranges:
+        if merged_ranges and start <= merged_ranges[-1][1]:
+            merged_ranges[-1] = (merged_ranges[-1][0], max(end, merged_ranges[-1][1]))
+        else:
+            merged_ranges.append((start, end))
+
+    result_segments: list[str] = []
+    last_end = 0
+
+    for start, end in merged_ranges:
+        unprotected_text = content[last_end:start]
+        for char, replacement in HTML_ESCAPE_CHARS.items():
+            unprotected_text = unprotected_text.replace(char, replacement)
+        result_segments.append(unprotected_text)
+
+        result_segments.append(content[start:end])
+        last_end = end
+
+    remainder_text = content[last_end:]
+    for char, replacement in HTML_ESCAPE_CHARS.items():
+        remainder_text = remainder_text.replace(char, replacement)
+    result_segments.append(remainder_text)
+
+    return "".join(result_segments)
+
+
+__all__ = ["HTML_ESCAPE_CHARS", "prepare_markdown_content"]

fast_agent/ui/markdown_truncator.py

@@ -57,6 +57,7 @@ class CodeBlockInfo:
     end_pos: int
     fence_line: int
     language: str
+    fence_text: str | None
     token: Token
 
 
@@ -89,6 +90,13 @@ class MarkdownTruncator:
         self._last_full_text: str | None = None
         self._last_truncated_text: str | None = None
         self._last_terminal_height: int | None = None
+        # Markdown parse cache
+        self._cache_source: str | None = None
+        self._cache_tokens: List[Token] | None = None
+        self._cache_lines: List[str] | None = None
+        self._cache_safe_points: List[TruncationPoint] | None = None
+        self._cache_code_blocks: List[CodeBlockInfo] | None = None
+        self._cache_tables: List[TableInfo] | None = None
 
     def truncate(
         self,
@@ -285,6 +293,18 @@ class MarkdownTruncator:
 
         return truncated_text
 
+    def _ensure_parse_cache(self, text: str) -> None:
+        if self._cache_source == text:
+            return
+
+        tokens = self.parser.parse(text)
+        self._cache_source = text
+        self._cache_tokens = tokens
+        self._cache_lines = text.split("\n")
+        self._cache_safe_points = None
+        self._cache_code_blocks = None
+        self._cache_tables = None
+
     def _find_safe_truncation_points(self, text: str) -> List[TruncationPoint]:
         """Find safe positions to truncate at (block boundaries).
 
@@ -294,11 +314,16 @@ class MarkdownTruncator:
         Returns:
             List of TruncationPoint objects representing safe truncation positions.
         """
-        tokens = self.parser.parse(text)
-        safe_points = []
+        self._ensure_parse_cache(text)
+        if self._cache_safe_points is not None:
+            return list(self._cache_safe_points)
 
-        # Don't flatten - we need to process top-level tokens
-        lines = text.split("\n")
+        assert self._cache_tokens is not None
+        assert self._cache_lines is not None
+
+        safe_points: List[TruncationPoint] = []
+        tokens = self._cache_tokens
+        lines = self._cache_lines
 
         for token in tokens:
             # We're interested in block-level tokens with map information
@@ -319,13 +344,13 @@ class MarkdownTruncator:
                             is_closing=(token.nesting == 0),  # Self-closing or block end
                         )
                     )
-
-        return safe_points
+        self._cache_safe_points = safe_points
+        return list(safe_points)
 
     def _get_code_block_info(self, text: str) -> List[CodeBlockInfo]:
         """Extract code block positions and metadata using markdown-it.
 
-        Uses same technique as _prepare_markdown_content in console_display.py:
+        Uses same technique as prepare_markdown_content in markdown_helpers.py:
         parse once with markdown-it, extract exact positions from tokens.
 
         Args:
@@ -334,9 +359,16 @@ class MarkdownTruncator:
         Returns:
             List of CodeBlockInfo objects with position and language metadata.
         """
-        tokens = self.parser.parse(text)
-        lines = text.split("\n")
-        code_blocks = []
+        self._ensure_parse_cache(text)
+        if self._cache_code_blocks is not None:
+            return list(self._cache_code_blocks)
+
+        assert self._cache_tokens is not None
+        assert self._cache_lines is not None
+
+        tokens = self._cache_tokens
+        lines = self._cache_lines
+        code_blocks: List[CodeBlockInfo] = []
 
         for token in self._flatten_tokens(tokens):
             if token.type in ("fence", "code_block") and token.map:
@@ -345,6 +377,9 @@ class MarkdownTruncator:
                 start_pos = sum(len(line) + 1 for line in lines[:start_line])
                 end_pos = sum(len(line) + 1 for line in lines[:end_line])
                 language = token.info or "" if hasattr(token, "info") else ""
+                fence_text: str | None = None
+                if token.type == "fence":
+                    fence_text = lines[start_line] if 0 <= start_line < len(lines) else None
 
                 code_blocks.append(
                     CodeBlockInfo(
@@ -352,11 +387,35 @@ class MarkdownTruncator:
                         end_pos=end_pos,
                         fence_line=start_line,
                         language=language,
+                        fence_text=fence_text,
                         token=token,
                     )
                 )
+        self._cache_code_blocks = code_blocks
+        return list(code_blocks)
+
+    def _build_code_block_prefix(self, block: CodeBlockInfo) -> str | None:
+        """Construct the opening fence text for a code block if applicable."""
+        token = block.token
+
+        if token.type == "fence":
+            if block.fence_text:
+                fence_line = block.fence_text
+            else:
+                markup = getattr(token, "markup", "") or "```"
+                info = (getattr(token, "info", "") or "").strip()
+                fence_line = f"{markup}{info}" if info else markup
+            return fence_line if fence_line.endswith("\n") else fence_line + "\n"
 
-        return code_blocks
+        if token.type == "code_block":
+            info = (getattr(token, "info", "") or "").strip()
+            if info:
+                return f"```{info}\n"
+            if block.language:
+                return f"```{block.language}\n"
+            return "```\n"
+
+        return None
 
     def _get_table_info(self, text: str) -> List[TableInfo]:
         """Extract table positions and metadata using markdown-it.
@@ -370,9 +429,16 @@ class MarkdownTruncator:
         Returns:
             List of TableInfo objects with position and header metadata.
         """
-        tokens = self.parser.parse(text)
-        lines = text.split("\n")
-        tables = []
+        self._ensure_parse_cache(text)
+        if self._cache_tables is not None:
+            return list(self._cache_tables)
+
+        assert self._cache_tokens is not None
+        assert self._cache_lines is not None
+
+        tokens = self._cache_tokens
+        lines = self._cache_lines
+        tables: List[TableInfo] = []
 
         for i, token in enumerate(tokens):
             if token.type == "table_open" and token.map:
@@ -435,8 +501,8 @@ class MarkdownTruncator:
                             header_lines=header_lines,
                         )
                     )
-
-        return tables
+        self._cache_tables = tables
+        return list(tables)
 
     def _find_best_truncation_point(
         self,
@@ -571,8 +637,8 @@ class MarkdownTruncator:
         # If truncation happened after the fence line, it scrolled off
         if truncation_point.char_position > code_block.start_pos:
             # Check if fence is already at the beginning (avoid duplicates)
-            fence = f"```{code_block.language}\n"
-            if not truncated_text.startswith(fence):
+            fence = self._build_code_block_prefix(code_block)
+            if fence and not truncated_text.startswith(fence):
                 # Fence scrolled off - prepend it
                 return fence + truncated_text
 
@@ -611,10 +677,8 @@ class MarkdownTruncator:
                 # Truncated within this code block
                 # Simple check: did truncation remove the fence?
                 if truncation_pos > block.start_pos:
-                    # Check if fence is already at the beginning (avoid duplicates)
-                    fence = f"```{block.language}\n"
-                    if not truncated_text.startswith(fence):
-                        # Fence scrolled off - prepend it
+                    fence = self._build_code_block_prefix(block)
+                    if fence and not truncated_text.startswith(fence):
                         return fence + truncated_text
                 # Fence still on screen or already prepended
                 return truncated_text
@@ -875,32 +939,26 @@ class MarkdownTruncator:
         if not truncated_text or truncated_text == original_text:
             return truncated_text
 
+        original_fragment = truncated_text
+
         # Find where the truncated text starts in the original
-        truncation_pos = original_text.rfind(truncated_text)
+        truncation_pos = original_text.rfind(original_fragment)
         if truncation_pos == -1:
-            # Can't find it, return as-is
-            return truncated_text
+            truncation_pos = max(0, len(original_text) - len(original_fragment))
+
+        code_blocks = self._get_code_block_info(original_text)
+        active_block = None
+        for block in code_blocks:
+            if block.start_pos <= truncation_pos < block.end_pos:
+                active_block = block
+
+        if active_block:
+            fence = self._build_code_block_prefix(active_block)
+            if fence and not truncated_text.startswith(fence):
+                truncated_text = fence + truncated_text
 
-        # Check for incomplete code blocks
-        original_fence_count = original_text[:truncation_pos].count('```')
-
-        # If we removed an odd number of fences, we're inside a code block
-        if original_fence_count % 2 == 1:
-            # Find the last opening fence before truncation point
-            import re
-            before_truncation = original_text[:truncation_pos]
-            fences = list(re.finditer(r'^```(\w*)', before_truncation, re.MULTILINE))
-            if fences:
-                last_fence = fences[-1]
-                language = last_fence.group(1) if last_fence.group(1) else ''
-                fence = f'```{language}\n'
-                if not truncated_text.startswith(fence):
-                    truncated_text = fence + truncated_text
-
-        # Check for incomplete tables
-        # Only if we're not inside a code block
-        if original_fence_count % 2 == 0 and '|' in truncated_text:
-            # Use the existing table header restoration logic
+        # Check for incomplete tables when not inside a code block
+        if active_block is None and '|' in truncated_text:
             tables = self._get_table_info(original_text)
             for table in tables:
                 if table.thead_end_pos <= truncation_pos < table.tbody_end_pos:

fast_agent/ui/message_primitives.py

@@ -0,0 +1,50 @@
+from __future__ import annotations
+
+from enum import Enum
+
+
+class MessageType(Enum):
+    """Types of messages that can be displayed."""
+
+    USER = "user"
+    ASSISTANT = "assistant"
+    SYSTEM = "system"
+    TOOL_CALL = "tool_call"
+    TOOL_RESULT = "tool_result"
+
+
+MESSAGE_CONFIGS: dict[MessageType, dict[str, str]] = {
+    MessageType.USER: {
+        "block_color": "blue",
+        "arrow": "▶",
+        "arrow_style": "dim blue",
+        "highlight_color": "blue",
+    },
+    MessageType.ASSISTANT: {
+        "block_color": "green",
+        "arrow": "◀",
+        "arrow_style": "dim green",
+        "highlight_color": "bright_green",
+    },
+    MessageType.SYSTEM: {
+        "block_color": "yellow",
+        "arrow": "●",
+        "arrow_style": "dim yellow",
+        "highlight_color": "bright_yellow",
+    },
+    MessageType.TOOL_CALL: {
+        "block_color": "magenta",
+        "arrow": "◀",
+        "arrow_style": "dim magenta",
+        "highlight_color": "magenta",
+    },
+    MessageType.TOOL_RESULT: {
+        "block_color": "magenta",
+        "arrow": "▶",
+        "arrow_style": "dim magenta",
+        "highlight_color": "magenta",
+    },
+}
+
+
+__all__ = ["MessageType", "MESSAGE_CONFIGS"]