remdb 0.3.202__py3-none-any.whl → 0.3.245__py3-none-any.whl

rem/api/routers/auth.py

@@ -3,11 +3,12 @@ Authentication Router.
 
 Supports multiple authentication methods:
 1. Email (passwordless): POST /api/auth/email/send-code, POST /api/auth/email/verify
-2. OAuth (Google, Microsoft): GET /api/auth/{provider}/login, GET /api/auth/{provider}/callback
+2. Pre-approved codes: POST /api/auth/email/verify (with pre-approved code, no send-code needed)
+3. OAuth (Google, Microsoft): GET /api/auth/{provider}/login, GET /api/auth/{provider}/callback
 
 Endpoints:
 - POST /api/auth/email/send-code     - Send login code to email
-- POST /api/auth/email/verify        - Verify code and create session
+- POST /api/auth/email/verify        - Verify code and create session (supports pre-approved codes)
 - GET  /api/auth/{provider}/login    - Initiate OAuth flow
 - GET  /api/auth/{provider}/callback - OAuth callback
 - POST /api/auth/logout              - Clear session
@@ -15,9 +16,39 @@ Endpoints:
 
 Supported providers:
 - email: Passwordless email login
+- preapproved: Pre-approved codes (bypass email, set via AUTH__PREAPPROVED_CODES)
 - google: Google OAuth 2.0 / OIDC
 - microsoft: Microsoft Entra ID OIDC
 
+=============================================================================
+Pre-Approved Code Authentication
+=============================================================================
+
+Pre-approved codes allow login without email verification. Useful for:
+- Demo accounts
+- Testing
+- Beta access codes
+- Admin provisioning
+
+Configuration:
+    AUTH__PREAPPROVED_CODES=A12345,A67890,B11111,B22222
+
+Code prefixes:
+    A = Admin role (e.g., A12345, AADMIN1)
+    B = Normal user role (e.g., B11111, BUSER1)
+
+Flow:
+    1. User enters email + pre-approved code (no send-code step needed)
+    2. POST /api/auth/email/verify with email and code
+    3. System validates code against AUTH__PREAPPROVED_CODES
+    4. Creates user if not exists, sets role based on prefix
+    5. Returns JWT tokens (same as email auth)
+
+Example:
+    curl -X POST http://localhost:8000/api/auth/email/verify \
+      -H "Content-Type: application/json" \
+      -d '{"email": "admin@example.com", "code": "A12345"}'
+
 =============================================================================
 Email Authentication Access Control
 =============================================================================
@@ -101,6 +132,8 @@ from authlib.integrations.starlette_client import OAuth
 from pydantic import BaseModel, EmailStr
 from loguru import logger
 
+from .common import ErrorResponse
+
 from ...settings import settings
 from ...services.postgres.service import PostgresService
 from ...services.user_service import UserService
@@ -159,7 +192,14 @@ class EmailVerifyRequest(BaseModel):
     code: str
 
 
-@router.post("/email/send-code")
+@router.post(
+    "/email/send-code",
+    responses={
+        400: {"model": ErrorResponse, "description": "Invalid request or email rejected"},
+        500: {"model": ErrorResponse, "description": "Failed to send login code"},
+        501: {"model": ErrorResponse, "description": "Email auth or database not configured"},
+    },
+)
 async def send_email_code(request: Request, body: EmailSendCodeRequest):
     """
     Send a login code to an email address.
@@ -221,11 +261,24 @@ async def send_email_code(request: Request, body: EmailSendCodeRequest):
         await db.disconnect()
 
 
-@router.post("/email/verify")
+@router.post(
+    "/email/verify",
+    responses={
+        400: {"model": ErrorResponse, "description": "Invalid or expired code"},
+        500: {"model": ErrorResponse, "description": "Failed to verify login code"},
+        501: {"model": ErrorResponse, "description": "Email auth or database not configured"},
+    },
+)
 async def verify_email_code(request: Request, body: EmailVerifyRequest):
     """
     Verify login code and create session with JWT tokens.
 
+    Supports two authentication methods:
+    1. Pre-approved codes: Codes from AUTH__PREAPPROVED_CODES bypass email verification.
+       - A prefix = admin role, B prefix = normal user role
+       - Creates user if not exists, logs in directly
+    2. Email verification: Standard 6-digit code sent via email
+
     Args:
         request: FastAPI request
         body: EmailVerifyRequest with email and code
@@ -233,12 +286,6 @@ async def verify_email_code(request: Request, body: EmailVerifyRequest):
     Returns:
         Success status with user info and JWT tokens
     """
-    if not settings.email.is_configured:
-        raise HTTPException(
-            status_code=501,
-            detail="Email authentication is not configured"
-        )
-
     if not settings.postgres.enabled:
         raise HTTPException(
             status_code=501,
@@ -248,6 +295,79 @@ async def verify_email_code(request: Request, body: EmailVerifyRequest):
     db = PostgresService()
     try:
         await db.connect()
+        user_service = UserService(db)
+
+        # Check for pre-approved code first
+        preapproved = settings.auth.check_preapproved_code(body.code)
+        if preapproved:
+            logger.info(f"Pre-approved code login attempt for {body.email} (role: {preapproved['role']})")
+
+            # Get or create user with pre-approved role
+            user_id = email_to_user_id(body.email)
+            user_entity = await user_service.get_user_by_id(user_id)
+
+            if not user_entity:
+                # Create new user with role from pre-approved code
+                user_entity = await user_service.get_or_create_user(
+                    email=body.email,
+                    name=body.email.split("@")[0],
+                    tenant_id="default",
+                )
+                # Update role based on pre-approved code prefix
+                user_entity.role = preapproved["role"]
+                from ...services.postgres.repository import Repository
+                from ...models.entities.user import User
+                user_repo = Repository(User, "users", db=db)
+                await user_repo.upsert(user_entity)
+                logger.info(f"Created user {body.email} with role={preapproved['role']} via pre-approved code")
+            else:
+                # Update existing user's role if admin code used
+                if preapproved["role"] == "admin" and user_entity.role != "admin":
+                    user_entity.role = "admin"
+                    from ...services.postgres.repository import Repository
+                    from ...models.entities.user import User
+                    user_repo = Repository(User, "users", db=db)
+                    await user_repo.upsert(user_entity)
+                    logger.info(f"Upgraded user {body.email} to admin via pre-approved code")
+
+            # Build user dict for session/JWT
+            user_dict = {
+                "id": str(user_entity.id),
+                "email": body.email,
+                "email_verified": True,
+                "name": user_entity.name or body.email.split("@")[0],
+                "provider": "preapproved",
+                "tenant_id": user_entity.tenant_id or "default",
+                "tier": user_entity.tier.value if user_entity.tier else "free",
+                "role": user_entity.role or preapproved["role"],
+                "roles": [user_entity.role or preapproved["role"]],
+            }
+
+            # Generate JWT tokens
+            jwt_service = get_jwt_service()
+            tokens = jwt_service.create_tokens(user_dict)
+
+            # Store user in session
+            request.session["user"] = user_dict
+
+            logger.info(f"User authenticated via pre-approved code: {body.email} (role: {user_dict['role']})")
+
+            return {
+                "success": True,
+                "message": "Successfully authenticated with pre-approved code!",
+                "user": user_dict,
+                "access_token": tokens["access_token"],
+                "refresh_token": tokens["refresh_token"],
+                "token_type": tokens["token_type"],
+                "expires_in": tokens["expires_in"],
+            }
+
+        # Standard email verification flow
+        if not settings.email.is_configured:
+            raise HTTPException(
+                status_code=501,
+                detail="Email authentication is not configured"
+            )
 
         # Initialize email auth provider
         email_auth = EmailAuthProvider()
@@ -272,7 +392,6 @@ async def verify_email_code(request: Request, body: EmailVerifyRequest):
         )
 
         # Fetch actual user data from database to get role/tier
-        user_service = UserService(db)
         try:
             user_entity = await user_service.get_user_by_id(result.user_id)
             if user_entity:
@@ -319,7 +438,13 @@ async def verify_email_code(request: Request, body: EmailVerifyRequest):
 # =============================================================================
 
 
-@router.get("/{provider}/login")
+@router.get(
+    "/{provider}/login",
+    responses={
+        400: {"model": ErrorResponse, "description": "Unknown OAuth provider"},
+        501: {"model": ErrorResponse, "description": "Authentication is disabled"},
+    },
+)
 async def login(provider: str, request: Request):
     """
     Initiate OAuth flow with provider.
@@ -361,7 +486,13 @@ async def login(provider: str, request: Request):
     return await client.authorize_redirect(request, redirect_uri)
 
 
-@router.get("/{provider}/callback")
+@router.get(
+    "/{provider}/callback",
+    responses={
+        400: {"model": ErrorResponse, "description": "Authentication failed or unknown provider"},
+        501: {"model": ErrorResponse, "description": "Authentication is disabled"},
+    },
+)
 async def callback(provider: str, request: Request):
     """
     OAuth callback endpoint.
@@ -498,7 +629,12 @@ async def logout(request: Request):
     return {"message": "Logged out successfully"}
 
 
-@router.get("/me")
+@router.get(
+    "/me",
+    responses={
+        401: {"model": ErrorResponse, "description": "Not authenticated"},
+    },
+)
 async def me(request: Request):
     """
     Get current user information from session or JWT.
@@ -536,7 +672,12 @@ class TokenRefreshRequest(BaseModel):
     refresh_token: str
 
 
-@router.post("/token/refresh")
+@router.post(
+    "/token/refresh",
+    responses={
+        401: {"model": ErrorResponse, "description": "Invalid or expired refresh token"},
+    },
+)
 async def refresh_token(body: TokenRefreshRequest):
     """
     Refresh access token using refresh token.
@@ -601,7 +742,12 @@ async def refresh_token(body: TokenRefreshRequest):
     return result
 
 
-@router.post("/token/verify")
+@router.post(
+    "/token/verify",
+    responses={
+        401: {"model": ErrorResponse, "description": "Missing, invalid, or expired token"},
+    },
+)
 async def verify_token(request: Request):
     """
     Verify an access token is valid.
@@ -665,7 +811,12 @@ def verify_dev_token(token: str) -> bool:
     return token == expected
 
 
-@router.get("/dev/token")
+@router.get(
+    "/dev/token",
+    responses={
+        401: {"model": ErrorResponse, "description": "Dev tokens not available in production"},
+    },
+)
 async def get_dev_token(request: Request):
     """
     Get a development token for testing (non-production only).
@@ -701,7 +852,13 @@ async def get_dev_token(request: Request):
     }
 
 
-@router.get("/dev/mock-code/{email}")
+@router.get(
+    "/dev/mock-code/{email}",
+    responses={
+        401: {"model": ErrorResponse, "description": "Mock codes not available in production"},
+        404: {"model": ErrorResponse, "description": "No code found for email"},
+    },
+)
 async def get_mock_code(email: str, request: Request):
     """
     Get the mock login code for testing (non-production only).

rem/api/routers/chat/child_streaming.py

@@ -0,0 +1,394 @@
+"""
+Child Agent Event Handling.
+
+Handles events from child agents during multi-agent orchestration.
+
+Event Flow:
+```
+Parent Agent (Siggy)
+      │
+      ▼
+  ask_agent tool
+      │
+      ├──────────────────────────────────┐
+      ▼                                  │
+  Child Agent (intake_diverge)           │
+      │                                  │
+      ├── child_tool_start ──────────────┼──► Event Sink (Queue)
+      ├── child_content ─────────────────┤
+      └── child_tool_result ─────────────┘
+                                         │
+                                         ▼
+                            drain_child_events()
+                                         │
+                                         ├── SSE to client
+                                         └── DB persistence
+```
+
+IMPORTANT: When child_content is streamed, parent text output should be SKIPPED
+to prevent content duplication.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import uuid
+from typing import TYPE_CHECKING, Any, AsyncGenerator
+
+from loguru import logger
+
+from .streaming_utils import StreamingState, build_content_chunk
+from .sse_events import MetadataEvent, ToolCallEvent, format_sse_event
+from ....services.session import SessionMessageStore
+from ....settings import settings
+from ....utils.date_utils import to_iso, utc_now
+
+if TYPE_CHECKING:
+    from ....agentic.context import AgentContext
+
+
+async def handle_child_tool_start(
+    state: StreamingState,
+    child_agent: str,
+    tool_name: str,
+    arguments: dict | str | None,
+    session_id: str | None,
+    user_id: str | None,
+) -> AsyncGenerator[str, None]:
+    """
+    Handle child_tool_start event.
+
+    Actions:
+    1. Log the tool call
+    2. Emit SSE event
+    3. Save to database (with tool_arguments in metadata for consistency with parent)
+    """
+    full_tool_name = f"{child_agent}:{tool_name}"
+    tool_id = f"call_{uuid.uuid4().hex[:8]}"
+
+    # Normalize arguments - may come as JSON string from ToolCallPart.args
+    if isinstance(arguments, str):
+        try:
+            arguments = json.loads(arguments)
+        except json.JSONDecodeError:
+            arguments = None
+    elif not isinstance(arguments, dict):
+        arguments = None
+
+    # 1. LOG
+    logger.info(f"🔧 {full_tool_name}")
+
+    # 2. EMIT SSE
+    yield format_sse_event(ToolCallEvent(
+        tool_name=full_tool_name,
+        tool_id=tool_id,
+        status="started",
+        arguments=arguments,
+    ))
+
+    # 3. SAVE TO DB - content contains args as JSON (pydantic_messages.py parses it)
+    if session_id and settings.postgres.enabled:
+        try:
+            store = SessionMessageStore(
+                user_id=user_id or settings.test.effective_user_id
+            )
+            tool_msg = {
+                "role": "tool",
+                # Content is the tool call args as JSON - this is what the agent sees on reload
+                # and what pydantic_messages.py parses for ToolCallPart.args
+                "content": json.dumps(arguments) if arguments else "",
+                "timestamp": to_iso(utc_now()),
+                "tool_call_id": tool_id,
+                "tool_name": full_tool_name,
+            }
+            await store.store_session_messages(
+                session_id=session_id,
+                messages=[tool_msg],
+                user_id=user_id,
+                compress=False,
+            )
+        except Exception as e:
+            logger.warning(f"Failed to save child tool call: {e}")
+
+
+def handle_child_content(
+    state: StreamingState,
+    child_agent: str,
+    content: str,
+) -> str | None:
+    """
+    Handle child_content event.
+
+    CRITICAL: Sets state.child_content_streamed = True
+    This flag is used to skip parent text output and prevent duplication.
+
+    Returns:
+        SSE chunk or None if content is empty
+    """
+    if not content:
+        return None
+
+    # Track that child content was streamed
+    # Parent text output should be SKIPPED when this is True
+    state.child_content_streamed = True
+    state.responding_agent = child_agent
+
+    return build_content_chunk(state, content)
+
+
+async def handle_child_tool_result(
+    state: StreamingState,
+    child_agent: str,
+    result: Any,
+    message_id: str | None,
+    session_id: str | None,
+    agent_schema: str | None,
+) -> AsyncGenerator[str, None]:
+    """
+    Handle child_tool_result event.
+
+    Actions:
+    1. Log metadata if present
+    2. Emit metadata event if present
+    3. Emit tool completion event
+    """
+    # Check for metadata registration
+    if isinstance(result, dict) and result.get("_metadata_event"):
+        risk = result.get("risk_level", "")
+        conf = result.get("confidence", "")
+        logger.info(f"📊 {child_agent} metadata: risk={risk}, confidence={conf}")
+
+        # Update responding agent from child
+        if result.get("agent_schema"):
+            state.responding_agent = result.get("agent_schema")
+
+        # Build extra dict with risk fields
+        extra_data = {}
+        if risk:
+            extra_data["risk_level"] = risk
+
+        yield format_sse_event(MetadataEvent(
+            message_id=message_id,
+            session_id=session_id,
+            agent_schema=agent_schema,
+            responding_agent=state.responding_agent,
+            confidence=result.get("confidence"),
+            extra=extra_data if extra_data else None,
+        ))
+
+    # Emit tool completion
+    # Preserve full result dict if it contains an artifact (e.g. finalize_intake)
+    # This is needed for frontend to extract artifact URLs for download
+    if isinstance(result, dict) and result.get("artifact"):
+        result_for_sse = result  # Full dict with artifact
+    else:
+        result_for_sse = str(result)[:200] if result else None
+
+    yield format_sse_event(ToolCallEvent(
+        tool_name=f"{child_agent}:tool",
+        tool_id=f"call_{uuid.uuid4().hex[:8]}",
+        status="completed",
+        result=result_for_sse,
+    ))
+
+
+async def drain_child_events(
+    event_sink: asyncio.Queue,
+    state: StreamingState,
+    session_id: str | None = None,
+    user_id: str | None = None,
+    message_id: str | None = None,
+    agent_schema: str | None = None,
+) -> AsyncGenerator[str, None]:
+    """
+    Drain all pending child events from the event sink.
+
+    This is called during tool execution to process events
+    pushed by child agents via ask_agent.
+
+    IMPORTANT: When child_content events are processed, this sets
+    state.child_content_streamed = True. Callers should check this
+    flag and skip parent text output to prevent duplication.
+    """
+    while not event_sink.empty():
+        try:
+            child_event = event_sink.get_nowait()
+            async for chunk in process_child_event(
+                child_event, state, session_id, user_id, message_id, agent_schema
+            ):
+                yield chunk
+        except Exception as e:
+            logger.warning(f"Error processing child event: {e}")
+
+
+async def process_child_event(
+    child_event: dict,
+    state: StreamingState,
+    session_id: str | None = None,
+    user_id: str | None = None,
+    message_id: str | None = None,
+    agent_schema: str | None = None,
+) -> AsyncGenerator[str, None]:
+    """Process a single child event and yield SSE chunks."""
+    event_type = child_event.get("type", "")
+    child_agent = child_event.get("agent_name", "child")
+
+    if event_type == "child_tool_start":
+        async for chunk in handle_child_tool_start(
+            state=state,
+            child_agent=child_agent,
+            tool_name=child_event.get("tool_name", "tool"),
+            arguments=child_event.get("arguments"),
+            session_id=session_id,
+            user_id=user_id,
+        ):
+            yield chunk
+
+    elif event_type == "child_content":
+        chunk = handle_child_content(
+            state=state,
+            child_agent=child_agent,
+            content=child_event.get("content", ""),
+        )
+        if chunk:
+            yield chunk
+
+    elif event_type == "child_tool_result":
+        async for chunk in handle_child_tool_result(
+            state=state,
+            child_agent=child_agent,
+            result=child_event.get("result"),
+            message_id=message_id,
+            session_id=session_id,
+            agent_schema=agent_schema,
+        ):
+            yield chunk
+
+
+async def stream_with_child_events(
+    tools_stream,
+    child_event_sink: asyncio.Queue,
+    state: StreamingState,
+    session_id: str | None = None,
+    user_id: str | None = None,
+    message_id: str | None = None,
+    agent_schema: str | None = None,
+) -> AsyncGenerator[tuple[str, Any], None]:
+    """
+    Multiplex tool events with child events using asyncio.wait().
+
+    This is the key fix for child agent streaming - instead of draining
+    the queue synchronously during tool event iteration, we concurrently
+    listen to both sources and yield events as they arrive.
+
+    Yields:
+        Tuples of (event_type, event_data) where event_type is either
+        "tool" or "child", allowing the caller to handle each appropriately.
+    """
+    tool_iter = tools_stream.__aiter__()
+
+    # Create initial tasks
+    pending_tool: asyncio.Task | None = None
+    pending_child: asyncio.Task | None = None
+
+    try:
+        pending_tool = asyncio.create_task(tool_iter.__anext__())
+    except StopAsyncIteration:
+        # No tool events, just drain any remaining child events
+        while not child_event_sink.empty():
+            try:
+                child_event = child_event_sink.get_nowait()
+                yield ("child", child_event)
+            except asyncio.QueueEmpty:
+                break
+        return
+
+    # Start listening for child events with a short timeout
+    pending_child = asyncio.create_task(
+        _get_child_event_with_timeout(child_event_sink, timeout=0.05)
+    )
+
+    try:
+        while True:
+            # Wait for either source to produce an event
+            tasks = {t for t in [pending_tool, pending_child] if t is not None}
+            if not tasks:
+                break
+
+            done, _ = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
+
+            for task in done:
+                try:
+                    result = task.result()
+                except asyncio.TimeoutError:
+                    # Child queue timeout - restart listener
+                    if task is pending_child:
+                        pending_child = asyncio.create_task(
+                            _get_child_event_with_timeout(child_event_sink, timeout=0.05)
+                        )
+                    continue
+                except StopAsyncIteration:
+                    # Tool stream exhausted
+                    if task is pending_tool:
+                        pending_tool = None
+                        # Final drain of any remaining child events
+                        if pending_child:
+                            pending_child.cancel()
+                            try:
+                                await pending_child
+                            except asyncio.CancelledError:
+                                pass
+                        while not child_event_sink.empty():
+                            try:
+                                child_event = child_event_sink.get_nowait()
+                                yield ("child", child_event)
+                            except asyncio.QueueEmpty:
+                                break
+                        return
+                    continue
+
+                if task is pending_child and result is not None:
+                    # Got a child event
+                    yield ("child", result)
+                    # Restart child listener
+                    pending_child = asyncio.create_task(
+                        _get_child_event_with_timeout(child_event_sink, timeout=0.05)
+                    )
+                elif task is pending_tool:
+                    # Got a tool event
+                    yield ("tool", result)
+                    # Get next tool event
+                    try:
+                        pending_tool = asyncio.create_task(tool_iter.__anext__())
+                    except StopAsyncIteration:
+                        pending_tool = None
+                elif task is pending_child and result is None:
+                    # Timeout with no event - restart listener
+                    pending_child = asyncio.create_task(
+                        _get_child_event_with_timeout(child_event_sink, timeout=0.05)
+                    )
+    finally:
+        # Cleanup any pending tasks
+        for task in [pending_tool, pending_child]:
+            if task and not task.done():
+                task.cancel()
+                try:
+                    await task
+                except asyncio.CancelledError:
+                    pass
+
+
+async def _get_child_event_with_timeout(
+    queue: asyncio.Queue, timeout: float = 0.05
+) -> dict | None:
+    """
+    Get an event from the queue with a timeout.
+
+    Returns None on timeout (no event available).
+    This allows the multiplexer to check for tool events regularly.
+    """
+    try:
+        return await asyncio.wait_for(queue.get(), timeout=timeout)
+    except asyncio.TimeoutError:
+        return None