Persist in-progress chat sessions

This commit is contained in:
武阳
2026-05-07 15:50:30 +08:00
parent d84b6c568c
commit 7a8152fba8
2 changed files with 119 additions and 4 deletions
+71 -4
View File
@@ -39,6 +39,9 @@ from src.session_store import (
StoredAgentSession,
deserialize_runtime_config,
load_agent_session,
save_agent_session,
serialize_model_config,
serialize_runtime_config,
)
from src.token_budget import calculate_token_budget
@@ -531,8 +534,11 @@ def create_app(state: AgentState) -> FastAPI:
@app.get('/api/sessions/{session_id}')
async def get_session(session_id: str, account_id: str | None = None) -> dict[str, Any]:
directory = state.account_paths(account_id)['sessions']
safe_id = _safe_session_id(session_id)
if safe_id is None:
raise HTTPException(status_code=404, detail='Session not found')
try:
stored = load_agent_session(session_id, directory=directory)
stored = load_agent_session(safe_id, directory=directory)
except FileNotFoundError:
raise HTTPException(status_code=404, detail='Session not found')
return _serialize_stored_session(stored)
@@ -568,8 +574,11 @@ def create_app(state: AgentState) -> FastAPI:
with state.lock():
if session_id:
directory = state.account_paths(account_id)['sessions']
safe_id = _safe_session_id(session_id)
if safe_id is None:
raise HTTPException(status_code=404, detail='Session not found')
try:
stored = load_agent_session(session_id, directory=directory)
stored = load_agent_session(safe_id, directory=directory)
except FileNotFoundError:
raise HTTPException(status_code=404, detail='Session not found')
session = _session_state_from_stored(stored)
@@ -607,9 +616,14 @@ def create_app(state: AgentState) -> FastAPI:
)
started_at = time.perf_counter()
if request.resume_session_id is not None:
if requested_session_id is None:
raise HTTPException(
status_code=404,
detail='Session to resume not found',
)
try:
stored = load_agent_session(
request.resume_session_id,
requested_session_id,
directory=session_directory,
)
except FileNotFoundError:
@@ -624,9 +638,15 @@ def create_app(state: AgentState) -> FastAPI:
event_sink=event_sink,
)
else:
_save_in_progress_session(
directory=session_directory,
agent=agent,
session_id=requested_session_id,
prompt=request.prompt.strip(),
)
result = agent.run(
request.prompt.strip(),
session_id=_safe_session_id(request.session_id),
session_id=requested_session_id,
runtime_context=request.runtime_context,
event_sink=event_sink,
)
@@ -770,6 +790,53 @@ def _serialize_stored_session(stored: StoredAgentSession) -> dict[str, Any]:
}
def _save_in_progress_session(
*,
directory: Path,
agent: LocalCodingAgent,
session_id: str | None,
prompt: str,
) -> None:
safe_id = _safe_session_id(session_id)
if safe_id is None or not prompt:
return
if _session_json_path(directory, safe_id).exists():
return
# 新会话的第一轮执行可能很久。先写一个运行中占位,避免刷新页面后找不到会话。
scratchpad_directory = (
agent.runtime_config.scratchpad_root / safe_id / 'scratchpad'
).resolve()
try:
scratchpad_directory.mkdir(parents=True, exist_ok=True)
session = agent.build_session(None, scratchpad_directory=scratchpad_directory)
session.append_user(
prompt,
metadata={'status': 'running', 'placeholder': True},
message_id='user_pending_0',
)
stored = StoredAgentSession(
session_id=safe_id,
model_config=serialize_model_config(agent.model_config),
runtime_config=serialize_runtime_config(agent.runtime_config),
system_prompt_parts=session.system_prompt_parts,
user_context=dict(session.user_context),
system_context=dict(session.system_context),
messages=session.transcript(),
turns=0,
tool_calls=0,
usage={},
total_cost_usd=0.0,
file_history=(),
budget_state={'status': 'running'},
plugin_state={},
scratchpad_directory=str(scratchpad_directory),
)
save_agent_session(stored, directory=directory)
except OSError:
return
def _session_state_from_stored(stored: StoredAgentSession) -> AgentSessionState:
return AgentSessionState(
system_prompt_parts=stored.system_prompt_parts,
+48
View File
@@ -17,6 +17,7 @@ from pathlib import Path
from fastapi.testclient import TestClient
from backend.api.server import AgentState, create_app
from src.agent_types import AgentRunResult
from src.session_store import StoredAgentSession, load_agent_session, save_agent_session
@@ -161,6 +162,53 @@ class GuiServerTests(unittest.TestCase):
payload = response.json()
self.assertIsNone(payload['session_id'])
def test_chat_persists_in_progress_session_before_agent_finishes(self) -> None:
with tempfile.TemporaryDirectory() as d:
root = Path(d)
client, state = _build_client(root)
agent = state.agent_for('alice')
def fake_run(
prompt: str,
session_id: str | None = None,
*,
runtime_context: str | None = None,
event_sink: object | None = None,
) -> AgentRunResult:
stored = load_agent_session(
'pending-1',
directory=root / 'accounts' / 'alice' / 'sessions',
)
self.assertEqual(stored.session_id, 'pending-1')
self.assertEqual(stored.budget_state.get('status'), 'running')
self.assertEqual(stored.messages[-1]['role'], 'user')
self.assertEqual(stored.messages[-1]['content'], 'hello')
return AgentRunResult(
final_output='done',
turns=1,
tool_calls=0,
transcript=stored.messages
+ ({'role': 'assistant', 'content': 'done'},),
session_id=session_id,
)
original_run = agent.run
agent.run = fake_run # type: ignore[method-assign]
try:
response = client.post(
'/api/chat',
json={
'prompt': 'hello',
'account_id': 'alice',
'session_id': ' pending-1 ',
},
)
finally:
agent.run = original_run # type: ignore[method-assign]
self.assertEqual(response.status_code, 200)
self.assertEqual(response.json()['session_id'], 'pending-1')
def test_chat_uses_account_scoped_model_config(self) -> None:
with tempfile.TemporaryDirectory() as d:
client, _ = _build_client(Path(d))