fix: stabilize background operations and large embedding results

This commit is contained in:
2026-09-06 16:26:17 +08:00
parent 874e916106
commit 3b9490e3fb
73 changed files with 3847 additions and 149 deletions
+51
View File
@@ -0,0 +1,51 @@
"""Serialize and batch durable Trace writes off the asyncio event loop."""
import asyncio
from contextvars import copy_context
class AsyncTraceWriter:
def __init__(self, repository):
self.repository = repository
self.queue = asyncio.Queue(maxsize=1024)
self.worker = None
async def submit(self, operation, *args):
future = asyncio.get_running_loop().create_future()
await self.queue.put((operation, args, future))
if self.worker is None or self.worker.done():
self.worker = asyncio.create_task(self._drain())
# Cancellation must not let an older snapshot commit after cancellation.
cancelled = False
while not future.done():
try:
await asyncio.shield(future)
except asyncio.CancelledError:
cancelled = True
future.result()
return cancelled
async def _drain(self):
while not self.queue.empty():
batch = []
while len(batch) < 64 and not self.queue.empty():
batch.append(self.queue.get_nowait())
try:
work = asyncio.get_running_loop().run_in_executor(
None, copy_context().run, self.repository.write_batch, [(op, args) for op, args, _ in batch])
# asyncio.run/shutdown may cancel every Task simultaneously. The
# executor Future survives; finish it and release all waiters.
while not work.done():
try:
await asyncio.shield(work)
except asyncio.CancelledError:
pass
work.result()
except Exception as exc:
for _, _, future in batch:
future.set_exception(exc)
else:
for _, _, future in batch:
future.set_result(None)
finally:
for _ in batch:
self.queue.task_done()
+128 -77
View File
@@ -4,6 +4,8 @@ from __future__ import annotations
import asyncio
import json
from app.agent.async_trace import AsyncTraceWriter
from app.operation_logs import log_event, agent_run_id
from collections.abc import AsyncIterator
from dataclasses import dataclass, field
from datetime import datetime, timezone
@@ -65,6 +67,9 @@ class RunRecord:
subscribers: set[asyncio.Queue[AgentEvent]] = field(default_factory=set)
task: asyncio.Task[None] | None = None
next_sequence: int = 0
publish_lock: asyncio.Lock = field(default_factory=asyncio.Lock)
cancel_lock: asyncio.Lock = field(default_factory=asyncio.Lock)
persisted_run: AgentRun | None = None
class AgentRuntime:
@@ -84,6 +89,7 @@ class AgentRuntime:
self.skills = skills
self.trace_repository = trace_repository or AgentTraceRepository()
self._records: dict[str, RunRecord] = {}
self._writer = AsyncTraceWriter(self.trace_repository)
async def create_run(self, request: AgentRunCreateRequest) -> AgentRun:
self._prune_records()
@@ -122,19 +128,25 @@ class AgentRuntime:
skill_config=skill_config,
allowed_tools=allowed_tools,
)
self.trace_repository.create_run(
run,
request,
self._config_snapshot(record),
)
# Reserve capacity before yielding to concurrent creators.
self._records[run.run_id] = record
try:
cancelled = await self._writer.submit('create', run.model_copy(deep=True), request.model_copy(deep=True), self._config_snapshot(record))
except BaseException:
self._records.pop(run.run_id, None)
raise
record.persisted_run = run.model_copy(deep=True)
log_event('agent', 'run.created', run_id=run.run_id, provider_id=run.provider_id, model=run.model)
if cancelled:
await self._finish_cancelled(record)
raise asyncio.CancelledError
record.task = asyncio.create_task(self._execute(record), name=run.run_id)
return run.model_copy(deep=True)
def get_run(self, run_id: str) -> AgentRun:
record = self._records.get(run_id)
if record is not None:
return record.run.model_copy(deep=True)
return (record.persisted_run or record.run).model_copy(deep=True)
run = self.trace_repository.recover_interrupted(run_id)
if run is None:
raise AgentRunNotFoundError(run_id)
@@ -145,7 +157,7 @@ class AgentRuntime:
recovered = [
self.trace_repository.recover_interrupted(item.run_id) or item
if item.run_id not in self._records
else self._records[item.run_id].run.model_copy(deep=True)
else (self._records[item.run_id].persisted_run or self._records[item.run_id].run).model_copy(deep=True)
for item in items
]
return recovered, total
@@ -154,25 +166,24 @@ class AgentRuntime:
record = self._records.get(run_id)
if record is None:
return self.get_run(run_id)
if record.run.status in TERMINAL_STATUSES:
return record.run.model_copy(deep=True)
record.run.cancelled = True
record.run.status = AgentRunStatus.cancelled
record.run.updated_at = datetime.now(timezone.utc)
self.permissions.cancel_run(run_id)
self._publish(record, AgentEventType.run_cancelled, {})
if record.task and not record.task.done():
record.task.cancel()
return record.run.model_copy(deep=True)
async with record.cancel_lock:
if record.task and not record.task.done():
if record.run.status not in TERMINAL_STATUSES:
record.task.cancel()
self.permissions.cancel_run(run_id)
await asyncio.gather(record.task, return_exceptions=True)
if record.run.status not in TERMINAL_STATUSES:
await self._finish_cancelled(record)
return (record.persisted_run or record.run).model_copy(deep=True)
def resolve_permission(self, run_id: str, request_id: str, decision: str) -> bool:
async def resolve_permission(self, run_id: str, request_id: str, decision: str) -> bool:
record = self._records.get(run_id)
if record is None:
return False
ticket = self.permissions.get_ticket(run_id, request_id)
resolved = self.permissions.resolve(run_id, request_id, decision)
if resolved:
self._publish(
await self._publish(
record,
AgentEventType.permission_resolved,
{
@@ -189,23 +200,24 @@ class AgentRuntime:
record = self._records.get(run_id)
run = self.get_run(run_id)
if record is None:
for event in self.trace_repository.list_events(
for event in await asyncio.to_thread(self.trace_repository.list_events,
run_id, after_sequence=after_sequence
):
yield event
return
# 先注册订阅再读持久化历史;同一事件循环内没有 await,不会丢失交界事件
# 先注册再异步读取历史;历史与实时队列的交界用 sequence 去重
queue: asyncio.Queue[AgentEvent] = asyncio.Queue()
record.subscribers.add(queue)
history = self.trace_repository.list_events(
run_id, after_sequence=after_sequence
)
last_sequence = after_sequence
try:
history = await asyncio.to_thread(self.trace_repository.list_events,
run_id, after_sequence=after_sequence)
for event in history:
last_sequence = event.sequence
yield event
if event.event in {AgentEventType.run_completed, AgentEventType.run_failed, AgentEventType.run_cancelled}:
return
if run.status in TERMINAL_STATUSES:
return
while True:
@@ -232,7 +244,7 @@ class AgentRuntime:
await asyncio.shield(record.task)
except asyncio.CancelledError:
pass
return record.run.model_copy(deep=True)
return (record.persisted_run or record.run).model_copy(deep=True)
def get_trace(
self, run_id: str, *, after_sequence: int, limit: int
@@ -246,23 +258,35 @@ class AgentRuntime:
return trace
async def _execute(self, record: RunRecord) -> None:
token = agent_run_id.set(record.run.run_id)
try:
async with asyncio.timeout(record.request.run_timeout_seconds):
await self._run_loop(record)
except asyncio.CancelledError:
if record.run.status != AgentRunStatus.cancelled:
self._finish_cancelled(record)
await self._finish_cancelled(record)
except TimeoutError:
self._fail(record, "AGENT_TIMEOUT", "Agent run exceeded its timeout.")
await self._fail(record, "AGENT_TIMEOUT", "Agent run exceeded its timeout.")
except ProviderError as exc:
self._fail(record, exc.code, exc.message)
await self._fail(record, exc.code, exc.message)
except Exception as exc:
self._fail(record, "AGENT_FAILED", str(exc))
log_event('agent', 'execution.failed', level='ERROR', error=exc, run_id=record.run.run_id)
await self._fail(record, "AGENT_FAILED", str(exc))
finally:
self.permissions.cancel_run(record.run.run_id)
agent_run_id.reset(token)
async def shutdown(self) -> None:
results = await asyncio.gather(*(self.cancel(run_id) for run_id in list(self._records)), return_exceptions=True)
for result in results:
if isinstance(result, BaseException):
log_event('agent', 'shutdown.failed', level='ERROR', error=result)
await self._writer.queue.join()
async def _run_loop(self, record: RunRecord) -> None:
record.run.status = AgentRunStatus.running
record.run.updated_at = datetime.now(timezone.utc)
self._publish(
await self._publish(
record,
AgentEventType.run_started,
{"provider_id": record.request.provider_id, "model": record.request.model},
@@ -277,7 +301,7 @@ class AgentRuntime:
record.run.updated_at = datetime.now(timezone.utc)
model_call_id = f"model_call_{uuid4().hex}"
started_at = perf_counter()
self._publish(
await self._publish(
record,
AgentEventType.model_call_started,
{
@@ -299,7 +323,7 @@ class AgentRuntime:
)
)
except Exception as exc:
self._publish(
await self._publish(
record,
AgentEventType.model_call_failed,
{
@@ -309,7 +333,7 @@ class AgentRuntime:
},
)
raise
self._publish(
await self._publish(
record,
AgentEventType.model_call_completed,
{
@@ -322,7 +346,7 @@ class AgentRuntime:
},
)
record.run.token_usage += turn.input_tokens + turn.output_tokens
self._publish(
await self._publish(
record,
AgentEventType.usage,
{"token_usage": record.run.token_usage},
@@ -331,12 +355,12 @@ class AgentRuntime:
record.request.token_budget is not None
and record.run.token_usage > record.request.token_budget
):
self._fail(record, "TOKEN_BUDGET_EXCEEDED", "Agent token budget exceeded.")
await self._fail(record, "TOKEN_BUDGET_EXCEEDED", "Agent token budget exceeded.")
return
if turn.tool_calls:
if len(turn.tool_calls) > MAX_TOOL_CALLS_PER_TURN:
self._fail(
await self._fail(
record,
"TOO_MANY_TOOL_CALLS",
f"Provider requested more than {MAX_TOOL_CALLS_PER_TURN} tools in one turn.",
@@ -360,10 +384,17 @@ class AgentRuntime:
async with semaphore:
return await self._execute_tool(record, call, model_call_id)
results = await asyncio.gather(*(execute(call) for call in calls))
executions = [asyncio.create_task(execute(call)) for call in calls]
try:
results = await asyncio.gather(*executions)
finally:
for execution in executions:
if not execution.done():
execution.cancel()
await asyncio.gather(*executions, return_exceptions=True)
for call, result in zip(calls, results):
record.run.tool_results.append(result)
self._collect_citations(record, result)
await self._collect_citations(record, result)
messages.append(
Message(
role=MessageRole.tool,
@@ -376,20 +407,20 @@ class AgentRuntime:
if turn.text is not None:
record.run.output = turn.text
self._publish(record, AgentEventType.text_delta, {"text": turn.text})
await self._publish(record, AgentEventType.text_delta, {"text": turn.text})
record.run.status = AgentRunStatus.completed
record.run.updated_at = datetime.now(timezone.utc)
self._publish(
await self._publish(
record,
AgentEventType.run_completed,
{"output": turn.text, "token_usage": record.run.token_usage},
)
return
self._fail(record, "EMPTY_MODEL_RESPONSE", "Provider returned no text or tool call.")
await self._fail(record, "EMPTY_MODEL_RESPONSE", "Provider returned no text or tool call.")
return
self._fail(record, "MAX_STEPS_EXCEEDED", "Agent reached its maximum step count.")
await self._fail(record, "MAX_STEPS_EXCEEDED", "Agent reached its maximum step count.")
async def _execute_tool(
self, record: RunRecord, call: ToolCall, parent_model_call_id: str
@@ -397,7 +428,7 @@ class AgentRuntime:
started_at = perf_counter()
call_data = call.model_dump(mode="json")
call_data["parent_model_call_id"] = parent_model_call_id
self._publish(record, AgentEventType.tool_call, call_data)
await self._publish(record, AgentEventType.tool_call, call_data)
try:
registered = self.tools.get(call.name)
except ToolNotFoundError:
@@ -411,7 +442,7 @@ class AgentRuntime:
error_code="TOOL_NOT_ALLOWED",
error_message="Tool is not included in allowed_tools.",
)
self._publish_tool_result(
await self._publish_tool_result(
record, result, parent_model_call_id, started_at
)
return result
@@ -425,7 +456,7 @@ class AgentRuntime:
error_code="NETWORK_NOT_ALLOWED",
error_message="Agent run does not allow network tools.",
)
self._publish_tool_result(
await self._publish_tool_result(
record, result, parent_model_call_id, started_at
)
return result
@@ -436,7 +467,7 @@ class AgentRuntime:
# 运行状态必须在等待期间可见,前端才能展示并处理权限确认卡片。
ticket = self.permissions.create_ticket(record.run.run_id, permission)
record.run.status = AgentRunStatus.waiting_permission
self._publish(
await self._publish(
record,
AgentEventType.permission_required,
{
@@ -458,13 +489,18 @@ class AgentRuntime:
error_code="PERMISSION_TIMEOUT",
error_message="Tool permission confirmation timed out.",
)
self._publish_tool_result(
await self._publish_tool_result(
record, result, parent_model_call_id, started_at
)
return result
record.run.status = AgentRunStatus.running
record.run.updated_at = datetime.now(timezone.utc)
self.trace_repository.save_run(record.run)
async with record.publish_lock:
snapshot = record.run.model_copy(deep=True)
cancelled = await self._writer.submit('save', snapshot)
record.persisted_run = snapshot
if cancelled:
raise asyncio.CancelledError
result = (
await self._invoke_tool(record, call)
if decision in {"allow_once", "allow_session"}
@@ -473,10 +509,10 @@ class AgentRuntime:
else:
result = await self._invoke_tool(record, call)
self._publish_tool_result(record, result, parent_model_call_id, started_at)
await self._publish_tool_result(record, result, parent_model_call_id, started_at)
return result
def _publish_tool_result(
async def _publish_tool_result(
self,
record: RunRecord,
result: ToolResult,
@@ -486,7 +522,7 @@ class AgentRuntime:
data = result.model_dump(mode="json")
data["parent_model_call_id"] = parent_model_call_id
data["duration_ms"] = int((perf_counter() - started_at) * 1000)
self._publish(record, AgentEventType.tool_result, data)
await self._publish(record, AgentEventType.tool_result, data)
async def _invoke_tool(self, record: RunRecord, call: ToolCall) -> ToolResult:
try:
@@ -519,45 +555,60 @@ class AgentRuntime:
error_message="Tool permission was denied.",
)
def _finish_cancelled(self, record: RunRecord) -> None:
async def _finish_cancelled(self, record: RunRecord) -> None:
record.run.cancelled = True
record.run.status = AgentRunStatus.cancelled
record.run.updated_at = datetime.now(timezone.utc)
self._publish(record, AgentEventType.run_cancelled, {})
await self._publish(record, AgentEventType.run_cancelled, {})
def _fail(self, record: RunRecord, code: str, message: str) -> None:
if record.run.status in TERMINAL_STATUSES:
async def _fail(self, record: RunRecord, code: str, message: str) -> None:
log_event('agent', 'run.error', level='ERROR', run_id=record.run.run_id, error_code=code)
if record.persisted_run and record.persisted_run.status in TERMINAL_STATUSES:
return
record.run.status = AgentRunStatus.failed
record.run.error_code = code
record.run.error_message = message
record.run.updated_at = datetime.now(timezone.utc)
self._publish(
await self._publish(
record,
AgentEventType.run_failed,
{"code": code, "message": message},
)
def _publish(
async def _publish(
self, record: RunRecord, event_type: AgentEventType, data: dict[str, object]
) -> None:
sanitized = sanitize_trace_value(data)
assert isinstance(sanitized, dict)
event = AgentEvent(
event=event_type,
run_id=record.run.run_id,
sequence=record.next_sequence,
data=sanitized,
timestamp=datetime.now(timezone.utc),
)
record.next_sequence += 1
record.events.append(event)
self.trace_repository.append_event(record.run, event)
# 内存只保留实时订阅窗口;完整审计轨迹由 SQLite 保存。
if len(record.events) > MAX_EVENTS_PER_RUN:
del record.events[: len(record.events) - MAX_EVENTS_PER_RUN]
for queue in record.subscribers:
queue.put_nowait(event)
async with record.publish_lock:
sanitized = sanitize_trace_value(data)
assert isinstance(sanitized, dict)
event = AgentEvent(
event=event_type,
run_id=record.run.run_id,
sequence=record.next_sequence,
data=sanitized,
timestamp=datetime.now(timezone.utc),
)
snapshot = record.run.model_copy(deep=True)
try:
cancelled = await self._writer.submit('event', snapshot, event)
except Exception as exc:
log_event('agent', 'trace.write_failed', level='ERROR', error=exc, run_id=record.run.run_id)
raise
record.next_sequence += 1
record.persisted_run = snapshot
record.events.append(event)
log_event('agent', event_type.value,
level='ERROR' if event_type.value.endswith('Failed') or data.get('success') is False else 'INFO',
run_id=record.run.run_id, provider_id=record.run.provider_id, model=record.run.model,
sequence=event.sequence, step=record.run.current_step, status=snapshot.status.value,
tool=data.get('name'), error_code=data.get('code') or data.get('error_code'))
# 内存只保留实时订阅窗口;完整审计轨迹由 SQLite 保存。
if len(record.events) > MAX_EVENTS_PER_RUN:
del record.events[: len(record.events) - MAX_EVENTS_PER_RUN]
for queue in record.subscribers:
queue.put_nowait(event)
if cancelled:
raise asyncio.CancelledError
@staticmethod
def _request_metadata(record: RunRecord) -> dict[str, object]:
@@ -583,7 +634,7 @@ class AgentRuntime:
"metadata": record.request.metadata,
}
def _collect_citations(self, record: RunRecord, result: ToolResult) -> None:
async def _collect_citations(self, record: RunRecord, result: ToolResult) -> None:
if not result.success or not isinstance(result.output, dict):
return
items = result.output.get("items")
@@ -601,7 +652,7 @@ class AgentRuntime:
continue
known.add(citation.citation_id)
record.run.citations.append(citation)
self._publish(record, AgentEventType.citation, citation.model_dump(mode="json"))
await self._publish(record, AgentEventType.citation, citation.model_dump(mode="json"))
def _get_record(self, run_id: str) -> RunRecord:
try:
@@ -618,7 +669,7 @@ class AgentRuntime:
(
record
for record in self._records.values()
if record.run.status in TERMINAL_STATUSES
if record.run.status in TERMINAL_STATUSES and (record.task is None or record.task.done())
),
key=lambda record: record.run.updated_at,
)
+30 -11
View File
@@ -6,6 +6,7 @@ SQLite 中的事件是 SSE、前端 Trace 和 Benchmark 的共同事实来源。
from __future__ import annotations
from contextlib import nullcontext
import json
import re
from datetime import datetime, timezone
@@ -94,15 +95,30 @@ def sanitize_trace_value(
class AgentTraceRepository:
def write_batch(self, jobs):
conn = connect()
try:
with transaction(conn):
for operation, args in jobs:
if operation == 'create':
self.create_run(*args, _conn=conn)
elif operation == 'save':
self.save_run(*args, _conn=conn)
else:
self.append_event(*args, _conn=conn)
finally:
conn.close()
def create_run(
self,
run: AgentRun,
request: AgentRunCreateRequest,
config_snapshot: dict[str, Any],
*, _conn=None,
) -> None:
conn = connect()
conn = _conn or connect()
try:
with transaction(conn):
with transaction(conn) if _conn is None else nullcontext():
conn.execute(
"""
INSERT INTO agent_runs(
@@ -126,22 +142,24 @@ class AgentTraceRepository:
),
)
finally:
conn.close()
if _conn is None:
conn.close()
def save_run(self, run: AgentRun) -> None:
conn = connect()
def save_run(self, run: AgentRun, *, _conn=None) -> None:
conn = _conn or connect()
try:
with transaction(conn):
with transaction(conn) if _conn is None else nullcontext():
self._update_run(conn, run)
finally:
conn.close()
if _conn is None:
conn.close()
def append_event(self, run: AgentRun, event: AgentEvent) -> None:
def append_event(self, run: AgentRun, event: AgentEvent, *, _conn=None) -> None:
"""在同一事务中保存最新 Run 和事件;复写同一序号时保持幂等。"""
conn = connect()
conn = _conn or connect()
try:
with transaction(conn):
with transaction(conn) if _conn is None else nullcontext():
self._update_run(conn, run)
conn.execute(
"""
@@ -158,7 +176,8 @@ class AgentTraceRepository:
),
)
finally:
conn.close()
if _conn is None:
conn.close()
def get_run(self, run_id: str) -> AgentRun | None:
conn = connect()
+5
View File
@@ -1132,6 +1132,11 @@ class TranscriptNoteRequest(Contract):
class IndexStatus(Contract):
running_jobs: int = 0
active_searches: int = 0
completed_searches: int = 0
failed_searches: int = 0
cancelled_searches: int = 0
vector_refresh_required: bool = False
total_notes: int = 0
total_blocks: int = 0
+4
View File
@@ -25,6 +25,10 @@ class ApiError(Exception):
async def api_error_handler(_: Request, exc: ApiError) -> JSONResponse:
from app.operation_logs import log_event
log_event('api', 'operation.failed', level='ERROR' if exc.status_code >= 500 else 'WARNING',
error=exc, status=exc.status_code,
**{key: value for key, value in exc.details.items() if key in {'run_id', 'task_id', 'note_id', 'job_id', 'provider_id'}})
body = ErrorResponse(
error=ErrorDetail(code=exc.code, message=exc.message, details=exc.details)
)
+11
View File
@@ -0,0 +1,11 @@
"""Bound embedding result frames so large notes do not exceed pipe line limits."""
import json
def response_lines(response, operation):
if operation == 'embedding' and 'result' in response and 'error_code' not in response:
vectors = response['result']
for offset in range(0, len(vectors), 128):
yield json.dumps({'embedding_offset': offset, 'embedding_chunk': vectors[offset:offset + 128]}, allow_nan=False) + '\n'
response = {**response, 'result': [], 'embedding_count': len(vectors)}
yield json.dumps(response, ensure_ascii=False, allow_nan=False) + '\n'
+16 -1
View File
@@ -189,15 +189,30 @@ class Runtime:
await process.stdin.drain()
process.stdin.close()
final = None
vectors = []
while line := await process.stdout.readline():
message = json.loads(line)
if "progress" in message:
if "embedding_chunk" in message:
chunk = message['embedding_chunk']
if (operation != 'embedding' or not isinstance(chunk, list)
or message.get('embedding_offset') != len(vectors)
or len(vectors) + len(chunk) > len(payload.get('texts', []))):
raise ProviderError('LOCAL_MODEL_INVALID_RESPONSE', '本地向量传输顺序或数量无效。')
vectors.extend(chunk)
elif "progress" in message:
callback = runtime_progress.get()
if callback:
callback(message)
else:
final = message
await process.wait()
if isinstance(final, dict) and 'embedding_count' in final:
if (final['embedding_count'] != len(vectors)
or len(vectors) != len(payload.get('texts', []))):
raise ProviderError('LOCAL_MODEL_INVALID_RESPONSE', '本地向量传输不完整。')
final['result'] = vectors
elif vectors:
raise ProviderError('LOCAL_MODEL_INVALID_RESPONSE', '本地向量传输缺少结束标记。')
return final
try:
result = await asyncio.wait_for(receive(), config.timeout_seconds)
+3 -1
View File
@@ -216,4 +216,6 @@ if __name__ == "__main__":
response = {"error_code": "LOCAL_INFERENCE_FAILED", "message": "本地推理失败,请检查媒体格式、模型和设备配置。"}
if "error_code" in response:
response["diagnostics"] = {"requested_device": request["config"]["device"], "actual_device": request.get("_actual_device", "unknown")}
sys.stdout.buffer.write((json.dumps(response, ensure_ascii=False, allow_nan=False) + "\n").encode("utf-8"))
from protocol import response_lines
for line in response_lines(response, request['operation']):
sys.stdout.buffer.write(line.encode('utf-8'))
+11
View File
@@ -0,0 +1,11 @@
from fastapi import APIRouter, Query
from app.operation_logs import get_store
router = APIRouter(prefix='/api/logs', tags=['Diagnostics'])
@router.get('')
def list_logs(limit: int = Query(50, ge=1, le=200), before: int | None = Query(None, ge=1),
level: str = Query('', pattern='^(|INFO|WARNING|ERROR|CRITICAL)$'),
source: str = Query('', max_length=100), q: str = Query('', max_length=200)):
return get_store().query(limit=limit, before=before, level=level, source=source, q=q)
+36
View File
@@ -1,4 +1,7 @@
from contextlib import asynccontextmanager
import asyncio
from time import perf_counter
from uuid import uuid4
from fastapi import FastAPI
from fastapi.exceptions import RequestValidationError
@@ -14,17 +17,22 @@ from app.local_model_routes import router as local_model_router
from app.usage_routes import router as usage_router
from app.provider_preview_routes import router as provider_preview_router
from app.schemas import HealthResponse, ServiceStatusResponse
from app.log_routes import router as log_router
from app.operation_logs import install_logging, log_event, request_id, shutdown_logging
settings = get_settings()
@asynccontextmanager
async def lifespan(_: FastAPI):
install_logging()
log_event('system', 'service.started')
from app.services import transcription_service
transcription_service.recover_interrupted()
try:
yield
finally:
await container.agent.shutdown()
from app.services import index_service
await index_service.shutdown()
await transcription_service.shutdown()
@@ -35,6 +43,8 @@ async def lifespan(_: FastAPI):
await manager.cancel_download(key)
container.plugins.shutdown()
container.mcp_servers.shutdown()
log_event('system', 'service.stopped')
await asyncio.to_thread(shutdown_logging)
app = FastAPI(
@@ -60,6 +70,32 @@ app.include_router(media_router)
app.include_router(local_model_router)
app.include_router(usage_router)
app.include_router(provider_preview_router)
app.include_router(log_router)
@app.middleware('http')
async def operation_log(request, call_next):
token = request_id.set(uuid4().hex)
started = perf_counter()
status = 500
failure = None
try:
response = await call_next(request)
status = response.status_code
response.headers['X-Request-ID'] = request_id.get()
return response
except Exception as exc:
failure = exc
raise
finally:
# Do not record query strings, request/response bodies or arbitrary URLs.
route = getattr(request.scope.get('route'), 'path', 'unmatched')
if not route.startswith('/api/logs') and (request.method not in {'GET', 'HEAD', 'OPTIONS'} or status >= 400 or perf_counter() - started > 1):
log_event('http', 'request.finished', level='ERROR' if status >= 500 else 'WARNING' if status >= 400 else 'INFO',
error=failure, method=request.method, route=route, status=status,
duration_ms=round((perf_counter() - started) * 1000, 2),
**{k: v for k, v in request.path_params.items() if k in {'run_id', 'task_id', 'note_id', 'job_id', 'provider_id'}})
request_id.reset(token)
@app.get("/health", response_model=HealthResponse, tags=["System"])
+186
View File
@@ -0,0 +1,186 @@
"""Bounded, asynchronous operational diagnostics, separate from business/Trace data.
Only explicitly allowed metadata is stored. Never store prompts, tool arguments,
provider response bodies or raw exception messages in this diagnostic channel.
"""
from __future__ import annotations
import json
import logging
import math
import queue
import re
import sqlite3
import threading
import traceback
from contextvars import ContextVar
from contextlib import closing
from datetime import datetime, timezone
from pathlib import Path
from app.config import get_settings
request_id: ContextVar[str] = ContextVar('log_request_id', default='')
agent_run_id: ContextVar[str] = ContextVar('log_agent_run_id', default='')
_allowed = {'run_id', 'task_id', 'note_id', 'job_id', 'provider_id', 'model',
'device', 'error_code', 'error_type', 'status', 'duration_ms', 'count',
'step', 'sequence', 'tool', 'method', 'route', 'request_id', 'fallback',
'frames', 'source', 'changed_fields'}
_safe = re.compile(r'[^\w .:/@{}\[\],()=+\-]', re.UNICODE)
def metadata(values: dict) -> dict:
result = {}
for key, value in values.items():
if key not in _allowed or value is None:
continue
if isinstance(value, (int, float, bool)):
if not isinstance(value, float) or math.isfinite(value):
result[key] = value
else:
text = str(value)
text = re.sub(r'(?i)(?:bearer\s+\S+|sk-[\w-]+)', '[REDACTED]', text)
result[key] = _safe.sub('', text)[:500]
return result
class LogStore:
def __init__(self, path: Path, *, retain: int = 20_000):
self.path = path
self.retain = retain
self.queue: queue.Queue = queue.Queue(maxsize=4096)
self.dropped = 0
self.failed = 0
self.closed = False
self.state_lock = threading.Lock()
self.thread = threading.Thread(target=self._write, name='operation-logs', daemon=True)
path.parent.mkdir(parents=True, exist_ok=True)
with closing(self._connect()) as conn, conn:
conn.execute('CREATE TABLE IF NOT EXISTS logs (id INTEGER PRIMARY KEY, timestamp TEXT NOT NULL, level TEXT NOT NULL, source TEXT NOT NULL, event TEXT NOT NULL, details TEXT NOT NULL)')
conn.execute('CREATE INDEX IF NOT EXISTS logs_level_id ON logs(level, id)')
conn.execute('CREATE INDEX IF NOT EXISTS logs_source_id ON logs(source, id)')
self.thread.start()
def _connect(self):
conn = sqlite3.connect(self.path, timeout=5)
conn.row_factory = sqlite3.Row
return conn
def emit(self, level: str, source: str, event: str, details: dict):
row = (datetime.now(timezone.utc).isoformat(), level, source[:100], event[:160], json.dumps(metadata(details), ensure_ascii=False))
with self.state_lock:
if self.closed:
return
try:
self.queue.put_nowait(row)
except queue.Full:
self.dropped += 1
def _write(self):
while True:
first = self.queue.get()
batch = [first]
while len(batch) < 128:
try:
batch.append(self.queue.get_nowait())
except queue.Empty:
break
stop = None in batch
rows = [row for row in batch if row is not None]
try:
if rows:
with closing(self._connect()) as conn, conn:
conn.executemany('INSERT INTO logs(timestamp,level,source,event,details) VALUES(?,?,?,?,?)', rows)
conn.execute('DELETE FROM logs WHERE id <= (SELECT id FROM logs ORDER BY id DESC LIMIT 1 OFFSET ?)', (self.retain,))
except Exception:
self.failed += len(rows)
finally:
for _ in batch:
self.queue.task_done()
if stop:
return
def query(self, *, limit=50, before=None, level='', source='', q=''):
clauses, args = [], []
for column, value in [('level', level), ('source', source)]:
if value:
clauses.append(f'{column} = ?')
args.append(value)
if before is not None:
clauses.append('id < ?')
args.append(before)
if q:
clauses.append('(instr(event, ?) > 0 OR instr(details, ?) > 0)')
args += [q, q]
where = ' WHERE ' + ' AND '.join(clauses) if clauses else ''
with closing(self._connect()) as conn, conn:
rows = conn.execute('SELECT * FROM logs' + where + ' ORDER BY id DESC LIMIT ?', (*args, limit + 1)).fetchall()
sources = [row[0] for row in conn.execute('SELECT DISTINCT source FROM logs ORDER BY source')]
items = [{**dict(row), 'details': json.loads(row['details'])} for row in rows[:limit]]
return {'items': items, 'next_cursor': items[-1]['id'] if len(rows) > limit else None,
'sources': sources, 'pending': self.queue.qsize(), 'dropped': self.dropped,
'write_failures': self.failed, 'retention': self.retain}
def close(self):
with self.state_lock:
if self.closed:
return
self.closed = True
self.queue.put(None)
self.thread.join(timeout=15)
_store: LogStore | None = None
_lock = threading.Lock()
def get_store() -> LogStore:
global _store
path = get_settings().data_dir / 'logs' / 'operations.sqlite3'
with _lock:
if _store is None or _store.path != path or _store.closed:
if _store is not None and not _store.closed:
_store.close()
_store = LogStore(path)
return _store
def log_event(module: str, event: str, *, level='INFO', error: BaseException | None = None, **details):
if request_id.get():
details.setdefault('request_id', request_id.get())
if agent_run_id.get():
details.setdefault('run_id', agent_run_id.get())
if error:
details['error_type'] = type(error).__name__
details.setdefault('error_code', getattr(error, 'code', None))
details['frames'] = '; '.join(f'{Path(f.filename).name}:{f.lineno}:{f.name}' for f in traceback.extract_tb(error.__traceback__)[-8:])
try:
get_store().emit(level, module, event, details)
except Exception:
# Logging must not turn a successful save/run into a business failure.
logging.getLogger('operation_log_storage').error('Operational log storage unavailable')
class ApplicationLogHandler(logging.Handler):
def emit(self, record):
if record.name == 'operation_log_storage' or getattr(record, '_notes_operation_logged', False):
return
record._notes_operation_logged = True
# Legacy log messages can include note text/credentials, even in f-strings.
# Preserve source location and error class; structured call sites carry IDs.
log_event(record.name, 'application.warning' if record.levelno < 40 else 'application.error',
level=record.levelname, error=record.exc_info[1] if record.exc_info else None,
frames=f'{Path(record.pathname).name}:{record.lineno}:{record.funcName}')
def install_logging():
# Uvicorn's default logger stops propagation before the root logger.
for name in ('', 'uvicorn'):
logger = logging.getLogger(name)
if not any(isinstance(h, ApplicationLogHandler) for h in logger.handlers):
logger.addHandler(ApplicationLogHandler(level=logging.WARNING))
def shutdown_logging():
if _store is not None and not _store.closed:
_store.close()
+30
View File
@@ -0,0 +1,30 @@
"""Process-local retrieval activity, shared by search, RAG and Agent callers."""
import asyncio
from functools import wraps
active = 0
completed = 0
failed = 0
cancelled = 0
def track_search(operation):
@wraps(operation)
async def wrapped(self, request):
global active, completed, failed, cancelled
if request.mode == 'fts':
return await operation(self, request)
active += 1
try:
result = await operation(self, request)
completed += 1
return result
except asyncio.CancelledError:
cancelled += 1
raise
except Exception:
failed += 1
raise
finally:
active -= 1
return wrapped
+2
View File
@@ -10,6 +10,7 @@ from __future__ import annotations
from datetime import datetime, timezone
from app import repository
from app.retrieval.activity import track_search
from app.contracts import (
Citation,
PageMeta,
@@ -52,6 +53,7 @@ class RetrievalEngine:
# remain authoritative, including monkeypatches on the singleton.
self._routed_defaults = (embedding, vector_store) if route_embeddings else None
@track_search
async def search(self, request: SearchRequest) -> SearchResponse:
if request.mode == SearchMode.fts:
return self._search_fts(request)
+3
View File
@@ -20,6 +20,7 @@ from typing import Protocol
from app.database.db import connect, transaction
from app.errors import ApiError
from app.operation_logs import log_event
from app.retrieval.vectorstore import VectorHit
from app.retrieval.provenance import record_embedding
from app.retrieval.hybrid import rrf_fuse
@@ -101,6 +102,8 @@ async def embed_remote(texts: list[str], *, accept_local=False, strict=False, lo
source=result.source,
)
except Exception as exc:
log_event('vectors', 'embedding.failed', level='ERROR' if strict else 'WARNING', error=exc,
count=len(texts), fallback='none' if strict else 'local_index')
# Avoid logging provider exceptions containing credentials or note text.
record_embedding(fallback_reason="REMOTE_EMBEDDING_UNAVAILABLE")
logger.warning("Remote embedding unavailable (%s); using local index", type(exc).__name__)
+13 -8
View File
@@ -11,6 +11,7 @@ from fastapi.responses import StreamingResponse
from app.agent import AgentCapacityError, AgentRunNotFoundError
from app.container import container
from app.config import get_settings
from app.operation_logs import log_event
from app.extensions.archive import MAX_ZIP_BYTES, install_zip
from app.services.persona_settings import PersonaSettings, load_persona, save_persona
from app.contracts import (
@@ -455,11 +456,15 @@ async def chat(request: ChatRequest) -> StreamingResponse:
usage = {"input_tokens": input_tokens, "output_tokens": output_tokens,
"total_tokens": input_tokens + output_tokens}
elif event.event == ModelEventType.error:
log_event('chat', 'model.error', level='ERROR', provider_id=request.provider_id,
model=request.model, error_code=event.data.get('code'))
if assistant_content:
assistant_content += "\n\n"
assistant_content += str(event.data.get("message", "Model generation failed."))
yield as_sse(event.event.value, event.model_dump_json())
except Exception as exc:
log_event('chat', 'chat.failed', level='ERROR', error=exc,
provider_id=request.provider_id, model=request.model)
failure_message = exc.message if isinstance(exc, ApiError) else "知识库检索或模型生成失败,请检查服务状态。"
if assistant_content:
assistant_content += "\n\n"
@@ -498,7 +503,7 @@ async def chat(request: ChatRequest) -> StreamingResponse:
async def list_agent_runs(
limit: int = Query(default=50, ge=1, le=100), offset: int = Query(default=0, ge=0)
) -> AgentRunListResponse:
items, total = container.agent.list_runs(limit=limit, offset=offset)
items, total = await asyncio.to_thread(container.agent.list_runs, limit=limit, offset=offset)
return AgentRunListResponse(
items=items,
page=PageMeta(total=total, limit=limit, offset=offset),
@@ -603,7 +608,7 @@ async def get_agent_trace(
limit: int = Query(default=200, ge=1, le=500),
) -> AgentTraceResponse:
try:
return container.agent.get_trace(
return await asyncio.to_thread(container.agent.get_trace,
run_id, after_sequence=after_sequence, limit=limit
)
except AgentRunNotFoundError as exc:
@@ -624,7 +629,7 @@ async def decide_agent_permission(
run_id: str, request_id: str, request: PermissionDecisionRequest
) -> OperationResponse:
agent_run_or_404(run_id)
if not container.agent.resolve_permission(run_id, request_id, request.decision):
if not await container.agent.resolve_permission(run_id, request_id, request.decision):
raise ApiError(
404,
"PERMISSION_REQUEST_NOT_FOUND",
@@ -1223,7 +1228,7 @@ async def test_provider(request: ProviderTestRequest) -> ProviderTestResponse:
async def list_tasks(
limit: int = Query(default=50, ge=1, le=100), offset: int = Query(default=0, ge=0)
) -> TaskListResponse:
items, total = task_service.list_tasks(limit=limit, offset=offset)
items, total = await asyncio.to_thread(task_service.list_tasks, limit=limit, offset=offset)
return TaskListResponse(
items=items, page=PageMeta(total=total, limit=limit, offset=offset)
)
@@ -1231,12 +1236,12 @@ async def list_tasks(
@router.post("/tasks", response_model=Task, tags=["Tasks"])
async def create_task(request: TaskCreateRequest) -> Task:
return task_service.create_task(**request.model_dump())
return await task_service.write_in_background(task_service.create_task, **request.model_dump())
@router.get("/tasks/{task_id}", response_model=Task, tags=["Tasks"])
async def get_task(task_id: str) -> Task:
task = task_service.get_task(task_id)
task = await asyncio.to_thread(task_service.get_task, task_id)
if task is None:
raise ApiError(
404, "RESOURCE_NOT_FOUND", "task not found", {"task_id": task_id}
@@ -1246,7 +1251,7 @@ async def get_task(task_id: str) -> Task:
@router.patch("/tasks/{task_id}", response_model=Task, tags=["Tasks"])
async def update_task(task_id: str, request: TaskUpdateRequest) -> Task:
return task_service.update_task(task_id, request.model_dump(exclude_unset=True))
return await task_service.write_in_background(task_service.update_task, task_id, request.model_dump(exclude_unset=True))
@router.delete(
@@ -1255,7 +1260,7 @@ async def update_task(task_id: str, request: TaskUpdateRequest) -> Task:
tags=["Tasks"],
)
async def delete_task(task_id: str) -> OperationResponse:
if not task_service.delete_task(task_id):
if not await task_service.write_in_background(task_service.delete_task, task_id):
raise ApiError(
404, "RESOURCE_NOT_FOUND", "task not found", {"task_id": task_id}
)
+24 -5
View File
@@ -4,6 +4,7 @@ from __future__ import annotations
import asyncio
import logging
from app.operation_logs import log_event
from datetime import datetime, timezone
from pathlib import Path
@@ -25,6 +26,7 @@ vector_store = SqliteVecStore()
_jobs: dict[str, IndexJob] = {}
_active_job_id: str | None = None
_active_scope: str | None = None
_last_completed_at: datetime | None = None
_last_error: str | None = None
MAX_JOBS = 100
@@ -33,6 +35,7 @@ _logger = logging.getLogger(__name__)
def _remember_job(job: IndexJob) -> None:
log_event('vectors', 'index.' + job.status, job_id=job.job_id, status=job.status)
_jobs[job.job_id] = job
while len(_jobs) > MAX_JOBS:
oldest = next(iter(_jobs))
@@ -64,7 +67,7 @@ def _scan_vault() -> list[tuple[str, str, str, datetime, datetime]]:
async def rebuild(request: IndexRebuildRequest) -> IndexJob:
global _active_job_id, _last_completed_at, _last_error
global _active_job_id, _active_scope, _last_completed_at, _last_error
if _active_job_id is not None:
raise ApiError(409, "INDEX_BUSY", "索引正在后台计算,请稍后重试。")
job_id = "job_" + uuid4().hex[:12]
@@ -82,6 +85,7 @@ async def rebuild(request: IndexRebuildRequest) -> IndexJob:
saved_paths = {record.file_path: record for record in saved_records.values() if record is not None}
_active_job_id = job_id
_active_scope = 'all'
_last_error = None
_remember_job(IndexJob(
job_id=job_id, status="running", scope=request.scope,
@@ -148,6 +152,7 @@ async def rebuild(request: IndexRebuildRequest) -> IndexJob:
finally:
conn.close()
except BaseException as exc:
log_event('vectors', 'index.failed', level='WARNING' if isinstance(exc, asyncio.CancelledError) else 'ERROR', error=exc, job_id=job_id)
_remember_job(IndexJob(
job_id=job_id, status="failed", scope=request.scope,
created_at=datetime.now(timezone.utc),
@@ -156,6 +161,7 @@ async def rebuild(request: IndexRebuildRequest) -> IndexJob:
raise
finally:
_active_job_id = None
_active_scope = None
job = IndexJob(job_id=job_id, status="completed", scope=request.scope, created_at=datetime.now(timezone.utc))
_remember_job(job)
@@ -166,16 +172,26 @@ async def rebuild(request: IndexRebuildRequest) -> IndexJob:
def get_status() -> IndexStatus:
from app.retrieval import activity
counts = repository.stats()
vector_refresh_required = repository.get_index_meta().get('workspace_vectors_pending') == '1' or bool(_pending_notes())
workspace_pending = repository.get_index_meta().get('workspace_vectors_pending') == '1'
notes_pending = len(_pending_notes())
vector_refresh_required = workspace_pending or bool(notes_pending)
running = int(_active_job_id is not None)
# An entire-vault rebuild is one job, not one job per block/note.
pending = 1 if running and _active_scope == 'all' else (1 + running if workspace_pending else max(notes_pending, running))
activity_fields = dict(running_jobs=running, active_searches=activity.active,
completed_searches=activity.completed, failed_searches=activity.failed,
cancelled_searches=activity.cancelled)
if _active_job_id is not None:
return IndexStatus(status="running", pending_jobs=0, active_job_id=_active_job_id, vector_refresh_required=vector_refresh_required,
return IndexStatus(**activity_fields, status="running", pending_jobs=pending, active_job_id=_active_job_id, vector_refresh_required=vector_refresh_required,
total_notes=counts["notes"], total_blocks=counts["blocks"])
return IndexStatus(
**activity_fields,
vector_refresh_required=vector_refresh_required,
total_notes=counts["notes"], total_blocks=counts["blocks"],
status="failed" if _last_error else "idle",
pending_jobs=0,
pending_jobs=pending,
last_completed_at=_last_completed_at,
error_message=_last_error,
)
@@ -227,7 +243,7 @@ def _pending_notes() -> list[str]:
async def _refresh_saved_note(note_id: str) -> None:
global _active_job_id, _last_error, _last_completed_at
global _active_job_id, _active_scope, _last_error, _last_completed_at
record = repository.get_note_record(note_id)
key = f'note_vectors_pending:{note_id}'
if record is None:
@@ -240,6 +256,7 @@ async def _refresh_saved_note(note_id: str) -> None:
parsed.title = record.title
job_id = 'job_' + uuid4().hex[:12]
_active_job_id = job_id
_active_scope = 'note'
_last_error = None
_remember_job(IndexJob(job_id=job_id, status='running', scope='all', created_at=datetime.now(timezone.utc)))
try:
@@ -267,8 +284,10 @@ async def _refresh_saved_note(note_id: str) -> None:
_last_completed_at = datetime.now(timezone.utc)
_remember_job(IndexJob(job_id=job_id, status='completed', scope='all', created_at=_last_completed_at))
except BaseException as exc:
log_event('vectors', 'index.failed', level='WARNING' if isinstance(exc, asyncio.CancelledError) else 'ERROR', error=exc, job_id=job_id)
_last_error = str(exc) or '后台向量计算已中断,笔记已保存。'
_remember_job(IndexJob(job_id=job_id, status='failed', scope='all', created_at=datetime.now(timezone.utc)))
raise
finally:
_active_job_id = None
_active_scope = None
@@ -19,6 +19,13 @@ def connection():
def record(**values):
from app.operation_logs import log_event
log_event('models', 'model.' + str(values.get('operation', 'inference')),
level='ERROR' if values.get('status') == 'failed' else 'WARNING' if values.get('status') == 'fallback' else 'INFO',
model=values.get('model'), source=values.get('source'), status=values.get('status'),
device=values.get('actual_device') or values.get('attempted_device'),
error_code=values.get('error_code'), fallback=values.get('fallback_reason'),
duration_ms=round(values.get('elapsed_seconds', 0) * 1000, 2))
safe = {key: value[:240] for key, value in values.items() if key in TEXT and isinstance(value, str)}
safe.update({key: value for key, value in values.items()
if key in NUMBERS and type(value) in (float, int) and math.isfinite(value) and value >= 0})
+29
View File
@@ -2,11 +2,37 @@ from __future__ import annotations
from datetime import datetime, timezone
from uuid import uuid4
import asyncio
from contextvars import copy_context
from functools import partial
from weakref import WeakKeyDictionary
from app import repository
from app.contracts import Task, TaskStatus
from app.database.db import connect, transaction
from app.errors import ApiError
from app.operation_logs import log_event
_write_locks = WeakKeyDictionary()
async def write_in_background(operation, *args, **kwargs):
# SQLite has one writer. Queue cooperatively instead of letting many worker
# threads fight over the file lock and starve unrelated model work.
loop = asyncio.get_running_loop()
lock = _write_locks.setdefault(loop, asyncio.Lock())
async with lock:
work = loop.run_in_executor(None, copy_context().run, partial(operation, *args, **kwargs))
cancelled = False
while not work.done():
try:
await asyncio.shield(work)
except asyncio.CancelledError:
cancelled = True
result = work.result()
if cancelled:
raise asyncio.CancelledError
return result
def _now() -> datetime:
@@ -49,6 +75,7 @@ def create_task(
),
)
row = conn.execute("SELECT * FROM tasks WHERE task_id = ?", (task_id,)).fetchone()
log_event('tasks', 'task.created', task_id=task_id, note_id=note_id, status='todo')
return _task_from_row(row)
finally:
conn.close()
@@ -111,6 +138,7 @@ def update_task(task_id: str, values: dict[str, object]) -> Task:
params,
)
row = conn.execute("SELECT * FROM tasks WHERE task_id = ?", (task_id,)).fetchone()
log_event('tasks', 'task.updated', task_id=task_id, status=row['status'], changed_fields=','.join(values))
return _task_from_row(row)
finally:
conn.close()
@@ -121,6 +149,7 @@ def delete_task(task_id: str) -> bool:
try:
with transaction(conn):
cursor = conn.execute("DELETE FROM tasks WHERE task_id = ?", (task_id,))
log_event('tasks', 'task.deleted' if cursor.rowcount else 'task.not_found', task_id=task_id)
return cursor.rowcount > 0
finally:
conn.close()
+5
View File
@@ -98,6 +98,11 @@ class UsageAttempt:
reasoning_tokens=first("output_tokens_details.reasoning_tokens", "completion_tokens_details.reasoning_tokens"))
def persist(self):
from app.operation_logs import log_event
log_event('providers', 'model.request_finished', level='INFO' if self.completed else 'WARNING',
provider_id=self.provider_id, model=self.model, run_id=self.run_id,
request_id=self.request_id, source=self.source,
status='completed' if self.completed else 'incomplete')
try:
with closing(connection()) as conn:
conn.execute("INSERT OR REPLACE INTO model_usage VALUES (?,?,?,?,?,?,?,?,?,?,?)", (
+250
View File
@@ -0,0 +1,250 @@
"""Offline Agent/runtime and task API load test; all state lives in a temporary directory."""
from __future__ import annotations
import argparse
import asyncio
import json
import inspect
import math
import os
import pathlib
import platform
import sys
import tempfile
import time
sys.path.insert(0, str(pathlib.Path(__file__).resolve().parents[1]))
def stats(values):
values = sorted(values)
return {"count": len(values), "median_ms": round(values[len(values)//2], 2),
"p95_ms": round(values[min(len(values)-1, math.ceil(len(values)*.95)-1)], 2),
"max_ms": round(values[-1], 2)} if values else {"count": 0}
async def main(output):
# Set before importing any app modules: container has import-time initialization.
with tempfile.TemporaryDirectory(prefix="notes-agent-task-stress-") as directory:
root = pathlib.Path(directory)
os.environ.update(APP_DATA_DIR=str(root / 'data'), APP_DB_PATH=str(root / 'app.db'),
APP_VAULT_PATH=str(root / 'vault'))
from app.container import container
from app.contracts import AgentRunCreateRequest, AgentRunStatus
from app.agent.runtime import AgentRuntime, AgentCapacityError
from app.agent.permissions import PermissionMode
from app.main import app
import httpx
report = {"python": platform.python_version(), "platform": platform.platform(),
"provider": "mock with 50 ms injected delay per model turn; no network", "results": []}
def save(name, value):
report['results'].append({"scenario": name, **value})
output.write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding='utf-8')
print(json.dumps(report['results'][-1], ensure_ascii=False), flush=True)
async def measured(operation):
delays = []
async def heartbeat():
while True:
start = time.perf_counter()
await asyncio.sleep(.01)
delays.append(max(0, (time.perf_counter()-start)*1000-10))
pulse = asyncio.create_task(heartbeat())
await asyncio.sleep(0)
start = time.perf_counter()
try:
result = await operation()
await asyncio.sleep(.02)
return {**result, "elapsed_ms": round((time.perf_counter()-start)*1000, 2),
"event_loop_lag": stats(delays)}
finally:
pulse.cancel()
await asyncio.gather(pulse, return_exceptions=True)
adapter = container.providers.get('mock').adapter
original = adapter.complete
async def delayed(request):
await asyncio.sleep(.05)
return await original(request)
adapter.complete = delayed
runtime = container.agent
for concurrency in (1, 10, 50, 200):
async def batch():
durations, sequences, statuses = [], [], []
semaphore = asyncio.Semaphore(concurrency)
async def one(index):
async with semaphore:
start = time.perf_counter()
created = await runtime.create_run(AgentRunCreateRequest(
input=f'/tool system.echo {{"text":"pressure-{index}"}}',
provider_id='mock', model='mock-1', allowed_tools=['system.echo']))
events = [event async for event in runtime.events(created.run_id)]
done = await runtime.wait(created.run_id)
durations.append((time.perf_counter()-start)*1000)
statuses.append(done.status.value)
sequence = [event.sequence for event in events]
sequences.append(sequence == list(range(len(sequence))))
assert done.status == AgentRunStatus.completed
assert done.tool_results[0].output == {"text": f"pressure-{index}"}
replay = [event.sequence async for event in runtime.events(created.run_id, after_sequence=2)]
assert replay == sequence[3:]
return created.run_id
ids = await asyncio.gather(*(one(i) for i in range(max(20, concurrency))))
assert all(sequences)
recovered = AgentRuntime(container.providers, container.tools, container.permissions,
trace_repository=runtime.trace_repository)
assert all(recovered.get_run(run_id).status == AgentRunStatus.completed for run_id in ids)
assert not any(record.subscribers for record in runtime._records.values())
return {"concurrency": concurrency, "runs": len(ids), "latency": stats(durations),
"completed": statuses.count('completed'), "ordered_events_and_replay": True,
"terminal_recovery": True, "retained_records": len(runtime._records)}
save('agent_tool_runs', await measured(batch))
# Hold model calls so all 200 records remain active while testing admission.
gate = asyncio.Event()
async def blocked(request):
await gate.wait()
return await original(request)
adapter.complete = blocked
async def capacity():
ids = [(await runtime.create_run(AgentRunCreateRequest(input='capacity', provider_id='mock', model='mock-1'))).run_id for _ in range(200)]
rejected = False
try:
await runtime.create_run(AgentRunCreateRequest(input='overflow', provider_id='mock', model='mock-1'))
except AgentCapacityError:
rejected = True
await asyncio.sleep(0)
latencies = []
for run_id in ids:
start = time.perf_counter()
await runtime.cancel(run_id)
latencies.append((time.perf_counter()-start)*1000)
states = await asyncio.gather(*(runtime.wait(run_id) for run_id in ids))
assert rejected and all(run.status == AgentRunStatus.cancelled for run in states)
assert all(not record.task or record.task.done() for record in runtime._records.values())
return {"active_limit": 200, "overflow_rejected": rejected, "cancelled": len(states), "cancel_latency": stats(latencies)}
save('capacity_and_cancel', await measured(capacity))
adapter.complete = delayed
async def permissions():
tool = container.tools.get('system.echo')
previous = tool.definition.permission
tool.definition.permission = 'stress.confirm'
container.permissions.policy.set_rule('stress.confirm', PermissionMode.confirm)
async def one(index):
created = await runtime.create_run(AgentRunCreateRequest(input='/tool system.echo {"text":"permission"}',
provider_id='mock', model='mock-1', allowed_tools=['system.echo'], tool_timeout_seconds=30))
stream = runtime.events(created.run_id)
try:
async for event in stream:
if event.event.value == 'PermissionRequired':
if index % 2: await runtime.cancel(created.run_id)
else: assert await runtime.resolve_permission(created.run_id, str(event.data['request_id']), 'allow')
break
finally:
await stream.aclose()
return (await runtime.wait(created.run_id)).status.value
try:
states = await asyncio.wait_for(asyncio.gather(*(one(i) for i in range(20))), 60)
assert states.count('completed') == states.count('cancelled') == 10
assert not any(record.subscribers for record in runtime._records.values())
return {"runs": 20, "approved_completed": 10, "cancelled_waiting_permission": 10, "subscribers_released": True}
finally:
tool.definition.permission = previous
save('permission_wait', await measured(permissions))
async def failures():
active = 0
async def injected(request):
text = request.messages[-1].content
if text == 'inject-provider-error': raise RuntimeError('injected offline provider failure')
if text == 'inject-model-timeout': await asyncio.sleep(60)
return await delayed(request)
tool = container.tools.get('system.echo')
executor = tool.executor
async def slow_tool(arguments, context):
nonlocal active
active += 1
try:
if arguments.text == 'slow': await asyncio.sleep(60)
value = executor(arguments, context)
return await value if inspect.isawaitable(value) else value
finally: active -= 1
adapter.complete = injected
tool.executor = slow_tool
async def one(index):
mode = index % 4
text = ['normal', 'inject-provider-error', 'inject-model-timeout', '/tool system.echo {"text":"slow"}'][mode]
created = await runtime.create_run(AgentRunCreateRequest(input=text, provider_id='mock', model='mock-1',
allowed_tools=['system.echo'], run_timeout_seconds=1 if mode == 2 else 30, tool_timeout_seconds=1))
result = await runtime.wait(created.run_id)
if mode == 0: assert result.status == AgentRunStatus.completed
elif mode == 1: assert result.status == AgentRunStatus.failed and result.error_code == 'AGENT_FAILED'
elif mode == 2: assert result.status == AgentRunStatus.failed and result.error_code == 'AGENT_TIMEOUT'
else: assert result.tool_results[0].error_code == 'TOOL_TIMEOUT'
try:
await asyncio.wait_for(asyncio.gather(*(one(i) for i in range(20))), 45)
assert active == 0
return {"runs": 20, "success": 5, "provider_errors": 5, "model_timeouts": 5,
"tool_timeouts": 5, "remaining_tool_executors": active}
finally:
adapter.complete = delayed
tool.executor = executor
save('failure_and_timeout_isolation', await measured(failures))
async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url='http://stress.local') as client:
for count in (100, 1000):
async def tasks():
timings = {name: [] for name in ('create', 'update', 'list', 'delete')}
async def call(method, path, **kwargs):
start = time.perf_counter()
response = await client.request(method, path, **kwargs)
response.raise_for_status()
return response.json(), (time.perf_counter()-start)*1000
semaphore = asyncio.Semaphore(20)
async def create(index):
async with semaphore:
body, duration = await call('POST', '/api/tasks', json={'title': f'压测任务 {index}', 'description': '独立测试数据'})
timings['create'].append(duration)
return body['task_id']
ids = await asyncio.gather(*(create(i) for i in range(count)))
first, _ = await call('GET', '/api/tasks')
seen = []
for offset in range(0, count, 100):
body, duration = await call('GET', f'/api/tasks?limit=100&offset={offset}')
timings['list'].append(duration)
seen.extend(item['task_id'] for item in body['items'])
assert len(set(seen)) == count and set(seen) == set(ids)
async def change(run_id):
async with semaphore:
body, duration = await call('PATCH', f'/api/tasks/{run_id}', json={'status': 'done'})
assert body['status'] == 'done'
timings['update'].append(duration)
_, duration = await call('DELETE', f'/api/tasks/{run_id}')
timings['delete'].append(duration)
await asyncio.gather(*(change(run_id) for run_id in ids))
final, _ = await call('GET', '/api/tasks')
assert final['page']['total'] == 0
return {"tasks": count, "client_concurrency": 20, "latencies": {key: stats(value) for key, value in timings.items()},
"default_page_count": len(first['items']), "default_total": first['page']['total'],
"pagination_complete": True, "final_total": 0}
save('task_api_crud', await measured(tasks))
report['database_bytes'] = (root / 'app.db').stat().st_size
report['complete'] = True
output.write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding='utf-8')
container.plugins.shutdown()
container.mcp_servers.shutdown()
from app.operation_logs import shutdown_logging
await asyncio.to_thread(shutdown_logging)
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--output', type=pathlib.Path, required=True)
args = parser.parse_args()
args.output.parent.mkdir(parents=True, exist_ok=True)
asyncio.run(main(args.output))
+96
View File
@@ -0,0 +1,96 @@
"""Real loopback HTTP task load with a separate, temporary Uvicorn process."""
import argparse
import asyncio
import json
import math
import os
from pathlib import Path
import socket
import subprocess
import sys
import tempfile
from time import perf_counter
import httpx
def stats(values):
values = sorted(values)
return {'count': len(values), 'p95_ms': round(values[math.ceil(len(values)*.95)-1], 2),
'max_ms': round(values[-1], 2)} if values else {'count': 0}
async def main(args):
with tempfile.TemporaryDirectory(prefix='notes-task-http-') as directory:
root = Path(directory)
env = {**os.environ, 'APP_DATA_DIR': str(root/'data'), 'APP_DB_PATH': str(root/'app.db'), 'APP_VAULT_PATH': str(root/'vault')}
with socket.socket() as sock:
sock.bind(('127.0.0.1', 0)); port = sock.getsockname()[1]
process = subprocess.Popen([sys.executable, '-m', 'uvicorn', 'app.main:app', '--host', '127.0.0.1', '--port', str(port), '--log-level', 'error'],
cwd=Path(__file__).resolve().parents[1], env=env, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
creationflags=subprocess.CREATE_NO_WINDOW if os.name == 'nt' else 0)
try:
async with httpx.AsyncClient(base_url=f'http://127.0.0.1:{port}', timeout=30) as client:
for _ in range(150):
if process.poll() is not None: raise RuntimeError('Isolated Uvicorn exited')
try:
(await client.get('/health')).raise_for_status(); break
except httpx.HTTPError: await asyncio.sleep(.1)
else: raise TimeoutError('Isolated Uvicorn startup')
sem = asyncio.Semaphore(args.concurrency)
timings = {key: [] for key in ('create', 'update', 'list', 'delete', 'health')}
errors = []
async def request(method, path, kind, **kwargs):
start = perf_counter()
response = await client.request(method, path, **kwargs)
timings[kind].append((perf_counter()-start)*1000)
response.raise_for_status()
return response.json()
async def health():
while True:
try: await request('GET', '/health', 'health')
except httpx.HTTPError as error: errors.append(type(error).__name__)
await asyncio.sleep(.05)
heartbeat = asyncio.create_task(health())
start = perf_counter()
try:
async def create(index):
async with sem:
return (await request('POST', '/api/tasks', 'create', json={'title': f'HTTP 压测 {index}'}))['task_id']
ids = await asyncio.gather(*(create(i) for i in range(args.count)))
seen = []
for offset in range(0, args.count, 100):
page = await request('GET', f'/api/tasks?limit=100&offset={offset}', 'list')
seen.extend(item['task_id'] for item in page['items'])
assert set(seen) == set(ids) and len(seen) == args.count
async def change(task_id):
async with sem:
updated = await request('PATCH', f'/api/tasks/{task_id}', 'update', json={'status': 'done'})
assert updated['status'] == 'done'
await request('DELETE', f'/api/tasks/{task_id}', 'delete')
await asyncio.gather(*(change(task_id) for task_id in ids))
remaining = await request('GET', '/api/tasks', 'list')
assert remaining['page']['total'] == 0
finally:
heartbeat.cancel(); await asyncio.gather(heartbeat, return_exceptions=True)
report = {'transport': 'real loopback HTTP, separate Uvicorn process', 'tasks': args.count,
'concurrency': args.concurrency, 'elapsed_ms': round((perf_counter()-start)*1000, 2),
'latencies': {key: stats(value) for key,value in timings.items()}, 'health_errors': errors,
'pagination_complete': True, 'final_total': 0}
args.output.parent.mkdir(parents=True, exist_ok=True)
args.output.write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding='utf-8')
print(json.dumps(report, ensure_ascii=False))
finally:
process.terminate()
try: process.wait(timeout=10)
except subprocess.TimeoutExpired: process.kill(); process.wait()
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--output', type=Path, required=True)
parser.add_argument('--count', type=int, default=1000)
parser.add_argument('--concurrency', type=int, default=20)
args = parser.parse_args()
if args.count < 1 or args.concurrency < 1: parser.error('count and concurrency must be positive')
asyncio.run(main(args))
+4 -3
View File
@@ -112,7 +112,7 @@ def test_permission_confirmation_resumes_agent() -> None:
break
assert request_id is not None
assert container.agent.resolve_permission(
assert await container.agent.resolve_permission(
created.run_id, request_id, "allow_once"
)
completed = await container.agent.wait(created.run_id)
@@ -370,8 +370,9 @@ def test_cancelling_permission_wait_cancels_run() -> None:
)
async with asyncio.timeout(2):
while container.agent.get_run(created.run_id).status != AgentRunStatus.waiting_permission:
await asyncio.sleep(0)
async for event in container.agent.events(created.run_id):
if event.event == AgentEventType.permission_required:
break
cancelled = await container.agent.cancel(created.run_id)
await container.agent.wait(created.run_id)
+68
View File
@@ -0,0 +1,68 @@
import asyncio
from types import SimpleNamespace
import pytest
from app.retrieval import activity
from app.services import index_service
@pytest.fixture(autouse=True)
def reset_activity(monkeypatch):
for field in ('active', 'completed', 'failed', 'cancelled'):
monkeypatch.setattr(activity, field, 0)
monkeypatch.setattr(index_service, '_active_job_id', None)
monkeypatch.setattr(index_service, '_active_scope', None)
monkeypatch.setattr(index_service, '_last_error', None)
monkeypatch.setattr(index_service.repository, 'stats', lambda: {'notes': 16, 'blocks': 3787})
monkeypatch.setattr(index_service.repository, 'get_index_meta', lambda: {})
def test_pending_rebuild_and_incremental_jobs(monkeypatch):
meta = {'workspace_vectors_pending': '1', 'note_vectors_pending:a': '1', 'note_vectors_pending:b': '1'}
monkeypatch.setattr(index_service.repository, 'get_index_meta', lambda: meta)
status = index_service.get_status()
assert status.vector_refresh_required and status.pending_jobs == 1
assert status.running_jobs == 0
monkeypatch.setattr(index_service, '_active_job_id', 'job_test')
monkeypatch.setattr(index_service, '_active_scope', 'all')
status = index_service.get_status()
assert (status.status, status.pending_jobs, status.running_jobs) == ('running', 1, 1)
monkeypatch.setattr(index_service, '_active_scope', 'note')
assert index_service.get_status().pending_jobs == 2
meta.pop('workspace_vectors_pending')
assert index_service.get_status().pending_jobs == 2
monkeypatch.setattr(index_service, '_active_job_id', None)
monkeypatch.setattr(index_service, '_last_error', 'failed')
status = index_service.get_status()
assert (status.status, status.pending_jobs, status.running_jobs) == ('failed', 2, 0)
meta.clear()
assert index_service.get_status().pending_jobs == 0
def test_search_activity_covers_completion_failure_and_cancellation():
async def scenario():
gate = asyncio.Event()
@activity.track_search
async def search(_self, request):
await gate.wait()
if request.query == 'fail':
raise ValueError('failure')
return 'ok'
tasks = [asyncio.create_task(search(None, SimpleNamespace(mode=mode, query=query)))
for mode, query in [('vector', 'ok'), ('hybrid', 'fail'), ('vector', 'cancel'), ('fts', 'ok')]]
await asyncio.sleep(0)
assert index_service.get_status().active_searches == 3
tasks[2].cancel()
with pytest.raises(asyncio.CancelledError):
await tasks[2]
gate.set()
results = await asyncio.gather(*tasks, return_exceptions=True)
assert results[0] == results[3] == 'ok'
status = index_service.get_status()
assert (status.active_searches, status.completed_searches, status.failed_searches, status.cancelled_searches) == (0, 1, 1, 1)
assert status.pending_jobs == 0
asyncio.run(scenario())
+36
View File
@@ -12,6 +12,42 @@ from app.local_models.runtime import Runtime
from app.providers.base import ProviderError
@pytest.mark.parametrize('threaded', [False, True])
def test_large_embedding_result_crosses_pipe_limit_without_truncation(monkeypatch, tmp_path, threaded):
import app.local_models.runtime as module
import app.local_models.process as process_module
from app.local_models.protocol import response_lines
vector = [0.012345678901234567] * 384
result = {'result': [vector] * 2111}
assert len(json.dumps(result).encode()) > 16 * 1024 * 1024
assert max(map(len, response_lines(result, 'embedding'))) < 16 * 1024 * 1024
worker = tmp_path / 'large_worker.py'
protocol_dir = Path(module.__file__).parent
worker.write_text(
'import sys,json\n'
f'sys.path.insert(0, {str(protocol_dir)!r})\n'
'from protocol import response_lines\n'
'request=json.load(sys.stdin)\n'
'vector=[0.012345678901234567]*384\n'
'for line in response_lines({"result":[vector]*len(request["payload"]["texts"])}, "embedding"):\n'
' sys.stdout.write(line)\n', encoding='utf-8')
monkeypatch.setattr(module, 'read_state', lambda key: {'status': 'installed'})
monkeypatch.setattr(module, 'interpreter', lambda *_: Path(sys.executable))
original_async = asyncio.create_subprocess_exec
original_threaded = process_module.ThreadedProcess
async def spawn(*args, **kwargs):
if threaded:
raise NotImplementedError
return await original_async(sys.executable, str(worker), **kwargs)
monkeypatch.setattr(module.asyncio, 'create_subprocess_exec', spawn)
monkeypatch.setattr(process_module, 'ThreadedProcess',
lambda args, **kwargs: original_threaded((sys.executable, str(worker)), **kwargs))
actual = asyncio.run(Runtime().infer('bekko', 'embedding', {'texts': ['test'] * 2111}))
assert actual == result['result']
def test_download_resumes_partial_and_checks_digest(monkeypatch):
payload = b'verified-model-weights'
entry = {'path':'model.safetensors','size':len(payload),'hash':hashlib.sha256(payload).hexdigest(),
+151
View File
@@ -0,0 +1,151 @@
import asyncio
import json
import logging
import threading
from pathlib import Path
from app.operation_logs import LogStore, ApplicationLogHandler, get_store, log_event, shutdown_logging
def test_logs_persist_filter_cursor_and_retention(tmp_path):
store = LogStore(tmp_path / 'logs.db', retain=3)
try:
for i in range(6):
store.emit('ERROR' if i % 2 else 'INFO', 'vectors', 'embedding.failed', {'run_id': f'run_{i}'})
store.queue.join()
first = store.query(limit=2)
assert len(first['items']) == 2 and first['next_cursor']
assert len(store.query(before=first['next_cursor'])['items']) == 1
assert len(store.query(level='ERROR')['items']) == 2
assert len(store.query(q='run_5')['items']) == 1
assert not store.query(source='tasks')['items']
finally:
store.close()
reopened = LogStore(tmp_path / 'logs.db', retain=3)
try:
assert len(reopened.query()['items']) == 3
finally:
reopened.close()
def test_logs_exclude_content_and_legacy_exception_messages():
try:
log_event('vectors', 'embedding.failed', level='ERROR',
error=ValueError('private note and secret'), model='embedding-v1',
prompt='private note', arguments={'api_key': 'secret'}, api_key='secret')
handler = ApplicationLogHandler()
record = logging.LogRecord('app.sample', logging.ERROR, __file__, 1,
'private note and secret %s', ('credentials',), None)
handler.emit(record)
handler.emit(record) # a logger propagated to another installed handler
store = get_store()
store.queue.join()
data = json.dumps(store.query())
assert len(store.query()['items']) == 2
assert 'private note' not in data and 'credentials' not in data and 'api_key' not in data
assert 'ValueError' in data and 'embedding-v1' in data
finally:
shutdown_logging()
def test_http_log_correlates_task_operations_without_body():
import httpx
from app.main import app
async def scenario():
async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url='http://test') as client:
response = await client.post('/api/tasks', json={'title': 'private title'})
assert response.status_code == 200
rid = response.headers['x-request-id']
store = get_store()
await asyncio.to_thread(store.queue.join)
result = (await client.get('/api/logs', params={'q': rid})).json()
assert len(result['items']) >= 2
assert 'private title' not in json.dumps(result)
assert any(item['event'] == 'task.created' for item in result['items'])
assert (await client.get('/api/logs', params={'limit': 201})).status_code == 422
try:
asyncio.run(scenario())
finally:
shutdown_logging()
def test_trace_writer_batches_off_loop_and_survives_cancel():
from app.agent.async_trace import AsyncTraceWriter
started, release = threading.Event(), threading.Event()
class Repository:
def write_batch(self, jobs):
assert threading.current_thread() is not threading.main_thread()
started.set()
assert release.wait(2)
self.jobs = jobs
async def scenario():
repo = Repository()
writer = AsyncTraceWriter(repo)
pending = asyncio.create_task(writer.submit('save', 'snapshot'))
while not started.is_set():
await asyncio.sleep(.001)
pending.cancel()
writer.worker.cancel() # simultaneous application shutdown
await asyncio.sleep(.005)
assert not pending.done()
release.set()
assert await asyncio.wait_for(pending, 2) is True
assert repo.jobs == [('save', ('snapshot',))]
assert writer.queue.empty()
asyncio.run(scenario())
def test_trace_write_failure_is_reported_and_next_submission_recovers():
from app.agent.async_trace import AsyncTraceWriter
class Repository:
fail = True
def write_batch(self, jobs):
if self.fail:
self.fail = False
raise OSError('disk unavailable')
async def scenario():
import pytest
writer = AsyncTraceWriter(Repository())
with pytest.raises(OSError):
await writer.submit('save', 'first')
assert not await writer.submit('save', 'second')
await writer.queue.join()
asyncio.run(scenario())
def test_log_write_failure_does_not_stall_queue(tmp_path, monkeypatch):
store = LogStore(tmp_path / 'failed.db')
try:
def broken():
raise OSError('disk unavailable')
monkeypatch.setattr(store, '_connect', broken)
store.emit('ERROR', 'vectors', 'embedding.failed', {'duration_ms': float('nan')})
store.queue.join()
assert store.failed == 1
finally:
store.close()
def test_cancelled_task_write_keeps_its_slot_until_commit():
from app.services.task_service import write_in_background
started, release = threading.Event(), threading.Event()
order = []
def first():
started.set()
assert release.wait(2)
order.append('first')
async def scenario():
import pytest
pending = asyncio.create_task(write_in_background(first))
while not started.is_set():
await asyncio.sleep(.001)
pending.cancel()
second = asyncio.create_task(write_in_background(lambda: order.append('second')))
await asyncio.sleep(.01)
assert order == []
release.set()
with pytest.raises(asyncio.CancelledError):
await pending
await second
assert order == ['first', 'second']
asyncio.run(scenario())