Add session scoped runs and cancellation
This commit is contained in:
武阳
+304
-48
@@ -19,6 +19,7 @@ from pathlib import Path
|
||||
from threading import Lock, RLock
|
||||
from typing import Any
|
||||
from urllib import error, request
|
||||
from uuid import uuid4
|
||||
|
||||
from fastapi import FastAPI, HTTPException
|
||||
from fastapi.responses import FileResponse, JSONResponse, StreamingResponse
|
||||
@@ -167,6 +168,129 @@ class AgentInstanceConfig:
|
||||
allow_write: bool
|
||||
|
||||
|
||||
class RunProcessRegistry:
|
||||
def __init__(self) -> None:
|
||||
self._lock = Lock()
|
||||
self._processes: set[Any] = set()
|
||||
|
||||
def add(self, process: Any) -> None:
|
||||
with self._lock:
|
||||
self._processes.add(process)
|
||||
|
||||
def discard(self, process: Any) -> None:
|
||||
with self._lock:
|
||||
self._processes.discard(process)
|
||||
|
||||
def kill_all(self) -> None:
|
||||
with self._lock:
|
||||
processes = list(self._processes)
|
||||
for process in processes:
|
||||
try:
|
||||
if process.poll() is None:
|
||||
process.terminate()
|
||||
except OSError:
|
||||
continue
|
||||
deadline = time.monotonic() + 1.0
|
||||
for process in processes:
|
||||
try:
|
||||
remaining = max(0.0, deadline - time.monotonic())
|
||||
process.wait(timeout=remaining)
|
||||
except Exception:
|
||||
try:
|
||||
process.kill()
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
|
||||
@dataclass
|
||||
class RunRecord:
|
||||
run_id: str
|
||||
account_key: str
|
||||
session_id: str
|
||||
status: str
|
||||
started_at: float
|
||||
updated_at: float
|
||||
cancel_event: threading.Event
|
||||
process_registry: RunProcessRegistry
|
||||
current_stage: str = ''
|
||||
error: str = ''
|
||||
|
||||
|
||||
class RunManager:
|
||||
def __init__(self) -> None:
|
||||
self._lock = RLock()
|
||||
self._runs: dict[str, RunRecord] = {}
|
||||
self._latest_by_session: dict[tuple[str, str], str] = {}
|
||||
|
||||
def start(self, account_key: str, session_id: str) -> RunRecord:
|
||||
now = time.time()
|
||||
record = RunRecord(
|
||||
run_id=uuid4().hex,
|
||||
account_key=account_key,
|
||||
session_id=session_id,
|
||||
status='queued',
|
||||
started_at=now,
|
||||
updated_at=now,
|
||||
cancel_event=threading.Event(),
|
||||
process_registry=RunProcessRegistry(),
|
||||
)
|
||||
with self._lock:
|
||||
self._runs[record.run_id] = record
|
||||
self._latest_by_session[(account_key, session_id)] = record.run_id
|
||||
return record
|
||||
|
||||
def update(self, run_id: str, *, status: str | None = None, stage: str | None = None, error: str | None = None) -> None:
|
||||
with self._lock:
|
||||
record = self._runs.get(run_id)
|
||||
if record is None:
|
||||
return
|
||||
if status is not None:
|
||||
record.status = status
|
||||
if stage is not None:
|
||||
record.current_stage = stage
|
||||
if error is not None:
|
||||
record.error = error
|
||||
record.updated_at = time.time()
|
||||
|
||||
def finish(self, run_id: str, status: str) -> None:
|
||||
self.update(run_id, status=status)
|
||||
|
||||
def cancel_latest(self, account_key: str, session_id: str) -> bool:
|
||||
with self._lock:
|
||||
run_id = self._latest_by_session.get((account_key, session_id))
|
||||
record = self._runs.get(run_id or '')
|
||||
return self.cancel_record(record)
|
||||
|
||||
def cancel_run(self, run_id: str) -> bool:
|
||||
with self._lock:
|
||||
record = self._runs.get(run_id)
|
||||
return self.cancel_record(record)
|
||||
|
||||
def cancel_record(self, record: RunRecord | None) -> bool:
|
||||
if record is None or record.status in {'completed', 'failed', 'cancelled'}:
|
||||
return False
|
||||
record.cancel_event.set()
|
||||
record.process_registry.kill_all()
|
||||
self.update(record.run_id, status='cancelled', stage='用户已取消')
|
||||
return True
|
||||
|
||||
def snapshot_latest(self, account_key: str, session_id: str) -> dict[str, Any] | None:
|
||||
with self._lock:
|
||||
run_id = self._latest_by_session.get((account_key, session_id))
|
||||
record = self._runs.get(run_id or '')
|
||||
if record is None:
|
||||
return None
|
||||
return {
|
||||
'run_id': record.run_id,
|
||||
'session_id': record.session_id,
|
||||
'status': record.status,
|
||||
'current_stage': record.current_stage,
|
||||
'started_at': record.started_at,
|
||||
'updated_at': record.updated_at,
|
||||
'error': record.error,
|
||||
}
|
||||
|
||||
|
||||
class AgentState:
|
||||
"""Holds account-scoped agent instances, config, and execution locks."""
|
||||
|
||||
@@ -186,6 +310,7 @@ class AgentState:
|
||||
self._agents: dict[str, LocalCodingAgent] = {}
|
||||
self._run_locks: dict[str, Lock] = {}
|
||||
self._account_configs: dict[str, AgentInstanceConfig] = {}
|
||||
self.run_manager = RunManager()
|
||||
self._default_config = AgentInstanceConfig(
|
||||
cwd=cwd.resolve(),
|
||||
model=model,
|
||||
@@ -222,6 +347,10 @@ class AgentState:
|
||||
def _account_key(self, account_id: str | None) -> str:
|
||||
return _safe_account_id(account_id) if account_id else '__default__'
|
||||
|
||||
def _session_key(self, account_id: str | None, session_id: str | None = None) -> str:
|
||||
session_part = _safe_session_id(session_id) or '__shared__'
|
||||
return f'{self._account_key(account_id)}:{session_part}'
|
||||
|
||||
def _config_for(self, account_id: str | None) -> AgentInstanceConfig:
|
||||
if not account_id:
|
||||
return self._default_config
|
||||
@@ -298,17 +427,25 @@ class AgentState:
|
||||
runtime_config=runtime_config,
|
||||
)
|
||||
|
||||
def agent_for(self, account_id: str | None = None) -> LocalCodingAgent:
|
||||
def agent_for(
|
||||
self,
|
||||
account_id: str | None = None,
|
||||
session_id: str | None = None,
|
||||
) -> LocalCodingAgent:
|
||||
with self._lock:
|
||||
key = self._account_key(account_id)
|
||||
key = self._session_key(account_id, session_id)
|
||||
agent = self._agents.get(key)
|
||||
if agent is None:
|
||||
agent = self._build_agent(account_id)
|
||||
self._agents[key] = agent
|
||||
return agent
|
||||
|
||||
def run_lock_for(self, account_id: str | None = None) -> Lock:
|
||||
key = self._account_key(account_id)
|
||||
def run_lock_for(
|
||||
self,
|
||||
account_id: str | None = None,
|
||||
session_id: str | None = None,
|
||||
) -> Lock:
|
||||
key = self._session_key(account_id, session_id)
|
||||
with self._lock:
|
||||
lock = self._run_locks.get(key)
|
||||
if lock is None:
|
||||
@@ -345,7 +482,10 @@ class AgentState:
|
||||
if allow_write is not None:
|
||||
config = replace(config, allow_write=allow_write)
|
||||
self._set_config_for(account_id, config)
|
||||
self._agents.pop(self._account_key(account_id), None)
|
||||
account_prefix = f'{self._account_key(account_id)}:'
|
||||
for key in list(self._agents):
|
||||
if key.startswith(account_prefix):
|
||||
self._agents.pop(key, None)
|
||||
|
||||
def snapshot(self, account_id: str | None = None) -> dict[str, Any]:
|
||||
with self._lock:
|
||||
@@ -391,6 +531,12 @@ class ChatRequest(BaseModel):
|
||||
session_id: str | None = None
|
||||
|
||||
|
||||
class RunCancelRequest(BaseModel):
|
||||
session_id: str = Field(min_length=1)
|
||||
account_id: str | None = None
|
||||
run_id: str | None = None
|
||||
|
||||
|
||||
class StateUpdate(BaseModel):
|
||||
model: str | None = None
|
||||
base_url: str | None = None
|
||||
@@ -598,57 +744,136 @@ def create_app(state: AgentState) -> FastAPI:
|
||||
)
|
||||
return _serialize_token_budget(snapshot)
|
||||
|
||||
@app.get('/api/runs/latest')
|
||||
async def latest_run(
|
||||
session_id: str,
|
||||
account_id: str | None = None,
|
||||
) -> dict[str, Any]:
|
||||
safe_id = _safe_session_id(session_id)
|
||||
if safe_id is None:
|
||||
raise HTTPException(status_code=400, detail='Invalid session id')
|
||||
snapshot = state.run_manager.snapshot_latest(
|
||||
state._account_key(account_id),
|
||||
safe_id,
|
||||
)
|
||||
return snapshot or {'session_id': safe_id, 'status': 'idle'}
|
||||
|
||||
@app.post('/api/runs/cancel')
|
||||
async def cancel_run(payload: RunCancelRequest) -> dict[str, Any]:
|
||||
safe_id = _safe_session_id(payload.session_id)
|
||||
if safe_id is None:
|
||||
raise HTTPException(status_code=400, detail='Invalid session id')
|
||||
cancelled = (
|
||||
state.run_manager.cancel_run(payload.run_id)
|
||||
if payload.run_id
|
||||
else state.run_manager.cancel_latest(
|
||||
state._account_key(payload.account_id),
|
||||
safe_id,
|
||||
)
|
||||
)
|
||||
return {'session_id': safe_id, 'cancelled': cancelled}
|
||||
|
||||
def _run_chat_payload(
|
||||
request: ChatRequest,
|
||||
event_sink: Any | None = None,
|
||||
) -> dict[str, Any]:
|
||||
run_lock = state.run_lock_for(request.account_id)
|
||||
with run_lock:
|
||||
agent = state.agent_for(request.account_id)
|
||||
config = state.config_for(request.account_id)
|
||||
session_directory = state.account_paths(request.account_id)['sessions']
|
||||
requested_session_id = _safe_session_id(
|
||||
request.resume_session_id or request.session_id
|
||||
requested_session_id = _safe_session_id(
|
||||
request.resume_session_id or request.session_id
|
||||
) or uuid4().hex
|
||||
account_key = state._account_key(request.account_id)
|
||||
run_record = state.run_manager.start(account_key, requested_session_id)
|
||||
run_lock = state.run_lock_for(request.account_id, requested_session_id)
|
||||
agent = state.agent_for(request.account_id, requested_session_id)
|
||||
config = state.config_for(request.account_id)
|
||||
session_directory = state.account_paths(request.account_id)['sessions']
|
||||
if request.resume_session_id is None:
|
||||
_save_in_progress_session(
|
||||
directory=session_directory,
|
||||
agent=agent,
|
||||
session_id=requested_session_id,
|
||||
prompt=request.prompt.strip(),
|
||||
)
|
||||
previous_title = _read_session_title(
|
||||
session_directory,
|
||||
requested_session_id,
|
||||
previous_title = _read_session_title(
|
||||
session_directory,
|
||||
requested_session_id,
|
||||
)
|
||||
if run_lock.locked():
|
||||
_emit_runtime_event(
|
||||
event_sink,
|
||||
{
|
||||
'type': 'run_queued',
|
||||
'run_id': run_record.run_id,
|
||||
'session_id': requested_session_id or '',
|
||||
'account_id': request.account_id or '',
|
||||
},
|
||||
)
|
||||
with run_lock:
|
||||
if run_record.cancel_event.is_set():
|
||||
state.run_manager.finish(run_record.run_id, 'cancelled')
|
||||
return {
|
||||
'final_output': '已取消排队中的请求。',
|
||||
'turns': 0,
|
||||
'tool_calls': 0,
|
||||
'transcript': (),
|
||||
'session_id': requested_session_id,
|
||||
'usage': {},
|
||||
'total_cost_usd': 0.0,
|
||||
'stop_reason': 'cancelled',
|
||||
}
|
||||
state.run_manager.update(run_record.run_id, status='running', stage='后端开始执行')
|
||||
_emit_runtime_event(
|
||||
event_sink,
|
||||
{
|
||||
'type': 'run_started',
|
||||
'run_id': run_record.run_id,
|
||||
'session_id': requested_session_id or '',
|
||||
'account_id': request.account_id or '',
|
||||
},
|
||||
)
|
||||
agent.tool_context = replace(
|
||||
agent.tool_context,
|
||||
cancel_event=run_record.cancel_event,
|
||||
process_registry=run_record.process_registry,
|
||||
)
|
||||
started_at = time.perf_counter()
|
||||
if request.resume_session_id is not None:
|
||||
if requested_session_id is None:
|
||||
raise HTTPException(
|
||||
status_code=404,
|
||||
detail='Session to resume not found',
|
||||
)
|
||||
try:
|
||||
try:
|
||||
stored = load_agent_session(
|
||||
requested_session_id,
|
||||
directory=session_directory,
|
||||
if request.resume_session_id is not None:
|
||||
try:
|
||||
stored = load_agent_session(
|
||||
requested_session_id,
|
||||
directory=session_directory,
|
||||
)
|
||||
except FileNotFoundError:
|
||||
raise HTTPException(
|
||||
status_code=404,
|
||||
detail='Session to resume not found',
|
||||
)
|
||||
result = agent.resume(
|
||||
request.prompt.strip(),
|
||||
stored,
|
||||
runtime_context=request.runtime_context,
|
||||
event_sink=event_sink,
|
||||
)
|
||||
else:
|
||||
result = agent.run(
|
||||
request.prompt.strip(),
|
||||
session_id=requested_session_id,
|
||||
runtime_context=request.runtime_context,
|
||||
event_sink=event_sink,
|
||||
)
|
||||
except Exception as exc:
|
||||
state.run_manager.update(
|
||||
run_record.run_id,
|
||||
status='failed',
|
||||
error=str(exc),
|
||||
)
|
||||
except FileNotFoundError:
|
||||
raise HTTPException(
|
||||
status_code=404,
|
||||
detail='Session to resume not found',
|
||||
)
|
||||
result = agent.resume(
|
||||
request.prompt.strip(),
|
||||
stored,
|
||||
runtime_context=request.runtime_context,
|
||||
event_sink=event_sink,
|
||||
)
|
||||
else:
|
||||
_save_in_progress_session(
|
||||
directory=session_directory,
|
||||
agent=agent,
|
||||
session_id=requested_session_id,
|
||||
prompt=request.prompt.strip(),
|
||||
)
|
||||
result = agent.run(
|
||||
request.prompt.strip(),
|
||||
session_id=requested_session_id,
|
||||
runtime_context=request.runtime_context,
|
||||
event_sink=event_sink,
|
||||
raise
|
||||
finally:
|
||||
agent.tool_context = replace(
|
||||
agent.tool_context,
|
||||
cancel_event=None,
|
||||
process_registry=None,
|
||||
)
|
||||
elapsed_ms = max(0, int((time.perf_counter() - started_at) * 1000))
|
||||
payload = _serialize_run_result(result)
|
||||
@@ -668,6 +893,10 @@ def create_app(state: AgentState) -> FastAPI:
|
||||
previous_title=previous_title,
|
||||
)
|
||||
_annotate_transcript_elapsed(payload, elapsed_ms)
|
||||
state.run_manager.finish(
|
||||
run_record.run_id,
|
||||
'cancelled' if run_record.cancel_event.is_set() else 'completed',
|
||||
)
|
||||
return payload
|
||||
|
||||
# ------------- chat ------------------------------------------------------
|
||||
@@ -732,9 +961,27 @@ def create_app(state: AgentState) -> FastAPI:
|
||||
|
||||
threading.Thread(target=worker, daemon=True).start()
|
||||
|
||||
stream_started = time.perf_counter()
|
||||
|
||||
def generate() -> Any:
|
||||
while True:
|
||||
item = event_queue.get()
|
||||
try:
|
||||
item = event_queue.get(timeout=15)
|
||||
except queue.Empty:
|
||||
yield json.dumps(
|
||||
{
|
||||
'kind': 'event',
|
||||
'event': {
|
||||
'type': 'server_heartbeat',
|
||||
'elapsed_ms': max(
|
||||
0,
|
||||
int((time.perf_counter() - stream_started) * 1000),
|
||||
),
|
||||
},
|
||||
},
|
||||
ensure_ascii=False,
|
||||
) + '\n'
|
||||
continue
|
||||
if item is None:
|
||||
break
|
||||
yield json.dumps(item, ensure_ascii=False) + '\n'
|
||||
@@ -790,6 +1037,15 @@ def _serialize_stored_session(stored: StoredAgentSession) -> dict[str, Any]:
|
||||
}
|
||||
|
||||
|
||||
def _emit_runtime_event(
|
||||
event_sink: Any | None,
|
||||
event: dict[str, object],
|
||||
) -> None:
|
||||
if event_sink is None:
|
||||
return
|
||||
event_sink(event)
|
||||
|
||||
|
||||
def _save_in_progress_session(
|
||||
*,
|
||||
directory: Path,
|
||||
|
||||
Reference in New Issue
Block a user