From b0e3208fa3fe45fa0b6afd537b3d42060f2f88ce Mon Sep 17 00:00:00 2001
From: Abhimanyu Saharan <asaharan@onemindservices.com>
Date: Wed, 4 Feb 2026 14:58:20 +0530
Subject: [PATCH] feat: enhance agent management with session handling and UI
 improvements

---
 .../d3e4f5a6b7c8_add_agent_token_hash.py      | 38 ++++++++++++++
 backend/app/core/agent_auth.py                | 51 +++++++++++++++++++
 backend/app/core/agent_tokens.py              | 47 +++++++++++++++++
 3 files changed, 136 insertions(+)
 create mode 100644 backend/alembic/versions/d3e4f5a6b7c8_add_agent_token_hash.py
 create mode 100644 backend/app/core/agent_auth.py
 create mode 100644 backend/app/core/agent_tokens.py

diff --git a/backend/alembic/versions/d3e4f5a6b7c8_add_agent_token_hash.py b/backend/alembic/versions/d3e4f5a6b7c8_add_agent_token_hash.py
new file mode 100644
index 00000000..edc1eeea
--- /dev/null
+++ b/backend/alembic/versions/d3e4f5a6b7c8_add_agent_token_hash.py
@@ -0,0 +1,38 @@
+"""add agent token hash
+
+Revision ID: d3e4f5a6b7c8
+Revises: c7f0a2b1d4e3
+Create Date: 2026-02-04 06:50:00.000000
+
+"""
+
+from __future__ import annotations
+
+from alembic import op
+import sqlalchemy as sa
+import sqlmodel
+
+
+# revision identifiers, used by Alembic.
+revision = "d3e4f5a6b7c8"
+down_revision = "c7f0a2b1d4e3"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    op.add_column(
+        "agents",
+        sa.Column("agent_token_hash", sqlmodel.sql.sqltypes.AutoString(), nullable=True),
+    )
+    op.create_index(
+        op.f("ix_agents_agent_token_hash"),
+        "agents",
+        ["agent_token_hash"],
+        unique=False,
+    )
+
+
+def downgrade() -> None:
+    op.drop_index(op.f("ix_agents_agent_token_hash"), table_name="agents")
+    op.drop_column("agents", "agent_token_hash")
diff --git a/backend/app/core/agent_auth.py b/backend/app/core/agent_auth.py
new file mode 100644
index 00000000..ca05038b
--- /dev/null
+++ b/backend/app/core/agent_auth.py
@@ -0,0 +1,51 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Literal
+
+from fastapi import Depends, Header, HTTPException, status
+from sqlmodel import Session, col, select
+
+from app.core.agent_tokens import verify_agent_token
+from app.db.session import get_session
+from app.models.agents import Agent
+
+
+@dataclass
+class AgentAuthContext:
+    actor_type: Literal["agent"]
+    agent: Agent
+
+
+def _find_agent_for_token(session: Session, token: str) -> Agent | None:
+    agents = list(
+        session.exec(select(Agent).where(col(Agent.agent_token_hash).is_not(None)))
+    )
+    for agent in agents:
+        if agent.agent_token_hash and verify_agent_token(token, agent.agent_token_hash):
+            return agent
+    return None
+
+
+def get_agent_auth_context(
+    agent_token: str | None = Header(default=None, alias="X-Agent-Token"),
+    session: Session = Depends(get_session),
+) -> AgentAuthContext:
+    if not agent_token:
+        raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED)
+    agent = _find_agent_for_token(session, agent_token)
+    if agent is None:
+        raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED)
+    return AgentAuthContext(actor_type="agent", agent=agent)
+
+
+def get_agent_auth_context_optional(
+    agent_token: str | None = Header(default=None, alias="X-Agent-Token"),
+    session: Session = Depends(get_session),
+) -> AgentAuthContext | None:
+    if not agent_token:
+        return None
+    agent = _find_agent_for_token(session, agent_token)
+    if agent is None:
+        raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED)
+    return AgentAuthContext(actor_type="agent", agent=agent)
diff --git a/backend/app/core/agent_tokens.py b/backend/app/core/agent_tokens.py
new file mode 100644
index 00000000..13fe1a93
--- /dev/null
+++ b/backend/app/core/agent_tokens.py
@@ -0,0 +1,47 @@
+from __future__ import annotations
+
+import base64
+import hashlib
+import hmac
+import secrets
+
+ITERATIONS = 200_000
+SALT_BYTES = 16
+
+
+def generate_agent_token() -> str:
+    return secrets.token_urlsafe(32)
+
+
+def _b64encode(value: bytes) -> str:
+    return base64.urlsafe_b64encode(value).decode("utf-8").rstrip("=")
+
+
+def _b64decode(value: str) -> bytes:
+    padding = "=" * (-len(value) % 4)
+    return base64.urlsafe_b64decode(value + padding)
+
+
+def hash_agent_token(token: str) -> str:
+    salt = secrets.token_bytes(SALT_BYTES)
+    digest = hashlib.pbkdf2_hmac("sha256", token.encode("utf-8"), salt, ITERATIONS)
+    return f"pbkdf2_sha256${ITERATIONS}${_b64encode(salt)}${_b64encode(digest)}"
+
+
+def verify_agent_token(token: str, stored_hash: str) -> bool:
+    try:
+        algorithm, iterations, salt_b64, digest_b64 = stored_hash.split("$")
+    except ValueError:
+        return False
+    if algorithm != "pbkdf2_sha256":
+        return False
+    try:
+        iterations_int = int(iterations)
+    except ValueError:
+        return False
+    salt = _b64decode(salt_b64)
+    expected_digest = _b64decode(digest_b64)
+    candidate = hashlib.pbkdf2_hmac(
+        "sha256", token.encode("utf-8"), salt, iterations_int
+    )
+    return hmac.compare_digest(candidate, expected_digest)