diff --git a/src/fast_agent/cli/runtime/model_bootstrap.py b/src/fast_agent/cli/runtime/model_bootstrap.py index 8c2c6210b..1688f8f81 100644 --- a/src/fast_agent/cli/runtime/model_bootstrap.py +++ b/src/fast_agent/cli/runtime/model_bootstrap.py @@ -445,7 +445,8 @@ def activate_model_picker_provider(action: ProviderActivation) -> bool: from fast_agent.ui import console handler = get_oauth_provider(action.provider.config_name) - if handler.status().get("present"): + status = handler.status() + if status.get("present") and not status.get("expired"): return True typer.echo(f"Starting {handler.display_name} OAuth login...", err=True) diff --git a/src/fast_agent/history/process_poll_folding.py b/src/fast_agent/history/process_poll_folding.py index 066827da2..dca8c4da3 100644 --- a/src/fast_agent/history/process_poll_folding.py +++ b/src/fast_agent/history/process_poll_folding.py @@ -144,6 +144,10 @@ def _process_status(exchange: PollExchange) -> str | None: return status if isinstance(status, str) else None +def _has_resource_observation(exchange: PollExchange) -> bool: + return isinstance(exchange.metadata.get("resource_observation"), str) + + def _output_line_count(exchange: PollExchange) -> int | None: return _non_negative_int(exchange.metadata.get("output_line_count")) @@ -845,6 +849,8 @@ def fold_managed_process_poll_history( process_status = _process_status(current) if process_status not in _FOLDABLE_PROCESS_STATUSES: return None + if _has_resource_observation(current): + return None reverse_exchanges = [current] cursor = len(history) - 2 @@ -854,6 +860,8 @@ def fold_managed_process_poll_history( exchange = _exchange(request, result, request_index=cursor - 1) if exchange is None or exchange.process_id != current.process_id: break + if _has_resource_observation(exchange): + break reverse_exchanges.append(exchange) cursor -= 2 diff --git a/src/fast_agent/llm/provider/openai/codex_oauth.py b/src/fast_agent/llm/provider/openai/codex_oauth.py index ada0cc645..d7a7f39de 100644 --- a/src/fast_agent/llm/provider/openai/codex_oauth.py +++ b/src/fast_agent/llm/provider/openai/codex_oauth.py @@ -332,16 +332,16 @@ def _load_codex_tokens_with_source() -> tuple[CodexOAuthTokens | None, str | Non stored = load_oauth_credential("codex") return (_tokens_from_credential(stored.credential), stored.source) if stored else (None, None) + # Codex CLI credentials are external and read-only. Prefer them before touching + # the OS keyring so an existing auth.json is sufficient on its own. + tokens = _load_codex_cli_tokens() + if tokens: + return tokens, "auth.json" + stored = load_oauth_credential("codex") if stored: return _tokens_from_credential(stored.credential), stored.source - # Codex CLI credentials are an external, read-only fallback. Fast-agent-owned - # credentials take precedence so refreshed tokens can persist without modifying - # the CLI's auth.json. - tokens = _load_codex_cli_tokens() - if tokens: - return tokens, "auth.json" _warn_about_legacy_codex_keyring_credentials() return None, None diff --git a/src/fast_agent/llm/provider/openai/xai_oauth.py b/src/fast_agent/llm/provider/openai/xai_oauth.py index 183c9c713..48e54e54c 100644 --- a/src/fast_agent/llm/provider/openai/xai_oauth.py +++ b/src/fast_agent/llm/provider/openai/xai_oauth.py @@ -39,6 +39,10 @@ class XaiDeviceCode: interval: float +class _InvalidGrantError(ProviderKeyError): + pass + + def _oauth_error(action: str, response: httpx.Response) -> ProviderKeyError: try: payload = response.json() @@ -54,6 +58,15 @@ def _oauth_error(action: str, response: httpx.Response) -> ProviderKeyError: ) +def _oauth_error_code(response: httpx.Response) -> str | None: + try: + payload = response.json() + except ValueError: + return None + error = payload.get("error") if isinstance(payload, dict) else None + return error if isinstance(error, str) else None + + def _required_string(payload: dict[str, Any], field: str) -> str: value = payload.get(field) if not isinstance(value, str) or not value: @@ -190,7 +203,10 @@ def refresh_xai_credential( }, ) if not response.is_success: - raise _oauth_error("token refresh", response) + error = _oauth_error("token refresh", response) + if _oauth_error_code(response) == "invalid_grant": + raise _InvalidGrantError(error.message, error.details) + raise error payload = response.json() if not isinstance(payload, dict): raise ProviderKeyError("Invalid xAI OAuth response", "Expected a JSON object.") @@ -243,7 +259,11 @@ def get_xai_access_token( credential = current.credential expired = credential.expires_at is not None and time.time() >= credential.expires_at if force_refresh or expired: - credential = refresh_xai_credential(credential, client=client) + try: + credential = refresh_xai_credential(credential, client=client) + except _InvalidGrantError: + delete_oauth_credential(XAI_PROVIDER_ID) + return None save_oauth_credential( XAI_PROVIDER_ID, credential, diff --git a/src/fast_agent/tools/process_resources.py b/src/fast_agent/tools/process_resources.py new file mode 100644 index 000000000..b195c4610 --- /dev/null +++ b/src/fast_agent/tools/process_resources.py @@ -0,0 +1,510 @@ +"""Best-effort, bounded resource observations for managed shell processes.""" + +from __future__ import annotations + +import asyncio +import os +import queue +import shutil +import sys +import threading +import time +from dataclasses import dataclass, field +from pathlib import Path +from typing import TYPE_CHECKING, TypedDict + +if TYPE_CHECKING: + from collections.abc import Callable + +_GIB = 1024**3 +_SAMPLE_TIMEOUT_SECONDS = 0.05 +_DISK_LOW_RATIO = 0.20 +_DISK_RECOVERED_RATIO = 0.25 +_DISK_LOW_BYTES = 2 * _GIB +_DISK_RECOVERED_BYTES = 3 * _GIB +_MEMORY_HIGH_RATIO = 0.80 +_MEMORY_RECOVERED_RATIO = 0.75 +_CPU_HIGH_RATIO = 0.90 +_CPU_RECOVERED_RATIO = 0.75 +_LARGE_BYTE_CHANGE = _GIB +_LARGE_DISK_RATIO_CHANGE = 0.10 + + +class ProcessResourceSnapshotMetadata(TypedDict, total=False): + """Serializable resource fields attached to managed-process metadata.""" + + sampled_at: float + disk_total_bytes: int + disk_free_bytes: int + memory_used_bytes: int + memory_limit_bytes: int + process_rss_bytes: int + process_cpu_seconds: float + process_count: int + cpu_capacity: float + + +@dataclass(frozen=True, slots=True) +class ProcessResourceSnapshot: + sampled_at: float + disk_total_bytes: int | None = None + disk_free_bytes: int | None = None + memory_used_bytes: int | None = None + memory_limit_bytes: int | None = None + process_rss_bytes: int | None = None + process_cpu_seconds: float | None = None + process_count: int | None = None + cpu_capacity: float | None = None + + def metadata(self) -> ProcessResourceSnapshotMetadata: + metadata: ProcessResourceSnapshotMetadata = { + "sampled_at": self.sampled_at, + } + if self.disk_total_bytes is not None: + metadata["disk_total_bytes"] = self.disk_total_bytes + if self.disk_free_bytes is not None: + metadata["disk_free_bytes"] = self.disk_free_bytes + if self.memory_used_bytes is not None: + metadata["memory_used_bytes"] = self.memory_used_bytes + if self.memory_limit_bytes is not None: + metadata["memory_limit_bytes"] = self.memory_limit_bytes + if self.process_rss_bytes is not None: + metadata["process_rss_bytes"] = self.process_rss_bytes + if self.process_cpu_seconds is not None: + metadata["process_cpu_seconds"] = self.process_cpu_seconds + if self.process_count is not None: + metadata["process_count"] = self.process_count + if self.cpu_capacity is not None: + metadata["cpu_capacity"] = self.cpu_capacity + return metadata + + +@dataclass(slots=True) +class ProcessResourceObservationState: + baseline: ProcessResourceSnapshot | None = None + previous: ProcessResourceSnapshot | None = None + last_reported: ProcessResourceSnapshot | None = None + active_warnings: set[str] = field(default_factory=set) + + +def _read_int(path: Path) -> int | None: + try: + value = path.read_text(encoding="utf-8").strip() + return int(value) + except (OSError, ValueError): + return None + + +def _process_cgroup_path(pid: int) -> Path | None: + try: + lines = Path(f"/proc/{pid}/cgroup").read_text(encoding="utf-8").splitlines() + except OSError: + return None + for line in lines: + fields = line.split(":", 2) + if len(fields) == 3 and fields[0] == "0" and fields[1] == "": + return Path("/sys/fs/cgroup") / fields[2].lstrip("/") + return None + + +def _memory_metrics(pid: int) -> tuple[int | None, int | None]: + cgroup_path = _process_cgroup_path(pid) + if cgroup_path is not None: + used = _read_int(cgroup_path / "memory.current") + try: + raw_limit = (cgroup_path / "memory.max").read_text(encoding="utf-8").strip() + limit = None if raw_limit == "max" else int(raw_limit) + except (OSError, ValueError): + limit = None + if used is not None and limit is not None and limit > 0: + return used, limit + + try: + lines = Path("/proc/meminfo").read_text(encoding="utf-8").splitlines() + except OSError: + return None, None + values: dict[str, int] = {} + for line in lines: + name, separator, raw_value = line.partition(":") + if not separator: + continue + number = raw_value.strip().split(maxsplit=1)[0] + try: + values[name] = int(number) * 1024 + except ValueError: + continue + total = values.get("MemTotal") + available = values.get("MemAvailable") + if total is None or available is None: + return None, None + return max(total - available, 0), total + + +def _cpu_capacity(pid: int) -> float: + cgroup_path = _process_cgroup_path(pid) + if cgroup_path is not None: + try: + quota_text, period_text = ( + cgroup_path / "cpu.max" + ).read_text(encoding="utf-8").split(maxsplit=1) + if quota_text != "max": + quota = int(quota_text) + period = int(period_text) + if quota > 0 and period > 0: + return quota / period + except (OSError, ValueError): + pass + return float(os.cpu_count() or 1) + + +def _process_children(pid: int) -> tuple[int, ...]: + try: + raw = Path(f"/proc/{pid}/task/{pid}/children").read_text(encoding="utf-8") + except OSError: + return () + children: list[int] = [] + for token in raw.split(): + try: + children.append(int(token)) + except ValueError: + continue + return tuple(children) + + +def _process_tree(pid: int) -> tuple[int, ...]: + pending = [pid] + observed: set[int] = set() + while pending: + current = pending.pop() + if current in observed: + continue + observed.add(current) + pending.extend(_process_children(current)) + return tuple(observed) + + +def _process_stat(pid: int) -> tuple[float, int] | None: + try: + raw = Path(f"/proc/{pid}/stat").read_text(encoding="utf-8") + except OSError: + return None + command_end = raw.rfind(")") + if command_end < 0: + return None + fields = raw[command_end + 2 :].split() + if len(fields) < 22: + return None + try: + ticks = int(fields[11]) + int(fields[12]) + rss_pages = int(fields[21]) + clock_ticks = os.sysconf("SC_CLK_TCK") + page_size = os.sysconf("SC_PAGE_SIZE") + except (OSError, ValueError): + return None + return ticks / clock_ticks, max(rss_pages, 0) * page_size + + +def _linux_process_metrics( + pid: int, +) -> tuple[ + int | None, + int | None, + float | None, + int | None, + int | None, + float | None, +]: + memory_used, memory_limit = _memory_metrics(pid) + total_cpu = 0.0 + total_rss = 0 + process_count = 0 + for process_id in _process_tree(pid): + stat = _process_stat(process_id) + if stat is None: + continue + cpu_seconds, rss_bytes = stat + total_cpu += cpu_seconds + total_rss += rss_bytes + process_count += 1 + return ( + memory_used, + memory_limit, + total_cpu if process_count else None, + total_rss if process_count else None, + process_count if process_count else None, + _cpu_capacity(pid), + ) + + +def _collect_process_resource_snapshot( + working_directory: str, + pid: int | None, + *, + platform_name: str = sys.platform, +) -> ProcessResourceSnapshot: + disk_total: int | None = None + disk_free: int | None = None + try: + disk = shutil.disk_usage(working_directory) + disk_total = disk.total + disk_free = disk.free + except OSError: + pass + + memory_used: int | None = None + memory_limit: int | None = None + process_cpu: float | None = None + process_rss: int | None = None + process_count: int | None = None + cpu_capacity: float | None = None + if pid is not None and platform_name.startswith("linux"): + ( + memory_used, + memory_limit, + process_cpu, + process_rss, + process_count, + cpu_capacity, + ) = _linux_process_metrics(pid) + + return ProcessResourceSnapshot( + sampled_at=time.monotonic(), + disk_total_bytes=disk_total, + disk_free_bytes=disk_free, + memory_used_bytes=memory_used, + memory_limit_bytes=memory_limit, + process_rss_bytes=process_rss, + process_cpu_seconds=process_cpu, + process_count=process_count, + cpu_capacity=cpu_capacity, + ) + + +@dataclass(frozen=True, slots=True) +class _SampleRequest: + working_directory: str + pid: int | None + loop: asyncio.AbstractEventLoop + future: asyncio.Future[ProcessResourceSnapshot | None] + + +class _ProcessResourceSampler: + """Run observations on one daemon worker so a blocked syscall cannot block exit.""" + + def __init__( + self, + collector: Callable[[str, int | None], ProcessResourceSnapshot] | None = None, + ) -> None: + self._collector = collector or _collect_process_resource_snapshot + self._requests: queue.Queue[_SampleRequest] = queue.Queue(maxsize=1) + self._thread: threading.Thread | None = None + self._thread_lock = threading.Lock() + + def _ensure_started(self) -> None: + with self._thread_lock: + if self._thread is not None: + return + self._thread = threading.Thread( + target=self._run, + name="fast-agent-resource-observer", + daemon=True, + ) + self._thread.start() + + @staticmethod + def _deliver( + future: asyncio.Future[ProcessResourceSnapshot | None], + snapshot: ProcessResourceSnapshot | None, + ) -> None: + if not future.done(): + future.set_result(snapshot) + + def _run(self) -> None: + while True: + request = self._requests.get() + try: + snapshot = self._collector(request.working_directory, request.pid) + except Exception: + snapshot = None + try: + request.loop.call_soon_threadsafe( + self._deliver, + request.future, + snapshot, + ) + except RuntimeError: + continue + + async def sample( + self, + working_directory: str, + pid: int | None, + *, + timeout_seconds: float = _SAMPLE_TIMEOUT_SECONDS, + ) -> ProcessResourceSnapshot | None: + self._ensure_started() + loop = asyncio.get_running_loop() + future: asyncio.Future[ProcessResourceSnapshot | None] = loop.create_future() + try: + self._requests.put_nowait( + _SampleRequest( + working_directory=working_directory, + pid=pid, + loop=loop, + future=future, + ) + ) + except queue.Full: + return None + try: + async with asyncio.timeout(timeout_seconds): + return await future + except TimeoutError: + return None + + +_SAMPLER = _ProcessResourceSampler() + + +async def sample_process_resources( + working_directory: str, + pid: int | None, +) -> ProcessResourceSnapshot | None: + """Return a resource snapshot or None within a small fixed time budget.""" + return await _SAMPLER.sample(working_directory, pid) + + +def _gib(value: int) -> str: + return f"{value / _GIB:.1f} GiB" + + +def _disk_observation( + state: ProcessResourceObservationState, + snapshot: ProcessResourceSnapshot, + comparison: ProcessResourceSnapshot, +) -> str | None: + total = snapshot.disk_total_bytes + free = snapshot.disk_free_bytes + if total is None or free is None or total <= 0: + return None + ratio = free / total + pressure = ratio <= _DISK_LOW_RATIO or free <= _DISK_LOW_BYTES + newly_low = pressure and "disk" not in state.active_warnings + if pressure: + state.active_warnings.add("disk") + elif ratio >= _DISK_RECOVERED_RATIO and free >= _DISK_RECOVERED_BYTES: + state.active_warnings.discard("disk") + + prior_free = comparison.disk_free_bytes + decline = max(prior_free - free, 0) if prior_free is not None else 0 + large_decline = decline >= _LARGE_BYTE_CHANGE or decline / total >= _LARGE_DISK_RATIO_CHANGE + if not newly_low and not large_decline: + return None + detail = f"disk free {_gib(free)}/{_gib(total)} ({ratio:.0%}" + if decline: + detail += f", down {_gib(decline)}" + return f"{detail})" + + +def _memory_observation( + state: ProcessResourceObservationState, + snapshot: ProcessResourceSnapshot, +) -> str | None: + used = snapshot.memory_used_bytes + limit = snapshot.memory_limit_bytes + if used is None or limit is None or limit <= 0: + return None + ratio = used / limit + pressure = ratio >= _MEMORY_HIGH_RATIO + newly_high = pressure and "memory" not in state.active_warnings + if pressure: + state.active_warnings.add("memory") + elif ratio <= _MEMORY_RECOVERED_RATIO: + state.active_warnings.discard("memory") + if not newly_high: + return None + return f"memory used {_gib(used)}/{_gib(limit)} ({ratio:.0%})" + + +def _rss_observation( + snapshot: ProcessResourceSnapshot, + comparison: ProcessResourceSnapshot, +) -> str | None: + rss = snapshot.process_rss_bytes + prior_rss = comparison.process_rss_bytes + if rss is None or prior_rss is None: + return None + growth = rss - prior_rss + if growth < _LARGE_BYTE_CHANGE: + return None + process_label = ( + f", {snapshot.process_count} processes" + if snapshot.process_count is not None + else "" + ) + return f"process tree RSS {_gib(rss)} (up {_gib(growth)}{process_label})" + + +def _cpu_observation( + state: ProcessResourceObservationState, + snapshot: ProcessResourceSnapshot, +) -> str | None: + previous = state.previous + if ( + previous is None + or previous.process_cpu_seconds is None + or snapshot.process_cpu_seconds is None + or snapshot.cpu_capacity is None + or snapshot.cpu_capacity <= 0 + ): + return None + elapsed = snapshot.sampled_at - previous.sampled_at + if elapsed <= 0: + return None + cpu_ratio = max(snapshot.process_cpu_seconds - previous.process_cpu_seconds, 0) / ( + elapsed * snapshot.cpu_capacity + ) + newly_high = cpu_ratio >= _CPU_HIGH_RATIO and "cpu" not in state.active_warnings + if cpu_ratio >= _CPU_HIGH_RATIO: + state.active_warnings.add("cpu") + elif cpu_ratio <= _CPU_RECOVERED_RATIO: + state.active_warnings.discard("cpu") + if not newly_high: + return None + process_label = ( + f", {snapshot.process_count} processes" + if snapshot.process_count is not None + else "" + ) + return ( + f"process tree CPU {cpu_ratio:.0%} of {snapshot.cpu_capacity:.1f}-core capacity" + f"{process_label}" + ) + + +def observe_resource_changes( + state: ProcessResourceObservationState, + snapshot: ProcessResourceSnapshot, +) -> str | None: + """Return one compact observation only for warnings or large changes.""" + if state.baseline is None: + state.baseline = snapshot + state.previous = snapshot + return None + + comparison = state.last_reported or state.baseline + observations = [ + observation + for observation in ( + _disk_observation(state, snapshot, comparison), + _memory_observation(state, snapshot), + _rss_observation(snapshot, comparison), + _cpu_observation(state, snapshot), + ) + if observation is not None + ] + state.previous = snapshot + if not observations: + return None + state.last_reported = snapshot + return "; ".join(observations) diff --git a/src/fast_agent/tools/shell_runtime.py b/src/fast_agent/tools/shell_runtime.py index ce380d023..d95627af6 100644 --- a/src/fast_agent/tools/shell_runtime.py +++ b/src/fast_agent/tools/shell_runtime.py @@ -44,6 +44,13 @@ ) from fast_agent.tools.local_shell_executor import LocalShellExecutor from fast_agent.tools.output_truncation import format_output_truncation_notice +from fast_agent.tools.process_resources import ( + ProcessResourceObservationState, + ProcessResourceSnapshot, + ProcessResourceSnapshotMetadata, + observe_resource_changes, + sample_process_resources, +) from fast_agent.tools.tool_sources import SHELL_TOOL_SOURCE, set_tool_source from fast_agent.ui import console from fast_agent.ui.console_display import ConsoleDisplay @@ -85,6 +92,7 @@ def _default_max_process_poll_seconds() -> int: _TERMINATE_PROCESS_ARGUMENTS = frozenset({"process_id"}) _PROCESS_OUTPUT_DEBOUNCE_SECONDS = 2.0 _PROCESS_PROGRESS_EMIT_INTERVAL_SECONDS = 1.0 +_RESOURCE_OBSERVATION_TIMEOUT_SECONDS = 0.075 def _text_result(message: str, *, is_error: bool) -> CallToolResult: @@ -112,6 +120,8 @@ class ProcessResultMetadata(TypedDict, total=False): poll_wake_on_output: bool poll_elapsed_seconds: float poll_deadline_overshoot_seconds: float + resource_snapshot: ProcessResourceSnapshotMetadata + resource_observation: str def process_result_metadata(result: CallToolResult) -> ProcessResultMetadata | None: @@ -185,7 +195,11 @@ class _ShellRuntimeCallbacks: async def on_started(self, process_id: int | None) -> None: self.os_process_id = process_id - self.started_event.set() + try: + if self.process is not None: + await self.runtime._capture_process_resource_baseline(self.process) + finally: + self.started_event.set() async def on_stdout(self, text: str) -> None: self.runtime._record_stream_output( @@ -277,6 +291,9 @@ class _ManagedShellProcess: terminated: bool = False buffered_result_recorded: bool = False active_poll: _ActiveProcessPoll | None = None + resource_observations: ProcessResourceObservationState = field( + default_factory=ProcessResourceObservationState + ) @dataclass(frozen=True, slots=True) @@ -355,6 +372,7 @@ def __init__( process_poll_default_wait_seconds, self._max_process_poll_seconds, ) + self._resource_observations_enabled = self.runtime_info().kind == "local" if self.enabled: # Detect the shell early so we can include it in the tool description @@ -1493,6 +1511,27 @@ async def _get_managed_process(self, process_id: str) -> _ManagedShellProcess | async with self._processes_lock: return self._managed_processes.get(process_id) + async def _sample_managed_process_resources( + self, + process: _ManagedShellProcess, + ) -> ProcessResourceSnapshot | None: + if not self._resource_observations_enabled: + return None + pid = process.callbacks.os_process_id if not process.task.done() else None + try: + async with asyncio.timeout(_RESOURCE_OBSERVATION_TIMEOUT_SECONDS): + return await sample_process_resources(process.working_directory, pid) + except Exception: + return None + + async def _capture_process_resource_baseline( + self, + process: _ManagedShellProcess, + ) -> None: + snapshot = await self._sample_managed_process_resources(process) + if snapshot is not None: + observe_resource_changes(process.resource_observations, snapshot) + async def _wait_for_initial_process_result( self, process: _ManagedShellProcess, @@ -1584,6 +1623,16 @@ def _append_poll_output_activity( block.text = f"{block.text}\n{line}" return + @staticmethod + def _append_resource_observation( + result: CallToolResult, + observation: str, + ) -> None: + for block in result.content: + if isinstance(block, TextContent): + block.text = f"{block.text}\nresource_observation: {observation}" + return + def _managed_process_result( self, process: _ManagedShellProcess, @@ -1778,6 +1827,9 @@ async def poll_process( with suppress(asyncio.CancelledError): await active_poll.heartbeat_task + poll_elapsed_seconds = time.monotonic() - poll_started_at + resource_snapshot = await self._sample_managed_process_resources(process) + async with process.lock: if process.task.done(): poll_yield_reason = "completion" @@ -1813,13 +1865,21 @@ async def poll_process( metadata["has_observed_output"] = output_observed metadata["poll_wait_sec"] = parsed.wait_sec metadata["poll_wake_on_output"] = parsed.wake_on_output - poll_elapsed_seconds = time.monotonic() - poll_started_at metadata["poll_elapsed_seconds"] = poll_elapsed_seconds if poll_yield_reason == "deadline": metadata["poll_deadline_overshoot_seconds"] = max( poll_elapsed_seconds - parsed.wait_sec, 0.0, ) + if resource_snapshot is not None: + metadata["resource_snapshot"] = resource_snapshot.metadata() + observation = observe_resource_changes( + process.resource_observations, + resource_snapshot, + ) + if observation is not None: + metadata["resource_observation"] = observation + self._append_resource_observation(result, observation) self._append_poll_output_activity( result, output_bytes=output_bytes_since_last_poll, diff --git a/tests/unit/fast_agent/commands/test_runtime_model_picker_bootstrap.py b/tests/unit/fast_agent/commands/test_runtime_model_picker_bootstrap.py index 1764ab4fd..67c4152b5 100644 --- a/tests/unit/fast_agent/commands/test_runtime_model_picker_bootstrap.py +++ b/tests/unit/fast_agent/commands/test_runtime_model_picker_bootstrap.py @@ -41,6 +41,7 @@ from fast_agent.ui.model_picker import ModelPickerResult from fast_agent.ui.model_picker_common import ( LLAMACPP_PROVIDER_KEY, + ProviderActivation, normalize_generic_model_spec, ) @@ -495,6 +496,49 @@ async def fake_run_model_picker_async(**kwargs): assert selected == "haikutiny" +@pytest.mark.asyncio +async def test_select_model_from_picker_reauthenticates_expired_oauth_provider( + monkeypatch, +) -> None: + request = _make_request() + logins = 0 + + async def fake_run_model_picker_async(**kwargs): + del kwargs + return ModelPickerResult( + provider=Provider.XAI.config_name, + provider_available=False, + selected_model="xai.grok-4", + resolved_model="xai.grok-4", + source="curated", + refer_to_docs=False, + activation_action=ProviderActivation(Provider.XAI), + ) + + def login() -> None: + nonlocal logins + logins += 1 + + handler = SimpleNamespace( + display_name="xAI", + status=lambda: {"present": True, "expired": True}, + login=login, + ) + monkeypatch.setattr( + "fast_agent.ui.model_picker.run_model_picker_async", + fake_run_model_picker_async, + ) + monkeypatch.setattr( + "fast_agent.auth.providers.get_oauth_provider", + lambda provider: handler, + ) + + selected = await _select_model_from_picker(request, config_payload={}) + + assert selected == "xai.grok-4" + assert logins == 1 + + @pytest.mark.asyncio async def test_select_model_from_picker_passes_config_start_path( monkeypatch, tmp_path: Path diff --git a/tests/unit/fast_agent/history/test_process_poll_folding.py b/tests/unit/fast_agent/history/test_process_poll_folding.py index 47bd85c8c..05d2caabc 100644 --- a/tests/unit/fast_agent/history/test_process_poll_folding.py +++ b/tests/unit/fast_agent/history/test_process_poll_folding.py @@ -509,6 +509,48 @@ def test_quiet_running_process_folds_earlier_polls() -> None: assert "Current status: running." in first.text +def test_resource_observation_poll_is_not_folded() -> None: + history = _history_before_terminal(4) + warning = _poll_result(4, output_line_count=0) + _update_result_metadata( + warning, + resource_observation="disk free 1.0 GiB/10.0 GiB (10%)", + ) + + assert fold_completed_process_poll_history(history, warning) is None + + +def test_quiet_polls_after_resource_observation_fold_without_removing_warning() -> None: + warning = _poll_result(1, output_line_count=0) + _update_result_metadata( + warning, + resource_observation="disk free 1.0 GiB/10.0 GiB (10%)", + ) + history = [ + PromptMessageExtended( + role="user", + content=[TextContent(type="text", text="build the project")], + ), + _poll_request(1), + warning, + ] + for index in range(2, 5): + history.extend( + [_poll_request(index), _poll_result(index, output_line_count=0)] + ) + history.append(_poll_request(5)) + + folded = fold_completed_process_poll_history( + history, + _poll_result(5, output_line_count=0), + ) + + assert folded is not None + assert warning in folded.history + assert len(folded.history) == 4 + assert folded.metadata["polls_folded"] == 3 + + def test_narrated_poll_suffix_folds_and_preserves_removed_updates_exactly() -> None: history = _history_before_terminal(4) narrations = [ diff --git a/tests/unit/fast_agent/llm/providers/test_codex_oauth.py b/tests/unit/fast_agent/llm/providers/test_codex_oauth.py index 1b852b2ab..2bb147c6e 100644 --- a/tests/unit/fast_agent/llm/providers/test_codex_oauth.py +++ b/tests/unit/fast_agent/llm/providers/test_codex_oauth.py @@ -33,7 +33,11 @@ def test_explicit_auth_json_path_overrides_keyring(monkeypatch, tmp_path: Path) monkeypatch.delenv("FAST_AGENT_AUTH_FILE", raising=False) monkeypatch.setenv("CODEX_AUTH_JSON_PATH", str(auth_path)) monkeypatch.delenv("CODEX_HOME", raising=False) - monkeypatch.setattr(codex_oauth, "load_oauth_credential", lambda provider: None) + monkeypatch.setattr( + codex_oauth, + "load_oauth_credential", + lambda provider: pytest.fail("auth.json must not access the fast-agent credential store"), + ) tokens, source = codex_oauth._load_codex_tokens_with_source() @@ -63,7 +67,11 @@ def test_default_auth_json_path_is_used_when_fast_agent_store_is_empty( monkeypatch.delenv("CODEX_AUTH_JSON_PATH", raising=False) monkeypatch.delenv("CODEX_HOME", raising=False) monkeypatch.setattr(codex_oauth.Path, "home", lambda: tmp_path) - monkeypatch.setattr(codex_oauth, "load_oauth_credential", lambda provider: None) + monkeypatch.setattr( + codex_oauth, + "load_oauth_credential", + lambda provider: pytest.fail("auth.json must not access the fast-agent credential store"), + ) tokens, source = codex_oauth._load_codex_tokens_with_source() @@ -72,7 +80,7 @@ def test_default_auth_json_path_is_used_when_fast_agent_store_is_empty( assert tokens.access_token == "cli-token" -def test_fast_agent_owned_credential_precedes_codex_cli(monkeypatch) -> None: +def test_codex_cli_credential_precedes_fast_agent_owned_credential(monkeypatch) -> None: monkeypatch.delenv("FAST_AGENT_AUTH_FILE", raising=False) monkeypatch.setattr( codex_oauth, @@ -89,9 +97,9 @@ def test_fast_agent_owned_credential_precedes_codex_cli(monkeypatch) -> None: tokens, source = codex_oauth._load_codex_tokens_with_source() - assert source == "keyring" + assert source == "auth.json" assert tokens is not None - assert tokens.access_token == "fast-agent-token" + assert tokens.access_token == "cli-token" def test_fast_agent_auth_file_is_authoritative_over_codex_cli( diff --git a/tests/unit/fast_agent/llm/providers/test_xai_oauth.py b/tests/unit/fast_agent/llm/providers/test_xai_oauth.py index 68267136e..de17671f8 100644 --- a/tests/unit/fast_agent/llm/providers/test_xai_oauth.py +++ b/tests/unit/fast_agent/llm/providers/test_xai_oauth.py @@ -109,3 +109,34 @@ def test_expired_xai_credential_refreshes_and_preserves_refresh_token( assert access_token == "refreshed-access" assert simulator.refreshes == 1 assert document["providers"]["xai"]["refresh_token"] == "keep-refresh" + + +def test_revoked_xai_refresh_token_is_cleared( + monkeypatch, + tmp_path: Path, +) -> None: + auth_path = tmp_path / "xai.auth.json" + monkeypatch.setenv("FAST_AGENT_AUTH_FILE", str(auth_path)) + save_oauth_credential( + "xai", + OAuthCredential( + access_token="expired", + refresh_token="revoked", + expires_at=time.time() - 1, + ), + ) + + def revoked_refresh(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 400, + json={ + "error": "invalid_grant", + "error_description": "Refresh token has been revoked.", + }, + ) + + with httpx.Client(transport=httpx.MockTransport(revoked_refresh)) as client: + assert get_xai_access_token(client=client) is None + + document = json.loads(auth_path.read_text()) + assert "xai" not in document["providers"] diff --git a/tests/unit/fast_agent/tools/test_process_resources.py b/tests/unit/fast_agent/tools/test_process_resources.py new file mode 100644 index 000000000..2d1f42fa2 --- /dev/null +++ b/tests/unit/fast_agent/tools/test_process_resources.py @@ -0,0 +1,157 @@ +import time +from pathlib import Path + +import pytest + +import fast_agent.tools.process_resources as process_resources +from fast_agent.tools.process_resources import ( + ProcessResourceObservationState, + ProcessResourceSnapshot, + observe_resource_changes, +) + +GIB = 1024**3 + + +def _snapshot( + *, + sampled_at: float = 1.0, + disk_total: int = 10 * GIB, + disk_free: int = 8 * GIB, + memory_used: int = 2 * GIB, + memory_limit: int = 10 * GIB, + rss: int = GIB, + cpu: float = 1.0, + process_count: int = 1, + cpu_capacity: float = 2.0, +) -> ProcessResourceSnapshot: + return ProcessResourceSnapshot( + sampled_at=sampled_at, + disk_total_bytes=disk_total, + disk_free_bytes=disk_free, + memory_used_bytes=memory_used, + memory_limit_bytes=memory_limit, + process_rss_bytes=rss, + process_cpu_seconds=cpu, + process_count=process_count, + cpu_capacity=cpu_capacity, + ) + + +def test_resource_observations_are_silent_for_small_changes() -> None: + state = ProcessResourceObservationState() + + assert observe_resource_changes(state, _snapshot()) is None + assert ( + observe_resource_changes( + state, + _snapshot( + sampled_at=2.0, + disk_free=15 * GIB // 2, + memory_used=3 * GIB, + rss=GIB + 100, + cpu=1.5, + ), + ) + is None + ) + + +def test_resource_observation_reports_large_disk_decline_once() -> None: + state = ProcessResourceObservationState() + observe_resource_changes(state, _snapshot()) + + observation = observe_resource_changes( + state, + _snapshot(sampled_at=2.0, disk_free=6 * GIB, cpu=1.2), + ) + repeated = observe_resource_changes( + state, + _snapshot(sampled_at=3.0, disk_free=6 * GIB, cpu=1.4), + ) + + assert observation is not None + assert "disk free 6.0 GiB/10.0 GiB" in observation + assert "down 2.0 GiB" in observation + assert repeated is None + + +def test_resource_observation_reports_pressure_and_high_cpu() -> None: + state = ProcessResourceObservationState() + observe_resource_changes(state, _snapshot()) + + observation = observe_resource_changes( + state, + _snapshot( + sampled_at=2.0, + disk_free=GIB, + memory_used=9 * GIB, + cpu=3.0, + process_count=5, + ), + ) + + assert observation is not None + assert "disk free 1.0 GiB/10.0 GiB (10%" in observation + assert "memory used 9.0 GiB/10.0 GiB (90%)" in observation + assert "process tree CPU 100% of 2.0-core capacity, 5 processes" in observation + + +@pytest.mark.asyncio +async def test_sampler_timeout_never_waits_for_blocked_collector() -> None: + def blocked_collector( + working_directory: str, + pid: int | None, + ) -> ProcessResourceSnapshot: + del working_directory, pid + while True: + time.sleep(1) + + sampler = process_resources._ProcessResourceSampler(blocked_collector) + started = time.monotonic() + + snapshot = await sampler.sample("/workspace", 1, timeout_seconds=0.01) + + assert snapshot is None + assert time.monotonic() - started < 0.2 + + +@pytest.mark.asyncio +async def test_sampler_swallows_collector_exception() -> None: + def failed_collector( + working_directory: str, + pid: int | None, + ) -> ProcessResourceSnapshot: + del working_directory, pid + raise OSError("unavailable") + + sampler = process_resources._ProcessResourceSampler(failed_collector) + + assert await sampler.sample("/workspace", 1, timeout_seconds=0.1) is None + + +def test_windows_collection_has_disk_only_fallback( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + monkeypatch.setattr( + process_resources.shutil, + "disk_usage", + lambda path: process_resources.shutil._ntuple_diskusage( + 10 * GIB, + 4 * GIB, + 6 * GIB, + ), + ) + + snapshot = process_resources._collect_process_resource_snapshot( + str(tmp_path), + 1234, + platform_name="win32", + ) + + assert snapshot.disk_total_bytes == 10 * GIB + assert snapshot.disk_free_bytes == 6 * GIB + assert snapshot.process_rss_bytes is None + assert snapshot.process_cpu_seconds is None + assert snapshot.memory_limit_bytes is None diff --git a/tests/unit/fast_agent/tools/test_shell_runtime.py b/tests/unit/fast_agent/tools/test_shell_runtime.py index def97fa35..dbf5af31f 100644 --- a/tests/unit/fast_agent/tools/test_shell_runtime.py +++ b/tests/unit/fast_agent/tools/test_shell_runtime.py @@ -32,6 +32,7 @@ ShellRuntimeInfo, ) from fast_agent.tools.local_shell_executor import LocalShellExecutor +from fast_agent.tools.process_resources import ProcessResourceSnapshot from fast_agent.tools.shell_runtime import ShellRuntime from fast_agent.ui import console from fast_agent.ui.display_suppression import suppress_interactive_display @@ -288,6 +289,11 @@ async def close(self) -> None: return None +class _LocalManagedShellEnvironment(_ManagedShellEnvironment): + def runtime_info(self) -> ShellRuntimeInfo: + return ShellRuntimeInfo(name="bash", kind="local", provider="managed-test") + + class _ActiveManagedShellEnvironment(_ManagedShellEnvironment): async def execute( self, @@ -1051,6 +1057,82 @@ async def test_silent_command_yields_alive_then_poll_reports_completion() -> Non assert getattr(poll_result, "output_line_count", None) == 1 +@pytest.mark.asyncio +async def test_poll_resource_warning_is_anchored_to_same_result( + monkeypatch: pytest.MonkeyPatch, +) -> None: + snapshots = iter( + [ + ProcessResourceSnapshot( + sampled_at=1.0, + disk_total_bytes=10 * 1024**3, + disk_free_bytes=8 * 1024**3, + ), + ProcessResourceSnapshot( + sampled_at=2.0, + disk_total_bytes=10 * 1024**3, + disk_free_bytes=1024**3, + ), + ] + ) + + async def sample(working_directory: str, pid: int | None): + del working_directory, pid + return next(snapshots) + + monkeypatch.setattr(shell_runtime_module, "sample_process_resources", sample) + environment = _LocalManagedShellEnvironment() + runtime = ShellRuntime( + activation_reason="test", + logger=logging.getLogger("shell-runtime-test"), + shell_environment=environment, + ) + await runtime.execute({"command": "large install", "background": True}) + + result = await runtime.poll_process({"process_id": "process-1"}) + + assert result.content is not None + assert isinstance(result.content[0], TextContent) + assert "resource_observation: disk free 1.0 GiB/10.0 GiB" in result.content[0].text + metadata = (result.meta or {})[FAST_AGENT_SHELL_PROCESS_METADATA] + assert metadata["resource_observation"].startswith("disk free 1.0 GiB") + assert metadata["resource_snapshot"]["disk_free_bytes"] == 1024**3 + await runtime.close() + + +@pytest.mark.asyncio +async def test_resource_sampler_timeout_and_error_do_not_delay_poll( + monkeypatch: pytest.MonkeyPatch, +) -> None: + calls = 0 + + async def sample(working_directory: str, pid: int | None): + nonlocal calls + del working_directory, pid + calls += 1 + if calls == 1: + raise OSError("sampling unavailable") + await asyncio.Event().wait() + + monkeypatch.setattr(shell_runtime_module, "sample_process_resources", sample) + environment = _LocalManagedShellEnvironment() + runtime = ShellRuntime( + activation_reason="test", + logger=logging.getLogger("shell-runtime-test"), + shell_environment=environment, + ) + await runtime.execute({"command": "large install", "background": True}) + started = time.monotonic() + + result = await runtime.poll_process({"process_id": "process-1"}) + + assert result.isError is False + assert time.monotonic() - started < 0.2 + metadata = (result.meta or {})[FAST_AGENT_SHELL_PROCESS_METADATA] + assert "resource_snapshot" not in metadata + await runtime.close() + + @pytest.mark.asyncio async def test_continuous_output_still_yields_at_foreground_ceiling() -> None: environment = _ActiveManagedShellEnvironment()