cocalc-api: improve api key type discovery

Author: haraldschilly

src/packages/conat/hub/api/system.ts

@@ -32,8 +32,8 @@ export interface System {
   getCustomize: (fields?: string[]) => Promise<Customize>;
   // ping server and get back the current time
   ping: () => { now: number };
-  // test API key and return scope information (account_id or project_id) and server time
-  test: () => Promise<{ account_id?: string; project_id?: string; server_time: number }>;
+  // test API key and return scope information (account_id) and server time
+  test: () => Promise<{ account_id: string; server_time: number }>;
   // terminate a service:
   //   - only admin can do this.
   //   - useful for development

src/packages/next/pages/api/conat/project.ts

@@ -58,11 +58,6 @@ export default async function handle(req, res) {
       args,
       timeout,
     });
-    // For project-scoped API keys, include the project_id in the response
-    // so the client can discover it
-    if (project_id0 && !resp.project_id) {
-      resp.project_id = project_id0;
-    }
     res.json(resp);
   } catch (err) {
     res.json({ error: err.message });

src/packages/server/api/project-bridge.ts

@@ -4,6 +4,7 @@ import { projectSubject } from "@cocalc/conat/names";
 import { conat } from "@cocalc/backend/conat";
 import { type Client as ConatClient } from "@cocalc/conat/core/client";
 import { getProject } from "@cocalc/server/projects/control";
+
 const DEFAULT_TIMEOUT = 15000;
 
 let client: ConatClient | null = null;
@@ -58,10 +59,17 @@ async function callProject({
       await project.start();
     }
 
-    // For system.test(), inject project_id into args[0] if not already present
+    // For discovery-style calls, inject identifiers so the project can report scope
     let finalArgs = args;
-    if (name === "system.test" && (!args || args.length === 0)) {
-      finalArgs = [{ project_id }];
+    if (name === "system.test") {
+      if (!args || args.length === 0 || typeof args[0] !== "object") {
+        finalArgs = [{}];
+      }
+      if (finalArgs[0] == null || typeof finalArgs[0] !== "object") {
+        finalArgs = [{ project_id }];
+      } else {
+        finalArgs = [{ ...finalArgs[0], project_id }];
+      }
     }
     const data = { name, args: finalArgs };
     // we use waitForInterest because often the project hasn't

src/packages/server/conat/api/system.ts

@@ -20,21 +20,13 @@ export function ping() {
 
 export async function test({
   account_id,
-  project_id,
-}: { account_id?: string; project_id?: string } = {}) {
+}: { account_id?: string } = {}) {
   // Return API key scope information and server time
-  // The authFirst decorator determines the scope from the API key and injects
-  // either account_id (for account-scoped keys) or project_id (for project-scoped keys)
-  // into this parameter object.
-  const response: { account_id?: string; project_id?: string; server_time: number } = {
+  // The authFirst decorator determines the scope from the API key and injects account_id.
+  const response: { account_id: string; server_time: number } = {
+    account_id: account_id ?? "",
     server_time: Date.now(),
   };
-  if (account_id) {
-    response.account_id = account_id;
-  }
-  if (project_id) {
-    response.project_id = project_id;
-  }
   return response;
 }
 

src/packages/server/projects/control/single-user.ts

@@ -156,7 +156,9 @@ class Project extends BaseProject {
 
       // First attempt: graceful shutdown with SIGTERM
       // This allows the process to clean up child processes (e.g., Jupyter kernels)
-      let usedSigterm = false;
+      const stopStartedAt = Date.now();
+      const SIGKILL_GRACE_MS = 5000;
+      let sigkillSent = false;
       const killProject = (signal: NodeJS.Signals = "SIGTERM") => {
         try {
           logger.debug(`stop: sending kill -${pid} with ${signal}`);
@@ -169,16 +171,15 @@ class Project extends BaseProject {
 
       // Try SIGTERM first for graceful shutdown
       killProject("SIGTERM");
-      usedSigterm = true;
 
       await this.wait({
         until: async () => {
           if (await isProjectRunning(this.HOME)) {
-            // After 5 seconds, escalate to SIGKILL
-            if (usedSigterm) {
+            // After a grace period, escalate to SIGKILL
+            if (!sigkillSent && Date.now() - stopStartedAt >= SIGKILL_GRACE_MS) {
               logger.debug("stop: escalating to SIGKILL");
               killProject("SIGKILL");
-              usedSigterm = false;
+              sigkillSent = true;
             }
             return false;
           } else {

src/python/cocalc-api/src/cocalc_api/mcp/DEVELOPMENT.md

@@ -174,13 +174,15 @@ except Exception as e:
 
 - `COCALC_API_KEY` - API key (required)
 - `COCALC_HOST` - CoCalc instance URL (optional, defaults to `https://cocalc.com`)
-- `COCALC_PROJECT_ID` - Project ID for project-scoped keys (optional, embedded in key)
+- `COCALC_PROJECT_ID` - Optional project ID used only with **account-scoped** keys to target a specific project. Ignored for project-scoped keys (project_id comes from the key itself).
 
 ### API Key Scope Detection
 
-The server calls `hub.system.test()` to determine scope:
-- If returns `account_id` → Account-scoped key
-- If returns `project_id` → Project-scoped key
+The server detects scope in two steps:
+- Call `hub.system.test()` (account-scoped keys only) → returns `account_id`.
+- If that fails, call `project.system.test()` → returns `project_id` for project-scoped keys.
+
+If `COCALC_PROJECT_ID` is provided with an account-scoped key, it is used as the default project target; it is ignored for project-scoped keys.
 
 ## Available CoCalc APIs
 

src/python/cocalc-api/src/cocalc_api/mcp/README.md

@@ -9,6 +9,8 @@ A Model Context Protocol (MCP) server that provides LLMs (Claude, etc.) with dir
 ```bash
 export COCALC_API_KEY="sk-your-api-key"  # Account or project-scoped
 export COCALC_HOST="http://localhost:5000"  # Optional, defaults to https://cocalc.com
+# Optional: only used with account-scoped keys to target a specific project
+# export COCALC_PROJECT_ID="your-project-uuid"
 ```
 
 ### 2. Run the Server
@@ -18,6 +20,7 @@ uv run cocalc-mcp-server
 ```
 
 The server will detect your API key type and automatically register the appropriate tools/resources.
+If you supply `COCALC_PROJECT_ID` with an account-scoped key, the MCP server will also prepare a project client for that project. For project-scoped keys, `COCALC_PROJECT_ID` is ignored because the project is embedded in the key.
 
 ## Setup with Claude Code
 

src/python/cocalc-api/src/cocalc_api/mcp/mcp_server.py

@@ -40,7 +40,7 @@
 import os
 import sys
 import time
-from typing import Optional
+from typing import Optional, TypedDict, Union, cast
 
 from mcp.server.fastmcp import FastMCP
 
@@ -93,7 +93,18 @@ def get_config() -> tuple[str, str, Optional[str]]:
     return api_key, host, project_id
 
 
-def check_api_key_scope(api_key: str, host: str) -> dict[str, str]:
+class AccountScope(TypedDict):
+    account_id: str
+
+
+class ProjectScope(TypedDict):
+    project_id: str
+
+
+Scope = Union[AccountScope, ProjectScope]
+
+
+def check_api_key_scope(api_key: str, host: str) -> Scope:
     """
     Check if the API key is account-scoped or project-scoped.
 
@@ -107,29 +118,26 @@ def check_api_key_scope(api_key: str, host: str) -> dict[str, str]:
     Raises:
         RuntimeError: If the API key is invalid or scope cannot be determined
     """
+    # Try account scope first; hub.system.test only works for account-scoped keys
     try:
-        hub = Hub(api_key=api_key, host=host)
-
-        # Try the hub.system.test() method (only works for account-scoped keys)
-        result = hub.system.test()
-
-        # Check which scope is returned
-        if "account_id" in result and result["account_id"]:
-            return {"account_id": result["account_id"]}
-        elif "project_id" in result and result["project_id"]:
-            return {"project_id": result["project_id"]}
-        else:
-            raise RuntimeError("API key test returned neither account_id nor project_id")
-
+        result = Hub(api_key=api_key, host=host).system.test()
+        account_id = result.get("account_id")
+        if account_id:
+            return {"account_id": account_id}
+    except Exception:
+        pass
+
+    # Fall back to project scope
+    try:
+        result = Project(api_key=api_key, host=host).system.test()
+        project_id = result.get("project_id")
+        if project_id:
+            return {"project_id": project_id}
     except Exception as e:
-        # Check if this looks like a project-scoped key error
-        error_msg = str(e)
-        if "must be signed in and MUST provide an api key" in error_msg:
-            raise RuntimeError("API key appears to be project-scoped. "
-                               "Project-scoped keys require the project_id to be specified at the OS level. "
-                               "Please set the COCALC_PROJECT_ID environment variable and try again.") from e
         raise RuntimeError(f"API key validation failed: {e}") from e
 
+    raise RuntimeError("API key test returned neither account_id nor project_id")
+
 
 # Initialize FastMCP server with instructions and documentation
 mcp = FastMCP(
@@ -168,7 +176,7 @@ def check_api_key_scope(api_key: str, host: str) -> dict[str, str]:
 # Configuration (initialized at startup)
 _api_key: Optional[str] = None
 _host: Optional[str] = None
-_api_key_scope: Optional[dict[str, str]] = None  # Either {"account_id": ...} or {"project_id": ...}
+_api_key_scope: Optional[Scope] = None  # Either {"account_id": ...} or {"project_id": ...}
 
 # Lazy-initialized project clients map: project_id -> Project
 _project_clients: dict[str, Project] = {}
@@ -191,41 +199,31 @@ def _initialize_config() -> None:
 
     # Validate API key and determine scope
     try:
-        try:
-            _api_key_scope = check_api_key_scope(_api_key, _host)
-        except RuntimeError as check_error:
-            # If it's a project-scoped key error, use a placeholder project_id
-            # Project-scoped keys have the project_id embedded in the key itself
-            if "project-scoped" in str(check_error):
-                # Use empty string as project_id - the Project client will extract it from the API key
-                _api_key_scope = {"project_id": ""}
-                print("✓ Connected with project-scoped API key", file=sys.stderr)
-            else:
-                raise
+        _api_key_scope = check_api_key_scope(_api_key, _host)
 
-        if "account_id" in _api_key_scope:
-            account_id = _api_key_scope["account_id"]
+        scope = _api_key_scope
+        if scope is None:
+            raise RuntimeError("Could not determine API key scope")
+
+        if "account_id" in scope:
+            account_id = cast(AccountScope, scope)["account_id"]
             print(f"✓ Connected with account-scoped API key (account: {account_id})", file=sys.stderr)
-        elif "project_id" in _api_key_scope:
-            project_id = _api_key_scope["project_id"]
-            # For project-scoped keys with empty/None project_id, the Project client will extract it from the API key
-            if project_id:
-                print(f"✓ Connected with project-scoped API key (project: {project_id})", file=sys.stderr)
-                # For project-scoped keys, eagerly create the project client
-                client = Project(api_key=_api_key, project_id=project_id, host=_host)
-                _project_clients[project_id] = client
-            else:
-                # Project-scoped key with empty project_id - will be discovered on first use
-                print("✓ Connected with project-scoped API key (project ID will be discovered on first use)", file=sys.stderr)
-        else:
-            # If we got here with no project_id but it might be project-scoped, check if COCALC_PROJECT_ID was provided
+            # If a project_id is explicitly provided via env, prepare a client for it
             if project_id_config:
-                _api_key_scope = {"project_id": project_id_config}
-                print(f"✓ Using project-scoped API key with explicitly provided project_id (project: {project_id_config})", file=sys.stderr)
                 client = Project(api_key=_api_key, project_id=project_id_config, host=_host)
                 _project_clients[project_id_config] = client
-            else:
-                raise RuntimeError("Could not determine API key scope")
+                print(
+                    f"✓ Using account-scoped API key with explicitly provided project_id (project: {project_id_config})",
+                    file=sys.stderr,
+                )
+        elif "project_id" in scope:
+            project_id = cast(ProjectScope, scope)["project_id"]
+            print(f"✓ Connected with project-scoped API key (project: {project_id})", file=sys.stderr)
+            # For project-scoped keys, eagerly create the project client
+            client = Project(api_key=_api_key, project_id=project_id, host=_host)
+            _project_clients[project_id] = client
+        else:
+            raise RuntimeError("Could not determine API key scope")
 
     except RuntimeError as e:
         print(f"Error: {e}", file=sys.stderr)
@@ -255,8 +253,9 @@ def get_project_client(project_id: Optional[str] = None) -> Project:
     # Determine which project_id to use
     if project_id is None:
         # If no project_id provided, try to use the one from project-scoped key
-        if _api_key_scope and "project_id" in _api_key_scope:
-            project_id = _api_key_scope["project_id"]
+        scope = _api_key_scope
+        if scope and "project_id" in scope:
+            project_id = cast(ProjectScope, scope)["project_id"]
         else:
             # Account-scoped key requires explicit project_id
             raise RuntimeError("Account-scoped API key requires an explicit project_id argument. "