chatPanelProvider.ts

import * as vscode from "vscode";

import { type CoderApi } from "../../api/coderApi";
import { type Logger } from "../../logging/logger";
import { getNonce } from "../util";

/**
 * Provides a webview that embeds the Coder agent chat UI.
 * Authentication flows through postMessage:
 *
 *   1. The iframe loads /agents/{id}/embed on the Coder server.
 *   2. The embed page detects the user is signed out and sends
 *      { type: "coder:vscode-ready" } to window.parent.
 *   3. Our webview relays this to the extension host.
 *   4. The extension host replies with the session token.
 *   5. The webview forwards { type: "coder:vscode-auth-bootstrap" }
 *      with the token back into the iframe.
 *   6. The embed page calls API.setSessionToken(token), re-fetches
 *      the authenticated user, and renders the chat UI.
 */
export class ChatPanelProvider
	implements vscode.WebviewViewProvider, vscode.Disposable
{
	public static readonly viewType = "coder.chatPanel";

	private view?: vscode.WebviewView;
	private disposables: vscode.Disposable[] = [];
	private chatId: string | undefined;
	private authRetryTimer: ReturnType<typeof setTimeout> | undefined;

	constructor(
		private readonly client: Pick<CoderApi, "getHost" | "getSessionToken">,
		private readonly logger: Logger,
	) {}

	private getTheme(): "light" | "dark" {
		const kind = vscode.window.activeColorTheme.kind;
		return kind === vscode.ColorThemeKind.Light ||
			kind === vscode.ColorThemeKind.HighContrastLight
			? "light"
			: "dark";
	}

	private sendTheme(): void {
		this.view?.webview.postMessage({
			type: "coder:set-theme",
			theme: this.getTheme(),
		});
	}

	/**
	 * Opens the chat panel for the given chat ID.
	 * Called after a deep link reload via the persisted
	 * pendingChatId, or directly for testing.
	 */
	public openChat(chatId: string): void {
		if (this.chatId === chatId && this.view) {
			this.view.show(true);
			return;
		}
		this.chatId = chatId;
		// No-op if unresolved; the focus command triggers resolveWebviewView().
		this.refresh();
		void vscode.commands.executeCommand(`${ChatPanelProvider.viewType}.focus`);
	}

	resolveWebviewView(
		webviewView: vscode.WebviewView,
		_context: vscode.WebviewViewResolveContext,
		_token: vscode.CancellationToken,
	): void {
		// Clean up state from a previous view instance to avoid
		// duplicates if VS Code re-resolves the view.
		this.disposeView();
		this.view = webviewView;
		webviewView.webview.options = { enableScripts: true };
		this.disposables.push(
			webviewView.webview.onDidReceiveMessage((msg: unknown) => {
				this.handleMessage(msg);
			}),
			vscode.window.onDidChangeActiveColorTheme(() => {
				this.sendTheme();
			}),
		);
		this.renderView();
		this.disposables.push(webviewView.onDidDispose(() => this.disposeView()));
	}

	public refresh(): void {
		if (!this.view) {
			return;
		}
		this.renderView();
	}

	private renderView(): void {
		if (!this.view) {
			throw new Error("renderView called before resolveWebviewView");
		}
		const webview = this.view.webview;

		if (!this.chatId) {
			webview.html = this.getNoAgentHtml();
			return;
		}

		const coderUrl = this.client.getHost();
		if (!coderUrl) {
			webview.html = this.getNoAgentHtml();
			return;
		}

		const embedUrl = `${coderUrl}/agents/${this.chatId}/embed?theme=${this.getTheme()}`;
		webview.html = this.getIframeHtml(embedUrl, coderUrl);
	}

	private handleMessage(message: unknown): void {
		if (typeof message !== "object" || message === null) {
			return;
		}
		const msg = message as { type?: string; payload?: { url?: string } };
		switch (msg.type) {
			case "coder:vscode-ready":
				this.sendAuthToken();
				break;
			case "coder:chat-ready":
				this.sendTheme();
				break;
			case "coder:navigate": {
				const url = msg.payload?.url;
				const coderUrl = this.client.getHost();
				if (url && coderUrl) {
					try {
						const resolved = new URL(url, coderUrl);
						const expected = new URL(coderUrl);
						if (resolved.origin === expected.origin) {
							void vscode.env.openExternal(
								vscode.Uri.parse(resolved.toString()),
							);
						}
					} catch {
						this.logger.warn(`Chat: invalid navigate URL: ${url}`);
					}
				}
				break;
			}
			default:
				break;
		}
	}

	/**
	 * Attempt to forward the session token to the chat iframe.
	 * The token may not be available immediately after a reload
	 * (e.g. deployment setup is still in progress), so we retry
	 * with exponential backoff before giving up.
	 */
	private static readonly MAX_AUTH_RETRIES = 5;
	private static readonly AUTH_RETRY_BASE_MS = 500;

	private sendAuthToken(attempt = 0): void {
		clearTimeout(this.authRetryTimer);
		const token = this.client.getSessionToken();
		if (!token) {
			if (attempt < ChatPanelProvider.MAX_AUTH_RETRIES) {
				const delay = ChatPanelProvider.AUTH_RETRY_BASE_MS * 2 ** attempt;
				this.logger.info(
					`Chat: no session token yet, retrying in ${delay}ms ` +
						`(attempt ${attempt + 1}/${ChatPanelProvider.MAX_AUTH_RETRIES})`,
				);
				this.authRetryTimer = setTimeout(
					() => this.sendAuthToken(attempt + 1),
					delay,
				);
				return;
			}
			this.logger.warn(
				"Chat iframe requested auth but no session token available " +
					"after all retries",
			);
			this.view?.webview.postMessage({
				type: "coder:auth-error",
				error: "No session token available. Please sign in and retry.",
			});
			return;
		}
		this.logger.info("Chat: forwarding token to iframe");
		this.view?.webview.postMessage({
			type: "coder:auth-bootstrap-token",
			token,
		});
	}

	private getIframeHtml(embedUrl: string, allowedOrigin: string): string {
		const nonce = getNonce();

		return /* html */ `<!DOCTYPE html>
<html lang="en">
<head>
  <meta charset="UTF-8">
  <meta http-equiv="Content-Security-Policy"
        content="default-src 'none';
                 frame-src ${allowedOrigin};
                 script-src 'nonce-${nonce}';
                 style-src 'unsafe-inline';">
  <meta name="viewport" content="width=device-width, initial-scale=1.0">
  <title>Coder Chat</title>
  <style>
    html, body {
      margin: 0; padding: 0;
      width: 100%; height: 100%;
      overflow: hidden;
      background: var(--vscode-editor-background, #1e1e1e);
    }
    iframe { border: none; width: 100%; height: 100%; }
    #status {
      color: var(--vscode-foreground, #ccc);
      font-family: var(--vscode-font-family, sans-serif);
      font-size: 13px; padding: 16px; text-align: center;
    }
    #retry-btn {
      margin-top: 12px; padding: 6px 16px;
      background: var(--vscode-button-background, #0e639c);
      color: var(--vscode-button-foreground, #fff);
      border: none; border-radius: 2px; cursor: pointer;
      font-family: var(--vscode-font-family, sans-serif);
      font-size: 13px;
    }
    #retry-btn:hover {
      background: var(--vscode-button-hoverBackground, #1177bb);
    }
  </style>
</head>
<body>
  <div id="status">Loading chat…</div>
  <iframe id="chat-frame" src="${embedUrl}" allow="clipboard-write"
          style="display:none;"></iframe>
  <script nonce="${nonce}">
    (function () {
      const vscode = acquireVsCodeApi();
      const iframe = document.getElementById('chat-frame');
      const status = document.getElementById('status');

      iframe.addEventListener('load', () => {
        iframe.style.display = 'block';
        status.style.display = 'none';
      });

      window.addEventListener('message', (event) => {
        const data = event.data;
        if (!data || typeof data !== 'object') return;

        if (event.source === iframe.contentWindow) {
          if (data.type === 'coder:vscode-ready') {
            status.textContent = 'Authenticating…';
            vscode.postMessage({ type: 'coder:vscode-ready' });
          }
          if (data.type === 'coder:chat-ready') {
            vscode.postMessage({ type: 'coder:chat-ready' });
          }
          if (data.type === 'coder:navigate') {
            vscode.postMessage(data);
          }
          return;
        }

        if (data.type === 'coder:auth-bootstrap-token') {
          status.textContent = 'Signing in…';
          iframe.contentWindow.postMessage({
            type: 'coder:vscode-auth-bootstrap',
            payload: { token: data.token },
          }, '${allowedOrigin}');
        }

        if (data.type === 'coder:set-theme') {
          iframe.contentWindow.postMessage({
            type: 'coder:set-theme',
            payload: { theme: data.theme },
          }, '${allowedOrigin}');
        }

        if (data.type === 'coder:auth-error') {
          status.textContent = '';
          status.appendChild(document.createTextNode(data.error || 'Authentication failed.'));
          const btn = document.createElement('button');
          btn.id = 'retry-btn';
          btn.textContent = 'Retry';
          btn.addEventListener('click', () => {
            status.textContent = 'Authenticating…';
            vscode.postMessage({ type: 'coder:vscode-ready' });
          });
          status.appendChild(document.createElement('br'));
          status.appendChild(btn);
          status.style.display = 'block';
          iframe.style.display = 'none';
        }
      });
    })();
  </script>
</body>
</html>`;
	}

	private getNoAgentHtml(): string {
		return /* html */ `<!DOCTYPE html>
<html lang="en"><head><meta charset="UTF-8">
<style>body{display:flex;align-items:center;justify-content:center;
height:100vh;margin:0;padding:16px;box-sizing:border-box;
font-family:var(--vscode-font-family);color:var(--vscode-foreground);
text-align:center;}</style></head>
<body><p>No active chat session. Open a chat from the Agents tab on your Coder deployment.</p></body></html>`;
	}

	private disposeView(): void {
		clearTimeout(this.authRetryTimer);
		for (const d of this.disposables) {
			d.dispose();
		}
		this.disposables = [];
	}

	dispose(): void {
		this.disposeView();
	}
}