diff --git a/backend/application/chat/utilities/file_utils.py b/backend/application/chat/utilities/file_utils.py index 3c6a889..001455e 100644 --- a/backend/application/chat/utilities/file_utils.py +++ b/backend/application/chat/utilities/file_utils.py @@ -455,5 +455,6 @@ def build_files_manifest(session_context: Dict[str, Any]) -> Optional[Dict[str, f"{file_list}\n\n" "(You can ask to open or analyze any of these by name. " "Large contents are not fully in this prompt unless user or tools provided excerpts.)" + "The user may refer to these files in their requests as session files or attachments." ) } \ No newline at end of file diff --git a/backend/tests/test_attach_file_flow.py b/backend/tests/test_attach_file_flow.py new file mode 100644 index 0000000..b1bd23a --- /dev/null +++ b/backend/tests/test_attach_file_flow.py @@ -0,0 +1,209 @@ +import base64 +import uuid +import pytest + +from application.chat.service import ChatService +from modules.file_storage.manager import FileManager +from modules.file_storage.mock_s3_client import MockS3StorageClient + + +class FakeLLM: + async def call_plain(self, model_name, messages, temperature=0.7): + return "ok" + + async def call_with_tools(self, model_name, messages, tools_schema, tool_choice="auto", temperature=0.7): + from interfaces.llm import LLMResponse + return LLMResponse(content="ok", tool_calls=None, model_used=model_name) + + async def call_with_rag(self, model_name, messages, data_sources, user_email, temperature=0.7): + return "ok" + + async def call_with_rag_and_tools(self, model_name, messages, data_sources, tools_schema, user_email, tool_choice="auto", temperature=0.7): + from interfaces.llm import LLMResponse + return LLMResponse(content="ok", tool_calls=None, model_used=model_name) + + +@pytest.fixture +def file_manager(): + # Use in-process mock S3 for deterministic tests + return FileManager(s3_client=MockS3StorageClient()) + + +@pytest.fixture +def chat_service(file_manager): + # Minimal ChatService wiring for file/session operations + return ChatService(llm=FakeLLM(), file_manager=file_manager) + + +@pytest.mark.asyncio +async def test_handle_attach_file_success_creates_session_and_emits_update(chat_service, file_manager): + user_email = "user1@example.com" + session_id = uuid.uuid4() + + # Seed a file into the mock storage for this user + filename = "report.txt" + content_b64 = base64.b64encode(b"hello world").decode() + upload_meta = await file_manager.s3_client.upload_file( + user_email=user_email, + filename=filename, + content_base64=content_b64, + content_type="text/plain", + tags={"source": "user"}, + source_type="user", + ) + s3_key = upload_meta["key"] + + updates = [] + + async def capture_update(msg): + updates.append(msg) + + # Act: attach the file to a brand new session (auto-creates session) + resp = await chat_service.handle_attach_file( + session_id=session_id, + s3_key=s3_key, + user_email=user_email, + update_callback=capture_update, + ) + + # Assert: success response and files_update emitted + assert resp.get("type") == "file_attach" + assert resp.get("success") is True + assert resp.get("filename") == filename + + assert any( + u.get("type") == "intermediate_update" and u.get("update_type") == "files_update" + for u in updates + ), "Expected a files_update intermediate update to be emitted" + + # Session context should include the file by filename + session = chat_service.sessions.get(session_id) + assert session is not None + assert filename in session.context.get("files", {}) + assert session.context["files"][filename]["key"] == s3_key + + +@pytest.mark.asyncio +async def test_handle_attach_file_not_found_returns_error(chat_service): + user_email = "user1@example.com" + session_id = uuid.uuid4() + + # Non-existent S3 key for the same user + bad_key = f"users/{user_email}/uploads/does_not_exist_12345.txt" + resp = await chat_service.handle_attach_file( + session_id=session_id, + s3_key=bad_key, + user_email=user_email, + update_callback=None, + ) + + assert resp.get("type") == "file_attach" + assert resp.get("success") is False + assert "File not found" in resp.get("error", "") + + +@pytest.mark.asyncio +async def test_handle_attach_file_unauthorized_other_user_key(chat_service, file_manager): + # Upload under user1 + owner_email = "owner@example.com" + other_email = "other@example.com" + session_id = uuid.uuid4() + + filename = "secret.pdf" + content_b64 = base64.b64encode(b"top-secret").decode() + upload_meta = await file_manager.s3_client.upload_file( + user_email=owner_email, + filename=filename, + content_base64=content_b64, + content_type="application/pdf", + tags={"source": "user"}, + source_type="user", + ) + s3_key = upload_meta["key"] + + # Attempt to attach with a different user should fail + resp = await chat_service.handle_attach_file( + session_id=session_id, + s3_key=s3_key, + user_email=other_email, + update_callback=None, + ) + + assert resp.get("type") == "file_attach" + assert resp.get("success") is False + assert "Access denied" in resp.get("error", "") + + +@pytest.mark.asyncio +async def test_handle_reset_session_reinitializes(chat_service): + user_email = "user1@example.com" + session_id = uuid.uuid4() + + # Create a session first + await chat_service.create_session(session_id, user_email) + assert chat_service.sessions.get(session_id) is not None + + # Reset the session + resp = await chat_service.handle_reset_session(session_id=session_id, user_email=user_email) + + assert resp.get("type") == "session_reset" + # After reset, a fresh active session should exist for the same id + new_session = chat_service.sessions.get(session_id) + assert new_session is not None + assert new_session.active is True + + +@pytest.mark.asyncio +async def test_handle_download_file_success_after_attach(chat_service, file_manager): + user_email = "user1@example.com" + session_id = uuid.uuid4() + + # Upload and then attach to session + filename = "notes.md" + content_bytes = b"### Title\nSome content." + content_b64 = base64.b64encode(content_bytes).decode() + upload_meta = await file_manager.s3_client.upload_file( + user_email=user_email, + filename=filename, + content_base64=content_b64, + content_type="text/markdown", + tags={"source": "user"}, + source_type="user", + ) + s3_key = upload_meta["key"] + + await chat_service.handle_attach_file( + session_id=session_id, + s3_key=s3_key, + user_email=user_email, + update_callback=None, + ) + + # Act: download by filename (from session context) + resp = await chat_service.handle_download_file( + session_id=session_id, + filename=filename, + user_email=user_email, + ) + + assert resp.get("type") is not None + # content_base64 should match uploaded content + returned_b64 = resp.get("content_base64") + assert isinstance(returned_b64, str) and len(returned_b64) > 0 + assert base64.b64decode(returned_b64) == content_bytes + + +@pytest.mark.asyncio +async def test_handle_download_file_not_in_session_returns_error(chat_service): + user_email = "user1@example.com" + session_id = uuid.uuid4() + filename = "missing.txt" + + # No attach performed; should error that file isn't in session + resp = await chat_service.handle_download_file( + session_id=session_id, + filename=filename, + user_email=user_email, + ) + + assert resp.get("error") == "Session or file manager not available" or resp.get("error") == "File not found in session" diff --git a/frontend/src/components/AllFilesView.jsx b/frontend/src/components/AllFilesView.jsx index da0df53..59c4fc1 100644 --- a/frontend/src/components/AllFilesView.jsx +++ b/frontend/src/components/AllFilesView.jsx @@ -17,7 +17,7 @@ import { useChat } from '../contexts/ChatContext' import { useWS } from '../contexts/WSContext' const AllFilesView = () => { - const { token, user: userEmail } = useChat() + const { token, user: userEmail, ensureSession, addSystemEvent, addPendingFileEvent, attachments } = useChat() const { sendMessage } = useWS() const [allFiles, setAllFiles] = useState([]) const [filteredFiles, setFilteredFiles] = useState([]) @@ -211,17 +211,36 @@ const AllFilesView = () => { } } - const handleLoadToSession = async (file) => { + const handleAddToSession = async (file) => { try { + // Check if file is already attached + if (attachments.has(file.key)) { + addSystemEvent('file-attached', `'${file.filename}' is already in this session.`) + return + } + + // Ensure session exists + await ensureSession() + + // Add "attaching" system event and track it as pending + const eventId = addSystemEvent('file-attaching', `Adding '${file.filename}' to this session...`, { + fileId: file.key, + fileName: file.filename, + source: 'library' + }) + + // Track this as a pending file event + addPendingFileEvent(file.key, eventId) + + // Send attach_file message (WebSocket handler will resolve the pending event) sendMessage({ type: 'attach_file', s3_key: file.key, user: userEmail }) - showNotification(`File "${file.filename}" loaded to current session`, 'success') } catch (error) { - console.error('Error loading file to session:', error) - showNotification('Failed to load file to session', 'error') + console.error('Error adding file to session:', error) + addSystemEvent('file-attach-error', `Failed to add '${file.filename}' to session: ${error.message}`) } } @@ -380,9 +399,9 @@ const AllFilesView = () => { {/* Action Buttons */}
diff --git a/frontend/src/components/Message.jsx b/frontend/src/components/Message.jsx index 22cfdcf..23044a2 100644 --- a/frontend/src/components/Message.jsx +++ b/frontend/src/components/Message.jsx @@ -50,7 +50,7 @@ hljs.registerLanguage('sh', bash) const processFileReferences = (content) => { return content.replace( /@file\s+([^\s]+)/g, - '📎 @file $1' + '@file $1' ) } @@ -629,24 +629,24 @@ const renderContent = () => {
)} - {/* Result Section */} - {message.result && ( -
-
- {!toolOutputCollapsed && ( @@ -768,6 +768,31 @@ const renderContent = () => { } if (isUser || isSystem) { + // Handle file attachment system events + if (message.type === 'system' && message.subtype) { + switch (message.subtype) { + case 'file-attaching': + return ( +
+ {message.text} +
+ ) + case 'file-attached': + return ( +
+ {message.text} +
+ ) + case 'file-attach-error': + return ( +
+ {message.text} +
+ ) + default: + return
{message.content}
+ } + } return
{message.content}
} diff --git a/frontend/src/contexts/ChatContext.jsx b/frontend/src/contexts/ChatContext.jsx index 18f2b1b..8cf172b 100644 --- a/frontend/src/contexts/ChatContext.jsx +++ b/frontend/src/contexts/ChatContext.jsx @@ -9,6 +9,13 @@ import { useFiles } from '../hooks/chat/useFiles' import { useSettings } from '../hooks/useSettings' import { createWebSocketHandler } from '../handlers/chat/websocketHandlers' +// Generate cryptographically secure random string +const generateSecureRandomString = (length = 9) => { + const array = new Uint8Array(length) + crypto.getRandomValues(array) + return Array.from(array, byte => byte.toString(36)).join('').slice(0, length) +} + const ChatContext = createContext(null) export const useChat = () => { @@ -29,6 +36,38 @@ export const ChatProvider = ({ children }) => { const [isWelcomeVisible, setIsWelcomeVisible] = useState(true) const [isThinking, setIsThinking] = useState(false) + const [sessionId, setSessionId] = useState(null) + const [attachments, setAttachments] = useState(new Set()) + const [pendingFileEvents, setPendingFileEvents] = useState(new Map()) + + // Method to add a file to attachments + const addAttachment = useCallback((fileId) => { + setAttachments(prev => new Set([...prev, fileId])) + }, []) + + // Methods to manage pending file events + const addPendingFileEvent = useCallback((fileKey, eventId) => { + setPendingFileEvents(prev => new Map(prev.set(fileKey, eventId))) + }, []) + + const resolvePendingFileEvent = useCallback((fileKey, newSubtype, newText) => { + setPendingFileEvents(prev => { + const eventId = prev.get(fileKey) + if (eventId) { + // Update the message in-place + mapMessages(messages => messages.map(msg => + msg.id === eventId + ? { ...msg, subtype: newSubtype, text: newText } + : msg + )) + // Remove from pending + const next = new Map(prev) + next.delete(fileKey) + return next + } + return prev + }) + }, [mapMessages]) const { sendMessage, addMessageHandler } = useWS() const { currentModel } = config @@ -59,10 +98,13 @@ export const ChatProvider = ({ children }) => { setCustomUIContent: files.setCustomUIContent, setSessionFiles: files.setSessionFiles, getFileType: files.getFileType, - triggerFileDownload + triggerFileDownload, + addAttachment, + addPendingFileEvent, + resolvePendingFileEvent }) return addMessageHandler(handler) - }, [addMessageHandler, addMessage, mapMessages, agent.setCurrentAgentStep, files, triggerFileDownload]) + }, [addMessageHandler, addMessage, mapMessages, agent.setCurrentAgentStep, files, triggerFileDownload, addAttachment, addPendingFileEvent, resolvePendingFileEvent]) const selectAllServerTools = useCallback((server) => { const group = config.tools.find(t => t.server === server); if (!group) return @@ -230,7 +272,7 @@ export const ChatProvider = ({ children }) => { }, [selections, selectedTools, selectedPrompts, config.tools, config.prompts]) // Flatten ragServers into a single list of data source objects for easier consumption - const ragSources = config.ragServers.flatMap(server => + const ragSources = config.ragServers.flatMap(server => server.sources.map(source => ({ ...source, serverName: server.server, @@ -239,6 +281,42 @@ export const ChatProvider = ({ children }) => { })) ) + // ensureSession: ensures a session exists, returns sessionId once ready + const ensureSession = useCallback(() => { + return new Promise((resolve) => { + if (sessionId) { + resolve(sessionId) + return + } + + // Create a temporary session ID for frontend tracking + const tempSessionId = `session_${Date.now()}_${generateSecureRandomString()}` + setSessionId(tempSessionId) + + // Send reset_session to create a new session on backend + sendMessage({ type: 'reset_session', user: config.user }) + + // For now, resolve immediately since backend handles session creation + // In a more robust implementation, we'd wait for session confirmation + resolve(tempSessionId) + }) + }, [sessionId, sendMessage, config.user]) + + // addSystemEvent: adds a system event message to the chat timeline + const addSystemEvent = useCallback((subtype, text, meta = {}) => { + const eventId = `system_${Date.now()}_${generateSecureRandomString()}` + addMessage({ + role: 'system', + type: 'system', + subtype, + text, + meta, + timestamp: new Date().toISOString(), + id: eventId + }) + return eventId + }, [addMessage]) + const value = { appName: config.appName, user: config.user, @@ -305,6 +383,13 @@ export const ChatProvider = ({ children }) => { taggedFiles: files.taggedFiles, toggleFileTag: files.toggleFileTag, clearTaggedFiles: files.clearTaggedFiles, + sessionId, + attachments, + addAttachment, + addPendingFileEvent, + resolvePendingFileEvent, + ensureSession, + addSystemEvent, settings, } diff --git a/frontend/src/handlers/chat/websocketHandlers.js b/frontend/src/handlers/chat/websocketHandlers.js index 988c8b5..e42b458 100644 --- a/frontend/src/handlers/chat/websocketHandlers.js +++ b/frontend/src/handlers/chat/websocketHandlers.js @@ -14,7 +14,9 @@ export function createWebSocketHandler(deps) { setCustomUIContent, setSessionFiles, getFileType, - triggerFileDownload + triggerFileDownload, + addAttachment, + resolvePendingFileEvent } = deps const handleAgentUpdate = (data) => { @@ -138,7 +140,12 @@ export function createWebSocketHandler(deps) { setSessionFiles(() => { if (updateData.files && updateData.files.length > 0) { const viewableFiles = updateData.files.filter(f => ['png','jpg','jpeg','gif','svg','pdf','html'].includes(f.filename.toLowerCase().split('.').pop())) - if (viewableFiles.length > 0) { + // Only auto-open canvas for tool-generated files or when explicitly requested + // Skip auto-open for library attachments (source === 'user') + const shouldAutoOpenCanvas = viewableFiles.length > 0 && + !updateData.files.every(f => f.source === 'user') + + if (shouldAutoOpenCanvas) { const cFiles = viewableFiles.map(f => ({ ...f, type: getFileType(f.filename) })) setCanvasFiles(cFiles) setCurrentCanvasFileIndex(0) @@ -262,6 +269,46 @@ export function createWebSocketHandler(deps) { // Could show a toast notification here } break + case 'file_attach': + // Handle file attachment response + if (data.success) { + // File was successfully attached - add to attachments state + if (typeof addAttachment === 'function' && data.s3_key) { + addAttachment(data.s3_key) + } + + // Try to update pending event in-place, fallback to adding new message + if (typeof resolvePendingFileEvent === 'function' && data.s3_key) { + resolvePendingFileEvent(data.s3_key, 'file-attached', `Added '${data.filename || 'file'}' to this session.`) + } else { + addMessage({ + role: 'system', + type: 'system', + subtype: 'file-attached', + text: `Added '${data.filename || 'file'}' to this session.`, + meta: { fileId: data.s3_key, fileName: data.filename, source: 'library' }, + timestamp: new Date().toISOString(), + id: `file_attach_${Date.now()}` + }) + } + } else { + // File attachment failed + // Try to update pending event in-place, fallback to adding new message + if (typeof resolvePendingFileEvent === 'function' && data.s3_key) { + resolvePendingFileEvent(data.s3_key, 'file-attach-error', `Failed to add file to session: ${data.error || 'Unknown error'}`) + } else { + addMessage({ + role: 'system', + type: 'system', + subtype: 'file-attach-error', + text: `Failed to add file to session: ${data.error || 'Unknown error'}`, + meta: { fileId: data.s3_key, source: 'library' }, + timestamp: new Date().toISOString(), + id: `file_attach_error_${Date.now()}` + }) + } + } + break case 'intermediate_update': handleIntermediateUpdate(data) break diff --git a/frontend/src/test/basic.test.js b/frontend/src/test/basic.test.js index bedba4a..3e3cf20 100644 --- a/frontend/src/test/basic.test.js +++ b/frontend/src/test/basic.test.js @@ -2,7 +2,7 @@ * Basic test to verify test infrastructure is working */ -import { describe, it, expect } from 'vitest'; +import { describe, it, expect, vi } from 'vitest'; describe('Basic Test Infrastructure', () => { it('should run a simple test', () => { @@ -18,4 +18,41 @@ describe('Basic Test Infrastructure', () => { expect(arr.length).toBe(3); expect(arr).toContain(2); }); -}); \ No newline at end of file +}); + +describe('File Attachment System', () => { + it('should handle ensureSession promise resolution', async () => { + // Mock sendMessage function + const mockSendMessage = vi.fn(); + + // Simulate ensureSession logic + const ensureSession = () => { + return new Promise((resolve) => { + const tempSessionId = `session_${Date.now()}_${Math.random().toString(36).substr(2, 9)}`; + mockSendMessage({ type: 'reset_session', user: 'test@example.com' }); + resolve(tempSessionId); + }); + }; + + const sessionId = await ensureSession(); + + expect(sessionId).toMatch(/^session_\d+_[a-z0-9]+$/); + expect(mockSendMessage).toHaveBeenCalledWith({ + type: 'reset_session', + user: 'test@example.com' + }); + }); + + it('should handle file attachment system events', () => { + const events = [ + { subtype: 'file-attaching', text: 'Adding test.pdf to this session...' }, + { subtype: 'file-attached', text: 'Added test.pdf to this session.' }, + { subtype: 'file-attach-error', text: 'Failed to add file: error message' } + ]; + + events.forEach(event => { + expect(event.subtype).toMatch(/^file-(attaching|attached|attach-error)$/); + expect(event.text).toBeDefined(); + }); + }); +});