feat: add user_management plugin with session management and API key …

plugins/user_management/.gitignore

@@ -0,0 +1,4 @@
+__pycache__/
+*.pyc
+*.pyo
+execute_record.json

plugins/user_management/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 kgobakis
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

plugins/user_management/README.md

@@ -0,0 +1,32 @@
+# User Management Plugin for Agent Zero
+
+Multi-user authentication, chat isolation, and token usage tracking with PostgreSQL.
+
+## Features
+
+- **PostgreSQL Database** — Store users, sessions, and token usage
+- **Multi-user Auth** — Replace single login with per-user accounts
+- **Chat Isolation** — Tag each context with `user_id`, filter in UI
+- **Token Tracker** — Hook into LLM calls, log input/output tokens per user
+- **Admin Panel** — WebUI modal to view usage stats + export
+- **Session Manager** — Track active sessions per user with expiry
+- **API Key Store** — Generate & store per-user API keys in PostgreSQL for automation
+- **API Auth Middleware** — Validate incoming API key, resolve to user_id
+- **POST /api/chat Endpoint** — Accept external messages via API key for automations
+- **Rate Limiter** — Throttle API requests per key to prevent abuse
+
+## Installation
+
+Install via the Agent Zero Plugin Hub, or clone into your `usr/plugins/` directory:
+
+```bash
+git clone https://github.com/kgobakis/a0-user-management.git usr/plugins/user_management
+```
+
+## Configuration
+
+Configure PostgreSQL connection and user settings in the Agent Zero settings panel under **External** settings.
+
+## License
+
+MIT

plugins/user_management/api/login.py

@@ -0,0 +1,79 @@
+from helpers.api import ApiHandler, Input, Output, Request, Response
+from usr.plugins.user_management.helpers.auth import (
+    verify_user,
+    set_session_user,
+    clear_session_user,
+    get_current_user_from_session,
+)
+from usr.plugins.user_management.helpers.db import execute_query
+
+
+class Login(ApiHandler):
+    """Handle user login and logout."""
+
+    @classmethod
+    def requires_auth(cls) -> bool:
+        # Login endpoint must be accessible without existing auth
+        return False
+
+    @classmethod
+    def requires_csrf(cls) -> bool:
+        return False
+
+    async def process(self, input: Input, request: Request) -> Output:
+        action = str(input.get("action", "login")).strip().lower()
+
+        if action == "login":
+            username = str(input.get("username", "")).strip()
+            password = str(input.get("password", "")).strip()
+            if not username or not password:
+                return Response("Missing username or password", 400)
+
+            user = verify_user(username, password)
+            if not user:
+                return Response("Invalid credentials", 401)
+
+            set_session_user(user)
+            return {
+                "ok": True,
+                "user": {
+                    "id": user["id"],
+                    "username": user["username"],
+                    "role": user["role"],
+                },
+            }
+
+        elif action == "logout":
+            clear_session_user()
+            return {"ok": True}
+
+        elif action == "status":
+            user = get_current_user_from_session()
+            if user:
+                return {
+                    "ok": True,
+                    "logged_in": True,
+                    "user": {
+                        "id": user["id"],
+                        "username": user["username"],
+                        "role": user["role"],
+                    },
+                }
+            return {"ok": True, "logged_in": False, "user": None}
+
+        elif action == "owned_contexts":
+            from flask import session
+            user_id = session.get("um_user_id")
+            role = session.get("um_role")
+            if not user_id:
+                return {"ok": True, "context_ids": []}
+            # Admins see all - no filtering needed
+            if role == "admin":
+                return {"ok": True, "context_ids": [], "is_admin": True}
+            rows = execute_query(
+                "SELECT context_id FROM um_context_ownership WHERE user_id = %s",
+                (user_id,),
+            )
+            return {"ok": True, "context_ids": [r["context_id"] for r in rows]}
+
+        return Response("Unknown action", 400)

plugins/user_management/api/usage.py

@@ -0,0 +1,69 @@
+from helpers.api import ApiHandler, Input, Output, Request, Response
+from flask import session, send_file
+import io
+from usr.plugins.user_management.helpers.token_logger import (
+    get_usage,
+    get_usage_summary,
+    export_to_excel,
+)
+
+
+class Usage(ApiHandler):
+    """Token usage queries and Excel export."""
+
+    async def process(self, input: Input, request: Request) -> Output:
+        action = str(input.get("action", "query")).strip().lower()
+        role = session.get("um_role")
+        current_user_id = session.get("um_user_id")
+
+        # Non-admins can only see their own usage
+        user_id = input.get("user_id")
+        if role != "admin":
+            user_id = current_user_id
+
+        from_date = input.get("from_date")
+        to_date = input.get("to_date")
+
+        if action == "query":
+            data = get_usage(
+                user_id=user_id,
+                from_date=from_date,
+                to_date=to_date,
+                model=input.get("model"),
+                context_id=input.get("context_id"),
+                limit=int(input.get("limit", 500)),
+            )
+            # Serialize datetimes
+            for row in data:
+                if row.get("timestamp"):
+                    row["timestamp"] = str(row["timestamp"])
+            return {"ok": True, "data": data}
+
+        elif action == "summary":
+            data = get_usage_summary(
+                user_id=user_id,
+                group_by=str(input.get("group_by", "day")),
+                from_date=from_date,
+                to_date=to_date,
+            )
+            for row in data:
+                if row.get("group_key"):
+                    row["group_key"] = str(row["group_key"])
+            return {"ok": True, "data": data}
+
+        elif action == "export":
+            excel_bytes = export_to_excel(
+                user_id=user_id,
+                from_date=from_date,
+                to_date=to_date,
+            )
+            return Response(
+                response=excel_bytes,
+                status=200,
+                mimetype="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
+                headers={
+                    "Content-Disposition": "attachment; filename=token_usage.xlsx",
+                },
+            )
+
+        return Response("Unknown action", 400)

plugins/user_management/api/users.py

@@ -0,0 +1,84 @@
+from helpers.api import ApiHandler, Input, Output, Request, Response
+from flask import session
+from usr.plugins.user_management.helpers.auth import (
+    create_user,
+    list_users,
+    get_user,
+    delete_user,
+    update_user,
+)
+
+
+class Users(ApiHandler):
+    """CRUD operations for users (admin only)."""
+
+    async def process(self, input: Input, request: Request) -> Output:
+        # Admin check
+        role = session.get("um_role")
+        if role != "admin":
+            return Response("Admin access required", 403)
+
+        action = str(input.get("action", "")).strip().lower()
+
+        if action == "list":
+            users = list_users()
+            # Serialize datetimes
+            for u in users:
+                if u.get("created_at"):
+                    u["created_at"] = str(u["created_at"])
+            return {"ok": True, "users": users}
+
+        elif action == "get":
+            user_id = input.get("user_id")
+            if not user_id:
+                return Response("Missing user_id", 400)
+            user = get_user(int(user_id))
+            if not user:
+                return Response("User not found", 404)
+            if user.get("created_at"):
+                user["created_at"] = str(user["created_at"])
+            return {"ok": True, "user": user}
+
+        elif action == "create":
+            username = str(input.get("username", "")).strip()
+            password = str(input.get("password", "")).strip()
+            role_val = str(input.get("role", "user")).strip()
+            if not username or not password:
+                return Response("Missing username or password", 400)
+            if role_val not in ("admin", "user"):
+                return Response("Invalid role", 400)
+            try:
+                user = create_user(username, password, role_val)
+                if user and user.get("created_at"):
+                    user["created_at"] = str(user["created_at"])
+                return {"ok": True, "user": user}
+            except Exception as e:
+                if "duplicate" in str(e).lower() or "unique" in str(e).lower():
+                    return Response("Username already exists", 409)
+                raise
+
+        elif action == "update":
+            user_id = input.get("user_id")
+            if not user_id:
+                return Response("Missing user_id", 400)
+            user = update_user(
+                int(user_id),
+                username=input.get("username"),
+                password=input.get("password"),
+                role=input.get("role"),
+            )
+            if user and user.get("created_at"):
+                user["created_at"] = str(user["created_at"])
+            return {"ok": True, "user": user}
+
+        elif action == "delete":
+            user_id = input.get("user_id")
+            if not user_id:
+                return Response("Missing user_id", 400)
+            # Prevent self-deletion
+            if int(user_id) == session.get("um_user_id"):
+                return Response("Cannot delete your own account", 400)
+            delete_user(int(user_id))
+            return {"ok": True}
+
+        return Response("Unknown action", 400)

plugins/user_management/configs/users.json

@@ -0,0 +1,4 @@
+[
+    {"username": "admin", "password": "admin123", "role": "admin"},
+    {"username": "kostas", "password": "kostas123", "role": "admin"}
+]

plugins/user_management/default_config.yaml

@@ -0,0 +1,3 @@
+db_url: "postgresql://postgres:HiowQBnoeHWplBkJSivq8D9KFqtJuxooD1sPXr97S7LCqLPJS6aeFdqtHz6DHrJ9@ssi2b2w61m4qo6hc1tq5bi20:5432/postgres"
+admin_username: "admin"
+admin_default_password: "admin123"