app.py

"""Chainlit chat UI for StarShell with visible tool calls."""

import json
import os
import tempfile
from contextlib import AsyncExitStack

import chainlit as cl
from openai.types.responses import ResponseTextDeltaEvent
from agents import Agent, Runner, ItemHelpers
from agents.mcp import MCPServerStdio, ToolFilterStatic

from starshell.config import dotenv, build_mcp_env, build_playwright_config, build_playwright_login_url
from starshell.prompt import build_instructions
from starshell.agents.api_call import _make_api_call
from starshell.agents.bash import _make_terminal
from starshell.agents.mcp import _mcp_server_params
from starshell.agents import TruncatingMCPServer
from starshell.model import get_model, get_model_settings

# Read agent type, domain, and model from env.
_agent_type = dotenv.get("STARSHELL_AGENT_TYPE", os.environ.get("STARSHELL_AGENT_TYPE", "bash"))
_domain = dotenv.get("STARSHELL_DOMAIN", os.environ.get("STARSHELL_DOMAIN", "servicenow"))
_model_id = dotenv.get("STARSHELL_MODEL", os.environ.get("STARSHELL_MODEL", None))
_model = get_model(_model_id)
_model_settings = get_model_settings(_model_id)


@cl.on_chat_start
async def on_chat_start():
    stack = AsyncExitStack()
    agent = None

    if _agent_type == "bash":
        agent = Agent(
            name="StarShell",
            model=_model,
            model_settings=_model_settings,
            instructions=build_instructions(_domain, "bash"),
            tools=[_make_terminal(_domain)],
        )

    elif _agent_type == "mcp":
        mcp_env = build_mcp_env(_domain)
        server_name, params = _mcp_server_params(_domain, mcp_env)
        server = await stack.enter_async_context(
            MCPServerStdio(name=server_name, params=params, client_session_timeout_seconds=120)
        )
        agent = Agent(
            name="StarShell MCP",
            model=_model,
            model_settings=_model_settings,
            instructions=build_instructions(_domain, "mcp"),
            mcp_servers=[TruncatingMCPServer(server)],
        )

    elif _agent_type == "playwright":
        pw_config = build_playwright_config(_domain)
        config_file = tempfile.NamedTemporaryFile(
            mode="w", suffix=".json", prefix="playwright_mcp_", delete=False
        )
        json.dump(pw_config, config_file)
        config_file.close()
        cl.user_session.set("pw_config_path", config_file.name)

        server = await stack.enter_async_context(
            MCPServerStdio(
                name="Playwright",
                params={
                    "command": "npx",
                    "args": [
                        "@playwright/mcp",
                        "--headless",
                        "--config", config_file.name,
                    ],
                },
                client_session_timeout_seconds=120,
                tool_filter=ToolFilterStatic(blocked_tool_names=["browser_install"]),
            )
        )
        # Log in before the agent starts
        login_url = build_playwright_login_url(_domain)
        if login_url:
            await server.call_tool("browser_navigate", {"url": login_url})

        agent = Agent(
            name="StarShell Playwright",
            model=_model,
            model_settings=_model_settings,
            instructions=build_instructions(_domain, "playwright"),
            mcp_servers=[TruncatingMCPServer(server)],
        )

    elif _agent_type == "api_call":
        agent = Agent(
            name="StarShell API",
            model=_model,
            model_settings=_model_settings,
            instructions=build_instructions(_domain, "api_call"),
            tools=[_make_api_call(_domain)],
        )

    else:
        raise ValueError(f"Unknown agent type: {_agent_type!r}")

    cl.user_session.set("agent", agent)
    cl.user_session.set("exit_stack", stack)


@cl.on_chat_end
async def on_chat_end():
    stack = cl.user_session.get("exit_stack")
    if stack:
        await stack.aclose()
    # Clean up playwright config file
    pw_path = cl.user_session.get("pw_config_path")
    if pw_path:
        try:
            os.unlink(pw_path)
        except OSError:
            pass


@cl.set_starters
async def set_starters():
    if _domain != "servicenow":
        return []
    return [
        cl.Starter(
            label="Create Incident",
            message='Create a new incident with a value of "Defect tracking tool is down." for field "Short description", a value of "David Miller" for field "Caller", a value of "" for field "Assignment group", a value of "While launching the defect tracking base URL, it is redirecting to an error page." for field "Description", a value of "" for field "Configuration item", a value of "" for field "Problem", a value of "false" for field "Knowledge", a value of "System Administrator" for field "Resolved by", a value of "Resolved by caller" for field "Resolution code", a value of "" for field "Service offering", a value of "" for field "Change Request", a value of "1 - High" for field "Urgency", a value of "INC0009004" for field "Number", a value of "Reverted to a previous version." for field "Resolution notes", a value of "" for field "Service", a value of "Walk-in" for field "Channel", a value of "" for field "Parent Incident", a value of "3 - Low" for field "Impact", and a value of "Software" for field "Category".',
        ),
        cl.Starter(
            label="Order Laptop",
            message="Go to the hardware store and order 1 \"Standard Laptop\" with configuration {'Additional software requirements': 'Salesforce, QuickBooks, Microsoft Office 365, Google Workspace, Asana', 'Adobe Acrobat': False, 'Adobe Photoshop': True}",
        ),
        cl.Starter(
            label="Filter Change Request",
            message='Create a filter for the change request list to extract all entries where:\n    - "Assigned to" is "Fred Luddy" or \n    - "Number" is "CHG0000068" or \n    - "Risk" is "High" or \n    - "Requested by" is "System Administrator" or \n    - "State" is "Implement"',
        ),
    ]


@cl.on_message
async def on_message(message: cl.Message):
    agent = cl.user_session.get("agent")
    history = cl.user_session.get("history", [])
    history.append({"role": "user", "content": message.content})

    result = Runner.run_streamed(agent, history, max_turns=100)

    # Single collapsible "Steps" container for all thinking + tool calls
    steps_container = cl.Step(name="Steps", type="run", parent_id=message.id)
    await steps_container.send()
    container_id = steps_container.id

    # Track current thinking step so we can stream CoT text into it
    thinking_step = None
    thinking_text = ""
    final_text = ""
    step_count = 0

    async def ensure_thinking():
        """Create a new collapsed thinking step if one isn't active."""
        nonlocal thinking_step, thinking_text
        if thinking_step is None:
            thinking_text = ""
            thinking_step = cl.Step(
                name="Thinking",
                type="llm",
                parent_id=container_id,
            )
            await thinking_step.send()
        return thinking_step

    async def finalize_thinking():
        """Close the current thinking step."""
        nonlocal thinking_step, thinking_text
        if thinking_step is not None:
            thinking_step.output = thinking_text
            await thinking_step.update()
            thinking_step = None
            thinking_text = ""

    async for event in result.stream_events():
        # Stream text tokens into the current thinking step
        if event.type == "raw_response_event" and isinstance(event.data, ResponseTextDeltaEvent):
            await ensure_thinking()
            thinking_text += event.data.delta
            thinking_step.output = thinking_text
            await thinking_step.update()
            continue

        if event.type != "run_item_stream_event":
            continue

        item = event.item

        if item.type == "tool_call_item":
            # Close thinking before showing tool call
            await finalize_thinking()
            step_count += 1

            tool_name = getattr(item.raw_item, "name", "tool")
            try:
                args = json.loads(item.raw_item.arguments)
            except (json.JSONDecodeError, TypeError):
                args = {}

            # For bash terminal: show command directly; otherwise show tool name + args
            if tool_name == "terminal" and "command" in args:
                display_name = "Terminal"
                display_input = args["command"]
            else:
                display_name = tool_name
                display_input = json.dumps(args, indent=2)

            step = cl.Step(name=display_name, type="tool", parent_id=container_id)
            step.input = display_input
            await step.send()
            cl.user_session.set("current_step", step)

        elif item.type == "tool_call_output_item":
            step = cl.user_session.get("current_step")
            if step:
                raw = item.output or "[no output]"
                # MCP tools may return structured content (dict/list) instead of str
                if not isinstance(raw, str):
                    text = json.dumps(raw, indent=2, default=str)
                else:
                    text = raw
                lines = text.splitlines()
                if len(lines) > 20:
                    text = "\n".join(lines[:20]) + f"\n... ({len(lines) - 20} more lines)"
                step.output = text
                await step.update()

        elif item.type == "message_output_item":
            await finalize_thinking()
            final_text = ItemHelpers.text_message_output(item)

    # Final cleanup
    await finalize_thinking()
    final_text = final_text if final_text else (result.final_output or "")

    # Update container summary
    steps_container.output = f"{step_count} tool call(s)"
    await steps_container.update()

    # Send the final answer as a separate message after all steps
    final_msg = cl.Message(content=final_text)
    await final_msg.send()

    # Persist history for multi-turn
    history.append({"role": "assistant", "content": final_text})
    cl.user_session.set("history", history)