kiwi-code 0.0.4__py3-none-any.whl

kiwi_tui/screens/dashboard.py

@@ -0,0 +1,608 @@
+"""Dashboard screen for Autobots TUI."""
+
+from textual.app import ComposeResult
+from textual.screen import Screen
+from textual.widgets import Header, Footer, Input, Static
+from textual.containers import Vertical, VerticalScroll
+from textual.worker import Worker, WorkerState
+from loguru import logger
+import json
+import asyncio
+import re
+import html
+
+
+class DashboardScreen(Screen):
+    """Main dashboard screen showing overview and stats."""
+
+    DEFAULT_ACTION_ID = "69295914ccbcd104b2e7446f"
+
+    def __init__(self, *args, **kwargs):
+        """Initialize the dashboard screen."""
+        super().__init__(*args, **kwargs)
+        self.current_action_id = self.DEFAULT_ACTION_ID
+        self.current_run_id = None  # Track the current conversation run_id
+        self.active_stream_tasks = []  # Track active SSE stream tasks to cancel them
+
+    CSS = """
+    DashboardScreen {
+        background: $surface;
+    }
+
+    #messages {
+        height: 1fr;
+        width: 100%;
+    }
+
+    .message {
+        width: 100%;
+        height: auto;
+        padding: 0 1;
+        margin: 0;
+    }
+
+    .user-message {
+        color: #00d4d4;
+    }
+
+    .assistant-message {
+        color: $text;
+    }
+
+    .error-message {
+        color: #ff5555;
+    }
+
+    .info-message {
+        color: #4db8e8;
+        text-style: italic;
+    }
+
+    #input-bar {
+        dock: bottom;
+        height: 3;
+        width: 100%;
+    }
+
+    Input {
+        width: 100%;
+        border: solid #00d4d4;
+    }
+
+    Input:focus {
+        border: solid #00ffff;
+    }
+    """
+
+    def compose(self) -> ComposeResult:
+        """Compose dashboard widgets."""
+        yield Header()
+
+        with VerticalScroll(id="messages"):
+            pass
+
+        with Vertical(id="input-bar"):
+            yield Input(placeholder="Message...", id="chat-input")
+
+        yield Footer()
+
+    def on_mount(self) -> None:
+        """Called when screen is mounted."""
+        self.query_one("#chat-input", Input).focus()
+
+    def on_input_submitted(self, event: Input.Submitted) -> None:
+        """Handle message submission."""
+        message = event.value.strip()
+        if not message:
+            return
+
+        # Clear input immediately for responsive feel
+        event.input.value = ""
+
+        # Handle "/" commands
+        if message.startswith("/"):
+            self.handle_slash_command(message)
+            return
+
+        # Add user message immediately to chat
+        self.add_message(f"You: {message}", "user")
+
+        # Process command asynchronously
+        self.process_message(message)
+
+    def handle_slash_command(self, command: str) -> None:
+        """Dispatch slash commands — TUI-specific first, then shared commands."""
+        parts = command.strip().split()
+        cmd = parts[0].lower()
+        args = parts[1:]
+
+        # --- TUI-specific commands ---
+
+        if cmd == "/cancel":
+            # Cancel the active stream worker and reset for new input
+            workers = self.workers
+            for worker in workers:
+                if not worker.is_finished:
+                    worker.cancel()
+            self.add_message("Cancelled active request.", "info")
+            return
+
+        if cmd == "/new":
+            # Cancel any active stream first
+            for worker in self.workers:
+                if not worker.is_finished:
+                    worker.cancel()
+            self.current_action_id = self.DEFAULT_ACTION_ID
+            self.current_run_id = None
+            self.add_message(f"Starting new conversation with default action ({self.DEFAULT_ACTION_ID})...", "info")
+            return
+
+        if cmd == "/use":
+            # /use <action_id> — switch active action
+            if not args:
+                self.add_message("Usage: /use <action_id>", "info")
+                return
+            action_id = args[0]
+            self.add_message(f"> /use {action_id}", "user")
+            # Validate action exists
+            api_client = self._get_api_client()
+            if not api_client:
+                return
+            from ..commands import actions_get
+            lines = actions_get(api_client, id=action_id)
+            if lines and lines[0].startswith("Error"):
+                self.add_message(f"Action not found: {action_id}", "error")
+                return
+            self.current_action_id = action_id
+            self.current_run_id = None
+            self.add_message("\n".join(["Switched to action:"] + lines), "info")
+            return
+
+        if cmd == "/continue":
+            # /continue <run_id> — attach to an existing run
+            if not args:
+                self.add_message("Usage: /continue <run_id>", "info")
+                return
+            run_id = args[0]
+            self.add_message(f"> /continue {run_id}", "user")
+            # Validate run exists and get its action_id
+            api_client = self._get_api_client()
+            if not api_client:
+                return
+            from ..commands import runs_get
+            lines = runs_get(api_client, id=run_id)
+            if lines and lines[0].startswith("Error"):
+                self.add_message(f"Run not found: {run_id}", "error")
+                return
+            self.current_run_id = run_id
+            self.add_message("\n".join(["Continuing run:"] + lines), "info")
+            return
+
+        if cmd == "/status":
+            # Show current action and run
+            self.add_message(f"> /status", "user")
+            info = [
+                f"Action ID:  {self.current_action_id}",
+                f"Run ID:     {self.current_run_id or '(none — next message starts new run)'}",
+            ]
+            # Fetch action name
+            api_client = self._get_api_client()
+            if api_client:
+                from ..commands import actions_get
+                action_lines = actions_get(api_client, id=self.current_action_id)
+                if action_lines and not action_lines[0].startswith("Error"):
+                    info.append("")
+                    info.extend(action_lines)
+            self.add_message("\n".join(info), "assistant")
+            return
+
+        # --- Shared commands (dispatch to commands.py) ---
+
+        self.add_message(f"> {command}", "user")
+
+        api_client = self._get_api_client()
+        if not api_client:
+            return
+
+        from ..commands import dispatch
+        try:
+            lines = dispatch(command, api_client)
+            self.add_message("\n".join(lines), "assistant")
+        except Exception as e:
+            self.add_message(f"Command error: {e}", "error")
+
+    def _get_api_client(self):
+        """Get the AuthenticatedClient for API calls. Returns None with error message if not available."""
+        from autobots_client import AuthenticatedClient
+
+        if not hasattr(self.app, 'autobots_client'):
+            self.add_message("Error: Client not initialized", "error")
+            return None
+
+        api_client = self.app.autobots_client.client
+        if not isinstance(api_client, AuthenticatedClient):
+            self.add_message("Error: Not authenticated", "error")
+            return None
+        return api_client
+
+    def add_message(self, text: str, msg_type: str = "assistant") -> None:
+        """Add a message to the chat."""
+        messages = self.query_one("#messages", VerticalScroll)
+        css_class = f"message {msg_type}-message"
+        messages.mount(Static(text, classes=css_class))
+        messages.scroll_end(animate=False)
+
+    def process_message(self, message: str) -> None:
+        """Process user message by running action."""
+        if not hasattr(self.app, 'autobots_client'):
+            self.add_message("Error: Client not initialized", "error")
+            return
+
+        # Run action async and poll for result
+        self.run_action_with_polling(message)
+
+    def run_action_with_polling(self, user_input: str) -> None:
+        """Run action and stream results via SSE."""
+        client = self.app.autobots_client
+
+        # Send the message to the action
+        # If continuing a conversation, pass the current_run_id as action_result_id
+        success, run_id, message = client.run_action_async(
+            self.current_action_id,
+            user_input,
+            action_result_id=self.current_run_id
+        )
+
+        if not success:
+            self.add_message(f"Error starting action: {message}", "error")
+            return
+
+        # Check if this is continuing an existing conversation
+        if self.current_run_id and run_id == self.current_run_id:
+            logger.info(f"Continuing conversation with run_id: {run_id}")
+        else:
+            # New conversation started
+            self.current_run_id = run_id
+            logger.info(f"Started new conversation with run_id: {run_id}")
+
+        # exclusive=True cancels any previous stream worker entirely
+        self.run_worker(self.stream_results(run_id), exclusive=True, group="stream")
+
+    async def stream_results(self, run_id: str) -> None:
+        """Stream action status and display final result from results array.
+
+        Runs SSE streaming and result-polling concurrently.  Whichever
+        detects a terminal state first (success *or* error) wins.
+        """
+        client = self.app.autobots_client
+        got_final_result = False
+
+        # Track status widget for streaming transitional messages
+        status_widget_container = [None]
+
+        logger.info(f"Starting stream_results for {run_id}")
+
+        def _remove_status_widget() -> None:
+            """Remove the streaming status widget if it exists."""
+            if status_widget_container[0]:
+                try:
+                    status_widget_container[0].remove()
+                    status_widget_container[0] = None
+                except Exception as e:
+                    logger.warning(f"Failed to remove status widget: {e}")
+
+        def _try_fetch_final_result() -> bool:
+            """Attempt to fetch and display the final result. Returns True on success/error."""
+            nonlocal got_final_result
+            if got_final_result:
+                return True
+            success, final_result, message = client.get_action_result(run_id)
+            if not success or not final_result:
+                return False
+
+            status = final_result.get("status", "").lower()
+            logger.info(f"Polled run status: {status}")
+
+            # Handle error/failed states
+            if status in ("error", "failed"):
+                logger.info(f"Run failed with status: {status}")
+                _remove_status_widget()
+                error_msg = final_result.get("message", "") or final_result.get("error", "") or "Action failed"
+                self.add_message(f"Error: {error_msg}", "error")
+                got_final_result = True
+                return True
+
+            # Only treat as complete if status indicates completion
+            if status not in ("completed", "success", "finished"):
+                return False
+
+            action_doc = final_result.get("result", {})
+            results_list = action_doc.get("results", []) if isinstance(action_doc, dict) else []
+            # Verify the last result actually has output
+            if results_list:
+                last = results_list[-1]
+                if isinstance(last, dict) and last.get("output"):
+                    logger.info("Final result fetched successfully")
+                    _remove_status_widget()
+                    self.display_final_result(final_result)
+                    got_final_result = True
+                    return True
+            return False
+
+        def handle_status_message(data: dict) -> None:
+            """Handle SSE messages — status updates and completion signals."""
+            if not isinstance(data, dict) or got_final_result:
+                return
+
+            # Plain-text status messages (type="status" set by client.py)
+            if data.get("type") == "status":
+                text = data.get("text", "")
+                text_lower = text.lower()
+                logger.debug(f"SSE status: {text}")
+
+                # Show all non-empty status messages as progress
+                if text.strip():
+                    status_widget_container[0] = self.update_streaming_message(
+                        {"blocks": [{"text": text}]},
+                        status_widget_container[0],
+                    )
+
+                # Detect completion from text signals (server sends these as plain text)
+                if any(kw in text_lower for kw in ["finishing", "completed", "finished"]):
+                    logger.info(f"Completion signal from SSE text: {text}")
+                    _try_fetch_final_result()
+                return
+
+            # JSON status messages
+            status = data.get("status", "").lower()
+            if status:
+                logger.info(f"SSE JSON status: {status}")
+            if status in ["completed", "success", "finished", "error", "failed"]:
+                _try_fetch_final_result()
+
+        # ---- Concurrent polling task ----
+        async def _poll_until_done() -> None:
+            """Poll the run result every few seconds until it reaches a terminal state."""
+            while not got_final_result:
+                await asyncio.sleep(5)
+                if got_final_result:
+                    return
+                try:
+                    if _try_fetch_final_result():
+                        return
+                except Exception as e:
+                    logger.warning(f"Poll error: {e}")
+
+        poll_task = asyncio.create_task(_poll_until_done())
+
+        try:
+            status_task = asyncio.create_task(
+                asyncio.wait_for(
+                    client.stream_action_result(run_id, handle_status_message),
+                    timeout=300.0  # 5 minutes — avoids indefinite hangs
+                )
+            )
+
+            try:
+                # Wait for EITHER the SSE stream to end OR polling to find the result
+                done, pending = await asyncio.wait(
+                    [status_task, poll_task],
+                    return_when=asyncio.FIRST_COMPLETED,
+                )
+                # Cancel whichever is still running
+                for task in pending:
+                    task.cancel()
+                    try:
+                        await task
+                    except (asyncio.CancelledError, Exception):
+                        pass
+                # Check for exceptions in completed tasks
+                for task in done:
+                    if task.exception() and not isinstance(task.exception(), (asyncio.CancelledError, asyncio.TimeoutError)):
+                        logger.error(f"Stream/poll error: {task.exception()}")
+            except asyncio.CancelledError:
+                logger.info(f"Stream cancelled for {run_id}")
+                status_task.cancel()
+                poll_task.cancel()
+                _remove_status_widget()
+                return
+            except asyncio.TimeoutError:
+                logger.warning(f"SSE timeout for {run_id}")
+                poll_task.cancel()
+                _remove_status_widget()
+                self.add_message("Action timed out waiting for response", "error")
+                return
+        except asyncio.CancelledError:
+            logger.info(f"Stream worker cancelled for {run_id}")
+            poll_task.cancel()
+            _remove_status_widget()
+            return
+        except Exception as e:
+            logger.error(f"SSE error for {run_id}: {e}")
+            poll_task.cancel()
+
+        # Always clean up status widget when SSE stream ends
+        _remove_status_widget()
+
+        # Final fallback — if neither SSE nor polling found the result
+        if not got_final_result:
+            logger.info(f"Final fallback poll for {run_id}")
+            _try_fetch_final_result()
+
+        if not got_final_result:
+            self.add_message("Could not get result. Use /new to start over.", "error")
+
+    def update_streaming_message(self, output: any, widget_ref: any = None) -> Static:
+        """Update or create a streaming message widget with new output.
+
+        Returns the widget being updated/created for future updates.
+        """
+        text_content = self.extract_text_from_output(output)
+        if not text_content:
+            return widget_ref
+
+        messages = self.query_one("#messages", VerticalScroll)
+
+        if widget_ref is None:
+            widget_ref = Static(text_content, classes="message assistant-message", markup=False, expand=True)
+            messages.mount(widget_ref)
+        else:
+            try:
+                widget_ref.update(text_content)
+            except Exception as e:
+                logger.warning(f"Failed to update widget: {e}")
+                widget_ref = Static(text_content, classes="message assistant-message", markup=False, expand=True)
+                messages.mount(widget_ref)
+
+        messages.scroll_end(animate=False)
+        return widget_ref
+
+    def extract_text_from_output(self, output: any) -> str:
+        """Extract text content from output blocks structure and clean it for display."""
+        if not isinstance(output, dict):
+            return ""
+
+        text_parts = []
+
+        if "blocks" in output:
+            blocks = output["blocks"]
+            if isinstance(blocks, list):
+                for block in blocks:
+                    if isinstance(block, dict):
+                        text = block.get("text", "").strip()
+                        if text:
+                            text_parts.append(self._clean_text_for_display(text))
+
+        elif "text" in output:
+            text = output["text"]
+            if text:
+                text = str(text) if not isinstance(text, str) else text
+                text_parts.append(self._clean_text_for_display(text))
+
+        return "\n".join(text_parts)
+
+    def _clean_text_for_display(self, text: str) -> str:
+        """Clean text for safe display in Textual (remove HTML, unescape entities, etc.)."""
+        if not text:
+            return ""
+
+        # Unescape HTML entities (e.g., &amp; -> &, &lt; -> <)
+        text = html.unescape(text)
+
+        # Remove HTML tags if present
+        text = re.sub(r'<[^>]+>', '', text)
+
+        # Unescape literal \n to actual newlines if present
+        if '\\n' in text:
+            text = text.replace('\\n', '\n')
+
+        return text.strip()
+
+    def display_final_result(self, result: dict) -> None:
+        """Display the final result from results array, showing only the latest output.
+
+        Args:
+            result: The result dictionary from get_action_result containing result.result.results[]
+        """
+        logger.info(f"Displaying final result from results array")
+
+        # Extract results array from result.result.results
+        if "result" not in result or not isinstance(result["result"], dict):
+            logger.warning(f"No result.result field found")
+            self.add_message("Action completed (no output)", "info")
+            return
+
+        action_doc = result["result"]
+        logger.info(f"ActionDoc keys: {action_doc.keys()}")
+
+        if "results" not in action_doc or not isinstance(action_doc["results"], list):
+            logger.warning(f"No results array found in ActionDoc")
+            self.add_message("Action completed (no output)", "info")
+            return
+
+        results_list = action_doc["results"]
+        logger.info(f"Found {len(results_list)} items in results array")
+
+        if len(results_list) == 0:
+            logger.warning(f"Results array is empty")
+            self.add_message("Action completed (no output)", "info")
+            return
+
+        # Only display the LAST result item's output.
+        # The results array contains full conversation history; the user's input
+        # was already shown when they typed it, so we only need the latest response.
+        last_result = results_list[-1]
+        if not isinstance(last_result, dict):
+            logger.warning(f"Last result is not a dict")
+            self.add_message("Action completed (no output)", "info")
+            return
+
+        logger.info(f"Processing last result, keys: {last_result.keys()}")
+
+        # Extract and display only the output (skip input — already shown)
+        if "output" in last_result and last_result["output"]:
+            output_data = last_result["output"]
+            output_text = self.extract_text_from_output(output_data)
+            if output_text:
+                logger.info(f"Last result output: {len(output_text)} chars, {len(output_text.splitlines())} lines")
+                self.add_message(output_text, "assistant")
+            else:
+                logger.warning(f"Last result output extraction returned empty")
+                self.add_message("Action completed (no output)", "info")
+        else:
+            logger.warning(f"Last result has no output field")
+            self.add_message("Action completed (no output)", "info")
+
+    def format_and_display_output(self, output: any) -> None:
+        """Format and display output, extracting text and files from blocks."""
+        if isinstance(output, dict):
+            # Check for blocks structure
+            if "blocks" in output:
+                blocks = output["blocks"]
+                if isinstance(blocks, list):
+                    for block in blocks:
+                        if isinstance(block, dict):
+                            # Extract text
+                            text = block.get("text", "").strip()
+                            if text:
+                                # Unescape unicode characters
+                                text = text.encode().decode('unicode_escape')
+                                self.add_message(text, "assistant")
+
+                            # Extract files
+                            files = block.get("files", [])
+                            if files and isinstance(files, list):
+                                for file_info in files:
+                                    if isinstance(file_info, dict):
+                                        file_name = file_info.get("name", "unnamed file")
+                                        self.add_message(f"📎 File: {file_name}", "info")
+                                    elif isinstance(file_info, str):
+                                        self.add_message(f"📎 File: {file_info}", "info")
+                return
+
+            # Check for direct text field
+            if "text" in output:
+                text = output["text"]
+                if text:
+                    text = text.encode().decode('unicode_escape') if isinstance(text, str) else str(text)
+                    self.add_message(text, "assistant")
+                return
+
+        # Fallback: display as-is
+        if isinstance(output, dict):
+            self.add_message(json.dumps(output, indent=2), "assistant")
+        else:
+            self.add_message(str(output), "assistant")
+
+    def update_last_assistant_message(self, text: str) -> None:
+        """Update the last assistant message with new text."""
+        messages = self.query_one("#messages", VerticalScroll)
+        assistant_messages = messages.query(".assistant-message")
+
+        if assistant_messages:
+            # Update the last assistant message
+            last_msg = assistant_messages[-1]
+            last_msg.update(text)
+        else:
+            # No existing message, create new one
+            self.add_message(text, "assistant")