fix: support cursor auth

jiashengguo · jiashengguo · commit 188ff5f3912e · 2025-09-11T14:15:18.000+08:00
diff --git a/packages/schema/package.json b/packages/schema/package.json
@@ -27,21 +27,12 @@
         "linkDirectory": true
     },
     "engines": {
-        "vscode": "^1.102.0"
+        "vscode": "^1.99.3"
     },
     "categories": [
         "Programming Languages"
     ],
     "contributes": {
-        "languageModelTools": [
-            {
-                "name": "zmodel_mermaid_generator",
-                "displayName": "ZModel Mermaid Generator",
-                "modelDescription": "Generate Mermaid charts from ZModel schema files. This tool analyzes the current ZModel file and creates comprehensive entity-relationship diagrams showing all models and their relationships.",
-                "canBeReferencedInPrompt": true,
-                "toolReferenceName": "zmodel_mermaid_generator"
-            }
-        ],
         "languages": [
             {
                 "id": "zmodel",
@@ -102,7 +93,7 @@
         "commands": [
             {
                 "command": "zenstack.preview-zmodel",
-                "title": "ZenStack: Preview ZModel",
+                "title": "ZenStack: Preview ZModel Documentation",
                 "icon": "$(preview)"
             },
             {
@@ -183,7 +174,7 @@
         "@types/strip-color": "^0.1.0",
         "@types/tmp": "^0.2.3",
         "@types/uuid": "^8.3.4",
-        "@types/vscode": "^1.102.0",
+        "@types/vscode": "^1.99.3",
         "@vscode/vsce": "^3.5.0",
         "@zenstackhq/runtime": "workspace:*",
         "dotenv": "^16.0.3",
diff --git a/packages/schema/src/documentation-cache.ts b/packages/schema/src/documentation-cache.ts
@@ -53,10 +53,10 @@ export class DocumentationCache implements vscode.Disposable {
     /**
      * Generate a cache key from request body with normalized content
      */
-    private generateCacheKey(requestBody: { models: string[] }): string {
+    private generateCacheKey(models: string[]): string {
         // Remove ALL whitespace characters from each model string for cache key generation
         // This ensures identical content with different formatting uses the same cache
-        const normalizedModels = requestBody.models.map((model) => model.replace(/\s/g, '')).sort();
+        const normalizedModels = models.map((model) => model.replace(/\s/g, '')).sort();
         const hash = createHash('sha512')
             .update(JSON.stringify({ models: normalizedModels }))
             .digest('hex');
@@ -73,8 +73,8 @@ export class DocumentationCache implements vscode.Disposable {
     /**
      * Get cached response if available and valid
      */
-    async getCachedResponse(requestBody: { models: string[] }): Promise<string | null> {
-        const cacheKey = this.generateCacheKey(requestBody);
+    async getCachedResponse(models: string[]): Promise<string | null> {
+        const cacheKey = this.generateCacheKey(models);
         const entry = this.extensionContext.globalState.get<CacheEntry>(cacheKey);
 
         if (entry && this.isCacheValid(entry)) {
@@ -93,8 +93,8 @@ export class DocumentationCache implements vscode.Disposable {
     /**
      * Cache a response for future use
      */
-    async setCachedResponse(requestBody: { models: string[] }, data: string): Promise<void> {
-        const cacheKey = this.generateCacheKey(requestBody);
+    async setCachedResponse(models: string[], data: string): Promise<void> {
+        const cacheKey = this.generateCacheKey(models);
         const cacheEntry: CacheEntry = {
             data,
             timestamp: Date.now(),
diff --git a/packages/schema/src/zenstack-auth-provider.ts b/packages/schema/src/zenstack-auth-provider.ts
@@ -10,6 +10,7 @@ interface JWTClaims {
 
 export const AUTH_PROVIDER_ID = 'ZenStack';
 export const AUTH_URL = 'https://accounts.zenstack.dev';
+export const API_URL = 'https://api.zenstack.dev';
 
 export class ZenStackAuthenticationProvider implements vscode.AuthenticationProvider, vscode.Disposable {
     private _onDidChangeSessions =
@@ -20,7 +21,6 @@ export class ZenStackAuthenticationProvider implements vscode.AuthenticationProv
     private _context: vscode.ExtensionContext;
     private _disposable: vscode.Disposable;
     private pendingAuth?: {
-        state: string;
         resolve: (session: vscode.AuthenticationSession) => void;
         reject: (error: Error) => void;
         scopes: readonly string[];
@@ -95,16 +95,22 @@ export class ZenStackAuthenticationProvider implements vscode.AuthenticationProv
                         reject(new Error('User Cancelled'));
                     });
 
-                    // Generate a unique state parameter for security
-                    const state = this.generateState();
-                    // Construct the ZenStack sign-in URL for implicit flow (returns access_token directly)
-                    const signInUrl = new URL('/sign-in', AUTH_URL);
+                    const isCursor = vscode.env.appName == 'Cursor';
 
+                    let signInUrl = vscode.Uri.parse(new URL('/sign-in', AUTH_URL).toString());
+
+                    if (isCursor) {
+                        signInUrl = signInUrl.with({
+                            query: `redirect_url=${API_URL}/oauth/oauth_callback?vscodeapp=cursor`,
+                        });
+                    }
+
+                    console.log('ZenStack sign-in URL:', signInUrl.toString());
                     // Store the state and resolve function for later use
-                    this.pendingAuth = { state, resolve, reject, scopes };
+                    this.pendingAuth = { resolve, reject, scopes };
 
                     // Open the ZenStack sign-in page in the user's default browser
-                    vscode.env.openExternal(vscode.Uri.parse(signInUrl.toString())).then(
+                    vscode.env.openExternal(signInUrl).then(
                         () => {
                             console.log('Opened ZenStack sign-in page in browser');
                             progress.report({ message: 'Waiting for return from browser...' });
diff --git a/packages/schema/src/zmodel-preview.ts b/packages/schema/src/zmodel-preview.ts
@@ -1,10 +1,12 @@
 import * as vscode from 'vscode';
 import * as path from 'path';
 import * as os from 'os';
+import { z } from 'zod';
 import { LanguageClient } from 'vscode-languageclient/node';
 import { URI } from 'vscode-uri';
 import { DocumentationCache } from './documentation-cache';
 import { requireAuth } from './extension';
+import { API_URL } from './zenstack-auth-provider';
 
 /**
  * ZModelPreview class handles ZModel file preview functionality
@@ -13,6 +15,25 @@ export class ZModelPreview implements vscode.Disposable {
     private documentationCache: DocumentationCache;
     private languageClient: LanguageClient;
 
+    // Schema for validating the request body
+    private static DocRequestSchema = z.object({
+        models: z.array(
+            z.object({
+                path: z.string().optional(),
+                content: z.string(),
+            })
+        ),
+        environments: z
+            .object({
+                vscodeAppName: z.string(),
+                vscodeVersion: z.string(),
+                vscodeAppHost: z.string(),
+                osRelease: z.string(),
+                osType: z.string(),
+            })
+            .optional(),
+    });
+
     constructor(context: vscode.ExtensionContext, client: LanguageClient, cache: DocumentationCache) {
         this.documentationCache = cache;
         this.languageClient = client;
@@ -141,50 +162,53 @@ export class ZModelPreview implements vscode.Disposable {
             const importedURIs = astInfo?.importedURIs;
 
             // get vscode document from importedURIs
-            const importedTexts = await Promise.all(
+            const importedModels = await Promise.all(
                 importedURIs.map(async (uri) => {
                     try {
                         const fileUri = vscode.Uri.file(uri.path);
                         const fileContent = await vscode.workspace.fs.readFile(fileUri);
-                        return Buffer.from(fileContent).toString('utf8');
+                        const filePath = fileUri.path;
+                        return { content: Buffer.from(fileContent).toString('utf8').trim(), path: filePath };
                     } catch (error) {
-                        console.warn(`Could not read file for URI ${uri}:`, error);
-                        return null;
+                        throw new Error(
+                            `Failed to read imported ZModel file at ${uri.path}: ${
+                                error instanceof Error ? error.message : String(error)
+                            }`
+                        );
                     }
                 })
             );
 
-            const zmodelContent = [document.getText(), ...importedTexts.filter((text) => text !== null)];
-
-            // Trim whitespace from each model string
-            const trimmedZmodelContent = zmodelContent.map((content) => content.trim());
+            const allModels = [{ content: document.getText().trim(), path: document.uri.path }, ...importedModels];
 
             const session = await requireAuth();
             if (!session) {
                 throw new Error('Authentication required to generate documentation');
             }
 
             // Prepare request body
-            const requestBody = {
-                models: trimmedZmodelContent,
+            const requestBody: z.infer<typeof ZModelPreview.DocRequestSchema> = {
+                models: allModels,
                 environments: {
-                    editorName: vscode.env.appName,
+                    vscodeAppName: vscode.env.appName,
                     vscodeVersion: vscode.version,
-                    appHost: vscode.env.appHost,
+                    vscodeAppHost: vscode.env.appHost,
                     osRelease: os.release(),
                     osType: os.type(),
                 },
             };
 
+            const allModelsContent = allModels.map((m) => m.content);
+
             // Check cache first
-            const cachedResponse = await this.documentationCache.getCachedResponse(requestBody);
+            const cachedResponse = await this.documentationCache.getCachedResponse(allModelsContent);
             if (cachedResponse) {
                 return cachedResponse;
             }
 
             // record the time spent
             const startTime = Date.now();
-            const apiResponse = await fetch('https://api.zenstack.dev/api/doc', {
+            const apiResponse = await fetch(`${API_URL}/api/doc`, {
                 method: 'POST',
                 headers: {
                     'Content-Type': 'application/json',
@@ -202,7 +226,7 @@ export class ZModelPreview implements vscode.Disposable {
             const responseText = await apiResponse.text();
 
             // Cache the response
-            await this.documentationCache.setCachedResponse(requestBody, responseText);
+            await this.documentationCache.setCachedResponse(allModelsContent, responseText);
 
             return responseText;
         } catch (error) {
diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml
