storybook/zed-storybook/src/lib.rs

use zed_extension_api::{self as zed, Result};
use std::{env, fs};

struct StorybookExtension {
    cached_binary_path: Option<String>,
}

impl StorybookExtension {
    fn language_server_binary_path(
        &mut self,
        language_server_id: &zed::LanguageServerId,
        worktree: &zed::Worktree,
    ) -> Result<String> {
        // Return cached path if we have it
        if let Some(path) = &self.cached_binary_path {
            if fs::metadata(path).is_ok() {
                return Ok(path.clone());
            }
        }

        zed::set_language_server_installation_status(
            language_server_id,
            &zed::LanguageServerInstallationStatus::CheckingForUpdate,
        );

        // 1. Check system PATH (works for local dev when installed via cargo)
        if let Some(path) = worktree.which("storybook-lsp") {
            self.cached_binary_path = Some(path.clone());

            zed::set_language_server_installation_status(
                language_server_id,
                &zed::LanguageServerInstallationStatus::None,
            );

            return Ok(path);
        }

        // 2. Check extension's bin directory (for bundled binaries)
        if let Ok(current_dir) = env::current_dir() {
            let bundled_binary = current_dir.join("bin/storybook-lsp");
            if bundled_binary.exists() {
                let path_str = bundled_binary.to_string_lossy().to_string();
                self.cached_binary_path = Some(path_str.clone());

                zed::set_language_server_installation_status(
                    language_server_id,
                    &zed::LanguageServerInstallationStatus::None,
                );

                return Ok(path_str);
            }
        }

        // 3. TODO: Download from GitHub releases

        zed::set_language_server_installation_status(
            language_server_id,
            &zed::LanguageServerInstallationStatus::Failed(
                "storybook-lsp not found. Install with: cargo install --path storybook --bin storybook-lsp".to_string()
            ),
        );

        Err("storybook-lsp binary not found".into())
    }
}

impl zed::Extension for StorybookExtension {
    fn new() -> Self {
        Self {
            cached_binary_path: None,
        }
    }

    fn language_server_command(
        &mut self,
        language_server_id: &zed::LanguageServerId,
        worktree: &zed::Worktree,
    ) -> Result<zed::Command> {
        let binary_path = self.language_server_binary_path(language_server_id, worktree)?;

        Ok(zed::Command {
            command: binary_path,
            args: vec![],
            env: Default::default(),
        })
    }
}

zed::register_extension!(StorybookExtension);