aiqtoolkit 1.2.0.dev0__py3-none-any.whl → 1.2.0rc1__py3-none-any.whl

aiq/front_ends/fastapi/fastapi_front_end_plugin_worker.py

@@ -16,11 +16,13 @@
 import asyncio
 import logging
 import os
+import time
 import typing
 from abc import ABC
 from abc import abstractmethod
+from collections.abc import Awaitable
+from collections.abc import Callable
 from contextlib import asynccontextmanager
-from functools import partial
 from pathlib import Path
 
 from fastapi import BackgroundTasks
@@ -28,10 +30,13 @@ from fastapi import Body
 from fastapi import FastAPI
 from fastapi import Request
 from fastapi import Response
+from fastapi import UploadFile
 from fastapi.exceptions import HTTPException
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import StreamingResponse
 from pydantic import BaseModel
+from pydantic import Field
+from starlette.websockets import WebSocket
 
 from aiq.builder.workflow_builder import WorkflowBuilder
 from aiq.data_models.api_server import AIQChatRequest
@@ -39,20 +44,28 @@ from aiq.data_models.api_server import AIQChatResponse
 from aiq.data_models.api_server import AIQChatResponseChunk
 from aiq.data_models.api_server import AIQResponseIntermediateStep
 from aiq.data_models.config import AIQConfig
+from aiq.data_models.object_store import KeyAlreadyExistsError
+from aiq.data_models.object_store import NoSuchKeyError
 from aiq.eval.config import EvaluationRunOutput
 from aiq.eval.evaluate import EvaluationRun
 from aiq.eval.evaluate import EvaluationRunConfig
+from aiq.front_ends.fastapi.auth_flow_handlers.http_flow_handler import HTTPAuthenticationFlowHandler
+from aiq.front_ends.fastapi.auth_flow_handlers.websocket_flow_handler import FlowState
+from aiq.front_ends.fastapi.auth_flow_handlers.websocket_flow_handler import WebSocketAuthenticationFlowHandler
+from aiq.front_ends.fastapi.fastapi_front_end_config import AIQAsyncGenerateResponse
+from aiq.front_ends.fastapi.fastapi_front_end_config import AIQAsyncGenerationStatusResponse
 from aiq.front_ends.fastapi.fastapi_front_end_config import AIQEvaluateRequest
 from aiq.front_ends.fastapi.fastapi_front_end_config import AIQEvaluateResponse
 from aiq.front_ends.fastapi.fastapi_front_end_config import AIQEvaluateStatusResponse
 from aiq.front_ends.fastapi.fastapi_front_end_config import FastApiFrontEndConfig
 from aiq.front_ends.fastapi.job_store import JobInfo
 from aiq.front_ends.fastapi.job_store import JobStore
+from aiq.front_ends.fastapi.message_handler import WebSocketMessageHandler
 from aiq.front_ends.fastapi.response_helpers import generate_single_response
 from aiq.front_ends.fastapi.response_helpers import generate_streaming_response_as_str
 from aiq.front_ends.fastapi.response_helpers import generate_streaming_response_full_as_str
 from aiq.front_ends.fastapi.step_adaptor import StepAdaptor
-from aiq.front_ends.fastapi.websocket import AIQWebSocket
+from aiq.object_store.models import ObjectStoreItem
 from aiq.runtime.session import AIQSessionManager
 
 logger = logging.getLogger(__name__)
@@ -68,13 +81,16 @@ class FastApiFrontEndPluginWorkerBase(ABC):
 
         self._front_end_config = config.general.front_end
 
+        self._cleanup_tasks: list[str] = []
+        self._cleanup_tasks_lock = asyncio.Lock()
+        self._http_flow_handler: HTTPAuthenticationFlowHandler | None = HTTPAuthenticationFlowHandler()
+
     @property
     def config(self) -> AIQConfig:
         return self._config
 
     @property
     def front_end_config(self) -> FastApiFrontEndConfig:
-
         return self._front_end_config
 
     def build_app(self) -> FastAPI:
@@ -92,17 +108,30 @@ class FastApiFrontEndPluginWorkerBase(ABC):
                 yield
 
                 # If a cleanup task is running, cancel it
-                cleanup_task = getattr(starting_app.state, "cleanup_task", None)
-                if cleanup_task:
-                    logger.info("Cancelling cleanup task")
-                    cleanup_task.cancel()
+                async with self._cleanup_tasks_lock:
+
+                    # Cancel all cleanup tasks
+                    for task_name in self._cleanup_tasks:
+                        cleanup_task: asyncio.Task | None = getattr(starting_app.state, task_name, None)
+                        if cleanup_task is not None:
+                            logger.info("Cancelling %s cleanup task", task_name)
+                            cleanup_task.cancel()
+                        else:
+                            logger.warning("No cleanup task found for %s", task_name)
+
+                    self._cleanup_tasks.clear()
 
             logger.debug("Closing AIQ Toolkit server from process %s", os.getpid())
 
         aiq_app = FastAPI(lifespan=lifespan)
 
+        # Configure app CORS.
         self.set_cors_config(aiq_app)
 
+        @aiq_app.middleware("http")
+        async def authentication_log_filter(request: Request, call_next: Callable[[Request], Awaitable[Response]]):
+            return await self._suppress_authentication_logs(request, call_next)
+
         return aiq_app
 
     def set_cors_config(self, aiq_app: FastAPI) -> None:
@@ -137,6 +166,26 @@ class FastApiFrontEndPluginWorkerBase(ABC):
             **cors_kwargs,
         )
 
+    async def _suppress_authentication_logs(self, request: Request,
+                                            call_next: Callable[[Request], Awaitable[Response]]) -> Response:
+        """
+        Intercepts authentication request and supreses logs that contain sensitive data.
+        """
+        from aiq.utils.log_utils import LogFilter
+
+        logs_to_suppress: list[str] = []
+
+        if (self.front_end_config.oauth2_callback_path):
+            logs_to_suppress.append(self.front_end_config.oauth2_callback_path)
+
+        logging.getLogger("uvicorn.access").addFilter(LogFilter(logs_to_suppress))
+        try:
+            response = await call_next(request)
+        finally:
+            logging.getLogger("uvicorn.access").removeFilter(LogFilter(logs_to_suppress))
+
+        return response
+
     @abstractmethod
     async def configure(self, app: FastAPI, builder: WorkflowBuilder):
         pass
@@ -153,6 +202,38 @@ class RouteInfo(BaseModel):
 
 class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
 
+    def __init__(self, config: AIQConfig):
+        super().__init__(config)
+
+        self._outstanding_flows: dict[str, FlowState] = {}
+        self._outstanding_flows_lock = asyncio.Lock()
+
+    @staticmethod
+    async def _periodic_cleanup(name: str, job_store: JobStore, sleep_time_sec: int = 300):
+        while True:
+            try:
+                job_store.cleanup_expired_jobs()
+                logger.debug("Expired %s jobs cleaned up", name)
+            except Exception as e:
+                logger.error("Error during %s job cleanup: %s", name, e)
+            await asyncio.sleep(sleep_time_sec)
+
+    async def create_cleanup_task(self, app: FastAPI, name: str, job_store: JobStore, sleep_time_sec: int = 300):
+        # Schedule periodic cleanup of expired jobs on first job creation
+        attr_name = f"{name}_cleanup_task"
+
+        # Cheap check, if it doesn't exist, we will need to re-check after we acquire the lock
+        if not hasattr(app.state, attr_name):
+            async with self._cleanup_tasks_lock:
+                if not hasattr(app.state, attr_name):
+                    logger.info("Starting %s periodic cleanup task", name)
+                    setattr(
+                        app.state,
+                        attr_name,
+                        asyncio.create_task(
+                            self._periodic_cleanup(name=name, job_store=job_store, sleep_time_sec=sleep_time_sec)))
+                    self._cleanup_tasks.append(attr_name)
+
     def get_step_adaptor(self) -> StepAdaptor:
 
         return StepAdaptor(self.front_end_config.step_adaptor)
@@ -168,6 +249,8 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
 
         await self.add_default_route(app, AIQSessionManager(builder.build()))
         await self.add_evaluate_route(app, AIQSessionManager(builder.build()))
+        await self.add_static_files_route(app, builder)
+        await self.add_authorization_route(app)
 
         for ep in self.front_end_config.endpoints:
 
@@ -198,21 +281,6 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
         # Don't run multiple evaluations at the same time
         evaluation_lock = asyncio.Lock()
 
-        async def periodic_cleanup(job_store: JobStore):
-            while True:
-                try:
-                    job_store.cleanup_expired_jobs()
-                    logger.debug("Expired jobs cleaned up")
-                except Exception as e:
-                    logger.error("Error during job cleanup: %s", str(e))
-                await asyncio.sleep(300)  # every 5 minutes
-
-        def create_cleanup_task():
-            # Schedule periodic cleanup of expired jobs on first job creation
-            if not hasattr(app.state, "cleanup_task"):
-                logger.info("Starting periodic cleanup task")
-                app.state.cleanup_task = asyncio.create_task(periodic_cleanup(job_store))
-
         async def run_evaluation(job_id: str, config_file: str, reps: int, session_manager: AIQSessionManager):
             """Background task to run the evaluation."""
             async with evaluation_lock:
@@ -250,7 +318,7 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
                         return AIQEvaluateResponse(job_id=job.job_id, status=job.status)
 
                 job_id = job_store.create_job(request.config_file, request.job_id, request.expiry_seconds)
-                create_cleanup_task()
+                await self.create_cleanup_task(app=app, name="async_evaluation", job_store=job_store)
                 background_tasks.add_task(run_evaluation, job_id, request.config_file, request.reps, session_manager)
 
                 return AIQEvaluateResponse(job_id=job_id, status="submitted")
@@ -276,7 +344,7 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
                 if not job:
                     logger.warning("Job %s not found", job_id)
                     raise HTTPException(status_code=404, detail=f"Job {job_id} not found")
-                logger.info(f"Found job {job_id} with status {job.status}")
+                logger.info("Found job %s with status %s", job_id, job.status)
                 return translate_job_to_response(job)
 
         async def get_last_job_status(http_request: Request) -> AIQEvaluateStatusResponse:
@@ -355,6 +423,100 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
                 responses={500: response_500},
             )
 
+    async def add_static_files_route(self, app: FastAPI, builder: WorkflowBuilder):
+
+        if not self.front_end_config.object_store:
+            logger.debug("No object store configured, skipping static files route")
+            return
+
+        object_store_client = await builder.get_object_store_client(self.front_end_config.object_store)
+
+        def sanitize_path(path: str) -> str:
+            sanitized_path = os.path.normpath(path.strip("/"))
+            if sanitized_path == ".":
+                raise HTTPException(status_code=400, detail="Invalid file path.")
+            filename = os.path.basename(sanitized_path)
+            if not filename:
+                raise HTTPException(status_code=400, detail="Filename cannot be empty.")
+            return sanitized_path
+
+        # Upload static files to the object store; if key is present, it will fail with 409 Conflict
+        async def add_static_file(file_path: str, file: UploadFile):
+            sanitized_file_path = sanitize_path(file_path)
+            file_data = await file.read()
+
+            try:
+                await object_store_client.put_object(sanitized_file_path,
+                                                     ObjectStoreItem(data=file_data, content_type=file.content_type))
+            except KeyAlreadyExistsError as e:
+                raise HTTPException(status_code=409, detail=str(e)) from e
+
+            return {"filename": sanitized_file_path}
+
+        # Upsert static files to the object store; if key is present, it will overwrite the file
+        async def upsert_static_file(file_path: str, file: UploadFile):
+            sanitized_file_path = sanitize_path(file_path)
+            file_data = await file.read()
+
+            await object_store_client.upsert_object(sanitized_file_path,
+                                                    ObjectStoreItem(data=file_data, content_type=file.content_type))
+
+            return {"filename": sanitized_file_path}
+
+        # Get static files from the object store
+        async def get_static_file(file_path: str):
+
+            try:
+                file_data = await object_store_client.get_object(file_path)
+            except NoSuchKeyError as e:
+                raise HTTPException(status_code=404, detail=str(e)) from e
+
+            filename = file_path.split("/")[-1]
+
+            async def reader():
+                yield file_data.data
+
+            return StreamingResponse(reader(),
+                                     media_type=file_data.content_type,
+                                     headers={"Content-Disposition": f"attachment; filename={filename}"})
+
+        async def delete_static_file(file_path: str):
+            try:
+                await object_store_client.delete_object(file_path)
+            except NoSuchKeyError as e:
+                raise HTTPException(status_code=404, detail=str(e)) from e
+
+            return Response(status_code=204)
+
+        # Add the static files route to the FastAPI app
+        app.add_api_route(
+            path="/static/{file_path:path}",
+            endpoint=add_static_file,
+            methods=["POST"],
+            description="Upload a static file to the object store",
+        )
+
+        app.add_api_route(
+            path="/static/{file_path:path}",
+            endpoint=upsert_static_file,
+            methods=["PUT"],
+            description="Upsert a static file to the object store",
+        )
+
+        app.add_api_route(
+            path="/static/{file_path:path}",
+            endpoint=get_static_file,
+            methods=["GET"],
+            description="Get a static file from the object store",
+        )
+
+        app.add_api_route(
+            path="/static/{file_path:path}",
+            endpoint=delete_static_file,
+            methods=["DELETE"],
+            description="Delete a static file from the object store",
+        )
+
     async def add_route(self,
                         app: FastAPI,
                         endpoint: FastApiFrontEndConfig.EndpointBase,
@@ -362,17 +524,32 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
 
         workflow = session_manager.workflow
 
-        if (endpoint.websocket_path):
-            app.add_websocket_route(endpoint.websocket_path,
-                                    partial(AIQWebSocket, session_manager, self.get_step_adaptor()))
-
         GenerateBodyType = workflow.input_schema  # pylint: disable=invalid-name
         GenerateStreamResponseType = workflow.streaming_output_schema  # pylint: disable=invalid-name
         GenerateSingleResponseType = workflow.single_output_schema  # pylint: disable=invalid-name
 
+        # Append job_id and expiry_seconds to the input schema, this effectively makes these reserved keywords
+        # Consider prefixing these with "aiq_" to avoid conflicts
+        class AIQAsyncGenerateRequest(GenerateBodyType):
+            job_id: str | None = Field(default=None, description="Unique identifier for the evaluation job")
+            sync_timeout: int = Field(
+                default=0,
+                ge=0,
+                le=300,
+                description="Attempt to perform the job synchronously up until `sync_timeout` sectonds, "
+                "if the job hasn't been completed by then a job_id will be returned with a status code of 202.")
+            expiry_seconds: int = Field(default=JobStore.DEFAULT_EXPIRY,
+                                        ge=JobStore.MIN_EXPIRY,
+                                        le=JobStore.MAX_EXPIRY,
+                                        description="Optional time (in seconds) before the job expires. "
+                                        "Clamped between 600 (10 min) and 86400 (24h).")
+
         # Ensure that the input is in the body. POD types are treated as query parameters
         if (not issubclass(GenerateBodyType, BaseModel)):
             GenerateBodyType = typing.Annotated[GenerateBodyType, Body()]
+        else:
+            logger.info("Expecting generate request payloads in the following format: %s",
+                        GenerateBodyType.model_fields)
 
         response_500 = {
             "description": "Internal Server Error",
@@ -385,13 +562,20 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
             },
         }
 
+        # Create job store for tracking async generation jobs
+        job_store = JobStore()
+
+        # Run up to max_running_async_jobs jobs at the same time
+        async_job_concurrency = asyncio.Semaphore(self._front_end_config.max_running_async_jobs)
+
         def get_single_endpoint(result_type: type | None):
 
             async def get_single(response: Response, request: Request):
 
                 response.headers["Content-Type"] = "application/json"
 
-                async with session_manager.session(request=request):
+                async with session_manager.session(request=request,
+                                                   user_authentication_callback=self._http_flow_handler.authenticate):
 
                     return await generate_single_response(None, session_manager, result_type=result_type)
 
@@ -401,7 +585,8 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
 
             async def get_stream(request: Request):
 
-                async with session_manager.session(request=request):
+                async with session_manager.session(request=request,
+                                                   user_authentication_callback=self._http_flow_handler.authenticate):
 
                     return StreamingResponse(headers={"Content-Type": "text/event-stream; charset=utf-8"},
                                              content=generate_streaming_response_as_str(
@@ -435,7 +620,8 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
 
                 response.headers["Content-Type"] = "application/json"
 
-                async with session_manager.session(request=request):
+                async with session_manager.session(request=request,
+                                                   user_authentication_callback=self._http_flow_handler.authenticate):
 
                     return await generate_single_response(payload, session_manager, result_type=result_type)
 
@@ -448,7 +634,8 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
 
             async def post_stream(request: Request, payload: request_type):
 
-                async with session_manager.session(request=request):
+                async with session_manager.session(request=request,
+                                                   user_authentication_callback=self._http_flow_handler.authenticate):
 
                     return StreamingResponse(headers={"Content-Type": "text/event-stream; charset=utf-8"},
                                              content=generate_streaming_response_as_str(
@@ -482,7 +669,206 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
 
             return post_stream
 
+        def post_openai_api_compatible_endpoint(request_type: type):
+            """
+            OpenAI-compatible endpoint that handles both streaming and non-streaming
+            based on the 'stream' parameter in the request.
+            """
+
+            async def post_openai_api_compatible(response: Response, request: Request, payload: request_type):
+                # Check if streaming is requested
+                stream_requested = getattr(payload, 'stream', False)
+
+                async with session_manager.session(request=request):
+                    if stream_requested:
+                        # Return streaming response
+                        return StreamingResponse(headers={"Content-Type": "text/event-stream; charset=utf-8"},
+                                                 content=generate_streaming_response_as_str(
+                                                     payload,
+                                                     session_manager=session_manager,
+                                                     streaming=True,
+                                                     step_adaptor=self.get_step_adaptor(),
+                                                     result_type=AIQChatResponseChunk,
+                                                     output_type=AIQChatResponseChunk))
+                    else:
+                        # Return single response - check if workflow supports non-streaming
+                        try:
+                            response.headers["Content-Type"] = "application/json"
+                            return await generate_single_response(payload, session_manager, result_type=AIQChatResponse)
+                        except ValueError as e:
+                            if "Cannot get a single output value for streaming workflows" in str(e):
+                                # Workflow only supports streaming, but client requested non-streaming
+                                # Fall back to streaming and collect the result
+                                chunks = []
+                                async for chunk_str in generate_streaming_response_as_str(
+                                        payload,
+                                        session_manager=session_manager,
+                                        streaming=True,
+                                        step_adaptor=self.get_step_adaptor(),
+                                        result_type=AIQChatResponseChunk,
+                                        output_type=AIQChatResponseChunk):
+                                    if chunk_str.startswith("data: ") and not chunk_str.startswith("data: [DONE]"):
+                                        chunk_data = chunk_str[6:].strip()  # Remove "data: " prefix
+                                        if chunk_data:
+                                            try:
+                                                chunk_json = AIQChatResponseChunk.model_validate_json(chunk_data)
+                                                if (chunk_json.choices and len(chunk_json.choices) > 0
+                                                        and chunk_json.choices[0].delta
+                                                        and chunk_json.choices[0].delta.content is not None):
+                                                    chunks.append(chunk_json.choices[0].delta.content)
+                                            except Exception:
+                                                continue
+
+                                # Create a single response from collected chunks
+                                content = "".join(chunks)
+                                single_response = AIQChatResponse.from_string(content)
+                                response.headers["Content-Type"] = "application/json"
+                                return single_response
+                            else:
+                                raise
+
+            return post_openai_api_compatible
+
+        async def run_generation(job_id: str,
+                                 payload: typing.Any,
+                                 session_manager: AIQSessionManager,
+                                 result_type: type):
+            """Background task to run the evaluation."""
+            async with async_job_concurrency:
+                try:
+                    result = await generate_single_response(payload=payload,
+                                                            session_manager=session_manager,
+                                                            result_type=result_type)
+                    job_store.update_status(job_id, "success", output=result)
+                except Exception as e:
+                    logger.error("Error in evaluation job %s: %s", job_id, e)
+                    job_store.update_status(job_id, "failure", error=str(e))
+
+        def _job_status_to_response(job: JobInfo) -> AIQAsyncGenerationStatusResponse:
+            job_output = job.output
+            if job_output is not None:
+                job_output = job_output.model_dump()
+            return AIQAsyncGenerationStatusResponse(job_id=job.job_id,
+                                                    status=job.status,
+                                                    error=job.error,
+                                                    output=job_output,
+                                                    created_at=job.created_at,
+                                                    updated_at=job.updated_at,
+                                                    expires_at=job_store.get_expires_at(job))
+
+        def post_async_generation(request_type: type, final_result_type: type):
+
+            async def start_async_generation(
+                    request: request_type, background_tasks: BackgroundTasks, response: Response,
+                    http_request: Request) -> AIQAsyncGenerateResponse | AIQAsyncGenerationStatusResponse:
+                """Handle async generation requests."""
+
+                async with session_manager.session(request=http_request):
+
+                    # if job_id is present and already exists return the job info
+                    if request.job_id:
+                        job = job_store.get_job(request.job_id)
+                        if job:
+                            return AIQAsyncGenerateResponse(job_id=job.job_id, status=job.status)
+
+                    job_id = job_store.create_job(job_id=request.job_id, expiry_seconds=request.expiry_seconds)
+                    await self.create_cleanup_task(app=app, name="async_generation", job_store=job_store)
+
+                    # The fastapi/starlette background tasks won't begin executing until after the response is sent
+                    # to the client, so we need to wrap the task in a function, alowing us to start the task now,
+                    # and allowing the background task function to await the results.
+                    task = asyncio.create_task(
+                        run_generation(job_id=job_id,
+                                       payload=request,
+                                       session_manager=session_manager,
+                                       result_type=final_result_type))
+
+                    async def wrapped_task(t: asyncio.Task):
+                        return await t
+
+                    background_tasks.add_task(wrapped_task, task)
+
+                    now = time.time()
+                    sync_timeout = now + request.sync_timeout
+                    while time.time() < sync_timeout:
+                        job = job_store.get_job(job_id)
+                        if job is not None and job.status not in job_store.ACTIVE_STATUS:
+                            # If the job is done, return the result
+                            response.status_code = 200
+                            return _job_status_to_response(job)
+
+                        # Sleep for a short time before checking again
+                        await asyncio.sleep(0.1)
+
+                    response.status_code = 202
+                    return AIQAsyncGenerateResponse(job_id=job_id, status="submitted")
+
+            return start_async_generation
+
+        async def get_async_job_status(job_id: str, http_request: Request) -> AIQAsyncGenerationStatusResponse:
+            """Get the status of an async job."""
+            logger.info("Getting status for job %s", job_id)
+
+            async with session_manager.session(request=http_request):
+
+                job = job_store.get_job(job_id)
+                if not job:
+                    logger.warning("Job %s not found", job_id)
+                    raise HTTPException(status_code=404, detail=f"Job {job_id} not found")
+
+                logger.info("Found job %s with status %s", job_id, job.status)
+                return _job_status_to_response(job)
+
+        async def websocket_endpoint(websocket: WebSocket):
+
+            # Universal cookie handling: works for both cross-origin and same-origin connections
+            session_id = websocket.query_params.get("session")
+            if session_id:
+                headers = list(websocket.scope.get("headers", []))
+                cookie_header = f"aiqtoolkit-session={session_id}"
+
+                # Check if the session cookie already exists to avoid duplicates
+                cookie_exists = False
+                existing_session_cookie = False
+
+                for i, (name, value) in enumerate(headers):
+                    if name == b"cookie":
+                        cookie_exists = True
+                        cookie_str = value.decode()
+
+                        # Check if aiqtoolkit-session already exists in cookies
+                        if "aiqtoolkit-session=" in cookie_str:
+                            existing_session_cookie = True
+                            logger.info("WebSocket: Session cookie already present in headers (same-origin)")
+                        else:
+                            # Append to existing cookie header (cross-origin case)
+                            headers[i] = (name, f"{cookie_str}; {cookie_header}".encode())
+                            logger.info("WebSocket: Added session cookie to existing cookie header: %s",
+                                        session_id[:10] + "...")
+                        break
+
+                # Add new cookie header only if no cookies exist and no session cookie found
+                if not cookie_exists and not existing_session_cookie:
+                    headers.append((b"cookie", cookie_header.encode()))
+                    logger.info("WebSocket: Added new session cookie header: %s", session_id[:10] + "...")
+
+                # Update the websocket scope with the modified headers
+                websocket.scope["headers"] = headers
+
+            async with WebSocketMessageHandler(websocket, session_manager, self.get_step_adaptor()) as handler:
+
+                flow_handler = WebSocketAuthenticationFlowHandler(self._add_flow, self._remove_flow, handler)
+
+                # Ugly hack to set the flow handler on the message handler. Both need eachother to be set.
+                handler.set_flow_handler(flow_handler)
+
+                await handler.run()
+
+        if (endpoint.websocket_path):
+            app.add_websocket_route(endpoint.websocket_path, websocket_endpoint)
+
         if (endpoint.path):
+
             if (endpoint.method == "GET"):
 
                 app.add_api_route(
@@ -554,9 +940,31 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
                     responses={500: response_500},
                 )
 
+                app.add_api_route(
+                    path=f"{endpoint.path}/async",
+                    endpoint=post_async_generation(request_type=AIQAsyncGenerateRequest,
+                                                   final_result_type=GenerateSingleResponseType),
+                    methods=[endpoint.method],
+                    response_model=AIQAsyncGenerateResponse | AIQAsyncGenerationStatusResponse,
+                    description="Start an async generate job",
+                    responses={500: response_500},
+                )
             else:
                 raise ValueError(f"Unsupported method {endpoint.method}")
 
+            app.add_api_route(
+                path=f"{endpoint.path}/async/job/{{job_id}}",
+                endpoint=get_async_job_status,
+                methods=["GET"],
+                response_model=AIQAsyncGenerationStatusResponse,
+                description="Get the status of an async job",
+                responses={
+                    404: {
+                        "description": "Job not found"
+                    }, 500: response_500
+                },
+            )
+
         if (endpoint.openai_api_path):
             if (endpoint.method == "GET"):
 
@@ -582,26 +990,103 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
 
             elif (endpoint.method == "POST"):
 
-                app.add_api_route(
-                    path=endpoint.openai_api_path,
-                    endpoint=post_single_endpoint(request_type=AIQChatRequest, result_type=AIQChatResponse),
-                    methods=[endpoint.method],
-                    response_model=AIQChatResponse,
-                    description=endpoint.description,
-                    responses={500: response_500},
-                )
-
-                app.add_api_route(
-                    path=f"{endpoint.openai_api_path}/stream",
-                    endpoint=post_streaming_endpoint(request_type=AIQChatRequest,
-                                                     streaming=True,
-                                                     result_type=AIQChatResponseChunk,
-                                                     output_type=AIQChatResponseChunk),
-                    methods=[endpoint.method],
-                    response_model=AIQChatResponseChunk | AIQResponseIntermediateStep,
-                    description=endpoint.description,
-                    responses={500: response_500},
-                )
+                # Check if OpenAI v1 compatible endpoint is configured
+                openai_v1_path = getattr(endpoint, 'openai_api_v1_path', None)
+
+                # Always create legacy endpoints for backward compatibility (unless they conflict with v1 path)
+                if not openai_v1_path or openai_v1_path != endpoint.openai_api_path:
+                    # <openai_api_path> = non-streaming (legacy behavior)
+                    app.add_api_route(
+                        path=endpoint.openai_api_path,
+                        endpoint=post_single_endpoint(request_type=AIQChatRequest, result_type=AIQChatResponse),
+                        methods=[endpoint.method],
+                        response_model=AIQChatResponse,
+                        description=endpoint.description,
+                        responses={500: response_500},
+                    )
+
+                    # <openai_api_path>/stream = streaming (legacy behavior)
+                    app.add_api_route(
+                        path=f"{endpoint.openai_api_path}/stream",
+                        endpoint=post_streaming_endpoint(request_type=AIQChatRequest,
+                                                         streaming=True,
+                                                         result_type=AIQChatResponseChunk,
+                                                         output_type=AIQChatResponseChunk),
+                        methods=[endpoint.method],
+                        response_model=AIQChatResponseChunk | AIQResponseIntermediateStep,
+                        description=endpoint.description,
+                        responses={500: response_500},
+                    )
+
+                # Create OpenAI v1 compatible endpoint if configured
+                if openai_v1_path:
+                    # OpenAI v1 Compatible Mode: Create single endpoint that handles both streaming and non-streaming
+                    app.add_api_route(
+                        path=openai_v1_path,
+                        endpoint=post_openai_api_compatible_endpoint(request_type=AIQChatRequest),
+                        methods=[endpoint.method],
+                        response_model=AIQChatResponse | AIQChatResponseChunk,
+                        description=f"{endpoint.description} (OpenAI Chat Completions API compatible)",
+                        responses={500: response_500},
+                    )
 
             else:
                 raise ValueError(f"Unsupported method {endpoint.method}")
+
+    async def add_authorization_route(self, app: FastAPI):
+
+        from fastapi.responses import HTMLResponse
+
+        from aiq.front_ends.fastapi.html_snippets.auth_code_grant_success import AUTH_REDIRECT_SUCCESS_HTML
+
+        async def redirect_uri(request: Request):
+            """
+            Handle the redirect URI for OAuth2 authentication.
+            Args:
+                request: The FastAPI request object containing query parameters.
+
+            Returns:
+                HTMLResponse: A response indicating the success of the authentication flow.
+            """
+            state = request.query_params.get("state")
+
+            async with self._outstanding_flows_lock:
+                if not state or state not in self._outstanding_flows:
+                    return "Invalid state. Please restart the authentication process."
+
+                flow_state = self._outstanding_flows[state]
+
+            config = flow_state.config
+            verifier = flow_state.verifier
+            client = flow_state.client
+
+            try:
+                res = await client.fetch_token(url=config.token_url,
+                                               authorization_response=str(request.url),
+                                               code_verifier=verifier,
+                                               state=state)
+                flow_state.future.set_result(res)
+            except Exception as e:
+                flow_state.future.set_exception(e)
+
+            return HTMLResponse(content=AUTH_REDIRECT_SUCCESS_HTML,
+                                status_code=200,
+                                headers={
+                                    "Content-Type": "text/html; charset=utf-8", "Cache-Control": "no-cache"
+                                })
+
+        if (self.front_end_config.oauth2_callback_path):
+            # Add the redirect URI route
+            app.add_api_route(
+                path=self.front_end_config.oauth2_callback_path,
+                endpoint=redirect_uri,
+                methods=["GET"],
+                description="Handles the authorization code and state returned from the Authorization Code Grant Flow.")
+
+    async def _add_flow(self, state: str, flow_state: FlowState):
+        async with self._outstanding_flows_lock:
+            self._outstanding_flows[state] = flow_state
+
+    async def _remove_flow(self, state: str):
+        async with self._outstanding_flows_lock:
+            del self._outstanding_flows[state]