Add session scoped runs and cancellation

This commit is contained in:
武阳
2026-05-07 16:26:20 +08:00
parent a4a0677ce4
commit 9d7d496443
6 changed files with 542 additions and 76 deletions
+279 -23
View File
@@ -19,6 +19,7 @@ from pathlib import Path
from threading import Lock, RLock from threading import Lock, RLock
from typing import Any from typing import Any
from urllib import error, request from urllib import error, request
from uuid import uuid4
from fastapi import FastAPI, HTTPException from fastapi import FastAPI, HTTPException
from fastapi.responses import FileResponse, JSONResponse, StreamingResponse from fastapi.responses import FileResponse, JSONResponse, StreamingResponse
@@ -167,6 +168,129 @@ class AgentInstanceConfig:
allow_write: bool 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: class AgentState:
"""Holds account-scoped agent instances, config, and execution locks.""" """Holds account-scoped agent instances, config, and execution locks."""
@@ -186,6 +310,7 @@ class AgentState:
self._agents: dict[str, LocalCodingAgent] = {} self._agents: dict[str, LocalCodingAgent] = {}
self._run_locks: dict[str, Lock] = {} self._run_locks: dict[str, Lock] = {}
self._account_configs: dict[str, AgentInstanceConfig] = {} self._account_configs: dict[str, AgentInstanceConfig] = {}
self.run_manager = RunManager()
self._default_config = AgentInstanceConfig( self._default_config = AgentInstanceConfig(
cwd=cwd.resolve(), cwd=cwd.resolve(),
model=model, model=model,
@@ -222,6 +347,10 @@ class AgentState:
def _account_key(self, account_id: str | None) -> str: def _account_key(self, account_id: str | None) -> str:
return _safe_account_id(account_id) if account_id else '__default__' 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: def _config_for(self, account_id: str | None) -> AgentInstanceConfig:
if not account_id: if not account_id:
return self._default_config return self._default_config
@@ -298,17 +427,25 @@ class AgentState:
runtime_config=runtime_config, 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: with self._lock:
key = self._account_key(account_id) key = self._session_key(account_id, session_id)
agent = self._agents.get(key) agent = self._agents.get(key)
if agent is None: if agent is None:
agent = self._build_agent(account_id) agent = self._build_agent(account_id)
self._agents[key] = agent self._agents[key] = agent
return agent return agent
def run_lock_for(self, account_id: str | None = None) -> Lock: def run_lock_for(
key = self._account_key(account_id) self,
account_id: str | None = None,
session_id: str | None = None,
) -> Lock:
key = self._session_key(account_id, session_id)
with self._lock: with self._lock:
lock = self._run_locks.get(key) lock = self._run_locks.get(key)
if lock is None: if lock is None:
@@ -345,7 +482,10 @@ class AgentState:
if allow_write is not None: if allow_write is not None:
config = replace(config, allow_write=allow_write) config = replace(config, allow_write=allow_write)
self._set_config_for(account_id, config) 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]: def snapshot(self, account_id: str | None = None) -> dict[str, Any]:
with self._lock: with self._lock:
@@ -391,6 +531,12 @@ class ChatRequest(BaseModel):
session_id: str | None = None 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): class StateUpdate(BaseModel):
model: str | None = None model: str | None = None
base_url: str | None = None base_url: str | None = None
@@ -598,29 +744,101 @@ def create_app(state: AgentState) -> FastAPI:
) )
return _serialize_token_budget(snapshot) 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( def _run_chat_payload(
request: ChatRequest, request: ChatRequest,
event_sink: Any | None = None, event_sink: Any | None = None,
) -> dict[str, Any]: ) -> 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( requested_session_id = _safe_session_id(
request.resume_session_id or request.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( previous_title = _read_session_title(
session_directory, session_directory,
requested_session_id, requested_session_id,
) )
started_at = time.perf_counter() if run_lock.locked():
if request.resume_session_id is not None: _emit_runtime_event(
if requested_session_id is None: event_sink,
raise HTTPException( {
status_code=404, 'type': 'run_queued',
detail='Session to resume not found', '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()
try:
try:
if request.resume_session_id is not None:
try: try:
stored = load_agent_session( stored = load_agent_session(
requested_session_id, requested_session_id,
@@ -638,18 +856,25 @@ def create_app(state: AgentState) -> FastAPI:
event_sink=event_sink, event_sink=event_sink,
) )
else: else:
_save_in_progress_session(
directory=session_directory,
agent=agent,
session_id=requested_session_id,
prompt=request.prompt.strip(),
)
result = agent.run( result = agent.run(
request.prompt.strip(), request.prompt.strip(),
session_id=requested_session_id, session_id=requested_session_id,
runtime_context=request.runtime_context, runtime_context=request.runtime_context,
event_sink=event_sink, event_sink=event_sink,
) )
except Exception as exc:
state.run_manager.update(
run_record.run_id,
status='failed',
error=str(exc),
)
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)) elapsed_ms = max(0, int((time.perf_counter() - started_at) * 1000))
payload = _serialize_run_result(result) payload = _serialize_run_result(result)
payload['elapsed_ms'] = elapsed_ms payload['elapsed_ms'] = elapsed_ms
@@ -668,6 +893,10 @@ def create_app(state: AgentState) -> FastAPI:
previous_title=previous_title, previous_title=previous_title,
) )
_annotate_transcript_elapsed(payload, elapsed_ms) _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 return payload
# ------------- chat ------------------------------------------------------ # ------------- chat ------------------------------------------------------
@@ -732,9 +961,27 @@ def create_app(state: AgentState) -> FastAPI:
threading.Thread(target=worker, daemon=True).start() threading.Thread(target=worker, daemon=True).start()
stream_started = time.perf_counter()
def generate() -> Any: def generate() -> Any:
while True: 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: if item is None:
break break
yield json.dumps(item, ensure_ascii=False) + '\n' 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( def _save_in_progress_session(
*, *,
directory: Path, directory: Path,
+92 -16
View File
@@ -29,6 +29,7 @@ type ClawChatResponse = {
type ClawRuntimeEvent = { type ClawRuntimeEvent = {
type?: string; type?: string;
run_id?: string;
tool_name?: string; tool_name?: string;
tool_call_id?: string; tool_call_id?: string;
arguments?: unknown; arguments?: unknown;
@@ -36,6 +37,7 @@ type ClawRuntimeEvent = {
metadata?: Record<string, unknown>; metadata?: Record<string, unknown>;
ok?: boolean; ok?: boolean;
delta?: string; delta?: string;
elapsed_ms?: number;
}; };
type ClawStreamItem = type ClawStreamItem =
@@ -43,6 +45,13 @@ type ClawStreamItem =
| { kind: "result"; payload?: ClawChatResponse } | { kind: "result"; payload?: ClawChatResponse }
| { kind: "error"; error?: string; detail?: string; status_code?: number }; | { kind: "error"; error?: string; detail?: string; status_code?: number };
type StreamState = {
textStarted: boolean;
textStreamed: boolean;
phase: "waiting" | "queued" | "running" | "done";
runId?: string;
};
type ChatRequestBody = { type ChatRequestBody = {
id?: string; id?: string;
messages: UIMessage[]; messages: UIMessage[];
@@ -92,26 +101,25 @@ export async function POST(req: Request) {
const stream = createUIMessageStream({ const stream = createUIMessageStream({
execute: async ({ writer }) => { execute: async ({ writer }) => {
const streamState = { textStarted: false, textStreamed: false }; const streamState: StreamState = {
let heartbeatCount = 0; textStarted: false,
const heartbeat = setInterval(() => { textStreamed: false,
heartbeatCount += 1; phase: "waiting",
writer.write({ };
type: "reasoning-delta",
id: "reasoning-1",
delta: `\n仍在处理请求,已等待 ${formatDuration(heartbeatCount * 15000)}`,
});
}, 15000);
writer.write({ type: "start" }); writer.write({ type: "start" });
writer.write({ type: "start-step" }); writer.write({ type: "start-step" });
try { try {
const cancelOnAbort = () => {
void cancelClawRun(account.id, sessionId, streamState.runId);
};
req.signal.addEventListener("abort", cancelOnAbort, { once: true });
const startedAt = Date.now(); const startedAt = Date.now();
writer.write({ type: "reasoning-start", id: "reasoning-1" }); writer.write({ type: "reasoning-start", id: "reasoning-1" });
writer.write({ writer.write({
type: "reasoning-delta", type: "reasoning-delta",
id: "reasoning-1", id: "reasoning-1",
delta: "思考中,正在判断是否需要调用工具。", delta: "请求已发送,等待后端开始处理。",
}); });
const payload = await callClawBackendStream( const payload = await callClawBackendStream(
userPrompt, userPrompt,
@@ -121,7 +129,9 @@ export async function POST(req: Request) {
resumeId, resumeId,
writer, writer,
streamState, streamState,
req.signal,
); );
req.signal.removeEventListener("abort", cancelOnAbort);
const elapsedMs = payload.elapsed_ms ?? Date.now() - startedAt; const elapsedMs = payload.elapsed_ms ?? Date.now() - startedAt;
const text = formatClawResponse(payload); const text = formatClawResponse(payload);
@@ -152,8 +162,22 @@ export async function POST(req: Request) {
usage: payload.usage, usage: payload.usage,
}, },
}); });
} finally { } catch (err) {
clearInterval(heartbeat); const text =
err instanceof Error && err.name === "AbortError"
? "请求已取消。"
: err instanceof Error
? err.message
: "请求失败。";
writer.write({ type: "reasoning-end", id: "reasoning-1" });
writer.write({ type: "text-start", id: "text-1" });
writer.write({ type: "text-delta", id: "text-1", delta: text });
writer.write({ type: "text-end", id: "text-1" });
writer.write({ type: "finish-step" });
writer.write({
type: "finish",
finishReason: "error",
});
} }
}, },
}); });
@@ -326,12 +350,14 @@ async function callClawBackendStream(
sessionId: string, sessionId: string,
resumeSessionId?: string, resumeSessionId?: string,
writer?: UIMessageStreamWriter<UIMessage>, writer?: UIMessageStreamWriter<UIMessage>,
streamState?: { textStarted: boolean; textStreamed: boolean }, streamState?: StreamState,
signal?: AbortSignal,
): Promise<ClawChatResponse> { ): Promise<ClawChatResponse> {
try { try {
const response = await fetch(`${CLAW_API_URL}/api/chat/stream`, { const response = await fetch(`${CLAW_API_URL}/api/chat/stream`, {
method: "POST", method: "POST",
headers: { "content-type": "application/json" }, headers: { "content-type": "application/json" },
signal,
body: JSON.stringify({ body: JSON.stringify({
prompt, prompt,
runtime_context: runtimeContext, runtime_context: runtimeContext,
@@ -364,10 +390,30 @@ async function callClawBackendStream(
} }
} }
async function cancelClawRun(
accountId: string,
sessionId: string,
runId?: string,
) {
try {
await fetch(`${CLAW_API_URL}/api/runs/cancel`, {
method: "POST",
headers: { "content-type": "application/json" },
body: JSON.stringify({
account_id: accountId,
session_id: sessionId,
...(runId ? { run_id: runId } : {}),
}),
});
} catch {
// 浏览器关闭或网络中断时,取消请求本身也可能失败;后端流断开仍会结束本次响应。
}
}
async function consumeClawStream( async function consumeClawStream(
body: ReadableStream<Uint8Array>, body: ReadableStream<Uint8Array>,
writer?: UIMessageStreamWriter<UIMessage>, writer?: UIMessageStreamWriter<UIMessage>,
streamState?: { textStarted: boolean; textStreamed: boolean }, streamState?: StreamState,
) { ) {
const reader = body.getReader(); const reader = body.getReader();
const decoder = new TextDecoder(); const decoder = new TextDecoder();
@@ -430,8 +476,38 @@ function parseStreamLine(line: string): ClawStreamItem | undefined {
function writeRuntimeEvent( function writeRuntimeEvent(
writer: UIMessageStreamWriter<UIMessage>, writer: UIMessageStreamWriter<UIMessage>,
event: ClawRuntimeEvent, event: ClawRuntimeEvent,
streamState?: { textStarted: boolean; textStreamed: boolean }, streamState?: StreamState,
) { ) {
if (event.type === "run_queued") {
if (streamState) streamState.phase = "queued";
if (streamState && event.run_id) streamState.runId = event.run_id;
writer.write({
type: "reasoning-delta",
id: "reasoning-1",
delta: "\n当前会话已有任务在执行,本轮正在排队。",
});
}
if (event.type === "run_started") {
if (streamState) streamState.phase = "running";
if (streamState && event.run_id) streamState.runId = event.run_id;
writer.write({
type: "reasoning-delta",
id: "reasoning-1",
delta: "\n后端已开始执行本轮任务。",
});
}
if (event.type === "server_heartbeat") {
const elapsed =
typeof event.elapsed_ms === "number"
? `,已等待 ${formatDuration(event.elapsed_ms)}`
: "";
const phase = streamState?.phase === "queued" ? "仍在排队" : "后端仍在处理";
writer.write({
type: "reasoning-delta",
id: "reasoning-1",
delta: `\n${phase}${elapsed}`,
});
}
if (event.type === "final_text_start") { if (event.type === "final_text_start") {
if (!streamState?.textStarted) { if (!streamState?.textStarted) {
writer.write({ type: "text-start", id: "text-1" }); writer.write({ type: "text-start", id: "text-1" });
+24
View File
@@ -1060,6 +1060,30 @@ class LocalCodingAgent:
tool_result = update.result tool_result = update.result
if tool_result is None: if tool_result is None:
raise RuntimeError(f'Tool executor returned no final result for {tool_call.name}') raise RuntimeError(f'Tool executor returned no final result for {tool_call.name}')
if (
self.tool_context.cancel_event is not None
and self.tool_context.cancel_event.is_set()
):
result = AgentRunResult(
final_output='已取消当前请求。',
turns=turn_index,
tool_calls=tool_calls,
transcript=session.transcript(),
events=tuple(stream_events),
usage=total_usage,
total_cost_usd=total_cost_usd,
stop_reason='cancelled',
file_history=tuple(file_history),
session_id=session_id,
scratchpad_directory=(
str(scratchpad_directory)
if scratchpad_directory is not None
else None
),
)
result = self._persist_session(session, result)
self.last_run_result = result
return result
if self.plugin_runtime is not None: if self.plugin_runtime is not None:
self.plugin_runtime.record_tool_result( self.plugin_runtime.record_tool_result(
tool_call.name, tool_call.name,
+6
View File
@@ -40,6 +40,8 @@ class ToolExecutionContext:
scratchpad_directory: Path | None = None scratchpad_directory: Path | None = None
python_env_dir: Path | None = None python_env_dir: Path | None = None
extra_env: dict[str, str] = field(default_factory=dict) extra_env: dict[str, str] = field(default_factory=dict)
cancel_event: Any | None = None
process_registry: Any | None = None
tool_registry: dict[str, 'AgentTool'] | None = None tool_registry: dict[str, 'AgentTool'] | None = None
search_runtime: 'SearchRuntime | None' = None search_runtime: 'SearchRuntime | None' = None
account_runtime: 'AccountRuntime | None' = None account_runtime: 'AccountRuntime | None' = None
@@ -127,6 +129,8 @@ def build_tool_context(
scratchpad_directory: Path | None = None, scratchpad_directory: Path | None = None,
python_env_dir: Path | None = None, python_env_dir: Path | None = None,
extra_env: dict[str, str] | None = None, extra_env: dict[str, str] | None = None,
cancel_event: Any | None = None,
process_registry: Any | None = None,
tool_registry: dict[str, AgentTool] | None = None, tool_registry: dict[str, AgentTool] | None = None,
search_runtime: 'SearchRuntime | None' = None, search_runtime: 'SearchRuntime | None' = None,
account_runtime: 'AccountRuntime | None' = None, account_runtime: 'AccountRuntime | None' = None,
@@ -156,6 +160,8 @@ def build_tool_context(
else None else None
), ),
extra_env=dict(extra_env or {}), extra_env=dict(extra_env or {}),
cancel_event=cancel_event,
process_registry=process_registry,
tool_registry=tool_registry, tool_registry=tool_registry,
search_runtime=search_runtime, search_runtime=search_runtime,
account_runtime=account_runtime, account_runtime=account_runtime,
+86 -11
View File
@@ -1825,6 +1825,7 @@ def _grep_search(arguments: dict[str, Any], context: ToolExecutionContext) -> st
def _run_python_exec(arguments: dict[str, Any], context: ToolExecutionContext) -> str: def _run_python_exec(arguments: dict[str, Any], context: ToolExecutionContext) -> str:
_ensure_process_execution_allowed(context, 'Python execution') _ensure_process_execution_allowed(context, 'Python execution')
_raise_if_cancelled(context)
code = _optional_string(arguments, 'code') code = _optional_string(arguments, 'code')
script_path = _optional_string(arguments, 'script_path') script_path = _optional_string(arguments, 'script_path')
if bool(code) == bool(script_path): if bool(code) == bool(script_path):
@@ -1856,19 +1857,40 @@ def _run_python_exec(arguments: dict[str, Any], context: ToolExecutionContext) -
command = [interpreter, str(script), *raw_args] command = [interpreter, str(script), *raw_args]
mode = 'script' mode = 'script'
process: subprocess.Popen[str] | None = None
try: try:
completed = subprocess.run( process = subprocess.Popen(
command, command,
cwd=context.root, cwd=context.root,
input=stdin if stdin else None, stdin=subprocess.PIPE if stdin else None,
capture_output=True, stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
text=True, text=True,
timeout=timeout_seconds,
env=_build_subprocess_env(context), env=_build_subprocess_env(context),
) )
except subprocess.TimeoutExpired as exc: _register_child_process(context, process)
stdout = exc.stdout if isinstance(exc.stdout, str) else '' _register_child_process(context, process)
stderr = exc.stderr if isinstance(exc.stderr, str) else '' if stdin and process.stdin is not None:
process.stdin.write(stdin)
process.stdin.close()
process.stdin = None
deadline = time.monotonic() + timeout_seconds
while True:
_raise_if_cancelled(context, process)
remaining = deadline - time.monotonic()
if remaining <= 0:
raise subprocess.TimeoutExpired(command, timeout_seconds)
try:
stdout, stderr = process.communicate(timeout=min(remaining, 0.2))
break
except subprocess.TimeoutExpired:
continue
except subprocess.TimeoutExpired:
if process is not None:
_terminate_process(process)
stdout, stderr = process.communicate()
else:
stdout, stderr = '', ''
payload = [ payload = [
'timed_out=true', 'timed_out=true',
f'timeout_seconds={timeout_seconds:g}', f'timeout_seconds={timeout_seconds:g}',
@@ -1894,11 +1916,17 @@ def _run_python_exec(arguments: dict[str, Any], context: ToolExecutionContext) -
'output_preview': _snapshot_text('\n'.join(payload).strip()), 'output_preview': _snapshot_text('\n'.join(payload).strip()),
}, },
) )
except ToolExecutionError:
raise
finally:
if process is not None:
_unregister_child_process(context, process)
stdout = completed.stdout or '' stdout = stdout or ''
stderr = completed.stderr or '' stderr = stderr or ''
returncode = process.returncode if process is not None else 1
payload = [ payload = [
f'exit_code={completed.returncode}', f'exit_code={returncode}',
f'interpreter={interpreter}', f'interpreter={interpreter}',
f'scratchpad={context.scratchpad_directory or ""}', f'scratchpad={context.scratchpad_directory or ""}',
f'python_env={context.python_env_dir or ""}', f'python_env={context.python_env_dir or ""}',
@@ -1915,7 +1943,7 @@ def _run_python_exec(arguments: dict[str, Any], context: ToolExecutionContext) -
'interpreter': interpreter, 'interpreter': interpreter,
'scratchpad_directory': str(context.scratchpad_directory) if context.scratchpad_directory else '', 'scratchpad_directory': str(context.scratchpad_directory) if context.scratchpad_directory else '',
'python_env_dir': str(context.python_env_dir) if context.python_env_dir else '', 'python_env_dir': str(context.python_env_dir) if context.python_env_dir else '',
'exit_code': completed.returncode, 'exit_code': returncode,
'stdout_preview': _snapshot_text(stdout), 'stdout_preview': _snapshot_text(stdout),
'stderr_preview': _snapshot_text(stderr), 'stderr_preview': _snapshot_text(stderr),
'output_preview': _snapshot_text('\n'.join(payload).strip()), 'output_preview': _snapshot_text('\n'.join(payload).strip()),
@@ -1923,6 +1951,48 @@ def _run_python_exec(arguments: dict[str, Any], context: ToolExecutionContext) -
) )
def _raise_if_cancelled(
context: ToolExecutionContext,
process: subprocess.Popen[Any] | None = None,
) -> None:
cancel_event = context.cancel_event
if cancel_event is not None and cancel_event.is_set():
if process is not None:
_terminate_process(process)
raise ToolExecutionError('Run cancelled by user')
def _register_child_process(
context: ToolExecutionContext,
process: subprocess.Popen[Any],
) -> None:
registry = context.process_registry
if registry is not None:
registry.add(process)
def _unregister_child_process(
context: ToolExecutionContext,
process: subprocess.Popen[Any],
) -> None:
registry = context.process_registry
if registry is not None:
registry.discard(process)
def _terminate_process(process: subprocess.Popen[Any]) -> None:
if process.poll() is not None:
return
try:
process.terminate()
process.wait(timeout=1)
except (OSError, subprocess.TimeoutExpired):
try:
process.kill()
except OSError:
pass
def _run_python_package(arguments: dict[str, Any], context: ToolExecutionContext) -> str: def _run_python_package(arguments: dict[str, Any], context: ToolExecutionContext) -> str:
_ensure_process_execution_allowed(context, 'Python package management') _ensure_process_execution_allowed(context, 'Python package management')
action = _require_string(arguments, 'action') action = _require_string(arguments, 'action')
@@ -3364,6 +3434,10 @@ def _stream_bash(
try: try:
while selector.get_map(): while selector.get_map():
if context.cancel_event is not None and context.cancel_event.is_set():
timeout_error = f'Command cancelled by user: {command}'
_terminate_process(process)
break
remaining = deadline - time.monotonic() remaining = deadline - time.monotonic()
if remaining <= 0: if remaining <= 0:
timeout_error = ( timeout_error = (
@@ -3398,6 +3472,7 @@ def _stream_bash(
stream=stream_name, stream=stream_name,
) )
finally: finally:
_unregister_child_process(context, process)
try: try:
selector.close() selector.close()
except OSError: except OSError:
+30 -1
View File
@@ -104,6 +104,14 @@ class GuiServerTests(unittest.TestCase):
_, state = _build_client(Path(d)) _, state = _build_client(Path(d))
self.assertIs(state.run_lock_for('alice'), state.run_lock_for('alice')) self.assertIs(state.run_lock_for('alice'), state.run_lock_for('alice'))
self.assertIsNot(state.run_lock_for('alice'), state.run_lock_for('bob')) self.assertIsNot(state.run_lock_for('alice'), state.run_lock_for('bob'))
self.assertIs(
state.run_lock_for('alice', 's1'),
state.run_lock_for('alice', 's1'),
)
self.assertIsNot(
state.run_lock_for('alice', 's1'),
state.run_lock_for('alice', 's2'),
)
def test_state_update_rejects_missing_cwd(self) -> None: def test_state_update_rejects_missing_cwd(self) -> None:
with tempfile.TemporaryDirectory() as d: with tempfile.TemporaryDirectory() as d:
@@ -166,7 +174,7 @@ class GuiServerTests(unittest.TestCase):
with tempfile.TemporaryDirectory() as d: with tempfile.TemporaryDirectory() as d:
root = Path(d) root = Path(d)
client, state = _build_client(root) client, state = _build_client(root)
agent = state.agent_for('alice') agent = state.agent_for('alice', 'pending-1')
def fake_run( def fake_run(
prompt: str, prompt: str,
@@ -209,6 +217,27 @@ class GuiServerTests(unittest.TestCase):
self.assertEqual(response.status_code, 200) self.assertEqual(response.status_code, 200)
self.assertEqual(response.json()['session_id'], 'pending-1') self.assertEqual(response.json()['session_id'], 'pending-1')
def test_cancel_run_marks_latest_session_cancelled(self) -> None:
with tempfile.TemporaryDirectory() as d:
client, state = _build_client(Path(d))
record = state.run_manager.start('alice', 'thread-1')
self.assertFalse(record.cancel_event.is_set())
response = client.post(
'/api/runs/cancel',
json={'account_id': 'alice', 'session_id': 'thread-1'},
)
self.assertEqual(response.status_code, 200)
self.assertTrue(response.json()['cancelled'])
self.assertTrue(record.cancel_event.is_set())
latest = client.get(
'/api/runs/latest',
params={'account_id': 'alice', 'session_id': 'thread-1'},
)
self.assertEqual(latest.status_code, 200)
self.assertEqual(latest.json()['status'], 'cancelled')
def test_chat_uses_account_scoped_model_config(self) -> None: def test_chat_uses_account_scoped_model_config(self) -> None:
with tempfile.TemporaryDirectory() as d: with tempfile.TemporaryDirectory() as d:
client, _ = _build_client(Path(d)) client, _ = _build_client(Path(d))