diff --git a/.flocks/flockshub/index.json b/.flocks/flockshub/index.json index 7ac3a1117..3c07841ed 100644 --- a/.flocks/flockshub/index.json +++ b/.flocks/flockshub/index.json @@ -14640,7 +14640,7 @@ "name": "SOC Workspace WebUI", "description": "SOC workspace pages for posture, overview, and alert investigation.", "descriptionCn": "SOC 工作区页面,包含态势、SOC 总览和告警调查。", - "version": "1.0.0", + "version": "1.1.4", "category": "workflow-automation", "tags": [ "siem", diff --git a/.flocks/flockshub/plugins/webuis/soc_ui/access/soc_alerts_operations.py b/.flocks/flockshub/plugins/webuis/soc_ui/access/soc_alerts_operations.py index 463212ca4..d8fd8b430 100644 --- a/.flocks/flockshub/plugins/webuis/soc_ui/access/soc_alerts_operations.py +++ b/.flocks/flockshub/plugins/webuis/soc_ui/access/soc_alerts_operations.py @@ -84,6 +84,8 @@ "_source_type", "net_type", "direction", + "threat_severity", + "threat_level", "threat_name", "threat_type", "threat_phase", @@ -228,7 +230,7 @@ def _incident_from_row(row: InternalDataRow) -> dict[str, Any]: record = row.raw record_id = _record_id(record) observed_at = _observed_at(record) - threat_name = _first_text(record, "threat_name", "_threat_type", "report_title") or "SOC alert" + threat_name = _first_text(record, "threat_name", "report_title") or "SOC alert" threat_msg = _first_text(record, "threat_msg", "report_title", "threat_type") verdict = _verdict_bucket(record) table_cells = _table_cells(record) @@ -425,7 +427,7 @@ def _verdict_bucket(record: dict[str, Any]) -> str: return "success" raw = " ".join( _text(record.get(key)).lower() - for key in ("attack_verdict", "threat_result", "risk_level", "threat_level") + for key in ("attack_verdict", "threat_result") ) if any(marker in raw for marker in ("success", "attack_success", "succeeded")): return "success" diff --git a/.flocks/flockshub/plugins/webuis/soc_ui/manifest.json b/.flocks/flockshub/plugins/webuis/soc_ui/manifest.json index c0f50cefc..21a8d74fa 100644 --- a/.flocks/flockshub/plugins/webuis/soc_ui/manifest.json +++ b/.flocks/flockshub/plugins/webuis/soc_ui/manifest.json @@ -5,7 +5,7 @@ "name": "SOC Workspace WebUI", "description": "SOC workspace pages for posture, overview, and alert investigation.", "descriptionCn": "SOC 工作区页面,包含态势、SOC 总览和告警调查。", - "version": "1.0.0", + "version": "1.1.4", "author": "Flocks Team", "license": "MIT", "homepage": "", diff --git a/.flocks/flockshub/plugins/webuis/soc_ui/soc_alerts/src/filterValues.ts b/.flocks/flockshub/plugins/webuis/soc_ui/soc_alerts/src/filterValues.ts new file mode 100644 index 000000000..9ae0a4f8c --- /dev/null +++ b/.flocks/flockshub/plugins/webuis/soc_ui/soc_alerts/src/filterValues.ts @@ -0,0 +1,94 @@ +type Translate = (text: string) => string; + +const identityTr: Translate = (text) => text; + +const FILTER_VALUE_TEXT: Record> = { + direction: { + in: '入站', + inbound: '入站', + ingress: '入站', + out: '出站', + outbound: '出站', + egress: '出站', + lateral: '横向', + }, + threat_severity: { + critical: '严重', + severe: '严重', + high: '高危', + medium: '中危', + low: '低危', + info: '信息', + informational: '信息', + }, + threat_level: { + critical: '严重', + severe: '严重', + high: '高危', + medium: '中危', + low: '低危', + info: '信息', + informational: '信息', + }, + threat_phase: { + recon: '侦察', + reconnaissance: '侦察', + access: '访问', + initial_access: '初始访问', + execution: '执行', + persistence: '持久化', + privilege_escalation: '权限提升', + defense_evasion: '防御规避', + credential_access: '凭据访问', + discovery: '发现', + lateral_movement: '横向移动', + collection: '收集', + command_and_control: '命令与控制', + exfiltration: '数据渗出', + impact: '影响', + exploit: '利用', + exploitation: '利用', + }, + threat_result: { + attack_success: '攻击成功', + success: '成功', + succeeded: '成功', + attack_failed: '攻击失败', + failed: '失败', + blocked: '已阻断', + detected: '已检测', + attack: '攻击行为', + benign: '安全', + safe: '安全', + normal: '正常', + unknown: '未知', + }, +}; + +function normalized(value: string) { + return value.trim().toLowerCase(); +} + +function localizedValueText(key: string, value: string) { + return FILTER_VALUE_TEXT[key]?.[normalized(value)] || ''; +} + +export function filterOptionText(key: string, value: string, tr: Translate = identityTr) { + if (key === 'rsp_status_code') return value || tr('未知响应'); + if (key === '_source_type' || key === 'net_type') return value || 'unknown'; + const localizedText = localizedValueText(key, value); + return localizedText ? tr(localizedText) : value || tr('空值'); +} + +export function matchesFilterOptionSearch( + key: string, + value: string, + search: string, + tr: Translate = identityTr, +) { + const query = normalized(search); + if (!query) return true; + const localizedText = localizedValueText(key, value); + return [value, localizedText, localizedText ? tr(localizedText) : ''] + .some((candidate) => normalized(candidate).includes(query)); +} diff --git a/.flocks/flockshub/plugins/webuis/soc_ui/soc_alerts/src/index.tsx b/.flocks/flockshub/plugins/webuis/soc_ui/soc_alerts/src/index.tsx index ab12870ab..c681d4aa9 100644 --- a/.flocks/flockshub/plugins/webuis/soc_ui/soc_alerts/src/index.tsx +++ b/.flocks/flockshub/plugins/webuis/soc_ui/soc_alerts/src/index.tsx @@ -1,8 +1,9 @@ import { Fragment, useCallback, useEffect, useMemo, useRef, useState } from 'react'; import { api } from '@flocks/webui-contract-sdk'; +import { filterOptionText, matchesFilterOptionSearch } from './filterValues'; type Tone = 'red' | 'orange' | 'blue' | 'green' | 'purple' | 'slate'; -type FilterKey = '_source_type' | 'net_type' | 'direction' | 'threat_name' | 'threat_type' | 'threat_phase' | 'threat_result' | 'rsp_status_code' | 'sip' | 'dport' | 'dip' | 'req_host' | 'threat_rule_id'; +type FilterKey = '_source_type' | 'net_type' | 'direction' | 'threat_severity' | 'threat_level' | 'threat_name' | 'threat_type' | 'threat_phase' | 'threat_result' | 'rsp_status_code' | 'sip' | 'dport' | 'dip' | 'req_host' | 'threat_rule_id'; type TimeRangeKey = '15m' | '1h' | '2h' | '24h' | 'today' | '7d' | '30d'; type TimeFilterMode = 'relative' | 'custom'; type TimePanelTab = 'auto' | 'custom'; @@ -201,6 +202,8 @@ const EN_TEXT: Record = { '数据源': 'Data Source', '协议类型': 'Protocol', '流量方向': 'Traffic Direction', + '严重等级': 'Severity', + '威胁级别': 'Threat Level', '威胁名称': 'Threat Name', '威胁类型': 'Threat Type', '攻击阶段': 'Attack Stage', @@ -261,6 +264,34 @@ const EN_TEXT: Record = { '展开趋势图': 'Expand timeline', '攻击成功': 'Attack Success', '攻击失败': 'Attack Failed', + '入站': 'Inbound', + '出站': 'Outbound', + '横向': 'Lateral', + '严重': 'Critical', + '高危': 'High', + '中危': 'Medium', + '低危': 'Low', + '信息': 'Informational', + '侦察': 'Reconnaissance', + '初始访问': 'Initial Access', + '执行': 'Execution', + '持久化': 'Persistence', + '权限提升': 'Privilege Escalation', + '防御规避': 'Defense Evasion', + '凭据访问': 'Credential Access', + '发现': 'Discovery', + '横向移动': 'Lateral Movement', + '收集': 'Collection', + '命令与控制': 'Command and Control', + '数据渗出': 'Exfiltration', + '影响': 'Impact', + '利用': 'Exploitation', + '成功': 'Success', + '失败': 'Failed', + '已阻断': 'Blocked', + '已检测': 'Detected', + '安全': 'Benign', + '正常': 'Normal', '未知': 'Unknown', '暂无可展示的告警数据。': 'No alerts to display.', '显示 {start}-{end} / {total} 条,每页 {pageSize} 条': 'Showing {start}-{end} / {total}, {pageSize} per page', @@ -356,6 +387,8 @@ const BASE_FILTER_CONFIGS: FilterConfig[] = [ { key: '_source_type', label: '数据源' }, { key: 'net_type', label: '协议类型' }, { key: 'direction', label: '流量方向' }, + { key: 'threat_severity', label: '严重等级' }, + { key: 'threat_level', label: '威胁级别' }, { key: 'threat_name', label: '威胁名称' }, ]; @@ -377,6 +410,8 @@ const DEFAULT_FILTER_VALUES: Record = { _source_type: ['tdp'], net_type: ['http'], direction: [], + threat_severity: [], + threat_level: [], threat_name: [], threat_type: [], threat_phase: [], @@ -990,9 +1025,7 @@ function readFilterValue(incident: IncidentCluster, key: FilterKey) { } function optionText(key: FilterKey, value: string, tr: Translate = identityTr) { - if (key === 'rsp_status_code') return value || tr('未知响应'); - if (key === '_source_type' || key === 'net_type') return value || 'unknown'; - return value || tr('空值'); + return filterOptionText(key, value, tr); } function optionLabel(key: FilterKey, value: string | string[], tr: Translate = identityTr) { @@ -1118,7 +1151,7 @@ function FilterDropdown({ }, [open, value]); const choices = options.filter((option) => option && option !== ALL_FILTER_VALUE); - const visibleChoices = choices.filter((choice) => optionText(config.key, choice, tr).toLowerCase().includes(search.trim().toLowerCase())); + const visibleChoices = choices.filter((choice) => matchesFilterOptionSearch(config.key, choice, search, tr)); const selected = new Set(draft.map(normalized)); const toggleChoice = (choice: string) => { diff --git a/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/api/handlers.py b/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/api/handlers.py index 575d978a0..fb0fe6bc7 100644 --- a/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/api/handlers.py +++ b/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/api/handlers.py @@ -21,7 +21,7 @@ FACTS_TABLE = "soc_dashboard_alert_facts" ACTIVITY_TABLE = "soc_dashboard_activity" META_TABLE = "soc_dashboard_meta" -SCHEMA_VERSION = "2" +SCHEMA_VERSION = "3" ACTIVITY_DEFAULT_LIMIT = 20 ACTIVITY_MAX_LIMIT = 50 ACTIVITY_WINDOW_MS = 3000 @@ -66,6 +66,7 @@ class _RecordSource: "event_time", "source_type", "threat_name", + "threat_type", "is_duplicate", "phase", "direction", @@ -97,6 +98,7 @@ def _fact_expressions(prefix): source_type_fallback = _json_value(prefix, "source_type") threat_name = _json_value(prefix, "threat_name") threat_type = _json_value(prefix, "_threat_type") + threat_type_fallback = _json_value(prefix, "threat_type") phase = _json_value(prefix, "threat_phase") attack_phase = _json_value(prefix, "attack_phase") kill_chain = _json_value(prefix, "kill_chain_phase") @@ -108,7 +110,6 @@ def _fact_expressions(prefix): app_protocol = _json_value(prefix, "net_app_proto") protocol_fallback = _json_value(prefix, "protocol") severity = _json_value(prefix, "threat_severity") - threat_level = _json_value(prefix, "threat_level") risk_level = _json_value(prefix, "risk_level") response = _json_value(prefix, "rsp_status_code") status_code = _json_value(prefix, "status_code") @@ -134,15 +135,15 @@ def _fact_expressions(prefix): f"{prefix}.event_time", f"COALESCE(NULLIF({prefix}.source_type, ''), NULLIF({source_type}, ''), " f"NULLIF({source_type_fallback}, ''), 'unknown')", - f"COALESCE(NULLIF({prefix}.threat_name, ''), NULLIF({threat_name}, ''), " - f"NULLIF({threat_type}, ''), 'unknown')", + f"COALESCE(NULLIF({prefix}.threat_name, ''), NULLIF({threat_name}, ''), 'unknown')", + f"COALESCE(NULLIF({threat_type}, ''), NULLIF({threat_type_fallback}, ''), 'unknown')", f"COALESCE({prefix}.is_duplicate, 0)", f"COALESCE(NULLIF({phase}, ''), NULLIF({attack_phase}, ''), NULLIF({kill_chain}, ''), 'unknown')", f"COALESCE(NULLIF({direction}, ''), NULLIF({traffic_direction}, ''), 'unknown')", f"COALESCE(NULLIF({result}, ''), NULLIF({verdict}, ''), 'unknown')", f"COALESCE(NULLIF({protocol}, ''), NULLIF({app_protocol}, ''), " f"NULLIF({protocol_fallback}, ''), 'unknown')", - f"COALESCE(NULLIF({severity}, ''), NULLIF({threat_level}, ''), NULLIF({risk_level}, ''), 'unknown')", + f"COALESCE(NULLIF({severity}, ''), 'unknown')", f"COALESCE(NULLIF({response}, ''), NULLIF({status_code}, ''), 'unknown')", f"COALESCE(NULLIF({destination_port}, ''), NULLIF({destination_port_fallback}, ''), " f"NULLIF({destination_port_legacy}, ''), 'unknown')", @@ -151,7 +152,7 @@ def _fact_expressions(prefix): f"COALESCE({triage_status}, '')", f"COALESCE({triage_source}, '')", f"COALESCE({verdict}, 'unknown')", - f"COALESCE({risk_level}, {threat_level}, {severity}, 'unknown')", + f"COALESCE(NULLIF({risk_level}, ''), 'unknown')", f"COALESCE(CAST({triage_ms} AS INTEGER), 0)", f"CASE WHEN {attack_success} IN (1, '1', 'true') THEN 1 ELSE 0 END", ) @@ -360,6 +361,7 @@ def _ensure_sqlite_schema(): event_time INTEGER, source_type TEXT, threat_name TEXT, + threat_type TEXT, is_duplicate INTEGER NOT NULL DEFAULT 0, phase TEXT, direction TEXT, @@ -911,10 +913,7 @@ def _get_workflow_recent_events( raw_count = metrics["rawCount"] unique_count = metrics["uniqueCount"] threat_name = str( - preview.get("threat_name") - or preview.get("_threat_type") - or preview.get("threat_type") - or f"降噪批次 · 原始 {raw_count} 条" + preview.get("threat_name") or f"降噪批次 · 原始 {raw_count} 条" ) events.append( { @@ -1374,12 +1373,12 @@ def _activity_event(row): "alert": { "id": _first_activity_text(record, "id", "record_id", "uuid", "event_id", "dedup_key"), "sourceType": _first_activity_text(record, "_source_type", "source_type", "device_type"), - "threatName": _first_activity_text(record, "threat_name", "_threat_type", "threat_type") or "未知告警", + "threatName": _first_activity_text(record, "threat_name") or "未知告警", "srcIp": _first_activity_text(record, "sip", "src_ip", "source_ip"), "dstIp": _first_activity_text(record, "dip", "dst_ip", "destination_ip"), "requestUri": _first_activity_text(record, "req_http_url", "uri", "url"), "threatPhase": _first_activity_text(record, "threat_phase"), - "threatType": _first_activity_text(record, "threat_type", "_threat_type"), + "threatType": _first_activity_text(record, "_threat_type", "threat_type"), }, } if stage == "denoise": @@ -1396,7 +1395,8 @@ def _activity_event(row): "durationMs": _safe_int(record.get("triage_ms")), "verdict": verdict, "verdictLabel": RESULT_LABELS.get(verdict, "待确认"), - "riskLevel": _first_activity_text(record, "risk_level", "threat_level"), + "threatSeverity": _first_activity_text(record, "threat_severity"), + "riskLevel": _first_activity_text(record, "risk_level"), "reportTitle": _first_activity_text(record, "report_title"), "hasReport": bool(record.get("triage_report")), } @@ -1623,7 +1623,14 @@ def _get_stats(params): {"key": "benign", "label": "良性", "value": triage["benign"], "color": "#58a6ff"}, {"key": "unknown", "label": "未知", "value": triage["unknown"], "color": "#9b8cff"}, ], - "topThreats": _counter_items(triage["threatCounter"] or denoise["threatCounter"], 14), + "topThreatTypes": _counter_items( + triage["threatTypeCounter"] or denoise["threatTypeCounter"], + 14, + ), + "severityLevels": _counter_items( + _profile_counter(denoise, triage, "severityCounter"), + 8, + ), "riskLevels": _counter_items(triage["riskCounter"], 5), "timeline": { "labels": denoise.get("_timelineLabels") @@ -1844,7 +1851,7 @@ def _read_denoise(paths, workflow_call_count: int = 0): parse_errors = 0 headers = [] source_counter = Counter() - threat_counter = Counter() + threat_type_counter = Counter() profile_counters = _new_profile_counters() event_start = None event_end = None @@ -1865,7 +1872,7 @@ def _read_denoise(paths, workflow_call_count: int = 0): if obj.get("is_duplicate") is True: file_duplicates += 1 source_counter[_norm(obj.get("_source_type") or obj.get("source_type") or obj.get("device_type"))] += 1 - threat_counter[_norm(obj.get("_threat_type") or obj.get("threat_name") or obj.get("threat_type"))] += 1 + threat_type_counter[_norm(obj.get("_threat_type") or obj.get("threat_type"))] += 1 _update_profile_counters(obj, profile_counters) event_start, event_end = _merge_record_time(event_start, event_end, obj) total_raw += file_raw @@ -1889,7 +1896,7 @@ def _read_denoise(paths, workflow_call_count: int = 0): "eventStart": _format_event_time(event_start), "eventEnd": _format_event_time(event_end), "sourceCounter": source_counter, - "threatCounter": threat_counter, + "threatTypeCounter": threat_type_counter, **profile_counters, "seriesRaw": series_raw, "seriesUnique": series_unique, @@ -1928,7 +1935,7 @@ def _read_sqlite_denoise(paths, workflow_call_count): f"WHERE {where_clause} GROUP BY \"source_type\"", query_params, ).fetchall() - profile_counters, threat_counter = _sqlite_detail_counters( + profile_counters, threat_type_counter = _sqlite_detail_counters( conn, settings, where_clause, @@ -1969,7 +1976,7 @@ def _read_sqlite_denoise(paths, workflow_call_count): "eventStart": _format_event_time(min(event_values) if event_values else None), "eventEnd": _format_event_time(max(event_values) if event_values else None), "sourceCounter": Counter({_norm(key): _safe_int(value) for key, value in source_rows}), - "threatCounter": threat_counter, + "threatTypeCounter": threat_type_counter, **profile_counters, "seriesRaw": timeline["raw"], "seriesUnique": timeline["unique"], @@ -2025,8 +2032,12 @@ def _sqlite_timeline(conn, settings, where_clause, query_params, dates, start_ti f"COUNT(*) AS raw_count, " f"COALESCE(SUM(CASE WHEN is_duplicate = 0 THEN 1 ELSE 0 END), 0) AS unique_count, " f"COALESCE(SUM(has_triage), 0) AS triage_count, " - f"COALESCE(SUM(CASE WHEN has_triage = 1 AND LOWER(verdict) IN " - f"('attack_success', 'attack', 'attack_failed') THEN 1 ELSE 0 END), 0) AS attack_count " + f"COALESCE(SUM(CASE WHEN has_triage = 1 " + f"AND LOWER(triage_status) NOT IN ('failed', 'error') " + f"AND (LOWER(verdict) IN ('attack_success', 'attack', 'attack_failed') " + f"OR (attack_success = 1 AND LOWER(verdict) NOT IN " + f"('attack_success', 'attack', 'attack_failed', 'benign'))) " + f"THEN 1 ELSE 0 END), 0) AS attack_count " f"FROM {settings['facts_table']} WHERE {where_clause} AND event_time IS NOT NULL " f"GROUP BY bucket_index ORDER BY bucket_index", (bucket_start, bucket_seconds, *query_params), @@ -2079,12 +2090,12 @@ def _sqlite_detail_counters(conn, settings, where_clause, query_params): ): return ( {key: Counter(value) for key, value in cached["profileCounters"].items()}, - Counter(cached["threatCounter"]), + Counter(cached["threatTypeCounter"]), ) rows = conn.execute( f"SELECT phase, direction, result, protocol, severity, response_code, " - f"port, threat_name, COUNT(*) AS profile_count FROM {table} " + f"port, threat_type, COUNT(*) AS profile_count FROM {table} " f"WHERE {where_clause} " f"GROUP BY 1, 2, 3, 4, 5, 6, 7, 8", query_params, @@ -2098,7 +2109,7 @@ def _sqlite_detail_counters(conn, settings, where_clause, query_params): "severityCounter", "responseCounter", ) - threat_counter = Counter() + threat_type_counter = Counter() for row in rows: count = _safe_int(row[8]) for index, key in enumerate(profile_keys): @@ -2106,19 +2117,19 @@ def _sqlite_detail_counters(conn, settings, where_clause, query_params): port_value = row[6] port = str(_safe_int(port_value)) if _safe_int(port_value) > 0 else _norm(port_value) profile_counters["portCounter"][port] += count - threat_counter[_norm(row[7])] += count + threat_type_counter[_norm(row[7])] += count with _cache_lock: _denoise_detail_cache[cache_key] = { "lastRowId": latest_row_id, "lastTriagePersistedAt": latest_triage_at, "updatedAt": time.monotonic(), "profileCounters": {key: Counter(value) for key, value in profile_counters.items()}, - "threatCounter": Counter(threat_counter), + "threatTypeCounter": Counter(threat_type_counter), } _denoise_detail_cache.move_to_end(cache_key) while len(_denoise_detail_cache) > _DENOISE_DETAIL_CACHE_MAX: _denoise_detail_cache.popitem(last=False) - return profile_counters, threat_counter + return profile_counters, threat_type_counter def _read_sqlite_triage(paths): @@ -2135,24 +2146,37 @@ def _read_sqlite_triage(paths): where_clause += " AND event_time BETWEEN ? AND ?" query_params.extend((start_time, end_time)) triage_where = f"{where_clause} AND has_triage = 1" + cache_condition = "(LOWER(triage_source) = 'cache' OR LOWER(triage_status) = 'cached')" + follower_condition = ( + "(LOWER(triage_source) IN ('follower', 'followers', 'follower_reused') " + "OR LOWER(triage_status) = 'follower_reused')" + ) + failed_condition = "LOWER(triage_status) IN ('failed', 'error')" + resolved_condition = f"NOT ({failed_condition})" + new_triage_condition = ( + f"NOT {cache_condition} AND NOT {follower_condition} AND NOT ({failed_condition})" + ) try: with sqlite3.connect(settings["db_path"]) as conn: row = conn.execute( f"SELECT COUNT(*), " - f"COALESCE(SUM(CASE WHEN LOWER(triage_source) = 'cache' " - f"OR LOWER(triage_status) = 'cached' THEN 1 ELSE 0 END), 0), " - f"COALESCE(SUM(CASE WHEN LOWER(triage_source) IN " - f"('follower', 'followers', 'follower_reused') " - f"OR LOWER(triage_status) = 'follower_reused' THEN 1 ELSE 0 END), 0), " - f"COALESCE(SUM(CASE WHEN LOWER(triage_status) IN ('failed', 'error') " + f"COALESCE(SUM(CASE WHEN {cache_condition} THEN 1 ELSE 0 END), 0), " + f"COALESCE(SUM(CASE WHEN {follower_condition} THEN 1 ELSE 0 END), 0), " + f"COALESCE(SUM(CASE WHEN {failed_condition} " + f"THEN 1 ELSE 0 END), 0), " + f"COALESCE(SUM(CASE WHEN {new_triage_condition} " + f"THEN 1 ELSE 0 END), 0), " + f"COALESCE(SUM(CASE WHEN {resolved_condition} AND " + f"(LOWER(verdict) = 'attack_success' OR (attack_success = 1 AND LOWER(verdict) NOT IN " + f"('attack_success', 'attack', 'attack_failed', 'benign'))) THEN 1 ELSE 0 END), 0), " + f"COALESCE(SUM(CASE WHEN {resolved_condition} AND LOWER(verdict) = 'attack' " + f"THEN 1 ELSE 0 END), 0), " + f"COALESCE(SUM(CASE WHEN {resolved_condition} AND LOWER(verdict) = 'attack_failed' " + f"THEN 1 ELSE 0 END), 0), " + f"COALESCE(SUM(CASE WHEN {resolved_condition} AND LOWER(verdict) = 'benign' " f"THEN 1 ELSE 0 END), 0), " - f"COALESCE(SUM(CASE WHEN LOWER(verdict) = 'attack_success' THEN 1 ELSE 0 END), 0), " - f"COALESCE(SUM(CASE WHEN LOWER(verdict) = 'attack' THEN 1 ELSE 0 END), 0), " - f"COALESCE(SUM(CASE WHEN LOWER(verdict) = 'attack_failed' THEN 1 ELSE 0 END), 0), " - f"COALESCE(SUM(CASE WHEN LOWER(verdict) = 'benign' THEN 1 ELSE 0 END), 0), " - f"COALESCE(SUM(CASE WHEN LOWER(verdict) NOT IN " - f"('attack_success', 'attack', 'attack_failed', 'benign') THEN 1 ELSE 0 END), 0), " - f"COALESCE(SUM(CASE WHEN attack_success = 1 AND LOWER(verdict) <> 'attack_success' " + f"COALESCE(SUM(CASE WHEN {resolved_condition} AND LOWER(verdict) NOT IN " + f"('attack_success', 'attack', 'attack_failed', 'benign') AND attack_success <> 1 " f"THEN 1 ELSE 0 END), 0), MIN(event_time), MAX(event_time), " f"COALESCE(ROUND(AVG(CASE WHEN triage_ms > 0 THEN triage_ms END)), 0) " f"FROM {settings['facts_table']} WHERE {triage_where}", @@ -2163,9 +2187,9 @@ def _read_sqlite_triage(paths): f"WHERE {triage_where} GROUP BY source_type", query_params, ).fetchall() - threat_rows = conn.execute( - f"SELECT threat_name, COUNT(*) FROM {settings['facts_table']} " - f"WHERE {triage_where} GROUP BY threat_name", + threat_type_rows = conn.execute( + f"SELECT threat_type, COUNT(*) FROM {settings['facts_table']} " + f"WHERE {triage_where} GROUP BY threat_type", query_params, ).fetchall() risk_rows = conn.execute( @@ -2191,11 +2215,12 @@ def _read_sqlite_triage(paths): cache_hit = _safe_int(row[1]) followers_reused = _safe_int(row[2]) triage_failed = _safe_int(row[3]) - attack_success = _safe_int(row[4]) + _safe_int(row[9]) - attack = _safe_int(row[5]) - attack_failed = _safe_int(row[6]) - benign = _safe_int(row[7]) - unknown = _safe_int(row[8]) + new_triaged = _safe_int(row[4]) + attack_success = _safe_int(row[5]) + attack = _safe_int(row[6]) + attack_failed = _safe_int(row[7]) + benign = _safe_int(row[8]) + unknown = _safe_int(row[9]) attack_total = attack_success + attack + attack_failed avg_triage_ms = _safe_int(row[12]) profile_counters = _new_profile_counters() @@ -2217,7 +2242,7 @@ def _read_sqlite_triage(paths): return { "totalRecords": total_records, "batchTotal": 0, - "newTriaged": max(total_records - cache_hit - followers_reused - triage_failed, 0), + "newTriaged": new_triaged, "cacheHit": cache_hit, "triageFailed": triage_failed, "followersReused": followers_reused, @@ -2238,7 +2263,9 @@ def _read_sqlite_triage(paths): "eventStart": _format_event_time(_parse_event_time(row[10])), "eventEnd": _format_event_time(_parse_event_time(row[11])), "sourceCounter": Counter({_norm(key): _safe_int(value) for key, value in source_rows}), - "threatCounter": Counter({_norm(key): _safe_int(value) for key, value in threat_rows}), + "threatTypeCounter": Counter( + {_norm(key): _safe_int(value) for key, value in threat_type_rows} + ), "riskCounter": Counter({_norm(key): _safe_int(value) for key, value in risk_rows}), "statusCounter": Counter({_norm(key): _safe_int(value) for key, value in status_rows}), **profile_counters, @@ -2257,7 +2284,7 @@ def _read_triage(paths): headers = [] verdict_counter = Counter() source_counter = Counter() - threat_counter = Counter() + threat_type_counter = Counter() risk_counter = Counter() status_counter = Counter() profile_counters = _new_profile_counters() @@ -2268,7 +2295,6 @@ def _read_triage(paths): fallback_cache = 0 fallback_failed = 0 fallback_followers = 0 - extra_success = 0 series_total = [] series_attack = [] triage_ms_total = 0 @@ -2298,16 +2324,10 @@ def _read_triage(paths): verdict = _norm(obj.get("attack_verdict") or "unknown") if verdict not in {"attack_success", "attack", "attack_failed", "benign", "unknown"}: verdict = "unknown" - verdict_counter[verdict] += 1 - if obj.get("attack_success") is True and verdict != "attack_success": - extra_success += 1 - if verdict in {"attack_success", "attack", "attack_failed"}: - file_attack += 1 - source = _norm(obj.get("_source_type") or obj.get("source_type") or obj.get("device_type")) source_counter[source] += 1 - threat_counter[_norm(obj.get("_threat_type") or obj.get("threat_name") or obj.get("threat_type"))] += 1 - risk_counter[_norm(obj.get("risk_level") or obj.get("threat_level") or obj.get("threat_severity"))] += 1 + threat_type_counter[_norm(obj.get("_threat_type") or obj.get("threat_type"))] += 1 + risk_counter[_norm(obj.get("risk_level"))] += 1 triage_ms = _safe_int(obj.get("triage_ms")) if triage_ms > 0: triage_ms_total += triage_ms @@ -2317,14 +2337,22 @@ def _read_triage(paths): triage_source = _norm(obj.get("triage_source")) triage_status = _norm(obj.get("triage_status")) status_counter[triage_status or triage_source] += 1 + triage_failed = triage_status in {"failed", "error"} + + if not triage_failed: + if verdict == "unknown" and obj.get("attack_success") is True: + verdict = "attack_success" + verdict_counter[verdict] += 1 + if verdict in {"attack_success", "attack", "attack_failed"}: + file_attack += 1 + else: + fallback_failed += 1 if triage_source == "cache" or triage_status == "cached": fallback_cache += 1 elif triage_source in {"follower", "followers", "follower_reused"} or triage_status == "follower_reused": fallback_followers += 1 - elif triage_status in {"failed", "error"}: - fallback_failed += 1 - else: + elif not triage_failed: fallback_new += 1 series_total.append(file_total) @@ -2339,7 +2367,7 @@ def _read_triage(paths): triage_failed = fallback_failed followers_reused = fallback_followers - attack_success = verdict_counter["attack_success"] + extra_success + attack_success = verdict_counter["attack_success"] attack = verdict_counter["attack"] attack_failed = verdict_counter["attack_failed"] attack_total = attack_success + attack + attack_failed @@ -2373,7 +2401,7 @@ def _read_triage(paths): "eventStart": _format_event_time(event_start), "eventEnd": _format_event_time(event_end), "sourceCounter": source_counter, - "threatCounter": threat_counter, + "threatTypeCounter": threat_type_counter, "riskCounter": risk_counter, "statusCounter": status_counter, **profile_counters, @@ -2399,7 +2427,7 @@ def _update_profile_counters(obj, counters): counters["directionCounter"][_norm(obj.get("direction") or obj.get("traffic_direction"))] += 1 counters["resultCounter"][_norm(obj.get("threat_result") or obj.get("attack_verdict"))] += 1 counters["protocolCounter"][_norm(obj.get("net_type") or obj.get("net_app_proto") or obj.get("protocol"))] += 1 - counters["severityCounter"][_norm(obj.get("threat_severity") or obj.get("threat_level") or obj.get("risk_level"))] += 1 + counters["severityCounter"][_norm(obj.get("threat_severity"))] += 1 counters["responseCounter"][_norm(obj.get("rsp_status_code") or obj.get("status_code"))] += 1 port_value = obj.get("dport") or obj.get("dst_port") or obj.get("destination_port") diff --git a/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/src/Page.tsx b/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/src/Page.tsx index 08092a6c4..9ebc10b18 100644 --- a/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/src/Page.tsx +++ b/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/src/Page.tsx @@ -1,3 +1,5 @@ +import { severityKey, severityRows } from './severityValues'; + function getSdk() { const sdk = globalThis.__FLOCKS_WEBUI_CONTRACT_SDK__; if (!sdk || !sdk.React || !sdk.api) { @@ -74,7 +76,8 @@ const EMPTY_STATS = { closedLoop: { autoClosed: 0, resolved: 0, manualDecision: 0, pending: 0, resolutionRate: 0 }, verdicts: [], attackProfile: [], - topThreats: [], + topThreatTypes: [], + severityLevels: [], riskLevels: [], timeline: { denoiseRaw: [], denoiseUnique: [], triageTotal: [], triageAttack: [] }, }; @@ -620,7 +623,7 @@ function CenterColumn({ stats, activity }) { } function RightColumn({ stats }) { - const threatRows = stats.topThreats || []; + const threatRows = stats.topThreatTypes || []; const threatRankLabel = threatRows.length ? `Top ${threatRows.length}` : '排行'; const threatTotal = threatRows.reduce((sum, item) => sum + (item.value || 0), 0); const threatBase = stats.triage.totalRecords || stats.triage.attackTotal || threatTotal; @@ -631,7 +634,7 @@ function RightColumn({ stats }) { h('div', { className: 'loop-diagram', key: 'diagram' }, [ h('div', { className: 'loop-node primary', key: 'valid' }, [h('b', { key: 'v' }, compactNumber(stats.triage.attackTotal)), h('span', { key: 'l' }, '有效事件')]), h('div', { className: 'loop-node', key: 'auto' }, [h('b', { key: 'v' }, compactNumber(stats.closedLoop.autoClosed)), h('span', { key: 'l' }, '自动闭环')]), - h('div', { className: 'loop-node warn', key: 'manual' }, [h('b', { key: 'v' }, compactNumber(stats.closedLoop.manualDecision)), h('span', { key: 'l' }, '人工决策')]), + h('div', { className: 'loop-node warn', key: 'manual' }, [h('b', { key: 'v' }, compactNumber(stats.closedLoop.manualDecision)), h('span', { key: 'l' }, 'AI转人工')]), h('div', { className: 'loop-node hot', key: 'pending' }, [h('b', { key: 'v' }, compactNumber(stats.closedLoop.pending)), h('span', { key: 'l' }, '待处理')]), ]), h(Gauge, { label: '闭环率', value: stats.closedLoop.resolutionRate, color: '#2ee6a6', key: 'gauge' }), @@ -864,7 +867,11 @@ function AiCore({ stats, activity }) { h('small', { key: 'label' }, coreLabel), activeEvent ? h('div', { className: 'ai-operation-window', key: `operation-${activeEvent.eventId}` }, [ h('div', { className: 'ai-operation-track', key: 'track' }, [...operations, operations[0]].map((operation, index) => h('span', { key: `${operation}-${index}` }, operation))), - ]) : h('div', { className: 'ai-operation-idle', key: 'operation-idle' }, '等待新的处理任务'), + ]) : h('div', { + className: 'ai-operation-idle', + title: 'AI 新完成且未使用缓存、复用或失败的研判数量', + key: 'operation-idle', + }, `AI新增研判 ${compactNumber(stats.triage.newTriaged)}`), ]), activeEvent ? h('div', { className: 'ai-evidence-field', key: `evidence-${activeEvent.eventId}` }, evidenceItems.map((item, index) => h('div', { className: `ai-evidence-card evidence-${index + 1}`, @@ -1375,26 +1382,12 @@ function CommandActivityLane({ kind, lane }) { ]); } -function severityRows(stats) { - return [ - { key: 'critical', label: '严重', value: stats.triage.attackSuccess || 0, tone: 'critical' }, - { key: 'high', label: '高危', value: stats.triage.attack || 0, tone: 'high' }, - { key: 'medium', label: '中危', value: stats.triage.attackFailed || 0, tone: 'medium' }, - { key: 'low', label: '低危', value: stats.triage.benign || 0, tone: 'low' }, - ]; -} - function CommandGraph({ stats, activity }) { const denoiseActive = Boolean(activity.denoise.current); const triageActive = Boolean(activity.triage.current); const severityToneFor = (event) => { if (!event) return ''; - if (event.result?.verdict === 'attack_success') return 'critical'; - const risk = String(event.result?.riskLevel || '').toLowerCase(); - if (risk === 'high') return 'high'; - if (risk === 'medium') return 'medium'; - if (risk === 'low' || event.result?.verdict === 'benign') return 'low'; - return ''; + return severityKey(event.result?.threatSeverity); }; const activeSeverityTone = severityToneFor(activity.triage.current); const recentSeverityTone = severityToneFor(activity.triage.last); @@ -1417,9 +1410,27 @@ function CommandGraph({ stats, activity }) { h(AiCore, { stats, activity, key: 'sphere' }), ]), h('div', { className: 'outcome-stack', key: 'outcomes' }, [ - h('div', { title: 'AI 新完成且未使用缓存、复用或失败的研判数量', key: 'auto' }, [h(AnimatedNumber, { tag: 'b', value: stats.triage.newTriaged, key: 'value' }), h('span', { key: 'label' }, 'AI自主研判')]), - h('div', { className: cx('primary', triageActive && 'processing'), title: '研判结论为攻击成功、攻击行为或攻击失败的事件数量', key: 'events' }, [h(AnimatedNumber, { tag: 'b', value: stats.triage.attackTotal, key: 'value' }), h('span', { key: 'label' }, triageActive ? '结果生成中' : '安全事件')]), - h('div', { key: 'manual' }, [h(AnimatedNumber, { tag: 'b', value: stats.closedLoop.manualDecision, key: 'value' }), h('span', { key: 'label' }, '人工研判')]), + h('div', { className: cx('primary', triageActive && 'processing'), title: 'AI 研判结论为攻击成功、攻击行为或攻击失败的事件数量', key: 'events' }, [ + h(AnimatedNumber, { tag: 'b', value: stats.triage.attackTotal, key: 'value' }), + h('span', { className: cx('outcome-label', triageActive && 'is-processing'), key: 'label' }, [ + h('i', { key: 'ai' }, triageActive ? 'AI研判' : 'AI判定'), + h('em', { key: 'text' }, triageActive ? '处理中' : '安全事件'), + ]), + ]), + h('div', { title: 'AI 研判结论为良性的事件数量', key: 'benign' }, [ + h(AnimatedNumber, { tag: 'b', value: stats.triage.benign, key: 'value' }), + h('span', { className: 'outcome-label', key: 'label' }, [ + h('i', { key: 'ai' }, 'AI判定'), + h('em', { key: 'text' }, '非安全事件'), + ]), + ]), + h('div', { title: 'AI 判定需要进入人工复核的事件数量', key: 'manual' }, [ + h(AnimatedNumber, { tag: 'b', value: stats.closedLoop.pending, key: 'value' }), + h('span', { className: 'outcome-label', key: 'label' }, [ + h('i', { key: 'ai' }, 'AI判定'), + h('em', { key: 'text' }, '待人工复核'), + ]), + ]), ]), h('div', { className: 'severity-stack', key: 'severity' }, severities.map((item) => h('div', { className: cx(`severity-node severity-${item.tone}`, item.tone === activeSeverityTone && 'active-target', !activeSeverityTone && item.tone === recentSeverityTone && 'recent-target'), @@ -3746,6 +3757,28 @@ const CSS = ` .outcome-stack > div { display: flex; flex-direction: column; min-width: 90px; } .outcome-stack b { color: #eef1f0; font-size: 23px; line-height: 1; } .outcome-stack span { margin-top: 6px; color: #78837f; font-size: 12px; } +.outcome-stack span.outcome-label { + display: flex; + align-items: center; + gap: 6px; + white-space: nowrap; +} +.outcome-label i { + padding: 1px 4px; + border: 1px solid rgba(43,231,255,.3); + border-radius: 3px; + color: #62cbe8; + background: rgba(43,231,255,.07); + font-size: 9px; + font-style: normal; + line-height: 14px; +} +.outcome-label em { color: #86a7b8; font-size: 12px; font-style: normal; } +.outcome-label.is-processing i { + border-color: rgba(155,140,255,.38); + color: #aa9cff; + background: rgba(155,140,255,.1); +} .outcome-stack .primary b { color: #4bdbae; } .severity-stack { position: absolute; diff --git a/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/src/severityValues.ts b/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/src/severityValues.ts new file mode 100644 index 000000000..d2fcd6f2a --- /dev/null +++ b/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/src/severityValues.ts @@ -0,0 +1,27 @@ +type SeverityItem = { + label?: string; + value?: number; +}; + +export function severityKey(value: unknown) { + const key = String(value || '').trim().toLowerCase(); + if (key === 'critical' || key === 'severe' || key === '4') return 'critical'; + if (key === 'high' || key === '3') return 'high'; + if (key === 'medium' || key === '2') return 'medium'; + if (key === 'low' || key === '1') return 'low'; + return ''; +} + +export function severityRows(stats: { severityLevels?: SeverityItem[] }) { + const counts = { critical: 0, high: 0, medium: 0, low: 0 }; + for (const item of stats.severityLevels || []) { + const key = severityKey(item.label); + if (key) counts[key] += Number(item.value || 0); + } + return [ + { key: 'critical', label: '严重', value: counts.critical, tone: 'critical' }, + { key: 'high', label: '高危', value: counts.high, tone: 'high' }, + { key: 'medium', label: '中危', value: counts.medium, tone: 'medium' }, + { key: 'low', label: '低危', value: counts.low, tone: 'low' }, + ]; +} diff --git a/.flocks/flockshub/plugins/webuis/soc_ui/soc_overview/api/handlers.py b/.flocks/flockshub/plugins/webuis/soc_ui/soc_overview/api/handlers.py index f9c70edcc..6cc44b672 100644 --- a/.flocks/flockshub/plugins/webuis/soc_ui/soc_overview/api/handlers.py +++ b/.flocks/flockshub/plugins/webuis/soc_ui/soc_overview/api/handlers.py @@ -271,8 +271,8 @@ def _simulate_triage_from_denoise(paths): if verdict in {"attack_success", "attack", "attack_failed"}: file_attack += 1 source_counter[_norm(obj.get("_source_type") or obj.get("source_type") or obj.get("device_type"))] += 1 - threat_counter[_norm(obj.get("_threat_type") or obj.get("threat_name") or obj.get("threat_type"))] += 1 - risk_counter[_norm(obj.get("threat_level") or obj.get("threat_severity") or obj.get("risk_level"))] += 1 + threat_counter[_norm(obj.get("threat_name"))] += 1 + risk_counter[_norm(obj.get("risk_level"))] += 1 _update_profile_counters(obj, profile_counters) event_start, event_end = _merge_record_time(event_start, event_end, obj) series_total.append(file_total) @@ -627,7 +627,7 @@ def _read_denoise(paths, workflow_call_count: int = 0): if obj.get("is_duplicate") is True: file_duplicates += 1 source_counter[_norm(obj.get("_source_type") or obj.get("source_type") or obj.get("device_type"))] += 1 - threat_counter[_norm(obj.get("_threat_type") or obj.get("threat_name") or obj.get("threat_type"))] += 1 + threat_counter[_norm(obj.get("threat_name"))] += 1 _update_profile_counters(obj, profile_counters) event_start, event_end = _merge_record_time(event_start, event_end, obj) total_raw += file_raw @@ -712,8 +712,8 @@ def _read_triage(paths): source = _norm(obj.get("_source_type") or obj.get("source_type") or obj.get("device_type")) source_counter[source] += 1 - threat_counter[_norm(obj.get("_threat_type") or obj.get("threat_name") or obj.get("threat_type"))] += 1 - risk_counter[_norm(obj.get("risk_level") or obj.get("threat_level") or obj.get("threat_severity"))] += 1 + threat_counter[_norm(obj.get("threat_name"))] += 1 + risk_counter[_norm(obj.get("risk_level"))] += 1 _update_profile_counters(obj, profile_counters) event_start, event_end = _merge_record_time(event_start, event_end, obj) triage_source = _norm(obj.get("triage_source")) @@ -800,7 +800,7 @@ def _update_profile_counters(obj, counters): counters["directionCounter"][_norm(obj.get("direction") or obj.get("traffic_direction"))] += 1 counters["resultCounter"][_norm(obj.get("threat_result") or obj.get("attack_verdict"))] += 1 counters["protocolCounter"][_norm(obj.get("net_type") or obj.get("net_app_proto") or obj.get("protocol"))] += 1 - counters["severityCounter"][_norm(obj.get("threat_severity") or obj.get("threat_level") or obj.get("risk_level"))] += 1 + counters["severityCounter"][_norm(obj.get("threat_severity"))] += 1 counters["responseCounter"][_norm(obj.get("rsp_status_code") or obj.get("status_code"))] += 1 port_value = obj.get("dport") or obj.get("dst_port") or obj.get("destination_port") @@ -892,7 +892,7 @@ def _build_field_stats(paths): status_code = _field_text(obj.get("rsp_status_code")) direction = _norm(obj.get("direction") or "unknown") protocol = _norm(obj.get("net_type") or obj.get("net_app_proto") or "unknown") - threat_type = _norm(obj.get("threat_type") or obj.get("_threat_type") or obj.get("threat_name")) + threat_type = _norm(obj.get("_threat_type") or obj.get("threat_type")) threat_result = _norm(obj.get("threat_result") or "unknown") threat_phase = _norm(obj.get("threat_phase") or "unknown") diff --git a/.flocks/flockshub/plugins/webuis/soc_ui/soc_overview/src/index.tsx b/.flocks/flockshub/plugins/webuis/soc_ui/soc_overview/src/index.tsx index dafc71b62..cbf4d3600 100644 --- a/.flocks/flockshub/plugins/webuis/soc_ui/soc_overview/src/index.tsx +++ b/.flocks/flockshub/plugins/webuis/soc_ui/soc_overview/src/index.tsx @@ -106,24 +106,19 @@ const EMPTY_STATS: Required = { fieldStats: EMPTY_FIELD_STATS, }; -const LABELS: Record = { +const PHASE_LABELS: Record = { exploit: '漏洞利用', recon: '侦察探测', post_exploit: '后渗透', control: '控制通信', - tunneling: '隧道通信', - file: '文件风险', - c2: '控制通信', - trojan: '木马', - ransom: '勒索', - shell: '命令执行', - botnet: '僵尸网络', - success: '攻击成功', - failed: '攻击失败', - unknown: '待确认', + unknown: '未知阶段', +}; + +const DIRECTION_LABELS: Record = { in: '入站', out: '出站', lateral: '横向', + unknown: '未知方向', }; const TIME_RANGE_OPTIONS: ChoiceOption[] = [ @@ -271,8 +266,8 @@ function refreshLabel(value: RefreshKey) { return REFRESH_OPTIONS.find((option) => option.value === value)?.label || '关闭'; } -function labelOf(item: CounterItem) { - return LABELS[item.key || item.label] || LABELS[item.label] || item.label; +function labelOf(item: CounterItem, labels?: Record) { + return labels?.[item.key || item.label] || labels?.[item.label] || item.label; } function truncate(value: string, max = 34) { @@ -415,8 +410,8 @@ export default function SocOverviewPage() {
- - + +
@@ -563,13 +558,13 @@ function Panel({ title, children }: { title: string; children: React.ReactNode } ); } -function RankList({ rows, mono = false }: { rows: CounterItem[]; mono?: boolean }) { +function RankList({ rows, mono = false, labels }: { rows: CounterItem[]; mono?: boolean; labels?: Record }) { return (
{rows.length ? rows.map((item, index) => (
{String(index + 1).padStart(2, '0')} - {truncate(labelOf(item), mono ? 44 : 30)} + {truncate(labelOf(item, labels), mono ? 44 : 30)} {formatNumber(item.value)}
)) :
暂无数据
} @@ -582,7 +577,7 @@ function TileList({ rows }: { rows: CounterItem[] }) {
{rows.slice(0, 9).map((item) => (
- {labelOf(item)} + {item.label} {formatNumber(item.value)}
))} @@ -591,13 +586,13 @@ function TileList({ rows }: { rows: CounterItem[] }) { ); } -function ProgressList({ rows }: { rows: CounterItem[] }) { +function ProgressList({ rows, labels }: { rows: CounterItem[]; labels?: Record }) { const total = Math.max(1, rows.reduce((value, item) => value + item.value, 0)); return (
{rows.map((item) => (
- {labelOf(item)} + {labelOf(item, labels)} {formatNumber(item.value)}
@@ -607,10 +602,10 @@ function ProgressList({ rows }: { rows: CounterItem[] }) { ); } -function SplitRanks({ leftTitle, leftRows, rightTitle, rightRows, mono = false }: { leftTitle: string; leftRows: CounterItem[]; rightTitle: string; rightRows: CounterItem[]; mono?: boolean }) { +function SplitRanks({ leftTitle, leftRows, leftLabels, rightTitle, rightRows, mono = false }: { leftTitle: string; leftRows: CounterItem[]; leftLabels?: Record; rightTitle: string; rightRows: CounterItem[]; mono?: boolean }) { return (
-

{leftTitle}

+

{leftTitle}

{rightTitle}

); diff --git a/.flocks/flockshub/plugins/webuis/soc_ui/workspace.json b/.flocks/flockshub/plugins/webuis/soc_ui/workspace.json index 4988a545e..b17ffc05a 100644 --- a/.flocks/flockshub/plugins/webuis/soc_ui/workspace.json +++ b/.flocks/flockshub/plugins/webuis/soc_ui/workspace.json @@ -1,5 +1,6 @@ { "id": "soc_ui", + "version": "1.1.4", "title": "SOC 工作区", "titleEn": "SOC Workspace", "icon": "ShieldCheck", diff --git a/.flocks/plugins/skills/sangfor-edr-use/SKILL.md b/.flocks/plugins/skills/sangfor-edr-use/SKILL.md index 479c8a18a..616a5b9ec 100644 --- a/.flocks/plugins/skills/sangfor-edr-use/SKILL.md +++ b/.flocks/plugins/skills/sangfor-edr-use/SKILL.md @@ -1,37 +1,81 @@ --- name: sangfor-edr-use -description: 用于处理深信服 EDR(终端检测与响应)相关任务,通过 Flocks browser/CDP 完成登录态复用、终端状态查询、概况统计、失陷设备排查和设备运行状态查看。只要用户提到深信服 EDR、EDR 或 sangfor EDR,必须先加载本 skill;本 skill 是 EDR 任务的唯一入口,未阅读前不要直接使用 browser-use。 +description: 深信服 EDR 登录态管理与首页仪表盘 API 采集。用户提到深信服 EDR、EDR 或 sangfor EDR 时必须先加载本 skill。 --- # 深信服 EDR Use -本 skill 只负责任务入口和登录分流;CDP 操作、浏览器启动、验证码、selector、tab/iframe 处理和页面数据提取见 [references/cdp-workflow.md](references/cdp-workflow.md)。 +## 核心功能 -## 登录分流 +实现边界:`sangfor_edr_http_login.py` 负责 HTTP 登录、认证探测、 +`auth-state.json` Cookie 及 Secret Manager token bundle; +`sangfor_edr_dashboard_api.py` 负责仪表盘 API 请求,只读取 HTTP 登录模块 +验证过的同一套 Cookie/token,不从其他状态源拼接凭据。 -需要打开 EDR 页面、抓取页面请求或调用 EDR API 采集工具时,按以下顺序执行: +- 管理同一次登录产生的 Cookie 与 `login_token`。 +- 默认使用 HTTP 登录,开始前必须向用户索取并保存 EDR 地址、用户名和密码。 +- 仅当用户明确选择“打开页面后手动登录”时,才可不索取账密并直接使用 browser/CDP。 +- HTTP 登录连续 3 次失败后,按“browser/CDP 自动化登录(仍需账密)→保留页面供用户手动登录(不需账密)”顺序降级。 +- 每次登录或数据采集前探测现有认证,认证有效则跳过登录,失效则重新登录并更新存储。 +- 通过 API 采集首页终端概况、受影响终端、漏洞、勒索防护、实时病毒、Top 5 终端和设备资源使用率。 -1. 调用 `sangfor_edr_auth(action=status_auth_state)`,确认 `validation.valid`、`can_auto_refresh` 和 `has_saved_token`。不要先向用户索要账密。 -2. `validation.valid=true`: - - 仅浏览器/CDP任务:直接继续; - - 需要 API token:`has_saved_token=true` 才继续,否则调用 `refresh_auth_state` 补齐 token。 -3. `validation.valid=false` 且 `can_auto_refresh=true`:调用 `ensure_auth_state` 自动登录。 -4. 没有可用 state 或账密时:调用 `ensure_auth_state` 打开登录页;返回 `manual_login_required` 后,让用户在该工具打开的同一个 EDR tab 中完成登录、MFA 或 UKey。 -5. 用户完成手动登录后,调用 `complete_manual_login`;只有返回 `manual_login_captured_auth_state` 且 `token_saved=true`,才认为浏览器登录和 API 登录准备均完成。 -6. 返回 `browser_daemon_not_ready` 或 `auth_state_load_failed_browser_daemon_not_ready`:依次执行 `flocks browser --setup`、`flocks browser --doctor`,再重试原 action。 +## 输入与输出 -认证工具会在登录提交前监听 `launch_login.php` 的 fetch/XHR 响应,将 `data.token` 保存到 Secret Manager;token 不得回显、写入日志或写入 `auth-state.json`。后续 API 工具应从 Secret Manager 读取对应 token,并同时使用同一 EDR 的 cookies。 +### 认证工具 -`bu.port` 是 Flocks browser daemon 的 IPC 端口文件,不是 Chrome 的 remote-debugging 端口;禁止手工创建或修改它。 +调用 `sangfor_edr_auth`: -## 任务边界 +- `status_auth_state`:返回认证文件、凭据和认证探测状态,不返回敏感值。 +- `ensure_auth_state`、`refresh_auth_state`、`http_login`:执行“探测后按需 HTTP 登录”;HTTP 连续 3 次失败时按降级策略进入 browser/CDP,再失败则转为手动登录。 +- `browser_login`:用户明确选择自动化登录时,执行“探测后按需 browser/CDP 登录”;该路径需要账密,除非用户明确要求打开页面后自行手动登录。 +- `validate_auth_state`:仅验证浏览器 state。 +- `complete_manual_login`:保存用户在已打开浏览器中完成的登录。 -- 当前 `sangfor_edr_v1_0_0` 的页面业务仍通过浏览器/CDP完成;需要 API 采集时,必须先完成 token readiness 检查。 -- 首页仪表盘可用于设备 CPU/内存/硬盘、终端概况和失陷统计。 -- 查询失陷终端清单时,必须进入“威胁资产分析”并选择“已失陷终端”,不能使用默认“全部”列表。 -- 设备状态抓取脚本位于 `references/fetch_edr_system_state.py`;执行前必须完成上述登录分流。 +成功输出包含 `status`、`valid`、`login_skipped` 和非敏感探测结果;失败输出包含稳定的 `reason` 与恢复建议。 + +### 仪表盘工具 + +调用 `sangfor_edr_dashboard`,可输入: + +- `sections`:可选采集项;省略时采集全部仪表盘数据。 +- `days`:统计时间范围,允许 1–90 天。 +- `base_url`、`auth_state_path`:可选运行时覆盖。 + +输出包含 `data`、`errors`、`sections`、`days` 和本次认证是复用还是重登;不得输出 Cookie、密码或 `login_token`。 + +## 关键配置 + +- `base_url`:从用户提供的 EDR 地址提取 scheme、host 和 port;不得使用固定示例地址。 +- `username`、`password`:从设备配置或 Secret Manager 读取。 +- `auth_state_path`:Cookie state 文件位置。 +- `auto_ocr_code`、`max_captcha_retry`:HTTP/browser 登录验证码配置。 +- selector 和页面路径配置仅用于 browser/CDP 登录。 + +## EDR 交互协议 + +1. 从 Secret Manager 与 `auth-state.json` 加载成套认证,并校验 base URL 和 Cookie 指纹。 +2. 使用 Cookie 与 `login_token` 调用威胁终端概览接口 `get_agent_overview`: + - HTTP 200、无登录页重定向、响应成功且包含终端概览数据:认证有效,跳过登录。 + - Cookie 缺失或不匹配、401/403、重定向、非 200、响应无终端概览数据:认证失效。 +3. 默认重新登录流程:访问登录页,获取 RSA 公钥和验证码,提交 `dlogin`,调用 `launch_login.php`,再 GET `/ui`;HTTP 登录连续 3 次失败后进入 browser/CDP 自动化登录,自动化登录仍需账密,自动化登录失败后保留页面供用户手动登录。 +4. 登录成功后将 Cookie 写入 `auth-state.json`,将 `login_token` 写入 Secret Manager,并更新配对指纹。 +5. 仪表盘 API 只能使用通过上述探测的同一套 Cookie/token;禁止从不同 state 或 Secret 拼接。 + +## 错误处理 + +- 缺少地址或账密:返回缺失字段,向用户索取后保存到配置或 Secret Manager。 +- 验证码或 HTTP 登录失败:最多进行 3 次独立 HTTP 登录尝试;仍失败则切换 browser/CDP 自动化登录,切换时仍需账密。 +- browser/CDP 自动化登录失败或用户明确选择打开页面后手动登录:保留浏览器供用户完成登录,不再索取账密,再调用 `complete_manual_login`。 +- 认证探测失败:禁止继续业务 API;先执行 HTTP 重登并再次探测,连续 3 次 HTTP 仍失败则按 browser/CDP 自动化登录→手动登录降级。 +- 仪表盘部分接口失败:保留成功数据,在 `errors` 中按采集项返回失败原因。 +- Cookie、密码和 `login_token` 不得回显、记录日志或混入业务输出。 ## 执行约束 -- 运行 skill 内 Python 脚本时必须使用 Flocks 虚拟环境;不要使用系统 Python。 -- 需要具体 CDP 命令、平台启动方式、验证码/selector 配置、tab/iframe 处理或页面关键词时,读取 `references/cdp-workflow.md`,不要在本文件重复展开。 +- 运行本 Skill 提供的 Python 脚本时,必须使用 Flocks 虚拟环境;禁止使用系统 Python。 +- 不得假设 Flocks 项目、插件或虚拟环境的绝对路径;代码必须通过当前运行时加载的模块、`Path.home()`、`~/.flocks` 或显式配置/环境变量解析路径。 +- 需要具体 CDP 命令、浏览器启动方式、验证码识别、selector、tab/iframe 处理或页面关键词时,必须先阅读 [references/cdp-workflow.md](references/cdp-workflow.md),不要在本文件重复展开。 +- `bu.port` 是 Flocks browser daemon 的 IPC 端口文件,不是 Chrome remote-debugging 端口;禁止手工创建或修改。 +- 默认认证和仪表盘采集开始时不得启动 browser daemon;只有用户明确选择 `browser_login`、执行 `validate_auth_state`/`complete_manual_login`,或 HTTP 登录连续 3 次失败进入降级流程时,才允许使用 browser/CDP。 +- HTTP 登录连续 3 次失败后必须按 browser/CDP 自动化登录→手动登录顺序降级;自动化登录阶段仍需账密,手动登录阶段不得要求账密。 +- 任何 API 采集前必须完成认证探测;认证探测失败时不得继续调用业务接口。 diff --git a/.flocks/plugins/tools/device/sangfor_edr_webcli/_provider.yaml b/.flocks/plugins/tools/device/sangfor_edr_webcli/_provider.yaml index 5a065f634..f750abcfc 100644 --- a/.flocks/plugins/tools/device/sangfor_edr_webcli/_provider.yaml +++ b/.flocks/plugins/tools/device/sangfor_edr_webcli/_provider.yaml @@ -4,13 +4,11 @@ service_id: sangfor_edr version: "1.0.0" integration_type: device description: > - Sangfor EDR WebCLI-backed integration. It can reuse an existing browser - storage state, or refresh auth-state by driving the real EDR login page - through browser daemon / CDP when credentials are configured. + Sangfor EDR integration with default HTTP login, explicitly selected + browser/CDP login, cookie/token validation, and dashboard API collection. description_cn: > - Sangfor EDR WebCLI device integration. It reuses full browser auth-state - from manual login, or refreshes full browser auth-state through CDP-assisted - browser login when credentials are configured. + 深信服 EDR 集成。默认通过 HTTP 登录,仅用户明确选择时使用 browser/CDP; + 每次操作验证成套 Cookie/token,并通过 API 采集首页仪表盘。 credential_fields: - key: base_url label: Base URL @@ -33,7 +31,7 @@ credential_fields: secret_id: sangfor_edr_username input_type: text required: false - description: Optional. Stored for CDP-assisted browser login when auth-state is missing or expired. + description: Optional. Stored for default HTTP login and explicitly selected browser/CDP login. - key: password label: Password storage: secret @@ -41,7 +39,7 @@ credential_fields: secret_id: sangfor_edr_password input_type: password required: false - description: Optional. Stored as a secret and used only to refresh browser auth-state. + description: Optional. Stored as a secret and used only to refresh EDR authentication. - key: auto_ocr_code label: Auto OCR Captcha storage: config @@ -114,10 +112,10 @@ defaults: category: custom verify_ssl: false notes: | - The existing manual-login flow saves a full browser storageState containing - cookies and localStorage. When credentials are available, the automatic flow - also drives the real EDR login page through browser daemon / CDP and then - saves the same full browser storageState format. + All login methods save cookies to auth-state.json and login_token to Secret + Manager. Pairing metadata prevents dashboard APIs from mixing credentials + produced by different logins. Device URLs are normalized to scheme, host, + and port before use. If captcha OCR, MFA, page selector matching, or login success detection fails, keep the browser/manual login recovery path and save the resulting full diff --git a/.flocks/plugins/tools/device/sangfor_edr_webcli/sangfor_edr.handler.py b/.flocks/plugins/tools/device/sangfor_edr_webcli/sangfor_edr.handler.py index 2680c140e..7abff1e6e 100644 --- a/.flocks/plugins/tools/device/sangfor_edr_webcli/sangfor_edr.handler.py +++ b/.flocks/plugins/tools/device/sangfor_edr_webcli/sangfor_edr.handler.py @@ -1,38 +1,39 @@ -"""Sangfor EDR browser-state authentication helper. - -EDR has no stable Open API in this integration. This handler follows the same -browser workflow used by TDP / OneSEC / SkyEye / Qingteng skills: - -1. try to load the saved browser auth-state; -2. if it is still valid, reuse it; -3. if it is missing or expired, open the real EDR login page through CDP; -4. fill username, password and captcha in the browser page; -5. after login succeeds, save the full browser auth-state again. - -The handler only replaces the "wait for the user to log in manually" step with -CDP-assisted form login. It does not implement EDR business APIs. -""" +"""Sangfor EDR authentication and dashboard API helpers.""" from __future__ import annotations import base64 +import hashlib import json import os +import secrets +import sys import time +from datetime import datetime, timedelta from pathlib import Path from typing import Any, Optional from urllib.parse import urljoin, urlparse +import requests +import urllib3 + from flocks.browser import helpers from flocks.browser.admin import ensure_daemon from flocks.config.config_writer import ConfigWriter from flocks.tool.registry import ToolContext, ToolResult +_PLUGIN_DIR = str(Path(__file__).resolve().parent) +if _PLUGIN_DIR not in sys.path: + sys.path.insert(0, _PLUGIN_DIR) +import sangfor_edr_dashboard_api as _dashboard_api_module # noqa: E402 +import sangfor_edr_http_login as _http_login_module # noqa: E402 + SERVICE_ID = "sangfor_edr_v1_0_0" LEGACY_SERVICE_ID = "sangfor_edr" USERNAME_SECRET_ID = "sangfor_edr_username" PASSWORD_SECRET_ID = "sangfor_edr_password" TOKEN_SECRET_ID = "sangfor_edr_token" +TOKEN_BUNDLE_SECRET_ID = "sangfor_edr_token_bundle" DEFAULT_AUTH_STATE_PATH = "~/.flocks/browser/sangfor-edr/auth-state.json" DEFAULT_LOGIN_PATH = "/ui/login.php" DEFAULT_INDEX_PATH = "/ui/#/index" @@ -241,13 +242,26 @@ def _saved_auto_login_status(params: dict[str, Any]) -> dict[str, Any]: has_base_url = bool(str(base_url or "").strip()) has_username = bool(str(username or "").strip()) has_password = bool(str(password or "").strip()) + has_saved_token = bool(secrets.get(TOKEN_SECRET_ID)) + pair_verified = False + auth_probe: dict[str, Any] | None = None + if has_base_url and auth_state_path.exists(): + try: + cfg = _resolve_runtime_config({**params, "persist_credentials": False}) + auth_probe = _probe_auth_pair(cfg) + pair_verified = bool(auth_probe.get("valid")) + except Exception as exc: + auth_probe = {"valid": False, "reason": "auth_probe_failed", "error": str(exc)} + pair_verified = False return { "auth_state_path": str(auth_state_path), "auth_state_exists": auth_state_path.exists(), "has_base_url": has_base_url, "has_saved_username": has_username, "has_saved_password": has_password, - "has_saved_token": bool(secrets.get(TOKEN_SECRET_ID)), + "has_saved_token": has_saved_token, + "has_verified_auth_pair": pair_verified, + "auth_probe": auth_probe, "can_auto_refresh": has_base_url and has_username and has_password, } @@ -313,6 +327,196 @@ def _now_ms() -> str: return str(int(time.time() * 1000)) +def _query_id() -> str: + return f"Query_{_now_ms()}" + + +def _safe_error(exc: Exception, *sensitive_values: str) -> str: + message = str(exc) + for value in sensitive_values: + if value: + message = message.replace(value, "") + return message + + +def _http_headers(cfg: RuntimeConfig, *, image: bool = False) -> dict[str, str]: + headers = { + "Pragma": "no-cache", + "Cache-Control": "no-cache", + "User-Agent": ( + "Mozilla/5.0 (Windows NT 10.0; Win64; x64) " + "AppleWebKit/537.36 (KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36" + ), + "Accept-Language": "zh-CN,zh;q=0.9", + "Referer": _login_url(cfg) if image else _url(cfg, "/ui/"), + } + if image: + headers["Accept"] = "image/avif,image/webp,image/apng,image/svg+xml,image/*,*/*;q=0.8" + else: + headers.update( + { + "Accept": "application/json, text/plain, */*", + "Content-Type": "application/json", + "Origin": cfg.base_url, + } + ) + return headers + + +def _read_auth_state(path: Path) -> dict[str, Any]: + state = json.loads(path.read_text(encoding="utf-8")) + if not isinstance(state, dict) or not isinstance(state.get("cookies"), list): + raise ValueError("EDR auth-state does not contain a cookie list.") + return state + + +def _cookie_fingerprint(cookies: list[dict[str, Any]], base_url: str) -> str: + normalized = sorted( + ( + str(cookie.get("name") or ""), + str(cookie.get("value") or ""), + str(cookie.get("domain") or ""), + str(cookie.get("path") or "/"), + ) + for cookie in cookies + if isinstance(cookie, dict) and cookie.get("name") + ) + payload = json.dumps( + {"base_url": _normalise_base_url(base_url), "cookies": normalized}, + ensure_ascii=False, + separators=(",", ":"), + ) + return hashlib.sha256(payload.encode("utf-8")).hexdigest() + + +def _save_auth_pair(cfg: RuntimeConfig, state: dict[str, Any], token: str) -> dict[str, Any]: + cookies = state.get("cookies") + if not isinstance(cookies, list) or not cookies or not token: + raise ValueError("EDR login did not return a complete cookie/token pair.") + cfg.auth_state_path.parent.mkdir(parents=True, exist_ok=True) + temp_path = cfg.auth_state_path.with_name(f"{cfg.auth_state_path.name}.tmp") + temp_path.write_text(json.dumps(state, ensure_ascii=False, indent=2), encoding="utf-8") + temp_path.replace(cfg.auth_state_path) + fingerprint = _cookie_fingerprint(cookies, cfg.base_url) + secret_manager = _get_secret_manager() + secret_manager.set(TOKEN_SECRET_ID, token) + secret_manager.set( + TOKEN_BUNDLE_SECRET_ID, + json.dumps( + { + "token": token, + "base_url": cfg.base_url, + "cookie_fingerprint": fingerprint, + }, + ensure_ascii=False, + separators=(",", ":"), + ), + ) + return { + "auth_state_path": str(cfg.auth_state_path), + "cookie_count": len(cookies), + "pair_verified": True, + } + + +def _load_verified_auth_pair(cfg: RuntimeConfig) -> tuple[dict[str, Any], str]: + state = _read_auth_state(cfg.auth_state_path) + raw_bundle = _get_secret_manager().get(TOKEN_BUNDLE_SECRET_ID) + if not raw_bundle: + raise RuntimeError("EDR cookie/token pair metadata is missing; refresh authentication.") + try: + bundle = json.loads(raw_bundle) + except (TypeError, json.JSONDecodeError) as exc: + raise RuntimeError("EDR cookie/token pair metadata is invalid; refresh authentication.") from exc + token = str(bundle.get("token") or "") + base_url = str(bundle.get("base_url") or "") + fingerprint = str(bundle.get("cookie_fingerprint") or "") + actual = _cookie_fingerprint(state["cookies"], cfg.base_url) + if ( + not token + or _normalise_base_url(base_url) != cfg.base_url + or not secrets.compare_digest(fingerprint, actual) + ): + raise RuntimeError("EDR cookies and login_token are not from the same login; refresh authentication.") + return state, token + + +def _response_contains_agent_overview(payload: Any) -> bool: + """Return whether a successful threat-terminal response has data content.""" + if not isinstance(payload, dict) or not payload.get("success"): + return False + return any( + key in payload and payload.get(key) is not None + for key in ("data", "result", "agent_overview", "agent_total", "total", "count") + ) + + +def _unix_date_range(days: int = 7) -> dict[str, int]: + end = datetime.now().replace(hour=23, minute=59, second=59, microsecond=0) + start = (end - timedelta(days=max(1, days) - 1)).replace(hour=0, minute=0, second=0) + return {"start": int(start.timestamp()), "end": int(end.timestamp())} + + +def _probe_auth_pair(cfg: RuntimeConfig) -> dict[str, Any]: + try: + state, token = _load_verified_auth_pair(cfg) + except Exception as exc: + return {"valid": False, "reason": "auth_pair_missing_or_mismatched", "error": str(exc)} + + session = _dashboard_session(cfg, state) + try: + response = session.post( + _url(cfg, f"/launch.php?s={token}&opr=get_agent_overview"), + headers=_http_headers(cfg), + json={ + "app_args": {"name": "app.web.event_center.head", "options": {}}, + "auto": 1, + "opr": "get_agent_overview", + "date_range": _unix_date_range(), + "query_id": _query_id(), + }, + timeout=cfg.timeout, + allow_redirects=False, + ) + except Exception as exc: + return {"valid": False, "reason": "auth_probe_request_failed", "error": _safe_error(exc, token)} + + location = str(response.headers.get("Location") or "") + if response.status_code in {301, 302, 303, 307, 308}: + return { + "valid": False, + "reason": "auth_probe_redirected", + "http_status": response.status_code, + "location": location, + } + if response.status_code in {401, 403}: + return { + "valid": False, + "reason": "auth_probe_unauthorized", + "http_status": response.status_code, + } + if response.status_code != 200: + return { + "valid": False, + "reason": "auth_probe_http_error", + "http_status": response.status_code, + } + try: + payload = response.json() + except Exception as exc: + return {"valid": False, "reason": "auth_probe_invalid_json", "error": str(exc)} + if not isinstance(payload, dict) or not payload.get("success"): + return {"valid": False, "reason": "auth_probe_rejected"} + if not _response_contains_agent_overview(payload): + return {"valid": False, "reason": "auth_probe_expected_agent_data_missing"} + return { + "valid": True, + "reason": "auth_probe_succeeded", + "http_status": 200, + "agent_overview_verified": True, + } + + def _login_url(cfg: RuntimeConfig) -> str: return _url(cfg, cfg.login_path) @@ -415,18 +619,34 @@ def _install_login_token_capture() -> bool: return False -def _save_captured_login_token(timeout: float = 2.0) -> bool: +def _captured_login_token(timeout: float = 2.0) -> str: deadline = time.time() + timeout while time.time() < deadline: try: token = str(helpers.js("sessionStorage.getItem('__flocks_sangfor_edr_token') || ''") or "").strip() if token: - _get_secret_manager().set(TOKEN_SECRET_ID, token) - return True + return token except Exception: - return False + return "" time.sleep(0.1) - return False + return "" + + +def _save_captured_login_token(timeout: float = 2.0) -> bool: + token = _captured_login_token(timeout) + if not token: + return False + _get_secret_manager().set(TOKEN_SECRET_ID, token) + return True + + +def _save_browser_auth_pair(cfg: RuntimeConfig) -> dict[str, Any]: + token = _captured_login_token() + if not token: + raise RuntimeError("EDR login succeeded but login_token was not captured.") + saved = _save_auth_state(cfg) + pair = _save_auth_pair(cfg, _read_auth_state(cfg.auth_state_path), token) + return {"state": saved, "pair": pair} def _page_text() -> str: @@ -929,8 +1149,7 @@ def _refresh_auth_state_with_cdp_login(cfg: RuntimeConfig, captcha_code: str = " filled = _set_login_form_values(cfg, code) if _wait_for_login_success(cfg): - token_saved = _save_captured_login_token() - saved = _save_auth_state(cfg) + saved = _save_browser_auth_pair(cfg) return { "success": True, "status": "browser_cdp_login_refreshed_auth_state", @@ -939,7 +1158,7 @@ def _refresh_auth_state_with_cdp_login(cfg: RuntimeConfig, captcha_code: str = " "form": form_state, "filled": {key: bool(value) for key, value in filled.items()}, "saved": saved, - "token_saved": token_saved, + "token_saved": True, } last_error = "login_success_check_timeout" except Exception as exc: @@ -967,14 +1186,27 @@ def _complete_manual_login(cfg: RuntimeConfig) -> dict[str, Any]: _open_page(_index_url(cfg)) if not _is_logged_in(cfg): return _manual_login_result(cfg, "manual_login_not_completed") - token_saved = _save_captured_login_token() + saved = _save_browser_auth_pair(cfg) + probe = _probe_auth_pair(cfg) + if not probe.get("valid"): + return { + "success": False, + "valid": False, + "status": "manual_login_capture_probe_failed", + "reason": str(probe.get("reason") or "auth_probe_failed"), + "error": "browser login state was saved but HTTP authentication probe failed", + "saved": saved, + "token_saved": True, + "probe": probe, + } return { "success": True, "valid": True, "status": "manual_login_captured_auth_state", "auth_state_path": str(cfg.auth_state_path), - "saved": _save_auth_state(cfg), - "token_saved": token_saved, + "saved": saved, + "token_saved": True, + "probe": probe, } except Exception as exc: return { @@ -988,6 +1220,99 @@ def _complete_manual_login(cfg: RuntimeConfig) -> dict[str, Any]: } +def _dashboard_session(cfg: RuntimeConfig, state: dict[str, Any]) -> requests.Session: + session = requests.Session() + session.verify = False + host = urlparse(cfg.base_url).hostname or "" + for cookie in state["cookies"]: + if not isinstance(cookie, dict) or not cookie.get("name"): + continue + domain = str(cookie.get("domain") or host) + session.cookies.set( + str(cookie["name"]), + str(cookie.get("value") or ""), + domain=domain, + path=str(cookie.get("path") or "/"), + ) + urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning) + return session + + +def _dashboard_requests( + cfg: RuntimeConfig, + token: str, + *, + days: int, +) -> dict[str, tuple[str, dict[str, Any]]]: + end = datetime.now().replace(hour=23, minute=59, second=59, microsecond=0) + start = (end - timedelta(days=days - 1)).replace(hour=0, minute=0, second=0) + unix_range = _unix_date_range(days) + text_range = { + "start": start.strftime("%Y-%m-%d %H:%M:%S"), + "end": end.strftime("%Y-%m-%d %H:%M:%S"), + } + + def launch(opr: str, app_name: str, extra: Optional[dict[str, Any]] = None): + payload = { + "app_args": {"name": app_name, "options": {}}, + "auto": 1, + "opr": opr, + "query_id": _query_id(), + } + payload.update(extra or {}) + return f"/launch.php?s={token}&opr={opr}", payload + + return { + "agent_overview": launch( + "get_agent_overview", + "app.web.event_center.head", + {"date_range": unix_range}, + ), + "influenced_agent_overview": launch( + "get_influenced_agent_overview", + "app.web.event_center.head", + ), + "vulnerability_overview": ( + f"/api/edrgoweb/v1/vulner/list/homepageVulner?_method=post&s={token}&req_type=polling", + { + "recentCheckTime": unix_range, + "uuid": "sf-id-6", + "tid": "0", + "uid": cfg.username, + "token": token, + }, + ), + "ransomware_defense": launch( + "ransom_virus_defense_interface", + "app.web.event_center.ransom_virus_protection", + ), + "realtime_virus": launch( + "real_time_virus", + "app.web.event_center.real_time_events", + {"date": text_range, "day_sum": days}, + ), + "top_agents": launch( + "get_top5_agents", + "app.web.event_center.head", + {"day_sum": days, "date_range": unix_range}, + ), + "resource_usage": launch( + "get_realtime_resource_usage", + "app.web.event_center.head", + ), + } + + +def _collect_dashboard( + cfg: RuntimeConfig, + *, + sections: list[str], + days: int, +) -> dict[str, Any]: + """Compatibility entry point; dashboard API implementation is standalone.""" + return _dashboard_api_module.collect_dashboard(cfg, sections=sections, days=days) + + # ── Tool actions ───────────────────────────────────────────────────────────── def _auth_state_loaded_output(validation: dict[str, Any]) -> dict[str, Any]: @@ -1003,27 +1328,12 @@ async def handle(ctx: ToolContext) -> ToolResult: action = str(params.get("action") or "ensure_auth_state").strip() try: - if action == "status_auth_state": - status = _saved_auto_login_status(params) - validation: dict[str, Any] | None = None - if status.get("has_base_url"): - try: - validation = _validate_auth_state(_resolve_runtime_config({**params, "persist_credentials": False})) - except Exception as exc: - validation = { - "valid": False, - "reason": "auth_state_validate_failed", - "error": str(exc), - "auth_state_path": status["auth_state_path"], - } + if action in {"status_auth_state", "ensure_auth_state", "refresh_auth_state", "http_login"}: + result = _http_login_module.run_auth_action(params) return ToolResult( - success=True, - output={ - "success": True, - "status": "saved_auto_login_status", - **status, - "validation": validation, - }, + success=bool(result.get("success")), + output=result, + error=None if result.get("success") else str(result.get("error") or result.get("reason")), ) cfg = _resolve_runtime_config(params) @@ -1044,26 +1354,68 @@ async def handle(ctx: ToolContext) -> ToolResult: error=None if result.get("success") else result.get("reason"), ) + if action == "browser_login": + probe = _probe_auth_pair(cfg) + if probe.get("valid"): + result = { + "success": True, + "valid": True, + "status": "browser_login_skipped_valid_auth_pair", + "login_skipped": True, + "probe": probe, + } + else: + result = _refresh_auth_state_with_cdp_login( + cfg, + captcha_code=str(params.get("captcha_code") or ""), + ) + result["previous_probe"] = probe + if result.get("success"): + confirmation = _probe_auth_pair(cfg) + result["probe"] = confirmation + if not confirmation.get("valid"): + result.update( + { + "success": False, + "valid": False, + "status": "browser_login_probe_failed", + "reason": str(confirmation.get("reason") or "auth_probe_failed"), + } + ) + return ToolResult( + success=bool(result.get("success")), + output=result, + error=None if result.get("success") else str(result.get("error") or result.get("reason")), + ) + if action not in {"ensure_auth_state", "refresh_auth_state"}: return ToolResult( success=False, - error="Unsupported Sangfor EDR auth action. Use status_auth_state, ensure_auth_state, validate_auth_state, refresh_auth_state, or complete_manual_login.", + error=( + "Unsupported Sangfor EDR auth action. Use status_auth_state, " + "ensure_auth_state, validate_auth_state, refresh_auth_state, " + "http_login, browser_login, or complete_manual_login." + ), ) - force_refresh = action == "refresh_auth_state" or _coerce_bool(params.get("force_refresh"), default=False) - if not force_refresh: - validation = _validate_auth_state(cfg) - if validation.get("valid"): - return ToolResult(success=True, output=_auth_state_loaded_output(validation)) - - result = _refresh_auth_state_with_cdp_login( - cfg, - captcha_code=str(params.get("captcha_code") or ""), + result = _http_login_module.run_auth_action(params) + return ToolResult( + success=bool(result.get("success")), + output=result, + error=None if result.get("success") else str(result.get("error") or result.get("reason")), ) + except Exception as exc: + return ToolResult(success=False, error=str(exc)) + + +async def handle_dashboard(ctx: ToolContext) -> ToolResult: + params = dict(ctx.params) + try: + result = _dashboard_api_module.run_dashboard(params) return ToolResult( success=bool(result.get("success")), output=result, - error=None if result.get("success") else result.get("reason"), + error=None if result.get("success") else "dashboard_api_partial_failure", ) except Exception as exc: return ToolResult(success=False, error=str(exc)) diff --git a/.flocks/plugins/tools/device/sangfor_edr_webcli/sangfor_edr_auth.yaml b/.flocks/plugins/tools/device/sangfor_edr_webcli/sangfor_edr_auth.yaml index 1b3d060ba..3aed97339 100644 --- a/.flocks/plugins/tools/device/sangfor_edr_webcli/sangfor_edr_auth.yaml +++ b/.flocks/plugins/tools/device/sangfor_edr_webcli/sangfor_edr_auth.yaml @@ -1,15 +1,14 @@ name: sangfor_edr_auth description: > - Manage Sangfor EDR browser authentication state. Use status_auth_state or - validate_auth_state before opening pages; use ensure_auth_state to reuse a - valid saved auth-state or automatically refresh full browser state through - browser daemon / CDP when credentials are configured. Browser actions start - or recover the daemon automatically and return structured recovery guidance. + Manage a matched Sangfor EDR cookie/login_token pair. ensure_auth_state and + refresh_auth_state use HTTP login by default. browser_login is used only when + the user explicitly selects browser/CDP automation. The implementation is in + sangfor_edr_http_login.py; every action probes the saved authentication first + and skips login while it remains valid. description_cn: > - 管理深信服 EDR 浏览器登录态。先用 status_auth_state 或 validate_auth_state - 检查已保存 state;当 state 缺失或失效且已配置账密时,ensure_auth_state - 通过 browser daemon / CDP 驱动真实登录页自动登录并重新保存完整浏览器 state。 - 需要浏览器的 action 会自动启动或恢复 daemon,失败时返回结构化恢复指引。 + 管理深信服 EDR 成套的 Cookie 与 login_token。ensure_auth_state 和 + refresh_auth_state 默认使用 HTTP 登录;仅用户明确选择 browser_login 时 + 使用 browser/CDP 自动化。每次操作先探测现有认证,有效时跳过登录。 category: custom enabled: true requires_confirmation: false @@ -19,13 +18,14 @@ inputSchema: properties: action: type: string - enum: [status_auth_state, ensure_auth_state, validate_auth_state, refresh_auth_state, complete_manual_login] + enum: [status_auth_state, ensure_auth_state, validate_auth_state, refresh_auth_state, http_login, browser_login, complete_manual_login] default: ensure_auth_state description: > Authentication action. status_auth_state returns non-sensitive saved state/credential availability; ensure_auth_state reuses existing state when it appears valid and refreshes when needed; validate_auth_state - checks only; refresh_auth_state always performs CDP-assisted browser login; + checks only; refresh_auth_state uses HTTP login without browser fallback; + http_login and browser_login explicitly select one login method; complete_manual_login validates a user-completed login and saves its browser state. captcha_code: type: string @@ -39,6 +39,16 @@ inputSchema: password: type: string description: Optional password to save as a secret before running auth. + auto_ocr_code: + type: boolean + default: true + description: Whether to recognize the HTTP/browser captcha with OCR when captcha_code is omitted. + max_captcha_retry: + type: integer + minimum: 1 + maximum: 10 + default: 5 + description: Maximum number of HTTP/browser login attempts when captcha validation fails. auth_state_path: type: string description: Optional auth-state path. Defaults to ~/.flocks/browser/sangfor-edr/auth-state.json. @@ -61,10 +71,6 @@ inputSchema: type: string default: "#button,input[name='button'],.login-opr-btn,#login,#submit,.login-btn,.btn-login,button[type='submit'],input[type='submit']" description: Optional comma-separated selectors for the EDR login submit control. - force_refresh: - type: boolean - default: false - description: Force CDP-assisted browser login even if auth-state exists. required: [action] handler: type: script diff --git a/.flocks/plugins/tools/device/sangfor_edr_webcli/sangfor_edr_dashboard.yaml b/.flocks/plugins/tools/device/sangfor_edr_webcli/sangfor_edr_dashboard.yaml new file mode 100644 index 000000000..4a8c27ce3 --- /dev/null +++ b/.flocks/plugins/tools/device/sangfor_edr_webcli/sangfor_edr_dashboard.yaml @@ -0,0 +1,46 @@ +name: sangfor_edr_dashboard +description: > + Collect Sangfor EDR home dashboard data through authenticated HTTP APIs. + Probes the saved cookie/login_token pair before every collection, skips login + when valid, and performs HTTP re-login when missing or expired. API request + definitions are implemented in sangfor_edr_dashboard_api.py, which reads the + verified pair owned by sangfor_edr_http_login.py. +description_cn: > + 通过深信服 EDR HTTP API 采集首页仪表盘。每次采集前探测 Cookie/token, + 有效时复用,缺失或失效时通过 HTTP 重新登录并更新认证。 +category: custom +enabled: true +requires_confirmation: false +provider: sangfor_edr +inputSchema: + type: object + properties: + sections: + type: array + items: + type: string + enum: + - agent_overview + - influenced_agent_overview + - vulnerability_overview + - ransomware_defense + - realtime_virus + - top_agents + - resource_usage + description: Optional dashboard sections. Omit to collect every section. + days: + type: integer + minimum: 1 + maximum: 90 + default: 7 + description: Time range in days for dashboard APIs that accept a range. + base_url: + type: string + description: Optional EDR device URL; its scheme, host, and port become the base URL. + auth_state_path: + type: string + description: Optional auth-state path. Defaults to ~/.flocks/browser/sangfor-edr/auth-state.json. +handler: + type: script + script_file: sangfor_edr.handler.py + function: handle_dashboard diff --git a/.flocks/plugins/tools/device/sangfor_edr_webcli/sangfor_edr_dashboard_api.py b/.flocks/plugins/tools/device/sangfor_edr_webcli/sangfor_edr_dashboard_api.py new file mode 100644 index 000000000..ea6a49a45 --- /dev/null +++ b/.flocks/plugins/tools/device/sangfor_edr_webcli/sangfor_edr_dashboard_api.py @@ -0,0 +1,138 @@ +"""Sangfor EDR dashboard API collection using the HTTP auth pair.""" + +from __future__ import annotations + +from datetime import datetime, timedelta +from typing import Any, Optional + +import sangfor_edr_http_login as auth + + +def _dashboard_requests( + cfg: auth.RuntimeConfig, + token: str, + *, + days: int, +) -> dict[str, tuple[str, dict[str, Any]]]: + end = datetime.now().replace(hour=23, minute=59, second=59, microsecond=0) + start = (end - timedelta(days=max(1, days) - 1)).replace(hour=0, minute=0, second=0) + unix_range = auth._unix_date_range(days) + text_range = { + "start": start.strftime("%Y-%m-%d %H:%M:%S"), + "end": end.strftime("%Y-%m-%d %H:%M:%S"), + } + + def launch(opr: str, app_name: str, extra: Optional[dict[str, Any]] = None): + payload = { + "app_args": {"name": app_name, "options": {}}, + "auto": 1, + "opr": opr, + "query_id": auth._query_id(), + } + payload.update(extra or {}) + return f"/launch.php?s={token}&opr={opr}", payload + + return { + "agent_overview": launch( + "get_agent_overview", + "app.web.event_center.head", + {"date_range": unix_range}, + ), + "influenced_agent_overview": launch( + "get_influenced_agent_overview", + "app.web.event_center.head", + ), + "vulnerability_overview": ( + f"/api/edrgoweb/v1/vulner/list/homepageVulner?_method=post&s={token}&req_type=polling", + { + "recentCheckTime": unix_range, + "uuid": "sf-id-6", + "tid": "0", + "uid": cfg.username, + "token": token, + }, + ), + "ransomware_defense": launch( + "ransom_virus_defense_interface", + "app.web.event_center.ransom_virus_protection", + ), + "realtime_virus": launch( + "real_time_virus", + "app.web.event_center.real_time_events", + {"date": text_range, "day_sum": days}, + ), + "top_agents": launch( + "get_top5_agents", + "app.web.event_center.head", + {"day_sum": days, "date_range": unix_range}, + ), + "resource_usage": launch( + "get_realtime_resource_usage", + "app.web.event_center.head", + ), + } + + +def collect_dashboard( + cfg: auth.RuntimeConfig, + *, + sections: list[str], + days: int, +) -> dict[str, Any]: + auth_result = auth.ensure_http_auth_pair(cfg) + if not auth_result.get("success"): + raise RuntimeError( + "EDR authentication refresh failed: " + f"{auth_result.get('error') or auth_result.get('reason') or auth_result.get('status')}" + ) + state, token = auth.load_verified_auth_pair(cfg) + session = auth.dashboard_session(cfg, state) + definitions = _dashboard_requests(cfg, token, days=days) + selected = sections or list(definitions) + unknown = sorted(set(selected) - set(definitions)) + if unknown: + raise ValueError(f"Unsupported EDR dashboard sections: {', '.join(unknown)}") + + data: dict[str, Any] = {} + errors: dict[str, str] = {} + for section in selected: + path, payload = definitions[section] + try: + response = session.post( + auth._url(cfg, path), + headers=auth._http_headers(cfg), + json=payload, + timeout=cfg.timeout, + ) + response.raise_for_status() + data[section] = response.json() + except Exception as exc: + errors[section] = auth._safe_error(exc, token) + + return { + "success": not errors, + "status": "dashboard_api_collected" if not errors else "dashboard_api_partially_collected", + "base_url": cfg.base_url, + "days": days, + "sections": selected, + "data": data, + "errors": errors, + "auth_pair_verified": True, + "authentication": { + "status": auth_result.get("status"), + "login_skipped": bool(auth_result.get("login_skipped")), + }, + } + + +def run_dashboard(params: dict[str, Any]) -> dict[str, Any]: + cfg = auth.resolve_runtime_config({**params, "persist_credentials": False}) + raw_sections = params.get("sections") + if isinstance(raw_sections, str): + sections = [item.strip() for item in raw_sections.split(",") if item.strip()] + elif isinstance(raw_sections, list): + sections = [str(item).strip() for item in raw_sections if str(item).strip()] + else: + sections = [] + days = max(1, min(90, auth._coerce_int(params.get("days"), 7))) + return collect_dashboard(cfg, sections=sections, days=days) diff --git a/.flocks/plugins/tools/device/sangfor_edr_webcli/sangfor_edr_http_login.py b/.flocks/plugins/tools/device/sangfor_edr_webcli/sangfor_edr_http_login.py new file mode 100644 index 000000000..149dcac01 --- /dev/null +++ b/.flocks/plugins/tools/device/sangfor_edr_webcli/sangfor_edr_http_login.py @@ -0,0 +1,523 @@ +"""Standalone Sangfor EDR HTTP authentication and auth-pair storage.""" + +from __future__ import annotations + +import hashlib +import json +import os +import secrets +import time +from datetime import datetime, timedelta +from http.cookiejar import CookieJar +from pathlib import Path +from typing import Any, Optional +from urllib.parse import urljoin, urlparse + +import requests +import urllib3 + +from flocks.config.config_writer import ConfigWriter + +SERVICE_ID = "sangfor_edr_v1_0_0" +LEGACY_SERVICE_ID = "sangfor_edr" +USERNAME_SECRET_ID = "sangfor_edr_username" +PASSWORD_SECRET_ID = "sangfor_edr_password" +TOKEN_SECRET_ID = "sangfor_edr_token" +TOKEN_BUNDLE_SECRET_ID = "sangfor_edr_token_bundle" +DEFAULT_AUTH_STATE_PATH = "~/.flocks/browser/sangfor-edr/auth-state.json" +DEFAULT_LOGIN_PATH = "/ui/login.php" +DEFAULT_TIMEOUT = 25 +PUBLIC_EXPONENT = 0x10001 +CONFIG_KEYS = ( + "base_url", + "auth_state_path", + "auto_ocr_code", + "max_captcha_retry", + "login_path", +) + + +class RuntimeConfig: + def __init__( + self, + *, + base_url: str, + auth_state_path: Path, + username: str, + password: str, + login_path: str, + timeout: int, + auto_ocr_code: bool, + max_captcha_retry: int, + ) -> None: + self.base_url = base_url + self.auth_state_path = auth_state_path + self.username = username + self.password = password + self.login_path = login_path + self.timeout = timeout + self.auto_ocr_code = auto_ocr_code + self.max_captcha_retry = max_captcha_retry + + +def _get_secret_manager(): + from flocks.security import get_secret_manager + + return get_secret_manager() + + +def _resolve_ref(value: Any) -> Optional[str]: + if value is None: + return None + if not isinstance(value, str): + return str(value) + if value.startswith("{secret:") and value.endswith("}"): + return _get_secret_manager().get(value[8:-1]) + if value.startswith("{env:") and value.endswith("}"): + return os.getenv(value[5:-1]) + return value + + +def _coerce_bool(value: Any, default: bool = False) -> bool: + if value is None: + return default + if isinstance(value, bool): + return value + return str(value).strip().lower() in {"1", "true", "yes", "y", "on"} + + +def _coerce_int(value: Any, default: int) -> int: + try: + return int(str(value).strip()) + except (TypeError, ValueError): + return default + + +def _normalise_base_url(value: str) -> str: + candidate = value.strip() + if not candidate: + raise ValueError("Sangfor EDR base_url is required.") + if "://" not in candidate: + candidate = f"https://{candidate}" + parsed = urlparse(candidate) + if parsed.scheme not in {"http", "https"} or not parsed.hostname: + raise ValueError(f"Invalid Sangfor EDR base_url: {value!r}") + host = parsed.hostname + if ":" in host and not host.startswith("["): + host = f"[{host}]" + port = f":{parsed.port}" if parsed.port else "" + return f"{parsed.scheme}://{host}{port}".rstrip("/") + + +def _has_device_context() -> bool: + try: + from flocks.tool.credential_context import get_active_device_id + + return bool(get_active_device_id()) + except Exception: + return False + + +def _load_service_config() -> dict[str, Any]: + primary = ConfigWriter.get_api_service_raw(SERVICE_ID) + primary = dict(primary) if isinstance(primary, dict) else {} + if _has_device_context(): + return primary + legacy = ConfigWriter.list_api_services_raw() + fallback = legacy.get(LEGACY_SERVICE_ID) if isinstance(legacy, dict) else {} + for key, value in (fallback or {}).items(): + if primary.get(key) in (None, "") and value not in (None, ""): + primary[key] = value + return primary + + +def _save_params_to_service(params: dict[str, Any]) -> dict[str, Any]: + service = _load_service_config() + persist = _coerce_bool(params.get("persist_credentials"), default=True) + for key in CONFIG_KEYS: + if params.get(key) not in (None, ""): + service[key] = params[key] + for key, secret_id in (("username", USERNAME_SECRET_ID), ("password", PASSWORD_SECRET_ID)): + value = params.get(key) + if isinstance(value, str) and value: + if persist: + _get_secret_manager().set(secret_id, value) + service[key] = f"{{secret:{secret_id}}}" + else: + service[key] = value + if persist and any(params.get(key) not in (None, "") for key in ("base_url", "auth_state_path", "username", "password")): + ConfigWriter.set_api_service(SERVICE_ID, service) + return service + + +def resolve_runtime_config(params: dict[str, Any]) -> RuntimeConfig: + raw = _save_params_to_service(params) + secrets_store = _get_secret_manager() + base_url = _normalise_base_url( + _resolve_ref(raw.get("base_url")) + or _resolve_ref(raw.get("host")) + or os.getenv("SANGFOR_EDR_BASE_URL") + or "" + ) + state_path = Path( + _resolve_ref(raw.get("auth_state_path")) + or os.getenv("SANGFOR_EDR_AUTH_STATE") + or DEFAULT_AUTH_STATE_PATH + ).expanduser() + username = ( + _resolve_ref(raw.get("username")) + or secrets_store.get(USERNAME_SECRET_ID) + or secrets_store.get(f"{SERVICE_ID}_username") + or secrets_store.get(f"{LEGACY_SERVICE_ID}_username") + or os.getenv("SANGFOR_EDR_USERNAME") + or "" + ).strip() + password = ( + _resolve_ref(raw.get("password")) + or secrets_store.get(PASSWORD_SECRET_ID) + or secrets_store.get(f"{SERVICE_ID}_password") + or secrets_store.get(f"{LEGACY_SERVICE_ID}_password") + or os.getenv("SANGFOR_EDR_PASSWORD") + or "" + ).strip() + return RuntimeConfig( + base_url=base_url, + auth_state_path=state_path, + username=username, + password=password, + login_path=str(raw.get("login_path") or DEFAULT_LOGIN_PATH), + timeout=max(5, _coerce_int(raw.get("timeout"), DEFAULT_TIMEOUT)), + auto_ocr_code=_coerce_bool(raw.get("auto_ocr_code"), default=True), + max_captcha_retry=max(1, min(10, _coerce_int(raw.get("max_captcha_retry"), 5))), + ) + + +def _url(cfg: RuntimeConfig, path: str) -> str: + return urljoin(cfg.base_url + "/", path.lstrip("/")) + + +def _now_ms() -> str: + return str(int(time.time() * 1000)) + + +def _query_id() -> str: + return f"Query_{_now_ms()}" + + +def _safe_error(exc: Exception, *sensitive_values: str) -> str: + message = str(exc) + for value in sensitive_values: + if value: + message = message.replace(value, "") + return message + + +def _http_headers(cfg: RuntimeConfig, *, image: bool = False) -> dict[str, str]: + headers = { + "Pragma": "no-cache", + "Cache-Control": "no-cache", + "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 Chrome/131.0.0.0 Safari/537.36", + "Accept-Language": "zh-CN,zh;q=0.9", + "Referer": _url(cfg, cfg.login_path) if image else _url(cfg, "/ui/"), + } + if image: + headers["Accept"] = "image/avif,image/webp,image/apng,image/svg+xml,image/*,*/*;q=0.8" + else: + headers.update({"Accept": "application/json, text/plain, */*", "Content-Type": "application/json", "Origin": cfg.base_url}) + return headers + + +def http_session(cfg: RuntimeConfig) -> requests.Session: + session = requests.Session() + session.verify = False + session.headers.update({"Accept-Language": "zh-CN,zh;q=0.9"}) + session.cookies.set("hadSetUkey", "0", domain=urlparse(cfg.base_url).hostname, path="/") + urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning) + return session + + +def _cookie_state(cookies: CookieJar, cfg: RuntimeConfig) -> list[dict[str, Any]]: + domain = urlparse(cfg.base_url).hostname or "" + result = [] + for cookie in cookies: + item = { + "name": cookie.name, + "value": cookie.value, + "domain": cookie.domain or domain, + "path": cookie.path or "/", + "secure": bool(cookie.secure), + "httpOnly": bool(cookie.has_nonstandard_attr("HttpOnly")), + "sameSite": "Lax", + } + if cookie.expires and cookie.expires > 0: + item["expires"] = float(cookie.expires) + result.append(item) + return result + + +def _cookie_fingerprint(cookies: list[dict[str, Any]], base_url: str) -> str: + normalized = sorted( + ( + str(cookie.get("name") or ""), + str(cookie.get("value") or ""), + str(cookie.get("domain") or ""), + str(cookie.get("path") or "/"), + ) + for cookie in cookies + if isinstance(cookie, dict) and cookie.get("name") + ) + raw = json.dumps({"base_url": _normalise_base_url(base_url), "cookies": normalized}, separators=(",", ":")) + return hashlib.sha256(raw.encode()).hexdigest() + + +def _read_auth_state(path: Path) -> dict[str, Any]: + state = json.loads(path.read_text(encoding="utf-8")) + if not isinstance(state, dict) or not isinstance(state.get("cookies"), list): + raise ValueError("EDR auth-state does not contain a cookie list.") + return state + + +def _save_auth_pair(cfg: RuntimeConfig, state: dict[str, Any], token: str) -> dict[str, Any]: + cookies = state.get("cookies") + if not isinstance(cookies, list) or not cookies or not token: + raise ValueError("EDR login did not return a complete cookie/token pair.") + cfg.auth_state_path.parent.mkdir(parents=True, exist_ok=True) + temp_path = cfg.auth_state_path.with_name(f"{cfg.auth_state_path.name}.tmp") + temp_path.write_text(json.dumps(state, ensure_ascii=False, indent=2), encoding="utf-8") + temp_path.replace(cfg.auth_state_path) + fingerprint = _cookie_fingerprint(cookies, cfg.base_url) + manager = _get_secret_manager() + manager.set(TOKEN_SECRET_ID, token) + manager.set(TOKEN_BUNDLE_SECRET_ID, json.dumps({"token": token, "base_url": cfg.base_url, "cookie_fingerprint": fingerprint}, separators=(",", ":"))) + return {"auth_state_path": str(cfg.auth_state_path), "cookie_count": len(cookies), "pair_verified": True} + + +def load_verified_auth_pair(cfg: RuntimeConfig) -> tuple[dict[str, Any], str]: + state = _read_auth_state(cfg.auth_state_path) + raw_bundle = _get_secret_manager().get(TOKEN_BUNDLE_SECRET_ID) + if not raw_bundle: + raise RuntimeError("EDR cookie/token pair metadata is missing; refresh authentication.") + try: + bundle = json.loads(raw_bundle) + except (TypeError, json.JSONDecodeError) as exc: + raise RuntimeError("EDR cookie/token pair metadata is invalid; refresh authentication.") from exc + token = str(bundle.get("token") or "") + if ( + not token + or _normalise_base_url(str(bundle.get("base_url") or "")) != cfg.base_url + or not secrets.compare_digest(str(bundle.get("cookie_fingerprint") or ""), _cookie_fingerprint(state["cookies"], cfg.base_url)) + ): + raise RuntimeError("EDR cookies and login_token are not from the same login; refresh authentication.") + return state, token + + +def dashboard_session(cfg: RuntimeConfig, state: dict[str, Any]) -> requests.Session: + session = requests.Session() + session.verify = False + host = urlparse(cfg.base_url).hostname or "" + for cookie in state.get("cookies", []): + if isinstance(cookie, dict) and cookie.get("name"): + session.cookies.set(str(cookie["name"]), str(cookie.get("value") or ""), domain=str(cookie.get("domain") or host), path=str(cookie.get("path") or "/")) + urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning) + return session + + +def _unix_date_range(days: int = 7) -> dict[str, int]: + end = datetime.now().replace(hour=23, minute=59, second=59, microsecond=0) + start = (end - timedelta(days=max(1, days) - 1)).replace(hour=0, minute=0, second=0) + return {"start": int(start.timestamp()), "end": int(end.timestamp())} + + +def _response_contains_agent_overview(payload: Any) -> bool: + return isinstance(payload, dict) and bool(payload.get("success")) and any( + key in payload and payload.get(key) is not None + for key in ("data", "result", "agent_overview", "agent_total", "total", "count") + ) + + +def probe_auth_pair(cfg: RuntimeConfig) -> dict[str, Any]: + try: + state, token = load_verified_auth_pair(cfg) + except Exception as exc: + return {"valid": False, "reason": "auth_pair_missing_or_mismatched", "error": str(exc)} + session = dashboard_session(cfg, state) + try: + response = session.post( + _url(cfg, f"/launch.php?s={token}&opr=get_agent_overview"), + headers=_http_headers(cfg), + json={"app_args": {"name": "app.web.event_center.head", "options": {}}, "auto": 1, "opr": "get_agent_overview", "date_range": _unix_date_range(), "query_id": _query_id()}, + timeout=cfg.timeout, + allow_redirects=False, + ) + except Exception as exc: + return {"valid": False, "reason": "auth_probe_request_failed", "error": _safe_error(exc, token)} + if response.status_code in {301, 302, 303, 307, 308}: + return {"valid": False, "reason": "auth_probe_redirected", "http_status": response.status_code, "location": str(response.headers.get("Location") or "")} + if response.status_code in {401, 403}: + return {"valid": False, "reason": "auth_probe_unauthorized", "http_status": response.status_code} + if response.status_code != 200: + return {"valid": False, "reason": "auth_probe_http_error", "http_status": response.status_code} + try: + payload = response.json() + except Exception as exc: + return {"valid": False, "reason": "auth_probe_invalid_json", "error": str(exc)} + if not isinstance(payload, dict) or not payload.get("success"): + return {"valid": False, "reason": "auth_probe_rejected"} + if not _response_contains_agent_overview(payload): + return {"valid": False, "reason": "auth_probe_expected_agent_data_missing"} + return {"valid": True, "reason": "auth_probe_succeeded", "http_status": 200, "agent_overview_verified": True} + + +def _rsa_encrypt_password(rsa_key_hex: str, password: str) -> str: + modulus = int(rsa_key_hex.strip().lower().removeprefix("0x"), 16) + key_size = (modulus.bit_length() + 7) // 8 + message = password.encode("utf-8") + if key_size < len(message) + 11: + raise ValueError("EDR password is too long for the login RSA key.") + padding = bytearray() + while len(padding) < key_size - len(message) - 3: + padding.extend(byte for byte in secrets.token_bytes(key_size) if byte) + encoded = b"\x00\x02" + bytes(padding[: key_size - len(message) - 3]) + b"\x00" + message + return f"{pow(int.from_bytes(encoded, 'big'), PUBLIC_EXPONENT, modulus):0{key_size * 2}x}" + + +def _ocr_verify_code(image: bytes) -> str: + try: + import ddddocr + except Exception as exc: + raise RuntimeError("ddddocr is required for automatic Sangfor EDR captcha recognition.") from exc + return ddddocr.DdddOcr(show_ad=False).classification(image).strip()[:4] + + +def _is_captcha_failure(result: dict[str, Any]) -> bool: + message = str(result.get("msg") or "") + data = result.get("data") if isinstance(result.get("data"), dict) else {} + return "验证码" in message or data.get("code") is True + + +def _http_login(cfg: RuntimeConfig, captcha_code: str = "") -> dict[str, Any]: + missing = [key for key, value in (("username", cfg.username), ("password", cfg.password)) if not value] + if missing: + return {"success": False, "status": "http_login_credentials_required", "reason": "missing_http_login_credentials", "missing": missing} + session = http_session(cfg) + phase = "session_init" + try: + phase = "login_page" + session.get(_url(cfg, cfg.login_path), headers=_http_headers(cfg), timeout=cfg.timeout).raise_for_status() + phase = "rsakey" + rsa_response = session.post(_url(cfg, "/login"), headers=_http_headers(cfg), json={"opr": "rsakey"}, timeout=cfg.timeout) + rsa_response.raise_for_status() + rsa_result = rsa_response.json() + rsa_key = str(rsa_result.get("key") or "") + if not rsa_result.get("success") or not rsa_key: + raise RuntimeError("EDR RSA key request failed.") + last_error = "" + for attempt in range(1, cfg.max_captcha_retry + 1): + code = captcha_code.strip() + if not code: + if not cfg.auto_ocr_code: + return {"success": False, "status": "http_login_captcha_required", "reason": "captcha_code_required"} + phase = "captcha" + captcha = session.get(_url(cfg, f"/ui/randcode.php?{_now_ms()}"), headers=_http_headers(cfg, image=True), timeout=cfg.timeout) + captcha.raise_for_status() + code = _ocr_verify_code(captcha.content) + phase = "dlogin" + login_response = session.post( + _url(cfg, "/login"), + headers=_http_headers(cfg), + json={"opr": "dlogin", "data": {"auth_type": "pwd", "user_name": cfg.username, "code": code, "pwd": _rsa_encrypt_password(rsa_key, cfg.password)}}, + timeout=cfg.timeout, + ) + login_response.raise_for_status() + login_result = login_response.json() + if not login_result.get("success") or login_result.get("key") in (None, ""): + last_error = str(login_result.get("msg") or "EDR dlogin failed.") + if captcha_code or not _is_captcha_failure(login_result): + raise RuntimeError(last_error) + continue + phase = "launch_login" + launch_response = session.post( + _url(cfg, "/launch_login.php"), + headers=_http_headers(cfg), + json={"opr": "dlogin", "app_args": {"name": "app.web.auth.login", "options": {}}, "data": {"key": login_result["key"], "user_aggreement_status": "true"}}, + timeout=cfg.timeout, + ) + launch_response.raise_for_status() + launch_result = launch_response.json() + token = str((launch_result.get("data") or {}).get("token") or "") + if launch_result.get("success") and token: + phase = "ui" + session.get( + _url(cfg, "/ui"), + headers=_http_headers(cfg), + timeout=cfg.timeout, + ).raise_for_status() + cookies = _cookie_state(session.cookies, cfg) + if launch_result.get("success") and token and any(cookie["name"].lower() == "sessionid" for cookie in cookies): + saved = _save_auth_pair(cfg, {"cookies": cookies, "origins": []}, token) + return {"success": True, "valid": True, "status": "http_login_refreshed_auth_state", "attempt": attempt, "token_saved": True, "saved": saved} + last_error = str(launch_result.get("msg") or "incomplete EDR cookie/token response") + if captcha_code: + break + raise RuntimeError(last_error or "EDR HTTP login retry limit exceeded.") + except Exception as exc: + response = getattr(exc, "response", None) + status = getattr(response, "status_code", None) + suffix = f"; http_status={status}" if status is not None else "" + return {"success": False, "valid": False, "status": "http_login_failed", "reason": "http_login_failed", "error": f"phase={phase}{suffix}; detail={_safe_error(exc, cfg.username, cfg.password)}", "phase": phase} + + +def ensure_http_auth_pair(cfg: RuntimeConfig, captcha_code: str = "") -> dict[str, Any]: + probe = probe_auth_pair(cfg) + if probe.get("valid"): + return {"success": True, "valid": True, "status": "http_auth_pair_reused", "login_skipped": True, "probe": probe} + result = _http_login(cfg, captcha_code=captcha_code) + result["previous_probe"] = probe + if result.get("success"): + confirmation = probe_auth_pair(cfg) + result["probe"] = confirmation + if not confirmation.get("valid"): + result.update({"success": False, "valid": False, "status": "http_login_probe_failed", "reason": str(confirmation.get("reason") or "http_login_probe_failed")}) + return result + + +def status_auth_state(params: dict[str, Any]) -> dict[str, Any]: + service = _save_params_to_service({**params, "persist_credentials": False}) + manager = _get_secret_manager() + base_url = _resolve_ref(service.get("base_url")) or os.getenv("SANGFOR_EDR_BASE_URL") or "" + path = Path(_resolve_ref(service.get("auth_state_path")) or os.getenv("SANGFOR_EDR_AUTH_STATE") or DEFAULT_AUTH_STATE_PATH).expanduser() + username = _resolve_ref(service.get("username")) or manager.get(USERNAME_SECRET_ID) or "" + password = _resolve_ref(service.get("password")) or manager.get(PASSWORD_SECRET_ID) or "" + probe = None + if base_url and path.exists(): + try: + probe = probe_auth_pair(resolve_runtime_config({**params, "persist_credentials": False})) + except Exception as exc: + probe = {"valid": False, "reason": "auth_probe_failed", "error": str(exc)} + return { + "auth_state_path": str(path), + "auth_state_exists": path.exists(), + "has_base_url": bool(str(base_url).strip()), + "has_saved_username": bool(str(username).strip()), + "has_saved_password": bool(str(password).strip()), + "has_saved_token": bool(manager.get(TOKEN_SECRET_ID)), + "has_verified_auth_pair": bool(probe and probe.get("valid")), + "auth_probe": probe, + "can_auto_refresh": bool(str(base_url).strip() and str(username).strip() and str(password).strip()), + } + + +def run_auth_action(params: dict[str, Any]) -> dict[str, Any]: + action = str(params.get("action") or "ensure_auth_state") + if action == "status_auth_state": + status = status_auth_state(params) + return { + "success": True, + "status": "saved_auto_login_status", + **status, + "validation": status.get("auth_probe"), + } + if action not in {"ensure_auth_state", "refresh_auth_state", "http_login"}: + raise ValueError(f"Unsupported HTTP EDR auth action: {action}") + cfg = resolve_runtime_config(params) + return ensure_http_auth_pair(cfg, captcha_code=str(params.get("captcha_code") or "")) diff --git a/.github/workflows/dispatch-autotest-windows-upgrade.yml b/.github/workflows/dispatch-autotest-windows-upgrade.yml index 59029c4bf..f300a31b1 100644 --- a/.github/workflows/dispatch-autotest-windows-upgrade.yml +++ b/.github/workflows/dispatch-autotest-windows-upgrade.yml @@ -1,17 +1,16 @@ name: Dispatch autotest Windows upgrade on: - push: - branches: - - main + release: + types: [published] workflow_dispatch: inputs: release_tag: - description: "Optional Flocks version tag. Defaults to pyproject.toml version." + description: "Optional expected version. It must match project.version in pyproject.toml." required: false type: string release_url: - description: "Optional source URL. Defaults to the release or commit URL." + description: "Optional source URL. Defaults to the published release URL or commit URL." required: false type: string rollback_version: @@ -48,6 +47,9 @@ jobs: INPUT_RELEASE_URL: ${{ inputs.release_url }} INPUT_ROLLBACK_VERSION: ${{ inputs.rollback_version }} INPUT_RUN_FORCE_FALLBACK: ${{ inputs.run_force_fallback }} + RELEASE_EVENT_TAG: ${{ github.event.release.tag_name }} + RELEASE_EVENT_URL: ${{ github.event.release.html_url }} + RELEASE_EVENT_ASSETS_JSON: ${{ toJson(github.event.release.assets) }} GITHUB_SERVER_URL_VALUE: ${{ github.server_url }} GITHUB_REPOSITORY_VALUE: ${{ github.repository }} GITHUB_SHA_VALUE: ${{ github.sha }} @@ -64,21 +66,37 @@ jobs: PY )" - release_tag="${INPUT_RELEASE_TAG:-}" - release_url="${INPUT_RELEASE_URL:-}" + requested_release_tag="${INPUT_RELEASE_TAG:-${RELEASE_EVENT_TAG:-}}" + release_url="${INPUT_RELEASE_URL:-${RELEASE_EVENT_URL:-}}" + installer_asset_name="$(python3 - <<'PY' + import json + import os + + assets_json = os.environ.get("RELEASE_EVENT_ASSETS_JSON") or "[]" + try: + assets = json.loads(assets_json) + except json.JSONDecodeError: + assets = [] + if not isinstance(assets, list): + assets = [] + + for asset in assets: + name = str(asset.get("name") or "") + if name.lower().endswith(".exe"): + print(name) + break + PY + )" - if [ -z "$release_tag" ]; then - release_tag="$pyproject_version" + if [ -n "$requested_release_tag" ] && [ "$requested_release_tag" != "$pyproject_version" ]; then + echo "::error::Release version mismatch: requested=$requested_release_tag pyproject=$pyproject_version" + exit 1 fi + release_tag="$pyproject_version" if [ -z "$release_url" ]; then release_url="${GITHUB_SERVER_URL_VALUE}/${GITHUB_REPOSITORY_VALUE}/commit/${GITHUB_SHA_VALUE}" fi - if [ -z "$release_tag" ]; then - echo "release_tag is empty" >&2 - exit 1 - fi - artifact_time="$(printf "%02d" "${GITHUB_RUN_ATTEMPT_VALUE:-1}")" artifact_name="flocks-windows-upgrade-${release_tag}-time${artifact_time}" @@ -88,16 +106,19 @@ jobs: echo "rollback_version=${INPUT_ROLLBACK_VERSION:-}" echo "run_force_fallback=${INPUT_RUN_FORCE_FALLBACK:-false}" echo "artifact_name=$artifact_name" + echo "installer_asset_name=$installer_asset_name" } >> "$GITHUB_OUTPUT" { echo "### flocks_autotest dispatch payload" echo "" echo "- Event: ${EVENT_NAME}" - echo "- Release tag: ${release_tag}" + echo "- Release tag: ${requested_release_tag:-not provided}" + echo "- pyproject.toml version: ${release_tag}" echo "- Release URL: ${release_url}" echo "- Source ref: ${GITHUB_REF_NAME_VALUE}" echo "- Source sha: ${GITHUB_SHA_VALUE}" + echo "- Installer asset: ${installer_asset_name}" echo "- Rollback version: ${INPUT_ROLLBACK_VERSION:-}" echo "- Run force fallback: ${INPUT_RUN_FORCE_FALLBACK:-false}" echo "- Expected autotest artifact: ${artifact_name}" @@ -111,6 +132,7 @@ jobs: ROLLBACK_VERSION: ${{ steps.payload.outputs.rollback_version }} RUN_FORCE_FALLBACK: ${{ steps.payload.outputs.run_force_fallback }} ARTIFACT_NAME: ${{ steps.payload.outputs.artifact_name }} + INSTALLER_ASSET_NAME: ${{ steps.payload.outputs.installer_asset_name }} SOURCE_REPOSITORY: ${{ github.repository }} SOURCE_RUN_ID: ${{ github.run_id }} SOURCE_SHA: ${{ github.sha }} @@ -123,11 +145,11 @@ jobs: import os payload = { - "event_type": "flocks_main_updated", + "event_type": "flocks_release_published", "client_payload": { "release_tag": os.environ["RELEASE_TAG"], "release_url": os.environ["RELEASE_URL"], - "installer_asset_name": "", + "installer_asset_name": os.environ.get("INSTALLER_ASSET_NAME", ""), "source_repository": os.environ["SOURCE_REPOSITORY"], "source_run_id": os.environ["SOURCE_RUN_ID"], "source_sha": os.environ["SOURCE_SHA"], @@ -183,6 +205,7 @@ jobs: python3 - <<'PY' import json import os + import re import sys import time import urllib.error @@ -263,8 +286,15 @@ jobs: f"{api_root}/repos/{autotest_repo}/actions/workflows/" f"{workflow_ref}/runs?per_page=50" ) - version_title_prefix = f"Windows upgrade {release_tag} " - version_artifact_prefix = f"flocks-windows-upgrade-{release_tag}-" + version_title_pattern = re.compile( + rf"(?= lookup_deadline: break available_text = ", ".join(available_artifact_names) if available_artifact_names else "(none)" - expected_text = exact_name or f"{name_prefix}*" + expected_text = exact_name or f"artifact containing exact version {release_tag}" print( f"Waiting for artifact {expected_text}. Available artifacts: {available_text}", flush=True, @@ -406,7 +445,7 @@ jobs: fallback_run_id = fallback_run["id"] artifact, available_names = find_artifact( fallback_run_id, - name_prefix=version_artifact_prefix, + match_release_version=True, retry_seconds=120, ) if artifact: @@ -467,7 +506,7 @@ jobs: f"- Conclusion: {conclusion}", f"- Source run: {source_run_url or 'not found'}", f"- Source status: {source_failure or 'ok'}", - f"- Version fallback prefix: {version_artifact_prefix}", + f"- Version fallback match: exact version {release_tag} anywhere in artifact name", f"- Expected artifact: {expected_artifact_name}", ] if artifact_name: diff --git a/flocks/acp/agent.py b/flocks/acp/agent.py index ae2ed6f6f..3025fa1e9 100644 --- a/flocks/acp/agent.py +++ b/flocks/acp/agent.py @@ -1334,9 +1334,13 @@ async def _restart_session_state(self, session_id: str) -> None: from flocks.session.message import Message from flocks.session.features.todo import Todo from flocks.session.core.status import SessionStatus + from flocks.session.session import Session - await Message.clear(session_id) - await Todo.clear(session_id) + async def clear_persisted_state() -> None: + await Message.clear(session_id) + await Todo.clear(session_id) + + await Session.run_active_write(session_id, clear_persisted_state) SessionStatus.clear(session_id) session = self._session_manager.get(session_id) diff --git a/flocks/agent/agents/rex/prompt_builder.py b/flocks/agent/agents/rex/prompt_builder.py index 281234bbc..a4bc0a01c 100644 --- a/flocks/agent/agents/rex/prompt_builder.py +++ b/flocks/agent/agents/rex/prompt_builder.py @@ -256,6 +256,11 @@ def _build_command_guidance_section() -> str: Use `flocks --help` to inspect Flocks CLI commands and usage. run_slash_command tool with help command to get the latest slash command guidance. + +### Safe Flocks Restart + +When you need to restart the running Flocks service yourself, use `flocks restart --server-only`. +永远不要直接执行 `flocks restart`,这将会导致你杀死自己并且无法自启动。 """ diff --git a/flocks/auth/context.py b/flocks/auth/context.py index c95d7b6af..7ce35f985 100644 --- a/flocks/auth/context.py +++ b/flocks/auth/context.py @@ -10,6 +10,9 @@ from pydantic import BaseModel, Field +API_TOKEN_SERVICE_USER_ID = "api-token-service" + + class AuthUser(BaseModel): """Current authenticated local user.""" diff --git a/flocks/auth/service.py b/flocks/auth/service.py index 3c5d761c0..3db79323d 100644 --- a/flocks/auth/service.py +++ b/flocks/auth/service.py @@ -15,7 +15,7 @@ import aiosqlite from pydantic import BaseModel, Field -from flocks.auth.context import AuthUser +from flocks.auth.context import API_TOKEN_SERVICE_USER_ID, AuthUser from flocks.extensions import ensure_callable_methods from flocks.storage.storage import Storage from flocks.utils.id import Identifier @@ -256,6 +256,8 @@ async def _create_user_internal( normalized_username = username.strip() if not normalized_username: raise ValueError("用户名不能为空") + if normalized_username == API_TOKEN_SERVICE_USER_ID: + raise ValueError("该用户名为系统保留用户名") if len(password) < 8: raise ValueError("密码长度至少 8 位") @@ -635,6 +637,7 @@ async def reassign_orphan_sessions( await Session.update( project_id=session.project_id, session_id=session.id, + allow_inactive=True, owner_user_id=admin_user_id, owner_username=admin_user.username, ) @@ -672,6 +675,7 @@ async def migrate_legacy_sessions_to_admin(cls, admin_user_id: str) -> None: await Session.update( project_id=session.project_id, session_id=session.id, + allow_inactive=True, owner_user_id=admin_user_id, owner_username=admin_username, ) diff --git a/flocks/channel/inbound/dispatcher.py b/flocks/channel/inbound/dispatcher.py index 44a88f566..1e516c91b 100644 --- a/flocks/channel/inbound/dispatcher.py +++ b/flocks/channel/inbound/dispatcher.py @@ -1053,16 +1053,38 @@ async def _get_channel_config(channel_id: str) -> ChannelConfig: except Exception: return ChannelConfig() + @staticmethod + async def _run_session_loop(binding: Any, loop_callbacks: Any) -> Any: + """Run an IM turn with the persisted model preference for its session.""" + from flocks.session.session import ( + Session, + is_model_auto_session_category, + ) + from flocks.session.session_loop import SessionLoop + + session = await Session.get_by_id(binding.session_id) + auto_failover = bool( + session + and is_model_auto_session_category( + getattr(session, "category", "user") + ) + and getattr(session, "model_auto", False) + ) + return await SessionLoop.run( + session_id=binding.session_id, + agent_name=binding.agent_id, + callbacks=loop_callbacks, + auto_failover=auto_failover, + ) + @staticmethod async def _run_agent(binding, callbacks: ChannelDeliveryCallbacks) -> None: """Run Agent and deliver the final assistant reply.""" try: - from flocks.session.session_loop import SessionLoop loop_callbacks = callbacks.to_loop_callbacks() - result = await SessionLoop.run( - session_id=binding.session_id, - agent_name=binding.agent_id, - callbacks=loop_callbacks, + result = await InboundDispatcher._run_session_loop( + binding, + loop_callbacks, ) if result.last_message: @@ -1147,11 +1169,9 @@ async def _on_text_delta(delta: str) -> None: runner_cbs = RunnerCallbacks(on_text_delta=_on_text_delta) loop_callbacks = callbacks.to_loop_callbacks(runner_callbacks=runner_cbs) - from flocks.session.session_loop import SessionLoop - result = await SessionLoop.run( - session_id=binding.session_id, - agent_name=binding.agent_id, - callbacks=loop_callbacks, + result = await InboundDispatcher._run_session_loop( + binding, + loop_callbacks, ) final_text = None @@ -1181,6 +1201,31 @@ async def _append_user_message( channel_config: Optional[ChannelConfig] = None, model: Optional[dict] = None, agent: Optional[str] = None, + ) -> None: + """Append an inbound message only while its bound session is active.""" + + from flocks.session.session import Session + + await Session.run_active_write( + session_id, + lambda: InboundDispatcher._append_user_message_unchecked( + session_id, + text, + msg, + channel_config, + model, + agent, + ), + ) + + @staticmethod + async def _append_user_message_unchecked( + session_id: str, + text: str, + msg: InboundMessage, + channel_config: Optional[ChannelConfig] = None, + model: Optional[dict] = None, + agent: Optional[str] = None, ) -> None: import mimetypes import os diff --git a/flocks/channel/inbound/session_binding.py b/flocks/channel/inbound/session_binding.py index c9e49cd2b..3f023a358 100644 --- a/flocks/channel/inbound/session_binding.py +++ b/flocks/channel/inbound/session_binding.py @@ -86,8 +86,7 @@ async def resolve_channel_session_owner_kwargs(source_session=None) -> dict[str, ``Session.create`` cannot infer the owner from ``current_auth_user``. When an existing channel session is being replaced, preserve its owner. Otherwise, attach new channel sessions to the local admin if one exists. - Installs without local accounts remain ownerless for backward-compatible - no-login operation. + Installs without local accounts use the explicit system identity. """ owner_user_id = getattr(source_session, "owner_user_id", None) if source_session else None owner_username = getattr(source_session, "owner_username", None) if source_session else None @@ -103,7 +102,12 @@ async def resolve_channel_session_owner_kwargs(source_session=None) -> dict[str, from flocks.auth.service import AuthService if not await AuthService.has_users(): - return {} + from flocks.auth.context import API_TOKEN_SERVICE_USER_ID + + return { + "owner_user_id": API_TOKEN_SERVICE_USER_ID, + "owner_username": API_TOKEN_SERVICE_USER_ID, + } users = await AuthService.list_users() except Exception as exc: log.warn("channel.owner.resolve_failed", {"error": str(exc)}) @@ -288,23 +292,30 @@ async def resolve_or_create( existing = await self._find_binding( msg.channel_id, msg.account_id, chat_id, thread_id, ) + replaced_session = None if existing: - # Verify the bound session still exists (user may have deleted it via WebUI) + # Archived sessions are immutable history, not live conversation + # targets. Replace stale or inactive bindings before the dispatcher + # persists the inbound message. from flocks.session.session import Session as _Session - still_alive = await _Session.get_by_id(existing.session_id) - if still_alive: + bound_session = await _Session.get_by_id_unfiltered(existing.session_id) + if bound_session and bound_session.status == "active": await self._touch(existing.session_id) return existing - # Session was deleted — remove stale binding and fall through to create a new one + replaced_session = bound_session log.info("channel.binding.stale", { "channel": msg.channel_id, "chat_id": chat_id, "old_session_id": existing.session_id, + "status": getattr(bound_session, "status", "missing"), }) await self.unbind(existing.session_id) session_id = await self._create_session( - msg, default_agent=default_agent, directory=directory, + msg, + default_agent=default_agent, + directory=directory, + source_session=replaced_session, ) now = time.time() binding = SessionBinding( @@ -364,8 +375,11 @@ async def bind_session( ValueError: if *session_id* does not exist. """ from flocks.session.session import Session as _Session - if not await _Session.get_by_id(session_id): + session = await _Session.get_by_id_unfiltered(session_id) + if not session: raise ValueError(f"Session '{session_id}' not found") + if session.status != "active": + raise ValueError(f"Session '{session_id}' is not active") now = time.time() binding = SessionBinding( @@ -556,6 +570,7 @@ async def _create_session( msg: InboundMessage, default_agent: Optional[str] = None, directory: Optional[str] = None, + source_session=None, ) -> str: """Create a new Flocks Session and return its ID. @@ -568,7 +583,7 @@ async def _create_session( from flocks.session.session import Session title = _build_title(msg) - owner_kwargs = await resolve_channel_session_owner_kwargs() + owner_kwargs = await resolve_channel_session_owner_kwargs(source_session) session = await Session.create( project_id="channel", directory=_resolve_session_directory(directory), diff --git a/flocks/cli/commands/session.py b/flocks/cli/commands/session.py index fb13111f3..c23f1384f 100644 --- a/flocks/cli/commands/session.py +++ b/flocks/cli/commands/session.py @@ -208,21 +208,20 @@ async def _show_session(session_id: str, project_id: Optional[str]): @session_app.command("delete") def session_delete( - session_id: str = typer.Argument(..., help="Session ID to delete"), + session_id: str = typer.Argument(..., help="Session ID to permanently delete"), project: Optional[str] = typer.Option( None, "-p", "--project", help="Project ID (uses current project if not specified)" ), force: bool = typer.Option( False, "-f", "--force", - help="Skip confirmation prompt" + help="Skip permanent deletion confirmation" ), ): """ - Delete a session - - This performs a soft delete. The session data is marked as deleted - but not permanently removed from storage. + Permanently delete a session and its persisted history. + + This operation cannot be restored. """ asyncio.run(_delete_session(session_id, project, force)) @@ -241,7 +240,7 @@ async def _delete_session(session_id: str, project_id: Optional[str], force: boo # Confirm deletion if not force: confirm = typer.confirm( - f"Delete session '{session.title}'?", + f"Permanently delete session '{session.title}' and all of its messages and history?", default=False ) if not confirm: @@ -252,7 +251,7 @@ async def _delete_session(session_id: str, project_id: Optional[str], force: boo success = await Session.delete(project_id, session_id) if success: - console.print(f"[green]Deleted session: {session_id}[/green]") + console.print(f"[green]Permanently deleted session: {session_id}[/green]") else: console.print(f"[red]Failed to delete session: {session_id}[/red]") raise typer.Exit(1) @@ -317,9 +316,16 @@ async def _restore_session(session_id: str, project_id: Optional[str]): console.print(f"[red]Failed to restore session (not found or not archived): {session_id}[/red]") raise typer.Exit(1) - project_id, _ = resolved + project_id, session = resolved + if session.parent_id is not None: + console.print(f"[red]Restore the root session for this session tree: {session_id}[/red]") + raise typer.Exit(1) - success = await Session.unarchive(project_id, session_id) + success = await Session.restore( + project_id, + session_id, + project_owner_id=session.owner_user_id, + ) if success: console.print(f"[green]Restored session: {session_id}[/green]") diff --git a/flocks/cli/main.py b/flocks/cli/main.py index 8fce87905..626b9f5a7 100644 --- a/flocks/cli/main.py +++ b/flocks/cli/main.py @@ -41,6 +41,7 @@ ServiceError, resolve_flocks_cli_command, restart_all, + restart_server, runtime_paths, show_logs, show_status, @@ -264,6 +265,11 @@ def stop(): @app.command() def restart( + server_only: bool = typer.Option( + False, + "--server-only", + help="Restart only the backend server without stopping the supervisor daemon", + ), no_browser: bool = typer.Option(False, "--no-browser", help="Do not open WebUI in a browser"), skip_webui_build: bool = typer.Option( False, @@ -281,6 +287,9 @@ def restart( Restart Flocks service. """ try: + if server_only: + restart_server(console) + return restart_all( _restart_service_config( no_browser=no_browser, diff --git a/flocks/cli/service_manager.py b/flocks/cli/service_manager.py index cd52842ef..7e90c41c4 100644 --- a/flocks/cli/service_manager.py +++ b/flocks/cli/service_manager.py @@ -31,6 +31,7 @@ read_logs, read_supervisor_status, request_restart, + request_restart_backend, request_stop, stream_logs, supervisor_is_running, @@ -60,6 +61,7 @@ "src\\win\\async.c", "src/win/async.c", ) +WEBUI_BUILD_IGNORED_DIRS = frozenset({"dist", "node_modules", ".vite"}) WATCHDOG_PID_FILENAME = "watchdog.pid" SUPERVISOR_START_TIMEOUT_SECONDS = 180.0 @@ -313,6 +315,8 @@ def get_node_major_version() -> int | None: check=False, capture_output=True, text=True, + encoding="utf-8", + errors="replace", ) except OSError: return None @@ -513,6 +517,8 @@ def _windows_tasklist_process_name(pid: int) -> str | None: check=False, capture_output=True, text=True, + encoding="utf-8", + errors="replace", ) if completed.returncode != 0: return None @@ -886,6 +892,8 @@ def _process_list_pids() -> list[int]: check=False, capture_output=True, text=True, + encoding="utf-8", + errors="replace", ) else: completed = subprocess.run( @@ -1105,16 +1113,47 @@ def _build_webui_dist(root: Path, config: ServiceConfig, console) -> None: raise ServiceError("WebUI 构建失败。") +def _webui_needs_build(webui_dir: Path) -> bool: + """Return whether WebUI sources are newer than the production bundle.""" + index_path = webui_dir / "dist" / "index.html" + if not index_path.is_file(): + return True + + built_at = index_path.stat().st_mtime_ns + if webui_dir.stat().st_mtime_ns > built_at: + return True + + for current, directories, files in os.walk(webui_dir): + directories[:] = [name for name in directories if name not in WEBUI_BUILD_IGNORED_DIRS] + current_dir = Path(current) + try: + if any((current_dir / name).stat().st_mtime_ns > built_at for name in directories + files): + return True + except FileNotFoundError: + # A concurrent source edit is itself sufficient reason to rebuild. + return True + return False + + def _ensure_webui_dist(root: Path, config: ServiceConfig, console) -> None: """Ensure the FastAPI process can serve the production WebUI bundle.""" from flocks.server.static_webui import WebUIDistMissingError, ensure_webui_dist_dir + webui_dir = root / "webui" try: - ensure_webui_dist_dir() - return + dist_dir = ensure_webui_dist_dir() except WebUIDistMissingError: if config.skip_frontend_build: raise + else: + source_dist_dir = webui_dir / "dist" + if ( + config.skip_frontend_build + or not (webui_dir / "package.json").is_file() + or dist_dir != source_dist_dir.resolve() + or not _webui_needs_build(webui_dir) + ): + return _build_webui_dist(root, config, console) ensure_webui_dist_dir() @@ -1561,6 +1600,21 @@ def restart_all(config: ServiceConfig, console) -> None: _start_all_unlocked(config, console, paths=paths) +def restart_server(console) -> None: + """Restart only the backend through the running supervisor daemon.""" + paths = ensure_runtime_dirs() + with service_lock(paths): + if not supervisor_is_running(paths): + raise ServiceError("Flocks daemon 未运行;请执行 `flocks restart` 进行全量重启。") + try: + status = request_restart_backend(paths=paths) + except Exception as error: + raise ServiceError(f"Flocks server 重启请求失败:{error}") from error + _print_status_payload(status.raw, console, include_daemon_step=False) + if not _startup_payload_is_ready(status.raw): + raise ServiceError(_startup_failure_message(status.raw)) + + def _print_static_port_migration_hint(config: ServiceConfig, console) -> None: """Explain legacy server-port behavior when it differs from public WebUI port.""" if ( @@ -2105,6 +2159,8 @@ def _run_windows_netstat(port: int) -> str: check=False, capture_output=True, text=True, + encoding="utf-8", + errors="replace", ) if completed.returncode != 0: return "" diff --git a/flocks/cli/service_process.py b/flocks/cli/service_process.py index a64d82810..aff0627ba 100644 --- a/flocks/cli/service_process.py +++ b/flocks/cli/service_process.py @@ -54,7 +54,7 @@ def probe(self, process: subprocess.Popen | None, host: str, port: int) -> Servi restart=True, ) if not tcp_port_accepts_connections(host, port): - return ServiceProbeResult(healthy=False, reason=f"port {port} is not listening", restart=True) + return ServiceProbeResult(healthy=False, reason=f"port {port} is not listening") return ServiceProbeResult(healthy=True, reason="liveness check passed") diff --git a/flocks/cli/service_supervisor.py b/flocks/cli/service_supervisor.py index 4a615b59a..937558891 100644 --- a/flocks/cli/service_supervisor.py +++ b/flocks/cli/service_supervisor.py @@ -27,8 +27,8 @@ ) from flocks.cli.service_process import BackendProcessAdapter, ProcessAdapter -SUPERVISOR_CHECK_INTERVAL_SECONDS = 5.0 -SUPERVISOR_HEALTH_FAILURE_THRESHOLD = 2 +SUPERVISOR_CHECK_INTERVAL_SECONDS = 30.0 +SUPERVISOR_HEALTH_FAILURE_THRESHOLD = 10 SUPERVISOR_BACKOFF_SECONDS = (1.0, 2.0, 5.0, 10.0, 30.0) _CLIENT_DISCONNECT_ERRORS = (BrokenPipeError, ConnectionResetError, ConnectionAbortedError) diff --git a/flocks/cli/session_runner.py b/flocks/cli/session_runner.py index ff57c4b1d..994042259 100644 --- a/flocks/cli/session_runner.py +++ b/flocks/cli/session_runner.py @@ -361,7 +361,7 @@ async def _process_message( ) async def _clear_history() -> None: - await Message.clear(self._session.id) + await Message.clear_active(self._session.id) await self._clear_screen() self.console.print("[dim]Conversation history cleared.[/dim]") diff --git a/flocks/config/config.py b/flocks/config/config.py index a5f234f9c..99fa1f3af 100644 --- a/flocks/config/config.py +++ b/flocks/config/config.py @@ -1205,7 +1205,7 @@ async def load_text(cls, text: str, filepath: Path) -> ConfigInfo: # Remove comments properly # 1. Remove /* */ block comments first text_no_comments = re.sub(r'/\*.*?\*/', '', text, flags=re.DOTALL) - + # 2. Remove // line comments, but NOT in strings! # We need to be careful not to remove // inside quoted strings (like URLs) # This regex matches // that are NOT inside quotes @@ -1218,30 +1218,30 @@ async def load_text(cls, text: str, filepath: Path) -> ConfigInfo: in_string = False escape_next = False comment_start = -1 - + for i, char in enumerate(line): if escape_next: escape_next = False continue - + if char == '\\': escape_next = True continue - + if char == '"' and not escape_next: in_string = not in_string - + if not in_string and i < len(line) - 1 and line[i:i+2] == '//': comment_start = i break - + if comment_start >= 0: line = line[:comment_start] - + cleaned_lines.append(line) - + text_no_comments = '\n'.join(cleaned_lines) - + # Parse JSON data = json.loads(text_no_comments) except json.JSONDecodeError as e: diff --git a/flocks/contracts/webui/models.py b/flocks/contracts/webui/models.py index 912bf6a24..a39205321 100644 --- a/flocks/contracts/webui/models.py +++ b/flocks/contracts/webui/models.py @@ -29,6 +29,7 @@ class WebUIWorkspaceManifest(BaseModel): model_config = ConfigDict(populate_by_name=True) id: str = Field(..., description="Stable workspace identifier") + version: str = Field("0.0.0", description="Workspace package version") title: str = Field(..., description="Navigation label") titleEn: Optional[str] = Field(None, description="English navigation label", alias="titleEn") icon: str = Field("LayoutDashboard", description="Lucide icon name") @@ -117,6 +118,7 @@ class WebUIWorkspaceListItem(BaseModel): model_config = ConfigDict(populate_by_name=True, by_alias=True) id: str + version: str = "0.0.0" title: str titleEn: Optional[str] = Field(None, alias="titleEn") route: str diff --git a/flocks/contracts/webui/store.py b/flocks/contracts/webui/store.py index f47a2f8d9..7222ba9a9 100644 --- a/flocks/contracts/webui/store.py +++ b/flocks/contracts/webui/store.py @@ -288,6 +288,7 @@ def list_workspaces(self, *, enabled_only: bool = False) -> list[WebUIWorkspaceL workspaces.append( WebUIWorkspaceListItem( id=manifest.id, + version=manifest.version, title=manifest.title, titleEn=manifest.titleEn, route=webui_contract_workspace_route(manifest.id), @@ -677,6 +678,8 @@ def _iter_page_dirs(self, root: Path) -> list[tuple[Path, str]]: page_dir = manifest_path.parent if page_dir == root: continue + if self._is_hidden_relative(page_dir, root): + continue page_id = self._manifest_page_id_at(manifest_path) if page_id is None: continue @@ -701,6 +704,8 @@ def _iter_workspace_dirs(self, root: Path) -> list[tuple[Path, WebUIWorkspaceMan workspace_dir = manifest_path.parent if workspace_dir == root: continue + if self._is_hidden_relative(workspace_dir, root): + continue manifest = self._read_workspace_manifest_at(workspace_dir) if manifest is None: continue @@ -764,6 +769,23 @@ def _page_dir_in_root(self, root: Path, page_id: str) -> Path: self._assert_inside_root(page_path, root) return page_path + @staticmethod + def _is_hidden_relative(path: Path, root: Path) -> bool: + """True when *path* lives under a dot-prefixed dir relative to *root*. + + The Hub installer stages WebUI packages into sibling scratch dirs + named ``..`` / ``..bak`` inside this very root + before the atomic swap. On Windows that swap can fail (WinError 5) + while the page watcher holds the tree open, leaving those dot-dirs + behind. Scans must skip them so half-written or stale copies never + surface as real pages (and duplicate the live ones). + """ + try: + rel = path.relative_to(root) + except ValueError: + return False + return any(part.startswith(".") for part in rel.parts) + @staticmethod def _assert_inside_root(path: Path, root: Path) -> None: try: diff --git a/flocks/contracts/webui/watcher.py b/flocks/contracts/webui/watcher.py index 9cbb413f0..14eea47fd 100644 --- a/flocks/contracts/webui/watcher.py +++ b/flocks/contracts/webui/watcher.py @@ -143,6 +143,14 @@ def _classify_event( return None if not rel.parts: return None + # Ignore events under dot-prefixed dirs. The Hub installer stages + # WebUI packages into sibling scratch dirs (``..`` / + # ``..bak``) inside this watched root before its atomic + # swap. Reacting to those writes would race the installer's own + # build — on Windows the watcher's file handles block the swap + # with a WinError 5 access-denied — and surface half-built pages. + if any(part.startswith(".") for part in rel.parts): + return None if rel.name == WORKSPACE_MANIFEST_FILE: workspace_id = self._store.workspace_id_for_path(src) diff --git a/flocks/hooks/pipeline.py b/flocks/hooks/pipeline.py index 817406c57..47ecd0d7c 100644 --- a/flocks/hooks/pipeline.py +++ b/flocks/hooks/pipeline.py @@ -2,12 +2,16 @@ Hook Pipeline Provides a lightweight hook registry and execution pipeline that mirrors -oh-my-opencode's lifecycle stages: -- chat.message +agent lifecycle stages: +- user.prompt.submit +- session.start - llm.call.before - llm.call.after - tool.execute.before - tool.execute.after +- turn.finish +- subagent.start +- subagent.stop - event """ @@ -26,11 +30,15 @@ class HookStage: - CHAT_MESSAGE = "chat.message" + USER_PROMPT_SUBMIT = "user.prompt.submit" + SESSION_START = "session.start" LLM_BEFORE = "llm.call.before" LLM_AFTER = "llm.call.after" TOOL_BEFORE = "tool.execute.before" TOOL_AFTER = "tool.execute.after" + TURN_FINISH = "turn.finish" + SUBAGENT_START = "subagent.start" + SUBAGENT_STOP = "subagent.stop" EVENT = "event" CHANNEL_INBOUND = "channel.inbound" CHANNEL_OUTBOUND_BEFORE = "channel.outbound.before" @@ -38,11 +46,15 @@ class HookStage: _DEFAULT_STAGE_TIMEOUTS: Dict[str, float] = { - HookStage.CHAT_MESSAGE: 5.0, + HookStage.USER_PROMPT_SUBMIT: 5.0, + HookStage.SESSION_START: 5.0, HookStage.LLM_BEFORE: 5.0, HookStage.LLM_AFTER: 5.0, HookStage.TOOL_BEFORE: 5.0, HookStage.TOOL_AFTER: 5.0, + HookStage.TURN_FINISH: 5.0, + HookStage.SUBAGENT_START: 5.0, + HookStage.SUBAGENT_STOP: 5.0, HookStage.CHANNEL_INBOUND: 5.0, HookStage.CHANNEL_OUTBOUND_BEFORE: 5.0, HookStage.CHANNEL_OUTBOUND_AFTER: 5.0, @@ -58,7 +70,10 @@ class HookContext: class HookBase: - async def chat_message(self, ctx: HookContext) -> None: # pragma: no cover - default no-op + async def user_prompt_submit(self, ctx: HookContext) -> None: # pragma: no cover - default no-op + return None + + async def session_start(self, ctx: HookContext) -> None: # pragma: no cover - default no-op return None async def llm_before(self, ctx: HookContext) -> None: # pragma: no cover - default no-op @@ -73,6 +88,15 @@ async def tool_before(self, ctx: HookContext) -> None: # pragma: no cover - def async def tool_after(self, ctx: HookContext) -> None: # pragma: no cover - default no-op return None + async def turn_finish(self, ctx: HookContext) -> None: # pragma: no cover - default no-op + return None + + async def subagent_start(self, ctx: HookContext) -> None: # pragma: no cover - default no-op + return None + + async def subagent_stop(self, ctx: HookContext) -> None: # pragma: no cover - default no-op + return None + async def event(self, ctx: HookContext) -> None: # pragma: no cover - default no-op return None @@ -218,12 +242,20 @@ async def ensure_initialized(cls, project_dir: Optional[Path] = None) -> None: cls._loaded_project_dir = str(load_project_dir.resolve(strict=False)) @classmethod - async def run_chat_message( + async def run_user_prompt_submit( cls, input_data: Dict[str, Any], output_data: Optional[Dict[str, Any]] = None, ) -> HookContext: - return await cls._run_stage(HookStage.CHAT_MESSAGE, input_data, output_data) + return await cls._run_stage(HookStage.USER_PROMPT_SUBMIT, input_data, output_data) + + @classmethod + async def run_session_start( + cls, + input_data: Dict[str, Any], + output_data: Optional[Dict[str, Any]] = None, + ) -> HookContext: + return await cls._run_stage(HookStage.SESSION_START, input_data, output_data) @classmethod async def run_llm_before( @@ -257,6 +289,30 @@ async def run_tool_after( ) -> HookContext: return await cls._run_stage(HookStage.TOOL_AFTER, input_data, output_data) + @classmethod + async def run_turn_finish( + cls, + input_data: Dict[str, Any], + output_data: Optional[Dict[str, Any]] = None, + ) -> HookContext: + return await cls._run_stage(HookStage.TURN_FINISH, input_data, output_data) + + @classmethod + async def run_subagent_start( + cls, + input_data: Dict[str, Any], + output_data: Optional[Dict[str, Any]] = None, + ) -> HookContext: + return await cls._run_stage(HookStage.SUBAGENT_START, input_data, output_data) + + @classmethod + async def run_subagent_stop( + cls, + input_data: Dict[str, Any], + output_data: Optional[Dict[str, Any]] = None, + ) -> HookContext: + return await cls._run_stage(HookStage.SUBAGENT_STOP, input_data, output_data) + @classmethod async def run_event( cls, @@ -411,22 +467,29 @@ async def _invoke_handler(handler: Callable[[HookContext], Awaitable[None]], ctx @staticmethod def _resolve_handler(hook: HookBase, stage: str) -> Optional[Callable[[HookContext], Awaitable[None]]]: - if stage == HookStage.CHAT_MESSAGE: - return getattr(hook, "chat_message", None) - if stage == HookStage.LLM_BEFORE: - return getattr(hook, "llm_before", None) - if stage == HookStage.LLM_AFTER: - return getattr(hook, "llm_after", None) - if stage == HookStage.TOOL_BEFORE: - return getattr(hook, "tool_before", None) - if stage == HookStage.TOOL_AFTER: - return getattr(hook, "tool_after", None) - if stage == HookStage.EVENT: - return getattr(hook, "event", None) - if stage == HookStage.CHANNEL_INBOUND: - return getattr(hook, "channel_inbound", None) - if stage == HookStage.CHANNEL_OUTBOUND_BEFORE: - return getattr(hook, "channel_outbound_before", None) - if stage == HookStage.CHANNEL_OUTBOUND_AFTER: - return getattr(hook, "channel_outbound_after", None) - return None + method_name = { + HookStage.USER_PROMPT_SUBMIT: "user_prompt_submit", + HookStage.SESSION_START: "session_start", + HookStage.LLM_BEFORE: "llm_before", + HookStage.LLM_AFTER: "llm_after", + HookStage.TOOL_BEFORE: "tool_before", + HookStage.TOOL_AFTER: "tool_after", + HookStage.TURN_FINISH: "turn_finish", + HookStage.SUBAGENT_START: "subagent_start", + HookStage.SUBAGENT_STOP: "subagent_stop", + HookStage.EVENT: "event", + HookStage.CHANNEL_INBOUND: "channel_inbound", + HookStage.CHANNEL_OUTBOUND_BEFORE: "channel_outbound_before", + HookStage.CHANNEL_OUTBOUND_AFTER: "channel_outbound_after", + }.get(stage) + if method_name is None: + return None + + handler = getattr(hook, method_name, None) + if not callable(handler): + return None + base_handler = getattr(HookBase, method_name, None) + concrete_handler = getattr(type(hook), method_name, None) + if base_handler is not None and concrete_handler is base_handler: + return None + return handler diff --git a/flocks/hub/catalog.py b/flocks/hub/catalog.py index 87a45ce51..a999a12b3 100644 --- a/flocks/hub/catalog.py +++ b/flocks/hub/catalog.py @@ -86,6 +86,8 @@ def _plugin_manifest_signature(plugin_type: PluginType, root: Path) -> tuple[tup candidates = [root / "agent.yaml"] elif plugin_type == "workflow": candidates = [root / "workflow.json", root / "workflow.md"] + elif plugin_type == "webui": + candidates = [root / "workspace.json", root / "manifest.json"] else: try: candidates = [ @@ -741,6 +743,7 @@ def _version_tuple(value: str) -> tuple[int, ...]: def _catalog_install_state( + plugin_type: PluginType, install_path: Optional[Path], record: Optional[local.InstalledPluginRecord], available_version: str, @@ -748,10 +751,15 @@ def _catalog_install_state( if install_path is None: return "available", None if record is None: - # An inferred install has no trustworthy package version because Hub - # manifests are not copied into the install directory. Reconcile it - # through the update path instead of assuming it is already current. - return "updateAvailable", None + installed_version = local.installed_payload_version(plugin_type, install_path) + if installed_version is None: + # Legacy inferred installs have no trustworthy package version. + # Reconcile them through the update path instead of assuming they + # are already current. + return "updateAvailable", None + if _version_tuple(installed_version) < _version_tuple(available_version): + return "updateAvailable", installed_version + return "installed", installed_version if _version_tuple(record.version) < _version_tuple(available_version): return "updateAvailable", record.version return "installed", record.version @@ -776,6 +784,7 @@ def _entry_from_manifest(manifest: HubPluginManifest) -> HubCatalogEntry: ) state, installed_version = _catalog_install_state( + manifest.type, install_path, record, manifest.version, @@ -838,6 +847,7 @@ def _entry_from_index( ) state, installed_version = _catalog_install_state( + item.type, install_path, record, item.version, @@ -920,6 +930,7 @@ def _entry_from_bundled_tool( ) state, installed_version = _catalog_install_state( + manifest.type, install_path, record, manifest.version, diff --git a/flocks/hub/installer.py b/flocks/hub/installer.py index 07b9a30b0..eb894126c 100644 --- a/flocks/hub/installer.py +++ b/flocks/hub/installer.py @@ -4,7 +4,9 @@ import asyncio import shutil +import sys import tempfile +import time from pathlib import Path from typing import Awaitable, Callable @@ -97,17 +99,64 @@ def _remove_path(path: Path) -> None: path.unlink() +def _purge_stale_scratch(parent: Path, name: str) -> None: + """Remove leftover ``..`` / ``..bak`` staging dirs. + + A failed atomic swap (see :func:`_replace_prepared_path`) can leave + scratch and backup dirs behind next to *parent*/*name*. They are never + valid installs, but on Windows a lingering ``..bak`` blocks the + next swap, so we clear both before staging a fresh copy. + """ + if not parent.is_dir(): + return + for entry in parent.iterdir(): + stale = entry.name.startswith(f".{name}.") or entry.name == f".{name}.bak" + if not stale: + continue + try: + if entry.is_dir() and not entry.is_symlink(): + shutil.rmtree(entry, ignore_errors=True) + else: + entry.unlink() + except OSError: + pass + + +def _replace_with_retry(src: Path, dst: Path) -> None: + """``src.replace(dst)`` with a Windows access-denied backoff. + + On Windows an antivirus scan or a directory watcher (e.g. the WebUI + page watcher over ``~/.flocks/plugins/contracts/webui``) can hold a + transient handle on the freshly written tree, making the atomic swap + fail with ``PermissionError`` (WinError 5 / 32). Elsewhere the rename + is atomic and never needs retrying. + """ + if sys.platform != "win32": + src.replace(dst) + return + delay = 0.1 + for attempt in range(6): + try: + src.replace(dst) + return + except PermissionError: + if attempt == 5: + raise + time.sleep(delay) + delay = min(delay * 2, 1.0) + + def _replace_prepared_path(prepared: Path, dst: Path) -> Path | None: backup: Path | None = None if dst.exists() or dst.is_symlink(): backup = dst.parent / f".{dst.name}.bak" _remove_path(backup) - dst.replace(backup) + _replace_with_retry(dst, backup) try: - prepared.replace(dst) + _replace_with_retry(prepared, dst) except Exception: if backup is not None and (backup.exists() or backup.is_symlink()): - backup.replace(dst) + _replace_with_retry(backup, dst) raise return backup @@ -124,12 +173,13 @@ def _commit_replacement(backup: Path | None) -> None: def _rollback_replacement(dst: Path, backup: Path | None) -> None: _remove_path(dst) if backup is not None and (backup.exists() or backup.is_symlink()): - backup.replace(dst) + _replace_with_retry(backup, dst) def _copy_package(src: Path, dst: Path, *, retain_backup: bool = False) -> Path | None: parent = dst.parent parent.mkdir(parents=True, exist_ok=True) + _purge_stale_scratch(parent, dst.name) tmp = Path(tempfile.mkdtemp(prefix=f".{dst.name}.", dir=str(parent))) try: _copy_package_contents(src, tmp) @@ -218,6 +268,7 @@ def _copy_webui_package_with_build( ) -> Path | None: parent = dst.parent parent.mkdir(parents=True, exist_ok=True) + _purge_stale_scratch(parent, dst.name) tmp = Path(tempfile.mkdtemp(prefix=f".{dst.name}.", dir=str(parent))) try: _copy_package_contents(src, tmp) diff --git a/flocks/hub/local.py b/flocks/hub/local.py index e7f5d2580..963d87e9a 100644 --- a/flocks/hub/local.py +++ b/flocks/hub/local.py @@ -122,6 +122,26 @@ def has_install_payload(plugin_type: PluginType, path: Path) -> bool: return path.exists() +def installed_payload_version(plugin_type: PluginType, path: Path) -> Optional[str]: + """Read a version embedded in an installed payload when one is available.""" + if plugin_type != "webui" or not path.is_dir(): + return None + + import json + + workspace_path = path / "workspace.json" + if not workspace_path.is_file(): + return None + try: + raw = json.loads(workspace_path.read_text(encoding="utf-8")) + except Exception: + return None + version = raw.get("version") if isinstance(raw, dict) else None + if not isinstance(version, str) or not version.strip(): + return None + return version.strip() + + def get_record(plugin_type: PluginType, plugin_id: str) -> Optional[InstalledPluginRecord]: return load_installed_records().get(f"{plugin_type}:{plugin_id}") diff --git a/flocks/input/events.py b/flocks/input/events.py index d0fd4e54c..4714c1097 100644 --- a/flocks/input/events.py +++ b/flocks/input/events.py @@ -8,6 +8,7 @@ from flocks.command.command import CommandInfo from flocks.input.types import InputSourceType, surface_for_source +from flocks.session.execution_mode import SessionExecutionMode class UserInputEvent(BaseModel): @@ -31,6 +32,10 @@ class UserInputEvent(BaseModel): mock_reply: Optional[str] = Field(None, alias="mockReply") system: Optional[str] = None tools: Optional[Dict[str, bool]] = None + execution_mode: SessionExecutionMode = Field( + SessionExecutionMode.BUILD, + alias="executionMode", + ) @property def surface(self) -> str: diff --git a/flocks/permission/next.py b/flocks/permission/next.py index 18bfc51af..1bd23a826 100644 --- a/flocks/permission/next.py +++ b/flocks/permission/next.py @@ -7,7 +7,7 @@ import asyncio from datetime import datetime -from typing import Optional, Dict, Any, List, Callable, Awaitable +from typing import Optional, Dict, Any, List, Callable, Awaitable, Iterable from pydantic import BaseModel, Field @@ -59,6 +59,7 @@ class PermissionNext: _session_permissions: Dict[str, Dict[str, str]] = {} _permanent_rules: Dict[str, str] = {} _state_loaded: bool = False + _persistence_tasks: set[asyncio.Task[Any]] = set() _PENDING_PREFIX = "permission_pending:" _REPLY_PREFIX = "permission_reply:" @@ -114,10 +115,19 @@ async def _ensure_persisted_state_loaded(cls) -> None: def _schedule_persist(cls, coro: Awaitable[Any]) -> None: try: loop = asyncio.get_running_loop() - loop.create_task(coro) + task = loop.create_task(coro) + cls._persistence_tasks.add(task) + task.add_done_callback(cls._persistence_tasks.discard) except RuntimeError: pass + @classmethod + async def _flush_persistence(cls) -> None: + """Wait for permission writes already scheduled in this process.""" + + while cls._persistence_tasks: + await asyncio.gather(*tuple(cls._persistence_tasks)) + @classmethod async def _persist_pending_request(cls, request_info: PermissionRequestInfo) -> None: await Storage.set( @@ -184,6 +194,53 @@ async def _persist_session_rules(cls, session_id: str) -> None: "permission_session", ) + @classmethod + async def deletion_storage_keys(cls, session_ids: Iterable[str]) -> List[str]: + """Return every persisted permission key owned by the given sessions.""" + + ids = {session_id for session_id in session_ids if session_id} + if not ids: + return [] + + await cls._flush_persistence() + keys = {f"{cls._SESSION_PREFIX}{session_id}" for session_id in ids} + request_ids = { + request_id + for request_id, pending in cls._pending.items() + if getattr(pending.get("info"), "session_id", None) in ids + } + + for key, value in await Storage.list_entries(prefix=cls._PENDING_PREFIX): + if isinstance(value, dict) and value.get("sessionID") in ids: + keys.add(key) + request_ids.add(key.removeprefix(cls._PENDING_PREFIX)) + + for key, value in await Storage.list_entries(prefix=cls._REPLY_PREFIX): + if isinstance(value, dict) and value.get("sessionID") in ids: + keys.add(key) + request_ids.add(key.removeprefix(cls._REPLY_PREFIX)) + + for request_id in request_ids: + keys.add(f"{cls._PENDING_PREFIX}{request_id}") + keys.add(f"{cls._REPLY_PREFIX}{request_id}") + return sorted(keys) + + @classmethod + def clear_session_runtime(cls, session_ids: Iterable[str]) -> None: + """Drop session permission caches and cancel their pending requests.""" + + ids = {session_id for session_id in session_ids if session_id} + for session_id in ids: + cls._session_permissions.pop(session_id, None) + + for request_id, pending in list(cls._pending.items()): + if getattr(pending.get("info"), "session_id", None) not in ids: + continue + future = pending.get("future") + if isinstance(future, asyncio.Future) and not future.done(): + future.cancel() + cls._pending.pop(request_id, None) + @classmethod async def list_pending_infos(cls) -> List[PermissionRequestInfo]: pending_infos = [ diff --git a/flocks/project/project.py b/flocks/project/project.py index edfa1aa19..8654a8576 100644 --- a/flocks/project/project.py +++ b/flocks/project/project.py @@ -8,6 +8,7 @@ from __future__ import annotations import asyncio +import contextvars import hashlib import json import os @@ -16,13 +17,15 @@ import tempfile import time import uuid -from contextlib import contextmanager +import weakref +from contextlib import asynccontextmanager, contextmanager from datetime import datetime from pathlib import Path -from typing import Any, ClassVar, Dict, Iterator, List, Literal, Optional, Tuple +from typing import AsyncIterator, Any, ClassVar, Dict, Iterator, List, Literal, Optional, Tuple from pydantic import BaseModel, ConfigDict, Field +from flocks.config.config import Config from flocks.utils.log import Log log = Log.create(service="project") @@ -58,6 +61,24 @@ def _platform_file_unlock(fd: int) -> None: fcntl.flock(fd, fcntl.LOCK_UN) +async def _acquire_file_lock(fd: int) -> None: + """Acquire a blocking OS lock without leaking it when the waiter is cancelled.""" + + worker = asyncio.create_task(asyncio.to_thread(_platform_file_lock, fd)) + try: + await asyncio.shield(worker) + except asyncio.CancelledError: + while not worker.done(): + try: + await asyncio.shield(worker) + except asyncio.CancelledError: + continue + if not worker.cancelled(): + worker.result() + _platform_file_unlock(fd) + raise + + @contextmanager def _registry_cross_process_lock(registry_path: Path) -> Iterator[None]: """Serialize registry read-modify-write operations across processes.""" @@ -148,6 +169,7 @@ class ProjectRegistryEntry(BaseModel): updated_at: int = Field(alias="updatedAt") owner_user_id: Optional[str] = Field(None, alias="ownerUserID") shared_local: bool = Field(False, alias="sharedLocal") + removed_at: Optional[int] = Field(None, alias="removedAt") class ProjectRegistry(BaseModel): @@ -163,6 +185,12 @@ class Project: """User-scoped registry of explicitly created projects.""" _lock = asyncio.Lock() + _lifecycle_locks: ClassVar[weakref.WeakValueDictionary[str, asyncio.Lock]] = ( + weakref.WeakValueDictionary() + ) + _held_lifecycle_guards: ClassVar[ + contextvars.ContextVar[Dict[str, asyncio.Task[Any]]] + ] = contextvars.ContextVar("project_lifecycle_guards", default={}) _session_stats_cache: ClassVar[ Dict[Tuple[str, str], Tuple[float, Dict[str, Tuple[int, int, Optional[int]]]]] ] = {} @@ -172,6 +200,89 @@ class Project: def _now_ms() -> int: return int(datetime.now().timestamp() * 1000) + @classmethod + def lifecycle_lock(cls, project_id: str) -> asyncio.Lock: + """Serialize session creation, project removal, and restoration per project.""" + + lock = cls._lifecycle_locks.get(project_id) + if lock is None: + lock = asyncio.Lock() + cls._lifecycle_locks[project_id] = lock + return lock + + @classmethod + @asynccontextmanager + async def lifecycle_guard(cls, project_id: str) -> AsyncIterator[None]: + """Serialize a project lifecycle transition across tasks and processes. + + The context is re-entrant for the current asyncio task so project deletion + can archive several roots through the same public session APIs. ContextVar + state records the owning task explicitly because child tasks inherit + context and must acquire their own lock. + """ + + task = asyncio.current_task() + if task is None: + raise RuntimeError("Project lifecycle guard requires an asyncio task") + + held = cls._held_lifecycle_guards.get() + if held.get(project_id) is task: + yield + return + + async with cls.lifecycle_lock(project_id): + lock_dir = Config.get_data_path() / "locks" / "project-lifecycle" + lock_dir.mkdir(parents=True, exist_ok=True) + lock_name = hashlib.sha256(project_id.encode("utf-8")).hexdigest() + fd = os.open(str(lock_dir / f"{lock_name}.lock"), os.O_RDWR | os.O_CREAT, 0o600) + acquired = False + try: + await _acquire_file_lock(fd) + acquired = True + token = cls._held_lifecycle_guards.set({**held, project_id: task}) + try: + yield + finally: + cls._held_lifecycle_guards.reset(token) + finally: + try: + if acquired: + _platform_file_unlock(fd) + finally: + os.close(fd) + + @classmethod + def registry_state( + cls, + project_id: str, + *, + owner_id: str, + ) -> Literal["virtual", "missing", "active", "removed"]: + """Return the registry state used by session create/restore guards.""" + + if project_id in {DEFAULT_PROJECT_ID, TASK_SESSION_GROUP_ID}: + return "virtual" + entry = next( + ( + item + for item in cls._read_registry(owner_id).projects + if item.id == project_id + ), + None, + ) + if entry is None: + return "missing" + return "removed" if entry.removed_at is not None else "active" + + @classmethod + def is_removed(cls, project_id: str) -> bool: + """Return whether any local registry marks this project as removed.""" + + return any( + entry.id == project_id and entry.removed_at is not None + for entry in cls._all_registry_entries(include_removed=True) + ) + @staticmethod def _flocks_root() -> Path: return Path(os.getenv("FLOCKS_ROOT", str(Path.home() / ".flocks"))).expanduser() @@ -416,6 +527,8 @@ async def create( ) for entry in registry.projects: + if entry.removed_at is not None: + continue if cls._normalized_worktree(entry.worktree) == normalized_worktree: raise ProjectPathConflictError(cls._entry_to_info(entry)) @@ -424,10 +537,35 @@ async def create( raise ValueError("Project name cannot be empty") if normalized_name.casefold() in _RESERVED_PROJECT_NAMES: raise ProjectNameConflictError("Project name is reserved") - if any(entry.name.strip().casefold() == normalized_name.casefold() for entry in registry.projects): + if any( + entry.removed_at is None + and entry.name.strip().casefold() == normalized_name.casefold() + for entry in registry.projects + ): raise ProjectNameConflictError(f"Project name '{normalized_name}' already exists") now = cls._now_ms() + removed_entry = next( + ( + entry + for entry in registry.projects + if entry.removed_at is not None + and cls._normalized_worktree(entry.worktree) == normalized_worktree + ), + None, + ) + if removed_entry is not None: + removed_entry.name = normalized_name + removed_entry.worktree = normalized_worktree + removed_entry.owner_user_id = owner_id + removed_entry.shared_local = False + removed_entry.removed_at = None + removed_entry.updated_at = now + cls._write_registry(owner_id, registry) + cls.invalidate_session_stats(owner_id) + log.info("project.restored", {"id": removed_entry.id}) + return cls._entry_to_info(removed_entry) + project_id = f"prj_{uuid.uuid4()}" entry = ProjectRegistryEntry( id=project_id, @@ -450,12 +588,20 @@ async def list( owner_id: str, ) -> List[ProjectInfo]: registry = await cls.ensure_registry(owner_id) - projects = [cls._entry_to_info(entry) for entry in registry.projects] + projects = [ + cls._entry_to_info(entry) + for entry in registry.projects + if entry.removed_at is None + ] projects.sort(key=lambda item: item.time.updated, reverse=True) return projects @classmethod - def _all_registry_entries(cls) -> List[ProjectRegistryEntry]: + def _all_registry_entries( + cls, + *, + include_removed: bool = False, + ) -> List[ProjectRegistryEntry]: """Return valid entries across local user registries.""" registry_dir = cls._flocks_root() / "projects" @@ -465,7 +611,11 @@ def _all_registry_entries(cls) -> List[ProjectRegistryEntry]: for path in registry_dir.glob("*.json"): registry = cls._read_registry_file(path) if registry is not None: - entries.extend(registry.projects) + entries.extend( + entry + for entry in registry.projects + if include_removed or entry.removed_at is None + ) return entries @classmethod @@ -505,7 +655,12 @@ def is_local_shared(cls, owner_id: Optional[str], project_id: Optional[str]) -> if not owner_id or not project_id or project_id == DEFAULT_PROJECT_ID: return False registry = cls._read_registry(owner_id) - return any(entry.id == project_id and entry.shared_local for entry in registry.projects) + return any( + entry.id == project_id + and entry.removed_at is None + and entry.shared_local + for entry in registry.projects + ) @classmethod async def set_local_shared( @@ -522,7 +677,14 @@ async def set_local_shared( async with cls._lock: with _registry_cross_process_lock(cls.registry_path(owner_id)): registry = cls._read_registry(owner_id) - entry = next((item for item in registry.projects if item.id == project_id), None) + entry = next( + ( + item + for item in registry.projects + if item.id == project_id and item.removed_at is None + ), + None, + ) if entry is None: raise ValueError(f"Project {project_id} not found") entry.owner_user_id = owner_id @@ -544,7 +706,14 @@ async def get( owner_id: str, ) -> Optional[ProjectInfo]: registry = await cls.ensure_registry(owner_id) - entry = next((item for item in registry.projects if item.id == project_id), None) + entry = next( + ( + item + for item in registry.projects + if item.id == project_id and item.removed_at is None + ), + None, + ) return cls._entry_to_info(entry) if entry else None @classmethod @@ -562,11 +731,20 @@ async def update(cls, project_id: str, *, owner_id: str, name: str) -> ProjectIn async with cls._lock: with _registry_cross_process_lock(cls.registry_path(owner_id)): registry = cls._read_registry(owner_id) - entry = next((item for item in registry.projects if item.id == project_id), None) + entry = next( + ( + item + for item in registry.projects + if item.id == project_id and item.removed_at is None + ), + None, + ) if entry is None: raise ValueError(f"Project {project_id} not found") if any( - item.id != project_id and item.name.strip().casefold() == normalized_name.casefold() + item.id != project_id + and item.removed_at is None + and item.name.strip().casefold() == normalized_name.casefold() for item in registry.projects ): raise ProjectNameConflictError(f"Project name '{normalized_name}' already exists") @@ -580,26 +758,114 @@ async def update(cls, project_id: str, *, owner_id: str, name: str) -> ProjectIn @classmethod async def delete(cls, project_id: str, *, owner_id: str) -> bool: - """Remove a project registration without changing sessions or files.""" + """Hide a project registration while preserving metadata for restoration.""" + + async with cls.lifecycle_guard(project_id): + return await cls._delete_locked(project_id, owner_id=owner_id) + + @classmethod + async def _delete_locked(cls, project_id: str, *, owner_id: str) -> bool: + """Update the project registry while its lifecycle guard is held.""" if project_id == DEFAULT_PROJECT_ID: raise ProjectDeletionError("The default project cannot be deleted") async with cls._lock: with _registry_cross_process_lock(cls.registry_path(owner_id)): registry = cls._read_registry(owner_id) - remaining = [entry for entry in registry.projects if entry.id != project_id] - if len(remaining) == len(registry.projects): + entry = next( + ( + item + for item in registry.projects + if item.id == project_id and item.removed_at is None + ), + None, + ) + if entry is None: raise ValueError(f"Project {project_id} not found") - registry.projects = remaining + now = cls._now_ms() + entry.removed_at = now + entry.updated_at = now + entry.shared_local = False cls._write_registry(owner_id, registry) cls.invalidate_session_stats(owner_id) - log.info("project.deleted", {"id": project_id}) + log.info("project.removed", {"id": project_id}) return True + @classmethod + async def restore( + cls, + project_id: str, + *, + owner_id: str, + ) -> Optional[ProjectInfo]: + """Restore a soft-removed project when its original directory is usable.""" + + async with cls.lifecycle_guard(project_id): + return await cls._restore_locked(project_id, owner_id=owner_id) + + @classmethod + async def _restore_locked( + cls, + project_id: str, + *, + owner_id: str, + ) -> Optional[ProjectInfo]: + """Restore registry metadata while its lifecycle guard is held.""" + + if project_id == DEFAULT_PROJECT_ID: + return None + async with cls._lock: + with _registry_cross_process_lock(cls.registry_path(owner_id)): + registry = cls._read_registry(owner_id) + entry = next((item for item in registry.projects if item.id == project_id), None) + if entry is None: + raise ProjectDeletionError( + f"Project {project_id} restoration metadata is unavailable" + ) + if entry.removed_at is None: + return cls._entry_to_info(entry) + if cls._path_status(entry.worktree) != "available": + raise ProjectDeletionError( + f"Project '{entry.name}' directory is unavailable; the archived task was not restored" + ) + + normalized_worktree = cls.validate_worktree( + entry.worktree, + create_if_missing=False, + ) + for other in registry.projects: + if other.id == entry.id or other.removed_at is not None: + continue + if cls._normalized_worktree(other.worktree) == normalized_worktree: + raise ProjectPathConflictError(cls._entry_to_info(other)) + if other.name.strip().casefold() == entry.name.strip().casefold(): + raise ProjectNameConflictError( + f"Project name '{entry.name}' already exists" + ) + + entry.worktree = normalized_worktree + entry.owner_user_id = owner_id + entry.shared_local = False + entry.removed_at = None + entry.updated_at = cls._now_ms() + cls._write_registry(owner_id, registry) + cls.invalidate_session_stats(owner_id) + log.info("project.restored", {"id": project_id}) + return cls._entry_to_info(entry) + @classmethod def registered_project_ids(cls, owner_id: str) -> set[str]: registry = cls._read_registry(owner_id) - return {entry.id for entry in registry.projects} + return {entry.id for entry in registry.projects if entry.removed_at is None} + + @classmethod + def registered_project_names(cls) -> Dict[str, str]: + """Return project display names across local user registries.""" + + return { + entry.id: entry.name + for entry in cls._all_registry_entries(include_removed=True) + } @classmethod def effective_project_id(cls, owner_id: str, stored_project_id: Optional[str]) -> str: diff --git a/flocks/provider/catalog.json b/flocks/provider/catalog.json index d6af63c94..69f35518f 100644 --- a/flocks/provider/catalog.json +++ b/flocks/provider/catalog.json @@ -1229,11 +1229,12 @@ "MOONSHOT_API_KEY" ], "models": { - "kimi-k2.5": { - "name": "Kimi K2.5", - "family": "kimi-k2.5", + "kimi-k3": { + "name": "Kimi K3", + "family": "kimi-k3", "capabilities": { "supports_tools": true, + "supports_vision": true, "supports_reasoning": true, "interleaved": { "field": "reasoning_content", @@ -1244,18 +1245,19 @@ "supports_streaming": true }, "limits": { - "context_window": 262144, - "max_output_tokens": 262144 + "context_window": 1048576, + "max_output_tokens": 131072 }, "pricing": { - "input": 4.0, - "output": 21.0, + "input": 20.0, + "output": 100.0, + "cache_read": 2.0, "currency": "CNY" } }, - "kimi-k2.6": { - "name": "Kimi K2.6", - "family": "kimi-k2.6", + "kimi-k2.7-code": { + "name": "Kimi K2.7 Code", + "family": "kimi-k2.7-code", "capabilities": { "supports_tools": true, "supports_vision": true, @@ -1269,9 +1271,8 @@ "supports_streaming": true }, "limits": { - "context_window": 256000, - "max_input_tokens": 224000, - "max_output_tokens": 16000 + "context_window": 262144, + "max_output_tokens": 32768 }, "pricing": { "input": 6.5, @@ -1280,9 +1281,35 @@ "currency": "CNY" } }, - "kimi-k2-thinking": { - "name": "Kimi K2 Thinking", - "family": "kimi-k2", + "kimi-k2.7-code-highspeed": { + "name": "Kimi K2.7 Code HighSpeed", + "family": "kimi-k2.7-code", + "capabilities": { + "supports_tools": true, + "supports_vision": true, + "supports_reasoning": true, + "interleaved": { + "field": "reasoning_content", + "echo": "tool_calls", + "placeholder": " ", + "cross_provider_policy": "placeholder" + }, + "supports_streaming": true + }, + "limits": { + "context_window": 262144, + "max_output_tokens": 32768 + }, + "pricing": { + "input": 13.0, + "output": 54.0, + "cache_read": 2.6, + "currency": "CNY" + } + }, + "kimi-k2.5": { + "name": "Kimi K2.5", + "family": "kimi-k2.5", "capabilities": { "supports_tools": true, "supports_reasoning": true, @@ -1300,24 +1327,34 @@ }, "pricing": { "input": 4.0, - "output": 16.0, + "output": 21.0, "currency": "CNY" } }, - "kimi-k2": { - "name": "Kimi K2 0711", - "family": "kimi-k2", + "kimi-k2.6": { + "name": "Kimi K2.6", + "family": "kimi-k2.6", "capabilities": { "supports_tools": true, + "supports_vision": true, + "supports_reasoning": true, + "interleaved": { + "field": "reasoning_content", + "echo": "tool_calls", + "placeholder": " ", + "cross_provider_policy": "placeholder" + }, "supports_streaming": true }, "limits": { - "context_window": 131072, - "max_output_tokens": 131072 + "context_window": 256000, + "max_input_tokens": 224000, + "max_output_tokens": 16000 }, "pricing": { - "input": 4.0, - "output": 16.0, + "input": 6.5, + "output": 27.0, + "cache_read": 1.3, "currency": "CNY" } } diff --git a/flocks/provider/interleaved.py b/flocks/provider/interleaved.py index 727628a4e..897833b62 100644 --- a/flocks/provider/interleaved.py +++ b/flocks/provider/interleaved.py @@ -6,6 +6,11 @@ REASONING_TRANSPORT_GENERIC_CHAT = "generic_chat" REASONING_TRANSPORT_ANTHROPIC_MESSAGES = "anthropic_messages" +KIMI_K27_CODE_MODELS = frozenset({ + "kimi-k2.7-code", + "kimi-k2.7-code-highspeed", +}) +KIMI_K3_MODELS = frozenset({"kimi-k3"}) _STRICT_REASONING_CONTENT = { @@ -87,6 +92,18 @@ def _matches_any(text: str, *tokens: str) -> bool: return any(token in text for token in tokens) +def normalize_model_id(model_id: str) -> str: + return model_id.strip().lower().replace("_", "-") + + +def is_kimi_k27_code_model(model_id: str) -> bool: + return normalize_model_id(model_id) in KIMI_K27_CODE_MODELS + + +def is_kimi_k3_model(model_id: str) -> bool: + return normalize_model_id(model_id) in KIMI_K3_MODELS + + def infer_interleaved_capability( *, provider_id: str, @@ -100,7 +117,7 @@ def infer_interleaved_capability( works without user-visible toggles. """ pid = _lower(provider_id) - mid = _lower(model_id) + mid = normalize_model_id(model_id) burl = _lower(base_url) if "minimax" in mid or pid == "minimax": @@ -122,6 +139,8 @@ def infer_interleaved_capability( ): return dict(_PROMOTE_REASONING_CONTENT) + if is_kimi_k27_code_model(mid) or is_kimi_k3_model(mid): + return dict(_STRICT_REASONING_CONTENT) if _matches_any(mid, *_STRICT_REASONING_CONTENT_TOKENS): return dict(_STRICT_REASONING_CONTENT) if ( diff --git a/flocks/provider/options.py b/flocks/provider/options.py index 07cb66def..733093df5 100644 --- a/flocks/provider/options.py +++ b/flocks/provider/options.py @@ -14,6 +14,8 @@ from flocks.provider.interleaved import ( REASONING_TRANSPORT_ANTHROPIC_MESSAGES, REASONING_TRANSPORT_GENERIC_CHAT, + is_kimi_k27_code_model, + is_kimi_k3_model, resolve_interleaved_capability, resolve_reasoning_transport, ) @@ -27,6 +29,8 @@ # --------------------------------------------------------------------------- DEFAULT_THINKING_BUDGET = 16000 DEFAULT_OUTPUT_BUFFER = 8192 +DEFAULT_KIMI_K3_REASONING_EFFORT = "max" +KIMI_K3_REASONING_EFFORTS = frozenset({"low", "high", "max"}) _GENERIC_CHAT_REASONING_EXTRA_BODY_KEYS = { "reasoning_content": "enable_thinking", @@ -71,6 +75,26 @@ def _resolve_reasoning_enabled(provider_id: str, model_id: str) -> Optional[bool return None +def _resolve_reasoning_effort(provider_id: str, model_id: str) -> Optional[str]: + """Read a model-level reasoning effort from flocks.json.""" + try: + from flocks.provider.model_manager import get_model_manager + + setting = get_model_manager().get_setting(provider_id, model_id) + if not setting: + return None + + value = (setting.default_parameters or {}).get("reasoning_effort") + return value.strip().lower() if isinstance(value, str) else None + except Exception as exc: + log.debug("options.reasoning_effort_setting_lookup_failed", { + "provider_id": provider_id, + "model_id": model_id, + "error": str(exc), + }) + return None + + def _resolve_default_extra_body(provider_id: str, model_id: str) -> Optional[Dict[str, Any]]: """Read model-level OpenAI-compatible extra_body from flocks.json.""" try: @@ -201,12 +225,27 @@ def _build_generic_chat_extra_body( model_id: str, interleaved_capability: Optional[Dict[str, Any]], reasoning_enabled: Optional[bool], + reasoning_effort: Optional[str], ) -> Optional[Dict[str, Any]]: """Build OpenAI-compatible reasoning params for the active replay field.""" provider_lower = provider_id.lower() model_lower = model_id.lower() enabled = reasoning_enabled is not False + if is_kimi_k3_model(model_id): + effort = reasoning_effort or DEFAULT_KIMI_K3_REASONING_EFFORT + if effort not in KIMI_K3_REASONING_EFFORTS: + log.warning( + "options.kimi_k3.invalid_reasoning_effort", + { + "model_id": model_id, + "reasoning_effort": effort, + "fallback": DEFAULT_KIMI_K3_REASONING_EFFORT, + }, + ) + effort = DEFAULT_KIMI_K3_REASONING_EFFORT + return {"reasoning_effort": effort} + if "deepseek" in model_lower or provider_lower == "deepseek": return { "thinking": ( @@ -234,6 +273,11 @@ def _build_generic_chat_extra_body( ) } + if is_kimi_k27_code_model(model_id): + # K2.7 cannot disable thinking. Send the enabled form explicitly + # because OpenAI-compatible gateways may not apply Moonshot's default. + return {"thinking": {"type": "enabled"}} + if "kimi" in model_lower: return { "thinking": ( @@ -260,6 +304,7 @@ def build_provider_options( model_id: str, *, reasoning_enabled: Optional[bool] = None, + reasoning_effort: Optional[str] = None, thinking_budget: int = DEFAULT_THINKING_BUDGET, resolve_max_tokens: bool = True, ) -> Dict[str, Any]: @@ -273,6 +318,8 @@ def build_provider_options( The model being called (e.g. ``"claude-sonnet-4"``). thinking_budget: Token budget for extended-thinking / reasoning where applicable. + reasoning_effort: + Kimi K3 reasoning effort (``low``, ``high``, or ``max``). resolve_max_tokens: If *True*, fall back to the model's configured ``max_tokens`` when no provider-specific logic has already set it. @@ -285,11 +332,17 @@ def build_provider_options( model_lower = model_id.lower() interleaved_capability = _resolve_interleaved_capability(provider_id, model_id) reasoning_transport = _resolve_reasoning_transport(provider_id, model_id) + reasoning_effort_explicit = isinstance(reasoning_effort, str) reasoning_enabled = ( _coerce_optional_bool(reasoning_enabled) if reasoning_enabled is not None else _resolve_reasoning_enabled(provider_id, model_id) ) + reasoning_effort = ( + reasoning_effort.strip().lower() + if isinstance(reasoning_effort, str) + else _resolve_reasoning_effort(provider_id, model_id) + ) configured_extra_body = _resolve_default_extra_body(provider_id, model_id) interleaved_enabled = interleaved_capability is not None if interleaved_enabled and reasoning_enabled is None: @@ -348,13 +401,37 @@ def build_provider_options( configured_extra_body or interleaved_enabled or reasoning_enabled is True + or is_kimi_k27_code_model(model_id) + or is_kimi_k3_model(model_id) ): - extra_body = configured_extra_body or _build_generic_chat_extra_body( + automatic_extra_body = _build_generic_chat_extra_body( provider_id, model_id, interleaved_capability, reasoning_enabled, + reasoning_effort, ) + extra_body = dict(configured_extra_body or automatic_extra_body or {}) + if is_kimi_k3_model(model_id): + # K3 is always a reasoning model and uses the top-level + # reasoning_effort field instead of the K2.x thinking object. + extra_body.pop("thinking", None) + if reasoning_effort_explicit: + extra_body["reasoning_effort"] = (automatic_extra_body or {}).get( + "reasoning_effort", + DEFAULT_KIMI_K3_REASONING_EFFORT, + ) + else: + configured_effort = extra_body.get("reasoning_effort") + if configured_effort not in KIMI_K3_REASONING_EFFORTS: + extra_body["reasoning_effort"] = (automatic_extra_body or {}).get( + "reasoning_effort", + DEFAULT_KIMI_K3_REASONING_EFFORT, + ) + elif is_kimi_k27_code_model(model_id): + # K2.7 rejects disabled thinking. Normalize configured values so + # direct and gateway-backed endpoints always request reasoning. + extra_body["thinking"] = {"type": "enabled"} if extra_body: options["extra_body"] = extra_body log.debug("options.thinking_params.resolved", { diff --git a/flocks/provider/provider.py b/flocks/provider/provider.py index 89c815c71..65465ee5c 100644 --- a/flocks/provider/provider.py +++ b/flocks/provider/provider.py @@ -4,11 +4,13 @@ Manages different AI model providers (Anthropic, OpenAI, Google, etc.) """ -from typing import Dict, List, Optional, Any, AsyncIterator, Union -from pydantic import BaseModel, Field, PrivateAttr from enum import Enum +import json import os import threading +from typing import Any, AsyncIterator, Dict, List, Optional, Union + +from pydantic import BaseModel, Field, PrivateAttr from flocks.utils.log import Log from flocks.config.config import Config @@ -52,6 +54,11 @@ def _model_info_signature(model: "ModelInfo") -> tuple: if pricing is not None else None ) + custom_settings = json.dumps( + getattr(model, "custom_settings", None) or {}, + default=str, + sort_keys=True, + ) explicit = tuple(sorted(getattr(model, "_explicit_keys", set()) or set())) return ( getattr(model, "id", None), @@ -59,6 +66,7 @@ def _model_info_signature(model: "ModelInfo") -> tuple: getattr(model, "provider_id", None), cap_sig, pricing_sig, + custom_settings, explicit, ) @@ -124,6 +132,7 @@ class ModelInfo(BaseModel): provider_id: str = Field(..., description="Provider ID") capabilities: ModelCapabilities = Field(default_factory=ModelCapabilities) pricing: Optional[Dict[str, Any]] = Field(None, description="Pricing info") + custom_settings: Dict[str, Any] = Field(default_factory=dict) _explicit_keys: set = PrivateAttr(default_factory=set) """Field names explicitly present in flocks.json (not defaults).""" @@ -737,66 +746,72 @@ async def apply_config(cls, config: Optional[Any] = None, provider_id: Optional[ if not provider: continue + # Provider credentials/options are optional. Model definitions and + # display names below still need to load when credentials are + # unresolved or supplied outside the main config. options = getattr(pconfig, "options", None) - if not options: - continue - + options_data: Optional[Dict[str, Any]] = None if hasattr(options, "model_dump"): - options_data = options.model_dump(exclude_none=True, by_alias=False) + options_data = options.model_dump( + exclude_none=True, + by_alias=False, + ) elif isinstance(options, dict): - options_data = {k: v for k, v in options.items() if v is not None} - else: - continue - - # Handle both Python-style (api_key, base_url) and JS-style (apiKey, baseURL) - api_key = ( - options_data.pop("api_key", None) - or options_data.pop("apiKey", None) - ) - base_url = ( - options_data.pop("base_url", None) - or options_data.pop("baseURL", None) - ) - - # Treat empty strings as None (e.g. unresolved {secret:xxx}) - if isinstance(api_key, str) and not api_key.strip(): - api_key = None - if isinstance(base_url, str) and not base_url.strip(): - base_url = None - - # Also filter out remaining options that resolved to empty strings - options_data = { - k: v for k, v in options_data.items() - if not (isinstance(v, str) and not v.strip()) - } - - if api_key is None and base_url is None and not options_data: - continue + options_data = { + key: value + for key, value in options.items() + if value is not None + } + + if options_data is not None: + # Handle both Python-style (api_key, base_url) and JS-style + # (apiKey, baseURL). + api_key = ( + options_data.pop("api_key", None) + or options_data.pop("apiKey", None) + ) + base_url = ( + options_data.pop("base_url", None) + or options_data.pop("baseURL", None) + ) - # ----- Idempotent ProviderConfig update ------------------------------- - # ``apply_config`` is called from many hot paths: every session - # step (``session.runner._step``), every workflow ``llm.ask``, - # the ``/session/*`` HTTP routes, plus startup. When session and - # workflow run concurrently on different event loops they would - # otherwise rewrite the same ``provider._config`` repeatedly and - # race on the ``_config_models`` rebuild. Skip mutation whenever - # the desired config already matches. - desired_cfg = ProviderConfig( - provider_id=pid, - api_key=api_key, - base_url=base_url, - custom_settings=options_data, - ) - current_cfg = provider._config - current_unchanged = ( - current_cfg is not None - and getattr(current_cfg, "api_key", None) == desired_cfg.api_key - and getattr(current_cfg, "base_url", None) == desired_cfg.base_url - and (getattr(current_cfg, "custom_settings", None) or {}) - == (desired_cfg.custom_settings or {}) - ) - if not current_unchanged: - provider.configure(desired_cfg) + # Treat empty strings as None (e.g. unresolved {secret:xxx}). + if isinstance(api_key, str) and not api_key.strip(): + api_key = None + if isinstance(base_url, str) and not base_url.strip(): + base_url = None + + # Also filter out remaining options that resolved to empty strings. + options_data = { + key: value + for key, value in options_data.items() + if not (isinstance(value, str) and not value.strip()) + } + + if api_key is not None or base_url is not None or options_data: + # ----- Idempotent ProviderConfig update ------------------- + # ``apply_config`` is called from many hot paths: every + # session step, every workflow ``llm.ask``, HTTP routes, + # and startup. Skip mutation whenever the desired config + # already matches. + desired_cfg = ProviderConfig( + provider_id=pid, + api_key=api_key, + base_url=base_url, + custom_settings=options_data, + ) + current_cfg = provider._config + current_unchanged = ( + current_cfg is not None + and getattr(current_cfg, "api_key", None) + == desired_cfg.api_key + and getattr(current_cfg, "base_url", None) + == desired_cfg.base_url + and (getattr(current_cfg, "custom_settings", None) or {}) + == (desired_cfg.custom_settings or {}) + ) + if not current_unchanged: + provider.configure(desired_cfg) # Update provider display name from flocks.json only for providers # that support custom naming (openai-compatible instances and custom-* providers). @@ -855,6 +870,16 @@ async def apply_config(cls, config: Optional[Any] = None, provider_id: Optional[ context_window=model_dict.get("context_window"), ), pricing=_pricing, + custom_settings={ + key: model_dict[key] + for key in ( + "stream_first_chunk_timeout_s", + "streamFirstChunkTimeoutSeconds", + "stream_ongoing_chunk_timeout_s", + "streamOngoingChunkTimeoutSeconds", + ) + if key in model_dict + }, ) model_info._explicit_keys = _explicit_keys desired_models.append(model_info) @@ -1036,7 +1061,10 @@ def __init__(self, provider_id: str, name: str): def configure(self, config: ProviderConfig) -> None: """Configure the provider""" + config_changed = self._config != config self._config = config + if config_changed and hasattr(self, "_client"): + self._client = None self.log.info("provider.configured", { "provider_id": self.id, "has_api_key": config.api_key is not None, diff --git a/flocks/provider/sdk/openai.py b/flocks/provider/sdk/openai.py index ad33a18f3..a360b6e71 100644 --- a/flocks/provider/sdk/openai.py +++ b/flocks/provider/sdk/openai.py @@ -18,6 +18,7 @@ ) from flocks.provider.sdk.openai_base import ( DEFAULT_HTTP_TIMEOUT, + _normalize_stream_usage, build_reasoning_metadata, _coerce_bool, extract_reasoning_content_with_source, @@ -219,15 +220,12 @@ async def chat_stream( # Track usage from final chunk (when stream_options.include_usage is set) stream_usage: Optional[Dict[str, int]] = None + usage_emitted = False async for chunk in stream: # Capture usage from the final chunk (OpenAI returns it in a chunk with no choices) if hasattr(chunk, 'usage') and chunk.usage: - stream_usage = { - "prompt_tokens": getattr(chunk.usage, 'prompt_tokens', 0) or 0, - "completion_tokens": getattr(chunk.usage, 'completion_tokens', 0) or 0, - "total_tokens": getattr(chunk.usage, 'total_tokens', 0) or 0, - } + stream_usage = _normalize_stream_usage(chunk.usage) if not chunk.choices: continue @@ -251,6 +249,7 @@ async def chat_stream( finish_reason=choice.finish_reason, usage=stream_usage, ) + usage_emitted = stream_usage is not None continue # Handle reasoning/thinking content (for o1/o3/gpt-5 models) @@ -315,6 +314,13 @@ async def chat_stream( finish_reason=choice.finish_reason, usage=stream_usage, ) + usage_emitted = stream_usage is not None + + # OpenAI sends the usage-only chunk after the terminal finish chunk. + # Surface it separately when it was not available on the terminal chunk + # so the runner can persist usage and nested reasoning tokens. + if stream_usage and not usage_emitted: + yield StreamChunk(delta="", finish_reason=None, usage=stream_usage) # Embeddings support (added for memory system) async def embed( diff --git a/flocks/provider/sdk/openai_base.py b/flocks/provider/sdk/openai_base.py index e99d36536..b56cb05b2 100644 --- a/flocks/provider/sdk/openai_base.py +++ b/flocks/provider/sdk/openai_base.py @@ -20,6 +20,7 @@ ModelInfo, StreamChunk, ) +from flocks.provider.interleaved import is_kimi_k27_code_model, is_kimi_k3_model from flocks.utils.log import Log log = Log.create(service="provider.openai_base") @@ -301,7 +302,26 @@ def _normalize_stream_usage(raw_usage: Any) -> Optional[Dict[str, int]]: prompt_tokens = (_pt if _pt is not None else getattr(raw_usage, "input_tokens", 0)) or 0 _ct = getattr(raw_usage, "completion_tokens", None) completion_tokens = (_ct if _ct is not None else getattr(raw_usage, "output_tokens", 0)) or 0 - reasoning_tokens = getattr(raw_usage, "reasoning_tokens", 0) or 0 + completion_details = getattr(raw_usage, "completion_tokens_details", None) + output_details = getattr(raw_usage, "output_tokens_details", None) + nested_reasoning_tokens = ( + getattr(completion_details, "reasoning_tokens", 0) + or getattr(output_details, "reasoning_tokens", 0) + or 0 + ) + reasoning_tokens = ( + getattr(raw_usage, "reasoning_tokens", 0) + or nested_reasoning_tokens + or 0 + ) + if nested_reasoning_tokens: + # OpenAI reports reasoning_tokens as a breakdown already included in + # completion_tokens / output_tokens. Flocks stores visible output and + # reasoning separately, so remove that nested subset before persisting. + completion_tokens = max( + 0, + completion_tokens - nested_reasoning_tokens, + ) total_tokens = getattr(raw_usage, "total_tokens", 0) or ( prompt_tokens + completion_tokens + reasoning_tokens ) @@ -390,6 +410,33 @@ def apply_openai_token_limit( params["max_tokens"] = max_tokens +def _is_effective_thinking_enabled( + model_id: str, + thinking: Any, + extra_body: Dict[str, Any], +) -> bool: + """Return the effective reasoning state represented by a request.""" + if is_kimi_k27_code_model(model_id) or is_kimi_k3_model(model_id): + return True + + if isinstance(thinking, dict): + return thinking.get("type") != "disabled" + if thinking: + return True + + extra_thinking = extra_body.get("thinking") + if isinstance(extra_thinking, dict): + return extra_thinking.get("type") != "disabled" + if extra_thinking: + return True + + if extra_body.get("reasoning_effort") is not None: + return True + if extra_body.get("reasoning_split") is True: + return True + return extra_body.get("enable_thinking") is True + + async def create_chat_completion_with_fallbacks( create_call, params: Dict[str, Any], @@ -934,7 +981,9 @@ async def chat( apply_openai_token_limit( params, max_tokens, - prefer_completion_tokens=self.PREFER_MAX_COMPLETION_TOKENS, + prefer_completion_tokens=( + self.PREFER_MAX_COMPLETION_TOKENS or is_kimi_k3_model(model_id) + ), completion_tokens_explicit=max_completion_tokens_explicit, ) if kwargs.get("tools"): @@ -945,7 +994,11 @@ async def chat( # these request logs are emitted for every model call in long sessions. log.info("openai_base.chat.request", { "model": model_id, - "thinking_enabled": bool(thinking), + "thinking_enabled": _is_effective_thinking_enabled( + model_id, + thinking, + extra_body, + ), "has_extra_body": "extra_body" in params, "has_tools": bool(kwargs.get("tools")), "max_tokens": max_tokens, @@ -1021,7 +1074,9 @@ async def chat_stream( apply_openai_token_limit( params, max_tokens, - prefer_completion_tokens=self.PREFER_MAX_COMPLETION_TOKENS, + prefer_completion_tokens=( + self.PREFER_MAX_COMPLETION_TOKENS or is_kimi_k3_model(model_id) + ), completion_tokens_explicit=max_completion_tokens_explicit, ) if kwargs.get("tools"): @@ -1031,7 +1086,11 @@ async def chat_stream( # without repeating the full message history on every turn. log.info("openai_base.stream.request", { "model": model_id, - "thinking_enabled": bool(thinking), + "thinking_enabled": _is_effective_thinking_enabled( + model_id, + thinking, + extra_body, + ), "has_extra_body": "extra_body" in params, "has_tools": bool(kwargs.get("tools")), "max_tokens": max_tokens, @@ -1048,6 +1107,7 @@ async def chat_stream( log_prefix="openai_base", ) tool_calls: Dict[int, Dict[str, Any]] = {} + started_tool_inputs: set[int] = set() emitted_substantive_chunk = False stream_usage: Optional[Dict[str, int]] = None usage_emitted = False @@ -1131,6 +1191,7 @@ async def chat_stream( delta_tcs = getattr(delta, "tool_calls", None) if delta_tcs: emitted_substantive_chunk = True + tool_input_markers: List[Dict[str, Any]] = [] for tc in delta_tcs: idx = tc.index if idx not in tool_calls: @@ -1148,6 +1209,28 @@ async def chat_stream( tool_calls[idx]["function"]["arguments"] += ( tc.function.arguments ) + accumulated_name = tool_calls[idx]["function"]["name"] + if accumulated_name and idx not in started_tool_inputs: + tool_input_markers.append({ + "index": idx, + "id": tool_calls[idx]["id"], + "type": "function", + "function": { + "name": accumulated_name, + "arguments": "", + }, + }) + started_tool_inputs.add(idx) + + # Surface the tool as soon as its name is known, but keep + # partial JSON private. The terminal chunk below publishes + # the complete input once the model finishes generating it. + if tool_input_markers: + yield StreamChunk( + delta="", + finish_reason=None, + tool_calls=tool_input_markers, + ) if choice.finish_reason: # Flush any remaining buffered content from the think-tag extractor @@ -1173,7 +1256,10 @@ async def chat_stream( yield StreamChunk(delta=seg_text, finish_reason=None) if tool_calls: - sorted_calls = [tool_calls[i] for i in sorted(tool_calls.keys())] + sorted_calls = [ + {"index": i, **tool_calls[i]} + for i in sorted(tool_calls.keys()) + ] tool_calls.clear() # Preserve real finish_reason (e.g. "length" when max_tokens # hit) so the runner can detect truncated tool arguments. diff --git a/flocks/provider/sdk/openai_compatible.py b/flocks/provider/sdk/openai_compatible.py index 5f21424f1..812131da3 100644 --- a/flocks/provider/sdk/openai_compatible.py +++ b/flocks/provider/sdk/openai_compatible.py @@ -28,6 +28,7 @@ build_reasoning_metadata, create_chat_completion_with_fallbacks, _coerce_bool, + _is_effective_thinking_enabled, _normalize_stream_usage, extract_reasoning_content_with_source, extract_reasoning_details, @@ -303,7 +304,11 @@ async def chat_stream( self.log.info("openai_compatible.stream.request", { "model": model_id, - "thinking_enabled": bool(thinking), + "thinking_enabled": _is_effective_thinking_enabled( + model_id, + thinking, + extra_body, + ), "has_tools": bool(tools), "max_tokens": max_tokens, "include_usage": True, diff --git a/flocks/security/__init__.py b/flocks/security/__init__.py index 434072168..41f97db7f 100644 --- a/flocks/security/__init__.py +++ b/flocks/security/__init__.py @@ -52,7 +52,7 @@ def _resolve_fofa_derived_secret(secret_id: str, secrets: SecretManager) -> Opti def resolve_secret_value(secret_id: str, secrets: Optional[SecretManager] = None) -> Optional[str]: - """Resolve a secret by id, including provider-specific derived secrets.""" + """Resolve a secret by id, including read-only legacy fallbacks.""" if secrets is None: secrets = get_secret_manager() @@ -60,6 +60,12 @@ def resolve_secret_value(secret_id: str, secrets: Optional[SecretManager] = None if value is not None: return value + if secret_id.endswith("_llm_key"): + provider_id = secret_id.removesuffix("_llm_key") + legacy_value = secrets.get(f"{provider_id}_api_key") + if legacy_value is not None: + return legacy_value + return _resolve_fofa_derived_secret(secret_id, secrets) diff --git a/flocks/server/auth.py b/flocks/server/auth.py index 046fd32f2..47b08e4fc 100644 --- a/flocks/server/auth.py +++ b/flocks/server/auth.py @@ -12,13 +12,17 @@ from fastapi import HTTPException, Request, Response, status from starlette.requests import HTTPConnection -from flocks.auth.context import AuthUser, reset_current_auth_user, set_current_auth_user +from flocks.auth.context import ( + API_TOKEN_SERVICE_USER_ID, + AuthUser, + reset_current_auth_user, + set_current_auth_user, +) from flocks.auth.service import AuthService from flocks.security import get_secret_manager SESSION_COOKIE_NAME = "flocks_session" API_TOKEN_SECRET_ID = "server_api_token" -API_TOKEN_SERVICE_USER_ID = "api-token-service" # Paths that never require auth. Everything else is protected by default. PUBLIC_PATHS = frozenset({ diff --git a/flocks/server/routes/default_model.py b/flocks/server/routes/default_model.py index cda00f077..5aa58b862 100644 --- a/flocks/server/routes/default_model.py +++ b/flocks/server/routes/default_model.py @@ -4,7 +4,7 @@ Provides endpoints to get/set default models per model type. """ -from typing import List, Optional +from typing import List from fastapi import APIRouter, HTTPException, status from pydantic import BaseModel, Field @@ -58,6 +58,7 @@ async def get_all_defaults() -> DefaultModelListResponse: async def get_resolved_default_model(): """Return the resolved default LLM model (provider_id + model_id).""" from flocks.config.config import Config + result = await Config.resolve_default_llm() if not result: raise HTTPException( diff --git a/flocks/server/routes/project.py b/flocks/server/routes/project.py index 9143f1f0d..ab6ec565e 100644 --- a/flocks/server/routes/project.py +++ b/flocks/server/routes/project.py @@ -2,6 +2,7 @@ from __future__ import annotations +import asyncio import os from pathlib import Path from typing import List, Optional @@ -19,7 +20,7 @@ ) from flocks.server.auth import require_user from flocks.session.policy import SessionPolicy -from flocks.session.session import Session +from flocks.session.session import Session, SessionInfo from flocks.utils.log import Log router = APIRouter() @@ -76,6 +77,8 @@ async def _list_project_summaries(user: AuthUser, search: Optional[str]) -> List shared_project_ids=shared_project_ids, ): continue + if session.status != "active": + continue metadata = session.metadata if isinstance(session.metadata, dict) else {} if metadata.get("hideFromSessionManager"): continue @@ -285,31 +288,56 @@ async def unshare_project_local(project_id: str, request: Request): @router.delete("/{project_id}", response_model=bool, summary="Delete project") async def delete_project(project_id: str, request: Request): - """Remove a project and its sessions while preserving project files.""" + """Archive project tasks and hide the project while preserving project files.""" user = require_user(request) try: - if await Project.get(project_id, owner_id=user.id) is None: - raise ValueError(f"Project {project_id} not found") - - project_sessions = [ - session - for session in await Session.list_all_unfiltered() - if session.project_id == project_id - ] - if any(not SessionPolicy.can_delete(session, user) for session in project_sessions): - raise HTTPException( - status_code=403, - detail="Only session owners can delete project sessions", - ) - - from flocks.server.routes.session import delete_session_for_user - - for session in project_sessions: - if await Session.get(project_id, session.id) is not None: - await delete_session_for_user(session.id, user) - - return await Project.delete(project_id, owner_id=user.id) + async with Project.lifecycle_guard(project_id): + if await Project.get(project_id, owner_id=user.id) is None: + raise ValueError(f"Project {project_id} not found") + + project_sessions = [ + session + for session in await Session.list_all_unfiltered() + if session.project_id == project_id + ] + if any(not SessionPolicy.can_delete(session, user) for session in project_sessions): + raise HTTPException( + status_code=403, + detail="Only session owners can delete project sessions", + ) + + root_sessions = [ + session + for session in project_sessions + if session.parent_id is None + ] + if project_sessions and not root_sessions: + raise ProjectDeletionError("Project sessions do not contain a restorable root task") + + from flocks.server.routes.session import archive_session_for_user + + newly_archived: list[SessionInfo] = [] + try: + for session in root_sessions: + await archive_session_for_user(session.id, user) + if session.status != "archived": + newly_archived.append(session) + + return await Project.delete(project_id, owner_id=user.id) + except BaseException: + for session in reversed(newly_archived): + try: + await asyncio.shield( + Session._unarchive_locked(project_id, session.id) + ) + except Exception as rollback_exc: + log.warn("project.delete.rollback_failed", { + "project_id": project_id, + "session_id": session.id, + "error": str(rollback_exc), + }) + raise except ProjectDeletionError as exc: raise HTTPException(status_code=409, detail=str(exc)) from exc except ValueError as exc: diff --git a/flocks/server/routes/provider.py b/flocks/server/routes/provider.py index 9ec50d8fe..c35f13b49 100644 --- a/flocks/server/routes/provider.py +++ b/flocks/server/routes/provider.py @@ -500,6 +500,7 @@ class ProviderInfo(BaseModel): source: str = Field(default="config", description="Provider source: env, config, custom, api") env: List[str] = Field(default_factory=list, description="Environment variable names") key: Optional[str] = Field(None, description="API key (if configured)") + configured: bool = Field(False, description="Whether runtime credentials are configured") options: Dict[str, Any] = Field(default_factory=dict, description="Provider options") # Flocks expects models as Dict[modelID, Model], not List[Model] models: Dict[str, Dict[str, Any]] = Field(default_factory=dict, description="Available models") @@ -596,6 +597,7 @@ async def list_providers() -> ProviderListResponse: source="config", # Provider source: env, config, custom, api env=[], # Environment variable names (can be enhanced later) key=None, # API key not exposed in list + configured=provider.is_configured(), options={}, models=models_dict, ) @@ -1936,6 +1938,10 @@ async def set_provider_credentials( config_dict["name"] = request.provider_name ConfigWriter.add_provider(provider_id, config_dict) + # Config.get() caches resolved secret values. Invalidate it after + # persisting credentials so a later session cannot reapply the old key. + Config.clear_cache() + # 3. Configure the provider runtime so is_configured() reflects the change await _ensure_provider_initialized() provider = Provider.get(provider_id) diff --git a/flocks/server/routes/session.py b/flocks/server/routes/session.py index 2f9f2acad..cbc6c3bdd 100644 --- a/flocks/server/routes/session.py +++ b/flocks/server/routes/session.py @@ -25,8 +25,20 @@ build_context_usage_snapshot, token_usage_to_dict, ) -from flocks.session.session import Session, SessionInfo as SessionModel +from flocks.session.background_tasks import ( + has_pending_session_tasks, + pending_background_tasks, + track_background_task, +) +from flocks.session.session import ( + Session, + SessionInactiveError, + SessionInfo as SessionModel, + SessionNotFoundError, + is_model_auto_session_category, +) from flocks.session.policy import SessionPolicy +from flocks.session.execution_mode import SessionExecutionMode from flocks.utils.log import Log from flocks.utils.json_repair import parse_json_robust, repair_truncated_json from flocks.utils.monitor import get_monitor @@ -38,7 +50,6 @@ # Default agent name constant DEFAULT_AGENT = "rex" DEFAULT_MESSAGE_PAGE_LIMIT = 50 -_DESCENDANT_ABORT_SCAN_LIMIT = 3 _CONTEXT_USAGE_CACHE_TTL_SECONDS = 5.0 _context_usage_cache: Dict[Tuple[str, int], Tuple[float, ContextUsageSnapshot]] = {} _context_usage_inflight: Dict[Tuple[str, int], asyncio.Task[ContextUsageSnapshot]] = {} @@ -63,13 +74,9 @@ def _session_uploads_dir(session_id: str) -> Path: """Return the application-owned upload directory for one session.""" - from flocks.config.config import Config + from flocks.session.files import session_uploads_dir - uploads_root = (Config.get_data_path() / "uploads").resolve() - target = (uploads_root / session_id).resolve() - if target == uploads_root or not target.is_relative_to(uploads_root): - raise ValueError(f"Invalid session ID for upload path: {session_id}") - return target + return session_uploads_dir(session_id) def _materialize_data_url_part( @@ -191,6 +198,10 @@ class SessionCreateRequest(BaseModel): title: Optional[str] = Field(None, description="Session title") permission: Optional[List[PermissionRule]] = Field(None, description="Permission rules") category: Optional[str] = Field(None, description="Session category (e.g. 'user', 'workflow')") + model_auto: bool = Field( + False, + description="Enable WebUI runtime model failover for this session", + ) class FileDiff(BaseModel): @@ -244,9 +255,11 @@ class SessionResponse(BaseModel): permission: Optional[List[Dict[str, Any]]] = Field(None, description="Permission rules") revert: Optional[Dict[str, Any]] = Field(None, description="Revert state") category: str = Field("user", description="Session category: user or task") + status: str = Field("active", description="Session status: active or archived") provider: Optional[str] = Field(None, description="Pinned provider ID") model: Optional[str] = Field(None, description="Pinned model ID") model_pinned: bool = Field(False, description="Whether provider/model are pinned for this session") + model_auto: bool = Field(False, description="Whether WebUI Auto mode is selected") ownerUserID: Optional[str] = Field(None, description="Session owner user id") ownerUsername: Optional[str] = Field(None, description="Session owner username") canWrite: bool = Field(False, description="Whether current user can continue this session") @@ -261,15 +274,20 @@ class SessionListItem(BaseModel): id: str projectID: str + projectName: Optional[str] = None effectiveProjectID: str directory: str title: str time: SessionTime category: str = "user" + status: str = "active" parentID: Optional[str] = None provider: Optional[str] = None model: Optional[str] = None model_pinned: bool = False + model_auto: bool = False + ownerUserID: Optional[str] = None + ownerUsername: Optional[str] = None canWrite: bool = False canDelete: bool = False isShared: bool = False @@ -284,7 +302,7 @@ def _session_to_response( Convert SessionModel to SessionResponse """ current_user = get_current_auth_user() - can_write = SessionPolicy.can_write(session, current_user) + can_write = session.status == "active" and SessionPolicy.can_write(session, current_user) can_delete = SessionPolicy.can_delete(session, current_user) is_shared = SessionPolicy.is_shared(session, shared_project_ids) from flocks.project.project import Project @@ -314,9 +332,11 @@ def _session_to_response( revert=session.revert.model_dump(by_alias=True) if session.revert else None, permission=[p.model_dump() for p in session.permission] if session.permission else None, category=session.category, + status=session.status, provider=session.provider, model=session.model, model_pinned=session.model_pinned, + model_auto=session.model_auto, ownerUserID=session.owner_user_id, ownerUsername=session.owner_username, canWrite=can_write, @@ -329,6 +349,7 @@ def _session_to_list_item( session: SessionModel, effective_project_id: Optional[str] = None, shared_project_ids: Optional[set[str]] = None, + project_name: Optional[str] = None, ) -> SessionListItem: """Convert a session to the lightweight manager-list response shape.""" current_user = get_current_auth_user() @@ -342,6 +363,7 @@ def _session_to_list_item( return SessionListItem( id=session.id, projectID=session.project_id, + projectName=project_name, effectiveProjectID=effective_project_id, directory=session.directory, title=session.title, @@ -352,11 +374,15 @@ def _session_to_list_item( archived=session.time.archived, ), category=session.category, + status=session.status, parentID=session.parent_id, provider=session.provider, model=session.model, model_pinned=session.model_pinned, - canWrite=SessionPolicy.can_write(session, current_user), + model_auto=session.model_auto, + ownerUserID=session.owner_user_id, + ownerUsername=session.owner_username, + canWrite=session.status == "active" and SessionPolicy.can_write(session, current_user), canDelete=SessionPolicy.can_delete(session, current_user), isShared=SessionPolicy.is_shared(session, shared_project_ids), ) @@ -394,6 +420,33 @@ def _require_session_read_access(session: SessionModel, user) -> None: def _require_session_write_access(session: SessionModel, user) -> None: if not SessionPolicy.can_write(session, user): raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail="仅会话所有者可写,受邀用户为只读") + if session.status != "active": + raise HTTPException(status_code=status.HTTP_409_CONFLICT, detail="已归档会话不可修改,请先恢复") + + +async def _persist_active_session_write( + session_id: str, + operation, + *, + expected_generation: Optional[int] = None, +): + """Linearize a durable route write with archive/delete transitions.""" + try: + return await Session.run_active_write( + session_id, + operation, + expected_generation=expected_generation, + ) + except SessionNotFoundError as exc: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Session {session_id} not found", + ) from exc + except SessionInactiveError as exc: + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail="已归档会话不可修改,请先恢复", + ) from exc async def _require_agent_usable_for_chat(agent_name: Optional[str]) -> None: @@ -449,11 +502,7 @@ def _share_metadata(session: SessionModel, *, shared: bool, actor_user_id: str) async def _get_session_by_id_unfiltered(session_id: str) -> Optional[SessionModel]: """Fetch session by id while bypassing policy filtering.""" - token = set_current_auth_user(None) - try: - return await Session.get_by_id(session_id) - finally: - reset_current_auth_user(token) + return await Session.get_by_id_unfiltered(session_id) async def _resolve_session_working_directory(session: SessionModel) -> str: @@ -474,6 +523,7 @@ async def _resolve_session_working_directory(session: SessionModel) -> str: detail="The session directory and current request directory are unavailable", ) fallback_directory = str(fallback_path.resolve()) + log.warn( "session.directory.fallback", { @@ -586,6 +636,11 @@ async def list_sessions( limit: Optional[int] = Query(None, ge=1, description="Maximum sessions to return"), offset: Optional[int] = Query(None, ge=0, description="Number of filtered sessions to skip"), category: Optional[str] = Query(None, description="Filter by category: user or task"), + session_status: Literal["active", "archived", "all"] = Query( + "active", + alias="status", + description="Filter by session status", + ), ) -> List[Union[SessionResponse, SessionListItem]]: """List all sessions with optional filters""" started_at = time.perf_counter() @@ -597,6 +652,11 @@ async def list_sessions( ) list_started_at = time.perf_counter() all_sessions = await Session.list_all_unfiltered() + if session_status == "archived": + all_sessions.sort( + key=lambda item: item.time.archived or item.time.updated, + reverse=True, + ) list_elapsed_ms = (time.perf_counter() - list_started_at) * 1000 visible_project_ids = Project.visible_project_ids(current_user.id) shared_project_ids = Project.shared_project_ids() @@ -605,10 +665,14 @@ async def list_sessions( effective_project_ids: Dict[str, str] = {} term = search.lower() if search else None manager_categories = {"user", "workflow", "entity-config"} + project_names = Project.registered_project_names() if view == "list" else {} skip_remaining = offset or 0 for session in all_sessions: - if not SessionPolicy.can_read( + if session.status == "archived": + if current_user.role != "admin" and not SessionPolicy.is_owner(session, current_user): + continue + elif not SessionPolicy.can_read( session, current_user, shared_project_ids=shared_project_ids, @@ -616,6 +680,8 @@ async def list_sessions( continue if _is_hidden_from_session_manager(session): continue + if session_status != "all" and session.status != session_status: + continue if directory is not None and session.directory != directory: continue effective_project_id = ( @@ -658,7 +724,12 @@ async def list_sessions( if view == "list": response = [ - _session_to_list_item(s, effective_project_ids[s.id], shared_project_ids) + _session_to_list_item( + s, + effective_project_ids[s.id], + shared_project_ids, + project_names.get(s.project_id), + ) for s in filtered ] log_route_timing(log, "session.list.light.complete", started_at=started_at, extra={ @@ -670,6 +741,7 @@ async def list_sessions( "offset": offset, "search": bool(search), "category": category, + "status": session_status, "projectID": projectID, "list_ms": round(list_elapsed_ms, 2), }) @@ -691,6 +763,7 @@ async def list_sessions( "offset": offset, "search": bool(search), "category": category, + "status": session_status, "projectID": projectID, "list_ms": round(list_elapsed_ms, 2), }) @@ -710,17 +783,42 @@ async def create_session(http_request: Request, request: Optional[SessionCreateR await assert_license_active(feature="session_create") if request is None: request = SessionCreateRequest() + if request.model_auto: + category = request.category if request.category is not None else "user" + if not is_model_auto_session_category(category): + raise HTTPException( + status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, + detail=( + "Auto mode is only available for user, entity " + "configuration, and workflow sessions" + ), + ) + if current_user.id == API_TOKEN_SERVICE_USER_ID: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail="Auto mode can only be enabled from the WebUI", + ) + from flocks.session.session_loop import SessionLoop + + auto_available, auto_reason = await SessionLoop.validate_auto_configuration() + if not auto_available: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail=f"Auto mode is unavailable: {auto_reason}", + ) from flocks.project.instance import Instance from flocks.project.project import ( DEFAULT_PROJECT_ID, Project, + ProjectDeletionError, TASK_SESSION_GROUP_ID, ) directory = Instance.get_directory() or os.getcwd() instance_project = Instance.get_project() project_id = instance_project.id if instance_project else DEFAULT_PROJECT_ID + parent_session: Optional[SessionModel] = None # "default" is accepted for compatibility with older WebUI clients, but # ordinary sessions use the current request context rather than a virtual @@ -742,6 +840,26 @@ async def create_session(http_request: Request, request: Optional[SessionCreateR ) project_id = project.id directory = project.worktree + + if request.parentID: + parent_session = await _get_session_by_id_unfiltered(request.parentID) + if parent_session is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Parent session {request.parentID} not found", + ) + _require_session_write_access(parent_session, current_user) + if request.projectID and request.projectID not in { + DEFAULT_PROJECT_ID, + TASK_SESSION_GROUP_ID, + parent_session.project_id, + }: + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail="Parent and child sessions must belong to the same project", + ) + project_id = parent_session.project_id + directory = parent_session.directory # Trigger command:new hook if creating from parent (like /new command) if request.parentID: @@ -785,18 +903,34 @@ async def create_session(http_request: Request, request: Optional[SessionCreateR for p in request.permission ] - is_api_token_client = current_user.id == API_TOKEN_SERVICE_USER_ID - - session = await Session.create( - project_id=project_id, - directory=directory, - title=request.title, - parent_id=request.parentID, - permission=permission, - owner_user_id=None if is_api_token_client else current_user.id, - owner_username=None if is_api_token_client else current_user.username, - **({"category": request.category} if request.category else {}), - ) + if request.projectID and request.projectID not in { + DEFAULT_PROJECT_ID, + TASK_SESSION_GROUP_ID, + }: + project = await Project.get(request.projectID, owner_id=current_user.id) + if project is None: + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail=f"Project {request.projectID} is no longer available", + ) + try: + session = await Session.create( + project_id=project_id, + directory=directory, + title=request.title, + parent_id=request.parentID, + permission=permission, + owner_user_id=(parent_session.owner_user_id if parent_session else current_user.id), + owner_username=(parent_session.owner_username if parent_session else current_user.username), + model_auto=request.model_auto, + model_pinned=False, + **({"category": request.category} if request.category else {}), + ) + except ProjectDeletionError as exc: + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail=str(exc), + ) from exc Project.invalidate_session_stats() log.info("session.created", {"session_id": session.id}) @@ -838,6 +972,11 @@ async def get_session(sessionID: str, request: Request) -> SessionResponse: detail=f"Session {sessionID} not found" ) _require_session_read_access(session, _current_user) + if session.status != "active": + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Session {sessionID} not found", + ) return await _session_to_response_with_goal(session) @@ -941,31 +1080,169 @@ async def update_session_todos(sessionID: str, todos: List[TodoInfo], request: R detail=f"Session {sessionID} not found" ) _require_session_write_access(session, _current_user) + lifecycle_generation = Session.lifecycle_generation(sessionID) try: - await Todo.update( + normalized_todos = [ + SessionTodoInfo(**todo.model_dump(exclude_none=True)) + for todo in todos + ] + await _persist_active_session_write( sessionID, - [SessionTodoInfo(**t.model_dump(exclude_none=True)) for t in todos], + lambda: Todo.update(sessionID, normalized_todos), + expected_generation=lifecycle_generation, ) return todos + except HTTPException: + raise except Exception as e: log.error("session.todo.write_error", {"sessionID": sessionID, "error": str(e)}) raise HTTPException(status_code=500, detail=str(e)) +@router.post( + "/{sessionID}/archive", + response_model=SessionResponse, + summary="Archive session", + description="Archive a session and its descendants without deleting persisted data", +) +async def archive_session(sessionID: str, request: Request) -> SessionResponse: + current_user = require_user(request) + return await archive_session_for_user(sessionID, current_user) + + +@router.post( + "/{sessionID}/restore", + response_model=SessionResponse, + summary="Restore session", + description="Restore an archived session and its descendants", +) +async def restore_session(sessionID: str, request: Request) -> SessionResponse: + current_user = require_user(request) + return await restore_session_for_user(sessionID, current_user) + + +async def _manageable_session_tree(session: SessionModel, current_user: AuthUser) -> List[SessionModel]: + tree = await Session.collect_tree(session.project_id, session.id) + if any(not SessionPolicy.can_delete(item, current_user) for item in tree): + raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail="仅会话所有者可管理归档会话") + return tree + + +async def archive_session_for_user(session_id: str, current_user: AuthUser) -> SessionResponse: + """Stop and archive a complete session tree while retaining persisted data.""" + session = await _get_session_by_id_unfiltered(session_id) + if not session: + raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=f"Session {session_id} not found") + + tree = await _manageable_session_tree(session, current_user) + + auth_token = set_current_auth_user(current_user) + try: + archived_ok = await Session.archive(session.project_id, session.id) + finally: + reset_current_auth_user(auth_token) + if not archived_ok: + raise HTTPException(status_code=status.HTTP_409_CONFLICT, detail="会话归档失败") + + from flocks.server.routes.event import publish_event + + try: + await publish_event("session.updated", {"id": session_id, "status": "archived"}) + except Exception as exc: + log.warn("session.archive.event_error", {"session_id": session_id, "error": str(exc)}) + try: + await emit_audit_event("session_action", { + "action": "archive", + "actor_id": current_user.username, + "actor_name": current_user.username, + "user_name": current_user.username, + "username": current_user.username, + "session_id": session_id, + "owner_user_id": current_user.id, + "project_id": session.project_id, + "affected_sessions": len(tree), + }) + except Exception: + pass + + archived = await _get_session_by_id_unfiltered(session_id) + if archived is None: + raise HTTPException(status_code=status.HTTP_409_CONFLICT, detail="会话归档状态不可用") + return await _session_to_response_with_goal(archived) + + +async def restore_session_for_user(session_id: str, current_user: AuthUser) -> SessionResponse: + """Restore a complete archived session tree.""" + session = await _get_session_by_id_unfiltered(session_id) + if not session: + raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=f"Session {session_id} not found") + + tree = await _manageable_session_tree(session, current_user) + if session.parent_id is not None: + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail="只能从根任务恢复完整任务树", + ) + project_owner_id = session.owner_user_id or current_user.id + auth_token = set_current_auth_user(current_user) + try: + try: + restored_ok = await Session.restore( + session.project_id, + session.id, + project_owner_id=project_owner_id, + ) + except ValueError as exc: + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail=str(exc), + ) from exc + finally: + reset_current_auth_user(auth_token) + if not restored_ok: + raise HTTPException(status_code=status.HTTP_409_CONFLICT, detail="会话恢复失败") + + from flocks.server.routes.event import publish_event + + try: + await publish_event("session.updated", {"id": session_id, "status": "active"}) + except Exception as exc: + log.warn("session.restore.event_error", {"session_id": session_id, "error": str(exc)}) + try: + await emit_audit_event("session_action", { + "action": "restore", + "actor_id": current_user.username, + "actor_name": current_user.username, + "user_name": current_user.username, + "username": current_user.username, + "session_id": session_id, + "owner_user_id": current_user.id, + "project_id": session.project_id, + "affected_sessions": len(tree), + }) + except Exception: + pass + + restored = await _get_session_by_id_unfiltered(session_id) + if restored is None: + raise HTTPException(status_code=status.HTTP_409_CONFLICT, detail="会话恢复状态不可用") + return await _session_to_response_with_goal(restored) + + @router.delete( "/{sessionID}", status_code=status.HTTP_200_OK, - summary="Delete session", - description="Delete session by ID", + summary="Permanently delete session", + description="Permanently delete a session tree and its persisted history", ) async def delete_session(sessionID: str, request: Request) -> bool: - """Delete session by ID (returns true)""" + """Permanently delete a session by ID (returns true).""" current_user = require_user(request) return await delete_session_for_user(sessionID, current_user) async def delete_session_for_user(session_id: str, current_user: AuthUser) -> bool: - """Delete one session using the normal lifecycle cleanup.""" + """Permanently delete a session tree using the normal lifecycle cleanup.""" session = await _get_session_by_id_unfiltered(session_id) if not session: @@ -974,49 +1251,22 @@ async def delete_session_for_user(session_id: str, current_user: AuthUser) -> bo detail=f"Session {session_id} not found" ) - if not SessionPolicy.can_delete(session, current_user): - raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail="仅会话所有者可删除会话") - - from flocks.session.goal import GoalManager - from flocks.session.interaction_queue import InteractionQueue - - await _abort_session_processing(session_id) - await InteractionQueue.clear(session_id) - await GoalManager.clear(session_id) - await _wait_for_sessions_idle([session_id]) - await _abort_and_wait_descendant_sessions(session.project_id, session_id) - await Session.delete(session.project_id, session_id) - from flocks.project.project import Project + tree = await _manageable_session_tree(session, current_user) + tree_ids = [item.id for item in tree] - Project.invalidate_session_stats() - - # Best-effort cleanup of any image/file uploads materialised for this - # session via ``_materialize_data_url_part`` (see prompt_async). - # The session DB row is gone, so the on-disk bytes are now orphaned — - # remove them to keep application data tidy. We deliberately swallow any - # filesystem errors: deletion of the session record is the contract, - # the upload cleanup is incidental. + auth_token = set_current_auth_user(current_user) try: - import shutil - uploads_root = _session_uploads_dir(session_id) - if uploads_root.exists() and uploads_root.is_dir(): - shutil.rmtree(uploads_root, ignore_errors=True) - log.info("session.uploads.cleaned", { - "session_id": session_id, - "path": str(uploads_root), - }) - except Exception as exc: - log.warn("session.uploads.cleanup_failed", { - "session_id": session_id, - "error": str(exc), - }) - + deleted_ok = await Session.delete(session.project_id, session_id) + finally: + reset_current_auth_user(auth_token) + if not deleted_ok: + raise HTTPException(status_code=status.HTTP_409_CONFLICT, detail="会话永久删除失败") log.info("session.deleted", {"session_id": session_id}) try: await emit_audit_event( "session_action", { - "action": "delete", + "action": "permanent_delete", "actor_id": current_user.username, "actor_name": current_user.username, "user_name": current_user.username, @@ -1024,6 +1274,7 @@ async def delete_session_for_user(session_id: str, current_user: AuthUser) -> bo "session_id": session_id, "owner_user_id": current_user.id, "project_id": session.project_id, + "affected_sessions": len(tree_ids), }, ) except Exception: @@ -1031,56 +1282,6 @@ async def delete_session_for_user(session_id: str, current_user: AuthUser) -> bo return True -async def _collect_descendant_session_ids(project_id: str, session_id: str) -> List[str]: - """Return child session IDs in the same order Session.delete will recurse.""" - sessions = await Session.list(project_id) - children_by_parent: Dict[str, List[str]] = {} - for child in sessions: - if child.parent_id is None: - continue - children_by_parent.setdefault(child.parent_id, []).append(child.id) - - descendants: List[str] = [] - seen: set[str] = set() - - def visit(parent_id: str) -> None: - for child_id in children_by_parent.get(parent_id, []): - if child_id in seen: - continue - seen.add(child_id) - descendants.append(child_id) - visit(child_id) - - visit(session_id) - return descendants - - -async def _abort_and_wait_descendant_sessions(project_id: str, session_id: str) -> None: - """Abort descendants that exist after the parent stops and wait as a batch.""" - known: set[str] = set() - latest: List[str] = [] - - for _ in range(_DESCENDANT_ABORT_SCAN_LIMIT): - latest = await _collect_descendant_session_ids(project_id, session_id) - new_ids = [sid for sid in latest if sid not in known] - for descendant_id in new_ids: - await _abort_session_processing(descendant_id) - known.update(new_ids) - - if latest: - await _wait_for_sessions_idle(latest) - - refreshed = await _collect_descendant_session_ids(project_id, session_id) - if set(refreshed).issubset(known): - return - latest = refreshed - - log.warn("session.delete.descendants_unstable", { - "session_id": session_id, - "descendants": latest, - }) - - class SessionUpdateRequest(BaseModel): """Request to update session""" model_config = ConfigDict(populate_by_name=True) @@ -1090,6 +1291,15 @@ class SessionUpdateRequest(BaseModel): provider: Optional[str] = Field(None, description="Pinned provider ID") model: Optional[str] = Field(None, description="Pinned model ID") model_pinned: Optional[bool] = Field(None, description="Whether provider/model are pinned for this session") + model_auto: Optional[bool] = Field(None, description="Whether WebUI Auto mode is selected") + + +class SessionMoveProjectRequest(BaseModel): + """Request to move a complete task tree to a project.""" + + model_config = ConfigDict(populate_by_name=True) + + project_id: str = Field(..., alias="projectID", description="Target project ID") @router.patch( @@ -1114,6 +1324,43 @@ async def update_session( current_user = require_user(http_request) _require_session_write_access(existing, current_user) + if request.model_auto is True and ( + request.provider is not None + or request.model is not None + or request.model_pinned is True + ): + raise HTTPException( + status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, + detail="Auto mode cannot be combined with a pinned provider/model", + ) + if request.model_auto is True: + if not is_model_auto_session_category(existing.category): + raise HTTPException( + status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, + detail=( + "Auto mode is only available for user, entity " + "configuration, and workflow sessions" + ), + ) + if current_user.id == API_TOKEN_SERVICE_USER_ID: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail="Auto mode can only be enabled from the WebUI", + ) + from flocks.session.session_loop import SessionLoop + + auto_available, auto_reason = await SessionLoop.validate_auto_configuration() + if not auto_available: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail=f"Auto mode is unavailable: {auto_reason}", + ) + if (request.provider is None) != (request.model is None): + raise HTTPException( + status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, + detail="provider and model must be updated together", + ) + updates = {} if request.title is not None: updates["title"] = request.title @@ -1125,6 +1372,15 @@ async def update_session( updates["model"] = request.model if request.model_pinned is not None: updates["model_pinned"] = request.model_pinned + if request.model_pinned: + updates["model_auto"] = False + if request.model_auto is not None: + updates["model_auto"] = request.model_auto + if request.model_auto: + updates["model_pinned"] = False + if request.provider is not None and request.model is not None: + updates["model_auto"] = False + updates["model_pinned"] = True session = await Session.update( project_id=existing.project_id, @@ -1137,11 +1393,119 @@ async def update_session( status_code=status.HTTP_404_NOT_FOUND, detail=f"Session {sessionID} not found" ) + + if ( + request.model_auto is False + or request.model_pinned is True + or (request.provider is not None and request.model is not None) + ): + from flocks.session.session_loop import SessionLoop + + SessionLoop.clear_auto_failover_state(sessionID) log.info("session.updated", {"session_id": sessionID}) return await _session_to_response_with_goal(session) +@router.patch( + "/{sessionID}/project", + response_model=SessionResponse, + summary="Move session to project", + description="Move a root session and all descendants to another project", +) +async def move_session_to_project( + sessionID: str, + request: SessionMoveProjectRequest, + http_request: Request, +) -> SessionResponse: + """Move an idle root task tree into a writable project.""" + + existing = await _get_session_by_id_unfiltered(sessionID) + if existing is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Session {sessionID} not found", + ) + + current_user = require_user(http_request) + _require_session_write_access(existing, current_user) + if existing.parent_id is not None: + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail="只能移动根任务及其完整任务树", + ) + + from flocks.project.project import Project + + target_project = await Project.get(request.project_id, owner_id=current_user.id) + if target_project is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Project {request.project_id} not found", + ) + if target_project.path_status != "available": + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail="目标项目目录不可用", + ) + if existing.project_id == target_project.id: + return await _session_to_response_with_goal(existing) + + tree = await Session.collect_tree(existing.project_id, existing.id) + if any(not SessionPolicy.can_write(item, current_user) for item in tree): + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail="仅任务所有者可移动完整任务树", + ) + + if any(_is_session_busy_for_move(item.id) for item in tree): + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail="任务正在运行,请等待任务完成或停止后再移动", + ) + if any(item.revert is not None for item in tree): + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail="任务正在历史恢复状态,请先完成或取消恢复后再移动", + ) + + auth_token = set_current_auth_user(current_user) + try: + moved = await Session.move_to_project( + existing.project_id, + existing.id, + target_project_id=target_project.id, + target_directory=target_project.worktree, + target_owner_id=current_user.id, + additional_busy_check=_is_session_busy_for_move, + ) + finally: + reset_current_auth_user(auth_token) + if moved is None: + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail="任务移动失败,请刷新后重试", + ) + + from flocks.server.routes.event import publish_event + + try: + await publish_event("session.updated", { + "id": sessionID, + "projectID": target_project.id, + }) + except Exception as exc: + log.warn("session.move_project.event_error", { + "session_id": sessionID, + "error": str(exc), + }) + + return await _session_to_response_with_goal( + moved, + effective_project_id=target_project.id, + ) + + @router.post( "/{sessionID}/share-local", response_model=SessionResponse, @@ -1457,8 +1821,11 @@ async def _run_in_background(): }, }) - import asyncio - asyncio.create_task(_run_in_background()) + _schedule_background_coro( + _run_in_background(), + session_id=sessionID, + action="session.summarize", + ) log.info("session.summarized", {"session_id": sessionID}) return True @@ -1489,11 +1856,17 @@ async def revert_session(sessionID: str, request: RevertRequest, http_request: R ) _require_session_write_access(session, current_user) - updated = await SessionRevert.revert( - session_id=sessionID, - message_id=request.messageID, - part_id=request.partID, - ) + try: + updated = await SessionRevert.revert( + session_id=sessionID, + message_id=request.messageID, + part_id=request.partID, + ) + except (ValueError, SessionInactiveError) as exc: + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail=str(exc), + ) from exc log.info("session.reverted", {"session_id": sessionID, "message_id": request.messageID}) return await _session_to_response_with_goal(updated) @@ -1584,6 +1957,11 @@ class PromptRequest(BaseModel): tools: Optional[Dict[str, bool]] = Field(None, description="Tool settings (deprecated)") system: Optional[str] = Field(None, description="System prompt override") variant: Optional[str] = Field(None, description="Model variant") + execution_mode: SessionExecutionMode = Field( + SessionExecutionMode.BUILD, + alias="executionMode", + description="Execution mode for this user turn", + ) class UserMessageInfo(BaseModel): @@ -1612,6 +1990,7 @@ class UserMessageInfo(BaseModel): tools: Optional[Dict[str, bool]] = None variant: Optional[str] = None compacted: Optional[bool] = None + executionMode: SessionExecutionMode = SessionExecutionMode.BUILD class AssistantMessageInfo(BaseModel): @@ -1768,6 +2147,11 @@ async def _message_to_response_info(msg: Any, *, cwd: str) -> MessageInfo: agent=getattr(msg, "agent", None) or DEFAULT_AGENT, model=model_info, compacted=getattr(msg, "compacted", None), + executionMode=getattr( + msg, + "executionMode", + SessionExecutionMode.BUILD, + ), ) tokens_raw = getattr(msg, "tokens", None) @@ -1914,9 +2298,10 @@ async def get_session_messages( return result except Exception as e: log.error("session.messages.error", {"error": str(e), "sessionID": sessionID}) - if page or before is not None: - return MessagePage(sessionID=sessionID, items=[], hasMore=False, nextBefore=None) - return [] + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail="Failed to load session messages", + ) from e @router.get( @@ -2087,12 +2472,14 @@ async def _publish_text_part_update( }) -def _track_background_task(task: "asyncio.Task[Any]") -> None: +def _track_background_task( + task: "asyncio.Task[Any]", + *, + session_id: Optional[str] = None, +) -> None: """Keep background tasks alive until completion.""" - if not hasattr(router, "_pending_tasks"): - router._pending_tasks = set() - router._pending_tasks.add(task) - task.add_done_callback(lambda t: router._pending_tasks.discard(t)) + track_background_task(task, session_id=session_id) + router._pending_tasks = pending_background_tasks() def _schedule_background_coro( @@ -2135,13 +2522,13 @@ async def _guarded_coro() -> None: }) task = asyncio.get_running_loop().create_task(_guarded_coro()) - _track_background_task(task) + _track_background_task(task, session_id=session_id) async def _prepare_replay_runtime( session_id: str, user_message, -) -> Dict[str, str]: +) -> Dict[str, Any]: """Resolve replay runtime state before mutating session history.""" from flocks.agent.registry import Agent from flocks.config.config import Config @@ -2149,21 +2536,40 @@ async def _prepare_replay_runtime( agent_name = getattr(user_message, "agent", None) or await Agent.default_agent() agent = await Agent.get(agent_name) or await Agent.get(DEFAULT_AGENT) - # Replay should follow the model that is active *now* for this session - # (current session pin / current default / current agent override), not the - # historical model stored on the original user message being replayed. - dummy_request = type( - "_MessageReplayRequest", - (), - {"model": None, "agent": agent_name}, - )() - provider_id, model_id, _ = await _resolve_model(dummy_request, agent, session_id) + session = await Session.get_by_id(session_id) + auto_failover = bool( + session + and is_model_auto_session_category(getattr(session, "category", "user")) + and getattr(session, "model_auto", False) + ) + if auto_failover: + default_llm = await Config.resolve_default_llm() + if not default_llm: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail="Auto mode requires a configured default LLM", + ) + provider_id = default_llm["provider_id"] + model_id = default_llm["model_id"] + else: + # Replay follows the model active *now*, not the model stored on the + # historical user message. + dummy_request = type( + "_MessageReplayRequest", + (), + {"model": None, "agent": agent_name}, + )() + provider_id, model_id, _ = await _resolve_model( + dummy_request, + agent, + session_id, + ) Provider._ensure_initialized() config = await Config.get() await Provider.apply_config(config, provider_id=provider_id) provider = Provider.get(provider_id) - if not provider: + if not provider and not auto_failover: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail=f"Provider {provider_id} not found", @@ -2173,6 +2579,7 @@ async def _prepare_replay_runtime( "agent_name": agent_name, "provider_id": provider_id, "model_id": model_id, + "auto_failover": auto_failover, } @@ -2181,7 +2588,7 @@ async def _run_existing_user_message( session, user_message, working_directory: str, - runtime: Optional[Dict[str, str]] = None, + runtime: Optional[Dict[str, Any]] = None, ): """Run SessionLoop using an already-persisted user message.""" from flocks.server.routes.event import publish_event @@ -2214,6 +2621,7 @@ async def _on_error(error: str): agent_name=agent_name, callbacks=loop_callbacks, working_directory=working_directory, + auto_failover=bool(runtime.get("auto_failover", False)), ) if result.action == "queued": @@ -2233,6 +2641,8 @@ async def _on_error(error: str): assistant_message_id = None created_ms = end_ms final_tokens = {"input": 0, "output": 0, "reasoning": 0, "cache": {"read": 0, "write": 0}} + actual_provider_id = result.provider_id or provider_id + actual_model_id = result.model_id or model_id if result.last_message: assistant_message_id = result.last_message.id @@ -2241,6 +2651,14 @@ async def _on_error(error: str): finish = getattr(result.last_message, "finish", None) if finish: finish_reason = finish + actual_provider_id = ( + getattr(result.last_message, "providerID", None) + or actual_provider_id + ) + actual_model_id = ( + getattr(result.last_message, "modelID", None) + or actual_model_id + ) result_time = getattr(result.last_message, "time", None) if isinstance(result_time, dict): created_ms = result_time.get("created", created_ms) @@ -2260,8 +2678,8 @@ async def _on_error(error: str): "role": "assistant", "time": {"created": created_ms, "completed": end_ms}, "parentID": user_message.id, - "modelID": model_id, - "providerID": provider_id, + "modelID": actual_model_id, + "providerID": actual_provider_id, "mode": agent_name, "agent": agent_name, "path": {"cwd": working_directory, "root": working_directory}, @@ -2274,8 +2692,8 @@ async def _on_error(error: str): publish_event, session_id, session=session, - provider_id=provider_id, - model_id=model_id, + provider_id=actual_provider_id, + model_id=actual_model_id, ) log.info("session.message.replay.completed", { @@ -2343,6 +2761,14 @@ async def resend_session_message( detail="Session is currently generating a response", ) + try: + await SessionRevert.ensure_replayable(session, messageID) + except ValueError as exc: + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail=str(exc), + ) from exc + working_directory = await _resolve_session_working_directory(session) async def _handle_resend() -> None: @@ -2438,6 +2864,14 @@ async def regenerate_session_message( detail="Session is currently generating a response", ) + try: + await SessionRevert.ensure_replayable(session, parent_message_id) + except ValueError as exc: + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail=str(exc), + ) from exc + working_directory = await _resolve_session_working_directory(session) async def _handle_regenerate() -> None: @@ -2502,8 +2936,28 @@ async def send_session_message(sessionID: str, request: PromptRequest, http_requ ) current_user = require_user(http_request) _require_session_write_access(session, current_user) + lifecycle_generation = Session.lifecycle_generation(sessionID) working_directory = await _resolve_session_working_directory(session) + _validate_execution_mode_request(request) + if request.execution_mode == SessionExecutionMode.GOAL: + from flocks.session.goal import GoalManager + + objective = _extract_text_from_parts(request.parts).strip() + state = await GoalManager.set_goal(sessionID, objective) + await publish_event("session.goal.updated", { + "sessionID": sessionID, + "status": state.status, + "objective": state.objective, + "reason": state.last_reason, + }) + request = request.model_copy(update={ + "parts": _replace_text_parts( + request.parts, + GoalManager.goal_prompt(state.objective), + ), + "display_text": request.display_text or objective, + }) log.info("session.message.send.processing", { "sessionID": sessionID, @@ -2517,7 +2971,13 @@ async def send_session_message(sessionID: str, request: PromptRequest, http_requ result = await Instance.provide( directory=working_directory, init=instance_bootstrap, - fn=lambda: _process_session_message(sessionID, session, request, working_directory) + fn=lambda: _process_session_message( + sessionID, + session, + request, + working_directory, + lifecycle_generation=lifecycle_generation, + ) ) log.info("session.message.send.complete", {"sessionID": sessionID}) return result @@ -2818,6 +3278,8 @@ async def _process_session_message( session, request: PromptRequest, working_directory: str, + *, + lifecycle_generation: Optional[int] = None, ): """ Process session message within Instance context. @@ -2837,6 +3299,9 @@ async def _process_session_message( from flocks.session.runner import RunnerCallbacks import time import os + + if lifecycle_generation is None: + lifecycle_generation = Session.lifecycle_generation(sessionID) # Clean up revert state before processing (Flocks compatibility) await SessionRevert.cleanup(session) @@ -2873,9 +3338,27 @@ async def _process_session_message( agent_name = request.agent or await Agent.default_agent() agent = await Agent.get(agent_name) or await Agent.get(DEFAULT_AGENT) - provider_id, model_id, model_source = await _resolve_model( - request, agent, sessionID + auto_failover = bool( + is_model_auto_session_category(getattr(session, "category", "user")) + and getattr(session, "model_auto", False) + and not request.model ) + if auto_failover: + from flocks.config.config import Config + + default_llm = await Config.resolve_default_llm() + if not default_llm: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail="Auto mode requires a configured default LLM", + ) + provider_id = default_llm["provider_id"] + model_id = default_llm["model_id"] + model_source = "auto_primary" + else: + provider_id, model_id, model_source = await _resolve_model( + request, agent, sessionID + ) log.info("session.message.model", { "provider_id": provider_id, @@ -2895,15 +3378,17 @@ async def _process_session_message( session.provider = provider_id session.model = model_id session.model_pinned = True + session.model_auto = False + SessionLoop.clear_auto_failover_state(sessionID) # Ensure providers are initialized and configured Provider._ensure_initialized() from flocks.config.config import Config config = await Config.get() await Provider.apply_config(config, provider_id=provider_id) - + provider = Provider.get(provider_id) - if not provider: + if not provider and not auto_failover: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail=f"Provider {provider_id} not found" @@ -2925,17 +3410,22 @@ async def _process_session_message( display_metadata = {"displayText": display_text} if display_text else None _is_no_reply = bool(request.noReply) - user_message = await Message.create( - session_id=sessionID, - role=MessageRole.USER, - content=text_content, - id=user_message_id, - time={"created": now_ms}, - agent=agent_name, - model={"providerID": provider_id, "modelID": model_id}, - part_id=user_part_id, - part_metadata=display_metadata, - synthetic=True if _is_no_reply else None, + user_message = await _persist_active_session_write( + sessionID, + lambda: Message.create( + session_id=sessionID, + role=MessageRole.USER, + content=text_content, + id=user_message_id, + time={"created": now_ms}, + agent=agent_name, + model={"providerID": provider_id, "modelID": model_id}, + executionMode=request.execution_mode, + part_id=user_part_id, + part_metadata=display_metadata, + synthetic=True if _is_no_reply else None, + ), + expected_generation=lifecycle_generation, ) user_message_id = user_message.id @@ -2947,6 +3437,7 @@ async def _process_session_message( "time": {"created": now_ms}, "agent": agent_name, "model": {"providerID": provider_id, "modelID": model_id}, + "executionMode": request.execution_mode.value, } }) _part_event: dict = { @@ -3008,7 +3499,11 @@ async def _process_session_message( filename=raw_part.get("filename"), url=url, ) - await Message.add_part(sessionID, user_message_id, file_part) + await _persist_active_session_write( + sessionID, + lambda: Message.add_part(sessionID, user_message_id, file_part), + expected_generation=lifecycle_generation, + ) await publish_event("message.part.updated", { "part": { "id": file_part_id, @@ -3033,15 +3528,19 @@ async def _process_session_message( mock_msg_id = Identifier.ascending("message") mock_part_id = Identifier.ascending("part") mock_now = int(time.time() * 1000) - await Message.create( - session_id=sessionID, - role=MessageRole.ASSISTANT, - content=request.mockReply, - id=mock_msg_id, - time={"created": mock_now, "completed": mock_now}, - parentID=user_message_id, - modelID="mock", - part_id=mock_part_id, + await _persist_active_session_write( + sessionID, + lambda: Message.create( + session_id=sessionID, + role=MessageRole.ASSISTANT, + content=request.mockReply, + id=mock_msg_id, + time={"created": mock_now, "completed": mock_now}, + parentID=user_message_id, + modelID="mock", + part_id=mock_part_id, + ), + expected_generation=lifecycle_generation, ) await publish_event("message.updated", { "info": { @@ -3102,6 +3601,7 @@ async def _on_error(error: str): agent_name=agent_name, callbacks=loop_callbacks, working_directory=working_directory, + auto_failover=auto_failover, ) # ------------------------------------------------------------------ @@ -3131,6 +3631,8 @@ async def _on_error(error: str): final_content = "" assistant_message_id = None final_tokens = {"input": 0, "output": 0, "reasoning": 0, "cache": {"read": 0, "write": 0}} + actual_provider_id = result.provider_id or provider_id + actual_model_id = result.model_id or model_id if result.last_message: assistant_message_id = result.last_message.id @@ -3139,6 +3641,14 @@ async def _on_error(error: str): finish = getattr(result.last_message, 'finish', None) if finish: finish_reason = finish + actual_provider_id = ( + getattr(result.last_message, "providerID", None) + or actual_provider_id + ) + actual_model_id = ( + getattr(result.last_message, "modelID", None) + or actual_model_id + ) if result.action == "error": finish_reason = "error" @@ -3156,8 +3666,8 @@ async def _on_error(error: str): "role": "assistant", "time": {"created": now_ms, "completed": end_ms}, "parentID": user_message_id, - "modelID": model_id, - "providerID": provider_id, + "modelID": actual_model_id, + "providerID": actual_provider_id, "mode": agent_name, "agent": agent_name, "path": {"cwd": working_directory, "root": working_directory}, @@ -3170,8 +3680,8 @@ async def _on_error(error: str): publish_event, sessionID, session=session, - provider_id=provider_id, - model_id=model_id, + provider_id=actual_provider_id, + model_id=actual_model_id, ) # Collect parts for the response @@ -3199,8 +3709,8 @@ async def _on_error(error: str): title_task = loop.create_task( SessionTitle.generate_title_after_first_message( session_id=sessionID, - model_id=model_id, - provider_id=provider_id, + model_id=actual_model_id, + provider_id=actual_provider_id, event_publish_callback=publish_event, ) ) @@ -3218,8 +3728,8 @@ async def _on_error(error: str): "role": "assistant", "time": {"created": now_ms, "completed": end_ms}, "parentID": user_message_id, - "modelID": model_id, - "providerID": provider_id, + "modelID": actual_model_id, + "providerID": actual_provider_id, "mode": agent_name, "agent": agent_name, "path": {"cwd": working_directory, "root": working_directory}, @@ -3353,6 +3863,32 @@ def _extract_text_from_parts(parts: List[Dict[str, Any]]) -> str: return "".join(part.get("text", "") for part in parts if part.get("type") == "text") +def _validate_execution_mode_request(request: PromptRequest) -> None: + if request.execution_mode != SessionExecutionMode.GOAL: + return + objective = _extract_text_from_parts(request.parts).strip() + if not objective: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail="Goal mode requires a non-empty text objective", + ) + if any(part.get("type") != "text" for part in request.parts): + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail="Goal mode does not support attachments", + ) + + +def _event_text_for_execution_mode( + parts: List[Dict[str, Any]], + execution_mode: SessionExecutionMode, +) -> str: + text = _extract_text_from_parts(parts) + if execution_mode == SessionExecutionMode.GOAL: + return f"/goal {text.strip()}" + return text + + def _replace_text_parts( parts: Optional[List[Dict[str, Any]]], text: str, @@ -3417,6 +3953,19 @@ def _is_prompt_chain_active(session_id: str) -> bool: return session_id in getattr(router, "_prompt_queue_active_sessions", set()) +def _is_session_busy_for_move(session_id: str) -> bool: + """Cover loop execution and route-owned background work during a move.""" + + from flocks.session.session_loop import SessionLoop + + return ( + Session.has_active_operations(session_id) + or SessionLoop.is_running(session_id) + or _is_prompt_chain_active(session_id) + or has_pending_session_tasks(session_id) + ) + + def _set_prompt_chain_active(session_id: str, active: bool) -> None: if not hasattr(router, "_prompt_queue_active_sessions"): router._prompt_queue_active_sessions = set() @@ -3458,17 +4007,21 @@ def _event_from_queued_prompt(item, working_directory: str): return UserInputEvent( source_type="webui", sessionID=item.sessionID, - text=_extract_text_from_parts(item.parts), + text=_event_text_for_execution_mode(item.parts, item.executionMode), parts=[dict(part) for part in item.parts], agent=item.agent, model=item.model, variant=item.variant, + metadata={ + "_sessionLifecycleGeneration": Session.lifecycle_generation(item.sessionID), + }, display_text=item.display_text, messageID=item.messageID, noReply=item.noReply, mockReply=item.mockReply, tools=item.tools, system=item.system, + executionMode=item.executionMode, working_directory=working_directory, ) @@ -3605,19 +4158,23 @@ def _build_prompt_request_from_event(event, prompt_text: str, display_text: Opti noReply=event.no_reply, tools=event.tools, system=event.system, + execution_mode=event.execution_mode, ) async def _dispatch_sse_input(sessionID: str, session, event, working_directory: str) -> None: import time as _time - from flocks.input.dispatcher import dispatch_user_input + from flocks.input.dispatcher import dispatch_user_input, parse_slash_command from flocks.input.output import SSEOutputSink from flocks.server.routes.event import publish_event from flocks.session.message import Message, MessageRole from flocks.utils.id import Identifier agent_name = event.agent or "rex" + lifecycle_generation = event.metadata.get("_sessionLifecycleGeneration") + if not isinstance(lifecycle_generation, int): + lifecycle_generation = Session.lifecycle_generation(sessionID) async def _create_user_message( user_text: str, @@ -3629,15 +4186,20 @@ async def _create_user_message( user_msg_id = event.message_id or Identifier.create("message") user_part_id = Identifier.create("part") message_agent = agent_override or agent_name - await Message.create( - session_id=sessionID, - role=MessageRole.USER, - content=user_text, - id=user_msg_id, - time={"created": now_ms}, - agent=message_agent, - **({"model": model_info} if model_info else {}), - part_id=user_part_id, + await _persist_active_session_write( + sessionID, + lambda: Message.create( + session_id=sessionID, + role=MessageRole.USER, + content=user_text, + id=user_msg_id, + time={"created": now_ms}, + agent=message_agent, + executionMode=event.execution_mode, + **({"model": model_info} if model_info else {}), + part_id=user_part_id, + ), + expected_generation=lifecycle_generation, ) await publish_event("message.updated", { "info": { @@ -3646,6 +4208,7 @@ async def _create_user_message( "role": "user", "time": {"created": now_ms}, "agent": message_agent, + "executionMode": event.execution_mode.value, **({"model": model_info} if model_info else {}), } }) @@ -3667,18 +4230,22 @@ async def _publish_direct_response(output_event, text: str) -> None: asst_now = int(_time.time() * 1000) asst_msg_id = Identifier.ascending("message") asst_part_id = Identifier.ascending("part") - await Message.create( - session_id=sessionID, - role=MessageRole.ASSISTANT, - content=text, - id=asst_msg_id, - time={"created": asst_now, "completed": asst_now}, - parentID=parent_msg_id, - modelID="command", - providerID="builtin", - agent=agent_name, - finish="stop", - part_id=asst_part_id, + await _persist_active_session_write( + sessionID, + lambda: Message.create( + session_id=sessionID, + role=MessageRole.ASSISTANT, + content=text, + id=asst_msg_id, + time={"created": asst_now, "completed": asst_now}, + parentID=parent_msg_id, + modelID="command", + providerID="builtin", + agent=agent_name, + finish="stop", + part_id=asst_part_id, + ), + expected_generation=lifecycle_generation, ) await publish_event("message.updated", { "info": { @@ -3719,8 +4286,26 @@ async def _publish_direct_response(output_event, text: str) -> None: ) async def _run_llm(output_event, prompt_text: str, display_text: Optional[str] = None) -> None: + parsed = parse_slash_command(output_event.text, output_event.metadata) + if parsed is not None and parsed.canonical_name == "goal": + from flocks.session.goal import GoalManager + + state = await GoalManager.get(sessionID) + if state is not None: + await publish_event("session.goal.updated", { + "sessionID": sessionID, + "status": state.status, + "objective": state.objective, + "reason": state.last_reason, + }) request = _build_prompt_request_from_event(output_event, prompt_text, display_text) - await _process_session_message(sessionID, session, request, working_directory) + await _process_session_message( + sessionID, + session, + request, + working_directory, + lifecycle_generation=lifecycle_generation, + ) async def _clear_history() -> None: await _clear_session_history(sessionID) @@ -3763,18 +4348,7 @@ async def _run_session_control(output_event, parsed) -> bool: session_control=_run_session_control, clear_history=_clear_history, ) - result = await dispatch_user_input(event, sink) - if result.command_name == "goal" and result.action == "llm": - from flocks.session.goal import GoalManager - - state = await GoalManager.get(sessionID) - if state is not None: - await publish_event("session.goal.updated", { - "sessionID": sessionID, - "status": state.status, - "objective": state.objective, - "reason": state.last_reason, - }) + await dispatch_user_input(event, sink) class PromptQueueUpdateRequest(BaseModel): @@ -3784,24 +4358,34 @@ class PromptQueueUpdateRequest(BaseModel): async def _enqueue_prompt_request( session_id: str, request: PromptRequest, + *, + expected_generation: Optional[int] = None, ): from flocks.session.interaction_queue import InteractionQueue + _validate_execution_mode_request(request) + if expected_generation is None: + expected_generation = Session.lifecycle_generation(session_id) await _require_agent_usable_for_chat(request.agent) model = request.model.model_dump(by_alias=True) if request.model else None parts = _materialize_queued_parts(session_id, [dict(part) for part in request.parts]) - return await InteractionQueue.enqueue( + return await _persist_active_session_write( session_id, - parts=parts, - agent=request.agent, - model=model, - variant=request.variant, - display_text=request.display_text, - message_id=request.messageID, - no_reply=request.noReply, - mock_reply=request.mockReply, - tools=request.tools, - system=request.system, + lambda: InteractionQueue.enqueue( + session_id, + parts=parts, + agent=request.agent, + model=model, + variant=request.variant, + display_text=request.display_text, + message_id=request.messageID, + no_reply=request.noReply, + mock_reply=request.mockReply, + tools=request.tools, + system=request.system, + execution_mode=request.execution_mode, + ), + expected_generation=expected_generation, ) @@ -3948,8 +4532,10 @@ async def send_session_message_async( if http_request is not None: current_user = require_user(http_request) _require_session_write_access(session, current_user) + lifecycle_generation = Session.lifecycle_generation(sessionID) working_directory = await _resolve_session_working_directory(session) + _validate_execution_mode_request(request) await _require_agent_usable_for_chat(request.agent) log.info("session.prompt_async.accepted", { @@ -3957,27 +4543,44 @@ async def send_session_message_async( "directory": working_directory, }) + event_text = _event_text_for_execution_mode( + request.parts, + request.execution_mode, + ) + event_display_text = request.display_text + if request.execution_mode == SessionExecutionMode.GOAL: + event_display_text = ( + event_display_text + or _extract_text_from_parts(request.parts).strip() + ) + event = UserInputEvent( source_type="webui", sessionID=sessionID, - text=_extract_text_from_parts(request.parts), + text=event_text, parts=[dict(part) for part in request.parts], agent=request.agent, model=request.model.model_dump(by_alias=True) if request.model else None, variant=request.variant, - display_text=request.display_text, + metadata={"_sessionLifecycleGeneration": lifecycle_generation}, + display_text=event_display_text, messageID=request.messageID, noReply=request.noReply, mockReply=request.mockReply, tools=request.tools, system=request.system, + executionMode=request.execution_mode, working_directory=working_directory, ) existing_queue = await InteractionQueue.list(sessionID) if SessionLoop.is_running(sessionID) or existing_queue or _is_prompt_chain_active(sessionID): try: - item = await _enqueue_prompt_request(sessionID, request) + item = await _enqueue_prompt_request( + sessionID, + request, + expected_generation=lifecycle_generation, + ) except QueueFullError as exc: raise HTTPException(status_code=status.HTTP_409_CONFLICT, detail=str(exc)) from exc await _publish_prompt_queue(sessionID) @@ -4045,6 +4648,7 @@ async def send_session_command(sessionID: str, request: CommandRequest, http_req if http_request is not None: current_user = require_user(http_request) _require_session_write_access(session, current_user) + lifecycle_generation = Session.lifecycle_generation(sessionID) working_directory = await _resolve_session_working_directory(session) await _require_agent_usable_for_chat(request.agent) @@ -4052,6 +4656,7 @@ async def send_session_command(sessionID: str, request: CommandRequest, http_req if not raw_arguments and request.arguments_json is not None: raw_arguments = json.dumps(request.arguments_json, ensure_ascii=False) command_metadata: Dict[str, Any] = {} + command_metadata["_sessionLifecycleGeneration"] = lifecycle_generation if request.arguments_json is not None: command_metadata["commandArgumentsJson"] = request.arguments_json @@ -4107,10 +4712,7 @@ async def _handle_command() -> None: loop = asyncio.get_running_loop() task = loop.create_task(_handle_command()) - if not hasattr(router, "_pending_tasks"): - router._pending_tasks = set() - router._pending_tasks.add(task) - task.add_done_callback(lambda t: router._pending_tasks.discard(t)) + _track_background_task(task, session_id=sessionID) log.info("session.command.accepted", { "sessionID": sessionID, @@ -4149,12 +4751,24 @@ async def run_shell_command(sessionID: str, request: ShellRequest, http_request: if request.model: model = {"providerID": request.model.providerID, "modelID": request.model.modelID} - result = await SessionRunner.shell( - session_id=sessionID, - agent=request.agent, - command=request.command, - model=model, - ) + try: + async with Session.active_operation(sessionID): + result = await SessionRunner.shell( + session_id=sessionID, + agent=request.agent, + command=request.command, + model=model, + ) + except SessionNotFoundError as exc: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=str(exc), + ) from exc + except SessionInactiveError as exc: + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail=str(exc), + ) from exc log.info("session.shell.executed", { "sessionID": sessionID, @@ -4369,7 +4983,11 @@ async def get_session_statistics(sessionID: str): raise HTTPException(status_code=500, detail=f"Failed to get session statistics: {str(e)}") -async def _clear_session_history(sessionID: str) -> int: +async def _clear_session_history( + sessionID: str, + *, + expected_generation: Optional[int] = None, +) -> int: """Clear stored messages for a session and notify subscribed UIs.""" session_info = await _get_session_by_id_unfiltered(sessionID) if not session_info: @@ -4385,14 +5003,23 @@ async def _clear_session_history(sessionID: str) -> int: await abort_session(sessionID) await InteractionQueue.clear(sessionID) - await GoalManager.clear(sessionID) try: await _publish_prompt_queue(sessionID) except Exception as exc: log.warn("session.clear.prompt_queue_event_error", {"sessionID": sessionID, "error": str(exc)}) await _wait_for_session_idle(sessionID) - deleted_count = await Message.clear(sessionID) + async def clear_persisted_history() -> int: + await GoalManager.clear(sessionID) + deleted_count = await Message.clear(sessionID) + await Session._clear_project_move_metadata_locked(sessionID) + return deleted_count + + deleted_count = await _persist_active_session_write( + sessionID, + clear_persisted_history, + expected_generation=expected_generation, + ) log.info("session.cleared", {"sessionID": sessionID, "deleted": deleted_count}) try: @@ -4423,7 +5050,10 @@ async def clear_session(sessionID: str, http_request: Request): current_user = require_user(http_request) _require_session_write_access(session_info, current_user) - deleted_count = await _clear_session_history(sessionID) + deleted_count = await _clear_session_history( + sessionID, + expected_generation=Session.lifecycle_generation(sessionID), + ) return { "status": "success", "sessionID": sessionID, diff --git a/flocks/session/background_tasks.py b/flocks/session/background_tasks.py new file mode 100644 index 000000000..3608580a2 --- /dev/null +++ b/flocks/session/background_tasks.py @@ -0,0 +1,60 @@ +"""Track cancellable asynchronous work associated with sessions.""" + +import asyncio +from typing import Any, Optional + + +_pending_tasks: set[asyncio.Task[Any]] = set() +_tasks_by_session: dict[str, set[asyncio.Task[Any]]] = {} + + +def track_background_task( + task: asyncio.Task[Any], + *, + session_id: Optional[str] = None, +) -> None: + """Keep a task alive and optionally associate it with a session.""" + + _pending_tasks.add(task) + if session_id: + _tasks_by_session.setdefault(session_id, set()).add(task) + + def discard(completed: asyncio.Task[Any]) -> None: + _pending_tasks.discard(completed) + if not session_id: + return + session_tasks = _tasks_by_session.get(session_id) + if session_tasks is None: + return + session_tasks.discard(completed) + if not session_tasks: + _tasks_by_session.pop(session_id, None) + + task.add_done_callback(discard) + + +async def cancel_session_background_tasks(session_id: str) -> None: + """Cancel and join tracked work before a lifecycle transition commits.""" + + current_task = asyncio.current_task() + tasks = [ + task + for task in _tasks_by_session.get(session_id, ()) + if task is not current_task and not task.done() + ] + for task in tasks: + task.cancel() + if tasks: + await asyncio.gather(*tasks, return_exceptions=True) + + +def pending_background_tasks() -> set[asyncio.Task[Any]]: + """Return the live task set for compatibility with route-level drains.""" + + return _pending_tasks + + +def has_pending_session_tasks(session_id: str) -> bool: + """Return whether tracked asynchronous work is still active for a session.""" + + return any(not task.done() for task in _tasks_by_session.get(session_id, ())) diff --git a/flocks/session/callable_state.py b/flocks/session/callable_state.py index d4e0097e9..e3dbd2a0f 100644 --- a/flocks/session/callable_state.py +++ b/flocks/session/callable_state.py @@ -73,5 +73,10 @@ async def clear_session_callable_tools(session_id: str) -> None: await Storage.delete(f"{_CALLABLE_PREFIX}{session_id}") +def invalidate_session_callable_tools_cache(session_id: str) -> None: + """Drop process-local callable-tool state after an external transaction.""" + _cache.pop(session_id, None) + + async def session_can_call_tool(session_id: str, tool_name: str) -> bool: return tool_name in await get_session_callable_tools(session_id) diff --git a/flocks/session/execution_mode.py b/flocks/session/execution_mode.py new file mode 100644 index 000000000..fbad85d6f --- /dev/null +++ b/flocks/session/execution_mode.py @@ -0,0 +1,196 @@ +"""Session execution-mode policy derived from OpenCode and Codex.""" + +from __future__ import annotations + +from enum import Enum +from typing import Any, Iterable, Optional + +from flocks.session.plan_file import ( + SessionPlanFile, + is_current_plan_path, + plan_edit_patterns_allowed, + plan_file_prompt, +) + + +class SessionExecutionMode(str, Enum): + """Execution mode selected for a user turn.""" + + BUILD = "build" + PLAN = "plan" + GOAL = "goal" + + +PLAN_ONLY_TOOL_NAMES = frozenset({"plan_exit"}) +PLAN_DENIED_TOOL_NAMES = frozenset( + { + # Explicit slash commands keep their existing direct user-only path. + "run_slash_command", + } +) +PLAN_DELEGATION_TOOL_NAMES = frozenset({"delegate_task", "task"}) +PLAN_DELEGATABLE_AGENT_NAMES = frozenset({"explore", "librarian"}) +PLAN_PATH_SCOPED_TOOL_NAMES = frozenset({"apply_patch", "edit", "write"}) + +PLAN_MODE_PROMPT = """# Plan Mode + +You are in a planning turn. You may inspect files, configuration, types, +tests, and documentation. Bash is available only for read-only exploration +and validation. Do not use shell commands to modify files, configuration, +services, dependencies, version control state, or any other system state. + +The only file you may modify is the session plan file named below. The runtime +enforces this boundary for file-editing tools. + +Follow this workflow: + +1. Explore first. Ground the plan in the existing environment and resolve + discoverable facts through inspection before asking the user. + Delegation is limited to the `explore` and `librarian` subagents. +2. Use the question tool only for material ambiguities, preferences, or + trade-offs that cannot be resolved from the environment. After the user + answers, continue exploring and planning as needed. +3. Review the proposed approach for remaining gaps. Ask another focused + question if a decision is still required. +4. Write the decision-complete implementation plan to the session plan file. + The plan must be detailed enough for another engineer to execute without + making additional design decisions. +5. Present the final plan to the user, then immediately call plan_exit. That + tool asks the user whether to start implementation. If approved, it switches + the next turn to Build and starts implementing the approved plan. If + declined, remain in Plan and use the feedback to refine it. + +Do not ask for implementation approval with ordinary prose or the question +tool; plan_exit owns that transition. A Plan turn may end only by asking a +material clarification question or by calling plan_exit after the final plan. +""" + + +def coerce_execution_mode(value: object) -> SessionExecutionMode: + """Return a valid execution mode, defaulting legacy values to Build.""" + + if isinstance(value, SessionExecutionMode): + return value + try: + return SessionExecutionMode(str(value or SessionExecutionMode.BUILD.value)) + except ValueError: + return SessionExecutionMode.BUILD + + +def runtime_execution_mode(value: object) -> SessionExecutionMode: + """Resolve the permission mode used while executing a turn.""" + + mode = coerce_execution_mode(value) + if mode == SessionExecutionMode.GOAL: + return SessionExecutionMode.BUILD + return mode + + +def is_tool_allowed(value: object, tool_name: str) -> bool: + """Evaluate tool visibility against OpenCode-style Plan permissions.""" + + mode = runtime_execution_mode(value) + if tool_name in PLAN_ONLY_TOOL_NAMES: + return mode == SessionExecutionMode.PLAN + if mode == SessionExecutionMode.BUILD: + return True + return tool_name not in PLAN_DENIED_TOOL_NAMES + + +def tool_call_denial_reason( + value: object, + tool_name: str, + arguments: dict[str, Any], + ctx: Any, +) -> Optional[str]: + """Return a hard Plan-mode denial reason for a concrete tool call.""" + + if runtime_execution_mode(value) != SessionExecutionMode.PLAN: + return None + if tool_name in PLAN_DELEGATION_TOOL_NAMES: + subagent_type = str(arguments.get("subagent_type") or "").strip().lower() + if ( + subagent_type in PLAN_DELEGATABLE_AGENT_NAMES + and not arguments.get("category") + and not arguments.get("session_id") + ): + return None + allowed = ", ".join(sorted(PLAN_DELEGATABLE_AGENT_NAMES)) + return ( + f"Tool {tool_name!r} may only delegate to {allowed} via " + "subagent_type while Plan mode is active." + ) + if tool_name not in PLAN_PATH_SCOPED_TOOL_NAMES: + return None + + if tool_name in {"edit", "write"}: + paths = [arguments.get("filePath")] + else: + try: + from flocks.tool.file.apply_patch import parse_patch + + hunks = parse_patch(str(arguments.get("patchText") or "")) + except Exception: + hunks = [] + paths = [ + path + for hunk in hunks + for path in (getattr(hunk, "path", None), getattr(hunk, "move_path", None)) + if path + ] + + if paths and all(is_current_plan_path(ctx, path) for path in paths): + return None + return ( + f"Tool {tool_name!r} may only edit the current session plan file " + "while Plan mode is active." + ) + + +def is_permission_allowed( + value: object, + permission: str, + patterns: Iterable[object], + ctx: Any, +) -> bool: + """Apply the same Plan file boundary at the permission entry point.""" + + if runtime_execution_mode(value) != SessionExecutionMode.PLAN: + return True + if permission != "edit": + return True + return plan_edit_patterns_allowed(ctx, patterns) + + +def is_plan_file_edit(value: object, ctx: Any, path: object) -> bool: + """Return whether a read-only sandbox may allow this Plan artifact edit.""" + + return ( + runtime_execution_mode(value) == SessionExecutionMode.PLAN + and is_current_plan_path(ctx, path) + ) + + +def filter_tool_names(value: object, tool_names: Iterable[str]) -> list[str]: + """Return only tool names allowed by the selected execution mode.""" + + return [name for name in tool_names if is_tool_allowed(value, name)] + + +def execution_mode_prompt( + value: object, + *, + session: Any = None, + plan_file: Optional[SessionPlanFile] = None, +) -> str: + """Return the per-turn developer guidance for a mode.""" + + mode = runtime_execution_mode(value) + if mode == SessionExecutionMode.PLAN: + file_prompt = ( + plan_file_prompt(session, plan=plan_file) + if session is not None + else "" + ) + return f"{PLAN_MODE_PROMPT.rstrip()}\n\n{file_prompt}".strip() + return "" diff --git a/flocks/session/features/todo.py b/flocks/session/features/todo.py index f941ac331..a34812ec6 100644 --- a/flocks/session/features/todo.py +++ b/flocks/session/features/todo.py @@ -91,23 +91,37 @@ async def update(cls, session_id: str, todos: List[TodoInfo]) -> None: for todo in todos ] - # Store in storage await Storage.set( f"todo:{session_id}", [todo.model_dump(exclude_none=True) for todo in validated_todos], - "todo" + "todo", ) - - # Publish event await Bus.publish(cls.Updated, { "sessionID": session_id, "todos": validated_todos, }) - log.info("todo.updated", { "session_id": session_id, "count": len(validated_todos), }) + + @classmethod + async def update_active( + cls, + session_id: str, + todos: List[TodoInfo], + *, + expected_generation: Optional[int] = None, + ) -> None: + """Update todos only while the owning session remains active.""" + + from flocks.session.session import Session + + await Session.run_active_write( + session_id, + lambda: cls.update(session_id, todos), + expected_generation=expected_generation, + ) @classmethod async def get(cls, session_id: str) -> List[TodoInfo]: diff --git a/flocks/session/files.py b/flocks/session/files.py new file mode 100644 index 000000000..3abf90642 --- /dev/null +++ b/flocks/session/files.py @@ -0,0 +1,26 @@ +"""Application-owned files associated with persisted sessions.""" + +from pathlib import Path +import shutil + +from flocks.config.config import Config + + +def session_uploads_dir(session_id: str) -> Path: + """Return the upload directory for one session, constrained to app data.""" + + uploads_root = (Config.get_data_path() / "uploads").resolve() + target = (uploads_root / session_id).resolve() + if target == uploads_root or not target.is_relative_to(uploads_root): + raise ValueError(f"Invalid session ID for upload path: {session_id}") + return target + + +def remove_session_uploads(session_id: str) -> bool: + """Remove one session's upload directory when it exists.""" + + target = session_uploads_dir(session_id) + if not target.is_dir(): + return False + shutil.rmtree(target) + return True diff --git a/flocks/session/goal.py b/flocks/session/goal.py index 29c1b5b80..a43572e70 100644 --- a/flocks/session/goal.py +++ b/flocks/session/goal.py @@ -199,7 +199,6 @@ async def judge_goal_with_model( ], **provider_options, max_tokens=JUDGE_MAX_TOKENS, - temperature=0, ) payload = _extract_json_object(response.content) diff --git a/flocks/session/interaction_queue.py b/flocks/session/interaction_queue.py index b92464f9e..ba240d4f7 100644 --- a/flocks/session/interaction_queue.py +++ b/flocks/session/interaction_queue.py @@ -8,6 +8,7 @@ from pydantic import BaseModel, Field +from flocks.session.execution_mode import SessionExecutionMode from flocks.utils.id import Identifier @@ -35,6 +36,7 @@ class QueuedPrompt(BaseModel): mockReply: Optional[str] = None tools: Optional[Dict[str, bool]] = None system: Optional[str] = None + executionMode: SessionExecutionMode = SessionExecutionMode.BUILD status: str = "pending" createdAt: int = Field(default_factory=lambda: int(time.time() * 1000)) updatedAt: int = Field(default_factory=lambda: int(time.time() * 1000)) @@ -45,6 +47,7 @@ class InteractionQueue: _queues: Dict[str, List[QueuedPrompt]] = {} _locks: Dict[str, asyncio.Lock] = {} + _paused: set[str] = set() @classmethod def _lock_for(cls, session_id: str) -> asyncio.Lock: @@ -69,6 +72,7 @@ async def enqueue( mock_reply: Optional[str] = None, tools: Optional[Dict[str, bool]] = None, system: Optional[str] = None, + execution_mode: SessionExecutionMode = SessionExecutionMode.BUILD, ) -> QueuedPrompt: async with cls._lock_for(session_id): queue = cls._queues.setdefault(session_id, []) @@ -87,6 +91,7 @@ async def enqueue( mockReply=mock_reply, tools=dict(tools) if tools else None, system=system, + executionMode=execution_mode, ) queue.append(item) return item @@ -133,6 +138,8 @@ async def remove(cls, session_id: str, item_id: str) -> QueuedPrompt: @classmethod async def pop_next(cls, session_id: str) -> Optional[QueuedPrompt]: async with cls._lock_for(session_id): + if session_id in cls._paused: + return None queue = cls._queues.get(session_id, []) if not queue: return None @@ -159,6 +166,18 @@ async def promote(cls, session_id: str, item_id: str) -> QueuedPrompt: async def clear(cls, session_id: str) -> None: async with cls._lock_for(session_id): cls._queues.pop(session_id, None) + cls._paused.discard(session_id) + + @classmethod + async def pause(cls, session_id: str) -> None: + async with cls._lock_for(session_id): + if cls._queues.get(session_id): + cls._paused.add(session_id) + + @classmethod + async def resume(cls, session_id: str) -> None: + async with cls._lock_for(session_id): + cls._paused.discard(session_id) @classmethod def _find_locked(cls, session_id: str, item_id: str) -> QueuedPrompt: diff --git a/flocks/session/lifecycle/revert.py b/flocks/session/lifecycle/revert.py index 3fc7fc84f..ebfaec0cb 100644 --- a/flocks/session/lifecycle/revert.py +++ b/flocks/session/lifecycle/revert.py @@ -32,6 +32,29 @@ class SessionRevert: Provides simplified API for route handlers. """ + + @classmethod + async def ensure_replayable(cls, session: SessionInfo, message_id: str) -> None: + """Reject replay across the latest project-move boundary.""" + + metadata = session.metadata if isinstance(session.metadata, dict) else {} + move_metadata = metadata.get("projectMove") + if not isinstance(move_metadata, dict): + return + boundary_message_id = move_metadata.get("boundaryMessageID") + if not boundary_message_id: + return + + messages = await Message.list(session.id, include_archived=True) + message_ids = [message.id for message in messages] + if boundary_message_id not in message_ids or message_id not in message_ids: + raise ValueError( + "移动项目前的历史消息不能重发或重新生成" + ) + if message_ids.index(message_id) <= message_ids.index(boundary_message_id): + raise ValueError( + "移动项目前的历史消息不能重发或重新生成" + ) @classmethod async def revert( @@ -54,18 +77,36 @@ async def revert( session = await Session.get_by_id(session_id) if not session: raise ValueError(f"Session {session_id} not found") - - input_obj = RevertInput( - session_id=session_id, - message_id=message_id, - part_id=part_id, - ) - - return await SessionRevertManager.revert( - project_id=session.project_id, - input=input_obj, - worktree=session.directory, - ) + + async with Session.active_operation(session_id) as active_session: + if part_id is not None: + parts = await Message.parts(message_id, session_id) + part_belongs_to_message = any( + ( + part.get("id") + if isinstance(part, dict) + else getattr(part, "id", None) + ) == part_id + for part in parts + ) + if not part_belongs_to_message: + raise ValueError( + f"Part {part_id} does not belong to message {message_id}" + ) + + await cls.ensure_replayable(active_session, message_id) + + input_obj = RevertInput( + session_id=session_id, + message_id=message_id, + part_id=part_id, + ) + + return await SessionRevertManager.revert( + project_id=active_session.project_id, + input=input_obj, + worktree=active_session.directory, + ) @classmethod async def unrevert(cls, session_id: str) -> SessionInfo: @@ -164,9 +205,15 @@ async def revert( continue # Check if this is the revert point - if (msg.id == input.message_id and not input.part_id) or \ - (hasattr(part, "id") and part.id == input.part_id) or \ - (isinstance(part, dict) and part.get("id") == input.part_id): + is_requested_part = ( + input.part_id is not None + and msg.id == input.message_id + and ( + (hasattr(part, "id") and part.id == input.part_id) + or (isinstance(part, dict) and part.get("id") == input.part_id) + ) + ) + if (msg.id == input.message_id and not input.part_id) or is_requested_part: # Check if remaining parts have useful content has_useful = any( (isinstance(p, dict) and p.get("type") in ["text", "tool"]) or diff --git a/flocks/session/message.py b/flocks/session/message.py index 7fd6dcbe2..0ca768a1a 100644 --- a/flocks/session/message.py +++ b/flocks/session/message.py @@ -21,6 +21,7 @@ from flocks.utils.id import Identifier from flocks.storage.storage import Storage from flocks.session.recorder import Recorder +from flocks.session.execution_mode import SessionExecutionMode log = Log.create(service="message") @@ -376,6 +377,10 @@ class UserMessageInfo(BaseModel): tools: Optional[Dict[str, bool]] = Field(None, description="Tool availability") variant: Optional[str] = Field(None, description="Prompt variant") compacted: Optional[bool] = Field(None, description="Archived by compaction (soft-deleted)") + executionMode: SessionExecutionMode = Field( + SessionExecutionMode.BUILD, + description="Execution mode used for this user turn", + ) class AssistantMessageInfo(BaseModel): @@ -466,6 +471,25 @@ def _cancel_parts_flush_task(cls, session_id: str) -> None: if task and not task.done(): task.cancel() + @classmethod + async def quiesce_parts(cls, session_id: str, *, persist: bool) -> None: + """Stop a delayed parts flush, optionally persisting its latest cache.""" + task = cls._parts_flush_tasks.pop(session_id, None) + if task and not task.done(): + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + + async with _session_locks.get(session_id): + if persist and session_id in cls._parts_cache: + if cls._parts_storage_format.get(session_id) == "legacy": + await cls._persist_parts(session_id) + else: + for message_id in list(cls._parts_cache[session_id]): + await cls._persist_parts(session_id, message_id=message_id) + @classmethod def _cache_token(cls, session_id: str) -> tuple[int, int]: return cls._cache_epoch, cls._session_cache_generations.get(session_id, 0) @@ -942,6 +966,11 @@ def _normalize_stored_message( if not isinstance(model_raw, dict): model_raw = {} normalized["agent"] = normalized.get("agent") or "rex" + normalized["executionMode"] = ( + normalized.get("executionMode") + or normalized.get("execution_mode") + or SessionExecutionMode.BUILD.value + ) normalized["model"] = { "providerID": model_raw.get("providerID") or normalized.get("providerID") @@ -1863,19 +1892,73 @@ async def delete(cls, session_id: str, message_id: str) -> bool: return False messages = cls._messages_cache.get(session_id, []) if idx < len(messages) and messages[idx].id == message_id: - messages.pop(idx) + missing = object() + removed_message = messages.pop(idx) cls._rebuild_id_index(session_id) - if session_id in cls._parts_cache: - cls._parts_cache[session_id].pop(message_id, None) - cls._parts_revision_cache.get(session_id, {}).pop(message_id, None) - cls._parts_serialized_cache.get(session_id, {}).pop(message_id, None) + parts_cache = cls._parts_cache.get(session_id) + removed_parts = ( + parts_cache.pop(message_id, missing) + if parts_cache is not None + else missing + ) + parts_revision_cache = cls._parts_revision_cache.get(session_id) + removed_revision = ( + parts_revision_cache.pop(message_id, missing) + if parts_revision_cache is not None + else missing + ) + parts_serialized_cache = cls._parts_serialized_cache.get(session_id) + removed_serialized = ( + parts_serialized_cache.pop(message_id, missing) + if parts_serialized_cache is not None + else missing + ) + had_pending_parts_flush = session_id in cls._parts_flush_tasks cls._cancel_parts_flush_task(session_id) - await cls._persist_messages(session_id) - if cls._parts_storage_format.get(session_id) == "legacy": - await cls._persist_parts(session_id) - else: - await Storage.delete(cls._parts_item_key(session_id, message_id)) - cls._parts_persisted_mids.setdefault(session_id, set()).discard(message_id) + try: + await cls._persist_messages(session_id) + except BaseException: + # Message metadata is the deletion commit point. Restore + # every in-memory index/cache if it was not persisted so a + # later retry or process restart sees the same message. + messages.insert(idx, removed_message) + cls._rebuild_id_index(session_id) + if parts_cache is not None and removed_parts is not missing: + parts_cache[message_id] = removed_parts + if ( + parts_revision_cache is not None + and removed_revision is not missing + ): + parts_revision_cache[message_id] = removed_revision + if ( + parts_serialized_cache is not None + and removed_serialized is not missing + ): + parts_serialized_cache[message_id] = removed_serialized + if had_pending_parts_flush: + cls._schedule_parts_flush( + session_id, + message_id=message_id, + ) + raise + + try: + if cls._parts_storage_format.get(session_id) == "legacy": + await cls._persist_parts(session_id) + else: + await Storage.delete(cls._parts_item_key(session_id, message_id)) + cls._parts_persisted_mids.setdefault(session_id, set()).discard( + message_id + ) + except Exception as exc: + # Metadata deletion has committed. Orphaned parts are not + # user-visible and can be cleaned later; restoring the + # message here would make cache and durable metadata diverge. + log.warn("message.delete.parts_cleanup_failed", { + "session_id": session_id, + "message_id": message_id, + "error": str(exc), + }) log.info("message.deleted", {"id": message_id, "session_id": session_id}) return True return False @@ -1924,7 +2007,7 @@ async def clear(cls, session_id: str) -> int: Number of messages cleared """ await cls._ensure_message_cache(session_id) - + async with _session_locks.get(session_id): count = len(cls._messages_cache.get(session_id, [])) cls._messages_cache[session_id] = [] @@ -1935,18 +2018,33 @@ async def clear(cls, session_id: str) -> int: cls._parts_storage_format[session_id] = "per_message" cls._parts_fully_loaded.add(session_id) cls._cancel_parts_flush_task(session_id) - - # Persist changes + await cls._persist_messages(session_id) await Storage.clear(prefix=cls._parts_item_prefix(session_id)) await Storage.delete(cls._parts_blob_key(session_id)) - + log.info("messages.cleared", { "session_id": session_id, "count": count, }) - return count + + @classmethod + async def clear_active( + cls, + session_id: str, + *, + expected_generation: Optional[int] = None, + ) -> int: + """Clear history only while the owning session remains active.""" + + from flocks.session.session import Session + + return await Session.run_active_write( + session_id, + lambda: cls.clear(session_id), + expected_generation=expected_generation, + ) @classmethod def invalidate_cache(cls, session_id: Optional[str] = None) -> None: @@ -2479,7 +2577,7 @@ def delete(cls, session_id: str, message_id: str) -> bool: @classmethod def clear(cls, session_id: str) -> int: """Sync version""" - return cls._run_async(Message.clear(session_id)) + return cls._run_async(Message.clear_active(session_id)) @classmethod def get_text_content(cls, message: MessageInfo) -> str: diff --git a/flocks/session/plan_file.py b/flocks/session/plan_file.py new file mode 100644 index 000000000..ee650ac6d --- /dev/null +++ b/flocks/session/plan_file.py @@ -0,0 +1,180 @@ +"""Session-scoped plan files modeled after OpenCode's plan artifacts.""" + +from __future__ import annotations + +import os +import re +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Iterable, Optional + + +PLAN_DIRECTORY = Path(".flocks") / "plans" +_SAFE_COMPONENT_RE = re.compile(r"[^A-Za-z0-9._-]+") + + +@dataclass(frozen=True) +class SessionPlanFile: + """Resolved plan artifact for one session.""" + + path: Path + relative_path: str + permission_path: str + + +def _safe_component(value: object, fallback: str) -> str: + normalized = _SAFE_COMPONENT_RE.sub("-", str(value or "")).strip(".-") + return normalized or fallback + + +def session_plan_file( + session: Any, + *, + worktree: Optional[str] = None, +) -> SessionPlanFile: + """Return the stable, project-local plan file for a session.""" + + session_root = Path(str(session.directory)).expanduser().resolve(strict=False) + root = ( + Path(worktree).expanduser().resolve(strict=False) + if worktree and Path(worktree) != Path("/") + else session_root + ) + created = int(getattr(getattr(session, "time", None), "created", 0) or 0) + slug = _safe_component(getattr(session, "slug", None), "session") + filename = f"{created}-{slug}.md" + path = root / PLAN_DIRECTORY / filename + relative_path = Path(os.path.relpath(path, session_root)).as_posix() + return SessionPlanFile( + path=path, + relative_path=relative_path, + permission_path=(PLAN_DIRECTORY / filename).as_posix(), + ) + + +def context_plan_file(ctx: Any) -> Optional[SessionPlanFile]: + """Resolve the plan artifact in the tool's host or sandbox workspace.""" + + extra = getattr(ctx, "extra", {}) or {} + relative_path = str(extra.get("plan_relative_path") or "").strip() + permission_path = str(extra.get("plan_permission_path") or "").strip() + sandbox = extra.get("sandbox") + if isinstance(sandbox, dict) and sandbox.get("workspace_dir") and relative_path: + path = ( + Path(str(sandbox["workspace_dir"])).expanduser() + / Path(relative_path) + ).resolve(strict=False) + else: + raw_path = str(extra.get("plan_file_path") or "").strip() + if not raw_path: + return None + path = Path(raw_path).expanduser().absolute() + if not relative_path: + return None + return SessionPlanFile( + path=path, + relative_path=relative_path, + permission_path=permission_path, + ) + + +def _validation_root(ctx: Any) -> Path: + extra = getattr(ctx, "extra", {}) or {} + sandbox = extra.get("sandbox") + if isinstance(sandbox, dict) and sandbox.get("workspace_dir"): + return Path(str(sandbox["workspace_dir"])).expanduser().resolve(strict=False) + workspace_dir = extra.get("workspace_dir") + if workspace_dir: + return Path(str(workspace_dir)).expanduser().resolve(strict=False) + return Path.cwd().resolve(strict=False) + + +def _expected_plan_path(ctx: Any) -> Optional[Path]: + extra = getattr(ctx, "extra", {}) or {} + sandbox = extra.get("sandbox") + relative_path = str(extra.get("plan_relative_path") or "").strip() + if isinstance(sandbox, dict) and sandbox.get("workspace_dir") and relative_path: + return _validation_root(ctx) / Path(relative_path) + absolute_path = str(extra.get("plan_file_path") or "").strip() + if absolute_path: + return Path(absolute_path).expanduser().absolute() + return None + + +def _normalize_tool_path(ctx: Any, raw_path: object) -> Optional[Path]: + value = str(raw_path or "").strip() + if not value: + return None + path = Path(value).expanduser() + if not path.is_absolute(): + path = _validation_root(ctx) / path + return path.resolve(strict=False) + + +def is_current_plan_path(ctx: Any, raw_path: object) -> bool: + """Return whether a tool path is exactly the current session plan file.""" + + expected = _expected_plan_path(ctx) + candidate = _normalize_tool_path(ctx, raw_path) + if expected is None or candidate is None: + return False + # Keep the expected path lexical. If .flocks/plans or the file itself is a + # symlink, resolving the candidate moves it elsewhere and the comparison + # fails instead of granting an external write. + return os.path.normcase(str(candidate)) == os.path.normcase(str(expected.absolute())) + + +def plan_edit_patterns_allowed(ctx: Any, patterns: Iterable[object]) -> bool: + """Validate resolved edit permission patterns for Plan mode.""" + + values = list(patterns) + expected = _expected_plan_path(ctx) + expected_is_safe = bool( + expected + and os.path.normcase(str(expected.resolve(strict=False))) + == os.path.normcase(str(expected.absolute())) + ) + extra = getattr(ctx, "extra", {}) or {} + accepted_relative_paths = { + Path(str(value)).as_posix() + for value in ( + extra.get("plan_relative_path"), + extra.get("plan_permission_path"), + ) + if value + } + return bool(values) and all( + ( + expected_is_safe + and Path(str(value)).as_posix() in accepted_relative_paths + ) + or is_current_plan_path(ctx, value) + for value in values + ) + + +def plan_file_prompt( + session: Any, + *, + plan: Optional[SessionPlanFile] = None, +) -> str: + """Build the OpenCode-style per-turn plan file reminder.""" + + plan = plan or session_plan_file(session) + if plan.path.is_file(): + file_guidance = ( + f"A plan file already exists at `{plan.relative_path}`. Read it and " + "update it incrementally with the edit or write tool." + ) + else: + file_guidance = ( + f"No plan file exists yet. Create it at `{plan.relative_path}` with " + "the write tool." + ) + return f"""## Plan File + +{file_guidance} + +This is the only file you may edit in Plan mode. Keep the decision-complete +implementation plan in this file, then call plan_exit. +""" diff --git a/flocks/session/policy.py b/flocks/session/policy.py index 998779b70..89c07556f 100644 --- a/flocks/session/policy.py +++ b/flocks/session/policy.py @@ -11,6 +11,8 @@ from typing import Optional, TYPE_CHECKING +from flocks.auth.context import API_TOKEN_SERVICE_USER_ID + if TYPE_CHECKING: from flocks.auth.context import AuthUser from flocks.session.session import SessionInfo @@ -34,6 +36,11 @@ def _resolve_user(user: Optional["AuthUser"]) -> Optional["AuthUser"]: def is_owner(session: "SessionInfo", user: Optional["AuthUser"]) -> bool: if user is None: return False + if session.owner_user_id == API_TOKEN_SERVICE_USER_ID or ( + not session.owner_user_id + and session.owner_username == API_TOKEN_SERVICE_USER_ID + ): + return user.id == API_TOKEN_SERVICE_USER_ID if session.owner_user_id and session.owner_user_id == user.id: return True if session.owner_username and session.owner_username == user.username: @@ -55,6 +62,17 @@ def is_local_shared(session: "SessionInfo") -> bool: def _has_no_owner(session: "SessionInfo") -> bool: return not session.owner_user_id and not session.owner_username + @staticmethod + def _is_system_owned(session: "SessionInfo") -> bool: + return session.owner_user_id == API_TOKEN_SERVICE_USER_ID or ( + not session.owner_user_id + and session.owner_username == API_TOKEN_SERVICE_USER_ID + ) + + @classmethod + def _is_admin_managed(cls, session: "SessionInfo") -> bool: + return cls._has_no_owner(session) or cls._is_system_owned(session) + @classmethod def is_shared( cls, @@ -128,14 +146,14 @@ def can_read( - No auth context (CLI/internal runtime): keep legacy permissive behaviour. - Logged-in users: owner, local-shared readers, shared readers, or admins - managing ownerless legacy/channel sessions. + managing system-owned and ownerless legacy/channel sessions. """ resolved = cls._resolve_user(user) if resolved is None: return True if cls.is_owner(session, resolved): return True - if cls._has_no_owner(session) and cls.is_admin(resolved): + if cls._is_admin_managed(session) and cls.is_admin(resolved): return True return cls.is_shared_read_only(session, resolved, shared_project_ids) @@ -144,15 +162,15 @@ def can_write(cls, session: "SessionInfo", user: Optional["AuthUser"] = None) -> """ Session write permission. - Owner can always write. Admins may repair/manage ownerless sessions - accumulated before local ownership was available. + Owner can always write. Admins may manage system-owned sessions and + repair ownerless sessions accumulated before ownership was available. """ resolved = cls._resolve_user(user) if resolved is None: return False if cls.is_owner(session, resolved): return True - if cls._has_no_owner(session) and cls.is_admin(resolved): + if cls._is_admin_managed(session) and cls.is_admin(resolved): return True return False @@ -163,6 +181,6 @@ def can_delete(cls, session: "SessionInfo", user: Optional["AuthUser"]) -> bool: return False if cls.is_owner(session, resolved): return True - if cls._has_no_owner(session) and cls.is_admin(resolved): + if cls._is_admin_managed(session) and cls.is_admin(resolved): return True return False diff --git a/flocks/session/prompt.py b/flocks/session/prompt.py index 56466f8cf..21fb0ebb3 100644 --- a/flocks/session/prompt.py +++ b/flocks/session/prompt.py @@ -1077,6 +1077,7 @@ async def build_system_prompts( agent_prompt: Optional[str], provider_id: str, model_id: str, + execution_mode_prompt: Optional[str] = None, prompt_tool_names: Iterable[str] = (), tool_revision: Optional[int] = None, memory_bootstrap_data: Optional[Dict[str, Any]] = None, @@ -1165,6 +1166,13 @@ async def build_custom_context() -> Optional[str]: digest_inputs={"agent_name": agent_name, "agent_prompt": agent_prompt or ""}, builder=lambda: cls._normalize_prompt_text(agent_prompt), ), + cls._build_cached_prompt_block( + static_cache=static_cache, + name="execution_mode", + cache_scope="runtime", + digest_inputs={"prompt": execution_mode_prompt or ""}, + builder=lambda: cls._normalize_prompt_text(execution_mode_prompt), + ), cls._build_cached_prompt_block( static_cache=static_cache, name="memory_snapshot", diff --git a/flocks/session/runner.py b/flocks/session/runner.py index 69aec5144..ab388b073 100644 --- a/flocks/session/runner.py +++ b/flocks/session/runner.py @@ -47,6 +47,11 @@ ReasoningDeltaEvent, ReasoningEndEvent, ) +from flocks.session.streaming.timeouts import ( + DEFAULT_FIRST_CHUNK_TIMEOUT_S, + DEFAULT_ONGOING_CHUNK_TIMEOUT_S, + resolve_llm_stream_timeouts, +) from flocks.session.callable_schema import list_session_callable_tool_infos from flocks.agent.registry import Agent from flocks.agent.agent import AgentInfo @@ -66,6 +71,13 @@ is_text_extractable_mime, extract_file_text, ) +from flocks.session.execution_mode import ( + SessionExecutionMode, + execution_mode_prompt, + is_tool_allowed, + runtime_execution_mode, +) +from flocks.session.plan_file import session_plan_file log = Log.create(service="session.runner") @@ -102,13 +114,13 @@ def _annotate_with_provider_version(tool_info: Any, description: Optional[str]) # Maximum seconds to wait for the *first* chunk from the LLM stream. # If the model never starts responding, the stream times out and the session # surfaces a clear error rather than hanging forever. -LLM_STREAM_FIRST_CHUNK_TIMEOUT_S = 60 +LLM_STREAM_FIRST_CHUNK_TIMEOUT_S = DEFAULT_FIRST_CHUNK_TIMEOUT_S # Once the stream has started (at least one chunk received), allow a much # longer gap between chunks. Some models pause for extended periods between # reasoning and content generation phases; a tight inter-chunk timeout causes # spurious failures in those cases. -LLM_STREAM_ONGOING_CHUNK_TIMEOUT_S = 300 +LLM_STREAM_ONGOING_CHUNK_TIMEOUT_S = DEFAULT_ONGOING_CHUNK_TIMEOUT_S _RETRYABLE_TRANSPORT_EXCEPTIONS = ( httpx.TransportError, @@ -204,6 +216,43 @@ class ToolCall: name: str arguments: Dict[str, Any] + +@dataclass +class LlmAttemptState: + """Observable side effects accumulated across retries for one model.""" + + received_chunk: bool = False + observable_output_started: bool = False + tool_execution_started: bool = False + + @property + def replay_safe(self) -> bool: + """Whether the same logical LLM call can safely run on another model.""" + return not self.observable_output_started and not self.tool_execution_started + + +@dataclass(frozen=True) +class FailoverDecision: + """Hermes-aligned retry/failover classification for a provider error.""" + + eligible: bool + reason: str + same_model_retries: int = 3 + + +@dataclass +class StepFailure: + """Failure details returned to SessionLoop when finalization is deferred.""" + + message: str + error_data: Dict[str, Any] + assistant_message_id: Optional[str] + reason: str + allow_fallback: bool + attempt_state: LlmAttemptState + attempts: int = 0 + + @dataclass class StepResult: """Result of a single processing step.""" @@ -212,6 +261,7 @@ class StepResult: tool_calls: List[ToolCall] = field(default_factory=list) error: Optional[str] = None usage: Optional[Dict[str, int]] = None + failure: Optional[StepFailure] = None @dataclass @@ -257,6 +307,10 @@ def __init__( session_ctx: Optional[Any] = None, # SessionContext interface memory_bootstrap_data: Optional[Dict[str, Any]] = None, static_cache: Optional[Dict[str, Any]] = None, + defer_step_errors: bool = False, + failover_available: bool = False, + turn_additional_context: Optional[str] = None, + session_start_pending: bool = False, ): self.session = session from flocks.session.core.defaults import fallback_provider_id, fallback_model_id @@ -271,6 +325,12 @@ def __init__( self.session_ctx = session_ctx # SessionContext interface for decoupled access self._memory_bootstrap_data: Optional[Dict[str, Any]] = memory_bootstrap_data self._static_cache = static_cache if static_cache is not None else {} + self._defer_step_errors = defer_step_errors + self._failover_available = failover_available + self._turn_additional_context = turn_additional_context + self._session_start_pending = session_start_pending + self._session_start_fired = False + self._attempt_state = LlmAttemptState() @staticmethod def _canonical_tool_signature(tool_name: str, arguments: Dict[str, Any]) -> str: @@ -308,6 +368,26 @@ def _should_warn_about_tool_loop(self, *, last_user_id: str) -> bool: state = self._get_tool_loop_guard_state(last_user_id=last_user_id) return int(state.get("exact_count", 0)) >= max(2, REPEATED_EXACT_TOOL_CALL_HALT_THRESHOLD - 1) + async def _run_session_start_hook(self, agent: Any) -> None: + """Run SessionStart after the first system prompt has been built.""" + if not self._session_start_pending or self._session_start_fired: + return + self._session_start_fired = True + try: + from flocks.hooks.pipeline import HookPipeline + + await HookPipeline.run_session_start({ + "sessionID": self.session.id, + "workspace": self.session.directory, + "agent": agent.name, + "model": { + "providerID": self.provider_id, + "modelID": self.model_id, + }, + }) + except Exception as exc: + log.debug("runner.hook.session_start.error", {"error": str(exc)}) + def _build_tool_loop_halt_message( self, *, @@ -498,13 +578,43 @@ async def _list_callable_tool_infos_for_turn( agent: AgentInfo, messages: List[MessageInfo], ) -> Tuple[List[Any], Dict[str, Any]]: + execution_mode = self._execution_mode_from_messages(messages) result = await list_session_callable_tool_infos( session_id=self.session.id, declared_tool_names=getattr(agent, "tools", None), step=self._step, event_publish_callback=self.callbacks.event_publish_callback, ) - return result.tool_infos, dict(result.metadata) + tool_infos = [ + tool_info + for tool_info in result.tool_infos + if is_tool_allowed(execution_mode, tool_info.name) + ] + if ( + execution_mode == SessionExecutionMode.PLAN + and all(tool_info.name != "plan_exit" for tool_info in tool_infos) + ): + plan_exit = ToolRegistry.get("plan_exit") + if plan_exit is not None and getattr(plan_exit.info, "enabled", True): + tool_infos.append(plan_exit.info) + metadata = dict(result.metadata) + metadata["executionMode"] = execution_mode.value + metadata["modeAllowedToolNames"] = sorted( + tool_info.name for tool_info in tool_infos + ) + return tool_infos, metadata + + @staticmethod + def _execution_mode_from_messages( + messages: Optional[List[MessageInfo]], + ) -> SessionExecutionMode: + for message in reversed(messages or []): + if getattr(message, "role", None) != MessageRole.USER: + continue + return runtime_execution_mode( + getattr(message, "executionMode", None) + ) + return runtime_execution_mode(None) @staticmethod def _get_prompt_tool_names_from_schema(tools: List[Dict[str, Any]]) -> Tuple[str, ...]: @@ -1093,6 +1203,122 @@ def is_aborted(self) -> bool: if self._external_abort is not None and self._external_abort.is_set(): return True return False + + @staticmethod + def classify_failover_error(error: Dict[str, Any]) -> FailoverDecision: + """Classify a provider failure using Hermes-compatible switch timing. + + The classifier deliberately requires an API-shaped error (status code, + APIError marker, or a known provider response pattern). Local Python, + storage, hook, and tool failures must never move a user turn to another + model. + """ + data = error.get("data") or {} + status_code = data.get("statusCode") + try: + status_code = int(status_code) if status_code is not None else None + except (TypeError, ValueError): + status_code = None + + message = str(data.get("message") or error.get("message") or "") + lowered = message.lower() + error_name = str(error.get("name") or "") + + if status_code == 413 or any(pattern in lowered for pattern in ( + "context length", "context_length", "context window", "prompt is too long", + "request entity too large", "payload too large", + )): + return FailoverDecision(False, "context_overflow") + if error_name in {"CancelledError", "MessageAbortedError", "AbortedError"}: + return FailoverDecision(False, "cancelled") + + quota_or_rate_limited = any(pattern in lowered for pattern in ( + "rate limit", "too many requests", "quota exceeded", "resource exhausted", + "insufficient quota", "billing limit", + )) + # Some providers report exhausted quota as HTTP 401/403 rather than + # 429. Classify the semantic error before the generic auth branch so + # the primary receives the same cooldown as other quota failures. + if status_code == 429 or quota_or_rate_limited: + reason = "billing" if any( + pattern in lowered for pattern in ("billing", "insufficient quota") + ) else "rate_limit" + return FailoverDecision(True, reason, 0) + if status_code in {401, 403}: + return FailoverDecision(True, "auth", 0) + if status_code == 402: + return FailoverDecision(True, "billing", 0) + + if "model" in lowered and any(pattern in lowered for pattern in ( + "not found", "model_not_found", "unknown model", "no such model", + )): + return FailoverDecision(True, "model_not_found", 0) + + if status_code == 404: + if any(pattern in lowered for pattern in ( + "model not found", "model_not_found", "unknown model", "no such model", + )): + return FailoverDecision(True, "model_not_found", 0) + return FailoverDecision(True, "unknown_api", 3) + + if status_code in {408, 504} or data.get("isConnectionError") is True or any( + pattern in lowered + for pattern in ("timeout", "timed out", "connection error", "connection reset") + ): + return FailoverDecision(True, "timeout", 1) + if status_code in {503, 529} or any( + pattern in lowered for pattern in ("overloaded", "temporarily unavailable") + ): + return FailoverDecision(True, "overloaded", 1) + if status_code in {500, 502}: + return FailoverDecision(True, "server_error", 3) + + if any(pattern in lowered for pattern in ( + "content policy", "content filter", "content_filter", "safety policy", + "policy violation", + )): + return FailoverDecision(True, "content_policy", 0) + if error_name == "JSONDecodeError" or any( + pattern in lowered for pattern in ( + "malformed response", "invalid response", "empty choices", + "returned choice with null", "null message", + ) + ): + return FailoverDecision(True, "invalid_response", 0) + + if status_code is not None and 400 <= status_code < 500: + return FailoverDecision(True, "provider_request", 0) + if error_name == "APIError" or data.get("isRetryable") is True: + return FailoverDecision(True, "unknown_api", 3) + return FailoverDecision(False, "local_error") + + def _deferred_failure_result( + self, + *, + message: str, + error_data: Dict[str, Any], + assistant_message_id: Optional[str], + decision: FailoverDecision, + attempts: int, + ) -> StepResult: + state = LlmAttemptState( + received_chunk=self._attempt_state.received_chunk, + observable_output_started=self._attempt_state.observable_output_started, + tool_execution_started=self._attempt_state.tool_execution_started, + ) + return StepResult( + action="stop", + error=message, + failure=StepFailure( + message=message, + error_data=error_data, + assistant_message_id=assistant_message_id, + reason=decision.reason, + allow_fallback=decision.eligible and state.replay_safe, + attempt_state=state, + attempts=attempts, + ), + ) async def _process_step( self, @@ -1100,6 +1326,17 @@ async def _process_step( last_user: MessageInfo, ) -> StepResult: """Process a single step in the loop with retry logic.""" + self._attempt_state = LlmAttemptState() + turn_execution_mode = runtime_execution_mode( + getattr(last_user, "executionMode", None) + ) + self._turn_execution_mode = turn_execution_mode + from flocks.project.instance import Instance + + self._turn_plan_file = session_plan_file( + self.session, + worktree=Instance.get_worktree(), + ) # Check for CLI callbacks (if running in CLI mode) # Only use CLI fallback if no callbacks were explicitly provided via constructor has_explicit_callbacks = any([ @@ -1137,6 +1374,18 @@ async def _process_step( provider = Provider.get(self.provider_id) if not provider: error = f"Provider {self.provider_id} not found" + if self._defer_step_errors: + return self._deferred_failure_result( + message=error, + error_data={ + "name": "ProviderUnavailableError", + "message": error, + "data": {"message": error}, + }, + assistant_message_id=None, + decision=FailoverDecision(True, "provider_unavailable", 0), + attempts=0, + ) error_dict = self._build_session_error_dict( error, name="ProviderUnavailableError", @@ -1165,6 +1414,18 @@ async def _process_step( if not provider.is_configured(): error = f"Provider {self.provider_id} not configured" + if self._defer_step_errors: + return self._deferred_failure_result( + message=error, + error_data={ + "name": "ProviderUnavailableError", + "message": error, + "data": {"message": error}, + }, + assistant_message_id=None, + decision=FailoverDecision(True, "provider_unavailable", 0), + attempts=0, + ) error_dict = self._build_session_error_dict( error, name="ProviderConfigurationError", @@ -1212,6 +1473,11 @@ async def device_asset_prompt_factory() -> Optional[str]: agent_prompt=getattr(agent, "prompt", None), provider_id=self.provider_id, model_id=self.model_id, + execution_mode_prompt=execution_mode_prompt( + turn_execution_mode, + session=self.session, + plan_file=self._turn_plan_file, + ), prompt_tool_names=prompt_tool_names, tool_revision=ToolRegistry.revision(), memory_bootstrap_data=self._memory_bootstrap_data, @@ -1225,6 +1491,11 @@ async def device_asset_prompt_factory() -> Optional[str]: ) self._log_perf("runner.process_step.system_prompts_ready", prompts_started_at, prompt_count=len(system_prompts)) + await self._run_session_start_hook(agent) + + if self._turn_additional_context: + system_prompts.append(self._turn_additional_context) + if self._should_use_text_tool_call_mode() and tools: text_tool_catalog = self._build_text_tool_call_catalog_prompt(tools) if text_tool_catalog: @@ -1259,29 +1530,6 @@ async def device_asset_prompt_factory() -> Optional[str]: from flocks.session.prompt_strings import PROMPT_REPEATED_TOOL_CALLS system_prompts.append(PROMPT_REPEATED_TOOL_CALLS) - # Hook pipeline: chat.message stage - try: - from flocks.hooks.pipeline import HookPipeline - user_text = await Message.get_text_content(last_user) - hook_input = { - "sessionID": self.session.id, - "workspace": self.session.directory, - "agent": agent.name, - "model": {"providerID": self.provider_id, "modelID": self.model_id}, - "message": { - "id": last_user.id, - "role": "user", - "content": user_text, - }, - } - hook_output = {"message": {"variant": getattr(last_user, "variant", None)}} - ctx = await HookPipeline.run_chat_message(hook_input, hook_output) - variant = ctx.output.get("message", {}).get("variant") if ctx else None - if variant: - await Message.update(self.session.id, last_user.id, variant=variant) - except Exception as e: - log.debug("runner.hook.chat_message.error", {"error": str(e)}) - # Convert messages to chat format with error handling try: queued_user_message_ids = self._get_queued_user_message_ids(messages) @@ -1397,7 +1645,7 @@ async def device_asset_prompt_factory() -> Optional[str]: try: # Set status to busy SessionStatus.set(self.session.id, SessionStatusBusy()) - + # Call LLM with tools result = await self._call_llm( provider=provider, @@ -1429,7 +1677,11 @@ async def device_asset_prompt_factory() -> Optional[str]: if (result.action == "stop" and not result.error and not result.content and not result.tool_calls): empty_attempt += 1 - if empty_attempt <= MAX_EMPTY_RETRIES: + unsafe_auto_replay = ( + self._defer_step_errors + and not self._attempt_state.replay_safe + ) + if empty_attempt <= MAX_EMPTY_RETRIES and not unsafe_auto_replay: # Record usage for this empty attempt even though we are # about to retry – the provider may have already charged # for the tokens returned in this response. @@ -1456,10 +1708,16 @@ async def device_asset_prompt_factory() -> Optional[str]: # All retries exhausted — surface a clear error so the # user knows the model is incompatible, rather than # silently hanging or showing a blank response. - empty_error_msg = ( - f"Model '{self.model_id}' returned an empty response " - f"after {MAX_EMPTY_RETRIES} retries." - ) + if unsafe_auto_replay: + empty_error_msg = ( + f"Model '{self.model_id}' returned no final content after " + "starting observable output; the call was not replayed." + ) + else: + empty_error_msg = ( + f"Model '{self.model_id}' returned an empty response " + f"after {MAX_EMPTY_RETRIES} retries." + ) log.error("runner.step.empty_response_exhausted", { "session_id": self.session.id, "model": self.model_id, @@ -1474,6 +1732,14 @@ async def device_asset_prompt_factory() -> Optional[str]: "attempts": empty_attempt, }, } + if self._defer_step_errors: + return self._deferred_failure_result( + message=empty_error_msg, + error_data=empty_error_dict, + assistant_message_id=assistant_msg.id, + decision=FailoverDecision(True, "empty_response", 3), + attempts=empty_attempt, + ) if self.callbacks.on_error: await self.callbacks.on_error(empty_error_msg) await Message.update( @@ -1522,12 +1788,12 @@ async def device_asset_prompt_factory() -> Optional[str]: finish = "tool-calls" if result.tool_calls else "stop" await Message.update(self.session.id, assistant_msg.id, finish=finish) await self._record_usage_if_available(result.usage, message_id=assistant_msg.id) - + # Note: Compaction check is now done in the main loop (run()) before processing step # This matches Flocks's logic: check lastFinished.tokens at loop start return result - + except Exception as e: error_attempt += 1 error_log_context = { @@ -1541,7 +1807,24 @@ async def device_asset_prompt_factory() -> Optional[str]: # Check if retryable retry_message = SessionRetry.retryable(error_dict) - will_retry = retry_message is not None and error_attempt <= MAX_ERROR_RETRIES + failover_decision = self.classify_failover_error(error_dict) + retry_limit = MAX_ERROR_RETRIES + if ( + self._defer_step_errors + and failover_decision.eligible + ): + retry_limit = failover_decision.same_model_retries + will_retry = error_attempt <= retry_limit + if will_retry and retry_message is None: + retry_message = ( + f"Provider error ({failover_decision.reason}), retrying..." + ) + else: + will_retry = retry_message is not None and error_attempt <= retry_limit + if self._defer_step_errors and not self._attempt_state.replay_safe: + # Retrying after text/reasoning/tool activity can duplicate + # visible output or execute a tool twice. + will_retry = False if will_retry: # Error is retryable and we have budget left @@ -1558,7 +1841,7 @@ async def device_asset_prompt_factory() -> Optional[str]: "attempt": error_attempt, "delay_ms": delay_ms, "reason": retry_message, - "max_retries": MAX_ERROR_RETRIES, + "max_retries": retry_limit, }) # Set retry status @@ -1582,7 +1865,7 @@ async def device_asset_prompt_factory() -> Optional[str]: log.error("runner.step.max_retries_exceeded", { **error_log_context, "attempt": error_attempt, - "max_retries": MAX_ERROR_RETRIES, + "max_retries": retry_limit, }) else: log.error("runner.step.not_retryable", { @@ -1595,6 +1878,15 @@ async def device_asset_prompt_factory() -> Optional[str]: final_error_message = CONNECTION_ERROR_DISPLAY_MESSAGE error_dict["data"]["displayMessage"] = CONNECTION_ERROR_DISPLAY_MESSAGE + if self._defer_step_errors: + return self._deferred_failure_result( + message=final_error_message, + error_data=error_dict, + assistant_message_id=assistant_msg.id, + decision=failover_decision, + attempts=error_attempt, + ) + if self.callbacks.on_error: await self.callbacks.on_error(final_error_message) @@ -1634,6 +1926,66 @@ def _build_tokens_update(stream_usage: Optional[Dict[str, int]]) -> Optional[Dic }, } + @staticmethod + def _serialize_chat_message_for_langfuse(message: ChatMessage) -> Dict[str, Any]: + """Serialize the exact provider-bound message payload for Langfuse.""" + if hasattr(message, "model_dump"): + return message.model_dump(exclude_none=True) + return { + "role": message.role, + "content": message.content, + "reasoning": getattr(message, "reasoning", None), + "tool_calls": getattr(message, "tool_calls", None), + "tool_call_id": getattr(message, "tool_call_id", None), + "name": getattr(message, "name", None), + "custom_settings": getattr(message, "custom_settings", {}), + } + + @classmethod + def _build_langfuse_request_payload( + cls, + *, + step: int, + messages: List[ChatMessage], + request_tools: Optional[List[Dict[str, Any]]], + available_tools: List[Dict[str, Any]], + provider_options: Dict[str, Any], + ) -> Dict[str, Any]: + payload: Dict[str, Any] = { + "step": step, + "messages": [cls._serialize_chat_message_for_langfuse(message) for message in messages], + "provider_options": provider_options, + } + if request_tools is not None: + payload["request_tools"] = request_tools + if available_tools: + payload["available_tools"] = available_tools + return payload + + @staticmethod + def _build_langfuse_response_payload( + *, + action: str, + content: str, + reasoning: str, + finish_reason: Optional[str], + tool_calls: List[ToolCall], + ) -> Dict[str, Any]: + return { + "action": action, + "content": content, + "reasoning": reasoning, + "finish_reason": finish_reason, + "tool_calls": [ + { + "id": tool_call.id, + "name": tool_call.name, + "arguments": tool_call.arguments, + } + for tool_call in tool_calls + ], + } + def _resolve_usage_pricing(self) -> Optional[Any]: """Resolve pricing config for the current provider/model pair.""" from flocks.provider.usage_service import resolve_usage_pricing @@ -2017,9 +2369,42 @@ def _exception_to_error_dict(self, exception: Exception) -> Dict[str, Any]: "transportExceptionModule": type(transport_exception).__module__, }) - # Check if it's an API error with specific attributes - if hasattr(exception, 'status_code'): - status_code = getattr(exception, 'status_code') + # Provider SDKs expose HTTP status through several shapes. Walk the + # normal exception chain so lightweight wrapper errors do not hide it. + status_code = None + status_exception = None + seen: set[int] = set() + current: Optional[BaseException] = exception + while current is not None and id(current) not in seen: + seen.add(id(current)) + response = getattr(current, "response", None) + status_values = ( + getattr(current, "status_code", None), + getattr(response, "status_code", None), + getattr(current, "code", None), + ) + for value in status_values: + if callable(value): + try: + value = value() + except TypeError: + continue + value = getattr(value, "value", value) + if isinstance(value, tuple) and value: + value = value[0] + try: + normalized = int(value) + except (TypeError, ValueError): + continue + if 100 <= normalized <= 599: + status_code = normalized + status_exception = current + break + if status_code is not None: + break + current = current.__cause__ or current.__context__ + + if status_code is not None: error_dict["name"] = "APIError" error_dict["data"]["statusCode"] = status_code @@ -2028,9 +2413,13 @@ def _exception_to_error_dict(self, exception: Exception) -> Dict[str, Any]: error_dict["data"]["isRetryable"] = is_retryable # Extract response headers if available - if hasattr(exception, 'response') and hasattr(exception.response, 'headers'): - headers = dict(exception.response.headers) - error_dict["data"]["responseHeaders"] = headers + response = getattr(status_exception, "response", None) + headers = getattr(response, "headers", None) + if headers is not None: + try: + error_dict["data"]["responseHeaders"] = dict(headers) + except (TypeError, ValueError): + pass # Check for common retryable error patterns error_msg = str(exception).lower() @@ -2402,7 +2791,7 @@ async def _to_chat_messages( "type": "text", "text": "The following tool was executed by the user", }) - + if user_content_blocks and any( block.get("type") == "image" for block in user_content_blocks @@ -2724,6 +3113,19 @@ def _build_llm_response_payload( except Exception as e: log.debug("runner.sandbox_context_init_failed", {"error": str(e)}) + async def _on_tool_execution_start( + tool_name: str, + tool_input: Dict[str, Any], + ) -> None: + # Mark before hooks/callbacks/execution (including parallel tools) + # so a concurrent provider error can never replay side effects. + self._attempt_state.tool_execution_started = True + if self.callbacks.on_tool_start: + await self.callbacks.on_tool_start(tool_name, tool_input) + + turn_plan_file = getattr(self, "_turn_plan_file", None) + if turn_plan_file is None: + turn_plan_file = session_plan_file(self.session) processor = StreamProcessor( session_id=self.session.id, assistant_message=assistant_msg, @@ -2732,7 +3134,7 @@ def _build_llm_response_payload( permission_callback=self._handle_permission, text_delta_callback=self.callbacks.on_text_delta, reasoning_delta_callback=self.callbacks.on_reasoning_delta, - tool_start_callback=self.callbacks.on_tool_start, + tool_start_callback=_on_tool_execution_start, tool_end_callback=self.callbacks.on_tool_end, event_publish_callback=self.callbacks.event_publish_callback, session_key=self.session.id, @@ -2740,11 +3142,22 @@ def _build_llm_response_payload( workspace_dir=self.session.directory, langfuse_generation=None, step_index=self._step, + execution_mode=runtime_execution_mode( + getattr( + self, + "_turn_execution_mode", + SessionExecutionMode.BUILD, + ) + ).value, + plan_file_path=str(turn_plan_file.path), + plan_relative_path=turn_plan_file.relative_path, + plan_permission_path=turn_plan_file.permission_path, ) # Build provider options (thinking / reasoning / max_tokens) from flocks.provider.options import build_provider_options provider_options = build_provider_options(self.provider_id, self.model_id) + provider_tools = None if self._should_use_text_tool_call_mode() else (tools if tools else None) # Clean up any leftover reasoning state from a previous (failed) call if hasattr(self, '_current_reasoning_id'): @@ -2766,13 +3179,6 @@ def _build_llm_response_payload( generation_ctx = None if langfuse_is_active(): try: - input_preview = [] - for _msg in messages[-12:]: - _mc = _msg.content or "" - input_preview.append( - {"role": _msg.role, "chars": len(_mc), "preview": _mc[:240]} - ) - trace_tags = [ f"session:{self.session.id}", f"step:{self._step}", @@ -2780,36 +3186,47 @@ def _build_llm_response_payload( f"agent:{agent.name}", f"provider:{self.provider_id}", ] + request_payload = self._build_langfuse_request_payload( + step=self._step, + messages=messages, + request_tools=provider_tools, + available_tools=tools, + provider_options=provider_options, + ) trace_ctx = trace_scope( name="SessionRunner.step", session_id=self.session.id, tags=trace_tags, - input={ - "step": self._step, - "message_count": len(messages), - "tool_count": len(tools), - "last_user_preview": next( - ((m.content or "")[:280] for m in reversed(messages) if m.role == "user"), - "", - ), - }, + input=request_payload, metadata={ "provider_id": self.provider_id, "model_id": self.model_id, "agent": agent.name, "workspace": self.session.directory, + "message_count": len(messages), + "available_tool_count": len(tools), + "request_tool_count": len(provider_tools or []), + "tool_transport": ( + "provider_param" if provider_tools is not None else "text_prompt" + ), }, ) generation_ctx = generation_scope( parent=trace_ctx.observation, name="LLM.generate", model=self.model_id, - input=input_preview, + input=request_payload, metadata={ "provider_id": self.provider_id, "session_id": self.session.id, "step": self._step, - "tool_names": [t.get("function", {}).get("name", "") for t in tools][:50], + "agent": agent.name, + "workspace": self.session.directory, + "available_tool_count": len(tools), + "request_tool_count": len(provider_tools or []), + "tool_transport": ( + "provider_param" if provider_tools is not None else "text_prompt" + ), }, ) processor._langfuse_generation = generation_ctx.observation @@ -2842,7 +3259,6 @@ def _build_llm_response_payload( stream_usage: Optional[Dict[str, int]] = None # Stream response and convert chunks to events - provider_tools = None if self._should_use_text_tool_call_mode() else (tools if tools else None) if provider_tools is None and tools: log.info("runner.text_tool_call_mode.enabled", { "session_id": self.session.id, @@ -2904,6 +3320,14 @@ def _build_llm_response_payload( llm_call_started_at = time.perf_counter() first_chunk_logged = False aborted_during_stream = False + stream_timeouts = resolve_llm_stream_timeouts(provider, self.model_id) + log.debug("runner.llm.stream_timeouts", { + "provider_id": self.provider_id, + "model_id": self.model_id, + "first_chunk_timeout_s": stream_timeouts.first_chunk_s, + "ongoing_chunk_timeout_s": stream_timeouts.ongoing_chunk_s, + "local_endpoint": stream_timeouts.is_local, + }) try: async for chunk in _iter_with_chunk_timeout( provider.chat_stream( @@ -2917,10 +3341,11 @@ def _build_llm_response_payload( session_id=self.session.id, **provider_options, ), - first_chunk_timeout_s=LLM_STREAM_FIRST_CHUNK_TIMEOUT_S, - ongoing_chunk_timeout_s=LLM_STREAM_ONGOING_CHUNK_TIMEOUT_S, + first_chunk_timeout_s=stream_timeouts.first_chunk_s, + ongoing_chunk_timeout_s=stream_timeouts.ongoing_chunk_s, ): chunk_counts["total"] += 1 + self._attempt_state.received_chunk = True if not first_chunk_logged: first_chunk_logged = True self._log_perf( @@ -2929,20 +3354,20 @@ def _build_llm_response_payload( provider_id=self.provider_id, model_id=self.model_id, ) - + chunk_finish = getattr(chunk, 'finish_reason', None) if chunk_finish: stream_finish_reason = chunk_finish - + # Capture usage from chunk (providers may include it in the final chunk) if hasattr(chunk, 'usage') and chunk.usage: stream_usage = chunk.usage - + # Check for abort if self.is_aborted: aborted_during_stream = True break - + # Determine event type from chunk. A single chunk may carry any # combination of reasoning / text / tool_calls (e.g. Gemini bundles # them). We must not drop non-reasoning content when reasoning is @@ -2959,6 +3384,7 @@ def _build_llm_response_payload( self._current_reasoning_metadata = current_metadata if event_type == "reasoning-start" and not hasattr(self, '_current_reasoning_id'): + self._attempt_state.observable_output_started = True reasoning_id_counter += 1 self._current_reasoning_id = f"reasoning-{reasoning_id_counter}" self._current_reasoning_metadata = dict(chunk_metadata) @@ -2999,6 +3425,7 @@ def _build_llm_response_payload( # 1) Process reasoning delta (start reasoning block on first sight). if chunk_reasoning or (event_type == 'reasoning' and has_reasoning_metadata): + self._attempt_state.observable_output_started = True reasoning_text = chunk_reasoning or "" chunk_counts["reasoning"] += 1 log.debug("runner.reasoning.received", { @@ -3035,6 +3462,7 @@ def _build_llm_response_payload( # 3) Process text delta. if chunk_text: + self._attempt_state.observable_output_started = True chunk_counts["text"] += 1 if not text_started: await processor.process_event(TextStartEvent()) @@ -3046,6 +3474,9 @@ def _build_llm_response_payload( # 4) Process tool calls. if chunk_tool_calls: + # Tool fragments are persisted/accumulated and may become an + # executable call in this same await, so any replay is unsafe. + self._attempt_state.observable_output_started = True chunk_counts["tool"] += 1 for tc in chunk_tool_calls: await tool_accumulator.feed_chunk(tc) @@ -3159,6 +3590,7 @@ def _build_llm_response_payload( ) for tc_state in list(processor.tool_calls.values()) ] + finish_reason = processor.get_finish_reason() result_action = "continue" if tool_calls_for_result else "stop" response_payload = _build_llm_response_payload( content=content, @@ -3196,26 +3628,23 @@ def _build_llm_response_payload( log.debug("runner.hook.llm_after.error", {"error": str(exc)}) if tool_calls_for_result: + response_payload = self._build_langfuse_response_payload( + action="continue", + content=content, + reasoning=reasoning, + finish_reason=finish_reason, + tool_calls=tool_calls_for_result, + ) self._end_observability( generation_ctx, trace_ctx, - output={ - "content_preview": content[:600], - "content_chars": len(content), - "reasoning_chars": len(reasoning), - "tool_calls": [{"id": tc.id, "name": tc.name} for tc in tool_calls_for_result[:30]], - }, + output=response_payload, usage=stream_usage, metadata={ - "finish_reason": processor.get_finish_reason(), + "finish_reason": finish_reason, "status": "continue_with_tools", "tool_call_count": len(tool_calls_for_result), }, - trace_output={ - "status": "ok", - "next_action": "continue", - "finish_reason": processor.get_finish_reason(), - "tool_call_count": len(tool_calls_for_result), - }, + trace_output=response_payload, ) return StepResult( action=result_action, @@ -3224,24 +3653,23 @@ def _build_llm_response_payload( usage=stream_usage, ) + response_payload = self._build_langfuse_response_payload( + action="stop", + content=content, + reasoning=reasoning, + finish_reason=finish_reason, + tool_calls=tool_calls_for_result, + ) self._end_observability( generation_ctx, trace_ctx, - output={ - "content_preview": content[:600], - "content_chars": len(content), - "reasoning_chars": len(reasoning), - }, + output=response_payload, usage=stream_usage, metadata={ - "finish_reason": processor.get_finish_reason(), + "finish_reason": finish_reason, "status": "stop", "tool_call_count": 0, }, - trace_output={ - "status": "ok", - "next_action": "stop", - "finish_reason": processor.get_finish_reason(), - }, + trace_output=response_payload, ) return StepResult(action=result_action, content=content, usage=stream_usage) diff --git a/flocks/session/session.py b/flocks/session/session.py index b83d898c2..18ab7c2df 100644 --- a/flocks/session/session.py +++ b/flocks/session/session.py @@ -5,13 +5,17 @@ Based on Flocks' ported src/session/index.ts """ +import asyncio import contextvars import re -from typing import List, Dict, Any, Optional +import weakref +from contextlib import AsyncExitStack, asynccontextmanager +from typing import AsyncIterator, Awaitable, Callable, List, Dict, Any, Optional, TypeVar from datetime import datetime from pydantic import BaseModel, Field, ConfigDict from flocks.auth.context import ( + API_TOKEN_SERVICE_USER_ID, get_current_auth_user, reset_current_auth_user, set_current_auth_user, @@ -19,7 +23,7 @@ from flocks.storage.storage import Storage from flocks.utils.log import Log from flocks.utils.id import Identifier -from flocks.session.message import Message, MessageInfo, AssistantMessageInfo +from flocks.session.message import Message, MessageInfo, MessageRole, AssistantMessageInfo # Sentinel for explicitly setting a field to None via Session.update() _UNSET = object() @@ -30,6 +34,21 @@ # Title prefix patterns for default title detection PARENT_TITLE_PREFIX = "New session - " CHILD_TITLE_PREFIX = "Child session - " +MODEL_AUTO_SESSION_CATEGORIES = frozenset({"user", "entity-config", "workflow"}) +_WriteResult = TypeVar("_WriteResult") + + +class SessionNotFoundError(ValueError): + """Raised when a lifecycle-aware write cannot find its session.""" + + +class SessionInactiveError(RuntimeError): + """Raised when a lifecycle-aware write targets a non-active session.""" + + +def is_model_auto_session_category(category: Optional[str]) -> bool: + """Return whether a session category may use WebUI Auto mode.""" + return category in MODEL_AUTO_SESSION_CATEGORIES class SessionChangeStats(BaseModel): @@ -94,6 +113,13 @@ class SessionInfo(BaseModel): "Unpinned sessions follow the normal default-model resolution chain." ), ) + model_auto: bool = Field( + False, + description=( + "Whether WebUI Auto runtime failover was explicitly selected for " + "this session. Other session entry points ignore this flag." + ), + ) # Session hierarchy parent_id: Optional[str] = Field(None, alias="parentID", description="Parent session for branching") @@ -141,6 +167,15 @@ class Session: _id_index: Dict[str, str] = {} # Hot-path cache for repeatedly listing sessions in the UI. _all_sessions_cache: Optional[List[SessionInfo]] = None + # Serializes tree topology changes only. Ordinary writes use a keyed lock, + # so unrelated sessions no longer block each other. + _tree_lock = asyncio.Lock() + _lifecycle_locks: weakref.WeakValueDictionary[str, asyncio.Lock] = ( + weakref.WeakValueDictionary() + ) + _lifecycle_transition_ids: set[str] = set() + _lifecycle_generations: Dict[str, int] = {} + _active_operation_counts: Dict[str, int] = {} @staticmethod def _sort_sessions(sessions: List[SessionInfo]) -> List[SessionInfo]: @@ -175,12 +210,127 @@ def _sync_list_cache(cls, session: SessionInfo) -> None: remaining.append(session) cls._all_sessions_cache = cls._sort_sessions(remaining) + @classmethod + def _remove_from_list_cache(cls, session_id: str) -> None: + """Remove a permanently deleted session from derived caches.""" + if cls._all_sessions_cache is not None: + cls._all_sessions_cache = [ + session + for session in cls._all_sessions_cache + if session.id != session_id + ] + @classmethod def invalidate_cache(cls) -> None: """Clear in-memory indexes when the underlying storage changes.""" cls._id_index.clear() cls._all_sessions_cache = None + @classmethod + def is_lifecycle_transitioning(cls, session_id: str) -> bool: + """Return whether archive/delete currently owns this session lifecycle.""" + return session_id in cls._lifecycle_transition_ids + + @classmethod + def lifecycle_generation(cls, session_id: str) -> int: + """Return a token that changes whenever a lifecycle transition begins.""" + return cls._lifecycle_generations.get(session_id, 0) + + @classmethod + def lifecycle_lock(cls, session_id: str) -> asyncio.Lock: + """Return the lock that linearizes durable writes for one session.""" + + lock = cls._lifecycle_locks.get(session_id) + if lock is None: + lock = asyncio.Lock() + cls._lifecycle_locks[session_id] = lock + return lock + + @classmethod + def has_active_operations(cls, session_id: str) -> bool: + """Return whether a synchronous operation currently owns this session.""" + + return cls._active_operation_counts.get(session_id, 0) > 0 + + @classmethod + @asynccontextmanager + async def active_operation(cls, session_id: str) -> AsyncIterator[SessionInfo]: + """Prevent lifecycle transitions while a synchronous operation is running.""" + + async with cls.lifecycle_lock(session_id): + session = await cls.get_by_id_unfiltered(session_id) + if session is None: + raise SessionNotFoundError(f"Session {session_id} not found") + if session.status != "active" or cls.is_lifecycle_transitioning(session_id): + raise SessionInactiveError(f"Session {session_id} is not active") + cls._active_operation_counts[session_id] = ( + cls._active_operation_counts.get(session_id, 0) + 1 + ) + + try: + yield session + finally: + remaining = cls._active_operation_counts.get(session_id, 0) - 1 + if remaining > 0: + cls._active_operation_counts[session_id] = remaining + else: + cls._active_operation_counts.pop(session_id, None) + + @classmethod + async def get_by_id_unfiltered(cls, session_id: str) -> Optional[SessionInfo]: + """Get a session by ID without applying the ambient auth policy.""" + + token = set_current_auth_user(None) + try: + return await cls.get_by_id(session_id) + finally: + reset_current_auth_user(token) + + @classmethod + async def run_active_write( + cls, + session_id: str, + operation: Callable[[], Awaitable[_WriteResult]], + *, + expected_generation: Optional[int] = None, + ) -> _WriteResult: + """Run one durable write atomically with lifecycle transitions.""" + + async with cls.lifecycle_lock(session_id): + session = await cls.get_by_id_unfiltered(session_id) + if session is None: + raise SessionNotFoundError(f"Session {session_id} not found") + generation_changed = ( + expected_generation is not None + and cls.lifecycle_generation(session_id) != expected_generation + ) + if ( + session.status != "active" + or cls.is_lifecycle_transitioning(session_id) + or generation_changed + ): + raise SessionInactiveError(f"Session {session_id} is not active") + return await operation() + + @classmethod + async def _clear_project_move_metadata_locked(cls, session_id: str) -> bool: + """Remove a stale replay boundary while the session write lock is held.""" + + session = await cls.get_by_id_unfiltered(session_id) + if session is None: + raise SessionNotFoundError(f"Session {session_id} not found") + metadata = dict(getattr(session, "metadata", {}) or {}) + if "projectMove" not in metadata: + return False + + metadata.pop("projectMove") + updated_session = session.model_copy(update={"metadata": metadata}) + storage_key = f"session:{session.project_id}:{session.id}" + await Storage.set(storage_key, updated_session, "session") + cls._id_index[session.id] = storage_key + cls._sync_list_cache(updated_session) + return True + @staticmethod def has_pinned_model(session: Optional[SessionInfo]) -> bool: """Return whether a session has an explicit model lock.""" @@ -198,11 +348,21 @@ def explicit_model_updates(provider_id: str, model_id: str) -> Dict[str, Any]: "provider": provider_id, "model": model_id, "model_pinned": True, + "model_auto": False, } @classmethod def inherited_model_kwargs(cls, session: Optional[SessionInfo]) -> Dict[str, Any]: - """Return pinned model kwargs that should propagate to a child session.""" + """Return model preference kwargs that should propagate to a new session.""" + if ( + session + and is_model_auto_session_category(getattr(session, "category", "user")) + and getattr(session, "model_auto", False) + ): + return { + "model_auto": True, + "model_pinned": False, + } if not cls.has_pinned_model(session): return {} return { @@ -278,27 +438,77 @@ async def create( except Exception as e: log.warn("session.memory.default.error", {"error": str(e)}) - # Bind ownership from current auth context unless explicitly provided. - if "owner_user_id" not in kwargs or "owner_username" not in kwargs: + # Bind root ownership here; children inherit ownership from their parent below. + if parent_id is None and ( + "owner_user_id" not in kwargs or "owner_username" not in kwargs + ): current_user = get_current_auth_user() if current_user: kwargs.setdefault("owner_user_id", current_user.id) kwargs.setdefault("owner_username", current_user.username) + else: + kwargs.setdefault("owner_user_id", API_TOKEN_SERVICE_USER_ID) + kwargs.setdefault("owner_username", API_TOKEN_SERVICE_USER_ID) - session = SessionInfo( - project_id=project_id, - directory=directory, - title=title or cls._create_default_title(is_child), - parent_id=parent_id, - permission=permission, - **kwargs - ) - - # Save to storage - storage_key = f"session:{project_id}:{session.id}" - await Storage.set(storage_key, session, "session") - cls._id_index[session.id] = storage_key - cls._sync_list_cache(session) + async def persist(parent: Optional[SessionInfo] = None) -> SessionInfo: + if parent_id is not None: + if parent is None: + raise ValueError(f"Parent session {parent_id} not found") + child_owner_id = kwargs.get("owner_user_id") + child_owner_username = kwargs.get("owner_username") + if parent.owner_user_id and child_owner_id and parent.owner_user_id != child_owner_id: + raise ValueError("Child session owner must match its parent") + if parent.owner_username and child_owner_username and parent.owner_username != child_owner_username: + raise ValueError("Child session owner must match its parent") + if child_owner_id is None: + kwargs["owner_user_id"] = parent.owner_user_id + if child_owner_username is None: + kwargs["owner_username"] = parent.owner_username + + created = SessionInfo( + project_id=project_id, + directory=directory, + title=title or cls._create_default_title(is_child), + parent_id=parent_id, + permission=permission, + **kwargs + ) + + # The parent status check and child row creation share the archive + # lock, so an archive cannot commit between them in this process. + storage_key = f"session:{project_id}:{created.id}" + await Storage.set(storage_key, created, "session") + cls._id_index[created.id] = storage_key + cls._sync_list_cache(created) + return created + + from flocks.project.project import Project, ProjectDeletionError + + # Every creation path participates in the project lifecycle claim. If + # deletion wins the race, a waiting creator observes the removed marker + # and cannot recreate an orphan task under the hidden project. + async with Project.lifecycle_guard(project_id): + if project_id.startswith("prj_") and Project.is_removed(project_id): + raise ProjectDeletionError(f"Project {project_id} is no longer available") + + if parent_id is None: + session = await persist() + else: + async with cls._tree_lock: + async with cls.lifecycle_lock(parent_id): + parent = await Storage.get( + f"session:{project_id}:{parent_id}", + SessionInfo, + ) + if parent is None or parent.status == "deleted": + raise ValueError(f"Parent session {parent_id} not found") + if cls.is_lifecycle_transitioning(parent_id): + raise ValueError( + f"Parent session {parent_id} is changing lifecycle state" + ) + if parent.status != "active": + raise ValueError(f"Parent session {parent_id} is not active") + session = await persist(parent) try: from flocks.agent.registry import Agent @@ -506,6 +716,8 @@ async def update( cls, project_id: str, session_id: str, + *, + allow_inactive: bool = False, **updates ) -> Optional[SessionInfo]: """ @@ -519,58 +731,55 @@ async def update( Returns: Updated session info or None """ - session = await cls.get(project_id, session_id) - if not session: - return None - - # Field alias mapping (snake_case -> camelCase) - alias_map = { - "project_id": "projectID", - "parent_id": "parentID", - "owner_user_id": "ownerUserID", - "owner_username": "ownerUsername", - } - - # Update fields. - # Use ``_UNSET`` sentinel to explicitly set a field to None - # (plain ``None`` is skipped to preserve existing values). - update_data = session.model_dump(by_alias=True) - for key, value in updates.items(): - # Sentinel means "explicitly clear this field" - if value is _UNSET: - alias_key = alias_map.get(key, key) - if alias_key in update_data: - update_data[alias_key] = None - elif key in update_data: - update_data[key] = None - continue + async with cls.lifecycle_lock(session_id): + if cls.is_lifecycle_transitioning(session_id): + return None + session = await cls.get(project_id, session_id) + if not session or (session.status != "active" and not allow_inactive): + return None - if value is not None: - # Handle nested updates for summary, revert, time - if key == "summary" and isinstance(value, dict): - if update_data.get("summary"): - update_data["summary"].update(value) - else: - update_data["summary"] = value - elif key == "revert" and isinstance(value, dict): - update_data["revert"] = value - else: - # Check both original key and aliased key + alias_map = { + "project_id": "projectID", + "parent_id": "parentID", + "owner_user_id": "ownerUserID", + "owner_username": "ownerUsername", + } + + # Use ``_UNSET`` to explicitly clear a field; ordinary ``None`` + # remains a no-op for backward compatibility. + update_data = session.model_dump(by_alias=True) + for key, value in updates.items(): + if value is _UNSET: alias_key = alias_map.get(key, key) if alias_key in update_data: - update_data[alias_key] = value + update_data[alias_key] = None elif key in update_data: - update_data[key] = value - - # Update timestamp - if "time" not in update_data: - update_data["time"] = {} - update_data["time"]["updated"] = int(datetime.now().timestamp() * 1000) - - updated_session = SessionInfo(**update_data) - await Storage.set(f"session:{project_id}:{session_id}", updated_session, "session") - cls._id_index[session_id] = f"session:{project_id}:{session_id}" - cls._sync_list_cache(updated_session) + update_data[key] = None + continue + + if value is not None: + if key == "summary" and isinstance(value, dict): + if update_data.get("summary"): + update_data["summary"].update(value) + else: + update_data["summary"] = value + elif key == "revert" and isinstance(value, dict): + update_data["revert"] = value + else: + alias_key = alias_map.get(key, key) + if alias_key in update_data: + update_data[alias_key] = value + elif key in update_data: + update_data[key] = value + + if "time" not in update_data: + update_data["time"] = {} + update_data["time"]["updated"] = int(datetime.now().timestamp() * 1000) + + updated_session = SessionInfo(**update_data) + await Storage.set(f"session:{project_id}:{session_id}", updated_session, "session") + cls._id_index[session_id] = f"session:{project_id}:{session_id}" + cls._sync_list_cache(updated_session) log.info("session.updated", { "id": session_id, @@ -578,11 +787,125 @@ async def update( }) return updated_session + + @classmethod + async def move_to_project( + cls, + source_project_id: str, + session_id: str, + *, + target_project_id: str, + target_directory: str, + target_owner_id: Optional[str] = None, + additional_busy_check: Optional[Callable[[str], bool]] = None, + ) -> Optional[SessionInfo]: + """Move a complete active root session tree to another project.""" + + if source_project_id == target_project_id: + return await cls.get(source_project_id, session_id) + + from flocks.project.project import Project + + async with AsyncExitStack() as project_guards: + for project_id in sorted({source_project_id, target_project_id}): + await project_guards.enter_async_context(Project.lifecycle_guard(project_id)) + + if ( + target_owner_id is not None + and Project.registry_state(target_project_id, owner_id=target_owner_id) != "active" + ): + return None + + sessions = await cls._begin_lifecycle_transition( + source_project_id, + session_id, + require_root=True, + ) + if not sessions: + return None + + try: + if any(session.status != "active" for session in sessions): + return None + if any(session.revert is not None for session in sessions): + return None + + from flocks.session.session_loop import SessionLoop + + if any( + SessionLoop.is_running(session.id) + or ( + additional_busy_check is not None + and additional_busy_check(session.id) + ) + for session in sessions + ): + return None + + updated_at = int(datetime.now().timestamp() * 1000) + move_boundaries: Dict[str, Optional[str]] = {} + for session in sessions: + messages = await Message.list(session.id, include_archived=True) + move_boundaries[session.id] = messages[-1].id if messages else None + + moved_sessions = [ + session.model_copy(update={ + "project_id": target_project_id, + "directory": target_directory, + "metadata": { + **( + session.metadata + if isinstance(session.metadata, dict) + else {} + ), + "projectMove": { + "sourceProjectID": source_project_id, + "targetProjectID": target_project_id, + "boundaryMessageID": move_boundaries[session.id], + "movedAt": updated_at, + }, + }, + "time": session.time.model_copy(update={"updated": updated_at}), + }) + for session in sessions + ] + await Storage.mutate_many( + set_entries=[ + (f"session:{target_project_id}:{session.id}", session, "session") + for session in moved_sessions + ], + delete_keys=[ + f"session:{source_project_id}:{session.id}" + for session in sessions + ], + ) + for session in moved_sessions: + cls._id_index[session.id] = f"session:{target_project_id}:{session.id}" + cls._sync_list_cache(session) + + log.info("session.project_moved", { + "id": session_id, + "source_project_id": source_project_id, + "target_project_id": target_project_id, + "affected_sessions": len(moved_sessions), + }) + return moved_sessions[0] + finally: + await cls._end_lifecycle_transition(sessions) @classmethod async def delete(cls, project_id: str, session_id: str) -> bool: + """Permanently delete a session tree and all application-owned data.""" + + from flocks.project.project import Project + + async with Project.lifecycle_guard(project_id): + return await cls._delete_locked(project_id, session_id) + + @classmethod + async def _delete_locked(cls, project_id: str, session_id: str) -> bool: """ - Delete a session (soft delete) + Permanently delete a session tree and all application-owned data. Also deletes child sessions. @@ -593,52 +916,109 @@ async def delete(cls, project_id: str, session_id: str) -> bool: Returns: True if deleted """ - session = await cls.get(project_id, session_id) - if not session: + sessions = await cls._begin_lifecycle_transition(project_id, session_id) + if not sessions: return False - - # Delete child sessions first - children = await cls.children(project_id, session_id) - for child in children: - await cls.delete(project_id, child.id) - - # Soft delete - await cls.update(project_id, session_id, status="deleted") - cls._id_index.pop(session_id, None) - - # Clear messages - await Message.clear(session_id) + deleted = False try: - from flocks.session.callable_state import clear_session_callable_tools + if not await cls._stop_session_tree_for_archive(sessions): + return False + for session in sessions: + await Message.quiesce_parts(session.id, persist=False) - await clear_session_callable_tools(session_id) - except Exception as e: - log.warn("session.callable_tools.clear_error", {"id": session_id, "error": str(e)}) - - log.info("session.deleted", { - "id": session_id, - "project_id": project_id, - }) + refreshed = await cls.collect_tree(project_id, session_id) + if not refreshed: + return False + sessions = refreshed + session_ids = [session.id for session in sessions] + + from flocks.permission.next import PermissionNext + + permission_keys = await PermissionNext.deletion_storage_keys(session_ids) + await Storage.mutate_many( + delete_keys=[ + key + for session in sessions + for key in ( + f"session:{session.project_id}:{session.id}", + f"message:{session.id}", + f"message_parts:{session.id}", + f"todo:{session.id}", + f"goal:{session.id}", + f"session_diff:{session.id}", + f"session_callable_tools:{session.id}", + ) + ] + permission_keys, + delete_prefixes=[ + prefix + for session in sessions + for prefix in ( + f"message_parts:{session.id}:", + f"message_diff:{session.id}:", + f"system_prompts:{session.id}:", + ) + ], + ) + PermissionNext.clear_session_runtime(session_ids) - # Flocks compatibility: clear state and publish event - try: - from flocks.session.core.session_state import get_main_session_id, set_main_session, remove_subagent_session - if get_main_session_id() == session_id: - set_main_session(None) - remove_subagent_session(session_id) - except Exception as e: - log.warn("session.state.error", {"error": str(e)}) + from flocks.session.session_loop import SessionLoop + from flocks.session.callable_state import invalidate_session_callable_tools_cache + from flocks.project.project import Project - try: - from flocks.bus.bus import Bus - from flocks.bus.events import SessionDeleted - await Bus.publish(SessionDeleted, { - "sessionID": session_id, + Project.invalidate_session_stats() + + for session in sessions: + cls._id_index.pop(session.id, None) + cls._remove_from_list_cache(session.id) + SessionLoop.clear_auto_failover_state(session.id) + Message.invalidate_cache(session.id) + invalidate_session_callable_tools_cache(session.id) + + try: + from flocks.session.files import remove_session_uploads + + if await asyncio.to_thread(remove_session_uploads, session.id): + log.info("session.uploads.cleaned", {"session_id": session.id}) + except Exception as exc: + log.warn("session.uploads.cleanup_failed", { + "session_id": session.id, + "error": str(exc), + }) + + for session in reversed(sessions): + try: + from flocks.session.core.session_state import ( + get_main_session_id, + remove_subagent_session, + set_main_session, + ) + + if get_main_session_id() == session.id: + set_main_session(None) + remove_subagent_session(session.id) + except Exception as e: + log.warn("session.state.error", {"id": session.id, "error": str(e)}) + + try: + from flocks.bus.bus import Bus + from flocks.bus.events import SessionDeleted + + await Bus.publish(SessionDeleted, {"sessionID": session.id}) + except Exception as e: + log.warn("session.deleted.event_error", {"id": session.id, "error": str(e)}) + + log.info("session.deleted", { + "id": session_id, + "project_id": project_id, + "count": len(sessions), }) - except Exception as e: - log.warn("session.deleted.event_error", {"error": str(e)}) - - return True + deleted = True + return True + finally: + await cls._end_lifecycle_transition(sessions) + if deleted: + for session in sessions: + cls._lifecycle_generations.pop(session.id, None) @classmethod async def retain_deleted_user_sessions(cls, user_id: str, username: str) -> int: @@ -659,15 +1039,89 @@ async def retain_deleted_user_sessions(cls, user_id: str, username: str) -> int: await cls.update( project_id=session.project_id, session_id=session.id, + allow_inactive=True, owner_user_id=_UNSET, owner_username=username, ) migrated += 1 return migrated + + @classmethod + async def _stop_session_tree_for_archive( + cls, + sessions: List[SessionInfo], + *, + timeout_s: float = 5.0, + clear_prompt_queue: bool = True, + ) -> bool: + """Stop persisted and in-memory work before committing archive state.""" + from flocks.session.interaction_queue import InteractionQueue + from flocks.session.runner import SessionRunner + from flocks.session.session_loop import SessionLoop + + session_ids = [session.id for session in sessions] + for session_id in session_ids: + SessionLoop.abort(session_id) + SessionRunner.cancel(session_id) + if clear_prompt_queue: + await InteractionQueue.clear(session_id) + try: + from flocks.server.routes.question import reject_session_questions + + await reject_session_questions(session_id) + except Exception as exc: + log.warn("session.archive.question_cleanup_error", { + "id": session_id, + "error": str(exc), + }) + try: + from flocks.session.background_tasks import cancel_session_background_tasks + + await cancel_session_background_tasks(session_id) + except Exception as exc: + log.warn("session.archive.background_task_cleanup_error", { + "id": session_id, + "error": str(exc), + }) + try: + from flocks.task.background import get_background_manager + + get_background_manager().cancel_by_parent_session_id(session_id) + except Exception as exc: + log.warn("session.archive.background_cleanup_error", { + "id": session_id, + "error": str(exc), + }) + + deadline = asyncio.get_running_loop().time() + timeout_s + pending = set(session_ids) + while pending: + running = {session_id for session_id in pending if SessionLoop.is_running(session_id)} + if not running: + return True + now = asyncio.get_running_loop().time() + if now >= deadline: + log.warn("session.archive.wait_idle_timeout", { + "ids": sorted(running), + "timeout_s": timeout_s, + }) + return False + pending = running + await asyncio.sleep(min(0.05, deadline - now)) + return True @classmethod async def archive(cls, project_id: str, session_id: str) -> bool: + """Archive a root session and its descendants.""" + + from flocks.project.project import Project + + async with Project.lifecycle_guard(project_id): + return await cls._archive_locked(project_id, session_id) + + @classmethod + async def _archive_locked(cls, project_id: str, session_id: str) -> bool: """ Archive a session @@ -678,23 +1132,111 @@ async def archive(cls, project_id: str, session_id: str) -> bool: Returns: True if archived """ - session = await cls.get(project_id, session_id) - if not session: + sessions = await cls._begin_lifecycle_transition( + project_id, + session_id, + require_root=True, + ) + if not sessions: return False + prompt_queue_paused = False + archive_committed = False + try: + if any(session.status != "archived" for session in sessions): + from flocks.session.interaction_queue import InteractionQueue + + prompt_queue_paused = True + for session in sessions: + await InteractionQueue.pause(session.id) + if not await cls._stop_session_tree_for_archive( + sessions, + clear_prompt_queue=False, + ): + return False + for session in sessions: + await Message.quiesce_parts(session.id, persist=True) + + # Lifecycle claims prevent new in-process children and writes. Read + # once more so every child committed before the claim is included. + refreshed = await cls.collect_tree(project_id, session_id) + if not refreshed: + return False + sessions = refreshed + + archived_ts = int(datetime.now().timestamp() * 1000) + updated_sessions: List[SessionInfo] = [] + for session in sessions: + if session.status == "archived": + updated_sessions.append(session) + continue + updated_sessions.append(session.model_copy(update={ + "status": "archived", + "time": session.time.model_copy(update={ + "archived": session.time.archived or archived_ts, + }), + })) + + original_status = {session.id: session.status for session in sessions} + changed = [ + session + for session in updated_sessions + if original_status[session.id] != session.status + ] + if changed: + await Storage.set_many([ + (f"session:{session.project_id}:{session.id}", session, "session") + for session in changed + ]) + for session in changed: + cls._id_index[session.id] = f"session:{session.project_id}:{session.id}" + cls._sync_list_cache(session) + archive_committed = True - archived_ts = int(datetime.now().timestamp() * 1000) - time_data = session.time.model_dump() - time_data["archived"] = archived_ts - await cls.update(project_id, session_id, status="archived", time=time_data) + try: + from flocks.session.session_loop import SessionLoop + from flocks.session.core.session_state import ( + get_main_session_id, + remove_subagent_session, + set_main_session, + ) + + for session in sessions: + SessionLoop.clear_auto_failover_state(session.id) + Message.invalidate_cache(session.id) + if get_main_session_id() == session.id: + set_main_session(None) + remove_subagent_session(session.id) + except Exception as exc: + log.warn("session.archive.runtime_cleanup_error", { + "id": session_id, + "error": str(exc), + }) + + log.info("session.archived", { + "id": session_id, + "project_id": project_id, + "count": len(sessions), + }) + return True + finally: + if prompt_queue_paused and not archive_committed: + from flocks.session.interaction_queue import InteractionQueue - log.info("session.archived", { - "id": session_id, - "project_id": project_id, - }) - return True + for session in sessions: + await InteractionQueue.resume(session.id) + await cls._end_lifecycle_transition(sessions) @classmethod async def unarchive(cls, project_id: str, session_id: str) -> bool: + """Restore an archived root session and its descendants.""" + + from flocks.project.project import Project + + async with Project.lifecycle_guard(project_id): + return await cls._unarchive_locked(project_id, session_id) + + @classmethod + async def _unarchive_locked(cls, project_id: str, session_id: str) -> bool: """ Restore an archived session @@ -705,27 +1247,174 @@ async def unarchive(cls, project_id: str, session_id: str) -> bool: Returns: True if restored """ - session = await cls.get(project_id, session_id) - if not session: - # get() filters deleted; manually check archived - raw = await Storage.get(f"session:{project_id}:{session_id}", SessionInfo) - if not raw or raw.status != "archived": - return False - session = raw - - if session.status != "archived": + sessions = await cls._begin_lifecycle_transition( + project_id, + session_id, + require_root=True, + ) + if not sessions: return False + try: + changed = [ + session.model_copy(update={ + "status": "active", + "time": session.time.model_copy(update={ + "archived": None, + }), + }) + for session in sessions + if session.status == "archived" + ] + if changed: + await Storage.set_many([ + (f"session:{session.project_id}:{session.id}", session, "session") + for session in changed + ]) + for session in changed: + cls._id_index[session.id] = f"session:{session.project_id}:{session.id}" + cls._sync_list_cache(session) + + from flocks.session.interaction_queue import InteractionQueue + + for session in sessions: + await InteractionQueue.resume(session.id) + + log.info("session.unarchived", { + "id": session_id, + "project_id": project_id, + "count": len(sessions), + }) + return True + finally: + await cls._end_lifecycle_transition(sessions) - time_data = session.time.model_dump() - time_data["archived"] = None - await cls.update(project_id, session_id, status="active", time=time_data) - - log.info("session.unarchived", { - "id": session_id, - "project_id": project_id, - }) - - return True + @classmethod + async def restore( + cls, + project_id: str, + session_id: str, + *, + project_owner_id: Optional[str], + ) -> bool: + """Restore a session tree and its removed project as one operation.""" + + from flocks.project.project import ( + DEFAULT_PROJECT_ID, + Project, + ProjectDeletionError, + TASK_SESSION_GROUP_ID, + ) + + async with Project.lifecycle_guard(project_id): + project_state = ( + Project.registry_state(project_id, owner_id=project_owner_id) + if project_owner_id + else ( + "virtual" + if project_id in {DEFAULT_PROJECT_ID, TASK_SESSION_GROUP_ID} + else "missing" + ) + ) + if project_state == "missing" and project_id.startswith("prj_"): + raise ProjectDeletionError( + f"Project {project_id} restoration metadata is unavailable" + ) + + project_was_removed = project_state == "removed" + if project_state in {"active", "removed"} and project_owner_id is not None: + await Project.restore(project_id, owner_id=project_owner_id) + + try: + restored = await cls.unarchive(project_id, session_id) + except BaseException: + if project_was_removed and project_owner_id is not None: + try: + await asyncio.shield( + Project._delete_locked(project_id, owner_id=project_owner_id) + ) + except Exception as rollback_exc: + log.error("session.restore.project_rollback_failed", { + "project_id": project_id, + "error": str(rollback_exc), + }) + raise + + if not restored and project_was_removed and project_owner_id is not None: + await Project.delete(project_id, owner_id=project_owner_id) + return restored + + @classmethod + async def collect_tree(cls, project_id: str, session_id: str) -> List[SessionInfo]: + """Return a root session and all descendants in parent-first order.""" + token = set_current_auth_user(None) + try: + root = await cls.get(project_id, session_id) + finally: + reset_current_auth_user(token) + if root is None: + return [] + + children_by_parent: Dict[str, List[SessionInfo]] = {} + for session in await cls.list_all_unfiltered(): + if session.project_id != project_id: + continue + if session.parent_id: + children_by_parent.setdefault(session.parent_id, []).append(session) + + tree: List[SessionInfo] = [] + seen: set[str] = set() + + def visit(session: SessionInfo) -> None: + if session.id in seen: + return + seen.add(session.id) + tree.append(session) + for child in children_by_parent.get(session.id, []): + visit(child) + + visit(root) + return tree + + @classmethod + async def _begin_lifecycle_transition( + cls, + project_id: str, + session_id: str, + *, + require_root: bool = False, + ) -> List[SessionInfo]: + """Claim a complete tree so ordinary writes cannot overtake its transition.""" + async with cls._tree_lock: + sessions = await cls.collect_tree(project_id, session_id) + if not sessions: + return [] + if require_root and sessions[0].parent_id is not None: + log.warn("session.lifecycle.non_root_rejected", {"id": session_id}) + return [] + + async with AsyncExitStack() as locks: + for item_id in sorted(session.id for session in sessions): + await locks.enter_async_context(cls.lifecycle_lock(item_id)) + + # A child may have committed before the topology lock was + # acquired. Re-read under all known keyed locks before claiming. + sessions = await cls.collect_tree(project_id, session_id) + ids = {session.id for session in sessions} + if ( + not sessions + or ids & cls._lifecycle_transition_ids + or any(cls.has_active_operations(item_id) for item_id in ids) + ): + return [] + cls._lifecycle_transition_ids.update(ids) + for item_id in ids: + cls._lifecycle_generations[item_id] = cls.lifecycle_generation(item_id) + 1 + return sessions + + @classmethod + async def _end_lifecycle_transition(cls, sessions: List[SessionInfo]) -> None: + async with cls._tree_lock: + cls._lifecycle_transition_ids.difference_update(session.id for session in sessions) @classmethod async def children(cls, project_id: str, parent_id: str) -> List[SessionInfo]: @@ -760,70 +1449,92 @@ async def fork( Returns: New forked session """ - # Get original session - original = await cls.get(project_id, session_id) - if not original: - raise ValueError(f"Session {session_id} not found") - - # Create new session with parent_id set - new_session = await cls.create( - project_id=project_id, - directory=original.directory, - ) - - # Update parent_id after creation - new_session = await cls.update( - project_id=project_id, - session_id=new_session.id, - parent_id=session_id, - ) - - # Copy messages with all parts (include archived so fork preserves full history) - messages = await Message.list(session_id, include_archived=True) - id_map: Dict[str, str] = {} - - for msg in messages: - if message_id and msg.id >= message_id: - break - - new_id = Identifier.ascending("message") - id_map[msg.id] = new_id - - # Get text content for the initial message creation - content = await Message.get_text_content(msg) - parent_ref = None - if isinstance(msg, AssistantMessageInfo): - parent_ref = id_map.get(msg.parentID) - - # Create the message (this also creates an initial TextPart) - await Message.create( - session_id=new_session.id, - role=msg.role, - content=content, - id=new_id, - parentID=parent_ref or "", + from flocks.project.project import Project + + # Keep the parent, new child, and copied history in one project-lifecycle + # interval. A concurrent move can only happen before or after the fork. + async with Project.lifecycle_guard(project_id): + original = await cls.get(project_id, session_id) + if not original: + raise ValueError(f"Session {session_id} not found") + + messages = await Message.list(session_id, include_archived=True) + messages_to_copy: List[MessageInfo] = [] + for msg in messages: + if message_id and msg.id >= message_id: + break + messages_to_copy.append(msg) + id_map = { + msg.id: Identifier.ascending("message") + for msg in messages_to_copy + } + + fork_move_metadata: Optional[Dict[str, Any]] = None + move_metadata = original.metadata.get("projectMove") + if isinstance(move_metadata, dict) and move_metadata.get("boundaryMessageID"): + original_boundary = move_metadata["boundaryMessageID"] + mapped_boundary = id_map.get(original_boundary) + if mapped_boundary is None and messages_to_copy: + # A fork ending before the move boundary contains only unsafe + # imported history, so protect the complete copied prefix. + mapped_boundary = id_map[messages_to_copy[-1].id] + if mapped_boundary is not None: + fork_move_metadata = { + **move_metadata, + "boundaryMessageID": mapped_boundary, + } + + metadata = ( + {"projectMove": fork_move_metadata} + if fork_move_metadata is not None + else None ) - - # Copy non-text parts (tool calls, files, patches, etc.) - original_parts = await Message.parts(msg.id, session_id) - for part in original_parts: - if part.type == "text": - continue # Already created by Message.create above - # Clone part with updated session/message IDs - part_data = part.model_dump() - part_data["id"] = Identifier.ascending("part") - part_data["sessionID"] = new_session.id - part_data["messageID"] = new_id - cloned_part = Message.deserialize_part(part_data) - await Message.store_part(new_session.id, new_id, cloned_part) - - log.info("session.forked", { - "from": session_id, - "to": new_session.id, - "messages": len(id_map), - }) - - return new_session + new_session = await cls.create( + project_id=project_id, + directory=original.directory, + parent_id=session_id, + **({"metadata": metadata} if metadata is not None else {}), + ) + + # Copy messages with all parts (include archived so fork preserves full history) + for msg in messages_to_copy: + new_id = id_map[msg.id] + + # Get text content for the initial message creation + content = await Message.get_text_content(msg) + parent_ref = None + if isinstance(msg, AssistantMessageInfo): + parent_ref = id_map.get(msg.parentID) + + # Create the message (this also creates an initial TextPart) + await Message.create( + session_id=new_session.id, + role=MessageRole(msg.role), + content=content, + id=new_id, + parentID=parent_ref or "", + ) + + # Copy non-text parts (tool calls, files, patches, etc.) + original_parts = await Message.parts(msg.id, session_id) + for part in original_parts: + if part.type == "text": + continue # Already created by Message.create above + # Clone part with updated session/message IDs + part_data = part.model_dump() + part_data["id"] = Identifier.ascending("part") + part_data["sessionID"] = new_session.id + part_data["messageID"] = new_id + cloned_part = Message.deserialize_part(part_data) + await Message.store_part(new_session.id, new_id, cloned_part) + + log.info("session.forked", { + "from": session_id, + "to": new_session.id, + "messages": len(id_map), + }) + + return new_session @classmethod async def set_revert( diff --git a/flocks/session/session_loop.py b/flocks/session/session_loop.py index adb1e112e..d163d663b 100644 --- a/flocks/session/session_loop.py +++ b/flocks/session/session_loop.py @@ -13,6 +13,7 @@ """ import asyncio +import hashlib import inspect import time from typing import Optional, List, Dict, Any, Callable, Awaitable @@ -21,7 +22,11 @@ from flocks.utils.log import Log from flocks.utils.id import Identifier -from flocks.session.session import Session, SessionInfo +from flocks.session.session import ( + Session, + SessionInfo, + is_model_auto_session_category, +) from flocks.session.message import Message, MessageInfo, MessageRole from flocks.session.core.status import SessionStatus, SessionStatusBusy, SessionStatusIdle from flocks.session.core.task_utils import fire_and_forget @@ -47,6 +52,26 @@ MAX_OVERFLOW_COMPACTION_ATTEMPTS = 3 POST_COMPACTION_COOLDOWN_STEPS = 2 +RATE_LIMIT_COOLDOWN_SECONDS = 60.0 +CHAIN_EXHAUSTION_COOLDOWN_SECONDS = 5.0 + + +@dataclass(frozen=True) +class RuntimeModel: + """Concrete provider/model candidate used by Auto failover.""" + + provider_id: str + model_id: str + + +@dataclass +class AutoFailoverCooldown: + """Process-local Hermes-style starting candidate cooldown.""" + + model: RuntimeModel + primary: RuntimeModel + expires_at: float + reason: str @dataclass @@ -81,6 +106,16 @@ class LoopContext: # is what the upstream will actually bill us for on the next turn # (matches the "observed value wins" rule from docs/design/context-compaction-v2.md §B3). last_observed_prompt_tokens: int = 0 + auto_failover: bool = False + # Entrypoint authorization is separate from persisted model_auto. Only a + # WebUI message route may set this bit; non-WebUI entrypoints use the default. + auto_failover_allowed: bool = False + model_candidates: List[RuntimeModel] = field(default_factory=list) + candidate_index: int = 0 + turn_user_id: Optional[str] = None + turn_additional_context: Optional[str] = None + stop_hook_active: bool = False + session_start_pending: bool = False @property def trace_step(self) -> int: @@ -120,6 +155,8 @@ class LoopResult: action: str # "stop", "continue", "compact", "error", "queued" last_message: Optional[MessageInfo] = None error: Optional[str] = None + provider_id: Optional[str] = None + model_id: Optional[str] = None metadata: Dict[str, Any] = field(default_factory=dict) @@ -137,6 +174,209 @@ class SessionLoop: # Active loop contexts by session ID _active_loops: Dict[str, LoopContext] = {} + _auto_failover_cooldowns: Dict[str, AutoFailoverCooldown] = {} + + @classmethod + def clear_auto_failover_state(cls, session_id: str) -> None: + """Clear process-local routing state when WebUI Auto is disabled.""" + cls._auto_failover_cooldowns.pop(session_id, None) + + @classmethod + async def validate_runtime_model( + cls, + provider_id: str, + model_id: str, + *, + config: Optional[Any] = None, + ) -> tuple[bool, str]: + """Validate a configured LLM candidate without a network health probe.""" + from flocks.config.config import Config + from flocks.provider.model_manager import get_model_manager + from flocks.provider.types import ModelType + + Provider._ensure_initialized() + config = config or await Config.get() + if provider_id in (getattr(config, "disabled_providers", None) or []): + return False, "provider_disabled" + enabled_providers = getattr(config, "enabled_providers", None) or [] + if enabled_providers and provider_id not in enabled_providers: + return False, "provider_disabled" + try: + await Provider.apply_config(config, provider_id=provider_id) + except Exception as exc: + log.warn("session.model.candidate_config_failed", { + "provider_id": provider_id, + "model_id": model_id, + "error": str(exc), + }) + return False, "provider_config_error" + + provider = Provider.get(provider_id) + if provider is None: + return False, "provider_not_found" + + definition = get_model_manager().get_model(provider_id, model_id) + if definition is None: + return False, "model_not_found" + if getattr(definition, "model_type", None) != ModelType.LLM: + return False, "not_llm" + + setting = get_model_manager().get_setting(provider_id, model_id) + if setting is not None and not setting.enabled: + return False, "model_disabled" + if not provider.is_configured(): + return False, "provider_not_configured" + return True, "available" + + @classmethod + async def _build_model_candidates( + cls, + primary: RuntimeModel, + *, + route_seed: str, + preferred: Optional[RuntimeModel] = None, + ) -> List[RuntimeModel]: + """Build a stable per-turn primary, same-provider, cross-provider chain.""" + from flocks.config.config import Config + from flocks.provider.model_manager import get_model_manager + from flocks.provider.types import ModelType + + config = await Config.get() + await Provider.apply_config(config) + definitions = get_model_manager().list_models( + model_type=ModelType.LLM, + enabled_only=True, + ) + discovered = { + RuntimeModel(definition.provider_id, definition.id) + for definition in definitions + } + discovered.discard(primary) + + same_provider: List[RuntimeModel] = [] + other_providers: List[RuntimeModel] = [] + for candidate in sorted( + discovered, + key=lambda item: (item.provider_id, item.model_id), + ): + available, reason = await cls.validate_runtime_model( + candidate.provider_id, + candidate.model_id, + config=config, + ) + if not available: + log.debug("session.model.fallback_skipped", { + "provider_id": candidate.provider_id, + "model_id": candidate.model_id, + "reason": reason, + }) + continue + + if candidate.provider_id == primary.provider_id: + same_provider.append(candidate) + else: + other_providers.append(candidate) + + candidates = [primary] + for tier, pool in ( + ("same_provider", same_provider), + ("other_provider", other_providers), + ): + if not pool: + continue + selected = ( + preferred + if preferred is not None and preferred in pool + else cls._stable_candidate_choice(pool, route_seed, tier) + ) + candidates.append(selected) + return candidates + + @staticmethod + def _stable_candidate_choice( + candidates: List[RuntimeModel], + route_seed: str, + tier: str, + ) -> RuntimeModel: + """Choose pseudo-randomly without Python's process-randomized hash().""" + ordered = sorted( + candidates, + key=lambda item: (item.provider_id, item.model_id), + ) + digest = hashlib.sha256( + f"{route_seed}\0{tier}".encode("utf-8") + ).digest() + index = int.from_bytes(digest[:8], "big") % len(ordered) + return ordered[index] + + @classmethod + async def validate_auto_configuration(cls) -> tuple[bool, str]: + """Validate that a newly selected Auto mode has a usable chain.""" + from flocks.config.config import Config + + default_llm = await Config.resolve_default_llm() + if not default_llm: + return False, "default_model_missing" + primary = RuntimeModel( + default_llm["provider_id"], + default_llm["model_id"], + ) + available, reason = await cls.validate_runtime_model( + primary.provider_id, + primary.model_id, + ) + if not available: + return False, f"primary_{reason}" + return True, "available" + + @classmethod + def _active_cooldown_model( + cls, + session_id: str, + primary: RuntimeModel, + ) -> Optional[RuntimeModel]: + """Return a still-valid cooldown target for the current primary.""" + cooldown = cls._auto_failover_cooldowns.get(session_id) + if cooldown is None: + return None + if cooldown.expires_at <= time.monotonic() or cooldown.primary != primary: + cls._auto_failover_cooldowns.pop(session_id, None) + return None + return cooldown.model + + @classmethod + def _cooldown_candidate_index( + cls, + session_id: str, + candidates: List[RuntimeModel], + ) -> int: + if not candidates: + return 0 + cooldown_model = cls._active_cooldown_model(session_id, candidates[0]) + if cooldown_model is None: + return 0 + try: + return candidates.index(cooldown_model) + except ValueError: + cls._auto_failover_cooldowns.pop(session_id, None) + return 0 + + @classmethod + def _select_candidate(cls, ctx: LoopContext, index: int) -> None: + candidate = ctx.model_candidates[index] + ctx.candidate_index = index + ctx.provider_id = candidate.provider_id + ctx.model_id = candidate.model_id + ctx.session.provider = candidate.provider_id + ctx.session.model = candidate.model_id + # Prompt and model-capability caches are keyed in most places, but a + # fresh dict makes the runtime rebuild guarantee explicit. The tool + # loop guard is turn state rather than model state, so it must survive + # a provider switch to keep repeated-tool protection effective. + tool_loop_guard = ctx.runner_static_cache.get("tool_loop_guard") + ctx.runner_static_cache.clear() + if tool_loop_guard is not None: + ctx.runner_static_cache["tool_loop_guard"] = tool_loop_guard @classmethod def is_running(cls, session_id: str) -> bool: @@ -268,7 +508,7 @@ async def _detect_queued_user_message( _session_id: str, post_messages: List[MessageInfo], current_user_id: str, - last_message: Optional[MessageInfo], + _last_message: Optional[MessageInfo], ) -> Optional[MessageInfo]: if not post_messages: return None @@ -283,11 +523,11 @@ async def _detect_queued_user_message( return None if newest_user.id <= current_user_id: return None - if last_message is None: - return newest_user - if newest_user.id > last_message.id: - return newest_user - return None + # A fallback assistant is created after a user message that arrived + # while the primary model was running. Its newer ID must not make that + # user message look handled; the current turn's user ID is the stable + # boundary for queued work. + return newest_user @classmethod async def run( @@ -298,6 +538,7 @@ async def run( agent_name: Optional[str] = None, callbacks: Optional[LoopCallbacks] = None, working_directory: Optional[str] = None, + auto_failover: bool = False, ) -> LoopResult: """ Run session loop @@ -327,6 +568,15 @@ async def run( # next iteration once it finishes the current step. if cls.is_running(session_id): log.info("loop.already_running", {"session_id": session_id}) + if auto_failover: + active_ctx = cls._active_loops.get(session_id) + if ( + active_ctx is not None + and is_model_auto_session_category( + getattr(active_ctx.session, "category", "user") + ) + ): + active_ctx.auto_failover_allowed = True return LoopResult( action="queued", error="Loop already running", @@ -340,6 +590,15 @@ async def run( action="error", error=f"Session {session_id} not found", ) + if session.status != "active": + log.warning("loop.session_not_active", { + "session_id": session_id, + "status": session.status, + }) + return LoopResult( + action="error", + error=f"Session {session_id} is {session.status}", + ) if working_directory: session = session.model_copy(update={"directory": working_directory}) @@ -351,6 +610,19 @@ async def run( provider_id = provider_id or resolved_provider model_id = model_id or resolved_model + primary_model = RuntimeModel( + provider_id=provider_id, + model_id=model_id, + ) + model_candidates = [primary_model] + candidate_index = 0 + auto_failover = bool( + auto_failover + and is_model_auto_session_category( + getattr(session, "category", "user") + ) + ) + # Keep the in-memory session aligned with the runtime model so # downstream helpers (title generation, compaction checks, etc.) see # the model actually selected for this loop iteration. Unpinned @@ -382,10 +654,46 @@ async def run( agent_name=agent_name or session.agent or "rex", session_ctx=session_ctx, trace_step_offset=trace_offset, + auto_failover=auto_failover, + auto_failover_allowed=auto_failover, + model_candidates=model_candidates, + candidate_index=candidate_index, + session_start_pending=trace_offset == 0, ) - # Register context - cls._active_loops[session_id] = ctx + # Register under the same lock used by archive/delete. This closes the + # gap where archival could commit after the status check above but + # before the loop became visible to the lifecycle stop logic. + async with Session.lifecycle_lock(session_id): + latest_session = await Session.get_by_id(session_id) + if latest_session is None: + log.warning("loop.session_not_found_before_register", { + "session_id": session_id, + }) + return LoopResult( + action="error", + error=f"Session {session_id} not found", + ) + if latest_session.status != "active": + log.warning("loop.session_not_active_before_register", { + "session_id": session_id, + "status": latest_session.status, + }) + return LoopResult( + action="error", + error=f"Session {session_id} is {latest_session.status}", + ) + if Session.is_lifecycle_transitioning(session_id): + return LoopResult( + action="error", + error=f"Session {session_id} is changing lifecycle state", + ) + if cls.is_running(session_id): + return LoopResult( + action="queued", + error="Loop already running", + ) + cls._active_loops[session_id] = ctx # Set status to busy SessionStatus.set(session_id, SessionStatusBusy()) @@ -424,7 +732,12 @@ async def run( }) except Exception as exc: log.warn("loop.error.event_error", {"error": str(exc)}) - return LoopResult(action="error", error=str(e)) + return LoopResult( + action="error", + error=str(e), + provider_id=ctx.provider_id, + model_id=ctx.model_id, + ) finally: # Clean up if session_id in cls._active_loops: @@ -562,6 +875,484 @@ async def _resolve_model( return resolved_provider, resolved_model, source return resolved_provider, resolved_model + @classmethod + async def _prepare_auto_turn( + cls, + ctx: LoopContext, + last_user: MessageInfo, + ) -> bool: + """Synchronize routing when the loop advances to a real WebUI turn. + + Returns: + True when ``last_user`` starts a new non-synthetic user turn. + """ + if last_user.id == ctx.turn_user_id: + return False + + parts = await Message.parts(last_user.id, ctx.session.id) + if any(bool(getattr(part, "synthetic", False)) for part in parts): + return False + + if ctx.turn_user_id is None: + ctx.turn_user_id = last_user.id + if ctx.auto_failover and ctx.auto_failover_allowed: + primary = ctx.model_candidates[0] + preferred = cls._active_cooldown_model( + ctx.session.id, + primary, + ) + ctx.model_candidates = await cls._build_model_candidates( + primary, + route_seed=f"{ctx.session.id}:{last_user.id}", + preferred=preferred, + ) + next_index = cls._cooldown_candidate_index( + ctx.session.id, + ctx.model_candidates, + ) + cls._select_candidate(ctx, next_index) + return True + + ctx.turn_user_id = last_user.id + persisted_session = await Session.get_by_id(ctx.session.id) + persisted_model_auto = bool( + persisted_session + and is_model_auto_session_category( + getattr(persisted_session, "category", "user") + ) + and getattr(persisted_session, "model_auto", False) + ) + persisted_auto = persisted_model_auto and ctx.auto_failover_allowed + + user_model = getattr(last_user, "model", None) + user_provider_id = None + user_model_id = None + if isinstance(user_model, dict): + user_provider_id = user_model.get("providerID") or user_model.get("provider_id") + user_model_id = user_model.get("modelID") or user_model.get("model_id") + + if not persisted_auto: + ctx.auto_failover = False + if not persisted_model_auto: + cls.clear_auto_failover_state(ctx.session.id) + ctx.auto_failover_allowed = False + provider_id = ( + getattr(persisted_session, "provider", None) + if Session.has_pinned_model(persisted_session) + else user_provider_id + ) or ctx.provider_id + model_id = ( + getattr(persisted_session, "model", None) + if Session.has_pinned_model(persisted_session) + else user_model_id + ) or ctx.model_id + ctx.model_candidates = [RuntimeModel(provider_id, model_id)] + cls._select_candidate(ctx, 0) + log.info("session.model.auto_disabled_for_turn", { + "session_id": ctx.session.id, + "provider_id": provider_id, + "model_id": model_id, + }) + return True + + from flocks.config.config import Config + + previous = RuntimeModel(ctx.provider_id, ctx.model_id) + default_llm = await Config.resolve_default_llm() + primary = RuntimeModel( + provider_id=(default_llm or {}).get("provider_id") or user_provider_id or ctx.provider_id, + model_id=(default_llm or {}).get("model_id") or user_model_id or ctx.model_id, + ) + # Rebuild once for every real turn. The user message ID makes the + # pseudo-random choices stable throughout that turn, while an active + # cooldown keeps its valid target in the newly sampled tier. + preferred = cls._active_cooldown_model(ctx.session.id, primary) + ctx.model_candidates = await cls._build_model_candidates( + primary, + route_seed=f"{ctx.session.id}:{last_user.id}", + preferred=preferred, + ) + ctx.auto_failover = True + next_index = cls._cooldown_candidate_index( + ctx.session.id, + ctx.model_candidates, + ) + cls._select_candidate(ctx, next_index) + active = ctx.model_candidates[next_index] + log.info("session.model.auto_turn_reset", { + "session_id": ctx.session.id, + "from_provider_id": previous.provider_id, + "from_model_id": previous.model_id, + "to_provider_id": active.provider_id, + "to_model_id": active.model_id, + "cooldown_active": next_index > 0, + }) + return True + + @classmethod + async def _run_user_prompt_submit_hook( + cls, + ctx: LoopContext, + last_user: MessageInfo, + ) -> None: + """Run UserPromptSubmit once for a newly observed real user turn.""" + try: + from flocks.hooks.pipeline import HookPipeline + + prompt = await Message.get_text_content(last_user) + hook_ctx = await HookPipeline.run_user_prompt_submit({ + "sessionID": ctx.session.id, + "workspace": ctx.session.directory, + "agent": getattr(last_user, "agent", None) or ctx.agent_name, + "model": { + "providerID": ctx.provider_id, + "modelID": ctx.model_id, + }, + "messageID": last_user.id, + "prompt": prompt, + }) + additional_context = hook_ctx.output.get("additionalContext") + if isinstance(additional_context, str) and additional_context.strip(): + ctx.turn_additional_context = additional_context.strip() + except Exception as exc: + log.debug("loop.hook.user_prompt_submit.error", { + "session_id": ctx.session.id, + "message_id": last_user.id, + "error": str(exc), + }) + + @classmethod + async def _run_turn_finish_hook( + cls, + ctx: LoopContext, + callbacks: LoopCallbacks, + last_user: MessageInfo, + last_message: MessageInfo, + ) -> bool: + """Run TurnFinish and continue the loop when the hook blocks stopping.""" + try: + from flocks.hooks.pipeline import HookPipeline + + hook_user = last_user + if ctx.turn_user_id: + hook_user = ( + await Message.get(ctx.session.id, ctx.turn_user_id) + or last_user + ) + user_text = await Message.get_text_content(hook_user) + assistant_text = await Message.get_text_content(last_message) + hook_ctx = await HookPipeline.run_turn_finish({ + "sessionID": ctx.session.id, + "workspace": ctx.session.directory, + "agent": getattr(last_message, "agent", None) or ctx.agent_name, + "model": { + "providerID": ctx.provider_id, + "modelID": ctx.model_id, + }, + "step": ctx.trace_step, + "userMessage": { + "id": hook_user.id, + "content": user_text, + }, + "assistantMessage": { + "id": last_message.id, + "content": assistant_text, + }, + "finishReason": "stop", + "stopHookActive": ctx.stop_hook_active, + }) + except Exception as exc: + log.debug("loop.hook.turn_finish.error", { + "session_id": ctx.session.id, + "message_id": getattr(last_message, "id", None), + "error": str(exc), + }) + return False + + decision = str(hook_ctx.output.get("decision") or "").strip().lower() + reason = str(hook_ctx.output.get("reason") or "").strip() + if decision != "block": + return False + if not reason: + log.warn("loop.hook.turn_finish.missing_reason", { + "session_id": ctx.session.id, + "message_id": last_message.id, + }) + return False + if ctx.should_abort(): + log.info("loop.hook.turn_finish.ignored_after_abort", { + "session_id": ctx.session.id, + "message_id": last_message.id, + }) + return False + + try: + if ctx.session_ctx: + post_hook_messages = await ctx.session_ctx.get_messages() + else: + post_hook_messages = await Message.list(ctx.session.id) + queued_user = await cls._detect_queued_user_message( + ctx.session.id, + post_hook_messages, + last_user.id, + last_message, + ) + except Exception as exc: + queued_user = None + log.debug("loop.hook.turn_finish.queued_recheck_error", { + "session_id": ctx.session.id, + "error": str(exc), + }) + if queued_user is not None: + turn_state = set_turn_state( + ctx.session.id, + step=ctx.step, + status="continued", + continue_reason="queued_message", + queued_message_detected=True, + ) + await cls._publish_runtime_event(callbacks, "turn.continued", { + **turn_state.model_dump(by_alias=True), + "queuedUserMessageID": queued_user.id, + }) + log.info("loop.hook.turn_finish.queued_message_won", { + "session_id": ctx.session.id, + "queued_user_id": queued_user.id, + "source_assistant_message_id": last_message.id, + }) + return True + + from flocks.agent.registry import Agent + from flocks.session.core.defaults import DEFAULT_MAX_TOOL_STEPS + + try: + agent = await Agent.get( + getattr(last_message, "agent", None) or ctx.agent_name + ) + except Exception as exc: + log.debug("loop.hook.turn_finish.agent_load_error", { + "session_id": ctx.session.id, + "error": str(exc), + }) + agent = None + max_steps = ( + agent.steps + if agent is not None and getattr(agent, "steps", None) is not None + else DEFAULT_MAX_TOOL_STEPS + ) + if ctx.trace_step >= max_steps: + log.warn("loop.hook.turn_finish.step_limit", { + "session_id": ctx.session.id, + "step": ctx.trace_step, + "max_steps": max_steps, + }) + return False + + try: + continuation = await Message.create( + session_id=ctx.session.id, + role=MessageRole.USER, + content=reason, + agent=getattr(hook_user, "agent", None) or ctx.agent_name, + model={ + "providerID": ctx.provider_id, + "modelID": ctx.model_id, + }, + synthetic=True, + part_metadata={ + "turnFinishContinuation": True, + "stopHookActive": True, + "sourceAssistantMessageID": last_message.id, + }, + ) + except Exception as exc: + log.error("loop.hook.turn_finish.continuation_error", { + "session_id": ctx.session.id, + "error": str(exc), + }) + return False + ctx.stop_hook_active = True + turn_state = set_turn_state( + ctx.session.id, + step=ctx.step, + status="continued", + continue_reason="turn_finish_hook", + queued_message_detected=False, + ) + await cls._publish_runtime_event(callbacks, "turn.continued", { + **turn_state.model_dump(by_alias=True), + "turnFinishMessageID": continuation.id, + }) + log.info("loop.continuing_for_turn_finish_hook", { + "session_id": ctx.session.id, + "continuation_message_id": continuation.id, + "source_assistant_message_id": last_message.id, + }) + return True + + @classmethod + async def _finalize_deferred_failure( + cls, + ctx: LoopContext, + failure: Any, + last_user: MessageInfo, + ) -> None: + """Persist only the final Auto candidate failure.""" + if not failure.assistant_message_id: + assistant = await Message.create( + session_id=ctx.session.id, + role=MessageRole.ASSISTANT, + content="", + agent=getattr(last_user, "agent", None) or ctx.agent_name or "rex", + model_id=ctx.model_id, + provider_id=ctx.provider_id, + parent_id=last_user.id, + error=failure.error_data, + finish="error", + ) + failure.assistant_message_id = assistant.id + return + await Message.update( + ctx.session.id, + failure.assistant_message_id, + error=failure.error_data, + finish="error", + ) + + @classmethod + async def _process_step_with_failover( + cls, + ctx: LoopContext, + callbacks: LoopCallbacks, + messages: List[MessageInfo], + last_user: MessageInfo, + ) -> Any: + """Run one logical step, moving across candidates without replaying output.""" + from flocks.session.runner import RunnerCallbacks, SessionRunner + + while True: + runner_cbs = callbacks.runner_callbacks + if runner_cbs is None: + runner_cbs = RunnerCallbacks() + if callbacks.event_publish_callback and not runner_cbs.event_publish_callback: + runner_cbs.event_publish_callback = callbacks.event_publish_callback + + runner = SessionRunner( + session=ctx.session, + provider_id=ctx.provider_id, + model_id=ctx.model_id, + agent_name=ctx.agent_name, + abort_event=ctx.abort_event, + callbacks=runner_cbs, + session_ctx=ctx.session_ctx, + memory_bootstrap_data=ctx.memory_bootstrap_data, + static_cache=ctx.runner_static_cache, + defer_step_errors=ctx.auto_failover, + failover_available=( + ctx.auto_failover + and ctx.candidate_index + 1 < len(ctx.model_candidates) + ), + turn_additional_context=ctx.turn_additional_context, + session_start_pending=ctx.session_start_pending, + ) + runner._step = ctx.trace_step + + step_result = await runner._process_step(messages, last_user) + if runner._session_start_fired: + ctx.session_start_pending = False + failure = step_result.failure + if not ctx.auto_failover or failure is None: + return step_result + + next_index = ctx.candidate_index + 1 + has_next = next_index < len(ctx.model_candidates) + if not failure.allow_fallback or not has_next: + if ( + failure.allow_fallback + and not has_next + and ctx.candidate_index > 0 + and failure.reason not in {"rate_limit", "billing"} + ): + expires_at = time.monotonic() + CHAIN_EXHAUSTION_COOLDOWN_SECONDS + existing_cooldown = cls._auto_failover_cooldowns.get(ctx.session.id) + if not ( + existing_cooldown + and existing_cooldown.expires_at > expires_at + ): + cls._auto_failover_cooldowns[ctx.session.id] = AutoFailoverCooldown( + model=ctx.model_candidates[ctx.candidate_index], + primary=ctx.model_candidates[0], + expires_at=expires_at, + reason="chain_exhausted", + ) + await cls._finalize_deferred_failure(ctx, failure, last_user) + return step_result + + # A candidate may be removed only while its attempt is completely + # replay-safe. Failure to delete stops the switch to avoid leaving + # two assistant cards for one logical response. + if failure.assistant_message_id: + try: + deleted = await Message.delete( + ctx.session.id, + failure.assistant_message_id, + ) + except Exception as exc: + deleted = False + log.error("session.model.fallback_cleanup_failed", { + "session_id": ctx.session.id, + "message_id": failure.assistant_message_id, + "error": str(exc), + }) + if not deleted: + await cls._finalize_deferred_failure(ctx, failure, last_user) + return step_result + await cls._publish_runtime_event(callbacks, "message.removed", { + "sessionID": ctx.session.id, + "messageID": failure.assistant_message_id, + }) + + previous = ctx.model_candidates[ctx.candidate_index] + next_candidate = ctx.model_candidates[next_index] + + if ctx.candidate_index == 0 and failure.reason in {"rate_limit", "billing"}: + cls._auto_failover_cooldowns[ctx.session.id] = AutoFailoverCooldown( + model=next_candidate, + primary=ctx.model_candidates[0], + expires_at=time.monotonic() + RATE_LIMIT_COOLDOWN_SECONDS, + reason=failure.reason, + ) + else: + cooldown = cls._auto_failover_cooldowns.get(ctx.session.id) + if cooldown and cooldown.expires_at > time.monotonic(): + cooldown.model = next_candidate + + cls._select_candidate(ctx, next_index) + event_payload = { + "sessionID": ctx.session.id, + "from": { + "providerID": previous.provider_id, + "modelID": previous.model_id, + }, + "to": { + "providerID": next_candidate.provider_id, + "modelID": next_candidate.model_id, + }, + "reason": failure.reason, + "candidateIndex": next_index, + } + log.warn("session.model.fallback", { + "from": event_payload["from"], + "to": event_payload["to"], + "reason": event_payload["reason"], + "candidateIndex": event_payload["candidateIndex"], + }) + await cls._publish_runtime_event( + callbacks, + "session.model.fallback", + event_payload, + ) + @classmethod async def _run_loop( cls, @@ -581,6 +1372,7 @@ async def _run_loop( 7. Loop until complete """ last_message: Optional[MessageInfo] = None + loop_error: Optional[str] = None while not ctx.should_abort(): # Set status to busy @@ -678,7 +1470,7 @@ async def _run_loop( stop_reason="no_user_message", ) break - + last_assistant_parts = ( await Message.parts(last_assistant.id, ctx.session.id) if last_assistant @@ -699,6 +1491,11 @@ async def _run_loop( }) last_message = last_assistant break + + if await cls._prepare_auto_turn(ctx, last_user): + ctx.turn_additional_context = None + ctx.stop_hook_active = False + await cls._run_user_prompt_submit_hook(ctx, last_user) # Bootstrap memory on first step (once per loop, stored in ctx) if ctx.step == 1 and ctx.session.memory_enabled and ctx.memory_bootstrap_data is None: @@ -720,7 +1517,7 @@ async def _run_loop( # may cancel this task before it finishes). # generate_title_after_first_message is idempotent: if this task saves # the title first, the safety-net call returns immediately. - if ctx.step == 1: + if ctx.step == 1 and not ctx.auto_failover: try: from flocks.session.lifecycle.title import SessionTitle # UserMessageInfo.model is Dict[str, str] {"providerID": ..., "modelID": ...} @@ -1207,34 +2004,16 @@ async def progress_callback_overflow(stage: str, data: dict) -> None: except Exception as e: log.error("loop.compaction_overflow_check_error", {"error": str(e)}) - # Process step - delegate to runner (matching TUI SessionProcessor.process) - from flocks.session.runner import SessionRunner, RunnerCallbacks - - # Build runner callbacks from loop callbacks - runner_cbs = callbacks.runner_callbacks - if runner_cbs is None: - runner_cbs = RunnerCallbacks() - # Ensure event_publish_callback is propagated - if callbacks.event_publish_callback and not runner_cbs.event_publish_callback: - runner_cbs.event_publish_callback = callbacks.event_publish_callback - - runner = SessionRunner( - session=ctx.session, - provider_id=ctx.provider_id, - model_id=ctx.model_id, - agent_name=ctx.agent_name, - abort_event=ctx.abort_event, - callbacks=runner_cbs, - session_ctx=ctx.session_ctx, - memory_bootstrap_data=ctx.memory_bootstrap_data, - static_cache=ctx.runner_static_cache, - ) - # Use session-cumulative step number for observability. - runner._step = ctx.trace_step - # Process single step — wrap in a Task so abort() can cancel it immediately # rather than waiting for the current tool call to finish. - step_task = asyncio.create_task(runner._process_step(messages, last_user)) + step_task = asyncio.create_task( + cls._process_step_with_failover( + ctx, + callbacks, + messages, + last_user, + ) + ) ctx._current_step_task = step_task step_started_at = asyncio.get_event_loop().time() try: @@ -1256,6 +2035,7 @@ async def progress_callback_overflow(stage: str, data: dict) -> None: # Handle result if step_result.action == "stop": + loop_error = step_result.error # Report error if step failed if step_result.error and callbacks.on_error: await callbacks.on_error(step_result.error) @@ -1266,7 +2046,13 @@ async def progress_callback_overflow(stage: str, data: dict) -> None: else: post_messages = await Message.list(ctx.session.id) for msg in reversed(post_messages): - if msg.role == MessageRole.ASSISTANT: + if ( + msg.role == MessageRole.ASSISTANT + and ( + not ctx.auto_failover + or getattr(msg, "parentID", None) == last_user.id + ) + ): last_message = msg break @@ -1375,6 +2161,20 @@ async def progress_callback_overflow(stage: str, data: dict) -> None: }) continue + if ( + not step_result.error + and not ctx.should_abort() + and last_message is not None + and getattr(last_message, "finish", None) == "stop" + and await cls._run_turn_finish_hook( + ctx, + callbacks, + last_user, + last_message, + ) + ): + continue + stop_reason = step_result.error or (getattr(last_message, "finish", None) if last_message else None) or "stop" turn_state = set_turn_state( ctx.session.id, @@ -1429,12 +2229,16 @@ async def progress_callback_overflow(stage: str, data: dict) -> None: # Return result return LoopResult( - action="stop", + action="error" if ctx.auto_failover and loop_error else "stop", last_message=last_message, + error=loop_error if ctx.auto_failover else None, + provider_id=ctx.provider_id, + model_id=ctx.model_id, metadata={ "steps": ctx.step, "session_id": ctx.session.id, "last_compaction_step": ctx.last_compaction_step, + **({"aborted": True} if ctx.should_abort() else {}), }, ) @@ -1709,6 +2513,7 @@ async def _execute_subtask( agent=last_user.agent if hasattr(last_user, 'agent') else agent_name, model=last_user.model if hasattr(last_user, 'model') else model_id, provider=last_user.provider if hasattr(last_user, 'provider') else provider_id, + synthetic=True, ) log.info("loop.subtask.completed", { diff --git a/flocks/session/streaming/stream_processor.py b/flocks/session/streaming/stream_processor.py index 1835e0ff8..ca1666474 100644 --- a/flocks/session/streaming/stream_processor.py +++ b/flocks/session/streaming/stream_processor.py @@ -119,6 +119,10 @@ def __init__( workspace_dir: Optional[str] = None, langfuse_generation: Optional[Any] = None, step_index: Optional[int] = None, + execution_mode: str = "build", + plan_file_path: Optional[str] = None, + plan_relative_path: Optional[str] = None, + plan_permission_path: Optional[str] = None, ): self.session_id = session_id self.assistant_message = assistant_message @@ -136,6 +140,10 @@ def __init__( self._workspace_dir = workspace_dir self._langfuse_generation = langfuse_generation self._step_index = step_index + self._execution_mode = execution_mode + self._plan_file_path = plan_file_path + self._plan_relative_path = plan_relative_path + self._plan_permission_path = plan_permission_path self._sandbox_runtime_cache = None self._sandbox_config_cache = None self._sandbox_context_cache = None @@ -657,9 +665,11 @@ async def _handle_tool_call(self, event: ToolCallEvent) -> None: except Exception as e: log.error("stream.tool_start_callback.error", {"error": str(e)}) - # Hook pipeline: tool.execute.before + # Hook pipeline: tool_before + post_hook_fired = False try: from flocks.hooks.pipeline import HookPipeline + hook_ctx = await HookPipeline.run_tool_before({ "sessionID": self.session_id, "workspace": self._workspace_dir, @@ -674,10 +684,28 @@ async def _handle_tool_call(self, event: ToolCallEvent) -> None: updated = hook_ctx.input.get("tool", {}).get("input") if isinstance(updated, dict): tool_input = updated - hook_skip = hook_ctx.output.get("skip") if hook_ctx else False + decision = str( + (hook_ctx.output.get("decision") if hook_ctx else "") or "" + ).strip().lower() + hook_blocked = decision == "block" + hook_block_reason = str( + (hook_ctx.output.get("reason") if hook_ctx else "") or "" + ).strip() + except asyncio.CancelledError: + await self._finalize_interrupted_tool_call( + tool_state=tool_state, + tool_name=tool_name, + tool_input=tool_input, + tool_call_id=tool_call_id, + tool_start_time=tool_start_time, + post_hook_fired=False, + ) + raise except Exception as e: log.error("stream.tool_before_hook.error", {"error": str(e)}) - hook_skip = False + hook_blocked = False + hook_block_reason = "" + tool_state.input = tool_input # Execute tool synchronously tool_span_ctx = None @@ -699,10 +727,11 @@ async def _handle_tool_call(self, event: ToolCallEvent) -> None: except Exception as exc: log.debug("stream.tool_span.init_failed", {"error": str(exc)}) try: - if hook_skip: + if hook_blocked: result = ToolResult( success=False, - error="Tool execution blocked by hook", + error=hook_block_reason or "Tool execution blocked by hook", + metadata={"blocked_by_hook": True}, ) else: sandbox_meta = await self._resolve_sandbox_meta(tool_name) @@ -818,6 +847,26 @@ def _mark_finished() -> None: _cb.mark_finished = _mark_finished return _cb + tool_extra = { + **sandbox_meta["extra"], + "execution_mode": self._execution_mode, + "workspace_dir": self._workspace_dir, + "model": { + "providerID": getattr( + self.assistant_message, + "providerID", + None, + ), + "modelID": getattr( + self.assistant_message, + "modelID", + None, + ), + }, + "plan_file_path": self._plan_file_path, + "plan_relative_path": self._plan_relative_path, + "plan_permission_path": self._plan_permission_path, + } ctx = ToolContext( session_id=self.session_id, message_id=self.assistant_message.id, @@ -825,7 +874,7 @@ def _mark_finished() -> None: call_id=tool_call_id, abort_event=self.abort_event, permission_callback=self.permission_callback, - extra=sandbox_meta["extra"], + extra=tool_extra, metadata_callback=_make_metadata_cb(), event_publish_callback=self.event_publish_callback, ) @@ -843,26 +892,23 @@ def _mark_finished() -> None: if cb and hasattr(cb, 'mark_finished'): cb.mark_finished() - # Hook pipeline: tool.execute.after - try: - from flocks.hooks.pipeline import HookPipeline - hook_ctx = await HookPipeline.run_tool_after({ - "sessionID": self.session_id, - "workspace": self._workspace_dir, - "agent": self.agent.name, - "tool": { - "name": tool_name, - "input": tool_input, - "callID": tool_call_id, - }, - "result": result.model_dump(), - }) - if hook_ctx and isinstance(hook_ctx.output, dict): - override = hook_ctx.output.get("result") - if isinstance(override, dict): - result = ToolResult(**override) - except Exception as e: - log.error("stream.tool_after_hook.error", {"error": str(e)}) + if result.success: + hook_status = "completed" + elif (result.metadata or {}).get("blocked_by_hook"): + hook_status = "blocked" + elif (result.metadata or {}).get("blocked_by_policy"): + hook_status = "blocked" + else: + hook_status = "error" + post_hook_fired = True + result = await self._run_tool_after_hook( + tool_name=tool_name, + tool_input=tool_input, + tool_call_id=tool_call_id, + result=result, + status=hook_status, + tool_start_time=tool_start_time, + ) # Update tool state tool_state.status = "completed" if result.success else "error" @@ -874,14 +920,10 @@ def _mark_finished() -> None: "tool_name": tool_name, "success": result.success, }) - output_preview = result.output if result.success else result.error - if isinstance(output_preview, str): - output_preview = output_preview[:600] - try: if tool_span_ctx is not None: end_kwargs: Dict[str, Any] = { - "output": output_preview, + "output": result.output if result.success else result.error, "metadata": { "success": result.success, "title": result.title, @@ -968,65 +1010,15 @@ def _mark_finished() -> None: log.error("stream.tool_end_callback.error", {"error": str(e)}) except asyncio.CancelledError: - interrupt_msg = "Tool execution was interrupted" - log.info("stream.tool_call.cancelled", { - "tool_call_id": tool_call_id, - "tool_name": tool_name, - }) - try: - if tool_span_ctx is not None: - tool_span_ctx.end( - output=interrupt_msg, - metadata={"success": False}, - level="ERROR", - status_message="tool_cancelled", - ) - except Exception as _span_err: - log.debug("stream.tool_span.cancel_end_failed", {"error": str(_span_err)}) - - tool_state.status = "error" - tool_state.error = interrupt_msg - - try: - tool_end_time = int(datetime.now().timestamp() * 1000) - error_state = ToolStateError( - status="error", - input=tool_input, - error=interrupt_msg, - time={"start": tool_start_time if 'tool_start_time' in locals() else tool_end_time, "end": tool_end_time}, - ) - - error_part = ToolPart( - id=tool_state.part_id, - sessionID=self.session_id, - messageID=self.assistant_message.id, - type="tool", - callID=tool_call_id, - tool=tool_name, - state=error_state, - ) - await Message.store_part(self.session_id, self.assistant_message.id, error_part) - - if self.event_publish_callback: - await self.event_publish_callback("message.part.updated", { - "part": { - "id": tool_state.part_id, - "messageID": self.assistant_message.id, - "sessionID": self.session_id, - "type": "tool", - "callID": tool_call_id, - "tool": tool_name, - "state": { - "status": "error", - "input": tool_input, - "error": interrupt_msg, - "time": {"start": tool_start_time if 'tool_start_time' in locals() else tool_end_time, "end": tool_end_time}, - } - } - }) - except Exception as store_e: - log.error("stream.tool_call.cancelled_update_failed", {"error": str(store_e)}) - + await self._finalize_interrupted_tool_call( + tool_state=tool_state, + tool_name=tool_name, + tool_input=tool_input, + tool_call_id=tool_call_id, + tool_start_time=tool_start_time, + post_hook_fired=post_hook_fired, + tool_span_ctx=tool_span_ctx, + ) raise except Exception as e: log.error("stream.tool_call.error", { @@ -1034,6 +1026,21 @@ def _mark_finished() -> None: "tool_name": tool_name, "error": str(e), }) + if not post_hook_fired: + post_hook_fired = True + exception_result = ToolResult( + success=False, + error=str(e), + ) + await self._run_tool_after_hook( + tool_name=tool_name, + tool_input=tool_input, + tool_call_id=tool_call_id, + result=exception_result, + status="error", + tool_start_time=tool_start_time, + allow_result_override=False, + ) try: if tool_span_ctx is not None: tool_span_ctx.end( @@ -1107,6 +1114,148 @@ def _mark_finished() -> None: except Exception as e2: log.error("stream.tool_end_callback.error", {"error": str(e2)}) + async def _finalize_interrupted_tool_call( + self, + *, + tool_state: ToolCallState, + tool_name: str, + tool_input: Dict[str, Any], + tool_call_id: str, + tool_start_time: int, + post_hook_fired: bool, + tool_span_ctx: Optional[Any] = None, + ) -> None: + """Emit and persist the terminal state for an interrupted tool call.""" + interrupt_msg = "Tool execution was interrupted" + log.info("stream.tool_call.cancelled", { + "tool_call_id": tool_call_id, + "tool_name": tool_name, + }) + try: + if not post_hook_fired: + interrupted_result = ToolResult( + success=False, + error=interrupt_msg, + metadata={"interrupted": True}, + ) + await self._run_tool_after_hook( + tool_name=tool_name, + tool_input=tool_input, + tool_call_id=tool_call_id, + result=interrupted_result, + status="interrupted", + tool_start_time=tool_start_time, + allow_result_override=False, + ) + finally: + try: + if tool_span_ctx is not None: + tool_span_ctx.end( + output=interrupt_msg, + metadata={"success": False}, + level="ERROR", + status_message="tool_cancelled", + ) + except Exception as span_error: + log.debug( + "stream.tool_span.cancel_end_failed", + {"error": str(span_error)}, + ) + + tool_state.status = "error" + tool_state.error = interrupt_msg + tool_end_time = int(datetime.now().timestamp() * 1000) + error_state = ToolStateError( + status="error", + input=tool_input, + error=interrupt_msg, + time={"start": tool_start_time, "end": tool_end_time}, + ) + error_part = ToolPart( + id=tool_state.part_id, + sessionID=self.session_id, + messageID=self.assistant_message.id, + type="tool", + callID=tool_call_id, + tool=tool_name, + state=error_state, + ) + try: + await Message.store_part( + self.session_id, + self.assistant_message.id, + error_part, + ) + if self.event_publish_callback: + await self.event_publish_callback( + "message.part.updated", + { + "part": { + "id": tool_state.part_id, + "messageID": self.assistant_message.id, + "sessionID": self.session_id, + "type": "tool", + "callID": tool_call_id, + "tool": tool_name, + "state": { + "status": "error", + "input": tool_input, + "error": interrupt_msg, + "time": { + "start": tool_start_time, + "end": tool_end_time, + }, + }, + } + }, + ) + except Exception as store_error: + log.error( + "stream.tool_call.cancelled_update_failed", + {"error": str(store_error)}, + ) + + async def _run_tool_after_hook( + self, + *, + tool_name: str, + tool_input: Dict[str, Any], + tool_call_id: str, + result: ToolResult, + status: str, + tool_start_time: int, + allow_result_override: bool = True, + ) -> ToolResult: + """Run tool_after with a normalized result for every tool outcome.""" + try: + from flocks.hooks.pipeline import HookPipeline + + duration_ms = max( + 0, + int(datetime.now().timestamp() * 1000) - tool_start_time, + ) + hook_ctx = await HookPipeline.run_tool_after({ + "sessionID": self.session_id, + "workspace": self._workspace_dir, + "agent": self.agent.name, + "tool": { + "name": tool_name, + "input": tool_input, + "callID": tool_call_id, + }, + "status": status, + "durationMs": duration_ms, + "result": result.model_dump(), + "error": result.error if not result.success else None, + }) + if allow_result_override and isinstance(hook_ctx.output, dict): + override = hook_ctx.output.get("result") + if isinstance(override, dict): + return ToolResult(**override) + except Exception as exc: + log.error("stream.tool_after_hook.error", {"error": str(exc)}) + return result + async def _load_config_data(self) -> Dict[str, Any]: """Load and cache config as plain dict.""" if isinstance(self._config_data, dict): diff --git a/flocks/session/streaming/timeouts.py b/flocks/session/streaming/timeouts.py new file mode 100644 index 000000000..4cc4babc7 --- /dev/null +++ b/flocks/session/streaming/timeouts.py @@ -0,0 +1,156 @@ +"""Resolve adaptive timeout budgets for LLM streaming responses.""" + +from __future__ import annotations + +import ipaddress +import os +from dataclasses import dataclass +from typing import Any, Mapping, Optional +from urllib.parse import urlparse + + +DEFAULT_FIRST_CHUNK_TIMEOUT_S = 120.0 +DEFAULT_LOCAL_FIRST_CHUNK_TIMEOUT_S = 1800.0 +DEFAULT_ONGOING_CHUNK_TIMEOUT_S = 300.0 + +FIRST_CHUNK_TIMEOUT_ENV = "FLOCKS_LLM_STREAM_FIRST_CHUNK_TIMEOUT_S" +ONGOING_CHUNK_TIMEOUT_ENV = "FLOCKS_LLM_STREAM_ONGOING_CHUNK_TIMEOUT_S" + +_FIRST_CHUNK_SETTING_KEYS = ( + "stream_first_chunk_timeout_s", + "streamFirstChunkTimeoutSeconds", +) +_ONGOING_CHUNK_SETTING_KEYS = ( + "stream_ongoing_chunk_timeout_s", + "streamOngoingChunkTimeoutSeconds", +) +_LOCAL_PROVIDER_IDS = frozenset({"local", "ollama"}) +_LOCAL_HOSTNAMES = frozenset({"localhost", "host.docker.internal"}) + + +@dataclass(frozen=True) +class LlmStreamTimeouts: + """Effective timeout budgets for one LLM stream.""" + + first_chunk_s: float + ongoing_chunk_s: float + is_local: bool + + +def _positive_float(value: Any) -> Optional[float]: + """Return a positive finite float, or ``None`` for an invalid value.""" + try: + parsed = float(value) + except (TypeError, ValueError): + return None + if parsed <= 0 or parsed == float("inf") or parsed != parsed: + return None + return parsed + + +def _setting( + settings: Mapping[str, Any], + keys: tuple[str, ...], +) -> Optional[float]: + for key in keys: + if key in settings: + value = _positive_float(settings[key]) + if value is not None: + return value + return None + + +def _provider_base_url(provider: Any) -> str: + config = getattr(provider, "_config", None) + candidates = ( + getattr(config, "base_url", None), + getattr(provider, "_base_url", None), + getattr(provider, "DEFAULT_BASE_URL", None), + ) + return next( + (value.strip() for value in candidates if isinstance(value, str) and value.strip()), + "", + ) + + +def _is_local_endpoint(provider: Any) -> bool: + provider_id = str(getattr(provider, "id", "") or "").strip().lower() + if provider_id in _LOCAL_PROVIDER_IDS: + return True + + hostname = (urlparse(_provider_base_url(provider)).hostname or "").lower() + if not hostname: + return False + if hostname in _LOCAL_HOSTNAMES or hostname.endswith(".local"): + return True + + try: + address = ipaddress.ip_address(hostname) + except ValueError: + return False + return address.is_loopback or address.is_private or address.is_link_local + + +def _provider_settings(provider: Any) -> Mapping[str, Any]: + config = getattr(provider, "_config", None) + settings = getattr(config, "custom_settings", None) + return settings if isinstance(settings, Mapping) else {} + + +def _model_settings(provider: Any, model_id: str) -> Mapping[str, Any]: + try: + models = provider.get_models() + except Exception: + return {} + for model in models or []: + if getattr(model, "id", None) != model_id: + continue + settings = getattr(model, "custom_settings", None) + return settings if isinstance(settings, Mapping) else {} + return {} + + +def _resolve_timeout( + *, + model_settings: Mapping[str, Any], + provider_settings: Mapping[str, Any], + setting_keys: tuple[str, ...], + environment_name: str, + default: float, +) -> float: + candidates = ( + _setting(model_settings, setting_keys), + _setting(provider_settings, setting_keys), + _positive_float(os.getenv(environment_name)), + ) + return next((value for value in candidates if value is not None), default) + + +def resolve_llm_stream_timeouts(provider: Any, model_id: str) -> LlmStreamTimeouts: + """Resolve model, provider, environment, and endpoint-aware timeouts. + + Precedence is model configuration, provider configuration, environment, + then the built-in endpoint-aware default. + """ + is_local = _is_local_endpoint(provider) + provider_settings = _provider_settings(provider) + model_settings = _model_settings(provider, model_id) + first_chunk_default = DEFAULT_LOCAL_FIRST_CHUNK_TIMEOUT_S if is_local else DEFAULT_FIRST_CHUNK_TIMEOUT_S + + return LlmStreamTimeouts( + first_chunk_s=_resolve_timeout( + model_settings=model_settings, + provider_settings=provider_settings, + setting_keys=_FIRST_CHUNK_SETTING_KEYS, + environment_name=FIRST_CHUNK_TIMEOUT_ENV, + default=first_chunk_default, + ), + ongoing_chunk_s=_resolve_timeout( + model_settings=model_settings, + provider_settings=provider_settings, + setting_keys=_ONGOING_CHUNK_SETTING_KEYS, + environment_name=ONGOING_CHUNK_TIMEOUT_ENV, + default=DEFAULT_ONGOING_CHUNK_TIMEOUT_S, + ), + is_local=is_local, + ) diff --git a/flocks/session/streaming/tool_accumulator.py b/flocks/session/streaming/tool_accumulator.py index 1957ed259..6e266e95d 100644 --- a/flocks/session/streaming/tool_accumulator.py +++ b/flocks/session/streaming/tool_accumulator.py @@ -43,12 +43,17 @@ def __init__(self, processor: Any) -> None: async def feed_chunk(self, tc: dict[str, Any]) -> None: """Process a single tool-call chunk from the provider stream.""" tc_index = tc.get("index", 0) - tc_id = tc.get("id") + provider_tc_id = tc.get("id") - if tc_id: - self._index_to_id[tc_index] = tc_id - elif tc_index in self._index_to_id: + # Keep one stable internal id for the lifetime of a streamed tool call. + # Some providers reveal the tool name before their call id; replacing a + # generated id later would disconnect the already-published pending UI + # part from the eventual ToolCallEvent. + if tc_index in self._index_to_id: tc_id = self._index_to_id[tc_index] + elif provider_tc_id: + tc_id = provider_tc_id + self._index_to_id[tc_index] = tc_id else: tc_id = Identifier.create("call") self._index_to_id[tc_index] = tc_id @@ -79,6 +84,17 @@ async def feed_chunk(self, tc: dict[str, Any]) -> None: accumulated = self._accumulator[tc_id]["arguments_str"] final_name = self._accumulator[tc_id]["name"] + + # Publish the tool step as soon as its name is known. Parameters can be + # large (especially write/edit content), so waiting for valid JSON here + # would hide the tool until input generation — and often execution — + # was nearly complete. + if final_name and not self._accumulator[tc_id].get("input_started"): + await self._processor.process_event( + ToolInputStartEvent(id=tc_id, tool_name=final_name) + ) + self._accumulator[tc_id]["input_started"] = True + if accumulated and final_name: arguments, ok = _parse_json_robust(accumulated) if ok: @@ -89,12 +105,6 @@ async def feed_chunk(self, tc: dict[str, Any]) -> None: self._accumulator[tc_id]["awaiting_required"] = True return - if not self._accumulator[tc_id].get("input_started") and final_name: - await self._processor.process_event( - ToolInputStartEvent(id=tc_id, tool_name=final_name) - ) - self._accumulator[tc_id]["input_started"] = True - if final_name: await self._processor.process_event( ToolCallEvent( diff --git a/flocks/storage/storage.py b/flocks/storage/storage.py index debf1bf16..43a8724cc 100644 --- a/flocks/storage/storage.py +++ b/flocks/storage/storage.py @@ -13,7 +13,7 @@ from contextlib import asynccontextmanager from pathlib import Path import sqlite3 -from typing import Any, AsyncIterator, Awaitable, Callable, Dict, List, Optional, Tuple, Type, TypeVar +from typing import Any, AsyncIterator, Awaitable, Callable, Dict, List, Optional, Sequence, Tuple, Type, TypeVar import json import aiosqlite from datetime import datetime @@ -1543,6 +1543,105 @@ async def _write() -> None: cls._log.debug("storage.set", {"key": key, "type": value_type}) + @classmethod + async def set_many(cls, entries: Sequence[Tuple[str, Any, str]]) -> None: + """Store multiple values atomically in one SQLite transaction.""" + await cls.mutate_many(set_entries=entries) + + @classmethod + async def mutate_many( + cls, + *, + set_entries: Sequence[Tuple[str, Any, str]] = (), + delete_keys: Sequence[str] = (), + delete_prefixes: Sequence[str] = (), + ) -> int: + """Apply related set/delete operations in one SQLite transaction.""" + entries = list(set_entries) + keys_to_delete = list(delete_keys) + prefixes_to_delete = list(delete_prefixes) + if not entries and not keys_to_delete and not prefixes_to_delete: + return 0 + + routing_paths = { + *(cls.route_db_path_for_key(key) for key, _value, _value_type in entries), + *(cls.route_db_path_for_key(key) for key in keys_to_delete), + *(cls.route_db_path_for_prefix(prefix) for prefix in prefixes_to_delete), + } + if len(routing_paths) != 1: + raise ValueError("Storage.mutate_many operations must target the same database") + + if not entries: + serialized_entries = [] + else: + serialized_entries = [ + ( + key, + value.model_dump_json() if isinstance(value, BaseModel) else json.dumps(value), + value_type, + ) + for key, value, value_type in entries + ] + + await cls._ensure_init() + db_path = next(iter(routing_paths)) + + from datetime import UTC + + now = datetime.now(UTC).isoformat() + + async def _write() -> int: + async with cls.connect(db_path) as db: + try: + await db.execute("BEGIN IMMEDIATE") + for key, serialized, value_type in serialized_entries: + await db.execute( + """ + INSERT OR REPLACE INTO storage (key, value, type, created_at, updated_at) + VALUES (?, ?, ?, + COALESCE((SELECT created_at FROM storage WHERE key = ?), ?), + ?) + """, + (key, serialized, value_type, key, now, now), + ) + deleted = 0 + for key in keys_to_delete: + cursor = await db.execute("DELETE FROM storage WHERE key = ?", (key,)) + deleted += max(cursor.rowcount, 0) + for prefix in prefixes_to_delete: + cursor = await db.execute( + f"DELETE FROM storage WHERE {cls._like_prefix_clause()}", + (cls._like_prefix_pattern(prefix),), + ) + deleted += max(cursor.rowcount, 0) + await db.commit() + return deleted + except BaseException: + await db.rollback() + raise + + deleted = await cls._run_with_corruption_recovery( + lambda: cls._run_write_with_retry( + _write, + action="mutate_many", + target=f"{len(entries)} sets/{len(keys_to_delete)} keys/{len(prefixes_to_delete)} prefixes", + ), + db_path=db_path, + action="mutate_many", + ) + + for key in keys_to_delete: + cls._invalidate_runtime_caches(key) + for prefix in prefixes_to_delete: + cls._invalidate_runtime_caches(prefix) + cls._log.debug("storage.mutate_many", { + "set_count": len(entries), + "delete_key_count": len(keys_to_delete), + "delete_prefix_count": len(prefixes_to_delete), + "deleted": deleted, + }) + return deleted + @classmethod async def get(cls, key: str, model: Optional[Type[T]] = None) -> Optional[T | Any]: """ diff --git a/flocks/tool/agent/delegate_task.py b/flocks/tool/agent/delegate_task.py index 88c42d52b..cf00616be 100644 --- a/flocks/tool/agent/delegate_task.py +++ b/flocks/tool/agent/delegate_task.py @@ -4,6 +4,8 @@ from __future__ import annotations +import asyncio +import time from typing import Optional, List, Dict, Any from flocks.tool.registry import ( @@ -26,12 +28,152 @@ from flocks.agent.registry import is_delegatable from flocks.skill.skill import Skill from flocks.config.config import Config -from flocks.tool.subagent_result import format_sync_subagent_result +from flocks.tool.subagent_result import ( + _extract_message_error, + format_sync_subagent_result, +) from flocks.utils.log import Log log = Log.create(service="tool.delegate_task") +async def _run_subagent_with_hooks( + *, + ctx: ToolContext, + child_session_id: str, + child_agent: str, + workspace: str, + prompt: str, + description: str, + resumed: bool, + provider_id: Optional[str] = None, + model_id: Optional[str] = None, + callbacks: Optional[Any] = None, +) -> Any: + """Run one child session with paired SubagentStart/SubagentStop hooks.""" + from flocks.hooks.pipeline import HookPipeline + + common_payload = { + "sessionID": ctx.session_id, + "workspace": workspace, + "parentSessionID": ctx.session_id, + "parentMessageID": ctx.message_id, + "childSessionID": child_session_id, + "agentType": child_agent, + "prompt": prompt, + "description": description, + "resumed": resumed, + } + try: + await HookPipeline.run_subagent_start(common_payload) + except Exception as exc: + log.debug("delegate_task.hook.subagent_start.error", { + "child_session_id": child_session_id, + "error": str(exc), + }) + + started_at = time.perf_counter() + try: + result = await SessionLoop.run( + child_session_id, + provider_id=provider_id, + model_id=model_id, + callbacks=callbacks, + ) + except asyncio.CancelledError: + duration_ms = int((time.perf_counter() - started_at) * 1000) + try: + await HookPipeline.run_subagent_stop({ + **common_payload, + "status": "interrupted", + "durationMs": duration_ms, + "summary": None, + "error": "Sub-agent execution was interrupted", + }) + except Exception as exc: + log.debug("delegate_task.hook.subagent_stop.error", { + "child_session_id": child_session_id, + "error": str(exc), + }) + raise + except Exception as exc: + duration_ms = int((time.perf_counter() - started_at) * 1000) + try: + await HookPipeline.run_subagent_stop({ + **common_payload, + "status": "error", + "durationMs": duration_ms, + "summary": None, + "error": str(exc), + }) + except Exception as hook_exc: + log.debug("delegate_task.hook.subagent_stop.error", { + "child_session_id": child_session_id, + "error": str(hook_exc), + }) + raise + + summary = None + last_message = getattr(result, "last_message", None) + if last_message is not None: + try: + summary = await Message.get_text_content(last_message) + except Exception as exc: + log.debug("delegate_task.hook.subagent_summary.error", { + "child_session_id": child_session_id, + "error": str(exc), + }) + result_error = getattr(result, "error", None) + message_error = ( + _extract_message_error(last_message) + if last_message is not None + else None + ) + message_finish = ( + getattr(last_message, "finish", None) + if last_message is not None + else None + ) + result_metadata = getattr(result, "metadata", None) + interrupted = ( + isinstance(result_metadata, dict) + and bool(result_metadata.get("aborted")) + ) + if interrupted: + status = "interrupted" + stop_error = result_error or "Sub-agent execution was interrupted" + elif ( + getattr(result, "action", None) == "error" + or result_error + or message_error + or message_finish == "error" + ): + status = "error" + stop_error = ( + result_error + or message_error + or "Sub-agent execution failed" + ) + else: + status = "completed" + stop_error = None + duration_ms = int((time.perf_counter() - started_at) * 1000) + try: + await HookPipeline.run_subagent_stop({ + **common_payload, + "status": status, + "durationMs": duration_ms, + "summary": summary, + "error": stop_error, + }) + except Exception as exc: + log.debug("delegate_task.hook.subagent_stop.error", { + "child_session_id": child_session_id, + "error": str(exc), + }) + return result + + async def _subagent_session_permissions(agent_name: str) -> list: """Build session permission rules for a delegated subagent.""" from flocks.agent.registry import Agent @@ -395,8 +537,15 @@ async def delegate_task_tool( agent=session.agent or ctx.agent, ) from flocks.session.session_loop import LoopCallbacks - result = await SessionLoop.run( - session.id, + + result = await _run_subagent_with_hooks( + ctx=ctx, + child_session_id=session.id, + child_agent=session.agent or ctx.agent, + workspace=getattr(session, "directory", None) or "", + prompt=prompt, + description=description, + resumed=True, callbacks=LoopCallbacks( event_publish_callback=ctx.event_publish_callback, ), @@ -494,8 +643,14 @@ async def delegate_task_tool( description=description, ) ctx.metadata({"title": description, "metadata": {"sessionId": created.id, "status": "running"}}) - result = await SessionLoop.run( - created.id, + result = await _run_subagent_with_hooks( + ctx=ctx, + child_session_id=created.id, + child_agent=agent_to_use, + workspace=runtime_directory, + prompt=full_prompt, + description=description, + resumed=False, provider_id=(category_model or {}).get("providerID"), model_id=(category_model or {}).get("modelID"), callbacks=forwarder.build_callbacks( diff --git a/flocks/tool/file/apply_patch.py b/flocks/tool/file/apply_patch.py index af46ef698..e7ef476e1 100644 --- a/flocks/tool/file/apply_patch.py +++ b/flocks/tool/file/apply_patch.py @@ -293,7 +293,12 @@ async def apply_patch_tool( ) sandbox = ctx.extra.get("sandbox") if ctx.extra else None - if isinstance(sandbox, dict) and sandbox.get("workspace_access") == "ro": + sandbox_read_only = ( + isinstance(sandbox, dict) + and sandbox.get("workspace_access") == "ro" + ) + execution_mode = ctx.extra.get("execution_mode") if ctx.extra else None + if sandbox_read_only and execution_mode != "plan": return ToolResult( success=False, error=( @@ -397,6 +402,25 @@ async def apply_patch_tool( success=False, error=f"Failed to process hunk for {hunk.path}: {str(e)}" ) + + if sandbox_read_only: + from flocks.session.execution_mode import is_plan_file_edit + + if not all( + is_plan_file_edit(execution_mode, ctx, change["filePath"]) + and ( + not change.get("movePath") + or is_plan_file_edit(execution_mode, ctx, change["movePath"]) + ) + for change in file_changes + ): + return ToolResult( + success=False, + error=( + "Patch is blocked in sandbox read-only workspace mode. " + "Only the current session plan file may be changed in Plan mode." + ), + ) # Request permission await ctx.ask( diff --git a/flocks/tool/file/edit.py b/flocks/tool/file/edit.py index abc0566cc..377ab212d 100644 --- a/flocks/tool/file/edit.py +++ b/flocks/tool/file/edit.py @@ -526,6 +526,20 @@ async def edit_tool( sandbox = ctx.extra.get("sandbox") if ctx.extra else None if isinstance(sandbox, dict) and sandbox.get("workspace_access") == "ro": + from flocks.session.execution_mode import is_plan_file_edit + + plan_file_edit = is_plan_file_edit( + ctx.extra.get("execution_mode"), + ctx, + filepath, + ) + else: + plan_file_edit = False + if ( + isinstance(sandbox, dict) + and sandbox.get("workspace_access") == "ro" + and not plan_file_edit + ): return ToolResult( success=False, error=( diff --git a/flocks/tool/file/write.py b/flocks/tool/file/write.py index eee0379cf..50aee858e 100644 --- a/flocks/tool/file/write.py +++ b/flocks/tool/file/write.py @@ -273,6 +273,20 @@ async def write_tool( sandbox = ctx.extra.get("sandbox") if ctx.extra else None if isinstance(sandbox, dict) and sandbox.get("workspace_access") == "ro": + from flocks.session.execution_mode import is_plan_file_edit + + plan_file_edit = is_plan_file_edit( + ctx.extra.get("execution_mode"), + ctx, + filepath, + ) + else: + plan_file_edit = False + if ( + isinstance(sandbox, dict) + and sandbox.get("workspace_access") == "ro" + and not plan_file_edit + ): return ToolResult( success=False, error=( diff --git a/flocks/tool/registry.py b/flocks/tool/registry.py index b172a847d..35d1c3cd7 100644 --- a/flocks/tool/registry.py +++ b/flocks/tool/registry.py @@ -267,6 +267,20 @@ async def ask( always: Always-allow patterns metadata: Additional metadata """ + execution_mode = self.extra.get("execution_mode") + if execution_mode: + from flocks.session.execution_mode import is_permission_allowed + + if not is_permission_allowed( + execution_mode, + permission, + patterns, + self, + ): + raise PermissionError( + "Plan mode may only edit the current session plan file." + ) + request = PermissionRequest( permission=permission, patterns=patterns, @@ -870,6 +884,40 @@ async def execute( message_id="default" ) + execution_mode = ctx.extra.get("execution_mode") + if execution_mode: + from flocks.session.execution_mode import ( + is_tool_allowed, + tool_call_denial_reason, + ) + + if not is_tool_allowed(execution_mode, tool_name): + log.warn("tool.execute.execution_mode_denied", { + "name": tool_name, + "execution_mode": str(execution_mode), + "session_id": ctx.session_id, + }) + return ToolResult( + success=False, + error=( + f"Tool {tool_name!r} is not available in " + f"{str(execution_mode)!r} execution mode." + ), + ) + denial_reason = tool_call_denial_reason( + execution_mode, + tool_name, + kwargs, + ctx, + ) + if denial_reason: + log.warn("tool.execute.execution_mode_call_denied", { + "name": tool_name, + "execution_mode": str(execution_mode), + "session_id": ctx.session_id, + }) + return ToolResult(success=False, error=denial_reason) + log.info("tool.execute", { "name": tool_name, "params": list(kwargs.keys()), @@ -1616,7 +1664,7 @@ def _register_builtin_tools(cls) -> None: # security/ — SSH forensics + threat intelligence (optional: asyncssh) ("flocks.tool.security", ["ssh_host_cmd", "ssh_run_script"]), # system/ — questions, model config, memory, MCP management, session management, slash commands - ("flocks.tool.system", ["question", "model_config", "memory", "flocks_mcp", "session_manage", "slash_command", "tool_search"]), + ("flocks.tool.system", ["question", "plan_exit", "model_config", "memory", "flocks_mcp", "session_manage", "slash_command", "tool_search"]), # skill/ — skill management (search, install, status, deps, remove, load) ("flocks.tool.skill", ["flocks_skills", "skill_load"]), # device/ — security device asset context and status probes @@ -1650,16 +1698,43 @@ def _register_builtin_tools(cls) -> None: if "get_time" not in cls._tools: @cls.register_function( name="get_time", - description="Get current date and time", + description="Get current date and time in ISO 8601 or Unix timestamp format", category=ToolCategory.SYSTEM, native=True, - parameters=[] + parameters=[ + ToolParameter( + name="format", + type=ParameterType.STRING, + description=( + "Output format: 'iso' for ISO 8601, 'unix' for Unix seconds, " + "or 'unix_ms' for Unix milliseconds. Defaults to 'iso'." + ), + required=False, + default="iso", + enum=["iso", "unix", "unix_ms"], + ) + ] ) - async def get_time(ctx: ToolContext) -> ToolResult: + async def get_time(ctx: ToolContext, format: str = "iso") -> ToolResult: from datetime import datetime + + if format not in {"iso", "unix", "unix_ms"}: + return ToolResult( + success=False, + error="format must be one of: iso, unix, unix_ms", + ) + + now = datetime.now() + if format == "unix": + output = str(int(now.timestamp())) + elif format == "unix_ms": + output = str(int(now.timestamp() * 1000)) + else: + output = now.isoformat() + return ToolResult( success=True, - output=datetime.now().isoformat() + output=output, ) @classmethod diff --git a/flocks/tool/system/plan_exit.py b/flocks/tool/system/plan_exit.py new file mode 100644 index 000000000..2dacda008 --- /dev/null +++ b/flocks/tool/system/plan_exit.py @@ -0,0 +1,186 @@ +"""Plan completion and Build handoff modeled after OpenCode's plan_exit tool.""" + +from __future__ import annotations + +from typing import Any + +from flocks.session.execution_mode import SessionExecutionMode +from flocks.session.message import Message, MessageRole +from flocks.session.plan_file import context_plan_file +from flocks.tool.registry import ( + ToolCategory, + ToolContext, + ToolRegistry, + ToolResult, +) +from flocks.tool.system.question import question_tool + + +START_IMPLEMENTING = "开始实施" +CONTINUE_PLANNING = "调整计划" + +DESCRIPTION = """Finish a completed plan and ask the user whether to implement it. + +Call this only after presenting a decision-complete plan. Approval starts a new +Build turn in the current session loop to implement the approved plan. +Declining keeps the session in Plan so the plan can be refined. +""" + + +async def _publish(ctx: ToolContext, event_type: str, properties: dict[str, Any]) -> None: + if ctx.event_publish_callback: + await ctx.event_publish_callback(event_type, properties) + + +async def _turn_model_and_variant( + ctx: ToolContext, +) -> tuple[dict[str, str] | None, str | None]: + """Copy the active Plan turn model like OpenCode's synthetic Build turn.""" + + try: + messages = await Message.list(ctx.session_id) + last_user = next( + ( + message + for message in reversed(messages) + if getattr(message, "role", None) == MessageRole.USER + ), + None, + ) + except Exception: + last_user = None + + model = getattr(last_user, "model", None) + if not isinstance(model, dict) or not all( + model.get(key) for key in ("providerID", "modelID") + ): + model = ctx.extra.get("model") + if not isinstance(model, dict) or not all( + model.get(key) for key in ("providerID", "modelID") + ): + model = None + return model, getattr(last_user, "variant", None) + + +@ToolRegistry.register_function( + name="plan_exit", + description=DESCRIPTION, + category=ToolCategory.SYSTEM, + parameters=[], +) +async def plan_exit_tool(ctx: ToolContext) -> ToolResult: + """Ask for plan approval and continue immediately in Build mode.""" + + plan = context_plan_file(ctx) + if plan is None or not plan.path.is_file(): + return ToolResult( + success=False, + error="Write the session plan file before calling plan_exit.", + ) + try: + if not plan.path.read_text(encoding="utf-8").strip(): + return ToolResult( + success=False, + error="The session plan file is empty. Complete it before calling plan_exit.", + ) + except OSError as exc: + return ToolResult(success=False, error=f"Could not read the session plan file: {exc}") + + confirmation = await question_tool( + ctx, + questions=[ + { + "header": "Plan complete", + "question": ( + f"The plan at {plan.relative_path} is complete. Would you like " + "to switch to Build and start implementing?" + ), + "type": "choice", + "options": [ + { + "label": START_IMPLEMENTING, + "description": "Switch to Build and implement the approved plan now.", + }, + { + "label": CONTINUE_PLANNING, + "description": "Stay in Plan and describe what should be changed.", + "allowText": True, + }, + ], + "multiple": False, + "custom": False, + } + ], + ) + if not confirmation.success: + return confirmation + if confirmation.metadata.get("deferred"): + return confirmation + + answers = confirmation.metadata.get("answers") or [] + selected = answers[0] if answers else [] + if START_IMPLEMENTING not in selected: + feedback = "\n".join( + str(value).strip() + for value in selected + if str(value).strip() and value != CONTINUE_PLANNING + ) + output = ( + "The user chose to remain in Plan. Continue refining the plan " + "using their feedback." + ) + metadata = { + "approved": False, + "executionMode": SessionExecutionMode.PLAN.value, + } + if feedback: + output = f"{output}\n\nUser feedback:\n{feedback}" + metadata["feedback"] = feedback + return ToolResult( + success=True, + output=output, + title="Continue planning", + metadata=metadata, + ) + + build_model, build_variant = await _turn_model_and_variant(ctx) + build_message = await Message.create( + session_id=ctx.session_id, + role=MessageRole.USER, + content=( + f"The plan at {plan.relative_path} has been approved. " + "Switch to Build mode, read that file, and implement it now." + ), + agent=ctx.agent, + model=build_model, + variant=build_variant, + executionMode=SessionExecutionMode.BUILD, + synthetic=True, + part_metadata={ + "planImplementation": True, + "planPath": plan.relative_path, + }, + ) + await _publish( + ctx, + "session.execution_mode.changed", + { + "sessionID": ctx.session_id, + "executionMode": SessionExecutionMode.BUILD.value, + "reason": "plan-approved", + }, + ) + return ToolResult( + success=True, + output=( + "The plan was approved. Continue immediately in Build mode and " + "implement the approved plan." + ), + title="Plan approved", + metadata={ + "approved": True, + "executionMode": SessionExecutionMode.BUILD.value, + "buildMessageID": build_message.id, + "planPath": plan.relative_path, + }, + ) diff --git a/flocks/tool/system/question.py b/flocks/tool/system/question.py index 4fb80273e..e83af3201 100644 --- a/flocks/tool/system/question.py +++ b/flocks/tool/system/question.py @@ -107,7 +107,7 @@ def _first_non_empty_string(data: Dict[str, Any], keys: tuple[str, ...]) -> str: return "" -def normalize_question_option(opt: Any) -> Optional[Dict[str, str]]: +def normalize_question_option(opt: Any) -> Optional[Dict[str, Any]]: """Normalize LLM-produced choice options into the UI's label/description shape.""" if isinstance(opt, str): label = opt.strip() @@ -123,7 +123,13 @@ def normalize_question_option(opt: Any) -> Optional[Dict[str, str]]: label, description = description, "" if not label: return None - return {"label": label, "description": description} + normalized: Dict[str, Any] = { + "label": label, + "description": description, + } + if opt.get("allowText") is True: + normalized["allowText"] = True + return normalized def _format_channel_question_text(questions: List[Dict[str, Any]]) -> str: @@ -296,6 +302,13 @@ async def default_question_handler( "properties": { "label": {"type": "string"}, "description": {"type": "string"}, + "allowText": { + "type": "boolean", + "description": ( + "Show a text input for this option and " + "return both its label and entered text." + ), + }, }, "required": ["label"], "additionalProperties": False, diff --git a/flocks/tool/system/session_manage.py b/flocks/tool/system/session_manage.py index 7bf2968d2..9eb070f67 100644 --- a/flocks/tool/system/session_manage.py +++ b/flocks/tool/system/session_manage.py @@ -40,8 +40,8 @@ model, provider, and memory_enabled. For requests like "change this session to gpt-5" or "update the session model", use action=update with model, and set provider too when the user specifies one. -- delete: soft-delete a session and its child sessions; requires session_id and - confirmation. +- delete: permanently delete a session, its child sessions, messages, and + history; requires session_id and confirmation. This cannot be restored. - archive: archive or restore a session; requires session_id. Set archive=false to restore an archived session to active. """ @@ -221,7 +221,11 @@ async def session_manage( permission="session_manage", patterns=[f"delete:{session_id}"], always=[], - metadata={"action": "delete", "session_id": session_id}, + metadata={ + "action": "permanent_delete", + "session_id": session_id, + "destructive": True, + }, ) return await _session_delete_impl(ctx, session_id=session_id) if action == "archive": @@ -482,7 +486,7 @@ async def _session_delete_impl(ctx: ToolContext, session_id: str) -> ToolResult: return ToolResult( success=True, - output=f"Session '{session_id}'({session.title})已删除", + output=f"Session '{session_id}'({session.title})已永久删除", ) @@ -495,27 +499,32 @@ async def _session_archive_impl( session_id: str, archive: Optional[bool] = True, ) -> ToolResult: - from flocks.session.session import Session, SessionInfo - from flocks.storage.storage import Storage + from flocks.session.session import Session - # get_by_id 会跳过 archived session,需直接扫 Storage - session = None - keys = await Storage.list_keys(prefix="session:") - for key in keys: - try: - s = await Storage.get(key, SessionInfo) - if s and s.id == session_id and s.status != "deleted": - session = s - break - except Exception: - continue + session = await Session.get_by_id_unfiltered(session_id) if not session: return ToolResult(success=False, error=f"未找到 session '{session_id}'") + if archive is not False and session_id == ctx.session_id: + return ToolResult( + success=False, + error="不能在当前会话的工具调用尚未结束时归档自身,请从工作台归档该会话", + ) + if archive is False and getattr(session, "parent_id", None) is not None: + return ToolResult( + success=False, + error="只能从根任务恢复完整任务树", + ) + try: if archive is False: - ok = await Session.unarchive(session.project_id, session_id) + owner_user_id = getattr(session, "owner_user_id", None) + ok = await Session.restore( + session.project_id, + session_id, + project_owner_id=owner_user_id, + ) action = "取消归档" else: ok = await Session.archive(session.project_id, session_id) diff --git a/flocks/tool/task/todo.py b/flocks/tool/task/todo.py index a9e5f04b0..dd275858d 100644 --- a/flocks/tool/task/todo.py +++ b/flocks/tool/task/todo.py @@ -237,9 +237,12 @@ async def todo_tool( old_todos = await Todo.get(ctx.session_id) normalized_todos = _normalize_todos(todos) if _all_terminal(normalized_todos): - await Todo.update(ctx.session_id, []) + await Todo.update_active(ctx.session_id, []) else: - await Todo.update(ctx.session_id, normalized_todos) + await Todo.update_active( + ctx.session_id, + normalized_todos, + ) old_serialized = _serialize_todos(old_todos) new_serialized = _serialize_todos(normalized_todos) diff --git a/flocks/utils/langfuse.py b/flocks/utils/langfuse.py index 53564a43b..52d7939f9 100644 --- a/flocks/utils/langfuse.py +++ b/flocks/utils/langfuse.py @@ -42,17 +42,47 @@ def end(self, **_: Any) -> None: "langfuse_current_observation", default=None, ) +_CAPTURE_MODE_ENV = "FLOCKS_LANGFUSE_CAPTURE_MODE" +_MAX_CHARS_ENV = "FLOCKS_LANGFUSE_MAX_CHARS" +_DEFAULT_CAPTURE_MODE = "full" +_DEFAULT_MAX_CHARS = 8000 +_VALID_CAPTURE_MODES = {"full", "truncated"} +_TRACE_NAME_ATTR = "langfuse.trace.name" +_TRACE_USER_ID_ATTR = "user.id" +_TRACE_SESSION_ID_ATTR = "session.id" +_TRACE_TAGS_ATTR = "langfuse.trace.tags" def _filter_none(values: Dict[str, Any]) -> Dict[str, Any]: return {k: v for k, v in values.items() if v is not None} -def _truncate_value(value: Any, max_chars: int = 8000) -> Any: +def _get_capture_mode() -> str: + mode = os.getenv(_CAPTURE_MODE_ENV, _DEFAULT_CAPTURE_MODE).strip().lower() + if mode in _VALID_CAPTURE_MODES: + return mode + return _DEFAULT_CAPTURE_MODE + + +def _get_capture_max_chars() -> int: + raw_value = os.getenv(_MAX_CHARS_ENV, str(_DEFAULT_MAX_CHARS)).strip() + try: + max_chars = int(raw_value) + except (TypeError, ValueError): + return _DEFAULT_MAX_CHARS + if max_chars <= 0: + return _DEFAULT_MAX_CHARS + return max_chars + + +def _truncate_value(value: Any, max_chars: Optional[int] = None) -> Any: if isinstance(value, str): - if len(value) <= max_chars: + if _get_capture_mode() == "full": return value - return value[:max_chars] + f"...[truncated:{len(value) - max_chars}]" + limit = max_chars if max_chars is not None else _get_capture_max_chars() + if len(value) <= limit: + return value + return value[:limit] + f"...[truncated:{len(value) - limit}]" return value @@ -67,6 +97,58 @@ def _sanitize_payload(value: Any) -> Any: return _truncate_value(value) +def _set_trace_attribute(observation: Any, key: str, value: Any) -> None: + otel_span = getattr(observation, "_otel_span", None) + if otel_span is None or not hasattr(otel_span, "is_recording"): + return + try: + if otel_span.is_recording(): + otel_span.set_attribute(key, value) + except Exception: + return + + +def _propagate_trace_dimensions( + observation: Any, + *, + trace_name: Optional[str] = None, + session_id: Optional[str] = None, + user_id: Optional[str] = None, + tags: Optional[list[str]] = None, + parent: Any = None, +) -> Any: + inherited_trace_name = trace_name + inherited_session_id = session_id + inherited_user_id = user_id + inherited_tags = tags + + parent_span = getattr(parent, "_otel_span", None) + parent_attrs = getattr(parent_span, "attributes", None) if parent_span is not None else None + if parent_attrs: + if inherited_trace_name is None: + inherited_trace_name = parent_attrs.get(_TRACE_NAME_ATTR) + if inherited_session_id is None: + inherited_session_id = parent_attrs.get(_TRACE_SESSION_ID_ATTR) + if inherited_user_id is None: + inherited_user_id = parent_attrs.get(_TRACE_USER_ID_ATTR) + if inherited_tags is None: + parent_tags = parent_attrs.get(_TRACE_TAGS_ATTR) + if isinstance(parent_tags, list): + inherited_tags = parent_tags + elif isinstance(parent_tags, tuple): + inherited_tags = list(parent_tags) + + if inherited_trace_name is not None: + _set_trace_attribute(observation, _TRACE_NAME_ATTR, inherited_trace_name) + if inherited_session_id is not None: + _set_trace_attribute(observation, _TRACE_SESSION_ID_ATTR, inherited_session_id) + if inherited_user_id is not None: + _set_trace_attribute(observation, _TRACE_USER_ID_ATTR, inherited_user_id) + if inherited_tags is not None: + _set_trace_attribute(observation, _TRACE_TAGS_ATTR, inherited_tags) + return observation + + def initialize() -> None: """Initialize Langfuse client once (no-op when unavailable).""" global _client, _initialized @@ -140,26 +222,48 @@ def create_trace( ) try: # Old SDKs may expose .trace(), newer SDKs are OTEL-native and expose - # .start_span() / .start_observation(). + # .start_observation(). if hasattr(client, "trace"): return client.trace(**payload) - trace_obs = client.start_span( - name=name, - input=payload.get("input"), - metadata=payload.get("metadata"), - ) - # Best-effort enrich current trace with session/user dimensions. - try: - client.update_current_trace( + if hasattr(client, "start_observation"): + trace_metadata = dict(payload.get("metadata") or {}) + if session_id is not None: + trace_metadata.setdefault("session_id", session_id) + if user_id is not None: + trace_metadata.setdefault("user_id", user_id) + if tags is not None: + trace_metadata.setdefault("tags", tags) + trace_obs = client.start_observation( + name=name, + as_type="span", + input=payload.get("input"), + metadata=trace_metadata or None, + ) + return _propagate_trace_dimensions( + trace_obs, + trace_name=name, session_id=session_id, user_id=user_id, tags=tags, + ) + if hasattr(client, "start_span"): + trace_obs = client.start_span( + name=name, input=payload.get("input"), metadata=payload.get("metadata"), ) - except Exception: - pass - return trace_obs + # Best-effort enrich current trace with session/user dimensions. + try: + client.update_current_trace( + session_id=session_id, + user_id=user_id, + tags=tags, + input=payload.get("input"), + metadata=payload.get("metadata"), + ) + except Exception: + pass + return trace_obs except Exception as exc: log.warn("langfuse.trace_failed", {"error": str(exc), "name": name}) return _NoopObservation("trace") @@ -189,15 +293,21 @@ def create_generation( ) try: if parent and hasattr(parent, "generation"): - return parent.generation(**payload) + return _propagate_trace_dimensions(parent.generation(**payload), parent=parent) if parent and hasattr(parent, "start_generation"): - return parent.start_generation(**payload) + return _propagate_trace_dimensions(parent.start_generation(**payload), parent=parent) if parent and hasattr(parent, "start_observation"): - return parent.start_observation(as_type="generation", **payload) + return _propagate_trace_dimensions( + parent.start_observation(as_type="generation", **payload), + parent=parent, + ) if hasattr(client, "start_generation"): - return client.start_generation(**payload) + return _propagate_trace_dimensions(client.start_generation(**payload), parent=parent) if hasattr(client, "start_observation"): - return client.start_observation(as_type="generation", **payload) + return _propagate_trace_dimensions( + client.start_observation(as_type="generation", **payload), + parent=parent, + ) except Exception as exc: log.warn("langfuse.generation_failed", {"error": str(exc), "name": name}) return _NoopObservation("generation") @@ -225,15 +335,21 @@ def create_span( ) try: if parent and hasattr(parent, "span"): - return parent.span(**payload) + return _propagate_trace_dimensions(parent.span(**payload), parent=parent) if parent and hasattr(parent, "start_span"): - return parent.start_span(**payload) + return _propagate_trace_dimensions(parent.start_span(**payload), parent=parent) if parent and hasattr(parent, "start_observation"): - return parent.start_observation(as_type="span", **payload) + return _propagate_trace_dimensions( + parent.start_observation(as_type="span", **payload), + parent=parent, + ) if hasattr(client, "start_span"): - return client.start_span(**payload) + return _propagate_trace_dimensions(client.start_span(**payload), parent=parent) if hasattr(client, "start_observation"): - return client.start_observation(as_type="span", **payload) + return _propagate_trace_dimensions( + client.start_observation(as_type="span", **payload), + parent=parent, + ) except Exception as exc: log.warn("langfuse.span_failed", {"error": str(exc), "name": name}) return _NoopObservation("span") @@ -267,20 +383,32 @@ def end_observation( observation.end(**payload) return except TypeError: - try: - observation.end() - return - except Exception: - pass + pass except Exception: pass try: if hasattr(observation, "update"): - observation.update(**payload) + update_payload = dict(payload) + if usage: + update_payload["usage_details"] = usage + try: + observation.update(**update_payload) + except TypeError: + fallback_payload = dict(payload) + if usage: + fallback_payload["usage"] = usage + fallback_payload.pop("usage_details", None) + observation.update(**fallback_payload) except Exception as exc: log.debug("langfuse.end_fallback_update_failed", {"error": str(exc)}) + try: + if hasattr(observation, "end"): + observation.end() + except Exception as exc: + log.debug("langfuse.end_fallback_end_failed", {"error": str(exc)}) + def is_active() -> bool: """Return True when Langfuse is initialized and has a live client.""" diff --git a/flocks/workflow/tool_context.py b/flocks/workflow/tool_context.py index e06ad3b42..b803f9810 100644 --- a/flocks/workflow/tool_context.py +++ b/flocks/workflow/tool_context.py @@ -127,23 +127,7 @@ async def cleanup_workflow_tool_context(tool_context: Optional[ToolContext]) -> if extra.get("workflow_child_session_created") is True: return False - await Session.delete(session.project_id, session.id) - - # Session.delete is intentionally a soft delete. Trigger parents without - # child tasks are implementation details, so remove their residual rows - # to keep high-frequency trigger traffic storage-bounded. - from flocks.storage.storage import Storage - - await Storage.delete(f"session:{session.project_id}:{session.id}") - await Storage.delete(f"message:{session.id}") - await Storage.delete(f"todo:{session.id}") - await Storage.delete(f"goal:{session.id}") - await Storage.delete(f"session_diff:{session.id}") - await Storage.clear(prefix=f"message_diff:{session.id}:") - await Storage.clear(prefix=f"system_prompts:{session.id}:") - - Message.invalidate_cache(session.id) - return True + return await Session.delete(session.project_id, session.id) except Exception as exc: log.warning( "workflow.tool_context.cleanup_failed", diff --git a/pyproject.toml b/pyproject.toml index 939c3b233..384249fa7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "flocks" -version = "v2026.7.23" +version = "v2026.7.29" description = "AI-Native SecOps platform with multi-agent collaboration" authors = [ {name = "Flocks Team", email = "team@example.com"} diff --git a/scripts/dev.sh b/scripts/dev.sh index 80474efb6..03a4fab5b 100644 --- a/scripts/dev.sh +++ b/scripts/dev.sh @@ -25,7 +25,6 @@ if [ "${BACKEND_ACCESS_HOST}" = "0.0.0.0" ] || [ "${BACKEND_ACCESS_HOST}" = "::" fi BACKEND_BASE_URL="http://${BACKEND_ACCESS_HOST}:${BACKEND_PORT}" -BACKEND_WS_URL="ws://${BACKEND_ACCESS_HOST}:${BACKEND_PORT}" BACKEND_PID="" BACKEND_PGID="" CLEANUP_DONE=0 @@ -223,8 +222,7 @@ start_frontend() { echo -e "${GREEN}🎨 启动前端服务: http://${FRONTEND_HOST}:${FRONTEND_PORT}${NC}" cd "${PROJECT_ROOT}/webui" - VITE_API_BASE_URL="${BACKEND_BASE_URL}" \ - VITE_WS_BASE_URL="${BACKEND_WS_URL}" \ + FLOCKS_API_PROXY_TARGET="${BACKEND_BASE_URL}" \ npm run dev -- --host "${FRONTEND_HOST}" --port "${FRONTEND_PORT}" } diff --git a/scripts/validate_flockshub.py b/scripts/validate_flockshub.py index 7068f7cdc..ffdd3208e 100644 --- a/scripts/validate_flockshub.py +++ b/scripts/validate_flockshub.py @@ -62,6 +62,8 @@ def main() -> int: manifest = load_json(manifest_path) if manifest.get("id") != plugin_id or manifest.get("type") != plugin_type: fail(f"Manifest id/type mismatch: {manifest_rel}") + if manifest.get("version") != entry.get("version"): + fail(f"Index/manifest version mismatch: {manifest_rel}") if manifest.get("category") not in categories: fail(f"Unknown category in {manifest_rel}: {manifest.get('category')}") unknown_tags = set(manifest.get("tags", [])) - tags @@ -75,6 +77,11 @@ def main() -> int: fail(f"Unknown risk level in {manifest_rel}: {risk_level}") package_dir = manifest_path.parent + workspace_path = package_dir / "workspace.json" + if plugin_type == "webui" and workspace_path.is_file(): + workspace = load_json(workspace_path) + if workspace.get("version") != manifest.get("version"): + fail(f"WebUI workspace/manifest version mismatch: {manifest_rel}") for entrypoint in manifest.get("entrypoints", []): ensure_relative(entrypoint) if not (package_dir / entrypoint).exists(): diff --git a/tests/agent/test_unified_session_loop.py b/tests/agent/test_unified_session_loop.py index 255882ccd..44aaa8fa0 100644 --- a/tests/agent/test_unified_session_loop.py +++ b/tests/agent/test_unified_session_loop.py @@ -8,6 +8,8 @@ 4. _resolve_model implements 5-level priority correctly """ +import asyncio + import pytest from unittest.mock import AsyncMock, MagicMock, patch from dataclasses import dataclass @@ -279,7 +281,226 @@ async def test_process_session_message_pins_explicit_request_model(self, monkeyp provider="anthropic", model="claude-sonnet-4-5", model_pinned=True, + model_auto=False, + ) + + @pytest.mark.asyncio + @pytest.mark.parametrize("category", ["user", "entity-config", "workflow"]) + async def test_webui_auto_uses_default_primary_and_returns_actual_model( + self, + monkeypatch, + category, + ): + """Supported WebUI chats use Auto and report the recovered model.""" + from types import SimpleNamespace + + from flocks.server.routes import session as session_routes + from flocks.session.session_loop import LoopResult, SessionLoop + + request = session_routes.PromptRequest( + parts=[{"type": "text", "text": "hello"}], ) + session = SimpleNamespace( + id="ses_auto", + project_id="proj", + directory="/tmp/project", + agent="rex", + provider="stale", + model="stale-model", + model_pinned=False, + model_auto=True, + category=category, + ) + agent = SimpleNamespace( + name="rex", + model={"providerID": "agent-provider", "modelID": "agent-model"}, + ) + fallback_message = SimpleNamespace( + id="msg_assistant", + providerID="fallback", + modelID="fallback-model", + finish="stop", + tokens=None, + ) + loop_run = AsyncMock(return_value=LoopResult( + action="stop", + last_message=fallback_message, + provider_id="fallback", + model_id="fallback-model", + )) + message_create = AsyncMock(return_value=SimpleNamespace(id="msg_user")) + context_usage_update = AsyncMock() + title_generation = AsyncMock() + + monkeypatch.setattr(session_routes, "_require_agent_usable_for_chat", AsyncMock()) + monkeypatch.setattr( + "flocks.agent.registry.Agent.default_agent", + AsyncMock(return_value="rex"), + ) + monkeypatch.setattr("flocks.agent.registry.Agent.get", AsyncMock(return_value=agent)) + monkeypatch.setattr( + "flocks.config.config.Config.resolve_default_llm", + AsyncMock(return_value={ + "provider_id": "primary", + "model_id": "primary-model", + }), + ) + monkeypatch.setattr( + "flocks.config.config.Config.get", + AsyncMock(return_value=SimpleNamespace()), + ) + monkeypatch.setattr( + SessionLoop, + "validate_runtime_model", + AsyncMock(return_value=(True, "available")), + ) + monkeypatch.setattr(SessionLoop, "run", loop_run) + monkeypatch.setattr("flocks.provider.provider.Provider._ensure_initialized", lambda: None) + monkeypatch.setattr("flocks.provider.provider.Provider.apply_config", AsyncMock()) + monkeypatch.setattr("flocks.provider.provider.Provider.get", lambda _provider_id: object()) + monkeypatch.setattr("flocks.tool.registry.ToolRegistry.init", lambda: None) + monkeypatch.setattr( + "flocks.session.lifecycle.revert.SessionRevert.cleanup", + AsyncMock(), + ) + monkeypatch.setattr("flocks.session.message.Message.create", message_create) + monkeypatch.setattr( + "flocks.session.message.Message.get_text_content", + AsyncMock(return_value="recovered"), + ) + monkeypatch.setattr("flocks.session.message.Message.parts", AsyncMock(return_value=[])) + monkeypatch.setattr("flocks.server.routes.event.publish_event", AsyncMock()) + monkeypatch.setattr( + session_routes, + "_publish_context_usage_update", + context_usage_update, + ) + monkeypatch.setattr( + "flocks.session.lifecycle.title.SessionTitle.generate_title_after_first_message", + title_generation, + ) + + response = await session_routes._process_session_message( + session.id, + session, + request, + session.directory, + ) + + assert message_create.await_args.kwargs["model"] == { + "providerID": "primary", + "modelID": "primary-model", + } + assert loop_run.await_args.kwargs["auto_failover"] is True + assert loop_run.await_args.kwargs["provider_id"] == "primary" + assert response["info"]["providerID"] == "fallback" + assert response["info"]["modelID"] == "fallback-model" + assert context_usage_update.await_args.kwargs["provider_id"] == "fallback" + assert context_usage_update.await_args.kwargs["model_id"] == "fallback-model" + await asyncio.sleep(0) + assert title_generation.await_args.kwargs["provider_id"] == "fallback" + assert title_generation.await_args.kwargs["model_id"] == "fallback-model" + + @pytest.mark.asyncio + async def test_unsupported_session_ignores_auto_flag(self, monkeypatch): + """A corrupt task Auto flag cannot activate WebUI Auto.""" + from types import SimpleNamespace + + from flocks.server.routes import session as session_routes + from flocks.session.session_loop import SessionLoop + + request = session_routes.PromptRequest( + parts=[{"type": "text", "text": "task input"}], + noReply=True, + ) + session = SimpleNamespace( + id="ses_task", + project_id="proj", + directory="/tmp/project", + agent="rex", + provider="direct", + model="direct-model", + model_pinned=True, + model_auto=True, + category="task", + ) + agent = SimpleNamespace(name="rex", model=None) + resolve = AsyncMock(return_value=("direct", "direct-model", "session")) + default_llm = AsyncMock() + validate = AsyncMock() + message_create = AsyncMock(return_value=SimpleNamespace(id="msg_user")) + + monkeypatch.setattr( + session_routes, + "_require_agent_usable_for_chat", + AsyncMock(), + ) + monkeypatch.setattr( + "flocks.agent.registry.Agent.default_agent", + AsyncMock(return_value="rex"), + ) + monkeypatch.setattr( + "flocks.agent.registry.Agent.get", + AsyncMock(return_value=agent), + ) + monkeypatch.setattr(session_routes, "_resolve_model", resolve) + monkeypatch.setattr( + "flocks.config.config.Config.resolve_default_llm", + default_llm, + ) + monkeypatch.setattr( + "flocks.config.config.Config.get", + AsyncMock(return_value=SimpleNamespace()), + ) + monkeypatch.setattr(SessionLoop, "validate_runtime_model", validate) + monkeypatch.setattr( + "flocks.provider.provider.Provider._ensure_initialized", + lambda: None, + ) + monkeypatch.setattr( + "flocks.provider.provider.Provider.apply_config", + AsyncMock(), + ) + monkeypatch.setattr( + "flocks.provider.provider.Provider.get", + lambda _provider_id: object(), + ) + monkeypatch.setattr("flocks.tool.registry.ToolRegistry.init", lambda: None) + monkeypatch.setattr( + "flocks.session.lifecycle.revert.SessionRevert.cleanup", + AsyncMock(), + ) + monkeypatch.setattr( + "flocks.session.message.Message.create", + message_create, + ) + monkeypatch.setattr( + "flocks.server.routes.event.publish_event", + AsyncMock(), + ) + context_usage = AsyncMock() + monkeypatch.setattr( + session_routes, + "_publish_context_usage_update", + context_usage, + ) + + await session_routes._process_session_message( + session.id, + session, + request, + session.directory, + ) + + resolve.assert_awaited_once() + default_llm.assert_not_awaited() + validate.assert_not_awaited() + assert message_create.await_args.kwargs["model"] == { + "providerID": "direct", + "modelID": "direct-model", + } + assert context_usage.await_args.kwargs["provider_id"] == "direct" + assert context_usage.await_args.kwargs["model_id"] == "direct-model" @pytest.mark.asyncio async def test_display_text_does_not_replace_model_prompt(self, monkeypatch): diff --git a/tests/channel/test_channel.py b/tests/channel/test_channel.py index 34a7830cf..5b6c57430 100644 --- a/tests/channel/test_channel.py +++ b/tests/channel/test_channel.py @@ -342,6 +342,62 @@ def test_empty_input(self): class TestFeishuNativeCommands: + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("category", "model_auto", "expected"), + [ + ("user", True, True), + ("user", False, False), + ("task", True, False), + ], + ) + async def test_im_loop_uses_supported_persisted_auto_mode( + self, + monkeypatch, + category, + model_auto, + expected, + ): + from flocks.channel.inbound.dispatcher import InboundDispatcher + from flocks.channel.inbound.session_binding import SessionBinding + + binding = SessionBinding( + channel_id="slack", + account_id="T1", + chat_id="C123", + chat_type=ChatType.CHANNEL, + thread_id=None, + session_id="session_1", + agent_id="rex", + created_at=0, + last_message_at=0, + ) + monkeypatch.setattr( + "flocks.session.session.Session.get_by_id", + AsyncMock( + return_value=SimpleNamespace( + id="session_1", + category=category, + model_auto=model_auto, + ) + ), + ) + run_mock = AsyncMock(return_value=SimpleNamespace(last_message=None)) + monkeypatch.setattr( + "flocks.session.session_loop.SessionLoop.run", + run_mock, + ) + + loop_callbacks = SimpleNamespace() + await InboundDispatcher._run_session_loop(binding, loop_callbacks) + + run_mock.assert_awaited_once_with( + session_id="session_1", + agent_name="rex", + callbacks=loop_callbacks, + auto_failover=expected, + ) + @pytest.mark.asyncio async def test_status_command_reports_session_state(self, monkeypatch): from flocks.channel.inbound.dispatcher import InboundDispatcher @@ -496,6 +552,7 @@ async def fake_update(project_id, session_id, **updates): "provider": "anthropic", "model": "claude-sonnet-4-20250514", "model_pinned": True, + "model_auto": False, }, ) ] @@ -599,6 +656,90 @@ async def fake_deliver(ctx, session_id=None): assert update_mock.await_args.args == ("channel", "session_old") assert update_mock.await_args.kwargs["status"] == "archived" + @pytest.mark.asyncio + async def test_new_command_inherits_auto_model_mode(self, monkeypatch): + from flocks.channel.inbound.dispatcher import InboundDispatcher + from flocks.channel.inbound.session_binding import SessionBinding + + dispatcher = InboundDispatcher() + dispatcher._trigger_command_hook = AsyncMock() + dispatcher.binding_service.rebind = AsyncMock( + return_value=SessionBinding( + channel_id="wecom", + account_id="default", + chat_id="room_1", + chat_type=ChatType.DIRECT, + thread_id=None, + session_id="session_new", + agent_id="rex", + created_at=0, + last_message_at=0, + ) + ) + binding = SessionBinding( + channel_id="wecom", + account_id="default", + chat_id="room_1", + chat_type=ChatType.DIRECT, + thread_id=None, + session_id="session_old", + agent_id="rex", + created_at=0, + last_message_at=0, + ) + msg = InboundMessage( + channel_id="wecom", + account_id="default", + message_id="msg_1", + sender_id="user_1", + chat_id="room_1", + chat_type=ChatType.DIRECT, + text="/new", + mention_text="/new", + ) + monkeypatch.setattr( + "flocks.channel.outbound.deliver.OutboundDelivery.deliver", + AsyncMock(), + ) + monkeypatch.setattr( + "flocks.session.session.Session.get_by_id", + AsyncMock( + return_value=SimpleNamespace( + id="session_old", + project_id="channel", + directory="/tmp/project", + agent="rex", + category="user", + model_auto=True, + model_pinned=False, + provider=None, + model=None, + ) + ), + ) + create_mock = AsyncMock( + return_value=SimpleNamespace(id="session_new", agent="rex") + ) + monkeypatch.setattr("flocks.session.session.Session.create", create_mock) + monkeypatch.setattr( + "flocks.session.session.Session.update", + AsyncMock(return_value=None), + ) + + await dispatcher._handle_session_command( + binding=binding, + msg=msg, + callbacks=dispatcher._build_callbacks(binding, msg), + scope_override=None, + channel_config=ChannelConfig(enabled=True), + ) + + create_kwargs = create_mock.await_args.kwargs + assert create_kwargs["model_auto"] is True + assert create_kwargs["model_pinned"] is False + assert "provider" not in create_kwargs + assert "model" not in create_kwargs + @pytest.mark.asyncio async def test_new_command_with_args_sends_args_to_new_session(self, monkeypatch): from flocks.channel.inbound.dispatcher import InboundDispatcher diff --git a/tests/channel/test_session_binding.py b/tests/channel/test_session_binding.py index 45e7ad366..0a09d8e43 100644 --- a/tests/channel/test_session_binding.py +++ b/tests/channel/test_session_binding.py @@ -3,7 +3,8 @@ import pytest -from flocks.channel.inbound.session_binding import SessionBindingService +from flocks.channel.base import ChatType, InboundMessage +from flocks.channel.inbound.session_binding import SessionBinding, SessionBindingService @pytest.mark.asyncio @@ -25,3 +26,52 @@ async def test_latest_active_user_binding_returns_none_when_channel_is_ambiguous result = await service.latest_active_user_binding(channel_id="wecom") assert result is None + + +@pytest.mark.asyncio +async def test_resolve_or_create_replaces_archived_binding() -> None: + service = SessionBindingService() + existing = SessionBinding( + channel_id="slack", + account_id="default", + chat_id="chat-1", + chat_type=ChatType.DIRECT, + thread_id=None, + session_id="ses_archived", + agent_id="rex", + created_at=1, + last_message_at=2, + ) + archived = SimpleNamespace( + status="archived", + owner_user_id="usr_1", + owner_username="alice", + ) + service._find_binding = AsyncMock(return_value=existing) + service.unbind = AsyncMock() + service._create_session = AsyncMock(return_value="ses_replacement") + service._insert = AsyncMock() + msg = InboundMessage( + channel_id="slack", + account_id="default", + message_id="msg-1", + sender_id="alice", + chat_id="chat-1", + chat_type=ChatType.DIRECT, + text="hello again", + ) + + with patch( + "flocks.session.session.Session.get_by_id_unfiltered", + AsyncMock(return_value=archived), + ): + binding = await service.resolve_or_create(msg, default_agent="rex") + + assert binding.session_id == "ses_replacement" + service.unbind.assert_awaited_once_with("ses_archived") + service._create_session.assert_awaited_once_with( + msg, + default_agent="rex", + directory=None, + source_session=archived, + ) diff --git a/tests/channel/test_unified_prompt_context.py b/tests/channel/test_unified_prompt_context.py index ac09461b3..b1c105b41 100644 --- a/tests/channel/test_unified_prompt_context.py +++ b/tests/channel/test_unified_prompt_context.py @@ -257,7 +257,7 @@ async def _fake_create(**kwargs): assert captured["owner_username"] == "admin" @pytest.mark.asyncio - async def test_create_session_stays_ownerless_without_local_accounts(self): + async def test_create_session_uses_system_owner_without_local_accounts(self): captured = {} class _StubSession: @@ -277,8 +277,8 @@ async def _fake_create(**kwargs): ) assert sid == "ses_ownerless" - assert "owner_user_id" not in captured - assert "owner_username" not in captured + assert captured["owner_user_id"] == "api-token-service" + assert captured["owner_username"] == "api-token-service" list_users.assert_not_awaited() @pytest.mark.asyncio @@ -407,6 +407,15 @@ async def _boom(_session, _p, _m): class TestAppendUserMessagePersistsModel: + @pytest.fixture(autouse=True) + def _active_session_write(self, monkeypatch): + from flocks.session.session import Session + + async def run_active_write(_cls, _session_id, operation, **_kwargs): + return await operation() + + monkeypatch.setattr(Session, "run_active_write", classmethod(run_active_write)) + @pytest.mark.asyncio async def test_model_is_passed_to_message_create(self): captured = {} @@ -498,6 +507,29 @@ async def _fake_create(**kwargs): assert "agent" not in captured +@pytest.mark.asyncio +async def test_append_user_message_rejects_inactive_session_before_persisting(monkeypatch): + from flocks.session.message import Message + from flocks.session.session import Session, SessionInactiveError + + create = AsyncMock() + monkeypatch.setattr(Message, "create", create) + monkeypatch.setattr( + Session, + "run_active_write", + AsyncMock(side_effect=SessionInactiveError("archived")), + ) + + with pytest.raises(SessionInactiveError, match="archived"): + await InboundDispatcher._append_user_message( + "ses_archived", + "must not be stored", + _msg(), + ) + + create.assert_not_awaited() + + # --------------------------------------------------------------------------- # 4. Legacy channel_bindings.agent_id one-shot normalisation # --------------------------------------------------------------------------- diff --git a/tests/cli/test_service_manager.py b/tests/cli/test_service_manager.py index 42f3736ad..3ab93e061 100644 --- a/tests/cli/test_service_manager.py +++ b/tests/cli/test_service_manager.py @@ -1,4 +1,5 @@ import json +import os import shutil import subprocess import sys @@ -17,6 +18,8 @@ make_short_runtime_root, ) +_REAL_ENSURE_WEBUI_DIST = service_manager._ensure_webui_dist + class DummyConsole: def __init__(self) -> None: @@ -520,6 +523,87 @@ def test_run_windows_netstat_handles_missing_stdout(monkeypatch) -> None: assert service_manager._run_windows_netstat(5173) == "" +def test_run_windows_netstat_decodes_with_replacement(monkeypatch) -> None: + """netstat output is decoded as utf-8/replace so GBK console output on + Chinese Windows never raises UnicodeDecodeError in the reader thread.""" + seen: dict[str, object] = {} + + def fake_run(*_args, **kwargs): + seen.update(kwargs) + return SimpleNamespace(returncode=0, stdout=" TCP 127.0.0.1:5173 0.0.0.0:0 LISTENING 42\n") + + monkeypatch.setattr(service_manager.subprocess, "run", fake_run) + + assert "42" in service_manager._run_windows_netstat(5173) + assert seen.get("encoding") == "utf-8" + assert seen.get("errors") == "replace" + + +def test_windows_tasklist_process_name_decodes_with_replacement(monkeypatch) -> None: + monkeypatch.setattr(service_manager.sys, "platform", "win32") + seen: dict[str, object] = {} + + def fake_run(*_args, **kwargs): + seen.update(kwargs) + return SimpleNamespace(returncode=0, stdout='"python.exe","9436"\n') + + monkeypatch.setattr(service_manager.subprocess, "run", fake_run) + + assert service_manager._windows_tasklist_process_name(9436) == "python.exe" + assert seen.get("encoding") == "utf-8" + assert seen.get("errors") == "replace" + + +def test_process_list_pids_windows_decodes_with_replacement(monkeypatch) -> None: + monkeypatch.setattr(service_manager.sys, "platform", "win32") + seen: dict[str, object] = {} + + def fake_run(*_args, **kwargs): + seen.update(kwargs) + return SimpleNamespace(returncode=0, stdout="9436\n36056\n") + + monkeypatch.setattr(service_manager.subprocess, "run", fake_run) + + assert service_manager._process_list_pids() == [9436, 36056] + assert seen.get("encoding") == "utf-8" + assert seen.get("errors") == "replace" + + +def test_get_node_major_version_decodes_with_replacement(monkeypatch) -> None: + monkeypatch.setattr(service_manager, "resolve_node_executable", lambda: "node") + seen: dict[str, object] = {} + + def fake_run(*_args, **kwargs): + seen.update(kwargs) + return SimpleNamespace(returncode=0, stdout="v20.11.1\n") + + monkeypatch.setattr(service_manager.subprocess, "run", fake_run) + + assert service_manager.get_node_major_version() == 20 + assert seen.get("encoding") == "utf-8" + assert seen.get("errors") == "replace" + + +def test_windows_process_probes_survive_gbk_bytes(monkeypatch) -> None: + """End-to-end: real GBK-encoded bytes flowing through the decode path + must not raise UnicodeDecodeError (the restart-time reader-thread crash).""" + monkeypatch.setattr(service_manager.sys, "platform", "win32") + # 0xbb is the byte that crashed utf-8 decode in the original bug report. + gbk_stdout = "映像名称: python.exe 拒绝访问".encode("gbk").decode("utf-8", errors="replace") + + def fake_run(*_args, **kwargs): + assert kwargs.get("encoding") == "utf-8" + assert kwargs.get("errors") == "replace" + return SimpleNamespace(returncode=0, stdout=gbk_stdout) + + monkeypatch.setattr(service_manager.subprocess, "run", fake_run) + + # None of these should raise, even with mojibake stdout. + service_manager._windows_tasklist_process_name(123) + service_manager._run_windows_netstat(5173) + service_manager._process_list_pids() + + def test_port_owner_pids_warns_when_no_tool_found(monkeypatch) -> None: monkeypatch.setattr(service_manager.sys, "platform", "linux") monkeypatch.setattr(service_manager, "which", lambda _name: None) @@ -1112,6 +1196,40 @@ def test_restart_all_stops_then_starts_daemon(monkeypatch) -> None: assert call_order == ["stop", "start"] +def test_restart_server_requests_backend_restart(monkeypatch) -> None: + calls: list[str] = [] + console = DummyConsole() + paths = _make_runtime_paths(Path("/tmp/flocks-test")) + status = _supervisor_status(_supervisor_status_payload()) + + monkeypatch.setattr(service_manager, "ensure_runtime_dirs", lambda: paths) + monkeypatch.setattr(service_manager, "supervisor_is_running", lambda _paths: True) + monkeypatch.setattr( + service_manager, + "request_restart_backend", + lambda **_kwargs: calls.append("backend") or status, + ) + monkeypatch.setattr( + service_manager, + "_print_status_payload", + lambda *_args, **_kwargs: calls.append("status"), + ) + + service_manager.restart_server(console) + + assert calls == ["backend", "status"] + + +def test_restart_server_requires_running_supervisor(monkeypatch) -> None: + paths = _make_runtime_paths(Path("/tmp/flocks-test")) + + monkeypatch.setattr(service_manager, "ensure_runtime_dirs", lambda: paths) + monkeypatch.setattr(service_manager, "supervisor_is_running", lambda _paths: False) + + with pytest.raises(service_manager.ServiceError, match="请执行 `flocks restart` 进行全量重启"): + service_manager.restart_server(DummyConsole()) + + def test_start_all_without_stop_starts_supervisor_daemon(monkeypatch, tmp_path: Path) -> None: paths = _make_runtime_paths(tmp_path) calls: list[str] = [] @@ -1518,7 +1636,10 @@ def _fake_process(pid: int, args: list[str] | None = None, returncode: int | Non return SimpleNamespace(pid=pid, args=args or [str(pid)], returncode=returncode, poll=lambda: returncode) -def test_supervisor_recovers_backend_when_port_disappears(monkeypatch, tmp_path: Path) -> None: +def test_supervisor_restarts_backend_after_tenth_consecutive_port_failure( + monkeypatch, + tmp_path: Path, +) -> None: paths = _make_runtime_paths(tmp_path) calls: list[str] = [] monkeypatch.setattr(service_manager, "ensure_runtime_dirs", lambda: paths) @@ -1535,12 +1656,90 @@ def test_supervisor_recovers_backend_when_port_disappears(monkeypatch, tmp_path: lambda *_args, **_kwargs: calls.append("start:backend") or _fake_process(333, ["backend-new"]), ) - daemon.tick() + assert daemon.interval == 30.0 + assert daemon.failure_threshold == 10 + + for expected_failure_count in range(1, 10): + daemon.tick() + assert calls == [] + assert daemon.backend.state == "degraded" + assert daemon.backend.health_failure_count == expected_failure_count + daemon.tick() assert calls == ["stop:后端", "start:backend"] assert daemon.backend.pid == 333 +def test_supervisor_resets_backend_port_failures_after_recovery( + monkeypatch, + tmp_path: Path, +) -> None: + paths = _make_runtime_paths(tmp_path) + calls: list[str] = [] + port_states = iter([False] * 9 + [True] + [False] * 9) + monkeypatch.setattr(service_manager, "ensure_runtime_dirs", lambda: paths) + daemon = service_supervisor.SupervisorDaemon(service_manager.ServiceConfig(backend_port=9995)) + daemon.paths = paths + daemon.backend.process = _fake_process(111, ["backend"]) + + monkeypatch.setattr( + service_process, + "tcp_port_accepts_connections", + lambda *_args: next(port_states), + ) + monkeypatch.setattr( + service_manager, + "_terminate_process", + lambda *_args, **_kwargs: calls.append("stop"), + ) + + for _ in range(9): + daemon.tick() + + assert calls == [] + assert daemon.backend.health_failure_count == 9 + + daemon.tick() + + assert daemon.backend.state == "healthy" + assert daemon.backend.health_failure_count == 0 + + for _ in range(9): + daemon.tick() + + assert calls == [] + assert daemon.backend.state == "degraded" + assert daemon.backend.health_failure_count == 9 + + +def test_supervisor_restarts_backend_on_first_probe_after_process_exits( + monkeypatch, + tmp_path: Path, +) -> None: + paths = _make_runtime_paths(tmp_path) + calls: list[str] = [] + monkeypatch.setattr(service_manager, "ensure_runtime_dirs", lambda: paths) + daemon = service_supervisor.SupervisorDaemon(service_manager.ServiceConfig(backend_port=9995)) + daemon.paths = paths + daemon.backend.process = _fake_process(111, ["backend"], returncode=7) + + monkeypatch.setattr( + service_manager, + "_terminate_process", + lambda *_args, **_kwargs: calls.append("stop"), + ) + monkeypatch.setattr( + service_manager, + "_start_backend_process", + lambda *_args, **_kwargs: calls.append("start") or _fake_process(333, ["backend-new"]), + ) + + daemon.tick() + + assert calls == ["stop", "start"] + assert daemon.backend.pid == 333 + + def test_supervisor_liveness_probe_keeps_backend_healthy(monkeypatch, tmp_path: Path) -> None: """Liveness-only probe: process alive + TCP port open = healthy, no HTTP check needed.""" paths = _make_runtime_paths(tmp_path) @@ -1655,6 +1854,108 @@ def fake_run(command, **_kwargs): assert build_calls[0][0] == r"C:\Users\flocks\AppData\Local\Programs\Flocks\tools\node\npm.cmd" +def test_webui_needs_build_when_source_is_newer_than_dist(tmp_path: Path) -> None: + webui_dir = tmp_path / "webui" + source_path = webui_dir / "src" / "main.tsx" + index_path = webui_dir / "dist" / "index.html" + source_path.parent.mkdir(parents=True) + index_path.parent.mkdir() + source_path.write_text("before", encoding="utf-8") + index_path.write_text("", encoding="utf-8") + built_at = index_path.stat().st_mtime_ns + os.utime(source_path, ns=(built_at + 1_000_000_000, built_at + 1_000_000_000)) + + assert service_manager._webui_needs_build(webui_dir) is True + + +def test_webui_needs_build_ignores_generated_directories(tmp_path: Path) -> None: + webui_dir = tmp_path / "webui" + source_path = webui_dir / "src" / "main.tsx" + index_path = webui_dir / "dist" / "index.html" + generated_path = webui_dir / "node_modules" / "package" / "index.js" + source_path.parent.mkdir(parents=True) + index_path.parent.mkdir() + generated_path.parent.mkdir(parents=True) + source_path.write_text("source", encoding="utf-8") + index_path.write_text("", encoding="utf-8") + generated_path.write_text("generated", encoding="utf-8") + built_at = max(path.stat().st_mtime_ns for path in (webui_dir, source_path.parent, source_path)) + os.utime(index_path, ns=(built_at + 1_000_000_000, built_at + 1_000_000_000)) + os.utime(generated_path, ns=(built_at + 2_000_000_000, built_at + 2_000_000_000)) + + assert service_manager._webui_needs_build(webui_dir) is False + + +def test_webui_needs_build_when_source_is_deleted(tmp_path: Path) -> None: + webui_dir = tmp_path / "webui" + source_dir = webui_dir / "src" + source_path = source_dir / "main.tsx" + index_path = webui_dir / "dist" / "index.html" + source_dir.mkdir(parents=True) + index_path.parent.mkdir() + source_path.write_text("source", encoding="utf-8") + index_path.write_text("", encoding="utf-8") + built_at = max(path.stat().st_mtime_ns for path in (webui_dir, source_dir, source_path)) + os.utime(index_path, ns=(built_at + 1_000_000_000, built_at + 1_000_000_000)) + source_path.unlink() + os.utime(source_dir, ns=(built_at + 2_000_000_000, built_at + 2_000_000_000)) + + assert service_manager._webui_needs_build(webui_dir) is True + + +def test_ensure_webui_dist_rebuilds_when_source_is_newer(monkeypatch, tmp_path: Path) -> None: + from flocks.server import static_webui + + webui_dir = tmp_path / "webui" + source_path = webui_dir / "src" / "main.tsx" + index_path = webui_dir / "dist" / "index.html" + source_path.parent.mkdir(parents=True) + index_path.parent.mkdir() + (webui_dir / "package.json").write_text("{}", encoding="utf-8") + source_path.write_text("before", encoding="utf-8") + index_path.write_text("", encoding="utf-8") + built_at = index_path.stat().st_mtime_ns + os.utime(source_path, ns=(built_at + 1_000_000_000, built_at + 1_000_000_000)) + builds: list[Path] = [] + monkeypatch.setattr(static_webui, "ensure_webui_dist_dir", lambda: index_path.parent.resolve()) + monkeypatch.setattr( + service_manager, + "_build_webui_dist", + lambda root, _config, _console: builds.append(root), + ) + + _REAL_ENSURE_WEBUI_DIST(tmp_path, service_manager.ServiceConfig(), DummyConsole()) + + assert builds == [tmp_path] + + +def test_ensure_webui_dist_skips_stale_source_when_requested(monkeypatch, tmp_path: Path) -> None: + from flocks.server import static_webui + + webui_dir = tmp_path / "webui" + source_path = webui_dir / "src" / "main.tsx" + index_path = webui_dir / "dist" / "index.html" + source_path.parent.mkdir(parents=True) + index_path.parent.mkdir() + (webui_dir / "package.json").write_text("{}", encoding="utf-8") + source_path.write_text("before", encoding="utf-8") + index_path.write_text("", encoding="utf-8") + built_at = index_path.stat().st_mtime_ns + os.utime(source_path, ns=(built_at + 1_000_000_000, built_at + 1_000_000_000)) + builds: list[Path] = [] + monkeypatch.setattr(static_webui, "ensure_webui_dist_dir", lambda: index_path.parent.resolve()) + monkeypatch.setattr( + service_manager, + "_build_webui_dist", + lambda root, _config, _console: builds.append(root), + ) + + config = service_manager.ServiceConfig(skip_frontend_build=True) + _REAL_ENSURE_WEBUI_DIST(tmp_path, config, DummyConsole()) + + assert builds == [] + + def test_start_backend_raises_when_port_has_listener(monkeypatch, tmp_path: Path) -> None: paths = service_manager.RuntimePaths( root=tmp_path, diff --git a/tests/cli/test_session_commands.py b/tests/cli/test_session_commands.py index a664095c7..4bd75558a 100644 --- a/tests/cli/test_session_commands.py +++ b/tests/cli/test_session_commands.py @@ -1,3 +1,5 @@ +from unittest.mock import AsyncMock + import pytest from typer.testing import CliRunner @@ -83,6 +85,30 @@ async def fake_delete(project_id: str, session_id: str): await session_cmd._delete_session(session.id, None, force=True) +@pytest.mark.asyncio +async def test_delete_session_confirmation_calls_out_permanent_deletion(monkeypatch) -> None: + session = _build_session(project_id="proj_delete_prompt", title="Delete Forever") + prompts: list[str] = [] + + async def fake_get_by_id(_session_id: str): + return session + + def fake_confirm(prompt: str, *, default: bool): + prompts.append(prompt) + assert default is False + return False + + monkeypatch.setattr(session_cmd.Storage, "init", _noop_storage_init) + monkeypatch.setattr(session_cmd.Session, "get_by_id", fake_get_by_id) + monkeypatch.setattr(session_cmd.typer, "confirm", fake_confirm) + + await session_cmd._delete_session(session.id, None, force=False) + + assert prompts == [ + "Permanently delete session 'Delete Forever' and all of its messages and history?" + ] + + @pytest.mark.asyncio async def test_archive_session_resolves_project_from_session(monkeypatch) -> None: session = _build_session(project_id="proj_archive") @@ -105,19 +131,34 @@ async def fake_archive(project_id: str, session_id: str): @pytest.mark.asyncio async def test_restore_session_resolves_project_from_session(monkeypatch) -> None: - session = _build_session(project_id="proj_restore") + session = _build_session(project_id="proj_restore").model_copy( + update={"owner_user_id": "usr_cli"} + ) async def fake_get_by_id(session_id: str): assert session_id == session.id return session - async def fake_unarchive(project_id: str, session_id: str): + async def fake_restore( + project_id: str, + session_id: str, + *, + project_owner_id: str, + ): assert project_id == session.project_id assert session_id == session.id + assert project_owner_id == "usr_cli" return True monkeypatch.setattr(session_cmd.Storage, "init", _noop_storage_init) monkeypatch.setattr(session_cmd.Session, "get_by_id", fake_get_by_id) - monkeypatch.setattr(session_cmd.Session, "unarchive", fake_unarchive) + restore_session = AsyncMock(side_effect=fake_restore) + monkeypatch.setattr(session_cmd.Session, "restore", restore_session) await session_cmd._restore_session(session.id, None) + + restore_session.assert_awaited_once_with( + session.project_id, + session.id, + project_owner_id="usr_cli", + ) diff --git a/tests/conftest.py b/tests/conftest.py index ecb8e12c9..651e714fc 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -3,6 +3,7 @@ """ import os +from pathlib import Path import pytest @@ -37,3 +38,51 @@ def pytest_runtest_setup(item): for marker_name, (env_var, reason) in _API_KEY_MARKERS.items(): if item.get_closest_marker(marker_name) and not os.getenv(env_var): pytest.skip(reason) + + +@pytest.fixture(autouse=True) +def _home_honors_env(monkeypatch): + """Make ``Path.home()`` follow the ``HOME`` env var on every platform. + + Many suites isolate their filesystem by ``monkeypatch.setenv("HOME", ...)`` + and then rely on production code that resolves ``Path.home() / ".flocks"`` + (hub installs, plugin roots, the WebUI contract store, ...). That works on + Linux CI, where ``Path.home()`` honors ``$HOME`` — but on Windows + ``Path.home()`` reads ``%USERPROFILE%`` and ignores ``HOME`` entirely, so + those tests silently leak into (and assert against) the *real* user home. + + This autouse fixture patches ``pathlib.Path.home`` to honor ``HOME`` (then + ``USERPROFILE``) at call time, aligning Windows with Linux CI. It is + behavior-preserving when ``HOME`` already points at the real home (the + default), so tests that never touch ``HOME`` see no change. + """ + def _home() -> Path: + candidate = os.environ.get("HOME") or os.environ.get("USERPROFILE") + if candidate: + return Path(candidate) + return Path(os.path.expanduser("~")) + + monkeypatch.setattr(Path, "home", staticmethod(_home)) + + # ``Path.home()`` covers explicit home lookups, but ``~``/``~user`` tilde + # expansion goes through ``os.path.expanduser`` (and ``Path.expanduser``), + # which on Windows keys off ``%USERPROFILE%``/``%HOMEDRIVE%%HOMEPATH%`` and + # likewise ignores ``HOME``. Route a bare ``~`` / ``~/...`` through the same + # ``HOME``-aware root so tilde-based tests are hermetic too. Anything else + # (``~otheruser``) falls back to the real implementation untouched. + _real_expanduser = os.path.expanduser + + def _expanduser(path): + text = os.fspath(path) + # Only rewrite plain ``str`` tilde paths; leave bytes and ``~user`` + # forms to the real implementation so we never change its type + # contract or raise on a bytes path. + if isinstance(text, str): + home = os.environ.get("HOME") or os.environ.get("USERPROFILE") + if home and (text == "~" or text.startswith("~/") or text.startswith("~\\")): + return home + text[1:] + return _real_expanduser(path) + + monkeypatch.setattr(os.path, "expanduser", _expanduser) + + diff --git a/tests/contracts/webui/test_store.py b/tests/contracts/webui/test_store.py index d88f35175..78d4c9ede 100644 --- a/tests/contracts/webui/test_store.py +++ b/tests/contracts/webui/test_store.py @@ -50,6 +50,7 @@ def _write_workspace( workspace_id: str, title: str, order: int = 100, + version: str | None = None, default_page_id: str | None = None, sections: list[dict] | None = None, ) -> None: @@ -63,6 +64,8 @@ def _write_workspace( "enabled": True, "placement": "sceneWorkspace", } + if version is not None: + payload["version"] = version if default_page_id is not None: payload["defaultPageId"] = default_page_id if sections is not None: @@ -109,6 +112,40 @@ def test_list_pages_scans_user_and_project_roots_with_user_priority(tmp_path): assert store.page_dir("project-page").is_relative_to(project_root) +def test_list_pages_skips_installer_scratch_dirs(tmp_path): + """Pages that exist only under the Hub installer's ``..`` / + ``..bak`` scratch dirs must never surface as real pages. + + Reproduces the Windows WinError 5 aftermath where a failed atomic swap + left the SOC pages' manifests stuck inside scratch dirs while the real + install was incomplete. + """ + user_root = tmp_path / "user" / "contracts" / "webui" + _write_page(user_root, "real-page", "Real Page") + # Manifests stranded in leftover scratch/backup dirs (no real install). + _write_page_at(user_root, ".soc_ui.abc123/soc_overview", "soc-overview", "Stranded Overview") + _write_page_at(user_root, ".soc_ui.bak/soc_dashboard", "soc-dashboard", "Stranded Dashboard") + + store = WebUIPagesStore(root=user_root, project_root=None, legacy_root=None) + + pages = store.list_pages() + assert [page.id for page in pages] == ["real-page"] + assert pages[0].title == "Real Page" + + +def test_list_workspaces_skips_installer_scratch_dirs(tmp_path): + user_root = tmp_path / "user" / "contracts" / "webui" + _write_workspace(user_root, "real_ws", "Real Workspace") + # A workspace manifest stranded under a scratch dir must be ignored. + _write_workspace(user_root / ".soc_ui.abc123", "soc_ui", "Stranded Workspace") + + store = WebUIPagesStore(root=user_root, project_root=None, legacy_root=None) + + workspaces = store.list_workspaces() + assert [workspace.id for workspace in workspaces] == ["real_ws"] + assert workspaces[0].title == "Real Workspace" + + def test_grouped_page_directory_uses_manifest_id_for_lookup(tmp_path): user_root = tmp_path / "user" / "contracts" / "webui" _write_workspace(user_root, "scene_workspace", "场景工作区") @@ -140,6 +177,7 @@ def test_list_workspaces_returns_grouped_pages(tmp_path): "scene_workspace", "场景工作区", order=5, + version="2.3.4", default_page_id="ops-overview", sections=[ { @@ -159,6 +197,7 @@ def test_list_workspaces_returns_grouped_pages(tmp_path): workspaces = store.list_workspaces() assert [workspace.id for workspace in workspaces] == ["scene_workspace"] + assert workspaces[0].version == "2.3.4" assert workspaces[0].title == "场景工作区" assert workspaces[0].route == "/contracts/webui/workspaces/scene_workspace" assert workspaces[0].placement == "sceneWorkspace" diff --git a/tests/contracts/webui/test_watcher.py b/tests/contracts/webui/test_watcher.py index a5043c6a1..271bbfc48 100644 --- a/tests/contracts/webui/test_watcher.py +++ b/tests/contracts/webui/test_watcher.py @@ -104,3 +104,39 @@ def test_watcher_classifies_workspace_manifest_change(tmp_path): assert page_id == "scene_workspace" assert pending.manifest_changed + + +def test_watcher_ignores_installer_scratch_dir_events(tmp_path): + """Events under the Hub installer's ``..`` / ``..bak`` + scratch dirs must be dropped *before* any store lookup — reacting there + races the installer's atomic swap and on Windows blocks it with a + WinError 5 access-denied. + + The store is stubbed to raise if consulted, proving the dot-dir guard + short-circuits ahead of path resolution rather than relying on the + store's own scratch-dir filtering. + """ + root = tmp_path / "webui_pages" + root.mkdir(parents=True) + + class _RaisingStore: + def workspace_id_for_path(self, _path): + raise AssertionError("store should not be consulted for scratch-dir events") + + def page_id_for_path(self, _path): + raise AssertionError("store should not be consulted for scratch-dir events") + + watcher = WebUIPagesWatcher( + store=_RaisingStore(), builder=_BuilderStub(), api_runtime=_RuntimeStub() + ) + + for candidate in ( + root / ".soc_ui.abc123" / "soc_overview" / "src" / "index.tsx", + root / ".soc_ui.abc123" / "soc_overview" / "manifest.json", + root / ".soc_ui.abc123" / "soc_overview" / "dist" / "page.js", + root / ".soc_ui.bak" / "soc_overview" / "workspace.json", + ): + assert ( + watcher._classify_event(candidate, root, event_type="modified", is_directory=False) + is None + ) diff --git a/tests/hooks/test_pipeline.py b/tests/hooks/test_pipeline.py index 807a43853..5e285b2bb 100644 --- a/tests/hooks/test_pipeline.py +++ b/tests/hooks/test_pipeline.py @@ -10,7 +10,7 @@ import pytest -from flocks.hooks.pipeline import HookBase, HookPipeline +from flocks.hooks.pipeline import HookBase, HookPipeline, HookStage from flocks.plugin.loader import PluginLoader @@ -111,3 +111,60 @@ async def test_pipeline_resolves_project_dir_from_session_when_workspace_missing ) assert "hook-session" in set(HookPipeline.list_hooks()) + + +def test_hook_base_exposes_only_canonical_lifecycle_names() -> None: + assert not hasattr(HookBase, "chat_message") + assert hasattr(HookBase, "tool_before") + assert hasattr(HookBase, "tool_after") + assert not hasattr(HookBase, "pre_tool_use") + assert not hasattr(HookBase, "post_tool_use") + assert not hasattr(HookPipeline, "run_chat_message") + assert hasattr(HookPipeline, "run_tool_before") + assert hasattr(HookPipeline, "run_tool_after") + assert not hasattr(HookPipeline, "run_pre_tool_use") + assert not hasattr(HookPipeline, "run_post_tool_use") + + +def test_default_hook_methods_are_not_stage_handlers() -> None: + hook = HookBase() + + assert HookPipeline._resolve_handler( + hook, + HookStage.USER_PROMPT_SUBMIT, + ) is None + assert HookPipeline._resolve_handler(hook, HookStage.LLM_BEFORE) is None + assert HookPipeline._resolve_handler(hook, HookStage.TOOL_BEFORE) is None + + +@pytest.mark.asyncio +async def test_pipeline_runs_canonical_lifecycle_handlers() -> None: + seen: list[str] = [] + + class _CanonicalHook(HookBase): + async def user_prompt_submit(self, ctx) -> None: + seen.append(ctx.stage) + + async def tool_before(self, ctx) -> None: + seen.append(ctx.stage) + + async def tool_after(self, ctx) -> None: + seen.append(ctx.stage) + + async def turn_finish(self, ctx) -> None: + seen.append(ctx.stage) + + HookPipeline.register("canonical-hook", _CanonicalHook()) + HookPipeline._initialized = True + + await HookPipeline.run_user_prompt_submit({"sessionID": "ses_test"}) + await HookPipeline.run_tool_before({"sessionID": "ses_test"}) + await HookPipeline.run_tool_after({"sessionID": "ses_test"}) + await HookPipeline.run_turn_finish({"sessionID": "ses_test"}) + + assert seen == [ + HookStage.USER_PROMPT_SUBMIT, + HookStage.TOOL_BEFORE, + HookStage.TOOL_AFTER, + HookStage.TURN_FINISH, + ] diff --git a/tests/hub/test_hub_catalog.py b/tests/hub/test_hub_catalog.py index bdb479b2e..0eaab0640 100644 --- a/tests/hub/test_hub_catalog.py +++ b/tests/hub/test_hub_catalog.py @@ -270,6 +270,45 @@ def test_catalog_ignores_stale_record_when_legacy_install_is_inferred( assert local.get_record("device", plugin_id) is None +def test_catalog_uses_webui_workspace_version_for_inferred_installs( + isolated_hub_env, +): + from flocks.hub.catalog import clear_catalog_caches + + webui_dir = ( + isolated_hub_env["home"] + / ".flocks" + / "plugins" + / "contracts" + / "webui" + / "soc_ui" + ) + webui_dir.mkdir(parents=True) + workspace_path = webui_dir / "workspace.json" + workspace = { + "id": "soc_ui", + "version": "1.0.0", + "title": "SOC 工作区", + "placement": "sceneWorkspace", + } + workspace_path.write_text(json.dumps(workspace), encoding="utf-8") + clear_catalog_caches() + + entry = {item.id: item for item in list_catalog(plugin_type="webui")}["soc_ui"] + + assert entry.version == "1.1.4" + assert entry.state == "updateAvailable" + assert entry.installedVersion == "1.0.0" + + workspace["version"] = "1.1.4" + workspace_path.write_text(json.dumps(workspace), encoding="utf-8") + + refreshed = {item.id: item for item in list_catalog(plugin_type="webui")}["soc_ui"] + + assert refreshed.state == "installed" + assert refreshed.installedVersion == "1.1.4" + + def test_pentest_agents_are_listed_in_agent_catalog(): entries = list_catalog(plugin_type="agent") ids = {entry.id for entry in entries} @@ -381,6 +420,7 @@ async def noop_refresh(_plugin_type, _changed_path=None): webui_dir = home_plugins / "contracts" / "webui" / "soc_ui" access_dir = home_plugins / "contracts" / "access" / "soc_ui" dashboard_manifest = json.loads((webui_dir / "soc_dashboard" / "manifest.json").read_text(encoding="utf-8")) + workspace_manifest = json.loads((webui_dir / "workspace.json").read_text(encoding="utf-8")) assert (webui_dir / "workspace.json").is_file() assert (webui_dir / "soc_alerts" / "dist" / "page.js").is_file() @@ -390,6 +430,7 @@ async def noop_refresh(_plugin_type, _changed_path=None): assert (access_dir / "soc_alerts_operations.py").is_file() assert set(built_pages) == {"soc-alerts", "soc-dashboard", "soc-overview"} assert dashboard_manifest["id"] == "soc-dashboard" + assert workspace_manifest["version"] == record.version == load_manifest("webui", "soc_ui").version assert record.installPath == str(webui_dir) removed = await uninstall_plugin("webui", "soc_ui") @@ -475,6 +516,8 @@ async def noop_refresh(_plugin_type, _changed_path=None): "threat_rule_id": "D1181087257", "threat_name": "SQL injection", "threat_msg": "Detected SQL injection attempt.", + "threat_severity": "critical", + "threat_level": "high", "threat_phase": "exploit", "threat_type": "exploit", "threat_result": "failed", @@ -545,6 +588,44 @@ async def noop_refresh(_plugin_type, _changed_path=None): assert response.body["summary"]["attackFailed"] == 1 assert response.body["incidents"][0]["id"] == "alert-1" assert response.body["incidents"][0]["tableCells"]["_source_type"]["value"] == "tdp" + assert response.body["incidents"][0]["tableCells"]["threat_severity"]["value"] == "critical" + assert response.body["incidents"][0]["tableCells"]["threat_level"]["value"] == "high" + + filtered = runtime.execute( + page_id="soc-alerts", + contract_id="soc.alerts.operations", + operation_name="list", + payload={ + "params": { + "filters": { + "threat_severity": ["critical"], + "threat_level": ["high"], + }, + "limit": 10, + } + }, + principal=AuthUser(id="u1", username="admin", role="admin"), + ) + + assert filtered.status_code == 200 + assert filtered.body["summary"]["representativeCount"] == 1 + assert [incident["id"] for incident in filtered.body["incidents"]] == ["alert-1"] + + for filters in ( + {"threat_severity": ["low"]}, + {"threat_level": ["low"]}, + ): + excluded = runtime.execute( + page_id="soc-alerts", + contract_id="soc.alerts.operations", + operation_name="list", + payload={"params": {"filters": filters, "limit": 10}}, + principal=AuthUser(id="u1", username="admin", role="admin"), + ) + + assert excluded.status_code == 200 + assert excluded.body["summary"]["representativeCount"] == 0 + assert excluded.body["incidents"] == [] async def test_hub_installs_soc_workspace_component_children(isolated_hub_env, monkeypatch: pytest.MonkeyPatch): diff --git a/tests/hub/test_installer_swap.py b/tests/hub/test_installer_swap.py new file mode 100644 index 000000000..2ce6bbbbd --- /dev/null +++ b/tests/hub/test_installer_swap.py @@ -0,0 +1,156 @@ +"""Unit tests for the Hub installer's Windows-safe directory swap helpers. + +These exercise the atomic-swap path in isolation (explicit ``tmp_path`` +dirs, no ``Path.home()`` dependency), covering the WinError 5 aftermath +where a directory watcher or AV scan holds a transient handle on the +freshly staged tree and blocks ``src.replace(dst)``. +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +from flocks.hub import installer + + +def _make_tree(path: Path, marker: str) -> None: + path.mkdir(parents=True, exist_ok=True) + (path / "manifest.json").write_text(marker, encoding="utf-8") + + +def test_replace_prepared_path_swaps_into_place(tmp_path): + src = tmp_path / ".soc_ui.scratch" + dst = tmp_path / "soc_ui" + _make_tree(src, "new") + + backup = installer._replace_prepared_path(src, dst) + + assert backup is None + assert (dst / "manifest.json").read_text(encoding="utf-8") == "new" + assert not src.exists() + assert not (tmp_path / ".soc_ui.bak").exists() + + +def test_replace_prepared_path_overwrites_existing_and_returns_backup(tmp_path): + src = tmp_path / ".soc_ui.scratch" + dst = tmp_path / "soc_ui" + _make_tree(src, "new") + _make_tree(dst, "old") + + backup = installer._replace_prepared_path(src, dst) + + assert (dst / "manifest.json").read_text(encoding="utf-8") == "new" + assert backup == tmp_path / ".soc_ui.bak" + assert (backup / "manifest.json").read_text(encoding="utf-8") == "old" + + installer._commit_replacement(backup) + + assert not (tmp_path / ".soc_ui.bak").exists() + + +def test_replace_prepared_path_restores_existing_after_new_swap_fails(tmp_path, monkeypatch): + monkeypatch.setattr(installer.sys, "platform", "win32") + monkeypatch.setattr(installer.time, "sleep", lambda _s: None) + + src = tmp_path / ".soc_ui.scratch" + dst = tmp_path / "soc_ui" + _make_tree(src, "new") + _make_tree(dst, "old") + + real_replace = Path.replace + failed_attempts = 0 + + def deny_new_swap(self, target): + nonlocal failed_attempts + if self == src and Path(target) == dst: + failed_attempts += 1 + raise PermissionError("[WinError 5] Access is denied") + return real_replace(self, target) + + monkeypatch.setattr(Path, "replace", deny_new_swap) + with pytest.raises(PermissionError): + installer._replace_prepared_path(src, dst) + + assert failed_attempts == 6 + assert (dst / "manifest.json").read_text(encoding="utf-8") == "old" + assert (src / "manifest.json").read_text(encoding="utf-8") == "new" + assert not (tmp_path / ".soc_ui.bak").exists() + + +def test_replace_with_retry_recovers_from_transient_permission_error(tmp_path, monkeypatch): + """A first ``PermissionError`` (WinError 5) is retried, not surfaced.""" + monkeypatch.setattr(installer.sys, "platform", "win32") + monkeypatch.setattr(installer.time, "sleep", lambda _s: None) + + src = tmp_path / ".soc_ui.scratch" + dst = tmp_path / "soc_ui" + _make_tree(src, "new") + + real_replace = Path.replace + calls = {"n": 0} + + def flaky_replace(self, target): + calls["n"] += 1 + if calls["n"] == 1: + raise PermissionError("[WinError 5] Access is denied") + return real_replace(self, target) + + monkeypatch.setattr(Path, "replace", flaky_replace) + installer._replace_with_retry(src, dst) + + assert calls["n"] == 2 + assert (dst / "manifest.json").read_text(encoding="utf-8") == "new" + + +def test_replace_with_retry_reraises_after_exhausting_attempts(tmp_path, monkeypatch): + monkeypatch.setattr(installer.sys, "platform", "win32") + monkeypatch.setattr(installer.time, "sleep", lambda _s: None) + + src = tmp_path / ".soc_ui.scratch" + dst = tmp_path / "soc_ui" + _make_tree(src, "new") + + def always_denied(self, target): + raise PermissionError("[WinError 5] Access is denied") + + monkeypatch.setattr(Path, "replace", always_denied) + with pytest.raises(PermissionError): + installer._replace_with_retry(src, dst) + + +def test_purge_stale_scratch_removes_leftovers(tmp_path): + parent = tmp_path + _make_tree(parent / ".soc_ui.55ram7wo" / "soc_overview", "stranded") + _make_tree(parent / ".soc_ui.bak" / "soc_dashboard", "stranded") + _make_tree(parent / "soc_ui", "live") + # Unrelated dot-dir for a different plugin must be left untouched. + _make_tree(parent / ".other.bak", "keep") + + installer._purge_stale_scratch(parent, "soc_ui") + + assert not (parent / ".soc_ui.55ram7wo").exists() + assert not (parent / ".soc_ui.bak").exists() + assert (parent / "soc_ui" / "manifest.json").read_text(encoding="utf-8") == "live" + assert (parent / ".other.bak").exists() + + +def test_copy_package_purges_stale_scratch_before_staging(tmp_path): + """A prior failed install's leftovers self-heal on the next install.""" + src = tmp_path / "bundled" / "soc_ui" + _make_tree(src, "payload") + (src / "manifest.json").write_text('{"id": "soc_ui"}', encoding="utf-8") + (src / "src").mkdir() + (src / "src" / "index.tsx").write_text("export default 1;\n", encoding="utf-8") + + dst = tmp_path / "install" / "soc_ui" + dst.parent.mkdir(parents=True) + _make_tree(dst.parent / ".soc_ui.leftover" / "soc_overview", "stranded") + + installer._copy_package(src, dst) + + assert (dst / "src" / "index.tsx").is_file() + assert not (dst.parent / ".soc_ui.leftover").exists() + # ``manifest.json`` is intentionally not copied by the installer. + assert not (dst / "manifest.json").exists() diff --git a/tests/hub/test_soc_dashboard_schema.py b/tests/hub/test_soc_dashboard_schema.py index fc4257162..30e87ce49 100644 --- a/tests/hub/test_soc_dashboard_schema.py +++ b/tests/hub/test_soc_dashboard_schema.py @@ -30,11 +30,59 @@ def _load_dashboard_handlers(): return module +def _load_overview_handlers(): + handler_path = ( + Path(__file__).resolve().parents[2] + / ".flocks" + / "flockshub" + / "plugins" + / "webuis" + / "soc_ui" + / "soc_overview" + / "api" + / "handlers.py" + ) + spec = importlib.util.spec_from_file_location("soc_overview_fields_test", handler_path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + previous = sys.dont_write_bytecode + try: + sys.dont_write_bytecode = True + spec.loader.exec_module(module) + finally: + sys.dont_write_bytecode = previous + return module + + +def _load_alert_operations(): + operations_path = ( + Path(__file__).resolve().parents[2] + / ".flocks" + / "flockshub" + / "plugins" + / "webuis" + / "soc_ui" + / "access" + / "soc_alerts_operations.py" + ) + spec = importlib.util.spec_from_file_location("soc_alert_operations_test", operations_path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + previous = sys.dont_write_bytecode + try: + sys.dont_write_bytecode = True + spec.loader.exec_module(module) + finally: + sys.dont_write_bytecode = previous + return module + + def test_soc_dashboard_migrates_legacy_alert_records_schema(tmp_path: Path): db_path = tmp_path / "soc.db" first_record = { "_source_type": "tdp", "threat_name": "SQL injection", + "_threat_type": "web_attack", "triage_status": "ok", "_triage_persisted_at": "2026-07-14T13:00:00", "attack_verdict": "attack", @@ -74,6 +122,7 @@ def test_soc_dashboard_migrates_legacy_alert_records_schema(tmp_path: Path): second_record = { "_source_type": "hids", "threat_name": "Malware download", + "threat_type": "malware", "triage_status": "ok", "_triage_persisted_at": "2026-07-14T13:01:00", } @@ -127,7 +176,7 @@ def test_soc_dashboard_migrates_legacy_alert_records_schema(tmp_path: Path): columns = {row[1] for row in conn.execute("PRAGMA table_info(alert_records)")} indexes = {row[1] for row in conn.execute("PRAGMA index_list(alert_records)")} facts = conn.execute( - "SELECT alert_row_id, row_key, source_type, threat_name, has_triage " + "SELECT alert_row_id, row_key, source_type, threat_name, threat_type, has_triage " "FROM soc_dashboard_alert_facts ORDER BY alert_row_id" ).fetchall() source_rows = conn.execute( @@ -163,12 +212,189 @@ def test_soc_dashboard_migrates_legacy_alert_records_schema(tmp_path: Path): } <= indexes assert source_rows == [(1, "1", 0), (2, "2", 0)] assert facts == [ - (1, "1", "tdp", "SQL injection", 1), - (2, "2", "hids", "Malware download", 1), + (1, "1", "tdp", "SQL injection", "web_attack", 1), + (2, "2", "hids", "Malware download", "malware", 1), ] assert "COALESCE(NULLIF(NEW.row_id, ''), CAST(NEW.rowid AS TEXT))" in trigger_sql assert updated_fact == ("2026-07-14T13:02:00", "attack") - assert schema_version == "2" + assert schema_version == "3" + + +def test_soc_dashboard_triage_outcomes_partition_records(tmp_path: Path): + db_path = tmp_path / "soc.db" + asset_date = "2026-07-14" + records = [ + {"triage_status": "ok", "attack_verdict": "attack_success", "attack_success": True}, + {"triage_status": "ok", "attack_verdict": "attack"}, + {"triage_status": "ok", "attack_verdict": "attack_failed"}, + {"triage_status": "ok", "attack_verdict": "benign"}, + {"triage_status": "ok", "attack_verdict": "unknown"}, + {"triage_status": "failed", "attack_verdict": "unknown"}, + {"triage_status": "failed", "attack_verdict": "benign"}, + {"triage_status": "ok", "attack_verdict": "legacy", "attack_success": True}, + ] + severity_values = ["low", "critical", "high", "medium", "low", "critical", "high", "medium"] + for index, record in enumerate(records): + record.update( + { + "threat_name": f"threat-name-{index}", + "_threat_type": f"threat-type-{index}", + "threat_type": f"fallback-type-{index}", + "threat_severity": severity_values[index], + "threat_level": "ignored-threat-level", + "risk_level": "High", + } + ) + with sqlite3.connect(db_path) as conn: + conn.execute( + """ + CREATE TABLE alert_records ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + record_json TEXT NOT NULL, + asset_date TEXT NOT NULL, + event_time INTEGER NOT NULL + ) + """ + ) + conn.executemany( + "INSERT INTO alert_records(record_json, asset_date, event_time) VALUES (?, ?, ?)", + [ + (json.dumps(record), asset_date, 1784014800 + index) + for index, record in enumerate(records) + ], + ) + conn.commit() + + handlers = _load_dashboard_handlers() + handlers.DEFAULT_SQLITE_DB = db_path + handlers._schema_ready.clear() + + assert handlers._ensure_sqlite_schema() is True + sources = handlers._find_sqlite_sources( + asset_date, + asset_date, + 1784014800, + 1784014800 + len(records), + ) + triage = handlers._read_triage(sources) + closed_loop = handlers._build_closed_loop(triage) + with sqlite3.connect(db_path) as conn: + timeline = handlers._sqlite_timeline( + conn, + handlers._sqlite_settings(), + "asset_date = ? AND event_time BETWEEN ? AND ?", + [asset_date, 1784014800, 1784014800 + len(records)], + [asset_date], + 1784014800, + 1784014800 + len(records), + ) + first_fact = conn.execute( + "SELECT threat_name, threat_type, severity, risk_level " + "FROM soc_dashboard_alert_facts ORDER BY alert_row_id LIMIT 1" + ).fetchone() + + assert triage["totalRecords"] == 8 + assert triage["newTriaged"] == 6 + assert triage["attackSuccess"] == 2 + assert triage["attack"] == 1 + assert triage["attackFailed"] == 1 + assert triage["attackTotal"] == 4 + assert triage["benign"] == 1 + assert triage["unknown"] == 1 + assert triage["triageFailed"] == 2 + assert first_fact == ("threat-name-0", "threat-type-0", "low", "High") + assert dict(triage["threatTypeCounter"]) == { + f"threat-type-{index}": 1 for index in range(len(records)) + } + assert dict(triage["severityCounter"]) == { + "critical": 2, + "high": 2, + "low": 2, + "medium": 2, + } + assert dict(triage["riskCounter"]) == {"high": 8} + assert closed_loop["pending"] == 3 + assert triage["attackTotal"] + triage["benign"] + closed_loop["pending"] == 8 + assert sum(timeline["attack"]) == triage["attackTotal"] + + +def test_soc_overview_keeps_threat_names_and_types_separate(tmp_path: Path): + db_path = tmp_path / "soc.db" + asset_date = "2026-07-14" + records = [ + { + "threat_name": "SQL injection attempt", + "_threat_type": "web_attack", + "threat_type": "ignored_fallback", + }, + { + "threat_name": "Suspicious crawler", + }, + ] + with sqlite3.connect(db_path) as conn: + conn.execute( + "CREATE TABLE alert_records " + "(record_json TEXT NOT NULL, asset_date TEXT NOT NULL, event_time INTEGER NOT NULL)" + ) + conn.executemany( + "INSERT INTO alert_records VALUES (?, ?, ?)", + [ + (json.dumps(record), asset_date, 1784014800 + index) + for index, record in enumerate(records) + ], + ) + conn.commit() + + handlers = _load_overview_handlers() + handlers.DEFAULT_SQLITE_DB = db_path + source = handlers._RecordSource( + path=db_path, + role="denoise", + date=asset_date, + data_source="sqlite", + ) + denoise = handlers._read_denoise([source]) + field_stats = handlers._build_field_stats([source]) + + assert dict(denoise["threatCounter"]) == { + "sql injection attempt": 1, + "suspicious crawler": 1, + } + assert { + item["label"]: item["value"] for item in field_stats["threatTypes"] + } == {"web_attack": 1, "unknown": 1} + + +def test_soc_dashboard_activity_does_not_mix_name_type_or_risk_fields(): + handlers = _load_dashboard_handlers() + row = { + "activity_row_id": 1, + "activity_event_time": 1784014800, + "sample_count": 1, + "record_json": json.dumps( + { + "_threat_type": "web_attack", + "triage_status": "ok", + "threat_level": "critical", + } + ), + } + + event = handlers._activity_event(row) + + assert event["alert"]["threatName"] == "未知告警" + assert event["alert"]["threatType"] == "web_attack" + assert event["result"]["threatSeverity"] == "" + assert event["result"]["riskLevel"] == "" + + +def test_soc_alert_verdict_does_not_fall_back_to_risk_or_threat_level(): + operations = _load_alert_operations() + + assert operations._verdict_bucket( + {"risk_level": "attack_success", "threat_level": "attack_failed"} + ) == "unknown" + assert operations._verdict_bucket({"attack_verdict": "attack_success"}) == "success" def test_soc_dashboard_activity_exposes_live_denoise_workflow_progress(tmp_path: Path): diff --git a/tests/observability/test_langfuse_observability.py b/tests/observability/test_langfuse_observability.py index 8473ac802..8dcc02ef8 100644 --- a/tests/observability/test_langfuse_observability.py +++ b/tests/observability/test_langfuse_observability.py @@ -27,6 +27,15 @@ def trace(self, **kwargs): return _FakeObservation("trace", kwargs) +class _NewSdkTraceClient: + def __init__(self): + self.start_observation_payload = None + + def start_observation(self, **kwargs): + self.start_observation_payload = kwargs + return _TrackingObservation("trace", kwargs) + + class _FakeObservation: """Fake Langfuse observation with end/generation/span support.""" @@ -45,6 +54,41 @@ def end(self, **kwargs): self.end_payload = kwargs +class _TrackingSpan: + def __init__(self) -> None: + self.attributes = {} + + def is_recording(self): + return True + + def set_attribute(self, key, value): + self.attributes[key] = value + + +class _TrackingObservation(_FakeObservation): + def __init__(self, kind: str, payload: dict): + super().__init__(kind, payload) + self._otel_span = _TrackingSpan() + + def generation(self, **kwargs): + return _TrackingObservation("generation", kwargs) + + def span(self, **kwargs): + return _TrackingObservation("span", kwargs) + + +class _NewSdkLikeObservation: + def __init__(self) -> None: + self.update_payload = None + self.end_calls = 0 + + def update(self, **kwargs): + self.update_payload = kwargs + + def end(self): + self.end_calls += 1 + + def test_create_span_uses_current_observation(monkeypatch): monkeypatch.setattr(lf, "_get_client", lambda: object()) parent = _FakeParent() @@ -92,6 +136,55 @@ def test_create_trace_forwards_tags(monkeypatch): assert client.trace_payload["tags"] == ["session:s1", "step:2", "session_step:s1:2"] +def test_create_trace_uses_start_observation_for_new_sdk(monkeypatch): + client = _NewSdkTraceClient() + monkeypatch.setattr(lf, "_get_client", lambda: client) + + obs = lf.create_trace( + name="SessionRunner.step", + session_id="s1", + user_id="u1", + tags=["session:s1", "step:2"], + input={"step": 2}, + metadata={"provider_id": "openai"}, + ) + + assert obs.kind == "trace" + assert client.start_observation_payload is not None + assert client.start_observation_payload["as_type"] == "span" + assert client.start_observation_payload["input"] == {"step": 2} + assert client.start_observation_payload["metadata"]["provider_id"] == "openai" + assert client.start_observation_payload["metadata"]["session_id"] == "s1" + assert client.start_observation_payload["metadata"]["user_id"] == "u1" + assert client.start_observation_payload["metadata"]["tags"] == ["session:s1", "step:2"] + assert obs._otel_span.attributes["langfuse.trace.name"] == "SessionRunner.step" + assert obs._otel_span.attributes["session.id"] == "s1" + assert obs._otel_span.attributes["user.id"] == "u1" + assert obs._otel_span.attributes["langfuse.trace.tags"] == ["session:s1", "step:2"] + + +def test_generation_and_span_inherit_trace_dimensions_from_parent(monkeypatch): + monkeypatch.setattr(lf, "_get_client", lambda: object()) + + parent = _TrackingObservation("trace", {"name": "trace"}) + parent._otel_span.attributes["langfuse.trace.name"] = "SessionRunner.step" + parent._otel_span.attributes["session.id"] = "s1" + parent._otel_span.attributes["user.id"] = "u1" + parent._otel_span.attributes["langfuse.trace.tags"] = ["session:s1", "step:2"] + + gen = lf.create_generation(parent=parent, name="LLM.generate", model="gpt-5", input={"x": 1}) + span = lf.create_span(parent=parent, name="Tool.execute.read", input={"path": "/tmp/a"}) + + assert gen._otel_span.attributes["langfuse.trace.name"] == "SessionRunner.step" + assert gen._otel_span.attributes["session.id"] == "s1" + assert gen._otel_span.attributes["user.id"] == "u1" + assert gen._otel_span.attributes["langfuse.trace.tags"] == ["session:s1", "step:2"] + assert span._otel_span.attributes["langfuse.trace.name"] == "SessionRunner.step" + assert span._otel_span.attributes["session.id"] == "s1" + assert span._otel_span.attributes["user.id"] == "u1" + assert span._otel_span.attributes["langfuse.trace.tags"] == ["session:s1", "step:2"] + + def test_initialize_supports_langfuse_base_url(monkeypatch): class _FakeLangfuseClient: def __init__(self, **kwargs): @@ -156,6 +249,29 @@ def test_end_observation_passes_usage(monkeypatch): assert gen_obs.end_payload.get("output") == "result" +def test_end_observation_updates_before_end_for_new_sdk(): + """Langfuse v4-style observations require update(...), then end().""" + obs = _NewSdkLikeObservation() + usage = {"prompt_tokens": 100, "completion_tokens": 50} + + lf.end_observation( + obs, + output={"content": "result"}, + metadata={"status": "ok"}, + usage=usage, + level="ERROR", + status_message="done", + ) + + assert obs.update_payload is not None + assert obs.update_payload["output"] == {"content": "result"} + assert obs.update_payload["metadata"] == {"status": "ok"} + assert obs.update_payload["usage_details"] == usage + assert obs.update_payload["level"] == "ERROR" + assert obs.update_payload["status_message"] == "done" + assert obs.end_calls == 1 + + def test_scope_end_is_idempotent(): """Calling end() twice on a scope should not raise.""" noop = lf._NoopObservation("test") @@ -164,11 +280,21 @@ def test_scope_end_is_idempotent(): scope.end(output="second") -def test_sanitize_truncates_long_strings(): +def test_sanitize_keeps_full_strings_in_full_mode(monkeypatch): + long_str = "a" * 10000 + monkeypatch.setenv("FLOCKS_LANGFUSE_CAPTURE_MODE", "full") + result = lf._sanitize_payload(long_str) + assert result == long_str + + +def test_sanitize_truncates_long_strings_in_truncated_mode(monkeypatch): long_str = "a" * 10000 + monkeypatch.setenv("FLOCKS_LANGFUSE_CAPTURE_MODE", "truncated") + monkeypatch.setenv("FLOCKS_LANGFUSE_MAX_CHARS", "128") result = lf._sanitize_payload(long_str) assert len(result) < 10000 assert "truncated" in result + assert result.startswith("a" * 128) def test_observation_scope_exception_handling(): diff --git a/tests/project/test_project.py b/tests/project/test_project.py index bd9e5e7a2..364e97c1e 100644 --- a/tests/project/test_project.py +++ b/tests/project/test_project.py @@ -1,5 +1,7 @@ import asyncio import json +import os +import sys import uuid from unittest.mock import AsyncMock, patch @@ -13,6 +15,8 @@ ProjectPathConflictError, TASK_SESSION_GROUP_ID, ) +from flocks.config.config import Config +from flocks.session.session import Session from flocks.storage.storage import Storage @@ -24,6 +28,49 @@ def project_root(tmp_path, monkeypatch): return tmp_path +@pytest.mark.asyncio +async def test_lifecycle_guard_is_reentrant_and_cross_process(tmp_path, monkeypatch): + monkeypatch.setattr( + Config, + "get_data_path", + classmethod(lambda _cls: tmp_path), + ) + marker = tmp_path / "child-acquired" + project_id = "prj_cross_process_guard" + script = """ +import os +import sys +from pathlib import Path +from flocks.project.project import _platform_file_lock, _platform_file_unlock + +fd = os.open(sys.argv[1], os.O_RDWR) +try: + _platform_file_lock(fd) + Path(sys.argv[2]).touch() + _platform_file_unlock(fd) +finally: + os.close(fd) +""" + + async with Project.lifecycle_guard(project_id): + async with Project.lifecycle_guard(project_id): + pass + lock_path = next((tmp_path / "locks" / "project-lifecycle").glob("*.lock")) + process = await asyncio.create_subprocess_exec( + sys.executable, + "-c", + script, + str(lock_path), + str(marker), + cwd=os.getcwd(), + ) + await asyncio.sleep(0.1) + assert not marker.exists() + + assert await asyncio.wait_for(process.wait(), timeout=5) == 0 + assert marker.exists() + + @pytest.mark.asyncio async def test_list_projects_uses_json_registry_without_virtual_default(project_root): labs = project_root / "labs" @@ -115,7 +162,7 @@ async def test_update_renames_registered_project(project_root): @pytest.mark.asyncio -async def test_delete_only_removes_registry_entry(project_root): +async def test_delete_soft_removes_registry_entry_and_restore_preserves_identity(project_root): labs = project_root / "labs" labs.mkdir() created = await Project.create(owner_id="user-1", name="Labs", worktree=str(labs)) @@ -126,8 +173,51 @@ async def test_delete_only_removes_registry_entry(project_root): assert result is True assert labs.exists() assert await Project.get(created.id, owner_id="user-1") is None + assert await Project.list(owner_id="user-1") == [] + assert Project.registered_project_names()[created.id] == "Labs" + registry = json.loads(Project.registry_path("user-1").read_text(encoding="utf-8")) + assert registry["projects"][0]["id"] == created.id + assert registry["projects"][0]["removedAt"] > 0 storage_delete.assert_not_awaited() + restored = await Project.restore(created.id, owner_id="user-1") + + assert restored is not None + assert restored.id == created.id + assert restored.name == "Labs" + assert [project.id for project in await Project.list(owner_id="user-1")] == [created.id] + + +@pytest.mark.asyncio +async def test_restore_removed_project_requires_original_directory(project_root): + labs = project_root / "labs" + labs.mkdir() + created = await Project.create(owner_id="user-1", name="Labs", worktree=str(labs)) + await Project.delete(created.id, owner_id="user-1") + labs.rmdir() + + with pytest.raises(ProjectDeletionError, match="directory is unavailable"): + await Project.restore(created.id, owner_id="user-1") + + assert await Project.get(created.id, owner_id="user-1") is None + + +@pytest.mark.asyncio +async def test_removed_project_rejects_direct_session_creation(project_root): + labs = project_root / "removed" + labs.mkdir() + created = await Project.create(owner_id="user-1", name="Removed", worktree=str(labs)) + await Project.delete(created.id, owner_id="user-1") + + with pytest.raises(ProjectDeletionError, match="no longer available"): + await Session.create(project_id=created.id, directory=str(labs)) + + +@pytest.mark.asyncio +async def test_restore_missing_project_metadata_is_rejected(project_root): + with pytest.raises(ProjectDeletionError, match="metadata is unavailable"): + await Project.restore("prj_missing", owner_id="user-1") + @pytest.mark.asyncio async def test_delete_rejects_default_project(project_root): diff --git a/tests/provider/test_chinese_providers.py b/tests/provider/test_chinese_providers.py index af892d611..1245b52fe 100644 --- a/tests/provider/test_chinese_providers.py +++ b/tests/provider/test_chinese_providers.py @@ -169,12 +169,43 @@ def test_alibaba_catalog(self): def test_moonshot_catalog(self): models = get_provider_model_definitions("moonshot") assert {m.id for m in models} == { + "kimi-k3", + "kimi-k2.7-code", + "kimi-k2.7-code-highspeed", "kimi-k2.5", "kimi-k2.6", - "kimi-k2-thinking", - "kimi-k2", } + k3 = next(m for m in models if m.id == "kimi-k3") + assert k3.capabilities.supports_vision is True + assert k3.capabilities.supports_reasoning is True + assert k3.capabilities.interleaved["field"] == "reasoning_content" + assert k3.pricing.input == 20.0 + assert k3.pricing.output == 100.0 + assert k3.pricing.cache_read == 2.0 + assert k3.limits.context_window == 1048576 + assert k3.limits.max_output_tokens == 131072 + + k27 = next(m for m in models if m.id == "kimi-k2.7-code") + assert k27.capabilities.supports_vision is True + assert k27.capabilities.supports_reasoning is True + assert k27.capabilities.interleaved["field"] == "reasoning_content" + assert k27.pricing.input == 6.5 + assert k27.pricing.output == 27.0 + assert k27.pricing.cache_read == 1.3 + assert k27.limits.context_window == 262144 + assert k27.limits.max_output_tokens == 32768 + + k27_highspeed = next(m for m in models if m.id == "kimi-k2.7-code-highspeed") + assert k27_highspeed.capabilities.supports_vision is True + assert k27_highspeed.capabilities.supports_reasoning is True + assert k27_highspeed.capabilities.interleaved["field"] == "reasoning_content" + assert k27_highspeed.pricing.input == 13.0 + assert k27_highspeed.pricing.output == 54.0 + assert k27_highspeed.pricing.cache_read == 2.6 + assert k27_highspeed.limits.context_window == 262144 + assert k27_highspeed.limits.max_output_tokens == 32768 + k26 = next(m for m in models if m.id == "kimi-k2.6") assert k26.capabilities.supports_vision is True assert k26.capabilities.supports_reasoning is True @@ -184,8 +215,6 @@ def test_moonshot_catalog(self): assert k26.pricing.cache_read == 1.3 assert k26.limits.context_window == 256000 - thinking = next(m for m in models if m.id == "kimi-k2-thinking") - assert thinking.capabilities.supports_reasoning is True k25 = next(m for m in models if m.id == "kimi-k2.5") assert k25.capabilities.supports_reasoning is True assert k25.capabilities.interleaved["field"] == "reasoning_content" diff --git a/tests/provider/test_openai_base_provider.py b/tests/provider/test_openai_base_provider.py index 5c2202411..6f7fbd72f 100644 --- a/tests/provider/test_openai_base_provider.py +++ b/tests/provider/test_openai_base_provider.py @@ -624,6 +624,54 @@ async def test_chat_logs_compact_message_summary(self): "assistant": 1, } + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("model_id", "extra_body"), + [ + ("kimi-k2.6", {"thinking": {"type": "enabled"}}), + ("kimi-k2.7-code", {"thinking": {"type": "enabled"}}), + ("kimi-k2.7-code-highspeed", {"thinking": {"type": "enabled"}}), + ("kimi-k3", {"reasoning_effort": "max"}), + ], + ) + async def test_chat_logs_effective_extra_body_thinking_state( + self, + model_id, + extra_body, + ): + provider, create = self._build_provider_with_client() + create.return_value = self._mock_chat_response() + info = Mock() + + from flocks.provider.provider import ChatMessage + + with patch.object(openai_base_module, "log", Mock(info=info)): + await provider.chat( + model_id, + [ChatMessage(role="user", content="hello")], + extra_body=extra_body, + ) + + _event, payload = info.call_args.args + assert payload["thinking_enabled"] is True + + @pytest.mark.asyncio + async def test_kimi_k3_uses_max_completion_tokens(self): + provider, create = self._build_provider_with_client() + create.return_value = self._mock_chat_response() + + from flocks.provider.provider import ChatMessage + + await provider.chat( + "kimi-k3", + [ChatMessage(role="user", content="hello")], + max_tokens=131072, + ) + + kwargs = create.await_args.kwargs + assert kwargs["max_completion_tokens"] == 131072 + assert "max_tokens" not in kwargs + class TestExtractReasoningContent: """Regression: some proxies send stream chunks with ``delta is None``.""" @@ -657,6 +705,108 @@ async def _stream_from_chunks(*chunks): yield chunk +class TestOpenAIBaseProviderStreamingToolCalls: + @pytest.mark.asyncio + async def test_chat_stream_emits_name_only_marker_before_complete_tool_input(self): + provider = MockProviderWithoutCatalog() + create = AsyncMock() + provider._client = MagicMock() + provider._client.chat.completions.create = create + + first_arguments = '{"filePath":"/tmp/report.md","content":"' + remaining_arguments = f'{"long content " * 500}"}}' + tool_start_chunk = SimpleNamespace( + choices=[ + SimpleNamespace( + delta=SimpleNamespace( + content=None, + tool_calls=[ + SimpleNamespace( + index=0, + id="call_write", + function=SimpleNamespace( + name="write", + arguments=first_arguments, + ), + ) + ], + ), + finish_reason=None, + ) + ], + usage=None, + ) + tool_arguments_chunk = SimpleNamespace( + choices=[ + SimpleNamespace( + delta=SimpleNamespace( + content=None, + tool_calls=[ + SimpleNamespace( + index=0, + id=None, + function=SimpleNamespace( + name=None, + arguments=remaining_arguments, + ), + ) + ], + ), + finish_reason=None, + ) + ], + usage=None, + ) + finish_chunk = SimpleNamespace( + choices=[ + SimpleNamespace( + delta=SimpleNamespace(content=None, tool_calls=None), + finish_reason="tool_calls", + ) + ], + usage=None, + ) + create.return_value = _stream_from_chunks( + tool_start_chunk, + tool_arguments_chunk, + finish_chunk, + ) + + from flocks.provider.provider import ChatMessage + + chunks = [ + chunk + async for chunk in provider.chat_stream( + "kimi-k2.7-code", + [ChatMessage(role="user", content="write a long file")], + tools=[{"type": "function", "function": {"name": "write"}}], + ) + ] + + assert len(chunks) == 2 + assert chunks[0].finish_reason is None + assert chunks[0].tool_calls == [ + { + "index": 0, + "id": "call_write", + "type": "function", + "function": {"name": "write", "arguments": ""}, + } + ] + assert chunks[1].finish_reason == "tool_calls" + assert chunks[1].tool_calls == [ + { + "index": 0, + "id": "call_write", + "type": "function", + "function": { + "name": "write", + "arguments": first_arguments + remaining_arguments, + }, + } + ] + + class TestOpenAIBaseProviderStreamingUsage: @staticmethod def _build_provider_with_stream(): @@ -716,6 +866,53 @@ async def test_chat_stream_includes_usage_and_attaches_to_terminal_chunk(self): "total_tokens": 18, } + @pytest.mark.asyncio + async def test_chat_stream_preserves_nested_reasoning_tokens(self): + provider, create = self._build_provider_with_stream() + + content_chunk = SimpleNamespace( + choices=[ + SimpleNamespace( + delta=SimpleNamespace(content="hello", tool_calls=None), + finish_reason=None, + ) + ], + usage=None, + ) + usage_chunk = SimpleNamespace( + choices=[], + usage=SimpleNamespace( + prompt_tokens=11, + completion_tokens=7, + total_tokens=18, + completion_tokens_details=SimpleNamespace(reasoning_tokens=5), + ), + ) + finish_chunk = SimpleNamespace( + choices=[SimpleNamespace(delta=None, finish_reason="stop")], + usage=None, + ) + create.return_value = _stream_from_chunks( + content_chunk, usage_chunk, finish_chunk + ) + + from flocks.provider.provider import ChatMessage + + chunks = [ + chunk + async for chunk in provider.chat_stream( + "gpt-5.6-luna", + [ChatMessage(role="user", content="hello")], + ) + ] + + assert chunks[-1].usage == { + "prompt_tokens": 11, + "completion_tokens": 2, + "total_tokens": 18, + "reasoning_tokens": 5, + } + @pytest.mark.asyncio async def test_chat_stream_emits_trailing_usage_when_usage_only_chunk_arrives_after_finish(self): provider, create = self._build_provider_with_stream() diff --git a/tests/provider/test_openai_compatible_provider.py b/tests/provider/test_openai_compatible_provider.py index ee8389fe6..92efbe651 100644 --- a/tests/provider/test_openai_compatible_provider.py +++ b/tests/provider/test_openai_compatible_provider.py @@ -183,6 +183,49 @@ async def test_chat_does_not_fallback_for_completion_token_value_errors(self): assert create.await_count == 1 + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("model_id", "extra_body"), + [ + ("kimi-k2.7-code-highspeed", {"thinking": {"type": "enabled"}}), + ("kimi-k3", {"reasoning_effort": "max"}), + ], + ) + async def test_stream_logs_effective_kimi_thinking_state( + self, + model_id, + extra_body, + ): + provider, create = _build_provider_with_client() + provider.log = MagicMock() + create.return_value = _stream_from_chunks( + SimpleNamespace( + choices=[ + SimpleNamespace( + delta=SimpleNamespace(content="ok", tool_calls=None), + finish_reason="stop", + ) + ], + usage=None, + ) + ) + + _ = [ + chunk + async for chunk in provider.chat_stream( + model_id, + [ChatMessage(role="user", content="hello")], + extra_body=extra_body, + ) + ] + + request_logs = [ + call.args + for call in provider.log.info.call_args_list + if call.args[0] == "openai_compatible.stream.request" + ] + assert request_logs[0][1]["thinking_enabled"] is True + class TestOpenAICompatibleProviderMiniMaxFallback: def test_is_minimax_empty_response_target_matches_all_minimax_aliases(self): diff --git a/tests/provider/test_openai_provider.py b/tests/provider/test_openai_provider.py index cc78a6972..10f6211d4 100644 --- a/tests/provider/test_openai_provider.py +++ b/tests/provider/test_openai_provider.py @@ -1,9 +1,17 @@ -from unittest.mock import MagicMock, patch +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock, patch -from flocks.provider.provider import ProviderConfig +import pytest + +from flocks.provider.provider import ChatMessage, ProviderConfig from flocks.provider.sdk.openai import OpenAIProvider +async def _stream_from_chunks(*chunks): + for chunk in chunks: + yield chunk + + class TestOpenAIProviderConfiguration: @patch("httpx.AsyncClient") @patch("openai.AsyncOpenAI") @@ -40,3 +48,57 @@ def test_get_client_respects_verify_ssl_false(self, mock_async_openai, mock_http base_url="https://gateway.internal/v1", http_client=http_client, ) + + def test_configure_invalidates_existing_client_when_credentials_change(self): + provider = OpenAIProvider() + stale_client = MagicMock() + provider._client = stale_client + + provider.configure( + ProviderConfig( + provider_id=provider.id, + api_key="new-api-key", + base_url="https://new-gateway.internal/v1", + ) + ) + + assert provider._client is None + + +class TestOpenAIProviderStreamingUsage: + @pytest.mark.asyncio + async def test_chat_stream_preserves_nested_reasoning_tokens(self): + provider = OpenAIProvider() + create = AsyncMock() + provider._client = MagicMock() + provider._client.chat.completions.create = create + + usage_chunk = SimpleNamespace( + choices=[], + usage=SimpleNamespace( + prompt_tokens=11, + completion_tokens=7, + total_tokens=18, + completion_tokens_details=SimpleNamespace(reasoning_tokens=5), + ), + ) + finish_chunk = SimpleNamespace( + choices=[SimpleNamespace(delta=None, finish_reason="stop")], + usage=None, + ) + create.return_value = _stream_from_chunks(finish_chunk, usage_chunk) + + chunks = [ + chunk + async for chunk in provider.chat_stream( + "gpt-5.6-luna", + [ChatMessage(role="user", content="hello")], + ) + ] + + assert chunks[-1].usage == { + "prompt_tokens": 11, + "completion_tokens": 2, + "total_tokens": 18, + "reasoning_tokens": 5, + } diff --git a/tests/provider/test_provider.py b/tests/provider/test_provider.py index 6309bb141..c4d9fa843 100644 --- a/tests/provider/test_provider.py +++ b/tests/provider/test_provider.py @@ -221,6 +221,7 @@ def test_resolve_model_does_not_infer_interleaved_for_non_reasoning_model(monkey [ ("openai-compatible", "qwen3-235b-a22b-thinking", "https://api.example.com/v1", "reasoning_content"), ("openai-compatible", "kimi-k2-thinking-turbo", "https://api.example.com/v1", "reasoning_content"), + ("openai-compatible", "kimi-k3", "https://api.example.com/v1", "reasoning_content"), ("openai-compatible", "deepseek-v4-pro", "https://api.deepseek.com/v1", "reasoning_content"), ("openai-compatible", "glm-4.7", "https://api.example.com/v1", "reasoning_content"), ("openai-compatible", "minimax-m3", "https://api.example.com/v1", "reasoning_details"), diff --git a/tests/provider/test_provider_apply_config_idempotent.py b/tests/provider/test_provider_apply_config_idempotent.py index b00af163d..3f8aa54c8 100644 --- a/tests/provider/test_provider_apply_config_idempotent.py +++ b/tests/provider/test_provider_apply_config_idempotent.py @@ -67,7 +67,11 @@ def __exit__(self, *_args, **_kwargs): self.real_provider.configure = self._original_configure -def _build_fake_config(provider_id: str) -> Any: +def _build_fake_config( + provider_id: str, + *, + model_first_chunk_timeout_s: int = 480, +) -> Any: """Return a SimpleNamespace mimicking ``ConfigInfo`` shape needed by ``apply_config``: ``.provider`` dict -> ``.options`` (with model_dump) and ``.models``. @@ -88,6 +92,7 @@ def _build_fake_config(provider_id: str) -> Any: "supports_vision": False, "supports_reasoning": False, "max_tokens": 4096, + "stream_first_chunk_timeout_s": model_first_chunk_timeout_s, } ) } @@ -128,6 +133,9 @@ async def test_apply_config_is_idempotent_for_unchanged_input() -> None: "apply_config must not rebuild _config_models when the desired " "model list matches the existing one" ) + assert provider.get_models()[0].custom_settings == { + "stream_first_chunk_timeout_s": 480, + } @pytest.mark.asyncio @@ -171,3 +179,87 @@ async def test_apply_config_still_mutates_when_input_changes() -> None: assert rec.configure_calls == 1, ( f"expected exactly 1 configure call on api_key change, got {rec.configure_calls}" ) + + +@pytest.mark.asyncio +async def test_apply_config_rebuilds_models_when_stream_timeout_changes() -> None: + """A model timeout change must invalidate the model-list signature.""" +@pytest.mark.parametrize( + "options", + [ + None, + SimpleNamespace( + model_dump=lambda exclude_none, by_alias: { + "api_key": " ", + "base_url": "", + } + ), + ], + ids=["missing-options", "empty-credentials"], +) +async def test_apply_config_loads_name_and_models_without_credentials( + options: Any, +) -> None: + """Provider metadata must not depend on resolved credentials.""" + Provider._ensure_initialized() + provider = Provider.get("openai-compatible") + assert provider is not None + + first_cfg = _build_fake_config( + "openai-compatible", + model_first_chunk_timeout_s=480, + ) + with patch("flocks.provider.provider.Config.get", return_value=first_cfg): + await Provider.apply_config(provider_id="openai-compatible") + + second_cfg = _build_fake_config( + "openai-compatible", + model_first_chunk_timeout_s=600, + ) + with patch( + "flocks.provider.provider.Config.get", + return_value=second_cfg, + ), _Recorder(provider) as rec: + await Provider.apply_config(provider_id="openai-compatible") + + assert rec.models_assignments == 1 + assert provider.get_models()[0].custom_settings == { + "stream_first_chunk_timeout_s": 600, + } + original_config = provider._config + original_models = list(getattr(provider, "_config_models", []) or []) + original_name = provider.name + model_id = f"credentialless-model-{id(options)}" + fake_cfg = SimpleNamespace( + provider={ + "openai-compatible": SimpleNamespace( + name="Credentialless Provider", + options=options, + models={ + model_id: { + "name": "Credentialless Model", + "supports_tools": True, + } + }, + ) + } + ) + + try: + with _Recorder(provider) as first: + await Provider.apply_config(fake_cfg, provider_id="openai-compatible") + + assert first.configure_calls == 0 + assert first.models_assignments == 1 + assert provider.name == "Credentialless Provider" + assert [model.id for model in provider._config_models] == [model_id] + + with _Recorder(provider) as second: + await Provider.apply_config(fake_cfg, provider_id="openai-compatible") + + assert second.configure_calls == 0 + assert second.models_assignments == 0 + finally: + provider._config = original_config + provider._config_models = original_models + provider.name = original_name diff --git a/tests/provider/test_provider_options.py b/tests/provider/test_provider_options.py index 4e8f11be0..45f2af82b 100644 --- a/tests/provider/test_provider_options.py +++ b/tests/provider/test_provider_options.py @@ -58,6 +58,125 @@ def test_moonshot_kimi_hybrid_models_use_official_thinking_payload(self): assert options["extra_body"] == KIMI_THINKING_EXTRA_BODY + def test_moonshot_kimi_k27_forces_official_thinking_payload(self): + options = provider_options.build_provider_options( + "moonshot", + "kimi-k2.7-code", + reasoning_enabled=False, + resolve_max_tokens=False, + ) + + assert options["extra_body"] == KIMI_THINKING_EXTRA_BODY + + def test_moonshot_kimi_k27_highspeed_forces_official_thinking_payload(self): + options = provider_options.build_provider_options( + "moonshot", + "kimi-k2.7-code-highspeed", + reasoning_enabled=False, + resolve_max_tokens=False, + ) + + assert options["extra_body"] == KIMI_THINKING_EXTRA_BODY + + def test_moonshot_kimi_k3_uses_default_reasoning_effort(self): + options = provider_options.build_provider_options( + "moonshot", + "kimi-k3", + reasoning_enabled=False, + resolve_max_tokens=False, + ) + + assert options["extra_body"] == {"reasoning_effort": "max"} + + def test_kimi_k27_forces_thinking_even_when_toggle_is_disabled(self): + options = provider_options.build_provider_options( + "threatbook-cn-llm", + "kimi-k2.7-code-highspeed", + reasoning_enabled=False, + resolve_max_tokens=False, + ) + + assert options["extra_body"] == KIMI_THINKING_EXTRA_BODY + + def test_kimi_k27_overrides_configured_disabled_thinking(self, monkeypatch): + monkeypatch.setattr( + provider_options, + "_resolve_default_extra_body", + lambda *_args: { + "thinking": {"type": "disabled"}, + "custom_option": True, + }, + ) + + options = provider_options.build_provider_options( + "openai-compatible", + "kimi-k2.7-code", + resolve_max_tokens=False, + ) + + assert options["extra_body"] == { + "thinking": {"type": "enabled"}, + "custom_option": True, + } + + def test_kimi_k3_uses_reasoning_effort_instead_of_thinking(self): + options = provider_options.build_provider_options( + "openai-compatible", + "kimi-k3", + reasoning_enabled=False, + resolve_max_tokens=False, + ) + + assert options["extra_body"] == {"reasoning_effort": "max"} + assert "thinking" not in options["extra_body"] + + def test_kimi_k3_respects_supported_reasoning_effort(self): + options = provider_options.build_provider_options( + "openai-compatible", + "kimi-k3", + reasoning_effort="low", + resolve_max_tokens=False, + ) + + assert options["extra_body"] == {"reasoning_effort": "low"} + + def test_kimi_k3_preserves_configured_extra_body_reasoning_effort( + self, + monkeypatch, + ): + monkeypatch.setattr( + provider_options, + "_resolve_default_extra_body", + lambda *_args: {"reasoning_effort": "high"}, + ) + + options = provider_options.build_provider_options( + "openai-compatible", + "kimi-k3", + resolve_max_tokens=False, + ) + + assert options["extra_body"] == {"reasoning_effort": "high"} + + def test_kimi_k3_explicit_reasoning_effort_overrides_configured_default( + self, + monkeypatch, + ): + monkeypatch.setattr( + provider_options, + "_resolve_default_extra_body", + lambda *_args: {"reasoning_effort": "high"}, + ) + + options = provider_options.build_provider_options( + "openai-compatible", + "kimi-k3", + reasoning_effort="low", + resolve_max_tokens=False, + ) + + assert options["extra_body"] == {"reasoning_effort": "low"} + def test_openai_compatible_qwen_models_enable_thinking_by_default(self, monkeypatch): monkeypatch.setattr( provider_options, diff --git a/tests/provider/test_thinking_params.py b/tests/provider/test_thinking_params.py index db82214f1..d9c638214 100644 --- a/tests/provider/test_thinking_params.py +++ b/tests/provider/test_thinking_params.py @@ -47,6 +47,7 @@ from flocks.provider import model_catalog from flocks.provider import options as provider_options +from flocks.provider.interleaved import is_kimi_k27_code_model, is_kimi_k3_model DEEPSEEK_THINKING_EXTRA_BODY = {"thinking": {"type": "enabled"}} GLM_THINKING_EXTRA_BODY = {"thinking": {"type": "enabled", "clear_thinking": False}} @@ -98,6 +99,10 @@ def _expected_generic_chat_extra_body( return GLM_THINKING_EXTRA_BODY if "mimo" in model_lower: return MIMO_THINKING_EXTRA_BODY + if is_kimi_k3_model(model_id): + return {"reasoning_effort": "max"} + if is_kimi_k27_code_model(model_id): + return KIMI_THINKING_EXTRA_BODY if "kimi" in model_lower: return KIMI_THINKING_EXTRA_BODY if "minimax" in model_lower or provider_lower == "minimax": @@ -437,6 +442,9 @@ def test_anthropic_transport_still_uses_thinking_field( ("qwen3-7b-uncatalogued", {"enable_thinking": True}), ("glm-5-uncatalogued", GLM_THINKING_EXTRA_BODY), ("kimi-k2.6-uncatalogued", KIMI_THINKING_EXTRA_BODY), + ("kimi-k2.7-code", KIMI_THINKING_EXTRA_BODY), + ("kimi-k2.7-code-highspeed", KIMI_THINKING_EXTRA_BODY), + ("kimi-k3", {"reasoning_effort": "max"}), ("mimo-v2.5-pro-uncatalogued", MIMO_THINKING_EXTRA_BODY), ("minimax-m4-uncatalogued", {"reasoning_split": True}), ("step-3.5-flash-uncatalogued", {"enable_thinking": True}), @@ -465,6 +473,16 @@ def test_series_token_fallback_emits_expected_extra_body( f"emitted {expected_extra_body}. options={options!r}" ) + def test_kimi_k27_dispatch_does_not_match_unknown_suffix(self) -> None: + options = provider_options.build_provider_options( + "openai-compatible", + "kimi-k2.7-code-future", + reasoning_enabled=False, + resolve_max_tokens=False, + ) + + assert "extra_body" not in options + @pytest.mark.parametrize( "provider_id,model_id,expected_extra_body", [ diff --git a/tests/sandbox/test_sandbox_runtime_integration.py b/tests/sandbox/test_sandbox_runtime_integration.py index 734e0959c..7223845af 100644 --- a/tests/sandbox/test_sandbox_runtime_integration.py +++ b/tests/sandbox/test_sandbox_runtime_integration.py @@ -297,6 +297,13 @@ async def fake_resolve_sandbox_context(**_kwargs): assert len(captured_extras) == 2 assert captured_extras[0][0] == "bash" - assert captured_extras[0][1] == {} + assert captured_extras[0][1] == { + "execution_mode": "build", + "workspace_dir": "/tmp", + "model": {"providerID": "test-provider", "modelID": "test-model"}, + "plan_file_path": None, + "plan_relative_path": None, + "plan_permission_path": None, + } assert "sandbox" in captured_extras[1][1] assert captured_extras[1][1]["sandbox"]["container_name"] == "flocks-sbx-test" diff --git a/tests/security/test_secret_resolution.py b/tests/security/test_secret_resolution.py index c29d261e1..dd8d54f7a 100644 --- a/tests/security/test_secret_resolution.py +++ b/tests/security/test_secret_resolution.py @@ -25,6 +25,20 @@ def test_resolve_secret_value_prefers_real_fofa_split_secrets(): assert resolve_secret_value("fofa_api_key", secrets) == "override-api-key" +def test_resolve_secret_value_reads_legacy_llm_key_without_mutating_secrets(): + secrets = MagicMock() + secrets.get.side_effect = lambda key: { + "custom-codex_api_key": "legacy-provider-key", + }.get(key) + + assert ( + resolve_secret_value("custom-codex_llm_key", secrets) + == "legacy-provider-key" + ) + secrets.set.assert_not_called() + secrets.delete.assert_not_called() + + def test_resolve_secret_refs_returns_empty_for_invalid_fofa_canonical_secret(): secrets = MagicMock() secrets.get.side_effect = lambda key: { diff --git a/tests/server/routes/test_admin_users_routes.py b/tests/server/routes/test_admin_users_routes.py index 1ed56409e..b813d7f12 100644 --- a/tests/server/routes/test_admin_users_routes.py +++ b/tests/server/routes/test_admin_users_routes.py @@ -96,3 +96,16 @@ async def test_admin_routes_delete_user_not_allowed(client: AsyncClient): response = await client.delete(f"/api/admin/users/{user_id}") assert response.status_code == 404, response.text + + +@pytest.mark.asyncio +async def test_service_owner_username_is_reserved(): + from flocks.auth.context import API_TOKEN_SERVICE_USER_ID + from flocks.auth.service import AuthService + + with pytest.raises(ValueError, match="系统保留"): + await AuthService._create_user_internal( + username=API_TOKEN_SERVICE_USER_ID, + password="Password123!", + role="member", + ) diff --git a/tests/server/routes/test_project_routes.py b/tests/server/routes/test_project_routes.py index f147ebf13..f269b879b 100644 --- a/tests/server/routes/test_project_routes.py +++ b/tests/server/routes/test_project_routes.py @@ -1,7 +1,8 @@ +import asyncio from pathlib import Path import pytest -from fastapi import status +from fastapi import HTTPException, status from httpx import AsyncClient from flocks.auth.context import AuthUser @@ -48,6 +49,12 @@ async def test_project_crud_uses_registry_without_project_database_rows( ) assert session_response.status_code == status.HTTP_200_OK session_id = session_response.json()["id"] + other_session_response = await client.post( + "/api/session", + json={"title": "Keep archived", "projectID": project["id"]}, + ) + assert other_session_response.status_code == status.HTTP_200_OK + other_session_id = other_session_response.json()["id"] delete_response = await client.delete(f"/api/project/{project['id']}") assert delete_response.status_code == status.HTTP_200_OK @@ -55,6 +62,25 @@ async def test_project_crud_uses_registry_without_project_database_rows( assert retained_file.read_text(encoding="utf-8") == "project files stay" assert (await client.get("/api/project")).json() == [] assert (await client.get(f"/api/session/{session_id}")).status_code == status.HTTP_404_NOT_FOUND + archived = ( + await client.get( + "/api/session", + params={"status": "archived", "view": "list", "manager": "true", "roots": "true"}, + ) + ).json() + archived_by_id = {item["id"]: item for item in archived} + assert archived_by_id[session_id]["projectName"] == "Security Labs" + assert other_session_id in archived_by_id + assert (await Session.get(project["id"], session_id)).status == "archived" + + restore_response = await client.post(f"/api/session/{session_id}/restore") + assert restore_response.status_code == status.HTTP_200_OK + assert restore_response.json()["status"] == "active" + restored_projects = (await client.get("/api/project")).json() + assert [item["id"] for item in restored_projects] == [project["id"]] + assert restored_projects[0]["name"] == "Security Labs" + assert (await Session.get(project["id"], session_id)).status == "active" + assert (await Session.get(project["id"], other_session_id)).status == "archived" @pytest.mark.asyncio @@ -74,6 +100,309 @@ async def test_create_project_creates_missing_directory( assert worktree.is_dir() +@pytest.mark.asyncio +async def test_move_regular_task_into_project( + client: AsyncClient, + tmp_path: Path, +): + worktree = tmp_path / "move-target" + worktree.mkdir() + project = ( + await client.post( + "/api/project", + json={"name": "Move Target", "worktree": str(worktree)}, + ) + ).json() + session = ( + await client.post( + "/api/session", + json={"title": "Move me"}, + ) + ).json() + + response = await client.patch( + f"/api/session/{session['id']}/project", + json={"projectID": project["id"]}, + ) + + assert response.status_code == status.HTTP_200_OK + assert response.json()["projectID"] == project["id"] + assert response.json()["effectiveProjectID"] == project["id"] + assert response.json()["directory"] == str(worktree.resolve()) + assert await Session.get(session["projectID"], session["id"]) is None + moved = await Session.get(project["id"], session["id"]) + assert moved is not None and moved.directory == str(worktree.resolve()) + + +@pytest.mark.asyncio +async def test_move_rejects_active_prompt_chain( + client: AsyncClient, + tmp_path: Path, +): + from flocks.server.routes import session as session_routes + + worktree = tmp_path / "busy-move-target" + worktree.mkdir() + project = ( + await client.post( + "/api/project", + json={"name": "Busy Move Target", "worktree": str(worktree)}, + ) + ).json() + session = (await client.post("/api/session", json={"title": "Busy"})).json() + + session_routes._set_prompt_chain_active(session["id"], True) + try: + response = await client.patch( + f"/api/session/{session['id']}/project", + json={"projectID": project["id"]}, + ) + finally: + session_routes._set_prompt_chain_active(session["id"], False) + + assert response.status_code == status.HTTP_409_CONFLICT + assert await Session.get(session["projectID"], session["id"]) is not None + + +@pytest.mark.asyncio +async def test_move_blocks_replay_of_messages_from_previous_project( + client: AsyncClient, + tmp_path: Path, +): + worktree = tmp_path / "history-move-target" + worktree.mkdir() + project = ( + await client.post( + "/api/project", + json={"name": "History Move Target", "worktree": str(worktree)}, + ) + ).json() + session = (await client.post("/api/session", json={"title": "History"})).json() + assert ( + await client.post( + f"/api/session/{session['id']}/message", + json={ + "parts": [{"type": "text", "text": "Before moving"}], + "noReply": True, + "mockReply": "Old project reply", + }, + ) + ).status_code == status.HTTP_200_OK + messages = (await client.get(f"/api/session/{session['id']}/message")).json() + user_message = next(item for item in messages if item["info"]["role"] == "user") + text_part = next(part for part in user_message["parts"] if part["type"] == "text") + assert ( + await client.patch( + f"/api/session/{session['id']}/project", + json={"projectID": project["id"]}, + ) + ).status_code == status.HTTP_200_OK + + response = await client.post( + f"/api/session/{session['id']}/message/{user_message['info']['id']}/resend", + json={"text": "Do not replay", "partID": text_part["id"]}, + ) + + assert response.status_code == status.HTTP_409_CONFLICT + assert "移动项目前" in response.text + revert_response = await client.post( + f"/api/session/{session['id']}/revert", + json={"messageID": user_message["info"]["id"]}, + ) + assert revert_response.status_code == status.HTTP_409_CONFLICT + + +@pytest.mark.asyncio +async def test_clear_resets_move_boundary_for_new_message_replay( + client: AsyncClient, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +): + from flocks.server.routes import session as session_routes + + worktree = tmp_path / "clear-move-target" + worktree.mkdir() + project = ( + await client.post( + "/api/project", + json={"name": "Clear Move Target", "worktree": str(worktree)}, + ) + ).json() + session = (await client.post("/api/session", json={"title": "Clear History"})).json() + assert ( + await client.post( + f"/api/session/{session['id']}/message", + json={"parts": [{"type": "text", "text": "Before moving"}], "noReply": True}, + ) + ).status_code == status.HTTP_200_OK + assert ( + await client.patch( + f"/api/session/{session['id']}/project", + json={"projectID": project["id"]}, + ) + ).status_code == status.HTTP_200_OK + assert ( + await client.post(f"/api/session/{session['id']}/clear") + ).status_code == status.HTTP_200_OK + + moved = await Session.get(project["id"], session["id"]) + assert moved is not None + assert "projectMove" not in moved.metadata + + assert ( + await client.post( + f"/api/session/{session['id']}/message", + json={"parts": [{"type": "text", "text": "After clearing"}], "noReply": True}, + ) + ).status_code == status.HTTP_200_OK + messages = (await client.get(f"/api/session/{session['id']}/message")).json() + user_message = next(item for item in messages if item["info"]["role"] == "user") + text_part = next(part for part in user_message["parts"] if part["type"] == "text") + scheduled_coroutines = [] + + def close_scheduled_coro(coro, **_kwargs): + scheduled_coroutines.append(coro) + coro.close() + + monkeypatch.setattr(session_routes, "_schedule_background_coro", close_scheduled_coro) + + replay_response = await client.post( + f"/api/session/{session['id']}/message/{user_message['info']['id']}/resend", + json={"text": "Replay after clearing", "partID": text_part["id"]}, + ) + + assert replay_response.status_code == status.HTTP_202_ACCEPTED + assert len(scheduled_coroutines) == 1 + + +@pytest.mark.asyncio +async def test_restore_archived_task_keeps_project_removed_when_directory_is_missing( + client: AsyncClient, + tmp_path: Path, +): + worktree = tmp_path / "missing-after-remove" + worktree.mkdir() + project = ( + await client.post( + "/api/project", + json={"name": "Missing Project", "worktree": str(worktree)}, + ) + ).json() + session = ( + await client.post( + "/api/session", + json={"title": "Still Archived", "projectID": project["id"]}, + ) + ).json() + assert (await client.delete(f"/api/project/{project['id']}")).status_code == status.HTTP_200_OK + worktree.rmdir() + + restore_response = await client.post(f"/api/session/{session['id']}/restore") + + assert restore_response.status_code == status.HTTP_409_CONFLICT + assert "directory is unavailable" in restore_response.text + assert (await client.get("/api/project")).json() == [] + stored = await Session.get(project["id"], session["id"]) + assert stored is not None and stored.status == "archived" + + +@pytest.mark.asyncio +async def test_project_removal_waits_for_inflight_session_creation( + client: AsyncClient, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +): + worktree = tmp_path / "concurrent-remove" + worktree.mkdir() + project = ( + await client.post( + "/api/project", + json={"name": "Concurrent", "worktree": str(worktree)}, + ) + ).json() + create_started = asyncio.Event() + allow_create = asyncio.Event() + original_set = Storage.set + + async def blocked_set(key, value, *args, **kwargs): + if key.startswith(f"session:{project['id']}:"): + create_started.set() + await allow_create.wait() + return await original_set(key, value, *args, **kwargs) + + monkeypatch.setattr(Storage, "set", blocked_set) + create_task = asyncio.create_task( + client.post( + "/api/session", + json={"title": "Concurrent task", "projectID": project["id"]}, + ) + ) + await create_started.wait() + delete_task = asyncio.create_task(client.delete(f"/api/project/{project['id']}")) + + with pytest.raises(asyncio.TimeoutError): + await asyncio.wait_for(asyncio.shield(delete_task), timeout=0.05) + + allow_create.set() + create_response, delete_response = await asyncio.gather(create_task, delete_task) + + assert create_response.status_code == status.HTTP_200_OK + assert delete_response.status_code == status.HTTP_200_OK + created = await Session.get(project["id"], create_response.json()["id"]) + assert created is not None and created.status == "archived" + + +@pytest.mark.asyncio +async def test_project_delete_restores_earlier_tasks_when_later_archive_fails( + client: AsyncClient, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +): + from flocks.server.routes import session as session_routes + + worktree = tmp_path / "archive-rollback" + worktree.mkdir() + project = ( + await client.post( + "/api/project", + json={"name": "Archive rollback", "worktree": str(worktree)}, + ) + ).json() + session_ids = [ + ( + await client.post( + "/api/session", + json={"title": title, "projectID": project["id"]}, + ) + ).json()["id"] + for title in ("First", "Second") + ] + original_archive = session_routes.archive_session_for_user + calls = 0 + + async def fail_second_archive(session_id, user): + nonlocal calls + calls += 1 + if calls == 2: + raise HTTPException(status_code=409, detail="archive failed") + return await original_archive(session_id, user) + + monkeypatch.setattr( + session_routes, + "archive_session_for_user", + fail_second_archive, + ) + + response = await client.delete(f"/api/project/{project['id']}") + + assert response.status_code == status.HTTP_409_CONFLICT + assert [item["id"] for item in (await client.get("/api/project")).json()] == [ + project["id"] + ] + stored = [await Session.get(project["id"], session_id) for session_id in session_ids] + assert all(session is not None and session.status == "active" for session in stored) + + @pytest.mark.asyncio async def test_duplicate_project_directory_returns_existing_project( client: AsyncClient, @@ -273,6 +602,13 @@ async def test_project_counts_and_session_lists_are_user_scoped( "Second Alice session", } + archive_response = await client.post( + f"/api/session/{alice_session.json()['id']}/archive", + ) + assert archive_response.status_code == status.HTTP_200_OK + after_archive = (await client.get("/api/project")).json() + assert after_archive[0]["sessionCount"] == 1 + @pytest.mark.asyncio async def test_session_title_update_invalidates_project_search_stats( diff --git a/tests/server/routes/test_provider_optional_api_key.py b/tests/server/routes/test_provider_optional_api_key.py index 5aa342202..dd244acef 100644 --- a/tests/server/routes/test_provider_optional_api_key.py +++ b/tests/server/routes/test_provider_optional_api_key.py @@ -32,6 +32,8 @@ def patched_runtime(monkeypatch: pytest.MonkeyPatch): monkeypatch.setattr("flocks.security.get_secret_manager", lambda: fake_secrets) monkeypatch.setattr(provider_routes.Provider, "_ensure_initialized", MagicMock()) monkeypatch.setattr(provider_routes.Provider, "get", lambda _pid: runtime_provider) + clear_config_cache = MagicMock() + monkeypatch.setattr(provider_routes.Config, "clear_cache", clear_config_cache) # Pretend the provider already exists in flocks.json so set_provider_credentials # follows the "update existing" path and never calls add_provider() with real @@ -58,10 +60,30 @@ def patched_runtime(monkeypatch: pytest.MonkeyPatch): lambda _provider: {}, ) - return {"secrets": fake_secrets, "provider": runtime_provider} + return { + "secrets": fake_secrets, + "provider": runtime_provider, + "clear_config_cache": clear_config_cache, + } class TestOptionalApiKey: + @pytest.mark.asyncio + async def test_saving_credentials_invalidates_resolved_config_cache( + self, patched_runtime + ): + """The next request must not reapply a secret cached before this update.""" + result = await provider_routes.set_provider_credentials( + "custom-local-gateway", + provider_routes.ProviderCredentialRequest( + api_key="new-secret", + base_url="http://127.0.0.1:8317/v1", + ), + ) + + assert result["success"] is True + patched_runtime["clear_config_cache"].assert_called_once_with() + @pytest.mark.asyncio async def test_openai_compatible_accepts_empty_api_key(self, patched_runtime): """openai-compatible: empty api_key -> success, placeholder persisted.""" diff --git a/tests/server/routes/test_provider_route_responsiveness.py b/tests/server/routes/test_provider_route_responsiveness.py index ba6967e37..ba1f4a09f 100644 --- a/tests/server/routes/test_provider_route_responsiveness.py +++ b/tests/server/routes/test_provider_route_responsiveness.py @@ -211,3 +211,37 @@ async def _config(): assert heartbeat_ticks >= 3 assert response.all == [] + + +async def test_list_providers_exposes_runtime_credential_state( + monkeypatch: pytest.MonkeyPatch, +): + from types import SimpleNamespace + + from flocks.server.routes import provider as provider_routes + + async def _initialized() -> None: + return None + + async def _config(): + raise RuntimeError("no merged config needed for focused route test") + + runtime_provider = SimpleNamespace( + name="Missing Credentials", + is_configured=lambda: False, + ) + monkeypatch.setattr(provider_routes, "_ensure_provider_initialized", _initialized) + monkeypatch.setattr(provider_routes.Config, "get", _config) + monkeypatch.setattr( + provider_routes.ConfigWriter, + "list_provider_ids", + lambda: ["missing-credentials"], + ) + monkeypatch.setattr(provider_routes.Provider, "get", lambda _provider_id: runtime_provider) + monkeypatch.setattr(provider_routes.Provider, "list_models", lambda _provider_id: []) + + response = await provider_routes.list_providers() + + assert len(response.all) == 1 + assert response.all[0].configured is False + assert response.connected == ["missing-credentials"] diff --git a/tests/server/routes/test_session_routes.py b/tests/server/routes/test_session_routes.py index 05d0470f1..e20a70d07 100644 --- a/tests/server/routes/test_session_routes.py +++ b/tests/server/routes/test_session_routes.py @@ -14,12 +14,12 @@ import asyncio from types import SimpleNamespace -from unittest.mock import AsyncMock +from unittest.mock import AsyncMock, MagicMock import pytest from fastapi import HTTPException, status from httpx import AsyncClient -from flocks.auth.context import AuthUser +from flocks.auth.context import API_TOKEN_SERVICE_USER_ID, AuthUser from flocks.session.core.status import SessionStatus, SessionStatusBusy from flocks.session.message import ( Message, @@ -31,6 +31,16 @@ from flocks.session.orphan_tools import INTERRUPTED_TOOL_ERROR from flocks.session.session import Session + +def _use_webui_admin(monkeypatch: pytest.MonkeyPatch) -> AuthUser: + """Authenticate a route test as a browser user instead of the API token.""" + from flocks.server.routes import session as session_routes + + user = AuthUser(id="usr_admin", username="admin", role="admin", status="active") + monkeypatch.setattr(session_routes, "require_user", lambda _request: user) + return user + + # =========================================================================== # CRUD # =========================================================================== @@ -114,15 +124,110 @@ async def test_create_session_with_category(self, client: AsyncClient): assert resp.json()["category"] == "workflow" @pytest.mark.asyncio - async def test_create_session_with_api_token_is_ownerless(self, client: AsyncClient): - """Sessions created by API-token clients remain manageable by WebUI admins.""" + @pytest.mark.parametrize("category", [None, "entity-config", "workflow"]) + async def test_create_session_with_manual_auto_mode( + self, + client: AsyncClient, + monkeypatch: pytest.MonkeyPatch, + category: str | None, + ): + """Auto is persisted for supported WebUI conversation sessions.""" + from flocks.session.session_loop import SessionLoop + + _use_webui_admin(monkeypatch) + monkeypatch.setattr( + SessionLoop, + "validate_auto_configuration", + AsyncMock(return_value=(True, "available")), + ) + payload = {"title": "Auto Session", "model_auto": True} + if category is not None: + payload["category"] = category + resp = await client.post("/api/session", json=payload) + + assert resp.status_code == status.HTTP_200_OK + assert resp.json()["model_auto"] is True + assert resp.json()["model_pinned"] is False + assert resp.json()["category"] == (category or "user") + + @pytest.mark.asyncio + async def test_create_unsupported_session_rejects_auto( + self, + client: AsyncClient, + monkeypatch: pytest.MonkeyPatch, + ): + from flocks.session.session_loop import SessionLoop + + validate_auto = AsyncMock(return_value=(True, "available")) + monkeypatch.setattr( + SessionLoop, + "validate_auto_configuration", + validate_auto, + ) + + resp = await client.post( + "/api/session", + json={"category": "task", "model_auto": True}, + ) + + assert resp.status_code == status.HTTP_422_UNPROCESSABLE_ENTITY + assert ( + "only available for user, entity configuration, and workflow sessions" + in str(resp.json()) + ) + validate_auto.assert_not_awaited() + + @pytest.mark.asyncio + async def test_create_session_rejects_unavailable_auto( + self, + client: AsyncClient, + monkeypatch: pytest.MonkeyPatch, + ): + from flocks.session.session_loop import SessionLoop + + _use_webui_admin(monkeypatch) + monkeypatch.setattr( + SessionLoop, + "validate_auto_configuration", + AsyncMock(return_value=(False, "fallback_unavailable")), + ) + + resp = await client.post("/api/session", json={"model_auto": True}) + + assert resp.status_code == status.HTTP_400_BAD_REQUEST + assert "fallback_unavailable" in str(resp.json()) + + @pytest.mark.asyncio + async def test_api_token_cannot_create_auto_session( + self, + client: AsyncClient, + monkeypatch: pytest.MonkeyPatch, + ): + from flocks.session.session_loop import SessionLoop + + validate_auto = AsyncMock(return_value=(True, "available")) + monkeypatch.setattr( + SessionLoop, + "validate_auto_configuration", + validate_auto, + ) + + resp = await client.post("/api/session", json={"model_auto": True}) + + assert resp.status_code == status.HTTP_403_FORBIDDEN + assert "only be enabled from the WebUI" in str(resp.json()) + validate_auto.assert_not_awaited() + + @pytest.mark.asyncio + async def test_create_session_with_api_token_is_system_owned(self, client: AsyncClient): + """API-token sessions have an explicit owner manageable by WebUI admins.""" resp = await client.post("/api/session", json={"title": "TUI Session"}) assert resp.status_code == status.HTTP_200_OK session = await Session.get_by_id(resp.json()["id"]) assert session is not None - assert session.owner_user_id is None - assert session.owner_username is None + assert session.owner_user_id == API_TOKEN_SERVICE_USER_ID + assert session.owner_username == API_TOKEN_SERVICE_USER_ID @pytest.mark.asyncio async def test_create_session_with_local_user_keeps_owner( @@ -176,6 +281,16 @@ async def test_webui_admin_can_manage_api_token_session( ) assert message_resp.status_code == status.HTTP_200_OK + child_resp = await client.post( + "/api/session", + json={"title": "WebUI child", "parentID": session_id}, + ) + assert child_resp.status_code == status.HTTP_200_OK + child = await Session.get_by_id(child_resp.json()["id"]) + assert child is not None + assert child.owner_user_id == API_TOKEN_SERVICE_USER_ID + assert child.owner_username == API_TOKEN_SERVICE_USER_ID + rename_resp = await client.patch( f"/api/session/{session_id}", json={"title": "Renamed in WebUI"}, @@ -284,24 +399,255 @@ async def test_list_sessions_light_manager_filters_and_omits_heavy_fields(self, assert set(row) == { "id", "projectID", + "projectName", "effectiveProjectID", "directory", "title", "time", "category", + "status", "parentID", "provider", "model", "model_pinned", + "model_auto", + "ownerUserID", + "ownerUsername", "canWrite", "canDelete", "isShared", } assert row["projectID"] assert row["directory"] + assert "ownerUsername" in row assert "goal" not in row assert "summary" not in row + @pytest.mark.asyncio + async def test_archive_hides_session_preserves_history_and_restores_tree(self, client: AsyncClient): + parent_resp = await client.post("/api/session", json={"title": "Archive Parent"}) + parent_id = parent_resp.json()["id"] + project_id = parent_resp.json()["projectID"] + child_resp = await client.post( + "/api/session", + json={"title": "Archive Child", "parentID": parent_id}, + ) + child_id = child_resp.json()["id"] + await client.post( + f"/api/session/{parent_id}/message", + json={"parts": [{"type": "text", "text": "keep this history"}], "noReply": True}, + ) + + archive_resp = await client.post(f"/api/session/{parent_id}/archive") + + assert archive_resp.status_code == status.HTTP_200_OK + assert archive_resp.json()["status"] == "archived" + archived_parent = await Session.get(project_id, parent_id) + archived_child = await Session.get(project_id, child_id) + assert archived_parent is not None and archived_parent.status == "archived" + assert archived_child is not None and archived_child.status == "archived" + assert len(await Message.list(parent_id)) == 1 + + active_list = await client.get("/api/session", params={"status": "active"}) + archived_list = await client.get( + "/api/session", + params={"status": "archived", "view": "list", "manager": "true", "roots": "true"}, + ) + assert parent_id not in {item["id"] for item in active_list.json()} + assert parent_id in {item["id"] for item in archived_list.json()} + assert child_id not in {item["id"] for item in archived_list.json()} + archived_row = next(item for item in archived_list.json() if item["id"] == parent_id) + assert archived_row["canWrite"] is False + + get_resp = await client.get(f"/api/session/{parent_id}") + update_resp = await client.patch(f"/api/session/{parent_id}", json={"title": "Blocked"}) + message_resp = await client.post( + f"/api/session/{parent_id}/message", + json={"parts": [{"type": "text", "text": "blocked"}], "noReply": True}, + ) + assert get_resp.status_code == status.HTTP_404_NOT_FOUND + assert update_resp.status_code == status.HTTP_409_CONFLICT + assert message_resp.status_code == status.HTTP_409_CONFLICT + + child_restore_resp = await client.post(f"/api/session/{child_id}/restore") + late_child_resp = await client.post( + "/api/session", + json={"title": "Late Child", "parentID": parent_id}, + ) + assert child_restore_resp.status_code == status.HTTP_409_CONFLICT + assert late_child_resp.status_code == status.HTTP_409_CONFLICT + + restore_resp = await client.post(f"/api/session/{parent_id}/restore") + assert restore_resp.status_code == status.HTTP_200_OK + assert restore_resp.json()["status"] == "active" + restored_child = await Session.get(project_id, child_id) + assert restored_child is not None and restored_child.status == "active" + assert len(await Message.list(parent_id)) == 1 + + @pytest.mark.asyncio + async def test_archive_cancels_route_background_work( + self, + client: AsyncClient, + session_id: str, + ): + from flocks.server.routes import session as session_routes + + started = asyncio.Event() + stopped = asyncio.Event() + + async def background_work() -> None: + started.set() + try: + await asyncio.Event().wait() + finally: + stopped.set() + + session_routes._schedule_background_coro( + background_work(), + session_id=session_id, + action="test.background", + ) + await started.wait() + + response = await client.post(f"/api/session/{session_id}/archive") + + assert response.status_code == status.HTTP_200_OK + assert stopped.is_set() + + @pytest.mark.asyncio + async def test_list_status_filter_applies_before_pagination(self, client: AsyncClient): + archived_resp = await client.post("/api/session", json={"title": "Archived First"}) + active_resp = await client.post("/api/session", json={"title": "Active Second"}) + assert (await client.post(f"/api/session/{archived_resp.json()['id']}/archive")).status_code == 200 + + response = await client.get( + "/api/session", + params={"status": "active", "limit": "1", "offset": "0"}, + ) + + assert response.status_code == status.HTTP_200_OK + assert [item["id"] for item in response.json()] == [active_resp.json()["id"]] + + @pytest.mark.asyncio + async def test_archived_manager_admin_sees_all_and_member_only_owns( + self, + client: AsyncClient, + monkeypatch: pytest.MonkeyPatch, + ): + from flocks.server.routes import session as session_routes + + alice = AuthUser(id="usr_alice", username="alice", role="member", status="active") + bob = AuthUser(id="usr_bob", username="bob", role="member", status="active") + admin = AuthUser(id="usr_admin", username="admin", role="admin", status="active") + alice_session = await Session.create( + project_id="archive_permissions", + directory="/tmp", + title="Alice archive", + owner_user_id=alice.id, + owner_username=alice.username, + ) + bob_session = await Session.create( + project_id="archive_permissions", + directory="/tmp", + title="Bob shared archive", + owner_user_id=bob.id, + owner_username=bob.username, + metadata={"shared_local": True}, + ) + assert await Session.archive(alice_session.project_id, alice_session.id) is True + assert await Session.archive(bob_session.project_id, bob_session.id) is True + + params = { + "status": "archived", + "view": "list", + "manager": "true", + "roots": "true", + } + monkeypatch.setattr(session_routes, "require_user", lambda _request: admin) + admin_response = await client.get("/api/session", params=params) + + assert admin_response.status_code == status.HTTP_200_OK + admin_rows = {row["id"]: row for row in admin_response.json()} + assert {alice_session.id, bob_session.id}.issubset(admin_rows) + assert admin_rows[alice_session.id]["ownerUsername"] == "alice" + assert admin_rows[bob_session.id]["ownerUsername"] == "bob" + + monkeypatch.setattr(session_routes, "require_user", lambda _request: alice) + member_response = await client.get("/api/session", params=params) + + assert member_response.status_code == status.HTTP_200_OK + assert [row["id"] for row in member_response.json()] == [alice_session.id] + + ordinary_archived_response = await client.get( + "/api/session", + params={"status": "archived", "view": "list", "roots": "true"}, + ) + assert ordinary_archived_response.status_code == status.HTTP_200_OK + assert [row["id"] for row in ordinary_archived_response.json()] == [ + alice_session.id + ] + + all_status_response = await client.get( + "/api/session", + params={"status": "all", "view": "list", "roots": "true"}, + ) + assert bob_session.id not in { + row["id"] for row in all_status_response.json() + } + + @pytest.mark.asyncio + async def test_archive_rejects_tree_with_descendant_owned_by_another_user( + self, + client: AsyncClient, + monkeypatch: pytest.MonkeyPatch, + ): + from flocks.auth.context import reset_current_auth_user, set_current_auth_user + from flocks.server.routes import session as session_routes + from flocks.session.session import SessionInfo + from flocks.storage.storage import Storage + + alice = AuthUser(id="usr_alice", username="alice", role="member", status="active") + parent = await Session.create( + project_id="archive_mixed_owner_tree", + directory="/tmp", + title="Alice parent", + owner_user_id=alice.id, + owner_username=alice.username, + ) + child = SessionInfo( + projectID=parent.project_id, + directory=parent.directory, + title="Bob child", + parentID=parent.id, + ownerUserID="usr_bob", + ownerUsername="bob", + ) + await Storage.set( + f"session:{child.project_id}:{child.id}", + child, + "session", + ) + Session.invalidate_cache() + monkeypatch.setattr(session_routes, "require_user", lambda _request: alice) + + token = set_current_auth_user(alice) + try: + response = await client.post(f"/api/session/{parent.id}/archive") + finally: + reset_current_auth_user(token) + + assert response.status_code == status.HTTP_403_FORBIDDEN + stored_parent = await Storage.get( + f"session:{parent.project_id}:{parent.id}", + SessionInfo, + ) + stored_child = await Storage.get( + f"session:{child.project_id}:{child.id}", + SessionInfo, + ) + assert stored_parent is not None and stored_parent.status == "active" + assert stored_child is not None and stored_child.status == "active" + @pytest.mark.asyncio async def test_create_session_in_user_managed_project( self, @@ -333,6 +679,7 @@ async def test_create_session_in_user_managed_project( ) row = next(item for item in list_resp.json() if item["id"] == session_resp.json()["id"]) assert row["projectID"] == project["id"] + assert row["projectName"] == "Labs" assert row["effectiveProjectID"] == project["id"] assert row["directory"] == project["worktree"] @@ -368,12 +715,12 @@ async def test_legacy_session_is_grouped_under_tasks_without_rewrite( assert stored.project_id == "legacy-git-project" @pytest.mark.asyncio - async def test_deleted_project_sessions_are_removed( + async def test_deleted_project_sessions_are_archived_and_restorable( self, client: AsyncClient, tmp_path, ): - """Deleting a project removes its sessions but preserves its directory.""" + """Deleting a project archives its sessions and restoring one revives the project.""" worktree = tmp_path / "removable-project" worktree.mkdir() project_response = await client.post( @@ -399,7 +746,65 @@ async def test_deleted_project_sessions_are_removed( params={"view": "list", "manager": "true", "projectID": "tasks"}, ) assert all(item["id"] != session_id for item in tasks_response.json()) - assert await Session.get(project["id"], session_id) is None + archived = await Session.get(project["id"], session_id) + assert archived is not None and archived.status == "archived" + + restore_response = await client.post(f"/api/session/{session_id}/restore") + assert restore_response.status_code == status.HTTP_200_OK + assert restore_response.json()["status"] == "active" + assert [item["id"] for item in (await client.get("/api/project")).json()] == [project["id"]] + + @pytest.mark.asyncio + async def test_restore_rolls_project_back_when_session_restore_fails( + self, + client: AsyncClient, + tmp_path, + monkeypatch: pytest.MonkeyPatch, + ): + worktree = tmp_path / "restore-rollback" + worktree.mkdir() + project = ( + await client.post( + "/api/project", + json={"name": "Rollback", "worktree": str(worktree)}, + ) + ).json() + session_id = ( + await client.post( + "/api/session", + json={"title": "Rollback Me", "projectID": project["id"]}, + ) + ).json()["id"] + assert (await client.delete(f"/api/project/{project['id']}")).status_code == 200 + + async def fail_unarchive(_project_id: str, _session_id: str) -> bool: + return False + + monkeypatch.setattr(Session, "unarchive", fail_unarchive) + response = await client.post(f"/api/session/{session_id}/restore") + + assert response.status_code == status.HTTP_409_CONFLICT + assert (await client.get("/api/project")).json() == [] + archived = await Session.get(project["id"], session_id) + assert archived is not None and archived.status == "archived" + + @pytest.mark.asyncio + async def test_restore_rejects_missing_managed_project_metadata( + self, + client: AsyncClient, + ): + session = await Session.create( + project_id="prj_missing_restore_metadata", + directory="/tmp", + title="Missing project metadata", + ) + assert await Session.archive(session.project_id, session.id) is True + + response = await client.post(f"/api/session/{session.id}/restore") + + assert response.status_code == status.HTTP_409_CONFLICT + archived = await Session.get(session.project_id, session.id) + assert archived is not None and archived.status == "archived" @pytest.mark.asyncio async def test_get_session(self, client: AsyncClient, session_id: str): @@ -554,193 +959,213 @@ async def test_update_session_title(self, client: AsyncClient, session_id: str): assert resp.json()["title"] == "Updated Title" @pytest.mark.asyncio - async def test_update_session_not_found(self, client: AsyncClient): - """PATCH for unknown session returns 404.""" - resp = await client.patch( - "/api/session/ses_nonexistent00000000000000", - json={"title": "X"}, - ) - assert resp.status_code == status.HTTP_404_NOT_FOUND - - @pytest.mark.asyncio - async def test_delete_session(self, client: AsyncClient, session_id: str): - """DELETE /api/session/{id} removes the session.""" - resp = await client.delete(f"/api/session/{session_id}") - assert resp.status_code == status.HTTP_200_OK - assert resp.json() is True - - # Confirm it is gone - get_resp = await client.get(f"/api/session/{session_id}") - assert get_resp.status_code == status.HTTP_404_NOT_FOUND - - @staticmethod - def _patch_delete_session_dependencies( - monkeypatch, - session_routes, - *, + async def test_update_session_auto_and_concrete_model_are_exclusive( + self, + client: AsyncClient, session_id: str, - order: list[str], - session_list, - ) -> None: - async def fake_abort_session_processing(abort_session_id: str) -> bool: - order.append(f"abort:{abort_session_id}") - return True + monkeypatch: pytest.MonkeyPatch, + ): + from flocks.session.session_loop import SessionLoop - async def fake_wait_for_sessions_idle(session_ids: list[str], timeout_s: float = 5.0) -> None: - order.append(f"wait:{','.join(session_ids)}") + _use_webui_admin(monkeypatch) + monkeypatch.setattr( + SessionLoop, + "validate_auto_configuration", + AsyncMock(return_value=(True, "available")), + ) + auto_resp = await client.patch( + f"/api/session/{session_id}", + json={"model_auto": True}, + ) + assert auto_resp.status_code == status.HTTP_200_OK + assert auto_resp.json()["model_auto"] is True + assert auto_resp.json()["model_pinned"] is False - async def fake_interaction_queue_clear(_session_id: str) -> None: - order.append("queue_clear") + invalid_resp = await client.patch( + f"/api/session/{session_id}", + json={ + "model_auto": True, + "provider": "openai", + "model": "gpt-4o", + "model_pinned": True, + }, + ) + assert invalid_resp.status_code == status.HTTP_422_UNPROCESSABLE_ENTITY - async def fake_goal_clear(_session_id: str) -> None: - order.append("goal_clear") + concrete_resp = await client.patch( + f"/api/session/{session_id}", + json={"provider": "openai", "model": "gpt-4o"}, + ) + assert concrete_resp.status_code == status.HTTP_200_OK + assert concrete_resp.json()["model_auto"] is False + assert concrete_resp.json()["model_pinned"] is True - async def fake_session_delete(_project_id: str, delete_session_id: str) -> bool: - assert delete_session_id == session_id - order.append("delete") - return True + @pytest.mark.asyncio + @pytest.mark.parametrize("category", ["entity-config", "workflow"]) + async def test_update_supported_sidebar_session_allows_auto( + self, + client: AsyncClient, + monkeypatch: pytest.MonkeyPatch, + category: str, + ): + from flocks.session.session_loop import SessionLoop - monkeypatch.setattr(session_routes.Session, "list", session_list) - monkeypatch.setattr( - session_routes, - "_abort_session_processing", - fake_abort_session_processing, - ) - monkeypatch.setattr( - session_routes, - "_wait_for_sessions_idle", - fake_wait_for_sessions_idle, + create_resp = await client.post( + "/api/session", + json={"title": "Sidebar Chat", "category": category}, ) + assert create_resp.status_code == status.HTTP_200_OK + + _use_webui_admin(monkeypatch) monkeypatch.setattr( - "flocks.session.interaction_queue.InteractionQueue.clear", - fake_interaction_queue_clear, + SessionLoop, + "validate_auto_configuration", + AsyncMock(return_value=(True, "available")), ) - monkeypatch.setattr( - "flocks.session.goal.GoalManager.clear", - fake_goal_clear, + resp = await client.patch( + f"/api/session/{create_resp.json()['id']}", + json={"model_auto": True}, ) - monkeypatch.setattr(session_routes.Session, "delete", fake_session_delete) + + assert resp.status_code == status.HTTP_200_OK + assert resp.json()["category"] == category + assert resp.json()["model_auto"] is True + assert resp.json()["model_pinned"] is False @pytest.mark.asyncio - async def test_delete_session_aborts_and_waits_before_delete( + async def test_update_session_rejects_unavailable_auto( self, client: AsyncClient, session_id: str, - monkeypatch, + monkeypatch: pytest.MonkeyPatch, ): - """DELETE waits for active processing to stop before clearing messages.""" - from flocks.server.routes import session as session_routes + from flocks.session.session_loop import SessionLoop - order: list[str] = [] - - async def fake_session_list(_project_id: str): - return [] - - self._patch_delete_session_dependencies( - monkeypatch, - session_routes, - session_id=session_id, - order=order, - session_list=fake_session_list, + _use_webui_admin(monkeypatch) + monkeypatch.setattr( + SessionLoop, + "validate_auto_configuration", + AsyncMock(return_value=(False, "primary_provider_not_configured")), ) - resp = await client.delete(f"/api/session/{session_id}") + resp = await client.patch( + f"/api/session/{session_id}", + json={"model_auto": True}, + ) - assert resp.status_code == status.HTTP_200_OK - assert resp.json() is True - assert order == [ - f"abort:{session_id}", - "queue_clear", - "goal_clear", - f"wait:{session_id}", - "delete", - ] + assert resp.status_code == status.HTTP_400_BAD_REQUEST + assert "primary_provider_not_configured" in str(resp.json()) @pytest.mark.asyncio - async def test_delete_session_waits_for_descendant_loops_before_delete( + async def test_pinning_existing_auto_session_disables_auto( self, client: AsyncClient, session_id: str, - monkeypatch, + monkeypatch: pytest.MonkeyPatch, ): - """DELETE waits for child and grandchild loops before recursive delete.""" - from flocks.server.routes import session as session_routes + from flocks.session.session_loop import SessionLoop - child_id = "ses_delete_child_wait" - grandchild_id = "ses_delete_grandchild_wait" - order: list[str] = [] + _use_webui_admin(monkeypatch) + monkeypatch.setattr( + SessionLoop, + "validate_auto_configuration", + AsyncMock(return_value=(True, "available")), + ) + clear_state = MagicMock() + monkeypatch.setattr(SessionLoop, "clear_auto_failover_state", clear_state) + auto_resp = await client.patch( + f"/api/session/{session_id}", + json={"model_auto": True}, + ) + assert auto_resp.status_code == status.HTTP_200_OK - async def fake_session_list(_project_id: str): - return [ - SimpleNamespace(id=child_id, parent_id=session_id), - SimpleNamespace(id=grandchild_id, parent_id=child_id), - ] + pin_resp = await client.patch( + f"/api/session/{session_id}", + json={"model_pinned": True}, + ) - self._patch_delete_session_dependencies( - monkeypatch, - session_routes, - session_id=session_id, - order=order, - session_list=fake_session_list, + assert pin_resp.status_code == status.HTTP_200_OK + assert pin_resp.json()["model_pinned"] is True + assert pin_resp.json()["model_auto"] is False + clear_state.assert_called_once_with(session_id) + + @pytest.mark.asyncio + async def test_api_token_cannot_enable_auto_on_existing_session( + self, + client: AsyncClient, + session_id: str, + monkeypatch: pytest.MonkeyPatch, + ): + from flocks.session.session_loop import SessionLoop + + validate_auto = AsyncMock(return_value=(True, "available")) + monkeypatch.setattr( + SessionLoop, + "validate_auto_configuration", + validate_auto, ) - resp = await client.delete(f"/api/session/{session_id}") + resp = await client.patch( + f"/api/session/{session_id}", + json={"model_auto": True}, + ) - assert resp.status_code == status.HTTP_200_OK - assert resp.json() is True - assert order == [ - f"abort:{session_id}", - "queue_clear", - "goal_clear", - f"wait:{session_id}", - f"abort:{child_id}", - f"abort:{grandchild_id}", - f"wait:{child_id},{grandchild_id}", - "delete", - ] + assert resp.status_code == status.HTTP_403_FORBIDDEN + assert "only be enabled from the WebUI" in str(resp.json()) + validate_auto.assert_not_awaited() @pytest.mark.asyncio - async def test_delete_session_aborts_descendant_that_appears_after_parent_wait( + async def test_update_unsupported_session_rejects_auto( self, client: AsyncClient, - session_id: str, - monkeypatch, + monkeypatch: pytest.MonkeyPatch, ): - """DELETE re-collects descendants after parent abort to catch late children.""" - from flocks.server.routes import session as session_routes + from flocks.session.session_loop import SessionLoop - child_id = "ses_delete_late_child_wait" - list_calls = 0 - order: list[str] = [] + create_resp = await client.post( + "/api/session", + json={"title": "Non-user Session", "category": "task"}, + ) + assert create_resp.status_code == status.HTTP_200_OK + session_id = create_resp.json()["id"] + validate_auto = AsyncMock(return_value=(True, "available")) + monkeypatch.setattr( + SessionLoop, + "validate_auto_configuration", + validate_auto, + ) - async def fake_session_list(_project_id: str): - nonlocal list_calls - list_calls += 1 - if list_calls == 1: - return [] - return [SimpleNamespace(id=child_id, parent_id=session_id)] + resp = await client.patch( + f"/api/session/{session_id}", + json={"model_auto": True}, + ) - self._patch_delete_session_dependencies( - monkeypatch, - session_routes, - session_id=session_id, - order=order, - session_list=fake_session_list, + assert resp.status_code == status.HTTP_422_UNPROCESSABLE_ENTITY + assert ( + "only available for user, entity configuration, and workflow sessions" + in str(resp.json()) ) + validate_auto.assert_not_awaited() - resp = await client.delete(f"/api/session/{session_id}") + @pytest.mark.asyncio + async def test_update_session_not_found(self, client: AsyncClient): + """PATCH for unknown session returns 404.""" + resp = await client.patch( + "/api/session/ses_nonexistent00000000000000", + json={"title": "X"}, + ) + assert resp.status_code == status.HTTP_404_NOT_FOUND + @pytest.mark.asyncio + async def test_delete_session(self, client: AsyncClient, session_id: str): + """DELETE /api/session/{id} removes the session.""" + resp = await client.delete(f"/api/session/{session_id}") assert resp.status_code == status.HTTP_200_OK assert resp.json() is True - assert order == [ - f"abort:{session_id}", - "queue_clear", - "goal_clear", - f"wait:{session_id}", - f"abort:{child_id}", - f"wait:{child_id}", - "delete", - ] + + # Confirm it is gone + get_resp = await client.get(f"/api/session/{session_id}") + assert get_resp.status_code == status.HTTP_404_NOT_FOUND @pytest.mark.asyncio async def test_delete_session_not_found(self, client: AsyncClient): @@ -765,6 +1190,24 @@ async def test_list_messages_empty(self, client: AsyncClient, session_id: str): assert isinstance(data, list) assert len(data) == 0 + @pytest.mark.asyncio + async def test_list_messages_reports_storage_failure( + self, + client: AsyncClient, + session_id: str, + monkeypatch: pytest.MonkeyPatch, + ): + """A read failure must not look like a successfully empty history.""" + monkeypatch.setattr( + Message, + "list_with_parts", + AsyncMock(side_effect=RuntimeError("storage unavailable")), + ) + + resp = await client.get(f"/api/session/{session_id}/message") + + assert resp.status_code == status.HTTP_500_INTERNAL_SERVER_ERROR + @pytest.mark.asyncio async def test_send_message_noReply(self, client: AsyncClient, session_id: str): """POST /api/session/{id}/message with noReply=True stores without triggering LLM.""" @@ -1057,6 +1500,47 @@ async def test_message_to_unknown_session_returns_404(self, client: AsyncClient) ) assert resp.status_code == status.HTTP_404_NOT_FOUND + @pytest.mark.asyncio + async def test_message_is_not_persisted_if_session_archives_after_preflight( + self, + client: AsyncClient, + session_id: str, + monkeypatch: pytest.MonkeyPatch, + ): + from flocks.provider.provider import Provider + from flocks.server.routes import session as session_routes + + apply_config_entered = asyncio.Event() + release_apply_config = asyncio.Event() + original_apply_config = Provider.apply_config + + async def blocked_apply_config(*args, **kwargs): + await original_apply_config(*args, **kwargs) + apply_config_entered.set() + await release_apply_config.wait() + + monkeypatch.setattr(Provider, "apply_config", blocked_apply_config) + + message_task = asyncio.create_task( + client.post( + f"/api/session/{session_id}/message", + json={ + "parts": [{"type": "text", "text": "must not persist"}], + "noReply": True, + }, + ) + ) + await apply_config_entered.wait() + archive_response = await client.post(f"/api/session/{session_id}/archive") + restore_response = await client.post(f"/api/session/{session_id}/restore") + release_apply_config.set() + message_response = await message_task + + assert archive_response.status_code == status.HTTP_200_OK + assert restore_response.status_code == status.HTTP_200_OK + assert message_response.status_code == status.HTTP_409_CONFLICT + assert await Message.list(session_id) == [] + @pytest.mark.asyncio async def test_resend_user_message_updates_text_and_truncates_followups( self, @@ -1444,8 +1928,191 @@ async def test_prepare_replay_runtime_uses_current_model_resolution( "agent_name": "rex", "provider_id": "openai", "model_id": "gpt-4.1", + "auto_failover": False, } + @pytest.mark.asyncio + async def test_auto_replay_reports_actual_fallback_model( + self, + monkeypatch: pytest.MonkeyPatch, + ): + """Replay passes Auto authorization and publishes the recovered model.""" + from flocks.server.routes import session as session_routes + from flocks.session.session_loop import LoopResult, SessionLoop + + user_message = SimpleNamespace(id="msg_user", agent="rex") + session = SimpleNamespace(id="ses_auto", directory="/tmp/project") + assistant = SimpleNamespace( + id="msg_assistant", + providerID="fallback", + modelID="fallback-model", + finish="stop", + tokens=None, + time={"created": 123}, + ) + run = AsyncMock(return_value=LoopResult( + action="stop", + last_message=assistant, + provider_id="fallback", + model_id="fallback-model", + )) + publish = AsyncMock() + context_usage = AsyncMock() + + monkeypatch.setattr(SessionLoop, "run", run) + monkeypatch.setattr( + "flocks.session.lifecycle.revert.SessionRevert.cleanup", + AsyncMock(), + ) + monkeypatch.setattr( + "flocks.session.message.Message.get_text_content", + AsyncMock(return_value="recovered"), + ) + monkeypatch.setattr("flocks.server.routes.event.publish_event", publish) + monkeypatch.setattr( + session_routes, + "_publish_context_usage_update", + context_usage, + ) + + result = await session_routes._run_existing_user_message( + session.id, + session, + user_message, + session.directory, + runtime={ + "agent_name": "rex", + "provider_id": "primary", + "model_id": "primary-model", + "auto_failover": True, + }, + ) + + assert result["status"] == "completed" + assert run.await_args.kwargs["auto_failover"] is True + completion = next( + call.args[1] + for call in publish.await_args_list + if call.args[0] == "message.updated" + ) + assert completion["info"]["providerID"] == "fallback" + assert completion["info"]["modelID"] == "fallback-model" + assert context_usage.await_args.kwargs["provider_id"] == "fallback" + assert context_usage.await_args.kwargs["model_id"] == "fallback-model" + + @pytest.mark.asyncio + @pytest.mark.parametrize("category", ["user", "entity-config", "workflow"]) + async def test_prepare_auto_replay_defers_unavailable_primary_to_failover( + self, + monkeypatch: pytest.MonkeyPatch, + category: str, + ): + from flocks.server.routes import session as session_routes + from flocks.session.session_loop import SessionLoop + + user_message = SimpleNamespace(agent="rex") + monkeypatch.setattr( + "flocks.agent.registry.Agent.get", + AsyncMock(return_value=SimpleNamespace(name="rex", model=None)), + ) + monkeypatch.setattr( + "flocks.session.session.Session.get_by_id", + AsyncMock(return_value=SimpleNamespace( + model_auto=True, + category=category, + )), + ) + monkeypatch.setattr( + "flocks.config.config.Config.resolve_default_llm", + AsyncMock(return_value={ + "provider_id": "primary", + "model_id": "primary-model", + }), + ) + monkeypatch.setattr( + "flocks.config.config.Config.get", + AsyncMock(return_value=SimpleNamespace()), + ) + validate = AsyncMock(return_value=(False, "provider_not_configured")) + monkeypatch.setattr(SessionLoop, "validate_runtime_model", validate) + monkeypatch.setattr("flocks.provider.provider.Provider._ensure_initialized", lambda: None) + monkeypatch.setattr("flocks.provider.provider.Provider.apply_config", AsyncMock()) + monkeypatch.setattr("flocks.provider.provider.Provider.get", lambda _provider_id: None) + resolve = AsyncMock() + monkeypatch.setattr(session_routes, "_resolve_model", resolve) + + runtime = await session_routes._prepare_replay_runtime("ses_auto", user_message) + + assert runtime == { + "agent_name": "rex", + "provider_id": "primary", + "model_id": "primary-model", + "auto_failover": True, + } + validate.assert_not_awaited() + resolve.assert_not_awaited() + + @pytest.mark.asyncio + async def test_prepare_replay_ignores_auto_on_unsupported_session( + self, + monkeypatch: pytest.MonkeyPatch, + ): + from flocks.server.routes import session as session_routes + from flocks.session.session_loop import SessionLoop + + user_message = SimpleNamespace(agent="rex") + monkeypatch.setattr( + "flocks.agent.registry.Agent.get", + AsyncMock(return_value=SimpleNamespace(name="rex", model=None)), + ) + monkeypatch.setattr( + "flocks.session.session.Session.get_by_id", + AsyncMock(return_value=SimpleNamespace( + model_auto=True, + category="task", + )), + ) + resolve = AsyncMock(return_value=("direct", "direct-model", "session")) + monkeypatch.setattr(session_routes, "_resolve_model", resolve) + default_llm = AsyncMock() + monkeypatch.setattr( + "flocks.config.config.Config.resolve_default_llm", + default_llm, + ) + monkeypatch.setattr( + "flocks.config.config.Config.get", + AsyncMock(return_value=SimpleNamespace()), + ) + validate = AsyncMock() + monkeypatch.setattr(SessionLoop, "validate_runtime_model", validate) + monkeypatch.setattr( + "flocks.provider.provider.Provider._ensure_initialized", + lambda: None, + ) + monkeypatch.setattr( + "flocks.provider.provider.Provider.apply_config", + AsyncMock(), + ) + monkeypatch.setattr( + "flocks.provider.provider.Provider.get", + lambda _provider_id: object(), + ) + + runtime = await session_routes._prepare_replay_runtime( + "ses_workflow", + user_message, + ) + + assert runtime == { + "agent_name": "rex", + "provider_id": "direct", + "model_id": "direct-model", + "auto_failover": False, + } + resolve.assert_awaited_once() + default_llm.assert_not_awaited() + validate.assert_not_awaited() + @pytest.mark.asyncio async def test_resend_uses_current_model_for_replay( self, @@ -1679,7 +2346,11 @@ async def fake_publish_event(_event: str, _payload: dict) -> None: monkeypatch.setattr( session_routes.Session, "get_by_id", - AsyncMock(return_value=SimpleNamespace(id=session_id, directory="/tmp/project")), + AsyncMock(return_value=SimpleNamespace( + id=session_id, + directory="/tmp/project", + status="active", + )), ) monkeypatch.setattr(session_routes, "abort_session", fake_abort_session) monkeypatch.setattr(session_routes, "_wait_for_session_idle", fake_wait_for_session_idle) diff --git a/tests/server/routes/test_workflow_poller_routes.py b/tests/server/routes/test_workflow_poller_routes.py index f683b4057..66bf57eab 100644 --- a/tests/server/routes/test_workflow_poller_routes.py +++ b/tests/server/routes/test_workflow_poller_routes.py @@ -23,6 +23,13 @@ async def _fake_restart(workflow_id: str) -> dict[str, Any]: assert workflow_id == "wf-1" return {"workflowId": workflow_id, "state": "running", "lastStatus": None} + async def _fake_persist( + _workflow_id: str, + workflow_data: dict[str, Any], + _triggers: list[Any], + ) -> dict[str, Any]: + return workflow_data + monkeypatch.setattr( workflow_routes, "_read_workflow_from_fs", @@ -31,6 +38,7 @@ async def _fake_restart(workflow_id: str) -> dict[str, Any]: ), ) monkeypatch.setattr(workflow_routes.WorkflowStore, "put_config", _fake_put_config) + monkeypatch.setattr(workflow_routes, "_persist_workflow_triggers", _fake_persist) monkeypatch.setattr( "flocks.workflow.poller_manager.default_manager", SimpleNamespace(restart_workflow=_fake_restart), diff --git a/tests/server/test_auth_service_reassign.py b/tests/server/test_auth_service_reassign.py index 1bcb37307..159d94e9a 100644 --- a/tests/server/test_auth_service_reassign.py +++ b/tests/server/test_auth_service_reassign.py @@ -60,7 +60,8 @@ async def test_reassign_orphan_sessions_skips_owned_and_rewrites_orphans(monkeyp async def _list_all(): return listed - async def _update(*, project_id, session_id, owner_user_id, owner_username): + async def _update(*, project_id, session_id, owner_user_id, owner_username, allow_inactive=False): + assert allow_inactive is True update_calls.append({ "project_id": project_id, "session_id": session_id, @@ -119,7 +120,8 @@ async def test_reassign_orphan_sessions_continues_on_partial_failure(monkeypatch async def _list_all(): return listed - async def _update(*, project_id, session_id, owner_user_id, owner_username): + async def _update(*, project_id, session_id, owner_user_id, owner_username, allow_inactive=False): + assert allow_inactive is True update_calls.append(session_id) if session_id == "ses_b": raise RuntimeError("storage write failed") diff --git a/tests/server/test_input_dispatcher.py b/tests/server/test_input_dispatcher.py index 17a11d8f0..3a2c767cc 100644 --- a/tests/server/test_input_dispatcher.py +++ b/tests/server/test_input_dispatcher.py @@ -288,6 +288,63 @@ async def test_channel_unsafe_command_is_rejected(self): class TestSessionRoutesUseDispatcher: + @pytest.mark.asyncio + async def test_goal_mode_publishes_active_goal_before_llm(self, monkeypatch): + from flocks.input.events import UserInputEvent + from flocks.server.routes import session as session_routes + + order = [] + goal_state = SimpleNamespace( + status="active", + objective="fix tests", + last_reason=None, + ) + + async def publish(event_type, _properties): + if event_type == "session.goal.updated": + order.append("goal") + + async def process(*_args, **_kwargs): + order.append("llm") + + monkeypatch.setattr( + "flocks.command.direct.GoalManager.set_goal", + AsyncMock(return_value=goal_state), + ) + monkeypatch.setattr( + "flocks.command.direct.GoalManager.goal_prompt", + MagicMock(return_value="goal prompt"), + ) + monkeypatch.setattr( + "flocks.session.goal.GoalManager.get", + AsyncMock(return_value=goal_state), + ) + monkeypatch.setattr( + "flocks.server.routes.event.publish_event", + publish, + ) + monkeypatch.setattr( + session_routes, + "_process_session_message", + process, + ) + + await session_routes._dispatch_sse_input( + "ses_goal_mode", + SimpleNamespace(id="ses_goal_mode"), + UserInputEvent( + source_type="webui", + sessionID="ses_goal_mode", + text="/goal fix tests", + parts=[{"type": "text", "text": "fix tests"}], + display_text="fix tests", + executionMode="goal", + ), + "/tmp/project", + ) + + assert order == ["goal", "llm"] + @pytest.mark.asyncio async def test_prompt_async_routes_through_dispatcher(self, monkeypatch): from flocks.server.routes import session as session_routes @@ -389,7 +446,9 @@ def test_session_upload_path_rejects_traversal(self, monkeypatch, tmp_path): session_routes._session_uploads_dir("../outside") @pytest.mark.asyncio - async def test_prompt_async_queues_when_session_running_without_creating_message(self, monkeypatch): + async def test_prompt_async_queues_when_session_running_without_creating_message( + self, monkeypatch, tmp_path + ): from flocks.server.routes import session as session_routes from flocks.session.interaction_queue import InteractionQueue @@ -399,7 +458,13 @@ async def test_prompt_async_queues_when_session_running_without_creating_message message_create = AsyncMock() monkeypatch.setattr( "flocks.session.session.Session.get_by_id", - AsyncMock(return_value=SimpleNamespace(id=session_id, directory="/tmp/project")), + AsyncMock( + return_value=SimpleNamespace( + id=session_id, + directory=str(tmp_path), + status="active", + ) + ), ) monkeypatch.setattr("flocks.session.session_loop.SessionLoop.is_running", lambda _sid: True) monkeypatch.setattr("flocks.session.message.Message.create", message_create) @@ -417,7 +482,7 @@ async def test_prompt_async_queues_when_session_running_without_creating_message message_create.assert_not_called() @pytest.mark.asyncio - async def test_prompt_queue_rejects_when_full(self, monkeypatch): + async def test_prompt_queue_rejects_when_full(self, monkeypatch, tmp_path): from fastapi import HTTPException from flocks.server.routes import session as session_routes @@ -427,7 +492,13 @@ async def test_prompt_queue_rejects_when_full(self, monkeypatch): await InteractionQueue.clear(session_id) monkeypatch.setattr( "flocks.session.session.Session.get_by_id", - AsyncMock(return_value=SimpleNamespace(id=session_id, directory="/tmp/project")), + AsyncMock( + return_value=SimpleNamespace( + id=session_id, + directory=str(tmp_path), + status="active", + ) + ), ) monkeypatch.setattr("flocks.session.session_loop.SessionLoop.is_running", lambda _sid: True) monkeypatch.setattr(session_routes, "_publish_prompt_queue", AsyncMock()) @@ -445,7 +516,7 @@ async def test_prompt_queue_rejects_when_full(self, monkeypatch): assert exc_info.value.status_code == 409 @pytest.mark.asyncio - async def test_run_now_aborts_and_schedules_drain(self, monkeypatch): + async def test_run_now_aborts_and_schedules_drain(self, monkeypatch, tmp_path): from flocks.server.routes import session as session_routes from flocks.session.interaction_queue import InteractionQueue @@ -461,7 +532,13 @@ async def test_run_now_aborts_and_schedules_drain(self, monkeypatch): drain_mock = AsyncMock() monkeypatch.setattr( "flocks.session.session.Session.get_by_id", - AsyncMock(return_value=SimpleNamespace(id=session_id, directory="/tmp/project")), + AsyncMock( + return_value=SimpleNamespace( + id=session_id, + directory=str(tmp_path), + status="active", + ) + ), ) monkeypatch.setattr("flocks.session.session_loop.SessionLoop.is_running", lambda _sid: True) monkeypatch.setattr(session_routes, "abort_session", abort_mock) @@ -474,7 +551,7 @@ async def test_run_now_aborts_and_schedules_drain(self, monkeypatch): assert resp["status"] == "accepted" abort_mock.assert_awaited_once_with(session_id) wait_mock.assert_awaited_once_with(session_id) - drain_mock.assert_awaited_once_with(session_id, "/tmp/project") + drain_mock.assert_awaited_once_with(session_id, str(tmp_path.resolve())) @pytest.mark.asyncio async def test_scheduled_drain_retries_until_session_idle(self, monkeypatch): diff --git a/tests/server/test_server_port_config.py b/tests/server/test_server_port_config.py index 3964ff3d2..fb6b1f867 100644 --- a/tests/server/test_server_port_config.py +++ b/tests/server/test_server_port_config.py @@ -307,6 +307,22 @@ def fake_restart_all(config, _console): assert captured["config"].frontend_port == 5273 assert captured["config"].legacy_backend_port == 9100 + def test_restart_server_only_does_not_restart_daemon(self, monkeypatch): + """Test server-only restart leaves the supervisor daemon running.""" + calls = [] + + monkeypatch.setattr(cli_main, "restart_server", lambda _console: calls.append("server")) + monkeypatch.setattr( + cli_main, + "restart_all", + lambda _config, _console: calls.append("all"), + ) + + result = CliRunner().invoke(cli_main.app, ["restart", "--server-only"]) + + assert result.exit_code == 0 + assert calls == ["server"] + def test_restart_accepts_public_host_and_port(self, monkeypatch): """Test restart command accepts the unified public host/port options.""" captured = {} diff --git a/tests/session/test_auto_model_failover.py b/tests/session/test_auto_model_failover.py new file mode 100644 index 000000000..98541ec3c --- /dev/null +++ b/tests/session/test_auto_model_failover.py @@ -0,0 +1,1380 @@ +"""Focused tests for WebUI Auto runtime model failover.""" + +import time +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from flocks.session.message import Message, MessageRole +from flocks.session.runner import ( + LlmAttemptState, + SessionRunner, + StepFailure, + StepResult, +) +from flocks.session.session import Session, SessionInfo +from flocks.session.session_loop import ( + AutoFailoverCooldown, + LoopCallbacks, + LoopContext, + LoopResult, + RuntimeModel, + SessionLoop, +) + + +def _session(**updates) -> SessionInfo: + values = { + "id": "ses_auto", + "projectID": "project", + "directory": "/tmp/project", + "agent": "rex", + "provider": "primary", + "model": "primary-model", + "model_pinned": False, + "model_auto": True, + } + values.update(updates) + return SessionInfo.model_construct(**values) + + +def _ctx( + *, + auto: bool = True, + index: int = 0, + category: str = "user", +) -> LoopContext: + candidates = [ + RuntimeModel("primary", "primary-model"), + RuntimeModel("fallback", "fallback-model"), + ] + active = candidates[index] + return LoopContext( + session=_session( + provider=active.provider_id, + model=active.model_id, + category=category, + ), + provider_id=active.provider_id, + model_id=active.model_id, + agent_name="rex", + auto_failover=auto, + auto_failover_allowed=auto, + model_candidates=candidates if auto else [active], + candidate_index=index if auto else 0, + ) + + +def _failure( + *, + assistant_id: str, + reason: str = "server_error", + safe: bool = True, +) -> StepResult: + state = LlmAttemptState(observable_output_started=not safe) + message = "provider failed" + return StepResult( + action="stop", + error=message, + failure=StepFailure( + message=message, + error_data={"name": "APIError", "data": {"message": message}}, + assistant_message_id=assistant_id, + reason=reason, + allow_fallback=safe, + attempt_state=state, + attempts=1, + ), + ) + + +@pytest.fixture(autouse=True) +def _clear_cooldowns(): + SessionLoop._auto_failover_cooldowns.clear() + yield + SessionLoop._auto_failover_cooldowns.clear() + + +@pytest.mark.parametrize( + ("status_code", "message", "reason", "same_model_retries"), + [ + (401, "Unauthorized", "auth", 0), + (402, "Payment required", "billing", 0), + (429, "Too many requests", "rate_limit", 0), + (403, "Quota exceeded", "rate_limit", 0), + (403, "Insufficient quota", "billing", 0), + (408, "Request timeout", "timeout", 1), + (404, "Route not found", "unknown_api", 3), + (500, "Internal server error", "server_error", 3), + (502, "Bad gateway", "server_error", 3), + (529, "Provider overloaded", "overloaded", 1), + ], +) +def test_failover_classifier_retry_thresholds( + status_code: int, + message: str, + reason: str, + same_model_retries: int, +): + decision = SessionRunner.classify_failover_error({ + "name": "APIError", + "data": {"message": message, "statusCode": status_code}, + }) + + assert decision.eligible is True + assert decision.reason == reason + assert decision.same_model_retries == same_model_retries + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("status_code", "expected_calls"), + [ + (401, 1), + (402, 1), + (429, 1), + (408, 2), + (404, 4), + (500, 4), + (502, 4), + (529, 2), + ], +) +async def test_runner_applies_failover_retry_thresholds( + monkeypatch, + status_code: int, + expected_calls: int, +): + runner = SessionRunner( + session=_session(), + provider_id="primary", + model_id="primary-model", + defer_step_errors=True, + failover_available=True, + ) + last_user = SimpleNamespace(id="msg_user", agent="rex", role="user") + provider = MagicMock() + provider.is_configured.return_value = True + assistant = SimpleNamespace(id="msg_assistant") + failure = RuntimeError(f"Provider HTTP {status_code}") + failure.status_code = status_code + call_llm = AsyncMock(side_effect=failure) + + monkeypatch.setattr( + "flocks.session.runner.Agent.get", + AsyncMock(return_value=SimpleNamespace( + name="rex", + steps=None, + mode="primary", + prompt="", + tools=[], + )), + ) + monkeypatch.setattr("flocks.session.runner.Provider.get", lambda _provider_id: provider) + monkeypatch.setattr("flocks.session.runner.Provider.apply_config", AsyncMock()) + monkeypatch.setattr( + "flocks.session.runner.SessionPrompt.build_system_prompts", + AsyncMock(return_value=[]), + ) + monkeypatch.setattr(runner, "_build_callable_tool_schema", AsyncMock(return_value=[])) + monkeypatch.setattr( + runner, + "_to_chat_messages", + AsyncMock(return_value=[SimpleNamespace(role="user", content="hello")]), + ) + monkeypatch.setattr(Message, "get_text_content", AsyncMock(return_value="hello")) + monkeypatch.setattr(Message, "parts", AsyncMock(return_value=[])) + monkeypatch.setattr(Message, "create", AsyncMock(return_value=assistant)) + monkeypatch.setattr(Message, "update", AsyncMock()) + monkeypatch.setattr(runner, "_call_llm", call_llm) + monkeypatch.setattr("flocks.session.runner.SessionRetry.sleep", AsyncMock()) + + result = await runner._process_step([last_user], last_user) + + assert result.failure is not None + assert call_llm.await_count == expected_calls + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("status_code", "expected_calls"), + [ + (429, 1), + (503, 2), + ], +) +async def test_last_auto_candidate_keeps_hermes_retry_thresholds( + monkeypatch, + status_code: int, + expected_calls: int, +): + """The last candidate still has Auto's per-model retry budget.""" + runner = SessionRunner( + session=_session(), + provider_id="fallback", + model_id="fallback-model", + defer_step_errors=True, + failover_available=False, + ) + last_user = SimpleNamespace(id="msg_user", agent="rex", role="user") + provider = MagicMock() + provider.is_configured.return_value = True + assistant = SimpleNamespace(id="msg_assistant") + failure = RuntimeError(f"Provider HTTP {status_code}") + failure.status_code = status_code + call_llm = AsyncMock(side_effect=failure) + + monkeypatch.setattr( + "flocks.session.runner.Agent.get", + AsyncMock(return_value=SimpleNamespace( + name="rex", + steps=None, + mode="primary", + prompt="", + tools=[], + )), + ) + monkeypatch.setattr("flocks.session.runner.Provider.get", lambda _provider_id: provider) + monkeypatch.setattr("flocks.session.runner.Provider.apply_config", AsyncMock()) + monkeypatch.setattr( + "flocks.session.runner.SessionPrompt.build_system_prompts", + AsyncMock(return_value=[]), + ) + monkeypatch.setattr(runner, "_build_callable_tool_schema", AsyncMock(return_value=[])) + monkeypatch.setattr( + runner, + "_to_chat_messages", + AsyncMock(return_value=[SimpleNamespace(role="user", content="hello")]), + ) + monkeypatch.setattr(Message, "get_text_content", AsyncMock(return_value="hello")) + monkeypatch.setattr(Message, "parts", AsyncMock(return_value=[])) + monkeypatch.setattr(Message, "create", AsyncMock(return_value=assistant)) + monkeypatch.setattr(Message, "update", AsyncMock()) + monkeypatch.setattr(runner, "_call_llm", call_llm) + monkeypatch.setattr("flocks.session.runner.SessionRetry.sleep", AsyncMock()) + + result = await runner._process_step([last_user], last_user) + + assert result.failure is not None + assert call_llm.await_count == expected_calls + + +@pytest.mark.parametrize( + ("exception", "status_code", "reason"), + [ + ( + type("GoogleSdkError", (RuntimeError,), {"code": 429})( + "Resource exhausted" + ), + 429, + "rate_limit", + ), + ( + type( + "ResponseSdkError", + (RuntimeError,), + { + "response": SimpleNamespace( + status_code=503, + headers={"retry-after": "1"}, + ) + }, + )("Service unavailable"), + 503, + "overloaded", + ), + ], +) +def test_exception_status_is_normalized_from_sdk_shapes( + exception: Exception, + status_code: int, + reason: str, +): + runner = SessionRunner( + session=_session(), + provider_id="primary", + model_id="primary-model", + ) + + error = runner._exception_to_error_dict(exception) + + assert error["data"]["statusCode"] == status_code + assert SessionRunner.classify_failover_error(error).reason == reason + + +def test_exception_status_is_normalized_from_cause_chain(): + inner = type("GoogleSdkError", (RuntimeError,), {"code": 401})( + "Unauthenticated" + ) + outer = RuntimeError("Provider wrapper failed") + outer.__cause__ = inner + runner = SessionRunner( + session=_session(), + provider_id="primary", + model_id="primary-model", + ) + + error = runner._exception_to_error_dict(outer) + + assert error["data"]["statusCode"] == 401 + assert SessionRunner.classify_failover_error(error).reason == "auth" + + +def test_local_validation_error_never_fails_over(): + decision = SessionRunner.classify_failover_error({ + "name": "ValidationError", + "data": {"message": "Local prompt schema validation failed"}, + }) + + assert decision.eligible is False + assert decision.reason == "local_error" + + +def test_model_not_found_without_status_fails_over(): + decision = SessionRunner.classify_failover_error({ + "name": "ValueError", + "data": {"message": "Model acme-v2 not found for provider custom"}, + }) + + assert decision.eligible is True + assert decision.reason == "model_not_found" + assert decision.same_model_retries == 0 + + +def test_content_filter_error_fails_over_immediately(): + decision = SessionRunner.classify_failover_error({ + "name": "BadRequestError", + "data": {"message": "Response blocked by content_filter"}, + }) + + assert decision.eligible is True + assert decision.reason == "content_policy" + assert decision.same_model_retries == 0 + + +def test_candidate_switch_keeps_tool_loop_guard_only(): + ctx = _ctx() + tool_loop_guard = { + "last_user_id": "msg_user", + "signature": "same-tool-call", + "count": 2, + } + ctx.runner_static_cache.update({ + "tool_loop_guard": tool_loop_guard, + "tool_schema_cache": {"primary": "schema"}, + "chat_context_cache": {"primary": "context"}, + "system_prompt": "primary prompt", + }) + + SessionLoop._select_candidate(ctx, 1) + + assert ctx.runner_static_cache == {"tool_loop_guard": tool_loop_guard} + assert ctx.runner_static_cache["tool_loop_guard"] is tool_loop_guard + + +@pytest.mark.asyncio +async def test_reasoning_only_empty_response_is_not_replayed(monkeypatch): + runner = SessionRunner( + session=_session(), + provider_id="primary", + model_id="primary-model", + defer_step_errors=True, + failover_available=True, + ) + last_user = SimpleNamespace(id="msg_user", agent="rex", role="user") + provider = MagicMock() + provider.is_configured.return_value = True + assistant = SimpleNamespace(id="msg_assistant") + call_count = 0 + + async def call_llm(*_args, **_kwargs): + nonlocal call_count + call_count += 1 + runner._attempt_state.observable_output_started = True + return StepResult(action="stop", content="") + + monkeypatch.setattr( + "flocks.session.runner.Agent.get", + AsyncMock(return_value=SimpleNamespace( + name="rex", + steps=None, + mode="primary", + prompt="", + tools=[], + )), + ) + monkeypatch.setattr("flocks.session.runner.Provider.get", lambda _provider_id: provider) + monkeypatch.setattr("flocks.session.runner.Provider.apply_config", AsyncMock()) + monkeypatch.setattr( + "flocks.session.runner.SessionPrompt.build_system_prompts", + AsyncMock(return_value=[]), + ) + monkeypatch.setattr(runner, "_build_callable_tool_schema", AsyncMock(return_value=[])) + monkeypatch.setattr( + runner, + "_to_chat_messages", + AsyncMock(return_value=[SimpleNamespace(role="user", content="hello")]), + ) + monkeypatch.setattr(Message, "get_text_content", AsyncMock(return_value="hello")) + monkeypatch.setattr(Message, "parts", AsyncMock(return_value=[])) + monkeypatch.setattr(Message, "create", AsyncMock(return_value=assistant)) + monkeypatch.setattr(Message, "update", AsyncMock()) + monkeypatch.setattr(runner, "_call_llm", call_llm) + sleep = AsyncMock() + monkeypatch.setattr("flocks.session.runner.SessionRetry.sleep", sleep) + + result = await runner._process_step([last_user], last_user) + + assert call_count == 1 + assert result.failure is not None + assert result.failure.allow_fallback is False + sleep.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("chunk_kind", ["text", "reasoning", "tool"]) +async def test_real_stream_activity_prevents_retry_and_fallback( + monkeypatch, + chunk_kind: str, +): + """Exercise the real stream loop, then fail after an observable fragment.""" + + class FakeStreamProcessor: + def __init__(self, *args, tool_start_callback=None, **kwargs): + self.tool_start_callback = tool_start_callback + self.tool_calls = {} + self._text = [] + self._reasoning = [] + + async def process_event(self, event): + event_name = type(event).__name__ + if event_name == "TextDeltaEvent": + self._text.append(event.text) + elif event_name == "ReasoningDeltaEvent": + self._reasoning.append(event.text) + elif event_name == "ToolCallEvent": + if self.tool_start_callback: + await self.tool_start_callback(event.tool_name, event.input) + self.tool_calls[event.tool_call_id] = SimpleNamespace( + id=event.tool_call_id, + name=event.tool_name, + input=event.input, + ) + + def get_text_content(self): + return "".join(self._text) + + def get_reasoning_content(self): + return "".join(self._reasoning) + + if chunk_kind == "text": + chunk = SimpleNamespace( + delta="visible text", + reasoning=None, + event_type=None, + metadata={}, + tool_calls=None, + finish_reason=None, + usage=None, + ) + elif chunk_kind == "reasoning": + chunk = SimpleNamespace( + delta="visible reasoning", + reasoning=None, + event_type="reasoning", + metadata={}, + tool_calls=None, + finish_reason=None, + usage=None, + ) + else: + chunk = SimpleNamespace( + delta="", + reasoning=None, + event_type=None, + metadata={}, + tool_calls=[{ + "index": 0, + "id": "call_1", + "function": {"name": "example_tool", "arguments": "{}"}, + }], + finish_reason=None, + usage=None, + ) + + class FailingStreamProvider: + def __init__(self): + self.calls = 0 + + def is_configured(self): + return True + + def chat_stream(self, **_kwargs): + self.calls += 1 + + async def stream(): + yield chunk + failure = RuntimeError("Provider HTTP 500") + failure.status_code = 500 + raise failure + + return stream() + + provider = FailingStreamProvider() + runner = SessionRunner( + session=_session(), + provider_id="primary", + model_id="primary-model", + defer_step_errors=True, + failover_available=True, + ) + last_user = SimpleNamespace(id="msg_user", agent="rex", role="user") + assistant = SimpleNamespace(id="msg_assistant") + + monkeypatch.setattr( + "flocks.session.runner.Agent.get", + AsyncMock(return_value=SimpleNamespace( + name="rex", + steps=None, + mode="primary", + prompt="", + tools=[], + )), + ) + monkeypatch.setattr("flocks.session.runner.Provider.get", lambda _provider_id: provider) + monkeypatch.setattr("flocks.session.runner.Provider.apply_config", AsyncMock()) + monkeypatch.setattr( + "flocks.session.runner.SessionPrompt.build_system_prompts", + AsyncMock(return_value=[]), + ) + monkeypatch.setattr(runner, "_build_callable_tool_schema", AsyncMock(return_value=[])) + monkeypatch.setattr( + runner, + "_to_chat_messages", + AsyncMock(return_value=[SimpleNamespace(role="user", content="hello")]), + ) + monkeypatch.setattr(runner, "_should_use_text_tool_call_mode", lambda: False) + monkeypatch.setattr(Message, "get_text_content", AsyncMock(return_value="hello")) + monkeypatch.setattr(Message, "parts", AsyncMock(return_value=[])) + monkeypatch.setattr(Message, "create", AsyncMock(return_value=assistant)) + monkeypatch.setattr(Message, "update", AsyncMock()) + monkeypatch.setattr("flocks.session.runner.StreamProcessor", FakeStreamProcessor) + monkeypatch.setattr( + "flocks.session.runner.HookPipeline.has_stage_handlers", + AsyncMock(return_value=False), + ) + monkeypatch.setattr("flocks.session.runner.langfuse_is_active", lambda: False) + monkeypatch.setattr( + "flocks.provider.options.build_provider_options", + lambda _provider_id, _model_id: {}, + ) + sleep = AsyncMock() + monkeypatch.setattr("flocks.session.runner.SessionRetry.sleep", sleep) + + result = await runner._process_step([last_user], last_user) + + assert provider.calls == 1 + assert result.failure is not None + assert result.failure.allow_fallback is False + assert result.failure.attempt_state.received_chunk is True + assert result.failure.attempt_state.observable_output_started is True + assert result.failure.attempt_state.tool_execution_started is ( + chunk_kind == "tool" + ) + sleep.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_safe_failure_switches_candidate_and_removes_blank_message(monkeypatch): + ctx = _ctx() + last_user = SimpleNamespace(id="msg_user", agent="rex") + events = [] + + async def process_step(runner, _messages, _last_user): + if runner.provider_id == "primary": + return _failure(assistant_id="msg_failed") + return StepResult(action="stop", content="recovered") + + monkeypatch.setattr(SessionRunner, "_process_step", process_step) + delete = AsyncMock(return_value=True) + monkeypatch.setattr(Message, "delete", delete) + + async def publish(event, payload): + events.append((event, payload)) + + result = await SessionLoop._process_step_with_failover( + ctx, + LoopCallbacks(event_publish_callback=publish), + [last_user], + last_user, + ) + + assert result.content == "recovered" + assert (ctx.provider_id, ctx.model_id) == ("fallback", "fallback-model") + delete.assert_awaited_once_with("ses_auto", "msg_failed") + assert any(event == "message.removed" for event, _ in events) + assert any(event == "session.model.fallback" for event, _ in events) + + +@pytest.mark.asyncio +async def test_queued_user_is_detected_before_replacement_assistant(): + current_user = SimpleNamespace(id="msg_001", role=MessageRole.USER) + queued_user = SimpleNamespace(id="msg_002", role=MessageRole.USER) + replacement_assistant = SimpleNamespace( + id="msg_003", + role=MessageRole.ASSISTANT, + ) + + detected = await SessionLoop._detect_queued_user_message( + "ses_auto", + [current_user, queued_user, replacement_assistant], + current_user.id, + replacement_assistant, + ) + + assert detected is queued_user + + +@pytest.mark.asyncio +async def test_preflight_chain_exhaustion_persists_final_error(monkeypatch): + ctx = _ctx() + last_user = SimpleNamespace( + id="msg_user", + agent="rex", + role=MessageRole.USER, + ) + + async def preflight_failure(_runner, _messages, _last_user): + message = "Provider not configured" + return StepResult( + action="stop", + error=message, + failure=StepFailure( + message=message, + error_data={ + "name": "ProviderUnavailableError", + "data": {"message": message}, + }, + assistant_message_id=None, + reason="provider_unavailable", + allow_fallback=True, + attempt_state=LlmAttemptState(), + attempts=0, + ), + ) + + final_assistant = SimpleNamespace(id="msg_final_error") + create = AsyncMock(return_value=final_assistant) + monkeypatch.setattr(SessionRunner, "_process_step", preflight_failure) + monkeypatch.setattr(Message, "create", create) + + result = await SessionLoop._process_step_with_failover( + ctx, + LoopCallbacks(), + [last_user], + last_user, + ) + + assert result.error == "Provider not configured" + assert result.failure is not None + assert result.failure.assistant_message_id == final_assistant.id + create.assert_awaited_once_with( + session_id="ses_auto", + role=MessageRole.ASSISTANT, + content="", + agent="rex", + model_id="fallback-model", + provider_id="fallback", + parent_id=last_user.id, + error={ + "name": "ProviderUnavailableError", + "data": {"message": "Provider not configured"}, + }, + finish="error", + ) + + +@pytest.mark.asyncio +async def test_failed_blank_message_deletion_stops_switch(monkeypatch): + ctx = _ctx() + last_user = SimpleNamespace(id="msg_user", agent="rex") + monkeypatch.setattr( + SessionRunner, + "_process_step", + AsyncMock(return_value=_failure(assistant_id="msg_failed")), + ) + monkeypatch.setattr(Message, "delete", AsyncMock(return_value=False)) + update = AsyncMock() + monkeypatch.setattr(Message, "update", update) + + result = await SessionLoop._process_step_with_failover( + ctx, + LoopCallbacks(), + [last_user], + last_user, + ) + + assert result.error == "provider failed" + assert ctx.provider_id == "primary" + update.assert_awaited_once_with( + "ses_auto", + "msg_failed", + error={"name": "APIError", "data": {"message": "provider failed"}}, + finish="error", + ) + + +@pytest.mark.asyncio +async def test_fallbacks_are_attempted_in_candidate_order(monkeypatch): + ctx = _ctx() + ctx.model_candidates = [ + RuntimeModel("primary", "primary-model"), + RuntimeModel("fallback-1", "model-1"), + RuntimeModel("fallback-2", "model-2"), + ] + last_user = SimpleNamespace(id="msg_user", agent="rex") + attempts = [] + + async def process_step(runner, _messages, _last_user): + attempts.append((runner.provider_id, runner.model_id)) + if runner.provider_id != "fallback-2": + return _failure(assistant_id=f"msg_{runner.provider_id}") + return StepResult(action="stop", content="recovered") + + monkeypatch.setattr(SessionRunner, "_process_step", process_step) + delete = AsyncMock(return_value=True) + monkeypatch.setattr(Message, "delete", delete) + + result = await SessionLoop._process_step_with_failover( + ctx, + LoopCallbacks(), + [last_user], + last_user, + ) + + assert result.content == "recovered" + assert attempts == [ + ("primary", "primary-model"), + ("fallback-1", "model-1"), + ("fallback-2", "model-2"), + ] + assert delete.await_count == 2 + + +@pytest.mark.asyncio +async def test_chain_exhaustion_finalizes_only_last_candidate(monkeypatch): + ctx = _ctx() + ctx.model_candidates = [ + RuntimeModel("primary", "primary-model"), + RuntimeModel("fallback-1", "model-1"), + RuntimeModel("fallback-2", "model-2"), + ] + last_user = SimpleNamespace(id="msg_user", agent="rex") + + async def process_step(runner, _messages, _last_user): + return _failure(assistant_id=f"msg_{runner.provider_id}") + + monkeypatch.setattr(SessionRunner, "_process_step", process_step) + delete = AsyncMock(return_value=True) + update = AsyncMock() + monkeypatch.setattr(Message, "delete", delete) + monkeypatch.setattr(Message, "update", update) + + result = await SessionLoop._process_step_with_failover( + ctx, + LoopCallbacks(), + [last_user], + last_user, + ) + + assert result.error == "provider failed" + assert delete.await_count == 2 + update.assert_awaited_once() + assert update.await_args.args[1] == "msg_fallback-2" + cooldown = SessionLoop._auto_failover_cooldowns[ctx.session.id] + assert cooldown.model == RuntimeModel("fallback-2", "model-2") + assert cooldown.reason == "chain_exhausted" + + +@pytest.mark.asyncio +async def test_full_loop_reports_chain_exhaustion_once(monkeypatch): + ctx = _ctx() + ctx.session.memory_enabled = False + user = SimpleNamespace( + id="msg_user", + role=MessageRole.USER, + agent="rex", + model={"providerID": "primary", "modelID": "primary-model"}, + ) + # This test exercises exhaustion of an already fixed per-turn chain. + ctx.turn_user_id = user.id + final_assistant = SimpleNamespace( + id="msg_fallback", + role=MessageRole.ASSISTANT, + parentID=user.id, + finish="error", + ) + ctx.session_ctx = SimpleNamespace( + get_messages=AsyncMock(side_effect=[ + [user], + [user, final_assistant], + ]) + ) + attempts = [] + + async def process_step(runner, _messages, _last_user): + attempts.append((runner.provider_id, runner.model_id)) + return _failure(assistant_id=f"msg_{runner.provider_id}") + + monkeypatch.setattr(SessionRunner, "_process_step", process_step) + monkeypatch.setattr(Message, "parts", AsyncMock(return_value=[])) + monkeypatch.setattr(Message, "delete", AsyncMock(return_value=True)) + update = AsyncMock() + monkeypatch.setattr(Message, "update", update) + on_error = AsyncMock() + + result = await SessionLoop._run_loop( + ctx, + LoopCallbacks( + on_error=on_error, + event_publish_callback=AsyncMock(), + ), + ) + + assert result.action == "error" + assert result.error == "provider failed" + assert result.last_message is final_assistant + assert (result.provider_id, result.model_id) == ( + "fallback", + "fallback-model", + ) + assert attempts == [ + ("primary", "primary-model"), + ("fallback", "fallback-model"), + ] + on_error.assert_awaited_once_with("provider failed") + update.assert_awaited_once() + assert update.await_args.args[1] == "msg_fallback" + + +@pytest.mark.asyncio +async def test_observable_failure_is_finalized_without_replay(monkeypatch): + ctx = _ctx() + last_user = SimpleNamespace(id="msg_user", agent="rex") + monkeypatch.setattr( + SessionRunner, + "_process_step", + AsyncMock(return_value=_failure(assistant_id="msg_partial", safe=False)), + ) + delete = AsyncMock(return_value=True) + update = AsyncMock() + monkeypatch.setattr(Message, "delete", delete) + monkeypatch.setattr(Message, "update", update) + + result = await SessionLoop._process_step_with_failover( + ctx, + LoopCallbacks(), + [last_user], + last_user, + ) + + assert result.error == "provider failed" + assert ctx.provider_id == "primary" + delete.assert_not_awaited() + update.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_rate_limit_switch_sets_primary_cooldown(monkeypatch): + ctx = _ctx() + last_user = SimpleNamespace(id="msg_user", agent="rex") + + async def process_step(runner, _messages, _last_user): + if runner.provider_id == "primary": + return _failure(assistant_id="msg_rate", reason="rate_limit") + return StepResult(action="stop", content="recovered") + + monkeypatch.setattr(SessionRunner, "_process_step", process_step) + monkeypatch.setattr(Message, "delete", AsyncMock(return_value=True)) + + await SessionLoop._process_step_with_failover( + ctx, + LoopCallbacks(), + [last_user], + last_user, + ) + + cooldown = SessionLoop._auto_failover_cooldowns[ctx.session.id] + assert cooldown.model == RuntimeModel("fallback", "fallback-model") + assert cooldown.reason == "rate_limit" + assert SessionLoop._cooldown_candidate_index( + ctx.session.id, + ctx.model_candidates, + ) == 1 + + +@pytest.mark.asyncio +async def test_403_quota_failure_sets_primary_cooldown(monkeypatch): + decision = SessionRunner.classify_failover_error({ + "name": "APIError", + "data": {"message": "Quota exceeded", "statusCode": 403}, + }) + ctx = _ctx() + last_user = SimpleNamespace(id="msg_user", agent="rex") + + async def process_step(runner, _messages, _last_user): + if runner.provider_id == "primary": + return _failure( + assistant_id="msg_quota", + reason=decision.reason, + ) + return StepResult(action="stop", content="recovered") + + monkeypatch.setattr(SessionRunner, "_process_step", process_step) + monkeypatch.setattr(Message, "delete", AsyncMock(return_value=True)) + + await SessionLoop._process_step_with_failover( + ctx, + LoopCallbacks(), + [last_user], + last_user, + ) + + cooldown = SessionLoop._auto_failover_cooldowns[ctx.session.id] + assert cooldown.reason == "rate_limit" + assert cooldown.model == RuntimeModel("fallback", "fallback-model") + assert cooldown.expires_at > time.monotonic() + 50 + + +@pytest.mark.asyncio +async def test_chain_exhaustion_does_not_shorten_rate_limit_cooldown(monkeypatch): + ctx = _ctx() + last_user = SimpleNamespace(id="msg_user", agent="rex") + + async def process_step(runner, _messages, _last_user): + reason = "rate_limit" if runner.provider_id == "primary" else "server_error" + return _failure( + assistant_id=f"msg_{runner.provider_id}", + reason=reason, + ) + + monkeypatch.setattr(SessionRunner, "_process_step", process_step) + monkeypatch.setattr(Message, "delete", AsyncMock(return_value=True)) + monkeypatch.setattr(Message, "update", AsyncMock()) + + await SessionLoop._process_step_with_failover( + ctx, + LoopCallbacks(), + [last_user], + last_user, + ) + + cooldown = SessionLoop._auto_failover_cooldowns[ctx.session.id] + assert cooldown.reason == "rate_limit" + assert cooldown.model == RuntimeModel("fallback", "fallback-model") + # A 5s anti-replay window must not replace the primary's 60s cooldown. + assert cooldown.expires_at > time.monotonic() + 50 + + +@pytest.mark.asyncio +async def test_candidate_builder_discovers_one_model_per_tier_with_stable_seed( + monkeypatch, +): + config = SimpleNamespace() + model_manager = MagicMock() + model_manager.list_models.return_value = [ + SimpleNamespace(provider_id="primary", id="primary-model"), + SimpleNamespace(provider_id="primary", id="same-a"), + SimpleNamespace(provider_id="primary", id="same-b"), + SimpleNamespace(provider_id="other-a", id="other-a-model"), + SimpleNamespace(provider_id="other-b", id="other-b-model"), + SimpleNamespace(provider_id="missing", id="missing-model"), + ] + monkeypatch.setattr( + "flocks.config.config.Config.get", + AsyncMock(return_value=config), + ) + monkeypatch.setattr( + "flocks.provider.provider.Provider.apply_config", + AsyncMock(), + ) + monkeypatch.setattr( + "flocks.provider.model_manager.get_model_manager", + lambda: model_manager, + ) + + async def validate(provider_id, _model_id, **_kwargs): + available = provider_id != "missing" + return available, "available" if available else "provider_not_configured" + + monkeypatch.setattr(SessionLoop, "validate_runtime_model", validate) + primary = RuntimeModel("primary", "primary-model") + + first = await SessionLoop._build_model_candidates( + primary, + route_seed="ses_auto:msg_1", + ) + repeated = await SessionLoop._build_model_candidates( + primary, + route_seed="ses_auto:msg_1", + ) + + assert first == repeated + assert first[0] == primary + assert len(first) == 3 + assert first[1].provider_id == "primary" + assert first[1].model_id in {"same-a", "same-b"} + assert first[2].provider_id in {"other-a", "other-b"} + assert all(candidate.provider_id != "missing" for candidate in first) + + selections = { + tuple(await SessionLoop._build_model_candidates( + primary, + route_seed=f"ses_auto:msg_{index}", + )) + for index in range(12) + } + assert len(selections) > 1 + + +@pytest.mark.asyncio +async def test_candidate_builder_keeps_active_cooldown_model_in_its_tier( + monkeypatch, +): + config = SimpleNamespace() + model_manager = MagicMock() + model_manager.list_models.return_value = [ + SimpleNamespace(provider_id="primary", id="primary-model"), + SimpleNamespace(provider_id="primary", id="same-a"), + SimpleNamespace(provider_id="primary", id="same-b"), + SimpleNamespace(provider_id="other", id="other-a"), + SimpleNamespace(provider_id="other", id="other-b"), + ] + monkeypatch.setattr( + "flocks.config.config.Config.get", + AsyncMock(return_value=config), + ) + monkeypatch.setattr( + "flocks.provider.provider.Provider.apply_config", + AsyncMock(), + ) + monkeypatch.setattr( + "flocks.provider.model_manager.get_model_manager", + lambda: model_manager, + ) + monkeypatch.setattr( + SessionLoop, + "validate_runtime_model", + AsyncMock(return_value=(True, "available")), + ) + primary = RuntimeModel("primary", "primary-model") + cooldown_model = RuntimeModel("other", "other-b") + + candidates = await SessionLoop._build_model_candidates( + primary, + route_seed="ses_auto:new-turn", + preferred=cooldown_model, + ) + + assert candidates[0] == primary + assert candidates[1].provider_id == "primary" + assert candidates[2] == cooldown_model + + +@pytest.mark.asyncio +async def test_auto_configuration_only_requires_available_primary(monkeypatch): + monkeypatch.setattr( + "flocks.config.config.Config.resolve_default_llm", + AsyncMock(return_value={ + "provider_id": "primary", + "model_id": "primary-model", + }), + ) + monkeypatch.setattr( + SessionLoop, + "validate_runtime_model", + AsyncMock(return_value=(True, "available")), + ) + build_candidates = AsyncMock() + monkeypatch.setattr(SessionLoop, "_build_model_candidates", build_candidates) + + assert await SessionLoop.validate_auto_configuration() == (True, "available") + build_candidates.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_candidate_builder_allows_primary_only_chain(monkeypatch): + model_manager = MagicMock() + model_manager.list_models.return_value = [ + SimpleNamespace(provider_id="primary", id="primary-model"), + ] + monkeypatch.setattr( + "flocks.config.config.Config.get", + AsyncMock(return_value=SimpleNamespace()), + ) + monkeypatch.setattr( + "flocks.provider.provider.Provider.apply_config", + AsyncMock(), + ) + monkeypatch.setattr( + "flocks.provider.model_manager.get_model_manager", + lambda: model_manager, + ) + + primary = RuntimeModel("primary", "primary-model") + + assert await SessionLoop._build_model_candidates( + primary, + route_seed="ses_auto:msg_primary_only", + ) == [primary] + + +def test_cooldown_is_cleared_when_primary_changes(): + candidates = [ + RuntimeModel("new-primary", "new-model"), + RuntimeModel("fallback", "fallback-model"), + ] + SessionLoop._auto_failover_cooldowns["ses_auto"] = AutoFailoverCooldown( + model=RuntimeModel("fallback", "fallback-model"), + primary=RuntimeModel("old-primary", "old-model"), + expires_at=float("inf"), + reason="rate_limit", + ) + + assert SessionLoop._cooldown_candidate_index("ses_auto", candidates) == 0 + assert "ses_auto" not in SessionLoop._auto_failover_cooldowns + + +@pytest.mark.asyncio +async def test_synthetic_subtask_continuation_keeps_fallback(monkeypatch): + ctx = _ctx(index=1) + ctx.turn_user_id = "msg_real" + synthetic_user = SimpleNamespace( + id="msg_subtask_continue", + model={"providerID": "primary", "modelID": "primary-model"}, + ) + monkeypatch.setattr( + Message, + "parts", + AsyncMock(return_value=[SimpleNamespace(synthetic=True)]), + ) + + await SessionLoop._prepare_auto_turn(ctx, synthetic_user) + + assert ctx.auto_failover is True + assert ctx.turn_user_id == "msg_real" + assert (ctx.provider_id, ctx.model_id) == ("fallback", "fallback-model") + + +@pytest.mark.asyncio +async def test_first_real_turn_builds_stable_chain_from_user_id(monkeypatch): + ctx = _ctx() + ctx.turn_user_id = None + ctx.model_candidates = [RuntimeModel("primary", "primary-model")] + first_user = SimpleNamespace( + id="msg_first", + model={"providerID": "primary", "modelID": "primary-model"}, + ) + rebuilt = [ + RuntimeModel("primary", "primary-model"), + RuntimeModel("primary", "same-provider-model"), + RuntimeModel("other", "other-provider-model"), + ] + build = AsyncMock(return_value=rebuilt) + monkeypatch.setattr(SessionLoop, "_build_model_candidates", build) + + await SessionLoop._prepare_auto_turn(ctx, first_user) + + assert ctx.turn_user_id == "msg_first" + assert ctx.model_candidates == rebuilt + build.assert_awaited_once_with( + RuntimeModel("primary", "primary-model"), + route_seed="ses_auto:msg_first", + preferred=None, + ) + + +@pytest.mark.asyncio +async def test_queued_explicit_model_disables_auto(monkeypatch): + ctx = _ctx(index=1) + ctx.turn_user_id = "msg_real" + queued_user = SimpleNamespace( + id="msg_explicit", + model={"providerID": "explicit", "modelID": "explicit-model"}, + ) + persisted = _session( + provider="explicit", + model="explicit-model", + model_pinned=True, + model_auto=False, + ) + monkeypatch.setattr(Message, "parts", AsyncMock(return_value=[])) + monkeypatch.setattr( + "flocks.session.session.Session.get_by_id", + AsyncMock(return_value=persisted), + ) + + await SessionLoop._prepare_auto_turn(ctx, queued_user) + + assert ctx.auto_failover is False + assert ctx.model_candidates == [RuntimeModel("explicit", "explicit-model")] + assert (ctx.provider_id, ctx.model_id) == ("explicit", "explicit-model") + + +@pytest.mark.asyncio +async def test_non_webui_loop_cannot_activate_persisted_auto(monkeypatch): + ctx = _ctx(auto=False) + ctx.turn_user_id = "msg_real" + ctx.auto_failover_allowed = False + queued_user = SimpleNamespace( + id="msg_non_webui", + model={"providerID": "direct", "modelID": "direct-model"}, + ) + monkeypatch.setattr(Message, "parts", AsyncMock(return_value=[])) + monkeypatch.setattr( + "flocks.session.session.Session.get_by_id", + AsyncMock(return_value=_session(model_auto=True)), + ) + + await SessionLoop._prepare_auto_turn(ctx, queued_user) + + assert ctx.auto_failover is False + assert ctx.auto_failover_allowed is False + assert ctx.model_candidates == [RuntimeModel("direct", "direct-model")] + + +@pytest.mark.asyncio +async def test_queued_webui_turn_rebuilds_auto_chain(monkeypatch): + ctx = _ctx(auto=False) + ctx.turn_user_id = "msg_real" + ctx.auto_failover_allowed = True + queued_user = SimpleNamespace( + id="msg_auto", + model={"providerID": "primary", "modelID": "primary-model"}, + ) + rebuilt = [ + RuntimeModel("primary", "primary-model"), + RuntimeModel("new-fallback", "new-fallback-model"), + ] + monkeypatch.setattr(Message, "parts", AsyncMock(return_value=[])) + monkeypatch.setattr( + "flocks.session.session.Session.get_by_id", + AsyncMock(return_value=_session(model_auto=True)), + ) + monkeypatch.setattr( + "flocks.config.config.Config.resolve_default_llm", + AsyncMock(return_value={ + "provider_id": "primary", + "model_id": "primary-model", + }), + ) + build = AsyncMock(return_value=rebuilt) + monkeypatch.setattr(SessionLoop, "_build_model_candidates", build) + + await SessionLoop._prepare_auto_turn(ctx, queued_user) + + assert ctx.auto_failover is True + assert ctx.model_candidates == rebuilt + build.assert_awaited_once_with( + RuntimeModel("primary", "primary-model"), + route_seed="ses_auto:msg_auto", + preferred=None, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("category", ["user", "entity-config", "workflow"]) +async def test_queued_webui_auto_authorizes_active_loop(category): + ctx = _ctx(auto=False, category=category) + SessionLoop._active_loops[ctx.session.id] = ctx + try: + result = await SessionLoop.run(ctx.session.id, auto_failover=True) + finally: + SessionLoop._active_loops.pop(ctx.session.id, None) + + assert result.action == "queued" + assert ctx.auto_failover_allowed is True + + +@pytest.mark.asyncio +async def test_unsupported_session_loop_ignores_auto_authorization( + monkeypatch, +): + task_session = _session(category="task") + captured_ctx = None + + async def run_loop(ctx, _callbacks): + nonlocal captured_ctx + captured_ctx = ctx + return LoopResult(action="stop") + + build_candidates = AsyncMock() + monkeypatch.setattr( + "flocks.session.session.Session.get_by_id", + AsyncMock(return_value=task_session), + ) + monkeypatch.setattr(SessionLoop, "_build_model_candidates", build_candidates) + monkeypatch.setattr(SessionLoop, "_run_loop", run_loop) + monkeypatch.setattr(SessionLoop, "_publish_session_status", AsyncMock()) + monkeypatch.setattr(Message, "list", AsyncMock(return_value=[])) + monkeypatch.setattr( + "flocks.session.orphan_tools.abort_orphan_running_parts", + AsyncMock(), + ) + monkeypatch.setattr( + "flocks.session.session.Session.touch", + AsyncMock(), + ) + monkeypatch.setattr("flocks.bus.bus.Bus.publish", AsyncMock()) + + await SessionLoop.run( + task_session.id, + provider_id="primary", + model_id="primary-model", + auto_failover=True, + ) + + assert captured_ctx is not None + assert captured_ctx.auto_failover is False + assert captured_ctx.auto_failover_allowed is False + assert captured_ctx.model_candidates == [ + RuntimeModel("primary", "primary-model") + ] + build_candidates.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_active_unsupported_loop_rejects_auto_authorization(): + ctx = _ctx(auto=False, category="task") + SessionLoop._active_loops[ctx.session.id] = ctx + try: + result = await SessionLoop.run(ctx.session.id, auto_failover=True) + finally: + SessionLoop._active_loops.pop(ctx.session.id, None) + + assert result.action == "queued" + assert ctx.auto_failover_allowed is False + + +@pytest.mark.asyncio +async def test_session_delete_clears_auto_failover_cooldown(monkeypatch): + session = _session() + SessionLoop._auto_failover_cooldowns[session.id] = AutoFailoverCooldown( + model=RuntimeModel("fallback", "fallback-model"), + primary=RuntimeModel("primary", "primary-model"), + expires_at=float("inf"), + reason="rate_limit", + ) + monkeypatch.setattr(Session, "get", AsyncMock(return_value=session)) + monkeypatch.setattr(Session, "children", AsyncMock(return_value=[])) + monkeypatch.setattr(Session, "update", AsyncMock(return_value=session)) + monkeypatch.setattr(Message, "clear", AsyncMock(return_value=0)) + monkeypatch.setattr( + "flocks.session.callable_state.clear_session_callable_tools", + AsyncMock(), + ) + monkeypatch.setattr("flocks.bus.bus.Bus.publish", AsyncMock()) + + assert await Session.delete("project", session.id) is True + assert session.id not in SessionLoop._auto_failover_cooldowns diff --git a/tests/session/test_execution_mode.py b/tests/session/test_execution_mode.py new file mode 100644 index 000000000..005042d3c --- /dev/null +++ b/tests/session/test_execution_mode.py @@ -0,0 +1,482 @@ +from types import SimpleNamespace + +import pytest +from fastapi import HTTPException +from pydantic import ValidationError + +from flocks.server.routes.session import ( + PromptRequest, + _event_text_for_execution_mode, + _validate_execution_mode_request, +) +from flocks.session.execution_mode import ( + SessionExecutionMode, + execution_mode_prompt, + is_tool_allowed, + runtime_execution_mode, + tool_call_denial_reason, +) +from flocks.session.interaction_queue import InteractionQueue +from flocks.session.plan_file import is_current_plan_path, session_plan_file +from flocks.session.session import SessionInfo, SessionTime +from flocks.tool.registry import ( + Tool, + ToolCategory, + ToolContext, + ToolInfo, + ToolParameter, + ToolRegistry, + ToolResult, + ParameterType, +) +from flocks.tool.file.write import write_tool + + +def test_prompt_request_defaults_to_build_and_accepts_plan() -> None: + default_request = PromptRequest(parts=[{"type": "text", "text": "hello"}]) + plan_request = PromptRequest.model_validate({ + "parts": [{"type": "text", "text": "hello"}], + "executionMode": "plan", + }) + + assert default_request.execution_mode == SessionExecutionMode.BUILD + assert plan_request.execution_mode == SessionExecutionMode.PLAN + + +def test_prompt_request_rejects_removed_ask_mode() -> None: + with pytest.raises(ValidationError): + PromptRequest.model_validate({ + "parts": [{"type": "text", "text": "hello"}], + "executionMode": "ask", + }) + + +def test_goal_transport_uses_build_permissions_and_slash_dispatch() -> None: + parts = [{"type": "text", "text": " finish the feature "}] + + assert runtime_execution_mode("goal") == SessionExecutionMode.BUILD + assert _event_text_for_execution_mode( + parts, + SessionExecutionMode.GOAL, + ) == "/goal finish the feature" + + +def test_goal_requires_text_only_objective() -> None: + empty = PromptRequest.model_validate({ + "parts": [], + "executionMode": "goal", + }) + attachment = PromptRequest.model_validate({ + "parts": [ + {"type": "text", "text": "inspect this"}, + {"type": "file", "url": "file:///tmp/report.txt"}, + ], + "executionMode": "goal", + }) + + with pytest.raises(HTTPException, match="non-empty text objective"): + _validate_execution_mode_request(empty) + with pytest.raises(HTTPException, match="does not support attachments"): + _validate_execution_mode_request(attachment) + + +def test_plan_uses_read_only_permission_rules() -> None: + assert is_tool_allowed(SessionExecutionMode.PLAN, "read") + assert is_tool_allowed(SessionExecutionMode.PLAN, "grep") + assert is_tool_allowed(SessionExecutionMode.PLAN, "question") + assert is_tool_allowed(SessionExecutionMode.PLAN, "plan_exit") + assert is_tool_allowed(SessionExecutionMode.PLAN, "bash") + assert is_tool_allowed(SessionExecutionMode.PLAN, "edit") + assert is_tool_allowed(SessionExecutionMode.PLAN, "write") + assert is_tool_allowed(SessionExecutionMode.PLAN, "unknown_plugin_tool") + assert is_tool_allowed(SessionExecutionMode.PLAN, "task") + assert is_tool_allowed(SessionExecutionMode.PLAN, "delegate_task") + assert not is_tool_allowed(SessionExecutionMode.PLAN, "run_slash_command") + + assert is_tool_allowed(SessionExecutionMode.BUILD, "bash") + assert not is_tool_allowed(SessionExecutionMode.BUILD, "plan_exit") + assert "decision-complete implementation plan" in execution_mode_prompt("plan") + assert "material clarification question" in execution_mode_prompt("plan") + assert "call plan_exit" in execution_mode_prompt("plan") + assert "`explore` and `librarian`" in execution_mode_prompt("plan") + assert execution_mode_prompt("build") == "" + + +@pytest.mark.parametrize("tool_name", ["task", "delegate_task"]) +def test_plan_delegation_only_allows_explore_and_librarian(tool_name) -> None: + ctx = ToolContext(session_id="session-1", message_id="message-1") + + for subagent_type in ("explore", "librarian"): + assert tool_call_denial_reason( + SessionExecutionMode.PLAN, + tool_name, + {"subagent_type": subagent_type}, + ctx, + ) is None + + for arguments in ( + {"subagent_type": "general"}, + {"category": "quick"}, + {"session_id": "child-session"}, + {}, + ): + reason = tool_call_denial_reason( + SessionExecutionMode.PLAN, + tool_name, + arguments, + ctx, + ) + assert reason is not None + assert "explore, librarian" in reason + + +def test_plan_file_is_stable_and_session_scoped(tmp_path) -> None: + first = SessionInfo( + id="session-1", + slug="first-plan", + projectID="project-1", + directory=str(tmp_path), + time=SessionTime(created=1234, updated=1234), + ) + second = first.model_copy(update={ + "slug": "second-plan", + "time": SessionTime(created=5678, updated=5678), + }) + + first_plan = session_plan_file(first) + second_plan = session_plan_file(second) + + assert first_plan.path == tmp_path / ".flocks" / "plans" / "1234-first-plan.md" + assert first_plan.relative_path == ".flocks/plans/1234-first-plan.md" + assert first_plan.permission_path == ".flocks/plans/1234-first-plan.md" + assert session_plan_file(first) == first_plan + assert second_plan.path != first_plan.path + + +def test_plan_file_uses_worktree_root_and_directory_relative_tool_path( + tmp_path, +) -> None: + worktree = tmp_path / "repo" + directory = worktree / "packages" / "app" + directory.mkdir(parents=True) + session = SessionInfo( + slug="nested-plan", + projectID="project-1", + directory=str(directory), + time=SessionTime(created=1234, updated=1234), + ) + + plan = session_plan_file(session, worktree=str(worktree)) + + assert plan.path == worktree / ".flocks" / "plans" / "1234-nested-plan.md" + assert plan.relative_path == "../../.flocks/plans/1234-nested-plan.md" + assert plan.permission_path == ".flocks/plans/1234-nested-plan.md" + + +def test_plan_prompt_describes_create_then_incremental_edit(tmp_path) -> None: + session = SessionInfo( + slug="prompt-plan", + projectID="project-1", + directory=str(tmp_path), + time=SessionTime(created=1234, updated=1234), + ) + plan = session_plan_file(session) + + create_prompt = execution_mode_prompt("plan", session=session) + assert plan.relative_path in create_prompt + assert "No plan file exists yet" in create_prompt + assert "Bash is available only for read-only exploration" in create_prompt + + plan.path.parent.mkdir(parents=True) + plan.path.write_text("# Plan\n", encoding="utf-8") + + edit_prompt = execution_mode_prompt("plan", session=session) + assert plan.relative_path in edit_prompt + assert "already exists" in edit_prompt + assert "update it incrementally" in edit_prompt + + +def test_plan_path_guard_rejects_symlink_escape(tmp_path) -> None: + plan_relative = ".flocks/plans/1234-plan.md" + plan_path = tmp_path / plan_relative + external = tmp_path / "external.md" + external.write_text("outside\n", encoding="utf-8") + plan_path.parent.mkdir(parents=True) + plan_path.symlink_to(external) + ctx = ToolContext( + session_id="session-1", + message_id="message-1", + extra={ + "execution_mode": "plan", + "workspace_dir": str(tmp_path), + "plan_file_path": str(plan_path), + "plan_relative_path": plan_relative, + "plan_permission_path": plan_relative, + }, + ) + + assert not is_current_plan_path(ctx, plan_relative) + + +@pytest.mark.asyncio +async def test_prompt_queue_preserves_execution_mode() -> None: + session_id = "execution-mode-queue" + await InteractionQueue.clear(session_id) + + item = await InteractionQueue.enqueue( + session_id, + parts=[{"type": "text", "text": "plan this"}], + execution_mode=SessionExecutionMode.PLAN, + ) + + queued = await InteractionQueue.list(session_id) + assert item.executionMode == SessionExecutionMode.PLAN + assert queued[0].executionMode == SessionExecutionMode.PLAN + + await InteractionQueue.clear(session_id) + + +@pytest.mark.asyncio +async def test_registry_scopes_plan_delegation_before_handler(monkeypatch) -> None: + calls: list[dict] = [] + + async def handler(_ctx, **_kwargs): + calls.append(_kwargs) + return ToolResult(success=True, output="ok") + + tool = Tool( + info=ToolInfo( + name="task", + description="Delegation test tool", + category=ToolCategory.FILE, + ), + handler=handler, + ) + monkeypatch.setattr( + ToolRegistry, + "get", + classmethod(lambda _cls, _name: tool), + ) + + explore = await ToolRegistry.execute( + "task", + ctx=ToolContext( + session_id="session-1", + message_id="message-1", + extra={"execution_mode": "plan"}, + ), + subagent_type="explore", + ) + denied = await ToolRegistry.execute( + "delegate_task", + ctx=ToolContext( + session_id="session-1", + message_id="message-1", + extra={"execution_mode": "plan"}, + ), + subagent_type="general", + ) + + assert explore.success + assert not denied.success + assert "explore, librarian" in (denied.error or "") + assert calls == [{"subagent_type": "explore"}] + + +@pytest.mark.asyncio +async def test_registry_scopes_plan_edits_to_current_plan_file( + monkeypatch, + tmp_path, +) -> None: + calls: list[str] = [] + + async def handler(_ctx, **kwargs): + calls.append(kwargs["filePath"]) + return ToolResult(success=True, output="ok") + + tool = Tool( + info=ToolInfo( + name="write", + description="Write test", + category=ToolCategory.FILE, + parameters=[ + ToolParameter( + name="filePath", + type=ParameterType.STRING, + required=True, + ) + ], + ), + handler=handler, + ) + monkeypatch.setattr( + ToolRegistry, + "get", + classmethod(lambda _cls, _name: tool), + ) + plan_relative = ".flocks/plans/1234-plan.md" + ctx = ToolContext( + session_id="session-1", + message_id="message-1", + extra={ + "execution_mode": "plan", + "workspace_dir": str(tmp_path), + "plan_file_path": str(tmp_path / plan_relative), + "plan_relative_path": plan_relative, + "plan_permission_path": plan_relative, + }, + ) + + allowed = await ToolRegistry.execute( + "write", + ctx=ctx, + filePath=plan_relative, + ) + denied = await ToolRegistry.execute( + "write", + ctx=ctx, + filePath=".flocks/plans/other.md", + ) + traversal = await ToolRegistry.execute( + "write", + ctx=ctx, + filePath=".flocks/plans/../outside.md", + ) + + assert allowed.success + assert not denied.success + assert not traversal.success + assert calls == [plan_relative] + + +@pytest.mark.asyncio +async def test_tool_context_rechecks_plan_edit_permission(tmp_path) -> None: + plan_relative = ".flocks/plans/1234-plan.md" + ctx = ToolContext( + session_id="session-1", + message_id="message-1", + extra={ + "execution_mode": "plan", + "workspace_dir": str(tmp_path), + "plan_file_path": str(tmp_path / plan_relative), + "plan_relative_path": plan_relative, + "plan_permission_path": plan_relative, + }, + ) + + await ctx.ask(permission="edit", patterns=[plan_relative]) + with pytest.raises(PermissionError, match="current session plan file"): + await ctx.ask(permission="edit", patterns=["src/main.py"]) + + +@pytest.mark.asyncio +async def test_read_only_sandbox_allows_only_plan_artifact_write(tmp_path) -> None: + plan_relative = ".flocks/plans/1234-plan.md" + plan_path = tmp_path / plan_relative + ctx = ToolContext( + session_id="session-1", + message_id="message-1", + extra={ + "execution_mode": "plan", + "workspace_dir": str(tmp_path), + "plan_file_path": str(plan_path), + "plan_relative_path": plan_relative, + "plan_permission_path": plan_relative, + "sandbox": { + "workspace_dir": str(tmp_path), + "workspace_access": "ro", + }, + }, + ) + + allowed = await write_tool(ctx, "# Plan\n", plan_relative) + denied = await write_tool(ctx, "bad\n", "src/main.py") + + assert allowed.success + assert plan_path.read_text(encoding="utf-8") == "# Plan\n" + assert not denied.success + assert not (tmp_path / "src" / "main.py").exists() + + +@pytest.mark.asyncio +async def test_runner_filters_tools_with_message_mode(monkeypatch) -> None: + from flocks.session.runner import SessionRunner + + runner = object.__new__(SessionRunner) + runner.session = SimpleNamespace(id="session-1") + runner._step = 1 + runner.callbacks = SimpleNamespace(event_publish_callback=None) + agent = SimpleNamespace( + tools=[ + "read", + "bash", + "write", + "edit", + "task", + "delegate_task", + "run_slash_command", + ] + ) + + result = SimpleNamespace( + tool_infos=[ + SimpleNamespace(name="read"), + SimpleNamespace(name="bash"), + SimpleNamespace(name="write"), + SimpleNamespace(name="edit"), + SimpleNamespace(name="task"), + SimpleNamespace(name="delegate_task"), + SimpleNamespace(name="run_slash_command"), + ], + metadata={}, + ) + + async def list_tools(**_kwargs): + return result + + monkeypatch.setattr( + "flocks.session.runner.list_session_callable_tool_infos", + list_tools, + ) + monkeypatch.setattr( + ToolRegistry, + "get", + classmethod( + lambda _cls, name: ( + SimpleNamespace(info=SimpleNamespace(name="plan_exit", enabled=True)) + if name == "plan_exit" + else None + ) + ), + ) + messages = [ + SimpleNamespace( + role="user", + executionMode=SessionExecutionMode.PLAN, + ) + ] + + tools, metadata = await runner._list_callable_tool_infos_for_turn( + agent, + messages, + ) + + assert [tool.name for tool in tools] == [ + "read", + "bash", + "write", + "edit", + "task", + "delegate_task", + "plan_exit", + ] + assert metadata["executionMode"] == "plan" + assert metadata["modeAllowedToolNames"] == [ + "bash", + "delegate_task", + "edit", + "plan_exit", + "read", + "task", + "write", + ] diff --git a/tests/session/test_goal.py b/tests/session/test_goal.py index 7e845d7f3..9a4c189ad 100644 --- a/tests/session/test_goal.py +++ b/tests/session/test_goal.py @@ -291,7 +291,7 @@ async def test_goal_model_judge_uses_provider_options_without_main_token_budget( kwargs = provider.chat.await_args.kwargs assert kwargs["extra_body"] == {"reasoning_split": True} assert kwargs["max_tokens"] == JUDGE_MAX_TOKENS - assert kwargs["temperature"] == 0 + assert "temperature" not in kwargs assert decision.verdict == "complete" diff --git a/tests/session/test_lifecycle_hooks.py b/tests/session/test_lifecycle_hooks.py new file mode 100644 index 000000000..b08c3581e --- /dev/null +++ b/tests/session/test_lifecycle_hooks.py @@ -0,0 +1,699 @@ +"""Focused tests for the Python lifecycle-hook seams.""" + +from __future__ import annotations + +import asyncio +from types import SimpleNamespace +from unittest.mock import ANY, AsyncMock, MagicMock, patch + +import pytest + +from flocks.hooks.pipeline import HookContext, HookStage +from flocks.session.goal import GoalDecision +from flocks.session.runner import SessionRunner, StepResult +from flocks.session.session import SessionInfo +from flocks.session.session_loop import LoopCallbacks, LoopContext, SessionLoop + + +def _session(session_id: str = "ses_lifecycle_hooks") -> SessionInfo: + return SessionInfo.model_construct( + id=session_id, + slug="hooks", + project_id="project", + directory="/tmp/project", + title="Lifecycle Hooks", + agent="rex", + ) + + +def _loop_context(session_id: str = "ses_lifecycle_hooks") -> LoopContext: + return LoopContext( + session=_session(session_id), + provider_id="test-provider", + model_id="test-model", + agent_name="rex", + ) + + +@pytest.mark.asyncio +async def test_real_user_turn_is_detected_once_and_synthetic_is_ignored() -> None: + ctx = _loop_context() + first_user = SimpleNamespace(id="msg_user", model={}) + synthetic_user = SimpleNamespace(id="msg_synthetic", model={}) + + with ( + patch( + "flocks.session.session_loop.Message.parts", + AsyncMock( + side_effect=[ + [], + [SimpleNamespace(synthetic=True)], + ] + ), + ), + patch.object( + SessionLoop, + "_run_user_prompt_submit_hook", + AsyncMock(), + ) as submit_hook, + ): + for user in (first_user, first_user, synthetic_user): + if await SessionLoop._prepare_auto_turn(ctx, user): + await SessionLoop._run_user_prompt_submit_hook(ctx, user) + + assert ctx.turn_user_id == first_user.id + submit_hook.assert_awaited_once_with(ctx, first_user) + + +@pytest.mark.asyncio +async def test_user_prompt_submit_adds_ephemeral_turn_context() -> None: + ctx = _loop_context() + user = SimpleNamespace(id="msg_user", agent="rex") + run_hook = AsyncMock( + return_value=HookContext( + stage=HookStage.USER_PROMPT_SUBMIT, + input={}, + output={"additionalContext": " current sprint context "}, + ) + ) + + with ( + patch( + "flocks.session.session_loop.Message.get_text_content", + AsyncMock(return_value="implement hooks"), + ), + patch( + "flocks.hooks.pipeline.HookPipeline.run_user_prompt_submit", + run_hook, + ), + ): + await SessionLoop._run_user_prompt_submit_hook(ctx, user) + + assert ctx.turn_additional_context == "current sprint context" + payload = run_hook.await_args.args[0] + assert payload["messageID"] == user.id + assert payload["prompt"] == "implement hooks" + assert payload["model"] == { + "providerID": "test-provider", + "modelID": "test-model", + } + + +@pytest.mark.asyncio +async def test_session_start_runs_only_when_pending() -> None: + runner = SessionRunner( + session=_session("ses_session_start"), + provider_id="test-provider", + model_id="test-model", + session_start_pending=True, + ) + run_hook = AsyncMock() + + with patch( + "flocks.session.runner.HookPipeline.run_session_start", + run_hook, + ): + await runner._run_session_start_hook(SimpleNamespace(name="rex")) + await runner._run_session_start_hook(SimpleNamespace(name="rex")) + + run_hook.assert_awaited_once() + assert runner._session_start_fired is True + assert run_hook.await_args.args[0]["sessionID"] == "ses_session_start" + + +@pytest.mark.asyncio +async def test_turn_finish_block_creates_synthetic_continuation() -> None: + ctx = _loop_context("ses_turn_finish") + ctx.turn_user_id = "msg_user" + user = SimpleNamespace( + id="msg_user", + agent="rex", + model={"providerID": "test-provider", "modelID": "test-model"}, + ) + assistant = SimpleNamespace( + id="msg_assistant", + agent="rex", + finish="stop", + ) + continuation = SimpleNamespace(id="msg_continuation") + callbacks = LoopCallbacks(event_publish_callback=AsyncMock()) + create_message = AsyncMock(return_value=continuation) + run_hook = AsyncMock( + return_value=HookContext( + stage=HookStage.TURN_FINISH, + input={}, + output={ + "decision": "block", + "reason": "Run the test suite before finishing.", + }, + ) + ) + + with ( + patch( + "flocks.session.session_loop.Message.get", + AsyncMock(return_value=user), + ), + patch( + "flocks.session.session_loop.Message.get_text_content", + AsyncMock(side_effect=["implement hooks", "implementation complete"]), + ), + patch( + "flocks.session.session_loop.Message.create", + create_message, + ), + patch( + "flocks.hooks.pipeline.HookPipeline.run_turn_finish", + run_hook, + ), + patch( + "flocks.agent.registry.Agent.get", + AsyncMock(return_value=SimpleNamespace(steps=10)), + ), + ): + continued = await SessionLoop._run_turn_finish_hook( + ctx, + callbacks, + user, + assistant, + ) + + assert continued is True + assert ctx.stop_hook_active is True + assert create_message.await_args.kwargs["content"] == ("Run the test suite before finishing.") + assert create_message.await_args.kwargs["synthetic"] is True + assert create_message.await_args.kwargs["part_metadata"] == { + "turnFinishContinuation": True, + "stopHookActive": True, + "sourceAssistantMessageID": assistant.id, + } + hook_payload = run_hook.await_args.args[0] + assert hook_payload["finishReason"] == "stop" + assert hook_payload["stopHookActive"] is False + callbacks.event_publish_callback.assert_awaited_once() + assert callbacks.event_publish_callback.await_args.args[0] == "turn.continued" + + +@pytest.mark.asyncio +async def test_queued_prompt_arriving_during_turn_finish_wins() -> None: + ctx = _loop_context("ses_turn_finish_queue_race") + ctx.turn_user_id = "msg_001" + user = SimpleNamespace(id="msg_001", agent="rex", role="user") + assistant = SimpleNamespace( + id="msg_002", + agent="rex", + role="assistant", + finish="stop", + ) + queued_user = SimpleNamespace(id="msg_003", agent="rex", role="user") + ctx.session_ctx = SimpleNamespace( + get_messages=AsyncMock( + return_value=[user, assistant, queued_user], + ) + ) + callbacks = LoopCallbacks(event_publish_callback=AsyncMock()) + create_message = AsyncMock() + + with ( + patch( + "flocks.session.session_loop.Message.get", + AsyncMock(return_value=user), + ), + patch( + "flocks.session.session_loop.Message.get_text_content", + AsyncMock(side_effect=["prompt", "response"]), + ), + patch( + "flocks.session.session_loop.Message.create", + create_message, + ), + patch( + "flocks.hooks.pipeline.HookPipeline.run_turn_finish", + AsyncMock( + return_value=HookContext( + stage=HookStage.TURN_FINISH, + input={}, + output={"decision": "block", "reason": "continue"}, + ) + ), + ), + ): + continued = await SessionLoop._run_turn_finish_hook( + ctx, + callbacks, + user, + assistant, + ) + + assert continued is True + create_message.assert_not_awaited() + callbacks.event_publish_callback.assert_awaited_once() + event_name, payload = callbacks.event_publish_callback.await_args.args + assert event_name == "turn.continued" + assert payload["queuedUserMessageID"] == queued_user.id + + +@pytest.mark.asyncio +async def test_turn_finish_block_is_ignored_at_agent_step_limit() -> None: + ctx = _loop_context("ses_turn_finish_limit") + ctx.turn_user_id = "msg_user" + ctx.trace_step_offset = 2 + ctx.step = 1 + user = SimpleNamespace(id="msg_user", agent="rex") + assistant = SimpleNamespace( + id="msg_assistant", + agent="rex", + finish="stop", + ) + create_message = AsyncMock() + + with ( + patch( + "flocks.session.session_loop.Message.get", + AsyncMock(return_value=user), + ), + patch( + "flocks.session.session_loop.Message.get_text_content", + AsyncMock(side_effect=["prompt", "response"]), + ), + patch( + "flocks.session.session_loop.Message.create", + create_message, + ), + patch( + "flocks.hooks.pipeline.HookPipeline.run_turn_finish", + AsyncMock( + return_value=HookContext( + stage=HookStage.TURN_FINISH, + input={}, + output={"decision": "block", "reason": "continue"}, + ) + ), + ), + patch( + "flocks.agent.registry.Agent.get", + AsyncMock(return_value=SimpleNamespace(steps=3)), + ), + ): + continued = await SessionLoop._run_turn_finish_hook( + ctx, + LoopCallbacks(), + user, + assistant, + ) + + assert continued is False + create_message.assert_not_awaited() + + +def _message( + message_id: str, + role: str, + *, + finish: str | None = None, +) -> SimpleNamespace: + return SimpleNamespace( + id=message_id, + role=role, + finish=finish, + tokens=None, + summary=False, + agent="rex", + model={"providerID": "test-provider", "modelID": "test-model"}, + ) + + +@pytest.mark.asyncio +async def test_turn_finish_runs_only_after_persisted_stop() -> None: + ctx = _loop_context("ses_turn_finish_integration") + user = _message("msg_001", "user") + assistant = _message("msg_002", "assistant", finish="stop") + ctx.session_ctx = SimpleNamespace( + get_messages=AsyncMock( + side_effect=[ + [user], + [user, assistant], + ] + ) + ) + run_turn_finish = AsyncMock(return_value=False) + + with ( + patch( + "flocks.session.session_loop.Message.parts", + AsyncMock(return_value=[]), + ), + patch( + "flocks.session.session_loop.Message.get_text_content", + AsyncMock(return_value="final response"), + ), + patch( + "flocks.session.session_loop.Provider.resolve_model_info", + return_value=(0, 0, None), + ), + patch( + "flocks.session.session_loop.GoalManager.evaluate_after_turn", + AsyncMock( + return_value=GoalDecision( + status="inactive", + verdict="inactive", + ) + ), + ), + patch( + "flocks.session.session_loop.SessionLoop._run_user_prompt_submit_hook", + AsyncMock(), + ), + patch( + "flocks.session.session_loop.SessionLoop._run_turn_finish_hook", + run_turn_finish, + ), + patch( + "flocks.session.runner.SessionRunner._process_step", + AsyncMock(return_value=StepResult(action="stop")), + ), + patch( + "flocks.session.lifecycle.title.SessionTitle.ensure_title", + MagicMock(return_value=None), + ), + patch( + "flocks.session.session_loop.fire_and_forget", + MagicMock(), + ), + ): + result = await SessionLoop._run_loop(ctx, LoopCallbacks()) + + assert result.action == "stop" + run_turn_finish.assert_awaited_once_with( + ctx, + ANY, + user, + assistant, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("step_result", "assistant_finish"), + [ + (StepResult(action="stop", error="provider failed"), "error"), + (StepResult(action="continue"), "tool-calls"), + ], +) +async def test_turn_finish_skips_errors_and_tool_calls( + step_result: StepResult, + assistant_finish: str, +) -> None: + ctx = _loop_context(f"ses_turn_finish_{assistant_finish}") + user = _message("msg_001", "user") + assistant = _message("msg_002", "assistant", finish=assistant_finish) + ctx.session_ctx = SimpleNamespace( + get_messages=AsyncMock( + side_effect=[ + [user], + [user, assistant], + ] + ) + ) + run_turn_finish = AsyncMock(return_value=False) + + async def process_step(*_args, **_kwargs): + if step_result.action == "continue": + ctx.signal_abort() + return step_result + + with ( + patch( + "flocks.session.session_loop.Message.parts", + AsyncMock(return_value=[]), + ), + patch( + "flocks.session.session_loop.Provider.resolve_model_info", + return_value=(0, 0, None), + ), + patch( + "flocks.session.session_loop.SessionLoop._run_user_prompt_submit_hook", + AsyncMock(), + ), + patch( + "flocks.session.session_loop.SessionLoop._run_turn_finish_hook", + run_turn_finish, + ), + patch( + "flocks.session.runner.SessionRunner._process_step", + AsyncMock(side_effect=process_step), + ), + patch( + "flocks.session.lifecycle.title.SessionTitle.ensure_title", + MagicMock(return_value=None), + ), + patch( + "flocks.session.session_loop.fire_and_forget", + MagicMock(), + ), + ): + result = await SessionLoop._run_loop(ctx, LoopCallbacks()) + + assert result.action == "stop" + run_turn_finish.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_queued_user_message_takes_priority_over_turn_finish() -> None: + ctx = _loop_context("ses_turn_finish_queue") + user = _message("msg_001", "user") + assistant = _message("msg_002", "assistant", finish="stop") + queued_user = _message("msg_003", "user") + ctx.session_ctx = SimpleNamespace( + get_messages=AsyncMock( + side_effect=[ + [user], + [user, assistant, queued_user], + ] + ) + ) + run_turn_finish = AsyncMock(return_value=False) + + async def process_step(*_args, **_kwargs): + ctx.signal_abort() + return StepResult(action="stop") + + with ( + patch( + "flocks.session.session_loop.Message.parts", + AsyncMock(return_value=[]), + ), + patch( + "flocks.session.session_loop.Provider.resolve_model_info", + return_value=(0, 0, None), + ), + patch( + "flocks.session.session_loop.SessionLoop._run_user_prompt_submit_hook", + AsyncMock(), + ), + patch( + "flocks.session.session_loop.SessionLoop._run_turn_finish_hook", + run_turn_finish, + ), + patch( + "flocks.session.runner.SessionRunner._process_step", + AsyncMock(side_effect=process_step), + ), + patch( + "flocks.session.lifecycle.title.SessionTitle.ensure_title", + MagicMock(return_value=None), + ), + patch( + "flocks.session.session_loop.fire_and_forget", + MagicMock(), + ), + ): + await SessionLoop._run_loop(ctx, LoopCallbacks()) + + run_turn_finish.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_goal_continuation_takes_priority_over_turn_finish() -> None: + ctx = _loop_context("ses_turn_finish_goal") + user = _message("msg_001", "user") + assistant = _message("msg_002", "assistant", finish="stop") + goal_user = _message("msg_003", "user") + ctx.session_ctx = SimpleNamespace( + get_messages=AsyncMock( + side_effect=[ + [user], + [user, assistant], + ] + ) + ) + run_turn_finish = AsyncMock(return_value=False) + + async def process_step(*_args, **_kwargs): + ctx.signal_abort() + return StepResult(action="stop") + + with ( + patch( + "flocks.session.session_loop.Message.parts", + AsyncMock(return_value=[]), + ), + patch( + "flocks.session.session_loop.Message.get_text_content", + AsyncMock(return_value="not done"), + ), + patch( + "flocks.session.session_loop.Message.create", + AsyncMock(return_value=goal_user), + ), + patch( + "flocks.session.session_loop.Provider.resolve_model_info", + return_value=(0, 0, None), + ), + patch( + "flocks.session.session_loop.GoalManager.evaluate_after_turn", + AsyncMock( + return_value=GoalDecision( + status="active", + verdict="continue", + should_continue=True, + continuation_prompt="continue the goal", + ) + ), + ), + patch( + "flocks.session.session_loop.SessionLoop._run_user_prompt_submit_hook", + AsyncMock(), + ), + patch( + "flocks.session.session_loop.SessionLoop._run_turn_finish_hook", + run_turn_finish, + ), + patch( + "flocks.session.runner.SessionRunner._process_step", + AsyncMock(side_effect=process_step), + ), + patch( + "flocks.session.lifecycle.title.SessionTitle.ensure_title", + MagicMock(return_value=None), + ), + patch( + "flocks.session.session_loop.fire_and_forget", + MagicMock(), + ), + ): + await SessionLoop._run_loop(ctx, LoopCallbacks()) + + run_turn_finish.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_abort_does_not_trigger_turn_finish() -> None: + ctx = _loop_context("ses_turn_finish_abort") + user = _message("msg_001", "user") + ctx.session_ctx = SimpleNamespace(get_messages=AsyncMock(return_value=[user])) + run_turn_finish = AsyncMock(return_value=False) + + with ( + patch( + "flocks.session.session_loop.Message.parts", + AsyncMock(return_value=[]), + ), + patch( + "flocks.session.session_loop.Provider.resolve_model_info", + return_value=(0, 0, None), + ), + patch( + "flocks.session.session_loop.SessionLoop._run_user_prompt_submit_hook", + AsyncMock(), + ), + patch( + "flocks.session.session_loop.SessionLoop._run_turn_finish_hook", + run_turn_finish, + ), + patch( + "flocks.session.runner.SessionRunner._process_step", + AsyncMock(side_effect=asyncio.CancelledError()), + ), + patch( + "flocks.session.lifecycle.title.SessionTitle.ensure_title", + MagicMock(return_value=None), + ), + patch( + "flocks.session.session_loop.fire_and_forget", + MagicMock(), + ), + ): + await SessionLoop._run_loop(ctx, LoopCallbacks()) + + run_turn_finish.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_late_abort_after_step_completion_skips_turn_finish() -> None: + ctx = _loop_context("ses_turn_finish_late_abort") + user = _message("msg_001", "user") + assistant = _message("msg_002", "assistant", finish="stop") + ctx.session_ctx = SimpleNamespace( + get_messages=AsyncMock( + side_effect=[ + [user], + [user, assistant], + ] + ) + ) + run_turn_finish = AsyncMock(return_value=False) + + async def abort_after_step(_step: int) -> None: + ctx.abort_event.set() + + with ( + patch( + "flocks.session.session_loop.Message.parts", + AsyncMock(return_value=[]), + ), + patch( + "flocks.session.session_loop.Message.get_text_content", + AsyncMock(return_value="final response"), + ), + patch( + "flocks.session.session_loop.Provider.resolve_model_info", + return_value=(0, 0, None), + ), + patch( + "flocks.session.session_loop.GoalManager.evaluate_after_turn", + AsyncMock( + return_value=GoalDecision( + status="inactive", + verdict="inactive", + ) + ), + ), + patch( + "flocks.session.session_loop.SessionLoop._run_user_prompt_submit_hook", + AsyncMock(), + ), + patch( + "flocks.session.session_loop.SessionLoop._run_turn_finish_hook", + run_turn_finish, + ), + patch( + "flocks.session.runner.SessionRunner._process_step", + AsyncMock(return_value=StepResult(action="stop")), + ), + patch( + "flocks.session.lifecycle.title.SessionTitle.ensure_title", + MagicMock(return_value=None), + ), + patch( + "flocks.session.session_loop.fire_and_forget", + MagicMock(), + ), + ): + result = await SessionLoop._run_loop( + ctx, + LoopCallbacks(on_step_end=abort_after_step), + ) + + run_turn_finish.assert_not_awaited() + assert result.metadata["aborted"] is True diff --git a/tests/session/test_message_parts_persistence.py b/tests/session/test_message_parts_persistence.py index 4cb058253..61dafc1ff 100644 --- a/tests/session/test_message_parts_persistence.py +++ b/tests/session/test_message_parts_persistence.py @@ -1,6 +1,7 @@ """Persistence tests for message parts storage formats.""" import asyncio +from unittest.mock import AsyncMock import pytest @@ -180,6 +181,61 @@ async def test_delete_removes_parts_using_session_storage_format() -> None: assert await Storage.get(f"message_parts:{per_message_session_id}") is None +@pytest.mark.asyncio +async def test_delete_restores_caches_when_message_persistence_fails( + monkeypatch: pytest.MonkeyPatch, +) -> None: + session_id = "ses_parts_delete_message_failure" + await Message.create( + session_id, + MessageRole.USER, + "keep me", + id="msg_a", + part_id="part_a", + ) + persist_messages = AsyncMock( + side_effect=RuntimeError("message storage unavailable") + ) + monkeypatch.setattr(Message, "_persist_messages", persist_messages) + + with pytest.raises(RuntimeError, match="message storage unavailable"): + await Message.delete(session_id, "msg_a") + + restored = await Message.get_with_parts_lazy(session_id, "msg_a") + assert restored is not None + assert restored.info.id == "msg_a" + assert [part.text for part in restored.parts] == ["keep me"] + stored_messages = await Storage.get(f"message:{session_id}") + assert [message["id"] for message in stored_messages] == ["msg_a"] + + +@pytest.mark.asyncio +async def test_delete_commits_when_parts_cleanup_fails( + monkeypatch: pytest.MonkeyPatch, +) -> None: + session_id = "ses_parts_delete_parts_failure" + await Message.create( + session_id, + MessageRole.USER, + "delete me", + id="msg_a", + part_id="part_a", + ) + original_delete = Storage.delete + + async def fail_parts_delete(key: str) -> None: + if key == f"message_parts:{session_id}:msg_a": + raise RuntimeError("parts storage unavailable") + await original_delete(key) + + monkeypatch.setattr(Storage, "delete", fail_parts_delete) + + assert await Message.delete(session_id, "msg_a") is True + assert await Message.get(session_id, "msg_a") is None + assert await Storage.get(f"message:{session_id}") == [] + assert await Storage.get(f"message_parts:{session_id}:msg_a") is not None + + @pytest.mark.asyncio async def test_clear_removes_legacy_blob_and_per_message_keys() -> None: legacy_session_id = "ses_parts_clear_legacy" diff --git a/tests/session/test_runner_langfuse_payloads.py b/tests/session/test_runner_langfuse_payloads.py new file mode 100644 index 000000000..129ee8b3a --- /dev/null +++ b/tests/session/test_runner_langfuse_payloads.py @@ -0,0 +1,81 @@ +from flocks.provider.provider import ChatMessage +from flocks.session.runner import SessionRunner, ToolCall + + +def test_build_langfuse_request_payload_keeps_full_messages_and_system_prompt() -> None: + tools = [ + { + "type": "function", + "function": { + "name": "read_file", + "parameters": {"type": "object"}, + }, + } + ] + messages = [ + ChatMessage(role="system", content="system prompt"), + ChatMessage(role="user", content="user question"), + ChatMessage( + role="assistant", + content="calling tool", + tool_calls=[ + { + "id": "call_1", + "type": "function", + "function": { + "name": "read_file", + "arguments": "{\"path\":\"/tmp/demo.txt\"}", + }, + } + ], + ), + ] + + payload = SessionRunner._build_langfuse_request_payload( + step=3, + messages=messages, + request_tools=tools, + available_tools=tools, + provider_options={"temperature": 0.1, "max_tokens": 1024}, + ) + + assert payload["step"] == 3 + assert payload["messages"][0]["role"] == "system" + assert payload["messages"][0]["content"] == "system prompt" + assert payload["messages"][1]["content"] == "user question" + assert payload["messages"][2]["tool_calls"][0]["function"]["arguments"] == "{\"path\":\"/tmp/demo.txt\"}" + assert payload["request_tools"] == tools + assert payload["available_tools"] == tools + assert payload["provider_options"] == {"temperature": 0.1, "max_tokens": 1024} + + +def test_build_langfuse_response_payload_keeps_full_content_reasoning_and_tool_arguments() -> None: + full_content = "assistant output " * 80 + full_reasoning = "reasoning output " * 60 + tool_calls = [ + ToolCall( + id="call_1", + name="read_file", + arguments={"path": "/tmp/demo.txt", "offset": 1, "limit": 1000}, + ) + ] + + payload = SessionRunner._build_langfuse_response_payload( + action="continue", + content=full_content, + reasoning=full_reasoning, + finish_reason="tool_calls", + tool_calls=tool_calls, + ) + + assert payload["action"] == "continue" + assert payload["content"] == full_content + assert payload["reasoning"] == full_reasoning + assert payload["finish_reason"] == "tool_calls" + assert payload["tool_calls"] == [ + { + "id": "call_1", + "name": "read_file", + "arguments": {"path": "/tmp/demo.txt", "offset": 1, "limit": 1000}, + } + ] diff --git a/tests/session/test_runner_llm_hooks.py b/tests/session/test_runner_llm_hooks.py index de691d02f..7a3233e36 100644 --- a/tests/session/test_runner_llm_hooks.py +++ b/tests/session/test_runner_llm_hooks.py @@ -59,6 +59,9 @@ def get_reasoning_content(self) -> str: def get_finish_reason(self): return self.finish_reason + async def drain_parallel_tool_calls(self) -> None: + return None + class _FakeToolAccumulator: def __init__(self, processor): diff --git a/tests/session/test_session.py b/tests/session/test_session.py index 84892fecb..b470ec3b5 100644 --- a/tests/session/test_session.py +++ b/tests/session/test_session.py @@ -7,7 +7,12 @@ from unittest.mock import AsyncMock import pytest -from flocks.auth.context import AuthUser, reset_current_auth_user, set_current_auth_user +from flocks.auth.context import ( + API_TOKEN_SERVICE_USER_ID, + AuthUser, + reset_current_auth_user, + set_current_auth_user, +) from flocks.session.session import Session from flocks.session.message import Message, MessageInfo, MessageRole, TokenUsage from flocks.session.callable_state import add_session_callable_tools, get_session_callable_tools @@ -29,6 +34,8 @@ async def test_session_create(): assert session.directory == "/test/dir" assert session.title == "Test Session" assert session.status == "active" + assert session.owner_user_id == API_TOKEN_SERVICE_USER_ID + assert session.owner_username == API_TOKEN_SERVICE_USER_ID @pytest.mark.asyncio diff --git a/tests/session/test_session_advanced.py b/tests/session/test_session_advanced.py index d6dcf0bcd..f2004e487 100644 --- a/tests/session/test_session_advanced.py +++ b/tests/session/test_session_advanced.py @@ -11,15 +11,29 @@ - SessionInfo model fields """ +import asyncio +from types import SimpleNamespace +from unittest.mock import AsyncMock + import pytest +from flocks.auth.context import ( + API_TOKEN_SERVICE_USER_ID, + AuthUser, + reset_current_auth_user, + set_current_auth_user, +) +from flocks.session.message import Message, MessageRole, ToolPart, ToolStatePending +from flocks.session.features.todo import Todo, TodoInfo from flocks.session.session import ( PermissionRule, Session, + SessionInactiveError, SessionInfo, SessionRevert, SessionTime, ) +from flocks.storage.storage import Storage # --------------------------------------------------------------------------- @@ -41,7 +55,9 @@ async def test_archive_sets_status(self): result = await Session.archive("proj_arch_1", session.id) assert result is True raw = await Session.get("proj_arch_1", session.id) - assert raw is None or (raw and raw.status == "archived") + assert raw is not None + assert raw.status == "archived" + assert raw.time.archived is not None @pytest.mark.asyncio async def test_archive_nonexistent_returns_false(self): @@ -51,22 +67,540 @@ async def test_archive_nonexistent_returns_false(self): @pytest.mark.asyncio async def test_unarchive_restores_session(self): session = await _create(project_id="proj_arch_2") + await asyncio.sleep(0.01) await Session.archive("proj_arch_2", session.id) + archived = await Session.get("proj_arch_2", session.id) + assert archived is not None + assert archived.time.updated == session.time.updated result = await Session.unarchive("proj_arch_2", session.id) assert result is True + restored = await Session.get("proj_arch_2", session.id) + assert restored is not None + assert restored.status == "active" + assert restored.time.archived is None + assert restored.time.updated == session.time.updated @pytest.mark.asyncio - async def test_unarchive_nonarchived_returns_false(self): + async def test_unarchive_nonarchived_is_idempotent(self): session = await _create(project_id="proj_arch_3") - # Active session should not be unarchiveable result = await Session.unarchive("proj_arch_3", session.id) - assert result is False + assert result is True @pytest.mark.asyncio async def test_unarchive_nonexistent_returns_false(self): result = await Session.unarchive("proj_x", "ses_does_not_exist") assert result is False + @pytest.mark.asyncio + async def test_archive_and_unarchive_apply_to_descendants(self): + parent = await _create(project_id="proj_arch_tree", title="Parent") + child = await Session.fork("proj_arch_tree", parent.id) + + assert await Session.archive("proj_arch_tree", parent.id) is True + archived_parent = await Session.get("proj_arch_tree", parent.id) + archived_child = await Session.get("proj_arch_tree", child.id) + assert archived_parent is not None and archived_parent.status == "archived" + assert archived_child is not None and archived_child.status == "archived" + + assert await Session.unarchive("proj_arch_tree", parent.id) is True + restored_parent = await Session.get("proj_arch_tree", parent.id) + restored_child = await Session.get("proj_arch_tree", child.id) + assert restored_parent is not None and restored_parent.status == "active" + assert restored_child is not None and restored_child.status == "active" + + @pytest.mark.asyncio + async def test_archive_and_unarchive_preserve_prompt_queues(self): + from flocks.session.interaction_queue import InteractionQueue + + parent = await _create(project_id="proj_arch_queue", title="Parent") + child = await Session.create( + project_id=parent.project_id, + directory=parent.directory, + title="Child", + parent_id=parent.id, + ) + queued_text = { + parent.id: ["1", "11", "2", "22", "3"], + child.id: ["child-1", "child-2"], + } + + try: + before = {} + for queued_session_id, prompts in queued_text.items(): + for prompt in prompts: + await InteractionQueue.enqueue( + queued_session_id, + parts=[{"type": "text", "text": prompt}], + ) + before[queued_session_id] = [ + item.model_dump() + for item in await InteractionQueue.list(queued_session_id) + ] + + assert await Session.archive(parent.project_id, parent.id) is True + for queued_session_id in queued_text: + assert await InteractionQueue.pop_next(queued_session_id) is None + assert [ + item.model_dump() + for item in await InteractionQueue.list(queued_session_id) + ] == before[queued_session_id] + + assert await Session.unarchive(parent.project_id, parent.id) is True + for queued_session_id, prompts in queued_text.items(): + restored = await InteractionQueue.list(queued_session_id) + assert [item.model_dump() for item in restored] == before[queued_session_id] + assert [ + item.parts[0]["text"] + for item in restored + ] == prompts + first = await InteractionQueue.pop_next(queued_session_id) + assert first is not None + assert first.id == before[queued_session_id][0]["id"] + assert [ + item.id + for item in await InteractionQueue.list(queued_session_id) + ] == [ + item["id"] + for item in before[queued_session_id][1:] + ] + finally: + await InteractionQueue.clear(parent.id) + await InteractionQueue.clear(child.id) + + @pytest.mark.asyncio + async def test_archive_is_idempotent_and_preserves_first_timestamp(self): + session = await _create(project_id="proj_arch_idempotent") + + assert await Session.archive(session.project_id, session.id) is True + first = await Session.get(session.project_id, session.id) + assert first is not None and first.time.archived is not None + + assert await Session.archive(session.project_id, session.id) is True + second = await Session.get(session.project_id, session.id) + assert second is not None + assert second.time.archived == first.time.archived + + @pytest.mark.asyncio + async def test_unarchive_child_is_rejected_to_preserve_tree_state(self): + parent = await _create(project_id="proj_arch_child_restore", title="Parent") + child = await Session.create( + project_id=parent.project_id, + directory=parent.directory, + title="Child", + parent_id=parent.id, + ) + assert await Session.archive(parent.project_id, parent.id) is True + + assert await Session.unarchive(parent.project_id, child.id) is False + archived_parent = await Session.get(parent.project_id, parent.id) + archived_child = await Session.get(parent.project_id, child.id) + assert archived_parent is not None and archived_parent.status == "archived" + assert archived_child is not None and archived_child.status == "archived" + + @pytest.mark.asyncio + async def test_create_child_under_archived_parent_is_rejected(self): + parent = await _create(project_id="proj_arch_child_create", title="Parent") + assert await Session.archive(parent.project_id, parent.id) is True + + with pytest.raises(ValueError, match="is not active"): + await Session.create( + project_id=parent.project_id, + directory=parent.directory, + title="Late child", + parent_id=parent.id, + ) + + @pytest.mark.asyncio + async def test_child_inherits_parent_owner_when_explicit_values_are_none(self): + parent = await Session.create( + project_id="proj_arch_child_owner", + directory="/tmp", + title="Parent", + owner_user_id="owner-1", + owner_username="alice", + ) + + child = await Session.create( + project_id=parent.project_id, + directory=parent.directory, + title="Child", + parent_id=parent.id, + owner_user_id=None, + owner_username=None, + ) + + assert child.owner_user_id == parent.owner_user_id + assert child.owner_username == parent.owner_username + + @pytest.mark.asyncio + async def test_child_inherits_system_parent_owner_under_admin_context(self): + parent = await Session.create( + project_id="proj_system_parent_owner", + directory="/tmp", + title="System parent", + ) + admin = AuthUser(id="usr_admin", username="admin", role="admin") + token = set_current_auth_user(admin) + try: + child = await Session.create( + project_id=parent.project_id, + directory=parent.directory, + title="Child", + parent_id=parent.id, + ) + finally: + reset_current_auth_user(token) + + assert child.owner_user_id == API_TOKEN_SERVICE_USER_ID + assert child.owner_username == API_TOKEN_SERVICE_USER_ID + + @pytest.mark.asyncio + async def test_archived_session_loop_cannot_restart(self): + from flocks.session.session_loop import SessionLoop + + session = await _create(project_id="proj_arch_loop_guard") + assert await Session.archive(session.project_id, session.id) is True + + result = await SessionLoop.run(session.id) + + assert result.action == "error" + assert "archived" in (result.error or "") + + @pytest.mark.asyncio + async def test_lifecycle_guard_rejects_todo_write_after_archive(self): + session = await _create(project_id="proj_arch_todo_guard") + assert await Session.archive(session.project_id, session.id) is True + + with pytest.raises(SessionInactiveError): + await Todo.update_active( + session.id, + [TodoInfo(id="late", content="must not persist")], + ) + + assert await Todo.get(session.id) == [] + + @pytest.mark.asyncio + async def test_archive_wins_before_session_loop_registration(self, monkeypatch): + from flocks.session.session_loop import SessionLoop + + session = await _create(project_id="proj_arch_loop_race") + + async def archive_during_loop_start(_session_id: str): + assert await Session.archive(session.project_id, session.id) is True + return [] + + monkeypatch.setattr(Message, "list", archive_during_loop_start) + + result = await SessionLoop.run( + session.id, + provider_id="test-provider", + model_id="test-model", + ) + + assert result.action == "error" + assert "archived" in (result.error or "") + assert SessionLoop.is_running(session.id) is False + + @pytest.mark.asyncio + async def test_archive_cannot_be_overwritten_by_an_inflight_update(self, monkeypatch): + session = await _create(project_id="proj_arch_update_race") + storage_key = f"session:{session.project_id}:{session.id}" + update_reached_write = asyncio.Event() + release_update = asyncio.Event() + original_set = Storage.set + + async def delayed_set(key, value, *args, **kwargs): + if key == storage_key and getattr(value, "title", None) == "Racing update": + update_reached_write.set() + await release_update.wait() + return await original_set(key, value, *args, **kwargs) + + monkeypatch.setattr(Storage, "set", delayed_set) + + update_task = asyncio.create_task( + Session.update(session.project_id, session.id, title="Racing update") + ) + await update_reached_write.wait() + archive_task = asyncio.create_task(Session.archive(session.project_id, session.id)) + await asyncio.sleep(0) + assert archive_task.done() is False + release_update.set() + + assert await update_task is not None + assert await archive_task is True + stored = await Storage.get(storage_key, SessionInfo) + assert stored is not None + assert stored.status == "archived" + assert stored.title == "Racing update" + + @pytest.mark.asyncio + async def test_unrelated_session_updates_do_not_share_a_global_lock(self, monkeypatch): + first = await _create(project_id="proj_keyed_locks", title="First") + second = await _create(project_id="proj_keyed_locks", title="Second") + first_write_started = asyncio.Event() + release_first_write = asyncio.Event() + original_set = Storage.set + + async def delayed_set(key, value, *args, **kwargs): + if key.endswith(f":{first.id}") and getattr(value, "title", None) == "Slow": + first_write_started.set() + await release_first_write.wait() + return await original_set(key, value, *args, **kwargs) + + monkeypatch.setattr(Storage, "set", delayed_set) + first_update = asyncio.create_task( + Session.update(first.project_id, first.id, title="Slow") + ) + await first_write_started.wait() + + try: + second_update = await asyncio.wait_for( + Session.update(second.project_id, second.id, title="Fast"), + timeout=0.5, + ) + assert second_update is not None and second_update.title == "Fast" + finally: + release_first_write.set() + await first_update + + @pytest.mark.asyncio + async def test_archive_flushes_debounced_tool_parts_before_cache_invalidation(self): + session = await _create(project_id="proj_arch_parts_flush") + message = await Message.create( + session_id=session.id, + role=MessageRole.ASSISTANT, + content="working", + ) + tool_part = ToolPart( + sessionID=session.id, + messageID=message.id, + callID="call-pending", + tool="read", + state=ToolStatePending(input={"path": "/tmp/a"}, raw='{"path":"/tmp/a"}'), + ) + await Message.store_part(session.id, message.id, tool_part) + + assert session.id in Message._parts_flush_tasks + assert await Session.archive(session.project_id, session.id) is True + + reloaded_parts = await Message.parts(message.id, session.id) + reloaded_tool = next(part for part in reloaded_parts if part.id == tool_part.id) + assert isinstance(reloaded_tool, ToolPart) + assert reloaded_tool.state.status == "pending" + + @pytest.mark.asyncio + async def test_permanent_delete_clears_prompt_queue(self): + from flocks.session.interaction_queue import InteractionQueue + + session = await _create(project_id="proj_delete_queue") + await InteractionQueue.enqueue( + session.id, + parts=[{"type": "text", "text": "discard on delete"}], + ) + + assert await Session.delete(session.project_id, session.id) is True + assert await InteractionQueue.list(session.id) == [] + + @pytest.mark.asyncio + async def test_permanent_delete_removes_tree_data_in_one_mutation( + self, + tmp_path, + monkeypatch, + ): + from flocks.config.config import Config + from flocks.permission.next import PermissionNext, PermissionRequestInfo + from flocks.session.files import session_uploads_dir + + monkeypatch.setattr( + Config, + "get_data_path", + classmethod(lambda _cls: tmp_path), + ) + parent = await _create(project_id="proj_delete_atomic", title="Parent") + child = await Session.create( + project_id=parent.project_id, + directory=parent.directory, + title="Child", + parent_id=parent.id, + ) + for session in (parent, child): + upload_dir = session_uploads_dir(session.id) + upload_dir.mkdir(parents=True) + (upload_dir / "attachment.txt").write_text("remove", encoding="utf-8") + await Storage.set(f"message:{session.id}", [{"id": "message"}], "message") + await Storage.set(f"message_parts:{session.id}", {"legacy": []}, "message_parts") + await Storage.set(f"message_parts:{session.id}:message", [], "message_parts") + await Storage.set(f"todo:{session.id}", [{"content": "remove"}], "todo") + await Storage.set(f"goal:{session.id}", {"objective": "retain until delete"}, "goal") + await Storage.set(f"session_diff:{session.id}", {"files": []}, "session_diff") + await Storage.set(f"message_diff:{session.id}:message", {"diff": "remove"}, "message_diff") + await Storage.set(f"system_prompts:{session.id}:default", {"text": "remove"}, "system_prompt") + await Storage.set(f"session_callable_tools:{session.id}", {"tools": ["read"]}, "session_callable_tools") + await Storage.set( + f"{PermissionNext._SESSION_PREFIX}{session.id}", + {"bash": "allow"}, + "permission_session", + ) + PermissionNext._session_permissions[session.id] = {"bash": "allow"} + + pending = PermissionRequestInfo( + id="per_delete_tree", + sessionID=child.id, + permission="write", + patterns=["*"], + ) + pending_future = asyncio.get_running_loop().create_future() + PermissionNext._pending[pending.id] = { + "info": pending, + "future": pending_future, + } + await Storage.set( + f"{PermissionNext._PENDING_PREFIX}{pending.id}", + pending.model_dump(by_alias=True), + "permission_pending", + ) + await Storage.set( + f"{PermissionNext._REPLY_PREFIX}{pending.id}", + {"reply": "allow", "sessionID": child.id}, + "permission_reply", + ) + + assert await Session.delete(parent.project_id, parent.id) is True + + for session in (parent, child): + deleted_session = await Storage.get( + f"session:{session.project_id}:{session.id}", + SessionInfo, + ) + assert deleted_session is None + assert await Storage.get(f"message:{session.id}") is None + assert await Storage.get(f"message_parts:{session.id}") is None + assert await Storage.list_keys(prefix=f"message_parts:{session.id}:") == [] + assert await Storage.get(f"todo:{session.id}") is None + assert await Storage.get(f"goal:{session.id}") is None + assert await Storage.get(f"session_diff:{session.id}") is None + assert await Storage.list_keys(prefix=f"message_diff:{session.id}:") == [] + assert await Storage.list_keys(prefix=f"system_prompts:{session.id}:") == [] + assert await Storage.get(f"session_callable_tools:{session.id}") is None + assert await Storage.get(f"{PermissionNext._SESSION_PREFIX}{session.id}") is None + assert session.id not in PermissionNext._session_permissions + assert not session_uploads_dir(session.id).exists() + assert await Storage.get(f"{PermissionNext._PENDING_PREFIX}{pending.id}") is None + assert await Storage.get(f"{PermissionNext._REPLY_PREFIX}{pending.id}") is None + assert pending.id not in PermissionNext._pending + assert pending_future.cancelled() + + +class TestMoveToProject: + @pytest.mark.asyncio + async def test_move_updates_complete_tree_and_removes_source_keys(self): + parent = await _create(project_id="proj_move_source", title="Parent", directory="/tmp/source") + child = await Session.create( + project_id=parent.project_id, + directory=parent.directory, + title="Child", + parent_id=parent.id, + ) + + moved = await Session.move_to_project( + parent.project_id, + parent.id, + target_project_id="proj_move_target", + target_directory="/tmp/target", + ) + + assert moved is not None + assert moved.project_id == "proj_move_target" + assert moved.directory == "/tmp/target" + assert await Session.get("proj_move_source", parent.id) is None + assert await Session.get("proj_move_source", child.id) is None + moved_parent = await Session.get("proj_move_target", parent.id) + moved_child = await Session.get("proj_move_target", child.id) + assert moved_parent is not None and moved_parent.directory == "/tmp/target" + assert moved_child is not None and moved_child.directory == "/tmp/target" + assert moved_child.parent_id == moved_parent.id + + @pytest.mark.asyncio + async def test_move_rejects_a_child_as_the_tree_root(self): + parent = await _create(project_id="proj_move_child_source", title="Parent") + child = await Session.create( + project_id=parent.project_id, + directory=parent.directory, + title="Child", + parent_id=parent.id, + ) + + moved = await Session.move_to_project( + parent.project_id, + child.id, + target_project_id="proj_move_child_target", + target_directory="/tmp/target", + ) + + assert moved is None + assert await Session.get(parent.project_id, parent.id) is not None + assert await Session.get(parent.project_id, child.id) is not None + + @pytest.mark.asyncio + async def test_move_records_replay_boundary(self, monkeypatch): + parent = await _create(project_id="proj_move_history_source") + monkeypatch.setattr( + Message, + "list", + AsyncMock(return_value=[SimpleNamespace(id="msg_old")]), + ) + + moved = await Session.move_to_project( + parent.project_id, + parent.id, + target_project_id="proj_move_history_target", + target_directory="/tmp/target", + ) + + assert moved is not None + assert moved.revert is None + assert moved.metadata["projectMove"] == { + "sourceProjectID": "proj_move_history_source", + "targetProjectID": "proj_move_history_target", + "boundaryMessageID": "msg_old", + "movedAt": moved.time.updated, + } + + @pytest.mark.asyncio + async def test_move_rejects_active_revert(self): + parent = await _create(project_id="proj_move_revert_source") + await Session.update( + parent.project_id, + parent.id, + revert={"messageID": "msg_old", "snapshot": "tree_old"}, + ) + + moved = await Session.move_to_project( + parent.project_id, + parent.id, + target_project_id="proj_move_revert_target", + target_directory="/tmp/target", + ) + + assert moved is None + original = await Session.get(parent.project_id, parent.id) + assert original is not None and original.revert is not None + + @pytest.mark.asyncio + async def test_move_rejects_active_synchronous_operation(self): + parent = await _create(project_id="proj_move_active_source") + + async with Session.active_operation(parent.id): + moved = await Session.move_to_project( + parent.project_id, + parent.id, + target_project_id="proj_move_active_target", + target_directory="/tmp/target", + ) + + assert moved is None + assert await Session.get(parent.project_id, parent.id) is not None + # --------------------------------------------------------------------------- # Fork / Children @@ -81,6 +615,94 @@ async def test_fork_creates_child_session(self): assert child.parent_id == parent.id assert child.project_id == parent.project_id + @pytest.mark.asyncio + async def test_fork_preserves_project_move_replay_boundary(self): + from flocks.session.lifecycle.revert import SessionRevert as LifecycleSessionRevert + + parent = await _create(project_id="proj_fork_move_source") + await Message.create( + parent.id, + MessageRole.USER, + "First message before moving", + ) + boundary_message = await Message.create( + parent.id, + MessageRole.USER, + "Last message before moving", + ) + moved = await Session.move_to_project( + parent.project_id, + parent.id, + target_project_id="proj_fork_move_target", + target_directory="/tmp/fork-target", + ) + assert moved is not None + await Message.create(moved.id, MessageRole.USER, "After moving") + + child = await Session.fork(moved.project_id, moved.id) + child_messages = await Message.list(child.id, include_archived=True) + + assert child.metadata["projectMove"]["boundaryMessageID"] == child_messages[1].id + with pytest.raises(ValueError, match="移动项目前"): + await LifecycleSessionRevert.ensure_replayable(child, child_messages[0].id) + with pytest.raises(ValueError, match="移动项目前"): + await LifecycleSessionRevert.ensure_replayable(child, child_messages[1].id) + await LifecycleSessionRevert.ensure_replayable(child, child_messages[2].id) + + prefix_child = await Session.fork( + moved.project_id, + moved.id, + message_id=boundary_message.id, + ) + prefix_messages = await Message.list(prefix_child.id, include_archived=True) + assert len(prefix_messages) == 1 + assert ( + prefix_child.metadata["projectMove"]["boundaryMessageID"] + == prefix_messages[0].id + ) + with pytest.raises(ValueError, match="移动项目前"): + await LifecycleSessionRevert.ensure_replayable( + prefix_child, + prefix_messages[0].id, + ) + + @pytest.mark.asyncio + async def test_concurrent_move_waits_for_complete_fork(self, monkeypatch): + parent = await _create(project_id="proj_fork_race_source") + await Message.create(parent.id, MessageRole.USER, "Copy me") + copy_started = asyncio.Event() + allow_copy = asyncio.Event() + original_get_text_content = Message.get_text_content + + async def blocking_get_text_content(message): + copy_started.set() + await allow_copy.wait() + return await original_get_text_content(message) + + monkeypatch.setattr(Message, "get_text_content", blocking_get_text_content) + + fork_task = asyncio.create_task(Session.fork(parent.project_id, parent.id)) + await asyncio.wait_for(copy_started.wait(), timeout=1) + move_task = asyncio.create_task(Session.move_to_project( + parent.project_id, + parent.id, + target_project_id="proj_fork_race_target", + target_directory="/tmp/fork-race-target", + )) + await asyncio.sleep(0) + assert not move_task.done() + + allow_copy.set() + child = await asyncio.wait_for(fork_task, timeout=1) + moved = await asyncio.wait_for(move_task, timeout=1) + + assert moved is not None + assert await Session.get(parent.project_id, parent.id) is None + assert await Session.get(parent.project_id, child.id) is None + moved_child = await Session.get("proj_fork_race_target", child.id) + assert moved_child is not None + assert moved_child.parent_id == parent.id + @pytest.mark.asyncio async def test_fork_nonexistent_raises_or_returns_none(self): # fork() may raise ValueError or return None for nonexistent sessions @@ -134,6 +756,24 @@ async def test_set_revert_persisted(self): assert updated.revert is not None assert updated.revert.message_id == "msg_003" + @pytest.mark.asyncio + async def test_revert_rejects_part_from_a_different_message(self): + from flocks.session.lifecycle.revert import SessionRevert as LifecycleSessionRevert + + session = await _create(project_id="proj_revert_part_binding") + first_message = await Message.create(session.id, MessageRole.USER, "First") + second_message = await Message.create(session.id, MessageRole.USER, "Second") + first_parts = await Message.parts(first_message.id, session.id) + + with pytest.raises(ValueError, match="does not belong"): + await LifecycleSessionRevert.revert( + session.id, + second_message.id, + first_parts[0].id, + ) + + assert not Session.has_active_operations(session.id) + # --------------------------------------------------------------------------- # set_current / get_current diff --git a/tests/session/test_session_policy.py b/tests/session/test_session_policy.py index eaa698c5d..3281eace9 100644 --- a/tests/session/test_session_policy.py +++ b/tests/session/test_session_policy.py @@ -2,7 +2,7 @@ from __future__ import annotations -from flocks.auth.context import AuthUser +from flocks.auth.context import API_TOKEN_SERVICE_USER_ID, AuthUser from flocks.session.session import SessionInfo from flocks.session.policy import SessionPolicy @@ -57,6 +57,22 @@ def test_can_delete_requires_owner_only(): assert SessionPolicy.can_delete(session, stranger) is False +def test_service_owned_session_cannot_be_claimed_by_matching_username(): + session = _make_session( + owner_user_id=API_TOKEN_SERVICE_USER_ID, + owner_username=API_TOKEN_SERVICE_USER_ID, + ) + colliding_member = AuthUser( + id="usr_real_member", + username=API_TOKEN_SERVICE_USER_ID, + role="member", + ) + + assert SessionPolicy.is_owner(session, colliding_member) is False + assert SessionPolicy.can_read(session, colliding_member) is False + assert SessionPolicy.can_write(session, colliding_member) is False + + def test_can_read_requires_owner_for_private_session(): owner = _make_user() admin = _make_user(user_id="usr_admin", username="root", role="admin") @@ -94,3 +110,20 @@ def test_ownerless_session_admin_can_manage_but_member_cannot(): assert SessionPolicy.can_read(session, member) is False assert SessionPolicy.can_write(session, member) is False assert SessionPolicy.can_delete(session, member) is False + + +def test_system_owned_session_admin_can_manage_but_member_cannot(): + admin = _make_user(user_id="usr_admin", username="admin", role="admin") + member = _make_user(user_id="usr_member", username="member", role="member") + session = _make_session( + owner_user_id=API_TOKEN_SERVICE_USER_ID, + owner_username=API_TOKEN_SERVICE_USER_ID, + ) + + assert SessionPolicy.can_read(session, admin) is True + assert SessionPolicy.can_write(session, admin) is True + assert SessionPolicy.can_delete(session, admin) is True + + assert SessionPolicy.can_read(session, member) is False + assert SessionPolicy.can_write(session, member) is False + assert SessionPolicy.can_delete(session, member) is False diff --git a/tests/session/test_stream_processor.py b/tests/session/test_stream_processor.py index 186de7c3f..914f16d30 100644 --- a/tests/session/test_stream_processor.py +++ b/tests/session/test_stream_processor.py @@ -16,6 +16,7 @@ import pytest from unittest.mock import AsyncMock, MagicMock, patch +from flocks.hooks.pipeline import HookContext, HookStage from flocks.tool.registry import ToolResult from flocks.session.streaming.stream_processor import StreamProcessor, ToolCallState from flocks.session.streaming.stream_events import ( @@ -399,6 +400,333 @@ async def test_tool_call_state_created(self): # --------------------------------------------------------------------------- class TestToolCallExecution: + @pytest.mark.asyncio + async def test_tool_before_can_update_input_and_tool_after_can_replace_result(self): + proc = _make_processor() + execute = AsyncMock(return_value=ToolResult( + success=True, + output="original result", + )) + + async def update_input(payload): + payload["tool"]["input"] = {"command": "echo changed"} + return HookContext( + stage=HookStage.TOOL_BEFORE, + input=payload, + ) + + post_hook = AsyncMock(return_value=HookContext( + stage=HookStage.TOOL_AFTER, + input={}, + output={ + "result": { + "success": True, + "output": "hook result", + }, + }, + )) + + with ( + patch( + "flocks.session.streaming.stream_processor.Message.store_part", + new=AsyncMock(), + ), + patch( + "flocks.hooks.pipeline.HookPipeline.run_tool_before", + new=AsyncMock(side_effect=update_input), + ) as pre_hook, + patch( + "flocks.hooks.pipeline.HookPipeline.run_tool_after", + new=post_hook, + ), + patch( + "flocks.session.streaming.stream_processor.ToolRegistry.execute", + new=execute, + ), + ): + await proc.process_event( + ToolInputStartEvent(id="tc_changed", tool_name="bash") + ) + await proc.process_event(ToolCallEvent( + tool_call_id="tc_changed", + tool_name="bash", + input={"command": "echo original"}, + )) + + pre_hook.assert_awaited_once() + assert execute.await_args.kwargs["command"] == "echo changed" + post_hook.assert_awaited_once() + post_payload = post_hook.await_args.args[0] + assert post_payload["status"] == "completed" + assert post_payload["tool"]["input"] == {"command": "echo changed"} + assert post_payload["result"]["output"] == "original result" + assert post_payload["durationMs"] >= 0 + assert proc.tool_calls["tc_changed"].input == { + "command": "echo changed" + } + assert proc.tool_calls["tc_changed"].output == "hook result" + + @pytest.mark.asyncio + async def test_tool_error_still_emits_tool_after(self): + proc = _make_processor() + post_hook = AsyncMock(return_value=HookContext( + stage=HookStage.TOOL_AFTER, + input={}, + )) + + with ( + patch( + "flocks.session.streaming.stream_processor.Message.store_part", + new=AsyncMock(), + ), + patch( + "flocks.hooks.pipeline.HookPipeline.run_tool_before", + new=AsyncMock(return_value=HookContext( + stage=HookStage.TOOL_BEFORE, + input={ + "tool": { + "name": "bash", + "input": {"command": "false"}, + "callID": "tc_error", + } + }, + )), + ), + patch( + "flocks.hooks.pipeline.HookPipeline.run_tool_after", + new=post_hook, + ), + patch( + "flocks.session.streaming.stream_processor.ToolRegistry.execute", + new=AsyncMock(return_value=ToolResult( + success=False, + error="command failed", + )), + ), + ): + await proc.process_event( + ToolInputStartEvent(id="tc_error", tool_name="bash") + ) + await proc.process_event(ToolCallEvent( + tool_call_id="tc_error", + tool_name="bash", + input={"command": "false"}, + )) + + post_hook.assert_awaited_once() + post_payload = post_hook.await_args.args[0] + assert post_payload["status"] == "error" + assert post_payload["error"] == "command failed" + + @pytest.mark.asyncio + async def test_tool_before_can_block_and_tool_after_reports_blocked(self): + proc = _make_processor() + execute = AsyncMock( + return_value=ToolResult(success=True, output="should not run") + ) + post_hook = AsyncMock(return_value=HookContext( + stage=HookStage.TOOL_AFTER, + input={}, + )) + + with ( + patch( + "flocks.session.streaming.stream_processor.Message.store_part", + new=AsyncMock(), + ), + patch( + "flocks.hooks.pipeline.HookPipeline.run_tool_before", + new=AsyncMock(return_value=HookContext( + stage=HookStage.TOOL_BEFORE, + input={ + "tool": { + "name": "bash", + "input": {"command": "rm -rf target"}, + "callID": "tc_blocked", + } + }, + output={ + "decision": "block", + "reason": "Destructive command is not allowed", + }, + )), + ), + patch( + "flocks.hooks.pipeline.HookPipeline.run_tool_after", + new=post_hook, + ), + patch( + "flocks.session.streaming.stream_processor.ToolRegistry.execute", + new=execute, + ), + ): + await proc.process_event( + ToolInputStartEvent(id="tc_blocked", tool_name="bash") + ) + await proc.process_event(ToolCallEvent( + tool_call_id="tc_blocked", + tool_name="bash", + input={"command": "rm -rf target"}, + )) + + execute.assert_not_awaited() + assert proc.tool_calls["tc_blocked"].status == "error" + assert proc.tool_calls["tc_blocked"].error == ( + "Destructive command is not allowed" + ) + post_payload = post_hook.await_args.args[0] + assert post_payload["status"] == "blocked" + assert post_payload["error"] == "Destructive command is not allowed" + + @pytest.mark.asyncio + async def test_sandbox_block_still_emits_tool_after(self): + proc = _make_processor() + post_hook = AsyncMock(return_value=HookContext( + stage=HookStage.TOOL_AFTER, + input={}, + )) + + with ( + patch( + "flocks.session.streaming.stream_processor.Message.store_part", + new=AsyncMock(), + ), + patch( + "flocks.hooks.pipeline.HookPipeline.run_tool_before", + new=AsyncMock(return_value=HookContext( + stage=HookStage.TOOL_BEFORE, + input={ + "tool": { + "name": "bash", + "input": {"command": "pwd"}, + "callID": "tc_sandbox", + } + }, + )), + ), + patch( + "flocks.hooks.pipeline.HookPipeline.run_tool_after", + new=post_hook, + ), + patch.object( + proc, + "_resolve_sandbox_meta", + new=AsyncMock(return_value={ + "blocked": True, + "error": "Sandbox denied tool", + "extra": {}, + }), + ), + patch( + "flocks.session.streaming.stream_processor.ToolRegistry.execute", + new=AsyncMock(), + ) as execute, + ): + await proc.process_event( + ToolInputStartEvent(id="tc_sandbox", tool_name="bash") + ) + await proc.process_event(ToolCallEvent( + tool_call_id="tc_sandbox", + tool_name="bash", + input={"command": "pwd"}, + )) + + execute.assert_not_awaited() + assert post_hook.await_args.args[0]["status"] == "blocked" + + @pytest.mark.asyncio + async def test_cancelled_tool_emits_interrupted_tool_after(self): + proc = _make_processor() + post_hook = AsyncMock(return_value=HookContext( + stage=HookStage.TOOL_AFTER, + input={}, + )) + + with ( + patch( + "flocks.session.streaming.stream_processor.Message.store_part", + new=AsyncMock(), + ), + patch( + "flocks.hooks.pipeline.HookPipeline.run_tool_before", + new=AsyncMock(return_value=HookContext( + stage=HookStage.TOOL_BEFORE, + input={ + "tool": { + "name": "bash", + "input": {"command": "sleep 10"}, + "callID": "tc_interrupted", + } + }, + )), + ), + patch( + "flocks.hooks.pipeline.HookPipeline.run_tool_after", + new=post_hook, + ), + patch( + "flocks.session.streaming.stream_processor.ToolRegistry.execute", + new=AsyncMock(side_effect=asyncio.CancelledError()), + ), + ): + await proc.process_event( + ToolInputStartEvent(id="tc_interrupted", tool_name="bash") + ) + with pytest.raises(asyncio.CancelledError): + await proc.process_event(ToolCallEvent( + tool_call_id="tc_interrupted", + tool_name="bash", + input={"command": "sleep 10"}, + )) + + post_hook.assert_awaited_once() + assert post_hook.await_args.args[0]["status"] == "interrupted" + + @pytest.mark.asyncio + async def test_cancelled_tool_before_persists_interrupted_state(self): + proc = _make_processor() + store_part = AsyncMock() + post_hook = AsyncMock(return_value=HookContext( + stage=HookStage.TOOL_AFTER, + input={}, + )) + + with ( + patch( + "flocks.session.streaming.stream_processor.Message.store_part", + new=store_part, + ), + patch( + "flocks.hooks.pipeline.HookPipeline.run_tool_before", + new=AsyncMock(side_effect=asyncio.CancelledError()), + ), + patch( + "flocks.hooks.pipeline.HookPipeline.run_tool_after", + new=post_hook, + ), + patch( + "flocks.session.streaming.stream_processor.ToolRegistry.execute", + new=AsyncMock(), + ) as execute, + ): + await proc.process_event( + ToolInputStartEvent(id="tc_before_cancel", tool_name="bash") + ) + with pytest.raises(asyncio.CancelledError): + await proc.process_event(ToolCallEvent( + tool_call_id="tc_before_cancel", + tool_name="bash", + input={"command": "sleep 10"}, + )) + + execute.assert_not_awaited() + post_hook.assert_awaited_once() + assert post_hook.await_args.args[0]["status"] == "interrupted" + assert proc.tool_calls["tc_before_cancel"].status == "error" + final_part = store_part.await_args.args[2] + assert final_part.state.status == "error" + assert final_part.state.error == "Tool execution was interrupted" + @pytest.mark.asyncio async def test_tool_call_executes_tool(self): proc = _make_processor() @@ -659,6 +987,68 @@ async def test_tool_start_callback_called(self): callback.assert_called_once() + @pytest.mark.asyncio + async def test_tool_span_records_full_output(self): + proc = _make_processor() + proc._langfuse_generation = object() + long_output = "tool output " * 120 + span_ctx = MagicMock() + + successful_result = ToolResult( + success=True, + output=long_output, + title="bash", + metadata={}, + ) + + with ( + patch("flocks.session.streaming.stream_processor.Message.store_part", new=AsyncMock()), + patch("flocks.session.streaming.stream_processor.Message.update_part", new=AsyncMock()), + patch( + "flocks.session.streaming.stream_processor.ToolRegistry.execute", + new=AsyncMock(return_value=successful_result), + ), + patch("flocks.session.streaming.stream_processor.span_scope", return_value=span_ctx), + ): + await proc.process_event(ToolInputStartEvent(id="tc_span_output", tool_name="bash")) + await proc.process_event( + ToolCallEvent(tool_call_id="tc_span_output", tool_name="bash", input={"command": "ls -la"}) + ) + + assert span_ctx.end.call_count == 1 + assert span_ctx.end.call_args.kwargs["output"] == long_output + + @pytest.mark.asyncio + async def test_tool_span_records_full_error(self): + proc = _make_processor() + proc._langfuse_generation = object() + long_error = "tool error " * 120 + span_ctx = MagicMock() + + failed_result = ToolResult( + success=False, + error=long_error, + metadata={}, + ) + + with ( + patch("flocks.session.streaming.stream_processor.Message.store_part", new=AsyncMock()), + patch("flocks.session.streaming.stream_processor.Message.update_part", new=AsyncMock()), + patch( + "flocks.session.streaming.stream_processor.ToolRegistry.execute", + new=AsyncMock(return_value=failed_result), + ), + patch("flocks.session.streaming.stream_processor.span_scope", return_value=span_ctx), + ): + await proc.process_event(ToolInputStartEvent(id="tc_span_error", tool_name="bash")) + await proc.process_event( + ToolCallEvent(tool_call_id="tc_span_error", tool_name="bash", input={"command": "false"}) + ) + + assert span_ctx.end.call_count == 1 + assert span_ctx.end.call_args.kwargs["output"] == long_error + assert span_ctx.end.call_args.kwargs["level"] == "ERROR" + @pytest.mark.asyncio async def test_tool_error_falls_back_to_metadata_output(self): event_callback = AsyncMock() diff --git a/tests/session/test_stream_timeouts.py b/tests/session/test_stream_timeouts.py new file mode 100644 index 000000000..9aedb7b09 --- /dev/null +++ b/tests/session/test_stream_timeouts.py @@ -0,0 +1,111 @@ +"""Tests for adaptive LLM stream timeout resolution.""" + +from types import SimpleNamespace + +from flocks.session.streaming.timeouts import resolve_llm_stream_timeouts + + +def _provider( + *, + provider_id: str = "openai", + base_url: str = "https://api.example.com/v1", + provider_settings: dict | None = None, + model_settings: dict | None = None, +): + model = SimpleNamespace( + id="test-model", + custom_settings=model_settings or {}, + ) + return SimpleNamespace( + id=provider_id, + _base_url=base_url, + _config=SimpleNamespace( + base_url=base_url, + custom_settings=provider_settings or {}, + ), + get_models=lambda: [model], + ) + + +def test_cloud_provider_uses_safe_defaults(monkeypatch): + monkeypatch.delenv("FLOCKS_LLM_STREAM_FIRST_CHUNK_TIMEOUT_S", raising=False) + monkeypatch.delenv("FLOCKS_LLM_STREAM_ONGOING_CHUNK_TIMEOUT_S", raising=False) + + timeouts = resolve_llm_stream_timeouts(_provider(), "test-model") + + assert timeouts.first_chunk_s == 120.0 + assert timeouts.ongoing_chunk_s == 300.0 + assert timeouts.is_local is False + + +def test_local_provider_gets_long_prefill_budget(monkeypatch): + monkeypatch.delenv("FLOCKS_LLM_STREAM_FIRST_CHUNK_TIMEOUT_S", raising=False) + + timeouts = resolve_llm_stream_timeouts( + _provider(base_url="http://127.0.0.1:11434/v1"), + "test-model", + ) + + assert timeouts.first_chunk_s == 1800.0 + assert timeouts.is_local is True + + +def test_private_network_endpoint_is_treated_as_local(monkeypatch): + monkeypatch.delenv("FLOCKS_LLM_STREAM_FIRST_CHUNK_TIMEOUT_S", raising=False) + + timeouts = resolve_llm_stream_timeouts( + _provider(base_url="http://192.168.1.20:8000/v1"), + "test-model", + ) + + assert timeouts.first_chunk_s == 1800.0 + assert timeouts.is_local is True + + +def test_environment_overrides_default(monkeypatch): + monkeypatch.setenv("FLOCKS_LLM_STREAM_FIRST_CHUNK_TIMEOUT_S", "240") + monkeypatch.setenv("FLOCKS_LLM_STREAM_ONGOING_CHUNK_TIMEOUT_S", "420") + + timeouts = resolve_llm_stream_timeouts(_provider(), "test-model") + + assert timeouts.first_chunk_s == 240.0 + assert timeouts.ongoing_chunk_s == 420.0 + + +def test_provider_config_overrides_environment(monkeypatch): + monkeypatch.setenv("FLOCKS_LLM_STREAM_FIRST_CHUNK_TIMEOUT_S", "240") + provider = _provider( + provider_settings={"stream_first_chunk_timeout_s": 360}, + ) + + timeouts = resolve_llm_stream_timeouts(provider, "test-model") + + assert timeouts.first_chunk_s == 360.0 + + +def test_model_config_overrides_provider_config(monkeypatch): + monkeypatch.setenv("FLOCKS_LLM_STREAM_FIRST_CHUNK_TIMEOUT_S", "240") + provider = _provider( + provider_settings={"stream_first_chunk_timeout_s": 360}, + model_settings={ + "stream_first_chunk_timeout_s": 480, + "stream_ongoing_chunk_timeout_s": 600, + }, + ) + + timeouts = resolve_llm_stream_timeouts(provider, "test-model") + + assert timeouts.first_chunk_s == 480.0 + assert timeouts.ongoing_chunk_s == 600.0 + + +def test_invalid_overrides_fall_back_to_defaults(monkeypatch): + monkeypatch.setenv("FLOCKS_LLM_STREAM_FIRST_CHUNK_TIMEOUT_S", "invalid") + provider = _provider( + provider_settings={"stream_first_chunk_timeout_s": -1}, + model_settings={"stream_first_chunk_timeout_s": 0}, + ) + + timeouts = resolve_llm_stream_timeouts(provider, "test-model") + + assert timeouts.first_chunk_s == 120.0 diff --git a/tests/session/test_tool_accumulator.py b/tests/session/test_tool_accumulator.py index e470253bc..0ef761f6f 100644 --- a/tests/session/test_tool_accumulator.py +++ b/tests/session/test_tool_accumulator.py @@ -108,18 +108,20 @@ async def test_index_to_id_mapping_without_explicit_id(self): class TestFeedChunkIncremental: @pytest.mark.asyncio - async def test_partial_json_does_not_fire(self): + async def test_partial_json_emits_start_without_executing_tool(self): acc, proc = _make_accumulator() with patch("flocks.session.streaming.tool_accumulator.ToolRegistry") as mock_reg: mock_reg.get_schema.return_value = None - # Incomplete JSON - no closing brace await acc.feed_chunk(_make_chunk( - tc_id="call_inc", name="my_tool", arguments='{"command": "ls' + tc_id="call_inc", name="write", arguments='{"filePath": "/tmp/f", "content": "long' )) - # Incomplete JSON should not fire - assert proc.process_event.call_count == 0 + events = [call.args[0] for call in proc.process_event.call_args_list] + assert len(events) == 1 + assert isinstance(events[0], ToolInputStartEvent) + assert events[0].id == "call_inc" + assert events[0].tool_name == "write" @pytest.mark.asyncio async def test_incremental_chunks_complete_and_fire(self): @@ -128,12 +130,46 @@ async def test_incremental_chunks_complete_and_fire(self): mock_reg.get_schema.return_value = None await acc.feed_chunk(_make_chunk(tc_id="call_inc2", name="tool_x", arguments='{"key":')) - assert proc.process_event.call_count == 0 + assert proc.process_event.call_count == 1 + assert isinstance(proc.process_event.call_args_list[0].args[0], ToolInputStartEvent) await acc.feed_chunk(_make_chunk(tc_id="call_inc2", arguments='"value"}')) - event_types = [c.args[0].type for c in proc.process_event.call_args_list] - assert "tool-call" in event_types + events = [call.args[0] for call in proc.process_event.call_args_list] + assert [event.type for event in events] == ["tool-input-start", "tool-call"] + assert events[1].input == {"key": "value"} + + @pytest.mark.asyncio + async def test_name_only_chunk_emits_start_before_arguments_arrive(self): + acc, proc = _make_accumulator() + with patch("flocks.session.streaming.tool_accumulator.ToolRegistry") as mock_reg: + mock_reg.get_schema.return_value = None + + await acc.feed_chunk(_make_chunk(tc_id="call_name_first", name="edit")) + + events = [call.args[0] for call in proc.process_event.call_args_list] + assert len(events) == 1 + assert isinstance(events[0], ToolInputStartEvent) + assert events[0].tool_name == "edit" + + @pytest.mark.asyncio + async def test_generated_id_stays_stable_when_provider_id_arrives_late(self): + acc, proc = _make_accumulator() + with patch("flocks.session.streaming.tool_accumulator.ToolRegistry") as mock_reg: + mock_reg.get_schema.return_value = None + + await acc.feed_chunk(_make_chunk(index=3, name="write")) + start_event = proc.process_event.call_args_list[0].args[0] + + await acc.feed_chunk(_make_chunk( + index=3, + tc_id="provider_call_id", + arguments='{"filePath": "/tmp/f", "content": "done"}', + )) + + events = [call.args[0] for call in proc.process_event.call_args_list] + assert [event.type for event in events] == ["tool-input-start", "tool-call"] + assert events[1].tool_call_id == start_event.id @pytest.mark.asyncio async def test_completed_call_ignored_on_re_feed(self): diff --git a/tests/storage/test_storage.py b/tests/storage/test_storage.py index 063731c2b..36f644d2d 100644 --- a/tests/storage/test_storage.py +++ b/tests/storage/test_storage.py @@ -70,6 +70,85 @@ async def test_storage_with_model(storage): assert retrieved.value == 42 +@pytest.mark.asyncio +async def test_storage_set_many_commits_all_entries(storage): + """Multiple values should be persisted by one batch operation.""" + await storage.set_many([ + ("batch-set:key1", {"value": 1}, "test"), + ("batch-set:key2", StorageTestModel(id="m2", name="Beta", value=2), "test_model"), + ]) + + assert await storage.get("batch-set:key1") == {"value": 1} + model = await storage.get("batch-set:key2", StorageTestModel) + assert model == StorageTestModel(id="m2", name="Beta", value=2) + + +@pytest.mark.asyncio +async def test_storage_set_many_rolls_back_all_entries(storage, monkeypatch): + """A failed entry must not leave an earlier entry committed.""" + original_connect = Storage.connect + + @asynccontextmanager + async def failing_connect(db_path=None): + async with original_connect(db_path) as connection: + original_execute = connection.execute + calls = 0 + + async def fail_second_insert(sql, parameters=None): + nonlocal calls + if "INSERT OR REPLACE INTO storage" in sql: + calls += 1 + if calls == 2: + raise sqlite3.OperationalError("injected batch failure") + return await original_execute(sql, parameters) + + monkeypatch.setattr(connection, "execute", fail_second_insert) + yield connection + + with monkeypatch.context() as patcher: + patcher.setattr(Storage, "connect", failing_connect) + with pytest.raises(sqlite3.OperationalError, match="injected batch failure"): + await storage.set_many([ + ("rollback:key1", {"value": 1}, "test"), + ("rollback:key2", {"value": 2}, "test"), + ]) + + assert await storage.get("rollback:key1") is None + assert await storage.get("rollback:key2") is None + + +@pytest.mark.asyncio +async def test_storage_mutate_many_rolls_back_sets_and_deletes(storage, monkeypatch): + """A failed mixed mutation must preserve every pre-transaction value.""" + await storage.set("mixed:replace", {"value": "before"}, "test") + await storage.set("mixed:delete", {"value": "keep"}, "test") + original_connect = Storage.connect + + @asynccontextmanager + async def failing_connect(db_path=None): + async with original_connect(db_path) as connection: + original_execute = connection.execute + + async def fail_delete(sql, parameters=None): + if "DELETE FROM storage" in sql: + raise sqlite3.OperationalError("injected mixed mutation failure") + return await original_execute(sql, parameters) + + monkeypatch.setattr(connection, "execute", fail_delete) + yield connection + + with monkeypatch.context() as patcher: + patcher.setattr(Storage, "connect", failing_connect) + with pytest.raises(sqlite3.OperationalError, match="injected mixed mutation failure"): + await storage.mutate_many( + set_entries=[("mixed:replace", {"value": "after"}, "test")], + delete_keys=["mixed:delete"], + ) + + assert await storage.get("mixed:replace") == {"value": "before"} + assert await storage.get("mixed:delete") == {"value": "keep"} + + @pytest.mark.asyncio async def test_storage_delete(storage): """Test delete operation""" diff --git a/tests/tool/test_plan_exit.py b/tests/tool/test_plan_exit.py new file mode 100644 index 000000000..47e257252 --- /dev/null +++ b/tests/tool/test_plan_exit.py @@ -0,0 +1,173 @@ +from __future__ import annotations + +from pathlib import Path + +import pytest + +from flocks.session.execution_mode import SessionExecutionMode +from flocks.session.interaction_queue import InteractionQueue +from flocks.session.message import Message, MessageRole +from flocks.tool.registry import ToolContext, ToolResult +from flocks.tool.system import plan_exit + + +@pytest.fixture(autouse=True) +async def clear_queue(): + session_id = "plan-exit-session" + await InteractionQueue.clear(session_id) + await Message.clear(session_id) + yield + await InteractionQueue.clear(session_id) + await Message.clear(session_id) + + +def _context( + events: list[tuple[str, dict]], + tmp_path, + *, + plan_content: str = "# Approved plan\n", +) -> ToolContext: + async def publish(event_type: str, properties: dict) -> None: + events.append((event_type, properties)) + + plan_path = tmp_path / ".flocks" / "plans" / "1234-plan.md" + plan_path.parent.mkdir(parents=True, exist_ok=True) + plan_path.write_text(plan_content, encoding="utf-8") + return ToolContext( + session_id="plan-exit-session", + message_id="message-1", + agent="rex", + extra={ + "execution_mode": "plan", + "workspace_dir": str(tmp_path), + "model": {"providerID": "openai", "modelID": "gpt-test"}, + "plan_file_path": str(plan_path), + "plan_relative_path": ".flocks/plans/1234-plan.md", + "plan_permission_path": ".flocks/plans/1234-plan.md", + }, + event_publish_callback=publish, + ) + + +@pytest.mark.asyncio +async def test_plan_exit_approval_continues_immediately_in_build( + monkeypatch, + tmp_path, +) -> None: + events: list[tuple[str, dict]] = [] + + async def approve(*_args, **_kwargs): + return ToolResult( + success=True, + output="approved", + metadata={"answers": [[plan_exit.START_IMPLEMENTING]]}, + ) + + monkeypatch.setattr(plan_exit, "question_tool", approve) + + result = await plan_exit.plan_exit_tool(_context(events, tmp_path)) + messages = await Message.list("plan-exit-session") + build_message = messages[-1] + parts = await Message.parts(build_message.id, "plan-exit-session") + + assert result.success + assert result.metadata["approved"] is True + assert await InteractionQueue.list("plan-exit-session") == [] + assert build_message.role == MessageRole.USER + assert build_message.executionMode == SessionExecutionMode.BUILD + assert build_message.agent == "rex" + assert build_message.model == {"providerID": "openai", "modelID": "gpt-test"} + assert parts[0].synthetic is True + assert ".flocks/plans/1234-plan.md" in parts[0].text + assert result.metadata["planPath"] == ".flocks/plans/1234-plan.md" + assert [event_type for event_type, _ in events] == [ + "session.execution_mode.changed", + ] + + +@pytest.mark.asyncio +async def test_plan_exit_decline_stays_in_plan(monkeypatch, tmp_path) -> None: + async def decline(*_args, **_kwargs): + return ToolResult( + success=True, + output="declined", + metadata={"answers": [[plan_exit.CONTINUE_PLANNING]]}, + ) + + monkeypatch.setattr(plan_exit, "question_tool", decline) + + result = await plan_exit.plan_exit_tool(_context([], tmp_path)) + + assert result.success + assert result.metadata == {"approved": False, "executionMode": "plan"} + assert await InteractionQueue.list("plan-exit-session") == [] + + +@pytest.mark.asyncio +async def test_plan_exit_returns_continue_planning_feedback( + monkeypatch, + tmp_path, +) -> None: + async def provide_feedback(*_args, **_kwargs): + return ToolResult( + success=True, + output="feedback", + metadata={ + "answers": [ + [ + plan_exit.CONTINUE_PLANNING, + "Keep the public API unchanged.", + ] + ] + }, + ) + + monkeypatch.setattr(plan_exit, "question_tool", provide_feedback) + + result = await plan_exit.plan_exit_tool(_context([], tmp_path)) + + assert result.success + assert result.metadata == { + "approved": False, + "executionMode": "plan", + "feedback": "Keep the public API unchanged.", + } + assert "Keep the public API unchanged." in result.output + assert await InteractionQueue.list("plan-exit-session") == [] + + +@pytest.mark.asyncio +async def test_plan_exit_does_not_approve_deferred_channel_question( + monkeypatch, + tmp_path, +) -> None: + async def deferred(*_args, **_kwargs): + return ToolResult( + success=True, + output="sent", + metadata={"deferred": True}, + ) + + monkeypatch.setattr(plan_exit, "question_tool", deferred) + + result = await plan_exit.plan_exit_tool(_context([], tmp_path)) + + assert result.metadata["deferred"] is True + assert await InteractionQueue.list("plan-exit-session") == [] + + +@pytest.mark.asyncio +async def test_plan_exit_requires_non_empty_plan_file(tmp_path) -> None: + missing_ctx = _context([], tmp_path) + missing_path = missing_ctx.extra["plan_file_path"] + Path(missing_path).unlink() + + missing = await plan_exit.plan_exit_tool(missing_ctx) + empty = await plan_exit.plan_exit_tool( + _context([], tmp_path, plan_content=" \n") + ) + + assert not missing.success + assert "Write the session plan file" in (missing.error or "") + assert not empty.success + assert "empty" in (empty.error or "") diff --git a/tests/tool/test_question_channel.py b/tests/tool/test_question_channel.py index 50d5b6fee..963335f24 100644 --- a/tests/tool/test_question_channel.py +++ b/tests/tool/test_question_channel.py @@ -21,6 +21,14 @@ def test_normalize_question_option_accepts_common_llm_shapes() -> None: "label": "Only descriptive text", "description": "", } + assert normalize_question_option({ + "label": "调整计划", + "allowText": True, + }) == { + "label": "调整计划", + "description": "", + "allowText": True, + } assert normalize_question_option({"label": ""}) is None diff --git a/tests/tool/test_sangfor_edr_handler.py b/tests/tool/test_sangfor_edr_handler.py index 3eb286367..dd18988ab 100644 --- a/tests/tool/test_sangfor_edr_handler.py +++ b/tests/tool/test_sangfor_edr_handler.py @@ -120,8 +120,8 @@ def test_complete_manual_login_saves_state(tmp_path, monkeypatch): monkeypatch.setattr(handler, "_ensure_browser_daemon", lambda: None) monkeypatch.setattr(handler.helpers, "page_info", lambda: {"url": "https://edr.example.com/ui/#/index"}) monkeypatch.setattr(handler, "_is_logged_in", lambda cfg: True) - monkeypatch.setattr(handler, "_save_auth_state", lambda cfg: {"cookies": 1}) - monkeypatch.setattr(handler, "_save_captured_login_token", lambda: True) + monkeypatch.setattr(handler, "_save_browser_auth_pair", lambda cfg: {"cookies": 1}) + monkeypatch.setattr(handler, "_probe_auth_pair", lambda cfg: {"valid": True}) result = handler._complete_manual_login(cfg) @@ -181,3 +181,310 @@ def test_login_token_capture_hooks_fetch_and_xhr(monkeypatch): assert "/launch_login.php" in scripts[0] assert "window.fetch" in scripts[0] assert "XMLHttpRequest.prototype" in scripts[0] + + +def test_normalise_base_url_extracts_origin_from_page_url(): + handler = _load_handler() + + assert ( + handler._normalise_base_url("https://edr.example.com:8443/ui/#/index") + == "https://edr.example.com:8443" + ) + assert handler._normalise_base_url("edr.example.com/ui/login.php") == "https://edr.example.com" + + +def test_auth_pair_rejects_cookie_token_mismatch(tmp_path, monkeypatch): + handler = _load_handler() + state_path = tmp_path / "auth-state.json" + state_path.write_text( + '{"cookies":[{"name":"sessionid","value":"new","domain":"edr.example.com","path":"/"}],"origins":[]}', + encoding="utf-8", + ) + old_cookies = [{"name": "sessionid", "value": "old", "domain": "edr.example.com", "path": "/"}] + bundle = { + "token": "old-token", + "base_url": "https://edr.example.com", + "cookie_fingerprint": handler._cookie_fingerprint(old_cookies, "https://edr.example.com"), + } + monkeypatch.setattr( + handler, + "_get_secret_manager", + lambda: type("Secrets", (), {"get": lambda self, key: __import__("json").dumps(bundle)})(), + ) + + try: + handler._load_verified_auth_pair(_cfg(handler, state_path)) + except RuntimeError as exc: + assert "not from the same login" in str(exc) + else: + raise AssertionError("mismatched EDR authentication must be rejected") + + +def test_http_login_saves_matched_cookie_and_token(tmp_path, monkeypatch): + handler = _load_handler() + http = handler._http_login_module + cfg = _cfg(handler, tmp_path / "auth-state.json") + + class Response: + def __init__(self, data=None, content=b"captcha"): + self._data = data or {} + self.content = content + + def raise_for_status(self): + return None + + def json(self): + return self._data + + class Session: + def __init__(self): + self.cookies = http.requests.cookies.RequestsCookieJar() + self.cookies.set("sessionid", "cookie-value", domain="edr.example.com", path="/") + self.posts = [] + + def get(self, url, **kwargs): + return Response() + + def post(self, url, **kwargs): + self.posts.append((url, kwargs.get("json"))) + payload = kwargs.get("json") or {} + if payload.get("opr") == "rsakey": + return Response({"success": True, "key": "c7" * 64}) + if url.endswith("/login"): + return Response({"success": True, "key": 7}) + return Response({"success": True, "data": {"token": "token-value"}}) + + saved = {} + monkeypatch.setattr(http, "http_session", lambda cfg: Session()) + monkeypatch.setattr(http, "_ocr_verify_code", lambda content: "1234") + monkeypatch.setattr( + http, + "_save_auth_pair", + lambda cfg, state, token: saved.update({"state": state, "token": token}) or {"pair_verified": True}, + ) + + result = http._http_login(cfg) + + assert result["status"] == "http_login_refreshed_auth_state" + assert saved["token"] == "token-value" + assert saved["state"]["cookies"][0]["value"] == "cookie-value" + + +def test_dashboard_request_definitions_use_dynamic_base_inputs(tmp_path): + handler = _load_handler() + cfg = _cfg(handler, tmp_path / "auth-state.json") + + definitions = handler._dashboard_requests(cfg, "token-value", days=7) + + agent_path, agent_payload = definitions["agent_overview"] + assert agent_path.endswith("opr=get_agent_overview") + assert agent_payload["app_args"]["name"] == "app.web.event_center.head" + vulner_path, vulner_payload = definitions["vulnerability_overview"] + assert "s=token-value" in vulner_path + assert vulner_payload["uid"] == "admin" + assert vulner_payload["token"] == "token-value" + + +def test_auth_probe_requires_http_200_and_agent_overview_data(tmp_path, monkeypatch): + handler = _load_handler() + cfg = _cfg(handler, tmp_path / "auth-state.json") + monkeypatch.setattr(handler, "_load_verified_auth_pair", lambda cfg: ({"cookies": []}, "token")) + + class Response: + status_code = 200 + headers = {} + + def json(self): + return {"success": True, "data": {"total": 12, "online": 10}} + + class Session: + def post(self, *args, **kwargs): + assert kwargs["allow_redirects"] is False + assert args[0].endswith("opr=get_agent_overview") + assert kwargs["json"]["app_args"]["options"] == {} + assert kwargs["json"]["opr"] == "get_agent_overview" + assert "date_range" in kwargs["json"] + return Response() + + monkeypatch.setattr(handler, "_dashboard_session", lambda cfg, state: Session()) + + result = handler._probe_auth_pair(cfg) + + assert result == { + "valid": True, + "reason": "auth_probe_succeeded", + "http_status": 200, + "agent_overview_verified": True, + } + + +def test_auth_probe_rejects_success_without_agent_overview_data(tmp_path, monkeypatch): + handler = _load_handler() + cfg = _cfg(handler, tmp_path / "auth-state.json") + monkeypatch.setattr(handler, "_load_verified_auth_pair", lambda cfg: ({"cookies": []}, "token")) + + class Response: + status_code = 200 + headers = {} + + def json(self): + return {"success": True, "message": "ok"} + + class Session: + def post(self, *args, **kwargs): + return Response() + + monkeypatch.setattr(handler, "_dashboard_session", lambda cfg, state: Session()) + + result = handler._probe_auth_pair(cfg) + + assert result == { + "valid": False, + "reason": "auth_probe_expected_agent_data_missing", + } + + +def test_auth_probe_rejects_login_redirect(tmp_path, monkeypatch): + handler = _load_handler() + cfg = _cfg(handler, tmp_path / "auth-state.json") + monkeypatch.setattr(handler, "_load_verified_auth_pair", lambda cfg: ({"cookies": []}, "token")) + + class Response: + status_code = 302 + headers = {"Location": "/ui/login.php"} + + class Session: + def post(self, *args, **kwargs): + return Response() + + monkeypatch.setattr(handler, "_dashboard_session", lambda cfg, state: Session()) + + result = handler._probe_auth_pair(cfg) + + assert result["valid"] is False + assert result["reason"] == "auth_probe_redirected" + + +def test_http_auth_reuses_valid_pair_without_login(tmp_path, monkeypatch): + handler = _load_handler() + http = handler._http_login_module + cfg = _cfg(handler, tmp_path / "auth-state.json") + monkeypatch.setattr( + http, + "probe_auth_pair", + lambda cfg: {"valid": True, "reason": "auth_probe_succeeded"}, + ) + monkeypatch.setattr( + http, + "_http_login", + lambda *args, **kwargs: (_ for _ in ()).throw(AssertionError("login must be skipped")), + ) + + result = http.ensure_http_auth_pair(cfg) + + assert result["status"] == "http_auth_pair_reused" + assert result["login_skipped"] is True + + +def test_http_auth_relogs_and_confirms_when_probe_fails(tmp_path, monkeypatch): + handler = _load_handler() + http = handler._http_login_module + cfg = _cfg(handler, tmp_path / "auth-state.json") + probes = iter( + [ + {"valid": False, "reason": "auth_probe_unauthorized"}, + {"valid": True, "reason": "auth_probe_succeeded"}, + ] + ) + monkeypatch.setattr(http, "probe_auth_pair", lambda cfg: next(probes)) + monkeypatch.setattr( + http, + "_http_login", + lambda cfg, captcha_code="": {"success": True, "status": "http_login_refreshed_auth_state"}, + ) + + result = http.ensure_http_auth_pair(cfg) + + assert result["success"] is True + assert result["previous_probe"]["reason"] == "auth_probe_unauthorized" + assert result["probe"]["valid"] is True + + +def test_dashboard_error_redacts_login_token(): + handler = _load_handler() + + error = handler._safe_error( + RuntimeError("401 Client Error for url: https://edr.test/launch.php?s=secret-token&opr=check"), + "secret-token", + ) + + assert "secret-token" not in error + assert "" in error + + +def test_http_login_reports_failure_phase(tmp_path, monkeypatch): + handler = _load_handler() + http = handler._http_login_module + cfg = _cfg(handler, tmp_path / "auth-state.json") + + class Session: + def get(self, *args, **kwargs): + raise http.requests.ConnectionError("connection refused") + + monkeypatch.setattr(http, "http_session", lambda cfg: Session()) + + result = http._http_login(cfg) + + assert result["status"] == "http_login_failed" + assert result["phase"] == "login_page" + assert "phase=login_page" in result["error"] + assert "connection refused" in result["error"] + + +def test_http_login_retries_captcha_dlogin_failure(tmp_path, monkeypatch): + handler = _load_handler() + http = handler._http_login_module + cfg = _cfg(handler, tmp_path / "auth-state.json") + cfg.max_captcha_retry = 2 + + class Response: + def __init__(self, payload): + self.payload = payload + self.content = b"captcha" + + def raise_for_status(self): + return None + + def json(self): + return self.payload + + class Session: + def __init__(self): + self.cookies = http.requests.cookies.RequestsCookieJar() + self.cookies.set("sessionid", "cookie-value", domain="edr.example.com", path="/") + self.dlogin_count = 0 + + def get(self, url, **kwargs): + return Response({}) + + def post(self, url, **kwargs): + payload = kwargs.get("json") or {} + if payload.get("opr") == "rsakey": + return Response({"success": True, "key": "c7" * 64}) + if url.endswith("/login"): + self.dlogin_count += 1 + if self.dlogin_count == 1: + return Response({"success": False, "msg": "验证码错误"}) + return Response({"success": True, "key": 7}) + return Response({"success": True, "data": {"token": "token-value"}}) + + session = Session() + monkeypatch.setattr(http, "http_session", lambda cfg: session) + monkeypatch.setattr(http, "_ocr_verify_code", lambda content: "1234") + monkeypatch.setattr(http, "_save_auth_pair", lambda cfg, state, token: {"pair_verified": True}) + + result = http._http_login(cfg) + + assert result["success"] is True + assert result["attempt"] == 2 + assert session.dlogin_count == 2 diff --git a/tests/tool/test_session_manage_tool.py b/tests/tool/test_session_manage_tool.py index cae00eb23..dff530441 100644 --- a/tests/tool/test_session_manage_tool.py +++ b/tests/tool/test_session_manage_tool.py @@ -1,10 +1,11 @@ +from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, patch import pytest from flocks.tool.registry import ToolContext, ToolRegistry, ToolResult import flocks.tool.system.session_manage # noqa: F401 - ensure tool registration -from flocks.tool.system.session_manage import session_manage +from flocks.tool.system.session_manage import _session_archive_impl, session_manage def make_ctx() -> ToolContext: @@ -83,5 +84,62 @@ async def test_session_manage_delete_requests_confirmation(): ctx.ask.assert_awaited_once() ask_kwargs = ctx.ask.await_args.kwargs assert ask_kwargs["permission"] == "session_manage" - assert ask_kwargs["metadata"] == {"action": "delete", "session_id": "ses_123"} + assert ask_kwargs["metadata"] == { + "action": "permanent_delete", + "session_id": "ses_123", + "destructive": True, + } delete_impl.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_session_manage_rejects_archiving_the_current_running_session(): + ctx = make_ctx() + ctx.session_id = "ses_current" + session = SimpleNamespace( + id="ses_current", + project_id="project", + title="Current", + status="active", + ) + + with ( + patch("flocks.storage.storage.Storage.list_keys", AsyncMock(return_value=["session:project:ses_current"])), + patch("flocks.storage.storage.Storage.get", AsyncMock(return_value=session)), + patch("flocks.session.session.Session.archive", AsyncMock()) as archive, + ): + result = await _session_archive_impl(ctx, "ses_current", True) + + assert result.success is False + assert "不能在当前会话" in (result.error or "") + archive.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_session_manage_restore_revives_removed_project_first(): + ctx = make_ctx() + ctx.session_id = "ses_operator" + session = SimpleNamespace( + id="ses_archived", + project_id="project_removed", + title="Archived", + status="archived", + owner_user_id="usr_owner", + ) + + with ( + patch( + "flocks.storage.storage.Storage.list_keys", + AsyncMock(return_value=["session:project_removed:ses_archived"]), + ), + patch("flocks.storage.storage.Storage.get", AsyncMock(return_value=session)), + patch("flocks.session.session.Session.restore", AsyncMock(return_value=True)) as restore_session, + ): + result = await _session_archive_impl(ctx, session.id, False) + + assert result.success is True + restore_session.assert_awaited_once_with( + session.project_id, + session.id, + project_owner_id="usr_owner", + ) diff --git a/tests/tool/test_subagent_hooks.py b/tests/tool/test_subagent_hooks.py new file mode 100644 index 000000000..11b83aad0 --- /dev/null +++ b/tests/tool/test_subagent_hooks.py @@ -0,0 +1,261 @@ +"""Tests for paired delegate_task subagent lifecycle hooks.""" + +from __future__ import annotations + +import asyncio +from types import SimpleNamespace +from unittest.mock import AsyncMock, patch + +import pytest + +from flocks.tool.agent.delegate_task import _run_subagent_with_hooks +from flocks.tool.registry import ToolContext + + +@pytest.mark.asyncio +@pytest.mark.parametrize("resumed", [False, True]) +async def test_subagent_hooks_wrap_child_run(resumed: bool) -> None: + ctx = ToolContext( + session_id="ses_parent", + message_id="msg_parent", + agent="rex", + ) + last_message = SimpleNamespace(id="msg_child_final") + loop_result = SimpleNamespace( + action="stop", + error=None, + last_message=last_message, + ) + start_hook = AsyncMock() + stop_hook = AsyncMock() + + with ( + patch( + "flocks.tool.agent.delegate_task.SessionLoop.run", + AsyncMock(return_value=loop_result), + ) as run_loop, + patch( + "flocks.tool.agent.delegate_task.Message.get_text_content", + AsyncMock(return_value="child summary"), + ), + patch( + "flocks.hooks.pipeline.HookPipeline.run_subagent_start", + start_hook, + ), + patch( + "flocks.hooks.pipeline.HookPipeline.run_subagent_stop", + stop_hook, + ), + ): + result = await _run_subagent_with_hooks( + ctx=ctx, + child_session_id="ses_child", + child_agent="explore", + workspace="/tmp/project", + prompt="inspect hooks", + description="Inspect hooks", + resumed=resumed, + ) + + assert result is loop_result + run_loop.assert_awaited_once_with( + "ses_child", + provider_id=None, + model_id=None, + callbacks=None, + ) + start_payload = start_hook.await_args.args[0] + assert start_payload["parentSessionID"] == "ses_parent" + assert start_payload["childSessionID"] == "ses_child" + assert start_payload["agentType"] == "explore" + assert start_payload["resumed"] is resumed + stop_payload = stop_hook.await_args.args[0] + assert stop_payload["status"] == "completed" + assert stop_payload["summary"] == "child summary" + assert stop_payload["durationMs"] >= 0 + + +@pytest.mark.asyncio +async def test_subagent_stop_reports_child_error() -> None: + ctx = ToolContext( + session_id="ses_parent", + message_id="msg_parent", + agent="rex", + ) + start_hook = AsyncMock() + stop_hook = AsyncMock() + + with ( + patch( + "flocks.tool.agent.delegate_task.SessionLoop.run", + AsyncMock(side_effect=RuntimeError("child failed")), + ), + patch( + "flocks.hooks.pipeline.HookPipeline.run_subagent_start", + start_hook, + ), + patch( + "flocks.hooks.pipeline.HookPipeline.run_subagent_stop", + stop_hook, + ), + ): + with pytest.raises(RuntimeError, match="child failed"): + await _run_subagent_with_hooks( + ctx=ctx, + child_session_id="ses_child", + child_agent="explore", + workspace="/tmp/project", + prompt="inspect hooks", + description="Inspect hooks", + resumed=False, + ) + + start_hook.assert_awaited_once() + stop_hook.assert_awaited_once() + stop_payload = stop_hook.await_args.args[0] + assert stop_payload["status"] == "error" + assert stop_payload["summary"] is None + assert stop_payload["error"] == "child failed" + + +@pytest.mark.asyncio +async def test_subagent_stop_reports_persisted_message_error() -> None: + ctx = ToolContext( + session_id="ses_parent", + message_id="msg_parent", + agent="rex", + ) + last_message = SimpleNamespace( + id="msg_child_error", + finish="error", + error={"message": "provider authentication failed"}, + ) + loop_result = SimpleNamespace( + action="stop", + error=None, + last_message=last_message, + metadata={}, + ) + stop_hook = AsyncMock() + + with ( + patch( + "flocks.tool.agent.delegate_task.SessionLoop.run", + AsyncMock(return_value=loop_result), + ), + patch( + "flocks.tool.agent.delegate_task.Message.get_text_content", + AsyncMock(return_value=""), + ), + patch( + "flocks.hooks.pipeline.HookPipeline.run_subagent_start", + AsyncMock(), + ), + patch( + "flocks.hooks.pipeline.HookPipeline.run_subagent_stop", + stop_hook, + ), + ): + await _run_subagent_with_hooks( + ctx=ctx, + child_session_id="ses_child", + child_agent="explore", + workspace="/tmp/project", + prompt="inspect hooks", + description="Inspect hooks", + resumed=False, + ) + + stop_hook.assert_awaited_once() + stop_payload = stop_hook.await_args.args[0] + assert stop_payload["status"] == "error" + assert stop_payload["summary"] == "" + assert stop_payload["error"] == "provider authentication failed" + + +@pytest.mark.asyncio +async def test_subagent_stop_reports_interruption() -> None: + ctx = ToolContext( + session_id="ses_parent", + message_id="msg_parent", + agent="rex", + ) + stop_hook = AsyncMock() + + with ( + patch( + "flocks.tool.agent.delegate_task.SessionLoop.run", + AsyncMock(side_effect=asyncio.CancelledError()), + ), + patch( + "flocks.hooks.pipeline.HookPipeline.run_subagent_start", + AsyncMock(), + ), + patch( + "flocks.hooks.pipeline.HookPipeline.run_subagent_stop", + stop_hook, + ), + ): + with pytest.raises(asyncio.CancelledError): + await _run_subagent_with_hooks( + ctx=ctx, + child_session_id="ses_child", + child_agent="explore", + workspace="/tmp/project", + prompt="inspect hooks", + description="Inspect hooks", + resumed=True, + ) + + stop_hook.assert_awaited_once() + stop_payload = stop_hook.await_args.args[0] + assert stop_payload["status"] == "interrupted" + assert stop_payload["summary"] is None + assert stop_payload["error"] == "Sub-agent execution was interrupted" + + +@pytest.mark.asyncio +async def test_subagent_stop_reports_normalized_loop_abort() -> None: + ctx = ToolContext( + session_id="ses_parent", + message_id="msg_parent", + agent="rex", + ) + stop_hook = AsyncMock() + loop_result = SimpleNamespace( + action="stop", + error=None, + last_message=None, + metadata={"aborted": True}, + ) + + with ( + patch( + "flocks.tool.agent.delegate_task.SessionLoop.run", + AsyncMock(return_value=loop_result), + ), + patch( + "flocks.hooks.pipeline.HookPipeline.run_subagent_start", + AsyncMock(), + ), + patch( + "flocks.hooks.pipeline.HookPipeline.run_subagent_stop", + stop_hook, + ), + ): + result = await _run_subagent_with_hooks( + ctx=ctx, + child_session_id="ses_child", + child_agent="explore", + workspace="/tmp/project", + prompt="inspect hooks", + description="Inspect hooks", + resumed=False, + ) + + assert result is loop_result + stop_hook.assert_awaited_once() + stop_payload = stop_hook.await_args.args[0] + assert stop_payload["status"] == "interrupted" + assert stop_payload["summary"] is None + assert stop_payload["error"] == "Sub-agent execution was interrupted" diff --git a/tests/tool/test_tools.py b/tests/tool/test_tools.py index 895cf0f80..dfba07484 100644 --- a/tests/tool/test_tools.py +++ b/tests/tool/test_tools.py @@ -17,6 +17,7 @@ import os import tempfile import shutil +import time import uuid from pathlib import Path from typing import Dict, Any, List @@ -920,6 +921,46 @@ async def test_get_time_tool_after_init(self, tool_context): # Should return ISO format datetime assert "T" in result.output # ISO format contains 'T' + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("output_format", "scale"), + [("unix", 1), ("unix_ms", 1000)], + ) + async def test_get_time_tool_timestamp_formats( + self, + tool_context, + output_format, + scale, + ): + """Test Unix timestamp output in seconds and milliseconds.""" + ToolRegistry.init() + before = int(time.time() * scale) + + result = await ToolRegistry.execute( + "get_time", + ctx=tool_context, + format=output_format, + ) + + after = int(time.time() * scale) + assert result.success + assert result.output.isdigit() + assert before <= int(result.output) <= after + + @pytest.mark.asyncio + async def test_get_time_tool_rejects_invalid_format(self, tool_context): + """Test invalid output formats return a clear error.""" + ToolRegistry.init() + + result = await ToolRegistry.execute( + "get_time", + ctx=tool_context, + format="invalid", + ) + + assert not result.success + assert result.error == "format must be one of: iso, unix, unix_ms" + # ============================================================================= # ToolContext Tests diff --git a/uv.lock b/uv.lock index 12a019c9d..5eb156e92 100644 --- a/uv.lock +++ b/uv.lock @@ -553,7 +553,7 @@ wheels = [ [[package]] name = "flocks" -version = "2026.7.23" +version = "2026.7.29" source = { editable = "." } dependencies = [ { name = "aiofiles" }, diff --git a/webui/src/api/provider.ts b/webui/src/api/provider.ts index e0f3062ef..de9467005 100644 --- a/webui/src/api/provider.ts +++ b/webui/src/api/provider.ts @@ -249,6 +249,7 @@ export const defaultModelAPI = { /** Delete default model for a type */ delete: (modelType: string) => client.delete(`/api/default-model/${modelType}`), + }; // ==================== Usage API ==================== diff --git a/webui/src/api/session.test.ts b/webui/src/api/session.test.ts index 3303e361e..8f10c48f2 100644 --- a/webui/src/api/session.test.ts +++ b/webui/src/api/session.test.ts @@ -94,4 +94,14 @@ describe('sessionApi message actions', () => { expect(mockDelete).toHaveBeenCalledWith('/api/session/session-1/prompt_queue/queue-1'); expect(mockPost).toHaveBeenCalledWith('/api/session/session-1/prompt_queue/queue-2/run_now'); }); + + it('calls archive and restore endpoints', async () => { + const { sessionApi } = await import('./session'); + + await sessionApi.archive('session-1'); + await sessionApi.restore('session-1'); + + expect(mockPost).toHaveBeenCalledWith('/api/session/session-1/archive'); + expect(mockPost).toHaveBeenCalledWith('/api/session/session-1/restore'); + }); }); diff --git a/webui/src/api/session.ts b/webui/src/api/session.ts index a421cb219..fd7fd63cb 100644 --- a/webui/src/api/session.ts +++ b/webui/src/api/session.ts @@ -1,4 +1,6 @@ import client from './client'; +import type { Session } from '@/types'; +import type { SessionExecutionMode } from '@/utils/sessionExecutionMode'; export interface SessionMessagePartPayload { id: string; @@ -29,6 +31,7 @@ export interface QueuedPrompt { status: 'pending' | 'executing' | string; createdAt: number; updatedAt: number; + executionMode?: SessionExecutionMode; } export interface PromptQueueResponse { @@ -82,6 +85,7 @@ export interface SessionListParams { start?: number; search?: string; category?: string; + status?: 'active' | 'archived' | 'all'; } export interface SessionMessagePage { @@ -102,8 +106,8 @@ export const sessionApi = { /** * 获取会话列表 */ - list: async (params?: SessionListParams) => { - const response = await client.get('/api/session', { params }); + list: async (params?: SessionListParams): Promise => { + const response = await client.get('/api/session', { params }); return response.data; }, @@ -126,27 +130,57 @@ export const sessionApi = { /** * 创建会话 */ - create: async (data?: { title?: string; parentID?: string; projectID?: string }) => { + create: async (data?: { title?: string; parentID?: string; projectID?: string; model_auto?: boolean }) => { const response = await client.post('/api/session', data || {}); return response.data; }, /** - * 删除会话 + * 永久删除会话(普通工作台应使用 archive) */ - delete: async (sessionId: string) => { - const response = await client.delete(`/api/session/${sessionId}`); + delete: async (sessionId: string): Promise => { + const response = await client.delete(`/api/session/${sessionId}`); + return response.data; + }, + + /** + * 归档会话并保留全部持久化数据 + */ + archive: async (sessionId: string): Promise => { + const response = await client.post(`/api/session/${sessionId}/archive`); + return response.data; + }, + + /** + * 恢复已归档会话 + */ + restore: async (sessionId: string): Promise => { + const response = await client.post(`/api/session/${sessionId}/restore`); return response.data; }, /** * 更新会话 */ - update: async (sessionId: string, data: { title?: string; provider?: string; model?: string; model_pinned?: boolean }) => { + update: async (sessionId: string, data: { + title?: string; + provider?: string; + model?: string; + model_pinned?: boolean; + model_auto?: boolean; + }) => { const response = await client.patch(`/api/session/${sessionId}`, data); return response.data; }, + /** + * 将任务及其子任务移动到指定项目 + */ + moveToProject: async (sessionId: string, projectID: string): Promise => { + const response = await client.patch(`/api/session/${sessionId}/project`, { projectID }); + return response.data; + }, + /** * 本地共享会话(所有本地账号可见,只读) */ @@ -215,6 +249,7 @@ export const sessionApi = { model?: Record; variant?: string; displayText?: string; + executionMode?: SessionExecutionMode; }) => { const response = await client.post(`/api/session/${sessionId}/prompt_queue`, data); return response.data; diff --git a/webui/src/api/webuiContractPages.ts b/webui/src/api/webuiContractPages.ts index 5cb3f3846..abdee30af 100644 --- a/webui/src/api/webuiContractPages.ts +++ b/webui/src/api/webuiContractPages.ts @@ -29,6 +29,7 @@ export interface WebUIContractWorkspaceSection { export interface WebUIContractWorkspaceListItem { id: string; + version?: string; title: string; titleEn?: string | null; route: string; diff --git a/webui/src/components/common/ChatPromptSelectors.test.tsx b/webui/src/components/common/ChatPromptSelectors.test.tsx index 60ea93cbe..8bbe976fb 100644 --- a/webui/src/components/common/ChatPromptSelectors.test.tsx +++ b/webui/src/components/common/ChatPromptSelectors.test.tsx @@ -37,6 +37,7 @@ vi.mock('react-i18next', () => ({ 'modelPicker.empty': '暂无模型', 'modelPicker.count': `${params?.count ?? 0}`, 'modelPicker.vision': '视觉', + 'modelPicker.auto': 'Auto', loading: '加载中', }; return translations[key] ?? key; @@ -104,7 +105,7 @@ beforeEach(() => { }); describe('ChatModelPicker', () => { - it('opens the model menu toward the left edge of the trigger', async () => { + it('opens the model menu from the left edge of the trigger', async () => { const user = userEvent.setup(); render( @@ -120,13 +121,157 @@ describe('ChatModelPicker', () => { const menu = screen.getByText('选择模型').closest('.absolute'); expect(menu).not.toBeNull(); - expect(menu).toHaveClass('right-0'); + expect(menu).toHaveClass('left-0'); expect(menu).toHaveClass('bottom-full'); - expect(menu).not.toHaveClass('left-0'); + expect(menu).not.toHaveClass('right-0'); + expect(menu).toHaveStyle({ transform: 'translateX(0px)' }); + }); + + it('shifts the left-anchored menu only when the right edge would overflow', async () => { + const user = userEvent.setup(); + + render( + , + ); + + const trigger = screen.getByRole('button', { name: /minimax-m3/i }); + const selector = trigger.closest('[data-model-selector]'); + vi.spyOn(selector!, 'getBoundingClientRect').mockReturnValue({ + bottom: 0, + height: 0, + left: 700, + right: 700, + top: 0, + width: 0, + x: 700, + y: 0, + toJSON: () => ({}), + }); + const originalViewportWidth = window.innerWidth; + Object.defineProperty(window, 'innerWidth', { configurable: true, value: 1000 }); + + await user.click(trigger); + + const menu = screen.getByText('选择模型').closest('.absolute'); + expect(menu).toHaveStyle({ transform: 'translateX(-36px)' }); + Object.defineProperty(window, 'innerWidth', { configurable: true, value: originalViewportWidth }); + }); + + it('renders Auto as an opt-in single-line item with its hint in the info tooltip', async () => { + const user = userEvent.setup(); + const onSelectAuto = vi.fn(); + + render( + , + ); + + await user.click(screen.getByRole('button', { name: /minimax-m3/i })); + const autoButton = screen.getByRole('button', { name: 'Auto' }); + expect(screen.queryByText('Primary then fallback')).not.toBeInTheDocument(); + + const info = autoButton.querySelector('.lucide-info')?.parentElement; + expect(info).toBeInTheDocument(); + await user.hover(info!); + expect(await screen.findByText('Primary then fallback')).toBeInTheDocument(); + await user.click(autoButton); + expect(onSelectAuto).toHaveBeenCalledOnce(); + }); + + it('shows only Auto as selected when Auto mode is active', async () => { + const user = userEvent.setup(); + + render( + , + ); + + const trigger = screen.getByRole('button', { name: /^Auto/i }); + expect(trigger).toHaveAttribute('title', 'Auto: Primary then fallback'); + await user.click(trigger); + + expect(screen.getAllByRole('button', { name: 'Auto' })[1]).toHaveClass('shadow-[inset_2px_0_0_#a1a1aa]'); + expect(screen.getByRole('button', { name: /minimax-m3/i })).not.toHaveClass('shadow-[inset_2px_0_0_#a1a1aa]'); + }); + + it('does not show Auto unless the caller opts in', async () => { + const user = userEvent.setup(); + + render( + , + ); + + await user.click(screen.getByRole('button', { name: /minimax-m3/i })); + expect(screen.queryByRole('button', { name: 'Auto' })).not.toBeInTheDocument(); }); }); describe('useChatModelOptions', () => { + it('keeps Auto opt-in and clears it when a concrete model is selected', async () => { + listDefinitionsMock.mockResolvedValue({ + data: { models: [makeModelDefinition()] }, + }); + getResolvedMock.mockResolvedValue({ + data: { provider_id: 'provider-1', model_id: 'model-1' }, + }); + + const { result } = renderHook(() => useChatModelOptions({ enableAuto: true })); + + await waitFor(() => { + expect(result.current.canSelectAuto).toBe(true); + expect(result.current.selectedModelKey).toBe('provider-1::model-1'); + }); + expect(result.current.selectedModelAuto).toBe(false); + expect(result.current.selectedPromptModel).toEqual({ + providerID: 'provider-1', + modelID: 'model-1', + }); + + act(() => result.current.selectAuto()); + + expect(result.current.selectedModelAuto).toBe(true); + expect(result.current.selectedPromptModel).toBeNull(); + expect(result.current.effectiveModelOption).toEqual(result.current.primaryModelOption); + + act(() => result.current.selectModelKey('provider-1::model-1')); + + expect(result.current.selectedModelAuto).toBe(false); + expect(result.current.selectedPromptModel).toEqual({ + providerID: 'provider-1', + modelID: 'model-1', + }); + }); + it('shares enabled model and default model requests across concurrent hook instances', async () => { let resolveDefinitions: (value: { data: { models: any[] } }) => void = () => {}; listDefinitionsMock.mockReturnValue(new Promise((resolve) => { diff --git a/webui/src/components/common/ChatPromptSelectors.tsx b/webui/src/components/common/ChatPromptSelectors.tsx index 0df85d411..c39b6bec2 100644 --- a/webui/src/components/common/ChatPromptSelectors.tsx +++ b/webui/src/components/common/ChatPromptSelectors.tsx @@ -1,5 +1,5 @@ -import { useCallback, useEffect, useMemo, useState } from 'react'; -import { Bot, ChevronDown, Cpu, Info } from 'lucide-react'; +import { useCallback, useEffect, useMemo, useRef, useState } from 'react'; +import { Bot, ChevronDown, Cpu, Info, Sparkles } from 'lucide-react'; import { useTranslation } from 'react-i18next'; import type { Agent } from '@/api/agent'; @@ -33,6 +33,13 @@ export type ChatModelProviderGroup = { models: ChatModelOption[]; }; +export type ChatModelPickerAutoOption = { + selected: boolean; + disabled: boolean; + statusLabel: string; + onSelect: () => void; +}; + type SelectorTooltip = { title: string; lines: string[]; @@ -40,6 +47,23 @@ type SelectorTooltip = { y: number; }; +const MODEL_MENU_WIDTH_PX = 320; +const MODEL_MENU_VIEWPORT_PADDING_PX = 16; + +export function getAnchoredMenuLeftOffset( + anchorLeft: number, + viewportWidth: number, + menuWidth = MODEL_MENU_WIDTH_PX, + viewportPadding = MODEL_MENU_VIEWPORT_PADDING_PX, +): number { + const constrainedWidth = Math.min( + menuWidth, + Math.max(0, viewportWidth - viewportPadding * 2), + ); + const rightOverflow = anchorLeft + constrainedWidth + viewportPadding - viewportWidth; + return rightOverflow > 0 ? -rightOverflow : 0; +} + export function __resetChatModelOptionsResourcesForTesting(): void { __resetChatModelResourcesForTesting(); } @@ -73,7 +97,7 @@ export function useChatAgentOptions(options: { allowedAgentNames?: string[] } = }; } -export function useChatModelOptions() { +export function useChatModelOptions({ enableAuto = false }: { enableAuto?: boolean } = {}) { const { t } = useTranslation('session'); const { providers, loading: loadingProviders } = useProviders(); const { @@ -81,6 +105,7 @@ export function useChatModelOptions() { loading: loadingEnabledModels, } = useEnabledChatModelDefinitions(); const [selectedModelKey, setSelectedModelKey] = useState(null); + const [selectedModelAuto, setSelectedModelAuto] = useState(false); const options = useMemo(() => { const providerById = new Map( @@ -158,6 +183,34 @@ export function useChatModelOptions() { data: resolvedDefaultModel, initialized: resolvedDefaultModelInitialized, } = useResolvedDefaultModel(options.length > 0); + const primaryModelOption = useMemo(() => { + if (!resolvedDefaultModel) return null; + return options.find((option) => ( + option.providerID === resolvedDefaultModel.providerID + && option.modelID === resolvedDefaultModel.modelID + )) ?? null; + }, [options, resolvedDefaultModel]); + const canSelectAuto = Boolean(enableAuto && primaryModelOption); + const effectiveModelOption = selectedModelAuto ? primaryModelOption : selectedModelOption; + const selectModelKey = useCallback((key: string) => { + setSelectedModelAuto(false); + setSelectedModelKey(key); + }, []); + const selectAuto = useCallback(() => { + if (canSelectAuto) setSelectedModelAuto(true); + }, [canSelectAuto]); + const modelPickerAutoOption = useMemo(() => ( + enableAuto + ? { + selected: selectedModelAuto, + disabled: !canSelectAuto, + statusLabel: canSelectAuto + ? t('modelPicker.autoHint') + : t('modelPicker.autoUnavailable'), + onSelect: selectAuto, + } + : undefined + ), [canSelectAuto, enableAuto, selectAuto, selectedModelAuto, t]); useEffect(() => { if (selectedModelKey || options.length === 0 || !resolvedDefaultModelInitialized) return; @@ -178,11 +231,19 @@ export function useChatModelOptions() { groupedOptions, loading: loadingProviders || loadingEnabledModels, options, + canSelectAuto, + effectiveModelOption, + modelPickerAutoOption, + primaryModelOption, + selectAuto, + selectModelKey, + selectedModelAuto, selectedModelKey, selectedModelOption, - selectedPromptModel: selectedModelOption + selectedPromptModel: !selectedModelAuto && selectedModelOption ? { providerID: selectedModelOption.providerID, modelID: selectedModelOption.modelID } : null, + setSelectedModelAuto, setSelectedModelKey, }; } @@ -417,16 +478,28 @@ export function ChatModelPicker({ loading, selectedModelOption, onSelectModel, + autoOption, }: { groupedOptions: ChatModelProviderGroup[]; loading: boolean; selectedModelOption: ChatModelOption | null; onSelectModel: (option: ChatModelOption) => void; + autoOption?: ChatModelPickerAutoOption; }) { const { t } = useTranslation('session'); const [open, setOpen] = useState(false); + const selectorRef = useRef(null); + const [menuLeftOffset, setMenuLeftOffset] = useState(0); const { tooltip, showTooltip, hideTooltip } = useSelectorTooltip(); const hasOptions = groupedOptions.some((group) => group.models.length > 0); + const updateMenuLeftOffset = useCallback(() => { + const selector = selectorRef.current; + if (!selector) return; + setMenuLeftOffset(getAnchoredMenuLeftOffset( + selector.getBoundingClientRect().left, + window.innerWidth, + )); + }, []); useEffect(() => { if (!open) return; @@ -442,23 +515,44 @@ export function ChatModelPicker({ if (!open) hideTooltip(); }, [hideTooltip, open]); + useEffect(() => { + if (!open) return; + updateMenuLeftOffset(); + window.addEventListener('resize', updateMenuLeftOffset); + return () => window.removeEventListener('resize', updateMenuLeftOffset); + }, [open, updateMenuLeftOffset]); + return ( -
+
{open && ( -
+
{t('modelPicker.title')}
{t('modelPicker.hint')}
@@ -466,60 +560,98 @@ export function ChatModelPicker({
{loading ? (
{t('loading')}
- ) : groupedOptions.length > 0 ? ( - groupedOptions.map((group) => ( -
-
- {group.providerName} - - {t('modelPicker.count', { count: group.models.length })} - + ) : ( + <> + {autoOption && ( +
+
-
- {group.models.map((option) => ( -
- - ))} + + ))} +
-
- )) - ) : ( -
{t('modelPicker.empty')}
+ )) : ( +
{t('modelPicker.empty')}
+ )} + )}
diff --git a/webui/src/components/common/DelegateDetailSheet.test.tsx b/webui/src/components/common/DelegateDetailSheet.test.tsx new file mode 100644 index 000000000..4074768b5 --- /dev/null +++ b/webui/src/components/common/DelegateDetailSheet.test.tsx @@ -0,0 +1,51 @@ +import { render } from '@testing-library/react'; +import { describe, expect, it, vi } from 'vitest'; + +import DelegateDetailSheet from './DelegateDetailSheet'; + +const { sessionChatPropsMock } = vi.hoisted(() => ({ + sessionChatPropsMock: vi.fn(), +})); + +vi.mock('react-i18next', () => ({ + useTranslation: () => ({ + t: (key: string) => key, + }), +})); + +vi.mock('./SessionChat', () => ({ + default: (props: Record) => { + sessionChatPropsMock(props); + return null; + }, +})); + +describe('DelegateDetailSheet', () => { + it('uses the session-management process timeline inside the child conversation', () => { + render( + , + ); + + expect(sessionChatPropsMock).toHaveBeenCalledWith(expect.objectContaining({ + sessionId: 'ses-child', + live: true, + hideInput: true, + agentName: 'Librarian', + display: { + compact: false, + pageCanvas: true, + showTimestamp: true, + collapseIntermediateSteps: true, + processGroupsDefaultOpen: false, + processGroupsOpenWhileActive: true, + }, + })); + }); +}); diff --git a/webui/src/components/common/DelegateDetailSheet.tsx b/webui/src/components/common/DelegateDetailSheet.tsx index 420c981b3..274eb6922 100644 --- a/webui/src/components/common/DelegateDetailSheet.tsx +++ b/webui/src/components/common/DelegateDetailSheet.tsx @@ -178,7 +178,15 @@ export default function DelegateDetailSheet({ sessionId={sessionId} live hideInput - display={{ compact: false, showTimestamp: true }} + display={{ + compact: false, + pageCanvas: true, + showTimestamp: true, + collapseIntermediateSteps: true, + processGroupsDefaultOpen: false, + processGroupsOpenWhileActive: true, + }} + agentName={agentName} className="h-full" emptyText={t('delegate.emptyChat')} /> diff --git a/webui/src/components/common/DelegateTaskCard.test.tsx b/webui/src/components/common/DelegateTaskCard.test.tsx index a50dbbefb..3db628e49 100644 --- a/webui/src/components/common/DelegateTaskCard.test.tsx +++ b/webui/src/components/common/DelegateTaskCard.test.tsx @@ -1,8 +1,48 @@ -import { describe, expect, it } from 'vitest'; +import { fireEvent, render, screen } from '@testing-library/react'; +import { describe, expect, it, vi } from 'vitest'; -import { shouldRenderDelegateTaskCard } from './DelegateTaskCard'; +import DelegateTaskCard, { shouldRenderDelegateTaskCard } from './DelegateTaskCard'; import type { MessagePart } from '../../types'; +vi.mock('react-i18next', () => ({ + useTranslation: () => ({ + t: (key: string, options?: Record) => { + const translations: Record = { + 'delegate.delegatedTo': '委派给 {{agent}}', + 'delegate.pending': '等待中', + 'delegate.running': '执行中', + 'delegate.completed': '已完成', + 'delegate.background': '后台', + 'delegate.elapsedRunning': '已运行', + 'delegate.elapsedDone': '耗时', + 'delegate.steps': '步', + 'delegate.resultSummary': '结果摘要', + 'delegate.viewExecution': '查看执行过程', + 'delegate.subTask': '子任务', + }; + return (translations[key] ?? key).replace( + /\{\{(\w+)\}\}/g, + (_, name: string) => String(options?.[name] ?? ''), + ); + }, + }), +})); + +vi.mock('./DelegateDetailSheet', async () => { + const ReactModule = await import('react'); + return { + default: ({ + open, + sessionId, + }: { + open: boolean; + sessionId: string; + }) => open + ? ReactModule.createElement('div', { 'data-testid': 'delegate-detail-sheet' }, sessionId) + : null, + }; +}); + describe('shouldRenderDelegateTaskCard', () => { it('does not treat generic tool category fields as delegate tasks', () => { const part = { @@ -84,3 +124,48 @@ describe('shouldRenderDelegateTaskCard', () => { expect(shouldRenderDelegateTaskCard(part)).toBe(false); }); }); + +describe('DelegateTaskCard process step', () => { + it('shows delegated agent context and opens the child execution sheet', () => { + const part = { + id: 'part-delegate', + type: 'tool', + tool: 'delegate_task', + state: { + status: 'running', + input: { + description: '调研 OpenClaw 最新版本', + prompt: '调研 GitHub 上的最新发布', + subagent_type: 'librarian', + run_in_background: true, + }, + metadata: { + sessionId: 'ses-child', + status: 'running', + background: true, + steps: [ + { + tool: 'websearch', + title: '搜索最新版本', + status: 'running', + }, + ], + stepCount: 1, + }, + }, + } as MessagePart; + + render(); + + const processStep = screen.getByTestId('chat-process-delegate-step'); + expect(processStep).toHaveTextContent('委派给 Librarian'); + expect(processStep).toHaveTextContent('调研 OpenClaw 最新版本'); + expect(processStep).toHaveTextContent('执行中'); + expect(processStep).toHaveTextContent('后台'); + expect(processStep.querySelector('summary')).toHaveClass('text-sm'); + + fireEvent.click(screen.getByRole('button', { name: /查看执行过程/ })); + + expect(screen.getByTestId('delegate-detail-sheet')).toHaveTextContent('ses-child'); + }); +}); diff --git a/webui/src/components/common/DelegateTaskCard.tsx b/webui/src/components/common/DelegateTaskCard.tsx index 9110fb61f..1b0a28278 100644 --- a/webui/src/components/common/DelegateTaskCard.tsx +++ b/webui/src/components/common/DelegateTaskCard.tsx @@ -6,7 +6,7 @@ */ import { useState, useEffect, useRef } from 'react'; -import { ChevronRight, ExternalLink, XCircle } from 'lucide-react'; +import { Bot, ChevronDown, ChevronRight, ExternalLink, XCircle } from 'lucide-react'; import { useTranslation } from 'react-i18next'; import type { MessagePart, ToolState } from '@/types'; import DelegateDetailSheet from './DelegateDetailSheet'; @@ -211,9 +211,10 @@ const STATUS_STYLE: Record = {}; let info: DelegateInfo; @@ -253,6 +254,123 @@ export default function DelegateTaskCard({ part }: DelegateTaskCardProps) { setElapsed(info.durationMs); }, [info.status, state.time?.start, state.time?.end, info.durationMs]); + if (processStep) { + return ( + <> +
+ + + + + + {t('delegate.delegatedTo', { agent: info.agentName })} + + + {info.description} + + {info.isBackground && ( + + {t('delegate.background')} + + )} + + {cfg.pulse && ( + + )} + {t(`delegate.${info.status}`, { defaultValue: info.status })} + + + + +
+ {elapsed !== null && elapsed > 0 && ( +
+ {info.status === 'running' ? t('delegate.elapsedRunning') : t('delegate.elapsedDone')} + {' '} + {formatDuration(elapsed)} + {info.stepCount > 0 && ` · ${info.stepCount} ${t('delegate.steps')}`} +
+ )} + + {info.status === 'running' && info.steps.length > 0 && ( +
+ {info.steps.map((step, index) => ( +
+ + {step.status === 'completed' ? '✓' : step.status === 'error' ? '✗' : '◌'} + + + {step.tool} + + {step.title} +
+ ))} + {info.currentText && ( +
+ ⋯ {info.currentText.slice(-80)} +
+ )} +
+ )} + + {info.status === 'error' && info.error && ( +
+ + {info.error} +
+ )} + + {info.status === 'completed' && info.output && ( +
+ + {t('delegate.resultSummary')} + +
+ {truncateOutput(info.output)} +
+
+ )} + + +
+
+ + {info.childSessionId && ( + setSheetOpen(false)} + sessionId={info.childSessionId} + agentName={info.agentName} + description={info.description} + status={info.status} + /> + )} + + ); + } + return ( <>
diff --git a/webui/src/components/common/EntitySheet.test.tsx b/webui/src/components/common/EntitySheet.test.tsx index ffde9c541..a8a85b394 100644 --- a/webui/src/components/common/EntitySheet.test.tsx +++ b/webui/src/components/common/EntitySheet.test.tsx @@ -200,6 +200,48 @@ describe('EntitySheet', () => { ); }); + it('propagates manually selected Auto mode to session creation and sending', async () => { + const createAndSend = vi.fn().mockResolvedValue('rex-session-auto'); + mockUseSessionChat.mockReturnValue({ + sessionId: null, + loading: false, + error: null, + create: vi.fn().mockResolvedValue(undefined), + createAndSend, + retry: vi.fn().mockResolvedValue(undefined), + reset: vi.fn(), + }); + + render( + +
Form content
+
, + ); + + expect(mockUseSessionChat).toHaveBeenCalledWith(expect.objectContaining({ + category: 'entity-config', + modelAuto: true, + })); + const sessionChatProps = vi.mocked(SessionChat).mock.calls.at(-1)?.[0] as any; + expect(sessionChatProps).toEqual(expect.objectContaining({ + model: null, + modelAuto: true, + })); + + await sessionChatProps.onCreateAndSend('hello', [], undefined, undefined); + expect(createAndSend).toHaveBeenCalledWith(expect.objectContaining({ + text: 'hello', + model: null, + modelAuto: true, + })); + }); + it('renders extract from Rex as a guide action instead of a standalone footer action', async () => { const user = userEvent.setup(); const onExtractFromRex = vi.fn().mockResolvedValue(undefined); @@ -327,6 +369,30 @@ describe('EntitySheet', () => { }); }); + it('restores Auto selection from a persisted Rex session', async () => { + const onRexModelAutoChange = vi.fn(); + window.localStorage.setItem( + 'flocks:entity-sheet:rex-session:v1:agent-edit:auto-agent', + 'persisted-auto-session', + ); + mockClientGet.mockResolvedValueOnce({ data: { model_auto: true } }); + + render( + +
Form content
+
, + ); + + await waitFor(() => { + expect(onRexModelAutoChange).toHaveBeenCalledWith(true); + }); + }); + it('clears a stored Rex session when validation reports it missing', async () => { window.localStorage.setItem( 'flocks:entity-sheet:rex-session:v1:agent-edit:audit-agent', diff --git a/webui/src/components/common/EntitySheet.tsx b/webui/src/components/common/EntitySheet.tsx index 37e39b752..355f73ff0 100644 --- a/webui/src/components/common/EntitySheet.tsx +++ b/webui/src/components/common/EntitySheet.tsx @@ -164,6 +164,8 @@ export interface EntitySheetProps { rexAgentName?: string; rexMentionAgents?: Agent[]; rexModel?: { providerID: string; modelID: string } | null; + rexModelAuto?: boolean; + onRexModelAutoChange?: (enabled: boolean) => void; rexSupportsVision?: boolean | null; rexContextWindowTokens?: number | null; /** Persist and resume the Rex conversation across refreshes when provided. */ @@ -212,6 +214,8 @@ export default function EntitySheet({ rexAgentName, rexMentionAgents, rexModel, + rexModelAuto = false, + onRexModelAutoChange, rexSupportsVision, rexContextWindowTokens, rexSessionStorageKey, @@ -267,6 +271,7 @@ export default function EntitySheet({ } = useSessionChat({ title: `${title} — ${t('entity.rexAssist')}`, category: 'entity-config', + modelAuto: rexModelAuto, contextMessage: rexSystemContext, welcomeMessage: rexWelcomeMessage, initialSessionId: storedRexSessionId, @@ -288,8 +293,9 @@ export default function EntitySheet({ (async () => { try { - await client.get(`/api/session/${stored}`); + const response = await client.get(`/api/session/${stored}`); if (cancelled) return; + onRexModelAutoChange?.(Boolean(response.data?.model_auto)); setStoredRexSessionId(stored); } catch { if (cancelled) return; @@ -305,7 +311,7 @@ export default function EntitySheet({ return () => { cancelled = true; }; - }, [rexSessionStorageKey]); + }, [onRexModelAutoChange, rexSessionStorageKey]); useEffect(() => { if (!rexSessionHydrated || !sessionId) return; @@ -374,6 +380,7 @@ export default function EntitySheet({ useEffect(() => { if (!open) { + onRexModelAutoChange?.(false); setActiveTab(getDefaultTab()); if (!rexSessionStorageKey) { resetRexSession(); @@ -387,7 +394,7 @@ export default function EntitySheet({ setTestPrompt(effectiveDefaultTestPrompt); setDrawerWidth(resolvedInitialWidth()); } - }, [open, mode, defaultTestPrompt, resetRexSession, initialWidth, showTabs, hideRex, hideForm, initialTab, rexSessionStorageKey]); + }, [open, mode, defaultTestPrompt, resetRexSession, initialWidth, showTabs, hideRex, hideForm, initialTab, onRexModelAutoChange, rexSessionStorageKey]); // ── Tab handling ────────────────────────────────────────────────────────── @@ -432,18 +439,15 @@ export default function EntitySheet({ const openRex = useCallback( (msg?: string) => { setActiveTab('rex'); - if (activeRexSessionId && msg) { - const payload: Record = { - parts: [{ type: 'text', text: msg }], - }; - if (rexAgentName) payload.agent = rexAgentName; - if (rexModel) payload.model = rexModel; - client.post(`/api/session/${activeRexSessionId}/prompt_async`, payload); - } else if (msg) { - createAndSendRex({ text: msg, agent: rexAgentName, model: rexModel }).catch(() => {}); - } + if (!msg) return; + createAndSendRex({ + text: msg, + agent: rexAgentName, + model: rexModel, + modelAuto: rexModelAuto, + }).catch(() => {}); }, - [activeRexSessionId, createAndSendRex, rexAgentName, rexModel], + [createAndSendRex, rexAgentName, rexModel, rexModelAuto], ); // ── openTest (exposed via context) ──────────────────────────────────────── @@ -499,9 +503,10 @@ export default function EntitySheet({ text: prompt, agent: rexAgentName, model: rexModel, + modelAuto: rexModelAuto, displayText: buildInstructionDisplayText(label), }).catch(() => {}); - }, [createAndSendRex, handleExtract, rexAgentName, rexModel]); + }, [createAndSendRex, handleExtract, rexAgentName, rexModel, rexModelAuto]); if (!open) return null; @@ -698,6 +703,7 @@ export default function EntitySheet({ agentName={rexAgentName} mentionAgents={rexMentionAgents} model={rexModel} + modelAuto={rexModelAuto} supportsVision={rexSupportsVision ?? supportsVision} contextWindowTokens={rexContextWindowTokens} toolbarSlot={rexToolbarSlot} @@ -709,6 +715,7 @@ export default function EntitySheet({ imageParts, agent: agentOverride || rexAgentName, model: modelOverride === undefined ? rexModel : modelOverride, + modelAuto: rexModelAuto, displayText: options?.displayText, }) : undefined} welcomeContent={( diff --git a/webui/src/components/common/QuestionTool.test.tsx b/webui/src/components/common/QuestionTool.test.tsx index 5399e521b..4be6c0cc6 100644 --- a/webui/src/components/common/QuestionTool.test.tsx +++ b/webui/src/components/common/QuestionTool.test.tsx @@ -195,6 +195,37 @@ describe('QuestionTool', () => { expect(screen.queryByRole('button', { name: /自定义 \/ 补充说明/ })).not.toBeInTheDocument(); }); + it('accepts feedback through the continue-planning option', async () => { + const user = userEvent.setup(); + const onAnswer = vi.fn().mockResolvedValue(undefined); + + render( + , + ); + + await user.click(screen.getByRole('button', { name: /调整计划/ })); + await user.type(screen.getByRole('textbox'), 'Keep the public API unchanged.'); + await user.click(screen.getByRole('button', { name: /确认/ })); + + expect(onAnswer).toHaveBeenCalledWith([ + ['调整计划', 'Keep the public API unchanged.'], + ]); + }); + it('falls back to text input when a choice question has no visible options', async () => { const user = userEvent.setup(); const onAnswer = vi.fn().mockResolvedValue(undefined); diff --git a/webui/src/components/common/QuestionTool.tsx b/webui/src/components/common/QuestionTool.tsx index 99ac8f631..79117e314 100644 --- a/webui/src/components/common/QuestionTool.tsx +++ b/webui/src/components/common/QuestionTool.tsx @@ -23,6 +23,8 @@ export type QuestionType = 'choice' | 'text' | 'number' | 'file' | 'confirm' | ' export interface QuestionOption { label?: string; description?: string; + /** Show a text input and submit both the option label and entered text. */ + allowText?: boolean; [key: string]: unknown; } @@ -95,10 +97,14 @@ function optionDescription(opt: QuestionOption | string): string { return ''; } +function optionAllowsText(opt: QuestionOption | string): boolean { + return typeof opt !== 'string' && opt.allowText === true; +} + const CUSTOM_CHOICE_PREFIX = '__flocks_custom_choice__:'; function isCustomChoiceLabel(label: string): boolean { - return /^(其他|其它|自定义|补充)|\b(other|custom)\b|请补充|补充说明|type your answer/i.test(label.trim()); + return /^(其他|其它|自定义|补充)|\b(other|custom|feedback)\b|请补充|补充说明|type your answer/i.test(label.trim()); } function customChoiceValue(text: string): string { @@ -187,22 +193,26 @@ function ChoiceInput({ .map(opt => ({ label: optionLabel(opt), description: optionDescription(opt), + allowText: optionAllowsText(opt), custom: false, })) .filter(opt => opt.label); - const hasProvidedCustomOption = visibleOptions.some(opt => isCustomChoiceLabel(opt.label)); + const hasProvidedCustomOption = visibleOptions.some( + opt => opt.allowText || isCustomChoiceLabel(opt.label), + ); const options = shouldOfferCustomChoice(q) && !hasProvidedCustomOption ? [ ...visibleOptions, { label: t('question.customAnswer'), description: t('question.textPlaceholder'), + allowText: false, custom: true, }, ] : visibleOptions.map(opt => ({ ...opt, - custom: isCustomChoiceLabel(opt.label), + custom: opt.allowText || isCustomChoiceLabel(opt.label), })); const customSelected = hasCustomChoice(answer); const customText = customChoiceText(answer); @@ -213,7 +223,15 @@ function ChoiceInput({ onChange([label]); } }; - const toggleCustom = () => { + const toggleCustom = (label: string, preserveLabel: boolean) => { + if (preserveLabel) { + onChange( + customSelected && answer.includes(label) + ? [] + : [label, customChoiceValue(customText)], + ); + return; + } if (multiple) { if (customSelected) { onChange(answer.filter(value => !isCustomChoiceValue(value))); @@ -224,8 +242,12 @@ function ChoiceInput({ } onChange(customSelected ? [] : [customChoiceValue(customText)]); }; - const setCustomText = (text: string) => { + const setCustomText = (label: string, preserveLabel: boolean, text: string) => { const nextCustom = customChoiceValue(text); + if (preserveLabel) { + onChange([label, nextCustom]); + return; + } if (multiple) { const withoutCustom = answer.filter(value => !isCustomChoiceValue(value)); onChange([...withoutCustom, nextCustom]); @@ -244,11 +266,17 @@ function ChoiceInput({ {options.map(opt => { const label = opt.label; const desc = opt.description; - const selected = opt.custom ? customSelected : answer.includes(label); + const selected = opt.custom + ? customSelected && (!opt.allowText || answer.includes(label)) + : answer.includes(label); return (
{open && ( @@ -817,27 +847,21 @@ export function getMessageBubbleClassName({ isUser: boolean; isEditing: boolean; }): string { + if (!isUser) { + const typographyClass = compact ? 'text-sm' : 'text-[15px]'; + + return `w-full max-w-full min-w-0 bg-transparent py-1 ${typographyClass} text-[#34393e] break-words dark:text-zinc-100`; + } + if (compact) { - const widthClass = isUser - ? (isEditing ? 'w-full max-w-full' : 'max-w-full') - : 'w-full max-w-full'; - - return `${widthClass} min-w-0 px-4 py-3 rounded-[20px] text-sm break-words shadow-sm ${ - isUser - ? 'bg-sky-50 border border-sky-100 text-zinc-900 dark:border-zinc-700 dark:bg-zinc-800 dark:text-zinc-50 dark:shadow-none' - : 'bg-white border border-zinc-200/90 dark:border-zinc-800 dark:bg-zinc-900 dark:text-zinc-100 dark:shadow-none' - }`; + const widthClass = isEditing ? 'w-full max-w-full' : 'max-w-full'; + + return `${widthClass} min-w-0 px-4 py-3 rounded-[20px] text-sm break-words shadow-sm border border-black/[0.07] bg-zinc-50 text-[#30343a] dark:border-white/[0.08] dark:bg-[#303842] dark:text-zinc-50 dark:shadow-none`; } - const widthClass = isUser - ? (isEditing ? 'w-full' : 'w-auto') - : 'w-full'; + const widthClass = isEditing ? 'w-full' : 'w-auto'; - return `${widthClass} min-w-0 max-w-full px-5 py-4 rounded-[24px] text-sm break-words shadow-sm ${ - isUser - ? 'bg-sky-50 border border-sky-100 text-zinc-900 dark:border-zinc-700 dark:bg-zinc-800 dark:text-zinc-50 dark:shadow-none' - : 'bg-white border border-zinc-200/90 dark:border-zinc-800 dark:bg-zinc-900 dark:text-zinc-100 dark:shadow-none' - }`; + return `${widthClass} min-w-0 max-w-full px-5 py-3 rounded-[18px] text-sm break-words shadow-sm border border-black/[0.09] bg-zinc-50 text-[#30343a] dark:border-white/[0.10] dark:bg-[#303842] dark:text-zinc-50 dark:shadow-none`; } export function getInstructionDisplayBubbleClassName(compact: boolean): string { @@ -901,6 +925,54 @@ export function hasActiveToolPart(parts?: Array= 0; index -= 1) { + if (messages[index]?.info?.role === 'user') { + latestUserIndex = index; + break; + } + } + + let turnParentID = latestUserIndex >= 0 + ? messages[latestUserIndex]?.info?.id + : undefined; + let turnStartIndex = latestUserIndex + 1; + + // A single tool-heavy turn can exceed the latest-message page. If its user + // message is outside the page, recover the turn from the newest assistant's + // parent instead of falling back to every historical tool in the page. + if (!turnParentID) { + for (let index = messages.length - 1; index >= 0; index -= 1) { + const info = messages[index]?.info; + if (info?.role === 'assistant' && info.parentID) { + turnParentID = info.parentID; + turnStartIndex = 0; + break; + } + } + } + + if (!turnParentID) return null; + return messages.slice(turnStartIndex).filter((message) => ( + message.info?.role === 'assistant' + && message.info.parentID === turnParentID + )); +} + export function isActiveSessionStatus(status?: { type?: string } | null): boolean { return status?.type === 'busy' || status?.type === 'compacting' || status?.type === 'retry'; } @@ -1417,15 +1489,19 @@ function findMentionTrigger(text: string, cursor: number): { start: number; end: }; } -function resolveMentionAgentName(text: string, agents: Agent[]): string | null { +function resolveReferencedAgentName(text: string, agents: Agent[]): string | null { const sorted = [...agents].sort((a, b) => b.name.length - a.name.length); for (const agent of sorted) { - const pattern = new RegExp(`(^|\\s)@${escapeRegExp(agent.name)}(?=$|\\s|[,.!?;:,。!?;:])`, 'i'); + const pattern = new RegExp(`(^|\\s)subagent:${escapeRegExp(agent.name)}(?=$|\\s|[,.!?;,。!?;])`, 'i'); if (pattern.test(text)) return agent.name; } return null; } +function formatComposerReference(reference: ComposerReference): string { + return `${reference.kind}:${reference.value}`; +} + export default function SessionChat({ sessionId, live = false, @@ -1441,6 +1517,9 @@ export default function SessionChat({ initialDisplayText, agentName, model, + executionMode = 'build', + onExecutionModeAccepted, + modelAuto = false, display, welcomeContent, conversationBottomSlot, @@ -1452,6 +1531,8 @@ export default function SessionChat({ onInitialMessageConsumed, supportsVision, toolbarSlot, + composerAddMenuSlot, + onComposerAddMenuOpenChange, composerTextareaMinHeight, composerTextareaMaxHeight, centerToolbarSlot, @@ -1462,6 +1543,7 @@ export default function SessionChat({ const toast = useToast(); const compact = display?.compact ?? true; const fullWidth = display?.fullWidth ?? false; + const pageCanvas = display?.pageCanvas ?? false; const showActions = display?.showActions ?? false; const showTimestamp = display?.showTimestamp ?? false; const collapseIntermediateSteps = display?.collapseIntermediateSteps ?? false; @@ -1471,15 +1553,29 @@ export default function SessionChat({ const effectiveComposerTextareaMaxHeight = composerTextareaMaxHeight ?? (compact ? 96 : 200); const effectivePlaceholder = placeholder ?? t('chat.placeholder'); const effectiveEmptyText = emptyText ?? t('chat.emptyText'); + const autoModelSessionRef = useRef(null); + useEffect(() => { + if (!modelAuto) autoModelSessionRef.current = null; + }, [modelAuto]); + const ensureAutoModelSession = useCallback(async () => { + if (!sessionId || !modelAuto || autoModelSessionRef.current === sessionId) return; + await sessionApi.update(sessionId, { + model_auto: true, + model_pinned: false, + }); + autoModelSessionRef.current = sessionId; + }, [modelAuto, sessionId]); // Restore any persisted draft on first mount so navigating away (e.g. // sidebar → Agents → back to Sessions) doesn't wipe the user's half-typed // message. Subsequent session changes are re-hydrated by the effect below. const [input, setInput] = useState(() => readChatDraft(sessionId)); + const [composerReferences, setComposerReferences] = useState([]); const [sending, setSending] = useState(false); const [isStreaming, setIsStreaming] = useState(false); const activeToolPartIdsRef = useRef>(new Set()); const [attachments, setAttachments] = useState([]); const [isDragOver, setIsDragOver] = useState(false); + const [showComposerAddMenu, setShowComposerAddMenu] = useState(false); // Lightbox preview for composer thumbnails. Shares the same overlay // component used by message bubbles so the click-to-enlarge gesture is // consistent across the upload tray and the rendered chat history. @@ -1581,6 +1677,7 @@ export default function SessionChat({ const initialMessageSentRef = useRef(''); const abortingRef = useRef(false); const sessionBusyRef = useRef(false); + const sessionStatusRevisionRef = useRef(0); const goalHydrationVersionRef = useRef(0); // ID of the assistant message that was aborted; used to ignore its finish event const abortedMessageIdRef = useRef(null); @@ -1604,6 +1701,45 @@ export default function SessionChat({ const fileInputRef = useRef(null); const isComposingRef = useRef(false); + const closeComposerAddMenu = useCallback(() => { + setShowComposerAddMenu(false); + onComposerAddMenuOpenChange?.(false); + }, [onComposerAddMenuOpenChange]); + + const openComposerAddMenu = useCallback(() => { + setShowComposerAddMenu(true); + onComposerAddMenuOpenChange?.(true); + }, [onComposerAddMenuOpenChange]); + + const toggleComposerAddMenu = useCallback(() => { + setShowComposerAddMenu((open) => { + const nextOpen = !open; + onComposerAddMenuOpenChange?.(nextOpen); + return nextOpen; + }); + }, [onComposerAddMenuOpenChange]); + + useEffect(() => { + if (!showComposerAddMenu) return; + const handlePointerDown = (event: MouseEvent) => { + const target = event.target as HTMLElement; + if (!target.closest('[data-composer-add-menu]')) closeComposerAddMenu(); + }; + const handleKeyDown = (event: KeyboardEvent) => { + if (event.key === 'Escape') closeComposerAddMenu(); + }; + document.addEventListener('mousedown', handlePointerDown); + document.addEventListener('keydown', handleKeyDown); + return () => { + document.removeEventListener('mousedown', handlePointerDown); + document.removeEventListener('keydown', handleKeyDown); + }; + }, [closeComposerAddMenu, showComposerAddMenu]); + + useEffect(() => { + if (sending && showComposerAddMenu) closeComposerAddMenu(); + }, [closeComposerAddMenu, sending, showComposerAddMenu]); + // Slash command autocomplete state const [commands, setCommands] = useState([]); const [showCommandDropdown, setShowCommandDropdown] = useState(false); @@ -1630,7 +1766,12 @@ export default function SessionChat({ ); const hasUploadingFiles = attachments.some((attachment) => attachment.status === 'uploading'); const canSend = !sending && !hasUploadingFiles && - (!!input.trim() || successfulDocAttachments.length > 0 || successfulImageAttachments.length > 0); + ( + !!input.trim() + || composerReferences.length > 0 + || successfulDocAttachments.length > 0 + || successfulImageAttachments.length > 0 + ); const filteredMentionAgents = useMemo(() => { const q = mentionQuery.trim().toLowerCase(); return mentionAgents @@ -1654,9 +1795,6 @@ export default function SessionChat({ } }, []); - const loadOlderMessagesRef = useRef<(() => Promise) | null>(null); - const hasMoreMessagesRef = useRef(false); - const loadingOlderMessagesRef = useRef(false); const rafScheduledRef = useRef(false); const handleScroll = useCallback(() => { if (rafScheduledRef.current) return; @@ -1665,18 +1803,6 @@ export default function SessionChat({ const el = scrollContainerRef.current; if (el) { isAtBottomRef.current = el.scrollHeight - el.scrollTop - el.clientHeight <= SCROLL_BOTTOM_THRESHOLD_PX; - if (el.scrollTop <= 80 && hasMoreMessagesRef.current && !loadingOlderMessagesRef.current) { - const previousHeight = el.scrollHeight; - const previousTop = el.scrollTop; - const loadPromise = loadOlderMessagesRef.current?.(); - if (loadPromise) void loadPromise.finally(() => { - requestAnimationFrame(() => { - const current = scrollContainerRef.current; - if (!current) return; - current.scrollTop = current.scrollHeight - previousHeight + previousTop; - }); - }); - } } rafScheduledRef.current = false; }); @@ -1685,21 +1811,18 @@ export default function SessionChat({ const { messages, loading, - loadingOlder, - hasMore: hasMoreMessages, + error: messagesError, refetch, - loadOlder, addMessage, updateMessage, updateMessagePart, + removeMessage, + clearMessages, replaceMessageText, markMessageStopped, truncateAfterMessage, } = useSessionMessages(sessionId || undefined); - useEffect(() => { loadOlderMessagesRef.current = loadOlder; }, [loadOlder]); - useEffect(() => { hasMoreMessagesRef.current = hasMoreMessages; }, [hasMoreMessages]); - useEffect(() => { loadingOlderMessagesRef.current = loadingOlder; }, [loadingOlder]); const contextUsageMessages = contextUsageRefreshing && !contextUsageSnapshot ? [] : messages; const contextUsageBreakdown = useMemo( () => buildContextUsageBreakdown(contextUsageMessages, input, contextUsageSnapshot), @@ -1725,6 +1848,8 @@ export default function SessionChat({ // Keep a ref to latest messages so handleAbort can read it without stale closure const messagesRef = useRef(messages); useEffect(() => { messagesRef.current = messages; }, [messages]); + const pendingQuestionsRef = useRef(pendingQuestions); + useEffect(() => { pendingQuestionsRef.current = pendingQuestions; }, [pendingQuestions]); const hasUserMessage = useMemo(() => messages.some((m) => m.role === 'user'), [messages]); @@ -1767,6 +1892,7 @@ export default function SessionChat({ case 'ignore': return; case 'session-cleared': + sessionStatusRevisionRef.current += 1; abortingRef.current = false; sessionBusyRef.current = false; activeToolPartIdsRef.current.clear(); @@ -1775,10 +1901,11 @@ export default function SessionChat({ setIsStreaming(false); setGoalBanner(null); setDismissedGoalKey(''); - refetch(); + clearMessages(); void refreshContextUsage({ clear: true }); return; case 'session-status': + sessionStatusRevisionRef.current += 1; if (action.statusType === 'busy') { sessionBusyRef.current = true; if ( @@ -1854,6 +1981,17 @@ export default function SessionChat({ } return; } + case 'message-removed': { + const removedMessage = messagesRef.current.find((message) => message.id === action.messageID); + removedMessage?.parts.forEach((part) => { + if (part.id) activeToolPartIdsRef.current.delete(part.id); + }); + if (abortedMessageIdRef.current === action.messageID) { + abortedMessageIdRef.current = null; + } + removeMessage(action.messageID); + return; + } case 'message-part-updated': { const part = action.part as Pick; if (part.id) { @@ -1931,6 +2069,8 @@ export default function SessionChat({ sessionId, updateMessage, updateMessagePart, + removeMessage, + clearMessages, refetch, refreshContextUsage, applyContextUsagePushSnapshot, @@ -1977,11 +2117,52 @@ export default function SessionChat({ [onError, submitReject, t, toast], ); + const reconcileSessionStatusAfterReconnect = useCallback(async () => { + if (!sessionId) return; + const statusRevision = sessionStatusRevisionRef.current; + + try { + const response = await client.get('/api/session/status'); + if (statusRevision !== sessionStatusRevisionRef.current) return; + + const status = response.data?.[sessionId]; + if (isActiveSessionStatus(status)) { + sessionBusyRef.current = true; + if (!abortingRef.current && !suppressStreamingUntilIdleRef.current) { + setIsStreaming(true); + } + if (status?.type === 'compacting') { + setIsCompacting(true); + isCompactingRef.current = true; + setCompactingMessage(status.message || t('chat.compacting')); + } else { + setIsCompacting(false); + isCompactingRef.current = false; + setCompactingMessage(''); + setCompactionStages([]); + } + return; + } + + sessionBusyRef.current = false; + activeToolPartIdsRef.current.clear(); + setSending(false); + setIsStreaming(false); + setIsCompacting(false); + isCompactingRef.current = false; + setCompactingMessage(''); + setCompactionStages([]); + } catch { + // Keep the current activity state when reconnect status cannot be verified. + } + }, [sessionId, t]); + const { status: sseStatus } = useSSE({ url: `${getApiBase()}/api/event`, onEvent: handleSSEEvent, onReconnect: () => { if (!sessionId) return; + void reconcileSessionStatusAfterReconnect(); refetch(); refreshContextUsage(); fetchPromptQueue(); @@ -2047,11 +2228,11 @@ export default function SessionChat({ setMentionRange(null); setMentionQuery(''); setSelectedMentionIndex(0); - setPendingAgentName(agentName || 'rex'); abortingRef.current = false; abortedMessageIdRef.current = null; suppressStreamingUntilIdleRef.current = false; sessionBusyRef.current = false; + sessionStatusRevisionRef.current += 1; statusCheckedRef.current = null; isAtBottomRef.current = true; clearPendingQuestions(); @@ -2059,8 +2240,9 @@ export default function SessionChat({ // don't force a remount (Session/index.tsx does, but other consumers // such as WorkflowDetail/ChatTab may swap sessionId without a remount). setInput(readChatDraft(sessionId)); + setComposerReferences([]); setProcessGroupOpenState(readProcessGroupOpenState(sessionId)); - }, [sessionId, agentName, clearPendingQuestions]); + }, [sessionId, clearPendingQuestions]); const handleProcessGroupOpenChange = useCallback((key: string, open: boolean) => { setProcessGroupOpenState(prev => { @@ -2422,20 +2604,22 @@ export default function SessionChat({ setIsStreaming(true); const displayText = args ? `/${command} ${args}` : `/${command}`; - const tempId = `temp-${Date.now()}`; + const messageId = createMessageId(); addMessage({ - id: tempId, + id: messageId, sessionID: sessionId, role: 'user', - parts: [{ id: `${tempId}-part`, type: 'text', text: displayText }], + parts: [{ id: `temp-${messageId}-part`, type: 'text', text: displayText }], timestamp: Date.now(), } as Message); try { + await ensureAutoModelSession(); await client.post(`/api/session/${sessionId}/command`, { command, arguments: args, agent: agentName, + messageID: messageId, }); if (command === 'goal' && args.trim()) { goalHydrationVersionRef.current += 1; @@ -2445,6 +2629,7 @@ export default function SessionChat({ } } catch (err: unknown) { setIsStreaming(false); + removeMessage(messageId); const axiosErr = err as any; if (axiosErr?.response?.status === 404) { onError?.('Session not found. Please start a new session.'); @@ -2465,6 +2650,7 @@ export default function SessionChat({ options?: PromptDisplayOptions, ) => { if (!sessionId) return; + await ensureAutoModelSession(); const effectiveAgent = agentOverride || agentName; const visibleText = options?.displayText || text; // Clear abort state immediately so SSE events for the new stream are not suppressed @@ -2477,18 +2663,18 @@ export default function SessionChat({ setIsStreaming(true); setPendingAgentName(effectiveAgent || 'rex'); - const tempId = `temp-${Date.now()}`; + const messageId = createMessageId(); const tempParts: MessagePart[] = []; - if (visibleText) tempParts.push({ id: `${tempId}-text`, type: 'text', text: visibleText }); + if (visibleText) tempParts.push({ id: `temp-${messageId}-text`, type: 'text', text: visibleText }); imageParts.forEach((img, i) => { - tempParts.push({ id: `${tempId}-img-${i}`, type: 'file', url: img.url, mime: img.mime, filename: img.filename }); + tempParts.push({ id: `temp-${messageId}-img-${i}`, type: 'file', url: img.url, mime: img.mime, filename: img.filename }); }); addMessage({ - id: tempId, + id: messageId, sessionID: sessionId, role: 'user', - parts: tempParts.length > 0 ? tempParts : [{ id: `${tempId}-part`, type: 'text', text: visibleText }], + parts: tempParts.length > 0 ? tempParts : [{ id: `temp-${messageId}-part`, type: 'text', text: visibleText }], timestamp: Date.now(), agent: effectiveAgent, } as Message); @@ -2496,14 +2682,24 @@ export default function SessionChat({ try { const payload: Record = { parts: buildPromptParts(text, imageParts), + messageID: messageId, }; if (effectiveAgent) payload.agent = effectiveAgent; if (model) payload.model = model; if (options?.displayText) payload.displayText = options.displayText; + payload.executionMode = executionMode; await client.post(`/api/session/${sessionId}/prompt_async`, payload); + if (executionMode === 'goal' && text.trim()) { + goalHydrationVersionRef.current += 1; + writeDismissedGoalKey(sessionId, ''); + setGoalBanner({ objective: text.trim(), status: 'active' }); + setDismissedGoalKey(''); + } + onExecutionModeAccepted?.(executionMode); } catch (err: unknown) { setIsStreaming(false); + removeMessage(messageId); const axiosErr = err as any; if (axiosErr?.response?.status === 404) { onError?.(`Session not found. Please start a new session.`); @@ -2525,12 +2721,15 @@ export default function SessionChat({ if (!sessionId) return; const effectiveAgent = agentOverride || agentName; try { + await ensureAutoModelSession(); await enqueuePrompt({ parts: buildPromptParts(text, imageParts), ...(effectiveAgent ? { agent: effectiveAgent } : {}), ...(model ? { model } : {}), ...(options?.displayText ? { displayText: options.displayText } : {}), + executionMode, }); + onExecutionModeAccepted?.(executionMode); } catch (err: any) { const statusCode = err?.response?.status; const detail = err?.response?.data?.detail; @@ -2569,7 +2768,15 @@ export default function SessionChat({ setSending(true); try { setPendingAgentName(agentName || 'rex'); - await onCreateAndSend(trimmed, [], agentName, model, options); + await onCreateAndSend( + trimmed, + [], + agentName, + model, + options, + executionMode, + ); + onExecutionModeAccepted?.(executionMode); } catch { setInput(trimmed); } finally { @@ -2587,16 +2794,24 @@ export default function SessionChat({ const handleSend = async () => { if (!canSend) return; - const rawText = input.trim(); + const draftText = input.trim(); + const referencesToSend = [...composerReferences]; + const referenceText = referencesToSend.map(formatComposerReference).join(' '); + const rawText = [referenceText, draftText].filter(Boolean).join(' '); const docAttachmentsToSend = [...successfulDocAttachments]; const imageAttachmentsToSend = [...successfulImageAttachments]; const text = buildMessageText(rawText, docAttachmentsToSend); - const mentionedAgent = resolveMentionAgentName(rawText, mentionAgents); + const mentionedAgent = resolveReferencedAgentName(rawText, mentionAgents); + const restoreDraft = () => { + setInput(draftText); + setComposerReferences(referencesToSend); + }; // Need either text content or image attachments if (!text && imageAttachmentsToSend.length === 0) return; setInput(''); + setComposerReferences([]); setShowCommandDropdown(false); setMentionRange(null); @@ -2621,7 +2836,7 @@ export default function SessionChat({ await enqueueText(text, imageParts, mentionedAgent || undefined); setAttachments([]); } catch { - setInput(rawText); + restoreDraft(); setAttachments([...docAttachmentsToSend, ...imageAttachmentsToSend]); } return; @@ -2631,13 +2846,13 @@ export default function SessionChat({ if (parsed) { if (!sessionId) { // Slash commands need an existing session; restore input and do nothing - setInput(rawText); + restoreDraft(); return; } try { await sendCommand(parsed.command, parsed.args); } catch { - setInput(rawText); + restoreDraft(); } return; } @@ -2648,13 +2863,21 @@ export default function SessionChat({ try { const effectiveAgent = mentionedAgent || agentName; setPendingAgentName(effectiveAgent || 'rex'); - await onCreateAndSend(text, imageParts, effectiveAgent || undefined, model); + await onCreateAndSend( + text, + imageParts, + effectiveAgent || undefined, + model, + undefined, + executionMode, + ); + onExecutionModeAccepted?.(executionMode); setAttachments([]); } catch { // Restore both the text and the attachment list so the user can // retry without re-uploading images. Image data URLs are already // in memory, so restoring the array is safe and cheap. - setInput(rawText); + restoreDraft(); setAttachments(imageAttachmentsToSend); } finally { setSending(false); @@ -2667,7 +2890,7 @@ export default function SessionChat({ await sendText(text, imageParts, mentionedAgent || undefined); setAttachments([]); } catch { - setInput(rawText); + restoreDraft(); setAttachments(imageAttachmentsToSend); } }; @@ -2706,6 +2929,43 @@ export default function SessionChat({ }); }, [input, mentionRange]); + const insertAgentMention = useCallback((name: string) => { + setComposerReferences((current) => [ + { kind: 'subagent', value: name }, + ...current.filter((reference) => reference.kind !== 'subagent'), + ]); + setMentionRange(null); + setMentionQuery(''); + setSelectedMentionIndex(0); + requestAnimationFrame(() => { + textareaRef.current?.focus(); + }); + }, []); + + const insertComposerReference = useCallback(( + value: string, + kind: 'workflow' | 'skill', + ) => { + setComposerReferences((current) => { + if (kind === 'skill') { + if (current.some((reference) => reference.kind === 'skill' && reference.value === value)) { + return current; + } + return [{ kind, value }, ...current]; + } + return [ + { kind, value }, + ...current.filter((reference) => reference.kind !== kind), + ]; + }); + setMentionRange(null); + setMentionQuery(''); + setSelectedMentionIndex(0); + requestAnimationFrame(() => { + textareaRef.current?.focus(); + }); + }, []); + const handleKeyDown = (e: React.KeyboardEvent) => { const currentValue = e.currentTarget instanceof HTMLTextAreaElement ? e.currentTarget.value : input; const activeMention = mentionRange @@ -2864,24 +3124,81 @@ export default function SessionChat({ // Fallback polling to detect completion when SSE events are missed useEffect(() => { if (!isStreaming || !sessionId) return; + let questionRecoveryInFlight = false; const timer = setInterval(async () => { try { const res = await client.get(`/api/session/${sessionId}/message`, { params: { page: true, limit: 50, include_archived: true }, }); - const msgs: any[] = Array.isArray(res.data) ? res.data : (res.data?.items || []); + const msgs: FetchedMessageWithParts[] = Array.isArray(res.data) + ? res.data + : (res.data?.items || []); + const currentTurnAssistantMessages = getCurrentTurnAssistantMessages(msgs); + const shouldRecoverQuestionPart = (part: MessagePart) => ( + isQuestionToolName(part.tool || '') + && ( + isActiveToolPart(part) + || !!(part.callID && pendingQuestionsRef.current[part.callID]) + ) + ); + const hasFetchedPendingQuestion = currentTurnAssistantMessages?.some((message) => ( + (message.parts || []).some((part) => ( + shouldRecoverQuestionPart(part) + )) + )) ?? false; + if (currentTurnAssistantMessages && hasFetchedPendingQuestion) { + const localQuestionPartIds = new Set(); + const localQuestionCallIds = new Set(); + for (const message of messagesRef.current) { + for (const part of message.parts || []) { + if ( + (part.type !== 'tool' && part.type !== 'toolCall') + || !isQuestionToolName(part.tool || '') + ) continue; + if (part.id) localQuestionPartIds.add(part.id); + if (part.callID) localQuestionCallIds.add(part.callID); + } + } + const hasMissingFetchedQuestion = currentTurnAssistantMessages.some((msg) => ( + (msg.parts || []).some((part: MessagePart) => ( + shouldRecoverQuestionPart(part) + && !( + (part.id && localQuestionPartIds.has(part.id)) + || (part.callID && localQuestionCallIds.has(part.callID)) + ) + )) + )); + if (hasMissingFetchedQuestion && !questionRecoveryInFlight) { + questionRecoveryInFlight = true; + try { + await refetch(); + } finally { + questionRecoveryInFlight = false; + } + } + } + const lastMsg = msgs[msgs.length - 1]; if (lastMsg?.info?.role === 'assistant' && (lastMsg.info.finish || lastMsg.info.time?.completed)) { - const hasFetchedActiveTool = msgs.some((msg) => hasActiveToolPart(msg.parts)); - if (hasFetchedActiveTool) { - return; - } + const currentTurnMessages = currentTurnAssistantMessages + ? new Set(currentTurnAssistantMessages) + : null; + const hasFetchedActiveTool = msgs.some((msg) => ( + (msg.parts || []).some((part: MessagePart) => { + if (!isQuestionToolName(part.tool || '')) return isActiveToolPart(part); + const isPendingQuestion = isActiveToolPart(part) + || !!(part.callID && pendingQuestionsRef.current[part.callID]); + if (!isPendingQuestion) return false; + return currentTurnMessages === null || currentTurnMessages.has(msg); + }) + )); + if (hasFetchedActiveTool) return; + activeToolPartIdsRef.current.clear(); const statusRes = await client.get('/api/session/status'); const status = statusRes.data?.[sessionId]; - if (isActiveSessionStatus(status)) { - return; - } + if (isActiveSessionStatus(status)) return; + refetch(); setIsStreaming(false); } @@ -3050,11 +3367,17 @@ export default function SessionChat({ // ── Styling based on compact mode ── const msgAreaClass = compact ? 'relative flex flex-col flex-1 min-h-0 overflow-y-auto bg-gray-50 px-4 py-4 dark:bg-zinc-950' - : 'relative flex flex-col flex-1 min-h-0 overflow-y-auto bg-gray-50 py-6 dark:bg-zinc-950'; + : pageCanvas + ? 'relative flex flex-col flex-1 min-h-0 overflow-y-auto bg-transparent py-5' + : 'relative flex flex-col flex-1 min-h-0 overflow-y-auto bg-gray-50 py-6 dark:bg-zinc-950'; const msgListClass = compact ? fullWidth ? 'space-y-3 w-full px-4' : 'space-y-3' - : fullWidth ? 'space-y-5 w-full px-5' : 'space-y-5 w-[min(76%,64rem)] mx-auto px-6'; + : fullWidth + ? 'space-y-5 w-full px-12' + : pageCanvas + ? 'space-y-[18px] w-full max-w-[760px] mx-auto px-7' + : 'space-y-5 w-[min(76%,64rem)] mx-auto px-6'; const visibleGoalBanner = goalBanner && getGoalBannerKey(goalBanner) !== dismissedGoalKey ? goalBanner : null; @@ -3075,7 +3398,26 @@ export default function SessionChat({ > {loading && messages.length === 0 ? (
- + +
+ ) : messagesError && messages.length === 0 ? ( +
+
+ + {t('chat.loadFailed')} + + +
) : messages.length === 0 ? ( welcomeContent ? ( @@ -3093,19 +3435,6 @@ export default function SessionChat({ ) ) : (
- {hasMoreMessages && ( -
- -
- )}
- {formatAgentName(pendingAgentName)} + {formatAgentName(pendingAgentName)}
@@ -3228,7 +3557,7 @@ export default function SessionChat({
- {formatAgentName(pendingAgentName)} + {formatAgentName(pendingAgentName)}
@@ -3273,8 +3602,26 @@ export default function SessionChat({ {/* Follow-up input */} {!hideInput && ( -
-
+
+
{conversationBottomSlot && (
{typeof conversationBottomSlot === 'function' @@ -3348,14 +3695,16 @@ export default function SessionChat({ onDragOver={handleComposerDragOver} onDragLeave={handleComposerDragLeave} onDrop={handleComposerDrop} - className={`rounded-2xl border transition-all ${ + className={`${pageCanvas && !compact ? 'min-h-[124px] rounded-[20px] shadow-[0_3px_12px_rgba(22,27,34,0.045)]' : 'rounded-2xl'} border transition-all ${ isCompacting ? 'border-amber-200 bg-amber-50/30 dark:border-amber-500/35 dark:bg-amber-950/25' : isDragOver ? 'border-sky-300 bg-sky-50/60 ring-4 ring-sky-100 dark:border-sky-500/50 dark:bg-sky-950/35 dark:ring-sky-500/10' : isStreaming ? 'border-zinc-200 bg-zinc-50 dark:border-zinc-800 dark:bg-zinc-900/70' - : 'border-zinc-200 bg-zinc-50 hover:border-zinc-300 focus-within:border-zinc-300 focus-within:bg-white focus-within:ring-4 focus-within:ring-zinc-100 dark:border-zinc-800 dark:bg-zinc-900/70 dark:hover:border-zinc-700 dark:focus-within:border-zinc-700 dark:focus-within:bg-zinc-900 dark:focus-within:ring-zinc-800/60' + : pageCanvas && !compact + ? 'border-black/[0.09] bg-zinc-50 hover:border-black/[0.14] focus-within:border-black/[0.14] focus-within:bg-white focus-within:ring-4 focus-within:ring-black/[0.025] dark:border-white/[0.10] dark:bg-[#303842] dark:hover:border-white/[0.16] dark:focus-within:border-white/[0.16] dark:focus-within:bg-[#343d48] dark:focus-within:ring-white/[0.03]' + : 'border-zinc-200 bg-zinc-50 hover:border-zinc-300 focus-within:border-zinc-300 focus-within:bg-white focus-within:ring-4 focus-within:ring-zinc-100 dark:border-zinc-800 dark:bg-zinc-900/70 dark:hover:border-zinc-700 dark:focus-within:border-zinc-700 dark:focus-within:bg-zinc-900 dark:focus-within:ring-zinc-800/60' }`} > {/* Node reference chip */} @@ -3477,14 +3826,77 @@ export default function SessionChat({ {t('chat.upload.dropHint')}
)} -
+ {composerReferences.length > 0 && ( +
+ {composerReferences.map((reference) => { + const referenceStyle = reference.kind === 'subagent' + ? 'border-sky-200/80 bg-sky-50 text-sky-700 dark:border-sky-400/20 dark:bg-sky-400/[0.10] dark:text-sky-200' + : reference.kind === 'skill' + ? 'border-emerald-200/80 bg-emerald-50 text-emerald-700 dark:border-emerald-400/20 dark:bg-emerald-400/[0.10] dark:text-emerald-200' + : 'border-amber-200/80 bg-amber-50 text-amber-700 dark:border-amber-400/20 dark:bg-amber-400/[0.10] dark:text-amber-200'; + const ReferenceIcon = reference.kind === 'subagent' + ? Bot + : reference.kind === 'skill' + ? BookOpen + : WorkflowIcon; + return ( + + + + {t(`chat.references.${reference.kind}`)} + + {reference.value} + + + ); + })} +
+ )} +
0 ? 'pt-2' : 'pt-3'}`}>