[codex] Fix index and document ingestion reliability

backend/alembic/env.py

@@ -2,12 +2,12 @@
 
 from logging.config import fileConfig
 
-from alembic import context
 from sqlalchemy import engine_from_config, pool
 
+from alembic import context
 from chatmaster.config import get_settings
-from chatmaster.db.base import Base
 from chatmaster.db import models  # noqa: F401
+from chatmaster.db.base import Base
 
 config = context.config
 

backend/alembic/versions/0001_baseline.py

@@ -7,9 +7,10 @@
 
 from __future__ import annotations
 
-from alembic import op
 import sqlalchemy as sa
 
+from alembic import op
+
 revision = "0001_baseline"
 down_revision = None
 branch_labels = None

backend/alembic/versions/0002_reliability_fields.py

@@ -7,9 +7,9 @@
 
 from __future__ import annotations
 
-from alembic import op
 import sqlalchemy as sa
 
+from alembic import op
 
 revision = "0002_reliability_fields"
 down_revision = "0001_baseline"

backend/alembic/versions/0003_persona_management.py

@@ -7,9 +7,9 @@
 
 from __future__ import annotations
 
-from alembic import op
 import sqlalchemy as sa
 
+from alembic import op
 
 revision = "0003_persona_management"
 down_revision = "0002_reliability_fields"

backend/chatmaster/ai/models.py

@@ -115,6 +115,10 @@ def build_embeddings(identity: IdentityConfig | None = None) -> Embeddings:
     provider = cfg.provider.lower()
 
     if provider == "huggingface":
+        validate_provider_url(
+            cfg.huggingface_endpoint,
+            allow_private_network=settings.allow_private_provider_urls,
+        )
         return _build_hf_embeddings(model, cfg.huggingface_endpoint)
     if provider in _OPENAI_COMPATIBLE:
         api_key = cfg.api_key or settings.openai_api_key

backend/chatmaster/ai/providers.py

@@ -13,12 +13,11 @@
 import threading
 import uuid
 from functools import lru_cache
-
-from sqlalchemy import select
-from sqlalchemy.exc import SQLAlchemyError
 from typing import Literal
 
 from pydantic import BaseModel, Field
+from sqlalchemy import select
+from sqlalchemy.exc import SQLAlchemyError
 
 from chatmaster.config import get_settings
 

backend/chatmaster/chat/graph.py

@@ -16,8 +16,8 @@
 from chatmaster.conversations.service import load_history as load_history_from_db
 from chatmaster.db.models import Conversation, Message, utc_now
 from chatmaster.db.session import SessionLocal
-from chatmaster.identities.service import get_identity_config
 from chatmaster.identities.schema import IdentityConfig
+from chatmaster.identities.service import get_identity_config
 from chatmaster.retrieval.retriever import retrieve
 from chatmaster.retrieval.schemas import RetrievedChunk
 

backend/chatmaster/chat/service.py

@@ -78,9 +78,7 @@ def _begin_turn(
             conversation = db.get(Conversation, conversation_id)
             if conversation is not None and conversation.title == "新对话":
                 has_messages = db.scalars(
-                    select(Message.id)
-                    .where(Message.conversation_id == conversation_id)
-                    .limit(1)
+                    select(Message.id).where(Message.conversation_id == conversation_id).limit(1)
                 ).first()
                 if has_messages is None:
                     conversation.title = title_from_message(message)

backend/chatmaster/db/init_db.py

@@ -15,13 +15,13 @@ def init_db() -> None:
 def migrate_db() -> None:
     """Upgrade the runtime database, adopting legacy create_all databases once."""
     import shutil
-    from datetime import datetime
+    from datetime import datetime, timezone
     from pathlib import Path
 
-    from alembic import command
     from alembic.config import Config
     from sqlalchemy import inspect
 
+    from alembic import command
     from chatmaster.config import get_settings
 
     settings = get_settings()
@@ -34,7 +34,7 @@ def migrate_db() -> None:
         if settings.database_url.startswith("sqlite:///"):
             db_path = Path(settings.database_url.removeprefix("sqlite:///"))
             if db_path.exists():
-                suffix = datetime.now().strftime("%Y%m%d%H%M%S")
+                suffix = datetime.now(timezone.utc).strftime("%Y%m%d%H%M%S")
                 shutil.copy2(db_path, db_path.with_suffix(db_path.suffix + f".{suffix}.bak"))
         message_columns = (
             {column["name"] for column in inspector.get_columns("messages")}

backend/chatmaster/db/models.py

@@ -9,7 +9,7 @@
 from datetime import datetime, timezone
 from typing import Any
 
-from sqlalchemy import Boolean, ForeignKey, Integer, JSON, String, Text, UniqueConstraint
+from sqlalchemy import JSON, Boolean, ForeignKey, Integer, String, Text, UniqueConstraint
 from sqlalchemy.orm import Mapped, mapped_column, relationship
 
 from chatmaster.db.base import Base

backend/chatmaster/documents/jobs.py

@@ -6,11 +6,13 @@
 import threading
 from concurrent.futures import ThreadPoolExecutor
 
-from sqlalchemy import select
+from sqlalchemy import select, update
+from sqlalchemy.orm import aliased
 
 from chatmaster.config import get_settings
 from chatmaster.db.models import Document, IngestJob
 from chatmaster.db.session import SessionLocal
+from chatmaster.documents.service import document_operation_lock
 from chatmaster.services.ingest_service import ingest
 
 logger = logging.getLogger(__name__)
@@ -29,46 +31,80 @@ def _executor() -> ThreadPoolExecutor:
         return _EXECUTOR
 
 
+def _claim_job(db, *, job_id: str, document_id: str) -> bool:
+    """Atomically claim a pending job when no sibling job is already running."""
+    other_job = aliased(IngestJob)
+    claimed = db.execute(
+        update(IngestJob)
+        .where(
+            IngestJob.id == job_id,
+            IngestJob.status == "pending",
+            ~select(other_job.id)
+            .where(
+                other_job.document_id == document_id,
+                other_job.status == "running",
+                other_job.id != job_id,
+            )
+            .exists(),
+        )
+        .values(status="running", error=None)
+    ).rowcount
+    return claimed == 1
+
+
 def _run_job(job_id: str) -> None:
     with SessionLocal() as db:
-        job = db.get(IngestJob, job_id)
-        if job is None or job.status not in {"pending", "running"}:
+        initial_job = db.get(IngestJob, job_id)
+        if initial_job is None:
             return
-        document = db.get(Document, job.document_id)
-        if document is None:
-            job.status = "failed"
-            job.error = "Document not found"
+        document_id = initial_job.document_id
+        with document_operation_lock(initial_job.document_id):
+            if not _claim_job(db, job_id=job_id, document_id=document_id):
+                db.rollback()
+                return
             db.commit()
-            return
-        job.status = "running"
-        document.status = "ingesting"
-        db.commit()
-        try:
-            result = ingest(
-                document.identity_id,
-                [__import__("pathlib").Path(document.storage_path)],
-                target=document.namespace,
-                workspace_id=document.workspace_id,
-                db=db,
-            )
-            errors = [item.error for item in result.files if item.error]
-            if errors or result.total_chunks == 0:
-                raise RuntimeError("; ".join(errors) or "No indexable content was produced")
-            job.status = "completed"
-            job.total_chunks = result.total_chunks
-            job.error = None
-            document.status = "indexed"
-        except Exception as exc:  # noqa: BLE001
-            logger.exception("Ingest job failed job_id=%s document_id=%s", job.id, document.id)
-            db.rollback()
             job = db.get(IngestJob, job_id)
-            document = db.get(Document, job.document_id) if job is not None else None
-            if job is not None:
+            if job is None:
+                return
+            document = db.get(Document, job.document_id)
+            if document is None:
                 job.status = "failed"
-                job.error = f"{type(exc).__name__}: {exc}"
-            if document is not None:
-                document.status = "failed"
-        db.commit()
+                job.error = "Document not found"
+                db.commit()
+                return
+            document.status = "ingesting"
+            db.commit()
+            try:
+                result = ingest(
+                    document.identity_id,
+                    [__import__("pathlib").Path(document.storage_path)],
+                    target=document.namespace,
+                    workspace_id=document.workspace_id,
+                    db=db,
+                )
+                db.expire_all()
+                document = db.get(Document, job.document_id)
+                job = db.get(IngestJob, job_id)
+                if document is None or job is None or job.status != "running":
+                    raise RuntimeError("Document was removed while it was being indexed")
+                errors = [item.error for item in result.files if item.error]
+                if errors or result.total_chunks == 0:
+                    raise RuntimeError("; ".join(errors) or "No indexable content was produced")
+                job.status = "completed"
+                job.total_chunks = result.total_chunks
+                job.error = None
+                document.status = "indexed"
+            except Exception as exc:
+                logger.exception("Ingest job failed job_id=%s document_id=%s", job_id, document_id)
+                db.rollback()
+                job = db.get(IngestJob, job_id)
+                document = db.get(Document, job.document_id) if job is not None else None
+                if job is not None:
+                    job.status = "failed"
+                    job.error = f"{type(exc).__name__}: {exc}"
+                if document is not None:
+                    document.status = "failed"
+            db.commit()
 
 
 def enqueue_job(job_id: str) -> None:
@@ -109,7 +145,7 @@ def run() -> None:
                     identity_id=identity_id,
                     target=target,
                 )
-            except Exception:  # noqa: BLE001
+            except Exception:
                 logger.exception(
                     "Index rebuild failed workspace_id=%s identity_id=%s target=%s",
                     workspace_id,