code-puppy 0.0.325__py3-none-any.whl → 0.0.341__py3-none-any.whl

code_puppy/agents/base_agent.py

@@ -4,7 +4,6 @@ import asyncio
 import json
 import math
 import signal
-import sys
 import threading
 import uuid
 from abc import ABC, abstractmethod
@@ -914,6 +913,11 @@ class BaseAgent(ABC):
         """
         Truncate message history to manage token usage.
 
+        Protects:
+        - The first message (system prompt) - always kept
+        - The second message if it contains a ThinkingPart (extended thinking context)
+        - The most recent messages up to protected_tokens
+
         Args:
             messages: List of messages to truncate
             protected_tokens: Number of tokens to protect
@@ -925,12 +929,30 @@ class BaseAgent(ABC):
 
         emit_info("Truncating message history to manage token usage")
         result = [messages[0]]  # Always keep the first message (system prompt)
+
+        # Check if second message exists and contains a ThinkingPart
+        # If so, protect it (extended thinking context shouldn't be lost)
+        skip_second = False
+        if len(messages) > 1:
+            second_msg = messages[1]
+            has_thinking = any(
+                isinstance(part, ThinkingPart) for part in second_msg.parts
+            )
+            if has_thinking:
+                result.append(second_msg)
+                skip_second = True
+
         num_tokens = 0
         stack = queue.LifoQueue()
 
+        # Determine which messages to consider for the recent-tokens window
+        # Skip first message (already added), and skip second if it has thinking
+        start_idx = 2 if skip_second else 1
+        messages_to_scan = messages[start_idx:]
+
         # Put messages in reverse order (most recent first) into the stack
         # but break when we exceed protected_tokens
-        for idx, msg in enumerate(reversed(messages[1:])):  # Skip the first message
+        for msg in reversed(messages_to_scan):
             num_tokens += self.estimate_tokens_for_message(msg)
             if num_tokens > protected_tokens:
                 break
@@ -1354,7 +1376,6 @@ class BaseAgent(ABC):
             ToolCallPartDelta,
         )
         from rich.console import Console
-        from rich.markdown import Markdown
         from rich.markup import escape
 
         from code_puppy.messaging.spinner import pause_all_spinners
@@ -1376,22 +1397,28 @@ class BaseAgent(ABC):
         text_parts: set[int] = set()  # Track which parts are text
         tool_parts: set[int] = set()  # Track which parts are tool calls
         banner_printed: set[int] = set()  # Track if banner was already printed
-        text_buffer: dict[int, list[str]] = {}  # Buffer text for final markdown render
         token_count: dict[int, int] = {}  # Track token count per text/tool part
         did_stream_anything = False  # Track if we streamed any content
 
+        # Termflow streaming state for text parts
+        from termflow import Parser as TermflowParser
+        from termflow import Renderer as TermflowRenderer
+
+        termflow_parsers: dict[int, TermflowParser] = {}
+        termflow_renderers: dict[int, TermflowRenderer] = {}
+        termflow_line_buffers: dict[int, str] = {}  # Buffer incomplete lines
+
         def _print_thinking_banner() -> None:
             """Print the THINKING banner with spinner pause and line clear."""
             nonlocal did_stream_anything
-            import sys
             import time
 
             from code_puppy.config import get_banner_color
 
             pause_all_spinners()
             time.sleep(0.1)  # Delay to let spinner fully clear
-            sys.stdout.write("\r\x1b[K")  # Clear line
-            sys.stdout.flush()
+            # Clear line and print newline before banner
+            console.print(" " * 50, end="\r")
             console.print()  # Newline before banner
             # Bold banner with configurable color and lightning bolt
             thinking_color = get_banner_color("thinking")
@@ -1401,21 +1428,19 @@ class BaseAgent(ABC):
                 ),
                 end="",
             )
-            sys.stdout.flush()
             did_stream_anything = True
 
         def _print_response_banner() -> None:
             """Print the AGENT RESPONSE banner with spinner pause and line clear."""
             nonlocal did_stream_anything
-            import sys
             import time
 
             from code_puppy.config import get_banner_color
 
             pause_all_spinners()
             time.sleep(0.1)  # Delay to let spinner fully clear
-            sys.stdout.write("\r\x1b[K")  # Clear line
-            sys.stdout.flush()
+            # Clear line and print newline before banner
+            console.print(" " * 50, end="\r")
             console.print()  # Newline before banner
             response_color = get_banner_color("agent_response")
             console.print(
@@ -1423,7 +1448,6 @@ class BaseAgent(ABC):
                     f"[bold white on {response_color}] AGENT RESPONSE [/bold white on {response_color}]"
                 )
             )
-            sys.stdout.flush()
             did_stream_anything = True
 
         async for event in events:
@@ -1442,13 +1466,17 @@ class BaseAgent(ABC):
                 elif isinstance(part, TextPart):
                     streaming_parts.add(event.index)
                     text_parts.add(event.index)
-                    text_buffer[event.index] = []  # Initialize buffer
-                    token_count[event.index] = 0  # Initialize token counter
-                    # Buffer initial content if present
+                    # Initialize termflow streaming for this text part
+                    termflow_parsers[event.index] = TermflowParser()
+                    termflow_renderers[event.index] = TermflowRenderer(
+                        output=console.file, width=console.width
+                    )
+                    termflow_line_buffers[event.index] = ""
+                    # Handle initial content if present
                     if part.content and part.content.strip():
-                        text_buffer[event.index].append(part.content)
-                        # Use len(content) / 3 for token estimation (more accurate than chunk counting)
-                        token_count[event.index] += len(part.content) // 3
+                        _print_response_banner()
+                        banner_printed.add(event.index)
+                        termflow_line_buffers[event.index] = part.content
                 elif isinstance(part, ToolCallPart):
                     streaming_parts.add(event.index)
                     tool_parts.add(event.index)
@@ -1464,26 +1492,29 @@ class BaseAgent(ABC):
                     delta = event.delta
                     if isinstance(delta, (TextPartDelta, ThinkingPartDelta)):
                         if delta.content_delta:
-                            # For text parts, show token counter then render at end
+                            # For text parts, stream markdown with termflow
                             if event.index in text_parts:
-                                import sys
-
                                 # Print banner on first content
                                 if event.index not in banner_printed:
                                     _print_response_banner()
                                     banner_printed.add(event.index)
-                                # Accumulate text for final markdown render
-                                text_buffer[event.index].append(delta.content_delta)
-                                # Use len(content) / 3 for token estimation
-                                token_count[event.index] += (
-                                    len(delta.content_delta) // 3
-                                )
-                                # Update token counter in place (single line)
-                                count = token_count[event.index]
-                                sys.stdout.write(
-                                    f"\r\x1b[K  ⏳ Receiving... {count} tokens"
+
+                                # Add content to line buffer
+                                termflow_line_buffers[event.index] += (
+                                    delta.content_delta
                                 )
-                                sys.stdout.flush()
+
+                                # Process complete lines
+                                parser = termflow_parsers[event.index]
+                                renderer = termflow_renderers[event.index]
+                                buffer = termflow_line_buffers[event.index]
+
+                                while "\n" in buffer:
+                                    line, buffer = buffer.split("\n", 1)
+                                    events_to_render = parser.parse_line(line)
+                                    renderer.render_all(events_to_render)
+
+                                termflow_line_buffers[event.index] = buffer
                             else:
                                 # For thinking parts, stream immediately (dim)
                                 if event.index not in banner_printed:
@@ -1492,48 +1523,51 @@ class BaseAgent(ABC):
                                 escaped = escape(delta.content_delta)
                                 console.print(f"[dim]{escaped}[/dim]", end="")
                     elif isinstance(delta, ToolCallPartDelta):
-                        import sys
-
-                        # For tool calls, show token counter (use string repr for estimation)
-                        token_count[event.index] += len(str(delta)) // 3
+                        # For tool calls, count chunks received
+                        token_count[event.index] += 1
                         # Get tool name if available
                         tool_name = getattr(delta, "tool_name_delta", "")
                         count = token_count[event.index]
                         # Display with tool wrench icon and tool name
                         if tool_name:
-                            sys.stdout.write(
-                                f"\r\x1b[K  🔧 Calling {tool_name}... {count} tokens"
+                            console.print(
+                                f"  🔧 Calling {tool_name}... {count} chunks   ",
+                                end="\r",
                             )
                         else:
-                            sys.stdout.write(
-                                f"\r\x1b[K  🔧 Calling tool... {count} tokens"
+                            console.print(
+                                f"  🔧 Calling tool... {count} chunks   ",
+                                end="\r",
                             )
-                        sys.stdout.flush()
 
             # PartEndEvent - finish the streaming with a newline
             elif isinstance(event, PartEndEvent):
                 if event.index in streaming_parts:
-                    import sys
-
-                    # For text parts, clear counter line and render markdown
+                    # For text parts, finalize termflow rendering
                     if event.index in text_parts:
-                        # Clear the token counter line
-                        sys.stdout.write("\r\x1b[K")
-                        sys.stdout.flush()
-                        # Render the final markdown nicely
-                        if event.index in text_buffer:
-                            try:
-                                final_content = "".join(text_buffer[event.index])
-                                if final_content.strip():
-                                    console.print(Markdown(final_content))
-                            except Exception:
-                                pass
-                            del text_buffer[event.index]
-                    # For tool parts, clear the token counter line
+                        # Render any remaining buffered content
+                        if event.index in termflow_parsers:
+                            parser = termflow_parsers[event.index]
+                            renderer = termflow_renderers[event.index]
+                            remaining = termflow_line_buffers.get(event.index, "")
+
+                            # Parse and render any remaining partial line
+                            if remaining.strip():
+                                events_to_render = parser.parse_line(remaining)
+                                renderer.render_all(events_to_render)
+
+                            # Finalize the parser to close any open blocks
+                            final_events = parser.finalize()
+                            renderer.render_all(final_events)
+
+                            # Clean up termflow state
+                            del termflow_parsers[event.index]
+                            del termflow_renderers[event.index]
+                            del termflow_line_buffers[event.index]
+                    # For tool parts, clear the chunk counter line
                     elif event.index in tool_parts:
-                        # Clear the token counter line
-                        sys.stdout.write("\r\x1b[K")
-                        sys.stdout.flush()
+                        # Clear the chunk counter line by printing spaces and returning
+                        console.print(" " * 50, end="\r")
                     # For thinking parts, just print newline
                     elif event.index in banner_printed:
                         console.print()  # Final newline after streaming
@@ -1952,74 +1986,35 @@ class BaseAgent(ABC):
         def graceful_sigint_handler(_sig, _frame):
             # When using keyboard-based cancel, SIGINT should be a no-op
             # (just show a hint to user about the configured cancel key)
-            import sys
-
+            # Also reset terminal to prevent bricking on Windows+uvx
             from code_puppy.keymap import get_cancel_agent_display_name
+            from code_puppy.terminal_utils import reset_windows_terminal_full
+
+            # Reset terminal state first to prevent bricking
+            reset_windows_terminal_full()
 
             cancel_key = get_cancel_agent_display_name()
-            if sys.platform == "win32":
-                # On Windows, we use keyboard listener, so SIGINT might still fire
-                # but we handle cancellation via the key listener
-                pass  # Silent on Windows - the key listener handles it
-            else:
-                emit_info(f"Use {cancel_key} to cancel the agent task.")
+            emit_info(f"Use {cancel_key} to cancel the agent task.")
 
         original_handler = None
         key_listener_stop_event = None
         _key_listener_thread = None
-        _windows_ctrl_handler = None  # Store reference to prevent garbage collection
 
         try:
-            if sys.platform == "win32":
-                # Windows: Use SetConsoleCtrlHandler for reliable Ctrl+C handling
-                import ctypes
-
-                # Define the handler function type
-                HANDLER_ROUTINE = ctypes.WINFUNCTYPE(ctypes.c_bool, ctypes.c_ulong)
-
-                def windows_ctrl_handler(ctrl_type):
-                    """Handle Windows console control events."""
-                    CTRL_C_EVENT = 0
-                    CTRL_BREAK_EVENT = 1
-
-                    if ctrl_type in (CTRL_C_EVENT, CTRL_BREAK_EVENT):
-                        # Check if we're awaiting user input
-                        if is_awaiting_user_input():
-                            return False  # Let default handler run
-
-                        # Schedule agent cancellation
-                        schedule_agent_cancel()
-                        return True  # We handled it, don't terminate
-
-                    return False  # Let other handlers process it
-
-                # Create the callback - must keep reference alive!
-                _windows_ctrl_handler = HANDLER_ROUTINE(windows_ctrl_handler)
-
-                # Register the handler
-                kernel32 = ctypes.windll.kernel32
-                if not kernel32.SetConsoleCtrlHandler(_windows_ctrl_handler, True):
-                    emit_warning("Failed to set Windows Ctrl+C handler")
-
-                # Also spawn keyboard listener for Ctrl+X (shell cancel) and other keys
-                key_listener_stop_event = threading.Event()
-                _key_listener_thread = self._spawn_ctrl_x_key_listener(
-                    key_listener_stop_event,
-                    on_escape=lambda: None,  # Ctrl+X handled by command_runner
-                    on_cancel_agent=None,  # Ctrl+C handled by SetConsoleCtrlHandler above
-                )
-            elif cancel_agent_uses_signal():
-                # Unix with Ctrl+C: Use SIGINT-based cancellation
+            if cancel_agent_uses_signal():
+                # Use SIGINT-based cancellation (default Ctrl+C behavior)
                 original_handler = signal.signal(
                     signal.SIGINT, keyboard_interrupt_handler
                 )
             else:
-                # Unix with different cancel key: Use keyboard listener
+                # Use keyboard listener for agent cancellation
+                # Set a graceful SIGINT handler that shows a hint
                 original_handler = signal.signal(signal.SIGINT, graceful_sigint_handler)
+                # Spawn keyboard listener with the cancel agent callback
                 key_listener_stop_event = threading.Event()
                 _key_listener_thread = self._spawn_ctrl_x_key_listener(
                     key_listener_stop_event,
-                    on_escape=lambda: None,
+                    on_escape=lambda: None,  # Ctrl+X handled by command_runner
                     on_cancel_agent=schedule_agent_cancel,
                 )
 
@@ -2044,17 +2039,8 @@ class BaseAgent(ABC):
             # Stop keyboard listener if it was started
             if key_listener_stop_event is not None:
                 key_listener_stop_event.set()
-
-            # Unregister Windows Ctrl handler
-            if sys.platform == "win32" and _windows_ctrl_handler is not None:
-                try:
-                    import ctypes
-
-                    kernel32 = ctypes.windll.kernel32
-                    kernel32.SetConsoleCtrlHandler(_windows_ctrl_handler, False)
-                except Exception:
-                    pass  # Best effort cleanup
-
-            # Restore original signal handler (Unix)
-            if original_handler is not None:
+            # Restore original signal handler
+            if (
+                original_handler is not None
+            ):  # Explicit None check - SIG_DFL can be 0/falsy!
                 signal.signal(signal.SIGINT, original_handler)

code_puppy/claude_cache_client.py

@@ -9,11 +9,19 @@ serialization, avoiding httpx/Pydantic internals.
 
 from __future__ import annotations
 
+import base64
 import json
-from typing import Any, Callable
+import logging
+import time
+from typing import Any, Callable, MutableMapping
 
 import httpx
 
+logger = logging.getLogger(__name__)
+
+# Refresh token if it's older than 1 hour (3600 seconds)
+TOKEN_MAX_AGE_SECONDS = 3600
+
 try:
     from anthropic import AsyncAnthropic
 except ImportError:  # pragma: no cover - optional dep
@@ -21,9 +29,108 @@ except ImportError:  # pragma: no cover - optional dep
 
 
 class ClaudeCacheAsyncClient(httpx.AsyncClient):
+    def _get_jwt_age_seconds(self, token: str | None) -> float | None:
+        """Decode a JWT and return its age in seconds.
+
+        Returns None if the token can't be decoded or has no timestamp claims.
+        Uses 'iat' (issued at) if available, otherwise calculates from 'exp'.
+        """
+        if not token:
+            return None
+
+        try:
+            # JWT format: header.payload.signature
+            # We only need the payload (second part)
+            parts = token.split(".")
+            if len(parts) != 3:
+                return None
+
+            # Decode the payload (base64url encoded)
+            payload_b64 = parts[1]
+            # Add padding if needed (base64url doesn't require padding)
+            padding = 4 - len(payload_b64) % 4
+            if padding != 4:
+                payload_b64 += "=" * padding
+
+            payload_bytes = base64.urlsafe_b64decode(payload_b64)
+            payload = json.loads(payload_bytes.decode("utf-8"))
+
+            now = time.time()
+
+            # Prefer 'iat' (issued at) claim if available
+            if "iat" in payload:
+                iat = float(payload["iat"])
+                age = now - iat
+                return age
+
+            # Fall back to calculating from 'exp' claim
+            # Assume tokens are typically valid for 1 hour
+            if "exp" in payload:
+                exp = float(payload["exp"])
+                # If exp is in the future, calculate how long until expiry
+                # and assume the token was issued 1 hour before expiry
+                time_until_exp = exp - now
+                # If token has less than 1 hour left, it's "old"
+                age = TOKEN_MAX_AGE_SECONDS - time_until_exp
+                return max(0, age)
+
+            return None
+        except Exception as exc:
+            logger.debug("Failed to decode JWT age: %s", exc)
+            return None
+
+    def _extract_bearer_token(self, request: httpx.Request) -> str | None:
+        """Extract the bearer token from request headers."""
+        auth_header = request.headers.get("Authorization") or request.headers.get(
+            "authorization"
+        )
+        if auth_header and auth_header.lower().startswith("bearer "):
+            return auth_header[7:]  # Strip "Bearer " prefix
+        return None
+
+    def _should_refresh_token(self, request: httpx.Request) -> bool:
+        """Check if the token in the request is older than 1 hour."""
+        token = self._extract_bearer_token(request)
+        if not token:
+            return False
+
+        age = self._get_jwt_age_seconds(token)
+        if age is None:
+            return False
+
+        should_refresh = age >= TOKEN_MAX_AGE_SECONDS
+        if should_refresh:
+            logger.info(
+                "JWT token is %.1f seconds old (>= %d), will refresh proactively",
+                age,
+                TOKEN_MAX_AGE_SECONDS,
+            )
+        return should_refresh
+
     async def send(
         self, request: httpx.Request, *args: Any, **kwargs: Any
     ) -> httpx.Response:  # type: ignore[override]
+        # Proactive token refresh: check JWT age before every request
+        if not request.extensions.get("claude_oauth_refresh_attempted"):
+            try:
+                if self._should_refresh_token(request):
+                    refreshed_token = self._refresh_claude_oauth_token()
+                    if refreshed_token:
+                        logger.info("Proactively refreshed token before request")
+                        # Rebuild request with new token
+                        headers = dict(request.headers)
+                        self._update_auth_headers(headers, refreshed_token)
+                        body_bytes = self._extract_body_bytes(request)
+                        request = self.build_request(
+                            method=request.method,
+                            url=request.url,
+                            headers=headers,
+                            content=body_bytes,
+                        )
+                        request.extensions["claude_oauth_refresh_attempted"] = True
+            except Exception as exc:
+                logger.debug("Error during proactive token refresh check: %s", exc)
+
         try:
             if request.url.path.endswith("/v1/messages"):
                 body_bytes = self._extract_body_bytes(request)
@@ -56,7 +163,47 @@ class ClaudeCacheAsyncClient(httpx.AsyncClient):
         except Exception:
             # Swallow wrapper errors; do not break real calls.
             pass
-        return await super().send(request, *args, **kwargs)
+        response = await super().send(request, *args, **kwargs)
+        try:
+            # Check for both 401 and 400 - Anthropic/Cloudflare may return 400 for auth errors
+            # Also check if it's a Cloudflare HTML error response
+            if response.status_code in (400, 401) and not request.extensions.get(
+                "claude_oauth_refresh_attempted"
+            ):
+                # Determine if this is an auth error (including Cloudflare HTML errors)
+                is_auth_error = response.status_code == 401
+
+                if response.status_code == 400:
+                    # Check if this is a Cloudflare HTML error
+                    is_auth_error = self._is_cloudflare_html_error(response)
+                    if is_auth_error:
+                        logger.info(
+                            "Detected Cloudflare 400 error (likely auth-related), attempting token refresh"
+                        )
+
+                if is_auth_error:
+                    refreshed_token = self._refresh_claude_oauth_token()
+                    if refreshed_token:
+                        logger.info("Token refreshed successfully, retrying request")
+                        await response.aclose()
+                        body_bytes = self._extract_body_bytes(request)
+                        headers = dict(request.headers)
+                        self._update_auth_headers(headers, refreshed_token)
+                        retry_request = self.build_request(
+                            method=request.method,
+                            url=request.url,
+                            headers=headers,
+                            content=body_bytes,
+                        )
+                        retry_request.extensions["claude_oauth_refresh_attempted"] = (
+                            True
+                        )
+                        return await super().send(retry_request, *args, **kwargs)
+                    else:
+                        logger.warning("Token refresh failed, returning original error")
+        except Exception as exc:
+            logger.debug("Error during token refresh attempt: %s", exc)
+        return response
 
     @staticmethod
     def _extract_body_bytes(request: httpx.Request) -> bytes | None:
@@ -78,6 +225,65 @@ class ClaudeCacheAsyncClient(httpx.AsyncClient):
 
         return None
 
+    @staticmethod
+    def _update_auth_headers(
+        headers: MutableMapping[str, str], access_token: str
+    ) -> None:
+        bearer_value = f"Bearer {access_token}"
+        if "Authorization" in headers or "authorization" in headers:
+            headers["Authorization"] = bearer_value
+        elif "x-api-key" in headers or "X-API-Key" in headers:
+            headers["x-api-key"] = access_token
+        else:
+            headers["Authorization"] = bearer_value
+
+    @staticmethod
+    def _is_cloudflare_html_error(response: httpx.Response) -> bool:
+        """Check if this is a Cloudflare HTML error response.
+
+        Cloudflare often returns HTML error pages with status 400 when
+        there are authentication issues.
+        """
+        # Check content type
+        content_type = response.headers.get("content-type", "")
+        if "text/html" not in content_type.lower():
+            return False
+
+        # Check if body contains Cloudflare markers
+        try:
+            # Read response body if not already consumed
+            if hasattr(response, "_content") and response._content:
+                body = response._content.decode("utf-8", errors="ignore")
+            else:
+                # Try to read the text (this might be already consumed)
+                try:
+                    body = response.text
+                except Exception:
+                    return False
+
+            # Look for Cloudflare and 400 Bad Request markers
+            body_lower = body.lower()
+            return "cloudflare" in body_lower and "400 bad request" in body_lower
+        except Exception as exc:
+            logger.debug("Error checking for Cloudflare error: %s", exc)
+            return False
+
+    def _refresh_claude_oauth_token(self) -> str | None:
+        try:
+            from code_puppy.plugins.claude_code_oauth.utils import refresh_access_token
+
+            logger.info("Attempting to refresh Claude Code OAuth token...")
+            refreshed_token = refresh_access_token(force=True)
+            if refreshed_token:
+                self._update_auth_headers(self.headers, refreshed_token)
+                logger.info("Successfully refreshed Claude Code OAuth token")
+            else:
+                logger.warning("Token refresh returned None")
+            return refreshed_token
+        except Exception as exc:
+            logger.error("Exception during token refresh: %s", exc)
+            return None
+
     @staticmethod
     def _inject_cache_control(body: bytes) -> bytes | None:
         try: