diff --git a/app/api/auth.py b/app/api/auth.py index cdd1bd5d..0a0e70d1 100644 --- a/app/api/auth.py +++ b/app/api/auth.py @@ -4,7 +4,7 @@ import os import time import uuid -from datetime import datetime +from datetime import datetime, timedelta, UTC import email_validator from typing import Optional, List, Union @@ -36,7 +36,7 @@ from app.core.worker import generate_pricing_url from app.db.database import get_db -from app.db.models import DBUser, DBAPIToken, DBRegion, DBTeam +from app.db.models import DBUser, DBAPIToken, DBRegion, DBTeam, DBAPITokenExpiryOption from app.services.litellm import LiteLLMService from app.services.dynamodb import DynamoDBService @@ -50,6 +50,7 @@ APIToken, APITokenCreate, APITokenResponse, + APITokenExpiryOption, UserUpdate, EmailValidation, LoginData, @@ -519,6 +520,20 @@ def generate_api_token() -> str: return secrets.token_urlsafe(32) +@router.get("/token/expiry-options", response_model=List[APITokenExpiryOption]) +async def list_expiry_options( + current_user=Depends(get_current_user_from_auth), + db: Session = Depends(get_db), +): + """List available API token expiry options""" + return ( + db.query(DBAPITokenExpiryOption) + .filter(DBAPITokenExpiryOption.is_active) + .order_by(DBAPITokenExpiryOption.id) + .all() + ) + + @router.post("/token", response_model=APIToken) async def create_token( token_create: APITokenCreate, @@ -552,8 +567,34 @@ async def create_token( # Create token for the current user user_id = current_user.id + # Fetch expiry option from DB (only active options are valid) + expiry_slug = token_create.expiry or "forever" + db_expiry_opt = ( + db.query(DBAPITokenExpiryOption) + .filter( + DBAPITokenExpiryOption.slug == expiry_slug, + DBAPITokenExpiryOption.is_active, + ) + .first() + ) + + if not db_expiry_opt: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail=f"Invalid or inactive expiry option: {expiry_slug}", + ) + + # Calculate expiration date + expires_at = None + if db_expiry_opt.days is not None: + expires_at = datetime.now(UTC) + timedelta(days=db_expiry_opt.days) + db_token = DBAPIToken( - name=token_create.name, token=generate_api_token(), user_id=user_id + name=token_create.name, + token=generate_api_token(), + user_id=user_id, + expires_at=expires_at, + expiry_option=db_expiry_opt.slug, ) db.add(db_token) db.commit() diff --git a/app/api/public.py b/app/api/public.py index 00c35666..4915b8de 100644 --- a/app/api/public.py +++ b/app/api/public.py @@ -119,7 +119,9 @@ def _to_display_name(model_id: str, aliases: list[str] | None = None) -> str: # Replace hyphenated number sequences with dotted versions before splitting. # Use word-boundary anchors so that e.g. "3-5" does not corrupt "123-5". modified_id = model_id - for hyphenated, dotted in sorted(dot_replacements.items(), key=lambda x: -len(x[0])): + for hyphenated, dotted in sorted( + dot_replacements.items(), key=lambda x: -len(x[0]) + ): modified_id = re.sub( r"(? None: + # Create api_token_expiry_options table + op.create_table( + "api_token_expiry_options", + sa.Column("id", sa.Integer(), nullable=False), + sa.Column("name", sa.String(), nullable=False), + sa.Column("slug", sa.String(), nullable=False), + sa.Column("days", sa.Integer(), nullable=True), + sa.Column("is_active", sa.Boolean(), nullable=False, server_default=sa.true()), + sa.PrimaryKeyConstraint("id"), + ) + op.create_index( + op.f("ix_api_token_expiry_options_id"), + "api_token_expiry_options", + ["id"], + unique=False, + ) + op.create_index( + op.f("ix_api_token_expiry_options_slug"), + "api_token_expiry_options", + ["slug"], + unique=True, + ) + + # Seed default expiry options so the API is usable immediately after migration + expiry_options_table = sa.table( + "api_token_expiry_options", + sa.column("name", sa.String()), + sa.column("slug", sa.String()), + sa.column("days", sa.Integer()), + sa.column("is_active", sa.Boolean()), + ) + op.bulk_insert( + expiry_options_table, + [ + {"name": "1 day", "slug": "1_day", "days": 1, "is_active": True}, + {"name": "1 week", "slug": "1_week", "days": 7, "is_active": True}, + {"name": "1 month", "slug": "1_month", "days": 30, "is_active": True}, + {"name": "2 months", "slug": "2_months", "days": 60, "is_active": True}, + {"name": "3 months", "slug": "3_months", "days": 90, "is_active": True}, + {"name": "4 months", "slug": "4_months", "days": 120, "is_active": True}, + {"name": "5 months", "slug": "5_months", "days": 150, "is_active": True}, + {"name": "6 months", "slug": "6_months", "days": 180, "is_active": True}, + {"name": "7 months", "slug": "7_months", "days": 210, "is_active": True}, + {"name": "8 months", "slug": "8_months", "days": 240, "is_active": True}, + {"name": "9 months", "slug": "9_months", "days": 270, "is_active": True}, + {"name": "10 months", "slug": "10_months", "days": 300, "is_active": True}, + {"name": "11 months", "slug": "11_months", "days": 330, "is_active": True}, + {"name": "1 year", "slug": "1_year", "days": 365, "is_active": True}, + {"name": "forever", "slug": "forever", "days": None, "is_active": True}, + ], + ) + + # Update api_tokens table - these might already exist in some environments but let's ensure they are there + # Check if columns exist first to be safe + conn = op.get_bind() + inspector = sa.inspect(conn) + columns = [c["name"] for c in inspector.get_columns("api_tokens")] + + if "expires_at" not in columns: + op.add_column( + "api_tokens", + sa.Column("expires_at", sa.DateTime(timezone=True), nullable=True), + ) + if "expiry_option" not in columns: + op.add_column( + "api_tokens", + sa.Column( + "expiry_option", sa.String(), nullable=False, server_default="forever" + ), + ) + + +def downgrade() -> None: + conn = op.get_bind() + inspector = sa.inspect(conn) + columns = [c["name"] for c in inspector.get_columns("api_tokens")] + + if "expiry_option" in columns: + op.drop_column("api_tokens", "expiry_option") + if "expires_at" in columns: + op.drop_column("api_tokens", "expires_at") + op.drop_index( + op.f("ix_api_token_expiry_options_slug"), table_name="api_token_expiry_options" + ) + op.drop_index( + op.f("ix_api_token_expiry_options_id"), table_name="api_token_expiry_options" + ) + op.drop_table("api_token_expiry_options") diff --git a/app/schemas/models.py b/app/schemas/models.py index e266b5ae..0e5456b6 100644 --- a/app/schemas/models.py +++ b/app/schemas/models.py @@ -68,12 +68,22 @@ class User(UserBase): audit_logs: ClassVar = relationship("AuditLog", back_populates="user") +class APITokenExpiryOption(BaseModel): + id: int + name: str + slug: str + days: Optional[int] = None + is_active: bool + model_config = ConfigDict(from_attributes=True) + + class APITokenBase(BaseModel): name: str class APITokenCreate(APITokenBase): user_id: Optional[int] = None + expiry: Optional[str] = "forever" class APIToken(APITokenBase): @@ -81,6 +91,8 @@ class APIToken(APITokenBase): token: str created_at: datetime last_used_at: Optional[datetime] = None + expires_at: Optional[datetime] = None + expiry_option: str user_id: int model_config = ConfigDict(from_attributes=True) @@ -89,6 +101,8 @@ class APITokenResponse(APITokenBase): id: int created_at: datetime last_used_at: Optional[datetime] = None + expires_at: Optional[datetime] = None + expiry_option: str user_id: int model_config = ConfigDict(from_attributes=True) diff --git a/frontend/src/app/auth/token/page.tsx b/frontend/src/app/auth/token/page.tsx index d8799564..2e5f38e0 100644 --- a/frontend/src/app/auth/token/page.tsx +++ b/frontend/src/app/auth/token/page.tsx @@ -1,7 +1,7 @@ "use client"; import { Loader2, X } from "lucide-react"; -import { useState, useEffect, useCallback } from "react"; +import { useState } from "react"; import { Alert, AlertDescription } from "@/components/ui/alert"; import { Button } from "@/components/ui/button"; import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card"; @@ -11,41 +11,66 @@ import { useToast } from "@/hooks/use-toast"; import { get, post, del } from "@/utils/api"; import { useQuery, useMutation, useQueryClient } from "@tanstack/react-query"; +import { + Select, + SelectContent, + SelectItem, + SelectTrigger, + SelectValue, +} from "@/components/ui/select"; + interface APIToken { id: string; name: string; token: string; created_at: string; last_used_at?: string; + expires_at?: string; + expiry_option: string; +} + +interface ExpiryOption { + id: number; + name: string; + slug: string; + days: number | null; } export default function APITokensPage() { const { toast } = useToast(); const queryClient = useQueryClient(); const [newTokenName, setNewTokenName] = useState(""); + const [selectedExpiry, setSelectedExpiry] = useState("forever"); const [showNewToken, setShowNewToken] = useState(null); - const [tokens, setTokens] = useState([]); - const { isLoading: queryLoading } = useQuery({ + const { data: expiryOptions } = useQuery({ + queryKey: ["expiry-options"], + queryFn: async () => { + const response = await get("/auth/token/expiry-options"); + return response.json() as Promise; + }, + }); + + const { isLoading: tokensLoading, data: tokens = [] } = useQuery({ queryKey: ["tokens"], queryFn: async () => { const response = await get("/auth/token"); const data = await response.json(); - return data; + return data as APIToken[]; }, }); const createMutation = useMutation({ - mutationFn: async (name: string) => { - const response = await post("/auth/token", { name }); + mutationFn: async (payload: { name: string; expiry: string }) => { + const response = await post("/auth/token", payload); const data = await response.json(); return data; }, onSuccess: (newToken) => { queryClient.invalidateQueries({ queryKey: ["tokens"] }); - queryClient.refetchQueries({ queryKey: ["tokens"], exact: true }); setShowNewToken(newToken); setNewTokenName(""); + setSelectedExpiry("forever"); toast({ title: "Success", description: "Token created successfully", @@ -66,7 +91,6 @@ export default function APITokensPage() { }, onSuccess: () => { queryClient.invalidateQueries({ queryKey: ["tokens"] }); - queryClient.refetchQueries({ queryKey: ["tokens"], exact: true }); toast({ title: "Success", description: "Token deleted successfully", @@ -81,50 +105,17 @@ export default function APITokensPage() { }, }); - const fetchTokens = useCallback(async () => { - try { - const response = await get("auth/token", { credentials: "include" }); - const data = await response.json(); - setTokens(data); - } catch (error) { - console.error("Error fetching tokens:", error); - toast({ - title: "Error", - description: "Failed to fetch tokens", - variant: "destructive", - }); - } - }, [toast, setTokens]); - const handleCreateToken = async (e: React.FormEvent) => { e.preventDefault(); if (!newTokenName.trim()) return; - createMutation.mutate(newTokenName); + createMutation.mutate({ name: newTokenName, expiry: selectedExpiry }); }; const handleDeleteToken = async (tokenId: string) => { - try { - await del("/auth/token/" + tokenId, { credentials: "include" }); - setTokens(tokens.filter((token) => token.id !== tokenId)); - toast({ - title: "Success", - description: "Token deleted successfully", - }); - } catch (error) { - console.error("Error deleting token:", error); - toast({ - title: "Error", - description: "Failed to delete token", - variant: "destructive", - }); - } + deleteMutation.mutate(tokenId); }; - useEffect(() => { - void fetchTokens(); - }, [fetchTokens]); - - if (queryLoading) { + if (tokensLoading) { return (
@@ -144,14 +135,42 @@ export default function APITokensPage() { Create New Token -
- setNewTokenName(e.target.value)} - placeholder="Token name" - className="max-w-sm" - /> + +
+ + setNewTokenName(e.target.value)} + placeholder="Token name" + /> +
+
+ + +