diff --git a/plugins/user_management/.gitignore b/plugins/user_management/.gitignore new file mode 100644 index 0000000000..2d47491edf --- /dev/null +++ b/plugins/user_management/.gitignore @@ -0,0 +1,4 @@ +__pycache__/ +*.pyc +*.pyo +execute_record.json diff --git a/plugins/user_management/LICENSE b/plugins/user_management/LICENSE new file mode 100644 index 0000000000..a76dba14f6 --- /dev/null +++ b/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. diff --git a/plugins/user_management/README.md b/plugins/user_management/README.md new file mode 100644 index 0000000000..618fb787c9 --- /dev/null +++ b/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 diff --git a/plugins/user_management/__init__.py b/plugins/user_management/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/plugins/user_management/api/__init__.py b/plugins/user_management/api/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/plugins/user_management/api/login.py b/plugins/user_management/api/login.py new file mode 100644 index 0000000000..8ad6b956bd --- /dev/null +++ b/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) diff --git a/plugins/user_management/api/usage.py b/plugins/user_management/api/usage.py new file mode 100644 index 0000000000..b0e1f36856 --- /dev/null +++ b/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) diff --git a/plugins/user_management/api/users.py b/plugins/user_management/api/users.py new file mode 100644 index 0000000000..889e2ddc32 --- /dev/null +++ b/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) diff --git a/plugins/user_management/configs/users.json b/plugins/user_management/configs/users.json new file mode 100644 index 0000000000..c6bffec6cb --- /dev/null +++ b/plugins/user_management/configs/users.json @@ -0,0 +1,4 @@ +[ + {"username": "admin", "password": "admin123", "role": "admin"}, + {"username": "kostas", "password": "kostas123", "role": "admin"} +] diff --git a/plugins/user_management/default_config.yaml b/plugins/user_management/default_config.yaml new file mode 100644 index 0000000000..66136dade3 --- /dev/null +++ b/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" diff --git a/plugins/user_management/execute.py b/plugins/user_management/execute.py new file mode 100644 index 0000000000..62f552717c --- /dev/null +++ b/plugins/user_management/execute.py @@ -0,0 +1,121 @@ +"""Execute script for user_management plugin. +Called from the Plugin UI. Installs dependencies and initializes the database. +Must be standalone — no framework imports. +""" +import subprocess +import sys +import os + + +def main(): + print("[user_management] Running execute.py...") + + # 1. Install Python dependencies + deps = ["psycopg2-binary", "bcrypt", "openpyxl"] + print(f"[user_management] Installing: {', '.join(deps)}") + try: + result = subprocess.run( + [sys.executable, "-m", "pip", "install", *deps, + "--quiet", "--disable-pip-version-check"], + text=True, capture_output=True, + ) + if result.returncode != 0: + print(f"[user_management] pip install failed: {result.stderr}") + return 1 + print("[user_management] Dependencies installed.") + except Exception as e: + print(f"[user_management] pip install failed: {e}") + return 1 + + # 2. Initialize database + try: + import psycopg2 + import bcrypt + except ImportError as e: + print(f"[user_management] Missing dependency after install: {e}") + return 1 + + # Resolve DB URL from env var or default + db_url = os.environ.get( + "USER_MGMT_DB_URL", + "postgresql://a0:a0@localhost:5432/a0_user_mgmt" + ) + print(f"[user_management] Connecting to: {db_url}") + + try: + conn = psycopg2.connect(db_url) + conn.autocommit = True + cur = conn.cursor() + + # Create tables + cur.execute(""" + CREATE TABLE IF NOT EXISTS um_users ( + id SERIAL PRIMARY KEY, + username VARCHAR(100) UNIQUE NOT NULL, + password_hash VARCHAR(255) NOT NULL, + role VARCHAR(20) DEFAULT 'user' CHECK (role IN ('admin', 'user')), + created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW() + ); + """) + print("[user_management] Table um_users: OK") + + cur.execute(""" + CREATE TABLE IF NOT EXISTS um_token_usage ( + id SERIAL PRIMARY KEY, + user_id INTEGER REFERENCES um_users(id) ON DELETE SET NULL, + context_id VARCHAR(100) NOT NULL, + model VARCHAR(200) NOT NULL, + input_tokens INTEGER DEFAULT 0, + output_tokens INTEGER DEFAULT 0, + timestamp TIMESTAMP WITH TIME ZONE DEFAULT NOW() + ); + """) + print("[user_management] Table um_token_usage: OK") + + cur.execute(""" + CREATE TABLE IF NOT EXISTS um_context_ownership ( + context_id VARCHAR(100) PRIMARY KEY, + user_id INTEGER REFERENCES um_users(id) ON DELETE CASCADE, + created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW() + ); + """) + print("[user_management] Table um_context_ownership: OK") + + # Create indexes + cur.execute("CREATE INDEX IF NOT EXISTS idx_um_token_usage_user ON um_token_usage(user_id);") + cur.execute("CREATE INDEX IF NOT EXISTS idx_um_token_usage_ts ON um_token_usage(timestamp);") + cur.execute("CREATE INDEX IF NOT EXISTS idx_um_token_usage_ctx ON um_token_usage(context_id);") + cur.execute("CREATE INDEX IF NOT EXISTS idx_um_ctx_ownership_user ON um_context_ownership(user_id);") + print("[user_management] Indexes: OK") + + # Create admin user if not exists + cur.execute("SELECT id FROM um_users WHERE username = 'admin'") + if cur.fetchone() is None: + admin_pw = os.environ.get("USER_MGMT_ADMIN_PASSWORD", "admin123") + pw_hash = bcrypt.hashpw(admin_pw.encode(), bcrypt.gensalt()).decode() + cur.execute( + "INSERT INTO um_users (username, password_hash, role) VALUES (%s, %s, %s)", + ('admin', pw_hash, 'admin') + ) + print(f"[user_management] Admin user created (admin / {admin_pw})") + else: + print("[user_management] Admin user already exists.") + + cur.close() + conn.close() + print("[user_management] Database initialized successfully.") + + except Exception as e: + print(f"[user_management] DB init failed: {e}") + print( + "[user_management] Make sure PostgreSQL is running and " + "USER_MGMT_DB_URL env var or plugin config db_url is correct." + ) + return 1 + + print("[user_management] Setup complete!") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/plugins/user_management/extensions/python/_functions/agent/AgentContext/__init__/end/_50_tag_user.py b/plugins/user_management/extensions/python/_functions/agent/AgentContext/__init__/end/_50_tag_user.py new file mode 100644 index 0000000000..1dc1a6b63e --- /dev/null +++ b/plugins/user_management/extensions/python/_functions/agent/AgentContext/__init__/end/_50_tag_user.py @@ -0,0 +1,101 @@ +"""Backup tagger: try to tag context at creation time. + +This fires on AgentContext.__init__ end. Flask session is usually NOT +available here (contexts are created from WS events or startup loading), +so this is a best-effort backup. The primary tagging happens in +communicate/start/_50_tag_context.py. + +Also checks DB for previously recorded ownership (handles contexts +loaded from disk that were tagged in a previous session). +""" +from __future__ import annotations + +from helpers.extension import Extension +from helpers.print_style import PrintStyle + + +class TagUser(Extension): + + def execute(self, **kwargs): + try: + data = kwargs.get("data", {}) + args = data.get("args", ()) + if not args: + return + + context = args[0] # AgentContext being initialized + + # Already tagged via persisted data? Great, just ensure DB ownership. + if context.data.get("um_user_id"): + self._ensure_db_ownership(context) + return + + # Strategy 1: Try Flask session (works if created from HTTP context) + user_id, username = self._try_flask_session() + if user_id: + context.data["um_user_id"] = user_id + context.data["um_username"] = username + self._ensure_db_ownership(context) + PrintStyle.debug( + f"[user_management] tag_user __init__: tagged {context.id[:8]}... " + f"from session (user={username})" + ) + return + + # Strategy 2: Check DB for existing ownership + user_id, username = self._try_db_lookup(context.id) + if user_id: + context.data["um_user_id"] = user_id + context.data["um_username"] = username + PrintStyle.debug( + f"[user_management] tag_user __init__: tagged {context.id[:8]}... " + f"from DB (user={username})" + ) + return + + # No user info available — will be tagged at communicate() time + + except Exception as e: + PrintStyle.debug(f"[user_management] tag_user: {e}") + + @staticmethod + def _try_flask_session(): + try: + from flask import session as flask_session, has_request_context + if not has_request_context(): + return None, None + return flask_session.get("um_user_id"), flask_session.get("um_username") + except Exception: + return None, None + + @staticmethod + def _try_db_lookup(context_id): + try: + from usr.plugins.user_management.helpers.db import execute_query + rows = execute_query( + "SELECT co.user_id, u.username " + "FROM um_context_ownership co " + "LEFT JOIN um_users u ON co.user_id = u.id " + "WHERE co.context_id = %s LIMIT 1", + (context_id,), + ) + if rows: + return rows[0]["user_id"], rows[0].get("username", "unknown") + except Exception: + pass + return None, None + + @staticmethod + def _ensure_db_ownership(context): + try: + user_id = context.data.get("um_user_id") + if not user_id: + return + from usr.plugins.user_management.helpers.db import execute_write + execute_write( + "INSERT INTO um_context_ownership (context_id, user_id) " + "VALUES (%s, %s) ON CONFLICT (context_id) DO NOTHING", + (context.id, user_id), + ) + except Exception: + pass diff --git a/plugins/user_management/extensions/python/_functions/agent/AgentContext/communicate/start/_50_tag_context.py b/plugins/user_management/extensions/python/_functions/agent/AgentContext/communicate/start/_50_tag_context.py new file mode 100644 index 0000000000..e7de8b34c4 --- /dev/null +++ b/plugins/user_management/extensions/python/_functions/agent/AgentContext/communicate/start/_50_tag_context.py @@ -0,0 +1,82 @@ +"""Tag AgentContext with user info at communicate() time. + +Primary tagger — Flask session IS available here. +""" +from __future__ import annotations +import datetime + +from helpers.extension import Extension +from helpers.print_style import PrintStyle + +DEBUG_LOG = "/a0/tmp/um_debug.log" + +def _dbg(msg): + try: + ts = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S") + line = f"[{ts}] [tag_context] {msg}\n" + PrintStyle.standard(line.strip()) + with open(DEBUG_LOG, "a") as f: + f.write(line) + except Exception: + pass + + +class TagContext(Extension): + + def execute(self, **kwargs): + try: + _dbg(">>> communicate/start hook FIRED") + + data = kwargs.get("data", {}) + args = data.get("args", ()) + if not args: + _dbg("no args in data") + return + + context = args[0] # AgentContext instance + _dbg(f"context id={context.id[:12]}... existing um_user_id={context.data.get('um_user_id', 'NOT SET')}") + + # Already tagged? Skip. + if context.data.get("um_user_id"): + _dbg("already tagged, skipping") + return + + # Try Flask session + try: + from flask import session as flask_session, has_request_context + has_ctx = has_request_context() + _dbg(f"has_request_context={has_ctx}") + if not has_ctx: + return + + user_id = flask_session.get("um_user_id") + username = flask_session.get("um_username") + _dbg(f"session um_user_id={user_id} um_username={username}") + _dbg(f"session keys={list(flask_session.keys())}") + + if not user_id: + _dbg("no um_user_id in session - NOT TAGGED") + return + except Exception as e: + _dbg(f"session read error: {e}") + return + + # Tag context data + context.data["um_user_id"] = user_id + context.data["um_username"] = username + _dbg(f"SUCCESS: tagged context with user_id={user_id} username={username}") + + # Persist ownership in DB + try: + from usr.plugins.user_management.helpers.db import execute_write + execute_write( + "INSERT INTO um_context_ownership (context_id, user_id) " + "VALUES (%s, %s) ON CONFLICT (context_id) DO NOTHING", + (context.id, user_id), + ) + _dbg("DB ownership recorded") + except Exception as e: + _dbg(f"DB write error: {e}") + + except Exception as e: + _dbg(f"EXCEPTION: {e}") diff --git a/plugins/user_management/extensions/python/_functions/helpers/ui_server/UiRouteHandlers/login_handler/start/_50_um_login.py b/plugins/user_management/extensions/python/_functions/helpers/ui_server/UiRouteHandlers/login_handler/start/_50_um_login.py new file mode 100644 index 0000000000..1cee34ff57 --- /dev/null +++ b/plugins/user_management/extensions/python/_functions/helpers/ui_server/UiRouteHandlers/login_handler/start/_50_um_login.py @@ -0,0 +1,83 @@ +"""Intercept A0 native login to authenticate against um_users table. + +On POST (login attempt): + - Verify credentials against um_users + - If valid: set both A0 session auth AND um_* session vars, redirect to main UI + - If invalid: show login page with error + +On GET: let the original handler render the login page normally. + +Fallback: if the extension errors (e.g. DB down), the original handler +proceeds and checks .env AUTH_LOGIN/AUTH_PASSWORD as emergency access. +""" +from __future__ import annotations + +import asyncio + +from helpers.extension import Extension +from helpers.print_style import PrintStyle + + +class UmLogin(Extension): + + async def execute(self, **kwargs): + try: + from flask import redirect, render_template_string, request, session, url_for + from helpers import files, login + + data = kwargs.get("data", {}) + + # Only intercept POST requests (login attempts) + if request.method != "POST": + return # Let original handler show the login page + + username = request.form.get("username", "").strip() + password = request.form.get("password", "").strip() + + if not username or not password: + await asyncio.sleep(1) + login_page = files.read_file("webui/login.html") + data["result"] = render_template_string( + login_page, error="Please enter username and password." + ) + return + + # Authenticate against um_users table + from usr.plugins.user_management.helpers.auth import ( + verify_user, + set_session_user, + ) + + user = verify_user(username, password) + if user: + # Set A0 native session auth so requires_auth decorator passes + cred_hash = login.get_credentials_hash() + if cred_hash: + session["authentication"] = cred_hash + else: + # Edge case: no .env credentials configured + session["authentication"] = True + + # Set plugin session vars so token tracking works automatically + set_session_user(user) + + PrintStyle.standard( + f"[user_management] User '{username}' logged in " + f"(id={user['id']}, role={user['role']})" + ) + + # Short-circuit original handler: redirect to main UI + data["result"] = redirect(url_for("serve_index")) + return + + # Authentication failed + await asyncio.sleep(1) + login_page = files.read_file("webui/login.html") + data["result"] = render_template_string( + login_page, error="Invalid credentials. Please try again." + ) + + except Exception as e: + # On error (e.g. DB unreachable), DON'T set data["result"] + # so the original handler proceeds as fallback (.env credentials) + PrintStyle.error(f"[user_management] login extension error: {e}") diff --git a/plugins/user_management/extensions/python/_functions/helpers/ui_server/UiRouteHandlers/logout_handler/end/_50_um_logout.py b/plugins/user_management/extensions/python/_functions/helpers/ui_server/UiRouteHandlers/logout_handler/end/_50_um_logout.py new file mode 100644 index 0000000000..fa8b6226d1 --- /dev/null +++ b/plugins/user_management/extensions/python/_functions/helpers/ui_server/UiRouteHandlers/logout_handler/end/_50_um_logout.py @@ -0,0 +1,18 @@ +"""Clear um_* session vars when a user logs out via the A0 native logout.""" +from __future__ import annotations + +from helpers.extension import Extension +from helpers.print_style import PrintStyle + + +class UmLogout(Extension): + + async def execute(self, **kwargs): + try: + from flask import session + username = session.get("um_username", "unknown") + for key in ("um_user_id", "um_username", "um_role"): + session.pop(key, None) + PrintStyle.standard(f"[user_management] User '{username}' logged out.") + except Exception as e: + PrintStyle.debug(f"[user_management] logout extension error: {e}") diff --git a/plugins/user_management/extensions/python/monologue_start/_50_ensure_user_tag.py b/plugins/user_management/extensions/python/monologue_start/_50_ensure_user_tag.py new file mode 100644 index 0000000000..80d4d92942 --- /dev/null +++ b/plugins/user_management/extensions/python/monologue_start/_50_ensure_user_tag.py @@ -0,0 +1,51 @@ +"""Belt-and-suspenders: ensure context is tagged before monologue. + +Runs at the start of each agent monologue loop. If context.data still +doesn't have um_user_id (e.g. communicate/start was missed), check +the DB for a recorded ownership. + +This runs inside DeferredTask (no Flask session), so it can only +recover from DB records — it cannot discover the user from session. +""" +from __future__ import annotations + +from helpers.extension import Extension +from helpers.print_style import PrintStyle +from agent import LoopData + + +class EnsureUserTag(Extension): + + async def execute(self, loop_data: LoopData = LoopData(), **kwargs): + try: + if not self.agent or self.agent.number != 0: + return # Only track top-level agent + + context = self.agent.context + + # Already tagged? + if context.data.get("um_user_id"): + return + + # Try DB lookup (only source available in async context) + try: + from usr.plugins.user_management.helpers.db import execute_query + rows = execute_query( + "SELECT co.user_id, u.username " + "FROM um_context_ownership co " + "LEFT JOIN um_users u ON co.user_id = u.id " + "WHERE co.context_id = %s LIMIT 1", + (context.id,), + ) + if rows: + context.data["um_user_id"] = rows[0]["user_id"] + context.data["um_username"] = rows[0].get("username", "unknown") + PrintStyle.debug( + f"[user_management] monologue_start: recovered tag for " + f"{context.id[:8]}... from DB" + ) + except Exception: + pass + + except Exception as e: + PrintStyle.debug(f"[user_management] ensure_user_tag: {e}") diff --git a/plugins/user_management/extensions/python/process_chain_end/_50_log_tokens.py b/plugins/user_management/extensions/python/process_chain_end/_50_log_tokens.py new file mode 100644 index 0000000000..2f64286bc6 --- /dev/null +++ b/plugins/user_management/extensions/python/process_chain_end/_50_log_tokens.py @@ -0,0 +1,149 @@ +"""Log token usage after each agent processing chain. +""" +from __future__ import annotations +import datetime + +from helpers.extension import Extension +from helpers.print_style import PrintStyle + +DEBUG_LOG = "/a0/tmp/um_debug.log" + +def _dbg(msg): + try: + ts = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S") + line = f"[{ts}] [log_tokens] {msg}\n" + PrintStyle.standard(line.strip()) + with open(DEBUG_LOG, "a") as f: + f.write(line) + except Exception: + pass + + +class LogTokens(Extension): + + async def execute(self, **kwargs): + try: + _dbg(">>> process_chain_end hook FIRED") + + if not self.agent: + _dbg("no self.agent") + return + + _dbg(f"agent.number={self.agent.number}") + if self.agent.number != 0: + _dbg("not agent0, skipping") + return + + context = self.agent.context + user_id = context.data.get("um_user_id") + username = context.data.get("um_username", "?") + _dbg(f"context.id={context.id[:12]}... um_user_id={user_id} um_username={username}") + _dbg(f"context.data keys={list(context.data.keys())[:15]}") + + if not user_id: + _dbg("NO um_user_id - CANNOT LOG TOKENS") + return + + # Get model info + model_name = "unknown" + try: + from helpers import settings as settings_module + s = settings_module.get_settings() + model_name = s.get("chat_model_name", "unknown") + except Exception as e: + _dbg(f"settings error: {e}") + + # Count tokens + input_tokens, output_tokens = self._count_exchange_tokens() + _dbg(f"token count: input={input_tokens} output={output_tokens} model={model_name}") + + if input_tokens <= 0 and output_tokens <= 0: + _dbg("no tokens to log") + return + + # Log to DB + from usr.plugins.user_management.helpers.token_logger import log_tokens + log_tokens( + user_id=user_id, + context_id=context.id, + model=model_name, + input_tokens=input_tokens, + output_tokens=output_tokens, + ) + _dbg(f"SUCCESS: logged {input_tokens}+{output_tokens} tokens for user {username} model={model_name}") + + except Exception as e: + _dbg(f"EXCEPTION: {e}") + import traceback + _dbg(traceback.format_exc()) + + def _count_exchange_tokens(self): + input_tokens = 0 + output_tokens = 0 + + try: + from helpers.tokens import approximate_tokens + + history = self.agent.history + if not history: + _dbg("no history") + return 0, 0 + + messages = [] + if hasattr(history, "messages"): + messages = history.messages + elif hasattr(history, "output"): + messages = history.output() + + _dbg(f"history has {len(messages)} messages") + + if not messages: + return 0, 0 + + # Count the most recent exchange + last_user_idx = -1 + for i in range(len(messages) - 1, -1, -1): + msg = messages[i] + is_ai = getattr(msg, "ai", None) + if is_ai is None: + msg_type = getattr(msg, "type", "") + is_ai = msg_type in ("ai", "assistant") + if not is_ai: + last_user_idx = i + break + + _dbg(f"last_user_idx={last_user_idx}") + + if last_user_idx < 0: + for msg in messages: + content = self._get_content(msg) + output_tokens += approximate_tokens(content) + return 0, output_tokens + + # Input = the user message + user_msg = messages[last_user_idx] + input_tokens = approximate_tokens(self._get_content(user_msg)) + + # Output = all AI messages after the user message + for msg in messages[last_user_idx + 1:]: + content = self._get_content(msg) + output_tokens += approximate_tokens(content) + + except Exception as e: + _dbg(f"token count error: {e}") + + return input_tokens, output_tokens + + @staticmethod + def _get_content(msg): + content = getattr(msg, "content", "") + if isinstance(content, str): + return content + if isinstance(content, dict): + parts = [] + if "response" in content: + parts.append(str(content["response"])) + if "reasoning" in content: + parts.append(str(content["reasoning"])) + return " ".join(parts) if parts else str(content) + return str(content) diff --git a/plugins/user_management/extensions/webui/_sidebar-quick-actions-main-start/um-admin-button.html b/plugins/user_management/extensions/webui/_sidebar-quick-actions-main-start/um-admin-button.html new file mode 100644 index 0000000000..f1783112c6 --- /dev/null +++ b/plugins/user_management/extensions/webui/_sidebar-quick-actions-main-start/um-admin-button.html @@ -0,0 +1,51 @@ +
+ +
diff --git a/plugins/user_management/extensions/webui/apply_snapshot_before/um-filter-contexts.js b/plugins/user_management/extensions/webui/apply_snapshot_before/um-filter-contexts.js new file mode 100644 index 0000000000..1906b92477 --- /dev/null +++ b/plugins/user_management/extensions/webui/apply_snapshot_before/um-filter-contexts.js @@ -0,0 +1,41 @@ +/** + * Before each snapshot is applied, filter contexts by ownership. + * + * IMPORTANT: If the um-store has not yet finished initializing, + * we hide ALL contexts to prevent non-admin users from briefly + * seeing other users' chats. Once init completes, + * _reapplyContextFilter() in um-store.js will show the correct ones. + */ +export default function umFilterContextsBeforeSnapshot(ctx) { + const store = globalThis.Alpine?.store?.("umUser"); + if (!store) return; + + const snapshot = ctx?.snapshot; + if (!snapshot || !Array.isArray(snapshot.contexts)) return; + + // If store has not finished initializing yet, hide all contexts. + // This prevents the flash of all chats before ownership is known. + // The correct contexts will appear after init + _reapplyContextFilter. + if (!store._initialized) { + snapshot.contexts = []; + return; + } + + // Not logged in via plugin (shouldn't happen with unified login) + if (!store.isLoggedIn) return; + + // Admins see everything + if (store.currentUser && store.currentUser.role === "admin") return; + + // Non-admin: filter to only owned contexts + if (store.ownedContextIds instanceof Set) { + const chatsStore = globalThis.Alpine?.store?.("chats"); + const selectedId = chatsStore?.selected; + + snapshot.contexts = snapshot.contexts.filter( + (ctx) => + store.ownedContextIds.has(ctx.id) || + (selectedId && ctx.id === selectedId) + ); + } +} diff --git a/plugins/user_management/extensions/webui/initFw_end/um-init.html b/plugins/user_management/extensions/webui/initFw_end/um-init.html new file mode 100644 index 0000000000..bee70ec4b2 --- /dev/null +++ b/plugins/user_management/extensions/webui/initFw_end/um-init.html @@ -0,0 +1,20 @@ + + + + + + +
+ +
+ + + + diff --git a/plugins/user_management/helpers/__init__.py b/plugins/user_management/helpers/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/plugins/user_management/helpers/auth.py b/plugins/user_management/helpers/auth.py new file mode 100644 index 0000000000..2f4c6acd1b --- /dev/null +++ b/plugins/user_management/helpers/auth.py @@ -0,0 +1,119 @@ +import bcrypt +from usr.plugins.user_management.helpers.db import execute_query, execute_write + + +def hash_password(password: str) -> str: + return bcrypt.hashpw(password.encode("utf-8"), bcrypt.gensalt()).decode("utf-8") + + +def verify_password(password: str, password_hash: str) -> bool: + return bcrypt.checkpw(password.encode("utf-8"), password_hash.encode("utf-8")) + + +def create_user(username: str, password: str, role: str = "user"): + pw_hash = hash_password(password) + return execute_write( + "INSERT INTO um_users (username, password_hash, role) " + "VALUES (%s, %s, %s) RETURNING id, username, role, created_at", + (username, pw_hash, role), + returning=True, + ) + + +def verify_user(username: str, password: str): + rows = execute_query( + "SELECT id, username, password_hash, role FROM um_users WHERE username = %s", + (username,), + ) + if not rows: + return None + user = rows[0] + if verify_password(password, user["password_hash"]): + return {"id": user["id"], "username": user["username"], "role": user["role"]} + return None + + +def get_user(user_id: int): + rows = execute_query( + "SELECT id, username, role, created_at FROM um_users WHERE id = %s", + (user_id,), + ) + return rows[0] if rows else None + + +def list_users(): + return execute_query( + "SELECT id, username, role, created_at FROM um_users ORDER BY id" + ) + + +def delete_user(user_id: int): + execute_write("DELETE FROM um_users WHERE id = %s", (user_id,)) + + +def update_user(user_id: int, username=None, password=None, role=None): + parts, params = [], [] + if username is not None: + parts.append("username = %s") + params.append(username) + if password is not None: + parts.append("password_hash = %s") + params.append(hash_password(password)) + if role is not None: + parts.append("role = %s") + params.append(role) + if not parts: + return get_user(user_id) + params.append(user_id) + execute_write( + "UPDATE um_users SET " + ", ".join(parts) + " WHERE id = %s", + params, + ) + return get_user(user_id) + + +def get_current_user_from_session(): + """Get current user from Flask session (safe for non-request contexts).""" + try: + from flask import session, has_request_context + if not has_request_context(): + return None + user_id = session.get("um_user_id") + if not user_id: + return None + return get_user(user_id) + except Exception: + return None + + +def set_session_user(user: dict): + from flask import session + session["um_user_id"] = user["id"] + session["um_username"] = user["username"] + session["um_role"] = user["role"] + + +def clear_session_user(): + from flask import session + for key in ("um_user_id", "um_username", "um_role"): + session.pop(key, None) + + +def ensure_admin_user(): + """Create the initial admin account if it does not exist.""" + admin_username = "admin" + admin_password = "admin123" + try: + from helpers.plugins import get_plugin_config + config = get_plugin_config("user_management") + if config: + admin_username = config.get("admin_username", admin_username) + admin_password = config.get("admin_default_password", admin_password) + except Exception: + pass + + existing = execute_query( + "SELECT id FROM um_users WHERE username = %s", (admin_username,) + ) + if not existing: + create_user(admin_username, admin_password, role="admin") diff --git a/plugins/user_management/helpers/db.py b/plugins/user_management/helpers/db.py new file mode 100644 index 0000000000..14f8c965b0 --- /dev/null +++ b/plugins/user_management/helpers/db.py @@ -0,0 +1,109 @@ +import os +import threading +from contextlib import contextmanager + +import psycopg2 +import psycopg2.pool +import psycopg2.extras + +_pool = None +_pool_lock = threading.Lock() + +SCHEMA_SQL = """ +CREATE TABLE IF NOT EXISTS um_users ( + id SERIAL PRIMARY KEY, + username VARCHAR(100) UNIQUE NOT NULL, + password_hash VARCHAR(255) NOT NULL, + role VARCHAR(20) DEFAULT 'user' CHECK (role IN ('admin', 'user')), + created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW() +); + +CREATE TABLE IF NOT EXISTS um_token_usage ( + id SERIAL PRIMARY KEY, + user_id INTEGER REFERENCES um_users(id) ON DELETE SET NULL, + context_id VARCHAR(100) NOT NULL, + model VARCHAR(200) NOT NULL, + input_tokens INTEGER DEFAULT 0, + output_tokens INTEGER DEFAULT 0, + timestamp TIMESTAMP WITH TIME ZONE DEFAULT NOW() +); + +CREATE TABLE IF NOT EXISTS um_context_ownership ( + context_id VARCHAR(100) PRIMARY KEY, + user_id INTEGER REFERENCES um_users(id) ON DELETE CASCADE, + created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW() +); + +CREATE INDEX IF NOT EXISTS idx_um_token_usage_user ON um_token_usage(user_id); +CREATE INDEX IF NOT EXISTS idx_um_token_usage_ts ON um_token_usage(timestamp); +CREATE INDEX IF NOT EXISTS idx_um_token_usage_ctx ON um_token_usage(context_id); +CREATE INDEX IF NOT EXISTS idx_um_ctx_ownership_user ON um_context_ownership(user_id); +""" + + +def _get_db_url(): + url = os.environ.get("USER_MGMT_DB_URL") + if url: + return url + try: + from helpers.plugins import get_plugin_config + config = get_plugin_config("user_management") + if config and config.get("db_url"): + return config["db_url"] + except Exception: + pass + return "postgresql://a0:a0@localhost:5432/a0_user_mgmt" + + +def _get_pool(): + global _pool + if _pool is not None: + return _pool + with _pool_lock: + if _pool is not None: + return _pool + _pool = psycopg2.pool.ThreadedConnectionPool( + minconn=1, + maxconn=10, + dsn=_get_db_url(), + ) + return _pool + + +@contextmanager +def get_conn(): + """Get a connection from the pool (context manager).""" + pool = _get_pool() + conn = pool.getconn() + try: + yield conn + finally: + pool.putconn(conn) + + +def execute_query(query, params=None): + """Execute a SELECT query, return list of dicts.""" + with get_conn() as conn: + with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur: + cur.execute(query, params) + return [dict(row) for row in cur.fetchall()] + + +def execute_write(query, params=None, returning=False): + """Execute INSERT/UPDATE/DELETE. Returns first row if RETURNING clause used.""" + with get_conn() as conn: + with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur: + cur.execute(query, params) + conn.commit() + if returning: + row = cur.fetchone() + return dict(row) if row else None + return None + + +def init_db(): + """Create tables and indexes.""" + with get_conn() as conn: + with conn.cursor() as cur: + cur.execute(SCHEMA_SQL) + conn.commit() diff --git a/plugins/user_management/helpers/token_logger.py b/plugins/user_management/helpers/token_logger.py new file mode 100644 index 0000000000..4920059adc --- /dev/null +++ b/plugins/user_management/helpers/token_logger.py @@ -0,0 +1,130 @@ +import io +from usr.plugins.user_management.helpers.db import execute_query, execute_write + + +def log_tokens(user_id, context_id, model, input_tokens=0, output_tokens=0): + """Insert a token usage record.""" + execute_write( + "INSERT INTO um_token_usage " + "(user_id, context_id, model, input_tokens, output_tokens) " + "VALUES (%s, %s, %s, %s, %s)", + (user_id, context_id, model, input_tokens, output_tokens), + ) + + +def get_usage( + user_id=None, from_date=None, to_date=None, + model=None, context_id=None, limit=1000, +): + """Query token usage with optional filters.""" + conds, params = [], [] + if user_id is not None: + conds.append("tu.user_id = %s"); params.append(user_id) + if from_date: + conds.append("tu.timestamp >= %s"); params.append(from_date) + if to_date: + conds.append("tu.timestamp <= %s"); params.append(to_date) + if model: + conds.append("tu.model = %s"); params.append(model) + if context_id: + conds.append("tu.context_id = %s"); params.append(context_id) + where = ("WHERE " + " AND ".join(conds)) if conds else "" + params.append(limit) + return execute_query( + f"""SELECT tu.id, tu.user_id, u.username, tu.context_id, tu.model, + tu.input_tokens, tu.output_tokens, tu.timestamp + FROM um_token_usage tu + LEFT JOIN um_users u ON tu.user_id = u.id + {where} + ORDER BY tu.timestamp DESC + LIMIT %s""", + params, + ) + + +def get_usage_summary(user_id=None, group_by="day", from_date=None, to_date=None): + """Aggregate token usage grouped by day, user, or model.""" + conds, params = [], [] + if user_id is not None: + conds.append("tu.user_id = %s"); params.append(user_id) + if from_date: + conds.append("tu.timestamp >= %s"); params.append(from_date) + if to_date: + conds.append("tu.timestamp <= %s"); params.append(to_date) + where = ("WHERE " + " AND ".join(conds)) if conds else "" + + if group_by == "user": + gcol, scol = "u.username", "u.username AS group_key" + elif group_by == "model": + gcol, scol = "tu.model", "tu.model AS group_key" + else: + gcol, scol = "DATE(tu.timestamp)", "DATE(tu.timestamp)::text AS group_key" + + return execute_query( + f"""SELECT {scol}, + COUNT(*) AS request_count, + COALESCE(SUM(tu.input_tokens), 0) AS total_input_tokens, + COALESCE(SUM(tu.output_tokens), 0) AS total_output_tokens, + COALESCE(SUM(tu.input_tokens + tu.output_tokens), 0) AS total_tokens + FROM um_token_usage tu + LEFT JOIN um_users u ON tu.user_id = u.id + {where} + GROUP BY {gcol} + ORDER BY {gcol}""", + params, + ) + + +def export_to_excel(user_id=None, from_date=None, to_date=None): + """Generate an Excel workbook with summary + detail sheets.""" + import openpyxl + from openpyxl.styles import Font, PatternFill + + data = get_usage(user_id=user_id, from_date=from_date, to_date=to_date, limit=50000) + summary = get_usage_summary( + user_id=user_id, from_date=from_date, to_date=to_date, group_by="day" + ) + + wb = openpyxl.Workbook() + hf = Font(bold=True, color="FFFFFF") + hfill = PatternFill(start_color="4472C4", end_color="4472C4", fill_type="solid") + + # --- Summary sheet --- + ws = wb.active + ws.title = "Summary" + for ci, h in enumerate( + ["Date", "Requests", "Input Tokens", "Output Tokens", "Total Tokens"], 1 + ): + c = ws.cell(row=1, column=ci, value=h) + c.font, c.fill = hf, hfill + for ri, row in enumerate(summary, 2): + ws.cell(row=ri, column=1, value=str(row["group_key"])) + ws.cell(row=ri, column=2, value=row["request_count"]) + ws.cell(row=ri, column=3, value=row["total_input_tokens"]) + ws.cell(row=ri, column=4, value=row["total_output_tokens"]) + ws.cell(row=ri, column=5, value=row["total_tokens"]) + for ci in range(1, 6): + ws.column_dimensions[chr(64 + ci)].width = 18 + + # --- Detail sheet --- + ws2 = wb.create_sheet("Details") + for ci, h in enumerate( + ["ID", "Username", "Context", "Model", "Input Tokens", "Output Tokens", "Timestamp"], 1 + ): + c = ws2.cell(row=1, column=ci, value=h) + c.font, c.fill = hf, hfill + for ri, row in enumerate(data, 2): + ws2.cell(row=ri, column=1, value=row["id"]) + ws2.cell(row=ri, column=2, value=row.get("username") or "N/A") + ws2.cell(row=ri, column=3, value=row["context_id"]) + ws2.cell(row=ri, column=4, value=row["model"]) + ws2.cell(row=ri, column=5, value=row["input_tokens"]) + ws2.cell(row=ri, column=6, value=row["output_tokens"]) + ws2.cell(row=ri, column=7, value=str(row["timestamp"])) + for ci in range(1, 8): + ws2.column_dimensions[chr(64 + ci)].width = 18 + + buf = io.BytesIO() + wb.save(buf) + buf.seek(0) + return buf.getvalue() diff --git a/plugins/user_management/hooks.py b/plugins/user_management/hooks.py new file mode 100644 index 0000000000..8f204d02e5 --- /dev/null +++ b/plugins/user_management/hooks.py @@ -0,0 +1,30 @@ +import subprocess +import sys +from helpers.print_style import PrintStyle + + +def install(): + """Called when plugin is installed or enabled.""" + PrintStyle.standard("[user_management] Installing dependencies...") + try: + subprocess.check_call([ + sys.executable, "-m", "pip", "install", + "psycopg2-binary", "bcrypt", "openpyxl", + "--quiet", "--disable-pip-version-check", + ]) + except Exception as e: + PrintStyle.error(f"[user_management] pip install failed: {e}") + return + + try: + from usr.plugins.user_management.helpers.db import init_db + from usr.plugins.user_management.helpers.auth import ensure_admin_user + init_db() + ensure_admin_user() + PrintStyle.standard("[user_management] Database initialized successfully.") + except Exception as e: + PrintStyle.error(f"[user_management] DB init failed: {e}") + PrintStyle.error( + "[user_management] Ensure PostgreSQL is running and " + "USER_MGMT_DB_URL env var or plugin config is correct." + ) diff --git a/plugins/user_management/plugin.yaml b/plugins/user_management/plugin.yaml new file mode 100644 index 0000000000..76e110efe4 --- /dev/null +++ b/plugins/user_management/plugin.yaml @@ -0,0 +1,8 @@ +name: user_management +title: User Management +description: Multi-user authentication, chat isolation, and token usage tracking with PostgreSQL. +version: 1.0.0 +settings_sections: + - external +per_project_config: false +per_agent_config: false diff --git a/plugins/user_management/scripts/manage_users.py b/plugins/user_management/scripts/manage_users.py new file mode 100644 index 0000000000..dc67a8cd24 --- /dev/null +++ b/plugins/user_management/scripts/manage_users.py @@ -0,0 +1,193 @@ +#!/usr/bin/env python3 +"""CLI tool for managing um_users. + +Usage: + python manage_users.py list + python manage_users.py create [--role admin|user] + python manage_users.py update [--password ] [--role admin|user] + python manage_users.py delete + python manage_users.py sync # Sync from config file +""" +import argparse +import json +import os +import sys + +# Ensure the A0 root is on the path +sys.path.insert(0, "/a0") + +from usr.plugins.user_management.helpers.db import init_db, execute_query +from usr.plugins.user_management.helpers.auth import ( + create_user, + update_user, + delete_user, + list_users, + hash_password, +) + +SYNC_FILE = os.environ.get( + "UM_USERS_SYNC_FILE", + "/a0/usr/plugins/user_management/configs/users.json", +) + + +def cmd_list(args): + users = list_users() + if not users: + print("No users found.") + return + print(f"{'ID':<5} {'Username':<20} {'Role':<10} {'Created'}") + print("-" * 60) + for u in users: + print(f"{u['id']:<5} {u['username']:<20} {u['role']:<10} {u.get('created_at', '')}") + + +def cmd_create(args): + role = args.role or "admin" + try: + result = create_user(args.username, args.password, role=role) + print(f"Created user '{args.username}' (role={role}, id={result['id']})") + except Exception as e: + if "duplicate key" in str(e).lower() or "unique" in str(e).lower(): + print(f"User '{args.username}' already exists. Use 'update' instead.") + else: + print(f"Error: {e}") + sys.exit(1) + + +def cmd_update(args): + # Find user by username + users = execute_query( + "SELECT id FROM um_users WHERE username = %s", (args.username,) + ) + if not users: + print(f"User '{args.username}' not found.") + sys.exit(1) + + uid = users[0]["id"] + kwargs = {} + if args.password: + kwargs["password"] = args.password + if args.role: + kwargs["role"] = args.role + + if not kwargs: + print("Nothing to update. Use --password and/or --role.") + return + + result = update_user(uid, **kwargs) + print(f"Updated user '{args.username}': {result}") + + +def cmd_delete(args): + users = execute_query( + "SELECT id FROM um_users WHERE username = %s", (args.username,) + ) + if not users: + print(f"User '{args.username}' not found.") + sys.exit(1) + + delete_user(users[0]["id"]) + print(f"Deleted user '{args.username}'.") + + +def cmd_sync(args): + """Sync users from a JSON config file. + + File format: + [ + {"username": "admin", "password": "secret", "role": "admin"}, + {"username": "kostas", "password": "secret", "role": "admin"} + ] + + - New users are created + - Existing users get password/role updated if specified + - Users NOT in the file are NOT deleted (safe sync) + """ + sync_file = args.file or SYNC_FILE + if not os.path.isfile(sync_file): + print(f"Sync file not found: {sync_file}") + print(f"Create it with the expected JSON array format.") + sys.exit(1) + + with open(sync_file) as f: + desired_users = json.load(f) + + if not isinstance(desired_users, list): + print("Sync file must contain a JSON array of user objects.") + sys.exit(1) + + existing = {u["username"]: u for u in list_users()} + + for entry in desired_users: + uname = entry.get("username", "").strip() + pw = entry.get("password", "").strip() + role = entry.get("role", "admin").strip() + + if not uname: + continue + + if uname in existing: + # Update if password or role changed + kwargs = {} + if pw: + kwargs["password"] = pw + if role and role != existing[uname]["role"]: + kwargs["role"] = role + if kwargs: + update_user(existing[uname]["id"], **kwargs) + print(f" Updated: {uname} (role={role})") + else: + print(f" Unchanged: {uname}") + else: + if not pw: + print(f" Skipped: {uname} (no password for new user)") + continue + create_user(uname, pw, role=role) + print(f" Created: {uname} (role={role})") + + print("Sync complete.") + + +def main(): + # Ensure DB tables exist + init_db() + + parser = argparse.ArgumentParser(description="User management CLI") + sub = parser.add_subparsers(dest="command") + + sub.add_parser("list", help="List all users") + + p_create = sub.add_parser("create", help="Create a user") + p_create.add_argument("username") + p_create.add_argument("password") + p_create.add_argument("--role", default="admin", choices=["admin", "user"]) + + p_update = sub.add_parser("update", help="Update a user") + p_update.add_argument("username") + p_update.add_argument("--password", default=None) + p_update.add_argument("--role", default=None, choices=["admin", "user"]) + + p_delete = sub.add_parser("delete", help="Delete a user") + p_delete.add_argument("username") + + p_sync = sub.add_parser("sync", help="Sync users from JSON config") + p_sync.add_argument("--file", default=None, help="Path to sync file") + + args = parser.parse_args() + if not args.command: + parser.print_help() + sys.exit(1) + + cmds = { + "list": cmd_list, + "create": cmd_create, + "update": cmd_update, + "delete": cmd_delete, + "sync": cmd_sync, + } + cmds[args.command](args) + + +if __name__ == "__main__": + main() diff --git a/plugins/user_management/webui/admin-panel.html b/plugins/user_management/webui/admin-panel.html new file mode 100644 index 0000000000..1a4c185a65 --- /dev/null +++ b/plugins/user_management/webui/admin-panel.html @@ -0,0 +1,211 @@ + + + User Management - Admin Panel + + + + +
+ +
+ + + + diff --git a/plugins/user_management/webui/admin-store.js b/plugins/user_management/webui/admin-store.js new file mode 100644 index 0000000000..afe96004e7 --- /dev/null +++ b/plugins/user_management/webui/admin-store.js @@ -0,0 +1,191 @@ +import { createStore } from "/js/AlpineStore.js"; +import { callJsonApi } from "/js/api.js"; +import { + toastFrontendError, + toastFrontendSuccess, +} from "/components/notifications/notification-store.js"; + +const model = { + // State + users: [], + usageData: [], + usageSummary: [], + currentUser: null, + isLoading: false, + activeTab: "users", + + // Filters + filterUserId: null, + filterGroupBy: "day", + filterFromDate: "", + filterToDate: "", + + // New user form + newUsername: "", + newPassword: "", + newRole: "user", + + async init() { + await this.checkStatus(); + await this.loadUsers(); + }, + + async checkStatus() { + try { + const res = await callJsonApi("/plugins/user_management/login", { + action: "status", + }); + if (res.ok && res.logged_in) { + this.currentUser = res.user; + } + } catch (e) { + console.error("[user_management] status check failed:", e); + } + }, + + async loadUsers() { + if (!this.currentUser || this.currentUser.role !== "admin") return; + try { + this.isLoading = true; + const res = await callJsonApi("/plugins/user_management/users", { + action: "list", + }); + if (res.ok) { + this.users = res.users; + } + } catch (e) { + toastFrontendError("Failed to load users"); + } finally { + this.isLoading = false; + } + }, + + async createUser() { + if (!this.newUsername || !this.newPassword) { + toastFrontendError("Username and password are required"); + return; + } + try { + this.isLoading = true; + const res = await callJsonApi("/plugins/user_management/users", { + action: "create", + username: this.newUsername, + password: this.newPassword, + role: this.newRole, + }); + if (res.ok) { + toastFrontendSuccess(`User "${this.newUsername}" created`); + this.newUsername = ""; + this.newPassword = ""; + this.newRole = "user"; + await this.loadUsers(); + } + } catch (e) { + toastFrontendError("Failed to create user: " + (e.message || e)); + } finally { + this.isLoading = false; + } + }, + + async deleteUser(userId, username) { + if (!confirm(`Delete user "${username}"? This cannot be undone.`)) return; + try { + this.isLoading = true; + const res = await callJsonApi("/plugins/user_management/users", { + action: "delete", + user_id: userId, + }); + if (res.ok) { + toastFrontendSuccess(`User "${username}" deleted`); + await this.loadUsers(); + } + } catch (e) { + toastFrontendError("Failed to delete user: " + (e.message || e)); + } finally { + this.isLoading = false; + } + }, + + async loadUsage() { + try { + this.isLoading = true; + const params = { + action: "query", + limit: 500, + }; + if (this.filterUserId) params.user_id = this.filterUserId; + if (this.filterFromDate) params.from_date = this.filterFromDate; + if (this.filterToDate) params.to_date = this.filterToDate; + + const res = await callJsonApi("/plugins/user_management/usage", params); + if (res.ok) { + this.usageData = res.data; + } + } catch (e) { + toastFrontendError("Failed to load usage data"); + } finally { + this.isLoading = false; + } + }, + + async loadUsageSummary() { + try { + this.isLoading = true; + const params = { + action: "summary", + group_by: this.filterGroupBy, + }; + if (this.filterUserId) params.user_id = this.filterUserId; + if (this.filterFromDate) params.from_date = this.filterFromDate; + if (this.filterToDate) params.to_date = this.filterToDate; + + const res = await callJsonApi("/plugins/user_management/usage", params); + if (res.ok) { + this.usageSummary = res.data; + } + } catch (e) { + toastFrontendError("Failed to load usage summary"); + } finally { + this.isLoading = false; + } + }, + + async exportExcel() { + try { + const params = new URLSearchParams({ action: "export" }); + if (this.filterUserId) params.append("user_id", this.filterUserId); + if (this.filterFromDate) params.append("from_date", this.filterFromDate); + if (this.filterToDate) params.append("to_date", this.filterToDate); + + const resp = await fetch("/api/plugins/user_management/usage", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ + action: "export", + user_id: this.filterUserId || undefined, + from_date: this.filterFromDate || undefined, + to_date: this.filterToDate || undefined, + }), + }); + + if (!resp.ok) throw new Error("Export failed"); + + const blob = await resp.blob(); + const url = URL.createObjectURL(blob); + const a = document.createElement("a"); + a.href = url; + a.download = "token_usage.xlsx"; + a.click(); + URL.revokeObjectURL(url); + toastFrontendSuccess("Excel exported"); + } catch (e) { + toastFrontendError("Export failed: " + (e.message || e)); + } + }, + + formatNumber(n) { + return (n || 0).toLocaleString(); + }, +}; + +export const store = createStore("umAdmin", model); diff --git a/plugins/user_management/webui/config.html b/plugins/user_management/webui/config.html new file mode 100644 index 0000000000..7afb813b59 --- /dev/null +++ b/plugins/user_management/webui/config.html @@ -0,0 +1,68 @@ + + + User Management + + + +
+ +
+ + diff --git a/plugins/user_management/webui/main.html b/plugins/user_management/webui/main.html new file mode 100644 index 0000000000..497099c299 --- /dev/null +++ b/plugins/user_management/webui/main.html @@ -0,0 +1,56 @@ + + + + + + User Management + + + +
+ + + + + + + + + + +
+ + diff --git a/plugins/user_management/webui/um-store.js b/plugins/user_management/webui/um-store.js new file mode 100644 index 0000000000..80d1150c4b --- /dev/null +++ b/plugins/user_management/webui/um-store.js @@ -0,0 +1,199 @@ +import { createStore } from "/js/AlpineStore.js"; +import { callJsonApi } from "/js/api.js"; +import { + toastFrontendError, + toastFrontendSuccess, +} from "/components/notifications/notification-store.js"; + +const model = { + currentUser: null, + isLoggedIn: false, + ownedContextIds: null, // null = not loaded, Set = loaded + _initialized: false, + _patched: false, + _chatHooked: false, + + async init() { + if (this._initialized) return; + this._initialized = true; + await this.checkLoginStatus(); + if (this.isLoggedIn) { + await this.loadOwnedContexts(); + this._patchApplyContexts(); + this._hookNewChat(); + // Force re-apply of current snapshot so filter catches existing contexts + this._reapplyContextFilter(); + } + }, + + async checkLoginStatus() { + try { + const res = await callJsonApi("/plugins/user_management/login", { + action: "status", + }); + if (res.ok && res.logged_in) { + this.currentUser = res.user; + this.isLoggedIn = true; + } else { + this.currentUser = null; + this.isLoggedIn = false; + } + } catch (e) { + // Plugin may not be active or DB not connected + console.debug("[user_management] status check:", e); + } + }, + + async login(username, password) { + try { + const res = await callJsonApi("/plugins/user_management/login", { + action: "login", + username, + password, + }); + if (res.ok) { + this.currentUser = res.user; + this.isLoggedIn = true; + toastFrontendSuccess(`Logged in as ${res.user.username}`); + await this.loadOwnedContexts(); + this._patchApplyContexts(); + this._hookNewChat(); + this._reapplyContextFilter(); + return true; + } + } catch (e) { + toastFrontendError("Login failed: " + (e.message || "Invalid credentials")); + } + return false; + }, + + async logout() { + // Only admins can logout via the UI + if (!this.isAdmin()) { + toastFrontendError("Logout is restricted to administrators."); + return; + } + try { + await callJsonApi("/plugins/user_management/login", { + action: "logout", + }); + } catch (e) { + // ignore + } + this.currentUser = null; + this.isLoggedIn = false; + this.ownedContextIds = null; + // Reload page to clear all chat data and show login overlay + window.location.reload(); + }, + + async loadOwnedContexts() { + // Load the context IDs this user owns for filtering + try { + const res = await callJsonApi("/plugins/user_management/login", { + action: "owned_contexts", + }); + if (res.ok) { + if (res.is_admin) { + // Admins see everything + this.ownedContextIds = null; + } else if (Array.isArray(res.context_ids)) { + this.ownedContextIds = new Set(res.context_ids); + } + } + } catch (e) { + console.debug("[user_management] failed to load owned contexts:", e); + this.ownedContextIds = null; + } + }, + + /** + * Force re-apply context filtering on already-loaded chat list. + * Fixes timing issue where first snapshot arrives before init() completes. + */ + _reapplyContextFilter() { + if (this.isAdmin()) return; // Admins see everything + if (!(this.ownedContextIds instanceof Set)) return; + + const chatsStore = globalThis.Alpine?.store?.("chats"); + if (!chatsStore || !chatsStore.contexts) return; + + const selectedId = chatsStore.selected; + const filtered = chatsStore.contexts.filter( + (ctx) => + this.ownedContextIds.has(ctx.id) || + (selectedId && ctx.id === selectedId) + ); + chatsStore.contexts = filtered; + }, + + // Patch the chats store to filter contexts by ownership + _patchApplyContexts() { + if (this._patched) return; + const chatsStore = globalThis.Alpine?.store?.("chats"); + if (!chatsStore) return; + + const original = chatsStore.applyContexts.bind(chatsStore); + const self = this; + + chatsStore.applyContexts = function (contextsList) { + let filtered = contextsList; + if (self.isLoggedIn && self.currentUser && self.currentUser.role !== "admin") { + if (self.ownedContextIds instanceof Set) { + // Always allow the currently selected context through + // so newly created chats remain visible immediately + const selectedId = chatsStore.selected; + filtered = contextsList.filter( + (ctx) => + self.ownedContextIds.has(ctx.id) || + (selectedId && ctx.id === selectedId) + ); + } + } + original(filtered); + }; + + this._patched = true; + }, + + /** + * Register a newly created context ID so the ownership Set is + * immediately up-to-date without waiting for a full DB refresh. + */ + registerOwnedContext(contextId) { + if (this.ownedContextIds instanceof Set && contextId) { + this.ownedContextIds.add(contextId); + } + }, + + /** + * Hook into the chats store newChat method to auto-register + * ownership of chats created while logged in. + */ + _hookNewChat() { + if (this._chatHooked) return; + const chatsStore = globalThis.Alpine?.store?.("chats"); + if (!chatsStore) return; + + const originalNewChat = chatsStore.newChat.bind(chatsStore); + const self = this; + + chatsStore.newChat = async function () { + await originalNewChat(); + // After newChat completes, the chats store selected ID + // is already set to the new context — register it + const newId = chatsStore.selected; + if (newId) { + self.registerOwnedContext(newId); + } + }; + + this._chatHooked = true; + }, + + isAdmin() { + return this.currentUser && this.currentUser.role === "admin"; + }, +}; + +export const store = createStore("umUser", model);