mirror of
https://github.com/pewdiepie-archdaemon/odysseus.git
synced 2026-11-02 19:22:22 +01:00
Compare commits
10
Commits
6efc07c5fc
...
d16b849c3e
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
d16b849c3e | ||
|
|
6b8f84553c | ||
|
|
b571c7ddc1 | ||
|
|
7f3fd77121 | ||
|
|
732b20776c | ||
|
|
801c3a2ff1 | ||
|
|
2531ba401c | ||
|
|
890d6a0220 | ||
|
|
1f8687abeb | ||
|
|
e6d9d68729 |
@@ -655,7 +655,12 @@ app.include_router(setup_emoji_routes())
|
||||
# Sessions
|
||||
from routes.session_routes import setup_session_routes
|
||||
session_config = {"REQUEST_TIMEOUT": REQUEST_TIMEOUT, "OPENAI_API_KEY": OPENAI_API_KEY, "SESSIONS_FILE": SESSIONS_FILE}
|
||||
app.include_router(setup_session_routes(session_manager, session_config, webhook_manager=webhook_manager))
|
||||
app.include_router(setup_session_routes(
|
||||
session_manager,
|
||||
session_config,
|
||||
webhook_manager=webhook_manager,
|
||||
upload_handler=upload_handler,
|
||||
))
|
||||
|
||||
# Admin Danger Zone wipes (Settings → System → Danger Zone)
|
||||
from routes.admin_wipe_routes import setup_admin_wipe_routes
|
||||
@@ -684,7 +689,7 @@ app.include_router(setup_research_routes(research_handler, session_manager=sessi
|
||||
|
||||
# History
|
||||
from routes.history.history_routes import setup_history_routes
|
||||
app.include_router(setup_history_routes(session_manager))
|
||||
app.include_router(setup_history_routes(session_manager, upload_handler=upload_handler))
|
||||
|
||||
# Search
|
||||
from routes.search_routes import setup_search_routes
|
||||
@@ -763,7 +768,7 @@ app.include_router(setup_assistant_routes(task_scheduler))
|
||||
|
||||
# Calendar (CalDAV)
|
||||
from routes.calendar_routes import setup_calendar_routes
|
||||
calendar_router = setup_calendar_routes()
|
||||
calendar_router = setup_calendar_routes(upload_handler=upload_handler)
|
||||
app.include_router(calendar_router)
|
||||
|
||||
# Shell (user-facing command execution)
|
||||
@@ -826,7 +831,7 @@ logger.info("Webhook & API token routes initialized")
|
||||
|
||||
# Notes (Google Keep-style notes/todos)
|
||||
from routes.note_routes import setup_note_routes
|
||||
app.include_router(setup_note_routes(task_scheduler))
|
||||
app.include_router(setup_note_routes(task_scheduler, upload_handler=upload_handler))
|
||||
|
||||
# Email
|
||||
from routes.email_routes import setup_email_routes
|
||||
|
||||
+55
-5
@@ -1904,6 +1904,20 @@ def _migrate_chat_messages_fts():
|
||||
conn = None
|
||||
try:
|
||||
conn = sqlite3.connect(db_path)
|
||||
fts_content_expr_new = (
|
||||
"CASE WHEN instr(COALESCE(new.content, ''), ';base64,') > 0 "
|
||||
"OR instr(COALESCE(new.content, ''), 'data:image/') > 0 "
|
||||
"OR instr(COALESCE(new.content, ''), 'data:audio/') > 0 "
|
||||
"THEN '[inline media omitted from search index]' "
|
||||
"ELSE COALESCE(new.content, '') END"
|
||||
)
|
||||
fts_content_expr_cm = (
|
||||
"CASE WHEN instr(COALESCE(cm.content, ''), ';base64,') > 0 "
|
||||
"OR instr(COALESCE(cm.content, ''), 'data:image/') > 0 "
|
||||
"OR instr(COALESCE(cm.content, ''), 'data:audio/') > 0 "
|
||||
"THEN '[inline media omitted from search index]' "
|
||||
"ELSE COALESCE(cm.content, '') END"
|
||||
)
|
||||
try:
|
||||
conn.execute("CREATE VIRTUAL TABLE IF NOT EXISTS temp._odysseus_fts5_probe USING fts5(content)")
|
||||
conn.execute("DROP TABLE IF EXISTS temp._odysseus_fts5_probe")
|
||||
@@ -1912,7 +1926,7 @@ def _migrate_chat_messages_fts():
|
||||
return
|
||||
|
||||
conn.executescript(
|
||||
"""
|
||||
f"""
|
||||
CREATE VIRTUAL TABLE IF NOT EXISTS chat_messages_fts USING fts5(
|
||||
content,
|
||||
message_id UNINDEXED,
|
||||
@@ -1920,10 +1934,14 @@ def _migrate_chat_messages_fts():
|
||||
role UNINDEXED
|
||||
);
|
||||
|
||||
DROP TRIGGER IF EXISTS chat_messages_fts_ai;
|
||||
DROP TRIGGER IF EXISTS chat_messages_fts_ad;
|
||||
DROP TRIGGER IF EXISTS chat_messages_fts_au;
|
||||
|
||||
CREATE TRIGGER IF NOT EXISTS chat_messages_fts_ai
|
||||
AFTER INSERT ON chat_messages BEGIN
|
||||
INSERT INTO chat_messages_fts(content, message_id, session_id, role)
|
||||
VALUES (COALESCE(new.content, ''), new.id, new.session_id, new.role);
|
||||
VALUES ({fts_content_expr_new}, new.id, new.session_id, new.role);
|
||||
END;
|
||||
|
||||
CREATE TRIGGER IF NOT EXISTS chat_messages_fts_ad
|
||||
@@ -1935,14 +1953,14 @@ def _migrate_chat_messages_fts():
|
||||
AFTER UPDATE ON chat_messages BEGIN
|
||||
DELETE FROM chat_messages_fts WHERE message_id = old.id;
|
||||
INSERT INTO chat_messages_fts(content, message_id, session_id, role)
|
||||
VALUES (COALESCE(new.content, ''), new.id, new.session_id, new.role);
|
||||
VALUES ({fts_content_expr_new}, new.id, new.session_id, new.role);
|
||||
END;
|
||||
"""
|
||||
)
|
||||
conn.execute(
|
||||
"""
|
||||
f"""
|
||||
INSERT INTO chat_messages_fts(content, message_id, session_id, role)
|
||||
SELECT COALESCE(cm.content, ''), cm.id, cm.session_id, cm.role
|
||||
SELECT {fts_content_expr_cm}, cm.id, cm.session_id, cm.role
|
||||
FROM chat_messages cm
|
||||
WHERE NOT EXISTS (
|
||||
SELECT 1 FROM chat_messages_fts fts
|
||||
@@ -1950,6 +1968,7 @@ def _migrate_chat_messages_fts():
|
||||
)
|
||||
"""
|
||||
)
|
||||
_scrub_legacy_chat_message_fts_media(conn)
|
||||
conn.commit()
|
||||
except Exception as e:
|
||||
logging.getLogger(__name__).warning(f"chat_messages FTS migration failed: {e}")
|
||||
@@ -1960,6 +1979,37 @@ def _migrate_chat_messages_fts():
|
||||
pass
|
||||
|
||||
|
||||
def _scrub_legacy_chat_message_fts_media(conn) -> None:
|
||||
"""Replace already-indexed inline media rows with searchable text only."""
|
||||
try:
|
||||
from src.attachment_refs import search_index_text
|
||||
except Exception as e:
|
||||
logging.getLogger(__name__).warning(f"chat_messages FTS media scrub skipped: {e}")
|
||||
return
|
||||
|
||||
try:
|
||||
rows = conn.execute(
|
||||
"""
|
||||
SELECT id, session_id, role, content
|
||||
FROM chat_messages
|
||||
WHERE instr(COALESCE(content, ''), ';base64,') > 0
|
||||
OR instr(COALESCE(content, ''), 'data:image/') > 0
|
||||
OR instr(COALESCE(content, ''), 'data:audio/') > 0
|
||||
"""
|
||||
).fetchall()
|
||||
for message_id, session_id, role, content in rows:
|
||||
conn.execute("DELETE FROM chat_messages_fts WHERE message_id = ?", (message_id,))
|
||||
conn.execute(
|
||||
"""
|
||||
INSERT INTO chat_messages_fts(content, message_id, session_id, role)
|
||||
VALUES (?, ?, ?, ?)
|
||||
""",
|
||||
(search_index_text(content), message_id, session_id, role),
|
||||
)
|
||||
except Exception as e:
|
||||
logging.getLogger(__name__).warning(f"chat_messages FTS media scrub failed: {e}")
|
||||
|
||||
|
||||
def _migrate_add_email_smtp_security():
|
||||
"""Add explicit SMTP security mode for Proton Bridge/custom local SMTP."""
|
||||
import sqlite3
|
||||
|
||||
+47
-21
@@ -16,6 +16,8 @@ from typing import Dict, Optional
|
||||
|
||||
from .database import Session as DbSession, ChatMessage as DbChatMessage, Document as DbDocument, SessionLocal, utcnow_naive
|
||||
from .models import Session, ChatMessage
|
||||
from src.attachment_refs import persistable_message_content
|
||||
from src.upload_handler import reserve_message_upload_references
|
||||
|
||||
# Re-export singleton accessors from models for convenience
|
||||
from .models import set_session_manager_instance, get_session_manager_instance
|
||||
@@ -72,6 +74,7 @@ class SessionManager:
|
||||
def __init__(self, sessions_file: str = None):
|
||||
# sessions_file kept for backward compat, not used
|
||||
self.sessions: Dict[str, Session] = {}
|
||||
self.upload_handler = None
|
||||
self.load_sessions()
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
@@ -230,17 +233,26 @@ class SessionManager:
|
||||
logger.warning("Dropping message for deleted session %s", session_id)
|
||||
return
|
||||
|
||||
missing_upload_id = reserve_message_upload_references(
|
||||
getattr(self, "upload_handler", None),
|
||||
getattr(db_session, "owner", None),
|
||||
message.content,
|
||||
message.metadata,
|
||||
)
|
||||
if missing_upload_id:
|
||||
raise ValueError(
|
||||
f"Referenced upload is no longer available: {missing_upload_id}"
|
||||
)
|
||||
|
||||
msg_id = str(uuid.uuid4())
|
||||
msg_time = datetime.utcnow()
|
||||
if message.metadata is None:
|
||||
message.metadata = {}
|
||||
message.metadata.setdefault('timestamp', _message_timestamp_iso(msg_time))
|
||||
# Multimodal content (image/audio attachments) is a list — serialize
|
||||
# to JSON so the Text column can store it. On reload, _db_to_session
|
||||
# detects the JSON-array prefix and parses it back.
|
||||
_content = message.content
|
||||
if isinstance(_content, list):
|
||||
_content = json.dumps(_content)
|
||||
# Multimodal content may contain provider data URLs for the live
|
||||
# model call. Persist only readable text plus attachment references
|
||||
# so chat_messages/FTS do not duplicate upload bytes.
|
||||
_content = persistable_message_content(message.content, message.metadata)
|
||||
db_message = DbChatMessage(
|
||||
id=msg_id,
|
||||
session_id=session_id,
|
||||
@@ -322,6 +334,28 @@ class SessionManager:
|
||||
session = self.get_session(session_id)
|
||||
db = SessionLocal()
|
||||
try:
|
||||
db_session = db.query(DbSession).filter(DbSession.id == session_id).first()
|
||||
if db_session is None:
|
||||
logger.warning("Cannot replace history for missing session %s", session_id)
|
||||
return False
|
||||
|
||||
# Reserve every incoming attachment before removing any durable
|
||||
# message row. reserve_upload() shares the upload lifecycle lock
|
||||
# with cleanup, so an upload cannot be deleted between this
|
||||
# ownership check/access touch and the replacement transaction.
|
||||
# A failed reservation must leave the existing transcript intact.
|
||||
for message in messages:
|
||||
missing_upload_id = reserve_message_upload_references(
|
||||
getattr(self, "upload_handler", None),
|
||||
getattr(db_session, "owner", None),
|
||||
message.content,
|
||||
message.metadata,
|
||||
)
|
||||
if missing_upload_id:
|
||||
raise ValueError(
|
||||
f"Referenced upload is no longer available: {missing_upload_id}"
|
||||
)
|
||||
|
||||
db.query(DbChatMessage).filter(DbChatMessage.session_id == session_id).delete()
|
||||
now = datetime.now(timezone.utc)
|
||||
for i, message in enumerate(messages):
|
||||
@@ -330,15 +364,9 @@ class SessionManager:
|
||||
id=msg_id,
|
||||
session_id=session_id,
|
||||
role=message.role,
|
||||
# Multimodal content (image/audio attachments) is a list;
|
||||
# serialize to JSON so the Text column round-trips via
|
||||
# _parse_msg_content. Storing the raw list let SQLAlchemy
|
||||
# bind its single-quoted repr, which _parse_msg_content
|
||||
# cannot parse (it looks for double-quoted "type"), so the
|
||||
# attachment was destroyed on reload. Mirrors _persist_message.
|
||||
content=(json.dumps(message.content)
|
||||
if isinstance(message.content, list)
|
||||
else message.content),
|
||||
# Mirrors _persist_message: keep raw media bytes out of the
|
||||
# persisted transcript and search index.
|
||||
content=persistable_message_content(message.content, message.metadata),
|
||||
meta_data=json.dumps(message.metadata) if message.metadata else None,
|
||||
timestamp=now + timedelta(microseconds=i),
|
||||
)
|
||||
@@ -347,12 +375,10 @@ class SessionManager:
|
||||
message.metadata = {}
|
||||
message.metadata["_db_id"] = msg_id
|
||||
|
||||
db_session = db.query(DbSession).filter(DbSession.id == session_id).first()
|
||||
if db_session:
|
||||
db_session.message_count = len(messages)
|
||||
db_session.updated_at = now
|
||||
db_session.last_accessed = now
|
||||
db_session.last_message_at = now
|
||||
db_session.message_count = len(messages)
|
||||
db_session.updated_at = now
|
||||
db_session.last_accessed = now
|
||||
db_session.last_message_at = now
|
||||
|
||||
db.commit()
|
||||
session.history = list(messages)
|
||||
|
||||
@@ -0,0 +1,85 @@
|
||||
# Attachment References and Upload Storage
|
||||
|
||||
Odysseus stores uploaded bytes once under the configured upload directory and
|
||||
passes stable references through chat history, tools, and future artifact work.
|
||||
The goal is to avoid duplicating large inline media payloads in
|
||||
`chat_messages.content` or the SQLite FTS index.
|
||||
|
||||
## Reference Shape
|
||||
|
||||
Attachment references use this minimum shape:
|
||||
|
||||
```json
|
||||
{
|
||||
"type": "attachment_ref",
|
||||
"attachment_id": "32hex-or-32hex.ext",
|
||||
"name": "original-filename.png",
|
||||
"mime": "image/png",
|
||||
"size": 12345,
|
||||
"checksum_sha256": "hex-digest",
|
||||
"created_at": "2026-07-09T12:00:00"
|
||||
}
|
||||
```
|
||||
|
||||
Optional fields such as `width`, `height`, `vision`, `vision_model`, and
|
||||
`gallery_id` may be present when the uploader or preprocessing path knows them.
|
||||
|
||||
## Persistence
|
||||
|
||||
The live model call may still receive provider-specific multimodal blocks for
|
||||
the current turn. Persistence is different:
|
||||
|
||||
- `chat_messages.content` stores readable text plus compact attachment reference
|
||||
lines, never raw `data:*;base64,...` upload bytes.
|
||||
- `chat_messages.metadata.attachments` stores structured attachment reference
|
||||
metadata for UI reloads and future processing.
|
||||
- The SQLite FTS migration recreates chat-message FTS triggers so new rows do
|
||||
not index inline media payloads, and it scrubs legacy rows that were already
|
||||
indexed with data URLs.
|
||||
|
||||
## Tool Access
|
||||
|
||||
Agent/tool context receives upload entries as `attachment_ref` manifests with an
|
||||
`odysseus://attachment/<id>` URI and `read_policy: "owner_checked_upload"`.
|
||||
|
||||
For compatibility with existing built-in tools, a local `path` may be included
|
||||
only after all of these checks pass:
|
||||
|
||||
- the upload ID resolves through `UploadHandler.resolve_upload`;
|
||||
- the requested owner is allowed to read the upload;
|
||||
- the file remains inside the configured upload directory;
|
||||
- the file path is inside the tool-readable roots.
|
||||
|
||||
External MCP/custom tools should treat the URI and attachment ID as the stable
|
||||
contract and request bytes through an owner-checked server path, not by assuming
|
||||
host filesystem layout.
|
||||
|
||||
## Retention and Deletion
|
||||
|
||||
Current retention behavior is conservative:
|
||||
|
||||
- uploads are indexed in `uploads.json` with owner, checksum, MIME type, size,
|
||||
and creation time;
|
||||
- admin cleanup first scans persisted chat metadata/content, document versions,
|
||||
PDF source markers, gallery hashes, notes, and calendar records for live
|
||||
references;
|
||||
- cleanup fails closed if that reference scan cannot complete, and the lower-level
|
||||
cleanup API removes nothing unless it receives a complete reference snapshot;
|
||||
- expired, unreferenced uploads are removed during the completed scan, while
|
||||
attachment-bearing writers must first take an owner-checked reservation that
|
||||
serializes with deletion and refreshes the upload's access timestamp;
|
||||
- deliberate removal atomically drops matching `uploads.json` rows before deleting
|
||||
the bytes and restores those rows if filesystem removal fails;
|
||||
- deleting a chat removes the chat rows but does not immediately delete shared
|
||||
upload bytes, because the same upload may also be referenced by gallery items,
|
||||
documents, duplicate-upload rows, or future artifact records.
|
||||
|
||||
There is no distinct artifact table in the current schema. Artifact-like upload
|
||||
references persisted in chat or document text are covered by the canonical
|
||||
attachment-ID scan; any future artifact store must be added to reference discovery
|
||||
before cleanup is allowed to consider its uploads unreferenced.
|
||||
|
||||
Cleanup and write reservations share the upload-index lock. This closes the
|
||||
scan/write/delete race in the documented single-worker deployment; a future
|
||||
multi-process deployment must add an inter-process lock or move lifecycle state
|
||||
into the database before enabling destructive cleanup in more than one worker.
|
||||
@@ -13,8 +13,9 @@ from sqlalchemy import or_, and_
|
||||
from dateutil.rrule import rrulestr
|
||||
|
||||
from core.database import SessionLocal, CalendarCal, CalendarDeletedEvent, CalendarEvent
|
||||
from src.auth_helpers import require_user
|
||||
from src.auth_helpers import effective_user, require_user
|
||||
from src.upload_limits import read_upload_limited, ICS_MAX_BYTES
|
||||
from src.upload_handler import reserve_upload_references
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -697,9 +698,18 @@ def _expand_rrule(
|
||||
|
||||
# ── Routes ──
|
||||
|
||||
def setup_calendar_routes() -> APIRouter:
|
||||
def setup_calendar_routes(upload_handler=None) -> APIRouter:
|
||||
router = APIRouter(prefix="/api/calendar", tags=["calendar"])
|
||||
|
||||
def _reserve_calendar_uploads(request: Request, *values) -> None:
|
||||
missing_id = reserve_upload_references(
|
||||
upload_handler,
|
||||
effective_user(request),
|
||||
*values,
|
||||
)
|
||||
if missing_id:
|
||||
raise HTTPException(409, f"Referenced upload is no longer available: {missing_id}")
|
||||
|
||||
# ── CalDAV multi-account helpers ─────────────────────────────────────────
|
||||
|
||||
def _get_caldav_accounts(owner: str) -> list:
|
||||
@@ -913,7 +923,24 @@ def setup_calendar_routes() -> APIRouter:
|
||||
'</d:prop></d:propfind>'
|
||||
)
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=8.0, follow_redirects=False, trust_env=False) as cx:
|
||||
# Build an SSL context that trusts the operator's custom CA bundle
|
||||
# (SSL_CERT_FILE / REQUESTS_CA_BUNDLE) so self-signed CalDAV servers
|
||||
# pass the pre-flight the same way they pass the real sync.
|
||||
# trust_env=False is kept to block proxy/auth env leakage; the CA
|
||||
# bundle is loaded explicitly instead.
|
||||
import ssl as _ssl
|
||||
_ssl_ctx = _ssl.create_default_context()
|
||||
# Disable VERIFY_X509_STRICT so certs without a keyUsage extension
|
||||
# (common in self-signed setups) are accepted, matching the
|
||||
# requests/urllib3 behavior used by the CalDAV sync path.
|
||||
_ssl_ctx.verify_flags &= ~_ssl.VERIFY_X509_STRICT
|
||||
_ca_bundle = _os.environ.get("SSL_CERT_FILE") or _os.environ.get("REQUESTS_CA_BUNDLE")
|
||||
if _ca_bundle:
|
||||
if _os.path.isfile(_ca_bundle):
|
||||
_ssl_ctx.load_verify_locations(_ca_bundle)
|
||||
else:
|
||||
logger.warning("CalDAV test: CA bundle %s not found, using system CAs", _ca_bundle)
|
||||
async with httpx.AsyncClient(timeout=8.0, follow_redirects=False, trust_env=False, verify=_ssl_ctx) as cx:
|
||||
r = await cx.request(
|
||||
"PROPFIND", url,
|
||||
auth=(user, pw),
|
||||
@@ -1070,6 +1097,7 @@ def setup_calendar_routes() -> APIRouter:
|
||||
@router.post("/events")
|
||||
async def create_event(request: Request, data: EventCreate):
|
||||
owner = _require_user(request)
|
||||
_reserve_calendar_uploads(request, data.color, data.description, data.location)
|
||||
db = SessionLocal()
|
||||
try:
|
||||
cal = None
|
||||
@@ -1131,6 +1159,7 @@ def setup_calendar_routes() -> APIRouter:
|
||||
@router.put("/events/{uid}")
|
||||
async def update_event(request: Request, uid: str, data: EventUpdate):
|
||||
owner = _require_user(request)
|
||||
_reserve_calendar_uploads(request, data.color, data.description, data.location)
|
||||
try:
|
||||
base_uid = _resolve_base_uid(uid)
|
||||
except ValueError as e:
|
||||
@@ -1224,6 +1253,7 @@ def setup_calendar_routes() -> APIRouter:
|
||||
@router.post("/calendars")
|
||||
async def create_calendar(request: Request, name: str = "Imported", color: str = "#5b8abf"):
|
||||
owner = _require_user(request)
|
||||
_reserve_calendar_uploads(request, color)
|
||||
db = SessionLocal()
|
||||
try:
|
||||
cal = CalendarCal(
|
||||
@@ -1246,6 +1276,7 @@ def setup_calendar_routes() -> APIRouter:
|
||||
@router.put("/calendars/{cal_id}")
|
||||
async def update_calendar(request: Request, cal_id: str, name: str = None, color: str = None):
|
||||
owner = _require_user(request)
|
||||
_reserve_calendar_uploads(request, color)
|
||||
db = SessionLocal()
|
||||
try:
|
||||
cal = _get_or_404_calendar(db, cal_id, owner)
|
||||
|
||||
@@ -17,6 +17,7 @@ from src.context_compactor import maybe_compact, trim_for_context
|
||||
from src.model_context import estimate_tokens
|
||||
from src.auth_helpers import effective_user
|
||||
from src.prompt_security import untrusted_context_message
|
||||
from src.attachment_refs import attachment_ref
|
||||
from routes.prefs_routes import _load_for_user as load_prefs_for_user
|
||||
|
||||
from fastapi import HTTPException
|
||||
@@ -418,13 +419,16 @@ def build_uploaded_file_manifest(att_ids: list, upload_handler, owner: Optional[
|
||||
except Exception:
|
||||
path = None
|
||||
|
||||
manifest.append({
|
||||
"id": info.get("id") or str(att_id),
|
||||
"name": info.get("name") or info.get("original_name") or str(att_id),
|
||||
"mime": info.get("mime", ""),
|
||||
"size": info.get("size", 0),
|
||||
ref = attachment_ref({**info, "id": info.get("id") or str(att_id)})
|
||||
ref.update({
|
||||
"id": ref["attachment_id"],
|
||||
"uri": f"odysseus://attachment/{ref['attachment_id']}",
|
||||
"read_policy": "owner_checked_upload",
|
||||
# Transitional compatibility: existing built-in tools can still use
|
||||
# this path, but only after owner, upload-root, and tool-root checks.
|
||||
"path": path,
|
||||
})
|
||||
manifest.append(ref)
|
||||
return manifest
|
||||
|
||||
|
||||
|
||||
@@ -1450,7 +1450,9 @@ def setup_chat_routes(
|
||||
"tool_start", "tool_output", "agent_step",
|
||||
"doc_stream_open", "doc_stream_delta",
|
||||
"doc_update", "doc_suggestions", "ui_control",
|
||||
"rounds_exhausted",
|
||||
"rounds_exhausted", "budget_exceeded",
|
||||
"loop_breaker_triggered",
|
||||
"intent_nudge_exhausted",
|
||||
"ask_user",
|
||||
"plan_update",
|
||||
):
|
||||
|
||||
@@ -12,6 +12,7 @@ from core.database import SessionLocal, Document, DocumentVersion
|
||||
from core.database import Session as DbSession
|
||||
from src.auth_helpers import get_current_user, _auth_disabled
|
||||
from src.constants import MAIL_ATTACHMENTS_DIR
|
||||
from src.upload_handler import reserve_upload_references
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -78,6 +79,14 @@ from routes.document_helpers import (
|
||||
def setup_document_routes(session_manager, upload_handler=None) -> APIRouter:
|
||||
router = APIRouter(tags=["documents"])
|
||||
|
||||
def _reserve_document_uploads(user: Optional[str], content: str) -> None:
|
||||
missing_id = reserve_upload_references(upload_handler, user, content)
|
||||
if missing_id:
|
||||
raise HTTPException(
|
||||
409,
|
||||
f"Referenced upload is no longer available: {missing_id}",
|
||||
)
|
||||
|
||||
def _locate_current_user_upload(request: Request, upload_id: str, user: Optional[str]):
|
||||
if upload_handler is None:
|
||||
return None
|
||||
@@ -124,6 +133,7 @@ def setup_document_routes(session_manager, upload_handler=None) -> APIRouter:
|
||||
if _looks_like_email_document(req.content, req.title):
|
||||
language = "email"
|
||||
|
||||
_reserve_document_uploads(user, req.content)
|
||||
_assert_pdf_marker_upload_owned(request, req.content, user, upload_handler)
|
||||
|
||||
# Reply drafts are keyed to the source email. If a UI/tool path tries
|
||||
@@ -636,6 +646,7 @@ def setup_document_routes(session_manager, upload_handler=None) -> APIRouter:
|
||||
if doc.current_content == incoming_content and not req.force_version:
|
||||
return _doc_to_dict(doc)
|
||||
|
||||
_reserve_document_uploads(user, incoming_content)
|
||||
_assert_pdf_marker_upload_owned(request, incoming_content, user, upload_handler)
|
||||
|
||||
# Check if we can coalesce with the latest version
|
||||
|
||||
@@ -10,7 +10,9 @@ from fastapi import APIRouter, Request, HTTPException
|
||||
|
||||
from core.models import ChatMessage
|
||||
from core.database import SessionLocal, ChatMessage as DbChatMessage, Session as DbSession
|
||||
from src.auth_helpers import effective_user
|
||||
from src.topic_analyzer import analyze_topics
|
||||
from src.upload_handler import reserve_message_upload_references
|
||||
from routes.session_routes import (
|
||||
_message_role,
|
||||
_message_text,
|
||||
@@ -98,9 +100,29 @@ def _merge_continue_rows_to_delete(db_messages, db1, db2):
|
||||
return to_delete
|
||||
|
||||
|
||||
def setup_history_routes(session_manager) -> APIRouter:
|
||||
def setup_history_routes(session_manager, upload_handler=None) -> APIRouter:
|
||||
router = APIRouter(tags=["history"])
|
||||
|
||||
def _reserve_message_uploads(
|
||||
request: Request,
|
||||
content: Any,
|
||||
metadata: Any = None,
|
||||
) -> None:
|
||||
try:
|
||||
missing_id = reserve_message_upload_references(
|
||||
upload_handler,
|
||||
effective_user(request),
|
||||
content,
|
||||
metadata,
|
||||
)
|
||||
except (TypeError, ValueError) as exc:
|
||||
raise HTTPException(400, "Invalid message attachment metadata") from exc
|
||||
if missing_id:
|
||||
raise HTTPException(
|
||||
409,
|
||||
f"Referenced upload is no longer available: {missing_id}",
|
||||
)
|
||||
|
||||
def _db_history_entry(m: DbChatMessage) -> Dict[str, Any]:
|
||||
entry = {"role": m.role, "content": _history_display_content(m.content)}
|
||||
meta = {}
|
||||
@@ -251,7 +273,9 @@ def setup_history_routes(session_manager) -> APIRouter:
|
||||
content = body.get("content", "")
|
||||
if not content:
|
||||
raise HTTPException(400, "content is required")
|
||||
msg = ChatMessage(role=role, content=content, metadata=body.get("metadata"))
|
||||
metadata = body.get("metadata")
|
||||
_reserve_message_uploads(request, content, metadata)
|
||||
msg = ChatMessage(role=role, content=content, metadata=metadata)
|
||||
session_manager.add_message(session_id, msg)
|
||||
return {"status": "ok"}
|
||||
except KeyError:
|
||||
@@ -331,6 +355,8 @@ def setup_history_routes(session_manager) -> APIRouter:
|
||||
if not msg_id or content is None:
|
||||
raise HTTPException(400, "msg_id and content are required")
|
||||
|
||||
_reserve_message_uploads(request, content)
|
||||
|
||||
session = session_manager.get_session(session_id)
|
||||
db = SessionLocal()
|
||||
try:
|
||||
|
||||
+21
-1
@@ -13,6 +13,7 @@ from core.database import SessionLocal, Note
|
||||
from core.middleware import INTERNAL_TOOL_USER
|
||||
from src.auth_helpers import require_user
|
||||
from src.constants import DATA_DIR
|
||||
from src.upload_handler import reserve_upload_references
|
||||
from sqlalchemy.orm.attributes import flag_modified
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -574,7 +575,7 @@ async def dispatch_reminder(
|
||||
# Router factory
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def setup_note_routes(task_scheduler=None):
|
||||
def setup_note_routes(task_scheduler=None, upload_handler=None):
|
||||
# Expose the scheduler to module-level `dispatch_reminder` so reminders
|
||||
# can also push to the in-app notification queue (the polling system
|
||||
# turns each entry into a real browser Notification + the existing
|
||||
@@ -596,6 +597,11 @@ def setup_note_routes(task_scheduler=None):
|
||||
# did not.
|
||||
return require_user(request) or None
|
||||
|
||||
def _reserve_note_uploads(owner: Optional[str], *values) -> None:
|
||||
missing_id = reserve_upload_references(upload_handler, owner, *values)
|
||||
if missing_id:
|
||||
raise HTTPException(409, f"Referenced upload is no longer available: {missing_id}")
|
||||
|
||||
def _is_admin_or_single_user(request: Request, user: str | None) -> bool:
|
||||
if user == INTERNAL_TOOL_USER:
|
||||
return True
|
||||
@@ -645,6 +651,13 @@ def setup_note_routes(task_scheduler=None):
|
||||
@router.post("")
|
||||
def create_note(request: Request, body: NoteCreate):
|
||||
user = _owner(request)
|
||||
_reserve_note_uploads(
|
||||
user,
|
||||
body.image_url,
|
||||
body.color,
|
||||
body.content,
|
||||
json.dumps(body.items) if body.items is not None else None,
|
||||
)
|
||||
db = SessionLocal()
|
||||
try:
|
||||
note = Note(
|
||||
@@ -702,6 +715,13 @@ def setup_note_routes(task_scheduler=None):
|
||||
if user is not None and note.owner != user:
|
||||
raise HTTPException(404, "Note not found")
|
||||
|
||||
_reserve_note_uploads(
|
||||
user,
|
||||
body.image_url,
|
||||
body.color,
|
||||
body.content,
|
||||
json.dumps(body.items) if body.items is not None else None,
|
||||
)
|
||||
if body.title is not None:
|
||||
note.title = body.title
|
||||
if body.content is not None:
|
||||
|
||||
@@ -13,6 +13,7 @@ from src.request_models import SessionResponse
|
||||
from core.database import Session as DbSession, SessionLocal, Document, GalleryImage, utcnow_naive
|
||||
from src.auth_helpers import effective_user, _auth_disabled, owner_filter
|
||||
from src.session_actions import is_session_recently_active
|
||||
from src.upload_handler import reserve_message_upload_references
|
||||
|
||||
|
||||
def _sanitize_export_filename(name: str) -> str:
|
||||
@@ -203,7 +204,12 @@ def _pick_endpoint_for_sort(owner=None):
|
||||
return url, model, headers
|
||||
return None, None, None
|
||||
|
||||
def setup_session_routes(session_manager: SessionManager, config: dict, webhook_manager=None):
|
||||
def setup_session_routes(
|
||||
session_manager: SessionManager,
|
||||
config: dict,
|
||||
webhook_manager=None,
|
||||
upload_handler=None,
|
||||
):
|
||||
"""Setup session routes with the provided manager and config"""
|
||||
|
||||
REQUEST_TIMEOUT = config.get("REQUEST_TIMEOUT", 20)
|
||||
@@ -537,6 +543,22 @@ def setup_session_routes(session_manager: SessionManager, config: dict, webhook_
|
||||
body = await request.json()
|
||||
messages = body.get("messages", [])
|
||||
from core.models import ChatMessage
|
||||
owner = effective_user(request)
|
||||
try:
|
||||
for message in messages:
|
||||
missing_id = reserve_message_upload_references(
|
||||
upload_handler,
|
||||
owner,
|
||||
message.get("content"),
|
||||
message.get("metadata"),
|
||||
)
|
||||
if missing_id:
|
||||
raise HTTPException(
|
||||
409,
|
||||
f"Referenced upload is no longer available: {missing_id}",
|
||||
)
|
||||
except (AttributeError, TypeError, ValueError) as exc:
|
||||
raise HTTPException(400, "Invalid message attachment metadata") from exc
|
||||
for m in messages:
|
||||
sess.add_message(ChatMessage(m["role"], m["content"], metadata=m.get("metadata")))
|
||||
session_manager.save_sessions()
|
||||
|
||||
+145
-3
@@ -10,16 +10,140 @@ from fastapi import APIRouter, Request, File, UploadFile, HTTPException, Form
|
||||
from typing import List, Optional
|
||||
import logging
|
||||
from core.middleware import require_admin
|
||||
from core.database import SessionLocal, GalleryImage, Session as DbSession
|
||||
from core.database import (
|
||||
SessionLocal,
|
||||
ChatMessage as DbChatMessage,
|
||||
CalendarCal,
|
||||
CalendarEvent,
|
||||
Document,
|
||||
DocumentVersion,
|
||||
GalleryImage,
|
||||
Note,
|
||||
Session as DbSession,
|
||||
)
|
||||
from src.auth_helpers import effective_user
|
||||
from src.attachment_refs import attachment_refs_from_metadata
|
||||
from src.constants import GENERATED_IMAGES_DIR
|
||||
from src.upload_handler import count_recent_uploads
|
||||
from src.upload_handler import (
|
||||
UploadCleanupSafetyError,
|
||||
count_recent_uploads,
|
||||
extract_upload_ids,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
router = APIRouter(prefix="/api/upload", tags=["upload"])
|
||||
UPLOAD_RESPONSE_HEADERS = {"X-Content-Type-Options": "nosniff"}
|
||||
|
||||
def _upload_ids_from_persisted_text(value: object) -> set[str]:
|
||||
"""Return canonical upload IDs embedded in persisted text.
|
||||
|
||||
This covers attachment reference lines/URIs and the PDF source markers
|
||||
stored by the document editor. False positives are intentionally
|
||||
conservative: retaining an extra upload is safer than deleting referenced
|
||||
bytes.
|
||||
"""
|
||||
return extract_upload_ids(value)
|
||||
|
||||
|
||||
def _upload_ids_from_message_metadata(raw_metadata: object) -> set[str]:
|
||||
"""Extract attachment IDs from a persisted chat metadata JSON value.
|
||||
|
||||
Malformed metadata raises instead of being treated as an empty reference
|
||||
set. The admin cleanup route catches that failure and aborts cleanup.
|
||||
"""
|
||||
if raw_metadata in (None, ""):
|
||||
return set()
|
||||
if isinstance(raw_metadata, str):
|
||||
metadata = json.loads(raw_metadata)
|
||||
else:
|
||||
metadata = raw_metadata
|
||||
if not isinstance(metadata, dict):
|
||||
raise ValueError("chat message metadata must be a JSON object")
|
||||
|
||||
attachments = metadata.get("attachments")
|
||||
if attachments is not None:
|
||||
if not isinstance(attachments, list) or any(
|
||||
not isinstance(item, dict) for item in attachments
|
||||
):
|
||||
raise ValueError("chat message attachments metadata is malformed")
|
||||
|
||||
ids = {
|
||||
str(ref["attachment_id"])
|
||||
for ref in attachment_refs_from_metadata(metadata)
|
||||
if ref.get("attachment_id")
|
||||
}
|
||||
# Preserve canonical IDs even in older metadata shapes not normalized by
|
||||
# attachment_refs_from_metadata().
|
||||
ids.update(_upload_ids_from_persisted_text(json.dumps(metadata)))
|
||||
return ids
|
||||
|
||||
|
||||
def _collect_persisted_upload_references() -> tuple[set[str], set[str]]:
|
||||
"""Collect upload IDs/hashes still referenced by durable application data.
|
||||
|
||||
The caller must treat any exception as an incomplete scan and fail closed.
|
||||
There is no distinct artifact table in the current schema; artifact-like
|
||||
attachment references persisted in chat/document text are covered by the
|
||||
canonical-ID scan.
|
||||
"""
|
||||
referenced_ids: set[str] = set()
|
||||
referenced_hashes: set[str] = set()
|
||||
db = SessionLocal()
|
||||
try:
|
||||
for content, raw_metadata in db.query(
|
||||
DbChatMessage.content,
|
||||
DbChatMessage.meta_data,
|
||||
).yield_per(500):
|
||||
referenced_ids.update(_upload_ids_from_persisted_text(content))
|
||||
referenced_ids.update(_upload_ids_from_message_metadata(raw_metadata))
|
||||
|
||||
for (content,) in db.query(Document.current_content).yield_per(500):
|
||||
referenced_ids.update(_upload_ids_from_persisted_text(content))
|
||||
|
||||
for (content,) in db.query(DocumentVersion.content).yield_per(500):
|
||||
referenced_ids.update(_upload_ids_from_persisted_text(content))
|
||||
|
||||
for filename, file_hash in db.query(
|
||||
GalleryImage.filename,
|
||||
GalleryImage.file_hash,
|
||||
).yield_per(500):
|
||||
referenced_ids.update(_upload_ids_from_persisted_text(filename))
|
||||
if file_hash:
|
||||
referenced_hashes.add(str(file_hash))
|
||||
|
||||
for image_url, color, content, items in db.query(
|
||||
Note.image_url,
|
||||
Note.color,
|
||||
Note.content,
|
||||
Note.items,
|
||||
).yield_per(500):
|
||||
for value in (image_url, color, content, items):
|
||||
referenced_ids.update(_upload_ids_from_persisted_text(value))
|
||||
|
||||
for (color,) in db.query(CalendarCal.color).yield_per(500):
|
||||
referenced_ids.update(_upload_ids_from_persisted_text(color))
|
||||
|
||||
for color, description, location in db.query(
|
||||
CalendarEvent.color,
|
||||
CalendarEvent.description,
|
||||
CalendarEvent.location,
|
||||
).yield_per(500):
|
||||
for value in (color, description, location):
|
||||
referenced_ids.update(_upload_ids_from_persisted_text(value))
|
||||
|
||||
return referenced_ids, referenced_hashes
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
|
||||
def _run_reference_safe_cleanup(upload_handler) -> int:
|
||||
referenced_ids, referenced_hashes = _collect_persisted_upload_references()
|
||||
return upload_handler.cleanup_old_uploads(
|
||||
referenced_upload_ids=referenced_ids,
|
||||
referenced_upload_hashes=referenced_hashes,
|
||||
)
|
||||
|
||||
def setup_upload_routes(upload_handler):
|
||||
"""Setup upload routes with the provided handler"""
|
||||
|
||||
@@ -172,7 +296,9 @@ def setup_upload_routes(upload_handler):
|
||||
"mime": meta["mime"],
|
||||
"size": meta["size"],
|
||||
"hash": meta["hash"],
|
||||
"checksum_sha256": meta.get("checksum_sha256") or meta["hash"],
|
||||
"uploaded_at": meta["uploaded_at"],
|
||||
"created_at": meta.get("created_at") or meta["uploaded_at"],
|
||||
"width": meta.get("width"),
|
||||
"height": meta.get("height"),
|
||||
"is_duplicate": meta.get("is_duplicate", False)
|
||||
@@ -195,7 +321,23 @@ def setup_upload_routes(upload_handler):
|
||||
async def manual_cleanup(request: Request):
|
||||
"""Manually trigger cleanup of old uploads."""
|
||||
require_admin(request)
|
||||
cleaned_count = upload_handler.cleanup_old_uploads()
|
||||
try:
|
||||
cleaned_count = await asyncio.to_thread(
|
||||
_run_reference_safe_cleanup,
|
||||
upload_handler,
|
||||
)
|
||||
except UploadCleanupSafetyError:
|
||||
logger.exception("Upload cleanup aborted because index safety checks failed")
|
||||
raise HTTPException(
|
||||
503,
|
||||
"Upload cleanup aborted because upload index integrity could not be verified",
|
||||
)
|
||||
except Exception:
|
||||
logger.exception("Upload cleanup skipped because reference discovery failed")
|
||||
raise HTTPException(
|
||||
503,
|
||||
"Upload cleanup skipped because persisted references could not be verified",
|
||||
)
|
||||
return {"status": "success", "files_cleaned": cleaned_count}
|
||||
|
||||
@router.get("/stats")
|
||||
|
||||
+66
-12
@@ -3828,9 +3828,8 @@ async def stream_agent_loop(
|
||||
and _intent_match is not None
|
||||
and len(_intent_text) < 400
|
||||
and "```" not in _intent_text
|
||||
and _intent_nudge_count < _MAX_INTENT_NUDGES
|
||||
)
|
||||
if _looks_like_promise:
|
||||
if _looks_like_promise and _intent_nudge_count < _MAX_INTENT_NUDGES:
|
||||
_intent_nudge_count += 1
|
||||
_matched_phrase = _intent_match.group(0).strip()
|
||||
logger.info(f"[agent] intent-without-action nudge #{_intent_nudge_count} on round {round_num}: {_matched_phrase!r}")
|
||||
@@ -3859,6 +3858,31 @@ async def stream_agent_loop(
|
||||
# Visible signal in the stream so the user knows we caught it.
|
||||
yield f'data: {json.dumps({"type": "agent_step", "round": round_num + 1})}\n\n'
|
||||
continue
|
||||
if _looks_like_promise:
|
||||
_matched_phrase = _intent_match.group(0).strip()
|
||||
_guard_message = (
|
||||
"The agent stopped because it repeatedly announced a tool "
|
||||
"action without making the tool call."
|
||||
)
|
||||
logger.warning(
|
||||
"[agent] intent-without-action guard exhausted on round %d after %d nudges: %r",
|
||||
round_num,
|
||||
_intent_nudge_count,
|
||||
_matched_phrase,
|
||||
)
|
||||
yield (
|
||||
"data: "
|
||||
+ json.dumps({
|
||||
"type": "intent_nudge_exhausted",
|
||||
"reason": "intent_without_action_nudge_cap",
|
||||
"message": _guard_message,
|
||||
"round": round_num,
|
||||
"nudges": _intent_nudge_count,
|
||||
"matched": _matched_phrase,
|
||||
})
|
||||
+ "\n\n"
|
||||
)
|
||||
break
|
||||
break # no tools — done
|
||||
|
||||
# ── Loop-breaker (Terminus-style stall detector) ──────────────
|
||||
@@ -3895,6 +3919,21 @@ async def stream_agent_loop(
|
||||
reason = (f"calling {_runaway} with identical arguments over and over" if _runaway
|
||||
else "repeating the same tool calls without new progress")
|
||||
logger.warning(f"[agent] loop-breaker tripped on round {round_num} ({reason}); sig={_sig[:80]!r}")
|
||||
yield (
|
||||
"data: "
|
||||
+ json.dumps({
|
||||
"type": "loop_breaker_triggered",
|
||||
"reason": "loop_breaker_stall",
|
||||
"message": (
|
||||
"The loop-breaker detected repeated tool calls without "
|
||||
"new progress, so the agent is being forced to stop "
|
||||
"using tools and give its best final answer."
|
||||
),
|
||||
"round": round_num,
|
||||
"detail": reason,
|
||||
})
|
||||
+ "\n\n"
|
||||
)
|
||||
# The model has been executing tools, so its results are already
|
||||
# in context. Force ONE tool-free round to converge: write the
|
||||
# answer from what it has, or state plainly what's blocking it.
|
||||
@@ -4013,16 +4052,31 @@ async def stream_agent_loop(
|
||||
await _progress_q.put(None)
|
||||
|
||||
_tool_task = asyncio.create_task(_run_tool())
|
||||
# Drain progress events as they arrive — block until the
|
||||
# next event OR the tool finishes (sentinel = None).
|
||||
while True:
|
||||
evt = await _progress_q.get()
|
||||
if evt is None:
|
||||
break
|
||||
yield (
|
||||
f'data: {json.dumps({"type": "tool_progress", "tool": block.tool_type, "round": round_num, **evt})}\n\n'
|
||||
)
|
||||
desc, result = await _tool_task
|
||||
try:
|
||||
# Drain progress events as they arrive — block until the
|
||||
# next event OR the tool finishes (sentinel = None).
|
||||
while True:
|
||||
evt = await _progress_q.get()
|
||||
if evt is None:
|
||||
break
|
||||
yield (
|
||||
f'data: {json.dumps({"type": "tool_progress", "tool": block.tool_type, "round": round_num, **evt})}\n\n'
|
||||
)
|
||||
desc, result = await _tool_task
|
||||
finally:
|
||||
# If the SSE client disconnects (or this generator is
|
||||
# otherwise closed) while we're awaiting a progress event
|
||||
# above, GeneratorExit is thrown in right here and the
|
||||
# `await _tool_task` on the line above never runs — the
|
||||
# task (and any subprocess execute_tool_block spawned for
|
||||
# bash/python tools) would otherwise keep running
|
||||
# orphaned with nothing left to await or cancel it.
|
||||
if not _tool_task.done():
|
||||
_tool_task.cancel()
|
||||
try:
|
||||
await _tool_task
|
||||
except (asyncio.CancelledError, Exception):
|
||||
pass
|
||||
|
||||
# A skill the model just loaded can prescribe tools that weren't
|
||||
# RAG-selected this turn (declared via requires_toolsets in its
|
||||
|
||||
@@ -2,10 +2,16 @@ from typing import Any, Dict, List, Optional
|
||||
import logging
|
||||
import re
|
||||
from src.constants import MAX_READ_CHARS
|
||||
from src.tool_utils import _parse_tool_args
|
||||
from src.tool_utils import _parse_tool_args, get_upload_handler
|
||||
from src.upload_handler import reserve_upload_references
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _missing_document_upload(owner: Optional[str], content: Any) -> Optional[str]:
|
||||
"""Reserve explicit upload URLs before an agent persists document text."""
|
||||
return reserve_upload_references(get_upload_handler(), owner, content)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Active document state
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -384,6 +390,13 @@ class CreateDocumentTool:
|
||||
return {"error": "Cannot create document in another user's session"}
|
||||
_owner = _sess.owner if _sess else None
|
||||
|
||||
missing_id = _missing_document_upload(_owner, content)
|
||||
if missing_id:
|
||||
return {
|
||||
"error": f"Referenced upload is no longer available: {missing_id}",
|
||||
"exit_code": 1,
|
||||
}
|
||||
|
||||
doc = Document(
|
||||
id=doc_id,
|
||||
session_id=session_id,
|
||||
@@ -455,6 +468,13 @@ class UpdateDocumentTool:
|
||||
if is_email_doc:
|
||||
doc.language = "email"
|
||||
|
||||
missing_id = _missing_document_upload(owner, new_content)
|
||||
if missing_id:
|
||||
return {
|
||||
"error": f"Referenced upload is no longer available: {missing_id}",
|
||||
"exit_code": 1,
|
||||
}
|
||||
|
||||
if not is_email_doc and _pdf_source_upload_id(doc.current_content or ""):
|
||||
return _create_pdf_text_derivative(
|
||||
db,
|
||||
@@ -532,6 +552,12 @@ class EditDocumentTool:
|
||||
applied = 1
|
||||
skipped = max(0, len(edits) - 1)
|
||||
doc.language = "email"
|
||||
missing_id = _missing_document_upload(owner, updated_content)
|
||||
if missing_id:
|
||||
return {
|
||||
"error": f"Referenced upload is no longer available: {missing_id}",
|
||||
"exit_code": 1,
|
||||
}
|
||||
new_ver = doc.version_count + 1
|
||||
ver = DocumentVersion(
|
||||
id=str(uuid.uuid4()),
|
||||
@@ -584,6 +610,13 @@ class EditDocumentTool:
|
||||
if applied == 0:
|
||||
return {"error": f"No edits applied — none of the FIND blocks matched the document content (skipped {skipped})"}
|
||||
|
||||
missing_id = _missing_document_upload(owner, updated_content)
|
||||
if missing_id:
|
||||
return {
|
||||
"error": f"Referenced upload is no longer available: {missing_id}",
|
||||
"exit_code": 1,
|
||||
}
|
||||
|
||||
if _pdf_source_upload_id(doc.current_content or ""):
|
||||
return _create_pdf_text_derivative(
|
||||
db,
|
||||
|
||||
@@ -21,6 +21,7 @@ from src.model_discovery import ModelDiscovery
|
||||
from src.chat_handler import ChatHandler
|
||||
from src.research_handler import ResearchHandler
|
||||
from src.upload_handler import UploadHandler
|
||||
from src.tool_utils import set_upload_handler
|
||||
from src.search import update_search_config
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -49,6 +50,8 @@ def initialize_managers(base_dir: str, rag_manager=None) -> Dict[str, Any]:
|
||||
session_manager = SessionManager(SESSIONS_FILE)
|
||||
set_session_manager(session_manager) # Enable Session.add_message() persistence
|
||||
upload_handler = UploadHandler(base_dir, UPLOAD_DIR)
|
||||
session_manager.upload_handler = upload_handler
|
||||
set_upload_handler(upload_handler)
|
||||
personal_docs_manager = PersonalDocsManager(PERSONAL_DIR, rag_manager)
|
||||
api_key_manager = APIKeyManager(DATA_DIR)
|
||||
preset_manager = PresetManager(DATA_DIR)
|
||||
|
||||
@@ -0,0 +1,164 @@
|
||||
"""Attachment reference helpers for chat storage and tool manifests.
|
||||
|
||||
Live model calls may need provider-specific data URLs for the current turn.
|
||||
Persisted history and search indexes should keep stable upload references and
|
||||
human-readable text instead of duplicating raw media bytes.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import re
|
||||
from typing import Any, Iterable
|
||||
|
||||
|
||||
DATA_URL_RE = re.compile(
|
||||
r"data:[^;,\s\"']+;base64,[A-Za-z0-9+/=]+",
|
||||
re.IGNORECASE,
|
||||
)
|
||||
|
||||
MEDIA_BLOCK_TYPES = {
|
||||
"image",
|
||||
"image_url",
|
||||
"input_image",
|
||||
"audio",
|
||||
"input_audio",
|
||||
"file",
|
||||
}
|
||||
|
||||
|
||||
def strip_inline_data_urls(text: str) -> str:
|
||||
"""Replace inline data URLs with a compact marker."""
|
||||
if not isinstance(text, str) or ";base64," not in text:
|
||||
return text
|
||||
return DATA_URL_RE.sub("[inline media omitted from persisted history]", text)
|
||||
|
||||
|
||||
def attachment_ref(info: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Return the stable attachment reference shape used outside raw uploads."""
|
||||
upload_id = str(info.get("id") or info.get("attachment_id") or "").strip()
|
||||
try:
|
||||
size = int(info.get("size") or 0)
|
||||
except (TypeError, ValueError):
|
||||
size = 0
|
||||
ref = {
|
||||
"type": "attachment_ref",
|
||||
"attachment_id": upload_id,
|
||||
"name": info.get("name") or info.get("original_name") or upload_id,
|
||||
"mime": info.get("mime") or "application/octet-stream",
|
||||
"size": size,
|
||||
}
|
||||
checksum = info.get("checksum_sha256") or info.get("sha256") or info.get("hash")
|
||||
if checksum:
|
||||
ref["checksum_sha256"] = checksum
|
||||
created_at = info.get("created_at") or info.get("uploaded_at")
|
||||
if created_at:
|
||||
ref["created_at"] = created_at
|
||||
for key in ("width", "height", "vision", "vision_model", "gallery_id"):
|
||||
value = info.get(key)
|
||||
if value is not None:
|
||||
ref[key] = value
|
||||
return ref
|
||||
|
||||
|
||||
def attachment_refs_from_metadata(metadata: dict[str, Any] | None) -> list[dict[str, Any]]:
|
||||
"""Extract attachment refs from message metadata."""
|
||||
attachments = (metadata or {}).get("attachments") or []
|
||||
if not isinstance(attachments, list):
|
||||
return []
|
||||
refs: list[dict[str, Any]] = []
|
||||
for item in attachments:
|
||||
if isinstance(item, dict):
|
||||
ref = attachment_ref(item)
|
||||
if ref.get("attachment_id"):
|
||||
refs.append(ref)
|
||||
return refs
|
||||
|
||||
|
||||
def _ref_line(ref: dict[str, Any]) -> str:
|
||||
parts = [f"Attachment: {ref.get('name') or ref.get('attachment_id') or 'upload'}"]
|
||||
if ref.get("attachment_id"):
|
||||
parts.append(f"id={ref['attachment_id']}")
|
||||
if ref.get("mime"):
|
||||
parts.append(f"mime={ref['mime']}")
|
||||
if ref.get("size"):
|
||||
parts.append(f"size={ref['size']} bytes")
|
||||
if ref.get("checksum_sha256"):
|
||||
parts.append(f"sha256={ref['checksum_sha256']}")
|
||||
line = "[" + " | ".join(parts) + "]"
|
||||
if ref.get("vision"):
|
||||
line += f"\n[Attachment description: {str(ref['vision']).strip()}]"
|
||||
return line
|
||||
|
||||
|
||||
def _text_from_blocks(blocks: Iterable[Any]) -> str:
|
||||
lines: list[str] = []
|
||||
omitted_media = 0
|
||||
for block in blocks:
|
||||
if isinstance(block, str):
|
||||
lines.append(strip_inline_data_urls(block))
|
||||
continue
|
||||
if not isinstance(block, dict):
|
||||
continue
|
||||
block_type = block.get("type")
|
||||
if block_type == "text":
|
||||
text = block.get("text")
|
||||
if isinstance(text, str) and text:
|
||||
lines.append(strip_inline_data_urls(text))
|
||||
elif block_type == "attachment_ref":
|
||||
lines.append(_ref_line(block))
|
||||
elif block_type in MEDIA_BLOCK_TYPES:
|
||||
omitted_media += 1
|
||||
else:
|
||||
try:
|
||||
encoded = json.dumps(block, ensure_ascii=True, sort_keys=True)
|
||||
except TypeError:
|
||||
encoded = str(block)
|
||||
lines.append(strip_inline_data_urls(encoded))
|
||||
if omitted_media:
|
||||
plural = "s" if omitted_media != 1 else ""
|
||||
lines.append(f"[{omitted_media} inline media payload{plural} omitted]")
|
||||
return "\n".join(line for line in lines if line).strip()
|
||||
|
||||
|
||||
def persistable_message_content(
|
||||
content: Any,
|
||||
metadata: dict[str, Any] | None = None,
|
||||
) -> str:
|
||||
"""Return content safe for DB persistence and FTS indexing.
|
||||
|
||||
Multimodal provider blocks are collapsed to readable text plus stable
|
||||
attachment reference lines from metadata. This avoids storing base64 media
|
||||
in ``chat_messages.content`` while preserving enough context for reloads,
|
||||
search, and later turns.
|
||||
"""
|
||||
if isinstance(content, list):
|
||||
text = _text_from_blocks(content)
|
||||
refs = attachment_refs_from_metadata(metadata)
|
||||
ref_lines = [_ref_line(ref) for ref in refs]
|
||||
if ref_lines:
|
||||
text = "\n".join([part for part in (text, "\n".join(ref_lines)) if part]).strip()
|
||||
return text
|
||||
if isinstance(content, str):
|
||||
return strip_inline_data_urls(content)
|
||||
try:
|
||||
return strip_inline_data_urls(json.dumps(content, ensure_ascii=True, sort_keys=True))
|
||||
except TypeError:
|
||||
return strip_inline_data_urls(str(content))
|
||||
|
||||
|
||||
def search_index_text(content: Any) -> str:
|
||||
"""Best-effort searchable text for legacy stored content."""
|
||||
if isinstance(content, str):
|
||||
raw = content.strip()
|
||||
if raw.startswith("[") and '"type"' in raw:
|
||||
try:
|
||||
parsed = json.loads(content)
|
||||
except (TypeError, ValueError):
|
||||
parsed = None
|
||||
if isinstance(parsed, list):
|
||||
return _text_from_blocks(parsed)
|
||||
return strip_inline_data_urls(content)
|
||||
if isinstance(content, list):
|
||||
return _text_from_blocks(content)
|
||||
return persistable_message_content(content)
|
||||
@@ -1125,14 +1125,14 @@ async def action_learn_sender_signatures(owner: str, **kwargs) -> Tuple[str, boo
|
||||
conn = _imap_connect(None, owner=owner)
|
||||
try:
|
||||
conn.select("INBOX", readonly=True)
|
||||
status, data = conn.search(None, "ALL")
|
||||
status, data = conn.uid("SEARCH", None, "ALL")
|
||||
if status != "OK" or not data or not data[0]:
|
||||
return results
|
||||
uids = data[0].split()[-300:][::-1] # newest 300
|
||||
for uid in uids:
|
||||
try:
|
||||
st, msg_data = conn.fetch(
|
||||
uid, "(BODY.PEEK[HEADER.FIELDS (FROM)])"
|
||||
st, msg_data = conn.uid(
|
||||
"FETCH", uid, "(BODY.PEEK[HEADER.FIELDS (FROM)])"
|
||||
)
|
||||
if st != "OK" or not msg_data or not msg_data[0]:
|
||||
continue
|
||||
@@ -1214,7 +1214,7 @@ async def action_learn_sender_signatures(owner: str, **kwargs) -> Tuple[str, boo
|
||||
conn2.select("INBOX", readonly=True)
|
||||
for mm in _msgs:
|
||||
try:
|
||||
st, data = conn2.fetch(mm["uid"], "(BODY.PEEK[TEXT])")
|
||||
st, data = conn2.uid("FETCH", mm["uid"], "(BODY.PEEK[TEXT])")
|
||||
if st != "OK" or not data or not data[0]:
|
||||
continue
|
||||
raw = data[0][1] if isinstance(data[0], tuple) else None
|
||||
@@ -1356,13 +1356,13 @@ async def action_daily_brief(owner: str, **kwargs) -> Tuple[str, bool]:
|
||||
conn = _imap_connect(None)
|
||||
try:
|
||||
conn.select("INBOX", readonly=True)
|
||||
status, data = conn.search(None, "UNSEEN")
|
||||
status, data = conn.uid("SEARCH", None, "UNSEEN")
|
||||
uids = (data[0].split() if status == "OK" and data and data[0] else [])
|
||||
unread_count = len(uids)
|
||||
# Grab headers for the most recent 5 unread (UIDs increase with arrival)
|
||||
for uid in uids[-5:][::-1]:
|
||||
try:
|
||||
_, msg_data = conn.fetch(uid, "(BODY.PEEK[HEADER.FIELDS (FROM SUBJECT)])")
|
||||
_, msg_data = conn.uid("FETCH", uid, "(BODY.PEEK[HEADER.FIELDS (FROM SUBJECT)])")
|
||||
if not msg_data or not msg_data[0]:
|
||||
continue
|
||||
hdr = msg_data[0][1] if isinstance(msg_data[0], tuple) else msg_data[0]
|
||||
|
||||
+16
-1
@@ -104,6 +104,21 @@ def _spawn_bg(coro) -> asyncio.Task:
|
||||
return task
|
||||
|
||||
|
||||
def builtin_python_env(base_dir: str) -> dict[str, str]:
|
||||
"""Environment for built-in Python MCP subprocesses.
|
||||
|
||||
The app root must be importable so mcp_servers can import local modules, but
|
||||
replacing PYTHONPATH entirely hides site-packages in container/dev launches
|
||||
that rely on PYTHONPATH for their active environment.
|
||||
"""
|
||||
existing = os.environ.get("PYTHONPATH", "")
|
||||
parts = [base_dir]
|
||||
for item in existing.split(os.pathsep):
|
||||
if item and item not in parts:
|
||||
parts.append(item)
|
||||
return {"PYTHONPATH": os.pathsep.join(parts)}
|
||||
|
||||
|
||||
async def register_builtin_servers(mcp_manager):
|
||||
"""Connect all built-in MCP servers to the manager."""
|
||||
if MCP_DISABLED:
|
||||
@@ -121,7 +136,7 @@ async def register_builtin_servers(mcp_manager):
|
||||
transport="stdio",
|
||||
command=python,
|
||||
args=[script_path],
|
||||
env={"PYTHONPATH": base_dir},
|
||||
env=builtin_python_env(base_dir),
|
||||
)
|
||||
if ok:
|
||||
logger.info(f"Built-in MCP server registered: {name}")
|
||||
|
||||
+190
-188
@@ -280,214 +280,216 @@ def _sync_blocking(owner: str, url: str, username: str, password: str, account_i
|
||||
result = {"calendars": 0, "events": 0, "deleted": 0, "errors": []}
|
||||
|
||||
client = _build_dav_client(url, username, password)
|
||||
|
||||
# Discovery: try principal → calendars first; if the server doesn't
|
||||
# support discovery (or the URL points directly at a calendar), fall
|
||||
# back to treating the URL as a single calendar.
|
||||
calendars = []
|
||||
try:
|
||||
principal = client.principal()
|
||||
calendars = principal.calendars()
|
||||
except (AuthorizationError, NotFoundError) as e:
|
||||
result["errors"].append(f"Discovery failed: {e}")
|
||||
return result
|
||||
except Exception as e:
|
||||
logger.info(f"CalDAV principal discovery failed, trying URL as calendar: {e}")
|
||||
# Discovery: try principal → calendars first; if the server doesn't
|
||||
# support discovery (or the URL points directly at a calendar), fall
|
||||
# back to treating the URL as a single calendar.
|
||||
calendars = []
|
||||
try:
|
||||
calendars = [_open_url_as_calendar(client, url)]
|
||||
except Exception as e2:
|
||||
result["errors"].append(f"Could not open URL as calendar: {e2}")
|
||||
return result
|
||||
|
||||
if not calendars:
|
||||
try:
|
||||
calendars = [_open_url_as_calendar(client, url)]
|
||||
principal = client.principal()
|
||||
calendars = principal.calendars()
|
||||
except (AuthorizationError, NotFoundError) as e:
|
||||
result["errors"].append(f"Discovery failed: {e}")
|
||||
return result # outer finally will call client.close()
|
||||
except Exception as e:
|
||||
result["errors"].append(f"No calendars and URL fallback failed: {e}")
|
||||
return result
|
||||
|
||||
start = datetime.utcnow() - timedelta(days=_LOOKBACK_DAYS)
|
||||
end = datetime.utcnow() + timedelta(days=_LOOKAHEAD_DAYS)
|
||||
|
||||
db = SessionLocal()
|
||||
try:
|
||||
for remote_cal in calendars:
|
||||
logger.info(f"CalDAV principal discovery failed, trying URL as calendar: {e}")
|
||||
try:
|
||||
remote_url = str(remote_cal.url)
|
||||
cal_id = _stable_cal_id(remote_url, owner=owner, account_id=account_id)
|
||||
display_name = (remote_cal.name or "").strip() or "CalDAV"
|
||||
calendars = [_open_url_as_calendar(client, url)]
|
||||
except Exception as e2:
|
||||
result["errors"].append(f"Could not open URL as calendar: {e2}")
|
||||
return result # outer finally will call client.close()
|
||||
|
||||
local_cal = db.query(CalendarCal).filter(
|
||||
CalendarCal.id == cal_id,
|
||||
CalendarCal.owner == owner,
|
||||
).first()
|
||||
if not local_cal:
|
||||
local_cal = CalendarCal(
|
||||
id=cal_id,
|
||||
owner=owner,
|
||||
name=display_name,
|
||||
color="#5b8abf",
|
||||
source="caldav",
|
||||
account_id=account_id or None,
|
||||
caldav_base_url=remote_url,
|
||||
)
|
||||
db.add(local_cal)
|
||||
db.commit()
|
||||
else:
|
||||
# Refresh display name and stamp CalDAV metadata if missing.
|
||||
changed = False
|
||||
if local_cal.name != display_name:
|
||||
local_cal.name = display_name
|
||||
changed = True
|
||||
if account_id and not local_cal.account_id:
|
||||
local_cal.account_id = account_id
|
||||
changed = True
|
||||
if local_cal.caldav_base_url != remote_url:
|
||||
local_cal.caldav_base_url = remote_url
|
||||
changed = True
|
||||
if changed:
|
||||
db.commit()
|
||||
result["calendars"] += 1
|
||||
if not calendars:
|
||||
try:
|
||||
calendars = [_open_url_as_calendar(client, url)]
|
||||
except Exception as e:
|
||||
result["errors"].append(f"No calendars and URL fallback failed: {e}")
|
||||
return result # outer finally will call client.close()
|
||||
|
||||
# Fetch events in window. `date_search` returns CalendarObject
|
||||
# resources; each may contain one VEVENT (most servers) or
|
||||
# several (rare).
|
||||
from icalendar import Calendar as iCal
|
||||
start = datetime.utcnow() - timedelta(days=_LOOKBACK_DAYS)
|
||||
end = datetime.utcnow() + timedelta(days=_LOOKAHEAD_DAYS)
|
||||
|
||||
seen_uids = set()
|
||||
# Track events added to the session but not yet committed so
|
||||
# duplicate UIDs within the same batch are updated, not re-inserted
|
||||
# (which would violate the UNIQUE constraint on commit).
|
||||
pending: dict = {}
|
||||
parse_failed = False
|
||||
db = SessionLocal() # if this raises, outer finally still calls client.close()
|
||||
try:
|
||||
for remote_cal in calendars:
|
||||
try:
|
||||
objs = remote_cal.date_search(start=start, end=end, expand=False)
|
||||
except Exception as e:
|
||||
result["errors"].append(f"{display_name}: date_search failed ({e})")
|
||||
continue
|
||||
remote_url = str(remote_cal.url)
|
||||
cal_id = _stable_cal_id(remote_url, owner=owner, account_id=account_id)
|
||||
display_name = (remote_cal.name or "").strip() or "CalDAV"
|
||||
|
||||
for obj in objs:
|
||||
local_cal = db.query(CalendarCal).filter(
|
||||
CalendarCal.id == cal_id,
|
||||
CalendarCal.owner == owner,
|
||||
).first()
|
||||
if not local_cal:
|
||||
local_cal = CalendarCal(
|
||||
id=cal_id,
|
||||
owner=owner,
|
||||
name=display_name,
|
||||
color="#5b8abf",
|
||||
source="caldav",
|
||||
account_id=account_id or None,
|
||||
caldav_base_url=remote_url,
|
||||
)
|
||||
db.add(local_cal)
|
||||
db.commit()
|
||||
else:
|
||||
# Refresh display name and stamp CalDAV metadata if missing.
|
||||
changed = False
|
||||
if local_cal.name != display_name:
|
||||
local_cal.name = display_name
|
||||
changed = True
|
||||
if account_id and not local_cal.account_id:
|
||||
local_cal.account_id = account_id
|
||||
changed = True
|
||||
if local_cal.caldav_base_url != remote_url:
|
||||
local_cal.caldav_base_url = remote_url
|
||||
changed = True
|
||||
if changed:
|
||||
db.commit()
|
||||
result["calendars"] += 1
|
||||
|
||||
# Fetch events in window. `date_search` returns CalendarObject
|
||||
# resources; each may contain one VEVENT (most servers) or
|
||||
# several (rare).
|
||||
from icalendar import Calendar as iCal
|
||||
|
||||
seen_uids = set()
|
||||
# Track events added to the session but not yet committed so
|
||||
# duplicate UIDs within the same batch are updated, not re-inserted
|
||||
# (which would violates the UNIQUE constraint on commit).
|
||||
pending: dict = {}
|
||||
parse_failed = False
|
||||
try:
|
||||
ical = iCal.from_ical(obj.data)
|
||||
objs = remote_cal.date_search(start=start, end=end, expand=False)
|
||||
except Exception as e:
|
||||
result["errors"].append(f"{display_name}: parse failed ({e})")
|
||||
parse_failed = True
|
||||
result["errors"].append(f"{display_name}: date_search failed ({e})")
|
||||
continue
|
||||
|
||||
for comp in ical.walk():
|
||||
if comp.name != "VEVENT":
|
||||
for obj in objs:
|
||||
try:
|
||||
ical = iCal.from_ical(obj.data)
|
||||
except Exception as e:
|
||||
result["errors"].append(f"{display_name}: parse failed ({e})")
|
||||
parse_failed = True
|
||||
continue
|
||||
uid_val = str(comp.get("uid", "")) or str(uuid.uuid4())
|
||||
seen_uids.add(uid_val)
|
||||
|
||||
dtstart_p = comp.get("dtstart")
|
||||
if not dtstart_p:
|
||||
continue
|
||||
start_dt, all_day = _to_utc_naive(dtstart_p.dt)
|
||||
|
||||
dtend_p = comp.get("dtend")
|
||||
if dtend_p:
|
||||
end_dt, _ = _to_utc_naive(dtend_p.dt)
|
||||
elif all_day:
|
||||
end_dt = start_dt + timedelta(days=1)
|
||||
else:
|
||||
end_dt = start_dt + timedelta(hours=1)
|
||||
# A synced event with DTEND <= DTSTART (e.g. a single-day
|
||||
# all-day event whose source wrote DTEND equal to DTSTART)
|
||||
# would be stored zero-duration and silently dropped by the
|
||||
# list_events overlap filter. Clamp to a positive span.
|
||||
end_dt = _ensure_positive_duration(start_dt, end_dt, all_day)
|
||||
|
||||
# is_utc reflects whether the source carried a TZ
|
||||
# we converted from. All-day = no TZ semantics.
|
||||
row_is_utc = (
|
||||
not all_day
|
||||
and isinstance(dtstart_p.dt, datetime)
|
||||
and dtstart_p.dt.tzinfo is not None
|
||||
)
|
||||
|
||||
summary = str(comp.get("summary", ""))
|
||||
description = str(comp.get("description", ""))
|
||||
location = str(comp.get("location", ""))
|
||||
rrule = (
|
||||
comp.get("rrule").to_ical().decode()
|
||||
if comp.get("rrule")
|
||||
else ""
|
||||
)
|
||||
|
||||
existing = _find_existing_event(db, pending, uid_val, local_cal.id)
|
||||
if existing:
|
||||
if existing.caldav_sync_pending in {"create", "update"}:
|
||||
result["events"] += 1
|
||||
for comp in ical.walk():
|
||||
if comp.name != "VEVENT":
|
||||
continue
|
||||
existing.calendar_id = local_cal.id
|
||||
existing.summary = summary
|
||||
existing.description = description
|
||||
existing.location = location
|
||||
existing.dtstart = start_dt
|
||||
existing.dtend = end_dt
|
||||
existing.all_day = all_day
|
||||
existing.is_utc = row_is_utc
|
||||
existing.rrule = rrule
|
||||
existing.origin = "caldav"
|
||||
existing.remote_href = str(getattr(obj, "url", "") or "") or None
|
||||
existing.remote_etag = _event_etag(obj) or None
|
||||
existing.caldav_sync_pending = None
|
||||
else:
|
||||
new_ev = CalendarEvent(
|
||||
uid=uid_val,
|
||||
calendar_id=local_cal.id,
|
||||
summary=summary,
|
||||
description=description,
|
||||
location=location,
|
||||
dtstart=start_dt,
|
||||
dtend=end_dt,
|
||||
all_day=all_day,
|
||||
is_utc=row_is_utc,
|
||||
rrule=rrule,
|
||||
origin="caldav",
|
||||
remote_href=str(getattr(obj, "url", "") or "") or None,
|
||||
remote_etag=_event_etag(obj) or None,
|
||||
uid_val = str(comp.get("uid", "")) or str(uuid.uuid4())
|
||||
seen_uids.add(uid_val)
|
||||
|
||||
dtstart_p = comp.get("dtstart")
|
||||
if not dtstart_p:
|
||||
continue
|
||||
start_dt, all_day = _to_utc_naive(dtstart_p.dt)
|
||||
|
||||
dtend_p = comp.get("dtend")
|
||||
if dtend_p:
|
||||
end_dt, _ = _to_utc_naive(dtend_p.dt)
|
||||
elif all_day:
|
||||
end_dt = start_dt + timedelta(days=1)
|
||||
else:
|
||||
end_dt = start_dt + timedelta(hours=1)
|
||||
# A synced event with DTEND <= DTSTART (e.g. a single-day
|
||||
# all-day event whose source wrote DTEND equal to DTSTART)
|
||||
# would be stored zero-duration and silently dropped by the
|
||||
# list_events overlap filter. Clamp to a positive span.
|
||||
end_dt = _ensure_positive_duration(start_dt, end_dt, all_day)
|
||||
|
||||
# is_utc reflects whether the source carried a TZ
|
||||
# we converted from. All-day = no TZ semantics.
|
||||
row_is_utc = (
|
||||
not all_day
|
||||
and isinstance(dtstart_p.dt, datetime)
|
||||
and dtstart_p.dt.tzinfo is not None
|
||||
)
|
||||
db.add(new_ev)
|
||||
pending[uid_val] = new_ev
|
||||
result["events"] += 1
|
||||
db.commit()
|
||||
|
||||
# Prune locally-cached CalDAV events that vanished
|
||||
# upstream (only within our sync window — events outside
|
||||
# the window aren't in `objs`, so we'd false-delete them).
|
||||
# Only rows we previously pulled from the server (origin=="caldav")
|
||||
# are prunable; locally-created events (agent / email triage / a
|
||||
# UI event whose write-back failed) carry origin NULL and must
|
||||
# never be deleted just because the server didn't return them.
|
||||
# Skip the prune on any parse failure: seen_uids is then an
|
||||
# incomplete view of the server, so pruning against it would
|
||||
# delete events that still exist upstream but could not be read
|
||||
# (the empty-seen_uids case wipes the whole window; a partial
|
||||
# failure deletes just the unreadable rows).
|
||||
if _should_prune_window(seen_uids, parse_failed):
|
||||
stale = db.query(CalendarEvent).filter(
|
||||
CalendarEvent.calendar_id == local_cal.id,
|
||||
CalendarEvent.origin == "caldav",
|
||||
CalendarEvent.dtstart >= start,
|
||||
CalendarEvent.dtstart <= end,
|
||||
CalendarEvent.remote_href.isnot(None),
|
||||
CalendarEvent.caldav_sync_pending.is_(None),
|
||||
~CalendarEvent.uid.in_(seen_uids) if seen_uids else CalendarEvent.uid.isnot(None),
|
||||
).all()
|
||||
for ev in stale:
|
||||
db.delete(ev)
|
||||
result["deleted"] += len(stale)
|
||||
summary = str(comp.get("summary", ""))
|
||||
description = str(comp.get("description", ""))
|
||||
location = str(comp.get("location", ""))
|
||||
rrule = (
|
||||
comp.get("rrule").to_ical().decode()
|
||||
if comp.get("rrule")
|
||||
else ""
|
||||
)
|
||||
|
||||
existing = _find_existing_event(db, pending, uid_val, local_cal.id)
|
||||
if existing:
|
||||
if existing.caldav_sync_pending in {"create", "update"}:
|
||||
result["events"] += 1
|
||||
continue
|
||||
existing.calendar_id = local_cal.id
|
||||
existing.summary = summary
|
||||
existing.description = description
|
||||
existing.location = location
|
||||
existing.dtstart = start_dt
|
||||
existing.dtend = end_dt
|
||||
existing.all_day = all_day
|
||||
existing.is_utc = row_is_utc
|
||||
existing.rrule = rrule
|
||||
existing.origin = "caldav"
|
||||
existing.remote_href = str(getattr(obj, "url", "") or "") or None
|
||||
existing.remote_etag = _event_etag(obj) or None
|
||||
existing.caldav_sync_pending = None
|
||||
else:
|
||||
new_ev = CalendarEvent(
|
||||
uid=uid_val,
|
||||
calendar_id=local_cal.id,
|
||||
summary=summary,
|
||||
description=description,
|
||||
location=location,
|
||||
dtstart=start_dt,
|
||||
dtend=end_dt,
|
||||
all_day=all_day,
|
||||
is_utc=row_is_utc,
|
||||
rrule=rrule,
|
||||
origin="caldav",
|
||||
remote_href=str(getattr(obj, "url", "") or "") or None,
|
||||
remote_etag=_event_etag(obj) or None,
|
||||
)
|
||||
db.add(new_ev)
|
||||
pending[uid_val] = new_ev
|
||||
result["events"] += 1
|
||||
db.commit()
|
||||
except Exception as e:
|
||||
logger.exception("CalDAV sync failed for one calendar")
|
||||
result["errors"].append(str(e)[:200])
|
||||
db.rollback()
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
return result
|
||||
# Prune locally-cached CalDAV events that vanished
|
||||
# upstream (only within our sync window — events outside
|
||||
# the window aren't in `objs`, so we'd false-delete them).
|
||||
# Only rows we previously pulled from the server (origin=="caldav")
|
||||
# are prunable; locally-created events (agent / email triage / a
|
||||
# UI event whose write-back failed) carry origin NULL and must
|
||||
# never be deleted just because the server didn't return them.
|
||||
# Skip the prune on any parse failure: seen_uids is then an
|
||||
# incomplete view of the server, so pruning against it would
|
||||
# delete events that still exist upstream but could not be read
|
||||
# (the empty-seen_uids case wipes the whole window; a partial
|
||||
# failure deletes just the unreadable rows).
|
||||
if _should_prune_window(seen_uids, parse_failed):
|
||||
stale = db.query(CalendarEvent).filter(
|
||||
CalendarEvent.calendar_id == local_cal.id,
|
||||
CalendarEvent.origin == "caldav",
|
||||
CalendarEvent.dtstart >= start,
|
||||
CalendarEvent.dtstart <= end,
|
||||
CalendarEvent.remote_href.isnot(None),
|
||||
CalendarEvent.caldav_sync_pending.is_(None),
|
||||
~CalendarEvent.uid.in_(seen_uids) if seen_uids else CalendarEvent.uid.isnot(None),
|
||||
).all()
|
||||
for ev in stale:
|
||||
db.delete(ev)
|
||||
result["deleted"] += len(stale)
|
||||
db.commit()
|
||||
except Exception as e:
|
||||
logger.exception("CalDAV sync failed for one calendar")
|
||||
result["errors"].append(str(e)[:200])
|
||||
db.rollback()
|
||||
finally:
|
||||
db.close() # NOT client.close() here anymore
|
||||
|
||||
return result
|
||||
finally:
|
||||
client.close() # always called
|
||||
|
||||
|
||||
def _event_payload(ev) -> dict:
|
||||
|
||||
@@ -192,11 +192,14 @@ def _writeback_blocking(local_cal_id, ev, delete, url, username, password,
|
||||
# Redirects disabled here too: the write-back path opens its own DAVClient,
|
||||
# so it needs the same SSRF-via-redirect protection as the pull path.
|
||||
client = _build_dav_client(url, username, password)
|
||||
calendars = _discover_calendars(client)
|
||||
if not calendars:
|
||||
return {"ok": False, "error": "no remote calendars discovered"}
|
||||
return push_event(calendars, local_cal_id, ev, delete=delete,
|
||||
owner=owner, account_id=account_id)
|
||||
try:
|
||||
calendars = _discover_calendars(client)
|
||||
if not calendars:
|
||||
return {"ok": False, "error": "no remote calendars discovered"}
|
||||
return push_event(calendars, local_cal_id, ev, delete=delete,
|
||||
owner=owner, account_id=account_id)
|
||||
finally:
|
||||
client.close()
|
||||
|
||||
|
||||
def _persist_writeback_result(owner: str, calendar_id: str, uid: str, result: dict, *, delete: bool) -> None:
|
||||
|
||||
@@ -191,6 +191,8 @@ class ChatHandler:
|
||||
"name": fi.get("name") or fi.get("original_name") or fi["id"],
|
||||
"mime": fi.get("mime", ""),
|
||||
"size": fi.get("size", 0),
|
||||
"checksum_sha256": fi.get("checksum_sha256") or fi.get("hash"),
|
||||
"created_at": fi.get("created_at") or fi.get("uploaded_at"),
|
||||
"width": fi.get("width"),
|
||||
"height": fi.get("height"),
|
||||
})
|
||||
|
||||
+32
-2
@@ -766,6 +766,36 @@ def apply_kimi_code_headers(headers: Optional[Dict], url: str) -> Dict[str, str]
|
||||
return h
|
||||
|
||||
|
||||
async def apply_kimi_code_headers_async(client, headers: Optional[Dict], url: str) -> Dict[str, str]:
|
||||
"""Pick a Kimi Code User-Agent without blocking the event loop."""
|
||||
h = dict(headers or {})
|
||||
if not _is_kimi_code_url(url):
|
||||
return h
|
||||
base_key = _kimi_code_base_key(url)
|
||||
cached = _kimi_code_ua_cache.get(base_key)
|
||||
if cached:
|
||||
h["User-Agent"] = cached
|
||||
return h
|
||||
models_url = base_key.rstrip("/") + "/models"
|
||||
for ua in KIMI_CODE_USER_AGENTS:
|
||||
trial = dict(h)
|
||||
trial["User-Agent"] = ua
|
||||
try:
|
||||
r = await client.get(models_url, headers=trial, timeout=8)
|
||||
except Exception:
|
||||
continue
|
||||
if _is_kimi_code_access_denied(r.status_code, r.content):
|
||||
logger.debug("Kimi Code rejected User-Agent %s (403), trying next", ua)
|
||||
continue
|
||||
if r.status_code < 400:
|
||||
_remember_kimi_code_user_agent(url, ua)
|
||||
h["User-Agent"] = ua
|
||||
return h
|
||||
break
|
||||
h.setdefault("User-Agent", KIMI_CODE_USER_AGENT)
|
||||
return h
|
||||
|
||||
|
||||
def httpx_get_kimi_aware(url: str, headers: Optional[Dict], **kwargs):
|
||||
h = apply_kimi_code_headers(headers, url)
|
||||
if not _is_kimi_code_url(url):
|
||||
@@ -799,7 +829,7 @@ def httpx_post_kimi_aware(url: str, headers: Optional[Dict], **kwargs):
|
||||
|
||||
|
||||
async def httpx_post_kimi_aware_async(client, url: str, headers: Optional[Dict], **kwargs):
|
||||
h = apply_kimi_code_headers(headers, url)
|
||||
h = await apply_kimi_code_headers_async(client, headers, url)
|
||||
if not _is_kimi_code_url(url):
|
||||
return await client.post(url, headers=h, **kwargs)
|
||||
last = None
|
||||
@@ -2466,9 +2496,9 @@ async def _stream_llm_inner(url: str, model: str, messages: List[Dict], temperat
|
||||
events.append(_stream_delta_event(part))
|
||||
return events
|
||||
|
||||
h = apply_kimi_code_headers(h, target_url)
|
||||
try:
|
||||
client = _get_http_client()
|
||||
h = await apply_kimi_code_headers_async(client, h, target_url)
|
||||
async with client.stream('POST', target_url, json=payload, headers=h, timeout=stream_timeout) as r:
|
||||
_clear_host_dead(target_url)
|
||||
if r.status_code != 200:
|
||||
|
||||
+2
-2
@@ -504,7 +504,7 @@ class McpManager:
|
||||
async def _reconnect_builtin(self, server_id: str) -> bool:
|
||||
"""Tear down and reconnect a crashed builtin MCP server."""
|
||||
import sys
|
||||
from src.builtin_mcp import _BUILTIN_SERVERS
|
||||
from src.builtin_mcp import _BUILTIN_SERVERS, builtin_python_env
|
||||
|
||||
if server_id not in _BUILTIN_SERVERS:
|
||||
return False
|
||||
@@ -523,7 +523,7 @@ class McpManager:
|
||||
transport="stdio",
|
||||
command=sys.executable,
|
||||
args=[script_path],
|
||||
env={"PYTHONPATH": base_dir},
|
||||
env=builtin_python_env(base_dir),
|
||||
)
|
||||
if ok:
|
||||
logger.info(f"Reconnected builtin MCP server: {name}")
|
||||
|
||||
@@ -9,6 +9,7 @@ import json
|
||||
from src.constants import MAX_OUTPUT_CHARS
|
||||
|
||||
_mcp_manager = None
|
||||
_upload_handler = None
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# MCP Manager singleton
|
||||
@@ -23,6 +24,21 @@ def get_mcp_manager():
|
||||
"""Get the global MCP manager instance."""
|
||||
return _mcp_manager
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Shared upload lifecycle handler
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def set_upload_handler(handler):
|
||||
"""Register the process's UploadHandler without importing app modules."""
|
||||
global _upload_handler
|
||||
_upload_handler = handler
|
||||
|
||||
|
||||
def get_upload_handler():
|
||||
"""Return the shared UploadHandler used by route and agent writers."""
|
||||
return _upload_handler
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
+29
-2
@@ -10,6 +10,8 @@ import re
|
||||
from typing import Dict, Optional
|
||||
|
||||
from src.tools._common import _parse_tool_args
|
||||
from src.tool_utils import get_upload_handler
|
||||
from src.upload_handler import reserve_upload_references
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -408,11 +410,25 @@ async def do_manage_calendar(content: str, owner: Optional[str] = None) -> Dict:
|
||||
importance = args.get("importance") or "normal"
|
||||
minutes_before = _reminder_minutes(args)
|
||||
|
||||
event_description = _event_description(args, minutes_before)
|
||||
event_location = args.get("location", "") or ""
|
||||
missing_id = reserve_upload_references(
|
||||
get_upload_handler(),
|
||||
owner,
|
||||
event_description,
|
||||
event_location,
|
||||
)
|
||||
if missing_id:
|
||||
return {
|
||||
"error": f"Referenced upload is no longer available: {missing_id}",
|
||||
"exit_code": 1,
|
||||
}
|
||||
|
||||
uid = str(_uuid.uuid4())
|
||||
ev = CalendarEvent(
|
||||
uid=uid, calendar_id=cal.id, summary=summary,
|
||||
description=_event_description(args, minutes_before),
|
||||
location=args.get("location", "") or "",
|
||||
description=event_description,
|
||||
location=event_location,
|
||||
dtstart=dtstart, dtend=dtend, all_day=all_day,
|
||||
is_utc=dtstart_is_utc and not all_day,
|
||||
rrule=args.get("rrule", "") or "",
|
||||
@@ -465,6 +481,17 @@ async def do_manage_calendar(content: str, owner: Optional[str] = None) -> Dict:
|
||||
ev = _event_query().filter(CalendarEvent.uid == base_uid).first()
|
||||
if not ev:
|
||||
return {"error": f"Event {uid} not found", "exit_code": 1}
|
||||
missing_id = reserve_upload_references(
|
||||
get_upload_handler(),
|
||||
owner,
|
||||
args.get("description"),
|
||||
args.get("location"),
|
||||
)
|
||||
if missing_id:
|
||||
return {
|
||||
"error": f"Referenced upload is no longer available: {missing_id}",
|
||||
"exit_code": 1,
|
||||
}
|
||||
if args.get("summary") is not None:
|
||||
ev.summary = args["summary"]
|
||||
if args.get("description") is not None:
|
||||
|
||||
@@ -10,6 +10,8 @@ import re
|
||||
from typing import Dict, Optional
|
||||
|
||||
from src.tools._common import _parse_tool_args
|
||||
from src.tool_utils import get_upload_handler
|
||||
from src.upload_handler import reserve_upload_references
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -198,6 +200,18 @@ async def do_manage_notes(content: str, owner: Optional[str] = None) -> Dict:
|
||||
"duplicate": True,
|
||||
"exit_code": 0,
|
||||
}
|
||||
missing_id = reserve_upload_references(
|
||||
get_upload_handler(),
|
||||
owner,
|
||||
content_raw,
|
||||
args.get("color"),
|
||||
items_json,
|
||||
)
|
||||
if missing_id:
|
||||
return {
|
||||
"error": f"Referenced upload is no longer available: {missing_id}",
|
||||
"exit_code": 1,
|
||||
}
|
||||
note = Note(
|
||||
id=str(_uuid.uuid4()),
|
||||
owner=owner,
|
||||
@@ -235,6 +249,19 @@ async def do_manage_notes(content: str, owner: Optional[str] = None) -> Dict:
|
||||
return {"error": f"Note '{note_id}' not found", "exit_code": 1}
|
||||
if not _note_visible_to_owner(note, owner):
|
||||
return {"error": "Note not found", "exit_code": 1}
|
||||
missing_id = reserve_upload_references(
|
||||
get_upload_handler(),
|
||||
owner,
|
||||
args.get("content"),
|
||||
args.get("color"),
|
||||
args.get("checklist_items"),
|
||||
args.get("items"),
|
||||
)
|
||||
if missing_id:
|
||||
return {
|
||||
"error": f"Referenced upload is no longer available: {missing_id}",
|
||||
"exit_code": 1,
|
||||
}
|
||||
for field in ("title", "content", "note_type", "color", "label"):
|
||||
if field in args and args[field] is not None:
|
||||
setattr(note, field, args[field])
|
||||
|
||||
+652
-109
@@ -35,11 +35,34 @@ import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class UploadCleanupSafetyError(RuntimeError):
|
||||
"""Raised when cleanup cannot prove that destructive work is safe."""
|
||||
|
||||
# The extension is optional: save_upload builds the id as `{uuid.hex}{ext}`,
|
||||
# and a file with no extension (Dockerfile, README, ...) yields a bare 32-hex
|
||||
# id. Requiring `.ext` made those ids fail validation, so the stored file
|
||||
# could never be resolved or downloaded again.
|
||||
UPLOAD_ID_RE = re.compile(r"^[0-9a-fA-F]{32}(?:\.[A-Za-z0-9]+)?$")
|
||||
UPLOAD_ID_TOKEN_RE = re.compile(
|
||||
r"(?<![0-9a-fA-F])([0-9a-fA-F]{32}(?:\.[A-Za-z0-9]+)?)(?![A-Za-z0-9])"
|
||||
)
|
||||
INTERNAL_UPLOAD_URL_RE = re.compile(
|
||||
r"(?:odysseus://attachment/|/api/upload/)"
|
||||
r"([0-9a-fA-F]{32}(?:\.[A-Za-z0-9]+)?)"
|
||||
r"(?=$|[\s\"'<>\[\](){},;!?:&#]|\.(?![A-Za-z0-9]))"
|
||||
)
|
||||
PDF_SOURCE_UPLOAD_RE = re.compile(
|
||||
r"<!--\s*pdf(?:_form)?_source\b[^>]*\bupload_id="
|
||||
r"[\"']([0-9a-fA-F]{32}(?:\.[A-Za-z0-9]+)?)[\"'][^>]*-->",
|
||||
re.IGNORECASE,
|
||||
)
|
||||
ATTACHMENT_REFERENCE_LINE_RE = re.compile(
|
||||
r"\[Attachment:[^\]\r\n]*\|\s*id="
|
||||
r"([0-9a-fA-F]{32}(?:\.[A-Za-z0-9]+)?)"
|
||||
r"(?:\s*\||\s*\])",
|
||||
re.IGNORECASE,
|
||||
)
|
||||
|
||||
|
||||
def is_valid_upload_id(upload_id: str) -> bool:
|
||||
@@ -47,6 +70,110 @@ def is_valid_upload_id(upload_id: str) -> bool:
|
||||
return UPLOAD_ID_RE.fullmatch(upload_id or "") is not None
|
||||
|
||||
|
||||
def extract_upload_ids(value: Any) -> set[str]:
|
||||
"""Return canonical upload IDs embedded in a persisted URL/text value."""
|
||||
if not isinstance(value, str) or not value:
|
||||
return set()
|
||||
return set(UPLOAD_ID_TOKEN_RE.findall(value))
|
||||
|
||||
|
||||
def extract_internal_upload_ids(value: Any) -> set[str]:
|
||||
"""Return IDs from explicit internal upload references only.
|
||||
|
||||
Cleanup intentionally uses :func:`extract_upload_ids` conservatively, but
|
||||
write-time reservation must not treat an arbitrary 32-hex checksum in note
|
||||
or calendar text as an upload reference. Nested JSON-like values are
|
||||
supported because note checklist items are persisted as structured data.
|
||||
"""
|
||||
if isinstance(value, dict):
|
||||
found: set[str] = set()
|
||||
for nested in value.values():
|
||||
found.update(extract_internal_upload_ids(nested))
|
||||
return found
|
||||
if isinstance(value, (list, tuple, set)):
|
||||
found: set[str] = set()
|
||||
for nested in value:
|
||||
found.update(extract_internal_upload_ids(nested))
|
||||
return found
|
||||
if not isinstance(value, str) or not value:
|
||||
return set()
|
||||
return (
|
||||
set(INTERNAL_UPLOAD_URL_RE.findall(value))
|
||||
| set(PDF_SOURCE_UPLOAD_RE.findall(value))
|
||||
| set(ATTACHMENT_REFERENCE_LINE_RE.findall(value))
|
||||
)
|
||||
|
||||
|
||||
def reserve_upload_references(
|
||||
upload_handler: Any,
|
||||
owner: Optional[str],
|
||||
*values: Any,
|
||||
) -> Optional[str]:
|
||||
"""Reserve upload IDs in values before a caller persists references.
|
||||
|
||||
Returns the first ID that cannot be owner-checked/reserved, otherwise
|
||||
``None``. A missing handler is treated as no-op for backward-compatible
|
||||
route factories; production wires the shared UploadHandler instance.
|
||||
"""
|
||||
if upload_handler is None:
|
||||
return None
|
||||
upload_ids: set[str] = set()
|
||||
for value in values:
|
||||
upload_ids.update(extract_internal_upload_ids(value))
|
||||
return reserve_upload_ids(upload_handler, owner, upload_ids)
|
||||
|
||||
|
||||
def reserve_upload_ids(
|
||||
upload_handler: Any,
|
||||
owner: Optional[str],
|
||||
upload_ids: Any,
|
||||
) -> Optional[str]:
|
||||
"""Owner-reserve canonical IDs from a trusted structured reference field."""
|
||||
if upload_handler is None:
|
||||
return None
|
||||
canonical_ids = {
|
||||
str(upload_id).strip()
|
||||
for upload_id in (upload_ids or [])
|
||||
if is_valid_upload_id(str(upload_id).strip())
|
||||
}
|
||||
for upload_id in sorted(canonical_ids):
|
||||
try:
|
||||
resolved = upload_handler.reserve_upload(
|
||||
upload_id,
|
||||
owner=owner,
|
||||
allow_admin=False,
|
||||
)
|
||||
except Exception:
|
||||
resolved = None
|
||||
if not resolved:
|
||||
return upload_id
|
||||
return None
|
||||
|
||||
|
||||
def reserve_message_upload_references(
|
||||
upload_handler: Any,
|
||||
owner: Optional[str],
|
||||
content: Any,
|
||||
metadata: Any = None,
|
||||
) -> Optional[str]:
|
||||
"""Reserve explicit chat references, including structured attachment IDs."""
|
||||
upload_ids = extract_internal_upload_ids(content)
|
||||
if metadata not in (None, ""):
|
||||
if isinstance(metadata, str):
|
||||
metadata = json.loads(metadata)
|
||||
if not isinstance(metadata, dict):
|
||||
raise ValueError("message metadata must be a JSON object")
|
||||
upload_ids.update(extract_internal_upload_ids(metadata))
|
||||
from src.attachment_refs import attachment_refs_from_metadata
|
||||
|
||||
upload_ids.update(
|
||||
str(ref.get("attachment_id") or "").strip()
|
||||
for ref in attachment_refs_from_metadata(metadata)
|
||||
if ref.get("attachment_id")
|
||||
)
|
||||
return reserve_upload_ids(upload_handler, owner, upload_ids)
|
||||
|
||||
|
||||
def _build_upload_id(safe_filename: str) -> str:
|
||||
"""Build a unique upload id whose extension matches UPLOAD_ID_RE.
|
||||
|
||||
@@ -249,43 +376,295 @@ class UploadHandler:
|
||||
|
||||
return True
|
||||
|
||||
def cleanup_old_uploads(self):
|
||||
"""Remove uploaded files older than CLEANUP_DAYS days."""
|
||||
@staticmethod
|
||||
def _parse_upload_timestamp(value: Any) -> Optional[datetime]:
|
||||
if not isinstance(value, str) or not value.strip():
|
||||
return None
|
||||
try:
|
||||
cutoff_date = datetime.now() - timedelta(days=self.cleanup_days)
|
||||
parsed = datetime.fromisoformat(value.strip().replace("Z", "+00:00"))
|
||||
if parsed.tzinfo is not None:
|
||||
parsed = parsed.astimezone().replace(tzinfo=None)
|
||||
return parsed
|
||||
except (TypeError, ValueError):
|
||||
return None
|
||||
|
||||
@classmethod
|
||||
def _upload_metadata_is_recent(cls, info: Dict[str, Any], cutoff_date: datetime) -> bool:
|
||||
"""Return True when upload metadata records activity inside retention."""
|
||||
for field in ("last_accessed", "created_at", "uploaded_at"):
|
||||
parsed = cls._parse_upload_timestamp(info.get(field))
|
||||
if parsed is None:
|
||||
continue
|
||||
if parsed >= cutoff_date:
|
||||
return True
|
||||
return False
|
||||
|
||||
@classmethod
|
||||
def _upload_index_keys_for_file(
|
||||
cls,
|
||||
upload_index: Dict[str, Any],
|
||||
upload_id: str,
|
||||
file_path: str,
|
||||
) -> list[str]:
|
||||
"""Find a coherent set of index rows for one physical upload.
|
||||
|
||||
Every related row must agree on ID, canonical path, owner, and a
|
||||
non-empty checksum. Each row must also contain the complete lifecycle
|
||||
timestamps written for new uploads. Ambiguous or incomplete index
|
||||
state cannot authorize destructive cleanup.
|
||||
"""
|
||||
target_path = os.path.normcase(os.path.realpath(file_path))
|
||||
matches: list[str] = []
|
||||
owners: set[str] = set()
|
||||
checksums: set[str] = set()
|
||||
for key, info in upload_index.items():
|
||||
if not isinstance(info, dict):
|
||||
continue
|
||||
stored_path = info.get("path")
|
||||
stored_real_path = (
|
||||
os.path.normcase(os.path.realpath(stored_path))
|
||||
if isinstance(stored_path, str) and stored_path
|
||||
else None
|
||||
)
|
||||
same_id = info.get("id") == upload_id
|
||||
same_path = stored_real_path == target_path
|
||||
if not same_id and not same_path:
|
||||
continue
|
||||
if not same_id or not same_path:
|
||||
logger.warning(
|
||||
"Skipping ambiguous cleanup candidate %s: related row has id=%r path=%r",
|
||||
file_path,
|
||||
info.get("id"),
|
||||
stored_path,
|
||||
)
|
||||
return []
|
||||
|
||||
owner = info.get("owner")
|
||||
if not isinstance(owner, str) or not owner.strip():
|
||||
logger.warning(
|
||||
"Skipping incomplete cleanup candidate %s: matching row has no owner",
|
||||
file_path,
|
||||
)
|
||||
return []
|
||||
|
||||
row_checksums = {
|
||||
str(info.get(field)).strip().lower()
|
||||
for field in ("hash", "checksum_sha256")
|
||||
if info.get(field) is not None and str(info.get(field)).strip()
|
||||
}
|
||||
if not row_checksums:
|
||||
logger.warning(
|
||||
"Skipping incomplete cleanup candidate %s: matching row has no checksum",
|
||||
file_path,
|
||||
)
|
||||
return []
|
||||
if len(row_checksums) != 1:
|
||||
logger.warning(
|
||||
"Skipping ambiguous cleanup candidate %s: matching row has conflicting checksums",
|
||||
file_path,
|
||||
)
|
||||
return []
|
||||
|
||||
lifecycle_fields = ("uploaded_at", "created_at", "last_accessed")
|
||||
if any(
|
||||
cls._parse_upload_timestamp(info.get(field)) is None
|
||||
for field in lifecycle_fields
|
||||
):
|
||||
logger.warning(
|
||||
"Skipping incomplete cleanup candidate %s: matching row lacks lifecycle timestamps",
|
||||
file_path,
|
||||
)
|
||||
return []
|
||||
|
||||
matches.append(key)
|
||||
owners.add(owner)
|
||||
checksums.update(row_checksums)
|
||||
|
||||
if len(owners) > 1 or len(checksums) > 1:
|
||||
logger.warning(
|
||||
"Skipping ambiguous cleanup candidate %s: matching rows disagree on owner or checksum",
|
||||
file_path,
|
||||
)
|
||||
return []
|
||||
return matches
|
||||
|
||||
def cleanup_old_uploads(
|
||||
self,
|
||||
referenced_upload_ids: Optional[set[str]] = None,
|
||||
referenced_upload_hashes: Optional[set[str]] = None,
|
||||
):
|
||||
"""Remove expired uploads proven unreferenced by a complete snapshot.
|
||||
|
||||
``None`` means reference discovery was not completed, so cleanup fails
|
||||
closed and removes nothing. The admin route supplies both sets after
|
||||
scanning persisted chats, documents, and gallery records.
|
||||
"""
|
||||
if referenced_upload_ids is None or referenced_upload_hashes is None:
|
||||
logger.warning("Upload cleanup skipped: persisted reference snapshot unavailable")
|
||||
return 0
|
||||
|
||||
try:
|
||||
cleanup_started_at = datetime.now()
|
||||
cutoff_date = cleanup_started_at - timedelta(days=self.cleanup_days)
|
||||
cleaned_count = 0
|
||||
|
||||
for root, dirs, files in os.walk(self.upload_dir):
|
||||
if root == self.upload_dir:
|
||||
continue
|
||||
|
||||
path_parts = root.split(os.sep)
|
||||
if len(path_parts) >= 4:
|
||||
|
||||
referenced_ids = {str(value) for value in referenced_upload_ids}
|
||||
referenced_hashes = {str(value) for value in referenced_upload_hashes}
|
||||
uploads_db_path = os.path.join(self.upload_dir, "uploads.json")
|
||||
|
||||
# Keep index mutation and file removal serialized with upload writes.
|
||||
# Each row removal is atomically persisted before the bytes are
|
||||
# deleted; if deletion fails, the previous index is restored.
|
||||
with self._index_lock:
|
||||
current_index = dict(self._load_upload_index(fail_on_error=True))
|
||||
|
||||
for root, dirs, files in os.walk(self.upload_dir, followlinks=False):
|
||||
is_junction = getattr(os.path, "isjunction", lambda _path: False)
|
||||
dirs[:] = [
|
||||
directory
|
||||
for directory in dirs
|
||||
if not os.path.islink(os.path.join(root, directory))
|
||||
and not is_junction(os.path.join(root, directory))
|
||||
]
|
||||
if root == self.upload_dir:
|
||||
continue
|
||||
if not self._inside_upload_dir(root):
|
||||
dirs[:] = []
|
||||
continue
|
||||
|
||||
path_parts = root.split(os.sep)
|
||||
if len(path_parts) < 4:
|
||||
continue
|
||||
try:
|
||||
dir_date = datetime(int(path_parts[-3]), int(path_parts[-2]), int(path_parts[-1]))
|
||||
if dir_date < cutoff_date:
|
||||
for file in files:
|
||||
file_path = os.path.join(root, file)
|
||||
try:
|
||||
os.remove(file_path)
|
||||
cleaned_count += 1
|
||||
logger.info(f"Cleaned up old upload: {file_path}")
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to remove {file_path}: {e}")
|
||||
|
||||
try:
|
||||
os.rmdir(root)
|
||||
logger.info(f"Removed empty upload directory: {root}")
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to remove directory {root}: {e}")
|
||||
except (ValueError, IndexError):
|
||||
continue
|
||||
|
||||
if dir_date >= cutoff_date:
|
||||
continue
|
||||
|
||||
for file in files:
|
||||
# Reference discovery only understands canonical upload
|
||||
# IDs; unknown files fail closed instead of being swept.
|
||||
if not self.validate_upload_id(file):
|
||||
continue
|
||||
|
||||
file_path = os.path.join(root, file)
|
||||
if not self._inside_upload_dir(file_path):
|
||||
logger.warning(
|
||||
"Skipping cleanup candidate outside upload directory: %s",
|
||||
file_path,
|
||||
)
|
||||
continue
|
||||
matching_keys = self._upload_index_keys_for_file(
|
||||
current_index,
|
||||
file,
|
||||
file_path,
|
||||
)
|
||||
matching_rows = [
|
||||
current_index[key]
|
||||
for key in matching_keys
|
||||
if isinstance(current_index.get(key), dict)
|
||||
]
|
||||
|
||||
# Files without authoritative live index rows are not
|
||||
# eligible for destructive cleanup. Reference hashes,
|
||||
# recency, and ownership cannot be proven for them.
|
||||
if not matching_rows:
|
||||
continue
|
||||
|
||||
is_referenced = file in referenced_ids or any(
|
||||
str(info.get("id") or "") in referenced_ids
|
||||
or str(info.get("hash") or "") in referenced_hashes
|
||||
or str(info.get("checksum_sha256") or "") in referenced_hashes
|
||||
for info in matching_rows
|
||||
)
|
||||
metadata_is_recent = any(
|
||||
self._upload_metadata_is_recent(info, cutoff_date)
|
||||
for info in matching_rows
|
||||
)
|
||||
if is_referenced or metadata_is_recent:
|
||||
continue
|
||||
|
||||
reduced_index = {
|
||||
key: value
|
||||
for key, value in current_index.items()
|
||||
if key not in matching_keys
|
||||
}
|
||||
if matching_keys:
|
||||
try:
|
||||
self._atomic_write_json(
|
||||
uploads_db_path,
|
||||
reduced_index,
|
||||
sync_backup=True,
|
||||
)
|
||||
except Exception as e:
|
||||
try:
|
||||
self._atomic_write_json(
|
||||
uploads_db_path,
|
||||
current_index,
|
||||
sync_backup=True,
|
||||
)
|
||||
except Exception:
|
||||
logger.exception(
|
||||
"Failed to restore upload indexes after reconciliation failed for %s",
|
||||
file_path,
|
||||
)
|
||||
raise UploadCleanupSafetyError(
|
||||
"upload index rollback failed before file removal"
|
||||
) from e
|
||||
logger.warning(
|
||||
"Failed to reconcile upload index before removing %s: %s",
|
||||
file_path,
|
||||
e,
|
||||
)
|
||||
continue
|
||||
|
||||
try:
|
||||
os.remove(file_path)
|
||||
except FileNotFoundError:
|
||||
# The bytes are already absent. Keep the reduced
|
||||
# lifecycle index instead of recreating a stale row.
|
||||
current_index = reduced_index
|
||||
logger.info(
|
||||
"Reconciled missing expired upload from index: %s",
|
||||
file_path,
|
||||
)
|
||||
continue
|
||||
except Exception as e:
|
||||
if matching_keys:
|
||||
try:
|
||||
self._atomic_write_json(
|
||||
uploads_db_path,
|
||||
current_index,
|
||||
sync_backup=True,
|
||||
)
|
||||
except Exception:
|
||||
logger.exception(
|
||||
"Failed to restore upload index after removal failed for %s",
|
||||
file_path,
|
||||
)
|
||||
raise UploadCleanupSafetyError(
|
||||
"upload index rollback failed after file removal was refused"
|
||||
) from e
|
||||
logger.warning(f"Failed to remove {file_path}: {e}")
|
||||
continue
|
||||
|
||||
current_index = reduced_index
|
||||
cleaned_count += 1
|
||||
logger.info(f"Cleaned up old unreferenced upload: {file_path}")
|
||||
|
||||
try:
|
||||
if not os.listdir(root):
|
||||
os.rmdir(root)
|
||||
logger.info(f"Removed empty upload directory: {root}")
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to inspect/remove directory {root}: {e}")
|
||||
|
||||
logger.info(f"Upload cleanup completed: {cleaned_count} files removed")
|
||||
return cleaned_count
|
||||
except Exception as e:
|
||||
logger.error(f"Upload cleanup failed: {e}")
|
||||
return 0
|
||||
raise UploadCleanupSafetyError("upload cleanup safety checks failed") from e
|
||||
|
||||
def validate_upload_id(self, upload_id: str) -> bool:
|
||||
"""Validate that the upload ID matches the expected pattern."""
|
||||
@@ -293,71 +672,103 @@ class UploadHandler:
|
||||
|
||||
def _inside_upload_dir(self, path: str) -> bool:
|
||||
"""Check if path is inside the upload directory."""
|
||||
base = os.path.realpath(self.upload_dir)
|
||||
p = os.path.realpath(path)
|
||||
base = os.path.normcase(os.path.realpath(self.upload_dir))
|
||||
p = os.path.normcase(os.path.realpath(path))
|
||||
try:
|
||||
return os.path.commonpath([base, p]) == base
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
def _atomic_write_json(self, path: str, data: dict) -> None:
|
||||
def _atomic_write_json(
|
||||
self,
|
||||
path: str,
|
||||
data: dict,
|
||||
*,
|
||||
sync_backup: bool = False,
|
||||
) -> None:
|
||||
"""Write `data` to `path` atomically: write to a temp file in the
|
||||
same directory, then `os.replace` onto the target. The kernel
|
||||
guarantees `os.replace` is atomic on POSIX, so a reader either
|
||||
sees the old contents or the new contents, never a half-written
|
||||
file. Also keeps a `.bak` sibling of the previous good state.
|
||||
file. Normally `.bak` retains the previous good state. Destructive
|
||||
lifecycle transitions use ``sync_backup=True`` so recovery cannot
|
||||
resurrect metadata for bytes that were deliberately removed.
|
||||
"""
|
||||
directory = os.path.dirname(path) or "."
|
||||
fd, tmp = tempfile.mkstemp(prefix=".uploads-", suffix=".tmp", dir=directory)
|
||||
try:
|
||||
with os.fdopen(fd, "w", encoding="utf-8") as f:
|
||||
json.dump(data, f, indent=2)
|
||||
f.flush()
|
||||
os.fsync(f.fileno())
|
||||
if os.path.exists(path):
|
||||
bak = path + ".bak"
|
||||
|
||||
def _replace_json(target: str) -> None:
|
||||
fd, tmp = tempfile.mkstemp(
|
||||
prefix=".uploads-",
|
||||
suffix=".tmp",
|
||||
dir=directory,
|
||||
)
|
||||
try:
|
||||
with os.fdopen(fd, "w", encoding="utf-8") as f:
|
||||
json.dump(data, f, indent=2)
|
||||
f.flush()
|
||||
os.fsync(f.fileno())
|
||||
os.replace(tmp, target)
|
||||
except Exception:
|
||||
try:
|
||||
shutil.copy2(path, bak)
|
||||
os.unlink(tmp)
|
||||
except OSError:
|
||||
pass
|
||||
os.replace(tmp, path)
|
||||
# Update cache if this is the main index
|
||||
if path.endswith("uploads.json"):
|
||||
self._index_cache = data
|
||||
try:
|
||||
self._index_mtime = os.path.getmtime(path)
|
||||
except OSError:
|
||||
self._index_mtime = time.time()
|
||||
except Exception:
|
||||
raise
|
||||
|
||||
if sync_backup:
|
||||
_replace_json(path + ".bak")
|
||||
elif os.path.exists(path):
|
||||
try:
|
||||
os.unlink(tmp)
|
||||
shutil.copy2(path, path + ".bak")
|
||||
except OSError:
|
||||
pass
|
||||
raise
|
||||
|
||||
def _load_upload_index(self) -> Dict[str, Any]:
|
||||
_replace_json(path)
|
||||
# Update cache if this is the main index
|
||||
if path.endswith("uploads.json"):
|
||||
self._index_cache = data
|
||||
try:
|
||||
self._index_mtime = os.path.getmtime(path)
|
||||
except OSError:
|
||||
self._index_mtime = time.time()
|
||||
|
||||
def _load_upload_index(self, *, fail_on_error: bool = False) -> Dict[str, Any]:
|
||||
"""Load the upload index from disk/cache. Uses mtime-based validation
|
||||
to avoid redundant parsing on hot paths.
|
||||
to avoid redundant parsing on hot paths. When ``fail_on_error`` is
|
||||
true, a missing, malformed, or unreadable live index raises so
|
||||
destructive callers cannot mistake corruption for an empty store.
|
||||
"""
|
||||
uploads_db_path = os.path.join(self.upload_dir, "uploads.json")
|
||||
if not os.path.exists(uploads_db_path):
|
||||
candidates = (uploads_db_path, uploads_db_path + ".bak")
|
||||
if fail_on_error:
|
||||
# A backup is intentionally the previous snapshot. It is useful for
|
||||
# non-destructive reads, but cannot authorize deletion when the live
|
||||
# index is missing or corrupt.
|
||||
if not os.path.exists(uploads_db_path):
|
||||
raise ValueError("live uploads database is missing")
|
||||
existing_candidates = [uploads_db_path]
|
||||
else:
|
||||
existing_candidates = [path for path in candidates if os.path.exists(path)]
|
||||
if not existing_candidates:
|
||||
self._index_cache = {}
|
||||
self._index_mtime = 0.0
|
||||
return {}
|
||||
|
||||
# Check cache validity
|
||||
try:
|
||||
mtime = os.path.getmtime(uploads_db_path)
|
||||
if self._index_cache is not None and mtime <= self._index_mtime:
|
||||
mtime = max(os.path.getmtime(path) for path in existing_candidates)
|
||||
if (
|
||||
not fail_on_error
|
||||
and self._index_cache is not None
|
||||
and mtime <= self._index_mtime
|
||||
):
|
||||
return self._index_cache
|
||||
except OSError:
|
||||
mtime = 0.0
|
||||
|
||||
# Try the live file first, fall back to the .bak sibling if the
|
||||
# live file is truncated/corrupted.
|
||||
for candidate in (uploads_db_path, uploads_db_path + ".bak"):
|
||||
if not os.path.exists(candidate):
|
||||
continue
|
||||
for candidate in existing_candidates:
|
||||
try:
|
||||
with open(candidate, "r", encoding="utf-8") as f:
|
||||
data = json.load(f)
|
||||
@@ -369,6 +780,8 @@ class UploadHandler:
|
||||
logger.warning(f"Failed to read uploads database ({candidate}): {e}")
|
||||
continue
|
||||
|
||||
if fail_on_error:
|
||||
raise ValueError("live uploads database is unreadable")
|
||||
self._index_cache = {}
|
||||
return {}
|
||||
|
||||
@@ -381,6 +794,144 @@ class UploadHandler:
|
||||
return dict(info)
|
||||
return None
|
||||
|
||||
def reserve_upload(
|
||||
self,
|
||||
upload_id: str,
|
||||
*,
|
||||
owner: Optional[str],
|
||||
auth_manager: Any = None,
|
||||
allow_admin: bool = False,
|
||||
) -> Optional[Dict[str, Any]]:
|
||||
"""Owner-check and reserve an indexed upload against cleanup.
|
||||
|
||||
The live index lookup, ownership/path validation, and access touch all
|
||||
occur under the cleanup lock. A durable-reference writer must not
|
||||
commit when this returns ``None``.
|
||||
"""
|
||||
if not self.validate_upload_id(upload_id):
|
||||
return None
|
||||
|
||||
auth_configured = bool(auth_manager and getattr(auth_manager, "is_configured", False))
|
||||
if auth_configured and not owner:
|
||||
return None
|
||||
|
||||
uploads_db_path = os.path.join(self.upload_dir, "uploads.json")
|
||||
with self._index_lock:
|
||||
try:
|
||||
current = dict(self._load_upload_index(fail_on_error=True))
|
||||
except Exception:
|
||||
logger.warning("Cannot reserve upload %s without a valid live index", upload_id)
|
||||
return None
|
||||
matching_keys = [
|
||||
key
|
||||
for key, info in current.items()
|
||||
if isinstance(info, dict) and info.get("id") == upload_id
|
||||
]
|
||||
if not matching_keys:
|
||||
return None
|
||||
|
||||
matching_rows = [dict(current[key]) for key in matching_keys]
|
||||
row_owners = {
|
||||
str(row.get("owner")) if row.get("owner") is not None else None
|
||||
for row in matching_rows
|
||||
}
|
||||
row_hashes = {
|
||||
str(row.get("hash") or row.get("checksum_sha256"))
|
||||
for row in matching_rows
|
||||
if row.get("hash") or row.get("checksum_sha256")
|
||||
}
|
||||
if len(row_owners) != 1 or len(row_hashes) > 1:
|
||||
logger.warning(
|
||||
"Cannot reserve ambiguous upload index rows for %s",
|
||||
upload_id,
|
||||
)
|
||||
return None
|
||||
|
||||
is_admin = False
|
||||
if allow_admin and owner and auth_manager and hasattr(auth_manager, "is_admin"):
|
||||
try:
|
||||
is_admin = bool(auth_manager.is_admin(owner))
|
||||
except Exception:
|
||||
is_admin = False
|
||||
|
||||
now = datetime.now()
|
||||
current_info = matching_rows[0]
|
||||
if owner and not is_admin and current_info.get("owner") != owner:
|
||||
return None
|
||||
if not owner and current_info.get("owner") is not None:
|
||||
return None
|
||||
|
||||
existing_paths: set[str] = set()
|
||||
for row in matching_rows:
|
||||
stored_path = row.get("path")
|
||||
if not stored_path:
|
||||
continue
|
||||
if not self._inside_upload_dir(stored_path):
|
||||
logger.warning(
|
||||
"Cannot reserve upload %s with an out-of-root index path",
|
||||
upload_id,
|
||||
)
|
||||
return None
|
||||
if os.path.isfile(stored_path):
|
||||
if os.path.basename(stored_path) != upload_id:
|
||||
return None
|
||||
existing_paths.add(os.path.normcase(os.path.realpath(stored_path)))
|
||||
if len(existing_paths) > 1:
|
||||
logger.warning("Cannot reserve upload %s with multiple indexed paths", upload_id)
|
||||
return None
|
||||
path = next(iter(existing_paths), None) or self._find_upload_path(upload_id)
|
||||
if not path or not os.path.isfile(path) or not self._inside_upload_dir(path):
|
||||
return None
|
||||
|
||||
last_accessed = self._parse_upload_timestamp(current_info.get("last_accessed"))
|
||||
path_changed = current_info.get("path") != path
|
||||
needs_write = (
|
||||
path_changed
|
||||
or last_accessed is None
|
||||
or last_accessed < now - timedelta(minutes=5)
|
||||
)
|
||||
if needs_write:
|
||||
accessed_at = now.isoformat()
|
||||
updated_index = dict(current)
|
||||
for key in matching_keys:
|
||||
updated = dict(updated_index[key])
|
||||
updated["path"] = path
|
||||
updated["last_accessed"] = accessed_at
|
||||
updated_index[key] = updated
|
||||
try:
|
||||
self._atomic_write_json(
|
||||
uploads_db_path,
|
||||
updated_index,
|
||||
sync_backup=True,
|
||||
)
|
||||
except Exception:
|
||||
try:
|
||||
self._atomic_write_json(
|
||||
uploads_db_path,
|
||||
current,
|
||||
sync_backup=True,
|
||||
)
|
||||
except Exception:
|
||||
logger.exception(
|
||||
"Failed to restore upload indexes after reservation failed for %s",
|
||||
upload_id,
|
||||
)
|
||||
logger.exception("Failed to reserve upload %s against cleanup", upload_id)
|
||||
return None
|
||||
current_info = dict(updated_index[matching_keys[0]])
|
||||
|
||||
resolved = dict(current_info)
|
||||
resolved.setdefault("id", upload_id)
|
||||
resolved["path"] = path
|
||||
resolved.setdefault("name", os.path.basename(path))
|
||||
resolved.setdefault("original_name", resolved["name"])
|
||||
resolved.setdefault("mime", mimetypes.guess_type(path)[0] or "application/octet-stream")
|
||||
if resolved.get("hash") and not resolved.get("checksum_sha256"):
|
||||
resolved["checksum_sha256"] = resolved["hash"]
|
||||
if resolved.get("uploaded_at") and not resolved.get("created_at"):
|
||||
resolved["created_at"] = resolved["uploaded_at"]
|
||||
return resolved
|
||||
|
||||
def _renamed_upload_index_key(self, key: str, info: Dict[str, Any], old_owner: str, new_owner: str) -> str:
|
||||
"""Return the storage key to use after renaming an owned upload row.
|
||||
|
||||
@@ -475,16 +1026,32 @@ class UploadHandler:
|
||||
if not self.validate_upload_id(upload_id):
|
||||
return None
|
||||
|
||||
candidates: list[str] = []
|
||||
direct = os.path.join(self.upload_dir, upload_id)
|
||||
if os.path.exists(direct) and self._inside_upload_dir(direct):
|
||||
return direct
|
||||
if os.path.isfile(direct) and self._inside_upload_dir(direct):
|
||||
candidates.append(os.path.realpath(direct))
|
||||
|
||||
for root, _dirs, files in os.walk(self.upload_dir, followlinks=False):
|
||||
for root, dirs, files in os.walk(self.upload_dir, followlinks=False):
|
||||
is_junction = getattr(os.path, "isjunction", lambda _path: False)
|
||||
dirs[:] = [
|
||||
directory
|
||||
for directory in dirs
|
||||
if not os.path.islink(os.path.join(root, directory))
|
||||
and not is_junction(os.path.join(root, directory))
|
||||
]
|
||||
if upload_id in files:
|
||||
path = os.path.join(root, upload_id)
|
||||
if self._inside_upload_dir(path):
|
||||
return path
|
||||
return None
|
||||
if os.path.isfile(path) and self._inside_upload_dir(path):
|
||||
real_path = os.path.realpath(path)
|
||||
if real_path not in candidates:
|
||||
candidates.append(real_path)
|
||||
if len(candidates) > 1:
|
||||
logger.warning(
|
||||
"Upload ID %s resolves to multiple physical files",
|
||||
upload_id,
|
||||
)
|
||||
return None
|
||||
return candidates[0] if candidates else None
|
||||
|
||||
def resolve_upload(
|
||||
self,
|
||||
@@ -493,52 +1060,20 @@ class UploadHandler:
|
||||
auth_manager: Any = None,
|
||||
allow_admin: bool = True,
|
||||
) -> Optional[Dict[str, Any]]:
|
||||
"""Resolve an upload ID to metadata only if the caller may read it.
|
||||
"""Resolve and reserve an upload only if the caller may read it.
|
||||
|
||||
This is the owner-aware lookup used by internal processors. Public
|
||||
download routes already perform owner checks; chat/document paths must
|
||||
do the same before reading file bytes server-side.
|
||||
do the same before reading file bytes server-side. Reservation shares
|
||||
cleanup's lifecycle lock and prevents a newly persisted reference from
|
||||
racing final deletion.
|
||||
"""
|
||||
if not self.validate_upload_id(upload_id):
|
||||
logger.warning(f"Invalid upload ID format: {upload_id}")
|
||||
return None
|
||||
|
||||
auth_configured = bool(auth_manager and getattr(auth_manager, "is_configured", False))
|
||||
if auth_configured and not owner:
|
||||
return None
|
||||
|
||||
info = self.get_upload_info(upload_id) or {}
|
||||
is_admin = False
|
||||
if allow_admin and owner and auth_manager and hasattr(auth_manager, "is_admin"):
|
||||
try:
|
||||
is_admin = bool(auth_manager.is_admin(owner))
|
||||
except Exception:
|
||||
is_admin = False
|
||||
|
||||
if owner and not is_admin:
|
||||
if info.get("owner") != owner:
|
||||
logger.warning("Upload %s denied for owner %s", upload_id, owner)
|
||||
return None
|
||||
if not owner and info.get("owner") is not None:
|
||||
logger.warning("Upload %s denied without an authenticated owner", upload_id)
|
||||
return None
|
||||
|
||||
path = info.get("path")
|
||||
if not path or not os.path.exists(path) or not self._inside_upload_dir(path):
|
||||
path = self._find_upload_path(upload_id)
|
||||
if not path:
|
||||
return None
|
||||
if not self._inside_upload_dir(path):
|
||||
logger.warning(f"Upload path outside upload directory: {path}")
|
||||
return None
|
||||
|
||||
resolved = dict(info)
|
||||
resolved.setdefault("id", upload_id)
|
||||
resolved["path"] = path
|
||||
resolved.setdefault("name", os.path.basename(path))
|
||||
resolved.setdefault("original_name", resolved["name"])
|
||||
resolved.setdefault("mime", mimetypes.guess_type(path)[0] or "application/octet-stream")
|
||||
return resolved
|
||||
return self.reserve_upload(
|
||||
upload_id,
|
||||
owner=owner,
|
||||
auth_manager=auth_manager,
|
||||
allow_admin=allow_admin,
|
||||
)
|
||||
|
||||
def cleanup_rate_limits(self):
|
||||
"""Remove stale entries from upload_rate_log."""
|
||||
@@ -706,6 +1241,9 @@ class UploadHandler:
|
||||
# fresh-insert path below; release the lock first.
|
||||
raise LookupError("upload entry vanished mid-dedupe")
|
||||
existing_file["last_accessed"] = datetime.now().isoformat()
|
||||
existing_file.setdefault("checksum_sha256", file_hash)
|
||||
if existing_file.get("uploaded_at"):
|
||||
existing_file.setdefault("created_at", existing_file["uploaded_at"])
|
||||
current[live_key] = existing_file
|
||||
self._atomic_write_json(uploads_db_path, current)
|
||||
except LookupError:
|
||||
@@ -721,7 +1259,9 @@ class UploadHandler:
|
||||
"size": existing_file["size"],
|
||||
"name": existing_file["original_name"],
|
||||
"hash": file_hash,
|
||||
"checksum_sha256": existing_file.get("checksum_sha256") or file_hash,
|
||||
"uploaded_at": existing_file["uploaded_at"],
|
||||
"created_at": existing_file.get("created_at") or existing_file["uploaded_at"],
|
||||
"owner": existing_file.get("owner"),
|
||||
"width": existing_file.get("width"),
|
||||
"height": existing_file.get("height"),
|
||||
@@ -744,6 +1284,7 @@ class UploadHandler:
|
||||
raise HTTPException(status_code=500, detail=f"Failed to save file: {str(e)}")
|
||||
|
||||
# Create file metadata
|
||||
created_at = datetime.now().isoformat()
|
||||
file_metadata = {
|
||||
"id": file_id,
|
||||
"path": file_path,
|
||||
@@ -751,9 +1292,11 @@ class UploadHandler:
|
||||
"size": file_size,
|
||||
"name": safe_filename,
|
||||
"hash": file_hash,
|
||||
"checksum_sha256": file_hash,
|
||||
"original_name": original_filename,
|
||||
"uploaded_at": datetime.now().isoformat(),
|
||||
"last_accessed": datetime.now().isoformat(),
|
||||
"uploaded_at": created_at,
|
||||
"created_at": created_at,
|
||||
"last_accessed": created_at,
|
||||
"client_ip": client_ip,
|
||||
"owner": owner,
|
||||
}
|
||||
|
||||
+17
-1
@@ -1822,7 +1822,7 @@ import { wireArrowUpRecall, getLastUserMessageFromChatHistory } from './composer
|
||||
typewriterInto(roundHolder.querySelector('.body'), errMsg);
|
||||
break;
|
||||
}
|
||||
if (json.delta || json.type === 'agent_prep' || json.type === 'tool_start' || json.type === 'tool_output' || json.type === 'tool_progress' || json.type === 'agent_step' || json.type === 'doc_stream_open' || json.type === 'doc_stream_delta' || json.type === 'research_progress') {
|
||||
if (json.delta || json.type === 'agent_prep' || json.type === 'tool_start' || json.type === 'tool_output' || json.type === 'tool_progress' || json.type === 'agent_step' || json.type === 'loop_breaker_triggered' || json.type === 'intent_nudge_exhausted' || json.type === 'doc_stream_open' || json.type === 'doc_stream_delta' || json.type === 'research_progress') {
|
||||
clearResponseTimeout();
|
||||
clearProcessingProbe();
|
||||
clearFirstTokenWaitTimers();
|
||||
@@ -2852,6 +2852,22 @@ import { wireArrowUpRecall, getLastUserMessageFromChatHistory } from './composer
|
||||
const chatBox = document.getElementById('chat-history');
|
||||
chatBox.appendChild(budgetDiv);
|
||||
|
||||
} else if (json.type === 'loop_breaker_triggered' || json.type === 'intent_nudge_exhausted') {
|
||||
if (_isBg) continue;
|
||||
_cancelThinkingTimer();
|
||||
_removeThinkingSpinner();
|
||||
const guardDiv = document.createElement('div');
|
||||
guardDiv.className = 'stopped-indicator';
|
||||
const guardLabel = document.createElement('span');
|
||||
guardLabel.textContent = `[Agent guard: ${json.message || json.reason || 'internal stop'}]`;
|
||||
guardDiv.appendChild(guardLabel);
|
||||
const targetBody = roundHolder && roundHolder.querySelector('.body');
|
||||
if (targetBody) targetBody.appendChild(guardDiv);
|
||||
else {
|
||||
const chatBox = document.getElementById('chat-history');
|
||||
if (chatBox) chatBox.appendChild(guardDiv);
|
||||
}
|
||||
|
||||
} else if (json.type === 'teacher_takeover') {
|
||||
if (_isBg) continue;
|
||||
_cancelThinkingTimer();
|
||||
|
||||
@@ -952,8 +952,20 @@ function _libCachePut(key, value) {
|
||||
}
|
||||
}
|
||||
|
||||
function _resetBulkSelectionForContextChange({ rerender = false } = {}) {
|
||||
const hadSelection = !!(state._selectedUids && state._selectedUids.size);
|
||||
const wasSelectMode = !!state._selectMode;
|
||||
if (state._selectedUids) state._selectedUids.clear();
|
||||
state._selectMode = false;
|
||||
if (hadSelection || wasSelectMode) {
|
||||
_updateBulkBar();
|
||||
if (rerender) _renderGrid();
|
||||
}
|
||||
}
|
||||
|
||||
function _resetEmailListForFreshLoad() {
|
||||
_exitEmailReaderModeForList();
|
||||
_resetBulkSelectionForContextChange();
|
||||
state._libOffset = 0;
|
||||
state._libEmails = [];
|
||||
state._libTotal = 0;
|
||||
@@ -2507,6 +2519,7 @@ function _clearFilterPillSideEffect() {
|
||||
|
||||
function _addSearchPill(pill) {
|
||||
if (!pill) return;
|
||||
_resetBulkSelectionForContextChange({ rerender: true });
|
||||
if (!Array.isArray(state._libSearchPills)) state._libSearchPills = [];
|
||||
// Dedup by email (contact), text (text pill), or filter value.
|
||||
if (pill.type === 'contact') {
|
||||
@@ -2541,6 +2554,7 @@ function _searchQueryFromPills() {
|
||||
|
||||
function _removeSearchPillAt(idx) {
|
||||
if (!Array.isArray(state._libSearchPills)) return;
|
||||
_resetBulkSelectionForContextChange({ rerender: true });
|
||||
const removed = state._libSearchPills[idx];
|
||||
state._libSearchPills.splice(idx, 1);
|
||||
if (removed && removed.type === 'filter') _clearFilterPillSideEffect();
|
||||
@@ -2718,6 +2732,7 @@ async function _initEmailSearchChipBar() {
|
||||
// directly.
|
||||
let _libSearchTypingTimer = null;
|
||||
input.addEventListener('input', async () => {
|
||||
_resetBulkSelectionForContextChange({ rerender: true });
|
||||
state._libSearchDraft = input.value;
|
||||
await _refreshSuggestions();
|
||||
if (_libSearchTypingTimer) clearTimeout(_libSearchTypingTimer);
|
||||
@@ -2853,6 +2868,7 @@ window.addEventListener('click', (e) => {
|
||||
|
||||
async function _doSearch() {
|
||||
_exitEmailReaderModeForList();
|
||||
_resetBulkSelectionForContextChange({ rerender: true });
|
||||
const seq = ++_libSearchSeq;
|
||||
const derived = _deriveSearchScope(state._libSearch);
|
||||
const q = derived.q;
|
||||
|
||||
@@ -661,8 +661,11 @@ export function mdToHtml(src, opts) {
|
||||
return placeholder;
|
||||
} catch (e) { return match; }
|
||||
});
|
||||
// Inline math: $...$ (not preceded/followed by $ or digit, not spanning multiple lines)
|
||||
s = s.replace(/(?<!\$)\$(?!\$)([^\$\n]+?)\$(?!\$)/g, (match, math) => {
|
||||
// Inline math: $...$ — single line only, and Pandoc-style delimiter rules so
|
||||
// currency doesn't render as math ("$5 to $10"): the opening $ must be
|
||||
// immediately followed by a non-space, the closing $ must be immediately
|
||||
// preceded by a non-space and not followed by a digit.
|
||||
s = s.replace(/(?<![\$\d])\$(?!\$)(?=\S)([^\$\n]+?)(?<=\S)\$(?!\$|\d)/g, (match, math) => {
|
||||
try {
|
||||
const raw = math.replace(/&/g, '&').replace(/</g, '<').replace(/>/g, '>');
|
||||
const placeholder = `___MATH_BLOCK_${mathBlocks.length}___`;
|
||||
|
||||
@@ -2070,6 +2070,9 @@ body.bg-pattern-sparkles {
|
||||
overflow-wrap: break-word;
|
||||
overflow: hidden;
|
||||
}
|
||||
.msg-user:has(textarea) {
|
||||
width: 85%;
|
||||
}
|
||||
.msg-ai {
|
||||
align-items: flex-start;
|
||||
margin-right: auto;
|
||||
|
||||
@@ -68,3 +68,24 @@ def test_no_rounds_exhausted_on_normal_finish(monkeypatch):
|
||||
# A plain answer (no tool block) -> done-break on round 1 -> no event.
|
||||
events = _run_loop(monkeypatch, "All done, here is your answer.", max_rounds=2)
|
||||
assert not any(e.get("type") == "rounds_exhausted" for e in events), events
|
||||
|
||||
|
||||
def test_emits_intent_nudge_exhausted_when_cap_is_exhausted(monkeypatch):
|
||||
_patch_common(monkeypatch)
|
||||
|
||||
events = _run_loop(monkeypatch, "Let me check the logs", max_rounds=5)
|
||||
|
||||
guard = next((e for e in events if e.get("type") == "intent_nudge_exhausted"), None)
|
||||
assert guard is not None, events
|
||||
assert guard["reason"] == "intent_without_action_nudge_cap"
|
||||
assert guard["nudges"] == 2
|
||||
|
||||
|
||||
def test_emits_loop_breaker_triggered_when_loop_breaker_trips(monkeypatch):
|
||||
_patch_common(monkeypatch)
|
||||
|
||||
events = _run_loop(monkeypatch, "```bash\necho hi\n```", max_rounds=6)
|
||||
|
||||
guard = next((e for e in events if e.get("type") == "loop_breaker_triggered"), None)
|
||||
assert guard is not None, events
|
||||
assert guard["reason"] == "loop_breaker_stall"
|
||||
|
||||
@@ -0,0 +1,75 @@
|
||||
import json
|
||||
|
||||
from src.attachment_refs import (
|
||||
attachment_ref,
|
||||
persistable_message_content,
|
||||
search_index_text,
|
||||
)
|
||||
|
||||
|
||||
def test_persistable_message_content_replaces_inline_media_with_attachment_ref():
|
||||
metadata = {
|
||||
"attachments": [
|
||||
{
|
||||
"id": "abc123.png",
|
||||
"name": "diagram.png",
|
||||
"mime": "image/png",
|
||||
"size": 42,
|
||||
"checksum_sha256": "sha256-digest",
|
||||
"created_at": "2026-07-09T12:00:00",
|
||||
"vision": "A small architecture diagram.",
|
||||
}
|
||||
]
|
||||
}
|
||||
content = [
|
||||
{"type": "text", "text": "Please inspect this."},
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {"url": "data:image/png;base64," + ("A" * 5000)},
|
||||
},
|
||||
]
|
||||
|
||||
stored = persistable_message_content(content, metadata)
|
||||
|
||||
assert "base64" not in stored
|
||||
assert "A" * 100 not in stored
|
||||
assert "Please inspect this." in stored
|
||||
assert "Attachment: diagram.png" in stored
|
||||
assert "id=abc123.png" in stored
|
||||
assert "sha256=sha256-digest" in stored
|
||||
assert "A small architecture diagram." in stored
|
||||
|
||||
|
||||
def test_search_index_text_strips_legacy_serialized_data_url_blocks():
|
||||
legacy = json.dumps([
|
||||
{"type": "text", "text": "Find this useful caption"},
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {"url": "data:image/jpeg;base64," + ("B" * 4096)},
|
||||
},
|
||||
])
|
||||
|
||||
indexed = search_index_text(legacy)
|
||||
|
||||
assert indexed == "Find this useful caption\n[1 inline media payload omitted]"
|
||||
|
||||
|
||||
def test_attachment_ref_normalizes_hash_aliases():
|
||||
ref = attachment_ref({
|
||||
"id": "file-id",
|
||||
"original_name": "report.pdf",
|
||||
"mime": "application/pdf",
|
||||
"size": 99,
|
||||
"hash": "abc",
|
||||
"uploaded_at": "2026-07-09T12:00:00",
|
||||
})
|
||||
|
||||
assert ref == {
|
||||
"type": "attachment_ref",
|
||||
"attachment_id": "file-id",
|
||||
"name": "report.pdf",
|
||||
"mime": "application/pdf",
|
||||
"size": 99,
|
||||
"checksum_sha256": "abc",
|
||||
"created_at": "2026-07-09T12:00:00",
|
||||
}
|
||||
@@ -109,10 +109,9 @@ async def test_learn_sender_signatures_resolves_llm_for_task_owner(monkeypatch):
|
||||
def select(self, *_args, **_kwargs):
|
||||
return "OK", []
|
||||
|
||||
def search(self, *_args, **_kwargs):
|
||||
return "OK", [b"1 2 3"]
|
||||
|
||||
def fetch(self, _uid, _query):
|
||||
def uid(self, command, *_args):
|
||||
if command == "SEARCH":
|
||||
return "OK", [b"1 2 3"]
|
||||
return "OK", [(None, b"From: Writer <writer@example.com>\r\n\r\n")]
|
||||
|
||||
def logout(self):
|
||||
@@ -171,11 +170,10 @@ async def test_learn_sender_signatures_writes_owner_scoped_cache(monkeypatch, tm
|
||||
def select(self, *_args, **_kwargs):
|
||||
return "OK", []
|
||||
|
||||
def search(self, *_args, **_kwargs):
|
||||
return "OK", [b"1 2 3"]
|
||||
|
||||
def fetch(self, uid, query):
|
||||
if "HEADER.FIELDS" in query:
|
||||
def uid(self, command, uid=None, query=None):
|
||||
if command == "SEARCH":
|
||||
return "OK", [b"1 2 3"]
|
||||
if query and "HEADER.FIELDS" in query:
|
||||
return "OK", [(None, b"From: Writer <writer@example.com>\r\n\r\n")]
|
||||
return "OK", [
|
||||
(
|
||||
|
||||
@@ -0,0 +1,22 @@
|
||||
import os
|
||||
|
||||
from src.builtin_mcp import builtin_python_env
|
||||
|
||||
|
||||
def test_builtin_python_env_preserves_existing_pythonpath(monkeypatch):
|
||||
monkeypatch.setenv(
|
||||
"PYTHONPATH",
|
||||
os.pathsep.join(["/app/venv/lib/python3.13/site-packages", "/app", "/extra"]),
|
||||
)
|
||||
|
||||
env = builtin_python_env("/app")
|
||||
|
||||
assert env == {
|
||||
"PYTHONPATH": os.pathsep.join(["/app", "/app/venv/lib/python3.13/site-packages", "/extra"])
|
||||
}
|
||||
|
||||
|
||||
def test_builtin_python_env_uses_app_root_without_existing_pythonpath(monkeypatch):
|
||||
monkeypatch.delenv("PYTHONPATH", raising=False)
|
||||
|
||||
assert builtin_python_env("/srv/odysseus") == {"PYTHONPATH": "/srv/odysseus"}
|
||||
@@ -0,0 +1,145 @@
|
||||
"""Issue #4593 — the CalDAV DAVClient must be closed on every path.
|
||||
|
||||
`_sync_blocking` (src/caldav_sync.py) and `_writeback_blocking`
|
||||
(src/caldav_writeback.py) each open their own DAVClient. The client holds an
|
||||
HTTP session with pooled connections; if it is never closed those connections
|
||||
leak for the lifetime of the process. These tests pin that the client is
|
||||
closed on the discovery early-returns, the normal return, and the
|
||||
write-back paths, using a fake client so no network or `caldav` install is
|
||||
needed.
|
||||
"""
|
||||
|
||||
import sys
|
||||
import types
|
||||
|
||||
import pytest
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
|
||||
def _stub_sync_deps(monkeypatch):
|
||||
"""Make `_sync_blocking`'s lazy imports resolve without a real caldav/db."""
|
||||
err_mod = types.ModuleType("caldav.lib.error")
|
||||
|
||||
class AuthorizationError(Exception):
|
||||
pass
|
||||
|
||||
class NotFoundError(Exception):
|
||||
pass
|
||||
|
||||
err_mod.AuthorizationError = AuthorizationError
|
||||
err_mod.NotFoundError = NotFoundError
|
||||
monkeypatch.setitem(sys.modules, "caldav", types.ModuleType("caldav"))
|
||||
monkeypatch.setitem(sys.modules, "caldav.lib", types.ModuleType("caldav.lib"))
|
||||
monkeypatch.setitem(sys.modules, "caldav.lib.error", err_mod)
|
||||
|
||||
db_mod = types.ModuleType("core.database")
|
||||
db_mod.CalendarCal = MagicMock()
|
||||
db_mod.CalendarEvent = MagicMock()
|
||||
db_mod.CalendarDeletedEvent = MagicMock()
|
||||
db_mod.SessionLocal = MagicMock()
|
||||
if "core" not in sys.modules:
|
||||
monkeypatch.setitem(sys.modules, "core", types.ModuleType("core"))
|
||||
monkeypatch.setitem(sys.modules, "core.database", db_mod)
|
||||
|
||||
# Stub routes.calendar_routes so the lazy import of _ensure_positive_duration
|
||||
# inside _sync_blocking doesn't drag in dateutil / FastAPI / SQLAlchemy.
|
||||
routes_mod = types.ModuleType("routes")
|
||||
cal_routes_mod = types.ModuleType("routes.calendar_routes")
|
||||
cal_routes_mod._ensure_positive_duration = lambda start, end, all_day: end
|
||||
if "routes" not in sys.modules:
|
||||
monkeypatch.setitem(sys.modules, "routes", routes_mod)
|
||||
monkeypatch.setitem(sys.modules, "routes.calendar_routes", cal_routes_mod)
|
||||
|
||||
return AuthorizationError
|
||||
|
||||
|
||||
def test_sync_closes_client_on_discovery_auth_failure(monkeypatch):
|
||||
import src.caldav_sync as sync
|
||||
|
||||
AuthorizationError = _stub_sync_deps(monkeypatch)
|
||||
client = MagicMock()
|
||||
client.principal.side_effect = AuthorizationError("bad credentials")
|
||||
monkeypatch.setattr(sync, "_build_dav_client", lambda *a, **k: client)
|
||||
|
||||
result = sync._sync_blocking("alice", "https://dav.example.com/", "u", "p")
|
||||
|
||||
client.close.assert_called_once()
|
||||
assert any("Discovery failed" in e for e in result["errors"])
|
||||
|
||||
|
||||
def test_sync_closes_client_when_url_fallback_fails(monkeypatch):
|
||||
import src.caldav_sync as sync
|
||||
|
||||
_stub_sync_deps(monkeypatch)
|
||||
client = MagicMock()
|
||||
# principal() raises a generic error -> the URL-as-calendar fallback is
|
||||
# tried; make that fail too so the function hits the early return.
|
||||
client.principal.side_effect = RuntimeError("no principal endpoint")
|
||||
monkeypatch.setattr(sync, "_build_dav_client", lambda *a, **k: client)
|
||||
monkeypatch.setattr(
|
||||
sync, "_open_url_as_calendar",
|
||||
MagicMock(side_effect=RuntimeError("not a calendar")),
|
||||
)
|
||||
|
||||
result = sync._sync_blocking("alice", "https://dav.example.com/", "u", "p")
|
||||
|
||||
client.close.assert_called_once()
|
||||
assert result["errors"]
|
||||
|
||||
|
||||
def test_writeback_closes_client_when_no_calendars(monkeypatch):
|
||||
import src.caldav_sync as sync
|
||||
import src.caldav_writeback as wb
|
||||
|
||||
client = MagicMock()
|
||||
monkeypatch.setattr(sync, "_build_dav_client", lambda *a, **k: client)
|
||||
monkeypatch.setattr(wb, "_discover_calendars", lambda c: [])
|
||||
|
||||
result = wb._writeback_blocking(
|
||||
"caldav-1", {"uid": "evt-1"}, False, "https://dav.example.com/", "u", "p"
|
||||
)
|
||||
|
||||
client.close.assert_called_once()
|
||||
assert result["ok"] is False
|
||||
|
||||
|
||||
def test_writeback_closes_client_on_success(monkeypatch):
|
||||
import src.caldav_sync as sync
|
||||
import src.caldav_writeback as wb
|
||||
|
||||
client = MagicMock()
|
||||
monkeypatch.setattr(sync, "_build_dav_client", lambda *a, **k: client)
|
||||
monkeypatch.setattr(wb, "_discover_calendars", lambda c: [MagicMock()])
|
||||
monkeypatch.setattr(wb, "push_event", lambda *a, **k: {"ok": True})
|
||||
|
||||
result = wb._writeback_blocking(
|
||||
"caldav-1", {"uid": "evt-1"}, False, "https://dav.example.com/", "u", "p"
|
||||
)
|
||||
|
||||
client.close.assert_called_once()
|
||||
assert result["ok"] is True
|
||||
|
||||
|
||||
def test_sync_closes_client_when_session_local_raises(monkeypatch):
|
||||
import src.caldav_sync as sync
|
||||
|
||||
AuthorizationError = _stub_sync_deps(monkeypatch)
|
||||
|
||||
# Give principal() a working response so discovery passes
|
||||
mock_principal = MagicMock()
|
||||
mock_cal = MagicMock()
|
||||
mock_cal.url = "https://dav.example.com/alice/home/"
|
||||
mock_principal.calendars.return_value = [mock_cal]
|
||||
|
||||
client = MagicMock()
|
||||
client.principal.return_value = mock_principal
|
||||
monkeypatch.setattr(sync, "_build_dav_client", lambda *a, **k: client)
|
||||
|
||||
# Make SessionLocal blow up before any DB work
|
||||
import sys
|
||||
sys.modules["core.database"].SessionLocal.side_effect = RuntimeError("DB unavailable")
|
||||
|
||||
with pytest.raises(RuntimeError, match="DB unavailable"):
|
||||
sync._sync_blocking("alice", "https://dav.example.com/", "u", "p")
|
||||
|
||||
client.close.assert_called_once()
|
||||
@@ -93,6 +93,10 @@ class _FakeClient:
|
||||
def calendar(self, url=None):
|
||||
return _FakeCalendar(url)
|
||||
|
||||
def close(self):
|
||||
# Mirror the real DAVClient: sync now closes the client on every path.
|
||||
self.closed = True
|
||||
|
||||
|
||||
def _install_fake_caldav(monkeypatch):
|
||||
fake = types.ModuleType("caldav")
|
||||
|
||||
@@ -0,0 +1,194 @@
|
||||
"""Regression: CalDAV test_connection must trust the operator's CA bundle.
|
||||
|
||||
The pre-flight used httpx with trust_env=False, which ignored
|
||||
SSL_CERT_FILE/REQUESTS_CA_BUNDLE. Self-signed CalDAV servers that the
|
||||
real sync accepts (via caldav lib -> requests -> honors bundle) were
|
||||
rejected by the test with CERTIFICATE_VERIFY_FAILED.
|
||||
|
||||
These tests exercise the *route handler* directly (via ASGI TestClient)
|
||||
and capture the verify= kwarg passed to httpx.AsyncClient, ensuring the
|
||||
route code — not a test-side duplicate — builds the SSL context correctly.
|
||||
"""
|
||||
import os
|
||||
import ssl
|
||||
import sys
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
# No module-level sys.modules stubbing here: conftest pre-imports the real
|
||||
# sqlalchemy/core.database, and stubbing extras (e.g. caldav) at collection
|
||||
# time leaks MagicMocks into later tests in the same process — it made
|
||||
# test_caldav_redirect_hardening's real DAVClient a mock that never sent
|
||||
# the PROPFIND. The route's lazy imports are patched per-request instead.
|
||||
|
||||
|
||||
def _fake_response(status_code=207, headers=None):
|
||||
resp = MagicMock()
|
||||
resp.status_code = status_code
|
||||
resp.headers = headers or {}
|
||||
return resp
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def client():
|
||||
from fastapi import FastAPI
|
||||
from fastapi.testclient import TestClient
|
||||
from routes.calendar_routes import setup_calendar_routes
|
||||
|
||||
with patch("routes.calendar_routes._require_user", return_value="test-owner"):
|
||||
router = setup_calendar_routes()
|
||||
app = FastAPI()
|
||||
app.include_router(router)
|
||||
return TestClient(app)
|
||||
|
||||
|
||||
def _make_fake_async_client(captured):
|
||||
"""Return a fake httpx.AsyncClient class that captures constructor kwargs."""
|
||||
class FakeAsyncClient:
|
||||
def __init__(self, **kwargs):
|
||||
captured.update(kwargs)
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *a):
|
||||
pass
|
||||
|
||||
async def request(self, *a, **kw):
|
||||
return _fake_response(207)
|
||||
|
||||
return FakeAsyncClient
|
||||
|
||||
|
||||
def _post_test(client, captured, env=None):
|
||||
"""POST /api/calendar/test with credentials in body so no DB lookup needed.
|
||||
|
||||
Patches httpx.AsyncClient at the real module level so the route's
|
||||
``import httpx; httpx.AsyncClient(...)`` picks up the fake class.
|
||||
Also stubs validate_caldav_url (lazy-imported from src.caldav_sync).
|
||||
"""
|
||||
fake_cls = _make_fake_async_client(captured)
|
||||
|
||||
# Stub the caldav_sync module so the lazy `from src.caldav_sync import validate_caldav_url`
|
||||
# inside the route body resolves to a pass-through.
|
||||
caldav_sync_stub = MagicMock()
|
||||
caldav_sync_stub.validate_caldav_url = lambda u: u
|
||||
|
||||
ctx_managers = [
|
||||
patch.object(httpx, "AsyncClient", fake_cls),
|
||||
patch.dict(sys.modules, {"src.caldav_sync": caldav_sync_stub}),
|
||||
patch("routes.calendar_routes._require_user", return_value="test-owner"),
|
||||
]
|
||||
if env is not None:
|
||||
ctx_managers.append(patch.dict(os.environ, env))
|
||||
|
||||
# Enter all context managers
|
||||
for cm in ctx_managers:
|
||||
cm.__enter__()
|
||||
try:
|
||||
return client.post(
|
||||
"/api/calendar/test",
|
||||
json={"url": "https://cal.example.com", "username": "u", "password": "p"},
|
||||
)
|
||||
finally:
|
||||
for cm in reversed(ctx_managers):
|
||||
cm.__exit__(None, None, None)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Route-level tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def test_route_passes_ssl_context_with_correct_flags(client):
|
||||
"""The route must pass an ssl.SSLContext to httpx.AsyncClient(verify=...)
|
||||
with trust_env=False, follow_redirects=False, and VERIFY_X509_STRICT cleared."""
|
||||
captured = {}
|
||||
resp = _post_test(client, captured)
|
||||
|
||||
assert resp.status_code == 200
|
||||
assert isinstance(captured.get("verify"), ssl.SSLContext), (
|
||||
f"verify= should be an ssl.SSLContext, got {type(captured.get('verify'))}"
|
||||
)
|
||||
assert captured.get("trust_env") is False
|
||||
assert captured.get("follow_redirects") is False
|
||||
ctx = captured["verify"]
|
||||
assert not (ctx.verify_flags & ssl.VERIFY_X509_STRICT), (
|
||||
"VERIFY_X509_STRICT must be cleared for self-signed CA compat"
|
||||
)
|
||||
|
||||
|
||||
def test_route_ssl_cert_file_takes_precedence(client, tmp_path):
|
||||
"""SSL_CERT_FILE is the exact bundle loaded when both variables are set."""
|
||||
bundle_a = tmp_path / "ssl-cert-file.pem"
|
||||
bundle_b = tmp_path / "requests-ca-bundle.pem"
|
||||
bundle_a.write_text("ssl-cert-file", encoding="utf-8")
|
||||
bundle_b.write_text("requests-ca-bundle", encoding="utf-8")
|
||||
|
||||
loaded = []
|
||||
|
||||
class FakeSSLContext:
|
||||
def __init__(self):
|
||||
self.verify_flags = ssl.VERIFY_X509_STRICT
|
||||
|
||||
def load_verify_locations(self, cafile=None, capath=None, cadata=None):
|
||||
loaded.append(
|
||||
{
|
||||
"cafile": cafile,
|
||||
"capath": capath,
|
||||
"cadata": cadata,
|
||||
}
|
||||
)
|
||||
|
||||
ssl_context = FakeSSLContext()
|
||||
captured = {}
|
||||
env = {
|
||||
"SSL_CERT_FILE": str(bundle_a),
|
||||
"REQUESTS_CA_BUNDLE": str(bundle_b),
|
||||
}
|
||||
|
||||
with patch.object(
|
||||
ssl,
|
||||
"create_default_context",
|
||||
return_value=ssl_context,
|
||||
):
|
||||
resp = _post_test(client, captured, env=env)
|
||||
|
||||
assert resp.status_code == 200
|
||||
assert resp.json() == {"ok": True}
|
||||
assert loaded == [
|
||||
{
|
||||
"cafile": str(bundle_a),
|
||||
"capath": None,
|
||||
"cadata": None,
|
||||
}
|
||||
]
|
||||
assert captured.get("verify") is ssl_context
|
||||
assert captured.get("trust_env") is False
|
||||
assert captured.get("follow_redirects") is False
|
||||
assert not (
|
||||
ssl_context.verify_flags & ssl.VERIFY_X509_STRICT
|
||||
)
|
||||
|
||||
|
||||
def test_route_missing_bundle_does_not_crash(client):
|
||||
"""A nonexistent CA bundle path must not crash -- fall back to system CAs."""
|
||||
captured = {}
|
||||
resp = _post_test(client, captured, env={"SSL_CERT_FILE": "/nonexistent/ca-bundle.pem"})
|
||||
|
||||
assert resp.status_code == 200
|
||||
ctx = captured["verify"]
|
||||
assert isinstance(ctx, ssl.SSLContext)
|
||||
assert not (ctx.verify_flags & ssl.VERIFY_X509_STRICT)
|
||||
|
||||
|
||||
def test_route_empty_env_vars_use_system_defaults(client):
|
||||
"""Empty SSL_CERT_FILE and REQUESTS_CA_BUNDLE should not crash."""
|
||||
captured = {}
|
||||
resp = _post_test(client, captured, env={"SSL_CERT_FILE": "", "REQUESTS_CA_BUNDLE": ""})
|
||||
|
||||
assert resp.status_code == 200
|
||||
ctx = captured["verify"]
|
||||
assert isinstance(ctx, ssl.SSLContext)
|
||||
assert not (ctx.verify_flags & ssl.VERIFY_X509_STRICT)
|
||||
@@ -233,6 +233,10 @@ def test_build_uploaded_file_manifest_filters_and_nulls_unreadable_paths(monkeyp
|
||||
)
|
||||
|
||||
assert [item["id"] for item in manifest] == ["good", "outside", "missing"]
|
||||
assert manifest[0]["type"] == "attachment_ref"
|
||||
assert manifest[0]["attachment_id"] == "good"
|
||||
assert manifest[0]["uri"] == "odysseus://attachment/good"
|
||||
assert manifest[0]["read_policy"] == "owner_checked_upload"
|
||||
assert os.path.realpath(manifest[0]["path"]) == os.path.realpath(good)
|
||||
assert manifest[1]["path"] is None
|
||||
assert manifest[2]["path"] is None
|
||||
|
||||
@@ -12,6 +12,16 @@ def _bulk_action_source() -> str:
|
||||
return text[start:end]
|
||||
|
||||
|
||||
def _function_source(name: str) -> str:
|
||||
text = _EMAIL_LIBRARY.read_text(encoding="utf-8")
|
||||
start = text.index(f"function {name}")
|
||||
next_function = text.find("\nfunction ", start + 1)
|
||||
next_async = text.find("\nasync function ", start + 1)
|
||||
candidates = [idx for idx in (next_function, next_async) if idx != -1]
|
||||
end = min(candidates) if candidates else len(text)
|
||||
return text[start:end]
|
||||
|
||||
|
||||
def test_email_bulk_read_unread_calls_provider_write_routes():
|
||||
"""Bulk read/unread must persist to IMAP/provider, not only mutate UI state.
|
||||
|
||||
@@ -34,3 +44,33 @@ def test_email_bulk_read_unread_checks_backend_success_before_syncing_cache():
|
||||
assert "data?.success === false" in src
|
||||
assert "throw new Error(data?.error" in src
|
||||
assert "_libCacheWriteBack()" in src
|
||||
|
||||
|
||||
def test_email_context_changes_clear_bulk_selection_state():
|
||||
"""IMAP UIDs are folder/account scoped, so stale bulk selections must die.
|
||||
|
||||
Folder, account, filter, quick-filter, attachment, and search basis changes
|
||||
must exit select mode before the next list/search view can run bulk actions.
|
||||
"""
|
||||
text = _EMAIL_LIBRARY.read_text(encoding="utf-8")
|
||||
reset_src = _function_source("_resetBulkSelectionForContextChange")
|
||||
fresh_src = _function_source("_resetEmailListForFreshLoad")
|
||||
add_pill_src = _function_source("_addSearchPill")
|
||||
remove_pill_src = _function_source("_removeSearchPillAt")
|
||||
search_src = text[text.index("async function _doSearch()"):text.index("// Custom dropdown", text.index("async function _doSearch()"))]
|
||||
|
||||
assert "state._selectedUids.clear()" in reset_src
|
||||
assert "state._selectMode = false" in reset_src
|
||||
assert "_updateBulkBar()" in reset_src
|
||||
|
||||
assert "_resetBulkSelectionForContextChange()" in fresh_src
|
||||
assert "_resetBulkSelectionForContextChange({ rerender: true })" in add_pill_src
|
||||
assert "_resetBulkSelectionForContextChange({ rerender: true })" in remove_pill_src
|
||||
assert "_resetBulkSelectionForContextChange({ rerender: true })" in search_src
|
||||
|
||||
assert "state._libFolder = e.target.value;" in text
|
||||
assert "state._libFilter = e.target.value;" in text
|
||||
assert "state._libHasAttachments = !state._libHasAttachments;" in text
|
||||
assert "state._libAccountId = btn.dataset.accId || null;" in text
|
||||
assert text.count("_loadEmailsFresh();") >= 5
|
||||
assert "state._libSearchDraft = input.value;" in text
|
||||
|
||||
@@ -0,0 +1,108 @@
|
||||
"""Regression: IMAP calls must use uid() not search()/fetch().
|
||||
|
||||
conn.search() / conn.fetch() operate on volatile positional sequence
|
||||
numbers that shift whenever messages are deleted or expunged. The
|
||||
sig-learner and daily-brief actions must use conn.uid("SEARCH", ...)
|
||||
and conn.uid("FETCH", ...) which address messages by their persistent
|
||||
RFC 3501 UID (§2.3.1.1, §6.4.8).
|
||||
"""
|
||||
import pytest
|
||||
|
||||
|
||||
class _SpyImap:
|
||||
"""IMAP stub that records uid() calls and raises on search()/fetch()."""
|
||||
|
||||
def __init__(self, uid_list=b"1 2 3"):
|
||||
self._uid_list = uid_list
|
||||
self.uid_calls: list[tuple] = []
|
||||
|
||||
def select(self, *args, **kwargs):
|
||||
return "OK", []
|
||||
|
||||
def uid(self, command, *args):
|
||||
self.uid_calls.append((command,) + args)
|
||||
if command == "SEARCH":
|
||||
return "OK", [self._uid_list]
|
||||
if command == "FETCH":
|
||||
query = args[1] if len(args) > 1 else ""
|
||||
if "HEADER.FIELDS" in query:
|
||||
return "OK", [(None, b"From: Writer <writer@example.com>\r\n"
|
||||
b"Subject: Hello\r\n\r\n")]
|
||||
return "OK", [(None, b"Body text\r\n\r\nRegards,\r\nThe Writer\r\n")]
|
||||
return "OK", []
|
||||
|
||||
def search(self, *args):
|
||||
raise AssertionError("conn.search() called — must use conn.uid('SEARCH', ...) instead")
|
||||
|
||||
def fetch(self, *args):
|
||||
raise AssertionError("conn.fetch() called — must use conn.uid('FETCH', ...) instead")
|
||||
|
||||
def logout(self):
|
||||
pass
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sig_learner_uses_uid_search(monkeypatch):
|
||||
"""_pull_headers must call conn.uid('SEARCH', ...) not conn.search()."""
|
||||
from routes import email_helpers
|
||||
from src import task_endpoint
|
||||
from src.builtin_actions import action_learn_sender_signatures
|
||||
|
||||
spy = _SpyImap()
|
||||
monkeypatch.setattr(email_helpers, "_imap_connect", lambda *a, **kw: spy)
|
||||
monkeypatch.setattr(task_endpoint, "resolve_task_candidates", lambda *a, **kw: [])
|
||||
|
||||
message, ok = await action_learn_sender_signatures("alice")
|
||||
|
||||
assert ok is False # no LLM candidates — stops before LLM, after IMAP
|
||||
assert any(c[0] == "SEARCH" for c in spy.uid_calls), "uid('SEARCH', ...) was not called"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sig_learner_uses_uid_fetch(monkeypatch):
|
||||
"""_pull_headers must call conn.uid('FETCH', ...) not conn.fetch()."""
|
||||
from routes import email_helpers
|
||||
from src import task_endpoint
|
||||
from src.builtin_actions import action_learn_sender_signatures
|
||||
|
||||
spy = _SpyImap()
|
||||
monkeypatch.setattr(email_helpers, "_imap_connect", lambda *a, **kw: spy)
|
||||
monkeypatch.setattr(task_endpoint, "resolve_task_candidates", lambda *a, **kw: [])
|
||||
|
||||
await action_learn_sender_signatures("alice")
|
||||
|
||||
assert any(c[0] == "FETCH" for c in spy.uid_calls), "uid('FETCH', ...) was not called"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_daily_brief_uses_uid_commands(monkeypatch):
|
||||
"""action_daily_brief email section must use uid() not search()/fetch()."""
|
||||
from core import database
|
||||
from core import auth as _auth_mod
|
||||
from routes import email_helpers
|
||||
from src.builtin_actions import action_daily_brief
|
||||
|
||||
class _Q:
|
||||
def filter(self, *a, **kw): return self
|
||||
def join(self, *a, **kw): return self
|
||||
def order_by(self, *a): return self
|
||||
def all(self): return []
|
||||
|
||||
class _Db:
|
||||
def query(self, *a): return _Q()
|
||||
def close(self): pass
|
||||
|
||||
class _FakeAuth:
|
||||
is_configured = False
|
||||
|
||||
monkeypatch.setattr(database, "SessionLocal", _Db)
|
||||
monkeypatch.setattr(_auth_mod, "AuthManager", lambda: _FakeAuth())
|
||||
|
||||
spy = _SpyImap(uid_list=b"10 20 30")
|
||||
monkeypatch.setattr(email_helpers, "_imap_connect", lambda *a, **kw: spy)
|
||||
|
||||
message, ok = await action_daily_brief("")
|
||||
|
||||
assert ok is True
|
||||
assert any(c[0] == "SEARCH" for c in spy.uid_calls), "uid('SEARCH', ...) was not called"
|
||||
assert any(c[0] == "FETCH" for c in spy.uid_calls), "uid('FETCH', ...) was not called"
|
||||
@@ -1,4 +1,7 @@
|
||||
"""Kimi Code User-Agent fallback list and 403 detection."""
|
||||
import pytest
|
||||
|
||||
from src import llm_core
|
||||
from src.llm_core import (
|
||||
KIMI_CODE_USER_AGENTS,
|
||||
KIMI_CODE_USER_AGENT,
|
||||
@@ -12,6 +15,35 @@ from src.llm_core import (
|
||||
)
|
||||
|
||||
|
||||
KIMI_CHAT_URL = "https://api.kimi.com/coding/v1/chat/completions"
|
||||
|
||||
|
||||
class _Resp:
|
||||
def __init__(self, status, text="{}"):
|
||||
self.status_code = status
|
||||
self.content = text.encode()
|
||||
self.text = text
|
||||
|
||||
|
||||
class _FakeStreamResp(_Resp):
|
||||
async def aiter_lines(self):
|
||||
yield "data: [DONE]"
|
||||
|
||||
async def aread(self):
|
||||
return b""
|
||||
|
||||
|
||||
class _FakeStreamCtx:
|
||||
def __init__(self, response):
|
||||
self.response = response
|
||||
|
||||
async def __aenter__(self):
|
||||
return self.response
|
||||
|
||||
async def __aexit__(self, *args):
|
||||
return False
|
||||
|
||||
|
||||
class TestKimiCodeUserAgents:
|
||||
def test_default_is_first_fallback(self):
|
||||
assert KIMI_CODE_USER_AGENT == KIMI_CODE_USER_AGENTS[0]
|
||||
@@ -29,9 +61,8 @@ class TestKimiCodeUserAgents:
|
||||
|
||||
def test_ua_candidates_prefers_cache(self):
|
||||
_kimi_code_ua_cache.clear()
|
||||
url = "https://api.kimi.com/coding/v1/chat/completions"
|
||||
_remember_kimi_code_user_agent(url, "Kilo-Code/1.0")
|
||||
candidates = _kimi_code_ua_candidates(url)
|
||||
_remember_kimi_code_user_agent(KIMI_CHAT_URL, "Kilo-Code/1.0")
|
||||
candidates = _kimi_code_ua_candidates(KIMI_CHAT_URL)
|
||||
assert candidates[0] == "Kilo-Code/1.0"
|
||||
assert len(candidates) == len(KIMI_CODE_USER_AGENTS)
|
||||
_kimi_code_ua_cache.clear()
|
||||
@@ -48,22 +79,134 @@ class TestKimiCodeUserAgents:
|
||||
_kimi_code_ua_cache.clear()
|
||||
calls = []
|
||||
|
||||
class _Resp:
|
||||
def __init__(self, status, text=""):
|
||||
self.status_code = status
|
||||
self.content = text.encode()
|
||||
self.text = text
|
||||
|
||||
def fake_post(url, headers=None, **kwargs):
|
||||
calls.append(headers.get("User-Agent"))
|
||||
if headers.get("User-Agent") == KIMI_CODE_USER_AGENTS[0]:
|
||||
return _Resp(403, '{"error":{"type":"access_terminated_error"}}')
|
||||
return _Resp(200, "{}")
|
||||
|
||||
monkeypatch.setattr(llm_core.httpx, "get", lambda *a, **k: (_ for _ in ()).throw(RuntimeError()))
|
||||
monkeypatch.setattr("src.llm_core.httpx.post", fake_post)
|
||||
url = "https://api.kimi.com/coding/v1/chat/completions"
|
||||
r = httpx_post_kimi_aware(url, {"Authorization": "Bearer x"}, json={})
|
||||
r = httpx_post_kimi_aware(KIMI_CHAT_URL, {"Authorization": "Bearer x"}, json={})
|
||||
assert r.status_code == 200
|
||||
assert calls[0] == KIMI_CODE_USER_AGENTS[0]
|
||||
assert calls[1] == KIMI_CODE_USER_AGENTS[1]
|
||||
_kimi_code_ua_cache.clear()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_post_uses_async_probe_not_sync_httpx_get(self, monkeypatch):
|
||||
_kimi_code_ua_cache.clear()
|
||||
|
||||
class FakeClient:
|
||||
def __init__(self):
|
||||
self.get_user_agents = []
|
||||
self.post_user_agents = []
|
||||
|
||||
async def get(self, url, headers=None, **kwargs):
|
||||
self.get_user_agents.append(headers.get("User-Agent"))
|
||||
if headers.get("User-Agent") == KIMI_CODE_USER_AGENTS[0]:
|
||||
return _Resp(403, '{"error":{"type":"access_terminated_error"}}')
|
||||
return _Resp(200)
|
||||
|
||||
async def post(self, url, headers=None, **kwargs):
|
||||
self.post_user_agents.append(headers.get("User-Agent"))
|
||||
return _Resp(200)
|
||||
|
||||
def forbidden_sync_get(*args, **kwargs):
|
||||
raise AssertionError("async Kimi path must not call sync httpx.get")
|
||||
|
||||
client = FakeClient()
|
||||
monkeypatch.setattr(llm_core.httpx, "get", forbidden_sync_get)
|
||||
|
||||
r = await llm_core.httpx_post_kimi_aware_async(
|
||||
client,
|
||||
KIMI_CHAT_URL,
|
||||
{"Authorization": "Bearer x"},
|
||||
json={},
|
||||
)
|
||||
|
||||
assert r.status_code == 200
|
||||
assert client.get_user_agents == [KIMI_CODE_USER_AGENTS[0], KIMI_CODE_USER_AGENTS[1]]
|
||||
assert client.post_user_agents == [KIMI_CODE_USER_AGENTS[1]]
|
||||
assert _kimi_code_ua_cache[_kimi_code_base_key(KIMI_CHAT_URL)] == KIMI_CODE_USER_AGENTS[1]
|
||||
_kimi_code_ua_cache.clear()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_post_preserves_fallback_when_probe_fails(self, monkeypatch):
|
||||
_kimi_code_ua_cache.clear()
|
||||
|
||||
class FakeClient:
|
||||
def __init__(self):
|
||||
self.post_user_agents = []
|
||||
|
||||
async def get(self, url, headers=None, **kwargs):
|
||||
raise RuntimeError("models probe unavailable")
|
||||
|
||||
async def post(self, url, headers=None, **kwargs):
|
||||
self.post_user_agents.append(headers.get("User-Agent"))
|
||||
if headers.get("User-Agent") == KIMI_CODE_USER_AGENTS[0]:
|
||||
return _Resp(403, '{"error":{"type":"access_terminated_error"}}')
|
||||
return _Resp(200)
|
||||
|
||||
def forbidden_sync_get(*args, **kwargs):
|
||||
raise AssertionError("async Kimi path must not call sync httpx.get")
|
||||
|
||||
client = FakeClient()
|
||||
monkeypatch.setattr(llm_core.httpx, "get", forbidden_sync_get)
|
||||
|
||||
r = await llm_core.httpx_post_kimi_aware_async(
|
||||
client,
|
||||
KIMI_CHAT_URL,
|
||||
{"Authorization": "Bearer x"},
|
||||
json={},
|
||||
)
|
||||
|
||||
assert r.status_code == 200
|
||||
assert client.post_user_agents == [KIMI_CODE_USER_AGENTS[0], KIMI_CODE_USER_AGENTS[1]]
|
||||
assert _kimi_code_ua_cache[_kimi_code_base_key(KIMI_CHAT_URL)] == KIMI_CODE_USER_AGENTS[1]
|
||||
_kimi_code_ua_cache.clear()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_uses_async_kimi_probe_not_sync_httpx_get(self, monkeypatch):
|
||||
_kimi_code_ua_cache.clear()
|
||||
|
||||
class FakeClient:
|
||||
def __init__(self):
|
||||
self.get_user_agents = []
|
||||
self.stream_headers = []
|
||||
|
||||
async def get(self, url, headers=None, **kwargs):
|
||||
self.get_user_agents.append(headers.get("User-Agent"))
|
||||
if headers.get("User-Agent") == KIMI_CODE_USER_AGENTS[0]:
|
||||
return _Resp(403, '{"error":{"type":"access_terminated_error"}}')
|
||||
return _Resp(200)
|
||||
|
||||
def stream(self, method, url, **kwargs):
|
||||
self.stream_headers.append(kwargs.get("headers") or {})
|
||||
return _FakeStreamCtx(_FakeStreamResp(200))
|
||||
|
||||
def forbidden_sync_get(*args, **kwargs):
|
||||
raise AssertionError("streaming Kimi path must not call sync httpx.get")
|
||||
|
||||
client = FakeClient()
|
||||
monkeypatch.setattr(llm_core.httpx, "get", forbidden_sync_get)
|
||||
monkeypatch.setattr(llm_core, "_get_http_client", lambda: client)
|
||||
monkeypatch.setattr(llm_core, "_is_host_dead", lambda url: False)
|
||||
monkeypatch.setattr(llm_core, "note_model_activity", lambda *args, **kwargs: None)
|
||||
monkeypatch.setattr(llm_core, "_clear_host_dead", lambda *args, **kwargs: None)
|
||||
|
||||
chunks = [
|
||||
chunk
|
||||
async for chunk in llm_core.stream_llm(
|
||||
KIMI_CHAT_URL,
|
||||
"kimi-for-coding",
|
||||
[{"role": "user", "content": "hi"}],
|
||||
headers={"Authorization": "Bearer x"},
|
||||
)
|
||||
]
|
||||
|
||||
assert chunks == ["data: [DONE]\n\n"]
|
||||
assert client.get_user_agents == [KIMI_CODE_USER_AGENTS[0], KIMI_CODE_USER_AGENTS[1]]
|
||||
assert client.stream_headers[0]["User-Agent"] == KIMI_CODE_USER_AGENTS[1]
|
||||
assert _kimi_code_ua_cache[_kimi_code_base_key(KIMI_CHAT_URL)] == KIMI_CODE_USER_AGENTS[1]
|
||||
_kimi_code_ua_cache.clear()
|
||||
|
||||
@@ -18,12 +18,24 @@ def node_available():
|
||||
pytest.skip("node binary not on PATH")
|
||||
|
||||
|
||||
def _run_markdown_case(markdown: str, render_expr: str = "mod.mdToHtml(input)"):
|
||||
def _run_markdown_case(markdown: str, render_expr: str = "mod.mdToHtml(input)", with_katex: bool = False):
|
||||
script = textwrap.dedent(
|
||||
r"""
|
||||
import fs from 'node:fs';
|
||||
|
||||
globalThis.window = { location: { origin: 'http://localhost' }, katex: null };
|
||||
if (__WITH_KATEX__) {
|
||||
// Minimal stand-in for the CDN katex global: wraps the source so tests
|
||||
// can assert what was (or wasn't) handed to KaTeX.
|
||||
const katexStub = {
|
||||
renderToString(src, opts) {
|
||||
const display = !!(opts && opts.displayMode);
|
||||
return `<span class="katex" data-display="${display}">${src}</span>`;
|
||||
},
|
||||
};
|
||||
globalThis.window.katex = katexStub;
|
||||
globalThis.katex = katexStub;
|
||||
}
|
||||
globalThis.document = {
|
||||
readyState: 'loading',
|
||||
addEventListener() {},
|
||||
@@ -77,7 +89,9 @@ def _run_markdown_case(markdown: str, render_expr: str = "mod.mdToHtml(input)"):
|
||||
const input = JSON.parse(process.argv[1]);
|
||||
console.log(JSON.stringify({ html: __RENDER_EXPR__ }));
|
||||
"""
|
||||
).replace("__RENDER_EXPR__", render_expr)
|
||||
).replace("__RENDER_EXPR__", render_expr).replace(
|
||||
"__WITH_KATEX__", "true" if with_katex else "false"
|
||||
)
|
||||
result = subprocess.run(
|
||||
["node", "--input-type=module", "-e", script, json.dumps(markdown)],
|
||||
cwd=_REPO,
|
||||
@@ -200,6 +214,33 @@ def test_inline_code_content_is_html_escaped(node_available):
|
||||
assert "<b>" not in html
|
||||
|
||||
|
||||
def test_currency_dollar_amounts_are_not_rendered_as_math(node_available):
|
||||
# "$5 to $10" used to pair the two dollar signs as inline-math delimiters
|
||||
# and render "5 to" through KaTeX. Pandoc-style rules now reject it: the
|
||||
# closing $ is preceded by a space and followed by a digit.
|
||||
html = _run_markdown_case(
|
||||
"The price rose from $5 to $10 overnight.", with_katex=True
|
||||
)
|
||||
|
||||
assert 'class="katex"' not in html
|
||||
assert "$5" in html
|
||||
assert "$10" in html
|
||||
|
||||
|
||||
def test_inline_math_still_renders_through_katex(node_available):
|
||||
html = _run_markdown_case("Pythagoras: $x^2 + y^2 = z^2$ holds.", with_katex=True)
|
||||
|
||||
assert '<span class="katex" data-display="false">x^2 + y^2 = z^2</span>' in html
|
||||
assert "$" not in html
|
||||
|
||||
|
||||
def test_display_math_still_renders_through_katex(node_available):
|
||||
html = _run_markdown_case("$$\\frac{a}{b}$$", with_katex=True)
|
||||
|
||||
assert 'data-display="true"' in html
|
||||
assert "$$" not in html
|
||||
|
||||
|
||||
def test_dotted_python_import_paths_are_not_autolinked(node_available):
|
||||
html = _run_markdown_case(
|
||||
"from imblearn.combine import SMOTETomek\n"
|
||||
|
||||
@@ -1,13 +1,9 @@
|
||||
"""A plain text message that merely *looks* like a JSON array of objects must
|
||||
NOT be silently re-parsed into a list on reload.
|
||||
"""Persistence contracts for JSON-like text and multimodal chat content.
|
||||
|
||||
_parse_msg_content de-serializes multimodal (image/audio) content back into a
|
||||
list of content blocks. The old heuristic accepted ANY string that started
|
||||
with "[{" and contained the substring '"type"'. A user who pasted an API
|
||||
schema / sample such as `[{"type": "object", "name": "foo"}]` therefore had
|
||||
their text message permanently corrupted into a Python list on the next
|
||||
session hydration. The fix restricts the round-trip to lists whose elements
|
||||
are all recognized content-block types (text/image_url/audio/...).
|
||||
Plain text that resembles a JSON content-block list must remain an exact
|
||||
string. Real provider multimodal blocks follow the durable attachment
|
||||
contract: readable text plus stable attachment metadata is persisted, while
|
||||
raw inline media bytes are omitted.
|
||||
"""
|
||||
import tempfile
|
||||
import uuid
|
||||
@@ -65,16 +61,48 @@ def test_jsonlike_user_string_not_corrupted(manager):
|
||||
assert reloaded.history[0].content == text
|
||||
|
||||
|
||||
def test_real_multimodal_content_still_round_trips(manager):
|
||||
def test_real_multimodal_content_persists_reference_without_base64(manager):
|
||||
sid = "sess-" + uuid.uuid4().hex[:8]
|
||||
_make_session(sid)
|
||||
attachment_id = "a" * 32 + ".png"
|
||||
multimodal = [
|
||||
{"type": "text", "text": "what is this?"},
|
||||
{"type": "image_url", "image_url": {"url": "data:image/png;base64,AAAA"}},
|
||||
]
|
||||
msgs = [ChatMessage(role="user", content=multimodal)]
|
||||
metadata = {
|
||||
"attachments": [
|
||||
{
|
||||
"id": attachment_id,
|
||||
"name": "diagram.png",
|
||||
"mime": "image/png",
|
||||
"size": 4,
|
||||
"checksum_sha256": "sha256-digest",
|
||||
}
|
||||
]
|
||||
}
|
||||
msgs = [ChatMessage(role="user", content=multimodal, metadata=metadata)]
|
||||
assert manager.replace_messages(sid, msgs) is True
|
||||
|
||||
expected = (
|
||||
"what is this?\n"
|
||||
"[1 inline media payload omitted]\n"
|
||||
f"[Attachment: diagram.png | id={attachment_id} | mime=image/png | "
|
||||
"size=4 bytes | sha256=sha256-digest]"
|
||||
)
|
||||
|
||||
db = _TS()
|
||||
try:
|
||||
stored = db.query(cdb.ChatMessage).filter_by(session_id=sid).one()
|
||||
assert stored.content == expected
|
||||
assert "what is this?" in stored.content
|
||||
assert attachment_id in stored.content
|
||||
assert "data:image/png;base64,AAAA" not in stored.content
|
||||
assert "base64" not in stored.content
|
||||
assert "AAAA" not in stored.content
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
manager.sessions.clear()
|
||||
reloaded = manager.get_session(sid)
|
||||
assert reloaded.history[0].content == multimodal
|
||||
assert reloaded.history[0].content == expected
|
||||
assert reloaded.history[0].metadata["attachments"][0]["id"] == attachment_id
|
||||
|
||||
@@ -1,14 +1,9 @@
|
||||
"""replace_messages must JSON-serialize multimodal (list) content.
|
||||
"""replace_messages must persist readable, path-free multimodal history.
|
||||
|
||||
A chat with an image/audio attachment carries list content. When such a
|
||||
chat is compacted, the manual-compaction path calls replace_messages with
|
||||
the retained messages. replace_messages wrote message.content straight into
|
||||
the Text column, so SQLAlchemy bound the list\'s single-quoted repr. On
|
||||
reload _parse_msg_content only de-serializes a string that contains the
|
||||
double-quoted "type", so the repr failed the check and the message came
|
||||
back as a corrupted string blob - the attachment was destroyed. The
|
||||
sibling _persist_message json.dumps-es list content; replace_messages did
|
||||
not.
|
||||
Live model input may contain provider-specific media blocks and inline data
|
||||
URLs. Compaction uses replace_messages for the retained transcript, which must
|
||||
store readable text plus stable structured attachment references without
|
||||
copying raw base64 payloads into ChatMessage.content.
|
||||
"""
|
||||
import uuid
|
||||
|
||||
@@ -27,6 +22,7 @@ def manager(monkeypatch):
|
||||
monkeypatch.setattr(sm, "SessionLocal", _TS)
|
||||
mgr = sm.SessionManager.__new__(sm.SessionManager)
|
||||
mgr.sessions = {}
|
||||
mgr.upload_handler = None
|
||||
return mgr
|
||||
|
||||
|
||||
@@ -41,33 +37,71 @@ def _make_session(sid, owner="alice"):
|
||||
db.close()
|
||||
|
||||
|
||||
def test_multimodal_content_round_trips_through_replace_messages(manager):
|
||||
def test_multimodal_content_persists_text_and_attachment_ref_without_payload(manager):
|
||||
sid = "sess-" + uuid.uuid4().hex[:8]
|
||||
_make_session(sid)
|
||||
|
||||
upload_id = "a" * 32 + ".png"
|
||||
multimodal = [
|
||||
{"type": "text", "text": "what is this?"},
|
||||
{"type": "image_url", "image_url": {"url": "data:image/png;base64,AAAA"}},
|
||||
]
|
||||
msgs = [ChatMessage(role="user", content=multimodal)]
|
||||
msgs = [ChatMessage(
|
||||
role="user",
|
||||
content=multimodal,
|
||||
metadata={
|
||||
"attachments": [{
|
||||
"id": upload_id,
|
||||
"name": "diagram.png",
|
||||
"mime": "image/png",
|
||||
"size": 4,
|
||||
"checksum_sha256": "sha256-digest",
|
||||
}]
|
||||
},
|
||||
)]
|
||||
assert manager.replace_messages(sid, msgs) is True
|
||||
|
||||
expected = (
|
||||
"what is this?\n"
|
||||
"[1 inline media payload omitted]\n"
|
||||
f"[Attachment: diagram.png | id={upload_id} | mime=image/png | "
|
||||
"size=4 bytes | sha256=sha256-digest]"
|
||||
)
|
||||
|
||||
db = _TS()
|
||||
try:
|
||||
stored = db.query(cdb.ChatMessage).filter_by(session_id=sid).one()
|
||||
assert stored.content == expected
|
||||
assert "data:image/png;base64,AAAA" not in stored.content
|
||||
assert "base64" not in stored.content
|
||||
assert "AAAA" not in stored.content
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
# Drop the in-memory cache so the next read hydrates from the DB.
|
||||
manager.sessions.clear()
|
||||
reloaded = manager.get_session(sid)
|
||||
assert len(reloaded.history) == 1
|
||||
# Content must come back as the original list, not a repr string blob.
|
||||
assert reloaded.history[0].content == multimodal
|
||||
persisted = reloaded.history[0].content
|
||||
assert isinstance(persisted, str)
|
||||
assert persisted == expected
|
||||
assert reloaded.history[0].metadata["attachments"][0]["id"] == upload_id
|
||||
assert (
|
||||
reloaded.history[0].metadata["attachments"][0]["checksum_sha256"]
|
||||
== "sha256-digest"
|
||||
)
|
||||
|
||||
|
||||
def test_plain_string_content_still_round_trips(manager):
|
||||
def test_jsonlike_plain_string_content_still_round_trips(manager):
|
||||
sid = "sess-" + uuid.uuid4().hex[:8]
|
||||
_make_session(sid)
|
||||
msgs = [ChatMessage(role="user", content="just text")]
|
||||
text = '[{"type": "object", "name": "foo"}]'
|
||||
msgs = [ChatMessage(role="user", content=text)]
|
||||
assert manager.replace_messages(sid, msgs) is True
|
||||
manager.sessions.clear()
|
||||
reloaded = manager.get_session(sid)
|
||||
assert reloaded.history[0].content == "just text"
|
||||
assert isinstance(reloaded.history[0].content, str)
|
||||
assert reloaded.history[0].content == text
|
||||
|
||||
|
||||
def test_replace_messages_keeps_history_alias_for_context_messages(manager):
|
||||
|
||||
@@ -0,0 +1,259 @@
|
||||
"""Upload lifecycle guarantees for compaction's replace_messages path."""
|
||||
|
||||
import concurrent.futures
|
||||
import json
|
||||
import os
|
||||
import threading
|
||||
import uuid
|
||||
|
||||
import pytest
|
||||
from sqlalchemy import event
|
||||
|
||||
import core.database as cdb
|
||||
import core.session_manager as session_manager_module
|
||||
from core.models import ChatMessage
|
||||
from src.upload_handler import UploadHandler
|
||||
from tests.helpers.sqlite_db import make_temp_sqlite
|
||||
|
||||
|
||||
OLD_TIMESTAMP = "2000-01-01T00:00:00"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def manager_db(monkeypatch):
|
||||
SessionLocal, engine, tmpfile = make_temp_sqlite(cdb.Base.metadata)
|
||||
monkeypatch.setattr(session_manager_module, "SessionLocal", SessionLocal)
|
||||
manager = session_manager_module.SessionManager.__new__(
|
||||
session_manager_module.SessionManager
|
||||
)
|
||||
manager.sessions = {}
|
||||
manager.upload_handler = None
|
||||
try:
|
||||
yield manager, SessionLocal, engine
|
||||
finally:
|
||||
engine.dispose()
|
||||
tmpfile.close()
|
||||
try:
|
||||
os.unlink(tmpfile.name)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
|
||||
def _seed_session(SessionLocal, *, owner="alice", content="existing durable history"):
|
||||
session_id = "replace-" + uuid.uuid4().hex
|
||||
db = SessionLocal()
|
||||
try:
|
||||
db.add(cdb.Session(
|
||||
id=session_id,
|
||||
owner=owner,
|
||||
name="Compaction reservation regression",
|
||||
model="test-model",
|
||||
endpoint_url="http://localhost:11434",
|
||||
archived=False,
|
||||
message_count=1,
|
||||
))
|
||||
db.add(cdb.ChatMessage(
|
||||
id="message-" + uuid.uuid4().hex,
|
||||
session_id=session_id,
|
||||
role="user",
|
||||
content=content,
|
||||
meta_data=json.dumps({"source": "before-replacement"}),
|
||||
))
|
||||
db.commit()
|
||||
finally:
|
||||
db.close()
|
||||
return session_id
|
||||
|
||||
|
||||
def _attachment_message(upload_id, text):
|
||||
return ChatMessage(
|
||||
role="user",
|
||||
content=text,
|
||||
metadata={
|
||||
"attachments": [{
|
||||
"id": upload_id,
|
||||
"name": f"{text}.txt",
|
||||
"mime": "text/plain",
|
||||
"size": len(text),
|
||||
}]
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def _durable_messages(SessionLocal, session_id):
|
||||
db = SessionLocal()
|
||||
try:
|
||||
return [
|
||||
(message.role, message.content, message.meta_data)
|
||||
for message in db.query(cdb.ChatMessage)
|
||||
.filter(cdb.ChatMessage.session_id == session_id)
|
||||
.order_by(cdb.ChatMessage.timestamp, cdb.ChatMessage.id)
|
||||
.all()
|
||||
]
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
|
||||
def test_replace_messages_reserves_every_incoming_attachment_before_delete(
|
||||
manager_db,
|
||||
monkeypatch,
|
||||
):
|
||||
manager, SessionLocal, engine = manager_db
|
||||
session_id = _seed_session(SessionLocal)
|
||||
manager.upload_handler = object()
|
||||
incoming = [
|
||||
_attachment_message("1" * 32 + ".txt", "first"),
|
||||
_attachment_message("2" * 32 + ".txt", "second"),
|
||||
]
|
||||
message_mutations = []
|
||||
reservation_calls = []
|
||||
|
||||
def record_sql(_conn, _cursor, statement, _parameters, _context, _executemany):
|
||||
normalized = statement.lstrip().upper()
|
||||
if normalized.startswith(("DELETE FROM CHAT_MESSAGES", "INSERT INTO CHAT_MESSAGES")):
|
||||
message_mutations.append(normalized.split(maxsplit=1)[0])
|
||||
|
||||
def reserve(handler, owner, content, metadata):
|
||||
assert message_mutations == []
|
||||
reservation_calls.append((handler, owner, content, metadata))
|
||||
return None
|
||||
|
||||
event.listen(engine, "before_cursor_execute", record_sql)
|
||||
monkeypatch.setattr(
|
||||
session_manager_module,
|
||||
"reserve_message_upload_references",
|
||||
reserve,
|
||||
)
|
||||
try:
|
||||
assert manager.replace_messages(session_id, incoming) is True
|
||||
finally:
|
||||
event.remove(engine, "before_cursor_execute", record_sql)
|
||||
|
||||
assert [call[2] for call in reservation_calls] == ["first", "second"]
|
||||
assert all(call[0] is manager.upload_handler for call in reservation_calls)
|
||||
assert all(call[1] == "alice" for call in reservation_calls)
|
||||
assert message_mutations == ["DELETE", "INSERT"]
|
||||
|
||||
|
||||
def test_replace_messages_reservation_failure_leaves_durable_history_unchanged(
|
||||
manager_db,
|
||||
monkeypatch,
|
||||
):
|
||||
manager, SessionLocal, engine = manager_db
|
||||
session_id = _seed_session(SessionLocal)
|
||||
before = _durable_messages(SessionLocal, session_id)
|
||||
manager.upload_handler = object()
|
||||
missing_upload_id = "4" * 32 + ".txt"
|
||||
incoming = [
|
||||
_attachment_message("3" * 32 + ".txt", "available"),
|
||||
_attachment_message(missing_upload_id, "missing"),
|
||||
]
|
||||
calls = []
|
||||
message_mutations = []
|
||||
|
||||
def record_sql(_conn, _cursor, statement, _parameters, _context, _executemany):
|
||||
normalized = statement.lstrip().upper()
|
||||
if normalized.startswith(("DELETE FROM CHAT_MESSAGES", "INSERT INTO CHAT_MESSAGES")):
|
||||
message_mutations.append(normalized.split(maxsplit=1)[0])
|
||||
|
||||
def reserve(_handler, _owner, content, _metadata):
|
||||
calls.append(content)
|
||||
return missing_upload_id if content == "missing" else None
|
||||
|
||||
event.listen(engine, "before_cursor_execute", record_sql)
|
||||
monkeypatch.setattr(
|
||||
session_manager_module,
|
||||
"reserve_message_upload_references",
|
||||
reserve,
|
||||
)
|
||||
try:
|
||||
assert manager.replace_messages(session_id, incoming) is False
|
||||
finally:
|
||||
event.remove(engine, "before_cursor_execute", record_sql)
|
||||
|
||||
assert calls == ["available", "missing"]
|
||||
assert message_mutations == []
|
||||
assert _durable_messages(SessionLocal, session_id) == before
|
||||
assert [message.content for message in manager.sessions[session_id].history] == [
|
||||
"existing durable history"
|
||||
]
|
||||
assert all("_db_id" not in (message.metadata or {}) for message in incoming)
|
||||
|
||||
|
||||
def test_cleanup_cannot_delete_attachment_during_concurrent_compaction_replacement(
|
||||
manager_db,
|
||||
monkeypatch,
|
||||
tmp_path,
|
||||
):
|
||||
manager, SessionLocal, _engine = manager_db
|
||||
session_id = _seed_session(SessionLocal)
|
||||
base_dir = tmp_path / "base"
|
||||
upload_dir = tmp_path / "uploads"
|
||||
base_dir.mkdir()
|
||||
upload_dir.mkdir()
|
||||
handler = UploadHandler(str(base_dir), str(upload_dir))
|
||||
manager.upload_handler = handler
|
||||
|
||||
upload_id = "5" * 32 + ".txt"
|
||||
upload_hash = "6" * 64
|
||||
dated_dir = upload_dir / "2000" / "01" / "01"
|
||||
dated_dir.mkdir(parents=True)
|
||||
upload_path = dated_dir / upload_id
|
||||
upload_path.write_text("attachment retained by compaction", encoding="utf-8")
|
||||
upload_row = {
|
||||
"id": upload_id,
|
||||
"path": str(upload_path),
|
||||
"mime": "text/plain",
|
||||
"size": upload_path.stat().st_size,
|
||||
"name": "compaction.txt",
|
||||
"original_name": "compaction.txt",
|
||||
"hash": upload_hash,
|
||||
"checksum_sha256": upload_hash,
|
||||
"uploaded_at": OLD_TIMESTAMP,
|
||||
"created_at": OLD_TIMESTAMP,
|
||||
"last_accessed": OLD_TIMESTAMP,
|
||||
"owner": "alice",
|
||||
}
|
||||
(upload_dir / "uploads.json").write_text(
|
||||
json.dumps({f"alice:{upload_hash}": upload_row}),
|
||||
encoding="utf-8",
|
||||
)
|
||||
handler._index_cache = None
|
||||
|
||||
reservation_write_entered = threading.Event()
|
||||
release_reservation_write = threading.Event()
|
||||
real_atomic_write = handler._atomic_write_json
|
||||
|
||||
def block_reservation_write(path, data, *, sync_backup=False):
|
||||
refreshed = any(
|
||||
isinstance(row, dict)
|
||||
and row.get("id") == upload_id
|
||||
and row.get("last_accessed") != OLD_TIMESTAMP
|
||||
for row in data.values()
|
||||
)
|
||||
if sync_backup and refreshed and not reservation_write_entered.is_set():
|
||||
reservation_write_entered.set()
|
||||
assert release_reservation_write.wait(5)
|
||||
return real_atomic_write(path, data, sync_backup=sync_backup)
|
||||
|
||||
monkeypatch.setattr(handler, "_atomic_write_json", block_reservation_write)
|
||||
incoming = [_attachment_message(upload_id, "retained after compaction")]
|
||||
|
||||
with concurrent.futures.ThreadPoolExecutor(max_workers=2) as pool:
|
||||
replace_future = pool.submit(manager.replace_messages, session_id, incoming)
|
||||
assert reservation_write_entered.wait(5)
|
||||
cleanup_future = pool.submit(handler.cleanup_old_uploads, set(), set())
|
||||
try:
|
||||
with pytest.raises(concurrent.futures.TimeoutError):
|
||||
cleanup_future.result(timeout=0.1)
|
||||
finally:
|
||||
release_reservation_write.set()
|
||||
|
||||
assert replace_future.result(timeout=5) is True
|
||||
assert cleanup_future.result(timeout=5) == 0
|
||||
|
||||
assert upload_path.is_file()
|
||||
assert handler.resolve_upload(upload_id, owner="alice") is not None
|
||||
durable = _durable_messages(SessionLocal, session_id)
|
||||
assert len(durable) == 1
|
||||
assert json.loads(durable[0][2])["attachments"][0]["id"] == upload_id
|
||||
@@ -0,0 +1,92 @@
|
||||
"""Regression: the tool-execution task inside stream_agent_loop must be
|
||||
cancelled (not orphaned) when the SSE consumer stops draining the generator
|
||||
early — e.g. a client disconnect mid tool-call.
|
||||
|
||||
The drain loop in stream_agent_loop:
|
||||
|
||||
_tool_task = asyncio.create_task(_run_tool())
|
||||
while True:
|
||||
evt = await _progress_q.get()
|
||||
if evt is None:
|
||||
break
|
||||
yield ...
|
||||
desc, result = await _tool_task
|
||||
|
||||
used to have no try/finally around it. If the generator is closed while
|
||||
suspended on `await _progress_q.get()` (which is exactly what Starlette does
|
||||
via `aclose()` when an SSE client disconnects), GeneratorExit is thrown at
|
||||
that point and `_tool_task` is abandoned mid-flight — never awaited, never
|
||||
cancelled. For a long-running `bash`/`python` tool this orphans the
|
||||
subprocess server-side with nothing left to reap it.
|
||||
|
||||
The fix wraps the drain loop in try/finally and cancels+awaits `_tool_task`
|
||||
on early exit. This test drives the real stream_agent_loop with a fake tool
|
||||
handler that sleeps until cancelled, closes the generator mid-tool-call (the
|
||||
same way a dropped SSE connection would), and asserts the fake handler
|
||||
actually observed cancellation.
|
||||
"""
|
||||
import asyncio
|
||||
import json
|
||||
|
||||
import src.agent_loop as al
|
||||
|
||||
|
||||
def test_tool_task_cancelled_on_generator_close(monkeypatch):
|
||||
cancelled = {"v": False}
|
||||
|
||||
async def _slow_exec(block, *a, progress_cb=None, **k):
|
||||
if progress_cb:
|
||||
await progress_cb({"elapsed_s": 1, "tail": "running"})
|
||||
try:
|
||||
await asyncio.sleep(60)
|
||||
except asyncio.CancelledError:
|
||||
cancelled["v"] = True
|
||||
raise
|
||||
return ("bash", {"output": "ok", "exit_code": 0})
|
||||
|
||||
monkeypatch.setattr(al, "get_setting", lambda key, default=None: default, raising=False)
|
||||
monkeypatch.setattr(al, "get_mcp_manager", lambda: None, raising=False)
|
||||
monkeypatch.setattr(al, "estimate_tokens", lambda *a, **k: 10, raising=False)
|
||||
monkeypatch.setattr(al, "execute_tool_block", _slow_exec, raising=False)
|
||||
|
||||
native_calls = [{"name": "bash", "arguments": json.dumps({"command": "sleep 60"})}]
|
||||
|
||||
async def _fake_stream(_candidates, messages, **kwargs):
|
||||
yield f'data: {json.dumps({"delta": "Running it now."})}\n\n'
|
||||
yield f'data: {json.dumps({"type": "tool_calls", "calls": native_calls})}\n\n'
|
||||
yield "data: [DONE]\n\n"
|
||||
|
||||
monkeypatch.setattr(al, "stream_llm_with_fallback", _fake_stream, raising=False)
|
||||
|
||||
async def _run():
|
||||
gen = al.stream_agent_loop(
|
||||
"https://api.openai.com/v1", "gpt-4o",
|
||||
[{"role": "user", "content": "run sleep 60"}],
|
||||
max_rounds=2,
|
||||
relevant_tools={"bash"},
|
||||
)
|
||||
saw_tool_start = False
|
||||
saw_tool_progress = False
|
||||
async for chunk in gen:
|
||||
if '"type": "tool_start"' in chunk:
|
||||
saw_tool_start = True
|
||||
elif '"type": "tool_progress" ' in chunk or '"type": "tool_progress"' in chunk:
|
||||
saw_tool_progress = True
|
||||
break
|
||||
assert saw_tool_start, "expected a tool_start event before the tool ran"
|
||||
assert saw_tool_progress, "expected a tool_progress event once the fake tool started (task must exist by now)"
|
||||
# Simulate an SSE client disconnecting mid tool-call: close the
|
||||
# generator while it is suspended awaiting the next progress event.
|
||||
await gen.aclose()
|
||||
# Assert *inside* this coroutine, immediately after aclose() returns.
|
||||
# asyncio.run()'s own shutdown sequence cancels any tasks still
|
||||
# pending once _run() itself completes — checking after asyncio.run()
|
||||
# returns would pass even with the bug, because that unrelated
|
||||
# cleanup would cancel the orphaned task anyway and mask the fix.
|
||||
assert cancelled["v"] is True, (
|
||||
"tool task must be cancelled by stream_agent_loop's own cleanup "
|
||||
"on generator close, not left running until asyncio.run() tears "
|
||||
"down the loop"
|
||||
)
|
||||
|
||||
asyncio.run(_run())
|
||||
@@ -0,0 +1,831 @@
|
||||
import asyncio
|
||||
import concurrent.futures
|
||||
import json
|
||||
import os
|
||||
import threading
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
from fastapi import HTTPException
|
||||
|
||||
from core.database import (
|
||||
Base,
|
||||
ChatMessage as DbChatMessage,
|
||||
CalendarCal,
|
||||
CalendarEvent,
|
||||
Document,
|
||||
DocumentVersion,
|
||||
GalleryImage,
|
||||
Note,
|
||||
Session as DbSession,
|
||||
)
|
||||
from src.upload_handler import (
|
||||
UploadCleanupSafetyError,
|
||||
UploadHandler,
|
||||
extract_internal_upload_ids,
|
||||
reserve_message_upload_references,
|
||||
reserve_upload_references,
|
||||
)
|
||||
from tests.helpers.sqlite_db import make_temp_sqlite
|
||||
|
||||
|
||||
OLD_TIMESTAMP = "2000-01-01T00:00:00"
|
||||
|
||||
|
||||
class _AdminAuth:
|
||||
is_configured = True
|
||||
|
||||
@staticmethod
|
||||
def is_admin(user):
|
||||
return user == "admin"
|
||||
|
||||
|
||||
class _AdminRequest:
|
||||
headers = {}
|
||||
state = SimpleNamespace(current_user="admin")
|
||||
app = SimpleNamespace(state=SimpleNamespace(auth_manager=_AdminAuth()))
|
||||
|
||||
|
||||
def _make_handler(tmp_path: Path) -> UploadHandler:
|
||||
base_dir = tmp_path / "base"
|
||||
upload_dir = tmp_path / "uploads"
|
||||
base_dir.mkdir()
|
||||
upload_dir.mkdir()
|
||||
return UploadHandler(str(base_dir), str(upload_dir))
|
||||
|
||||
|
||||
def _seed_old_uploads(handler: UploadHandler, rows: list[dict]) -> dict[str, Path]:
|
||||
dated_dir = Path(handler.upload_dir) / "2000" / "01" / "01"
|
||||
dated_dir.mkdir(parents=True)
|
||||
index = {}
|
||||
paths = {}
|
||||
for row in rows:
|
||||
upload_id = row["id"]
|
||||
path = dated_dir / upload_id
|
||||
path.write_bytes(row.get("bytes", upload_id.encode("ascii")))
|
||||
info = {
|
||||
"id": upload_id,
|
||||
"path": str(path),
|
||||
"mime": row.get("mime", "application/octet-stream"),
|
||||
"size": path.stat().st_size,
|
||||
"name": row.get("name", upload_id),
|
||||
"original_name": row.get("name", upload_id),
|
||||
"hash": row["hash"],
|
||||
"checksum_sha256": row["hash"],
|
||||
"uploaded_at": row.get("uploaded_at", OLD_TIMESTAMP),
|
||||
"created_at": row.get("created_at", OLD_TIMESTAMP),
|
||||
"last_accessed": row.get("last_accessed", OLD_TIMESTAMP),
|
||||
"owner": row.get("owner", "alice"),
|
||||
}
|
||||
index[f"{info['owner']}:{info['hash']}"] = info
|
||||
paths[upload_id] = path
|
||||
|
||||
Path(handler.upload_dir, "uploads.json").write_text(
|
||||
json.dumps(index),
|
||||
encoding="utf-8",
|
||||
)
|
||||
handler._index_cache = None
|
||||
return paths
|
||||
|
||||
|
||||
def _manual_cleanup_endpoint(handler: UploadHandler, monkeypatch):
|
||||
import fastapi.dependencies.utils as dependency_utils
|
||||
from routes.upload_routes import router, setup_upload_routes
|
||||
|
||||
monkeypatch.setattr(dependency_utils, "ensure_multipart_is_installed", lambda: None)
|
||||
before = len(router.routes)
|
||||
setup_upload_routes(handler)
|
||||
return {
|
||||
route.endpoint.__name__: route.endpoint
|
||||
for route in router.routes[before:]
|
||||
}["manual_cleanup"]
|
||||
|
||||
|
||||
def _reference_database(monkeypatch, *, upload_id: str, gallery_hash: str = None):
|
||||
from routes import upload_routes
|
||||
|
||||
SessionLocal, engine, tmpfile = make_temp_sqlite(Base.metadata)
|
||||
db = SessionLocal()
|
||||
try:
|
||||
db.add(DbSession(
|
||||
id="session-1",
|
||||
name="Cleanup regression",
|
||||
endpoint_url="http://localhost",
|
||||
model="test-model",
|
||||
owner="alice",
|
||||
))
|
||||
db.add(DbChatMessage(
|
||||
id="message-1",
|
||||
session_id="session-1",
|
||||
role="user",
|
||||
content=f"[Attachment: retained.png | id={upload_id} | mime=image/png]",
|
||||
meta_data=json.dumps({
|
||||
"attachments": [{
|
||||
"id": upload_id,
|
||||
"name": "retained.png",
|
||||
"mime": "image/png",
|
||||
"size": 8,
|
||||
}]
|
||||
}),
|
||||
))
|
||||
if gallery_hash:
|
||||
db.add(GalleryImage(
|
||||
id="gallery-cleanup-reference",
|
||||
filename="abcdef123456.png",
|
||||
prompt="Chat upload",
|
||||
owner="alice",
|
||||
file_hash=gallery_hash,
|
||||
))
|
||||
db.commit()
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
monkeypatch.setattr(upload_routes, "SessionLocal", SessionLocal)
|
||||
return engine, tmpfile
|
||||
|
||||
|
||||
def test_admin_cleanup_preserves_referenced_upload_and_reconciles_deleted_row(
|
||||
tmp_path,
|
||||
monkeypatch,
|
||||
):
|
||||
handler = _make_handler(tmp_path)
|
||||
referenced_id = "a" * 32 + ".png"
|
||||
unreferenced_id = "b" * 32 + ".txt"
|
||||
gallery_id = "7" * 32 + ".png"
|
||||
gallery_hash = "7" * 64
|
||||
paths = _seed_old_uploads(handler, [
|
||||
{
|
||||
"id": referenced_id,
|
||||
"hash": "1" * 64,
|
||||
"mime": "image/png",
|
||||
},
|
||||
{
|
||||
"id": unreferenced_id,
|
||||
"hash": "2" * 64,
|
||||
"mime": "text/plain",
|
||||
},
|
||||
{
|
||||
"id": gallery_id,
|
||||
"hash": gallery_hash,
|
||||
"mime": "image/png",
|
||||
},
|
||||
])
|
||||
engine, tmpfile = _reference_database(
|
||||
monkeypatch,
|
||||
upload_id=referenced_id,
|
||||
gallery_hash=gallery_hash,
|
||||
)
|
||||
|
||||
try:
|
||||
response = asyncio.run(
|
||||
_manual_cleanup_endpoint(handler, monkeypatch)(_AdminRequest())
|
||||
)
|
||||
finally:
|
||||
engine.dispose()
|
||||
tmpfile.close()
|
||||
try:
|
||||
os.unlink(tmpfile.name)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
assert response == {"status": "success", "files_cleaned": 1}
|
||||
assert paths[referenced_id].is_file()
|
||||
referenced_info = handler.get_upload_info(referenced_id)
|
||||
assert referenced_info is not None
|
||||
assert handler.resolve_upload(referenced_id, owner="alice") is not None
|
||||
assert paths[gallery_id].is_file()
|
||||
assert handler.get_upload_info(gallery_id) is not None
|
||||
|
||||
assert not paths[unreferenced_id].exists()
|
||||
assert handler.get_upload_info(unreferenced_id) is None
|
||||
assert handler.resolve_upload(unreferenced_id, owner="alice") is None
|
||||
|
||||
live_index = json.loads(
|
||||
Path(handler.upload_dir, "uploads.json").read_text(encoding="utf-8")
|
||||
)
|
||||
assert {info["id"] for info in live_index.values()} == {
|
||||
referenced_id,
|
||||
gallery_id,
|
||||
}
|
||||
backup_index = json.loads(
|
||||
Path(handler.upload_dir, "uploads.json.bak").read_text(encoding="utf-8")
|
||||
)
|
||||
assert {info["id"] for info in backup_index.values()} == {
|
||||
referenced_id,
|
||||
gallery_id,
|
||||
}
|
||||
|
||||
# Recovery must not resurrect the deliberately deleted row.
|
||||
Path(handler.upload_dir, "uploads.json").write_text("{broken", encoding="utf-8")
|
||||
handler._index_cache = None
|
||||
assert handler.get_upload_info(unreferenced_id) is None
|
||||
assert paths[referenced_id].parent.is_dir()
|
||||
|
||||
|
||||
def test_cleanup_retains_upload_and_all_rows_when_index_rows_disagree(tmp_path):
|
||||
handler = _make_handler(tmp_path)
|
||||
upload_id = "c" * 32 + ".txt"
|
||||
path = _seed_old_uploads(handler, [{
|
||||
"id": upload_id,
|
||||
"hash": "1" * 64,
|
||||
"mime": "text/plain",
|
||||
"owner": "alice",
|
||||
}])[upload_id]
|
||||
index_path = Path(handler.upload_dir, "uploads.json")
|
||||
index = json.loads(index_path.read_text(encoding="utf-8"))
|
||||
alice_row = next(iter(index.values()))
|
||||
index["bob:" + "2" * 64] = {
|
||||
**alice_row,
|
||||
"owner": "bob",
|
||||
"hash": "2" * 64,
|
||||
"checksum_sha256": "2" * 64,
|
||||
}
|
||||
index_path.write_text(json.dumps(index), encoding="utf-8")
|
||||
handler._index_cache = None
|
||||
|
||||
assert handler.cleanup_old_uploads(set(), set()) == 0
|
||||
assert path.is_file()
|
||||
assert json.loads(index_path.read_text(encoding="utf-8")) == index
|
||||
|
||||
|
||||
def test_cleanup_retains_lone_row_without_authoritative_lifecycle_metadata(tmp_path):
|
||||
handler = _make_handler(tmp_path)
|
||||
upload_id = "6" * 32 + ".txt"
|
||||
path = _seed_old_uploads(handler, [{
|
||||
"id": upload_id,
|
||||
"hash": "6" * 64,
|
||||
"mime": "text/plain",
|
||||
}])[upload_id]
|
||||
index_path = Path(handler.upload_dir, "uploads.json")
|
||||
index = json.loads(index_path.read_text(encoding="utf-8"))
|
||||
row = next(iter(index.values()))
|
||||
for field in (
|
||||
"owner",
|
||||
"hash",
|
||||
"checksum_sha256",
|
||||
"uploaded_at",
|
||||
"created_at",
|
||||
"last_accessed",
|
||||
):
|
||||
row.pop(field)
|
||||
index_path.write_text(json.dumps(index), encoding="utf-8")
|
||||
handler._index_cache = None
|
||||
|
||||
assert handler.cleanup_old_uploads(set(), set()) == 0
|
||||
assert path.is_file()
|
||||
assert json.loads(index_path.read_text(encoding="utf-8")) == index
|
||||
|
||||
|
||||
def test_reservation_and_cleanup_are_serialized_without_dangling_references(
|
||||
tmp_path,
|
||||
monkeypatch,
|
||||
):
|
||||
# Writer wins: reservation holds the shared index lock, refreshes access,
|
||||
# then cleanup observes the refreshed row and preserves the file.
|
||||
writer_root = tmp_path / "writer-wins"
|
||||
writer_root.mkdir()
|
||||
writer_handler = _make_handler(writer_root)
|
||||
upload_id = "2" * 32 + ".txt"
|
||||
writer_path = _seed_old_uploads(writer_handler, [{
|
||||
"id": upload_id,
|
||||
"hash": "2" * 64,
|
||||
"mime": "text/plain",
|
||||
}])[upload_id]
|
||||
write_entered = threading.Event()
|
||||
release_write = threading.Event()
|
||||
real_atomic_write = writer_handler._atomic_write_json
|
||||
|
||||
def blocking_reservation_write(path, data, *, sync_backup=False):
|
||||
refreshed = any(
|
||||
isinstance(row, dict) and row.get("last_accessed") != OLD_TIMESTAMP
|
||||
for row in data.values()
|
||||
)
|
||||
if sync_backup and refreshed and not write_entered.is_set():
|
||||
write_entered.set()
|
||||
assert release_write.wait(5)
|
||||
return real_atomic_write(path, data, sync_backup=sync_backup)
|
||||
|
||||
monkeypatch.setattr(writer_handler, "_atomic_write_json", blocking_reservation_write)
|
||||
with concurrent.futures.ThreadPoolExecutor(max_workers=2) as pool:
|
||||
reserve_future = pool.submit(
|
||||
writer_handler.reserve_upload,
|
||||
upload_id,
|
||||
owner="alice",
|
||||
)
|
||||
assert write_entered.wait(5)
|
||||
cleanup_future = pool.submit(writer_handler.cleanup_old_uploads, set(), set())
|
||||
release_write.set()
|
||||
assert reserve_future.result(timeout=5) is not None
|
||||
assert cleanup_future.result(timeout=5) == 0
|
||||
assert writer_path.is_file()
|
||||
|
||||
# Cleanup wins: reservation cannot pass the same lock until the row and
|
||||
# bytes are gone, then fails so a caller cannot commit a dangling reference.
|
||||
cleanup_root = tmp_path / "cleanup-wins"
|
||||
cleanup_root.mkdir()
|
||||
cleanup_handler = _make_handler(cleanup_root)
|
||||
cleanup_path = _seed_old_uploads(cleanup_handler, [{
|
||||
"id": upload_id,
|
||||
"hash": "3" * 64,
|
||||
"mime": "text/plain",
|
||||
}])[upload_id]
|
||||
remove_entered = threading.Event()
|
||||
release_remove = threading.Event()
|
||||
real_remove = os.remove
|
||||
|
||||
def blocking_remove(candidate):
|
||||
if os.path.realpath(candidate) == os.path.realpath(cleanup_path):
|
||||
remove_entered.set()
|
||||
assert release_remove.wait(5)
|
||||
return real_remove(candidate)
|
||||
|
||||
monkeypatch.setattr("src.upload_handler.os.remove", blocking_remove)
|
||||
with concurrent.futures.ThreadPoolExecutor(max_workers=2) as pool:
|
||||
cleanup_future = pool.submit(cleanup_handler.cleanup_old_uploads, set(), set())
|
||||
assert remove_entered.wait(5)
|
||||
reserve_future = pool.submit(
|
||||
cleanup_handler.reserve_upload,
|
||||
upload_id,
|
||||
owner="alice",
|
||||
)
|
||||
release_remove.set()
|
||||
assert cleanup_future.result(timeout=5) == 1
|
||||
assert reserve_future.result(timeout=5) is None
|
||||
assert not cleanup_path.exists()
|
||||
|
||||
|
||||
def test_admin_cleanup_reference_discovery_failure_returns_503_without_deleting(
|
||||
tmp_path,
|
||||
monkeypatch,
|
||||
):
|
||||
from routes import upload_routes
|
||||
|
||||
handler = _make_handler(tmp_path)
|
||||
upload_id = "d" * 32 + ".png"
|
||||
path = _seed_old_uploads(handler, [
|
||||
{"id": upload_id, "hash": "4" * 64, "mime": "image/png"},
|
||||
])[upload_id]
|
||||
|
||||
def fail_reference_scan():
|
||||
raise RuntimeError("database unavailable")
|
||||
|
||||
monkeypatch.setattr(
|
||||
upload_routes,
|
||||
"_collect_persisted_upload_references",
|
||||
fail_reference_scan,
|
||||
)
|
||||
endpoint = _manual_cleanup_endpoint(handler, monkeypatch)
|
||||
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
asyncio.run(endpoint(_AdminRequest()))
|
||||
|
||||
assert exc.value.status_code == 503
|
||||
assert path.is_file()
|
||||
assert handler.get_upload_info(upload_id) is not None
|
||||
|
||||
|
||||
def test_cleanup_restores_index_when_file_removal_fails(tmp_path, monkeypatch):
|
||||
handler = _make_handler(tmp_path)
|
||||
upload_id = "e" * 32 + ".txt"
|
||||
path = _seed_old_uploads(handler, [
|
||||
{
|
||||
"id": upload_id,
|
||||
"hash": "5" * 64,
|
||||
"mime": "text/plain",
|
||||
},
|
||||
])[upload_id]
|
||||
|
||||
real_remove = os.remove
|
||||
|
||||
def fail_target_remove(candidate):
|
||||
if os.path.realpath(candidate) == os.path.realpath(path):
|
||||
raise PermissionError("file is in use")
|
||||
return real_remove(candidate)
|
||||
|
||||
monkeypatch.setattr("src.upload_handler.os.remove", fail_target_remove)
|
||||
|
||||
assert handler.cleanup_old_uploads(set(), set()) == 0
|
||||
assert path.is_file()
|
||||
assert handler.get_upload_info(upload_id) is not None
|
||||
assert any(
|
||||
info["id"] == upload_id
|
||||
for info in json.loads(
|
||||
Path(handler.upload_dir, "uploads.json").read_text(encoding="utf-8")
|
||||
).values()
|
||||
)
|
||||
assert any(
|
||||
info["id"] == upload_id
|
||||
for info in json.loads(
|
||||
Path(handler.upload_dir, "uploads.json.bak").read_text(encoding="utf-8")
|
||||
).values()
|
||||
)
|
||||
|
||||
|
||||
def test_admin_cleanup_with_corrupt_index_returns_503_and_fails_closed(
|
||||
tmp_path,
|
||||
monkeypatch,
|
||||
):
|
||||
from routes import upload_routes
|
||||
|
||||
handler = _make_handler(tmp_path)
|
||||
upload_id = "9" * 32 + ".png"
|
||||
path = _seed_old_uploads(handler, [
|
||||
{"id": upload_id, "hash": "9" * 64, "mime": "image/png"},
|
||||
])[upload_id]
|
||||
Path(handler.upload_dir, "uploads.json").write_text(
|
||||
'{"alice:broken": {',
|
||||
encoding="utf-8",
|
||||
)
|
||||
handler._index_cache = None
|
||||
|
||||
monkeypatch.setattr(
|
||||
upload_routes,
|
||||
"_collect_persisted_upload_references",
|
||||
lambda: (set(), set()),
|
||||
)
|
||||
endpoint = _manual_cleanup_endpoint(handler, monkeypatch)
|
||||
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
asyncio.run(endpoint(_AdminRequest()))
|
||||
|
||||
assert exc.value.status_code == 503
|
||||
assert path.is_file()
|
||||
|
||||
|
||||
def test_cleanup_with_missing_live_index_fails_closed(tmp_path):
|
||||
handler = _make_handler(tmp_path)
|
||||
upload_id = "8" * 32 + ".png"
|
||||
dated_dir = Path(handler.upload_dir) / "2000" / "01" / "01"
|
||||
dated_dir.mkdir(parents=True)
|
||||
path = dated_dir / upload_id
|
||||
path.write_bytes(b"unindexed bytes")
|
||||
|
||||
with pytest.raises(UploadCleanupSafetyError):
|
||||
handler.cleanup_old_uploads(set(), set())
|
||||
|
||||
assert path.is_file()
|
||||
|
||||
|
||||
def test_reference_discovery_covers_all_durable_upload_stores(
|
||||
monkeypatch,
|
||||
):
|
||||
from routes import upload_routes
|
||||
|
||||
document_id = "f" * 32 + ".pdf"
|
||||
version_id = "1" * 32 + ".pdf"
|
||||
note_upload_id = "3" * 32 + ".png"
|
||||
note_color_id = "2" * 32 + ".png"
|
||||
calendar_upload_id = "4" * 32 + ".png"
|
||||
event_upload_id = "5" * 32 + ".png"
|
||||
event_description_id = "7" * 32 + ".txt"
|
||||
event_location_id = "8" * 32 + ".png"
|
||||
gallery_hash = "6" * 64
|
||||
SessionLocal, engine, tmpfile = make_temp_sqlite(Base.metadata)
|
||||
db = SessionLocal()
|
||||
try:
|
||||
db.add(DbSession(
|
||||
id="session-2",
|
||||
name="Reference sources",
|
||||
endpoint_url="http://localhost",
|
||||
model="test-model",
|
||||
owner="alice",
|
||||
))
|
||||
db.add(Document(
|
||||
id="document-1",
|
||||
session_id="session-2",
|
||||
title="PDF",
|
||||
current_content=f'<!-- pdf_source upload_id="{document_id}" -->',
|
||||
owner="alice",
|
||||
))
|
||||
db.add(DocumentVersion(
|
||||
id="version-1",
|
||||
document_id="document-1",
|
||||
version_number=1,
|
||||
content=f'<!-- pdf_form_source upload_id="{version_id}" fields="1" -->',
|
||||
))
|
||||
db.add(GalleryImage(
|
||||
id="gallery-1",
|
||||
# Gallery filenames are normally generated 12-hex names, so this
|
||||
# record proves retention comes from its stored content hash.
|
||||
filename="abcdef123456.png",
|
||||
prompt="Chat upload",
|
||||
owner="alice",
|
||||
file_hash=gallery_hash,
|
||||
))
|
||||
db.add(Note(
|
||||
id="note-1",
|
||||
owner="alice",
|
||||
title="Photo note",
|
||||
image_url=f"/api/upload/{note_upload_id}",
|
||||
color=f"odysseus://attachment/{note_color_id}",
|
||||
))
|
||||
db.add(CalendarCal(
|
||||
id="calendar-1",
|
||||
owner="alice",
|
||||
name="Personal",
|
||||
color=f"/api/upload/{calendar_upload_id}",
|
||||
))
|
||||
db.add(CalendarEvent(
|
||||
uid="event-1",
|
||||
calendar_id="calendar-1",
|
||||
summary="Photo event",
|
||||
dtstart=datetime(2026, 7, 10, 12, 0),
|
||||
dtend=datetime(2026, 7, 10, 13, 0),
|
||||
color=f"/api/upload/{event_upload_id}",
|
||||
description=f"Notes: odysseus://attachment/{event_description_id}",
|
||||
location=f"/api/upload/{event_location_id}",
|
||||
))
|
||||
db.commit()
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
monkeypatch.setattr(upload_routes, "SessionLocal", SessionLocal)
|
||||
try:
|
||||
referenced_ids, referenced_hashes = (
|
||||
upload_routes._collect_persisted_upload_references()
|
||||
)
|
||||
finally:
|
||||
engine.dispose()
|
||||
tmpfile.close()
|
||||
try:
|
||||
os.unlink(tmpfile.name)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
assert {
|
||||
document_id,
|
||||
version_id,
|
||||
note_upload_id,
|
||||
note_color_id,
|
||||
calendar_upload_id,
|
||||
event_upload_id,
|
||||
event_description_id,
|
||||
event_location_id,
|
||||
} <= referenced_ids
|
||||
assert gallery_hash in referenced_hashes
|
||||
|
||||
|
||||
def test_write_reservation_extracts_only_explicit_internal_references():
|
||||
upload_id = "a" * 32 + ".png"
|
||||
checksum_like_text = "b" * 32
|
||||
|
||||
assert extract_internal_upload_ids(checksum_like_text) == set()
|
||||
assert extract_internal_upload_ids(f"sha={checksum_like_text}") == set()
|
||||
assert extract_internal_upload_ids({
|
||||
"image": f"/api/upload/{upload_id}",
|
||||
"nested": [f"odysseus://attachment/{upload_id}"],
|
||||
}) == {upload_id}
|
||||
assert extract_internal_upload_ids(
|
||||
f'<!-- pdf_source upload_id="{upload_id}" -->'
|
||||
) == {upload_id}
|
||||
assert extract_internal_upload_ids(
|
||||
f"[Attachment: photo.png | id={upload_id} | mime=image/png]"
|
||||
) == {upload_id}
|
||||
extensionless_id = "c" * 32
|
||||
assert extract_internal_upload_ids(
|
||||
f"See /api/upload/{extensionless_id}. Then continue."
|
||||
) == {extensionless_id}
|
||||
assert extract_internal_upload_ids(
|
||||
f"Attachment: odysseus://attachment/{extensionless_id}: ready"
|
||||
) == {extensionless_id}
|
||||
assert extract_internal_upload_ids(f"/api/upload/{upload_id}/extra") == set()
|
||||
|
||||
|
||||
def test_reservation_never_uses_admin_override(tmp_path):
|
||||
handler = _make_handler(tmp_path)
|
||||
upload_id = "c" * 32 + ".txt"
|
||||
_seed_old_uploads(handler, [{
|
||||
"id": upload_id,
|
||||
"hash": "c" * 64,
|
||||
"mime": "text/plain",
|
||||
"owner": "alice",
|
||||
}])
|
||||
|
||||
assert reserve_upload_references(
|
||||
handler,
|
||||
"alice",
|
||||
f"odysseus://attachment/{upload_id}",
|
||||
) is None
|
||||
assert reserve_upload_references(
|
||||
handler,
|
||||
"admin",
|
||||
f"odysseus://attachment/{upload_id}",
|
||||
) == upload_id
|
||||
assert handler.reserve_upload(
|
||||
upload_id,
|
||||
owner="admin",
|
||||
auth_manager=_AdminAuth(),
|
||||
allow_admin=False,
|
||||
) is None
|
||||
assert reserve_message_upload_references(
|
||||
handler,
|
||||
"admin",
|
||||
"legacy attachment metadata",
|
||||
{"attachments": [{"id": upload_id, "name": "owned.txt"}]},
|
||||
) == upload_id
|
||||
|
||||
|
||||
def test_remaining_durable_writers_reserve_before_commit(monkeypatch):
|
||||
import core.database as database
|
||||
import core.session_manager as session_manager_module
|
||||
import src.database as legacy_database
|
||||
from core.models import ChatMessage
|
||||
from core.session_manager import SessionManager
|
||||
from src import tool_utils
|
||||
from src.agent_tools.document_tools import EditDocumentTool
|
||||
from src.tools.calendar import do_manage_calendar
|
||||
from src.tools.notes import do_manage_notes
|
||||
|
||||
upload_id = "6" * 32 + ".png"
|
||||
|
||||
class RejectingHandler:
|
||||
def reserve_upload(self, _candidate, **_kwargs):
|
||||
return None
|
||||
|
||||
handler = RejectingHandler()
|
||||
monkeypatch.setattr(tool_utils, "_upload_handler", handler)
|
||||
|
||||
SessionLocal, engine, tmpfile = make_temp_sqlite(Base.metadata)
|
||||
monkeypatch.setattr(database, "SessionLocal", SessionLocal)
|
||||
monkeypatch.setattr(legacy_database, "SessionLocal", SessionLocal)
|
||||
monkeypatch.setattr(legacy_database, "Document", Document, raising=False)
|
||||
monkeypatch.setattr(
|
||||
legacy_database,
|
||||
"DocumentVersion",
|
||||
DocumentVersion,
|
||||
raising=False,
|
||||
)
|
||||
monkeypatch.setattr(session_manager_module, "SessionLocal", SessionLocal)
|
||||
db = SessionLocal()
|
||||
try:
|
||||
db.add(DbSession(
|
||||
id="writer-session",
|
||||
name="Writer coverage",
|
||||
endpoint_url="http://localhost",
|
||||
model="test-model",
|
||||
owner="alice",
|
||||
))
|
||||
db.add(Document(
|
||||
id="email-document",
|
||||
session_id="writer-session",
|
||||
title="New Email",
|
||||
language="email",
|
||||
current_content="To: team@example.test\nSubject: Status\n---\nOld body",
|
||||
version_count=1,
|
||||
owner="alice",
|
||||
))
|
||||
db.commit()
|
||||
|
||||
manager = SessionManager()
|
||||
manager.upload_handler = handler
|
||||
manager._persist_message(
|
||||
"writer-session",
|
||||
ChatMessage(
|
||||
"user",
|
||||
"attachment",
|
||||
metadata={"attachments": [{"id": upload_id}]},
|
||||
),
|
||||
)
|
||||
|
||||
document_result = asyncio.run(EditDocumentTool().execute(
|
||||
"<<<FIND>>>\n\n<<<REPLACE>>>\n"
|
||||
f"See /api/upload/{upload_id}\n<<<END>>>",
|
||||
{"doc_id": "email-document", "owner": "alice"},
|
||||
))
|
||||
assert document_result["exit_code"] == 1
|
||||
assert "no longer available" in document_result["error"]
|
||||
|
||||
calendar_result = asyncio.run(do_manage_calendar(
|
||||
json.dumps({
|
||||
"action": "create_event",
|
||||
"summary": "Attachment review",
|
||||
"dtstart": "2026-07-12T12:00:00",
|
||||
"description": f"See /api/upload/{upload_id}",
|
||||
}),
|
||||
owner="alice",
|
||||
))
|
||||
assert calendar_result["exit_code"] == 1
|
||||
assert "no longer available" in calendar_result["error"]
|
||||
|
||||
note_result = asyncio.run(do_manage_notes(
|
||||
json.dumps({
|
||||
"action": "add",
|
||||
"title": "Attachment note",
|
||||
"content": f"See /api/upload/{upload_id}",
|
||||
}),
|
||||
owner="alice",
|
||||
))
|
||||
assert note_result["exit_code"] == 1
|
||||
assert "no longer available" in note_result["error"]
|
||||
|
||||
verify = SessionLocal()
|
||||
try:
|
||||
assert verify.query(DbChatMessage).count() == 0
|
||||
stored_doc = verify.query(Document).filter(Document.id == "email-document").one()
|
||||
assert stored_doc.current_content.endswith("Old body")
|
||||
assert verify.query(CalendarEvent).count() == 0
|
||||
assert verify.query(Note).count() == 0
|
||||
finally:
|
||||
verify.close()
|
||||
finally:
|
||||
db.close()
|
||||
engine.dispose()
|
||||
tmpfile.close()
|
||||
try:
|
||||
os.unlink(tmpfile.name)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
|
||||
def test_note_calendar_and_document_routes_reserve_before_database_writes(monkeypatch):
|
||||
from routes.calendar_routes import EventCreate, setup_calendar_routes
|
||||
from routes import document_routes
|
||||
from routes.document_helpers import DocumentCreate
|
||||
from routes.note_routes import NoteCreate, setup_note_routes
|
||||
from src import auth_helpers
|
||||
|
||||
upload_id = "d" * 32 + ".png"
|
||||
|
||||
class RejectingHandler:
|
||||
def __init__(self):
|
||||
self.calls = []
|
||||
|
||||
def reserve_upload(self, candidate, **kwargs):
|
||||
self.calls.append((candidate, kwargs))
|
||||
return None
|
||||
|
||||
request = SimpleNamespace(
|
||||
state=SimpleNamespace(current_user="alice", api_token=False),
|
||||
app=SimpleNamespace(state=SimpleNamespace()),
|
||||
)
|
||||
|
||||
note_handler = RejectingHandler()
|
||||
note_router = setup_note_routes(upload_handler=note_handler)
|
||||
create_note = next(
|
||||
route.endpoint for route in note_router.routes
|
||||
if route.endpoint.__name__ == "create_note"
|
||||
)
|
||||
with pytest.raises(HTTPException) as note_error:
|
||||
create_note(
|
||||
request,
|
||||
NoteCreate(image_url=f"/api/upload/{upload_id}"),
|
||||
)
|
||||
assert note_error.value.status_code == 409
|
||||
assert note_handler.calls == [
|
||||
(upload_id, {"owner": "alice", "allow_admin": False})
|
||||
]
|
||||
|
||||
calendar_handler = RejectingHandler()
|
||||
calendar_router = setup_calendar_routes(upload_handler=calendar_handler)
|
||||
create_event = next(
|
||||
route.endpoint for route in calendar_router.routes
|
||||
if route.endpoint.__name__ == "create_event"
|
||||
)
|
||||
with pytest.raises(HTTPException) as calendar_error:
|
||||
asyncio.run(create_event(
|
||||
request,
|
||||
EventCreate(
|
||||
summary="Photo",
|
||||
dtstart="2026-07-10T12:00:00",
|
||||
color=f"odysseus://attachment/{upload_id}",
|
||||
),
|
||||
))
|
||||
assert calendar_error.value.status_code == 409
|
||||
assert calendar_handler.calls == [
|
||||
(upload_id, {"owner": "alice", "allow_admin": False})
|
||||
]
|
||||
|
||||
class EmptyDb:
|
||||
@staticmethod
|
||||
def close():
|
||||
return None
|
||||
|
||||
document_handler = RejectingHandler()
|
||||
monkeypatch.setattr(document_routes, "SessionLocal", EmptyDb)
|
||||
monkeypatch.setattr(
|
||||
auth_helpers,
|
||||
"require_privilege",
|
||||
lambda _request, _privilege: "alice",
|
||||
)
|
||||
document_router = document_routes.setup_document_routes(
|
||||
SimpleNamespace(),
|
||||
document_handler,
|
||||
)
|
||||
create_document = next(
|
||||
route.endpoint for route in document_router.routes
|
||||
if route.endpoint.__name__ == "create_document"
|
||||
)
|
||||
with pytest.raises(HTTPException) as document_error:
|
||||
asyncio.run(create_document(
|
||||
request,
|
||||
DocumentCreate(
|
||||
language="markdown",
|
||||
content=f"",
|
||||
),
|
||||
))
|
||||
assert document_error.value.status_code == 409
|
||||
assert document_handler.calls == [
|
||||
(upload_id, {"owner": "alice", "allow_admin": False})
|
||||
]
|
||||
Reference in New Issue
Block a user