diff --git a/.env.template b/.env.template index 3d96b431..67cb0a61 100644 --- a/.env.template +++ b/.env.template @@ -1,5 +1,15 @@ CONNECTION_URI=postgresql+psycopg://testuser:testpwd@localhost:5432/honcho # sample for local database -DATABASE_SCHEMA=honcho +# CONNECTION_URI=postgresql+psycopg://testuser:testpwd@database:5432/honcho # sample for docker-compose database + +# Use something unique here if you want to share a database with other projects. +# Leave blank for public. Make sure to avoid `-` in name. +DATABASE_SCHEMA= + +# Auth +# Set to true to enable API authorization. Blank is equivalent to false. +USE_AUTH=false +# Required if USE_AUTH is true. Generate with scripts/generate_jwt_secret.py +AUTH_JWT_SECRET= OPENAI_API_KEY= OPENAI_BASE_URL=https://api.openai.com/v1 @@ -19,19 +29,10 @@ AZURE_OPENAI_DEPLOYMENT= OPENAI_COMPATIBLE_BASE_URL= OPENAI_COMPATIBLE_API_KEY= -# Logging -LOGFIRE_TOKEN= # optional logfire config - -# Auth -USE_AUTH_SERVICE=false -AUTH_SERVICE_URL= - # Sentry SENTRY_ENABLED=false SENTRY_DSN= -OPENTELEMETRY_ENABLED=false - # Deriver DERIVER_WORKERS=1 diff --git a/.github/workflows/fly-deploy.yml b/.github/workflows/fly-deploy.yml new file mode 100644 index 00000000..b0c246ed --- /dev/null +++ b/.github/workflows/fly-deploy.yml @@ -0,0 +1,18 @@ +# See https://fly.io/docs/app-guides/continuous-deployment-with-github-actions/ + +name: Fly Deploy +on: + push: + branches: + - main +jobs: + deploy: + name: Deploy app + runs-on: ubuntu-latest + concurrency: deploy-group # optional: ensure only one action runs at a time + steps: + - uses: actions/checkout@v4 + - uses: superfly/flyctl-actions/setup-flyctl@master + - run: flyctl deploy --remote-only + env: + FLY_API_TOKEN: ${{ secrets.FLY_API_TOKEN }} diff --git a/.github/workflows/unittest.yml b/.github/workflows/unittest.yml index 4b7b7574..733215e7 100644 --- a/.github/workflows/unittest.yml +++ b/.github/workflows/unittest.yml @@ -47,7 +47,7 @@ jobs: run: uv run pytest -x env: CONNECTION_URI: postgresql+psycopg://postgres:postgres@localhost:5432/test_db - USE_AUTH_SERVICE: false + USE_AUTH: false SENTRY_ENABLED: false OPENTELEMETRY_ENABLED: false OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }} diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index ca8e2676..44d6302f 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -91,10 +91,24 @@ that they are disabled you can verify the following environment variables are set to false. ```env -USE_AUTH_SERVICE=false +USE_AUTH=false SENTRY_ENABLED=false ``` +If you set `USE_AUTH` to true you will need to generate a JWT secret. You can +do this with the following command: + +```bash +python scripts/generate_jwt_secret.py +``` + +This will generate a JWT secret and print it to the console. You can then set +the `AUTH_JWT_SECRET` environment variable. This is required for `USE_AUTH`. + +```env +AUTH_JWT_SECRET= +``` + 5. Launch the API With the dependencies installed, a database setup and enabled with `pgvector`, diff --git a/docs/contributing/self-hosting.mdx b/docs/contributing/self-hosting.mdx index 0f752b56..c9d36812 100644 --- a/docs/contributing/self-hosting.mdx +++ b/docs/contributing/self-hosting.mdx @@ -18,7 +18,7 @@ The minimum poetry version is `0.4.9` ### Setup Once the dependencies are installed on the system run the following steps to get -the local project setup. +the local project setup. 1. Clone the repository @@ -49,14 +49,14 @@ source honcho/.venv/bin/activate 3. Set up a database -Honcho utilized [Postgres](https://www.postgresql.org/) for its database with +Honcho utilized [Postgres](https://www.postgresql.org/) for its database with pgvector. An easy way to get started with a postgresdb is to create a project with [Supabase](https://supabase.com/) A `docker-compose` template is also available with a database configuration -available. +available. -4. Edit the environment variables. +4. Edit the environment variables. Honcho uses a `.env` file for managing runtime environment variables. A `.env.template` file is included for convenience. Several of the configurations @@ -75,10 +75,10 @@ ANTHROPIC_API_KEY= # API Key for Anthropic used for the deriver and dialectic AP The template has the additional functionality disabled by default. To ensure that they are disabled you can verify the following environment variables are -set to false. +set to false. ```env -USE_AUTH_SERVICE=false +USE_AUTH=false SENTRY_ENABLED=false ``` @@ -99,7 +99,7 @@ necessary tables for Honcho to operate. As mentioned earlier a `docker-compose` template is included for running Honcho. As an alternative to running Honcho locally it can also be run with the compose -template. +template. Copy the template and update the appropriate environment variables before launching the service. diff --git a/pyproject.toml b/pyproject.toml index a0c38012..74252985 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,6 +23,7 @@ dependencies = [ "nanoid>=2.0.0", "alembic>=1.14.0", "langfuse>=2.57.1", + "pyjwt>=2.10.0", ] [tool.uv] dev-dependencies = [ diff --git a/scripts/generate_jwt_secret.py b/scripts/generate_jwt_secret.py new file mode 100755 index 00000000..097f7d49 --- /dev/null +++ b/scripts/generate_jwt_secret.py @@ -0,0 +1,38 @@ +#!/usr/bin/env python +""" +Utility script to generate a JWT secret for use in the .env file. +This uses the same logic as the automatically generated version in security.py. +""" + +import argparse +import secrets + + +def generate_jwt_secret(): + """Generate a random JWT secret using the secrets module.""" + return secrets.token_hex(32) + + +def main(): + parser = argparse.ArgumentParser( + description="Generate a JWT secret for authentication." + ) + parser.add_argument( + "--print-only", + action="store_true", + help="Only print the secret without instructions", + ) + args = parser.parse_args() + + secret = generate_jwt_secret() + + if args.print_only: + print(secret) + else: + print(f"Generated JWT secret: {secret}") + print("\nAdd this to your .env file as:") + print(f"AUTH_JWT_SECRET={secret}") + + +if __name__ == "__main__": + main() diff --git a/src/crud.py b/src/crud.py index d32c6ee7..313eae58 100644 --- a/src/crud.py +++ b/src/crud.py @@ -56,6 +56,29 @@ async def get_app(db: AsyncSession, app_id: str) -> models.App: return app +async def get_all_apps( + db: AsyncSession, + reverse: Optional[bool] = False, + filter: Optional[dict] = None, +) -> Select: + """ + Get all apps. + + Args: + db: Database session + reverse: Whether to reverse the order of the apps + filter: Filter the apps by a dictionary of metadata + """ + stmt = select(models.App) + if reverse: + stmt = stmt.order_by(models.App.id.desc()) + else: + stmt = stmt.order_by(models.App.id) + if filter is not None: + stmt = stmt.where(models.App.h_metadata.contains(filter)) + return stmt + + async def get_app_by_name(db: AsyncSession, name: str) -> models.App: """ Get an app by its name. @@ -391,7 +414,7 @@ async def create_session( """ try: # This will raise ResourceNotFoundException if user not found - honcho_user = await get_user(db, app_id=app_id, user_id=user_id) + _honcho_user = await get_user(db, app_id=app_id, user_id=user_id) honcho_session = models.Session( user_id=user_id, diff --git a/src/db.py b/src/db.py index c3498ba1..f26e8e12 100644 --- a/src/db.py +++ b/src/db.py @@ -1,7 +1,7 @@ import os from dotenv import load_dotenv -from sqlalchemy import MetaData, create_engine, inspect +from sqlalchemy import MetaData, create_engine, inspect, text from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine from sqlalchemy.orm import declarative_base @@ -57,6 +57,11 @@ def scaffold_db(): # Create inspector to check if database exists inspector = inspect(engine) + if table_schema: + with engine.connect() as connection: + connection.execute(text(f'CREATE SCHEMA IF NOT EXISTS "{table_schema}"')) + connection.commit() + print(inspector.get_table_names(Base.metadata.schema)) # If no tables exist, create them with SQLAlchemy diff --git a/src/exceptions.py b/src/exceptions.py index c62bb9e9..3088ce82 100644 --- a/src/exceptions.py +++ b/src/exceptions.py @@ -50,3 +50,10 @@ class AuthorizationException(HonchoException): status_code = 403 detail = "Not authorized to access this resource" + + +class DisabledException(HonchoException): + """Exception raised when a feature is disabled.""" + + status_code = 405 + detail = "Feature is disabled" diff --git a/src/main.py b/src/main.py index e7f75271..0d7f520a 100644 --- a/src/main.py +++ b/src/main.py @@ -10,18 +10,20 @@ from fastapi_pagination import add_pagination from sentry_sdk.integrations.fastapi import FastApiIntegration from sentry_sdk.integrations.starlette import StarletteIntegration +from src.db import engine, scaffold_db from src.exceptions import HonchoException from src.routers import ( apps, collections, documents, + keys, messages, metamessages, sessions, users, ) +from src.security import create_admin_jwt -from .db import engine def get_log_level(env_var="LOG_LEVEL", default="INFO"): @@ -56,8 +58,14 @@ logging.basicConfig( ) logger = logging.getLogger(__name__) -# Sentry Setup +# JWT Setup +async def setup_admin_jwt(): + token = create_admin_jwt() + print(f"\n ADMIN JWT: {token}\n") + + +# Sentry Setup SENTRY_ENABLED = os.getenv("SENTRY_ENABLED", "False").lower() == "true" if SENTRY_ENABLED: sentry_sdk.init( @@ -85,8 +93,8 @@ async def lifespan(app: FastAPI): app = FastAPI( lifespan=lifespan, servers=[ - {"url": "http://127.0.0.1:8000", "description": "Local Development Server"}, - {"url": "https:/demo.honcho.dev", "description": "Demo Server"}, + {"url": "http://localhost:8000", "description": "Local Development Server"}, + {"url": "https://demo.honcho.dev", "description": "Demo Server"}, ], title="Honcho API", summary="An API for adding personalization to AI Apps", @@ -126,6 +134,7 @@ app.include_router(messages.router, prefix="/v1") app.include_router(metamessages.router, prefix="/v1") app.include_router(collections.router, prefix="/v1") app.include_router(documents.router, prefix="/v1") +app.include_router(keys.router, prefix="/v1") # Global exception handlers diff --git a/src/routers/apps.py b/src/routers/apps.py index 43580441..aa180a20 100644 --- a/src/routers/apps.py +++ b/src/routers/apps.py @@ -1,30 +1,76 @@ import logging +from typing import Optional from fastapi import APIRouter, Depends +from fastapi_pagination import Page +from fastapi_pagination.ext.sqlalchemy import paginate from src import crud, schemas from src.dependencies import db -from src.exceptions import ResourceNotFoundException -from src.security import auth +from src.exceptions import AuthenticationException, ResourceNotFoundException +from src.security import require_auth logger = logging.getLogger(__name__) router = APIRouter( prefix="/apps", tags=["apps"], - dependencies=[Depends(auth)], ) +jwt_params = Depends(require_auth(app_id="app_id")) -@router.get("/{app_id}", response_model=schemas.App) + +@router.get( + "", + response_model=schemas.App, +) +async def get_app_from_token(jwt_params=jwt_params, db=db): + """ + Get an App by ID from the app_id provided in the JWT. + If no app_id is provided, return a 401 Unauthorized error. + """ + if jwt_params.ap is None: + raise AuthenticationException("App not found in JWT") + return await crud.get_app(db, app_id=jwt_params.ap) + + +@router.post( + "/list", + response_model=Page[schemas.App], + dependencies=[Depends(require_auth(admin=True))], +) +async def get_all_apps( + options: schemas.AppGet, + reverse: Optional[bool] = False, + db=db, +): + """Get all Apps""" + return await paginate( + db, + await crud.get_all_apps( + db, + reverse=reverse, + filter=options.filter, + ), + ) + + +@router.get( + "/{app_id}", + response_model=schemas.App, + dependencies=[Depends(require_auth(app_id="app_id"))], +) async def get_app(app_id: str, db=db): """Get an App by ID""" - # ResourceNotFoundException will be caught by global handler if app not found app = await crud.get_app(db, app_id=app_id) return app -@router.get("/name/{name}", response_model=schemas.App) +@router.get( + "/name/{name}", + response_model=schemas.App, + dependencies=[Depends(require_auth(admin=True))], +) async def get_app_by_name(name: str, db=db): """Get an App by Name""" # ResourceNotFoundException will be caught by global handler if app not found @@ -32,14 +78,20 @@ async def get_app_by_name(name: str, db=db): return app -@router.post("", response_model=schemas.App) +@router.post( + "", response_model=schemas.App, dependencies=[Depends(require_auth(admin=True))] +) async def create_app(app: schemas.AppCreate, db=db): """Create a new App""" honcho_app = await crud.create_app(db, app=app) return honcho_app -@router.get("/get_or_create/{name}", response_model=schemas.App) +@router.get( + "/get_or_create/{name}", + response_model=schemas.App, + dependencies=[Depends(require_auth(admin=True))], +) async def get_or_create_app(name: str, db=db): """Get or Create an App""" try: @@ -51,7 +103,11 @@ async def get_or_create_app(name: str, db=db): return app -@router.put("/{app_id}", response_model=schemas.App) +@router.put( + "/{app_id}", + response_model=schemas.App, + dependencies=[Depends(require_auth(app_id="app_id"))], +) async def update_app( app_id: str, app: schemas.AppUpdate, diff --git a/src/routers/collections.py b/src/routers/collections.py index 65a42bd5..d57ff2d3 100644 --- a/src/routers/collections.py +++ b/src/routers/collections.py @@ -6,16 +6,43 @@ from fastapi_pagination.ext.sqlalchemy import paginate from src import crud, schemas from src.dependencies import db -from src.security import auth +from src.exceptions import AuthenticationException +from src.security import require_auth router = APIRouter( prefix="/apps/{app_id}/users/{user_id}/collections", tags=["collections"], - dependencies=[Depends(auth)], ) +jwt_params = Depends(require_auth(collection_id="collection_id")) -@router.post("/list", response_model=Page[schemas.Collection]) + +@router.get( + "", + response_model=schemas.Collection, +) +async def get_collection_from_token( + app_id: str, + user_id: str, + jwt_params=jwt_params, + db=db, +): + """ + Get a specific collection for a user by collection_id provided in the JWT. + If no collection_id is provided, return a 401 Unauthorized error. + """ + if jwt_params.co is None: + raise AuthenticationException("Collection not found in JWT") + return await crud.get_collection_by_id( + db, app_id=app_id, collection_id=jwt_params.co, user_id=user_id + ) + + +@router.post( + "/list", + response_model=Page[schemas.Collection], + dependencies=[Depends(require_auth(app_id="app_id", user_id="user_id"))], +) async def get_collections( app_id: str, user_id: str, @@ -32,7 +59,11 @@ async def get_collections( ) -@router.get("/name/{name}", response_model=schemas.Collection) +@router.get( + "/name/{name}", + response_model=schemas.Collection, + dependencies=[Depends(require_auth(app_id="app_id", user_id="user_id"))], +) async def get_collection_by_name( app_id: str, user_id: str, @@ -46,7 +77,17 @@ async def get_collection_by_name( return honcho_collection -@router.get("/{collection_id}", response_model=schemas.Collection) +@router.get( + "/{collection_id}", + response_model=schemas.Collection, + dependencies=[ + Depends( + require_auth( + app_id="app_id", user_id="user_id", collection_id="collection_id" + ) + ) + ], +) async def get_collection_by_id( app_id: str, user_id: str, @@ -60,7 +101,11 @@ async def get_collection_by_id( return honcho_collection -@router.post("", response_model=schemas.Collection) +@router.post( + "", + response_model=schemas.Collection, + dependencies=[Depends(require_auth(app_id="app_id", user_id="user_id"))], +) async def create_collection( app_id: str, user_id: str, @@ -75,7 +120,17 @@ async def create_collection( ) -@router.put("/{collection_id}", response_model=schemas.Collection) +@router.put( + "/{collection_id}", + response_model=schemas.Collection, + dependencies=[ + Depends( + require_auth( + app_id="app_id", user_id="user_id", collection_id="collection_id" + ) + ) + ], +) async def update_collection( app_id: str, user_id: str, @@ -96,7 +151,16 @@ async def update_collection( return honcho_collection -@router.delete("/{collection_id}") +@router.delete( + "/{collection_id}", + dependencies=[ + Depends( + require_auth( + app_id="app_id", user_id="user_id", collection_id="collection_id" + ) + ) + ], +) async def delete_collection( app_id: str, user_id: str, diff --git a/src/routers/documents.py b/src/routers/documents.py index fa11fcba..c067c8f5 100644 --- a/src/routers/documents.py +++ b/src/routers/documents.py @@ -9,14 +9,20 @@ from fastapi_pagination.ext.sqlalchemy import paginate from src import crud, schemas from src.dependencies import db from src.exceptions import ResourceNotFoundException, ValidationException -from src.security import auth +from src.security import require_auth logger = logging.getLogger(__name__) router = APIRouter( prefix="/apps/{app_id}/users/{user_id}/collections/{collection_id}/documents", tags=["documents"], - dependencies=[Depends(auth)], + dependencies=[ + Depends( + require_auth( + app_id="app_id", user_id="user_id", collection_id="collection_id" + ) + ) + ], ) @@ -50,10 +56,7 @@ async def get_documents( ) from e -@router.get( - "/{document_id}", - response_model=schemas.Document, -) +@router.get("/{document_id}", response_model=schemas.Document) async def get_document( app_id: str, user_id: str, diff --git a/src/routers/keys.py b/src/routers/keys.py new file mode 100644 index 00000000..aaadd82e --- /dev/null +++ b/src/routers/keys.py @@ -0,0 +1,51 @@ +import logging +import os + +from fastapi import APIRouter, Depends + +from src.exceptions import DisabledException, ValidationException +from src.security import ( + JWTParams, + create_jwt, + require_auth, +) + +logger = logging.getLogger(__name__) + +USE_AUTH = os.getenv("USE_AUTH", "False").lower() == "true" + +router = APIRouter( + prefix="/keys", + tags=["keys"], + dependencies=[Depends(require_auth(admin=True))], +) + + +@router.post("") +async def create_key( + app_id: str | None = None, + user_id: str | None = None, + session_id: str | None = None, + collection_id: str | None = None, +): + """Create a new Key""" + if not USE_AUTH: + raise DisabledException() + + # Validate that at least one parameter is provided for proper scoping + if not any([app_id, user_id, session_id, collection_id]): + raise ValidationException( + "At least one of app_id, user_id, session_id, or collection_id must be provided" + ) + + key_str = create_jwt( + JWTParams( + ap=app_id, + us=user_id, + se=session_id, + co=collection_id, + ) + ) + return { + "key": key_str, + } diff --git a/src/routers/messages.py b/src/routers/messages.py index 7bdb0604..43dfd8d1 100644 --- a/src/routers/messages.py +++ b/src/routers/messages.py @@ -12,14 +12,18 @@ from src.db import SessionLocal from src.dependencies import db from src.exceptions import ResourceNotFoundException from src.models import QueueItem -from src.security import auth +from src.security import require_auth logger = logging.getLogger(__name__) router = APIRouter( prefix="/apps/{app_id}/users/{user_id}/sessions/{session_id}/messages", tags=["messages"], - dependencies=[Depends(auth)], + dependencies=[Depends(require_auth( + app_id="app_id", + user_id="user_id", + session_id="session_id" + ))], ) @@ -47,6 +51,11 @@ async def enqueue(payload: dict | list[dict]): user_id=payload[0]["user_id"], session_id=payload[0]["session_id"], ) + if not session: + logger.warning( + f"Session {payload[0]['session_id']} not found, skipping enqueue" + ) + return except ResourceNotFoundException: logger.warning( f"Session {payload[0]['session_id']} not found, skipping enqueue" @@ -97,6 +106,11 @@ async def enqueue(payload: dict | list[dict]): user_id=payload["user_id"], session_id=payload["session_id"], ) + if not session: + logger.warning( + f"Session {payload['session_id']} not found, skipping enqueue" + ) + return except ResourceNotFoundException: logger.warning( f"Session {payload['session_id']} not found, skipping enqueue" diff --git a/src/routers/metamessages.py b/src/routers/metamessages.py index d541ae01..083a8a57 100644 --- a/src/routers/metamessages.py +++ b/src/routers/metamessages.py @@ -8,14 +8,17 @@ from fastapi_pagination.ext.sqlalchemy import paginate from src import crud, schemas from src.dependencies import db from src.exceptions import ResourceNotFoundException, ValidationException -from src.security import auth +from src.security import require_auth logger = logging.getLogger(__name__) router = APIRouter( prefix="/apps/{app_id}/users/{user_id}/metamessages", tags=["metamessages"], - dependencies=[Depends(auth)], + dependencies=[Depends(require_auth( + app_id="app_id", + user_id="user_id" + ))], ) @@ -75,7 +78,6 @@ async def get_metamessages( filter=options.filter, reverse=reverse, ) - return await paginate(db, metamessages_query) except (ResourceNotFoundException, ValidationException) as e: logger.warning(f"Failed to get metamessages: {str(e)}") diff --git a/src/routers/sessions.py b/src/routers/sessions.py index c4110614..0b104efb 100644 --- a/src/routers/sessions.py +++ b/src/routers/sessions.py @@ -9,19 +9,49 @@ from fastapi_pagination.ext.sqlalchemy import paginate from src import agent, crud, schemas from src.dependencies import db -from src.exceptions import ResourceNotFoundException, ValidationException -from src.security import auth +from src.exceptions import ( + AuthenticationException, + ResourceNotFoundException, + ValidationException, +) +from src.security import require_auth logger = logging.getLogger(__name__) router = APIRouter( prefix="/apps/{app_id}/users/{user_id}/sessions", tags=["sessions"], - dependencies=[Depends(auth)], ) +jwt_params = Depends(require_auth(session_id="session_id")) -@router.post("/list", response_model=Page[schemas.Session]) + +@router.get( + "", + response_model=schemas.Session, +) +async def get_session_from_token( + app_id: str, + user_id: str, + jwt_params=jwt_params, + db=db, +): + """ + Get a specific session for a user by session_id provided in the JWT. + If no session_id is provided, return a 401 Unauthorized error. + """ + if jwt_params.se is None: + raise AuthenticationException("Session not found in JWT") + return await crud.get_session( + db, app_id=app_id, session_id=jwt_params.se, user_id=user_id + ) + + +@router.post( + "/list", + response_model=Page[schemas.Session], + dependencies=[Depends(require_auth(app_id="app_id", user_id="user_id"))], +) async def get_sessions( app_id: str, user_id: str, @@ -43,7 +73,11 @@ async def get_sessions( ) -@router.post("", response_model=schemas.Session) +@router.post( + "", + response_model=schemas.Session, + dependencies=[Depends(require_auth(app_id="app_id", user_id="user_id"))], +) async def create_session( app_id: str, user_id: str, @@ -62,7 +96,15 @@ async def create_session( raise ValidationException(str(e)) from e -@router.put("/{session_id}", response_model=schemas.Session) +@router.put( + "/{session_id}", + response_model=schemas.Session, + dependencies=[ + Depends( + require_auth(app_id="app_id", user_id="user_id", session_id="session_id") + ) + ], +) async def update_session( app_id: str, user_id: str, @@ -82,7 +124,14 @@ async def update_session( raise ResourceNotFoundException("Session not found") from e -@router.delete("/{session_id}") +@router.delete( + "/{session_id}", + dependencies=[ + Depends( + require_auth(app_id="app_id", user_id="user_id", session_id="session_id") + ) + ], +) async def delete_session( app_id: str, user_id: str, @@ -101,7 +150,15 @@ async def delete_session( raise ResourceNotFoundException("Session not found") from e -@router.get("/{session_id}", response_model=schemas.Session) +@router.get( + "/{session_id}", + response_model=schemas.Session, + dependencies=[ + Depends( + require_auth(app_id="app_id", user_id="user_id", session_id="session_id") + ) + ], +) async def get_session( app_id: str, user_id: str, @@ -118,7 +175,15 @@ async def get_session( return honcho_session -@router.post("/{session_id}/chat", response_model=schemas.AgentChat) +@router.post( + "/{session_id}/chat", + response_model=schemas.AgentChat, + dependencies=[ + Depends( + require_auth(app_id="app_id", user_id="user_id", session_id="session_id") + ) + ], +) async def chat( app_id: str, user_id: str, @@ -141,6 +206,11 @@ async def chat( }, } }, + dependencies=[ + Depends( + require_auth(app_id="app_id", user_id="user_id", session_id="session_id") + ) + ], ) async def get_chat_stream( app_id: str, @@ -168,7 +238,15 @@ async def get_chat_stream( ) -@router.get("/{session_id}/clone", response_model=schemas.Session) +@router.get( + "/{session_id}/clone", + response_model=schemas.Session, + dependencies=[ + Depends( + require_auth(app_id="app_id", user_id="user_id", session_id="session_id") + ) + ], +) async def clone_session( app_id: str, user_id: str, diff --git a/src/routers/users.py b/src/routers/users.py index ee64ce13..878923cc 100644 --- a/src/routers/users.py +++ b/src/routers/users.py @@ -7,20 +7,40 @@ from fastapi_pagination.ext.sqlalchemy import paginate from src import crud, schemas from src.dependencies import db from src.exceptions import ( + AuthenticationException, ResourceNotFoundException, ) -from src.security import auth +from src.security import require_auth logger = logging.getLogger(__name__) router = APIRouter( prefix="/apps/{app_id}/users", tags=["users"], - dependencies=[Depends(auth)], ) +jwt_params = Depends(require_auth(user_id="user_id")) -@router.post("", response_model=schemas.User) + +@router.get( + "", + response_model=schemas.User, +) +async def get_user_from_token(app_id: str, jwt_params=jwt_params, db=db): + """ + Get a User by ID from the user_id provided in the JWT. + If no user_id is provided, return a 401 Unauthorized error. + """ + if jwt_params.us is None: + raise AuthenticationException("User not found in JWT") + return await crud.get_user(db, app_id=app_id, user_id=jwt_params.us) + + +@router.post( + "", + response_model=schemas.User, + dependencies=[Depends(require_auth(app_id="app_id", user_id="user_id"))], +) async def create_user( app_id: str, user: schemas.UserCreate, @@ -31,7 +51,11 @@ async def create_user( return user_obj -@router.post("/list", response_model=Page[schemas.User]) +@router.post( + "/list", + response_model=Page[schemas.User], + dependencies=[Depends(require_auth(app_id="app_id", user_id="user_id"))], +) async def get_users( app_id: str, options: schemas.UserGet, @@ -45,7 +69,17 @@ async def get_users( ) -@router.get("/name/{name}", response_model=schemas.User) +@router.get( + "/name/{name}", + response_model=schemas.User, + dependencies=[ + Depends( + require_auth( + app_id="app_id", + ) + ) + ], +) async def get_user_by_name( app_id: str, name: str, @@ -56,7 +90,11 @@ async def get_user_by_name( return user -@router.get("/{user_id}", response_model=schemas.User) +@router.get( + "/{user_id}", + response_model=schemas.User, + dependencies=[Depends(require_auth(app_id="app_id", user_id="user_id"))], +) async def get_user( app_id: str, user_id: str, @@ -67,7 +105,17 @@ async def get_user( return user -@router.get("/get_or_create/{name}", response_model=schemas.User) +@router.get( + "/get_or_create/{name}", + response_model=schemas.User, + dependencies=[ + Depends( + require_auth( + app_id="app_id", + ) + ) + ], +) async def get_or_create_user(app_id: str, name: str, db=db): """Get a User or create a new one by the input name""" try: @@ -81,7 +129,11 @@ async def get_or_create_user(app_id: str, name: str, db=db): return user -@router.put("/{user_id}", response_model=schemas.User) +@router.put( + "/{user_id}", + response_model=schemas.User, + dependencies=[Depends(require_auth(app_id="app_id", user_id="user_id"))], +) async def update_user( app_id: str, user_id: str, diff --git a/src/schemas.py b/src/schemas.py index c4786431..9df273a3 100644 --- a/src/schemas.py +++ b/src/schemas.py @@ -13,6 +13,10 @@ class AppCreate(AppBase): metadata: dict = {} +class AppGet(AppBase): + filter: dict | None = None + + class AppUpdate(AppBase): name: str | None = None metadata: dict | None = None diff --git a/src/security.py b/src/security.py index c12eb163..f0aae28c 100644 --- a/src/security.py +++ b/src/security.py @@ -1,28 +1,187 @@ +import datetime import logging import os -from typing import Annotated +from typing import Annotated, Optional -from fastapi import Depends +import jwt +from fastapi import Depends, Request from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer +from pydantic import BaseModel +from sqlalchemy.ext.asyncio import AsyncSession + +from src.dependencies import get_db from .exceptions import AuthenticationException logger = logging.getLogger(__name__) -USE_AUTH_SERVICE = os.getenv("USE_AUTH_SERVICE", "False").lower() == "true" -SECRET_KEY = os.getenv("SECRET_KEY", "test") +USE_AUTH = os.getenv("USE_AUTH", "False").lower() == "true" +AUTH_JWT_SECRET = os.getenv("AUTH_JWT_SECRET", "") if USE_AUTH else "" + +if USE_AUTH and AUTH_JWT_SECRET == "": + print( + "\n ERROR: No JWT secret provided. Set the AUTH_JWT_SECRET environment variable.\n" + ) + exit(1) security = HTTPBearer( auto_error=False, ) +# +# jwt params +# all optional, used to produce tokens valid for different routes +# hierarchy: app > user > ( session / collection ) +# routes that involve a 'name' parameter require permissions for the parent object +# name routes are considered 'queries' as names are mutable properties +# +# note: add routes without parameters that assume the most immediately scoped key is providing +# +class JWTParams(BaseModel): + """ + JWT parameters used to produce tokens valid for different routes. + Hierarchy: app > user > (session / collection) + + All routers require at least the most tightly scoped parameter. + Routes will accept a JWT with a scope higher in the hierarchy. + + Names shortened to minimize token size. Timestamp is included + so that many unique tokens can be generated for the same resource. + Note that the timestamp itself is not used for security, and can + be omitted, such as when Honcho generates the initial admin JWT. + + Fields (all optional other than `t`): + + `t`: a string timestamp of when the JWT was created + `ad`: a boolean flag indicating if the JWT is an admin JWT + `ap`: (string) app id + `us`: (string) user id + `se`: (string) session id + `co`: (string) collection id + """ + + t: str = datetime.datetime.now().isoformat() + ad: Optional[bool] = None + ap: Optional[str] = None + us: Optional[str] = None + se: Optional[str] = None + co: Optional[str] = None + + +def create_admin_jwt() -> str: + """Create a JWT for admin operations.""" + params = JWTParams(t="", ad=True) + key = create_jwt(params) + return key + + +def create_jwt(params: JWTParams) -> str: + """Create a JWT token from the given parameters.""" + payload = {k: v for k, v in params.__dict__.items() if v is not None} + return jwt.encode(payload, AUTH_JWT_SECRET.encode("utf-8"), algorithm="HS256") + + +async def verify_jwt(token: str) -> JWTParams: + """Verify a JWT token and return the decoded parameters.""" + + params = JWTParams() + try: + decoded = jwt.decode( + token, AUTH_JWT_SECRET.encode("utf-8"), algorithms=["HS256"] + ) + if "t" in decoded: + params.t = decoded["t"] + if "ad" in decoded: + params.ad = decoded["ad"] + if "ap" in decoded: + params.ap = decoded["ap"] + if "us" in decoded: + params.us = decoded["us"] + if "se" in decoded: + params.se = decoded["se"] + if "co" in decoded: + params.co = decoded["co"] + return params + except jwt.PyJWTError: + raise AuthenticationException("Invalid JWT") from None + + +def require_auth( + admin: Optional[bool] = None, + app_id: Optional[str] = None, + user_id: Optional[str] = None, + session_id: Optional[str] = None, + collection_id: Optional[str] = None, +): + """ + Generate a dependency that requires authentication for the given parameters. + """ + + async def auth_dependency( + request: Request, + credentials: HTTPAuthorizationCredentials = Depends(security), + db: AsyncSession = Depends(get_db), + ): + app_id_param = request.path_params.get(app_id) if app_id else None + user_id_param = request.path_params.get(user_id) if user_id else None + session_id_param = request.path_params.get(session_id) if session_id else None + collection_id_param = ( + request.path_params.get(collection_id) if collection_id else None + ) + + return await auth( + credentials=credentials, + admin=admin, + app_id=app_id_param, + user_id=user_id_param, + session_id=session_id_param, + collection_id=collection_id_param, + ) + + return auth_dependency + + async def auth( credentials: Annotated[HTTPAuthorizationCredentials, Depends(security)], -): - if not USE_AUTH_SERVICE: - return True - if not credentials or credentials.credentials != SECRET_KEY: - logger.warning("Invalid access token attempt") - raise AuthenticationException("Invalid access token") - return {"message": "OK"} + admin: Optional[bool] = None, + app_id: Optional[str] = None, + user_id: Optional[str] = None, + session_id: Optional[str] = None, + collection_id: Optional[str] = None, +) -> JWTParams: + """Authenticate the given JWT and return the decoded parameters.""" + if not USE_AUTH: + return JWTParams(t="", ad=True) + if not credentials or not credentials.credentials: + logger.warning("No access token provided") + raise AuthenticationException("No access token provided") + + jwt_params = await verify_jwt(credentials.credentials) + + # based on api operation, verify api key based on that key's permissions + if jwt_params.ad: + return jwt_params + if admin: + raise AuthenticationException("Resource requires admin privileges") + + # Check if the JWT has direct access to the requested resource + # For session or collection level access + if session_id and jwt_params.se == session_id: + return jwt_params + if collection_id and jwt_params.co == collection_id: + return jwt_params + + # For user level access - can access all sessions/collections under this user + if user_id and jwt_params.us == user_id: + return jwt_params + + # For app level access - can access all users/sessions/collections under this app + if app_id and jwt_params.ap == app_id: + return jwt_params + + if any([session_id, collection_id, user_id, app_id]): + raise AuthenticationException("JWT not permissioned for this resource") + + # Route did not specify any parameters, so it should parse parameters itself + return jwt_params diff --git a/tests/conftest.py b/tests/conftest.py index 28a9d9fa..3286bd61 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,6 +1,7 @@ import logging # noqa: I001 import os import sys +import jwt from nanoid import generate as generate_nanoid from unittest.mock import patch, MagicMock @@ -19,6 +20,7 @@ from src import models from src.db import Base from src.dependencies import get_db from src.exceptions import HonchoException +from src.security import create_admin_jwt, create_jwt, JWTParams from src.main import app # Create a custom handler that doesn't get closed prematurely @@ -46,6 +48,10 @@ CONNECTION_URI = make_url(os.getenv("CONNECTION_URI", "postgresql+psycopg://post TEST_DB_URL = CONNECTION_URI.set(database="test_db") DEFAULT_DB_URL = str(CONNECTION_URI.set(database="postgres")) +# Test API authorization +USE_AUTH = os.getenv("USE_AUTH", "False").lower() == "true" +AUTH_JWT_SECRET = os.getenv("AUTH_JWT_SECRET", "test-secret") + def create_test_database(db_url): """Helper function create a database if it does not already exist @@ -129,7 +135,7 @@ async def db_session(db_engine): @pytest.fixture(scope="function") -def client(db_session): +async def client(db_session): """Create a FastAPI TestClient for the scope of a single test function""" # Register exception handlers for tests @@ -145,9 +151,49 @@ def client(db_session): app.dependency_overrides[get_db] = override_get_db with TestClient(app) as c: + if USE_AUTH: + # give the test client the admin JWT + c.headers["Authorization"] = f"Bearer {create_admin_jwt()}" yield c +def create_invalid_jwt() -> str: + return jwt.encode({"ad": "invalid"}, "this is not the secret", algorithm="HS256") + + +@pytest.fixture( + params=[ + ("none", None), # No auth + ("invalid", create_invalid_jwt), # Invalid JWT + ("empty", lambda: create_jwt(JWTParams())), # Empty JWT + ("admin", create_admin_jwt), # Admin JWT + ] +) +def auth_client(client, request, monkeypatch): + """ + Fixture that provides a client with different authentication states. + Always ensures USE_AUTH is set to True. + """ + # Ensure USE_AUTH is always True for this fixture + import src.routers.keys as keys_module + import src.security as security + + monkeypatch.setattr(keys_module, "USE_AUTH", "true") + monkeypatch.setattr(security, "USE_AUTH", "true") + + # Clear any existing Authorization header + client.headers.pop("Authorization", None) + + auth_type, token_func = request.param + client.auth_type = auth_type + + if token_func is not None: + token = token_func() + client.headers["Authorization"] = f"Bearer {token}" + + return client + + @pytest_asyncio.fixture(scope="function") async def sample_data(db_session): """Helper function to create test data""" diff --git a/tests/routes/test_apps.py b/tests/routes/test_apps.py index 1f5463fa..7c670d5c 100644 --- a/tests/routes/test_apps.py +++ b/tests/routes/test_apps.py @@ -1,3 +1,4 @@ +import pytest from nanoid import generate as generate_nanoid @@ -80,6 +81,39 @@ def test_get_app_by_id(client, sample_data): assert data["id"] == str(test_app.public_id) +@pytest.mark.asyncio +async def test_get_all_apps(client, db_session, sample_data): + test_app, test_user = sample_data + + # create a test app with metadata + response = client.post( + "/v1/apps", + json={ + "name": "test_app", + "metadata": {"test_key": "test_value"}, + }, + ) + + response = client.post( + "/v1/apps/list", + json={}, + ) + assert response.status_code == 200 + data = response.json() + assert "items" in data + assert len(data["items"]) > 0 + + response = client.post( + "/v1/apps/list", + json={"filter": {"test_key": "test_value"}}, + ) + assert response.status_code == 200 + data = response.json() + assert "items" in data + assert len(data["items"]) > 0 + assert data["items"][0]["metadata"] == {"test_key": "test_value"} + + def test_get_app_by_name(client, sample_data): test_app, _ = sample_data response = client.get(f"/v1/apps/name/{test_app.name}") diff --git a/tests/routes/test_keys.py b/tests/routes/test_keys.py new file mode 100644 index 00000000..5e0ec01c --- /dev/null +++ b/tests/routes/test_keys.py @@ -0,0 +1,44 @@ +def test_create_key_no_params(auth_client): + """Test creating a key with no parameters""" + response = auth_client.post("/v1/keys") + + # Only admin JWT should be allowed + if auth_client.auth_type == "admin": + # key with no params should fail + assert response.status_code == 422 + else: + assert response.status_code == 401 + + +def test_create_key_with_params(auth_client, sample_data): + """Test creating a key with specific parameters""" + test_app, test_user = sample_data + + if auth_client.auth_type != "admin": + return # Skip test if not admin authentication + + # Test with app_id + response = auth_client.post("/v1/keys", params={"app_id": test_app.public_id}) + assert response.status_code == 200 + assert "key" in response.json() + + # Test with app_id and user_id + response = auth_client.post( + "/v1/keys", + params={"app_id": test_app.public_id, "user_id": test_user.public_id}, + ) + assert response.status_code == 200 + assert "key" in response.json() + + # Test with session_id and collection_id + response = auth_client.post( + "/v1/keys", + params={ + "app_id": test_app.public_id, + "user_id": test_user.public_id, + "session_id": "test-session", + "collection_id": "test-collection", + }, + ) + assert response.status_code == 200 + assert "key" in response.json() diff --git a/tests/routes/test_scoped_api.py b/tests/routes/test_scoped_api.py new file mode 100644 index 00000000..0e4a99e1 --- /dev/null +++ b/tests/routes/test_scoped_api.py @@ -0,0 +1,552 @@ +from nanoid import generate as generate_nanoid + +from src.security import JWTParams, create_jwt + + +def test_create_app_with_auth(auth_client): + name = str(generate_nanoid()) + + response = auth_client.post( + "/v1/apps", json={"name": name, "metadata": {"key": "value"}} + ) + + # Check expected behavior based on auth type + if auth_client.auth_type != "admin": + assert response.status_code == 401 + return + + assert response.status_code == 200 + + +def test_auth_response_time(auth_client): + name = str(generate_nanoid()) + + import time + + start_time = time.time() + + response = auth_client.post( + "/v1/apps", json={"name": name, "metadata": {"key": "value"}} + ) + + end_time = time.time() + response_time = end_time - start_time + print( + f"Server response time for client {auth_client.auth_type}: {response_time:.6f} seconds" + ) + + # Check expected behavior based on auth type + if auth_client.auth_type != "admin": + assert response.status_code == 401 + return + + assert response.status_code == 200 + + +def test_get_or_create_app_with_auth(auth_client): + name = str(generate_nanoid()) + # Should return a ResourceNotFoundException with 404 status + response = auth_client.get(f"/v1/apps/name/{name}") + + if auth_client.auth_type != "admin": + assert response.status_code == 401 + return + + assert response.status_code == 404 + + response = auth_client.get(f"/v1/apps/get_or_create/{name}") + + if auth_client.auth_type != "admin": + assert response.status_code == 401 + return + + assert response.status_code == 200 + + +def test_get_app_by_id_with_auth(auth_client, sample_data): + test_app, _ = sample_data + + if auth_client.auth_type == "empty": + # For non-admin, include the app_id in the JWT + auth_client.headers["Authorization"] = ( + f"Bearer {create_jwt(JWTParams(ap=test_app.public_id))}" + ) + + response = auth_client.get(f"/v1/apps/{test_app.public_id}") + + # Admin JWT or JWT with matching app_id should be allowed + if auth_client.auth_type in ["admin", "empty"]: + assert response.status_code == 200 + else: + assert response.status_code == 401 + + +def test_get_app_from_token(auth_client, sample_data): + test_app, _ = sample_data + + if auth_client.auth_type == "empty": + # For non-admin, include the app_id in the JWT + auth_client.headers["Authorization"] = ( + f"Bearer {create_jwt(JWTParams(ap=test_app.public_id))}" + ) + + response = auth_client.get("/v1/apps") + + if auth_client.auth_type == "empty": + assert response.status_code == 200 + assert response.json()["id"] == test_app.public_id + else: + assert response.status_code == 401 + + +def test_get_app_by_name_with_auth(auth_client, sample_data): + test_app, _ = sample_data + + if auth_client.auth_type == "empty": + # For non-admin, include the app_id in the JWT + # Note that this will still fail because name route requires admin + auth_client.headers["Authorization"] = ( + f"Bearer {create_jwt(JWTParams(ap=test_app.public_id))}" + ) + + response = auth_client.get(f"/v1/apps/name/{test_app.name}") + + # Only admin JWT should be allowed + if auth_client.auth_type == "admin": + assert response.status_code == 200 + else: + assert response.status_code == 401 + + +def test_update_app_with_auth(auth_client, sample_data): + test_app, _ = sample_data + + if auth_client.auth_type == "empty": + # For non-admin, include the app_id in the JWT + # Note that this will still fail because name route requires admin + auth_client.headers["Authorization"] = ( + f"Bearer {create_jwt(JWTParams(ap=test_app.public_id))}" + ) + + new_name = str(generate_nanoid()) + response = auth_client.put( + f"/v1/apps/{test_app.public_id}", + json={"name": new_name, "metadata": {"new_key": "new_value"}}, + ) + + # Only admin JWT or JWT with matching app_id should be allowed + if auth_client.auth_type in ["admin", "empty"]: + assert response.status_code == 200 + else: + assert response.status_code == 401 + + +def test_update_app_with_wrong_auth(auth_client, sample_data): + test_app, _ = sample_data + + different_app = str(generate_nanoid()) + + if auth_client.auth_type == "empty": + # For non-admin, include the *wrong* app_id in the JWT + auth_client.headers["Authorization"] = ( + f"Bearer {create_jwt(JWTParams(ap=different_app))}" + ) + + new_name = str(generate_nanoid()) + response = auth_client.put( + f"/v1/apps/{test_app.public_id}", + json={"name": new_name, "metadata": {"new_key": "new_value"}}, + ) + + # Only admin JWT or JWT with matching app_id should be allowed + if auth_client.auth_type == "admin": + assert response.status_code == 200 + else: + # wrong app_id should be rejected + assert response.status_code == 401 + + +def test_create_user_with_auth(auth_client, sample_data): + test_app, _ = sample_data + + if auth_client.auth_type == "empty": + # For non-admin, include the app_id in the JWT + auth_client.headers["Authorization"] = ( + f"Bearer {create_jwt(JWTParams(ap=test_app.public_id))}" + ) + + name = str(generate_nanoid()) + response = auth_client.post( + f"/v1/apps/{test_app.public_id}/users", + json={"name": name, "metadata": {"user_key": "user_value"}}, + ) + + # Only admin JWT or JWT with matching app_id should be allowed + if auth_client.auth_type in ["admin", "empty"]: + assert response.status_code == 200 + else: + assert response.status_code == 401 + + +def test_get_user_by_id_with_auth(auth_client, sample_data): + test_app, test_user = sample_data + + if auth_client.auth_type == "empty": + # For non-admin, include the app_id in the JWT + auth_client.headers["Authorization"] = ( + f"Bearer {create_jwt(JWTParams(ap=test_app.public_id))}" + ) + + response = auth_client.get( + f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}" + ) + + # Admin JWT or JWT with matching app_id should be allowed + if auth_client.auth_type in ["admin", "empty"]: + assert response.status_code == 200 + else: + assert response.status_code == 401 + + # Test with user-scoped JWT + if auth_client.auth_type == "empty": + auth_client.headers["Authorization"] = ( + f"Bearer {create_jwt(JWTParams(us=test_user.public_id))}" + ) + + response = auth_client.get( + f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}" + ) + + assert response.status_code == 200 + + +def test_get_user_by_name_with_auth(auth_client, sample_data): + test_app, test_user = sample_data + + if auth_client.auth_type == "empty": + # For non-admin, include the app_id in the JWT + auth_client.headers["Authorization"] = ( + f"Bearer {create_jwt(JWTParams(ap=test_app.public_id))}" + ) + + response = auth_client.get( + f"/v1/apps/{test_app.public_id}/users/name/{test_user.name}" + ) + + # Admin JWT or JWT with matching app_id should be allowed + if auth_client.auth_type in ["admin", "empty"]: + assert response.status_code == 200 + else: + assert response.status_code == 401 + + +def test_update_user_with_auth(auth_client, sample_data): + test_app, test_user = sample_data + + if auth_client.auth_type == "empty": + # For non-admin, include the app_id in the JWT + auth_client.headers["Authorization"] = ( + f"Bearer {create_jwt(JWTParams(ap=test_app.public_id))}" + ) + + new_name = str(generate_nanoid()) + response = auth_client.put( + f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}", + json={"name": new_name, "metadata": {"updated_key": "updated_value"}}, + ) + + # Admin JWT or JWT with matching app_id should be allowed + if auth_client.auth_type in ["admin", "empty"]: + assert response.status_code == 200 + else: + assert response.status_code == 401 + + # Test with user-scoped JWT + if auth_client.auth_type == "empty": + auth_client.headers["Authorization"] = ( + f"Bearer {create_jwt(JWTParams(us=test_user.public_id))}" + ) + + response = auth_client.put( + f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}", + json={ + "name": str(generate_nanoid()), + "metadata": {"user_key": "user_value"}, + }, + ) + + assert response.status_code == 200 + + +def test_create_session_with_auth(auth_client, sample_data): + test_app, test_user = sample_data + + if auth_client.auth_type == "empty": + # For non-admin, include the app_id and user_id in the JWT + auth_client.headers["Authorization"] = ( + f"Bearer {create_jwt(JWTParams(ap=test_app.public_id, us=test_user.public_id))}" + ) + + response = auth_client.post( + f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/sessions", + json={}, + ) + + # Only admin JWT or JWT with matching app_id and user_id should be allowed + if auth_client.auth_type in ["admin", "empty"]: + assert response.status_code == 200 + else: + assert response.status_code == 401 + + # Remove app_id from header and make sure user-scoped key works too + if auth_client.auth_type == "empty": + auth_client.headers["Authorization"] = ( + f"Bearer {create_jwt(JWTParams(us=test_user.public_id))}" + ) + + response = auth_client.post( + f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/sessions", + json={}, + ) + + assert response.status_code == 200 + + +def test_get_session_by_id_with_auth(auth_client, sample_data): + test_app, test_user = sample_data + + # First create a session + if auth_client.auth_type == "empty": + auth_client.headers["Authorization"] = ( + f"Bearer {create_jwt(JWTParams(ap=test_app.public_id, us=test_user.public_id))}" + ) + + create_response = auth_client.post( + f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/sessions", + json={}, + ) + + if auth_client.auth_type not in ["admin", "empty"]: + assert create_response.status_code == 401 + return + + assert create_response.status_code == 200 + session_id = create_response.json()["id"] + + # Test with app and user scoped JWT + response = auth_client.get( + f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/sessions/{session_id}" + ) + assert response.status_code == 200 + + # Test with session-scoped JWT + if auth_client.auth_type == "empty": + auth_client.headers["Authorization"] = ( + f"Bearer {create_jwt(JWTParams(se=session_id))}" + ) + + response = auth_client.get( + f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/sessions/{session_id}" + ) + assert response.status_code == 200 + + +def test_create_collection(auth_client, sample_data) -> None: + test_app, test_user = sample_data + + if auth_client.auth_type == "empty": + # For non-admin, include the app_id and user_id in the JWT + auth_client.headers["Authorization"] = ( + f"Bearer {create_jwt(JWTParams(ap=test_app.public_id, us=test_user.public_id))}" + ) + + response = auth_client.post( + f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections", + json={"name": "test_collection", "metadata": {}}, + ) + + # Only admin JWT or JWT with matching app_id and user_id should be allowed + if auth_client.auth_type in ["admin", "empty"]: + assert response.status_code == 200 + else: + assert response.status_code == 401 + + # Remove app_id from header and make sure user-scoped key works too + if auth_client.auth_type == "empty": + auth_client.headers["Authorization"] = ( + f"Bearer {create_jwt(JWTParams(us=test_user.public_id))}" + ) + + response = auth_client.post( + f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections", + json={"name": "test_collection2", "metadata": {}}, + ) + + assert response.status_code == 200 + + # Remove user_id from header and make sure app-scoped key works too + if auth_client.auth_type == "empty": + auth_client.headers["Authorization"] = ( + f"Bearer {create_jwt(JWTParams(ap=test_app.public_id))}" + ) + + response = auth_client.post( + f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections", + json={"name": "test_collection3", "metadata": {}}, + ) + + assert response.status_code == 200 + + +def test_get_collection_by_id_with_auth(auth_client, sample_data) -> None: + test_app, test_user = sample_data + + # First create a collection + if auth_client.auth_type == "empty": + auth_client.headers["Authorization"] = ( + f"Bearer {create_jwt(JWTParams(ap=test_app.public_id, us=test_user.public_id))}" + ) + + create_response = auth_client.post( + f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections", + json={"name": "test_collection_get", "metadata": {}}, + ) + + if auth_client.auth_type not in ["admin", "empty"]: + assert create_response.status_code == 401 + return + + assert create_response.status_code == 200 + collection_id = create_response.json()["id"] + + # Test with app and user scoped JWT + response = auth_client.get( + f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections/{collection_id}" + ) + assert response.status_code == 200 + + # Test with collection-scoped JWT + if auth_client.auth_type == "empty": + auth_client.headers["Authorization"] = ( + f"Bearer {create_jwt(JWTParams(co=collection_id))}" + ) + + response = auth_client.get( + f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections/{collection_id}" + ) + assert response.status_code == 200 + + +def test_get_collection_by_name_with_auth(auth_client, sample_data) -> None: + test_app, test_user = sample_data + collection_name = f"test_collection_{generate_nanoid()}" + + # First create a collection + if auth_client.auth_type == "empty": + auth_client.headers["Authorization"] = ( + f"Bearer {create_jwt(JWTParams(ap=test_app.public_id, us=test_user.public_id))}" + ) + + create_response = auth_client.post( + f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections", + json={"name": collection_name, "metadata": {}}, + ) + + if auth_client.auth_type not in ["admin", "empty"]: + assert create_response.status_code == 401 + return + + assert create_response.status_code == 200 + + # Test with app and user scoped JWT + response = auth_client.get( + f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections/name/{collection_name}" + ) + assert response.status_code == 200 + + +def test_create_document_with_auth(auth_client, sample_data) -> None: + test_app, test_user = sample_data + + # First create a collection + if auth_client.auth_type == "empty": + auth_client.headers["Authorization"] = ( + f"Bearer {create_jwt(JWTParams(ap=test_app.public_id, us=test_user.public_id))}" + ) + + create_collection_response = auth_client.post( + f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections", + json={"name": "test_collection_docs", "metadata": {}}, + ) + + if auth_client.auth_type not in ["admin", "empty"]: + assert create_collection_response.status_code == 401 + return + + assert create_collection_response.status_code == 200 + collection_id = create_collection_response.json()["id"] + + # Create document with app and user scoped JWT + response = auth_client.post( + f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections/{collection_id}/documents", + json={"content": "Test document content", "metadata": {"doc_key": "doc_value"}}, + ) + assert response.status_code == 200 + + # Test with collection-scoped JWT + if auth_client.auth_type == "empty": + auth_client.headers["Authorization"] = ( + f"Bearer {create_jwt(JWTParams(co=collection_id))}" + ) + + response = auth_client.post( + f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections/{collection_id}/documents", + json={"content": "Test document with collection JWT", "metadata": {}}, + ) + assert response.status_code == 200 + + +def test_get_document_with_auth(auth_client, sample_data) -> None: + test_app, test_user = sample_data + + # First create a collection and document + if auth_client.auth_type == "empty": + auth_client.headers["Authorization"] = ( + f"Bearer {create_jwt(JWTParams(ap=test_app.public_id, us=test_user.public_id))}" + ) + + create_collection_response = auth_client.post( + f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections", + json={"name": "test_collection_get_doc", "metadata": {}}, + ) + + if auth_client.auth_type not in ["admin", "empty"]: + assert create_collection_response.status_code == 401 + return + + assert create_collection_response.status_code == 200 + collection_id = create_collection_response.json()["id"] + + create_doc_response = auth_client.post( + f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections/{collection_id}/documents", + json={"content": "Test document for retrieval", "metadata": {}}, + ) + assert create_doc_response.status_code == 200 + document_id = create_doc_response.json()["id"] + + # Get document with app and user scoped JWT + response = auth_client.get( + f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections/{collection_id}/documents/{document_id}" + ) + assert response.status_code == 200 + + # Test with collection-scoped JWT + if auth_client.auth_type == "empty": + auth_client.headers["Authorization"] = ( + f"Bearer {create_jwt(JWTParams(co=collection_id))}" + ) + + response = auth_client.get( + f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections/{collection_id}/documents/{document_id}" + ) + assert response.status_code == 200 diff --git a/uv.lock b/uv.lock index c0001b04..230a1618 100644 --- a/uv.lock +++ b/uv.lock @@ -430,6 +430,7 @@ dependencies = [ { name = "openai" }, { name = "pgvector" }, { name = "psycopg", extra = ["binary"] }, + { name = "pyjwt" }, { name = "python-dotenv" }, { name = "rich" }, { name = "sentry-sdk", extra = ["anthropic", "fastapi", "sqlalchemy"] }, @@ -460,6 +461,7 @@ requires-dist = [ { name = "openai", specifier = ">=1.43.0" }, { name = "pgvector", specifier = ">=0.2.5" }, { name = "psycopg", extras = ["binary"], specifier = ">=3.1.19" }, + { name = "pyjwt", specifier = ">=2.10.0" }, { name = "python-dotenv", specifier = ">=1.0.0" }, { name = "rich", specifier = ">=13.7.1" }, { name = "sentry-sdk", extras = ["anthropic", "fastapi", "sqlalchemy"], specifier = ">=2.3.1" }, @@ -1119,6 +1121,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/f7/3f/01c8b82017c199075f8f788d0d906b9ffbbc5a47dc9918a945e13d5a2bda/pygments-2.18.0-py3-none-any.whl", hash = "sha256:b8e6aca0523f3ab76fee51799c488e38782ac06eafcf95e7ba832985c8e7b13a", size = 1205513 }, ] +[[package]] +name = "pyjwt" +version = "2.10.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/e7/46/bd74733ff231675599650d3e47f361794b22ef3e3770998dda30d3b63726/pyjwt-2.10.1.tar.gz", hash = "sha256:3cc5772eb20009233caf06e9d8a0577824723b44e6648ee0a2aedb6cf9381953", size = 87785 } +wheels = [ + { url = "https://files.pythonhosted.org/packages/61/ad/689f02752eeec26aed679477e80e632ef1b682313be70793d798c1d5fc8f/PyJWT-2.10.1-py3-none-any.whl", hash = "sha256:dcdd193e30abefd5debf142f9adfcdd2b58004e644f25406ffaebd50bd98dacb", size = 22997 }, +] + [[package]] name = "pytest" version = "8.3.4"