wuyang/zk-data-agent
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Refactor session runtime persistence

Browse Source
This commit is contained in:
wuyang6
2026-06-12 16:17:58 +08:00
parent 77d360c1e8
commit 953126e1d3
15 changed files with 717 additions and 253 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/api/server.py
+247 -141
View File
@@ -37,7 +37,7 @@ from fastapi.responses import FileResponse, JSONResponse, Response, StreamingRes
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel, Field
from src.agent_session import AgentMessage, AgentSessionState
from src.agent_session import AgentSessionState
from src.agent_runtime import LocalCodingAgent
from src.agent_slash_commands import get_slash_command_specs
from src.agent_types import (
@@ -71,7 +71,9 @@ from src.run_state_store import ACTIVE_RUN_STATUSES, RunStateStore
from src.session_store import (
DEFAULT_AGENT_SESSION_DIR,
StoredAgentSession,
delete_agent_session,
deserialize_runtime_config,
list_agent_sessions,
load_agent_session,
save_agent_session,
serialize_model_config,
@@ -435,6 +437,8 @@ class RunManager:
normalized = _normalize_run_event(event)
if normalized is None:
return None
if normalized.get('type') in {'content_delta', 'tool_delta'}:
return None
normalized['recorded_at'] = time.time()
with self._lock:
record = self._runs.get(run_id)
@@ -1665,35 +1669,20 @@ def create_app(state: AgentState) -> FastAPI:
include_children: bool = False,
) -> list[dict[str, Any]]:
directory = state.account_paths(account_id)['sessions']
if not directory.exists():
return []
results: list[dict[str, Any]] = []
for path in sorted(_iter_session_files(directory), key=_session_mtime, reverse=True):
try:
data = json.loads(path.read_text(encoding='utf-8'))
except (OSError, json.JSONDecodeError):
continue
session_id = str(data.get('session_id', _session_id_from_path(path)))
session_metadata = (
data.get('session_metadata')
if isinstance(data.get('session_metadata'), dict)
else {}
)
for stored, updated_at in list_agent_sessions(directory):
session_id = stored.session_id
session_metadata = stored.session_metadata or {}
if (
not include_children
and session_metadata.get('visibility') == 'child'
):
continue
usage = data.get('usage') if isinstance(data.get('usage'), dict) else {}
model_config = (
data.get('model_config')
if isinstance(data.get('model_config'), dict)
else {}
)
messages = data.get('messages') or []
title = data.get('title')
title = session_metadata.get('title')
if not isinstance(title, str) or not title.strip():
title = None
preview = ''
for msg in messages:
for msg in _stored_display_messages(stored):
if isinstance(msg, dict) and msg.get('role') == 'user':
content = msg.get('content', '')
if isinstance(content, str) and not _is_internal_message(content):
@@ -1702,13 +1691,13 @@ def create_app(state: AgentState) -> FastAPI:
results.append(
{
'session_id': session_id,
'turns': data.get('turns', 0),
'tool_calls': data.get('tool_calls', 0),
'turns': stored.turns,
'tool_calls': stored.tool_calls,
'preview': title if isinstance(title, str) and title.strip() else preview,
'modified_at': _session_mtime(path),
'model': model_config.get('model'),
'usage': usage,
'is_training': bool(data.get('is_training', False)),
'modified_at': updated_at,
'model': stored.model_config.get('model'),
'usage': stored.usage,
'is_training': stored.is_training,
'session_metadata': session_metadata,
}
)
@@ -1730,7 +1719,7 @@ def create_app(state: AgentState) -> FastAPI:
async def delete_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 or not _delete_session_files(directory, safe_id):
if safe_id is None or not delete_agent_session(safe_id, directory=directory):
raise HTTPException(status_code=404, detail='Session not found')
return {'deleted': True, 'session_id': safe_id}
@@ -2356,20 +2345,26 @@ def create_app(state: AgentState) -> FastAPI:
def _run_chat_payload(
request: ChatRequest,
event_sink: Any | None = None,
run_record: RunRecord | None = None,
) -> dict[str, Any]:
requested_session_id = _safe_session_id(
request.resume_session_id or request.session_id
) or uuid4().hex
account_key = state._account_key(request.account_id)
prompt = request.prompt.strip()
run_record = state.run_manager.start(account_key, requested_session_id, prompt)
state.run_state_store.start(
run_id=run_record.run_id,
account_key=account_key,
session_id=requested_session_id,
pending_prompt=prompt,
started_at=run_record.started_at,
)
if run_record is None:
run_record = state.run_manager.start(
account_key,
requested_session_id,
prompt,
)
state.run_state_store.start(
run_id=run_record.run_id,
account_key=account_key,
session_id=requested_session_id,
pending_prompt=prompt,
started_at=run_record.started_at,
)
run_lock = state.run_lock_for(request.account_id, requested_session_id)
agent = state.agent_for(request.account_id, requested_session_id)
config = state.config_for(request.account_id)
@@ -2627,6 +2622,53 @@ def create_app(state: AgentState) -> FastAPI:
)
return payload
def _start_chat_run(request: ChatRequest) -> dict[str, Any]:
requested_session_id = _safe_session_id(
request.resume_session_id or request.session_id
) or uuid4().hex
account_key = state._account_key(request.account_id)
prompt = request.prompt.strip()
run_record = state.run_manager.start(account_key, requested_session_id, prompt)
state.run_state_store.start(
run_id=run_record.run_id,
account_key=account_key,
session_id=requested_session_id,
pending_prompt=prompt,
started_at=run_record.started_at,
)
# Persist the submitted user message before the worker starts so the UI
# can immediately reload this exact session from DB without depending on
# browser-local optimistic state.
try:
agent = state.agent_for(request.account_id, requested_session_id)
_save_in_progress_session(
directory=state.account_paths(request.account_id)['sessions'],
agent=agent,
session_id=requested_session_id,
prompt=prompt,
)
except Exception:
# The worker will surface the real failure through run_state_store.
pass
def worker() -> None:
try:
_run_chat_payload(request, run_record=run_record)
except Exception as exc: # noqa: BLE001
print(
f'[chat-start] background run failed '
f'session={requested_session_id} run={run_record.run_id}: {exc}',
flush=True,
)
threading.Thread(target=worker, daemon=True).start()
return {
'session_id': requested_session_id,
'run_id': run_record.run_id,
'status': 'queued',
'started_at': run_record.started_at,
}
# ------------- chat ------------------------------------------------------
@app.post('/api/chat')
async def chat(request: ChatRequest) -> dict[str, Any]:
@@ -2651,6 +2693,24 @@ def create_app(state: AgentState) -> FastAPI:
)
return payload
@app.post('/api/chat/start')
async def chat_start(request: ChatRequest) -> dict[str, Any]:
prompt = request.prompt.strip()
if not prompt:
raise HTTPException(status_code=400, detail='Prompt is empty')
try:
return _start_chat_run(request)
except HTTPException:
raise
except Exception as exc:
return JSONResponse(
status_code=500,
content={
'error': str(exc),
'error_type': type(exc).__name__,
},
)
@app.post('/api/chat/stream')
async def chat_stream(request: ChatRequest) -> StreamingResponse:
prompt = request.prompt.strip()
@@ -2775,11 +2835,12 @@ def _truncate_api_text(text: str, limit: int) -> str:
def _serialize_stored_session(stored: StoredAgentSession) -> dict[str, Any]:
display_messages = _stored_display_messages(stored)
return {
'session_id': stored.session_id,
'turns': stored.turns,
'tool_calls': stored.tool_calls,
'messages': [_normalize_transcript_entry(dict(m)) for m in stored.messages],
'messages': [_normalize_transcript_entry(dict(m)) for m in display_messages],
'usage': stored.usage,
'total_cost_usd': stored.total_cost_usd,
'model': stored.model_config.get('model'),
@@ -2788,6 +2849,14 @@ def _serialize_stored_session(stored: StoredAgentSession) -> dict[str, Any]:
}
def _stored_display_messages(
stored: StoredAgentSession,
) -> tuple[dict[str, Any], ...]:
if stored.display_messages:
return tuple(dict(message) for message in stored.display_messages)
return tuple(dict(message) for message in stored.messages)
def _stored_session_has_incomplete_tail(stored: StoredAgentSession) -> bool:
_, changed = _trim_incomplete_message_tail(stored.messages)
if changed:
@@ -2855,35 +2924,45 @@ def _mark_session_interrupted(
except (FileNotFoundError, OSError, json.JSONDecodeError):
return
trimmed, changed = _trim_incomplete_message_tail(stored.messages)
display_trimmed, display_changed = _trim_incomplete_message_tail(
_stored_display_messages(stored)
)
budget_state = (
dict(stored.budget_state)
if isinstance(stored.budget_state, dict)
else {}
)
if not changed and budget_state.get('status') not in {'queued', 'running'}:
if (
not changed
and not display_changed
and budget_state.get('status') not in {'queued', 'running'}
):
return
messages = list(trimmed)
display_messages = list(display_trimmed)
run_status_message = {
'role': 'assistant',
'content': _interrupted_session_message(status),
'state': 'final',
'stop_reason': status,
'metadata': {
'kind': 'run_status',
'status': status,
'detail': detail[:500],
'created_at_ms': int(time.time() * 1000),
},
}
if not _last_message_is_run_status(messages, status):
messages.append(
{
'role': 'assistant',
'content': _interrupted_session_message(status),
'state': 'final',
'stop_reason': status,
'metadata': {
'kind': 'run_status',
'status': status,
'detail': detail[:500],
'created_at_ms': int(time.time() * 1000),
},
}
)
messages.append(dict(run_status_message))
if not _last_message_is_run_status(display_messages, status):
display_messages.append(dict(run_status_message))
budget_state['status'] = status
budget_state['interrupted_at'] = int(time.time())
save_agent_session(
replace(
stored,
messages=tuple(messages),
display_messages=tuple(display_messages),
budget_state=budget_state,
),
directory=directory,
@@ -2900,7 +2979,16 @@ def _sanitize_stored_session_for_resume(
messages = _without_run_status_messages(trimmed)
if len(messages) != len(trimmed):
changed = True
if not changed:
display_trimmed, display_changed = _trim_incomplete_message_tail(
_stored_display_messages(stored)
)
if display_trimmed and _is_pending_user_message(display_trimmed[-1]):
display_trimmed = display_trimmed[:-1]
display_changed = True
display_messages = _without_run_status_messages(display_trimmed)
if len(display_messages) != len(display_trimmed):
display_changed = True
if not changed and not display_changed:
return stored
budget_state = (
dict(stored.budget_state)
@@ -2911,6 +2999,7 @@ def _sanitize_stored_session_for_resume(
return replace(
stored,
messages=messages,
display_messages=display_messages,
budget_state=budget_state,
)
@@ -3146,17 +3235,22 @@ def _save_in_progress_session(
except (FileNotFoundError, OSError, json.JSONDecodeError):
return None
messages = list(stored.messages)
display_messages = list(_stored_display_messages(stored))
if messages and _is_pending_user_message(dict(messages[-1])):
messages.pop()
if display_messages and _is_pending_user_message(dict(display_messages[-1])):
display_messages.pop()
pending_message = {
'role': 'user',
'content': prompt,
'state': 'final',
'metadata': pending_metadata,
'message_id': f'user_pending_{int(time.time() * 1000)}',
}
messages.append(
{
'role': 'user',
'content': prompt,
'state': 'final',
'metadata': pending_metadata,
'message_id': f'user_pending_{int(time.time() * 1000)}',
}
dict(pending_message)
)
display_messages.append(dict(pending_message))
budget_state = (
dict(stored.budget_state)
if isinstance(stored.budget_state, dict)
@@ -3168,6 +3262,7 @@ def _save_in_progress_session(
replace(
stored,
messages=tuple(messages),
display_messages=tuple(display_messages),
budget_state=budget_state,
),
directory=directory,
@@ -3188,6 +3283,15 @@ def _save_in_progress_session(
metadata=pending_metadata,
message_id='user_pending_0',
)
session_metadata: dict[str, Any] = {}
if initial_title:
session_metadata.update(
{
'title': initial_title,
'title_source': 'first_message',
'title_generated_at': int(time.time()),
}
)
stored = StoredAgentSession(
session_id=safe_id,
model_config=serialize_model_config(agent.model_config),
@@ -3195,7 +3299,8 @@ def _save_in_progress_session(
system_prompt_parts=session.system_prompt_parts,
user_context=dict(session.user_context),
system_context=dict(session.system_context),
messages=session.transcript(),
messages=session.model_transcript(),
display_messages=session.display_transcript(),
turns=0,
tool_calls=0,
usage={},
@@ -3204,17 +3309,9 @@ def _save_in_progress_session(
budget_state={'status': 'running'},
plugin_state={},
scratchpad_directory=str(scratchpad_directory),
session_metadata=session_metadata,
)
saved_path = save_agent_session(stored, directory=directory)
if initial_title:
try:
data = json.loads(saved_path.read_text(encoding='utf-8'))
except (OSError, json.JSONDecodeError):
return initial_title
data['title'] = initial_title
data['title_source'] = 'first_message'
data['title_generated_at'] = int(time.time())
_write_session_metadata(saved_path, data)
save_agent_session(stored, directory=directory)
return initial_title
except OSError:
return None
@@ -3232,15 +3329,12 @@ def _derive_initial_session_title(prompt: str) -> str | None:
def _session_state_from_stored(stored: StoredAgentSession) -> AgentSessionState:
return AgentSessionState(
return AgentSessionState.from_persisted(
system_prompt_parts=stored.system_prompt_parts,
user_context=stored.user_context,
system_context=stored.system_context,
messages=[
AgentMessage.from_openai_message(message)
for message in stored.messages
if isinstance(message, dict)
],
messages=stored.messages,
display_messages=stored.display_messages,
)
@@ -3432,32 +3526,39 @@ def _annotate_last_assistant_elapsed(
except (OSError, json.JSONDecodeError):
return
messages = data.get('messages')
if not isinstance(messages, list):
display_messages = data.get('display_messages')
if not isinstance(display_messages, list):
display_messages = messages
if not isinstance(messages, list) and not isinstance(display_messages, list):
return
for message in reversed(messages):
if not isinstance(message, dict) or message.get('role') != 'assistant':
for message_list in (messages, display_messages):
if not isinstance(message_list, list):
continue
metadata = message.get('metadata')
if not isinstance(metadata, dict):
metadata = {}
message['metadata'] = metadata
metadata['elapsed_ms'] = elapsed_ms
try:
path.write_text(json.dumps(data, ensure_ascii=False, indent=2), encoding='utf-8')
except OSError:
return
for message in reversed(message_list):
if not isinstance(message, dict) or message.get('role') != 'assistant':
continue
metadata = message.get('metadata')
if not isinstance(metadata, dict):
metadata = {}
message['metadata'] = metadata
metadata['elapsed_ms'] = elapsed_ms
break
try:
path.write_text(json.dumps(data, ensure_ascii=False, indent=2), encoding='utf-8')
except OSError:
return
return
def _read_session_title(directory: Path, session_id: str | None) -> str | None:
if not session_id:
return None
path = _session_json_path(directory, session_id)
try:
data = json.loads(path.read_text(encoding='utf-8'))
except (OSError, json.JSONDecodeError):
stored = load_agent_session(session_id, directory=directory)
except FileNotFoundError:
return None
title = data.get('title')
metadata = stored.session_metadata or {}
title = metadata.get('title')
if isinstance(title, str) and title.strip():
return title.strip()
return None
@@ -3473,38 +3574,37 @@ def _ensure_session_title(
previous_title: str | None = None,
fallback_title: str | None = None,
) -> None:
path = _session_json_path(directory, session_id)
try:
data = json.loads(path.read_text(encoding='utf-8'))
except (OSError, json.JSONDecodeError):
stored = load_agent_session(session_id, directory=directory)
except FileNotFoundError:
return
metadata = dict(stored.session_metadata or {})
if previous_title:
data['title'] = previous_title
_write_session_metadata(path, data)
metadata['title'] = previous_title
save_agent_session(
replace(stored, session_metadata=metadata),
directory=directory,
)
return
title = data.get('title')
title_source = data.get('title_source')
title = metadata.get('title')
title_source = metadata.get('title_source')
if (
isinstance(title, str)
and title.strip()
and title_source != 'first_message'
):
return
messages = data.get('messages')
if not isinstance(messages, list):
if fallback_title and not (isinstance(title, str) and title.strip()):
data['title'] = fallback_title
data['title_source'] = 'first_message'
data['title_generated_at'] = int(time.time())
_write_session_metadata(path, data)
return
messages = [dict(message) for message in _stored_display_messages(stored)]
user_messages = _session_user_messages(messages)
if not user_messages:
if fallback_title and not (isinstance(title, str) and title.strip()):
data['title'] = fallback_title
data['title_source'] = 'first_message'
data['title_generated_at'] = int(time.time())
_write_session_metadata(path, data)
metadata['title'] = fallback_title
metadata['title_source'] = 'first_message'
metadata['title_generated_at'] = int(time.time())
save_agent_session(
replace(stored, session_metadata=metadata),
directory=directory,
)
return
generated = _generate_session_title(
user_messages,
@@ -3514,15 +3614,21 @@ def _ensure_session_title(
)
if not generated:
if fallback_title and not (isinstance(title, str) and title.strip()):
data['title'] = fallback_title
data['title_source'] = 'first_message'
data['title_generated_at'] = int(time.time())
_write_session_metadata(path, data)
metadata['title'] = fallback_title
metadata['title_source'] = 'first_message'
metadata['title_generated_at'] = int(time.time())
save_agent_session(
replace(stored, session_metadata=metadata),
directory=directory,
)
return
data['title'] = generated
data['title_source'] = 'llm'
data['title_generated_at'] = int(time.time())
_write_session_metadata(path, data)
metadata['title'] = generated
metadata['title_source'] = 'llm'
metadata['title_generated_at'] = int(time.time())
save_agent_session(
replace(stored, session_metadata=metadata),
directory=directory,
)
def _session_user_messages(messages: list[Any]) -> list[str]:
@@ -5069,7 +5175,7 @@ def _admin_session_tool_call_count(data: dict[str, Any]) -> int:
file_history = data.get('file_history')
file_history_count = len(file_history) if isinstance(file_history, list) else 0
message_count = max(
_admin_count_tool_calls_in_items(data.get('messages')),
_admin_count_tool_calls_in_items(data.get('display_messages') or data.get('messages')),
_admin_count_tool_calls_in_items(data.get('turns')),
)
derived_count = max(file_history_count, message_count)
@@ -5303,30 +5409,30 @@ def _delete_session_files(directory: Path, session_id: str) -> bool:
def _update_session_title(directory: Path, session_id: str, title: str) -> bool:
path = _session_json_path(directory, session_id)
if not path.exists():
return False
try:
data = json.loads(path.read_text(encoding='utf-8'))
except (OSError, json.JSONDecodeError):
stored = load_agent_session(session_id, directory=directory)
except FileNotFoundError:
return False
data['title'] = title
data['title_source'] = 'manual'
data['title_updated_at'] = int(time.time())
_write_session_metadata(path, data)
metadata = dict(stored.session_metadata or {})
metadata['title'] = title
metadata['title_source'] = 'manual'
metadata['title_updated_at'] = int(time.time())
save_agent_session(
replace(stored, session_metadata=metadata),
directory=directory,
)
return True
def _update_session_training(directory: Path, session_id: str, is_training: bool) -> bool:
path = _session_json_path(directory, session_id)
if not path.exists():
return False
try:
data = json.loads(path.read_text(encoding='utf-8'))
except (OSError, json.JSONDecodeError):
stored = load_agent_session(session_id, directory=directory)
except FileNotFoundError:
return False
data['is_training'] = is_training
_write_session_metadata(path, data)
save_agent_session(
replace(stored, is_training=is_training),
directory=directory,
)
return True
@@ -7022,7 +7128,7 @@ def _last_assistant_text(sessions_dir: Path, session_id: str | None) -> str | No
data = json.loads(path.read_text(encoding='utf-8'))
except (OSError, json.JSONDecodeError):
return None
messages = data.get('messages')
messages = data.get('display_messages') or data.get('messages')
if not isinstance(messages, list):
return None
for msg in reversed(messages):
@@ -7732,7 +7838,7 @@ def _extract_trigger_from_session(
data = json.loads(path.read_text(encoding='utf-8'))
except (OSError, json.JSONDecodeError):
return None
messages = data.get('messages')
messages = data.get('display_messages') or data.get('messages')
if not isinstance(messages, list):
return None
for msg in messages:
frontend/app/app/api/chat/route.ts
+10 -13
View File
@@ -21,6 +21,8 @@ export const maxDuration = 3600;
type ClawChatResponse = {
final_output?: string;
session_id?: string;
run_id?: string;
status?: string;
tool_calls?: number;
stop_reason?: string;
elapsed_ms?: number;
@@ -130,16 +132,14 @@ export async function POST(req: Request) {
writer.write({
type: "reasoning-delta",
id: "reasoning-1",
delta: "请求已发送，等待后端开始处理。",
delta: "任务已提交到后端运行队列。",
});
const payload = await callClawBackendStream(
const payload = await startClawBackendRun(
userPrompt,
runtimeContext,
account.id,
sessionId,
resumeId,
writer,
streamState,
);
const elapsedMs = payload.elapsed_ms ?? Date.now() - startedAt;
const text = formatClawResponse(payload);
@@ -148,13 +148,13 @@ export async function POST(req: Request) {
writer.write({
type: "reasoning-delta",
id: "reasoning-1",
delta: `\n思考并执行完成，用时 ${formatDuration(elapsedMs)}。`,
delta: `\n后端已接管执行，用时 ${formatDuration(elapsedMs)}。`,
});
writer.write({ type: "reasoning-end", id: "reasoning-1" });
streamState.reasoningEnded = true;
}
writeToolTrace(writer, payload.transcript, streamState);
if (!streamState.textStreamed) {
if (!streamState.textStreamed && text) {
writeTextDelta(writer, text, streamState);
}
endTextPart(writer, streamState);
@@ -164,6 +164,7 @@ export async function POST(req: Request) {
finishReason: payload.error ? "error" : "stop",
messageMetadata: {
sessionId: payload.session_id,
runId: payload.run_id,
toolCalls: payload.tool_calls,
stopReason: payload.stop_reason,
elapsedMs,
@@ -411,17 +412,15 @@ function safeFilename(value: string) {
return value.replace(/[^a-zA-Z0-9._-]+/g, "_").slice(0, 120) || "attachment";
}
async function callClawBackendStream(
async function startClawBackendRun(
prompt: string,
runtimeContext: string,
accountId: string,
sessionId: string,
resumeSessionId?: string,
writer?: UIMessageStreamWriter<UIMessage>,
streamState?: StreamState,
): Promise<ClawChatResponse> {
try {
const response = await fetch(`${CLAW_API_URL}/api/chat/stream`, {
const response = await fetch(`${CLAW_API_URL}/api/chat/start`, {
method: "POST",
headers: { "content-type": "application/json" },
body: JSON.stringify({
@@ -443,9 +442,7 @@ async function callClawBackendStream(
`Claw backend returned ${response.status}`,
};
}
if (!response.body)
return { error: "Claw backend returned an empty stream" };
return await consumeClawStream(response.body, writer, streamState);
return (await response.json()) as ClawChatResponse;
} catch (err) {
return {
error:
frontend/app/app/assistant.tsx
+6 -8
View File
@@ -71,9 +71,9 @@ export const Assistant = ({ initialSessionId }: AssistantProps = {}) => {
: null;
const selectedResumeSessionId = selectedSessionId;
const outgoingSessionId =
lastSessionId ??
selectedResumeSessionId ??
pendingWorkspaceSessionId ??
lastSessionId ??
options.id;
const lastUserMessage = [...messages]
.reverse()
@@ -87,7 +87,7 @@ export const Assistant = ({ initialSessionId }: AssistantProps = {}) => {
id: outgoingSessionId,
messages: lastUserMessage ? [lastUserMessage] : [],
resumeSessionId:
lastSessionId ?? selectedResumeSessionId ?? undefined,
selectedResumeSessionId ?? lastSessionId ?? undefined,
},
};
},
@@ -102,15 +102,13 @@ export const Assistant = ({ initialSessionId }: AssistantProps = {}) => {
console.log("[replay-session] called", { sessionId });
writeActiveSessionId(sessionId);
pushSessionUrl(sessionId);
// 后端已经落盘/结束后，用回放结果接管当前会话。
// 先取消本地仍挂起的 assistant-ui run，避免 UI 残留“运行中”且无法停止。
if (runtime.thread.getState().isRunning) {
runtime.thread.cancelRun();
}
runtime.thread.import(repository);
},
[runtime],
);
const clearSession = useCallback(() => {
runtime.thread.import({ headId: null, messages: [] });
}, [runtime]);
useEffect(() => {
const sessionId = initialSessionId?.trim();
@@ -162,7 +160,7 @@ export const Assistant = ({ initialSessionId }: AssistantProps = {}) => {
return (
<AssistantRuntimeProvider runtime={runtime}>
<ClawSessionReplayProvider value={{ replaySession }}>
<ClawSessionReplayProvider value={{ replaySession, clearSession }}>
<ActivityProvider>
<SidebarProvider>
<div className="flex h-dvh w-full">
frontend/app/components/assistant-ui/thread-list.tsx
+20 -13
View File
@@ -124,24 +124,22 @@ export const ThreadList: FC = () => {
};
const ThreadListNew: FC = () => {
const { clearSession } = useClawSessionReplay();
return (
<div
onClickCapture={() => {
<Button
type="button"
variant="outline"
className="aui-thread-list-new h-9 w-full justify-start gap-2 rounded-lg px-3 text-sm hover:bg-muted data-active:bg-muted"
onClick={() => {
clearActiveSessionId();
pushHomeUrl();
clearSession();
window.dispatchEvent(new Event("claw-active-session-cleared"));
}}
>
<ThreadListPrimitive.New asChild>
<Button
variant="outline"
className="aui-thread-list-new h-9 w-full justify-start gap-2 rounded-lg px-3 text-sm hover:bg-muted data-active:bg-muted"
>
<PlusIcon className="size-4" />
New Task
</Button>
</ThreadListPrimitive.New>
</div>
<PlusIcon className="size-4" />
New Task
</Button>
);
};
@@ -754,6 +752,7 @@ export function toReplayRepository(
}
const id = `${fallbackSessionId}-run-${runStatus.run_id ?? "active"}`;
const runStartedAtMs = resolveRunStartedAtMs(runStatus);
const runCreatedAt = new Date(runStartedAtMs ?? Date.now());
const parts = buildActiveRunParts(runStatus) as UIMessage["parts"];
const uiMessage: UIMessage = {
id,
@@ -763,12 +762,17 @@ export function toReplayRepository(
sessionId: session.session_id ?? fallbackSessionId,
runId: runStatus.run_id,
runStartedAtMs,
custom: {
sessionId: session.session_id ?? fallbackSessionId,
runId: runStatus.run_id,
runStartedAtMs,
},
},
} as UIMessage;
const threadMessage = {
id,
role: "assistant",
createdAt: new Date(),
createdAt: runCreatedAt,
content: toThreadMessageContent(parts),
status: { type: "running" },
metadata: {
@@ -776,6 +780,9 @@ export function toReplayRepository(
unstable_annotations: [],
unstable_data: [],
steps: [],
sessionId: session.session_id ?? fallbackSessionId,
runId: runStatus.run_id,
runStartedAtMs,
custom: {
sessionId: session.session_id ?? fallbackSessionId,
runId: runStatus.run_id,
frontend/app/components/assistant-ui/thread.tsx
+35 -23
View File
@@ -434,11 +434,15 @@ function useCurrentThreadSessionId({
clearPendingWorkspaceSessionId();
}
}, [currentRun.sessionId, pendingSessionId]);
const replayRunSessionId = currentRun.active ? currentRun.sessionId : null;
const replayRunSessionId =
currentRun.active &&
(!activeSessionId || currentRun.sessionId === activeSessionId)
? currentRun.sessionId
: null;
return (
replayRunSessionId ??
(includePending ? pendingSessionId : null) ??
(includeActive ? activeSessionId : null) ??
(includePending ? pendingSessionId : null) ??
replayRunSessionId ??
currentRun.sessionId
);
}
@@ -894,6 +898,10 @@ const Composer: FC = () => {
includePending: true,
includeActive: true,
});
const replayedRunIsCurrent =
replayedRun.active &&
Boolean(workspaceSessionId) &&
replayedRun.sessionId === workspaceSessionId;
return (
<ComposerPrimitive.Root className="aui-composer-root relative flex w-full flex-col">
<ComposerPrimitive.AttachmentDropzone asChild>
@@ -909,7 +917,7 @@ const Composer: FC = () => {
rows={1}
autoFocus
aria-label="Message input"
disabled={replayedRun.active}
disabled={replayedRunIsCurrent}
/>
<ComposerAction />
</div>
@@ -960,9 +968,18 @@ const ComposerAction: FC = () => {
includeActive: true,
});
const [runtimeCancelling, setRuntimeCancelling] = useState(false);
const cancelSessionId = replayedRun.sessionId ?? currentSessionId;
const replayedRunIsCurrent =
replayedRun.active &&
Boolean(currentSessionId) &&
replayedRun.sessionId === currentSessionId;
const cancelSessionId = replayedRunIsCurrent
? replayedRun.sessionId
: currentSessionId;
const backendRunStatus = useBackendActiveRunStatus(cancelSessionId);
const cancelRunId = replayedRun.runId ?? backendRunStatus?.run_id ?? null;
const cancelRunId =
(replayedRunIsCurrent ? replayedRun.runId : null) ??
backendRunStatus?.run_id ??
null;
useEffect(() => {
if (!runtimeRunning) setRuntimeCancelling(false);
}, [runtimeRunning]);
@@ -973,7 +990,7 @@ const ComposerAction: FC = () => {
<ComposerAssistButtons />
</div>
<ComposerContextStatus />
{!runtimeRunning && !replayedRun.active && !backendRunStatus ? (
{!runtimeRunning && !replayedRunIsCurrent && !backendRunStatus ? (
<ComposerPrimitive.Send asChild>
<TooltipIconButton
tooltip="Send message"
@@ -1012,7 +1029,7 @@ const ComposerAction: FC = () => {
<SquareIcon className="aui-composer-cancel-icon size-3 fill-current" />
</Button>
) : null}
{!runtimeRunning && (replayedRun.active || backendRunStatus) ? (
{!runtimeRunning && (replayedRunIsCurrent || backendRunStatus) ? (
<ReplayedRunCancelButton
sessionId={cancelSessionId}
runId={cancelRunId}
@@ -1323,16 +1340,6 @@ function useComposerContextStatus() {
setModel: async () => {},
});
useEffect(() => {
const cleanLatestSessionId = normalizeSessionId(latestSessionId);
if (cleanLatestSessionId && !isRunning) {
writeActiveSessionId(cleanLatestSessionId);
window.setTimeout(() => {
window.dispatchEvent(new Event("claw-sessions-changed"));
}, 0);
}
}, [isRunning, latestSessionId]);
useEffect(() => {
if (wasRunningRef.current && !isRunning) {
window.dispatchEvent(new Event("claw-sessions-changed"));
@@ -1359,11 +1366,8 @@ function useComposerContextStatus() {
);
const latestMessageSessionId = normalizeSessionId(latestSessionId);
let sessionId = normalizeSessionId(
isRunning
? activeSessionId || latestMessageSessionId
: latestMessageSessionId || activeSessionId,
activeSessionId || latestMessageSessionId,
);
if (sessionId) writeActiveSessionId(sessionId);
let sessionPayload: ClawStoredSession | null = null;
let contextBudget: ContextBudget | undefined;
if (sessionId) {
@@ -2063,10 +2067,10 @@ const ActivityChainSummary: FC<{
metadataRunStartedAtMs ?? contentRunStartedAtMs ?? messageCreatedAtMs;
const liveStartedAtRef = useRef<number | null>(initialStartedAtMs);
const [elapsedMs, setElapsedMs] = useState<number | null>(() => {
if (storedElapsedMs !== null) return storedElapsedMs;
if (initialStartedAtMs !== null) {
return Math.max(0, Date.now() - initialStartedAtMs);
}
if (storedElapsedMs !== null) return storedElapsedMs;
return null;
});
const label = useAuiState((s) => {
@@ -2095,6 +2099,14 @@ const ActivityChainSummary: FC<{
liveStartedAtRef.current = authoritativeStartedAt;
} else if (liveStartedAtRef.current === null && storedElapsedMs !== null) {
liveStartedAtRef.current = Date.now() - storedElapsedMs;
} else if (
liveStartedAtRef.current !== null &&
storedElapsedMs !== null
) {
liveStartedAtRef.current = Math.min(
liveStartedAtRef.current,
Date.now() - storedElapsedMs,
);
} else if (
liveStartedAtRef.current === null &&
messageCreatedAtMs !== null
frontend/app/components/assistant-ui/threadlist-sidebar.tsx
+7 -9
View File
@@ -1,5 +1,4 @@
import type { ExportedMessageRepository } from "@assistant-ui/core";
import { ThreadListPrimitive } from "@assistant-ui/react";
import {
BarChart3Icon,
BotIcon,
@@ -185,6 +184,7 @@ function CollapsedSidebarRail({
onLogout: () => void;
}) {
const { setOpen } = useSidebar();
const { clearSession } = useClawSessionReplay();
return (
<div className="flex h-full w-full flex-col items-center border-r bg-sidebar py-3 text-sidebar-foreground">
@@ -200,19 +200,17 @@ function CollapsedSidebarRail({
</button>
<div className="flex flex-col items-center gap-3">
<div
onClickCapture={() => {
<CollapsedIconButton
label="新任务"
onClick={() => {
clearActiveSessionId();
pushHomeUrl();
clearSession();
window.dispatchEvent(new Event("claw-active-session-cleared"));
}}
>
<ThreadListPrimitive.New asChild>
<CollapsedIconButton label="新任务">
<SquarePenIcon className="size-4" />
</CollapsedIconButton>
</ThreadListPrimitive.New>
</div>
<SquarePenIcon className="size-4" />
</CollapsedIconButton>
<CollapsedSessionSearch />
<CollapsedRecentSessions />
</div>
frontend/app/lib/claw-session-replay.tsx
+1
View File
@@ -8,6 +8,7 @@ type ClawSessionReplayContextValue = {
sessionId: string,
repository: ExportedMessageRepository,
) => void;
clearSession: () => void;
};
const ClawSessionReplayContext =
src/agent_runtime.py
+9 -1
View File
@@ -569,6 +569,7 @@ class LocalCodingAgent:
user_context=stored_session.user_context,
system_context=stored_session.system_context,
messages=stored_session.messages,
display_messages=stored_session.display_messages,
)
self._append_file_history_replay_if_needed(
session,
@@ -916,6 +917,7 @@ class LocalCodingAgent:
'continuation_index': continuation_count,
},
message_id=f'continuation_{turn_index}',
display=False,
)
stream_events.append(
{
@@ -1073,6 +1075,7 @@ class LocalCodingAgent:
'continuation_index': continuation_count,
},
message_id=f'continuation_{turn_index}',
display=False,
)
stream_events.append(
{
@@ -1444,6 +1447,7 @@ class LocalCodingAgent:
'plugin_preflight_count': len(plugin_preflight_messages),
},
message_id=f'plugin_tool_runtime_{tool_call.id}',
display=False,
)
stream_events.append(
{
@@ -3862,6 +3866,7 @@ class LocalCodingAgent:
'file_history_snapshot_count': snapshot_count,
},
message_id=f'file_history_replay_{replay_count}',
display=False,
)
def _render_file_history_replay(
@@ -3977,6 +3982,7 @@ class LocalCodingAgent:
message_id=(
f'compaction_replay_{len(compact_messages)}_{len(snipped_messages)}'
),
display=False,
)
def _render_compaction_replay(
@@ -4194,6 +4200,7 @@ class LocalCodingAgent:
'message_count': len(persist_messages),
},
message_id=f'plugin_persist_{result.session_id}',
display=False,
)
persist_events.append(
{
@@ -4239,7 +4246,8 @@ class LocalCodingAgent:
system_prompt_parts=session.system_prompt_parts,
user_context=dict(session.user_context),
system_context=dict(session.system_context),
messages=session.transcript(),
messages=session.model_transcript(),
display_messages=session.display_transcript(),
turns=previous_turns + result.turns,
tool_calls=previous_tool_calls + result.tool_calls,
usage=result.usage.to_dict(),
src/agent_session.py
+128 -33
View File
@@ -102,8 +102,16 @@ class AgentSessionState:
user_context: dict[str, str] = field(default_factory=dict)
system_context: dict[str, str] = field(default_factory=dict)
messages: list[AgentMessage] = field(default_factory=list)
display_messages: list[AgentMessage] = field(default_factory=list)
mutation_serial: int = 0
def __post_init__(self) -> None:
# Backward compatibility for tests and older callers that construct an
# AgentSessionState directly. Runtime-created sessions append messages
# through helpers and explicitly decide whether each message is visible.
if self.messages and not self.display_messages:
self.display_messages = list(self.messages)
@classmethod
def create(
cls,
@@ -118,42 +126,48 @@ class AgentSessionState:
user_context=dict(user_context or {}),
system_context=dict(system_context or {}),
)
state.messages.append(
state._append_message(
AgentMessage(
role='system',
content='\n\n'.join(
_append_system_context(system_prompt_parts, state.system_context)
),
blocks=_text_blocks('\n\n'.join(_append_system_context(system_prompt_parts, state.system_context))),
message_id='system_0',
metadata=_initialize_message_metadata(
role='system',
message_id='system_0',
),
)
),
display=False,
)
if state.user_context:
state.messages.append(
state._append_message(
AgentMessage(
role='user',
content=_render_user_context_reminder(state.user_context),
blocks=_text_blocks(_render_user_context_reminder(state.user_context)),
message_id='user_context_0',
metadata=_initialize_message_metadata(
role='user',
message_id='user_context_0',
),
)
),
display=False,
)
if user_prompt is not None:
state.messages.append(
state._append_message(
AgentMessage(
role='user',
content=user_prompt,
blocks=_text_blocks(user_prompt),
message_id='user_0',
metadata=_initialize_message_metadata(
role='user',
message_id='user_0',
),
)
),
display=True,
)
return state
@@ -165,41 +179,47 @@ class AgentSessionState:
message_id: str | None = None,
stop_reason: str | None = None,
usage: UsageStats | None = None,
display: bool = True,
) -> None:
self.messages.append(
actual_message_id = message_id or f'assistant_{len(self.messages)}'
self._append_message(
AgentMessage(
role='assistant',
content=content,
tool_calls=tool_calls,
blocks=_assistant_blocks(content, tool_calls),
message_id=message_id,
message_id=actual_message_id,
stop_reason=stop_reason,
usage=usage or UsageStats(),
metadata=_initialize_message_metadata(
role='assistant',
message_id=message_id or f'assistant_{len(self.messages)}',
message_id=actual_message_id,
),
)
),
display=display,
)
def start_assistant(
self,
*,
message_id: str | None = None,
display: bool = True,
) -> int:
self.messages.append(
actual_message_id = message_id or f'assistant_{len(self.messages)}'
self._append_message(
AgentMessage(
role='assistant',
content='',
tool_calls=(),
blocks=(),
message_id=message_id,
message_id=actual_message_id,
state='streaming',
metadata=_initialize_message_metadata(
role='assistant',
message_id=message_id or f'assistant_{len(self.messages)}',
message_id=actual_message_id,
),
)
),
display=display,
)
return len(self.messages) - 1
@@ -214,12 +234,14 @@ class AgentSessionState:
mutation_serial=self._next_mutation_serial(),
)
merged_metadata = _advance_lineage_revision(merged_metadata)
self.messages[index] = replace(
updated = replace(
message,
content=message.content + delta,
blocks=_assistant_blocks(message.content + delta, message.tool_calls),
metadata=merged_metadata,
)
self.messages[index] = updated
self._replace_display_message(message, updated)
def merge_assistant_tool_call_delta(
self,
@@ -261,12 +283,14 @@ class AgentSessionState:
mutation_serial=self._next_mutation_serial(),
)
merged_metadata = _advance_lineage_revision(merged_metadata)
self.messages[index] = replace(
updated = replace(
message,
tool_calls=tuple(tool_calls),
blocks=_assistant_blocks(message.content, tuple(tool_calls)),
metadata=merged_metadata,
)
self.messages[index] = updated
self._replace_display_message(message, updated)
def finalize_assistant(
self,
@@ -285,7 +309,7 @@ class AgentSessionState:
mutation_serial=self._next_mutation_serial(),
)
merged_metadata = _advance_lineage_revision(merged_metadata)
self.messages[index] = replace(
updated = replace(
message,
state='final',
stop_reason=finish_reason,
@@ -293,6 +317,8 @@ class AgentSessionState:
blocks=_assistant_blocks(message.content, message.tool_calls),
metadata=merged_metadata,
)
self.messages[index] = updated
self._replace_display_message(message, updated)
def append_user(
self,
@@ -301,8 +327,10 @@ class AgentSessionState:
model_content: str | None = None,
metadata: dict[str, Any] | None = None,
message_id: str | None = None,
display: bool = True,
) -> None:
self.messages.append(
actual_message_id = message_id or f'user_{len(self.messages)}'
self._append_message(
AgentMessage(
role='user',
content=content,
@@ -310,27 +338,38 @@ class AgentSessionState:
blocks=_text_blocks(content),
metadata=_initialize_message_metadata(
role='user',
message_id=message_id or f'user_{len(self.messages)}',
message_id=actual_message_id,
metadata=dict(metadata or {}),
),
message_id=message_id,
)
message_id=actual_message_id,
),
display=display,
)
def append_tool(self, name: str, tool_call_id: str, content: str) -> None:
self.messages.append(
def append_tool(
self,
name: str,
tool_call_id: str,
content: str,
*,
display: bool = True,
) -> None:
actual_message_id = f'tool_{len(self.messages)}'
self._append_message(
AgentMessage(
role='tool',
content=content,
name=name,
tool_call_id=tool_call_id,
blocks=_tool_blocks(name, tool_call_id, content),
message_id=actual_message_id,
metadata=_initialize_message_metadata(
role='tool',
message_id=f'tool_{len(self.messages)}',
message_id=actual_message_id,
metadata={'tool_name': name, 'tool_call_id': tool_call_id},
),
)
),
display=display,
)
def start_tool(
@@ -340,26 +379,29 @@ class AgentSessionState:
tool_call_id: str,
message_id: str | None = None,
metadata: dict[str, Any] | None = None,
display: bool = True,
) -> int:
self.messages.append(
actual_message_id = message_id or f'tool_{len(self.messages)}'
self._append_message(
AgentMessage(
role='tool',
content='',
name=name,
tool_call_id=tool_call_id,
blocks=(),
message_id=message_id,
message_id=actual_message_id,
state='streaming',
metadata=_initialize_message_metadata(
role='tool',
message_id=message_id or f'tool_{len(self.messages)}',
message_id=actual_message_id,
metadata={
'tool_name': name,
'tool_call_id': tool_call_id,
**dict(metadata or {}),
},
),
)
),
display=display,
)
return len(self.messages) - 1
@@ -383,12 +425,14 @@ class AgentSessionState:
merged_metadata = _advance_lineage_revision(merged_metadata)
if metadata:
merged_metadata.update(metadata)
self.messages[index] = replace(
updated = replace(
message,
content=message.content + delta,
blocks=_tool_blocks(message.name, message.tool_call_id, message.content + delta),
metadata=merged_metadata,
)
self.messages[index] = updated
self._replace_display_message(message, updated)
def finalize_tool(
self,
@@ -413,7 +457,7 @@ class AgentSessionState:
merged_metadata = _advance_lineage_revision(merged_metadata)
if metadata:
merged_metadata.update(metadata)
self.messages[index] = replace(
updated = replace(
message,
content=content,
blocks=_tool_blocks(message.name, message.tool_call_id, content),
@@ -421,6 +465,8 @@ class AgentSessionState:
stop_reason=stop_reason,
metadata=merged_metadata,
)
self.messages[index] = updated
self._replace_display_message(message, updated)
def update_message(
self,
@@ -453,7 +499,7 @@ class AgentSessionState:
merged_metadata = _advance_lineage_revision(merged_metadata)
if metadata:
merged_metadata.update(metadata)
self.messages[index] = replace(
updated = replace(
message,
content=new_content,
blocks=_derive_blocks(
@@ -468,6 +514,8 @@ class AgentSessionState:
stop_reason=new_stop_reason,
metadata=merged_metadata,
)
self.messages[index] = updated
self._replace_display_message(message, updated)
def tombstone_message(
self,
@@ -491,8 +539,40 @@ class AgentSessionState:
return [message.to_openai_message() for message in self.messages]
def transcript(self) -> tuple[JSONDict, ...]:
return self.display_transcript()
def model_transcript(self) -> tuple[JSONDict, ...]:
return tuple(message.to_transcript_entry() for message in self.messages)
def display_transcript(self) -> tuple[JSONDict, ...]:
return tuple(message.to_transcript_entry() for message in self.display_messages)
def _append_message(self, message: AgentMessage, *, display: bool) -> None:
self.messages.append(message)
if display:
self.display_messages.append(message)
def _replace_display_message(
self,
previous_message: AgentMessage,
updated_message: AgentMessage,
) -> None:
index = self._find_display_message_index(previous_message)
if index is not None:
self.display_messages[index] = updated_message
def _find_display_message_index(self, message: AgentMessage) -> int | None:
if message.message_id:
for index in range(len(self.display_messages) - 1, -1, -1):
if self.display_messages[index].message_id == message.message_id:
return index
lineage_id = message.metadata.get('lineage_id')
if isinstance(lineage_id, str) and lineage_id:
for index in range(len(self.display_messages) - 1, -1, -1):
if self.display_messages[index].metadata.get('lineage_id') == lineage_id:
return index
return None
def _next_mutation_serial(self) -> int:
self.mutation_serial += 1
return self.mutation_serial
@@ -505,12 +585,27 @@ class AgentSessionState:
user_context: dict[str, str] | None,
system_context: dict[str, str] | None,
messages: tuple[JSONDict, ...] | list[JSONDict],
display_messages: tuple[JSONDict, ...] | list[JSONDict] | None = None,
) -> 'AgentSessionState':
model_messages = [
AgentMessage.from_openai_message(message)
for message in messages
if isinstance(message, dict)
]
display_source = messages if display_messages is None else display_messages
visible_messages = [
AgentMessage.from_openai_message(message)
for message in display_source
if isinstance(message, dict)
]
if display_messages is not None and not visible_messages:
visible_messages = list(model_messages)
return cls(
system_prompt_parts=tuple(system_prompt_parts),
user_context=dict(user_context or {}),
system_context=dict(system_context or {}),
messages=[AgentMessage.from_openai_message(message) for message in messages],
messages=model_messages,
display_messages=visible_messages,
mutation_serial=max(
(
int(message.get('metadata', {}).get('last_mutation_serial', 0))
src/run_state_store.py
+6 -1
View File
@@ -275,6 +275,11 @@ class RunStateStore:
continue
if isinstance(event, dict):
events.append(event)
elapsed_ms = row['elapsed_ms']
if row['status'] in ACTIVE_RUN_STATUSES:
started_at = row['started_at']
if isinstance(started_at, (int, float)):
elapsed_ms = max(0, int((time.time() - float(started_at)) * 1000))
return {
'run_id': row['run_id'],
'session_id': row['session_id'],
@@ -283,7 +288,7 @@ class RunStateStore:
'started_at': row['started_at'],
'updated_at': row['updated_at'],
'finished_at': row['finished_at'],
'elapsed_ms': row['elapsed_ms'],
'elapsed_ms': elapsed_ms,
'pending_prompt': row['pending_prompt'] or '',
'error': row['error'] or '',
'cancellable': bool(row['cancellable']),
src/session_store.py
+192 -9
View File
@@ -1,6 +1,8 @@
from __future__ import annotations
import json
import sqlite3
import time
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Any
@@ -27,6 +29,7 @@ class StoredSession:
DEFAULT_SESSION_DIR = Path('.port_sessions')
DEFAULT_AGENT_SESSION_DIR = DEFAULT_SESSION_DIR / 'agent'
AGENT_SESSION_DB_FILENAME = 'sessions.db'
def save_session(session: StoredSession, directory: Path | None = None) -> Path:
@@ -70,6 +73,7 @@ class StoredAgentSession:
scratchpad_directory: str | None = None
is_training: bool = False
session_metadata: JSONDict | None = None
display_messages: tuple[JSONDict, ...] = ()
def save_agent_session(session: StoredAgentSession, directory: Path | None = None) -> Path:
@@ -78,16 +82,96 @@ def save_agent_session(session: StoredAgentSession, directory: Path | None = Non
session_dir = target_dir / session.session_id
session_dir.mkdir(parents=True, exist_ok=True)
path = session_dir / 'session.json'
path.write_text(json.dumps(asdict(session), indent=2), encoding='utf-8')
payload = asdict(session)
_write_agent_session_db_payload(target_dir, session.session_id, payload, path)
path.write_text(json.dumps(payload, indent=2, ensure_ascii=False), encoding='utf-8')
return path
def load_agent_session(session_id: str, directory: Path | None = None) -> StoredAgentSession:
target_dir = directory or DEFAULT_AGENT_SESSION_DIR
data = _read_agent_session_db_payload(target_dir, session_id)
if data is not None:
return _stored_agent_session_from_payload(data)
path = target_dir / session_id / 'session.json'
if not path.exists():
path = target_dir / f'{session_id}.json'
data = json.loads(path.read_text(encoding='utf-8'))
# Backfill SQLite on first read so future live UI reads do not depend on
# repeatedly parsing JSON snapshots.
_write_agent_session_db_payload(target_dir, session_id, data, path)
return _stored_agent_session_from_payload(data)
def list_agent_sessions(
directory: Path | None = None,
) -> list[tuple[StoredAgentSession, float]]:
target_dir = directory or DEFAULT_AGENT_SESSION_DIR
sessions: dict[str, tuple[StoredAgentSession, float]] = {}
for data, updated_at in _iter_agent_session_db_payloads(target_dir):
try:
stored = _stored_agent_session_from_payload(data)
except (KeyError, TypeError, ValueError):
continue
sessions[stored.session_id] = (stored, updated_at)
for path in _iter_agent_session_snapshot_files(target_dir):
session_id = _session_id_from_snapshot_path(path)
if session_id in sessions:
continue
try:
stored = load_agent_session(session_id, directory=target_dir)
except (FileNotFoundError, json.JSONDecodeError, KeyError, TypeError, ValueError):
continue
try:
updated_at = path.stat().st_mtime
except OSError:
updated_at = 0.0
sessions[stored.session_id] = (stored, updated_at)
return sorted(sessions.values(), key=lambda item: item[1], reverse=True)
def delete_agent_session(session_id: str, directory: Path | None = None) -> bool:
target_dir = directory or DEFAULT_AGENT_SESSION_DIR
deleted = False
with _connect_agent_session_db(target_dir) as conn:
cursor = conn.execute(
'delete from agent_sessions where session_id = ?',
(session_id,),
)
deleted = cursor.rowcount > 0
nested = target_dir / session_id
if nested.exists():
import shutil
shutil.rmtree(nested)
deleted = True
legacy = target_dir / f'{session_id}.json'
if legacy.exists():
legacy.unlink()
deleted = True
return deleted
def _stored_agent_session_from_payload(data: JSONDict) -> StoredAgentSession:
messages = tuple(
message for message in data['messages'] if isinstance(message, dict)
)
display_messages = tuple(
message
for message in data.get('display_messages', messages)
if isinstance(message, dict)
)
if not display_messages:
display_messages = messages
session_metadata = (
dict(data.get('session_metadata', {}))
if isinstance(data.get('session_metadata'), dict)
else {}
)
for key in ('title', 'title_source', 'title_updated_at'):
if key in data and key not in session_metadata:
session_metadata[key] = data[key]
return StoredAgentSession(
session_id=data['session_id'],
model_config=dict(data['model_config']),
@@ -95,9 +179,7 @@ def load_agent_session(session_id: str, directory: Path | None = None) -> Stored
system_prompt_parts=tuple(data['system_prompt_parts']),
user_context=dict(data['user_context']),
system_context=dict(data['system_context']),
messages=tuple(
message for message in data['messages'] if isinstance(message, dict)
),
messages=messages,
turns=int(data['turns']),
tool_calls=int(data['tool_calls']),
usage=dict(data.get('usage', {})),
@@ -121,14 +203,115 @@ def load_agent_session(session_id: str, directory: Path | None = None) -> Stored
else None
),
is_training=bool(data.get('is_training', False)),
session_metadata=(
dict(data.get('session_metadata', {}))
if isinstance(data.get('session_metadata'), dict)
else {}
),
session_metadata=session_metadata,
display_messages=display_messages,
)
def _write_agent_session_db_payload(
directory: Path,
session_id: str,
payload: JSONDict,
snapshot_path: Path,
) -> None:
updated_at = time.time()
payload_json = json.dumps(payload, ensure_ascii=False, sort_keys=True)
with _connect_agent_session_db(directory) as conn:
conn.execute(
"""
insert into agent_sessions (
session_id, updated_at, snapshot_path, payload_json
)
values (?, ?, ?, ?)
on conflict(session_id) do update set
updated_at=excluded.updated_at,
snapshot_path=excluded.snapshot_path,
payload_json=excluded.payload_json
""",
(session_id, updated_at, str(snapshot_path), payload_json),
)
def _read_agent_session_db_payload(
directory: Path,
session_id: str,
) -> JSONDict | None:
with _connect_agent_session_db(directory) as conn:
row = conn.execute(
'select payload_json from agent_sessions where session_id = ?',
(session_id,),
).fetchone()
if row is None:
return None
try:
payload = json.loads(row['payload_json'])
except (TypeError, json.JSONDecodeError):
return None
return payload if isinstance(payload, dict) else None
def _iter_agent_session_db_payloads(directory: Path) -> list[tuple[JSONDict, float]]:
with _connect_agent_session_db(directory) as conn:
rows = conn.execute(
"""
select payload_json, updated_at
from agent_sessions
order by updated_at desc
"""
).fetchall()
result: list[tuple[JSONDict, float]] = []
for row in rows:
try:
payload = json.loads(row['payload_json'])
except (TypeError, json.JSONDecodeError):
continue
if isinstance(payload, dict):
result.append((payload, float(row['updated_at'] or 0.0)))
return result
def _connect_agent_session_db(directory: Path) -> sqlite3.Connection:
directory.mkdir(parents=True, exist_ok=True)
conn = sqlite3.connect(directory.parent / AGENT_SESSION_DB_FILENAME, timeout=10)
conn.row_factory = sqlite3.Row
conn.execute('pragma journal_mode=wal')
conn.execute('pragma busy_timeout=5000')
conn.execute(
"""
create table if not exists agent_sessions (
session_id text primary key,
updated_at real not null,
snapshot_path text default '',
payload_json text not null
)
"""
)
conn.execute(
"""
create index if not exists idx_agent_sessions_updated_at
on agent_sessions(updated_at)
"""
)
return conn
def _iter_agent_session_snapshot_files(directory: Path) -> list[Path]:
if not directory.exists():
return []
by_session_id: dict[str, Path] = {}
for path in directory.glob('*.json'):
by_session_id[_session_id_from_snapshot_path(path)] = path
for path in directory.glob('*/session.json'):
by_session_id[_session_id_from_snapshot_path(path)] = path
return list(by_session_id.values())
def _session_id_from_snapshot_path(path: Path) -> str:
if path.name == 'session.json':
return path.parent.name
return path.stem
def serialize_model_config(model_config: ModelConfig) -> JSONDict:
return {
'model': model_config.model,
tests/test_agent_runtime.py
+7 -2
View File
@@ -1025,8 +1025,13 @@ class AgentRuntimeTests(unittest.TestCase):
message for message in result.transcript
if message.get('metadata', {}).get('kind') == 'compact_boundary'
]
self.assertEqual(len(transcript_compact_messages), 1)
compact_metadata = transcript_compact_messages[0].get('metadata', {})
self.assertEqual(len(transcript_compact_messages), 0)
model_compact_messages = [
message for message in (agent.last_session.messages if agent.last_session else [])
if message.metadata.get('kind') == 'compact_boundary'
]
self.assertEqual(len(model_compact_messages), 1)
compact_metadata = model_compact_messages[0].metadata
self.assertEqual(compact_metadata.get('compaction_depth'), 1)
self.assertEqual(compact_metadata.get('nested_compaction_count'), 0)
self.assertIn('preserved_tail_ids', compact_metadata)
tests/test_compact.py
+4
View File
@@ -255,6 +255,10 @@ class TestCompactConversation(unittest.TestCase):
self.assertNotIn('<analysis>', result.summary_text)
# Summary should contain the actual summary content
self.assertIn('User wanted to test compaction', result.summary_text)
# UI display transcript must remain append-only and not be replaced by
# compact_summary / compact_boundary model-context messages.
display_contents = [m.content for m in agent.last_session.display_messages]
self.assertEqual(display_contents, [m.content for m in msgs])
def test_api_error_returns_compaction_error(self) -> None:
with tempfile.TemporaryDirectory() as tmp_dir:
tests/test_gui_server.py
+19
View File
@@ -483,6 +483,21 @@ class GuiServerTests(unittest.TestCase):
},
{'role': 'user', 'content': 'continue'},
),
display_messages=(
{'role': 'user', 'content': 'hello'},
{
'role': 'assistant',
'content': '上一次任务已中断，后台没有正在执行的进程。',
'state': 'final',
'stop_reason': 'interrupted',
'metadata': {'kind': 'run_status', 'status': 'interrupted'},
},
{
'role': 'user',
'content': 'pending',
'metadata': {'status': 'running', 'placeholder': True},
},
),
turns=1,
tool_calls=0,
usage={},
@@ -498,6 +513,10 @@ class GuiServerTests(unittest.TestCase):
[message['content'] for message in sanitized.messages],
['hello', 'continue'],
)
self.assertEqual(
[message['content'] for message in sanitized.display_messages],
['hello'],
)
def test_runtime_event_stage_names_long_running_work(self) -> None:
self.assertEqual(
tests/test_session_store.py
+26
View File
@@ -80,6 +80,7 @@ class TestStoredAgentSessionRoundTrip(unittest.TestCase):
'user_context': {'lang': 'en'},
'system_context': {'os': 'linux'},
'messages': ({'role': 'user', 'content': 'hi'},),
'display_messages': ({'role': 'user', 'content': 'visible hi'},),
'turns': 3,
'tool_calls': 7,
'usage': {'input_tokens': 500, 'output_tokens': 300},
@@ -106,6 +107,7 @@ class TestStoredAgentSessionRoundTrip(unittest.TestCase):
self.assertEqual(loaded.user_context, session.user_context)
self.assertEqual(loaded.system_context, session.system_context)
self.assertEqual(loaded.messages, session.messages)
self.assertEqual(loaded.display_messages, session.display_messages)
self.assertEqual(loaded.turns, session.turns)
self.assertEqual(loaded.tool_calls, session.tool_calls)
self.assertEqual(loaded.usage, session.usage)
@@ -157,6 +159,30 @@ class TestStoredAgentSessionRoundTrip(unittest.TestCase):
self.assertEqual(len(loaded.messages), 2)
self.assertEqual(loaded.messages[0]['role'], 'user')
self.assertEqual(loaded.messages[1]['role'], 'assistant')
self.assertEqual(loaded.display_messages, loaded.messages)
def test_load_agent_session_falls_back_to_messages_for_display(self) -> None:
with tempfile.TemporaryDirectory() as td:
directory = Path(td)
path = directory / 'legacy.json'
data = {
'session_id': 'legacy',
'model_config': {},
'runtime_config': {'cwd': '/'},
'system_prompt_parts': [],
'user_context': {},
'system_context': {},
'messages': [{'role': 'user', 'content': 'legacy hi'}],
'turns': 0,
'tool_calls': 0,
'usage': {},
'total_cost_usd': 0.0,
'file_history': [],
}
path.write_text(json.dumps(data))
loaded = load_agent_session('legacy', directory=directory)
self.assertEqual(loaded.display_messages, loaded.messages)
def test_load_defaults_for_missing_optional_fields(self) -> None:
"""Missing optional fields get sensible defaults."""