diff --git a/backend/src/handlers/fileHandler.ts b/backend/src/handlers/fileHandler.ts index 3f8e25d..a92c8cf 100644 --- a/backend/src/handlers/fileHandler.ts +++ b/backend/src/handlers/fileHandler.ts @@ -224,9 +224,12 @@ const triggerExtraction = async (req: AuthenticatedRequest, res: Response): Prom // Reset state so the UI immediately reflects "queued" — even before // the worker picks the message up. Also flush the cached chat index: // a re-extracted file should be chunked from the new text, not the old. + // Reset extractionAttempts too — a manual retry deserves a fresh + // budget, not to inherit a stale count from the previous failure run. fileRecord.extractionStatus = 'pending'; fileRecord.extractionError = undefined; fileRecord.chatReady = false; + fileRecord.extractionAttempts = 0; await fileRecord.save(); // Best-effort: wipe per-chunk vectors so the next chat re-prepares them. wipeChunksForFile(userId, fileId).catch((err) => @@ -324,6 +327,7 @@ const togglePrivacy = async (req: AuthenticatedRequest, res: Response): Promise< file.isPrivate = isPrivate; file.extractionStatus = 'pending'; file.chatReady = false; + file.extractionAttempts = 0; await file.save(); await publishFileMetadata(file); diff --git a/backend/src/models/userFileModel.ts b/backend/src/models/userFileModel.ts index dd27e23..ed83d85 100644 --- a/backend/src/models/userFileModel.ts +++ b/backend/src/models/userFileModel.ts @@ -25,6 +25,13 @@ export interface UserFileType extends mongoose.Document { /** Extraction progress for UI ("extracting page 3 of 12"). Worker * updates this periodically during processing. Cleared when status='done'. */ extractionProgress?: { current: number; total: number; stage: string } | null; + /** Self-tracked retry counter, incremented by the worker on each Pub/Sub + * delivery. Once it exceeds MAX_DELIVERY_ATTEMPTS the worker marks the + * file 'failed' and stops redelivery — prevents infinite retry loops + * when the subscription has no dead-letter policy configured (Pub/Sub's + * own delivery_attempt field is 0 in that case and can't be used). + * Reset to 0 whenever the backend explicitly re-queues extraction. */ + extractionAttempts?: number; createdAt: Date; updatedAt: Date; } @@ -78,6 +85,7 @@ const userFileSchema = new mongoose.Schema( total: { type: Number, default: 0 }, stage: { type: String, default: "" }, }, + extractionAttempts: { type: Number, default: 0 }, }, { timestamps: true, diff --git a/smartdrive_core/src/smartdrive_core/mongo_status.py b/smartdrive_core/src/smartdrive_core/mongo_status.py index a511909..6f7889a 100644 --- a/smartdrive_core/src/smartdrive_core/mongo_status.py +++ b/smartdrive_core/src/smartdrive_core/mongo_status.py @@ -123,6 +123,65 @@ def update_progress(file_id: str, stage: str, current: int = 0, total: int = 0) return False +def increment_attempt_and_check(file_id: str, max_attempts: int) -> tuple[int, bool]: + """Atomically increment the self-tracked `extractionAttempts` counter and + report whether it has now exceeded `max_attempts`. + + Why self-tracked instead of relying on Pub/Sub's `delivery_attempt`: + that field is ONLY populated when the subscription has a dead-letter + policy configured. Without one (the common case unless explicitly set + up), `delivery_attempt` is always 0 and any cap based on it silently + never triggers — the exact bug that caused an infinite retry loop here. + A counter we own works regardless of subscription configuration. + + Returns (attempt_count_after_increment, exceeded). + Fails open (0, False) on any Mongo error — never block real processing + because status-tracking hiccupped. + """ + if not file_id: + return (0, False) + collection = _get_collection() + if collection is None: + return (0, False) + try: + from bson import ObjectId + oid = ObjectId(file_id) + except Exception: + return (0, False) + try: + from pymongo import ReturnDocument + doc = collection.find_one_and_update( + {"_id": oid}, + {"$inc": {"extractionAttempts": 1}, "$set": {"updatedAt": datetime.now(timezone.utc)}}, + return_document=ReturnDocument.AFTER, + projection={"extractionAttempts": 1}, + ) + if not doc: + return (0, False) + count = int(doc.get("extractionAttempts", 1)) + return (count, count > max_attempts) + except Exception as e: + logger.warning(f"increment_attempt_and_check({file_id}) failed: {e}") + return (0, False) + + +def reset_attempts(file_id: str) -> None: + """Clear the attempt counter. Called whenever a file is explicitly + re-queued (manual retry, privacy toggle) so it gets a fresh budget + instead of inheriting a stale count from a previous failure run.""" + if not file_id: + return + collection = _get_collection() + if collection is None: + return + try: + from bson import ObjectId + oid = ObjectId(file_id) + collection.update_one({"_id": oid}, {"$set": {"extractionAttempts": 0}}) + except Exception as e: + logger.warning(f"reset_attempts({file_id}) failed: {e}") + + def sweep_orphaned_files(stale_minutes: int = 10) -> int: """Find files stuck in `processing` for too long (worker died mid-extraction from OOM, timeout, deploy, etc.) and reset them to `pending` so they get diff --git a/smartdrive_core/src/smartdrive_core/pubsub_pull.py b/smartdrive_core/src/smartdrive_core/pubsub_pull.py index 6b53877..a8951c3 100644 --- a/smartdrive_core/src/smartdrive_core/pubsub_pull.py +++ b/smartdrive_core/src/smartdrive_core/pubsub_pull.py @@ -52,16 +52,16 @@ def pull_and_process(handler, subscription_id: str | None = None): logger.info("No messages in queue.") return 0 - # E8 — max retry count. Pub/Sub's `delivery_attempt` is set when the - # subscription has a dead-letter policy. Beyond MAX_RETRIES we ack the - # message (so Pub/Sub stops redelivering) AND mark the file `failed` - # permanently so the UI shows the right state instead of "processing". + # E8 — max retry count, self-tracked in Mongo (NOT via Pub/Sub's + # `delivery_attempt` — that field is only populated when the subscription + # has a dead-letter policy configured; without one it's always 0 and any + # cap based on it silently never fires, which is what caused the + # infinite-retry loop this replaces). max_retries = int(os.getenv("MAX_DELIVERY_ATTEMPTS", "3")) for rm in resp.received_messages: ack_id = rm.ack_id msg = rm.message - delivery_attempt = getattr(rm, "delivery_attempt", 0) or 0 # Extend ack deadline upfront if processing may take time subscriber.modify_ack_deadline( @@ -77,33 +77,32 @@ def pull_and_process(handler, subscription_id: str | None = None): subscriber.acknowledge(request={"subscription": sub_path, "ack_ids": [ack_id]}) continue - # E8 — give up after MAX retries. Mark the file failed so the - # user sees a clear status instead of "processing" forever. - if delivery_attempt > max_retries: - logger.warning( - f"Message exceeded max retries ({delivery_attempt}/{max_retries}) " - f"for {data.get('fileName')} — marking failed and acking" - ) - try: - from smartdrive_core.mongo_status import update_status - file_id = data.get("_id") - if file_id: - update_status( - str(file_id), - "failed", - error=f"Extraction failed after {delivery_attempt} attempts", - ) - except Exception as inner: - logger.warning(f"Failed to mark file as failed: {inner}") - subscriber.acknowledge(request={"subscription": sub_path, "ack_ids": [ack_id]}) - continue + # E8 — self-tracked attempt counter. Increment BEFORE processing + # so a crash mid-handler (OOM, timeout) still counts as an attempt. + file_id = data.get("_id") + attempt_count = 0 + if file_id: + from smartdrive_core.mongo_status import increment_attempt_and_check, update_status + attempt_count, exceeded = increment_attempt_and_check(str(file_id), max_retries) + if exceeded: + logger.warning( + f"File {file_id} ('{data.get('fileName')}') exceeded max attempts " + f"({attempt_count}/{max_retries}) — marking failed and acking to stop redelivery" + ) + update_status( + str(file_id), + "failed", + error=f"Extraction failed after {attempt_count} attempts", + ) + subscriber.acknowledge(request={"subscription": sub_path, "ack_ids": [ack_id]}) + continue handler(data) # <-- service-specific work subscriber.acknowledge(request={"subscription": sub_path, "ack_ids": [ack_id]}) logger.info( f"✅ Successfully processed {data['fileName']}" - + (f" (attempt {delivery_attempt})" if delivery_attempt > 1 else "") + + (f" (attempt {attempt_count})" if attempt_count > 1 else "") ) except json.JSONDecodeError: @@ -111,12 +110,9 @@ def pull_and_process(handler, subscription_id: str | None = None): subscriber.acknowledge(request={"subscription": sub_path, "ack_ids": [ack_id]}) except Exception as e: - logger.error( - f"❌ Error processing message (attempt {delivery_attempt}): {e}", - exc_info=True, - ) + logger.error(f"❌ Error processing message: {e}", exc_info=True) # IMPORTANT: do NOT ack on transient failure -> nack by setting deadline to 0 - # Pub/Sub will redeliver with delivery_attempt incremented. + # Pub/Sub will redeliver; our own counter (incremented above) caps retries. subscriber.modify_ack_deadline( request={"subscription": sub_path, "ack_ids": [ack_id], "ack_deadline_seconds": 0} ) diff --git a/smartdrive_core/src/smartdrive_core/weaviate_client.py b/smartdrive_core/src/smartdrive_core/weaviate_client.py index 74f8690..a742ad4 100644 --- a/smartdrive_core/src/smartdrive_core/weaviate_client.py +++ b/smartdrive_core/src/smartdrive_core/weaviate_client.py @@ -63,6 +63,20 @@ def _get_collection(collection_name: str): _collections[collection_name] = col return col +def _build_vector_index_config(): + """Weaviate Cloud plans differ in which vector index types they allow: + - Standard plans: hnsw (default) + - Serverless/hobby plans: only hfresh (managed by Weaviate) + Client-side we can't specify hfresh directly — it's server-selected. + `Configure.VectorIndex.dynamic()` tells the server "you pick" and works + across plans. Fall back to None (server default) if the SDK version + doesn't expose dynamic().""" + try: + return wvc.config.Configure.VectorIndex.dynamic() + except AttributeError: + return None + + def ensure_collection(collection_name: str, properties: list[wvc.config.Property]): client = get_weaviate_client() if client.collections.exists(collection_name): @@ -72,11 +86,35 @@ def ensure_collection(collection_name: str, properties: list[wvc.config.Property return logger.info(f"Creating collection '{collection_name}'...") - client.collections.create( - name=collection_name, - properties=properties, - vectorizer_config=wvc.config.Configure.Vectorizer.none(), - ) + + # Try the plan-adaptive vector index first, then fall back to server + # default if that gets rejected. Servers that require hfresh reject + # anything else with 422; catching lets us retry cleanly. + create_kwargs = { + "name": collection_name, + "properties": properties, + "vectorizer_config": wvc.config.Configure.Vectorizer.none(), + } + vic = _build_vector_index_config() + if vic is not None: + create_kwargs["vector_index_config"] = vic + + try: + client.collections.create(**create_kwargs) + except Exception as e: + msg = str(e).lower() + # If the cluster rejects our explicit vector_index_config, retry + # without one and let the server pick its default. + if "not allowed" in msg and "vector_index_type" in msg: + logger.warning( + f"Cluster rejected explicit vector index config for '{collection_name}' " + f"({e}). Retrying with server default (hfresh)." + ) + create_kwargs.pop("vector_index_config", None) + client.collections.create(**create_kwargs) + else: + raise + logger.info(f"Collection '{collection_name}' created.") @@ -98,6 +136,18 @@ def _ensure_properties(collection_name: str, properties: list[wvc.config.Propert logger.warning(f"_ensure_properties for {collection_name} failed: {e}") def check_file_exists(collection_name: str, file_id: str, user_id: str) -> bool: + # Cheap upfront check: if the collection doesn't exist yet, there's no + # file to find and no need to query. Prevents a noisy GRPC UNKNOWN + # traceback on first-ever ingestion into a fresh cluster. + try: + client = get_weaviate_client() + if not client.collections.exists(collection_name): + logger.info(f"check_file_exists: collection '{collection_name}' doesn't exist yet") + return False + except Exception as e: + logger.warning(f"check_file_exists preflight failed: {e}") + # Fall through to the query; may still succeed. + try: col = _get_collection(collection_name) filters = Filter.all_of([ @@ -107,8 +157,7 @@ def check_file_exists(collection_name: str, file_id: str, user_id: str) -> bool: resp = col.query.fetch_objects(limit=1, filters=filters) return bool(resp and resp.objects) except Exception as e: - logger.error(f"check_file_exists failed: {e}", exc_info=True) - # If exists-check fails, return False so pipeline can proceed (or you can return True to be conservative) + logger.warning(f"check_file_exists query failed (treating as absent): {e}") return False def upload(collection_name: str, properties_to_save: dict, embedding: list[float]) -> dict: