import asyncio import base64 import json import logging from io import BytesIO import pytest from fastapi import HTTPException from starlette.datastructures import Headers, UploadFile from investibook import config as cfg from investibook.api import image_attachments as image_api from investibook.core import image_attachments as images from investibook.db import queries as dbq def _init_tmp_db(tmp_path, monkeypatch): db_path = tmp_path / "investibook.db" workspaces = tmp_path / "workspaces" workspaces.mkdir() monkeypatch.setattr(dbq, "DB_PATH", db_path) monkeypatch.setattr(cfg, "DB_PATH", db_path) monkeypatch.setattr(cfg, "WORKSPACES_DIR", workspaces) dbq._db_init() return workspaces def test_image_attachment_db_and_message_serialization(tmp_path, monkeypatch): _init_tmp_db(tmp_path, monkeypatch) item = dbq._db_create_image_attachment( attachment_id="img_1", workspace_id="ws-1", user_id="u-1", filename="cpu.png", mime_type="image/png", rel_path="user_uploads/images/img_1/original.png", size_bytes=12, sha256="abc", ) assert item["status"] == "uploaded" updated = dbq._db_update_image_attachment( workspace_id="ws-1", attachment_id="img_1", status="ready", kind="metric_chart", analysis_path="user_uploads/images/img_1/analysis.json", summary_path="user_uploads/images/img_1/summary_for_agent.txt", ) assert updated["kind"] == "metric_chart" dbq._db_append_messages_v2( "ws-1", [ { "id": "user-1", "role": "user", "content": "look at this", "attachments": [images.image_attachment_message_view(updated)], } ], ) messages = dbq._db_get_messages_v2("ws-1") assert messages[0]["attachments"][0]["id"] == "img_1" assert messages[0]["attachments"][0]["kind"] == "metric_chart" def test_upload_in_vision_mode_returns_before_parse_completes(tmp_path, monkeypatch): """Vision-parse mode returns 'uploaded' synchronously and runs parse in the background — the upload response must not block on the LLM call.""" workspaces = _init_tmp_db(tmp_path, monkeypatch) monkeypatch.setattr(cfg, "IMAGE_PARSE_PROVIDER", "bedrock") monkeypatch.setattr(cfg, "IMAGE_PARSE_BEDROCK_MODEL_ID", "anthropic.test-model") async def fake_require_workspace_editor(_workspace_id, _current_user): return {"id": _workspace_id} monkeypatch.setattr(image_api, "require_workspace_editor", fake_require_workspace_editor) # Stub vision model so the background parse task (if it runs) doesn't # actually hit Bedrock. We don't assert it's called — asyncio.run cancels # pending tasks on exit, so the background parse may or may not get a tick. def stub_vision(*_args, **_kwargs): return ({"kind": "metric_chart", "confidence": "medium"}, {"inputTokens": 0, "outputTokens": 0}) monkeypatch.setattr(images, "_call_vision_model", stub_vision) upload = UploadFile( BytesIO(b"png-bytes"), filename="cpu.png", headers=Headers({"content-type": "image/png"}), ) response = asyncio.run( image_api.upload_image_attachment( "ws-upload", file=upload, current_user={"id": "u-1"}, ) ) # Upload response itself reflects pre-parse state; status flip to ready # happens asynchronously and the frontend picks it up via polling. attachment = response["attachment"] assert attachment["status"] == "uploaded" assert not attachment["analysis_path"] assert not attachment["summary_path"] upload_dir = workspaces / "ws-upload" / images.image_attachment_rel_dir(attachment["id"]) assert (upload_dir / "original.png").exists() def test_upload_in_passthrough_mode_resolves_ready_inline(tmp_path, monkeypatch): """Passthrough mode flips status to ready during the upload call itself, so the frontend's ready-gate doesn't block sends.""" workspaces = _init_tmp_db(tmp_path, monkeypatch) monkeypatch.setattr(cfg, "IMAGE_PARSE_PROVIDER", "none") images._passthrough_fallback_warned.clear() async def fake_require_workspace_editor(_workspace_id, _current_user): return {"id": _workspace_id} def fail_vision(*_args, **_kwargs): raise AssertionError("passthrough must not call vision") monkeypatch.setattr(image_api, "require_workspace_editor", fake_require_workspace_editor) monkeypatch.setattr(images, "_call_vision_model", fail_vision) upload = UploadFile( BytesIO(b"png-bytes"), filename="cpu.png", headers=Headers({"content-type": "image/png"}), ) response = asyncio.run( image_api.upload_image_attachment( "ws-pt-upload", file=upload, current_user={"id": "u-1"}, ) ) attachment = response["attachment"] assert attachment["status"] == "ready" assert attachment["analysis_path"] == "" assert attachment["summary_path"] == "" upload_dir = workspaces / "ws-pt-upload" / images.image_attachment_rel_dir(attachment["id"]) assert (upload_dir / "original.png").exists() assert not (upload_dir / "analysis.json").exists() assert not (upload_dir / "summary_for_agent.txt").exists() def test_build_image_summary_prompt_block_parses_uploaded_on_send(tmp_path, monkeypatch): workspaces = _init_tmp_db(tmp_path, monkeypatch) monkeypatch.setattr(cfg, "IMAGE_PARSE_PROVIDER", "bedrock") monkeypatch.setattr(cfg, "IMAGE_PARSE_BEDROCK_MODEL_ID", "anthropic.test-model") image_rel = "user_uploads/images/img_2/original.png" image_path = workspaces / "ws-2" / image_rel image_path.parent.mkdir(parents=True) image_path.write_bytes(b"not-a-real-png") dbq._db_create_image_attachment( attachment_id="img_2", workspace_id="ws-2", user_id="u-1", filename="cpu.png", mime_type="image/png", rel_path=image_rel, size_bytes=12, sha256="abc", ) def fake_vision(_image_bytes, _mime_type, _rel_path): return ( { "kind": "metric_chart", "confidence": "high", "chart_summary": "CPU rises", }, {"inputTokens": 1, "outputTokens": 2}, ) monkeypatch.setattr(images, "_call_vision_model", fake_vision) prompt, message_items = asyncio.run( images.build_image_summary_prompt_block("ws-2", ["img_2"]) ) item = dbq._db_get_image_attachment("ws-2", "img_2") assert item["status"] == "ready" assert item["summary_path"] == "user_uploads/images/img_2/summary_for_agent.txt" assert "[image_attachment:img_2]" in prompt assert "CPU rises" in prompt assert message_items[0]["summary_path"] == item["summary_path"] def test_parse_image_attachment_writes_analysis_and_summary(tmp_path, monkeypatch): workspaces = _init_tmp_db(tmp_path, monkeypatch) monkeypatch.setattr(cfg, "IMAGE_PARSE_PROVIDER", "bedrock") monkeypatch.setattr(cfg, "IMAGE_PARSE_BEDROCK_MODEL_ID", "anthropic.test-model") image_rel = "user_uploads/images/img_3/original.png" image_path = workspaces / "ws-3" / image_rel image_path.parent.mkdir(parents=True) image_path.write_bytes(b"not-a-real-png") dbq._db_create_image_attachment( attachment_id="img_3", workspace_id="ws-3", user_id="u-1", filename="cpu.png", mime_type="image/png", rel_path=image_rel, size_bytes=14, sha256="abc", ) def fake_vision(_image_bytes, _mime_type, _rel_path): return ( { "kind": "metric_chart", "confidence": "medium", "ocr_summary": "TiDB CPU panel", "chart_summary": "CPU rises near the incident window", "trends": [{"summary": "upward trend"}], }, {"inputTokens": 10, "outputTokens": 20}, ) monkeypatch.setattr(images, "_call_vision_model", fake_vision) asyncio.run(images.parse_image_attachment("ws-3", "img_3")) item = dbq._db_get_image_attachment("ws-3", "img_3") assert item["status"] == "ready" assert item["kind"] == "metric_chart" analysis = json.loads((workspaces / "ws-3" / item["analysis_path"]).read_text()) summary = (workspaces / "ws-3" / item["summary_path"]).read_text() assert analysis["ocr_summary"] == "TiDB CPU panel" assert "CPU rises" in summary usage = dbq._db_get_workspace_token_usage("ws-3") assert len(usage) == 1 assert usage[0]["channel"] == "image-send" assert usage[0]["session_id"] == "image-parse:img_3" assert usage[0]["input_tokens"] == 10 assert usage[0]["output_tokens"] == 20 assert usage[0]["cost_currency"] == "USD" assert usage[0]["cost_amount"] == pytest.approx(0.00033) # --------------------------------------------------------------------------- # Passthrough mode tests # --------------------------------------------------------------------------- def _create_uploaded_attachment(workspaces, ws_id, attachment_id, *, payload=b"bytes"): rel = f"user_uploads/images/{attachment_id}/original.png" abs_path = workspaces / ws_id / rel abs_path.parent.mkdir(parents=True, exist_ok=True) abs_path.write_bytes(payload) dbq._db_create_image_attachment( attachment_id=attachment_id, workspace_id=ws_id, user_id="u-1", filename="cpu.png", mime_type="image/png", rel_path=rel, size_bytes=len(payload), sha256="abc", ) return abs_path def test_passthrough_short_circuits_parse_image_attachment(tmp_path, monkeypatch): workspaces = _init_tmp_db(tmp_path, monkeypatch) monkeypatch.setattr(cfg, "IMAGE_PARSE_PROVIDER", "none") images._passthrough_fallback_warned.clear() def fail_vision(*_args, **_kwargs): raise AssertionError("passthrough must not call vision") monkeypatch.setattr(images, "_call_vision_model", fail_vision) _create_uploaded_attachment(workspaces, "ws-pt", "img_pt") asyncio.run(images.parse_image_attachment("ws-pt", "img_pt")) item = dbq._db_get_image_attachment("ws-pt", "img_pt") assert item["status"] == "ready" assert item["analysis_path"] == "" assert item["summary_path"] == "" rel_dir = workspaces / "ws-pt" / images.image_attachment_rel_dir("img_pt") assert not (rel_dir / "analysis.json").exists() assert not (rel_dir / "summary_for_agent.txt").exists() assert dbq._db_get_workspace_token_usage("ws-pt") == [] def test_passthrough_build_prompt_returns_empty_text(tmp_path, monkeypatch): workspaces = _init_tmp_db(tmp_path, monkeypatch) monkeypatch.setattr(cfg, "IMAGE_PARSE_PROVIDER", "none") images._passthrough_fallback_warned.clear() _create_uploaded_attachment(workspaces, "ws-pt2", "img_pt2") prompt, message_items = asyncio.run( images.build_image_summary_prompt_block("ws-pt2", ["img_pt2"]) ) assert prompt == "" assert len(message_items) == 1 assert message_items[0]["id"] == "img_pt2" assert message_items[0]["status"] == "ready" def test_passthrough_load_image_blocks_returns_acp_image_block(tmp_path, monkeypatch): workspaces = _init_tmp_db(tmp_path, monkeypatch) monkeypatch.setattr(cfg, "IMAGE_PARSE_PROVIDER", "none") images._passthrough_fallback_warned.clear() payload = b"\x89PNG\r\n\x1a\nfake-bytes" _create_uploaded_attachment(workspaces, "ws-pt3", "img_pt3", payload=payload) # Promote to ready via passthrough parse. asyncio.run(images.parse_image_attachment("ws-pt3", "img_pt3")) blocks = asyncio.run( images.load_image_blocks_for_prompt("ws-pt3", ["img_pt3"]) ) assert len(blocks) == 1 block = blocks[0] assert block.mime_type == "image/png" assert base64.b64decode(block.data) == payload def test_bedrock_fallback_when_model_id_unset(tmp_path, monkeypatch, caplog): workspaces = _init_tmp_db(tmp_path, monkeypatch) monkeypatch.setattr(cfg, "IMAGE_PARSE_PROVIDER", "bedrock") monkeypatch.setattr(cfg, "IMAGE_PARSE_BEDROCK_MODEL_ID", "") images._passthrough_fallback_warned.clear() def fail_vision(*_args, **_kwargs): raise AssertionError("fallback path must not call vision") monkeypatch.setattr(images, "_call_vision_model", fail_vision) _create_uploaded_attachment(workspaces, "ws-pt4", "img_pt4") with caplog.at_level(logging.WARNING, logger="uvicorn.error"): asyncio.run(images.parse_image_attachment("ws-pt4", "img_pt4")) # Second call must not duplicate the warning. _create_uploaded_attachment(workspaces, "ws-pt4", "img_pt4b") asyncio.run(images.parse_image_attachment("ws-pt4", "img_pt4b")) item = dbq._db_get_image_attachment("ws-pt4", "img_pt4") assert item["status"] == "ready" assert item["summary_path"] == "" fallback_warnings = [ r for r in caplog.records if r.levelno == logging.WARNING and "falling back to passthrough" in r.getMessage() ] assert len(fallback_warnings) == 1 def test_build_prompt_waits_for_in_flight_processing(tmp_path, monkeypatch): """If a background parse is already in flight (status='processing' when chat_stream arrives), build_image_summary_prompt_block must poll until it terminates rather than raising a 'not ready' error.""" workspaces = _init_tmp_db(tmp_path, monkeypatch) monkeypatch.setattr(cfg, "IMAGE_PARSE_PROVIDER", "bedrock") monkeypatch.setattr(cfg, "IMAGE_PARSE_BEDROCK_MODEL_ID", "anthropic.test-model") image_rel = "user_uploads/images/img_proc/original.png" image_path = workspaces / "ws-proc" / image_rel image_path.parent.mkdir(parents=True) image_path.write_bytes(b"not-a-real-png") dbq._db_create_image_attachment( attachment_id="img_proc", workspace_id="ws-proc", user_id="u-1", filename="cpu.png", mime_type="image/png", rel_path=image_rel, size_bytes=14, sha256="abc", ) # Simulate an in-flight background parse: status already flipped to # 'processing' but the LLM hasn't returned yet. dbq._db_update_image_attachment( workspace_id="ws-proc", attachment_id="img_proc", status="processing", ) # After a couple of polls, complete the parse from a side task. async def background_finish(): await asyncio.sleep(0.2) dbq._db_update_image_attachment( workspace_id="ws-proc", attachment_id="img_proc", status="ready", kind="metric_chart", analysis_path="user_uploads/images/img_proc/analysis.json", summary_path="user_uploads/images/img_proc/summary_for_agent.txt", ) # Provide the summary file the prompt builder reads. summary = workspaces / "ws-proc" / "user_uploads/images/img_proc/summary_for_agent.txt" summary.parent.mkdir(parents=True, exist_ok=True) summary.write_text("# User Image Summary\nfake summary\n") async def driver(): finish_task = asyncio.create_task(background_finish()) prompt, items = await images.build_image_summary_prompt_block( "ws-proc", ["img_proc"] ) await finish_task return prompt, items prompt, items = asyncio.run(driver()) assert "[image_attachment:img_proc]" in prompt assert items[0]["status"] == "ready" # --------------------------------------------------------------------------- # Generic workspace attachment contract # --------------------------------------------------------------------------- def test_canonical_and_legacy_attachment_routes_share_handlers(): route_handlers = { (route.path, method): route.endpoint for route in image_api.router.routes for method in route.methods } assert route_handlers[ ("/api/workspace/{workspace_id}/attachments", "POST") ] is route_handlers[ ("/api/workspace/{workspace_id}/image-attachments", "POST") ] assert route_handlers[ ("/api/workspace/{workspace_id}/attachments", "GET") ] is route_handlers[ ("/api/workspace/{workspace_id}/image-attachments", "GET") ] assert route_handlers[ ("/api/workspace/{workspace_id}/attachments/{attachment_id}", "GET") ] is route_handlers[ ("/api/workspace/{workspace_id}/image-attachments/{attachment_id}", "GET") ] def test_parse_false_image_is_ready_without_vision_parse(tmp_path, monkeypatch): workspaces = _init_tmp_db(tmp_path, monkeypatch) monkeypatch.setattr(cfg, "IMAGE_PARSE_PROVIDER", "bedrock") monkeypatch.setattr(cfg, "IMAGE_PARSE_BEDROCK_MODEL_ID", "anthropic.test-model") async def fake_require_workspace_editor(_workspace_id, _current_user): return {"id": _workspace_id} def fail_vision(*_args, **_kwargs): raise AssertionError("parse=false must not call vision") monkeypatch.setattr(image_api, "require_workspace_editor", fake_require_workspace_editor) monkeypatch.setattr(images, "_call_vision_model", fail_vision) payload = b"\x89PNG\r\n\x1a\nraw-image" upload = UploadFile( BytesIO(payload), filename="cpu.png", headers=Headers({"content-type": "image/png"}), ) response = asyncio.run( image_api.upload_image_attachment( "ws-image-no-parse", file=upload, parse=False, current_user={"id": "u-1"}, ) ) attachment = response["attachment"] assert attachment["status"] == "ready" assert attachment["kind"] == "image" assert attachment["summary_path"] == "" stored = workspaces / "ws-image-no-parse" / attachment["rel_path"] assert stored.read_bytes() == payload blocks = asyncio.run( images.load_image_blocks_for_prompt( "ws-image-no-parse", [attachment["id"]] ) ) assert base64.b64decode(blocks[0].data) == payload @pytest.mark.parametrize( ("filename", "mime_type", "payload"), [ ("notes.txt", "text/plain", b"plain text"), ("diagnostic.pdf", "application/pdf", b"%PDF-raw"), ( "report.docx", "application/vnd.openxmlformats-officedocument.wordprocessingml.document", b"PK-docx", ), ("bundle.zip", "application/zip", b"PK-archive"), ], ) def test_generic_file_upload_is_ready_and_byte_exact( tmp_path, monkeypatch, filename, mime_type, payload ): workspaces = _init_tmp_db(tmp_path, monkeypatch) monkeypatch.setattr(cfg, "IMAGE_PARSE_PROVIDER", "bedrock") monkeypatch.setattr(cfg, "IMAGE_PARSE_BEDROCK_MODEL_ID", "anthropic.test-model") async def fake_require_workspace_editor(_workspace_id, _current_user): return {"id": _workspace_id} def fail_vision(*_args, **_kwargs): raise AssertionError("non-image attachments must not call vision") monkeypatch.setattr(image_api, "require_workspace_editor", fake_require_workspace_editor) monkeypatch.setattr(images, "_call_vision_model", fail_vision) upload = UploadFile( BytesIO(payload), filename=f"../customer {filename}", headers=Headers({"content-type": mime_type}), ) response = asyncio.run( image_api.upload_image_attachment( "ws-generic", file=upload, parse=True, current_user={"id": "u-1"}, ) ) attachment = response["attachment"] assert attachment["status"] == "ready" assert attachment["kind"] == "file" assert attachment["filename"] == f"customer_{filename}" assert attachment["mime_type"] == mime_type assert attachment["analysis_path"] == "" assert attachment["summary_path"] == "" stored = workspaces / "ws-generic" / attachment["rel_path"] assert stored.parent.name == attachment["id"] assert stored.read_bytes() == payload # Defensive parse calls also keep generic files out of the image parser. asyncio.run(images.parse_image_attachment("ws-generic", attachment["id"])) assert dbq._db_get_image_attachment("ws-generic", attachment["id"])["kind"] == "file" def test_generic_file_chat_references_workspace_path_without_inlining(tmp_path, monkeypatch): workspaces = _init_tmp_db(tmp_path, monkeypatch) monkeypatch.setattr(cfg, "IMAGE_PARSE_PROVIDER", "bedrock") monkeypatch.setattr(cfg, "IMAGE_PARSE_BEDROCK_MODEL_ID", "anthropic.test-model") async def fake_require_workspace_editor(_workspace_id, _current_user): return {"id": _workspace_id} def fail_vision(*_args, **_kwargs): raise AssertionError("generic file chat must not call vision") monkeypatch.setattr(image_api, "require_workspace_editor", fake_require_workspace_editor) monkeypatch.setattr(images, "_call_vision_model", fail_vision) secret_payload = b"RAW-CONTENT-MUST-NOT-BE-INLINED" response = asyncio.run( image_api.upload_image_attachment( "ws-generic-chat", file=UploadFile( BytesIO(secret_payload), filename="diagnostics.tar.gz", headers=Headers({"content-type": "application/gzip"}), ), parse=False, current_user={"id": "u-1"}, ) ) attachment = response["attachment"] prompt, message_items = asyncio.run( images.build_image_summary_prompt_block( "ws-generic-chat", [attachment["id"]] ) ) blocks = asyncio.run( images.load_image_blocks_for_prompt( "ws-generic-chat", [attachment["id"]] ) ) assert f"[workspace_attachment:{attachment['id']}]" in prompt assert f"filename: {attachment['filename']}" in prompt assert f"mime_type: {attachment['mime_type']}" in prompt assert f"workspace_relative_path: {attachment['rel_path']}" in prompt assert secret_payload.decode() not in prompt assert message_items[0]["rel_path"] == attachment["rel_path"] assert blocks == [] assert ( workspaces / "ws-generic-chat" / attachment["rel_path"] ).read_bytes() == secret_payload def test_empty_attachment_is_rejected(tmp_path, monkeypatch): _init_tmp_db(tmp_path, monkeypatch) async def fake_require_workspace_editor(_workspace_id, _current_user): return {"id": _workspace_id} monkeypatch.setattr(image_api, "require_workspace_editor", fake_require_workspace_editor) with pytest.raises(HTTPException) as exc_info: asyncio.run( image_api.upload_image_attachment( "ws-empty", file=UploadFile( BytesIO(b""), filename="empty.bin", headers=Headers({"content-type": "application/octet-stream"}), ), parse=False, current_user={"id": "u-1"}, ) ) assert exc_info.value.status_code == 400 assert exc_info.value.detail == "Attachment file is empty" def test_generic_attachment_respects_configured_size_limit(tmp_path, monkeypatch): workspaces = _init_tmp_db(tmp_path, monkeypatch) monkeypatch.setattr(cfg, "IMAGE_ATTACHMENT_MAX_FILE_MB", 1) async def fake_require_workspace_editor(_workspace_id, _current_user): return {"id": _workspace_id} monkeypatch.setattr(image_api, "require_workspace_editor", fake_require_workspace_editor) with pytest.raises(HTTPException) as exc_info: asyncio.run( image_api.upload_image_attachment( "ws-too-large", file=UploadFile( BytesIO(b"x" * (1024 * 1024 + 1)), filename="large.bin", headers=Headers({"content-type": "application/octet-stream"}), ), parse=False, current_user={"id": "u-1"}, ) ) assert exc_info.value.status_code == 400 assert exc_info.value.detail == "Attachment exceeds maximum size (1 MB)" assert not (workspaces / "ws-too-large").exists()