From 844214ada1c9568f69891c46dc415745e554aca4 Mon Sep 17 00:00:00 2001 From: Dan Zuckerberg Date: Sat, 11 Jul 2026 00:16:11 -0400 Subject: [PATCH 001/134] feat: add managed-memory agent panel via app-templates submodule + Vercel-native SPN proxy - Vendor databricks/app-templates as a sparse submodule (agent-openai-agents-sdk, e2e-chatbot-app-next) with an agent/ overlay + assemble_agent.sh. - Add a slide-out Agent panel (store + trigger) to the sso-spn org view. - Embed the managed-memory Databricks App via a same-origin Vercel route (api/agent-proxy) that mints the current user's SPN token (guest/BYOD) and injects it as a bearer, streaming responses (chat SSE). No Go proxy / Cloud Run needed for the agent. - Chat UI overlay patches: base:"./" + a mount-prefix fetch shim so assets and /api/ calls resolve under the proxy subpath. --- .gitignore | 7 + .gitmodules | 3 + agent/agent_server/agent.py | 168 +++++ agent/agent_server/genie_tools.py | 109 +++ agent/agent_server/start_server.py | 29 + agent/agent_server/utils.py | 67 ++ agent/agent_server/utils_memory.py | 205 ++++++ agent/databricks.yml | 76 +++ .../client/src/components/chat.tsx | 355 ++++++++++ .../src/components/elements/mcp-tool.tsx | 261 ++++++++ .../src/components/genie-attribution.tsx | 95 +++ .../client/src/components/message.tsx | 534 +++++++++++++++ .../client/src/contexts/AppConfigContext.tsx | 70 ++ .../client/src/lib/genie-attribution.ts | 121 ++++ .../client/src/lib/utils.ts | 129 ++++ .../e2e-chatbot-app-next/client/src/main.tsx | 44 ++ .../client/vite.config.ts | 54 ++ .../server/src/lib/sanitize-ui-messages.ts | 50 ++ .../server/src/routes/chat.ts | 620 ++++++++++++++++++ .../server/src/routes/config.ts | 103 +++ agent/scripts/start_app.py | 367 +++++++++++ scripts/assemble_agent.sh | 32 + src/app/api/agent-proxy/[[...path]]/route.ts | 185 ++++++ src/app/sso-spn/[orgId]/layout.tsx | 2 + src/components/agent/agent-panel.tsx | 152 +++++ src/components/sso-spn-top-nav.tsx | 2 + src/stores/agent-panel-store.ts | 64 ++ vendor/app-templates | 1 + 28 files changed, 3905 insertions(+) create mode 100644 .gitmodules create mode 100644 agent/agent_server/agent.py create mode 100644 agent/agent_server/genie_tools.py create mode 100644 agent/agent_server/start_server.py create mode 100644 agent/agent_server/utils.py create mode 100644 agent/agent_server/utils_memory.py create mode 100644 agent/databricks.yml create mode 100644 agent/patches/e2e-chatbot-app-next/client/src/components/chat.tsx create mode 100644 agent/patches/e2e-chatbot-app-next/client/src/components/elements/mcp-tool.tsx create mode 100644 agent/patches/e2e-chatbot-app-next/client/src/components/genie-attribution.tsx create mode 100644 agent/patches/e2e-chatbot-app-next/client/src/components/message.tsx create mode 100644 agent/patches/e2e-chatbot-app-next/client/src/contexts/AppConfigContext.tsx create mode 100644 agent/patches/e2e-chatbot-app-next/client/src/lib/genie-attribution.ts create mode 100644 agent/patches/e2e-chatbot-app-next/client/src/lib/utils.ts create mode 100644 agent/patches/e2e-chatbot-app-next/client/src/main.tsx create mode 100644 agent/patches/e2e-chatbot-app-next/client/vite.config.ts create mode 100644 agent/patches/e2e-chatbot-app-next/server/src/lib/sanitize-ui-messages.ts create mode 100644 agent/patches/e2e-chatbot-app-next/server/src/routes/chat.ts create mode 100644 agent/patches/e2e-chatbot-app-next/server/src/routes/config.ts create mode 100644 agent/scripts/start_app.py create mode 100755 scripts/assemble_agent.sh create mode 100644 src/app/api/agent-proxy/[[...path]]/route.ts create mode 100644 src/components/agent/agent-panel.tsx create mode 100644 src/stores/agent-panel-store.ts create mode 160000 vendor/app-templates diff --git a/.gitignore b/.gitignore index 5d8be11..73698db 100644 --- a/.gitignore +++ b/.gitignore @@ -62,3 +62,10 @@ go/databricks-proxy go/proxy # Local skills and personal tooling .local/ + +# Assembled agent (ephemeral, from vendor submodule + agent/ overlay) +agent-build/ + +# Python build artifacts +__pycache__/ +*.pyc diff --git a/.gitmodules b/.gitmodules new file mode 100644 index 0000000..8fc3f1f --- /dev/null +++ b/.gitmodules @@ -0,0 +1,3 @@ +[submodule "vendor/app-templates"] + path = vendor/app-templates + url = https://github.com/databricks/app-templates.git diff --git a/agent/agent_server/agent.py b/agent/agent_server/agent.py new file mode 100644 index 0000000..ec279d6 --- /dev/null +++ b/agent/agent_server/agent.py @@ -0,0 +1,168 @@ +import logging +import os +from contextlib import AsyncExitStack +from datetime import datetime +from typing import AsyncGenerator + +import mlflow +from agents import Agent, Runner, function_tool, set_default_openai_api, set_default_openai_client +from agents.tracing import set_trace_processors +from databricks.sdk import WorkspaceClient +from databricks_openai import AsyncDatabricksOpenAI +from databricks_openai.agents import McpServer +from fastapi import HTTPException +from mlflow.genai.agent_server import invoke, stream +from mlflow.types.responses import ( + ResponsesAgentRequest, + ResponsesAgentResponse, + ResponsesAgentStreamEvent, +) + +from agent_server.utils import ( + build_mcp_url, + get_session_id, + process_agent_stream_events, +) +from agent_server.genie_tools import GENIE_INSTRUCTIONS, GENIE_TOOLS +from agent_server.utils_memory import ( + MEMORY_INSTRUCTIONS, + MEMORY_TOOLS, + MemoryContext, + resolve_scope, +) + +# Skill-first: utils_memory.py stays a verbatim copy of the managed-memory skill; +# the Genie deviation is composed here at wire time. +INSTRUCTIONS = MEMORY_INSTRUCTIONS + "\n\n" + GENIE_INSTRUCTIONS + +logger = logging.getLogger(__name__) + +set_default_openai_client(AsyncDatabricksOpenAI()) +set_default_openai_api("chat_completions") +set_trace_processors([]) +mlflow.openai.autolog() +logging.getLogger("mlflow.utils.autologging_utils").setLevel(logging.ERROR) + + +@function_tool +def get_current_time() -> str: + """Get the current date and time.""" + return datetime.now().isoformat() + + +def _app_workspace_client() -> WorkspaceClient: + """App service principal (Databricks Apps OAuth M2M), not the proxy caller token.""" + return WorkspaceClient() + + +GENIE_ONE_MCP_PATH = "/api/2.0/mcp/genie" + + +def _genie_mcp_path() -> str: + mode = os.environ.get("GENIE_MCP_MODE", "one").strip().lower() + explicit = os.environ.get("GENIE_MCP_URL", "").strip() + + if mode == "space": + space_id = os.environ.get("GENIE_SPACE_ID", "").strip() + if not space_id: + raise ValueError("GENIE_MCP_MODE=space requires GENIE_SPACE_ID") + return f"/api/2.0/mcp/genie/{space_id}" + + if explicit and not explicit.rstrip("/").endswith("/api/2.0/mcp/genie"): + return explicit if explicit.startswith("/") else f"/{explicit}" + return GENIE_ONE_MCP_PATH + + +def _genie_mcp_url(app_wc: WorkspaceClient) -> str: + """Genie One managed MCP when GENIE_MCP_MODE=one (default).""" + path_or_url = _genie_mcp_path() + if path_or_url.startswith("http"): + return path_or_url.rstrip("/") + if not path_or_url.startswith("/"): + path_or_url = f"/{path_or_url}" + return build_mcp_url(path_or_url.rstrip("/"), workspace_client=app_wc) + + +async def init_genie_mcp_server() -> McpServer: + app_wc = _app_workspace_client() + url = _genie_mcp_url(app_wc) + logger.info("Genie MCP mode=%s url=%s", os.environ.get("GENIE_MCP_MODE", "one"), url) + return McpServer( + url=url, + name="Genie", + workspace_client=app_wc, + timeout=60.0, + ) + + +async def connect_healthy_mcp_servers( + stack: AsyncExitStack, servers: list[McpServer] +) -> tuple[list[McpServer], list[str]]: + healthy: list[McpServer] = [] + unavailable: list[str] = [] + for server in servers: + name = getattr(server, "name", "MCP server") + try: + connected = await stack.enter_async_context(server) + await connected.list_tools() + healthy.append(connected) + except Exception: + logger.warning("MCP server %r unavailable; continuing without it.", name, exc_info=True) + unavailable.append(name) + return healthy, unavailable + + +def create_agent(mcp_servers: list[McpServer] | None = None) -> Agent: + return Agent( + name="Agent", + instructions=INSTRUCTIONS, + model="databricks-gpt-5-2", + tools=[get_current_time, *GENIE_TOOLS, *MEMORY_TOOLS], + mcp_servers=mcp_servers or [], + ) + + +def _require_scope(request: ResponsesAgentRequest) -> str: + scope = resolve_scope(request) + if not scope: + raise HTTPException( + status_code=401, + detail="No end-user identity — refusing a shared memory scope.", + ) + return scope + + +async def _mcp_servers_for_request(stack: AsyncExitStack) -> list[McpServer]: + servers, _ = await connect_healthy_mcp_servers(stack, [await init_genie_mcp_server()]) + return servers + + +@invoke() +async def invoke_handler(request: ResponsesAgentRequest) -> ResponsesAgentResponse: + if session_id := get_session_id(request): + mlflow.update_current_trace(metadata={"mlflow.trace.session": session_id}) + scope = _require_scope(request) + async with AsyncExitStack() as stack: + agent = create_agent(mcp_servers=await _mcp_servers_for_request(stack)) + messages = [i.model_dump() for i in request.input] + result = await Runner.run(agent, messages, context=MemoryContext(scope=scope)) + return ResponsesAgentResponse( + output=[item.to_input_item() for item in result.new_items] + ) + + +@stream() +async def stream_handler( + request: ResponsesAgentRequest, +) -> AsyncGenerator[ResponsesAgentStreamEvent, None]: + if session_id := get_session_id(request): + mlflow.update_current_trace(metadata={"mlflow.trace.session": session_id}) + scope = _require_scope(request) + async with AsyncExitStack() as stack: + agent = create_agent(mcp_servers=await _mcp_servers_for_request(stack)) + messages = [i.model_dump() for i in request.input] + result = Runner.run_streamed( + agent, input=messages, context=MemoryContext(scope=scope) + ) + async for event in process_agent_stream_events(result.stream_events()): + yield event diff --git a/agent/agent_server/genie_tools.py b/agent/agent_server/genie_tools.py new file mode 100644 index 0000000..ab8064e --- /dev/null +++ b/agent/agent_server/genie_tools.py @@ -0,0 +1,109 @@ +"""Genie One MCP helpers — ask + poll until completed.""" + +import json +import logging +import time + +from agents import function_tool +from databricks.sdk import WorkspaceClient + +logger = logging.getLogger(__name__) + +GENIE_MCP_PATH = "/api/2.0/mcp/genie" +MAX_POLLS = 45 +POLL_INTERVAL_SEC = 2.0 + + +def _app_workspace_client() -> WorkspaceClient: + return WorkspaceClient() + + +def _mcp_tool_call(name: str, arguments: dict) -> dict: + body = { + "jsonrpc": "2.0", + "id": 1, + "method": "tools/call", + "params": {"name": name, "arguments": arguments}, + } + raw = _app_workspace_client().api_client.do( + "POST", GENIE_MCP_PATH, body=body + ) + if not isinstance(raw, dict): + return {} + result = raw.get("result") or {} + structured = result.get("structuredContent") + if structured: + return structured + content = result.get("content") or [] + if content and isinstance(content[0], dict) and content[0].get("text"): + try: + return json.loads(content[0]["text"]) + except json.JSONDecodeError: + return {"text": content[0]["text"]} + return result + + +def _poll_genie(conversation_id: str, response_id: str) -> dict: + for _ in range(MAX_POLLS): + payload = _mcp_tool_call( + "genie_poll_response", + { + "conversation_id": conversation_id, + "response_id": response_id, + }, + ) + status = payload.get("status") + if status == "completed": + return payload + if status == "failed": + return payload + time.sleep(POLL_INTERVAL_SEC) + return {"status": "timeout", "message": "Genie poll timed out"} + + +_BROAD_DATA_PROMPTS = ( + "tell me about my data", + "what is my data", + "what's my data", + "describe my data", + "my data", +) + +_TABLE_DETAIL_SUFFIX = ( + " List the catalogs, schemas, and tables available. For each table, give its " + "purpose, key columns, and approximate row count. Prioritize concrete tables " + "and their columns over dashboards or Genie spaces." +) + + +def _augment_broad_question(question: str) -> str: + normalized = question.strip().lower().rstrip("?.! ") + if normalized in _BROAD_DATA_PROMPTS: + return question.strip() + _TABLE_DETAIL_SUFFIX + return question + + +@function_tool +def ask_genie_one(question: str) -> str: + """Query Genie One over workspace data. Use for any question about tables, catalogs, dashboards, or 'my data'. For broad questions, it automatically requests table- and column-level detail. Polls until Genie completes.""" + question = _augment_broad_question(question) + ask = _mcp_tool_call("genie_ask", {"question": question}) + if ask.get("status") == "completed" and ask.get("final_answer"): + return ask["final_answer"] + conversation_id = ask.get("conversation_id") + response_id = ask.get("response_id") + if not conversation_id or not response_id: + return f"Genie ask failed: {json.dumps(ask)[:2000]}" + result = _poll_genie(conversation_id, response_id) + if result.get("status") == "completed": + return result.get("final_answer") or json.dumps(result)[:8000] + return json.dumps(result)[:2000] + + +GENIE_TOOLS = [ask_genie_one] + +# Composed onto MEMORY_INSTRUCTIONS in agent.py so utils_memory.py stays a pristine, +# regenerable copy of the managed-memory skill (this is the local Genie deviation). +GENIE_INSTRUCTIONS = """For any question about workspace data, tables, catalogs, dashboards, metrics, or phrases like "my data", you MUST call ask_genie_one first with the user's question — do not ask the user to clarify catalog/schema first. When the user's question is broad (e.g. "tell me about my data"), phrase the ask_genie_one question to request concrete data assets: the catalogs, schemas, and tables available, plus each table's purpose, key columns, and row counts where possible (for example: "What catalogs, schemas, and tables do I have? For each table, give its purpose, key columns, and approximate row count."). If the first Genie answer stays high-level (only dashboards or Genie spaces, no tables/columns), call ask_genie_one again asking specifically for the tables and their columns in the catalogs that were returned. Use get_current_time for the current date and time. + +After ask_genie_one or memory tools return, answer in clear markdown prose (headings, bullets, tables). When Genie returns dashboard or asset links, include them as markdown links. Prefer surfacing concrete tables and their key columns over a list of dashboards. Never paste raw tool JSON in the final reply.""" diff --git a/agent/agent_server/start_server.py b/agent/agent_server/start_server.py new file mode 100644 index 0000000..731cf4a --- /dev/null +++ b/agent/agent_server/start_server.py @@ -0,0 +1,29 @@ +import os +from pathlib import Path + +from dotenv import load_dotenv +from mlflow.genai.agent_server import AgentServer, setup_mlflow_git_based_version_tracking + +# Load env vars from .env before importing the agent for proper auth +load_dotenv(dotenv_path=Path(__file__).parent.parent / ".env", override=True) + +# Need to import the agent to register the functions with the server +import agent_server.agent # noqa: E402 + +_enable_chat_proxy = os.environ.get("ENABLE_CHAT_PROXY", "true").lower() not in ( + "0", + "false", + "no", +) +agent_server = AgentServer("ResponsesAgent", enable_chat_proxy=_enable_chat_proxy) +# Define the app as a module level variable to enable multiple workers +app = agent_server.app # noqa: F841 +setup_mlflow_git_based_version_tracking() + + +def main(): + agent_server.run(app_import_string="agent_server.start_server:app") + + +if __name__ == "__main__": + main() diff --git a/agent/agent_server/utils.py b/agent/agent_server/utils.py new file mode 100644 index 0000000..c0bb8fa --- /dev/null +++ b/agent/agent_server/utils.py @@ -0,0 +1,67 @@ +import logging +from typing import AsyncGenerator, AsyncIterator, Optional +from uuid import uuid4 + +from agents.result import StreamEvent +from databricks.sdk import WorkspaceClient +from mlflow.genai.agent_server import get_request_headers +from mlflow.types.responses import ResponsesAgentRequest, ResponsesAgentStreamEvent + + +def get_session_id(request: ResponsesAgentRequest) -> str | None: + if request.context and request.context.conversation_id: + return request.context.conversation_id + if request.custom_inputs and isinstance(request.custom_inputs, dict): + return request.custom_inputs.get("session_id") + return None + + +def get_databricks_host(workspace_client: WorkspaceClient | None = None) -> Optional[str]: + workspace_client = workspace_client or WorkspaceClient() + try: + return workspace_client.config.host + except Exception as e: + logging.exception(f"Error getting databricks host from env: {e}") + return None + + +def build_mcp_url(path: str, workspace_client: WorkspaceClient | None = None) -> str: + if not path.startswith("/"): + return path + hostname = get_databricks_host(workspace_client) + return f"{hostname}{path}" + + +def get_user_workspace_client() -> WorkspaceClient: + headers = get_request_headers() + token = headers.get("x-forwarded-access-token") + if not token: + auth = headers.get("authorization") or headers.get("Authorization") or "" + if auth.lower().startswith("bearer "): + token = auth[7:].strip() + if not token: + return WorkspaceClient() + return WorkspaceClient(token=token, auth_type="pat") + + +async def process_agent_stream_events( + async_stream: AsyncIterator[StreamEvent], +) -> AsyncGenerator[ResponsesAgentStreamEvent, None]: + curr_item_id = str(uuid4()) + async for event in async_stream: + if event.type == "raw_response_event": + event_data = event.data.model_dump() + if event_data["type"] == "response.output_item.added": + curr_item_id = str(uuid4()) + event_data["item"]["id"] = curr_item_id + elif event_data.get("item") is not None and event_data["item"].get("id") is not None: + event_data["item"]["id"] = curr_item_id + elif event_data.get("item_id") is not None: + event_data["item_id"] = curr_item_id + yield event_data + elif event.type == "run_item_stream_event" and event.item.type == "tool_call_output_item": + # Required for chat history: Edit/Regenerate must replay tool_call + output pairs. + yield ResponsesAgentStreamEvent( + type="response.output_item.done", + item=event.item.to_input_item(), + ) diff --git a/agent/agent_server/utils_memory.py b/agent/agent_server/utils_memory.py new file mode 100644 index 0000000..dfadaa7 --- /dev/null +++ b/agent/agent_server/utils_memory.py @@ -0,0 +1,205 @@ +# GENERATED from vendor/app-templates/agent-openai-agents-sdk/.claude/skills/managed-memory/SKILL.md +# submodule pin: 2a4c79296aae89c8a05e7825c61e0d94ad34c944 +# Regenerate verbatim by re-running the managed-memory skill; do NOT hand-edit. +# local deviations: NONE (Genie instructions composed in agent.py via genie_tools.GENIE_INSTRUCTIONS). +"""Databricks managed memory (UC memory-store) tools for the OpenAI Agents SDK.""" + +import os + +from agents import RunContextWrapper, function_tool +from databricks.sdk import WorkspaceClient +from databricks.sdk.errors import DatabricksError +from dataclasses import dataclass +from mlflow.genai.agent_server import get_request_headers +from mlflow.types.responses import ResponsesAgentRequest + +from agent_server.utils import get_user_workspace_client + +_client: WorkspaceClient | None = None + + +def _ws() -> WorkspaceClient: + global _client + if _client is None: + _client = WorkspaceClient() + return _client + + +def _entries(suffix: str = "") -> str: + store = os.getenv("DATABRICKS_MEMORY_STORE") + if not store: + raise RuntimeError( + "DATABRICKS_MEMORY_STORE is not set — it must be the full catalog.schema.name." + ) + return f"/api/2.1/unity-catalog/memory-stores/{store}/entries{suffix}" + + +def resolve_scope(request: ResponsesAgentRequest | None = None) -> str | None: + """End-user id used as memory scope; fail closed when unknown in production.""" + headers = get_request_headers() or {} + if headers.get("x-forwarded-access-token"): + return get_user_workspace_client().current_user.me().id + if os.getenv("DATABRICKS_APP_NAME"): + return None + ci = dict(getattr(request, "custom_inputs", None) or {}) + return headers.get("x-forwarded-user") or ci.get("user_id") + + +def _save(scope: str, path: str, description: str, contents: str = "") -> str: + try: + _ws().api_client.do( + "POST", + _entries(), + query={"scope": scope}, + body={ + "path": path, + "contents": contents, + "description": description, + "creation_reason": "CREATION_REASON_AGENT_INFERRED", + "creation_source": "CREATION_SOURCE_ONLINE_AGENT", + }, + ) + except DatabricksError as e: + if e.error_code == "ALREADY_EXISTS": + return f"A memory already exists at {path}; use update_memory to revise it." + return f"Could not save {path}: {getattr(e, 'message', str(e))}" + return f"Saved memory at {path}." + + +def _get(scope: str, path: str) -> str: + try: + entry = _ws().api_client.do( + "GET", _entries(":get"), query={"scope": scope, "path": path} + ) + except DatabricksError as e: + if e.error_code == "NOT_FOUND": + return f"No memory at {path}." + return f"Could not read {path}: {getattr(e, 'message', str(e))}" + return entry.get("contents") or entry.get("description") or f"(empty memory at {path})" + + +def _list(scope: str) -> str: + try: + resp = _ws().api_client.do("GET", _entries(), query={"scope": scope}) + except DatabricksError as e: + return f"Could not list memories: {getattr(e, 'message', str(e))}" + items = resp.get("entries", []) + if not items: + return "No memories yet." + lines = [ + ("[has_contents] " if e.get("has_contents") else "") + + f"- {e['path']}: {e.get('description', '')}" + for e in items + ] + return f"{len(items)} memories total:\n" + "\n".join(lines) + + +def _update( + scope: str, path: str, op: dict | None = None, description: str | None = None +) -> str: + op = op or {} + if len(op) > 1: + return "Pass at most one contents edit (str_replace / insert / replace_all)." + if not op and description is None: + return "Provide a new description and/or one contents edit." + body: dict = {"scope": scope, "path": path, **op} + if description is not None: + body["description"] = description + try: + _ws().api_client.do("PATCH", _entries(), body=body) + except DatabricksError as e: + if e.error_code == "NOT_FOUND": + return f"No memory at {path} to update — check list_memories or save it first." + return f"Could not update {path}: {getattr(e, 'message', str(e))}" + return f"Updated {path}." + + +def _delete(scope: str, path: str) -> str: + try: + _ws().api_client.do( + "DELETE", _entries(), query={"scope": scope, "path": path} + ) + except DatabricksError as e: + if e.error_code == "NOT_FOUND": + return f"No memory at {path} (already gone)." + return f"Could not delete {path}: {getattr(e, 'message', str(e))}" + return f"Deleted {path}." + + +@dataclass +class MemoryContext: + """Per-request run context; scope partitions memories by end user.""" + + scope: str + + +def _scope(ctx: RunContextWrapper[MemoryContext]) -> str: + if not ctx.context.scope: + raise RuntimeError("No end-user scope for this request — refusing a shared memory bucket.") + return ctx.context.scope + + +@function_tool(strict_mode=False) +async def save_memory( + ctx: RunContextWrapper[MemoryContext], + path: str, + description: str, + contents: str = "", +) -> str: + """Create ONE durable memory at a stable /memories/... path. Check list_memories first.""" + return _save(_scope(ctx), path, description, contents) + + +@function_tool +async def get_memory(ctx: RunContextWrapper[MemoryContext], path: str) -> str: + """Read the full contents of one memory by exact path from list_memories.""" + return _get(_scope(ctx), path) + + +@function_tool +async def list_memories(ctx: RunContextWrapper[MemoryContext]) -> str: + """List saved memories (path + description only). Call before save or recall.""" + return _list(_scope(ctx)) + + +@function_tool(strict_mode=False) +async def update_memory( + ctx: RunContextWrapper[MemoryContext], + path: str, + description: str | None = None, + str_replace: dict | None = None, + insert: dict | None = None, + replace_all: dict | None = None, +) -> str: + """Revise an existing memory. Pass description and/or exactly one contents edit op.""" + op = { + k: v + for k, v in ( + ("str_replace", str_replace), + ("insert", insert), + ("replace_all", replace_all), + ) + if v + } + return _update(_scope(ctx), path, op, description) + + +@function_tool +async def delete_memory(ctx: RunContextWrapper[MemoryContext], path: str) -> str: + """Permanently remove one memory by exact path.""" + return _delete(_scope(ctx), path) + + +MEMORY_TOOLS = [save_memory, get_memory, list_memories, update_memory, delete_memory] + +MEMORY_INSTRUCTIONS = """You have durable, cross-session memory about whoever (or whatever) this conversation is scoped to. Use it deliberately, not by reflex. + +Recall whenever the answer is about the user or calls for personalized information — anything that might draw on preferences, decisions, or workflows they've shared before — and you don't already have it from this conversation; also list once before saving, to find the right existing topic. Don't tell the user you don't know their preferences without checking — list_memories first. Skip memory only when the answer truly doesn't depend on who's asking (general knowledge, math, coding) or you already have what you need. A `[has_contents]` entry has a body to get_memory; one without is fully captured by its description. Open a memory with get_memory before you state its specifics, and never assert a fact that isn't stored — if nothing relevant is stored, just answer without it. Don't re-list what you've already seen this turn. + +Save only what will still matter in a future, unrelated conversation — a stable preference, fact, decision, or ongoing project the user actually stated or decided. Don't save your own suggestions or guesses, passing chatter, secrets, or anything scoped to this chat ("for now", a one-off label). +- Write each memory so it stands on its own out of context, under one broad, stable /memories/... topic per subject with the specifics inside it. +- Check the list first and update_memory an existing topic instead of minting a near-duplicate. +- For a very broad question that touches many memories, summarize from the list's descriptions; reserve get_memory for the specific entry you actually need. +- If the user's info changes or contradicts what's stored, update or replace it rather than keeping both — but don't rewrite a memory that already says the same thing. +- delete_memory what's stale. +- Briefly tell the user whenever you save, update, or delete.""" diff --git a/agent/databricks.yml b/agent/databricks.yml new file mode 100644 index 0000000..744df84 --- /dev/null +++ b/agent/databricks.yml @@ -0,0 +1,76 @@ +bundle: + name: firefly_openai_managed_mem + +sync: + exclude: + - vendor-wheels/** + - uv.lock + - pyproject.toml + - databricks.yml + - .python-version + - .vendor-req-export.txt + - .vendor-req-filtered.txt + - requirements.txt.volume.bak + +resources: + apps: + agent_openai_agents_sdk: + name: "firefly-openai-managed-mem-v2" + description: "OpenAI Agents SDK agent with managed memory and Genie MCP" + source_code_path: ./ + config: + command: ["python", "scripts/start_app.py"] + env: + - name: MLFLOW_TRACKING_URI + value: "databricks" + - name: MLFLOW_REGISTRY_URI + value: "databricks-uc" + - name: API_PROXY + value: "http://localhost:8000/invocations" + - name: CHAT_APP_PORT + value: "3000" + - name: CHAT_PROXY_TIMEOUT_SECONDS + value: "300" + - name: MLFLOW_EXPERIMENT_ID + value_from: "experiment" + - name: LAKEBASE_AUTOSCALING_ENDPOINT + value_from: "postgres" + - name: DATABRICKS_MEMORY_STORE + value: "main.default.firefly_managed_memory" + - name: DATABRICKS_HOST + value: "https://dbc-acc143fd-4ea9.cloud.databricks.com" + - name: DATABRICKS_WORKSPACE_ID + value: "7280163305785459" + - name: GENIE_ONE_URL + value: "https://dbc-acc143fd-4ea9.cloud.databricks.com/one?o=7280163305785459" + - name: GENIE_MCP_URL + value: "/api/2.0/mcp/genie" + - name: GENIE_MCP_MODE + value: "one" + resources: + - name: wheels_volume + uc_securable: + securable_full_name: main.default.firefly_wheels + securable_type: VOLUME + permission: READ_VOLUME + - name: experiment + experiment: + experiment_id: "123237888438046" + permission: CAN_MANAGE + - name: 'postgres' + postgres: + branch: "projects/firefly-openai-managed-mem-v2-ui-lb-202607090900/branches/firefly-openai-managed-mem-v2-ui-lb-202607090900-branch" + database: "projects/firefly-openai-managed-mem-v2-ui-lb-202607090900/branches/firefly-openai-managed-mem-v2-ui-lb-202607090900-branch/databases/databricks-postgres" + permission: 'CAN_CONNECT_AND_CREATE' + +targets: + dev: + mode: development + default: true + + prod: + mode: production + resources: + apps: + agent_openai_agents_sdk: + name: firefly-openai-managed-mem diff --git a/agent/patches/e2e-chatbot-app-next/client/src/components/chat.tsx b/agent/patches/e2e-chatbot-app-next/client/src/components/chat.tsx new file mode 100644 index 0000000..81cbb75 --- /dev/null +++ b/agent/patches/e2e-chatbot-app-next/client/src/components/chat.tsx @@ -0,0 +1,355 @@ +import type { DataUIPart, LanguageModelUsage, UIMessageChunk } from 'ai'; +import { useChat } from '@ai-sdk/react'; +import { useCallback, useEffect, useMemo, useRef, useState } from 'react'; +import { useSWRConfig } from 'swr'; +import { ChatHeader } from '@/components/chat-header'; +import { fetchWithErrorHandlers, generateUUID } from '@/lib/utils'; +import { MultimodalInput } from './multimodal-input'; +import { Messages } from './messages'; +import type { + Attachment, + ChatMessage, + CustomUIDataTypes, + FeedbackMap, + VisibilityType, +} from '@chat-template/core'; +import { unstable_serialize } from 'swr/infinite'; +import { getChatHistoryPaginationKey } from './sidebar-history'; +import { toast } from './toast'; +import { useSearchParams } from 'react-router-dom'; +import { useChatVisibility } from '@/hooks/use-chat-visibility'; +import { ChatSDKError } from '@chat-template/core/errors'; +import { useDataStream } from './data-stream-provider'; +import { isCredentialErrorMessage } from '@/lib/oauth-error-utils'; +import { ChatTransport } from '../lib/ChatTransport'; +import type { ClientSession } from '@chat-template/auth'; +import { softNavigateToChatId } from '@/lib/navigation'; +import { useAppConfig } from '@/contexts/AppConfigContext'; +import { Greeting } from './greeting'; +import { GenieAttribution } from './genie-attribution'; + +export function Chat({ + id, + initialMessages, + initialChatModel, + initialVisibilityType, + isReadonly, + initialLastContext, + feedback = {}, + title, +}: { + id: string; + initialMessages: ChatMessage[]; + initialChatModel: string; + initialVisibilityType: VisibilityType; + isReadonly: boolean; + session: ClientSession; + initialLastContext?: LanguageModelUsage; + feedback?: FeedbackMap; + title?: string; +}) { + const { visibilityType } = useChatVisibility({ + chatId: id, + initialVisibilityType, + }); + + const { mutate } = useSWRConfig(); + const { setDataStream } = useDataStream(); + const { chatHistoryEnabled, genieConfig } = useAppConfig(); + + const [input, setInput] = useState(''); + const [_usage, setUsage] = useState( + initialLastContext, + ); + + const [lastPart, setLastPart] = useState(); + const lastPartRef = useRef(lastPart); + lastPartRef.current = lastPart; + + // Single counter for resume attempts - reset when stream parts are received + const resumeAttemptCountRef = useRef(0); + const maxResumeAttempts = 3; + + const abortController = useRef(new AbortController()); + useEffect(() => { + return () => { + abortController.current?.abort('ABORT_SIGNAL'); + }; + }, []); + + const fetchWithAbort = useMemo(() => { + return async (input: RequestInfo | URL, init?: RequestInit) => { + // useChat does not cancel /stream requests when the component is unmounted + const signal = abortController.current?.signal; + return fetchWithErrorHandlers(input, { ...init, signal }); + }; + }, []); + + const stop = useCallback(() => { + abortController.current?.abort('USER_ABORT_SIGNAL'); + }, []); + + const isNewChat = initialMessages.length === 0; + const didFetchHistoryOnNewChat = useRef(false); + const fetchChatHistory = useCallback(() => { + mutate(unstable_serialize(getChatHistoryPaginationKey)); + }, [mutate]); + + // For new chats, the title arrives via a `data-title` stream part + // once backend title generation completes — no separate fetch needed. + const [streamTitle, setStreamTitle] = useState(); + const [titlePending, setTitlePending] = useState(false); + const displayTitle = title ?? streamTitle; + + const { + messages, + setMessages, + sendMessage, + status, + resumeStream, + clearError, + addToolApprovalResponse, + regenerate, + } = useChat({ + id, + messages: initialMessages, + experimental_throttle: 100, + generateId: generateUUID, + resume: id !== undefined && initialMessages.length > 0, // Enable automatic stream resumption + transport: new ChatTransport({ + onStreamPart: (part) => { + if (isNewChat && !didFetchHistoryOnNewChat.current) { + fetchChatHistory(); + if (chatHistoryEnabled) { + setTitlePending(true); + } + didFetchHistoryOnNewChat.current = true; + } + // Reset resume attempts when we successfully receive stream parts + resumeAttemptCountRef.current = 0; + setLastPart(part); + }, + api: '/api/chat', + fetch: fetchWithAbort, + prepareSendMessagesRequest({ messages, id, body }) { + const lastMessage = messages.at(-1); + const isUserMessage = lastMessage?.role === 'user'; + + // For continuations (non-user messages like tool results), we must always + // send previousMessages because the tool result only exists client-side + // and hasn't been saved to the database yet. + const needsPreviousMessages = !chatHistoryEnabled || !isUserMessage; + + return { + body: { + id, + // Only include message field for user messages (new messages) + // For continuation (assistant messages with tool results), omit message field + ...(isUserMessage ? { message: lastMessage } : {}), + selectedChatModel: initialChatModel, + selectedVisibilityType: visibilityType, + nextMessageId: generateUUID(), + // Send previous messages when: + // 1. Database is disabled (ephemeral mode) - always need client-side messages + // 2. Continuation request (tool results) - tool result only exists client-side + ...(needsPreviousMessages + ? { + previousMessages: isUserMessage + ? messages.slice(0, -1) + : messages, + } + : {}), + ...body, + }, + }; + }, + prepareReconnectToStreamRequest({ id }) { + return { + api: `/api/chat/${id}/stream`, + credentials: 'include', + }; + }, + }), + onData: (dataPart) => { + setDataStream((ds) => + ds ? [...ds, dataPart as DataUIPart] : [], + ); + if (dataPart.type === 'data-usage') { + setUsage(dataPart.data as LanguageModelUsage); + } + if (dataPart.type === 'data-title') { + setStreamTitle(dataPart.data as string); + setTitlePending(false); + fetchChatHistory(); + } + }, + onFinish: ({ + isAbort, + isDisconnect, + isError, + messages: finishedMessages, + }) => { + didFetchHistoryOnNewChat.current = false; + setTitlePending(false); + + // If user aborted, don't try to resume + if (isAbort) { + console.log('[Chat onFinish] Stream was aborted by user, not resuming'); + fetchChatHistory(); + return; + } + + // Check if the last message contains an OAuth credential error + // If so, don't try to resume - the user needs to authenticate first + const lastMessage = finishedMessages?.at(-1); + const hasOAuthError = lastMessage?.parts?.some( + (part) => + part.type === 'data-error' && + typeof part.data === 'string' && + isCredentialErrorMessage(part.data), + ); + + if (hasOAuthError) { + console.log( + '[Chat onFinish] OAuth credential error detected, not resuming', + ); + fetchChatHistory(); + clearError(); + return; + } + + // Determine if we should attempt to resume: + // 1. Stream didn't end with a 'finish' part (incomplete) + // 2. It was a disconnect/error that terminated the stream + // 3. We haven't exceeded max resume attempts + const streamIncomplete = lastPartRef.current?.type !== 'finish'; + const shouldResume = + streamIncomplete && + (isDisconnect || isError || lastPartRef.current === undefined); + + if (shouldResume && resumeAttemptCountRef.current < maxResumeAttempts) { + console.log( + '[Chat onFinish] Resuming stream. Attempt:', + resumeAttemptCountRef.current + 1, + ); + resumeAttemptCountRef.current++; + // Ref: https://github.com/vercel/ai/issues/8477#issuecomment-3603209884 + queueMicrotask(() => { + resumeStream(); + }) + } else { + // Stream completed normally or we've exhausted resume attempts + if (resumeAttemptCountRef.current >= maxResumeAttempts) { + console.warn('[Chat onFinish] Max resume attempts reached'); + } + fetchChatHistory(); + } + }, + onError: (error) => { + console.log('[Chat onError] Error occurred:', error); + + // Only show toast for explicit ChatSDKError (backend validation errors) + // Other errors (network, schema validation) are handled silently or in message parts + if (error instanceof ChatSDKError) { + toast({ + type: 'error', + description: error.message, + }); + } else { + // Non-ChatSDKError: Could be network error or in-stream error + // Log but don't toast - errors during streaming may be informational + console.warn('[Chat onError] Error during streaming:', error.message); + } + // Note: We don't call resumeStream here because onError can be called + // while the stream is still active (e.g., for data-error parts). + // Resume logic is handled exclusively in onFinish. + }, + }); + + const [searchParams] = useSearchParams(); + const query = searchParams.get('query'); + + const [hasAppendedQuery, setHasAppendedQuery] = useState(false); + + useEffect(() => { + if (query && !hasAppendedQuery) { + sendMessage({ + role: 'user' as const, + parts: [{ type: 'text', text: query }], + }); + + setHasAppendedQuery(true); + softNavigateToChatId(id, chatHistoryEnabled); + } + }, [query, sendMessage, hasAppendedQuery, id, chatHistoryEnabled]); + + const [attachments, setAttachments] = useState>([]); + + const inputElement = + + const genieFooter = genieConfig ? ( + + ) : null; + + if (messages.length === 0) { + return ( +
+ +
+
+ + {inputElement} + {genieFooter} +
+
+
+ ); + } + + return ( + <> +
+ + + + + + +
+ {!isReadonly && ( + inputElement + )} + {genieFooter} +
+
+ + ); +} diff --git a/agent/patches/e2e-chatbot-app-next/client/src/components/elements/mcp-tool.tsx b/agent/patches/e2e-chatbot-app-next/client/src/components/elements/mcp-tool.tsx new file mode 100644 index 0000000..45862d2 --- /dev/null +++ b/agent/patches/e2e-chatbot-app-next/client/src/components/elements/mcp-tool.tsx @@ -0,0 +1,261 @@ +import { Button } from '@/components/ui/button'; +import { CollapsibleTrigger } from '@/components/ui/collapsible'; +import { cn } from '@/lib/utils'; +import { + ChevronDownIcon, + ServerIcon, + ShieldAlertIcon, + ShieldCheckIcon, + ShieldXIcon, + Sparkles, +} from 'lucide-react'; +import { + ToolContainer, + ToolContent, + ToolInput, + ToolOutput, + ToolStatusBadge, + type ToolState, +} from './tool'; + +// MCP-specific container with distinct styling +type McpToolProps = Parameters[0]; + +export const McpTool = ({ className, ...props }: McpToolProps) => ( + +); + +// Re-export shared components for convenience +export { + ToolContent as McpToolContent, + ToolInput as McpToolInput, + ToolOutput as McpToolOutput, +}; + +// MCP-specific header with banner +type McpToolHeaderProps = { + serverName?: string; + toolName: string; + state: ToolState; + // Used when state is 'approval-responded' to determine approval outcome + approved?: boolean; + className?: string; +}; + +// Badge component for approval status in the banner +// Uses AI SDK native tool states directly +type ApprovalStatusBadgeProps = { + state: ToolState; + // Used when state is 'approval-responded' to determine approval outcome + approved?: boolean; +}; + +const ApprovalStatusBadge = ({ state, approved }: ApprovalStatusBadgeProps) => { + // Pending: waiting for user approval + if (state === 'approval-requested') { + return ( + + + Pending + + ); + } + + // Allowed: tool executed successfully or user approved (waiting for execution) + if ( + state === 'output-available' || + (state === 'approval-responded' && approved === true) + ) { + return ( + + + Allowed + + ); + } + + // Denied: user explicitly denied the tool + if ( + state === 'output-denied' || + (state === 'approval-responded' && approved === false) + ) { + return ( + + + Denied + + ); + } + + // Fallback for any other state - show as pending + return ( + + + Pending + + ); +}; + +function isGenieServer(serverName?: string): boolean { + return !!serverName && /genie/i.test(serverName); +} + +export const McpToolHeader = ({ + className, + serverName, + toolName, + state, + approved, +}: McpToolHeaderProps) => ( +
+ {/* MCP Banner */} +
+ + + Tool Call Request + + {isGenieServer(serverName) ? ( + <> + • + + + Powered by Genie + + + ) : ( + serverName && ( + <> + • + {serverName} + + ) + )} + • + +
+ {/* Tool header */} + +
+ {toolName} +
+
+ {/* Only show tool status badge when tool is running/completed (approved) */} + {(state === 'output-available' || + (state === 'approval-responded' && approved === true)) && ( + + )} + +
+ +
+); + +// MCP-specific approval actions +type McpApprovalActionsProps = { + onApprove: () => void; + onDeny: () => void; + isSubmitting: boolean; +}; + +export const McpApprovalActions = ({ + onApprove, + onDeny, + isSubmitting, +}: McpApprovalActionsProps) => ( +
+
+ +

+ This tool requires your permission to run. +

+
+
+ + +
+
+); + +// MCP-specific approval status display +type McpApprovalStatusProps = { + approved: boolean; + reason?: string; +}; + +export const McpApprovalStatus = ({ + approved, + reason, +}: McpApprovalStatusProps) => ( +
+ {approved ? ( + + ) : ( + + )} + + {approved ? 'Allowed' : 'Denied'} + + {reason && ( + <> + • + {reason} + + )} +
+); diff --git a/agent/patches/e2e-chatbot-app-next/client/src/components/genie-attribution.tsx b/agent/patches/e2e-chatbot-app-next/client/src/components/genie-attribution.tsx new file mode 100644 index 0000000..83885a0 --- /dev/null +++ b/agent/patches/e2e-chatbot-app-next/client/src/components/genie-attribution.tsx @@ -0,0 +1,95 @@ +import { ExternalLink, Sparkles } from 'lucide-react'; +import { cn } from '@/lib/utils'; +import { buildGenieOneUrl } from '@/lib/genie-attribution'; + +export type GenieAttributionProps = { + genieOneUrl?: string | null; + workspaceHost: string; + workspaceId?: string; + links?: string[]; + variant?: 'inline' | 'footer' | 'compact'; + className?: string; +}; + +function linkLabel(url: string, index: number): string { + try { + const parsed = new URL(url); + if (parsed.pathname.includes('/dashboard')) { + return 'AI/BI dashboard'; + } + if (parsed.pathname.includes('/genie/')) { + return 'Genie Agent'; + } + const path = parsed.pathname.split('/').filter(Boolean).slice(-2).join('/'); + return path || parsed.hostname || `Related link ${index + 1}`; + } catch { + return `Related link ${index + 1}`; + } +} + +export function GenieAttribution({ + genieOneUrl, + workspaceHost, + workspaceId, + links = [], + variant = 'inline', + className, +}: GenieAttributionProps) { + const oneUrl = + genieOneUrl ?? buildGenieOneUrl(workspaceHost, workspaceId) ?? null; + const uniqueLinks = [...new Set(links.filter((link) => link.startsWith('http')))]; + + return ( +
+
+ + + Powered by Genie + + {oneUrl && ( + <> + • + + Genie One + + + + )} +
+ {uniqueLinks.length > 0 && ( + + )} +
+ ); +} diff --git a/agent/patches/e2e-chatbot-app-next/client/src/components/message.tsx b/agent/patches/e2e-chatbot-app-next/client/src/components/message.tsx new file mode 100644 index 0000000..2765bb9 --- /dev/null +++ b/agent/patches/e2e-chatbot-app-next/client/src/components/message.tsx @@ -0,0 +1,534 @@ +import React, { memo, useState } from 'react'; +import { AnimatedAssistantIcon } from './animation-assistant-icon'; +import { Response } from './elements/response'; +import { MessageContent } from './elements/message'; +import { + Tool, + ToolHeader, + ToolContent, + ToolInput, + ToolOutput, + type ToolState, +} from './elements/tool'; +import { + McpTool, + McpToolHeader, + McpToolContent, + McpToolInput, + McpApprovalActions, +} from './elements/mcp-tool'; +import { MessageActions } from './message-actions'; +import { PreviewAttachment } from './preview-attachment'; +import equal from 'fast-deep-equal'; +import { cn, sanitizeText, formatToolOutput } from '@/lib/utils'; +import { MessageEditor } from './message-editor'; +import { MessageReasoning } from './message-reasoning'; +import { Shimmer } from './ui/shimmer'; +import type { UseChatHelpers } from '@ai-sdk/react'; +import type { ChatMessage, Feedback } from '@chat-template/core'; +import { useDataStream } from './data-stream-provider'; +import { + createMessagePartSegments, + formatNamePart, + isNamePart, + joinMessagePartSegments, +} from './databricks-message-part-transformers'; +import { MessageError } from './message-error'; +import { MessageOAuthError } from './message-oauth-error'; +import { isCredentialErrorMessage } from '@/lib/oauth-error-utils'; +import { + groupConsecutiveToolSegments, + type ChatPart, + type ToolPart, +} from '@/lib/tool-group-segments'; +import { Streamdown } from 'streamdown'; +import { useApproval } from '@/hooks/use-approval'; +import { useAppConfig } from '@/contexts/AppConfigContext'; +import { GenieAttribution } from './genie-attribution'; +import { + collectGenieLinksFromMessage, + extractGenieLinks, + isGenieToolPart, +} from '@/lib/genie-attribution'; + +const PurePreviewMessage = ({ + message, + allMessages, + isLoading, + setMessages, + addToolApprovalResponse, + sendMessage, + regenerate, + isReadonly, + requiresScrollPadding, + initialFeedback, +}: { + message: ChatMessage; + allMessages: ChatMessage[]; + isLoading: boolean; + setMessages: UseChatHelpers['setMessages']; + addToolApprovalResponse: UseChatHelpers['addToolApprovalResponse']; + sendMessage: UseChatHelpers['sendMessage']; + regenerate: UseChatHelpers['regenerate']; + isReadonly: boolean; + requiresScrollPadding: boolean; + initialFeedback?: Feedback; +}) => { + const [mode, setMode] = useState<'view' | 'edit'>('view'); + const [showErrors, setShowErrors] = useState(false); + + // Hook for handling MCP approval requests + const { submitApproval, isSubmitting, pendingApprovalId } = useApproval({ + addToolApprovalResponse, + sendMessage, + }); + const { genieConfig } = useAppConfig(); + const genieLinks = React.useMemo( + () => collectGenieLinksFromMessage(message.parts), + [message.parts], + ); + const showGenieAttribution = + message.role === 'assistant' && !!genieConfig; + + const attachmentsFromMessage = message.parts.filter( + (part) => part.type === 'file', + ); + + // Extract non-OAuth error parts separately (OAuth errors are rendered inline) + const errorParts = React.useMemo( + () => + message.parts + .filter((part) => part.type === 'data-error') + .filter((part) => { + // OAuth errors are rendered inline, not in the error section + return !isCredentialErrorMessage(part.data); + }), + [message.parts], + ); + + useDataStream(); + + const partSegments = React.useMemo( + /** + * We segment message parts into segments that can be rendered as a single component. + * Used to render citations as part of the associated text. + * Note: OAuth errors are included here for inline rendering, non-OAuth errors are filtered out. + */ + () => + createMessagePartSegments( + message.parts.filter( + (part) => + part.type !== 'data-error' || isCredentialErrorMessage(part.data), + ), + ), + [message.parts], + ); + + const renderBlocks = React.useMemo( + () => groupConsecutiveToolSegments(partSegments), + [partSegments], + ); + + // Check if message only contains non-OAuth errors (no other content) + const hasOnlyErrors = React.useMemo(() => { + const nonErrorParts = message.parts.filter( + (part) => part.type !== 'data-error', + ); + // Only consider non-OAuth errors for this check + return errorParts.length > 0 && nonErrorParts.length === 0; + }, [message.parts, errorParts.length]); + + return ( +
+
+ {partSegments.length === 0 && errorParts.length === 0 && message.role === 'assistant' && ( + + )} + +
+ {attachmentsFromMessage.length > 0 && ( +
+ {attachmentsFromMessage.map((attachment) => ( + + ))} +
+ )} + + {renderBlocks.map((block) => { + if (block.kind === 'tool-group') { + return ( + + ); + } + + const parts = block.parts; + const index = block.index; + const [part] = parts; + const { type } = part; + const key = `message-${message.id}-part-${index}`; + + if (type === 'reasoning' && part.text?.trim().length > 0) { + return ( + + ); + } + + if (type === 'text') { + if (isNamePart(part)) { + return ( + {`# ${formatNamePart(part)}`} + ); + } + if (mode === 'view') { + return ( +
+ + + {sanitizeText(joinMessagePartSegments(parts))} + + +
+ ); + } + + if (mode === 'edit') { + return ( +
+
+
+ +
+
+ ); + } + } + + // dynamic-tool parts are rendered by MessageToolGroup above. + + // Support for citations/annotations + if (type === 'source-url') { + return ( + + [{part.title || part.url}] + + ); + } + + // Render OAuth errors inline + if (type === 'data-error' && isCredentialErrorMessage(part.data)) { + return ( + + ); + } + })} + + {showGenieAttribution && genieConfig && ( + + )} + + {!isReadonly && !hasOnlyErrors && ( + setShowErrors(!showErrors)} + initialFeedback={initialFeedback} + /> + )} + + {errorParts.length > 0 && (hasOnlyErrors || showErrors) && ( +
+ {errorParts.map((part, index) => ( + + ))} +
+ )} +
+
+
+ ); +}; + +export const PreviewMessage = memo( + PurePreviewMessage, + (prevProps, nextProps) => { + if (prevProps.isLoading !== nextProps.isLoading) return false; + // While streaming, re-render whenever the AI SDK produces a new message + // object (each throttled update). We use reference equality rather than + // deep-equal on parts because fast-deep-equal short-circuits on identical + // references — and the SDK may mutate parts in place during streaming. + if (nextProps.isLoading && prevProps.message !== nextProps.message) + return false; + + if (prevProps.message.id !== nextProps.message.id) return false; + if (prevProps.requiresScrollPadding !== nextProps.requiresScrollPadding) + return false; + if (!equal(prevProps.message.parts, nextProps.message.parts)) return false; + if (prevProps.initialFeedback?.feedbackType !== nextProps.initialFeedback?.feedbackType) + return false; + + return true; // Props are equal, skip re-render + }, +); + + +const MessageToolGroup = ({ + tools, + isLoading, + submitApproval, + isSubmitting, + pendingApprovalId, +}: { + tools: ToolPart[]; + isLoading: boolean; + submitApproval: ReturnType['submitApproval']; + isSubmitting: boolean; + pendingApprovalId: string | null; +}) => { + const isMultiple = tools.length > 1; + return ( +
+ {tools.map((tool) => ( + + ))} +
+ ); +}; + +const ToolPartRenderer = ({ + part, + isLoading, + submitApproval, + isSubmitting, + pendingApprovalId, +}: { + part: ToolPart; + isLoading: boolean; + submitApproval: ReturnType['submitApproval']; + isSubmitting: boolean; + pendingApprovalId: string | null; +}) => { + const { genieConfig } = useAppConfig(); + const { toolCallId, input, state, errorText, output, toolName } = part; + + const isMcpApproval = + part.callProviderMetadata?.databricks?.approvalRequestId != null; + const mcpServerName = + part.callProviderMetadata?.databricks?.mcpServerName?.toString(); + const isGenieTool = isGenieToolPart(part); + const genieLinks = + isGenieTool && output != null ? extractGenieLinks(output) : []; + + const approved: boolean | undefined = + 'approval' in part ? part.approval?.approved : undefined; + + const effectiveState: ToolState = (() => { + if (part.providerExecuted && !isLoading && state === 'input-available') { + return 'output-available'; + } + return state; + })(); + + if (isMcpApproval) { + return ( + + + + + {state === 'approval-requested' && ( + + submitApproval({ approvalRequestId: toolCallId, approve: true }) + } + onDeny={() => + submitApproval({ + approvalRequestId: toolCallId, + approve: false, + }) + } + isSubmitting={isSubmitting && pendingApprovalId === toolCallId} + /> + )} + {state === 'output-available' && output != null && ( + + Error: {errorText} +
+ ) : ( +
+
+ {formatToolOutput(output)} +
+ {isGenieTool && genieConfig && ( + + )} +
+ ) + } + errorText={undefined} + /> + )} + + + ); + } + + return ( + + + + + {state === 'output-available' && ( + + Error: {errorText} + + ) : ( +
+
+ {formatToolOutput(output)} +
+ {isGenieTool && genieConfig && ( + + )} +
+ ) + } + errorText={undefined} + /> + )} + + + ); +}; + +export const AwaitingResponseMessage = () => { + const role = 'assistant'; + + return ( +
+
+ Generating response +
+
+ ); +}; diff --git a/agent/patches/e2e-chatbot-app-next/client/src/contexts/AppConfigContext.tsx b/agent/patches/e2e-chatbot-app-next/client/src/contexts/AppConfigContext.tsx new file mode 100644 index 0000000..4067df3 --- /dev/null +++ b/agent/patches/e2e-chatbot-app-next/client/src/contexts/AppConfigContext.tsx @@ -0,0 +1,70 @@ +import { createContext, useContext, type ReactNode } from 'react'; +import useSWR from 'swr'; +import { fetcher } from '@/lib/utils'; + +export interface GenieConfig { + workspaceHost: string; + workspaceId?: string; + genieOneUrl?: string; +} + +interface ConfigResponse { + features: { + chatHistory: boolean; + feedback: boolean; + }; + obo?: { + missingScopes: string[]; + }; + genie?: GenieConfig | null; +} + +interface AppConfigContextType { + config: ConfigResponse | undefined; + isLoading: boolean; + error: Error | undefined; + chatHistoryEnabled: boolean; + feedbackEnabled: boolean; + oboMissingScopes: string[]; + genieConfig: GenieConfig | null; +} + +const AppConfigContext = createContext( + undefined, +); + +export function AppConfigProvider({ children }: { children: ReactNode }) { + const { data, error, isLoading } = useSWR( + '/api/config', + fetcher, + { + revalidateOnFocus: false, + revalidateOnReconnect: false, + dedupingInterval: 60000, + }, + ); + + const value: AppConfigContextType = { + config: data, + isLoading, + error, + chatHistoryEnabled: data?.features.chatHistory ?? true, + feedbackEnabled: data?.features.feedback ?? false, + oboMissingScopes: data?.obo?.missingScopes ?? [], + genieConfig: data?.genie ?? null, + }; + + return ( + + {children} + + ); +} + +export function useAppConfig() { + const context = useContext(AppConfigContext); + if (context === undefined) { + throw new Error('useAppConfig must be used within an AppConfigProvider'); + } + return context; +} diff --git a/agent/patches/e2e-chatbot-app-next/client/src/lib/genie-attribution.ts b/agent/patches/e2e-chatbot-app-next/client/src/lib/genie-attribution.ts new file mode 100644 index 0000000..63e1171 --- /dev/null +++ b/agent/patches/e2e-chatbot-app-next/client/src/lib/genie-attribution.ts @@ -0,0 +1,121 @@ +import type { ChatMessage } from '@chat-template/core'; + +const URL_RE = /https?:\/\/[^\s"'<>)\]]+/gi; +const MARKDOWN_LINK_RE = /\[[^\]]*\]\((https?:\/\/[^)]+)\)/gi; + +export function isGenieServerName(name?: string): boolean { + if (!name) return false; + return /genie/i.test(name); +} + +export function isGenieToolPart( + part: ChatMessage['parts'][number], +): part is Extract { + if (part.type !== 'dynamic-tool') return false; + const server = part.callProviderMetadata?.databricks?.mcpServerName?.toString(); + if (isGenieServerName(server)) return true; + const toolName = part.toolName ?? ''; + return /genie/i.test(toolName); +} + +export function messageUsesGenie(parts: ChatMessage['parts']): boolean { + return parts.some(isGenieToolPart); +} + +export function buildGenieOneUrl( + workspaceHost: string, + workspaceId?: string, + explicitUrl?: string, +): string | null { + if (explicitUrl?.trim()) { + return explicitUrl.trim(); + } + if (!workspaceHost) return null; + const host = workspaceHost.replace(/\/$/, ''); + if (workspaceId?.trim()) { + return `${host}/one?o=${workspaceId.trim()}`; + } + return `${host}/one`; +} + +function normalizeUrl(url: string): string { + return url.replace(/[.,;]+$/g, ''); +} + +function isNavigableGenieLink(url: string): boolean { + try { + const parsed = new URL(url); + return parsed.protocol === 'http:' || parsed.protocol === 'https:'; + } catch { + return false; + } +} + +export function extractMarkdownLinks(text: string): string[] { + const urls: string[] = []; + for (const match of text.matchAll(MARKDOWN_LINK_RE)) { + const url = match[1]; + if (url && isNavigableGenieLink(url)) { + urls.push(normalizeUrl(url)); + } + } + return urls; +} + +function collectUrls(value: unknown, out: Set, depth = 0): void { + if (depth > 12 || value == null) return; + + if (typeof value === 'string') { + extractMarkdownLinks(value).forEach((url) => out.add(url)); + const matches = value.match(URL_RE); + matches?.forEach((match) => { + const normalized = normalizeUrl(match); + if (isNavigableGenieLink(normalized)) { + out.add(normalized); + } + }); + return; + } + + if (Array.isArray(value)) { + value.forEach((item) => collectUrls(item, out, depth + 1)); + return; + } + + if (typeof value === 'object') { + for (const [key, nested] of Object.entries(value as Record)) { + if ( + /(url|link|href|dashboard|visualization|citation)/i.test(key) && + typeof nested === 'string' && + nested.startsWith('http') + ) { + out.add(normalizeUrl(nested)); + } + collectUrls(nested, out, depth + 1); + } + } +} + +export function extractGenieLinks(output: unknown): string[] { + const urls = new Set(); + collectUrls(output, urls); + return [...urls]; +} + +export function collectGenieLinksFromMessage( + parts: ChatMessage['parts'], +): string[] { + const urls = new Set(); + for (const part of parts) { + if (isGenieToolPart(part) && part.output != null) { + extractGenieLinks(part.output).forEach((url) => urls.add(url)); + } + if (part.type === 'text' && typeof part.text === 'string') { + extractGenieLinks(part.text).forEach((url) => urls.add(url)); + } + if (part.type === 'source-url' && typeof part.url === 'string') { + urls.add(normalizeUrl(part.url)); + } + } + return [...urls]; +} diff --git a/agent/patches/e2e-chatbot-app-next/client/src/lib/utils.ts b/agent/patches/e2e-chatbot-app-next/client/src/lib/utils.ts new file mode 100644 index 0000000..e130743 --- /dev/null +++ b/agent/patches/e2e-chatbot-app-next/client/src/lib/utils.ts @@ -0,0 +1,129 @@ +import type { UIMessagePart } from 'ai'; +import { type ClassValue, clsx } from 'clsx'; +import { twMerge } from 'tailwind-merge'; +import type { DBMessage } from '@chat-template/db'; +import { ChatSDKError, type ErrorCode } from '@chat-template/core/errors'; +import type { + ChatMessage, + ChatTools, + CustomUIDataTypes, +} from '@chat-template/core'; +import { formatISO } from 'date-fns'; + +export function cn(...inputs: ClassValue[]) { + return twMerge(clsx(inputs)); +} + +export const fetcher = async (url: string) => { + const response = await fetch(url); + + if (!response.ok) { + const { code, cause } = await response.json(); + throw new ChatSDKError(code as ErrorCode, cause); + } + + // Handle 204 No Content - return empty chat history response + if (response.status === 204) { + return { chats: [], hasMore: false }; + } + + return response.json(); +}; + +export async function fetchWithErrorHandlers( + input: RequestInfo | URL, + init?: RequestInit, +) { + try { + const response = await fetch(input, init); + + if (!response.ok) { + const parsedResponse = await response.json(); + console.log('parsedResponse', parsedResponse); + const { code, cause } = parsedResponse; + throw new ChatSDKError(code as ErrorCode, cause); + } + + return response; + } catch (error: unknown) { + if (typeof navigator !== 'undefined' && !navigator.onLine) { + throw new ChatSDKError('offline:chat'); + } + + throw error; + } +} + +export function generateUUID(): string { + return 'xxxxxxxx-xxxx-4xxx-yxxx-xxxxxxxxxxxx'.replace(/[xy]/g, (c) => { + const r = (Math.random() * 16) | 0; + const v = c === 'x' ? r : (r & 0x3) | 0x8; + return v.toString(16); + }); +} + +export function sanitizeText(text: string) { + return text.replace('', ''); +} + +export function convertToUIMessages(messages: DBMessage[]): ChatMessage[] { + return messages.map((message) => ({ + id: message.id, + role: message.role as 'user' | 'assistant' | 'system', + parts: message.parts as UIMessagePart[], + metadata: { + createdAt: formatISO(message.createdAt), + }, + })); +} + +export function getTextFromMessage(message: ChatMessage): string { + return message.parts + .filter((part) => part.type === 'text') + .map((part) => part.text) + .join(''); +} + +/** Best-effort human-readable rendering for MCP / Genie tool outputs. */ +export function formatToolOutput(output: unknown): string { + if (output == null) return ''; + if (typeof output === 'string') { + const trimmed = output.trim(); + if (trimmed.startsWith('{') || trimmed.startsWith('[')) { + try { + return formatToolOutput(JSON.parse(trimmed)); + } catch { + return output; + } + } + return output; + } + + if (Array.isArray(output)) { + const parts = output + .map((item) => formatToolOutput(item)) + .filter((s) => s.length > 0); + return parts.join('\n\n'); + } + + if (typeof output === 'object') { + const record = output as Record; + if (typeof record.text === 'string' && record.text.trim()) { + return record.text; + } + if (typeof record.message === 'string' && record.message.trim()) { + return record.message; + } + if (typeof record.result === 'string' && record.result.trim()) { + return record.result; + } + if (Array.isArray(record.content)) { + return formatToolOutput(record.content); + } + if (record.type === 'text' && typeof record.text === 'string') { + return record.text; + } + } + + return JSON.stringify(output, null, 2); +} diff --git a/agent/patches/e2e-chatbot-app-next/client/src/main.tsx b/agent/patches/e2e-chatbot-app-next/client/src/main.tsx new file mode 100644 index 0000000..86c17d7 --- /dev/null +++ b/agent/patches/e2e-chatbot-app-next/client/src/main.tsx @@ -0,0 +1,44 @@ +import ReactDOM from 'react-dom/client'; +import { BrowserRouter } from 'react-router-dom'; +import App from './App'; +import './index.css'; + +// Firefly proxy-embed shim (source-side; keeps the reverse proxy stock). +// When embedded, the iframe is served under a dynamic mount: +// - Vercel-native route: /api/agent-proxy/... +// - Go proxy (optional): /app-proxy/{toolId}/... +// Derive that mount prefix from the current path and (1) drive the router +// basename and (2) prefix absolute /api/ fetches so they route back through +// the proxy instead of hitting the frontend origin root. +declare global { + interface Window { + __FIREFLY_PROXY_BASENAME__?: string; + } +} + +const proxyPrefix = + window.location.pathname.match(/^\/(?:app-proxy\/[^/]+|api\/agent-proxy)/)?.[0] ?? + ''; +const routerBasename = proxyPrefix || '/'; + +if (proxyPrefix) { + window.__FIREFLY_PROXY_BASENAME__ = proxyPrefix; + const originalFetch = window.fetch.bind(window); + window.fetch = (input, init) => { + if (typeof input === 'string' && input.startsWith('/api/')) { + input = proxyPrefix + input; + } + return originalFetch(input, init); + }; +} + +const rootElement = document.getElementById('root'); +if (!rootElement) { + throw new Error('Failed to find the root element with ID "root"'); +} + +ReactDOM.createRoot(rootElement).render( + + + , +); diff --git a/agent/patches/e2e-chatbot-app-next/client/vite.config.ts b/agent/patches/e2e-chatbot-app-next/client/vite.config.ts new file mode 100644 index 0000000..de3798f --- /dev/null +++ b/agent/patches/e2e-chatbot-app-next/client/vite.config.ts @@ -0,0 +1,54 @@ +import { defineConfig } from "vite"; +import react from "@vitejs/plugin-react"; +import path from "node:path"; +import type { ProxyOptions } from "vite"; + +export function simulateNetworkError( + _timeout: number, +): ProxyOptions["configure"] { + return (proxy, _options) => { + proxy.on("proxyReq", (proxyReq, _req, res) => { + setTimeout(() => { + // Destroy the socket connection to the browser + res.socket?.destroy(new Error("simulated network error")); + + // Also clean up the backend connection + proxyReq.socket?.destroy(); + }, 2000); + }); + }; +} + +const proxyTarget = "http://localhost:3001"; + +// https://vitejs.dev/config/ +export default defineConfig({ + // Firefly proxy-embed: emit relative asset URLs so they resolve under the + // dynamic /app-proxy/{toolId}/ mount instead of the proxy origin root. + base: "./", + plugins: [react()], + resolve: { + alias: { + "@": path.resolve(__dirname, "./src"), + }, + }, + server: { + port: process.env.PORT ? Number.parseInt(process.env.PORT) : 3000, + proxy: { + "/api/chat": { + target: proxyTarget, + changeOrigin: true, + // Uncomment this to test situations where the stream will time out. + // configure: simulateNetworkError(2000), + }, + "/api": { + target: process.env.BACKEND_URL || "http://localhost:3001", + changeOrigin: true, + }, + }, + }, + build: { + outDir: "dist", + sourcemap: false, + }, +}); diff --git a/agent/patches/e2e-chatbot-app-next/server/src/lib/sanitize-ui-messages.ts b/agent/patches/e2e-chatbot-app-next/server/src/lib/sanitize-ui-messages.ts new file mode 100644 index 0000000..941f22c --- /dev/null +++ b/agent/patches/e2e-chatbot-app-next/server/src/lib/sanitize-ui-messages.ts @@ -0,0 +1,50 @@ +import type { ChatMessage } from '@chat-template/core'; + +const COMPLETED_TOOL_STATES = new Set([ + 'output-available', + 'output-denied', + 'output-error', +]); + +/** + * Drop assistant tool parts that never received outputs. Orphan function_call + * items break OpenAI replay (Edit / Regenerate) with "tool_call_ids did not + * have response messages". + */ +export function sanitizeUiMessagesForModel( + messages: ChatMessage[], +): ChatMessage[] { + return messages + .map((message) => { + if (message.role !== 'assistant' || !message.parts?.length) { + return message; + } + + const parts = message.parts.filter((part) => { + if (part.type !== 'dynamic-tool') { + return true; + } + return COMPLETED_TOOL_STATES.has(part.state); + }); + + if (parts.length === message.parts.length) { + return message; + } + + return { ...message, parts }; + }) + .filter((message) => { + if (message.role !== 'assistant') { + return true; + } + if (!message.parts?.length) { + return false; + } + return message.parts.some( + (part) => + part.type === 'text' || + part.type === 'reasoning' || + part.type === 'dynamic-tool', + ); + }); +} diff --git a/agent/patches/e2e-chatbot-app-next/server/src/routes/chat.ts b/agent/patches/e2e-chatbot-app-next/server/src/routes/chat.ts new file mode 100644 index 0000000..dab5761 --- /dev/null +++ b/agent/patches/e2e-chatbot-app-next/server/src/routes/chat.ts @@ -0,0 +1,620 @@ +import { + Router, + type Request, + type Response, + type Router as RouterType, +} from 'express'; +import { + convertToModelMessages, + createUIMessageStream, + streamText, + generateText, + type LanguageModelUsage, + pipeUIMessageStreamToResponse, +} from 'ai'; +import type { LanguageModelV3Usage } from '@ai-sdk/provider'; + +// Convert ai's LanguageModelUsage to @ai-sdk/provider's LanguageModelV3Usage +function toV3Usage(usage: LanguageModelUsage): LanguageModelV3Usage { + return { + inputTokens: { + total: usage.inputTokens, + noCache: undefined, + cacheRead: undefined, + cacheWrite: undefined, + }, + outputTokens: { + total: usage.outputTokens, + text: undefined, + reasoning: undefined, + }, + }; +} +import { + authMiddleware, + requireAuth, + requireChatAccess, + getIdFromRequest, +} from '../middleware/auth'; +import { + deleteChatById, + getMessagesByChatId, + saveChat, + saveMessages, + updateChatLastContextById, + updateChatVisiblityById, + isDatabaseAvailable, + updateChatTitleById, +} from '@chat-template/db'; +import { + type ChatMessage, + checkChatAccess, + convertToUIMessages, + generateUUID, + myProvider, + postRequestBodySchema, + type PostRequestBody, + StreamCache, + type VisibilityType, + CONTEXT_HEADER_CONVERSATION_ID, + CONTEXT_HEADER_USER_ID, +} from '@chat-template/core'; +import { ChatSDKError } from '@chat-template/core/errors'; +import { storeMessageMeta } from '../lib/message-meta-store'; +import { drainStreamToWriter, fallbackToGenerateText } from '../lib/stream-fallback'; +import { sanitizeUiMessagesForModel } from '../lib/sanitize-ui-messages'; + +export const chatRouter: RouterType = Router(); + +const streamCache = new StreamCache(); +// Apply auth middleware to all chat routes +chatRouter.use(authMiddleware); + +/** + * POST /api/chat - Send a message and get streaming response + * + * Note: Works in ephemeral mode when database is disabled. + * Streaming continues normally, but no chat/message persistence occurs. + */ +chatRouter.post('/', requireAuth, async (req: Request, res: Response) => { + const dbAvailable = isDatabaseAvailable(); + if (!dbAvailable) { + console.log('[Chat] Running in ephemeral mode - no persistence'); + } + + let requestBody: PostRequestBody; + + try { + requestBody = postRequestBodySchema.parse(req.body); + } catch (_) { + console.error('Error parsing request body:', _); + const error = new ChatSDKError('bad_request:api'); + const response = error.toResponse(); + return res.status(response.status).json(response.json); + } + + try { + const { + id, + message, + selectedChatModel, + selectedVisibilityType, + }: { + id: string; + message?: ChatMessage; + selectedChatModel: string; + selectedVisibilityType: VisibilityType; + } = requestBody; + + const session = req.session; + if (!session) { + const error = new ChatSDKError('unauthorized:chat'); + const response = error.toResponse(); + return res.status(response.status).json(response.json); + } + + const { chat, allowed, reason } = await checkChatAccess( + id, + session?.user.id, + ); + + if (reason !== 'not_found' && !allowed) { + const error = new ChatSDKError('forbidden:chat'); + const response = error.toResponse(); + return res.status(response.status).json(response.json); + } + + let titlePromise: Promise | undefined; + + if (!chat) { + // Only create new chat if we have a message (not a continuation) + if (isDatabaseAvailable() && message) { + await saveChat({ + id, + userId: session.user.id, + title: 'New chat', + visibility: selectedVisibilityType, + }); + + titlePromise = generateTitleFromUserMessage({ message }) + .then(async (title) => { + await updateChatTitleById({ chatId: id, title }); + return title; + }) + .catch(async (error) => { + console.error('Error generating title:', error); + const textFromUserMessage = message?.parts.find( + (part) => part.type === 'text', + )?.text; + if (textFromUserMessage) { + const fallback = truncatePreserveWords( + textFromUserMessage, + 128, + ); + await updateChatTitleById({ chatId: id, title: fallback }); + return fallback; + } + return null; + }); + } + } else { + if (chat.userId !== session.user.id) { + const error = new ChatSDKError('forbidden:chat'); + const response = error.toResponse(); + return res.status(response.status).json(response.json); + } + } + + const messagesFromDb = await getMessagesByChatId({ id }); + + // Use previousMessages from request body when: + // 1. Ephemeral mode (DB not available) - always use client-side messages + // 2. Continuation request (no message) - tool results only exist client-side + const useClientMessages = + !dbAvailable || (!message && requestBody.previousMessages); + const previousMessages = useClientMessages + ? (requestBody.previousMessages ?? []) + : convertToUIMessages(messagesFromDb); + + // If message is provided, add it to the list and save it + // If not (continuation/regeneration), just use previous messages + let uiMessages: ChatMessage[]; + if (message) { + uiMessages = [...previousMessages, message]; + await saveMessages({ + messages: [ + { + chatId: id, + id: message.id, + role: 'user', + parts: message.parts, + attachments: [], + createdAt: new Date(), + traceId: null, + }, + ], + }); + } else { + // Continuation: use existing messages without adding new user message + uiMessages = previousMessages as ChatMessage[]; + + // For continuations with database enabled, save any updated assistant messages + // This ensures tool-result parts (like MCP approval responses) are persisted + if (dbAvailable && requestBody.previousMessages) { + const assistantMessages = requestBody.previousMessages.filter( + (m: ChatMessage) => m.role === 'assistant', + ); + if (assistantMessages.length > 0) { + await saveMessages({ + messages: assistantMessages.map((m: ChatMessage) => ({ + chatId: id, + id: m.id, + role: m.role, + parts: m.parts, + attachments: [], + createdAt: m.metadata?.createdAt + ? new Date(m.metadata.createdAt) + : new Date(), + traceId: null, + })), + }); + + // Check if this is an MCP denial - if so, we're done (no need to call LLM) + // Denial is indicated by a dynamic-tool part with state 'output-denied' + // or with approval.approved === false + const hasMcpDenial = requestBody.previousMessages?.some( + (m: ChatMessage) => + m.parts?.some( + (p) => + p.type === 'dynamic-tool' && + (p.state === 'output-denied' || + ('approval' in p && p.approval?.approved === false)), + ), + ); + + if (hasMcpDenial) { + // We don't need to call the LLM because the user has denied the tool call + res.end(); + return; + } + } + } + } + + // Clear any previous active stream for this chat + streamCache.clearActiveStream(id); + + let finalUsage: LanguageModelUsage | undefined; + let traceId: string | null = null; + const streamId = generateUUID(); + + const model = await myProvider.languageModel(selectedChatModel); + const modelMessages = await convertToModelMessages( + sanitizeUiMessagesForModel(uiMessages), + ); + const requestHeaders = { + [CONTEXT_HEADER_CONVERSATION_ID]: id, + [CONTEXT_HEADER_USER_ID]: session.user.email ?? session.user.id, + // Forward OBO user token to the backend/serving endpoint + ...(req.headers['x-forwarded-access-token'] + ? { 'x-forwarded-access-token': req.headers['x-forwarded-access-token'] as string } + : {}), + }; + + const result = streamText({ + model, + messages: modelMessages, + providerOptions: { + databricks: { includeTrace: true }, + }, + includeRawChunks: true, + headers: requestHeaders, + onChunk: ({ chunk }) => { + if (chunk.type === 'raw') { + const raw = chunk.rawValue as any; + // Extract trace in Databricks serving endpoint output format, if present + if (raw?.type === 'response.output_item.done') { + const traceIdFromChunk = + raw?.databricks_output?.trace?.info?.trace_id; + if (typeof traceIdFromChunk === 'string') { + traceId = traceIdFromChunk; + } + } + // Extract trace from MLflow AgentServer output format, if present + if (!traceId && typeof raw?.trace_id === 'string') { + traceId = raw.trace_id; + } + } + }, + onFinish: ({ usage }) => { + finalUsage = usage; + }, + }); + + /** + * We manually read from toUIMessageStream instead of using writer.merge + * so the execute promise (and thus the outer stream) stays alive if we + * need to fall back to generateText after a streaming error. + */ + const stream = createUIMessageStream({ + // Pass originalMessages so that continuation responses reuse the existing + // assistant message ID. Without this, handleUIMessageStreamFinish generates + // a fresh ID, causing the client to push a second assistant message instead + // of replacing the existing one. + originalMessages: uiMessages, + // The DB Message.id column is typed as uuid, so we must generate UUIDs + // rather than the AI SDK's default short-id format (e.g. "Xt8nZiQRj1fS4yiU"). + generateId: generateUUID, + execute: async ({ writer }) => { + // Manually drain the AI stream so we can append the traceId data part + // after all model chunks are processed (traceId is captured via onChunk). + // result.toUIMessageStream() converts TextStreamPart → UIMessageChunk: + // - text-delta: maps TextStreamPart.text → UIMessageChunk.delta + // - start-step/finish-step: strips extra fields + // - finish: strips rawFinishReason/totalUsage + // - raw: dropped (trace_id captured via onChunk above) + const aiStream = result.toUIMessageStream({ + sendReasoning: true, + sendSources: true, + sendFinish: false, + onError: (error) => { + const msg = + error instanceof Error ? error.message : String(error); + writer.onError?.(error); + return msg; + }, + }); + + const { failed } = await drainStreamToWriter(aiStream, writer); + + if (failed) { + console.log('Streaming failed, falling back to generateText...'); + const fallbackResult = await fallbackToGenerateText( + { model, messages: modelMessages, headers: requestHeaders }, + writer, + ); + + finalUsage = fallbackResult?.usage; + traceId = fallbackResult?.traceId ?? null; + } + if (titlePromise) { + const generatedTitle = await titlePromise; + if (generatedTitle) { + writer.write({ type: 'data-title', data: generatedTitle }); + } + } + + // Write traceId so the client knows whether feedback is supported. + writer.write({ type: 'data-traceId', data: traceId }); + }, + onFinish: async ({ responseMessage }) => { + // Store in-memory for ephemeral mode (also useful when DB is available) + storeMessageMeta(responseMessage.id, id, traceId); + + try { + await saveMessages({ + messages: [ + { + id: responseMessage.id, + role: responseMessage.role, + parts: responseMessage.parts, + createdAt: new Date(), + attachments: [], + chatId: id, + traceId, // Store trace ID for feedback + }, + ], + }); + } catch (err) { + console.error('[onFinish] Failed to save assistant message:', err); + } + + if (finalUsage) { + try { + await updateChatLastContextById({ + chatId: id, + context: toV3Usage(finalUsage), + }); + } catch (err) { + console.warn('Unable to persist last usage for chat', id, err); + } + } + + streamCache.clearActiveStream(id); + }, + }); + + pipeUIMessageStreamToResponse({ + stream, + response: res, + consumeSseStream({ stream }) { + streamCache.storeStream({ + streamId, + chatId: id, + stream, + }); + }, + }); + } catch (error) { + console.error('[Chat] Caught error in chat API:', { + errorType: error?.constructor?.name, + message: error instanceof Error ? error.message : String(error), + stack: error instanceof Error ? error.stack : undefined, + error, + }); + + if (error instanceof ChatSDKError) { + const response = error.toResponse(); + return res.status(response.status).json(response.json); + } + + const chatError = new ChatSDKError('offline:chat'); + const response = chatError.toResponse(); + return res.status(response.status).json(response.json); + } +}); + +/** + * DELETE /api/chat?id=:id - Delete a chat + */ +chatRouter.delete( + '/:id', + [requireAuth, requireChatAccess], + async (req: Request, res: Response) => { + const id = getIdFromRequest(req); + if (!id) return; + + const deletedChat = await deleteChatById({ id }); + return res.status(200).json(deletedChat); + }, +); + +/** + * GET /api/chat/:id + */ + +chatRouter.get( + '/:id', + [requireAuth, requireChatAccess], + async (req: Request, res: Response) => { + const id = getIdFromRequest(req); + if (!id) return; + + const { chat } = await checkChatAccess(id, req.session?.user.id); + + return res.status(200).json(chat); + }, +); + +/** + * GET /api/chat/:id/stream - Resume a stream + */ +chatRouter.get( + '/:id/stream', + [requireAuth], + async (req: Request, res: Response) => { + const chatId = getIdFromRequest(req); + if (!chatId) return; + const cursor = req.headers['x-resume-stream-cursor'] as string; + + console.log(`[Stream Resume] Cursor: ${cursor}`); + + console.log(`[Stream Resume] GET request for chat ${chatId}`); + + // Check if there's an active stream for this chat first + const streamId = streamCache.getActiveStreamId(chatId); + + if (!streamId) { + console.log(`[Stream Resume] No active stream for chat ${chatId}`); + const streamError = new ChatSDKError('empty:stream'); + const response = streamError.toResponse(); + return res.status(response.status).json(response.json); + } + + const { allowed, reason } = await checkChatAccess( + chatId, + req.session?.user.id, + ); + + // If chat doesn't exist in DB, it's a temporary chat from the homepage - allow it + if (reason === 'not_found') { + console.log( + `[Stream Resume] Resuming stream for temporary chat ${chatId} (not yet in DB)`, + ); + } else if (!allowed) { + console.log( + `[Stream Resume] User ${req.session?.user.id} does not have access to chat ${chatId} (reason: ${reason})`, + ); + const streamError = new ChatSDKError('forbidden:chat', reason); + const response = streamError.toResponse(); + return res.status(response.status).json(response.json); + } + + // Get all cached chunks for this stream + const stream = streamCache.getStream(streamId, { + cursor: cursor ? Number.parseInt(cursor) : undefined, + }); + + if (!stream) { + console.log(`[Stream Resume] No stream found for ${streamId}`); + const streamError = new ChatSDKError('empty:stream'); + const response = streamError.toResponse(); + return res.status(response.status).json(response.json); + } + + console.log(`[Stream Resume] Resuming stream ${streamId}`); + + // Set headers for SSE + res.setHeader('Content-Type', 'text/event-stream'); + res.setHeader('Cache-Control', 'no-cache'); + res.setHeader('Connection', 'keep-alive'); + + // Pipe the cached stream directly to the response + stream.pipe(res); + + // Handle stream errors + stream.on('error', (error) => { + console.error('[Stream Resume] Stream error:', error); + if (!res.headersSent) { + res.status(500).end(); + } + }); + }, +); + +/** + * POST /api/chat/title - Generate title from message + */ +chatRouter.post('/title', requireAuth, async (req: Request, res: Response) => { + try { + const { message } = req.body; + const title = await generateTitleFromUserMessage({ message }); + res.json({ title }); + } catch (error) { + console.error('Error generating title:', error); + res.status(500).json({ error: 'Failed to generate title' }); + } +}); + +/** + * PATCH /api/chat/:id/visibility - Update chat visibility + */ +chatRouter.patch( + '/:id/visibility', + [requireAuth, requireChatAccess], + async (req: Request, res: Response) => { + try { + const id = getIdFromRequest(req); + if (!id) return; + const { visibility } = req.body; + + if (!visibility || !['public', 'private'].includes(visibility)) { + return res.status(400).json({ error: 'Invalid visibility type' }); + } + + await updateChatVisiblityById({ chatId: id, visibility }); + res.json({ success: true }); + } catch (error) { + console.error('Error updating visibility:', error); + res.status(500).json({ error: 'Failed to update visibility' }); + } + }, +); + +// Helper function to generate title from user message +async function generateTitleFromUserMessage({ + message, + maxMessageLength = 256, +}: { + message: ChatMessage; + maxMessageLength?: number; +}) { + const model = await myProvider.languageModel('title-model'); + + // Truncate each text part to the maxMessageLength + const truncatedMessage = { + ...message, + parts: message.parts.map((part) => + part.type === 'text' + ? { ...part, text: part.text.slice(0, maxMessageLength) } + : part, + ), + }; + + const { text: title } = await generateText({ + model, + system: `\n + - you will generate a short title based on the first message a user begins a conversation with + - ensure it is not more than 80 characters long + - the title should be a summary of the user's message + - do not use quotes or colons. do not include other expository content ("I'll help...")`, + prompt: JSON.stringify(truncatedMessage), + }); + + return title; +} + +function truncatePreserveWords(input: string, maxLength: number): string { + if (maxLength <= 0) return ''; + if (input.length <= maxLength) return input; + + // Take the raw slice first + const slice = input.slice(0, maxLength); + + // Find the last whitespace within the slice + const lastSpaceIndex = slice.lastIndexOf(' '); + + // If no whitespace found, we must break mid-word + if (lastSpaceIndex === -1) { + return slice; + } + + // If the whitespace is too close to the start (e.g., leading space), + // fallback to mid-word break to avoid returning an empty string + if (lastSpaceIndex === 0) { + return slice; + } + + return slice.slice(0, lastSpaceIndex); +} + diff --git a/agent/patches/e2e-chatbot-app-next/server/src/routes/config.ts b/agent/patches/e2e-chatbot-app-next/server/src/routes/config.ts new file mode 100644 index 0000000..e28f194 --- /dev/null +++ b/agent/patches/e2e-chatbot-app-next/server/src/routes/config.ts @@ -0,0 +1,103 @@ +import { + Router, + type Request, + type Response, + type Router as RouterType, +} from 'express'; +import { isDatabaseAvailable } from '@chat-template/db'; +import { getEndpointOboInfo } from '@chat-template/ai-sdk-providers'; + +export const configRouter: RouterType = Router(); + +/** + * Extract OAuth scopes from a JWT token (without verification). + * Databricks tokens use 'scope' (space-separated string) or 'scp' (array). + */ +function getScopesFromToken(token: string): string[] { + try { + const parts = token.split('.'); + if (parts.length !== 3) return []; + const payload = JSON.parse(Buffer.from(parts[1], 'base64url').toString('utf-8')); + if (typeof payload.scope === 'string') return payload.scope.split(' '); + if (Array.isArray(payload.scp)) return payload.scp as string[]; + return []; + } catch { + return []; + } +} + +function workspaceHostFromEnv(): string { + const host = process.env.DATABRICKS_HOST?.trim(); + if (host) { + return host.replace(/\/$/, ''); + } + const mcpUrl = process.env.GENIE_MCP_URL?.trim(); + if (!mcpUrl) { + return ''; + } + try { + return new URL(mcpUrl).origin; + } catch { + return ''; + } +} + +function genieOneUrlFromEnv(workspaceHost: string, workspaceId?: string): string { + const explicit = process.env.GENIE_ONE_URL?.trim(); + if (explicit) { + return explicit; + } + if (!workspaceHost) { + return ''; + } + if (workspaceId) { + return `${workspaceHost}/one?o=${workspaceId}`; + } + return `${workspaceHost}/one`; +} + +function genieConfigFromEnv() { + const workspaceHost = workspaceHostFromEnv(); + const workspaceId = process.env.DATABRICKS_WORKSPACE_ID?.trim(); + const genieOneUrl = genieOneUrlFromEnv(workspaceHost, workspaceId); + if (!genieOneUrl && !workspaceHost) { + return null; + } + return { + workspaceHost, + workspaceId: workspaceId || undefined, + genieOneUrl: genieOneUrl || undefined, + }; +} + +/** + * GET /api/config - Get application configuration + * Returns feature flags and OBO status based on environment configuration. + * If the user's OBO token is present, decodes it to check which required + * scopes are missing — the banner only shows missing scopes. + */ +configRouter.get('/', async (req: Request, res: Response) => { + const oboInfo = await getEndpointOboInfo(); + + let missingScopes = oboInfo.endpointRequiredScopes; + + const userToken = req.headers['x-forwarded-access-token'] as string | undefined; + if (userToken && oboInfo.isEndpointOboEnabled) { + const tokenScopes = getScopesFromToken(userToken); + missingScopes = oboInfo.endpointRequiredScopes.filter(required => { + const parent = required.split('.')[0]; + return !tokenScopes.some(ts => ts === required || ts === parent); + }); + } + + res.json({ + features: { + chatHistory: isDatabaseAvailable(), + feedback: !!process.env.MLFLOW_EXPERIMENT_ID, + }, + obo: { + missingScopes, + }, + genie: genieConfigFromEnv(), + }); +}); diff --git a/agent/scripts/start_app.py b/agent/scripts/start_app.py new file mode 100644 index 0000000..6717c21 --- /dev/null +++ b/agent/scripts/start_app.py @@ -0,0 +1,367 @@ +#!/usr/bin/env python3 +""" +Start script for running frontend and backend processes concurrently. + +Requirements: +1. Not reporting ready until BOTH frontend and backend processes are ready +2. Exiting as soon as EITHER process fails +3. Printing error logs if either process fails + +Usage: + start-app [OPTIONS] + +All options are passed through to the backend server (start-server). +See 'uv run start-server --help' for available options. +""" + +import argparse +import os +import re +import shutil +import socket +import subprocess +import sys +import threading +import time +from pathlib import Path + +from dotenv import load_dotenv + +# Readiness patterns +BACKEND_READY = [r"Uvicorn running on", r"Application startup complete", r"Started server process"] +FRONTEND_READY = [r"Server is running on http://localhost"] + + +def backend_command(backend_args: list[str] | None = None) -> list[str]: + """Console script on Apps (pip install); uv locally.""" + if shutil.which("start-server"): + cmd = ["start-server"] + else: + cmd = ["uv", "run", "start-server"] + if backend_args: + cmd.extend(backend_args) + return cmd + + +def check_port_available(port: int) -> bool: + """Check if a port is available (nothing is actively listening on it).""" + try: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.settimeout(1) + s.connect(("localhost", port)) + return False # Something is listening + except (ConnectionRefusedError, OSError): + return True # Nothing listening = available + + +class ProcessManager: + def __init__(self, port=8000, no_ui=False): + self.backend_process = None + self.frontend_process = None + self.backend_ready = False + self.frontend_ready = False + self.failed = threading.Event() + self.backend_log = None + self.frontend_log = None + self.port = port + self.no_ui = no_ui + + def check_ports(self): + """Check that required ports are available before starting processes.""" + backend_port = self.port + + errors = [] + if not check_port_available(backend_port): + errors.append( + f"Port {backend_port} (backend) is already in use.\n" + f" To free it: lsof -ti :{backend_port} | xargs kill -9" + ) + + if not self.no_ui: + frontend_port = int(os.environ.get("CHAT_APP_PORT", os.environ.get("PORT", "3000"))) + + if backend_port == frontend_port: + print( + f"ERROR: Backend and frontend are both configured to use port {backend_port}." + ) + print(" Set CHAT_APP_PORT in .env to a different port (e.g., CHAT_APP_PORT=3000).") + sys.exit(1) + + if not check_port_available(frontend_port): + port_source = ( + "CHAT_APP_PORT" + if os.environ.get("CHAT_APP_PORT") + else "PORT" + if os.environ.get("PORT") + else "default" + ) + errors.append( + f"Port {frontend_port} (frontend, source: {port_source}) is already in use.\n" + f" To free it: lsof -ti :{frontend_port} | xargs kill -9\n" + f" Or set a different port: CHAT_APP_PORT= in .env" + ) + + if errors: + print("ERROR: Port(s) already in use:\n") + for error in errors: + print(f" {error}\n") + sys.exit(1) + + def monitor_process(self, process, name, log_file, patterns): + is_ready = False + try: + for line in iter(process.stdout.readline, ""): + if not line: + break + + line = line.rstrip() + log_file.write(line + "\n") + print(f"[{name}] {line}") + + # Check readiness + if not is_ready and any(re.search(p, line, re.IGNORECASE) for p in patterns): + is_ready = True + if name == "backend": + self.backend_ready = True + else: + self.frontend_ready = True + print(f"✓ {name.capitalize()} is ready!") + + if self.no_ui and self.backend_ready: + print("\n" + "=" * 50) + print("✓ Backend is ready! (running without UI)") + print(f"✓ API available at http://localhost:{self.port}") + print("=" * 50 + "\n") + elif self.backend_ready and self.frontend_ready: + print("\n" + "=" * 50) + print("✓ Both frontend and backend are ready!") + print(f"✓ Open the frontend at http://localhost:{self.port}") + print("=" * 50 + "\n") + + process.wait() + if process.returncode != 0: + self.failed.set() + + except Exception as e: + print(f"Error monitoring {name}: {e}") + self.failed.set() + + def apply_chat_ui_patches(self): + """Overlay Firefly-specific chat UI fixes onto the cloned template.""" + patch_root = Path(__file__).resolve().parent.parent / "patches" / "e2e-chatbot-app-next" + target_root = Path("e2e-chatbot-app-next") + if not patch_root.is_dir() or not target_root.is_dir(): + return + for src in patch_root.rglob("*"): + if not src.is_file(): + continue + rel = src.relative_to(patch_root) + dest = target_root / rel + dest.parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(src, dest) + print("Applied chat UI patches from patches/e2e-chatbot-app-next") + + def clone_frontend_if_needed(self): + if Path("e2e-chatbot-app-next").exists(): + self.apply_chat_ui_patches() + return True + + print("Cloning e2e-chatbot-app-next...") + for url in [ + "https://github.com/databricks/app-templates.git", + "git@github.com:databricks/app-templates.git", + ]: + try: + subprocess.run( + ["git", "clone", "--filter=blob:none", "--sparse", url, "temp-app-templates"], + check=True, + capture_output=True, + ) + break + except subprocess.CalledProcessError: + continue + else: + print("ERROR: Failed to clone repository.") + print( + "Manually download from: https://download-directory.github.io/?url=https://github.com/databricks/app-templates/tree/main/e2e-chatbot-app-next" + ) + return False + + subprocess.run( + ["git", "sparse-checkout", "set", "e2e-chatbot-app-next"], + cwd="temp-app-templates", + check=True, + ) + Path("temp-app-templates/e2e-chatbot-app-next").rename("e2e-chatbot-app-next") + shutil.rmtree("temp-app-templates", ignore_errors=True) + self.apply_chat_ui_patches() + return True + + def start_process(self, cmd, name, log_file, patterns, cwd=None, env=None): + print(f"Starting {name}...") + process = subprocess.Popen( + cmd, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + bufsize=1, + cwd=cwd, + env=env, + ) + + thread = threading.Thread( + target=self.monitor_process, args=(process, name, log_file, patterns), daemon=True + ) + thread.start() + return process + + def print_logs(self, log_path): + print(f"\nLast 50 lines of {log_path}:") + print("-" * 40) + try: + lines = Path(log_path).read_text().splitlines() + print("\n".join(lines[-50:])) + except FileNotFoundError: + print(f"(no {log_path} found)") + print("-" * 40) + + def cleanup(self): + print("\n" + "=" * 42) + print("Shutting down..." if self.no_ui else "Shutting down both processes...") + print("=" * 42) + + for proc in [self.backend_process, self.frontend_process]: + if proc: + try: + proc.terminate() + proc.wait(timeout=5) + except (subprocess.TimeoutExpired, Exception): + proc.kill() + + if self.backend_log: + self.backend_log.close() + if self.frontend_log: + self.frontend_log.close() + + def run(self, backend_args=None): + load_dotenv(dotenv_path=Path(__file__).parent.parent / ".env", override=True) + if not os.environ.get("DATABRICKS_APP_NAME"): + self.check_ports() + + if not self.no_ui: + if not self.clone_frontend_if_needed(): + print("WARNING: Failed to clone frontend. Continuing with backend only.") + self.no_ui = True + else: + # Set API_PROXY environment variable for frontend to connect to backend + os.environ["API_PROXY"] = f"http://localhost:{self.port}/invocations" + + # Open log files + self.backend_log = open("backend.log", "w", buffering=1) + if not self.no_ui: + self.frontend_log = open("frontend.log", "w", buffering=1) + + try: + backend_cmd = backend_command(backend_args) + backend_env = os.environ.copy() + if self.no_ui: + backend_env["ENABLE_CHAT_PROXY"] = "false" + + # Start backend + self.backend_process = self.start_process( + backend_cmd, "backend", self.backend_log, BACKEND_READY, env=backend_env + ) + + if not self.no_ui: + # Setup and start frontend + frontend_dir = Path("e2e-chatbot-app-next") + for cmd, desc in [("npm install", "install"), ("npm run build", "build")]: + print(f"Running npm {desc}...") + result = subprocess.run( + cmd.split(), cwd=frontend_dir, capture_output=True, text=True + ) + if result.returncode != 0: + print(f"npm {desc} failed: {result.stderr}") + return 1 + + self.frontend_process = self.start_process( + ["npm", "run", "start"], + "frontend", + self.frontend_log, + FRONTEND_READY, + cwd=frontend_dir, + ) + + print( + f"\nMonitoring processes (Backend PID: {self.backend_process.pid}, Frontend PID: {self.frontend_process.pid})\n" + ) + else: + print(f"\nMonitoring backend process (PID: {self.backend_process.pid})\n") + + # Wait for failure + while not self.failed.is_set(): + time.sleep(0.1) + if self.backend_process.poll() is not None: + self.failed.set() + break + if ( + not self.no_ui + and self.frontend_process + and self.frontend_process.poll() is not None + ): + self.failed.set() + break + + # Determine which failed + if self.no_ui or self.backend_process.poll() is not None: + failed_name = "backend" + failed_proc = self.backend_process + else: + failed_name = "frontend" + failed_proc = self.frontend_process + exit_code = failed_proc.returncode if failed_proc else 1 + + print( + f"\n{'=' * 42}\nERROR: {failed_name} process exited with code {exit_code}\n{'=' * 42}" + ) + self.print_logs("backend.log") + if not self.no_ui: + self.print_logs("frontend.log") + return exit_code + + except KeyboardInterrupt: + print("\nInterrupted") + return 0 + + finally: + self.cleanup() + + +def main(): + parser = argparse.ArgumentParser( + description="Start agent frontend and backend", + usage="%(prog)s [OPTIONS]\n\nAll options are passed through to start-server. " + "Use 'uv run start-server --help' for available options.", + ) + parser.add_argument( + "--no-ui", + action="store_true", + help="Run backend only, skip frontend UI", + ) + args, backend_args = parser.parse_known_args() + + # Extract port from backend_args if specified + port = 8000 + for i, arg in enumerate(backend_args): + if arg == "--port" and i + 1 < len(backend_args): + try: + port = int(backend_args[i + 1]) + except ValueError: + pass + break + + sys.exit(ProcessManager(port=port, no_ui=args.no_ui).run(backend_args)) + + +if __name__ == "__main__": + main() diff --git a/scripts/assemble_agent.sh b/scripts/assemble_agent.sh new file mode 100755 index 0000000..a6c24b5 --- /dev/null +++ b/scripts/assemble_agent.sh @@ -0,0 +1,32 @@ +#!/usr/bin/env bash +# Merge vendor/app-templates (submodule) + agent/ overlay -> agent-build/ (deploy source). +# The submodule is a pristine upstream pin; all Firefly deltas live in agent/. +set -euo pipefail + +ROOT="$(cd "$(dirname "$0")/.." && pwd)" +TEMPLATE="$ROOT/vendor/app-templates/agent-openai-agents-sdk" +CHAT="$ROOT/vendor/app-templates/e2e-chatbot-app-next" +OVERLAY="$ROOT/agent" +BUILD="$ROOT/agent-build" + +[[ -d "$TEMPLATE" ]] || { echo "Missing $TEMPLATE — run: git submodule update --init"; exit 1; } +[[ -d "$CHAT" ]] || { echo "Missing $CHAT — check sparse-checkout"; exit 1; } + +rm -rf "$BUILD" +mkdir -p "$BUILD" + +# 1) upstream agent template +cp -R "$TEMPLATE"/. "$BUILD"/ +# 2) pre-vendor the chat UI so start_app.py doesn't clone at runtime +cp -R "$CHAT" "$BUILD"/e2e-chatbot-app-next +# 3) overlay: our agent_server deltas (agent.py, utils.py, start_server.py, genie_tools.py, utils_memory.py) +cp -R "$OVERLAY"/agent_server/. "$BUILD"/agent_server/ +# 4) overlay: bundle config + startup +[[ -f "$OVERLAY/databricks.yml" ]] && cp "$OVERLAY/databricks.yml" "$BUILD"/ +[[ -f "$OVERLAY/scripts/start_app.py" ]] && cp "$OVERLAY/scripts/start_app.py" "$BUILD"/scripts/ +# 5) overlay: chat UI patches (Genie attribution, proxy-friendly tweaks) +if [[ -d "$OVERLAY/patches/e2e-chatbot-app-next" ]]; then + cp -R "$OVERLAY/patches/e2e-chatbot-app-next/." "$BUILD"/e2e-chatbot-app-next/ +fi + +echo "Assembled agent at $BUILD" diff --git a/src/app/api/agent-proxy/[[...path]]/route.ts b/src/app/api/agent-proxy/[[...path]]/route.ts new file mode 100644 index 0000000..0a7aef6 --- /dev/null +++ b/src/app/api/agent-proxy/[[...path]]/route.ts @@ -0,0 +1,185 @@ +import { NextRequest, NextResponse } from "next/server"; +import { headers } from "next/headers"; +import { eq } from "drizzle-orm"; +import { getAuthInstance } from "@/lib/auth-dynamic"; +import { getDatabricksSpnToken } from "@/lib/databricks-spn-authtoken"; +import { db } from "@/db"; +import { organization } from "@/db/schema"; + +// Vercel-native reverse proxy for the managed-memory agent Databricks App. +// +// Mirrors what the stock Go proxy does (inject a bearer, strip forwarded +// headers, relax frame headers) but resolves the token from the *current +// user's* mapped SPN (guest / BYOD supported) via getDatabricksSpnToken, so a +// guest never sees the Databricks OAuth wall. The agent iframe loads +// /api/agent-proxy/ (same origin) and every asset/api request is caught by this +// catch-all and forwarded to the App with the injected token. +// +// The agent chat UI (built with base:"./" + the fetch shim) emits relative +// asset URLs and prefixes /api/ calls with /api/agent-proxy, so both static +// assets and the streaming /api/chat SSE route back through here. + +export const dynamic = "force-dynamic"; +export const runtime = "nodejs"; +// Chat responses stream via SSE; allow a long-running function. +export const maxDuration = 300; + +const AGENT_APP_URL = (process.env.DATABRICKS_AGENT_APP_URL ?? "").replace( + /\/$/, + "", +); + +// Request headers we must not forward upstream. +const STRIPPED_REQUEST_HEADERS = new Set([ + "host", + "cookie", + "authorization", + "connection", + "content-length", + "x-forwarded-for", + "x-forwarded-host", + "x-forwarded-proto", + "x-forwarded-port", + "x-real-ip", +]); + +// Response headers we must not pass back (Node fetch already decoded the body, +// and we set our own framing policy). +const STRIPPED_RESPONSE_HEADERS = new Set([ + "content-encoding", + "content-length", + "transfer-encoding", + "connection", + "content-security-policy", + "x-frame-options", +]); + +async function resolveToken(): Promise< + | { ok: true; token: string } + | { ok: false; status: number; body: unknown } +> { + const auth = await getAuthInstance(); + const session = await auth.api.getSession({ headers: await headers() }); + + const activeOrgId = session?.session?.activeOrganizationId; + if (!session || !activeOrgId) { + return { + ok: false, + status: 401, + body: { error: "No active organization in session" }, + }; + } + + const [org] = await db + .select() + .from(organization) + .where(eq(organization.id, activeOrgId)) + .limit(1); + + if (!org?.workspaceUrl) { + return { + ok: false, + status: 400, + body: { error: "No workspace URL configured for this organization" }, + }; + } + + const workspaceUrl = org.workspaceUrl.replace(/\/$/, ""); + const tokenResult = await getDatabricksSpnToken( + workspaceUrl, + undefined, + session.user.email, + activeOrgId, + ); + + if (!tokenResult.success) { + return { + ok: false, + status: tokenResult.error.status, + body: { + error: tokenResult.error.error, + details: tokenResult.error.details, + }, + }; + } + + return { ok: true, token: tokenResult.data.accessToken }; +} + +async function proxy( + req: NextRequest, + context: { params: Promise<{ path?: string[] }> }, +): Promise { + if (!AGENT_APP_URL) { + return NextResponse.json( + { error: "DATABRICKS_AGENT_APP_URL is not configured" }, + { status: 503 }, + ); + } + + const auth = await resolveToken(); + if (!auth.ok) { + return NextResponse.json(auth.body, { status: auth.status }); + } + + const { path } = await context.params; + const suffix = (path ?? []).join("/"); + const search = req.nextUrl.search; + const targetUrl = `${AGENT_APP_URL}/${suffix}${search}`; + + // Copy inbound headers minus the ones we strip, then inject the bearer. + const outboundHeaders = new Headers(); + req.headers.forEach((value, key) => { + if (!STRIPPED_REQUEST_HEADERS.has(key.toLowerCase())) { + outboundHeaders.set(key, value); + } + }); + outboundHeaders.set("Authorization", `Bearer ${auth.token}`); + + const method = req.method.toUpperCase(); + const hasBody = method !== "GET" && method !== "HEAD"; + + const init: RequestInit & { duplex?: "half" } = { + method, + headers: outboundHeaders, + redirect: "manual", + }; + if (hasBody) { + init.body = req.body; + init.duplex = "half"; + } + + let upstream: Response; + try { + upstream = await fetch(targetUrl, init); + } catch (err) { + return NextResponse.json( + { error: "Upstream agent app request failed", details: String(err) }, + { status: 502 }, + ); + } + + const responseHeaders = new Headers(); + upstream.headers.forEach((value, key) => { + if (!STRIPPED_RESPONSE_HEADERS.has(key.toLowerCase())) { + responseHeaders.set(key, value); + } + }); + // Same-origin iframe: allow the frontend to embed the proxied app. + responseHeaders.set("X-Frame-Options", "SAMEORIGIN"); + responseHeaders.set("Content-Security-Policy", "frame-ancestors 'self'"); + + return new Response(upstream.body, { + status: upstream.status, + statusText: upstream.statusText, + headers: responseHeaders, + }); +} + +export const GET = proxy; +export const POST = proxy; +export const PUT = proxy; +export const PATCH = proxy; +export const DELETE = proxy; +export const HEAD = proxy; +export const OPTIONS = proxy; diff --git a/src/app/sso-spn/[orgId]/layout.tsx b/src/app/sso-spn/[orgId]/layout.tsx index 7501b80..edb1186 100644 --- a/src/app/sso-spn/[orgId]/layout.tsx +++ b/src/app/sso-spn/[orgId]/layout.tsx @@ -10,6 +10,7 @@ import { SsoSpnUserStoreInitializer } from "@/components/sso-spn-user-store-init import { Spinner } from "@/components/ui/spinner"; import { Toaster } from "@/components/ui/sonner"; import { Separator } from "@/components/ui/separator"; +import { AgentPanel } from "@/components/agent/agent-panel"; export default function OrgLayout({ children, @@ -69,6 +70,7 @@ export default function OrgLayout({ {children} + diff --git a/src/components/agent/agent-panel.tsx b/src/components/agent/agent-panel.tsx new file mode 100644 index 0000000..98e6632 --- /dev/null +++ b/src/components/agent/agent-panel.tsx @@ -0,0 +1,152 @@ +"use client"; + +import { useEffect } from "react"; +import { Bot, ChevronLeft, Minus, X } from "lucide-react"; +import { Button } from "@/components/ui/button"; +import { cn } from "@/lib/utils"; +import { useAgentPanelStore } from "@/stores/agent-panel-store"; + +// The agent chat UI (managed-memory Databricks App) is embedded via the +// same-origin Vercel proxy route at /api/agent-proxy/. That route mints the +// current user's SPN token (guest/BYOD supported) and injects it as a bearer, +// so a guest never hits the Databricks OAuth wall. No client-side token/tool id +// is needed — the route resolves the org + token from the session cookie. +const AGENT_ENABLED = + process.env.NEXT_PUBLIC_AGENT_ENABLED?.trim().toLowerCase() === "true"; +// Trailing slash matters: the app is built with base:"./", so relative assets +// resolve under /api/agent-proxy/ instead of the site root. +const AGENT_PROXY_SRC = "/api/agent-proxy/"; + +function agentPanelEnabled() { + return AGENT_ENABLED; +} + +function AgentChatFrame({ orgId: _orgId }: { orgId: string }) { + return ( +
+
+