Add assistant-ui data agent frontend

tests/test_gui_server.py

@@ -11,11 +11,13 @@ from __future__ import annotations
 
 import tempfile
 import unittest
+import json
 from pathlib import Path
 
 from fastapi.testclient import TestClient
 
 from backend.api.server import AgentState, create_app
+from src.session_store import StoredAgentSession, load_agent_session, save_agent_session
 
 
 def _build_client(tmp: Path) -> tuple[TestClient, AgentState]:
@@ -63,6 +65,16 @@ class GuiServerTests(unittest.TestCase):
             self.assertTrue(data['allow_shell'])
             self.assertEqual(data['model'], 'other-model')
 
+    def test_state_snapshot_can_scope_to_account(self) -> None:
+        with tempfile.TemporaryDirectory() as d:
+            client, _ = _build_client(Path(d))
+            response = client.get('/api/state', params={'account_id': 'alice'})
+            self.assertEqual(response.status_code, 200)
+            payload = response.json()
+            self.assertEqual(payload['account_id'], 'alice')
+            self.assertIn('/accounts/alice/sessions', payload['session_directory'])
+            self.assertIn('/accounts/alice/sessions', payload['upload_directory'])
+
     def test_state_update_rejects_missing_cwd(self) -> None:
         with tempfile.TemporaryDirectory() as d:
             client, _ = _build_client(Path(d))
@@ -104,6 +116,47 @@ class GuiServerTests(unittest.TestCase):
             self.assertIn('/help', payload['final_output'])
             self.assertIn('total_tokens', payload['usage'])
 
+    def test_chat_uses_requested_session_directory(self) -> None:
+        with tempfile.TemporaryDirectory() as d:
+            root = Path(d)
+            client, _ = _build_client(root)
+            response = client.post(
+                '/api/chat',
+                json={
+                    'prompt': '/help',
+                    'account_id': 'alice',
+                    'session_id': 'thread-1',
+                },
+            )
+            self.assertEqual(response.status_code, 200)
+            payload = response.json()
+            self.assertIsNone(payload['session_id'])
+
+    def test_agent_session_store_uses_session_subdirectory(self) -> None:
+        with tempfile.TemporaryDirectory() as d:
+            root = Path(d)
+            stored = StoredAgentSession(
+                session_id='thread-1',
+                model_config={'model': 'test-model'},
+                runtime_config={'cwd': str(root)},
+                system_prompt_parts=(),
+                user_context={},
+                system_context={},
+                messages=({'role': 'user', 'content': 'hello'},),
+                turns=1,
+                tool_calls=0,
+                usage={},
+                total_cost_usd=0.0,
+                file_history=(),
+                budget_state={},
+                plugin_state={},
+                scratchpad_directory=str(root / 'sessions' / 'thread-1' / 'scratchpad'),
+            )
+            path = save_agent_session(stored, directory=root / 'sessions')
+            self.assertEqual(path, root / 'sessions' / 'thread-1' / 'session.json')
+            loaded = load_agent_session('thread-1', directory=root / 'sessions')
+            self.assertEqual(loaded.session_id, 'thread-1')
+
     def test_chat_rejects_blank_prompt(self) -> None:
         with tempfile.TemporaryDirectory() as d:
             client, _ = _build_client(Path(d))
@@ -126,6 +179,66 @@ class GuiServerTests(unittest.TestCase):
             self.assertEqual(response.status_code, 200)
             self.assertEqual(response.json(), [])
 
+    def test_sessions_are_isolated_by_account(self) -> None:
+        with tempfile.TemporaryDirectory() as d:
+            root = Path(d)
+            client, _ = _build_client(root)
+            alice_sessions = root / 'accounts' / 'alice' / 'sessions'
+            alice_sessions.mkdir(parents=True)
+            (alice_sessions / 's1.json').write_text(
+                json.dumps(
+                    {
+                        'session_id': 's1',
+                        'turns': 1,
+                        'tool_calls': 0,
+                        'messages': [{'role': 'user', 'content': 'alice only'}],
+                    }
+                ),
+                encoding='utf-8',
+            )
+
+            alice = client.get('/api/sessions', params={'account_id': 'alice'})
+            bob = client.get('/api/sessions', params={'account_id': 'bob'})
+            self.assertEqual(alice.status_code, 200)
+            self.assertEqual(bob.status_code, 200)
+            self.assertEqual(alice.json()[0]['session_id'], 's1')
+            self.assertEqual(bob.json(), [])
+
+    def test_session_list_deduplicates_legacy_and_nested_files(self) -> None:
+        with tempfile.TemporaryDirectory() as d:
+            root = Path(d)
+            client, _ = _build_client(root)
+            session_root = root / 'accounts' / 'alice' / 'sessions'
+            session_root.mkdir(parents=True)
+            legacy = {
+                'session_id': 'same-id',
+                'turns': 1,
+                'tool_calls': 0,
+                'messages': [{'role': 'user', 'content': 'legacy'}],
+            }
+            nested = {
+                'session_id': 'same-id',
+                'turns': 2,
+                'tool_calls': 1,
+                'messages': [{'role': 'user', 'content': 'nested'}],
+            }
+            (session_root / 'same-id.json').write_text(
+                json.dumps(legacy),
+                encoding='utf-8',
+            )
+            (session_root / 'same-id').mkdir()
+            (session_root / 'same-id' / 'session.json').write_text(
+                json.dumps(nested),
+                encoding='utf-8',
+            )
+
+            response = client.get('/api/sessions', params={'account_id': 'alice'})
+            self.assertEqual(response.status_code, 200)
+            payload = response.json()
+            self.assertEqual(len(payload), 1)
+            self.assertEqual(payload[0]['session_id'], 'same-id')
+            self.assertEqual(payload[0]['preview'], 'nested')
+
     def test_session_detail_404_when_missing(self) -> None:
         with tempfile.TemporaryDirectory() as d:
             client, _ = _build_client(Path(d))