planar 0.5.0__py3-none-any.whl → 0.8.0__py3-none-any.whl

planar/app.py

@@ -1,6 +1,8 @@
 import asyncio
+import signal
 from asyncio import CancelledError
 from contextlib import asynccontextmanager
+from types import FrameType
 from typing import Any, Callable, Coroutine, Type
 
 from fastapi import APIRouter, FastAPI, HTTPException, Request
@@ -11,7 +13,7 @@ from sqlalchemy.ext.asyncio import AsyncEngine
 from typing_extensions import TypeVar
 
 from planar.ai import Agent
-from planar.config import PlanarConfig, load_environment_aware_config
+from planar.config import Environment, PlanarConfig, load_environment_aware_config
 from planar.db import DatabaseManager
 from planar.files.storage.base import Storage
 from planar.files.storage.config import create_from_config
@@ -30,8 +32,8 @@ from planar.routers.entity_router import create_entities_router
 from planar.routers.object_config_router import create_object_config_router
 from planar.routers.rule import create_rule_router
 from planar.rules.decorator import RULE_REGISTRY
+from planar.security.auth_middleware import AuthMiddleware
 from planar.security.authorization import PolicyService, policy_service_context
-from planar.security.jwt_middleware import JWTMiddleware
 from planar.session import config_var, session_context
 from planar.sse.proxy import SSEProxy
 from planar.workflows import (
@@ -92,7 +94,7 @@ class PlanarApp:
         setup_orchestrator_middleware(self)
         setup_workflow_notification_middleware(self)
         setup_tracer_middleware(self)
-        setup_jwt_middleware(self)
+        setup_auth_middleware(self)
         setup_http_exception_handler(self)
         setup_authorization_policy_service(self)
 
@@ -202,6 +204,27 @@ class PlanarApp:
 
     @asynccontextmanager
     async def _lifespan(self, app: FastAPI):
+        # We manually capture SIGINT/SIGTERM to trigger our own graceful shutdown.
+        # This is necessary because long-lived connections, such as from the SSE
+        # proxy, can cause uvicorn's default graceful shutdown to hang, preventing
+        # the lifespan shutdown logic (after the yield) from ever being reached.
+        # Our handler starts the shutdown of these components and then chains to the
+        # original uvicorn handler to allow it to proceed with its own shutdown.
+        original_handlers = {
+            signal.SIGINT: signal.getsignal(signal.SIGINT),
+            signal.SIGTERM: signal.getsignal(signal.SIGTERM),
+        }
+
+        def terminate_now(signum: int, frame: FrameType | None = None):
+            asyncio.create_task(self.graceful_shutdown())
+            handler = original_handlers.get(signal.Signals(signum))
+            if callable(handler):
+                handler(signum, frame)
+
+        signal.signal(signal.SIGINT, terminate_now)
+        signal.signal(signal.SIGTERM, terminate_now)
+
+        # Begin the normal lifespan logic
         self.db_manager.connect()
         await self.db_manager.migrate(
             self.config.use_alembic if self.config.use_alembic is not None else True
@@ -240,6 +263,10 @@ class PlanarApp:
             config_var.reset(config_tok)
 
         await self.db_manager.disconnect()
+
+        if self.storage:
+            await self.storage.close()
+
         logger.info("stopping sse")
         await self.stop_sse()
         logger.info("lifespan completed")
@@ -435,15 +462,15 @@ def setup_http_exception_handler(app: PlanarApp):
 
 def setup_cors_middleware(app: PlanarApp):
     opts = {
-        "allow_headers": app.config.cors.allow_headers,
-        "allow_methods": app.config.cors.allow_methods,
-        "allow_credentials": app.config.cors.allow_credentials,
+        "allow_headers": app.config.security.cors.allow_headers,
+        "allow_methods": app.config.security.cors.allow_methods,
+        "allow_credentials": app.config.security.cors.allow_credentials,
     }
 
-    if isinstance(app.config.cors.allow_origins, str):
-        opts["allow_origin_regex"] = app.config.cors.allow_origins
+    if isinstance(app.config.security.cors.allow_origins, str):
+        opts["allow_origin_regex"] = app.config.security.cors.allow_origins
     else:
-        opts["allow_origins"] = app.config.cors.allow_origins
+        opts["allow_origins"] = app.config.security.cors.allow_origins
 
     app.fastapi.add_middleware(
         CORSMiddleware,
@@ -451,32 +478,49 @@ def setup_cors_middleware(app: PlanarApp):
     )
 
 
-def setup_jwt_middleware(app: PlanarApp):
-    if app.config.jwt and app.config.jwt.enabled and app.config.jwt.client_id:
-        client_id = app.config.jwt.client_id
-        org_id = app.config.jwt.org_id
-        additional_exclusion_paths = app.config.jwt.additional_exclusion_paths
+def setup_auth_middleware(app: PlanarApp):
+    if (
+        app.config.security
+        and app.config.security.jwt
+        and app.config.security.jwt.client_id
+        and app.config.security.jwt.org_id
+    ):
+        client_id = app.config.security.jwt.client_id
+        org_id = app.config.security.jwt.org_id
+        additional_exclusion_paths = app.config.security.jwt.additional_exclusion_paths
         app.fastapi.add_middleware(
-            JWTMiddleware,  # type: ignore
+            AuthMiddleware,  # type: ignore
             client_id,
             org_id,
             additional_exclusion_paths,
+            service_token=app.config.security.service_token.token
+            if app.config.security.service_token
+            and app.config.security.service_token.token
+            else None,
         )
         logger.info(
-            "jwt middleware enabled",
+            "Auth middleware enabled",
             client_id=client_id,
             org_id=org_id,
             additional_exclusion_paths=additional_exclusion_paths,
         )
+    elif app.config.environment == Environment.PROD:
+        raise ValueError(
+            "Auth middleware is required in production. Please set the JWT config and optionally service token config."
+        )
     else:
-        logger.warning("JWT middleware disabled")
+        logger.warning("Auth middleware disabled")
 
 
 def setup_authorization_policy_service(app: PlanarApp):
-    if app.config.authz and app.config.authz.enabled:
+    if (
+        app.config.security
+        and app.config.security.authz
+        and app.config.security.authz.enabled
+    ):
         app.policy_service = PolicyService(
-            policy_file_path=app.config.authz.policy_file
-            if app.config.authz.policy_file
+            policy_file_path=app.config.security.authz.policy_file
+            if app.config.security.authz.policy_file
             else None
         )
         logger.info(

planar/cli.py

@@ -13,28 +13,6 @@ from planar.config import Environment
 app = typer.Typer(help="Planar CLI tool")
 
 
-class PlanarServer(uvicorn.Server):
-    """Intercept SIGINT/SIGTERM to trigger early shutdown on the app."""
-
-    def __init__(self, config: uvicorn.Config, app_import_string: str):
-        super().__init__(config)
-        self.app_import_string = app_import_string
-
-    def handle_exit(self, sig, frame):
-        # Import the PlanarApp instance and fire its early-shutdown hook
-        import asyncio
-        import importlib
-
-        module_name, var_name = self.app_import_string.split(":")
-        app_module = importlib.import_module(module_name)
-        planar_app = getattr(app_module, var_name, None)
-        if planar_app and hasattr(planar_app, "graceful_shutdown"):
-            asyncio.create_task(planar_app.graceful_shutdown())
-
-        # Continue with Uvicorn's normal shutdown procedure
-        super().handle_exit(sig, frame)
-
-
 def find_default_app_path() -> Path:
     """Checks for default app file paths (app.py, then main.py)."""
     for filename in ["app.py", "main.py"]:
@@ -94,9 +72,25 @@ def dev_command(
         "--script",
         help="Run as a script with 'uv run' instead of starting a server",
     ),
+    ssl_keyfile: str | None = typer.Option(
+        None, "--ssl-keyfile", help="Path to SSL key file"
+    ),
+    ssl_certfile: str | None = typer.Option(
+        None, "--ssl-certfile", help="Path to SSL cert file"
+    ),
 ):
     """Run Planar in development mode"""
-    run_command(Environment.DEV, port, host, config, path, app_name, script)
+    run_command(
+        Environment.DEV,
+        port,
+        host,
+        config,
+        path,
+        app_name,
+        script,
+        ssl_keyfile,
+        ssl_certfile,
+    )
 
 
 @app.command("prod")
@@ -119,9 +113,25 @@ def prod_command(
         "--script",
         help="Run as a script with 'uv run' instead of starting a server",
     ),
+    ssl_keyfile: str | None = typer.Option(
+        None, "--ssl-keyfile", help="Path to SSL key file"
+    ),
+    ssl_certfile: str | None = typer.Option(
+        None, "--ssl-certfile", help="Path to SSL cert file"
+    ),
 ):
     """Run Planar in production mode"""
-    run_command(Environment.PROD, port, host, config, path, app_name, script)
+    run_command(
+        Environment.PROD,
+        port,
+        host,
+        config,
+        path,
+        app_name,
+        script,
+        ssl_keyfile,
+        ssl_certfile,
+    )
 
 
 def run_command(
@@ -132,6 +142,8 @@ def run_command(
     path: Path | None,
     app_name: str,
     script: bool = False,
+    ssl_keyfile: str | None = None,
+    ssl_certfile: str | None = None,
 ):
     """Common logic for both dev and prod commands"""
     os.environ["PLANAR_ENV"] = env.value
@@ -188,15 +200,15 @@ def run_command(
         typer.echo(f"Starting Planar in {env.value} mode")
 
         try:
-            config = uvicorn.Config(
+            uvicorn.run(
                 app_import_string,
                 host=host or ("127.0.0.1" if env == Environment.DEV else "0.0.0.0"),
                 port=port or 8000,
                 reload=True if env == Environment.DEV else False,
                 timeout_graceful_shutdown=4,
+                ssl_keyfile=ssl_keyfile,
+                ssl_certfile=ssl_certfile,
             )
-
-            PlanarServer(config, app_import_string).run()
         except Exception as e:
             # Provide more context on import errors
             if isinstance(e, (ImportError, AttributeError)):

planar/config.py

@@ -5,13 +5,14 @@ import os
 import sys
 from enum import Enum
 from pathlib import Path
-from typing import Annotated, Any, Dict, Literal, Optional
+from typing import Annotated, Any, Dict, Literal
 
 import boto3
 import yaml
 from dotenv import load_dotenv
 from pydantic import (
     BaseModel,
+    ConfigDict,
     Field,
     HttpUrl,
     SecretStr,
@@ -45,8 +46,8 @@ class LogLevel(str, Enum):
 
 class LoggerConfig(BaseModel):
     level: LogLevel = LogLevel.INFO
-    propagate: Optional[bool] = False
-    file: Optional[str] = None
+    propagate: bool | None = False
+    file: str | None = None
 
 
 class SQLiteConfig(BaseModel):
@@ -66,11 +67,11 @@ class PostgreSQLConfig(BaseModel):
     driver: Literal["postgresql", "postgresql+asyncpg"] = (
         "postgresql+asyncpg"  # Allow async PostgreSQL
     )
-    host: Optional[str] = None
-    port: Optional[int] = None
-    user: Optional[str] = None
-    password: Optional[str] = None
-    db: Optional[str]
+    host: str | None = None
+    port: int | None = None
+    user: str | None = None
+    password: str | None = None
+    db: str | None
 
     def connection_url(self) -> URL:
         driver = self.driver
@@ -92,15 +93,15 @@ class OpenAIConfig(BaseModel):
     """Configuration for OpenAI provider."""
 
     api_key: SecretStr
-    base_url: Optional[str] = None
-    organization: Optional[str] = None
+    base_url: str | None = None
+    organization: str | None = None
 
 
 class AnthropicConfig(BaseModel):
     """Configuration for Anthropic provider."""
 
     api_key: SecretStr
-    base_url: Optional[str] = None
+    base_url: str | None = None
 
 
 class GeminiConfig(BaseModel):
@@ -112,9 +113,9 @@ class GeminiConfig(BaseModel):
 class AIProvidersConfig(BaseModel):
     """Configuration for AI providers."""
 
-    openai: Optional[OpenAIConfig] = None
-    anthropic: Optional[AnthropicConfig] = None
-    gemini: Optional[GeminiConfig] = None
+    openai: OpenAIConfig | None = None
+    anthropic: AnthropicConfig | None = None
+    gemini: GeminiConfig | None = None
 
 
 DatabaseConfig = Annotated[
@@ -124,7 +125,7 @@ DatabaseConfig = Annotated[
 
 class AppConfig(BaseModel):
     db_connection: str
-    max_db_conflict_retries: Optional[int] = None
+    max_db_conflict_retries: int | None = None
 
 
 def default_storage_config() -> StorageConfig:
@@ -162,31 +163,27 @@ PROD_CORS_CONFIG = CorsConfig(
 
 
 class JWTConfig(BaseModel):
-    enabled: bool = False
     client_id: str | None = None
     org_id: str | None = None
     additional_exclusion_paths: list[str] | None = Field(default_factory=list)
 
     @model_validator(mode="after")
     def validate_client_id(cls, instance):
-        if instance.enabled and not instance.client_id:
-            raise ValueError("client_id is required when JWT is enabled")
-        if instance.client_id and not instance.enabled:
-            raise ValueError(
-                "You cannot specify a client_id without enabling JWT - did you mean to set enabled=True?"
-            )
+        if not instance.client_id or not instance.org_id:
+            raise ValueError("Both client_id and org_id required to enable JWT")
         return instance
 
 
-JWT_DISABLED_CONFIG = JWTConfig(enabled=False)
+# Coplane ORG JWT config
 JWT_COPLANE_CONFIG = JWTConfig(
-    enabled=True, client_id="client_01JSJHJP9Q8GZDK5Y856FEHTB0", org_id=None
+    client_id="client_01JSJHJPKG09TMSK6NHJP0S180",
+    org_id="org_01JY4QP57Y7H4EQ7HT3BGN7TNK",
 )
 
 
 class OtelConfig(BaseModel):
     collector_endpoint: HttpUrl
-    resource_attributes: Optional[dict[str, str]] = None
+    resource_attributes: dict[str, str] | None = None
 
 
 def install_otel_provider(otel_config: OtelConfig):
@@ -206,19 +203,31 @@ class AuthzConfig(BaseModel):
     policy_file: str | None = None
 
 
+class ServiceTokenConfig(BaseModel):
+    token: str | None = Field(None, min_length=1)
+
+
+class SecurityConfig(BaseModel):
+    cors: CorsConfig = PROD_CORS_CONFIG
+    jwt: JWTConfig | None = None
+    service_token: ServiceTokenConfig | None = None
+    authz: AuthzConfig | None = None
+
+
 class PlanarConfig(BaseModel):
     db_connections: Dict[str, DatabaseConfig | str]
     app: AppConfig
-    ai_providers: Optional[AIProvidersConfig] = None
-    storage: Optional[StorageConfig] = default_storage_config()
+    ai_providers: AIProvidersConfig | None = None
+    storage: StorageConfig | None = default_storage_config()
     sse_hub: str | bool = False
-    cors: CorsConfig = PROD_CORS_CONFIG
     environment: Environment = Environment.DEV
-    jwt: JWTConfig | None = None
-    logging: Optional[dict[str, LoggerConfig]] = None
+    security: SecurityConfig = SecurityConfig()
+    logging: dict[str, LoggerConfig] | None = None
     use_alembic: bool | None = True
-    otel: Optional[OtelConfig] = None
-    authz: AuthzConfig | None = None
+    otel: OtelConfig | None = None
+
+    # forbid extra keys in the config to prevent accidental misconfiguration
+    model_config = ConfigDict(extra="forbid")
 
     @model_validator(mode="after")
     def validate_db_connection_reference(cls, instance):
@@ -467,14 +476,14 @@ def load_environment_aware_config[ConfigClass]() -> PlanarConfig:
 
     if env == "dev":
         base_config = sqlite_config(db_path="planar_dev.db")
-        base_config.cors = LOCAL_CORS_CONFIG
+        base_config.security = SecurityConfig(cors=LOCAL_CORS_CONFIG)
         base_config.environment = Environment.DEV
-        base_config.jwt = JWT_DISABLED_CONFIG
     else:
         base_config = sqlite_config(db_path="planar.db")
-        base_config.cors = PROD_CORS_CONFIG
         base_config.environment = Environment.PROD
-        base_config.jwt = JWT_COPLANE_CONFIG
+        base_config.security = SecurityConfig(
+            cors=PROD_CORS_CONFIG, jwt=JWT_COPLANE_CONFIG
+        )
 
     # Convert base config to dict for merging
     # Use by_alias=False to work with Python field names before validation

planar/db/db.py

@@ -176,7 +176,8 @@ class DatabaseManager:
 
     def _create_sqlite_engine(self, url: URL) -> AsyncEngine:
         # in practice this high timeout is only use
-        timeout = int(str(url.query.get("timeout", 10)))
+        timeout = int(str(url.query.get("timeout", 60)))
+        logger.info("Setting up SQLite engine with timeout", timeout=timeout)
 
         engine = create_async_engine(
             url,