fix: stabilize background operations and large embedding results
This commit is contained in:
@@ -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
@@ -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,
|
||||
)
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
)
|
||||
|
||||
@@ -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'
|
||||
@@ -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)
|
||||
|
||||
@@ -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'))
|
||||
|
||||
@@ -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)
|
||||
@@ -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"])
|
||||
|
||||
@@ -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()
|
||||
@@ -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
|
||||
@@ -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)
|
||||
|
||||
@@ -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
@@ -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}
|
||||
)
|
||||
|
||||
@@ -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})
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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 (?,?,?,?,?,?,?,?,?,?,?)", (
|
||||
|
||||
@@ -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))
|
||||
@@ -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))
|
||||
@@ -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)
|
||||
|
||||
@@ -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())
|
||||
@@ -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(),
|
||||
|
||||
@@ -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())
|
||||
Reference in New Issue
Block a user